feat: add Proto Columns support in JDBC (#1252)

* feat: add code changes to support DML for Proto Columns

* feat: support DML for Proto Columns

* feat: add untyped null when inserting null value in array of proto columns

* feat: add code changes to support DDL for Proto Columns

* feat: code changes to support getArray and getResultSet in JdbcArray for Proto columns

* feat: add unit tests for DML and DQL

* feat: add integration tests for Proto Columns DDL

* feat: add integration tests for Proto columns DML and DQL

* feat: lint format

* feat: code refactoring to throw exceptions and handle null values in JdbcArray

* feat: update tests to validate null in JdbcArray

* feat: Integration test refactoring

* fix: add copyright header

* feat: update junit assertions

* feat: move array conversion logic for protos to seperate methods in JdbcTypeConverter

* feat: add review suggestions to JdbcArray

* feat: add review suggestion

* feat: add review suggestions in JdbcParameterStore file

* feat: add untyped null integration test

* feat: add inter compatibilty and lint fix

* feat: update schema and base64 protodescriptors files

* feat: nit

* chore: update java-spanner version

* chore: lint fix

* chore: skip tests on graalvm

* chore: nit fixes

clirr-ignored-differences.xml

@@ -57,5 +57,20 @@
     <className>com/google/cloud/spanner/jdbc/CloudSpannerJdbcConnection</className>
     <method>void setMaxPartitions(int)</method>
   </difference>
+  <difference>
+    <differenceType>7012</differenceType>
+    <className>com/google/cloud/spanner/jdbc/CloudSpannerJdbcConnection</className>
+    <method>byte[] getProtoDescriptors()</method>
+  </difference>
+  <difference>
+    <differenceType>7012</differenceType>
+    <className>com/google/cloud/spanner/jdbc/CloudSpannerJdbcConnection</className>
+    <method>void setProtoDescriptors(byte[])</method>
+  </difference>
+  <difference>
+    <differenceType>7012</differenceType>
+    <className>com/google/cloud/spanner/jdbc/CloudSpannerJdbcConnection</className>
+    <method>void setProtoDescriptors(java.io.InputStream)</method>
+  </difference>
 
 </differences>

src/main/java/com/google/cloud/spanner/jdbc/AbstractJdbcWrapper.java

@@ -48,6 +48,7 @@ static int extractColumnType(Type type) {
       case BOOL:
         return Types.BOOLEAN;
       case BYTES:
+      case PROTO:
         return Types.BINARY;
       case DATE:
         return Types.DATE;
@@ -56,6 +57,7 @@ static int extractColumnType(Type type) {
       case FLOAT64:
         return Types.DOUBLE;
       case INT64:
+      case ENUM:
         return Types.BIGINT;
       case NUMERIC:
       case PG_NUMERIC:
@@ -145,6 +147,7 @@ static String getClassName(Type type) {
       case BOOL:
         return Boolean.class.getName();
       case BYTES:
+      case PROTO:
         return byte[].class.getName();
       case DATE:
         return Date.class.getName();
@@ -153,6 +156,7 @@ static String getClassName(Type type) {
       case FLOAT64:
         return Double.class.getName();
       case INT64:
+      case ENUM:
         return Long.class.getName();
       case NUMERIC:
       case PG_NUMERIC:
@@ -168,6 +172,7 @@ static String getClassName(Type type) {
           case BOOL:
             return Boolean[].class.getName();
           case BYTES:
+          case PROTO:
             return byte[][].class.getName();
           case DATE:
             return Date[].class.getName();
@@ -176,6 +181,7 @@ static String getClassName(Type type) {
           case FLOAT64:
             return Double[].class.getName();
           case INT64:
+          case ENUM:
             return Long[].class.getName();
           case NUMERIC:
           case PG_NUMERIC:

src/main/java/com/google/cloud/spanner/jdbc/CloudSpannerJdbcConnection.java

@@ -28,11 +28,14 @@
 import com.google.cloud.spanner.connection.AutocommitDmlMode;
 import com.google.cloud.spanner.connection.SavepointSupport;
 import com.google.cloud.spanner.connection.TransactionMode;
+import java.io.IOException;
+import java.io.InputStream;
 import java.sql.Connection;
 import java.sql.ResultSet;
 import java.sql.SQLException;
 import java.sql.Timestamp;
 import java.util.Iterator;
+import javax.annotation.Nonnull;
 
 /**
  * JDBC connection with a number of additional Cloud Spanner specific methods. JDBC connections that
@@ -448,4 +451,35 @@ Iterator<com.google.cloud.spanner.jdbc.TransactionRetryListener> getTransactionR
    */
   Iterator<com.google.cloud.spanner.connection.TransactionRetryListener>
       getTransactionRetryListenersFromConnection() throws SQLException;
+
+  /**
+   * Sets the proto descriptors to use for the next DDL statement (single or batch) that will be
+   * executed. The proto descriptor is automatically cleared after the statement is executed.
+   *
+   * @param protoDescriptors The proto descriptors to use with the next DDL statement (single or
+   *     batch) that will be executed on this connection.
+   */
+  default void setProtoDescriptors(@Nonnull byte[] protoDescriptors) throws SQLException {
+    throw new UnsupportedOperationException();
+  }
+
+  /**
+   * Sets the proto descriptors to use for the next DDL statement (single or batch) that will be
+   * executed. The proto descriptor is automatically cleared after the statement is executed.
+   *
+   * @param protoDescriptors The proto descriptors to use with the next DDL statement (single or
+   *     batch) that will be executed on this connection.
+   */
+  default void setProtoDescriptors(@Nonnull InputStream protoDescriptors)
+      throws SQLException, IOException {
+    throw new UnsupportedOperationException();
+  }
+
+  /**
+   * @return The proto descriptor that will be used with the next DDL statement (single or batch)
+   *     that is executed on this connection.
+   */
+  default byte[] getProtoDescriptors() throws SQLException {
+    throw new UnsupportedOperationException();
+  }
 }

src/main/java/com/google/cloud/spanner/jdbc/JdbcArray.java

@@ -24,6 +24,11 @@
 import com.google.cloud.spanner.Value;
 import com.google.cloud.spanner.ValueBinder;
 import com.google.common.collect.ImmutableList;
+import com.google.protobuf.AbstractMessage;
+import com.google.protobuf.Descriptors;
+import com.google.protobuf.Descriptors.Descriptor;
+import com.google.protobuf.Message;
+import com.google.protobuf.ProtocolMessageEnum;
 import com.google.rpc.Code;
 import java.math.BigDecimal;
 import java.sql.Array;
@@ -78,7 +83,16 @@ static JdbcArray createArray(JdbcDataType type, List<?> elements) {
   private JdbcArray(JdbcDataType type, Object[] elements) throws SQLException {
     this.type = type;
     if (elements != null) {
-      this.data = java.lang.reflect.Array.newInstance(type.getJavaClass(), elements.length);
+      if ((type.getCode() == Type.Code.PROTO
+              && AbstractMessage[].class.isAssignableFrom(elements.getClass()))
+          || (type.getCode() == Type.Code.ENUM
+              && ProtocolMessageEnum[].class.isAssignableFrom(elements.getClass()))) {
+        this.data =
+            java.lang.reflect.Array.newInstance(
+                elements.getClass().getComponentType(), elements.length);
+      } else {
+        this.data = java.lang.reflect.Array.newInstance(type.getJavaClass(), elements.length);
+      }
       try {
         System.arraycopy(elements, 0, this.data, 0, elements.length);
       } catch (Exception e) {
@@ -138,9 +152,17 @@ public Object getArray(long index, int count) throws SQLException {
   @Override
   public Object getArray(long index, int count, Map<String, Class<?>> map) throws SQLException {
     checkFree();
-    if (data != null) {
-      Object res = java.lang.reflect.Array.newInstance(type.getJavaClass(), count);
-      System.arraycopy(data, (int) index - 1, res, 0, count);
+    if (this.data != null) {
+      Object res;
+      if ((this.type.getCode() == Type.Code.PROTO
+              && AbstractMessage[].class.isAssignableFrom(this.data.getClass()))
+          || (this.type.getCode() == Type.Code.ENUM
+              && ProtocolMessageEnum[].class.isAssignableFrom(this.data.getClass()))) {
+        res = java.lang.reflect.Array.newInstance(this.data.getClass().getComponentType(), count);
+      } else {
+        res = java.lang.reflect.Array.newInstance(this.type.getJavaClass(), count);
+      }
+      System.arraycopy(this.data, (int) index - 1, res, 0, count);
       return res;
     }
     return null;
@@ -167,24 +189,37 @@ public ResultSet getResultSet(long startIndex, int count) throws SQLException {
     JdbcPreconditions.checkArgument(startIndex >= 1L, "Start index must be >= 1");
     JdbcPreconditions.checkArgument(count >= 0, "Count must be >= 0");
     checkFree();
+    Type spannerTypeForProto = getSpannerTypeForProto();
+    Type spannerType =
+        spannerTypeForProto == null ? this.type.getSpannerType() : spannerTypeForProto;
+
     ImmutableList.Builder<Struct> rows = ImmutableList.builder();
     int added = 0;
-    if (data != null) {
+    if (this.data != null) {
       // Note that array index in JDBC is base-one.
       for (int index = (int) startIndex;
-          added < count && index <= ((Object[]) data).length;
+          added < count && index <= ((Object[]) this.data).length;
           index++) {
-        Object value = ((Object[]) data)[index - 1];
+        Object value = ((Object[]) this.data)[index - 1];
         ValueBinder<Struct.Builder> binder =
             Struct.newBuilder().set("INDEX").to(index).set("VALUE");
         Struct.Builder builder;
-        switch (type.getCode()) {
+        switch (this.type.getCode()) {
           case BOOL:
             builder = binder.to((Boolean) value);
             break;
           case BYTES:
             builder = binder.to(ByteArray.copyFrom((byte[]) value));
             break;
+          case PROTO:
+            if (value == null && AbstractMessage[].class.isAssignableFrom(this.data.getClass())) {
+              builder = binder.to((ByteArray) null, spannerType.getProtoTypeFqn());
+            } else if (value instanceof AbstractMessage) {
+              builder = binder.to((AbstractMessage) value);
+            } else {
+              builder = binder.to(value != null ? ByteArray.copyFrom((byte[]) value) : null);
+            }
+            break;
           case DATE:
             builder = binder.to(JdbcTypeConverter.toGoogleDate((Date) value));
             break;
@@ -197,6 +232,16 @@ public ResultSet getResultSet(long startIndex, int count) throws SQLException {
           case INT64:
             builder = binder.to((Long) value);
             break;
+          case ENUM:
+            if (value == null
+                && ProtocolMessageEnum[].class.isAssignableFrom(this.data.getClass())) {
+              builder = binder.to((Long) null, spannerType.getProtoTypeFqn());
+            } else if (value instanceof ProtocolMessageEnum) {
+              builder = binder.to((ProtocolMessageEnum) value);
+            } else {
+              builder = binder.to((Long) value);
+            }
+            break;
           case NUMERIC:
             builder = binder.to((BigDecimal) value);
             break;
@@ -217,7 +262,8 @@ public ResultSet getResultSet(long startIndex, int count) throws SQLException {
           default:
             throw new SQLFeatureNotSupportedException(
                 String.format(
-                    "Array of type %s cannot be converted to a ResultSet", type.getCode().name()));
+                    "Array of type %s cannot be converted to a ResultSet",
+                    this.type.getCode().name()));
         }
         rows.add(builder.build());
         added++;
@@ -226,14 +272,54 @@ public ResultSet getResultSet(long startIndex, int count) throws SQLException {
         }
       }
     }
+
     return JdbcResultSet.of(
         ResultSets.forRows(
             Type.struct(
-                StructField.of("INDEX", Type.int64()),
-                StructField.of("VALUE", type.getSpannerType())),
+                StructField.of("INDEX", Type.int64()), StructField.of("VALUE", spannerType)),
             rows.build()));
   }
 
+  // Returns null if the type is not a PROTO or ENUM
+  private Type getSpannerTypeForProto() throws SQLException {
+    Type spannerType = null;
+    if (this.data != null) {
+      if (this.type.getCode() == Type.Code.PROTO
+          && AbstractMessage[].class.isAssignableFrom(this.data.getClass())) {
+        spannerType = createSpannerProtoType();
+      } else if (this.type.getCode() == Type.Code.ENUM
+          && ProtocolMessageEnum[].class.isAssignableFrom(this.data.getClass())) {
+        spannerType = createSpannerProtoEnumType();
+      }
+    }
+    return spannerType;
+  }
+
+  private Type createSpannerProtoType() throws SQLException {
+    Class<?> componentType = this.data.getClass().getComponentType();
+    try {
+      Message.Builder builder =
+          (Message.Builder) componentType.getMethod("newBuilder").invoke(null);
+      Descriptor msgDescriptor = builder.getDescriptorForType();
+      return Type.proto(msgDescriptor.getFullName());
+    } catch (Exception e) {
+      throw JdbcSqlExceptionFactory.of(
+          "Error occurred when getting proto message descriptor from data", Code.UNKNOWN, e);
+    }
+  }
+
+  private Type createSpannerProtoEnumType() throws SQLException {
+    Class<?> componentType = this.data.getClass().getComponentType();
+    try {
+      Descriptors.EnumDescriptor enumDescriptor =
+          (Descriptors.EnumDescriptor) componentType.getMethod("getDescriptor").invoke(null);
+      return Type.protoEnum(enumDescriptor.getFullName());
+    } catch (Exception e) {
+      throw JdbcSqlExceptionFactory.of(
+          "Error occurred when getting proto enum descriptor from data", Code.UNKNOWN, e);
+    }
+  }
+
   @Override
   public ResultSet getResultSet(long index, int count, Map<String, Class<?>> map)
       throws SQLException {

src/main/java/com/google/cloud/spanner/jdbc/JdbcConnection.java

@@ -20,6 +20,7 @@
 import static com.google.cloud.spanner.jdbc.JdbcStatement.isNullOrEmpty;
 
 import com.google.api.client.util.Preconditions;
+import com.google.cloud.ByteArray;
 import com.google.cloud.spanner.CommitResponse;
 import com.google.cloud.spanner.DatabaseId;
 import com.google.cloud.spanner.Mutation;
@@ -38,6 +39,8 @@
 import io.opentelemetry.api.OpenTelemetry;
 import io.opentelemetry.api.common.Attributes;
 import io.opentelemetry.api.common.AttributesBuilder;
+import java.io.IOException;
+import java.io.InputStream;
 import java.sql.Array;
 import java.sql.Blob;
 import java.sql.Clob;
@@ -867,4 +870,34 @@ public Iterator<TransactionRetryListener> getTransactionRetryListeners() throws
     checkClosed();
     return getSpannerConnection().getTransactionRetryListeners();
   }
+
+  @Override
+  public void setProtoDescriptors(@Nonnull byte[] protoDescriptors) throws SQLException {
+    Preconditions.checkNotNull(protoDescriptors);
+    checkClosed();
+    try {
+      getSpannerConnection().setProtoDescriptors(protoDescriptors);
+    } catch (SpannerException e) {
+      throw JdbcSqlExceptionFactory.of(e);
+    }
+  }
+
+  @Override
+  public void setProtoDescriptors(@Nonnull InputStream protoDescriptors)
+      throws SQLException, IOException {
+    Preconditions.checkNotNull(protoDescriptors);
+    checkClosed();
+    try {
+      getSpannerConnection()
+          .setProtoDescriptors(ByteArray.copyFrom(protoDescriptors).toByteArray());
+    } catch (SpannerException e) {
+      throw JdbcSqlExceptionFactory.of(e);
+    }
+  }
+
+  @Override
+  public byte[] getProtoDescriptors() throws SQLException {
+    checkClosed();
+    return getSpannerConnection().getProtoDescriptors();
+  }
 }