143 lines
5.1 KiB
Python
143 lines
5.1 KiB
Python
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
|