python - 在 psycopg2 中使用 with 语句创建事务

Question

我正在尝试使用 psycopg2向表中添加一些新列。 PostgreSQL 缺少 ALTER TABLE table ADD COLUMN IF NOT EXISTS，所以我在它自己的事务中添加每一列。如果该列存在，将出现 python 和 postgres 错误，没关系，我希望我的程序继续并尝试添加下一列。目标是实现幂等性，因此它可以连续运行多次。

目前看起来是这样的:

def main():
    # <snip>
    with psycopg2.connect("") as connection:
        create_columns(connection, args.table)

def create_columns(connection, table_name):
    def sql(sql):
        with connection.cursor() as cursor:
            cursor.execute(sql.format(table_name=table_name))

    sql("ALTER TABLE {table_name} ADD COLUMN my_new_col numeric(10,0);")
    sql("ALTER TABLE {table_name} ADD COLUMN another_new_col INTEGER NOT NULL;")

但是，如果my_new_col存在，则会出现异常ProgrammingError('column "parent_osm_id"of relation "relations"already exists\n',)，即符合预期，但是当它尝试添加 another_new_col 时，出现异常 InternalError('current transaction is aborted, commands ignored until end of transaction block\n',)。

psycogpg2 document for the with statement暗示 with connection.cursor() as cursor: 会将代码包装在事务中。这显然没有发生。实验表明我需要 2 个级别的 with 语句，包括 pscyopg2.connect 调用，然后我得到一个事务。

我怎样才能传递 connection object并在自己的事务中运行查询以允许这种“优雅的错误处理”？我想以“干净的架构”风格将 postgres 连接代码分开。这可能吗？

Answer 1

The psycogpg2 document for the with statement implies that the with connection.cursor() as cursor: will wrap that code in a transaction.

这实际上不是真的，它说:

with psycopg2.connect(DSN) as conn:
    with conn.cursor() as curs:
       curs.execute(SQL)

When a connection exits the with block, if no exception has been raised by the block, the transaction is committed. In case of exception the transaction is rolled back. In no case the connection is closed: a connection can be used in more than a with statement and each with block is effectively wrapped in a transaction.

所以这不是由 with 处理的游标对象，而是连接对象

还值得注意的是，当我们离开 with 子句时，cursor 持有的所有资源都将被释放

When a cursor exits the with block it is closed, releasing any resource eventually associated with it. The state of the transaction is not affected.

所以回到你的代码，你可能会重写它更像是:

def main():
    # <snip>
    with psycopg2.connect("") as connection:
        create_columns(connection, args.table)

def create_columns(con, table_name):
    def sql(connection, sql):
        with connection:
            with connection.cursor() as cursor:
                cursor.execute(sql.format(table_name=table_name))

    sql(con, "ALTER TABLE {table_name} ADD COLUMN my_new_col numeric(10,0);")
    sql(con, "ALTER TABLE {table_name} ADD COLUMN another_new_col INTEGER NOT NULL;")

确保您执行的每个查询的连接都包含在 with 中，因此如果它失败，连接上下文管理器将恢复事务