[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