""" Base Lakehouse Operator for MERICA Data Lakehouse Studio Provides standard lineage extraction, OpenMetadata 2.0 lineage emission, and shared catalog resolution across all lakehouse operators. """ import os import json import logging import urllib.request import urllib.parse from typing import Any, Dict, List, Optional, Sequence, Union try: from airflow.sdk.bases.operator import BaseOperator from airflow.sdk.definitions.context import Context except ImportError: try: from airflow.models.baseoperator import BaseOperator from airflow.utils.context import Context except ImportError: class BaseOperator: def __init__(self, *args, **kwargs): self.task_id = kwargs.get("task_id", "lakehouse_operator") self.inlets = kwargs.get("inlets", []) self.outlets = kwargs.get("outlets", []) self.log = logging.getLogger(self.__class__.__name__) class Context(dict): pass logger = logging.getLogger(__name__) ENABLE_OPENMETADATA_SYNC = os.getenv("ENABLE_OPENMETADATA_SYNC", "true").lower() in ("true", "1", "yes") class BaseLakehouseOperator(BaseOperator): """ Abstract base operator for all MERICA Lakehouse ingestion and transformation tasks. Enforces standardized dataset metadata (inlets/outlets) and automated lineage emission to OpenMetadata 2.0 REST API upon successful execution. """ template_fields: Sequence[str] = ("layer", "pipeline_id") def __init__( self, *, layer: str = "bronze", engine: str = "generic", pipeline_id: Optional[str] = None, inlets: Optional[List[str]] = None, outlets: Optional[List[str]] = None, lakehouse_catalog: Optional[str] = None, **kwargs ): super().__init__(**kwargs) self.layer = layer.lower() self.engine = engine.lower() self.pipeline_id = pipeline_id self.custom_inlets = inlets or [] self.custom_outlets = outlets or [] self.lakehouse_catalog = lakehouse_catalog or os.getenv("ICEBERG_CATALOG_NAME", "dlh_lakehouse") # Populate Airflow native inlets/outlets if supported if inlets: self.inlets = inlets if outlets: self.outlets = outlets def post_execute(self, context: Context, result: Any = None): """ Airflow post-execution lifecycle hook. Automatically emits lineage edge to OpenMetadata 2.0 when task succeeds. """ super().post_execute(context, result) dag_id = self.dag_id or (context.get("dag").dag_id if context.get("dag") else self.pipeline_id) task_id = self.task_id logger.info( "Task '%s' in DAG '%s' (Layer: %s, Engine: %s) completed successfully.", task_id, dag_id, self.layer, self.engine ) if ENABLE_OPENMETADATA_SYNC: self._emit_openmetadata_lineage(context, dag_id, task_id) def _emit_openmetadata_lineage(self, context: Context, dag_id: str, task_id: str): """ Resolve OpenMetadata JWT token and register lineage edges for all inlets -> outlets. """ try: token, active_host = self._get_openmetadata_auth() if not token or not active_host: logger.debug("OpenMetadata host/token unavailable; skipping lineage emission.") return pipeline_fqn = f"dlh_airflow_service.{dag_id}" pipeline_id_uuid = self._resolve_entity_uuid(active_host, token, "pipeline", pipeline_fqn) # Emit lineage for each inlet -> outlet pair for inlet in self.custom_inlets: inlet_fqn = self._normalize_dataset_fqn(inlet) inlet_uuid = self._resolve_entity_uuid(active_host, token, "table", inlet_fqn) if not inlet_uuid: inlet_uuid = self._ensure_table_entity(active_host, token, inlet_fqn) for outlet in self.custom_outlets: outlet_fqn = self._normalize_dataset_fqn(outlet) outlet_uuid = self._resolve_entity_uuid(active_host, token, "table", outlet_fqn) if not outlet_uuid: outlet_uuid = self._ensure_table_entity(active_host, token, outlet_fqn) if inlet_uuid and outlet_uuid: payload = { "edge": { "fromEntity": {"id": inlet_uuid, "type": "table"}, "toEntity": {"id": outlet_uuid, "type": "table"}, } } if pipeline_id_uuid: payload["edge"]["lineageDetails"] = { "pipeline": {"id": pipeline_id_uuid, "type": "pipeline"} } req = urllib.request.Request( f"{active_host}/api/v1/lineage", data=json.dumps(payload).encode("utf-8"), headers={"Authorization": f"Bearer {token}", "Content-Type": "application/json"}, method="PUT" ) with urllib.request.urlopen(req, timeout=4) as resp: logger.info( "✅ Lineage emitted: %s -> %s via pipeline %s", inlet_fqn, outlet_fqn, dag_id ) except Exception as e: logger.warning("OpenMetadata Lineage Emission Note: %s", e) def _normalize_dataset_fqn(self, dataset_str: str) -> str: """ Convert dataset URI/table into OpenMetadata FQN format. e.g. 's3://dlh-bronze/db_cdc/environmental_permits' -> 'dlh_s3_storage.dlh-bronze.db_cdc.environmental_permits' e.g. 'silver.environmental_permits_clean' -> 'dlh_postgres_service.dlh_lakehouse.silver.environmental_permits_clean' """ clean = dataset_str.strip() if clean.startswith("s3://"): parts = clean.replace("s3://", "").strip("/").split("/") return f"dlh_minio_s3.{'.'.join(parts)}" elif "." in clean: return f"dlh_postgres_service.dlh_lakehouse.{clean}" return f"dlh_postgres_service.dlh_lakehouse.public.{clean}" def _get_openmetadata_auth(self) -> tuple: """Authenticate and get access token.""" hosts = [ os.getenv("OPENMETADATA_HOST", "http://openmetadata:8585"), "http://localhost:8585", "http://127.0.0.1:8585" ] for h in hosts: try: login_req = urllib.request.Request( f"{h}/api/v1/users/login", data=json.dumps({"email": "admin@open-metadata.org", "password": "YWRtaW4="}).encode("utf-8"), headers={"Content-Type": "application/json"}, method="POST" ) with urllib.request.urlopen(login_req, timeout=3) as resp: token = json.loads(resp.read().decode("utf-8")).get("accessToken") if token: return token, h except Exception: continue return None, None def _resolve_entity_uuid(self, host: str, token: str, entity_type: str, fqn: str) -> Optional[str]: """Fetch entity UUID from OpenMetadata given FQN.""" try: endpoint = "tables" if entity_type == "table" else f"{entity_type}s" url = f"{host}/api/v1/{endpoint}/name/{urllib.parse.quote(fqn, safe='')}" req = urllib.request.Request( url, headers={"Authorization": f"Bearer {token}", "Accept": "application/json"}, method="GET" ) with urllib.request.urlopen(req, timeout=3) as resp: data = json.loads(resp.read().decode("utf-8")) return data.get("id") except Exception: return None def _ensure_table_entity(self, host: str, token: str, fqn: str) -> Optional[str]: """Auto-provision a missing table entity in OpenMetadata.""" try: parts = fqn.split(".") table_name = parts[-1] schema_fqn = ".".join(parts[:-1]) if len(parts) > 1 else "dlh_postgres_service.dlh_lakehouse.public" payload = { "name": table_name, "displayName": table_name.replace("_", " ").title(), "databaseSchema": schema_fqn, "description": f"Managed Lakehouse Table '{table_name}'", "columns": [ {"name": "id", "dataType": "BIGINT", "description": "Primary Key"}, {"name": "data", "dataType": "JSON", "description": "Record Payload"}, {"name": "updated_at", "dataType": "TIMESTAMP", "description": "Timestamp"} ] } req = urllib.request.Request( f"{host}/api/v1/tables", data=json.dumps(payload).encode("utf-8"), headers={"Authorization": f"Bearer {token}", "Content-Type": "application/json"}, method="PUT" ) with urllib.request.urlopen(req, timeout=4) as resp: data = json.loads(resp.read().decode("utf-8")) return data.get("id") except Exception: return None