fix: (Spanner) include schema name in table name (#2510)

olavloite · web-flow · commit 5dfe2260c602 · 2024-01-11T11:01:39.000-05:00
Spanner databases contain both system schemas and a user schema. The
current Spring Data Spanner implementation assumes that all table names
in a database are unique, but this is not true if a user creates a table
that is equal to the name of a system table in one of the system
schemas.

Additionally, Spanner could introduce support for user-defined named
schemas in the future. This would increase the probability that a single
database contains multiple tables with the same name, but in different
schemas.

This change fixes any problems that might occur if a database contains
multiple tables with the same name in different schemas. The change does
not attempt to implement full support for using multiple different
schemas with Spring Cloud Data Spanner.
diff --git a/spring-cloud-gcp-data-spanner/src/main/java/com/google/cloud/spring/data/spanner/core/admin/SpannerDatabaseAdminTemplate.java b/spring-cloud-gcp-data-spanner/src/main/java/com/google/cloud/spring/data/spanner/core/admin/SpannerDatabaseAdminTemplate.java
@@ -20,13 +20,15 @@
 import com.google.cloud.spanner.DatabaseAdminClient;
 import com.google.cloud.spanner.DatabaseClient;
 import com.google.cloud.spanner.DatabaseId;
+import com.google.cloud.spanner.Dialect;
 import com.google.cloud.spanner.ResultSet;
 import com.google.cloud.spanner.Statement;
 import com.google.cloud.spanner.Struct;
 import com.google.cloud.spring.data.spanner.core.mapping.SpannerDataException;
 import java.util.HashMap;
 import java.util.HashSet;
 import java.util.Map;
+import java.util.Objects;
 import java.util.Set;
 import java.util.concurrent.ExecutionException;
 import java.util.function.Supplier;
@@ -40,13 +42,17 @@
  */
 public class SpannerDatabaseAdminTemplate {
 
+  private static final String TABLE_SCHEMA_COL_NAME = "table_schema";
+  
   private static final String TABLE_NAME_COL_NAME = "table_name";
 
   private static final String PARENT_TABLE_NAME_COL_NAME = "parent_table_name";
 
   private static final Statement TABLE_AND_PARENT_QUERY =
       Statement.of(
           "SELECT t."
+              + TABLE_SCHEMA_COL_NAME
+              + ", t."
               + TABLE_NAME_COL_NAME
               + ", t."
               + PARENT_TABLE_NAME_COL_NAME
@@ -151,14 +157,36 @@ public Map<String, String> getChildParentTablesMap() {
       while (results.next()) {
         Struct row = results.getCurrentRowAsStruct();
         relationships.put(
-            row.getString(TABLE_NAME_COL_NAME),
+            getQualifiedTableName(
+                row.getString(TABLE_SCHEMA_COL_NAME),
+                row.getString(TABLE_NAME_COL_NAME)),
             row.isNull(PARENT_TABLE_NAME_COL_NAME)
                 ? null
-                : row.getString(PARENT_TABLE_NAME_COL_NAME));
+                : getQualifiedTableName(
+                    // Parent/child tables must be in the same schema, and so there is no separate
+                    // schema name column for the parent table.
+                    row.getString(TABLE_SCHEMA_COL_NAME),
+                    row.getString(PARENT_TABLE_NAME_COL_NAME)));
       }
       return relationships;
     }
   }
+  
+  private String getQualifiedTableName(String schema, String table) {
+    if (schema == null || Objects.equals(schema, getDefaultSchemaName())) {
+      return table;
+    }
+    return schema + "." + table;
+  }
+  
+  private String getDefaultSchemaName() {
+    // TODO: Get the default schema directly from the dialect when this is supported in the Spanner
+    //       Java client.
+    if (databaseClientProvider.get().getDialect() == Dialect.POSTGRESQL) {
+      return "public";
+    }
+    return "";
+  }
 
   /**
    * Return true if the given table names are interleaved as ancestor and descendant. These may be
diff --git a/spring-cloud-gcp-data-spanner/src/test/java/com/google/cloud/spring/data/spanner/core/admin/SpannerDatabaseAdminTemplateTests.java b/spring-cloud-gcp-data-spanner/src/test/java/com/google/cloud/spring/data/spanner/core/admin/SpannerDatabaseAdminTemplateTests.java
@@ -76,39 +76,104 @@ void setup() {
 
   @Test
   void getTableRelationshipsTest() {
+    when(this.mockDatabasePage.getValues())
+        .thenReturn(
+            List.of(new Database(this.databaseId, State.READY, this.databaseAdminClient)));
     ReadContext readContext = mock(ReadContext.class);
 
     Struct s1 =
         Struct.newBuilder()
+            .set("table_schema")
+            .to("")
             .set("table_name")
             .to(Value.string("grandpa"))
             .set("parent_table_name")
             .to(Value.string(null))
             .build();
     Struct s2 =
         Struct.newBuilder()
+            .set("table_schema")
+            .to("")
             .set("table_name")
             .to(Value.string("parent_a"))
             .set("parent_table_name")
             .to(Value.string("grandpa"))
             .build();
     Struct s3 =
         Struct.newBuilder()
+            .set("table_schema")
+            .to("")
             .set("table_name")
             .to(Value.string("parent_b"))
             .set("parent_table_name")
             .to(Value.string("grandpa"))
             .build();
     Struct s4 =
         Struct.newBuilder()
+            .set("table_schema")
+            .to("")
             .set("table_name")
             .to(Value.string("child"))
             .set("parent_table_name")
             .to(Value.string("parent_a"))
             .build();
+    Struct s5 =
+        Struct.newBuilder()
+            .set("table_schema")
+            .to("INFORMATION_SCHEMA")
+            .set("table_name")
+            .to(Value.string("TABLES"))
+            .set("parent_table_name")
+            .to(Value.string(null))
+            .build();
+    Struct s6 =
+        Struct.newBuilder()
+            .set("table_schema")
+            .to("")
+            .set("table_name")
+            .to(Value.string("TABLES"))
+            .set("parent_table_name")
+            .to(Value.string(null))
+            .build();
+    Struct s7 =
+        Struct.newBuilder()
+            .set("table_schema")
+            .to("")
+            .set("table_name")
+            .to(Value.string("CHILD_TABLES"))
+            .set("parent_table_name")
+            .to(Value.string("TABLES"))
+            .build();
+    Struct s8 =
+        Struct.newBuilder()
+            .set("table_schema")
+            .to("my_schema")
+            .set("table_name")
+            .to(Value.string("grandpa"))
+            .set("parent_table_name")
+            .to(Value.string(null))
+            .build();
+    Struct s9 =
+        Struct.newBuilder()
+            .set("table_schema")
+            .to("my_schema")
+            .set("table_name")
+            .to(Value.string("dad"))
+            .set("parent_table_name")
+            .to(Value.string("grandpa"))
+            .build();
+    Struct s10 =
+        Struct.newBuilder()
+            .set("table_schema")
+            .to("my_schema")
+            .set("table_name")
+            .to(Value.string("child"))
+            .set("parent_table_name")
+            .to(Value.string("dad"))
+            .build();
 
     MockResults mockResults = new MockResults();
-    mockResults.structs = Arrays.asList(s1, s2, s3, s4);
+    mockResults.structs = Arrays.asList(s1, s2, s3, s4, s5, s6, s7, s8, s9, s10);
     ResultSet results = mock(ResultSet.class);
     when(results.next()).thenAnswer(invocation -> mockResults.next());
     when(results.getCurrentRowAsStruct()).thenAnswer(invocation -> mockResults.getCurrent());
@@ -118,9 +183,27 @@ void getTableRelationshipsTest() {
     Map<String, Set<String>> relationships =
         this.spannerDatabaseAdminTemplate.getParentChildTablesMap();
 
-    assertThat(relationships).hasSize(2);
+    assertThat(this.spannerDatabaseAdminTemplate.tableExists("some-random-table")).isFalse();
+    assertThat(this.spannerDatabaseAdminTemplate.tableExists("grandpa")).isTrue();
+    assertThat(this.spannerDatabaseAdminTemplate.tableExists("parent_a")).isTrue();
+    assertThat(this.spannerDatabaseAdminTemplate.tableExists("parent_b")).isTrue();
+    assertThat(this.spannerDatabaseAdminTemplate.tableExists("child")).isTrue();
+    assertThat(this.spannerDatabaseAdminTemplate.tableExists("child")).isTrue();
+    assertThat(this.spannerDatabaseAdminTemplate.tableExists("INFORMATION_SCHEMA.TABLES")).isTrue();
+    assertThat(this.spannerDatabaseAdminTemplate.tableExists("TABLES")).isTrue();
+    assertThat(this.spannerDatabaseAdminTemplate.tableExists("CHILD_TABLES")).isTrue();
+    assertThat(this.spannerDatabaseAdminTemplate.tableExists("my_schema.grandpa")).isTrue();
+    assertThat(this.spannerDatabaseAdminTemplate.tableExists("my_schema.dad")).isTrue();
+    assertThat(this.spannerDatabaseAdminTemplate.tableExists("my_schema.child")).isTrue();
+
+    assertThat(relationships).hasSize(5);
     assertThat(relationships.get("grandpa")).containsExactlyInAnyOrder("parent_a", "parent_b");
     assertThat(relationships.get("parent_a")).containsExactlyInAnyOrder("child");
+    assertThat(relationships.get("child")).isNull();
+    assertThat(relationships.get("TABLES")).containsExactlyInAnyOrder("CHILD_TABLES");
+    assertThat(relationships.get("INFORMATION_SCHEMA.TABLES")).isNull();
+    assertThat(relationships.get("my_schema.grandpa")).containsExactlyInAnyOrder("my_schema.dad");
+    assertThat(relationships.get("my_schema.dad")).containsExactlyInAnyOrder("my_schema.child");
 
     assertThat(this.spannerDatabaseAdminTemplate.isInterleaved("grandpa", "child"))
         .as("verify grand-child relationship")
@@ -140,6 +223,19 @@ void getTableRelationshipsTest() {
     assertThat(this.spannerDatabaseAdminTemplate.isInterleaved("parent_b", "child"))
         .as("verify not parent-child relationship")
         .isFalse();
+    
+    assertThat(this.spannerDatabaseAdminTemplate
+        .isInterleaved("my_schema.grandpa", "my_schema.dad"))
+        .as("verify my-schema grand-child relationship")
+        .isTrue();
+    assertThat(this.spannerDatabaseAdminTemplate
+        .isInterleaved("my_schema.grandpa", "my_schema.child"))
+        .as("verify my-schema grand-child relationship")
+        .isTrue();
+    assertThat(this.spannerDatabaseAdminTemplate
+        .isInterleaved("my_schema.grandpa", "child"))
+        .as("verify not my-schema grand-child relationship")
+        .isFalse();
   }
 
   @Test
