matrix_shop_bot/db.py

   1 import logging
   2 import os
   3 from datetime import datetime as dt
   4 from decimal import Decimal
   5 from typing import List
   6
   7 from sqlmodel import Field, Session, SQLModel, create_engine, select, text, Relationship, func
   8
   9 from sqlite3 import IntegrityError
  10
  11 MODULE='shop'
  12
  13 DATE_FORMAT = '%Y-%m-%d'
  14
  15 def today():
  16     return dt.now().strftime(DATE_FORMAT)
  17
  18
  19 logger = logging.getLogger(__name__)
  20
  21 class ItemStoreLink(SQLModel, table=True):
  22     item_id: int | None = Field(default=None, foreign_key="item.id", primary_key=True)
  23     store_id: int | None = Field(default=None, foreign_key="store.id", primary_key=True)
  24
  25
  26 class Item(SQLModel, table=True):
  27     '''An item to shop for.'''
  28     id: int | None = Field(default=None, primary_key=True)
  29     name: str = Field(default=None, unique=True)
  30     price: Decimal | None = Field(decimal_places=2)
  31     stores: list["Store"] = Relationship(back_populates="shopping_list", link_model=ItemStoreLink)
  32     purchase_date: dt | None = Field(default=None)
  33     def __str__(self):
  34         return f'{self.name} (#{self.id})'
  35
  36 class ItemZone(SQLModel, table=True):
  37     '''Where an item is usually found.
  38
  39     Matched only by loose text matching, since items come and go, over time.
  40     '''
  41     id: int | None = Field(default=None, primary_key=True)
  42     name: str = Field(default=None, unique=True)  # milk
  43     zone: str = Field(default=None, unique=False) # dairy section
  44
  45
  46 class PreferStore(SQLModel, table=True):
  47     '''Tracks which store we are shopping at, at the moment.'''
  48     store_id: int | None = Field(default=None, foreign_key="store.id", primary_key=True)
  49
  50 class Store(SQLModel, table=True):
  51     '''A place to shop at.'''
  52     id: int | None = Field(default=None, primary_key=True)
  53     name: str = Field(default=None, unique=True)
  54     shopping_list: list["Item"] = Relationship(back_populates="stores", link_model=ItemStoreLink)
  55     # purchase_history: list["PurchaseRecord"] = Relationship(back_populates="store")
  56     def __str__(self):
  57         return f'{self.name} (#{self.id})'
  58
  59     def get_pretty_shopping_list(self) -> List[str]:
  60         results = []
  61         results.append(f' ## {self.name}')
  62         items = self.shopping_list
  63         if not items:
  64             results.append(f'-- Nothing to shop for at {self.name} --')
  65         else:
  66             for item in items:
  67                 if(item.purchase_date):
  68                     results.append(f' - [x] {item}')
  69                 else:
  70                     results.append(f' - [ ] {item}')
  71         return results
  72
  73
  74 class PurchaseRecord(SQLModel, table=True):
  75     '''Tracks when and where items were purchased.'''
  76     id: int | None = Field(default=None, primary_key=True)
  77     name: str = Field(default=None)
  78     price: Decimal | None = Field(decimal_places=2)
  79     store: str | None = Field(default=None)
  80     purchase_date: dt | None = Field(default=None)
  81     def __str__(self):
  82         date = self.purchase_date.strftime(DATE_FORMAT)
  83         return f'Bought {self.name} at {self.store} on {date}'
  84
  85 def _get_with_id(cls, item_id):
  86     return select(cls).where(cls.id == item_id)
  87
  88 def _get_with_name(cls, name):
  89     return select(cls).where(func.lower(cls.name) == func.lower(name))
  90
  91
  92 class Data():
  93     def __init__(self, db_file=None):
  94         if not db_file:
  95             db_file = os.path.join(os.environ.get('MATRIX_DATA_DIR', ''), f'{MODULE}.db')
  96
  97         self.file = db_file
  98         logger.debug(f'Shop bot is using db file {self.file}')
  99         self.engine = self.get_engine()
 100         self.last_store_id = None
 101
 102     def _set_pref_store_id(self, store_id):
 103         with Session(self.engine) as session:
 104             prefs = session.exec(select(PreferStore)).all()
 105             for pref in prefs:
 106                 session.delete(pref)
 107             pref = PreferStore()
 108             pref.store_id = store_id
 109             session.add(pref)
 110             session.commit()
 111
 112     def _get_pref_store_id(self):
 113         with Session(self.engine) as session:
 114             pref = session.exec(select(PreferStore)).one()
 115             return pref.store_id
 116
 117     def _add_with_name(self, cls, name):
 118         item_id = -1
 119         try:
 120             with Session(self.engine) as session:
 121                 new = cls(name=name)
 122                 session.add(new)
 123                 session.commit()
 124                 item_id = new.id
 125         except:
 126             # duplicate add
 127             # recover by returning the item_id
 128             item = session.exec(select(Item).where(Item.name == name)).first()
 129             item_id = item.id
 130
 131         return item_id
 132
 133     def add_store(self, name):
 134         store_id = self._add_with_name(Store, name)
 135         self._set_pref_store_id(store_id)
 136         return store_id
 137
 138     def add_item(self, name):
 139         names = name.split(',')
 140         for name in names:
 141             item_id = self._add_with_name(Item, name.strip())
 142             store_id = self._get_pref_store_id()
 143             self.assign_store(item_id, store_id)
 144         return item_id
 145
 146     def assign_last_store(self, item_id):
 147         store_id = self._get_pref_store_id()
 148         self.assign_store(item_id, store_id)
 149         return store_id
 150
 151     def assign_zone(self, item_name, zone_name):
 152         item_id = -1
 153         with Session(self.engine) as session:
 154             try:
 155                 new = ItemZone(name=item_name)
 156                 new.zone = zone_name
 157                 session.add(new)
 158                 session.commit()
 159                 item_id = new.id
 160             except:
 161                 pass # Already assigned
 162         return item_id
 163
 164     def assign_store(self, item_id, store_id):
 165         self._set_pref_store_id(store_id)
 166
 167         with Session(self.engine) as session:
 168             store = session.exec(_get_with_id(Store, store_id)).first()
 169             item = session.exec(_get_with_id(Item, item_id)).first()
 170             try:
 171                 item.stores.append(store)
 172                 session.commit()
 173             except:
 174                 pass # Already assigned
 175         return ['Success']
 176
 177     def list_purchases(self, store_name):
 178         with Session(self.engine) as session:
 179             cmd = select(PurchaseRecord).where(func.lower(PurchaseRecord.store) == func.lower(store_name))
 180             records = session.exec(cmd).all()
 181             return records
 182
 183     def purchase_items(self, item_ids):
 184         results = []
 185         with Session(self.engine) as session:
 186             # Mark items as purchased
 187             for item_id in item_ids.split(' '):
 188                 item = session.exec(_get_with_id(Item, item_id)).first()
 189                 if item:
 190                     item.purchase_date = dt.today()
 191
 192                 # Track purchase history (locally only!)
 193                 # to generate recommended shopping lists
 194                 record = PurchaseRecord()
 195                 record.name = item.name
 196                 record.price = item.price
 197                 record.purchase_date = item.purchase_date
 198
 199                 shopping_at = self._get_pref_store_id() # We bought it here
 200                 cmd = select(Store).where(Store.id == shopping_at)
 201                 store = session.exec(cmd).one()
 202                 record.store = store.name
 203                 session.add(record)
 204
 205                 results.append(str(record))
 206                 # results.append(str(item))
 207
 208             session.commit()
 209         return results
 210
 211
 212     def clear_list(self, store_name, clear_all=False):
 213         results = []
 214         with Session(self.engine) as session:
 215             store = session.exec(_get_with_name(Store, store_name)).first()
 216             if not store:
 217                 return [f'Store {store_name} not found.']
 218
 219             for item in store.shopping_list:
 220                 if clear_all or item.purchase_date:
 221                     session.delete(item)
 222
 223             session.commit()
 224             results = store.get_pretty_shopping_list()
 225         return results
 226
 227     def get_stores(self):
 228         with Session(self.engine) as session:
 229             cmd = select(Store)
 230             return session.exec(cmd).all()
 231
 232     def get_unsorted_list(self):
 233
 234         results = []
 235         with Session(self.engine) as session:
 236             cmd = select(Item).where(Item.stores == None)
 237             items = session.exec(cmd).all()
 238
 239             if not items:
 240                 return [f'No unsorted items!']
 241             return items
 242
 243     def get_last_shopping_list(self) -> List[str]:
 244         store_id = self._get_pref_store_id()
 245         with Session(self.engine) as session:
 246             store = session.exec(_get_with_id(Store, store_id)).first()
 247             if store:
 248                 return store.get_pretty_shopping_list()
 249         return [f'Last shopping list store not found.']
 250
 251     def get_shopping_list(self, store_name) -> List[str]:
 252         retsults = []
 253         with Session(self.engine) as session:
 254             store = session.exec(_get_with_name(Store, store_name)).first()
 255             if not store:
 256                 return [f'Store {store_name} not found.']
 257             results = store.get_pretty_shopping_list()
 258         self._set_pref_store_id(store.id)
 259         return results
 260
 261     def remove_item(self, item_id):
 262         with Session(self.engine) as session:
 263             item = session.exec(_get_with_id(Item, item_id)).first()
 264             session.delete(item)
 265             session.commit()
 266
 267     def remove_store(self, store_id):
 268         with Session(self.engine) as session:
 269             store = session.exec(_get_with_id(Store, store_id)).first()
 270             session.delete(store)
 271             session.commit()
 272
 273     def get_engine(self):
 274         return create_engine(f"sqlite:///{self.file}")
 275
 276     def setup(self):
 277         SQLModel.metadata.create_all(self.engine)
 278         return 'Success'
 279
 280
 281
 282