diff --git a/sdks/java/build-tools/src/main/resources/beam/checkstyle/suppressions.xml b/sdks/java/build-tools/src/main/resources/beam/checkstyle/suppressions.xml index 833fd30cf218..474cf3d93a2a 100644 --- a/sdks/java/build-tools/src/main/resources/beam/checkstyle/suppressions.xml +++ b/sdks/java/build-tools/src/main/resources/beam/checkstyle/suppressions.xml @@ -23,11 +23,11 @@ - + - + @@ -43,6 +43,7 @@ + diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/BigtableIO.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/BigtableIO.java index 02236469e325..448d7cffdfe4 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/BigtableIO.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/BigtableIO.java @@ -2258,7 +2258,7 @@ public ReadChangeStream withStartTime(Instant startTime) { /** Used only for integration tests. Unsafe to use in production. */ @VisibleForTesting - ReadChangeStream withEndTime(Instant endTime) { + public ReadChangeStream withEndTime(Instant endTime) { return toBuilder().setEndTime(endTime).build(); } diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/BigtableChangeStreamReadSchemaTransformProvider.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/BigtableChangeStreamReadSchemaTransformProvider.java new file mode 100644 index 000000000000..1137a255e71b --- /dev/null +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/BigtableChangeStreamReadSchemaTransformProvider.java @@ -0,0 +1,501 @@ +/* + * 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.gcp.bigtable.changestreams; + +import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkArgument; + +import com.google.auto.service.AutoService; +import com.google.auto.value.AutoValue; +import com.google.cloud.bigtable.data.v2.models.AddToCell; +import com.google.cloud.bigtable.data.v2.models.ChangeStreamMutation; +import com.google.cloud.bigtable.data.v2.models.DeleteCells; +import com.google.cloud.bigtable.data.v2.models.DeleteFamily; +import com.google.cloud.bigtable.data.v2.models.Entry; +import com.google.cloud.bigtable.data.v2.models.MergeToCell; +import com.google.cloud.bigtable.data.v2.models.Range; +import com.google.cloud.bigtable.data.v2.models.SetCell; +import com.google.cloud.bigtable.data.v2.models.Value; +import com.google.protobuf.ByteString; +import java.io.Serializable; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import org.apache.beam.sdk.io.gcp.bigtable.BigtableIO; +import org.apache.beam.sdk.schemas.AutoValueSchema; +import org.apache.beam.sdk.schemas.Schema; +import org.apache.beam.sdk.schemas.annotations.DefaultSchema; +import org.apache.beam.sdk.schemas.annotations.SchemaFieldDescription; +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.transforms.DoFn; +import org.apache.beam.sdk.transforms.ParDo; +import org.apache.beam.sdk.values.KV; +import org.apache.beam.sdk.values.PCollection; +import org.apache.beam.sdk.values.PCollectionRowTuple; +import org.apache.beam.sdk.values.Row; +import org.checkerframework.checker.nullness.qual.Nullable; +import org.joda.time.Instant; + +/** + * A {@link TypedSchemaTransformProvider} for reading Cloud Bigtable change streams. + * + *

Internal only. + */ +@AutoService(SchemaTransformProvider.class) +public class BigtableChangeStreamReadSchemaTransformProvider + extends TypedSchemaTransformProvider< + BigtableChangeStreamReadSchemaTransformProvider.BigtableChangeStreamReadConfiguration> { + + private static final String OUTPUT_TAG = "output"; + + static final Schema VALUE_SCHEMA = + Schema.builder() + .addStringField("type") + .addNullableInt64Field("int_value") + .addNullableInt64Field("raw_timestamp_micros") + .addNullableByteArrayField("raw_value") + .build(); + + static final Schema TIMESTAMP_RANGE_SCHEMA = + Schema.builder() + .addStringField("start_bound") + .addNullableInt64Field("start_timestamp_micros") + .addStringField("end_bound") + .addNullableInt64Field("end_timestamp_micros") + .build(); + + static final Schema ENTRY_SCHEMA = + Schema.builder() + .addStringField("type") + .addStringField("family_name") + .addNullableByteArrayField("column_qualifier") + .addNullableInt64Field("timestamp_micros") + .addNullableByteArrayField("value") + .addNullableRowField("timestamp_range", TIMESTAMP_RANGE_SCHEMA) + .addNullableRowField("value_qualifier", VALUE_SCHEMA) + .addNullableRowField("value_timestamp", VALUE_SCHEMA) + .addNullableRowField("value_input", VALUE_SCHEMA) + .build(); + + static final Schema CHANGE_STREAM_MUTATION_SCHEMA = + Schema.builder() + .addByteArrayField("key") + .addStringField("mutation_type") + .addStringField("source_cluster_id") + .addDateTimeField("commit_timestamp") + .addInt32Field("tie_breaker") + .addStringField("token") + .addDateTimeField("estimated_low_watermark") + .addArrayField("entries", Schema.FieldType.row(ENTRY_SCHEMA)) + .build(); + + static final Schema FLATTENED_CHANGE_STREAM_MUTATION_SCHEMA = + Schema.builder() + .addByteArrayField("key") + .addStringField("mutation_type") + .addStringField("source_cluster_id") + .addDateTimeField("commit_timestamp") + .addInt32Field("tie_breaker") + .addStringField("token") + .addDateTimeField("estimated_low_watermark") + .addStringField("type") + .addStringField("family_name") + .addNullableByteArrayField("column_qualifier") + .addNullableInt64Field("timestamp_micros") + .addNullableByteArrayField("value") + .addNullableRowField("timestamp_range", TIMESTAMP_RANGE_SCHEMA) + .addNullableRowField("value_qualifier", VALUE_SCHEMA) + .addNullableRowField("value_timestamp", VALUE_SCHEMA) + .addNullableRowField("value_input", VALUE_SCHEMA) + .build(); + + @Override + protected Class configurationClass() { + return BigtableChangeStreamReadConfiguration.class; + } + + @Override + protected SchemaTransform from(BigtableChangeStreamReadConfiguration configuration) { + return new BigtableChangeStreamReadSchemaTransform(configuration); + } + + @Override + public String identifier() { + return "beam:schematransform:org.apache.beam:bigtable_cdc_read:v1"; + } + + @Override + public String description() { + return "Reads change stream records from a Google Cloud Bigtable table.\n" + + "The transform requires the project ID, instance ID, and table ID parameters.\n" + + "Optionally, an app profile, start timestamp, change stream name, and flatten mode can be provided.\n" + + "By default, the output is flattened to one row per change stream entry.\n" + + "Example usage:\n" + + " - type: ReadFromBigtableCDC\n" + + " config:\n" + + " project: \"my-gcp-project\"\n" + + " instance: \"my-bigtable-instance\"\n" + + " table: \"my-table\"\n" + + " flatten: true\n"; + } + + @Override + public List inputCollectionNames() { + return Collections.emptyList(); + } + + @Override + public List outputCollectionNames() { + return Collections.singletonList(OUTPUT_TAG); + } + + /** Configuration for reading a Bigtable change stream. */ + @DefaultSchema(AutoValueSchema.class) + @AutoValue + public abstract static class BigtableChangeStreamReadConfiguration implements Serializable { + + public void validate() { + checkArgument(!getTableId().isEmpty(), "Bigtable table ID must not be empty."); + checkArgument(!getInstanceId().isEmpty(), "Bigtable instance ID must not be empty."); + checkArgument(!getProjectId().isEmpty(), "Bigtable project ID must not be empty."); + } + + @SchemaFieldDescription("Google Cloud project ID containing the Bigtable instance.") + public abstract String getProjectId(); + + @SchemaFieldDescription("Bigtable instance ID to connect to.") + public abstract String getInstanceId(); + + @SchemaFieldDescription("Bigtable table ID whose change stream should be read.") + public abstract String getTableId(); + + @SchemaFieldDescription("Bigtable app profile used to read the change stream.") + public abstract @Nullable String getAppProfileId(); + + @SchemaFieldDescription("Timestamp from which to start reading the change stream.") + public abstract @Nullable String getStartAtTimestamp(); + + @SchemaFieldDescription("Name used to identify the Bigtable change stream pipeline.") + public abstract @Nullable String getChangeStreamName(); + + @SchemaFieldDescription( + "Timestamp at which to stop reading the change stream. Intended for integration testing.") + public abstract @Nullable String getEndAtTimestamp(); + + @SchemaFieldDescription( + "If true or omitted, outputs one row per change stream entry. " + + "If false, outputs one row per mutation with nested entries.") + public abstract @Nullable Boolean getFlatten(); + + public static Builder builder() { + return new AutoValue_BigtableChangeStreamReadSchemaTransformProvider_BigtableChangeStreamReadConfiguration + .Builder() + .setFlatten(true); + } + + /** Builder for {@link BigtableChangeStreamReadConfiguration}. */ + @AutoValue.Builder + public abstract static class Builder { + + public abstract Builder setProjectId(String projectId); + + public abstract Builder setInstanceId(String instanceId); + + public abstract Builder setTableId(String tableId); + + public abstract Builder setAppProfileId(String appProfileId); + + public abstract Builder setStartAtTimestamp(String startAtTimestamp); + + public abstract Builder setChangeStreamName(String changeStreamName); + + public abstract Builder setEndAtTimestamp(@Nullable String endAtTimestamp); + + public abstract Builder setFlatten(@Nullable Boolean flatten); + + public abstract BigtableChangeStreamReadConfiguration build(); + } + } + + /** SchemaTransform implementation for Bigtable change stream reads. */ + private static class BigtableChangeStreamReadSchemaTransform extends SchemaTransform { + + private final BigtableChangeStreamReadConfiguration configuration; + + BigtableChangeStreamReadSchemaTransform(BigtableChangeStreamReadConfiguration configuration) { + configuration.validate(); + this.configuration = configuration; + } + + @Override + public PCollectionRowTuple expand(PCollectionRowTuple input) { + checkArgument( + input.getAll().isEmpty(), + String.format( + "Input to %s is expected to be empty, but is not.", getClass().getSimpleName())); + + BigtableIO.ReadChangeStream readChangeStream = + BigtableIO.readChangeStream() + .withProjectId(configuration.getProjectId()) + .withInstanceId(configuration.getInstanceId()) + .withTableId(configuration.getTableId()); + + @Nullable String appProfileId = configuration.getAppProfileId(); + if (appProfileId != null) { + readChangeStream = readChangeStream.withAppProfileId(appProfileId); + } + + @Nullable String startAtTimestamp = configuration.getStartAtTimestamp(); + if (startAtTimestamp != null) { + readChangeStream = readChangeStream.withStartTime(Instant.parse(startAtTimestamp)); + } + + @Nullable String endAtTimestamp = configuration.getEndAtTimestamp(); + if (endAtTimestamp != null) { + readChangeStream = readChangeStream.withEndTime(Instant.parse(endAtTimestamp)); + } + + @Nullable String changeStreamName = configuration.getChangeStreamName(); + if (changeStreamName != null) { + readChangeStream = readChangeStream.withChangeStreamName(changeStreamName); + } + + PCollection> mutations = + input.getPipeline().apply(readChangeStream); + + PCollection rows; + + if (Boolean.FALSE.equals(configuration.getFlatten())) { + rows = + mutations + .apply("ConvertToBeamRows", ParDo.of(new ChangeStreamMutationToRowDoFn())) + .setRowSchema(CHANGE_STREAM_MUTATION_SCHEMA); + } else { + rows = + mutations + .apply( + "ConvertToFlattenedBeamRows", + ParDo.of(new ChangeStreamMutationToFlattenedRowsDoFn())) + .setRowSchema(FLATTENED_CHANGE_STREAM_MUTATION_SCHEMA); + } + + return PCollectionRowTuple.of(OUTPUT_TAG, rows); + } + } + + private static class ChangeStreamMutationToRowDoFn + extends DoFn, Row> { + + @ProcessElement + public void processElement( + @Element KV element, OutputReceiver out) { + out.output(mutationToRow(element.getValue())); + } + } + + private static class ChangeStreamMutationToFlattenedRowsDoFn + extends DoFn, Row> { + + @ProcessElement + public void processElement( + @Element KV element, OutputReceiver out) { + + ChangeStreamMutation mutation = element.getValue(); + + for (Entry entry : mutation.getEntries()) { + out.output(flattenedMutationEntryToRow(mutation, entry)); + } + } + } + + static Row mutationToRow(ChangeStreamMutation mutation) { + List entries = new ArrayList<>(); + + for (Entry entry : mutation.getEntries()) { + entries.add(entryToRow(entry)); + } + + return Row.withSchema(CHANGE_STREAM_MUTATION_SCHEMA) + .addValue(mutation.getRowKey().toByteArray()) + .addValue(mutation.getType().name()) + .addValue(mutation.getSourceClusterId()) + .addValue(new Instant(mutation.getCommitTime().toEpochMilli())) + .addValue(mutation.getTieBreaker()) + .addValue(mutation.getToken()) + .addValue(new Instant(mutation.getEstimatedLowWatermarkTime().toEpochMilli())) + .addValue(entries) + .build(); + } + + static Row flattenedMutationEntryToRow(ChangeStreamMutation mutation, Entry entry) { + + Row entryRow = entryToRow(entry); + + return Row.withSchema(FLATTENED_CHANGE_STREAM_MUTATION_SCHEMA) + .addValue(mutation.getRowKey().toByteArray()) + .addValue(mutation.getType().name()) + .addValue(mutation.getSourceClusterId()) + .addValue(new Instant(mutation.getCommitTime().toEpochMilli())) + .addValue(mutation.getTieBreaker()) + .addValue(mutation.getToken()) + .addValue(new Instant(mutation.getEstimatedLowWatermarkTime().toEpochMilli())) + .addValue(entryRow.getString("type")) + .addValue(entryRow.getString("family_name")) + .addValue(entryRow.getBytes("column_qualifier")) + .addValue(entryRow.getInt64("timestamp_micros")) + .addValue(entryRow.getBytes("value")) + .addValue(entryRow.getRow("timestamp_range")) + .addValue(entryRow.getRow("value_qualifier")) + .addValue(entryRow.getRow("value_timestamp")) + .addValue(entryRow.getRow("value_input")) + .build(); + } + + static Row entryToRow(Entry entry) { + if (entry instanceof SetCell) { + SetCell setCell = (SetCell) entry; + + return Row.withSchema(ENTRY_SCHEMA) + .addValue("SET_CELL") + .addValue(setCell.getFamilyName()) + .addValue(setCell.getQualifier().toByteArray()) + .addValue(setCell.getTimestamp()) + .addValue(setCell.getValue().toByteArray()) + .addValue(null) + .addValue(null) + .addValue(null) + .addValue(null) + .build(); + } + + if (entry instanceof DeleteCells) { + DeleteCells deleteCells = (DeleteCells) entry; + + return Row.withSchema(ENTRY_SCHEMA) + .addValue("DELETE_CELLS") + .addValue(deleteCells.getFamilyName()) + .addValue(deleteCells.getQualifier().toByteArray()) + .addValue(null) + .addValue(null) + .addValue(timestampRangeToRow(deleteCells.getTimestampRange())) + .addValue(null) + .addValue(null) + .addValue(null) + .build(); + } + + if (entry instanceof DeleteFamily) { + DeleteFamily deleteFamily = (DeleteFamily) entry; + + return Row.withSchema(ENTRY_SCHEMA) + .addValue("DELETE_FAMILY") + .addValue(deleteFamily.getFamilyName()) + .addValue(null) + .addValue(null) + .addValue(null) + .addValue(null) + .addValue(null) + .addValue(null) + .addValue(null) + .build(); + } + + if (entry instanceof AddToCell) { + AddToCell addToCell = (AddToCell) entry; + + return Row.withSchema(ENTRY_SCHEMA) + .addValue("ADD_TO_CELL") + .addValue(addToCell.getFamily()) + .addValue(null) + .addValue(null) + .addValue(null) + .addValue(null) + .addValue(valueToRow(addToCell.getQualifier())) + .addValue(valueToRow(addToCell.getTimestamp())) + .addValue(valueToRow(addToCell.getInput())) + .build(); + } + + if (entry instanceof MergeToCell) { + MergeToCell mergeToCell = (MergeToCell) entry; + + return Row.withSchema(ENTRY_SCHEMA) + .addValue("MERGE_TO_CELL") + .addValue(mergeToCell.getFamily()) + .addValue(null) + .addValue(null) + .addValue(null) + .addValue(null) + .addValue(valueToRow(mergeToCell.getQualifier())) + .addValue(valueToRow(mergeToCell.getTimestamp())) + .addValue(valueToRow(mergeToCell.getInput())) + .build(); + } + + throw new IllegalArgumentException( + "Unsupported Bigtable change stream entry: " + entry.getClass().getName()); + } + + static Row timestampRangeToRow(Range.TimestampRange range) { + @Nullable Long startTimestampMicros = + range.getStartBound() == Range.BoundType.UNBOUNDED ? null : range.getStart(); + + @Nullable Long endTimestampMicros = + range.getEndBound() == Range.BoundType.UNBOUNDED ? null : range.getEnd(); + + return Row.withSchema(TIMESTAMP_RANGE_SCHEMA) + .addValue(range.getStartBound().name()) + .addValue(startTimestampMicros) + .addValue(range.getEndBound().name()) + .addValue(endTimestampMicros) + .build(); + } + + static Row valueToRow(Value value) { + switch (value.getValueType()) { + case Int64: + return Row.withSchema(VALUE_SCHEMA) + .addValue("INT64") + .addValue(((Value.IntValue) value).getValue()) + .addValue(null) + .addValue(null) + .build(); + + case RawTimestamp: + return Row.withSchema(VALUE_SCHEMA) + .addValue("RAW_TIMESTAMP") + .addValue(null) + .addValue(((Value.RawTimestamp) value).getValue()) + .addValue(null) + .build(); + + case RawValue: + return Row.withSchema(VALUE_SCHEMA) + .addValue("RAW_VALUE") + .addValue(null) + .addValue(null) + .addValue(((Value.RawValue) value).getValue().toByteArray()) + .build(); + + default: + throw new IllegalArgumentException( + "Unsupported Bigtable Value type: " + value.getValueType()); + } + } +} diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/BigtableChangeStreamReadSchemaTransformProviderTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/BigtableChangeStreamReadSchemaTransformProviderTest.java new file mode 100644 index 000000000000..efd7dbc4a5a3 --- /dev/null +++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/BigtableChangeStreamReadSchemaTransformProviderTest.java @@ -0,0 +1,392 @@ +/* + * 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.gcp.bigtable.changestreams; + +import static java.nio.charset.StandardCharsets.UTF_8; +import static org.junit.Assert.assertArrayEquals; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertNull; +import static org.junit.Assert.assertThrows; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +import com.google.cloud.bigtable.data.v2.models.AddToCell; +import com.google.cloud.bigtable.data.v2.models.ChangeStreamMutation; +import com.google.cloud.bigtable.data.v2.models.DeleteCells; +import com.google.cloud.bigtable.data.v2.models.DeleteFamily; +import com.google.cloud.bigtable.data.v2.models.MergeToCell; +import com.google.cloud.bigtable.data.v2.models.Range; +import com.google.cloud.bigtable.data.v2.models.SetCell; +import com.google.cloud.bigtable.data.v2.models.Value; +import com.google.common.collect.ImmutableList; +import com.google.protobuf.ByteString; +import java.util.Collection; +import java.util.HashSet; +import java.util.Set; +import org.apache.beam.sdk.io.gcp.bigtable.changestreams.BigtableChangeStreamReadSchemaTransformProvider.BigtableChangeStreamReadConfiguration; +import org.apache.beam.sdk.schemas.Schema; +import org.apache.beam.sdk.values.Row; +import org.junit.Test; + +public class BigtableChangeStreamReadSchemaTransformProviderTest { + + @Test + public void testIntValueToRow() { + Row row = BigtableChangeStreamReadSchemaTransformProvider.valueToRow(Value.intValue(123L)); + + assertEquals("INT64", row.getString("type")); + assertEquals(Long.valueOf(123L), row.getInt64("int_value")); + assertNull(row.getInt64("raw_timestamp_micros")); + assertNull(row.getBytes("raw_value")); + } + + @Test + public void testRawTimestampValueToRow() { + Row row = + BigtableChangeStreamReadSchemaTransformProvider.valueToRow(Value.rawTimestamp(123456L)); + + assertEquals("RAW_TIMESTAMP", row.getString("type")); + assertNull(row.getInt64("int_value")); + assertEquals(Long.valueOf(123456L), row.getInt64("raw_timestamp_micros")); + assertNull(row.getBytes("raw_value")); + } + + @Test + public void testRawValueToRow() { + byte[] value = "value".getBytes(UTF_8); + + Row row = + BigtableChangeStreamReadSchemaTransformProvider.valueToRow( + Value.rawValue(ByteString.copyFrom(value))); + + assertEquals("RAW_VALUE", row.getString("type")); + assertNull(row.getInt64("int_value")); + assertNull(row.getInt64("raw_timestamp_micros")); + assertArrayEquals(value, row.getBytes("raw_value")); + } + + @Test + public void testTimestampRangeToRow() { + Range.TimestampRange range = Range.TimestampRange.create(100L, 200L); + + Row row = BigtableChangeStreamReadSchemaTransformProvider.timestampRangeToRow(range); + + assertEquals("CLOSED", row.getString("start_bound")); + assertEquals(Long.valueOf(100L), row.getInt64("start_timestamp_micros")); + assertEquals("OPEN", row.getString("end_bound")); + assertEquals(Long.valueOf(200L), row.getInt64("end_timestamp_micros")); + } + + @Test + public void testUnboundedTimestampRangeToRow() { + Range.TimestampRange range = Range.TimestampRange.unbounded(); + + Row row = BigtableChangeStreamReadSchemaTransformProvider.timestampRangeToRow(range); + + assertEquals("UNBOUNDED", row.getString("start_bound")); + assertNull(row.getInt64("start_timestamp_micros")); + assertEquals("UNBOUNDED", row.getString("end_bound")); + assertNull(row.getInt64("end_timestamp_micros")); + } + + @Test + public void testSetCellEntryToRow() { + SetCell setCell = + SetCell.create( + "family", + ByteString.copyFromUtf8("column_qualifier"), + 123L, + ByteString.copyFromUtf8("value")); + + Row row = BigtableChangeStreamReadSchemaTransformProvider.entryToRow(setCell); + + assertEquals("SET_CELL", row.getString("type")); + assertEquals("family", row.getString("family_name")); + assertArrayEquals("column_qualifier".getBytes(UTF_8), row.getBytes("column_qualifier")); + assertEquals(Long.valueOf(123L), row.getInt64("timestamp_micros")); + assertArrayEquals("value".getBytes(UTF_8), row.getBytes("value")); + + assertNull(row.getRow("timestamp_range")); + assertNull(row.getRow("value_qualifier")); + assertNull(row.getRow("value_timestamp")); + assertNull(row.getRow("value_input")); + } + + @Test + public void testDeleteFamilyEntryToRow() { + DeleteFamily deleteFamily = DeleteFamily.create("family"); + + Row row = BigtableChangeStreamReadSchemaTransformProvider.entryToRow(deleteFamily); + + assertEquals("DELETE_FAMILY", row.getString("type")); + assertEquals("family", row.getString("family_name")); + + assertNull(row.getBytes("column_qualifier")); + assertNull(row.getInt64("timestamp_micros")); + assertNull(row.getBytes("value")); + assertNull(row.getRow("timestamp_range")); + assertNull(row.getRow("value_qualifier")); + assertNull(row.getRow("value_timestamp")); + assertNull(row.getRow("value_input")); + } + + @Test + public void testConfigurationValidation() { + BigtableChangeStreamReadConfiguration configuration = + BigtableChangeStreamReadConfiguration.builder() + .setProjectId("project") + .setInstanceId("instance") + .setTableId("table") + .build(); + + configuration.validate(); + } + + @Test + public void testConfigurationRejectsEmptyProject() { + BigtableChangeStreamReadConfiguration configuration = + BigtableChangeStreamReadConfiguration.builder() + .setProjectId("") + .setInstanceId("instance") + .setTableId("table") + .build(); + + assertThrows(IllegalArgumentException.class, configuration::validate); + } + + @Test + public void testConfigurationRejectsEmptyInstance() { + BigtableChangeStreamReadConfiguration configuration = + BigtableChangeStreamReadConfiguration.builder() + .setProjectId("project") + .setInstanceId("") + .setTableId("table") + .build(); + + assertThrows(IllegalArgumentException.class, configuration::validate); + } + + @Test + public void testConfigurationRejectsEmptyTable() { + BigtableChangeStreamReadConfiguration configuration = + BigtableChangeStreamReadConfiguration.builder() + .setProjectId("project") + .setInstanceId("instance") + .setTableId("") + .build(); + + assertThrows(IllegalArgumentException.class, configuration::validate); + } + + @Test + public void testDeleteCellsEntryToRow() { + DeleteCells deleteCells = + DeleteCells.create( + "family", + ByteString.copyFromUtf8("column_qualifier"), + Range.TimestampRange.create(100L, 200L)); + + Row row = BigtableChangeStreamReadSchemaTransformProvider.entryToRow(deleteCells); + + assertEquals("DELETE_CELLS", row.getString("type")); + assertEquals("family", row.getString("family_name")); + assertArrayEquals("column_qualifier".getBytes(UTF_8), row.getBytes("column_qualifier")); + + Row range = row.getRow("timestamp_range"); + assertEquals("CLOSED", range.getString("start_bound")); + assertEquals(Long.valueOf(100L), range.getInt64("start_timestamp_micros")); + assertEquals("OPEN", range.getString("end_bound")); + assertEquals(Long.valueOf(200L), range.getInt64("end_timestamp_micros")); + } + + @Test + public void testAddToCellEntryToRow() { + AddToCell addToCell = + AddToCell.create( + "family", + Value.rawValue(ByteString.copyFromUtf8("column_qualifier")), + Value.rawTimestamp(123L), + Value.intValue(42L)); + + Row row = BigtableChangeStreamReadSchemaTransformProvider.entryToRow(addToCell); + + assertEquals("ADD_TO_CELL", row.getString("type")); + assertEquals("family", row.getString("family_name")); + + Row qualifier = row.getRow("value_qualifier"); + assertEquals("RAW_VALUE", qualifier.getString("type")); + assertArrayEquals("column_qualifier".getBytes(UTF_8), qualifier.getBytes("raw_value")); + + Row timestamp = row.getRow("value_timestamp"); + assertEquals("RAW_TIMESTAMP", timestamp.getString("type")); + assertEquals(Long.valueOf(123L), timestamp.getInt64("raw_timestamp_micros")); + + Row input = row.getRow("value_input"); + assertEquals("INT64", input.getString("type")); + assertEquals(Long.valueOf(42L), input.getInt64("int_value")); + } + + @Test + public void testMergeToCellEntryToRow() { + MergeToCell mergeToCell = + MergeToCell.create( + "family", + Value.rawValue(ByteString.copyFromUtf8("column_qualifier")), + Value.rawTimestamp(123L), + Value.rawValue(ByteString.copyFromUtf8("input"))); + + Row row = BigtableChangeStreamReadSchemaTransformProvider.entryToRow(mergeToCell); + + assertEquals("MERGE_TO_CELL", row.getString("type")); + assertEquals("family", row.getString("family_name")); + + Row input = row.getRow("value_input"); + assertEquals("RAW_VALUE", input.getString("type")); + assertArrayEquals("input".getBytes(UTF_8), input.getBytes("raw_value")); + } + + @Test + public void testMutationToRow() { + ChangeStreamMutation mutation = mock(ChangeStreamMutation.class); + + SetCell setCell = + SetCell.create( + "family", + ByteString.copyFromUtf8("column_qualifier"), + 123L, + ByteString.copyFromUtf8("value")); + + java.time.Instant commitTime = java.time.Instant.parse("2026-10-19T10:00:00.123Z"); + java.time.Instant estimatedLowWatermark = java.time.Instant.parse("2026-10-19T09:59:59.456Z"); + + when(mutation.getRowKey()).thenReturn(ByteString.copyFromUtf8("key")); + when(mutation.getType()).thenReturn(ChangeStreamMutation.MutationType.USER); + when(mutation.getSourceClusterId()).thenReturn("cluster-1"); + when(mutation.getCommitTime()).thenReturn(commitTime); + when(mutation.getTieBreaker()).thenReturn(42); + when(mutation.getToken()).thenReturn("token-1"); + when(mutation.getEstimatedLowWatermarkTime()).thenReturn(estimatedLowWatermark); + when(mutation.getEntries()).thenReturn(ImmutableList.of(setCell)); + + Row row = BigtableChangeStreamReadSchemaTransformProvider.mutationToRow(mutation); + + assertArrayEquals("key".getBytes(UTF_8), row.getBytes("key")); + assertEquals("USER", row.getString("mutation_type")); + assertEquals("cluster-1", row.getString("source_cluster_id")); + assertEquals( + new org.joda.time.Instant(commitTime.toEpochMilli()), row.getDateTime("commit_timestamp")); + assertEquals(Integer.valueOf(42), row.getInt32("tie_breaker")); + assertEquals("token-1", row.getString("token")); + assertEquals( + new org.joda.time.Instant(estimatedLowWatermark.toEpochMilli()), + row.getDateTime("estimated_low_watermark")); + + assertEquals(1, row.getArray("entries").size()); + + Collection entries = row.getArray("entries"); + assertEquals(1, entries.size()); + + Row entry = entries.iterator().next(); + assertEquals("SET_CELL", entry.getString("type")); + assertEquals("family", entry.getString("family_name")); + assertArrayEquals("column_qualifier".getBytes(UTF_8), entry.getBytes("column_qualifier")); + assertEquals(Long.valueOf(123L), entry.getInt64("timestamp_micros")); + assertArrayEquals("value".getBytes(UTF_8), entry.getBytes("value")); + } + + @Test + public void testFlattenedMutationEntryToRow() { + ChangeStreamMutation mutation = mock(ChangeStreamMutation.class); + + SetCell setCell = + SetCell.create( + "family", + ByteString.copyFromUtf8("column_qualifier"), + 123L, + ByteString.copyFromUtf8("value")); + + java.time.Instant commitTime = java.time.Instant.parse("2026-10-19T10:00:00.123Z"); + java.time.Instant estimatedLowWatermark = java.time.Instant.parse("2026-10-19T09:59:59.456Z"); + + when(mutation.getRowKey()).thenReturn(ByteString.copyFromUtf8("key")); + when(mutation.getType()).thenReturn(ChangeStreamMutation.MutationType.USER); + when(mutation.getSourceClusterId()).thenReturn("cluster-1"); + when(mutation.getCommitTime()).thenReturn(commitTime); + when(mutation.getTieBreaker()).thenReturn(42); + when(mutation.getToken()).thenReturn("token-1"); + when(mutation.getEstimatedLowWatermarkTime()).thenReturn(estimatedLowWatermark); + + Row row = + BigtableChangeStreamReadSchemaTransformProvider.flattenedMutationEntryToRow( + mutation, setCell); + + assertArrayEquals("key".getBytes(UTF_8), row.getBytes("key")); + assertEquals("USER", row.getString("mutation_type")); + assertEquals("cluster-1", row.getString("source_cluster_id")); + assertEquals( + new org.joda.time.Instant(commitTime.toEpochMilli()), row.getDateTime("commit_timestamp")); + assertEquals(Integer.valueOf(42), row.getInt32("tie_breaker")); + assertEquals("token-1", row.getString("token")); + assertEquals( + new org.joda.time.Instant(estimatedLowWatermark.toEpochMilli()), + row.getDateTime("estimated_low_watermark")); + + assertEquals("SET_CELL", row.getString("type")); + assertEquals("family", row.getString("family_name")); + assertArrayEquals("column_qualifier".getBytes(UTF_8), row.getBytes("column_qualifier")); + assertEquals(Long.valueOf(123L), row.getInt64("timestamp_micros")); + assertArrayEquals("value".getBytes(UTF_8), row.getBytes("value")); + + assertNull(row.getRow("timestamp_range")); + assertNull(row.getRow("value_qualifier")); + assertNull(row.getRow("value_timestamp")); + assertNull(row.getRow("value_input")); + } + + @Test + public void testFlattenDefaultsToTrue() { + BigtableChangeStreamReadConfiguration configuration = + BigtableChangeStreamReadConfiguration.builder() + .setProjectId("project") + .setInstanceId("instance") + .setTableId("table") + .build(); + + assertEquals(Boolean.TRUE, configuration.getFlatten()); + } + + @Test + public void testConfigurationSchema() { + BigtableChangeStreamReadSchemaTransformProvider provider = + new BigtableChangeStreamReadSchemaTransformProvider(); + + Schema schema = provider.configurationSchema(); + + assertEquals( + Set.of( + "project_id", + "instance_id", + "table_id", + "app_profile_id", + "start_at_timestamp", + "end_at_timestamp", + "change_stream_name", + "flatten"), + new HashSet<>(schema.getFieldNames())); + } +} diff --git a/sdks/python/apache_beam/yaml/extended_tests/e2e/bigtable_cdc.yaml b/sdks/python/apache_beam/yaml/extended_tests/e2e/bigtable_cdc.yaml new file mode 100644 index 000000000000..f6eae17ea192 --- /dev/null +++ b/sdks/python/apache_beam/yaml/extended_tests/e2e/bigtable_cdc.yaml @@ -0,0 +1,74 @@ +# +# 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: BT_CDC + type: "apache_beam.yaml.integration_tests.temp_bigtable_change_stream_table" + config: + project: "apache-beam-testing" + +pipelines: + - pipeline: + type: chain + transforms: + - type: ReadFromBigtableCDC + config: + project: "{BT_CDC[PROJECT]}" + instance: "{BT_CDC[INSTANCE]}" + table: "{BT_CDC[TABLE]}" + app_profile: "{BT_CDC[APP_PROFILE]}" + start_at: "{BT_CDC[START_AT]}" + end_at: "{BT_CDC[END_AT]}" + flatten: true + + - type: MapToFields + config: + language: python + fields: + key: + callable: | + def decode_key(row): + return row.key.decode("utf-8") + + mutation_type: + mutation_type + + type: + type + + family_name: + family_name + + column_qualifier: + callable: | + def decode_qualifier(row): + return row.column_qualifier.decode("utf-8") + + value: + callable: | + def decode_value(row): + return row.value.decode("utf-8") + + - type: AssertEqual + config: + elements: + - key: "row1" + mutation_type: "USER" + type: "SET_CELL" + family_name: "cf1" + column_qualifier: "cq1" + value: "value1" diff --git a/sdks/python/apache_beam/yaml/integration_tests.py b/sdks/python/apache_beam/yaml/integration_tests.py index 71a2c3770c9a..7da244493fcc 100644 --- a/sdks/python/apache_beam/yaml/integration_tests.py +++ b/sdks/python/apache_beam/yaml/integration_tests.py @@ -35,6 +35,7 @@ import unittest import uuid from datetime import datetime +from datetime import timedelta from datetime import timezone import mock @@ -79,7 +80,11 @@ def get_impl(self): from apitools.base.py.exceptions import HttpError from google.cloud import pubsub_v1 from google.cloud.bigtable import client +from google.cloud.bigtable_admin_v2.types import bigtable_table_admin from google.cloud.bigtable_admin_v2.types import instance +from google.cloud.bigtable_admin_v2.types import table as table_pb +from google.protobuf import duration_pb2 +from google.protobuf import field_mask_pb2 try: from google.cloud import firestore @@ -259,6 +264,111 @@ def instance_prefix(instance): return instance_id +@contextlib.contextmanager +def temp_bigtable_change_stream_table(project, prefix='yaml_bt_cdc_it_'): + instance_name = 'bt-cdc-tests' + table_id = 'test-table' + cluster_id = 'test-cluster' + app_profile_id = 'cdc-profile' + + instance_id = instance_prefix(instance_name) + + bigtable_client = client.Client(admin=True, project=project) + + bigtable_instance = bigtable_client.instance( + instance_id, + display_name=instance_name, + instance_type=instance.Instance.Type.DEVELOPMENT) + + cluster = bigtable_instance.cluster(cluster_id, 'us-central1-a') + + operation = bigtable_instance.create(clusters=[cluster]) + operation.result(timeout=500) + + _LOGGER.info( + 'Created Bigtable CDC instance [%s] in project [%s]', + instance_id, + project) + + table = bigtable_instance.table(table_id) + table.create() + + _LOGGER.info('Created Bigtable CDC table [%s]', table_id) + + column_family = table.column_family('cf1') + column_family.create() + + table_name = (f'projects/{project}/instances/{instance_id}/tables/{table_id}') + + change_stream_config = table_pb.ChangeStreamConfig( + retention_period=duration_pb2.Duration(seconds=24 * 60 * 60)) + + request = bigtable_table_admin.UpdateTableRequest( + table=table_pb.Table( + name=table_name, change_stream_config=change_stream_config), + update_mask=field_mask_pb2.FieldMask(paths=['change_stream_config'])) + + operation = bigtable_client.table_admin_client.update_table(request=request) + operation.result(timeout=500) + + _LOGGER.info('Enabled change stream for Bigtable table [%s]', table_id) + + app_profile = bigtable_instance.app_profile( + app_profile_id, + routing_policy_type='single-cluster', + cluster_id=cluster_id, + allow_transactional_writes=True) + + app_profile.create() + + _LOGGER.info('Created Bigtable CDC app profile [%s]', app_profile_id) + + # Keep some distance from the start timestamp to avoid clock skew. + time.sleep(0.5) + start_at = datetime.now(timezone.utc) + time.sleep(0.5) + row = table.row(b'row1') + row.set_cell('cf1', b'cq1', b'value1') + row.commit() + + _LOGGER.info('Wrote test mutation to Bigtable CDC table [%s]', table_id) + # The change stream needs a future end time so that it can catch up + # and terminate predictably. + end_at = datetime.now(timezone.utc) + timedelta(seconds=30) + + def format_timestamp(timestamp): + return timestamp.isoformat(timespec='milliseconds').replace('+00:00', 'Z') + + try: + yield { + 'PROJECT': project, + 'INSTANCE': instance_id, + 'TABLE': table_id, + 'APP_PROFILE': app_profile_id, + 'START_AT': format_timestamp(start_at), + 'END_AT': format_timestamp(end_at), + } + finally: + try: + _LOGGER.info('Disabling change stream for Bigtable table [%s]', table_id) + + disable_request = bigtable_table_admin.UpdateTableRequest( + table=table_pb.Table(name=table_name), + update_mask=field_mask_pb2.FieldMask(paths=['change_stream_config'])) + + operation = bigtable_client.table_admin_client.update_table( + request=disable_request) + operation.result(timeout=500) + + _LOGGER.info('Deleting Bigtable CDC table [%s]', table_id) + table.delete() + + _LOGGER.info('Deleting Bigtable CDC instance [%s]', instance_id) + bigtable_instance.delete() + except Exception: + _LOGGER.exception('Failed to clean up Bigtable CDC resources') + + @contextlib.contextmanager def temp_bigtable_table(project, prefix='yaml_bt_it_'): INSTANCE = "bt-write-tests" diff --git a/sdks/python/apache_beam/yaml/standard_io.yaml b/sdks/python/apache_beam/yaml/standard_io.yaml index 173236dc5e94..6ac61a3f53f7 100644 --- a/sdks/python/apache_beam/yaml/standard_io.yaml +++ b/sdks/python/apache_beam/yaml/standard_io.yaml @@ -591,6 +591,28 @@ config: gradle_target: 'sdks:java:io:google-cloud-platform:expansion-service:shadowJar' +# Bigtable CDC +- type: renaming + transforms: + 'ReadFromBigtableCDC': 'ReadFromBigtableCDC' + config: + mappings: + 'ReadFromBigtableCDC': + project: 'project_id' + instance: 'instance_id' + table: 'table_id' + app_profile: 'app_profile_id' + start_at: 'start_at_timestamp' + end_at: 'end_at_timestamp' + change_stream: 'change_stream_name' + flatten: 'flatten' + underlying_provider: + type: beamJar + transforms: + 'ReadFromBigtableCDC': 'beam:schematransform:org.apache.beam:bigtable_cdc_read:v1' + config: + gradle_target: 'sdks:java:io:google-cloud-platform:expansion-service:shadowJar' + #IcebergCDC - type: renaming transforms: diff --git a/sdks/python/apache_beam/yaml/yaml_provider_unit_test.py b/sdks/python/apache_beam/yaml/yaml_provider_unit_test.py index e1e3ee847d96..4b57feae5296 100644 --- a/sdks/python/apache_beam/yaml/yaml_provider_unit_test.py +++ b/sdks/python/apache_beam/yaml/yaml_provider_unit_test.py @@ -75,6 +75,14 @@ def test_parse_duration_with_missing_value(self): self.parse_duration('s', 'size') +class StandardProvidersTest(unittest.TestCase): + def test_bigtable_cdc_provider_is_registered(self): + providers = yaml_provider.standard_providers() + + self.assertIn('ReadFromBigtableCDC', providers) + self.assertTrue(providers['ReadFromBigtableCDC']) + + class ProviderParsingTest(unittest.TestCase): INLINE_PROVIDER = {'type': 'TEST', 'name': 'INLINED'}