[DebeziumIO] Support Transaction Metadata Events in DebeziumIO JSON Output

.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": 66
+  "modification": 67
 }

.github/trigger_files/beam_PostCommit_XVR_Direct.json

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

CHANGES.md

@@ -65,6 +65,7 @@
 ## I/Os
 
 * Support for X source added (Java/Python) ([#X](https://github.com/apache/beam/issues/X)).
+* (Java/Python/Go) DebeziumIO: `readAsJson()` now emits Debezium transaction metadata with `provide.transaction.metadata=true` (Python/Go via `connection_properties`) ([#40784](https://github.com/apache/beam/issues/40784)).
 
 ## New Features / Improvements
 

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

@@ -146,6 +146,8 @@ func MaxTimeToRun(r int64) readOption {
 
 // ConnectionProperties specifies properties of the debezium connection passed as
 // a string with format [propertyName=property;]*
+// Set provide.transaction.metadata=true to add a transaction key to records and
+// emit transaction BEGIN/END records.
 func ConnectionProperties(cp []string) readOption {
 	return func(cfg *debeziumConfig) {
 		cfg.readSchema.ConnectionProperties = cp

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

@@ -95,6 +95,8 @@ If you want to use this built-in JSON Mapper, you can do it by setting an instan
 
 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.
 
+With `provide.transaction.metadata=true`, each change event has a `transaction` object (`id`, `total_order`, `data_collection_order`); otherwise the key is absent. The transaction `BEGIN` and `END` events are mapped to `status`, `id`, `ts_ms`, `event_count` and `data_collections`, where the last two are only set on `END`.
+
 ## Quick Example
 
 The following example is how an actual setup would look like using a **MySQL Debezium Connector** and **SourceRecordJsonMapper** as the Serializable Function.
@@ -172,4 +174,4 @@ You can run Integration Tests using **gradlew**.
 Example of running the MySQL Connector Integration Test:
 ```
 ./gradlew integrationTest -p sdks/java/io/debezium/ --tests org.apache.beam.io.debezium.DebeziumIOMySqlConnectorIT -DintegrationTestRunner=direct
-```
+```

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

@@ -19,16 +19,20 @@
 
 import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
 import java.io.Serializable;
+import java.util.ArrayList;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.LinkedHashMap;
+import java.util.List;
 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.grpc.v1p69p0.com.google.gson.annotations.SerializedName;
 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.Schema;
 import org.apache.kafka.connect.data.Struct;
 import org.apache.kafka.connect.errors.DataException;
 import org.apache.kafka.connect.source.SourceRecord;
@@ -39,7 +43,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>
  *
@@ -57,6 +61,9 @@
  *         <li>schema - Database Schema (for MySQL, the binlog file name)
  *         <li>table - Table name
  *       </ul>
+ *   <li>transaction - {@link #loadTransaction}: the Debezium {@code transaction} block of the
+ *       change event ({@code id}, {@code total_order} and {@code data_collection_order}). The key
+ *       is absent unless the connector has {@code provide.transaction.metadata=true}
  * </ul>
  *
  * <h3>Full metadata</h3>
@@ -76,6 +83,14 @@
  *
  * <p>By default ({@code fullMetadata = false}) the output only contains the six keys above.
  *
+ * <h3>Transaction metadata</h3>
+ *
+ * <p>When the connector has {@code provide.transaction.metadata=true}, Debezium also emits a {@code
+ * BEGIN} and an {@code END} event around the change events of each transaction. These have no
+ * {@code source}, {@code before} or {@code after}, so they are mapped to their own payload instead
+ * ({@link #loadTransactionEvent}): {@code status}, {@code id}, {@code ts_ms}, {@code event_count}
+ * and {@code data_collections}.
+ *
  * <h3>Usage Example</h3>
  *
  * <p>Map each SourceRecord to JSON
@@ -98,10 +113,17 @@
    * 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");
+      ImmutableSet.of(
+          "connector", "version", "name", "db", "database", "schema", "table", "op", "transaction");
+
+  private static final String SCHEMA_TRANSACTION_METADATA_EVENT_KEY_NAME =
+      "io.debezium.connector.common.TransactionMetadataKey";
+
+  private static final Gson GSON = new GsonBuilder().serializeNulls().create();
 
   private final @Nullable Struct value;
   private final @Nullable Event event;
+  private final @Nullable TransactionEvent transactionEvent;
   private final Map<String, Object> metadataExtras;
 
   /**
@@ -128,14 +150,21 @@
     this.value = (Struct) sourceRecord.value();
 
     if (this.value == null) {
-      this.event = new Event(null, null, null);
+      this.event = new Event(null, null, null, null);
+      this.transactionEvent = null;
+      this.metadataExtras = Collections.emptyMap();
+    } else if (isTransactionMetadataEvent(sourceRecord.keySchema())) {
+      this.event = null;
+      this.transactionEvent = this.loadTransactionEvent();
       this.metadataExtras = Collections.emptyMap();
     } else {
       Event.Metadata metadata = this.loadMetadata();
       Event.Before before = this.loadBefore();
       Event.After after = this.loadAfter();
+      Event.Transaction transaction = this.loadTransaction();
 
-      this.event = new Event(metadata, before, after);
+      this.event = new Event(metadata, before, after, transaction);
+      this.transactionEvent = null;
       this.metadataExtras =
           fullMetadata && metadata != null ? this.loadMetadataExtras() : Collections.emptyMap();
     }
@@ -251,12 +280,64 @@
     return new Event.After(fields);
   }
 
+  /**
+   * Extracts the transaction field within SourceRecord, set when the connector has {@code
+   * provide.transaction.metadata=true}.
+   *
+   * @return Transaction
+   */
+  private Event.Transaction loadTransaction() {
+    @Nullable Struct transaction;
+    try {
+      transaction = (Struct) this.value.get("transaction");
+    } catch (DataException e) {
+      return null;
+    }
+    if (transaction == null) {
+      return null;
+    }
+
+    return new Event.Transaction(
+        getStringOrNull(transaction, "id"),
+        getInt64OrNull(transaction, "total_order"),
+        getInt64OrNull(transaction, "data_collection_order"));
+  }
+
+  /**
+   * Extracts the payload of a transaction BEGIN/END event.
+   *
+   * @return TransactionEvent
+   */
+  private TransactionEvent loadTransactionEvent() {
+    @Nullable List<TransactionEvent.DataCollection> dataCollections = null;
+    @Nullable List<Struct> dataCollectionStructs = getArrayOrNull(this.value, "data_collections");
+    if (dataCollectionStructs != null) {
+      dataCollections = new ArrayList<>();
+      for (Struct dataCollection : dataCollectionStructs) {
+        dataCollections.add(
+            new TransactionEvent.DataCollection(
+                getStringOrNull(dataCollection, "data_collection"),
+                getInt64OrNull(dataCollection, "event_count")));
+      }
+    }
+
+    return new TransactionEvent(
+        getStringOrNull(this.value, "status"),
+        getStringOrNull(this.value, "id"),
+        getInt64OrNull(this.value, "ts_ms"),
+        getInt64OrNull(this.value, "event_count"),
+        dataCollections);
+  }
+
   /**
    * Transforms the extracted data to a JSON string.
    *
    * @return JSON String
    */
   public String toJson() {
+    if (this.transactionEvent != null) {
+      return this.transactionEvent.toJson();
+    }
     return this.event.toJson(this.metadataExtras);
   }
 
@@ -291,21 +372,25 @@
     private final SourceRecordJson.Event.Metadata metadata;
     private final SourceRecordJson.Event.Before before;
     private final SourceRecordJson.Event.After after;
+    private final SourceRecordJson.Event.Transaction transaction;
 
     /**
      * Event Initializer.
      *
      * @param metadata Metadata retrieved from SourceRecord
      * @param before Before data retrieved from SourceRecord
      * @param after After data retrieved from SourceRecord
+     * @param transaction Transaction block retrieved from SourceRecord
      */
     public Event(
         SourceRecordJson.Event.Metadata metadata,
         SourceRecordJson.Event.Before before,
-        SourceRecordJson.Event.After after) {
+        SourceRecordJson.Event.After after,
+        SourceRecordJson.Event.Transaction transaction) {
       this.metadata = metadata;
       this.before = before;
       this.after = after;
+      this.transaction = transaction;
     }
 
     /**
@@ -319,23 +404,25 @@
 
     /**
      * Transforms the Event to a JSON string, appending the given extra keys to {@code metadata}
-     * after {@code table}.
+     * after {@code table}. The {@code transaction} key is omitted when the record has none.
      *
      * @param extras Additional metadata keys, in iteration order
      * @return JSON String
      */
     String toJson(Map<String, Object> extras) {
-      Gson gson = new GsonBuilder().serializeNulls().create();
+      JsonObject tree = GSON.toJsonTree(this).getAsJsonObject();
+      if (transaction == null) {
+        tree.remove("transaction");
+      }
       if (extras.isEmpty() || metadata == null) {
-        return gson.toJson(this);
+        return GSON.toJson(tree);
       }
       // 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()));
+        metadataJson.add(extra.getKey(), GSON.toJsonTree(extra.getValue()));
       }
-      return gson.toJson(tree);
+      return GSON.toJson(tree);
     }
 
     /** Depicts the metadata within a SourceRecord. It has valuable fields. */
@@ -401,5 +488,122 @@
         this.fields = fields;
       }
     }
+
+    /** Depicts the transaction field within SourceRecord. */
+    @SuppressWarnings("unused") // Unused params are validated as JSON
+    static class Transaction implements Serializable {
+      private final @Nullable String id;
+
+      @SerializedName("total_order")
+      private final @Nullable Long totalOrder;
+
+      @SerializedName("data_collection_order")
+      private final @Nullable Long dataCollectionOrder;
+
+      /**
+       * Transaction Initializer.
+       *
+       * @param id Transaction id
+       * @param totalOrder Position of the event among all events of the transaction
+       * @param dataCollectionOrder Position of the event among the events of its data collection
+       */
+      public Transaction(
+          @Nullable String id, @Nullable Long totalOrder, @Nullable Long dataCollectionOrder) {
+        this.id = id;
+        this.totalOrder = totalOrder;
+        this.dataCollectionOrder = dataCollectionOrder;
+      }
+    }
+  }
+
+  /** Depicts a transaction BEGIN/END SourceRecord as a TransactionEvent to be mapped as JSON. */
+  @SuppressWarnings("unused") // Unused params are validated as JSON
+  static class TransactionEvent implements Serializable {
+    private final @Nullable String status;
+    private final @Nullable String id;
+
+    @SerializedName("ts_ms")
+    private final @Nullable Long tsMs;
+
+    @SerializedName("event_count")
+    private final @Nullable Long eventCount;
+
+    @SerializedName("data_collections")
+    private final @Nullable List<DataCollection> dataCollections;
+
+    /**
+     * TransactionEvent Initializer.
+     *
+     * @param status {@code BEGIN} or {@code END}
+     * @param id Transaction id
+     * @param tsMs Time the transaction began or ended, in milliseconds
+     * @param eventCount Number of events in the transaction, only set on {@code END}
+     * @param dataCollections Number of events per data collection, only set on {@code END}
+     */
+    public TransactionEvent(
+        @Nullable String status,
+        @Nullable String id,
+        @Nullable Long tsMs,
+        @Nullable Long eventCount,
+        @Nullable List<DataCollection> dataCollections) {
+      this.status = status;
+      this.id = id;
+      this.tsMs = tsMs;
+      this.eventCount = eventCount;
+      this.dataCollections = dataCollections;
+    }
+
+    /**
+     * Transforms the TransactionEvent to a JSON string.
+     *
+     * @return JSON String
+     */
+    public String toJson() {
+      return GSON.toJson(this);
+    }
+
+    /** Depicts an entry of the data_collections field within SourceRecord. */
+    @SuppressWarnings("unused") // Unused params are validated as JSON
+    static class DataCollection implements Serializable {
+      @SerializedName("data_collection")
+      private final @Nullable String dataCollection;
+
+      @SerializedName("event_count")
+      private final @Nullable Long eventCount;
+
+      /**
+       * DataCollection Initializer.
+       *
+       * @param dataCollection Data collection (table) name
+       * @param eventCount Number of events of the transaction in this data collection
+       */
+      public DataCollection(@Nullable String dataCollection, @Nullable Long eventCount) {
+        this.dataCollection = dataCollection;
+        this.eventCount = eventCount;
+      }
+    }
+  }
+
+  /**
+   * Whether the record is a transaction BEGIN/END event, emitted when {@code
+   * provide.transaction.metadata} is enabled. These are identified by their key schema name.
+   */
+  private static boolean isTransactionMetadataEvent(@Nullable Schema keySchema) {
+    return keySchema != null && SCHEMA_TRANSACTION_METADATA_EVENT_KEY_NAME.equals(keySchema.name());
+  }
+
+  /** Returns the string field's value, or null if the struct's schema has no such field. */
+  private static @Nullable String getStringOrNull(Struct s, String name) {
+    return s.schema().field(name) != null ? s.getString(name) : null;
+  }
+
+  /** Returns the int64 field's value, or null if the struct's schema has no such field. */
+  private static @Nullable Long getInt64OrNull(Struct s, String name) {
+    return s.schema().field(name) != null ? s.getInt64(name) : null;
+  }
+
+  /** Returns the array field's value, or null if the struct's schema has no such field. */
+  private static <T> @Nullable List<T> getArrayOrNull(Struct s, String name) {
+    return s.schema().field(name) != null ? s.getArray(name) : null;
   }
 }