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 MODULE='shop'
  10
  11 DATE_FORMAT = '%Y-%m-%d'
  12
  13 def today():
  14     return dt.now().strftime(DATE_FORMAT)
  15
  16
  17 logger = logging.getLogger(__name__)
  18 logging.basicConfig(filename=f'{MODULE}.log', encoding='utf-8', level=logging.DEBUG)
  19
  20 class ItemStoreLink(SQLModel, table=True):
  21     item_id: int | None = Field(default=None, foreign_key="item.id", primary_key=True)
  22     store_id: int | None = Field(default=None, foreign_key="store.id", primary_key=True)
  23
  24
  25 class Item(SQLModel, table=True):
  26     id: int | None = Field(default=None, primary_key=True)
  27     name: str = Field(default=None, unique=True)
  28     price: Decimal | None = Field(decimal_places=2)
  29     stores: list["Store"] = Relationship(back_populates="shopping_list", link_model=ItemStoreLink)
  30     purchase_date: dt | None = Field(default=None)
  31     def __str__(self):
  32         return f'{self.name} (#{self.id})'
  33
  34 class PreferStore(SQLModel, table=True):
  35     store_id: int | None = Field(default=None, foreign_key="store.id", primary_key=True)
  36
  37 class Store(SQLModel, table=True):
  38     id: int | None = Field(default=None, primary_key=True)
  39     name: str = Field(default=None, unique=True)
  40     shopping_list: list["Item"] = Relationship(back_populates="stores", link_model=ItemStoreLink)
  41     # purchase_history: list["PurchaseRecord"] = Relationship(back_populates="store")
  42     def __str__(self):
  43         return f'{self.name} (#{self.id})'
  44
  45
  46
  47
  48
  49 # May get used later...
  50 #class PurchaseRecord(SQLModel, table=True):
  51 #    id: int | None = Field(default=None, primary_key=True)
  52 #    name: str = Field(default=None, unique=True)
  53 #    price: Decimal | None = Field(decimal_places=2)
  54 #    store: Store = Relationship(back_populates="purchase_history")
  55 #    purchase_date: dt | None = Field(default=None)
  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         self.engine = self.get_engine()
  68         self.last_store_id = None
  69
  70     def _set_pref_store_id(self, store_id):
  71         with Session(self.engine) as session:
  72             prefs = session.exec(select(PreferStore)).all()
  73             for pref in prefs:
  74                 session.delete(pref)
  75             pref = PreferStore()
  76             pref.store_id = store_id
  77             session.add(pref)
  78             session.commit()
  79
  80     def _get_pref_store_id(self):
  81         with Session(self.engine) as session:
  82             pref = session.exec(select(PreferStore)).one()
  83             return pref.store_id
  84
  85     def _add_with_name(self, cls, name):
  86         id = -1
  87         with Session(self.engine) as session:
  88             new = cls(name=name)
  89             session.add(new)
  90             session.commit()
  91             id = new.id
  92         return id
  93
  94     def add_store(self, name):
  95         store_id = self._add_with_name(Store, name)
  96         self.last_store_id = store_id
  97         return store_id
  98
  99     def add_item(self, name):
 100         item_id = self._add_with_name(Item, name)
 101         if isinstance(self.last_store_id, int):
 102             with Session(self.engine) as session:
 103                 store_id = self.last_store_id
 104                 store = session.exec(_get_with_id(Store, store_id)).first()
 105                 item = session.exec(_get_with_id(Item, item_id)).first()
 106                 item.stores.append(store)
 107                 session.commit()
 108         return item_id
 109
 110     def assign_store(self, item_id, store_id=None):
 111         if store_id:
 112             self.last_store_id = store_id
 113         else:
 114             store_id = self.last_store_id
 115
 116         with Session(self.engine) as session:
 117             store = session.exec(_get_with_id(Store, store_id)).first()
 118             item = session.exec(_get_with_id(Item, item_id)).first()
 119             item.stores.append(store)
 120             session.commit()
 121
 122     def purchase_item(self, item_id):
 123         with Session(self.engine) as session:
 124             item = session.exec(_get_with_id(Item, item_id)).first()
 125             item.purchase_date = dt.today()
 126             session.commit()
 127
 128     def clear_list(self, store_id, clear_all=False):
 129         with Session(self.engine) as session:
 130             store = session.exec(_get_with_id(Store, store_id)).first()
 131             for item in store.shopping_list:
 132                 if clear_all or item.purchase_date:
 133                     session.delete(item)
 134             session.commit()
 135
 136     def get_stores(self):
 137         with Session(self.engine) as session:
 138             cmd = select(Store)
 139             return session.exec(cmd).all()
 140
 141     def get_unsorted_list(self):
 142
 143         results = []
 144         with Session(self.engine) as session:
 145             cmd = select(Item).where(Item.stores == None)
 146             items = session.exec(cmd).all()
 147
 148             if not items:
 149                 return [f'No unsorted items!']
 150             return items
 151
 152
 153     def get_shopping_list_for_store(self, store_id) -> List[str]:
 154         self._set_pref_store_id(store_id)
 155         return self.get_shopping_list(store_id)
 156
 157     def get_last_shopping_list(self) -> List[str]:
 158         store_id = self._get_pref_store_id()
 159         return self.get_shopping_list(store_id)
 160
 161     def get_shopping_list(self, store_id) -> List[str]:
 162         results = []
 163         with Session(self.engine) as session:
 164             cmd = select(Store).where(Store.id == store_id)
 165             store = session.exec(cmd).first()
 166             if not store:
 167                 return [f'Store {store_id} not found.']
 168
 169             results.append(f' ## {store.name}')
 170
 171             items = store.shopping_list
 172             if not items:
 173                 results.append(f'-- Nothing to shop for at {store.name} --')
 174             else:
 175                 for item in items:
 176                     if(item.purchase_date):
 177                         results.append(f' - [x] {item}')
 178                     else:
 179                         results.append(f' - [ ] {item}')
 180         return results
 181
 182     def remove_item(self, item_id):
 183         with Session(self.engine) as session:
 184             item = session.exec(_get_with_id(Item, item_id)).first()
 185             session.delete(item)
 186             session.commit()
 187
 188     def remove_store(self, store_id):
 189         with Session(self.engine) as session:
 190             store = session.exec(_get_with_id(Store, store_id)).first()
 191             session.delete(store)
 192             session.commit()
 193
 194     def get_engine(self):
 195         return create_engine(f"sqlite:///{self.file}")
 196
 197     def setup(self):
 198         SQLModel.metadata.create_all(self.engine)
 199         return 'Success'
 200
 201
 202
 203