Repository: GoogleCloudPlatform/pgadapter
Pull request #390 (2022-10-13 08:39:42)

docs: document how to set a statement timeout
Add an FAQ entry for setting a statement timeout.

codecov[bot]: # [Codecov](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/390?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform) Report
> Merging [#390](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/390?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform) (401ce57) into [postgresql-dialect](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/commit/9e3e95284a34516967f2acdd2acfeb94ab50a2ea?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform) (9e3e952) will **not change** coverage.
> The diff coverage is `n/a`.

```diff
@@                  Coverage Diff                  @@
##             postgresql-dialect     #390   +/-   ##
=====================================================
  Coverage                 86.32%   86.32%           
  Complexity                 1903     1903           
=====================================================
  Files                       122      122           
  Lines                      6170     6170           
  Branches                    847      847           
=====================================================
  Hits                       5326     5326           
  Misses                      603      603           
  Partials                    241      241           
```

| Flag | Coverage Δ | |
|---|---|---|
| all_tests | `86.32% <ø> (ø)` | |

Flags with carried forward coverage won't be shown. [Click here](https://docs.codecov.io/docs/carryforward-flags?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform#carryforward-flags-in-the-pull-request-comment) to find out more.


:mega: We’re building smart automated test selection to slash your CI/CD build times. [Learn more](https://about.codecov.io/iterative-testing/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform)


Patch:
diff --git a/docs/faq.md b/docs/faq.md
index 3ccf1da4..ccc3a0d4 100644
--- a/docs/faq.md
+++ b/docs/faq.md
@@ -87,3 +87,18 @@ You can solve this problem by either:
 PGAdapter supports `COPY table_name FROM STDIN [BINARY]` and `COPY table_name TO STDOUT [BINARY]`.
 These commands can be used to copy data between different Cloud Spanner databases, or between a
 Cloud Spanner database and a PostgreSQL database. See [copy support](copy.md) for more information.
+
+## How can I set a statement timeout?
+You can set the statement timeout that a connection should use by executing a
+`set statement_timeout=<timeout>` statement. The timeout is specified in milliseconds.
+
+Example:
+
+```sql
+-- Use a 10 second (10,000 milliseconds) timeout.
+set statement_timeout=10000;
+select *
+from large_table
+order by my_col;
+```
+
diff --git a/src/main/java/com/google/cloud/spanner/pgadapter/statements/CopyToStatement.java b/src/main/java/com/google/cloud/spanner/pgadapter/statements/CopyToStatement.java
index bbdce312..fdd267ed 100644
--- a/src/main/java/com/google/cloud/spanner/pgadapter/statements/CopyToStatement.java
+++ b/src/main/java/com/google/cloud/spanner/pgadapter/statements/CopyToStatement.java
@@ -20,13 +20,13 @@ import com.google.cloud.spanner.Statement;
 import com.google.cloud.spanner.connection.AbstractStatementParser.ParsedStatement;
 import com.google.cloud.spanner.connection.AbstractStatementParser.StatementType;
 import com.google.cloud.spanner.pgadapter.ConnectionHandler;
-import com.google.cloud.spanner.pgadapter.ConnectionHandler.QueryMode;
 import com.google.cloud.spanner.pgadapter.ProxyServer.DataFormat;
 import com.google.cloud.spanner.pgadapter.metadata.DescribePortalMetadata;
 import com.google.cloud.spanner.pgadapter.metadata.OptionsMetadata;
 import com.google.cloud.spanner.pgadapter.parsers.Parser;
 import com.google.cloud.spanner.pgadapter.parsers.copy.CopyTreeParser.CopyOptions;
 import com.google.cloud.spanner.pgadapter.parsers.copy.CopyTreeParser.CopyOptions.Format;
+import com.google.cloud.spanner.pgadapter.utils.Converter;
 import com.google.cloud.spanner.pgadapter.wireoutput.CopyDataResponse;
 import com.google.cloud.spanner.pgadapter.wireoutput.CopyDoneResponse;
 import com.google.cloud.spanner.pgadapter.wireoutput.CopyOutResponse;
@@ -106,6 +106,10 @@ public class CopyToStatement extends IntermediatePortalStatement {
     return csvFormat;
   }
 
+  public boolean isBinary() {
+    return copyOptions.getFormat() == Format.BINARY;
+  }
+
   @Override
   public String getCommandTag() {
     return "COPY";
@@ -161,10 +165,10 @@ public class CopyToStatement extends IntermediatePortalStatement {
   }
 
   @Override
-  public CopyDataResponse createDataRowResponse(ResultSet resultSet, QueryMode mode) {
+  public CopyDataResponse createDataRowResponse(Converter converter) {
     return copyOptions.getFormat() == Format.BINARY
-        ? createBinaryDataResponse(resultSet)
-        : createDataResponse(resultSet);
+        ? createBinaryDataResponse(converter)
+        : createDataResponse(converter.getResultSet());
   }
 
   @Override
@@ -191,20 +195,7 @@ public class CopyToStatement extends IntermediatePortalStatement {
     return new CopyDataResponse(this.outputStream, row, csvFormat.getRecordSeparator().charAt(0));
   }
 
-  CopyDataResponse createBinaryDataResponse(ResultSet resultSet) {
-    // Multiply number of columns by 4, as each column has as 4-byte length value.
-    // In addition, each row has a 2-byte number of columns value.s
-    int length = 2 + resultSet.getColumnCount() * 4;
-    byte[][] data = new byte[resultSet.getColumnCount()][];
-    for (int col = 0; col < resultSet.getColumnCount(); col++) {
-      if (!resultSet.isNull(col)) {
-        Parser<?> parser = Parser.create(resultSet, resultSet.getColumnType(col), col);
-        data[col] = parser.parse(DataFormat.POSTGRESQL_BINARY);
-        if (data[col] != null) {
-          length += data[col].length;
-        }
-      }
-    }
-    return new CopyDataResponse(this.outputStream, length, data);
+  CopyDataResponse createBinaryDataResponse(Converter converter) {
+    return new CopyDataResponse(this.outputStream, converter);
   }
 }
diff --git a/src/main/java/com/google/cloud/spanner/pgadapter/utils/Converter.java b/src/main/java/com/google/cloud/spanner/pgadapter/utils/Converter.java
index d18ba9c5..71503321 100644
--- a/src/main/java/com/google/cloud/spanner/pgadapter/utils/Converter.java
+++ b/src/main/java/com/google/cloud/spanner/pgadapter/utils/Converter.java
@@ -16,7 +16,9 @@ package com.google.cloud.spanner.pgadapter.utils;
 
 import com.google.cloud.spanner.ResultSet;
 import com.google.cloud.spanner.Type;
+import com.google.cloud.spanner.pgadapter.ConnectionHandler.QueryMode;
 import com.google.cloud.spanner.pgadapter.ProxyServer.DataFormat;
+import com.google.cloud.spanner.pgadapter.metadata.OptionsMetadata;
 import com.google.cloud.spanner.pgadapter.parsers.ArrayParser;
 import com.google.cloud.spanner.pgadapter.parsers.BinaryParser;
 import com.google.cloud.spanner.pgadapter.parsers.BooleanParser;
@@ -27,11 +29,69 @@ import com.google.cloud.spanner.pgadapter.parsers.LongParser;
 import com.google.cloud.spanner.pgadapter.parsers.NumericParser;
 import com.google.cloud.spanner.pgadapter.parsers.StringParser;
 import com.google.cloud.spanner.pgadapter.parsers.TimestampParser;
+import com.google.cloud.spanner.pgadapter.statements.CopyToStatement;
+import com.google.cloud.spanner.pgadapter.statements.IntermediateStatement;
 import com.google.common.base.Preconditions;
+import java.io.ByteArrayOutputStream;
+import java.io.DataOutputStream;
+import java.io.IOException;
 
 /** Utility class for converting between generic PostgreSQL conversions. */
 public class Converter {
-  private Converter() {}
+  private final ByteArrayOutputStream buffer = new ByteArrayOutputStream(256);
+  private final DataOutputStream outputStream = new DataOutputStream(buffer);
+  private final IntermediateStatement statement;
+  private final QueryMode mode;
+  private final OptionsMetadata options;
+  private final ResultSet resultSet;
+
+  public Converter(
+      IntermediateStatement statement,
+      QueryMode mode,
+      OptionsMetadata options,
+      ResultSet resultSet) {
+    this.statement = statement;
+    this.mode = mode;
+    this.options = options;
+    this.resultSet = resultSet;
+  }
+
+  public ResultSet getResultSet() {
+    return resultSet;
+  }
+
+  public int convertResultSetRowToDataRowResponse() throws IOException {
+    DataFormat fixedFormat = null;
+    if (statement instanceof CopyToStatement) {
+      fixedFormat =
+          ((CopyToStatement) statement).isBinary()
+              ? DataFormat.POSTGRESQL_BINARY
+              : DataFormat.POSTGRESQL_TEXT;
+    }
+    buffer.reset();
+    outputStream.writeShort(resultSet.getColumnCount());
+    for (int column_index = 0; /* column indices start at 0 */
+        column_index < resultSet.getColumnCount();
+        column_index++) {
+      if (resultSet.isNull(column_index)) {
+        outputStream.writeInt(-1);
+      } else {
+        DataFormat format =
+            fixedFormat == null
+                ? DataFormat.getDataFormat(column_index, statement, mode, options)
+                : fixedFormat;
+        byte[] column = Converter.convertToPG(this.resultSet, column_index, format);
+        outputStream.writeInt(column.length);
+        outputStream.write(column);
+      }
+    }
+    return buffer.size();
+  }
+
+  public void writeBuffer(DataOutputStream outputStream) throws IOException {
+    buffer.writeTo(outputStream);
+  }
+
   /**
    * Return the data of the specified column of the {@link ResultSet} as a byte array. The column
    * may not contain a null value.
diff --git a/src/main/java/com/google/cloud/spanner/pgadapter/wireoutput/CopyDataResponse.java b/src/main/java/com/google/cloud/spanner/pgadapter/wireoutput/CopyDataResponse.java
index 8b256862..c91250ba 100644
--- a/src/main/java/com/google/cloud/spanner/pgadapter/wireoutput/CopyDataResponse.java
+++ b/src/main/java/com/google/cloud/spanner/pgadapter/wireoutput/CopyDataResponse.java
@@ -18,6 +18,9 @@ import static com.google.cloud.spanner.pgadapter.statements.CopyToStatement.COPY
 
 import com.google.api.core.InternalApi;
 import com.google.cloud.spanner.pgadapter.ProxyServer.DataFormat;
+import com.google.cloud.spanner.pgadapter.error.PGExceptionFactory;
+import com.google.cloud.spanner.pgadapter.error.SQLState;
+import com.google.cloud.spanner.pgadapter.utils.Converter;
 import java.io.DataOutputStream;
 import java.nio.charset.StandardCharsets;
 
@@ -32,9 +35,9 @@ public class CopyDataResponse extends WireOutput {
 
   private final ResponseType responseType;
   private final DataFormat format;
-  private final byte[][] binaryData;
   private final String stringData;
   private final char rowTerminator;
+  private final Converter converter;
 
   /** Creates a {@link CopyDataResponse} message containing the fixed binary COPY header. */
   @InternalApi
@@ -52,9 +55,9 @@ public class CopyDataResponse extends WireOutput {
     super(output, length + 4);
     this.responseType = responseType;
     this.format = DataFormat.POSTGRESQL_BINARY;
-    this.binaryData = null;
     this.stringData = null;
     this.rowTerminator = 0;
+    this.converter = null;
   }
 
   public CopyDataResponse(DataOutputStream output, String data, char rowTerminator) {
@@ -63,16 +66,23 @@ public class CopyDataResponse extends WireOutput {
     this.format = DataFormat.POSTGRESQL_TEXT;
     this.stringData = data;
     this.rowTerminator = rowTerminator;
-    this.binaryData = null;
+    this.converter = null;
   }
 
-  public CopyDataResponse(DataOutputStream output, int length, byte[][] data) {
-    super(output, length + 4);
+  public CopyDataResponse(DataOutputStream output, Converter converter) {
+    super(output, 0);
     this.responseType = ResponseType.ROW;
     this.format = DataFormat.POSTGRESQL_BINARY;
     this.stringData = null;
     this.rowTerminator = 0;
-    this.binaryData = data;
+    this.converter = converter;
+  }
+
+  public void send(boolean flush) throws Exception {
+    if (converter != null) {
+      this.length = 4 + converter.convertResultSetRowToDataRowResponse();
+    }
+    super.send(flush);
   }
 
   @Override
@@ -87,16 +97,11 @@ public class CopyDataResponse extends WireOutput {
         this.outputStream.writeInt(0); // header extension area length
       } else if (this.responseType == ResponseType.TRAILER) {
         this.outputStream.writeShort(-1);
+      } else if (this.converter != null) {
+        this.converter.writeBuffer(this.outputStream);
       } else {
-        this.outputStream.writeShort(this.binaryData.length);
-        for (int col = 0; col < this.binaryData.length; col++) {
-          if (this.binaryData[col] == null) {
-            this.outputStream.writeInt(-1);
-          } else {
-            this.outputStream.writeInt(this.binaryData[col].length);
-            this.outputStream.write(this.binaryData[col]);
-          }
-        }
+        // This should not happen.
+        throw PGExceptionFactory.newPGException("Invalid CopyDataResponse", SQLState.InternalError);
       }
     }
   }
diff --git a/src/main/java/com/google/cloud/spanner/pgadapter/wireoutput/DataRowResponse.java b/src/main/java/com/google/cloud/spanner/pgadapter/wireoutput/DataRowResponse.java
index 8eca0993..4f2fc5ab 100644
--- a/src/main/java/com/google/cloud/spanner/pgadapter/wireoutput/DataRowResponse.java
+++ b/src/main/java/com/google/cloud/spanner/pgadapter/wireoutput/DataRowResponse.java
@@ -15,64 +15,29 @@
 package com.google.cloud.spanner.pgadapter.wireoutput;
 
 import com.google.api.core.InternalApi;
-import com.google.cloud.spanner.ResultSet;
-import com.google.cloud.spanner.pgadapter.ConnectionHandler.QueryMode;
-import com.google.cloud.spanner.pgadapter.ProxyServer.DataFormat;
-import com.google.cloud.spanner.pgadapter.metadata.OptionsMetadata;
-import com.google.cloud.spanner.pgadapter.statements.IntermediateStatement;
 import com.google.cloud.spanner.pgadapter.utils.Converter;
 import java.io.DataOutputStream;
-import java.util.ArrayList;
-import java.util.List;
 
 /** Sends to the client specific row contents. */
 @InternalApi
 public class DataRowResponse extends WireOutput {
-
   private static final int HEADER_LENGTH = 4;
-  private static final int COLUMN_NUMBER_LENGTH = 2;
-  private static final int COLUMN_SIZE_LENGTH = 4;
 
-  private final ResultSet resultSet;
-  private final List<byte[]> columns;
+  private final Converter converter;
 
-  public DataRowResponse(
-      DataOutputStream output,
-      IntermediateStatement statement,
-      ResultSet resultSet,
-      OptionsMetadata options,
-      QueryMode mode) {
+  public DataRowResponse(DataOutputStream output, Converter converter) {
     super(output, 0);
-    this.resultSet = resultSet;
-    this.columns = new ArrayList<>(this.resultSet.getColumnCount());
-    this.length = HEADER_LENGTH + COLUMN_NUMBER_LENGTH;
-    // TODO profile this with immense rows/column count to see if it's an aread of optimization.
-    for (int column_index = 0; /* column indices start at 0 */
-        column_index < this.resultSet.getColumnCount();
-        column_index++) {
-      length += COLUMN_SIZE_LENGTH;
-      if (resultSet.isNull(column_index)) {
-        columns.add(null);
-      } else {
-        DataFormat format = DataFormat.getDataFormat(column_index, statement, mode, options);
-        byte[] column = Converter.convertToPG(this.resultSet, column_index, format);
-        length += column.length;
-        columns.add(column);
-      }
-    }
+    this.converter = converter;
+  }
+
+  public void send(boolean flush) throws Exception {
+    this.length = HEADER_LENGTH + converter.convertResultSetRowToDataRowResponse();
+    super.send(flush);
   }
 
   @Override
   protected void sendPayload() throws Exception {
-    this.outputStream.writeShort(this.resultSet.getColumnCount());
-    for (byte[] column_value : columns) {
-      if (column_value == null) {
-        this.outputStream.writeInt(-1);
-      } else {
-        this.outputStream.writeInt(column_value.length);
-        this.outputStream.write(column_value);
-      }
-    }
+    converter.writeBuffer(this.outputStream);
   }
 
   @Override
diff --git a/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/ControlMessage.java b/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/ControlMessage.java
index 354c3110..fb71b2c9 100644
--- a/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/ControlMessage.java
+++ b/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/ControlMessage.java
@@ -38,6 +38,7 @@ import com.google.cloud.spanner.pgadapter.error.Severity;
 import com.google.cloud.spanner.pgadapter.metadata.SendResultSetState;
 import com.google.cloud.spanner.pgadapter.statements.BackendConnection.PartitionQueryResult;
 import com.google.cloud.spanner.pgadapter.statements.IntermediateStatement;
+import com.google.cloud.spanner.pgadapter.utils.Converter;
 import com.google.cloud.spanner.pgadapter.wireoutput.CommandCompleteResponse;
 import com.google.cloud.spanner.pgadapter.wireoutput.EmptyQueryResponse;
 import com.google.cloud.spanner.pgadapter.wireoutput.ErrorResponse;
@@ -388,6 +389,7 @@ public abstract class ControlMessage extends WireMessage {
   static final class SendResultSetRunnable implements Callable<Long> {
     private final IntermediateStatement describedResult;
     private ResultSet resultSet;
+    private Converter converter;
     private final BatchReadOnlyTransaction batchReadOnlyTransaction;
     private final Partition partition;
     private final long maxRows;
@@ -428,6 +430,12 @@ public abstract class ControlMessage extends WireMessage {
         boolean hasData) {
       this.describedResult = describedResult;
       this.resultSet = resultSet;
+      this.converter =
+          new Converter(
+              describedResult,
+              mode,
+              describedResult.getConnectionHandler().getServer().getOptions(),
+              resultSet);
       this.batchReadOnlyTransaction = null;
       this.partition = null;
       this.maxRows = maxRows;
@@ -462,6 +470,12 @@ public abstract class ControlMessage extends WireMessage {
         // will be cleaned up at a later moment.
         try {
           resultSet = batchReadOnlyTransaction.execute(partition);
+          converter =
+              new Converter(
+                  describedResult,
+                  mode,
+                  describedResult.getConnectionHandler().getServer().getOptions(),
+                  resultSet);
           hasData = resultSet.next();
         } catch (Throwable t) {
           if (includePrefix) {
@@ -490,7 +504,7 @@ public abstract class ControlMessage extends WireMessage {
         if (Thread.interrupted()) {
           throw PGExceptionFactory.newQueryCancelledException();
         }
-        WireOutput wireOutput = describedResult.createDataRowResponse(resultSet, mode);
+        WireOutput wireOutput = describedResult.createDataRowResponse(converter);
         synchronized (describedResult) {
           wireOutput.send(false);
         }


Test patch:
diff --git a/src/main/java/com/google/cloud/spanner/pgadapter/statements/IntermediateStatement.java b/src/main/java/com/google/cloud/spanner/pgadapter/statements/IntermediateStatement.java
index effdeee5..604e4ab4 100644
--- a/src/main/java/com/google/cloud/spanner/pgadapter/statements/IntermediateStatement.java
+++ b/src/main/java/com/google/cloud/spanner/pgadapter/statements/IntermediateStatement.java
@@ -30,11 +30,11 @@ import com.google.cloud.spanner.connection.PostgreSQLStatementParser;
 import com.google.cloud.spanner.connection.StatementResult;
 import com.google.cloud.spanner.connection.StatementResult.ResultType;
 import com.google.cloud.spanner.pgadapter.ConnectionHandler;
-import com.google.cloud.spanner.pgadapter.ConnectionHandler.QueryMode;
 import com.google.cloud.spanner.pgadapter.error.PGExceptionFactory;
 import com.google.cloud.spanner.pgadapter.metadata.DescribeMetadata;
 import com.google.cloud.spanner.pgadapter.metadata.OptionsMetadata;
 import com.google.cloud.spanner.pgadapter.statements.BackendConnection.NoResult;
+import com.google.cloud.spanner.pgadapter.utils.Converter;
 import com.google.cloud.spanner.pgadapter.wireoutput.DataRowResponse;
 import com.google.cloud.spanner.pgadapter.wireoutput.WireOutput;
 import java.io.DataOutputStream;
@@ -333,8 +333,8 @@ public class IntermediateStatement {
     return EMPTY_WIRE_OUTPUT_ARRAY;
   }
 
-  public WireOutput createDataRowResponse(ResultSet resultSet, QueryMode mode) {
-    return new DataRowResponse(this.outputStream, this, resultSet, this.options, mode);
+  public WireOutput createDataRowResponse(Converter converter) {
+    return new DataRowResponse(this.outputStream, converter);
   }
 
   public WireOutput[] createResultSuffix() {