You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@beam.apache.org by pa...@apache.org on 2023/03/31 03:38:45 UTC
[beam] branch master updated: Add partition metadata in error logs in the restriction tracker for change streams and delete unused files (#26040)
This is an automated email from the ASF dual-hosted git repository.
pabloem pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git
The following commit(s) were added to refs/heads/master by this push:
new dc7a7a02c7a Add partition metadata in error logs in the restriction tracker for change streams and delete unused files (#26040)
dc7a7a02c7a is described below
commit dc7a7a02c7ac64887b7ada426c8385a221859af9
Author: ChangyuLi28 <35...@users.noreply.github.com>
AuthorDate: Thu Mar 30 20:38:33 2023 -0700
Add partition metadata in error logs in the restriction tracker for change streams and delete unused files (#26040)
---
.../changestreams/restriction/PartitionMode.java | 35 -
.../restriction/PartitionPosition.java | 90 ---
.../restriction/PartitionRestriction.java | 141 ----
.../restriction/PartitionRestrictionClaimer.java | 136 ----
.../restriction/PartitionRestrictionMetadata.java | 106 ---
.../PartitionRestrictionProgressChecker.java | 96 ---
.../restriction/PartitionRestrictionSplitter.java | 128 ---
.../restriction/PartitionRestrictionTracker.java | 159 ----
.../ReadChangeStreamPartitionRangeTracker.java | 2 +-
.../restriction/TimestampRangeTracker.java | 24 +
.../restriction/PartitionPositionTest.java | 70 --
.../PartitionRestrictionClaimerTest.java | 272 -------
.../PartitionRestrictionProgressCheckerTest.java | 230 ------
.../PartitionRestrictionSplitterTest.java | 168 ----
.../restriction/PartitionRestrictionTest.java | 78 --
.../PartitionRestrictionTrackerTest.java | 859 ---------------------
.../util/PartitionPositionGenerator.java | 54 --
17 files changed, 25 insertions(+), 2623 deletions(-)
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/restriction/PartitionMode.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/restriction/PartitionMode.java
deleted file mode 100644
index 341dc7aa38f..00000000000
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/restriction/PartitionMode.java
+++ /dev/null
@@ -1,35 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements. See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership. The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * "License"); you may not use this file except in compliance
- * with the License. You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-package org.apache.beam.sdk.io.gcp.spanner.changestreams.restriction;
-
-/** This enum contains the states that PartitionRestrictionTracker will go through. */
-public enum PartitionMode {
- // In this state, the restriction tracker will update the state of the input partition token
- // from SCHEDULED to RUNNING.
- UPDATE_STATE,
- // In this state, the restriction tracker will execute a change stream query.
- QUERY_CHANGE_STREAM,
- // In this state, the restriction tracker will wait for the child partition SDFs to start
- // running before terminating the SDF.
- WAIT_FOR_CHILD_PARTITIONS,
- // In this state, the restriction tracker will terminate the SDF.
- DONE,
-
- // Occurs when Dataflow checkpoints the current restriction.
- STOP
-}
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/restriction/PartitionPosition.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/restriction/PartitionPosition.java
deleted file mode 100644
index 1011fadc74a..00000000000
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/restriction/PartitionPosition.java
+++ /dev/null
@@ -1,90 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements. See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership. The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * "License"); you may not use this file except in compliance
- * with the License. You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-package org.apache.beam.sdk.io.gcp.spanner.changestreams.restriction;
-
-import com.google.cloud.Timestamp;
-import java.io.Serializable;
-import java.util.Objects;
-import java.util.Optional;
-import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting;
-import org.checkerframework.checker.nullness.qual.Nullable;
-
-/** Represents the current position of the running SDF within PartitionRestriction. */
-public class PartitionPosition implements Serializable {
-
- private static final long serialVersionUID = -9088898012221404492L;
-
- private final Optional<Timestamp> maybeTimestamp;
- private final PartitionMode mode;
-
- public static PartitionPosition updateState() {
- return new PartitionPosition(Optional.empty(), PartitionMode.UPDATE_STATE);
- }
-
- public static PartitionPosition queryChangeStream(Timestamp timestamp) {
- return new PartitionPosition(Optional.of(timestamp), PartitionMode.QUERY_CHANGE_STREAM);
- }
-
- public static PartitionPosition waitForChildPartitions() {
- return new PartitionPosition(Optional.empty(), PartitionMode.WAIT_FOR_CHILD_PARTITIONS);
- }
-
- public static PartitionPosition done() {
- return new PartitionPosition(Optional.empty(), PartitionMode.DONE);
- }
-
- public static PartitionPosition stop() {
- return new PartitionPosition(Optional.empty(), PartitionMode.STOP);
- }
-
- @VisibleForTesting
- public PartitionPosition(Optional<Timestamp> maybeTimestamp, PartitionMode mode) {
- this.maybeTimestamp = maybeTimestamp;
- this.mode = mode;
- }
-
- public Optional<Timestamp> getTimestamp() {
- return maybeTimestamp;
- }
-
- public PartitionMode getMode() {
- return mode;
- }
-
- @Override
- public boolean equals(@Nullable Object o) {
- if (this == o) {
- return true;
- }
- if (!(o instanceof PartitionPosition)) {
- return false;
- }
- PartitionPosition that = (PartitionPosition) o;
- return Objects.equals(maybeTimestamp, that.maybeTimestamp) && mode == that.mode;
- }
-
- @Override
- public int hashCode() {
- return Objects.hash(maybeTimestamp, mode);
- }
-
- @Override
- public String toString() {
- return "PartitionPosition{" + "maybeTimestamp=" + maybeTimestamp + ", mode=" + mode + '}';
- }
-}
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/restriction/PartitionRestriction.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/restriction/PartitionRestriction.java
deleted file mode 100644
index c02b5574f24..00000000000
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/restriction/PartitionRestriction.java
+++ /dev/null
@@ -1,141 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements. See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership. The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * "License"); you may not use this file except in compliance
- * with the License. You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-package org.apache.beam.sdk.io.gcp.spanner.changestreams.restriction;
-
-import static org.apache.beam.sdk.io.gcp.spanner.changestreams.restriction.PartitionMode.DONE;
-import static org.apache.beam.sdk.io.gcp.spanner.changestreams.restriction.PartitionMode.QUERY_CHANGE_STREAM;
-import static org.apache.beam.sdk.io.gcp.spanner.changestreams.restriction.PartitionMode.STOP;
-import static org.apache.beam.sdk.io.gcp.spanner.changestreams.restriction.PartitionMode.UPDATE_STATE;
-import static org.apache.beam.sdk.io.gcp.spanner.changestreams.restriction.PartitionMode.WAIT_FOR_CHILD_PARTITIONS;
-
-import com.google.cloud.Timestamp;
-import java.io.Serializable;
-import java.util.Objects;
-import org.checkerframework.checker.nullness.qual.Nullable;
-
-/** Represents the restriction for PartitionRestrictionTracker. */
-@SuppressWarnings("initialization.fields.uninitialized")
-public class PartitionRestriction implements Serializable {
-
- private static final long serialVersionUID = -7009236776208644264L;
-
- private final Timestamp startTimestamp;
- private final Timestamp endTimestamp;
- private final PartitionMode mode;
- private final @Nullable PartitionMode stoppedMode;
- private PartitionRestrictionMetadata metadata;
-
- public static PartitionRestriction updateState(Timestamp startTimestamp, Timestamp endTimestamp) {
- return new PartitionRestriction(startTimestamp, endTimestamp, UPDATE_STATE, null);
- }
-
- public static PartitionRestriction queryChangeStream(
- Timestamp startTimestamp, Timestamp endTimestamp) {
- return new PartitionRestriction(startTimestamp, endTimestamp, QUERY_CHANGE_STREAM, null);
- }
-
- public static PartitionRestriction waitForChildPartitions(
- Timestamp startTimestamp, Timestamp endTimestamp) {
- return new PartitionRestriction(startTimestamp, endTimestamp, WAIT_FOR_CHILD_PARTITIONS, null);
- }
-
- public static PartitionRestriction done(Timestamp startTimestamp, Timestamp endTimestamp) {
- return new PartitionRestriction(startTimestamp, endTimestamp, DONE, null);
- }
-
- public static PartitionRestriction stop(PartitionRestriction restriction) {
- return new PartitionRestriction(
- restriction.getStartTimestamp(),
- restriction.getEndTimestamp(),
- STOP,
- restriction.getMode())
- .withMetadata(restriction.getMetadata());
- }
-
- public PartitionRestriction(
- Timestamp startTimestamp,
- Timestamp endTimestamp,
- PartitionMode mode,
- @Nullable PartitionMode stoppedMode) {
- this.startTimestamp = startTimestamp;
- this.endTimestamp = endTimestamp;
- this.mode = mode;
- this.stoppedMode = stoppedMode;
- }
-
- public PartitionRestriction withMetadata(PartitionRestrictionMetadata metadata) {
- this.metadata = metadata;
- return this;
- }
-
- public Timestamp getStartTimestamp() {
- return startTimestamp;
- }
-
- public Timestamp getEndTimestamp() {
- return endTimestamp;
- }
-
- public PartitionMode getMode() {
- return mode;
- }
-
- public @Nullable PartitionMode getStoppedMode() {
- return stoppedMode;
- }
-
- public PartitionRestrictionMetadata getMetadata() {
- return metadata;
- }
-
- @Override
- public boolean equals(@Nullable Object o) {
- if (this == o) {
- return true;
- }
- if (!(o instanceof PartitionRestriction)) {
- return false;
- }
- PartitionRestriction that = (PartitionRestriction) o;
- return Objects.equals(startTimestamp, that.startTimestamp)
- && Objects.equals(endTimestamp, that.endTimestamp)
- && mode == that.mode
- && stoppedMode == that.stoppedMode;
- }
-
- @Override
- public int hashCode() {
- return Objects.hash(startTimestamp, endTimestamp, mode, stoppedMode);
- }
-
- @Override
- public String toString() {
- return "PartitionRestriction{"
- + "startTimestamp="
- + startTimestamp
- + ", endTimestamp="
- + endTimestamp
- + ", mode="
- + mode
- + ", stoppedMode="
- + stoppedMode
- + ", metadata="
- + metadata
- + '}';
- }
-}
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/restriction/PartitionRestrictionClaimer.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/restriction/PartitionRestrictionClaimer.java
deleted file mode 100644
index 5ab35bf2f8a..00000000000
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/restriction/PartitionRestrictionClaimer.java
+++ /dev/null
@@ -1,136 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements. See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership. The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * "License"); you may not use this file except in compliance
- * with the License. You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-package org.apache.beam.sdk.io.gcp.spanner.changestreams.restriction;
-
-import static org.apache.beam.sdk.io.gcp.spanner.changestreams.restriction.PartitionMode.DONE;
-import static org.apache.beam.sdk.io.gcp.spanner.changestreams.restriction.PartitionMode.QUERY_CHANGE_STREAM;
-import static org.apache.beam.sdk.io.gcp.spanner.changestreams.restriction.PartitionMode.STOP;
-import static org.apache.beam.sdk.io.gcp.spanner.changestreams.restriction.PartitionMode.UPDATE_STATE;
-import static org.apache.beam.sdk.io.gcp.spanner.changestreams.restriction.PartitionMode.WAIT_FOR_CHILD_PARTITIONS;
-import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
-
-import com.google.cloud.Timestamp;
-import java.util.Collections;
-import java.util.HashMap;
-import java.util.Map;
-import java.util.Optional;
-import java.util.Set;
-import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Sets;
-import org.checkerframework.checker.nullness.qual.Nullable;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
-
-/** The PartitionRestrictionClaimer class. */
-public class PartitionRestrictionClaimer {
-
- private static final Logger LOG = LoggerFactory.getLogger(PartitionRestrictionClaimer.class);
-
- private final Map<PartitionMode, Set<PartitionMode>> allowedTransitions;
-
- public PartitionRestrictionClaimer() {
- this.allowedTransitions = new HashMap<>();
- allowedTransitions.put(UPDATE_STATE, Sets.newHashSet(UPDATE_STATE, QUERY_CHANGE_STREAM));
- allowedTransitions.put(
- QUERY_CHANGE_STREAM, Sets.newHashSet(QUERY_CHANGE_STREAM, WAIT_FOR_CHILD_PARTITIONS));
- allowedTransitions.put(
- WAIT_FOR_CHILD_PARTITIONS, Sets.newHashSet(WAIT_FOR_CHILD_PARTITIONS, DONE));
- }
-
- public boolean tryClaim(
- PartitionRestriction restriction,
- @Nullable PartitionPosition lastClaimedPosition,
- PartitionPosition position) {
- final PartitionMode fromMode =
- Optional.ofNullable(lastClaimedPosition)
- .map(PartitionPosition::getMode)
- .orElse(restriction.getMode());
- final PartitionMode toMode = position.getMode();
- final String token =
- Optional.ofNullable(restriction.getMetadata())
- .map(PartitionRestrictionMetadata::getPartitionToken)
- .orElse("");
-
- if (fromMode == STOP) {
- LOG.debug(
- "["
- + token
- + "] Try claim from ("
- + restriction
- + ","
- + lastClaimedPosition
- + ", "
- + position
- + ") is false");
- return false;
- }
-
- checkArgument(
- allowedTransitions.getOrDefault(fromMode, Collections.emptySet()).contains(toMode),
- "Invalid partition mode transition from %s to %s",
- fromMode,
- toMode);
- checkArgument(
- toMode != QUERY_CHANGE_STREAM || position.getTimestamp().isPresent(),
- "%s mode must specify a timestamp (no value sent)",
- toMode);
-
- boolean tryClaimResult;
- switch (toMode) {
- case QUERY_CHANGE_STREAM:
- final Timestamp attemptedTimestamp = position.getTimestamp().get();
- final Timestamp endTimestamp =
- Optional.ofNullable(restriction.getEndTimestamp()).orElse(Timestamp.MAX_VALUE);
-
- checkArgument(
- lastClaimedPosition == null
- || !lastClaimedPosition.getTimestamp().isPresent()
- || attemptedTimestamp.compareTo(lastClaimedPosition.getTimestamp().get()) >= 0,
- "Trying to claim offset %s while last attempted was %s",
- position,
- lastClaimedPosition);
- checkArgument(
- attemptedTimestamp.compareTo(restriction.getStartTimestamp()) >= 0,
- "Trying to claim offset %s before the start timestamp %s",
- position,
- restriction.getStartTimestamp().toString());
-
- tryClaimResult = attemptedTimestamp.compareTo(endTimestamp) < 0;
- break;
- case UPDATE_STATE:
- case WAIT_FOR_CHILD_PARTITIONS:
- case DONE:
- tryClaimResult = true;
- break;
- case STOP:
- throw new IllegalArgumentException("Trying to claim STOP state is invalid");
- default:
- throw new IllegalArgumentException("Unknown mode " + toMode);
- }
-
- LOG.debug(
- "["
- + token
- + "] Try claim from ("
- + restriction
- + ", "
- + position
- + ") is "
- + tryClaimResult);
- return tryClaimResult;
- }
-}
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/restriction/PartitionRestrictionMetadata.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/restriction/PartitionRestrictionMetadata.java
deleted file mode 100644
index 6f42065149a..00000000000
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/restriction/PartitionRestrictionMetadata.java
+++ /dev/null
@@ -1,106 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements. See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership. The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * "License"); you may not use this file except in compliance
- * with the License. You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-package org.apache.beam.sdk.io.gcp.spanner.changestreams.restriction;
-
-import com.google.cloud.Timestamp;
-import java.io.Serializable;
-
-/** The PartitionRestrictionMetadata class. */
-@SuppressWarnings("initialization.fields.uninitialized")
-public class PartitionRestrictionMetadata implements Serializable {
-
- private static final long serialVersionUID = 4744539370804123350L;
-
- private final String partitionToken;
- private final Timestamp partitionStartTimestamp;
- private final Timestamp partitionEndTimestamp;
-
- public PartitionRestrictionMetadata(
- String partitionToken, Timestamp partitionStartTimestamp, Timestamp partitionEndTimestamp) {
- this.partitionToken = partitionToken;
- this.partitionStartTimestamp = partitionStartTimestamp;
- this.partitionEndTimestamp = partitionEndTimestamp;
- }
-
- public String getPartitionToken() {
- return partitionToken;
- }
-
- public Timestamp getPartitionStartTimestamp() {
- return partitionStartTimestamp;
- }
-
- public Timestamp getPartitionEndTimestamp() {
- return partitionEndTimestamp;
- }
-
- @Override
- public String toString() {
- return "PartitionRestrictionMetadata{"
- + "partitionToken='"
- + partitionToken
- + '\''
- + ", partitionStartTimestamp="
- + partitionStartTimestamp
- + ", partitionEndTimestamp="
- + partitionEndTimestamp
- + '}';
- }
-
- public static Builder newBuilder() {
- return new Builder();
- }
-
- public static Builder newBuilder(PartitionRestrictionMetadata metadata) {
- return new Builder(metadata);
- }
-
- public static class Builder {
- private String partitionToken;
- private Timestamp partitionStartTimestamp;
- private Timestamp partitionEndTimestamp;
-
- public Builder() {}
-
- public Builder(PartitionRestrictionMetadata metadata) {
- this.partitionToken = metadata.partitionToken;
- this.partitionStartTimestamp = metadata.partitionStartTimestamp;
- this.partitionEndTimestamp = metadata.partitionEndTimestamp;
- }
-
- public Builder withPartitionToken(String partitionToken) {
- this.partitionToken = partitionToken;
- return this;
- }
-
- public Builder withPartitionStartTimestamp(Timestamp partitionStartTimestamp) {
- this.partitionStartTimestamp = partitionStartTimestamp;
- return this;
- }
-
- public Builder withPartitionEndTimestamp(Timestamp partitionEndTimestamp) {
- this.partitionEndTimestamp = partitionEndTimestamp;
- return this;
- }
-
- public PartitionRestrictionMetadata build() {
- return new PartitionRestrictionMetadata(
- partitionToken, partitionStartTimestamp, partitionEndTimestamp);
- }
- }
-}
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/restriction/PartitionRestrictionProgressChecker.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/restriction/PartitionRestrictionProgressChecker.java
deleted file mode 100644
index cd762aa58ec..00000000000
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/restriction/PartitionRestrictionProgressChecker.java
+++ /dev/null
@@ -1,96 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements. See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership. The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * "License"); you may not use this file except in compliance
- * with the License. You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-package org.apache.beam.sdk.io.gcp.spanner.changestreams.restriction;
-
-import static org.apache.beam.sdk.io.gcp.spanner.changestreams.restriction.PartitionMode.QUERY_CHANGE_STREAM;
-import static org.apache.beam.sdk.io.gcp.spanner.changestreams.restriction.PartitionMode.STOP;
-import static org.apache.beam.sdk.io.gcp.spanner.changestreams.restriction.PartitionMode.UPDATE_STATE;
-
-import com.google.cloud.Timestamp;
-import java.math.BigDecimal;
-import java.util.Optional;
-import java.util.function.Supplier;
-import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker.Progress;
-import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting;
-import org.checkerframework.checker.nullness.qual.Nullable;
-
-public class PartitionRestrictionProgressChecker {
-
- protected Supplier<Timestamp> timeSupplier;
-
- /**
- * Indicates how many mode transitions have been completed for the current mode. The transitions
- * are as follows: * (1) UPDATE_STATE, (2) QUERY_CHANGE_STREAM, (3) WAIT_FOR_CHILD_PARTITIONS, (4)
- * DONE.
- *
- * <p>This is used to calculate the units of work left, meaning that 1 transition = 1 unit of
- * work.
- */
- public PartitionRestrictionProgressChecker() {
- this.timeSupplier = () -> Timestamp.now();
- }
-
- @VisibleForTesting
- public void setTimeSupplier(Supplier<Timestamp> timeSupplier) {
- this.timeSupplier = timeSupplier;
- }
-
- public Progress getProgress(
- PartitionRestriction restriction, @Nullable PartitionPosition lastClaimedPosition) {
- BigDecimal totalWork;
- if (restriction.getEndTimestamp().compareTo(Timestamp.MAX_VALUE) == 0) {
- // When the given end timestamp equals to Timestamp.MAX_VALUE, this means that
- // the end timestamp is not specified which should be a streaming job. So we
- // use now() as the end timestamp.
- totalWork = BigDecimal.valueOf(timeSupplier.get().getSeconds());
- } else {
- totalWork = BigDecimal.valueOf(restriction.getEndTimestamp().getSeconds());
- }
-
- final PartitionMode currentMode =
- Optional.ofNullable(lastClaimedPosition)
- .map(PartitionPosition::getMode)
- .orElse(restriction.getMode());
-
- if (currentMode == STOP) {
- // Return progress indicating that there is no more work to be done for this SDF.
- final double workCompleted = totalWork.doubleValue();
- final double workRemaining = 1D;
-
- return Progress.from(workCompleted, workRemaining);
- }
-
- // We get the current number of seconds.
- final long currentSecondsNum =
- Optional.ofNullable(lastClaimedPosition)
- .flatMap(PartitionPosition::getTimestamp)
- .map(Timestamp::getSeconds)
- .orElse(
- (currentMode == UPDATE_STATE || currentMode == QUERY_CHANGE_STREAM)
- ? restriction.getStartTimestamp().getSeconds()
- : restriction.getEndTimestamp().getSeconds());
- final BigDecimal workCompleted = BigDecimal.valueOf(currentSecondsNum);
-
- // The work remaining is the total work minus the work completed.
- // The remaining work must be greater than 0. Otherwise, it will cause an issue
- // that the watermark does not advance.
- final BigDecimal workRemaining = totalWork.subtract(workCompleted).max(BigDecimal.ONE);
-
- return Progress.from(workCompleted.doubleValue(), workRemaining.doubleValue());
- }
-}
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/restriction/PartitionRestrictionSplitter.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/restriction/PartitionRestrictionSplitter.java
deleted file mode 100644
index 23866858750..00000000000
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/restriction/PartitionRestrictionSplitter.java
+++ /dev/null
@@ -1,128 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements. See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership. The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * "License"); you may not use this file except in compliance
- * with the License. You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-package org.apache.beam.sdk.io.gcp.spanner.changestreams.restriction;
-
-import static java.math.MathContext.DECIMAL128;
-import static org.apache.beam.sdk.io.gcp.spanner.changestreams.restriction.TimestampUtils.toNanos;
-import static org.apache.beam.sdk.io.gcp.spanner.changestreams.restriction.TimestampUtils.toTimestamp;
-
-import com.google.cloud.Timestamp;
-import java.math.BigDecimal;
-import java.util.Optional;
-import org.apache.beam.sdk.io.gcp.spanner.changestreams.model.InitialPartition;
-import org.apache.beam.sdk.transforms.splittabledofn.SplitResult;
-import org.checkerframework.checker.nullness.qual.Nullable;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
-
-/** The PartitionRestrictionSplitter class. */
-public class PartitionRestrictionSplitter {
-
- private static final Logger LOG = LoggerFactory.getLogger(PartitionRestrictionSplitter.class);
-
- public @Nullable SplitResult<PartitionRestriction> trySplit(
- double fractionOfRemainder,
- @Nullable PartitionPosition lastClaimedPosition,
- PartitionRestriction restriction) {
- if (lastClaimedPosition == null) {
- return null;
- }
-
- final String token =
- Optional.ofNullable(restriction.getMetadata())
- .map(PartitionRestrictionMetadata::getPartitionToken)
- .orElse("");
- final PartitionMode positionMode = lastClaimedPosition.getMode();
- final Timestamp startTimestamp = restriction.getStartTimestamp();
- final Timestamp endTimestamp = restriction.getEndTimestamp();
-
- SplitResult<PartitionRestriction> splitResult;
- switch (positionMode) {
- case UPDATE_STATE:
- splitResult =
- SplitResult.of(
- PartitionRestriction.stop(restriction),
- PartitionRestriction.queryChangeStream(startTimestamp, endTimestamp)
- .withMetadata(restriction.getMetadata()));
- break;
- case QUERY_CHANGE_STREAM:
- if (token.equals(InitialPartition.PARTITION_TOKEN)) {
- splitResult = null;
- } else {
- final BigDecimal toInNanos = toNanos(endTimestamp);
- final BigDecimal currentInNanos = toNanos(lastClaimedPosition.getTimestamp().get());
- final BigDecimal nanosOffset =
- toInNanos
- .subtract(currentInNanos, DECIMAL128)
- .multiply(BigDecimal.valueOf(fractionOfRemainder), DECIMAL128)
- .max(BigDecimal.ONE);
-
- // splitPosition = current + max(1, (range.getTo() - current) * fractionOfRemainder)
- final BigDecimal splitPositionInNanos = currentInNanos.add(nanosOffset, DECIMAL128);
- final Timestamp splitPosition = toTimestamp(splitPositionInNanos);
-
- // if the split position is greater than or equal to the restriction's end timestamp,
- // which is the partition's end timestamp + 1, it means that we are done executing the
- // change stream query. Thus, we should return STOP as primary, and WAIT_FOR_CHILD_
- // PARTITIONS as residual.
- if (splitPosition.compareTo(endTimestamp) >= 0) {
- splitResult =
- SplitResult.of(
- PartitionRestriction.stop(restriction),
- PartitionRestriction.waitForChildPartitions(startTimestamp, endTimestamp)
- .withMetadata(restriction.getMetadata()));
- } else {
- splitResult =
- SplitResult.of(
- PartitionRestriction.queryChangeStream(startTimestamp, splitPosition)
- .withMetadata(restriction.getMetadata()),
- PartitionRestriction.queryChangeStream(splitPosition, endTimestamp)
- .withMetadata(restriction.getMetadata()));
- }
- }
- break;
- case WAIT_FOR_CHILD_PARTITIONS:
- splitResult =
- SplitResult.of(
- PartitionRestriction.stop(restriction),
- PartitionRestriction.waitForChildPartitions(startTimestamp, endTimestamp)
- .withMetadata(restriction.getMetadata()));
- break;
- case DONE:
- case STOP:
- splitResult = null;
- break;
- default:
- // TODO: See if we need to throw or do something else
- throw new IllegalArgumentException("Unknown mode " + positionMode);
- }
-
- LOG.debug(
- "["
- + token
- + "] Split result for ("
- + fractionOfRemainder
- + ", "
- + lastClaimedPosition
- + ", "
- + restriction
- + ") is "
- + splitResult);
- return splitResult;
- }
-}
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/restriction/PartitionRestrictionTracker.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/restriction/PartitionRestrictionTracker.java
deleted file mode 100644
index f986999a3f5..00000000000
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/restriction/PartitionRestrictionTracker.java
+++ /dev/null
@@ -1,159 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements. See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership. The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * "License"); you may not use this file except in compliance
- * with the License. You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-package org.apache.beam.sdk.io.gcp.spanner.changestreams.restriction;
-
-import static org.apache.beam.sdk.io.gcp.spanner.changestreams.restriction.PartitionMode.DONE;
-import static org.apache.beam.sdk.io.gcp.spanner.changestreams.restriction.PartitionMode.QUERY_CHANGE_STREAM;
-import static org.apache.beam.sdk.io.gcp.spanner.changestreams.restriction.PartitionMode.STOP;
-import static org.apache.beam.sdk.io.gcp.spanner.changestreams.restriction.TimestampUtils.next;
-import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkState;
-
-import com.google.cloud.Timestamp;
-import java.util.Optional;
-import java.util.function.Supplier;
-import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
-import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker.HasProgress;
-import org.apache.beam.sdk.transforms.splittabledofn.SplitResult;
-import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting;
-import org.checkerframework.checker.nullness.qual.Nullable;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
-
-// TODO: Add java docs
-// TODO: Implement duration waiting for returning false on try claim
-public class PartitionRestrictionTracker
- extends RestrictionTracker<PartitionRestriction, PartitionPosition> implements HasProgress {
-
- private static final Logger LOG = LoggerFactory.getLogger(PartitionRestrictionTracker.class);
- private final PartitionRestrictionSplitter splitter;
- private final PartitionRestrictionClaimer claimer;
- private final PartitionRestrictionProgressChecker progressChecker;
- protected PartitionRestriction restriction;
- private @Nullable PartitionPosition lastClaimedPosition;
- protected Supplier<Timestamp> timeSupplier;
-
- public PartitionRestrictionTracker(PartitionRestriction restriction) {
- this(
- restriction,
- new PartitionRestrictionSplitter(),
- new PartitionRestrictionClaimer(),
- new PartitionRestrictionProgressChecker());
- }
-
- PartitionRestrictionTracker(
- PartitionRestriction restriction,
- PartitionRestrictionSplitter splitter,
- PartitionRestrictionClaimer claimer,
- PartitionRestrictionProgressChecker progressChecker) {
- this.splitter = splitter;
- this.claimer = claimer;
- this.restriction = restriction;
- this.progressChecker = progressChecker;
- this.timeSupplier = () -> Timestamp.now();
- }
-
- @VisibleForTesting
- public void setTimeSupplier(Supplier<Timestamp> timeSupplier) {
- this.timeSupplier = timeSupplier;
- }
-
- @Override
- public @Nullable SplitResult<PartitionRestriction> trySplit(double fractionOfRemainder) {
- final SplitResult<PartitionRestriction> splitResult =
- splitter.trySplit(fractionOfRemainder, lastClaimedPosition, restriction);
- if (splitResult != null) {
- PartitionRestriction restrictionFromSplit =
- Optional.ofNullable(splitResult.getPrimary()).orElse(this.restriction);
- this.restriction = restrictionFromSplit;
- }
- return splitResult;
- }
-
- @Override
- public boolean tryClaim(PartitionPosition position) {
- final boolean canClaim = claimer.tryClaim(restriction, lastClaimedPosition, position);
-
- if (canClaim) {
- this.lastClaimedPosition = position;
- }
-
- return canClaim;
- }
-
- @Override
- public Progress getProgress() {
- final String token =
- Optional.ofNullable(restriction.getMetadata())
- .map(PartitionRestrictionMetadata::getPartitionToken)
- .orElse(null);
- final Progress progress = progressChecker.getProgress(restriction, lastClaimedPosition);
- LOG.debug("[" + token + "] Progress is " + progress);
- return progress;
- }
-
- @Override
- public PartitionRestriction currentRestriction() {
- return restriction;
- }
-
- @Override
- public void checkDone() throws IllegalStateException {
- if (restriction.getMode() == STOP) {
- return;
- }
- if (lastClaimedPosition != null) {
- final PartitionMode currentMode = lastClaimedPosition.getMode();
- if (currentMode == QUERY_CHANGE_STREAM) {
- // If the mode is QUERY_CHANGE_STREAM, and the end of the timestamp range wasn't attempted,
- // throw an exception
- final Timestamp nextPosition = next(lastClaimedPosition.getTimestamp().get());
- if (nextPosition.compareTo(restriction.getEndTimestamp()) < 0) {
- throw new IllegalStateException(
- String.format(
- "Last attempted key was %s in range %s, claiming work in [%s, %s) was not"
- + " attempted",
- nextPosition, restriction, nextPosition, restriction.getEndTimestamp()));
- }
- } else {
- checkState(
- currentMode == DONE, "Restriction %s does not have mode DONE", restriction.toString());
- }
- } else {
- // If nothing was attempted, throws an exception
- throw new IllegalStateException(
- String.format(
- "restriction is non-empty %s and no keys have been attempted.", restriction));
- }
- }
-
- @Override
- public IsBounded isBounded() {
- return IsBounded.UNBOUNDED;
- }
-
- @VisibleForTesting
- PartitionRestriction getRestriction() {
- return restriction;
- }
-
- @VisibleForTesting
- @Nullable
- PartitionPosition getLastClaimedPosition() {
- return lastClaimedPosition;
- }
-}
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/restriction/ReadChangeStreamPartitionRangeTracker.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/restriction/ReadChangeStreamPartitionRangeTracker.java
index ff3e43dc78d..177a1b7494a 100644
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/restriction/ReadChangeStreamPartitionRangeTracker.java
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/restriction/ReadChangeStreamPartitionRangeTracker.java
@@ -64,7 +64,7 @@ public class ReadChangeStreamPartitionRangeTracker extends TimestampRangeTracker
return true;
}
- return super.tryClaim(position);
+ return super.tryClaim(position, this.partition);
}
/**
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/restriction/TimestampRangeTracker.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/restriction/TimestampRangeTracker.java
index e2ac2773741..180070407f8 100644
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/restriction/TimestampRangeTracker.java
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/restriction/TimestampRangeTracker.java
@@ -28,6 +28,7 @@ import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Prec
import com.google.cloud.Timestamp;
import java.math.BigDecimal;
import java.util.function.Supplier;
+import org.apache.beam.sdk.io.gcp.spanner.changestreams.model.PartitionMetadata;
import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker.HasProgress;
import org.apache.beam.sdk.transforms.splittabledofn.SplitResult;
@@ -107,6 +108,29 @@ public class TimestampRangeTracker extends RestrictionTracker<TimestampRange, Ti
return true;
}
+ public boolean tryClaim(Timestamp position, PartitionMetadata partitionMetadata) {
+ checkArgument(
+ lastAttemptedPosition == null || position.compareTo(lastAttemptedPosition) > 0,
+ "Trying to claim offset %s while last attempted was %s. The partition metadata: %s",
+ position,
+ lastAttemptedPosition,
+ partitionMetadata.toString());
+ checkArgument(
+ position.compareTo(range.getFrom()) >= 0,
+ "Trying to claim offset %s before start of the range %s. The partition metadata: %s",
+ position,
+ range,
+ partitionMetadata.toString());
+ lastAttemptedPosition = position;
+
+ // It is ok to try to claim offsets beyond of the range, this will simply return false
+ if (position.compareTo(range.getTo()) >= 0) {
+ return false;
+ }
+ lastClaimedPosition = position;
+ return true;
+ }
+
/**
* Splits the restriction through the following algorithm:
*
diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/restriction/PartitionPositionTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/restriction/PartitionPositionTest.java
deleted file mode 100644
index ce6eeb17cce..00000000000
--- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/restriction/PartitionPositionTest.java
+++ /dev/null
@@ -1,70 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements. See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership. The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * "License"); you may not use this file except in compliance
- * with the License. You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-package org.apache.beam.sdk.io.gcp.spanner.changestreams.restriction;
-
-import static org.apache.beam.sdk.io.gcp.spanner.changestreams.restriction.PartitionMode.DONE;
-import static org.apache.beam.sdk.io.gcp.spanner.changestreams.restriction.PartitionMode.QUERY_CHANGE_STREAM;
-import static org.apache.beam.sdk.io.gcp.spanner.changestreams.restriction.PartitionMode.STOP;
-import static org.apache.beam.sdk.io.gcp.spanner.changestreams.restriction.PartitionMode.UPDATE_STATE;
-import static org.apache.beam.sdk.io.gcp.spanner.changestreams.restriction.PartitionMode.WAIT_FOR_CHILD_PARTITIONS;
-import static org.junit.Assert.assertEquals;
-
-import com.google.cloud.Timestamp;
-import java.util.Optional;
-import org.junit.Before;
-import org.junit.Test;
-
-public class PartitionPositionTest {
-
- private Timestamp timestamp;
-
- @Before
- public void setUp() {
- timestamp = Timestamp.now();
- }
-
- @Test
- public void testPositionUpdateState() {
- assertEquals(
- new PartitionPosition(Optional.empty(), UPDATE_STATE), PartitionPosition.updateState());
- }
-
- @Test
- public void testPositionQueryChangeStream() {
- assertEquals(
- new PartitionPosition(Optional.of(timestamp), QUERY_CHANGE_STREAM),
- PartitionPosition.queryChangeStream(timestamp));
- }
-
- @Test
- public void testPositionWaitForChildPartitions() {
- assertEquals(
- new PartitionPosition(Optional.empty(), WAIT_FOR_CHILD_PARTITIONS),
- PartitionPosition.waitForChildPartitions());
- }
-
- @Test
- public void testPositionDone() {
- assertEquals(new PartitionPosition(Optional.empty(), DONE), PartitionPosition.done());
- }
-
- @Test
- public void testPositionStop() {
- assertEquals(new PartitionPosition(Optional.empty(), STOP), PartitionPosition.stop());
- }
-}
diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/restriction/PartitionRestrictionClaimerTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/restriction/PartitionRestrictionClaimerTest.java
deleted file mode 100644
index e53d2692bf7..00000000000
--- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/restriction/PartitionRestrictionClaimerTest.java
+++ /dev/null
@@ -1,272 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements. See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership. The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * "License"); you may not use this file except in compliance
- * with the License. You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-package org.apache.beam.sdk.io.gcp.spanner.changestreams.restriction;
-
-import static org.apache.beam.sdk.io.gcp.spanner.changestreams.restriction.PartitionMode.DONE;
-import static org.apache.beam.sdk.io.gcp.spanner.changestreams.restriction.PartitionMode.QUERY_CHANGE_STREAM;
-import static org.apache.beam.sdk.io.gcp.spanner.changestreams.restriction.PartitionMode.STOP;
-import static org.apache.beam.sdk.io.gcp.spanner.changestreams.restriction.PartitionMode.UPDATE_STATE;
-import static org.apache.beam.sdk.io.gcp.spanner.changestreams.restriction.PartitionMode.WAIT_FOR_CHILD_PARTITIONS;
-import static org.junit.Assert.assertEquals;
-import static org.junit.Assert.assertFalse;
-import static org.junit.Assert.assertThrows;
-import static org.junit.Assert.assertTrue;
-
-import com.google.cloud.Timestamp;
-import java.util.Optional;
-import java.util.Set;
-import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Sets;
-import org.junit.Before;
-import org.junit.Test;
-
-public class PartitionRestrictionClaimerTest {
-
- private PartitionRestrictionClaimer claimer;
- private String partitionToken;
-
- @Before
- public void setUp() {
- partitionToken = "partitionToken";
- claimer = new PartitionRestrictionClaimer();
- }
-
- @Test
- public void testQueryChangeStreamEqualToLastClaimedTimestamp() {
- final boolean canClaim =
- claimer.tryClaim(
- PartitionRestriction.queryChangeStream(
- Timestamp.ofTimeSecondsAndNanos(10L, 0),
- Timestamp.ofTimeSecondsAndNanos(20L, 0))
- .withMetadata(
- PartitionRestrictionMetadata.newBuilder()
- .withPartitionToken(partitionToken)
- .build()),
- PartitionPosition.queryChangeStream(Timestamp.ofTimeSecondsAndNanos(15L, 0)),
- PartitionPosition.queryChangeStream(Timestamp.ofTimeSecondsAndNanos(15L, 0)));
-
- assertTrue(canClaim);
- }
-
- @Test
- public void testQueryChangeStreamWithinRestrictionRange() {
- final boolean canClaim =
- claimer.tryClaim(
- PartitionRestriction.queryChangeStream(
- Timestamp.ofTimeSecondsAndNanos(10L, 0),
- Timestamp.ofTimeSecondsAndNanos(20L, 0))
- .withMetadata(
- PartitionRestrictionMetadata.newBuilder()
- .withPartitionToken(partitionToken)
- .build()),
- PartitionPosition.queryChangeStream(Timestamp.ofTimeSecondsAndNanos(10L, 0)),
- PartitionPosition.queryChangeStream(Timestamp.ofTimeSecondsAndNanos(15L, 0)));
-
- assertTrue(canClaim);
- }
-
- @Test
- public void testQueryChangeStreamWithStartTimestamp() {
- final boolean canClaim =
- claimer.tryClaim(
- PartitionRestriction.queryChangeStream(
- Timestamp.ofTimeSecondsAndNanos(10L, 0),
- Timestamp.ofTimeSecondsAndNanos(20L, 0))
- .withMetadata(
- PartitionRestrictionMetadata.newBuilder()
- .withPartitionToken(partitionToken)
- .build()),
- PartitionPosition.queryChangeStream(Timestamp.ofTimeSecondsAndNanos(10L, 0)),
- PartitionPosition.queryChangeStream(Timestamp.ofTimeSecondsAndNanos(10L, 0)));
-
- assertTrue(canClaim);
- }
-
- @Test
- public void testQueryChangeStreamWithEndTimestamp() {
- final boolean canClaim =
- claimer.tryClaim(
- PartitionRestriction.queryChangeStream(
- Timestamp.ofTimeSecondsAndNanos(10L, 0),
- Timestamp.ofTimeSecondsAndNanos(20L, 0))
- .withMetadata(
- PartitionRestrictionMetadata.newBuilder()
- .withPartitionToken(partitionToken)
- .build()),
- PartitionPosition.queryChangeStream(Timestamp.ofTimeSecondsAndNanos(10L, 0)),
- PartitionPosition.queryChangeStream(Timestamp.ofTimeSecondsAndNanos(20L, 0)));
-
- assertFalse(canClaim);
- }
-
- @Test
- public void testQueryChangeStreamAfterEndTimestamp() {
- final boolean canClaim =
- claimer.tryClaim(
- PartitionRestriction.queryChangeStream(
- Timestamp.ofTimeSecondsAndNanos(10L, 0),
- Timestamp.ofTimeSecondsAndNanos(20L, 0))
- .withMetadata(
- PartitionRestrictionMetadata.newBuilder()
- .withPartitionToken(partitionToken)
- .build()),
- PartitionPosition.queryChangeStream(Timestamp.ofTimeSecondsAndNanos(10L, 0)),
- PartitionPosition.queryChangeStream(Timestamp.ofTimeSecondsAndNanos(20L, 1)));
-
- assertFalse(canClaim);
- }
-
- @Test(expected = IllegalArgumentException.class)
- public void testQueryChangeStreamWithoutTimestamp() {
- claimer.tryClaim(
- PartitionRestriction.queryChangeStream(Timestamp.MIN_VALUE, Timestamp.MAX_VALUE)
- .withMetadata(
- PartitionRestrictionMetadata.newBuilder()
- .withPartitionToken(partitionToken)
- .build()),
- PartitionPosition.queryChangeStream(Timestamp.MIN_VALUE),
- new PartitionPosition(Optional.empty(), QUERY_CHANGE_STREAM));
- }
-
- @Test
- public void testUpdateStateTransitions() {
- new TryClaimTestScenario(claimer)
- .from(UPDATE_STATE)
- .to(UPDATE_STATE, QUERY_CHANGE_STREAM)
- .withTryClaimResultAs(true)
- .run();
- }
-
- @Test
- public void testQueryChangeStreamTransitions() {
- new TryClaimTestScenario(claimer)
- .from(QUERY_CHANGE_STREAM)
- .to(QUERY_CHANGE_STREAM, WAIT_FOR_CHILD_PARTITIONS)
- .withTryClaimResultAs(true)
- .run();
- }
-
- @Test
- public void testWaitForChildPartitionsTransitions() {
- new TryClaimTestScenario(claimer)
- .from(WAIT_FOR_CHILD_PARTITIONS)
- .to(WAIT_FOR_CHILD_PARTITIONS, DONE)
- .withTryClaimResultAs(true)
- .run();
- }
-
- @Test
- public void testDoneTransitions() {
- new TryClaimTestScenario(claimer).from(DONE).to().run();
- }
-
- @Test
- public void testStopTransitionsAlwaysReturnsFalse() {
- new TryClaimTestScenario(claimer)
- .from(STOP)
- .to(PartitionMode.values())
- .withTryClaimResultAs(false)
- .run();
- }
-
- private static class TryClaimTestScenario {
-
- private final PartitionRestrictionClaimer claimer;
- private PartitionMode from;
- private Set<PartitionMode> validTransitions;
- private boolean expectedTryClaimResult;
-
- public TryClaimTestScenario(PartitionRestrictionClaimer claimer) {
- this.claimer = claimer;
- }
-
- public TryClaimTestScenario from(PartitionMode from) {
- this.from = from;
- return this;
- }
-
- public TryClaimTestScenario to(PartitionMode... validTransitions) {
- this.validTransitions = Sets.newHashSet(validTransitions);
- return this;
- }
-
- public TryClaimTestScenario withTryClaimResultAs(boolean tryClaimResultAs) {
- this.expectedTryClaimResult = tryClaimResultAs;
- return this;
- }
-
- public void run() {
- final PartitionRestrictionMetadata metadata =
- PartitionRestrictionMetadata.newBuilder().withPartitionToken("partitionToken").build();
- final PartitionRestriction restriction =
- partitionRestrictionFrom(from).withMetadata(metadata);
- final PartitionPosition lastClaimedPosition = partitionPositionFrom(from);
- // Tests with and without last claimed position
- for (PartitionMode to : PartitionMode.values()) {
- final PartitionPosition position = partitionPositionFrom(to);
- if (validTransitions.contains(to)) {
- assertEquals(
- expectedTryClaimResult, claimer.tryClaim(restriction, lastClaimedPosition, position));
- assertEquals(expectedTryClaimResult, claimer.tryClaim(restriction, null, position));
- } else {
- assertThrows(
- IllegalArgumentException.class,
- () -> claimer.tryClaim(restriction, lastClaimedPosition, position));
- assertThrows(
- IllegalArgumentException.class, () -> claimer.tryClaim(restriction, null, position));
- }
- }
- }
-
- private PartitionRestriction partitionRestrictionFrom(PartitionMode mode) {
- final Timestamp startTimestamp = Timestamp.MIN_VALUE;
- final Timestamp endTimestamp = Timestamp.MAX_VALUE;
- switch (mode) {
- case UPDATE_STATE:
- return PartitionRestriction.updateState(startTimestamp, endTimestamp);
- case QUERY_CHANGE_STREAM:
- return PartitionRestriction.queryChangeStream(startTimestamp, endTimestamp);
- case WAIT_FOR_CHILD_PARTITIONS:
- return PartitionRestriction.waitForChildPartitions(startTimestamp, endTimestamp);
- case DONE:
- return PartitionRestriction.done(startTimestamp, endTimestamp);
- case STOP:
- return PartitionRestriction.stop(
- PartitionRestriction.queryChangeStream(startTimestamp, endTimestamp));
- default:
- throw new IllegalArgumentException("Unknown mode " + mode);
- }
- }
-
- private PartitionPosition partitionPositionFrom(PartitionMode mode) {
- switch (mode) {
- case UPDATE_STATE:
- return PartitionPosition.updateState();
- case QUERY_CHANGE_STREAM:
- return PartitionPosition.queryChangeStream(Timestamp.ofTimeSecondsAndNanos(1L, 0));
- case WAIT_FOR_CHILD_PARTITIONS:
- return PartitionPosition.waitForChildPartitions();
- case DONE:
- return PartitionPosition.done();
- case STOP:
- return PartitionPosition.stop();
- default:
- throw new IllegalArgumentException("Unknown mode " + mode);
- }
- }
- }
-}
diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/restriction/PartitionRestrictionProgressCheckerTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/restriction/PartitionRestrictionProgressCheckerTest.java
deleted file mode 100644
index b3377dab651..00000000000
--- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/restriction/PartitionRestrictionProgressCheckerTest.java
+++ /dev/null
@@ -1,230 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements. See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership. The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * "License"); you may not use this file except in compliance
- * with the License. You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-package org.apache.beam.sdk.io.gcp.spanner.changestreams.restriction;
-
-import static org.junit.Assert.assertEquals;
-
-import com.google.cloud.Timestamp;
-import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker.Progress;
-import org.junit.Before;
-import org.junit.Test;
-
-public class PartitionRestrictionProgressCheckerTest {
-
- private PartitionRestrictionProgressChecker progressChecker;
-
- @Before
- public void setUp() {
- progressChecker = new PartitionRestrictionProgressChecker();
- }
-
- // ------------------------
- // UPDATE_STATE mode
- @Test
- public void testRestrictionUpdateStateAndLastClaimedPositionNull() {
- final PartitionRestriction restriction =
- PartitionRestriction.updateState(
- Timestamp.ofTimeSecondsAndNanos(0L, 0), Timestamp.ofTimeSecondsAndNanos(60L, 0));
-
- final Progress progress = progressChecker.getProgress(restriction, null);
-
- assertEquals(Progress.from(0D, 60D), progress);
- }
-
- @Test
- public void testRestrictionUpdateStateAndLastClaimedPositionUpdateState() {
- final PartitionRestriction restriction =
- PartitionRestriction.updateState(
- Timestamp.ofTimeSecondsAndNanos(0L, 0), Timestamp.ofTimeSecondsAndNanos(60L, 60));
-
- final PartitionPosition position = PartitionPosition.updateState();
-
- final Progress progress = progressChecker.getProgress(restriction, position);
-
- assertEquals(Progress.from(0D, 60D), progress);
- }
-
- // ------------------------
- // QUERY_CHANGE_STREAM mode
- @Test
- public void testRestrictionQueryChangeStreamAndLastClaimedPositionNull() {
- final PartitionRestriction restriction =
- PartitionRestriction.queryChangeStream(
- Timestamp.ofTimeSecondsAndNanos(0L, 10), Timestamp.ofTimeSecondsAndNanos(60L, 60));
-
- final Progress progress = progressChecker.getProgress(restriction, null);
-
- assertEquals(Progress.from(0D, 60D), progress);
- }
-
- @Test
- public void testRestrictionQueryChangeStreamAndLastClaimedPositionQueryChangeStream() {
- final PartitionRestriction restriction =
- PartitionRestriction.queryChangeStream(
- Timestamp.ofTimeSecondsAndNanos(0L, 10), Timestamp.ofTimeSecondsAndNanos(60L, 60));
- final PartitionPosition position =
- PartitionPosition.queryChangeStream(Timestamp.ofTimeSecondsAndNanos(30L, 30));
-
- final Progress progress = progressChecker.getProgress(restriction, position);
-
- assertEquals(Progress.from(30D, 30D), progress);
- }
-
- @Test
- public void testRestrictionQueryChangeStreamAndLastClaimedPositionEndOfQueryChangeStream() {
- final PartitionRestriction restriction =
- PartitionRestriction.queryChangeStream(
- Timestamp.ofTimeSecondsAndNanos(0L, 10), Timestamp.ofTimeSecondsAndNanos(60L, 60));
- final PartitionPosition position =
- PartitionPosition.queryChangeStream(Timestamp.ofTimeSecondsAndNanos(60L, 60));
-
- final Progress progress = progressChecker.getProgress(restriction, position);
-
- assertEquals(Progress.from(60D, 1D), progress);
- }
-
- @Test
- public void testRestrictionQueryChangeStreamAndLastClaimedPositionWaitForChildPartitions() {
- final PartitionRestriction restriction =
- PartitionRestriction.queryChangeStream(
- Timestamp.ofTimeSecondsAndNanos(0L, 10), Timestamp.ofTimeSecondsAndNanos(60L, 60));
- final PartitionPosition position = PartitionPosition.waitForChildPartitions();
-
- final Progress progress = progressChecker.getProgress(restriction, position);
-
- assertEquals(Progress.from(60D, 1D), progress);
- }
-
- @Test
- public void testRestrictionQueryChangeStreamAndLastClaimedPositionDone() {
- final PartitionRestriction restriction =
- PartitionRestriction.queryChangeStream(
- Timestamp.ofTimeSecondsAndNanos(0L, 10), Timestamp.ofTimeSecondsAndNanos(60L, 60));
- final PartitionPosition position = PartitionPosition.done();
-
- final Progress progress = progressChecker.getProgress(restriction, position);
-
- assertEquals(Progress.from(60D, 1D), progress);
- }
-
- // ------------------------------
- // WAIT_FOR_CHILD_PARTITIONS mode
- @Test
- public void testRestrictionWaitForChildPartitionsAndLastClaimedPositionNull() {
- final PartitionRestriction restriction =
- PartitionRestriction.waitForChildPartitions(
- Timestamp.ofTimeSecondsAndNanos(0L, 10), Timestamp.ofTimeSecondsAndNanos(60L, 60));
-
- final Progress progress = progressChecker.getProgress(restriction, null);
-
- assertEquals(Progress.from(60D, 1D), progress);
- }
-
- @Test
- public void testRestrictionWaitForChildPartitionsAndLastClaimedPositionWaitForChildPartitions() {
- final PartitionRestriction restriction =
- PartitionRestriction.waitForChildPartitions(
- Timestamp.ofTimeSecondsAndNanos(0L, 10), Timestamp.ofTimeSecondsAndNanos(60L, 60));
- final PartitionPosition position = PartitionPosition.waitForChildPartitions();
-
- final Progress progress = progressChecker.getProgress(restriction, position);
-
- assertEquals(Progress.from(60D, 1D), progress);
- }
-
- @Test
- public void testRestrictionWaitForChildPartitionsAndLastClaimedPositionDone() {
- final PartitionRestriction restriction =
- PartitionRestriction.waitForChildPartitions(
- Timestamp.ofTimeSecondsAndNanos(0L, 10), Timestamp.ofTimeSecondsAndNanos(60L, 60));
- final PartitionPosition position = PartitionPosition.done();
-
- final Progress progress = progressChecker.getProgress(restriction, position);
-
- assertEquals(Progress.from(60D, 1D), progress);
- }
-
- // ------------------------
- // DONE mode
- @Test
- public void testRestrictionDoneAndLastClaimedPositionDone() {
- final PartitionRestriction restriction =
- PartitionRestriction.done(
- Timestamp.ofTimeSecondsAndNanos(0L, 10), Timestamp.ofTimeSecondsAndNanos(60L, 60));
- final PartitionPosition position = PartitionPosition.done();
-
- final Progress progress = progressChecker.getProgress(restriction, position);
-
- assertEquals(Progress.from(60D, 1D), progress);
- }
-
- // ------------------------
- // STOP mode
- @Test
- public void testRestrictionStopQueryChangeStream() {
- final PartitionRestriction stoppedRestriction =
- PartitionRestriction.queryChangeStream(
- Timestamp.ofTimeSecondsAndNanos(0L, 10), Timestamp.ofTimeSecondsAndNanos(60L, 60));
- final PartitionRestriction restriction = PartitionRestriction.stop(stoppedRestriction);
-
- final Progress progress = progressChecker.getProgress(restriction, null);
-
- assertEquals(Progress.from(60D, 1D), progress);
- }
-
- @Test
- public void testRestrictionStopWaitForChildPartitions() {
- final PartitionRestriction stoppedRestriction =
- PartitionRestriction.waitForChildPartitions(
- Timestamp.ofTimeSecondsAndNanos(0L, 10), Timestamp.ofTimeSecondsAndNanos(60L, 60));
- final PartitionRestriction restriction = PartitionRestriction.stop(stoppedRestriction);
-
- final Progress progress = progressChecker.getProgress(restriction, null);
-
- assertEquals(Progress.from(60D, 1D), progress);
- }
-
- @Test
- public void testRestrictionStopDone() {
- final PartitionRestriction stoppedRestriction =
- PartitionRestriction.done(
- Timestamp.ofTimeSecondsAndNanos(0L, 10), Timestamp.ofTimeSecondsAndNanos(60L, 60));
- final PartitionRestriction restriction = PartitionRestriction.stop(stoppedRestriction);
-
- final Progress progress = progressChecker.getProgress(restriction, null);
-
- assertEquals(Progress.from(60D, 1D), progress);
- }
-
- @Test
- public void testGetProgressForStreaming() {
- final Timestamp from = Timestamp.ofTimeSecondsAndNanos(0, 0);
- final Timestamp current = Timestamp.ofTimeSecondsAndNanos(101, 0);
- final PartitionPosition position = PartitionPosition.queryChangeStream(current);
-
- final PartitionRestriction restriction =
- PartitionRestriction.queryChangeStream(from, Timestamp.MAX_VALUE);
-
- progressChecker.setTimeSupplier(
- () -> Timestamp.ofTimeSecondsAndNanos(position.getTimestamp().get().getSeconds() + 10, 0));
-
- final Progress progress = progressChecker.getProgress(restriction, position);
-
- assertEquals(Progress.from(101D, 10D), progress);
- }
-}
diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/restriction/PartitionRestrictionSplitterTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/restriction/PartitionRestrictionSplitterTest.java
deleted file mode 100644
index b284eca7d03..00000000000
--- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/restriction/PartitionRestrictionSplitterTest.java
+++ /dev/null
@@ -1,168 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements. See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership. The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * "License"); you may not use this file except in compliance
- * with the License. You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-package org.apache.beam.sdk.io.gcp.spanner.changestreams.restriction;
-
-import static org.junit.Assert.assertEquals;
-import static org.junit.Assert.assertNull;
-
-import com.google.cloud.Timestamp;
-import org.apache.beam.sdk.io.gcp.spanner.changestreams.model.InitialPartition;
-import org.apache.beam.sdk.transforms.splittabledofn.SplitResult;
-import org.junit.Before;
-import org.junit.Test;
-
-public class PartitionRestrictionSplitterTest {
-
- private Timestamp startTimestamp;
- private Timestamp endTimestamp;
- private PartitionRestriction restriction;
- private PartitionRestriction initialRestriction;
- private PartitionRestrictionSplitter splitter;
-
- @Before
- public void setUp() {
- startTimestamp = Timestamp.ofTimeSecondsAndNanos(0L, 0);
- endTimestamp = Timestamp.ofTimeSecondsAndNanos(100L, 50);
- restriction =
- PartitionRestriction.queryChangeStream(startTimestamp, endTimestamp)
- .withMetadata(
- PartitionRestrictionMetadata.newBuilder()
- .withPartitionToken("partitionToken")
- .build());
- initialRestriction =
- PartitionRestriction.queryChangeStream(startTimestamp, endTimestamp)
- .withMetadata(
- PartitionRestrictionMetadata.newBuilder()
- .withPartitionToken(InitialPartition.PARTITION_TOKEN)
- .build());
- splitter = new PartitionRestrictionSplitter();
- }
-
- @Test
- public void testLastClaimedPositionIsNull() {
- final SplitResult<PartitionRestriction> splitResult = splitter.trySplit(0D, null, restriction);
-
- assertNull(splitResult);
- }
-
- @Test
- public void testUpdateState() {
- final PartitionPosition position = PartitionPosition.updateState();
-
- final SplitResult<PartitionRestriction> splitResult =
- splitter.trySplit(0D, position, restriction);
-
- assertEquals(
- SplitResult.of(
- PartitionRestriction.stop(restriction),
- PartitionRestriction.queryChangeStream(startTimestamp, endTimestamp)),
- splitResult);
- }
-
- @Test
- public void testQueryChangeStreamInitialPartition() {
- final PartitionPosition position =
- PartitionPosition.queryChangeStream(Timestamp.ofTimeMicroseconds(1L));
-
- final SplitResult<PartitionRestriction> splitResult =
- splitter.trySplit(0D, position, initialRestriction);
-
- assertNull(splitResult);
- }
-
- @Test
- public void testQueryChangeStream() {
- final PartitionPosition position =
- PartitionPosition.queryChangeStream(Timestamp.ofTimeSecondsAndNanos(50L, 0));
-
- final SplitResult<PartitionRestriction> splitResult =
- splitter.trySplit(0D, position, restriction);
-
- assertEquals(
- SplitResult.of(
- PartitionRestriction.queryChangeStream(
- startTimestamp, Timestamp.ofTimeSecondsAndNanos(50L, 1)),
- PartitionRestriction.queryChangeStream(
- Timestamp.ofTimeSecondsAndNanos(50L, 1), endTimestamp)),
- splitResult);
- }
-
- @Test
- public void testQueryChangeStreamSplitEqualToEndTimestamp() {
- final PartitionPosition position =
- PartitionPosition.queryChangeStream(Timestamp.ofTimeSecondsAndNanos(100L, 49));
-
- final SplitResult<PartitionRestriction> splitResult =
- splitter.trySplit(0D, position, restriction);
-
- assertEquals(
- SplitResult.of(
- PartitionRestriction.stop(restriction),
- PartitionRestriction.waitForChildPartitions(startTimestamp, endTimestamp)),
- splitResult);
- }
-
- @Test
- public void testQueryChangeStreamSplitGreaterThanEndTimestamp() {
- final PartitionPosition position =
- PartitionPosition.queryChangeStream(Timestamp.ofTimeSecondsAndNanos(100L, 50));
-
- final SplitResult<PartitionRestriction> splitResult =
- splitter.trySplit(0D, position, restriction);
-
- assertEquals(
- SplitResult.of(
- PartitionRestriction.stop(restriction),
- PartitionRestriction.waitForChildPartitions(startTimestamp, endTimestamp)),
- splitResult);
- }
-
- @Test
- public void testWaitForChildPartitions() {
- final PartitionPosition position = PartitionPosition.waitForChildPartitions();
-
- final SplitResult<PartitionRestriction> splitResult =
- splitter.trySplit(0D, position, restriction);
-
- assertEquals(
- SplitResult.of(
- PartitionRestriction.stop(restriction),
- PartitionRestriction.waitForChildPartitions(startTimestamp, endTimestamp)),
- splitResult);
- }
-
- @Test
- public void testDone() {
- final PartitionPosition position = PartitionPosition.done();
-
- final SplitResult<PartitionRestriction> splitResult =
- splitter.trySplit(0D, position, restriction);
-
- assertNull(splitResult);
- }
-
- @Test
- public void testStop() {
- final PartitionPosition position = PartitionPosition.stop();
-
- final SplitResult<PartitionRestriction> splitResult =
- splitter.trySplit(0D, position, restriction);
-
- assertNull(splitResult);
- }
-}
diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/restriction/PartitionRestrictionTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/restriction/PartitionRestrictionTest.java
deleted file mode 100644
index 16e1552dfc7..00000000000
--- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/restriction/PartitionRestrictionTest.java
+++ /dev/null
@@ -1,78 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements. See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership. The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * "License"); you may not use this file except in compliance
- * with the License. You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-package org.apache.beam.sdk.io.gcp.spanner.changestreams.restriction;
-
-import static org.apache.beam.sdk.io.gcp.spanner.changestreams.restriction.PartitionMode.DONE;
-import static org.apache.beam.sdk.io.gcp.spanner.changestreams.restriction.PartitionMode.QUERY_CHANGE_STREAM;
-import static org.apache.beam.sdk.io.gcp.spanner.changestreams.restriction.PartitionMode.STOP;
-import static org.apache.beam.sdk.io.gcp.spanner.changestreams.restriction.PartitionMode.UPDATE_STATE;
-import static org.apache.beam.sdk.io.gcp.spanner.changestreams.restriction.PartitionMode.WAIT_FOR_CHILD_PARTITIONS;
-import static org.junit.Assert.assertEquals;
-
-import com.google.cloud.Timestamp;
-import org.junit.Before;
-import org.junit.Test;
-
-public class PartitionRestrictionTest {
-
- private Timestamp startTimestamp;
- private Timestamp endTimestamp;
-
- @Before
- public void setUp() throws Exception {
- startTimestamp = Timestamp.MIN_VALUE;
- endTimestamp = Timestamp.MAX_VALUE;
- }
-
- @Test
- public void testUpdateStateRestriction() {
- assertEquals(
- PartitionRestriction.updateState(startTimestamp, endTimestamp),
- new PartitionRestriction(startTimestamp, endTimestamp, UPDATE_STATE, null));
- }
-
- @Test
- public void testQueryChangeStreamRestriction() {
- assertEquals(
- PartitionRestriction.queryChangeStream(startTimestamp, endTimestamp),
- new PartitionRestriction(startTimestamp, endTimestamp, QUERY_CHANGE_STREAM, null));
- }
-
- @Test
- public void testWaitForChildPartitionsRestriction() {
- assertEquals(
- PartitionRestriction.waitForChildPartitions(startTimestamp, endTimestamp),
- new PartitionRestriction(startTimestamp, endTimestamp, WAIT_FOR_CHILD_PARTITIONS, null));
- }
-
- @Test
- public void testDoneRestriction() {
- assertEquals(
- PartitionRestriction.done(startTimestamp, endTimestamp),
- new PartitionRestriction(startTimestamp, endTimestamp, DONE, null));
- }
-
- @Test
- public void testStopRestriction() {
- final PartitionRestriction restriction =
- PartitionRestriction.queryChangeStream(startTimestamp, endTimestamp);
- assertEquals(
- PartitionRestriction.stop(restriction),
- new PartitionRestriction(startTimestamp, endTimestamp, STOP, QUERY_CHANGE_STREAM));
- }
-}
diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/restriction/PartitionRestrictionTrackerTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/restriction/PartitionRestrictionTrackerTest.java
deleted file mode 100644
index 7022546f041..00000000000
--- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/restriction/PartitionRestrictionTrackerTest.java
+++ /dev/null
@@ -1,859 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements. See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership. The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * "License"); you may not use this file except in compliance
- * with the License. You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-package org.apache.beam.sdk.io.gcp.spanner.changestreams.restriction;
-
-import static org.apache.beam.sdk.io.gcp.spanner.changestreams.restriction.PartitionMode.DONE;
-import static org.apache.beam.sdk.io.gcp.spanner.changestreams.restriction.PartitionMode.QUERY_CHANGE_STREAM;
-import static org.apache.beam.sdk.io.gcp.spanner.changestreams.restriction.PartitionMode.STOP;
-import static org.apache.beam.sdk.io.gcp.spanner.changestreams.restriction.PartitionMode.UPDATE_STATE;
-import static org.apache.beam.sdk.io.gcp.spanner.changestreams.restriction.PartitionMode.WAIT_FOR_CHILD_PARTITIONS;
-import static org.apache.beam.sdk.io.gcp.spanner.changestreams.restriction.TimestampUtils.next;
-import static org.hamcrest.Matchers.anyOf;
-import static org.hamcrest.Matchers.equalTo;
-import static org.hamcrest.Matchers.greaterThanOrEqualTo;
-import static org.hamcrest.Matchers.lessThan;
-import static org.hamcrest.Matchers.lessThanOrEqualTo;
-import static org.junit.Assert.assertEquals;
-import static org.junit.Assert.assertFalse;
-import static org.junit.Assert.assertThrows;
-import static org.junit.Assert.assertTrue;
-import static org.junit.Assume.assumeThat;
-
-import com.google.cloud.Timestamp;
-import com.pholser.junit.quickcheck.From;
-import com.pholser.junit.quickcheck.Property;
-import com.pholser.junit.quickcheck.runner.JUnitQuickcheck;
-import java.util.Optional;
-import org.apache.beam.sdk.io.gcp.spanner.changestreams.util.PartitionPositionGenerator;
-import org.apache.beam.sdk.io.gcp.spanner.changestreams.util.TimestampGenerator;
-import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker.Progress;
-import org.apache.beam.sdk.transforms.splittabledofn.SplitResult;
-import org.junit.runner.RunWith;
-
-@RunWith(JUnitQuickcheck.class)
-public class PartitionRestrictionTrackerTest {
-
- private static final String PARTITION_TOKEN = "partitionToken";
- private static final double DELTA = 1e-10;
-
- // Test that trySplit returns null when no position was claimed.
- @Property
- public void testTrySplitReturnsNullhenNoPositionWasClaimed(
- @From(TimestampGenerator.class) Timestamp from,
- @From(TimestampGenerator.class) Timestamp to) {
- assumeThat(from, lessThanOrEqualTo(to));
-
- final PartitionRestriction restriction = PartitionRestriction.queryChangeStream(from, to);
- final PartitionRestrictionTracker tracker = new PartitionRestrictionTracker(restriction);
-
- final SplitResult<PartitionRestriction> splitResult = tracker.trySplit(0D);
-
- assertEquals(null, splitResult);
- }
-
- // Test that checkDone fails if the mode is not STOP and no position was claimed.
- @Property
- public void testCheckDoneFailsWhenNoPositionClaimed(
- @From(TimestampGenerator.class) Timestamp from,
- @From(TimestampGenerator.class) Timestamp to) {
- assumeThat(from, lessThan(to));
-
- final PartitionRestriction restriction = PartitionRestriction.queryChangeStream(from, to);
- final PartitionRestrictionTracker tracker = new PartitionRestrictionTracker(restriction);
-
- assertThrows(IllegalStateException.class, tracker::checkDone);
- }
-
- // ---------------------------Restriction mode is UPDATE_STATE---------------------------
-
- // Test transitioning from UPDATE_STATE to UPDATE_STATE for tryClaim
- @Property
- public void testTryClaimPositionUpdateStateToUpdateState(
- @From(TimestampGenerator.class) Timestamp from,
- @From(TimestampGenerator.class) Timestamp to,
- @From(PartitionPositionGenerator.class) PartitionPosition position) {
- assumeThat(from, lessThan(to));
- assumeThat(position.getMode(), equalTo(UPDATE_STATE));
-
- final PartitionRestriction restriction =
- PartitionRestriction.updateState(from, to)
- .withMetadata(
- PartitionRestrictionMetadata.newBuilder()
- .withPartitionToken(PARTITION_TOKEN)
- .build());
- final PartitionRestrictionTracker tracker = new PartitionRestrictionTracker(restriction);
-
- assertTrue(tracker.tryClaim(position));
- assertEquals(position, tracker.getLastClaimedPosition());
- }
-
- // Test transitioning from UPDATE_STATE to QUERY_CHANGE_STREAM for tryClaim
- @Property
- public void testTryClaimPositionUpdateStateToQueryChangeStream(
- @From(TimestampGenerator.class) Timestamp from,
- @From(TimestampGenerator.class) Timestamp to,
- @From(PartitionPositionGenerator.class) PartitionPosition position) {
- assumeThat(from, lessThan(to));
- assumeThat(position.getMode(), equalTo(QUERY_CHANGE_STREAM));
- assumeThat(position.getTimestamp().get(), greaterThanOrEqualTo(from));
- assumeThat(position.getTimestamp().get(), lessThan(to));
-
- final PartitionRestriction restriction =
- PartitionRestriction.updateState(from, to)
- .withMetadata(
- PartitionRestrictionMetadata.newBuilder()
- .withPartitionToken(PARTITION_TOKEN)
- .build());
- final PartitionRestrictionTracker tracker = new PartitionRestrictionTracker(restriction);
-
- assertTrue(tracker.tryClaim(position));
- assertEquals(position, tracker.getLastClaimedPosition());
- }
-
- // Test transitioning from UPDATE_STATE to all other modes for tryClaim fail
- @Property
- public void testTryClaimPositionUpdateStateToOtherModes(
- @From(TimestampGenerator.class) Timestamp from,
- @From(TimestampGenerator.class) Timestamp to,
- @From(PartitionPositionGenerator.class) PartitionPosition position) {
- assumeThat(from, lessThan(to));
- assumeThat(
- position.getMode(),
- anyOf(equalTo(WAIT_FOR_CHILD_PARTITIONS), equalTo(DONE), equalTo(STOP)));
-
- final PartitionRestriction restriction =
- PartitionRestriction.updateState(from, to)
- .withMetadata(
- PartitionRestrictionMetadata.newBuilder()
- .withPartitionToken(PARTITION_TOKEN)
- .build());
- final PartitionRestrictionTracker tracker = new PartitionRestrictionTracker(restriction);
-
- assertThrows(IllegalArgumentException.class, () -> tracker.tryClaim(position));
- }
-
- // Test trySplit behavior when UPDATE_STATE is claimed.
- @Property
- public void testTrySplitReturnsQueryChangeStreamhenUpdateStateClaimed(
- @From(TimestampGenerator.class) Timestamp from,
- @From(TimestampGenerator.class) Timestamp to) {
- assumeThat(from, lessThanOrEqualTo(to));
-
- final PartitionRestriction restriction =
- PartitionRestriction.updateState(from, to)
- .withMetadata(
- PartitionRestrictionMetadata.newBuilder()
- .withPartitionToken(PARTITION_TOKEN)
- .build());
- final PartitionRestrictionTracker tracker = new PartitionRestrictionTracker(restriction);
-
- final PartitionPosition position = PartitionPosition.updateState();
- assertTrue(tracker.tryClaim(position));
-
- final SplitResult<PartitionRestriction> splitResult = tracker.trySplit(0D);
- final PartitionRestriction primary = splitResult.getPrimary();
- final PartitionRestriction residual = splitResult.getResidual();
-
- assertEquals(PartitionRestriction.stop(restriction), primary);
- assertEquals(PartitionRestriction.queryChangeStream(from, to), residual);
- assertEquals(primary, tracker.restriction);
- }
-
- // Test that checkDone fails with UPDATE_STATE is the last claimed position.
- @Property
- public void testCheckDoneFailsWithUpdateState(
- @From(TimestampGenerator.class) Timestamp from,
- @From(TimestampGenerator.class) Timestamp to) {
- assumeThat(from, lessThan(to));
-
- final PartitionRestriction restriction =
- PartitionRestriction.updateState(from, to)
- .withMetadata(
- PartitionRestrictionMetadata.newBuilder()
- .withPartitionToken(PARTITION_TOKEN)
- .build());
- final PartitionRestrictionTracker tracker = new PartitionRestrictionTracker(restriction);
-
- final PartitionPosition position = PartitionPosition.updateState();
- assertTrue(tracker.tryClaim(position));
-
- assertThrows(IllegalStateException.class, tracker::checkDone);
- }
-
- // ---------------------------Restriction mode is QUERY_CHANGE_STREAM---------------------------
-
- // Test that tryClaim succeeds with QUERY_CHANGE_STREAM if timestamp is within the range.
- @Property
- public void testTryClaimReturnsTrueWhenPositionIsWithinTheRange(
- @From(TimestampGenerator.class) Timestamp from,
- @From(TimestampGenerator.class) Timestamp to,
- @From(TimestampGenerator.class) Timestamp curTimestamp) {
- assumeThat(from, lessThan(to));
- assumeThat(curTimestamp, greaterThanOrEqualTo(from));
- assumeThat(curTimestamp, lessThan(to));
-
- final PartitionRestriction restriction =
- PartitionRestriction.queryChangeStream(from, to)
- .withMetadata(
- PartitionRestrictionMetadata.newBuilder()
- .withPartitionToken(PARTITION_TOKEN)
- .build());
- final PartitionRestrictionTracker tracker = new PartitionRestrictionTracker(restriction);
- final PartitionPosition position =
- new PartitionPosition(Optional.of(curTimestamp), QUERY_CHANGE_STREAM);
-
- assertTrue(tracker.tryClaim(position));
- assertEquals(position, tracker.getLastClaimedPosition());
- }
-
- // Test that tryClaim throws an error if timestamp is greater than the range end timestamp.
- @Property
- public void testTryClaimReturnsFalseWhenPositionIsGreaterThanRange(
- @From(TimestampGenerator.class) Timestamp from,
- @From(TimestampGenerator.class) Timestamp to,
- @From(TimestampGenerator.class) Timestamp curTimestamp) {
- assumeThat(from, lessThan(to));
- assumeThat(curTimestamp, greaterThanOrEqualTo(from));
- assumeThat(curTimestamp, greaterThanOrEqualTo(to));
-
- final PartitionRestriction restriction =
- PartitionRestriction.queryChangeStream(from, to)
- .withMetadata(
- PartitionRestrictionMetadata.newBuilder()
- .withPartitionToken(PARTITION_TOKEN)
- .build());
- final PartitionRestrictionTracker tracker = new PartitionRestrictionTracker(restriction);
- final PartitionPosition position =
- new PartitionPosition(Optional.of(curTimestamp), QUERY_CHANGE_STREAM);
-
- assertFalse(tracker.tryClaim(position));
- assertEquals(null, tracker.getLastClaimedPosition());
- }
-
- // Test that tryClaim throws an error if the position is before the start of the range.
- @Property
- public void testTryClaimThrowsErrorWhenPositionBeforeRange(
- @From(TimestampGenerator.class) Timestamp from,
- @From(TimestampGenerator.class) Timestamp to,
- @From(TimestampGenerator.class) Timestamp curTimestamp) {
- assumeThat(from, lessThan(to));
- assumeThat(curTimestamp, lessThan(from));
-
- final PartitionRestriction restriction =
- PartitionRestriction.queryChangeStream(from, to)
- .withMetadata(
- PartitionRestrictionMetadata.newBuilder()
- .withPartitionToken(PARTITION_TOKEN)
- .build());
- final PartitionRestrictionTracker tracker = new PartitionRestrictionTracker(restriction);
- final PartitionPosition position =
- new PartitionPosition(Optional.of(curTimestamp), QUERY_CHANGE_STREAM);
-
- assertThrows(IllegalArgumentException.class, () -> tracker.tryClaim(position));
- }
-
- // Test that tryClaim succeeds if the timestamp is greater than or equal to the previous position.
- @Property
- public void testTryClaimPositionGreaterThanOrEqualToPreviousPosition(
- @From(TimestampGenerator.class) Timestamp from,
- @From(TimestampGenerator.class) Timestamp to,
- @From(TimestampGenerator.class) Timestamp curTimestamp,
- @From(TimestampGenerator.class) Timestamp nextTimestamp) {
- assumeThat(from, lessThan(to));
- assumeThat(curTimestamp, greaterThanOrEqualTo(from));
- assumeThat(nextTimestamp, greaterThanOrEqualTo(curTimestamp));
- assumeThat(nextTimestamp, lessThan(to));
-
- final PartitionRestriction restriction =
- PartitionRestriction.queryChangeStream(from, to)
- .withMetadata(
- PartitionRestrictionMetadata.newBuilder()
- .withPartitionToken(PARTITION_TOKEN)
- .build());
- final PartitionRestrictionTracker tracker = new PartitionRestrictionTracker(restriction);
- final PartitionPosition position =
- new PartitionPosition(Optional.of(curTimestamp), QUERY_CHANGE_STREAM);
- final PartitionPosition nextPosition =
- new PartitionPosition(Optional.of(nextTimestamp), QUERY_CHANGE_STREAM);
-
- assertTrue(tracker.tryClaim(position));
- assertEquals(position, tracker.getLastClaimedPosition());
-
- assertTrue(tracker.tryClaim(nextPosition));
- assertEquals(nextPosition, tracker.getLastClaimedPosition());
- }
-
- // Test that tryClaim throws an error if it claims a timestamp less than the previous position.
- @Property
- public void testTryClaimPositionLessThanPreviousPosition(
- @From(TimestampGenerator.class) Timestamp from,
- @From(TimestampGenerator.class) Timestamp to,
- @From(TimestampGenerator.class) Timestamp prevTimestamp,
- @From(TimestampGenerator.class) Timestamp curTimestamp) {
- assumeThat(from, lessThan(to));
- assumeThat(prevTimestamp, greaterThanOrEqualTo(from));
- assumeThat(prevTimestamp, lessThan(curTimestamp));
- assumeThat(curTimestamp, greaterThanOrEqualTo(from));
- assumeThat(curTimestamp, lessThan(to));
-
- final PartitionRestriction restriction =
- PartitionRestriction.queryChangeStream(from, to)
- .withMetadata(
- PartitionRestrictionMetadata.newBuilder()
- .withPartitionToken(PARTITION_TOKEN)
- .build());
- final PartitionRestrictionTracker tracker = new PartitionRestrictionTracker(restriction);
- final PartitionPosition position =
- new PartitionPosition(Optional.of(curTimestamp), QUERY_CHANGE_STREAM);
- final PartitionPosition prevPosition =
- new PartitionPosition(Optional.of(prevTimestamp), QUERY_CHANGE_STREAM);
-
- assertTrue(tracker.tryClaim(position));
- assertEquals(position, tracker.getLastClaimedPosition());
-
- assertThrows(IllegalArgumentException.class, () -> tracker.tryClaim(prevPosition));
- }
-
- @Property
- public void testTryClaimPositionQueryChangeStreamToWaitForChildPartitions(
- @From(TimestampGenerator.class) Timestamp from,
- @From(TimestampGenerator.class) Timestamp to,
- @From(PartitionPositionGenerator.class) PartitionPosition position) {
- assumeThat(from, lessThan(to));
- assumeThat(position.getMode(), equalTo(WAIT_FOR_CHILD_PARTITIONS));
-
- final PartitionRestriction restriction =
- PartitionRestriction.queryChangeStream(from, to)
- .withMetadata(
- PartitionRestrictionMetadata.newBuilder()
- .withPartitionToken(PARTITION_TOKEN)
- .build());
- final PartitionRestrictionTracker tracker = new PartitionRestrictionTracker(restriction);
-
- assertTrue(tracker.tryClaim(position));
- assertEquals(position, tracker.getLastClaimedPosition());
- }
-
- // Test that tryClaim succeeds if it transitions to WAIT_FOR_CHILD_PARTITIONS.
- @Property
- public void testTryClaimPositionQueryChangeStreamToOtherModes(
- @From(TimestampGenerator.class) Timestamp from,
- @From(TimestampGenerator.class) Timestamp to,
- @From(PartitionPositionGenerator.class) PartitionPosition position) {
- assumeThat(from, lessThan(to));
- assumeThat(position.getMode(), anyOf(equalTo(UPDATE_STATE), equalTo(STOP), equalTo(STOP)));
-
- final PartitionRestriction restriction =
- PartitionRestriction.queryChangeStream(from, to)
- .withMetadata(
- PartitionRestrictionMetadata.newBuilder()
- .withPartitionToken(PARTITION_TOKEN)
- .build());
- final PartitionRestrictionTracker tracker = new PartitionRestrictionTracker(restriction);
-
- assertThrows(IllegalArgumentException.class, () -> tracker.tryClaim(position));
- }
-
- // Test trySplit behavior for QUERY_CHANGE_STREAM when the split point is within the range.
- @Property
- public void testTrySplitQueryChangeStreamTimestampWithinRange(
- @From(TimestampGenerator.class) Timestamp from,
- @From(TimestampGenerator.class) Timestamp to,
- @From(TimestampGenerator.class) Timestamp curTimestamp) {
- assumeThat(from, lessThan(to));
- assumeThat(curTimestamp, greaterThanOrEqualTo(from));
- assumeThat(curTimestamp, lessThan(to));
-
- final PartitionRestriction restriction =
- PartitionRestriction.queryChangeStream(from, to)
- .withMetadata(
- PartitionRestrictionMetadata.newBuilder()
- .withPartitionToken(PARTITION_TOKEN)
- .build());
- final PartitionRestrictionTracker tracker = new PartitionRestrictionTracker(restriction);
-
- final PartitionPosition position = PartitionPosition.queryChangeStream(curTimestamp);
- assertTrue(tracker.tryClaim(position));
-
- final SplitResult<PartitionRestriction> splitResult = tracker.trySplit(0D);
- final PartitionRestriction primary = splitResult.getPrimary();
- final PartitionRestriction residual = splitResult.getResidual();
-
- assertEquals(PartitionRestriction.queryChangeStream(from, next(curTimestamp)), primary);
- assertEquals(PartitionRestriction.queryChangeStream(next(curTimestamp), to), residual);
- assertEquals(primary, tracker.restriction);
- }
-
- // Test trySplit behavior for QUERY_CHANGE_STREAM when the split point is at end of the range.
- @Property
- public void testTrySplitQueryChangeStreamTimestampEndOfRange(
- @From(TimestampGenerator.class) Timestamp from,
- @From(TimestampGenerator.class) Timestamp to) {
- assumeThat(from, lessThan(to));
-
- final PartitionRestriction restriction =
- PartitionRestriction.queryChangeStream(from, next(to))
- .withMetadata(
- PartitionRestrictionMetadata.newBuilder()
- .withPartitionToken(PARTITION_TOKEN)
- .build());
- final PartitionRestrictionTracker tracker = new PartitionRestrictionTracker(restriction);
-
- final PartitionPosition position = PartitionPosition.queryChangeStream(to);
- assertTrue(tracker.tryClaim(position));
-
- final SplitResult<PartitionRestriction> splitResult = tracker.trySplit(0D);
- final PartitionRestriction primary = splitResult.getPrimary();
- final PartitionRestriction residual = splitResult.getResidual();
-
- assertEquals(PartitionRestriction.stop(restriction), primary);
- assertEquals(PartitionRestriction.waitForChildPartitions(from, next(to)), residual);
- assertEquals(primary, tracker.restriction);
- }
-
- // Test checkDone fails when mode is QUERY_CHANGE_STREAM and last timestamp was not attempted.
- @Property
- public void testCheckDoneFailsNotAtEndOfRange(
- @From(TimestampGenerator.class) Timestamp from,
- @From(TimestampGenerator.class) Timestamp to) {
- assumeThat(from, lessThan(to));
-
- final PartitionRestriction restriction =
- PartitionRestriction.queryChangeStream(from, next(to))
- .withMetadata(
- PartitionRestrictionMetadata.newBuilder()
- .withPartitionToken(PARTITION_TOKEN)
- .build());
- final PartitionRestrictionTracker tracker = new PartitionRestrictionTracker(restriction);
-
- final PartitionPosition position = PartitionPosition.queryChangeStream(from);
- assertTrue(tracker.tryClaim(position));
- assertThrows(IllegalStateException.class, tracker::checkDone);
- }
-
- // Test checkDone succeeds when mode is QUERY_CHANGE_STREAM and last timestamp was claimed.
- @Property
- public void testCheckDoneSucceedsAtEndOfRange(
- @From(TimestampGenerator.class) Timestamp from,
- @From(TimestampGenerator.class) Timestamp to) {
- assumeThat(from, lessThan(to));
-
- final PartitionRestriction restriction =
- PartitionRestriction.queryChangeStream(from, next(to))
- .withMetadata(
- PartitionRestrictionMetadata.newBuilder()
- .withPartitionToken(PARTITION_TOKEN)
- .build());
- final PartitionRestrictionTracker tracker = new PartitionRestrictionTracker(restriction);
-
- final PartitionPosition position = PartitionPosition.queryChangeStream(to);
- assertTrue(tracker.tryClaim(position));
- tracker.checkDone();
- }
-
- // ---------------------------Restriction mode is WAIT_FOR_CHILD_PARTITIONS---------------------
-
- // Test TryClaim transitioning from WAIT_FOR_CHILD_PARTITIONS to WAIT_FOR_CHILD_PARTITIONS.
- @Property
- public void testTryClaimPositionWaitForChildPartitions(
- @From(TimestampGenerator.class) Timestamp from,
- @From(TimestampGenerator.class) Timestamp to,
- @From(PartitionPositionGenerator.class) PartitionPosition position) {
- assumeThat(from, lessThan(to));
- assumeThat(position.getMode(), anyOf(equalTo(WAIT_FOR_CHILD_PARTITIONS), equalTo(DONE)));
-
- final PartitionRestriction restriction =
- PartitionRestriction.waitForChildPartitions(from, to)
- .withMetadata(
- PartitionRestrictionMetadata.newBuilder()
- .withPartitionToken(PARTITION_TOKEN)
- .build());
- final PartitionRestrictionTracker tracker = new PartitionRestrictionTracker(restriction);
-
- assertTrue(tracker.tryClaim(position));
- assertEquals(position, tracker.getLastClaimedPosition());
- }
-
- // Test TryClaim transitioning from WAIT_FOR_CHILD_PARTITIONS to invalid states.
- @Property
- public void testTryClaimPositionWaitForChildPartitionsToInvalidStates(
- @From(TimestampGenerator.class) Timestamp from,
- @From(TimestampGenerator.class) Timestamp to,
- @From(PartitionPositionGenerator.class) PartitionPosition position) {
- assumeThat(from, lessThan(to));
- assumeThat(
- position.getMode(),
- anyOf(equalTo(UPDATE_STATE), anyOf(equalTo(QUERY_CHANGE_STREAM), equalTo(STOP))));
-
- final PartitionRestriction restriction =
- PartitionRestriction.waitForChildPartitions(from, to)
- .withMetadata(
- PartitionRestrictionMetadata.newBuilder()
- .withPartitionToken(PARTITION_TOKEN)
- .build());
- final PartitionRestrictionTracker tracker = new PartitionRestrictionTracker(restriction);
-
- assertThrows(IllegalArgumentException.class, () -> tracker.tryClaim(position));
- }
-
- // Test trySplit when last claimed position was WAIT_FOR_CHILD_PARTITIONS.
- @Property
- public void testTrySplitWaitForChildPartitions(
- @From(TimestampGenerator.class) Timestamp from,
- @From(TimestampGenerator.class) Timestamp to) {
- assumeThat(from, lessThan(to));
-
- final PartitionRestriction restriction =
- PartitionRestriction.waitForChildPartitions(from, to)
- .withMetadata(
- PartitionRestrictionMetadata.newBuilder()
- .withPartitionToken(PARTITION_TOKEN)
- .build());
- final PartitionRestrictionTracker tracker = new PartitionRestrictionTracker(restriction);
-
- final PartitionPosition position = PartitionPosition.waitForChildPartitions();
- assertTrue(tracker.tryClaim(position));
-
- final SplitResult<PartitionRestriction> splitResult = tracker.trySplit(0D);
- final PartitionRestriction primary = splitResult.getPrimary();
- final PartitionRestriction residual = splitResult.getResidual();
-
- assertEquals(PartitionRestriction.stop(restriction), primary);
- assertEquals(PartitionRestriction.waitForChildPartitions(from, to), residual);
- assertEquals(primary, tracker.restriction);
- }
-
- // Test checkDone() behavior when waiting for child partitions.
- @Property
- public void testCheckDoneWaitForChildPartitions(
- @From(TimestampGenerator.class) Timestamp from,
- @From(TimestampGenerator.class) Timestamp to) {
- assumeThat(from, lessThan(to));
-
- final PartitionRestriction restriction =
- PartitionRestriction.waitForChildPartitions(from, to)
- .withMetadata(
- PartitionRestrictionMetadata.newBuilder()
- .withPartitionToken(PARTITION_TOKEN)
- .build());
- final PartitionRestrictionTracker tracker = new PartitionRestrictionTracker(restriction);
-
- final PartitionPosition position = PartitionPosition.waitForChildPartitions();
- assertTrue(tracker.tryClaim(position));
-
- assertThrows(IllegalStateException.class, tracker::checkDone);
- }
-
- // ---------------------------Restriction mode is WAIT_FOR_CHILD_PARTITIONS---------------------
-
- // Test claiming anything fails when DONE has been claimed.
- @Property
- public void testTryClaimPositionDone(
- @From(TimestampGenerator.class) Timestamp from,
- @From(TimestampGenerator.class) Timestamp to,
- @From(PartitionPositionGenerator.class) PartitionPosition position) {
- assumeThat(from, lessThan(to));
-
- final PartitionRestriction restriction =
- PartitionRestriction.done(from, to)
- .withMetadata(
- PartitionRestrictionMetadata.newBuilder()
- .withPartitionToken(PARTITION_TOKEN)
- .build());
- final PartitionRestrictionTracker tracker = new PartitionRestrictionTracker(restriction);
-
- assertThrows(IllegalArgumentException.class, () -> tracker.tryClaim(position));
- }
-
- // Test that after DONE has been claimed, a restriction cannot be split.
- @Property
- public void testTrySplitClaimedDone(
- @From(TimestampGenerator.class) Timestamp from,
- @From(TimestampGenerator.class) Timestamp to) {
- assumeThat(from, lessThan(to));
-
- final PartitionRestriction restriction =
- PartitionRestriction.waitForChildPartitions(from, to)
- .withMetadata(
- PartitionRestrictionMetadata.newBuilder()
- .withPartitionToken(PARTITION_TOKEN)
- .build());
- final PartitionRestrictionTracker tracker = new PartitionRestrictionTracker(restriction);
-
- final PartitionPosition position = PartitionPosition.done();
- assertTrue(tracker.tryClaim(position));
-
- final SplitResult<PartitionRestriction> splitResult = tracker.trySplit(0D);
- assertEquals(null, splitResult);
- }
-
- // Test checkDone succeeds with DONE.
- @Property
- public void testCheckDoneSucceedsWithDone(
- @From(TimestampGenerator.class) Timestamp from,
- @From(TimestampGenerator.class) Timestamp to) {
- assumeThat(from, lessThan(to));
-
- final PartitionRestriction restriction =
- PartitionRestriction.waitForChildPartitions(from, to)
- .withMetadata(
- PartitionRestrictionMetadata.newBuilder()
- .withPartitionToken(PARTITION_TOKEN)
- .build());
- final PartitionRestrictionTracker tracker = new PartitionRestrictionTracker(restriction);
-
- final PartitionPosition position = PartitionPosition.done();
- assertTrue(tracker.tryClaim(position));
-
- tracker.checkDone();
- }
-
- // ---------------------------Restriction mode is DONE-----------------------------------------
-
- // Test that tryClaim returns False when restriction has mode STOP.
- @Property
- public void testTryClaimPositionStop(
- @From(TimestampGenerator.class) Timestamp from,
- @From(TimestampGenerator.class) Timestamp to,
- @From(PartitionPositionGenerator.class) PartitionPosition position) {
- assumeThat(from, lessThan(to));
-
- final PartitionRestriction restriction =
- PartitionRestriction.waitForChildPartitions(from, to)
- .withMetadata(
- PartitionRestrictionMetadata.newBuilder()
- .withPartitionToken(PARTITION_TOKEN)
- .build());
-
- final PartitionRestriction stoppedRestriction = PartitionRestriction.stop(restriction);
- final PartitionRestrictionTracker tracker = new PartitionRestrictionTracker(stoppedRestriction);
-
- assertFalse(tracker.tryClaim(position));
- }
-
- // Test splitting fails when restriction has mode STOP.
- @Property
- public void testTrySplitStop(
- @From(TimestampGenerator.class) Timestamp from,
- @From(TimestampGenerator.class) Timestamp to) {
- assumeThat(from, lessThan(to));
-
- final PartitionRestriction restriction =
- PartitionRestriction.waitForChildPartitions(from, to)
- .withMetadata(
- PartitionRestrictionMetadata.newBuilder()
- .withPartitionToken(PARTITION_TOKEN)
- .build());
-
- final PartitionRestriction stoppedRestriction = PartitionRestriction.stop(restriction);
- final PartitionRestrictionTracker tracker = new PartitionRestrictionTracker(stoppedRestriction);
-
- final SplitResult<PartitionRestriction> splitResult = tracker.trySplit(0D);
- assertEquals(null, splitResult);
- }
-
- // Test checkDone succeeds when restriction has mode STOP.
- @Property
- public void testCheckDoneSucceedsWithStop(
- @From(TimestampGenerator.class) Timestamp from,
- @From(TimestampGenerator.class) Timestamp to) {
- assumeThat(from, lessThan(to));
-
- final PartitionRestriction restriction =
- PartitionRestriction.waitForChildPartitions(from, to)
- .withMetadata(
- PartitionRestrictionMetadata.newBuilder()
- .withPartitionToken(PARTITION_TOKEN)
- .build());
-
- final PartitionRestriction stoppedRestriction = PartitionRestriction.stop(restriction);
- final PartitionRestrictionTracker tracker = new PartitionRestrictionTracker(stoppedRestriction);
-
- tracker.checkDone();
- }
-
- // ---------------------------Test progress tracking-----------------------------------------
- @Property
- public void testGetProgressWorkCompletedAndWorkRemaining(
- @From(TimestampGenerator.class) Timestamp from,
- @From(TimestampGenerator.class) Timestamp to,
- @From(TimestampGenerator.class) Timestamp curTimestamp) {
- assumeThat(from, lessThan(to));
- assumeThat(curTimestamp, greaterThanOrEqualTo(from));
- assumeThat(curTimestamp, lessThan(to));
- assumeThat(from.getSeconds(), greaterThanOrEqualTo(1L));
-
- final PartitionRestriction restriction =
- PartitionRestriction.queryChangeStream(from, to)
- .withMetadata(
- PartitionRestrictionMetadata.newBuilder()
- .withPartitionToken(PARTITION_TOKEN)
- .build());
- final PartitionRestrictionTracker tracker = new PartitionRestrictionTracker(restriction);
-
- final PartitionPosition position = PartitionPosition.queryChangeStream(curTimestamp);
-
- assertTrue(tracker.tryClaim(position));
- final Progress progress = tracker.getProgress();
-
- assertEquals(position.getTimestamp().get().getSeconds(), progress.getWorkCompleted(), DELTA);
- assertEquals(
- to.getSeconds() - position.getTimestamp().get().getSeconds(),
- progress.getWorkRemaining(),
- DELTA);
- }
-
- @Property
- public void testGetProgressWaitForChildPartitions(
- @From(TimestampGenerator.class) Timestamp from,
- @From(TimestampGenerator.class) Timestamp to) {
- assumeThat(from, lessThan(to));
- assumeThat(from.getSeconds(), greaterThanOrEqualTo(1L));
-
- final PartitionRestriction restriction =
- PartitionRestriction.queryChangeStream(from, to)
- .withMetadata(
- PartitionRestrictionMetadata.newBuilder()
- .withPartitionToken(PARTITION_TOKEN)
- .build());
- final PartitionRestrictionTracker tracker = new PartitionRestrictionTracker(restriction);
-
- final PartitionPosition position = PartitionPosition.waitForChildPartitions();
-
- assertTrue(tracker.tryClaim(position));
- final Progress progress = tracker.getProgress();
-
- assertEquals(to.getSeconds(), progress.getWorkCompleted(), DELTA);
- assertEquals(1, progress.getWorkRemaining(), DELTA);
- }
-
- @Property
- public void testGetProgressWaitForChildPartitionsNoStateClaimed(
- @From(TimestampGenerator.class) Timestamp from,
- @From(TimestampGenerator.class) Timestamp to) {
- assumeThat(from, greaterThanOrEqualTo(Timestamp.MIN_VALUE));
- assumeThat(from, lessThan(to));
- assumeThat(from.getSeconds(), greaterThanOrEqualTo(1L));
-
- final PartitionRestriction restriction =
- PartitionRestriction.waitForChildPartitions(from, to)
- .withMetadata(
- PartitionRestrictionMetadata.newBuilder()
- .withPartitionToken(PARTITION_TOKEN)
- .build());
- final PartitionRestrictionTracker tracker = new PartitionRestrictionTracker(restriction);
-
- final Progress progress = tracker.getProgress();
-
- assertEquals(to.getSeconds(), progress.getWorkCompleted(), DELTA);
- assertEquals(1, progress.getWorkRemaining(), DELTA);
- }
-
- @Property
- public void testGetProgressUpdateState(
- @From(TimestampGenerator.class) Timestamp from,
- @From(TimestampGenerator.class) Timestamp to) {
- assumeThat(from, lessThan(to));
- assumeThat(from.getSeconds(), greaterThanOrEqualTo(1L));
-
- final PartitionRestriction restriction =
- PartitionRestriction.updateState(from, to)
- .withMetadata(
- PartitionRestrictionMetadata.newBuilder()
- .withPartitionToken(PARTITION_TOKEN)
- .build());
- final PartitionRestrictionTracker tracker = new PartitionRestrictionTracker(restriction);
-
- final PartitionPosition position = PartitionPosition.updateState();
-
- assertTrue(tracker.tryClaim(position));
- final Progress progress = tracker.getProgress();
-
- assertEquals(from.getSeconds(), progress.getWorkCompleted(), DELTA);
- assertEquals(to.getSeconds() - from.getSeconds(), progress.getWorkRemaining(), DELTA);
- }
-
- @Property
- public void testGetProgressUpdateStateNoPositionClaimed(
- @From(TimestampGenerator.class) Timestamp from,
- @From(TimestampGenerator.class) Timestamp to) {
- assumeThat(from, lessThan(to));
- assumeThat(from.getSeconds(), greaterThanOrEqualTo(1L));
-
- final PartitionRestriction restriction =
- PartitionRestriction.updateState(from, to)
- .withMetadata(
- PartitionRestrictionMetadata.newBuilder()
- .withPartitionToken(PARTITION_TOKEN)
- .build());
- final PartitionRestrictionTracker tracker = new PartitionRestrictionTracker(restriction);
-
- final Progress progress = tracker.getProgress();
-
- assertEquals(from.getSeconds(), progress.getWorkCompleted(), DELTA);
- assertEquals(to.getSeconds() - from.getSeconds(), progress.getWorkRemaining(), DELTA);
- }
-
- @Property
- public void testGetProgressDone(
- @From(TimestampGenerator.class) Timestamp from,
- @From(TimestampGenerator.class) Timestamp to) {
- assumeThat(from, lessThan(to));
- assumeThat(from.getSeconds(), greaterThanOrEqualTo(1L));
-
- final PartitionRestriction restriction =
- PartitionRestriction.waitForChildPartitions(from, to)
- .withMetadata(
- PartitionRestrictionMetadata.newBuilder()
- .withPartitionToken(PARTITION_TOKEN)
- .build());
- final PartitionRestrictionTracker tracker = new PartitionRestrictionTracker(restriction);
-
- final PartitionPosition position = PartitionPosition.done();
-
- assertTrue(tracker.tryClaim(position));
-
- final Progress progress = tracker.getProgress();
-
- assertEquals(to.getSeconds(), progress.getWorkCompleted(), DELTA);
- assertEquals(1, progress.getWorkRemaining(), DELTA);
- }
-
- @Property
- public void testGetProgressStop(
- @From(TimestampGenerator.class) Timestamp from,
- @From(TimestampGenerator.class) Timestamp to) {
- assumeThat(from, lessThan(to));
- assumeThat(from.getSeconds(), greaterThanOrEqualTo(1L));
-
- final PartitionRestriction restriction =
- PartitionRestriction.waitForChildPartitions(from, to)
- .withMetadata(
- PartitionRestrictionMetadata.newBuilder()
- .withPartitionToken(PARTITION_TOKEN)
- .build());
-
- final PartitionRestriction stoppedRestriction = PartitionRestriction.stop(restriction);
- final PartitionRestrictionTracker tracker = new PartitionRestrictionTracker(stoppedRestriction);
-
- final Progress progress = tracker.getProgress();
-
- assertEquals(to.getSeconds(), progress.getWorkCompleted(), DELTA);
- assertEquals(1, progress.getWorkRemaining(), DELTA);
- }
-}
diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/util/PartitionPositionGenerator.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/util/PartitionPositionGenerator.java
deleted file mode 100644
index 005962d4972..00000000000
--- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/util/PartitionPositionGenerator.java
+++ /dev/null
@@ -1,54 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements. See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership. The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * "License"); you may not use this file except in compliance
- * with the License. You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-package org.apache.beam.sdk.io.gcp.spanner.changestreams.util;
-
-import static org.apache.beam.sdk.io.gcp.spanner.changestreams.restriction.PartitionMode.QUERY_CHANGE_STREAM;
-
-import com.google.cloud.Timestamp;
-import com.pholser.junit.quickcheck.generator.GenerationStatus;
-import com.pholser.junit.quickcheck.generator.Generator;
-import com.pholser.junit.quickcheck.random.SourceOfRandomness;
-import java.util.Optional;
-import java.util.concurrent.TimeUnit;
-import org.apache.beam.sdk.io.gcp.spanner.changestreams.restriction.PartitionMode;
-import org.apache.beam.sdk.io.gcp.spanner.changestreams.restriction.PartitionPosition;
-
-public class PartitionPositionGenerator extends Generator<PartitionPosition> {
- private static final long MIN_SECONDS = Timestamp.MIN_VALUE.getSeconds();
- private static final long MAX_SECONDS = Timestamp.MAX_VALUE.getSeconds();
- private static final int NANOS_PER_SECOND = (int) TimeUnit.SECONDS.toNanos(1);
- private static final int TOTAL_TRANSITIONS = 3;
-
- public PartitionPositionGenerator() {
- super(PartitionPosition.class);
- }
-
- @Override
- public PartitionPosition generate(SourceOfRandomness random, GenerationStatus status) {
- final PartitionMode mode = PartitionMode.values()[random.nextInt(0, TOTAL_TRANSITIONS + 1)];
-
- if (mode == QUERY_CHANGE_STREAM) {
- final long seconds = random.nextLong(MIN_SECONDS, MAX_SECONDS);
- final int nanos = random.nextInt(0, NANOS_PER_SECOND);
-
- Timestamp timestamp = Timestamp.ofTimeSecondsAndNanos(seconds, nanos);
- return new PartitionPosition(Optional.of(timestamp), mode);
- }
- return new PartitionPosition(Optional.empty(), mode);
- }
-}