diff --git a/README.md b/README.md index 6883249..5f462dd 100644 --- 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 index 61be864..2d12e8a 100644 --- a/pydbzengine/handlers/iceberg.py +++ b/pydbzengine/handlers/iceberg.py @@ -1,4 +1,5 @@ import datetime +import io import json import logging import uuid @@ -6,6 +7,7 @@ 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.") diff --git a/tests/base_postgresql_test.py b/tests/base_postgresql.py similarity index 100% rename from tests/base_postgresql_test.py rename to tests/base_postgresql.py diff --git a/tests/mock_events.py b/tests/mock_events.py new file mode 100644 index 0000000..cc43cd5 --- /dev/null +++ b/tests/mock_events.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 diff --git a/tests/test_change_handler.py b/tests/test_change_handler.py index 13b3f0f..e994dca 100644 --- a/tests/test_change_handler.py +++ b/tests/test_change_handler.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 diff --git a/tests/test_dlt_handler.py b/tests/test_dlt_handler.py index 9583e1f..5f9f28b 100644 --- a/tests/test_dlt_handler.py +++ b/tests/test_dlt_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 diff --git a/tests/test_iceberg_handler.py b/tests/test_iceberg_handler.py index f9700ae..ee0c463 100644 --- a/tests/test_iceberg_handler.py +++ b/tests/test_iceberg_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 diff --git a/tests/test_iceberg_handlerv2.py b/tests/test_iceberg_handlerv2.py new file mode 100644 index 0000000..f0f0629 --- /dev/null +++ b/tests/test_iceberg_handlerv2.py @@ -0,0 +1,160 @@ +import io +import time + +import pandas as pd +import pyarrow as pa +import pyarrow.json as pj +from pyiceberg.catalog import load_catalog + +from base_postgresql import BasePostgresqlTest +from catalog_rest import CatalogRestContainer +from mock_events import MockChangeEvent +from pydbzengine.handlers.iceberg import IcebergChangeHandlerV2 +from s3_minio import S3Minio + + +class TestIcebergChangeHandlerV2(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.RESTCATALOG.stop() + self.clean_offset_file() + + def test_read_json_lines_example(self): + json_data = """ +{"id": 1, "name": "Alice", "age": 30} +{"id": 2, "name": "Bob", "age": 24} +{"id": 3, "name": "Charlie", "age": 35} + """.strip() # .strip() removes leading/trailing whitespace/newlines + json_buffer = io.BytesIO(json_data.encode('utf-8')) + json_buffer.seek(0) + # ============================= + table_inferred = pj.read_json(json_buffer) + print("\nInferred Schema:") + print(table_inferred.schema) + # ============================= + explicit_schema = pa.schema([ + pa.field('id', pa.int64()), # Integer type for 'id' + pa.field('name', pa.string()), # String type for 'name' + ]) + json_buffer.seek(0) + po = pj.ParseOptions(explicit_schema=explicit_schema) + table_explicit = pj.read_json(json_buffer, parse_options=po) + print("\nExplicit Schema:") + print(table_explicit.schema) + + def _apply_source_db_changes(self): + time.sleep(12) + 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 ;") + + def test_iceberg_handler(self): + dest_ns1_database = "my_warehouse" + dest_ns2_schema = "dbz_cdc_data" + catalog_conf = { + "uri": self.RESTCATALOG.get_uri(), + "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, + } + destination_namespace=(dest_ns1_database, dest_ns2_schema,) + handler = IcebergChangeHandlerV2(catalog=load_catalog(name="rest", **catalog_conf), + destination_namespace=destination_namespace, + event_flattening_enabled=True + ) + + # TEST + testing_catalog = load_catalog(name="rest", **catalog_conf) + testing_catalog.create_namespace(namespace=destination_namespace) + + test_dest = "inventory.customers" + records = [ + MockChangeEvent(key='{"id": 1}', value='{"id": 1, "name": "Alice", "age": 30}', destination=test_dest), + MockChangeEvent(key='{"id": 2}', value='{"id": 2, "name": "Bob", "age": 24}', destination=test_dest), + MockChangeEvent(key='{"id": 3}', value='{"id": 3, "name": "Charlie", "age": 35}', destination=test_dest) + ] + handler.handleJsonBatch(records) + + test_ns = (dest_ns1_database,) + print(testing_catalog.list_namespaces()) + self._wait_for_condition( + predicate=lambda: test_ns in testing_catalog.list_namespaces(), + failure_message=f"Namespace {test_ns} did not appear in the catalog" + ) + + test_tbl_ref = ('my_warehouse', 'dbz_cdc_data', 'inventory_customers') + test_tbl_ns = (dest_ns1_database, dest_ns2_schema,) + self._wait_for_condition( + predicate=lambda: test_tbl_ref in testing_catalog.list_tables(test_tbl_ns), + failure_message=f"Table {test_tbl_ref} did not appear in the tables" + ) + test_tbl_ref = ('my_warehouse', 'dbz_cdc_data', 'inventory_customers') + data = self.red_table(testing_catalog, test_tbl_ref) + self.pprint_table(data=data) + + self._wait_for_condition( + predicate=lambda: "Charlie" in str(self.red_table(testing_catalog, test_tbl_ref)), + failure_message=f"Expected row not consumed!" + ) + self._wait_for_condition( + predicate=lambda: self.red_table(testing_catalog, test_tbl_ref).num_rows >= 3, + failure_message=f"Rows not consumed" + ) + # # ================================================================= + # ## ==== PART 2 Add additional field ======================== + # # ================================================================= + records = [ + MockChangeEvent(key='{"id": 1}', value='{"id": 1, "name": "Alice", "age": 30, "lastname":"Wonder"}', destination=test_dest), + MockChangeEvent(key='{"id": 2}', value='{"id": 2, "name": "Bob", "age": 24, "lastname":"Sponge"}', destination=test_dest), + MockChangeEvent(key='{"id": 3}', value='{"id": 3, "name": "Charlie", "age": 35, "lastname":"Chaplin"}', destination=test_dest) + ] + handler.handleJsonBatch(records) + data = self.red_table(testing_catalog, test_tbl_ref) + self.pprint_table(data=data) + + def red_table(self, catalog, table_identifier) -> "pa.Table": + tbl = catalog.load_table(identifier=table_identifier) + data = tbl.scan().to_arrow() + self.pprint_table(data) + return data + + def pprint_table(self, data): + print("--- Iceberg Table Content ---") + print(data.to_pandas()) + print("---------------------------\n") + + def _wait_for_condition(self, predicate, failure_message: str, retries: int = 10, delay_seconds: int = 2): + attempts = 0 + while attempts < retries: + print(f"Attempt {attempts + 1}/{retries}: Checking condition...") + if predicate(): + print("Condition met.") + return + + attempts += 1 + # Avoid sleeping after the last attempt + if attempts < retries: + time.sleep(delay_seconds) + + raise TimeoutError(f"{failure_message} after {retries} attempts ({retries * delay_seconds} seconds).")