diff --git a/sdks/python/apache_beam/io/gcp/pubsub.py b/sdks/python/apache_beam/io/gcp/pubsub.py index 55856f504787..e15d4027f61a 100644 --- a/sdks/python/apache_beam/io/gcp/pubsub.py +++ b/sdks/python/apache_beam/io/gcp/pubsub.py @@ -32,6 +32,7 @@ # pytype: skip-file import logging +import random import re import time from typing import Any @@ -43,13 +44,16 @@ from apache_beam.io import iobase from apache_beam.io.iobase import Read from apache_beam.metrics.metric import Lineage +from apache_beam.transforms import Create from apache_beam.transforms import DoFn from apache_beam.transforms import Flatten from apache_beam.transforms import Map from apache_beam.transforms import ParDo from apache_beam.transforms import PTransform +from apache_beam.transforms import window from apache_beam.transforms.display import DisplayDataItem from apache_beam.utils.annotations import deprecated +from apache_beam.utils.timestamp import Timestamp try: from google.cloud import pubsub @@ -170,8 +174,8 @@ def _to_proto_str(self, for_publish=False): if not for_publish: if self.message_id: message_id = self.message_id - if self.publish_time: - publish_time = self.publish_time + if self.publish_time: + publish_time = self.publish_time if len(self.ordering_key) > 1024: raise ValueError( @@ -218,7 +222,8 @@ def __init__( subscription: Optional[str] = None, id_label: Optional[str] = None, with_attributes: bool = False, - timestamp_attribute: Optional[str] = None) -> None: + timestamp_attribute: Optional[str] = None, + max_read_time_seconds: Optional[int] = None) -> None: """Initializes ``ReadFromPubSub``. Args: @@ -250,19 +255,38 @@ def __init__( ``2015-10-29T23:41:41.123Z``. The sub-second component of the timestamp is optional, and digits beyond the first three (i.e., time units smaller than milliseconds) may be ignored. + max_read_time_seconds: Maximum time in seconds to read from Pub/Sub. + Default is forever. """ super().__init__() + if max_read_time_seconds is not None and max_read_time_seconds <= 0: + raise ValueError('max_read_time_seconds must be positive.') + self.max_read_time_seconds = max_read_time_seconds self.with_attributes = with_attributes self._source = _PubSubSource( topic=topic, subscription=subscription, id_label=id_label, with_attributes=self.with_attributes, - timestamp_attribute=timestamp_attribute) + timestamp_attribute=timestamp_attribute, + max_read_time_seconds=self.max_read_time_seconds) def expand(self, pvalue): - # TODO(BEAM-27443): Apply a proper transform rather than Read. - pcoll = pvalue.pipeline | Read(self._source) + if self.max_read_time_seconds is not None: + sub_project = None + if pvalue.pipeline and pvalue.pipeline.options: + from apache_beam.options.pipeline_options import GoogleCloudOptions + sub_project = pvalue.pipeline.options.view_as( + GoogleCloudOptions).project + if not sub_project: + sub_project = self._source.project + pcoll = ( + pvalue.pipeline + | Create([None]) + | ParDo(_PubSubBoundedReadDoFn(self._source, sub_project))) + else: + # TODO(BEAM-27443): Apply a proper transform rather than Read. + pcoll = pvalue.pipeline | Read(self._source) # explicit element_type required after native read, otherwise coder error pcoll.element_type = bytes return self.expand_continued(pcoll) @@ -523,7 +547,8 @@ def __init__( subscription: Optional[str] = None, id_label: Optional[str] = None, with_attributes: bool = False, - timestamp_attribute: Optional[str] = None): + timestamp_attribute: Optional[str] = None, + max_read_time_seconds: Optional[int] = None): self.coder = coders.BytesCoder() self.full_topic = topic self.full_subscription = subscription @@ -532,6 +557,7 @@ def __init__( self.id_label = id_label self.with_attributes = with_attributes self.timestamp_attribute = timestamp_attribute + self.max_read_time_seconds = max_read_time_seconds # Perform some validation on the topic and subscription. if not (topic or subscription): @@ -557,13 +583,155 @@ def display_data(self): 'timestamp_attribute': DisplayDataItem( self.timestamp_attribute, label='Timestamp Attribute').drop_if_none(), + 'max_read_time_seconds': DisplayDataItem( + self.max_read_time_seconds, + label='Max Read Time Seconds').drop_if_none(), } def default_output_coder(self): return self.coder def is_bounded(self): - return False + return self.max_read_time_seconds is not None + + +def _parse_pubsub_message( + message: Any, + timestamp_attribute: Optional[str] = None +) -> tuple[Timestamp, 'PubsubMessage']: + parsed_message = PubsubMessage._from_message(message) + if (timestamp_attribute and timestamp_attribute in parsed_message.attributes): + rfc3339_or_milli = parsed_message.attributes[timestamp_attribute] + try: + timestamp = Timestamp(micros=int(rfc3339_or_milli) * 1000) + except ValueError: + try: + timestamp = Timestamp.from_rfc3339(rfc3339_or_milli) + except ValueError as e: + raise ValueError('Bad timestamp value: %s' % e) + if timestamp.precision() > Timestamp.MICROS_PRECISION: + # Element timestamps are limited to microsecond resolution, so + # ignore sub-microsecond digits, as the Dataflow service does. + timestamp = timestamp.to_precision( + Timestamp.MICROS_PRECISION, allow_lossy_conversion=True) + else: + if message.publish_time is None: + raise ValueError('No publish time present in message: %s' % message) + try: + timestamp = Timestamp.from_utc_datetime(message.publish_time) + except ValueError as e: + raise ValueError('Bad timestamp value for message %s: %s', message, e) + + return timestamp, parsed_message + + +class _PubSubBoundedReadDoFn(DoFn): + """DoFn for bounded reading from Cloud Pub/Sub up to max_read_time_seconds.""" + MAX_MESSAGES_PER_PULL = 100 + DEFAULT_PULL_TIMEOUT_SECS = 30 + + def __init__( + self, source: _PubSubSource, sub_project: Optional[str] = None) -> None: + if source.id_label: + raise NotImplementedError( + 'id_label is not supported for bounded PubSub reads') + self._source = source + self.project = source.project + self.topic_name = source.topic_name + self.subscription_name = source.subscription_name + self.sub_project = sub_project or source.project + self.with_attributes = source.with_attributes + self.timestamp_attribute = source.timestamp_attribute + self.max_read_time_seconds = source.max_read_time_seconds + self._sub_client = None + + def display_data(self): + return self._source.display_data() + + def setup(self): + from google.cloud import pubsub + self._sub_client = pubsub.SubscriberClient() + + def process(self, unused_element): + if self._sub_client is None: + from google.cloud import pubsub + self._sub_client = pubsub.SubscriberClient() + + temp_sub_name = None + if self.subscription_name: + sub_name = self._sub_client.subscription_path( + self.project, self.subscription_name) + else: + sub_name = self._sub_client.subscription_path( + self.sub_project, + 'beam_%d_%x' % (int(time.time()), random.randrange(1 << 32))) + topic_name = self._sub_client.topic_path(self.project, self.topic_name) + self._sub_client.create_subscription(name=sub_name, topic=topic_name) + temp_sub_name = sub_name + + try: + start_time = time.time() + while True: + elapsed = time.time() - start_time + remaining = self.max_read_time_seconds - elapsed + if remaining <= 0: + break + timeout = min(self.DEFAULT_PULL_TIMEOUT_SECS, remaining) + try: + response = self._sub_client.pull( + subscription=sub_name, + max_messages=self.MAX_MESSAGES_PER_PULL, + timeout=timeout) + results = [ + _parse_pubsub_message(rm.message, self.timestamp_attribute) + for rm in response.received_messages + ] + ack_ids = [rm.ack_id for rm in response.received_messages] + if ack_ids: + self._sub_client.acknowledge(subscription=sub_name, ack_ids=ack_ids) + except Exception as e: + try: + from google.api_core import exceptions + if isinstance(e, exceptions.DeadlineExceeded) or type( + e).__name__ == 'DeadlineExceeded': + results = [] + else: + raise + except ImportError: + if type(e).__name__ == 'DeadlineExceeded': + results = [] + else: + raise + + for msg_timestamp, parsed_message in results: + if self.with_attributes: + payload = parsed_message._to_proto_str(for_publish=False) + else: + payload = parsed_message.data + yield window.TimestampedValue(payload, msg_timestamp) + + if not results: + remaining = self.max_read_time_seconds - (time.time() - start_time) + if remaining > 0: + time.sleep(min(0.1, remaining)) + finally: + if temp_sub_name: + try: + self._sub_client.delete_subscription(subscription=temp_sub_name) + except Exception: + logging.warning( + 'Failed to delete temporary Pub/Sub subscription %s', + temp_sub_name, + exc_info=True) + + def teardown(self): + if self._sub_client is not None: + try: + self._sub_client.close() + except Exception: + logging.warning( + 'Failed to close Pub/Sub subscriber client', exc_info=True) + self._sub_client = None class _PubSubWriteDoFn(DoFn): diff --git a/sdks/python/apache_beam/io/gcp/pubsub_test.py b/sdks/python/apache_beam/io/gcp/pubsub_test.py index 5a8f2bf617b8..749d70a3d7df 100644 --- a/sdks/python/apache_beam/io/gcp/pubsub_test.py +++ b/sdks/python/apache_beam/io/gcp/pubsub_test.py @@ -454,6 +454,35 @@ def test_display_data_no_subscription(self): hc.assert_that(dd.items, hc.contains_inanyorder(*expected_items)) + def test_display_data_max_read_time_seconds(self): + source = _PubSubSource( + 'projects/fakeprj/topics/a_topic', max_read_time_seconds=60) + dd = DisplayData.create_from(source) + expected_items = [ + DisplayDataItemMatcher('topic', 'projects/fakeprj/topics/a_topic'), + DisplayDataItemMatcher('with_attributes', False), + DisplayDataItemMatcher('max_read_time_seconds', 60), + ] + + hc.assert_that(dd.items, hc.contains_inanyorder(*expected_items)) + + def test_invalid_max_read_time_seconds(self): + with self.assertRaisesRegex(ValueError, + 'max_read_time_seconds must be positive'): + ReadFromPubSub('projects/fakeprj/topics/a_topic', max_read_time_seconds=0) + with self.assertRaisesRegex(ValueError, + 'max_read_time_seconds must be positive'): + ReadFromPubSub( + 'projects/fakeprj/topics/a_topic', max_read_time_seconds=-10) + + def test_is_bounded(self): + unbounded_source = _PubSubSource('projects/fakeprj/topics/a_topic') + self.assertFalse(unbounded_source.is_bounded()) + + bounded_source = _PubSubSource( + 'projects/fakeprj/topics/a_topic', max_read_time_seconds=30) + self.assertTrue(bounded_source.is_bounded()) + @unittest.skipIf(pubsub is None, 'GCP dependencies are not installed') class TestPubSubSink(unittest.TestCase): @@ -506,6 +535,7 @@ class TestReadFromPubSub(unittest.TestCase): def setUp(self): _PubSubReadEvaluator._subscription_cache.clear() _PubSubReadEvaluator._subscriber_client_cache.clear() + _PubSubReadEvaluator._start_times.clear() def test_subscriber_client_is_reused_for_transform(self, mock_pubsub): class Transform(object): @@ -959,6 +989,162 @@ def test_read_from_pubsub_no_overwrite(self, unused_mock): Lineage.query(p.result.metrics(), Lineage.SOURCE), set([f"pubsub:{test_case}:fakeprj.topic_or_sub"])) + def test_read_bounded_batch_mode_topic_success(self, mock_pubsub): + data = b'bounded_data' + publish_time_secs = 1520861821 + publish_time_nanos = 234567000 + ack_id = 'ack_id_1' + pull_response = test_utils.create_pull_response([ + test_utils.PullResponseMessage( + data, {}, publish_time_secs, publish_time_nanos, ack_id) + ]) + empty_response = test_utils.create_pull_response([]) + client = mock_pubsub.return_value + client.subscription_path.return_value = ( + 'projects/fakeprj/subscriptions/beam_temp_sub') + client.topic_path.return_value = 'projects/fakeprj/topics/a_topic' + client.pull.side_effect = ( + lambda **kwargs: pull_response + if client.pull.call_count == 1 else empty_response) + + options = PipelineOptions([]) + options.view_as(StandardOptions).streaming = False + with TestPipeline(options=options) as p: + pcoll = p | ReadFromPubSub( + topic='projects/fakeprj/topics/a_topic', max_read_time_seconds=1) + self.assertTrue(pcoll.is_bounded) + assert_that(pcoll, equal_to([data])) + + client.create_subscription.assert_called_once_with( + name='projects/fakeprj/subscriptions/beam_temp_sub', + topic='projects/fakeprj/topics/a_topic') + client.acknowledge.assert_called_once_with( + subscription='projects/fakeprj/subscriptions/beam_temp_sub', + ack_ids=[ack_id]) + client.delete_subscription.assert_called_once_with( + subscription='projects/fakeprj/subscriptions/beam_temp_sub') + client.close.assert_called_once_with() + + def test_read_bounded_batch_mode_subscription_with_attributes( + self, mock_pubsub): + data = b'bounded_msg' + attributes = {'key': 'value', 'time': '2018-03-12T13:37:01.234567Z'} + publish_time_secs = 1337000000 + publish_time_nanos = 133700000 + ack_id = 'ack_id_sub' + pull_response = test_utils.create_pull_response([ + test_utils.PullResponseMessage( + data, attributes, publish_time_secs, publish_time_nanos, ack_id) + ]) + empty_response = test_utils.create_pull_response([]) + client = mock_pubsub.return_value + client.subscription_path.return_value = ( + 'projects/fakeprj/subscriptions/existing_sub') + client.pull.side_effect = ( + lambda **kwargs: pull_response + if client.pull.call_count == 1 else empty_response) + + expected_elements = [ + TestWindowedValue( + PubsubMessage(data, attributes), + timestamp.Timestamp.from_rfc3339(attributes['time']), + [window.GlobalWindow()]), + ] + + options = PipelineOptions([]) + options.view_as(StandardOptions).streaming = False + with TestPipeline(options=options) as p: + pcoll = p | ReadFromPubSub( + subscription='projects/fakeprj/subscriptions/existing_sub', + with_attributes=True, + timestamp_attribute='time', + max_read_time_seconds=1) + self.assertTrue(pcoll.is_bounded) + assert_that(pcoll, equal_to(expected_elements), reify_windows=True) + + client.create_subscription.assert_not_called() + client.delete_subscription.assert_not_called() + client.acknowledge.assert_called_once_with( + subscription='projects/fakeprj/subscriptions/existing_sub', + ack_ids=[ack_id]) + + def test_read_bounded_deadline_exceeded(self, mock_pubsub): + class DeadlineExceeded(Exception): + pass + + client = mock_pubsub.return_value + client.subscription_path.return_value = ( + 'projects/fakeprj/subscriptions/existing_sub') + client.pull.side_effect = DeadlineExceeded(' deadline exceeded ') + + options = PipelineOptions([]) + options.view_as(StandardOptions).streaming = False + with TestPipeline(options=options) as p: + pcoll = p | ReadFromPubSub( + subscription='projects/fakeprj/subscriptions/existing_sub', + max_read_time_seconds=1) + assert_that(pcoll, equal_to([])) + + client.acknowledge.assert_not_called() + + def test_read_bounded_id_label_unsupported(self, unused_mock_pubsub): + options = PipelineOptions([]) + options.view_as(StandardOptions).streaming = False + with self.assertRaisesRegex(NotImplementedError, + r'id_label is not supported'): + with TestPipeline(options=options) as p: + _ = p | ReadFromPubSub( + topic='projects/fakeprj/topics/a_topic', + id_label='a_label', + max_read_time_seconds=10) + + def test_evaluator_max_read_time_seconds_timeout_and_completion( + self, mock_pubsub): + from apache_beam.runners.direct.watermark_manager import WatermarkManager + + class DeadlineExceeded(Exception): + pass + + mock_pubsub.subscription_path.return_value = ( + 'projects/fakeprj/subscriptions/a_sub') + client = mock_pubsub.return_value + client.pull.side_effect = DeadlineExceeded('deadline exceeded') + + source = _PubSubSource( + subscription='projects/fakeprj/subscriptions/a_sub', + max_read_time_seconds=10) + applied_ptransform = mock.MagicMock() + applied_ptransform.transform._source = source + applied_ptransform.inputs = [mock.MagicMock()] + evaluation_context = mock.MagicMock() + evaluation_context.pipeline_options.view_as.return_value.project = 'fakeprj' + + with mock.patch('apache_beam.runners.direct.transform_evaluator.time.time' + ) as (mock_time): + # Initialize at t=100; first pull at t=104 (6s remaining) -> not yet done + mock_time.return_value = 100.0 + evaluator = _PubSubReadEvaluator( + evaluation_context, applied_ptransform, None, []) + mock_time.return_value = 104.0 + result = evaluator.finish_bundle() + + client.pull.assert_called_once_with( + subscription='projects/fakeprj/subscriptions/a_sub', + max_messages=10, + timeout=6.0) + self.assertEqual(len(result.unprocessed_bundles), 1) + + # Next bundle at t=111 (11s elapsed >= 10s) -> completes with +inf + client.pull.reset_mock() + mock_time.return_value = 111.0 + result_done = evaluator.finish_bundle() + + client.pull.assert_not_called() + self.assertEqual(result_done.unprocessed_bundles, []) + self.assertEqual( + result_done.keyed_watermark_holds, + {None: WatermarkManager.WATERMARK_POS_INF}) + @unittest.skipIf(pubsub is None, 'GCP dependencies are not installed') @mock.patch('google.cloud.pubsub.PublisherClient') diff --git a/sdks/python/apache_beam/runners/direct/direct_runner.py b/sdks/python/apache_beam/runners/direct/direct_runner.py index 443dbd63063d..55a0ea61428b 100644 --- a/sdks/python/apache_beam/runners/direct/direct_runner.py +++ b/sdks/python/apache_beam/runners/direct/direct_runner.py @@ -146,6 +146,10 @@ def accept(self, pipeline, is_interactive): self.supported_by_prism_runner = False return self.supported_by_prism_runner + def enter_composite_transform(self, applied_ptransform): + if isinstance(applied_ptransform.transform, ReadFromPubSub): + self.supported_by_prism_runner = False + def visit_transform(self, applied_ptransform): transform = applied_ptransform.transform # Python SDK assumes the direct runner TestStream implementation is @@ -527,8 +531,13 @@ def _get_pubsub_transform_overrides(pipeline_options): class ReadFromPubSubOverride(PTransformOverride): def matches(self, applied_ptransform): - return isinstance( - applied_ptransform.transform, beam_pubsub.ReadFromPubSub) + if not isinstance(applied_ptransform.transform, + beam_pubsub.ReadFromPubSub): + return False + if (not pipeline_options.view_as(StandardOptions).streaming and + applied_ptransform.transform.max_read_time_seconds is not None): + return False + return True def get_replacement_transform_for_applied_ptransform( self, applied_ptransform): diff --git a/sdks/python/apache_beam/runners/direct/transform_evaluator.py b/sdks/python/apache_beam/runners/direct/transform_evaluator.py index d60c840e8c74..487cdfc35b13 100644 --- a/sdks/python/apache_beam/runners/direct/transform_evaluator.py +++ b/sdks/python/apache_beam/runners/direct/transform_evaluator.py @@ -618,6 +618,7 @@ class _PubSubReadEvaluator(_TransformEvaluator): _subscription_cache = weakref.WeakKeyDictionary() _subscriber_client_cache = weakref.WeakKeyDictionary() _subscriber_client_cache_lock = threading.Lock() + _start_times = weakref.WeakKeyDictionary() def __init__( self, @@ -637,6 +638,10 @@ def __init__( raise NotImplementedError( 'DirectRunner: id_label is not supported for PubSub reads') + with self._subscriber_client_cache_lock: + if self._applied_ptransform not in self._start_times: + self._start_times[self._applied_ptransform] = time.time() + sub_project = None if hasattr(self._evaluation_context, 'pipeline_options'): from apache_beam.options.pipeline_options import GoogleCloudOptions @@ -710,46 +715,48 @@ def process_element(self, element): def _read_from_pubsub( self, timestamp_attribute) -> list[tuple[Timestamp, 'PubsubMessage']]: - from apache_beam.io.gcp.pubsub import PubsubMessage - - def _get_element(message): - parsed_message = PubsubMessage._from_message(message) - if (timestamp_attribute and - timestamp_attribute in parsed_message.attributes): - rfc3339_or_milli = parsed_message.attributes[timestamp_attribute] - try: - timestamp = Timestamp(micros=int(rfc3339_or_milli) * 1000) - except ValueError: - try: - timestamp = Timestamp.from_rfc3339(rfc3339_or_milli) - except ValueError as e: - raise ValueError('Bad timestamp value: %s' % e) - if timestamp.precision() > Timestamp.MICROS_PRECISION: - # Element timestamps are limited to microsecond resolution, so - # ignore sub-microsecond digits, as the Dataflow service does. - timestamp = timestamp.to_precision( - Timestamp.MICROS_PRECISION, allow_lossy_conversion=True) - else: - if message.publish_time is None: - raise ValueError('No publish time present in message: %s' % message) - try: - timestamp = Timestamp.from_utc_datetime(message.publish_time) - except ValueError as e: - raise ValueError('Bad timestamp value for message %s: %s', message, e) - - return timestamp, parsed_message + from apache_beam.io.gcp.pubsub import _parse_pubsub_message + + # Check if timeout has elapsed. + timeout = 30 + max_read_time_seconds = getattr(self.source, 'max_read_time_seconds', None) + if max_read_time_seconds: + with self._subscriber_client_cache_lock: + start_time = self._start_times.get( + self._applied_ptransform, time.time()) + elapsed = time.time() - start_time + timeout = max(0, min(timeout, max_read_time_seconds - elapsed)) + if timeout <= 0: + return [] # Because of the AutoAck, we are not able to reread messages if this # evaluator fails with an exception before emitting a bundle. However, # the DirectRunner currently doesn't retry work items anyway, so the # pipeline would enter an inconsistent state on any error. sub_client = self._get_subscriber_client(self._applied_ptransform) - response = sub_client.pull( - subscription=self._sub_name, max_messages=10, timeout=30) - results = [_get_element(rm.message) for rm in response.received_messages] - ack_ids = [rm.ack_id for rm in response.received_messages] - if ack_ids: - sub_client.acknowledge(subscription=self._sub_name, ack_ids=ack_ids) + try: + response = sub_client.pull( + subscription=self._sub_name, max_messages=10, timeout=timeout) + results = [ + _parse_pubsub_message(rm.message, timestamp_attribute) + for rm in response.received_messages + ] + ack_ids = [rm.ack_id for rm in response.received_messages] + if ack_ids: + sub_client.acknowledge(subscription=self._sub_name, ack_ids=ack_ids) + except Exception as e: + try: + from google.api_core import exceptions + if isinstance(e, exceptions.DeadlineExceeded) or type( + e).__name__ == 'DeadlineExceeded': + results = [] + else: + raise + except ImportError: + if type(e).__name__ == 'DeadlineExceeded': + results = [] + else: + raise return results @@ -769,6 +776,17 @@ def finish_bundle(self) -> TransformResult: bundles = [bundle] else: bundles = [] + + max_read_time_seconds = getattr(self.source, 'max_read_time_seconds', None) + if max_read_time_seconds: + with self._subscriber_client_cache_lock: + start_time = self._start_times.get( + self._applied_ptransform, time.time()) + elapsed = time.time() - start_time + if elapsed >= max_read_time_seconds: + return TransformResult( + self, bundles, [], None, {None: WatermarkManager.WATERMARK_POS_INF}) + assert self._applied_ptransform.transform is not None if self._applied_ptransform.inputs: input_pvalue = self._applied_ptransform.inputs[0] diff --git a/sdks/python/apache_beam/yaml/extended_tests/messaging/pubsub.yaml b/sdks/python/apache_beam/yaml/extended_tests/messaging/pubsub.yaml index 41f739ac77e0..a7a27df31637 100644 --- a/sdks/python/apache_beam/yaml/extended_tests/messaging/pubsub.yaml +++ b/sdks/python/apache_beam/yaml/extended_tests/messaging/pubsub.yaml @@ -21,41 +21,50 @@ fixtures: config: project_id: "apache-beam-testing" +# One pipeline is required due to the lifespan of the pubsub messages in this +# test environment. pipelines: - # Pubsub write pipeline - pipeline: - type: chain + type: composite transforms: - type: Create + name: elements config: elements: - {value: "11a"} - {value: "37a"} - {value: "389a"} - type: WriteToPubSub + input: elements config: topic: "{PS_TOPIC}" format: "RAW" - + - type: ReadFromPubSub + name: read + config: + topic: "{PS_TOPIC}" + format: "RAW" + publish_time_field: "publish_time" + max_read_time_seconds: 30 # Need a minimum of this for test_pubsub_emulator to work correctly + - type: MapToFields + name: convertBytesToString + input: read + config: + language: python + fields: + value: + callable: "lambda row: row.payload.decode('utf-8')" + output_type: string + has_publish_time: + callable: "lambda row: row.publish_time is not None" + output_type: boolean + - type: AssertEqual + input: convertBytesToString + config: + elements: + - {value: "11a", has_publish_time: true} + - {value: "37a", has_publish_time: true} + - {value: "389a", has_publish_time: true} options: streaming: true - -# TODO: Current PubSubIO doesn't have a max_read_time_seconds parameter like -# Kafka does. Without it, the ReadFromPubSub will run forever. This is not a -# trival change. For now, we will live with the mocked tests located -# [here](https://github.com/apache/beam/blob/bea04446b41c86856c24d0a9761622092ed9936f/sdks/python/apache_beam/yaml/yaml_io_test.py#L83). - - # - pipeline: - # type: chain - # transforms: - # - type: ReadFromPubSub - # config: - # topic: "{PS_TOPIC}" - # format: "RAW" - # # ... - - - # options: - # streaming: true - diff --git a/sdks/python/apache_beam/yaml/yaml_io.py b/sdks/python/apache_beam/yaml/yaml_io.py index 3fa15044b60d..93a3e99b1403 100644 --- a/sdks/python/apache_beam/yaml/yaml_io.py +++ b/sdks/python/apache_beam/yaml/yaml_io.py @@ -331,7 +331,8 @@ def read_from_pubsub( attributes_map: Optional[str] = None, id_attribute: Optional[str] = None, timestamp_attribute: Optional[str] = None, - publish_time_field: Optional[str] = None): + publish_time_field: Optional[str] = None, + max_read_time_seconds: Optional[int] = None): """Reads messages from Cloud Pub/Sub. Args: @@ -383,11 +384,15 @@ def read_from_pubsub( units smaller than milliseconds) may be ignored. publish_time_field: Field to add to output messages with the Pub/Sub message publish time. If None, no such field is added. + max_read_time_seconds: Maximum time in seconds to read from Pub/Sub. Default + is forever. """ if topic and subscription: raise TypeError('Only one of topic and subscription may be specified.') elif not topic and not subscription: raise TypeError('One of topic or subscription may be specified.') + if max_read_time_seconds is not None and max_read_time_seconds <= 0: + raise ValueError('max_read_time_seconds must be positive.') if publish_time_field is not None and not publish_time_field.strip(): raise ValueError('publish_time_field must be a non-empty field name.') has_publish_time_field = publish_time_field is not None @@ -430,7 +435,8 @@ def mapper(msg): with_attributes=bool( attributes or attributes_map or has_publish_time_field), id_label=id_attribute, - timestamp_attribute=timestamp_attribute) + timestamp_attribute=timestamp_attribute, + max_read_time_seconds=max_read_time_seconds) | 'ParseMessage' >> beam.Map(mapper)) output.element_type = schemas.named_tuple_from_schema( schema_pb2.Schema(fields=list(payload_schema.fields) + extra_fields)) diff --git a/sdks/python/apache_beam/yaml/yaml_io_test.py b/sdks/python/apache_beam/yaml/yaml_io_test.py index 7276aa5f0728..331646821001 100644 --- a/sdks/python/apache_beam/yaml/yaml_io_test.py +++ b/sdks/python/apache_beam/yaml/yaml_io_test.py @@ -49,12 +49,14 @@ def __init__( messages, subscription=None, id_attribute=None, - timestamp_attribute=None): + timestamp_attribute=None, + max_read_time_seconds=None): self._topic = topic self._subscription = subscription self._messages = messages self._id_attribute = id_attribute self._timestamp_attribute = timestamp_attribute + self._max_read_time_seconds = max_read_time_seconds def __call__( self, @@ -63,11 +65,13 @@ def __call__( subscription, with_attributes, id_label, - timestamp_attribute): + timestamp_attribute, + max_read_time_seconds=None): assert topic == self._topic assert id_label == self._id_attribute assert timestamp_attribute == self._timestamp_attribute assert subscription == self._subscription + assert max_read_time_seconds == self._max_read_time_seconds if with_attributes: data = self._messages else: @@ -305,6 +309,44 @@ def test_read_with_id_attribute(self): result, equal_to([beam.Row(payload=b'msg1'), beam.Row(payload=b'msg2')])) + def test_read_with_max_read_time_seconds(self): + with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( + pickle_library='cloudpickle')) as p: + with mock.patch('apache_beam.io.ReadFromPubSub', + FakeReadFromPubSub( + topic='my_topic', + messages=[PubsubMessage(b'msg1', {'attr': 'value1'}), + PubsubMessage(b'msg2', {'attr': 'value2'})], + max_read_time_seconds=60)): + result = p | YamlTransform( + ''' + type: ReadFromPubSub + config: + topic: my_topic + format: RAW + max_read_time_seconds: 60 + ''') + assert_that( + result, + equal_to([beam.Row(payload=b'msg1'), beam.Row(payload=b'msg2')])) + + def test_read_with_invalid_max_read_time_seconds(self): + for invalid_val in (0, -10): + with self.subTest(invalid_val=invalid_val): + with self.assertRaisesRegex(ValueError, + 'max_read_time_seconds must be positive'): + with beam.Pipeline( + options=beam.options.pipeline_options.PipelineOptions( + pickle_library='cloudpickle')) as p: + _ = p | YamlTransform( + f''' + type: ReadFromPubSub + config: + topic: my_topic + format: RAW + max_read_time_seconds: {invalid_val} + ''') + _avro_schema = { 'type': 'record', 'name': 'ec',