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