feat: prepare for named schemas (#889)

* feat: prepare for named schemas

* fix: add schema to unique key and index

google-cloud-spanner-hibernate-dialect/src/main/java/com/google/cloud/spanner/hibernate/SpannerUniqueDelegate.java

@@ -18,6 +18,7 @@
 
 package com.google.cloud.spanner.hibernate;
 
+import com.google.common.base.Strings;
 import org.hibernate.boot.Metadata;
 import org.hibernate.boot.model.relational.SqlStringGenerationContext;
 import org.hibernate.dialect.Dialect;
@@ -52,6 +53,9 @@ public String getAlterTableToAddUniqueKeyCommand(UniqueKey uniqueKey, Metadata m
   public String getAlterTableToDropUniqueKeyCommand(UniqueKey uniqueKey, Metadata metadata,
       SqlStringGenerationContext context) {
     StringBuilder buf = new StringBuilder("DROP INDEX ");
+    if (!Strings.isNullOrEmpty(uniqueKey.getTable().getSchema())) {
+      buf.append(dialect.quote(uniqueKey.getTable().getSchema())).append('.');
+    }
     buf.append(dialect.quote(uniqueKey.getName()));
     return buf.toString();
   }

google-cloud-spanner-hibernate-dialect/src/main/java/com/google/cloud/spanner/hibernate/schema/SpannerDatabaseInfo.java

@@ -21,93 +21,112 @@
 import java.sql.DatabaseMetaData;
 import java.sql.ResultSet;
 import java.sql.SQLException;
+import java.util.HashMap;
 import java.util.HashSet;
+import java.util.Map;
 import java.util.Set;
+import org.hibernate.boot.model.naming.Identifier;
+import org.hibernate.boot.model.relational.Database;
+import org.hibernate.mapping.Table;
 
 /**
  * Helper class for extracting information from the {@link DatabaseMetaData} which contains
  * information about what tables and indices currently exist in the database.
  */
 public class SpannerDatabaseInfo {
+
+  private final Set<Table> tableNames;
 
-  private final Set<String> tableNames;
-
-  private final Set<String> indexNames;
+  private final Map<Table, Set<String>> indexNames;
 
   private final DatabaseMetaData databaseMetaData;
 
   /**
    * Constructs the {@link SpannerDatabaseInfo} by querying the Spanner database metadata.
    */
-  public SpannerDatabaseInfo(DatabaseMetaData databaseMetaData) throws SQLException {
-    this.tableNames = extractDatabaseTables(databaseMetaData);
-    this.indexNames = extractDatabaseIndices(databaseMetaData);
+  public SpannerDatabaseInfo(Database database, DatabaseMetaData databaseMetaData)
+      throws SQLException {
+    this.tableNames = extractDatabaseTables(database, databaseMetaData);
+    this.indexNames = extractDatabaseIndices(database, databaseMetaData);
     this.databaseMetaData = databaseMetaData;
   }
 
   /**
    * Returns the table names in the Spanner database.
    */
-  public Set<String> getAllTables() {
+  public Set<Table> getAllTables() {
     return tableNames;
   }
 
   /**
    * Returns the names of all the indices in the Spanner database.
    */
-  public Set<String> getAllIndices() {
+  public Map<Table, Set<String>> getAllIndices() {
     return indexNames;
   }
 
   /**
    * Returns the names of all the imported foreign keys for a specified {@code tableName}.
    */
-  public Set<String> getImportedForeignKeys(String tableName) {
+  public Set<String> getImportedForeignKeys(Table table) {
     try {
       HashSet<String> foreignKeys = new HashSet<>();
 
-      ResultSet rs = databaseMetaData.getImportedKeys(null, null, tableName);
+      ResultSet rs = databaseMetaData.getImportedKeys(
+          table.getCatalog(), table.getSchema(), table.getName());
       while (rs.next()) {
         foreignKeys.add(rs.getString("FK_NAME"));
       }
       rs.close();
       return foreignKeys;
     } catch (SQLException e) {
       throw new RuntimeException(
-          "Failed to lookup Spanner Database foreign keys for table: " + tableName, e);
+          "Failed to lookup Spanner Database foreign keys for table: " + table, e);
     }
   }
 
-  private static Set<String> extractDatabaseTables(DatabaseMetaData databaseMetaData)
-      throws SQLException {
-    HashSet<String> result = new HashSet<String>();
+  private static Set<Table> extractDatabaseTables(
+      Database database, DatabaseMetaData databaseMetaData) throws SQLException {
+    HashSet<Table> result = new HashSet<>();
 
     // Passing all null parameters will get all the tables and apply no filters.
-    ResultSet resultSet = databaseMetaData.getTables(
-        null, null, null, null);
-    while (resultSet.next()) {
-      String type = resultSet.getString("TABLE_TYPE");
-      if (type.equals("TABLE")) {
-        result.add(resultSet.getString("TABLE_NAME"));
+    try (ResultSet resultSet = databaseMetaData.getTables(
+        null, null, null, null)) {
+      while (resultSet.next()) {
+        String type = resultSet.getString("TABLE_TYPE");
+        if (type.equals("TABLE")) {
+          Table table = new Table("orm",
+              database.locateNamespace(
+                  Identifier.toIdentifier(resultSet.getString("TABLE_CAT")),
+                  Identifier.toIdentifier(resultSet.getString("TABLE_SCHEM"))),
+              Identifier.toIdentifier(resultSet.getString("TABLE_NAME")),
+              false);
+          result.add(table);
+        }
       }
     }
-    resultSet.close();
 
     return result;
   }
 
-  private static Set<String> extractDatabaseIndices(DatabaseMetaData databaseMetaData)
-      throws SQLException {
-    HashSet<String> result = new HashSet<>();
-    ResultSet indexResultSet = databaseMetaData.getIndexInfo(
-        null, null, null, false, false);
-
-    while (indexResultSet.next()) {
-      String name = indexResultSet.getString("INDEX_NAME");
-      result.add(name);
+  private static Map<Table, Set<String>> extractDatabaseIndices(
+      Database database, DatabaseMetaData databaseMetaData) throws SQLException {
+    HashMap<Table, Set<String>> result = new HashMap<>();
+    try (ResultSet indexResultSet = databaseMetaData.getIndexInfo(
+        null, null, null, false, false)) {
+
+      while (indexResultSet.next()) {
+        String name = indexResultSet.getString("INDEX_NAME");
+        Table table = new Table("orm",
+            database.locateNamespace(
+                Identifier.toIdentifier(indexResultSet.getString("TABLE_CAT")),
+                Identifier.toIdentifier(indexResultSet.getString("TABLE_SCHEM"))),
+            Identifier.toIdentifier(indexResultSet.getString("TABLE_NAME")),
+            false);
+        Set<String> tableIndices = result.computeIfAbsent(table, k -> new HashSet<>());
+        tableIndices.add(name);
+      }
     }
-    indexResultSet.close();
-
     return result;
   }
 }

google-cloud-spanner-hibernate-dialect/src/main/java/com/google/cloud/spanner/hibernate/schema/SpannerForeignKeyExporter.java

@@ -22,6 +22,7 @@
 import org.hibernate.boot.model.relational.SqlStringGenerationContext;
 import org.hibernate.dialect.Dialect;
 import org.hibernate.mapping.ForeignKey;
+import org.hibernate.mapping.Table;
 import org.hibernate.tool.schema.internal.StandardForeignKeyExporter;
 
 /**
@@ -55,7 +56,7 @@ public String[] getSqlDropStrings(ForeignKey foreignKey, Metadata metadata,
   }
 
   private boolean foreignKeyExists(ForeignKey foreignKey) {
-    String table = foreignKey.getTable().getName();
+    Table table = foreignKey.getTable();
     return spannerDatabaseInfo.getAllTables().contains(table)
         && spannerDatabaseInfo.getImportedForeignKeys(table).contains(foreignKey.getName());
   }

google-cloud-spanner-hibernate-dialect/src/main/java/com/google/cloud/spanner/hibernate/schema/SpannerSchemaCreator.java

@@ -59,7 +59,8 @@ public void doCreation(
     DdlTransactionIsolator isolator = tool.getDdlTransactionIsolator(options);
     try {
       Connection connection = isolator.getIsolatedConnection();
-      SpannerDatabaseInfo spannerDatabaseInfo = new SpannerDatabaseInfo(connection.getMetaData());
+      SpannerDatabaseInfo spannerDatabaseInfo =
+          new SpannerDatabaseInfo(metadata.getDatabase(), connection.getMetaData());
       tool.getSpannerTableExporter(options).init(metadata, spannerDatabaseInfo, Action.CREATE);
       tool.getForeignKeyExporter(options).init(spannerDatabaseInfo);
       schemaCreator.doCreation(

google-cloud-spanner-hibernate-dialect/src/main/java/com/google/cloud/spanner/hibernate/schema/SpannerSchemaDropper.java

@@ -61,7 +61,8 @@ public void doDrop(
     try {
       Connection connection = isolator.getIsolatedConnection();
       // Initialize exporters with drop table dependencies so tables are dropped in the right order.
-      SpannerDatabaseInfo spannerDatabaseInfo = new SpannerDatabaseInfo(connection.getMetaData());
+      SpannerDatabaseInfo spannerDatabaseInfo =
+          new SpannerDatabaseInfo(metadata.getDatabase(), connection.getMetaData());
       tool.getSpannerTableExporter(options).init(metadata, spannerDatabaseInfo, Action.DROP);
       tool.getForeignKeyExporter(options).init(spannerDatabaseInfo);
       schemaDropper.doDrop(
@@ -84,7 +85,8 @@ public DelayedDropAction buildDelayedAction(
     try {
       Connection connection = isolator.getIsolatedConnection();
       // Initialize exporters with drop table dependencies so tables are dropped in the right order.
-      SpannerDatabaseInfo spannerDatabaseInfo = new SpannerDatabaseInfo(connection.getMetaData());
+      SpannerDatabaseInfo spannerDatabaseInfo =
+          new SpannerDatabaseInfo(metadata.getDatabase(), connection.getMetaData());
       tool.getSpannerTableExporter(options).init(metadata, spannerDatabaseInfo, Action.DROP);
       tool.getForeignKeyExporter(options).init(spannerDatabaseInfo);
       return schemaDropper.buildDelayedAction(

google-cloud-spanner-hibernate-dialect/src/main/java/com/google/cloud/spanner/hibernate/schema/SpannerSchemaMigrator.java

@@ -58,7 +58,8 @@ public void doMigration(
     DdlTransactionIsolator isolator = tool.getDdlTransactionIsolator(options);
     try {
       Connection connection = isolator.getIsolatedConnection();
-      SpannerDatabaseInfo spannerDatabaseInfo = new SpannerDatabaseInfo(connection.getMetaData());
+      SpannerDatabaseInfo spannerDatabaseInfo =
+          new SpannerDatabaseInfo(metadata.getDatabase(), connection.getMetaData());
       tool.getSpannerTableExporter(options).init(metadata, spannerDatabaseInfo, Action.UPDATE);
       tool.getForeignKeyExporter(options).init(spannerDatabaseInfo);
       schemaMigrator.doMigration(metadata, options, contributableInclusionFilter, targetDescriptor);

google-cloud-spanner-hibernate-dialect/src/main/java/com/google/cloud/spanner/hibernate/schema/SpannerTableStatements.java

@@ -23,6 +23,7 @@
 import com.google.cloud.spanner.hibernate.Interleaved;
 import com.google.cloud.spanner.hibernate.SpannerDialect;
 import com.google.cloud.spanner.hibernate.types.SpannerArrayListType;
+import com.google.common.base.Strings;
 import com.google.common.collect.Sets;
 import java.sql.Types;
 import java.text.MessageFormat;
@@ -64,25 +65,36 @@ public void initializeSpannerDatabaseInfo(SpannerDatabaseInfo spannerDatabaseInf
   public List<String> dropTable(Table table) {
     ArrayList<String> dropStrings = new ArrayList<>();
 
-    for (String indexName : getTableIndices(table)) {
-      if (spannerDatabaseInfo.getAllIndices().contains(indexName)) {
-        dropStrings.add("drop index " + indexName);
+    Set<String> existingTableIndices = spannerDatabaseInfo.getAllIndices().get(table);
+    if (existingTableIndices != null) {
+      for (String indexName : getTableIndices(table)) {
+        if (existingTableIndices.contains(indexName)) {
+          dropStrings.add("drop index " + getQualifiedIndexName(table, indexName));
+        }
       }
     }
 
-    if (spannerDatabaseInfo.getAllTables().contains(table.getName())) {
-      dropStrings.add(this.spannerDialect.getDropTableString(table.getQuotedName()));
+    if (spannerDatabaseInfo.getAllTables().contains(table)) {
+      dropStrings.add(this.spannerDialect.getDropTableString(
+          table.getQualifiedTableName().quote().getObjectName().toString()));
     }
     return dropStrings;
   }
 
+  private String getQualifiedIndexName(Table table, String index) {
+    if (Strings.isNullOrEmpty(table.getSchema())) {
+      return index;
+    }
+    return table.getSchema() + "." + index;
+  }
+
   private Set<String> getTableIndices(Table table) {
     return Sets.union(table.getIndexes().keySet(), table.getUniqueKeys().keySet());
   }
 
   /** Generates the statements needed to create a table. */
   public List<String> createTable(Table table, Metadata metadata) {
-    if (spannerDatabaseInfo.getAllTables().contains(table.getName())) {
+    if (spannerDatabaseInfo.getAllTables().contains(table)) {
       return Collections.emptyList();
     }
 
@@ -133,7 +145,7 @@ private List<String> getCreateTableStrings(
     String createTableString =
         MessageFormat.format(
             CREATE_TABLE_TEMPLATE,
-            table.getQuotedName(),
+            table.getQualifiedTableName(),
             allColumnNames,
             primaryKeyColNames,
             getInterleavedClause(table, metadata));
@@ -184,7 +196,7 @@ private static String getInterleavedClause(Table table, Metadata metadata) {
     Interleaved interleaved = SchemaUtils.getInterleaveAnnotation(table, metadata);
     if (interleaved != null) {
       Table parentTable = SchemaUtils.getTable(interleaved.parentEntity(), metadata);
-      String interleaveClause = ", INTERLEAVE IN PARENT " + parentTable.getQuotedName();
+      String interleaveClause = ", INTERLEAVE IN PARENT " + parentTable.getQualifiedTableName();
       if (interleaved.cascadeDelete()) {
         interleaveClause += " ON DELETE CASCADE";
       }

google-cloud-spanner-hibernate-dialect/src/test/java/com/google/cloud/spanner/hibernate/GeneratedCreateTableStatementsTests.java

@@ -21,6 +21,8 @@
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertThrows;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
 
 import com.google.cloud.spanner.SpannerException;
 import com.google.cloud.spanner.hibernate.entities.Account;
@@ -38,8 +40,12 @@
 import org.hibernate.Session;
 import org.hibernate.boot.Metadata;
 import org.hibernate.boot.MetadataSources;
+import org.hibernate.boot.model.naming.Identifier;
+import org.hibernate.boot.model.relational.Namespace;
+import org.hibernate.boot.model.relational.Namespace.Name;
 import org.hibernate.boot.registry.StandardServiceRegistry;
 import org.hibernate.boot.registry.StandardServiceRegistryBuilder;
+import org.hibernate.mapping.Table;
 import org.junit.Before;
 import org.junit.Test;
 
@@ -270,10 +276,14 @@ public void testDropTables() throws SQLException {
           new MetadataSources(this.registry)
               .addAnnotatedClass(Employee.class)
               .buildMetadata();
+      Namespace namespace = mock(Namespace.class);
+      when(namespace.getPhysicalName()).thenReturn(
+          new Name(Identifier.toIdentifier(""), Identifier.toIdentifier("")));
+      Table table = new Table("orm", namespace, Identifier.toIdentifier("Employee"), false);
 
       this.connection.setMetaData(MockJdbcUtils.metaDataBuilder()
           .setTables("Employee", "Employee_Sequence")
-          .setIndices("name_index")
+          .setIndices(table, "name_index")
           .build());
 
       Session session = metadata.buildSessionFactory().openSession();
@@ -286,8 +296,8 @@ public void testDropTables() throws SQLException {
       assertThat(sqlStrings).startsWith(
           "START BATCH DDL",
           "drop index name_index",
-          "drop table Employee",
-          "drop table Employee_Sequence",
+          "drop table `Employee`",
+          "drop table `Employee_Sequence`",
           "RUN BATCH"
       );
     } finally {
@@ -301,10 +311,14 @@ public void testDropTablesWithSequencesEnabled() throws SQLException {
         new MetadataSources(this.registry)
             .addAnnotatedClass(Employee.class)
             .buildMetadata();
+    Namespace namespace = mock(Namespace.class);
+    when(namespace.getPhysicalName()).thenReturn(
+        new Name(Identifier.toIdentifier(""), Identifier.toIdentifier("")));
+    Table table = new Table("orm", namespace, Identifier.toIdentifier("Employee"), false);
 
     this.connection.setMetaData(MockJdbcUtils.metaDataBuilder()
         .setTables("Employee")
-        .setIndices("name_index")
+        .setIndices(table, "name_index")
         .build());
 
     Session session = metadata.buildSessionFactory().openSession();
@@ -317,7 +331,7 @@ public void testDropTablesWithSequencesEnabled() throws SQLException {
     assertThat(sqlStrings).startsWith(
         "START BATCH DDL",
         "drop index name_index",
-        "drop table Employee",
+        "drop table `Employee`",
         "drop sequence Employee_Sequence",
         "RUN BATCH"
     );

google-cloud-spanner-hibernate-dialect/src/test/java/com/google/cloud/spanner/hibernate/MockJdbcUtils.java

@@ -30,6 +30,7 @@
 import java.util.Map.Entry;
 import java.util.UUID;
 import java.util.stream.Collectors;
+import org.hibernate.mapping.Table;
 
 /**
  * Helper class to building mock objects for the mock JDBC driver.
@@ -105,11 +106,13 @@ private static ResultSet createTableMetadataResultSet(String... tableNames) {
   /**
    * Constructs a {@link MockResultSet} containing database metadata about indices for testing.
    */
-  private static ResultSet createIndexMetadataResultSet(String... indexNames) {
-    MockResultSet mockResultSet = initResultSet("INDEX_NAME");
+  private static ResultSet createIndexMetadataResultSet(Table table, String... indexNames) {
+    MockResultSet mockResultSet = initResultSet(
+        "TABLE_CAT", "TABLE_SCHEM", "TABLE_NAME", "INDEX_NAME");
 
     for (int i = 0; i < indexNames.length; i++) {
-      String[] row = new String[]{indexNames[i]};
+      String[] row =
+          new String[]{table.getCatalog(), table.getSchema(), table.getName(), indexNames[i]};
       mockResultSet.addRow(row);
     }
 
@@ -181,8 +184,8 @@ public MockDatabaseMetaDataBuilder setTables(Map<String, List<String>> tablesAnd
     /**
      * Sets which indices are present in the Spanner database.
      */
-    public MockDatabaseMetaDataBuilder setIndices(String... indices) {
-      this.indexInfo = createIndexMetadataResultSet(indices);
+    public MockDatabaseMetaDataBuilder setIndices(Table table, String... indices) {
+      this.indexInfo = createIndexMetadataResultSet(table, indices);
       return this;
     }