sakhteman/core/utils/router.py

import time
from functools import wraps
from typing import Callable, TypeVar, Generic

from fastapi import APIRouter, Request, Response, Body
from fastapi import Security
from fastapi.routing import APIRoute
from starlette.background import BackgroundTasks

from auth.utils import get_current_active_user
from config import conf
from core.utils import to_http_res, http_exception
from core.utils.auth import get_current_active_user


class BaseRoute(APIRoute):
    def get_route_handler(self) -> Callable:
        original_route_handler = super().get_route_handler()

        async def custom_route_handler(request: Request) -> Response:
            before = time.time()
            response: Response = await original_route_handler(request)
            duration = time.time() - before
            response.headers["X-Response-Time"] = str(duration)
            return response

        return custom_route_handler


def shrink(cal: Callable, name: str = None, bt: BackgroundTasks = None):
    def can_inner(func):
        @wraps(func)
        async def wrapper(*args, **kwargs):
            return await safe_route(name or cal.__name__, cal, *kwargs.values())
        return wrapper

    return can_inner


def back_task(cal: Callable, name: str = None):
    print(cal)

    def back_inner(func):
        print(func)

        @wraps(func)
        async def wrapper(*args, **kwargs):
            if cal.__name__ == 'create_salam':
                print(cal.__name__)
            return await safe_route(name or cal.__name__, cal, *kwargs.values())

        return wrapper

    return back_inner


def generate_router(
        prefix,
        dependencies=None,
        auth_dependencies=False,
        version='v1',
        auth_scopes=['account']):
    """
    Generate a router with a prefix and dependencies.
    Args:
        prefix: prefix of the router.
        dependencies: dependencies of the router.
        auth_dependencies:  if router has auth dependencies is True.
        version: version of the router.
        auth_scopes: auth scopes of the router.

    Returns: ApiRouter
    """
    if auth_dependencies:
        dependencies = dependencies or []
        if 'customer' in auth_scopes:
            dependencies.extend(
                [Security(get_current_active_customer, scopes=auth_scopes)])
        else:
            dependencies.extend(
                [Security(get_current_active_user, scopes=auth_scopes)])
    _pre = conf.server.prefix_url.replace('v1', version)
    return APIRouter(
        prefix=f"{_pre}/{prefix}",
        route_class=BaseRoute,
        dependencies=dependencies if dependencies else []
    )


async def safe_route(_type, fn, *args):
    try:
        return to_http_res(await fn(*args), _type)
    except Exception as e:
        return http_exception(e)


class CrudMethods:
    create = True
    read = True
    update = True
    delete = True
    get_all = True

    def __init__(self, create=True, read=True, update=True, delete=True, fetch=True):
        self.create = create
        self.read = read
        self.update = update
        self.delete = delete
        self.fetch = fetch

    def __iter__(self):
        r = []
        if self.create:
            r.append('create')
        if self.read:
            r.append('read')
        if self.update:
            r.append('update')
        if self.delete:
            r.append('delete')
        if self.fetch:
            r.append('fetch')
        return r


T = TypeVar("T")


class SetupCrud(Generic[T]):
    _url_prefix = {
        'create': '/create',
        'get': '/{_id}',
        'update': '/{_id}',
        'delete': '/{_id}',
        'fetch': ''
    }
    _methods = {
        'create': 'post',
        'fetch': 'post',
        'get': 'get',
        'update': 'put',
        'delete': 'delete',
    }

    def __init__(self,
                 router: APIRouter,
                 repository,
                 model: T,
                 tags=[],
                 methods: CrudMethods = CrudMethods()
                 ):

        for i in ['create', 'update', 'get', 'delete', 'fetch']:
            if i == 'create':
                @router.post(
                    self._url_prefix[i],
                    response_model=model,
                    tags=tags
                )
                async def create(body=Body(...)):
                    b = T(**body)
                    return await safe_route(i, getattr(repository, 'create'), b)
            # elif i == 'update':
            #     async def update(_id, body: T = Body(...)):
            #         return await safe_route(i, getattr(repository, i), _id, body)
            #
            #     router.add_api_route(
            #         _url_prefix[i],
            #         update,
            #         methods=[_methods[i]],
            #         response_model=model,
            #         tags=tags
            #     )
            # elif i == 'delete':
            #     async def delete(_id):
            #         return await safe_route(i, getattr(repository, i), _id)
            #
            #     router.add_api_route(
            #         _url_prefix[i],
            #         delete,
            #         methods=[_methods[i]],
            #         response_model=model,
            #         tags=tags
            #     )
            # elif i == 'get':
            #     async def get(_id):
            #         return await safe_route(i, getattr(repository, i), _id)
            #
            #     router.add_api_route(
            #         _url_prefix[i],
            #         get,
            #         methods=[_methods[i]],
            #         response_model=model,
            #         tags=tags
            #     )
            # elif i == 'fetch':
            #     async def fetch(_id, body: T = Body(...)):
            #         return await safe_route(i, getattr(repository, i), _id, body)
            #
            #     router.add_api_route(
            #         _url_prefix[i],
            #         fetch,
            #         methods=[_methods[i]],
            #         response_model=model,
            #         tags=tags
            #     )