munch-ease-backend/persons/db.py

from typing import AsyncIterator, ClassVar, List, Optional

from common import ApiModel


class Person(ApiModel):
    KEYS: ClassVar[List[str]] = ["id", "name"]

    id: int = -1
    name: str


async def create(conn):
    await conn.execute(
        """
        CREATE TABLE IF NOT EXISTS Person (
            id INTEGER PRIMARY KEY,
            name TEXT UNIQUE
        );"""
    )
    # Useful indexes for search and pagination
    await conn.execute("CREATE INDEX IF NOT EXISTS idx_person_name ON Person(name);")


async def search_by_name(conn, name: str) -> AsyncIterator[Person]:
    async with conn.execute(
        """
        SELECT id, name
        FROM Person
        WHERE name LIKE ?
    """,
        (f"%{name}%",),
    ) as cursor:
        async for row in cursor:
            yield Person(id=row[0], name=row[1])


async def get_by_name(conn, name: str) -> Optional[Person]:
    cursor = await conn.execute(
        """
        SELECT id, name
        FROM Person
        WHERE name = ?
    """,
        (name,),
    )
    row = await cursor.fetchone()
    if not row:
        return None
    return Person(id=row[0], name=row[1])


async def get_by_id(conn, id: int) -> Optional[Person]:
    cursor = await conn.execute(
        """
        SELECT id, name
        FROM Person
        WHERE id = ?
    """,
        (id,),
    )
    row = await cursor.fetchone()
    if not row:
        return None
    return Person(id=row[0], name=row[1])


async def get_by_ids(conn, ids: List[int]) -> dict[int, Person]:
    """Fetch many persons in a single query. Returns a dict id->Person.

    If ids is empty, returns {}.
    """
    if not ids:
        return {}
    placeholders = ",".join(["?"] * len(ids))
    query = f"""
        SELECT id, name
        FROM Person
        WHERE id IN ({placeholders})
    """
    result: dict[int, Person] = {}
    async with conn.execute(query, ids) as cursor:
        async for row in cursor:
            p = Person(id=row[0], name=row[1])
            result[p.id] = p
    return result


async def get_all(conn) -> AsyncIterator[Person]:
    async with conn.execute(
        """
        SELECT id, name
        FROM Person
    """
    ) as cursor:
        async for row in cursor:
            yield Person(id=row[0], name=row[1])


async def get_all_paged(conn, after_id: Optional[int], limit: int) -> AsyncIterator[Person]:
    after = after_id if after_id is not None else -1
    async with conn.execute(
        """
        SELECT id, name
        FROM Person
        WHERE id > ?
        ORDER BY id
        LIMIT ?
        """,
        (after, limit),
    ) as cursor:
        async for row in cursor:
            yield Person(id=row[0], name=row[1])


async def search_by_name_paged(
    conn, name: str, after_id: Optional[int], limit: int
) -> AsyncIterator[Person]:
    after = after_id if after_id is not None else -1
    async with conn.execute(
        """
        SELECT id, name
        FROM Person
        WHERE name LIKE ? AND id > ?
        ORDER BY id
        LIMIT ?
        """,
        (f"%{name}%", after, limit),
    ) as cursor:
        async for row in cursor:
            yield Person(id=row[0], name=row[1])


async def count_all(conn) -> int:
    cursor = await conn.execute(
        """
        SELECT COUNT(1)
        FROM Person
        """
    )
    row = await cursor.fetchone()
    return int(row[0]) if row else 0


async def count_by_name(conn, name: str) -> int:
    cursor = await conn.execute(
        """
        SELECT COUNT(1)
        FROM Person
        WHERE name LIKE ?
        """,
        (f"%{name}%",),
    )
    row = await cursor.fetchone()
    return int(row[0]) if row else 0


async def compute_prev_cursor(conn, first_id: int, limit: int, name: Optional[str] = None) -> Optional[str]:
    """Compute a prevCursor string for paginated persons, respecting optional name LIKE filter."""
    if limit <= 0:
        return None
    if name:
        query = """
            SELECT id
            FROM Person
            WHERE name LIKE ? AND id < ?
            ORDER BY id DESC
            LIMIT ?
        """
        params = (f"%{name}%", first_id, limit)
    else:
        query = """
            SELECT id
            FROM Person
            WHERE id < ?
            ORDER BY id DESC
            LIMIT ?
        """
        params = (first_id, limit)

    async with conn.execute(query, params) as c:
        prev_ids = [row[0] async for row in c]
    if len(prev_ids) == limit and prev_ids:
        return str(min(prev_ids) - 1)
    return None


async def insert_person(conn, person: Person) -> Person:
    cursor = await conn.execute(
        """
        INSERT INTO Person (name)
        VALUES (?)
    """,
        (person.name,),
    )
    person.id = cursor.lastrowid
    return person