munch-ease-backend/products/db.py

import json
from typing import AsyncIterator, ClassVar, List, Optional

from pydantic import BaseModel


class Product(BaseModel):
    KEYS: ClassVar[List[str]] = [
        "id",
        "product_id",
        "shop_code",
        "link",
        "name",
        "quantity",
        "unit",
        "img_small",
        "img_large",
    ]
    NON_INSERT_KEYS: ClassVar[List[str]] = ["id"]

    id: int = -1
    product_id: str
    shop_code: str
    link: str
    name: str
    quantity: int
    unit: str
    img_small: str
    img_large: str
    # Non-persisted field used in tests and insert helper
    raw_data: Optional[dict] = None


async def create(conn):
    await conn.execute(
        """
        CREATE TABLE IF NOT EXISTS Product (
            id INTEGER PRIMARY KEY,
            product_id TEXT UNIQUE NOT NULL,
            shop_code TEXT NOT NULL,
            link TEXT NOT NULL,
            name TEXT NOT NULL,
            quantity INTEGER NOT NULL,
            unit TEXT NOT NULL,
            img_small TEXT,
            img_large TEXT,
            raw_data TEXT
        );"""
    )

    await conn.execute(
        """       
        CREATE TABLE IF NOT EXISTS ProductTag (
            food_item_id INTEGER,
            tag TEXT COLLATE NOCASE,
            PRIMARY KEY (food_item_id, tag),
            FOREIGN KEY (food_item_id) REFERENCES Product(id)
        );"""
    )


async def find_product_by_tag(conn, tag: str) -> AsyncIterator[Product]:
    async with conn.execute(
        f"""
        SELECT {','.join(Product.KEYS)} FROM Product
        WHERE id IN (
            SELECT food_item_id FROM ProductTag
            WHERE tag = ?
        )
    """,
        (tag,),
    ) as cursor:
        async for row in cursor:
            yield Product(**{k: v for k, v in zip(Product.KEYS, row)})


async def find_product_by_id(conn, product_id: int) -> Optional[Product]:
    async with conn.execute(
        f"""
        SELECT {','.join(Product.KEYS)} FROM Product
        WHERE id = ?
        LIMIT 1
    """,
        (product_id,),
    ) as cursor:
        async for row in cursor:
            return Product(**{k: v for k, v in zip(Product.KEYS, row)})
    return None


async def find_product_by_key(conn, shop_code: str, product_id: str) -> Optional[Product]:
    async with conn.execute(
        f"""
        SELECT {','.join(Product.KEYS)} FROM Product
        WHERE shop_code = ? AND product_id = ?
        LIMIT 1
    """,
        (
            shop_code,
            product_id,
        ),
    ) as cursor:
        async for row in cursor:
            return Product(**{k: v for k, v in zip(Product.KEYS, row)})
    return None


async def insert_product(conn, product: Product, data: dict):
    insert_keys = [k for k in Product.KEYS if k not in Product.NON_INSERT_KEYS]
    insert_values = [getattr(product, k) for k in insert_keys]

    async with conn.execute(
        f"""
        INSERT INTO Product ({','.join(insert_keys)}, raw_data)
        VALUES ({','.join(['?'] * len(insert_keys))}, ?)
    """,
        (*insert_values, json.dumps(data)),
    ) as cursor:
        product.id = cursor.lastrowid

    await conn.commit()


async def add_tag(conn, product: Product, tag: str):
    await conn.execute(
        """
        INSERT INTO ProductTag (food_item_id, tag)
        VALUES (?, ?)
    """,
        (product.id, tag),
    )

    await conn.commit()


async def get_tags(conn, product: Product) -> AsyncIterator[str]:
    async with conn.execute(
        """
        SELECT tag FROM ProductTag
        WHERE food_item_id = ?
    """,
        (product.id,),
    ) as cursor:
        async for row in cursor:
            yield row[0]