Search in sources :

Example 1 with PubSubProtoToBigQueryOptions

use of com.google.cloud.teleport.v2.templates.PubsubProtoToBigQuery.PubSubProtoToBigQueryOptions in project DataflowTemplates by GoogleCloudPlatform.

the class PubsubProtoToBigQueryTest method testGetDescriptorWithInvalidMessageName.

@Test
public void testGetDescriptorWithInvalidMessageName() {
    PubSubProtoToBigQueryOptions options = getOptions();
    options.setProtoSchemaPath(GENERATED_PROTO_SCHEMA_PATH);
    String badMessageName = "invalid.message.Name";
    options.setFullMessageName(badMessageName);
    RuntimeException exception = assertThrows(RuntimeException.class, () -> PubsubProtoToBigQuery.getDescriptor(options));
    assertThat(exception).hasMessageThat().contains(GENERATED_PROTO_SCHEMA_PATH);
    assertThat(exception).hasMessageThat().contains(badMessageName);
}
Also used : PubSubProtoToBigQueryOptions(com.google.cloud.teleport.v2.templates.PubsubProtoToBigQuery.PubSubProtoToBigQueryOptions) Test(org.junit.Test)

Example 2 with PubSubProtoToBigQueryOptions

use of com.google.cloud.teleport.v2.templates.PubsubProtoToBigQuery.PubSubProtoToBigQueryOptions in project DataflowTemplates by GoogleCloudPlatform.

the class PubsubProtoToBigQueryTest method testApplyUdfWithPathButNoFunction.

@Test
public void testApplyUdfWithPathButNoFunction() {
    PubSubProtoToBigQueryOptions options = getOptions();
    options.setJavascriptTextTransformGcsPath("/some/path.js");
    PCollection<String> input = pipeline.apply(Create.of(""));
    assertThrows(IllegalArgumentException.class, () -> runUdf(input, options));
    options.setJavascriptTextTransformFunctionName("");
    assertThrows(IllegalArgumentException.class, () -> runUdf(input, options));
    pipeline.run();
}
Also used : PubSubProtoToBigQueryOptions(com.google.cloud.teleport.v2.templates.PubsubProtoToBigQuery.PubSubProtoToBigQueryOptions) Test(org.junit.Test)

Example 3 with PubSubProtoToBigQueryOptions

use of com.google.cloud.teleport.v2.templates.PubsubProtoToBigQuery.PubSubProtoToBigQueryOptions in project DataflowTemplates by GoogleCloudPlatform.

the class PubsubProtoToBigQueryTest method testWriteBigQueryWithInvalidJsonSchemaPath.

@Test
public void testWriteBigQueryWithInvalidJsonSchemaPath() {
    PubSubProtoToBigQueryOptions options = getOptions();
    String path = "/some/invalid/path.json";
    options.setBigQueryTableSchemaPath(path);
    IllegalArgumentException exception = assertThrows(// Can pass a null descriptor, since it shouldn't be used.
    IllegalArgumentException.class, () -> PubsubProtoToBigQuery.writeToBigQuery(options, null));
    assertThat(exception).hasMessageThat().contains(path);
}
Also used : PubSubProtoToBigQueryOptions(com.google.cloud.teleport.v2.templates.PubsubProtoToBigQuery.PubSubProtoToBigQueryOptions) Test(org.junit.Test)

Example 4 with PubSubProtoToBigQueryOptions

use of com.google.cloud.teleport.v2.templates.PubsubProtoToBigQuery.PubSubProtoToBigQueryOptions in project DataflowTemplates by GoogleCloudPlatform.

the class PubsubProtoToBigQuery method runUdf.

/**
 * Handles running the UDF.
 *
 * <p>If {@code options} is configured so as not to run the UDF, then the UDF will not be called.
 *
 * <p>This may add a branch to the pipeline for outputting failed UDF records to an unprocessed
 * topic.
 *
 * @param jsonCollection {@link PCollection} of JSON strings for use as input to the UDF
 * @param options the options containing info on running the UDF
 * @return the {@link PCollection} of UDF output as JSON or {@code jsonCollection} if UDF not
 *     called
 */
@VisibleForTesting
static PCollection<String> runUdf(PCollection<String> jsonCollection, PubSubProtoToBigQueryOptions options) {
    // intended, simply return the input as "success" output.
    if (Strings.isNullOrEmpty(options.getJavascriptTextTransformGcsPath())) {
        return jsonCollection;
    }
    // a value.
    if (Strings.isNullOrEmpty(options.getJavascriptTextTransformFunctionName())) {
        throw new IllegalArgumentException("JavaScript function name cannot be null or empty if file is set");
    }
    PCollectionTuple maybeSuccess = jsonCollection.apply("Run UDF", new RunUdf(options));
    maybeSuccess.get(UDF_FAILURE_TAG).setCoder(FAILSAFE_CODER).apply("Get UDF Failures", ConvertFailsafeElementToPubsubMessage.<String, String>builder().setOriginalPayloadSerializeFn(s -> ArrayUtils.toObject(s.getBytes(UTF_8))).setErrorMessageAttributeKey("udfErrorMessage").build()).apply("Write Failed UDF", writeUdfFailures(options));
    return maybeSuccess.get(UDF_SUCCESS_TAG).setCoder(FAILSAFE_CODER).apply("Get UDF Output", MapElements.into(TypeDescriptors.strings()).via(FailsafeElement::getPayload)).setCoder(NullableCoder.of(StringUtf8Coder.of()));
}
Also used : PCollectionTuple(org.apache.beam.sdk.values.PCollectionTuple) FailsafeElement(com.google.cloud.teleport.v2.values.FailsafeElement) VisibleForTesting(com.google.common.annotations.VisibleForTesting)

Example 5 with PubSubProtoToBigQueryOptions

use of com.google.cloud.teleport.v2.templates.PubsubProtoToBigQuery.PubSubProtoToBigQueryOptions in project DataflowTemplates by GoogleCloudPlatform.

the class PubsubProtoToBigQueryTest method testApplyUdfWithNoUdfPathSet.

@Test
public void testApplyUdfWithNoUdfPathSet() {
    PubSubProtoToBigQueryOptions options = getOptions();
    ImmutableList<String> inputs = ImmutableList.of("First", "Second", "Third");
    PCollection<String> pInput = pipeline.apply(Create.of(inputs));
    PAssert.that(runUdf(pInput, options)).containsInAnyOrder(inputs);
    pipeline.run();
}
Also used : PubSubProtoToBigQueryOptions(com.google.cloud.teleport.v2.templates.PubsubProtoToBigQuery.PubSubProtoToBigQueryOptions) Test(org.junit.Test)

Aggregations

PubSubProtoToBigQueryOptions (com.google.cloud.teleport.v2.templates.PubsubProtoToBigQuery.PubSubProtoToBigQueryOptions)6 Test (org.junit.Test)6 FailsafeElement (com.google.cloud.teleport.v2.values.FailsafeElement)1 VisibleForTesting (com.google.common.annotations.VisibleForTesting)1 PCollectionTuple (org.apache.beam.sdk.values.PCollectionTuple)1