Use list per transaction instead of one big list for persist_on_commit (#48)

Author: Dave Flerlage

temporal_sqlalchemy/bases.py

@@ -13,7 +13,7 @@
 import psycopg2.extras as psql_extras
 
 from temporal_sqlalchemy import nine
-from temporal_sqlalchemy.metadata import get_session_metadata
+from temporal_sqlalchemy.metadata import STRICT_MODE_KEY
 
 _ClockSet = collections.namedtuple('_ClockSet', ('effective', 'vclock'))
 
@@ -96,7 +96,7 @@ def record_history(self,
         """record all history for a given clocked object"""
         new_tick = self._get_new_tick(clocked)
 
-        is_strict_mode = get_session_metadata(session).get('strict_mode', False)
+        is_strict_mode = session.info[STRICT_MODE_KEY]
         vclock_history = attributes.get_history(clocked, 'vclock')
         is_vclock_unchanged = vclock_history.unchanged and new_tick == vclock_history.unchanged[0]
 

temporal_sqlalchemy/metadata.py

@@ -1,29 +1,14 @@
 import sqlalchemy.orm as orm
 
-TEMPORAL_METADATA_KEY = '__temporal'
-
 __all__ = [
-    'get_session_metadata',
-    'set_session_metadata',
+    'STRICT_MODE_KEY',
+    'CHANGESET_STACK_KEY',
+    'IS_COMMITTING_KEY',
+    'IS_VCLOCK_UNCHANGED_KEY',
 ]
 
 
-def set_session_metadata(session: orm.Session, metadata: dict):
-    if isinstance(session, orm.Session):
-        session.info[TEMPORAL_METADATA_KEY] = metadata
-    elif isinstance(session, orm.sessionmaker):
-        session.configure(info={TEMPORAL_METADATA_KEY: metadata})
-    else:
-        raise ValueError('Invalid session')
-
-
-def get_session_metadata(session: orm.Session) -> dict:
-    """
-    :return: metadata dictionary, or None if it was never installed
-    """
-    if isinstance(session, orm.Session):
-        return session.info.get(TEMPORAL_METADATA_KEY)
-    elif isinstance(session, orm.sessionmaker):
-        return session.kw.get('info', {}).get(TEMPORAL_METADATA_KEY)
-    else:
-        raise ValueError('Invalid session')
+STRICT_MODE_KEY = '__temporal_strict_mode'
+CHANGESET_STACK_KEY = '__temporal_changeset_stack'
+IS_COMMITTING_KEY = '__temporal_is_committing'
+IS_VCLOCK_UNCHANGED_KEY = '__temporal_is_vclock_unchanged'

temporal_sqlalchemy/session.py

@@ -1,26 +1,25 @@
 import datetime
 import itertools
 import typing
+import warnings
 
 import sqlalchemy.event as event
 import sqlalchemy.orm as orm
 
 from temporal_sqlalchemy.bases import TemporalOption, Clocked
 from temporal_sqlalchemy.metadata import (
-    get_session_metadata,
-    set_session_metadata
+    STRICT_MODE_KEY,
+    CHANGESET_STACK_KEY,
+    IS_COMMITTING_KEY,
+    IS_VCLOCK_UNCHANGED_KEY,
 )
 
 
-CHANGESET_KEY = 'changeset'
-IS_COMMITTING_KEY = 'is_committing'
-IS_VCLOCK_UNCHANGED = 'is_vclock_unchanged'
+def get_current_changeset(session):
+    stack = session.info[CHANGESET_STACK_KEY]
+    assert len(stack) > 0
 
-
-def _reset_flags(metadata):
-    metadata[CHANGESET_KEY] = {}
-    metadata[IS_COMMITTING_KEY] = False
-    metadata[IS_VCLOCK_UNCHANGED] = True
+    return stack[-1]
 
 
 def _temporal_models(session: orm.Session) -> typing.Iterable[Clocked]:
@@ -30,13 +29,17 @@ def _temporal_models(session: orm.Session) -> typing.Iterable[Clocked]:
 
 
 def _build_history(session, correlate_timestamp):
-    metadata = get_session_metadata(session)
+    # this shouldn't happen, but it might happen, log a warning and investigate
+    if not session.info.get(CHANGESET_STACK_KEY):
+        warnings.warn('changeset_stack is missing but we are in _build_history()')
+        return
 
-    items = list(metadata[CHANGESET_KEY].items())
-    metadata[CHANGESET_KEY].clear()
+    changeset = get_current_changeset(session)
+    items = list(changeset.items())
+    changeset.clear()
 
-    is_strict_mode = metadata.get('strict_mode', False)
-    is_vclock_unchanged = metadata.get(IS_VCLOCK_UNCHANGED, False)
+    is_strict_mode = session.info[STRICT_MODE_KEY]
+    is_vclock_unchanged = session.info[IS_VCLOCK_UNCHANGED_KEY]
     if items and is_strict_mode:
         assert not is_vclock_unchanged, \
             'commit() has triggered for a changed temporalized property without a clock tick'
@@ -49,45 +52,102 @@ def persist_history(session: orm.Session, flush_context, instances):
     if any(_temporal_models(session.deleted)):
         raise ValueError("Cannot delete temporal objects.")
 
+    # its possible the temporal session was initialized after the transaction has started
+    _initialize_metadata(session)
+
     correlate_timestamp = datetime.datetime.now(tz=datetime.timezone.utc)
     changed_rows = _temporal_models(itertools.chain(session.dirty, session.new))
 
-    metadata = get_session_metadata(session)
+    changeset = get_current_changeset(session)
     for obj in changed_rows:
         if obj.temporal_options.allow_persist_on_commit:
             new_changes, is_vclock_unchanged = obj.temporal_options.get_history(obj)
 
             if new_changes:
-                if obj not in metadata[CHANGESET_KEY]:
-                    metadata[CHANGESET_KEY][obj] = {}
+                if obj not in changeset:
+                    changeset[obj] = {}
 
-                old_changes = metadata[CHANGESET_KEY][obj]
+                old_changes = changeset[obj]
                 old_changes.update(new_changes)
 
-            metadata[IS_VCLOCK_UNCHANGED] = metadata[IS_VCLOCK_UNCHANGED] and is_vclock_unchanged
+            session.info[IS_VCLOCK_UNCHANGED_KEY] = session.info[IS_VCLOCK_UNCHANGED_KEY] and is_vclock_unchanged
         else:
             obj.temporal_options.record_history(obj, session, correlate_timestamp)
 
     # if this is the last flush, build all the history
-    if metadata[IS_COMMITTING_KEY]:
+    if session.info[IS_COMMITTING_KEY]:
         _build_history(session, correlate_timestamp)
 
-        _reset_flags(metadata)
+        session.info[IS_COMMITTING_KEY] = False
 
 
 def enable_is_committing_flag(session):
-    metadata = get_session_metadata(session)
-
-    metadata[IS_COMMITTING_KEY] = True
+    """before_commit happens before before_flush, and we need to make sure the history gets built
+    during the final one of these two events, so we need to use the gross IS_COMMITTING_KEY flag to
+    control this behavior"""
+    session.info[IS_COMMITTING_KEY] = True
 
+    # if the session is clean, a final flush won't happen, so try to build the history now
     if session._is_clean():
         correlate_timestamp = datetime.datetime.now(tz=datetime.timezone.utc)
         _build_history(session, correlate_timestamp)
 
     # building the history can cause the session to be dirtied, which will in turn call another
     # flush(), so we want to check this before reseting
+    # if there are more changes, flush will build them itself
     if session._is_clean():
-        _reset_flags(metadata)
+        session.info[IS_COMMITTING_KEY] = False
+
+
+def _get_transaction_stack_depth(transaction):
+    depth = 0
+
+    current = transaction
+    while current:
+        depth += 1
+        current = transaction.parent
+
+    return depth
+
+
+def _initialize_metadata(session):
+    if CHANGESET_STACK_KEY not in session.info:
+        session.info[CHANGESET_STACK_KEY] = []
+
+    if IS_COMMITTING_KEY not in session.info:
+        session.info[IS_COMMITTING_KEY] = False
+
+    if IS_VCLOCK_UNCHANGED_KEY not in session.info:
+        session.info[IS_VCLOCK_UNCHANGED_KEY] = True
+
+    # sometimes temporalize a session after a transaction has already been open, so we need to
+    # backfill any missing stack entries
+    if len(session.info[CHANGESET_STACK_KEY]) == 0:
+        depth = _get_transaction_stack_depth(session.transaction)
+        for _ in range(depth):
+            session.info[CHANGESET_STACK_KEY].append({})
+
+
+def start_transaction(session, transaction):
+    _initialize_metadata(session)
+
+    session.info[CHANGESET_STACK_KEY].append({})
+
+
+def end_transaction(session, transaction):
+    # there are some edge cases where no temporal changes actually happen, so don't bother
+    if not session.info.get(CHANGESET_STACK_KEY):
+        return
+
+    session.info[CHANGESET_STACK_KEY].pop()
+
+    # reset bookkeeping fields if we're ending a top most transaction
+    if transaction.parent is None:
+        session.info[IS_VCLOCK_UNCHANGED_KEY] = True
+        session.info[IS_COMMITTING_KEY] = False
+
+        # there should be no more changeset stacks at this point, otherwise there is a mismatch
+        assert len(session.info[CHANGESET_STACK_KEY]) == 0
 
 
 def temporal_session(session: typing.Union[orm.Session, orm.sessionmaker], strict_mode=False) -> orm.Session:
@@ -98,23 +158,26 @@ def temporal_session(session: typing.Union[orm.Session, orm.sessionmaker], stric
     :param strict_mode: if True, will raise exceptions when improper flush() calls are made (default is False)
     :return: temporalized SQLALchemy ORM session
     """
-    temporal_metadata = {
-        'strict_mode': strict_mode,
-    }
-    _reset_flags(temporal_metadata)
-
     # defer listening to the flush hook until after we update the metadata
     install_flush_hook = not is_temporal_session(session)
 
-    # update to the latest metadata
-    set_session_metadata(session, temporal_metadata)
+    if isinstance(session, orm.Session):
+        session.info[STRICT_MODE_KEY] = strict_mode
+    elif isinstance(session, orm.sessionmaker):
+        session.configure(info={STRICT_MODE_KEY: strict_mode})
+    else:
+        raise ValueError('Invalid session')
 
     if install_flush_hook:
         event.listen(session, 'before_flush', persist_history)
         event.listen(session, 'before_commit', enable_is_committing_flag)
 
+        # nested transaction handling
+        event.listen(session, 'after_transaction_create', start_transaction)
+        event.listen(session, 'after_transaction_end', end_transaction)
+
     return session
 
 
 def is_temporal_session(session: orm.Session) -> bool:
-    return isinstance(session, orm.Session) and get_session_metadata(session) is not None
+    return isinstance(session, orm.Session) and session.info.get(STRICT_MODE_KEY) is not None

tests/models.py

@@ -275,3 +275,10 @@ class PersistOnFlushTable(PersistenceStrategy):
 class PersistOnCommitTable(PersistenceStrategy):
     __tablename__ = 'persist_on_commit_table'
     __table_args__ = {'schema': SCHEMA}
+
+
+class NonTemporalTable(Base):
+    __tablename__ = 'non_temporal_table'
+    __table_args__ = {'schema': SCHEMA}
+
+    id = auto_uuid()