feat: add use_session() for per-transaction session binding
diff --git a/README.md b/README.md index 43dbe69..a1df14c 100644 --- a/README.md +++ b/README.md
@@ -143,6 +143,37 @@ await session.commit() ``` +### Binding a Session per Transaction: `use_session()` + +Passing `db_session` to the constructor binds the session for the whole life of the +adapter, so it forces you to build a new adapter (and a new enforcer, and reload the +policy) for every transaction. When the enforcer is a long-lived singleton — the usual +setup in a FastAPI/Starlette app — use `adapter.use_session()` instead. Inside the +block every write the adapter performs joins your transaction, and the adapter neither +commits nor rolls back, so policy changes live or die together with your own changes: + +```python +# adapter and enforcer are created once at startup and reused +async def create_user(db_session, ...): + async with db_session.begin(): + user = User(...) + db_session.add(user) + await db_session.flush() # user.id is available, nothing is committed yet + + async with adapter.use_session(db_session): + await e.add_role_for_user(str(user.id), "admin") + await e.add_policies([[str(user.id), "data1", "read"]]) + + # If anything below raises, the user *and* the policies are rolled back + await notify(user) +``` + +The binding is stored in a `contextvars.ContextVar`, so concurrent requests sharing the +same adapter each keep their own session; outside the block the adapter goes back to +opening and committing its own sessions. Note that the enforcer's in-memory policy is +updated as usual — if the transaction is rolled back, call `await e.load_policy()` to +resynchronise it with the database. + ### Batch Operations Example ```python
diff --git a/casbin_async_sqlalchemy_adapter/adapter.py b/casbin_async_sqlalchemy_adapter/adapter.py index 884103f..f3287f6 100644 --- a/casbin_async_sqlalchemy_adapter/adapter.py +++ b/casbin_async_sqlalchemy_adapter/adapter.py
@@ -12,7 +12,8 @@ # See the License for the specific language governing permissions and # limitations under the License. from contextlib import asynccontextmanager -from typing import List, Optional +from contextvars import ContextVar +from typing import Dict, List, Optional from casbin import persist from casbin.persist.adapters.asyncio import AsyncAdapter @@ -24,6 +25,11 @@ Base = declarative_base() +# Sessions bound through Adapter.use_session(). Keyed by id(adapter) so that several +# adapters can be bound independently, and stored in a ContextVar so that concurrent +# tasks (e.g. web requests) never see each other's session. +_bound_sessions: ContextVar[Dict[int, AsyncSession]] = ContextVar("casbin_bound_sessions", default={}) + class CasbinRule(Base): __tablename__ = "casbin_rule" @@ -138,11 +144,47 @@ self._filtered = filtered @asynccontextmanager + async def use_session(self, session: AsyncSession): + """Temporarily run this adapter on an externally managed session. + + Every write the adapter performs inside the block joins the caller's + transaction and is neither committed nor rolled back by the adapter, so + policy changes can be made atomic together with the caller's own changes:: + + async with db_session.begin(): + user = User(...) + db_session.add(user) + await db_session.flush() + + async with adapter.use_session(db_session): + await enforcer.add_role_for_user(str(user.id), "admin") + + The binding is stored in a :class:`~contextvars.ContextVar`, so a single + long-lived adapter (and enforcer) can be shared by concurrent tasks while + each task keeps its own session. + """ + sessions = dict(_bound_sessions.get()) + sessions[id(self)] = session + token = _bound_sessions.set(sessions) + try: + yield session + finally: + _bound_sessions.reset(token) + + def _current_external_session(self) -> Optional[AsyncSession]: + """Return the externally managed session in effect, if any.""" + session = _bound_sessions.get().get(id(self)) + if session is not None: + return session + return self._external_session + + @asynccontextmanager async def _session_scope(self): """Provide an asynchronous transactional scope around a series of operations.""" - if self._external_session is not None: + session = self._current_external_session() + if session is not None: # Use external session without automatic commit/rollback - yield self._external_session + yield session else: # Use internal session with automatic commit/rollback async with self.session_local() as session: @@ -301,6 +343,10 @@ if not rules: return + async with self._session_scope() as session: + await self._add_policies(session, ptype, rules) + + async def _add_policies(self, session, ptype, rules): # Build rows for executemany bulk insert rows = [] for rule in rules: @@ -309,9 +355,8 @@ row[f"v{i}"] = v rows.append(row) - async with self._session_scope() as session: - stmt = insert(self._db_class) - await session.execute(stmt, rows) + stmt = insert(self._db_class) + await session.execute(stmt, rows) async def remove_policy(self, sec, ptype, rule): """removes a policy rule from the storage.""" @@ -338,22 +383,25 @@ if not rules: return async with self._session_scope() as session: - if self.softdelete_attribute is None: - stmt = delete(self._db_class).where(self._db_class.ptype == ptype) - rules_zipped = zip(*rules) - for i, rule in enumerate(rules_zipped): - stmt = stmt.where(or_(getattr(self._db_class, "v{}".format(i)) == v for v in rule)) - await session.execute(stmt) - else: - stmt = select(self._db_class).where(self._db_class.ptype == ptype) - stmt = self._softdelete_query(stmt) - rules_zipped = zip(*rules) - for i, rule in enumerate(rules_zipped): - stmt = stmt.where(or_(getattr(self._db_class, "v{}".format(i)) == v for v in rule)) - result = await session.execute(stmt) - lines = result.scalars().all() - for line in lines: - setattr(line, self.softdelete_attribute.name, True) + await self._remove_policies(session, ptype, rules) + + async def _remove_policies(self, session, ptype, rules): + if self.softdelete_attribute is None: + stmt = delete(self._db_class).where(self._db_class.ptype == ptype) + rules_zipped = zip(*rules) + for i, rule in enumerate(rules_zipped): + stmt = stmt.where(or_(getattr(self._db_class, "v{}".format(i)) == v for v in rule)) + await session.execute(stmt) + else: + stmt = select(self._db_class).where(self._db_class.ptype == ptype) + stmt = self._softdelete_query(stmt) + rules_zipped = zip(*rules) + for i, rule in enumerate(rules_zipped): + stmt = stmt.where(or_(getattr(self._db_class, "v{}".format(i)) == v for v in rule)) + result = await session.execute(stmt) + lines = result.scalars().all() + for line in lines: + setattr(line, self.softdelete_attribute.name, True) async def remove_filtered_policy(self, sec, ptype, field_index, *field_values): """removes policy rules that match the filter from the storage. @@ -399,25 +447,28 @@ """ async with self._session_scope() as session: - stmt = select(self._db_class).where(self._db_class.ptype == ptype) - stmt = self._softdelete_query(stmt) + await self._update_policy(session, ptype, old_rule, new_rule) - # locate the old rule - for index, value in enumerate(old_rule): - v_value = getattr(self._db_class, "v{}".format(index)) - stmt = stmt.where(v_value == value) + async def _update_policy(self, session, ptype: str, old_rule: List[str], new_rule: List[str]) -> None: + stmt = select(self._db_class).where(self._db_class.ptype == ptype) + stmt = self._softdelete_query(stmt) - # need the length of the longest_rule to perform overwrite - longest_rule = old_rule if len(old_rule) > len(new_rule) else new_rule - result = await session.execute(stmt) - old_rule_line = result.scalar_one() + # locate the old rule + for index, value in enumerate(old_rule): + v_value = getattr(self._db_class, "v{}".format(index)) + stmt = stmt.where(v_value == value) - # overwrite the old rule with the new rule - for index in range(len(longest_rule)): - if index < len(new_rule): - setattr(old_rule_line, "v{}".format(index), new_rule[index]) - else: - setattr(old_rule_line, "v{}".format(index), None) + # need the length of the longest_rule to perform overwrite + longest_rule = old_rule if len(old_rule) > len(new_rule) else new_rule + result = await session.execute(stmt) + old_rule_line = result.scalar_one() + + # overwrite the old rule with the new rule + for index in range(len(longest_rule)): + if index < len(new_rule): + setattr(old_rule_line, "v{}".format(index), new_rule[index]) + else: + setattr(old_rule_line, "v{}".format(index), None) async def update_policies( self, @@ -436,8 +487,10 @@ :return: None """ - for i in range(len(old_rules)): - await self.update_policy(sec, ptype, old_rules[i], new_rules[i]) + # A single scope keeps the whole batch inside one transaction. + async with self._session_scope() as session: + for i in range(len(old_rules)): + await self._update_policy(session, ptype, old_rules[i], new_rules[i]) async def update_filtered_policies(self, sec, ptype, new_rules: List[List[str]], field_index, *field_values) -> List[List[str]]: """update_filtered_policies updates all the policies on the basis of the filter.""" @@ -482,11 +535,13 @@ # Delete old policies - await self.remove_policies("p", filter.ptype, old_rules) + if old_rules: + await self._remove_policies(session, filter.ptype, old_rules) # Insert new policies - await self.add_policies("p", filter.ptype, new_rules) + if new_rules: + await self._add_policies(session, filter.ptype, new_rules) # return deleted rules
diff --git a/tests/test_transaction.py b/tests/test_transaction.py new file mode 100644 index 0000000..011cbba --- /dev/null +++ b/tests/test_transaction.py
@@ -0,0 +1,166 @@ +# Copyright 2023 The casbin Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Unit tests for Adapter.use_session(), i.e. per-call transaction control.""" + +import asyncio +import os +import unittest +from unittest import IsolatedAsyncioTestCase + +import casbin +from sqlalchemy import func, select +from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine + +from casbin_async_sqlalchemy_adapter import Adapter, CasbinRule + + +def get_fixture(path): + dir_path = os.path.split(os.path.realpath(__file__))[0] + "/" + return os.path.abspath(dir_path + path) + + +async def get_enforcer_and_session_factory(): + """A long-lived adapter/enforcer, like an app would keep as a singleton.""" + engine = create_async_engine("sqlite+aiosqlite://", future=True) + adapter = Adapter(engine) + await adapter.create_table() + + e = casbin.AsyncEnforcer(get_fixture("rbac_model.conf"), adapter) + await e.load_policy() + + session_factory = async_sessionmaker(engine, expire_on_commit=False, class_=AsyncSession) + return e, adapter, session_factory + + +async def count_rules(session_factory): + async with session_factory() as session: + result = await session.execute(select(func.count()).select_from(CasbinRule)) + return result.scalar() + + +class TestUseSession(IsolatedAsyncioTestCase): + async def test_rollback_discards_policy_changes(self): + e, adapter, session_factory = await get_enforcer_and_session_factory() + + async with session_factory() as session: + async with adapter.use_session(session): + await e.add_policy("alice", "data1", "read") + await e.add_grouping_policy("alice", "data2_admin") + await session.rollback() + + self.assertEqual(0, await count_rules(session_factory)) + + async def test_commit_persists_policy_changes(self): + e, adapter, session_factory = await get_enforcer_and_session_factory() + + async with session_factory() as session: + async with adapter.use_session(session): + await e.add_policy("alice", "data1", "read") + await e.add_policies([["bob", "data2", "write"], ["carol", "data3", "read"]]) + await session.commit() + + self.assertEqual(3, await count_rules(session_factory)) + + # A brand-new enforcer sees the committed rules. + new_enforcer = casbin.AsyncEnforcer(get_fixture("rbac_model.conf"), adapter) + await new_enforcer.load_policy() + self.assertTrue(new_enforcer.enforce("alice", "data1", "read")) + self.assertTrue(new_enforcer.enforce("bob", "data2", "write")) + + async def test_binding_is_released_after_block(self): + e, adapter, session_factory = await get_enforcer_and_session_factory() + + async with session_factory() as session: + async with adapter.use_session(session): + await e.add_policy("alice", "data1", "read") + await session.rollback() + + self.assertIsNone(adapter._current_external_session()) + + # Back to the default behaviour: the adapter opens and commits its own session. + await e.add_policy("bob", "data2", "write") + self.assertEqual(1, await count_rules(session_factory)) + + async def test_removal_participates_in_the_transaction(self): + e, adapter, session_factory = await get_enforcer_and_session_factory() + + await e.add_policies([["alice", "data1", "read"], ["bob", "data2", "write"]]) + self.assertEqual(2, await count_rules(session_factory)) + + async with session_factory() as session: + async with adapter.use_session(session): + await e.remove_policy("alice", "data1", "read") + await e.remove_filtered_policy(0, "bob") + await session.rollback() + + self.assertEqual(2, await count_rules(session_factory)) + + async def test_update_policies_is_a_single_transaction(self): + e, adapter, session_factory = await get_enforcer_and_session_factory() + + await e.add_policies([["alice", "data1", "read"], ["bob", "data2", "write"]]) + + async with session_factory() as session: + async with adapter.use_session(session): + await e.update_policies( + [["alice", "data1", "read"], ["bob", "data2", "write"]], + [["alice", "data1", "write"], ["bob", "data2", "read"]], + ) + await session.rollback() + + new_enforcer = casbin.AsyncEnforcer(get_fixture("rbac_model.conf"), adapter) + await new_enforcer.load_policy() + self.assertTrue(new_enforcer.enforce("alice", "data1", "read")) + self.assertFalse(new_enforcer.enforce("alice", "data1", "write")) + + async def test_binding_does_not_leak_into_other_tasks(self): + _, adapter, session_factory = await get_enforcer_and_session_factory() + seen_by_other_task = [] + bound = asyncio.Event() + checked = asyncio.Event() + + async def other_task(): + await bound.wait() + seen_by_other_task.append(adapter._current_external_session()) + checked.set() + + task = asyncio.create_task(other_task()) + async with session_factory() as session: + async with adapter.use_session(session): + bound.set() + await checked.wait() + await task + + self.assertEqual([None], seen_by_other_task) + + async def test_constructor_session_still_works(self): + """db_session= passed to the constructor keeps its previous behaviour.""" + engine = create_async_engine("sqlite+aiosqlite://", future=True) + session_factory = async_sessionmaker(engine, expire_on_commit=False, class_=AsyncSession) + + async with session_factory() as session: + adapter = Adapter(engine, db_session=session) + await adapter.create_table() + e = casbin.AsyncEnforcer(get_fixture("rbac_model.conf"), adapter) + await e.load_policy() + + await e.add_policy("alice", "data1", "read") + await session.rollback() + + self.assertEqual(0, await count_rules(session_factory)) + + +if __name__ == "__main__": + unittest.main()