Skip to content

Commit

Permalink
Kff dag (#3915)
Browse files Browse the repository at this point in the history
# 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)

<img width="1606" alt="Screenshot 2025-01-17 at 11 24 49 AM"
src="https://github.com/user-attachments/assets/5e9fb2f2-0941-419e-8690-b02eece36c4e"
/>
<img width="1606" alt="Screenshot 2025-01-17 at 11 26 00 AM"
src="https://github.com/user-attachments/assets/6cebbb1f-2ea4-49a6-a86d-b0f0c9243c49"
/>



## 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 <[email protected]>
  • Loading branch information
JDemlow and benhammondmusic authored Jan 17, 2025
1 parent 292dfdf commit 7638bf3
Show file tree
Hide file tree
Showing 3 changed files with 7 additions and 17 deletions.
12 changes: 6 additions & 6 deletions airflow/dags/kff_vaccination.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
8 changes: 0 additions & 8 deletions python/datasources/kff_vaccination.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
4 changes: 1 addition & 3 deletions python/tests/datasources/test_kff_vaccination.py
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down

0 comments on commit 7638bf3

Please sign in to comment.