This is an automated email from the ASF dual-hosted git repository.
potiuk pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new df0aa3fbe1f feature: Add OpenLineage support for
PubSubPublishMessageOperator (#54764)
df0aa3fbe1f is described below
commit df0aa3fbe1f1580a5a4ad3edba40fda8a8bba92d
Author: pawelgrochowicz <[email protected]>
AuthorDate: Thu Aug 21 17:42:28 2025 +0200
feature: Add OpenLineage support for PubSubPublishMessageOperator (#54764)
---
.../providers/google/cloud/operators/pubsub.py | 17 ++++++++++--
.../unit/google/cloud/operators/test_pubsub.py | 32 ++++++++++++++++++++++
2 files changed, 46 insertions(+), 3 deletions(-)
diff --git
a/providers/google/src/airflow/providers/google/cloud/operators/pubsub.py
b/providers/google/src/airflow/providers/google/cloud/operators/pubsub.py
index 100c479593b..b55706e2ab8 100644
--- a/providers/google/src/airflow/providers/google/cloud/operators/pubsub.py
+++ b/providers/google/src/airflow/providers/google/cloud/operators/pubsub.py
@@ -713,17 +713,28 @@ class
PubSubPublishMessageOperator(GoogleCloudBaseOperator):
self.enable_message_ordering = enable_message_ordering
self.impersonation_chain = impersonation_chain
- def execute(self, context: Context) -> None:
- hook = PubSubHook(
+ @cached_property
+ def pubsub_hook(self):
+ return PubSubHook(
gcp_conn_id=self.gcp_conn_id,
impersonation_chain=self.impersonation_chain,
enable_message_ordering=self.enable_message_ordering,
)
+ def execute(self, context: Context) -> None:
self.log.info("Publishing to topic %s", self.topic)
- hook.publish(project_id=self.project_id, topic=self.topic,
messages=self.messages)
+ self.pubsub_hook.publish(project_id=self.project_id, topic=self.topic,
messages=self.messages)
self.log.info("Published to topic %s", self.topic)
+ def get_openlineage_facets_on_complete(self, _) -> OperatorLineage:
+ from airflow.providers.common.compat.openlineage.facet import Dataset
+ from airflow.providers.openlineage.extractors import OperatorLineage
+
+ project_id = self.project_id or self.pubsub_hook.project_id
+ output_dataset = [Dataset(namespace="pubsub",
name=f"topic:{project_id}:{self.topic}")]
+
+ return OperatorLineage(outputs=output_dataset)
+
class PubSubPullOperator(GoogleCloudBaseOperator):
"""
diff --git a/providers/google/tests/unit/google/cloud/operators/test_pubsub.py
b/providers/google/tests/unit/google/cloud/operators/test_pubsub.py
index d62ef1ae475..ea2865aefa1 100644
--- a/providers/google/tests/unit/google/cloud/operators/test_pubsub.py
+++ b/providers/google/tests/unit/google/cloud/operators/test_pubsub.py
@@ -329,6 +329,38 @@ class TestPubSubPublishOperator:
project_id=TEST_PROJECT, topic=TEST_TOPIC,
messages=TEST_MESSAGES_ORDERING_KEY
)
+ @pytest.mark.parametrize(
+ "project_id, expected_dataset",
+ [
+ # 1. project_id provided
+ (TEST_PROJECT, f"topic:{TEST_PROJECT}:{TEST_TOPIC}"),
+ # 2. project_id not provided (use project_id from connection)
+ (None, f"topic:connection-project:{TEST_TOPIC}"),
+ ],
+ )
+ @mock.patch("airflow.providers.google.cloud.operators.pubsub.PubSubHook")
+ def test_get_openlineage_facets(self, mock_hook, project_id,
expected_dataset):
+ operator = PubSubPublishMessageOperator(
+ task_id=TASK_ID,
+ project_id=project_id,
+ topic=TEST_TOPIC,
+ messages=TEST_MESSAGES,
+ )
+
+ operator.execute(None)
+ mock_hook.return_value.publish.assert_called_once_with(
+ project_id=project_id, topic=TEST_TOPIC, messages=TEST_MESSAGES
+ )
+ mock_hook.return_value.project_id = project_id or "connection-project"
+
+ result = operator.get_openlineage_facets_on_complete(operator)
+ assert not result.run_facets
+ assert not result.job_facets
+ assert len(result.inputs) == 0
+ assert len(result.outputs) == 1
+ assert result.outputs[0].namespace == "pubsub"
+ assert result.outputs[0].name == expected_dataset
+
class TestPubSubPullOperator:
def _generate_messages(self, count):
