Repository: GoogleCloudPlatform/pgadapter
Pull request #254 (2022-07-02 06:03:32)

fix: COPY null values caused NullPointerException
`COPY` operations that included `null` values would cause a `NullPointerException`.


Patch:
diff --git a/src/main/java/com/google/cloud/spanner/pgadapter/utils/MutationWriter.java b/src/main/java/com/google/cloud/spanner/pgadapter/utils/MutationWriter.java
index 986990e9..dc1579a1 100644
--- a/src/main/java/com/google/cloud/spanner/pgadapter/utils/MutationWriter.java
+++ b/src/main/java/com/google/cloud/spanner/pgadapter/utils/MutationWriter.java
@@ -30,6 +30,7 @@ import com.google.cloud.spanner.Value;
 import com.google.cloud.spanner.connection.Connection;
 import com.google.cloud.spanner.connection.StatementResult;
 import com.google.cloud.spanner.pgadapter.statements.BackendConnection.UpdateCount;
+import com.google.common.annotations.VisibleForTesting;
 import com.google.common.collect.ImmutableList;
 import com.google.common.util.concurrent.FutureCallback;
 import com.google.common.util.concurrent.Futures;
@@ -42,7 +43,6 @@ import java.io.IOException;
 import java.io.InputStreamReader;
 import java.io.PipedInputStream;
 import java.io.PipedOutputStream;
-import java.io.PrintWriter;
 import java.io.Reader;
 import java.nio.charset.StandardCharsets;
 import java.sql.Timestamp;
@@ -130,7 +130,6 @@ public class MutationWriter implements Callable<StatementResult>, Closeable {
   private final int pipeBufferSize;
   private final CSVFormat format;
   private final CSVParser parser;
-  private PrintWriter errorFileWriter;
   private final PipedOutputStream payload = new PipedOutputStream();
   private final AtomicBoolean commit = new AtomicBoolean(false);
   private final AtomicBoolean rollback = new AtomicBoolean(false);
@@ -187,6 +186,7 @@ public class MutationWriter implements Callable<StatementResult>, Closeable {
                     "copy_in_pipe_buffer_size", String.valueOf(DEFAULT_PIPE_BUFFER_SIZE))),
             1024);
     this.format = format;
+    // TODO(b/237831799): Support binary format for COPY.
     this.parser = createParser();
   }
 
@@ -431,6 +431,11 @@ public class MutationWriter implements Callable<StatementResult>, Closeable {
   static int calculateSize(Mutation mutation) {
     int size = 0;
     for (Value value : mutation.getValues()) {
+      if (value.isNull()) {
+        size++;
+        continue;
+      }
+
       switch (value.getType().getCode()) {
         case BOOL:
           size++;
@@ -503,7 +508,8 @@ public class MutationWriter implements Callable<StatementResult>, Closeable {
     return size;
   }
 
-  private Mutation buildMutation(CSVRecord record) {
+  @VisibleForTesting
+  Mutation buildMutation(CSVRecord record) {
     TimestampUtils timestampUtils = new TimestampUtils(false, () -> null);
     WriteBuilder builder;
     // The default is to use Insert, but PGAdapter also supports InsertOrUpdate. This can be very
@@ -520,7 +526,7 @@ public class MutationWriter implements Callable<StatementResult>, Closeable {
       TypeCode columnType = this.tableColumns.get(columnName);
       String recordValue = "";
       try {
-        recordValue = record.get(columnName).trim();
+        recordValue = record.get(columnName);
         switch (columnType) {
           case STRING:
             builder.set(columnName).to(recordValue);
@@ -529,30 +535,42 @@ public class MutationWriter implements Callable<StatementResult>, Closeable {
             builder.set(columnName).to(Value.json(recordValue));
             break;
           case BOOL:
-            builder.set(columnName).to(Boolean.parseBoolean(recordValue));
+            builder
+                .set(columnName)
+                .to(recordValue == null ? null : Boolean.parseBoolean(recordValue));
             break;
           case INT64:
-            builder.set(columnName).to(Long.parseLong(recordValue));
+            builder.set(columnName).to(recordValue == null ? null : Long.parseLong(recordValue));
             break;
           case FLOAT64:
-            builder.set(columnName).to(Double.parseDouble(recordValue));
+            builder
+                .set(columnName)
+                .to(recordValue == null ? null : Double.parseDouble(recordValue));
             break;
           case NUMERIC:
             builder.set(columnName).to(Value.pgNumeric(recordValue));
             break;
           case BYTES:
-            if (recordValue.startsWith("\\x")) {
+            if (recordValue == null) {
+              builder.set(columnName).to((ByteArray) null);
+            } else if (recordValue.startsWith("\\x")) {
               builder
                   .set(columnName)
                   .to(ByteArray.copyFrom(Hex.decodeHex(recordValue.substring(2))));
+            } else {
+              throw SpannerExceptionFactory.newSpannerException(
+                  ErrorCode.INVALID_ARGUMENT,
+                  "COPY only supports the Hex format for bytea columns");
             }
             break;
           case DATE:
-            builder.set(columnName).to(Date.parseDate(recordValue));
+            builder.set(columnName).to(recordValue == null ? null : Date.parseDate(recordValue));
             break;
           case TIMESTAMP:
             Timestamp timestamp = timestampUtils.toTimestamp(null, recordValue);
-            builder.set(columnName).to(com.google.cloud.Timestamp.of(timestamp));
+            builder
+                .set(columnName)
+                .to(timestamp == null ? null : com.google.cloud.Timestamp.of(timestamp));
             break;
         }
       } catch (NumberFormatException | DateTimeParseException e) {


Test patch:
diff --git a/src/test/java/com/google/cloud/spanner/pgadapter/CopyInMockServerTest.java b/src/test/java/com/google/cloud/spanner/pgadapter/CopyInMockServerTest.java
index 7362789e..702297b1 100644
--- a/src/test/java/com/google/cloud/spanner/pgadapter/CopyInMockServerTest.java
+++ b/src/test/java/com/google/cloud/spanner/pgadapter/CopyInMockServerTest.java
@@ -24,6 +24,7 @@ import com.google.cloud.spanner.MockSpannerServiceImpl.SimulatedExecutionTime;
 import com.google.cloud.spanner.MockSpannerServiceImpl.StatementResult;
 import com.google.cloud.spanner.pgadapter.metadata.OptionsMetadata;
 import com.google.protobuf.ListValue;
+import com.google.protobuf.NullValue;
 import com.google.protobuf.Value;
 import com.google.spanner.v1.CommitRequest;
 import com.google.spanner.v1.Mutation;
@@ -129,6 +130,37 @@ public class CopyInMockServerTest extends AbstractMockServerTest {
     assertEquals(3, mutation.getInsert().getValuesCount());
   }
 
+  @Test
+  public void testCopyInNullValues() throws SQLException, IOException {
+    setupCopyInformationSchemaResults();
+
+    try (Connection connection = DriverManager.getConnection(createUrl())) {
+      CopyManager copyManager = new CopyManager(connection.unwrap(BaseConnection.class));
+      copyManager.copyIn(
+          "COPY users FROM STDIN;", new StringReader("5\t\\N\t\\N\n6\t\\N\t\\N\n7\t\\N\t\\N\n"));
+
+      // Verify that we can use the connection for normal queries.
+      try (ResultSet resultSet = connection.createStatement().executeQuery("SELECT 1")) {
+        assertTrue(resultSet.next());
+        assertEquals(1L, resultSet.getLong(1));
+        assertFalse(resultSet.next());
+      }
+    }
+
+    List<CommitRequest> commitRequests = mockSpanner.getRequestsOfType(CommitRequest.class);
+    assertEquals(1, commitRequests.size());
+    CommitRequest commitRequest = commitRequests.get(0);
+    assertEquals(1, commitRequest.getMutationsCount());
+
+    Mutation mutation = commitRequest.getMutations(0);
+    assertEquals(OperationCase.INSERT, mutation.getOperationCase());
+    assertEquals(3, mutation.getInsert().getValuesCount());
+    Value nullValue = Value.newBuilder().setNullValue(NullValue.NULL_VALUE).build();
+    assertEquals("5", mutation.getInsert().getValues(0).getValues(0).getStringValue());
+    assertEquals(nullValue, mutation.getInsert().getValues(0).getValues(1));
+    assertEquals(nullValue, mutation.getInsert().getValues(0).getValues(2));
+  }
+
   @Test
   public void testCopyInWithExplicitTransaction() throws SQLException, IOException {
     setupCopyInformationSchemaResults();
@@ -201,6 +233,35 @@ public class CopyInMockServerTest extends AbstractMockServerTest {
     }
   }
 
+  @Test
+  public void testCopyIn_Nulls() throws SQLException, IOException {
+    setupCopyInformationSchemaResults();
+
+    try (Connection connection = DriverManager.getConnection(createUrl())) {
+      PGConnection pgConnection = connection.unwrap(PGConnection.class);
+      CopyManager copyManager = pgConnection.getCopyAPI();
+      long copyCount =
+          copyManager.copyIn(
+              "copy all_types from stdin;",
+              new FileInputStream("./src/test/resources/all_types_data_nulls.txt"));
+      assertEquals(1L, copyCount);
+    }
+
+    List<CommitRequest> commitRequests = mockSpanner.getRequestsOfType(CommitRequest.class);
+    assertEquals(1, commitRequests.size());
+    CommitRequest commitRequest = commitRequests.get(0);
+    assertEquals(1, commitRequest.getMutationsCount());
+
+    Mutation mutation = commitRequest.getMutations(0);
+    assertEquals(OperationCase.INSERT, mutation.getOperationCase());
+    assertEquals(1, mutation.getInsert().getValuesCount());
+    Value nullValue = Value.newBuilder().setNullValue(NullValue.NULL_VALUE).build();
+    // The first column is not null, as it is the primary key.
+    for (int col = 1; col < mutation.getInsert().getColumnsCount(); col++) {
+      assertEquals(nullValue, mutation.getInsert().getValues(0).getValues(col));
+    }
+  }
+
   @Test
   public void testCopyIn_Large_FailsWhenAtomic() throws Exception {
     setupCopyInformationSchemaResults();
diff --git a/src/test/java/com/google/cloud/spanner/pgadapter/ITJdbcTest.java b/src/test/java/com/google/cloud/spanner/pgadapter/ITJdbcTest.java
index 6e6f1b92..0dcf5d1f 100644
--- a/src/test/java/com/google/cloud/spanner/pgadapter/ITJdbcTest.java
+++ b/src/test/java/com/google/cloud/spanner/pgadapter/ITJdbcTest.java
@@ -494,6 +494,34 @@ public class ITJdbcTest implements IntegrationTest {
     }
   }
 
+  @Test
+  public void testCopyIn_Nulls() throws SQLException, IOException {
+    // Empty all data in the table.
+    String databaseId = database.getId().getDatabase();
+    testEnv.write(databaseId, Collections.singleton(Mutation.delete("all_types", KeySet.all())));
+
+    try (Connection connection = DriverManager.getConnection(getConnectionUrl())) {
+      PGConnection pgConnection = connection.unwrap(PGConnection.class);
+      CopyManager copyManager = pgConnection.getCopyAPI();
+      long copyCount =
+          copyManager.copyIn(
+              "copy all_types from stdin;",
+              new FileInputStream("./src/test/resources/all_types_data_nulls.txt"));
+      assertEquals(1L, copyCount);
+
+      // Verify that there is 1 row in the table, and that the values of all columns except the
+      // primary key are null.
+      try (ResultSet resultSet =
+          connection.createStatement().executeQuery("select * from all_types")) {
+        assertTrue(resultSet.next());
+        for (int col = 2; col <= resultSet.getMetaData().getColumnCount(); col++) {
+          assertNull(String.format("Col %d should be null", col), resultSet.getObject(col));
+        }
+        assertFalse(resultSet.next());
+      }
+    }
+  }
+
   @Test
   public void testCopyIn_Large_FailsWhenAtomic() throws SQLException {
     // Empty all data in the table.
diff --git a/src/test/java/com/google/cloud/spanner/pgadapter/utils/MutationWriterTest.java b/src/test/java/com/google/cloud/spanner/pgadapter/utils/MutationWriterTest.java
index 4edfe88d..8de12e1d 100644
--- a/src/test/java/com/google/cloud/spanner/pgadapter/utils/MutationWriterTest.java
+++ b/src/test/java/com/google/cloud/spanner/pgadapter/utils/MutationWriterTest.java
@@ -38,6 +38,7 @@ import com.google.cloud.spanner.connection.StatementResult;
 import com.google.cloud.spanner.pgadapter.utils.MutationWriter.CopyTransactionMode;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.ImmutableSet;
 import com.google.spanner.v1.TypeCode;
 import java.io.IOException;
 import java.io.InputStreamReader;
@@ -45,6 +46,7 @@ import java.io.OutputStreamWriter;
 import java.io.PipedInputStream;
 import java.io.PipedOutputStream;
 import java.io.Reader;
+import java.io.StringReader;
 import java.nio.charset.StandardCharsets;
 import java.util.Iterator;
 import java.util.List;
@@ -381,6 +383,10 @@ public class MutationWriterTest {
 
   @Test
   public void testCalculateSize() {
+    assertEquals(
+        1,
+        calculateSize(
+            Mutation.newInsertBuilder("my_table").set("my_col").to((String) null).build()));
     assertEquals(
         8, calculateSize(Mutation.newInsertBuilder("my_table").set("my_col").to(1L).build()));
     assertEquals(
@@ -493,4 +499,37 @@ public class MutationWriterTest {
                 .toJsonArray(ImmutableList.of("foo", "a"))
                 .build()));
   }
+
+  @Test
+  public void testBuildMutationNulls() throws IOException {
+    // TODO(b/237831798): Add support for ARRAY in COPY operations.
+    ImmutableSet<TypeCode> unsupportedTypeCodes =
+        ImmutableSet.of(
+            TypeCode.TYPE_CODE_UNSPECIFIED, TypeCode.ARRAY, TypeCode.STRUCT, TypeCode.UNRECOGNIZED);
+
+    Connection connection = mock(Connection.class);
+    for (TypeCode typeCode : TypeCode.values()) {
+      if (unsupportedTypeCodes.contains(typeCode)) {
+        continue;
+      }
+
+      MutationWriter mutationWriter =
+          new MutationWriter(
+              CopyTransactionMode.ImplicitAtomic,
+              connection,
+              "my_table",
+              ImmutableMap.of("col", typeCode),
+              0,
+              CSVFormat.POSTGRESQL_TEXT,
+              false);
+      CSVParser parser =
+          CSVParser.parse(
+              new StringReader("col\n\\N\n"), CSVFormat.POSTGRESQL_TEXT.withFirstRecordAsHeader());
+      CSVRecord record = parser.getRecords().get(0);
+
+      Mutation mutation = mutationWriter.buildMutation(record);
+
+      assertEquals(String.format("Type code: %s", typeCode), 1, mutation.asMap().size());
+    }
+  }
 }
diff --git a/src/test/resources/all_types_data_nulls.txt b/src/test/resources/all_types_data_nulls.txt
new file mode 100644
index 00000000..b304e0ac
--- /dev/null
+++ b/src/test/resources/all_types_data_nulls.txt
@@ -0,0 +1,1 @@
+98641651896415341	\N	\N	\N	\N	\N	\N	\N	\N