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
   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     id: int | None = Field(default=None, primary_key=True)
  28     name: str = Field(default=None, unique=True)
  29     price: Decimal | None = Field(decimal_places=2)
  30     stores: list["Store"] = Relationship(back_populates="shopping_list", link_model=ItemStoreLink)
  31     purchase_date: dt | None = Field(default=None)
  32     def __str__(self):
  33         return f'{self.name} (#{self.id})'
  34
  35 class PreferStore(SQLModel, table=True):
  36     store_id: int | None = Field(default=None, foreign_key="store.id", primary_key=True)
  37
  38 class Store(SQLModel, table=True):
  39     id: int | None = Field(default=None, primary_key=True)
  40     name: str = Field(default=None, unique=True)
  41     shopping_list: list["Item"] = Relationship(back_populates="stores", link_model=ItemStoreLink)
  42     # purchase_history: list["PurchaseRecord"] = Relationship(back_populates="store")
  43     def __str__(self):
  44         return f'{self.name} (#{self.id})'
  45
  46 # May get used later...
  47 class PurchaseRecord(SQLModel, table=True):
  48     id: int | None = Field(default=None, primary_key=True)
  49     name: str = Field(default=None)
  50     price: Decimal | None = Field(decimal_places=2)
  51     store: str | None = Field(default=None)
  52     purchase_date: dt | None = Field(default=None)
  53     def __str__(self):
  54         date = self.purchase_date.strftime(DATE_FORMAT)
  55         return f'{self.name} at {self.store} on {date}'
  56
  57 def _get_with_id(cls, item_id):
  58     return select(cls).where(cls.id == item_id)
  59
  60
  61 class Data():
  62     def __init__(self, db_file=None):
  63         if not db_file:
  64             db_file = os.path.join(os.environ.get('MATRIX_DATA_DIR', ''), f'{MODULE}.db')
  65
  66         self.file = db_file
  67         logger.debug(f'Shop bot is using db file {self.file}')
  68         self.engine = self.get_engine()
  69         self.last_store_id = None
  70
  71     def _set_pref_store_id(self, store_id):
  72         with Session(self.engine) as session:
  73             prefs = session.exec(select(PreferStore)).all()
  74             for pref in prefs:
  75                 session.delete(pref)
  76             pref = PreferStore()
  77             pref.store_id = store_id
  78             session.add(pref)
  79             session.commit()
  80
  81     def _get_pref_store_id(self):
  82         with Session(self.engine) as session:
  83             pref = session.exec(select(PreferStore)).one()
  84             return pref.store_id
  85
  86     def _add_with_name(self, cls, name):
  87         item_id = -1
  88         try:
  89             with Session(self.engine) as session:
  90                 new = cls(name=name)
  91                 session.add(new)
  92                 session.commit()
  93                 item_id = new.id
  94         except:
  95             # duplicate add
  96             # recover by returning the item_id
  97             item = session.exec(select(Item).where(Item.name == name)).first()
  98             item_id = item.id
  99
 100         return item_id
 101
 102     def add_store(self, name):
 103         store_id = self._add_with_name(Store, name)
 104         self._set_pref_store_id(store_id)
 105         return store_id
 106
 107     def add_item(self, name):
 108         names = name.split(',')
 109         for name in names:
 110             item_id = self._add_with_name(Item, name.strip())
 111             store_id = self._get_pref_store_id()
 112             self.assign_store(item_id, store_id)
 113         return item_id
 114
 115     def assign_last_store(self, item_id):
 116         store_id = self._get_pref_store_id()
 117         self.assign_store(item_id, store_id)
 118         return store_id
 119
 120     def assign_store(self, item_id, store_id):
 121         self._set_pref_store_id(store_id)
 122
 123         with Session(self.engine) as session:
 124             store = session.exec(_get_with_id(Store, store_id)).first()
 125             item = session.exec(_get_with_id(Item, item_id)).first()
 126             try:
 127                 item.stores.append(store)
 128                 session.commit()
 129             except:
 130                 pass # Already assigned
 131         return ['Success']
 132
 133     def list_purchases(self, store_name):
 134         with Session(self.engine) as session:
 135             cmd = select(PurchaseRecord).where(PurchaseRecord.store == store_name)
 136             records = session.exec(cmd).all()
 137             return records
 138
 139     def purchase_items(self, item_ids):
 140         results = []
 141         with Session(self.engine) as session:
 142             # Mark items as purchased
 143             for item_id in item_ids.split(' '):
 144                 item = session.exec(_get_with_id(Item, item_id)).first()
 145                 if item:
 146                     item.purchase_date = dt.today()
 147
 148                 # Track purchase history (locally only!)
 149                 # to generate recommended shopping lists
 150                 record = PurchaseRecord()
 151                 record.name = item.name
 152                 record.price = item.price
 153                 record.purchase_date = item.purchase_date
 154
 155                 shopping_at = self._get_pref_store_id() # We bought it here
 156                 cmd = select(Store).where(Store.id == shopping_at)
 157                 store = session.exec(cmd).one()
 158                 record.store = store.name
 159                 session.add(record)
 160
 161                 results.append(str(record))
 162                 # results.append(str(item))
 163
 164             session.commit()
 165         return results
 166
 167
 168     def clear_list(self, store_id, clear_all=False):
 169         with Session(self.engine) as session:
 170             store = session.exec(_get_with_id(Store, store_id)).first()
 171             for item in store.shopping_list:
 172                 if clear_all or item.purchase_date:
 173                     session.delete(item)
 174             session.commit()
 175
 176     def get_stores(self):
 177         with Session(self.engine) as session:
 178             cmd = select(Store)
 179             return session.exec(cmd).all()
 180
 181     def get_unsorted_list(self):
 182
 183         results = []
 184         with Session(self.engine) as session:
 185             cmd = select(Item).where(Item.stores == None)
 186             items = session.exec(cmd).all()
 187
 188             if not items:
 189                 return [f'No unsorted items!']
 190             return items
 191
 192
 193     def get_shopping_list_for_store(self, store_id) -> List[str]:
 194         self._set_pref_store_id(store_id)
 195         return self.get_shopping_list(store_id)
 196
 197     def get_last_shopping_list(self) -> List[str]:
 198         store_id = self._get_pref_store_id()
 199         return self.get_shopping_list(store_id)
 200
 201     def get_shopping_list(self, store_id) -> List[str]:
 202         results = []
 203         with Session(self.engine) as session:
 204             cmd = select(Store).where(Store.id == store_id)
 205             store = session.exec(cmd).first()
 206             if not store:
 207                 return [f'Store {store_id} not found.']
 208
 209             results.append(f' ## {store.name}')
 210
 211             items = store.shopping_list
 212             if not items:
 213                 results.append(f'-- Nothing to shop for at {store.name} --')
 214             else:
 215                 for item in items:
 216                     if(item.purchase_date):
 217                         results.append(f' - [x] {item}')
 218                     else:
 219                         results.append(f' - [ ] {item}')
 220         return results
 221
 222     def remove_item(self, item_id):
 223         with Session(self.engine) as session:
 224             item = session.exec(_get_with_id(Item, item_id)).first()
 225             session.delete(item)
 226             session.commit()
 227
 228     def remove_store(self, store_id):
 229         with Session(self.engine) as session:
 230             store = session.exec(_get_with_id(Store, store_id)).first()
 231             session.delete(store)
 232             session.commit()
 233
 234     def get_engine(self):
 235         return create_engine(f"sqlite:///{self.file}")
 236
 237     def setup(self):
 238         SQLModel.metadata.create_all(self.engine)
 239         return 'Success'
 240
 241
 242
 243