add filter to watermark query

Author: changliiu

sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/cache/AsyncWatermarkCache.java

@@ -42,13 +42,23 @@ public class AsyncWatermarkCache implements WatermarkCache {
   private static final Object MIN_WATERMARK_KEY = new Object();
   private final LoadingCache<Object, Optional<Timestamp>> cache;
 
+  private Timestamp lastCachedMinWatermark = Timestamp.MIN_VALUE;
+
   public AsyncWatermarkCache(PartitionMetadataDao dao, Duration refreshRate) {
     this.cache =
         CacheBuilder.newBuilder()
             .refreshAfterWrite(java.time.Duration.ofMillis(refreshRate.getMillis()))
             .build(
                 CacheLoader.asyncReloading(
-                    CacheLoader.from(key -> Optional.ofNullable(dao.getUnfinishedMinWatermark())),
+                    CacheLoader.from(
+                        key -> {
+                          Timestamp unfinishedMinTimes =
+                              dao.getUnfinishedMinWatermark(Optional.of(lastCachedMinWatermark));
+                          if (unfinishedMinTimes != null) {
+                            lastCachedMinWatermark = unfinishedMinTimes;
+                          }
+                          return Optional.ofNullable(unfinishedMinTimes);
+                        }),
                     Executors.newSingleThreadExecutor(
                         new ThreadFactoryBuilder().setNameFormat(THREAD_NAME_FORMAT).build())));
   }

sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/cache/NoOpWatermarkCache.java

@@ -18,6 +18,7 @@
 package org.apache.beam.sdk.io.gcp.spanner.changestreams.cache;
 
 import com.google.cloud.Timestamp;
+import java.util.Optional;
 import javax.annotation.Nullable;
 import org.apache.beam.sdk.io.gcp.spanner.changestreams.dao.PartitionMetadataDao;
 
@@ -35,6 +36,6 @@ public NoOpWatermarkCache(PartitionMetadataDao dao) {
 
   @Override
   public @Nullable Timestamp getUnfinishedMinWatermark() {
-    return dao.getUnfinishedMinWatermark();
+    return dao.getUnfinishedMinWatermark(Optional.empty());
   }
 }

sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/dao/PartitionMetadataDao.java

@@ -47,6 +47,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.function.Function;
+import java.util.Optional;
 import java.util.stream.Collectors;
 import javax.annotation.Nullable;
 import org.apache.beam.sdk.io.gcp.spanner.changestreams.model.PartitionMetadata;
@@ -178,47 +179,56 @@ public List<String> findAllTableIndexes() {
    *
    * @return the earliest partition watermark which is not in a {@link State#FINISHED} state.
    */
-  public @Nullable Timestamp getUnfinishedMinWatermark() {
+  public @Nullable Timestamp getUnfinishedMinWatermark(Optional<Timestamp> since) {
+    Timestamp sinceTimestamp = since.orElse(Timestamp.MIN_VALUE);
     Statement statement;
+    final String MIN_WATERMARK = "min_watermark";
     if (this.isPostgres()) {
       statement =
           Statement.newBuilder(
-                  "SELECT \""
+                  "SELECT min(\""
                       + COLUMN_WATERMARK
-                      + "\" FROM \""
+                      + "\") as "
+                      + MIN_WATERMARK
+                      + " FROM \""
                       + metadataTableName
                       + "\" WHERE \""
                       + COLUMN_STATE
                       + "\" != $1"
-                      + " ORDER BY \""
+                      + " AND \""
                       + COLUMN_WATERMARK
-                      + "\" ASC LIMIT 1")
+                      + "\" >= $2")
               .bind("p1")
               .to(State.FINISHED.name())
+              .bind("p2")
+              .to(sinceTimestamp)
               .build();
     } else {
       statement =
-          Statement.newBuilder(
-                  "SELECT "
-                      + COLUMN_WATERMARK
-                      + " FROM "
-                      + metadataTableName
-                      + " WHERE "
-                      + COLUMN_STATE
-                      + " != @state"
-                      + " ORDER BY "
-                      + COLUMN_WATERMARK
-                      + " ASC LIMIT 1")
+          "SELECT min("
+              + COLUMN_WATERMARK
+              + ") as "
+              + MIN_WATERMARK
+              + " FROM "
+              + metadataTableName
+              + " WHERE "
+              + COLUMN_STATE
+              + " != @state"
+              + " AND "
+              + COLUMN_WATERMARK
+              + " >= @since;")
               .bind("state")
-              .to(State.FINISHED.name())
-              .build();
+          .to(State.FINISHED.name())
+          .bind("since")
+          .to(sinceTimestamp)
+          .build();
     }
     try (ResultSet resultSet =
         databaseClient
             .singleUse()
             .executeQuery(statement, Options.tag("query=getUnfinishedMinWatermark"))) {
       if (resultSet.next()) {
-        return resultSet.getTimestamp(COLUMN_WATERMARK);
+        return resultSet.getTimestamp(MIN_WATERMARK);
       }
       return null;
     }

sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/SpannerChangeStreamErrorTest.java

@@ -460,19 +460,27 @@ private void mockGetPartitionsAfter(Timestamp timestamp, ResultSet getPartitionR
   }
 
   private void mockGetWatermark(Timestamp watermark) {
+    final String MIN_WATERMARK = "min_watermark";
+    // The query needs to sync with getUnfinishedMinWatermark() in PartitionMetadataDao file.
     Statement watermarkStatement =
         Statement.newBuilder(
-                "SELECT Watermark FROM my-metadata-table WHERE State != @state ORDER BY Watermark ASC LIMIT 1")
+                "SELECT min(Watermark) as "
+                    + MIN_WATERMARK
+                    + " FROM my-metadata-table "
+                    + "WHERE State != @state "
+                    + "AND Watermark >= @since;")
             .bind("state")
             .to(State.FINISHED.name())
+            .bind("since")
+            .to(Timestamp.MIN_VALUE)
             .build();
     ResultSetMetadata watermarkResultSetMetadata =
         ResultSetMetadata.newBuilder()
             .setRowType(
                 StructType.newBuilder()
                     .addFields(
                         Field.newBuilder()
-                            .setName("Watermark")
+                            .setName(MIN_WATERMARK)
                             .setType(Type.newBuilder().setCode(TypeCode.TIMESTAMP).build())
                             .build())
                     .build())