public static Optional<Environment> getEnvironment(String ptransformId, Components components) { try { PTransform ptransform = components.getTransformsOrThrow(ptransformId); String envId = KNOWN_URN_SPEC_EXTRACTORS .getOrDefault(ptransform.getSpec().getUrn(), DEFAULT_SPEC_EXTRACTOR) .getEnvironmentId(ptransform); if (Strings.isNullOrEmpty(envId)) { // Some PTransform payloads may have an unspecified (empty) Environment ID, for example a // WindowIntoPayload with a known WindowFn. Others will never have an Environment ID, such // as a GroupByKeyPayload, and the Default extractor returns null in this case. return Optional.empty(); } else { return Optional.of(components.getEnvironmentsOrThrow(envId)); } } catch (IOException e) { throw new RuntimeException(e); } }
@Test public void getEnvironmentRead() throws IOException { SdkComponents components = SdkComponents.create(); components.registerEnvironment(Environments.createDockerEnvironment("java")); ReadPayload payload = ReadTranslation.toProto(Read.from(CountingSource.unbounded()), components); RehydratedComponents rehydratedComponents = RehydratedComponents.forComponents(components.toComponents()); PTransform builder = PTransform.newBuilder() .setSpec( FunctionSpec.newBuilder() .setUrn(PTransformTranslation.COMBINE_PER_KEY_TRANSFORM_URN) .setPayload(payload.toByteString()) .build()) .build(); Environment env = Environments.getEnvironment(builder, rehydratedComponents).get(); assertThat( env, equalTo( components .toComponents() .getEnvironmentsOrThrow(payload.getSource().getEnvironmentId()))); }
@Test public void getEnvironmentCombine() throws IOException { SdkComponents components = SdkComponents.create(); components.registerEnvironment(Environments.createDockerEnvironment("java")); CombinePayload payload = CombinePayload.newBuilder() .setCombineFn(CombineTranslation.toProto(Sum.ofLongs(), components)) .build(); RehydratedComponents rehydratedComponents = RehydratedComponents.forComponents(components.toComponents()); PTransform builder = PTransform.newBuilder() .setSpec( FunctionSpec.newBuilder() .setUrn(PTransformTranslation.COMBINE_PER_KEY_TRANSFORM_URN) .setPayload(payload.toByteString()) .build()) .build(); Environment env = Environments.getEnvironment(builder, rehydratedComponents).get(); assertThat( env, equalTo( components .toComponents() .getEnvironmentsOrThrow(payload.getCombineFn().getEnvironmentId()))); } }
components .toComponents() .getEnvironmentsOrThrow(payload.getWindowFn().getEnvironmentId())));
components .toComponents() .getEnvironmentsOrThrow(payload.getDoFn().getEnvironmentId())));
@Test public void getEnvironmentWindowIntoKnown() throws IOException { SdkComponents components = SdkComponents.create(); components.registerEnvironment(Environments.createDockerEnvironment("java")); WindowIntoPayload payload = WindowIntoPayload.newBuilder() .setWindowFn( WindowingStrategyTranslation.toProto( FixedWindows.of(Duration.standardMinutes(5L)), components)) .build(); RehydratedComponents rehydratedComponents = RehydratedComponents.forComponents(components.toComponents()); PTransform builder = PTransform.newBuilder() .setSpec( FunctionSpec.newBuilder() .setUrn(PTransformTranslation.ASSIGN_WINDOWS_TRANSFORM_URN) .setPayload(payload.toByteString()) .build()) .build(); Environment env = Environments.getEnvironment(builder, rehydratedComponents).get(); assertThat( env, equalTo( components .toComponents() .getEnvironmentsOrThrow(payload.getWindowFn().getEnvironmentId()))); }
/** Returns the {@link Environment} associated with the given ID. */ public Environment getEnvironment(String environmentId) { return components.getEnvironmentsOrThrow(environmentId); }