From 2a5cfc78df01daaa7f488085b6c6407effb0e09f Mon Sep 17 00:00:00 2001 From: Derrick Williams Date: Thu, 29 Oct 2026 14:02:23 +0000 Subject: [PATCH 1/2] add pubsub to splunk e2e and example tests --- .../yaml/examples/testing/examples_test.py | 10 +- .../pipelines/pubsub_to_splunk.yaml | 87 +++++++++++ .../extended_tests/e2e/pubsub_to_splunk.yaml | 143 ++++++++++++++++++ .../apache_beam/yaml/integration_tests.py | 124 +++++++++++++++ 4 files changed, 361 insertions(+), 3 deletions(-) create mode 100644 sdks/python/apache_beam/yaml/examples/transforms/pipelines/pubsub_to_splunk.yaml create mode 100644 sdks/python/apache_beam/yaml/extended_tests/e2e/pubsub_to_splunk.yaml diff --git a/sdks/python/apache_beam/yaml/examples/testing/examples_test.py b/sdks/python/apache_beam/yaml/examples/testing/examples_test.py index 7c81966a2db3..caf79283e7eb 100644 --- a/sdks/python/apache_beam/yaml/examples/testing/examples_test.py +++ b/sdks/python/apache_beam/yaml/examples/testing/examples_test.py @@ -187,10 +187,12 @@ def test_pubsub_read( if topic == 'test-topic': pubsub_messages = PUBSUB_TOPICS['test-topic'] + schema = RowTypeConstraint.from_fields([('label', str), ('rank', int)]) return ( pcoll | beam.Create([json.loads(msg.data) for msg in pubsub_messages]) - | beam.Map(lambda element: beam.Row(**element))) + | + beam.Map(lambda element: beam.Row(**element)).with_output_types(schema)) elif topic == 'taxi-ride-topic': pubsub_messages = PUBSUB_TOPICS['taxi-ride-topic'] schema = input_data.TaxiRideEventSchema @@ -686,7 +688,8 @@ def _kafka_test_preprocessor( 'test_bigquery_write_yaml', 'test_mongodb_to_bigquery_yaml', 'test_gcs_avro_to_cloud_bigtable_yaml', - 'test_bigquery_to_bigtable_yaml' + 'test_bigquery_to_bigtable_yaml', + 'test_pubsub_to_splunk_yaml' ]) def _io_write_test_preprocessor( test_spec: dict, expected: list[str], env: TestEnvironment): @@ -961,7 +964,8 @@ def _enrichment_test_preprocessor( @YamlExamplesTestSuite.register_test_preprocessor([ 'test_pubsub_topic_to_bigquery_yaml', 'test_pubsub_subscription_to_bigquery_yaml', - 'test_pubsub_to_iceberg_yaml' + 'test_pubsub_to_iceberg_yaml', + 'test_pubsub_to_splunk_yaml' ]) def _pubsub_io_read_test_preprocessor( test_spec: dict, expected: list[str], env: TestEnvironment): diff --git a/sdks/python/apache_beam/yaml/examples/transforms/pipelines/pubsub_to_splunk.yaml b/sdks/python/apache_beam/yaml/examples/transforms/pipelines/pubsub_to_splunk.yaml new file mode 100644 index 000000000000..d34d6e4b5bd8 --- /dev/null +++ b/sdks/python/apache_beam/yaml/examples/transforms/pipelines/pubsub_to_splunk.yaml @@ -0,0 +1,87 @@ +# coding=utf-8 +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You 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 is an example of a Beam YAML pipeline that reads messages from a Pub/Sub +# subscription, transforms them into Splunk HEC events, writes them to Splunk, +# and routes failed messages to a Pub/Sub deadletter topic. This matches the +# Dataflow Template located here - +# https://cloud.google.com/dataflow/docs/guides/templates/provided/pubsub-to-splunk + +pipeline: + type: composite + transforms: + # Step 1: Read messages from a Pub/Sub subscription + - type: ReadFromPubSub + name: ReadMessages + config: + subscription: "projects/apache-beam-testing/subscriptions/my-subscription" + format: JSON + schema: + type: object + properties: + label: {type: string} + rank: {type: integer} + + # Step 2: Transform input messages and map to Splunk HEC event fields + - type: MapToFields + name: ConvertToSplunkEvents + input: ReadMessages + config: + language: python + fields: + time: 1700000000 + host: "'test-host'" + source: "'apache-beam'" + sourcetype: "'test-type'" + index: "'test-index'" + event: + callable: | + def format_event(row): + return "Event for label " + row.label + " with rank " + str(row.rank) + output_type: string + + # Step 3: Write events to Splunk's HTTP Event Collector (HEC) endpoint + - type: WriteToSplunk + name: WriteToSplunk + input: ConvertToSplunkEvents + config: + url: "https://splunk-hec-host:8088" + token: "my-splunk-hec-token" + batch_count: 10 + parallelism: 1 + disable_certificate_validation: false + enable_batch_logs: true + enable_gzip_http_compression: true + error_handling: + output: write_errors + + # Step 4: Write failed records from Splunk HEC writes to a Pub/Sub deadletter topic + - type: WriteToPubSub + name: WriteFailedRecords + input: WriteToSplunk.write_errors + config: + topic: "projects/apache-beam-testing/topics/deadletter-topic" + format: JSON + +options: + streaming: true + +# Expected: +# Row(time=1700000000, host='test-host', source='apache-beam', sourcetype='test-type', index='test-index', event='Event for label 37a with rank 1') +# Row(time=1700000000, host='test-host', source='apache-beam', sourcetype='test-type', index='test-index', event='Event for label 37b with rank 4') +# Row(time=1700000000, host='test-host', source='apache-beam', sourcetype='test-type', index='test-index', event='Event for label 37c with rank 3') +# Row(time=1700000000, host='test-host', source='apache-beam', sourcetype='test-type', index='test-index', event='Event for label 37d with rank 2') diff --git a/sdks/python/apache_beam/yaml/extended_tests/e2e/pubsub_to_splunk.yaml b/sdks/python/apache_beam/yaml/extended_tests/e2e/pubsub_to_splunk.yaml new file mode 100644 index 000000000000..9f9d8494c30e --- /dev/null +++ b/sdks/python/apache_beam/yaml/extended_tests/e2e/pubsub_to_splunk.yaml @@ -0,0 +1,143 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You 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. +# + +fixtures: + - name: PUBSUB_VARS + type: "apache_beam.yaml.integration_tests.temp_pubsub_subscriptions" + config: + project_id: "apache-beam-testing" + - name: FAKE_SPLUNK_SERVER + type: "apache_beam.yaml.test_utils.splunk_test_utils.temp_fake_splunk_server" + config: + expected_records: + - { time: 1700000000, host: "test-host", source: "apache-beam-11a", sourcetype: "test-type", index: "test-index", event: "Event for label 11a with rank 1" } + - { time: 1700000000, host: "test-host", source: "apache-beam-37a", sourcetype: "test-type", index: "test-index", event: "Event for label 37a with rank 2" } + - { time: 1700000000, host: "test-host", source: "apache-beam-389a", sourcetype: "test-type", index: "test-index", event: "Event for label 389a with rank 3" } + - { time: 1700000000, host: "test-host", source: "apache-beam-3821b", sourcetype: "test-type", index: "test-index", event: "Event for label 3821b with rank 4" } + +pipelines: + # Pipeline 1: Publish test messages (valid and invalid) to Pub/Sub input topic + - pipeline: + type: chain + transforms: + - type: Create + config: + elements: + - {label: "11a", rank: 1} + - {label: "37a", rank: 2} + - {label: "389a", rank: 3} + - {label: "3821b", rank: 4} + - {label: "broken", rank: -1} + - type: WriteToPubSub + config: + topic: "{PUBSUB_VARS[INPUT_TOPIC]}" + format: JSON + + # Pipeline 2: Read from Pub/Sub subscription, transform and write to Splunk HEC, + # and route failed records to Pub/Sub deadletter topic + - pipeline: + type: composite + transforms: + - type: ReadFromPubSub + name: ReadMessages + config: + subscription: "{PUBSUB_VARS[INPUT_SUBSCRIPTION]}" + format: JSON + schema: + type: object + properties: + label: {type: string} + rank: {type: integer} + required: + - label + - rank + + - type: MapToFields + name: ConvertToSplunkEvents + input: ReadMessages + config: + language: python + fields: + time: 1700000000 + host: "'test-host'" + source: + expression: "'apache-beam-' + label" + output_type: string + sourcetype: "'test-type'" + index: "'test-index'" + event: + callable: | + from typing import Optional + def format_event(row) -> Optional[str]: + if row.rank < 0: + return None + return str("Event for label " + row.label + " with rank " + str(row.rank)) + + - type: WriteToSplunk + name: WriteToSplunk + input: ConvertToSplunkEvents + config: + url: "{FAKE_SPLUNK_SERVER.url}" + token: "{FAKE_SPLUNK_SERVER.token}" + batch_count: 2 + parallelism: 1 + disable_certificate_validation: false + enable_batch_logs: true + enable_gzip_http_compression: true + error_handling: + output: write_errors + + - type: MapToFields + name: FormatWriteErrors + input: WriteToSplunk.write_errors + config: + language: python + fields: + failed_source: + expression: "failed_row.source" + output_type: string + error_message: + expression: "statusMessage" + output_type: string + + - type: WriteToPubSub + name: WriteFailedRecords + input: FormatWriteErrors + config: + topic: "{PUBSUB_VARS[DEADLETTER_TOPIC]}" + format: JSON + + # Pipeline 3: Read from Pub/Sub deadletter subscription and verify failed records + - pipeline: + type: chain + transforms: + - type: ReadFromPubSub + config: + subscription: "{PUBSUB_VARS[DEADLETTER_SUBSCRIPTION]}" + format: JSON + schema: + type: object + properties: + failed_source: {type: string} + error_message: {type: string} + required: + - failed_source + - error_message + - type: AssertEqual + config: + elements: + - {failed_source: "apache-beam-broken", error_message: "Event is required."} diff --git a/sdks/python/apache_beam/yaml/integration_tests.py b/sdks/python/apache_beam/yaml/integration_tests.py index 26d398bc2bab..c95d8e607519 100644 --- a/sdks/python/apache_beam/yaml/integration_tests.py +++ b/sdks/python/apache_beam/yaml/integration_tests.py @@ -1214,6 +1214,130 @@ def temp_pubsub_emulator(project_id="apache-beam-testing"): yield created_topic_object.name +@contextlib.contextmanager +def temp_pubsub_subscriptions(project_id="apache-beam-testing"): + """Context manager providing a temporary Pub/Sub emulator with input and + deadletter topics and subscriptions, plus bounded ReadFromPubSub and + WriteToPubSub transforms compatible with FnApiRunner for E2E testing. + """ + from google.api_core import exceptions as gcp_exceptions + + with PubSubContainer(project=project_id) as pubsub_container: + emulator_host = pubsub_container.get_pubsub_emulator_host() + with mock.patch.dict(os.environ, {"PUBSUB_EMULATOR_HOST": emulator_host}): + publisher = pubsub_v1.PublisherClient() + subscriber = pubsub_v1.SubscriberClient() + suffix = uuid.uuid4().hex[:8] + + input_topic = f"projects/{pubsub_container.project}/topics/input-{suffix}" + input_sub = ( + f"projects/{pubsub_container.project}/subscriptions/" + f"input-sub-{suffix}") + dlq_topic = ( + f"projects/{pubsub_container.project}/topics/deadletter-{suffix}") + dlq_sub = ( + f"projects/{pubsub_container.project}/subscriptions/" + f"deadletter-sub-{suffix}") + + publisher.create_topic(name=input_topic) + subscriber.create_subscription(name=input_sub, topic=input_topic) + publisher.create_topic(name=dlq_topic) + subscriber.create_subscription(name=dlq_sub, topic=dlq_topic) + + class _BoundedReadFromPubSub(beam.PTransform): + def __init__( + self, + topic=None, + subscription=None, + id_label=None, + with_attributes=False, + timestamp_attribute=None): + super().__init__() + self._topic = topic + self._subscription = subscription + self._with_attributes = with_attributes + + def expand(self, pbegin): + sub_name = self._subscription + with_attrs = self._with_attributes + + class _PullFn(beam.DoFn): + def process(self, _): + os.environ["PUBSUB_EMULATOR_HOST"] = emulator_host + sub_client = pubsub_v1.SubscriberClient() + total_pulled = 0 + try: + while True: + try: + response = sub_client.pull( + request={ + "subscription": sub_name, + "max_messages": 100, + }, + timeout=5.0 if total_pulled == 0 else 1.0) + except gcp_exceptions.DeadlineExceeded: + break + if not response.received_messages: + break + total_pulled += len(response.received_messages) + ack_ids = [m.ack_id for m in response.received_messages] + sub_client.acknowledge( + request={ + "subscription": sub_name, "ack_ids": ack_ids + }) + for received in response.received_messages: + msg = received.message + if with_attrs: + yield beam.io.gcp.pubsub.PubsubMessage( + data=msg.data, attributes=dict(msg.attributes)) + else: + yield msg.data + finally: + sub_client.close() + + return pbegin | beam.Impulse() | beam.ParDo(_PullFn()) + + class _BoundedWriteToPubSub(beam.PTransform): + def __init__( + self, + topic, + with_attributes=False, + id_label=None, + timestamp_attribute=None): + super().__init__() + self._topic = topic + self._with_attributes = with_attributes + + def expand(self, pcoll): + topic_name = self._topic + with_attrs = self._with_attributes + + class _PublishFn(beam.DoFn): + def setup(self): + os.environ["PUBSUB_EMULATOR_HOST"] = emulator_host + self._pub_client = pubsub_v1.PublisherClient() + + def process(self, element): + if with_attrs: + data = element.data + attrs = dict(element.attributes or {}) + else: + data = element + attrs = {} + self._pub_client.publish(topic_name, data, **attrs).result() + + return pcoll | beam.ParDo(_PublishFn()) + + with mock.patch("apache_beam.io.ReadFromPubSub", _BoundedReadFromPubSub), \ + mock.patch("apache_beam.io.WriteToPubSub", _BoundedWriteToPubSub): + yield { + "INPUT_TOPIC": input_topic, + "INPUT_SUBSCRIPTION": input_sub, + "DEADLETTER_TOPIC": dlq_topic, + "DEADLETTER_SUBSCRIPTION": dlq_sub, + } + + @contextlib.contextmanager def temp_delta_table(): with tempfile.TemporaryDirectory() as temp_dir: From 886623994105a31349013543e2afd7554ad91248 Mon Sep 17 00:00:00 2001 From: Derrick Williams Date: Thu, 29 Oct 2026 15:38:00 +0000 Subject: [PATCH 2/2] address gemini comments --- .../extended_tests/e2e/pubsub_to_splunk.yaml | 2 +- .../python/apache_beam/yaml/integration_tests.py | 16 ++++++++++------ 2 files changed, 11 insertions(+), 7 deletions(-) diff --git a/sdks/python/apache_beam/yaml/extended_tests/e2e/pubsub_to_splunk.yaml b/sdks/python/apache_beam/yaml/extended_tests/e2e/pubsub_to_splunk.yaml index 9f9d8494c30e..7bf414ed76a2 100644 --- a/sdks/python/apache_beam/yaml/extended_tests/e2e/pubsub_to_splunk.yaml +++ b/sdks/python/apache_beam/yaml/extended_tests/e2e/pubsub_to_splunk.yaml @@ -108,7 +108,7 @@ pipelines: language: python fields: failed_source: - expression: "failed_row.source" + expression: "failed_row.source if failed_row else ''" output_type: string error_message: expression: "statusMessage" diff --git a/sdks/python/apache_beam/yaml/integration_tests.py b/sdks/python/apache_beam/yaml/integration_tests.py index c95d8e607519..1dec87fd400b 100644 --- a/sdks/python/apache_beam/yaml/integration_tests.py +++ b/sdks/python/apache_beam/yaml/integration_tests.py @@ -1225,8 +1225,6 @@ def temp_pubsub_subscriptions(project_id="apache-beam-testing"): with PubSubContainer(project=project_id) as pubsub_container: emulator_host = pubsub_container.get_pubsub_emulator_host() with mock.patch.dict(os.environ, {"PUBSUB_EMULATOR_HOST": emulator_host}): - publisher = pubsub_v1.PublisherClient() - subscriber = pubsub_v1.SubscriberClient() suffix = uuid.uuid4().hex[:8] input_topic = f"projects/{pubsub_container.project}/topics/input-{suffix}" @@ -1239,10 +1237,12 @@ def temp_pubsub_subscriptions(project_id="apache-beam-testing"): f"projects/{pubsub_container.project}/subscriptions/" f"deadletter-sub-{suffix}") - publisher.create_topic(name=input_topic) - subscriber.create_subscription(name=input_sub, topic=input_topic) - publisher.create_topic(name=dlq_topic) - subscriber.create_subscription(name=dlq_sub, topic=dlq_topic) + with pubsub_v1.PublisherClient() as publisher, \ + pubsub_v1.SubscriberClient() as subscriber: + publisher.create_topic(name=input_topic) + subscriber.create_subscription(name=input_sub, topic=input_topic) + publisher.create_topic(name=dlq_topic) + subscriber.create_subscription(name=dlq_sub, topic=dlq_topic) class _BoundedReadFromPubSub(beam.PTransform): def __init__( @@ -1326,6 +1326,10 @@ def process(self, element): attrs = {} self._pub_client.publish(topic_name, data, **attrs).result() + def teardown(self): + if hasattr(self, "_pub_client") and self._pub_client is not None: + self._pub_client.transport.close() + return pcoll | beam.ParDo(_PublishFn()) with mock.patch("apache_beam.io.ReadFromPubSub", _BoundedReadFromPubSub), \