[beam] branch master updated: Add Spanner Integration tests to verify exception handling

[pa...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

pabloem pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git


The following commit(s) were added to refs/heads/master by this push:
     new dbb8f99e8b1 Add Spanner Integration tests to verify exception handling
     new ac9a662aab5 Merge pull request #21748 from Add Spanner Integration tests to verify exception handling
dbb8f99e8b1 is described below

commit dbb8f99e8b14560efd3d1dfd7483c42b430b9004
Author: Niel Markwick <ni...@google.com>
AuthorDate: Wed Jun 8 17:06:34 2022 +0200

    Add Spanner Integration tests to verify exception handling
    
    Exceptions during read should abort the pipeline.
    Added 2 integration test cases.
    
    issue #21564
---
 .../beam/sdk/io/gcp/spanner/SpannerReadIT.java     | 77 +++++++++++++++++++++-
 1 file changed, 76 insertions(+), 1 deletion(-)

diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/SpannerReadIT.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/SpannerReadIT.java
index 7740bc44558..e1ba6a43c0c 100644
--- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/SpannerReadIT.java
+++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/SpannerReadIT.java
@@ -18,6 +18,7 @@
 package org.apache.beam.sdk.io.gcp.spanner;
 
 import com.google.api.gax.longrunning.OperationFuture;
+import com.google.cloud.spanner.BatchClient;
 import com.google.cloud.spanner.Database;
 import com.google.cloud.spanner.DatabaseAdminClient;
 import com.google.cloud.spanner.DatabaseClient;
@@ -40,8 +41,12 @@ import org.apache.beam.sdk.testing.PAssert;
 import org.apache.beam.sdk.testing.TestPipeline;
 import org.apache.beam.sdk.testing.TestPipelineOptions;
 import org.apache.beam.sdk.transforms.Count;
+import org.apache.beam.sdk.transforms.Create;
 import org.apache.beam.sdk.transforms.MapElements;
+import org.apache.beam.sdk.transforms.ParDo;
 import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.SimpleFunction;
+import org.apache.beam.sdk.transforms.View;
 import org.apache.beam.sdk.values.PCollection;
 import org.apache.beam.sdk.values.PCollectionView;
 import org.apache.beam.sdk.values.TypeDescriptor;
@@ -50,6 +55,7 @@ import org.junit.After;
 import org.junit.Before;
 import org.junit.Rule;
 import org.junit.Test;
+import org.junit.rules.ExpectedException;
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
 
@@ -60,6 +66,7 @@ public class SpannerReadIT {
   private static final int MAX_DB_NAME_LENGTH = 30;
 
   @Rule public final transient TestPipeline p = TestPipeline.create();
+  @Rule public transient ExpectedException thrown = ExpectedException.none();
 
   /** Pipeline options for this test. */
   public interface SpannerTestPipelineOptions extends TestPipelineOptions {
@@ -129,6 +136,7 @@ public class SpannerReadIT {
                     + "  Value         STRING(MAX),"
                     + ") PRIMARY KEY (Key)"));
     op.get();
+    // PG-dialect databases need 2-steps to create: Create DB then update DDL.
     databaseAdminClient
         .createDatabase(
             databaseAdminClient
@@ -175,7 +183,74 @@ public class SpannerReadIT {
                 .withTransaction(tx));
     PAssert.thatSingleton(output.apply("Count rows", Count.<Struct>globally())).isEqualTo(5L);
 
-    p.run();
+    p.run().waitUntilFinish();
+  }
+
+  @Test
+  public void testReadFailsBadTable() throws Exception {
+
+    thrown.expect(new SpannerWriteIT.StackTraceContainsString("SpannerException"));
+    thrown.expect(new SpannerWriteIT.StackTraceContainsString("NOT_FOUND: Table not found"));
+
+    SpannerConfig spannerConfig = createSpannerConfig();
+
+    PCollectionView<Transaction> tx =
+        p.apply(
+            "Create tx",
+            SpannerIO.createTransaction()
+                .withSpannerConfig(spannerConfig)
+                .withTimestampBound(TimestampBound.strong()));
+    p.apply(
+        "read db",
+        SpannerIO.read()
+            .withSpannerConfig(spannerConfig)
+            .withTable("IncorrectTableName")
+            .withColumns("Key", "Value")
+            .withTransaction(tx));
+    p.run().waitUntilFinish();
+  }
+
+  private static class CloseTransactionFn extends SimpleFunction<Transaction, Transaction> {
+    private final SpannerConfig spannerConfig;
+
+    private CloseTransactionFn(SpannerConfig spannerConfig) {
+      this.spannerConfig = spannerConfig;
+    }
+
+    @Override
+    public Transaction apply(Transaction tx) {
+      BatchClient batchClient = SpannerAccessor.getOrCreate(spannerConfig).getBatchClient();
+      batchClient.batchReadOnlyTransaction(tx.transactionId()).cleanup();
+      return tx;
+    }
+  }
+
+  @Test
+  public void testReadFailsBadSession() throws Exception {
+
+    thrown.expect(new SpannerWriteIT.StackTraceContainsString("SpannerException"));
+    thrown.expect(new SpannerWriteIT.StackTraceContainsString("NOT_FOUND: Session not found"));
+
+    SpannerConfig spannerConfig = createSpannerConfig();
+
+    // This creates a transaction then closes the session.
+    // The (closed) transaction is then passed to SpannerIO.read() and should
+    // raise SessionNotFound errors.
+    PCollectionView<Transaction> tx =
+        p.apply("Transaction seed", Create.of(1))
+            .apply(
+                "Create transaction",
+                ParDo.of(new CreateTransactionFn(spannerConfig, TimestampBound.strong())))
+            .apply("Close Transaction", MapElements.via(new CloseTransactionFn(spannerConfig)))
+            .apply("As PCollectionView", View.asSingleton());
+    p.apply(
+        "read db",
+        SpannerIO.read()
+            .withSpannerConfig(spannerConfig)
+            .withTable(options.getTable())
+            .withColumns("Key", "Value")
+            .withTransaction(tx));
+    p.run().waitUntilFinish();
   }
 
   @Test