Skip to content

Remove LocalSpannerIO and uses Beam's SpannerIO for Spanner Templates #1429

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
Abacn wants to merge 2 commits into from
Closed

Remove LocalSpannerIO and uses Beam's SpannerIO for Spanner Templates #1429

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
c74bca7
Remove LocalSpannerIO
Abacn Apr 11, 2024
4362fca
revert tokenlist UUID and tokenlist test senario
Abacn Feb 14, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 3 additions & 3 deletions v1/src/main/java/com/google/cloud/teleport/spanner/ApplyDDLTransform.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,7 @@
import com.google.spanner.admin.database.v1.UpdateDatabaseDdlMetadata;
import java.util.List;
import java.util.concurrent.ExecutionException;
import org.apache.beam.sdk.io.gcp.spanner.LocalSpannerAccessor;
import org.apache.beam.sdk.io.gcp.spanner.SpannerAccessor;
import org.apache.beam.sdk.io.gcp.spanner.SpannerConfig;
import org.apache.beam.sdk.options.ValueProvider;
import org.apache.beam.sdk.transforms.DoFn;
Expand Down Expand Up @@ -66,11 +66,11 @@
ParDo.of(
new DoFn<Ddl, Ddl>() {

private transient LocalSpannerAccessor spannerAccessor;
private transient SpannerAccessor spannerAccessor;

@Setup
public void setup() {
spannerAccessor = LocalSpannerAccessor.getOrCreate(spannerConfig);
spannerAccessor = SpannerAccessor.getOrCreate(spannerConfig);

Check warning on line 73 in v1/src/main/java/com/google/cloud/teleport/spanner/ApplyDDLTransform.java

View check run for this annotation

Codecov / codecov/patch

v1/src/main/java/com/google/cloud/teleport/spanner/ApplyDDLTransform.java#L73 

Added line #L73 was not covered by tests
}

@Teardown
Expand Down
6 changes: 3 additions & 3 deletions v1/src/main/java/com/google/cloud/teleport/spanner/ExportTransform.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -80,9 +80,9 @@
import org.apache.beam.sdk.io.WriteFilesResult;
import org.apache.beam.sdk.io.fs.ResolveOptions;
import org.apache.beam.sdk.io.fs.ResourceId;
import org.apache.beam.sdk.io.gcp.spanner.LocalSpannerIO;
import org.apache.beam.sdk.io.gcp.spanner.ReadOperation;
import org.apache.beam.sdk.io.gcp.spanner.SpannerConfig;
import org.apache.beam.sdk.io.gcp.spanner.SpannerIO;
import org.apache.beam.sdk.io.gcp.spanner.Transaction;
import org.apache.beam.sdk.options.ValueProvider;
import org.apache.beam.sdk.transforms.Combine;
Expand Down Expand Up @@ -185,7 +185,7 @@

/*
* Allow users to specify read timestamp.
* CreateTransaction and CreateTransactionFn classes in LocalSpannerIO
* CreateTransaction and CreateTransactionFn classes in SpannerIO
* only take a timestamp object for exact staleness which works when
* parameters are provided during template compile time. They do not work with
* a Timestamp valueProvider which can take parameters at runtime. Hence a new
Expand Down Expand Up @@ -457,7 +457,7 @@
PCollection<Struct> rows =
tableReadOperations.apply(
"Read all rows from Spanner",
LocalSpannerIO.readAll().withTransaction(tx).withSpannerConfig(spannerConfig));
SpannerIO.readAll().withTransaction(tx).withSpannerConfig(spannerConfig));

Check warning on line 460 in v1/src/main/java/com/google/cloud/teleport/spanner/ExportTransform.java

View check run for this annotation

Codecov / codecov/patch

v1/src/main/java/com/google/cloud/teleport/spanner/ExportTransform.java#L460 

Added line #L460 was not covered by tests

ValueProvider<ResourceId> resource =
ValueProvider.NestedValueProvider.of(
Expand Down
12 changes: 6 additions & 6 deletions v1/src/main/java/com/google/cloud/teleport/spanner/ImportTransform.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -66,9 +66,9 @@
import org.apache.beam.sdk.io.fs.EmptyMatchTreatment;
import org.apache.beam.sdk.io.fs.MatchResult;
import org.apache.beam.sdk.io.fs.ResourceId;
import org.apache.beam.sdk.io.gcp.spanner.LocalSpannerAccessor;
import org.apache.beam.sdk.io.gcp.spanner.LocalSpannerIO;
import org.apache.beam.sdk.io.gcp.spanner.SpannerAccessor;
import org.apache.beam.sdk.io.gcp.spanner.SpannerConfig;
import org.apache.beam.sdk.io.gcp.spanner.SpannerIO;
import org.apache.beam.sdk.io.gcp.spanner.SpannerWriteResult;
import org.apache.beam.sdk.io.gcp.spanner.Transaction;
import org.apache.beam.sdk.options.ValueProvider;
Expand Down Expand Up @@ -173,7 +173,7 @@
schemas.apply("Build avro DDL", Combine.globally(AsList.fn()));

PCollectionView<Transaction> tx =
begin.apply(LocalSpannerIO.createTransaction().withSpannerConfig(spannerConfig));
begin.apply(SpannerIO.createTransaction().withSpannerConfig(spannerConfig));

Check warning on line 176 in v1/src/main/java/com/google/cloud/teleport/spanner/ImportTransform.java

View check run for this annotation

Codecov / codecov/patch

v1/src/main/java/com/google/cloud/teleport/spanner/ImportTransform.java#L176 

Added line #L176 was not covered by tests

PCollection<Ddl> informationSchemaDdl =
begin.apply(
Expand Down Expand Up @@ -273,7 +273,7 @@
SpannerWriteResult result =
mutations.apply(
"Write mutations " + depth,
LocalSpannerIO.write()
SpannerIO.write()

Check warning on line 276 in v1/src/main/java/com/google/cloud/teleport/spanner/ImportTransform.java

View check run for this annotation

Codecov / codecov/patch

v1/src/main/java/com/google/cloud/teleport/spanner/ImportTransform.java#L276 

Added line #L276 was not covered by tests
.withSchemaReadySignal(ddl)
.withSpannerConfig(spannerConfig)
.withCommitDeadline(Duration.standardMinutes(1))
Expand Down Expand Up @@ -408,7 +408,7 @@
private final ValueProvider<Boolean> earlyIndexCreateFlag;
private final ValueProvider<Integer> ddlCreationTimeoutInMinutes;

private transient LocalSpannerAccessor spannerAccessor;
private transient SpannerAccessor spannerAccessor;
private static final Logger LOG = LoggerFactory.getLogger(CreateTables.class);

/* If the schema has a lot of DDL changes after data load, it's preferable to create
Expand Down Expand Up @@ -464,7 +464,7 @@

@Setup
public void setup() {
spannerAccessor = LocalSpannerAccessor.getOrCreate(spannerConfig);
spannerAccessor = SpannerAccessor.getOrCreate(spannerConfig);

Check warning on line 467 in v1/src/main/java/com/google/cloud/teleport/spanner/ImportTransform.java

View check run for this annotation

Codecov / codecov/patch

v1/src/main/java/com/google/cloud/teleport/spanner/ImportTransform.java#L467 

Added line #L467 was not covered by tests
}

@Teardown
Expand Down
6 changes: 3 additions & 3 deletions v1/src/main/java/com/google/cloud/teleport/spanner/ReadDialect.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,7 @@

import com.google.cloud.spanner.DatabaseClient;
import com.google.cloud.spanner.Dialect;
import org.apache.beam.sdk.io.gcp.spanner.LocalSpannerAccessor;
import org.apache.beam.sdk.io.gcp.spanner.SpannerAccessor;
import org.apache.beam.sdk.io.gcp.spanner.SpannerConfig;
import org.apache.beam.sdk.transforms.Create;
import org.apache.beam.sdk.transforms.DoFn;
Expand All @@ -43,15 +43,15 @@

private static class ReadDialectFn extends DoFn<Void, Dialect> {
private final SpannerConfig spannerConfig;
private transient LocalSpannerAccessor spannerAccessor;
private transient SpannerAccessor spannerAccessor;

public ReadDialectFn(SpannerConfig spannerConfig) {
this.spannerConfig = spannerConfig;
}

@Setup
public void setup() throws Exception {
spannerAccessor = LocalSpannerAccessor.getOrCreate(spannerConfig);
spannerAccessor = SpannerAccessor.getOrCreate(spannerConfig);

Check warning on line 54 in v1/src/main/java/com/google/cloud/teleport/spanner/ReadDialect.java

View check run for this annotation

Codecov / codecov/patch

v1/src/main/java/com/google/cloud/teleport/spanner/ReadDialect.java#L54 

Added line #L54 was not covered by tests
}

@Teardown
Expand Down
6 changes: 3 additions & 3 deletions v1/src/main/java/com/google/cloud/teleport/spanner/ReadInformationSchema.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,7 @@
import com.google.cloud.spanner.Dialect;
import com.google.cloud.teleport.spanner.ddl.Ddl;
import com.google.cloud.teleport.spanner.ddl.InformationSchemaScanner;
import org.apache.beam.sdk.io.gcp.spanner.LocalSpannerAccessor;
import org.apache.beam.sdk.io.gcp.spanner.SpannerAccessor;
import org.apache.beam.sdk.io.gcp.spanner.SpannerConfig;
import org.apache.beam.sdk.io.gcp.spanner.Transaction;
import org.apache.beam.sdk.transforms.Create;
Expand Down Expand Up @@ -59,7 +59,7 @@

private static class ReadInformationSchemaFn extends DoFn<Void, Ddl> {
private final SpannerConfig spannerConfig;
private transient LocalSpannerAccessor spannerAccessor;
private transient SpannerAccessor spannerAccessor;
private final PCollectionView<Transaction> tx;
private final PCollectionView<Dialect> dialectView;

Expand All @@ -74,7 +74,7 @@

@Setup
public void setup() throws Exception {
spannerAccessor = LocalSpannerAccessor.getOrCreate(spannerConfig);
spannerAccessor = SpannerAccessor.getOrCreate(spannerConfig);

Check warning on line 77 in v1/src/main/java/com/google/cloud/teleport/spanner/ReadInformationSchema.java

View check run for this annotation

Codecov / codecov/patch

v1/src/main/java/com/google/cloud/teleport/spanner/ReadInformationSchema.java#L77 

Added line #L77 was not covered by tests
}

@Teardown
Expand Down
6 changes: 3 additions & 3 deletions v1/src/main/java/com/google/cloud/teleport/spanner/TextImportTransform.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -53,8 +53,8 @@
import org.apache.beam.sdk.io.fs.EmptyMatchTreatment;
import org.apache.beam.sdk.io.fs.MatchResult;
import org.apache.beam.sdk.io.fs.ResourceId;
import org.apache.beam.sdk.io.gcp.spanner.LocalSpannerIO;
import org.apache.beam.sdk.io.gcp.spanner.SpannerConfig;
import org.apache.beam.sdk.io.gcp.spanner.SpannerIO;
import org.apache.beam.sdk.io.gcp.spanner.SpannerWriteResult;
import org.apache.beam.sdk.io.gcp.spanner.Transaction;
import org.apache.beam.sdk.options.ValueProvider;
Expand Down Expand Up @@ -105,7 +105,7 @@
@Override
public PDone expand(PBegin begin) {
PCollectionView<Transaction> tx =
begin.apply(LocalSpannerIO.createTransaction().withSpannerConfig(spannerConfig));
begin.apply(SpannerIO.createTransaction().withSpannerConfig(spannerConfig));

Check warning on line 108 in v1/src/main/java/com/google/cloud/teleport/spanner/TextImportTransform.java

View check run for this annotation

Codecov / codecov/patch

v1/src/main/java/com/google/cloud/teleport/spanner/TextImportTransform.java#L108 

Added line #L108 was not covered by tests

PCollectionView<Dialect> dialectView =
begin
Expand Down Expand Up @@ -203,7 +203,7 @@
.apply("Wait for previous depth " + depth, Wait.on(previousComputation))
.apply(
"Write mutations " + depth,
LocalSpannerIO.write()
SpannerIO.write()

Check warning on line 206 in v1/src/main/java/com/google/cloud/teleport/spanner/TextImportTransform.java

View check run for this annotation

Codecov / codecov/patch

v1/src/main/java/com/google/cloud/teleport/spanner/TextImportTransform.java#L206 

Added line #L206 was not covered by tests
.withSpannerConfig(spannerConfig)
.withCommitDeadline(Duration.standardMinutes(1))
.withMaxCumulativeBackoff(Duration.standardHours(2))
Expand Down
10 changes: 5 additions & 5 deletions v1/src/main/java/com/google/cloud/teleport/templates/SpannerToText.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -31,9 +31,9 @@
import org.apache.beam.sdk.io.FileSystems;
import org.apache.beam.sdk.io.TextIO;
import org.apache.beam.sdk.io.fs.ResourceId;
import org.apache.beam.sdk.io.gcp.spanner.LocalSpannerIO;
import org.apache.beam.sdk.io.gcp.spanner.ReadOperation;
import org.apache.beam.sdk.io.gcp.spanner.SpannerConfig;
import org.apache.beam.sdk.io.gcp.spanner.SpannerIO;
import org.apache.beam.sdk.io.gcp.spanner.Transaction;
import org.apache.beam.sdk.options.PipelineOptions;
import org.apache.beam.sdk.options.PipelineOptionsFactory;
Expand Down Expand Up @@ -144,7 +144,7 @@
options.getTextWritePrefix(),
options.getSpannerSnapshotTime());

/* CreateTransaction and CreateTransactionFn classes in LocalSpannerIO
/* CreateTransaction and CreateTransactionFn classes in SpannerIO
* only take a timestamp object for exact staleness which works when
* parameters are provided during template compile time. They do not work with
* a Timestamp valueProvider which can take parameters at runtime. Hence a new
Expand All @@ -164,14 +164,14 @@
PCollection<String> csv =
pipeline
.apply("Create export", spannerExport)
// We need to use LocalSpannerIO.readAll() instead of LocalSpannerIO.read()
// We need to use SpannerIO.readAll() instead of SpannerIO.read()
// because ValueProvider parameters such as table name required for
// LocalSpannerIO.read() can be read only inside DoFn but LocalSpannerIO.read() is of
// SpannerIO.read() can be read only inside DoFn but SpannerIO.read() is of
// type PTransform<PBegin, Struct>, which prevents prepending it with DoFn that reads
// these parameters at the pipeline execution time.
.apply(
"Read all records",
LocalSpannerIO.readAll().withTransaction(tx).withSpannerConfig(spannerConfig))
SpannerIO.readAll().withTransaction(tx).withSpannerConfig(spannerConfig))

Check warning on line 174 in v1/src/main/java/com/google/cloud/teleport/templates/SpannerToText.java

View check run for this annotation

Codecov / codecov/patch

v1/src/main/java/com/google/cloud/teleport/templates/SpannerToText.java#L174 

Added line #L174 was not covered by tests
.apply(
"Struct To Csv",
MapElements.into(TypeDescriptors.strings())
Expand Down
10 changes: 5 additions & 5 deletions v1/src/main/java/com/google/cloud/teleport/templates/SpannerVectorEmbeddingExport.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -27,9 +27,9 @@
import org.apache.beam.sdk.Pipeline;
import org.apache.beam.sdk.io.FileSystems;
import org.apache.beam.sdk.io.TextIO;
import org.apache.beam.sdk.io.gcp.spanner.LocalSpannerIO;
import org.apache.beam.sdk.io.gcp.spanner.ReadOperation;
import org.apache.beam.sdk.io.gcp.spanner.SpannerConfig;
import org.apache.beam.sdk.io.gcp.spanner.SpannerIO;
import org.apache.beam.sdk.io.gcp.spanner.Transaction;
import org.apache.beam.sdk.options.Default;
import org.apache.beam.sdk.options.PipelineOptions;
Expand Down Expand Up @@ -264,7 +264,7 @@
options.getSpannerColumnsToExport(),
ValueProvider.StaticValueProvider.of(/* disable_schema_export= */ false));

/* CreateTransaction and CreateTransactionFn classes in LocalSpannerIO
/* CreateTransaction and CreateTransactionFn classes in SpannerIO
* only take a timestamp object for exact staleness which works when
* parameters are provided during template compile time. They do not work with
* a Timestamp valueProvider which can take parameters at runtime. Hence a new
Expand All @@ -284,14 +284,14 @@
PCollection<String> json =
pipeline
.apply("Create export", spannerExport)
// We need to use LocalSpannerIO.readAll() instead of LocalSpannerIO.read()
// We need to use SpannerIO.readAll() instead of SpannerIO.read()
// because ValueProvider parameters such as table name required for
// LocalSpannerIO.read() can be read only inside DoFn but LocalSpannerIO.read() is of
// SpannerIO.read() can be read only inside DoFn but SpannerIO.read() is of
// type PTransform<PBegin, Struct>, which prevents prepending it with DoFn that reads
// these parameters at the pipeline execution time.
.apply(
"Read all records",
LocalSpannerIO.readAll().withTransaction(tx).withSpannerConfig(spannerConfig))
SpannerIO.readAll().withTransaction(tx).withSpannerConfig(spannerConfig))

Check warning on line 294 in v1/src/main/java/com/google/cloud/teleport/templates/SpannerVectorEmbeddingExport.java

View check run for this annotation

Codecov / codecov/patch

v1/src/main/java/com/google/cloud/teleport/templates/SpannerVectorEmbeddingExport.java#L294 

Added line #L294 was not covered by tests
.apply(
"Struct To JSON",
MapElements.into(TypeDescriptors.strings())
Expand Down
14 changes: 7 additions & 7 deletions v1/src/main/java/com/google/cloud/teleport/templates/common/SpannerConverters.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -56,8 +56,8 @@
import java.util.stream.Collectors;
import kotlin.Pair;
import org.apache.beam.sdk.io.FileSystems;
import org.apache.beam.sdk.io.gcp.spanner.LocalSpannerAccessor;
import org.apache.beam.sdk.io.gcp.spanner.ReadOperation;
import org.apache.beam.sdk.io.gcp.spanner.SpannerAccessor;
import org.apache.beam.sdk.io.gcp.spanner.SpannerConfig;
import org.apache.beam.sdk.io.gcp.spanner.Transaction;
import org.apache.beam.sdk.options.Default;
Expand Down Expand Up @@ -261,16 +261,16 @@
return new AutoValue_SpannerConverters_ExportTransform.Builder();
}

private LocalSpannerAccessor spannerAccessor;
private SpannerAccessor spannerAccessor;
private DatabaseClient databaseClient;

// LocalSpannerAccessor is not serializable, thus can't be passed as a mock so we need to pass
// SpannerAccessor is not serializable, thus can't be passed as a mock so we need to pass
// mocked database client directly instead. We can't generate stub of ExportTransform because
// AutoValue generates a final class.
// TODO make LocalSpannerAccessor serializable
// TODO make SpannerAccessor serializable
DatabaseClient getDatabaseClient(SpannerConfig spannerConfig) {
if (databaseClient == null) {
this.spannerAccessor = LocalSpannerAccessor.getOrCreate(spannerConfig);
this.spannerAccessor = SpannerAccessor.getOrCreate(spannerConfig);

Check warning on line 273 in v1/src/main/java/com/google/cloud/teleport/templates/common/SpannerConverters.java

View check run for this annotation

Codecov / codecov/patch

v1/src/main/java/com/google/cloud/teleport/templates/common/SpannerConverters.java#L273 

Added line #L273 was not covered by tests
return this.spannerAccessor.getDatabaseClient();
} else {
return this.databaseClient;
Expand Down Expand Up @@ -817,11 +817,11 @@
this.spannerSnapshotTime = spannerSnapshotTime;
}

private transient LocalSpannerAccessor spannerAccessor;
private transient SpannerAccessor spannerAccessor;

@DoFn.Setup
public void setup() throws Exception {
spannerAccessor = LocalSpannerAccessor.getOrCreate(config);
spannerAccessor = SpannerAccessor.getOrCreate(config);

Check warning on line 824 in v1/src/main/java/com/google/cloud/teleport/templates/common/SpannerConverters.java

View check run for this annotation

Codecov / codecov/patch

v1/src/main/java/com/google/cloud/teleport/templates/common/SpannerConverters.java#L824 

Added line #L824 was not covered by tests
}

@Teardown
Expand Down
Loading
Loading