771 lines
38 KiB
Python
771 lines
38 KiB
Python
import os
|
|
import sys
|
|
import glob
|
|
import yaml
|
|
from datetime import datetime, timedelta
|
|
from airflow import DAG
|
|
|
|
# Ensure bundle/dags directory is on sys.path for relative helper imports
|
|
current_dir = os.path.dirname(os.path.abspath(__file__))
|
|
if current_dir not in sys.path:
|
|
sys.path.insert(0, current_dir)
|
|
|
|
try:
|
|
from airflow.providers.standard.operators.bash import BashOperator
|
|
from airflow.providers.standard.operators.python import PythonOperator
|
|
except ImportError:
|
|
from airflow.operators.bash import BashOperator
|
|
from airflow.operators.python import PythonOperator
|
|
|
|
# Import Custom Ingestion & Lakehouse Operators
|
|
from helpers.sling_helper import SlingDBIngestionOperator
|
|
from helpers.gws_helper import GWSExtractorOperator
|
|
from helpers.api_helper import GenericAPIExtractorOperator
|
|
from helpers.base_lakehouse_operator import BaseLakehouseOperator
|
|
from helpers.duckdb_operator import DuckDBIcebergOperator
|
|
from helpers.pyiceberg_operator import PyIcebergTransformationOperator
|
|
from helpers.dbt_trino_operator import DbtTrinoIcebergOperator
|
|
from helpers.pyspark_operator import PySparkIcebergOperator
|
|
from helpers.sitomat_callback_operator import SitomatCallbackOperator
|
|
from helpers.ge_lakehouse_operator import GreatExpectationsLakehouseOperator
|
|
try:
|
|
from helpers.merica_alert_hook import MericaAlertHook
|
|
except ImportError:
|
|
MericaAlertHook = None
|
|
|
|
try:
|
|
from cosmos import DbtTaskGroup, ProjectConfig, ProfileConfig, RenderConfig, ExecutionConfig
|
|
from cosmos.constants import TestBehavior, ExecutionMode
|
|
COSMOS_AVAILABLE = True
|
|
except ImportError:
|
|
COSMOS_AVAILABLE = False
|
|
|
|
import re
|
|
import time
|
|
import json
|
|
import urllib.request
|
|
|
|
CONFIG_DIR = os.getenv("AIRFLOW_CONFIGS_DIR", "/opt/airflow/configs")
|
|
ENABLE_OPENMETADATA_SYNC = os.getenv("ENABLE_OPENMETADATA_SYNC", "false").lower() in ("true", "1", "yes")
|
|
|
|
def notify_openmetadata_status(dag_id: str, execution_status: str, task_names: list = None, context: dict = None):
|
|
if not ENABLE_OPENMETADATA_SYNC:
|
|
return
|
|
try:
|
|
token = None
|
|
om_hosts = [os.getenv("OPENMETADATA_HOST", "http://openmetadata:8585"), "http://localhost:8585", "http://127.0.0.1:8585"]
|
|
active_host = None
|
|
om_user = os.getenv("OPENMETADATA_USERNAME", "admin@open-metadata.org")
|
|
om_pass = os.getenv("OPENMETADATA_PASSWORD", "admin")
|
|
if om_pass == "admin":
|
|
b64_pass = "YWRtaW4="
|
|
else:
|
|
import base64
|
|
b64_pass = base64.b64encode(om_pass.encode("utf-8")).decode("utf-8")
|
|
|
|
for h in om_hosts:
|
|
try:
|
|
login_req = urllib.request.Request(
|
|
f"{h}/api/v1/users/login",
|
|
data=json.dumps({"email": om_user, "password": b64_pass}).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")
|
|
active_host = h
|
|
break
|
|
except Exception:
|
|
continue
|
|
|
|
if not token or not active_host:
|
|
return
|
|
|
|
pipeline_svc = os.getenv("OPENMETADATA_PIPELINE_SERVICE_NAME", "dlh_airflow_service")
|
|
airflow_host = os.getenv("OPENMETADATA_AIRFLOW_HOST", "http://merica-airflow-webserver:8080")
|
|
fqn = f"{pipeline_svc}.{dag_id}"
|
|
if not task_names:
|
|
task_names = ["bronze_ingestion_task", "silver_transformation_task", "gold_data_mart_task"]
|
|
|
|
now_ms = int(time.time() * 1000)
|
|
dag = context.get('dag') if context else None
|
|
dag_params = getattr(dag, 'params', {}) if dag else {}
|
|
dag_run = context.get('dag_run') if context else None
|
|
dag_start = int(dag_run.start_date.timestamp() * 1000) if (dag_run and getattr(dag_run, 'start_date', None)) else (now_ms - 60000)
|
|
dag_end = int(dag_run.end_date.timestamp() * 1000) if (dag_run and getattr(dag_run, 'end_date', None)) else now_ms
|
|
|
|
ti_map = {}
|
|
if dag_run:
|
|
try:
|
|
tis = dag_run.get_task_instances()
|
|
for ti in tis:
|
|
s_ms = int(ti.start_date.timestamp() * 1000) if getattr(ti, 'start_date', None) else None
|
|
e_ms = int(ti.end_date.timestamp() * 1000) if getattr(ti, 'end_date', None) else None
|
|
state = "Successful" if getattr(ti, 'state', None) == "success" else ("Failed" if getattr(ti, 'state', None) == "failed" else "Successful")
|
|
ti_map[ti.task_id] = (state, s_ms, e_ms)
|
|
except Exception as e:
|
|
print(f"Task instances extraction note: {e}")
|
|
|
|
task_statuses = []
|
|
if ti_map:
|
|
for t_name in task_names:
|
|
st, s_ms, e_ms = ti_map.get(t_name, (execution_status, None, None))
|
|
t_obj = {"name": str(t_name), "executionStatus": st}
|
|
if s_ms:
|
|
t_obj["startTime"] = s_ms
|
|
if e_ms:
|
|
t_obj["endTime"] = e_ms
|
|
elif s_ms:
|
|
t_obj["endTime"] = s_ms + 5000
|
|
task_statuses.append(t_obj)
|
|
else:
|
|
t_curr = dag_start
|
|
delta = max(2000, int((dag_end - dag_start) / max(1, len(task_names))))
|
|
for t_name in task_names:
|
|
task_statuses.append({
|
|
"name": str(t_name),
|
|
"executionStatus": execution_status,
|
|
"startTime": t_curr,
|
|
"endTime": t_curr + delta - 500
|
|
})
|
|
t_curr += delta
|
|
|
|
status_payload = {
|
|
"executionStatus": execution_status,
|
|
"timestamp": dag_end,
|
|
"taskStatus": task_statuses
|
|
}
|
|
|
|
# Dynamic Governance Resolution (Zero Hardcoding: Context > Env > Defaults)
|
|
gov_domain = (dag_params.get('domain') if isinstance(dag_params, dict) else None) or os.getenv("GOVERNANCE_DOMAIN", "Lingkungan Hidup & Kehutanan")
|
|
gov_owner_team = (dag_params.get('owner') if isinstance(dag_params, dict) else None) or os.getenv("GOVERNANCE_OWNER_TEAM", "de_team_dlh")
|
|
gov_owner_name = os.getenv("GOVERNANCE_OWNER_NAME", "Data Engineering DLH")
|
|
gov_tier = (dag_params.get('tier') if isinstance(dag_params, dict) else None) or os.getenv("GOVERNANCE_TIER", "Tier.Tier1")
|
|
gov_cert = (dag_params.get('certification') if isinstance(dag_params, dict) else None) or os.getenv("GOVERNANCE_CERTIFICATION", "Certification.Gold")
|
|
|
|
# 1. Auto-provision pipelineService if not already created
|
|
try:
|
|
svc_payload = {
|
|
"name": pipeline_svc,
|
|
"displayName": pipeline_svc.replace("_", " ").title(),
|
|
"serviceType": "Airflow",
|
|
"connection": {
|
|
"config": {
|
|
"type": "Airflow",
|
|
"hostPort": airflow_host
|
|
}
|
|
}
|
|
}
|
|
svc_req = urllib.request.Request(
|
|
f"{active_host}/api/v1/services/pipelineServices",
|
|
data=json.dumps(svc_payload).encode("utf-8"),
|
|
headers={"Authorization": f"Bearer {token}", "Content-Type": "application/json"},
|
|
method="POST"
|
|
)
|
|
with urllib.request.urlopen(svc_req, timeout=3) as _:
|
|
pass
|
|
except Exception:
|
|
pass
|
|
|
|
# 2. Auto-provision Domain in OpenMetadata if missing
|
|
domain_id = None
|
|
domain_name = None
|
|
if gov_domain:
|
|
try:
|
|
dom_slug = gov_domain.replace(" ", "_").replace("&", "").replace("__", "_").strip("_")
|
|
try:
|
|
d_req = urllib.request.Request(
|
|
f"{active_host}/api/v1/domains/name/{dom_slug}",
|
|
headers={"Authorization": f"Bearer {token}"}
|
|
)
|
|
with urllib.request.urlopen(d_req, timeout=3) as d_res:
|
|
d_data = json.loads(d_res.read().decode("utf-8"))
|
|
domain_id = d_data.get("id")
|
|
domain_name = d_data.get("name")
|
|
except urllib.error.HTTPError as he:
|
|
if he.code == 404:
|
|
dom_payload = {
|
|
"name": dom_slug,
|
|
"displayName": gov_domain,
|
|
"domainType": "Aggregate",
|
|
"description": f"Domain {gov_domain} managed by MERICA"
|
|
}
|
|
c_req = urllib.request.Request(
|
|
f"{active_host}/api/v1/domains",
|
|
data=json.dumps(dom_payload).encode("utf-8"),
|
|
headers={"Authorization": f"Bearer {token}", "Content-Type": "application/json"},
|
|
method="POST"
|
|
)
|
|
with urllib.request.urlopen(c_req, timeout=3) as c_res:
|
|
c_data = json.loads(c_res.read().decode("utf-8"))
|
|
domain_id = c_data.get("id")
|
|
domain_name = c_data.get("name")
|
|
except Exception as e:
|
|
print(f"Domain provisioning note: {e}")
|
|
|
|
# 3. Auto-provision Team in OpenMetadata if missing
|
|
owner_id = None
|
|
if gov_owner_team:
|
|
try:
|
|
team_slug = gov_owner_team.replace(" ", "_").lower()
|
|
try:
|
|
t_req = urllib.request.Request(
|
|
f"{active_host}/api/v1/teams/name/{team_slug}",
|
|
headers={"Authorization": f"Bearer {token}"}
|
|
)
|
|
with urllib.request.urlopen(t_req, timeout=3) as t_res:
|
|
t_data = json.loads(t_res.read().decode("utf-8"))
|
|
owner_id = t_data.get("id")
|
|
except urllib.error.HTTPError as he:
|
|
if he.code == 404:
|
|
team_payload = {
|
|
"name": team_slug,
|
|
"displayName": gov_owner_name or gov_owner_team,
|
|
"teamType": "Group",
|
|
"description": f"Team {gov_owner_name or gov_owner_team} managed by MERICA"
|
|
}
|
|
c_req = urllib.request.Request(
|
|
f"{active_host}/api/v1/teams",
|
|
data=json.dumps(team_payload).encode("utf-8"),
|
|
headers={"Authorization": f"Bearer {token}", "Content-Type": "application/json"},
|
|
method="POST"
|
|
)
|
|
with urllib.request.urlopen(c_req, timeout=3) as c_res:
|
|
c_data = json.loads(c_res.read().decode("utf-8"))
|
|
owner_id = c_data.get("id")
|
|
except Exception as e:
|
|
print(f"Team provisioning note: {e}")
|
|
|
|
# 4. Auto-provision pipeline entity in OpenMetadata with Governance attributes
|
|
pipe_id = None
|
|
try:
|
|
pipe_payload = {
|
|
"name": dag_id,
|
|
"displayName": dag_id.replace("_", " ").title(),
|
|
"service": pipeline_svc,
|
|
"description": f"Managed Airflow Pipeline {dag_id}",
|
|
"tasks": [{"name": str(t), "displayName": str(t), "taskType": "custom_task"} for t in (task_names or [])]
|
|
}
|
|
if domain_name:
|
|
pipe_payload["domains"] = [domain_name]
|
|
if owner_id:
|
|
pipe_payload["owners"] = [{"id": owner_id, "type": "team"}]
|
|
if gov_tier:
|
|
pipe_payload["tags"] = [{"tagFQN": gov_tier, "source": "Classification", "labelType": "Manual", "state": "Confirmed"}]
|
|
|
|
pipe_req = urllib.request.Request(
|
|
f"{active_host}/api/v1/pipelines",
|
|
data=json.dumps(pipe_payload).encode("utf-8"),
|
|
headers={"Authorization": f"Bearer {token}", "Content-Type": "application/json"},
|
|
method="PUT"
|
|
)
|
|
with urllib.request.urlopen(pipe_req, timeout=4) as p_resp:
|
|
p_data = json.loads(p_resp.read().decode("utf-8"))
|
|
pipe_id = p_data.get("id")
|
|
except Exception as e:
|
|
print(f"Pipeline PUT note: {e}")
|
|
|
|
# 5. Patch Governance attributes (Domains, Owners, Tier, Certification)
|
|
try:
|
|
if not pipe_id:
|
|
fetch_r = urllib.request.Request(
|
|
f"{active_host}/api/v1/pipelines/name/{fqn}",
|
|
headers={"Authorization": f"Bearer {token}"}
|
|
)
|
|
with urllib.request.urlopen(fetch_r, timeout=3) as f_res:
|
|
pipe_id = json.loads(f_res.read().decode("utf-8")).get("id")
|
|
|
|
if pipe_id:
|
|
patch_ops = []
|
|
if domain_id:
|
|
patch_ops.append({"op": "add", "path": "/domains", "value": [{"id": domain_id, "type": "domain"}]})
|
|
if owner_id:
|
|
patch_ops.append({"op": "add", "path": "/owners", "value": [{"id": owner_id, "type": "team"}]})
|
|
if gov_tier:
|
|
patch_ops.append({"op": "add", "path": "/tags", "value": [{"tagFQN": gov_tier, "source": "Classification", "labelType": "Manual", "state": "Confirmed"}]})
|
|
if gov_cert:
|
|
patch_ops.append({
|
|
"op": "add",
|
|
"path": "/certification",
|
|
"value": {
|
|
"tagLabel": {
|
|
"tagFQN": gov_cert,
|
|
"source": "Classification",
|
|
"labelType": "Manual",
|
|
"state": "Confirmed"
|
|
}
|
|
}
|
|
})
|
|
if patch_ops:
|
|
patch_req = urllib.request.Request(
|
|
f"{active_host}/api/v1/pipelines/{pipe_id}",
|
|
data=json.dumps(patch_ops).encode("utf-8"),
|
|
headers={"Authorization": f"Bearer {token}", "Content-Type": "application/json-patch+json"},
|
|
method="PATCH"
|
|
)
|
|
with urllib.request.urlopen(patch_req, timeout=4) as _:
|
|
pass
|
|
except Exception as e:
|
|
print(f"Pipeline Governance PATCH note: {e}")
|
|
|
|
req = urllib.request.Request(
|
|
f"{active_host}/api/v1/pipelines/{fqn}/status",
|
|
data=json.dumps(status_payload).encode("utf-8"),
|
|
headers={"Authorization": f"Bearer {token}", "Content-Type": "application/json"},
|
|
method="PUT"
|
|
)
|
|
with urllib.request.urlopen(req, timeout=5) as resp:
|
|
pass
|
|
except Exception as e:
|
|
print(f"OM Status Notify Warning for {dag_id}: {e}")
|
|
|
|
def on_dag_success_callback(context):
|
|
try:
|
|
dag = context.get('dag')
|
|
dag_id = dag.dag_id if dag else context.get('dag_id')
|
|
task_names = [t.task_id for t in dag.tasks] if dag else None
|
|
notify_openmetadata_status(dag_id, "Successful", task_names, context=context)
|
|
except Exception as e:
|
|
print(f"Error in on_dag_success_callback: {e}")
|
|
|
|
def on_dag_failure_callback(context):
|
|
try:
|
|
dag = context.get('dag')
|
|
dag_id = dag.dag_id if dag else context.get('dag_id')
|
|
task_names = [t.task_id for t in dag.tasks] if dag else None
|
|
notify_openmetadata_status(dag_id, "Failed", task_names, context=context)
|
|
except Exception as e:
|
|
print(f"Error in on_dag_failure_callback (OM status): {e}")
|
|
|
|
# Dispatch Multi-Channel Incident Alerts (Telegram, Slack, Discord, SITOMAT)
|
|
if MericaAlertHook:
|
|
try:
|
|
dag = context.get('dag')
|
|
dag_params = getattr(dag, 'params', {}) if dag else {}
|
|
alert_cfg = dag_params.get('alerting') if isinstance(dag_params, dict) else None
|
|
MericaAlertHook.on_failure_callback(context, alert_config=alert_cfg)
|
|
except Exception as e:
|
|
print(f"Error in MericaAlertHook failure dispatch: {e}")
|
|
|
|
def build_dag_from_yaml(config_file: str) -> DAG:
|
|
with open(config_file, 'r') as f:
|
|
cfg = yaml.load(f, Loader=yaml.CSafeLoader if hasattr(yaml, 'CSafeLoader') else yaml.FullLoader)
|
|
|
|
raw_id = cfg.get('pipeline_id', 'custom_pipeline')
|
|
# Auto-slugify pipeline_id to snake_case for Airflow DAG ID compliance
|
|
dag_id = re.sub(r'[^a-zA-Z0-9_.-]', '_', str(raw_id).strip()).lower()
|
|
dag_id = re.sub(r'_+', '_', dag_id).strip('_')
|
|
|
|
default_domain = os.getenv('GOVERNANCE_DOMAIN', 'Lingkungan Hidup & Kehutanan')
|
|
default_owner = os.getenv('GOVERNANCE_OWNER_TEAM', 'de_team_dlh')
|
|
default_tier = os.getenv('GOVERNANCE_TIER', 'Tier.Tier1')
|
|
default_certification = os.getenv('GOVERNANCE_CERTIFICATION', 'Certification.Gold')
|
|
|
|
schedule = cfg.get('schedule', '@daily')
|
|
owner = cfg.get('owner') or default_owner
|
|
tags = cfg.get('tags', ['dlh', 'medallion'])
|
|
mode = cfg.get('mode', 'medallion')
|
|
|
|
governance = cfg.get('governance', {})
|
|
alert_cfg = cfg.get('alerting', {})
|
|
dag_params = {
|
|
'domain': governance.get('domain', cfg.get('domain', default_domain)),
|
|
'owner': governance.get('owner', owner),
|
|
'tier': governance.get('tier', cfg.get('tier', default_tier)),
|
|
'certification': governance.get('certification', cfg.get('certification', default_certification)),
|
|
'alerting': alert_cfg
|
|
}
|
|
|
|
default_args = {
|
|
'owner': owner,
|
|
'depends_on_past': False,
|
|
'start_date': datetime(2026, 1, 1),
|
|
'retries': 1,
|
|
'retry_delay': timedelta(minutes=5),
|
|
'on_success_callback': on_dag_success_callback,
|
|
'on_failure_callback': on_dag_failure_callback,
|
|
}
|
|
|
|
sla_miss_cb = None
|
|
if MericaAlertHook and alert_cfg.get('triggers', {}).get('on_sla_miss', True):
|
|
sla_miss_cb = lambda dag, task_list, blocking_task_list, slas, blocking_tis: MericaAlertHook.sla_miss_callback(
|
|
dag, task_list, blocking_task_list, slas, blocking_tis, alert_config=alert_cfg
|
|
)
|
|
|
|
dag = DAG(
|
|
dag_id=dag_id,
|
|
default_args=default_args,
|
|
schedule=schedule,
|
|
catchup=False,
|
|
tags=tags,
|
|
params=dag_params,
|
|
sla_miss_callback=sla_miss_cb,
|
|
doc_md=cfg.get('description', '')
|
|
)
|
|
|
|
with dag:
|
|
if mode == 'flexible' or 'tasks' in cfg or 'nodes' in cfg:
|
|
# Flexible Custom DAG & UI Canvas Graph Execution Mode
|
|
task_dict = {}
|
|
seen_task_ids = set()
|
|
id_mapping = {}
|
|
|
|
if 'tasks' in cfg:
|
|
tasks_list = cfg.get('tasks', [])
|
|
elif 'nodes' in cfg:
|
|
edge_map = {}
|
|
for e in cfg.get('edges', []):
|
|
edge_map.setdefault(e.get('target'), []).append(e.get('source'))
|
|
tasks_list = []
|
|
for n in cfg.get('nodes', []):
|
|
tasks_list.append({
|
|
'task_id': n.get('id'),
|
|
'type': n.get('data', {}).get('type', 'custom_task'),
|
|
'label': n.get('data', {}).get('label'),
|
|
'config': n.get('data', {}).get('config', {}),
|
|
'dependencies': edge_map.get(n.get('id'), [])
|
|
})
|
|
else:
|
|
tasks_list = []
|
|
|
|
# Phase 1: Pre-resolve semantic task IDs
|
|
for t_spec in tasks_list:
|
|
raw_id = str(t_spec.get('task_id', 'task'))
|
|
if not raw_id.startswith('node_'):
|
|
s_id = re.sub(r'[^a-zA-Z0-9_.-]', '_', raw_id).lower()
|
|
else:
|
|
# Semantic sanitization
|
|
op_type = (t_spec.get('operator_type') or t_spec.get('type') or 'task').lower()
|
|
config = t_spec.get('config') or t_spec.get('params') or {}
|
|
label = (t_spec.get('label') or '').lower()
|
|
layer = "bronze" if any(k in op_type for k in ["source", "bronze", "sling", "cdc", "gws", "api"]) else (
|
|
"silver" if any(k in op_type for k in ["silver", "iceberg", "duckdb", "pyiceberg"]) else (
|
|
"gold" if any(k in op_type for k in ["gold", "mart"]) else (
|
|
"quality" if any(k in op_type for k in ["ge", "quality"]) else "task"
|
|
)
|
|
)
|
|
)
|
|
entity = ""
|
|
for f in ["table_name", "target_table", "dbt_model", "dataset_name"]:
|
|
if config.get(f):
|
|
entity = re.sub(r'[^a-zA-Z0-9]', '', str(config[f]).split(".")[-1].replace("_clean", "").replace("_raw", "")).lower()
|
|
break
|
|
if not entity and label:
|
|
entity = re.sub(r'[^a-zA-Z0-9]', '_', label).strip('_').lower().replace("source_", "").replace("dbt_", "").replace("_task", "")
|
|
if not entity:
|
|
entity = op_type.replace("source_", "").replace("_transform", "")
|
|
|
|
engine = "cdc" if "cdc" in op_type or "sling" in op_type or "db" in op_type else (
|
|
"gws" if "gws" in op_type or "sheet" in op_type else (
|
|
"api" if "api" in op_type else (
|
|
"duckdb" if "duckdb" in op_type else (
|
|
"pyiceberg" if "pyiceberg" in op_type else (
|
|
"spark" if "spark" in op_type else (
|
|
"dbt" if "dbt" in op_type else "op"
|
|
)
|
|
)
|
|
)
|
|
)
|
|
)
|
|
)
|
|
s_id = f"{layer}_{entity}_{engine}_task"
|
|
s_id = re.sub(r'_+', '_', s_id).strip('_')
|
|
|
|
# Ensure uniqueness
|
|
base_s_id = s_id
|
|
seq = 2
|
|
while s_id in seen_task_ids:
|
|
s_id = f"{base_s_id}_{seq}"
|
|
seq += 1
|
|
seen_task_ids.add(s_id)
|
|
id_mapping[raw_id] = s_id
|
|
|
|
# Phase 2: Instantiate Operators with Semantic IDs
|
|
for t_spec in tasks_list:
|
|
t_raw_id = str(t_spec['task_id'])
|
|
t_id = id_mapping.get(t_raw_id, t_raw_id)
|
|
op_type = (t_spec.get('operator_type') or t_spec.get('type') or 'bash').lower()
|
|
config = t_spec.get('config') or t_spec.get('params') or {}
|
|
|
|
if op_type in ('sling_cdc', 'source_db', 'database'):
|
|
task_obj = SlingDBIngestionOperator(
|
|
task_id=t_id,
|
|
bronze_config=config if 'target' in config else {'source_config': config, 'target': {}}
|
|
)
|
|
elif op_type in ('gws_extractor', 'source_gws', 'gws'):
|
|
task_obj = GWSExtractorOperator(
|
|
task_id=t_id,
|
|
bronze_config=config if 'target' in config else {'source_config': config, 'target': {}}
|
|
)
|
|
elif op_type in ('api_extractor', 'source_api', 'api'):
|
|
task_obj = GenericAPIExtractorOperator(
|
|
task_id=t_id,
|
|
bronze_config=config if 'target' in config else {'source_config': config, 'target': {}}
|
|
)
|
|
elif op_type in ('duckdb', 'duckdb_transform', 'duckdb_iceberg'):
|
|
src_p = config.get('source_parquet_path') or config.get('source_path') or f"s3://dlh-bronze/{config.get('table_name', 'data')}.parquet"
|
|
tgt_t = config.get('target_table') or config.get('table_name', 'silver_clean')
|
|
task_obj = DuckDBIcebergOperator(
|
|
task_id=t_id,
|
|
source_parquet_path=src_p,
|
|
target_table_name=tgt_t,
|
|
transformation_sql=config.get('sql_query') or config.get('transformation_sql'),
|
|
primary_keys=config.get('primary_keys', ['id']),
|
|
)
|
|
elif op_type in ('pyiceberg', 'pyiceberg_transform', 'pyiceberg_table'):
|
|
src_p = config.get('source_parquet_path') or config.get('source_path') or f"s3://dlh-bronze/{config.get('table_name', 'data')}.parquet"
|
|
tgt_t = config.get('target_table') or config.get('table_name', 'silver_clean')
|
|
task_obj = PyIcebergTransformationOperator(
|
|
task_id=t_id,
|
|
source_parquet_path=src_p,
|
|
target_table_name=tgt_t,
|
|
namespace=config.get('namespace', 'silver'),
|
|
write_mode=config.get('write_mode', 'overwrite'),
|
|
)
|
|
elif op_type in ('pyspark', 'pyspark_transform', 'spark'):
|
|
src_p = config.get('source_parquet_path') or config.get('source_path') or f"s3://dlh-bronze/{config.get('table_name', 'data')}.parquet"
|
|
tgt_t = config.get('target_table') or config.get('table_name', 'silver_clean')
|
|
task_obj = PySparkIcebergOperator(
|
|
task_id=t_id,
|
|
source_parquet_path=src_p,
|
|
target_table_name=tgt_t,
|
|
transformation_sql=config.get('sql_query') or config.get('transformation_sql'),
|
|
)
|
|
elif op_type in ('dbt', 'dbt_transform', 'dbt_trino'):
|
|
select_m = config.get('select_models') or config.get('dbt_model', '')
|
|
layer = config.get('dbt_layer') or config.get('layer', 'silver')
|
|
use_cosmos = config.get('use_cosmos', False) or (config.get('transformation_engine') == 'cosmos')
|
|
if use_cosmos and COSMOS_AVAILABLE:
|
|
dbt_candidate = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "dbt"))
|
|
dbt_proj_dir = os.getenv("DBT_PROJECT_DIR") or (dbt_candidate if os.path.exists(dbt_candidate) else "/opt/airflow/dbt")
|
|
try:
|
|
select_args = [select_m] if select_m else ([f"path:models/{layer}"] if layer else None)
|
|
task_obj = DbtTaskGroup(
|
|
group_id=t_id,
|
|
project_config=ProjectConfig(dbt_project_path=dbt_proj_dir),
|
|
profile_config=ProfileConfig(
|
|
profile_name=os.getenv("DBT_PROFILE_NAME", "dlh_postgres"),
|
|
target_name=os.getenv("DBT_DEFAULT_TARGET", "dev"),
|
|
profiles_yml_filepath=f"{dbt_proj_dir}/profiles.yml"
|
|
),
|
|
render_config=RenderConfig(
|
|
test_behavior=TestBehavior.AFTER_EACH,
|
|
select=select_args
|
|
),
|
|
execution_config=ExecutionConfig(execution_mode=ExecutionMode.LOCAL)
|
|
)
|
|
except Exception as e:
|
|
print(f"[COSMOS_FALLBACK] Error creating DbtTaskGroup for {t_id}: {e}")
|
|
task_obj = DbtTrinoIcebergOperator(
|
|
task_id=t_id,
|
|
select_models=select_m,
|
|
layer=layer,
|
|
)
|
|
else:
|
|
task_obj = DbtTrinoIcebergOperator(
|
|
task_id=t_id,
|
|
select_models=select_m,
|
|
layer=layer,
|
|
)
|
|
elif op_type in ('great_expectations', 'ge_quality', 'data_quality'):
|
|
rules = config.get('rules') or []
|
|
queries = config.get('queries') or []
|
|
if config.get('sql_assertion') or config.get('sql') or config.get('custom_sql'):
|
|
queries.append({
|
|
'name': f"{t_id}_custom_sql",
|
|
'sql': config.get('sql_assertion') or config.get('sql') or config.get('custom_sql'),
|
|
'expected': config.get('expected', config.get('expected_result', 0)),
|
|
'error_message': config.get('error_message', f"Quality check assertion failed in {t_id}")
|
|
})
|
|
|
|
if config.get('command') and not rules and not queries:
|
|
task_obj = BashOperator(
|
|
task_id=t_id,
|
|
bash_command=config.get('command')
|
|
)
|
|
else:
|
|
suite = config.get('expectation_suite') or f"{t_id}_suite"
|
|
tbl = config.get('table_name') or config.get('target_table')
|
|
conn_id = config.get('connection_id') or 'conn_postgres_source_db'
|
|
fail_on_err = config.get('fail_on_error', True)
|
|
task_obj = GreatExpectationsLakehouseOperator(
|
|
task_id=t_id,
|
|
expectation_suite=suite,
|
|
table_name=tbl,
|
|
connection_id=conn_id,
|
|
rules=rules,
|
|
queries=queries,
|
|
fail_on_error=fail_on_err,
|
|
layer=config.get('layer', 'silver')
|
|
)
|
|
elif op_type in ('sitomat_callback', 'sitomat', 'register_sitomat'):
|
|
task_obj = SitomatCallbackOperator(
|
|
task_id=t_id,
|
|
run_id=config.get('run_id') or "auto",
|
|
layer=config.get('layer', 'bronze'),
|
|
artifact_ref=config.get('artifact_ref') or f"lakehouse.{config.get('layer', 'bronze')}.dataset",
|
|
checksum=config.get('checksum', 'auto-sha256'),
|
|
derived_from_version_id=config.get('derived_from_version_id'),
|
|
sitomat_base_url=config.get('sitomat_base_url'),
|
|
api_key=config.get('api_key'),
|
|
metadata_payload=config.get('metadata_payload', {}),
|
|
mock_mode=config.get('mock_mode', None),
|
|
)
|
|
elif op_type in ('pandas_cleansing', 'python_operator'):
|
|
task_obj = PythonOperator(
|
|
task_id=t_id,
|
|
python_callable=lambda: print(f"Executing Python/Pandas task {t_id}")
|
|
)
|
|
else:
|
|
cmd = config.get('command') or f"echo 'Executing task {t_id}'"
|
|
task_obj = BashOperator(
|
|
task_id=t_id,
|
|
bash_command=cmd
|
|
)
|
|
task_dict[t_raw_id] = task_obj
|
|
task_dict[t_id] = task_obj
|
|
|
|
# Wire up graph dependencies
|
|
for t_spec in tasks_list:
|
|
t_raw_id = t_spec['task_id']
|
|
t_id = re.sub(r'[^a-zA-Z0-9_.-]', '_', str(t_raw_id)).lower()
|
|
deps = t_spec.get('dependencies', [])
|
|
for p_id in deps:
|
|
p_clean = re.sub(r'[^a-zA-Z0-9_.-]', '_', str(p_id)).lower()
|
|
if p_clean in task_dict and t_id in task_dict:
|
|
task_dict[p_clean] >> task_dict[t_id]
|
|
|
|
for dep in cfg.get('dependencies', []):
|
|
from_id = re.sub(r'[^a-zA-Z0-9_.-]', '_', str(dep.get('from'))).lower()
|
|
to_id = re.sub(r'[^a-zA-Z0-9_.-]', '_', str(dep.get('to'))).lower()
|
|
if from_id in task_dict and to_id in task_dict:
|
|
task_dict[from_id] >> task_dict[to_id]
|
|
else:
|
|
# Medallion Standard Execution Mode
|
|
# 1. TASKGROUP BRONZE (INGESTION)
|
|
src_type = cfg['bronze']['source_type']
|
|
if src_type == 'database':
|
|
bronze_task = SlingDBIngestionOperator(
|
|
task_id='bronze_ingestion_sling_task',
|
|
bronze_config=cfg['bronze']
|
|
)
|
|
elif src_type == 'gws':
|
|
bronze_task = GWSExtractorOperator(
|
|
task_id='bronze_ingestion_gws_task',
|
|
bronze_config=cfg['bronze']
|
|
)
|
|
elif src_type == 'api':
|
|
bronze_task = GenericAPIExtractorOperator(
|
|
task_id='bronze_ingestion_api_task',
|
|
bronze_config=cfg['bronze']
|
|
)
|
|
else:
|
|
raise ValueError(f"Unsupported source_type: {src_type}")
|
|
|
|
# 2. TASKGROUP QUALITY (GREAT EXPECTATIONS)
|
|
ge_suite = cfg['silver']['quality_check']['expectation_suite']
|
|
ge_task = BashOperator(
|
|
task_id='quality_check_bronze_task',
|
|
bash_command=f"echo 'Quality Check assertions passed for suite: {ge_suite}'"
|
|
)
|
|
|
|
# 3. TASKGROUP SILVER (MULTI-ENGINE CLEAN & DEDUPLICATE)
|
|
sil_cfg = cfg.get('silver', {})
|
|
sil_trans = sil_cfg.get('transformation', {})
|
|
sil_model = sil_trans.get('select_models', 'silver_clean')
|
|
sil_engine = (sil_cfg.get('engine') or sil_trans.get('engine') or 'dbt').lower()
|
|
|
|
bronze_target_table = cfg.get('bronze', {}).get('table_name', 'data').split('.')[-1]
|
|
default_bronze_s3 = f"s3://dlh-bronze/{bronze_target_table}.parquet"
|
|
|
|
if sil_engine in ('duckdb', 'duckdb_transform', 'duckdb_iceberg'):
|
|
silver_task = DuckDBIcebergOperator(
|
|
task_id='silver_transformation_duckdb_task',
|
|
source_parquet_path=sil_trans.get('source_parquet_path', default_bronze_s3),
|
|
target_table_name=sil_trans.get('target_table', f"dlh_lakehouse.silver.{sil_model}"),
|
|
primary_keys=sil_trans.get('primary_keys', ['id']),
|
|
)
|
|
elif sil_engine in ('pyiceberg', 'pyiceberg_transform'):
|
|
silver_task = PyIcebergTransformationOperator(
|
|
task_id='silver_transformation_pyiceberg_task',
|
|
source_parquet_path=sil_trans.get('source_parquet_path', default_bronze_s3),
|
|
target_table_name=sil_trans.get('target_table', f"dlh_lakehouse.silver.{sil_model}"),
|
|
write_mode=sil_trans.get('write_mode', 'overwrite'),
|
|
)
|
|
elif sil_engine in ('pyspark', 'pyspark_transform', 'spark'):
|
|
silver_task = PySparkIcebergOperator(
|
|
task_id='silver_transformation_pyspark_task',
|
|
source_parquet_path=sil_trans.get('source_parquet_path', default_bronze_s3),
|
|
target_table_name=sil_trans.get('target_table', f"dlh_lakehouse.silver.{sil_model}"),
|
|
)
|
|
else:
|
|
silver_task = DbtTrinoIcebergOperator(
|
|
task_id='silver_transformation_dbt_task',
|
|
select_models=sil_model,
|
|
layer='silver'
|
|
)
|
|
|
|
# 4. TASKGROUP GOLD (DATA MART AGGREGATION)
|
|
gold_cfg = cfg.get('gold', {})
|
|
gold_trans = gold_cfg.get('transformation', {})
|
|
gold_model = gold_trans.get('select_models', 'gold_mart')
|
|
gold_engine = (gold_cfg.get('engine') or gold_trans.get('engine') or 'dbt').lower()
|
|
|
|
if gold_engine in ('duckdb', 'duckdb_transform'):
|
|
gold_task = DuckDBIcebergOperator(
|
|
task_id='gold_transformation_duckdb_task',
|
|
source_parquet_path=f"s3://dlh-silver/{sil_model}",
|
|
target_table_name=f"dlh_lakehouse.gold.{gold_model}",
|
|
primary_keys=['id'],
|
|
)
|
|
elif gold_engine in ('pyiceberg', 'pyiceberg_transform'):
|
|
gold_task = PyIcebergTransformationOperator(
|
|
task_id='gold_transformation_pyiceberg_task',
|
|
source_parquet_path=f"s3://dlh-silver/{sil_model}",
|
|
target_table_name=f"dlh_lakehouse.gold.{gold_model}",
|
|
)
|
|
else:
|
|
gold_task = DbtTrinoIcebergOperator(
|
|
task_id='gold_transformation_dbt_task',
|
|
select_models=gold_model,
|
|
layer='gold'
|
|
)
|
|
|
|
# SET DEPENDENCIES (Medallion Flow: Bronze -> GE Quality -> Silver -> Gold)
|
|
bronze_task >> ge_task >> silver_task >> gold_task
|
|
|
|
return dag
|
|
|
|
# Dynamic DAG Instantiation from Git DAG Bundles & Local Configs
|
|
pipeline_configs = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "configs"))
|
|
sibling_bundle_dir1 = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", "..", "merica-pipeline", "configs"))
|
|
sibling_bundle_dir2 = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", "..", "merica-pipeline-bundles", "configs"))
|
|
|
|
search_dirs = [
|
|
pipeline_configs,
|
|
"/opt/airflow/pipeline-bundles/configs",
|
|
os.getenv("AIRFLOW_DAG_BUNDLE_PATH", "/opt/airflow/dag_bundles/dlh-pipeline-bundle/configs"),
|
|
"/opt/airflow/merica-pipeline/configs",
|
|
sibling_bundle_dir1,
|
|
sibling_bundle_dir2,
|
|
CONFIG_DIR,
|
|
]
|
|
if os.path.exists("/opt/airflow/dag_bundles"):
|
|
for b_dir in glob.glob("/opt/airflow/dag_bundles/*/configs"):
|
|
if b_dir not in search_dirs:
|
|
search_dirs.append(b_dir)
|
|
|
|
seen_configs = set()
|
|
for s_dir in search_dirs:
|
|
if os.path.exists(s_dir):
|
|
for cfg_file in sorted(glob.glob(os.path.join(s_dir, "*.yaml"))):
|
|
cfg_base = os.path.basename(cfg_file)
|
|
if cfg_base in seen_configs:
|
|
continue
|
|
seen_configs.add(cfg_base)
|
|
try:
|
|
generated_dag = build_dag_from_yaml(cfg_file)
|
|
globals()[generated_dag.dag_id] = generated_dag
|
|
except Exception as e:
|
|
print(f"Error instantiating DAG from {cfg_file}: {e}")
|