时间戳

首页 > 解决方案 > 返回 ID 时批量保存和更新

python - 返回 ID 时批量保存和更新

问题描述

所以我正在使用我正在sqlalchemy做的一个项目。我遇到了一个问题，我最终每小时都会有数千条记录需要保存。这些记录可能会被插入或更新。我一直在使用bulk_save_objects它，效果很好。但是现在我必须为这些正在保存的记录引入历史记录，这意味着我需要返回 ID，以便我可以将这些条目链接到历史记录表中的条目。我知道使用return_defaults，并且有效。但是，它引入了一个问题，即我bulk_save_objects一次插入和更新一个条目，而不是批量，这消除了目的。是否有另一种选择，我可以同时批量插入和更新，但保留 ID？

标签： pythonsqlalchemy

解决方案

可以使用类似于此处答案中描述的技术来实现所需的结果，方法是将行上传到临时表，然后执行 UPDATE，然后执行返回插入的 ID 值的 INSERT。对于 SQL Server，这将是 INSERT 语句上的 OUTPUT 子句：

main_table = "team"

# <set up test environment>
with engine.begin() as conn:
    conn.execute(sa.text(f"DROP TABLE IF EXISTS [{main_table}]"))
    conn.execute(
        sa.text(
            f"""
            CREATE TABLE [dbo].[{main_table}](
                [id] [int] IDENTITY(1,1) NOT NULL,
                [prov] [varchar](2) NOT NULL,
                [city] [varchar](50) NOT NULL,
                [name] [varchar](50) NOT NULL,
                [comments] [varchar](max) NULL,
             CONSTRAINT [PK_team] PRIMARY KEY CLUSTERED 
            (
                [id] ASC
            )
            )
            """
        )
    )
    conn.execute(
        sa.text(
            f"""
            CREATE UNIQUE NONCLUSTERED INDEX [UX_team_prov_city] ON [dbo].[{main_table}]
            (
                [prov] ASC,
                [city] ASC
            )            
            """
        )
    )
    conn.execute(
        sa.text(
            f"""
            INSERT INTO [{main_table}] ([prov], [city], [name])
            VALUES ('AB', 'Calgary', 'Flames')            
            """
        )
    )

# <data for upsert>
df = pd.DataFrame(
    [
        ("AB", "Calgary", "Flames", "hard-working, handsome lads"),
        ("AB", "Edmonton", "Oilers", "ruffians and scalawags"),
    ],
    columns=["prov", "city", "name", "comments"],
)

# <perform upsert, returning IDs>
temp_table = "#so65525098"
with engine.begin() as conn:
    df.to_sql(temp_table, conn, index=False, if_exists="replace")
    conn.execute(
        sa.text(
            f"""
            UPDATE main SET main.name = temp.name, 
                main.comments = temp.comments
            FROM [{main_table}] main INNER JOIN [{temp_table}] temp
                ON main.prov = temp.prov AND main.city = temp.city
            """
        )
    )
    inserted = conn.execute(
        sa.text(
            f"""
            INSERT INTO [{main_table}] (prov, city, name, comments)
            OUTPUT INSERTED.prov, INSERTED.city, INSERTED.id
            SELECT prov, city, name, comments FROM [{temp_table}] temp
            WHERE NOT EXISTS (
                SELECT * FROM [{main_table}] main
                WHERE main.prov = temp.prov AND main.city = temp.city
            )
            """
        )
    ).fetchall()
    print(inserted)
    """console output:
    [('AB', 'Edmonton', 2)]
    """

# <check results>
with engine.begin() as conn:
    pprint(conn.execute(sa.text(f"SELECT * FROM {main_table}")).fetchall())
    """console output:
    [(1, 'AB', 'Calgary', 'Flames', 'hard-working, handsome lads'),
     (2, 'AB', 'Edmonton', 'Oilers', 'ruffians and scalawags')]
    """

推荐阅读