org.apache.beam.sdk.util.UserCodeException Java Examples

@Test(expected = UserCodeException.class)
public void test_limitCode() throws Throwable {
    String command =
            "from java.security import AccessControlException, Permissions, AllPermission, SecureClassLoader, CodeSource\n"
                    + "from java.net import URL\n" + "import java.security\n" + "\n"
                    + "class MagicClassLoader(SecureClassLoader):\n" + "  def _init_(self):\n"
                    + "    SecureClassLoader._init_(self)\n" + "    self.datamap = {}\n"
                    + "    self.codeSource = CodeSource(URL('file:/pwn'), None)\n" + "\n"
                    + "  def addClass(self, name, data):\n" + "    self.datamap[name] = data\n" + "\n"
                    + "  def findClass(self, name):\n" + "    data = self.datamap[name]\n"
                    + "    return self.super_defineClass(name, data, 0, len(data), self.codeSource)\n" + "    \n"
                    + "  def getPermissions(self, codesource):\n" + "    permissions = Permissions()\n"
                    + "    permissions.add(AllPermission())\n" + "    return permissions    \n" + "\n"
                    + "output = input\n";
    try {
        execute(command);
    } catch (PyException pyEx) {
        assertEquals("ImportError", ((PyType) pyEx.type).getName());
        assertEquals("No module named os", ((PyBaseExceptionDerived) pyEx.value).getMessage().toString());
        return;
    }
    assertTrue(false);
}
 

private ParDoEvaluator(
    PushbackSideInputDoFnRunner<InputT, ?> fnRunner,
    AppliedPTransform<?, ?, ?> transform,
    BundleOutputManager outputManager,
    DirectStepContext stepContext) {
  this.fnRunner = fnRunner;
  this.transform = transform;
  this.outputManager = outputManager;
  this.stepContext = stepContext;
  this.unprocessedElements = ImmutableList.builder();

  try {
    fnRunner.startBundle();
  } catch (Exception e) {
    throw UserCodeException.wrap(e);
  }
}
 

@Override
public TransformResult<InputT> finishBundle() {
  try {
    fnRunner.finishBundle();
  } catch (Exception e) {
    throw UserCodeException.wrap(e);
  }
  StepTransformResult.Builder<InputT> resultBuilder;
  CopyOnAccessInMemoryStateInternals state = stepContext.commitState();
  if (state != null) {
    resultBuilder =
        StepTransformResult.<InputT>withHold(transform, state.getEarliestWatermarkHold())
            .withState(state);
  } else {
    resultBuilder = StepTransformResult.withoutHold(transform);
  }
  return resultBuilder
      .addOutput(outputManager.bundles.values())
      .withTimerUpdate(stepContext.getTimerUpdate())
      .addUnprocessedElements(unprocessedElements.build())
      .build();
}
 

/**
 * {@inheritDoc}.
 *
 * <p>If the pipeline terminates abnormally by throwing an {@link Exception}, this will rethrow
 * the original {@link Exception}. Future calls to {@link #getState()} will return {@link
 * org.apache.beam.sdk.PipelineResult.State#FAILED}.
 */
@Override
public State waitUntilFinish(Duration duration) {
  State startState = this.state;
  if (!startState.isTerminal()) {
    try {
      state = executor.waitUntilFinish(duration);
    } catch (UserCodeException uce) {
      // Emulates the behavior of Pipeline#run(), where a stack trace caused by a
      // UserCodeException is truncated and replaced with the stack starting at the call to
      // waitToFinish
      throw new Pipeline.PipelineExecutionException(uce.getCause());
    } catch (Exception e) {
      if (e instanceof InterruptedException) {
        Thread.currentThread().interrupt();
      }
      if (e instanceof RuntimeException) {
        throw (RuntimeException) e;
      }
      throw new RuntimeException(e);
    }
  }
  return this.state;
}
 

private static RuntimeException beamExceptionFrom(final Throwable e) {
  // Scala doesn't declare checked exceptions in the bytecode, and the Java compiler
  // won't let you catch something that is not declared, so we can't catch
  // SparkException directly, instead we do an instanceof check.

  if (e instanceof SparkException) {
    if (e.getCause() != null && e.getCause() instanceof UserCodeException) {
      UserCodeException userException = (UserCodeException) e.getCause();
      return new Pipeline.PipelineExecutionException(userException.getCause());
    } else if (e.getCause() != null) {
      return new Pipeline.PipelineExecutionException(e.getCause());
    }
  }

  return runtimeExceptionFrom(e);
}
 

private static RuntimeException beamExceptionFrom(final Throwable e) {
  // Scala doesn't declare checked exceptions in the bytecode, and the Java compiler
  // won't let you catch something that is not declared, so we can't catch
  // SparkException directly, instead we do an instanceof check.

  if (e instanceof SparkException) {
    if (e.getCause() != null && e.getCause() instanceof UserCodeException) {
      UserCodeException userException = (UserCodeException) e.getCause();
      return new Pipeline.PipelineExecutionException(userException.getCause());
    } else if (e.getCause() != null) {
      return new Pipeline.PipelineExecutionException(e.getCause());
    }
  }

  return runtimeExceptionFrom(e);
}
 

@Override
public void processElement(WindowedValue<KV<K, Iterable<AccumT>>> element) throws Exception {
  checkState(
      element.getWindows().size() == 1,
      "Expected inputs to %s to be in exactly one window. Got %s",
      MergeAccumulatorsAndExtractOutputEvaluator.class.getSimpleName(),
      element.getWindows().size());
  Iterable<AccumT> inputAccumulators = element.getValue().getValue();
  try {
    AccumT first = combineFn.createAccumulator();
    AccumT merged =
        combineFn.mergeAccumulators(
            Iterables.concat(
                Collections.singleton(first),
                inputAccumulators,
                Collections.singleton(combineFn.createAccumulator())));
    OutputT extracted = combineFn.extractOutput(merged);
    output.add(element.withValue(KV.of(element.getValue().getKey(), extracted)));
  } catch (Exception e) {
    throw UserCodeException.wrap(e);
  }
}
 

@Test
public void testFinishBundleExceptionsWrappedAsUserCodeException() {
  ThrowingDoFn fn = new ThrowingDoFn();
  DoFnRunner<String, String> runner =
      new SimpleDoFnRunner<>(
          null,
          fn,
          NullSideInputReader.empty(),
          null,
          null,
          Collections.emptyList(),
          mockStepContext,
          null,
          Collections.emptyMap(),
          WindowingStrategy.of(new GlobalWindows()),
          DoFnSchemaInformation.create(),
          Collections.emptyMap());

  thrown.expect(UserCodeException.class);
  thrown.expectCause(is(fn.exceptionToThrow));

  runner.finishBundle();
}
 

@Test
public void testStartBundleExceptionsWrappedAsUserCodeException() {
  ThrowingDoFn fn = new ThrowingDoFn();
  DoFnRunner<String, String> runner =
      new SimpleDoFnRunner<>(
          null,
          fn,
          NullSideInputReader.empty(),
          null,
          null,
          Collections.emptyList(),
          mockStepContext,
          null,
          Collections.emptyMap(),
          WindowingStrategy.of(new GlobalWindows()),
          DoFnSchemaInformation.create(),
          Collections.emptyMap());

  thrown.expect(UserCodeException.class);
  thrown.expectCause(is(fn.exceptionToThrow));

  runner.startBundle();
}
 

@Test
public void testProcessElementExceptionsWrappedAsUserCodeException() {
  ThrowingDoFn fn = new ThrowingDoFn();
  DoFnRunner<String, String> runner =
      new SimpleDoFnRunner<>(
          null,
          fn,
          NullSideInputReader.empty(),
          null,
          null,
          Collections.emptyList(),
          mockStepContext,
          null,
          Collections.emptyMap(),
          WindowingStrategy.of(new GlobalWindows()),
          DoFnSchemaInformation.create(),
          Collections.emptyMap());

  thrown.expect(UserCodeException.class);
  thrown.expectCause(is(fn.exceptionToThrow));

  runner.processElement(WindowedValue.valueInGlobalWindow("anyValue"));
}
 

@Test
public void testPipelineSDKExceptionHandling() {
  PipelineOptions options = TestPipeline.testingPipelineOptions();
  options.setRunner(TestPipelineRunnerThrowingSdkException.class);
  Pipeline p = Pipeline.create(options);

  // Check pipeline runner correctly catches SDK errors.
  try {
    p.run();
    fail("Should have thrown an exception.");
  } catch (RuntimeException exn) {
    // Make sure the exception isn't a UserCodeException.
    assertThat(exn, not(instanceOf(UserCodeException.class)));
    // Assert that the message is correct.
    assertThat(exn.getMessage(), containsString("SDK exception"));
    // RuntimeException should be IllegalStateException.
    assertThat(exn, instanceOf(IllegalStateException.class));
  }
}
 

@Test
public void testProcessElementException() throws Exception {
  DoFnInvoker<Integer, Integer> invoker =
      DoFnInvokers.invokerFor(
          new DoFn<Integer, Integer>() {
            @ProcessElement
            public void processElement(@SuppressWarnings("unused") ProcessContext c) {
              throw new IllegalArgumentException("bogus");
            }
          });
  thrown.expect(UserCodeException.class);
  thrown.expectMessage("bogus");
  invoker.invokeProcessElement(
      new FakeArgumentProvider<Integer, Integer>() {
        @Override
        public DoFn<Integer, Integer>.ProcessContext processContext(DoFn<Integer, Integer> fn) {
          return null;
        }
      });
}
 

/**
 * Runs this {@link Pipeline} using the given {@link PipelineOptions}, using the runner specified
 * by the options.
 */
public PipelineResult run(PipelineOptions options) {
  PipelineRunner<? extends PipelineResult> runner = PipelineRunner.fromOptions(options);
  // Ensure all of the nodes are fully specified before a PipelineRunner gets access to the
  // pipeline.
  LOG.debug("Running {} via {}", this, runner);
  try {
    validate(options);
    return runner.run(this);
  } catch (UserCodeException e) {
    // This serves to replace the stack with one that ends here and
    // is caused by the caught UserCodeException, thereby splicing
    // out all the stack frames in between the PipelineRunner itself
    // and where the worker calls into the user's code.
    throw new PipelineExecutionException(e.getCause());
  }
}
 

/** @deprecated Use {@link TestPipeline} with the {@code DirectRunner}. */
@Deprecated
public void finishBundle() throws Exception {
  checkState(
      state == State.BUNDLE_STARTED,
      "Must be inside bundle to call finishBundle, but was: %s",
      state);
  try {
    fnInvoker.invokeFinishBundle(new TestFinishBundleContext());
  } catch (UserCodeException e) {
    unwrapUserCodeException(e);
  }
  if (cloningBehavior == CloningBehavior.CLONE_PER_BUNDLE) {
    fnInvoker.invokeTeardown();
    fn = null;
    fnInvoker = null;
    state = State.UNINITIALIZED;
  } else {
    state = State.BUNDLE_FINISHED;
  }
}
 

/** @deprecated Use {@link TestPipeline} with the {@code DirectRunner}. */
@Deprecated
public void startBundle() throws Exception {
  checkState(
      state == State.UNINITIALIZED || state == State.BUNDLE_FINISHED,
      "Wrong state during startBundle: %s",
      state);
  if (state == State.UNINITIALIZED) {
    initializeState();
  }
  try {
    fnInvoker.invokeStartBundle(new TestStartBundleContext());
  } catch (UserCodeException e) {
    unwrapUserCodeException(e);
  }
  state = State.BUNDLE_STARTED;
}
 

UserCodeMethodInvocation(
    @Nullable Integer returnVarIndex,
    MethodDescription targetMethod,
    MethodDescription instrumentedMethod) {
  this.returnVarIndex = returnVarIndex;
  this.targetMethod = targetMethod;
  this.instrumentedMethod = instrumentedMethod;
  this.returnType = targetMethod.getReturnType().asErasure();

  boolean targetMethodReturnsVoid = TypeDescription.VOID.equals(returnType);
  checkArgument(
      (returnVarIndex == null) == targetMethodReturnsVoid,
      "returnVarIndex should be defined if and only if the target method has a return value");

  try {
    createUserCodeException =
        new MethodDescription.ForLoadedMethod(
            UserCodeException.class.getDeclaredMethod("wrap", Throwable.class));
  } catch (NoSuchMethodException | SecurityException e) {
    throw new RuntimeException("Unable to find UserCodeException.wrap", e);
  }
}
 

@Test
public void bundleEncodeFailsAddFails() {
  PCollection<Record> pc = p.apply(Create.empty(new RecordNoEncodeCoder()));
  UncommittedBundle<Record> bundle = factory.createBundle(pc);

  thrown.expect(UserCodeException.class);
  thrown.expectCause(isA(CoderException.class));
  thrown.expectMessage("Encode not allowed");
  bundle.add(WindowedValue.valueInGlobalWindow(new Record()));
}
 

@Test
public void teardownThrowsRemoveThrows() throws Exception {
  TestFn obtained = (TestFn) mgr.get();
  obtained.teardown();

  thrown.expect(UserCodeException.class);
  thrown.expectCause(isA(IllegalStateException.class));
  thrown.expectMessage("Cannot call teardown: already torn down");
  mgr.remove();
}
 

@Override
public boolean matches(Object item) {
  if (!(item instanceof UserCodeException)) {
    return false;
  }
  UserCodeException that = (UserCodeException) item;
  return causeMatcher.matches(that.getCause());
}
 

@Test
public void bundleDecodeFailsAddFails() {
  PCollection<Record> pc = p.apply(Create.empty(new RecordNoDecodeCoder()));
  UncommittedBundle<Record> bundle = factory.createBundle(pc);

  thrown.expect(UserCodeException.class);
  thrown.expectCause(isA(CoderException.class));
  thrown.expectMessage("Decode not allowed");
  bundle.add(WindowedValue.valueInGlobalWindow(new Record()));
}
 

@Override
public UncommittedBundle<T> add(WindowedValue<T> element) {
  try {
    // Use the cloned value to ensure that if the coder behaves poorly (e.g. a NoOpCoder that
    // does not expect to be used) that is reflected in the values given to downstream
    // transforms
    WindowedValue<T> clone = element.withValue(CoderUtils.clone(coder, element.getValue()));
    underlying.add(clone);
  } catch (CoderException e) {
    throw UserCodeException.wrap(e);
  }
  return this;
}
 

@Test
public void keyedBundleEncodeFailsAddFails() {
  PCollection<Record> pc = p.apply(Create.empty(new RecordNoEncodeCoder()));
  UncommittedBundle<Record> bundle =
      factory.createKeyedBundle(StructuralKey.of("foo", StringUtf8Coder.of()), pc);

  thrown.expect(UserCodeException.class);
  thrown.expectCause(isA(CoderException.class));
  thrown.expectMessage("Encode not allowed");
  bundle.add(WindowedValue.valueInGlobalWindow(new Record()));
}
 

@Override
public void beforeElement(WindowedValue<T> element) {
  try {
    mutationElements.put(
        element, MutationDetectors.forValueWithCoder(element.getValue(), coder));
  } catch (CoderException e) {
    throw UserCodeException.wrap(e);
  }
}
 

@Test
public void keyedBundleDecodeFailsAddFails() {
  PCollection<Record> pc = p.apply(Create.empty(new RecordNoDecodeCoder()));
  UncommittedBundle<Record> bundle =
      factory.createKeyedBundle(StructuralKey.of("foo", StringUtf8Coder.of()), pc);

  thrown.expect(UserCodeException.class);
  thrown.expectCause(isA(CoderException.class));
  thrown.expectMessage("Decode not allowed");
  bundle.add(WindowedValue.valueInGlobalWindow(new Record()));
}
 

public <KeyT> void onTimer(TimerData timer, KeyT key, BoundedWindow window) {
  try {
    fnRunner.onTimer(
        timer.getTimerId(),
        timer.getTimerFamilyId(),
        key,
        window,
        timer.getTimestamp(),
        timer.getOutputTimestamp(),
        timer.getDomain());
  } catch (Exception e) {
    throw UserCodeException.wrap(e);
  }
}
 

@Override
public void processElement(WindowedValue<InputT> element) {
  try {
    Iterable<WindowedValue<InputT>> unprocessed = fnRunner.processElementInReadyWindows(element);
    unprocessedElements.addAll(unprocessed);
  } catch (Exception e) {
    throw UserCodeException.wrap(e);
  }
}
 

/** Outputs the given element to the specified set of consumers wrapping any exceptions. */
private <T> void outputTo(
    Collection<FnDataReceiver<WindowedValue<T>>> consumers, WindowedValue<T> output) {
  if (currentWatermarkEstimator instanceof TimestampObservingWatermarkEstimator) {
    ((TimestampObservingWatermarkEstimator) currentWatermarkEstimator)
        .observeTimestamp(output.getTimestamp());
  }
  try {
    for (FnDataReceiver<WindowedValue<T>> consumer : consumers) {
      consumer.accept(output);
    }
  } catch (Throwable t) {
    throw UserCodeException.wrap(t);
  }
}
 

/**
 * Demonstrates that attempting to output an element before the timestamp of the current element
 * plus the value of {@link DoFn#getAllowedTimestampSkew()} throws, but between that value and the
 * current timestamp succeeds.
 */
@Test
public void testSkew() {
  SkewingDoFn fn = new SkewingDoFn(Duration.standardMinutes(10L));
  DoFnRunner<Duration, Duration> runner =
      new SimpleDoFnRunner<>(
          null,
          fn,
          NullSideInputReader.empty(),
          new ListOutputManager(),
          new TupleTag<>(),
          Collections.emptyList(),
          mockStepContext,
          null,
          Collections.emptyMap(),
          WindowingStrategy.of(new GlobalWindows()),
          DoFnSchemaInformation.create(),
          Collections.emptyMap());

  runner.startBundle();
  // Outputting between "now" and "now - allowed skew" succeeds.
  runner.processElement(
      WindowedValue.timestampedValueInGlobalWindow(Duration.standardMinutes(5L), new Instant(0)));
  thrown.expect(UserCodeException.class);
  thrown.expectCause(isA(IllegalArgumentException.class));
  thrown.expectMessage("must be no earlier");
  thrown.expectMessage(
      String.format("timestamp of the current input (%s)", new Instant(0).toString()));
  thrown.expectMessage(
      String.format(
          "the allowed skew (%s)",
          PeriodFormat.getDefault().print(Duration.standardMinutes(10L).toPeriod())));
  // Outputting before "now - allowed skew" fails.
  runner.processElement(
      WindowedValue.timestampedValueInGlobalWindow(Duration.standardHours(1L), new Instant(0)));
}
 

public SerializableArrayViaCoder(Coder<T> elementCoder, T[] value) {
  this.coder = ListCoder.of(elementCoder);
  this.value = value;
  try {
    this.encodedValue = CoderUtils.encodeToByteArray(coder, Arrays.asList(value));
  } catch (CoderException exc) {
    throw UserCodeException.wrap(exc);
  }
}
 

@Test
public void testProcessElementExceptionWithReturn() throws Exception {
  thrown.expect(UserCodeException.class);
  thrown.expectMessage("bogus");
  DoFnInvokers.invokerFor(
          new DoFn<Integer, Integer>() {
            @ProcessElement
            public ProcessContinuation processElement(
                @SuppressWarnings("unused") ProcessContext c,
                RestrictionTracker<SomeRestriction, Void> tracker) {
              throw new IllegalArgumentException("bogus");
            }

            @GetInitialRestriction
            public SomeRestriction getInitialRestriction(@Element Integer element) {
              return null;
            }

            @NewTracker
            public SomeRestrictionTracker newTracker(@Restriction SomeRestriction restriction) {
              return null;
            }
          })
      .invokeProcessElement(
          new FakeArgumentProvider<Integer, Integer>() {
            @Override
            public DoFn.ProcessContext processContext(DoFn<Integer, Integer> doFn) {
              return null; // will not be touched
            }

            @Override
            public RestrictionTracker<?, ?> restrictionTracker() {
              return null; // will not be touched
            }
          });
}