import logging from contextlib import closing import pendulum import utils.lark as lark from airflow.datasets import Dataset from airflow.decorators import dag, task from airflow.exceptions import AirflowFailException from airflow.models.baseoperator import chain from airflow.models.taskinstance import TaskInstance from airflow.providers.mysql.hooks.mysql import MySqlHook from dotenv import dotenv_values from utils.database import get_records_from_db from utils.github import github_action config = dotenv_values() ENV = config["ENV"] LARK_ALERT_CHAT_ID = config["LARK_ALERT_CHAT_ID"] MYSQL_CONN_ID = "autonomous-serverless-instance" DB_INSTANCE = f"auto_deploy_dedicated_{ENV}" AWS_ROLE = "auto-deploy-cli" GIT_REPO = "tidbcloud/auto-deploy" GIT_REF = config["AUTO_DEPLOY_REPO_BRANCH"] # GIT_REPO="tidbcloud/autonomous-airflow" task_logger = logging.getLogger("airflow.task") @task.short_circuit def get_task_list(ti: TaskInstance): """ fetch destroy regional tasks from regional_destroy_tasks """ db = MySqlHook(mysql_conn_id=MYSQL_CONN_ID) tasks = get_records_from_db( db, f"select * from {DB_INSTANCE}.regional_destroy_tasks", ) for t in tasks: if not t["logical_region_name"]: t["logical_region_name"] = t["region_name"] # concurrency is not allowed at this time if len(tasks) > 1: task_logger.error(f"Too many tasks: {len(tasks)}") raise AirflowFailException("Too many tasks") elif len(tasks) == 1: ti.xcom_push(key="task", value=tasks[0]) return True else: return False @task def notify_task_begin(ti: TaskInstance): t = ti.xcom_pull(key="task") if not t or not t["trigger"]: return lark.send_card_message( lark.build_status_card( lark.Color.BLUE, f"start destroying region {t['logical_region_name']}", { "☁️ Cloud": t["provider"], "📍 Region": t["region_name"], "🪪 Name": t["logical_region_name"], }, { "Regional Service": "deleting", "Cluster Service": "deleting", "BR Service": "deleting", "CDC Service": "deleting", }, ), email=t["trigger"], ) @task def notify_task_end(ti: TaskInstance): t = ti.xcom_pull(key="task") if not t or not t["trigger"]: return lark.send_card_message( lark.build_status_card( lark.Color.GREEN, f"finish destroying region {t['logical_region_name']}", { "☁️ Cloud": t["provider"], "📍 Region": t["region_name"], "🪪 Name": t["logical_region_name"], "🏁 Status": "SUCCESS", }, { "Regional Service": "deleted", "Cluster Service": "deleted", "BR Service": "deleted", "CDC Service": "deleted", }, ), email=t["trigger"], ) @task def get_region_meta(ti: TaskInstance): """ filter task, if task is invalid update task result """ task = ti.xcom_pull(key="task") db = MySqlHook(mysql_conn_id=MYSQL_CONN_ID) meta = get_records_from_db( db, f""" select meta_uri,is_ng from {DB_INSTANCE}.regional_services where provider = %(provider)s and region_name = %(region_name)s and logical_region_name = %(logical_region_name)s """, task, ) if len(meta) == 0: task_logger.error("no meta for %(provider)s %(logical_region_name)s", task) raise AirflowFailException( f"no meta for {task['provider']} {task['logical_region_name']}" ) else: task_logger.info("get meta-uri: %(meta_uri)s", meta[0]) task["meta_uri"] = meta[0]["meta_uri"] task["is_ng"] = meta[0]["is_ng"] task["status"] = "success" ti.xcom_push(key="task", value=task) @task( outlets=[ Dataset(f"mysql://{MYSQL_CONN_ID}/{DB_INSTANCE}/regional_destroy_results"), Dataset(f"mysql://{MYSQL_CONN_ID}/{DB_INSTANCE}/regional_services"), ] ) def write_task_results(ti: TaskInstance = None): task = ti.xcom_pull(key="task") db = MySqlHook(mysql_conn_id=MYSQL_CONN_ID) with closing(db.get_conn()) as conn, closing(conn.cursor()) as cur: cur.execute( f""" insert into {DB_INSTANCE}.regional_destroy_results (request_id, status, provider, region_name, logical_region_name, trace_id) values (%(id)s, %(status)s, %(provider)s, %(region_name)s, %(logical_region_name)s, %(trace_id)s) """, task, ) for table in [ "regional_services", "cluster_services", "br_services", "cdc_services", "default_sku_services", ]: cur.execute( f""" delete from {DB_INSTANCE}.{table} where provider = %(provider)s and region_name = %(region_name)s and logical_region_name = %(logical_region_name)s """, task, ) conn.commit() @dag( dag_id=f"destroy_dedicated_region_{ENV}", start_date=pendulum.datetime(2024, 10, 25, tz="UTC"), schedule=[ Dataset(f"mysql://{MYSQL_CONN_ID}/{DB_INSTANCE}/regional_destroy_requests") ], tags=["destroy", "dedicated", "region", ENV], max_active_runs=1, default_args={ "on_failure_callback": lambda context: lark.send_error_message_handler( LARK_ALERT_CHAT_ID, context ) }, ) def destroy_dedicated_region(): def git_action_inpuf_fn_aws(ti: TaskInstance): task = ti.xcom_pull(key="task") if task["provider"] != "aws": return [] return [ { "tidb-cloud-instance": ENV, "region-name": task["logical_region_name"], "aws-region": task["region_name"], "meta-uri": task["meta_uri"], "deploy-type": ( "DEPLOY_TYPE_NEXTGEN" if task.get("is_ng") else "DEPLOY_TYPE_DEDICATED" ), "role-arn": f"arn:aws:iam::{task['m_account']}:role/{AWS_ROLE}", } ] def git_action_inpuf_fn_azure(ti: TaskInstance): task = ti.xcom_pull(key="task") if task["provider"] != "azure": return [] return [ { "tidb-cloud-instance": ENV, "logical-region-name": task["logical_region_name"], "cloud-region": task["region_name"], "meta-uri": task["meta_uri"], # aws_account:tenant_id:subscription_id:client_id "tenant-id": task["m_account"].split(":")[1], "subscription-id": task["m_account"].split(":")[2], "client-id": task["m_account"].split(":")[3], } ] chain( get_task_list(), notify_task_begin(), get_region_meta(), github_action( group_id="disable_dedicated_region", repo=GIT_REPO, workflow="update-region.yml", ref=GIT_REF, input_fn=lambda ti: [ { "env": ENV, "logical-name": ti.xcom_pull(key="task").get("logical_region_name"), "deploy-type": ( "DEPLOY_TYPE_NEXTGEN" if ti.xcom_pull(key="task").get("is_ng") else "DEPLOY_TYPE_DEDICATED" ), "options": '{"state": "INACTIVE"}', }, ], ), [ github_action( group_id="destroy_aws_dedicated_region", repo=GIT_REPO, workflow="destroy-aws-region-k8s.yml", ref=GIT_REF, input_fn=git_action_inpuf_fn_aws, ), github_action( group_id="destroy_azure_dedicated_region", repo=GIT_REPO, workflow="destroy-azure-region.yml", ref=GIT_REF, input_fn=git_action_inpuf_fn_azure, ), ], write_task_results(), notify_task_end(), ) destroy_dedicated_region()