diff --git a/synapse/storage/_base.py b/synapse/storage/_base.py
index 579ed56377..ccf9697fa3 100644
--- a/synapse/storage/_base.py
+++ b/synapse/storage/_base.py
@@ -185,12 +185,16 @@ class LoggingTransaction(object):
"""An object that almost-transparently proxies for the 'txn' object
passed to the constructor. Adds logging and metrics to the .execute()
method."""
- __slots__ = ["txn", "name", "database_engine"]
+ __slots__ = ["txn", "name", "database_engine", "after_callbacks"]
- def __init__(self, txn, name, database_engine):
+ def __init__(self, txn, name, database_engine, after_callbacks):
object.__setattr__(self, "txn", txn)
object.__setattr__(self, "name", name)
object.__setattr__(self, "database_engine", database_engine)
+ object.__setattr__(self, "after_callbacks", after_callbacks)
+
+ def call_after(self, callback, *args):
+ self.after_callbacks.append((callback, args))
def __getattr__(self, name):
return getattr(self.txn, name)
@@ -336,6 +340,8 @@ class SQLBaseStore(object):
start_time = time.time() * 1000
+ after_callbacks = []
+
def inner_func(conn, *args, **kwargs):
with LoggingContext("runInteraction") as context:
if self.database_engine.is_connection_closed(conn):
@@ -360,10 +366,10 @@ class SQLBaseStore(object):
while True:
try:
txn = conn.cursor()
- return func(
- LoggingTransaction(txn, name, self.database_engine),
- *args, **kwargs
+ txn = LoggingTransaction(
+ txn, name, self.database_engine, after_callbacks
)
+ return func(txn, *args, **kwargs)
except self.database_engine.module.OperationalError as e:
# This can happen if the database disappears mid
# transaction.
@@ -412,6 +418,8 @@ class SQLBaseStore(object):
result = yield self._db_pool.runWithConnection(
inner_func, *args, **kwargs
)
+ for after_callback, after_args in after_callbacks:
+ after_callback(*after_args)
defer.returnValue(result)
def cursor_to_dict(self, cursor):
|
