munch-ease-backend/shopping/db.py

from meals import Meal, find_meal_by_id
from ingredients import Ingredient, insert_ingredient
from persons import Person
from products import Product

from pydantic import BaseModel
from typing import AsyncIterator, List, ClassVar, Optional

from datetime import datetime

class ShoppingListRequest(BaseModel):
    KEYS: ClassVar[List[str]] = ['id', 'ingredient_id', 'list_id', 'person_id', 'meal_id', 'created_date']
    id: int = -1
    list_id: Optional[int] = None
    
    ingredient_id: Optional[int] = None
    ingredient: Optional[Ingredient] = None

    person_id: Optional[int] = None
    person: Optional[Person] = None

    meal_id: Optional[int] = None
    meal: Optional[Meal] = None

    created_date: datetime = datetime.now().astimezone()

class ShoppingListResult(BaseModel):
    KEYS: ClassVar[List[str]] = ['id', 'product_id', 'list_id', 'quantity', 'unit' ]
    id: int = -1
    list_id: int
    product_id: int
    product: Optional[Product] = None

    quantity: float
    unit: str

from enum import Enum

class StoreEnum(str, Enum):
    woolworths = 'woolworths'
    coles = 'coles'
    home = ''

class ShoppingList(BaseModel):
    KEYS: ClassVar[List[str]] = ['id', 'created_date', 'store_name']
    id: int = -1
    created_date: datetime = datetime.now().astimezone()
    store_name: StoreEnum = ''

    requests: List[ShoppingListRequest] = []
    results: List[ShoppingListResult] = []

async def create(conn):
    await conn.execute('''
        CREATE TABLE IF NOT EXISTS ShoppingList (
            id INTEGER PRIMARY KEY,
            created_date DATETIME,
            store_name TEXT
        );''')

    await conn.execute('''
        CREATE TABLE IF NOT EXISTS ShoppingListRequest (
            id INTEGER PRIMARY KEY,
            ingredient_id INTEGER,
            list_id INTEGER,
            person_id INTEGER,
            meal_id INTEGER,
            created_date TEXT,
            FOREIGN KEY(ingredient_id) REFERENCES Ingredient(id),
            FOREIGN KEY(person_id) REFERENCES Person(id),
            FOREIGN KEY(meal_id) REFERENCES Meal(id),
            FOREIGN KEY(list_id) REFERENCES ShoppingList(id)
        );''')

    await conn.execute('''
        CREATE TABLE IF NOT EXISTS ShoppingListResult (
            id INTEGER PRIMARY KEY,
            product_id INTEGER,
            list_id INTEGER,
            quantity REAL,
            unit TEXT,
            FOREIGN KEY(product_id) REFERENCES Product(id),
            FOREIGN KEY(list_id) REFERENCES ShoppingList(id)
        );''')

def validate_request(request: ShoppingListRequest) -> None:
    # A request must always have a list id
    if request.list_id < 0:
        raise ValueError('Request must have a list id')
    
    # A request must have either an ingredient or a meal, but not both
    if not request.ingredient and not request.meal:
        raise ValueError('Request must have either an ingredient or a meal')
    
    if request.ingredient and request.meal:
        raise ValueError('Request cannot have both an ingredient and a meal')
    
    # If an ingredient is provided, it must have a person
    if request.ingredient and not request.person:
        raise ValueError('Ingredient requests must have a person')

async def insert_shopping_list(conn, shopping_list: ShoppingList):
    shopping_list.created_date = datetime.now().astimezone()

    async with conn.execute('''
        INSERT INTO ShoppingList (created_date, store_name)
        VALUES (?, ?)
    ''', (shopping_list.created_date.isoformat(), shopping_list.store_name,)) as cursor:
        shopping_list.id = cursor.lastrowid

    for request in shopping_list.requests:
        request.list_id = shopping_list.id

        validate_request(request)

        if request.ingredient and request.ingredient.id < 0:
            await insert_ingredient(conn, request.ingredient)

        if request.ingredient:
            request.ingredient_id = request.ingredient.id

        if request.meal:
            request.meal_id = request.meal.id

        async with conn.execute('''
            INSERT INTO ShoppingListRequest (ingredient_id, list_id, person_id, meal_id, created_date)
            VALUES (?, ?, ?, ?, ?)
        ''', (request.ingredient_id, shopping_list.id, request.person_id, request.meal_id, request.created_date.isoformat())) as cursor:
            request.id = cursor.lastrowid

    for item in shopping_list.results:
        item.product_id = item.product.id
        item.list_id = shopping_list.id

        async with conn.execute('''
            INSERT INTO ShoppingListResult (product_id, list_id, quantity, unit)
            VALUES (?, ?, ?, ?)
        ''', (item.product_id, item.list_id, item.quantity, item.unit)) as cursor:
            item.id = cursor.lastrowid

async def remove_request(conn, request: ShoppingListRequest) -> None:
    if request.list_id != None:
        raise ValueError('Request is already completed')
    
    if request.meal and not request.meal_id:
        raise ValueError('Meal request must have a meal id')
    
    if request.meal_id != None:
        await conn.execute('''
            DELETE FROM ShoppingListRequest
            WHERE meal_id = ? AND list_id IS NULL
        ''', (request.meal_id,))

    elif request.person_id != None and request.ingredient_id != None:
        await conn.execute('''
            DELETE FROM ShoppingListRequest
            WHERE person_id = ? AND ingredient_id = ? AND list_id IS NULL
        ''', (request.person_id, request.ingredient_id))

    else:
        raise ValueError('Request is invalid')

async def find_requests_by_list_id(conn, list_id: Optional[int]) -> AsyncIterator[ShoppingListRequest]:
    # Join Ingredient and Product to also load ingredient and product
    ingredient_keys = [f'ingredient.{key}' for key in Ingredient.KEYS]
    product_keys = [f'product.{key}' for key in Product.KEYS]
    request_keys = [f'shoppinglistrequest.{key}' for key in ShoppingListRequest.KEYS]
    person_keys = [f'person.{key}' for key in Person.KEYS]

    select = f'''
        SELECT {','.join(ingredient_keys + product_keys + request_keys + person_keys)}
        FROM ShoppingListRequest
        LEFT JOIN Ingredient ON ShoppingListRequest.ingredient_id = Ingredient.id
        LEFT JOIN Product ON Ingredient.product_id = Product.id
        LEFT JOIN Person ON ShoppingListRequest.person_id = Person.id
    '''
    
    where, params = ' WHERE list_id IS NULL', ()
    if list_id is not None:
        where, params = ' WHERE list_id = ?', (list_id,)
    
    cursor = await conn.execute(select + where, params)
    
    async for row in cursor:
        product_keys = {k:v for k,v in zip(Product.KEYS, row[len(Ingredient.KEYS):len(Ingredient.KEYS) + len(Product.KEYS)])}
        product = Product(**product_keys) if product_keys['id'] else None

        ingredient_keys = {k:v for k,v in zip(Ingredient.KEYS, row[:len(Ingredient.KEYS)])}
        ingredient = Ingredient(**ingredient_keys, product=product) if ingredient_keys['id'] else None

        person_keys = {k:v for k,v in zip(Person.KEYS, row[-len(Person.KEYS):])}
        person = Person(**person_keys) if person_keys['id'] else None

        request_keys = {k:v for k,v in zip(ShoppingListRequest.KEYS, row[len(Ingredient.KEYS) + len(Product.KEYS):-len(Person.KEYS)])}
        request = ShoppingListRequest(**request_keys, ingredient=ingredient, person=person)

        if request.meal_id is not None:
            request.meal = await find_meal_by_id(conn, request.meal_id)

        yield request

async def find_results_by_list_id(conn, list_id: int) -> AsyncIterator[ShoppingListResult]:
    product_keys = [f'product.{key}' for key in Product.KEYS]
    result_keys = [f'shoppinglistresult.{key}' for key in ShoppingListResult.KEYS]

    async with conn.execute(f'''
        SELECT {','.join(product_keys + result_keys)}
        FROM ShoppingListResult
        LEFT JOIN Product ON ShoppingListResult.product_id = Product.id
        WHERE list_id = ?
    ''', (list_id,)) as cursor:
        async for row in cursor:
            product_keys = {k:v for k,v in zip(Product.KEYS, row[:len(Product.KEYS)])}
            product = Product(**product_keys) if product_keys['id'] else None

            result_keys = {k:v for k,v in zip(ShoppingListResult.KEYS, row[len(Product.KEYS):])}
            result = ShoppingListResult(**result_keys, product=product)
            yield result

async def fill_related(conn, shopping_list: ShoppingList) -> ShoppingList:
    async for request in find_requests_by_list_id(conn, shopping_list.id):
        shopping_list.requests.append(request)

    async for item in find_results_by_list_id(conn, shopping_list.id):
        shopping_list.results.append(item)

async def load_shopping_list(conn, id: int) -> ShoppingList:
    shopping_list = None
    async with conn.execute(f'''
        SELECT {','.join(ShoppingList.KEYS)} FROM ShoppingList
        WHERE id = ?
        LIMIT 1
    ''', (id,)) as cursor:
        async for row in cursor:
            shopping_list = ShoppingList(**{k:v for k,v in zip(ShoppingList.KEYS, row)})
            break

    if shopping_list:
        await fill_related(conn, shopping_list)

    return shopping_list

async def request_ingredient(conn, person: Person, ingredient: Ingredient) -> ShoppingListRequest:
    if ingredient.id >= 0:
        raise ValueError('How did you get an existing ingredient?')
    
    await insert_ingredient(conn, ingredient)

    request = ShoppingListRequest(ingredient_id=ingredient.id, ingredient=ingredient, person_id=person.id, created_date=datetime.now().astimezone())
    async with conn.execute('''
        INSERT INTO ShoppingListRequest (ingredient_id, person_id, created_date)
        VALUES (?, ?, ?)
    ''', (request.ingredient_id, request.person_id, request.created_date.isoformat())) as cursor:
        request.id = cursor.lastrowid

    return request

async def request_meal(conn, person: Person, meal: Meal) -> ShoppingListRequest:
    request = ShoppingListRequest(meal_id=meal.id, meal=meal, person_id=person.id, created_date=datetime.now().astimezone())
    
    async with conn.execute('''
        INSERT INTO ShoppingListRequest (meal_id, person_id, created_date)
        VALUES (?, ?, ?)
    ''', (request.meal_id, request.person_id, request.created_date.isoformat())) as cursor:
        request.id = cursor.lastrowid

    return request

async def unrequest_meal(conn, meal: Meal) -> None:
    await conn.execute('''
        DELETE FROM ShoppingListRequest
        WHERE list_id IS NULL AND meal_id = ?
    ''', (meal.id,))

def get_current_requests(conn) -> AsyncIterator[ShoppingListRequest]:
    return find_requests_by_list_id(conn, None)

async def sync_persons_requested_ingredients(conn, person: Person, requests: List[Ingredient]) -> AsyncIterator[ShoppingListRequest]:
    # Delete existing and insert all as new
    await conn.execute('''
        DELETE FROM ShoppingListRequest
        WHERE list_id IS NULL AND person_id = ? AND ingredient_id IS NOT NULL
    ''', (person.id,))

    for ingredient in requests:
        ingredient.id = -1
        yield await request_ingredient(conn, person, ingredient)

async def get_shopping_list_with_meal(conn, meal_id: int) -> AsyncIterator[ShoppingList]:
    async with conn.execute('''
        SELECT list_id FROM ShoppingListRequest
        WHERE meal_id = ? AND list_id IS NOT NULL
    ''', (meal_id,)) as cursor:
        async for row in cursor:
            list_id = row[0]
            yield await load_shopping_list(conn, list_id)