Thanks to visit codestin.com
Credit goes to github.com

Skip to content

Commit a3d9a80

Browse files
dmitriy-drogovozdmitriy-drogovoz
authored
Composer tutorial: add function and dag (#8860)
1 parent c954b52 commit a3d9a80

File tree

5 files changed

+294
-0
lines changed

5 files changed

+294
-0
lines changed

‎composer/functions/pubsub_publisher.py‎

Lines changed: 59 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,59 @@
1+
# Copyright 2023 Google LLC.
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
""" This Cloud Function example creates Pub/Sub messages.
15+
16+
Usage: Replace <PROJECT_ID> with the project ID of your project.
17+
"""
18+
19+
# [START composer_pubsub_publisher_function]
20+
from google.cloud import pubsub_v1
21+
22+
project = "<PROJECT_ID>"
23+
topic = "dag-topic-trigger"
24+
25+
26+
def pubsub_publisher(request):
27+
"""Publish message from HTTP request to Pub/Sub topic.
28+
Args:
29+
request (flask.Request): HTTP request object.
30+
Returns:
31+
The response text with message published into Pub/Sub topic
32+
Response object using
33+
`make_response <http://flask.pocoo.org/docs/1.0/api/#flask.Flask.make_response>`.
34+
"""
35+
request_json = request.get_json()
36+
print(request_json)
37+
if request.args and 'message' in request.args:
38+
data_str = request.args.get('message')
39+
elif request_json and 'message' in request_json:
40+
data_str = request_json['message']
41+
else:
42+
return "Message content not found! Use 'message' key to specify"
43+
44+
publisher = pubsub_v1.PublisherClient()
45+
# The `topic_path` method creates a fully qualified identifier
46+
# in the form `projects/{project_id}/topics/{topic_id}`
47+
topic_path = publisher.topic_path(project, topic)
48+
49+
# The required data format is a bytestring
50+
data = data_str.encode("utf-8")
51+
# When you publish a message, the client returns a future.
52+
message_length = len(data_str)
53+
future = publisher.publish(topic_path,
54+
data,
55+
message_length=str(message_length))
56+
print(future.result())
57+
58+
return f"Message {data} with message_length {message_length} published to {topic_path}."
59+
# [END composer_pubsub_publisher_function]

‎composer/functions/pubsub_publisher_test.py‎

Lines changed: 107 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,107 @@
1+
# Copyright 2023 Google LLC.
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
import mock
16+
import pytest
17+
18+
import pubsub_publisher
19+
20+
21+
@pytest.fixture()
22+
def dump_request_args():
23+
class Request:
24+
args = {"message": "test with args"}
25+
26+
def get_json(self):
27+
return self.args
28+
29+
return Request()
30+
31+
32+
@pytest.fixture()
33+
def dump_request():
34+
class Request:
35+
args = None
36+
37+
def get_json(self):
38+
return {"message": "test with no args"}
39+
40+
return Request()
41+
42+
43+
@pytest.fixture()
44+
def dump_request_no_message():
45+
class Request:
46+
args = None
47+
48+
def get_json(self):
49+
return {"no_message": "test with no message key"}
50+
51+
return Request()
52+
53+
54+
# Pass None, an input that is not valid request
55+
def test_request_with_none():
56+
request = None
57+
with pytest.raises(Exception):
58+
pubsub_publisher.pubsub_publisher(request)
59+
60+
61+
def test_content_not_found(dump_request_no_message):
62+
output = "Message content not found! Use 'message' key to specify"
63+
assert pubsub_publisher.pubsub_publisher(dump_request_no_message) == output, f"The function didn't return '{output}'"
64+
65+
66+
@mock.patch("pubsub_publisher.pubsub_v1.PublisherClient.publish")
67+
@mock.patch("pubsub_publisher.pubsub_v1.PublisherClient.topic_path")
68+
def test_topic_path_args(topic_path, _, dump_request_args):
69+
pubsub_publisher.pubsub_publisher(dump_request_args)
70+
71+
topic_path.assert_called_once_with(
72+
"<PROJECT_ID>",
73+
"dag-topic-trigger",
74+
)
75+
76+
77+
@mock.patch("pubsub_publisher.pubsub_v1.PublisherClient.publish")
78+
def test_publish_args(publish, dump_request_args):
79+
pubsub_publisher.pubsub_publisher(dump_request_args)
80+
81+
publish.assert_called_once_with(
82+
"projects/<PROJECT_ID>/topics/dag-topic-trigger",
83+
dump_request_args.args.get("message").encode("utf-8"),
84+
message_length=str(len(dump_request_args.args.get("message"))),
85+
)
86+
87+
88+
@mock.patch("pubsub_publisher.pubsub_v1.PublisherClient.publish")
89+
@mock.patch("pubsub_publisher.pubsub_v1.PublisherClient.topic_path")
90+
def test_topic_path(topic_path, _, dump_request):
91+
pubsub_publisher.pubsub_publisher(dump_request)
92+
93+
topic_path.assert_called_once_with(
94+
"<PROJECT_ID>",
95+
"dag-topic-trigger",
96+
)
97+
98+
99+
@mock.patch("pubsub_publisher.pubsub_v1.PublisherClient.publish")
100+
def test_publish(publish, dump_request):
101+
pubsub_publisher.pubsub_publisher(dump_request)
102+
103+
publish.assert_called_once_with(
104+
"projects/<PROJECT_ID>/topics/dag-topic-trigger",
105+
dump_request.get_json().get("message").encode("utf-8"),
106+
message_length=str(len(dump_request.get_json().get("message"))),
107+
)

‎composer/functions/requirements.txt‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,2 +1,3 @@
11
requests-toolbelt==0.10.0
22
google-auth==2.6.2
3+
google-cloud-pubsub==2.13.11

‎composer/workflows/pubsub_trigger_response_dag.py‎

Lines changed: 102 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,102 @@
1+
# Copyright 2023 Google LLC.
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
""" Two Airflow DAGs that demonstrate the mechanism of triggering DAGs with Pub/Sub messages
15+
16+
Usage: Replace <PROJECT_ID> with the project ID of your project
17+
"""
18+
19+
# [START composer_pubsub_trigger_dag]
20+
from __future__ import annotations
21+
22+
from datetime import datetime
23+
import time
24+
25+
from airflow import DAG
26+
from airflow import XComArg
27+
from airflow.operators.python import PythonOperator
28+
from airflow.operators.trigger_dagrun import TriggerDagRunOperator
29+
from airflow.providers.google.cloud.operators.pubsub import (
30+
PubSubCreateSubscriptionOperator,
31+
PubSubPullOperator,
32+
)
33+
34+
PROJECT_ID = "<PROJECT_ID>"
35+
TOPIC_ID = "dag-topic-trigger"
36+
SUBSCRIPTION = "trigger_dag_subscription"
37+
38+
39+
def handle_messages(pulled_messages, context):
40+
dag_ids = list()
41+
for idx, m in enumerate(pulled_messages):
42+
data = m.message.data.decode('utf-8')
43+
print(f'message {idx} data is {data}')
44+
dag_ids.append(data)
45+
return dag_ids
46+
47+
48+
# This DAG will run minutely and handle pub/sub messages by triggering target DAG
49+
with DAG('trigger_dag',
50+
start_date=datetime(2021, 1, 1),
51+
schedule_interval="* * * * *",
52+
max_active_runs=1,
53+
catchup=False) as trigger_dag:
54+
55+
# If subscription exists, we will use it. If not - create new one
56+
subscribe_task = PubSubCreateSubscriptionOperator(task_id="subscribe_task",
57+
project_id=PROJECT_ID,
58+
topic=TOPIC_ID,
59+
subscription=SUBSCRIPTION)
60+
61+
subscription = subscribe_task.output
62+
63+
# Proceed maximum 50 messages in callback function handle_messages
64+
# Here we acknowledge messages automatically. You can use PubSubHook.acknowledge to acknowledge in downstream tasks
65+
# 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
66+
pull_messages_operator = PubSubPullOperator(
67+
task_id="pull_messages_operator",
68+
project_id=PROJECT_ID,
69+
ack_messages=True,
70+
messages_callback=handle_messages,
71+
subscription=subscription,
72+
max_messages=50,
73+
)
74+
75+
# Here we use Dynamic Task Mapping to trigger DAGs according to messages content
76+
# https://airflow.apache.org/docs/apache-airflow/2.3.0/concepts/dynamic-task-mapping.html
77+
trigger_target_dag = TriggerDagRunOperator\
78+
.partial(task_id='trigger_target')\
79+
.expand(trigger_dag_id=XComArg(pull_messages_operator))
80+
81+
(subscribe_task >> pull_messages_operator >> trigger_target_dag)
82+
83+
84+
def _some_heavy_task():
85+
print('Do some operation...')
86+
time.sleep(1)
87+
print('Done!')
88+
89+
90+
# Simple target DAG
91+
with DAG(
92+
'target_dag',
93+
start_date=datetime(2022, 1, 1),
94+
# Not scheduled, trigger only
95+
schedule_interval=None,
96+
catchup=False) as target_dag:
97+
98+
some_heavy_task = PythonOperator(task_id='some_heavy_task',
99+
python_callable=_some_heavy_task)
100+
101+
(some_heavy_task)
102+
# [END composer_pubsub_trigger_dag]

‎composer/workflows/pubsub_trigger_response_dag_test.py‎

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,25 @@
1+
# Copyright 2023 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# https://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
import internal_unit_testing
16+
17+
18+
def test_dag_import():
19+
"""Test that the DAG file can be successfully imported.
20+
This tests that the DAG can be parsed, but does not run it in an Airflow
21+
environment. This is a recommended confidence check by the official Airflow
22+
docs: https://airflow.incubator.apache.org/tutorial.html#testing
23+
"""
24+
from . import pubsub_trigger_response_dag as module
25+
internal_unit_testing.assert_has_valid_dag(module)

0 commit comments

Comments
 (0)