com.amazon.redshift.core.RedshiftStream 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
/*
* Copyright (c) 2017, PostgreSQL Global Development Group
* See the LICENSE file in the project root for more information.
*/
package com.amazon.redshift.core;
import com.amazon.redshift.RedshiftProperty;
import com.amazon.redshift.logger.LogLevel;
import com.amazon.redshift.logger.RedshiftLogger;
import com.amazon.redshift.util.ByteStreamWriter;
import com.amazon.redshift.util.GT;
import com.amazon.redshift.util.HostSpec;
import com.amazon.redshift.util.RedshiftPropertyMaxResultBufferParser;
import com.amazon.redshift.util.RedshiftException;
import com.amazon.redshift.util.RedshiftState;
import java.io.BufferedOutputStream;
import java.io.Closeable;
import java.io.EOFException;
import java.io.FilterOutputStream;
import java.io.Flushable;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.io.Writer;
import java.net.InetSocketAddress;
import java.net.Socket;
import java.net.SocketTimeoutException;
import java.sql.SQLException;
import java.util.Properties;
import javax.net.SocketFactory;
import static com.amazon.redshift.jdbc.RedshiftConnectionImpl.getOptionalSetting;
/**
* Wrapper around the raw connection to the server that implements some basic primitives
* (reading/writing formatted data, doing string encoding, etc).
*
* In general, instances of RedshiftStream are not threadsafe; the caller must ensure that only one thread
* at a time is accessing a particular RedshiftStream instance.
*/
public class RedshiftStream implements Closeable, Flushable {
private final SocketFactory socketFactory;
private final HostSpec hostSpec;
private final byte[] int4Buf;
private final byte[] int2Buf;
private Socket connection;
private VisibleBufferedInputStream pgInput;
private CompressedInputStream pgCompressedInput;
private OutputStream pgOutput;
private byte[] streamBuffer;
private long nextStreamAvailableCheckTime;
// This is a workaround for SSL sockets: sslInputStream.available() might return 0
// so we perform "1ms reads" once in a while
private int minStreamAvailableCheckDelay = 1000;
private Encoding encoding;
private Writer encodingWriter;
private long maxResultBuffer = -1;
private long resultBufferByteCount = 0;
private RedshiftLogger logger;
/**
* Constructor: Connect to the Redshift back end and return a stream connection.
*
* @param socketFactory socket factory to use when creating sockets
* @param hostSpec the host and port to connect to
* @param timeout timeout in milliseconds, or 0 if no timeout set
* @param logger the logger to log the entry for debugging.
* @throws IOException if an IOException occurs below it.
*/
public RedshiftStream(SocketFactory socketFactory, HostSpec hostSpec, int timeout, RedshiftLogger logger, boolean disableCompressionForSSL, Properties info) throws IOException {
this.logger = logger;
this.socketFactory = socketFactory;
this.hostSpec = hostSpec;
Socket socket = socketFactory.createSocket();
if (!socket.isConnected()) {
// When using a SOCKS proxy, the host might not be resolvable locally,
// thus we defer resolution until the traffic reaches the proxy. If there
// is no proxy, we must resolve the host to an IP to connect the socket.
if(RedshiftLogger.isEnable())
logger.log(LogLevel.INFO, "hostspec host: " + hostSpec.getHost());
logger.log(LogLevel.INFO, "hostspec port: " + hostSpec.getPort());
InetSocketAddress address = hostSpec.shouldResolve()
? new InetSocketAddress(hostSpec.getHost(), hostSpec.getPort())
: InetSocketAddress.createUnresolved(hostSpec.getHost(), hostSpec.getPort());
socket.connect(address, timeout);
if(RedshiftLogger.isEnable())
logger.log(LogLevel.INFO, "address: " + address.getAddress());
logger.log(LogLevel.INFO, "port: " + address.getPort());
logger.log(LogLevel.INFO, "hostname: " + address.getHostName());
logger.log(LogLevel.INFO, "hoststring: " + address.getHostString());
}
changeSocket(socket, disableCompressionForSSL, info);
setEncoding(Encoding.getJVMEncoding("UTF-8", logger));
int2Buf = new byte[2];
int4Buf = new byte[4];
if(RedshiftLogger.isEnable())
logger.log(LogLevel.INFO, "Gets a new stream on a new socket");
}
/**
* Constructor: Connect to the Redshift back end and return a stream connection.
*
* @param socketFactory socket factory
* @param hostSpec the host and port to connect to
* @param logger the logger to log the entry for debugging.
* @throws IOException if an IOException occurs below it.
* @deprecated use {@link #RedshiftStream(SocketFactory, com.amazon.redshift.util.HostSpec, int, RedshiftLogger, boolean, Properties)}
*/
@Deprecated
public RedshiftStream(SocketFactory socketFactory, HostSpec hostSpec, RedshiftLogger logger, Properties info) throws IOException {
this(socketFactory, hostSpec, 0, logger, true, info);
}
public RedshiftLogger getLogger() {
return logger;
}
public HostSpec getHostSpec() {
return hostSpec;
}
public Socket getSocket() {
return connection;
}
public SocketFactory getSocketFactory() {
return socketFactory;
}
/**
* Check for pending backend messages without blocking. Might return false when there actually are
* messages waiting, depending on the characteristics of the underlying socket. This is used to
* detect asynchronous notifies from the backend, when available.
*
* @return true if there is a pending backend message
* @throws IOException if something wrong happens
*/
public boolean hasMessagePending() throws IOException {
boolean available = false;
// In certain cases, available returns 0, yet there are bytes
if (pgInput.available() > 0) {
return true;
}
long now = System.nanoTime() / 1000000;
if (now < nextStreamAvailableCheckTime && minStreamAvailableCheckDelay != 0) {
// Do not use ".peek" too often
return false;
}
int soTimeout = getNetworkTimeout();
connection.setSoTimeout(1);
try {
if (!pgInput.ensureBytes(1, false)) {
return false;
}
available = (pgInput.peek() != -1);
} catch (SocketTimeoutException e) {
return false;
} finally {
connection.setSoTimeout(soTimeout);
}
/*
If none available then set the next check time
In the event that there more async bytes available we will continue to get them all
see issue 1547 https://github.com/pgjdbc/pgjdbc/issues/1547
*/
if (!available) {
nextStreamAvailableCheckTime = now + minStreamAvailableCheckDelay;
}
return available;
}
public void setMinStreamAvailableCheckDelay(int delay) {
this.minStreamAvailableCheckDelay = delay;
}
/**
* Switch this stream to using a new socket. Any existing socket is not closed; it's
* assumed that we are changing to a new socket that delegates to the original socket (e.g. SSL).
*
* @param socket the new socket to change to
* @throws IOException if something goes wrong
*/
public void changeSocket(Socket socket, Boolean disableCompressionForSSL, Properties info) throws IOException {
this.connection = socket;
changeStream(disableCompressionForSSL, info);
}
public void changeStream(Boolean disableCompressionForSSL, Properties info) throws IOException {
// Submitted by Jason Venner . Disable Nagle
// as we are selective about flushing output only when we
// really need to.
connection.setTcpNoDelay(true);
// Buffer sizes submitted by Sverre H Huseby
InputStream connectionStream = connection.getInputStream();
String compressionMode = getOptionalSetting(RedshiftProperty.COMPRESSION.getName(), info);
compressionMode = null == compressionMode ? RedshiftProperty.COMPRESSION.getDefaultValue() : compressionMode;
// If doing SSL handshake or if compression is set to off by user, use regular input stream
if(disableCompressionForSSL || compressionMode.equalsIgnoreCase("off"))
{
if(RedshiftLogger.isEnable())
{
logger.logInfo("Compression is disabled. Creating regular input stream.");
}
pgInput = new VisibleBufferedInputStream(connectionStream, 8192);
}
else
{
// Use a compressed input stream
if(RedshiftLogger.isEnable())
{
logger.logInfo("Compression is enabled. Creating compressed input stream.");
}
pgCompressedInput = new CompressedInputStream(connectionStream, logger);
pgInput = new VisibleBufferedInputStream(pgCompressedInput, 8192);
}
pgOutput = new BufferedOutputStream(connection.getOutputStream(), 8192);
if (encoding != null) {
setEncoding(encoding);
}
}
public long getBytesFromStream()
{
if(null != pgCompressedInput)
{
return pgCompressedInput.getBytesReadFromStream();
}
return pgInput.getBytesReadFromStream();
}
public Encoding getEncoding() {
return encoding;
}
/**
* Change the encoding used by this connection.
*
* @param encoding the new encoding to use
* @throws IOException if something goes wrong
*/
public void setEncoding(Encoding encoding) throws IOException {
if (this.encoding != null && this.encoding.name().equals(encoding.name())) {
return;
}
// Close down any old writer.
if (encodingWriter != null) {
encodingWriter.close();
}
this.encoding = encoding;
// Intercept flush() downcalls from the writer; our caller
// will call RedshiftStream.flush() as needed.
OutputStream interceptor = new FilterOutputStream(pgOutput) {
public void flush() throws IOException {
}
public void close() throws IOException {
super.flush();
}
};
encodingWriter = encoding.getEncodingWriter(interceptor);
}
/**
* Get a Writer instance that encodes directly onto the underlying stream.
*
* The returned Writer should not be closed, as it's a shared object. Writer.flush needs to be
* called when switching between use of the Writer and use of the RedshiftStream write methods, but it
* won't actually flush output all the way out -- call {@link #flush} to actually ensure all
* output has been pushed to the server.
*
* @return the shared Writer instance
* @throws IOException if something goes wrong.
*/
public Writer getEncodingWriter() throws IOException {
if (encodingWriter == null) {
throw new IOException("No encoding has been set on this connection");
}
return encodingWriter;
}
/**
* Sends a single character to the back end.
*
* @param val the character to be sent
* @throws IOException if an I/O error occurs
*/
public void sendChar(int val) throws IOException {
pgOutput.write(val);
}
/**
* Sends a 4-byte integer to the back end.
*
* @param val the integer to be sent
* @throws IOException if an I/O error occurs
*/
public void sendInteger4(int val) throws IOException {
int4Buf[0] = (byte) (val >>> 24);
int4Buf[1] = (byte) (val >>> 16);
int4Buf[2] = (byte) (val >>> 8);
int4Buf[3] = (byte) (val);
pgOutput.write(int4Buf);
}
/**
* Sends a 2-byte integer (short) to the back end.
*
* @param val the integer to be sent
* @throws IOException if an I/O error occurs or {@code val} cannot be encoded in 2 bytes
*/
public void sendInteger2(int val) throws IOException {
if (val < Short.MIN_VALUE || val > Short.MAX_VALUE) {
throw new IOException("Tried to send an out-of-range integer as a 2-byte value: " + val);
}
int2Buf[0] = (byte) (val >>> 8);
int2Buf[1] = (byte) val;
pgOutput.write(int2Buf);
}
/**
* Send an array of bytes to the backend.
*
* @param buf The array of bytes to be sent
* @throws IOException if an I/O error occurs
*/
public void send(byte[] buf) throws IOException {
pgOutput.write(buf);
}
/**
* Send a fixed-size array of bytes to the backend. If {@code buf.length < siz}, pad with zeros.
* If {@code buf.lengh > siz}, truncate the array.
*
* @param buf the array of bytes to be sent
* @param siz the number of bytes to be sent
* @throws IOException if an I/O error occurs
*/
public void send(byte[] buf, int siz) throws IOException {
send(buf, 0, siz);
}
/**
* Send a fixed-size array of bytes to the backend. If {@code length < siz}, pad with zeros. If
* {@code length > siz}, truncate the array.
*
* @param buf the array of bytes to be sent
* @param off offset in the array to start sending from
* @param siz the number of bytes to be sent
* @throws IOException if an I/O error occurs
*/
public void send(byte[] buf, int off, int siz) throws IOException {
int bufamt = buf.length - off;
pgOutput.write(buf, off, bufamt < siz ? bufamt : siz);
for (int i = bufamt; i < siz; ++i) {
pgOutput.write(0);
}
}
/**
* Send a fixed-size array of bytes to the backend. If {@code length < siz}, pad with zeros. If
* {@code length > siz}, truncate the array.
*
* @param writer the stream writer to invoke to send the bytes
* @throws IOException if an I/O error occurs
*/
public void send(ByteStreamWriter writer) throws IOException {
final FixedLengthOutputStream fixedLengthStream = new FixedLengthOutputStream(writer.getLength(), pgOutput);
try {
writer.writeTo(new ByteStreamWriter.ByteStreamTarget() {
@Override
public OutputStream getOutputStream() {
return fixedLengthStream;
}
});
} catch (IOException ioe) {
throw ioe;
} catch (Exception re) {
throw new IOException("Error writing bytes to stream", re);
}
for (int i = fixedLengthStream.remaining(); i > 0; i--) {
pgOutput.write(0);
}
}
/**
* Receives a single character from the backend, without advancing the current protocol stream
* position.
*
* @return the character received
* @throws IOException if an I/O Error occurs
*/
public int peekChar() throws IOException {
int c = pgInput.peek();
if (c < 0) {
throw new EOFException();
}
return c;
}
/**
* Receives a single character from the backend.
*
* @return the character received
* @throws IOException if an I/O Error occurs
*/
public int receiveChar() throws IOException {
int c = pgInput.read();
if (c < 0) {
throw new EOFException("The server closed the connection.");
}
return c;
}
/**
* Receives a four byte integer from the backend.
*
* @return the integer received from the backend
* @throws IOException if an I/O error occurs
*/
public int receiveInteger4() throws IOException {
if (pgInput.read(int4Buf) != 4) {
throw new EOFException();
}
return (int4Buf[0] & 0xFF) << 24 | (int4Buf[1] & 0xFF) << 16 | (int4Buf[2] & 0xFF) << 8
| int4Buf[3] & 0xFF;
}
/**
* Receives a two byte integer from the backend.
*
* @return the integer received from the backend
* @throws IOException if an I/O error occurs
*/
public int receiveInteger2() throws IOException {
if (pgInput.read(int2Buf) != 2) {
throw new EOFException();
}
return (int2Buf[0] & 0xFF) << 8 | int2Buf[1] & 0xFF;
}
/**
* Receives a fixed-size string from the backend.
*
* @param len the length of the string to receive, in bytes.
* @return the decoded string
* @throws IOException if something wrong happens
*/
public String receiveString(int len) throws IOException {
if (!pgInput.ensureBytes(len)) {
throw new EOFException();
}
String res = encoding.decode(pgInput.getBuffer(), pgInput.getIndex(), len);
pgInput.skip(len);
return res;
}
/**
* Receives a fixed-size string from the backend, and tries to avoid "UTF-8 decode failed"
* errors.
*
* @param len the length of the string to receive, in bytes.
* @return the decoded string
* @throws IOException if something wrong happens
*/
public EncodingPredictor.DecodeResult receiveErrorString(int len) throws IOException {
if (!pgInput.ensureBytes(len)) {
throw new EOFException();
}
EncodingPredictor.DecodeResult res;
try {
String value = encoding.decode(pgInput.getBuffer(), pgInput.getIndex(), len);
// no autodetect warning as the message was converted on its own
res = new EncodingPredictor.DecodeResult(value, null);
} catch (IOException e) {
res = EncodingPredictor.decode(pgInput.getBuffer(), pgInput.getIndex(), len, logger);
if (res == null) {
Encoding enc = Encoding.defaultEncoding();
String value = enc.decode(pgInput.getBuffer(), pgInput.getIndex(), len);
res = new EncodingPredictor.DecodeResult(value, enc.name());
}
}
pgInput.skip(len);
return res;
}
/**
* Receives a null-terminated string from the backend. If we don't see a null, then we assume
* something has gone wrong.
*
* @return string from back end
* @throws IOException if an I/O error occurs, or end of file
*/
public String receiveString() throws IOException {
int len = pgInput.scanCStringLength();
String res = encoding.decode(pgInput.getBuffer(), pgInput.getIndex(), len - 1);
pgInput.skip(len);
return res;
}
/**
* Read a tuple from the back end. A tuple is a two dimensional array of bytes. This variant reads
* the V3 protocol's tuple representation.
*
* @return tuple from the back end
* @throws IOException if a data I/O error occurs
* @throws SQLException if read more bytes than set maxResultBuffer
*/
public Tuple receiveTupleV3() throws IOException, OutOfMemoryError, SQLException {
int messageSize = receiveInteger4(); // MESSAGE SIZE
int nf = receiveInteger2();
//size = messageSize - 4 bytes of message size - 2 bytes of field count - 4 bytes for each column length
int dataToReadSize = messageSize - 4 - 2 - 4 * nf;
byte[][] answer = new byte[nf][];
increaseByteCounter(dataToReadSize);
OutOfMemoryError oom = null;
for (int i = 0; i < nf; ++i) {
int size = receiveInteger4();
if (size != -1) {
try {
answer[i] = new byte[size];
receive(answer[i], 0, size);
} catch (OutOfMemoryError oome) {
oom = oome;
skip(size);
}
}
}
if (oom != null) {
throw oom;
}
return new Tuple(answer, dataToReadSize);
}
/**
* Reads in a given number of bytes from the backend.
*
* @param siz number of bytes to read
* @return array of bytes received
* @throws IOException if a data I/O error occurs
*/
public byte[] receive(int siz) throws IOException {
byte[] answer = new byte[siz];
receive(answer, 0, siz);
return answer;
}
/**
* Reads in a given number of bytes from the backend.
*
* @param buf buffer to store result
* @param off offset in buffer
* @param siz number of bytes to read
* @throws IOException if a data I/O error occurs
*/
public void receive(byte[] buf, int off, int siz) throws IOException {
int s = 0;
while (s < siz) {
int w = pgInput.read(buf, off + s, siz - s);
if (w < 0) {
throw new EOFException();
}
s += w;
}
}
public void skip(int size) throws IOException {
long s = 0;
while (s < size) {
s += pgInput.skip(size - s);
}
}
/**
* Copy data from an input stream to the connection.
*
* @param inStream the stream to read data from
* @param remaining the number of bytes to copy
* @throws IOException if a data I/O error occurs
*/
public void sendStream(InputStream inStream, int remaining) throws IOException {
int expectedLength = remaining;
if (streamBuffer == null) {
streamBuffer = new byte[8192];
}
while (remaining > 0) {
int count = (remaining > streamBuffer.length ? streamBuffer.length : remaining);
int readCount;
try {
readCount = inStream.read(streamBuffer, 0, count);
if (readCount < 0) {
throw new EOFException(
GT.tr("Premature end of input stream, expected {0} bytes, but only read {1}.",
expectedLength, expectedLength - remaining));
}
} catch (IOException ioe) {
while (remaining > 0) {
send(streamBuffer, count);
remaining -= count;
count = (remaining > streamBuffer.length ? streamBuffer.length : remaining);
}
throw new RedshiftBindException(ioe);
}
send(streamBuffer, readCount);
remaining -= readCount;
}
}
/**
* Flush any pending output to the backend.
*
* @throws IOException if an I/O error occurs
*/
@Override
public void flush() throws IOException {
if (encodingWriter != null) {
encodingWriter.flush();
}
pgOutput.flush();
}
/**
* Consume an expected EOF from the backend.
*
* @throws IOException if an I/O error occurs
* @throws SQLException if we get something other than an EOF
*/
public void receiveEOF() throws SQLException, IOException {
int c = pgInput.read();
if (c < 0) {
return;
}
throw new RedshiftException(GT.tr("Expected an EOF from server, got: {0}", c),
RedshiftState.COMMUNICATION_ERROR);
}
/**
* Closes the connection.
*
* @throws IOException if an I/O Error occurs
*/
@Override
public void close() throws IOException {
if(RedshiftLogger.isEnable())
logger.log(LogLevel.INFO, "Stream on a connected socket closed");
if (encodingWriter != null) {
encodingWriter.close();
}
pgOutput.close();
pgInput.close();
connection.close();
}
public void setNetworkTimeout(int milliseconds) throws IOException {
connection.setSoTimeout(milliseconds);
pgInput.setTimeoutRequested(milliseconds != 0);
}
public int getNetworkTimeout() throws IOException {
return connection.getSoTimeout();
}
/**
* Method to set MaxResultBuffer inside RedshiftStream.
*
* @param value value of new max result buffer as string (cause we can expect % or chars to use
* multiplier)
* @throws RedshiftException exception returned when occurred parsing problem.
*/
public void setMaxResultBuffer(String value) throws RedshiftException {
maxResultBuffer = RedshiftPropertyMaxResultBufferParser.parseProperty(value, RedshiftProperty.MAX_RESULT_BUFFER.getName());
}
/**
* Method to clear count of byte buffer.
*/
public void clearResultBufferCount() {
resultBufferByteCount = 0;
}
/**
* Method to increase actual count of buffer. If buffer count is bigger than max result buffer
* limit, then gonna return an exception.
*
* @param value size of bytes to add to byte buffer.
* @throws SQLException exception returned when result buffer count is bigger than max result
* buffer.
*/
private void increaseByteCounter(long value) throws SQLException {
if (maxResultBuffer != -1) {
resultBufferByteCount += value;
if (resultBufferByteCount > maxResultBuffer) {
throw new RedshiftException(GT.tr(
"Result set exceeded maxResultBuffer limit. Received: {0}; Current limit: {1}",
String.valueOf(resultBufferByteCount), String.valueOf(maxResultBuffer)),RedshiftState.COMMUNICATION_ERROR);
}
}
}
public boolean isClosed() {
return connection.isClosed();
}
}