feat: add dataflow-composer tutorial code samples (GoogleCloudPlatform#3290)

* feat: add dataflow operator sample and test

* feat: add composer tutorial samples and tests

* fix: remove files

* fix: respond to comments

* fix: respond to comments

* fix: run lint

* fix import error

* feat: add dataflow template operator tutorial samples

* fix: commented out  variables

* fix errors

* fix: trailing whitespace

* fix: errant underscore

* fix: shorten UUID since it is too long

* fix: assertion

* fix: uppercase to lowercase f

* fix: create table format

* fix: imports out of order

* fix: import functions from helper file to test

* fix: run lint

* fix imports

* fix: failing test

* fix: tests to have specific uuid

* fix: run linting, clean up code

* fix: variables and headers

* fix: linting, comments, variable names

* fix: add command-line option for Python sample

* feat: add comments for reader

* Remove argparser

* correct errors in dag

* fix import order

* address lint

* fix: typo in parameters

* feat: fix typos in tutorial

* fix: change typle result

* fix: change variable name

* add dataset_name variable

* fix: change template to batch vs. streaming

* fix: linting

Co-authored-by: Leah E. Cole <[email protected]>
Co-authored-by: Leah Cole <[email protected]>
Co-authored-by: David Cavazos <[email protected]>

Author: 3 people

composer/workflows/dataflowtemplateoperator_create_dataset_and_table_helper.py

@@ -0,0 +1,79 @@
+# Copyright 2020 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.
+
+# This script is a helper function to the Dataflow Template Operator Tutorial.
+# It helps the user set up a BigQuery dataset and table that is needed
+# for the tutorial.
+
+# [START composer_dataflow_dataset_table_creation]
+
+# Make sure to follow the quickstart setup instructions beforehand.
+# See instructions here:
+# https://cloud.google.com/bigquery/docs/quickstarts/quickstart-client-libraries
+
+# Before running the sample, be sure to install the bigquery library
+# in your local environment by running pip install google.cloud.bigquery
+
+from google.cloud import bigquery
+
+# TODO(developer): Replace with your values
+project = 'your-project'  # Your GCP Project
+location = 'US'  # the location where you want your BigQuery data to reside. For more info on possible locations see https://cloud.google.com/bigquery/docs/locations
+dataset_name = 'average_weather'
+
+
+def create_dataset_and_table(project, location, dataset_name):
+    # Construct a BigQuery client object.
+    client = bigquery.Client(project)
+
+    dataset_id = f"{project}.{dataset_name}"
+
+    # Construct a full Dataset object to send to the API.
+    dataset = bigquery.Dataset(dataset_id)
+
+    # Set the location to your desired location for the dataset.
+    # For more information, see this link:
+    # https://cloud.google.com/bigquery/docs/locations
+    dataset.location = location
+
+    # Send the dataset to the API for creation.
+    # Raises google.api_core.exceptions.Conflict if the Dataset already
+    # exists within the project.
+    dataset = client.create_dataset(dataset)  # Make an API request.
+
+    print(f"Created dataset {client.project}.{dataset.dataset_id}")
+
+    # Create a table from this dataset.
+
+    table_id = f"{client.project}.{dataset_name}.average_weather"
+
+    schema = [
+        bigquery.SchemaField("location", "GEOGRAPHY", mode="REQUIRED"),
+        bigquery.SchemaField("average_temperature", "INTEGER", mode="REQUIRED"),
+        bigquery.SchemaField("month", "STRING", mode="REQUIRED"),
+        bigquery.SchemaField("inches_of_rain", "NUMERIC", mode="NULLABLE"),
+        bigquery.SchemaField("is_current", "BOOLEAN", mode="NULLABLE"),
+        bigquery.SchemaField("latest_measurement", "DATE", mode="NULLABLE"),
+    ]
+
+    table = bigquery.Table(table_id, schema=schema)
+    table = client.create_table(table)  # Make an API request.
+    print(f"Created table {table.project}.{table.dataset_id}.{table.table_id}")
+
+    # [END composer_dataflow_dataset_table_creation]
+    return dataset, table
+
+
+if __name__ == "__main__":
+    create_dataset_and_table(project, location, "average_weather")

composer/workflows/dataflowtemplateoperator_create_dataset_and_table_test.py

@@ -0,0 +1,53 @@
+# Copyright 2020 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.
+
+# This script is a helper function to the Dataflow Template Operator Tutorial.
+# It helps the user set up a BigQuery dataset and table that is needed
+# for the tutorial.
+
+import os
+import uuid
+
+from google.cloud import bigquery
+
+from . import dataflowtemplateoperator_create_dataset_and_table_helper as helper
+
+PROJECT_ID = os.environ["GCLOUD_PROJECT"]
+
+client = bigquery.Client()
+
+dataset_UUID = str(uuid.uuid4()).split("-")[0]
+
+expected_schema = [
+    bigquery.SchemaField("location", "GEOGRAPHY", mode="REQUIRED"),
+    bigquery.SchemaField("average_temperature", "INTEGER", mode="REQUIRED"),
+    bigquery.SchemaField("month", "STRING", mode="REQUIRED"),
+    bigquery.SchemaField("inches_of_rain", "NUMERIC", mode="NULLABLE"),
+    bigquery.SchemaField("is_current", "BOOLEAN", mode="NULLABLE"),
+    bigquery.SchemaField("latest_measurement", "DATE", mode="NULLABLE"),
+]
+
+
+def test_creation():
+    try:
+        dataset, table = helper.create_dataset_and_table(PROJECT_ID, "US", dataset_UUID)
+
+        assert table.table_id == "average_weather"
+        assert dataset.dataset_id == dataset_UUID
+        assert table.schema == expected_schema
+
+    finally:
+
+        client.delete_dataset(dataset, delete_contents=True, not_found_ok=True)
+        client.delete_table(table, not_found_ok=True)

composer/workflows/dataflowtemplateoperator_tutorial.py

@@ -0,0 +1,90 @@
+# Copyright 2020 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.
+
+# [START composer_dataflow_dag]
+
+
+"""Example Airflow DAG that creates a Cloud Dataproc cluster, runs the Hadoop
+wordcount example, and deletes the cluster.
+
+This DAG relies on three Airflow variables
+https://airflow.apache.org/concepts.html#variables
+* project_id - Google Cloud Project ID to use for the Cloud Dataflow cluster.
+* gce_zone - Google Compute Engine zone where Cloud Dataflow cluster should be
+  created.
+* gce_region - Google Compute Engine region where Cloud Dataflow cluster should be
+  created.
+Learn more about the difference between the two here:
+https://cloud.google.com/compute/docs/regions-zones
+* bucket_path - Google Cloud Storage bucket where you've stored the User Defined
+Function (.js), the input file (.txt), and the JSON schema (.json).
+"""
+
+import datetime
+
+from airflow import models
+from airflow.contrib.operators.dataflow_operator import DataflowTemplateOperator
+from airflow.utils.dates import days_ago
+
+bucket_path = models.Variable.get("bucket_path")
+project_id = models.Variable.get("project_id")
+gce_zone = models.Variable.get("gce_zone")
+gce_region = models.Variable.get("gce_region")
+
+
+default_args = {
+    # Tell airflow to start one day ago, so that it runs as soon as you upload it
+    "start_date": days_ago(1),
+    "dataflow_default_options": {
+        "project": project_id,
+        # Set to your region
+        "region": gce_region,
+        # Set to your zone
+        "zone": gce_zone,
+        # This is a subfolder for storing temporary files, like the staged pipeline job.
+        "temp_location": bucket_path + "/tmp/",
+    },
+}
+
+# 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(
+    # The id you will see in the DAG airflow page
+    "composer_dataflow_dag",
+    default_args=default_args,
+    # The interval with which to schedule the DAG
+    schedule_interval=datetime.timedelta(days=1),  # Override to match your needs
+) as dag:
+
+    start_template_job = DataflowTemplateOperator(
+        # The task id of your job
+        task_id="dataflow_operator_transform_csv_to_bq",
+        # The name of the template that you're using.
+        # Below is a list of all the templates you can use.
+        # For versions in non-production environments, use the subfolder 'latest'
+        # https://cloud.google.com/dataflow/docs/guides/templates/provided-batch#gcstexttobigquery
+        template="gs://dataflow-templates/latest/GCS_Text_to_BigQuery",
+        # Use the link above to specify the correct parameters for your template.
+        parameters={
+            "javascriptTextTransformFunctionName": "transformCSVtoJSON",
+            "JSONPath": bucket_path + "/jsonSchema.json",
+            "javascriptTextTransformGcsPath": bucket_path + "/transformCSVtoJSON.js",
+            "inputFilePattern": bucket_path + "/inputFile.txt",
+            "outputTable": project_id + ":average_weather.average_weather",
+            "bigQueryLoadingTemporaryDirectory": bucket_path + "/tmp/",
+        },
+    )
+
+# [END composer_dataflow_dag]

composer/workflows/dataflowtemplateoperator_tutorial_test.py

@@ -0,0 +1,33 @@
+# Copyright 2020 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.
+
+from airflow import models
+
+from . import unit_testing
+
+
+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
+    """
+    models.Variable.set("bucket_path", "gs://example_bucket")
+    models.Variable.set("project_id", "example-project")
+    models.Variable.set("gce_zone", "us-central1-f")
+    models.Variable.set("gce_region", "us-central1-f")
+    from . import dataflowtemplateoperator_tutorial as module
+
+    unit_testing.assert_has_valid_dag(module)