kern/src/de/superx/spring/batch/tasklet/ExecuteDbtTasklet.java

package de.superx.spring.batch.tasklet;

import javax.sql.DataSource;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.batch.core.ExitStatus;
import org.springframework.batch.core.StepContribution;
import org.springframework.batch.core.StepExecution;
import org.springframework.batch.core.StepExecutionListener;
import org.springframework.batch.core.scope.context.ChunkContext;
import org.springframework.batch.core.step.tasklet.Tasklet;
import org.springframework.batch.repeat.RepeatStatus;
import de.superx.conf.ApplicationConfig;
import de.superx.conf.ApplicationConfig.Key;
import de.superx.dbt.DbtWrapper;
import de.superx.dbt.DbtExecutionState;
import de.superx.dbt.DbtAction;
import de.superx.dbt.DbtPermissionException;
import de.superx.dbt.DbtRunRequest;
import de.superx.dbt.DbtRunResult;
import de.superx.dbt.DbtUtils;
import de.superx.dbt.UnsupportedPythonVersionException;

public class ExecuteDbtTasklet implements Tasklet, StepExecutionListener, InternalProgress {
    static Logger logger = LoggerFactory.getLogger(ExecuteDbtTasklet.class);
    
    private final static long DBT_TIMEOUT_MILLISECONDS = ((Integer) ApplicationConfig.getWithType(Key.dbt_job_execution_timeout_seconds)).longValue() * 1000;
    
    private final String dbtProject;
    private final String dbtSelect;
    private final DataSource dataSource;
    private final boolean dbtTest;
    private int totalStepsExecuted = 0;
    private int totalSteps = 0;

    private DbtWrapper dbtWrapper;

    public ExecuteDbtTasklet(DbtWrapper dbtWrapper, DataSource dataSource, String dbtProject, String dbtSelect, boolean dbtTest) throws UnsupportedPythonVersionException {
        this.dbtWrapper = dbtWrapper;
        this.dbtProject = dbtProject;
        this.dataSource = dataSource;
        this.dbtSelect = dbtSelect;
        this.dbtTest = dbtTest;
        DbtUtils.checkPythonCompatability();
    }

    @Override
    public void beforeStep(StepExecution stepExecution) {
        try {
            DbtUtils.checkPythonCompatability();
            DbtUtils.checkDbtSchemaPermissions(dataSource);
        } catch (UnsupportedPythonVersionException | DbtPermissionException e) {
            throw new RuntimeException(e);
        }
    }

    @Override
    public RepeatStatus execute(StepContribution contribution, ChunkContext chunkContext) throws Exception {
        try {
            // TODO fix, maybe own tasklet for dbt test?
            DbtAction action = (this.dbtTest) ? DbtAction.TEST : DbtAction.BUILD;

            DbtRunRequest request = new DbtRunRequest(action, dbtProject, dbtSelect,
                    DBT_TIMEOUT_MILLISECONDS,
                    progressStep -> totalStepsExecuted = progressStep.intValue());

            DbtRunResult result = dbtWrapper.execute(request);
        
            if (result.getState() != DbtExecutionState.SUCCESS) {
                String errorMessage = String.join(System.lineSeparator(), result.getErrorMessages());
                contribution.setExitStatus(ExitStatus.FAILED.addExitDescription(errorMessage));
            } else {
                // generate dbt documentation only if dbt run was successful
                request = new DbtRunRequest(DbtAction.GENERATE_DOCS, dbtProject, dbtSelect);
                result = dbtWrapper.execute(request);
                if(result.getState() != DbtExecutionState.SUCCESS) {
                    String errorMessage = String.join(System.lineSeparator(), result.getErrorMessages());
                    contribution.setExitStatus(ExitStatus.FAILED.addExitDescription(errorMessage));
                }
            }

        } catch (Exception e) {
            String errorMsg = e.getMessage();
            Throwable cause = e.getCause();
            if (cause != null) {
                errorMsg = cause.getMessage();
            } else {
                logger.warn("No cause given for exception!");
            }
            ExitStatus exitStatus = ExitStatus.FAILED;
            exitStatus = exitStatus.addExitDescription(errorMsg);
            contribution.setExitStatus(exitStatus);
        }

        this.totalStepsExecuted = this.totalSteps + 1;
        return RepeatStatus.FINISHED;
    }

    @Override
    public ExitStatus afterStep(StepExecution stepExecution) {
        return null;
    }
    
    @Override
    public int getInternalStepsExecutedCount() {
        return totalStepsExecuted;
    }

    @Override
    public int getInternalTotalStepsCount() {
        DbtAction action = (this.dbtTest) ? DbtAction.NODE_COUNT_TEST : DbtAction.NODE_COUNT_BUILD;
        DbtRunRequest request = new DbtRunRequest(action, dbtProject, dbtSelect, DBT_TIMEOUT_MILLISECONDS);
        DbtRunResult result = dbtWrapper.execute(request);
        this.totalSteps = result.getNodeCount();
        return result.getNodeCount();
    }

}