Repository: pgjdbc/pgjdbc
Pull request #1702 (2020-02-11 09:47:24)

feat: Expose `ByteStreamWriter` in CopyManager
Continuing https://github.com/pgjdbc/pgjdbc/pull/953, let's expose the garbageless writer to the COPY protocol as well.

### All Submissions:

* [x] Have you followed the guidelines in our [Contributing](https://github.com/pgjdbc/pgjdbc/blob/master/CONTRIBUTING.md) document?
* [x] Have you checked to ensure there aren't other open [Pull Requests](../../pulls) for the same update/change?

<!-- You can erase any parts of this template not applicable to your Pull Request. -->

### New Feature Submissions:

1. [ ] Does your submission pass tests?
2. [x] Does mvn checkstyle:check pass?
3. [x] Have you added your new test classes to an existing test suite in alphabetical order?

### Changes to Existing Features:

* [ ] Does this break existing behaviour? If so please explain.
* [x] Have you added an explanation of what your changes do and why you'd like us to include them?
* [x] Have you written new tests for your core changes, as applicable?
* [ ] Have you successfully run tests with your changes locally?


davecramer: looks fine to me. I'll merge this later if nobody has any issues with it

Patch:
diff --git a/pgjdbc/src/main/java/org/postgresql/copy/CopyIn.java b/pgjdbc/src/main/java/org/postgresql/copy/CopyIn.java
index 73671bb0..b0cd5b4d 100644
--- a/pgjdbc/src/main/java/org/postgresql/copy/CopyIn.java
+++ b/pgjdbc/src/main/java/org/postgresql/copy/CopyIn.java
@@ -5,6 +5,8 @@
 
 package org.postgresql.copy;
 
+import org.postgresql.util.ByteStreamWriter;
+
 import java.sql.SQLException;
 
 /**
@@ -22,6 +24,14 @@ public interface CopyIn extends CopyOperation {
    */
   void writeToCopy(byte[] buf, int off, int siz) throws SQLException;
 
+  /**
+   * Writes a ByteStreamWriter to an open and writable copy operation.
+   *
+   * @param from the source of bytes, e.g. a ByteBufferByteStreamWriter
+   * @throws SQLException if the operation fails
+   */
+  void writeToCopy(ByteStreamWriter from) throws SQLException;
+
   /**
    * Force any buffered output to be sent over the network to the backend. In general this is a
    * useless operation as it will get pushed over in due time or when endCopy is called. Some
diff --git a/pgjdbc/src/main/java/org/postgresql/copy/CopyManager.java b/pgjdbc/src/main/java/org/postgresql/copy/CopyManager.java
index a152b7c1..e723dc47 100644
--- a/pgjdbc/src/main/java/org/postgresql/copy/CopyManager.java
+++ b/pgjdbc/src/main/java/org/postgresql/copy/CopyManager.java
@@ -8,6 +8,7 @@ package org.postgresql.copy;
 import org.postgresql.core.BaseConnection;
 import org.postgresql.core.Encoding;
 import org.postgresql.core.QueryExecutor;
+import org.postgresql.util.ByteStreamWriter;
 import org.postgresql.util.GT;
 import org.postgresql.util.PSQLException;
 import org.postgresql.util.PSQLState;
@@ -227,4 +228,26 @@ public class CopyManager {
       }
     }
   }
+
+  /**
+   * 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();
+      }
+    }
+  }
 }
diff --git a/pgjdbc/src/main/java/org/postgresql/copy/PGCopyOutputStream.java b/pgjdbc/src/main/java/org/postgresql/copy/PGCopyOutputStream.java
index 520c645e..84692ba7 100644
--- a/pgjdbc/src/main/java/org/postgresql/copy/PGCopyOutputStream.java
+++ b/pgjdbc/src/main/java/org/postgresql/copy/PGCopyOutputStream.java
@@ -6,6 +6,7 @@
 package org.postgresql.copy;
 
 import org.postgresql.PGConnection;
+import org.postgresql.util.ByteStreamWriter;
 import org.postgresql.util.GT;
 
 import java.io.IOException;
@@ -140,6 +141,10 @@ public class PGCopyOutputStream extends OutputStream implements CopyIn {
     }
   }
 
+  public void writeToCopy(ByteStreamWriter from) throws SQLException {
+    op.writeToCopy(from);
+  }
+
   public int getFormat() {
     return op.getFormat();
   }
diff --git a/pgjdbc/src/main/java/org/postgresql/core/v3/CopyDualImpl.java b/pgjdbc/src/main/java/org/postgresql/core/v3/CopyDualImpl.java
index 076321f7..62ebda9b 100644
--- a/pgjdbc/src/main/java/org/postgresql/core/v3/CopyDualImpl.java
+++ b/pgjdbc/src/main/java/org/postgresql/core/v3/CopyDualImpl.java
@@ -6,6 +6,7 @@
 package org.postgresql.core.v3;
 
 import org.postgresql.copy.CopyDual;
+import org.postgresql.util.ByteStreamWriter;
 import org.postgresql.util.PSQLException;
 
 import java.sql.SQLException;
@@ -19,6 +20,10 @@ public class CopyDualImpl extends CopyOperationImpl implements CopyDual {
     queryExecutor.writeToCopy(this, data, off, siz);
   }
 
+  public void writeToCopy(ByteStreamWriter from) throws SQLException {
+    queryExecutor.writeToCopy(this, from);
+  }
+
   public void flushCopy() throws SQLException {
     queryExecutor.flushCopy(this);
   }
diff --git a/pgjdbc/src/main/java/org/postgresql/core/v3/CopyInImpl.java b/pgjdbc/src/main/java/org/postgresql/core/v3/CopyInImpl.java
index f1ce995e..0da2c1e1 100644
--- a/pgjdbc/src/main/java/org/postgresql/core/v3/CopyInImpl.java
+++ b/pgjdbc/src/main/java/org/postgresql/core/v3/CopyInImpl.java
@@ -6,6 +6,7 @@
 package org.postgresql.core.v3;
 
 import org.postgresql.copy.CopyIn;
+import org.postgresql.util.ByteStreamWriter;
 import org.postgresql.util.GT;
 import org.postgresql.util.PSQLException;
 import org.postgresql.util.PSQLState;
@@ -37,6 +38,10 @@ public class CopyInImpl extends CopyOperationImpl implements CopyIn {
     queryExecutor.writeToCopy(this, data, off, siz);
   }
 
+  public void writeToCopy(ByteStreamWriter from) throws SQLException {
+    queryExecutor.writeToCopy(this, from);
+  }
+
   public void flushCopy() throws SQLException {
     queryExecutor.flushCopy(this);
   }
diff --git a/pgjdbc/src/main/java/org/postgresql/core/v3/QueryExecutorImpl.java b/pgjdbc/src/main/java/org/postgresql/core/v3/QueryExecutorImpl.java
index 4c1c32f1..ab06a88b 100644
--- a/pgjdbc/src/main/java/org/postgresql/core/v3/QueryExecutorImpl.java
+++ b/pgjdbc/src/main/java/org/postgresql/core/v3/QueryExecutorImpl.java
@@ -37,6 +37,7 @@ import org.postgresql.core.v3.replication.V3ReplicationProtocol;
 import org.postgresql.jdbc.AutoSave;
 import org.postgresql.jdbc.BatchResultHandler;
 import org.postgresql.jdbc.TimestampUtils;
+import org.postgresql.util.ByteStreamWriter;
 import org.postgresql.util.GT;
 import org.postgresql.util.PSQLException;
 import org.postgresql.util.PSQLState;
@@ -1034,7 +1035,34 @@ public class QueryExecutorImpl extends QueryExecutorBase {
       pgStream.sendChar('d');
       pgStream.sendInteger4(siz + 4);
       pgStream.send(data, off, siz);
+    } catch (IOException ioe) {
+      throw new PSQLException(GT.tr("Database connection failed when writing to copy"),
+          PSQLState.CONNECTION_FAILURE, ioe);
+    }
+  }
+
+  /**
+   * Sends data during a live COPY IN operation. Only unlocks the connection if server suddenly
+   * returns CommandComplete, which should not happen
+   *
+   * @param op   the CopyIn operation presumably currently holding lock on this connection
+   * @param from the source of bytes, e.g. a ByteBufferByteStreamWriter
+   * @throws SQLException on failure
+   */
+  public synchronized void writeToCopy(CopyOperationImpl op, ByteStreamWriter from)
+      throws SQLException {
+    if (!hasLock(op)) {
+      throw new PSQLException(GT.tr("Tried to write to an inactive copy operation"),
+          PSQLState.OBJECT_NOT_IN_STATE);
+    }
 
+    int siz = from.getLength();
+    LOGGER.log(Level.FINEST, " FE=> CopyData({0})", siz);
+
+    try {
+      pgStream.sendChar('d');
+      pgStream.sendInteger4(siz + 4);
+      pgStream.send(from);
     } catch (IOException ioe) {
       throw new PSQLException(GT.tr("Database connection failed when writing to copy"),
           PSQLState.CONNECTION_FAILURE, ioe);


Test patch:
diff --git a/pgjdbc/src/test/java/org/postgresql/test/jdbc2/CopyTest.java b/pgjdbc/src/test/java/org/postgresql/test/jdbc2/CopyTest.java
index fde54ab8..80157a1d 100644
--- a/pgjdbc/src/test/java/org/postgresql/test/jdbc2/CopyTest.java
+++ b/pgjdbc/src/test/java/org/postgresql/test/jdbc2/CopyTest.java
@@ -18,6 +18,7 @@ import org.postgresql.copy.CopyOut;
 import org.postgresql.copy.PGCopyOutputStream;
 import org.postgresql.core.ServerVersion;
 import org.postgresql.test.TestUtil;
+import org.postgresql.util.ByteBufferByteStreamWriter;
 import org.postgresql.util.PSQLState;
 
 import org.junit.After;
@@ -31,6 +32,7 @@ import java.io.InputStream;
 import java.io.OutputStream;
 import java.io.PrintStream;
 import java.io.StringReader;
+import java.nio.ByteBuffer;
 import java.sql.Connection;
 import java.sql.ResultSet;
 import java.sql.SQLException;
@@ -210,6 +212,14 @@ public class CopyTest {
     assertEquals(dataRows, rowCount);
   }
 
+  @Test
+  public void testCopyInFromByteStreamWriter() throws SQLException, IOException {
+    String sql = "COPY copytest FROM STDIN";
+    copyAPI.copyIn(sql, new ByteBufferByteStreamWriter(ByteBuffer.wrap(getData(origData))));
+    int rowCount = getCount();
+    assertEquals(dataRows, rowCount);
+  }
+
   @Test
   public void testSkipping() {
     String sql = "COPY copytest FROM STDIN";