diff --git a/sdks/java/io/components/src/main/java/org/apache/beam/sdk/io/components/ratelimiter/RateLimiter.java b/sdks/java/io/components/src/main/java/org/apache/beam/sdk/io/components/ratelimiter/RateLimiter.java new file mode 100644 index 000000000000..8c02654b3964 --- /dev/null +++ b/sdks/java/io/components/src/main/java/org/apache/beam/sdk/io/components/ratelimiter/RateLimiter.java @@ -0,0 +1,41 @@ +/* + * 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.components.ratelimiter; + +import java.io.IOException; +import java.io.Serializable; + +/** + * A RateLimiter allows to fetch permits from a rate limiter service and blocks execution when the + * rate limit is exceeded. + * + *

Implementations must be {@link Serializable} as they are passed to workers. + */ +public interface RateLimiter extends Serializable, AutoCloseable { + + /** + * Blocks until the specified number of permits are acquired and returns true if the request was + * allowed or false if the request was rejected. + * + * @param permits Number of permits to acquire. + * @return true if the request was allowed, false if it was rejected (and retries exceeded). + * @throws IOException if there is an error communicating with the rate limiter service. + * @throws InterruptedException if the thread is interrupted while waiting. + */ + boolean allow(int permits) throws IOException, InterruptedException; +} diff --git a/sdks/java/io/components/src/main/java/org/apache/beam/sdk/io/components/ratelimiter/RateLimiterContext.java b/sdks/java/io/components/src/main/java/org/apache/beam/sdk/io/components/ratelimiter/RateLimiterContext.java new file mode 100644 index 000000000000..6387bf5789e4 --- /dev/null +++ b/sdks/java/io/components/src/main/java/org/apache/beam/sdk/io/components/ratelimiter/RateLimiterContext.java @@ -0,0 +1,27 @@ +/* + * 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.components.ratelimiter; + +import java.io.Serializable; + +/** + * A marker interface for context data required to check ratelimit. + * + *

Implementations must be {@link Serializable}. + */ +public interface RateLimiterContext extends Serializable {} diff --git a/sdks/java/io/components/src/main/java/org/apache/beam/sdk/io/components/ratelimiter/RateLimiterFactory.java b/sdks/java/io/components/src/main/java/org/apache/beam/sdk/io/components/ratelimiter/RateLimiterFactory.java new file mode 100644 index 000000000000..b4330cd53db5 --- /dev/null +++ b/sdks/java/io/components/src/main/java/org/apache/beam/sdk/io/components/ratelimiter/RateLimiterFactory.java @@ -0,0 +1,57 @@ +/* + * 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.components.ratelimiter; + +import java.io.IOException; +import java.io.Serializable; + +/** + * A factory that manages connections to rate limit service and creates lightweight handles. + * + *

Implementations must be {@link Serializable} as they are passed to workers. The factory + * typically manages the heavy connection (e.g. gRPC stub) and is thread-safe. + */ +public interface RateLimiterFactory extends Serializable, AutoCloseable { + + /** + * Creates a lightweight ratelimiter handle bound to a specific context. + * + *

Use this when passing ratelimiter to IO components, which doesn't need to know about the + * configuration or the underlying ratelimiter service details. This is also useful in DoFns when + * you want to use the ratelimiter in a static way based on the compile time context. + * + * @param context The context for the ratelimit. + * @return A {@link RateLimiter} handle. + */ + RateLimiter getLimiter(RateLimiterContext context); + + /** + * Blocks until the specified number of permits are acquired and returns true if the request was + * allowed or false if the request was rejected. + * + *

Use this for when the ratelimit namespace or descriptors are not known at compile time. + * allows you to use the ratelimiter in a dynamic way based on the runtime data. + * + * @param context The context for the ratelimit. + * @param permits Number of permits to acquire. + * @return true if the request is allowed, false if rejected. + * @throws IOException if there is an error communicating with the ratelimiter service. + * @throws InterruptedException if the thread is interrupted while waiting. + */ + boolean allow(RateLimiterContext context, int permits) throws IOException, InterruptedException; +} diff --git a/sdks/java/io/components/src/main/java/org/apache/beam/sdk/io/components/ratelimiter/package-info.java b/sdks/java/io/components/src/main/java/org/apache/beam/sdk/io/components/ratelimiter/package-info.java new file mode 100644 index 000000000000..556447ad11c2 --- /dev/null +++ b/sdks/java/io/components/src/main/java/org/apache/beam/sdk/io/components/ratelimiter/package-info.java @@ -0,0 +1,20 @@ +/* + * 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. + */ + +/** Defines ratelimiter utilities for Beam DoFn and IO components. */ +package org.apache.beam.sdk.io.components.ratelimiter; diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/SpannerIO.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/SpannerIO.java index bbce5fad82f4..3a69d1177f4a 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/SpannerIO.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/SpannerIO.java @@ -1977,11 +1977,6 @@ && getInclusiveStartAt().toSqlTimestamp().after(getInclusiveEndAt().toSqlTimesta + changeStreamDatabaseId + " has dialect " + changeStreamDatabaseDialect); - LOG.info( - "The Spanner database " - + fullPartitionMetadataDatabaseId - + " has dialect " - + metadataDatabaseDialect); PartitionMetadataTableNames partitionMetadataTableNames = Optional.ofNullable(getMetadataTable()) .map( @@ -2005,6 +2000,7 @@ && getInclusiveStartAt().toSqlTimestamp().after(getInclusiveEndAt().toSqlTimesta final boolean isMutableChangeStream = isMutableChangeStream( spannerAccessor.getDatabaseClient(), changeStreamDatabaseDialect, changeStreamName); + LOG.info("The change stream " + changeStreamName + " is mutable: " + isMutableChangeStream); final DaoFactory daoFactory = new DaoFactory( changeStreamSpannerConfig, diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/action/ActionFactory.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/action/ActionFactory.java index e8749a836669..cd84168b23f7 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/action/ActionFactory.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/action/ActionFactory.java @@ -187,7 +187,8 @@ public synchronized QueryChangeStreamAction queryChangeStreamAction( PartitionStartRecordAction partitionStartRecordAction, PartitionEndRecordAction partitionEndRecordAction, PartitionEventRecordAction partitionEventRecordAction, - ChangeStreamMetrics metrics) { + ChangeStreamMetrics metrics, + boolean isMutableChangeStream) { if (queryChangeStreamActionInstance == null) { queryChangeStreamActionInstance = new QueryChangeStreamAction( @@ -201,7 +202,8 @@ public synchronized QueryChangeStreamAction queryChangeStreamAction( partitionStartRecordAction, partitionEndRecordAction, partitionEventRecordAction, - metrics); + metrics, + isMutableChangeStream); } return queryChangeStreamActionInstance; } diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/action/QueryChangeStreamAction.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/action/QueryChangeStreamAction.java index 8da9f3d09515..69e89e74a38b 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/action/QueryChangeStreamAction.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/action/QueryChangeStreamAction.java @@ -34,6 +34,7 @@ import org.apache.beam.sdk.io.gcp.spanner.changestreams.model.ChildPartitionsRecord; import org.apache.beam.sdk.io.gcp.spanner.changestreams.model.DataChangeRecord; import org.apache.beam.sdk.io.gcp.spanner.changestreams.model.HeartbeatRecord; +import org.apache.beam.sdk.io.gcp.spanner.changestreams.model.InitialPartition; import org.apache.beam.sdk.io.gcp.spanner.changestreams.model.PartitionEndRecord; import org.apache.beam.sdk.io.gcp.spanner.changestreams.model.PartitionEventRecord; import org.apache.beam.sdk.io.gcp.spanner.changestreams.model.PartitionMetadata; @@ -89,6 +90,7 @@ public class QueryChangeStreamAction { private final PartitionEndRecordAction partitionEndRecordAction; private final PartitionEventRecordAction partitionEventRecordAction; private final ChangeStreamMetrics metrics; + private final boolean isMutableChangeStream; /** * Constructs an action class for performing a change stream query for a given partition. @@ -106,6 +108,7 @@ public class QueryChangeStreamAction { * @param PartitionEndRecordAction action class to process {@link PartitionEndRecord}s * @param PartitionEventRecordAction action class to process {@link PartitionEventRecord}s * @param metrics metrics gathering class + * @param isMutableChangeStream whether the change stream is mutable or not */ QueryChangeStreamAction( ChangeStreamDao changeStreamDao, @@ -118,7 +121,8 @@ public class QueryChangeStreamAction { PartitionStartRecordAction partitionStartRecordAction, PartitionEndRecordAction partitionEndRecordAction, PartitionEventRecordAction partitionEventRecordAction, - ChangeStreamMetrics metrics) { + ChangeStreamMetrics metrics, + boolean isMutableChangeStream) { this.changeStreamDao = changeStreamDao; this.partitionMetadataDao = partitionMetadataDao; this.changeStreamRecordMapper = changeStreamRecordMapper; @@ -130,6 +134,7 @@ public class QueryChangeStreamAction { this.partitionEndRecordAction = partitionEndRecordAction; this.partitionEventRecordAction = partitionEventRecordAction; this.metrics = metrics; + this.isMutableChangeStream = isMutableChangeStream; } /** @@ -195,13 +200,23 @@ public ProcessContinuation run( final Timestamp endTimestamp = partition.getEndTimestamp(); final boolean isBoundedRestriction = !endTimestamp.equals(MAX_INCLUSIVE_END_AT); final Timestamp changeStreamQueryEndTimestamp = - isBoundedRestriction ? endTimestamp : getNextReadChangeStreamEndTimestamp(); + isBoundedRestriction + ? getBoundedQueryEndTimestamp(endTimestamp) + : getNextReadChangeStreamEndTimestamp(); // Once the changeStreamQuery completes we may need to resume reading from the partition if we // had an unbounded restriction for which we set an arbitrary query end timestamp and for which // we didn't encounter any indications that the partition is done (explicit end records or - // exceptions about being out of timestamp range). - boolean stopAfterQuerySucceeds = isBoundedRestriction; + // exceptions about being out of timestamp range). We also special case the InitialPartition, + // which always stops after the query succeeds. + boolean stopAfterQuerySucceeds = false; + if (InitialPartition.isInitialPartition(partition.getPartitionToken())) { + stopAfterQuerySucceeds = true; + } else { + stopAfterQuerySucceeds = + isBoundedRestriction && changeStreamQueryEndTimestamp.equals(endTimestamp); + } + try (ChangeStreamResultSet resultSet = changeStreamDao.changeStreamQuery( token, startTimestamp, changeStreamQueryEndTimestamp, partition.getHeartbeatMillis())) { @@ -379,4 +394,14 @@ private Timestamp getNextReadChangeStreamEndTimestamp() { final Timestamp current = Timestamp.now(); return Timestamp.ofTimeSecondsAndNanos(current.getSeconds() + 2 * 60, current.getNanos()); } + + // For Mutable Change Stream bounded queries, update the query end timestamp to be within 2 + // minutes in the future. + private Timestamp getBoundedQueryEndTimestamp(Timestamp endTimestamp) { + if (this.isMutableChangeStream) { + Timestamp nextTimestamp = getNextReadChangeStreamEndTimestamp(); + return nextTimestamp.compareTo(endTimestamp) < 0 ? nextTimestamp : endTimestamp; + } + return endTimestamp; + } } diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/dao/ChangeStreamResultSet.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/dao/ChangeStreamResultSet.java index 1268c739164f..846c80951293 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/dao/ChangeStreamResultSet.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/dao/ChangeStreamResultSet.java @@ -20,6 +20,7 @@ import com.google.cloud.Timestamp; import com.google.cloud.spanner.ResultSet; import com.google.cloud.spanner.Struct; +import com.google.protobuf.InvalidProtocolBufferException; import org.joda.time.Duration; /** @@ -113,6 +114,9 @@ public Struct getCurrentRowAsStruct() { * updates the timestamp at which the record was read. This function enhances the getProtoMessage * function but only focus on the ChangeStreamRecord type. * + *

Should only be used for GoogleSQL databases when the change stream record is delivered as + * proto. + * * @return a change stream record as a proto or null */ public com.google.spanner.v1.ChangeStreamRecord getProtoChangeStreamRecord() { @@ -128,6 +132,33 @@ public boolean isProtoChangeRecord() { && resultSet.getColumnType(0).getCode() == com.google.cloud.spanner.Type.Code.PROTO; } + /** + * Returns the change stream record at the current pointer by parsing the bytes column. It also + * updates the timestamp at which the record was read. + * + *

Should only be used for PostgreSQL databases when the change stream record is delivered as + * proto bytes. + * + * @return a change stream record as a proto or null + */ + public com.google.spanner.v1.ChangeStreamRecord getBytes(int index) { + recordReadAt = Timestamp.now(); + try { + // Use getBytes(0) for the BYTES column returned by read_proto_bytes_ TVF + return com.google.spanner.v1.ChangeStreamRecord.parseFrom( + resultSet.getBytes(index).toByteArray()); + } catch (InvalidProtocolBufferException e) { + throw new RuntimeException("Failed to parse the proto bytes to ChangeStreamRecord proto", e); + } + } + + /** Returns true if the result set at the current pointer contain only one bytes change record. */ + public boolean isProtoBytesChangeRecord() { + return resultSet.getColumnCount() == 1 + && !resultSet.isNull(0) + && resultSet.getColumnType(0).getCode() == com.google.cloud.spanner.Type.Code.BYTES; + } + /** * Returns the record at the current pointer as {@link JsonB}. It also updates the timestamp at * which the record was read. diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/dao/DaoFactory.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/dao/DaoFactory.java index 67b58bace70f..95bdbfed7ca5 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/dao/DaoFactory.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/dao/DaoFactory.java @@ -151,4 +151,8 @@ public synchronized ChangeStreamDao getChangeStreamDao() { } return changeStreamDaoInstance; } + + public boolean isMutableChangeStream() { + return this.isMutableChangeStream; + } } diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/dofn/ReadChangeStreamPartitionDoFn.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/dofn/ReadChangeStreamPartitionDoFn.java index 4f5631c468be..c3650b42761b 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/dofn/ReadChangeStreamPartitionDoFn.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/dofn/ReadChangeStreamPartitionDoFn.java @@ -73,6 +73,7 @@ public class ReadChangeStreamPartitionDoFn extends DoFn(); } @@ -215,7 +217,8 @@ public void setup() { partitionStartRecordAction, partitionEndRecordAction, partitionEventRecordAction, - metrics); + metrics, + isMutableChangeStream); } /** diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/mapper/ChangeStreamRecordMapper.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/mapper/ChangeStreamRecordMapper.java index 631538646669..368fec88918e 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/mapper/ChangeStreamRecordMapper.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/mapper/ChangeStreamRecordMapper.java @@ -218,18 +218,27 @@ public class ChangeStreamRecordMapper { * @param resultSet the change stream result set * @param resultSetMetadata the metadata generated when reading the change stream row * @return a {@link List} of {@link ChangeStreamRecord} subclasses + * @throws InvalidProtocolBufferException */ public List toChangeStreamRecords( PartitionMetadata partition, ChangeStreamResultSet resultSet, ChangeStreamResultSetMetadata resultSetMetadata) { if (this.isPostgres()) { - // In PostgresQL, change stream records are returned as JsonB. + // For `MUTABLE_KEY_RANGE` option, change stream records are returned as protos. + if (resultSet.isProtoBytesChangeRecord()) { + return Arrays.asList( + toChangeStreamRecord(partition, resultSet.getBytes(0), resultSetMetadata)); + } + + // For `IMMUTABLE_KEY_RANGE` option, change stream records are returned as + // JsonB. return Collections.singletonList( toChangeStreamRecordJson(partition, resultSet.getPgJsonb(0), resultSetMetadata)); } - // In GoogleSQL, for `MUTABLE_KEY_RANGE` option, change stream records are returned as Protos. + // In GoogleSQL, for `MUTABLE_KEY_RANGE` option, change stream records are + // returned as Protos. if (resultSet.isProtoChangeRecord()) { return Arrays.asList( toChangeStreamRecord( diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/action/QueryChangeStreamActionTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/action/QueryChangeStreamActionTest.java index cf4c047025c4..26ab41dff878 100644 --- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/action/QueryChangeStreamActionTest.java +++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/action/QueryChangeStreamActionTest.java @@ -21,7 +21,9 @@ import static org.apache.beam.sdk.io.gcp.spanner.changestreams.model.PartitionMetadata.State.SCHEDULED; import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertNotEquals; +import static org.junit.Assert.assertTrue; import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyLong; import static org.mockito.ArgumentMatchers.eq; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.never; @@ -116,7 +118,8 @@ public void setUp() throws Exception { partitionStartRecordAction, partitionEndRecordAction, partitionEventRecordAction, - metrics); + metrics, + false); final Struct row = mock(Struct.class); partition = PartitionMetadata.newBuilder() @@ -916,6 +919,121 @@ public void testQueryChangeStreamWithChildPartitionsRecordUnboundedRestriction() verify(partitionMetadataDao, never()).updateWatermark(any(), any()); } + @Test + public void testQueryChangeStreamWithMutableChangeStreamCappedEndTimestamp() { + // Initialize action with isMutableChangeStream = true + action = + new QueryChangeStreamAction( + changeStreamDao, + partitionMetadataDao, + changeStreamRecordMapper, + partitionMetadataMapper, + dataChangeRecordAction, + heartbeatRecordAction, + childPartitionsRecordAction, + partitionStartRecordAction, + partitionEndRecordAction, + partitionEventRecordAction, + metrics, + true); + + // Set endTimestamp to 60 minutes in the future + Timestamp now = Timestamp.now(); + Timestamp endTimestamp = + Timestamp.ofTimeSecondsAndNanos(now.getSeconds() + 60 * 60, now.getNanos()); + + partition = partition.toBuilder().setEndTimestamp(endTimestamp).build(); + when(restriction.getTo()).thenReturn(endTimestamp); + when(partitionMetadataMapper.from(any())).thenReturn(partition); + + final ChangeStreamResultSet resultSet = mock(ChangeStreamResultSet.class); + final ArgumentCaptor timestampCaptor = ArgumentCaptor.forClass(Timestamp.class); + when(changeStreamDao.changeStreamQuery( + eq(PARTITION_TOKEN), eq(PARTITION_START_TIMESTAMP), + timestampCaptor.capture(), eq(PARTITION_HEARTBEAT_MILLIS))) + .thenReturn(resultSet); + when(resultSet.next()).thenReturn(false); // Query finishes (reaches cap) + when(watermarkEstimator.currentWatermark()).thenReturn(WATERMARK); + when(restrictionTracker.tryClaim(any(Timestamp.class))).thenReturn(true); + + final ProcessContinuation result = + action.run( + partition, restrictionTracker, outputReceiver, watermarkEstimator, bundleFinalizer); + + // Verify query was capped at ~2 minutes + long diff = timestampCaptor.getValue().getSeconds() - now.getSeconds(); + assertTrue("Query should be capped at approx 2 minutes (120s)", Math.abs(diff - 120) < 10); + + // Crucial: Should RESUME to process the rest later + assertEquals(ProcessContinuation.resume(), result); + } + + @Test + public void testQueryChangeStreamWithMutableChangeStreamUncappedEndTimestamp() { + action = + new QueryChangeStreamAction( + changeStreamDao, + partitionMetadataDao, + changeStreamRecordMapper, + partitionMetadataMapper, + dataChangeRecordAction, + heartbeatRecordAction, + childPartitionsRecordAction, + partitionStartRecordAction, + partitionEndRecordAction, + partitionEventRecordAction, + metrics, + true); + + // Set endTimestamp to only 10 seconds in the future + Timestamp now = Timestamp.now(); + Timestamp endTimestamp = Timestamp.ofTimeSecondsAndNanos(now.getSeconds() + 10, now.getNanos()); + + partition = partition.toBuilder().setEndTimestamp(endTimestamp).build(); + when(restriction.getTo()).thenReturn(endTimestamp); + when(partitionMetadataMapper.from(any())).thenReturn(partition); + + final ChangeStreamResultSet resultSet = mock(ChangeStreamResultSet.class); + final ArgumentCaptor timestampCaptor = ArgumentCaptor.forClass(Timestamp.class); + when(changeStreamDao.changeStreamQuery( + eq(PARTITION_TOKEN), eq(PARTITION_START_TIMESTAMP), + timestampCaptor.capture(), eq(PARTITION_HEARTBEAT_MILLIS))) + .thenReturn(resultSet); + when(resultSet.next()).thenReturn(false); + when(watermarkEstimator.currentWatermark()).thenReturn(WATERMARK); + when(restrictionTracker.tryClaim(endTimestamp)).thenReturn(true); + + final ProcessContinuation result = + action.run( + partition, restrictionTracker, outputReceiver, watermarkEstimator, bundleFinalizer); + + // Should use the exact endTimestamp since it is within the limit (10s < 2m) + assertEquals(endTimestamp, timestampCaptor.getValue()); + + // Should STOP because we reached the actual requested endTimestamp + assertEquals(ProcessContinuation.stop(), result); + } + + @Test + public void testQueryChangeStreamUnboundedResumesCorrectly() { + // Unbounded restriction (streaming forever) + setupUnboundedPartition(); + + final ChangeStreamResultSet resultSet = mock(ChangeStreamResultSet.class); + when(changeStreamDao.changeStreamQuery(any(), any(), any(), anyLong())).thenReturn(resultSet); + when(resultSet.next()).thenReturn(false); + when(watermarkEstimator.currentWatermark()).thenReturn(WATERMARK); + when(restrictionTracker.tryClaim(any(Timestamp.class))).thenReturn(true); + + final ProcessContinuation result = + action.run( + partition, restrictionTracker, outputReceiver, watermarkEstimator, bundleFinalizer); + + // Should return RESUME to continue reading the stream every 2 minutes + assertEquals(ProcessContinuation.resume(), result); + verify(metrics).incQueryCounter(); + } + private static class BundleFinalizerStub implements BundleFinalizer { @Override public void afterBundleCommit(Instant callbackExpiry, Callback callback) { diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/dao/ChangeStreamResultSetTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/dao/ChangeStreamResultSetTest.java new file mode 100644 index 000000000000..d3408536c82a --- /dev/null +++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/dao/ChangeStreamResultSetTest.java @@ -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. + */ +package org.apache.beam.sdk.io.gcp.spanner.changestreams.dao; + +import static org.apache.beam.sdk.io.gcp.spanner.changestreams.util.TestProtoMapper.recordToProto; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertNotNull; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +import com.google.cloud.ByteArray; +import com.google.cloud.Timestamp; +import com.google.cloud.spanner.ResultSet; +import org.apache.beam.sdk.io.gcp.spanner.changestreams.model.HeartbeatRecord; +import org.junit.Test; + +public class ChangeStreamResultSetTest { + + @Test + public void testGetBytes() throws Exception { + // 1. Create an expected ChangeStreamRecord proto + Timestamp now = Timestamp.now(); + final HeartbeatRecord heartbeatRecord = + new HeartbeatRecord(Timestamp.ofTimeSecondsAndNanos(10L, 20), null); + com.google.spanner.v1.ChangeStreamRecord expectedRecord = recordToProto(heartbeatRecord); + assertNotNull(expectedRecord); + + // 2. Convert it to bytes (simulating how Spanner PostgreSQL returns it) + byte[] protoBytes = expectedRecord.toByteArray(); + + // 3. Mock the underlying Spanner ResultSet + ResultSet mockResultSet = mock(ResultSet.class); + // Simulate column 0 containing the BYTES representation of the proto + when(mockResultSet.getBytes(0)).thenReturn(ByteArray.copyFrom(protoBytes)); + + // 4. Initialize ChangeStreamResultSet with the mock + ChangeStreamResultSet changeStreamResultSet = new ChangeStreamResultSet(mockResultSet); + + // 5. Call the new method and assert it parses correctly + // (Note: This assumes you have added getBytes(0) to the class) + com.google.spanner.v1.ChangeStreamRecord actualRecord = changeStreamResultSet.getBytes(0); + + assertEquals(expectedRecord, actualRecord); + } +} diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/dofn/ReadChangeStreamPartitionDoFnTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/dofn/ReadChangeStreamPartitionDoFnTest.java index 62fa39eef55a..9e588de77a03 100644 --- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/dofn/ReadChangeStreamPartitionDoFnTest.java +++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/dofn/ReadChangeStreamPartitionDoFnTest.java @@ -20,6 +20,8 @@ import static org.apache.beam.sdk.io.gcp.spanner.changestreams.model.PartitionMetadata.State.SCHEDULED; import static org.junit.Assert.assertEquals; import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyBoolean; +import static org.mockito.ArgumentMatchers.eq; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.never; import static org.mockito.Mockito.verify; @@ -139,17 +141,18 @@ public void setUp() { when(actionFactory.partitionEventRecordAction(partitionMetadataDao, metrics)) .thenReturn(partitionEventRecordAction); when(actionFactory.queryChangeStreamAction( - changeStreamDao, - partitionMetadataDao, - changeStreamRecordMapper, - partitionMetadataMapper, - dataChangeRecordAction, - heartbeatRecordAction, - childPartitionsRecordAction, - partitionStartRecordAction, - partitionEndRecordAction, - partitionEventRecordAction, - metrics)) + eq(changeStreamDao), + eq(partitionMetadataDao), + eq(changeStreamRecordMapper), + eq(partitionMetadataMapper), + eq(dataChangeRecordAction), + eq(heartbeatRecordAction), + eq(childPartitionsRecordAction), + eq(partitionStartRecordAction), + eq(partitionEndRecordAction), + eq(partitionEventRecordAction), + eq(metrics), + anyBoolean())) .thenReturn(queryChangeStreamAction); doFn.setup(); diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/it/SpannerChangeStreamPlacementTablePostgresIT.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/it/SpannerChangeStreamPlacementTablePostgresIT.java new file mode 100644 index 000000000000..573ac8259101 --- /dev/null +++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/it/SpannerChangeStreamPlacementTablePostgresIT.java @@ -0,0 +1,279 @@ +/* + * 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.spanner.changestreams.it; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertTrue; +import static org.junit.Assert.fail; + +import com.google.cloud.Timestamp; +import com.google.cloud.spanner.DatabaseClient; +import com.google.cloud.spanner.ErrorCode; +import com.google.cloud.spanner.Key; +import com.google.cloud.spanner.Mutation; +import com.google.cloud.spanner.Options; +import com.google.cloud.spanner.ResultSet; +import com.google.cloud.spanner.SpannerException; +import com.google.cloud.spanner.Statement; +import com.google.gson.Gson; +import java.util.Collections; +import java.util.Map; +import java.util.Optional; +import org.apache.beam.sdk.io.gcp.spanner.SpannerConfig; +import org.apache.beam.sdk.io.gcp.spanner.SpannerIO; +import org.apache.beam.sdk.io.gcp.spanner.changestreams.model.DataChangeRecord; +import org.apache.beam.sdk.io.gcp.spanner.changestreams.model.Mod; +import org.apache.beam.sdk.options.ValueProvider; +import org.apache.beam.sdk.testing.PAssert; +import org.apache.beam.sdk.testing.TestPipeline; +import org.apache.beam.sdk.transforms.DoFn; +import org.apache.beam.sdk.transforms.ParDo; +import org.apache.beam.sdk.values.PCollection; +import org.apache.commons.lang3.tuple.Pair; +import org.joda.time.Instant; +import org.junit.Before; +import org.junit.BeforeClass; +import org.junit.ClassRule; +import org.junit.Rule; +import org.junit.Test; +import org.junit.rules.Timeout; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +/** End-to-end test of Cloud Spanner CDC Source. */ +@RunWith(JUnit4.class) +public class SpannerChangeStreamPlacementTablePostgresIT { + + @Rule public transient Timeout globalTimeout = Timeout.seconds(3600); + + @ClassRule + public static final IntegrationTestEnv ENV = + new IntegrationTestEnv( + /*isPostgres=*/ true, + /*isPlacementTableBasedChangeStream=*/ true, + /*host=*/ Optional.empty()); + + @Rule public final transient TestPipeline pipeline = TestPipeline.create(); + + private static String instanceId; + private static String projectId; + private static String databaseId; + private static String metadataTableName; + private static String changeStreamTableName; + private static String changeStreamName; + private static DatabaseClient databaseClient; + private static String host = "https://spanner.googleapis.com"; + + @BeforeClass + public static void beforeClass() throws Exception { + projectId = ENV.getProjectId(); + instanceId = ENV.getInstanceId(); + databaseId = ENV.getDatabaseId(); + + metadataTableName = ENV.getMetadataTableName(); + changeStreamTableName = ENV.createSingersTable(); + changeStreamName = ENV.createChangeStreamFor(changeStreamTableName); + databaseClient = ENV.getDatabaseClient(); + } + + @Before + public void before() { + pipeline.getOptions().as(ChangeStreamTestPipelineOptions.class).setStreaming(true); + pipeline.getOptions().as(ChangeStreamTestPipelineOptions.class).setBlockOnRun(false); + } + + @Test + public void testReadSpannerChangeStream() { + // Defines how many rows are going to be inserted / updated / deleted in the test + final int numRows = 5; + // Inserts numRows rows and uses the first commit timestamp as the startAt for reading the + // change stream + final Pair insertTimestamps = insertRows(numRows); + final Timestamp startAt = insertTimestamps.getLeft(); + // Updates the created rows + updateRows(numRows); + // Delete the created rows and uses the last commit timestamp as the endAt for reading the + // change stream + final Pair deleteTimestamps = deleteRows(numRows); + final Timestamp endAt = deleteTimestamps.getRight(); + + final SpannerConfig spannerConfig = + SpannerConfig.create() + .withProjectId(projectId) + .withInstanceId(instanceId) + .withDatabaseId(databaseId) + .withHost(ValueProvider.StaticValueProvider.of(host)); + + final PCollection tokens = + pipeline + .apply( + SpannerIO.readChangeStream() + .withSpannerConfig(spannerConfig) + .withChangeStreamName(changeStreamName) + .withMetadataDatabase(databaseId) + .withMetadataTable(metadataTableName) + .withInclusiveStartAt(startAt) + .withInclusiveEndAt(endAt)) + .apply(ParDo.of(new ModsToString())); + + // Each row is composed by the following data + // + PAssert.that(tokens) + .containsInAnyOrder( + "INSERT,1,null,null,First Name 1,Last Name 1", + "INSERT,2,null,null,First Name 2,Last Name 2", + "INSERT,3,null,null,First Name 3,Last Name 3", + "INSERT,4,null,null,First Name 4,Last Name 4", + "INSERT,5,null,null,First Name 5,Last Name 5", + "UPDATE,1,First Name 1,Last Name 1,Updated First Name 1,Updated Last Name 1", + "UPDATE,2,First Name 2,Last Name 2,Updated First Name 2,Updated Last Name 2", + "UPDATE,3,First Name 3,Last Name 3,Updated First Name 3,Updated Last Name 3", + "UPDATE,4,First Name 4,Last Name 4,Updated First Name 4,Updated Last Name 4", + "UPDATE,5,First Name 5,Last Name 5,Updated First Name 5,Updated Last Name 5", + "DELETE,1,Updated First Name 1,Updated Last Name 1,null,null", + "DELETE,2,Updated First Name 2,Updated Last Name 2,null,null", + "DELETE,3,Updated First Name 3,Updated Last Name 3,null,null", + "DELETE,4,Updated First Name 4,Updated Last Name 4,null,null", + "DELETE,5,Updated First Name 5,Updated Last Name 5,null,null"); + pipeline.run().waitUntilFinish(); + + assertMetadataTableHasBeenDropped(); + } + + private static void assertMetadataTableHasBeenDropped() { + try (ResultSet resultSet = + databaseClient + .singleUse() + .executeQuery(Statement.of("SELECT * FROM \"" + metadataTableName + "\""))) { + resultSet.next(); + fail( + "The metadata table " + + metadataTableName + + " should had been dropped, but it still exists"); + } catch (SpannerException e) { + assertEquals(ErrorCode.INVALID_ARGUMENT, e.getErrorCode()); + assertTrue( + "Error message must contain \"Table not found\"", + e.getMessage().contains("relation \"" + metadataTableName + "\" does not exist")); + } + } + + private static Pair insertRows(int n) { + final Timestamp firstCommitTimestamp = insertRow(1); + for (int i = 2; i < n; i++) { + insertRow(i); + } + final Timestamp lastCommitTimestamp = insertRow(n); + return Pair.of(firstCommitTimestamp, lastCommitTimestamp); + } + + private static Pair updateRows(int n) { + final Timestamp firstCommitTimestamp = updateRow(1); + for (int i = 2; i < n; i++) { + updateRow(i); + } + final Timestamp lastCommitTimestamp = updateRow(n); + return Pair.of(firstCommitTimestamp, lastCommitTimestamp); + } + + private static Pair deleteRows(int n) { + final Timestamp firstCommitTimestamp = deleteRow(1); + for (int i = 2; i < n; i++) { + deleteRow(i); + } + final Timestamp lastCommitTimestamp = deleteRow(n); + return Pair.of(firstCommitTimestamp, lastCommitTimestamp); + } + + private static Timestamp insertRow(int singerId) { + return databaseClient + .writeWithOptions( + Collections.singletonList( + Mutation.newInsertBuilder(changeStreamTableName) + .set("SingerId") + .to(singerId) + .set("FirstName") + .to("First Name " + singerId) + .set("LastName") + .to("Last Name " + singerId) + .build())) + .getCommitTimestamp(); + } + + private static Timestamp updateRow(int singerId) { + return databaseClient + .writeWithOptions( + Collections.singletonList( + Mutation.newUpdateBuilder(changeStreamTableName) + .set("SingerId") + .to(singerId) + .set("FirstName") + .to("Updated First Name " + singerId) + .set("LastName") + .to("Updated Last Name " + singerId) + .build()), + Options.tag("app=beam;action=update")) + .getCommitTimestamp(); + } + + private static Timestamp deleteRow(int singerId) { + return databaseClient + .writeWithOptions( + Collections.singletonList(Mutation.delete(changeStreamTableName, Key.of(singerId))), + Options.tag("app=beam;action=delete")) + .getCommitTimestamp(); + } + + private static class ModsToString extends DoFn { + + private transient Gson gson; + + @Setup + public void setup() { + gson = new Gson(); + } + + @ProcessElement + public void processElement( + @Element DataChangeRecord record, OutputReceiver outputReceiver) { + final Mod mod = record.getMods().get(0); + final Map keys = gson.fromJson(mod.getKeysJson(), Map.class); + final Map oldValues = + Optional.ofNullable(mod.getOldValuesJson()) + .map(nonNullValues -> gson.fromJson(nonNullValues, Map.class)) + .orElseGet(Collections::emptyMap); + final Map newValues = + Optional.ofNullable(mod.getNewValuesJson()) + .map(nonNullValues -> gson.fromJson(nonNullValues, Map.class)) + .orElseGet(Collections::emptyMap); + + final String modsAsString = + String.join( + ",", + record.getModType().toString(), + keys.get("SingerId"), + oldValues.get("FirstName"), + oldValues.get("LastName"), + newValues.get("FirstName"), + newValues.get("LastName")); + final Instant timestamp = new Instant(record.getRecordTimestamp().toSqlTimestamp()); + + outputReceiver.outputWithTimestamp(modsAsString, timestamp); + } + } +} diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/mapper/ChangeStreamRecordMapperTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/mapper/ChangeStreamRecordMapperTest.java index b3dd1bef049f..f73647beb002 100644 --- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/mapper/ChangeStreamRecordMapperTest.java +++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/mapper/ChangeStreamRecordMapperTest.java @@ -1039,4 +1039,109 @@ public void testMappingProtoRowToDataChangeRecord() { Collections.singletonList(dataChangeRecord), mapper.toChangeStreamRecords(partition, resultSet, resultSetMetadata)); } + + @Test + public void testMappingProtoBytesRowToPartitionStartRecord() { + final PartitionStartRecord partitionStartRecord = + new PartitionStartRecord( + Timestamp.MIN_VALUE, + "fakeRecordSequence", + Arrays.asList("partitionToken1", "partitionToken2"), + null); + com.google.spanner.v1.ChangeStreamRecord changeStreamRecordProto = + recordToProto(partitionStartRecord); + assertNotNull(changeStreamRecordProto); + ChangeStreamResultSet resultSet = mock(ChangeStreamResultSet.class); + + when(resultSet.isProtoBytesChangeRecord()).thenReturn(true); + when(resultSet.getBytes(0)).thenReturn(changeStreamRecordProto); + assertEquals( + Collections.singletonList(partitionStartRecord), + mapperPostgres.toChangeStreamRecords(partition, resultSet, resultSetMetadata)); + } + + @Test + public void testMappingProtoBytesRowToPartitionEndRecord() { + final PartitionEndRecord partitionEndChange = + new PartitionEndRecord(Timestamp.MIN_VALUE, "fakeRecordSequence", null); + com.google.spanner.v1.ChangeStreamRecord changeStreamRecordProto = + recordToProto(partitionEndChange); + assertNotNull(changeStreamRecordProto); + ChangeStreamResultSet resultSet = mock(ChangeStreamResultSet.class); + + when(resultSet.isProtoBytesChangeRecord()).thenReturn(true); + when(resultSet.getBytes(0)).thenReturn(changeStreamRecordProto); + assertEquals( + Collections.singletonList(partitionEndChange), + mapperPostgres.toChangeStreamRecords(partition, resultSet, resultSetMetadata)); + } + + @Test + public void testMappingProtoBytesRowToPartitionEventRecord() { + final PartitionEventRecord partitionEventRecord = + new PartitionEventRecord(Timestamp.MIN_VALUE, "fakeRecordSequence", null); + com.google.spanner.v1.ChangeStreamRecord changeStreamRecordProto = + recordToProto(partitionEventRecord); + assertNotNull(changeStreamRecordProto); + ChangeStreamResultSet resultSet = mock(ChangeStreamResultSet.class); + + when(resultSet.isProtoBytesChangeRecord()).thenReturn(true); + when(resultSet.getBytes(0)).thenReturn(changeStreamRecordProto); + assertEquals( + Collections.singletonList(partitionEventRecord), + mapperPostgres.toChangeStreamRecords(partition, resultSet, resultSetMetadata)); + } + + @Test + public void testMappingProtoBytesRowToHeartbeatRecord() { + final HeartbeatRecord heartbeatRecord = + new HeartbeatRecord(Timestamp.ofTimeSecondsAndNanos(10L, 20), null); + com.google.spanner.v1.ChangeStreamRecord changeStreamRecordProto = + recordToProto(heartbeatRecord); + assertNotNull(changeStreamRecordProto); + ChangeStreamResultSet resultSet = mock(ChangeStreamResultSet.class); + + when(resultSet.isProtoBytesChangeRecord()).thenReturn(true); + when(resultSet.getBytes(0)).thenReturn(changeStreamRecordProto); + assertEquals( + Collections.singletonList(heartbeatRecord), + mapperPostgres.toChangeStreamRecords(partition, resultSet, resultSetMetadata)); + } + + @Test + public void testMappingProtoBytesRowToDataChangeRecord() { + final DataChangeRecord dataChangeRecord = + new DataChangeRecord( + "partitionToken", + Timestamp.ofTimeSecondsAndNanos(10L, 20), + "serverTransactionId", + true, + "1", + "tableName", + Arrays.asList( + new ColumnType("column1", new TypeCode("{\"code\":\"INT64\"}"), true, 1L), + new ColumnType("column2", new TypeCode("{\"code\":\"BYTES\"}"), false, 2L)), + Collections.singletonList( + new Mod( + "{\"column1\":\"value1\"}", + "{\"column2\":\"oldValue2\"}", + "{\"column2\":\"newValue2\"}")), + ModType.UPDATE, + ValueCaptureType.OLD_AND_NEW_VALUES, + 10L, + 2L, + "transactionTag", + true, + null); + com.google.spanner.v1.ChangeStreamRecord changeStreamRecordProto = + recordToProto(dataChangeRecord); + assertNotNull(changeStreamRecordProto); + ChangeStreamResultSet resultSet = mock(ChangeStreamResultSet.class); + + when(resultSet.isProtoBytesChangeRecord()).thenReturn(true); + when(resultSet.getBytes(0)).thenReturn(changeStreamRecordProto); + assertEquals( + Collections.singletonList(dataChangeRecord), + mapperPostgres.toChangeStreamRecords(partition, resultSet, resultSetMetadata)); + } } diff --git a/sdks/python/apache_beam/options/pipeline_options_context.py b/sdks/python/apache_beam/options/pipeline_options_context.py new file mode 100644 index 000000000000..5f2475e334af --- /dev/null +++ b/sdks/python/apache_beam/options/pipeline_options_context.py @@ -0,0 +1,65 @@ +# +# 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. +# + +"""Context-scoped access to pipeline options during graph construction and +translation. + +This module provides thread-safe and async-safe access to globally-available +instances of PipelineOptions using contextvars, scoped by the current pipeline. +It allows components like transforms and coders to access the pipeline's +configuration without requiring explicit parameter passing through every +level of the call stack. + +For internal use only; no backwards-compatibility guarantees. +""" + +from contextlib import contextmanager +from contextvars import ContextVar +from typing import TYPE_CHECKING +from typing import Optional + +if TYPE_CHECKING: + from apache_beam.options.pipeline_options import PipelineOptions + +# The contextvar holding the current pipeline's options. +# Each thread and each asyncio task gets its own isolated copy. +_pipeline_options: ContextVar[Optional['PipelineOptions']] = ContextVar( + 'pipeline_options', default=None) + + +def get_pipeline_options() -> Optional['PipelineOptions']: + """Get the current pipeline's options from the context. + + Returns: + The PipelineOptions for the currently executing pipeline operation, + or None if called outside of a pipeline context. + """ + return _pipeline_options.get() + + +@contextmanager +def scoped_pipeline_options(options: Optional['PipelineOptions']): + """Context manager that sets pipeline options for the duration of a block. + + Args: + options: The PipelineOptions to make available during this scope. + """ + token = _pipeline_options.set(options) + try: + yield + finally: + _pipeline_options.reset(token) diff --git a/sdks/python/apache_beam/options/pipeline_options_context_test.py b/sdks/python/apache_beam/options/pipeline_options_context_test.py new file mode 100644 index 000000000000..da9c8fce4aab --- /dev/null +++ b/sdks/python/apache_beam/options/pipeline_options_context_test.py @@ -0,0 +1,335 @@ +# +# 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. +# + +"""Tests for pipeline_options_context module. + +These tests verify that the contextvar-based approach properly isolates +pipeline options across threads and async tasks, preventing race conditions. +""" + +import asyncio +import threading +import unittest + +import apache_beam as beam +from apache_beam.options.pipeline_options import PipelineOptions +from apache_beam.options.pipeline_options_context import get_pipeline_options +from apache_beam.options.pipeline_options_context import scoped_pipeline_options + + +class PipelineConstructionOptionsTest(unittest.TestCase): + def test_nested_scoping(self): + """Test that nested scopes properly restore outer options.""" + outer_options = PipelineOptions(['--runner=DirectRunner']) + inner_options = PipelineOptions(['--runner=DataflowRunner']) + + with scoped_pipeline_options(outer_options): + self.assertIs(get_pipeline_options(), outer_options) + + with scoped_pipeline_options(inner_options): + self.assertIs(get_pipeline_options(), inner_options) + + self.assertIs(get_pipeline_options(), outer_options) + + self.assertIsNone(get_pipeline_options()) + + def test_exception_in_scope_restores_options(self): + """Test that options are restored even when an exception is raised.""" + outer_options = PipelineOptions(['--runner=DirectRunner']) + inner_options = PipelineOptions(['--runner=DataflowRunner']) + + with scoped_pipeline_options(outer_options): + try: + with scoped_pipeline_options(inner_options): + self.assertIs(get_pipeline_options(), inner_options) + raise ValueError("Test exception") + except ValueError: + pass + + self.assertIs(get_pipeline_options(), outer_options) + + def test_different_threads_see_their_own_isolated_options(self): + """Test that different threads see their own isolated options.""" + results = {} + errors = [] + barrier = threading.Barrier(2) + + def thread_worker(thread_id, runner_name): + try: + options = PipelineOptions([f'--runner={runner_name}']) + with scoped_pipeline_options(options): + barrier.wait(timeout=5) + + current = get_pipeline_options() + results[thread_id] = current.get_all_options()['runner'] + import time + time.sleep(0.01) + + current_after = get_pipeline_options() + if current_after is not current: + errors.append( + f"Thread {thread_id}: options changed during execution") + except Exception as e: + errors.append(f"Thread {thread_id}: {e}") + + thread1 = threading.Thread(target=thread_worker, args=(1, 'DirectRunner')) + thread2 = threading.Thread(target=thread_worker, args=(2, 'DataflowRunner')) + + thread1.start() + thread2.start() + + thread1.join(timeout=5) + thread2.join(timeout=5) + + self.assertEqual(errors, []) + self.assertEqual(results[1], 'DirectRunner') + self.assertEqual(results[2], 'DataflowRunner') + + def test_asyncio_task_isolation(self): + """Test that different asyncio tasks see their own isolated options.""" + async def async_worker( + task_id, runner_name, results, ready_event, go_event): + options = PipelineOptions([f'--runner={runner_name}']) + with scoped_pipeline_options(options): + ready_event.set() + await go_event.wait() + current = get_pipeline_options() + results[task_id] = current.get_all_options()['runner'] + await asyncio.sleep(0.01) + current_after = get_pipeline_options() + assert current_after is current, \ + f"Task {task_id}: options changed during execution" + + async def run_test(): + results = {} + ready_events = [asyncio.Event() for _ in range(2)] + go_event = asyncio.Event() + + task1 = asyncio.create_task( + async_worker(1, 'DirectRunner', results, ready_events[0], go_event)) + task2 = asyncio.create_task( + async_worker(2, 'DataflowRunner', results, ready_events[1], go_event)) + + # Wait for both tasks to be ready + await asyncio.gather(*[e.wait() for e in ready_events]) + # Signal all tasks to proceed + go_event.set() + + await asyncio.gather(task1, task2) + return results + + results = asyncio.run(run_test()) + self.assertEqual(results[1], 'DirectRunner') + self.assertEqual(results[2], 'DataflowRunner') + + def test_transform_sees_pipeline_options(self): + """Test that a transform can access pipeline options during expand().""" + class OptionsCapturingTransform(beam.PTransform): + """Transform that captures pipeline options during expand().""" + def __init__(self, expected_job_name): + self.expected_job_name = expected_job_name + self.captured_options = None + + def expand(self, pcoll): + # This runs during pipeline construction + self.captured_options = get_pipeline_options() + return pcoll | beam.Map(lambda x: x) + + options = PipelineOptions(['--job_name=test_job_123']) + transform = OptionsCapturingTransform('test_job_123') + + with beam.Pipeline(options=options) as p: + _ = p | beam.Create([1, 2, 3]) | transform + + # Verify the transform saw the correct options + self.assertIsNotNone(transform.captured_options) + self.assertEqual( + transform.captured_options.get_all_options()['job_name'], + 'test_job_123') + + def test_coder_sees_correct_options_during_run(self): + """Test that coders see correct pipeline options during proto conversion. + + This tests the run path where as_deterministic_coder() is called during + to_runner_api() proto conversion. + """ + from apache_beam.coders import coders + from apache_beam.utils import shared + + errors = [] + + class WeakRefDict(dict): + pass + + class TestKey: + def __init__(self, value): + self.value = value + + def __eq__(self, other): + return isinstance(other, TestKey) and self.value == other.value + + def __hash__(self): + return hash(self.value) + + class OptionsCapturingKeyCoder(coders.Coder): + """Coder that captures pipeline options in as_deterministic_coder.""" + shared_handle = shared.Shared() + + def encode(self, value): + return str(value.value).encode('utf-8') + + def decode(self, encoded): + return TestKey(encoded.decode('utf-8')) + + def is_deterministic(self): + return False + + def as_deterministic_coder(self, step_label, error_message=None): + opts = get_pipeline_options() + if opts is not None: + results = OptionsCapturingKeyCoder.shared_handle.acquire(WeakRefDict) + job_name = opts.get_all_options().get('job_name') + results['Worker1'] = job_name + return self + + beam.coders.registry.register_coder(TestKey, OptionsCapturingKeyCoder) + + results = OptionsCapturingKeyCoder.shared_handle.acquire(WeakRefDict) + + job_name = 'gbk_job' + options = PipelineOptions([f'--job_name={job_name}']) + + with beam.Pipeline(options=options) as p: + _ = ( + p + | beam.Create([(TestKey(1), 'a'), (TestKey(2), 'b')]) + | beam.GroupByKey()) + + self.assertEqual(errors, [], f"Errors occurred: {errors}") + self.assertEqual( + results.get('Worker1'), + job_name, + f"Worker1 saw wrong options: {results}") + self.assertFalse(get_pipeline_options() == options) + + def test_barrier_inside_default_type_hints(self): + """Test race condition detection with barrier inside default_type_hints. + + This test reliably detects race conditions because: + 1. Both threads start pipeline construction simultaneously + 2. Inside default_type_hints (which is called during Pipeline.apply()), + both threads hit a barrier and wait for each other + 3. At this point, BOTH threads are inside scoped_pipeline_options + 4. When they continue, they read options - with a global var, they'd see + the wrong values because the last thread to set options would win + """ + + results = {} + errors = [] + inner_barrier = threading.Barrier(2) + + class BarrierTransform(beam.PTransform): + """Transform that synchronizes threads INSIDE default_type_hints.""" + def __init__(self, worker_id, results_dict, barrier): + self.worker_id = worker_id + self.results_dict = results_dict + self.barrier = barrier + + def expand(self, pcoll): + return pcoll | beam.Map(lambda x: x) + + def default_type_hints(self): + self.barrier.wait(timeout=5) + opts = get_pipeline_options() + if opts is not None: + job_name = opts.get_all_options().get('job_name') + if self.worker_id not in self.results_dict: + self.results_dict[self.worker_id] = job_name + + return super().default_type_hints() + + def construct_pipeline(worker_id): + try: + job_name = f'barrier_job_{worker_id}' + options = PipelineOptions([f'--job_name={job_name}']) + transform = BarrierTransform(worker_id, results, inner_barrier) + + with beam.Pipeline(options=options) as p: + _ = p | beam.Create([1, 2, 3]) | transform + + except Exception as e: + import traceback + errors.append(f"Worker {worker_id}: {e}\n{traceback.format_exc()}") + + thread1 = threading.Thread( + target=construct_pipeline, args=(1, )) + thread2 = threading.Thread( + target=construct_pipeline, args=(2, )) + + thread1.start() + thread2.start() + + thread1.join(timeout=10) + thread2.join(timeout=10) + + self.assertEqual(errors, [], f"Errors occurred: {errors}") + self.assertEqual( + results.get(1), + 'barrier_job_1', + f"Worker 1 saw wrong options: {results}") + self.assertEqual( + results.get(2), + 'barrier_job_2', + f"Worker 2 saw wrong options: {results}") + + +class PipelineSubclassApplyTest(unittest.TestCase): + def test_subclass_apply_called_on_recursive_paths(self): + """Test that Pipeline subclass overrides of apply() are respected. + + _apply_internal's recursive calls must go through self.apply(), not + self._apply_internal(), so that subclass interceptions are not skipped. + """ + apply_calls = [] + + class TrackingPipeline(beam.Pipeline): + def apply(self, transform, pvalueish=None, label=None): + apply_calls.append(label or transform.label) + return super().apply(transform, pvalueish, label) + + options = PipelineOptions(['--job_name=subclass_test']) + with TrackingPipeline(options=options) as p: + # "my_label" >> transform creates a _NamedPTransform, which triggers + # two recursive apply() calls: one to unwrap _NamedPTransform, and + # one to handle the label argument. + _ = p | beam.Create([1, 2, 3]) | "my_label" >> beam.Map(lambda x: x) + + # beam.Create goes through apply() once (no recursion). + # "my_label" >> Map triggers: apply(_NamedPTransform) -> apply(Map, + # label="my_label") -> apply(Map). That's 3 calls through apply(). + # Total: 1 (Create) + 3 (Map) = 4 calls minimum. + map_calls = [c for c in apply_calls if c == 'my_label' or c == 'Map'] + self.assertGreaterEqual( + len(map_calls), + 3, + f"Expected at least 3 apply() calls for the Map transform " + f"(NamedPTransform unwrap + label handling + final), " + f"got {len(map_calls)}. All calls: {apply_calls}") + + +if __name__ == '__main__': + unittest.main() diff --git a/sdks/python/apache_beam/pipeline.py b/sdks/python/apache_beam/pipeline.py index 6ef06abb7436..a6080f2f3e7f 100644 --- a/sdks/python/apache_beam/pipeline.py +++ b/sdks/python/apache_beam/pipeline.py @@ -81,6 +81,7 @@ from apache_beam.options.pipeline_options import StandardOptions from apache_beam.options.pipeline_options import StreamingOptions from apache_beam.options.pipeline_options import TypeOptions +from apache_beam.options.pipeline_options_context import scoped_pipeline_options from apache_beam.options.pipeline_options_validator import PipelineOptionsValidator from apache_beam.portability import common_urns from apache_beam.portability.api import beam_runner_api_pb2 @@ -559,6 +560,12 @@ def replace_all(self, replacements: Iterable['PTransformOverride']) -> None: def run(self, test_runner_api: Union[bool, str] = 'AUTO') -> 'PipelineResult': """Runs the pipeline. Returns whatever our runner returns after running.""" + with scoped_pipeline_options(self._options): + return self._run_internal(test_runner_api) + + def _run_internal( + self, test_runner_api: Union[bool, str] = 'AUTO') -> 'PipelineResult': + """Internal implementation of run(), called within scoped options.""" # All pipeline options are finalized at this point. # Call get_all_options to print warnings on invalid options. self.options.get_all_options( @@ -698,6 +705,15 @@ def apply( RuntimeError: if the transform object was already applied to this pipeline and needs to be cloned in order to apply again. """ + with scoped_pipeline_options(self._options): + return self._apply_internal(transform, pvalueish, label) + + def _apply_internal( + self, + transform: ptransform.PTransform, + pvalueish: Optional[pvalue.PValue] = None, + label: Optional[str] = None) -> pvalue.PValue: + """Internal implementation of apply(), called within scoped options.""" if isinstance(transform, ptransform._NamedPTransform): return self.apply( transform.transform, pvalueish, label or transform.label)