使用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文档with暗示with connection.cursor() as cursor:将在事务中包装该代码.这显然没有发生.实验告诉我,我需要2级with语句,包括pscyopg2.connect调用,然后我得到一个事务.

如何传递一个connection对象并在自己的事务中运行查询以允许这种"优雅的错误处理"？我想以"干净的架构"风格将postgres连接代码分开.这可能吗？

Answer 1

with语句的psycogpg2文档暗示with connection.cursor()as cursor:将该代码包装在事务中.

这实际上不是真的它说:

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

当连接退出with块时,如果块没有引发异常,则提交事务.如果发生异常,则回滚事务.在任何情况下,连接都不会关闭:连接可以在多个with语句中使用,每个with块都可以有效地包含在事务中.

所以它不是由with连接对象处理的游标对象

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

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为执行的每个查询包装您的连接,因此如果连接失败,则上下文管理器将还原事务