feat: add support for Proto Columns to Connection API

google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ClientSideStatementValueConverters.java

@@ -16,6 +16,7 @@
 
 package com.google.cloud.spanner.connection;
 
+import com.google.cloud.ByteArray;
 import com.google.cloud.spanner.ErrorCode;
 import com.google.cloud.spanner.Options.RpcPriority;
 import com.google.cloud.spanner.SpannerExceptionFactory;
@@ -28,6 +29,11 @@
 import com.google.protobuf.Duration;
 import com.google.protobuf.util.Durations;
 import com.google.spanner.v1.RequestOptions.Priority;
+import java.io.File;
+import java.io.FileInputStream;
+import java.io.IOException;
+import java.io.InputStream;
+import java.util.Base64;
 import java.util.EnumSet;
 import java.util.HashMap;
 import java.util.Locale;
@@ -494,4 +500,52 @@ public String convert(String value) {
       return value.substring(7).trim();
     }
   }
+
+  /** Converter for converting strings to byte[] */
+  static class ProtoDescriptorsConverter implements ClientSideStatementValueConverter<byte[]> {
+
+    public ProtoDescriptorsConverter(String allowedValues) {}
+
+    @Override
+    public Class<byte[]> getParameterClass() {
+      return byte[].class;
+    }
+
+    @Override
+    public byte[] convert(String value) {
+      if (value.length() == 0 || value.equalsIgnoreCase("null")) {
+        return null;
+      }
+      return Base64.getDecoder().decode(value);
+    }
+  }
+
+  /** Converter for converting strings that take in file path as input to byte[] */
+  static class ProtoDescriptorsFileConverter implements ClientSideStatementValueConverter<byte[]> {
+
+    public ProtoDescriptorsFileConverter(String allowedValues) {}
+
+    @Override
+    public Class<byte[]> getParameterClass() {
+      return byte[].class;
+    }
+
+    @Override
+    public byte[] convert(String filePath) {
+      if (filePath != null && filePath.length() > 0) {
+        try {
+          File protoDescriptorsFile = new File(filePath);
+          if (!protoDescriptorsFile.isFile()) {
+            throw new IOException("File does not exist.");
+          }
+          InputStream pdStream = new FileInputStream(protoDescriptorsFile);
+          return ByteArray.copyFrom(pdStream).toByteArray();
+        } catch (Exception e) {
+          throw SpannerExceptionFactory.newSpannerException(
+              ErrorCode.INVALID_ARGUMENT, e.getMessage());
+        }
+      }
+      return null;
+    }
+  }
 }

google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/Connection.java

@@ -42,6 +42,7 @@
 import java.util.Iterator;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.TimeUnit;
+import javax.annotation.Nonnull;
 
 /**
  * Internal connection API for Google Cloud Spanner. This interface may introduce breaking changes
@@ -382,6 +383,25 @@ default String getStatementTag() {
     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 byte[] protoDescriptors) {
+    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() {
+    throw new UnsupportedOperationException();
+  }
+
   /**
    * @return <code>true</code> if this connection will automatically retry read/write transactions
    *     that abort. This method may only be called when the connection is in read/write

google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ConnectionImpl.java

@@ -63,6 +63,7 @@
 import java.util.concurrent.ThreadFactory;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
+import javax.annotation.Nonnull;
 import org.threeten.bp.Instant;
 
 /** Implementation for {@link Connection}, the generic Spanner connection API (not JDBC). */
@@ -85,7 +86,7 @@ private LeakedConnectionException() {
     }
   }
 
-  private volatile LeakedConnectionException leakedException;;
+  private volatile LeakedConnectionException leakedException;
   private final SpannerPool spannerPool;
   private AbstractStatementParser statementParser;
   /**
@@ -221,6 +222,8 @@ static UnitOfWorkType of(TransactionMode transactionMode) {
   private String transactionTag;
   private String statementTag;
 
+  private byte[] protoDescriptors;
+
   /** Create a connection and register it in the SpannerPool. */
   ConnectionImpl(ConnectionOptions options) {
     Preconditions.checkNotNull(options);
@@ -278,6 +281,7 @@ Spanner getSpanner() {
   private DdlClient createDdlClient() {
     return DdlClient.newBuilder()
         .setDatabaseAdminClient(spanner.getDatabaseAdminClient())
+        .setProjectId(options.getProjectId())
         .setInstanceId(options.getInstanceId())
         .setDatabaseName(options.getDatabaseName())
         .build();
@@ -623,6 +627,21 @@ public void setStatementTag(String tag) {
     this.statementTag = tag;
   }
 
+  @Override
+  public byte[] getProtoDescriptors() {
+    ConnectionPreconditions.checkState(!isClosed(), CLOSED_ERROR_MSG);
+    return protoDescriptors;
+  }
+
+  @Override
+  public void setProtoDescriptors(@Nonnull byte[] protoDescriptors) {
+    Preconditions.checkNotNull(protoDescriptors);
+    ConnectionPreconditions.checkState(!isClosed(), CLOSED_ERROR_MSG);
+    ConnectionPreconditions.checkState(
+        !isBatchActive(), "Proto descriptors cannot be set when a batch is active");
+    this.protoDescriptors = protoDescriptors;
+  }
+
   /**
    * Throws an {@link SpannerException} with code {@link ErrorCode#FAILED_PRECONDITION} if the
    * current state of this connection does not allow changing the setting for retryAbortsInternally.
@@ -1379,6 +1398,7 @@ UnitOfWork createNewUnitOfWork() {
           .setReturnCommitStats(returnCommitStats)
           .setStatementTimeout(statementTimeout)
           .withStatementExecutor(statementExecutor)
+          .setProtoDescriptors(protoDescriptors)
           .build();
     } else {
       switch (getUnitOfWorkType()) {
@@ -1421,6 +1441,7 @@ UnitOfWork createNewUnitOfWork() {
               .setDatabaseClient(dbClient)
               .setStatementTimeout(statementTimeout)
               .withStatementExecutor(statementExecutor)
+              .setProtoDescriptors(protoDescriptors)
               .build();
         default:
       }
@@ -1444,7 +1465,11 @@ private void popUnitOfWorkFromTransactionStack() {
   }
 
   private ApiFuture<Void> executeDdlAsync(CallType callType, ParsedStatement ddl) {
-    return getCurrentUnitOfWorkOrStartNewUnitOfWork().executeDdlAsync(callType, ddl);
+    ApiFuture<Void> result =
+        getCurrentUnitOfWorkOrStartNewUnitOfWork().executeDdlAsync(callType, ddl);
+    // reset proto descriptors after executing a DDL statement
+    this.protoDescriptors = null;
+    return result;
   }
 
   @Override
@@ -1535,6 +1560,10 @@ public ApiFuture<long[]> runBatchAsync() {
       }
       return ApiFutures.immediateFuture(new long[0]);
     } finally {
+      if (isDdlBatchActive()) {
+        // reset proto descriptors after executing a DDL batch
+        this.protoDescriptors = null;
+      }
       this.batchMode = BatchMode.NONE;
       setDefaultTransactionOptions();
     }

google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ConnectionStatementExecutor.java

@@ -124,5 +124,9 @@ StatementResult statementSetPgSessionCharacteristicsTransactionMode(
 
   StatementResult statementShowTransactionIsolationLevel();
 
+  StatementResult statementSetProtoDescriptors(byte[] protoDescriptors);
+
+  StatementResult statementShowProtoDescriptors();
+
   StatementResult statementExplain(String sql);
 }

google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ConnectionStatementExecutorImpl.java

@@ -28,6 +28,7 @@
 import static com.google.cloud.spanner.connection.StatementResult.ClientSideStatementType.SET_DELAY_TRANSACTION_START_UNTIL_FIRST_WRITE;
 import static com.google.cloud.spanner.connection.StatementResult.ClientSideStatementType.SET_OPTIMIZER_STATISTICS_PACKAGE;
 import static com.google.cloud.spanner.connection.StatementResult.ClientSideStatementType.SET_OPTIMIZER_VERSION;
+import static com.google.cloud.spanner.connection.StatementResult.ClientSideStatementType.SET_PROTO_DESCRIPTORS;
 import static com.google.cloud.spanner.connection.StatementResult.ClientSideStatementType.SET_READONLY;
 import static com.google.cloud.spanner.connection.StatementResult.ClientSideStatementType.SET_READ_ONLY_STALENESS;
 import static com.google.cloud.spanner.connection.StatementResult.ClientSideStatementType.SET_RETRY_ABORTS_INTERNALLY;
@@ -45,6 +46,7 @@
 import static com.google.cloud.spanner.connection.StatementResult.ClientSideStatementType.SHOW_DELAY_TRANSACTION_START_UNTIL_FIRST_WRITE;
 import static com.google.cloud.spanner.connection.StatementResult.ClientSideStatementType.SHOW_OPTIMIZER_STATISTICS_PACKAGE;
 import static com.google.cloud.spanner.connection.StatementResult.ClientSideStatementType.SHOW_OPTIMIZER_VERSION;
+import static com.google.cloud.spanner.connection.StatementResult.ClientSideStatementType.SHOW_PROTO_DESCRIPTORS;
 import static com.google.cloud.spanner.connection.StatementResult.ClientSideStatementType.SHOW_READONLY;
 import static com.google.cloud.spanner.connection.StatementResult.ClientSideStatementType.SHOW_READ_ONLY_STALENESS;
 import static com.google.cloud.spanner.connection.StatementResult.ClientSideStatementType.SHOW_READ_TIMESTAMP;
@@ -493,6 +495,20 @@ public StatementResult statementShowTransactionIsolationLevel() {
     return resultSet("transaction_isolation", "serializable", SHOW_TRANSACTION_ISOLATION_LEVEL);
   }
 
+  @Override
+  public StatementResult statementSetProtoDescriptors(byte[] protoDescriptors) {
+    getConnection().setProtoDescriptors(protoDescriptors);
+    return noResult(SET_PROTO_DESCRIPTORS);
+  }
+
+  @Override
+  public StatementResult statementShowProtoDescriptors() {
+    return resultSet(
+        String.format("%sPROTO_DESCRIPTORS", getNamespace(connection.getDialect())),
+        getConnection().getProtoDescriptors(),
+        SHOW_PROTO_DESCRIPTORS);
+  }
+
   private String processQueryPlan(PlanNode planNode) {
     StringBuilder planNodeDescription = new StringBuilder(" : { ");
     com.google.protobuf.Struct metadata = planNode.getMetadata();

google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/DdlBatch.java

@@ -58,10 +58,12 @@ class DdlBatch extends AbstractBaseUnitOfWork {
   private final DatabaseClient dbClient;
   private final List<String> statements = new ArrayList<>();
   private UnitOfWorkState state = UnitOfWorkState.STARTED;
+  private final byte[] protoDescriptors;
 
   static class Builder extends AbstractBaseUnitOfWork.Builder<Builder, DdlBatch> {
     private DdlClient ddlClient;
     private DatabaseClient dbClient;
+    private byte[] protoDescriptors;
 
     private Builder() {}
 
@@ -77,6 +79,11 @@ Builder setDatabaseClient(DatabaseClient client) {
       return this;
     }
 
+    Builder setProtoDescriptors(byte[] protoDescriptors) {
+      this.protoDescriptors = protoDescriptors;
+      return this;
+    }
+
     @Override
     DdlBatch build() {
       Preconditions.checkState(ddlClient != null, "No DdlClient specified");
@@ -93,6 +100,7 @@ private DdlBatch(Builder builder) {
     super(builder);
     this.ddlClient = builder.ddlClient;
     this.dbClient = builder.dbClient;
+    this.protoDescriptors = builder.protoDescriptors;
   }
 
   @Override
@@ -239,7 +247,7 @@ public ApiFuture<long[]> runBatchAsync(CallType callType) {
         () -> {
           try {
             OperationFuture<Void, UpdateDatabaseDdlMetadata> operation =
-                ddlClient.executeDdl(statements);
+                ddlClient.executeDdl(statements, protoDescriptors);
             try {
               // Wait until the operation has finished.
               getWithStatementTimeout(operation, RUN_BATCH_STATEMENT);

google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/DdlClient.java

@@ -19,6 +19,7 @@
 import com.google.api.gax.longrunning.OperationFuture;
 import com.google.cloud.spanner.Database;
 import com.google.cloud.spanner.DatabaseAdminClient;
+import com.google.cloud.spanner.DatabaseId;
 import com.google.cloud.spanner.Dialect;
 import com.google.cloud.spanner.ErrorCode;
 import com.google.cloud.spanner.SpannerExceptionFactory;
@@ -35,11 +36,13 @@
  */
 class DdlClient {
   private final DatabaseAdminClient dbAdminClient;
+  private final String projectId;
   private final String instanceId;
   private final String databaseName;
 
   static class Builder {
     private DatabaseAdminClient dbAdminClient;
+    private String projectId;
     private String instanceId;
     private String databaseName;
 
@@ -51,6 +54,13 @@ Builder setDatabaseAdminClient(DatabaseAdminClient client) {
       return this;
     }
 
+    Builder setProjectId(String projectId) {
+      Preconditions.checkArgument(
+          !Strings.isNullOrEmpty(projectId), "Empty projectId is not allowed");
+      this.projectId = projectId;
+      return this;
+    }
+
     Builder setInstanceId(String instanceId) {
       Preconditions.checkArgument(
           !Strings.isNullOrEmpty(instanceId), "Empty instanceId is not allowed");
@@ -67,6 +77,7 @@ Builder setDatabaseName(String name) {
 
     DdlClient build() {
       Preconditions.checkState(dbAdminClient != null, "No DatabaseAdminClient specified");
+      Preconditions.checkState(!Strings.isNullOrEmpty(projectId), "No ProjectId specified");
       Preconditions.checkState(!Strings.isNullOrEmpty(instanceId), "No InstanceId specified");
       Preconditions.checkArgument(
           !Strings.isNullOrEmpty(databaseName), "No database name specified");
@@ -80,6 +91,7 @@ static Builder newBuilder() {
 
   private DdlClient(Builder builder) {
     this.dbAdminClient = builder.dbAdminClient;
+    this.projectId = builder.projectId;
     this.instanceId = builder.instanceId;
     this.databaseName = builder.databaseName;
   }
@@ -92,17 +104,24 @@ OperationFuture<Database, CreateDatabaseMetadata> executeCreateDatabase(
   }
 
   /** Execute a single DDL statement. */
-  OperationFuture<Void, UpdateDatabaseDdlMetadata> executeDdl(String ddl) {
-    return executeDdl(Collections.singletonList(ddl));
+  OperationFuture<Void, UpdateDatabaseDdlMetadata> executeDdl(String ddl, byte[] protoDescriptors) {
+    return executeDdl(Collections.singletonList(ddl), protoDescriptors);
   }
 
   /** Execute a list of DDL statements as one operation. */
-  OperationFuture<Void, UpdateDatabaseDdlMetadata> executeDdl(List<String> statements) {
+  OperationFuture<Void, UpdateDatabaseDdlMetadata> executeDdl(
+      List<String> statements, byte[] protoDescriptors) {
     if (statements.stream().anyMatch(DdlClient::isCreateDatabaseStatement)) {
       throw SpannerExceptionFactory.newSpannerException(
           ErrorCode.INVALID_ARGUMENT, "CREATE DATABASE is not supported in a DDL batch");
     }
-    return dbAdminClient.updateDatabaseDdl(instanceId, databaseName, statements, null);
+    Database.Builder dbBuilder =
+        dbAdminClient.newDatabaseBuilder(DatabaseId.of(projectId, instanceId, databaseName));
+    if (protoDescriptors != null) {
+      dbBuilder.setProtoDescriptors(protoDescriptors);
+    }
+    Database db = dbBuilder.build();
+    return dbAdminClient.updateDatabaseDdl(db, statements, null);
   }
 
   /** Returns true if the statement is a `CREATE DATABASE ...` statement. */