diff --git a/dags/dag_factory.py b/dags/dag_factory.py index 49b5eac..40a27bd 100644 --- a/dags/dag_factory.py +++ b/dags/dag_factory.py @@ -111,6 +111,50 @@ def notify_openmetadata_status(dag_id: str, execution_status: str, task_names: l "taskStatus": task_statuses } + # Auto-provision pipelineService if not already created + try: + svc_payload = { + "name": "dlh_airflow_service", + "displayName": "DLH Airflow Service", + "serviceType": "Airflow", + "connection": { + "config": { + "type": "Airflow", + "hostPort": "http://merica-airflow-webserver:8080" + } + } + } + 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 + + # Auto-provision pipeline entity in OpenMetadata + try: + pipe_payload = { + "name": dag_id, + "displayName": dag_id.replace("_", " ").title(), + "service": "dlh_airflow_service", + "description": f"Managed Airflow Pipeline {dag_id}", + "tasks": [{"name": str(t), "displayName": str(t), "taskType": "custom_task"} for t in (task_names or [])] + } + 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 _: + pass + except Exception: + pass + req = urllib.request.Request( f"{active_host}/api/v1/pipelines/{fqn}/status", data=json.dumps(status_payload).encode("utf-8"), diff --git a/dags/helpers/base_lakehouse_operator.py b/dags/helpers/base_lakehouse_operator.py index 1a0b1ec..5e90a1e 100644 --- a/dags/helpers/base_lakehouse_operator.py +++ b/dags/helpers/base_lakehouse_operator.py @@ -100,10 +100,14 @@ class BaseLakehouseOperator(BaseOperator): 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 = { @@ -183,3 +187,32 @@ class BaseLakehouseOperator(BaseOperator): 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