[DebeziumIO] Add opt-in full metadata (op and Debezium source fields) to readAsJson

.github/trigger_files/beam_PostCommit_Python.json

@@ -1,5 +1,5 @@
 {
   "comment": "Modify this file in a trivial way to cause this test suite to run.",
   "pr": "40241",
-  "modification": 65
+  "modification": 66
 }

.github/trigger_files/beam_PostCommit_XVR_Direct.json

@@ -1,4 +1,4 @@
 {
-  "modification": 2,
+  "modification": 3,
   "https://github.com/apache/beam/pull/39990": "removing dead code from FnApiDoFnRunner"
 }

sdks/go/pkg/beam/io/xlang/debeziumio/debezium.go

@@ -79,6 +79,7 @@ type readFromDebeziumSchema struct {
 	ConnectionProperties []string
 	StartOffset          []string
 	OffsetStoragePath    *string
+	FullMetadata         *bool
 }
 
 type debeziumConfig struct {
@@ -189,6 +190,22 @@ func OffsetStoragePath(path string) readOption {
 	}
 }
 
+// FullMetadata adds op and all Debezium source fields not already present
+// (e.g. MySQL file, pos, row, ts_ms; PostgreSQL lsn, txId, ts_ms) to the
+// metadata of each JSON record. By default, only connector, version, name,
+// database, schema and table are included.
+//
+// Example:
+//
+//	debeziumio.Read(s, username, password, host, port, debeziumio.MySQL,
+//	                reflectx.String, debeziumio.FullMetadata())
+func FullMetadata() readOption {
+	return func(cfg *debeziumConfig) {
+		b := true
+		cfg.readSchema.FullMetadata = &b
+	}
+}
+
 // ExpansionAddr sets the expansion service address to use for DebeziumIO cross-langauage transform.
 func ExpansionAddr(expansionAddr string) readOption {
 	return func(cfg *debeziumConfig) {

sdks/go/pkg/beam/io/xlang/debeziumio/debezium_test.go

@@ -0,0 +1,29 @@
+// 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 debeziumio
+
+import "testing"
+
+func TestFullMetadata(t *testing.T) {
+	dc := debeziumConfig{readSchema: &readFromDebeziumSchema{}}
+	if dc.readSchema.FullMetadata != nil {
+		t.Fatalf("FullMetadata = %v by default, want nil", *dc.readSchema.FullMetadata)
+	}
+	FullMetadata()(&dc)
+	if got := dc.readSchema.FullMetadata; got == nil || !*got {
+		t.Errorf("FullMetadata() set FullMetadata = %v, want true", got)
+	}
+}

sdks/java/io/debezium/src/README.md

@@ -93,6 +93,8 @@ If you want to use this built-in JSON Mapper, you can do it by setting an instan
 ```
 > **Note:** `SourceRecordJsonMapper`comes out of the box, but you may use any Format Function you prefer.
 
+By default, the JSON `metadata` object holds `connector`, `version`, `name`, `database`, `schema` and `table`. Call `withFullMetadata(true)` on `readAsJson()` to also add the change event's `op` (`c`, `u`, `d`, `r` or `t`) and every other field of the Debezium `source` block to `metadata` (for example `ts_ms`, MySQL `file`, `pos` and `row`, or PostgreSQL `lsn` and `txId`). Fields that `metadata` already carries aren't repeated. For MySQL, `file` is the actual binlog file name, which the default output exposes under `schema`. The option is off by default, so the output is unchanged unless you enable it.
+
 ## Quick Example
 
 The following example is how an actual setup would look like using a **MySQL Debezium Connector** and **SourceRecordJsonMapper** as the Serializable Function.

sdks/java/io/debezium/src/main/java/org/apache/beam/io/debezium/DebeziumIO.java

@@ -147,6 +147,8 @@
 
     abstract @Nullable OffsetRetainer getOffsetRetainer();
 
+    abstract @Nullable Boolean getFullMetadata();
+
     abstract @Nullable Coder<T> getCoder();
 
     abstract Builder<T> toBuilder();
@@ -169,6 +171,8 @@
 
       abstract Builder<T> setOffsetRetainer(OffsetRetainer retainer);
 
+      abstract Builder<T> setFullMetadata(Boolean fullMetadata);
+
       abstract Read<T> build();
     }
 
@@ -305,6 +309,44 @@
       return toBuilder().setOffsetRetainer(retainer).build();
     }
 
+    /**
+     * Adds {@code op} and the remaining Debezium {@code source} fields (e.g. MySQL {@code file},
+     * {@code pos}, {@code row}, {@code ts_ms}; PostgreSQL {@code lsn}, {@code txId}, {@code ts_ms})
+     * to {@code metadata} when reading as JSON.
+     *
+     * <p>Only applies to {@link DebeziumIO#readAsJson()} / {@link
+     * SourceRecordJson.SourceRecordJsonMapper}; combining it with a custom {@link
+     * #withFormatFunction(SourceRecordMapper)} fails with an {@link IllegalArgumentException} when
+     * the transform is expanded. Default false.
+     *
+     * @param fullMetadata Whether to emit the full Debezium metadata.
+     * @return PTransform {@link #read}
+     */
+    public Read<T> withFullMetadata(boolean fullMetadata) {
+      return toBuilder().setFullMetadata(fullMetadata).build();
+    }
+
+    /**
+     * Returns the format function to use, taking {@link #withFullMetadata(boolean)} into account.
+     *
+     * @throws IllegalArgumentException if full metadata is requested but the format function is not
+     *     a {@link SourceRecordJson.SourceRecordJsonMapper}.
+     */
+    SourceRecordMapper<T> resolveFormatFunction() {
+      SourceRecordMapper<T> formatFunction = getFormatFunction();
+      if (!Boolean.TRUE.equals(getFullMetadata())) {
+        return formatFunction;
+      }
+      if (formatFunction instanceof SourceRecordJson.SourceRecordJsonMapper) {
+        @SuppressWarnings("unchecked")
+        SourceRecordMapper<T> fullMetadataMapper =
+            (SourceRecordMapper<T>) new SourceRecordJson.SourceRecordJsonMapper(true);
+        return fullMetadataMapper;
+      }
+      throw new IllegalArgumentException(
+          "withFullMetadata(true) only applies to readAsJson() / SourceRecordJsonMapper");
+    }
+
     protected Schema getRecordSchema() {
       KafkaSourceConsumerFn<T> fn =
           new KafkaSourceConsumerFn<>(getConnectorConfiguration().getConnectorClass().get(), this);
@@ -410,9 +452,9 @@
      * <p>Currently supported connectors are:
      *
      * <ul>
      *   <li>{@link io.debezium.connector.mysql.MySqlConnector}
      *   <li>{@link io.debezium.connector.postgresql.PostgresConnector}
      *   <li>{@link io.debezium.connector.sqlserver.SqlServerConnector }
      * </ul>
      *
      * @param connectorClass Any of the supported connectors.
@@ -429,9 +471,9 @@
      * <p>Currently supported connectors are:
      *
      * <ul>
      *   <li>{@link io.debezium.connector.mysql.MySqlConnector}
      *   <li>{@link io.debezium.connector.postgresql.PostgresConnector}
      *   <li>{@link io.debezium.connector.sqlserver.SqlServerConnector }
      * </ul>
      *
      * @param connectorClass (as ValueProvider)

sdks/java/io/debezium/src/main/java/org/apache/beam/io/debezium/DebeziumTransformRegistrar.java

@@ -81,6 +81,7 @@ public static class Configuration extends CrossLanguageConfiguration {
       private @Nullable Long maxTimeToRun;
       private @Nullable List<String> startOffset;
       private @Nullable String offsetStoragePath;
+      private @Nullable Boolean fullMetadata;
 
       public void setConnectionProperties(@Nullable List<String> connectionProperties) {
         this.connectionProperties = connectionProperties;
@@ -101,6 +102,10 @@ public void setStartOffset(@Nullable List<String> startOffset) {
       public void setOffsetStoragePath(@Nullable String offsetStoragePath) {
         this.offsetStoragePath = offsetStoragePath;
       }
+
+      public void setFullMetadata(@Nullable Boolean fullMetadata) {
+        this.fullMetadata = fullMetadata;
+      }
     }
 
     @Override
@@ -161,6 +166,10 @@ public PTransform<PBegin, PCollection<String>> buildExternal(Configuration confi
                 FileSystemOffsetRetainer.of(configuration.offsetStoragePath));
       }
 
+      if (Boolean.TRUE.equals(configuration.fullMetadata)) {
+        readTransform = readTransform.withFullMetadata(true);
+      }
+
       return readTransform;
     }
   }

sdks/java/io/debezium/src/main/java/org/apache/beam/io/debezium/KafkaSourceConsumerFn.java

@@ -113,7 +113,7 @@ public class KafkaSourceConsumerFn<T> extends DoFn<Map<String, String>, T> {
     // this(connectorClass, fn, maxRecords, null);
     this.connectorClass = (Class<? extends SourceConnector>) connectorClass;
     this.spec = spec;
-    this.fn = spec.getFormatFunction();
+    this.fn = spec.resolveFormatFunction();
     this.pollingTimeOut =
         MoreObjects.firstNonNull(spec.getPollingTimeout(), DEFAULT_POLLING_TIMEOUT);
   }

sdks/java/io/debezium/src/main/java/org/apache/beam/io/debezium/SourceRecordJson.java

@@ -19,10 +19,15 @@
 
 import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
 import java.io.Serializable;
+import java.util.Collections;
 import java.util.HashMap;
+import java.util.LinkedHashMap;
 import java.util.Map;
+import java.util.Set;
 import org.apache.beam.vendor.grpc.v1p69p0.com.google.gson.Gson;
 import org.apache.beam.vendor.grpc.v1p69p0.com.google.gson.GsonBuilder;
+import org.apache.beam.vendor.grpc.v1p69p0.com.google.gson.JsonObject;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableSet;
 import org.apache.kafka.connect.data.Field;
 import org.apache.kafka.connect.data.Struct;
 import org.apache.kafka.connect.errors.DataException;
@@ -34,7 +39,7 @@
  *
  * <h3>What it does</h3>
  *
  * <p>It maps any SourceRecord retrieved from any supported {@link io.debezium.connector} to JSON
  *
  * <h3>How it works</h3>
  *
@@ -45,12 +50,32 @@
  *   <li>after - {@link #loadAfter}
  *   <li>metadata - {@link #loadMetadata}
  *       <ul>
- *         <li>schema - Database Schema
  *         <li>connector - Connector used
  *         <li>version - Connector version
+ *         <li>name - Connector name
+ *         <li>database - Database name
+ *         <li>schema - Database Schema (for MySQL, the binlog file name)
+ *         <li>table - Table name
  *       </ul>
  * </ul>
  *
+ * <h3>Full metadata</h3>
+ *
+ * <p>When constructed with {@code fullMetadata = true} (see {@link
+ * SourceRecordJsonMapper#SourceRecordJsonMapper(boolean)} and {@link
+ * DebeziumIO.Read#withFullMetadata(boolean)}), extra keys are appended to {@code metadata} after
+ * {@code table}:
+ *
+ * <ul>
+ *   <li>op - The Debezium operation ({@code c}, {@code u}, {@code d}, {@code r}, {@code t}), if the
+ *       record envelope has an {@code op} field
+ *   <li>Every other field of the Debezium {@code source} block, in schema order, that is not
+ *       already represented above (for example {@code ts_ms}, {@code snapshot}, MySQL {@code file},
+ *       {@code pos}, {@code row}, or PostgreSQL {@code lsn}, {@code txId})
+ * </ul>
+ *
+ * <p>By default ({@code fullMetadata = false}) the output only contains the six keys above.
+ *
  * <h3>Usage Example</h3>
  *
  * <p>Map each SourceRecord to JSON
@@ -68,15 +93,34 @@
             + " Making the class final or moving the throwing code"
             + " into a static factory method may fix it.")
 public class SourceRecordJson {
+  /**
+   * Source fields that are already represented in {@link Event.Metadata} (or would collide with an
+   * existing metadata key), so they are never added as full metadata extras.
+   */
+  private static final Set<String> RESERVED_METADATA_KEYS =
+      ImmutableSet.of("connector", "version", "name", "db", "database", "schema", "table", "op");
+
   private final @Nullable Struct value;
   private final @Nullable Event event;
+  private final Map<String, Object> metadataExtras;
 
   /**
    * Initializer.
    *
    * @param sourceRecord retrieved SourceRecord using a supported SourceConnector
    */
   public SourceRecordJson(@Nullable SourceRecord sourceRecord) {
+    this(sourceRecord, false);
+  }
+
+  /**
+   * Initializer.
+   *
+   * @param sourceRecord retrieved SourceRecord using a supported SourceConnector
+   * @param fullMetadata whether to add {@code op} and the remaining Debezium {@code source} fields
+   *     to {@code metadata}
+   */
+  public SourceRecordJson(@Nullable SourceRecord sourceRecord, boolean fullMetadata) {
     if (sourceRecord == null) {
       throw new IllegalArgumentException();
     }
@@ -85,12 +129,15 @@
 
     if (this.value == null) {
       this.event = new Event(null, null, null);
+      this.metadataExtras = Collections.emptyMap();
     } else {
       Event.Metadata metadata = this.loadMetadata();
       Event.Before before = this.loadBefore();
       Event.After after = this.loadAfter();
 
       this.event = new Event(metadata, before, after);
+      this.metadataExtras =
+          fullMetadata && metadata != null ? this.loadMetadataExtras() : Collections.emptyMap();
     }
   }
 
@@ -129,6 +176,33 @@
         source.getString("table"));
   }
 
+  /**
+   * Extracts the full metadata extras: the envelope's {@code op} (if the value schema has one),
+   * then every field of the Debezium {@code source} block, in schema order, that isn't already
+   * represented in {@link Event.Metadata}.
+   *
+   * @return extra metadata keys, in insertion order
+   */
+  private Map<String, Object> loadMetadataExtras() {
+    Struct source = (Struct) this.value.get("source");
+    Map<String, Object> extras = new LinkedHashMap<>();
+    if (this.value.schema().field("op") != null) {
+      extras.put("op", toJsonFriendly(this.value.get("op")));
+    }
+    for (Field field : source.schema().fields()) {
+      String fieldName = field.name();
+      if (!RESERVED_METADATA_KEYS.contains(fieldName)) {
+        extras.put(fieldName, toJsonFriendly(source.get(field)));
+      }
+    }
+    return extras;
+  }
+
+  private static @Nullable Object toJsonFriendly(@Nullable Object value) {
+    // Source fields are flat primitives today; stringify a nested Struct to keep metadata flat.
+    return value instanceof Struct ? String.valueOf(value) : value;
+  }
+
   /**
    * Extracts the before field within SourceRecord.
    *
@@ -183,14 +257,31 @@
    * @return JSON String
    */
   public String toJson() {
-    return this.event.toJson();
+    return this.event.toJson(this.metadataExtras);
   }
 
   /** {@link SourceRecordJson} implementation. */
   public static class SourceRecordJsonMapper implements SourceRecordMapper<String> {
+    private final boolean fullMetadata;
+
+    /** Creates a mapper that emits the default metadata. */
+    public SourceRecordJsonMapper() {
+      this(false);
+    }
+
+    /**
+     * Creates a mapper.
+     *
+     * @param fullMetadata whether to add {@code op} and the remaining Debezium {@code source}
+     *     fields to {@code metadata}
+     */
+    public SourceRecordJsonMapper(boolean fullMetadata) {
+      this.fullMetadata = fullMetadata;
+    }
+
     @Override
     public String mapSourceRecord(SourceRecord sourceRecord) throws Exception {
-      return new SourceRecordJson(sourceRecord).toJson();
+      return new SourceRecordJson(sourceRecord, fullMetadata).toJson();
     }
   }
 
@@ -223,8 +314,28 @@
      * @return JSON String
      */
     public String toJson() {
+      return toJson(Collections.emptyMap());
+    }
+
+    /**
+     * Transforms the Event to a JSON string, appending the given extra keys to {@code metadata}
+     * after {@code table}.
+     *
+     * @param extras Additional metadata keys, in iteration order
+     * @return JSON String
+     */
+    String toJson(Map<String, Object> extras) {
       Gson gson = new GsonBuilder().serializeNulls().create();
-      return gson.toJson(this);
+      if (extras.isEmpty() || metadata == null) {
+        return gson.toJson(this);
+      }
+      // Gson can't flatten a map into an object, so append the extras to the metadata tree.
+      JsonObject tree = gson.toJsonTree(this).getAsJsonObject();
+      JsonObject metadataJson = tree.getAsJsonObject("metadata");
+      for (Map.Entry<String, Object> extra : extras.entrySet()) {
+        metadataJson.add(extra.getKey(), gson.toJsonTree(extra.getValue()));
+      }
+      return gson.toJson(tree);
     }
 
     /** Depicts the metadata within a SourceRecord. It has valuable fields. */