CODEWITHSUNDEEP
pythonpostgresqlpython-2.7psycopg2
Michał Czapliński
Michał Czapliński

Reputation: 1342

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

I wish to set a custom runtime parameter in postgresql every time I get the cursor for a connection.

I added the variables in postgresql.conf like so:

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

Now, my solution so far was to create a custom cursor class and then pass it as argument to the cursor() function:

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()

This produces the 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

Upvotes: 2

Views: 736

Answers (1)

Serge Ballesta
Serge Ballesta

Reputation: 149195

The signature for connection.cursor does not allow any additional parameter. That explains your error.

It would be simpler to sub-class the connection itself, because that way you could easily override the method cursor. Here is my implementation, that delegates all processing to an embedded connection object, except for the cursor creation :

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()

It correctly prints michael

Upvotes: 3

Related Questions