Decorator to parse the results of a raw sqlalchemy query to a Pydantic model.

parse_sql_return.py

import inspect
from functools import partial, wraps
from typing import Union

from app import schemas
from app.db import database
from pydantic.main import ModelMetaclass
from shortuuid import uuid

# The following 2 functions parse raw results from the SQL queries
# and convert them to Pydantic models.

# The current implementation only allows for not nested models.
# That means that a model cannot has a field that references another model.

# I still have to read FastAPI's code better to understand how it's done.
# This?: https://github.com/tiangolo/fastapi/blob/3223de5598310359262ff3e3a0259a00851cbc58/fastapi/utils.py#L73
def parse_result(res, model, names):
    """Parse list to Pydantic models."""

    return model.parse_obj(dict(zip(names, res)))


def parse_return(func=None, return_list=True):
    """Decorator to parse SQL results into Pydantic models."""

    if func is None:
        return partial(parse_return, return_list=return_list)

    model = inspect.signature(func).return_annotation

    if not isinstance(model, ModelMetaclass):
        raise TypeError("Return type must be a Pydantic model")

    names = model.schema()["properties"].keys()

    parser = partial(parse_result, model=model, names=names)

    @wraps(func)
    async def decorated(*args, **kwargs):

        results = await func(*args, **kwargs)

        if not results:
            return None

        # the results.items() method is available because the RowProxy
        # provided by sqlachemy has it.
        # If not using sqlalchemy you need to parse the returned tuples
        # IMPORTANT: I think this only works when using .fetchone()
        results = dict(results.items())

        # if using fetch_one
        if isinstance(results, dict):
            return model.parse_obj(results)

        # this part is for the times when you fetch many results,
        # not .fetchone() or .fetchval()
        # my use case still does not need this, it works but you need
        # to remove the line:
        # `results = dict(results.items())`
        # I'll update the gist when it gets more polished.
        
        if return_list:
            return list(map(parser, results))

        return map(parser, func(*args, **kwargs))

    return decorated


# example
@parse_return
async def user_get(user_id: int) -> schemas.User:  # schemas.User is a Pydantic schema.

    query = """
SELECT * FROM users WHERE user_id = :user_id
    """

    values = {"user_id": user_id}

    user = await database.fetch_one(query=query, values=values)

    return user