[pull] master from apache:master

sdks/java/io/components/src/main/java/org/apache/beam/sdk/io/components/ratelimiter/RateLimiter.java

@@ -0,0 +1,41 @@
+/*
+ * 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.components.ratelimiter;
+
+import java.io.IOException;
+import java.io.Serializable;
+
+/**
+ * A RateLimiter allows to fetch permits from a rate limiter service and blocks execution when the
+ * rate limit is exceeded.
+ *
+ * <p>Implementations must be {@link Serializable} as they are passed to workers.
+ */
+public interface RateLimiter extends Serializable, AutoCloseable {
+
+  /**
+   * Blocks until the specified number of permits are acquired and returns true if the request was
+   * allowed or false if the request was rejected.
+   *
+   * @param permits Number of permits to acquire.
+   * @return true if the request was allowed, false if it was rejected (and retries exceeded).
+   * @throws IOException if there is an error communicating with the rate limiter service.
+   * @throws InterruptedException if the thread is interrupted while waiting.
+   */
+  boolean allow(int permits) throws IOException, InterruptedException;
+}

sdks/java/io/components/src/main/java/org/apache/beam/sdk/io/components/ratelimiter/RateLimiterContext.java

@@ -0,0 +1,27 @@
+/*
+ * 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.components.ratelimiter;
+
+import java.io.Serializable;
+
+/**
+ * A marker interface for context data required to check ratelimit.
+ *
+ * <p>Implementations must be {@link Serializable}.
+ */
+public interface RateLimiterContext extends Serializable {}

sdks/java/io/components/src/main/java/org/apache/beam/sdk/io/components/ratelimiter/RateLimiterFactory.java

@@ -0,0 +1,57 @@
+/*
+ * 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.components.ratelimiter;
+
+import java.io.IOException;
+import java.io.Serializable;
+
+/**
+ * A factory that manages connections to rate limit service and creates lightweight handles.
+ *
+ * <p>Implementations must be {@link Serializable} as they are passed to workers. The factory
+ * typically manages the heavy connection (e.g. gRPC stub) and is thread-safe.
+ */
+public interface RateLimiterFactory extends Serializable, AutoCloseable {
+
+  /**
+   * Creates a lightweight ratelimiter handle bound to a specific context.
+   *
+   * <p>Use this when passing ratelimiter to IO components, which doesn't need to know about the
+   * configuration or the underlying ratelimiter service details. This is also useful in DoFns when
+   * you want to use the ratelimiter in a static way based on the compile time context.
+   *
+   * @param context The context for the ratelimit.
+   * @return A {@link RateLimiter} handle.
+   */
+  RateLimiter getLimiter(RateLimiterContext context);
+
+  /**
+   * Blocks until the specified number of permits are acquired and returns true if the request was
+   * allowed or false if the request was rejected.
+   *
+   * <p>Use this for when the ratelimit namespace or descriptors are not known at compile time.
+   * allows you to use the ratelimiter in a dynamic way based on the runtime data.
+   *
+   * @param context The context for the ratelimit.
+   * @param permits Number of permits to acquire.
+   * @return true if the request is allowed, false if rejected.
+   * @throws IOException if there is an error communicating with the ratelimiter service.
+   * @throws InterruptedException if the thread is interrupted while waiting.
+   */
+  boolean allow(RateLimiterContext context, int permits) throws IOException, InterruptedException;
+}

sdks/java/io/components/src/main/java/org/apache/beam/sdk/io/components/ratelimiter/package-info.java

@@ -0,0 +1,20 @@
+/*
+ * 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.
+ */
+
+/** Defines ratelimiter utilities for Beam DoFn and IO components. */
+package org.apache.beam.sdk.io.components.ratelimiter;

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

@@ -1977,11 +1977,6 @@ && getInclusiveStartAt().toSqlTimestamp().after(getInclusiveEndAt().toSqlTimesta
               + changeStreamDatabaseId
               + " has dialect "
               + changeStreamDatabaseDialect);
-      LOG.info(
-          "The Spanner database "
-              + fullPartitionMetadataDatabaseId
-              + " has dialect "
-              + metadataDatabaseDialect);
       PartitionMetadataTableNames partitionMetadataTableNames =
           Optional.ofNullable(getMetadataTable())
               .map(
@@ -2005,6 +2000,7 @@ && getInclusiveStartAt().toSqlTimestamp().after(getInclusiveEndAt().toSqlTimesta
       final boolean isMutableChangeStream =
           isMutableChangeStream(
               spannerAccessor.getDatabaseClient(), changeStreamDatabaseDialect, changeStreamName);
+      LOG.info("The change stream " + changeStreamName + " is mutable: " + isMutableChangeStream);
       final DaoFactory daoFactory =
           new DaoFactory(
               changeStreamSpannerConfig,

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

@@ -187,7 +187,8 @@ public synchronized QueryChangeStreamAction queryChangeStreamAction(
       PartitionStartRecordAction partitionStartRecordAction,
       PartitionEndRecordAction partitionEndRecordAction,
       PartitionEventRecordAction partitionEventRecordAction,
-      ChangeStreamMetrics metrics) {
+      ChangeStreamMetrics metrics,
+      boolean isMutableChangeStream) {
     if (queryChangeStreamActionInstance == null) {
       queryChangeStreamActionInstance =
           new QueryChangeStreamAction(
@@ -201,7 +202,8 @@ public synchronized QueryChangeStreamAction queryChangeStreamAction(
               partitionStartRecordAction,
               partitionEndRecordAction,
               partitionEventRecordAction,
-              metrics);
+              metrics,
+              isMutableChangeStream);
     }
     return queryChangeStreamActionInstance;
   }

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

@@ -34,6 +34,7 @@
 import org.apache.beam.sdk.io.gcp.spanner.changestreams.model.ChildPartitionsRecord;
 import org.apache.beam.sdk.io.gcp.spanner.changestreams.model.DataChangeRecord;
 import org.apache.beam.sdk.io.gcp.spanner.changestreams.model.HeartbeatRecord;
+import org.apache.beam.sdk.io.gcp.spanner.changestreams.model.InitialPartition;
 import org.apache.beam.sdk.io.gcp.spanner.changestreams.model.PartitionEndRecord;
 import org.apache.beam.sdk.io.gcp.spanner.changestreams.model.PartitionEventRecord;
 import org.apache.beam.sdk.io.gcp.spanner.changestreams.model.PartitionMetadata;
@@ -89,6 +90,7 @@ public class QueryChangeStreamAction {
   private final PartitionEndRecordAction partitionEndRecordAction;
   private final PartitionEventRecordAction partitionEventRecordAction;
   private final ChangeStreamMetrics metrics;
+  private final boolean isMutableChangeStream;
 
   /**
    * Constructs an action class for performing a change stream query for a given partition.
@@ -106,6 +108,7 @@ public class QueryChangeStreamAction {
    * @param PartitionEndRecordAction action class to process {@link PartitionEndRecord}s
    * @param PartitionEventRecordAction action class to process {@link PartitionEventRecord}s
    * @param metrics metrics gathering class
+   * @param isMutableChangeStream whether the change stream is mutable or not
    */
   QueryChangeStreamAction(
       ChangeStreamDao changeStreamDao,
@@ -118,7 +121,8 @@ public class QueryChangeStreamAction {
       PartitionStartRecordAction partitionStartRecordAction,
       PartitionEndRecordAction partitionEndRecordAction,
       PartitionEventRecordAction partitionEventRecordAction,
-      ChangeStreamMetrics metrics) {
+      ChangeStreamMetrics metrics,
+      boolean isMutableChangeStream) {
     this.changeStreamDao = changeStreamDao;
     this.partitionMetadataDao = partitionMetadataDao;
     this.changeStreamRecordMapper = changeStreamRecordMapper;
@@ -130,6 +134,7 @@ public class QueryChangeStreamAction {
     this.partitionEndRecordAction = partitionEndRecordAction;
     this.partitionEventRecordAction = partitionEventRecordAction;
     this.metrics = metrics;
+    this.isMutableChangeStream = isMutableChangeStream;
   }
 
   /**
@@ -195,13 +200,23 @@ public ProcessContinuation run(
     final Timestamp endTimestamp = partition.getEndTimestamp();
     final boolean isBoundedRestriction = !endTimestamp.equals(MAX_INCLUSIVE_END_AT);
     final Timestamp changeStreamQueryEndTimestamp =
-        isBoundedRestriction ? endTimestamp : getNextReadChangeStreamEndTimestamp();
+        isBoundedRestriction
+            ? getBoundedQueryEndTimestamp(endTimestamp)
+            : getNextReadChangeStreamEndTimestamp();
 
     // Once the changeStreamQuery completes we may need to resume reading from the partition if we
     // had an unbounded restriction for which we set an arbitrary query end timestamp and for which
     // we didn't  encounter any indications that the partition is done (explicit end records or
-    // exceptions about being out of timestamp range).
-    boolean stopAfterQuerySucceeds = isBoundedRestriction;
+    // exceptions about being out of timestamp range). We also special case the InitialPartition,
+    // which always stops after the query succeeds.
+    boolean stopAfterQuerySucceeds = false;
+    if (InitialPartition.isInitialPartition(partition.getPartitionToken())) {
+      stopAfterQuerySucceeds = true;
+    } else {
+      stopAfterQuerySucceeds =
+          isBoundedRestriction && changeStreamQueryEndTimestamp.equals(endTimestamp);
+    }
+
     try (ChangeStreamResultSet resultSet =
         changeStreamDao.changeStreamQuery(
             token, startTimestamp, changeStreamQueryEndTimestamp, partition.getHeartbeatMillis())) {
@@ -379,4 +394,14 @@ private Timestamp getNextReadChangeStreamEndTimestamp() {
     final Timestamp current = Timestamp.now();
     return Timestamp.ofTimeSecondsAndNanos(current.getSeconds() + 2 * 60, current.getNanos());
   }
+
+  // For Mutable Change Stream bounded queries, update the query end timestamp to be within 2
+  // minutes in the future.
+  private Timestamp getBoundedQueryEndTimestamp(Timestamp endTimestamp) {
+    if (this.isMutableChangeStream) {
+      Timestamp nextTimestamp = getNextReadChangeStreamEndTimestamp();
+      return nextTimestamp.compareTo(endTimestamp) < 0 ? nextTimestamp : endTimestamp;
+    }
+    return endTimestamp;
+  }
 }

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

@@ -20,6 +20,7 @@
 import com.google.cloud.Timestamp;
 import com.google.cloud.spanner.ResultSet;
 import com.google.cloud.spanner.Struct;
+import com.google.protobuf.InvalidProtocolBufferException;
 import org.joda.time.Duration;
 
 /**
@@ -113,6 +114,9 @@ public Struct getCurrentRowAsStruct() {
    * updates the timestamp at which the record was read. This function enhances the getProtoMessage
    * function but only focus on the ChangeStreamRecord type.
    *
+   * <p>Should only be used for GoogleSQL databases when the change stream record is delivered as
+   * proto.
+   *
    * @return a change stream record as a proto or null
    */
   public com.google.spanner.v1.ChangeStreamRecord getProtoChangeStreamRecord() {
@@ -128,6 +132,33 @@ public boolean isProtoChangeRecord() {
         && resultSet.getColumnType(0).getCode() == com.google.cloud.spanner.Type.Code.PROTO;
   }
 
+  /**
+   * Returns the change stream record at the current pointer by parsing the bytes column. It also
+   * updates the timestamp at which the record was read.
+   *
+   * <p>Should only be used for PostgreSQL databases when the change stream record is delivered as
+   * proto bytes.
+   *
+   * @return a change stream record as a proto or null
+   */
+  public com.google.spanner.v1.ChangeStreamRecord getBytes(int index) {
+    recordReadAt = Timestamp.now();
+    try {
+      // Use getBytes(0) for the BYTES column returned by read_proto_bytes_ TVF
+      return com.google.spanner.v1.ChangeStreamRecord.parseFrom(
+          resultSet.getBytes(index).toByteArray());
+    } catch (InvalidProtocolBufferException e) {
+      throw new RuntimeException("Failed to parse the proto bytes to ChangeStreamRecord proto", e);
+    }
+  }
+
+  /** Returns true if the result set at the current pointer contain only one bytes change record. */
+  public boolean isProtoBytesChangeRecord() {
+    return resultSet.getColumnCount() == 1
+        && !resultSet.isNull(0)
+        && resultSet.getColumnType(0).getCode() == com.google.cloud.spanner.Type.Code.BYTES;
+  }
+
   /**
    * Returns the record at the current pointer as {@link JsonB}. It also updates the timestamp at
    * which the record was read.

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

@@ -151,4 +151,8 @@ public synchronized ChangeStreamDao getChangeStreamDao() {
     }
     return changeStreamDaoInstance;
   }
+
+  public boolean isMutableChangeStream() {
+    return this.isMutableChangeStream;
+  }
 }

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

@@ -73,6 +73,7 @@ public class ReadChangeStreamPartitionDoFn extends DoFn<PartitionMetadata, DataC
   private final MapperFactory mapperFactory;
   private final ActionFactory actionFactory;
   private final ChangeStreamMetrics metrics;
+  private final boolean isMutableChangeStream;
   /**
    * Needs to be set through the {@link
    * ReadChangeStreamPartitionDoFn#setThroughputEstimator(BytesThroughputEstimator)} call.
@@ -104,6 +105,7 @@ public ReadChangeStreamPartitionDoFn(
     this.mapperFactory = mapperFactory;
     this.actionFactory = actionFactory;
     this.metrics = metrics;
+    this.isMutableChangeStream = daoFactory.isMutableChangeStream();
     this.throughputEstimator = new NullThroughputEstimator<>();
   }
 
@@ -215,7 +217,8 @@ public void setup() {
             partitionStartRecordAction,
             partitionEndRecordAction,
             partitionEventRecordAction,
-            metrics);
+            metrics,
+            isMutableChangeStream);
   }
 
   /**

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

@@ -218,18 +218,27 @@ public class ChangeStreamRecordMapper {
    * @param resultSet the change stream result set
    * @param resultSetMetadata the metadata generated when reading the change stream row
    * @return a {@link List} of {@link ChangeStreamRecord} subclasses
+   * @throws InvalidProtocolBufferException
    */
   public List<ChangeStreamRecord> toChangeStreamRecords(
       PartitionMetadata partition,
       ChangeStreamResultSet resultSet,
       ChangeStreamResultSetMetadata resultSetMetadata) {
     if (this.isPostgres()) {
-      // In PostgresQL, change stream records are returned as JsonB.
+      // For `MUTABLE_KEY_RANGE` option, change stream records are returned as protos.
+      if (resultSet.isProtoBytesChangeRecord()) {
+        return Arrays.asList(
+            toChangeStreamRecord(partition, resultSet.getBytes(0), resultSetMetadata));
+      }
+
+      // For `IMMUTABLE_KEY_RANGE` option, change stream records are returned as
+      // JsonB.
       return Collections.singletonList(
           toChangeStreamRecordJson(partition, resultSet.getPgJsonb(0), resultSetMetadata));
     }
 
-    // In GoogleSQL, for `MUTABLE_KEY_RANGE` option, change stream records are returned as Protos.
+    // In GoogleSQL, for `MUTABLE_KEY_RANGE` option, change stream records are
+    // returned as Protos.
     if (resultSet.isProtoChangeRecord()) {
       return Arrays.asList(
           toChangeStreamRecord(