#
# 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.
#
# pytype: skip-file

from __future__ import absolute_import

import argparse
import logging
import signal
import sys
import typing

import grpc
from past.builtins import unicode

import apache_beam as beam
import apache_beam.transforms.combiners as combine
from apache_beam.coders import StrUtf8Coder
from apache_beam.pipeline import PipelineOptions
from apache_beam.portability.api import beam_expansion_api_pb2_grpc
from apache_beam.portability.api.external_transforms_pb2 import ExternalConfigurationPayload
from apache_beam.runners.portability import expansion_service
from apache_beam.transforms import ptransform
from apache_beam.transforms.external import ImplicitSchemaPayloadBuilder
from apache_beam.utils.thread_pool_executor import UnboundedThreadPoolExecutor

# This script provides an expansion service and example ptransforms for running
# external transform test cases. See external_test.py for details.

_LOGGER = logging.getLogger(__name__)

TEST_PREFIX_URN = "beam:transforms:xlang:test:prefix"
TEST_MULTI_URN = "beam:transforms:xlang:test:multi"
TEST_GBK_URN = "beam:transforms:xlang:test:gbk"
TEST_CGBK_URN = "beam:transforms:xlang:test:cgbk"
TEST_COMGL_URN = "beam:transforms:xlang:test:comgl"
TEST_COMPK_URN = "beam:transforms:xlang:test:compk"
TEST_FLATTEN_URN = "beam:transforms:xlang:test:flatten"
TEST_PARTITION_URN = "beam:transforms:xlang:test:partition"


@ptransform.PTransform.register_urn('beam:transforms:xlang:count', None)
class CountPerElementTransform(ptransform.PTransform):
  def expand(self, pcoll):
    return pcoll | combine.Count.PerElement()

  def to_runner_api_parameter(self, unused_context):
    return 'beam:transforms:xlang:count', None

  @staticmethod
  def from_runner_api_parameter(
      unused_ptransform, unused_parameter, unused_context):
    return CountPerElementTransform()


@ptransform.PTransform.register_urn(
    'beam:transforms:xlang:filter_less_than_eq', bytes)
class FilterLessThanTransform(ptransform.PTransform):
  def __init__(self, payload):
    self._payload = payload

  def expand(self, pcoll):
    return (
        pcoll | beam.Filter(
            lambda elem, target: elem <= target, int(ord(self._payload[0]))))

  def to_runner_api_parameter(self, unused_context):
    return (
        'beam:transforms:xlang:filter_less_than', self._payload.encode('utf8'))

  @staticmethod
  def from_runner_api_parameter(unused_ptransform, payload, unused_context):
    return FilterLessThanTransform(payload.decode('utf8'))


@ptransform.PTransform.register_urn(TEST_PREFIX_URN, None)
@beam.typehints.with_output_types(unicode)
class PrefixTransform(ptransform.PTransform):
  def __init__(self, payload):
    self._payload = payload

  def expand(self, pcoll):
    return pcoll | 'TestLabel' >> beam.Map(
        lambda x: '{}{}'.format(self._payload, x))

  def to_runner_api_parameter(self, unused_context):
    return TEST_PREFIX_URN, ImplicitSchemaPayloadBuilder(
        {'data': self._payload}).payload()

  @staticmethod
  def from_runner_api_parameter(unused_ptransform, payload, unused_context):
    return PrefixTransform(parse_string_payload(payload)['data'])


@ptransform.PTransform.register_urn(TEST_MULTI_URN, None)
class MutltiTransform(ptransform.PTransform):
  def expand(self, pcolls):
    return {
        'main': (pcolls['main1'], pcolls['main2'])
        | beam.Flatten()
        | beam.Map(lambda x, s: x + s, beam.pvalue.AsSingleton(
            pcolls['side'])).with_output_types(unicode),
        'side': pcolls['side']
        | beam.Map(lambda x: x + x).with_output_types(unicode),
    }

  def to_runner_api_parameter(self, unused_context):
    return TEST_MULTI_URN, None

  @staticmethod
  def from_runner_api_parameter(
      unused_ptransform, unused_parameter, unused_context):
    return MutltiTransform()


@ptransform.PTransform.register_urn(TEST_GBK_URN, None)
class GBKTransform(ptransform.PTransform):
  def expand(self, pcoll):
    return pcoll | 'TestLabel' >> beam.GroupByKey()

  def to_runner_api_parameter(self, unused_context):
    return TEST_GBK_URN, None

  @staticmethod
  def from_runner_api_parameter(
      unused_ptransform, unused_parameter, unused_context):
    return GBKTransform()


@ptransform.PTransform.register_urn(TEST_CGBK_URN, None)
class CoGBKTransform(ptransform.PTransform):
  class ConcatFn(beam.DoFn):
    def process(self, element):
      (k, v) = element
      return [(k, v['col1'] + v['col2'])]

  def expand(self, pcoll):
    return pcoll \
           | beam.CoGroupByKey() \
           | beam.ParDo(self.ConcatFn()).with_output_types(
               typing.Tuple[int, typing.Iterable[unicode]])

  def to_runner_api_parameter(self, unused_context):
    return TEST_CGBK_URN, None

  @staticmethod
  def from_runner_api_parameter(
      unused_ptransform, unused_parameter, unused_context):
    return CoGBKTransform()


@ptransform.PTransform.register_urn(TEST_COMGL_URN, None)
class CombineGloballyTransform(ptransform.PTransform):
  def expand(self, pcoll):
    return pcoll \
           | beam.CombineGlobally(sum).with_output_types(int)

  def to_runner_api_parameter(self, unused_context):
    return TEST_COMGL_URN, None

  @staticmethod
  def from_runner_api_parameter(
      unused_ptransform, unused_parameter, unused_context):
    return CombineGloballyTransform()


@ptransform.PTransform.register_urn(TEST_COMPK_URN, None)
class CombinePerKeyTransform(ptransform.PTransform):
  def expand(self, pcoll):
    return pcoll \
           | beam.CombinePerKey(sum).with_output_types(
               typing.Tuple[unicode, int])

  def to_runner_api_parameter(self, unused_context):
    return TEST_COMPK_URN, None

  @staticmethod
  def from_runner_api_parameter(
      unused_ptransform, unused_parameter, unused_context):
    return CombinePerKeyTransform()


@ptransform.PTransform.register_urn(TEST_FLATTEN_URN, None)
class FlattenTransform(ptransform.PTransform):
  def expand(self, pcoll):
    return pcoll.values() | beam.Flatten().with_output_types(int)

  def to_runner_api_parameter(self, unused_context):
    return TEST_FLATTEN_URN, None

  @staticmethod
  def from_runner_api_parameter(
      unused_ptransform, unused_parameter, unused_context):
    return FlattenTransform()


@ptransform.PTransform.register_urn(TEST_PARTITION_URN, None)
class PartitionTransform(ptransform.PTransform):
  def expand(self, pcoll):
    col1, col2 = pcoll | beam.Partition(
        lambda elem, n: 0 if elem % 2 == 0 else 1, 2)
    typed_col1 = col1 | beam.Map(lambda x: x).with_output_types(int)
    typed_col2 = col2 | beam.Map(lambda x: x).with_output_types(int)
    return {'0': typed_col1, '1': typed_col2}

  def to_runner_api_parameter(self, unused_context):
    return TEST_PARTITION_URN, None

  @staticmethod
  def from_runner_api_parameter(
      unused_ptransform, unused_parameter, unused_context):
    return PartitionTransform()


@ptransform.PTransform.register_urn('payload', bytes)
class PayloadTransform(ptransform.PTransform):
  def __init__(self, payload):
    self._payload = payload

  def expand(self, pcoll):
    return pcoll | beam.Map(lambda x, s: x + s, self._payload)

  def to_runner_api_parameter(self, unused_context):
    return b'payload', self._payload.encode('ascii')

  @staticmethod
  def from_runner_api_parameter(unused_ptransform, payload, unused_context):
    return PayloadTransform(payload.decode('ascii'))


@ptransform.PTransform.register_urn('fib', bytes)
class FibTransform(ptransform.PTransform):
  def __init__(self, level):
    self._level = level

  def expand(self, p):
    if self._level <= 2:
      return p | beam.Create([1])
    else:
      a = p | 'A' >> beam.ExternalTransform(
          'fib',
          str(self._level - 1).encode('ascii'),
          expansion_service.ExpansionServiceServicer())
      b = p | 'B' >> beam.ExternalTransform(
          'fib',
          str(self._level - 2).encode('ascii'),
          expansion_service.ExpansionServiceServicer())
      return ((a, b)
              | beam.Flatten()
              | beam.CombineGlobally(sum).without_defaults())

  def to_runner_api_parameter(self, unused_context):
    return 'fib', str(self._level).encode('ascii')

  @staticmethod
  def from_runner_api_parameter(unused_ptransform, level, unused_context):
    return FibTransform(int(level.decode('ascii')))


def parse_string_payload(input_byte):
  payload = ExternalConfigurationPayload()
  payload.ParseFromString(input_byte)
  coder = StrUtf8Coder()
  return {
      k: coder.decode_nested(v.payload)
      for k,
      v in payload.configuration.items()
  }


server = None


def cleanup(unused_signum, unused_frame):
  _LOGGER.info('Shutting down expansion service.')
  server.stop(None)


def main(unused_argv):
  parser = argparse.ArgumentParser()
  parser.add_argument(
      '-p', '--port', type=int, help='port on which to serve the job api')
  options = parser.parse_args()
  global server
  server = grpc.server(UnboundedThreadPoolExecutor())
  beam_expansion_api_pb2_grpc.add_ExpansionServiceServicer_to_server(
      expansion_service.ExpansionServiceServicer(
          PipelineOptions(
              ["--experiments", "beam_fn_api", "--sdk_location", "container"])),
      server)
  server.add_insecure_port('localhost:{}'.format(options.port))
  server.start()
  _LOGGER.info('Listening for expansion requests at %d', options.port)

  signal.signal(signal.SIGTERM, cleanup)
  signal.signal(signal.SIGINT, cleanup)
  # blocking main thread forever.
  signal.pause()


if __name__ == '__main__':
  logging.getLogger().setLevel(logging.INFO)
  main(sys.argv)

[((44, 10, 44, 37), 'logging.getLogger', 'logging.getLogger', ({(44, 28, 44, 36): '__name__'}, {}), '(__name__)', False, 'import logging\n'), ((56, 1, 56, 72), 'apache_beam.transforms.ptransform.PTransform.register_urn', 'ptransform.PTransform.register_urn', ({(56, 36, 56, 65): '"""beam:transforms:xlang:count"""', (56, 67, 56, 71): 'None'}, {}), "('beam:transforms:xlang:count', None)", False, 'from apache_beam.transforms import ptransform\n'), ((70, 1, 71, 55), 'apache_beam.transforms.ptransform.PTransform.register_urn', 'ptransform.PTransform.register_urn', ({(71, 4, 71, 47): '"""beam:transforms:xlang:filter_less_than_eq"""', (71, 49, 71, 54): 'bytes'}, {}), "('beam:transforms:xlang:filter_less_than_eq',\n    bytes)", False, 'from apache_beam.transforms import ptransform\n'), ((90, 1, 90, 58), 'apache_beam.transforms.ptransform.PTransform.register_urn', 'ptransform.PTransform.register_urn', ({(90, 36, 90, 51): 'TEST_PREFIX_URN', (90, 53, 90, 57): 'None'}, {}), '(TEST_PREFIX_URN, None)', False, 'from apache_beam.transforms import ptransform\n'), ((91, 1, 91, 42), 'apache_beam.typehints.with_output_types', 'beam.typehints.with_output_types', ({(91, 34, 91, 41): 'unicode'}, {}), '(unicode)', True, 'import apache_beam as beam\n'), ((109, 1, 109, 57), 'apache_beam.transforms.ptransform.PTransform.register_urn', 'ptransform.PTransform.register_urn', ({(109, 36, 109, 50): 'TEST_MULTI_URN', (109, 52, 109, 56): 'None'}, {}), '(TEST_MULTI_URN, None)', False, 'from apache_beam.transforms import ptransform\n'), ((130, 1, 130, 55), 'apache_beam.transforms.ptransform.PTransform.register_urn', 'ptransform.PTransform.register_urn', ({(130, 36, 130, 48): 'TEST_GBK_URN', (130, 50, 130, 54): 'None'}, {}), '(TEST_GBK_URN, None)', False, 'from apache_beam.transforms import ptransform\n'), ((144, 1, 144, 56), 'apache_beam.transforms.ptransform.PTransform.register_urn', 'ptransform.PTransform.register_urn', ({(144, 36, 144, 49): 'TEST_CGBK_URN', (144, 51, 144, 55): 'None'}, {}), '(TEST_CGBK_URN, None)', False, 'from apache_beam.transforms import ptransform\n'), ((166, 1, 166, 57), 'apache_beam.transforms.ptransform.PTransform.register_urn', 'ptransform.PTransform.register_urn', ({(166, 36, 166, 50): 'TEST_COMGL_URN', (166, 52, 166, 56): 'None'}, {}), '(TEST_COMGL_URN, None)', False, 'from apache_beam.transforms import ptransform\n'), ((181, 1, 181, 57), 'apache_beam.transforms.ptransform.PTransform.register_urn', 'ptransform.PTransform.register_urn', ({(181, 36, 181, 50): 'TEST_COMPK_URN', (181, 52, 181, 56): 'None'}, {}), '(TEST_COMPK_URN, None)', False, 'from apache_beam.transforms import ptransform\n'), ((197, 1, 197, 59), 'apache_beam.transforms.ptransform.PTransform.register_urn', 'ptransform.PTransform.register_urn', ({(197, 36, 197, 52): 'TEST_FLATTEN_URN', (197, 54, 197, 58): 'None'}, {}), '(TEST_FLATTEN_URN, None)', False, 'from apache_beam.transforms import ptransform\n'), ((211, 1, 211, 61), 'apache_beam.transforms.ptransform.PTransform.register_urn', 'ptransform.PTransform.register_urn', ({(211, 36, 211, 54): 'TEST_PARTITION_URN', (211, 56, 211, 60): 'None'}, {}), '(TEST_PARTITION_URN, None)', False, 'from apache_beam.transforms import ptransform\n'), ((229, 1, 229, 53), 'apache_beam.transforms.ptransform.PTransform.register_urn', 'ptransform.PTransform.register_urn', ({(229, 36, 229, 45): '"""payload"""', (229, 47, 229, 52): 'bytes'}, {}), "('payload', bytes)", False, 'from apache_beam.transforms import ptransform\n'), ((245, 1, 245, 49), 'apache_beam.transforms.ptransform.PTransform.register_urn', 'ptransform.PTransform.register_urn', ({(245, 36, 245, 41): '"""fib"""', (245, 43, 245, 48): 'bytes'}, {}), "('fib', bytes)", False, 'from apache_beam.transforms import ptransform\n'), ((275, 12, 275, 42), 'apache_beam.portability.api.external_transforms_pb2.ExternalConfigurationPayload', 'ExternalConfigurationPayload', ({}, {}), '()', False, 'from apache_beam.portability.api.external_transforms_pb2 import ExternalConfigurationPayload\n'), ((277, 10, 277, 24), 'apache_beam.coders.StrUtf8Coder', 'StrUtf8Coder', ({}, {}), '()', False, 'from apache_beam.coders import StrUtf8Coder\n'), ((294, 11, 294, 36), 'argparse.ArgumentParser', 'argparse.ArgumentParser', ({}, {}), '()', False, 'import argparse\n'), ((309, 2, 309, 40), 'signal.signal', 'signal.signal', ({(309, 16, 309, 30): 'signal.SIGTERM', (309, 32, 309, 39): 'cleanup'}, {}), '(signal.SIGTERM, cleanup)', False, 'import signal\n'), ((310, 2, 310, 39), 'signal.signal', 'signal.signal', ({(310, 16, 310, 29): 'signal.SIGINT', (310, 31, 310, 38): 'cleanup'}, {}), '(signal.SIGINT, cleanup)', False, 'import signal\n'), ((312, 2, 312, 16), 'signal.pause', 'signal.pause', ({}, {}), '()', False, 'import signal\n'), ((299, 23, 299, 52), 'apache_beam.utils.thread_pool_executor.UnboundedThreadPoolExecutor', 'UnboundedThreadPoolExecutor', ({}, {}), '()', False, 'from apache_beam.utils.thread_pool_executor import UnboundedThreadPoolExecutor\n'), ((59, 19, 59, 45), 'apache_beam.transforms.combiners.Count.PerElement', 'combine.Count.PerElement', ({}, {}), '()', True, 'import apache_beam.transforms.combiners as combine\n'), ((214, 25, 215, 53), 'apache_beam.Partition', 'beam.Partition', ({(215, 8, 215, 49): '(lambda elem, n: 0 if elem % 2 == 0 else 1)', (215, 51, 215, 52): '(2)'}, {}), '(lambda elem, n: 0 if elem % 2 == 0 else 1, 2)', True, 'import apache_beam as beam\n'), ((235, 19, 235, 62), 'apache_beam.Map', 'beam.Map', ({(235, 28, 235, 46): '(lambda x, s: x + s)', (235, 48, 235, 61): 'self._payload'}, {}), '(lambda x, s: x + s, self._payload)', True, 'import apache_beam as beam\n'), ((302, 10, 303, 78), 'apache_beam.pipeline.PipelineOptions', 'PipelineOptions', ({(303, 14, 303, 77): "['--experiments', 'beam_fn_api', '--sdk_location', 'container']"}, {}), "(['--experiments', 'beam_fn_api', '--sdk_location', 'container']\n    )", False, 'from apache_beam.pipeline import PipelineOptions\n'), ((316, 2, 316, 21), 'logging.getLogger', 'logging.getLogger', ({}, {}), '()', False, 'import logging\n'), ((133, 34, 133, 51), 'apache_beam.GroupByKey', 'beam.GroupByKey', ({}, {}), '()', True, 'import apache_beam as beam\n'), ((153, 13, 153, 32), 'apache_beam.CoGroupByKey', 'beam.CoGroupByKey', ({}, {}), '()', True, 'import apache_beam as beam\n'), ((252, 17, 252, 33), 'apache_beam.Create', 'beam.Create', ({(252, 29, 252, 32): '[1]'}, {}), '([1])', True, 'import apache_beam as beam\n'), ((101, 28, 102, 32), 'apache_beam.transforms.external.ImplicitSchemaPayloadBuilder', 'ImplicitSchemaPayloadBuilder', ({(102, 8, 102, 31): "{'data': self._payload}"}, {}), "({'data': self._payload})", False, 'from apache_beam.transforms.external import ImplicitSchemaPayloadBuilder\n'), ((114, 10, 114, 24), 'apache_beam.Flatten', 'beam.Flatten', ({}, {}), '()', True, 'import apache_beam as beam\n'), ((170, 13, 170, 38), 'apache_beam.CombineGlobally', 'beam.CombineGlobally', ({(170, 34, 170, 37): 'sum'}, {}), '(sum)', True, 'import apache_beam as beam\n'), ((185, 13, 185, 36), 'apache_beam.CombinePerKey', 'beam.CombinePerKey', ({(185, 32, 185, 35): 'sum'}, {}), '(sum)', True, 'import apache_beam as beam\n'), ((200, 28, 200, 42), 'apache_beam.Flatten', 'beam.Flatten', ({}, {}), '()', True, 'import apache_beam as beam\n'), ((216, 24, 216, 45), 'apache_beam.Map', 'beam.Map', ({(216, 33, 216, 44): '(lambda x: x)'}, {}), '(lambda x: x)', True, 'import apache_beam as beam\n'), ((217, 24, 217, 45), 'apache_beam.Map', 'beam.Map', ({(217, 33, 217, 44): '(lambda x: x)'}, {}), '(lambda x: x)', True, 'import apache_beam as beam\n'), ((263, 16, 263, 30), 'apache_beam.Flatten', 'beam.Flatten', ({}, {}), '()', True, 'import apache_beam as beam\n'), ((118, 10, 118, 35), 'apache_beam.Map', 'beam.Map', ({(118, 19, 118, 34): '(lambda x: x + x)'}, {}), '(lambda x: x + x)', True, 'import apache_beam as beam\n'), ((257, 10, 257, 54), 'apache_beam.runners.portability.expansion_service.ExpansionServiceServicer', 'expansion_service.ExpansionServiceServicer', ({}, {}), '()', False, 'from apache_beam.runners.portability import expansion_service\n'), ((261, 10, 261, 54), 'apache_beam.runners.portability.expansion_service.ExpansionServiceServicer', 'expansion_service.ExpansionServiceServicer', ({}, {}), '()', False, 'from apache_beam.runners.portability import expansion_service\n'), ((264, 16, 264, 41), 'apache_beam.CombineGlobally', 'beam.CombineGlobally', ({(264, 37, 264, 40): 'sum'}, {}), '(sum)', True, 'import apache_beam as beam\n'), ((115, 39, 116, 27), 'apache_beam.pvalue.AsSingleton', 'beam.pvalue.AsSingleton', ({(116, 12, 116, 26): "pcolls['side']"}, {}), "(pcolls['side'])", True, 'import apache_beam as beam\n')]