Skip to content

Commit 465f29d

Browse files
changliiushunping
changliiu
authored and
shunping
committed
add filter to watermark query (apache#35042)
1 parent 81e1b7c commit 465f29d

File tree

11 files changed

+72
-15
lines changed

11 files changed

+72
-15
lines changed

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

Lines changed: 20 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121
import java.util.Optional;
2222
import java.util.concurrent.ExecutionException;
2323
import java.util.concurrent.Executors;
24+
import java.util.concurrent.atomic.AtomicReference;
2425
import javax.annotation.Nullable;
2526
import org.apache.beam.sdk.io.gcp.spanner.changestreams.dao.PartitionMetadataDao;
2627
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.cache.CacheBuilder;
@@ -42,13 +43,31 @@ public class AsyncWatermarkCache implements WatermarkCache {
4243
private static final Object MIN_WATERMARK_KEY = new Object();
4344
private final LoadingCache<Object, Optional<Timestamp>> cache;
4445

46+
// This is to cache the result of getUnfinishedMinWatermark query and filter the query in the next
47+
// run. For the initial query, the value of this cache is min timestamp. If there is no partition
48+
// in the metadata table, then this cache will not be updated. If the getUnfinishedMinWatermark
49+
// query fails or times out, then this cache will not be updated.
50+
// Note that, all the reload operations on this key are serialized due to use of the single
51+
// threaded async reloading executor.
52+
private AtomicReference<Timestamp> lastCachedMinWatermark =
53+
new AtomicReference<>(Timestamp.MIN_VALUE);
54+
4555
public AsyncWatermarkCache(PartitionMetadataDao dao, Duration refreshRate) {
4656
this.cache =
4757
CacheBuilder.newBuilder()
4858
.refreshAfterWrite(java.time.Duration.ofMillis(refreshRate.getMillis()))
4959
.build(
5060
CacheLoader.asyncReloading(
51-
CacheLoader.from(key -> Optional.ofNullable(dao.getUnfinishedMinWatermark())),
61+
CacheLoader.from(
62+
key -> {
63+
Timestamp unfinishedMinTimes =
64+
dao.getUnfinishedMinWatermarkFrom(lastCachedMinWatermark.get());
65+
if (unfinishedMinTimes != null
66+
&& lastCachedMinWatermark.get().compareTo(unfinishedMinTimes) < 0) {
67+
lastCachedMinWatermark.set(unfinishedMinTimes);
68+
}
69+
return Optional.ofNullable(unfinishedMinTimes);
70+
}),
5271
Executors.newSingleThreadExecutor(
5372
new ThreadFactoryBuilder().setNameFormat(THREAD_NAME_FORMAT).build())));
5473
}

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,6 @@ public NoOpWatermarkCache(PartitionMetadataDao dao) {
3535

3636
@Override
3737
public @Nullable Timestamp getUnfinishedMinWatermark() {
38-
return dao.getUnfinishedMinWatermark();
38+
return dao.getUnfinishedMinWatermarkFrom(Timestamp.MIN_VALUE);
3939
}
4040
}

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

Lines changed: 20 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -178,47 +178,56 @@ public List<String> findAllTableIndexes() {
178178
*
179179
* @return the earliest partition watermark which is not in a {@link State#FINISHED} state.
180180
*/
181-
public @Nullable Timestamp getUnfinishedMinWatermark() {
181+
public @Nullable Timestamp getUnfinishedMinWatermarkFrom(Timestamp sinceTimestamp) {
182182
Statement statement;
183+
final String minWatermark = "min_watermark";
183184
if (this.isPostgres()) {
184185
statement =
185186
Statement.newBuilder(
186-
"SELECT \""
187+
"SELECT MIN(\""
187188
+ COLUMN_WATERMARK
188-
+ "\" FROM \""
189+
+ "\") as "
190+
+ minWatermark
191+
+ " FROM \""
189192
+ metadataTableName
190193
+ "\" WHERE \""
191194
+ COLUMN_STATE
192195
+ "\" != $1"
193-
+ " ORDER BY \""
196+
+ " AND \""
194197
+ COLUMN_WATERMARK
195-
+ "\" ASC LIMIT 1")
198+
+ "\" >= $2")
196199
.bind("p1")
197200
.to(State.FINISHED.name())
201+
.bind("p2")
202+
.to(sinceTimestamp)
198203
.build();
199204
} else {
200205
statement =
201206
Statement.newBuilder(
202-
"SELECT "
207+
"SELECT MIN("
203208
+ COLUMN_WATERMARK
209+
+ ") as "
210+
+ minWatermark
204211
+ " FROM "
205212
+ metadataTableName
206213
+ " WHERE "
207214
+ COLUMN_STATE
208215
+ " != @state"
209-
+ " ORDER BY "
216+
+ " AND "
210217
+ COLUMN_WATERMARK
211-
+ " ASC LIMIT 1")
218+
+ " >= @since;")
212219
.bind("state")
213220
.to(State.FINISHED.name())
221+
.bind("since")
222+
.to(sinceTimestamp)
214223
.build();
215224
}
216225
try (ResultSet resultSet =
217226
databaseClient
218227
.singleUse()
219-
.executeQuery(statement, Options.tag("query=getUnfinishedMinWatermark"))) {
220-
if (resultSet.next()) {
221-
return resultSet.getTimestamp(COLUMN_WATERMARK);
228+
.executeQuery(statement, Options.tag("query=getUnfinishedMinWatermarkFrom"))) {
229+
if (resultSet.next() && !resultSet.isNull(minWatermark)) {
230+
return resultSet.getTimestamp(minWatermark);
222231
}
223232
return null;
224233
}

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

Lines changed: 10 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -460,19 +460,27 @@ private void mockGetPartitionsAfter(Timestamp timestamp, ResultSet getPartitionR
460460
}
461461

462462
private void mockGetWatermark(Timestamp watermark) {
463+
final String minWatermark = "min_watermark";
464+
// The query needs to sync with getUnfinishedMinWatermark() in PartitionMetadataDao file.
463465
Statement watermarkStatement =
464466
Statement.newBuilder(
465-
"SELECT Watermark FROM my-metadata-table WHERE State != @state ORDER BY Watermark ASC LIMIT 1")
467+
"SELECT MIN(Watermark) as "
468+
+ minWatermark
469+
+ " FROM my-metadata-table "
470+
+ "WHERE State != @state "
471+
+ "AND Watermark >= @since;")
466472
.bind("state")
467473
.to(State.FINISHED.name())
474+
.bind("since")
475+
.to(Timestamp.MIN_VALUE)
468476
.build();
469477
ResultSetMetadata watermarkResultSetMetadata =
470478
ResultSetMetadata.newBuilder()
471479
.setRowType(
472480
StructType.newBuilder()
473481
.addFields(
474482
Field.newBuilder()
475-
.setName("Watermark")
483+
.setName(minWatermark)
476484
.setType(Type.newBuilder().setCode(TypeCode.TIMESTAMP).build())
477485
.build())
478486
.build())

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

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -60,6 +60,7 @@
6060
import org.junit.Rule;
6161
import org.junit.Test;
6262
import org.junit.rules.ExpectedException;
63+
import org.junit.rules.Timeout;
6364
import org.junit.runner.RunWith;
6465
import org.junit.runners.JUnit4;
6566

@@ -82,6 +83,8 @@ public class SpannerChangeStreamIT {
8283
private static String changeStreamName;
8384
private static DatabaseClient databaseClient;
8485

86+
@Rule public transient Timeout globalTimeout = Timeout.seconds(3600);
87+
8588
@BeforeClass
8689
public static void beforeClass() throws Exception {
8790
projectId = ENV.getProjectId();

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

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -58,6 +58,7 @@
5858
import org.junit.ClassRule;
5959
import org.junit.Rule;
6060
import org.junit.Test;
61+
import org.junit.rules.Timeout;
6162
import org.junit.runner.RunWith;
6263
import org.junit.runners.JUnit4;
6364
import org.slf4j.Logger;
@@ -83,6 +84,8 @@ public class SpannerChangeStreamOrderedByTimestampAndTransactionIdIT {
8384
private static String changeStreamName;
8485
private static DatabaseClient databaseClient;
8586

87+
@Rule public transient Timeout globalTimeout = Timeout.seconds(3600);
88+
8689
@BeforeClass
8790
public static void setup() throws InterruptedException, ExecutionException, TimeoutException {
8891
projectId = ENV.getProjectId();

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

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -55,6 +55,7 @@
5555
import org.junit.ClassRule;
5656
import org.junit.Rule;
5757
import org.junit.Test;
58+
import org.junit.rules.Timeout;
5859
import org.junit.runner.RunWith;
5960
import org.junit.runners.JUnit4;
6061
import org.slf4j.Logger;
@@ -67,6 +68,8 @@ public class SpannerChangeStreamOrderedWithinKeyGloballyIT {
6768
private static final Logger LOG =
6869
LoggerFactory.getLogger(SpannerChangeStreamOrderedWithinKeyGloballyIT.class);
6970

71+
@Rule public transient Timeout globalTimeout = Timeout.seconds(3600);
72+
7073
@ClassRule public static final IntegrationTestEnv ENV = new IntegrationTestEnv();
7174
@Rule public final transient TestPipeline pipeline = TestPipeline.create();
7275

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

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -49,6 +49,7 @@
4949
import org.junit.ClassRule;
5050
import org.junit.Rule;
5151
import org.junit.Test;
52+
import org.junit.rules.Timeout;
5253
import org.junit.runner.RunWith;
5354
import org.junit.runners.JUnit4;
5455
import org.slf4j.Logger;
@@ -71,6 +72,8 @@ public class SpannerChangeStreamOrderedWithinKeyIT {
7172
private static String changeStreamName;
7273
private static DatabaseClient databaseClient;
7374

75+
@Rule public transient Timeout globalTimeout = Timeout.seconds(3600);
76+
7477
@BeforeClass
7578
public static void setup() throws InterruptedException, ExecutionException, TimeoutException {
7679
projectId = ENV.getProjectId();

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

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -51,13 +51,16 @@
5151
import org.junit.ClassRule;
5252
import org.junit.Rule;
5353
import org.junit.Test;
54+
import org.junit.rules.Timeout;
5455
import org.junit.runner.RunWith;
5556
import org.junit.runners.JUnit4;
5657

5758
/** End-to-end test of Cloud Spanner CDC Source. */
5859
@RunWith(JUnit4.class)
5960
public class SpannerChangeStreamPostgresIT {
6061

62+
@Rule public transient Timeout globalTimeout = Timeout.seconds(3600);
63+
6164
@ClassRule
6265
public static final IntegrationTestEnv ENV = new IntegrationTestEnv(/*isPostgres=*/ true);
6366

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

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -51,6 +51,7 @@
5151
import org.junit.ClassRule;
5252
import org.junit.Rule;
5353
import org.junit.Test;
54+
import org.junit.rules.Timeout;
5455
import org.junit.runner.RunWith;
5556
import org.junit.runners.JUnit4;
5657
import org.slf4j.Logger;
@@ -60,6 +61,8 @@
6061
@RunWith(JUnit4.class)
6162
public class SpannerChangeStreamTransactionBoundariesIT {
6263

64+
@Rule public transient Timeout globalTimeout = Timeout.seconds(3600);
65+
6366
private static final Logger LOG =
6467
LoggerFactory.getLogger(SpannerChangeStreamTransactionBoundariesIT.class);
6568

0 commit comments

Comments
 (0)