GCP Composer中PostgreSQL批量导BigQuery仅首批次成功问题排查

通过PostgreSQL cursor的fetchmany方法分批（每批次1000行）将10万行数据导入BigQuery，本地运行正常，但部署到GCP Composer后仅首批次1000行插入成功，后续批次无数据写入，且日志无报错信息。数据包含8列，类型涵盖整数、字符串、日期。已尝试添加sleep延迟，问题仍存在。

附原代码片段：

with cursor:
    cursor.execute(sql_query)
    while True:
        rows = cursor.fetchmany(1000)
        if not rows:
            break
        logger.info(f"rows :{len(rows)}")
        column_names = [desc[0] for desc in cursor.description]
        logger.info(f"Column name: {column_names}")
        df = pd.DataFrame(rows, columns=column_names)
        df.reset_index(drop=True, inplace=True)
        if schema_dict is not None and selected_column is not None:
            df = df[selected_column]
            df = convert_pandas_datatype(df, schema_dict)
        client.load_table_from_dataframe(
            df,
            table_id,
            job_config=job_config
        )
        # from time import sleep
        # sleep(5)
        # print("sleeping............")
conn.close()

1. 等待BigQuery加载作业完成（核心修复） #

load_table_from_dataframe是异步方法，返回的Job对象需要调用result()方法等待作业执行完成。在Composer环境中，如果不等待作业完成，后续代码可能在作业提交后直接继续执行，甚至进程提前退出，导致后续批次的任务未正确触发或完成。

修改代码如下：

with cursor:
    cursor.execute(sql_query)
    # 提前获取列名，避免循环中重复读取cursor.description
    column_names = [desc[0] for desc in cursor.description]
    logger.info(f"Column name: {column_names}")
    while True:
        rows = cursor.fetchmany(1000)
        if not rows:
            break
        logger.info(f"Processing batch with {len(rows)} rows")
        df = pd.DataFrame(rows, columns=column_names)
        df.reset_index(drop=True, inplace=True)
        if schema_dict is not None and selected_column is not None:
            df = df[selected_column]
            df = convert_pandas_datatype(df, schema_dict)
        # 提交作业并等待完成
        job = client.load_table_from_dataframe(
            df,
            table_id,
            job_config=job_config
        )
        # 等待作业完成，捕获可能的异常并记录
        try:
            job.result()
            logger.info(f"Batch completed, inserted {len(rows)} rows")
        except Exception as e:
            logger.error(f"Batch failed: {str(e)}", exc_info=True)
            raise
conn.close()

2. 优化PostgreSQL连接稳定性 #

Composer环境中，PostgreSQL连接可能因超时或资源回收导致后续fetchmany返回空值：

• 增加连接超时参数，确保分批读取过程中连接不会被断开
• 避免在循环内重复获取cursor.description，减少对数据库连接的依赖（已在上述代码中优化）

3. 调整BigQuery作业配置 #

确保job_config的写入模式和schema配置正确：

• 如果是追加数据，需设置write_disposition=bigquery.WriteDisposition.WRITE_APPEND
• 检查schema配置与DataFrame字段类型完全匹配，避免隐性类型转换错误导致数据丢弃

示例配置：

job_config = bigquery.LoadJobConfig(
    write_disposition=bigquery.WriteDisposition.WRITE_APPEND,
    schema=your_table_schema,
)

4. 添加详细日志与异常捕获 #

在代码中补充作业ID、状态等日志，便于排查隐性问题：

logger.info(f"Submitted BigQuery job: {job.job_id}")
job.result()
logger.info(f"Job {job.job_id} completed, state: {job.state}")

内容的提问来源于stack exchange，提问作者Thái Lệ