fastapi-rtk 0.2.60__py3-none-any.whl → 1.0.18__py3-none-any.whl

fastapi_rtk/security/sqla/security_manager.py

@@ -3,6 +3,8 @@ import typing
 from datetime import datetime
 
 import fastapi_users.exceptions
+import sqlalchemy
+import sqlalchemy.orm
 from pydantic import BaseModel, Field
 from sqlalchemy import select
 from sqlalchemy.ext.asyncio import AsyncSession
@@ -11,7 +13,7 @@ from sqlalchemy.orm import Session
 from ...const import logger
 from ...globals import g
 from ...setting import Setting
-from ...utils import merge_schema, safe_call
+from ...utils import T, lazy, merge_schema, safe_call
 from .models import Api, Permission, PermissionApi, Role, User
 
 __all__ = ["SecurityManager"]
@@ -28,6 +30,21 @@ class SecurityManager:
     """
 
     toolkit: typing.Optional["FastAPIReactToolkit"]
+    builtin_roles = lazy(lambda: Setting.ROLES)
+    """
+    The built-in roles defined in the settings.
+    
+    Format:
+    ```python
+    {
+        "role_name": [
+            (api_name1|api_name2|..., permission_name1|permission_name2|...),
+            ...
+        ],
+        ...
+    }
+    ```
+    """
 
     def __init__(self, toolkit: typing.Optional["FastAPIReactToolkit"] = None) -> None:
         self.toolkit = toolkit
@@ -385,6 +402,273 @@ class SecurityManager:
     -----------------------------------------
     """
 
+    def has_access_in_builtin_roles(
+        self, role_name: str, api_name: str, permission_name: str
+    ):
+        """
+        Checks if the given role has access to the specified API and permission in the built-in roles.
+
+        Args:
+            role_name (str): The name of the role to check.
+            api_name (str): The name of the API to check.
+            permission_name (str): The name of the permission to check.
+
+        Returns:
+            bool: True if the role has access, False otherwise.
+        """
+        if role_name not in self.builtin_roles:
+            return False
+
+        for api, perm in self.builtin_roles[role_name]:
+            api_names = api.split("|")
+            perm_names = perm.split("|")
+            if api_name in api_names and permission_name in perm_names:
+                return True
+        return False
+
+    def get_roles_from_builtin_roles(self):
+        """
+        Retrieves the names of the built-in roles.
+
+        Returns:
+            list[str]: The list of built-in role names.
+        """
+        return list(self.builtin_roles.keys())
+
+    def get_api_permission_tuples_from_builtin_roles(self, role: str | None = None):
+        """
+        Retrieves the API-permission tuples from the built-in roles.
+
+        Args:
+            role (str | None, optional): The name of the role to filter by. If None, retrieves from all roles. Defaults to None.
+
+        Returns:
+            list[tuple[str, str]]: The list of API-permission tuples.
+        """
+        api_permission_tuples = list[tuple[str, str]]()
+        for role_name, role_api_permission_list in self.builtin_roles.items():
+            if role is not None and role != role_name:
+                continue
+
+            for api, perm in role_api_permission_list:
+                api_permission_tuples.append((api, perm))
+        return api_permission_tuples
+
+    def get_role_and_api_permission_tuples_from_builtin_roles(self):
+        """
+        Retrieves the role and API-permission tuples from the built-in roles.
+
+        Returns:
+            list[tuple[str, list[tuple[str, str]]]]: The list of role and API-permission tuples.
+        """
+        role_api_permission_tuples = list[tuple[str, list[tuple[str, str]]]]()
+        for role_name, role_api_permission_list in self.builtin_roles.items():
+            api_permission_list = list[tuple[str, str]]()
+            for api, perm in role_api_permission_list:
+                api_permission_list.append((api, perm))
+            role_api_permission_tuples.append((role_name, api_permission_list))
+        return role_api_permission_tuples
+
+    async def create_roles(
+        self,
+        roles: list[str],
+        *,
+        session: AsyncSession | Session = None,
+        raise_exception: bool = True,
+    ):
+        """
+        Creates new roles with the given names. Existing roles are not duplicated.
+
+        Args:
+            roles (list[str]): The names of the roles to create.
+            session (AsyncSession | Session, optional): The database session to use. If not given, a new session will be created. Defaults to None.
+            raise_exception (bool, optional): When set to True, raises an exception if an error occurs, otherwise returns None. Defaults to True.
+
+        Returns:
+            list[Role] | None: The created role objects if successful, else None.
+
+        Raises:
+            SomeException: Description of the exception raised, if any.
+        """
+        return await self._create_entities(
+            Role,
+            roles,
+            session=session,
+            raise_exception=raise_exception,
+            on_after_create=lambda role: logger.info(f"ADDING ROLE {role}"),
+        )
+
+    async def create_permissions(
+        self,
+        permissions: list[str],
+        *,
+        session: AsyncSession | Session = None,
+        raise_exception: bool = True,
+    ):
+        """
+        Creates new permissions with the given names. Existing permissions are not duplicated.
+
+        Args:
+            permissions (list[str]): The names of the permissions to create.
+            session (AsyncSession | Session, optional): The database session to use. If not given, a new session will be created. Defaults to None.
+            raise_exception (bool, optional): When set to True, raises an exception if an error occurs, otherwise returns None. Defaults to True.
+
+        Returns:
+            list[Permission] | None: The created permission objects if successful, else None.
+
+        Raises:
+            SomeException: Description of the exception raised, if any.
+        """
+        return await self._create_entities(
+            Permission,
+            permissions,
+            session=session,
+            raise_exception=raise_exception,
+            on_after_create=lambda permission: logger.info(
+                f"ADDING PERMISSION {permission}"
+            ),
+        )
+
+    async def create_apis(
+        self,
+        apis: list[str],
+        *,
+        session: AsyncSession | Session = None,
+        raise_exception: bool = True,
+    ):
+        """
+        Creates new APIs with the given names. Existing APIs are not duplicated.
+
+        Args:
+            apis (list[str]): The names of the APIs to create.
+            session (AsyncSession | Session, optional): The database session to use. If not given, a new session will be created. Defaults to None.
+            raise_exception (bool, optional): When set to True, raises an exception if an error occurs, otherwise returns None. Defaults to True.
+
+        Returns:
+            list[Api] | None: The created API objects if successful, else None.
+
+        Raises:
+            SomeException: Description of the exception raised, if any.
+        """
+        return await self._create_entities(
+            Api,
+            apis,
+            session=session,
+            raise_exception=raise_exception,
+            on_after_create=lambda api: logger.info(f"ADDING API {api}"),
+        )
+
+    async def associate_list_of_permission_with_api(
+        self,
+        permission_api_tuples: list[tuple[Permission, Api]],
+        *,
+        session: AsyncSession | Session = None,
+        raise_exception: bool = True,
+    ):
+        """
+        Associates a list of permissions with APIs. Existing associations are not duplicated.
+
+        Args:
+            permission_api_tuples (list[tuple[Permission, Api]]): A list of tuples containing Permission and Api objects to associate.
+            session (AsyncSession | Session, optional): The database session to use. If not given, a new session will be created. Defaults to None.
+            raise_exception (bool, optional): When set to True, raises an exception if an error occurs, otherwise returns None. Defaults to True.
+
+        Raises:
+            SomeException: Description of the exception raised, if any.
+        """
+        try:
+            if not session:
+                async with db.session() as session:
+                    await self.associate_list_of_permission_with_api(
+                        permission_api_tuples,
+                        session=session,
+                        raise_exception=raise_exception,
+                    )
+                    return
+
+            conditions = []
+            query = select(PermissionApi).options(
+                sqlalchemy.orm.joinedload(PermissionApi.api),
+                sqlalchemy.orm.joinedload(PermissionApi.permission),
+                sqlalchemy.orm.selectinload(PermissionApi.roles),
+            )
+            for permission, api in permission_api_tuples:
+                conditions.append(
+                    sqlalchemy.and_(
+                        PermissionApi.permission_id == permission.id,
+                        PermissionApi.api_id == api.id,
+                    )
+                )
+            query = query.where(sqlalchemy.or_(*conditions))
+            result = await safe_call(session.scalars(query))
+            existing_permission_apis = list(result.all())
+
+            new_tuples = list[tuple[Permission, Api]]()
+            for permission, api in permission_api_tuples:
+                exists = False
+                for permission_api in existing_permission_apis:
+                    if (
+                        permission_api.permission.id == permission.id
+                        and permission_api.api.id == api.id
+                    ):
+                        exists = True
+                        break
+                if not exists:
+                    new_tuples.append((permission, api))
+
+            new_permission_apis = list[PermissionApi]()
+            for permission, api in new_tuples:
+                permission_api = PermissionApi(permission=permission, api=api, roles=[])
+                session.add(permission_api)
+                new_permission_apis.append(permission_api)
+                logger.info(f"ASSOCIATING PERMISSION {permission} WITH API {api}")
+            await safe_call(session.commit())
+            return existing_permission_apis + new_permission_apis
+        except Exception as e:
+            if not raise_exception:
+                return
+            raise e
+
+    async def associate_list_of_role_with_permission_api(
+        self,
+        role_permission_api_tuples: list[tuple[Role, PermissionApi]],
+        *,
+        session: AsyncSession | Session = None,
+        raise_exception: bool = True,
+    ):
+        """
+        Associates a list of roles with permission APIs. Existing associations are not duplicated.
+
+        Args:
+            role_permission_api_tuples (list[tuple[Role, PermissionApi]]): A list of tuples containing Role and PermissionApi objects to associate.
+            session (AsyncSession | Session, optional): The database session to use. If not given, a new session will be created. Defaults to None.
+            raise_exception (bool, optional): When set to True, raises an exception if an error occurs, otherwise returns None. Defaults to True.
+
+        Raises:
+            SomeException: Description of the exception raised, if any.
+        """
+        try:
+            if not session:
+                async with db.session() as session:
+                    await self.associate_list_of_role_with_permission_api(
+                        role_permission_api_tuples,
+                        session=session,
+                        raise_exception=raise_exception,
+                    )
+                    return
+
+            for role, permission_api in role_permission_api_tuples:
+                if role not in permission_api.roles:
+                    permission_api.roles.append(role)
+                    logger.info(
+                        f"ASSOCIATING ROLE {role} WITH PERMISSION API {permission_api}"
+                    )
+            await safe_call(session.commit())
+        except Exception as e:
+            if not raise_exception:
+                return
+            raise e
+
     async def cleanup(self, *, session: AsyncSession | Session = None):
         """
         Cleanup unused permissions from apis and roles.
@@ -404,13 +688,12 @@ class SecurityManager:
                 "FastAPIReactToolkit instance not provided, you must provide it to use this function."
             )
 
-        api_permission_tuples = (Setting.ROLES).values()
+        api_permission_tuples = self.get_api_permission_tuples_from_builtin_roles()
         apis = [api.__class__.__name__ for api in self.toolkit.apis]
         permissions = self.toolkit.total_permissions()
-        for api_permission_tuple in api_permission_tuples:
-            for api, permission in api_permission_tuple:
-                apis.append(api)
-                permissions.append(permission)
+        for api, permission in api_permission_tuples:
+            apis.append(api)
+            permissions.append(permission)
 
         # Clean up unused permissions
         unused_permissions = await safe_call(
@@ -428,14 +711,20 @@ class SecurityManager:
             logger.info(f"DELETING API {api} AND ITS ASSOCIATIONS")
             await safe_call(session.delete(api))
 
-        roles = Setting.ROLES.keys()
+        roles = self.get_roles_from_builtin_roles()
         if g.admin_role is not None:
             roles.append(g.admin_role)
 
         # Clean up existing permission-apis, that are no longer connected to any roles
-        unused_permission_apis = await safe_call(session.scalars(select(PermissionApi)))
-        for permission_api in unused_permission_apis:
-            for role in list(permission_api.roles) or []:
+        permission_apis_in_db = await safe_call(
+            session.scalars(
+                select(PermissionApi).options(
+                    sqlalchemy.orm.selectinload(PermissionApi.roles)
+                )
+            )
+        )
+        for permission_api in permission_apis_in_db:
+            for role in permission_api.roles:
                 if role.name not in roles:
                     permission_api.roles.remove(role)
                     logger.info(
@@ -443,3 +732,71 @@ class SecurityManager:
                     )
 
         await safe_call(session.commit())
+
+    """
+    -----------------------------------------
+         HELPER FUNCTIONS
+    -----------------------------------------
+    """
+
+    async def _create_entities(
+        self,
+        entity_class: typing.Type[T],
+        names: list[str],
+        *,
+        session: AsyncSession | Session = None,
+        raise_exception: bool = True,
+        on_after_create: typing.Optional[
+            typing.Callable[[T], typing.Awaitable[None] | None]
+        ] = None,
+    ):
+        """
+        Helper function to create new entities with the given names.
+        Existing entities are not duplicated.
+
+        Args:
+            entity_class (typing.Type[T]): The entity class to create.
+            names (list[str]): The names of the entities to create.
+            session (AsyncSession | Session, optional): The database session to use. If not given, a new session will be created. Defaults to None.
+            raise_exception (bool, optional): When set to True, raises an exception if an error occurs, otherwise returns None. Defaults to True.
+            on_after_create (Optional[Callable[[T], Awaitable[None] | None]], optional): An optional callback function to be called after each entity is created. Defaults to None.
+
+        Returns:
+            list[T] | None: The created entity objects if successful, else None.
+
+        Raises:
+            Exception: If an error occurs and raise_exception is True.
+        """
+        try:
+            if not session:
+                async with db.session() as session:
+                    return await self._create_entities(
+                        entity_class,
+                        names,
+                        session=session,
+                        raise_exception=raise_exception,
+                    )
+
+            # Check for existing entities
+            result = await safe_call(
+                session.scalars(
+                    select(entity_class).where(entity_class.name.in_(names))
+                )
+            )
+            existing_entities = list(result.all())
+            existing_names = [entity.name for entity in existing_entities]
+
+            new_names = [name for name in names if name not in existing_names]
+            new_entities = list[T]()
+            for name in new_names:
+                entity = entity_class(name=name)
+                session.add(entity)
+                new_entities.append(entity)
+                if on_after_create:
+                    await safe_call(on_after_create(entity))
+            await safe_call(session.commit())
+            return existing_entities + new_entities
+        except Exception as e:
+            if not raise_exception:
+                return
+            raise e

fastapi_rtk/utils/async_task_runner.py

@@ -2,12 +2,15 @@ import asyncio
 import contextvars
 import functools
 import inspect
+import traceback as tb
 import typing
 
 from .prettify_dict import prettify_dict
 
 __all__ = ["AsyncTaskRunner"]
 
+T = typing.TypeVar("T")
+
 
 class CallerInfo(typing.TypedDict):
     """
@@ -35,26 +38,37 @@ class AsyncTaskException(AsyncTaskRunnerException):
         self.caller = caller
 
 
-def wrap_in_async_task_exception(
-    task: typing.Callable[[], typing.Coroutine | None] | typing.Coroutine,
-    /,
-    caller: CallerInfo | None = None,
-):
-    async def wrapper():
-        try:
-            tsk = task
-            if callable(tsk):
-                tsk = tsk()
-            if tsk is None:
-                return None
-            return await tsk
-        except Exception as e:
-            raise AsyncTaskException(str(e), original_exception=e, caller=caller) from e
+class ResultAccessor(typing.Generic[T]):
+    """Helper class that can be awaited OR accessed directly."""
+
+    def __init__(self, task: "AsyncTask[T]"):
+        self._task = task
+
+    def __await__(self):
+        """Allow awaiting: await task.result"""
+        return self._task.__await__()
+
+    def __repr__(self) -> str:
+        """Direct access without await - returns value or raises error."""
+        return repr(self._ensure_result())
 
-    return wrapper
+    def __str__(self) -> str:
+        """Direct access without await - returns value or raises error."""
+        return str(self._ensure_result())
 
+    # Make it behave like the actual value when accessed
+    def __getattr__(self, name):
+        return getattr(self._ensure_result(), name)
 
-class AsyncTask:
+    def _ensure_result(self):
+        if not hasattr(self._task, "_result"):
+            raise RuntimeError(
+                "Task has not been executed yet. Await the task to get the result."
+            )
+        return self._task._result
+
+
+class AsyncTask(typing.Generic[T]):
     """
     Represents a task to be run asynchronously.
 
@@ -63,19 +77,51 @@ class AsyncTask:
 
     def __init__(
         self,
-        task: typing.Callable[[], typing.Coroutine | None] | typing.Coroutine,
+        task: typing.Callable[..., typing.Awaitable[T]] | typing.Awaitable[T],
         /,
         caller: CallerInfo | None = None,
         tags: list[str] | None = None,
     ):
-        self.task = wrap_in_async_task_exception(task, caller=caller)
+        self.task = task
         self.caller = caller
         self.tags = tags or []
 
+    @property
+    def result(self):
+        """
+        Provides access to the result of the task.
+
+        Returns:
+            ResultAccessor[T]: An accessor that can be awaited or accessed directly.
+        """
+        return ResultAccessor(self)
+
     def __call__(self):
-        if callable(self.task):
-            return self.task()
-        return self.task
+        return self._run()
+
+    def __await__(self):
+        return self._run().__await__()
+
+    async def _run(self):
+        """
+        Runs the task and caches the result.
+
+        Returns:
+            T: The result of the task.
+        """
+        if hasattr(self, "_result"):
+            return self._result
+
+        try:
+            coro = self.task
+            if callable(coro):
+                coro = coro()
+            self._result = await coro
+            return self._result
+        except Exception as e:
+            raise AsyncTaskException(
+                str(e), original_exception=e, caller=self.caller
+            ) from e
 
 
 class AsyncTaskRunner:
@@ -143,19 +189,59 @@ class AsyncTaskRunner:
         """
         self.run_tasks_even_if_exception = run_tasks_even_if_exception
 
+    @typing.overload
     @classmethod
     def add_task(
         cls,
-        *tasks: typing.Callable[[], typing.Coroutine | None] | typing.Coroutine,
+        task: typing.Callable[..., typing.Awaitable[T]]
+        | typing.Awaitable[T]
+        | AsyncTask[T],
+        *,
+        tags: list[str] | None = None,
+        instance: "AsyncTaskRunner | None" = None,
+    ) -> AsyncTask[T]: ...
+
+    @typing.overload
+    @classmethod
+    def add_task(
+        cls,
+        task1: typing.Callable[..., typing.Awaitable[T]]
+        | typing.Awaitable[T]
+        | AsyncTask[T],
+        task2: typing.Callable[..., typing.Awaitable[T]]
+        | typing.Awaitable[T]
+        | AsyncTask[T],
+        *tasks: typing.Callable[..., typing.Awaitable[T]]
+        | typing.Awaitable[T]
+        | AsyncTask[T],
+        tags: list[str] | None = None,
+        instance: "AsyncTaskRunner | None" = None,
+    ) -> list[AsyncTask[T]]: ...
+    @classmethod
+    def add_task(
+        cls,
+        *tasks: typing.Callable[..., typing.Awaitable[T]]
+        | typing.Awaitable[T]
+        | AsyncTask[T],
         tags: list[str] | None = None,
         instance: "AsyncTaskRunner | None" = None,
     ):
+        """
+        Adds one or more tasks to the current context's task list. The tasks will be executed when exiting the context.
+
+        Args:
+            tags (list[str] | None, optional): Tags to associate with the tasks. Defaults to None.
+            instance (AsyncTaskRunner | None, optional): The AsyncTaskRunner instance to add tasks to. Defaults to None.
+
+        Returns:
+            AsyncTask[T] | list[AsyncTask[T]]: The added task(s).
+        """
         task_list = cls._get_current_task_list("add_task", instance=instance)
         # Get caller info
         frame = inspect.currentframe()
         caller = inspect.getouterframes(frame, 2)[1] if frame else None
 
-        async_tasks = [
+        async_tasks: list[AsyncTask[T]] = [
             AsyncTask(
                 task,
                 caller=CallerInfo(
@@ -170,6 +256,7 @@ class AsyncTaskRunner:
             for task in tasks
         ]
         task_list.extend(async_tasks)
+        return async_tasks if len(async_tasks) > 1 else async_tasks[0]
 
     @classmethod
     def remove_tasks_by_tag(
@@ -214,16 +301,16 @@ class AsyncTaskRunner:
             return
 
         if tasks:
-            exceptions: list[AsyncTaskException] = []
+            exceptions_with_index = list[tuple[AsyncTaskException, int]]()
             futures = await asyncio.gather(
                 *(task() for task in tasks), return_exceptions=True
             )
-            for future in futures:
+            for index, future in enumerate(futures):
                 if isinstance(future, AsyncTaskException):
                     # Handle exceptions from tasks
-                    exceptions.append(future)
+                    exceptions_with_index.append((future, index))
 
-            if exceptions:
+            if exceptions_with_index:
                 raise AsyncTaskRunnerException(
                     f"\n{
                         prettify_dict(
@@ -233,9 +320,11 @@ class AsyncTaskRunner:
                                     f'Task {index + 1}': {
                                         'message': str(exc),
                                         'caller': exc.caller,
-                                        'traceback': exc.original_exception.__traceback__,
+                                        'traceback': ''.join(
+                                            tb.format_exception(exc.original_exception)
+                                        ),
                                     }
-                                    for index, exc in enumerate(exceptions)
+                                    for exc, index in exceptions_with_index
                                 },
                             }
                         )