Examples with TokenNestedValueProvider com.google.cloud.teleport.util.TokenNestedValueProvider used on opensource projects

Example 1 with TokenNestedValueProvider

use of com.google.cloud.teleport.util.TokenNestedValueProvider in project DataflowTemplates by GoogleCloudPlatform.

the class PubSubToSplunk method run.

/**
 * Runs the pipeline to completion with the specified options. This method does not wait until the
 * pipeline is finished before returning. Invoke {@code result.waitUntilFinish()} on the result
 * object to block until the pipeline is finished running if blocking programmatic execution is
 * required.
 *
 * @param options The execution options.
 * @return The pipeline result.
 */
public static PipelineResult run(PubSubToSplunkOptions options) {
    Pipeline pipeline = Pipeline.create(options);
    // Register coders.
    CoderRegistry registry = pipeline.getCoderRegistry();
    registry.registerCoderForClass(SplunkEvent.class, SplunkEventCoder.of());
    registry.registerCoderForType(FAILSAFE_ELEMENT_CODER.getEncodedTypeDescriptor(), FAILSAFE_ELEMENT_CODER);
    /*
     * Steps:
     *  1) Read messages in from Pub/Sub
     *  2) Convert message to FailsafeElement for processing.
     *  3) Apply user provided UDF (if any) on the input strings.
     *  4) Convert successfully transformed messages into SplunkEvent objects
     *  5) Write SplunkEvents to Splunk's HEC end point.
     *  5a) Wrap write failures into a FailsafeElement.
     *  6) Collect errors from UDF transform (#3), SplunkEvent transform (#4)
     *     and writing to Splunk HEC (#5) and stream into a Pub/Sub deadletter topic.
     */
    // 1) Read messages in from Pub/Sub
    PCollection<String> stringMessages = pipeline.apply("ReadMessages", new ReadMessages(options.getInputSubscription(), options.getIncludePubsubMessage()));
    // 2) Convert message to FailsafeElement for processing.
    PCollectionTuple transformedOutput = stringMessages.apply("ConvertToFailsafeElement", MapElements.into(FAILSAFE_ELEMENT_CODER.getEncodedTypeDescriptor()).via(input -> FailsafeElement.of(input, input))).apply("ApplyUDFTransformation", FailsafeJavascriptUdf.<String>newBuilder().setFileSystemPath(options.getJavascriptTextTransformGcsPath()).setFunctionName(options.getJavascriptTextTransformFunctionName()).setLoggingEnabled(ValueProvider.StaticValueProvider.of(true)).setSuccessTag(UDF_OUT).setFailureTag(UDF_DEADLETTER_OUT).build());
    // 4) Convert successfully transformed messages into SplunkEvent objects
    PCollectionTuple convertToEventTuple = transformedOutput.get(UDF_OUT).apply("ConvertToSplunkEvent", SplunkConverters.failsafeStringToSplunkEvent(SPLUNK_EVENT_OUT, SPLUNK_EVENT_DEADLETTER_OUT));
    // 5) Write SplunkEvents to Splunk's HEC end point.
    PCollection<SplunkWriteError> writeErrors = convertToEventTuple.get(SPLUNK_EVENT_OUT).apply("WriteToSplunk", SplunkIO.writeBuilder().withToken(new TokenNestedValueProvider(options.getTokenSecretId(), options.getTokenKMSEncryptionKey(), options.getToken(), options.getTokenSource())).withUrl(options.getUrl()).withBatchCount(options.getBatchCount()).withParallelism(options.getParallelism()).withDisableCertificateValidation(options.getDisableCertificateValidation()).withRootCaCertificatePath(options.getRootCaCertificatePath()).withEnableBatchLogs(options.getEnableBatchLogs()).withEnableGzipHttpCompression(options.getEnableGzipHttpCompression()).build());
    // 5a) Wrap write failures into a FailsafeElement.
    PCollection<FailsafeElement<String, String>> wrappedSplunkWriteErrors = writeErrors.apply("WrapSplunkWriteErrors", ParDo.of(new DoFn<SplunkWriteError, FailsafeElement<String, String>>() {

        @ProcessElement
        public void processElement(ProcessContext context) {
            SplunkWriteError error = context.element();
            FailsafeElement<String, String> failsafeElement = FailsafeElement.of(error.payload(), error.payload());
            if (error.statusMessage() != null) {
                failsafeElement.setErrorMessage(error.statusMessage());
            }
            if (error.statusCode() != null) {
                failsafeElement.setErrorMessage(String.format("Splunk write status code: %d", error.statusCode()));
            }
            context.output(failsafeElement);
        }
    }));
    // 6) Collect errors from UDF transform (#4), SplunkEvent transform (#5)
    // and writing to Splunk HEC (#6) and stream into a Pub/Sub deadletter topic.
    PCollectionList.of(ImmutableList.of(convertToEventTuple.get(SPLUNK_EVENT_DEADLETTER_OUT), wrappedSplunkWriteErrors, transformedOutput.get(UDF_DEADLETTER_OUT))).apply("FlattenErrors", Flatten.pCollections()).apply("WriteFailedRecords", ErrorConverters.WriteStringMessageErrorsToPubSub.newBuilder().setErrorRecordsTopic(options.getOutputDeadletterTopic()).build());
    return pipeline.run();
}