feat: initial commit for merica-pipeline delivery
This commit is contained in:
MERICA Admin
@@ -0,0 +1,178 @@
|
||||
"""
|
||||
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
|
||||
Reference in New Issue
Block a user