com.streamsets.pipeline.sdk.ProcessorRunner Maven / Gradle / Ivy

Go to download
/*
 * Copyright (c) 2021 StreamSets Inc.
 */
package com.streamsets.pipeline.sdk;

import com.streamsets.datacollector.main.RuntimeInfo;
import com.streamsets.datacollector.runner.BatchImpl;
import com.streamsets.pipeline.api.BatchMaker;
import com.streamsets.pipeline.api.DeliveryGuarantee;
import com.streamsets.pipeline.api.ExecutionMode;
import com.streamsets.pipeline.api.OnRecordError;
import com.streamsets.pipeline.api.Processor;
import com.streamsets.pipeline.api.Record;
import com.streamsets.pipeline.api.Source;
import com.streamsets.pipeline.api.StageException;
import com.streamsets.pipeline.api.StageType;
import com.streamsets.pipeline.api.impl.Utils;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.List;
import java.util.Map;

public class ProcessorRunner extends StageRunner {
  private static final Logger LOG = LoggerFactory.getLogger(Processor.class);

  public ProcessorRunner(
    Class processorClass,
    Processor processor,
    Map configuration,
    List outputLanes,
    boolean isPreview,
    OnRecordError onRecordError,
    Map constants,
    Map stageSdcConf,
    ExecutionMode executionMode,
    DeliveryGuarantee deliveryGuarantee,
    String resourcesDir,
    RuntimeInfo runtimeInfo,
    List services
  ) {
    super(
      processorClass,
      processor,
      StageType.PROCESSOR,
      configuration,
      outputLanes,
      isPreview,
      onRecordError,
      constants,
      stageSdcConf,
      executionMode,
      deliveryGuarantee,
      resourcesDir,
      runtimeInfo,
      services
    );
  }

  public ProcessorRunner(
    Class processorClass,
    Map configuration,
    List outputLanes,
    boolean isPreview,
    OnRecordError onRecordError,
    Map constants,
    Map stageSdcConf,
    ExecutionMode executionMode,
    DeliveryGuarantee deliveryGuarantee,
    String resourcesDir,
    RuntimeInfo runtimeInfo,
    List services
  ) {
    super(
      processorClass,
      StageType.PROCESSOR,
      configuration,
      outputLanes,
      isPreview,
      onRecordError,
      constants,
      stageSdcConf,
      executionMode,
      deliveryGuarantee,
      resourcesDir,
      runtimeInfo,
      services
    );
  }

  public Output runProcess(List inputRecords) throws StageException {
    try {
      LOG.debug("Stage '{}' process starts", getInfo().getInstanceName());
      ensureStatus(Status.INITIALIZED);
      BatchImpl batch = new BatchImpl(getInfo().getInstanceName(), "sdk", "sourceOffset", inputRecords);
      BatchMakerImpl batchMaker = new BatchMakerImpl(((Source.Context) getContext()).getOutputLanes());
      getStage().process(batch, batchMaker);
      return new Output(Source.POLL_SOURCE_OFFSET_KEY, "sdk:sourceOffset", batchMaker.getOutput());
    } finally {
      LOG.debug("Stage '{}' process ends", getInfo().getInstanceName());
    }
  }

  public static class Builder extends StageRunner.Builder {

    public Builder(Class processorClass,  Processor processor) {
      super((Class)processorClass, processor);
    }

    @SuppressWarnings("unchecked")
    public Builder(Class processorClass) {
      super((Class) processorClass);
    }

    @Override
    public ProcessorRunner build() {
      Utils.checkState(!outputLanes.isEmpty(), "A Processor must have at least one output stream");
      if (stage != null) {
        return new ProcessorRunner(
          stageClass,
          stage,
          configs,
          outputLanes,
          isPreview,
          onRecordError,
          constants,
          stageSdcConf,
          executionMode,
          deliveryGuarantee,
          resourcesDir,
          runtimeInfo,
          services
        );
      } else {
        return new ProcessorRunner(
          stageClass,
          configs,
          outputLanes,
          isPreview,
          onRecordError,
          constants,
          stageSdcConf,
          executionMode,
          deliveryGuarantee,
          resourcesDir,
          runtimeInfo,
          services
        );
      }
    }

  }

  public static BatchMaker createTestBatchMaker(String... outputLanes) {
    return StageRunner.createTestBatchMaker(outputLanes);
  }

  public static Output getOutput(BatchMaker batchMaker) {
    return StageRunner.getOutput(batchMaker);
  }

}