Untitled

from collections import defaultdict
from typing import Any, List

from sqlalchemy import Column, or_, select
from sqlalchemy.orm import Session

from app.helpers.custom_logger import BackendLogger


class BaseLoader:
    session: Session
    model: any
    column: Column

    def __init__(self, session: Session):
        self.session = session
        self.data = defaultdict(lambda: None)

    async def __call__(self, keys: List[any]) -> List[Any]:
        klass = self.__class__
        BackendLogger.debug(
            f"BaseLoader {klass.__name__}({klass.model.__name__}) {keys=}"
        )

        where = []
        for i in range(0, len(keys), 1000):
            chunk = keys[i : i + 1000]
            where.append(klass.column.in_(chunk))

        stmt = select(klass.model).where(or_(*where))
        for item in self.session.scalars(stmt):
            self.data[getattr(item, klass.column.key)] = item

        return [self.data[key] for key in keys]