Updated Spanner Dataflow connector samples

dataflow/spanner-io/src/main/java/com/example/dataflow/EstimateSize.java

@@ -0,0 +1,86 @@
+/*
+ * Copyright 2018 Google Inc.
+ *
+ * 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.example.dataflow;
+
+import com.google.cloud.spanner.Struct;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.values.PCollection;
+
+/**
+ * Estimates the size of the {@code Struct}.
+ */
+public class EstimateSize extends PTransform<PCollection<Struct>, PCollection<Long>> {
+
+  public static EstimateSize create() {
+    return new EstimateSize();
+  }
+
+  private EstimateSize() {
+  }
+
+  @Override
+  public PCollection<Long> expand(PCollection<Struct> input) {
+    return input.apply(ParDo.of(new EstimateStructSizeFn()));
+  }
+
+  /**
+   * Estimates the size of a Spanner row. For simplicity, arrays and structs aren't supported.
+   */
+  public static class EstimateStructSizeFn extends DoFn<Struct, Long> {
+
+    @ProcessElement
+    public void processElement(ProcessContext c) throws Exception {
+      Struct row = c.element();
+      long sum = 0;
+      for (int i = 0; i < row.getColumnCount(); i++) {
+        if (row.isNull(i)) {
+          continue;
+        }
+
+        switch (row.getColumnType(i).getCode()) {
+          case BOOL:
+            sum += 1;
+            break;
+          case INT64:
+          case FLOAT64:
+            sum += 8;
+            break;
+          case TIMESTAMP:
+          case DATE:
+            sum += 12;
+            break;
+          case BYTES:
+            sum += row.getBytes(i).length();
+            break;
+          case STRING:
+            sum += row.getString(i).length();
+            break;
+          case ARRAY:
+            throw new IllegalArgumentException("Arrays are not supported :(");
+          case STRUCT:
+            throw new IllegalArgumentException("Structs are not supported :(");
+          default:
+            throw new IllegalArgumentException("Unsupported type :(");
+        }
+      }
+      c.output(sum);
+    }
+  }
+
+}

dataflow/spanner-io/src/main/java/com/example/dataflow/SpannerGroupWrite.java

@@ -0,0 +1,114 @@
+/*
+ * Copyright 2018 Google Inc.
+ *
+ * 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.example.dataflow;
+
+import com.google.cloud.Timestamp;
+import com.google.cloud.spanner.Mutation;
+import com.google.common.base.Charsets;
+import com.google.common.hash.Hashing;
+import org.apache.beam.sdk.Pipeline;
+import org.apache.beam.sdk.io.TextIO;
+import org.apache.beam.sdk.io.gcp.spanner.MutationGroup;
+import org.apache.beam.sdk.io.gcp.spanner.SpannerIO;
+import org.apache.beam.sdk.options.Default;
+import org.apache.beam.sdk.options.Description;
+import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.options.PipelineOptionsFactory;
+import org.apache.beam.sdk.options.Validation;
+import org.apache.beam.sdk.transforms.MapElements;
+import org.apache.beam.sdk.transforms.SimpleFunction;
+import org.apache.beam.sdk.values.PCollection;
+
+/**
+ * This sample demonstrates how to group together mutations when writing to the Cloud Spanner
+ * database.
+ */
+public class SpannerGroupWrite {
+  public interface Options extends PipelineOptions {
+
+    @Description("Spanner instance ID to write to")
+    @Validation.Required
+    String getInstanceId();
+
+    void setInstanceId(String value);
+
+    @Description("Spanner database name to write to")
+    @Validation.Required
+    String getDatabaseId();
+
+    void setDatabaseId(String value);
+
+    @Description("Singers output filename in the format: singer_id\tfirst_name\tlast_name")
+    @Default.String("data/usersids.txt")
+    String getSuspiciousUsersFile();
+
+    void setSuspiciousUserFile(String value);
+
+  }
+
+  public static void main(String[] args) {
+    Options options = PipelineOptionsFactory.fromArgs(args).withValidation().as(Options.class);
+    Pipeline p = Pipeline.create(options);
+
+    String instanceId = options.getInstanceId();
+    String databaseId = options.getDatabaseId();
+
+    String usersIdFile = options.getSuspiciousUsersFile();
+
+    PCollection<String> suspiciousUserIds = p.apply(TextIO.read().from(usersIdFile));
+
+    final Timestamp timestamp = Timestamp.now();
+
+    // [START spanner_dataflow_writegroup]
+    PCollection<MutationGroup> mutations = suspiciousUserIds
+        .apply(MapElements.via(new SimpleFunction<String, MutationGroup>() {
+
+          @Override
+          public MutationGroup apply(String userId) {
+            // Immediately block the user.
+            Mutation userMutation = Mutation.newUpdateBuilder("Users")
+                .set("id").to(userId)
+                .set("state").to("BLOCKED")
+                .build();
+            long generatedId = Hashing.sha1().newHasher()
+                .putString(userId, Charsets.UTF_8)
+                .putLong(timestamp.getSeconds())
+                .putLong(timestamp.getNanos())
+                .hash()
+                .asLong();
+
+            // Add an entry to pending review requests.
+            Mutation pendingReview = Mutation.newInsertOrUpdateBuilder("PendingReviews")
+                .set("id").to(generatedId)  // Must be deterministically generated.
+                .set("userId").to(userId)
+                .set("action").to("REVIEW ACCOUNT")
+                .set("node").to("Suspicious activity detected.")
+                .build();
+
+            return MutationGroup.create(userMutation, pendingReview);
+          }
+        }));
+
+    mutations.apply(SpannerIO.write()
+        .withInstanceId(instanceId)
+        .withDatabaseId(databaseId)
+        .grouped());
+    // [END spanner_dataflow_writegroup]
+
+  }
+
+}

dataflow/spanner-io/src/main/java/com/example/dataflow/SpannerRead.java

@@ -87,66 +87,26 @@ public interface Options extends PipelineOptions {
     void setOutput(String value);
   }
 
-  /**
-   * Estimates the size of a Spanner row. For simplicity, arrays and structs aren't supported.
-   */
-  public static class EstimateStructSizeFn extends DoFn<Struct, Long> {
-
-    @ProcessElement
-    public void processElement(ProcessContext c) throws Exception {
-      Struct row = c.element();
-      long sum = 0;
-      for (int i = 0; i < row.getColumnCount(); i++) {
-        if (row.isNull(i)) {
-          continue;
-        }
-
-        switch (row.getColumnType(i).getCode()) {
-          case BOOL:
-            sum += 1;
-            break;
-          case INT64:
-          case FLOAT64:
-            sum += 8;
-            break;
-          case TIMESTAMP:
-          case DATE:
-            sum += 12;
-            break;
-          case BYTES:
-            sum += row.getBytes(i).length();
-            break;
-          case STRING:
-            sum += row.getString(i).length();
-            break;
-          case ARRAY:
-            throw new IllegalArgumentException("Arrays are not supported :(");
-          case STRUCT:
-            throw new IllegalArgumentException("Structs are not supported :(");
-          default:
-            throw new IllegalArgumentException("Unsupported type :(");
-        }
-      }
-      c.output(sum);
-    }
-  }
 
   public static void main(String[] args) {
     Options options = PipelineOptionsFactory.fromArgs(args).withValidation().as(Options.class);
     Pipeline p = Pipeline.create(options);
 
     String instanceId = options.getInstanceId();
     String databaseId = options.getDatabaseId();
-    String query = "SELECT * FROM " + options.getTable();
-
-    PCollection<Long> tableEstimatedSize = p
-        // Query for all the columns and rows in the specified Spanner table
-        .apply(SpannerIO.read()
+    // [START spanner_dataflow_read]
+    // Query for all the columns and rows in the specified Spanner table
+    PCollection<Struct> records = p.apply(
+        SpannerIO.read()
             .withInstanceId(instanceId)
             .withDatabaseId(databaseId)
-            .withQuery(query))
+            .withQuery("SELECT * FROM " + options.getTable()));
+    // [START spanner_dataflow_read]
+
+
+    PCollection<Long> tableEstimatedSize = records
         // Estimate the size of every row
-        .apply(ParDo.of(new EstimateStructSizeFn()))
+        .apply(EstimateSize.create())
         // Sum all the row sizes to get the total estimated size of the table
         .apply(Sum.longsGlobally());
 

dataflow/spanner-io/src/main/java/com/example/dataflow/SpannerReadAll.java

@@ -0,0 +1,84 @@
+/*
+ * Copyright 2018 Google Inc.
+ *
+ * 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.example.dataflow;
+
+import com.google.cloud.spanner.Struct;
+import org.apache.beam.sdk.Pipeline;
+import org.apache.beam.sdk.io.TextIO;
+import org.apache.beam.sdk.io.gcp.spanner.ReadOperation;
+import org.apache.beam.sdk.io.gcp.spanner.SpannerIO;
+import org.apache.beam.sdk.options.Description;
+import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.options.PipelineOptionsFactory;
+import org.apache.beam.sdk.options.Validation;
+import org.apache.beam.sdk.transforms.MapElements;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.Sum;
+import org.apache.beam.sdk.transforms.ToString;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.TypeDescriptor;
+
+/**
+ * This sample demonstrates how to read all data from the Cloud Spanner database.
+ */
+public class SpannerReadAll {
+
+  public interface Options extends PipelineOptions {
+
+    @Description("Spanner instance ID to query from")
+    @Validation.Required
+    String getInstanceId();
+
+    void setInstanceId(String value);
+
+    @Description("Spanner database name to query from")
+    @Validation.Required
+    String getDatabaseId();
+
+    void setDatabaseId(String value);
+
+    @Description("Output filename for records size")
+    @Validation.Required
+    String getOutput();
+
+    void setOutput(String value);
+  }
+
+  public static void main(String[] args) {
+    Options options = PipelineOptionsFactory.fromArgs(args).withValidation().as(Options.class);
+    Pipeline p = Pipeline.create(options);
+
+    // [START spanner_dataflow_readall]
+    PCollection<Struct> allRecords = p.apply(SpannerIO.read().withQuery(
+        "SELECT t.table_name FROM information_schema.tables AS t WHERE t"
+            + ".table_catalog = '' AND t.table_schema = ''")).apply(
+        MapElements.into(TypeDescriptor.of(ReadOperation.class))
+            .via((SerializableFunction<Struct, ReadOperation>) input -> {
+              String tableName = input.getString(0);
+              return ReadOperation.create().withQuery("SELECT * FROM " + tableName);
+            })).apply(SpannerIO.readAll());
+    // [END spanner_dataflow_readall]
+
+    PCollection<Long> dbEstimatedSize = allRecords.apply(EstimateSize.create())
+        .apply(Sum.longsGlobally());
+
+    dbEstimatedSize.apply(ToString.elements()).apply(TextIO.write().to(options.getOutput()));
+
+    p.run().waitUntilFinish();
+  }
+
+}