[Cloud_Spanner_to_GCS_Avro] Refactoring InformationSchemaScannerIT (#3915)

darshan-sj · web-flow · commit 99150c5d5b00 · 2026-06-18T19:04:23.000+05:30
* Corrected and optimized InformationSchemaScannerIT

* Formatting changes

* Correcting SpannerResourceManagerTest

* Formatting changes

* Correcting build failure

* Fixing a flaky test - MutationKeyEncoderTest

* Addressing comments
diff --git a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/spanner/SpannerResourceManager.java b/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/spanner/SpannerResourceManager.java
@@ -23,10 +23,12 @@
 import static org.apache.beam.it.gcp.spanner.utils.SpannerResourceManagerUtils.generateInstanceId;
 
 import com.google.auth.Credentials;
+import com.google.cloud.spanner.BatchClient;
 import com.google.cloud.spanner.Database;
 import com.google.cloud.spanner.DatabaseAdminClient;
 import com.google.cloud.spanner.DatabaseClient;
 import com.google.cloud.spanner.DatabaseId;
+import com.google.cloud.spanner.DatabaseInfo;
 import com.google.cloud.spanner.Dialect;
 import com.google.cloud.spanner.Instance;
 import com.google.cloud.spanner.InstanceAdminClient;
@@ -49,7 +51,11 @@
 import com.google.monitoring.v3.Aggregation.Aligner;
 import com.google.monitoring.v3.TimeInterval;
 import com.google.protobuf.Timestamp;
+import com.google.spanner.admin.instance.v1.CreateInstancePartitionRequest;
 import com.google.spanner.admin.instance.v1.Instance.Edition;
+import com.google.spanner.admin.instance.v1.InstanceConfigName;
+import com.google.spanner.admin.instance.v1.InstanceName;
+import com.google.spanner.admin.instance.v1.InstancePartition;
 import dev.failsafe.Failsafe;
 import dev.failsafe.RetryPolicy;
 import java.time.Duration;
@@ -100,6 +106,11 @@ public final class SpannerResourceManager implements ResourceManager {
 
   private boolean hasInstance = false;
   private boolean hasDatabase = false;
+  private final String instancePartitionId;
+  private final String instancePartitionConfig;
+  private final com.google.cloud.spanner.admin.instance.v1.InstanceAdminClient
+      v1InstanceAdminClient;
+  private boolean hasInstancePartition = false;
 
   private final String projectId;
   private final String instanceId;
@@ -114,6 +125,7 @@ public final class SpannerResourceManager implements ResourceManager {
   private final InstanceAdminClient instanceAdminClient;
   private final DatabaseAdminClient databaseAdminClient;
   private final int nodeCount;
+  private final byte[] protoDescriptors;
   private Timestamp startTime;
   private MonitoringClient monitoringClient;
   private final boolean suppressVerboseLogs;
@@ -164,7 +176,11 @@ private SpannerResourceManager(Builder builder) {
     this.instanceAdminClient = spanner.getInstanceAdminClient();
     this.databaseAdminClient = spanner.getDatabaseAdminClient();
     this.nodeCount = builder.nodeCount;
+    this.protoDescriptors = builder.protoDescriptors;
     this.monitoringClient = builder.monitoringClient;
+    this.instancePartitionId = builder.instancePartitionId;
+    this.instancePartitionConfig = builder.instancePartitionConfig;
+    this.v1InstanceAdminClient = spanner.createInstanceAdminClient();
   }
 
   public static Builder builder(String testId, String projectId, String region) {
@@ -181,33 +197,62 @@ private synchronized void maybeCreateInstance() {
     if (usingStaticInstance) {
       LOG.info("Not creating Spanner instance - reusing static {}", instanceId);
       hasInstance = true;
-      return;
-    }
-
-    if (hasInstance) {
-      return;
+    } else if (!hasInstance) {
+      LOG.info("Creating instance {} in project {}.", instanceId, projectId);
+      try {
+        InstanceInfo instanceInfo =
+            InstanceInfo.newBuilder(InstanceId.of(projectId, instanceId))
+                .setInstanceConfigId(InstanceConfigId.of(projectId, "regional-" + region))
+                .setDisplayName(instanceId)
+                .setEdition(Edition.ENTERPRISE_PLUS) // Needed by Full Text Search.
+                .setNodeCount(nodeCount)
+                .build();
+
+        // Retry creation if there's a quota error
+        Instance instance =
+            Failsafe.with(
+                    retryOnQuotaException(5, Duration.ofMinutes(1), Duration.ofMinutes(2), 0.5))
+                .get(() -> instanceAdminClient.createInstance(instanceInfo).get());
+
+        hasInstance = true;
+        LOG.info("Successfully created instance {}: {}.", instanceId, instance.getState());
+      } catch (Exception e) {
+        cleanupAll();
+        throw new SpannerResourceManagerException("Failed to create instance.", e);
+      }
     }
 
-    LOG.info("Creating instance {} in project {}.", instanceId, projectId);
-    try {
-      InstanceInfo instanceInfo =
-          InstanceInfo.newBuilder(InstanceId.of(projectId, instanceId))
-              .setInstanceConfigId(InstanceConfigId.of(projectId, "regional-" + region))
-              .setDisplayName(instanceId)
-              .setEdition(Edition.ENTERPRISE_PLUS) // Needed by Full Text Search.
-              .setNodeCount(nodeCount)
-              .build();
-
-      // Retry creation if there's a quota error
-      Instance instance =
-          Failsafe.with(retryOnQuotaException(5, Duration.ofMinutes(1), Duration.ofMinutes(2), 0.5))
-              .get(() -> instanceAdminClient.createInstance(instanceInfo).get());
-
-      hasInstance = true;
-      LOG.info("Successfully created instance {}: {}.", instanceId, instance.getState());
-    } catch (Exception e) {
-      cleanupAll();
-      throw new SpannerResourceManagerException("Failed to create instance.", e);
+    if (!hasInstancePartition && instancePartitionId != null && instancePartitionConfig != null) {
+      LOG.info(
+          "Creating instance partition {} with config {}",
+          instancePartitionId,
+          instancePartitionConfig);
+      try {
+        InstancePartition instancePartition =
+            InstancePartition.newBuilder()
+                .setDisplayName(instancePartitionId)
+                .setNodeCount(1)
+                .setConfig(InstanceConfigName.of(projectId, instancePartitionConfig).toString())
+                .build();
+        v1InstanceAdminClient
+            .createInstancePartitionAsync(
+                CreateInstancePartitionRequest.newBuilder()
+                    .setParent(InstanceName.of(projectId, instanceId).toString())
+                    .setInstancePartitionId(instancePartitionId)
+                    .setInstancePartition(instancePartition)
+                    .build())
+            .get();
+        hasInstancePartition = true;
+        LOG.info("Successfully created instance partition {}.", instancePartitionId);
+      } catch (Exception e) {
+        if (e.getMessage() != null && e.getMessage().contains("ALREADY_EXISTS")) {
+          LOG.info("Instance partition {} already exists.", instancePartitionId);
+          hasInstancePartition = true;
+        } else {
+          cleanupAll();
+          throw new SpannerResourceManagerException("Failed to create instance partition.", e);
+        }
+      }
     }
   }
 
@@ -223,16 +268,18 @@ private synchronized void maybeCreateDatabase() {
       Database database =
           Failsafe.with(retryOnQuotaException())
               .get(
-                  () ->
-                      databaseAdminClient
-                          .createDatabase(
-                              databaseAdminClient
-                                  .newDatabaseBuilder(
-                                      DatabaseId.of(projectId, instanceId, databaseId))
-                                  .setDialect(dialect)
-                                  .build(),
-                              ImmutableList.of())
-                          .get());
+                  () -> {
+                    DatabaseInfo.Builder dbBuilder =
+                        databaseAdminClient
+                            .newDatabaseBuilder(DatabaseId.of(projectId, instanceId, databaseId))
+                            .setDialect(dialect);
+                    if (protoDescriptors != null) {
+                      dbBuilder.setProtoDescriptors(protoDescriptors);
+                    }
+                    return databaseAdminClient
+                        .createDatabase(dbBuilder.build(), ImmutableList.of())
+                        .get();
+                  });
 
       hasDatabase = true;
       LOG.info("Successfully created database {}: {}.", databaseId, database.getState());
@@ -327,6 +374,10 @@ public synchronized void executeDdlStatements(List<String> statements)
       throws IllegalStateException {
     ensureUsableAndCreateResources();
 
+    if (statements.isEmpty()) {
+      return;
+    }
+
     if (suppressVerboseLogs) {
       LOG.info("Executing {} DDL statements on database {}.", statements.size(), databaseId);
     } else {
@@ -336,11 +387,18 @@ public synchronized void executeDdlStatements(List<String> statements)
       // executeDdlStatments can fail for spanner staging because of failfast.
       Failsafe.with(retryOnQuotaException())
           .run(
-              () ->
-                  databaseAdminClient
-                      .updateDatabaseDdl(
-                          instanceId, databaseId, statements, /* operationId= */ null)
-                      .get());
+              () -> {
+                DatabaseInfo.Builder dbBuilder =
+                    databaseAdminClient.newDatabaseBuilder(
+                        DatabaseId.of(projectId, instanceId, databaseId));
+                if (protoDescriptors != null) {
+                  dbBuilder.setProtoDescriptors(protoDescriptors);
+                }
+                Database database = dbBuilder.build();
+                databaseAdminClient
+                    .updateDatabaseDdl(database, statements, /* operationId= */ null)
+                    .get();
+              });
       if (suppressVerboseLogs) {
         LOG.info(
             "Successfully executed {} DDL statements on database {}.",
@@ -372,6 +430,17 @@ public synchronized DatabaseClient getDatabaseClient() {
     return spanner.getDatabaseClient(DatabaseId.of(projectId, instanceId, databaseId));
   }
 
+  /**
+   * Creates and returns Spanner Batch Client.
+   *
+   * @return Spanner Batch Client
+   */
+  public synchronized BatchClient getBatchClient() {
+    checkIsUsable();
+    checkHasInstanceAndDatabase();
+    return spanner.getBatchClient(DatabaseId.of(projectId, instanceId, databaseId));
+  }
+
   /**
    * Writes a given record into a table. This method requires {@link
    * SpannerResourceManager#executeDdlStatement(String)} to be called for the target table
@@ -583,7 +652,9 @@ public synchronized void cleanupAll() {
           Failsafe.with(retryOnQuotaException())
               .run(() -> databaseAdminClient.dropDatabase(instanceId, databaseId));
         }
+
       } else {
+
         LOG.info("Deleting instance {}...", instanceId);
 
         if (instanceAdminClient != null) {
@@ -597,6 +668,10 @@ public synchronized void cleanupAll() {
     } catch (SpannerException e) {
       throw new SpannerResourceManagerException("Failed to delete instance.", e);
     } finally {
+      if (v1InstanceAdminClient != null && !v1InstanceAdminClient.isShutdown()) {
+        v1InstanceAdminClient.shutdown();
+      }
+
       if (!spanner.isClosed()) {
         spanner.close();
       }
@@ -663,8 +738,11 @@ public static final class Builder {
     private Credentials credentials;
     private String host;
     private int nodeCount;
+    private byte[] protoDescriptors;
     private MonitoringClient monitoringClient;
     private boolean suppressVerboseLogs;
+    private String instancePartitionId;
+    private String instancePartitionConfig;
 
     private Builder(String testId, String projectId, String region, Dialect dialect) {
       this.testId = testId;
@@ -677,6 +755,13 @@ private Builder(String testId, String projectId, String region, Dialect dialect)
       this.nodeCount = 1;
     }
 
+    public Builder setInstancePartition(
+        String instancePartitionId, String instancePartitionConfig) {
+      this.instancePartitionId = instancePartitionId;
+      this.instancePartitionConfig = instancePartitionConfig;
+      return this;
+    }
+
     public Builder setCredentials(Credentials credentials) {
       this.credentials = credentials;
       return this;
@@ -751,6 +836,17 @@ public Builder setNodeCount(int nodeCount) {
       return this;
     }
 
+    /**
+     * Configures proto descriptors for database creation.
+     *
+     * @param protoDescriptors
+     * @return
+     */
+    public Builder setProtoDescriptors(byte[] protoDescriptors) {
+      this.protoDescriptors = protoDescriptors;
+      return this;
+    }
+
     /**
      * Sets Monitoring Client instance to be used for getMetrics method.
      *
diff --git a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/spanner/SpannerResourceManagerTest.java b/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/spanner/SpannerResourceManagerTest.java
@@ -94,6 +94,7 @@ public final class SpannerResourceManagerTest {
   @Captor private ArgumentCaptor<Iterable<String>> statementCaptor;
   @Captor private ArgumentCaptor<String> instanceIdCaptor;
   @Captor private ArgumentCaptor<String> databaseIdCaptor;
+  @Captor private ArgumentCaptor<Database> databaseCaptor;
   @Captor private ArgumentCaptor<String> projectIdCaptor;
 
   @Before
@@ -153,7 +154,10 @@ public void testExecuteDdlStatementShouldThrowExceptionWhenSpannerUpdateDatabase
     // arrange
     prepareCreateInstanceMock();
     prepareCreateDatabaseMock();
-    when(spanner.getDatabaseAdminClient().updateDatabaseDdl(any(), any(), any(), any()).get())
+    when(spanner
+            .getDatabaseAdminClient()
+            .updateDatabaseDdl(any(Database.class), any(), any())
+            .get())
         .thenThrow(InterruptedException.class);
     String statement =
         "CREATE TABLE Singers (\n"
@@ -190,14 +194,10 @@ public void testExecuteDdlStatementShouldWorkWhenSpannerDoesntThrowAnyError()
     verify(spanner.getInstanceAdminClient(), times(2)).createInstance(any());
     verify(spanner.getDatabaseAdminClient(), times(2)).createDatabase(any(), any());
     verify(spanner.getDatabaseAdminClient(), times(2))
-        .updateDatabaseDdl(
-            instanceIdCaptor.capture(),
-            databaseIdCaptor.capture(),
-            statementCaptor.capture(),
-            any());
-
-    String actualInstanceId = instanceIdCaptor.getValue();
-    String actualDatabaseId = databaseIdCaptor.getValue();
+        .updateDatabaseDdl(databaseCaptor.capture(), statementCaptor.capture(), any());
+
+    String actualInstanceId = testManager.getInstanceId();
+    String actualDatabaseId = testManager.getDatabaseId();
     Iterable<String> actualStatement = statementCaptor.getValue();
 
     assertThat(actualInstanceId).matches(TEST_ID + "-\\d{8}-\\d{6}-[a-zA-Z0-9]{6}");
@@ -222,7 +222,10 @@ public void testExecuteDdlStatementShouldRetryOnResourceExhaustedError()
     RuntimeException resourceExhaustedException =
         new RuntimeException(
             "com.google.cloud.spanner.SpannerException: RESOURCE_EXHAUSTED: io.grpc.StatusRuntimeException: RESOURCE_EXHAUSTED: CPU overload detected");
-    when(spanner.getDatabaseAdminClient().updateDatabaseDdl(any(), any(), any(), any()).get())
+    when(spanner
+            .getDatabaseAdminClient()
+            .updateDatabaseDdl(any(Database.class), any(), any())
+            .get())
         .thenThrow(resourceExhaustedException)
         .thenReturn(null);
 
@@ -234,14 +237,10 @@ public void testExecuteDdlStatementShouldRetryOnResourceExhaustedError()
     verify(spanner.getInstanceAdminClient(), times(2)).createInstance(any());
     verify(spanner.getDatabaseAdminClient(), times(2)).createDatabase(any(), any());
     verify(spanner.getDatabaseAdminClient(), times(3))
-        .updateDatabaseDdl(
-            instanceIdCaptor.capture(),
-            databaseIdCaptor.capture(),
-            statementCaptor.capture(),
-            any());
-
-    String actualInstanceId = instanceIdCaptor.getValue();
-    String actualDatabaseId = databaseIdCaptor.getValue();
+        .updateDatabaseDdl(databaseCaptor.capture(), statementCaptor.capture(), any());
+
+    String actualInstanceId = testManager.getInstanceId();
+    String actualDatabaseId = testManager.getDatabaseId();
     Iterable<String> actualStatement = statementCaptor.getValue();
 
     assertThat(actualInstanceId).matches(TEST_ID + "-\\d{8}-\\d{6}-[a-zA-Z0-9]{6}");
@@ -826,7 +825,11 @@ private void prepareCreateDatabaseMock() throws ExecutionException, InterruptedE
 
   private void prepareUpdateDatabaseMock() throws ExecutionException, InterruptedException {
     Mockito.lenient()
-        .when(spanner.getDatabaseAdminClient().updateDatabaseDdl(any(), any(), any(), any()).get())
+        .when(
+            spanner
+                .getDatabaseAdminClient()
+                .updateDatabaseDdl(any(Database.class), any(), any())
+                .get())
         .thenReturn(null);
   }
 
diff --git a/v1/src/main/java/com/google/cloud/teleport/spanner/ddl/InformationSchemaScanner.java b/v1/src/main/java/com/google/cloud/teleport/spanner/ddl/InformationSchemaScanner.java
@@ -1336,9 +1336,9 @@ private void listPropertyGraphTables(Ddl.Builder builder, String tableType) {
       String tablesJson;
       try {
         tablesJson = resultSet.getJson(2);
-      } catch (Exception edgeTableException) {
-        LOG.debug(propertyGraphNameQualified + " does not contain any edge tables");
-        return;
+      } catch (Exception tableException) {
+        LOG.debug(propertyGraphNameQualified + " does not contain any {}", tableType);
+        continue;
       }
 
       LOG.debug("Schema PropertyGraph {}", propertyGraphNameQualified);
diff --git a/v1/src/main/java/com/google/cloud/teleport/spanner/spannerio/MutationKeyEncoder.java b/v1/src/main/java/com/google/cloud/teleport/spanner/spannerio/MutationKeyEncoder.java
@@ -59,6 +59,11 @@ public MutationKeyEncoder(SpannerSchema schema) {
     this.schema = schema;
   }
 
+  @VisibleForTesting
+  static void clearUnknownTablesWarningsMap() {
+    unknownTablesWarnings.clear();
+  }
+
   /**
    * Builds a lexicographically sortable binary key based on a primary key descriptor.
    *
diff --git a/v1/src/test/java/com/google/cloud/teleport/spanner/ddl/InformationSchemaScannerIT.java b/v1/src/test/java/com/google/cloud/teleport/spanner/ddl/InformationSchemaScannerIT.java
diff --git a/v1/src/test/java/com/google/cloud/teleport/spanner/spannerio/MutationKeyEncoderTest.java b/v1/src/test/java/com/google/cloud/teleport/spanner/spannerio/MutationKeyEncoderTest.java

Original file line number	Diff line number	Diff line change
`@@ -59,6 +59,11 @@ public MutationKeyEncoder(SpannerSchema schema) {`
`59`	`59`	`this.schema = schema;`
`60`	`60`	`}`
`61`	`61`
	`62`	`+ @VisibleForTesting`
	`63`	`+ static void clearUnknownTablesWarningsMap() {`
	`64`	`+ unknownTablesWarnings.clear();`
	`65`	`+ }`
	`66`	`+`
`62`	`67`	`/**`
`63`	`68`	`* Builds a lexicographically sortable binary key based on a primary key descriptor.`
`64`	`69`	`*`