chore(dao/command): Add transaction decorator to try to enforce "unit of work" (#24969)

tests/unit_tests/utils/lock_tests.py

@@ -22,8 +22,8 @@ from uuid import UUID
 
 import pytest
 from freezegun import freeze_time
-from sqlalchemy.orm import Session, sessionmaker
 
+from superset import db
 from superset.exceptions import CreateKeyValueDistributedLockFailedException
 from superset.key_value.types import JsonKeyValueCodec
 from superset.utils.lock import get_key, KeyValueDistributedLock
@@ -32,56 +32,51 @@ MAIN_KEY = get_key("ns", a=1, b=2)
 OTHER_KEY = get_key("ns2", a=1, b=2)
 
 
-def _get_lock(key: UUID, session: Session) -> Any:
+def _get_lock(key: UUID) -> Any:
     from superset.key_value.models import KeyValueEntry
 
-    entry = session.query(KeyValueEntry).filter_by(uuid=key).first()
+    entry = db.session.query(KeyValueEntry).filter_by(uuid=key).first()
     if entry is None or entry.is_expired():
         return None
 
     return JsonKeyValueCodec().decode(entry.value)
 
 
-def _get_other_session() -> Session:
-    # This session is used to simulate what another worker will find in the metastore
-    # during the locking process.
-    from superset import db
-
-    bind = db.session.get_bind()
-    SessionMaker = sessionmaker(bind=bind)
-    return SessionMaker()
-
-
 def test_key_value_distributed_lock_happy_path() -> None:
     """
     Test successfully acquiring and returning the distributed lock.
+
+    Note we use a nested transaction to ensure that the cleanup from the outer context
+    manager is correctly invoked, otherwise a partial rollback would occur leaving the
+    database in a fractured state.
     """
-    session = _get_other_session()
 
     with freeze_time("2021-01-01"):
-        assert _get_lock(MAIN_KEY, session) is None
+        assert _get_lock(MAIN_KEY) is None
+
         with KeyValueDistributedLock("ns", a=1, b=2) as key:
             assert key == MAIN_KEY
-            assert _get_lock(key, session) is True
-            assert _get_lock(OTHER_KEY, session) is None
-            with pytest.raises(CreateKeyValueDistributedLockFailedException):
-                with KeyValueDistributedLock("ns", a=1, b=2):
-                    pass
+            assert _get_lock(key) is True
+            assert _get_lock(OTHER_KEY) is None
 
-        assert _get_lock(MAIN_KEY, session) is None
+            with db.session.begin_nested():
+                with pytest.raises(CreateKeyValueDistributedLockFailedException):
+                    with KeyValueDistributedLock("ns", a=1, b=2):
+                        pass
+
+        assert _get_lock(MAIN_KEY) is None
 
 
 def test_key_value_distributed_lock_expired() -> None:
     """
     Test expiration of the distributed lock
     """
-    session = _get_other_session()
 
-    with freeze_time("2021-01-01T"):
-        assert _get_lock(MAIN_KEY, session) is None
+    with freeze_time("2021-01-01"):
+        assert _get_lock(MAIN_KEY) is None
         with KeyValueDistributedLock("ns", a=1, b=2):
-            assert _get_lock(MAIN_KEY, session) is True
-            with freeze_time("2022-01-01T"):
-                assert _get_lock(MAIN_KEY, session) is None
+            assert _get_lock(MAIN_KEY) is True
+            with freeze_time("2022-01-01"):
+                assert _get_lock(MAIN_KEY) is None
 
-        assert _get_lock(MAIN_KEY, session) is None
+        assert _get_lock(MAIN_KEY) is None