matrix_shop_bot/db.py

   1 import logging
   2 import os
   3 from datetime import datetime as dt
   4 from datetime import timedelta
   5 from decimal import Decimal
   6 from typing import List
   7
   8 from sqlmodel import Field, Session, SQLModel, create_engine, select, text, Relationship, func
   9
  10 from sqlite3 import IntegrityError
  11
  12 MODULE='shop'
  13
  14 DATE_FORMAT = '%Y-%m-%d'
  15
  16 def today():
  17     return dt.now().strftime(DATE_FORMAT)
  18
  19
  20 logger = logging.getLogger(__name__)
  21
  22 class ItemStoreLink(SQLModel, table=True):
  23     item_id: int | None = Field(default=None, foreign_key="item.id", primary_key=True)
  24     store_id: int | None = Field(default=None, foreign_key="store.id", primary_key=True)
  25
  26
  27 class Item(SQLModel, table=True):
  28     '''An item to shop for.'''
  29     id: int | None = Field(default=None, primary_key=True)
  30     name: str = Field(default=None, unique=True)
  31     price: Decimal | None = Field(decimal_places=2)
  32     stores: list["Store"] = Relationship(back_populates="shopping_list", link_model=ItemStoreLink)
  33     purchase_date: dt | None = Field(default=None)
  34     def __str__(self):
  35         if(self.purchase_date):
  36             return f'- [x] {self.name} (#{self.id})'
  37         else:
  38             return f'- [ ] {self.name} (#{self.id})'
  39
  40 class ItemZone(SQLModel, table=True):
  41     '''Where an item is usually found.
  42
  43     Matched only by loose text matching, since items come and go, over time.
  44     '''
  45     id: int | None = Field(default=None, primary_key=True)
  46     name: str = Field(default=None, unique=True)  # milk
  47     zone: str = Field(default=None, unique=False) # dairy section
  48
  49 class PreferStore(SQLModel, table=True):
  50     '''Tracks which store we are shopping at, at the moment.'''
  51     store_id: int | None = Field(default=None, foreign_key="store.id", primary_key=True)
  52
  53 def _get_item_zone(item_name, session) -> str:
  54     zone = session.exec(select(ItemZone).where(func.lower(item_name).contains(func.lower(ItemZone.name)))).first()
  55     if zone:
  56         return zone.zone
  57     return 'No Zone Assigned'
  58
  59 def _get_pretty_shopping_list(store_id, session) -> List[str]:
  60     store = session.exec(_get_with_id(Store, store_id)).first()
  61     results = []
  62     results.append(f' ## {store.name}')
  63     items = store.shopping_list
  64     if not items:
  65         results.append(f'-- Nothing to shop for at {store.name} --')
  66
  67     zones = {}
  68
  69     for item in items:
  70         zone = str(_get_item_zone(item.name, session)).title() # Title Case
  71         if not zone in zones:
  72             zones[zone] = []
  73         zones[zone].append(str(item).title())
  74     for zone in zones:
  75         zones[zone].sort()
  76         results.append(f' ### {zone}')
  77         for item in zones[zone]:
  78             results.append(item)
  79     return results
  80
  81 class Store(SQLModel, table=True):
  82     '''A place to shop at.'''
  83     id: int | None = Field(default=None, primary_key=True)
  84     name: str = Field(default=None, unique=True)
  85     shopping_list: list["Item"] = Relationship(back_populates="stores", link_model=ItemStoreLink)
  86     # purchase_history: list["PurchaseRecord"] = Relationship(back_populates="store")
  87     def __str__(self):
  88         return f'{self.name} (#{self.id})'
  89
  90
  91 class PurchaseRecord(SQLModel, table=True):
  92     '''Tracks when and where items were purchased.'''
  93     id: int | None = Field(default=None, primary_key=True)
  94     name: str = Field(default=None)
  95     price: Decimal | None = Field(decimal_places=2)
  96     store: str | None = Field(default=None)
  97     purchase_date: dt | None = Field(default=None)
  98     def __str__(self):
  99         date = self.purchase_date.strftime(DATE_FORMAT)
 100         return f'Bought {self.name} at {self.store} on {date}'
 101
 102 def _get_with_id(cls, item_id):
 103     return select(cls).where(cls.id == item_id)
 104
 105 def _get_with_name(cls, name):
 106     return select(cls).where(func.lower(cls.name) == func.lower(name))
 107
 108 def _rename(cls, session, old, new) -> bool:
 109     with session:
 110         record = session.exec(_get_with_name(cls, old)).first()
 111         record.name = new
 112         session.commit()
 113     return True
 114
 115 class Data():
 116     def __init__(self, db_file=None):
 117         if not db_file:
 118             db_file = os.path.join(os.environ.get('MATRIX_DATA_DIR', ''), f'{MODULE}.db')
 119
 120         self.file = db_file
 121         logger.debug(f'Shop bot is using db file {self.file}')
 122         self.engine = self.get_engine()
 123         self.last_store_id = None
 124
 125     def rename_store(self, old, new):
 126         with Session(self.engine) as session:
 127             return _rename(Store, session, old, new)
 128
 129     def rename_item(self, old, new):
 130         with Session(self.engine) as session:
 131             return _rename(Item, session, old, new)
 132
 133     def _set_pref_store_id(self, store_id):
 134         with Session(self.engine) as session:
 135             prefs = session.exec(select(PreferStore)).all()
 136             for pref in prefs:
 137                 session.delete(pref)
 138             pref = PreferStore()
 139             pref.store_id = store_id
 140             session.add(pref)
 141             session.commit()
 142
 143     def _get_pref_store_id(self):
 144         with Session(self.engine) as session:
 145             pref = session.exec(select(PreferStore)).one()
 146             return pref.store_id
 147
 148     def _add_with_name(self, cls, name):
 149         item_id = -1
 150         try:
 151             with Session(self.engine) as session:
 152                 new = cls(name=name)
 153                 session.add(new)
 154                 session.commit()
 155                 item_id = new.id
 156         except:
 157             # duplicate add
 158             # recover by returning the item_id
 159             item = session.exec(select(Item).where(Item.name == name)).first()
 160             item_id = item.id
 161
 162         return item_id
 163
 164     def add_store(self, name):
 165         store_id = self._add_with_name(Store, name)
 166         self._set_pref_store_id(store_id)
 167         return store_id
 168
 169     def add_item(self, name):
 170         names = name.split(',')
 171         for name in names:
 172             item_id = self._add_with_name(Item, name.strip())
 173             store_id = self._get_pref_store_id()
 174             self.assign_store(item_id, store_id)
 175         return item_id
 176
 177     def assign_last_store(self, item_id):
 178         store_id = self._get_pref_store_id()
 179         self.assign_store(item_id, store_id)
 180         return store_id
 181
 182     def assign_zone(self, item_name, zone_name):
 183         item_id = -1
 184         with Session(self.engine) as session:
 185             try:
 186                 new = ItemZone(name=item_name)
 187                 new.zone = zone_name
 188                 session.add(new)
 189                 session.commit()
 190                 item_id = new.id
 191             except:
 192                 pass # Already assigned
 193         return item_id
 194
 195     def assign_store(self, item_id, store_id):
 196         self._set_pref_store_id(store_id)
 197
 198         with Session(self.engine) as session:
 199             store = session.exec(_get_with_id(Store, store_id)).first()
 200             item = session.exec(_get_with_id(Item, item_id)).first()
 201             try:
 202                 item.stores.append(store)
 203                 session.commit()
 204             except:
 205                 pass # Already assigned
 206         return ['Success']
 207
 208     def list_purchases(self, store_name):
 209         '''List unique purchases at store in the past 90 days'''
 210         with Session(self.engine) as session:
 211             recent = dt.now() - timedelta(days=90)
 212             cmd = select(PurchaseRecord).where(func.lower(PurchaseRecord.store) == func.lower(store_name)).where(PurchaseRecord.purchase_date > recent)
 213             records = session.exec(cmd).all()
 214             bought = set()
 215             for item in records:
 216                  bought.add(item.name)
 217             return list(bought)
 218
 219     def purchase_items(self, item_ids):
 220         results = []
 221         with Session(self.engine) as session:
 222             # Mark items as purchased
 223             for item_id in item_ids.split(' '):
 224                 item = session.exec(_get_with_id(Item, item_id)).first()
 225                 if item:
 226                     item.purchase_date = dt.today()
 227
 228                 # Track purchase history (locally only!)
 229                 # to generate recommended shopping lists
 230                 record = PurchaseRecord()
 231                 record.name = item.name
 232                 record.price = item.price
 233                 record.purchase_date = item.purchase_date
 234
 235                 shopping_at = self._get_pref_store_id() # We bought it here
 236                 cmd = select(Store).where(Store.id == shopping_at)
 237                 store = session.exec(cmd).one()
 238                 record.store = store.name
 239                 session.add(record)
 240
 241                 results.append(str(record))
 242                 # results.append(str(item))
 243
 244             session.commit()
 245         return results
 246
 247
 248     def clear_list(self, store_name, clear_all=False):
 249         results = []
 250         with Session(self.engine) as session:
 251             store = session.exec(_get_with_name(Store, store_name)).first()
 252             if not store:
 253                 return [f'Store {store_name} not found.']
 254
 255             for item in store.shopping_list:
 256                 if clear_all or item.purchase_date:
 257                     session.delete(item)
 258
 259             session.commit()
 260             results = _get_pretty_shopping_list(store.id, session)
 261         return results
 262
 263     def get_stores(self):
 264         with Session(self.engine) as session:
 265             cmd = select(Store)
 266             return session.exec(cmd).all()
 267
 268     def get_unsorted_list(self):
 269
 270         results = []
 271         with Session(self.engine) as session:
 272             cmd = select(Item).where(Item.stores == None)
 273             items = session.exec(cmd).all()
 274
 275             if not items:
 276                 return [f'No unsorted items!']
 277             return items
 278
 279     def get_last_shopping_list(self) -> List[str]:
 280         store_id = self._get_pref_store_id()
 281         with Session(self.engine) as session:
 282             store = session.exec(_get_with_id(Store, store_id)).first()
 283             if store:
 284                 return _get_pretty_shopping_list(store.id, session)
 285         return [f'Last shopping list store not found.']
 286
 287     def get_shopping_list(self, store_name) -> List[str]:
 288         retsults = []
 289         with Session(self.engine) as session:
 290             store = session.exec(_get_with_name(Store, store_name)).first()
 291             if not store:
 292                 return [f'Store {store_name} not found.']
 293             results = _get_pretty_shopping_list(store.id, session)
 294         self._set_pref_store_id(store.id)
 295         return results
 296
 297     def remove_item(self, item_id):
 298         with Session(self.engine) as session:
 299             item = session.exec(_get_with_id(Item, item_id)).first()
 300             session.delete(item)
 301             session.commit()
 302
 303     def remove_store(self, store_id):
 304         with Session(self.engine) as session:
 305             store = session.exec(_get_with_id(Store, store_id)).first()
 306             session.delete(store)
 307             session.commit()
 308
 309     def get_engine(self):
 310         return create_engine(f"sqlite:///{self.file}")
 311
 312     def setup(self):
 313         SQLModel.metadata.create_all(self.engine)
 314         with Session(self.engine) as session:
 315             # Add some obvious store zones by default
 316             DEFAULT_ZONES = {
 317                 'Bakery': ['bread', 'cake', 'cookie', 'muffin', 'pie'],
 318                 'Baking': ['flour', 'sugar'],
 319                 'Books': ['book'],
 320                 'Clothes': ['bra', 'coat', 'jacket', 'panties', 'pants', 'shirt', 'shorts', 'socks', 'sweater', 'underwear'],
 321                 'Cooler': ['butter', 'cream', 'eggs', 'juice', 'milk', 'tea'],
 322                 'Deli': ['beef', 'chicken', 'salmon', 'steak', 'turkey'],
 323                 'Movies': ['movie', 'DVD'],
 324                 'Music': ['cd', 'album'],
 325                 'Produce': ['carrot', 'cucumber', 'greens', 'lettuce'],
 326                 'Shoes': ['shoe'],
 327                 'Television': ['tv', 'season'],
 328             }
 329             for zone in DEFAULT_ZONES:
 330                 logging.debug(f'Adding default items for {zone}')
 331                 for item in DEFAULT_ZONES[zone]:
 332                     try:
 333                         logging.debug(f'Attempt adding {item} to {zone}...')
 334                         new = ItemZone(name=item, zone=zone)
 335                         session.add(new)
 336                         session.commit()
 337                         logging.debug(f'Success adding {item} to {zone}...')
 338                     except Exception as ex:
 339                         session.rollback()
 340                         logging.debug(f'Failed adding {item} to {zone}: {ex}')
 341
 342         return 'Success'
 343
 344
 345
 346