diff --git a/CHANGES.md b/CHANGES.md index cb5eced42f7c..afae8ac2b660 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -64,6 +64,7 @@ ## I/Os +* Added a SchemaTransform for SplunkIO writes, available in Beam YAML as `WriteToSplunk` (Java/YAML) ([#40470](https://github.com/apache/beam/issues/40470)). * Support for X source added (Java/Python) ([#X](https://github.com/apache/beam/issues/X)). ## New Features / Improvements @@ -99,6 +100,7 @@ * (Python) Fixed stateful DoFns with side inputs sometimes taking the timer key coder from a side input instead of the main input, which could make the worker fail to decode timer keys with `Unknown type tag` ([#40374](https://github.com/apache/beam/issues/40374)). * (Python) `Duration` built from float seconds now rounds to the nearest microsecond instead of truncating, which could lose a microsecond ([#40263](https://github.com/apache/beam/issues/40263)). * (Python) `ReadFromText` with `escapechar` no longer skips a delimiter that directly follows an escaped delimiter, which merged two records into one ([#40459](https://github.com/apache/beam/issues/40459)). +* (Java) SplunkIO write errors no longer lose their status code, status message and payload when encoded. The schema inferred for `SplunkWriteError` had no fields, so `SplunkIO.Write` now uses a dedicated `SplunkWriteErrorCoder` for its output ([#40470](https://github.com/apache/beam/issues/40470)). * Fixed X (Java/Python) ([#X](https://github.com/apache/beam/issues/X)). ## Security Fixes diff --git a/sdks/java/io/expansion-service/build.gradle b/sdks/java/io/expansion-service/build.gradle index 28b4060a2b58..4de4b016c739 100644 --- a/sdks/java/io/expansion-service/build.gradle +++ b/sdks/java/io/expansion-service/build.gradle @@ -99,6 +99,7 @@ dependencies { runtimeOnly project(":sdks:java:io:datadog") runtimeOnly project(":sdks:java:io:mongodb") runtimeOnly project(":sdks:java:io:delta") + runtimeOnly project(":sdks:java:io:splunk") runtimeOnly library.java.kafka_clients runtimeOnly library.java.slf4j_jdk14 diff --git a/sdks/java/io/splunk/src/main/java/org/apache/beam/sdk/io/splunk/SplunkEvent.java b/sdks/java/io/splunk/src/main/java/org/apache/beam/sdk/io/splunk/SplunkEvent.java index 177900a2d09a..400e477057f9 100644 --- a/sdks/java/io/splunk/src/main/java/org/apache/beam/sdk/io/splunk/SplunkEvent.java +++ b/sdks/java/io/splunk/src/main/java/org/apache/beam/sdk/io/splunk/SplunkEvent.java @@ -43,6 +43,14 @@ @AutoValue public abstract class SplunkEvent { + public static final String TIME = "time"; + public static final String HOST = "host"; + public static final String SOURCE = "source"; + public static final String SOURCE_TYPE = "sourcetype"; + public static final String INDEX = "index"; + public static final String FIELDS = "fields"; + public static final String EVENT = "event"; + /** Provides a builder for creating {@link SplunkEvent} objects. */ public static Builder newBuilder() { return new AutoValue_SplunkEvent.Builder(); diff --git a/sdks/java/io/splunk/src/main/java/org/apache/beam/sdk/io/splunk/SplunkIO.java b/sdks/java/io/splunk/src/main/java/org/apache/beam/sdk/io/splunk/SplunkIO.java index 2127cc55752d..f4197a51a125 100644 --- a/sdks/java/io/splunk/src/main/java/org/apache/beam/sdk/io/splunk/SplunkIO.java +++ b/sdks/java/io/splunk/src/main/java/org/apache/beam/sdk/io/splunk/SplunkIO.java @@ -166,7 +166,8 @@ public PCollection expand(PCollection input) { // Return a PCollection return input .apply("Create KV pairs", CreateKeys.of(parallelism())) - .apply("Write Splunk events", ParDo.of(writer)); + .apply("Write Splunk events", ParDo.of(writer)) + .setCoder(SplunkWriteErrorCoder.of()); } /** A builder for creating {@link Write} objects. */ diff --git a/sdks/java/io/splunk/src/main/java/org/apache/beam/sdk/io/splunk/SplunkWriteErrorCoder.java b/sdks/java/io/splunk/src/main/java/org/apache/beam/sdk/io/splunk/SplunkWriteErrorCoder.java new file mode 100644 index 000000000000..89691f01afcd --- /dev/null +++ b/sdks/java/io/splunk/src/main/java/org/apache/beam/sdk/io/splunk/SplunkWriteErrorCoder.java @@ -0,0 +1,84 @@ +/* + * 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.splunk; + +import java.io.IOException; +import java.io.InputStream; +import java.io.OutputStream; +import org.apache.beam.sdk.coders.AtomicCoder; +import org.apache.beam.sdk.coders.CoderException; +import org.apache.beam.sdk.coders.NullableCoder; +import org.apache.beam.sdk.coders.StringUtf8Coder; +import org.apache.beam.sdk.coders.VarIntCoder; +import org.apache.beam.sdk.values.TypeDescriptor; + +/** + * A {@link org.apache.beam.sdk.coders.Coder} for {@link SplunkWriteError} objects. + * + *

{@link SplunkWriteError} accessors don't follow the getter naming convention, so the schema + * inferred for it has no fields and a schema coder would drop every value. + */ +public class SplunkWriteErrorCoder extends AtomicCoder { + + private static final SplunkWriteErrorCoder SPLUNK_WRITE_ERROR_CODER = new SplunkWriteErrorCoder(); + + private static final TypeDescriptor TYPE_DESCRIPTOR = + new TypeDescriptor() {}; + private static final NullableCoder STRING_NULLABLE_CODER = + NullableCoder.of(StringUtf8Coder.of()); + private static final NullableCoder INTEGER_NULLABLE_CODER = + NullableCoder.of(VarIntCoder.of()); + + public static SplunkWriteErrorCoder of() { + return SPLUNK_WRITE_ERROR_CODER; + } + + @Override + public void encode(SplunkWriteError value, OutputStream out) throws IOException { + INTEGER_NULLABLE_CODER.encode(value.statusCode(), out); + STRING_NULLABLE_CODER.encode(value.statusMessage(), out); + STRING_NULLABLE_CODER.encode(value.payload(), out); + } + + @Override + public SplunkWriteError decode(InputStream in) throws CoderException, IOException { + SplunkWriteError.Builder builder = SplunkWriteError.newBuilder(); + + Integer statusCode = INTEGER_NULLABLE_CODER.decode(in); + if (statusCode != null) { + builder.withStatusCode(statusCode); + } + + String statusMessage = STRING_NULLABLE_CODER.decode(in); + if (statusMessage != null) { + builder.withStatusMessage(statusMessage); + } + + String payload = STRING_NULLABLE_CODER.decode(in); + if (payload != null) { + builder.withPayload(payload); + } + + return builder.create(); + } + + @Override + public TypeDescriptor getEncodedTypeDescriptor() { + return TYPE_DESCRIPTOR; + } +} diff --git a/sdks/java/io/splunk/src/main/java/org/apache/beam/sdk/io/splunk/SplunkWriteSchemaTransformConfiguration.java b/sdks/java/io/splunk/src/main/java/org/apache/beam/sdk/io/splunk/SplunkWriteSchemaTransformConfiguration.java new file mode 100644 index 000000000000..8ae1a229ce22 --- /dev/null +++ b/sdks/java/io/splunk/src/main/java/org/apache/beam/sdk/io/splunk/SplunkWriteSchemaTransformConfiguration.java @@ -0,0 +1,116 @@ +/* + * 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.splunk; + +import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkArgument; + +import com.google.auto.value.AutoValue; +import org.apache.beam.sdk.schemas.AutoValueSchema; +import org.apache.beam.sdk.schemas.annotations.DefaultSchema; +import org.apache.beam.sdk.schemas.annotations.SchemaFieldDescription; +import org.apache.beam.sdk.schemas.transforms.providers.ErrorHandling; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Strings; +import org.checkerframework.checker.nullness.qual.Nullable; + +/** + * Configuration for writing to Splunk's Http Event Collector (HEC). + * + *

This class is meant to be used with {@link SplunkWriteSchemaTransformProvider}. + */ +@DefaultSchema(AutoValueSchema.class) +@AutoValue +public abstract class SplunkWriteSchemaTransformConfiguration { + + public void validate() { + String invalidConfigMessage = "Invalid Splunk Write configuration: "; + checkArgument(!getUrl().isEmpty(), invalidConfigMessage + "url must be specified."); + checkArgument(!getToken().isEmpty(), invalidConfigMessage + "token must be specified."); + Integer batchCount = getBatchCount(); + if (batchCount != null) { + checkArgument(batchCount > 0, invalidConfigMessage + "batchCount must be greater than 0."); + } + Integer parallelism = getParallelism(); + if (parallelism != null) { + checkArgument(parallelism > 0, invalidConfigMessage + "parallelism must be greater than 0."); + } + ErrorHandling errorHandling = getErrorHandling(); + if (errorHandling != null) { + checkArgument( + !Strings.isNullOrEmpty(errorHandling.getOutput()), + invalidConfigMessage + "Output must not be empty if error handling specified."); + } + } + + /** Instantiates a {@link SplunkWriteSchemaTransformConfiguration.Builder} instance. */ + public static SplunkWriteSchemaTransformConfiguration.Builder builder() { + return new AutoValue_SplunkWriteSchemaTransformConfiguration.Builder(); + } + + @SchemaFieldDescription("The Splunk HEC endpoint URL, e.g. https://splunk-host:8088.") + public abstract String getUrl(); + + @SchemaFieldDescription("The Splunk HEC authentication token.") + public abstract String getToken(); + + @SchemaFieldDescription("The number of events to batch together for each write.") + public abstract @Nullable Integer getBatchCount(); + + @SchemaFieldDescription("The number of parallel requests to the HEC endpoint.") + public abstract @Nullable Integer getParallelism(); + + @SchemaFieldDescription( + "Whether to disable SSL certificate validation, e.g. for self-signed certificates.") + public abstract @Nullable Boolean getDisableCertificateValidation(); + + @SchemaFieldDescription("Path to a root CA certificate used to validate the HEC endpoint.") + public abstract @Nullable String getRootCaCertificatePath(); + + @SchemaFieldDescription("Whether to log the result of each batch write.") + public abstract @Nullable Boolean getEnableBatchLogs(); + + @SchemaFieldDescription( + "Whether requests sent to the HEC endpoint should be GZIP encoded. Defaults to true.") + public abstract @Nullable Boolean getEnableGzipHttpCompression(); + + @SchemaFieldDescription("Specifies how to handle errors.") + public abstract @Nullable ErrorHandling getErrorHandling(); + + @AutoValue.Builder + public abstract static class Builder { + public abstract Builder setUrl(String url); + + public abstract Builder setToken(String token); + + public abstract Builder setBatchCount(Integer batchCount); + + public abstract Builder setParallelism(Integer parallelism); + + public abstract Builder setDisableCertificateValidation(Boolean disableCertificateValidation); + + public abstract Builder setRootCaCertificatePath(String rootCaCertificatePath); + + public abstract Builder setEnableBatchLogs(Boolean enableBatchLogs); + + public abstract Builder setEnableGzipHttpCompression(Boolean enableGzipHttpCompression); + + public abstract Builder setErrorHandling(@Nullable ErrorHandling errorHandling); + + /** Builds the {@link SplunkWriteSchemaTransformConfiguration} configuration. */ + public abstract SplunkWriteSchemaTransformConfiguration build(); + } +} diff --git a/sdks/java/io/splunk/src/main/java/org/apache/beam/sdk/io/splunk/SplunkWriteSchemaTransformProvider.java b/sdks/java/io/splunk/src/main/java/org/apache/beam/sdk/io/splunk/SplunkWriteSchemaTransformProvider.java new file mode 100644 index 000000000000..a0f294ffc8c5 --- /dev/null +++ b/sdks/java/io/splunk/src/main/java/org/apache/beam/sdk/io/splunk/SplunkWriteSchemaTransformProvider.java @@ -0,0 +1,319 @@ +/* + * 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.splunk; + +import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkNotNull; + +import com.google.auto.service.AutoService; +import com.google.gson.JsonArray; +import com.google.gson.JsonElement; +import com.google.gson.JsonObject; +import com.google.gson.JsonParser; +import com.google.gson.JsonPrimitive; +import java.net.HttpURLConnection; +import java.util.Collections; +import java.util.List; +import org.apache.beam.sdk.coders.RowCoder; +import org.apache.beam.sdk.schemas.Schema; +import org.apache.beam.sdk.schemas.transforms.SchemaTransform; +import org.apache.beam.sdk.schemas.transforms.SchemaTransformProvider; +import org.apache.beam.sdk.schemas.transforms.TypedSchemaTransformProvider; +import org.apache.beam.sdk.schemas.transforms.providers.ErrorHandling; +import org.apache.beam.sdk.transforms.Create; +import org.apache.beam.sdk.transforms.DoFn; +import org.apache.beam.sdk.transforms.Flatten; +import org.apache.beam.sdk.transforms.MapElements; +import org.apache.beam.sdk.transforms.ParDo; +import org.apache.beam.sdk.values.PCollection; +import org.apache.beam.sdk.values.PCollectionList; +import org.apache.beam.sdk.values.PCollectionRowTuple; +import org.apache.beam.sdk.values.PCollectionTuple; +import org.apache.beam.sdk.values.Row; +import org.apache.beam.sdk.values.TupleTag; +import org.apache.beam.sdk.values.TupleTagList; +import org.apache.beam.sdk.values.TypeDescriptors; + +/** + * An implementation of {@link TypedSchemaTransformProvider} for writing to Splunk's Http Event + * Collector (HEC) using {@link SplunkIO}. + * + *

Each input row is converted to a {@link SplunkEvent}. The row must contain a string {@code + * event} field, and may contain {@code time} (INT64), {@code host}, {@code source}, {@code + * sourcetype} and {@code index} (strings) and {@code fields} (a row or a JSON object string). Any + * other fields are ignored. + */ +@AutoService(SchemaTransformProvider.class) +public class SplunkWriteSchemaTransformProvider + extends TypedSchemaTransformProvider { + private static final String IDENTIFIER = "beam:schematransform:org.apache.beam:splunk_write:v1"; + static final String INPUT = "input"; + static final String OUTPUT = "output"; + static final String ERROR = "errors"; + public static final TupleTag ERROR_TAG = new TupleTag() {}; + public static final TupleTag EVENT_TAG = new TupleTag() {}; + + @Override + protected Class configurationClass() { + return SplunkWriteSchemaTransformConfiguration.class; + } + + /** Returns the expected {@link SchemaTransform} of the configuration. */ + @Override + protected SchemaTransform from(SplunkWriteSchemaTransformConfiguration configuration) { + return new SplunkWriteSchemaTransform(configuration); + } + + /** Implementation of the {@link TypedSchemaTransformProvider} identifier method. */ + @Override + public String identifier() { + return IDENTIFIER; + } + + /** Implementation of the {@link TypedSchemaTransformProvider} input collection names method. */ + @Override + public List inputCollectionNames() { + return Collections.singletonList(INPUT); + } + + /** Implementation of the {@link TypedSchemaTransformProvider} output collection names method. */ + @Override + public List outputCollectionNames() { + return Collections.singletonList(ERROR); + } + + /** + * An implementation of {@link SchemaTransform} for Splunk Write jobs configured using {@link + * SplunkWriteSchemaTransformConfiguration}. + */ + static class SplunkWriteSchemaTransform extends SchemaTransform { + private final SplunkWriteSchemaTransformConfiguration configuration; + + SplunkWriteSchemaTransform(SplunkWriteSchemaTransformConfiguration configuration) { + this.configuration = configuration; + } + + @Override + public PCollectionRowTuple expand(PCollectionRowTuple input) { + configuration.validate(); + + PCollection inputRows = input.get(INPUT); + boolean handleErrors = ErrorHandling.hasOutput(configuration.getErrorHandling()); + + Schema errorSchema = + Schema.builder() + .addNullableRowField("failed_row", inputRows.getSchema()) + .addNullableField("payload", Schema.FieldType.STRING) + .addNullableField("statusCode", Schema.FieldType.INT32) + .addNullableField("statusMessage", Schema.FieldType.STRING) + .build(); + + PCollectionTuple convertResult = + inputRows.apply( + "Convert to SplunkEvent", + ParDo.of(new RowToEventFn(handleErrors, ERROR_TAG, errorSchema)) + .withOutputTags(EVENT_TAG, TupleTagList.of(ERROR_TAG))); + + PCollection splunkEvents = + convertResult.get(EVENT_TAG).setCoder(SplunkEventCoder.of()); + PCollection conversionErrors = + convertResult.get(ERROR_TAG).setCoder(RowCoder.of(errorSchema)); + + SplunkIO.Write write = SplunkIO.write(configuration.getUrl(), configuration.getToken()); + Integer batchCount = configuration.getBatchCount(); + if (batchCount != null) { + write = write.withBatchCount(batchCount); + } + Integer parallelism = configuration.getParallelism(); + if (parallelism != null) { + write = write.withParallelism(parallelism); + } + Boolean disableCertificateValidation = configuration.getDisableCertificateValidation(); + if (disableCertificateValidation != null) { + write = write.withDisableCertificateValidation(disableCertificateValidation); + } + String rootCaCertificatePath = configuration.getRootCaCertificatePath(); + if (rootCaCertificatePath != null) { + write = write.withRootCaCertificatePath(rootCaCertificatePath); + } + Boolean enableBatchLogs = configuration.getEnableBatchLogs(); + if (enableBatchLogs != null) { + write = write.withEnableBatchLogs(enableBatchLogs); + } + Boolean enableGzipHttpCompression = configuration.getEnableGzipHttpCompression(); + if (enableGzipHttpCompression != null) { + write = write.withEnableGzipHttpCompression(enableGzipHttpCompression); + } + + PCollection writeErrors = splunkEvents.apply("Write To Splunk", write); + + ErrorHandling errorHandling = configuration.getErrorHandling(); + if (handleErrors && errorHandling != null) { + PCollection writeErrorRows = + writeErrors + .apply( + "Convert Write Errors to Rows", + MapElements.into(TypeDescriptors.rows()) + .via( + error -> + Row.withSchema(errorSchema) + .addValue(null) + .addValue(error.payload()) + .addValue(error.statusCode()) + .addValue(error.statusMessage()) + .build())) + .setCoder(RowCoder.of(errorSchema)); + + PCollection allErrors = + PCollectionList.of(conversionErrors) + .and(writeErrorRows) + .apply("Flatten Errors", Flatten.pCollections()) + .setCoder(RowCoder.of(errorSchema)); + + return PCollectionRowTuple.of(errorHandling.getOutput(), allErrors); + } else { + writeErrors.apply("Fail on Write Error", ParDo.of(new FailOnWriteErrorFn())); + PCollection emptyErrors = + input + .getPipeline() + .apply("Empty Errors Placeholder", Create.empty(RowCoder.of(errorSchema))); + return PCollectionRowTuple.of(ERROR, emptyErrors); + } + } + } + + static SplunkEvent rowToEvent(Row row) { + SplunkEvent.Builder builder = SplunkEvent.newBuilder(); + Schema schema = row.getSchema(); + + Long time = schema.hasField(SplunkEvent.TIME) ? row.getInt64(SplunkEvent.TIME) : null; + if (time != null) { + builder.withTime(time); + } + String host = schema.hasField(SplunkEvent.HOST) ? row.getString(SplunkEvent.HOST) : null; + if (host != null) { + builder.withHost(host); + } + String source = schema.hasField(SplunkEvent.SOURCE) ? row.getString(SplunkEvent.SOURCE) : null; + if (source != null) { + builder.withSource(source); + } + String sourceType = + schema.hasField(SplunkEvent.SOURCE_TYPE) ? row.getString(SplunkEvent.SOURCE_TYPE) : null; + if (sourceType != null) { + builder.withSourceType(sourceType); + } + String index = schema.hasField(SplunkEvent.INDEX) ? row.getString(SplunkEvent.INDEX) : null; + if (index != null) { + builder.withIndex(index); + } + Object fields = schema.hasField(SplunkEvent.FIELDS) ? row.getValue(SplunkEvent.FIELDS) : null; + if (fields instanceof Row) { + builder.withFields(toJsonElement(fields).getAsJsonObject()); + } else if (fields instanceof String) { + builder.withFields(JsonParser.parseString((String) fields).getAsJsonObject()); + } else if (fields != null) { + throw new IllegalArgumentException( + "The fields field must be a row or a JSON object string, but got: " + + schema.getField(SplunkEvent.FIELDS).getType()); + } + String event = schema.hasField(SplunkEvent.EVENT) ? row.getString(SplunkEvent.EVENT) : null; + builder.withEvent(checkNotNull(event, "Event is required.")); + + return builder.create(); + } + + private static JsonElement toJsonElement(Object value) { + if (value instanceof Row) { + Row row = (Row) value; + JsonObject json = new JsonObject(); + for (Schema.Field field : row.getSchema().getFields()) { + Object fieldValue = row.getValue(field.getName()); + if (fieldValue != null) { + json.add(field.getName(), toJsonElement(fieldValue)); + } + } + return json; + } else if (value instanceof Iterable) { + JsonArray json = new JsonArray(); + for (Object element : (Iterable) value) { + if (element != null) { + json.add(toJsonElement(element)); + } + } + return json; + } else if (value instanceof Number) { + return new JsonPrimitive((Number) value); + } else if (value instanceof Boolean) { + return new JsonPrimitive((Boolean) value); + } + return new JsonPrimitive(value.toString()); + } + + static class RowToEventFn extends DoFn { + private final boolean handleErrors; + private final TupleTag errorOutputTag; + private final Schema errorSchema; + + RowToEventFn(boolean handleErrors, TupleTag errorOutputTag, Schema errorSchema) { + this.handleErrors = handleErrors; + this.errorOutputTag = errorOutputTag; + this.errorSchema = errorSchema; + } + + @ProcessElement + public void processElement(ProcessContext c) { + try { + c.output(rowToEvent(c.element())); + } catch (Exception e) { + if (handleErrors) { + String rowString = c.element().toString(); + String payload = rowString.length() <= 1024 ? rowString : rowString.substring(0, 1024); + c.output( + errorOutputTag, + Row.withSchema(errorSchema) + .addValue(c.element()) + .addValue(payload) + .addValue(HttpURLConnection.HTTP_BAD_REQUEST) + .addValue(e.getMessage()) + .build()); + } else { + throw new RuntimeException(e); + } + } + } + } + + /** + * A {@link DoFn} that throws a {@link RuntimeException} when a write error is encountered, + * causing the pipeline to fail. This is the default error handling behavior when no error output + * is configured. + */ + static class FailOnWriteErrorFn extends DoFn { + @ProcessElement + public void processElement(@Element SplunkWriteError error) { + String message = error.statusMessage(); + if (error.statusCode() != null) { + throw new RuntimeException( + String.format( + "Splunk write failed with status code %d: %s", error.statusCode(), message)); + } else { + throw new RuntimeException("Splunk write failed: " + message); + } + } + } +} diff --git a/sdks/java/io/splunk/src/test/java/org/apache/beam/sdk/io/splunk/SplunkWriteErrorCoderTest.java b/sdks/java/io/splunk/src/test/java/org/apache/beam/sdk/io/splunk/SplunkWriteErrorCoderTest.java new file mode 100644 index 000000000000..0ac4ee81dea8 --- /dev/null +++ b/sdks/java/io/splunk/src/test/java/org/apache/beam/sdk/io/splunk/SplunkWriteErrorCoderTest.java @@ -0,0 +1,77 @@ +/* + * 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.splunk; + +import static org.junit.Assert.assertEquals; + +import java.io.ByteArrayInputStream; +import java.io.ByteArrayOutputStream; +import java.io.IOException; +import org.junit.Test; + +/** Unit tests for {@link SplunkWriteErrorCoder} class. */ +public class SplunkWriteErrorCoderTest { + + /** + * Test whether {@link SplunkWriteErrorCoder} is able to encode/decode a {@link SplunkWriteError} + * correctly. + * + * @throws IOException + */ + @Test + public void testEncodeDecode() throws IOException { + + SplunkWriteError expectedError = + SplunkWriteError.newBuilder() + .withPayload("test-payload") + .withStatusCode(404) + .withStatusMessage("test-message") + .create(); + + SplunkWriteErrorCoder coder = SplunkWriteErrorCoder.of(); + try (ByteArrayOutputStream bos = new ByteArrayOutputStream()) { + coder.encode(expectedError, bos); + try (ByteArrayInputStream bin = new ByteArrayInputStream(bos.toByteArray())) { + SplunkWriteError decodedError = coder.decode(bin); + assertEquals(expectedError, decodedError); + } + } + } + + /** + * Test whether {@link SplunkWriteErrorCoder} is able to encode/decode a {@link SplunkWriteError} + * with missing values. + * + * @throws IOException + */ + @Test + public void testEncodeDecodeNullValues() throws IOException { + + SplunkWriteError expectedError = + SplunkWriteError.newBuilder().withPayload("test-payload").create(); + + SplunkWriteErrorCoder coder = SplunkWriteErrorCoder.of(); + try (ByteArrayOutputStream bos = new ByteArrayOutputStream()) { + coder.encode(expectedError, bos); + try (ByteArrayInputStream bin = new ByteArrayInputStream(bos.toByteArray())) { + SplunkWriteError decodedError = coder.decode(bin); + assertEquals(expectedError, decodedError); + } + } + } +} diff --git a/sdks/java/io/splunk/src/test/java/org/apache/beam/sdk/io/splunk/SplunkWriteSchemaTransformProviderTest.java b/sdks/java/io/splunk/src/test/java/org/apache/beam/sdk/io/splunk/SplunkWriteSchemaTransformProviderTest.java new file mode 100644 index 000000000000..0e3624768267 --- /dev/null +++ b/sdks/java/io/splunk/src/test/java/org/apache/beam/sdk/io/splunk/SplunkWriteSchemaTransformProviderTest.java @@ -0,0 +1,459 @@ +/* + * 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.splunk; + +import static org.apache.beam.sdk.io.splunk.SplunkWriteSchemaTransformProvider.ERROR; +import static org.apache.beam.sdk.io.splunk.SplunkWriteSchemaTransformProvider.INPUT; +import static org.apache.beam.sdk.io.splunk.SplunkWriteSchemaTransformProvider.OUTPUT; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertNull; +import static org.junit.Assert.assertTrue; +import static org.junit.jupiter.api.Assertions.assertThrows; + +import com.google.gson.JsonArray; +import com.google.gson.JsonObject; +import java.util.Arrays; +import java.util.List; +import java.util.ServiceLoader; +import java.util.stream.Collectors; +import java.util.stream.StreamSupport; +import org.apache.beam.sdk.Pipeline.PipelineExecutionException; +import org.apache.beam.sdk.schemas.Schema; +import org.apache.beam.sdk.schemas.transforms.SchemaTransform; +import org.apache.beam.sdk.schemas.transforms.SchemaTransformProvider; +import org.apache.beam.sdk.schemas.transforms.providers.ErrorHandling; +import org.apache.beam.sdk.testing.NeedsRunner; +import org.apache.beam.sdk.testing.PAssert; +import org.apache.beam.sdk.testing.TestPipeline; +import org.apache.beam.sdk.transforms.Create; +import org.apache.beam.sdk.values.PCollection; +import org.apache.beam.sdk.values.PCollectionRowTuple; +import org.apache.beam.sdk.values.Row; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Joiner; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Lists; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Sets; +import org.junit.Rule; +import org.junit.Test; +import org.junit.experimental.categories.Category; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; +import org.mockserver.client.MockServerClient; +import org.mockserver.junit.MockServerRule; +import org.mockserver.model.HttpRequest; +import org.mockserver.model.HttpResponse; +import org.mockserver.verify.VerificationTimes; + +/** Unit tests for {@link SplunkWriteSchemaTransformProvider}. */ +@RunWith(JUnit4.class) +public class SplunkWriteSchemaTransformProviderTest { + + private static final String EXPECTED_PATH = "/" + HttpEventPublisher.HEC_URL_PATH; + + @Rule public final transient TestPipeline p = TestPipeline.create(); + + // We create a MockServerRule to simulate an actual Splunk HEC server. + @Rule public MockServerRule mockServerRule = new MockServerRule(this); + private MockServerClient mockServerClient; + + private static final Schema FIELDS_SCHEMA = + Schema.builder() + .addStringField("customfield") + .addInt64Field("count") + .addArrayField("tags", Schema.FieldType.STRING) + .build(); + + private static final Schema SCHEMA = + Schema.builder() + .addNullableField("time", Schema.FieldType.INT64) + .addStringField("host") + .addStringField("source") + .addStringField("sourcetype") + .addNullableField("index", Schema.FieldType.STRING) + .addNullableRowField("fields", FIELDS_SCHEMA) + .addStringField("event") + .build(); + + private static final List ROWS = + Arrays.asList( + Row.withSchema(SCHEMA) + .withFieldValue("time", 12345L) + .withFieldValue("host", "test-host-1") + .withFieldValue("source", "test-source-1") + .withFieldValue("sourcetype", "test-source-type-1") + .withFieldValue("index", "test-index-1") + .withFieldValue( + "fields", + Row.withSchema(FIELDS_SCHEMA) + .withFieldValue("customfield", "value") + .withFieldValue("count", 1L) + .withFieldValue("tags", Arrays.asList("a", "b")) + .build()) + .withFieldValue("event", "test-event-1") + .build(), + Row.withSchema(SCHEMA) + .withFieldValue("time", null) + .withFieldValue("host", "test-host-2") + .withFieldValue("source", "test-source-2") + .withFieldValue("sourcetype", "test-source-type-2") + .withFieldValue("index", null) + .withFieldValue("fields", null) + .withFieldValue("event", "test-event-2") + .build()); + + private static List expectedEvents() { + JsonArray tags = new JsonArray(); + tags.add("a"); + tags.add("b"); + JsonObject fields = new JsonObject(); + fields.addProperty("customfield", "value"); + fields.addProperty("count", 1L); + fields.add("tags", tags); + return Arrays.asList( + SplunkEvent.newBuilder() + .withTime(12345L) + .withHost("test-host-1") + .withSource("test-source-1") + .withSourceType("test-source-type-1") + .withIndex("test-index-1") + .withFields(fields) + .withEvent("test-event-1") + .create(), + SplunkEvent.newBuilder() + .withHost("test-host-2") + .withSource("test-source-2") + .withSourceType("test-source-type-2") + .withEvent("test-event-2") + .create()); + } + + private String mockServerUrl() { + return Joiner.on(':').join("http://localhost", mockServerRule.getPort()); + } + + private void mockServerListening(int statusCode) { + mockServerClient + .when(HttpRequest.request(EXPECTED_PATH)) + .respond(HttpResponse.response().withStatusCode(statusCode)); + } + + @Test + public void testWriteInvalidConfigurations() { + // token not set + assertThrows( + IllegalStateException.class, + () -> SplunkWriteSchemaTransformConfiguration.builder().setUrl(mockServerUrl()).build()); + + // url not set + assertThrows( + IllegalStateException.class, + () -> SplunkWriteSchemaTransformConfiguration.builder().setToken("test-token").build()); + + assertThrows( + IllegalArgumentException.class, + () -> + SplunkWriteSchemaTransformConfiguration.builder() + .setUrl("") + .setToken("test-token") + .build() + .validate()); + + assertThrows( + IllegalArgumentException.class, + () -> + SplunkWriteSchemaTransformConfiguration.builder() + .setUrl(mockServerUrl()) + .setToken("test-token") + .setBatchCount(0) + .build() + .validate()); + + assertThrows( + IllegalArgumentException.class, + () -> + SplunkWriteSchemaTransformConfiguration.builder() + .setUrl(mockServerUrl()) + .setToken("test-token") + .setParallelism(0) + .build() + .validate()); + } + + @Test + public void testWriteBuildTransformWithCorrectFields() { + ServiceLoader serviceLoader = + ServiceLoader.load(SchemaTransformProvider.class); + List providers = + StreamSupport.stream(serviceLoader.spliterator(), false) + .filter(provider -> provider.getClass() == SplunkWriteSchemaTransformProvider.class) + .collect(Collectors.toList()); + SchemaTransformProvider splunkProvider = providers.get(0); + assertEquals(splunkProvider.outputCollectionNames(), Lists.newArrayList(ERROR)); + + assertEquals( + Sets.newHashSet( + "url", + "token", + "batch_count", + "parallelism", + "disable_certificate_validation", + "root_ca_certificate_path", + "enable_batch_logs", + "enable_gzip_http_compression", + "error_handling"), + splunkProvider.configurationSchema().getFields().stream() + .map(field -> field.getName()) + .collect(Collectors.toSet())); + } + + @Test + public void testRowToSplunkEvent() { + List events = expectedEvents(); + for (int i = 0; i < ROWS.size(); i++) { + assertEquals(events.get(i), SplunkWriteSchemaTransformProvider.rowToEvent(ROWS.get(i))); + } + } + + @Test + public void testRowToSplunkEventWithOnlyEvent() { + Schema schema = Schema.builder().addStringField("event").build(); + Row row = Row.withSchema(schema).withFieldValue("event", "test-event").build(); + + assertEquals( + SplunkEvent.newBuilder().withEvent("test-event").create(), + SplunkWriteSchemaTransformProvider.rowToEvent(row)); + } + + @Test + public void testRowToSplunkEventWithJsonStringFields() { + Schema schema = Schema.builder().addStringField("fields").addStringField("event").build(); + Row row = + Row.withSchema(schema) + .withFieldValue("fields", "{\"customfield\": \"value\"}") + .withFieldValue("event", "test-event") + .build(); + + JsonObject fields = new JsonObject(); + fields.addProperty("customfield", "value"); + assertEquals( + SplunkEvent.newBuilder().withFields(fields).withEvent("test-event").create(), + SplunkWriteSchemaTransformProvider.rowToEvent(row)); + } + + @Test + public void testRowToSplunkEventWithExtraFields_DiscardsExtraFields() { + Schema schema = + Schema.builder() + .addStringField("host") + .addStringField("event") + .addStringField("extra_field") + .build(); + Row row = + Row.withSchema(schema) + .withFieldValue("host", "test-host") + .withFieldValue("event", "test-event") + .withFieldValue("extra_field", "extra_value") + .build(); + + assertEquals( + SplunkEvent.newBuilder().withHost("test-host").withEvent("test-event").create(), + SplunkWriteSchemaTransformProvider.rowToEvent(row)); + } + + @Test(expected = NullPointerException.class) + public void testRowToSplunkEventWithNullEvent() { + Schema schema = + Schema.builder() + .addStringField("host") + .addNullableField("event", Schema.FieldType.STRING) + .build(); + Row row = + Row.withSchema(schema) + .withFieldValue("host", "test-host") + .withFieldValue("event", null) + .build(); + + SplunkWriteSchemaTransformProvider.rowToEvent(row); + } + + @Test(expected = IllegalArgumentException.class) + public void testRowToSplunkEventWithWrongFieldsType() { + Schema schema = Schema.builder().addInt64Field("fields").addStringField("event").build(); + Row row = + Row.withSchema(schema) + .withFieldValue("fields", 1L) + .withFieldValue("event", "test-event") + .build(); + + SplunkWriteSchemaTransformProvider.rowToEvent(row); + } + + @Test + @Category(NeedsRunner.class) + public void testSuccessfulWrite() { + mockServerListening(200); + + SplunkWriteSchemaTransformConfiguration configuration = + SplunkWriteSchemaTransformConfiguration.builder() + .setUrl(mockServerUrl()) + .setToken("test-token") + .setBatchCount(ROWS.size()) + .setParallelism(1) + .setErrorHandling(ErrorHandling.builder().setOutput(OUTPUT).build()) + .build(); + SchemaTransform transform = new SplunkWriteSchemaTransformProvider().from(configuration); + + PCollection input = p.apply(Create.of(ROWS).withRowSchema(SCHEMA)); + PCollectionRowTuple output = transform.expand(PCollectionRowTuple.of(INPUT, input)); + + PAssert.that(output.get(OUTPUT)).empty(); + + p.run(); + + mockServerClient.verify(HttpRequest.request(EXPECTED_PATH), VerificationTimes.once()); + } + + @Test + @Category(NeedsRunner.class) + public void testWriteErrorsWithErrorHandling() { + mockServerListening(404); + + SplunkWriteSchemaTransformConfiguration configuration = + SplunkWriteSchemaTransformConfiguration.builder() + .setUrl(mockServerUrl()) + .setToken("test-token") + .setBatchCount(ROWS.size()) + .setParallelism(1) + .setErrorHandling(ErrorHandling.builder().setOutput(OUTPUT).build()) + .build(); + SchemaTransform transform = new SplunkWriteSchemaTransformProvider().from(configuration); + + PCollection input = p.apply(Create.of(ROWS).withRowSchema(SCHEMA)); + PCollectionRowTuple output = transform.expand(PCollectionRowTuple.of(INPUT, input)); + + PAssert.that(output.get(OUTPUT)) + .satisfies( + errors -> { + int count = 0; + for (Row error : errors) { + count++; + assertNull(error.getRow("failed_row")); + assertEquals((Integer) 404, error.getInt32("statusCode")); + assertTrue(error.getString("payload").contains("test-event-")); + } + assertEquals(ROWS.size(), count); + return null; + }); + + p.run(); + } + + @Test + @Category(NeedsRunner.class) + public void testConversionErrorsWithErrorHandling() { + mockServerListening(200); + + SplunkWriteSchemaTransformConfiguration configuration = + SplunkWriteSchemaTransformConfiguration.builder() + .setUrl(mockServerUrl()) + .setToken("test-token") + .setErrorHandling(ErrorHandling.builder().setOutput(OUTPUT).build()) + .build(); + SchemaTransform transform = new SplunkWriteSchemaTransformProvider().from(configuration); + + Schema schema = Schema.builder().addStringField("host").build(); + Row row = Row.withSchema(schema).withFieldValue("host", "test-host").build(); + + PCollection input = p.apply(Create.of(row).withRowSchema(schema)); + PCollectionRowTuple output = transform.expand(PCollectionRowTuple.of(INPUT, input)); + + PAssert.that(output.get(OUTPUT)) + .satisfies( + errors -> { + Row error = errors.iterator().next(); + assertEquals(1, errors.spliterator().getExactSizeIfKnown()); + assertEquals(row, error.getRow("failed_row")); + assertEquals(row.toString(), error.getString("payload")); + assertEquals( + (Integer) java.net.HttpURLConnection.HTTP_BAD_REQUEST, + error.getInt32("statusCode")); + assertTrue(error.getString("statusMessage").contains("Event is required.")); + return null; + }); + + p.run(); + + mockServerClient.verify(HttpRequest.request(EXPECTED_PATH), VerificationTimes.exactly(0)); + } + + @Test + @Category(NeedsRunner.class) + public void testWriteErrorFailsPipelineWithoutErrorHandling() { + mockServerListening(404); + + SplunkWriteSchemaTransformConfiguration configuration = + SplunkWriteSchemaTransformConfiguration.builder() + .setUrl(mockServerUrl()) + .setToken("test-token") + .setBatchCount(ROWS.size()) + .build(); + SchemaTransform transform = new SplunkWriteSchemaTransformProvider().from(configuration); + + PCollection input = p.apply(Create.of(ROWS).withRowSchema(SCHEMA)); + PCollectionRowTuple output = transform.expand(PCollectionRowTuple.of(INPUT, input)); + assertEquals(1, output.getAll().size()); + assertTrue(output.has(ERROR)); + + assertThrows(PipelineExecutionException.class, () -> p.run().waitUntilFinish()); + } + + @Test + @Category(NeedsRunner.class) + public void testBuildTransformFromRowConfiguration() { + mockServerListening(200); + + SplunkWriteSchemaTransformProvider provider = new SplunkWriteSchemaTransformProvider(); + Schema configSchema = provider.configurationSchema(); + Schema errorHandlingSchema = configSchema.getField("error_handling").getType().getRowSchema(); + + Row configRow = + Row.withSchema(configSchema) + .withFieldValue("url", mockServerUrl()) + .withFieldValue("token", "test-token") + .withFieldValue("batch_count", 1) + .withFieldValue("parallelism", 1) + .withFieldValue("disable_certificate_validation", false) + .withFieldValue("root_ca_certificate_path", null) + .withFieldValue("enable_batch_logs", true) + .withFieldValue("enable_gzip_http_compression", false) + .withFieldValue( + "error_handling", + Row.withSchema(errorHandlingSchema).withFieldValue("output", OUTPUT).build()) + .build(); + SchemaTransform transform = provider.from(configRow); + + PCollection input = p.apply(Create.of(ROWS).withRowSchema(SCHEMA)); + PCollectionRowTuple output = transform.expand(PCollectionRowTuple.of(INPUT, input)); + + PAssert.that(output.get(OUTPUT)).empty(); + + p.run(); + + mockServerClient.verify( + HttpRequest.request(EXPECTED_PATH), VerificationTimes.exactly(ROWS.size())); + } +} diff --git a/sdks/python/apache_beam/yaml/integration_tests.py b/sdks/python/apache_beam/yaml/integration_tests.py index becfaa793032..435a5b88af6a 100644 --- a/sdks/python/apache_beam/yaml/integration_tests.py +++ b/sdks/python/apache_beam/yaml/integration_tests.py @@ -41,6 +41,7 @@ import requests from apache_beam.coders import Coder +from apache_beam.coders import VarIntCoder from apache_beam.coders.coder_impl import CoderImpl @@ -63,12 +64,20 @@ def get_impl(self): return BigEndianIntegerCoderImpl() +def _java_serialized_coder(payload, components, context): + # Java's VarIntCoder (used to key SplunkIO's stateful writer) shares the same + # fallback URN. The payload is the snappy-compressed serialized Java coder, + # and the start of its class name stays uncompressed, so it can be used to + # tell the two apart. + if b'VarInt' in payload: + return VarIntCoder() + return BigEndianIntegerCoder() + + # Register the coder with the fallback URN used by the Java SDK for this coder. # This allows the Python FnApiRunner to handle data sharded by Java transforms -# using BigEndianIntegerCoder in integration tests. -Coder.register_urn( - 'beam:coders:javasdk:0.1', - None, lambda payload, components, context: BigEndianIntegerCoder()) +# using BigEndianIntegerCoder or VarIntCoder in integration tests. +Coder.register_urn('beam:coders:javasdk:0.1', None, _java_serialized_coder) import psycopg2 import pyarrow as pa diff --git a/sdks/python/apache_beam/yaml/standard_io.yaml b/sdks/python/apache_beam/yaml/standard_io.yaml index 5176907d2fc9..62cca7d75bc9 100644 --- a/sdks/python/apache_beam/yaml/standard_io.yaml +++ b/sdks/python/apache_beam/yaml/standard_io.yaml @@ -734,6 +734,29 @@ config: gradle_target: 'sdks:java:io:expansion-service:shadowJar' +#Splunk +- type: renaming + transforms: + 'WriteToSplunk': 'WriteToSplunk' + config: + mappings: + 'WriteToSplunk': + 'url': 'url' + 'token': 'token' + 'batch_count': 'batch_count' + 'parallelism': 'parallelism' + 'disable_certificate_validation': 'disable_certificate_validation' + 'root_ca_certificate_path': 'root_ca_certificate_path' + 'enable_batch_logs': 'enable_batch_logs' + 'enable_gzip_http_compression': 'enable_gzip_http_compression' + 'error_handling': 'error_handling' + underlying_provider: + type: beamJar + transforms: + 'WriteToSplunk': 'beam:schematransform:org.apache.beam:splunk_write:v1' + config: + gradle_target: 'sdks:java:io:expansion-service:shadowJar' + #MongoDB - type: renaming transforms: diff --git a/sdks/python/apache_beam/yaml/test_utils/splunk_test_utils.py b/sdks/python/apache_beam/yaml/test_utils/splunk_test_utils.py new file mode 100644 index 000000000000..4b4d7f2b9c09 --- /dev/null +++ b/sdks/python/apache_beam/yaml/test_utils/splunk_test_utils.py @@ -0,0 +1,151 @@ +# +# 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. +# + +"""Helper utilities for Splunk integration tests.""" + +import contextlib +import gzip +import http.server +import io +import json +import logging +import threading + +_LOGGER = logging.getLogger(__name__) + +_HEC_PATH = "/services/collector/event" + + +class SplunkConnection: + def __init__(self, url, token): + self.url = url + self.token = token + + +def _parse_hec_payload(payload): + """Parses a HEC request body, which is a sequence of JSON objects that are + concatenated rather than wrapped in a list.""" + decoder = json.JSONDecoder() + records = [] + pos = 0 + while True: + while pos < len(payload) and payload[pos].isspace(): + pos += 1 + if pos >= len(payload): + return records + record, pos = decoder.raw_decode(payload, pos) + records.append(record) + + +class MockSplunkHandler(http.server.BaseHTTPRequestHandler): + def do_POST(self): + if self.path != _HEC_PATH: + self.send_response(404) + self.end_headers() + return + + if self.headers.get('Authorization') != f"Splunk {self.server.token}": + self.send_response(403) + self.send_header('Content-Type', 'application/json') + self.end_headers() + self.wfile.write(b'{"text": "Invalid token", "code": 4}') + return + + is_chunked = self.headers.get('Transfer-Encoding', '').lower() == 'chunked' + is_gzip = self.headers.get('Content-Encoding', '').lower() == 'gzip' + content_len = int(self.headers.get('Content-Length', 0)) + + try: + raw_data = b'' + if is_chunked: + while True: + line = self.rfile.readline().strip() + if not line: + break + chunk_len = int(line, 16) + if chunk_len == 0: + self.rfile.readline() # Clear trail + break + raw_data += self.rfile.read(chunk_len) + self.rfile.readline() # Clear trail + elif content_len > 0: + raw_data = self.rfile.read(content_len) + + if raw_data and is_gzip: + with gzip.GzipFile(fileobj=io.BytesIO(raw_data)) as f: + raw_data = f.read() + + if raw_data: + records = _parse_hec_payload(raw_data.decode('utf-8')) + with self.server.record_lock: + self.server.received_records.extend(records) + except Exception as e: + logging.error("CRITICAL: Failure unpacking mock splunk payload: %s", e) + + self.send_response(200) + self.send_header('Content-Type', 'application/json') + self.end_headers() + self.wfile.write(b'{"text": "Success", "code": 0}') + + def log_message(self, format, *args): + pass + + +@contextlib.contextmanager +def temp_splunk_mock_server(received_records, token): + server = http.server.ThreadingHTTPServer(('localhost', 0), MockSplunkHandler) + server.received_records = received_records + server.record_lock = threading.Lock() + server.token = token + ip, port = server.server_address + thread = threading.Thread(target=server.serve_forever) + thread.daemon = True + thread.start() + try: + yield f"http://{ip}:{port}" + finally: + server.shutdown() + server.server_close() + thread.join() + + +@contextlib.contextmanager +def temp_fake_splunk_server(expected_records=None): + """Context manager to provide a temporary fake Splunk HEC server for testing. + """ + received = [] + token = "dummy_token_for_testing" + with temp_splunk_mock_server(received, token) as mock_url: + try: + yield SplunkConnection(url=mock_url, token=token) + except Exception as err: + logging.error( + "Error interacting with temporary fake Splunk server: %s", err) + raise err + finally: + if expected_records is not None: + + canonicalize = lambda rec: json.dumps(rec, sort_keys=True) + + actual_strs = sorted([canonicalize(r) for r in received]) + expected_strs = sorted([canonicalize(e) for e in expected_records]) + + assert actual_strs == expected_strs, ( + f"Mismatch in recorded Splunk events!\n" + f"Expected: {expected_strs}\n" + f"Actual: {actual_strs}" + ) diff --git a/sdks/python/apache_beam/yaml/tests/splunk.yaml b/sdks/python/apache_beam/yaml/tests/splunk.yaml new file mode 100644 index 000000000000..ea06d92c6ea0 --- /dev/null +++ b/sdks/python/apache_beam/yaml/tests/splunk.yaml @@ -0,0 +1,60 @@ +# +# 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. +# + +fixtures: + - name: FAKE_SPLUNK_SERVER + type: "apache_beam.yaml.test_utils.splunk_test_utils.temp_fake_splunk_server" + config: + expected_records: + - { time: 1700000000, host: "test-host", source: "apache-beam1", sourcetype: "test-type", index: "test-index", event: "Event for label 11a" } + - { host: "test-host", source: "apache-beam2", sourcetype: "test-type", index: "test-index", event: "Event for label 37a" } + - { host: "test-host", source: "apache-beam3", sourcetype: "test-type", event: "Event for label 389a" } + - { host: "test-host", source: "apache-beam4", event: "Event for label 3821b" } + +pipelines: + - pipeline: + type: composite + transforms: + - type: Create + config: + elements: + - { time: 1700000000, host: "test-host", source: "apache-beam1", sourcetype: "test-type", index: "test-index", event: "Event for label 11a" } + - { host: "test-host", source: "apache-beam2", sourcetype: "test-type", index: "test-index", event: "Event for label 37a" } + - { host: "test-host", source: "apache-beam3", sourcetype: "test-type", event: "Event for label 389a" } + - { host: "test-host", source: "apache-beam4", event: "Event for label 3821b" } + - { host: "test-host", source: "apache-beam-broken" } # Triggers mandatory field failure + - type: WriteToSplunk + input: Create + config: + url: "{FAKE_SPLUNK_SERVER.url}" + token: "{FAKE_SPLUNK_SERVER.token}" + batch_count: 2 + parallelism: 1 + error_handling: + output: error_output + - type: MapToFields + input: WriteToSplunk.error_output + config: + language: python + fields: + failed_source: "failed_row.source" + - type: AssertEqual + input: MapToFields + config: + elements: + - { failed_source: "apache-beam-broken" } + # Asserting good records is taken care of by the fixture diff --git a/sdks/standard_external_transforms.yaml b/sdks/standard_external_transforms.yaml index d5401c92ec5e..4421bf66d4aa 100644 --- a/sdks/standard_external_transforms.yaml +++ b/sdks/standard_external_transforms.yaml @@ -149,6 +149,51 @@ type: str identifier: beam:schematransform:org.apache.beam:mongodb_write:v1 name: MongodbWrite +- default_service: sdks:java:io:expansion-service:shadowJar + description: '' + destinations: + python: apache_beam/io + fields: + - description: The number of events to batch together for each write. + name: batch_count + nullable: true + type: int32 + - description: Whether to disable SSL certificate validation, e.g. for self-signed + certificates. + name: disable_certificate_validation + nullable: true + type: boolean + - description: Whether to log the result of each batch write. + name: enable_batch_logs + nullable: true + type: boolean + - description: Whether requests sent to the HEC endpoint should be GZIP encoded. + Defaults to true. + name: enable_gzip_http_compression + nullable: true + type: boolean + - description: Specifies how to handle errors. + name: error_handling + nullable: true + type: Row(output=) + - description: The number of parallel requests to the HEC endpoint. + name: parallelism + nullable: true + type: int32 + - description: Path to a root CA certificate used to validate the HEC endpoint. + name: root_ca_certificate_path + nullable: true + type: str + - description: The Splunk HEC authentication token. + name: token + nullable: false + type: str + - description: The Splunk HEC endpoint URL, e.g. https://splunk-host:8088. + name: url + nullable: false + type: str + identifier: beam:schematransform:org.apache.beam:splunk_write:v1 + name: SplunkWrite - default_service: sdks:java:io:expansion-service:shadowJar description: Reads records from an Amazon Kinesis stream and outputs Beam Rows with `data` (bytes) plus stream metadata fields.