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