feat!: Adds support for Airflow 2 Cloud Composer environment and oper…

…ators (#134) * new DAG operator imports for Airflow 2.0+ * support Airflow 2 operators when generating DAGs * version compatibility checks for DAGs and Airflow environments * tests to copy pipeline.yaml to dot folder (for airflow version spec) * tests for version compatibility checks * revised pipeline YAML template for Airflow 1 compat * feat!: Upgrade dependencies to Airflow 2.1.1 (#135) * revised dependendices for Airflow 2.1.1 * use proper variables file as required by Airflow 2 * changed README to use Airflow 2 requirement * uncommented deploy_dag_versioning check * feat: Upgrade `usa_names` pipeline to usse Airflow 2 operators and environment (#136) * default to Airflow 2
GoogleCloudPlatform · Aug 10, 2021 · b2749c6 · b2749c6
1 parent d5ef401
commit b2749c6
Show file tree

Hide file tree

Showing 30 changed files with 1,287 additions and 675 deletions.
diff --git a/Pipfile b/Pipfile
@@ -7,7 +7,9 @@ name = "pypi"
 beautifulsoup4 = "==4.9.3"
 
 [dev-packages]
-apache-airflow = {version = "==1.10.15", extras = ["google"]}
+apache-airflow = "==2.1.1"
+apache-airflow-providers-google = "*"
+apache-airflow-providers-cncf-kubernetes = "*"
 black = "==20.8b1"
 flake8 = "==3.9.2"
 isort = "*"
@@ -16,8 +18,8 @@ pandas-gbq = "==0.14.1"
 pytest-mock = "*"
 pytest = "*"
 "ruamel.yaml" = "==0.17.10"
-Jinja2 = "*"
-SQLAlchemy = "==1.3.15"
+Jinja2 = "==2.11.3"
+SQLAlchemy = "==1.3.18"
 
 [requires]
 python_version = "3.8"
diff --git a/Pipfile.lock b/Pipfile.lock
diff --git a/README.md b/README.md
@@ -8,11 +8,11 @@ Cloud-native, data pipeline architecture for onboarding public datasets to [Data
 
 # Requirements
 - Python `>=3.6.10,<3.9`. We currently use `3.8`. For more info, see the [Cloud Composer version list](https://cloud.google.com/composer/docs/concepts/versioning/composer-versions).
-- Familiarity with [Apache Airflow](https://airflow.apache.org/docs/apache-airflow/1.10.15/concepts.html) (>=v1.10.15)
+- Familiarity with [Apache Airflow](https://airflow.apache.org/docs/apache-airflow/stable/concepts/index.html) (>=v2.1)
 - [pipenv](https://pipenv-fork.readthedocs.io/en/latest/install.html#installing-pipenv) for creating similar Python environments via `Pipfile.lock`
 - [gcloud](https://cloud.google.com/sdk/gcloud) command-line tool with Google Cloud Platform credentials configured. Instructions can be found [here](https://cloud.google.com/sdk/docs/initializing).
 - [Terraform](https://learn.hashicorp.com/tutorials/terraform/install-cli) `>=v0.15.1`
-- [Google Cloud Composer](https://cloud.google.com/composer/docs/concepts/overview) environment running [Apache Airflow](https://airflow.apache.org/docs/apache-airflow/1.10.15/concepts.html) `>=v1.10.15,<2.0`. To create a new Cloud Composer environment, see [this guide](https://cloud.google.com/composer/docs/how-to/managing/creating).
+- [Google Cloud Composer](https://cloud.google.com/composer/docs/concepts/overview) environment running [Apache Airflow](https://airflow.apache.org/docs/apache-airflow/stable/concepts.html) `>=2.0`. To create a new Cloud Composer environment, see [this guide](https://cloud.google.com/composer/docs/how-to/managing/creating).
 
 # Environment Setup
 

diff --git a/datasets/bls/cpsaat18/pipeline.yaml b/datasets/bls/cpsaat18/pipeline.yaml
@@ -19,6 +19,7 @@ resources:
     description: "Current population survey 18: Employed persons by detailed industry, sex, race, and Hispanic or Latino ethnicity"
 
 dag:
+  airflow_version: 1
   initialize:
     dag_id: cpsaat18
     default_args:

diff --git a/datasets/covid19_tracking/city_level_cases_and_deaths/pipeline.yaml b/datasets/covid19_tracking/city_level_cases_and_deaths/pipeline.yaml
@@ -18,6 +18,7 @@ resources:
     table_id: city_level_cases_and_deaths
 
 dag:
+  airflow_version: 1
   initialize:
     dag_id: "city_level_cases_and_deaths"
     default_args:

diff --git a/datasets/covid19_tracking/covid_racial_data_tracker/pipeline.yaml b/datasets/covid19_tracking/covid_racial_data_tracker/pipeline.yaml
@@ -19,6 +19,7 @@ resources:
 
 dag:
   initialize:
+    airflow_version: 1
     dag_id: "covid_racial_data_tracker"
     default_args:
       owner: "Google"

diff --git a/datasets/covid19_tracking/national_testing_and_outcomes/pipeline.yaml b/datasets/covid19_tracking/national_testing_and_outcomes/pipeline.yaml
@@ -18,6 +18,7 @@ resources:
     table_id: national_testing_and_outcomes
 
 dag:
+  airflow_version: 1
   initialize:
     dag_id: "national_testing_and_outcomes"
     default_args:

diff --git a/datasets/covid19_tracking/state_facility_level_long_term_care/pipeline.yaml b/datasets/covid19_tracking/state_facility_level_long_term_care/pipeline.yaml
@@ -18,6 +18,7 @@ resources:
     table_id: state_facility_level_long_term_care
 
 dag:
+  airflow_version: 1
   initialize:
     dag_id: "state_facility_level_long_term_care"
     default_args:

diff --git a/datasets/covid19_tracking/state_level_aggregate_long_term_care/pipeline.yaml b/datasets/covid19_tracking/state_level_aggregate_long_term_care/pipeline.yaml
@@ -18,6 +18,7 @@ resources:
     table_id: state_level_aggregate_long_term_care
 
 dag:
+  airflow_version: 1
   initialize:
     dag_id: "state_level_aggregate_long_term_care"
     default_args:

diff --git a/datasets/covid19_tracking/state_level_cumulative_long_term_care/pipeline.yaml b/datasets/covid19_tracking/state_level_cumulative_long_term_care/pipeline.yaml
@@ -18,6 +18,7 @@ resources:
     table_id: state_level_cumulative_long_term_care
 
 dag:
+  airflow_version: 1
   initialize:
     dag_id: "state_level_cumulative_long_term_care"
     default_args:

diff --git a/datasets/covid19_tracking/state_level_current_outbreak_long_term_care/pipeline.yaml b/datasets/covid19_tracking/state_level_current_outbreak_long_term_care/pipeline.yaml
@@ -18,6 +18,7 @@ resources:
     table_id: state_level_current_outbreak_long_term_care
 
 dag:
+  airflow_version: 1
   initialize:
     dag_id: "state_level_current_outbreak_long_term_care"
     default_args:

diff --git a/datasets/covid19_tracking/state_screenshots/pipeline.yaml b/datasets/covid19_tracking/state_screenshots/pipeline.yaml
@@ -18,6 +18,7 @@ resources:
     table_id: state_screenshots
 
 dag:
+  airflow_version: 1
   initialize:
     dag_id: "state_screenshots"
     default_args:

diff --git a/datasets/covid19_tracking/state_testing_and_outcomes/pipeline.yaml b/datasets/covid19_tracking/state_testing_and_outcomes/pipeline.yaml
@@ -18,6 +18,7 @@ resources:
     table_id: state_testing_and_outcomes
 
 dag:
+  airflow_version: 1
   initialize:
     dag_id: "state_testing_and_outcomes"
     default_args:

diff --git a/datasets/covid19_vaccination_access/vaccination_access_to_bq/pipeline.yaml b/datasets/covid19_vaccination_access/vaccination_access_to_bq/pipeline.yaml
@@ -28,6 +28,7 @@ resources:
     description: "This table represents the boundaries of areas surrounding vaccination facilities from which people can reach the facility by walking within predetermined time periods."
 
 dag:
+  airflow_version: 1
   initialize:
     dag_id: vaccination_access_to_bq
     default_args:

diff --git a/...ets/covid19_vaccination_search_insights/covid19_vaccination_search_insights/pipeline.yaml b/...ets/covid19_vaccination_search_insights/covid19_vaccination_search_insights/pipeline.yaml
@@ -122,6 +122,7 @@ resources:
       ]
 
 dag:
+  airflow_version: 1
   initialize:
     dag_id: covid19_vaccination_search_insights
     default_args:

diff --git a/datasets/google_dei/diversity_annual_report/pipeline.yaml b/datasets/google_dei/diversity_annual_report/pipeline.yaml
@@ -65,6 +65,7 @@ resources:
     deletion_protection: False
 
 dag:
+  airflow_version: 1
   initialize:
     dag_id: diversity_annual_report
     default_args:

diff --git a/datasets/google_trends/top_terms/pipeline.yaml b/datasets/google_trends/top_terms/pipeline.yaml
@@ -93,6 +93,7 @@ resources:
       ]
 
 dag:
+  airflow_version: 1
   initialize:
     dag_id: top_terms
     default_args:

diff --git a/datasets/iowa_liquor_sales_forecasting/2020_sales_train/pipeline.yaml b/datasets/iowa_liquor_sales_forecasting/2020_sales_train/pipeline.yaml
@@ -36,6 +36,7 @@ dag:
   # The DAG acronym stands for directed acyclic graph. This block represents
   # your data pipeline along with every property and configuration it needs to
   # onboard your data.
+  airflow_version: 1
   initialize:
     dag_id: 2020_sales_train
     default_args:

diff --git a/datasets/iowa_liquor_sales_forecasting/2021_sales_predict/pipeline.yaml b/datasets/iowa_liquor_sales_forecasting/2021_sales_predict/pipeline.yaml
@@ -36,6 +36,7 @@ dag:
   # The DAG acronym stands for directed acyclic graph. This block represents
   # your data pipeline along with every property and configuration it needs to
   # onboard your data.
+  airflow_version: 1
   initialize:
     dag_id: 2021_sales_predict
     default_args:

diff --git a/datasets/ml_datasets/penguins/pipeline.yaml b/datasets/ml_datasets/penguins/pipeline.yaml
@@ -36,6 +36,7 @@ dag:
   # The DAG acronym stands for directed acyclic graph. This block represents
   # your data pipeline along with every property and configuration it needs to
   # onboard your data.
+  airflow_version: 1
   initialize:
     dag_id: penguins
     default_args:

diff --git a/datasets/usa_names/usa_1910_current/pipeline.yaml b/datasets/usa_names/usa_1910_current/pipeline.yaml
@@ -22,6 +22,7 @@ resources:
       source: http://www.ssa.gov/OACT/babynames/limits.html
 
 dag:
+  airflow_version: 2
   initialize:
     dag_id: "usa_1910_current"
     default_args:

diff --git a/datasets/usa_names/usa_1910_current/usa_1910_current_dag.py b/datasets/usa_names/usa_1910_current/usa_1910_current_dag.py
@@ -14,8 +14,8 @@
 
 
 from airflow import DAG
-from airflow.contrib.operators import gcs_to_bq
-from airflow.operators import bash_operator
+from airflow.operators import bash
+from airflow.providers.google.cloud.transfers import gcs_to_bigquery
 
 default_args = {
     "owner": "Google",
@@ -34,7 +34,7 @@
 ) as dag:
 
     # Task to copy `namesbystate.zip` from Social Security Administration to GCS
-    download_and_process_source_zip_file = bash_operator.BashOperator(
+    download_and_process_source_zip_file = bash.BashOperator(
         task_id="download_and_process_source_zip_file",
         bash_command="mkdir -p $data_dir/{{ ds }}\ncurl -o $data_dir/{{ ds }}/namesbystate.zip -L $zip_source_url\nunzip $data_dir/{{ ds }}/namesbystate.zip -d $data_dir/{{ ds }}\ncat $data_dir/{{ ds }}/*.TXT \u003e\u003e $data_dir/{{ ds }}/data.csv\n",
         env={
@@ -44,7 +44,7 @@
     )
 
     # Task to load the data from Airflow data folder to BigQuery
-    load_csv_file_to_bq_table = gcs_to_bq.GoogleCloudStorageToBigQueryOperator(
+    load_csv_file_to_bq_table = gcs_to_bigquery.GCSToBigQueryOperator(
         task_id="load_csv_file_to_bq_table",
         bucket="{{ var.json.shared.composer_bucket }}",
         source_objects=["data/usa_names/usa_1910_current/{{ ds }}/data.csv"],

diff --git a/datasets/vizgen_merfish/mouse_brain_map/pipeline.yaml b/datasets/vizgen_merfish/mouse_brain_map/pipeline.yaml
@@ -19,6 +19,7 @@ dag:
   # The DAG acronym stands for directed acyclic graph. This block represents
   # your data pipeline along with every property and configuration it needs to
   # onboard your data.
+  airflow_version: 1
   initialize:
     dag_id: mouse_brain_map
     default_args: