init

airflow/TestDags/IDMC_Airflow_Test.py

@@ -0,0 +1,120 @@
+import json
+import sys
+import time
+import re
+import requests
+from datetime import datetime, timedelta
+
+from airflow import DAG
+from airflow.models import Variable
+from airflow.operators.python_operator import PythonOperator
+
+
+# from infromatic team, : connect to infromatica (akash)
+
+# Utility to make task_id Airflow-safe
+def sanitize_task_id(task_id: str) -> str:
+    sanitized = re.sub(r'[^a-zA-Z0-9_]+', '_', task_id)
+    if not re.match(r'^[a-zA-Z0-9]', sanitized):
+        sanitized = 'task_' + sanitized
+    return sanitized
+
+# Fetch parameters from Airflow Variables
+iics_username = Variable.get("iics_username")
+iics_password = Variable.get("iics_password")
+task_type = Variable.get("task_type", default_var="MTT")
+base_url = Variable.get("iics_base_url", default_var="")
+
+# Task name
+CDI_task_name = "CDI_task"
+
+# Default DAG args
+default_args = {
+   'owner': 'infa',
+   'depends_on_past': False,
+   'email': ['airflow@example.com'],
+   'email_on_failure': False,
+   'email_on_retry': False,
+   'retries': 1,
+   'retry_delay': timedelta(minutes=1),
+   'start_date': datetime.now() - timedelta(seconds=10),
+}
+
+# API logic (same as before)
+def get_session_id(un, pw):
+    data = {'@type': 'login', 'username': un, 'password': pw}
+    headers = {'Content-Type': 'application/json', 'Accept': 'application/json'}
+    r = requests.post(base_url, data=json.dumps(data), headers=headers)
+    if r.status_code == 200:
+        return r.json()["icSessionId"], r.json()["serverUrl"]
+    else:
+        print('API call failed:', r.status_code)
+        print(r.text)
+        sys.exit(1)
+
+def start_job(session_id, server_url, taskname, taskType):
+    job_start_url = server_url + "/api/v2/job"
+    headers = {'Content-Type': 'application/json', 'icSessionId': session_id, 'Accept': 'application/json'}
+    data = {'@type': 'job', 'taskName': taskname, 'taskType': taskType}
+    r = requests.post(job_start_url, data=json.dumps(data), headers=headers)
+    if r.status_code == 200:
+        response_content = r.json()
+        print(f"Job {taskname} started successfully")
+        return response_content['taskId'], response_content['runId']
+    else:
+        print('Job failed to start:', r.status_code)
+        print(r.text)
+        sys.exit(1)
+
+def get_status(server_url, session_id, task_id, run_id):
+    job_activity_url = server_url + "/api/v2/activity/activityMonitor"
+    headers = {'Content-Type': 'application/json', 'icSessionId': session_id, 'Accept': 'application/json'}
+    r = requests.get(job_activity_url, headers=headers)
+    if r.status_code == 200:
+        for obj in r.json():
+            if obj['taskId'] == task_id and obj['runId'] == run_id:
+                return obj['executionState']
+    else:
+        print('Failed to get status:', r.status_code)
+        print(r.text)
+        sys.exit(1)
+
+def execute_task(task_name):
+    session_id, server_url = get_session_id(iics_username, iics_password)
+    task_id, run_id = start_job(session_id, server_url, task_name, task_type)
+    
+    log_url = f"{server_url}/api/v2/activity/activityLog/"
+    headers = {'Content-Type': 'application/json', 'icSessionId': session_id, 'Accept': 'application/json'}
+
+    while True:
+        time.sleep(15)
+        status = get_status(server_url, session_id, task_id, run_id)
+        print(f"Task status: {status}")
+        if status not in {"RUNNING", "INITIALIZED", "STOPPING", "QUEUED"}:
+            # Fetch logs on completion
+            url = f"{log_url}?taskId={task_id}&runId={run_id}"
+            r = requests.get(url, headers=headers)
+            logs = r.json()
+            for obj in logs:
+                log_id = obj['id']
+                log_detail = requests.get(f"{log_url}{log_id}/sessionLog", headers=headers)
+                print(log_detail.text)
+            break
+
+# DAG with no schedule (manual trigger)
+dag = DAG(
+   'IDMC_Airflow_Test',
+   default_args=default_args,
+   description='Simplified DAG with one CDI task',
+   schedule_interval=None,
+   catchup=False
+)
+
+safe_task_id = sanitize_task_id(CDI_task_name)
+
+run_cdi_task = PythonOperator(
+    task_id=safe_task_id,
+    python_callable=execute_task,
+    op_kwargs={'task_name': CDI_task_name},
+    dag=dag
+)