From 69cf369218d0ef885f64563618517adc490509aa Mon Sep 17 00:00:00 2001 From: Finn Zhang Date: Thu, 3 Sep 2026 23:20:55 +0000 Subject: [PATCH 1/6] feat(spanner): support for locality_group_options --- .../spanner/AvroSchemaToDdlConverter.java | 25 +++++ .../cloud/teleport/spanner/AvroUtil.java | 1 + .../spanner/DdlToAvroSchemaConverter.java | 18 ++++ .../teleport/spanner/ExportTransform.java | 37 +++++++ .../teleport/spanner/ImportTransform.java | 29 +++++ .../cloud/teleport/spanner/ddl/Ddl.java | 50 +++++++++ .../spanner/ddl/InformationSchemaScanner.java | 100 +++++++++++++++++ .../teleport/spanner/ddl/LocalityGroup.java | 101 ++++++++++++++++++ v1/src/main/proto/export.proto | 2 + .../spanner/AvroSchemaToDdlConverterTest.java | 55 ++++++++++ .../spanner/DdlToAvroSchemaConverterTest.java | 42 ++++++++ .../teleport/spanner/ExportTransformTest.java | 12 +++ .../cloud/teleport/spanner/ddl/DdlTest.java | 81 ++++++++++++++ .../ddl/InformationSchemaScannerIT.java | 60 ++++++++++- .../ddl/InformationSchemaScannerTest.java | 67 ++++++++++++ .../spanner/ddl/LocalityGroupTest.java | 99 +++++++++++++++++ 16 files changed, 778 insertions(+), 1 deletion(-) create mode 100644 v1/src/main/java/com/google/cloud/teleport/spanner/ddl/LocalityGroup.java create mode 100644 v1/src/test/java/com/google/cloud/teleport/spanner/ddl/LocalityGroupTest.java 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..9e18bebbad3 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,27 @@ 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); + + ImmutableList.Builder localityGroupOptions = ImmutableList.builder(); + for (int i = 0; ; i++) { + String spannerOption = schema.getProp(SPANNER_OPTION + i); + if (spannerOption == null) { + break; + } + localityGroupOptions.add(spannerOption); + } + builder.options(localityGroupOptions.build()); + + return builder.build(); + } + public Table toTable(String tableName, Schema schema) { if (tableName == null) { tableName = getSpannerObjectName(schema); 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..fa292e9a48c 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; @@ -426,6 +428,22 @@ 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/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/InformationSchemaScanner.java b/v1/src/main/java/com/google/cloud/teleport/spanner/ddl/InformationSchemaScanner.java index e9f6a70bb6d..83de2ac2dbe 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 @@ -116,6 +116,9 @@ public Ddl scan() { if (placementsSupported()) { listPlacements(builder); } + if (isLocalityGroupSupported()) { + listLocalityGroups(builder); + } if (isPropertyGraphSupported()) { listPropertyGraphs(builder); listPropertyGraphPropertyDeclarations(builder); @@ -2112,6 +2115,103 @@ 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_type, t.option_value" + + " 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) { + ResultSet resultSet = context.executeQuery(listLocalityGroupsSQL()); + + Map> localityGroupNameToOptions = Maps.newHashMap(); + while (resultSet.next()) { + String name = resultSet.getString(0); + if ("default".equalsIgnoreCase(name)) { + // Skip `default` locality group as this is not created by user DDL. + continue; + } + String optionName = resultSet.getString(1); + String optionType = resultSet.getString(2); + String optionValue = resultSet.getString(3); + LOG.info( + "locality group option name = " + + optionName + + ", optionType = " + + optionType + + ", optionValue = " + + optionValue); + + ImmutableList.Builder options = + localityGroupNameToOptions.computeIfAbsent(name, k -> ImmutableList.builder()); + + if (optionType.equalsIgnoreCase("STRING") || optionType.equalsIgnoreCase("STRING(MAX)")) { + options.add( + optionName + + "=" + + GSQL_LITERAL_QUOTE + + OPTION_STRING_ESCAPER.escape(optionValue) + + GSQL_LITERAL_QUOTE); + } else if (optionType.equalsIgnoreCase("character varying") + || optionType.equalsIgnoreCase("text")) { + options.add( + optionName + + "=" + + POSTGRESQL_LITERAL_QUOTE + + OPTION_STRING_ESCAPER.escape(optionValue) + + POSTGRESQL_LITERAL_QUOTE); + } else { + options.add(optionName + "=" + optionValue); + } + } + + 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..fd291506151 --- /dev/null +++ b/v1/src/main/java/com/google/cloud/teleport/spanner/ddl/LocalityGroup.java @@ -0,0 +1,101 @@ +/* + * Copyright (C) 2024 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())); + } + + appendable.append("CREATE LOCALITY GROUP ").append(quoteIdentifier(name(), dialect())); + + if (options() != null && !options().isEmpty()) { + String optionsString = String.join(", ", options()); + appendable.append("\n\tOPTIONS (").append(optionsString).append(")"); + } + } + + 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/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..05939de0f69 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,61 @@ public void placements() { + " OPTIONS (instance_partition='mr-partition', default_leader='us-east1')")); } + @Test + public void localityGroups() { + String avroString = + "{" + + " \"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'\"" + + "}"; + Collection schemas = new ArrayList<>(); + Schema.Parser parser = new Schema.Parser(); + schemas.add(parser.parse(avroString)); + + AvroSchemaToDdlConverter converter = new AvroSchemaToDdlConverter(); + Ddl ddl = converter.toDdl(schemas); + assertThat(ddl.localityGroups(), hasSize(1)); + assertThat( + ddl.prettyPrint(), + equalToCompressingWhiteSpace( + "\nCREATE LOCALITY GROUP `LocalityGroup1`\n\t" + + " OPTIONS (storage='ssd', ssd_to_hdd_spill_timespan='30d')")); + } + + @Test + public void pgLocalityGroups() { + String avroString = + "{" + + " \"type\" : \"record\"," + + " \"name\" : \"LocalityGroup1\"," + + " \"fields\" : []," + + " \"namespace\" : \"spannertest\"," + + " \"googleStorage\" : \"CloudSpanner\"," + + " \"spannerEntity\":\"LocalityGroup\", " + + " \"googleFormatVersion\" : \"booleans\"," + + " \"spannerOption_0\" : \"storage='ssd'\"" + + "}"; + Collection schemas = new ArrayList<>(); + Schema.Parser parser = new Schema.Parser(); + schemas.add(parser.parse(avroString)); + + AvroSchemaToDdlConverter converter = new AvroSchemaToDdlConverter(Dialect.POSTGRESQL); + Ddl ddl = converter.toDdl(schemas); + assertThat(ddl.localityGroups(), hasSize(1)); + assertThat( + ddl.prettyPrint(), + equalToCompressingWhiteSpace( + "\nCREATE LOCALITY GROUP \"LocalityGroup1\"\n\t" + + " OPTIONS (storage='ssd')")); + } + @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..a2d0703177d 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,47 @@ public void placements() { equalTo("default_leader=\"" + defaultLeader + "\"")); } + @Test + public void localityGroups() { + DdlToAvroSchemaConverter converter = + new DdlToAvroSchemaConverter("spannertest", "booleans", false); + String lgName1 = "lg1"; + String lgName2 = "lg2"; + 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() + .build(); + + Collection result = converter.convert(ddl); + assertThat(result, hasSize(2)); + + Iterator it = result.iterator(); + 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'")); + } + @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..ac03f53b77d 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,87 @@ public void pgSequences() { assertNotNull(ddl.hashCode()); } + @Test + public void localityGroups() { + Ddl ddl = + Ddl.builder() + .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( + "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(2, statements.size()); + assertThat( + statements.get(0), + equalToCompressingWhiteSpace("CREATE LOCALITY GROUP `lg1` OPTIONS (storage='ssd')")); + assertThat( + statements.get(1), + equalToCompressingWhiteSpace( + "CREATE LOCALITY GROUP `lg2` OPTIONS (storage='ssd', ssd_to_hdd_spill_timespan='30d')")); + assertNotNull(ddl.hashCode()); + assertEquals(2, ddl.localityGroups().size()); + assertNotNull(ddl.localityGroup("lg1")); + } + + @Test + public void pgLocalityGroups() { + Ddl ddl = + Ddl.builder(Dialect.POSTGRESQL) + .createLocalityGroup("lg1") + .options(ImmutableList.of("storage='ssd'")) + .endLocalityGroup() + .build(); + assertThat( + ddl.prettyPrint(), + equalToCompressingWhiteSpace("CREATE LOCALITY GROUP \"lg1\" OPTIONS (storage='ssd')")); + + List statements = ddl.statements(); + assertEquals(1, statements.size()); + assertThat( + statements.get(0), + equalToCompressingWhiteSpace("CREATE LOCALITY GROUP \"lg1\" OPTIONS (storage='ssd')")); + assertNotNull(ddl.hashCode()); + assertEquals(1, ddl.localityGroups().size()); + assertNotNull(ddl.localityGroup("lg1")); + } + + @Test + public void localityGroupsBeforeTables() { + Ddl ddl = + Ddl.builder() + .createTable("Users") + .column("id") + .int64() + .notNull() + .endColumn() + .primaryKey() + .asc("id") + .end() + .endTable() + .createLocalityGroup("lg1") + .options(ImmutableList.of("storage='ssd'")) + .endLocalityGroup() + .build(); + + List statements = ddl.statements(); + assertEquals(2, statements.size()); + assertThat( + statements.get(0), + equalToCompressingWhiteSpace("CREATE LOCALITY GROUP `lg1` OPTIONS (storage='ssd')")); + assertThat( + statements.get(1), + equalToCompressingWhiteSpace(ddl.table("Users").prettyPrint())); + } + @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..6669e8bb000 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 @@ -1761,6 +1761,62 @@ private SharedTestCase pgSequences() { }); } + private SharedTestCase localityGroups() { + DdlToAvroSchemaConverter converter = + new DdlToAvroSchemaConverter("spannertest", "booleans", true); + List statements = + Arrays.asList( + "CREATE LOCALITY GROUP `lg_it_gsql` OPTIONS (storage = 'ssd')"); + + 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 (storage = 'ssd')")); + + 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("storage=\"ssd\"")); + }); + } + + private SharedTestCase pgLocalityGroups() { + DdlToAvroSchemaConverter converter = + new DdlToAvroSchemaConverter("spannertest", "booleans", true); + List statements = + Arrays.asList( + "CREATE LOCALITY GROUP \"lg_it_pg\" OPTIONS (storage = 'ssd')"); + + 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\" OPTIONS (storage = 'ssd')")); + + 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("storage='ssd'")); + }); + } + private SharedTestCase propertyGraphOnViewGroupBy() { List statements = Arrays.asList( @@ -1963,6 +2019,7 @@ public void sharedInformationSchemaScannerTestGsql() throws Exception { databaseOptions(), changeStreams(), sequences(), + localityGroups(), propertyGraphOnViewGroupBy(), propertyGraphOnViewSimple(), propertyGraphOnViewMixedOrder(), @@ -2002,7 +2059,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..a3bfde9e865 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; @@ -241,4 +243,69 @@ 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_type, t.option_value" + + " 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_type, t.option_value" + + " 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 (skipped after getString(0)), Row 2: lg1 storage, Row 3: lg1 ssd_to_hdd_spill_timespan + when(mockResultSet.next()).thenReturn(true, true, true, false); + when(mockResultSet.getString(0)).thenReturn("default", "lg1", "lg1"); + when(mockResultSet.getString(1)).thenReturn("storage", "ssd_to_hdd_spill_timespan"); + when(mockResultSet.getString(2)).thenReturn("STRING", "STRING"); + when(mockResultSet.getString(3)).thenReturn("ssd", "30d"); + + 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)); + + // 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("lg_pg"); + when(mockResultSetPg.getString(1)).thenReturn("storage"); + when(mockResultSetPg.getString(2)).thenReturn("character varying"); + when(mockResultSetPg.getString(3)).thenReturn("ssd"); + + InformationSchemaScanner scannerPg = + new InformationSchemaScanner(mockContextPg, Dialect.POSTGRESQL); + Ddl.Builder builderPg = Ddl.builder(Dialect.POSTGRESQL); + scannerPg.listLocalityGroups(builderPg); + + Ddl ddlPg = builderPg.build(); + assertNotNull(ddlPg.localityGroup("lg_pg")); + assertEquals(1, ddlPg.localityGroup("lg_pg").options().size()); + assertEquals("storage='ssd'", ddlPg.localityGroup("lg_pg").options().get(0)); + } } 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..169550d2830 --- /dev/null +++ b/v1/src/test/java/com/google/cloud/teleport/spanner/ddl/LocalityGroupTest.java @@ -0,0 +1,99 @@ +/* + * Copyright (C) 2024 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("storage = 'ssd'")) + .build(); + + assertThat( + lg.prettyPrint(), + equalToCompressingWhiteSpace("CREATE LOCALITY GROUP \"lg1\"\n\tOPTIONS (storage = 'ssd')")); + } + + @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); + } +} From 2737d5bd047be5271ec034cabbf3cd51f4d3e4fa Mon Sep 17 00:00:00 2001 From: Finn Zhang Date: Thu, 15 Oct 2026 22:57:08 +0000 Subject: [PATCH 2/6] feat(spanner): complete locality group DDL and Avro import/export support --- .../spanner/AvroSchemaToDdlConverter.java | 2 + .../spanner/DdlToAvroSchemaConverter.java | 5 + .../cloud/teleport/spanner/ddl/Column.java | 26 ++- .../cloud/teleport/spanner/ddl/Index.java | 48 ++++- .../spanner/ddl/InformationSchemaScanner.java | 146 +++++++++---- .../teleport/spanner/ddl/LocalityGroup.java | 39 +++- .../cloud/teleport/spanner/ddl/Table.java | 31 +++ .../spanner/AvroSchemaToDdlConverterTest.java | 110 ++++++++-- .../spanner/DdlToAvroSchemaConverterTest.java | 41 +++- .../cloud/teleport/spanner/ddl/DdlTest.java | 166 +++++++++++++-- .../ddl/InformationSchemaScannerIT.java | 84 +++++++- .../ddl/InformationSchemaScannerTest.java | 199 ++++++++++++++++-- .../spanner/ddl/LocalityGroupTest.java | 52 +++-- 13 files changed, 824 insertions(+), 125 deletions(-) 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 9e18bebbad3..f1038e4b0f5 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 @@ -716,6 +716,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/DdlToAvroSchemaConverter.java b/v1/src/main/java/com/google/cloud/teleport/spanner/DdlToAvroSchemaConverter.java index fa292e9a48c..658ef6b7356 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 @@ -208,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()) { 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/Index.java b/v1/src/main/java/com/google/cloud/teleport/spanner/ddl/Index.java index d003e6696c3..ea06d32a183 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,50 @@ 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 (interleaveIn() != null) { + if (localityGroup != null) { + appendable + .append(" LOCALITY GROUP ") + .append( + "NULL".equalsIgnoreCase(localityGroup) + ? "NULL" + : quoteIdentifier(localityGroup, dialect())); + } + 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())); } } 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 83de2ac2dbe..f0a8f1fc4a3 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); @@ -662,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_") @@ -681,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); @@ -732,7 +737,66 @@ Statement listIndexOptionsSQL() { } } - private void listColumnOptions(Ddl.Builder builder) { + @VisibleForTesting + void listTableOptions(Ddl.Builder builder) { + Statement statement = listTableOptionsSQL(); + + ResultSet resultSet = context.executeQuery(statement); + + Map> allOptions = Maps.newHashMap(); + 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); @@ -744,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); @@ -783,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')" @@ -792,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')" @@ -2151,7 +2224,7 @@ Statement listLocalityGroupsSQL() { case GOOGLE_STANDARD_SQL: case POSTGRESQL: return Statement.of( - "SELECT t.locality_group_name, t.option_name, t.option_type, t.option_value" + "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: @@ -2166,42 +2239,39 @@ void listLocalityGroups(Ddl.Builder builder) { Map> localityGroupNameToOptions = Maps.newHashMap(); while (resultSet.next()) { String name = resultSet.getString(0); - if ("default".equalsIgnoreCase(name)) { - // Skip `default` locality group as this is not created by user DDL. + 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; } - String optionName = resultSet.getString(1); - String optionType = resultSet.getString(2); - String optionValue = resultSet.getString(3); + LOG.info( - "locality group option name = " + "locality group name = " + + name + + ", optionName = " + optionName - + ", optionType = " - + optionType + ", optionValue = " - + optionValue); + + optionValue + + ", isInherited = " + + isInherited); ImmutableList.Builder options = localityGroupNameToOptions.computeIfAbsent(name, k -> ImmutableList.builder()); - if (optionType.equalsIgnoreCase("STRING") || optionType.equalsIgnoreCase("STRING(MAX)")) { - options.add( - optionName - + "=" - + GSQL_LITERAL_QUOTE - + OPTION_STRING_ESCAPER.escape(optionValue) - + GSQL_LITERAL_QUOTE); - } else if (optionType.equalsIgnoreCase("character varying") - || optionType.equalsIgnoreCase("text")) { - options.add( - optionName - + "=" - + POSTGRESQL_LITERAL_QUOTE - + OPTION_STRING_ESCAPER.escape(optionValue) - + POSTGRESQL_LITERAL_QUOTE); - } else { - options.add(optionName + "=" + optionValue); - } + 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 : 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 index fd291506151..bdfcd264ae8 100644 --- 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 @@ -52,11 +52,44 @@ public void prettyPrint(Appendable appendable) throws IOException { throw new IllegalArgumentException(String.format("Unrecognized Dialect: %s.", dialect())); } - appendable.append("CREATE LOCALITY GROUP ").append(quoteIdentifier(name(), dialect())); + boolean isDefault = "default".equalsIgnoreCase(name()); + appendable + .append(isDefault ? "ALTER LOCALITY GROUP " : "CREATE LOCALITY GROUP ") + .append(quoteIdentifier(name(), dialect())); if (options() != null && !options().isEmpty()) { - String optionsString = String.join(", ", options()); - appendable.append("\n\tOPTIONS (").append(optionsString).append(")"); + 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); + } + } } } 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/test/java/com/google/cloud/teleport/spanner/AvroSchemaToDdlConverterTest.java b/v1/src/test/java/com/google/cloud/teleport/spanner/AvroSchemaToDdlConverterTest.java index 05939de0f69..2c2a32d3ba1 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 @@ -1508,7 +1508,18 @@ public void placements() { @Test public void localityGroups() { - String avroString = + 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\"," @@ -1517,26 +1528,68 @@ public void localityGroups() { + " \"googleStorage\" : \"CloudSpanner\"," + " \"spannerEntity\":\"LocalityGroup\", " + " \"googleFormatVersion\" : \"booleans\"," - + " \"spannerOption_0\" : \"storage='ssd'\"," - + " \"spannerOption_1\" : \"ssd_to_hdd_spill_timespan='30d'\"" + + " \"spannerOption_0\" : \"storage=\\\"ssd\\\"\"," + + " \"spannerOption_1\" : \"ssd_to_hdd_spill_timespan=\\\"30d\\\"\"" + + "}"; + String tableAvroString = + "{" + + " \"type\" : \"record\"," + + " \"name\" : \"Users\"," + + " \"namespace\" : \"spannertest\"," + + " \"fields\" : [ {" + + " \"name\" : \"id\"," + + " \"type\" : \"long\"," + + " \"sqlType\" : \"INT64\"" + + " }, {" + + " \"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(avroString)); + schemas.add(parser.parse(defaultLgAvroString)); + schemas.add(parser.parse(lgAvroString)); + schemas.add(parser.parse(tableAvroString)); AvroSchemaToDdlConverter converter = new AvroSchemaToDdlConverter(); Ddl ddl = converter.toDdl(schemas); - assertThat(ddl.localityGroups(), hasSize(1)); + assertThat(ddl.localityGroups(), hasSize(2)); + assertThat(ddl.allTables(), hasSize(1)); assertThat( ddl.prettyPrint(), equalToCompressingWhiteSpace( - "\nCREATE LOCALITY GROUP `LocalityGroup1`\n\t" - + " OPTIONS (storage='ssd', ssd_to_hdd_spill_timespan='30d')")); + "CREATE TABLE `Users` (" + + " `id` INT64 NOT NULL," + + " `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 `LocalityGroup1`" + + " OPTIONS (storage=\"ssd\", ssd_to_hdd_spill_timespan=\"30d\")")); } @Test public void pgLocalityGroups() { - String avroString = + 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\"," @@ -1545,20 +1598,51 @@ public void pgLocalityGroups() { + " \"googleStorage\" : \"CloudSpanner\"," + " \"spannerEntity\":\"LocalityGroup\", " + " \"googleFormatVersion\" : \"booleans\"," - + " \"spannerOption_0\" : \"storage='ssd'\"" + + " \"spannerOption_0\" : \"ssd_to_hdd_spill_timespan='30d'\"," + + " \"spannerOption_1\" : \"storage='ssd'\"" + + "}"; + String tableAvroString = + "{" + + " \"type\" : \"record\"," + + " \"name\" : \"Users\"," + + " \"namespace\" : \"spannertest\"," + + " \"fields\" : [ {" + + " \"name\" : \"id\"," + + " \"type\" : \"long\"," + + " \"sqlType\" : \"bigint\"" + + " }, {" + + " \"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(avroString)); + schemas.add(parser.parse(defaultLgAvroString)); + schemas.add(parser.parse(lgAvroString)); + schemas.add(parser.parse(tableAvroString)); AvroSchemaToDdlConverter converter = new AvroSchemaToDdlConverter(Dialect.POSTGRESQL); Ddl ddl = converter.toDdl(schemas); - assertThat(ddl.localityGroups(), hasSize(1)); + assertThat(ddl.localityGroups(), hasSize(2)); + assertThat(ddl.allTables(), hasSize(1)); assertThat( ddl.prettyPrint(), equalToCompressingWhiteSpace( - "\nCREATE LOCALITY GROUP \"LocalityGroup1\"\n\t" - + " OPTIONS (storage='ssd')")); + "CREATE TABLE \"Users\" (" + + " \"id\" bigint NOT NULL," + + " \"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 \"LocalityGroup1\" STORAGE 'ssd' SSD_TO_HDD_SPILL_TIMESPAN '30d'")); } @Test 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 a2d0703177d..bfe2f96214b 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 @@ -2395,17 +2395,46 @@ public void localityGroups() { Ddl ddl = Ddl.builder() .createLocalityGroup(lgName1) - .options(ImmutableList.of("storage='ssd'")) + .options(ImmutableList.of("storage=\"ssd\"")) .endLocalityGroup() .createLocalityGroup(lgName2) - .options(ImmutableList.of("storage='ssd'", "ssd_to_hdd_spill_timespan='30d'")) + .options(ImmutableList.of("storage=\"ssd\"", "ssd_to_hdd_spill_timespan=\"30d\"")) .endLocalityGroup() + .createTable("Users") + .column("id") + .int64() + .notNull() + .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(2)); + assertThat(result, hasSize(3)); 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("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")); @@ -2413,7 +2442,7 @@ public void localityGroups() { 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'")); + assertThat(avroSchema1.getProp(SPANNER_OPTION + "0"), equalTo("storage=\"ssd\"")); Schema avroSchema2 = it.next(); assertThat(avroSchema2.getNamespace(), equalTo("spannertest")); @@ -2422,9 +2451,9 @@ public void localityGroups() { 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 + "0"), equalTo("storage=\"ssd\"")); assertThat( - avroSchema2.getProp(SPANNER_OPTION + "1"), equalTo("ssd_to_hdd_spill_timespan='30d'")); + avroSchema2.getProp(SPANNER_OPTION + "1"), equalTo("ssd_to_hdd_spill_timespan=\"30d\"")); } @Test 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 ac03f53b77d..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 @@ -1568,30 +1568,38 @@ public void pgSequences() { public void localityGroups() { Ddl ddl = Ddl.builder() + .createLocalityGroup("default") + .options(ImmutableList.of("storage=\"hdd\"")) + .endLocalityGroup() .createLocalityGroup("lg1") - .options(ImmutableList.of("storage='ssd'")) + .options(ImmutableList.of("storage=\"ssd\"")) .endLocalityGroup() .createLocalityGroup("lg2") - .options(ImmutableList.of("storage='ssd'", "ssd_to_hdd_spill_timespan='30d'")) + .options(ImmutableList.of("storage=\"ssd\"", "ssd_to_hdd_spill_timespan=\"30d\"")) .endLocalityGroup() .build(); assertThat( ddl.prettyPrint(), equalToCompressingWhiteSpace( - "CREATE LOCALITY GROUP `lg1` OPTIONS (storage='ssd')" - + " CREATE LOCALITY GROUP `lg2` OPTIONS (storage='ssd', ssd_to_hdd_spill_timespan='30d')")); + "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(2, statements.size()); + assertEquals(3, statements.size()); assertThat( statements.get(0), - equalToCompressingWhiteSpace("CREATE LOCALITY GROUP `lg1` OPTIONS (storage='ssd')")); + 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')")); + "CREATE LOCALITY GROUP `lg2` OPTIONS (storage=\"ssd\", ssd_to_hdd_spill_timespan=\"30d\")")); assertNotNull(ddl.hashCode()); - assertEquals(2, ddl.localityGroups().size()); + assertEquals(3, ddl.localityGroups().size()); assertNotNull(ddl.localityGroup("lg1")); } @@ -1599,26 +1607,47 @@ public void localityGroups() { public void pgLocalityGroups() { Ddl ddl = Ddl.builder(Dialect.POSTGRESQL) + .createLocalityGroup("default") + .options(ImmutableList.of("storage='hdd'")) + .endLocalityGroup() .createLocalityGroup("lg1") - .options(ImmutableList.of("storage='ssd'")) + .options(ImmutableList.of("ssd_to_hdd_spill_timespan='30d'", "storage='ssd'")) .endLocalityGroup() .build(); assertThat( ddl.prettyPrint(), - equalToCompressingWhiteSpace("CREATE LOCALITY GROUP \"lg1\" OPTIONS (storage='ssd')")); + equalToCompressingWhiteSpace( + "ALTER LOCALITY GROUP \"default\" STORAGE 'hdd'" + + " CREATE LOCALITY GROUP \"lg1\" STORAGE 'ssd' SSD_TO_HDD_SPILL_TIMESPAN '30d'")); List statements = ddl.statements(); - assertEquals(1, statements.size()); + assertEquals(2, statements.size()); assertThat( statements.get(0), - equalToCompressingWhiteSpace("CREATE LOCALITY GROUP \"lg1\" OPTIONS (storage='ssd')")); + 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(1, ddl.localityGroups().size()); + 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") @@ -1626,9 +1655,96 @@ public void localityGroupsBeforeTables() { .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'")) @@ -1636,13 +1752,31 @@ public void localityGroupsBeforeTables() { .build(); List statements = ddl.statements(); - assertEquals(2, statements.size()); + assertEquals(5, statements.size()); assertThat( statements.get(0), - equalToCompressingWhiteSpace("CREATE LOCALITY GROUP `lg1` OPTIONS (storage='ssd')")); + equalToCompressingWhiteSpace("CREATE LOCALITY GROUP \"lg1\" STORAGE 'ssd'")); assertThat( statements.get(1), - equalToCompressingWhiteSpace(ddl.table("Users").prettyPrint())); + 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 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 6669e8bb000..c038b07e7b3 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; @@ -1764,9 +1765,19 @@ 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')"); + "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, @@ -1775,7 +1786,23 @@ private SharedTestCase localityGroups() { assertThat( ddl.localityGroup("lg_it_gsql").prettyPrint(), equalToCompressingWhiteSpace( - "CREATE LOCALITY GROUP `lg_it_gsql` OPTIONS (storage = 'ssd')")); + "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 = @@ -1785,16 +1812,36 @@ private SharedTestCase localityGroups() { assertThat(avroSchema, notNullValue()); assertThat(avroSchema.getName(), equalTo("lg_it_gsql")); assertThat(avroSchema.getProp("spannerEntity"), equalTo("LocalityGroup")); - assertThat(avroSchema.getProp("spannerOption_0"), equalTo("storage=\"ssd\"")); + assertThat( + avroSchema.getProp("spannerOption_0"), equalTo("ssd_to_hdd_spill_timespan=\"10m\"")); + assertThat(avroSchema.getProp("spannerOption_1"), equalTo("storage=\"ssd\"")); + + Ddl roundTripDdl = toDdlConverter.toDdl(result); + assertThat( + roundTripDdl.localityGroup("lg_it_gsql").prettyPrint(), + equalToCompressingWhiteSpace(ddl.localityGroup("lg_it_gsql").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\" OPTIONS (storage = 'ssd')"); + "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, @@ -1803,7 +1850,22 @@ private SharedTestCase pgLocalityGroups() { assertThat( ddl.localityGroup("lg_it_pg").prettyPrint(), equalToCompressingWhiteSpace( - "CREATE LOCALITY GROUP \"lg_it_pg\" OPTIONS (storage = 'ssd')")); + "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 = @@ -1813,7 +1875,17 @@ private SharedTestCase pgLocalityGroups() { assertThat(avroSchema, notNullValue()); assertThat(avroSchema.getName(), equalTo("lg_it_pg")); assertThat(avroSchema.getProp("spannerEntity"), equalTo("LocalityGroup")); - assertThat(avroSchema.getProp("spannerOption_0"), equalTo("storage='ssd'")); + assertThat( + avroSchema.getProp("spannerOption_0"), equalTo("ssd_to_hdd_spill_timespan='10m'")); + assertThat(avroSchema.getProp("spannerOption_1"), equalTo("storage='ssd'")); + + Ddl roundTripDdl = toDdlConverter.toDdl(result); + assertThat( + roundTripDdl.localityGroup("lg_it_pg").prettyPrint(), + equalToCompressingWhiteSpace(ddl.localityGroup("lg_it_pg").prettyPrint())); + assertThat( + roundTripDdl.table("t_pgLocalityGroups_Users").prettyPrint(), + equalToCompressingWhiteSpace(ddl.table("t_pgLocalityGroups_Users").prettyPrint())); }); } 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 a3bfde9e865..768708391e9 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 @@ -204,13 +204,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')" @@ -220,7 +287,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')" @@ -228,6 +295,79 @@ public void testListColumnOptionsSQL() { + " ORDER BY t.table_name, t.column_name")); } + @Test + public void testListColumnOptions() { + ReadContext mockContext = mock(ReadContext.class); + ResultSet mockResultSet = mock(ResultSet.class); + when(mockContext.executeQuery(any())).thenReturn(mockResultSet); + + // Row 1: inherited locality_group (skipped), Row 2: explicit locality_group (kept) + when(mockResultSet.next()).thenReturn(true, true, false); + when(mockResultSet.getString(0)).thenReturn("", ""); + when(mockResultSet.getString(1)).thenReturn("Users", "Users"); + when(mockResultSet.getString(2)).thenReturn("col1", "col2"); + when(mockResultSet.getString(3)).thenReturn("locality_group", "locality_group"); + when(mockResultSet.getString(4)).thenReturn("STRING", "STRING"); + when(mockResultSet.getString(5)).thenReturn("lg1", "default"); + when(mockResultSet.isNull(6)).thenReturn(false, false); + when(mockResultSet.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() + .endColumn() + .column("col2") + .int64() + .endColumn() + .endTable(); + 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)); + + // PostgreSQL test + ReadContext mockContextPg = mock(ReadContext.class); + ResultSet mockResultSetPg = mock(ResultSet.class); + when(mockContextPg.executeQuery(any())).thenReturn(mockResultSetPg); + + when(mockResultSetPg.next()).thenReturn(true, true, false); + when(mockResultSetPg.getString(0)).thenReturn("public", "public"); + when(mockResultSetPg.getString(1)).thenReturn("Users", "Users"); + when(mockResultSetPg.getString(2)).thenReturn("col1", "col2"); + when(mockResultSetPg.getString(3)).thenReturn("locality_group", "locality_group"); + when(mockResultSetPg.getString(4)).thenReturn("character varying", "character varying"); + when(mockResultSetPg.getString(5)).thenReturn("lg_pg", "default"); + when(mockResultSetPg.isNull(6)).thenReturn(false, false); + when(mockResultSetPg.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() + .endColumn() + .column("col2") + .pgInt8() + .endColumn() + .endTable(); + 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)); + } + @Test public void testListFunctionParametersSQL() { assertThat( @@ -249,14 +389,14 @@ public void testListLocalityGroupsSQL() { assertThat( googleSQLInfoScanner.listLocalityGroupsSQL().getSql(), equalToCompressingWhiteSpace( - "SELECT t.locality_group_name, t.option_name, t.option_type, t.option_value" + "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_type, t.option_value" + "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")); } @@ -267,12 +407,28 @@ public void testListLocalityGroups() { ResultSet mockResultSet = mock(ResultSet.class); when(mockContext.executeQuery(any())).thenReturn(mockResultSet); - // Row 1: default (skipped after getString(0)), Row 2: lg1 storage, Row 3: lg1 ssd_to_hdd_spill_timespan - when(mockResultSet.next()).thenReturn(true, true, true, false); - when(mockResultSet.getString(0)).thenReturn("default", "lg1", "lg1"); - when(mockResultSet.getString(1)).thenReturn("storage", "ssd_to_hdd_spill_timespan"); - when(mockResultSet.getString(2)).thenReturn("STRING", "STRING"); - when(mockResultSet.getString(3)).thenReturn("ssd", "30d"); + // 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=false -> 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, false, false, true); + when(mockResultSet.getBoolean(3)).thenReturn(true, false, false, true); InformationSchemaScanner scanner = new InformationSchemaScanner(mockContext, Dialect.GOOGLE_STANDARD_SQL); @@ -286,17 +442,25 @@ public void testListLocalityGroups() { 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 + // PostgreSQL test with altered default locality group 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("lg_pg"); - when(mockResultSetPg.getString(1)).thenReturn("storage"); - when(mockResultSetPg.getString(2)).thenReturn("character varying"); - when(mockResultSetPg.getString(3)).thenReturn("ssd"); + // Row 1: default storage='hdd' (is_inherited='NO' -> kept!) + // Row 2: default ssd_to_hdd_spill_timespan=null + // Row 3: lg_pg storage='ssd' (is_inherited='NO' -> kept) + when(mockResultSetPg.next()).thenReturn(true, true, true, false); + when(mockResultSetPg.getString(0)).thenReturn("default", "default", "lg_pg"); + when(mockResultSetPg.getString(1)) + .thenReturn("storage", "ssd_to_hdd_spill_timespan", "storage"); + when(mockResultSetPg.isNull(2)).thenReturn(false, true, false); + when(mockResultSetPg.getString(2)).thenReturn("hdd", "ssd"); + when(mockResultSetPg.isNull(3)).thenReturn(false, true, false); + when(mockResultSetPg.getString(3)).thenReturn("NO", "NO"); InformationSchemaScanner scannerPg = new InformationSchemaScanner(mockContextPg, Dialect.POSTGRESQL); @@ -304,6 +468,9 @@ public void testListLocalityGroups() { 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(1, ddlPg.localityGroup("lg_pg").options().size()); assertEquals("storage='ssd'", ddlPg.localityGroup("lg_pg").options().get(0)); 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 index 169550d2830..36152a37b0a 100644 --- 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 @@ -50,9 +50,7 @@ public void testLocalityGroupWithOptions() { LocalityGroup lg = LocalityGroup.builder() .name("lg1") - .options( - ImmutableList.of( - "storage = 'ssd'", "ssd_to_hdd_spill_timespan = '30d'")) + .options(ImmutableList.of("storage = 'ssd'", "ssd_to_hdd_spill_timespan = '30d'")) .build(); assertThat( @@ -66,31 +64,51 @@ public void testLocalityGroupWithOptionsPg() { LocalityGroup lg = LocalityGroup.builder(Dialect.POSTGRESQL) .name("lg1") - .options(ImmutableList.of("storage = 'ssd'")) + .options(ImmutableList.of("ssd_to_hdd_spill_timespan='30d'", "storage='ssd'")) .build(); assertThat( lg.prettyPrint(), - equalToCompressingWhiteSpace("CREATE LOCALITY GROUP \"lg1\"\n\tOPTIONS (storage = 'ssd')")); + equalToCompressingWhiteSpace( + "CREATE LOCALITY GROUP \"lg1\" STORAGE 'ssd' SSD_TO_HDD_SPILL_TIMESPAN '30d'")); } @Test - public void testEqualsAndHashCode() { - LocalityGroup lg1 = + public void testDefaultLocalityGroupWithOptions() { + LocalityGroup lg = LocalityGroup.builder() - .name("lg1") - .options(ImmutableList.of("storage = 'ssd'")) + .name("default") + .options(ImmutableList.of("storage=\"hdd\"")) .build(); - LocalityGroup lg2 = - LocalityGroup.builder() - .name("lg1") - .options(ImmutableList.of("storage = 'ssd'")) + + 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(); + LocalityGroup.builder().name("lg2").options(ImmutableList.of("storage = 'ssd'")).build(); assertEquals(lg1, lg2); assertEquals(lg1.hashCode(), lg2.hashCode()); From d0b54f3d96005415d05b481b29dae27461837c64 Mon Sep 17 00:00:00 2001 From: Finn Zhang Date: Fri, 16 Oct 2026 20:56:17 +0000 Subject: [PATCH 3/6] test(spanner): mock null is_inherited for spill timespan and guard empty default locality group --- .../spanner/DdlToAvroSchemaConverter.java | 3 ++- .../cloud/teleport/spanner/ddl/Index.java | 15 +++++------- .../teleport/spanner/ddl/LocalityGroup.java | 5 +++- .../ddl/InformationSchemaScannerTest.java | 24 ++++++++++--------- .../spanner/ddl/LocalityGroupTest.java | 2 +- 5 files changed, 26 insertions(+), 23 deletions(-) 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 658ef6b7356..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 @@ -437,7 +437,8 @@ public Collection convert(Ddl ddl) { for (LocalityGroup localityGroup : ddl.localityGroups()) { LOG.info("DdlToAvro localityGroup {}", localityGroup.name()); SchemaBuilder.RecordBuilder recordBuilder = - SchemaBuilder.record(generateAvroSchemaName(localityGroup.name())).namespace(this.namespace); + 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"); 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 ea06d32a183..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 @@ -159,16 +159,10 @@ private void prettyPrintPg(Appendable appendable) throws IOException { } ImmutableList otherOptions = otherOptionsBuilder.build(); - if (interleaveIn() != null) { - if (localityGroup != null) { - appendable - .append(" LOCALITY GROUP ") - .append( - "NULL".equalsIgnoreCase(localityGroup) - ? "NULL" - : quoteIdentifier(localityGroup, dialect())); + if ("SEARCH".equals(type())) { + if (interleaveIn() != null) { + appendable.append(" INTERLEAVE IN ").append(quoteIdentifier(interleaveIn(), dialect())); } - appendable.append(" INTERLEAVE IN ").append(quoteIdentifier(interleaveIn(), dialect())); if (!otherOptions.isEmpty()) { appendable.append(" WITH (").append(String.join(",", otherOptions)).append(")"); } @@ -184,6 +178,9 @@ private void prettyPrintPg(Appendable appendable) throws IOException { ? "NULL" : quoteIdentifier(localityGroup, dialect())); } + if (interleaveIn() != null) { + appendable.append(" INTERLEAVE IN ").append(quoteIdentifier(interleaveIn(), dialect())); + } } if (filter() != null && !filter().isEmpty()) { 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 index bdfcd264ae8..e421cdb5c29 100644 --- 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 @@ -1,5 +1,5 @@ /* - * Copyright (C) 2024 Google LLC + * 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 @@ -53,6 +53,9 @@ public void prettyPrint(Appendable appendable) throws IOException { } 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())); 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 768708391e9..a79d67c7bed 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 @@ -410,7 +410,7 @@ public void testListLocalityGroups() { // 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=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) @@ -427,8 +427,8 @@ public void testListLocalityGroups() { "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, false, false, true); - when(mockResultSet.getBoolean(3)).thenReturn(true, false, false, true); + 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); @@ -451,15 +451,16 @@ public void testListLocalityGroups() { when(mockContextPg.executeQuery(any())).thenReturn(mockResultSetPg); // Row 1: default storage='hdd' (is_inherited='NO' -> kept!) - // Row 2: default ssd_to_hdd_spill_timespan=null + // Row 2: default ssd_to_hdd_spill_timespan=null (is_inherited=null -> skipped) // Row 3: lg_pg storage='ssd' (is_inherited='NO' -> kept) - when(mockResultSetPg.next()).thenReturn(true, true, true, false); - when(mockResultSetPg.getString(0)).thenReturn("default", "default", "lg_pg"); + // 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"); - when(mockResultSetPg.isNull(2)).thenReturn(false, true, false); - when(mockResultSetPg.getString(2)).thenReturn("hdd", "ssd"); - when(mockResultSetPg.isNull(3)).thenReturn(false, true, false); + .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 = @@ -472,7 +473,8 @@ public void testListLocalityGroups() { assertEquals(1, ddlPg.localityGroup("default").options().size()); assertEquals("storage='hdd'", ddlPg.localityGroup("default").options().get(0)); assertNotNull(ddlPg.localityGroup("lg_pg")); - assertEquals(1, ddlPg.localityGroup("lg_pg").options().size()); + 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 index 36152a37b0a..51f81e4ec4d 100644 --- 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 @@ -1,5 +1,5 @@ /* - * Copyright (C) 2024 Google LLC + * 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 From 43bbbeaeb6d140ac564691faa3fa42cba51d5aac Mon Sep 17 00:00:00 2001 From: Finn Zhang Date: Tue, 20 Oct 2026 18:22:32 +0000 Subject: [PATCH 4/6] test(spanner): verify index and inherited column locality group DDL output in scanner tests --- .../ddl/InformationSchemaScannerTest.java | 146 ++++++++++++++---- 1 file changed, 117 insertions(+), 29 deletions(-) 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 a79d67c7bed..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 @@ -130,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); @@ -147,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); @@ -155,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 @@ -298,19 +339,30 @@ public void testListColumnOptionsSQL() { @Test public void testListColumnOptions() { ReadContext mockContext = mock(ReadContext.class); - ResultSet mockResultSet = mock(ResultSet.class); - when(mockContext.executeQuery(any())).thenReturn(mockResultSet); - - // Row 1: inherited locality_group (skipped), Row 2: explicit locality_group (kept) - when(mockResultSet.next()).thenReturn(true, true, false); - when(mockResultSet.getString(0)).thenReturn("", ""); - when(mockResultSet.getString(1)).thenReturn("Users", "Users"); - when(mockResultSet.getString(2)).thenReturn("col1", "col2"); - when(mockResultSet.getString(3)).thenReturn("locality_group", "locality_group"); - when(mockResultSet.getString(4)).thenReturn("STRING", "STRING"); - when(mockResultSet.getString(5)).thenReturn("lg1", "default"); - when(mockResultSet.isNull(6)).thenReturn(false, false); - when(mockResultSet.getBoolean(6)).thenReturn(true, false); + 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); @@ -319,11 +371,16 @@ public void testListColumnOptions() { .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(); @@ -331,21 +388,39 @@ public void testListColumnOptions() { 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 mockResultSetPg = mock(ResultSet.class); - when(mockContextPg.executeQuery(any())).thenReturn(mockResultSetPg); - - when(mockResultSetPg.next()).thenReturn(true, true, false); - when(mockResultSetPg.getString(0)).thenReturn("public", "public"); - when(mockResultSetPg.getString(1)).thenReturn("Users", "Users"); - when(mockResultSetPg.getString(2)).thenReturn("col1", "col2"); - when(mockResultSetPg.getString(3)).thenReturn("locality_group", "locality_group"); - when(mockResultSetPg.getString(4)).thenReturn("character varying", "character varying"); - when(mockResultSetPg.getString(5)).thenReturn("lg_pg", "default"); - when(mockResultSetPg.isNull(6)).thenReturn(false, false); - when(mockResultSetPg.getString(6)).thenReturn("YES", "NO"); + 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); @@ -354,11 +429,16 @@ public void testListColumnOptions() { .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(); @@ -366,6 +446,14 @@ public void testListColumnOptions() { 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 From 51b7f92df8afe154df10a5f3c50e6fa61c5b0d71 Mon Sep 17 00:00:00 2001 From: Finn Zhang Date: Wed, 21 Oct 2026 19:38:14 +0000 Subject: [PATCH 5/6] test(spanner): add Avro and round-trip assertions for optionless locality groups --- .../spanner/AvroSchemaToDdlConverterTest.java | 38 ++++++++++++++++++- .../spanner/DdlToAvroSchemaConverterTest.java | 19 +++++++++- .../ddl/InformationSchemaScannerIT.java | 18 +++++++++ 3 files changed, 72 insertions(+), 3 deletions(-) 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 2c2a32d3ba1..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 @@ -1531,6 +1531,16 @@ public void localityGroups() { + " \"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\"," @@ -1541,6 +1551,10 @@ public void localityGroups() { + " \"type\" : \"long\"," + " \"sqlType\" : \"INT64\"" + " }, {" + + " \"name\" : \"description\"," + + " \"type\" : [ \"null\", \"string\" ]," + + " \"sqlType\" : \"STRING(MAX)\"" + + " }, {" + " \"name\" : \"name\"," + " \"type\" : [ \"null\", \"string\" ]," + " \"sqlType\" : \"STRING(MAX)\"," @@ -1556,22 +1570,25 @@ public void localityGroups() { 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(2)); + 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\")")); } @@ -1601,6 +1618,16 @@ public void pgLocalityGroups() { + " \"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\"," @@ -1611,6 +1638,10 @@ public void pgLocalityGroups() { + " \"type\" : \"long\"," + " \"sqlType\" : \"bigint\"" + " }, {" + + " \"name\" : \"description\"," + + " \"type\" : [ \"null\", \"string\" ]," + + " \"sqlType\" : \"character varying\"" + + " }, {" + " \"name\" : \"name\"," + " \"type\" : [ \"null\", \"string\" ]," + " \"sqlType\" : \"character varying\"," @@ -1626,22 +1657,25 @@ public void pgLocalityGroups() { 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(2)); + 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'")); } 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 bfe2f96214b..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 @@ -2392,6 +2392,7 @@ public void localityGroups() { new DdlToAvroSchemaConverter("spannertest", "booleans", false); String lgName1 = "lg1"; String lgName2 = "lg2"; + String lgEmpty = "lg_empty"; Ddl ddl = Ddl.builder() .createLocalityGroup(lgName1) @@ -2400,11 +2401,17 @@ public void localityGroups() { .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() @@ -2421,7 +2428,7 @@ public void localityGroups() { .build(); Collection result = converter.convert(ddl); - assertThat(result, hasSize(3)); + assertThat(result, hasSize(4)); Iterator it = result.iterator(); Schema tableSchema = it.next(); @@ -2431,6 +2438,7 @@ public void localityGroups() { 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\"")); @@ -2454,6 +2462,15 @@ public void localityGroups() { 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 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 c038b07e7b3..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 @@ -1816,10 +1816,19 @@ private SharedTestCase localityGroups() { 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())); @@ -1879,10 +1888,19 @@ private SharedTestCase pgLocalityGroups() { 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())); From f464de404a3bf3a62f56c5f9686541b3ace3cead Mon Sep 17 00:00:00 2001 From: Finn Zhang Date: Wed, 21 Oct 2026 22:23:26 +0000 Subject: [PATCH 6/6] refactor(spanner): address PR review feedback on ResultSet and Avro option helper --- .../spanner/AvroSchemaToDdlConverter.java | 11 +- .../spanner/ddl/InformationSchemaScanner.java | 110 +++++++++--------- 2 files changed, 56 insertions(+), 65 deletions(-) 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 f1038e4b0f5..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 @@ -590,16 +590,7 @@ public LocalityGroup toLocalityGroup(String localityGroupName, Schema schema) { LOG.debug("Converting to Ddl localityGroupName {}", localityGroupName); LocalityGroup.Builder builder = LocalityGroup.builder(dialect).name(localityGroupName); - - ImmutableList.Builder localityGroupOptions = ImmutableList.builder(); - for (int i = 0; ; i++) { - String spannerOption = schema.getProp(SPANNER_OPTION + i); - if (spannerOption == null) { - break; - } - localityGroupOptions.add(spannerOption); - } - builder.options(localityGroupOptions.build()); + builder.options(getNumberedPropsWithPrefix(schema, SPANNER_OPTION)); return builder.build(); } 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 f0a8f1fc4a3..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 @@ -741,28 +741,28 @@ Statement listIndexOptionsSQL() { void listTableOptions(Ddl.Builder builder) { Statement statement = listTableOptionsSQL(); - ResultSet resultSet = context.executeQuery(statement); - Map> allOptions = Maps.newHashMap(); - 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); + 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); + } } } @@ -2234,44 +2234,44 @@ Statement listLocalityGroupsSQL() { @VisibleForTesting void listLocalityGroups(Ddl.Builder builder) { - ResultSet resultSet = context.executeQuery(listLocalityGroupsSQL()); - Map> localityGroupNameToOptions = Maps.newHashMap(); - 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; - } + 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()); + } - LOG.info( - "locality group name = " - + name - + ", optionName = " - + optionName - + ", optionValue = " - + optionValue - + ", isInherited = " - + isInherited); + if (optionValue == null || isInherited) { + continue; + } - ImmutableList.Builder options = - localityGroupNameToOptions.computeIfAbsent(name, k -> ImmutableList.builder()); + 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); + 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 :