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