[
https://issues.apache.org/jira/browse/BEAM-4275?focusedWorklogId=134639&page=com.atlassian.jira.plugin.system.issuetabpanels:worklog-tabpanel#worklog-134639
]
ASF GitHub Bot logged work on BEAM-4275:
----------------------------------------
Author: ASF GitHub Bot
Created on: 14/Aug/18 17:57
Start Date: 14/Aug/18 17:57
Worklog Time Spent: 10m
Work Description: charlesccychen closed pull request #6217: [BEAM-4275]
Implement integration tests for PubSub on DirectRunner
URL: https://github.com/apache/beam/pull/6217
This is a PR merged from a forked repository.
As GitHub hides the original diff on merge, it is displayed below for
the sake of provenance:
As this is a foreign pull request (from a fork), the diff is supplied
below (as it won't show otherwise due to GitHub magic):
diff --git a/sdks/python/apache_beam/io/gcp/pubsub_integration_test.py
b/sdks/python/apache_beam/io/gcp/pubsub_integration_test.py
index c3921bb2e55..9bb81fc645f 100644
--- a/sdks/python/apache_beam/io/gcp/pubsub_integration_test.py
+++ b/sdks/python/apache_beam/io/gcp/pubsub_integration_test.py
@@ -48,28 +48,53 @@ class PubSubIntegrationTest(unittest.TestCase):
ID_LABEL = 'id'
TIMESTAMP_ATTRIBUTE = 'timestamp'
- INPUT_MESSAGES = [
- # Use ID_LABEL attribute to deduplicate messages with the same ID.
- PubsubMessage('data001', {ID_LABEL: 'foo'}),
- PubsubMessage('data001', {ID_LABEL: 'foo'}),
- PubsubMessage('data001', {ID_LABEL: 'foo'}),
- # For those elements that have the TIMESTAMP_ATTRIBUTE attribute, the IT
- # pipeline writes back the timestamp of each element (as reported by
- # Beam), as a TIMESTAMP_ATTRIBUTE + '_out' attribute.
- PubsubMessage('data002', {
- TIMESTAMP_ATTRIBUTE: '2018-07-11T02:02:50.149000Z',
- }),
- ]
- EXPECTED_OUTPUT_MESSAGES = [
- PubsubMessage('data001-seen', {'processed': 'IT'}),
- PubsubMessage('data002-seen', {
- TIMESTAMP_ATTRIBUTE + '_out': '2018-07-11T02:02:50.149000Z',
- 'processed': 'IT',
- }),
- ]
+ INPUT_MESSAGES = {
+ # TODO(BEAM-4275): DirectRunner doesn't support reading or writing
+ # label_ids, nor writing timestamp attributes. Once these features exist,
+ # TestDirectRunner and TestDataflowRunner should behave identically.
+ 'TestDirectRunner': [
+ PubsubMessage('data001', {}),
+ # For those elements that have the TIMESTAMP_ATTRIBUTE attribute, the
+ # IT pipeline writes back the timestamp of each element (as reported
+ # by Beam), as a TIMESTAMP_ATTRIBUTE + '_out' attribute.
+ PubsubMessage('data002', {
+ TIMESTAMP_ATTRIBUTE: '2018-07-11T02:02:50.149000Z',
+ }),
+ ],
+ 'TestDataflowRunner': [
+ # Use ID_LABEL attribute to deduplicate messages with the same ID.
+ PubsubMessage('data001', {ID_LABEL: 'foo'}),
+ PubsubMessage('data001', {ID_LABEL: 'foo'}),
+ PubsubMessage('data001', {ID_LABEL: 'foo'}),
+ # For those elements that have the TIMESTAMP_ATTRIBUTE attribute, the
+ # IT pipeline writes back the timestamp of each element (as reported
+ # by Beam), as a TIMESTAMP_ATTRIBUTE + '_out' attribute.
+ PubsubMessage('data002', {
+ TIMESTAMP_ATTRIBUTE: '2018-07-11T02:02:50.149000Z',
+ })
+ ],
+ }
+ EXPECTED_OUTPUT_MESSAGES = {
+ 'TestDirectRunner': [
+ PubsubMessage('data001-seen', {'processed': 'IT'}),
+ PubsubMessage('data002-seen', {
+ TIMESTAMP_ATTRIBUTE: '2018-07-11T02:02:50.149000Z',
+ TIMESTAMP_ATTRIBUTE + '_out': '2018-07-11T02:02:50.149000Z',
+ 'processed': 'IT',
+ }),
+ ],
+ 'TestDataflowRunner': [
+ PubsubMessage('data001-seen', {'processed': 'IT'}),
+ PubsubMessage('data002-seen', {
+ TIMESTAMP_ATTRIBUTE + '_out': '2018-07-11T02:02:50.149000Z',
+ 'processed': 'IT',
+ }),
+ ],
+ }
def setUp(self):
self.test_pipeline = TestPipeline(is_integration_test=True)
+ self.runner_name = type(self.test_pipeline.runner).__name__
self.project = self.test_pipeline.get_option('project')
self.uuid = str(uuid.uuid4())
@@ -99,19 +124,26 @@ def _test_streaming(self, with_attributes):
True - Reads and writes message data and attributes. Also verifies
id_label and timestamp_attribute features.
"""
- # Build expected dataset.
- # Set extra options to the pipeline for test purpose
+ # Set on_success_matcher to verify pipeline state and pubsub output. These
+ # verifications run on a (remote) worker.
+
+ # Expect the state to be RUNNING since a streaming pipeline is usually
+ # never DONE. The test runner will cancel the pipeline after verification.
state_verifier = PipelineStateMatcher(PipelineState.RUNNING)
- expected_messages = self.EXPECTED_OUTPUT_MESSAGES
+ expected_messages = self.EXPECTED_OUTPUT_MESSAGES[self.runner_name]
if not with_attributes:
expected_messages = [pubsub_msg.data for pubsub_msg in expected_messages]
+ if self.runner_name == 'TestDirectRunner':
+ strip_attributes = None
+ else:
+ strip_attributes = [self.ID_LABEL, self.TIMESTAMP_ATTRIBUTE]
pubsub_msg_verifier = PubSubMessageMatcher(
self.project,
OUTPUT_SUB + self.uuid,
expected_messages,
timeout=MESSAGE_MATCHER_TIMEOUT_S,
with_attributes=with_attributes,
- strip_attributes=[self.ID_LABEL, self.TIMESTAMP_ATTRIBUTE])
+ strip_attributes=strip_attributes)
extra_opts = {'input_subscription': self.input_sub.full_name,
'output_topic': self.output_topic.full_name,
'wait_until_finish_duration': TEST_PIPELINE_DURATION_MS,
@@ -120,7 +152,7 @@ def _test_streaming(self, with_attributes):
# Generate input data and inject to PubSub.
test_utils.wait_for_subscriptions_created([self.input_sub])
- for msg in self.INPUT_MESSAGES:
+ for msg in self.INPUT_MESSAGES[self.runner_name]:
self.input_topic.publish(msg.data, **msg.attributes)
# Get pipeline options from command argument: --test-pipeline-options,
diff --git a/sdks/python/apache_beam/io/gcp/pubsub_it_pipeline.py
b/sdks/python/apache_beam/io/gcp/pubsub_it_pipeline.py
index 0c4535f1a29..f82714867c1 100644
--- a/sdks/python/apache_beam/io/gcp/pubsub_it_pipeline.py
+++ b/sdks/python/apache_beam/io/gcp/pubsub_it_pipeline.py
@@ -48,13 +48,20 @@ def run_pipeline(argv, with_attributes, id_label,
timestamp_attribute):
pipeline_options.view_as(SetupOptions).save_main_session = True
pipeline_options.view_as(StandardOptions).streaming = True
p = beam.Pipeline(options=pipeline_options)
+ runner_name = type(p.runner).__name__
# Read from PubSub into a PCollection.
- messages = p | beam.io.ReadFromPubSub(
- subscription=known_args.input_subscription,
- id_label=id_label,
- with_attributes=with_attributes,
- timestamp_attribute=timestamp_attribute)
+ if runner_name == 'TestDirectRunner':
+ messages = p | beam.io.ReadFromPubSub(
+ subscription=known_args.input_subscription,
+ with_attributes=with_attributes,
+ timestamp_attribute=timestamp_attribute)
+ else:
+ messages = p | beam.io.ReadFromPubSub(
+ subscription=known_args.input_subscription,
+ id_label=id_label,
+ with_attributes=with_attributes,
+ timestamp_attribute=timestamp_attribute)
def add_attribute(msg, timestamp=beam.DoFn.TimestampParam):
msg.data += '-seen'
@@ -72,10 +79,14 @@ def modify_data(data):
output = messages | 'modify_data' >> beam.Map(modify_data)
# Write to PubSub.
- _ = output | beam.io.WriteToPubSub(known_args.output_topic,
- id_label=id_label,
- with_attributes=with_attributes,
- timestamp_attribute=timestamp_attribute)
+ if runner_name == 'TestDirectRunner':
+ _ = output | beam.io.WriteToPubSub(known_args.output_topic,
+ with_attributes=with_attributes)
+ else:
+ _ = output | beam.io.WriteToPubSub(known_args.output_topic,
+ id_label=id_label,
+ with_attributes=with_attributes,
+ timestamp_attribute=timestamp_attribute)
result = p.run()
result.wait_until_finish()
diff --git a/sdks/python/apache_beam/io/gcp/pubsub_test.py
b/sdks/python/apache_beam/io/gcp/pubsub_test.py
index 44024cd4535..6e19950777f 100644
--- a/sdks/python/apache_beam/io/gcp/pubsub_test.py
+++ b/sdks/python/apache_beam/io/gcp/pubsub_test.py
@@ -450,7 +450,7 @@ def test_read_messages_success(self, mock_pubsub):
p.options.view_as(StandardOptions).streaming = True
pcoll = (p
| ReadFromPubSub('projects/fakeprj/topics/a_topic',
- None, 'a_label', with_attributes=True))
+ None, None, with_attributes=True))
assert_that(pcoll, equal_to(expected_elements), reify_windows=True)
p.run()
@@ -469,7 +469,7 @@ def test_read_strings_success(self, mock_pubsub):
p.options.view_as(StandardOptions).streaming = True
pcoll = (p
| ReadStringsFromPubSub('projects/fakeprj/topics/a_topic',
- None, 'a_label'))
+ None, None))
assert_that(pcoll, equal_to(expected_elements))
p.run()
@@ -486,8 +486,7 @@ def test_read_data_success(self, mock_pubsub):
p = TestPipeline()
p.options.view_as(StandardOptions).streaming = True
pcoll = (p
- | ReadFromPubSub('projects/fakeprj/topics/a_topic',
- None, 'a_label'))
+ | ReadFromPubSub('projects/fakeprj/topics/a_topic', None, None))
assert_that(pcoll, equal_to(expected_elements))
p.run()
@@ -513,7 +512,7 @@ def
test_read_messages_timestamp_attribute_milli_success(self, mock_pubsub):
p.options.view_as(StandardOptions).streaming = True
pcoll = (p
| ReadFromPubSub(
- 'projects/fakeprj/topics/a_topic', None, 'a_label',
+ 'projects/fakeprj/topics/a_topic', None, None,
with_attributes=True, timestamp_attribute='time'))
assert_that(pcoll, equal_to(expected_elements), reify_windows=True)
p.run()
@@ -540,31 +539,37 @@ def
test_read_messages_timestamp_attribute_rfc3339_success(self, mock_pubsub):
p.options.view_as(StandardOptions).streaming = True
pcoll = (p
| ReadFromPubSub(
- 'projects/fakeprj/topics/a_topic', None, 'a_label',
+ 'projects/fakeprj/topics/a_topic', None, None,
with_attributes=True, timestamp_attribute='time'))
assert_that(pcoll, equal_to(expected_elements), reify_windows=True)
p.run()
@mock.patch('google.cloud.pubsub')
- def test_read_messages_timestamp_attribute_fail_missing(self, mock_pubsub):
+ def test_read_messages_timestamp_attribute_missing(self, mock_pubsub):
data = 'data'
message_id = 'message_id'
- attributes = {'time': '1337'}
+ attributes = {}
publish_time = '2018-03-12T13:37:01.234567Z'
payloads = [
create_client_message(data, message_id, attributes, publish_time)]
+ expected_elements = [
+ TestWindowedValue(
+ PubsubMessage(data, attributes),
+ timestamp.Timestamp.from_rfc3339(publish_time),
+ [window.GlobalWindow()]),
+ ]
mock_pubsub.Client = functools.partial(FakePubsubClient, payloads)
mock_pubsub.subscription.AutoAck = FakeAutoAck
p = TestPipeline()
p.options.view_as(StandardOptions).streaming = True
- _ = (p
- | ReadFromPubSub(
- 'projects/fakeprj/topics/a_topic', None, 'a_label',
- with_attributes=True, timestamp_attribute='nonexistent'))
- with self.assertRaisesRegexp(KeyError, r'Timestamp.*nonexistent'):
- p.run()
+ pcoll = (p
+ | ReadFromPubSub(
+ 'projects/fakeprj/topics/a_topic', None, None,
+ with_attributes=True, timestamp_attribute='nonexistent'))
+ assert_that(pcoll, equal_to(expected_elements), reify_windows=True)
+ p.run()
@mock.patch('google.cloud.pubsub')
def test_read_messages_timestamp_attribute_fail_parse(self, mock_pubsub):
@@ -582,11 +587,31 @@ def
test_read_messages_timestamp_attribute_fail_parse(self, mock_pubsub):
p.options.view_as(StandardOptions).streaming = True
_ = (p
| ReadFromPubSub(
- 'projects/fakeprj/topics/a_topic', None, 'a_label',
+ 'projects/fakeprj/topics/a_topic', None, None,
with_attributes=True, timestamp_attribute='time'))
with self.assertRaisesRegexp(ValueError, r'parse'):
p.run()
+ @mock.patch('google.cloud.pubsub')
+ def test_read_message_id_label_unsupported(self, mock_pubsub):
+ # id_label is unsupported in DirectRunner.
+ data = 'data'
+ message_id = 'message_id'
+ attributes = {'time': '1337 unparseable'}
+ publish_time = '2018-03-12T13:37:01.234567Z'
+ payloads = [
+ create_client_message(data, message_id, attributes, publish_time)]
+
+ mock_pubsub.Client = functools.partial(FakePubsubClient, payloads)
+ mock_pubsub.subscription.AutoAck = FakeAutoAck
+
+ p = TestPipeline()
+ p.options.view_as(StandardOptions).streaming = True
+ _ = (p | ReadFromPubSub('projects/fakeprj/topics/a_topic', None,
'a_label'))
+ with self.assertRaisesRegexp(NotImplementedError,
+ r'id_label is not supported'):
+ p.run()
+
@unittest.skipIf(pubsub is None, 'GCP dependencies are not installed')
class TestWriteToPubSub(unittest.TestCase):
@@ -660,6 +685,35 @@ def test_write_messages_with_attributes_error(self,
mock_pubsub):
r'str.*has no attribute.*data'):
p.run()
+ @mock.patch('google.cloud.pubsub')
+ def test_write_messages_unsupported_features(self, mock_pubsub):
+ data = 'data'
+ attributes = {'key': 'value'}
+ payloads = [PubsubMessage(data, attributes)]
+ expected_payloads = [[data, attributes]]
+
+ mock_pubsub.Client = functools.partial(FakePubsubClient,
+ messages_write=expected_payloads)
+
+ p = TestPipeline()
+ p.options.view_as(StandardOptions).streaming = True
+ _ = (p
+ | Create(payloads)
+ | WriteToPubSub('projects/fakeprj/topics/a_topic',
+ id_label='a_label'))
+ with self.assertRaisesRegexp(NotImplementedError,
+ r'id_label is not supported'):
+ p.run()
+ p = TestPipeline()
+ p.options.view_as(StandardOptions).streaming = True
+ _ = (p
+ | Create(payloads)
+ | WriteToPubSub('projects/fakeprj/topics/a_topic',
+ timestamp_attribute='timestamp'))
+ with self.assertRaisesRegexp(NotImplementedError,
+ r'timestamp_attribute is not supported'):
+ p.run()
+
if __name__ == '__main__':
logging.getLogger().setLevel(logging.INFO)
diff --git a/sdks/python/apache_beam/runners/__init__.py
b/sdks/python/apache_beam/runners/__init__.py
index ad5c3f626fa..0f278d182d7 100644
--- a/sdks/python/apache_beam/runners/__init__.py
+++ b/sdks/python/apache_beam/runners/__init__.py
@@ -23,6 +23,7 @@
from __future__ import absolute_import
from apache_beam.runners.direct.direct_runner import DirectRunner
+from apache_beam.runners.direct.test_direct_runner import TestDirectRunner
from apache_beam.runners.runner import PipelineRunner
from apache_beam.runners.runner import PipelineState
from apache_beam.runners.runner import create_runner
diff --git a/sdks/python/apache_beam/runners/dataflow/dataflow_runner.py
b/sdks/python/apache_beam/runners/dataflow/dataflow_runner.py
index c36ae8ccd1e..30f33cb9381 100644
--- a/sdks/python/apache_beam/runners/dataflow/dataflow_runner.py
+++ b/sdks/python/apache_beam/runners/dataflow/dataflow_runner.py
@@ -1072,18 +1072,7 @@ def metrics(self):
def has_job(self):
return self._job is not None
- @property
- def state(self):
- """Return the current state of the remote job.
-
- Returns:
- A PipelineState object.
- """
- if not self.has_job:
- return PipelineState.UNKNOWN
-
- self._update_job()
-
+ def _get_job_state(self):
values_enum = dataflow_api.Job.CurrentStateValueValuesEnum
# TODO: Move this table to a another location.
@@ -1105,15 +1094,25 @@ def state(self):
return (api_jobstate_map[self._job.currentState] if self._job.currentState
else PipelineState.UNKNOWN)
+ @property
+ def state(self):
+ """Return the current state of the remote job.
+
+ Returns:
+ A PipelineState object.
+ """
+ if not self.has_job:
+ return PipelineState.UNKNOWN
+
+ self._update_job()
+
+ return self._get_job_state()
+
def is_in_terminal_state(self):
if not self.has_job:
return True
- values_enum = dataflow_api.Job.CurrentStateValueValuesEnum
- return self._job.currentState in [
- values_enum.JOB_STATE_STOPPED, values_enum.JOB_STATE_DONE,
- values_enum.JOB_STATE_FAILED, values_enum.JOB_STATE_CANCELLED,
- values_enum.JOB_STATE_UPDATED, values_enum.JOB_STATE_DRAINED]
+ return PipelineState.is_terminal(self._get_job_state())
def wait_until_finish(self, duration=None):
if not self.is_in_terminal_state():
diff --git a/sdks/python/apache_beam/runners/direct/direct_runner.py
b/sdks/python/apache_beam/runners/direct/direct_runner.py
index 0ca2f0d2fcc..4131c9b0c65 100644
--- a/sdks/python/apache_beam/runners/direct/direct_runner.py
+++ b/sdks/python/apache_beam/runners/direct/direct_runner.py
@@ -272,10 +272,11 @@ def __init__(self, sink):
# TODO(BEAM-4275): Add support for id_label and timestamp_attribute.
if sink.id_label:
- raise NotImplementedError('id_label is not supported in Direct Runner')
+ raise NotImplementedError('DirectRunner: id_label is not supported for '
+ 'PubSub writes')
if sink.timestamp_attribute:
- raise NotImplementedError('timestamp_attribute is not supported in
Direct'
- ' Runner')
+ raise NotImplementedError('DirectRunner: timestamp_attribute is not '
+ 'supported for PubSub writes')
def start_bundle(self):
from google.cloud import pubsub
@@ -423,11 +424,8 @@ def __del__(self):
'result.wait_until_finish() to wait for completion of pipeline '
'execution.')
- def _is_in_terminal_state(self):
- return self._state is not PipelineState.RUNNING
-
def wait_until_finish(self, duration=None):
- if not self._is_in_terminal_state():
+ if not PipelineState.is_terminal(self.state):
if duration:
raise NotImplementedError(
'DirectRunner does not support duration argument.')
@@ -444,3 +442,13 @@ def aggregated_values(self, aggregator_or_name):
def metrics(self):
return self._evaluation_context.metrics()
+
+ def cancel(self):
+ """Shuts down pipeline workers.
+
+ For testing use only. Does not properly wait for pipeline workers to shut
+ down.
+ """
+ self._state = PipelineState.CANCELLING
+ self._executor.shutdown()
+ self._state = PipelineState.CANCELLED
diff --git a/sdks/python/apache_beam/runners/direct/direct_runner_test.py
b/sdks/python/apache_beam/runners/direct/direct_runner_test.py
index f258168dd60..fde0883acd9 100644
--- a/sdks/python/apache_beam/runners/direct/direct_runner_test.py
+++ b/sdks/python/apache_beam/runners/direct/direct_runner_test.py
@@ -21,6 +21,9 @@
import unittest
import apache_beam as beam
+from apache_beam.runners import DirectRunner
+from apache_beam.runners import TestDirectRunner
+from apache_beam.runners import create_runner
from apache_beam.testing import test_pipeline
@@ -40,6 +43,14 @@ def test_waiting_on_result_stops_executor_threads(self):
new_threads = post_test_threads - pre_test_threads
self.assertEqual(len(new_threads), 0)
+ def test_create_runner(self):
+ self.assertTrue(
+ isinstance(create_runner('DirectRunner'),
+ DirectRunner))
+ self.assertTrue(
+ isinstance(create_runner('TestDirectRunner'),
+ TestDirectRunner))
+
if __name__ == '__main__':
unittest.main()
diff --git a/sdks/python/apache_beam/runners/direct/executor.py
b/sdks/python/apache_beam/runners/direct/executor.py
index 6fe3795df4e..32a6b32ea56 100644
--- a/sdks/python/apache_beam/runners/direct/executor.py
+++ b/sdks/python/apache_beam/runners/direct/executor.py
@@ -448,6 +448,9 @@ def await_completion(self):
self.executor_service.shutdown()
self.executor_service.await_completion()
+ def request_shutdown(self):
+ self.executor_service.shutdown()
+
def schedule_consumers(self, committed_bundle):
if committed_bundle.pcollection in self.value_to_consumers:
consumers = self.value_to_consumers[committed_bundle.pcollection]
diff --git a/sdks/python/apache_beam/runners/direct/test_direct_runner.py
b/sdks/python/apache_beam/runners/direct/test_direct_runner.py
new file mode 100644
index 00000000000..8facca8edc8
--- /dev/null
+++ b/sdks/python/apache_beam/runners/direct/test_direct_runner.py
@@ -0,0 +1,56 @@
+#
+# 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.
+#
+
+"""Wrapper of Beam runners that's built for running and verifying e2e tests."""
+
+from __future__ import absolute_import
+from __future__ import print_function
+
+from apache_beam.internal import pickler
+from apache_beam.options.pipeline_options import StandardOptions
+from apache_beam.options.pipeline_options import TestOptions
+from apache_beam.runners.direct.direct_runner import DirectRunner
+from apache_beam.runners.runner import PipelineState
+
+__all__ = ['TestDirectRunner']
+
+
+class TestDirectRunner(DirectRunner):
+ def run_pipeline(self, pipeline):
+ """Execute test pipeline and verify test matcher"""
+ options = pipeline._options.view_as(TestOptions)
+ on_success_matcher = options.on_success_matcher
+ is_streaming = options.view_as(StandardOptions).streaming
+
+ # [BEAM-1889] Do not send this to remote workers also, there is no need to
+ # send this option to remote executors.
+ options.on_success_matcher = None
+
+ self.result = super(TestDirectRunner, self).run_pipeline(pipeline)
+
+ try:
+ if not is_streaming:
+ self.result.wait_until_finish()
+
+ if on_success_matcher:
+ from hamcrest import assert_that as hc_assert_that
+ hc_assert_that(self.result, pickler.loads(on_success_matcher))
+ finally:
+ if not PipelineState.is_terminal(self.result.state):
+ self.result.cancel()
+
+ return self.result
diff --git a/sdks/python/apache_beam/runners/direct/transform_evaluator.py
b/sdks/python/apache_beam/runners/direct/transform_evaluator.py
index 308c11769be..22aedce2bb0 100644
--- a/sdks/python/apache_beam/runners/direct/transform_evaluator.py
+++ b/sdks/python/apache_beam/runners/direct/transform_evaluator.py
@@ -391,6 +391,9 @@ def __init__(self, evaluation_context, applied_ptransform,
side_inputs)
self.source = self._applied_ptransform.transform._source
+ if self.source.id_label:
+ raise NotImplementedError(
+ 'DirectRunner: id_label is not supported for PubSub reads')
self._subscription = _PubSubReadEvaluator.get_subscription(
self._applied_ptransform, self.source.project, self.source.topic_name,
self.source.subscription_name)
@@ -430,11 +433,9 @@ def _read_from_pubsub(self, timestamp_attribute):
max_messages=10) as results:
def _get_element(message):
parsed_message = PubsubMessage._from_message(message)
- if timestamp_attribute:
- try:
- rfc3339_or_milli = parsed_message.attributes[timestamp_attribute]
- except KeyError as e:
- raise KeyError('Timestamp attribute not found: %s' % e)
+ if (timestamp_attribute and
+ timestamp_attribute in parsed_message.attributes):
+ rfc3339_or_milli = parsed_message.attributes[timestamp_attribute]
try:
timestamp = Timestamp.from_rfc3339(rfc3339_or_milli)
except ValueError:
diff --git a/sdks/python/apache_beam/runners/runner.py
b/sdks/python/apache_beam/runners/runner.py
index f54ffe8a321..00ce3e6429b 100644
--- a/sdks/python/apache_beam/runners/runner.py
+++ b/sdks/python/apache_beam/runners/runner.py
@@ -50,7 +50,7 @@ def _get_runner_map(runner_names, module_path):
_KNOWN_DIRECT_RUNNERS = ('DirectRunner', 'BundleBasedDirectRunner',
'SwitchingDirectRunner')
_KNOWN_DATAFLOW_RUNNERS = ('DataflowRunner',)
-_KNOWN_TEST_RUNNERS = ('TestDataflowRunner',)
+_KNOWN_TEST_RUNNERS = ('TestDataflowRunner', 'TestDirectRunner')
_KNOWN_PORTABLE_RUNNERS = ('PortableRunner',)
_RUNNER_MAP = {}
@@ -76,8 +76,8 @@ def create_runner(runner_name):
Creates a runner instance from a runner class name.
Args:
- runner_name: Name of the pipeline runner. Possible values are:
- DirectRunner, DataflowRunner and TestDataflowRunner.
+ runner_name: Name of the pipeline runner. Possible values are listed in
+ _RUNNER_MAP above.
Returns:
A runner object.
@@ -333,6 +333,11 @@ class PipelineState(object):
CANCELLING = 'CANCELLING' # job has been explicitly cancelled and is
# in the process of stopping
+ @classmethod
+ def is_terminal(cls, state):
+ return state in [cls.STOPPED, cls.DONE, cls.FAILED, cls.CANCELLED,
+ cls.UPDATED, cls.DRAINED]
+
class PipelineResult(object):
"""A :class:`PipelineResult` provides access to info about a pipeline."""
diff --git a/sdks/python/apache_beam/runners/test/__init__.py
b/sdks/python/apache_beam/runners/test/__init__.py
index a52ea6e84df..26f13539d4a 100644
--- a/sdks/python/apache_beam/runners/test/__init__.py
+++ b/sdks/python/apache_beam/runners/test/__init__.py
@@ -27,6 +27,7 @@
try:
from apache_beam.runners.dataflow.test_dataflow_runner import
TestDataflowRunner
+ from apache_beam.runners.direct.test_direct_runner import TestDirectRunner
except ImportError:
pass
# pylint: enable=wrong-import-order, wrong-import-position
diff --git a/sdks/python/build.gradle b/sdks/python/build.gradle
index 240ed3ccec7..3ee9b66c980 100644
--- a/sdks/python/build.gradle
+++ b/sdks/python/build.gradle
@@ -187,12 +187,17 @@ task installGcpTest(dependsOn: 'setupVirtualenv') {
}
}
-task localWordCount(dependsOn: 'installGcpTest') {
+task directRunnerIT(dependsOn: 'installGcpTest') {
doLast {
exec {
executable 'sh'
- args '-c', ". ${envdir}/bin/activate && python -m
apache_beam.examples.wordcount --output /tmp/py-wordcount-direct"
- // TODO: Check that the output file is generated and runs.
+ args '-c', ". ${envdir}/bin/activate && ./scripts/run_postcommit.sh IT
batch TestDirectRunner"
+ }
+ }
+ doLast {
+ exec {
+ executable 'sh'
+ args '-c', ". ${envdir}/bin/activate && ./scripts/run_postcommit.sh IT
streaming TestDirectRunner"
}
}
}
@@ -281,7 +286,7 @@ flinkCompatibilityMatrix('Streaming')
task postCommit() {
dependsOn "preCommit"
- dependsOn "localWordCount"
+ dependsOn "directRunnerIT"
dependsOn "hdfsIntegrationTest"
dependsOn "postCommitITTests"
}
diff --git a/sdks/python/container/run_validatescontainer.sh
b/sdks/python/container/run_validatescontainer.sh
index e4da78326aa..f6c5deff392 100755
--- a/sdks/python/container/run_validatescontainer.sh
+++ b/sdks/python/container/run_validatescontainer.sh
@@ -77,7 +77,7 @@ SDK_LOCATION=$(find dist/apache-beam-*.tar.gz)
echo ">>> RUNNING DATAFLOW RUNNER VALIDATESCONTAINER TEST"
python setup.py nosetests \
--attr ValidatesContainer \
- --nocapture \
+ --nologcapture \
--processes=1 \
--process-timeout=900 \
--test-pipeline-options=" \
diff --git a/sdks/python/scripts/run_postcommit.sh
b/sdks/python/scripts/run_postcommit.sh
index c13ef505fa9..a228c7f2b58 100755
--- a/sdks/python/scripts/run_postcommit.sh
+++ b/sdks/python/scripts/run_postcommit.sh
@@ -30,9 +30,10 @@
# Usage check.
if (( $# < 2 )); then
- printf "Usage: \n$> ./scripts/run_postcommit.sh <test_type> <pipeline_type>
[gcp_location] [gcp_project]"
+ printf "Usage: \n$> ./scripts/run_postcommit.sh <test_type> <pipeline_type>
<runner_type> [gcp_location] [gcp_project]"
printf "\n\ttest_type: [required] ValidatesRunner or IT"
printf "\n\tpipeline_type: [required] streaming or batch"
+ printf "\n\trunner_type: [optional] TestDataflowRunner or TestDirectRunner"
printf "\n\tgcp_location: [optional] A gs:// path to stage artifacts and
output results"
printf "\n\tgcp_project: [optional] A GCP project to run Dataflow
pipelines\n"
exit 1
@@ -56,13 +57,15 @@ if [[ "*sdks/python" != $PWD ]]; then
cd $(pwd | sed 's/sdks\/python.*/sdks\/python/')
fi
+RUNNER=${3:-TestDataflowRunner}
+
# Where to store integration test outputs.
-GCS_LOCATION=${3:-gs://temp-storage-for-end-to-end-tests}
+GCS_LOCATION=${4:-gs://temp-storage-for-end-to-end-tests}
-PROJECT=${4:-apache-beam-testing}
+PROJECT=${5:-apache-beam-testing}
# Create a tarball
-python setup.py sdist
+python setup.py -q sdist
SDK_LOCATION=$(find dist/apache-beam-*.tar.gz)
@@ -70,9 +73,10 @@ SDK_LOCATION=$(find dist/apache-beam-*.tar.gz)
echo "pyhamcrest" > postcommit_requirements.txt
echo "mock" >> postcommit_requirements.txt
-# Options used to run testing pipeline on Cloud Dataflow Service.
+# Options used to run testing pipeline on Cloud Dataflow Service. Also used for
+# running on DirectRunner (some options ignored).
PIPELINE_OPTIONS=(
- "--runner=TestDataflowRunner"
+ "--runner=$RUNNER"
"--project=$PROJECT"
"--staging_location=$GCS_LOCATION/staging-it"
"--temp_location=$GCS_LOCATION/temp-it"
@@ -91,17 +95,23 @@ else
echo ">>> Set test pipeline to batch"
fi
+TESTS=""
+if [[ "$3" = "TestDirectRunner" ]]; then
+ TESTS="--tests=\
+apache_beam.examples.wordcount_it_test:WordCountIT.test_wordcount_it,\
+apache_beam.io.gcp.pubsub_integration_test:PubSubIntegrationTest"
+fi
###########################################################################
-# Run tests on the Google Cloud Dataflow service and validate that jobs
-# finish successfully.
+# Run tests and validate that jobs finish successfully.
JOINED_OPTS=$(IFS=" " ; echo "${PIPELINE_OPTIONS[*]}")
-echo ">>> RUNNING TEST DATAFLOW RUNNER $1 tests"
+echo ">>> RUNNING $RUNNER $1 tests"
python setup.py nosetests \
--attr $1 \
- --nocapture \
+ --nologcapture \
--processes=8 \
--process-timeout=3000 \
- --test-pipeline-options="$JOINED_OPTS"
+ --test-pipeline-options="$JOINED_OPTS" \
+ $TESTS
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
[email protected]
Issue Time Tracking
-------------------
Worklog Id: (was: 134639)
Time Spent: 1h (was: 50m)
> Pubsub: add DirectRunner support for id_label and timestamp_attribute in
> Python SDK
> -----------------------------------------------------------------------------------
>
> Key: BEAM-4275
> URL: https://issues.apache.org/jira/browse/BEAM-4275
> Project: Beam
> Issue Type: Bug
> Components: runner-direct, sdk-py-core
> Reporter: Udi Meiri
> Assignee: Udi Meiri
> Priority: Major
> Time Spent: 1h
> Remaining Estimate: 0h
>
> At least for publishing (and maybe pulling) messages, non-Dataflow-based
> sources and sinks for Pub/Sub use the [public
> API|https://cloud.google.com/pubsub/docs/publisher] for Pub/Sub, which
> doesn't support id_label and timestamp_attribute settings.
> Publishing:
> id_label - add an attribute to each message with a unique value
> timestamp_attribute - add an attribute to each message with the publishing
> time as its value
> Pulling:
> id_label - use the value of this message attribute to deduplicate messages
> timestamp_attribute - use the value of this message attribute as the
> element's timestamp
>
> Implementation details: could probably create a pubsubio.py module, for reuse
> with other runners (i.e. implement Pub/Sub IO as PTransforms and not
> NativeSinks and Sources).
--
This message was sent by Atlassian JIRA
(v7.6.3#76005)
