dataflow: update Dockerfile and modularize tests

dataflow/conftest.py

@@ -0,0 +1,378 @@
+# 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
+#
+#     http://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.
+
+from dataclasses import dataclass
+import itertools
+import json
+import multiprocessing as mp
+import os
+import subprocess
+import sys
+import time
+from typing import Any, Callable, Dict, Iterable, Optional
+import uuid
+
+import pytest
+
+# Default options.
+UUID = uuid.uuid4().hex[0:6]
+PROJECT = os.environ["GOOGLE_CLOUD_PROJECT"]
+REGION = "us-west1"
+ZONE = "us-west1-b"
+
+RETRY_MAX_TIME = 5 * 60  # 5 minutes in seconds
+
+
+@dataclass
+class Utils:
+    uuid: str = UUID
+    project: str = PROJECT
+    region: str = REGION
+    zone: str = ZONE
+
+    @staticmethod
+    def storage_bucket(bucket_name: str) -> str:
+        from google.cloud import storage
+
+        storage_client = storage.Client()
+        bucket_unique_name = f"{bucket_name}-{UUID}"
+        bucket = storage_client.create_bucket(bucket_unique_name)
+
+        print(f"storage_bucket: {bucket_unique_name}")
+        yield bucket_unique_name
+
+        bucket.delete(force=True)
+
+    @staticmethod
+    def bigquery_dataset(dataset_name: str, project: str = PROJECT) -> str:
+        from google.cloud import bigquery
+
+        bigquery_client = bigquery.Client()
+        dataset = bigquery_client.create_dataset(
+            bigquery.Dataset(f"{project}.{dataset_name.replace('-', '_')}_{UUID}")
+        )
+
+        print(f"bigquery_dataset: {dataset.full_dataset_id}")
+        yield dataset.full_dataset_id
+
+        bigquery_client.delete_dataset(
+            dataset.full_dataset_id.replace(":", "."), delete_contents=True
+        )
+
+    @staticmethod
+    def bigquery_query(query: str) -> Iterable[Dict[str, Any]]:
+        from google.cloud import bigquery
+
+        bigquery_client = bigquery.Client()
+        for row in bigquery_client.query(query):
+            yield dict(row)
+
+    @staticmethod
+    def pubsub_topic(topic_name: str, project: str = PROJECT) -> str:
+        from google.cloud import pubsub
+
+        publisher_client = pubsub.PublisherClient()
+        topic_path = publisher_client.topic_path(project, f"{topic_name}-{UUID}")
+        topic = publisher_client.create_topic(topic_path)
+
+        print(f"pubsub_topic: {topic.name}")
+        yield topic.name
+
+        # Due to the pinned library dependencies in apache-beam, client
+        # library throws an error upon deletion.
+        # We use gcloud for a workaround. See also:
+        # https://github.com/GoogleCloudPlatform/python-docs-samples/issues/4492
+        cmd = ["gcloud", "pubsub", "--project", project, "topics", "delete", topic.name]
+        print(cmd)
+        subprocess.run(cmd, check=True)
+
+    @staticmethod
+    def pubsub_subscription(
+        topic_path: str,
+        subscription_name: str,
+        project: str = PROJECT,
+    ) -> str:
+        from google.cloud import pubsub
+
+        subscriber = pubsub.SubscriberClient()
+        subscription_path = subscriber.subscription_path(
+            project, f"{subscription_name}-{UUID}"
+        )
+        subscription = subscriber.create_subscription(subscription_path, topic_path)
+
+        print(f"pubsub_subscription: {subscription.name}")
+        yield subscription.name
+
+        # Due to the pinned library dependencies in apache-beam, client
+        # library throws an error upon deletion.
+        # We use gcloud for a workaround. See also:
+        # https://github.com/GoogleCloudPlatform/python-docs-samples/issues/4492
+        cmd = [
+            "gcloud",
+            "pubsub",
+            "--project",
+            project,
+            "subscriptions",
+            "delete",
+            subscription.name,
+        ]
+        print(cmd)
+        subprocess.run(cmd, check=True)
+
+    @staticmethod
+    def pubsub_publisher(
+        topic_path: str,
+        new_msg: Callable[[int], str] = lambda i: json.dumps(
+            {"id": i, "content": f"message {i}"}
+        ),
+        sleep_sec: int = 1,
+    ) -> bool:
+        from google.cloud import pubsub
+
+        def _infinite_publish_job() -> None:
+            publisher_client = pubsub.PublisherClient()
+            for i in itertools.count():
+                msg = new_msg(i)
+                publisher_client.publish(topic_path, msg.encode("utf-8")).result()
+                time.sleep(sleep_sec)
+
+        # Start a subprocess in the background to do the publishing.
+        print(f"Starting publisher on {topic_path}")
+        p = mp.Process(target=_infinite_publish_job)
+        p.start()
+
+        yield p.is_alive()
+
+        # For cleanup, terminate the background process.
+        print("Stopping publisher")
+        p.join(timeout=0)
+        p.terminate()
+
+    @staticmethod
+    def container_image(
+        image_path: str,
+        project: str = PROJECT,
+        tag: str = "latest",
+    ) -> str:
+        image_name = f"gcr.io/{project}/{image_path}-{UUID}:{tag}"
+        cmd = ["gcloud", "auth", "configure-docker"]
+        print(cmd)
+        subprocess.run(cmd, check=True)
+        cmd = [
+            "gcloud",
+            "builds",
+            "submit",
+            f"--project={project}",
+            f"--tag={image_name}",
+            ".",
+        ]
+        print(cmd)
+        subprocess.run(cmd, check=True)
+
+        print(f"container_image: {image_name}")
+        yield image_name
+
+        cmd = [
+            "gcloud",
+            "container",
+            "images",
+            "delete",
+            image_name,
+            f"--project={project}",
+            "--quiet",
+        ]
+        print(cmd)
+        subprocess.run(cmd, check=True)
+
+    @staticmethod
+    def dataflow_job_id_from_job_name(
+        job_name: str,
+        project: str = PROJECT,
+    ) -> Optional[str]:
+        from googleapiclient.discovery import build
+
+        dataflow = build("dataflow", "v1b3")
+
+        # Only return the 50 most recent results - our job is likely to be in here.
+        # If the job is not found, first try increasing this number.[]''job_id
+        # For more info see:
+        #   https://cloud.google.com/dataflow/docs/reference/rest/v1b3/projects.jobs/list
+        jobs_request = (
+            dataflow.projects()
+            .jobs()
+            .list(
+                projectId=project,
+                filter="ACTIVE",
+                pageSize=50,
+            )
+        )
+        response = jobs_request.execute()
+
+        # Search for the job in the list that has our name (names are unique)
+        for job in response["jobs"]:
+            if job["name"] == job_name:
+                return job["id"]
+        return None
+
+    @staticmethod
+    def dataflow_jobs_wait(
+        job_id: str,
+        project: str = PROJECT,
+        status: str = "JOB_STATE_RUNNING",
+    ) -> bool:
+        from googleapiclient.discovery import build
+
+        dataflow = build("dataflow", "v1b3")
+
+        sleep_time_seconds = 30
+        max_sleep_time = 10 * 60
+
+        print(f"Waiting for Dataflow job ID: {job_id} (until status {status})")
+        for _ in range(0, max_sleep_time, sleep_time_seconds):
+            try:
+                # For more info see:
+                #   https://cloud.google.com/dataflow/docs/reference/rest/v1b3/projects.jobs/get
+                jobs_request = (
+                    dataflow.projects()
+                    .jobs()
+                    .get(
+                        projectId=project,
+                        jobId=job_id,
+                        view="JOB_VIEW_SUMMARY",
+                    )
+                )
+                response = jobs_request.execute()
+                print(response)
+                if response["currentState"] == status:
+                    return True
+            except:
+                pass
+            time.sleep(sleep_time_seconds)
+        return False
+
+    @staticmethod
+    def dataflow_jobs_cancel_by_job_id(
+        job_id: str, project: str = PROJECT, region: str = REGION
+    ) -> None:
+        print(f"Canceling Dataflow job ID: {job_id}")
+        # We get an error using the googleapiclient.discovery APIs, probably
+        # due to incompatible dependencies with apache-beam.
+        # We use gcloud instead to cancel the job.
+        #   https://cloud.google.com/sdk/gcloud/reference/dataflow/jobs/cancel
+        cmd = [
+            "gcloud",
+            f"--project={project}",
+            "dataflow",
+            "jobs",
+            "cancel",
+            job_id,
+            f"--region={region}",
+        ]
+        subprocess.run(cmd, check=True)
+
+    @staticmethod
+    def dataflow_jobs_cancel_by_job_name(
+        job_name: str, project: str = PROJECT, region: str = REGION
+    ) -> None:
+        # To cancel a dataflow job, we need its ID, not its name.
+        # If it doesn't, job_id will be equal to None.
+        job_id = Utils.dataflow_job_id_from_job_name(project, job_name)
+        if job_id is not None:
+            Utils.dataflow_jobs_cancel_by_job_id(job_id, project, region)
+
+    @staticmethod
+    def dataflow_flex_template_build(
+        bucket_name: str,
+        template_image: str,
+        metadata_file: str,
+        project: str = PROJECT,
+        template_file: str = "template.json",
+    ) -> str:
+        # https://cloud.google.com/sdk/gcloud/reference/dataflow/flex-template/build
+        template_gcs_path = f"gs://{bucket_name}/{template_file}"
+        cmd = [
+            "gcloud",
+            "dataflow",
+            "flex-template",
+            "build",
+            template_gcs_path,
+            f"--project={project}",
+            f"--image={template_image}",
+            "--sdk-language=PYTHON",
+            f"--metadata-file={metadata_file}",
+        ]
+        print(cmd)
+        subprocess.run(cmd, check=True)
+
+        print(f"dataflow_flex_template_build: {template_gcs_path}")
+        yield template_gcs_path
+        # The template file gets deleted when we delete the bucket.
+
+    @staticmethod
+    def dataflow_flex_template_run(
+        job_name: str,
+        template_path: str,
+        bucket_name: str,
+        parameters: Dict[str, str] = {},
+        project: str = PROJECT,
+        region: str = REGION,
+    ) -> str:
+        import yaml
+
+        # https://cloud.google.com/sdk/gcloud/reference/dataflow/flex-template/run
+        unique_job_name = f"{job_name}-{UUID}"
+        print(f"dataflow_job_name: {unique_job_name}")
+        cmd = [
+            "gcloud",
+            "dataflow",
+            "flex-template",
+            "run",
+            unique_job_name,
+            f"--template-file-gcs-location={template_path}",
+            f"--project={project}",
+            f"--region={region}",
+        ] + [
+            f"--parameters={name}={value}"
+            for name, value in {
+                **parameters,
+                "temp_location": f"gs://{bucket_name}/temp",
+            }.items()
+        ]
+        print(cmd)
+        try:
+            # The `capture_output` option was added in Python 3.7, so we must
+            # pass the `stdout` and `stderr` options explicitly to support 3.6.
+            # https://docs.python.org/3/library/subprocess.html#subprocess.run
+            p = subprocess.run(
+                cmd, check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE
+            )
+            stdout = p.stdout.decode("utf-8")
+            stderr = p.stderr.decode("utf-8")
+            print(f"Launched Dataflow Flex Template job: {unique_job_name}")
+        except subprocess.CalledProcessError as e:
+            print(e, file=sys.stderr)
+            stdout = stdout.decode("utf-8")
+            stderr = stderr.decode("utf-8")
+        finally:
+            print("--- stderr ---")
+            print(stderr)
+            print("--- stdout ---")
+            print(stdout)
+            print("--- end ---")
+        return yaml.safe_load(stdout)["job"]["id"]
+
+
+@pytest.fixture(scope="session")
+def utils() -> Utils:
+    print(f"Test unique identifier: {UUID}")
+    subprocess.run(["gcloud", "version"])
+    return Utils()

dataflow/flex-templates/__init__.py

@@ -0,0 +1 @@
+

dataflow/flex-templates/streaming_beam/.dockerignore

@@ -0,0 +1,4 @@
+# Ignore everything except for Python files and the requirements file.
+*
+!*.py
+!requirements.txt

dataflow/flex-templates/streaming_beam/.gcloudignore

@@ -0,0 +1,5 @@
+# Ignore everything except for Python files and the requirements file.
+*
+!Dockerfile
+!*.py
+!requirements.txt