Add IcebergChangeHandlerV2 handler for flattened data consuming, schema will be inferred by arrow (#23)

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

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

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

Author: ismailsimsek

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]`)
 

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
@@ -141,7 +143,8 @@ def load_table(self, table_identifier):
             table = self.catalog.create_table(identifier=table_identifier,
                                               schema=self._target_schema,
                                               partition_spec=self.DEBEZIUM_TABLE_PARTITION_SPEC)
-            self.log.info(f"Created iceberg table {'.'.join(table_identifier)} with daily partitioning on _consumed_at.")
+            self.log.info(
+                f"Created iceberg table {'.'.join(table_identifier)} with daily partitioning on _consumed_at.")
             return table
 
     @property
@@ -200,3 +203,208 @@ 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
+        )
+
+        self._handle_schema_changes(table=table, arrow_schema=enriched_arrow_data.schema)
+        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(), unexpected_field_behavior="infer")
+        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
+
+    def _handle_schema_changes(self, table: "Table", arrow_schema: "pa.Schema"):
+        with table.update_schema() as update:
+            update.union_by_name(new_schema=arrow_schema)
+        self.log.info(f"Schema for table {'.'.join(table.name())} has been updated.")

tests/base_postgresql_test.py renamed to tests/base_postgresql.py

@@ -0,0 +1,29 @@
+from pydbzengine import ChangeEvent
+
+
+class MockChangeEvent(ChangeEvent):
+    """
+    A concrete implementation of ChangeEvent for testing.
+    """
+
+    def __init__(self, key: str, value: str, destination: str, partition: int=1):
+        self._key = key
+        self._value = value
+        self._destination = destination
+        self._partition = partition
+
+    def key(self) -> str:
+        """Returns the record key."""
+        return self._key
+
+    def value(self) -> str:
+        """Returns the record value (payload)."""
+        return self._value
+
+    def destination(self) -> str:
+        """Returns the destination topic/table."""
+        return self._destination
+
+    def partition(self) -> int:
+        """Returns the partition the record belongs to."""
+        return self._partition

tests/mock_events.py

@@ -1,7 +1,7 @@
 import logging
 from typing import List
 
-from base_postgresql_test import BasePostgresqlTest
+from base_postgresql import BasePostgresqlTest
 from pydbzengine import ChangeEvent, BasePythonChangeHandler
 from pydbzengine import DebeziumJsonEngine
 

tests/test_change_handler.py

@@ -3,7 +3,7 @@
 import dlt
 import duckdb
 
-from base_postgresql_test import BasePostgresqlTest
+from base_postgresql import BasePostgresqlTest
 from pydbzengine import DebeziumJsonEngine
 from pydbzengine.handlers.dlt import DltChangeHandler
 from pydbzengine.helper import Utils

tests/test_dlt_handler.py

@@ -4,7 +4,7 @@
 from pyiceberg.schema import Schema
 from pyiceberg.types import LongType, NestedField, StringType
 
-from base_postgresql_test import BasePostgresqlTest
+from base_postgresql import BasePostgresqlTest
 from catalog_rest import CatalogRestContainer
 from pydbzengine import DebeziumJsonEngine
 from pydbzengine.handlers.iceberg import IcebergChangeHandler