[Spanner Change Stream] Use inclusive CreatedAt filter for partition discovery on split boundaries

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(
    *   <li>If there are no unfinished partitions, this function will stop and not be re-scheduled.
    *   <li>Updates the component's watermark to the min fetched.
    *   <li>Fetches the read timestamp from the restriction.
-   *   <li>Fetches all the partitions with a createdAt timestamp > read timestamp.
+   *   <li>Fetches all the partitions with a createdAt timestamp >= read timestamp.
    *   <li>Groups the partitions by createdAt timestamp.
    *   <li>Process the groups in ascending order of createdAt timestamp (oldest first)
    *   <li>For each group, updates the state to {@link State#SCHEDULED}.

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

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

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