186 lines
7.7 KiB
Python
186 lines
7.7 KiB
Python
|
"""
|
||
|
|
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)
|
||
|
|
|
||
|
|
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 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
|