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