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());
+ }
}