Rename to publish_with_ordering_key, gate Dataflow warning, fix test assertions

Author: Nikita Grover

sdks/python/apache_beam/io/gcp/pubsub.py

@@ -390,7 +390,7 @@ def __init__(
       with_attributes: bool = False,
       id_label: Optional[str] = None,
       timestamp_attribute: Optional[str] = None,
-      enable_message_ordering: bool = False) -> None:
+      publish_with_ordering_key: bool = False) -> None:
     """Initializes ``WriteToPubSub``.
 
     Args:
@@ -406,13 +406,13 @@ def __init__(
         in a ReadFromPubSub PTransform to deduplicate messages.
       timestamp_attribute: If set, will set an attribute for each Cloud Pub/Sub
         message with the given name and the message's publish time as the value.
-      enable_message_ordering: If True, enables message ordering on the
+      publish_with_ordering_key: If True, enables message ordering on the
         PublisherClient. Messages with an ordering_key will be delivered
         in order. Requires messages to have ordering_key set.
     """
     super().__init__()
     self.with_attributes = with_attributes
-    self.enable_message_ordering = enable_message_ordering
+    self.publish_with_ordering_key = publish_with_ordering_key
     self.id_label = id_label
     self.timestamp_attribute = timestamp_attribute
     self.project, self.topic_name = parse_topic(topic)
@@ -440,7 +440,7 @@ def expand(self, pcoll):
     # since _PubSubWriteDoFn._flush() is not used by Dataflow's implementation.
     runner = self.pipeline_options.get_all_options().get(
         'runner', '') if self.pipeline_options else ''
-    if 'Dataflow' in str(runner):
+    if 'Dataflow' in str(runner) and self.publish_with_ordering_key:
       logging.warning(
           'WriteToPubSub ordering_key support is not available on Dataflow '
           'via this transform. Use the XLang WriteToPubSub path instead: '
@@ -472,9 +472,9 @@ def display_data(self):
             True, label='With Attributes').drop_if_none(),
         'timestamp_attribute': DisplayDataItem(
             self.timestamp_attribute, label='Timestamp Attribute'),
-        'enable_message_ordering': DisplayDataItem(
-            self.enable_message_ordering,
-            label='Enable Message Ordering').drop_if_none(),
+        'publish_with_ordering_key': DisplayDataItem(
+            self.publish_with_ordering_key,
+            label='Publish With Ordering Key').drop_if_none(),
     }
 
 
@@ -581,7 +581,7 @@ def __init__(self, transform):
     self.id_label = transform.id_label
     self.timestamp_attribute = transform.timestamp_attribute
     self.with_attributes = transform.with_attributes
-    self.with_ordering = transform.enable_message_ordering
+    self.with_ordering = transform.publish_with_ordering_key
 
     # TODO(https://github.com/apache/beam/issues/18939): Add support for
     # id_label and timestamp_attribute.

sdks/python/apache_beam/io/gcp/pubsub_integration_test.py

@@ -319,9 +319,7 @@ def test_batch_write_with_ordering_key(self):
     if self.runner_name == 'TestDataflowRunner':
       self.skipTest(
           'Dataflow Native PubSub Sink does not support ordering_key '
-          '(see https://github.com/apache/beam/issues/36201). '
-          'Use apache_beam.io.external.gcp.pubsub.WriteToPubSub '
-          'with publish_with_ordering_key=True instead.')
+          '(see https://github.com/apache/beam/issues/36201).')
     from google.pubsub_v1.types import Subscription
 
     from apache_beam.options.pipeline_options import PipelineOptions
@@ -356,7 +354,9 @@ def test_batch_write_with_ordering_key(self):
       with TestPipeline(options=pipeline_options) as p:
         messages = p | 'CreateMessages' >> Create(test_messages)
         _ = messages | 'WriteToPubSub' >> WriteToPubSub(
-            ordering_topic.name, with_attributes=True)
+            ordering_topic.name,
+            with_attributes=True,
+            publish_with_ordering_key=True)
 
       time.sleep(10)
 

sdks/python/apache_beam/io/gcp/pubsub_test.py

@@ -467,6 +467,7 @@ def test_display_data(self):
         DisplayDataItemMatcher('id_label', 'id'),
         DisplayDataItemMatcher('with_attributes', True),
         DisplayDataItemMatcher('timestamp_attribute', 'time'),
+        DisplayDataItemMatcher('publish_with_ordering_key', False),
     ]
 
     hc.assert_that(dd.items, hc.contains_inanyorder(*expected_items))
@@ -1114,7 +1115,7 @@ def test_write_messages_with_ordering_key(self, mock_pubsub):
           | WriteToPubSub(
               'projects/fakeprj/topics/a_topic',
               with_attributes=True,
-              enable_message_ordering=True))
+              publish_with_ordering_key=True))
 
     # Verify that publish was called with ordering_key
     mock_pubsub.return_value.publish.assert_called()
@@ -1139,7 +1140,7 @@ def test_write_messages_with_ordering_key_no_attributes(self, mock_pubsub):
           | WriteToPubSub(
               'projects/fakeprj/topics/a_topic',
               with_attributes=True,
-              enable_message_ordering=True))
+              publish_with_ordering_key=True))
 
     # Verify that publish was called with ordering_key
     mock_pubsub.return_value.publish.assert_called()
@@ -1166,6 +1167,8 @@ def test_write_messages_without_ordering_key(self, mock_pubsub):
 
     # Verify that publish was called
     mock_pubsub.return_value.publish.assert_called()
+    call_args = mock_pubsub.return_value.publish.call_args
+    self.assertNotIn('ordering_key', call_args.kwargs)
 
 
 if __name__ == '__main__':