DirectTransformExecutor implements TransformExecutor { @Override public void run() { MetricsContainerImpl metricsContainer = new MetricsContainerImpl(transform.getFullName()); try (Closeable metricsScope = MetricsEnvironment.scopedMetricsContainer(metricsContainer)) { Collection<ModelEnforcement<T>> enforcements = new ArrayList<>(); for (ModelEnforcementFactory enforcementFactory : modelEnforcements) { ModelEnforcement<T> enforcement = enforcementFactory.forBundle(inputBundle, transform); enforcements.add(enforcement); } TransformEvaluator<T> evaluator = evaluatorRegistry.forApplication(transform, inputBundle); if (evaluator == null) { onComplete.handleEmpty(transform); return; } processElements(evaluator, metricsContainer, enforcements); finishBundle(evaluator, metricsContainer, enforcements); } catch (Exception e) { onComplete.handleException(inputBundle, e); if (e instanceof RuntimeException) { throw (RuntimeException) e; } throw new RuntimeException(e); } catch (Error err) { LOG.error("Error occurred within {}", this, err); onComplete.handleError(err); throw err; } finally { context.getMetrics().commitPhysical(inputBundle, metricsContainer.getCumulative()); transformEvaluationState.complete(this); } } @VisibleForTesting  DirectTransformExecutor(
      EvaluationContext context,
      TransformEvaluatorRegistry factory,
      Iterable<? extends ModelEnforcementFactory> modelEnforcements,
      CommittedBundle<T> inputBundle,
      AppliedPTransform<?, ?, ?> transform,
      CompletionCallback completionCallback,
      TransformExecutorService transformEvaluationState); @Override void run();  }
@Test public void callWithNullInputBundleFinishesBundleAndCompletes() throws Exception { final TransformResult<Object> result = StepTransformResult.withoutHold(createdProducer).build(); final AtomicBoolean finishCalled = new AtomicBoolean(false); TransformEvaluator<Object> evaluator = new TransformEvaluator<Object>() { @Override public void processElement(WindowedValue<Object> element) throws Exception { throw new IllegalArgumentException("Shouldn't be called"); } @Override public TransformResult<Object> finishBundle() throws Exception { finishCalled.set(true); return result; } }; when(registry.forApplication(createdProducer, null)).thenReturn(evaluator); DirectTransformExecutor<Object> executor = new DirectTransformExecutor<>( evaluationContext, registry, Collections.emptyList(), null, createdProducer, completionCallback, transformEvaluationState); executor.run(); assertThat(finishCalled.get(), is(true)); assertThat(completionCallback.handledResult, equalTo(result)); assertThat(completionCallback.handledException, is(nullValue())); }
@Test public void nullTransformEvaluatorTerminates() throws Exception { when(registry.forApplication(createdProducer, null)).thenReturn(null); DirectTransformExecutor<Object> executor = new DirectTransformExecutor<>( evaluationContext, registry, Collections.emptyList(), null, createdProducer, completionCallback, transformEvaluationState); executor.run(); assertThat(completionCallback.handledResult, is(nullValue())); assertThat(completionCallback.handledEmpty, equalTo(true)); assertThat(completionCallback.handledException, is(nullValue())); }
@Test public void callWithEnforcementAppliesEnforcement() throws Exception { final TransformResult<Object> result = StepTransformResult.withoutHold(downstreamProducer).build(); TransformEvaluator<Object> evaluator = new TransformEvaluator<Object>() { @Override public void processElement(WindowedValue<Object> element) throws Exception {} @Override public TransformResult<Object> finishBundle() throws Exception { return result; } }; WindowedValue<String> fooElem = WindowedValue.valueInGlobalWindow("foo"); WindowedValue<String> barElem = WindowedValue.valueInGlobalWindow("bar"); CommittedBundle<String> inputBundle = bundleFactory.createBundle(created).add(fooElem).add(barElem).commit(Instant.now()); when(registry.forApplication(downstreamProducer, inputBundle)).thenReturn(evaluator); TestEnforcementFactory enforcement = new TestEnforcementFactory(); DirectTransformExecutor<String> executor = new DirectTransformExecutor<>( evaluationContext, registry, Collections.<ModelEnforcementFactory>singleton(enforcement), inputBundle, downstreamProducer, completionCallback, transformEvaluationState); executor.run(); TestEnforcement<?> testEnforcement = enforcement.instance; assertThat(testEnforcement.beforeElements, containsInAnyOrder(barElem, fooElem)); assertThat(testEnforcement.afterElements, containsInAnyOrder(barElem, fooElem)); assertThat(testEnforcement.finishedBundles, Matchers.contains(result)); }