Add iceberg v2 consumer for flattened data consuming, schema will be inferred by arrow

ismailsimsek · ismailsimsek · commit 8a6c2dbc4c9a · 2025-07-22T19:22:53.000+02:00
diff --git a/README.md b/README.md
@@ -34,6 +34,13 @@ The Iceberg handlers are designed to stream CDC events directly into Apache Iceb
         *   **Automatic Table Creation & Partitioning**: It automatically creates a new Iceberg table for each source table and partitions it by day on the `_consumed_at` timestamp for efficient time-series queries.
         *   **Enriched Metadata**: It also adds `_consumed_at`, `_dbz_event_key`, and `_dbz_event_key_hash` columns for enhanced traceability.
 
+*   `IcebergChangeHandlerV2`: A more advanced handler that automatically infers the schema from the Debezium events and creates a well-structured Iceberg table accordingly.
+    *   **Use Case**: Ideal for scenarios where you want the pipeline to automatically create tables with native data types that mirror the source. This allows for direct querying of the data without needing to parse JSON.
+    *   **Schema and Features**:
+        *   **Automatic Schema Inference**: It inspects the first batch of records for a given table and infers the schema using PyArrow, preserving native data types (e.g., `LongType`, `TimestampType`).
+        *   **Robust Type Handling**: If a field's type cannot be inferred from the initial batch (e.g., it is always `null`), it safely falls back to `StringType` to prevent errors.
+        *   **Automatic Table Creation & Partitioning**: It automatically creates a new Iceberg table for each source table and partitions it by day on the `_consumed_at` timestamp for efficient time-series queries.
+        *   **Enriched Metadata**: It also adds `_consumed_at`, `_dbz_event_key`, and `_dbz_event_key_hash` columns for enhanced traceability.
 
 ### dlt (data load tool) Handler (`pydbzengine[dlt]`)
 
diff --git a/pydbzengine/handlers/iceberg.py b/pydbzengine/handlers/iceberg.py
@@ -1,11 +1,13 @@
 import datetime
+import io
 import json
 import logging
 import uuid
 from abc import abstractmethod
 from typing import List, Dict
 
 import pyarrow as pa
+from pyarrow import json as pa_json
 from pyiceberg.catalog import Catalog
 from pyiceberg.exceptions import NoSuchTableError
 from pyiceberg.partitioning import PartitionSpec, PartitionField
@@ -200,3 +202,202 @@ def _target_schema(self) -> Schema:
             ),
         )
 
+
+class IcebergChangeHandlerV2(BaseIcebergChangeHandler):
+    """
+    A change handler that uses Apache Iceberg to process Debezium change events.
+    This class receives batches of Debezium ChangeEvent objects and applies the changes
+    to the corresponding Iceberg tables.
+    """
+
+    def __init__(self, catalog: "Catalog", destination_namespace: tuple, supports_variant: bool = False,
+                 event_flattening_enabled=False):
+        super().__init__(catalog, destination_namespace, supports_variant)
+        self.event_flattening_enabled = event_flattening_enabled
+
+    def _handle_table_changes(self, destination: str, records: List[ChangeEvent]):
+        """
+        Handles changes for a specific table.
+        Args:
+            destination: The name of the table to apply the changes to.
+            records: A list of ChangeEvent objects for the specified table.
+        """
+
+        table = self.get_table(destination)
+        if table is None:
+            table_identifier: tuple = self.destination_to_table_identifier(destination=destination)
+            table = self._infer_and_create_table(records=records, table_identifier=table_identifier)
+        #
+        arrow_data = self._read_to_arrow_table(records=records, schema=table.schema())
+
+        # Populate all metadata columns (_consumed_at, _dbz_event_key, etc.)
+        enriched_arrow_data = self._enrich_arrow_table_with_metadata(
+            arrow_table=arrow_data,
+            records=records
+        )
+
+        table.append(enriched_arrow_data)
+        self.log.info(f"Appended {len(enriched_arrow_data)} records to table {'.'.join(table.name())}")
+
+    def _enrich_arrow_table_with_metadata(self, arrow_table: pa.Table, records: List[ChangeEvent]) -> pa.Table:
+        num_records = len(arrow_table)
+        dbz_event_keys = []
+        dbz_event_key_hashes = []
+
+        for record in records:
+            key = record.key()
+            dbz_event_keys.append(key)
+            key_hash = str(uuid.uuid5(uuid.NAMESPACE_DNS, key)) if key else None
+            dbz_event_key_hashes.append(key_hash)
+
+        # Create PyArrow arrays for each metadata column
+        consumed_at_array = pa.array([datetime.datetime.now(datetime.timezone.utc)] * num_records,
+                                     type=pa.timestamp('us', tz='UTC'))
+        dbz_event_key_array = pa.array(dbz_event_keys, type=pa.string())
+        dbz_event_key_hash_array = pa.array(dbz_event_key_hashes, type=pa.string())
+
+        # Replace the null columns in the Arrow table with the populated arrays.
+        # This uses set_column, which is efficient for replacing entire columns.
+        enriched_table = arrow_table.set_column(
+            arrow_table.schema.get_field_index("_consumed_at"),
+            "_consumed_at",
+            consumed_at_array
+        )
+        enriched_table = enriched_table.set_column(
+            enriched_table.schema.get_field_index("_dbz_event_key"),
+            "_dbz_event_key",
+            dbz_event_key_array
+        )
+        enriched_table = enriched_table.set_column(
+            enriched_table.schema.get_field_index("_dbz_event_key_hash"),
+            "_dbz_event_key_hash",
+            dbz_event_key_hash_array
+        )
+
+        return enriched_table
+
+    def _read_to_arrow_table(self, records, schema=None):
+        json_lines_buffer = io.BytesIO()
+        for record in records:
+            json_lines_buffer.write((record.value() + '\n').encode('utf-8'))
+        json_lines_buffer.seek(0)
+
+        parse_options = None
+        if schema:
+            # If an Iceberg schema is provided, convert it to a PyArrow schema and use it for parsing.
+            parse_options = pa_json.ParseOptions(explicit_schema=schema.as_arrow())
+        return pa_json.read_json(json_lines_buffer, parse_options=parse_options)
+
+    def get_table(self, destination: str) -> "Table":
+        table_identifier: tuple = self.destination_to_table_identifier(destination=destination)
+        return self.load_table(table_identifier=table_identifier)
+
+    def load_table(self, table_identifier):
+        try:
+            return self.catalog.load_table(identifier=table_identifier)
+        except NoSuchTableError:
+            return None
+
+    def _infer_and_create_table(self, records: List[ChangeEvent], table_identifier: tuple) -> Table:
+        """
+        Infers a schema from a batch of records, creates a new Iceberg table with that schema,
+        and sets up daily partitioning on the _consumed_at field.
+        """
+        arrow_table = self._read_to_arrow_table(records)
+        sanitized_fields = self._sanitize_schema_fields(data_schema=arrow_table.schema)
+
+        # Add metadata fields to the list of pyarrow fields
+        sanitized_fields.extend([
+            pa.field("_consumed_at", pa.timestamp('us', tz='UTC')),
+            pa.field("_dbz_event_key", pa.string()),
+            pa.field("_dbz_event_key_hash", pa.string())  # For UUIDType
+        ])
+
+        # Create a pyarrow schema first
+        sanitized_arrow_schema = pa.schema(sanitized_fields)
+
+        # Create the table
+        table = self.catalog.create_table(
+            identifier=table_identifier,
+            schema=sanitized_arrow_schema
+        )
+        # add partitioning
+        with table.update_spec() as update_spec:
+            update_spec.add_field(source_column_name="_consumed_at", transform=DayTransform(),
+                                  partition_field_name="_consumed_at_day")
+
+        if self.event_flattening_enabled and False:
+            # @TODO fix future. https://github.com/apache/iceberg-python/issues/1728
+            identifier_fields = self._get_identifier_fields(sample_event=records[0],
+                                                            table_identifier=table_identifier
+                                                            )
+            if identifier_fields:
+                with table.update_schema(allow_incompatible_changes=True) as update_schema:
+                    for field in identifier_fields:
+                        update_schema._set_column_requirement(path=field, required=True)
+                    update_schema.set_identifier_fields(*identifier_fields)
+        self.log.info(f"Created iceberg table {'.'.join(table_identifier)} with daily partitioning on _consumed_at.")
+        return table
+
+    def _sanitize_schema_fields(self, data_schema: pa.Schema) -> list:
+        """
+        Recursively traverses a PyArrow schema and replaces null types with a string fallback.
+        This is useful when a schema is inferred from JSON where some fields are always null.
+        """
+        new_fields = []
+        for field in data_schema:
+            field_type = field.type
+            if pa.types.is_null(field_type):
+                # Found a null type, replace it with a string as a fallback.
+                new_fields.append(field.with_type(pa.string()))
+            elif pa.types.is_struct(field_type):
+                # Found a struct, so we recurse on its fields to sanitize them.
+                # We can treat the struct's fields as a schema for the recursive call.
+                nested_schema = pa.schema(field_type)
+                sanitized_nested_schema = self._sanitize_schema_fields(nested_schema)
+                # Recreate the field with the new, sanitized struct type.
+                new_fields.append(field.with_type(pa.struct(sanitized_nested_schema)))
+            else:
+                # Not a null or struct, so we keep the field as is.
+                new_fields.append(field)
+        return new_fields
+
+    def _get_identifier_fields(self, sample_event: ChangeEvent, table_identifier: tuple) -> list:
+        """
+        Parses the Debezium event key to extract primary key field names.
+
+        This method uses a series of guard clauses to validate the key and returns
+        an empty list if any validation step fails.
+
+        Args:
+            sample_event: A sample change event to inspect for the key.
+            table_identifier: The identifier of the table, used for logging.
+
+        Returns:
+            A list of key field names, or an empty list if the key cannot be determined.
+        """
+        key_json_str = sample_event.key()
+        table_name_str = '.'.join(table_identifier)
+
+        if not key_json_str:
+            self.log.warning(f"Cannot determine identifier fields for {table_name_str}: event key is empty.")
+            return []
+
+        try:
+            key_data = json.loads(key_json_str)
+        except json.JSONDecodeError:
+            self.log.error(f"Failed to parse Debezium event key as JSON for table {table_name_str}: {key_json_str}")
+            return []
+
+        if not isinstance(key_data, dict):
+            self.log.warning(
+                f"Event key for {table_name_str} is not a JSON object, cannot infer primary key. Key: {key_json_str}")
+            return []
+
+        key_field_names = list(key_data.keys())
+        if not key_field_names:
+            self.log.warning(f"Event key for {table_name_str} is an empty JSON object, cannot infer primary key.")
+            return []
+
+        self.log.info(f"Found potential primary key fields {key_field_names} for table {table_name_str}")
+        return key_field_names
diff --git a/tests/test_iceberg_handlerv2.py b/tests/test_iceberg_handlerv2.py
@@ -0,0 +1,124 @@
+import unittest
+
+import pandas as pd
+from pyiceberg.catalog import load_catalog
+from pyiceberg.schema import Schema
+from pyiceberg.types import LongType, NestedField, StringType
+
+from base_postgresql_test import BasePostgresqlTest
+from catalog_rest import CatalogRestContainer
+from pydbzengine import DebeziumJsonEngine
+from pydbzengine.handlers.iceberg import IcebergChangeHandlerV2
+from pydbzengine.helper import Utils
+from s3_minio import S3Minio
+
+
+class TestIcebergChangeHandler(BasePostgresqlTest):
+    S3MiNIO = S3Minio()
+    RESTCATALOG = CatalogRestContainer()
+
+    def setUp(self):
+        print("setUp")
+        self.clean_offset_file()
+        self.SOURCEPGDB.start()
+        self.S3MiNIO.start()
+        self.RESTCATALOG.start(s3_endpoint=self.S3MiNIO.endpoint())
+        # Set pandas options to display all rows and columns, and prevent truncation of cell content
+        pd.set_option('display.max_rows', None)         # Show all rows
+        pd.set_option('display.max_columns', None)      # Show all columns
+        pd.set_option('display.width', None)            # Auto-detect terminal width
+        pd.set_option('display.max_colwidth', None)     # Do not truncate cell contents
+
+    def tearDown(self):
+        self.SOURCEPGDB.stop()
+        self.S3MiNIO.stop()
+        self.clean_offset_file()
+
+    @unittest.skip
+    def test_iceberg_catalog(self):
+        conf = {
+            "uri": self.RESTCATALOG.get_uri(),
+            # "s3.path-style.access": "true",
+            "warehouse": "warehouse",
+            "s3.endpoint": self.S3MiNIO.endpoint(),
+            "s3.access-key-id": S3Minio.AWS_ACCESS_KEY_ID,
+            "s3.secret-access-key": S3Minio.AWS_SECRET_ACCESS_KEY,
+        }
+        print(conf)
+        catalog = load_catalog(
+            name="rest",
+            **conf
+        )
+        catalog.create_namespace('my_warehouse')
+        debezium_event_schema = Schema(
+            NestedField(field_id=1, name="id", field_type=LongType(), required=True),
+            NestedField(field_id=2, name="data", field_type=StringType(), required=False),
+        )
+        table = catalog.create_table(identifier=("my_warehouse", "test_table",), schema=debezium_event_schema)
+        print(f"Created iceberg table {table.refs()}")
+
+    def test_iceberg_handler(self):
+        dest_ns1_database="my_warehouse"
+        dest_ns2_schema="dbz_cdc_data"
+        conf = {
+            "uri": self.RESTCATALOG.get_uri(),
+            # "s3.path-style.access": "true",
+            "warehouse": "warehouse",
+            "s3.endpoint": self.S3MiNIO.endpoint(),
+            "s3.access-key-id": S3Minio.AWS_ACCESS_KEY_ID,
+            "s3.secret-access-key": S3Minio.AWS_SECRET_ACCESS_KEY,
+        }
+        catalog = load_catalog(name="rest",**conf)
+
+        handler = IcebergChangeHandlerV2(catalog=catalog,
+                                         destination_namespace=(dest_ns1_database, dest_ns2_schema,),
+                                         event_flattening_enabled=True
+                                         )
+
+        dbz_props = self.debezium_engine_props(unwrap_messages=True)
+        engine = DebeziumJsonEngine(properties=dbz_props, handler=handler)
+        with self.assertLogs(IcebergChangeHandlerV2.LOGGER_NAME, level='INFO') as cm:
+            # run async then interrupt after timeout time to test the result!
+            Utils.run_engine_async(engine=engine, timeout_sec=44)
+
+        # for t in cm.output:
+        #     print(t)
+        self.assertRegex(text=str(cm.output), expected_regex='.*Created iceberg table.*')
+        self.assertRegex(text=str(cm.output), expected_regex='.*Appended.*records to table.*')
+
+        # catalog.create_namespace(dest_ns1_database)
+        namespaces = catalog.list_namespaces()
+        self.assertIn((dest_ns1_database,) , namespaces, msg="Namespace not found in catalog")
+
+        tables = catalog.list_tables((dest_ns1_database, dest_ns2_schema,))
+        print(tables)
+        self.assertIn(('my_warehouse', 'dbz_cdc_data', 'testc_inventory_customers'), tables, msg="Namespace not found in catalog")
+
+        tbl = catalog.load_table(identifier=('my_warehouse', 'dbz_cdc_data', 'testc_inventory_customers'))
+        data = tbl.scan().to_arrow()
+        self.assertIn("sally.thomas@acme.com", str(data))
+        self.assertIn("annek@noanswer.org", str(data))
+        self.assertEqual(data.num_rows, 4)
+        self.pprint_table(data=data)
+        #=================================================================
+        ## ==== PART 2 CONSUME CHANGES FROM BINLOG =======================
+        #=================================================================
+        self.execute_on_source_db("UPDATE inventory.customers SET first_name='George__UPDATE1' WHERE ID = 1002 ;")
+        # self.execute_on_source_db("ALTER TABLE inventory.customers DROP COLUMN email;")
+        self.execute_on_source_db("UPDATE inventory.customers SET first_name='George__UPDATE2'  WHERE ID = 1002 ;")
+        self.execute_on_source_db("DELETE FROM inventory.orders WHERE purchaser = 1002 ;")
+        self.execute_on_source_db("DELETE FROM inventory.customers WHERE id = 1002 ;")
+        self.execute_on_source_db("ALTER TABLE inventory.customers ADD birth_date date;")
+        self.execute_on_source_db("UPDATE inventory.customers SET birth_date = '2020-01-01'  WHERE id = 1001 ;")
+        # run
+        Utils.run_engine_async(engine=engine, timeout_sec=44)
+        # test
+        # @TODO test that new field is received and added to iceberg!
+        data = tbl.scan().to_arrow()
+        self.pprint_table(data=data)
+        self.assertEqual(data.num_rows, 4)
+
+    def pprint_table(self, data):
+        print("--- Iceberg Table Content ---")
+        print(data.to_pandas())
+        print("---------------------------\n")
