Search in sources :

Example 61 with Options

use of com.google.cloud.teleport.v2.transforms.DeleteBigQueryDataFn.Options in project DataflowTemplates by GoogleCloudPlatform.

the class GCSToElasticsearchTest method testGCSToElasticsearchUdfE2E.

/**
 * Tests the {@link GCSToElasticsearch} pipeline using a Udf to parse the Csv.
 */
@Test
public void testGCSToElasticsearchUdfE2E() {
    final String record = "007,CA,26.23";
    final String stringifiedJsonRecord = "{\"id\":\"007\",\"state\":\"CA\",\"price\":26.23}";
    final FailsafeElementCoder<String, String> coder = FailsafeElementCoder.of(NullableCoder.of(StringUtf8Coder.of()), NullableCoder.of(StringUtf8Coder.of()));
    CoderRegistry coderRegistry = pipeline.getCoderRegistry();
    coderRegistry.registerCoderForType(coder.getEncodedTypeDescriptor(), coder);
    GCSToElasticsearchOptions options = PipelineOptionsFactory.create().as(GCSToElasticsearchOptions.class);
    options.setJavascriptTextTransformGcsPath(TRANSFORM_FILE_PATH);
    options.setJavascriptTextTransformFunctionName("transform");
    options.setContainsHeaders(false);
    options.setInputFileSpec(NO_HEADER_CSV_FILE_PATH);
    options.setApiKey("key");
    // Build pipeline with no headers.
    PCollectionTuple readCsvOut = pipeline.apply("ReadCsv", CsvConverters.ReadCsv.newBuilder().setCsvFormat(options.getCsvFormat()).setDelimiter(options.getDelimiter()).setHasHeaders(options.getContainsHeaders()).setInputFileSpec(options.getInputFileSpec()).setHeaderTag(GCSToElasticsearch.CSV_HEADERS).setLineTag(GCSToElasticsearch.CSV_LINES).setFileEncoding(options.getCsvFileEncoding()).build()).apply("ConvertLine", CsvConverters.LineToFailsafeJson.newBuilder().setDelimiter(options.getDelimiter()).setUdfFileSystemPath(options.getJavascriptTextTransformGcsPath()).setUdfFunctionName(options.getJavascriptTextTransformFunctionName()).setJsonSchemaPath(options.getJsonSchemaPath()).setHeaderTag(GCSToElasticsearch.CSV_HEADERS).setLineTag(GCSToElasticsearch.CSV_LINES).setUdfOutputTag(GCSToElasticsearch.PROCESSING_OUT).setUdfDeadletterTag(GCSToElasticsearch.PROCESSING_DEADLETTER_OUT).build());
    // Assert
    PAssert.that(readCsvOut.get(GCSToElasticsearch.PROCESSING_OUT)).satisfies(collection -> {
        FailsafeElement element = collection.iterator().next();
        assertThat(element.getOriginalPayload(), is(equalTo(record)));
        assertThat(element.getPayload(), is(equalTo(stringifiedJsonRecord)));
        return null;
    });
    // Execute pipeline
    pipeline.run();
}
Also used : CoderRegistry(org.apache.beam.sdk.coders.CoderRegistry) GCSToElasticsearchOptions(com.google.cloud.teleport.v2.elasticsearch.options.GCSToElasticsearchOptions) PCollectionTuple(org.apache.beam.sdk.values.PCollectionTuple) FailsafeElement(com.google.cloud.teleport.v2.values.FailsafeElement) Test(org.junit.Test)

Example 62 with Options

use of com.google.cloud.teleport.v2.transforms.DeleteBigQueryDataFn.Options 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 63 with Options

use of com.google.cloud.teleport.v2.transforms.DeleteBigQueryDataFn.Options 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 64 with Options

use of com.google.cloud.teleport.v2.transforms.DeleteBigQueryDataFn.Options 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)

Example 65 with Options

use of com.google.cloud.teleport.v2.transforms.DeleteBigQueryDataFn.Options in project DataflowTemplates by GoogleCloudPlatform.

the class WriteToElasticsearchTest method testElasticsearchWriteOptionsRetryConfigMaxDuration.

/**
 * Tests that {@link WriteToElasticsearch} throws an exception if {@link
 * org.apache.beam.sdk.io.elasticsearch.ElasticsearchIO.RetryConfiguration} are invalid.
 */
@Test
public void testElasticsearchWriteOptionsRetryConfigMaxDuration() {
    exceptionRule.expect(IllegalArgumentException.class);
    ElasticsearchWriteOptions options = PipelineOptionsFactory.create().as(ElasticsearchWriteOptions.class);
    options.setConnectionUrl("https://host.domain");
    options.setApiKey("key");
    options.setMaxRetryDuration(null);
    options.setMaxRetryAttempts(3);
    pipeline.apply("CreateInput", Create.of("test")).apply("TestWriteToElasticsearchBadMaxDuration", WriteToElasticsearch.newBuilder().setOptions(options).build());
    pipeline.run();
}
Also used : ElasticsearchWriteOptions(com.google.cloud.teleport.v2.elasticsearch.options.ElasticsearchWriteOptions) Test(org.junit.Test)

Aggregations

Test (org.junit.Test)63 PCollectionTuple (org.apache.beam.sdk.values.PCollectionTuple)25 FailsafeElement (com.google.cloud.teleport.v2.values.FailsafeElement)20 Pipeline (org.apache.beam.sdk.Pipeline)19 CoderRegistry (org.apache.beam.sdk.coders.CoderRegistry)19 BigQueryTable (com.google.cloud.teleport.v2.values.BigQueryTable)15 GenericRecord (org.apache.avro.generic.GenericRecord)12 Category (org.junit.experimental.categories.Category)12 Filter (com.google.cloud.teleport.v2.utils.BigQueryMetadataLoader.Filter)10 BigQueryTablePartition (com.google.cloud.teleport.v2.values.BigQueryTablePartition)10 PubSubToElasticsearchOptions (com.google.cloud.teleport.v2.elasticsearch.options.PubSubToElasticsearchOptions)9 TableRow (com.google.api.services.bigquery.model.TableRow)8 DataplexClient (com.google.cloud.teleport.v2.clients.DataplexClient)8 FileFormatConversionOptions (com.google.cloud.teleport.v2.templates.FileFormatConversion.FileFormatConversionOptions)8 KV (org.apache.beam.sdk.values.KV)8 ArrayList (java.util.ArrayList)7 ElasticsearchWriteOptions (com.google.cloud.teleport.v2.elasticsearch.options.ElasticsearchWriteOptions)6 GCSToElasticsearchOptions (com.google.cloud.teleport.v2.elasticsearch.options.GCSToElasticsearchOptions)6 FileFormatConversionOptions (com.google.cloud.teleport.v2.templates.DataplexFileFormatConversion.FileFormatConversionOptions)6 PubSubProtoToBigQueryOptions (com.google.cloud.teleport.v2.templates.PubsubProtoToBigQuery.PubSubProtoToBigQueryOptions)6