code: /*
   * Copyright (c) 2022 Airbyte, Inc., all rights reserved.
   */
  
  package io.airbyte.workers.general;
  
  import com.fasterxml.jackson.databind.JsonNode;
  import com.google.common.base.Strings;
  import com.google.common.collect.ImmutableMap;
  import io.airbyte.commons.io.LineGobbler;
  import io.airbyte.commons.logging.LoggingHelper.Color;
  import io.airbyte.commons.logging.MdcScope;
  import io.airbyte.commons.logging.MdcScope.Builder;
  import io.airbyte.commons.resources.MoreResources;
  import io.airbyte.config.OperatorDbt;
  import io.airbyte.config.ResourceRequirements;
  import io.airbyte.workers.WorkerConfigs;
  import io.airbyte.workers.WorkerUtils;
  import io.airbyte.workers.exception.WorkerException;
  import io.airbyte.workers.normalization.NormalizationRunner;
  import io.airbyte.workers.process.AirbyteIntegrationLauncher;
  import io.airbyte.workers.process.ProcessFactory;
  import java.nio.file.Path;
  import java.util.ArrayList;
  import java.util.Collections;
  import java.util.List;
  import java.util.Map;
  import java.util.concurrent.TimeUnit;
  import org.apache.tools.ant.types.Commandline;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  public class DbtTransformationRunner implements AutoCloseable {
  
    private static final Logger LOGGER = LoggerFactory.getLogger(DbtTransformationRunner.class);
    private static final String DBT_ENTRYPOINT_SH = "entrypoint.sh";
    private static final MdcScope.Builder CONTAINER_LOG_MDC_BUILDER = new Builder()
        .setLogPrefix("dbt")
        .setPrefixColor(Color.PURPLE_BACKGROUND);
  
    private final WorkerConfigs workerConfigs;
    private final ProcessFactory processFactory;
    private final NormalizationRunner normalizationRunner;
    private Process process = null;
  
    public DbtTransformationRunner(final WorkerConfigs workerConfigs,
                                   final ProcessFactory processFactory,
                                   final NormalizationRunner normalizationRunner) {
      this.workerConfigs = workerConfigs;
      this.processFactory = processFactory;
      this.normalizationRunner = normalizationRunner;
    }
  
    public void start() throws Exception {
      normalizationRunner.start();
    }
  
    /**
     * The docker image used by the DbtTransformationRunner is provided by the User, so we can't ensure
     * to have the right python, dbt, dependencies etc software installed to successfully run our
     * transform-config scripts (to translate Airbyte Catalogs into Dbt profiles file). Thus, we depend
     * on the NormalizationRunner to configure the dbt project with the appropriate destination settings
     * and pull the custom git repository into the workspace.
     * <p>
     * Once the workspace folder/files is setup to run, we invoke the custom transformation command as
     * provided by the user to execute whatever extra transformation has been implemented.
     */
    public boolean run(final String jobId,
                       final int attempt,
                       final Path jobRoot,
                       final JsonNode config,
                       final ResourceRequirements resourceRequirements,
                       final OperatorDbt dbtConfig)
        throws Exception {
      if (!normalizationRunner.configureDbt(jobId, attempt, jobRoot, config, resourceRequirements, dbtConfig)) {
        return false;
      }
      return transform(jobId, attempt, jobRoot, config, resourceRequirements, dbtConfig);
    }
  
    public boolean transform(final String jobId,
                             final int attempt,
                             final Path jobRoot,
                             final JsonNode config,
                             final ResourceRequirements resourceRequirements,
                             final OperatorDbt dbtConfig)
        throws Exception {
      try {
        final Map<String, String> files = ImmutableMap.of(
            DBT_ENTRYPOINT_SH, MoreResources.readResource("dbt_transformation_entrypoint.sh"),
            "sshtunneling.sh", MoreResources.readResource("sshtunneling.sh"));
        final List<String> dbtArguments = new ArrayList<>();
        dbtArguments.add(DBT_ENTRYPOINT_SH);
        if (Strings.isNullOrEmpty(dbtConfig.getDbtArguments())) {
          throw new WorkerException("Dbt Arguments are required");
        }
        Collections.addAll(dbtArguments, Commandline.translateCommandline(dbtConfig.getDbtArguments()));
        process =
            processFactory.create(
                AirbyteIntegrationLauncher.CUSTOM_STEP,
                jobId,
                attempt,
                jobRoot,
                dbtConfig.getDockerImage(),
                false,
                files,
                "/bin/bash",
                resourceRequirements,
                Map.of(AirbyteIntegrationLauncher.JOB_TYPE, AirbyteIntegrationLauncher.SYNC_JOB, AirbyteIntegrationLauncher.SYNC_STEP,
                    AirbyteIntegrationLauncher.CUSTOM_STEP),
                Collections.emptyMap(),
                Collections.emptyMap(),
                dbtArguments.toArray(new String[0]));
        LineGobbler.gobble(process.getInputStream(), LOGGER::info, CONTAINER_LOG_MDC_BUILDER);
        LineGobbler.gobble(process.getErrorStream(), LOGGER::error, CONTAINER_LOG_MDC_BUILDER);
  
        WorkerUtils.wait(process);
  
        return process.exitValue() == 0;
      } catch (final Exception e) {
        // make sure we kill the process on failure to avoid zombies.
        if (process != null) {
          WorkerUtils.cancelProcess(process);
        }
        throw e;
      }
    }
  
    @Override
    public void close() throws Exception {
      normalizationRunner.close();
  
      if (process == null) {
        return;
      }
  
      LOGGER.debug("Closing dbt transformation process");
      WorkerUtils.gentleClose(workerConfigs, process, 1, TimeUnit.MINUTES);
      if (process.isAlive() || process.exitValue() != 0) {
        throw new WorkerException("Dbt transformation process wasn't successful");
      }
    }
  
  }
  