From e888d07e93b8c93f071190b2b52ebb2112b92fc0 Mon Sep 17 00:00:00 2001 From: Jiang Zhu Date: Thu, 29 Oct 2026 15:59:34 +0000 Subject: [PATCH] [Spanner Change Stream] Use inclusive CreatedAt filter for partition discovery on split boundaries --- .../action/DetectNewPartitionsAction.java | 2 +- .../dao/PartitionMetadataDao.java | 16 ++-- .../SpannerChangeStreamErrorTest.java | 4 +- .../dao/PartitionMetadataDaoTest.java | 88 +++++++++++++++++++ 4 files changed, 99 insertions(+), 11 deletions(-) diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/action/DetectNewPartitionsAction.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/action/DetectNewPartitionsAction.java index e3f291ead74b..ea5aa42fdd9f 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/action/DetectNewPartitionsAction.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/action/DetectNewPartitionsAction.java @@ -77,7 +77,7 @@ public DetectNewPartitionsAction( *
  • If there are no unfinished partitions, this function will stop and not be re-scheduled. *
  • Updates the component's watermark to the min fetched. *
  • Fetches the read timestamp from the restriction. - *
  • Fetches all the partitions with a createdAt timestamp > read timestamp. + *
  • Fetches all the partitions with a createdAt timestamp >= read timestamp. *
  • Groups the partitions by createdAt timestamp. *
  • Process the groups in ascending order of createdAt timestamp (oldest first) *
  • For each group, updates the state to {@link State#SCHEDULED}. diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/dao/PartitionMetadataDao.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/dao/PartitionMetadataDao.java index 7b0825cf615a..79086ee0886e 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/dao/PartitionMetadataDao.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/dao/PartitionMetadataDao.java @@ -269,8 +269,8 @@ public static String extractTvfName(String composedPartitionToken) { } /** - * Fetches all partitions with a {@link PartitionMetadataAdminDao#COLUMN_CREATED_AT} less than the - * given timestamp. The results are ordered by the {@link + * Fetches all partitions with a {@link PartitionMetadataAdminDao#COLUMN_CREATED_AT} greater than + * or equal to the given timestamp. The results are ordered by the {@link * PartitionMetadataAdminDao#COLUMN_CREATED_AT} and {@link * PartitionMetadataAdminDao#COLUMN_START_TIMESTAMP} columns in ascending order. */ @@ -283,7 +283,7 @@ public ResultSet getAllPartitionsCreatedAfter(Timestamp timestamp) { + metadataTableName + "\" WHERE \"" + COLUMN_CREATED_AT - + "\" > $1" + + "\" >= $1" + " ORDER BY \"" + COLUMN_CREATED_AT + "\" ASC" @@ -300,7 +300,7 @@ public ResultSet getAllPartitionsCreatedAfter(Timestamp timestamp) { + metadataTableName + " WHERE " + COLUMN_CREATED_AT - + " > @timestamp" + + " >= @timestamp" + " ORDER BY " + COLUMN_CREATED_AT + " ASC" @@ -317,8 +317,8 @@ public ResultSet getAllPartitionsCreatedAfter(Timestamp timestamp) { } /** - * Counts all partitions with a {@link PartitionMetadataAdminDao#COLUMN_CREATED_AT} less than the - * given timestamp. + * Counts all partitions with a {@link PartitionMetadataAdminDao#COLUMN_CREATED_AT} greater than + * or equal to the given timestamp. */ public long countPartitionsCreatedAfter(Timestamp timestamp) { Statement statement; @@ -329,7 +329,7 @@ public long countPartitionsCreatedAfter(Timestamp timestamp) { + metadataTableName + "\" WHERE \"" + COLUMN_CREATED_AT - + "\" > $1") + + "\" >= $1") .bind("p1") .to(timestamp) .build(); @@ -340,7 +340,7 @@ public long countPartitionsCreatedAfter(Timestamp timestamp) { + metadataTableName + " WHERE " + COLUMN_CREATED_AT - + " > @timestamp") + + " >= @timestamp") .bind("timestamp") .to(timestamp) .build(); diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/SpannerChangeStreamErrorTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/SpannerChangeStreamErrorTest.java index 9f09ab18c62e..f6a9ea8ed4ed 100644 --- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/SpannerChangeStreamErrorTest.java +++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/SpannerChangeStreamErrorTest.java @@ -490,7 +490,7 @@ private void mockInvalidChangeStreamRecordReceived(Timestamp now, Timestamp afte private void mockGetPartitionsAfter(Timestamp timestamp, ResultSet getPartitionResultSet) { Statement getPartitionsAfterStatement = Statement.newBuilder( - "SELECT * FROM my-metadata-table WHERE CreatedAt > @timestamp ORDER BY CreatedAt ASC, StartTimestamp ASC") + "SELECT * FROM my-metadata-table WHERE CreatedAt >= @timestamp ORDER BY CreatedAt ASC, StartTimestamp ASC") .bind("timestamp") .to(Timestamp.ofTimeSecondsAndNanos(timestamp.getSeconds(), timestamp.getNanos())) .build(); @@ -501,7 +501,7 @@ private void mockGetPartitionsAfter(Timestamp timestamp, ResultSet getPartitionR private void mockGetPartitionSize(Timestamp timestamp, long partitionSize) { Statement getPartitionsAfterStatement = Statement.newBuilder( - "SELECT COUNT(*) as count FROM my-metadata-table WHERE CreatedAt > @timestamp") + "SELECT COUNT(*) as count FROM my-metadata-table WHERE CreatedAt >= @timestamp") .bind("timestamp") .to(Timestamp.ofTimeSecondsAndNanos(timestamp.getSeconds(), timestamp.getNanos())) .build(); diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/dao/PartitionMetadataDaoTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/dao/PartitionMetadataDaoTest.java index 3b2b550ab1db..ac822767b01e 100644 --- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/dao/PartitionMetadataDaoTest.java +++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/dao/PartitionMetadataDaoTest.java @@ -21,6 +21,7 @@ import static org.junit.Assert.assertFalse; import static org.junit.Assert.assertNotNull; import static org.junit.Assert.assertNull; +import static org.junit.Assert.assertSame; import static org.junit.Assert.assertTrue; import static org.mockito.ArgumentMatchers.any; import static org.mockito.Mockito.doNothing; @@ -33,6 +34,7 @@ import com.google.cloud.spanner.DatabaseClient; import com.google.cloud.spanner.Dialect; import com.google.cloud.spanner.Mutation; +import com.google.cloud.spanner.ReadContext; import com.google.cloud.spanner.ReadOnlyTransaction; import com.google.cloud.spanner.ResultSet; import com.google.cloud.spanner.Statement; @@ -579,4 +581,90 @@ public void testInTransactionContextUpdateToScheduledPostgresMultipleTokens() { .get(PartitionMetadataAdminDao.COLUMN_STATE) .getString()); } + + @Test + public void testGetAllPartitionsCreatedAfterGoogleSql() { + ReadContext readContext = mock(ReadContext.class); + ResultSet resultSet = mock(ResultSet.class); + when(databaseClient.singleUse()).thenReturn(readContext); + when(readContext.executeQuery(any(Statement.class), any())).thenReturn(resultSet); + + assertSame(resultSet, partitionMetadataDao.getAllPartitionsCreatedAfter(CREATED_AT)); + + ArgumentCaptor statementCaptor = ArgumentCaptor.forClass(Statement.class); + verify(readContext).executeQuery(statementCaptor.capture(), any()); + Statement captured = statementCaptor.getValue(); + assertEquals( + "SELECT * FROM " + + METADATA_TABLE_NAME + + " WHERE CreatedAt >= @timestamp ORDER BY CreatedAt ASC, StartTimestamp ASC", + captured.getSql()); + assertEquals(CREATED_AT, captured.getParameters().get("timestamp").getTimestamp()); + } + + @Test + public void testGetAllPartitionsCreatedAfterPostgres() { + ReadContext readContext = mock(ReadContext.class); + ResultSet resultSet = mock(ResultSet.class); + when(databaseClient.singleUse()).thenReturn(readContext); + when(readContext.executeQuery(any(Statement.class), any())).thenReturn(resultSet); + + PartitionMetadataDao pgDao = + new PartitionMetadataDao(METADATA_TABLE_NAME, databaseClient, Dialect.POSTGRESQL); + assertSame(resultSet, pgDao.getAllPartitionsCreatedAfter(CREATED_AT)); + + ArgumentCaptor statementCaptor = ArgumentCaptor.forClass(Statement.class); + verify(readContext).executeQuery(statementCaptor.capture(), any()); + Statement captured = statementCaptor.getValue(); + assertEquals( + "SELECT * FROM \"" + + METADATA_TABLE_NAME + + "\" WHERE \"CreatedAt\" >= $1 ORDER BY \"CreatedAt\" ASC, \"StartTimestamp\" ASC", + captured.getSql()); + assertEquals(CREATED_AT, captured.getParameters().get("p1").getTimestamp()); + } + + @Test + public void testCountPartitionsCreatedAfterGoogleSql() { + ReadContext readContext = mock(ReadContext.class); + ResultSet resultSet = mock(ResultSet.class); + when(databaseClient.singleUse()).thenReturn(readContext); + when(readContext.executeQuery(any(Statement.class), any())).thenReturn(resultSet); + when(resultSet.next()).thenReturn(true).thenReturn(false); + when(resultSet.getLong("count")).thenReturn(5L); + + assertEquals(5L, partitionMetadataDao.countPartitionsCreatedAfter(CREATED_AT)); + assertEquals(0L, partitionMetadataDao.countPartitionsCreatedAfter(CREATED_AT)); + + ArgumentCaptor statementCaptor = ArgumentCaptor.forClass(Statement.class); + verify(readContext, times(2)).executeQuery(statementCaptor.capture(), any()); + Statement captured = statementCaptor.getValue(); + assertEquals( + "SELECT COUNT(*) as count FROM " + METADATA_TABLE_NAME + " WHERE CreatedAt >= @timestamp", + captured.getSql()); + assertEquals(CREATED_AT, captured.getParameters().get("timestamp").getTimestamp()); + } + + @Test + public void testCountPartitionsCreatedAfterPostgres() { + ReadContext readContext = mock(ReadContext.class); + ResultSet resultSet = mock(ResultSet.class); + when(databaseClient.singleUse()).thenReturn(readContext); + when(readContext.executeQuery(any(Statement.class), any())).thenReturn(resultSet); + when(resultSet.next()).thenReturn(true).thenReturn(false); + when(resultSet.getLong("count")).thenReturn(3L); + + PartitionMetadataDao pgDao = + new PartitionMetadataDao(METADATA_TABLE_NAME, databaseClient, Dialect.POSTGRESQL); + assertEquals(3L, pgDao.countPartitionsCreatedAfter(CREATED_AT)); + assertEquals(0L, pgDao.countPartitionsCreatedAfter(CREATED_AT)); + + ArgumentCaptor statementCaptor = ArgumentCaptor.forClass(Statement.class); + verify(readContext, times(2)).executeQuery(statementCaptor.capture(), any()); + Statement captured = statementCaptor.getValue(); + assertEquals( + "SELECT COUNT(*) as count FROM \"" + METADATA_TABLE_NAME + "\" WHERE \"CreatedAt\" >= $1", + captured.getSql()); + assertEquals(CREATED_AT, captured.getParameters().get("p1").getTimestamp()); + } }