540 lines
26 KiB
Python
540 lines
26 KiB
Python
|
import os
|
||
|
|
import glob
|
||
|
|
import yaml
|
||
|
|
from datetime import datetime, timedelta
|
||
|
|
from airflow import DAG
|
||
|
|
|
||
|
|
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
|
||
|
|
|
||
|
|
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
|
||
|
|
for h in om_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")
|
||
|
|
active_host = h
|
||
|
|
break
|
||
|
|
except Exception:
|
||
|
|
continue
|
||
|
|
|
||
|
|
if not token or not active_host:
|
||
|
|
return
|
||
|
|
|
||
|
|
fqn = f"dlh_airflow_service.{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_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
|
||
|
|
}
|
||
|
|
|
||
|
|
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('_')
|
||
|
|
|
||
|
|
schedule = cfg.get('schedule', '@daily')
|
||
|
|
owner = cfg.get('owner', 'de_team_dlh')
|
||
|
|
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', 'Lingkungan Hidup & Kehutanan')),
|
||
|
|
'owner': governance.get('owner', owner),
|
||
|
|
'tier': governance.get('tier', cfg.get('tier', 'Tier.Tier2')),
|
||
|
|
'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')
|
||
|
|
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}")
|