基于Apache Airflow实现多表BQ转Cloud SQL的方案咨询

1. 能否将三个算子整合为接收表名参数的自定义Operator？ #

完全可以。你可以继承Airflow的BaseOperator，把BigQuery查询/导出、GCS到Cloud SQL导入的逻辑封装成一个自定义Operator，通过接收table_name等参数实现复用。

示例代码：

from airflow.models.baseoperator import BaseOperator
from airflow.providers.google.cloud.operators.bigquery import BigQueryOperator, BigQueryToCloudStorageOperator
from airflow.providers.google.cloud.operators.cloud_sql import CloudSqlInstanceImportOperator

class BQToCloudSqlOperator(BaseOperator):
    def __init__(
        self,
        table_name: str,
        bq_project: str,
        bq_dataset: str,
        gcs_bucket: str,
        gcs_prefix: str,
        cloud_sql_instance: str,
        cloud_sql_db: str,
        cloud_sql_target_table: str,
        **kwargs
    ):
        super().__init__(**kwargs)
        self.table_name = table_name
        self.bq_project = bq_project
        self.bq_dataset = bq_dataset
        self.gcs_bucket = gcs_bucket
        self.gcs_prefix = gcs_prefix
        self.cloud_sql_instance = cloud_sql_instance
        self.cloud_sql_db = cloud_sql_db
        self.cloud_sql_target_table = cloud_sql_target_table

    def execute(self, context):
        # 1. BigQuery数据预处理（可选，比如生成临时表）
        bq_task = BigQueryOperator(
            task_id=f"preprocess_bq_{self.table_name}",
            sql=f"SELECT * FROM `{self.bq_project}.{self.bq_dataset}.{self.table_name}`",
            destination_dataset_table=f"{self.bq_project}.{self.bq_dataset}.temp_{self.table_name}",
            write_disposition="WRITE_TRUNCATE",
            use_legacy_sql=False,
            dag=self.dag
        )
        bq_task.execute(context)

        # 2. 导出BigQuery数据到GCS
        export_task = BigQueryToCloudStorageOperator(
            task_id=f"export_bq_to_gcs_{self.table_name}",
            source_project_dataset_table=f"{self.bq_project}.{self.bq_dataset}.temp_{self.table_name}",
            destination_cloud_storage_uris=[f"gs://{self.gcs_bucket}/{self.gcs_prefix}/{self.table_name}.csv"],
            export_format="CSV",
            field_delimiter=",",
            print_header=False,
            dag=self.dag
        )
        export_task.execute(context)

        # 3. 从GCS导入到Cloud SQL
        import_task = CloudSqlInstanceImportOperator(
            task_id=f"import_gcs_to_cloudsql_{self.table_name}",
            instance=self.cloud_sql_instance,
            body={
                "importContext": {
                    "fileType": "CSV",
                    "uri": f"gs://{self.gcs_bucket}/{self.gcs_prefix}/{self.table_name}.csv",
                    "database": self.cloud_sql_db,
                    "csvImportOptions": {
                        "table": self.cloud_sql_target_table,
                        "columns": []  # 按需指定列映射
                    }
                }
            },
            dag=self.dag
        )
        import_task.execute(context)

使用时只需循环表名列表，实例化这个自定义Operator即可，每个表对应一个独立的迁移任务。

2. 是否可通过运行时变量或环境变量实现？ #

可以，两种方式都能动态控制要迁移的表列表，无需修改DAG代码：

运行时变量（Airflow Variable） #

在Airflow UI的「Admin > Variables」中创建一个变量（比如bq_migrate_tables），值设为逗号分隔的表名（如table1,table2,table3），然后在DAG中读取并生成任务：

from airflow.models import Variable

# 读取变量并处理表名列表
tables = [t.strip() for t in Variable.get("bq_migrate_tables").split(",") if t.strip()]

for table in tables:
    BQToCloudSqlOperator(
        task_id=f"migrate_{table}",
        table_name=table,
        # 传入其他固定参数（bq_project、gcs_bucket等）
        dag=dag
    )

后续只需在UI中修改变量值，就能调整迁移的表范围。

环境变量 #

在Airflow部署的服务器/容器中设置环境变量（比如BQ_MIGRATE_TABLES="table1,table2"），然后在DAG中读取：

import os

# 读取环境变量并处理表名列表
tables = [t.strip() for t in os.getenv("BQ_MIGRATE_TABLES", "").split(",") if t.strip()]

# 循环生成迁移任务（同上）

适合CI/CD部署场景，可通过部署脚本动态传入表名列表。

其他可行方案建议 #

方案1：使用TaskGroup组织单表迁移流程 #

无需自定义Operator，直接用Airflow的TaskGroup将每个表的三个算子打包成一组，DAG界面更清晰，且便于单个表的逻辑调整：

from airflow.utils.task_group import TaskGroup

tables = ["table1", "table2", "table3"]

for table in tables:
    with TaskGroup(group_id=f"migrate_{table}_group") as table_group:
        # 1. BigQuery预处理任务
        bq_preprocess = BigQueryOperator(
            task_id="bq_preprocess",
            sql=f"SELECT * FROM `project.dataset.{table}`",
            destination_dataset_table=f"project.dataset.temp_{table}",
            write_disposition="WRITE_TRUNCATE",
            use_legacy_sql=False
        )
        # 2. 导出到GCS
        export_to_gcs = BigQueryToCloudStorageOperator(
            task_id="export_to_gcs",
            source_project_dataset_table=f"project.dataset.temp_{table}",
            destination_cloud_storage_uris=[f"gs://bucket/prefix/{table}.csv"],
            export_format="CSV"
        )
        # 3. 导入到Cloud SQL
        import_to_sql = CloudSqlInstanceImportOperator(
            task_id="import_to_cloudsql",
            instance="sql-instance",
            body={
                "importContext": {
                    "fileType": "CSV",
                    "uri": f"gs://bucket/prefix/{table}.csv",
                    "database": "sql-db",
                    "csvImportOptions": {
                        "table": table,
                        "columns": []
                    }
                }
            }
        )
        # 设置任务依赖
        bq_preprocess >> export_to_gcs >> import_to_sql
    # 将任务组加入DAG
    table_group

方案2：利用Dynamic Task Mapping（Airflow 2.2+） #

Airflow 2.2及以上支持动态任务映射，无需手动循环，直接通过expand方法批量生成任务，代码更简洁：

from airflow.decorators import dag, task_group
from datetime import datetime
from airflow.providers.google.cloud.operators.bigquery import BigQueryOperator, BigQueryToCloudStorageOperator
from airflow.providers.google.cloud.operators.cloud_sql import CloudSqlInstanceImportOperator

@dag(schedule_interval="@daily", start_date=datetime(2024, 1, 1), catchup=False)
def bq_cloudsql_batch_migrate():
    @task_group
    def migrate_single_table(table_name: str):
        # 定义单表迁移的三个任务
        bq_task = BigQueryOperator(
            task_id="bq_preprocess",
            sql=f"SELECT * FROM `project.dataset.{table_name}`",
            destination_dataset_table=f"project.dataset.temp_{table_name}",
            write_disposition="WRITE_TRUNCATE",
            use_legacy_sql=False
        )
        export_task = BigQueryToCloudStorageOperator(
            task_id="export_to_gcs",
            source_project_dataset_table=f"project.dataset.temp_{table_name}",
            destination_cloud_storage_uris=[f"gs://bucket/prefix/{table_name}.csv"],
            export_format="CSV"
        )
        import_task = CloudSqlInstanceImportOperator(
            task_id="import_to_cloudsql",
            instance="sql-instance",
            body={
                "importContext": {
                    "fileType": "CSV",
                    "uri": f"gs://bucket/prefix/{table_name}.csv",
                    "database": "sql-db",
                    "csvImportOptions": {
                        "table": table_name,
                        "columns": []
                    }
                }
            }
        )
        bq_task >> export_task >> import_task

    # 动态映射生成所有表的迁移任务组
    migrate_single_table.expand(table_name=["table1", "table2", "table3"])

# 实例化DAG
bq_cloudsql_batch_migrate()

方案3：添加校验与错误处理 #

在批量迁移流程中增加额外逻辑提升可靠性：

• 用BigQueryCheckOperator校验源表是否存在、数据量是否符合预期
• 用BranchPythonOperator跳过不存在的表
• 添加EmailOperator或SlackWebhookOperator在任务失败时发送告警
• 迁移完成后清理GCS临时文件（用GCSSynchronizeBucketsOperator或自定义Python函数）

内容的提问来源于stack exchange，提问作者AD90