Add samples for using GCP connections in an Airflow DAG (GoogleCloudPlatform#1590)

tswast · web-flow · commit b80895ed88ba · 2018-07-23T19:57:00.000-07:00
diff --git a/composer/workflows/bq_notify.py b/composer/workflows/bq_notify.py
@@ -80,7 +80,7 @@
 }
 
 with models.DAG(
-        'bq_notify',
+        'composer_sample_bq_notify',
         schedule_interval=datetime.timedelta(weeks=4),
         default_args=default_dag_args) as dag:
     # [END composer_notify_failure]
diff --git a/composer/workflows/connections.py b/composer/workflows/connections.py
@@ -0,0 +1,58 @@
+# Copyright 2018 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.
+
+"""Demonstrates how to use connections in an Airflow DAG."""
+
+import datetime
+
+from airflow import models
+from airflow.contrib.operators import bigquery_operator
+
+
+yesterday = datetime.datetime.combine(
+    datetime.datetime.today() - datetime.timedelta(1),
+    datetime.datetime.min.time())
+
+default_dag_args = {
+    # Setting start date as yesterday starts the DAG immediately when it is
+    # detected in the Cloud Storage bucket.
+    'start_date': yesterday,
+}
+
+# Define a DAG (directed acyclic graph) of tasks.
+# Any task you create within the context manager is automatically added to the
+# DAG object.
+with models.DAG(
+        'composer_sample_connections',
+        schedule_interval=datetime.timedelta(days=1),
+        default_args=default_dag_args) as dag:
+    # [START composer_connections_default]
+    task_default = bigquery_operator.BigQueryOperator(
+        task_id='task_default_connection',
+        bql='SELECT 1', use_legacy_sql=False)
+    # [END composer_connections_default]
+    # [START composer_connections_explicit]
+    task_explicit = bigquery_operator.BigQueryOperator(
+        task_id='task_explicit_connection',
+        bql='SELECT 1', use_legacy_sql=False,
+        # Composer creates a 'google_cloud_default' connection by default.
+        bigquery_conn_id='google_cloud_default')
+    # [END composer_connections_explicit]
+    # [START composer_connections_custom]
+    task_custom = bigquery_operator.BigQueryOperator(
+        task_id='task_custom_connection',
+        bql='SELECT 1', use_legacy_sql=False,
+        # Set a connection ID to use a connection that you have created.
+        bigquery_conn_id='my_gcp_connection')
+    # [END composer_connections_custom]
diff --git a/composer/workflows/connections_test.py b/composer/workflows/connections_test.py
@@ -0,0 +1,23 @@
+# Copyright 2018 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.
+
+
+def test_dag_import():
+    """Test that the DAG file can be successfully imported.
+
+    This tests that the DAG can be parsed, but does not run it in an Airflow
+    environment. This is a recommended sanity check by the official Airflow
+    docs: https://airflow.incubator.apache.org/tutorial.html#testing
+    """
+    from . import connections  # noqa
diff --git a/composer/workflows/kubernetes_pod_operator.py b/composer/workflows/kubernetes_pod_operator.py
@@ -60,7 +60,7 @@
 # will show the task as failed, as well as contain all of the task logs
 # required to debug.
 with models.DAG(
-        dag_id='kubernetes-pod-example',
+        dag_id='composer_sample_kubernetes_pod',
         schedule_interval=datetime.timedelta(days=1),
         start_date=YESTERDAY) as dag:
     # Only name, namespace, image, and task_id are required to create a
diff --git a/composer/workflows/quickstart.py b/composer/workflows/quickstart.py
@@ -64,7 +64,7 @@
 
 # [START composer_quickstart_schedule]
 with models.DAG(
-        'quickstart',
+        'composer_sample_quickstart',
         # Continue to run DAG once per day
         schedule_interval=datetime.timedelta(days=1),
         default_args=default_dag_args) as dag:
diff --git a/composer/workflows/simple.py b/composer/workflows/simple.py
@@ -41,7 +41,7 @@
 # Any task you create within the context manager is automatically added to the
 # DAG object.
 with models.DAG(
-        'simple_greeting',
+        'composer_sample_simple_greeting',
         schedule_interval=datetime.timedelta(days=1),
         default_args=default_dag_args) as dag:
     # [END composer_simple_define_dag]
diff --git a/composer/workflows/use_local_deps.py b/composer/workflows/use_local_deps.py
@@ -30,7 +30,9 @@
             datetime.datetime.min.time()),
 }
 
-with airflow.DAG('dependencies_dag', default_args=default_args) as dag:
+with airflow.DAG(
+        'composer_sample_dependencies_dag',
+        default_args=default_args) as dag:
     t1 = bash_operator.BashOperator(
         task_id='print_coin_result',
         bash_command='echo "{0}"'.format(coin_module.flip_coin()),

Original file line number	Diff line number	Diff line change
`@@ -80,7 +80,7 @@`
`80`	`80`	`}`
`81`	`81`
`82`	`82`	`with models.DAG(`
`83`		`- 'bq_notify',`
	`83`	`+ 'composer_sample_bq_notify',`
`84`	`84`	`schedule_interval=datetime.timedelta(weeks=4),`
`85`	`85`	`default_args=default_dag_args) as dag:`
`86`	`86`	`# [END composer_notify_failure]`