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 Store(SQLModel, table=True):
  35     id: int | None = Field(default=None, primary_key=True)
  36     name: str = Field(default=None, unique=True)
  37     shopping_list: list["Item"] = Relationship(back_populates="stores", link_model=ItemStoreLink)
  38     # purchase_history: list["PurchaseRecord"] = Relationship(back_populates="store")
  39     def __str__(self):
  40         return f'{self.name} (#{self.id})'
  41
  42
  43
  44
  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, unique=True)
  50 #    price: Decimal | None = Field(decimal_places=2)
  51 #    store: Store = Relationship(back_populates="purchase_history")
  52 #    purchase_date: dt | None = Field(default=None)
  53
  54 def _get_with_id(cls, item_id):
  55     return select(cls).where(cls.id == item_id)
  56
  57
  58 class Data():
  59     def __init__(self, db_file=None):
  60         if not db_file:
  61             db_file = os.path.join(os.environ.get('MATRIX_DATA_DIR', ''), f'{MODULE}.db')
  62
  63         self.file = db_file
  64         self.engine = self.get_engine()
  65         self.last_store_id = None
  66
  67     def _add_with_name(self, cls, name):
  68         id = -1
  69         with Session(self.engine) as session:
  70             new = cls(name=name)
  71             session.add(new)
  72             session.commit()
  73             id = new.id
  74         return id
  75
  76     def add_store(self, name):
  77         store_id = self._add_with_name(Store, name)
  78         self.last_store_id = store_id
  79         return store_id
  80
  81     def add_item(self, name):
  82         item_id = self._add_with_name(Item, name)
  83         if self.last_store_id:
  84             with Session(self.engine) as session:
  85                 store_id = self.last_store_id
  86                 store = session.exec(_get_with_id(Store, store_id)).first()
  87                 item = session.exec(_get_with_id(Item, item_id)).first()
  88                 item.stores.append(store)
  89                 session.commit()
  90         return item_id
  91
  92     def assign_store(self, item_id, store_id=None):
  93         if store_id:
  94             self.last_store_id = store_id
  95         else:
  96             store_id = self.last_store_id
  97
  98         with Session(self.engine) as session:
  99             store = session.exec(_get_with_id(Store, store_id)).first()
 100             item = session.exec(_get_with_id(Item, item_id)).first()
 101             item.stores.append(store)
 102             session.commit()
 103
 104     def purchase_item(self, item_id):
 105         with Session(self.engine) as session:
 106             item = session.exec(_get_with_id(Item, item_id)).first()
 107             item.purchase_date = dt.today()
 108             session.commit()
 109
 110     def clear_list(self, store_id, clear_all=False):
 111         with Session(self.engine) as session:
 112             store = session.exec(_get_with_id(Store, store_id)).first()
 113             for item in store.shopping_list:
 114                 if clear_all or item.purchase_date:
 115                     session.delete(item)
 116             session.commit()
 117
 118     def get_stores(self):
 119         with Session(self.engine) as session:
 120             cmd = select(Store)
 121             return session.exec(cmd).all()
 122
 123     def get_unsorted_list(self):
 124
 125         results = []
 126         with Session(self.engine) as session:
 127             cmd = select(Item).where(Item.stores == None)
 128             items = session.exec(cmd).all()
 129
 130             if not items:
 131                 return [f'No unsorted items!']
 132             return items
 133
 134     def get_shopping_list(self, store_id=None) -> List[str]:
 135         if not store_id:
 136             store_id = self.last_store_id
 137
 138         results = []
 139         with Session(self.engine) as session:
 140             cmd = select(Store).where(Store.id == store_id)
 141             store = session.exec(cmd).first()
 142             if not store:
 143                 return [f'Store {store_id} not found.']
 144
 145             results.append(f' ## {store.name}')
 146
 147             items = store.shopping_list
 148             if not items:
 149                 results.append(f'-- Nothing to shop for at {store.name} --')
 150             else:
 151                 for item in items:
 152                     if(item.purchase_date):
 153                         results.append(f' - [x] {item}')
 154                     else:
 155                         results.append(f' - [ ] {item}')
 156         return results
 157
 158     def remove_item(self, item_id):
 159         with Session(self.engine) as session:
 160             item = session.exec(_get_with_id(Item, item_id)).first()
 161             session.delete(item)
 162             session.commit()
 163
 164     def remove_store(self, store_id):
 165         with Session(self.engine) as session:
 166             store = session.exec(_get_with_id(Store, store_id)).first()
 167             session.delete(store)
 168             session.commit()
 169
 170     def get_engine(self):
 171         return create_engine(f"sqlite:///{self.file}")
 172
 173     def setup(self):
 174         SQLModel.metadata.create_all(self.engine)
 175         return 'Success'
 176
 177
 178
 179