cleanup: Upgrade Composer samples to Airflow 2.0

composer/blog/gcp-tech-blog/unit-test-dags-cloud-build/conftest.py

@@ -0,0 +1,39 @@
+# Copyright 2021 Google LLC
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     https://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import os
+
+import pytest
+
+
+# this fixture initializes the Airflow DB once per session
+# it is used by DAGs in this blog post code only
+@pytest.fixture(scope="session")
+def airflow_database():
+    import airflow.utils.db
+
+    # We use separate directory for local db path per session
+    # by setting AIRFLOW_HOME env var, which is done in noxfile_config.py.
+
+    assert('AIRFLOW_HOME' in os.environ)
+
+    airflow_home = os.environ["AIRFLOW_HOME"]
+    airflow_db = f"{airflow_home}/airflow.db"
+
+    # reset both resets and initializes a new database
+    airflow.utils.db.resetdb(rbac=None)
+
+    # Making sure we are using a data file there.
+    assert(os.path.isfile(airflow_db))

composer/conftest.py

@@ -1,4 +1,4 @@
-# Copyright 2019 Google Inc. All Rights Reserved.
+# Copyright 2019 Google LLC All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -19,6 +19,27 @@
 import pytest
 
 
+# this fixture initializes the Airflow DB once per session
+# it is used by DAGs in both the blogs and workflows directories, 
+# unless there exists a conftest at a lower level
+@pytest.fixture(scope="session")
+def airflow_database():
+    import airflow.utils.db
+
+    # We use separate directory for local db path per session
+    # by setting AIRFLOW_HOME env var, which is done in noxfile_config.py.
+
+    assert('AIRFLOW_HOME' in os.environ)
+
+    airflow_home = os.environ["AIRFLOW_HOME"]
+    airflow_db = f"{airflow_home}/airflow.db"
+
+    # reset both resets and initializes a new database
+    airflow.utils.db.resetdb()
+
+    # Making sure we are using a data file there.
+    assert(os.path.isfile(airflow_db))
+
 # this fixture initializes the Airflow DB once per session
 # it is used by DAGs in both the blogs and workflows directories
 @pytest.fixture(scope="session")
@@ -34,7 +55,7 @@ def airflow_database():
     airflow_db = f"{airflow_home}/airflow.db"
 
     # reset both resets and initializes a new database
-    airflow.utils.db.resetdb(rbac=None)  # this command will change in Airflow 2.0
+    airflow.utils.db.resetdb()
 
     # Making sure we are using a data file there.
     assert(os.path.isfile(airflow_db))

composer/workflows/airflow_db_cleanup.py

@@ -54,10 +54,10 @@
 import airflow
 from airflow import settings
 from airflow.configuration import conf
-from airflow.jobs import BaseJob
+from airflow.jobs.base_job import BaseJob
 from airflow.models import DAG, DagModel, DagRun, Log, SlaMiss, \
     TaskInstance, Variable, XCom
-from airflow.operators.python_operator import PythonOperator
+from airflow.operators.python import PythonOperator
 import dateutil.parser
 from sqlalchemy import and_, func
 from sqlalchemy.exc import ProgrammingError

composer/workflows/bq_copy_across_locations.py

@@ -40,10 +40,10 @@
 import logging
 
 from airflow import models
-from airflow.contrib.operators import bigquery_to_gcs
-from airflow.contrib.operators import gcs_to_bq
-from airflow.contrib.operators import gcs_to_gcs
-from airflow.operators import dummy_operator
+from airflow.operators import dummy
+from airflow.providers.google.cloud.transfers import bigquery_to_gcs
+from airflow.providers.google.cloud.transfers import gcs_to_bigquery
+from airflow.providers.google.cloud.transfers import gcs_to_gcs
 
 
 # --------------------------------------------------------------------------------
@@ -127,12 +127,12 @@ def read_table_list(table_list_file):
         'composer_sample_bq_copy_across_locations',
         default_args=default_args,
         schedule_interval=None) as dag:
-    start = dummy_operator.DummyOperator(
+    start = dummy.DummyOperator(
         task_id='start',
         trigger_rule='all_success'
     )
 
-    end = dummy_operator.DummyOperator(
+    end = dummy.DummyOperator(
         task_id='end',
         trigger_rule='all_success'
     )
@@ -148,7 +148,7 @@ def read_table_list(table_list_file):
         table_source = record['table_source']
         table_dest = record['table_dest']
 
-        BQ_to_GCS = bigquery_to_gcs.BigQueryToCloudStorageOperator(
+        BQ_to_GCS = bigquery_to_gcs.BigQueryToGCSOperator(
             # Replace ":" with valid character for Airflow task
             task_id='{}_BQ_to_GCS'.format(table_source.replace(":", "_")),
             source_project_dataset_table=table_source,
@@ -157,7 +157,7 @@ def read_table_list(table_list_file):
             export_format='AVRO'
         )
 
-        GCS_to_GCS = gcs_to_gcs.GoogleCloudStorageToGoogleCloudStorageOperator(
+        GCS_to_GCS = gcs_to_gcs.GCSToGCSOperator(
             # Replace ":" with valid character for Airflow task
             task_id='{}_GCS_to_GCS'.format(table_source.replace(":", "_")),
             source_bucket=source_bucket,
@@ -166,7 +166,7 @@ def read_table_list(table_list_file):
             # destination_object='{}-*.avro'.format(table_dest)
         )
 
-        GCS_to_BQ = gcs_to_bq.GoogleCloudStorageToBigQueryOperator(
+        GCS_to_BQ = gcs_to_bigquery.GCSToBigQueryOperator(
             # Replace ":" with valid character for Airflow task
             task_id='{}_GCS_to_BQ'.format(table_dest.replace(":", "_")),
             bucket=dest_bucket,

composer/workflows/bq_copy_eu_to_us_sample.csv

@@ -1,3 +1,3 @@
 Source, Target
-nyc-tlc:green.trips_2014,nyc_tlc_EU.trips_2014
-nyc-tlc:green.trips_2015,nyc_tlc_EU.trips_2015
+nyc-tlc.green.trips_2014,nyc_tlc_EU.trips_2014
+nyc-tlc.green.trips_2015,nyc_tlc_EU.trips_2015