[DebeziumIO] Enable Transaction Metadata Events in DebeziumIO JSON Output

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

@@ -172,4 +172,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/DebeziumIO.java

@@ -312,7 +312,8 @@
     /**
      * 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.
+     * to {@code metadata} when reading as JSON, plus the change event's {@code transaction} block
+     * as a nested object when the connector has {@code provide.transaction.metadata=true}.
      *
      * <p>Only applies to {@link DebeziumIO#readAsJson()} / {@link
      * SourceRecordJson.SourceRecordJsonMapper}; combining it with a custom {@link
@@ -452,9 +453,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.
@@ -471,9 +472,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/SourceRecordJson.java

@@ -19,16 +19,19 @@
 
 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.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 +42,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>
  *
@@ -72,10 +75,20 @@
  *   <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})
+ *   <li>transaction - The Debezium {@code transaction} block of the change event ({@code id},
+ *       {@code total_order} and {@code data_collection_order}) as a nested object, if the connector
+ *       has {@code provide.transaction.metadata=true}
  * </ul>
  *
  * <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:
+ * {@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,11 +111,18 @@
    * 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 Map<String, Object> metadataExtras;
+  private final @Nullable Map<String, Object> transactionMetadata;
 
   /**
    * Initializer.
@@ -130,7 +150,13 @@
     if (this.value == null) {
       this.event = new Event(null, null, null);
       this.metadataExtras = Collections.emptyMap();
+      this.transactionMetadata = null;
+    } else if (isTransactionMetadataEvent(sourceRecord.keySchema())) {
+      this.event = null;
+      this.metadataExtras = Collections.emptyMap();
+      this.transactionMetadata = this.loadTransactionMetadata();
     } else {
+      this.transactionMetadata = null;
       Event.Metadata metadata = this.loadMetadata();
       Event.Before before = this.loadBefore();
       Event.After after = this.loadAfter();
@@ -179,7 +205,7 @@
   /**
    * 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}.
+   * represented in {@link Event.Metadata}, then the envelope's {@code transaction} block (if set).
    *
    * @return extra metadata keys, in insertion order
    */
@@ -195,6 +221,18 @@
         extras.put(fieldName, toJsonFriendly(source.get(field)));
       }
     }
+
+    Struct transaction =
+        this.value.schema().field("transaction") != null
+            ? this.value.getStruct("transaction")
+            : null;
+    if (transaction != null) {
+      Map<String, Object> transactionFields = new LinkedHashMap<>();
+      for (Field field : transaction.schema().fields()) {
+        transactionFields.put(field.name(), transaction.get(field));
+      }
+      extras.put("transaction", transactionFields);
+    }
     return extras;
   }
 
@@ -203,6 +241,33 @@
     return value instanceof Struct ? String.valueOf(value) : value;
   }
 
+  /**
+   * Extracts the payload of a transaction BEGIN/END event.
+   *
+   * @return status, id, ts_ms, event_count and data_collections, in insertion order
+   */
+  private Map<String, Object> loadTransactionMetadata() {
+    List<Map<String, Object>> dataCollections = null;
+    List<Struct> dataCollectionStructs = this.value.getArray("data_collections");
+    if (dataCollectionStructs != null) {
+      dataCollections = new ArrayList<>();
+      for (Struct dataCollectionStruct : dataCollectionStructs) {
+        Map<String, Object> dataCollection = new LinkedHashMap<>();
+        dataCollection.put("data_collection", dataCollectionStruct.getString("data_collection"));
+        dataCollection.put("event_count", dataCollectionStruct.getInt64("event_count"));
+        dataCollections.add(dataCollection);
+      }
+    }
+
+    Map<String, Object> transactionMetadata = new LinkedHashMap<>();
+    transactionMetadata.put("status", this.value.getString("status"));
+    transactionMetadata.put("id", this.value.getString("id"));
+    transactionMetadata.put("ts_ms", this.value.getInt64("ts_ms"));
+    transactionMetadata.put("event_count", this.value.getInt64("event_count"));
+    transactionMetadata.put("data_collections", dataCollections);
+    return transactionMetadata;
+  }
+
   /**
    * Extracts the before field within SourceRecord.
    *
@@ -257,6 +322,9 @@
    * @return JSON String
    */
   public String toJson() {
+    if (this.transactionMetadata != null) {
+      return GSON.toJson(this.transactionMetadata);
+    }
     return this.event.toJson(this.metadataExtras);
   }
 
@@ -325,17 +393,16 @@
      * @return JSON String
      */
     String toJson(Map<String, Object> extras) {
-      Gson gson = new GsonBuilder().serializeNulls().create();
       if (extras.isEmpty() || metadata == null) {
-        return gson.toJson(this);
+        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 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. */
@@ -402,4 +469,12 @@
       }
     }
   }
+
+  /**
+   * 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());
+  }
 }

sdks/java/io/debezium/src/test/java/org/apache/beam/io/debezium/DebeziumIOPostgresSqlConnectorIT.java

@@ -18,10 +18,15 @@
 package org.apache.beam.io.debezium;
 
 import static org.hamcrest.MatcherAssert.assertThat;
+import static org.hamcrest.Matchers.allOf;
+import static org.hamcrest.Matchers.containsString;
 import static org.hamcrest.Matchers.equalTo;
 import static org.hamcrest.Matchers.hasItem;
 
 import io.debezium.connector.postgresql.PostgresConnector;
+import java.sql.Connection;
+import java.sql.SQLException;
+import java.sql.Statement;
 import javax.sql.DataSource;
 import org.apache.beam.sdk.Pipeline;
 import org.apache.beam.sdk.coders.StringUtf8Coder;
@@ -195,4 +200,82 @@ public void testDebeziumIOPostgresSql() {
     p.run().waitUntilFinish();
     POSTGRES_SQL_CONTAINER.stop();
   }
+
+  /**
+   * Debezium - PostgresSql connector Test with transaction metadata.
+   *
+   * <p>Tests that the transaction BEGIN/END events emitted with {@code
+   * provide.transaction.metadata=true}, which have no {@code source} block, are mapped to JSON
+   * alongside the data change events of the transaction, and that with full metadata each change
+   * event carries its {@code transaction} block.
+   */
+  @Test
+  public void testShouldProcessTransactionMetadata() throws SQLException {
+    POSTGRES_SQL_CONTAINER.start();
+
+    try (Connection connection = getPostgresDatasource().getConnection();
+        Statement statement = connection.createStatement()) {
+      statement.execute(
+          "SELECT pg_create_logical_replication_slot('beam_transaction_metadata', 'decoderbufs')");
+      connection.setAutoCommit(false);
+      statement.executeUpdate(
+          "INSERT INTO inventory.customers (id, first_name, last_name, email) "
+              + "VALUES (5001, 'Ann', 'One', 'ann.one@example.com')");
+      statement.executeUpdate(
+          "INSERT INTO inventory.customers (id, first_name, last_name, email) "
+              + "VALUES (5002, 'Bob', 'Two', 'bob.two@example.com')");
+      statement.executeUpdate("UPDATE inventory.customers SET last_name = 'Uno' WHERE id = 5001");
+      connection.commit();
+    }
+
+    String host = POSTGRES_SQL_CONTAINER.getHost();
+    String port = POSTGRES_SQL_CONTAINER.getMappedPort(5432).toString();
+
+    PipelineOptions options = PipelineOptionsFactory.create();
+    Pipeline p = Pipeline.create(options);
+    PCollection<String> results =
+        p.apply(
+            DebeziumIO.readAsJson()
+                .withConnectorConfiguration(
+                    DebeziumIO.ConnectorConfiguration.create()
+                        .withUsername("debezium")
+                        .withPassword("dbz")
+                        .withConnectorClass(PostgresConnector.class)
+                        .withHostName(host)
+                        .withPort(port)
+                        .withConnectionProperty("database.dbname", "inventory")
+                        .withConnectionProperty("table.include.list", "inventory.customers")
+                        .withConnectionProperty("snapshot.mode", "no_data")
+                        .withConnectionProperty("provide.transaction.metadata", "true")
+                        .withConnectionProperty("slot.name", "beam_transaction_metadata"))
+                .withFullMetadata(true)
+                .withMaxNumberOfRecords(5)
+                .withMaxTimeToRun(120_000L));
+
+    PAssert.that(results)
+        .satisfies(
+            (Iterable<String> res) -> {
+              assertThat(Lists.newArrayList(res).size(), equalTo(5));
+              assertThat(res, hasItem(containsString("{\"status\":\"BEGIN\",")));
+              assertThat(res, hasItem(containsString("\"email\":\"ann.one@example.com\"")));
+              assertThat(res, hasItem(containsString("\"email\":\"bob.two@example.com\"")));
+              assertThat(
+                  res,
+                  hasItem(
+                      allOf(
+                          containsString("\"last_name\":\"Uno\""),
+                          containsString("\"transaction\":{\"id\":\""),
+                          containsString("\"total_order\":3,\"data_collection_order\":3}"))));
+              assertThat(
+                  res,
+                  hasItem(
+                      containsString(
+                          "\"event_count\":3,\"data_collections\":"
+                              + "[{\"data_collection\":\"inventory.customers\",\"event_count\":3}]")));
+              return null;
+            });
+
+    p.run().waitUntilFinish();
+    POSTGRES_SQL_CONTAINER.stop();
+  }
 }

sdks/java/io/debezium/src/test/java/org/apache/beam/io/debezium/SourceRecordJsonTest.java

@@ -76,6 +76,39 @@ public void testFullMetadata() {
             .put("file", "mysql-bin.000003")
             .put("pos", 796L)
             .put("row", 0);
+
+    Schema transactionSchema =
+        SchemaBuilder.struct()
+            .optional()
+            .name("event.block")
+            .field("id", Schema.STRING_SCHEMA)
+            .field("total_order", Schema.INT64_SCHEMA)
+            .field("data_collection_order", Schema.INT64_SCHEMA)
+            .build();
+    Struct transaction =
+        new Struct(transactionSchema)
+            .put("id", "796:48630784")
+            .put("total_order", 1L)
+            .put("data_collection_order", 1L);
+    Schema rowSchema = SchemaBuilder.struct().optional().field("id", Schema.INT32_SCHEMA).build();
+    Struct source = postgresSource();
+    Schema envelopeSchema =
+        SchemaBuilder.struct()
+            .field("before", rowSchema)
+            .field("after", rowSchema)
+            .field("source", source.schema())
+            .field("op", Schema.STRING_SCHEMA)
+            .field("transaction", transactionSchema)
+            .build();
+    Struct envelope =
+        new Struct(envelopeSchema)
+            .put("after", new Struct(rowSchema).put("id", 1001))
+            .put("source", source)
+            .put("op", "c")
+            .put("transaction", transaction);
+    SourceRecord transactionRecord = new SourceRecord(null, null, null, envelopeSchema, envelope);
+
+    // MySQL has no "schema" source field, so the binlog "file" is also exposed as "schema".
     assertEquals(
         "{\"metadata\":{\"connector\":\"mysql\",\"version\":\"3.5.2.Final\",\"name\":\"n\","
             + "\"database\":\"inventory\",\"schema\":\"mysql-bin.000003\",\"table\":\"customers\","
@@ -90,6 +123,16 @@ public void testFullMetadata() {
             + "\"op\":\"u\",\"ts_ms\":1791689147000,\"lsn\":37094192},"
             + "\"before\":null,\"after\":{\"fields\":{\"id\":1001}}}",
         new SourceRecordJson(buildEnvelopeRecord(postgresSource(), "u"), true).toJson());
+
+    // With provide.transaction.metadata, the envelope's transaction block is included.
+    assertEquals(
+        "{\"metadata\":{\"connector\":\"postgresql\",\"version\":\"3.5.2.Final\",\"name\":\"n\","
+            + "\"database\":\"postgres\",\"schema\":\"inventory\",\"table\":\"customers\","
+            + "\"op\":\"c\",\"ts_ms\":1791689147000,\"lsn\":37094192,"
+            + "\"transaction\":{\"id\":\"796:48630784\",\"total_order\":1,"
+            + "\"data_collection_order\":1}},"
+            + "\"before\":null,\"after\":{\"fields\":{\"id\":1001}}}",
+        new SourceRecordJson(transactionRecord, true).toJson());
   }
 
   @Test