Question

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

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

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

现在，我的解决方案是创建一个自定义游标类，然后将其作为参数传递给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的签名不允许任何其他参数。这解释了你的错误。

对连接本身进行子类化会更简单，因为这样可以轻松覆盖方法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中获取db游标时设置自定义运行时参数

1 个答案: