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]