#

# 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.

#

"""Defines Transform whose expansion is implemented elsewhere."""

# pytype: skip-file

import contextlib

import copy

import functools

import glob

import logging

import re

import subprocess

import threading

import uuid

from collections import OrderedDict

from collections import namedtuple

import grpc

import yaml

from apache_beam import pvalue

from apache_beam.coders import RowCoder

from apache_beam.options import pipeline_options

from apache_beam.portability import common_urns

from apache_beam.portability.api import beam_artifact_api_pb2_grpc

from apache_beam.portability.api import beam_expansion_api_pb2

from apache_beam.portability.api import beam_expansion_api_pb2_grpc

from apache_beam.portability.api import beam_runner_api_pb2

from apache_beam.portability.api import external_transforms_pb2

from apache_beam.portability.api import schema_pb2

from apache_beam.portability.common_urns import ManagedTransforms

from apache_beam.runners import pipeline_context

from apache_beam.runners.portability import artifact_service

from apache_beam.transforms import environments

from apache_beam.transforms import ptransform

from apache_beam.typehints import WithTypeHints

from apache_beam.typehints import native_type_compatibility

from apache_beam.typehints import row_type

from apache_beam.typehints.schemas import named_fields_to_schema

from apache_beam.typehints.schemas import named_tuple_from_schema

from apache_beam.typehints.schemas import named_tuple_to_schema

from apache_beam.typehints.schemas import typing_from_runner_api

from apache_beam.typehints.trivial_inference import instance_to_type

from apache_beam.typehints.typehints import Union

from apache_beam.typehints.typehints import UnionConstraint

from apache_beam.utils import subprocess_server

from apache_beam.utils import transform_service_launcher

DEFAULT_EXPANSION_SERVICE = 'localhost:8097'

MANAGED_SCHEMA_TRANSFORM_IDENTIFIER = "beam:transform:managed:v1"

_IO_EXPANSION_SERVICE_JAR_TARGET = "sdks:java:io:expansion-service:shadowJar"

_GCP_EXPANSION_SERVICE_JAR_TARGET = (

"sdks:java:io:google-cloud-platform:expansion-service:shadowJar")

# A mapping from supported managed transforms URNs to expansion service jars

# that include the corresponding transforms.

MANAGED_TRANSFORM_URN_TO_JAR_TARGET_MAPPING = {

ManagedTransforms.Urns.ICEBERG_READ.urn: _IO_EXPANSION_SERVICE_JAR_TARGET,

ManagedTransforms.Urns.ICEBERG_WRITE.urn: _IO_EXPANSION_SERVICE_JAR_TARGET,

ManagedTransforms.Urns.ICEBERG_CDC_READ.urn: _IO_EXPANSION_SERVICE_JAR_TARGET, # pylint: disable=line-too-long

ManagedTransforms.Urns.KAFKA_READ.urn: _IO_EXPANSION_SERVICE_JAR_TARGET,

ManagedTransforms.Urns.KAFKA_WRITE.urn: _IO_EXPANSION_SERVICE_JAR_TARGET,

ManagedTransforms.Urns.BIGQUERY_READ.urn: _GCP_EXPANSION_SERVICE_JAR_TARGET,

ManagedTransforms.Urns.BIGQUERY_WRITE.urn: _GCP_EXPANSION_SERVICE_JAR_TARGET, # pylint: disable=line-too-long

ManagedTransforms.Urns.POSTGRES_READ.urn: _GCP_EXPANSION_SERVICE_JAR_TARGET,

ManagedTransforms.Urns.POSTGRES_WRITE.urn: _GCP_EXPANSION_SERVICE_JAR_TARGET, # pylint: disable=line-too-long

ManagedTransforms.Urns.MYSQL_READ.urn: _GCP_EXPANSION_SERVICE_JAR_TARGET,

ManagedTransforms.Urns.MYSQL_WRITE.urn: _GCP_EXPANSION_SERVICE_JAR_TARGET,

ManagedTransforms.Urns.SQL_SERVER_READ.urn: _GCP_EXPANSION_SERVICE_JAR_TARGET, # pylint: disable=line-too-long

ManagedTransforms.Urns.SQL_SERVER_WRITE.urn: _GCP_EXPANSION_SERVICE_JAR_TARGET, # pylint: disable=line-too-long

}

def convert_to_typing_type(type_):

if isinstance(type_, row_type.RowTypeConstraint):

return named_tuple_from_schema(named_fields_to_schema(type_._fields))

else:

return native_type_compatibility.convert_to_python_type(type_)

def _is_optional_or_none(typehint):

return (

type(None) in typehint.union_types if isinstance(

typehint, UnionConstraint) else typehint is type(None))

def _strip_optional(typehint):

if not _is_optional_or_none(typehint):

return typehint

new_types = typehint.union_types.difference({type(None)})

if len(new_types) == 1:

return list(new_types)[0]

return Union[new_types]

def iter_urns(coder, context=None):

yield coder.to_runner_api_parameter(context)[0]

for child in coder._get_component_coders():

for urn in iter_urns(child, context):

yield urn

class PayloadBuilder(object):

"""

Abstract base class for building payloads to pass to ExternalTransform.

"""

def build(self):

"""

:return: ExternalConfigurationPayload

"""

raise NotImplementedError

def payload(self):

"""

The serialized ExternalConfigurationPayload

:return: bytes

"""

return self.build().SerializeToString()

def _get_schema_proto_and_payload(self, **kwargs):

named_fields = []

fields_to_values = OrderedDict()

for key, value in kwargs.items():

if not key:

raise ValueError('Parameter name cannot be empty')

if value is None:

raise ValueError(

'Received value None for key %s. None values are currently not '

'supported' % key)

named_fields.append(

(key, convert_to_typing_type(instance_to_type(value))))

fields_to_values[key] = value

schema_proto = named_fields_to_schema(named_fields)

row = named_tuple_from_schema(schema_proto)(**fields_to_values)

schema = named_tuple_to_schema(type(row))

payload = RowCoder(schema).encode(row)

return (schema_proto, payload)

class SchemaBasedPayloadBuilder(PayloadBuilder):

"""

Base class for building payloads based on a schema that provides

type information for each configuration value to encode.

"""

def _get_named_tuple_instance(self):

raise NotImplementedError()

def build(self):

row = self._get_named_tuple_instance()

schema = named_tuple_to_schema(type(row))

return external_transforms_pb2.ExternalConfigurationPayload(

schema=schema, payload=RowCoder(schema).encode(row))

class ImplicitSchemaPayloadBuilder(SchemaBasedPayloadBuilder):

"""

Build a payload that generates a schema from the provided values.

"""

def __init__(self, values):

self._values = values

def _get_named_tuple_instance(self):

# omit fields with value=None since we can't infer their type

values = {

key: value

for key, value in self._values.items() if value is not None

}

schema = named_fields_to_schema([

(key, convert_to_typing_type(instance_to_type(value)))

for key, value in values.items()

])

return named_tuple_from_schema(schema)(**values)

class NamedTupleBasedPayloadBuilder(SchemaBasedPayloadBuilder):

"""

Build a payload based on a NamedTuple schema.

"""

def __init__(self, tuple_instance):

"""

:param tuple_instance: an instance of a typing.NamedTuple

"""

super().__init__()

self._tuple_instance = tuple_instance

def _get_named_tuple_instance(self):

return self._tuple_instance

class SchemaTransformPayloadBuilder(PayloadBuilder):

def __init__(self, identifier, **kwargs):

self._identifier = identifier

self._kwargs = kwargs

def identifier(self):

"""

The URN referencing this SchemaTransform

:return: str

"""

return self._identifier

def build(self):

schema_proto, payload = self._get_schema_proto_and_payload(**self._kwargs)

payload = external_transforms_pb2.SchemaTransformPayload(

identifier=self._identifier,

configuration_schema=schema_proto,

configuration_row=payload)

return payload

class ExplicitSchemaTransformPayloadBuilder(SchemaTransformPayloadBuilder):

def __init__(self, identifier, schema_proto, **kwargs):

self._identifier = identifier

self._schema_proto = schema_proto

self._kwargs = kwargs

def build(self):

def dict_to_row_recursive(field_type, py_value):

if py_value is None:

return None

type_info = field_type.WhichOneof('type_info')

if type_info == 'row_type':

return dict_to_row(field_type.row_type.schema, py_value)

elif type_info == 'array_type':

return [

dict_to_row_recursive(field_type.array_type.element_type, value)

for value in py_value

]

elif type_info == 'map_type':

return {

key: dict_to_row_recursive(field_type.map_type.value_type, value)

for key, value in py_value.items()

}

else:

return py_value

def dict_to_row(schema_proto, py_value):

row_type = named_tuple_from_schema(schema_proto)

if isinstance(py_value, dict):

extra = set(py_value.keys()) - set(row_type._fields)

if extra:

raise ValueError(

f"Transform '{self.identifier()}' was configured with unknown "

f"fields: {extra}. Valid fields: {set(row_type._fields)}")

return row_type(

*[

dict_to_row_recursive(

field.type, py_value.get(field.name, None))

for field in schema_proto.fields

])

else:

return row_type(py_value)

return external_transforms_pb2.SchemaTransformPayload(

identifier=self._identifier,

configuration_schema=self._schema_proto,

configuration_row=RowCoder(self._schema_proto).encode(

dict_to_row(self._schema_proto, self._kwargs)))

class JavaClassLookupPayloadBuilder(PayloadBuilder):

"""

Builds a payload for directly instantiating a Java transform using a

constructor and builder methods.

"""

IGNORED_ARG_FORMAT = 'ignore%d'

def __init__(self, class_name):

"""

:param class_name: fully qualified name of the transform class.

"""

if not class_name:

raise ValueError('Class name must not be empty')

self._class_name = class_name

self._constructor_method = None

self._constructor_param_args = None

self._constructor_param_kwargs = None

self._builder_methods_and_params = OrderedDict()

def _args_to_named_fields(self, args):

next_field_id = 0

named_fields = OrderedDict()

for value in args:

if value is None:

raise ValueError(

'Received value None. None values are currently not supported')

named_fields[(

JavaClassLookupPayloadBuilder.IGNORED_ARG_FORMAT %

next_field_id)] = value

next_field_id += 1

return named_fields

def build(self):

all_constructor_param_kwargs = self._args_to_named_fields(

self._constructor_param_args)

if self._constructor_param_kwargs:

all_constructor_param_kwargs.update(self._constructor_param_kwargs)

constructor_schema, constructor_payload = (

self._get_schema_proto_and_payload(**all_constructor_param_kwargs))

payload = external_transforms_pb2.JavaClassLookupPayload(

class_name=self._class_name,

constructor_schema=constructor_schema,

constructor_payload=constructor_payload)

if self._constructor_method:

payload.constructor_method = self._constructor_method

for builder_method_name, params in self._builder_methods_and_params.items():

builder_method_args, builder_method_kwargs = params

all_builder_method_kwargs = self._args_to_named_fields(

builder_method_args)

if builder_method_kwargs:

all_builder_method_kwargs.update(builder_method_kwargs)

builder_method_schema, builder_method_payload = (

self._get_schema_proto_and_payload(**all_builder_method_kwargs))

builder_method = external_transforms_pb2.BuilderMethod(

name=builder_method_name,

schema=builder_method_schema,

payload=builder_method_payload)

builder_method.name = builder_method_name

payload.builder_methods.append(builder_method)

return payload

def with_constructor(self, *args, **kwargs):

"""

Specifies the Java constructor to use.

Arguments provided using args and kwargs will be applied to the Java

transform constructor in the specified order.

:param args: parameter values of the constructor.

:param kwargs: parameter names and values of the constructor.

"""

if self._has_constructor():

raise ValueError(

'Constructor or constructor method can only be specified once')

self._constructor_param_args = args

self._constructor_param_kwargs = kwargs

def with_constructor_method(self, method_name, *args, **kwargs):

"""

Specifies the Java constructor method to use.

Arguments provided using args and kwargs will be applied to the Java

transform constructor method in the specified order.

:param method_name: name of the constructor method.

:param args: parameter values of the constructor method.

:param kwargs: parameter names and values of the constructor method.

"""

if self._has_constructor():

raise ValueError(

'Constructor or constructor method can only be specified once')

self._constructor_method = method_name

self._constructor_param_args = args

self._constructor_param_kwargs = kwargs

def add_builder_method(self, method_name, *args, **kwargs):

"""

Specifies a Java builder method to be invoked after instantiating the Java

transform class. Specified builder method will be applied in order.

Arguments provided using args and kwargs will be applied to the Java

transform builder method in the specified order.

:param method_name: name of the builder method.

:param args: parameter values of the builder method.

:param kwargs: parameter names and values of the builder method.

"""

self._builder_methods_and_params[method_name] = (args, kwargs)

def _has_constructor(self):

return (

self._constructor_method or self._constructor_param_args or

self._constructor_param_kwargs)

# Information regarding a SchemaTransform available in an external SDK.

SchemaTransformsConfig = namedtuple(

'SchemaTransformsConfig',

['identifier', 'configuration_schema', 'inputs', 'outputs', 'description'])

ManagedReplacement = namedtuple(

'ManagedReplacement',

['underlying_transform_identifier', 'update_compatibility_version'])

class SchemaAwareExternalTransform(ptransform.PTransform):

"""A proxy transform for SchemaTransforms implemented in external SDKs.

This allows Python pipelines to directly use existing SchemaTransforms

available to the expansion service without adding additional code in external

SDKs.

:param identifier: unique identifier of the SchemaTransform.

:param expansion_service: an expansion service to use. This should already be

available and the Schema-aware transforms to be used must already be

deployed.

:param rearrange_based_on_discovery: if this flag is set, the input kwargs

will be rearranged to match the order of fields in the external

SchemaTransform configuration. A discovery call will be made to fetch

the configuration.

:param classpath: (Optional) A list paths to additional jars to place on the

expansion service classpath.

:param managed_replacement: (Optional) a 'ManagedReplacement' namedtuple that

defines information needed to replace the transform with an equivalent

managed transform during the expansion. If an

'updateCompatibilityBeamVersion' pipeline option is provided, we will

only replace if the managed transform is update compatible with the

provided version.

:kwargs: field name to value mapping for configuring the schema transform.

keys map to the field names of the schema of the SchemaTransform

(in-order).

"""

def __init__(

self,

identifier,

expansion_service,

rearrange_based_on_discovery=False,

classpath=None,

managed_replacement=None,

**kwargs):

self._expansion_service = expansion_service

self._kwargs = kwargs

self._classpath = classpath

if managed_replacement:

assert isinstance(managed_replacement, ManagedReplacement)

self._managed_replacement = managed_replacement

_kwargs = kwargs

if rearrange_based_on_discovery:

config = SchemaAwareExternalTransform.discover_config(

self._expansion_service, identifier)

self._payload_builder = ExplicitSchemaTransformPayloadBuilder(

identifier,

named_tuple_to_schema(config.configuration_schema),

**_kwargs)

if self._managed_replacement:

# We have to do the replacement at the expansion instead of at

# construction

# since we don't have access to the PipelineOptions object at the

# construction.

underlying_transform_id = (

self._managed_replacement.underlying_transform_identifier)

if not (underlying_transform_id

in MANAGED_TRANSFORM_URN_TO_JAR_TARGET_MAPPING):

raise ValueError(

'Could not find an expansion service jar for the managed ' +

'transform ' + underlying_transform_id)

managed_expansion_service_jar = (

MANAGED_TRANSFORM_URN_TO_JAR_TARGET_MAPPING

)[underlying_transform_id]

self._managed_expansion_service = BeamJarExpansionService(

managed_expansion_service_jar)

managed_config = SchemaAwareExternalTransform.discover_config(

self._managed_expansion_service,

MANAGED_SCHEMA_TRANSFORM_IDENTIFIER)

yaml_config = yaml.dump(kwargs)

self._managed_payload_builder = (

ExplicitSchemaTransformPayloadBuilder(

MANAGED_SCHEMA_TRANSFORM_IDENTIFIER,

named_tuple_to_schema(managed_config.configuration_schema),

transform_identifier=underlying_transform_id,

config=yaml_config))

else:

self._payload_builder = SchemaTransformPayloadBuilder(

identifier, **_kwargs)

def expand(self, pcolls):

# Expand the transform using the expansion service.

payload_builder = self._payload_builder

expansion_service = self._expansion_service

if self._managed_replacement:

compat_version_prior_to_current = (

pcolls.pipeline._options.is_compat_version_prior_to(

self._managed_replacement.update_compatibility_version))

if not compat_version_prior_to_current:

payload_builder = self._managed_payload_builder

expansion_service = self._managed_expansion_service

return pcolls | self._payload_builder.identifier() >> ExternalTransform(

common_urns.schematransform_based_expand.urn,

payload_builder,

expansion_service)

@classmethod

@functools.lru_cache

def discover(cls, expansion_service, ignore_errors=False):

"""Discover all SchemaTransforms available to the given expansion service.

:return: a list of SchemaTransformsConfigs that represent the discovered

SchemaTransforms.

"""

return list(cls.discover_iter(expansion_service, ignore_errors))

@staticmethod

def discover_iter(expansion_service, ignore_errors=True):

with ExternalTransform.service(expansion_service) as service:

discover_response = service.DiscoverSchemaTransform(

beam_expansion_api_pb2.DiscoverSchemaTransformRequest())

for identifier in discover_response.schema_transform_configs:

proto_config = discover_response.schema_transform_configs[identifier]

try:

schema = named_tuple_from_schema(proto_config.config_schema)

except Exception as exn:

if ignore_errors:

truncated_schema = schema_pb2.Schema()

truncated_schema.CopyFrom(proto_config.config_schema)

for field in truncated_schema.fields:

try:

typing_from_runner_api(field.type)

except Exception:

if field.type.nullable:

# Set it to an empty placeholder type.

field.type.CopyFrom(

schema_pb2.FieldType(

nullable=True,

row_type=schema_pb2.RowType(

schema=schema_pb2.Schema())))

try:

schema = named_tuple_from_schema(truncated_schema)

except Exception as exn:

logging.info("Bad schema for %s: %s", identifier, str(exn)[:250])

continue

else:

raise

yield SchemaTransformsConfig(

identifier=identifier,

configuration_schema=schema,

inputs=proto_config.input_pcollection_names,

outputs=proto_config.output_pcollection_names,

description=proto_config.description)

@staticmethod

def discover_config(expansion_service, name):

"""Discover one SchemaTransform by name in the given expansion service.

:return: one SchemaTransformsConfig that represents the discovered

SchemaTransform

:raises:

ValueError: if more than one SchemaTransform is discovered, or if none

are discovered

"""

schematransforms = SchemaAwareExternalTransform.discover(

expansion_service, ignore_errors=True)

matched = []

for st in schematransforms:

if name in st.identifier:

matched.append(st)

if not matched:

raise ValueError(

"Did not discover any SchemaTransforms resembling the name '%s'" %

name)

elif len(matched) > 1:

raise ValueError(

"Found multiple SchemaTransforms with the name '%s':\n%s\n" %

(name, [st.identifier for st in matched]))

return matched[0]

class JavaExternalTransform(ptransform.PTransform):

"""A proxy for Java-implemented external transforms.

One builds these transforms just as one would in Java, e.g.::

transform = JavaExternalTransform('fully.qualified.ClassName'

)(contructorArg, ... ).builderMethod(...)

or::

JavaExternalTransform('fully.qualified.ClassName').staticConstructor(

...).builderMethod1(...).builderMethod2(...)

:param class_name: fully qualified name of the java class

:param expansion_service: (Optional) an expansion service to use. If none is

provided, a default expansion service will be started.

:param classpath: (Optional) A list paths to additional jars to place on the

expansion service classpath.

"""

def __init__(self, class_name, expansion_service=None, classpath=None):

if expansion_service and classpath:

raise ValueError(

f'Only one of expansion_service ({expansion_service}) '

f'or classpath ({classpath}) may be provided.')

self._payload_builder = JavaClassLookupPayloadBuilder(class_name)

self._classpath = classpath

self._expansion_service = expansion_service

# Beam explicitly looks for following attributes. Hence adding

# 'None' values here to prevent '__getattr__' from being called.

self.inputs = None

self._fn_api_payload = None

def __call__(self, *args, **kwargs):

self._payload_builder.with_constructor(*args, **kwargs)

return self

def __getattr__(self, name):

# Don't try to emulate special methods.

if name.startswith('__') and name.endswith('__'):

return super().__getattr__(name)

else:

return self[name]

def __getitem__(self, name):

# Use directly for keywords or attribute conflicts.

def construct(*args, **kwargs):

if self._payload_builder._has_constructor():

builder_method = self._payload_builder.add_builder_method

else:

builder_method = self._payload_builder.with_constructor_method

builder_method(name, *args, **kwargs)

return self

return construct

def expand(self, pcolls):

if self._expansion_service is None:

self._expansion_service = BeamJarExpansionService(

':sdks:java:expansion-service:app:shadowJar',

extra_args=['{{PORT}}', '--javaClassLookupAllowlistFile=*'],

classpath=self._classpath)

return pcolls | ExternalTransform(

common_urns.java_class_lookup.urn,

self._payload_builder,

self._expansion_service)

class AnnotationBasedPayloadBuilder(SchemaBasedPayloadBuilder):

"""

Build a payload based on an external transform's type annotations.

"""

def __init__(self, transform, **values):

"""

:param transform: a PTransform instance or class. type annotations will

be gathered from its __init__ method

:param values: values to encode

"""

self._transform = transform

self._values = values

def _get_named_tuple_instance(self):

schema = named_fields_to_schema([

(k, convert_to_typing_type(v))

for k, v in self._transform.__init__.__annotations__.items()

if k in self._values

])

return named_tuple_from_schema(schema)(**self._values)

class DataclassBasedPayloadBuilder(SchemaBasedPayloadBuilder):

"""

Build a payload based on an external transform that uses dataclasses.

"""

def __init__(self, transform):

"""

:param transform: a dataclass-decorated PTransform instance from which to

gather type annotations and values

"""

self._transform = transform

def _get_named_tuple_instance(self):

import dataclasses

schema = named_fields_to_schema([

(field.name, convert_to_typing_type(field.type))

for field in dataclasses.fields(self._transform)

])

return named_tuple_from_schema(schema)(

**dataclasses.asdict(self._transform))

class ExternalTransform(ptransform.PTransform):

"""

External provides a cross-language transform via expansion services in

foreign SDKs.

"""

_namespace_counter = 0

# Variable name _namespace conflicts with DisplayData._namespace so we use

# name _external_namespace here.

_external_namespace = threading.local()

_IMPULSE_PREFIX = 'impulse'

def __init__(self, urn, payload, expansion_service=None):

"""Wrapper for an external transform with the given urn and payload.

:param urn: the unique beam identifier for this transform

:param payload: the payload, either as a byte string or a PayloadBuilder

:param expansion_service: an expansion service implementing the beam

ExpansionService protocol, either as an object with an Expand method

or an address (as a str) to a grpc server that provides this method.

"""

expansion_service = expansion_service or DEFAULT_EXPANSION_SERVICE

if not urn and isinstance(payload, JavaClassLookupPayloadBuilder):

urn = common_urns.java_class_lookup.urn

self._urn = urn

self._payload = (

payload.payload() if isinstance(payload, PayloadBuilder) else payload)

self._expansion_service = expansion_service

self._external_namespace = self._fresh_namespace()

self._inputs: dict[str, pvalue.PCollection] = {}

self._outputs: dict[str, pvalue.PCollection] = {}

def with_output_types(self, *args, **kwargs):

return WithTypeHints.with_output_types(self, *args, **kwargs)

def replace_named_inputs(self, named_inputs):

self._inputs = named_inputs

def replace_named_outputs(self, named_outputs):

self._outputs = named_outputs

def __post_init__(self, expansion_service):

"""

This will only be invoked if ExternalTransform is used as a base class

for a class decorated with dataclasses.dataclass

"""

ExternalTransform.__init__(

self, self.URN, DataclassBasedPayloadBuilder(self), expansion_service)

def default_label(self):

return '%s(%s)' % (self.__class__.__name__, self._urn)

@classmethod

def get_local_namespace(cls):

return getattr(cls._external_namespace, 'value', 'external')

@classmethod

@contextlib.contextmanager

def outer_namespace(cls, namespace):

prev = cls.get_local_namespace()

try:

cls._external_namespace.value = namespace

yield

finally:

cls._external_namespace.value = prev

@classmethod

def _fresh_namespace(cls) -> str:

ExternalTransform._namespace_counter += 1

return '%s_%d' % (cls.get_local_namespace(), cls._namespace_counter)

def expand(self, pvalueish: pvalue.PCollection) -> pvalue.PCollection:

if isinstance(pvalueish, pvalue.PBegin):

self._inputs = {}

elif isinstance(pvalueish, (list, tuple)):

self._inputs = {str(ix): pvalue for ix, pvalue in enumerate(pvalueish)}

elif isinstance(pvalueish, dict):

self._inputs = pvalueish

else:

self._inputs = {'input': pvalueish}

pipeline = (

next(iter(self._inputs.values())).pipeline

if self._inputs else pvalueish.pipeline)

context = pipeline_context.PipelineContext(

component_id_map=pipeline.component_id_map)

transform_proto = beam_runner_api_pb2.PTransform(

unique_name=pipeline._current_transform().full_label,

spec=beam_runner_api_pb2.FunctionSpec(

urn=self._urn, payload=self._payload))

for tag, pcoll in self._inputs.items():

transform_proto.inputs[tag] = context.pcollections.get_id(pcoll)

# Conversion to/from proto assumes producers.

# TODO: Possibly loosen this.

context.transforms.put_proto(

'%s_%s' % (self._IMPULSE_PREFIX, tag),

beam_runner_api_pb2.PTransform(

unique_name='%s_%s' % (self._IMPULSE_PREFIX, tag),

spec=beam_runner_api_pb2.FunctionSpec(

urn=common_urns.primitives.IMPULSE.urn),

outputs={'out': transform_proto.inputs[tag]}))

# Retrieve type hints and store them in variables

# to avoid duplicate calls and AttributeError

hints = self.get_type_hints()

output_coders = None

if hints.output_types:

if hints.output_types[0]:

output_coders = dict((str(k), context.coder_id_from_element_type(v))

for (k, v) in enumerate(hints.output_types[0]))

elif hints.output_types[1]:

output_coders = {

k: context.coder_id_from_element_type(v)

for (k, v) in hints.output_types[1].items()

}

components = context.to_runner_api()

request = beam_expansion_api_pb2.ExpansionRequest(

components=components,

namespace=self._external_namespace, # type: ignore[arg-type]

transform=transform_proto,

output_coder_requests=output_coders,

pipeline_options=pipeline._options.to_runner_api())

expansion_service = _maybe_use_transform_service(

self._expansion_service, pipeline.options)

with ExternalTransform.service(expansion_service) as service:

response = service.Expand(request)

if response.error:

raise RuntimeError(_sanitize_java_traceback(response.error))

self._expanded_components = response.components

if any(e.dependencies

for env in self._expanded_components.environments.values()

for e in environments.expand_anyof_environments(env)):

self._expanded_components = self._resolve_artifacts(

self._expanded_components,

service.artifact_service(),

pipeline.local_tempdir)

self._expanded_transform = response.transform

self._expanded_requirements = response.requirements

result_context = pipeline_context.PipelineContext(response.components)

def fix_output(pcoll, tag):

pcoll.pipeline = pipeline

pcoll.tag = tag

return pcoll

self._outputs = {

tag: fix_output(result_context.pcollections.get_by_id(pcoll_id), tag)

for tag, pcoll_id in self._expanded_transform.outputs.items()

}

return self._output_to_pvalueish(self._outputs)

@staticmethod

@contextlib.contextmanager

def service(expansion_service):

if isinstance(expansion_service, str):

channel_options = [("grpc.max_receive_message_length", -1),

("grpc.max_send_message_length", -1)]

if hasattr(grpc, 'local_channel_credentials'):

# Some environments may not support insecure channels. Hence use a

# secure channel with local credentials here.

# TODO: update this to support secure non-local channels.

channel_factory_fn = functools.partial(

grpc.secure_channel,

expansion_service,

grpc.local_channel_credentials(),

options=channel_options)

else:

# local_channel_credentials is an experimental API which is unsupported

# by older versions of grpc which may be pulled in due to other project

# dependencies.

channel_factory_fn = functools.partial(

grpc.insecure_channel, expansion_service, options=channel_options)

with channel_factory_fn() as channel:

yield ExpansionAndArtifactRetrievalStub(channel)

elif hasattr(expansion_service, 'Expand'):

yield expansion_service

else:

with expansion_service as stub:

yield stub

def _resolve_artifacts(self, components, service, dest):

def _resolve_artifacts_for(env):

if env.urn == common_urns.environments.ANYOF.urn:

env.CopyFrom(

environments.AnyOfEnvironment.create_proto([

_resolve_artifacts_for(e)

for e in environments.expand_anyof_environments(env)

]))

elif env.dependencies:

resolved = list(

artifact_service.resolve_artifacts(env.dependencies, service, dest))

del env.dependencies[:]

env.dependencies.extend(resolved)

return env

for env in components.environments.values():

_resolve_artifacts_for(env)

return components

def _output_to_pvalueish(self, output_dict):

if len(output_dict) == 1:

return next(iter(output_dict.values()))

else:

return output_dict

def to_runner_api_transform(self, context, full_label):

pcoll_renames = {}

renamed_tag_seen = False

for tag, pcoll in self._inputs.items():

if tag not in self._expanded_transform.inputs:

if renamed_tag_seen:

raise RuntimeError(

'Ambiguity due to non-preserved tags: %s vs %s' % (

sorted(self._expanded_transform.inputs.keys()),

sorted(self._inputs.keys())))

else:

renamed_tag_seen = True

tag, = self._expanded_transform.inputs.keys()

pcoll_renames[self._expanded_transform.inputs[tag]] = (

context.pcollections.get_id(pcoll))

for tag, pcoll in self._outputs.items():

pcoll_renames[self._expanded_transform.outputs[tag]] = (

context.pcollections.get_id(pcoll))

def _equivalent(coder1, coder2):

return coder1 == coder2 or _normalize(coder1) == _normalize(coder2)

def _normalize(coder_proto):

normalized = copy.copy(coder_proto)

normalized.spec.environment_id = ''

# TODO(robertwb): Normalize components as well.

return normalized

for id, proto in self._expanded_components.coders.items():

if id.startswith(self._external_namespace):

context.coders.put_proto(id, proto)

elif id in context.coders:

if not _equivalent(context.coders._id_to_proto[id], proto):

raise RuntimeError(

'Re-used coder id: %s\n%s\n%s' %

(id, context.coders._id_to_proto[id], proto))

else:

context.coders.put_proto(id, proto)

for id, proto in self._expanded_components.windowing_strategies.items():

if id.startswith(self._external_namespace):

context.windowing_strategies.put_proto(id, proto)

for id, proto in self._expanded_components.environments.items():

if id.startswith(self._external_namespace):

context.environments.put_proto(id, proto)

for id, proto in self._expanded_components.pcollections.items():

id = pcoll_renames.get(id, id)

if id not in context.pcollections._id_to_obj.keys():

context.pcollections.put_proto(id, proto)

for id, proto in self._expanded_components.transforms.items():

if id.startswith(self._IMPULSE_PREFIX):

# Our fake inputs.

continue

assert id.startswith(

self._external_namespace), (id, self._external_namespace)

new_proto = beam_runner_api_pb2.PTransform(

unique_name=proto.unique_name,

# If URN is not set this is an empty spec.

spec=proto.spec if proto.spec.urn else None,

subtransforms=proto.subtransforms,

inputs={

tag: pcoll_renames.get(pcoll, pcoll)

for tag, pcoll in proto.inputs.items()

},

outputs={

tag: pcoll_renames.get(pcoll, pcoll)

for tag, pcoll in proto.outputs.items()

},

display_data=proto.display_data,

environment_id=proto.environment_id)

context.transforms.put_proto(id, new_proto)

for requirement in self._expanded_requirements:

context.add_requirement(requirement)

return beam_runner_api_pb2.PTransform(

unique_name=full_label,

spec=self._expanded_transform.spec,

subtransforms=self._expanded_transform.subtransforms,

inputs={

tag: pcoll_renames.get(pcoll, pcoll)

for tag, pcoll in self._expanded_transform.inputs.items()

},

outputs={

tag: pcoll_renames.get(pcoll, pcoll)

for tag, pcoll in self._expanded_transform.outputs.items()

},