如何在Polars+SQLAlchemy中设置查询超时并标记超时行

关键改进点 #

• 捕获查询超时异常并标记对应行
• 替换SQL字符串拼接为参数化查询，避免注入风险
• 新增状态列明确区分「超时」「无数据」「正常数据」
• 自动管理数据库连接，避免资源泄露

修改后的完整代码 #

import polars as pl
from sqlalchemy import create_engine
from sqlalchemy.exc import OperationalError

# 创建引擎，添加连接超时配置
E27V422024 = create_engine(
    'mssql://xd/WMBDP2024?trusted_connection=yes&driver=ODBC+Driver+17+for+SQL+Server',
    pool_size=0, max_overflow=0,
    connect_args={'timeout': 2}
)

concatenated_df = pl.DataFrame()

# 定义查询模板（参数化）
query_template = """
SELECT
    rut,
    codtran,
    coderr,
    monto,
    fecha
FROM database
WHERE rut = ? AND fecha BETWEEN ? AND ?
AND codtran IN ('MPMR80', 'MPMR22', 'SCULOE')
"""

for row in first_five_rows.to_dicts():
    rut = row['Rut']
    fechaI = row['FechaEnvio']
    fechaO = row['FechaTermino']
    temp_df = None

    try:
        # 使用with块自动管理连接，设置查询超时（单位：秒）
        with E27V422024.connect() as conn:
            temp_df = pl.read_database(
                query_template,
                connection=conn.execution_options(query_timeout=2),
                params=[rut, fechaI, fechaO]
            )

        # 处理查询无结果的情况
        if temp_df.is_empty():
            data = {
                "rut": rut,
                "codtran": 'nop',
                "coderr": 'nop',
                "monto": 8888,
                "fecha": fechaI,
                "status": 'no_data'
            }
            temp_df = pl.DataFrame(
                data,
                schema={
                    "rut": pl.Int64(),
                    "codtran": pl.String(),
                    "coderr": pl.String(),
                    "monto": pl.Decimal(),
                    "fecha": pl.Datetime(),
                    "status": pl.String()
                },
                strict=False
            )

    except OperationalError as e:
        # 判断是否为超时异常（根据MSSQL ODBC错误信息调整匹配条件）
        if "timeout" in str(e.orig).lower() or e.orig.args[0] in (1222, -2, 0x80004005):
            # 生成超时标记行
            data = {
                "rut": rut,
                "codtran": 'timeout',
                "coderr": 'timeout',
                "monto": None,
                "fecha": fechaI,
                "status": 'timeout'
            }
            temp_df = pl.DataFrame(
                data,
                schema={
                    "rut": pl.Int64(),
                    "codtran": pl.String(),
                    "coderr": pl.String(),
                    "monto": pl.Decimal(),
                    "fecha": pl.Datetime(),
                    "status": pl.String()
                },
                strict=False
            )
        else:
            # 非超时异常，打印信息后跳过当前行
            print(f"处理Rut {rut}时发生错误: {e}")
            continue

    # 合并结果
    if temp_df is not None:
        concatenated_df = pl.concat([concatenated_df, temp_df])

print(concatenated_df)

注意事项 #

1. 错误码匹配：MSSQL不同场景的超时对应不同错误码（如1222是锁超时，-2是连接超时），可通过print(e.orig)查看原始错误信息，调整except块中的判断条件。
2. 性能优化：循环单条查询效率较低，若业务允许，可将多个rut批量放入IN语句中一次查询，减少数据库交互次数。
3. 数据类型一致性：确保params中的参数类型与数据库列类型匹配，避免类型转换异常。

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