[pull] master from apache:master

.github/workflows/beam_Metrics_Report.yml

@@ -17,10 +17,11 @@
 
 name: Beam Metrics Report
 
+# DISABLED: Workflow disabled
 on:
-  schedule:
-    - cron: '0 11 * * 2'
-  workflow_dispatch:
+  # schedule:
+  #   - cron: '0 11 * * 2'
+  # workflow_dispatch:
 
 # This allows a subsequently queued workflow run to interrupt previous runs
 concurrency:

model/pipeline/src/main/proto/org/apache/beam/model/pipeline/v1/beam_runner_api.proto

@@ -1073,6 +1073,10 @@ message StandardCoders {
     // Components: None
     ROW = 13 [(beam_urn) = "beam:coder:row:v1"];
 
+    // Similar to ROW above, but for arbitrary types that can be converted
+    // to and from row objects, which can then be encoded with a schema.
+    SCHEMA = 18 [(beam_urn) = "beam:coder:schema:v1"];
+
     // Encodes a user key and a shard id which is an opaque byte string.
     //
     // The encoding for a sharded key consists of a shard id byte string and the

model/pipeline/src/main/proto/org/apache/beam/model/pipeline/v1/schema.proto

@@ -48,7 +48,7 @@ message Field {
   // OPTIONAL. Human readable description of this field, such as the query that generated it.
   string description = 2;
   FieldType type = 3;
-  
+
   int32 id = 4;
    // OPTIONAL. The position of this field's data when encoded, e.g. with beam:coder:row:v1.
    // Either no fields in a given row are have encoding position populated,
@@ -237,3 +237,23 @@ message MapTypeEntry {
 message LogicalTypeValue {
   FieldValue value = 1;
 }
+
+// Information needed to represent a Coder of type SCHEMA.
+message SchemaCoderPayload {
+  // The schema to use for encoding corresponding Row types.
+  Schema schema = 1;
+
+  // Function mapping from underlying object to Row type.
+  FunctionSpec to_row_fn = 2;
+
+  // Function mapping from Row type to underlying object.
+  FunctionSpec from_row_fn = 3;
+
+  // Any additional information SDKs need to encode/decode elements.
+  repeated AdditionalCoderInfo additional_coder_infos = 4;
+
+  message AdditionalCoderInfo {
+    string urn = 1;
+    bytes payload = 2;
+  }
+}

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

@@ -46,6 +46,8 @@
 import com.google.cloud.spanner.Options;
 import com.google.cloud.spanner.Options.RpcPriority;
 import com.google.cloud.spanner.PartitionOptions;
+import com.google.cloud.spanner.ReadOnlyTransaction;
+import com.google.cloud.spanner.ResultSet;
 import com.google.cloud.spanner.Spanner;
 import com.google.cloud.spanner.SpannerException;
 import com.google.cloud.spanner.SpannerOptions;
@@ -1998,6 +2000,11 @@ && getInclusiveStartAt().toSqlTimestamp().after(getInclusiveEndAt().toSqlTimesta
       final MapperFactory mapperFactory = new MapperFactory(changeStreamDatabaseDialect);
       final ChangeStreamMetrics metrics = new ChangeStreamMetrics();
       final RpcPriority rpcPriority = MoreObjects.firstNonNull(getRpcPriority(), RpcPriority.HIGH);
+      final SpannerAccessor spannerAccessor =
+          SpannerAccessor.getOrCreate(changeStreamSpannerConfig);
+      final boolean isMutableChangeStream =
+          isMutableChangeStream(
+              spannerAccessor.getDatabaseClient(), changeStreamDatabaseDialect, changeStreamName);
       final DaoFactory daoFactory =
           new DaoFactory(
               changeStreamSpannerConfig,
@@ -2007,7 +2014,8 @@ && getInclusiveStartAt().toSqlTimestamp().after(getInclusiveEndAt().toSqlTimesta
               rpcPriority,
               input.getPipeline().getOptions().getJobName(),
               changeStreamDatabaseDialect,
-              metadataDatabaseDialect);
+              metadataDatabaseDialect,
+              isMutableChangeStream);
       final ActionFactory actionFactory = new ActionFactory();
 
       final Duration watermarkRefreshRate =
@@ -2689,4 +2697,58 @@ static String resolveSpannerProjectId(SpannerConfig config) {
         ? SpannerOptions.getDefaultProjectId()
         : config.getProjectId().get();
   }
+
+  @VisibleForTesting
+  static boolean isMutableChangeStream(
+      DatabaseClient databaseClient, Dialect dialect, String changeStreamName) {
+    String fetchedPartitionMode = fetchPartitionMode(databaseClient, dialect, changeStreamName);
+    if (fetchedPartitionMode.isEmpty()
+        || fetchedPartitionMode.equalsIgnoreCase("IMMUTABLE_KEY_RANGE")) {
+      return false;
+    }
+    return true;
+  }
+
+  private static String fetchPartitionMode(
+      DatabaseClient databaseClient, Dialect dialect, String changeStreamName) {
+    try (ReadOnlyTransaction tx = databaseClient.readOnlyTransaction()) {
+      Statement statement;
+      if (dialect == Dialect.POSTGRESQL) {
+        statement =
+            Statement.newBuilder(
+                    "select option_value\n"
+                        + "from information_schema.change_stream_options\n"
+                        + "where change_stream_name = $1 and option_name = 'partition_mode'")
+                .bind("p1")
+                .to(changeStreamName)
+                .build();
+      } else {
+        statement =
+            Statement.newBuilder(
+                    "select option_value\n"
+                        + "from information_schema.change_stream_options\n"
+                        + "where change_stream_name = @changeStreamName and  option_name = 'partition_mode'")
+                .bind("changeStreamName")
+                .to(changeStreamName)
+                .build();
+      }
+      ResultSet resultSet = tx.executeQuery(statement);
+      while (resultSet.next()) {
+        String value = resultSet.getString(0);
+        if (value != null) {
+          return value;
+        }
+      }
+      return "";
+    } catch (RuntimeException e) {
+      // Log the failure (with stack trace) but rethrow so the caller still observes
+      // the error.
+      LOG.warn(
+          "Failed to fetch partition_mode for change stream '{}', dialect={} - will propagate exception",
+          changeStreamName,
+          dialect,
+          e);
+      throw e;
+    }
+  }
 }

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

@@ -31,12 +31,12 @@
  * as a {@link ResultSet}, which can be consumed until the stream is finished.
  */
 public class ChangeStreamDao {
-
   private final String changeStreamName;
   private final DatabaseClient databaseClient;
   private final RpcPriority rpcPriority;
   private final String jobName;
   private final Dialect dialect;
+  private final boolean isMutableChangeStream;
 
   /**
    * Constructs a change stream dao. All the queries performed by this class will be for the given
@@ -53,12 +53,14 @@ public class ChangeStreamDao {
       DatabaseClient databaseClient,
       RpcPriority rpcPriority,
       String jobName,
-      Dialect dialect) {
+      Dialect dialect,
+      boolean isMutableChangeStream) {
     this.changeStreamName = changeStreamName;
     this.databaseClient = databaseClient;
     this.rpcPriority = rpcPriority;
     this.jobName = jobName;
     this.dialect = dialect;
+    this.isMutableChangeStream = isMutableChangeStream;
   }
 
   /**
@@ -91,8 +93,18 @@ public ChangeStreamResultSet changeStreamQuery(
     String query = "";
     Statement statement;
     if (this.isPostgres()) {
-      query =
-          "SELECT * FROM \"spanner\".\"read_json_" + changeStreamName + "\"($1, $2, $3, $4, null)";
+      // Ensure we have determined whether change stream uses mutable key range
+      if (this.isMutableChangeStream) {
+        query =
+            "SELECT * FROM \"spanner\".\"read_proto_bytes_"
+                + changeStreamName
+                + "\"($1, $2, $3, $4, null)";
+      } else {
+        query =
+            "SELECT * FROM \"spanner\".\"read_json_"
+                + changeStreamName
+                + "\"($1, $2, $3, $4, null)";
+      }
       statement =
           Statement.newBuilder(query)
               .bind("p1")

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

@@ -49,6 +49,7 @@ public class DaoFactory implements Serializable {
   private final String jobName;
   private final Dialect spannerChangeStreamDatabaseDialect;
   private final Dialect metadataDatabaseDialect;
+  private final boolean isMutableChangeStream;
 
   /**
    * Constructs a {@link DaoFactory} with the configuration to be used for the underlying instances.
@@ -68,7 +69,8 @@ public DaoFactory(
       RpcPriority rpcPriority,
       String jobName,
       Dialect spannerChangeStreamDatabaseDialect,
-      Dialect metadataDatabaseDialect) {
+      Dialect metadataDatabaseDialect,
+      boolean isMutableChangeStream) {
     if (metadataSpannerConfig.getInstanceId() == null) {
       throw new IllegalArgumentException("Metadata instance can not be null");
     }
@@ -83,6 +85,7 @@ public DaoFactory(
     this.jobName = jobName;
     this.spannerChangeStreamDatabaseDialect = spannerChangeStreamDatabaseDialect;
     this.metadataDatabaseDialect = metadataDatabaseDialect;
+    this.isMutableChangeStream = isMutableChangeStream;
   }
 
   /**
@@ -143,7 +146,8 @@ public synchronized ChangeStreamDao getChangeStreamDao() {
               spannerAccessor.getDatabaseClient(),
               rpcPriority,
               jobName,
-              this.spannerChangeStreamDatabaseDialect);
+              this.spannerChangeStreamDatabaseDialect,
+              this.isMutableChangeStream);
     }
     return changeStreamDaoInstance;
   }