179 lines
6.6 KiB
Python
179 lines
6.6 KiB
Python
"""
|
|
PyIceberg Transformation Operator for MERICA Data Lakehouse Studio
|
|
Executes native Apache Iceberg table management using pyiceberg and PyArrow.
|
|
Applies Lazy Import Guardrail to protect Airflow Scheduler memory.
|
|
"""
|
|
import os
|
|
import logging
|
|
from typing import Any, Dict, List, Optional, Sequence
|
|
|
|
from helpers.base_lakehouse_operator import BaseLakehouseOperator
|
|
try:
|
|
from airflow.sdk.definitions.context import Context
|
|
except ImportError:
|
|
try:
|
|
from airflow.utils.context import Context
|
|
except ImportError:
|
|
class Context(dict):
|
|
pass
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class PyIcebergTransformationOperator(BaseLakehouseOperator):
|
|
"""
|
|
Apache PyIceberg-powered Lakehouse Operator.
|
|
Loads Bronze Parquet data, creates/evolves Silver Iceberg tables,
|
|
and commits atomic snapshots to the Iceberg Catalog.
|
|
"""
|
|
|
|
template_fields: Sequence[str] = (
|
|
"source_parquet_path",
|
|
"target_table_name"
|
|
)
|
|
|
|
def __init__(
|
|
self,
|
|
*,
|
|
source_parquet_path: str,
|
|
target_table_name: str,
|
|
namespace: str = "silver",
|
|
write_mode: str = "overwrite", # "append" or "overwrite"
|
|
catalog_name: Optional[str] = "dlh_lakehouse",
|
|
catalog_type: Optional[str] = "sql", # "sql" (PostgreSQL) or "rest"
|
|
catalog_uri: Optional[str] = None,
|
|
s3_endpoint: Optional[str] = None,
|
|
primary_keys: Optional[List[str]] = None,
|
|
**kwargs
|
|
):
|
|
inlets = kwargs.pop("inlets", [source_parquet_path])
|
|
clean_target = target_table_name if "." in target_table_name else f"{namespace}.{target_table_name}"
|
|
outlets = kwargs.pop("outlets", [clean_target])
|
|
|
|
super().__init__(
|
|
layer=namespace,
|
|
engine="pyiceberg",
|
|
inlets=inlets,
|
|
outlets=outlets,
|
|
**kwargs
|
|
)
|
|
|
|
self.source_parquet_path = source_parquet_path
|
|
self.target_table_name = target_table_name
|
|
self.namespace = namespace
|
|
self.write_mode = write_mode.lower()
|
|
self.catalog_name = catalog_name or "dlh_lakehouse"
|
|
self.catalog_type = catalog_type or "sql"
|
|
self.primary_keys = primary_keys or ["id"]
|
|
|
|
self.catalog_uri = catalog_uri or os.getenv(
|
|
"ICEBERG_CATALOG_URI",
|
|
"postgresql+psycopg2://dlh_user:dlh_password@postgres:5432/dlh_lakehouse"
|
|
)
|
|
self.s3_endpoint = s3_endpoint or os.getenv("MINIO_ENDPOINT", "http://minio:9000")
|
|
|
|
def execute(self, context: Context) -> Dict[str, Any]:
|
|
"""
|
|
Execute PyIceberg Table snapshot commit.
|
|
Lazy imports pyiceberg and pyarrow inside execute().
|
|
"""
|
|
logger.info(
|
|
"Starting PyIceberg Table transformation: %s -> %s.%s",
|
|
self.source_parquet_path, self.namespace, self.target_table_name
|
|
)
|
|
|
|
# Lazy Import Guardrail
|
|
try:
|
|
import pyarrow.parquet as pq
|
|
import pyarrow.fs as pafs
|
|
import pyarrow.csv as pacsv
|
|
from pyiceberg.catalog import load_catalog
|
|
except ImportError as e:
|
|
logger.error("PyIceberg or PyArrow is missing: %s", e)
|
|
raise
|
|
|
|
try:
|
|
# 1. Read Bronze dataset into PyArrow Table via S3FileSystem
|
|
logger.info("Reading source dataset from %s", self.source_parquet_path)
|
|
s3_access_key = os.getenv("MINIO_ROOT_USER", "minio_admin")
|
|
s3_secret_key = os.getenv("MINIO_ROOT_PASSWORD", "minio_password")
|
|
|
|
clean_endpoint = self.s3_endpoint
|
|
if not clean_endpoint.startswith("http"):
|
|
clean_endpoint = f"http://{clean_endpoint}"
|
|
|
|
s3_fs = pafs.S3FileSystem(
|
|
endpoint_override=clean_endpoint,
|
|
access_key=s3_access_key,
|
|
secret_key=s3_secret_key,
|
|
scheme="http"
|
|
)
|
|
|
|
if self.source_parquet_path.startswith("s3://"):
|
|
s3_key = self.source_parquet_path.replace("s3://", "")
|
|
if s3_key.lower().endswith(".csv"):
|
|
arrow_table = pacsv.read_csv(s3_key, filesystem=s3_fs)
|
|
else:
|
|
arrow_table = pq.read_table(s3_key, filesystem=s3_fs)
|
|
else:
|
|
if self.source_parquet_path.lower().endswith(".csv"):
|
|
arrow_table = pacsv.read_csv(self.source_parquet_path)
|
|
else:
|
|
arrow_table = pq.read_table(self.source_parquet_path)
|
|
|
|
row_count = arrow_table.num_rows
|
|
|
|
# 2. Connect to Iceberg Catalog
|
|
catalog_props = {
|
|
"type": self.catalog_type,
|
|
"uri": self.catalog_uri,
|
|
"s3.endpoint": clean_endpoint,
|
|
"s3.access-key-id": s3_access_key,
|
|
"s3.secret-access-key": s3_secret_key,
|
|
"s3.path-style-access": "true",
|
|
"s3.region": "us-east-1",
|
|
}
|
|
catalog = load_catalog(self.catalog_name, **catalog_props)
|
|
|
|
# Ensure namespace exists in catalog
|
|
clean_table = self.target_table_name.split(".")[-1]
|
|
full_identifier = f"{self.namespace}.{clean_table}"
|
|
try:
|
|
catalog.create_namespace_if_not_exists(self.namespace)
|
|
except Exception as ns_err:
|
|
logger.debug("Namespace check note: %s", ns_err)
|
|
|
|
# 3. Load or Create Table
|
|
try:
|
|
table = catalog.load_table(full_identifier)
|
|
logger.info("Loaded existing Iceberg table '%s'", full_identifier)
|
|
except Exception:
|
|
logger.info("Table '%s' not found. Creating new Iceberg table...", full_identifier)
|
|
table = catalog.create_table(
|
|
identifier=full_identifier,
|
|
schema=arrow_table.schema,
|
|
location=f"s3://dlh-silver/{clean_table}"
|
|
)
|
|
|
|
# 4. Commit snapshot (Append or Overwrite)
|
|
if self.write_mode == "append":
|
|
table.append(arrow_table)
|
|
else:
|
|
table.overwrite(arrow_table)
|
|
|
|
logger.info(
|
|
"✅ PyIceberg snapshot committed! %d records committed to table '%s'",
|
|
row_count, full_identifier
|
|
)
|
|
|
|
return {
|
|
"status": "success",
|
|
"engine": "pyiceberg",
|
|
"table": full_identifier,
|
|
"rows_committed": row_count,
|
|
"write_mode": self.write_mode
|
|
}
|
|
except Exception as e:
|
|
logger.error("PyIceberg Transformation Failed: %s", e)
|
|
raise
|