init

airflow/TestDags/IICS_Airflow_Demo.py

@@ -0,0 +1,142 @@
+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.dummy_operator import DummyOperator
+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:
+    # Replace invalid characters with underscores
+    sanitized = re.sub(r'[^a-zA-Z0-9_]+', '_', task_id)
+    # Ensure task_id starts with a letter or number
+    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="https://dm-us.informaticacloud.com/ma/api/v2/user/login")
+
+# Load task names from Airflow Variables
+CDI_task_name = json.loads(Variable.get("CDI_task_name", default_var='["Task_Date_Dim", "Task_Items", "Task_Store_Sales"]'))
+CDI_E_task_name = json.loads(Variable.get("CDI_E_task_name", default_var='["Task_Total_Store_Sales_IWDEMO"]'))
+
+# 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),
+   'schedule': '@daily'
+}
+
+# API logic
+def get_session_id(un, pw):
+    session_id = ''
+    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:
+        session_id = r.json()["icSessionId"]
+        server_url = r.json()["serverUrl"]
+    else:
+        print('API call failed:', r.status_code)
+        print(r.text)
+        sys.exit(1)
+    return session_id, server_url
+
+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("Job", taskname, "started successfully")
+        return response_content['taskId'], response_content['runId'], response_content['taskName']
+    else:
+        print('Job failed to start:', r.status_code)
+        print(r.text)
+
+def get_status(server_url, session_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():
+            return obj['taskId'], obj['executionState'], obj['taskName'], obj['runId']
+    else:
+        print('Failed to get status:', r.status_code)
+        print(r.text)
+
+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)
+        task_status = get_status(server_url, session_id)
+        if not task_status or task_status[1] not in {"RUNNING", "INITIALIZED", "STOPPING", "QUEUED"}:
+            # Fetch log
+            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
+
+# Define DAG
+dag = DAG(
+   'IICS_Airflow_Demo',
+   default_args=default_args,
+   description='A Sample IICS Airflow DAG',
+   schedule_interval='@daily',
+   catchup=False
+)
+
+# Task group 1: CDI Tasks
+cdi_start = DummyOperator(task_id='cdi_start', dag=dag)
+cdi_end = DummyOperator(task_id='cdi_end', dag=dag)
+
+for i in CDI_task_name:
+    safe_task_id = 'IICS_CDI_' + sanitize_task_id(i)
+    print(f"Creating task: {safe_task_id} for original task name: {i}")
+    cdi_task = PythonOperator(
+        task_id=safe_task_id,
+        python_callable=execute_task,
+        op_kwargs={'task_name': i},
+        dag=dag
+    )
+    cdi_start >> cdi_task >> cdi_end
+
+# Task group 2: CDI_E Tasks
+for j in CDI_E_task_name:
+    safe_task_id = 'IICS_CDI_E_' + sanitize_task_id(j)
+    print(f"Creating E task: {safe_task_id} for original task name: {j}")
+    cdi_e_task = PythonOperator(
+        task_id=safe_task_id,
+        python_callable=execute_task,
+        op_kwargs={'task_name': j},
+        dag=dag
+    )
+    cdi_end >> cdi_e_task