github.com/apache/beam/sdks/v2@v2.48.2/python/apache_beam/io/gcp/pubsublite/proto_api.py (about)

     1  #
     2  # Licensed to the Apache Software Foundation (ASF) under one or more
     3  # contributor license agreements.  See the NOTICE file distributed with
     4  # this work for additional information regarding copyright ownership.
     5  # The ASF licenses this file to You under the Apache License, Version 2.0
     6  # (the "License"); you may not use this file except in compliance with
     7  # the License.  You may obtain a copy of the License at
     8  #
     9  #    http://www.apache.org/licenses/LICENSE-2.0
    10  #
    11  # Unless required by applicable law or agreed to in writing, software
    12  # distributed under the License is distributed on an "AS IS" BASIS,
    13  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    14  # See the License for the specific language governing permissions and
    15  # limitations under the License.
    16  #
    17  
    18  from apache_beam.io.gcp.pubsublite.external import _ReadExternal
    19  from apache_beam.io.gcp.pubsublite.external import _WriteExternal
    20  from apache_beam.transforms import Map
    21  from apache_beam.transforms import PTransform
    22  
    23  try:
    24    from google.cloud import pubsublite
    25  except ImportError:
    26    pubsublite = None
    27  
    28  
    29  class ReadFromPubSubLite(PTransform):
    30    """
    31    A ``PTransform`` for reading from Pub/Sub Lite.
    32  
    33    Produces a PCollection of google.cloud.pubsublite.SequencedMessage
    34  
    35    Experimental; no backwards-compatibility guarantees.
    36    """
    37    def __init__(
    38        self,
    39        subscription_path,
    40        deduplicate=None,
    41        expansion_service=None,
    42    ):
    43      """Initializes ``ReadFromPubSubLite``.
    44  
    45      Args:
    46        subscription_path: Pub/Sub Lite Subscription in the form
    47            projects/<project>/locations/<location>/subscriptions/<subscription>
    48        deduplicate: Whether to deduplicate messages based on the value of
    49            the 'x-goog-pubsublite-dataflow-uuid' attribute. Defaults to False.
    50      """
    51      super().__init__()
    52      self._source = _ReadExternal(
    53          subscription_path=subscription_path,
    54          deduplicate=deduplicate,
    55          expansion_service=expansion_service,
    56      )
    57  
    58    def expand(self, pvalue):
    59      pcoll = pvalue.pipeline | self._source
    60      pcoll.element_type = bytes
    61      pcoll = pcoll | Map(pubsublite.SequencedMessage.deserialize)
    62      pcoll.element_type = pubsublite.SequencedMessage
    63      return pcoll
    64  
    65  
    66  class WriteToPubSubLite(PTransform):
    67    """
    68    A ``PTransform`` for writing to Pub/Sub Lite.
    69  
    70    Consumes a PCollection of google.cloud.pubsublite.PubSubMessage
    71  
    72    Experimental; no backwards-compatibility guarantees.
    73    """
    74    def __init__(
    75        self,
    76        topic_path,
    77        add_uuids=None,
    78        expansion_service=None,
    79    ):
    80      """Initializes ``WriteToPubSubLite``.
    81  
    82      Args:
    83        topic_path: A Pub/Sub Lite Topic path.
    84        add_uuids: Whether to add uuids to the 'x-goog-pubsublite-dataflow-uuid'
    85            uuid attribute. Defaults to False.
    86      """
    87      super().__init__()
    88      self._source = _WriteExternal(
    89          topic_path=topic_path,
    90          add_uuids=add_uuids,
    91          expansion_service=expansion_service,
    92      )
    93  
    94    @staticmethod
    95    def _message_to_proto_str(element: pubsublite.PubSubMessage):
    96      if not isinstance(element, pubsublite.PubSubMessage):
    97        raise TypeError(
    98            'Unexpected element. Type: %s (expected: PubSubMessage), '
    99            'value: %r' % (type(element), element))
   100      return pubsublite.PubSubMessage.serialize(element)
   101  
   102    def expand(self, pcoll):
   103      pcoll = pcoll | Map(WriteToPubSubLite._message_to_proto_str)
   104      pcoll.element_type = bytes
   105      pcoll = pcoll | self._source
   106      return pcoll