import hashlib from contextlib import closing from os.path import dirname, join, realpath import jinja2 import pendulum import utils.lark as lark import yaml from airflow.datasets import Dataset from airflow.decorators import dag, task from airflow.exceptions import AirflowFailException from airflow.io.path import ObjectStoragePath 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 ( get_first_availability_zone, 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"] LARK_ALERT_CHAT_ID = config["LARK_ALERT_CHAT_ID"] AUTO_DEPLOY_REPO_BRANCH = config["AUTO_DEPLOY_REPO_BRANCH"] @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 nextgen region {t['logical_region_name']}", { "☁️ Cloud": t["provider"], "📍 Region": t["region_name"], "🪪 Name": t["logical_region_name"], "🧩 AZ IDs": t["availability_zones"], }, { "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"nextgen region {t['logical_region_name']} created", { "☁️ Cloud": t["provider"], "📍 Region": t["region_name"], "🪪 Name": t["logical_region_name"], "🧩 AZ IDs": t["availability_zones"], "🏁 Status": "SUCCESS", }, { "Regional Service": "created", "Cluster Service": "created", "BR Service": "created", "CDC Service": "created", }, ), email=t["trigger"], ) @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}.nextgen_regional_create_tasks """, ) for t in tasks: t["is_byoc"] = False t["is_ng"] = True 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 t["meta_uri"] = ( f"s3://{META_BUCKET_AWS}/dedicated-regions/{t['logical_region_name']}/meta.yaml" ) t["rds_meta_uri"] = ( f"s3://{META_BUCKET_AWS}/region-shared-resoruces/ng-rds/{t['rds_instance_name']}/meta.yaml" ) ti.xcom_push(key="task_list", value=tasks) return len(tasks) > 0 @task def load_rds_info(ti: TaskInstance): task_list = ti.xcom_pull(key="task_list") new_task_list = [] for t in task_list: meta_uri = t["rds_meta_uri"] meta = ObjectStoragePath( meta_uri, conn_id=f"auto-deploy-cli-{t['provider']}-{ENV}" ) if not meta.is_file(): raise AirflowFailException(f"failed to load rds info: {meta_uri}") with meta.open() as f: x = yaml.safe_load(f) t["rds_full_name"] = x["full_name"] t["rds_azs"] = ",".join(x["azs"]) t["rds_endpoint"] = x["endpoint"] t["rds_password_secret_name"] = x["password_secret_name"] t["rds_port"] = x.get("port", "3306") t["rds_username"] = x.get("username", "admin") t["rds_publicly_accessible"] = x.get("publicly_accessible", "true") t["rds_rcfg_arn"] = x.get("rcfg_arn", "") t["rds_proxy_service"] = x.get("proxy_service", "") t["rds_proxy_name"] = x.get("proxy_name", "") new_task_list.append(t) ti.xcom_push("task_list", new_task_list) @task( outlets=[ Dataset( f"mysql://autonomous-serverless-instance/auto_deploy_dedicated_{ENV}/{tb}" ) for tb in [ "nextgen_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}.nextgen_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, is_byoc, is_ng, status, meta_uri, public_ips, trace_id, requested_at) values (%(provider)s, %(region_name)s, %(logical_region_name)s, %(is_byoc)s, %(is_ng)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_nextgen_region_{ENV}", start_date=pendulum.datetime(2024, 10, 25, tz="UTC"), schedule=[ Dataset( f"mysql://autonomous-serverless-instance/auto_deploy_dedicated_{ENV}/nextgen_regional_create_requests" ) ], tags=["deploy", "region", "nextgen", 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_nextgen_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_nextgen_{ENV}.yaml") rds_template = env.get_template(f"advance_options_aws_rds_{ENV}.json") chain( get_task_list(), notify_task_begin(), github_action( group_id="deploy_aws_ng_rds", repo="tidbcloud/auto-deploy", workflow="deploy-aws-ng-rds.yml", ref=AUTO_DEPLOY_REPO_BRANCH, input_fn=lambda ti: [ { "environment": ENV, "name": t["rds_instance_name"], "aws-region": t["region_name"], "aws-role-arn": f"arn:aws:iam::{t['m_account']}:role/auto-deploy-cli", "output-uri": t["rds_meta_uri"], "tf-state-s3-bucket": TF_STATE_BUCKET_AWS, "advance-options": rds_template.render( publicly_accessible=("true" if ENV != "prod" else "false"), availability_zones=t["availability_zones"], ), } for t in ti.xcom_pull(key="task_list") if t["provider"] == "aws" and t["create_rds"] ], ), load_rds_info(), github_action( group_id="deploy_aws_nextgen_regional_service", repo="tidbcloud/auto-deploy", workflow="deploy-aws-nextgen-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, "service-plan": "Premium", "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, availability_zones=t["availability_zones"], first_availability_zone=get_first_availability_zone( t["availability_zones"] ), default_dataplane_account=t["d_accounts"].split(",")[0], seed_instance_type=get_seed_instance_type( ENV, False, t["region_name"] ), rds={ "host": t["rds_endpoint"], "password_secret_name": t["rds_password_secret_name"], "publicly_accessible": ( '"true"' if t["rds_publicly_accessible"] else '"false"' ), "port": t["rds_port"], "username": t["rds_username"], "rcfg_arn": t["rds_rcfg_arn"], "availability_zones": t["rds_azs"], "proxy_service": t["rds_proxy_service"], "proxy_name": t["rds_proxy_name"], }, ), } for t in ti.xcom_pull(key="task_list") if t["provider"] == "aws" ], ), parse_k8s_public_ips(ENV), write_task_results(), notify_task_end(), ) create_dedicated_nextgen_region()