io.kestra.plugin.jdbc.AbstractJdbcQuery Maven / Gradle / Ivy
Go to download
Show more of this group Show more artifacts with this name
Show all versions of plugin-jdbc Show documentation
Show all versions of plugin-jdbc Show documentation
Abstract JDBC plugin for Kestra
package io.kestra.plugin.jdbc;
import com.fasterxml.jackson.databind.ObjectMapper;
import io.kestra.core.models.executions.metrics.Counter;
import io.kestra.core.models.tasks.Task;
import io.kestra.core.runners.RunContext;
import io.kestra.core.serializers.JacksonMapper;
import io.kestra.core.utils.Rethrow;
import io.swagger.v3.oas.annotations.media.Schema;
import lombok.*;
import lombok.experimental.SuperBuilder;
import org.slf4j.Logger;
import java.io.BufferedWriter;
import java.io.File;
import java.io.FileWriter;
import java.io.IOException;
import java.net.URI;
import java.sql.Connection;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.sql.Statement;
import java.time.ZoneId;
import java.util.*;
import java.util.function.Consumer;
@SuperBuilder
@ToString
@EqualsAndHashCode
@Getter
@NoArgsConstructor
public abstract class AbstractJdbcQuery extends Task implements JdbcQueryInterface {
private String url;
private String username;
private String password;
private String timeZoneId;
private String sql;
@Builder.Default
private boolean store = false;
@Builder.Default
private boolean fetchOne = false;
@Builder.Default
private boolean fetch = false;
@Builder.Default
protected Integer fetchSize = 10000;
@Builder.Default
@Getter(AccessLevel.NONE)
protected transient Map additionalVars = new HashMap<>();
private static final ObjectMapper MAPPER = JacksonMapper.ofIon();
protected abstract AbstractCellConverter getCellConverter(ZoneId zoneId);
protected Statement createStatement(Connection conn) throws SQLException {
return conn.createStatement(
ResultSet.TYPE_FORWARD_ONLY,
ResultSet.CONCUR_READ_ONLY
);
}
public AbstractJdbcQuery.Output run(RunContext runContext) throws Exception {
Logger logger = runContext.logger();
AbstractCellConverter cellConverter = getCellConverter(this.zoneId());
try (
Connection conn = this.connection(runContext);
Statement stmt = this.createStatement(conn);
) {
if (this instanceof AutoCommitInterface) {
if (this.store) {
conn.setAutoCommit(false);
} else {
conn.setAutoCommit(((AutoCommitInterface) this).getAutoCommit());
}
}
stmt.setFetchSize(this.getFetchSize());
String sql = runContext.render(this.sql, this.additionalVars);
logger.debug("Starting query: {}", sql);
boolean isResult = stmt.execute(sql);
try(ResultSet rs = stmt.getResultSet()) {
Output.OutputBuilder, ?> output = Output.builder();
long size = 0;
if (isResult) {
if (this.fetchOne) {
output
.row(fetchResult(rs, cellConverter, conn))
.size(1L);
size = 1;
} else if (this.store) {
File tempFile = runContext.tempFile(".ion").toFile();
BufferedWriter fileWriter = new BufferedWriter(new FileWriter(tempFile));
size = fetchToFile(stmt, rs, fileWriter, cellConverter, conn);
fileWriter.flush();
fileWriter.close();
output
.uri(runContext.putTempFile(tempFile))
.size(size);
} else if (this.fetch) {
List
© 2015 - 2025 Weber Informatics LLC | Privacy Policy