tonytanger commented on code in PR #40261: URL: https://github.com/apache/beam/pull/40261#discussion_r4447930984
########## 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. + * + * <p>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<BigtableChangeStreamReadConfiguration> 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<String> inputCollectionNames() { + return Collections.emptyList(); + } + + @Override + public List<String> 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()); Review Comment: **Missing metadata table and `ExistingPipelineOptions` configuration:** `BigtableChangeStreamReadConfiguration` exposes `change_stream_name` and `app_profile_id`, but omits two sets of options needed in production Bigtable CDC pipelines: 1. `existing_pipeline_options` (`withExistingPipelineOptions`): `BigtableIO.readChangeStream()` defaults to `ExistingPipelineOptions.FAIL_IF_EXISTS`. If a user configures `change_stream_name` in YAML and restarts their pipeline, the job will fail on startup unless `RESUME_OR_NEW` (or `SKIP_CLEANUP`) can be configured. 2. Metadata table configuration (`withMetadataTableProjectId`, `withMetadataTableInstanceId`, `withMetadataTableTableId`, `withMetadataTableAppProfileId`, `withCreateOrUpdateMetadataTable`): Production deployments commonly store the metadata table in a separate Bigtable instance or use a transactional read-write app profile separate from a read-only source table/app profile. **How to fix:** Add optional configuration fields for `existing_pipeline_options` and `metadata_table_*` to `BigtableChangeStreamReadConfiguration`, wire them to `BigtableIO.ReadChangeStream` here, and expose their mappings in `standard_io.yaml`. ########## 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 Review Comment: **Why `bigtable_cdc.yaml` times out here (and how to test this instead):** In `BigtableIO.readChangeStream()`, `DetectNewPartitionsDoFn` (DNP) and `ReadChangeStreamPartitionDoFn` (RCSP) are designed to run **concurrently** on a streaming runner: RCSP streams partitions and writes partition watermarks to the metadata table, while DNP tracks the global low watermark across all partitions in the metadata table and waits for all RCSPs to finish before terminating. The Python YAML `integration_tests.py` harness runs external xlang transforms on Python's local batch `FnApiRunner` (`fn_runner.py`), which executes stages sequentially rather than concurrently—it buffers DNP's output in `watermark_pending_inputs` at `MAX_TIMESTAMP` and never starts RCSP until DNP finishes, creating a deadlock. Since `end_at` shouldn't be exposed publicly on `ReadFromBigtableCDC` anyway (and `FnApiRunner` cannot run concurrent streaming SDF stages, which is also why `ReadFromSpannerCDC` and `ReadFromPubSub` don't run read pipelines in this harness), please remove `bigtable_cdc.yaml` / `temp_bigtable_change_stream_table` here and add a Java `SchemaTransform` integration test on a streaming runner instead (similar to `SpannerChangeStreamsSchemaTransformIT` / existing Bigtable CDC Java ITs). ########## sdks/python/apache_beam/yaml/integration_tests.py: ########## @@ -259,6 +264,111 @@ def instance_prefix(instance): return instance_id [email protected] +def temp_bigtable_change_stream_table(project, prefix='yaml_bt_cdc_it_'): Review Comment: **Resource leak if setup or table cleanup fails:** 1. `table.create()`, `column_family.create()`, `update_table()`, `app_profile.create()`, and `row.commit()` all run **before** the `try ... finally:` block (line 342). If any of those calls fail or time out after `bigtable_instance.create()` (line 285), `finally:` is never reached and the Bigtable instance is leaked. 2. Inside `finally:`, `bigtable_instance.delete()` is in the same `try` block after `update_table()` and `table.delete()`. If disabling the change stream or deleting the table raises an exception, `bigtable_instance.delete()` is skipped. 3. `prefix='yaml_bt_cdc_it_'` in the function signature is unused (`instance_name = 'bt-cdc-tests'` is hardcoded). **How to fix:** Move all post-instance-creation setup inside `try:`, and wrap `bigtable_instance.delete()` in its own `try ... except` block inside `finally:` so instance deletion is guaranteed even if table cleanup fails: ```python table = bigtable_instance.table(table_id) table_name = f'projects/{project}/instances/{instance_id}/tables/{table_id}' try: table.create() # ... column_family, update_table, app_profile, row.commit(), yield ... finally: try: disable_request = bigtable_table_admin.UpdateTableRequest( table=table_pb.Table(name=table_name), update_mask=field_mask_pb2.FieldMask(paths=['change_stream_config'])) bigtable_client.table_admin_client.update_table( request=disable_request).result(timeout=500) table.delete() except Exception: _LOGGER.warning('Failed to clean up Bigtable CDC table [%s]', table_id, exc_info=True) try: _LOGGER.info('Deleting Bigtable CDC instance [%s]', instance_id) bigtable_instance.delete() except Exception: _LOGGER.exception('Failed to delete Bigtable CDC instance [%s]', instance_id) ``` ########## 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) { Review Comment: **Do not expose `withEndTime` / `end_at` as a public user-facing option:** `withEndTime` is intentionally package-private and `@VisibleForTesting` (`"Used only for integration tests. Unsafe to use in production."`). Cloud Bigtable's low watermark is not guaranteed to be monotonic (it can go backwards), so terminating a change stream at `endTime` can lead to data loss—which is why `withEndTime` was purposefully disabled as a public option on `BigtableIO.ReadChangeStream`. Making `withEndTime` `public` (which also triggers ErrorProne `[VisibleForTesting]` warnings when called from production code in another package) and exposing `end_at` / `end_at_timestamp` in the public `ReadFromBigtableCDC` YAML configuration exposes an unsafe API to users. **How to fix:** Revert `BigtableIO.java` and remove `end_at_timestamp` / `end_at` from `BigtableChangeStreamReadConfiguration` and `standard_io.yaml`. ########## sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/BigtableChangeStreamReadSchemaTransformProvider.java: ########## @@ -0,0 +1,402 @@ +/* + * 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. + * + * <p>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("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 = Review Comment: **`flatten: true` breaks row-level atomicity and loses intra-`MutateRow` entry ordering:** In Bigtable, a `MutateRow` applies a list of mutations atomically to a single row, and per the `MutateRowRequest` contract, entries within a `MutateRow` are applied in order so later entries mask earlier ones (for example, `deleteRow()` / `deleteFamily()` followed by `setCell()` in an atomic row overwrite). `ChangeStreamMutation` intentionally groups all entries of a row mutation together to preserve both row-level atomicity and entry ordering. Splitting `mutation.getEntries()` into independent `PCollection<Row>` elements breaks row-level atomicity and—because all entries in the same `ChangeStreamMutation` share the exact same `key`, `commit_timestamp`, and `token` with no sequence index—loses intra-mutation ordering across any downstream shuffle (so a `DeleteFamily` followed by a `SetCell` in the same `MutateRow` can be replayed in reverse order downstream). **How to fix (two options):** 1. **(Preferred) Remove `flatten` mode:** Always output `CHANGE_STREAM_MUTATION_SCHEMA` (one Beam `Row` per `ChangeStreamMutation` with nested `entries` as `ARRAY<ROW>`), which preserves both row-level atomicity and entry ordering. 2. **If `flatten` is needed for YAML convenience:** Default `flatten` to `false` (opt-in rather than opt-out), and add an `.addInt32Field("entry_index")` (the 0-based index within `mutation.getEntries()`) to `FLATTENED_CHANGE_STREAM_MUTATION_SCHEMA` so downstream consumers can reconstruct the intra-mutation execution order. ########## sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/BigtableChangeStreamReadSchemaTransformProvider.java: ########## @@ -0,0 +1,402 @@ +/* + * 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. + * + * <p>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("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("row_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(); + + @Override + protected Class<BigtableChangeStreamReadConfiguration> 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."; + } + + @Override + public List<String> inputCollectionNames() { + return Collections.emptyList(); + } + + @Override + public List<String> 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(); + + public static Builder builder() { + return new AutoValue_BigtableChangeStreamReadSchemaTransformProvider_BigtableChangeStreamReadConfiguration + .Builder(); + } + + /** 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 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 changeStreamName = configuration.getChangeStreamName(); + if (changeStreamName != null) { + readChangeStream = readChangeStream.withChangeStreamName(changeStreamName); + } + + PCollection<KV<ByteString, ChangeStreamMutation>> mutations = + input.getPipeline().apply(readChangeStream); + + PCollection<Row> rows = + mutations + .apply("ConvertToBeamRows", ParDo.of(new ChangeStreamMutationToRowDoFn())) + .setRowSchema(CHANGE_STREAM_MUTATION_SCHEMA); + + return PCollectionRowTuple.of(OUTPUT_TAG, rows); + } + } + + private static class ChangeStreamMutationToRowDoFn + extends DoFn<KV<ByteString, ChangeStreamMutation>, Row> { + + @ProcessElement + public void processElement( + @Element KV<ByteString, ChangeStreamMutation> element, OutputReceiver<Row> out) { + out.output(mutationToRow(element.getValue())); + } + } + + static Row mutationToRow(ChangeStreamMutation mutation) { + List<Row> 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(); Review Comment: Bigtable's server-side replication clock (and `ChangeStreamMutation.getCommitTime()` / `getEstimatedLowWatermarkTime()`) has **microsecond** precision (`java.time.Instant`), so converting via `new org.joda.time.Instant(...toEpochMilli())` for `FieldType.DATETIME` (`millis_instant`) drops the sub-millisecond microseconds. Two distinct mutations on the same row within the same millisecond have distinct microsecond `commit_timestamp`s in Bigtable, which collide once truncated to milliseconds. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
