每次在 psycopg2 中获取数据库游标时设置自定义运行时参数

Question

我希望每次获取连接的游标时都在 postgresql 中设置自定义运行时参数。

我在 postgresql.conf 中添加了变量，如下所示：

currentuser.name = 'not set'
currentuser.reasonid = -1

现在，我目前的解决方案是创建自定义游标 class，然后将其作为参数传递给 cursor() 函数：

import psycopg2
import psycopg2.extensions

class mycursor(psycopg2.extensions.cursor):

    def __init__(self, username, *args, **kwargs):
        username = kwargs['username']
        del kwargs['username']
        super(mycursor, self).__init__(*args, **kwargs)
        self.execute("select set_config('currentuser.name', %s, false)", [username])


connection = psycopg2.connect(database="my_database", 
                              user="my_user", 
                              password="my_password", 
                              host='127.0.0.1')

cursor = connection.cursor(cursor_factory=mycursor, username='michael')

cursor.execute("select current_setting('currentuser.name')")
user = cursor.fetchall()[0][0]
print user
cursor.close()

这会产生 TypeError:

Traceback (most recent call last):
 File "customdb.py", line 22, in <module>
   cursor = connection.cursor(cursor_factory=mycursor, username='michael')
TypeError: 'username' is an invalid keyword argument for this function

Answer 1

connection.cursor 的签名不允许任何附加参数。这解释了你的错误。

子 class 连接本身会更简单，因为这样您可以轻松覆盖方法 cursor。这是我的实现，它将所有处理委托给嵌入式 connection 对象，但游标创建除外：

import psycopg2
import psycopg2.extensions

class myconnection(psycopg2.extensions.connection):
    def __init__(self, database, user, password, host = 'localhost',
                port=5432, **kwargs):
        self.__dict__['conn'] = psycopg2.connect(dsn = None, database = database, user=user,
                                password=password, host=host, port=port, **kwargs)
        self.__dict__['username'] = user

    def cursor(self, username=None, **kwargs):
        curs = self.conn.cursor(**kwargs)
        if username is None:
            username = self.username 
        curs.execute("select set_config('currentuser.name', %s, false)", [username])
        return curs

    def __getattr__(self, name):
        return getattr(self.conn, name)

    def __setattr__(self, name, value):
        setattr(self.conn, name, value)


connection = myconnection(database="my_database", 
                          user="my_user", 
                          password="my_password", 
                          host='127.0.0.1')


cursor = connection.cursor(username='michael')

cursor.execute("select current_setting('currentuser.name')")
user = cursor.fetchall()[0][0]
print user
cursor.close()

它正确打印 michael

每次在 psycopg2 中获取数据库游标时设置自定义运行时参数

Setting a custom runtime parameter every time I get the db cursor in psycopg2

python

postgresql

psycopg2

python-2.7