io.kestra.plugin.gcp.dataproc.batches.SparkSubmit Maven / Gradle / Ivy

Go to download
package io.kestra.plugin.gcp.dataproc.batches;

import com.google.cloud.dataproc.v1.Batch;
import com.google.cloud.dataproc.v1.SparkBatch;
import io.kestra.core.exceptions.IllegalVariableEvaluationException;
import io.kestra.core.models.annotations.Example;
import io.kestra.core.models.annotations.Plugin;
import io.kestra.core.models.annotations.PluginProperty;
import io.kestra.core.runners.RunContext;
import io.swagger.v3.oas.annotations.media.Schema;
import lombok.EqualsAndHashCode;
import lombok.Getter;
import lombok.NoArgsConstructor;
import lombok.ToString;
import lombok.experimental.SuperBuilder;

import jakarta.validation.constraints.NotNull;

@SuperBuilder
@ToString
@EqualsAndHashCode
@Getter
@NoArgsConstructor
@Schema(
    title = "Submit an [Apache Spark](https://spark.apache.org/) batch workload."
)
@Plugin(
    examples = @Example(
        full = true,
        code = """
            id: gcp_dataproc_spark_submit
            namespace: company.team
            tasks:
              - id: spark_submit
                type: io.kestra.plugin.gcp.dataproc.batches.SparkSubmit
                jarFileUris: 
                  - 'gs://spark-jobs-kestra/spark-examples.jar'
                mainClass: org.apache.spark.examples.SparkPi
                args:
                  - 1000
                name: test-spark
                region: europe-west3
            """
    )
)
public class SparkSubmit extends AbstractSparkSubmit {
    @Schema(
        title = "The name of the driver main class.",
        description = "The jar file that contains the class must be in the classpath or specified in `jarFileUris`"
    )
    @PluginProperty(dynamic = true)
    @NotNull
    private String mainClass;

    @Override
    protected void buildBatch(Batch.Builder builder, RunContext runContext) throws IllegalVariableEvaluationException {
        SparkBatch.Builder sparkBuilder = SparkBatch.newBuilder();

        if (this.mainClass != null) {
            sparkBuilder.setMainClass(runContext.render(this.mainClass));
        }

        if (this.jarFileUris != null) {
            sparkBuilder.addAllJarFileUris(runContext.render(this.jarFileUris));
        }

        if (this.fileUris != null) {
            sparkBuilder.addAllFileUris(runContext.render(this.fileUris));
        }

        if (this.archiveUris != null) {
            sparkBuilder.addAllArchiveUris(runContext.render(this.archiveUris));
        }

        if (this.args != null) {
            sparkBuilder.addAllArgs(runContext.render(this.args));
        }

        builder.setSparkBatch(sparkBuilder.build());
    }
}