Skip to content

Commit a5d05b6

Browse files
tianz101tianz101
committed
WIP
1 parent 4cdcf04 commit a5d05b6

25 files changed

Lines changed: 606 additions & 173 deletions

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

Lines changed: 74 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -527,6 +527,11 @@ public static Write write() {
527527
.build();
528528
}
529529

530+
/** Returns the given string without quotes. */
531+
public static String escapeQuotes(String str) {
532+
return str.replace("'", "").replace("\"", "").replace("`", "");
533+
}
534+
530535
/**
531536
* Creates an uninitialized instance of {@link ReadChangeStream}. Before use, the {@link
532537
* ReadChangeStream} must be configured with a {@link ReadChangeStream#withProjectId}, {@link
@@ -1746,6 +1751,8 @@ public abstract static class ReadChangeStream
17461751

17471752
abstract String getChangeStreamName();
17481753

1754+
abstract @Nullable List<String> getTvfNameList();
1755+
17491756
abstract @Nullable String getMetadataInstance();
17501757

17511758
abstract @Nullable String getMetadataDatabase();
@@ -1783,6 +1790,8 @@ abstract static class Builder {
17831790

17841791
abstract Builder setChangeStreamName(String changeStreamName);
17851792

1793+
abstract Builder setTvfNameList(List<String> tvfNameList);
1794+
17861795
abstract Builder setMetadataInstance(String metadataInstance);
17871796

17881797
abstract Builder setMetadataDatabase(String metadataDatabase);
@@ -1861,6 +1870,11 @@ public ReadChangeStream withChangeStreamName(String changeStreamName) {
18611870
return toBuilder().setChangeStreamName(changeStreamName).build();
18621871
}
18631872

1873+
/** Specifies the list of TVF names to query and union. */
1874+
public ReadChangeStream withTvfNameList(List<String> tvfNameList) {
1875+
return toBuilder().setTvfNameList(tvfNameList).build();
1876+
}
1877+
18641878
/** Specifies the metadata database. */
18651879
public ReadChangeStream withMetadataInstance(String metadataInstance) {
18661880
return toBuilder().setMetadataInstance(metadataInstance).build();
@@ -2042,6 +2056,7 @@ && getInclusiveStartAt().toSqlTimestamp().after(getInclusiveEndAt().toSqlTimesta
20422056
getInclusiveEndAt().compareTo(MAX_INCLUSIVE_END_AT) > 0
20432057
? MAX_INCLUSIVE_END_AT
20442058
: getInclusiveEndAt();
2059+
final List<String> tvfNameList = getTvfNameList();
20452060
final MapperFactory mapperFactory = new MapperFactory(changeStreamDatabaseDialect);
20462061
final ChangeStreamMetrics metrics = new ChangeStreamMetrics();
20472062
final RpcPriority rpcPriority = MoreObjects.firstNonNull(getRpcPriority(), RpcPriority.HIGH);
@@ -2051,10 +2066,24 @@ && getInclusiveStartAt().toSqlTimestamp().after(getInclusiveEndAt().toSqlTimesta
20512066
isMutableChangeStream(
20522067
spannerAccessor.getDatabaseClient(), changeStreamDatabaseDialect, changeStreamName);
20532068
LOG.info("The change stream {} is mutable: {}", changeStreamName, isMutableChangeStream);
2069+
List<String> quoteEscapedTvfNameList = null;
2070+
if (tvfNameList != null && !tvfNameList.isEmpty()) {
2071+
if (!isMutableChangeStream) {
2072+
throw new IllegalArgumentException(
2073+
"tvfNameList is only supported for change streams with MUTABLE_KEY_RANGE mode");
2074+
}
2075+
// TODO: if !per_placement_tvf=true, throw exception.
2076+
quoteEscapedTvfNameList = new ArrayList<>();
2077+
for (String tvfName : tvfNameList) {
2078+
quoteEscapedTvfNameList.add(escapeQuotes(tvfName));
2079+
}
2080+
checkTvfExistence(spannerAccessor.getDatabaseClient(), quoteEscapedTvfNameList);
2081+
}
20542082
final DaoFactory daoFactory =
20552083
new DaoFactory(
20562084
changeStreamSpannerConfig,
20572085
changeStreamName,
2086+
quoteEscapedTvfNameList,
20582087
partitionMetadataSpannerConfig,
20592088
partitionMetadataTableNames,
20602089
rpcPriority,
@@ -2754,6 +2783,51 @@ static String resolveSpannerProjectId(SpannerConfig config) {
27542783
: config.getProjectId().get();
27552784
}
27562785

2786+
@VisibleForTesting
2787+
static void checkTvfExistence(
2788+
DatabaseClient databaseClient, List<String> quoteEscapedTvfNameList) {
2789+
if (quoteEscapedTvfNameList == null || quoteEscapedTvfNameList.isEmpty()) {
2790+
return;
2791+
}
2792+
Dialect dialect = databaseClient.getDialect();
2793+
try (ReadOnlyTransaction tx = databaseClient.readOnlyTransaction()) {
2794+
StringBuilder sql =
2795+
new StringBuilder(
2796+
"SELECT routine_name FROM information_schema.routines WHERE routine_type LIKE '%FUNCTION' AND routine_name IN (");
2797+
for (int i = 0; i < quoteEscapedTvfNameList.size(); i++) {
2798+
if (dialect == Dialect.POSTGRESQL) {
2799+
sql.append("$").append(i + 1);
2800+
} else {
2801+
sql.append("@p").append(i);
2802+
}
2803+
if (i < quoteEscapedTvfNameList.size() - 1) {
2804+
sql.append(", ");
2805+
}
2806+
}
2807+
sql.append(")");
2808+
Statement.Builder builder = Statement.newBuilder(sql.toString());
2809+
for (int i = 0; i < quoteEscapedTvfNameList.size(); i++) {
2810+
if (dialect == Dialect.POSTGRESQL) {
2811+
builder.bind("p" + (i + 1)).to(quoteEscapedTvfNameList.get(i));
2812+
} else {
2813+
builder.bind("p" + i).to(quoteEscapedTvfNameList.get(i));
2814+
}
2815+
}
2816+
Statement statement = builder.build();
2817+
ResultSet resultSet = tx.executeQuery(statement);
2818+
java.util.Set<String> foundNames = new java.util.HashSet<>();
2819+
while (resultSet.next()) {
2820+
foundNames.add(resultSet.getString(0));
2821+
}
2822+
for (String tvfName : quoteEscapedTvfNameList) {
2823+
if (!foundNames.contains(tvfName)) {
2824+
throw new IllegalArgumentException(
2825+
"TVF specified: " + tvfName + " is not found in the existing TVF's: " + foundNames);
2826+
}
2827+
}
2828+
}
2829+
}
2830+
27572831
@VisibleForTesting
27582832
static boolean isMutableChangeStream(
27592833
DatabaseClient databaseClient, Dialect dialect, String changeStreamName) {

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

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@
2020
import com.google.cloud.Timestamp;
2121
import com.google.cloud.spanner.Options.RpcPriority;
2222
import java.util.Collections;
23+
import java.util.List;
2324
import org.apache.beam.sdk.io.gcp.spanner.changestreams.dofn.DetectNewPartitionsDoFn;
2425
import org.apache.beam.sdk.io.gcp.spanner.changestreams.model.PartitionMetadata;
2526
import org.apache.beam.sdk.io.gcp.spanner.changestreams.model.PartitionMetadata.State;
@@ -64,6 +65,18 @@ public class ChangeStreamsConstants {
6465
/** The sliding window size in seconds for throughput reporting. */
6566
public static final int THROUGHPUT_WINDOW_SECONDS = 10;
6667

68+
/**
69+
* The delimiter used to separate the partition token and the tvf name. Note this string does not
70+
* exist in the partition token itself.
71+
*/
72+
public static final String PARTITION_TOKEN_TVF_NAME_DELIMITER = "#";
73+
74+
/** The default tvf name for a change stream query is the empty {@link String}. */
75+
public static final String DEFAULT_TVF_NAME = "";
76+
77+
/** The default tvf name list to query and union is empty {@link Collections.emptyList()}. */
78+
public static final List<String> DEFAULT_TVF_NAME_LIST = Collections.emptyList();
79+
6780
/**
6881
* We use the following partition token to provide an estimate size of a partition token. A usual
6982
* partition token has around 140 characters.
@@ -85,6 +98,7 @@ public class ChangeStreamsConstants {
8598
.setState(State.CREATED)
8699
.setWatermark(Timestamp.now())
87100
.setCreatedAt(Timestamp.now())
101+
.setTvfName(DEFAULT_TVF_NAME)
88102
.build();
89103

90104
/**

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

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -155,13 +155,17 @@ private void processChildPartition(
155155
record.getStartTimestamp(),
156156
partition.getEndTimestamp(),
157157
partition.getHeartbeatMillis(),
158+
partition.getTvfName(),
158159
childPartition);
159160
LOG.debug("[{}] Inserting child partition token {}", partitionToken, childPartitionToken);
160161
final Boolean insertedRow =
161162
partitionMetadataDao
162163
.runInTransaction(
163164
transaction -> {
164-
if (transaction.getPartition(childPartitionToken) == null) {
165+
if (transaction.getPartition(
166+
PartitionMetadataDao.composePartitionTokenWithTvfName(
167+
childPartitionToken, partition.getTvfName()))
168+
== null) {
165169
transaction.insert(row);
166170
return true;
167171
} else {
@@ -188,13 +192,15 @@ private PartitionMetadata toPartitionMetadata(
188192
Timestamp startTimestamp,
189193
Timestamp endTimestamp,
190194
long heartbeatMillis,
195+
String tvfName,
191196
ChildPartition childPartition) {
192197
return PartitionMetadata.newBuilder()
193198
.setPartitionToken(childPartition.getToken())
194199
.setParentTokens(childPartition.getParentTokens())
195200
.setStartTimestamp(startTimestamp)
196201
.setEndTimestamp(endTimestamp)
197202
.setHeartbeatMillis(heartbeatMillis)
203+
.setTvfName(tvfName)
198204
.setState(CREATED)
199205
.setWatermark(startTimestamp)
200206
.build();

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

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -172,11 +172,14 @@ private ProcessContinuation schedulePartitions(
172172
}
173173

174174
private Timestamp updateBatchToScheduled(List<PartitionMetadata> batchPartitions) {
175-
final List<String> batchPartitionTokens =
175+
final List<String> batchComposedPartitionTokens =
176176
batchPartitions.stream()
177-
.map(PartitionMetadata::getPartitionToken)
177+
.map(
178+
partition ->
179+
PartitionMetadataDao.composePartitionTokenWithTvfName(
180+
partition.getPartitionToken(), partition.getTvfName()))
178181
.collect(Collectors.toList());
179-
return dao.updateToScheduled(batchPartitionTokens);
182+
return dao.updateToScheduled(batchComposedPartitionTokens);
180183
}
181184

182185
private void outputBatch(

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

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -138,6 +138,7 @@ private void processStartPartition(
138138
.setStartTimestamp(record.getStartTimestamp())
139139
.setEndTimestamp(partition.getEndTimestamp())
140140
.setHeartbeatMillis(partition.getHeartbeatMillis())
141+
.setTvfName(partition.getTvfName())
141142
.setState(CREATED)
142143
.setWatermark(record.getStartTimestamp())
143144
.build();
@@ -146,7 +147,10 @@ private void processStartPartition(
146147
partitionMetadataDao
147148
.runInTransaction(
148149
transaction -> {
149-
if (transaction.getPartition(startPartitionToken) == null) {
150+
if (transaction.getPartition(
151+
PartitionMetadataDao.composePartitionTokenWithTvfName(
152+
startPartitionToken, partition.getTvfName()))
153+
== null) {
150154
transaction.insert(row);
151155
return true;
152156
}

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

Lines changed: 23 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -185,11 +185,14 @@ public ProcessContinuation run(
185185
ManualWatermarkEstimator<Instant> watermarkEstimator,
186186
BundleFinalizer bundleFinalizer) {
187187
final String token = partition.getPartitionToken();
188+
final String tvfName = partition.getTvfName();
188189

189190
// TODO: Potentially we can avoid this fetch, by enriching the runningAt timestamp when the
190191
// ReadChangeStreamPartitionDoFn#processElement is called
191192
final PartitionMetadata updatedPartition =
192-
Optional.ofNullable(partitionMetadataDao.getPartition(token))
193+
Optional.ofNullable(
194+
partitionMetadataDao.getPartition(
195+
PartitionMetadataDao.composePartitionTokenWithTvfName(token, tvfName)))
193196
.map(partitionMetadataMapper::from)
194197
.orElseThrow(
195198
() ->
@@ -223,7 +226,11 @@ public ProcessContinuation run(
223226

224227
try (ChangeStreamResultSet resultSet =
225228
changeStreamDao.changeStreamQuery(
226-
token, startTimestamp, changeStreamQueryEndTimestamp, partition.getHeartbeatMillis())) {
229+
token,
230+
tvfName,
231+
startTimestamp,
232+
changeStreamQueryEndTimestamp,
233+
partition.getHeartbeatMillis())) {
227234

228235
metrics.incQueryCounter();
229236
while (resultSet.next()) {
@@ -298,7 +305,9 @@ public ProcessContinuation run(
298305
LOG.debug("[{}] Continuation present, returning {}", token, maybeContinuation);
299306
bundleFinalizer.afterBundleCommit(
300307
Instant.now().plus(BUNDLE_FINALIZER_TIMEOUT),
301-
updateWatermarkCallback(token, watermarkEstimator));
308+
updateWatermarkCallback(
309+
PartitionMetadataDao.composePartitionTokenWithTvfName(token, tvfName),
310+
watermarkEstimator));
302311
return maybeContinuation.get();
303312
}
304313
}
@@ -341,7 +350,9 @@ public ProcessContinuation run(
341350
}
342351
bundleFinalizer.afterBundleCommit(
343352
Instant.now().plus(BUNDLE_FINALIZER_TIMEOUT),
344-
updateWatermarkCallback(token, watermarkEstimator));
353+
updateWatermarkCallback(
354+
PartitionMetadataDao.composePartitionTokenWithTvfName(token, tvfName),
355+
watermarkEstimator));
345356
LOG.debug("[{}] Rescheduling partition to resume reading", token);
346357
return ProcessContinuation.resume();
347358
}
@@ -361,25 +372,27 @@ public ProcessContinuation run(
361372
LOG.debug("[{}] Finishing partition", token);
362373
// TODO: This should be performed after the commit succeeds. Since bundle finalizers are not
363374
// guaranteed to be called, this needs to be performed in a subsequent fused stage.
364-
partitionMetadataDao.updateToFinished(token);
375+
partitionMetadataDao.updateToFinished(
376+
PartitionMetadataDao.composePartitionTokenWithTvfName(token, tvfName));
365377
metrics.decActivePartitionReadCounter();
366378
LOG.info("[{}] After attempting to finish the partition", token);
367379
return ProcessContinuation.stop();
368380
}
369381

370382
private BundleFinalizer.Callback updateWatermarkCallback(
371-
String token, WatermarkEstimator<Instant> watermarkEstimator) {
383+
String composedToken, WatermarkEstimator<Instant> watermarkEstimator) {
372384
return () -> {
373385
final Instant watermark = watermarkEstimator.currentWatermark();
374-
LOG.debug("[{}] Updating current watermark to {}", token, watermark);
386+
LOG.debug("[{}] Updating current watermark to {}", composedToken, watermark);
375387
try {
376388
partitionMetadataDao.updateWatermark(
377-
token, Timestamp.ofTimeMicroseconds(watermark.getMillis() * 1_000L));
389+
composedToken, Timestamp.ofTimeMicroseconds(watermark.getMillis() * 1_000L));
378390
} catch (SpannerException e) {
379391
if (e.getErrorCode() == ErrorCode.NOT_FOUND) {
380-
LOG.debug("[{}] Unable to update the current watermark, partition NOT FOUND", token);
392+
LOG.debug(
393+
"[{}] Unable to update the current watermark, partition NOT FOUND", composedToken);
381394
} else {
382-
LOG.error("[{}] Error updating the current watermark", token, e);
395+
LOG.error("[{}] Error updating the current watermark", composedToken, e);
383396
}
384397
}
385398
};

0 commit comments

Comments
 (0)