import hashlib from contextlib import closing from os.path import dirname, join, realpath import jinja2 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 deploy.dedicated.common import ( add_infra_apiserver_source_ips, get_seed_instance_type, parse_k8s_public_ips, ) 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"] META_BUCKET_AWS = config["META_BUCKET_AWS"] TF_STATE_BUCKET_AWS = config["TF_STATE_BUCKET_AWS"] META_STORAGE_ACCOUNT = config["META_STORAGE_ACCOUNT"] LARK_ALERT_CHAT_ID = config["LARK_ALERT_CHAT_ID"] AUTO_DEPLOY_REPO_BRANCH = config["AUTO_DEPLOY_REPO_BRANCH"] @task.short_circuit def get_task_list(ti: TaskInstance): """ fetch create regional tasks from regional_create_tasks """ db = MySqlHook(mysql_conn_id="autonomous-serverless-instance") tasks = get_records_from_db( db, f""" select * from auto_deploy_dedicated_{ENV}.regional_create_tasks """, ) for t in tasks: t["is_ng"] = False if not t["logical_region_name"]: m = hashlib.md5() m.update(t["trace_id"].encode()) suffix = m.hexdigest()[-6:] t["logical_region_name"] = t["region_name"] + "-" + suffix if t["provider"] == "aws": t["meta_uri"] = ( f"s3://{META_BUCKET_AWS}/dedicated-regions/{t['logical_region_name']}/meta.yaml" ) elif t["provider"] == "azure": t["meta_uri"] = ( f"https://{META_STORAGE_ACCOUNT}.blob.core.windows.net/meta/dedicated-regions/{t['logical_region_name']}/meta.yaml" ) else: raise AirflowFailException(f"unknown provider: {t['provider']}") ti.xcom_push(key="task_list", value=tasks) return len(tasks) > 0 @task def notify_task_begin(ti: TaskInstance): task_list = ti.xcom_pull(key="task_list") for t in task_list: if not t["trigger"]: continue lark.send_card_message( lark.build_status_card( lark.Color.BLUE, f"start creating region {t['logical_region_name']}", { "☁️ Cloud": t["provider"], "📍 Region": t["region_name"], "🪪 Name": t["logical_region_name"], }, { "Regional Service": "creating", "Cluster Service": "creating", "BR Service": "creating", "CDC Service": "creating", }, ), email=t["trigger"], ) @task def notify_task_end(ti: TaskInstance): task_list = ti.xcom_pull(key="task_list") for t in task_list: if not t["trigger"]: continue lark.send_card_message( lark.build_status_card( lark.Color.GREEN, f"dedicated region {t['logical_region_name']} created", { "☁️ Cloud": t["provider"], "📍 Region": t["region_name"], "🪪 Name": t["logical_region_name"], "🏁 Status": "SUCCESS", }, { "Regional Service": "created", "Cluster Service": "created", "BR Service": "created", "CDC Service": "created", }, ), email=t["trigger"], ) @task( outlets=[ Dataset( f"mysql://autonomous-serverless-instance/auto_deploy_dedicated_{ENV}/{tb}" ) for tb in [ "regional_create_result", "regional_services", "cluster_services", "br_services", "cdc_services", ] ] ) def write_task_results(ti: TaskInstance = None): task_list = ti.xcom_pull(key="task_list") public_ips = ti.xcom_pull(key="public_ips") db = MySqlHook(mysql_conn_id="autonomous-serverless-instance") with closing(db.get_conn()) as conn, closing(conn.cursor()) as cur: for t in task_list: cur.execute( f""" insert into auto_deploy_dedicated_{ENV}.regional_create_results (request_id, status, meta_uri, provider, region_name, logical_region_name, trace_id) values (%(id)s, 'success', %(meta_uri)s, %(provider)s, %(region_name)s, %(logical_region_name)s, %(trace_id)s) """, t, ) cur.execute( f""" insert into auto_deploy_dedicated_{ENV}.regional_services (provider, region_name, logical_region_name, status, meta_uri, public_ips, trace_id, requested_at) values (%(provider)s, %(region_name)s, %(logical_region_name)s, %(status)s, %(meta_uri)s, %(public_ips)s, %(trace_id)s, %(requested_at)s) """, { **t, "requested_at": t["created_at"], "status": "active", "public_ips": public_ips[str(t["id"])], }, ) cur.execute( f""" insert into auto_deploy_dedicated_{ENV}.cluster_services (provider, region_name, logical_region_name, status, trace_id, requested_at) values (%(provider)s, %(region_name)s, %(logical_region_name)s, %(status)s, %(trace_id)s, %(requested_at)s) """, { **t, "requested_at": t["created_at"], "status": "active", }, ) if ENV != "staging": cur.execute( f""" insert into auto_deploy_dedicated_{ENV}.br_services (provider, region_name, logical_region_name, status, trace_id, requested_at) values (%(provider)s, %(region_name)s, %(logical_region_name)s, %(status)s, %(trace_id)s, %(requested_at)s) """, { **t, "requested_at": t["created_at"], "status": "active", }, ) cur.execute( f""" insert into auto_deploy_dedicated_{ENV}.cdc_services (provider, region_name, logical_region_name, status, trace_id, requested_at) values (%(provider)s, %(region_name)s, %(logical_region_name)s, %(status)s, %(trace_id)s, %(requested_at)s) """, { **t, "requested_at": t["created_at"], "status": "active", }, ) conn.commit() @dag( dag_id=f"create_dedicated_region_{ENV}", start_date=pendulum.datetime(2024, 10, 25, tz="UTC"), schedule=[ Dataset( f"mysql://autonomous-serverless-instance/auto_deploy_dedicated_{ENV}/regional_create_requests" ) ], tags=["deploy", "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 create_dedicated_region(): template_dir = realpath( join(dirname(__file__), "..", "..", "..", "..", "templates") ) env = jinja2.Environment( loader=jinja2.FileSystemLoader(template_dir), block_start_string="[%", block_end_string="%]", variable_start_string="[[", variable_end_string="]]", ) aws_template = env.get_template(f"advance_options_aws_{ENV}.yaml") azure_template = env.get_template(f"advance_options_azure_{ENV}.yaml") chain( get_task_list(), notify_task_begin(), [ github_action( group_id="deploy_aws_regional_service", repo="tidbcloud/auto-deploy", workflow="deploy-aws-region.yml", ref=AUTO_DEPLOY_REPO_BRANCH, input_fn=lambda ti: [ { "tidb-cloud-instance": ENV, "region-name": t["logical_region_name"], "aws-region": t["region_name"], "meta-uri": t["meta_uri"], "force-overwrite-resources": False, "advance-options": aws_template.render( m_account=t["m_account"], d_accounts=t["d_accounts"].split(","), region=t["region_name"], o11y_account=t["o11y_account"], tf_state_bucket=TF_STATE_BUCKET_AWS, seed_instance_type=get_seed_instance_type( ENV, False, t["region_name"] ), ), } for t in ti.xcom_pull(key="task_list") if t["provider"] == "aws" ], ), github_action( group_id="deploy_azure_regional_service", repo="tidbcloud/auto-deploy", workflow="deploy-azure-region.yml", ref=AUTO_DEPLOY_REPO_BRANCH, input_fn=lambda ti: [ { "tidb-cloud-instance": ENV, "logical-region-name": t["logical_region_name"], "cloud-region": t["region_name"], "meta-uri": t["meta_uri"], # aws_account:tenant_id:subscription_id:client_id "tenant-id": t["m_account"].split(":")[1], "subscription-id": t["m_account"].split(":")[2], "client-id": t["m_account"].split(":")[3], "force-overwrite-resources": False, "advance-options": azure_template.render( m_account=t["m_account"].split(":")[0], # subscription_1,subscription_2 d_account_subscription_ids=t["d_accounts"].split(","), ), } for t in ti.xcom_pull(key="task_list") if t["provider"] == "azure" ], ), ], parse_k8s_public_ips(ENV), *add_infra_apiserver_source_ips(ENV), write_task_results(), notify_task_end(), ) create_dedicated_region()