From 7638bf3b2bce1a0196aac6a6fc6aa8b4a86f9288 Mon Sep 17 00:00:00 2001 From: James Demlow Date: Fri, 17 Jan 2025 11:39:16 -0700 Subject: [PATCH] Kff dag (#3915) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit # Description and Motivation closes #3794 ## Has this been tested? How? Using infra-test ## Screenshots (if appropriate) ![Screenshot (146)](https://github.com/user-attachments/assets/e7f401bc-60b6-4680-80c1-a5243bd5827c) Screenshot 2025-01-17 at 11 24 49 AM Screenshot 2025-01-17 at 11 26 00 AM ## Types of changes (leave all that apply) - Bug fix - Refactor / chore ## New frontend preview link is below in the Netlify comment 😎 --------- Co-authored-by: Ben Hammond --- airflow/dags/kff_vaccination.py | 12 ++++++------ python/datasources/kff_vaccination.py | 8 -------- python/tests/datasources/test_kff_vaccination.py | 4 +--- 3 files changed, 7 insertions(+), 17 deletions(-) diff --git a/airflow/dags/kff_vaccination.py b/airflow/dags/kff_vaccination.py index e4a23f4f9e..71a40290eb 100644 --- a/airflow/dags/kff_vaccination.py +++ b/airflow/dags/kff_vaccination.py @@ -27,15 +27,15 @@ "kff_vaccination_to_bq", kff_vaccination_bq_payload, data_ingestion_dag ) -kff_vaccination_exporter_payload_race = {"dataset_name": _KFF_VACCINATION_DATASET_NAME, "demographic": "race"} +kff_vaccination_exporter_payload_race = { + "dataset_name": _KFF_VACCINATION_DATASET_NAME, + "should_export_as_alls": True, + "demographic": "race_and_ethnicity", +} kff_vaccination_exporter_operator_race = util.create_exporter_operator( "kff_vaccination_exporter_race", kff_vaccination_exporter_payload_race, data_ingestion_dag ) -kff_vaccination_exporter_payload_alls = {"dataset_name": _KFF_VACCINATION_DATASET_NAME, "demographic": "alls"} -kff_vaccination_exporter_operator_alls = util.create_exporter_operator( - "kff_vaccination_exporter_alls", kff_vaccination_exporter_payload_alls, data_ingestion_dag -) # Ingestion DAG -(kff_vaccination_bq_operator >> kff_vaccination_exporter_operator_race >> kff_vaccination_exporter_operator_alls) +(kff_vaccination_bq_operator >> kff_vaccination_exporter_operator_race) diff --git a/python/datasources/kff_vaccination.py b/python/datasources/kff_vaccination.py index ad537b20b0..c9c736c49e 100644 --- a/python/datasources/kff_vaccination.py +++ b/python/datasources/kff_vaccination.py @@ -316,14 +316,6 @@ def write_to_bq(self, dataset, gcs_bucket, write_local_instead_of_bq=False, **at table_id = gcs_to_bq_util.make_bq_table_id(std_col.RACE_OR_HISPANIC_COL, STATE_LEVEL, CURRENT) gcs_to_bq_util.add_df_to_bq(df, dataset, table_id, column_types=col_types) - # WRITE ALLS TABLE FOR SEX/AGE (get just the All rows from the race table and add needed cols) - df = df.copy() - df = df[df[std_col.RACE_CATEGORY_ID_COL] == std_col.Race.ALL.value] - df.loc[:, std_col.SEX_COL] = std_col.ALL_VALUE - df.loc[:, std_col.AGE_COL] = std_col.ALL_VALUE - col_types = gcs_to_bq_util.get_bq_column_types(df, float_cols) - gcs_to_bq_util.add_df_to_bq(df, dataset, "alls_state_current", column_types=col_types) - def clean_row(df, column): """Removes non float KFF data from the raw data and replaces diff --git a/python/tests/datasources/test_kff_vaccination.py b/python/tests/datasources/test_kff_vaccination.py index c0b54f4307..f29742a5ce 100644 --- a/python/tests/datasources/test_kff_vaccination.py +++ b/python/tests/datasources/test_kff_vaccination.py @@ -88,10 +88,8 @@ def testWriteToBq( kffVaccination.write_to_bq("dataset", "gcs_bucket", **kwargs) assert mock_json.call_count == 3 assert mock_csv_web.call_count == 1 - assert mock_bq.call_count == 2 + assert mock_bq.call_count == 1 assert mock_bq.call_args_list[0].args[2] == "race_and_ethnicity_state_current" - assert mock_bq.call_args_list[1].args[2] == "alls_state_current" - df = mock_bq.call_args_list[0].args[0] expected_df = pd.read_csv(