com.amazon.redshift.copy.CopyManager Maven / Gradle / Ivy
Go to download
Show more of this group Show more artifacts with this name
Show all versions of redshift-jdbc42 Show documentation
Show all versions of redshift-jdbc42 Show documentation
Java JDBC 4.2 (JRE 8+) driver for Redshift database
The newest version!
/*
* Copyright (c) 2009, PostgreSQL Global Development Group
* See the LICENSE file in the project root for more information.
*/
package com.amazon.redshift.copy;
import com.amazon.redshift.core.BaseConnection;
import com.amazon.redshift.core.Encoding;
import com.amazon.redshift.core.QueryExecutor;
import com.amazon.redshift.util.ByteStreamWriter;
import com.amazon.redshift.util.GT;
import com.amazon.redshift.util.RedshiftException;
import com.amazon.redshift.util.RedshiftState;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.io.Reader;
import java.io.Writer;
import java.sql.SQLException;
/**
* API for Redshift COPY bulk data transfer.
*/
public class CopyManager {
// I don't know what the best buffer size is, so we let people specify it if
// they want, and if they don't know, we don't make them guess, so that if we
// do figure it out we can just set it here and they reap the rewards.
// Note that this is currently being used for both a number of bytes and a number
// of characters.
static final int DEFAULT_BUFFER_SIZE = 65536;
private final Encoding encoding;
private final QueryExecutor queryExecutor;
private final BaseConnection connection;
public CopyManager(BaseConnection connection) throws SQLException {
this.encoding = connection.getEncoding();
this.queryExecutor = connection.getQueryExecutor();
this.connection = connection;
}
public CopyIn copyIn(String sql) throws SQLException {
CopyOperation op = queryExecutor.startCopy(sql, connection.getAutoCommit());
if (op == null || op instanceof CopyIn) {
return (CopyIn) op;
} else {
op.cancelCopy();
throw new RedshiftException(GT.tr("Requested CopyIn but got {0}", op.getClass().getName()),
RedshiftState.WRONG_OBJECT_TYPE);
}
}
public CopyOut copyOut(String sql) throws SQLException {
CopyOperation op = queryExecutor.startCopy(sql, connection.getAutoCommit());
if (op == null || op instanceof CopyOut) {
return (CopyOut) op;
} else {
op.cancelCopy();
throw new RedshiftException(GT.tr("Requested CopyOut but got {0}", op.getClass().getName()),
RedshiftState.WRONG_OBJECT_TYPE);
}
}
public CopyDual copyDual(String sql) throws SQLException {
CopyOperation op = queryExecutor.startCopy(sql, connection.getAutoCommit());
if (op == null || op instanceof CopyDual) {
return (CopyDual) op;
} else {
op.cancelCopy();
throw new RedshiftException(GT.tr("Requested CopyDual but got {0}", op.getClass().getName()),
RedshiftState.WRONG_OBJECT_TYPE);
}
}
/**
* Pass results of a COPY TO STDOUT query from database into a Writer.
*
* @param sql COPY TO STDOUT statement
* @param to the Writer to write the results to (row by row).
* The Writer is not closed at the end of the Copy Out operation.
* @return number of rows updated for server 8.2 or newer; -1 for older
* @throws SQLException on database usage errors
* @throws IOException upon writer or database connection failure
*/
public long copyOut(final String sql, Writer to) throws SQLException, IOException {
byte[] buf;
CopyOut cp = copyOut(sql);
try {
while ((buf = cp.readFromCopy()) != null) {
to.write(encoding.decode(buf));
}
return cp.getHandledRowCount();
} catch (IOException ioEX) {
// if not handled this way the close call will hang, at least in 8.2
if (cp.isActive()) {
cp.cancelCopy();
}
try { // read until exhausted or operation cancelled SQLException
while ((buf = cp.readFromCopy()) != null) {
}
} catch (SQLException sqlEx) {
} // typically after several kB
throw ioEX;
} finally { // see to it that we do not leave the connection locked
if (cp.isActive()) {
cp.cancelCopy();
}
}
}
/**
* Pass results of a COPY TO STDOUT query from database into an OutputStream.
*
* @param sql COPY TO STDOUT statement
* @param to the stream to write the results to (row by row)
* The stream is not closed at the end of the operation. This is intentional so the
* caller can continue to write to the output stream
* @return number of rows updated for server 8.2 or newer; -1 for older
* @throws SQLException on database usage errors
* @throws IOException upon output stream or database connection failure
*/
public long copyOut(final String sql, OutputStream to) throws SQLException, IOException {
byte[] buf;
CopyOut cp = copyOut(sql);
try {
while ((buf = cp.readFromCopy()) != null) {
to.write(buf);
}
return cp.getHandledRowCount();
} catch (IOException ioEX) {
// if not handled this way the close call will hang, at least in 8.2
if (cp.isActive()) {
cp.cancelCopy();
}
try { // read until exhausted or operation cancelled SQLException
while ((buf = cp.readFromCopy()) != null) {
}
} catch (SQLException sqlEx) {
} // typically after several kB
throw ioEX;
} finally { // see to it that we do not leave the connection locked
if (cp.isActive()) {
cp.cancelCopy();
}
}
}
/**
* Use COPY FROM STDIN for very fast copying from a Reader into a database table.
*
* @param sql COPY FROM STDIN statement
* @param from a CSV file or such
* @return number of rows updated for server 8.2 or newer; -1 for older
* @throws SQLException on database usage issues
* @throws IOException upon reader or database connection failure
*/
public long copyIn(final String sql, Reader from) throws SQLException, IOException {
return copyIn(sql, from, DEFAULT_BUFFER_SIZE);
}
/**
* Use COPY FROM STDIN for very fast copying from a Reader into a database table.
*
* @param sql COPY FROM STDIN statement
* @param from a CSV file or such
* @param bufferSize number of characters to buffer and push over network to server at once
* @return number of rows updated for server 8.2 or newer; -1 for older
* @throws SQLException on database usage issues
* @throws IOException upon reader or database connection failure
*/
public long copyIn(final String sql, Reader from, int bufferSize)
throws SQLException, IOException {
char[] cbuf = new char[bufferSize];
int len;
CopyIn cp = copyIn(sql);
try {
while ((len = from.read(cbuf)) >= 0) {
if (len > 0) {
byte[] buf = encoding.encode(new String(cbuf, 0, len));
cp.writeToCopy(buf, 0, buf.length);
}
}
return cp.endCopy();
} finally { // see to it that we do not leave the connection locked
if (cp.isActive()) {
cp.cancelCopy();
}
}
}
/**
* Use COPY FROM STDIN for very fast copying from an InputStream into a database table.
*
* @param sql COPY FROM STDIN statement
* @param from a CSV file or such
* @return number of rows updated for server 8.2 or newer; -1 for older
* @throws SQLException on database usage issues
* @throws IOException upon input stream or database connection failure
*/
public long copyIn(final String sql, InputStream from) throws SQLException, IOException {
return copyIn(sql, from, DEFAULT_BUFFER_SIZE);
}
/**
* Use COPY FROM STDIN for very fast copying from an InputStream into a database table.
*
* @param sql COPY FROM STDIN statement
* @param from a CSV file or such
* @param bufferSize number of bytes to buffer and push over network to server at once
* @return number of rows updated for server 8.2 or newer; -1 for older
* @throws SQLException on database usage issues
* @throws IOException upon input stream or database connection failure
*/
public long copyIn(final String sql, InputStream from, int bufferSize)
throws SQLException, IOException {
byte[] buf = new byte[bufferSize];
int len;
CopyIn cp = copyIn(sql);
try {
while ((len = from.read(buf)) >= 0) {
if (len > 0) {
cp.writeToCopy(buf, 0, len);
}
}
return cp.endCopy();
} finally { // see to it that we do not leave the connection locked
if (cp.isActive()) {
cp.cancelCopy();
}
}
}
/**
* Use COPY FROM STDIN for very fast copying from an ByteStreamWriter into a database table.
*
* @param sql COPY FROM STDIN statement
* @param from the source of bytes, e.g. a ByteBufferByteStreamWriter
* @return number of rows updated for server 8.2 or newer; -1 for older
* @throws SQLException on database usage issues
* @throws IOException upon input stream or database connection failure
*/
public long copyIn(String sql, ByteStreamWriter from)
throws SQLException, IOException {
CopyIn cp = copyIn(sql);
try {
cp.writeToCopy(from);
return cp.endCopy();
} finally { // see to it that we do not leave the connection locked
if (cp.isActive()) {
cp.cancelCopy();
}
}
}
}