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

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

import com.codahale.metrics.MetricRegistry;
import com.google.common.collect.Sets;
import com.streamsets.datacollector.el.RuntimeEL;
import com.streamsets.datacollector.main.RuntimeInfo;
import com.streamsets.datacollector.main.RuntimeModule;
import com.streamsets.datacollector.main.StandaloneRuntimeInfo;
import com.streamsets.pipeline.api.ConfigDef;
import com.streamsets.pipeline.api.impl.Utils;

import java.io.IOException;
import java.lang.reflect.Field;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.Set;

public abstract class ProtoRunner {

  protected Status status;

  enum Status {
    CREATED,
    INITIALIZED,
    DESTROYED
    ;

    public boolean isOneOf(Status ...statuses) {
      if(statuses == null) {
        return false;
      }

      for(Status t : statuses) {
        if(this == t) {
          return true;
        }
      }

      return false;
    }
  }

  static {
    RuntimeInfo runtimeInfo = new StandaloneRuntimeInfo(
        RuntimeInfo.SDC_PRODUCT,
        RuntimeModule.SDC_PROPERTY_PREFIX,
        new MetricRegistry(),
        Collections.singletonList(ProtoRunner.class.getClassLoader())
    );
    try {
      RuntimeEL.loadRuntimeConfiguration(runtimeInfo);
    } catch (IOException ex) {
      throw new RuntimeException(ex);
    }
  }

  void ensureStatus(Status ...statuses) {
    Utils.checkState(this.status.isOneOf(statuses), Utils.format("Current status '{}', expected one of '{}'", this.status, statuses));
  }

  // Configuration

  private boolean isConfigurationActive(ConfigDef configDef, Map configuration) {
    String dependsOn = configDef.dependsOn();
    if (!dependsOn.isEmpty()) {
      Object dependsOnValue = configuration.get(dependsOn);
      if (dependsOnValue != null) {
        String valueStr = dependsOnValue.toString();
        for (String trigger : configDef.triggeredByValue()) {
          if (valueStr.equals(trigger)) {
            return true;
          }
        }
        return false;
      }
      return false;
    }
    return true;
  }

  private Set getStageConfigurationFields(Class klass) throws Exception {
    Set names = new HashSet<>();
    for (Field field : klass.getFields()) {
      if (field.isAnnotationPresent(ConfigDef.class)) {
        names.add(field.getName());
      }
    }
    return names;
  }

  private Set filterNonActiveConfigurationsFromMissing(Object stage, Map configuration, Set missingConfigs) {
    missingConfigs = new HashSet<>(missingConfigs);
    Iterator it = missingConfigs.iterator();
    while (it.hasNext()) {
      String name = it.next();
      try {
        Field field = stage.getClass().getField(name);
        ConfigDef annotation = field.getAnnotation(ConfigDef.class);
        if (!annotation.required() || !isConfigurationActive(annotation, configuration)) {
          it.remove();
        }
      } catch (Exception ex) {
        throw new RuntimeException(ex);
      }
    }
    return missingConfigs;
  }

  @SuppressWarnings("unchecked")
  protected void configureObject(Object stage, Map configuration) {
    try {
      Set fields = getStageConfigurationFields(stage.getClass());
      Set configs = configuration.keySet();
      if (!fields.equals(configs)) {
        Set missingConfigs = Sets.difference(fields, configs);
        Set extraConfigs = Sets.difference(configs, fields);

        missingConfigs = filterNonActiveConfigurationsFromMissing(stage, configuration, missingConfigs);
        if (missingConfigs.size() + extraConfigs.size() > 0) { //x
          throw new RuntimeException(Utils.format(
              "Invalid stage configuration for '{}', Missing configurations '{}' and invalid configurations '{}'",
              stage.getClass().getName(), missingConfigs, extraConfigs));
        }
      }
      for (Field field : stage.getClass().getFields()) {
        if (field.isAnnotationPresent(ConfigDef.class)) {
          ConfigDef configDef = field.getAnnotation(ConfigDef.class);
          if (isConfigurationActive(configDef, configuration)) {
            if ( configDef.type() != ConfigDef.Type.MAP) {
              field.set(stage, configuration.get(field.getName()));
            } else {
              //we need to handle special case of List of Map elements with key/value entries
              Object value = configuration.get(field.getName());
              if (value != null && value instanceof List) {
                Map map = new HashMap();
                for (Map element : (List) value) {
                  if (!element.containsKey("key") || !element.containsKey("value")) {
                    throw new RuntimeException(Utils.format("Invalid stage configuration for '{}' Map as list must have" +
                                                            " a List of Maps all with 'key' and 'value' entries",
                                                            field.getName()));
                  }
                  String k = (String) element.get("key");
                  String v = (String) element.get("value");
                  map.put(k, v);
                }
                value = map;
              }
              field.set(stage, value);
            }
          }
        }
      }
    } catch (Exception ex) {
      if (ex instanceof RuntimeException) {
        throw (RuntimeException) ex;
      }
      throw new RuntimeException(ex);
    }
  }
}