github.com/apache/beam/sdks/v2@v2.48.2/python/apache_beam/testing/test_stream_service.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  # pytype: skip-file
    19  
    20  from concurrent.futures import ThreadPoolExecutor
    21  
    22  import grpc
    23  
    24  from apache_beam.portability.api import beam_runner_api_pb2_grpc
    25  
    26  
    27  class TestStreamServiceController(
    28      beam_runner_api_pb2_grpc.TestStreamServiceServicer):
    29    """A server that streams TestStreamPayload.Events from a single EventRequest.
    30  
    31    This server is used as a way for TestStreams to receive events from file.
    32    """
    33    def __init__(self, reader, endpoint=None, exception_handler=None):
    34      self._server = grpc.server(ThreadPoolExecutor(max_workers=10))
    35      self._server_started = False
    36      self._server_stopped = False
    37  
    38      if endpoint:
    39        self.endpoint = endpoint
    40        self._server.add_insecure_port(self.endpoint)
    41      else:
    42        port = self._server.add_insecure_port('localhost:0')
    43        self.endpoint = 'localhost:{}'.format(port)
    44  
    45      beam_runner_api_pb2_grpc.add_TestStreamServiceServicer_to_server(
    46          self, self._server)
    47      self._reader = reader
    48      self._exception_handler = exception_handler
    49      if not self._exception_handler:
    50        self._exception_handler = lambda _: False
    51  
    52    def start(self):
    53      # A server can only be started if never started and never stopped before.
    54      if self._server_started or self._server_stopped:
    55        return
    56      self._server_started = True
    57      self._server.start()
    58  
    59    def stop(self):
    60      # A server can only be stopped if already started and never stopped before.
    61      if not self._server_started or self._server_stopped:
    62        return
    63      self._server_started = False
    64      self._server_stopped = True
    65      self._server.stop(0)
    66      # This was introduced in grpcio 1.24 and might be gone in the future. Keep
    67      # this check in case the runtime is on a older, current or future grpcio.
    68      if hasattr(self._server, 'wait_for_termination'):
    69        self._server.wait_for_termination()
    70  
    71    def Events(self, request, context):
    72      """Streams back all of the events from the streaming cache."""
    73  
    74      # TODO(srohde): Once we get rid of the CacheManager, get rid of this 'full'
    75      # label.
    76      tags = [None if tag == 'None' else tag for tag in request.output_ids]
    77      try:
    78        reader = self._reader.read_multiple([('full', tag) for tag in tags])
    79        while True:
    80          e = next(reader)
    81          yield e
    82      except StopIteration:
    83        pass
    84      except Exception as e:
    85        if not self._exception_handler(e):
    86          raise e