org.nuiton.topia.persistence.script.SqlScriptConsumer Maven / Gradle / Ivy
Show all versions of topia-extension Show documentation
package org.nuiton.topia.persistence.script;
/*-
* #%L
* ObServe Toolkit :: ToPIA Extension
* %%
* Copyright (C) 2017 - 2018 IRD, Ultreia.io
* %%
* This program is free software: you can redistribute it and/or modify
* it under the terms of the GNU General Public License as
* published by the Free Software Foundation, either version 3 of the
* License, or (at your option) any later version.
*
* This program is distributed in the hope that it will be useful,
* but WITHOUT ANY WARRANTY; without even the implied warranty of
* MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
* GNU General Public License for more details.
*
* You should have received a copy of the GNU General Public
* License along with this program. If not, see
* .
* #L%
*/
import com.google.common.collect.ImmutableSet;
import io.ultreia.java4all.util.SingletonSupplier;
import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.nuiton.topia.persistence.TopiaException;
import org.nuiton.topia.persistence.support.TopiaSqlWork;
import javax.sql.rowset.serial.SerialBlob;
import java.io.Closeable;
import java.io.IOException;
import java.net.URL;
import java.nio.charset.Charset;
import java.nio.file.Path;
import java.sql.Connection;
import java.sql.PreparedStatement;
import java.sql.SQLException;
import java.sql.Statement;
import java.util.ArrayList;
import java.util.Collection;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.function.Consumer;
/**
* To consume an incoming sql script into a topia persistence database.
*
* Created by tchemit on 10/05/2018.
*
* @author Tony Chemit - [email protected]
*/
public class SqlScriptConsumer implements TopiaSqlWork, Closeable, Consumer {
private static final Log log = LogFactory.getLog(SqlScriptConsumer.class);
private final SingletonSupplier source;
private final ImmutableSet blobsContainers;
private final Integer batchSize;
private final List statements;
private long statementCount;
private SqlScriptConsumer(SingletonSupplier source, Integer batchSize, ImmutableSet blobsContainers) {
this.source = source;
this.batchSize = batchSize;
this.statements = new ArrayList<>(batchSize == null ? 1 : batchSize);
this.blobsContainers = blobsContainers;
}
public static SqlScriptConsumer of(String source) {
return builder(source).build();
}
public static SqlScriptConsumer of(byte[] source) {
return builder(source).build();
}
public static SqlScriptConsumer of(Path source) {
return builder(source).build();
}
public static SqlScriptConsumer of(URL source) {
return builder(source).build();
}
public static SqlScriptConsumer of(TopiaSqlScript source) {
return builder(source).build();
}
public static Builder builder(String source) {
return new Builder(SqlScriptReader.builder(source));
}
public static Builder builder(byte[] source) {
return new Builder(SqlScriptReader.builder(source));
}
public static Builder builder(Path source) {
return new Builder(SqlScriptReader.builder(source));
}
public static Builder builder(URL source) {
return new Builder(SqlScriptReader.builder(source));
}
public static Builder builder(TopiaSqlScript source) {
Builder builder = builder(source.getLocation());
if (source.withBlobs()) {
builder.blobs(source.getBlobsContainers());
}
return builder;
}
public static Builder builder(SqlScriptReader source) {
return new Builder(SqlScriptReader.builder(source));
}
@Override
public void close() throws IOException {
if (source.withValue()) {
source.get().close();
}
}
@Override
public void execute(Connection connection) throws SQLException {
boolean autoCommit = connection.getAutoCommit();
try {
connection.setAutoCommit(false);
try (Statement statement = connection.createStatement()) {
if (batchSize == null) {
executeOneByOne(statement);
} else {
executeByBatch(statement);
}
}
if (!blobsContainers.isEmpty()) {
importBlobs(connection);
}
} finally {
connection.setAutoCommit(autoCommit);
}
}
@Override
public void accept(Connection connection) {
try {
execute(connection);
} catch (Exception e) {
throw new TopiaException(e);
}
}
public long getStatementCount() {
return statementCount;
}
private void importBlobs(Connection connection) throws SQLException {
for (TopiaBlobsContainer blobsContainer : blobsContainers) {
String tableName = blobsContainer.getTableName();
String columnName = blobsContainer.getColumnName();
int batchSize = 0;
String sql = String.format("UPDATE %s SET %s = ? WHERE topiaId= ?", tableName, columnName);
log.debug(sql);
try (PreparedStatement statement = connection.prepareStatement(sql)) {
for (Map.Entry containerEntry : blobsContainer.getBlobsById().entrySet()) {
String topiaId = containerEntry.getKey();
byte[] content = containerEntry.getValue();
statement.clearParameters();
statement.setBlob(1, new SerialBlob(content));
statement.setString(2, topiaId);
statement.addBatch();
batchSize++;
statementCount++;
if (this.batchSize == null || batchSize % this.batchSize == 0) {
flush(statement);
}
}
flush(statement);
}
}
}
private void executeByBatch(Statement statement) throws SQLException {
for (String sqlStatement : source.get()) {
statementCount++;
try {
statement.addBatch(sqlStatement);
statements.add(sqlStatement);
} catch (Exception e) {
log.error(String.format("Can't add sql statement (%d) in batch: %s", statementCount, sqlStatement), e);
throw e;
}
if (statementCount > 0 && statementCount % batchSize == 0) {
flush(statement);
}
}
flush(statement);
}
private void executeOneByOne(Statement statement) throws SQLException {
for (String sqlStatement : source.get()) {
statementCount++;
try {
statement.execute(sqlStatement);
} catch (Exception e) {
log.error(String.format("Can't execute sql statement (%d): %s", statementCount, sqlStatement), e);
throw e;
}
}
}
private void flush(Statement statement) throws SQLException {
try {
statement.executeBatch();
} catch (Exception e) {
log.error(String.format("Can't execute sql statements (%d): %s", statementCount, statements), e);
throw e;
} finally {
statements.clear();
statement.clearBatch();
}
}
public static class Builder {
private final SqlScriptReader.Builder source;
private final ImmutableSet.Builder blobsContainers;
private Integer batchSize;
public Builder(SqlScriptReader.Builder source) {
this.source = Objects.requireNonNull(source);
this.blobsContainers = ImmutableSet.builder();
}
public Builder batchSize(int batchSize) {
this.batchSize = batchSize;
return this;
}
public Builder keepCommentLine() {
source.keepCommentLine();
return this;
}
public Builder keepEmptyLine() {
source.keepEmptyLine();
return this;
}
public Builder encoding(Charset encoding) {
source.encoding(Objects.requireNonNull(encoding));
return this;
}
public Builder blob(TopiaBlobsContainer blobsContainers) {
this.blobsContainers.add(Objects.requireNonNull(blobsContainers));
return this;
}
public Builder blobs(Collection blobsContainers) {
this.blobsContainers.addAll(Objects.requireNonNull(blobsContainers));
return this;
}
public SqlScriptConsumer build() {
return new SqlScriptConsumer(SingletonSupplier.of(source::build), batchSize, blobsContainers.build());
}
}
}