[POC] Added Support for AsyncSession in SQLAlchemy 1.4+

languages/python/sqlalchemy-oso/sqlalchemy_oso/__init__.py

@@ -1,12 +1,28 @@
 __version__ = "0.24.0"
 
-
 from .auth import register_models
 from .oso import SQLAlchemyOso
-from .session import authorized_sessionmaker
+from .session import authorized_sessionmaker, scoped_session
+from .compat import USING_SQLAlchemy_v1_4
+from .signal import do_orm_execute
 
 __all__ = [
  "register_models",
  "authorized_sessionmaker",
+ "scoped_session",
  "SQLAlchemyOso",
 ]
+
+try:
+ # Only load AsyncIO support is using SQLAlchemy => 1.4
+ if not USING_SQLAlchemy_v1_4:
+ raise ImportError
+
+ from .async_session import async_scoped_session, async_authorized_sessionmaker
+
+ __all__ += [
+ "async_scoped_session",
+ "async_authorized_sessionmaker"
+ ]
+except (ImportError, SyntaxError):
+ pass

languages/python/sqlalchemy-oso/sqlalchemy_oso/async_session.py

@@ -0,0 +1,136 @@
+"""SQLAlchemy async session classes and factories for oso."""
+from typing import Any, Callable, Optional, Type
+import logging
+
+from sqlalchemy.ext.asyncio import AsyncSession
+from sqlalchemy_oso.session import Permissions, AuthorizedSessionBase
+from sqlalchemy import orm
+
+from oso import Oso
+
+logger = logging.getLogger(__name__)
+
+
+def async_authorized_sessionmaker(
+ get_oso: Callable[[], Oso],
+ get_user: Callable[[], Any],
+ get_checked_permissions: Callable[[], Permissions],
+ class_: Type[AsyncSession] = None,
+ **kwargs,
+):
+ """AsyncSession factory for sessions with Oso authorization applied.
+
+ :param get_oso: Callable that returns the Oso instance to use for
+ authorization.
+ :param get_user: Callable that returns the user for an authorization
+ request.
+ :param get_checked_permissions: Callable that returns an optional map of
+ permissions (resource-action pairs) to
+ authorize for the session. If the callable
+ returns ``None``, no authorization will
+ be applied to the session. If a map of
+ permissions is provided, querying for
+ a SQLAlchemy model present in the map
+ will authorize results according to the
+ action specified as the value in the
+ map. E.g., providing a map of ``{Post:
+ "read", User: "view"}`` where ``Post`` and
+ ``User`` are SQLAlchemy models will apply
+ authorization to ``session.query(Post)``
+ and ``session.query(User)`` such that
+ only ``Post`` objects that the user can
+ ``"read"`` and ``User`` objects that the
+ user can ``"view"`` are fetched from the
+ database.
+ :param class_: Base class to use for sessions.
+
+ All other keyword arguments are passed through to
+ :py:func:`sqlalchemy.orm.session.sessionmaker` unchanged.
+
+ **Invariant**: the values returned by the `get_oso()`, `get_user()`, and
+ `get_checked_permissions()` callables provided to this function *must
+ remain fixed for a given session*. This prevents authorization responses
+ from changing, ensuring that the session's identity map never contains
+ unauthorized objects.
+ """
+ if class_ is None:
+ class_ = AsyncSession
+
+ # Oso, user, and checked permissions must remain unchanged for the entire
+ # session. This is to prevent unauthorized objects from ending up in the
+ # session's identity map.
+ class Sess(AuthorizedSessionBase, orm.Session): # type: ignore
+ def __init__(self, **options):
+ options.setdefault("oso", get_oso())
+ options.setdefault("user", get_user())
+ options.setdefault("checked_permissions", get_checked_permissions())
+ super().__init__(**options)
+
+ # We call sessionmaker here because sessionmaker adds a configure
+ # method to the returned session and we want to replicate that
+ # functionality.
+ return orm.sessionmaker(class_=class_, sync_session_class=Sess, **kwargs)
+
+
+def async_scoped_session(
+ get_oso: Callable[[], Oso],
+ get_user: Callable[[], Any],
+ get_checked_permissions: Callable[[], Permissions],
+ scopefunc: Optional[Callable[..., Any]] = None,
+ **kwargs,
+):
+ """Return a async scoped session maker that uses the Oso instance, user, and
+ checked permissions (resource-action pairs) as part of the scope function.
+
+ Use in place of sqlalchemy's scoped_session_.
+
+ Uses :py:func:`authorized_sessionmaker` as the factory.
+
+ :param get_oso: Callable that returns the Oso instance to use for
+ authorization.
+ :param get_user: Callable that returns the user for an authorization
+ request.
+ :param get_checked_permissions: Callable that returns an optional map of
+ permissions (resource-action pairs) to
+ authorize for the session. If the callable
+ returns ``None``, no authorization will
+ be applied to the session. If a map of
+ permissions is provided, querying for
+ a SQLAlchemy model present in the map
+ will authorize results according to the
+ action specified as the value in the
+ map. E.g., providing a map of ``{Post:
+ "read", User: "view"}`` where ``Post`` and
+ ``User`` are SQLAlchemy models will apply
+ authorization to ``session.query(Post)``
+ and ``session.query(User)`` such that
+ only ``Post`` objects that the user can
+ ``"read"`` and ``User`` objects that the
+ user can ``"view"`` are fetched from the
+ database.
+ :param scopefunc: Additional scope function to use for scoping sessions.
+ Output will be combined with the Oso, permissions
+ (resource-action pairs), and user objects.
+ :param kwargs: Additional keyword arguments to pass to
+ :py:func:`authorized_sessionmaker`.
+
+ NOTE: _baked_queries are disabled on SQLAlchemy 1.3 since the caching
+ mechanism can bypass authorization by using queries from the cache
+ that were previously baked without authorization applied. Note that
+ _baked_queries are deprecated as of SQLAlchemy 1.4.
+
+ .. _scoped_session: https://docs.sqlalchemy.org/en/13/orm/contextual.html
+
+ .. _baked_queries: https://docs.sqlalchemy.org/en/14/orm/extensions/baked.html
+ """
+ scopefunc = scopefunc or (lambda: None)
+
+ def _scopefunc():
+ checked_permissions = frozenset(get_checked_permissions().items())
+ return (get_oso(), checked_permissions, get_user(), scopefunc())
+
+ factory = async_authorized_sessionmaker(
+ get_oso, get_user, get_checked_permissions, **kwargs
+ )
+
+ return orm.scoped_session(factory, scopefunc=_scopefunc)

languages/python/sqlalchemy-oso/sqlalchemy_oso/compat.py

@@ -9,6 +9,7 @@
 
 version = parse(sqlalchemy.__version__) # type: ignore
 USING_SQLAlchemy_v1_3 = version >= parse("1.3") and version < parse("1.4")
+USING_SQLAlchemy_v1_4 = version >= parse("1.4")
 
 
 def iterate_model_classes(base_or_registry):

languages/python/sqlalchemy-oso/sqlalchemy_oso/session.py

@@ -2,15 +2,11 @@
 from typing import Any, Callable, Dict, Optional, Type
 import logging
 
-from sqlalchemy import event, inspect
 from sqlalchemy.orm import sessionmaker, Session
-from sqlalchemy.orm.util import AliasedClass
 from sqlalchemy import orm
-from sqlalchemy.sql import expression as expr
 
 from oso import Oso
 
-from sqlalchemy_oso.auth import authorize_model
 from sqlalchemy_oso.compat import USING_SQLAlchemy_v1_3
 
 logger = logging.getLogger(__name__)
@@ -271,105 +267,3 @@ class AuthorizedSession(AuthorizedSessionBase, Session):
  pass
 
 
-try:
- # TODO(gj): remove type ignore once we upgrade to 1.4-aware MyPy types.
- from sqlalchemy.orm import with_loader_criteria # type: ignore
- from sqlalchemy_oso.sqlalchemy_utils import all_entities_in_statement
-
- @event.listens_for(Session, "do_orm_execute")
- def do_orm_execute(execute_state):
- if not execute_state.is_select:
- return
-
- session = execute_state.session
-
- if not isinstance(session, AuthorizedSessionBase):
- return
- assert isinstance(session, Session)
-
- oso: Oso = session.oso_context["oso"]
- user = session.oso_context["user"]
- checked_permissions: Permissions = session.oso_context["checked_permissions"]
-
- # Early return if no authorization is to be applied.
- if checked_permissions is None:
- return
-
- entities = all_entities_in_statement(execute_state.statement)
- logger.info(f"Authorizing entities: {entities}")
- for entity in entities:
- action = checked_permissions.get(entity)
-
- # If permissions map does not specify an action to authorize for entity
- # or if the specified action is `None`, deny access.
- if action is None:
- logger.warning(f"No allowed action for entity {entity}")
- where = with_loader_criteria(entity, expr.false(), include_aliases=True)
- execute_state.statement = execute_state.statement.options(where)
- else:
- filter = authorize_model(oso, user, action, session, entity)
- if filter is not None:
- logger.info(f"Applying filter {filter} to entity {entity}")
- where = with_loader_criteria(entity, filter, include_aliases=True)
- execute_state.statement = execute_state.statement.options(where)
- else:
- logger.warning(f"Policy did not return filter for entity {entity}")
-
-
-except ImportError:
- from sqlalchemy.orm.query import Query
-
- @event.listens_for(Query, "before_compile", retval=True)
- def _before_compile(query):
- """Enable before compile hook."""
- return _authorize_query(query)
-
- def _authorize_query(query: Query) -> Optional[Query]:
- """Authorize an existing query with an Oso instance, user, and a
- permissions map indicating which actions to check for which SQLAlchemy
- models."""
- session = query.session
-
- # Early return if this isn't an authorized session.
- if not isinstance(session, AuthorizedSessionBase):
- return None
-
- oso: Oso = session.oso_context["oso"]
- user = session.oso_context["user"]
- checked_permissions: Permissions = session.oso_context["checked_permissions"]
-
- # Early return if no authorization is to be applied.
- if checked_permissions is None:
- return None
-
- # TODO (dhatch): This is necessary to allow ``authorize_query`` to work
- # on queries that have already been made. If a query has a LIMIT or OFFSET
- # applied, SQLAlchemy will by default throw an error if filters are applied.
- # This prevents these errors from occuring, but could result in some
- # incorrect queries. We should remove this if possible.
- query = query.enable_assertions(False) # type: ignore
-
- entities = {column["entity"] for column in query.column_descriptions}
- for entity in entities:
- # Only apply authorization to columns that represent a mapper entity.
- if entity is None:
- continue
-
- # If entity is an alias, get the action for the underlying class.
- if isinstance(entity, AliasedClass):
- action = checked_permissions.get(inspect(entity).class_) # type: ignore
- else:
- action = checked_permissions.get(entity)
-
- # If permissions map does not specify an action to authorize for entity
- # or if the specified action is `None`, deny access.
- if action is None:
- query = query.filter(expr.false()) # type: ignore
- continue
-
- assert isinstance(session, Session)
- authorized_filter = authorize_model(oso, user, action, session, entity)
- if authorized_filter is not None:
- query = query.filter(authorized_filter) # type: ignore
-
- return query