ru.yandex.clickhouse.Writer Maven / Gradle / Ivy
package ru.yandex.clickhouse;
import static ru.yandex.clickhouse.domain.ClickHouseFormat.Native;
import static ru.yandex.clickhouse.domain.ClickHouseFormat.RowBinary;
import static ru.yandex.clickhouse.domain.ClickHouseFormat.TabSeparated;
import java.io.File;
import java.io.FileInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.sql.SQLException;
import java.util.Objects;
import org.apache.http.HttpEntity;
import org.apache.http.entity.InputStreamEntity;
import ru.yandex.clickhouse.domain.ClickHouseCompression;
import ru.yandex.clickhouse.domain.ClickHouseFormat;
import ru.yandex.clickhouse.settings.ClickHouseQueryParam;
import ru.yandex.clickhouse.util.ClickHouseStreamCallback;
import ru.yandex.clickhouse.util.ClickHouseStreamHttpEntity;
public class Writer extends ConfigurableApi {
private ClickHouseFormat format = TabSeparated;
private ClickHouseCompression compression = null;
private String table = null;
private String sql = null;
private InputStreamProvider streamProvider = null;
Writer(ClickHouseStatementImpl statement) {
super(statement);
dataCompression(ClickHouseCompression.none);
}
/**
* Specifies format for further insert of data via send().
*
* @param format
* the format of the data to upload
* @return this writer instance
*/
public Writer format(ClickHouseFormat format) {
if (null == format) {
throw new NullPointerException("Format can not be null");
}
this.format = format;
return this;
}
/**
* Set table name for data insertion.
*
* @param table
* name of the table to upload the data to
* @return this writer instance
*/
public Writer table(String table) {
this.sql = null;
this.table = table;
return this;
}
/**
* Set SQL for data insertion.
*
* @param sql
* in a form "INSERT INTO table_name [(X,Y,Z)] VALUES "
* @return this writer instance
*/
public Writer sql(String sql) {
this.sql = sql;
this.table = null;
return this;
}
/**
* Specifies data input stream.
*
* @param stream
* a stream providing the data to upload
* @return this writer instance
*/
public Writer data(InputStream stream) {
streamProvider = new HoldingInputProvider(stream);
return this;
}
/**
* Specifies data input stream, and the format to use.
*
* @param stream
* a stream providing the data to upload
* @param format
* the format of the data to upload
* @return this writer instance
*/
public Writer data(InputStream stream, ClickHouseFormat format) {
return format(format).data(stream);
}
/**
* Shortcut method for specifying a file as an input.
*
* @param input
* the file to upload
* @return this writer instance
*/
public Writer data(File input) {
streamProvider = new FileInputProvider(input);
return this;
}
public Writer data(InputStream stream, ClickHouseFormat format, ClickHouseCompression compression) {
return dataCompression(compression).format(format).data(stream);
}
public Writer data(File input, ClickHouseFormat format, ClickHouseCompression compression) {
return dataCompression(compression).format(format).data(input);
}
public Writer dataCompression(ClickHouseCompression compression) {
this.compression = Objects.requireNonNull(compression, "Compression can not be null");
this.addDbParam(ClickHouseQueryParam.COMPRESS, String.valueOf(compression != ClickHouseCompression.none));
return this;
}
public Writer data(File input, ClickHouseFormat format) {
return format(format).data(input);
}
/**
* Method to call, when Writer is fully configured.
*/
public void send() throws SQLException {
HttpEntity entity;
try {
InputStream stream;
if (null == streamProvider || null == (stream = streamProvider.get())) {
throw new IOException("No input data specified");
}
entity = new InputStreamEntity(stream);
} catch (IOException err) {
throw new SQLException(err);
}
send(entity);
}
private void send(HttpEntity entity) throws SQLException {
statement.sendStream(this, entity);
}
/**
* Allows to send stream of data to ClickHouse.
*
* @param sql
* in a form of "INSERT INTO table_name (X,Y,Z) VALUES "
* @param data
* where to read data from
* @param format
* format of data in InputStream
* @throws SQLException
* if the upload fails
*/
public void send(String sql, InputStream data, ClickHouseFormat format) throws SQLException {
sql(sql).data(data).format(format).send();
}
/**
* Convenient method for importing the data into table.
*
* @param table
* table name
* @param data
* source data
* @param format
* format of data in InputStream
* @throws SQLException
* if the upload fails
*/
public void sendToTable(String table, InputStream data, ClickHouseFormat format) throws SQLException {
table(table).data(data).format(format).send();
}
/**
* Sends the data in {@link ClickHouseFormat#RowBinary RowBinary} or in
* {@link ClickHouseFormat#Native Native} format.
*
* @param sql
* the SQL statement to execute
* @param callback
* data source for the upload
* @param format
* the format to use, either {@link ClickHouseFormat#RowBinary
* RowBinary} or {@link ClickHouseFormat#Native Native}
* @throws SQLException
* if the upload fails
*/
public void send(String sql, ClickHouseStreamCallback callback, ClickHouseFormat format) throws SQLException {
if (!(RowBinary.equals(format) || Native.equals(format))) {
throw new SQLException("Wrong binary format - only RowBinary and Native are supported");
}
format(format).sql(sql).send(new ClickHouseStreamHttpEntity(callback, statement.getConnection().getTimeZone(), statement.properties));
}
String getSql() {
if (null != table) {
return "INSERT INTO " + table + " FORMAT " + format;
} else if (null != sql) {
String result = sql;
if (!ClickHouseFormat.containsFormat(result)) {
result += " FORMAT " + format;
}
return result;
} else {
throw new IllegalArgumentException("Neither table nor SQL clause are specified");
}
}
private interface InputStreamProvider {
InputStream get() throws IOException;
}
private static final class FileInputProvider implements InputStreamProvider {
private final File file;
private FileInputProvider(File file) {
this.file = file;
}
@Override
public InputStream get() throws IOException {
return new FileInputStream(file);
}
}
private static final class HoldingInputProvider implements InputStreamProvider {
private final InputStream stream;
private HoldingInputProvider(InputStream stream) {
this.stream = stream;
}
@Override
public InputStream get() throws IOException {
return stream;
}
}
public ClickHouseCompression getCompression() {
return compression;
}
}