from contextlib import closing 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.operators.empty import EmptyOperator from airflow.providers.mysql.hooks.mysql import MySqlHook from airflow.utils.trigger_rule import TriggerRule from dotenv import dotenv_values from utils.database import get_records_from_db from utils.github import github_action from utils.url import blob_url_to_abfs config = dotenv_values() ENV = config["ENV"] META_BUCKET_AWS = config["META_BUCKET_AWS"] META_STORAGE_ACCOUNT = config["META_STORAGE_ACCOUNT"] LARK_ALERT_CHAT_ID = config["LARK_ALERT_CHAT_ID"] O11Y_PROD_META_BUCKET_AWS = config["O11Y_PROD_META_BUCKET_AWS"] O11Y_PROD_META_STORAGE_ACCOUNT = config["O11Y_PROD_META_STORAGE_ACCOUNT"] AUTO_DEPLOY_REPO_BRANCH = config["AUTO_DEPLOY_REPO_BRANCH"] @task.short_circuit def get_task_list(ti: TaskInstance): db = MySqlHook(mysql_conn_id="autonomous-serverless-instance") tasks = get_records_from_db( db, f""" select * from auto_deploy_dedicated_{ENV}.o11y_config_tasks where provider in ('aws', 'azure') """, ) for t in tasks: if t["provider"] == "aws": t["meta_uri"] = ( f"s3://{META_BUCKET_AWS}/dedicated-regions/{t['logical_region_name']}/meta.yaml" ) t["role_arn"] = f"arn:aws:iam::{t['m_account']}:role/auto-deploy-cli" o11y_meta_file = ObjectStoragePath( f"s3://{O11Y_PROD_META_BUCKET_AWS}/o11y-regions/{t['o11y_region_name']}/meta.yaml", conn_id="auto-deploy-cli-aws-o11y-prod", ) 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" ) o11y_meta_file = ObjectStoragePath( f"abfs://{O11Y_PROD_META_STORAGE_ACCOUNT}/o11y-regions/{t['o11y_region_name']}/meta.yaml", conn_id="auto-deploy-cli-azure-o11y-prod", ) if o11y_meta_file.is_file(): with o11y_meta_file.open() as f: x = yaml.safe_load(f) t["o11y_cluster"] = x["regions"][t["o11y_region_name"]]["k8s"]["name"] else: # for o11y region not deployed by this tool t["o11y_cluster"] = f"o11y-prod-regional-{t['region_name']}" 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 config o11y for region {t['logical_region_name']}", { "☁️ Cloud": t["provider"], "📍 Region": t["region_name"], "🪪 Name": t["logical_region_name"], "🪪 O11Y": t["o11y_region_name"], }, { "IP Whitelist": "updating", "O11Y Agent": "enabling", "O11Y Roles for DP": "creating", "O11Y Config": "updating", "O11Y Regional Role": t["o11y_regional_role_name"], "Dedicated Public IPs": t["dedicated_public_ips"], "O11Y Public IPs": t["o11y_public_ips"], }, ), 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"finish config o11y for region {t['logical_region_name']}", { "☁️ Cloud": t["provider"], "📍 Region": t["region_name"], "🪪 Name": t["logical_region_name"], "🪪 O11Y": t["o11y_region_name"], "🏁 Status": "SUCCESS", }, { "IP Whitelist": "updated", "O11Y Agent": "enabled", "O11Y Roles for DP": "created", "O11Y Config": "updated", "O11Y Regional Role": t["o11y_regional_role_name"], "Dedicated Public IPs": t["dedicated_public_ips"], "O11Y Public IPs": t["o11y_public_ips"], }, ), email=t["trigger"], ) def parse_cluster_id(provider: str, output_uri: str) -> str: if provider == "azure": output_uri = blob_url_to_abfs(output_uri) file = ObjectStoragePath( output_uri, conn_id=f"auto-deploy-cli-{provider}-{ENV}", ) if not file.is_file(): raise AirflowFailException(f"job output uri {output_uri} not exist") with file.open() as f: x = yaml.safe_load(f) return x["o11y_cluster_id"] @task def load_o11y_cluster_id(ti: TaskInstance = None): task_list = ti.xcom_pull(key="task_list") new_task_list = [] for t in task_list: meta_uri = t["meta_uri"] if t["provider"] == "azure": meta_uri = blob_url_to_abfs(meta_uri) meta = ObjectStoragePath( meta_uri, conn_id=f"auto-deploy-cli-{t['provider']}-{ENV}" ) if not meta.is_file(): continue with meta.open() as f: x = yaml.safe_load(f) outputs = x["regions"][t["logical_region_name"]]["job-outputs"] if t["is_byoc"]: output_uri = outputs["controlplane-o11y-agent"]["output"] else: output_uri = outputs["o11y-agent"]["output"] t["o11y_cluster_id"] = parse_cluster_id(t["provider"], output_uri) 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}/o11y_config_results" ) ] ) def write_task_results(ti: TaskInstance = None): task_list = ti.xcom_pull(key="task_list") 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: t["status"] = "success" cur.execute( f""" insert into auto_deploy_dedicated_{ENV}.o11y_config_results (request_id, status, regional_service_record_id, o11y_record_id, meta_uri, trace_id) values (%(id)s, %(status)s, %(regional_service_record_id)s, %(o11y_record_id)s, %(meta_uri)s, %(trace_id)s) """, t, ) conn.commit() @dag( dag_id=f"config_o11y_for_dedicated_region_{ENV}", start_date=pendulum.datetime(2024, 10, 25, tz="UTC"), schedule=( Dataset( f"mysql://autonomous-serverless-instance/auto_deploy_dedicated_{ENV}/o11y_config_requests" ) | Dataset( f"mysql://autonomous-serverless-instance/auto_deploy_dedicated_{ENV}/regional_services" ) | Dataset( "mysql://autonomous-serverless-instance/auto_deploy_o11y_prod/o11y_regions" ) ), 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 config_o11y_for_dedicated_region(): chain( get_task_list(), notify_task_begin(), github_action( group_id="update-iplist", repo="tidbcloud/observability", workflow="new-region-update-iplist.yaml", input_fn=lambda ti: map( lambda t: { "vendor": t["provider"], "region": t["o11y_region_name"], "env": "prod", "cluster": f"o11y-prod-regional-{t['o11y_region_name']}", "rsIpWhitelist": ",".join( map(lambda ip: ip + "/32", t["dedicated_public_ips"].split(",")) ), "git-ref": "main", }, ti.xcom_pull(key="task_list"), ), ), [ github_action( group_id="enable-o11y-agent", repo="tidbcloud/auto-deploy", workflow="enable-o11y-agent.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"], "role-arn": t["role_arn"], "job-name": ( "controlplane-o11y-agent" if t["is_byoc"] else "o11y-agent" ), } for t in ti.xcom_pull(key="task_list") if t["provider"] == "aws" ], skip_on_empty=True, ), github_action( group_id="enable-o11y-agent-azure", repo="tidbcloud/auto-deploy", workflow="enable-azure-o11y-agent.yml", ref=AUTO_DEPLOY_REPO_BRANCH, input_fn=lambda ti: [ { "tidb-cloud-instance": ENV, "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], } for t in ti.xcom_pull(key="task_list") if t["provider"] == "azure" ], skip_on_empty=True, ), ], [ github_action( group_id="create-o11y-roles-on-dataplane", repo="tidbcloud/auto-deploy", workflow="create-o11y-roles-on-dataplane.yaml", ref=AUTO_DEPLOY_REPO_BRANCH, input_fn=lambda ti: [ { "env": ENV, "region": t["region_name"], "logical-region-name": t["o11y_region_name"], "d-accounts": t["d_accounts"], "o11y-account": t["o11y_account"], "o11y-cluster": t["o11y_cluster"], "o11y-regional-role-name": t["o11y_regional_role_name"], } for t in ti.xcom_pull(key="task_list") if t["provider"] == "aws" and not t["is_byoc"] ], skip_on_empty=True, ), github_action( group_id="assign-roles-to-o11y-azure-id-on-dataplane", repo="tidbcloud/auto-deploy", workflow="assign-roles-to-o11y-azure-id-on-dataplane.yaml", ref=AUTO_DEPLOY_REPO_BRANCH, input_fn=lambda ti: [ { "env": ENV, "d-accounts": t["d_accounts"], # 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], "o11y-role-client-id": t["o11y_regional_role_name"], "roles": "Contributor", } for t in ti.xcom_pull(key="task_list") if t["provider"] == "azure" ], skip_on_empty=True, ), ], [ EmptyOperator( task_id="continue_in_case_aws_skiped", trigger_rule=TriggerRule.NONE_FAILED, ), EmptyOperator( task_id="continue_in_case_azure_skiped", trigger_rule=TriggerRule.NONE_FAILED, ), ], load_o11y_cluster_id(), github_action( group_id="update-o11y-config", repo="tidbcloud/app-delivery", workflow="update-o11y-config-for-aws-region.yaml", input_fn=lambda ti: [ { "env": ENV, "cloud": t["provider"], "logicalRegionName": t["logical_region_name"], "o11y-regional-public-ips": ",".join( map(lambda ip: ip + "/32", t["o11y_public_ips"].split(",")) ), "o11y-regional-endpoint": f"https://www.rs.{t['o11y_region_name']}.{t['provider']}.observability.tidbcloud.com", "o11y-regional-iam-role": ( f"arn:aws:iam::{t['o11y_account']}:role/{t['o11y_regional_role_name']}" if t["provider"] == "aws" else "" ), "o11y-seed-cluster-id": t["o11y_cluster_id"], } for t in ti.xcom_pull(key="task_list") if not t["is_ng"] ], ), write_task_results(), notify_task_end(), ) config_o11y_for_dedicated_region()