Parameterising DDL creation timeout in Avro to Spanner Import Dataflo…

…w Template. PiperOrigin-RevId: 416244500
epidemicsound · Dec 14, 2021 · c43aab8 · c43aab8
1 parent 021c0b2
commit c43aab8
Show file tree

Hide file tree

Showing 6 changed files with 27 additions and 10 deletions.
diff --git a/src/main/java/com/google/cloud/teleport/spanner/ImportPipeline.java b/src/main/java/com/google/cloud/teleport/spanner/ImportPipeline.java
@@ -91,6 +91,12 @@ public interface Options extends PipelineOptions {
     void setSpannerProjectId(ValueProvider<String> value);
 
     void setWaitUntilFinish(boolean value);
+
+    @Description("DDL creation timeout.")
+    @Default.Integer(30)
+    ValueProvider<Integer> getDDLCreationTimeoutInMinutes();
+
+    void setDDLCreationTimeoutInMinutes(ValueProvider<Integer> value);
   }
 
   public static void main(String[] args) {
@@ -112,7 +118,8 @@ public static void main(String[] args) {
             options.getInputDir(),
             options.getWaitForIndexes(),
             options.getWaitForForeignKeys(),
-            options.getEarlyIndexCreateFlag()));
+            options.getEarlyIndexCreateFlag(),
+            options.getDDLCreationTimeoutInMinutes()));
 
     PipelineResult result = p.run();
     if (options.getWaitUntilFinish()

diff --git a/src/main/java/com/google/cloud/teleport/spanner/ImportTransform.java b/src/main/java/com/google/cloud/teleport/spanner/ImportTransform.java
@@ -101,18 +101,21 @@ public class ImportTransform extends PTransform<PBegin, PDone> {
   private final ValueProvider<Boolean> waitForIndexes;
   private final ValueProvider<Boolean> waitForForeignKeys;
   private final ValueProvider<Boolean> earlyIndexCreateFlag;
+  private final ValueProvider<Integer> ddlCreationTimeoutInMinutes;
 
   public ImportTransform(
       SpannerConfig spannerConfig,
       ValueProvider<String> importDirectory,
       ValueProvider<Boolean> waitForIndexes,
       ValueProvider<Boolean> waitForForeignKeys,
-      ValueProvider<Boolean> earlyIndexCreateFlag) {
+      ValueProvider<Boolean> earlyIndexCreateFlag,
+      ValueProvider<Integer> ddlCreationTimeoutInMinutes) {
     this.spannerConfig = spannerConfig;
     this.importDirectory = importDirectory;
     this.waitForIndexes = waitForIndexes;
     this.waitForForeignKeys = waitForForeignKeys;
     this.earlyIndexCreateFlag = earlyIndexCreateFlag;
+    this.ddlCreationTimeoutInMinutes = ddlCreationTimeoutInMinutes;
   }
 
   @Override
@@ -164,7 +167,8 @@ public void processElement(ProcessContext c) {
                 avroDdlView,
                 informationSchemaView,
                 manifestView,
-                earlyIndexCreateFlag));
+                earlyIndexCreateFlag,
+                ddlCreationTimeoutInMinutes));
 
     final PCollection<Ddl> ddl = createTableOutput.get(CreateTables.getDdlObjectTag());
     final PCollectionView<List<String>> pendingIndexes =
@@ -342,6 +346,7 @@ private static class CreateTables extends PTransform<PBegin, PCollectionTuple> {
     private final PCollectionView<Ddl> informationSchemaView;
     private final PCollectionView<Export> manifestView;
     private final ValueProvider<Boolean> earlyIndexCreateFlag;
+    private final ValueProvider<Integer> ddlCreationTimeoutInMinutes;
 
     private transient ExposedSpannerAccessor spannerAccessor;
 
@@ -372,12 +377,14 @@ public CreateTables(
         PCollectionView<List<KV<String, String>>> avroSchemasView,
         PCollectionView<Ddl> informationSchemaView,
         PCollectionView<Export> manifestView,
-        ValueProvider<Boolean> earlyIndexCreateFlag) {
+        ValueProvider<Boolean> earlyIndexCreateFlag,
+        ValueProvider<Integer> ddlCreationTimeoutInMinutes) {
       this.spannerConfig = spannerConfig;
       this.avroSchemasView = avroSchemasView;
       this.informationSchemaView = informationSchemaView;
       this.manifestView = manifestView;
       this.earlyIndexCreateFlag = earlyIndexCreateFlag;
+      this.ddlCreationTimeoutInMinutes = ddlCreationTimeoutInMinutes;
     }
 
     @Override
@@ -485,8 +492,7 @@ public void processElement(ProcessContext c) {
                                     ddlStatements,
                                     null);
                             try {
-                              // TODO: Wait till operation is complete.
-                              op.get(30, TimeUnit.MINUTES);
+                              op.get(ddlCreationTimeoutInMinutes.get(), TimeUnit.MINUTES);
                             } catch (InterruptedException
                                 | ExecutionException
                                 | TimeoutException e) {

diff --git a/src/test/java/com/google/cloud/teleport/spanner/CopyDbTest.java b/src/test/java/com/google/cloud/teleport/spanner/CopyDbTest.java
@@ -353,7 +353,8 @@ private void runTest() {
             source,
             ValueProvider.StaticValueProvider.of(true),
             ValueProvider.StaticValueProvider.of(true),
-            ValueProvider.StaticValueProvider.of(true)));
+            ValueProvider.StaticValueProvider.of(true),
+            ValueProvider.StaticValueProvider.of(30)));
     PipelineResult importResult = importPipeline.run();
     importResult.waitUntilFinish();
 

diff --git a/src/test/java/com/google/cloud/teleport/spanner/ExportRelatedTablesCheckTest.java b/src/test/java/com/google/cloud/teleport/spanner/ExportRelatedTablesCheckTest.java
@@ -1153,7 +1153,8 @@ private void exportAndImportDb(
             source,
             ValueProvider.StaticValueProvider.of(true),
             ValueProvider.StaticValueProvider.of(true),
-            ValueProvider.StaticValueProvider.of(true)));
+            ValueProvider.StaticValueProvider.of(true),
+            ValueProvider.StaticValueProvider.of(30)));
     PipelineResult importResult = importPipeline.run();
     importResult.waitUntilFinish();
   }

diff --git a/src/test/java/com/google/cloud/teleport/spanner/ExportTimestampTest.java b/src/test/java/com/google/cloud/teleport/spanner/ExportTimestampTest.java
@@ -299,7 +299,8 @@ private void exportAndImportDbAtTime(
             source,
             ValueProvider.StaticValueProvider.of(true),
             ValueProvider.StaticValueProvider.of(true),
-            ValueProvider.StaticValueProvider.of(true)));
+            ValueProvider.StaticValueProvider.of(true),
+            ValueProvider.StaticValueProvider.of(30)));
     PipelineResult importResult = importPipeline.run();
     importResult.waitUntilFinish();
   }

diff --git a/src/test/java/com/google/cloud/teleport/spanner/ImportFromAvroTest.java b/src/test/java/com/google/cloud/teleport/spanner/ImportFromAvroTest.java
@@ -467,7 +467,8 @@ private void runTest(Schema avroSchema, String spannerSchema, Iterable<GenericRe
             ValueProvider.StaticValueProvider.of(manifestFileLocation),
             ValueProvider.StaticValueProvider.of(true),
             ValueProvider.StaticValueProvider.of(true),
-            ValueProvider.StaticValueProvider.of(true)));
+            ValueProvider.StaticValueProvider.of(true),
+            ValueProvider.StaticValueProvider.of(30)));
     PipelineResult importResult = importPipeline.run();
     importResult.waitUntilFinish();