http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/harness/src/test/java/org/apache/beam/runners/core/BoundedSourceRunnerTest.java ---------------------------------------------------------------------- diff --git a/sdks/java/harness/src/test/java/org/apache/beam/runners/core/BoundedSourceRunnerTest.java b/sdks/java/harness/src/test/java/org/apache/beam/runners/core/BoundedSourceRunnerTest.java index 6c9a4cb..d8ed121 100644 --- a/sdks/java/harness/src/test/java/org/apache/beam/runners/core/BoundedSourceRunnerTest.java +++ b/sdks/java/harness/src/test/java/org/apache/beam/runners/core/BoundedSourceRunnerTest.java @@ -20,35 +20,25 @@ package org.apache.beam.runners.core; import static org.apache.beam.sdk.util.WindowedValue.valueInGlobalWindow; import static org.hamcrest.Matchers.contains; -import static org.hamcrest.Matchers.containsInAnyOrder; import static org.hamcrest.collection.IsEmptyCollection.empty; import static org.junit.Assert.assertThat; -import static org.junit.Assert.fail; -import com.google.common.base.Suppliers; -import com.google.common.collect.HashMultimap; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; -import com.google.common.collect.Iterables; -import com.google.common.collect.Multimap; import com.google.protobuf.Any; import com.google.protobuf.ByteString; import com.google.protobuf.BytesValue; import java.util.ArrayList; import java.util.Collection; import java.util.List; -import java.util.ServiceLoader; +import java.util.Map; import org.apache.beam.fn.harness.fn.ThrowingConsumer; -import org.apache.beam.fn.harness.fn.ThrowingRunnable; -import org.apache.beam.runners.core.PTransformRunnerFactory.Registrar; import org.apache.beam.sdk.common.runner.v1.RunnerApi; import org.apache.beam.sdk.io.BoundedSource; import org.apache.beam.sdk.io.CountingSource; import org.apache.beam.sdk.options.PipelineOptionsFactory; import org.apache.beam.sdk.util.SerializableUtils; import org.apache.beam.sdk.util.WindowedValue; -import org.hamcrest.Matchers; -import org.hamcrest.collection.IsMapContaining; import org.junit.Test; import org.junit.runner.RunWith; import org.junit.runners.JUnit4; @@ -56,25 +46,27 @@ import org.junit.runners.JUnit4; /** Tests for {@link BoundedSourceRunner}. */ @RunWith(JUnit4.class) public class BoundedSourceRunnerTest { - - public static final String URN = "urn:org.apache.beam:source:java:0.1"; - @Test public void testRunReadLoopWithMultipleSources() throws Exception { - List<WindowedValue<Long>> out1Values = new ArrayList<>(); + List<WindowedValue<Long>> out1ValuesA = new ArrayList<>(); + List<WindowedValue<Long>> out1ValuesB = new ArrayList<>(); List<WindowedValue<Long>> out2Values = new ArrayList<>(); - Collection<ThrowingConsumer<WindowedValue<Long>>> consumers = - ImmutableList.of(out1Values::add, out2Values::add); + Map<String, Collection<ThrowingConsumer<WindowedValue<Long>>>> outputMap = ImmutableMap.of( + "out1", ImmutableList.of(out1ValuesA::add, out1ValuesB::add), + "out2", ImmutableList.of(out2Values::add)); - BoundedSourceRunner<BoundedSource<Long>, Long> runner = new BoundedSourceRunner<>( + BoundedSourceRunner<BoundedSource<Long>, Long> runner = + new BoundedSourceRunner<>( PipelineOptionsFactory.create(), RunnerApi.FunctionSpec.getDefaultInstance(), - consumers); + outputMap); runner.runReadLoop(valueInGlobalWindow(CountingSource.upTo(2))); runner.runReadLoop(valueInGlobalWindow(CountingSource.upTo(1))); - assertThat(out1Values, + assertThat(out1ValuesA, + contains(valueInGlobalWindow(0L), valueInGlobalWindow(1L), valueInGlobalWindow(0L))); + assertThat(out1ValuesB, contains(valueInGlobalWindow(0L), valueInGlobalWindow(1L), valueInGlobalWindow(0L))); assertThat(out2Values, contains(valueInGlobalWindow(0L), valueInGlobalWindow(1L), valueInGlobalWindow(0L))); @@ -82,106 +74,40 @@ public class BoundedSourceRunnerTest { @Test public void testRunReadLoopWithEmptySource() throws Exception { - List<WindowedValue<Long>> outValues = new ArrayList<>(); - Collection<ThrowingConsumer<WindowedValue<Long>>> consumers = - ImmutableList.of(outValues::add); + List<WindowedValue<Long>> out1Values = new ArrayList<>(); + Map<String, Collection<ThrowingConsumer<WindowedValue<Long>>>> outputMap = ImmutableMap.of( + "out1", ImmutableList.of(out1Values::add)); - BoundedSourceRunner<BoundedSource<Long>, Long> runner = new BoundedSourceRunner<>( + BoundedSourceRunner<BoundedSource<Long>, Long> runner = + new BoundedSourceRunner<>( PipelineOptionsFactory.create(), RunnerApi.FunctionSpec.getDefaultInstance(), - consumers); + outputMap); runner.runReadLoop(valueInGlobalWindow(CountingSource.upTo(0))); - assertThat(outValues, empty()); + assertThat(out1Values, empty()); } @Test public void testStart() throws Exception { List<WindowedValue<Long>> outValues = new ArrayList<>(); - Collection<ThrowingConsumer<WindowedValue<Long>>> consumers = - ImmutableList.of(outValues::add); + Map<String, Collection<ThrowingConsumer<WindowedValue<Long>>>> outputMap = ImmutableMap.of( + "out", ImmutableList.of(outValues::add)); ByteString encodedSource = ByteString.copyFrom(SerializableUtils.serializeToByteArray(CountingSource.upTo(3))); - BoundedSourceRunner<BoundedSource<Long>, Long> runner = new BoundedSourceRunner<>( + BoundedSourceRunner<BoundedSource<Long>, Long> runner = + new BoundedSourceRunner<>( PipelineOptionsFactory.create(), - RunnerApi.FunctionSpec.newBuilder().setParameter( + RunnerApi.FunctionSpec.newBuilder().setParameter( Any.pack(BytesValue.newBuilder().setValue(encodedSource).build())).build(), - consumers); + outputMap); runner.start(); assertThat(outValues, contains(valueInGlobalWindow(0L), valueInGlobalWindow(1L), valueInGlobalWindow(2L))); } - - @Test - public void testCreatingAndProcessingSourceFromFactory() throws Exception { - List<WindowedValue<String>> outputValues = new ArrayList<>(); - - Multimap<String, ThrowingConsumer<WindowedValue<?>>> consumers = HashMultimap.create(); - consumers.put("outputPC", - (ThrowingConsumer) (ThrowingConsumer<WindowedValue<String>>) outputValues::add); - List<ThrowingRunnable> startFunctions = new ArrayList<>(); - List<ThrowingRunnable> finishFunctions = new ArrayList<>(); - - RunnerApi.FunctionSpec functionSpec = RunnerApi.FunctionSpec.newBuilder() - .setUrn("urn:org.apache.beam:source:java:0.1") - .setParameter(Any.pack(BytesValue.newBuilder() - .setValue(ByteString.copyFrom( - SerializableUtils.serializeToByteArray(CountingSource.upTo(3)))) - .build())) - .build(); - - RunnerApi.PTransform pTransform = RunnerApi.PTransform.newBuilder() - .setSpec(functionSpec) - .putInputs("input", "inputPC") - .putOutputs("output", "outputPC") - .build(); - - new BoundedSourceRunner.Factory<>().createRunnerForPTransform( - PipelineOptionsFactory.create(), - null /* beamFnDataClient */, - "pTransformId", - pTransform, - Suppliers.ofInstance("57L")::get, - ImmutableMap.of(), - ImmutableMap.of(), - consumers, - startFunctions::add, - finishFunctions::add); - - // This is testing a deprecated way of running sources and should be removed - // once all source definitions are instead propagated along the input edge. - Iterables.getOnlyElement(startFunctions).run(); - assertThat(outputValues, contains( - valueInGlobalWindow(0L), - valueInGlobalWindow(1L), - valueInGlobalWindow(2L))); - outputValues.clear(); - - // Check that when passing a source along as an input, the source is processed. - assertThat(consumers.keySet(), containsInAnyOrder("inputPC", "outputPC")); - Iterables.getOnlyElement(consumers.get("inputPC")).accept( - valueInGlobalWindow(CountingSource.upTo(2))); - assertThat(outputValues, contains( - valueInGlobalWindow(0L), - valueInGlobalWindow(1L))); - - assertThat(finishFunctions, Matchers.empty()); - } - - @Test - public void testRegistration() { - for (Registrar registrar : - ServiceLoader.load(Registrar.class)) { - if (registrar instanceof BoundedSourceRunner.Registrar) { - assertThat(registrar.getPTransformRunnerFactories(), IsMapContaining.hasKey(URN)); - return; - } - } - fail("Expected registrar not found."); - } }
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/harness/src/test/java/org/apache/beam/runners/core/FnApiDoFnRunnerTest.java ---------------------------------------------------------------------- diff --git a/sdks/java/harness/src/test/java/org/apache/beam/runners/core/FnApiDoFnRunnerTest.java b/sdks/java/harness/src/test/java/org/apache/beam/runners/core/FnApiDoFnRunnerTest.java deleted file mode 100644 index c4df77a..0000000 --- a/sdks/java/harness/src/test/java/org/apache/beam/runners/core/FnApiDoFnRunnerTest.java +++ /dev/null @@ -1,210 +0,0 @@ -/* - * Licensed to the Apache Software Foundation (ASF) under one - * or more contributor license agreements. See the NOTICE file - * distributed with this work for additional information - * regarding copyright ownership. The ASF licenses this file - * to you 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 org.apache.beam.runners.core; - -import static org.apache.beam.sdk.util.WindowedValue.timestampedValueInGlobalWindow; -import static org.apache.beam.sdk.util.WindowedValue.valueInGlobalWindow; -import static org.hamcrest.Matchers.contains; -import static org.hamcrest.Matchers.containsInAnyOrder; -import static org.junit.Assert.assertThat; -import static org.junit.Assert.fail; - -import com.fasterxml.jackson.databind.ObjectMapper; -import com.google.common.base.Suppliers; -import com.google.common.collect.HashMultimap; -import com.google.common.collect.ImmutableList; -import com.google.common.collect.ImmutableMap; -import com.google.common.collect.Iterables; -import com.google.common.collect.Multimap; -import com.google.protobuf.Any; -import com.google.protobuf.ByteString; -import com.google.protobuf.BytesValue; -import com.google.protobuf.Message; -import java.io.IOException; -import java.util.ArrayList; -import java.util.List; -import java.util.Map; -import java.util.ServiceLoader; -import org.apache.beam.fn.harness.fn.ThrowingConsumer; -import org.apache.beam.fn.harness.fn.ThrowingRunnable; -import org.apache.beam.runners.core.PTransformRunnerFactory.Registrar; -import org.apache.beam.runners.core.construction.ParDoTranslation; -import org.apache.beam.runners.dataflow.util.CloudObjects; -import org.apache.beam.runners.dataflow.util.DoFnInfo; -import org.apache.beam.sdk.coders.Coder; -import org.apache.beam.sdk.coders.StringUtf8Coder; -import org.apache.beam.sdk.common.runner.v1.RunnerApi; -import org.apache.beam.sdk.options.PipelineOptionsFactory; -import org.apache.beam.sdk.transforms.DoFn; -import org.apache.beam.sdk.transforms.windowing.BoundedWindow; -import org.apache.beam.sdk.transforms.windowing.GlobalWindow; -import org.apache.beam.sdk.util.SerializableUtils; -import org.apache.beam.sdk.util.WindowedValue; -import org.apache.beam.sdk.values.TupleTag; -import org.apache.beam.sdk.values.WindowingStrategy; -import org.hamcrest.collection.IsMapContaining; -import org.junit.Test; -import org.junit.runner.RunWith; -import org.junit.runners.JUnit4; - -/** Tests for {@link FnApiDoFnRunner}. */ -@RunWith(JUnit4.class) -public class FnApiDoFnRunnerTest { - - private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper(); - private static final Coder<WindowedValue<String>> STRING_CODER = - WindowedValue.getFullCoder(StringUtf8Coder.of(), GlobalWindow.Coder.INSTANCE); - private static final String STRING_CODER_SPEC_ID = "999L"; - private static final RunnerApi.Coder STRING_CODER_SPEC; - - static { - try { - STRING_CODER_SPEC = RunnerApi.Coder.newBuilder() - .setSpec(RunnerApi.SdkFunctionSpec.newBuilder() - .setSpec(RunnerApi.FunctionSpec.newBuilder() - .setParameter(Any.pack(BytesValue.newBuilder().setValue(ByteString.copyFrom( - OBJECT_MAPPER.writeValueAsBytes(CloudObjects.asCloudObject(STRING_CODER)))) - .build()))) - .build()) - .build(); - } catch (IOException e) { - throw new ExceptionInInitializerError(e); - } - } - - private static class TestDoFn extends DoFn<String, String> { - private static final TupleTag<String> mainOutput = new TupleTag<>("mainOutput"); - private static final TupleTag<String> additionalOutput = new TupleTag<>("output"); - - private BoundedWindow window; - - @ProcessElement - public void processElement(ProcessContext context, BoundedWindow window) { - context.output("MainOutput" + context.element()); - context.output(additionalOutput, "AdditionalOutput" + context.element()); - this.window = window; - } - - @FinishBundle - public void finishBundle(FinishBundleContext context) { - if (window != null) { - context.output("FinishBundle", window.maxTimestamp(), window); - window = null; - } - } - } - - /** - * Create a DoFn that has 3 inputs (inputATarget1, inputATarget2, inputBTarget) and 2 outputs - * (mainOutput, output). Validate that inputs are fed to the {@link DoFn} and that outputs - * are directed to the correct consumers. - */ - @Test - public void testCreatingAndProcessingDoFn() throws Exception { - Map<String, Message> fnApiRegistry = ImmutableMap.of(STRING_CODER_SPEC_ID, STRING_CODER_SPEC); - String pTransformId = "pTransformId"; - String mainOutputId = "101"; - String additionalOutputId = "102"; - - DoFnInfo<?, ?> doFnInfo = DoFnInfo.forFn( - new TestDoFn(), - WindowingStrategy.globalDefault(), - ImmutableList.of(), - StringUtf8Coder.of(), - Long.parseLong(mainOutputId), - ImmutableMap.of( - Long.parseLong(mainOutputId), TestDoFn.mainOutput, - Long.parseLong(additionalOutputId), TestDoFn.additionalOutput)); - RunnerApi.FunctionSpec functionSpec = RunnerApi.FunctionSpec.newBuilder() - .setUrn(ParDoTranslation.CUSTOM_JAVA_DO_FN_URN) - .setParameter(Any.pack(BytesValue.newBuilder() - .setValue(ByteString.copyFrom(SerializableUtils.serializeToByteArray(doFnInfo))) - .build())) - .build(); - RunnerApi.PTransform pTransform = RunnerApi.PTransform.newBuilder() - .setSpec(functionSpec) - .putInputs("inputA", "inputATarget") - .putInputs("inputB", "inputBTarget") - .putOutputs(mainOutputId, "mainOutputTarget") - .putOutputs(additionalOutputId, "additionalOutputTarget") - .build(); - - List<WindowedValue<String>> mainOutputValues = new ArrayList<>(); - List<WindowedValue<String>> additionalOutputValues = new ArrayList<>(); - Multimap<String, ThrowingConsumer<WindowedValue<?>>> consumers = HashMultimap.create(); - consumers.put("mainOutputTarget", - (ThrowingConsumer) (ThrowingConsumer<WindowedValue<String>>) mainOutputValues::add); - consumers.put("additionalOutputTarget", - (ThrowingConsumer) (ThrowingConsumer<WindowedValue<String>>) additionalOutputValues::add); - List<ThrowingRunnable> startFunctions = new ArrayList<>(); - List<ThrowingRunnable> finishFunctions = new ArrayList<>(); - - new FnApiDoFnRunner.Factory<>().createRunnerForPTransform( - PipelineOptionsFactory.create(), - null /* beamFnDataClient */, - pTransformId, - pTransform, - Suppliers.ofInstance("57L")::get, - ImmutableMap.of(), - ImmutableMap.of(), - consumers, - startFunctions::add, - finishFunctions::add); - - Iterables.getOnlyElement(startFunctions).run(); - mainOutputValues.clear(); - - assertThat(consumers.keySet(), containsInAnyOrder( - "inputATarget", "inputBTarget", "mainOutputTarget", "additionalOutputTarget")); - - Iterables.getOnlyElement(consumers.get("inputATarget")).accept(valueInGlobalWindow("A1")); - Iterables.getOnlyElement(consumers.get("inputATarget")).accept(valueInGlobalWindow("A2")); - Iterables.getOnlyElement(consumers.get("inputATarget")).accept(valueInGlobalWindow("B")); - assertThat(mainOutputValues, contains( - valueInGlobalWindow("MainOutputA1"), - valueInGlobalWindow("MainOutputA2"), - valueInGlobalWindow("MainOutputB"))); - assertThat(additionalOutputValues, contains( - valueInGlobalWindow("AdditionalOutputA1"), - valueInGlobalWindow("AdditionalOutputA2"), - valueInGlobalWindow("AdditionalOutputB"))); - mainOutputValues.clear(); - additionalOutputValues.clear(); - - Iterables.getOnlyElement(finishFunctions).run(); - assertThat( - mainOutputValues, - contains( - timestampedValueInGlobalWindow("FinishBundle", GlobalWindow.INSTANCE.maxTimestamp()))); - mainOutputValues.clear(); - } - - @Test - public void testRegistration() { - for (Registrar registrar : - ServiceLoader.load(Registrar.class)) { - if (registrar instanceof FnApiDoFnRunner.Registrar) { - assertThat(registrar.getPTransformRunnerFactories(), - IsMapContaining.hasKey(ParDoTranslation.CUSTOM_JAVA_DO_FN_URN)); - return; - } - } - fail("Expected registrar not found."); - } -} http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/amqp/pom.xml ---------------------------------------------------------------------- diff --git a/sdks/java/io/amqp/pom.xml b/sdks/java/io/amqp/pom.xml deleted file mode 100644 index 8da9448..0000000 --- a/sdks/java/io/amqp/pom.xml +++ /dev/null @@ -1,100 +0,0 @@ -<?xml version="1.0" encoding="UTF-8"?> -<!-- - Licensed to the Apache Software Foundation (ASF) under one or more - contributor license agreements. See the NOTICE file distributed with - this work for additional information regarding copyright ownership. - The ASF licenses this file to You 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. ---> -<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> - - <parent> - <groupId>org.apache.beam</groupId> - <artifactId>beam-sdks-java-io-parent</artifactId> - <version>2.2.0-SNAPSHOT</version> - <relativePath>../pom.xml</relativePath> - </parent> - - <artifactId>beam-sdks-java-io-amqp</artifactId> - <name>Apache Beam :: SDKs :: Java :: IO :: AMQP</name> - <description>IO to read and write using AMQP 1.0 protocol (http://www.amqp.org).</description> - - <dependencies> - <dependency> - <groupId>org.apache.beam</groupId> - <artifactId>beam-sdks-java-core</artifactId> - </dependency> - - <dependency> - <groupId>org.slf4j</groupId> - <artifactId>slf4j-api</artifactId> - </dependency> - - <dependency> - <groupId>joda-time</groupId> - <artifactId>joda-time</artifactId> - </dependency> - - <dependency> - <groupId>com.google.guava</groupId> - <artifactId>guava</artifactId> - </dependency> - - <dependency> - <groupId>com.google.code.findbugs</groupId> - <artifactId>jsr305</artifactId> - </dependency> - - <dependency> - <groupId>org.apache.qpid</groupId> - <artifactId>proton-j</artifactId> - <version>0.13.1</version> - </dependency> - - <!-- compile dependencies --> - <dependency> - <groupId>com.google.auto.value</groupId> - <artifactId>auto-value</artifactId> - <scope>provided</scope> - </dependency> - <dependency> - <groupId>com.google.auto.service</groupId> - <artifactId>auto-service</artifactId> - <optional>true</optional> - </dependency> - - <!-- test dependencies --> - <dependency> - <groupId>org.slf4j</groupId> - <artifactId>slf4j-jdk14</artifactId> - <scope>test</scope> - </dependency> - <dependency> - <groupId>junit</groupId> - <artifactId>junit</artifactId> - <scope>test</scope> - </dependency> - <dependency> - <groupId>org.hamcrest</groupId> - <artifactId>hamcrest-all</artifactId> - <scope>test</scope> - </dependency> - <dependency> - <groupId>org.apache.beam</groupId> - <artifactId>beam-runners-direct-java</artifactId> - <scope>test</scope> - </dependency> - </dependencies> - -</project> \ No newline at end of file http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/amqp/src/main/java/org/apache/beam/sdk/io/amqp/AmqpIO.java ---------------------------------------------------------------------- diff --git a/sdks/java/io/amqp/src/main/java/org/apache/beam/sdk/io/amqp/AmqpIO.java b/sdks/java/io/amqp/src/main/java/org/apache/beam/sdk/io/amqp/AmqpIO.java deleted file mode 100644 index 1f307b2..0000000 --- a/sdks/java/io/amqp/src/main/java/org/apache/beam/sdk/io/amqp/AmqpIO.java +++ /dev/null @@ -1,399 +0,0 @@ -/* - * Licensed to the Apache Software Foundation (ASF) under one - * or more contributor license agreements. See the NOTICE file - * distributed with this work for additional information - * regarding copyright ownership. The ASF licenses this file - * to you 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 org.apache.beam.sdk.io.amqp; - -import static com.google.common.base.Preconditions.checkArgument; -import static com.google.common.base.Preconditions.checkState; - -import com.google.auto.value.AutoValue; -import com.google.common.base.Joiner; - -import java.io.IOException; -import java.io.Serializable; -import java.util.ArrayList; -import java.util.List; -import java.util.NoSuchElementException; - -import javax.annotation.Nullable; - -import org.apache.beam.sdk.annotations.Experimental; -import org.apache.beam.sdk.coders.Coder; -import org.apache.beam.sdk.coders.SerializableCoder; -import org.apache.beam.sdk.io.UnboundedSource; -import org.apache.beam.sdk.options.PipelineOptions; -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.display.DisplayData; -import org.apache.beam.sdk.values.PBegin; -import org.apache.beam.sdk.values.PCollection; -import org.apache.beam.sdk.values.PDone; -import org.apache.qpid.proton.message.Message; -import org.apache.qpid.proton.messenger.Messenger; -import org.apache.qpid.proton.messenger.Tracker; -import org.joda.time.Duration; -import org.joda.time.Instant; - -/** - * AmqpIO supports AMQP 1.0 protocol using the Apache QPid Proton-J library. - * - * <p>It's also possible to use AMQP 1.0 protocol via Apache Qpid JMS connection factory and the - * Apache Beam JmsIO. - * - * <h3>Binding AMQP and receive messages</h3> - * - * <p>The {@link AmqpIO} {@link Read} can bind a AMQP listener endpoint and receive messages. It can - * also connect to a AMPQ broker (such as Apache Qpid or Apache ActiveMQ). - * - * <p>{@link AmqpIO} {@link Read} returns an unbounded {@link PCollection} of {@link Message} - * containing the received messages. - * - * <p>To configure a AMQP source, you have to provide a list of addresses where it will receive - * messages. An address has the following form: {@code - * [amqp[s]://][user[:password]@]domain[/[name]]} where {@code domain} can be one of {@code - * host | host:port | ip | ip:port | name}. NB: the {@code ~} character allows to bind a AMQP - * listener instead of connecting to a remote broker. For instance {@code amqp://~0.0.0.0:1234} - * will bind a AMQP listener on any network interface on the 1234 port number. - * - * <p>The following example illustrates how to configure a AMQP source: - * - * <pre>{@code - * - * pipeline.apply(AmqpIO.read() - * .withAddresses(Collections.singletonList("amqp://host:1234"))) - * - * }</pre> - * - * <h3>Sending messages to a AMQP endpoint</h3> - * - * <p>{@link AmqpIO} provides a sink to send {@link PCollection} elements as messages. - * - * <p>As for the {@link Read}, {@link AmqpIO} {@link Write} requires a list of addresses where to - * send messages. The following example illustrates how to configure the {@link AmqpIO} - * {@link Write}: - * - * <pre>{@code - * - * pipeline - * .apply(...) // provide PCollection<Message> - * .apply(AmqpIO.write()); - * - * }</pre> - */ -@Experimental(Experimental.Kind.SOURCE_SINK) -public class AmqpIO { - - public static Read read() { - return new AutoValue_AmqpIO_Read.Builder().setMaxNumRecords(Long.MAX_VALUE).build(); - } - - public static Write write() { - return new AutoValue_AmqpIO_Write(); - } - - private AmqpIO() { - } - - /** - * A {@link PTransform} to read/receive messages using AMQP 1.0 protocol. - */ - @AutoValue - public abstract static class Read extends PTransform<PBegin, PCollection<Message>> { - - @Nullable abstract List<String> addresses(); - abstract long maxNumRecords(); - @Nullable abstract Duration maxReadTime(); - - abstract Builder builder(); - - @AutoValue.Builder - abstract static class Builder { - abstract Builder setAddresses(List<String> addresses); - abstract Builder setMaxNumRecords(long maxNumRecords); - abstract Builder setMaxReadTime(Duration maxReadTime); - abstract Read build(); - } - - /** - * Define the AMQP addresses where to receive messages. - */ - public Read withAddresses(List<String> addresses) { - checkArgument(addresses != null, "AmqpIO.read().withAddresses(addresses) called with null" - + " addresses"); - checkArgument(!addresses.isEmpty(), "AmqpIO.read().withAddresses(addresses) called with " - + "empty addresses list"); - return builder().setAddresses(addresses).build(); - } - - /** - * Define the max number of records received by the {@link Read}. - * When the max number of records is lower than {@code Long.MAX_VALUE}, the {@link Read} will - * provide a bounded {@link PCollection}. - */ - public Read withMaxNumRecords(long maxNumRecords) { - checkArgument(maxReadTime() == null, - "maxNumRecord and maxReadTime are exclusive"); - return builder().setMaxNumRecords(maxNumRecords).build(); - } - - /** - * Define the max read time (duration) while the {@link Read} will receive messages. - * When this max read time is not null, the {@link Read} will provide a bounded - * {@link PCollection}. - */ - public Read withMaxReadTime(Duration maxReadTime) { - checkArgument(maxNumRecords() == Long.MAX_VALUE, - "maxNumRecord and maxReadTime are exclusive"); - return builder().setMaxReadTime(maxReadTime).build(); - } - - @Override - public void validate(PipelineOptions pipelineOptions) { - checkState(addresses() != null, "AmqIO.read() requires addresses list to be set via " - + "withAddresses(addresses)"); - checkState(!addresses().isEmpty(), "AmqIO.read() requires a non-empty addresses list to be" - + " set via withAddresses(addresses)"); - } - - @Override - public void populateDisplayData(DisplayData.Builder builder) { - builder.add(DisplayData.item("addresses", Joiner.on(" ").join(addresses()))); - } - - @Override - public PCollection<Message> expand(PBegin input) { - org.apache.beam.sdk.io.Read.Unbounded<Message> unbounded = - org.apache.beam.sdk.io.Read.from(new UnboundedAmqpSource(this)); - - PTransform<PBegin, PCollection<Message>> transform = unbounded; - - if (maxNumRecords() != Long.MAX_VALUE) { - transform = unbounded.withMaxNumRecords(maxNumRecords()); - } else if (maxReadTime() != null) { - transform = unbounded.withMaxReadTime(maxReadTime()); - } - - return input.getPipeline().apply(transform); - } - - } - - private static class AmqpCheckpointMark implements UnboundedSource.CheckpointMark, Serializable { - - private transient Messenger messenger; - private transient List<Tracker> trackers = new ArrayList<>(); - - public AmqpCheckpointMark() { - } - - @Override - public void finalizeCheckpoint() { - for (Tracker tracker : trackers) { - // flag as not cumulative - messenger.accept(tracker, 0); - } - trackers.clear(); - } - - // set an empty list to messages when deserialize - private void readObject(java.io.ObjectInputStream stream) - throws java.io.IOException, ClassNotFoundException { - trackers = new ArrayList<>(); - } - - } - - private static class UnboundedAmqpSource - extends UnboundedSource<Message, AmqpCheckpointMark> { - - private final Read spec; - - public UnboundedAmqpSource(Read spec) { - this.spec = spec; - } - - @Override - public List<UnboundedAmqpSource> split(int desiredNumSplits, - PipelineOptions pipelineOptions) { - // amqp is a queue system, so, it's possible to have multiple concurrent sources, even if - // they bind the listener - List<UnboundedAmqpSource> sources = new ArrayList<>(); - for (int i = 0; i < Math.max(1, desiredNumSplits); ++i) { - sources.add(new UnboundedAmqpSource(spec)); - } - return sources; - } - - @Override - public UnboundedReader<Message> createReader(PipelineOptions pipelineOptions, - AmqpCheckpointMark checkpointMark) { - return new UnboundedAmqpReader(this, checkpointMark); - } - - @Override - public Coder<Message> getDefaultOutputCoder() { - return new AmqpMessageCoder(); - } - - @Override - public Coder<AmqpCheckpointMark> getCheckpointMarkCoder() { - return SerializableCoder.of(AmqpCheckpointMark.class); - } - - @Override - public void validate() { - spec.validate(null); - } - - } - - private static class UnboundedAmqpReader extends UnboundedSource.UnboundedReader<Message> { - - private final UnboundedAmqpSource source; - - private Messenger messenger; - private Message current; - private Instant currentTimestamp; - private Instant watermark = new Instant(Long.MIN_VALUE); - private AmqpCheckpointMark checkpointMark; - - public UnboundedAmqpReader(UnboundedAmqpSource source, AmqpCheckpointMark checkpointMark) { - this.source = source; - this.current = null; - if (checkpointMark != null) { - this.checkpointMark = checkpointMark; - } else { - this.checkpointMark = new AmqpCheckpointMark(); - } - } - - @Override - public Instant getWatermark() { - return watermark; - } - - @Override - public Instant getCurrentTimestamp() { - if (current == null) { - throw new NoSuchElementException(); - } - return currentTimestamp; - } - - @Override - public Message getCurrent() { - if (current == null) { - throw new NoSuchElementException(); - } - return current; - } - - @Override - public UnboundedSource.CheckpointMark getCheckpointMark() { - return checkpointMark; - } - - @Override - public UnboundedAmqpSource getCurrentSource() { - return source; - } - - @Override - public boolean start() throws IOException { - Read spec = source.spec; - messenger = Messenger.Factory.create(); - messenger.start(); - for (String address : spec.addresses()) { - messenger.subscribe(address); - } - checkpointMark.messenger = messenger; - return advance(); - } - - @Override - public boolean advance() { - messenger.recv(); - if (messenger.incoming() <= 0) { - current = null; - return false; - } - Message message = messenger.get(); - Tracker tracker = messenger.incomingTracker(); - checkpointMark.trackers.add(tracker); - currentTimestamp = new Instant(message.getCreationTime()); - watermark = currentTimestamp; - current = message; - return true; - } - - @Override - public void close() { - if (messenger != null) { - messenger.stop(); - } - } - - } - - /** - * A {@link PTransform} to send messages using AMQP 1.0 protocol. - */ - @AutoValue - public abstract static class Write extends PTransform<PCollection<Message>, PDone> { - - @Override - public PDone expand(PCollection<Message> input) { - input.apply(ParDo.of(new WriteFn(this))); - return PDone.in(input.getPipeline()); - } - - private static class WriteFn extends DoFn<Message, Void> { - - private final Write spec; - - private transient Messenger messenger; - - public WriteFn(Write spec) { - this.spec = spec; - } - - @Setup - public void setup() throws Exception { - messenger = Messenger.Factory.create(); - messenger.start(); - } - - @ProcessElement - public void processElement(ProcessContext processContext) throws Exception { - Message message = processContext.element(); - messenger.put(message); - messenger.send(); - } - - @Teardown - public void teardown() throws Exception { - if (messenger != null) { - messenger.stop(); - } - } - - } - - } - -} http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/amqp/src/main/java/org/apache/beam/sdk/io/amqp/AmqpMessageCoder.java ---------------------------------------------------------------------- diff --git a/sdks/java/io/amqp/src/main/java/org/apache/beam/sdk/io/amqp/AmqpMessageCoder.java b/sdks/java/io/amqp/src/main/java/org/apache/beam/sdk/io/amqp/AmqpMessageCoder.java deleted file mode 100644 index 5a55260..0000000 --- a/sdks/java/io/amqp/src/main/java/org/apache/beam/sdk/io/amqp/AmqpMessageCoder.java +++ /dev/null @@ -1,79 +0,0 @@ -/* - * Licensed to the Apache Software Foundation (ASF) under one - * or more contributor license agreements. See the NOTICE file - * distributed with this work for additional information - * regarding copyright ownership. The ASF licenses this file - * to you 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 org.apache.beam.sdk.io.amqp; - -import com.google.common.io.ByteStreams; - -import java.io.IOException; -import java.io.InputStream; -import java.io.OutputStream; -import java.nio.BufferOverflowException; - -import org.apache.beam.sdk.coders.CoderException; -import org.apache.beam.sdk.coders.CustomCoder; -import org.apache.beam.sdk.util.VarInt; -import org.apache.qpid.proton.message.Message; - -/** - * A coder for AMQP message. - */ -public class AmqpMessageCoder extends CustomCoder<Message> { - - private static final int[] MESSAGE_SIZES = new int[]{ - 8 * 1024, - 64 * 1024, - 1 * 1024 * 1024, - 64 * 1024 * 1024 - }; - - static AmqpMessageCoder of() { - return new AmqpMessageCoder(); - } - - @Override - public void encode(Message value, OutputStream outStream) throws CoderException, IOException { - for (int maxMessageSize : MESSAGE_SIZES) { - try { - encode(value, outStream, maxMessageSize); - return; - } catch (Exception e) { - continue; - } - } - throw new CoderException("Message is larger than the max size supported by the coder"); - } - - private void encode(Message value, OutputStream outStream, int messageSize) throws - IOException, BufferOverflowException { - byte[] data = new byte[messageSize]; - int bytesWritten = value.encode(data, 0, data.length); - VarInt.encode(bytesWritten, outStream); - outStream.write(data, 0, bytesWritten); - } - - @Override - public Message decode(InputStream inStream) throws CoderException, IOException { - Message message = Message.Factory.create(); - int bytesToRead = VarInt.decodeInt(inStream); - byte[] encodedMessage = new byte[bytesToRead]; - ByteStreams.readFully(inStream, encodedMessage); - message.decode(encodedMessage, 0, encodedMessage.length); - return message; - } - -} http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/amqp/src/main/java/org/apache/beam/sdk/io/amqp/AmqpMessageCoderProviderRegistrar.java ---------------------------------------------------------------------- diff --git a/sdks/java/io/amqp/src/main/java/org/apache/beam/sdk/io/amqp/AmqpMessageCoderProviderRegistrar.java b/sdks/java/io/amqp/src/main/java/org/apache/beam/sdk/io/amqp/AmqpMessageCoderProviderRegistrar.java deleted file mode 100644 index bc3445c..0000000 --- a/sdks/java/io/amqp/src/main/java/org/apache/beam/sdk/io/amqp/AmqpMessageCoderProviderRegistrar.java +++ /dev/null @@ -1,44 +0,0 @@ -/* - * Licensed to the Apache Software Foundation (ASF) under one - * or more contributor license agreements. See the NOTICE file - * distributed with this work for additional information - * regarding copyright ownership. The ASF licenses this file - * to you 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 org.apache.beam.sdk.io.amqp; - -import com.google.auto.service.AutoService; -import com.google.common.collect.ImmutableList; - -import java.util.List; - -import org.apache.beam.sdk.coders.CoderProvider; -import org.apache.beam.sdk.coders.CoderProviderRegistrar; -import org.apache.beam.sdk.coders.CoderProviders; -import org.apache.beam.sdk.values.TypeDescriptor; -import org.apache.qpid.proton.message.Message; - -/** - * A {@link CoderProviderRegistrar} for standard types used with {@link AmqpIO}. - */ -@AutoService(CoderProviderRegistrar.class) -public class AmqpMessageCoderProviderRegistrar implements CoderProviderRegistrar { - - @Override - public List<CoderProvider> getCoderProviders() { - return ImmutableList.of( - CoderProviders.forCoder(TypeDescriptor.of(Message.class), - AmqpMessageCoder.of())); - } - -} http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/amqp/src/main/java/org/apache/beam/sdk/io/amqp/package-info.java ---------------------------------------------------------------------- diff --git a/sdks/java/io/amqp/src/main/java/org/apache/beam/sdk/io/amqp/package-info.java b/sdks/java/io/amqp/src/main/java/org/apache/beam/sdk/io/amqp/package-info.java deleted file mode 100644 index 091f234..0000000 --- a/sdks/java/io/amqp/src/main/java/org/apache/beam/sdk/io/amqp/package-info.java +++ /dev/null @@ -1,22 +0,0 @@ -/* - * Licensed to the Apache Software Foundation (ASF) under one - * or more contributor license agreements. See the NOTICE file - * distributed with this work for additional information - * regarding copyright ownership. The ASF licenses this file - * to you 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. - */ - -/** - * Transforms for reading and writing using AMQP 1.0 protocol. - */ -package org.apache.beam.sdk.io.amqp; http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/amqp/src/test/java/org/apache/beam/sdk/io/amqp/AmqpIOTest.java ---------------------------------------------------------------------- diff --git a/sdks/java/io/amqp/src/test/java/org/apache/beam/sdk/io/amqp/AmqpIOTest.java b/sdks/java/io/amqp/src/test/java/org/apache/beam/sdk/io/amqp/AmqpIOTest.java deleted file mode 100644 index c8fe4e8..0000000 --- a/sdks/java/io/amqp/src/test/java/org/apache/beam/sdk/io/amqp/AmqpIOTest.java +++ /dev/null @@ -1,148 +0,0 @@ -/* - * Licensed to the Apache Software Foundation (ASF) under one - * or more contributor license agreements. See the NOTICE file - * distributed with this work for additional information - * regarding copyright ownership. The ASF licenses this file - * to you 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 org.apache.beam.sdk.io.amqp; - -import static org.junit.Assert.assertEquals; -import static org.junit.Assert.assertTrue; - -import java.net.ServerSocket; -import java.util.ArrayList; -import java.util.Collections; -import java.util.List; - -import org.apache.beam.sdk.testing.PAssert; -import org.apache.beam.sdk.testing.TestPipeline; -import org.apache.beam.sdk.transforms.Count; -import org.apache.beam.sdk.transforms.Create; -import org.apache.beam.sdk.values.PCollection; -import org.apache.qpid.proton.amqp.messaging.AmqpValue; -import org.apache.qpid.proton.message.Message; -import org.apache.qpid.proton.messenger.Messenger; -import org.junit.Before; -import org.junit.Rule; -import org.junit.Test; -import org.junit.runner.RunWith; -import org.junit.runners.JUnit4; -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; - -/** - * Tests on {@link AmqpIO}. - */ -@RunWith(JUnit4.class) -public class AmqpIOTest { - - private static final Logger LOG = LoggerFactory.getLogger(AmqpIOTest.class); - - private int port; - - @Rule public TestPipeline pipeline = TestPipeline.create(); - - @Before - public void findFreeNetworkPort() throws Exception { - LOG.info("Finding free network port"); - ServerSocket socket = new ServerSocket(0); - port = socket.getLocalPort(); - socket.close(); - } - - @Test - public void testRead() throws Exception { - PCollection<Message> output = pipeline.apply(AmqpIO.read() - .withMaxNumRecords(100) - .withAddresses(Collections.singletonList("amqp://~localhost:" + port))); - PAssert.thatSingleton(output.apply(Count.<Message>globally())).isEqualTo(100L); - - Thread sender = new Thread() { - public void run() { - try { - Thread.sleep(500); - Messenger sender = Messenger.Factory.create(); - sender.start(); - for (int i = 0; i < 100; i++) { - Message message = Message.Factory.create(); - message.setAddress("amqp://localhost:" + port); - message.setBody(new AmqpValue("Test " + i)); - sender.put(message); - sender.send(); - } - sender.stop(); - } catch (Exception e) { - LOG.error("Sender error", e); - } - } - }; - try { - sender.start(); - pipeline.run(); - } finally { - sender.join(); - } - } - - @Test - public void testWrite() throws Exception { - final List<String> received = new ArrayList<>(); - Thread receiver = new Thread() { - @Override - public void run() { - try { - Messenger messenger = Messenger.Factory.create(); - messenger.start(); - messenger.subscribe("amqp://~localhost:" + port); - while (received.size() < 100) { - messenger.recv(); - while (messenger.incoming() > 0) { - Message message = messenger.get(); - LOG.info("Received: " + message.getBody().toString()); - received.add(message.getBody().toString()); - } - } - messenger.stop(); - } catch (Exception e) { - LOG.error("Receiver error", e); - } - } - }; - LOG.info("Starting AMQP receiver"); - receiver.start(); - - List<Message> data = new ArrayList<>(); - for (int i = 0; i < 100; i++) { - Message message = Message.Factory.create(); - message.setBody(new AmqpValue("Test " + i)); - message.setAddress("amqp://localhost:" + port); - message.setSubject("test"); - data.add(message); - } - pipeline.apply(Create.of(data).withCoder(AmqpMessageCoder.of())).apply(AmqpIO.write()); - LOG.info("Starting pipeline"); - try { - pipeline.run(); - } finally { - LOG.info("Join receiver thread"); - receiver.join(); - } - - assertEquals(100, received.size()); - for (int i = 0; i < 100; i++) { - assertTrue(received.contains("AmqpValue{Test " + i + "}")); - } - } - -} http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/amqp/src/test/java/org/apache/beam/sdk/io/amqp/AmqpMessageCoderTest.java ---------------------------------------------------------------------- diff --git a/sdks/java/io/amqp/src/test/java/org/apache/beam/sdk/io/amqp/AmqpMessageCoderTest.java b/sdks/java/io/amqp/src/test/java/org/apache/beam/sdk/io/amqp/AmqpMessageCoderTest.java deleted file mode 100644 index 7a8efeb..0000000 --- a/sdks/java/io/amqp/src/test/java/org/apache/beam/sdk/io/amqp/AmqpMessageCoderTest.java +++ /dev/null @@ -1,89 +0,0 @@ -/* - * Licensed to the Apache Software Foundation (ASF) under one - * or more contributor license agreements. See the NOTICE file - * distributed with this work for additional information - * regarding copyright ownership. The ASF licenses this file - * to you 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 org.apache.beam.sdk.io.amqp; - -import static org.junit.Assert.assertEquals; - -import com.google.common.base.Joiner; - -import java.util.Collections; - -import org.apache.beam.sdk.coders.CoderException; -import org.apache.beam.sdk.util.CoderUtils; -import org.apache.qpid.proton.amqp.messaging.AmqpValue; -import org.apache.qpid.proton.message.Message; -import org.junit.Rule; -import org.junit.Test; -import org.junit.rules.ExpectedException; -import org.junit.runner.RunWith; -import org.junit.runners.JUnit4; - -/** - * Test on {@link AmqpMessageCoder}. - */ -@RunWith(JUnit4.class) -public class AmqpMessageCoderTest { - - @Rule - public ExpectedException thrown = ExpectedException.none(); - - @Test - public void encodeDecode() throws Exception { - Message message = Message.Factory.create(); - message.setBody(new AmqpValue("body")); - message.setAddress("address"); - message.setSubject("test"); - AmqpMessageCoder coder = AmqpMessageCoder.of(); - - Message clone = CoderUtils.clone(coder, message); - - assertEquals("AmqpValue{body}", clone.getBody().toString()); - assertEquals("address", clone.getAddress()); - assertEquals("test", clone.getSubject()); - } - - @Test - public void encodeDecodeTooMuchLargerMessage() throws Exception { - thrown.expect(CoderException.class); - Message message = Message.Factory.create(); - message.setAddress("address"); - message.setSubject("subject"); - String body = Joiner.on("").join(Collections.nCopies(64 * 1024 * 1024, " ")); - message.setBody(new AmqpValue(body)); - - AmqpMessageCoder coder = AmqpMessageCoder.of(); - - byte[] encoded = CoderUtils.encodeToByteArray(coder, message); - } - - @Test - public void encodeDecodeLargeMessage() throws Exception { - Message message = Message.Factory.create(); - message.setAddress("address"); - message.setSubject("subject"); - String body = Joiner.on("").join(Collections.nCopies(32 * 1024 * 1024, " ")); - message.setBody(new AmqpValue(body)); - - AmqpMessageCoder coder = AmqpMessageCoder.of(); - - Message clone = CoderUtils.clone(coder, message); - - clone.getBody().toString().equals(message.getBody().toString()); - } - -} http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/cassandra/pom.xml ---------------------------------------------------------------------- diff --git a/sdks/java/io/cassandra/pom.xml b/sdks/java/io/cassandra/pom.xml index c74477e..8249f57 100644 --- a/sdks/java/io/cassandra/pom.xml +++ b/sdks/java/io/cassandra/pom.xml @@ -22,7 +22,7 @@ <parent> <groupId>org.apache.beam</groupId> <artifactId>beam-sdks-java-io-parent</artifactId> - <version>2.2.0-SNAPSHOT</version> + <version>2.1.0-SNAPSHOT</version> <relativePath>../pom.xml</relativePath> </parent> http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/cassandra/src/main/java/org/apache/beam/sdk/io/cassandra/CassandraIO.java ---------------------------------------------------------------------- diff --git a/sdks/java/io/cassandra/src/main/java/org/apache/beam/sdk/io/cassandra/CassandraIO.java b/sdks/java/io/cassandra/src/main/java/org/apache/beam/sdk/io/cassandra/CassandraIO.java index 32905b7..b6f4ef6 100644 --- a/sdks/java/io/cassandra/src/main/java/org/apache/beam/sdk/io/cassandra/CassandraIO.java +++ b/sdks/java/io/cassandra/src/main/java/org/apache/beam/sdk/io/cassandra/CassandraIO.java @@ -82,7 +82,7 @@ import org.slf4j.LoggerFactory; * .withEntity(Person.class)); * }</pre> */ -@Experimental(Experimental.Kind.SOURCE_SINK) +@Experimental public class CassandraIO { private static final Logger LOG = LoggerFactory.getLogger(CassandraIO.class); http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/common/pom.xml ---------------------------------------------------------------------- diff --git a/sdks/java/io/common/pom.xml b/sdks/java/io/common/pom.xml index df0d94b..f7525fd 100644 --- a/sdks/java/io/common/pom.xml +++ b/sdks/java/io/common/pom.xml @@ -22,7 +22,7 @@ <parent> <groupId>org.apache.beam</groupId> <artifactId>beam-sdks-java-io-parent</artifactId> - <version>2.2.0-SNAPSHOT</version> + <version>2.1.0-SNAPSHOT</version> <relativePath>../pom.xml</relativePath> </parent> http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/common/src/test/java/org/apache/beam/sdk/io/common/IOTestPipelineOptions.java ---------------------------------------------------------------------- diff --git a/sdks/java/io/common/src/test/java/org/apache/beam/sdk/io/common/IOTestPipelineOptions.java b/sdks/java/io/common/src/test/java/org/apache/beam/sdk/io/common/IOTestPipelineOptions.java index 25ab929..387fd22 100644 --- a/sdks/java/io/common/src/test/java/org/apache/beam/sdk/io/common/IOTestPipelineOptions.java +++ b/sdks/java/io/common/src/test/java/org/apache/beam/sdk/io/common/IOTestPipelineOptions.java @@ -71,7 +71,11 @@ public interface IOTestPipelineOptions extends TestPipelineOptions { Integer getElasticsearchHttpPort(); void setElasticsearchHttpPort(Integer value); - /* Cassandra */ + @Description("Tcp port for elasticsearch server") + @Default.Integer(9300) + Integer getElasticsearchTcpPort(); + void setElasticsearchTcpPort(Integer value); + @Description("Host for Cassandra server (host name/ip address)") @Default.String("cassandra-host") String getCassandraHost(); http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/elasticsearch/pom.xml ---------------------------------------------------------------------- diff --git a/sdks/java/io/elasticsearch/pom.xml b/sdks/java/io/elasticsearch/pom.xml index e0a7f21..03632ce 100644 --- a/sdks/java/io/elasticsearch/pom.xml +++ b/sdks/java/io/elasticsearch/pom.xml @@ -22,7 +22,7 @@ <parent> <groupId>org.apache.beam</groupId> <artifactId>beam-sdks-java-io-parent</artifactId> - <version>2.2.0-SNAPSHOT</version> + <version>2.1.0-SNAPSHOT</version> <relativePath>../pom.xml</relativePath> </parent> @@ -137,14 +137,6 @@ <scope>test</scope> </dependency> - <!-- This optional dependency is used by the test framework. Avoids a warning --> - <dependency> - <groupId>net.java.dev.jna</groupId> - <artifactId>jna</artifactId> - <version>4.1.0</version> - <scope>test</scope> - </dependency> - <dependency> <groupId>org.apache.beam</groupId> <artifactId>beam-runners-direct-java</artifactId> http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/elasticsearch/src/main/java/org/apache/beam/sdk/io/elasticsearch/ElasticsearchIO.java ---------------------------------------------------------------------- diff --git a/sdks/java/io/elasticsearch/src/main/java/org/apache/beam/sdk/io/elasticsearch/ElasticsearchIO.java b/sdks/java/io/elasticsearch/src/main/java/org/apache/beam/sdk/io/elasticsearch/ElasticsearchIO.java index 4d76887..f6ceef2 100644 --- a/sdks/java/io/elasticsearch/src/main/java/org/apache/beam/sdk/io/elasticsearch/ElasticsearchIO.java +++ b/sdks/java/io/elasticsearch/src/main/java/org/apache/beam/sdk/io/elasticsearch/ElasticsearchIO.java @@ -113,7 +113,7 @@ import org.elasticsearch.client.RestClientBuilder; * <p>Optionally, you can provide {@code withBatchSize()} and {@code withBatchSizeBytes()} * to specify the size of the write batch in number of documents or in bytes. */ -@Experimental(Experimental.Kind.SOURCE_SINK) +@Experimental public class ElasticsearchIO { public static Read read() { @@ -139,7 +139,7 @@ public class ElasticsearchIO { private static final ObjectMapper mapper = new ObjectMapper(); - static JsonNode parseResponse(Response response) throws IOException { + private static JsonNode parseResponse(Response response) throws IOException { return mapper.readValue(response.getEntity().getContent(), JsonNode.class); } @@ -264,7 +264,7 @@ public class ElasticsearchIO { builder.addIfNotNull(DisplayData.item("username", getUsername())); } - RestClient createClient() throws MalformedURLException { + private RestClient createClient() throws MalformedURLException { HttpHost[] hosts = new HttpHost[getAddresses().size()]; int i = 0; for (String address : getAddresses()) { @@ -455,7 +455,16 @@ public class ElasticsearchIO { while (shards.hasNext()) { Map.Entry<String, JsonNode> shardJson = shards.next(); String shardId = shardJson.getKey(); - sources.add(new BoundedElasticsearchSource(spec, shardId)); + JsonNode value = (JsonNode) shardJson.getValue(); + boolean isPrimaryShard = + value + .path(0) + .path("routing") + .path("primary") + .asBoolean(); + if (isPrimaryShard) { + sources.add(new BoundedElasticsearchSource(spec, shardId)); + } } checkArgument(!sources.isEmpty(), "No primary shard found"); return sources; http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/elasticsearch/src/test/java/org/apache/beam/sdk/io/elasticsearch/ElasticSearchIOTestUtils.java ---------------------------------------------------------------------- diff --git a/sdks/java/io/elasticsearch/src/test/java/org/apache/beam/sdk/io/elasticsearch/ElasticSearchIOTestUtils.java b/sdks/java/io/elasticsearch/src/test/java/org/apache/beam/sdk/io/elasticsearch/ElasticSearchIOTestUtils.java index 203963d..b0d161f 100644 --- a/sdks/java/io/elasticsearch/src/test/java/org/apache/beam/sdk/io/elasticsearch/ElasticSearchIOTestUtils.java +++ b/sdks/java/io/elasticsearch/src/test/java/org/apache/beam/sdk/io/elasticsearch/ElasticSearchIOTestUtils.java @@ -17,17 +17,19 @@ */ package org.apache.beam.sdk.io.elasticsearch; -import com.fasterxml.jackson.databind.JsonNode; import java.io.IOException; import java.util.ArrayList; -import java.util.Collections; import java.util.List; -import org.apache.http.HttpEntity; -import org.apache.http.entity.ContentType; -import org.apache.http.message.BasicHeader; -import org.apache.http.nio.entity.NStringEntity; -import org.elasticsearch.client.Response; -import org.elasticsearch.client.RestClient; +import org.elasticsearch.action.admin.indices.exists.indices.IndicesExistsRequest; +import org.elasticsearch.action.admin.indices.exists.indices.IndicesExistsResponse; +import org.elasticsearch.action.admin.indices.upgrade.post.UpgradeRequest; +import org.elasticsearch.action.bulk.BulkRequestBuilder; +import org.elasticsearch.action.bulk.BulkResponse; +import org.elasticsearch.action.search.SearchResponse; +import org.elasticsearch.client.Client; +import org.elasticsearch.client.IndicesAdminClient; +import org.elasticsearch.client.Requests; +import org.elasticsearch.index.IndexNotFoundException; /** Test utilities to use with {@link ElasticsearchIO}. */ class ElasticSearchIOTestUtils { @@ -39,68 +41,57 @@ class ElasticSearchIOTestUtils { } /** Deletes the given index synchronously. */ - static void deleteIndex(String index, RestClient restClient) throws IOException { - try { - restClient.performRequest("DELETE", String.format("/%s", index), new BasicHeader("", "")); - } catch (IOException e) { - // it is fine to ignore this expression as deleteIndex occurs in @before, - // so when the first tests is run, the index does not exist yet - if (!e.getMessage().contains("index_not_found_exception")){ - throw e; - } + static void deleteIndex(String index, Client client) throws Exception { + IndicesAdminClient indices = client.admin().indices(); + IndicesExistsResponse indicesExistsResponse = + indices.exists(new IndicesExistsRequest(index)).get(); + if (indicesExistsResponse.isExists()) { + indices.prepareClose(index).get(); + indices.delete(Requests.deleteIndexRequest(index)).get(); } } /** Inserts the given number of test documents into Elasticsearch. */ - static void insertTestDocuments(String index, String type, long numDocs, RestClient restClient) - throws IOException { + static void insertTestDocuments(String index, String type, long numDocs, Client client) + throws Exception { + final BulkRequestBuilder bulkRequestBuilder = client.prepareBulk().setRefresh(true); List<String> data = ElasticSearchIOTestUtils.createDocuments( numDocs, ElasticSearchIOTestUtils.InjectionMode.DO_NOT_INJECT_INVALID_DOCS); - StringBuilder bulkRequest = new StringBuilder(); for (String document : data) { - bulkRequest.append(String.format("{ \"index\" : {} }%n%s%n", document)); + bulkRequestBuilder.add(client.prepareIndex(index, type, null).setSource(document)); } - String endPoint = String.format("/%s/%s/_bulk", index, type); - HttpEntity requestBody = - new NStringEntity(bulkRequest.toString(), ContentType.APPLICATION_JSON); - Response response = restClient.performRequest("POST", endPoint, - Collections.singletonMap("refresh", "true"), requestBody, - new BasicHeader("", "")); - JsonNode searchResult = ElasticsearchIO.parseResponse(response); - boolean errors = searchResult.path("errors").asBoolean(); - if (errors){ + final BulkResponse bulkResponse = bulkRequestBuilder.execute().actionGet(); + if (bulkResponse.hasFailures()) { throw new IOException( - String.format("Failed to insert test documents in index %s", index)); + String.format( + "Cannot insert test documents in index %s : %s", + index, bulkResponse.buildFailureMessage())); } } /** - * Forces a refresh of the given index to make recently inserted documents available for search. + * Forces an upgrade of the given index to make recently inserted documents available for search. * * @return The number of docs in the index */ - static long refreshIndexAndGetCurrentNumDocs(String index, String type, RestClient restClient) - throws IOException { - long result = 0; + static long upgradeIndexAndGetCurrentNumDocs(String index, String type, Client client) { try { - String endPoint = String.format("/%s/_refresh", index); - restClient.performRequest("POST", endPoint, new BasicHeader("", "")); - - endPoint = String.format("/%s/%s/_search", index, type); - Response response = restClient.performRequest("GET", endPoint, new BasicHeader("", "")); - JsonNode searchResult = ElasticsearchIO.parseResponse(response); - result = searchResult.path("hits").path("total").asLong(); - } catch (IOException e) { + client.admin().indices().upgrade(new UpgradeRequest(index)).actionGet(); + SearchResponse response = + client.prepareSearch(index).setTypes(type).execute().actionGet(5000); + return response.getHits().getTotalHits(); // it is fine to ignore bellow exceptions because in testWriteWithBatchSize* sometimes, // we call upgrade before any doc have been written // (when there are fewer docs processed than batchSize). // In that cases index/type has not been created (created upon first doc insertion) - if (!e.getMessage().contains("index_not_found_exception")){ + } catch (IndexNotFoundException e) { + } catch (java.lang.IllegalArgumentException e) { + if (!e.getMessage().contains("No search type")) { throw e; } } - return result; + return 0; } /** http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/elasticsearch/src/test/java/org/apache/beam/sdk/io/elasticsearch/ElasticsearchIOIT.java ---------------------------------------------------------------------- diff --git a/sdks/java/io/elasticsearch/src/test/java/org/apache/beam/sdk/io/elasticsearch/ElasticsearchIOIT.java b/sdks/java/io/elasticsearch/src/test/java/org/apache/beam/sdk/io/elasticsearch/ElasticsearchIOIT.java index 7c37e87..2d6393a 100644 --- a/sdks/java/io/elasticsearch/src/test/java/org/apache/beam/sdk/io/elasticsearch/ElasticsearchIOIT.java +++ b/sdks/java/io/elasticsearch/src/test/java/org/apache/beam/sdk/io/elasticsearch/ElasticsearchIOIT.java @@ -32,7 +32,7 @@ import org.apache.beam.sdk.testing.TestPipeline; import org.apache.beam.sdk.transforms.Count; import org.apache.beam.sdk.transforms.Create; import org.apache.beam.sdk.values.PCollection; -import org.elasticsearch.client.RestClient; +import org.elasticsearch.client.transport.TransportClient; import org.junit.AfterClass; import org.junit.BeforeClass; import org.junit.Rule; @@ -57,7 +57,7 @@ import org.slf4j.LoggerFactory; */ public class ElasticsearchIOIT { private static final Logger LOG = LoggerFactory.getLogger(ElasticsearchIOIT.class); - private static RestClient restClient; + private static TransportClient client; private static IOTestPipelineOptions options; private static ElasticsearchIO.ConnectionConfiguration readConnectionConfiguration; @Rule public TestPipeline pipeline = TestPipeline.create(); @@ -66,16 +66,16 @@ public class ElasticsearchIOIT { public static void beforeClass() throws Exception { PipelineOptionsFactory.register(IOTestPipelineOptions.class); options = TestPipeline.testingPipelineOptions().as(IOTestPipelineOptions.class); + client = ElasticsearchTestDataSet.getClient(options); readConnectionConfiguration = ElasticsearchTestDataSet.getConnectionConfiguration( options, ElasticsearchTestDataSet.ReadOrWrite.READ); - restClient = readConnectionConfiguration.createClient(); } @AfterClass public static void afterClass() throws Exception { - ElasticsearchTestDataSet.deleteIndex(restClient, ElasticsearchTestDataSet.ReadOrWrite.WRITE); - restClient.close(); + ElasticsearchTestDataSet.deleteIndex(client, ElasticsearchTestDataSet.ReadOrWrite.WRITE); + client.close(); } @Test @@ -128,8 +128,8 @@ public class ElasticsearchIOIT { pipeline.run(); long currentNumDocs = - ElasticSearchIOTestUtils.refreshIndexAndGetCurrentNumDocs( - ElasticsearchTestDataSet.ES_INDEX, ElasticsearchTestDataSet.ES_TYPE, restClient); + ElasticSearchIOTestUtils.upgradeIndexAndGetCurrentNumDocs( + ElasticsearchTestDataSet.ES_INDEX, ElasticsearchTestDataSet.ES_TYPE, client); assertEquals(ElasticsearchTestDataSet.NUM_DOCS, currentNumDocs); } http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/elasticsearch/src/test/java/org/apache/beam/sdk/io/elasticsearch/ElasticsearchIOTest.java ---------------------------------------------------------------------- diff --git a/sdks/java/io/elasticsearch/src/test/java/org/apache/beam/sdk/io/elasticsearch/ElasticsearchIOTest.java b/sdks/java/io/elasticsearch/src/test/java/org/apache/beam/sdk/io/elasticsearch/ElasticsearchIOTest.java index b349a29..260af79 100644 --- a/sdks/java/io/elasticsearch/src/test/java/org/apache/beam/sdk/io/elasticsearch/ElasticsearchIOTest.java +++ b/sdks/java/io/elasticsearch/src/test/java/org/apache/beam/sdk/io/elasticsearch/ElasticsearchIOTest.java @@ -39,11 +39,11 @@ import org.apache.beam.sdk.transforms.Create; import org.apache.beam.sdk.transforms.DoFnTester; import org.apache.beam.sdk.values.PCollection; import org.elasticsearch.action.search.SearchResponse; -import org.elasticsearch.client.RestClient; import org.elasticsearch.common.settings.Settings; import org.elasticsearch.index.query.QueryBuilder; import org.elasticsearch.index.query.QueryBuilders; import org.elasticsearch.node.Node; +import org.elasticsearch.node.NodeBuilder; import org.hamcrest.CustomMatcher; import org.junit.AfterClass; import org.junit.Before; @@ -74,10 +74,9 @@ public class ElasticsearchIOTest implements Serializable { private static final long BATCH_SIZE_BYTES = 2048L; private static Node node; - private static RestClient restClient; private static ElasticsearchIO.ConnectionConfiguration connectionConfiguration; - @ClassRule public static final TemporaryFolder TEMPORARY_FOLDER = new TemporaryFolder(); + @ClassRule public static TemporaryFolder folder = new TemporaryFolder(); @Rule public TestPipeline pipeline = TestPipeline.create(); @@ -92,8 +91,8 @@ public class ElasticsearchIOTest implements Serializable { .put("cluster.name", "beam") .put("http.enabled", "true") .put("node.data", "true") - .put("path.data", TEMPORARY_FOLDER.getRoot().getPath()) - .put("path.home", TEMPORARY_FOLDER.getRoot().getPath()) + .put("path.data", folder.getRoot().getPath()) + .put("path.home", folder.getRoot().getPath()) .put("node.name", "beam") .put("network.host", ES_IP) .put("http.port", esHttpPort) @@ -101,29 +100,27 @@ public class ElasticsearchIOTest implements Serializable { // had problems with some jdk, embedded ES was too slow for bulk insertion, // and queue of 50 was full. No pb with real ES instance (cf testWrite integration test) .put("threadpool.bulk.queue_size", 100); - node = new Node(settingsBuilder.build()); + node = NodeBuilder.nodeBuilder().settings(settingsBuilder).build(); LOG.info("Elasticsearch node created"); node.start(); connectionConfiguration = ElasticsearchIO.ConnectionConfiguration.create( new String[] {"http://"; + ES_IP + ":" + esHttpPort}, ES_INDEX, ES_TYPE); - restClient = connectionConfiguration.createClient(); } @AfterClass - public static void afterClass() throws IOException{ - restClient.close(); + public static void afterClass() { node.close(); } @Before public void before() throws Exception { - ElasticSearchIOTestUtils.deleteIndex(ES_INDEX, restClient); + ElasticSearchIOTestUtils.deleteIndex(ES_INDEX, node.client()); } @Test public void testSizes() throws Exception { - ElasticSearchIOTestUtils.insertTestDocuments(ES_INDEX, ES_TYPE, NUM_DOCS, restClient); + ElasticSearchIOTestUtils.insertTestDocuments(ES_INDEX, ES_TYPE, NUM_DOCS, node.client()); PipelineOptions options = PipelineOptionsFactory.create(); ElasticsearchIO.Read read = ElasticsearchIO.read().withConnectionConfiguration(connectionConfiguration); @@ -137,7 +134,7 @@ public class ElasticsearchIOTest implements Serializable { @Test public void testRead() throws Exception { - ElasticSearchIOTestUtils.insertTestDocuments(ES_INDEX, ES_TYPE, NUM_DOCS, restClient); + ElasticSearchIOTestUtils.insertTestDocuments(ES_INDEX, ES_TYPE, NUM_DOCS, node.client()); PCollection<String> output = pipeline.apply( @@ -153,7 +150,7 @@ public class ElasticsearchIOTest implements Serializable { @Test public void testReadWithQuery() throws Exception { - ElasticSearchIOTestUtils.insertTestDocuments(ES_INDEX, ES_TYPE, NUM_DOCS, restClient); + ElasticSearchIOTestUtils.insertTestDocuments(ES_INDEX, ES_TYPE, NUM_DOCS, node.client()); String query = "{\n" @@ -188,7 +185,7 @@ public class ElasticsearchIOTest implements Serializable { pipeline.run(); long currentNumDocs = - ElasticSearchIOTestUtils.refreshIndexAndGetCurrentNumDocs(ES_INDEX, ES_TYPE, restClient); + ElasticSearchIOTestUtils.upgradeIndexAndGetCurrentNumDocs(ES_INDEX, ES_TYPE, node.client()); assertEquals(NUM_DOCS, currentNumDocs); QueryBuilder queryBuilder = QueryBuilders.queryStringQuery("Einstein").field("scientist"); @@ -261,8 +258,9 @@ public class ElasticsearchIOTest implements Serializable { if ((numDocsProcessed % 100) == 0) { // force the index to upgrade after inserting for the inserted docs // to be searchable immediately - long currentNumDocs = ElasticSearchIOTestUtils - .refreshIndexAndGetCurrentNumDocs(ES_INDEX, ES_TYPE, restClient); + long currentNumDocs = + ElasticSearchIOTestUtils.upgradeIndexAndGetCurrentNumDocs( + ES_INDEX, ES_TYPE, node.client()); if ((numDocsProcessed % BATCH_SIZE) == 0) { /* bundle end */ assertEquals( @@ -306,8 +304,8 @@ public class ElasticsearchIOTest implements Serializable { // force the index to upgrade after inserting for the inserted docs // to be searchable immediately long currentNumDocs = - ElasticSearchIOTestUtils.refreshIndexAndGetCurrentNumDocs( - ES_INDEX, ES_TYPE, restClient); + ElasticSearchIOTestUtils.upgradeIndexAndGetCurrentNumDocs( + ES_INDEX, ES_TYPE, node.client()); if (sizeProcessed / BATCH_SIZE_BYTES > batchInserted) { /* bundle end */ assertThat( @@ -329,7 +327,7 @@ public class ElasticsearchIOTest implements Serializable { @Test public void testSplit() throws Exception { - ElasticSearchIOTestUtils.insertTestDocuments(ES_INDEX, ES_TYPE, NUM_DOCS, restClient); + ElasticSearchIOTestUtils.insertTestDocuments(ES_INDEX, ES_TYPE, NUM_DOCS, node.client()); PipelineOptions options = PipelineOptionsFactory.create(); ElasticsearchIO.Read read = ElasticsearchIO.read().withConnectionConfiguration(connectionConfiguration); http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/elasticsearch/src/test/java/org/apache/beam/sdk/io/elasticsearch/ElasticsearchTestDataSet.java ---------------------------------------------------------------------- diff --git a/sdks/java/io/elasticsearch/src/test/java/org/apache/beam/sdk/io/elasticsearch/ElasticsearchTestDataSet.java b/sdks/java/io/elasticsearch/src/test/java/org/apache/beam/sdk/io/elasticsearch/ElasticsearchTestDataSet.java index 2a2dbe9..3a9aae6 100644 --- a/sdks/java/io/elasticsearch/src/test/java/org/apache/beam/sdk/io/elasticsearch/ElasticsearchTestDataSet.java +++ b/sdks/java/io/elasticsearch/src/test/java/org/apache/beam/sdk/io/elasticsearch/ElasticsearchTestDataSet.java @@ -17,11 +17,13 @@ */ package org.apache.beam.sdk.io.elasticsearch; +import static java.net.InetAddress.getByName; import java.io.IOException; import org.apache.beam.sdk.io.common.IOTestPipelineOptions; import org.apache.beam.sdk.options.PipelineOptionsFactory; -import org.elasticsearch.client.RestClient; +import org.elasticsearch.client.transport.TransportClient; +import org.elasticsearch.common.transport.InetSocketTransportAddress; /** * Manipulates test data used by the {@link ElasticsearchIO} @@ -49,6 +51,7 @@ public class ElasticsearchTestDataSet { * -Dexec.mainClass=org.apache.beam.sdk.io.elasticsearch.ElasticsearchTestDataSet \ * -Dexec.args="--elasticsearchServer=1.2.3.4 \ * --elasticsearchHttpPort=9200 \ + * --elasticsearchTcpPort=9300" \ * -Dexec.classpathScope=test * </pre> * @@ -59,20 +62,29 @@ public class ElasticsearchTestDataSet { PipelineOptionsFactory.register(IOTestPipelineOptions.class); IOTestPipelineOptions options = PipelineOptionsFactory.fromArgs(args).as(IOTestPipelineOptions.class); - createAndPopulateReadIndex(options); + + createAndPopulateIndex(getClient(options), ReadOrWrite.READ); } - private static void createAndPopulateReadIndex(IOTestPipelineOptions options) throws Exception { - RestClient restClient = getConnectionConfiguration(options, ReadOrWrite.READ).createClient(); + private static void createAndPopulateIndex(TransportClient client, ReadOrWrite rOw) + throws Exception { // automatically creates the index and insert docs - try { - ElasticSearchIOTestUtils.insertTestDocuments(ES_INDEX, ES_TYPE, NUM_DOCS, restClient); - } finally { - restClient.close(); - } + ElasticSearchIOTestUtils.insertTestDocuments( + (rOw == ReadOrWrite.READ) ? ES_INDEX : writeIndex, ES_TYPE, NUM_DOCS, client); + } + + public static TransportClient getClient(IOTestPipelineOptions options) throws Exception { + TransportClient client = + TransportClient.builder() + .build() + .addTransportAddress( + new InetSocketTransportAddress( + getByName(options.getElasticsearchServer()), + options.getElasticsearchTcpPort())); + return client; } - static ElasticsearchIO.ConnectionConfiguration getConnectionConfiguration( + public static ElasticsearchIO.ConnectionConfiguration getConnectionConfiguration( IOTestPipelineOptions options, ReadOrWrite rOw) throws IOException { ElasticsearchIO.ConnectionConfiguration connectionConfiguration = ElasticsearchIO.ConnectionConfiguration.create( @@ -87,9 +99,8 @@ public class ElasticsearchTestDataSet { return connectionConfiguration; } - static void deleteIndex(RestClient restClient, ReadOrWrite rOw) throws Exception { - ElasticSearchIOTestUtils - .deleteIndex((rOw == ReadOrWrite.READ) ? ES_INDEX : writeIndex, restClient); + public static void deleteIndex(TransportClient client, ReadOrWrite rOw) throws Exception { + ElasticSearchIOTestUtils.deleteIndex((rOw == ReadOrWrite.READ) ? ES_INDEX : writeIndex, client); } /** Enum that tells whether we use the index for reading or for writing. */ http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/google-cloud-platform/pom.xml ---------------------------------------------------------------------- diff --git a/sdks/java/io/google-cloud-platform/pom.xml b/sdks/java/io/google-cloud-platform/pom.xml index a1495f2..8b53820 100644 --- a/sdks/java/io/google-cloud-platform/pom.xml +++ b/sdks/java/io/google-cloud-platform/pom.xml @@ -22,7 +22,7 @@ <parent> <groupId>org.apache.beam</groupId> <artifactId>beam-sdks-java-io-parent</artifactId> - <version>2.2.0-SNAPSHOT</version> + <version>2.1.0-SNAPSHOT</version> <relativePath>../pom.xml</relativePath> </parent> @@ -93,12 +93,7 @@ <dependency> <groupId>com.google.api</groupId> - <artifactId>gax-grpc</artifactId> - </dependency> - - <dependency> - <groupId>com.google.cloud</groupId> - <artifactId>google-cloud-core-grpc</artifactId> + <artifactId>api-common</artifactId> </dependency> <dependency> @@ -258,6 +253,11 @@ <artifactId>proto-google-common-protos</artifactId> </dependency> + <dependency> + <groupId>org.apache.commons</groupId> + <artifactId>commons-lang3</artifactId> + </dependency> + <!-- Test dependencies --> <dependency> <groupId>org.apache.beam</groupId> http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BatchLoads.java ---------------------------------------------------------------------- diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BatchLoads.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BatchLoads.java index e46b1d3..4393a63 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BatchLoads.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BatchLoads.java @@ -32,7 +32,6 @@ import org.apache.beam.sdk.coders.Coder; import org.apache.beam.sdk.coders.KvCoder; import org.apache.beam.sdk.coders.ListCoder; import org.apache.beam.sdk.coders.NullableCoder; -import org.apache.beam.sdk.coders.ShardedKeyCoder; import org.apache.beam.sdk.coders.StringUtf8Coder; import org.apache.beam.sdk.coders.VoidCoder; import org.apache.beam.sdk.io.gcp.bigquery.BigQueryIO.Write.CreateDisposition; @@ -58,7 +57,6 @@ import org.apache.beam.sdk.values.PCollection; import org.apache.beam.sdk.values.PCollectionList; import org.apache.beam.sdk.values.PCollectionTuple; import org.apache.beam.sdk.values.PCollectionView; -import org.apache.beam.sdk.values.ShardedKey; import org.apache.beam.sdk.values.TupleTag; import org.apache.beam.sdk.values.TupleTagList; import org.apache.beam.sdk.values.TypeDescriptor; http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/DynamicDestinations.java ---------------------------------------------------------------------- diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/DynamicDestinations.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/DynamicDestinations.java index c5c2462..edb1e0d 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/DynamicDestinations.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/DynamicDestinations.java @@ -23,7 +23,8 @@ import static com.google.common.base.Preconditions.checkArgument; import com.google.api.services.bigquery.model.TableSchema; import com.google.common.collect.Lists; import java.io.Serializable; -import java.lang.reflect.TypeVariable; +import java.lang.reflect.ParameterizedType; +import java.lang.reflect.Type; import java.util.List; import javax.annotation.Nullable; import org.apache.beam.sdk.coders.CannotProvideCoderException; @@ -31,7 +32,6 @@ import org.apache.beam.sdk.coders.Coder; import org.apache.beam.sdk.coders.CoderRegistry; import org.apache.beam.sdk.transforms.DoFn; import org.apache.beam.sdk.values.PCollectionView; -import org.apache.beam.sdk.values.TypeDescriptor; import org.apache.beam.sdk.values.ValueInSingleWindow; /** @@ -158,16 +158,21 @@ public abstract class DynamicDestinations<T, DestinationT> implements Serializab } // If dynamicDestinations doesn't provide a coder, try to find it in the coder registry. // We must first use reflection to figure out what the type parameter is. - TypeDescriptor<?> superDescriptor = - TypeDescriptor.of(getClass()).getSupertype(DynamicDestinations.class); - if (!superDescriptor.getRawType().equals(DynamicDestinations.class)) { - throw new AssertionError( - "Couldn't find the DynamicDestinations superclass of " + this.getClass()); + for (Type superclass = getClass().getGenericSuperclass(); + superclass != null; + superclass = ((Class) superclass).getGenericSuperclass()) { + if (superclass instanceof ParameterizedType) { + ParameterizedType parameterized = (ParameterizedType) superclass; + if (parameterized.getRawType() == DynamicDestinations.class) { + // DestinationT is the second parameter. + Type parameter = parameterized.getActualTypeArguments()[1]; + @SuppressWarnings("unchecked") + Class<DestinationT> parameterClass = (Class<DestinationT>) parameter; + return registry.getCoder(parameterClass); + } + } } - TypeVariable typeVariable = superDescriptor.getTypeParameter("DestinationT"); - @SuppressWarnings("unchecked") - TypeDescriptor<DestinationT> descriptor = - (TypeDescriptor<DestinationT>) superDescriptor.resolveType(typeVariable); - return registry.getCoder(descriptor); + throw new AssertionError( + "Couldn't find the DynamicDestinations superclass of " + this.getClass()); } } http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/GenerateShardedTable.java ---------------------------------------------------------------------- diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/GenerateShardedTable.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/GenerateShardedTable.java index 55672ff..90d41a0 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/GenerateShardedTable.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/GenerateShardedTable.java @@ -23,7 +23,6 @@ import java.util.concurrent.ThreadLocalRandom; import org.apache.beam.sdk.transforms.DoFn; import org.apache.beam.sdk.transforms.windowing.BoundedWindow; import org.apache.beam.sdk.values.KV; -import org.apache.beam.sdk.values.ShardedKey; /** * Given a write to a specific table, assign that to one of the http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/ShardedKey.java ---------------------------------------------------------------------- diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/ShardedKey.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/ShardedKey.java new file mode 100644 index 0000000..c2b739f --- /dev/null +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/ShardedKey.java @@ -0,0 +1,67 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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 org.apache.beam.sdk.io.gcp.bigquery; + +import java.io.Serializable; +import java.util.Objects; + +/** + * A key and a shard number. + */ +class ShardedKey<K> implements Serializable { + private static final long serialVersionUID = 1L; + private final K key; + private final int shardNumber; + + public static <K> ShardedKey<K> of(K key, int shardNumber) { + return new ShardedKey<>(key, shardNumber); + } + + ShardedKey(K key, int shardNumber) { + this.key = key; + this.shardNumber = shardNumber; + } + + public K getKey() { + return key; + } + + public int getShardNumber() { + return shardNumber; + } + + @Override + public String toString() { + return "key: " + key + " shard: " + shardNumber; + } + + @Override + public boolean equals(Object o) { + if (!(o instanceof ShardedKey)) { + return false; + } + ShardedKey<K> other = (ShardedKey<K>) o; + return Objects.equals(key, other.key) && Objects.equals(shardNumber, other.shardNumber); + } + + @Override + public int hashCode() { + return Objects.hash(key, shardNumber); + } +}
