feat: Add more schema features

src/main/java/com/google/cloud/solutions/spannerddl/diff/DatabaseDefinition.java

@@ -24,6 +24,7 @@
 import com.google.cloud.solutions.spannerddl.parser.ASTcreate_change_stream_statement;
 import com.google.cloud.solutions.spannerddl.parser.ASTcreate_index_statement;
 import com.google.cloud.solutions.spannerddl.parser.ASTcreate_or_replace_statement;
+import com.google.cloud.solutions.spannerddl.parser.ASTcreate_locality_group_statement;
 import com.google.cloud.solutions.spannerddl.parser.ASTcreate_schema_statement;
 import com.google.cloud.solutions.spannerddl.parser.ASTcreate_search_index_statement;
 import com.google.cloud.solutions.spannerddl.parser.ASTcreate_table_statement;
@@ -64,6 +65,7 @@ public static DatabaseDefinition create(List<ASTddl_statement> statements) {
     LinkedHashMap<String, ASTcreate_change_stream_statement> changeStreams = new LinkedHashMap<>();
     LinkedHashMap<String, String> alterDatabaseOptions = new LinkedHashMap<>();
     LinkedHashMap<String, ASTcreate_schema_statement> schemas = new LinkedHashMap<>();
+    LinkedHashMap<String, ASTcreate_locality_group_statement> localityGroups = new LinkedHashMap<>();
 
     for (ASTddl_statement ddlStatement : statements) {
       final SimpleNode statement = (SimpleNode) ddlStatement.jjtGetChild(0);
@@ -92,6 +94,10 @@ public static DatabaseDefinition create(List<ASTddl_statement> statements) {
               ((ASTcreate_search_index_statement) statement).getName(),
               (ASTcreate_search_index_statement) statement);
           break;
+        case DdlParserTreeConstants.JJTCREATE_LOCALITY_GROUP_STATEMENT:
+          ASTcreate_locality_group_statement lg = (ASTcreate_locality_group_statement) statement;
+          localityGroups.put(lg.getNameOrDefault(), lg);
+          break;
         case DdlParserTreeConstants.JJTCREATE_INDEX_STATEMENT:
           indexes.put(
               ((ASTcreate_index_statement) statement).getIndexName(),
@@ -157,7 +163,8 @@ public static DatabaseDefinition create(List<ASTddl_statement> statements) {
         ImmutableMap.copyOf(ttls),
         ImmutableMap.copyOf(changeStreams),
         ImmutableMap.copyOf(alterDatabaseOptions),
-        ImmutableMap.copyOf(schemas));
+        ImmutableMap.copyOf(schemas),
+        ImmutableMap.copyOf(localityGroups));
   }
 
   public abstract ImmutableMap<String, ASTcreate_table_statement> tablesInCreationOrder();
@@ -175,4 +182,6 @@ public static DatabaseDefinition create(List<ASTddl_statement> statements) {
   abstract ImmutableMap<String, String> alterDatabaseOptions();
 
   abstract ImmutableMap<String, ASTcreate_schema_statement> schemas();
+
+  abstract ImmutableMap<String, ASTcreate_locality_group_statement> localityGroups();
 }

src/main/java/com/google/cloud/solutions/spannerddl/diff/DdlDiff.java

@@ -28,13 +28,15 @@
 import com.google.cloud.solutions.spannerddl.parser.ASTcreate_change_stream_statement;
 import com.google.cloud.solutions.spannerddl.parser.ASTcreate_index_statement;
 import com.google.cloud.solutions.spannerddl.parser.ASTcreate_or_replace_statement;
+import com.google.cloud.solutions.spannerddl.parser.ASTcreate_locality_group_statement;
 import com.google.cloud.solutions.spannerddl.parser.ASTcreate_schema_statement;
 import com.google.cloud.solutions.spannerddl.parser.ASTcreate_search_index_statement;
 import com.google.cloud.solutions.spannerddl.parser.ASTcreate_table_statement;
 import com.google.cloud.solutions.spannerddl.parser.ASTddl_statement;
 import com.google.cloud.solutions.spannerddl.parser.ASTforeign_key;
 import com.google.cloud.solutions.spannerddl.parser.ASToptions_clause;
 import com.google.cloud.solutions.spannerddl.parser.ASTrow_deletion_policy_clause;
+import com.google.cloud.solutions.spannerddl.parser.ASTtable_interleave_clause;
 import com.google.cloud.solutions.spannerddl.parser.DdlParser;
 import com.google.cloud.solutions.spannerddl.parser.DdlParserTreeConstants;
 import com.google.cloud.solutions.spannerddl.parser.ParseException;
@@ -54,6 +56,7 @@
 import java.util.Map;
 import java.util.Map.Entry;
 import java.util.Objects;
+import java.util.Optional;
 import java.util.Set;
 import java.util.TreeMap;
 import java.util.function.Function;
@@ -103,6 +106,8 @@ public class DdlDiff {
   private final MapDifference<String, ASTcreate_search_index_statement> searchIndexDifferences;
   private final String databaseName; // for alter Database
   private final MapDifference<String, ASTcreate_schema_statement> schemaDifferences;
+  private final MapDifference<String, ASTcreate_locality_group_statement>
+      localityGroupDifferences;
 
   private DdlDiff(DatabaseDefinition originalDb, DatabaseDefinition newDb, String databaseName)
       throws DdlDiffException {
@@ -122,6 +127,8 @@ private DdlDiff(DatabaseDefinition originalDb, DatabaseDefinition newDb, String
     this.searchIndexDifferences =
         Maps.difference(originalDb.searchIndexes(), newDb.searchIndexes());
     this.schemaDifferences = Maps.difference(originalDb.schemas(), newDb.schemas());
+    this.localityGroupDifferences =
+        Maps.difference(originalDb.localityGroups(), newDb.localityGroups());
 
     if (!alterDatabaseOptionsDifferences.areEqual() && Strings.isNullOrEmpty(databaseName)) {
       // should never happen, but...
@@ -197,6 +204,15 @@ public List<String> generateDifferenceStatements(Map<String, Boolean> options)
       }
     }
 
+    // Drop deleted locality groups.
+    if (options.get(ALLOW_DROP_STATEMENTS_OPT)) {
+      for (ASTcreate_locality_group_statement lg :
+          localityGroupDifferences.entriesOnlyOnLeft().values()) {
+        LOG.info("Dropping deleted locality group: {}", lg.getNameOrDefault());
+        output.add("DROP LOCALITY GROUP " + lg.getNameOrDefault());
+      }
+    }
+
     // drop deleted search indexes.
     if (options.get(ALLOW_DROP_STATEMENTS_OPT)) {
       for (String searchIndexName : searchIndexDifferences.entriesOnlyOnLeft().keySet()) {
@@ -273,6 +289,36 @@ public List<String> generateDifferenceStatements(Map<String, Boolean> options)
           generateAlterTableStatements(difference.leftValue(), difference.rightValue(), options));
     }
 
+    // update existing locality groups (options only)
+    for (ValueDifference<ASTcreate_locality_group_statement> lgDiff :
+        localityGroupDifferences.entriesDiffering().values()) {
+      ASTcreate_locality_group_statement left = lgDiff.leftValue();
+      ASTcreate_locality_group_statement right = lgDiff.rightValue();
+
+      // Only OPTIONS diffs are supported
+      ASToptions_clause leftOptionsClause = left.getOptionsClause();
+      ASToptions_clause rightOptionsClause = right.getOptionsClause();
+
+      Map<String, String> leftOptions =
+          leftOptionsClause == null ? Collections.emptyMap() : leftOptionsClause.getKeyValueMap();
+      Map<String, String> rightOptions =
+          rightOptionsClause == null ? Collections.emptyMap() : rightOptionsClause.getKeyValueMap();
+      MapDifference<String, String> optionsDiff = Maps.difference(leftOptions, rightOptions);
+
+      String updateText = generateOptionsUpdates(optionsDiff);
+      if (!Strings.isNullOrEmpty(updateText)) {
+        output.add(
+            "ALTER LOCALITY GROUP " + right.getNameOrDefault() + " SET OPTIONS (" + updateText + ")");
+      }
+    }
+
+    // Create new locality groups
+    for (ASTcreate_locality_group_statement lg :
+        localityGroupDifferences.entriesOnlyOnRight().values()) {
+      LOG.info("Creating new locality group: {}", lg.getNameOrDefault());
+      output.add(lg.toStringOptionalExistClause(false));
+    }
+
     // create schemas
     for (ASTcreate_schema_statement schema : schemaDifferences.entriesOnlyOnRight().values()) {
       LOG.info("creating schema: {}", schema.getName());
@@ -456,36 +502,51 @@ static List<String> generateAlterTableStatements(
     //   - change not null on column.
     // note that constraints need to be dropped before columns, and created after columns.
 
-    // Check interleaving has not changed.
-    if (left.getInterleaveClause().isPresent() != right.getInterleaveClause().isPresent()) {
-      throw new DdlDiffException("Cannot change interleaving on table " + left.getTableName());
-    }
-
-    if (left.getInterleaveClause().isPresent()
-        && !left.getInterleaveClause()
-            .get()
-            .getParentTableName()
-            .equals(right.getInterleaveClause().get().getParentTableName())) {
-      throw new DdlDiffException(
-          "Cannot change interleaved parent of table " + left.getTableName());
-    }
-
     // Check Key is same
     if (!left.getPrimaryKey().toString().equals(right.getPrimaryKey().toString())) {
       throw new DdlDiffException("Cannot change primary key of table " + left.getTableName());
     }
 
-    // On delete changed
-    if (left.getInterleaveClause().isPresent()
-        && !left.getInterleaveClause()
-            .get()
-            .getOnDelete()
-            .equals(right.getInterleaveClause().get().getOnDelete())) {
-      alterStatements.add(
-          "ALTER TABLE "
-              + left.getTableName()
-              + " SET "
-              + right.getInterleaveClause().get().getOnDelete());
+    // Interleaving changed (added, parent changed, or ON DELETE changed)
+    final Optional<ASTtable_interleave_clause> leftInterleave = left.getInterleaveClause();
+    final Optional<ASTtable_interleave_clause> rightInterleave = right.getInterleaveClause();
+
+    if (leftInterleave.isPresent() != rightInterleave.isPresent()) {
+      if (rightInterleave.isPresent()) {
+        // Added interleave
+        ASTtable_interleave_clause ri = rightInterleave.get();
+        alterStatements.add(
+            Joiner.on(" ")
+                .skipNulls()
+                .join(
+                    "ALTER TABLE",
+                    left.getTableName(),
+                    "SET INTERLEAVE IN",
+                    (ri.hasParentKeyword() ? "PARENT" : null),
+                    ri.getParentTableName(),
+                    ri.getOnDelete()));
+      } else {
+        // Removal not supported
+        throw new DdlDiffException(
+            "Cannot change interleaving on table " + left.getTableName());
+      }
+    } else if (leftInterleave.isPresent()) {
+      ASTtable_interleave_clause li = leftInterleave.get();
+      ASTtable_interleave_clause ri = rightInterleave.get();
+      if (!li.getParentTableName().equals(ri.getParentTableName())
+          || !li.getOnDelete().equals(ri.getOnDelete())
+          || li.hasParentKeyword() != ri.hasParentKeyword()) {
+        alterStatements.add(
+            Joiner.on(" ")
+                .skipNulls()
+                .join(
+                    "ALTER TABLE",
+                    left.getTableName(),
+                    "SET INTERLEAVE IN",
+                    (ri.hasParentKeyword() ? "PARENT" : null),
+                    ri.getParentTableName(),
+                    ri.getOnDelete()));
+      }
     }
 
     // compare columns.
@@ -815,6 +876,7 @@ public static List<ASTddl_statement> parseDdl(String original, boolean parseAnno
           case DdlParserTreeConstants.JJTALTER_DATABASE_STATEMENT:
           case DdlParserTreeConstants.JJTCREATE_CHANGE_STREAM_STATEMENT:
           case DdlParserTreeConstants.JJTCREATE_SEARCH_INDEX_STATEMENT:
+          case DdlParserTreeConstants.JJTCREATE_LOCALITY_GROUP_STATEMENT:
             // no-op - allowed
             break;
           case DdlParserTreeConstants.JJTCREATE_OR_REPLACE_STATEMENT:

src/main/java/com/google/cloud/solutions/spannerddl/parser/ASTcreate_table_statement.java

@@ -126,7 +126,8 @@ public String toStringOptionalExistClause(boolean includeExists) {
                     getOptionalChildByType(children, ASTtable_interleave_clause.class),
                     (withConstraints
                         ? getOptionalChildByType(children, ASTrow_deletion_policy_clause.class)
-                        : null)));
+                        : null),
+                    getOptionalChildByType(children, ASToptions_clause.class)));
   }
 
   private void validateChildren() {
@@ -141,6 +142,7 @@ private void validateChildren() {
             ASTprimary_key.class,
             ASTtable_interleave_clause.class,
             ASTrow_deletion_policy_clause.class,
+            ASToptions_clause.class,
             ASTannotation.class));
   }
 

src/main/java/com/google/cloud/solutions/spannerddl/parser/ASTtable_interleave_clause.java

@@ -30,20 +30,40 @@ public ASTtable_interleave_clause(DdlParser p, int id) {
   }
 
   public String getParentTableName() {
-    return AstTreeUtils.tokensToString((ASTinterleave_in) children[0]);
+    for (Node child : children) {
+      if (child instanceof ASTinterleave_in) {
+        return AstTreeUtils.tokensToString((ASTinterleave_in) child);
+      }
+    }
+    return null;
+  }
+
+  public boolean hasParentKeyword() {
+    for (Node child : children) {
+      if (child instanceof ASTparent) {
+        return true;
+      }
+    }
+    return false;
   }
 
   public String getOnDelete() {
-    if (children.length == 2) {
-      // verify child type
-      return children[1].toString();
-    } else {
-      return ASTon_delete_clause.ON_DELETE_NO_ACTION;
+    for (Node child : children) {
+      if (child instanceof ASTon_delete_clause) {
+        return child.toString();
+      }
     }
+    return ASTon_delete_clause.ON_DELETE_NO_ACTION;
   }
 
   @Override
   public String toString() {
-    return Joiner.on(" ").join("INTERLEAVE IN PARENT", getParentTableName(), getOnDelete());
+    return Joiner.on(" ")
+        .skipNulls()
+        .join(
+            "INTERLEAVE IN",
+            (hasParentKeyword() ? "PARENT" : null),
+            getParentTableName(),
+            getOnDelete());
   }
 }

src/main/jjtree-sources/ddl_keywords.jjt

@@ -149,6 +149,7 @@ TOKEN:
   | <INTERLEAVE:               "interleave">
   | <INVOKER:                  "invoker">
   | <JSON:                     "json">
+  | <LOCALITY:                 "locality">
   | <KEY:                      "key">
   | <LAST:                     "last">
   | <MAX:                      "max">
@@ -225,6 +226,7 @@ void pseudoReservedWord() #void :
   | <INTERLEAVE>
   | <INVOKER>
   | <JSON>
+  | <LOCALITY>
   | <KEY>
   | <LAST>
   | <MAX>

src/main/jjtree-sources/ddl_parser.jjt

@@ -111,6 +111,7 @@ void create_statement() #void :
     | create_or_replace_statement()
     | create_change_stream_statement()
     | create_sequence_statement()
+    | create_locality_group_statement()
     | create_role_statement()
   )
 }
@@ -204,6 +205,7 @@ void create_table_statement() :
   primary_key()
   [ LOOKAHEAD(2) "," table_interleave_clause() ]
   [ LOOKAHEAD(2) "," row_deletion_policy_clause() ]
+  [ LOOKAHEAD(2) "," options_clause() ]
 }
 
 void table_element() #void :
@@ -349,7 +351,7 @@ void synonym_clause() :
 void table_interleave_clause() :
 {}
 {
-  <INTERLEAVE> <IN> <PARENT> qualified_identifier() #interleave_in
+  <INTERLEAVE> <IN> [ ( <PARENT> #parent ) ] qualified_identifier() #interleave_in
     [ on_delete_clause() ]
 }
 
@@ -520,6 +522,7 @@ void target_type():
   | <CHANGE> <STREAM> #change_stream
   | <VIEW> #view
   | <SEQUENCE> #sequence
+  | <LOCALITY> <GROUP> #locality_group
 }
 
 void privilege_target():
@@ -561,6 +564,15 @@ void create_sequence_statement() :
   [ options_clause() ]
 }
 
+void create_locality_group_statement() :
+{}
+{
+  <LOCALITY> <GROUP>
+  [ if_not_exists() ]
+  ( <DEFAULTT> #defaultt | qualified_identifier() #name )
+  [ options_clause() ]
+}
+
 void key() #void :
 {}
 {
@@ -701,6 +713,7 @@ void drop_statement() :
       | <SCHEMA> #schema [ if_exists() ]
     ) qualified_identifier() #name
     | <PROTO> #proto_bundle <BUNDLE> #name
+    | <LOCALITY> <GROUP> #locality_group [ if_exists() ] ( <DEFAULTT> #defaultt | qualified_identifier() #name )
   )
 }
 
@@ -718,6 +731,7 @@ void alter_statement() #void :
     | alter_sequence_statement()
     | alter_model_statement()
     | alter_schema_statement()
+    | alter_locality_group_statement()
   )
 }
 
@@ -796,6 +810,15 @@ void alter_model_statement() :
   [<SET> options_clause()]
 }
 
+void alter_locality_group_statement() :
+{}
+{
+  <LOCALITY> <GROUP>
+  [ if_exists() ]
+  ( <DEFAULTT> #defaultt | qualified_identifier() #name )
+  [ <SET> options_clause() ]
+}
+
 void analyze_statement() :
 {}
 {