Rollback extraneous cahnge and fix tests.

Author: reuvenlax

runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/DataflowRunner.java

@@ -89,6 +89,7 @@
 import org.apache.beam.sdk.io.FileSystems;
 import org.apache.beam.sdk.io.Read;
 import org.apache.beam.sdk.io.UnboundedSource;
+import org.apache.beam.sdk.io.WriteFiles;
 import org.apache.beam.sdk.io.fs.ResourceId;
 import org.apache.beam.sdk.io.gcp.pubsub.PubsubMessage;
 import org.apache.beam.sdk.io.gcp.pubsub.PubsubMessageWithAttributesCoder;
@@ -1302,7 +1303,7 @@ public Map<PValue, ReplacementOutput> mapOutputs(
 
   @VisibleForTesting
   static class StreamingShardedWriteFactory<T>
-      implements PTransformOverrideFactory<PCollection<T>, PDone, Write<T>> {
+      implements PTransformOverrideFactory<PCollection<T>, PDone, WriteFiles<T>> {
     DataflowPipelineWorkerPoolOptions options;
 
     StreamingShardedWriteFactory(PipelineOptions options) {
@@ -1311,14 +1312,15 @@ static class StreamingShardedWriteFactory<T>
 
     @Override
     public PTransformReplacement<PCollection<T>, PDone> getReplacementTransform(
-        AppliedPTransform<PCollection<T>, PDone, Write<T>> transform) {
+        AppliedPTransform<PCollection<T>, PDone, WriteFiles<T>> transform) {
         return PTransformReplacement.of(
             PTransformReplacements.getSingletonMainInput(transform),
             transform.getTransform().withNumShards(options.getMaxNumWorkers() * 2));
     }
 
     @Override
-    public Map<PValue, ReplacementOutput> mapOutputs(Map<TupleTag<?>, PValue> outputs, PDone newOutput) {
+    public Map<PValue, ReplacementOutput> mapOutputs(Map<TupleTag<?>, PValue> outputs,
+                                                     PDone newOutput) {
       return Collections.emptyMap();
     }
   }

runners/google-cloud-dataflow-java/src/test/java/org/apache/beam/runners/dataflow/DataflowRunnerTest.java

@@ -18,14 +18,12 @@
 package org.apache.beam.runners.dataflow;
 
 import static org.apache.beam.runners.dataflow.DataflowRunner.getContainerImageForJob;
-import static org.hamcrest.Matchers.allOf;
 import static org.hamcrest.Matchers.both;
-import static org.hamcrest.Matchers.containsInAnyOrder;
 import static org.hamcrest.Matchers.containsString;
 import static org.hamcrest.Matchers.equalTo;
-import static org.hamcrest.Matchers.greaterThan;
 import static org.hamcrest.Matchers.hasItem;
 import static org.hamcrest.Matchers.is;
+import static org.hamcrest.Matchers.not;
 import static org.hamcrest.Matchers.startsWith;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
@@ -52,17 +50,14 @@
 import java.io.File;
 import java.io.FileNotFoundException;
 import java.io.IOException;
-import java.io.Reader;
 import java.net.URL;
 import java.net.URLClassLoader;
-import java.nio.CharBuffer;
 import java.nio.channels.FileChannel;
 import java.nio.channels.SeekableByteChannel;
 import java.nio.file.Files;
 import java.nio.file.StandardOpenOption;
 import java.util.ArrayList;
 import java.util.Arrays;
-import java.util.Collection;
 import java.util.Collections;
 import java.util.LinkedList;
 import java.util.List;
@@ -71,26 +66,27 @@
 import org.apache.beam.runners.dataflow.DataflowRunner.StreamingShardedWriteFactory;
 import org.apache.beam.runners.dataflow.options.DataflowPipelineDebugOptions;
 import org.apache.beam.runners.dataflow.options.DataflowPipelineOptions;
-import org.apache.beam.runners.dataflow.options.DataflowPipelineWorkerPoolOptions;
 import org.apache.beam.sdk.Pipeline;
 import org.apache.beam.sdk.Pipeline.PipelineVisitor;
 import org.apache.beam.sdk.coders.BigEndianIntegerCoder;
 import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.coders.VoidCoder;
 import org.apache.beam.sdk.extensions.gcp.auth.NoopCredentialFactory;
 import org.apache.beam.sdk.extensions.gcp.auth.TestCredential;
+import org.apache.beam.sdk.io.FileBasedSink;
 import org.apache.beam.sdk.io.FileSystems;
 import org.apache.beam.sdk.io.TextIO;
-import org.apache.beam.sdk.io.Write;
+import org.apache.beam.sdk.io.WriteFiles;
 import org.apache.beam.sdk.options.PipelineOptions;
 import org.apache.beam.sdk.options.PipelineOptions.CheckEnabled;
 import org.apache.beam.sdk.options.PipelineOptionsFactory;
 import org.apache.beam.sdk.options.ValueProvider;
 import org.apache.beam.sdk.options.ValueProvider.StaticValueProvider;
+import org.apache.beam.sdk.runners.AppliedPTransform;
 import org.apache.beam.sdk.runners.TransformHierarchy;
 import org.apache.beam.sdk.runners.TransformHierarchy.Node;
 import org.apache.beam.sdk.testing.ExpectedLogs;
 import org.apache.beam.sdk.testing.TestPipeline;
-import org.apache.beam.sdk.transforms.AppliedPTransform;
 import org.apache.beam.sdk.transforms.Create;
 import org.apache.beam.sdk.transforms.PTransform;
 import org.apache.beam.sdk.util.GcsUtil;
@@ -101,6 +97,7 @@
 import org.apache.beam.sdk.values.PDone;
 import org.apache.beam.sdk.values.PValue;
 import org.apache.beam.sdk.values.TimestampedValue;
+import org.apache.beam.sdk.values.TupleTag;
 import org.apache.beam.sdk.values.WindowingStrategy;
 import org.hamcrest.Description;
 import org.hamcrest.Matchers;
@@ -1130,59 +1127,31 @@ public void testWorkerHarnessContainerImage() {
     assertThat(getContainerImageForJob(options), equalTo("gcr.io/java/foo"));
   }
 
-  static class SimpleFilenamePolicy extends  FilenamePolicy {
-    String baseLocation;
-
-    SimpleFilenamePolicy(String baseLocation) {
-      this.baseLocation = baseLocation;
-    }
-
-    @Override
-    public String windowedFilename(WindowedContext c) {
-      return baseLocation + "-shard-" + c.getShardNumber();
-    }
-
-    @Override
-    public String unwindowedFilename(Context c) {
-      throw new UnsupportedOperationException("Unsupported");
-    }
-
-    @Override
-    public ValueProvider<String> getBaseOutputFilenameProvider() {
-      return StaticValueProvider.of(baseLocation);
-    }
-  };
-
-  PipelineOptions getOptions() {
-    PipelineOptions options = TestPipeline.testingPipelineOptions();
-    options.setRunner(DataflowRunner.class);
-    options.as(DataflowPipelineOptions.class).setStreaming(true);
-    options.as(DataflowPipelineWorkerPoolOptions.class).setMaxNumWorkers(10);
-    return options;
-  }
-  @Rule public final transient TestPipeline p = TestPipeline.fromOptions(getOptions());
-
   @Test
   public void testStreamingWriteWithNoShardingReturnsNewTransform() {
     TestPipeline p = TestPipeline.create();
     StreamingShardedWriteFactory<Object> factory = new StreamingShardedWriteFactory(p.getOptions());
-    Write<Object> original = Write.to(new TestSink());
+    WriteFiles<Object> original = WriteFiles.to(new TestSink(tmpFolder.toString()));
     PCollection<Object> objs = (PCollection) p.apply(Create.empty(VoidCoder.of()));
-    AppliedPTransform<PCollection<Object>, PDone, Write<Object>> originalApplication =
+    AppliedPTransform<PCollection<Object>, PDone, WriteFiles<Object>> originalApplication =
         AppliedPTransform.of(
-            "write", objs.expand(), Collections.<TupleTag<?>, PValue>emptyMap(), original, p);
+            "writefiles", objs.expand(), Collections.<TupleTag<?>, PValue>emptyMap(), original, p);
 
     assertThat(
         factory.getReplacementTransform(originalApplication).getTransform(),
         not(equalTo((Object) original)));
   }
 
-  private static class TestSink extends Sink<Object> {
+  private static class TestSink extends FileBasedSink<Object> {
     @Override
     public void validate(PipelineOptions options) {}
 
+    TestSink(String tmpFolder) {
+      super(StaticValueProvider.of(FileSystems.matchNewResource(tmpFolder, true)),
+          null);
+    }
     @Override
-    public WriteOperation<Object, ?> createWriteOperation(PipelineOptions options) {
+    public FileBasedWriteOperation<Object> createWriteOperation() {
       throw new IllegalArgumentException("Should not be used");
     }
   }

sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubIO.java

@@ -154,43 +154,6 @@ private static void populateCommonDisplayData(DisplayData.Builder builder,
     }
   }
 
-  /**
-   * Class representing a Pub/Sub message. Each message contains a single message payload and
-   * a map of attached attributes.
-   */
-  public static class PubsubMessage {
-    private byte[] message;
-    private Map<String, String> attributes;
-
-    public PubsubMessage(byte[] message, Map<String, String> attributes) {
-      this.message = message;
-      this.attributes = attributes;
-    }
-
-    /**
-     * Returns the main PubSub message.
-     */
-    public byte[] getMessage() {
-      return message;
-    }
-
-    /**
-     * Returns the given attribute value. If not such attribute exists, returns null.
-     */
-    @Nullable
-    public String getAttribute(String attribute) {
-      checkNotNull(attribute, "attribute");
-      return attributes.get(attribute);
-    }
-
-    /**
-     * Returns the full map of attributes. This is an unmodifiable map.
-     */
-    public Map<String, String> getAttributeMap() {
-      return attributes;
-    }
-  }
-
   /**
    * Class representing a Cloud Pub/Sub Subscription.
    */
@@ -890,7 +853,6 @@ public PDone expand(PCollection<T> input) {
               getTimestampAttribute(),
               getIdAttribute(),
               100 /* numShards */));
-
       }
       throw new RuntimeException(); // cases are exhaustive.
     }