Sqlite3 的装饰器

Question

我有一个简单的 sqlite3 函数

from sqlite3 import connect


def foo(name):
    conn = connect("data.db")
    curs = conn.cursor()
    curs.execute(f"CREATE TABLE IF NOT EXISTS {name}(test TEXT PRIMARY KEY);")
    conn.commit()
    conn.close()

我想要一个装饰器，这样我就可以写

from sqlite3 import connect

@db_connect
def foo(name):  # Don't know how to pass the args
    curs.execute(f"CREATE TABLE IF NOT EXISTS {name}(test TEXT PRIMARY KEY);")

目标是，我不必建立连接、关闭连接等。

我尝试过的：

def db_connect(func):
    def _db_connect(*args, **kwargs):
        conn = connect("data.db")
        curs = conn.cursor()
        result = func(*args, **kwargs)
        conn.commit()
        conn.close()
        return result
    return _db_connect

但现在我有点卡住了，因为如何将光标传递给函数以及我的装饰器会工作吗？

Answer 1

您真正需要的是上下文管理器，而不是装饰器。

import sqlite3
from contextlib import contextmanager

@contextmanager
def db_ops(db_name):
    conn = sqlite3.connect(db_name)
    try:
        cur = conn.cursor()
        yield cur
    except Exception as e:
        # do something with exception
        conn.rollback()
        raise e
    else:
        conn.commit()
    finally:
        conn.close()



with db_ops('db_path') as cur:
    cur.execute('create table if not exists temp (id int, name text)')

with db_ops('db_path') as cur:
    rows = [(1, 'a'), (2, 'b'), (3, 'c')]
    cur.executemany('insert into temp values (?, ?)', rows)

with db_ops('db_path') as cur:
    print(list(cur.execute('select * from temp')))

输出

[(1, 'a'), (2, 'b'), (3, 'c')]

如您所见，您不必再提交或创建连接。

值得注意的是，连接对象默认支持上下文管理器协议，这意味着你可以这样做

conn = sqlite3.connect(...)
with conn:
    ...

但这只是提交，并没有关闭连接，你仍然必须使用conn.close().