Python｜将 csv 文件转化为 sql 的 INSERT INTO 语句

最近有一个开发场景，需要将从 Hue 中导出的 csv 文件，转化为 INSERT INTO 语句并写入到测试数据库。在处理数据时，存在如下具体需求：

• 需要自动将 HUe 导出的 csv 中 “表名.列名” 的表头清洗为 “列名”
• 需要自动忽略部分字段（例如分区字段）

在具体的实现上：

• 为了兼容从文件中读取 TextIO 以及直接处理字符串，我们允许入参为 TextIO 或字符串，并将字符串类型的入参统一为 TextIO 类型
• 我们使用 csv.reader 解析 csv 文件，这是因为 reader 是 csv 模块中性能最高的
• 为了兼容需要跳过的字段名可能是 “表名.列名” 的格式，也有可能是 “列名” 的格式，所以我们在原始字段名和清洗后的字段名中分别查找需要过滤的字段名是否存在
• 为了避免使用保留字的列名引发问题，所以我们将所有列名都是用引号框住
• 为了避免不同类型的变量值引发问题，所以我们将所有的变量值都用引号框住

具体实现如下：

"""
将 csv 文件转化为 sql 的 INSERT INTO 语句
"""

import csv
from typing import List, Optional, TextIO, Union


def change_to_textio(s: Union[TextIO, str]) -> TextIO:
    """将字符串或 TextIO 转化为 TextIO

    Parameters
    ----------
    s : Union[TextIO, str]
        字符串或 TextIO

    Returns
    -------
    TextIO
    """
    if isinstance(s, str):
        return io.StringIO(s)  # 将字符串转化为 TextIO
    return s


def change_csv_to_sql(f: Union[TextIO, str],
                      table_name: Optional[str] = None,
                      ignore_columns: Optional[List[str]] = None
                      ) -> str:
    """将 csv 文件转化为 sql 的 INSERT INTO 语句

    1. 可以将 "表名.列名" 的表头清洗为 "列名"
    2. 可以设置忽略的字段列表（可以是原始字段名，也可以是清洗后的字段名）

    Parameters
    ----------
    f : Union[TextIO, str]
        字符串或 TextIO（打开的文件）
    table_name : Optional[str], default = None
        如果为 None，则 INSERT 语句中将使用 Table 代替表名
    ignore_columns : Optional[List[str]] = None
        需要忽略的字段列表，使用清洗后的字段名，如果为 None，则不忽略任何字段

    Returns
    -------
    str
        生成的 INSERT INTO 语句
    """
    f = change_to_textio(f)
    if table_name is None:
        table_name = "Table"
    if ignore_columns is None:
        ignore_columns = []

    reader = csv.reader(f)  # 解析 csv 文件

    # 解析字段名
    source_columns = next(reader)
    clean_columns = [column[column.index(".") + 1:] if "." in column else column for column in source_columns]

    # 计算需要忽略的字段下标
    ignore_idx_set = set()
    for ignore_column in ignore_columns:
        if ignore_column in source_columns:  # 先在原始字段名中检索
            ignore_idx = source_columns.index(ignore_column)
            ignore_idx_set.add(ignore_idx)
        elif ignore_column in clean_columns:  # 然后在清洗字段名中检索
            ignore_idx = clean_columns.index(ignore_column)
            ignore_idx_set.add(ignore_idx)

    # 格式化标题行
    columns = ", ".join([f"`{column}`" for i, column in enumerate(clean_columns) if i not in ignore_idx_set])

    # 格式化数据行
    values = ", ".join(["(" + ", ".join([f"'{v}'" if v != "NULL" else v
                                         for i, v in enumerate(row) if i not in ignore_idx_set]) + ")"
                        for row in csv.reader(f)])

    return f"INSERT INTO `{table_name}` ({columns}) VALUES {values};"

使用方法：

if __name__ == "__main__":
    with open(r"file_path", encoding="UTF-8") as file:
        print(change_csv_to_sql(file.read(), table_name="table_name", ignore_columns=["pt"]))