121 lines
4.1 KiB
Python
121 lines
4.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.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
|
|
)
|