Repository: GoogleCloudPlatform/pgadapter
Pull request #114 (2022-04-18 06:16:21)

fix: remember parameter types of parsed statement
Remember the parameter types that were given during the PARSE phase
of the extended query protocol, and use this when the statement is
described and when parameters are bound.


Patch:
diff --git a/src/main/java/com/google/cloud/spanner/pgadapter/statements/IntermediatePreparedStatement.java b/src/main/java/com/google/cloud/spanner/pgadapter/statements/IntermediatePreparedStatement.java
index eab578a4..d1855fc4 100644
--- a/src/main/java/com/google/cloud/spanner/pgadapter/statements/IntermediatePreparedStatement.java
+++ b/src/main/java/com/google/cloud/spanner/pgadapter/statements/IntermediatePreparedStatement.java
@@ -26,6 +26,7 @@ import com.google.cloud.spanner.pgadapter.metadata.DescribeStatementMetadata;
 import com.google.cloud.spanner.pgadapter.metadata.OptionsMetadata;
 import com.google.cloud.spanner.pgadapter.parsers.Parser;
 import com.google.cloud.spanner.pgadapter.parsers.Parser.FormatCode;
+import java.util.Arrays;
 import java.util.List;
 import java.util.Set;
 import org.postgresql.core.Oid;
@@ -35,7 +36,7 @@ import org.postgresql.core.Oid;
  */
 public class IntermediatePreparedStatement extends IntermediateStatement {
 
-  protected List<Integer> parameterDataTypes;
+  protected int[] parameterDataTypes;
   protected Statement statement;
 
   public IntermediatePreparedStatement(
@@ -53,18 +54,18 @@ public class IntermediatePreparedStatement extends IntermediateStatement {
    * @return The type of the item specified.
    */
   private int parseType(byte[][] parameters, int index) throws IllegalArgumentException {
-    if (this.parameterDataTypes.size() > index) {
-      return this.parameterDataTypes.get(index);
+    if (this.parameterDataTypes.length > index) {
+      return this.parameterDataTypes[index];
     } else {
       return Oid.UNSPECIFIED;
     }
   }
 
-  public List<Integer> getParameterDataTypes() {
+  public int[] getParameterDataTypes() {
     return this.parameterDataTypes;
   }
 
-  public void setParameterDataTypes(List<Integer> parameterDataTypes) {
+  public void setParameterDataTypes(int[] parameterDataTypes) {
     this.parameterDataTypes = parameterDataTypes;
   }
 
@@ -127,12 +128,15 @@ public class IntermediatePreparedStatement extends IntermediateStatement {
 
   /**
    * Returns the parameter types in the SQL string of this statement. The current implementation
-   * always returns Oid.UNSPECIFIED for all parameters, as we have no way to actually determine the
-   * parameter types.
+   * always returns any parameters that may have been specified in the PARSE message, and
+   * OID.Unspecified for all other parameters.
    */
   private int[] getParameterTypes() {
     Set<String> parameters =
         PARSER.getQueryParameters(this.parsedStatement.getSqlWithoutComments());
-    return new int[parameters.size()];
+    if (this.parameterDataTypes == null) {
+      return new int[parameters.size()];
+    }
+    return Arrays.copyOf(this.parameterDataTypes, parameters.size());
   }
 }
diff --git a/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/DescribeMessage.java b/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/DescribeMessage.java
index 238d7afa..f7442d20 100644
--- a/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/DescribeMessage.java
+++ b/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/DescribeMessage.java
@@ -29,9 +29,9 @@ public class DescribeMessage extends ControlMessage {
 
   protected static final char IDENTIFIER = 'D';
 
-  private PreparedType type;
-  private String name;
-  private IntermediateStatement statement;
+  private final PreparedType type;
+  private final String name;
+  private final IntermediateStatement statement;
 
   public DescribeMessage(ConnectionHandler connection) throws Exception {
     super(connection);
diff --git a/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/ParseMessage.java b/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/ParseMessage.java
index 90146072..a6e612a7 100644
--- a/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/ParseMessage.java
+++ b/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/ParseMessage.java
@@ -33,7 +33,7 @@ public class ParseMessage extends ControlMessage {
 
   private final String name;
   private final IntermediatePreparedStatement statement;
-  private final ImmutableList<Integer> parameterDataTypes;
+  private final int[] parameterDataTypes;
 
   public ParseMessage(ConnectionHandler connection) throws Exception {
     super(connection);
@@ -41,11 +41,10 @@ public class ParseMessage extends ControlMessage {
     ParsedStatement parsedStatement = PARSER.parse(Statement.of(this.readString()));
     ImmutableList.Builder<Integer> builder = ImmutableList.builder();
     short numberOfParameters = this.inputStream.readShort();
+    this.parameterDataTypes = new int[numberOfParameters];
     for (int i = 0; i < numberOfParameters; i++) {
-      int type = this.inputStream.readInt();
-      builder.add(type);
+      parameterDataTypes[i] = this.inputStream.readInt();
     }
-    this.parameterDataTypes = builder.build();
     this.statement =
         new IntermediatePreparedStatement(
             connection.getServer().getOptions(),


Test patch:
diff --git a/src/test/java/com/google/cloud/spanner/pgadapter/JdbcMockServerTest.java b/src/test/java/com/google/cloud/spanner/pgadapter/JdbcMockServerTest.java
index 486caf0d..f5a510e8 100644
--- a/src/test/java/com/google/cloud/spanner/pgadapter/JdbcMockServerTest.java
+++ b/src/test/java/com/google/cloud/spanner/pgadapter/JdbcMockServerTest.java
@@ -40,7 +40,6 @@ import java.sql.PreparedStatement;
 import java.sql.ResultSet;
 import java.sql.SQLException;
 import java.sql.Types;
-import java.time.Instant;
 import java.time.LocalDate;
 import java.time.LocalDateTime;
 import java.time.OffsetDateTime;
@@ -53,6 +52,7 @@ import org.junit.BeforeClass;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
+import org.postgresql.PGStatement;
 import org.postgresql.core.Oid;
 
 @RunWith(JUnit4.class)
@@ -161,9 +161,8 @@ public class JdbcMockServerTest extends AbstractMockServerTest {
         preparedStatement.setBytes(++index, "test".getBytes(StandardCharsets.UTF_8));
         preparedStatement.setDouble(++index, 3.14d);
         preparedStatement.setBigDecimal(++index, new BigDecimal("6.626"));
-        preparedStatement.setTimestamp(
-            ++index, java.sql.Timestamp.from(Instant.from(zonedDateTime)));
-        preparedStatement.setDate(++index, java.sql.Date.valueOf(LocalDate.of(2022, 3, 29)));
+        preparedStatement.setObject(++index, zonedDateTime);
+        preparedStatement.setObject(++index, LocalDate.of(2022, 3, 29));
         preparedStatement.setString(++index, "test");
         try (ResultSet resultSet = preparedStatement.executeQuery()) {
           assertTrue(resultSet.next());
@@ -406,4 +405,99 @@ public class JdbcMockServerTest extends AbstractMockServerTest {
     assertEquals(1, updateDatabaseDdlRequests.get(0).getStatementsCount());
     assertEquals(sql, updateDatabaseDdlRequests.get(0).getStatements(0));
   }
+
+  @Test
+  public void testPreparedStatement() throws SQLException {
+    mockSpanner.putStatementResult(
+        StatementResult.update(
+            Statement.newBuilder(
+                    "insert into all_types "
+                        + "(col_bigint, col_bool, col_bytea, col_float8, col_int, col_numeric, col_timestamptz, col_date, col_varchar) "
+                        + "values ($1, $2, $3, $4, $5, $6, $7, $8, $9)")
+                .bind("p1")
+                .to(2L)
+                .bind("p2")
+                .to((Boolean) null)
+                .bind("p3")
+                .to((ByteArray) null)
+                .bind("p4")
+                .to((Double) null)
+                .bind("p5")
+                .to((Long) null)
+                .bind("p6")
+                .to(com.google.cloud.spanner.Value.pgNumeric(null))
+                .bind("p7")
+                .to((Timestamp) null)
+                .bind("p8")
+                .to((Date) null)
+                .bind("p9")
+                .to((String) null)
+                .build(),
+            1L));
+    mockSpanner.putStatementResult(
+        StatementResult.update(
+            Statement.newBuilder(
+                    "insert into all_types "
+                        + "(col_bigint, col_bool, col_bytea, col_float8, col_int, col_numeric, col_timestamptz, col_date, col_varchar) "
+                        + "values ($1, $2, $3, $4, $5, $6, $7, $8, $9)")
+                .bind("p1")
+                .to(1L)
+                .bind("p2")
+                .to(true)
+                .bind("p3")
+                .to(ByteArray.copyFrom("test"))
+                .bind("p4")
+                .to(3.14d)
+                .bind("p5")
+                .to(100L)
+                .bind("p6")
+                .to(com.google.cloud.spanner.Value.pgNumeric("6.626"))
+                .bind("p7")
+                .to(Timestamp.parseTimestamp("2022-02-16T13:18:02.123457000"))
+                .bind("p8")
+                .to(Date.parseDate("2022-03-29"))
+                .bind("p9")
+                .to("test")
+                .build(),
+            1L));
+
+    OffsetDateTime zonedDateTime =
+        LocalDateTime.of(2022, 2, 16, 13, 18, 2, 123456789).atOffset(ZoneOffset.UTC);
+    try (Connection connection = DriverManager.getConnection(createUrl())) {
+      try (PreparedStatement statement =
+          connection.prepareStatement(
+              "insert into all_types "
+                  + "(col_bigint, col_bool, col_bytea, col_float8, col_int, col_numeric, col_timestamptz, col_date, col_varchar) "
+                  + "values (?, ?, ?, ?, ?, ?, ?, ?, ?)")) {
+        PGStatement pgStatement = statement.unwrap(PGStatement.class);
+        pgStatement.setPrepareThreshold(1);
+
+        int index = 0;
+        statement.setLong(++index, 1L);
+        statement.setBoolean(++index, true);
+        statement.setBytes(++index, "test".getBytes(StandardCharsets.UTF_8));
+        statement.setDouble(++index, 3.14d);
+        statement.setInt(++index, 100);
+        statement.setBigDecimal(++index, new BigDecimal("6.626"));
+        statement.setObject(++index, zonedDateTime);
+        statement.setObject(++index, LocalDate.of(2022, 3, 29));
+        statement.setString(++index, "test");
+
+        assertEquals(1, statement.executeUpdate());
+
+        index = 0;
+        statement.setLong(++index, 2);
+        statement.setNull(++index, Types.BOOLEAN);
+        statement.setNull(++index, Types.BINARY);
+        statement.setNull(++index, Types.DOUBLE);
+        statement.setNull(++index, Types.INTEGER);
+        statement.setNull(++index, Types.NUMERIC);
+        statement.setNull(++index, Types.TIMESTAMP_WITH_TIMEZONE);
+        statement.setNull(++index, Types.DATE);
+        statement.setNull(++index, Types.VARCHAR);
+
+        assertEquals(1, statement.executeUpdate());
+      }
+    }
+  }
 }
diff --git a/src/test/java/com/google/cloud/spanner/pgadapter/ProtocolTest.java b/src/test/java/com/google/cloud/spanner/pgadapter/ProtocolTest.java
index dcbca61b..34dd88e3 100644
--- a/src/test/java/com/google/cloud/spanner/pgadapter/ProtocolTest.java
+++ b/src/test/java/com/google/cloud/spanner/pgadapter/ProtocolTest.java
@@ -269,7 +269,7 @@ public class ProtocolTest {
             parameterCount,
             parameters);
 
-    List<Integer> expectedParameterDataTypes = Collections.singletonList(1002);
+    int[] expectedParameterDataTypes = new int[] {1002};
     String expectedSQL = "SELECT * FROM users WHERE name = $1";
     String expectedMessageName = "some statement";
 
@@ -288,7 +288,7 @@ public class ProtocolTest {
     assertEquals(ParseMessage.class, message.getClass());
     assertEquals(expectedMessageName, ((ParseMessage) message).getName());
     assertEquals(expectedSQL, ((ParseMessage) message).getStatement().getSql());
-    assertEquals(
+    assertArrayEquals(
         expectedParameterDataTypes,
         ((ParseMessage) message).getStatement().getParameterDataTypes());
 
@@ -330,7 +330,7 @@ public class ProtocolTest {
             parameterCount,
             parameters);
 
-    List<Integer> expectedParameterDataTypes = Collections.singletonList(1002);
+    int[] expectedParameterDataTypes = new int[] {1002};
     String expectedSQL = "SELECT * FROM users WHERE name = $1";
     String expectedMessageName = "some statement";
 
@@ -349,7 +349,7 @@ public class ProtocolTest {
     assertEquals(ParseMessage.class, message.getClass());
     assertEquals(expectedMessageName, ((ParseMessage) message).getName());
     assertEquals(expectedSQL, ((ParseMessage) message).getStatement().getSql());
-    assertEquals(
+    assertArrayEquals(
         expectedParameterDataTypes,
         ((ParseMessage) message).getStatement().getParameterDataTypes());
 
@@ -392,7 +392,7 @@ public class ProtocolTest {
             parameterCount,
             parameters);
 
-    List<Integer> expectedParameterDataTypes = Collections.singletonList(0);
+    int[] expectedParameterDataTypes = new int[] {0};
     String expectedSQL = "SELECT * FROM users WHERE name = $1";
     String expectedMessageName = "some statement";
 
@@ -411,7 +411,7 @@ public class ProtocolTest {
     assertEquals(ParseMessage.class, message.getClass());
     assertEquals(expectedMessageName, ((ParseMessage) message).getName());
     assertEquals(expectedSQL, ((ParseMessage) message).getStatement().getSql());
-    assertEquals(
+    assertArrayEquals(
         expectedParameterDataTypes,
         ((ParseMessage) message).getStatement().getParameterDataTypes());
 
@@ -439,7 +439,7 @@ public class ProtocolTest {
         Bytes.concat(
             messageMetadata, length, statementName.getBytes(), payload.getBytes(), intToBytes(0));
 
-    List<Integer> expectedParameterDataTypes = Collections.emptyList();
+    int[] expectedParameterDataTypes = new int[0];
     String expectedSQL = "SELECT * FROM users WHERE name = $1";
     String expectedMessageName = "some statement";
 
@@ -458,7 +458,7 @@ public class ProtocolTest {
     assertEquals(ParseMessage.class, message.getClass());
     assertEquals(expectedMessageName, ((ParseMessage) message).getName());
     assertEquals(expectedSQL, ((ParseMessage) message).getStatement().getSql());
-    assertEquals(
+    assertArrayEquals(
         expectedParameterDataTypes,
         ((ParseMessage) message).getStatement().getParameterDataTypes());
 
diff --git a/src/test/java/com/google/cloud/spanner/pgadapter/StatementTest.java b/src/test/java/com/google/cloud/spanner/pgadapter/StatementTest.java
index 792281d7..d64b833b 100644
--- a/src/test/java/com/google/cloud/spanner/pgadapter/StatementTest.java
+++ b/src/test/java/com/google/cloud/spanner/pgadapter/StatementTest.java
@@ -54,7 +54,6 @@ import java.io.File;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collections;
-import java.util.List;
 import org.junit.AfterClass;
 import org.junit.Rule;
 import org.junit.Test;
@@ -249,7 +248,7 @@ public class StatementTest {
   @Test
   public void testPreparedStatement() {
     String sqlStatement = "SELECT * FROM users WHERE age > $2 AND age < $3 AND name = $1";
-    List<Integer> parameterDataTypes = Arrays.asList(Oid.VARCHAR, Oid.INT8, Oid.INT4);
+    int[] parameterDataTypes = new int[] {Oid.VARCHAR, Oid.INT8, Oid.INT4};
 
     Statement statement =
         Statement.newBuilder(sqlStatement)
@@ -286,7 +285,7 @@ public class StatementTest {
   @Test
   public void testPreparedStatementIllegalTypeThrowsException() {
     String sqlStatement = "SELECT * FROM users WHERE metadata = $1";
-    List<Integer> parameterDataTypes = Collections.singletonList(Oid.JSON);
+    int[] parameterDataTypes = new int[] {Oid.JSON};
 
     IntermediatePreparedStatement intermediateStatement =
         new IntermediatePreparedStatement(options, parse(sqlStatement), connection);