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