Add Contextual Comments to SQLAlchemy Queries

Question

I'm trying to figure out a way to add comments to queries generated by SQLAlchemy that indicate its point of origin.

I'm aware of the prefix_with method but am unsatisfied with it since it will only place the prefix after 'SELECT' and I'd like my comment to be completely separated from the actual query logic.

I've stumbled across a solution that adds a cursor execution callback to handle this which is detailed here. My problem with this is that the type of comment I'd like to add is dependent on where the query is being executed from.

So, for example, I'd like to do something like the following...

class SomeTableCountMetric(object):

  def __init__(self):
    session = sessionmaker(bind=engine, info={ 'src' : self.__class__.__name__ })

  def get(self):
    return self.session.query(func.count(SomeTable.id)).all()

and in turn have the generated query be

SELECT COUNT(some_tables.id) FROM some_tables /* { 'src' : SomeTableCountMetric } */

So, basically, I'd like to be able to pass some arbitrary value into the session or the query object and then be able to retrieve it from one of the arguments available in the callback method (connection, cursor, or context).

Is it possible to do this or alternatively is there another way to implementing this?

Also, any thoughts on whether this seems like a good pattern?

Ivan Choo Ivan Choo · Accepted Answer · 2017-10-25T04:02:27

connection has an info dictionary where you can store user-defined data. You can store the comment info in the session, and pass it onto the associated connection via after_begin. Thereafter you can append the comment to the SQL statement prior to execution (as suggested in the other answer).

# Our session factory, typically is a scoped session
Session = sessionmaker(bind=engine)

# Pass `session` contextual info to `connection`
@event.listens_for(Session, 'after_begin')
def session_on_after_begin(session, transaction, connection):
  if 'src' in session.info:
    connection.info['src'] = session.info['src']


# Append comments to statement
@event.listens_for(engine, "before_cursor_execute", retval=True)
def comment_sql_calls(conn, cursor, statement, parameters, context, executemany):
    if 'src' in conn.info:
        statement = statement + " -- %s" % conn.info.pop('src')
    return statement, parameters


class SomeTableCountMetric(object):

  def __init__(self, session=None):
      if session is None:
          session = Session()
    self.session = session
    self.session.info['src'] = 'some info'

  def get(self):
    return self.session.query(func.count(SomeTable.id)).all()

try:
    qo = SomeTableCountMetric()
    qo.get()
    qo.session.close()
except:
    qo.session.rollback()

Add Contextual Comments to SQLAlchemy Queries

1 Answers