diff --git a/v1/src/main/java/com/google/cloud/teleport/spanner/AvroSchemaToDdlConverter.java b/v1/src/main/java/com/google/cloud/teleport/spanner/AvroSchemaToDdlConverter.java index 561b15155a7..8e4fe47edb1 100644 --- a/v1/src/main/java/com/google/cloud/teleport/spanner/AvroSchemaToDdlConverter.java +++ b/v1/src/main/java/com/google/cloud/teleport/spanner/AvroSchemaToDdlConverter.java @@ -29,6 +29,7 @@ import static com.google.cloud.teleport.spanner.AvroUtil.SPANNER_DYNAMIC_PROPERTIES; import static com.google.cloud.teleport.spanner.AvroUtil.SPANNER_EDGE_TABLE; import static com.google.cloud.teleport.spanner.AvroUtil.SPANNER_ENTITY; +import static com.google.cloud.teleport.spanner.AvroUtil.SPANNER_ENTITY_LOCALITY_GROUP; import static com.google.cloud.teleport.spanner.AvroUtil.SPANNER_ENTITY_MODEL; import static com.google.cloud.teleport.spanner.AvroUtil.SPANNER_ENTITY_PLACEMENT; import static com.google.cloud.teleport.spanner.AvroUtil.SPANNER_ENTITY_PROPERTY_GRAPH; @@ -73,6 +74,7 @@ import com.google.cloud.teleport.spanner.ddl.Ddl; import com.google.cloud.teleport.spanner.ddl.GraphElementTable; import com.google.cloud.teleport.spanner.ddl.GraphElementTable.Kind; +import com.google.cloud.teleport.spanner.ddl.LocalityGroup; import com.google.cloud.teleport.spanner.ddl.Model; import com.google.cloud.teleport.spanner.ddl.NamedSchema; import com.google.cloud.teleport.spanner.ddl.Placement; @@ -126,6 +128,8 @@ public Ddl toDdl(Collection avroSchemas) { builder.addSchema(toSchema(null, schema)); } else if (SPANNER_ENTITY_PLACEMENT.equals(schema.getProp(SPANNER_ENTITY))) { builder.addPlacement(toPlacement(null, schema)); + } else if (SPANNER_ENTITY_LOCALITY_GROUP.equals(schema.getProp(SPANNER_ENTITY))) { + builder.addLocalityGroup(toLocalityGroup(null, schema)); } else if (SPANNER_ENTITY_PROPERTY_GRAPH.equals(schema.getProp(SPANNER_ENTITY))) { builder.addPropertyGraph(toPropertyGraph(null, schema)); } else { @@ -579,6 +583,18 @@ public Placement toPlacement(String placementName, Schema schema) { return builder.build(); } + public LocalityGroup toLocalityGroup(String localityGroupName, Schema schema) { + if (localityGroupName == null) { + localityGroupName = getSpannerObjectName(schema); + } + LOG.debug("Converting to Ddl localityGroupName {}", localityGroupName); + + LocalityGroup.Builder builder = LocalityGroup.builder(dialect).name(localityGroupName); + builder.options(getNumberedPropsWithPrefix(schema, SPANNER_OPTION)); + + return builder.build(); + } + public Table toTable(String tableName, Schema schema) { if (tableName == null) { tableName = getSpannerObjectName(schema); @@ -691,6 +707,8 @@ public Table toTable(String tableName, Schema schema) { table.checkConstraints(getNumberedPropsWithPrefix(schema, SPANNER_CHECK_CONSTRAINT)); + table.options(getNumberedPropsWithPrefix(schema, SPANNER_OPTION)); + // Table parent options. String spannerParent = schema.getProp(SPANNER_PARENT); if (!Strings.isNullOrEmpty(spannerParent)) { diff --git a/v1/src/main/java/com/google/cloud/teleport/spanner/AvroUtil.java b/v1/src/main/java/com/google/cloud/teleport/spanner/AvroUtil.java index 4c3e6447a74..48b025786a1 100644 --- a/v1/src/main/java/com/google/cloud/teleport/spanner/AvroUtil.java +++ b/v1/src/main/java/com/google/cloud/teleport/spanner/AvroUtil.java @@ -39,6 +39,7 @@ private AvroUtil() {} public static final String SPANNER_ENTITY_MODEL = "Model"; public static final String SPANNER_ENTITY_PROPERTY_GRAPH = "PropertyGraph"; public static final String SPANNER_ENTITY_PLACEMENT = "Placement"; + public static final String SPANNER_ENTITY_LOCALITY_GROUP = "LocalityGroup"; public static final String SPANNER_FOREIGN_KEY = "spannerForeignKey_"; public static final String SPANNER_INDEX = "spannerIndex_"; public static final String SPANNER_ON_DELETE_ACTION = "spannerOnDeleteAction"; diff --git a/v1/src/main/java/com/google/cloud/teleport/spanner/DdlToAvroSchemaConverter.java b/v1/src/main/java/com/google/cloud/teleport/spanner/DdlToAvroSchemaConverter.java index a001fce20e2..ece776efd4a 100644 --- a/v1/src/main/java/com/google/cloud/teleport/spanner/DdlToAvroSchemaConverter.java +++ b/v1/src/main/java/com/google/cloud/teleport/spanner/DdlToAvroSchemaConverter.java @@ -31,6 +31,7 @@ import static com.google.cloud.teleport.spanner.AvroUtil.SPANNER_DYNAMIC_PROPERTIES; import static com.google.cloud.teleport.spanner.AvroUtil.SPANNER_EDGE_TABLE; import static com.google.cloud.teleport.spanner.AvroUtil.SPANNER_ENTITY; +import static com.google.cloud.teleport.spanner.AvroUtil.SPANNER_ENTITY_LOCALITY_GROUP; import static com.google.cloud.teleport.spanner.AvroUtil.SPANNER_ENTITY_MODEL; import static com.google.cloud.teleport.spanner.AvroUtil.SPANNER_ENTITY_PLACEMENT; import static com.google.cloud.teleport.spanner.AvroUtil.SPANNER_ENTITY_PROPERTY_GRAPH; @@ -73,6 +74,7 @@ import com.google.cloud.teleport.spanner.ddl.Ddl; import com.google.cloud.teleport.spanner.ddl.GraphElementTable; import com.google.cloud.teleport.spanner.ddl.IndexColumn; +import com.google.cloud.teleport.spanner.ddl.LocalityGroup; import com.google.cloud.teleport.spanner.ddl.Model; import com.google.cloud.teleport.spanner.ddl.ModelColumn; import com.google.cloud.teleport.spanner.ddl.NamedSchema; @@ -206,6 +208,11 @@ public Collection convert(Ddl ddl) { for (int i = 0; i < table.checkConstraints().size(); i++) { recordBuilder.prop(SPANNER_CHECK_CONSTRAINT + i, table.checkConstraints().get(i)); } + if (table.options() != null) { + for (int i = 0; i < table.options().size(); i++) { + recordBuilder.prop(SPANNER_OPTION + i, table.options().get(i)); + } + } SchemaBuilder.FieldAssembler fieldsAssembler = recordBuilder.fields(); int columnOrdinal = 0; for (Column cm : table.columns()) { @@ -426,6 +433,23 @@ public Collection convert(Ddl ddl) { } schemas.add(recordBuilder.fields().endRecord()); } + + for (LocalityGroup localityGroup : ddl.localityGroups()) { + LOG.info("DdlToAvro localityGroup {}", localityGroup.name()); + SchemaBuilder.RecordBuilder recordBuilder = + SchemaBuilder.record(generateAvroSchemaName(localityGroup.name())) + .namespace(this.namespace); + recordBuilder.prop(SPANNER_NAME, localityGroup.name()); + recordBuilder.prop(GOOGLE_FORMAT_VERSION, version); + recordBuilder.prop(GOOGLE_STORAGE, "CloudSpanner"); + recordBuilder.prop(SPANNER_ENTITY, SPANNER_ENTITY_LOCALITY_GROUP); + if (localityGroup.options() != null) { + for (int i = 0; i < localityGroup.options().size(); i++) { + recordBuilder.prop(SPANNER_OPTION + i, localityGroup.options().get(i)); + } + } + schemas.add(recordBuilder.fields().endRecord()); + } return schemas; } diff --git a/v1/src/main/java/com/google/cloud/teleport/spanner/ExportTransform.java b/v1/src/main/java/com/google/cloud/teleport/spanner/ExportTransform.java index a0bad4695ed..f1c5439390d 100644 --- a/v1/src/main/java/com/google/cloud/teleport/spanner/ExportTransform.java +++ b/v1/src/main/java/com/google/cloud/teleport/spanner/ExportTransform.java @@ -23,6 +23,7 @@ import com.google.cloud.spanner.TimestampBound; import com.google.cloud.teleport.spanner.ddl.ChangeStream; import com.google.cloud.teleport.spanner.ddl.Ddl; +import com.google.cloud.teleport.spanner.ddl.LocalityGroup; import com.google.cloud.teleport.spanner.ddl.Model; import com.google.cloud.teleport.spanner.ddl.NamedSchema; import com.google.cloud.teleport.spanner.ddl.Placement; @@ -435,6 +436,21 @@ public void processElement(ProcessContext c) { } })); + PCollection allLocalityGroupNames = + ddl.apply( + "List all locality group names", + ParDo.of( + new DoFn() { + + @ProcessElement + public void processElement(ProcessContext c) { + Ddl ddl = c.element(); + for (LocalityGroup localityGroup : ddl.localityGroups()) { + c.output(localityGroup.name()); + } + } + })); + PCollection allUdfNames = ddl.apply( "List all user-defined function names", @@ -675,6 +691,24 @@ public void processElement(ProcessContext c) { } })); + PCollection>> localityGroups = + allLocalityGroupNames.apply( + "Export locality groups", + ParDo.of( + new DoFn>>() { + + @ProcessElement + public void processElement(ProcessContext c) { + String localityGroupName = c.element(); + LOG.info("Exporting locality group: " + localityGroupName); + // This file will contain the schema definition for the locality group. + c.output( + KV.of( + localityGroupName, + Collections.singleton(localityGroupName + ".avro-00000-of-00001"))); + } + })); + PCollection>> udfs = allUdfNames.apply( "Export user-defined functions", @@ -700,6 +734,7 @@ public void processElement(ProcessContext c) { .and(sequences) .and(namedSchemas) .and(placements) + .and(localityGroups) .and(udfs) .and(propertyGraphs) .apply("Combine all empty schema files", Flatten.pCollections()); @@ -1046,6 +1081,8 @@ public void processElement( exportManifest.addSequences(obj); } else if (ddl.placement(obj.getName()) != null) { exportManifest.addPlacements(obj); + } else if (ddl.localityGroup(obj.getName()) != null) { + exportManifest.addLocalityGroups(obj); } else if (ddl.udf(obj.getName()) != null) { exportManifest.addUdfs(obj); } else { diff --git a/v1/src/main/java/com/google/cloud/teleport/spanner/ImportTransform.java b/v1/src/main/java/com/google/cloud/teleport/spanner/ImportTransform.java index d19d0f09cd1..5c266e28c78 100644 --- a/v1/src/main/java/com/google/cloud/teleport/spanner/ImportTransform.java +++ b/v1/src/main/java/com/google/cloud/teleport/spanner/ImportTransform.java @@ -23,6 +23,7 @@ import com.google.cloud.spanner.Mutation; import com.google.cloud.teleport.spanner.ddl.ChangeStream; import com.google.cloud.teleport.spanner.ddl.Ddl; +import com.google.cloud.teleport.spanner.ddl.LocalityGroup; import com.google.cloud.teleport.spanner.ddl.NamedSchema; import com.google.cloud.teleport.spanner.ddl.Placement; import com.google.cloud.teleport.spanner.ddl.PropertyGraph; @@ -519,6 +520,7 @@ public void processElement(ProcessContext c) { List> missingSequences = new ArrayList<>(); List> missingUdfs = new ArrayList<>(); List> missingPlacements = new ArrayList<>(); + List> missingLocalityGroups = new ArrayList<>(); List> missingPropertyGraphs = new ArrayList<>(); for (KV kv : avroSchemas) { if (informationSchemaDdl.schema(kv.getKey()) == null @@ -529,6 +531,7 @@ public void processElement(ProcessContext c) { && informationSchemaDdl.udf(kv.getKey()) == null && informationSchemaDdl.sequence(kv.getKey()) == null && informationSchemaDdl.placement(kv.getKey()) == null + && informationSchemaDdl.localityGroup(kv.getKey()) == null && informationSchemaDdl.propertyGraph(kv.getKey()) == null) { Schema schema = parser.parse(kv.getValue()); if (schema.getProp(AvroUtil.SPANNER_CHANGE_STREAM_FOR_CLAUSE) @@ -548,6 +551,9 @@ public void processElement(ProcessContext c) { missingNamedSchemas.add(KV.of(kv.getKey(), schema)); } else if ("Placement".equals(schema.getProp("spannerEntity"))) { missingPlacements.add(KV.of(kv.getKey(), schema)); + } else if (AvroUtil.SPANNER_ENTITY_LOCALITY_GROUP.equals( + schema.getProp("spannerEntity"))) { + missingLocalityGroups.add(KV.of(kv.getKey(), schema)); } else if (AvroUtil.SPANNER_ENTITY_PROPERTY_GRAPH.equals( schema.getProp("spannerEntity"))) { missingPropertyGraphs.add(KV.of(kv.getKey(), schema)); @@ -640,6 +646,18 @@ public void processElement(ProcessContext c) { ddlStatements.addAll(newDdl.createPlacementStatements()); } + if (!missingLocalityGroups.isEmpty()) { + Ddl.Builder builder = Ddl.builder(dialect); + for (KV kv : missingLocalityGroups) { + LocalityGroup localityGroup = + converter.toLocalityGroup(kv.getKey(), kv.getValue()); + builder.addLocalityGroup(localityGroup); + mergedDdl.addLocalityGroup(localityGroup); + } + Ddl newDdl = builder.build(); + ddlStatements.addAll(newDdl.createLocalityGroupStatements()); + } + if (!missingTables.isEmpty() || !missingModels.isEmpty() || !missingViews.isEmpty() @@ -855,6 +873,12 @@ public void processElement(ProcessContext c) { c.output(KV.of(placement.getName(), fullPath)); } } + for (Export.Table localityGroup : proto.getLocalityGroupsList()) { + for (String f : localityGroup.getDataFilesList()) { + String fullPath = GcsUtil.joinPath(importDirectory.get(), f); + c.output(KV.of(localityGroup.getName(), fullPath)); + } + } } })); @@ -892,6 +916,11 @@ public void processElement(ProcessContext c) { c.output(KV.of(placement.getName(), placement.getManifestFile())); } } + for (Export.Table localityGroup : proto.getLocalityGroupsList()) { + if (!Strings.isNullOrEmpty(localityGroup.getManifestFile())) { + c.output(KV.of(localityGroup.getName(), localityGroup.getManifestFile())); + } + } } })); diff --git a/v1/src/main/java/com/google/cloud/teleport/spanner/ddl/Column.java b/v1/src/main/java/com/google/cloud/teleport/spanner/ddl/Column.java index 483bd1f5aab..fa9a956e189 100644 --- a/v1/src/main/java/com/google/cloud/teleport/spanner/ddl/Column.java +++ b/v1/src/main/java/com/google/cloud/teleport/spanner/ddl/Column.java @@ -160,12 +160,30 @@ public void prettyPrint(Appendable appendable) throws IOException { if (isHidden()) { appendable.append(" HIDDEN"); } - if (columnOptions() == null) { + if (columnOptions() == null || columnOptions().isEmpty()) { return; } - String optionsString = String.join(",", columnOptions()); - if (!optionsString.isEmpty()) { - appendable.append(" OPTIONS (").append(optionsString).append(")"); + if (dialect() == Dialect.POSTGRESQL) { + for (String option : columnOptions()) { + String[] parts = option.split("=", 2); + if (parts.length == 2 && parts[0].trim().equalsIgnoreCase("locality_group")) { + String lgName = parts[1].trim(); + if (lgName.length() >= 2 + && ((lgName.startsWith("'") && lgName.endsWith("'")) + || (lgName.startsWith("\"") && lgName.endsWith("\"")))) { + lgName = lgName.substring(1, lgName.length() - 1); + } + appendable + .append(" LOCALITY GROUP ") + .append( + "NULL".equalsIgnoreCase(lgName) ? "NULL" : quoteIdentifier(lgName, dialect())); + } + } + } else { + String optionsString = String.join(",", columnOptions()); + if (!optionsString.isEmpty()) { + appendable.append(" OPTIONS (").append(optionsString).append(")"); + } } } diff --git a/v1/src/main/java/com/google/cloud/teleport/spanner/ddl/Ddl.java b/v1/src/main/java/com/google/cloud/teleport/spanner/ddl/Ddl.java index c5033ecf6ea..9965067df17 100644 --- a/v1/src/main/java/com/google/cloud/teleport/spanner/ddl/Ddl.java +++ b/v1/src/main/java/com/google/cloud/teleport/spanner/ddl/Ddl.java @@ -56,6 +56,7 @@ public class Ddl implements Serializable { private ImmutableSortedMap changeStreams; private ImmutableSortedMap sequences; private ImmutableSortedMap placements; + private ImmutableSortedMap localityGroups; private ImmutableSortedMap schemas; private TreeMultimap parents; // This is only populated by InformationSchemaScanner and not while reading from AVRO files. @@ -74,6 +75,7 @@ private Ddl( ImmutableSortedMap changeStreams, ImmutableSortedMap sequences, ImmutableSortedMap placements, + ImmutableSortedMap localityGroups, ImmutableSortedMap schemas, TreeMultimap parents, TreeMultimap referencedTables, @@ -89,6 +91,7 @@ private Ddl( this.changeStreams = changeStreams; this.sequences = sequences; this.placements = placements; + this.localityGroups = localityGroups; this.schemas = schemas; this.parents = parents; this.referencedTables = referencedTables; @@ -208,6 +211,14 @@ public Placement placement(String placementName) { return placements.get(placementName.toLowerCase()); } + public Collection localityGroups() { + return localityGroups.values(); + } + + public LocalityGroup localityGroup(String localityGroupName) { + return localityGroups.get(localityGroupName.toLowerCase()); + } + public Collection schemas() { return schemas.values(); } @@ -301,6 +312,11 @@ public void prettyPrint(Appendable appendable) throws IOException { appendable.append("\n"); placement.prettyPrint(appendable); } + + for (LocalityGroup localityGroup : localityGroups()) { + appendable.append("\n"); + localityGroup.prettyPrint(appendable); + } } public List statements() { @@ -313,6 +329,7 @@ public List statements() { builder .addAll(createNamedSchemaStatements()) .addAll(createSequenceStatements()) + .addAll(createLocalityGroupStatements()) .addAll(createTableStatements()) .addAll(createIndexStatements()) .addAll(addForeignKeyStatements()) @@ -436,6 +453,14 @@ public List createPlacementStatements() { return result; } + public List createLocalityGroupStatements() { + List result = new ArrayList<>(localityGroups.size()); + for (LocalityGroup localityGroup : localityGroups()) { + result.add(localityGroup.prettyPrint()); + } + return result; + } + public String createProtoBundleStatement() { StringBuilder appendable = new StringBuilder(); String quote = NameUtils.identifierQuote(dialect); @@ -537,6 +562,7 @@ public static class Builder { private Map changeStreams = Maps.newLinkedHashMap(); private Map sequences = Maps.newLinkedHashMap(); private Map placements = Maps.newLinkedHashMap(); + private Map localityGroups = Maps.newLinkedHashMap(); private Map schemas = Maps.newLinkedHashMap(); private TreeMultimap parents = TreeMultimap.create(); private TreeMultimap referencedTables = TreeMultimap.create(); @@ -685,6 +711,22 @@ public boolean hasPlacement(String name) { return placements.containsKey(name.toLowerCase()); } + public LocalityGroup.Builder createLocalityGroup(String name) { + LocalityGroup localityGroup = localityGroups.get(name.toLowerCase()); + if (localityGroup == null) { + return LocalityGroup.builder(dialect).name(name).ddlBuilder(this); + } + return localityGroup.toBuilder().ddlBuilder(this); + } + + public void addLocalityGroup(LocalityGroup localityGroup) { + localityGroups.put(localityGroup.name().toLowerCase(), localityGroup); + } + + public boolean hasLocalityGroup(String name) { + return localityGroups.containsKey(name.toLowerCase()); + } + public NamedSchema.Builder createSchema(String name) { NamedSchema schema = schemas.get(name.toLowerCase()); if (schema == null) { @@ -746,6 +788,7 @@ public Ddl build() { ImmutableSortedMap.copyOf(changeStreams), ImmutableSortedMap.copyOf(sequences), ImmutableSortedMap.copyOf(placements), + ImmutableSortedMap.copyOf(localityGroups), ImmutableSortedMap.copyOf(schemas), parents, referencedTables, @@ -767,6 +810,7 @@ public Builder toBuilder() { builder.changeStreams.putAll(changeStreams); builder.sequences.putAll(sequences); builder.placements.putAll(placements); + builder.localityGroups.putAll(localityGroups); builder.parents.putAll(parents); builder.referencedTables.putAll(referencedTables); builder.databaseOptions = databaseOptions; @@ -825,6 +869,11 @@ public boolean equals(Object o) { if (placements != null ? !placements.equals(ddl.placements) : ddl.placements != null) { return false; } + if (localityGroups != null + ? !localityGroups.equals(ddl.localityGroups) + : ddl.localityGroups != null) { + return false; + } if (protoDescriptors != null ? !protoDescriptors.equals(ddl.protoDescriptors) : ddl.protoDescriptors != null) { @@ -846,6 +895,7 @@ public int hashCode() { result = 31 * result + (changeStreams != null ? changeStreams.hashCode() : 0); result = 31 * result + (sequences != null ? sequences.hashCode() : 0); result = 31 * result + (placements != null ? placements.hashCode() : 0); + result = 31 * result + (localityGroups != null ? localityGroups.hashCode() : 0); result = 31 * result + (databaseOptions != null ? databaseOptions.hashCode() : 0); result = 31 * result + (protoBundle != null ? protoBundle.hashCode() : 0); result = 31 * result + (protoDescriptors != null ? protoDescriptors.hashCode() : 0); diff --git a/v1/src/main/java/com/google/cloud/teleport/spanner/ddl/Index.java b/v1/src/main/java/com/google/cloud/teleport/spanner/ddl/Index.java index d003e6696c3..bba598fd08c 100644 --- a/v1/src/main/java/com/google/cloud/teleport/spanner/ddl/Index.java +++ b/v1/src/main/java/com/google/cloud/teleport/spanner/ddl/Index.java @@ -139,14 +139,47 @@ private void prettyPrintPg(Appendable appendable) throws IOException { } } - if (interleaveIn() != null) { - appendable.append(" INTERLEAVE IN ").append(quoteIdentifier(interleaveIn(), dialect())); + String localityGroup = null; + ImmutableList.Builder otherOptionsBuilder = ImmutableList.builder(); + if (options() != null) { + for (String option : options()) { + String[] parts = option.split("=", 2); + if (parts.length == 2 && parts[0].trim().equalsIgnoreCase("locality_group")) { + String lgName = parts[1].trim(); + if (lgName.length() >= 2 + && ((lgName.startsWith("'") && lgName.endsWith("'")) + || (lgName.startsWith("\"") && lgName.endsWith("\"")))) { + lgName = lgName.substring(1, lgName.length() - 1); + } + localityGroup = lgName; + } else { + otherOptionsBuilder.add(option); + } + } } + ImmutableList otherOptions = otherOptionsBuilder.build(); - if (options() != null) { - String optionsString = String.join(",", options()); - if (!optionsString.isEmpty()) { - appendable.append(" WITH (").append(optionsString).append(")"); + if ("SEARCH".equals(type())) { + if (interleaveIn() != null) { + appendable.append(" INTERLEAVE IN ").append(quoteIdentifier(interleaveIn(), dialect())); + } + if (!otherOptions.isEmpty()) { + appendable.append(" WITH (").append(String.join(",", otherOptions)).append(")"); + } + } else { + if (!otherOptions.isEmpty()) { + appendable.append(" WITH (").append(String.join(",", otherOptions)).append(")"); + } + if (localityGroup != null) { + appendable + .append(" LOCALITY GROUP ") + .append( + "NULL".equalsIgnoreCase(localityGroup) + ? "NULL" + : quoteIdentifier(localityGroup, dialect())); + } + if (interleaveIn() != null) { + appendable.append(" INTERLEAVE IN ").append(quoteIdentifier(interleaveIn(), dialect())); } } 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 index e9f6a70bb6d..377331d1e8f 100644 --- 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 @@ -94,6 +94,7 @@ public Ddl scan() { } listColumns(builder); listColumnOptions(builder); + listTableOptions(builder); if (isModelSupported()) { listModels(builder); listModelOptions(builder); @@ -116,6 +117,9 @@ public Ddl scan() { if (placementsSupported()) { listPlacements(builder); } + if (isLocalityGroupSupported()) { + listLocalityGroups(builder); + } if (isPropertyGraphSupported()) { listPropertyGraphs(builder); listPropertyGraphPropertyDeclarations(builder); @@ -659,7 +663,10 @@ void listIndexOptions( Map, ImmutableList.Builder> allOptions = Maps.newHashMap(); while (resultSet.next()) { String tableName = getQualifiedName(resultSet.getString(0), resultSet.getString(1)); - String indexName = resultSet.getString(2); + String indexName = + dialect == Dialect.POSTGRESQL + ? resultSet.getString(2) + : getQualifiedName(resultSet.getString(0), resultSet.getString(2)); String indexType = resultSet.getString(3); String optionName = resultSet.getString(4); if (optionName.startsWith("internal_updated_") @@ -678,7 +685,8 @@ void listIndexOptions( dialect == Dialect.POSTGRESQL ? POSTGRESQL_LITERAL_QUOTE : GSQL_LITERAL_QUOTE; options.add( optionName + "=" + quoteChar + OPTION_STRING_ESCAPER.escape(optionValue) + quoteChar); - } else if (optionType.equalsIgnoreCase("character varying")) { + } else if (optionType.equalsIgnoreCase("character varying") + || optionType.equalsIgnoreCase("text")) { options.add(optionName + "='" + OPTION_STRING_ESCAPER.escape(optionValue) + "'"); } else { options.add(optionName + "=" + optionValue); @@ -729,7 +737,66 @@ Statement listIndexOptionsSQL() { } } - private void listColumnOptions(Ddl.Builder builder) { + @VisibleForTesting + void listTableOptions(Ddl.Builder builder) { + Statement statement = listTableOptionsSQL(); + + Map> allOptions = Maps.newHashMap(); + try (ResultSet resultSet = context.executeQuery(statement)) { + while (resultSet.next()) { + String tableName = getQualifiedName(resultSet.getString(0), resultSet.getString(1)); + String optionName = resultSet.getString(2); + String optionType = resultSet.getString(3); + String optionValue = resultSet.getString(4); + + ImmutableList.Builder options = + allOptions.computeIfAbsent(tableName, k -> ImmutableList.builder()); + + if (optionType.equalsIgnoreCase("STRING")) { + String quoteChar = + dialect == Dialect.POSTGRESQL ? POSTGRESQL_LITERAL_QUOTE : GSQL_LITERAL_QUOTE; + options.add( + optionName + "=" + quoteChar + OPTION_STRING_ESCAPER.escape(optionValue) + quoteChar); + } else if (optionType.equalsIgnoreCase("character varying") + || optionType.equalsIgnoreCase("text")) { + options.add(optionName + "='" + OPTION_STRING_ESCAPER.escape(optionValue) + "'"); + } else { + options.add(optionName + "=" + optionValue); + } + } + } + + for (Map.Entry> entry : allOptions.entrySet()) { + String tableName = entry.getKey(); + ImmutableList options = entry.getValue().build(); + builder.createTable(tableName).options(options).endTable(); + } + } + + @VisibleForTesting + Statement listTableOptionsSQL() { + switch (dialect) { + case GOOGLE_STANDARD_SQL: + return Statement.of( + "SELECT t.table_schema, t.table_name, t.option_name, t.option_type, t.option_value" + + " FROM information_schema.table_options AS t" + + " WHERE t.table_schema NOT IN" + + " ('INFORMATION_SCHEMA', 'SPANNER_SYS')" + + " ORDER BY t.table_name, t.option_name"); + case POSTGRESQL: + return Statement.of( + "SELECT t.table_schema, t.table_name, t.option_name, t.option_type, t.option_value" + + " FROM information_schema.table_options AS t" + + " WHERE t.table_schema NOT IN" + + " ('information_schema', 'spanner_sys', 'pg_catalog')" + + " ORDER BY t.table_name, t.option_name"); + default: + throw new IllegalArgumentException("Unrecognized dialect: " + dialect); + } + } + + @VisibleForTesting + void listColumnOptions(Ddl.Builder builder) { Statement statement = listColumnOptionsSQL(); ResultSet resultSet = context.executeQuery(statement); @@ -741,19 +808,28 @@ private void listColumnOptions(Ddl.Builder builder) { String optionName = resultSet.getString(3); String optionType = resultSet.getString(4); String optionValue = resultSet.getString(5); + boolean isInherited = + !resultSet.isNull(6) + && (dialect == Dialect.GOOGLE_STANDARD_SQL + ? resultSet.getBoolean(6) + : "YES".equalsIgnoreCase(resultSet.getString(6))); - KV kv = KV.of(tableName, columnName); - - if (optionName.equalsIgnoreCase("locality_group")) { + if (isInherited) { continue; } + KV kv = KV.of(tableName, columnName); + ImmutableList.Builder options = allOptions.computeIfAbsent(kv, k -> ImmutableList.builder()); if (optionType.equalsIgnoreCase("STRING")) { - options.add(optionName + "=\"" + OPTION_STRING_ESCAPER.escape(optionValue) + "\""); - } else if (optionType.equalsIgnoreCase("character varying")) { + String quoteChar = + dialect == Dialect.POSTGRESQL ? POSTGRESQL_LITERAL_QUOTE : GSQL_LITERAL_QUOTE; + options.add( + optionName + "=" + quoteChar + OPTION_STRING_ESCAPER.escape(optionValue) + quoteChar); + } else if (optionType.equalsIgnoreCase("character varying") + || optionType.equalsIgnoreCase("text")) { options.add(optionName + "='" + OPTION_STRING_ESCAPER.escape(optionValue) + "'"); } else { options.add(optionName + "=" + optionValue); @@ -780,7 +856,7 @@ Statement listColumnOptionsSQL() { case GOOGLE_STANDARD_SQL: return Statement.of( "SELECT t.table_schema, t.table_name, t.column_name, t.option_name, t.option_type," - + " t.option_value" + + " t.option_value, t.is_inherited" + " FROM information_schema.column_options AS t" + " WHERE t.table_schema NOT IN" + " ('INFORMATION_SCHEMA', 'SPANNER_SYS')" @@ -789,7 +865,7 @@ Statement listColumnOptionsSQL() { // Ignore the 'allow_commit_timestamp' option since it's not user-settable in POSTGRESQL. return Statement.of( "SELECT t.table_schema, t.table_name, t.column_name, t.option_name, t.option_type," - + " t.option_value" + + " t.option_value, t.is_inherited" + " FROM information_schema.column_options AS t" + " WHERE t.table_schema NOT IN " + " ('information_schema', 'spanner_sys', 'pg_catalog')" @@ -2112,6 +2188,100 @@ private void listPlacements(Ddl.Builder builder) { } } + @VisibleForTesting + boolean isLocalityGroupSupported() { + Statement statement; + switch (dialect) { + case GOOGLE_STANDARD_SQL: + statement = + Statement.of( + "SELECT COUNT(1)" + + " FROM INFORMATION_SCHEMA.TABLES t WHERE " + + " t.TABLE_SCHEMA = 'INFORMATION_SCHEMA'" + + " AND t.TABLE_NAME = 'LOCALITY_GROUP_OPTIONS'"); + break; + case POSTGRESQL: + statement = + Statement.of( + "SELECT COUNT(1)" + + " FROM INFORMATION_SCHEMA.TABLES t WHERE " + + " t.TABLE_SCHEMA = 'information_schema'" + + " AND t.TABLE_NAME = 'locality_group_options'"); + break; + default: + throw new IllegalArgumentException("Unrecognized dialect: " + dialect); + } + + try (ResultSet resultSet = context.executeQuery(statement)) { + resultSet.next(); + return resultSet.getLong(0) > 0; + } + } + + @VisibleForTesting + Statement listLocalityGroupsSQL() { + switch (dialect) { + case GOOGLE_STANDARD_SQL: + case POSTGRESQL: + return Statement.of( + "SELECT t.locality_group_name, t.option_name, t.option_value, t.is_inherited" + + " FROM information_schema.locality_group_options AS t" + + " ORDER BY t.locality_group_name, t.option_name"); + default: + throw new IllegalArgumentException("Unrecognized dialect: " + dialect); + } + } + + @VisibleForTesting + void listLocalityGroups(Ddl.Builder builder) { + Map> localityGroupNameToOptions = Maps.newHashMap(); + try (ResultSet resultSet = context.executeQuery(listLocalityGroupsSQL())) { + while (resultSet.next()) { + String name = resultSet.getString(0); + String optionName = resultSet.getString(1); + String optionValue = resultSet.isNull(2) ? null : resultSet.getString(2); + boolean isInherited = + !resultSet.isNull(3) + && (dialect == Dialect.GOOGLE_STANDARD_SQL + ? resultSet.getBoolean(3) + : "YES".equalsIgnoreCase(resultSet.getString(3))); + + if (!"default".equalsIgnoreCase(name)) { + localityGroupNameToOptions.computeIfAbsent(name, k -> ImmutableList.builder()); + } + + if (optionValue == null || isInherited) { + continue; + } + + LOG.info( + "locality group name = " + + name + + ", optionName = " + + optionName + + ", optionValue = " + + optionValue + + ", isInherited = " + + isInherited); + + ImmutableList.Builder options = + localityGroupNameToOptions.computeIfAbsent(name, k -> ImmutableList.builder()); + + String quoteChar = + dialect == Dialect.POSTGRESQL ? POSTGRESQL_LITERAL_QUOTE : GSQL_LITERAL_QUOTE; + options.add( + optionName + "=" + quoteChar + OPTION_STRING_ESCAPER.escape(optionValue) + quoteChar); + } + } + + for (Map.Entry> entry : + localityGroupNameToOptions.entrySet()) { + String localityGroupName = entry.getKey(); + ImmutableList options = entry.getValue().build(); + builder.createLocalityGroup(localityGroupName).options(options).endLocalityGroup(); + } + } + private boolean isUnknownType(DescriptorProto descriptor) { MessageOptions messageOptions = descriptor.getOptions(); // 14004 is the extension for placeholder descriptors for unknown types in spanner. diff --git a/v1/src/main/java/com/google/cloud/teleport/spanner/ddl/LocalityGroup.java b/v1/src/main/java/com/google/cloud/teleport/spanner/ddl/LocalityGroup.java new file mode 100644 index 00000000000..e421cdb5c29 --- /dev/null +++ b/v1/src/main/java/com/google/cloud/teleport/spanner/ddl/LocalityGroup.java @@ -0,0 +1,137 @@ +/* + * Copyright (C) 2026 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); you may not + * use this file except in compliance with the License. You may obtain a copy of + * the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT + * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the + * License for the specific language governing permissions and limitations under + * the License. + */ +package com.google.cloud.teleport.spanner.ddl; + +import static com.google.cloud.teleport.spanner.common.NameUtils.quoteIdentifier; + +import com.google.auto.value.AutoValue; +import com.google.cloud.spanner.Dialect; +import com.google.common.collect.ImmutableList; +import java.io.IOException; +import java.io.Serializable; +import javax.annotation.Nullable; + +/** Cloud Spanner locality group. */ +@AutoValue +public abstract class LocalityGroup implements Serializable { + + private static final long serialVersionUID = 1L; + + public abstract String name(); + + @Nullable + public abstract ImmutableList options(); + + public abstract Dialect dialect(); + + public abstract Builder toBuilder(); + + public static Builder builder(Dialect dialect) { + return new AutoValue_LocalityGroup.Builder().dialect(dialect); + } + + public static Builder builder() { + return builder(Dialect.GOOGLE_STANDARD_SQL); + } + + public void prettyPrint(Appendable appendable) throws IOException { + if (dialect() != Dialect.GOOGLE_STANDARD_SQL && dialect() != Dialect.POSTGRESQL) { + throw new IllegalArgumentException(String.format("Unrecognized Dialect: %s.", dialect())); + } + + boolean isDefault = "default".equalsIgnoreCase(name()); + if (isDefault && (options() == null || options().isEmpty())) { + return; + } + appendable + .append(isDefault ? "ALTER LOCALITY GROUP " : "CREATE LOCALITY GROUP ") + .append(quoteIdentifier(name(), dialect())); + + if (options() != null && !options().isEmpty()) { + if (dialect() == Dialect.GOOGLE_STANDARD_SQL) { + String optionsString = String.join(", ", options()); + appendable + .append(isDefault ? "\n\tSET OPTIONS (" : "\n\tOPTIONS (") + .append(optionsString) + .append(")"); + } else { + String storageClause = null; + String spillClause = null; + for (String option : options()) { + String[] parts = option.split("=", 2); + if (parts.length != 2) { + continue; + } + String key = parts[0].trim(); + String val = parts[1].trim(); + if (val.length() >= 2 && val.startsWith("\"") && val.endsWith("\"")) { + val = "'" + val.substring(1, val.length() - 1) + "'"; + } + if (key.equalsIgnoreCase("storage")) { + storageClause = "STORAGE " + val; + } else if (key.equalsIgnoreCase("ssd_to_hdd_spill_timespan")) { + spillClause = "SSD_TO_HDD_SPILL_TIMESPAN " + val; + } + } + if (storageClause != null) { + appendable.append(" ").append(storageClause); + } + if (spillClause != null) { + appendable.append(" ").append(spillClause); + } + } + } + } + + public String prettyPrint() { + StringBuilder sb = new StringBuilder(); + try { + prettyPrint(sb); + } catch (IOException e) { + throw new RuntimeException(e); + } + return sb.toString(); + } + + @Override + public String toString() { + return prettyPrint(); + } + + /** A builder for {@link LocalityGroup}. */ + @AutoValue.Builder + public abstract static class Builder { + private Ddl.Builder ddlBuilder; + + public Builder ddlBuilder(Ddl.Builder ddlBuilder) { + this.ddlBuilder = ddlBuilder; + return this; + } + + public abstract Builder name(String name); + + public abstract Builder options(ImmutableList options); + + abstract Builder dialect(Dialect dialect); + + public abstract LocalityGroup build(); + + public Ddl.Builder endLocalityGroup() { + ddlBuilder.addLocalityGroup(build()); + return ddlBuilder; + } + } +} diff --git a/v1/src/main/java/com/google/cloud/teleport/spanner/ddl/Table.java b/v1/src/main/java/com/google/cloud/teleport/spanner/ddl/Table.java index d111c3f9b5b..b8f42cd2fbe 100644 --- a/v1/src/main/java/com/google/cloud/teleport/spanner/ddl/Table.java +++ b/v1/src/main/java/com/google/cloud/teleport/spanner/ddl/Table.java @@ -59,6 +59,9 @@ public enum InterleaveType { public abstract ImmutableList checkConstraints(); + @Nullable + public abstract ImmutableList options(); + public abstract Builder autoToBuilder(); public abstract Dialect dialect(); @@ -81,6 +84,7 @@ public static Builder builder(Dialect dialect) { .indexes(ImmutableList.of()) .foreignKeys(ImmutableList.of()) .checkConstraints(ImmutableList.of()) + .options(ImmutableList.of()) .onDeleteCascade(false); } @@ -102,6 +106,15 @@ public void prettyPrint(Appendable appendable, boolean includeIndexes, boolean i } } + private static String unquoteOptionValue(String val) { + if (val.length() >= 2 + && ((val.startsWith("'") && val.endsWith("'")) + || (val.startsWith("\"") && val.endsWith("\"")))) { + return val.substring(1, val.length() - 1); + } + return val; + } + private void prettyPrintPg( Appendable appendable, boolean includeIndexes, boolean includeForeignKeys) throws IOException { @@ -123,6 +136,18 @@ private void prettyPrintPg( .collect(Collectors.joining(", ", "\n\tPRIMARY KEY (", ")"))); } appendable.append("\n)"); + if (options() != null && !options().isEmpty()) { + for (String option : options()) { + String[] parts = option.split("=", 2); + if (parts.length == 2 && parts[0].trim().equalsIgnoreCase("locality_group")) { + String lgName = unquoteOptionValue(parts[1].trim()); + appendable + .append(" LOCALITY GROUP ") + .append( + "NULL".equalsIgnoreCase(lgName) ? "NULL" : quoteIdentifier(lgName, dialect())); + } + } + } if (interleaveInParent() != null) { appendable .append(" \nINTERLEAVE IN ") @@ -172,6 +197,10 @@ private void prettyPrintGsql( appendable.append(" ON DELETE CASCADE"); } } + if (options() != null && !options().isEmpty()) { + String optionsString = String.join(", ", options()); + appendable.append(",\nOPTIONS (").append(optionsString).append(")"); + } if (includeIndexes) { appendable.append("\n"); appendable.append(String.join("\n", indexes())); @@ -230,6 +259,8 @@ Builder ddlBuilder(Ddl.Builder ddlBuilder) { public abstract Builder checkConstraints(ImmutableList checkConstraints); + public abstract Builder options(ImmutableList options); + abstract Builder dialect(Dialect dialect); public abstract Dialect dialect(); diff --git a/v1/src/main/proto/export.proto b/v1/src/main/proto/export.proto index 37f129531e0..bea6d5bf65f 100644 --- a/v1/src/main/proto/export.proto +++ b/v1/src/main/proto/export.proto @@ -63,6 +63,8 @@ message Export { repeated Table placements = 8; // The set of user-defined functions that were exported. repeated Table udfs = 9; + // The set of locality groups that were exported. + repeated Table locality_groups = 10; // Contains a protobuf-serialized google.protobuf.FileDescriptorSet for // proto_descriptors (from spanner_database_admin.proto) diff --git a/v1/src/test/java/com/google/cloud/teleport/spanner/AvroSchemaToDdlConverterTest.java b/v1/src/test/java/com/google/cloud/teleport/spanner/AvroSchemaToDdlConverterTest.java index 582019e6fbd..a3cf5844bb6 100644 --- a/v1/src/test/java/com/google/cloud/teleport/spanner/AvroSchemaToDdlConverterTest.java +++ b/v1/src/test/java/com/google/cloud/teleport/spanner/AvroSchemaToDdlConverterTest.java @@ -1506,6 +1506,179 @@ public void placements() { + " OPTIONS (instance_partition='mr-partition', default_leader='us-east1')")); } + @Test + public void localityGroups() { + String defaultLgAvroString = + "{" + + " \"type\" : \"record\"," + + " \"name\" : \"default\"," + + " \"fields\" : []," + + " \"namespace\" : \"spannertest\"," + + " \"googleStorage\" : \"CloudSpanner\"," + + " \"spannerEntity\":\"LocalityGroup\", " + + " \"googleFormatVersion\" : \"booleans\"," + + " \"spannerOption_0\" : \"storage=\\\"hdd\\\"\"" + + "}"; + String lgAvroString = + "{" + + " \"type\" : \"record\"," + + " \"name\" : \"LocalityGroup1\"," + + " \"fields\" : []," + + " \"namespace\" : \"spannertest\"," + + " \"googleStorage\" : \"CloudSpanner\"," + + " \"spannerEntity\":\"LocalityGroup\", " + + " \"googleFormatVersion\" : \"booleans\"," + + " \"spannerOption_0\" : \"storage=\\\"ssd\\\"\"," + + " \"spannerOption_1\" : \"ssd_to_hdd_spill_timespan=\\\"30d\\\"\"" + + "}"; + String emptyLgAvroString = + "{" + + " \"type\" : \"record\"," + + " \"name\" : \"lg_empty\"," + + " \"fields\" : []," + + " \"namespace\" : \"spannertest\"," + + " \"googleStorage\" : \"CloudSpanner\"," + + " \"spannerEntity\":\"LocalityGroup\", " + + " \"googleFormatVersion\" : \"booleans\"" + + "}"; + String tableAvroString = + "{" + + " \"type\" : \"record\"," + + " \"name\" : \"Users\"," + + " \"namespace\" : \"spannertest\"," + + " \"fields\" : [ {" + + " \"name\" : \"id\"," + + " \"type\" : \"long\"," + + " \"sqlType\" : \"INT64\"" + + " }, {" + + " \"name\" : \"description\"," + + " \"type\" : [ \"null\", \"string\" ]," + + " \"sqlType\" : \"STRING(MAX)\"" + + " }, {" + + " \"name\" : \"name\"," + + " \"type\" : [ \"null\", \"string\" ]," + + " \"sqlType\" : \"STRING(MAX)\"," + + " \"spannerOption_0\" : \"locality_group=\\\"default\\\"\"" + + " } ]," + + " \"googleStorage\" : \"CloudSpanner\"," + + " \"googleFormatVersion\" : \"booleans\"," + + " \"spannerPrimaryKey_0\" : \"`id` ASC\"," + + " \"spannerOption_0\" : \"locality_group=\\\"LocalityGroup1\\\"\"," + + " \"spannerIndex_0\" : \"CREATE INDEX `UsersByName` ON `Users`(`name` ASC) OPTIONS (locality_group=\\\"LocalityGroup1\\\")\"" + + "}"; + Collection schemas = new ArrayList<>(); + Schema.Parser parser = new Schema.Parser(); + schemas.add(parser.parse(defaultLgAvroString)); + schemas.add(parser.parse(lgAvroString)); + schemas.add(parser.parse(emptyLgAvroString)); + schemas.add(parser.parse(tableAvroString)); + + AvroSchemaToDdlConverter converter = new AvroSchemaToDdlConverter(); + Ddl ddl = converter.toDdl(schemas); + assertThat(ddl.localityGroups(), hasSize(3)); + assertThat(ddl.allTables(), hasSize(1)); + assertThat( + ddl.prettyPrint(), + equalToCompressingWhiteSpace( + "CREATE TABLE `Users` (" + + " `id` INT64 NOT NULL," + + " `description` STRING(MAX)," + + " `name` STRING(MAX) OPTIONS (locality_group=\"default\")," + + " ) PRIMARY KEY (`id` ASC)," + + " OPTIONS (locality_group=\"LocalityGroup1\")" + + " CREATE INDEX `UsersByName` ON `Users`(`name` ASC) OPTIONS (locality_group=\"LocalityGroup1\")" + + " ALTER LOCALITY GROUP `default` SET OPTIONS (storage=\"hdd\")" + + " CREATE LOCALITY GROUP `lg_empty`" + + " CREATE LOCALITY GROUP `LocalityGroup1`" + + " OPTIONS (storage=\"ssd\", ssd_to_hdd_spill_timespan=\"30d\")")); + } + + @Test + public void pgLocalityGroups() { + String defaultLgAvroString = + "{" + + " \"type\" : \"record\"," + + " \"name\" : \"default\"," + + " \"fields\" : []," + + " \"namespace\" : \"spannertest\"," + + " \"googleStorage\" : \"CloudSpanner\"," + + " \"spannerEntity\":\"LocalityGroup\", " + + " \"googleFormatVersion\" : \"booleans\"," + + " \"spannerOption_0\" : \"storage='hdd'\"" + + "}"; + String lgAvroString = + "{" + + " \"type\" : \"record\"," + + " \"name\" : \"LocalityGroup1\"," + + " \"fields\" : []," + + " \"namespace\" : \"spannertest\"," + + " \"googleStorage\" : \"CloudSpanner\"," + + " \"spannerEntity\":\"LocalityGroup\", " + + " \"googleFormatVersion\" : \"booleans\"," + + " \"spannerOption_0\" : \"ssd_to_hdd_spill_timespan='30d'\"," + + " \"spannerOption_1\" : \"storage='ssd'\"" + + "}"; + String emptyLgAvroString = + "{" + + " \"type\" : \"record\"," + + " \"name\" : \"lg_empty\"," + + " \"fields\" : []," + + " \"namespace\" : \"spannertest\"," + + " \"googleStorage\" : \"CloudSpanner\"," + + " \"spannerEntity\":\"LocalityGroup\", " + + " \"googleFormatVersion\" : \"booleans\"" + + "}"; + String tableAvroString = + "{" + + " \"type\" : \"record\"," + + " \"name\" : \"Users\"," + + " \"namespace\" : \"spannertest\"," + + " \"fields\" : [ {" + + " \"name\" : \"id\"," + + " \"type\" : \"long\"," + + " \"sqlType\" : \"bigint\"" + + " }, {" + + " \"name\" : \"description\"," + + " \"type\" : [ \"null\", \"string\" ]," + + " \"sqlType\" : \"character varying\"" + + " }, {" + + " \"name\" : \"name\"," + + " \"type\" : [ \"null\", \"string\" ]," + + " \"sqlType\" : \"character varying\"," + + " \"spannerOption_0\" : \"locality_group='default'\"" + + " } ]," + + " \"googleStorage\" : \"CloudSpanner\"," + + " \"googleFormatVersion\" : \"booleans\"," + + " \"spannerPrimaryKey_0\" : \"\\\"id\\\" ASC\"," + + " \"spannerOption_0\" : \"locality_group='LocalityGroup1'\"," + + " \"spannerIndex_0\" : \"CREATE INDEX \\\"UsersByName\\\" ON \\\"Users\\\"(\\\"name\\\" ASC) LOCALITY GROUP \\\"LocalityGroup1\\\"\"" + + "}"; + Collection schemas = new ArrayList<>(); + Schema.Parser parser = new Schema.Parser(); + schemas.add(parser.parse(defaultLgAvroString)); + schemas.add(parser.parse(lgAvroString)); + schemas.add(parser.parse(emptyLgAvroString)); + schemas.add(parser.parse(tableAvroString)); + + AvroSchemaToDdlConverter converter = new AvroSchemaToDdlConverter(Dialect.POSTGRESQL); + Ddl ddl = converter.toDdl(schemas); + assertThat(ddl.localityGroups(), hasSize(3)); + assertThat(ddl.allTables(), hasSize(1)); + assertThat( + ddl.prettyPrint(), + equalToCompressingWhiteSpace( + "CREATE TABLE \"Users\" (" + + " \"id\" bigint NOT NULL," + + " \"description\" character varying," + + " \"name\" character varying LOCALITY GROUP \"default\"," + + " PRIMARY KEY (\"id\")" + + " ) LOCALITY GROUP \"LocalityGroup1\"" + + " CREATE INDEX \"UsersByName\" ON \"Users\"(\"name\" ASC) LOCALITY GROUP \"LocalityGroup1\"" + + " ALTER LOCALITY GROUP \"default\" STORAGE 'hdd'" + + " CREATE LOCALITY GROUP \"lg_empty\"" + + " CREATE LOCALITY GROUP \"LocalityGroup1\" STORAGE 'ssd' SSD_TO_HDD_SPILL_TIMESPAN '30d'")); + } + @Test public void placementTable() { String placementKeyAsPrimaryKeyAvroString = diff --git a/v1/src/test/java/com/google/cloud/teleport/spanner/DdlToAvroSchemaConverterTest.java b/v1/src/test/java/com/google/cloud/teleport/spanner/DdlToAvroSchemaConverterTest.java index 3a946dfea76..519411c2c8e 100644 --- a/v1/src/test/java/com/google/cloud/teleport/spanner/DdlToAvroSchemaConverterTest.java +++ b/v1/src/test/java/com/google/cloud/teleport/spanner/DdlToAvroSchemaConverterTest.java @@ -31,6 +31,7 @@ import static com.google.cloud.teleport.spanner.AvroUtil.SPANNER_DYNAMIC_PROPERTIES; import static com.google.cloud.teleport.spanner.AvroUtil.SPANNER_EDGE_TABLE; import static com.google.cloud.teleport.spanner.AvroUtil.SPANNER_ENTITY; +import static com.google.cloud.teleport.spanner.AvroUtil.SPANNER_ENTITY_LOCALITY_GROUP; import static com.google.cloud.teleport.spanner.AvroUtil.SPANNER_ENTITY_MODEL; import static com.google.cloud.teleport.spanner.AvroUtil.SPANNER_ENTITY_PLACEMENT; import static com.google.cloud.teleport.spanner.AvroUtil.SPANNER_ENTITY_PROPERTY_GRAPH; @@ -2385,6 +2386,93 @@ public void placements() { equalTo("default_leader=\"" + defaultLeader + "\"")); } + @Test + public void localityGroups() { + DdlToAvroSchemaConverter converter = + new DdlToAvroSchemaConverter("spannertest", "booleans", false); + String lgName1 = "lg1"; + String lgName2 = "lg2"; + String lgEmpty = "lg_empty"; + Ddl ddl = + Ddl.builder() + .createLocalityGroup(lgName1) + .options(ImmutableList.of("storage=\"ssd\"")) + .endLocalityGroup() + .createLocalityGroup(lgName2) + .options(ImmutableList.of("storage=\"ssd\"", "ssd_to_hdd_spill_timespan=\"30d\"")) + .endLocalityGroup() + .createLocalityGroup(lgEmpty) + .endLocalityGroup() + .createTable("Users") + .column("id") + .int64() + .notNull() + .endColumn() + .column("description") + .string() + .max() + .endColumn() + .column("name") + .string() + .max() + .columnOptions(ImmutableList.of("locality_group=\"lg2\"")) + .endColumn() + .primaryKey() + .asc("id") + .end() + .options(ImmutableList.of("locality_group=\"lg1\"")) + .indexes( + ImmutableList.of( + "CREATE INDEX `UsersByName` ON `Users`(`name` ASC) OPTIONS (locality_group=\"lg2\")")) + .endTable() + .build(); + + Collection result = converter.convert(ddl); + assertThat(result, hasSize(4)); + + Iterator it = result.iterator(); + Schema tableSchema = it.next(); + assertThat(tableSchema.getName(), equalTo("Users")); + assertThat(tableSchema.getProp(SPANNER_OPTION + "0"), equalTo("locality_group=\"lg1\"")); + assertThat( + tableSchema.getProp(SPANNER_INDEX + "0"), + equalTo( + "CREATE INDEX `UsersByName` ON `Users`(`name` ASC) OPTIONS (locality_group=\"lg2\")")); + assertThat(tableSchema.getField("description").getProp(SPANNER_OPTION + "0"), nullValue()); + assertThat( + tableSchema.getField("name").getProp(SPANNER_OPTION + "0"), + equalTo("locality_group=\"lg2\"")); + + Schema avroSchema1 = it.next(); + assertThat(avroSchema1.getNamespace(), equalTo("spannertest")); + assertThat(avroSchema1.getProp(GOOGLE_FORMAT_VERSION), equalTo("booleans")); + assertThat(avroSchema1.getProp(GOOGLE_STORAGE), equalTo("CloudSpanner")); + assertThat(avroSchema1.getProp(SPANNER_ENTITY), equalTo(SPANNER_ENTITY_LOCALITY_GROUP)); + assertThat(avroSchema1.getFields(), empty()); + assertThat(avroSchema1.getName(), equalTo(lgName1)); + assertThat(avroSchema1.getProp(SPANNER_OPTION + "0"), equalTo("storage=\"ssd\"")); + + Schema avroSchema2 = it.next(); + assertThat(avroSchema2.getNamespace(), equalTo("spannertest")); + assertThat(avroSchema2.getProp(GOOGLE_FORMAT_VERSION), equalTo("booleans")); + assertThat(avroSchema2.getProp(GOOGLE_STORAGE), equalTo("CloudSpanner")); + assertThat(avroSchema2.getProp(SPANNER_ENTITY), equalTo(SPANNER_ENTITY_LOCALITY_GROUP)); + assertThat(avroSchema2.getFields(), empty()); + assertThat(avroSchema2.getName(), equalTo(lgName2)); + assertThat(avroSchema2.getProp(SPANNER_OPTION + "0"), equalTo("storage=\"ssd\"")); + assertThat( + avroSchema2.getProp(SPANNER_OPTION + "1"), equalTo("ssd_to_hdd_spill_timespan=\"30d\"")); + + Schema avroSchemaEmpty = it.next(); + assertThat(avroSchemaEmpty.getNamespace(), equalTo("spannertest")); + assertThat(avroSchemaEmpty.getProp(GOOGLE_FORMAT_VERSION), equalTo("booleans")); + assertThat(avroSchemaEmpty.getProp(GOOGLE_STORAGE), equalTo("CloudSpanner")); + assertThat(avroSchemaEmpty.getProp(SPANNER_ENTITY), equalTo(SPANNER_ENTITY_LOCALITY_GROUP)); + assertThat(avroSchemaEmpty.getFields(), empty()); + assertThat(avroSchemaEmpty.getName(), equalTo(lgEmpty)); + assertThat(avroSchemaEmpty.getProp(SPANNER_OPTION + "0"), nullValue()); + } + @Test public void placementTable() { DdlToAvroSchemaConverter converter = diff --git a/v1/src/test/java/com/google/cloud/teleport/spanner/ExportTransformTest.java b/v1/src/test/java/com/google/cloud/teleport/spanner/ExportTransformTest.java index 905e848799a..5f041629201 100644 --- a/v1/src/test/java/com/google/cloud/teleport/spanner/ExportTransformTest.java +++ b/v1/src/test/java/com/google/cloud/teleport/spanner/ExportTransformTest.java @@ -161,6 +161,8 @@ public void buildDatabaseManifestFile() throws InvalidProtocolBufferException { "function manifest", "placement", "placement manifest", + "localityGroup", + "localityGroup manifest", "propertyGraph1", "propertyGraph1 manifest"); @@ -201,6 +203,10 @@ public void buildDatabaseManifestFile() throws InvalidProtocolBufferException { .options( ImmutableList.of("instance_partition=\"mr-partition\"", "default_leader=\"us-east1\"")) .endPlacement(); + ddlBuilder + .createLocalityGroup("localityGroup") + .options(ImmutableList.of("storage=\"ssd\"")) + .endLocalityGroup(); ddlBuilder.mergeProtoBundle(protoBundle); ddlBuilder.mergeProtoDescriptors(protoDescriptors); Ddl ddl = ddlBuilder.build(); @@ -262,6 +268,12 @@ public void buildDatabaseManifestFile() throws InvalidProtocolBufferException { manifestProto.getPlacements(0).getManifestFile(), is("placement-manifest.json")); + assertThat(manifestProto.getLocalityGroupsCount(), is(1)); + assertThat(manifestProto.getLocalityGroups(0).getName(), is("localityGroup")); + assertThat( + manifestProto.getLocalityGroups(0).getManifestFile(), + is("localityGroup-manifest.json")); + assertThat(manifestProto.getSequencesCount(), is(1)); assertThat(manifestProto.getSequences(0).getName(), is("sequence")); assertThat( diff --git a/v1/src/test/java/com/google/cloud/teleport/spanner/ddl/DdlTest.java b/v1/src/test/java/com/google/cloud/teleport/spanner/ddl/DdlTest.java index 480bf82c342..9e1c2589a9a 100644 --- a/v1/src/test/java/com/google/cloud/teleport/spanner/ddl/DdlTest.java +++ b/v1/src/test/java/com/google/cloud/teleport/spanner/ddl/DdlTest.java @@ -1564,6 +1564,221 @@ public void pgSequences() { assertNotNull(ddl.hashCode()); } + @Test + public void localityGroups() { + Ddl ddl = + Ddl.builder() + .createLocalityGroup("default") + .options(ImmutableList.of("storage=\"hdd\"")) + .endLocalityGroup() + .createLocalityGroup("lg1") + .options(ImmutableList.of("storage=\"ssd\"")) + .endLocalityGroup() + .createLocalityGroup("lg2") + .options(ImmutableList.of("storage=\"ssd\"", "ssd_to_hdd_spill_timespan=\"30d\"")) + .endLocalityGroup() + .build(); + assertThat( + ddl.prettyPrint(), + equalToCompressingWhiteSpace( + "ALTER LOCALITY GROUP `default` SET OPTIONS (storage=\"hdd\")" + + " CREATE LOCALITY GROUP `lg1` OPTIONS (storage=\"ssd\")" + + " CREATE LOCALITY GROUP `lg2` OPTIONS (storage=\"ssd\", ssd_to_hdd_spill_timespan=\"30d\")")); + + List statements = ddl.statements(); + assertEquals(3, statements.size()); + assertThat( + statements.get(0), + equalToCompressingWhiteSpace( + "ALTER LOCALITY GROUP `default` SET OPTIONS (storage=\"hdd\")")); + assertThat( + statements.get(1), + equalToCompressingWhiteSpace("CREATE LOCALITY GROUP `lg1` OPTIONS (storage=\"ssd\")")); + assertThat( + statements.get(2), + equalToCompressingWhiteSpace( + "CREATE LOCALITY GROUP `lg2` OPTIONS (storage=\"ssd\", ssd_to_hdd_spill_timespan=\"30d\")")); + assertNotNull(ddl.hashCode()); + assertEquals(3, ddl.localityGroups().size()); + assertNotNull(ddl.localityGroup("lg1")); + } + + @Test + public void pgLocalityGroups() { + Ddl ddl = + Ddl.builder(Dialect.POSTGRESQL) + .createLocalityGroup("default") + .options(ImmutableList.of("storage='hdd'")) + .endLocalityGroup() + .createLocalityGroup("lg1") + .options(ImmutableList.of("ssd_to_hdd_spill_timespan='30d'", "storage='ssd'")) + .endLocalityGroup() + .build(); + assertThat( + ddl.prettyPrint(), + equalToCompressingWhiteSpace( + "ALTER LOCALITY GROUP \"default\" STORAGE 'hdd'" + + " CREATE LOCALITY GROUP \"lg1\" STORAGE 'ssd' SSD_TO_HDD_SPILL_TIMESPAN '30d'")); + + List statements = ddl.statements(); + assertEquals(2, statements.size()); + assertThat( + statements.get(0), + equalToCompressingWhiteSpace("ALTER LOCALITY GROUP \"default\" STORAGE 'hdd'")); + assertThat( + statements.get(1), + equalToCompressingWhiteSpace( + "CREATE LOCALITY GROUP \"lg1\" STORAGE 'ssd' SSD_TO_HDD_SPILL_TIMESPAN '30d'")); + assertNotNull(ddl.hashCode()); + assertEquals(2, ddl.localityGroups().size()); + assertNotNull(ddl.localityGroup("lg1")); + } + + @Test + public void localityGroupsBeforeTables() { + Index idx = + Index.builder() + .name("UsersByFirstName") + .table("Users") + .columns() + .create() + .name("first_name") + .asc() + .endIndexColumn() + .end() + .options(ImmutableList.of("locality_group=\"lg1\"")) + .build(); + Ddl ddl = + Ddl.builder() + .createTable("Users") + .column("id") + .int64() + .notNull() + .endColumn() + .column("first_name") + .string() + .max() + .columnOptions(ImmutableList.of("locality_group=\"default\"")) + .endColumn() + .primaryKey() + .asc("id") + .end() + .options(ImmutableList.of("locality_group=\"lg1\"")) + .indexes(ImmutableList.of(idx.prettyPrint())) + .endTable() + .createLocalityGroup("lg1") + .options(ImmutableList.of("storage=\"ssd\"")) + .endLocalityGroup() + .build(); + + List statements = ddl.statements(); + assertEquals(3, statements.size()); + assertThat( + statements.get(0), + equalToCompressingWhiteSpace("CREATE LOCALITY GROUP `lg1` OPTIONS (storage=\"ssd\")")); + assertThat( + statements.get(1), + equalToCompressingWhiteSpace( + "CREATE TABLE `Users` (" + + " `id` INT64 NOT NULL," + + " `first_name` STRING(MAX) OPTIONS (locality_group=\"default\")," + + " ) PRIMARY KEY (`id` ASC)," + + " OPTIONS (locality_group=\"lg1\")")); + assertThat( + statements.get(2), + equalToCompressingWhiteSpace( + "CREATE INDEX `UsersByFirstName` ON `Users`(`first_name` ASC) OPTIONS (locality_group=\"lg1\")")); + } + + @Test + public void pgLocalityGroupsWithTablesColumnsAndIndexes() { + Index idx1 = + Index.builder(Dialect.POSTGRESQL) + .name("UsersByFirstName") + .table("Users") + .columns() + .create() + .name("first_name") + .asc() + .endIndexColumn() + .end() + .options(ImmutableList.of("locality_group='lg1'")) + .build(); + Index idx2 = + Index.builder(Dialect.POSTGRESQL) + .name("UsersInterleavedIdx") + .table("Users") + .interleaveIn("ParentUsers") + .columns() + .create() + .name("id") + .asc() + .endIndexColumn() + .end() + .options(ImmutableList.of("locality_group='lg1'")) + .build(); + Ddl ddl = + Ddl.builder(Dialect.POSTGRESQL) + .createTable("ParentUsers") + .column("id") + .pgInt8() + .notNull() + .endColumn() + .primaryKey() + .asc("id") + .end() + .endTable() + .createTable("Users") + .column("id") + .pgInt8() + .notNull() + .endColumn() + .column("first_name") + .pgVarchar() + .columnOptions(ImmutableList.of("locality_group='default'")) + .endColumn() + .primaryKey() + .asc("id") + .end() + .interleaveInParent("ParentUsers") + .interleaveType(Table.InterleaveType.IN_PARENT) + .onDeleteCascade(true) + .options(ImmutableList.of("locality_group='lg1'")) + .indexes(ImmutableList.of(idx1.prettyPrint(), idx2.prettyPrint())) + .endTable() + .createLocalityGroup("lg1") + .options(ImmutableList.of("storage='ssd'")) + .endLocalityGroup() + .build(); + + List statements = ddl.statements(); + assertEquals(5, statements.size()); + assertThat( + statements.get(0), + equalToCompressingWhiteSpace("CREATE LOCALITY GROUP \"lg1\" STORAGE 'ssd'")); + assertThat( + statements.get(1), + equalToCompressingWhiteSpace( + "CREATE TABLE \"ParentUsers\" ( \"id\" bigint NOT NULL, PRIMARY KEY (\"id\") )")); + assertThat( + statements.get(2), + equalToCompressingWhiteSpace( + "CREATE TABLE \"Users\" (" + + " \"id\" bigint NOT NULL," + + " \"first_name\" character varying LOCALITY GROUP \"default\"," + + " PRIMARY KEY (\"id\")" + + " ) LOCALITY GROUP \"lg1\"" + + " INTERLEAVE IN PARENT \"ParentUsers\" ON DELETE CASCADE")); + assertThat( + statements.get(3), + equalToCompressingWhiteSpace( + "CREATE INDEX \"UsersByFirstName\" ON \"Users\"(\"first_name\" ASC) LOCALITY GROUP \"lg1\"")); + assertThat( + statements.get(4), + equalToCompressingWhiteSpace( + "CREATE INDEX \"UsersInterleavedIdx\" ON \"Users\"(\"id\" ASC) LOCALITY GROUP \"lg1\" INTERLEAVE IN \"ParentUsers\"")); + } + @Test public void protobundle() { Ddl.Builder builder = Ddl.builder(); 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 index 5fc57894af1..ed8c9605934 100644 --- 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 @@ -36,6 +36,7 @@ import com.google.cloud.spanner.TimestampBound; import com.google.cloud.teleport.metadata.SpannerStagingTest; import com.google.cloud.teleport.metadata.TemplateIntegrationTest; +import com.google.cloud.teleport.spanner.AvroSchemaToDdlConverter; import com.google.cloud.teleport.spanner.DdlToAvroSchemaConverter; import com.google.cloud.teleport.spanner.common.Type; import com.google.cloud.teleport.spanner.tests.TestMessage; @@ -1761,6 +1762,151 @@ private SharedTestCase pgSequences() { }); } + private SharedTestCase localityGroups() { + DdlToAvroSchemaConverter converter = + new DdlToAvroSchemaConverter("spannertest", "booleans", true); + AvroSchemaToDdlConverter toDdlConverter = + new AvroSchemaToDdlConverter(Dialect.GOOGLE_STANDARD_SQL); + List statements = + Arrays.asList( + "CREATE LOCALITY GROUP `lg_it_gsql` OPTIONS (storage = 'ssd', ssd_to_hdd_spill_timespan = '10m')", + "CREATE LOCALITY GROUP `lg_it_gsql_empty`", + "CREATE TABLE `t_localityGroups_Users` (" + + " `id` INT64 NOT NULL," + + " `col_inherited` STRING(MAX)," + + " `col_override` STRING(MAX) OPTIONS (locality_group = 'default')" + + ") PRIMARY KEY (`id` ASC), OPTIONS (locality_group = 'lg_it_gsql')", + "CREATE INDEX `idx_localityGroups_Users` ON `t_localityGroups_Users`(`col_inherited`)" + + " OPTIONS (locality_group = 'lg_it_gsql')"); + + return new SharedTestCase( + statements, + ddl -> { + assertThat(ddl.localityGroup("lg_it_gsql"), notNullValue()); + assertThat( + ddl.localityGroup("lg_it_gsql").prettyPrint(), + equalToCompressingWhiteSpace( + "CREATE LOCALITY GROUP `lg_it_gsql`" + + " OPTIONS (ssd_to_hdd_spill_timespan=\"10m\", storage=\"ssd\")")); + assertThat(ddl.localityGroup("lg_it_gsql_empty"), notNullValue()); + assertThat( + ddl.localityGroup("lg_it_gsql_empty").prettyPrint(), + equalToCompressingWhiteSpace("CREATE LOCALITY GROUP `lg_it_gsql_empty`")); + assertThat( + ddl.table("t_localityGroups_Users").prettyPrint(), + equalToCompressingWhiteSpace( + "CREATE TABLE `t_localityGroups_Users` (" + + " `id` INT64 NOT NULL," + + " `col_inherited` STRING(MAX)," + + " `col_override` STRING(MAX) OPTIONS (locality_group=\"default\")," + + " ) PRIMARY KEY (`id` ASC)," + + " OPTIONS (locality_group=\"lg_it_gsql\")" + + " CREATE INDEX `idx_localityGroups_Users` ON `t_localityGroups_Users`(`col_inherited` ASC)" + + " OPTIONS (locality_group=\"lg_it_gsql\")")); + + Collection result = converter.convert(ddl); + Map schemas = + result.stream().collect(Collectors.toMap(Schema::getName, s -> s)); + + Schema avroSchema = schemas.get("lg_it_gsql"); + assertThat(avroSchema, notNullValue()); + assertThat(avroSchema.getName(), equalTo("lg_it_gsql")); + assertThat(avroSchema.getProp("spannerEntity"), equalTo("LocalityGroup")); + assertThat( + avroSchema.getProp("spannerOption_0"), equalTo("ssd_to_hdd_spill_timespan=\"10m\"")); + assertThat(avroSchema.getProp("spannerOption_1"), equalTo("storage=\"ssd\"")); + + Schema emptyAvroSchema = schemas.get("lg_it_gsql_empty"); + assertThat(emptyAvroSchema, notNullValue()); + assertThat(emptyAvroSchema.getName(), equalTo("lg_it_gsql_empty")); + assertThat(emptyAvroSchema.getProp("spannerEntity"), equalTo("LocalityGroup")); + assertThat(emptyAvroSchema.getProp("spannerOption_0"), nullValue()); + + Ddl roundTripDdl = toDdlConverter.toDdl(result); + assertThat( + roundTripDdl.localityGroup("lg_it_gsql").prettyPrint(), + equalToCompressingWhiteSpace(ddl.localityGroup("lg_it_gsql").prettyPrint())); + assertThat( + roundTripDdl.localityGroup("lg_it_gsql_empty").prettyPrint(), + equalToCompressingWhiteSpace(ddl.localityGroup("lg_it_gsql_empty").prettyPrint())); + assertThat( + roundTripDdl.table("t_localityGroups_Users").prettyPrint(), + equalToCompressingWhiteSpace(ddl.table("t_localityGroups_Users").prettyPrint())); + }); + } + + private SharedTestCase pgLocalityGroups() { + DdlToAvroSchemaConverter converter = + new DdlToAvroSchemaConverter("spannertest", "booleans", true); + AvroSchemaToDdlConverter toDdlConverter = new AvroSchemaToDdlConverter(Dialect.POSTGRESQL); + List statements = + Arrays.asList( + "CREATE LOCALITY GROUP \"lg_it_pg\" STORAGE 'ssd' SSD_TO_HDD_SPILL_TIMESPAN '10m'", + "CREATE LOCALITY GROUP \"lg_it_pg_empty\"", + "CREATE TABLE \"t_pgLocalityGroups_Users\" (" + + " \"id\" bigint NOT NULL," + + " \"col_inherited\" character varying," + + " \"col_override\" character varying LOCALITY GROUP \"default\"," + + " PRIMARY KEY (\"id\")" + + ") LOCALITY GROUP \"lg_it_pg\"", + "CREATE INDEX \"idx_pgLocalityGroups_Users\" ON \"t_pgLocalityGroups_Users\"(\"col_inherited\")" + + " LOCALITY GROUP \"lg_it_pg\""); + + return new SharedTestCase( + statements, + ddl -> { + assertThat(ddl.localityGroup("lg_it_pg"), notNullValue()); + assertThat( + ddl.localityGroup("lg_it_pg").prettyPrint(), + equalToCompressingWhiteSpace( + "CREATE LOCALITY GROUP \"lg_it_pg\" STORAGE 'ssd' SSD_TO_HDD_SPILL_TIMESPAN '10m'")); + assertThat(ddl.localityGroup("lg_it_pg_empty"), notNullValue()); + assertThat( + ddl.localityGroup("lg_it_pg_empty").prettyPrint(), + equalToCompressingWhiteSpace("CREATE LOCALITY GROUP \"lg_it_pg_empty\"")); + assertThat( + ddl.table("t_pgLocalityGroups_Users").prettyPrint(), + equalToCompressingWhiteSpace( + "CREATE TABLE \"t_pgLocalityGroups_Users\" (" + + " \"id\" bigint NOT NULL," + + " \"col_inherited\" character varying," + + " \"col_override\" character varying LOCALITY GROUP \"default\"," + + " PRIMARY KEY (\"id\")" + + " ) LOCALITY GROUP \"lg_it_pg\"" + + " CREATE INDEX \"idx_pgLocalityGroups_Users\" ON \"t_pgLocalityGroups_Users\"(\"col_inherited\" ASC)" + + " LOCALITY GROUP \"lg_it_pg\"")); + + Collection result = converter.convert(ddl); + Map schemas = + result.stream().collect(Collectors.toMap(Schema::getName, s -> s)); + + Schema avroSchema = schemas.get("lg_it_pg"); + assertThat(avroSchema, notNullValue()); + assertThat(avroSchema.getName(), equalTo("lg_it_pg")); + assertThat(avroSchema.getProp("spannerEntity"), equalTo("LocalityGroup")); + assertThat( + avroSchema.getProp("spannerOption_0"), equalTo("ssd_to_hdd_spill_timespan='10m'")); + assertThat(avroSchema.getProp("spannerOption_1"), equalTo("storage='ssd'")); + + Schema emptyAvroSchema = schemas.get("lg_it_pg_empty"); + assertThat(emptyAvroSchema, notNullValue()); + assertThat(emptyAvroSchema.getName(), equalTo("lg_it_pg_empty")); + assertThat(emptyAvroSchema.getProp("spannerEntity"), equalTo("LocalityGroup")); + assertThat(emptyAvroSchema.getProp("spannerOption_0"), nullValue()); + + Ddl roundTripDdl = toDdlConverter.toDdl(result); + assertThat( + roundTripDdl.localityGroup("lg_it_pg").prettyPrint(), + equalToCompressingWhiteSpace(ddl.localityGroup("lg_it_pg").prettyPrint())); + assertThat( + roundTripDdl.localityGroup("lg_it_pg_empty").prettyPrint(), + equalToCompressingWhiteSpace(ddl.localityGroup("lg_it_pg_empty").prettyPrint())); + assertThat( + roundTripDdl.table("t_pgLocalityGroups_Users").prettyPrint(), + equalToCompressingWhiteSpace(ddl.table("t_pgLocalityGroups_Users").prettyPrint())); + }); + } + private SharedTestCase propertyGraphOnViewGroupBy() { List statements = Arrays.asList( @@ -1963,6 +2109,7 @@ public void sharedInformationSchemaScannerTestGsql() throws Exception { databaseOptions(), changeStreams(), sequences(), + localityGroups(), propertyGraphOnViewGroupBy(), propertyGraphOnViewSimple(), propertyGraphOnViewMixedOrder(), @@ -2002,7 +2149,8 @@ public void sharedInformationSchemaScannerTestPg() throws Exception { pgIdentityColumns(), pgDatabaseOptions(), pgChangeStreams(), - pgSequences()); + pgSequences(), + pgLocalityGroups()); List allStatements = new ArrayList<>(); for (SharedTestCase testCase : testCases) { diff --git a/v1/src/test/java/com/google/cloud/teleport/spanner/ddl/InformationSchemaScannerTest.java b/v1/src/test/java/com/google/cloud/teleport/spanner/ddl/InformationSchemaScannerTest.java index 7a649db2011..18c9313ef01 100644 --- a/v1/src/test/java/com/google/cloud/teleport/spanner/ddl/InformationSchemaScannerTest.java +++ b/v1/src/test/java/com/google/cloud/teleport/spanner/ddl/InformationSchemaScannerTest.java @@ -17,6 +17,8 @@ import static org.hamcrest.text.IsEqualCompressingWhiteSpace.equalToCompressingWhiteSpace; import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertNotNull; +import static org.junit.Assert.assertNull; import static org.junit.Assert.assertThat; import static org.junit.Assert.assertThrows; import static org.mockito.ArgumentMatchers.any; @@ -128,15 +130,16 @@ public void testListIndexOptions() { when(mockResultSet.next()).thenReturn(true, true, true, false); - // Row 1: Valid option, Row 2: internal_updated_, Row 3: system_optimized_ + // Row 1: locality_group, Row 2: internal_updated_, Row 3: system_optimized_ when(mockResultSet.getString(0)).thenReturn("", "", ""); // schema when(mockResultSet.getString(1)).thenReturn("table1", "table1", "table1"); // table when(mockResultSet.getString(2)).thenReturn("index1", "index1", "index1"); // index when(mockResultSet.getString(3)).thenReturn("INDEX", "INDEX", "INDEX"); // index type when(mockResultSet.getString(4)) - .thenReturn("valid_option", "internal_updated_opt", "system_optimized_opt"); // option name + .thenReturn( + "locality_group", "internal_updated_opt", "system_optimized_opt"); // option name when(mockResultSet.getString(5)).thenReturn("STRING", "STRING", "STRING"); // option type - when(mockResultSet.getString(6)).thenReturn("value1", "value2", "value3"); // option value + when(mockResultSet.getString(6)).thenReturn("lg1", "value2", "value3"); // option value InformationSchemaScanner scanner = new InformationSchemaScanner(mockContext, Dialect.GOOGLE_STANDARD_SQL); @@ -145,6 +148,7 @@ public void testListIndexOptions() { NavigableMap tableIndexes = Maps.newTreeMap(); Index.Builder indexBuilder = Index.builder(Dialect.GOOGLE_STANDARD_SQL).name("index1").table("table1"); + indexBuilder.columns().create().name("col1").asc().endIndexColumn().end(); tableIndexes.put("index1", indexBuilder); indexes.put("table1", tableIndexes); @@ -153,7 +157,46 @@ public void testListIndexOptions() { Index index = indexBuilder.build(); assertEquals(1, index.options().size()); - assertEquals("valid_option=\"value1\"", index.options().get(0)); + assertEquals("locality_group=\"lg1\"", index.options().get(0)); + assertThat( + index.prettyPrint(), + equalToCompressingWhiteSpace( + "CREATE INDEX `index1` ON `table1`(`col1` ASC) OPTIONS (locality_group=\"lg1\")")); + + // PostgreSQL test + ReadContext mockContextPg = mock(ReadContext.class); + ResultSet mockResultSetPg = mock(ResultSet.class); + when(mockContextPg.executeQuery(any())).thenReturn(mockResultSetPg); + + when(mockResultSetPg.next()).thenReturn(true, false); + when(mockResultSetPg.getString(0)).thenReturn("public"); + when(mockResultSetPg.getString(1)).thenReturn("table1"); + when(mockResultSetPg.getString(2)).thenReturn("index1"); + when(mockResultSetPg.getString(3)).thenReturn("INDEX"); + when(mockResultSetPg.getString(4)).thenReturn("locality_group"); + when(mockResultSetPg.getString(5)).thenReturn("character varying"); + when(mockResultSetPg.getString(6)).thenReturn("lg_pg"); + + InformationSchemaScanner scannerPg = + new InformationSchemaScanner(mockContextPg, Dialect.POSTGRESQL); + + Map> indexesPg = Maps.newHashMap(); + NavigableMap tableIndexesPg = Maps.newTreeMap(); + Index.Builder indexBuilderPg = Index.builder(Dialect.POSTGRESQL).name("index1").table("table1"); + indexBuilderPg.columns().create().name("col1").asc().endIndexColumn().end(); + tableIndexesPg.put("index1", indexBuilderPg); + indexesPg.put("table1", tableIndexesPg); + + Ddl.Builder builderPg = Ddl.builder(Dialect.POSTGRESQL); + scannerPg.listIndexOptions(builderPg, indexesPg); + + Index indexPg = indexBuilderPg.build(); + assertEquals(1, indexPg.options().size()); + assertEquals("locality_group='lg_pg'", indexPg.options().get(0)); + assertThat( + indexPg.prettyPrint(), + equalToCompressingWhiteSpace( + "CREATE INDEX \"index1\" ON \"table1\"(\"col1\" ASC) LOCALITY GROUP \"lg_pg\"")); } @Test @@ -202,13 +245,80 @@ public void testListIndexColumnsSQL() { + "ORDER BY t.table_name, t.index_name, t.ordinal_position")); } + @Test + public void testListTableOptionsSQL() { + assertThat( + googleSQLInfoScanner.listTableOptionsSQL().getSql(), + equalToCompressingWhiteSpace( + "SELECT t.table_schema, t.table_name, t.option_name, t.option_type, t.option_value" + + " FROM information_schema.table_options AS t" + + " WHERE t.table_schema NOT IN" + + " ('INFORMATION_SCHEMA', 'SPANNER_SYS')" + + " ORDER BY t.table_name, t.option_name")); + + assertThat( + postgresSQLInfoScanner.listTableOptionsSQL().getSql(), + equalToCompressingWhiteSpace( + "SELECT t.table_schema, t.table_name, t.option_name, t.option_type, t.option_value" + + " FROM information_schema.table_options AS t" + + " WHERE t.table_schema NOT IN" + + " ('information_schema', 'spanner_sys', 'pg_catalog')" + + " ORDER BY t.table_name, t.option_name")); + } + + @Test + public void testListTableOptions() { + ReadContext mockContext = mock(ReadContext.class); + ResultSet mockResultSet = mock(ResultSet.class); + when(mockContext.executeQuery(any())).thenReturn(mockResultSet); + + when(mockResultSet.next()).thenReturn(true, false); + when(mockResultSet.getString(0)).thenReturn(""); + when(mockResultSet.getString(1)).thenReturn("Users"); + when(mockResultSet.getString(2)).thenReturn("locality_group"); + when(mockResultSet.getString(3)).thenReturn("STRING"); + when(mockResultSet.getString(4)).thenReturn("lg1"); + + InformationSchemaScanner scanner = + new InformationSchemaScanner(mockContext, Dialect.GOOGLE_STANDARD_SQL); + Ddl.Builder builder = Ddl.builder(Dialect.GOOGLE_STANDARD_SQL); + builder.createTable("Users").endTable(); + scanner.listTableOptions(builder); + + Ddl ddl = builder.build(); + assertEquals(1, ddl.table("Users").options().size()); + assertEquals("locality_group=\"lg1\"", ddl.table("Users").options().get(0)); + + // PostgreSQL test + ReadContext mockContextPg = mock(ReadContext.class); + ResultSet mockResultSetPg = mock(ResultSet.class); + when(mockContextPg.executeQuery(any())).thenReturn(mockResultSetPg); + + when(mockResultSetPg.next()).thenReturn(true, false); + when(mockResultSetPg.getString(0)).thenReturn("public"); + when(mockResultSetPg.getString(1)).thenReturn("Users"); + when(mockResultSetPg.getString(2)).thenReturn("locality_group"); + when(mockResultSetPg.getString(3)).thenReturn("character varying"); + when(mockResultSetPg.getString(4)).thenReturn("lg_pg"); + + InformationSchemaScanner scannerPg = + new InformationSchemaScanner(mockContextPg, Dialect.POSTGRESQL); + Ddl.Builder builderPg = Ddl.builder(Dialect.POSTGRESQL); + builderPg.createTable("Users").endTable(); + scannerPg.listTableOptions(builderPg); + + Ddl ddlPg = builderPg.build(); + assertEquals(1, ddlPg.table("Users").options().size()); + assertEquals("locality_group='lg_pg'", ddlPg.table("Users").options().get(0)); + } + @Test public void testListColumnOptionsSQL() { assertThat( googleSQLInfoScanner.listColumnOptionsSQL().getSql(), equalToCompressingWhiteSpace( "SELECT t.table_schema, t.table_name, t.column_name, t.option_name, t.option_type," - + " t.option_value" + + " t.option_value, t.is_inherited" + " FROM information_schema.column_options AS t" + " WHERE t.table_schema NOT IN" + " ('INFORMATION_SCHEMA', 'SPANNER_SYS')" @@ -218,7 +328,7 @@ public void testListColumnOptionsSQL() { postgresSQLInfoScanner.listColumnOptionsSQL().getSql(), equalToCompressingWhiteSpace( "SELECT t.table_schema, t.table_name, t.column_name, t.option_name, t.option_type," - + " t.option_value" + + " t.option_value, t.is_inherited" + " FROM information_schema.column_options AS t" + " WHERE t.table_schema NOT IN " + " ('information_schema', 'spanner_sys', 'pg_catalog')" @@ -226,6 +336,126 @@ public void testListColumnOptionsSQL() { + " ORDER BY t.table_name, t.column_name")); } + @Test + public void testListColumnOptions() { + ReadContext mockContext = mock(ReadContext.class); + ResultSet mockTableOptionsResultSet = mock(ResultSet.class); + ResultSet mockColumnOptionsResultSet = mock(ResultSet.class); + when(mockContext.executeQuery(any())) + .thenReturn(mockTableOptionsResultSet, mockColumnOptionsResultSet); + + // Table Users has non-default locality_group "lg1" + when(mockTableOptionsResultSet.next()).thenReturn(true, false); + when(mockTableOptionsResultSet.getString(0)).thenReturn(""); + when(mockTableOptionsResultSet.getString(1)).thenReturn("Users"); + when(mockTableOptionsResultSet.getString(2)).thenReturn("locality_group"); + when(mockTableOptionsResultSet.getString(3)).thenReturn("STRING"); + when(mockTableOptionsResultSet.getString(4)).thenReturn("lg1"); + + // Row 1: col1 inherits locality_group "lg1" from table (skipped) + // Row 2: col2 explicitly sets locality_group "default" (kept) + when(mockColumnOptionsResultSet.next()).thenReturn(true, true, false); + when(mockColumnOptionsResultSet.getString(0)).thenReturn("", ""); + when(mockColumnOptionsResultSet.getString(1)).thenReturn("Users", "Users"); + when(mockColumnOptionsResultSet.getString(2)).thenReturn("col1", "col2"); + when(mockColumnOptionsResultSet.getString(3)).thenReturn("locality_group", "locality_group"); + when(mockColumnOptionsResultSet.getString(4)).thenReturn("STRING", "STRING"); + when(mockColumnOptionsResultSet.getString(5)).thenReturn("lg1", "default"); + when(mockColumnOptionsResultSet.isNull(6)).thenReturn(false, false); + when(mockColumnOptionsResultSet.getBoolean(6)).thenReturn(true, false); + + InformationSchemaScanner scanner = + new InformationSchemaScanner(mockContext, Dialect.GOOGLE_STANDARD_SQL); + Ddl.Builder builder = Ddl.builder(Dialect.GOOGLE_STANDARD_SQL); + builder + .createTable("Users") + .column("col1") + .int64() + .notNull() + .endColumn() + .column("col2") + .int64() + .endColumn() + .primaryKey() + .asc("col1") + .end() + .endTable(); + scanner.listTableOptions(builder); + scanner.listColumnOptions(builder); + + Ddl ddl = builder.build(); + assertEquals(0, ddl.table("Users").column("col1").columnOptions().size()); + assertEquals(1, ddl.table("Users").column("col2").columnOptions().size()); + assertEquals( + "locality_group=\"default\"", ddl.table("Users").column("col2").columnOptions().get(0)); + assertThat( + ddl.prettyPrint(), + equalToCompressingWhiteSpace( + "CREATE TABLE `Users` (" + + " `col1` INT64 NOT NULL," + + " `col2` INT64 OPTIONS (locality_group=\"default\")," + + " ) PRIMARY KEY (`col1` ASC)," + + " OPTIONS (locality_group=\"lg1\")")); + + // PostgreSQL test + ReadContext mockContextPg = mock(ReadContext.class); + ResultSet mockTableOptionsResultSetPg = mock(ResultSet.class); + ResultSet mockColumnOptionsResultSetPg = mock(ResultSet.class); + when(mockContextPg.executeQuery(any())) + .thenReturn(mockTableOptionsResultSetPg, mockColumnOptionsResultSetPg); + + when(mockTableOptionsResultSetPg.next()).thenReturn(true, false); + when(mockTableOptionsResultSetPg.getString(0)).thenReturn("public"); + when(mockTableOptionsResultSetPg.getString(1)).thenReturn("Users"); + when(mockTableOptionsResultSetPg.getString(2)).thenReturn("locality_group"); + when(mockTableOptionsResultSetPg.getString(3)).thenReturn("character varying"); + when(mockTableOptionsResultSetPg.getString(4)).thenReturn("lg_pg"); + + when(mockColumnOptionsResultSetPg.next()).thenReturn(true, true, false); + when(mockColumnOptionsResultSetPg.getString(0)).thenReturn("public", "public"); + when(mockColumnOptionsResultSetPg.getString(1)).thenReturn("Users", "Users"); + when(mockColumnOptionsResultSetPg.getString(2)).thenReturn("col1", "col2"); + when(mockColumnOptionsResultSetPg.getString(3)).thenReturn("locality_group", "locality_group"); + when(mockColumnOptionsResultSetPg.getString(4)) + .thenReturn("character varying", "character varying"); + when(mockColumnOptionsResultSetPg.getString(5)).thenReturn("lg_pg", "default"); + when(mockColumnOptionsResultSetPg.isNull(6)).thenReturn(false, false); + when(mockColumnOptionsResultSetPg.getString(6)).thenReturn("YES", "NO"); + + InformationSchemaScanner scannerPg = + new InformationSchemaScanner(mockContextPg, Dialect.POSTGRESQL); + Ddl.Builder builderPg = Ddl.builder(Dialect.POSTGRESQL); + builderPg + .createTable("Users") + .column("col1") + .pgInt8() + .notNull() + .endColumn() + .column("col2") + .pgInt8() + .endColumn() + .primaryKey() + .asc("col1") + .end() + .endTable(); + scannerPg.listTableOptions(builderPg); + scannerPg.listColumnOptions(builderPg); + + Ddl ddlPg = builderPg.build(); + assertEquals(0, ddlPg.table("Users").column("col1").columnOptions().size()); + assertEquals(1, ddlPg.table("Users").column("col2").columnOptions().size()); + assertEquals( + "locality_group='default'", ddlPg.table("Users").column("col2").columnOptions().get(0)); + assertThat( + ddlPg.prettyPrint(), + equalToCompressingWhiteSpace( + "CREATE TABLE \"Users\" (" + + " \"col1\" bigint NOT NULL," + + " \"col2\" bigint LOCALITY GROUP \"default\"," + + " PRIMARY KEY (\"col1\")" + + " ) LOCALITY GROUP \"lg_pg\"")); + } + @Test public void testListFunctionParametersSQL() { assertThat( @@ -241,4 +471,98 @@ public void testListFunctionParametersSQL() { IllegalArgumentException.class, () -> postgresSQLInfoScanner.listFunctionParametersSQL().getSql()); } + + @Test + public void testListLocalityGroupsSQL() { + assertThat( + googleSQLInfoScanner.listLocalityGroupsSQL().getSql(), + equalToCompressingWhiteSpace( + "SELECT t.locality_group_name, t.option_name, t.option_value, t.is_inherited" + + " FROM information_schema.locality_group_options AS t" + + " ORDER BY t.locality_group_name, t.option_name")); + + assertThat( + postgresSQLInfoScanner.listLocalityGroupsSQL().getSql(), + equalToCompressingWhiteSpace( + "SELECT t.locality_group_name, t.option_name, t.option_value, t.is_inherited" + + " FROM information_schema.locality_group_options AS t" + + " ORDER BY t.locality_group_name, t.option_name")); + } + + @Test + public void testListLocalityGroups() { + ReadContext mockContext = mock(ReadContext.class); + ResultSet mockResultSet = mock(ResultSet.class); + when(mockContext.executeQuery(any())).thenReturn(mockResultSet); + + // Row 1: default storage="ssd" (is_inherited=true -> skipped) + // Row 2: default ssd_to_hdd_spill_timespan=null (is_inherited=null -> skipped) + // Row 3: lg1 storage="ssd" (is_inherited=false -> kept) + // Row 4: lg1 ssd_to_hdd_spill_timespan="30d" (is_inherited=null -> kept) + // Row 5: lg_empty storage="ssd" (is_inherited=true -> lg_empty created with 0 options) + // Row 6: lg_empty ssd_to_hdd_spill_timespan=null (is_inherited=null -> lg_empty created with 0 + // options) + when(mockResultSet.next()).thenReturn(true, true, true, true, true, true, false); + when(mockResultSet.getString(0)) + .thenReturn("default", "default", "lg1", "lg1", "lg_empty", "lg_empty"); + when(mockResultSet.getString(1)) + .thenReturn( + "storage", + "ssd_to_hdd_spill_timespan", + "storage", + "ssd_to_hdd_spill_timespan", + "storage", + "ssd_to_hdd_spill_timespan"); + when(mockResultSet.isNull(2)).thenReturn(false, true, false, false, false, true); + when(mockResultSet.getString(2)).thenReturn("ssd", "ssd", "30d", "ssd"); + when(mockResultSet.isNull(3)).thenReturn(false, true, false, true, false, true); + when(mockResultSet.getBoolean(3)).thenReturn(true, false, true); + + InformationSchemaScanner scanner = + new InformationSchemaScanner(mockContext, Dialect.GOOGLE_STANDARD_SQL); + + Ddl.Builder builder = Ddl.builder(Dialect.GOOGLE_STANDARD_SQL); + scanner.listLocalityGroups(builder); + + Ddl ddl = builder.build(); + assertNull(ddl.localityGroup("default")); + assertNotNull(ddl.localityGroup("lg1")); + assertEquals(2, ddl.localityGroup("lg1").options().size()); + assertEquals("storage=\"ssd\"", ddl.localityGroup("lg1").options().get(0)); + assertEquals("ssd_to_hdd_spill_timespan=\"30d\"", ddl.localityGroup("lg1").options().get(1)); + assertNotNull(ddl.localityGroup("lg_empty")); + assertEquals(0, ddl.localityGroup("lg_empty").options().size()); + + // PostgreSQL test with altered default locality group + ReadContext mockContextPg = mock(ReadContext.class); + ResultSet mockResultSetPg = mock(ResultSet.class); + when(mockContextPg.executeQuery(any())).thenReturn(mockResultSetPg); + + // Row 1: default storage='hdd' (is_inherited='NO' -> kept!) + // Row 2: default ssd_to_hdd_spill_timespan=null (is_inherited=null -> skipped) + // Row 3: lg_pg storage='ssd' (is_inherited='NO' -> kept) + // Row 4: lg_pg ssd_to_hdd_spill_timespan='10d' (is_inherited=null -> kept) + when(mockResultSetPg.next()).thenReturn(true, true, true, true, false); + when(mockResultSetPg.getString(0)).thenReturn("default", "default", "lg_pg", "lg_pg"); + when(mockResultSetPg.getString(1)) + .thenReturn("storage", "ssd_to_hdd_spill_timespan", "storage", "ssd_to_hdd_spill_timespan"); + when(mockResultSetPg.isNull(2)).thenReturn(false, true, false, false); + when(mockResultSetPg.getString(2)).thenReturn("hdd", "ssd", "10d"); + when(mockResultSetPg.isNull(3)).thenReturn(false, true, false, true); + when(mockResultSetPg.getString(3)).thenReturn("NO", "NO"); + + InformationSchemaScanner scannerPg = + new InformationSchemaScanner(mockContextPg, Dialect.POSTGRESQL); + Ddl.Builder builderPg = Ddl.builder(Dialect.POSTGRESQL); + scannerPg.listLocalityGroups(builderPg); + + Ddl ddlPg = builderPg.build(); + assertNotNull(ddlPg.localityGroup("default")); + assertEquals(1, ddlPg.localityGroup("default").options().size()); + assertEquals("storage='hdd'", ddlPg.localityGroup("default").options().get(0)); + assertNotNull(ddlPg.localityGroup("lg_pg")); + assertEquals(2, ddlPg.localityGroup("lg_pg").options().size()); + assertEquals("storage='ssd'", ddlPg.localityGroup("lg_pg").options().get(0)); + assertEquals("ssd_to_hdd_spill_timespan='10d'", ddlPg.localityGroup("lg_pg").options().get(1)); + } } diff --git a/v1/src/test/java/com/google/cloud/teleport/spanner/ddl/LocalityGroupTest.java b/v1/src/test/java/com/google/cloud/teleport/spanner/ddl/LocalityGroupTest.java new file mode 100644 index 00000000000..51f81e4ec4d --- /dev/null +++ b/v1/src/test/java/com/google/cloud/teleport/spanner/ddl/LocalityGroupTest.java @@ -0,0 +1,117 @@ +/* + * Copyright (C) 2026 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); you may not + * use this file except in compliance with the License. You may obtain a copy of + * the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT + * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the + * License for the specific language governing permissions and limitations under + * the License. + */ +package com.google.cloud.teleport.spanner.ddl; + +import static org.hamcrest.text.IsEqualCompressingWhiteSpace.equalToCompressingWhiteSpace; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertNotEquals; +import static org.junit.Assert.assertThat; + +import com.google.cloud.spanner.Dialect; +import com.google.common.collect.ImmutableList; +import org.junit.Test; + +/** Unit tests for {@link LocalityGroup} class. */ +public class LocalityGroupTest { + + @Test + public void testBasicLocalityGroup() { + LocalityGroup lg = LocalityGroup.builder().name("lg1").build(); + + assertThat(lg.prettyPrint(), equalToCompressingWhiteSpace("CREATE LOCALITY GROUP `lg1`")); + assertThat( + lg.toBuilder().build().toString(), + equalToCompressingWhiteSpace("CREATE LOCALITY GROUP `lg1`")); + } + + @Test + public void testBasicLocalityGroupPg() { + LocalityGroup lg = LocalityGroup.builder(Dialect.POSTGRESQL).name("lg1").build(); + + assertThat(lg.prettyPrint(), equalToCompressingWhiteSpace("CREATE LOCALITY GROUP \"lg1\"")); + assertEquals(Dialect.POSTGRESQL, lg.dialect()); + } + + @Test + public void testLocalityGroupWithOptions() { + LocalityGroup lg = + LocalityGroup.builder() + .name("lg1") + .options(ImmutableList.of("storage = 'ssd'", "ssd_to_hdd_spill_timespan = '30d'")) + .build(); + + assertThat( + lg.prettyPrint(), + equalToCompressingWhiteSpace( + "CREATE LOCALITY GROUP `lg1`\n\tOPTIONS (storage = 'ssd', ssd_to_hdd_spill_timespan = '30d')")); + } + + @Test + public void testLocalityGroupWithOptionsPg() { + LocalityGroup lg = + LocalityGroup.builder(Dialect.POSTGRESQL) + .name("lg1") + .options(ImmutableList.of("ssd_to_hdd_spill_timespan='30d'", "storage='ssd'")) + .build(); + + assertThat( + lg.prettyPrint(), + equalToCompressingWhiteSpace( + "CREATE LOCALITY GROUP \"lg1\" STORAGE 'ssd' SSD_TO_HDD_SPILL_TIMESPAN '30d'")); + } + + @Test + public void testDefaultLocalityGroupWithOptions() { + LocalityGroup lg = + LocalityGroup.builder() + .name("default") + .options(ImmutableList.of("storage=\"hdd\"")) + .build(); + + assertThat( + lg.prettyPrint(), + equalToCompressingWhiteSpace( + "ALTER LOCALITY GROUP `default`\n\tSET OPTIONS (storage=\"hdd\")")); + } + + @Test + public void testDefaultLocalityGroupWithOptionsPg() { + LocalityGroup lg = + LocalityGroup.builder(Dialect.POSTGRESQL) + .name("default") + .options(ImmutableList.of("ssd_to_hdd_spill_timespan='10m'", "storage='ssd'")) + .build(); + + assertThat( + lg.prettyPrint(), + equalToCompressingWhiteSpace( + "ALTER LOCALITY GROUP \"default\" STORAGE 'ssd' SSD_TO_HDD_SPILL_TIMESPAN '10m'")); + } + + @Test + public void testEqualsAndHashCode() { + LocalityGroup lg1 = + LocalityGroup.builder().name("lg1").options(ImmutableList.of("storage = 'ssd'")).build(); + LocalityGroup lg2 = + LocalityGroup.builder().name("lg1").options(ImmutableList.of("storage = 'ssd'")).build(); + LocalityGroup lg3 = + LocalityGroup.builder().name("lg2").options(ImmutableList.of("storage = 'ssd'")).build(); + + assertEquals(lg1, lg2); + assertEquals(lg1.hashCode(), lg2.hashCode()); + assertNotEquals(lg1, lg3); + } +}