munch-ease-backend/product/db.py

from typing import List, ClassVar
from pydantic import BaseModel

class Product(BaseModel):
    KEYS: ClassVar[List[str]] = ['id', 'product_id', 'link', 'name', 'img_small', 'img_large', 'raw_data']
    id: int
    product_id: str
    link: str
    name: str
    img_small: str
    img_large: str
    raw_data: str

async def create(conn):
    await conn.execute('''
        CREATE TABLE IF NOT EXISTS Product (
            id INTEGER PRIMARY KEY,
            product_id TEXT UNIQUE,
            link TEXT,
            name TEXT,
            img_small TEXT,
            img_large TEXT,
            raw_data TEXT
        );''')
    
    await conn.execute('''       
        CREATE TABLE IF NOT EXISTS ProductTag (
            food_item_id INTEGER,
            tag TEXT,
            PRIMARY KEY (food_item_id, tag),
            FOREIGN KEY (food_item_id) REFERENCES Product(id)
        );''')

async def find_product_by_tag(conn, tag: str) -> List[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: str) -> 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)})

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

async def insert_product(conn, product: Product):
    async with conn.execute('''
        INSERT INTO Product (name, product_id, link, img_small, img_large, raw_data)
        VALUES (?, ?, ?, ?, ?, ?)
    ''', (product.name, product.product_id, product.link, product.img_small, product.img_large, product.raw_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) -> List[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]