Search in sources :

Example 6 with PubSubToElasticsearchOptions

use of com.google.cloud.teleport.v2.elasticsearch.options.PubSubToElasticsearchOptions in project DataflowTemplates by GoogleCloudPlatform.

the class PubSubToElasticsearch method run.

/**
 * Runs the pipeline with the supplied options.
 *
 * @param options The execution parameters to the pipeline.
 * @return The result of the pipeline execution.
 */
public static PipelineResult run(PubSubToElasticsearchOptions options) {
    // Create the pipeline
    Pipeline pipeline = Pipeline.create(options);
    // Register the coders for pipeline
    CoderRegistry coderRegistry = pipeline.getCoderRegistry();
    coderRegistry.registerCoderForType(FAILSAFE_ELEMENT_CODER.getEncodedTypeDescriptor(), FAILSAFE_ELEMENT_CODER);
    coderRegistry.registerCoderForType(CODER.getEncodedTypeDescriptor(), CODER);
    /*
     * Steps: 1) Read PubSubMessage with attributes from input PubSub subscription.
     *        2) Apply Javascript UDF if provided.
     *        3) Index Json string to output ES index.
     *
     */
    LOG.info("Reading from subscription: " + options.getInputSubscription());
    PCollectionTuple convertedPubsubMessages = pipeline.apply("ReadPubSubSubscription", PubsubIO.readMessagesWithAttributes().fromSubscription(options.getInputSubscription())).apply("ConvertMessageToJsonDocument", PubSubMessageToJsonDocument.newBuilder().setJavascriptTextTransformFunctionName(options.getJavascriptTextTransformFunctionName()).setJavascriptTextTransformGcsPath(options.getJavascriptTextTransformGcsPath()).build());
    /*
     * Step #3a: Write Json documents into Elasticsearch using {@link ElasticsearchTransforms.WriteToElasticsearch}.
     */
    convertedPubsubMessages.get(TRANSFORM_OUT).apply("GetJsonDocuments", MapElements.into(TypeDescriptors.strings()).via(FailsafeElement::getPayload)).apply("Insert metadata", new ProcessEventMetadata()).apply("WriteToElasticsearch", WriteToElasticsearch.newBuilder().setOptions(options.as(PubSubToElasticsearchOptions.class)).build());
    /*
     * Step 3b: Write elements that failed processing to error output PubSub topic via {@link PubSubIO}.
     */
    convertedPubsubMessages.get(TRANSFORM_ERROROUTPUT_OUT).apply(ParDo.of(new FailedPubsubMessageToPubsubTopicFn())).apply("writeFailureMessages", PubsubIO.writeMessages().to(options.getErrorOutputTopic()));
    // Execute the pipeline and return the result.
    return pipeline.run();
}
Also used : FailedPubsubMessageToPubsubTopicFn(com.google.cloud.teleport.v2.elasticsearch.transforms.FailedPubsubMessageToPubsubTopicFn) CoderRegistry(org.apache.beam.sdk.coders.CoderRegistry) ProcessEventMetadata(com.google.cloud.teleport.v2.elasticsearch.transforms.ProcessEventMetadata) PCollectionTuple(org.apache.beam.sdk.values.PCollectionTuple) Pipeline(org.apache.beam.sdk.Pipeline) PubSubToElasticsearchOptions(com.google.cloud.teleport.v2.elasticsearch.options.PubSubToElasticsearchOptions)

Example 7 with PubSubToElasticsearchOptions

use of com.google.cloud.teleport.v2.elasticsearch.options.PubSubToElasticsearchOptions in project DataflowTemplates by GoogleCloudPlatform.

the class PubSubToElasticsearchTest method testPubSubToElasticsearchUdfE2E.

/**
 * Tests the {@link PubSubToElasticsearch} pipeline end-to-end with a UDF.
 */
@Test
public void testPubSubToElasticsearchUdfE2E() {
    CoderRegistry coderRegistry = pipeline.getCoderRegistry();
    coderRegistry.registerCoderForType(PubSubToElasticsearch.FAILSAFE_ELEMENT_CODER.getEncodedTypeDescriptor(), PubSubToElasticsearch.FAILSAFE_ELEMENT_CODER);
    coderRegistry.registerCoderForType(PubSubToElasticsearch.CODER.getEncodedTypeDescriptor(), PubSubToElasticsearch.CODER);
    PubSubToElasticsearchOptions options = TestPipeline.testingPipelineOptions().as(PubSubToElasticsearchOptions.class);
    options.setErrorOutputTopic("projects/test/topics/test-error-topic");
    options.setJavascriptTextTransformFunctionName("transform");
    options.setJavascriptTextTransformGcsPath(TRANSFORM_FILE_PATH);
    options.setApiKey("key");
    PCollectionTuple pc = pipeline.apply(Create.of(goodTestMessages.get(0))).apply(PubSubMessageToJsonDocument.newBuilder().setJavascriptTextTransformFunctionName(options.getJavascriptTextTransformFunctionName()).setJavascriptTextTransformGcsPath(options.getJavascriptTextTransformGcsPath()).build());
    PAssert.that(pc.get(PubSubToElasticsearch.TRANSFORM_OUT)).satisfies(collection -> {
        FailsafeElement<PubsubMessage, String> element = collection.iterator().next();
        assertThat(element.getOriginalPayload().getPayload(), is(equalTo(goodTestMessages.get(0).getPayload())));
        return null;
    });
    // Execute pipeline
    pipeline.run(options);
}
Also used : CoderRegistry(org.apache.beam.sdk.coders.CoderRegistry) PCollectionTuple(org.apache.beam.sdk.values.PCollectionTuple) PubsubMessage(org.apache.beam.sdk.io.gcp.pubsub.PubsubMessage) PubSubToElasticsearchOptions(com.google.cloud.teleport.v2.elasticsearch.options.PubSubToElasticsearchOptions) Test(org.junit.Test)

Example 8 with PubSubToElasticsearchOptions

use of com.google.cloud.teleport.v2.elasticsearch.options.PubSubToElasticsearchOptions in project DataflowTemplates by GoogleCloudPlatform.

the class PubSubToElasticsearchTest method testPubSubToElasticsearchOnlyAttributesE2E.

/**
 * Tests the {@link PubSubToElasticsearch} pipeline end-to-end with an empty message payload but
 * attributes populated.
 */
@Test
public void testPubSubToElasticsearchOnlyAttributesE2E() {
    CoderRegistry coderRegistry = pipeline.getCoderRegistry();
    coderRegistry.registerCoderForType(PubSubToElasticsearch.FAILSAFE_ELEMENT_CODER.getEncodedTypeDescriptor(), PubSubToElasticsearch.FAILSAFE_ELEMENT_CODER);
    coderRegistry.registerCoderForType(PubSubToElasticsearch.CODER.getEncodedTypeDescriptor(), PubSubToElasticsearch.CODER);
    PubSubToElasticsearchOptions options = TestPipeline.testingPipelineOptions().as(PubSubToElasticsearchOptions.class);
    options.setErrorOutputTopic("projects/test/topics/test-error-topic");
    options.setApiKey("key");
    PCollectionTuple pc = pipeline.apply(Create.of(goodTestMessages.get(goodTestMessages.size() - 1))).apply(PubSubMessageToJsonDocument.newBuilder().setJavascriptTextTransformFunctionName(options.getJavascriptTextTransformFunctionName()).setJavascriptTextTransformGcsPath(options.getJavascriptTextTransformGcsPath()).build());
    PAssert.that(pc.get(PubSubToElasticsearch.TRANSFORM_OUT)).satisfies(collection -> {
        FailsafeElement<PubsubMessage, String> element = collection.iterator().next();
        assertThat(new Gson().fromJson(element.getPayload(), HashMap.class), is(equalTo(element.getOriginalPayload().getAttributeMap())));
        return null;
    });
    // Execute pipeline
    pipeline.run(options);
}
Also used : CoderRegistry(org.apache.beam.sdk.coders.CoderRegistry) HashMap(java.util.HashMap) PCollectionTuple(org.apache.beam.sdk.values.PCollectionTuple) Gson(com.google.gson.Gson) PubsubMessage(org.apache.beam.sdk.io.gcp.pubsub.PubsubMessage) PubSubToElasticsearchOptions(com.google.cloud.teleport.v2.elasticsearch.options.PubSubToElasticsearchOptions) Test(org.junit.Test)

Example 9 with PubSubToElasticsearchOptions

use of com.google.cloud.teleport.v2.elasticsearch.options.PubSubToElasticsearchOptions in project DataflowTemplates by GoogleCloudPlatform.

the class PubSubToElasticsearchTest method testPubSubToElasticsearchBadUdfE2E.

/**
 * Tests the {@link PubSubToElasticsearch} pipeline end-to-end with a bad UDF.
 */
@Test
public void testPubSubToElasticsearchBadUdfE2E() {
    CoderRegistry coderRegistry = pipeline.getCoderRegistry();
    coderRegistry.registerCoderForType(PubSubToElasticsearch.FAILSAFE_ELEMENT_CODER.getEncodedTypeDescriptor(), PubSubToElasticsearch.FAILSAFE_ELEMENT_CODER);
    coderRegistry.registerCoderForType(PubSubToElasticsearch.CODER.getEncodedTypeDescriptor(), PubSubToElasticsearch.CODER);
    PubSubToElasticsearchOptions options = TestPipeline.testingPipelineOptions().as(PubSubToElasticsearchOptions.class);
    options.setErrorOutputTopic("projects/test/topics/test-error-topic");
    options.setJavascriptTextTransformFunctionName("transformBad");
    options.setJavascriptTextTransformGcsPath(BAD_TRANSFORM_FILE_PATH);
    options.setApiKey("key");
    PCollectionTuple pc = pipeline.apply(Create.of(badTestMessages.get(0))).apply(PubSubMessageToJsonDocument.newBuilder().setJavascriptTextTransformFunctionName(options.getJavascriptTextTransformFunctionName()).setJavascriptTextTransformGcsPath(options.getJavascriptTextTransformGcsPath()).build());
    PAssert.that(pc.get(PubSubToElasticsearch.TRANSFORM_ERROROUTPUT_OUT)).satisfies(collection -> {
        FailsafeElement<PubsubMessage, String> element = collection.iterator().next();
        assertThat(element.getOriginalPayload().getPayload(), is(equalTo(badTestMessages.get(0).getPayload())));
        return null;
    });
    PAssert.that(pc.get(PubSubToElasticsearch.TRANSFORM_OUT)).empty();
    // Execute pipeline
    pipeline.run(options);
}
Also used : CoderRegistry(org.apache.beam.sdk.coders.CoderRegistry) PCollectionTuple(org.apache.beam.sdk.values.PCollectionTuple) PubsubMessage(org.apache.beam.sdk.io.gcp.pubsub.PubsubMessage) PubSubToElasticsearchOptions(com.google.cloud.teleport.v2.elasticsearch.options.PubSubToElasticsearchOptions) Test(org.junit.Test)

Aggregations

PubSubToElasticsearchOptions (com.google.cloud.teleport.v2.elasticsearch.options.PubSubToElasticsearchOptions)9 Test (org.junit.Test)7 CoderRegistry (org.apache.beam.sdk.coders.CoderRegistry)5 PCollectionTuple (org.apache.beam.sdk.values.PCollectionTuple)5 PubsubMessage (org.apache.beam.sdk.io.gcp.pubsub.PubsubMessage)4 JsonNode (com.fasterxml.jackson.databind.JsonNode)3 FailedPubsubMessageToPubsubTopicFn (com.google.cloud.teleport.v2.elasticsearch.transforms.FailedPubsubMessageToPubsubTopicFn)1 ProcessEventMetadata (com.google.cloud.teleport.v2.elasticsearch.transforms.ProcessEventMetadata)1 ElasticsearchIndex (com.google.cloud.teleport.v2.elasticsearch.utils.ElasticsearchIndex)1 Gson (com.google.gson.Gson)1 HashMap (java.util.HashMap)1 Pipeline (org.apache.beam.sdk.Pipeline)1