0

I'd like to find a good pattern to defer execution actions until after the current transaction successfully commits. For example, I may have:

with pool(...)  as connection:
    execute_steps(connection)

in which a sub-function start one of more any asynchronous tasks by submitting them to an external message queue (rabbit/AMQP, ...) ... ONLY if the transaction successfully commits. That is 

 def substep(connection): 
       submit_to_queue()

... may fail if the transaction subsequently fails (after submitting to the queue), or if the job placed in the queue executes before the transaction commits. So -- submit_to_queue MUST execute after connection.commit.

A natural way to do this is to collect a step of callbacks:

def substep(connection):

  def submit_to_queue(): ...
  ...
  connection.on_after_commit(submit_to_queue)
  ...

and have:

@contextmanager
def db_pool(db_name="lb2", current_connection=None):

    connection = connection_pool.getconn()
    connection.on_after_commit = []
    try:

        # Do the work
        yield connection

        connection.commit()

        # Reset the after_success list and execute
        for job in connection.after_success():
            job()
        connection.after_success = [], 

    except Exception:
        connection.rollback()
        connection.after_success = []
        raise

However, psycopg2's connection is an immutable object. So... is there another way to register post-execution callbacks with psycopg2?

Thanks.

user48956
  • 14,850
  • 19
  • 93
  • 154
  • Do your callbacks have to be attached to the `connection` object? – Blender Feb 03 '18 at 21:22
  • Only in the sense that they need to be executed after a successful connection.commit. However, since tons of my services already take this as the first parameter, its a convenient place. I could wrap it with a connection+callbacks object, however, this may mean changing the interface to many function that only care about the connection. – user48956 Feb 03 '18 at 21:25

1 Answers1

1

Subclassing psycopg2.extensions.connection will allow you to add attributes to the object (see the explanation here as to why) :

In [44]: import psycopg2

In [45]: import psycopg2.extensions

In [46]: conn = psycopg2.connect('dbname=postgres')

In [47]: conn.foo = 'bar'
---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
<ipython-input-47-3ab9f2c6e77e> in <module>()
----> 1 conn.foo = 'bar'

AttributeError: 'psycopg2.extensions.connection' object has no attribute 'foo'

In [48]: class MyConnection(psycopg2.extensions.connection):
   ...       pass
   ...

In [49]: conn2 = MyConnection('dbname=postgres')

In [50]: conn2.foo = 'bar'

In [51]: conn2.foo
Out[51]: 'bar'
Blender
  • 289,723
  • 53
  • 439
  • 496