Merge pull request GoogleCloudPlatform#1 from akvelon/dev-t1

add python code

Author: anton-chislov-akvelon

composer/functions/pubsub_function.zip

@@ -0,0 +1,57 @@
+# Copyright 2023 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.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+""" This Cloud Function example creates Pub/Sub messages.
+
+    Usage: Replace <PROJECT_ID> with the project ID of your project.
+"""
+
+from google.cloud import pubsub_v1
+
+project = "<PROJECT_ID>"
+topic = "dag-topic-trigger"
+
+
+def pubsub_publisher(request):
+    """Publish message from HTTP request to Pub/Sub topic.
+    Args:
+        request (flask.Request): HTTP request object.
+    Returns:
+        The response text with message published into Pub/Sub topic
+        Response object using
+        `make_response <http://flask.pocoo.org/docs/1.0/api/#flask.Flask.make_response>`.
+    """
+    request_json = request.get_json()
+    print(request_json)
+    if request.args and 'message' in request.args:
+        data_str = request.args.get('message')
+    elif request_json and 'message' in request_json:
+        data_str = request_json['message']
+    else:
+        return "Message content not found! Use 'message' key to specify"
+
+    publisher = pubsub_v1.PublisherClient()
+    # The `topic_path` method creates a fully qualified identifier
+    # in the form `projects/{project_id}/topics/{topic_id}`
+    topic_path = publisher.topic_path(project, topic)
+
+    # The required data format is a bytestring
+    data = data_str.encode("utf-8")
+    # When you publish a message, the client returns a future.
+    message_length = len(data_str)
+    future = publisher.publish(topic_path,
+                               data,
+                               message_length=str(message_length))
+    print(future.result())
+
+    return f"Message {data} with message_length {message_length} published to {topic_path}."

composer/functions/pubsub_publisher.py

@@ -0,0 +1,107 @@
+# Copyright 2023 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.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import mock
+import pytest
+
+import pubsub_publisher
+
+
+@pytest.fixture()
+def dump_request_args():
+    class Request:
+        args = {"message": "test with args"}
+
+        def get_json(self):
+            return self.args
+
+    return Request()
+
+
+@pytest.fixture()
+def dump_request():
+    class Request:
+        args = None
+
+        def get_json(self):
+            return {"message": "test with no args"}
+
+    return Request()
+
+
+@pytest.fixture()
+def dump_request_no_message():
+    class Request:
+        args = None
+
+        def get_json(self):
+            return {"no_message": "test with no message key"}
+
+    return Request()
+
+
+# Pass None, an input that is not valid request
+def test_request_with_none():
+    request = None
+    with pytest.raises(Exception):
+        pubsub_publisher.pubsub_publisher(request)
+
+
+def test_content_not_found(dump_request_no_message):
+    output = "Message content not found! Use 'message' key to specify"
+    assert pubsub_publisher.pubsub_publisher(dump_request_no_message) == output, f"The function didn't return '{output}'"
+
+
+@mock.patch("pubsub_publisher.pubsub_v1.PublisherClient.publish")
+@mock.patch("pubsub_publisher.pubsub_v1.PublisherClient.topic_path")
+def test_topic_path_args(topic_path, _, dump_request_args):
+    pubsub_publisher.pubsub_publisher(dump_request_args)
+
+    topic_path.assert_called_once_with(
+        "<PROJECT_ID>",
+        "dag-topic-trigger",
+    )
+
+
+@mock.patch("pubsub_publisher.pubsub_v1.PublisherClient.publish")
+def test_publish_args(publish, dump_request_args):
+    pubsub_publisher.pubsub_publisher(dump_request_args)
+
+    publish.assert_called_once_with(
+        "projects/<PROJECT_ID>/topics/dag-topic-trigger",
+        dump_request_args.args.get("message").encode("utf-8"),
+        message_length=str(len(dump_request_args.args.get("message"))),
+    )
+
+
+@mock.patch("pubsub_publisher.pubsub_v1.PublisherClient.publish")
+@mock.patch("pubsub_publisher.pubsub_v1.PublisherClient.topic_path")
+def test_topic_path(topic_path, _, dump_request):
+    pubsub_publisher.pubsub_publisher(dump_request)
+
+    topic_path.assert_called_once_with(
+        "<PROJECT_ID>",
+        "dag-topic-trigger",
+    )
+
+
+@mock.patch("pubsub_publisher.pubsub_v1.PublisherClient.publish")
+def test_publish(publish, dump_request):
+    pubsub_publisher.pubsub_publisher(dump_request)
+
+    publish.assert_called_once_with(
+        "projects/<PROJECT_ID>/topics/dag-topic-trigger",
+        dump_request.get_json().get("message").encode("utf-8"),
+        message_length=str(len(dump_request.get_json().get("message"))),
+    )

composer/functions/pubsub_publisher_test.py

@@ -1,2 +1,3 @@
 requests-toolbelt==0.10.0
 google-auth==2.6.2
+google-cloud-pubsub==2.13.11

composer/functions/requirements.txt

@@ -0,0 +1,100 @@
+# Copyright 2023 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.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+""" Two Airflow DAGs that demonstrate the mechanism of triggering DAGs with Pub/Sub messages
+
+    Usage: Replace <PROJECT_ID> with the project ID of your project
+"""
+
+from __future__ import annotations
+
+from datetime import datetime
+import time
+
+from airflow import DAG
+from airflow import XComArg
+from airflow.operators.python import PythonOperator
+from airflow.operators.trigger_dagrun import TriggerDagRunOperator
+from airflow.providers.google.cloud.operators.pubsub import (
+    PubSubCreateSubscriptionOperator,
+    PubSubPullOperator,
+)
+
+PROJECT_ID = "<PROJECT_ID>"
+TOPIC_ID = "dag-topic-trigger"
+SUBSCRIPTION = "trigger_dag_subscription"
+
+
+def handle_messages(pulled_messages, context):
+    dag_ids = list()
+    for idx, m in enumerate(pulled_messages):
+        data = m.message.data.decode('utf-8')
+        print(f'message {idx} data is {data}')
+        dag_ids.append(data)
+    return dag_ids
+
+
+# This DAG will run minutely and handle pub/sub messages by triggering target DAG
+with DAG('trigger_dag',
+         start_date=datetime(2021, 1, 1),
+         schedule_interval="* * * * *",
+         max_active_runs=1,
+         catchup=False) as trigger_dag:
+
+    # If subscription exists, we will use it. If not - create new one
+    subscribe_task = PubSubCreateSubscriptionOperator(task_id="subscribe_task",
+                                                      project_id=PROJECT_ID,
+                                                      topic=TOPIC_ID,
+                                                      subscription=SUBSCRIPTION)
+
+    subscription = subscribe_task.output
+
+    # Proceed maximum 50 messages in callback function handle_messages
+    # Here we acknowledge messages automatically. You can use PubSubHook.acknowledge to acknowledge in downstream tasks
+    # https://airflow.apache.org/docs/apache-airflow-providers-google/stable/_api/airflow/providers/google/cloud/hooks/pubsub/index.html#airflow.providers.google.cloud.hooks.pubsub.PubSubHook.acknowledge
+    pull_messages_operator = PubSubPullOperator(
+        task_id="pull_messages_operator",
+        project_id=PROJECT_ID,
+        ack_messages=True,
+        messages_callback=handle_messages,
+        subscription=subscription,
+        max_messages=50,
+    )
+
+    # Here we use Dynamic Task Mapping to trigger DAGs according to messages content
+    # https://airflow.apache.org/docs/apache-airflow/2.3.0/concepts/dynamic-task-mapping.html
+    trigger_target_dag = TriggerDagRunOperator\
+        .partial(task_id='trigger_target')\
+        .expand(trigger_dag_id=XComArg(pull_messages_operator))
+
+    (subscribe_task >> pull_messages_operator >> trigger_target_dag)
+
+
+def _some_heavy_task():
+    print('Do some operation...')
+    time.sleep(1)
+    print('Done!')
+
+
+# Simple target DAG
+with DAG(
+        'target_dag',
+        start_date=datetime(2022, 1, 1),
+        # Not scheduled, trigger only
+        schedule_interval=None,
+        catchup=False) as target_dag:
+
+    some_heavy_task = PythonOperator(task_id='some_heavy_task',
+                                     python_callable=_some_heavy_task)
+
+    (some_heavy_task)

composer/workflows/pubsub_trigger_response_dag.py

@@ -0,0 +1,25 @@
+# Copyright 2023 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 internal_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 confidence check by the official Airflow
+    docs: https://airflow.incubator.apache.org/tutorial.html#testing
+    """
+    from . import pubsub_trigger_response_dag as module
+    internal_unit_testing.assert_has_valid_dag(module)