Support connection_factory?

[bandophahita]

I'll start with the question as I am not entirely certain if it is even possible; Is there a way to create a queries.Session using a connection_factory? And if not, would it be possible to add such support?

Or....
perhaps there is another way to accomplish the following...?

In cases where I wanted to use a cursor which supported both a NamedTuple result and also using a LoggingCursor I would combine the two cursors but I would also pass in the connection_factory which sets up the connection to utilize logging for the cursor. Is there some way of doing this in queries?

from psycopg2.extras import NamedTupleCursor, LoggingCursor, LoggingConnection


class MyLoggingCursor(LoggingCursor):
    def execute(self, query, vars=None):
        self.timestamp = time.time()
        return super(MyLoggingCursor, self).execute(query, vars)

    def callproc(self, procname, vars=None):
        self.timestamp = time.time()
        return super(MyLoggingCursor, self).callproc(procname, vars)


class MyLoggingConnection(LoggingConnection):
    def filter(self, msg, curs):
        duration = int((time.time() - curs.timestamp) * 1000)
        output = f"{msg}  ==> {curs.rowcount} rows, {duration:d} ms"
        return output

    def cursor(self, *args, **kwargs):
        kwargs.setdefault('cursor_factory', MixinLoggedNamedTupleCursor)
        return LoggingConnection.cursor(self, *args, **kwargs)


class MixinLoggedNamedTupleCursor(MyLoggingCursor, NamedTupleCursor):
    pass


db_conn = psycopg2.connect(host=db_host, port=db_port,
                           user=db_user, password=db_pass,
                           database=db_name, 
                           connect_timeout=timeout,
                           connection_factory=MyLoggingConnection
                           )
db_conn.initialize(logger)

[ljluestc]

1. Set up your custom cursor and connection classes (MyLoggingCursor, MixinLoggedNamedTupleCursor, and MyLoggingConnection) using psycopg2.

2. Create a psycopg2 connection with your custom connection_factory.

3. Use the established psycopg2 connection to create a queries.Session object.

4. Continue using the queries.Session as usual for your queries.

Here's a modified version of your code to illustrate this approach:

import psycopg2
from psycopg2.extras import NamedTupleCursor, LoggingCursor, LoggingConnection
import time
import queries

class MyLoggingCursor(LoggingCursor):
    # ... your custom cursor logic ...

class MyLoggingConnection(LoggingConnection):
    # ... your custom connection logic ...

class MixinLoggedNamedTupleCursor(MyLoggingCursor, NamedTupleCursor):
    pass

# Create a psycopg2 connection with the custom connection factory
db_conn = psycopg2.connect(
    host=db_host,
    port=db_port,
    user=db_user,
    password=db_pass,
    database=db_name,
    connect_timeout=timeout,
    connection_factory=MyLoggingConnection
)

# Create a queries.Session using the established psycopg2 connection
db_session = queries.Session(db_conn)

# Now you can use db_session for your queries
for row in db_session.query('SELECT * FROM foo'):
    print(row)

# Don't forget to close the connection and session when done
db_session.close()
db_conn.close()