import contextlib from functools import partial from operator import itemgetter from typing import List, Tuple import pint import psycopg2 from sqlalchemy.orm import Session from sqlalchemy.orm.exc import NoResultFound from wrolpi.common import logger, Base from wrolpi.db import get_db_session, get_db_curs from wrolpi.errors import APIError from .models import Inventory, Item, InventoriesVersion logger = logger.getChild(__name__) __all__ = [ 'DEFAULT_CATEGORIES', 'DEFAULT_INVENTORIES', 'delete_inventory', 'delete_items', 'get_brands', 'get_categories', 'get_inventories', 'get_inventories_version', 'get_items', 'get_next_inventories_version', 'get_unit_registry', 'increment_inventories_version', 'logger', 'save_inventory', 'save_item', 'sort_categories', 'update_inventory', 'update_item', ] UNIT_REGISTRY: pint.UnitRegistry = None def get_unit_registry() -> pint.UnitRegistry: global UNIT_REGISTRY if UNIT_REGISTRY: return UNIT_REGISTRY UNIT_REGISTRY = pint.UnitRegistry() return UNIT_REGISTRY DEFAULT_CATEGORIES = [ ('salt', 'cooking ingredients'), ('dry milk', 'dairy'), ('evaporated milk', 'dairy'), ('freeze dried', 'dairy'), ('shortening', 'fats'), ('vegetable oil', 'fats'), ('canned', 'fruits'), ('dehydrated', 'fruits'), ('freeze-dried', 'fruits'), ('corn meal', 'grains'), ('flour', 'grains'), ('oats', 'grains'), ('pasta', 'grains'), ('rice', 'grains'), ('wheat', 'grains'), ('dry beans', 'legumes'), ('lentils', 'legumes'), ('canned', 'meats'), ('dried', 'meats'), ('freeze dried', 'meals'), ('brown sugar', 'sugars'), ('corn syrup', 'sugars'), ('honey', 'sugars'), ('juice mix', 'sugars'), ('white sugar', 'sugars'), ('canned', 'vegetables'), ('dehydrated', 'vegetables'), ('freeze dried', 'vegetables'), ('water', 'water'), ('bottled', 'water'), ('barrel', 'water'), ] DEFAULT_INVENTORIES = [ 'Food Storage', ] # Categories are sorted by category/sub-category. sort_categories = partial(sorted, key=itemgetter(1, 0)) def get_inventories() -> List[Inventory]: with get_db_session() as session: inventories = session.query(Inventory).filter( Inventory.deleted_at == None, # noqa ).order_by( Inventory.viewed_at.desc(), ).all() inventories = list(inventories) return inventories IGNORED_INVENTORY_KEYS = {'viewed_at', 'created_at', 'deleted_at'} def _remove_conflicting_deleted_inventory(inventory: dict, session: Session): """ Remove any Inventories that share the inventory's name, but are marked as deleted. """ inventories = session.query(Inventory).filter( Inventory.deleted_at != None, # noqa Inventory.name == inventory['name'], ).all() for i in inventories: session.query(Item).filter_by(inventory_id=i.id).delete() session.query(Inventory).filter_by(id=i.id).delete() def save_inventory(inventory): inventory = {k: v for k, v in inventory.items() if k not in IGNORED_INVENTORY_KEYS} # Cleanup the whitespace. inventory['name'] = inventory['name'].strip() with get_db_session(commit=True) as session: _remove_conflicting_deleted_inventory(inventory, session) session.add(Inventory(**inventory)) def update_inventory(inventory_id: int, inventory: dict): inventory = {k: v for k, v in inventory.items() if k not in IGNORED_INVENTORY_KEYS} # Cleanup the whitespace. inventory['name'] = inventory['name'].strip() with get_db_session(commit=True) as session: _remove_conflicting_deleted_inventory(inventory, session) i = session.query(Inventory).filter_by(id=inventory_id).one() for key, value in inventory.items(): setattr(i, key, value) def delete_inventory(inventory_id: int): with get_db_session(commit=True) as session: inventory = session.query(Inventory).filter_by(id=inventory_id).one() inventory.delete() def get_categories() -> List[Tuple[str, str]]: """ Get all distinct sets of subcategory and category. """ categories = DEFAULT_CATEGORIES.copy() with get_db_curs() as curs: curs.execute("SELECT DISTINCT COALESCE(subcategory, ''), COALESCE(category, '') FROM item ORDER BY 1, 2") try: db_categories = curs.fetchall() db_categories = [tuple(i) for i in db_categories] categories = set(categories + db_categories) except psycopg2.ProgrammingError: # No categories in DB. pass return sort_categories(categories) def get_brands() -> List[Tuple[int, str]]: with get_db_curs() as curs: curs.execute("SELECT DISTINCT brand FROM item WHERE brand IS NOT NULL AND brand != '' ORDER BY 1") try: brands = curs.fetchall() except psycopg2.ProgrammingError: # No brands return [] return list(brands) def get_items(inventory_id: int) -> List[Base]: """ Get all Items in an Inventory, except deleted items. """ with get_db_session() as session: results = session.query(Item).filter( Item.inventory_id == inventory_id, Item.deleted_at == None, # noqa ).all() return results def save_item(inventory_id: int, item: dict): with get_db_session(commit=True) as session: item = Item(inventory_id=inventory_id, **item) session.add(item) def update_item(item_id: int, item: dict): with get_db_session(commit=True) as session: i = session.query(Item).filter_by(id=item_id).one() del item['id'] del item['created_at'] for key, value in item.items(): setattr(i, key, value) def delete_items(items_ids: List[int]): with get_db_curs(commit=True) as curs: curs.execute('UPDATE item SET deleted_at=current_timestamp WHERE id = ANY(%s) RETURNING id', (items_ids,)) deleted_ids = {i['id'] for i in curs.fetchall()} if set(items_ids) != deleted_ids: raise APIError('Could not delete the items') def get_inventories_version(): """ The inventory_version table contains a single row with the Inventories version integer. """ with get_db_session() as session: try: version = session.query(InventoriesVersion).one() return version.version except NoResultFound: # No version is saved yet. pass def get_next_inventories_version(): """ Get the current inventories version, increment it if there is one. """ version = get_inventories_version() if version: return version + 1 return 1 @contextlib.contextmanager def increment_inventories_version(): """ Context manager that will increment the Inventories Version by 1 when exiting. """ version = get_next_inventories_version() yield version with get_db_session(commit=True) as session: if session.query(InventoriesVersion).count() == 0: session.add(InventoriesVersion(version=version)) else: session.query(InventoriesVersion).one().version = version