如何为Pandas to_sql()方法创建的数据库设置主键？

Pandas DataFrame的to_sql()方法非常实用，官方基础示例如下：

import pandas as pd
from sqlalchemy import create_engine, text

# 创建SQLite引擎
engine = create_engine('sqlite:///test.db')

# 用to_sql()创建users表
df = pd.DataFrame({'name' : ['User 1', 'User 2', 'User 3']})
df.to_sql('users', con=engine, if_exists='replace')

# 验证结果
with engine.connect() as conn:
    print(conn.execute(text("Select * from users")).fetchall())

如果希望将name列设置为主键，无需依赖SQL注入，完全可以用SQLAlchemy的面向对象接口实现，以下是两种可行方案：

方案1：提前定义SQLAlchemy表结构（推荐） #

这种方式从源头规范表结构，直接设置主键，是最符合SQLAlchemy设计理念的做法：

import pandas as pd
from sqlalchemy import create_engine, Table, Column, String, MetaData, text

engine = create_engine('sqlite:///test.db')
metadata = MetaData()

# 定义带主键的users表结构
users_table = Table(
    'users',
    metadata,
    Column('name', String, primary_key=True)
)

# 创建表到数据库
metadata.create_all(engine)

# 将DataFrame数据导入表（index=False避免生成默认索引列）
df = pd.DataFrame({'name': ['User 1', 'User 2', 'User 3']})
df.to_sql('users', con=engine, if_exists='append', index=False)

# 验证主键生效
with engine.connect() as conn:
    # 尝试插入重复name会触发主键约束报错
    try:
        conn.execute(text("INSERT INTO users (name) VALUES ('User 1')"))
        conn.commit()
    except Exception as e:
        print(f"主键约束生效：{e}")

方案2：为已创建的表添加主键（适配不同数据库） #

如果已经通过to_sql()创建了无主键的表，可以用SQLAlchemy的反射机制读取表结构，再通过重建表的方式添加主键（注意：部分数据库如SQLite不支持直接修改表添加主键，需重建）：

import pandas as pd
from sqlalchemy import create_engine, MetaData, Table, Column, String, text

engine = create_engine('sqlite:///test.db')

# 先创建无主键的users表
df = pd.DataFrame({'name': ['User 1', 'User 2', 'User 3']})
df.to_sql('users', con=engine, if_exists='replace', index=False)

# 反射读取现有表结构
metadata = MetaData()
metadata.reflect(bind=engine)
original_table = metadata.tables['users']

# 重建带主键的表
with engine.connect() as conn:
    # 创建临时表并设置主键
    temp_table = Table(
        'users_temp',
        metadata,
        Column('name', String, primary_key=True)
    )
    temp_table.create(conn)
    
    # 迁移数据到临时表
    conn.execute(text("INSERT INTO users_temp SELECT * FROM users"))
    
    # 删除原表并重命名临时表
    original_table.drop(conn)
    conn.execute(text("ALTER TABLE users_temp RENAME TO users"))
    
    conn.commit()

# 验证结果
with engine.connect() as conn:
    result = conn.execute(text("SELECT * FROM users")).fetchall()
    print(result)

方案1是优先选择，它既利用了SQLAlchemy的面向对象特性，又避免了后续修改表结构的繁琐操作。

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