Dataflow SpannerIO sample

dataflow/spanner-io/pom.xml

@@ -0,0 +1,93 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<project xmlns="http://maven.apache.org/POM/4.0.0"
+    xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
+    xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd">
+  <modelVersion>4.0.0</modelVersion>
+
+  <groupId>com.example.dataflow</groupId>
+  <artifactId>dataflow-spanner</artifactId>
+  <version>1.0-SNAPSHOT</version>
+  <packaging>jar</packaging>
+
+  <properties>
+    <project.build.sourceEncoding>UTF-8</project.build.sourceEncoding>
+    <java.version>1.8</java.version>
+    <apache_beam.version>2.3.0-SNAPSHOT</apache_beam.version>
+  </properties>
+
+  <build>
+    <plugins>
+      <plugin>
+        <groupId>org.apache.maven.plugins</groupId>
+        <artifactId>maven-compiler-plugin</artifactId>
+        <version>3.5</version>
+        <configuration>
+          <source>1.8</source>
+          <target>1.8</target>
+        </configuration>
+      </plugin>
+    </plugins>
+  </build>
+
+  <dependencies>
+
+    <!-- Apache Beam -->
+    <dependency>
+      <groupId>org.apache.beam</groupId>
+      <artifactId>beam-sdks-java-core</artifactId>
+      <version>${apache_beam.version}</version>
+    </dependency>
+
+    <dependency>
+      <groupId>org.apache.beam</groupId>
+      <artifactId>beam-sdks-java-io-google-cloud-platform</artifactId>
+      <version>${apache_beam.version}</version>
+      <exclusions>
+        <exclusion>
+          <groupId>com.google.api.grpc</groupId>
+          <artifactId>grpc-google-common-protos</artifactId>
+        </exclusion>
+      </exclusions>
+    </dependency>
+
+    <dependency>
+      <groupId>org.apache.beam</groupId>
+      <artifactId>beam-runners-direct-java</artifactId>
+      <version>${apache_beam.version}</version>
+    </dependency>
+
+    <dependency>
+      <groupId>org.apache.beam</groupId>
+      <artifactId>beam-runners-google-cloud-dataflow-java</artifactId>
+      <version>${apache_beam.version}</version>
+    </dependency>
+
+    <!-- Google Cloud -->
+    <dependency>
+      <groupId>com.google.api.grpc</groupId>
+      <artifactId>proto-google-common-protos</artifactId>
+      <version>0.1.9</version>
+    </dependency>
+
+    <dependency>
+      <groupId>com.google.cloud</groupId>
+      <artifactId>google-cloud-spanner</artifactId>
+      <version>0.20.0-beta</version>
+    </dependency>
+
+    <dependency>
+      <groupId>com.google.api.grpc</groupId>
+      <artifactId>proto-google-cloud-spanner-admin-database-v1</artifactId>
+      <version>0.1.9</version>
+    </dependency>
+
+    <!-- Misc -->
+    <dependency>
+      <groupId>org.slf4j</groupId>
+      <artifactId>slf4j-jdk14</artifactId>
+      <version>1.7.25</version>
+    </dependency>
+
+  </dependencies>
+
+</project>

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

@@ -0,0 +1,121 @@
+/*
+  Copyright 2016, 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.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.Count;
+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.transforms.Sum;
+import org.apache.beam.sdk.transforms.ToString;
+import org.apache.beam.sdk.values.PCollection;
+
+public class Spanner {
+
+  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("Spanner table name to query from")
+    @Validation.Required
+    String getTableName();
+    void setTableName(String value);
+
+    @Description("Output filename for records size")
+    @Validation.Required
+    String getOutput();
+    void setOutput(String value);
+  }
+
+  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 :(");
+        }
+      }
+      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.getTableName();
+
+    PCollection<Long> tableEstimatedSize = p
+        .apply(SpannerIO.read()
+            .withInstanceId(instanceId)
+            .withDatabaseId(databaseId)
+            .withQuery(query))
+        .apply(ParDo.of(new EstimateStructSizeFn()))
+        .apply(Sum.longsGlobally());
+
+    tableEstimatedSize
+        .apply(ToString.elements())
+        .apply(TextIO.write().to(options.getOutput()));
+
+    p.run().waitUntilFinish();
+  }
+}