CSV转SQL插入语句并写入文件的Python脚本

CSV转SQL插入语句并写入文件的Python脚本

import csv
import os
import time


def csv_to_sql_insert(csv_file, table_name, chunk_size):
    # 读取CSV文件
    with open(csv_file, 'r', newline='') as file:
        reader = csv.reader(file)

        # 读取并处理列名
        headers = next(reader)
        columns = ', '.join(map(lambda x: f'"{x}"', headers))
        insert_template = f'INSERT INTO {table_name} ({columns}) VALUES\n'

        # 处理CSV数据并生成SQL语句
        statements = []
        for row in reader:
            # 清除空值和NULL，并用空字符串替换
            cleaned_row = [value.replace('null', '').strip() if value != 'NULL' else 'NULL' for value in row]

            # # 组装SQL VALUES部分
            # row_values = ', '.join(map(lambda x: f'"{x}"' if x != 'NULL' else 'NULL', cleaned_row))

            # 组装SQL VALUES部分，保留NULL值
            # row_values = ', '.join(f'"{value}"' if value != 'NULL' else 'NULL' for value in row)

            # 组装SQL VALUES部分，如果值为null或NULL则不加双引号
            row_values = ', '.join(f'NULL' if value.upper() == 'NULL' else f'"{value}"' for value in row)

            # 添加到SQL语句列表中
            statements.append(f'({row_values}),\n')

            # 分组SQL语句
        chunked_statements = [statements[i:i + chunk_size] for i in range(0, len(statements), chunk_size)]

        # 组装完整的SQL INSERT语句
        sql_insert_statements = []
        for chunk in chunked_statements:
            sql_insert = insert_template + ''.join(chunk)[:-2] + ';\n'
            sql_insert_statements.append(sql_insert)

        return sql_insert_statements


def write_to_sql_files(sql_statements, base_path, table_name):
    # 写入SQL语句到文件
    for i, statement in enumerate(sql_statements, start=1):
        file_path = os.path.join(base_path, f'{table_name}-{i}.sql')
        with open(file_path, 'w', newline='') as f:
            f.write(statement)
            # print(statement)
        print(f'Wrote {file_path}')


# 测试代码
if __name__ == '__main__':
    # 记录开始时间
    start_time = time.time()

    base_path = 'E:\\20240421_test_python\\'
    file_name = 'tbl_test.csv'
    csv_file_path = os.path.join(base_path, file_name)
    table_name = 'tbl_test'
    chunk_size = 15

    # 转换CSV到SQL INSERT语句
    sql_statements = csv_to_sql_insert(csv_file_path, table_name, chunk_size)

    # 将SQL语句写入文件
    write_to_sql_files(sql_statements, base_path, table_name)

    # 记录结束时间
    end_time = time.time()

    # 计算并打印运行时间
    elapsed_time = end_time - start_time
    print(f'Elapsed time: {elapsed_time:.2f} seconds')