matrix_shop_bot/db.py

   1 import logging
   2 import os
   3 from datetime import datetime as dt
   4
   5 from sqlmodel import Field, Session, SQLModel, create_engine, select, text, Relationship
   6
   7 MODULE='shop'
   8
   9 DATE_FORMAT = '%Y-%m-%d'
  10
  11 def today():
  12     return dt.now().strftime(DATE_FORMAT)
  13
  14
  15 logger = logging.getLogger(__name__)
  16 logging.basicConfig(filename=f'{MODULE}.log', encoding='utf-8', level=logging.DEBUG)
  17
  18 class Item(SQLModel, table=True):
  19     id: int | None = Field(default=None, primary_key=True)
  20     name: str = Field(default=None, unique=True)
  21     price: Decimal | None = Field(decimal_places=2)
  22     stores: list["Store"] = Relationship(back_populates="shopping_list")
  23     purchase_date: dt | None = Field(default=None)
  24     def __str__(self):
  25         return f'- [ ] {self.name} (#{self.id})'
  26
  27 class Store(SQLModel, table=True):
  28     id: int | None = Field(default=None, primary_key=True)
  29     name: str = Field(default=None, unique=True)
  30     shopping_list: list["Item"] = Relationship(back_populates="stores")
  31     purchase_history: list["PurchaseRecord"] = Relationship(back_populates="store")
  32     def __str__(self):
  33         return f'{self.name} (#{self.id})'
  34
  35
  36 # May get used later...
  37 class PurchaseRecord(SQLModel, table=True):
  38     id: int | None = Field(default=None, primary_key=True)
  39     name: str = Field(default=None, unique=True)
  40     price: Decimal | None = Field(decimal_places=2)
  41     store: Store = Relationship(back_populates="purchase_history")
  42     purchase_date: dt | None = Field(default=None)
  43
  44
  45
  46
  47 class Data():
  48     def __init__(self, db_file=None):
  49         if not db_file:
  50             db_file = os.path.join(os.environ.get('MATRIX_DATA_DIR', ''), 'goals.db')
  51
  52         self.file = db_file
  53         self.engine = self.get_engine()
  54
  55     def _add_with_name(self, cls, name):
  56         id = -1
  57         with Session(self.engine) as session:
  58             new = cls(name=name)
  59             session.add(new)
  60             session.commit()
  61             id = new.id
  62         return id
  63
  64     def _get_with_id(self, cls, item_id):
  65         with Session(self.engine) as session:
  66             cmd = select(cls).where(cls.id == item_id)
  67             return session.exec(cmd).first()
  68
  69     def add_store(self, name):
  70         return self.add_with_name(Store, name)
  71
  72     def add_item(self, name):
  73         return self.add_with_name(Item, name)
  74
  75     def assign_store(self, store_id, item_id):
  76         with Session(self.engine) as session:
  77             store = self.get_with_id(Store, store_id)
  78             item = self.get_with_id(Item, item_id)
  79             item.stores.append(store)
  80             session.commit()
  81
  82     def purchase_item(self, item_id):
  83         with Session(self.engine) as session:
  84             item = self.get_with_id(Item, item_id)
  85             item.purchase_date = dt.today()
  86             session.commit()
  87
  88     def clear_list(self, store_id, clear_all=False):
  89         with Session(self.engine) as session:
  90             store = self.get_with_id(Store, store_id)
  91             for item in store.items:
  92                 if clear_all or item.purchased_date:
  93                     item.remove()
  94             session.commit()
  95
  96     def get_unsorted_list(self):
  97
  98         results = []
  99         with Session(self.engine) as session:
 100             cmd = select(Item).where(Item.store == None)
 101             items = session.exec(cmd).all()
 102
 103             if not items:
 104                 return [f'No unsorted items!']
 105             return items
 106
 107     def get_shopping_list(self, store_id) -> List(str):
 108         results = []
 109         with Session(self.engine) as session:
 110             cmd = select(Store).where(Store.id == store_id)
 111             store = session.exec(cmd).first()
 112             if not store:
 113                 return [f'Store {store_id} not found.']
 114
 115             results.append(f' ## {store.name}')
 116
 117             items = store.shopping_list
 118             if not items:
 119                 results.append('-- Nothing to shop for at {store.name} --')
 120
 121             for item in items:
 122                 if(item.purchase_date):
 123                     results.append(f' - [x] {item}')
 124                 else:
 125                     results.append(f' - [ ] {item}')
 126         return results
 127
 128     def get_engine(self):
 129         return create_engine(f"sqlite:///{self.file}")
 130
 131     def setup(self):
 132         SQLModel.metadata.create_all(self.engine)
 133         return 'Success'
 134
 135
 136
 137