matrix_chore_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
   6
   7 ## Todo List
   8 # - [ ] add Vision FK
   9 # - [ ] add notes to each chore
  10
  11 DATE_FORMAT = '%Y-%m-%d'
  12
  13 def today():
  14     return dt.now().strftime(DATE_FORMAT)
  15
  16 DEFAULT_DB_FILE = os.path.join(os.environ.get('MATRIX_DATA_DIR', ''), 'chores.db')
  17 DB_FILE = os.environ.get('MATRIX_CHORE_DB_FILE', DEFAULT_DB_FILE)
  18
  19 logger = logging.getLogger(__name__)
  20 logging.basicConfig(filename='chore.log', encoding='utf-8', level=logging.DEBUG)
  21
  22 class Chore(SQLModel, table=True):
  23     id: int | None = Field(default=None, primary_key=True)
  24     name: str = Field(default=None, unique=True)
  25
  26     def __str__(self):
  27         return f'{self.name} ({self.id})'
  28
  29
  30 class Record(SQLModel, table=True):
  31     id: int | None = Field(default=None, primary_key=True)
  32     chore_id: int | None = Field(default=None, foreign_key="chore.id")
  33     complete_date: dt = Field(default_factory=lambda: dt.now())
  34
  35
  36 def get_chore(session, chore_id):
  37     cmd = select(Chore).where(Chore.id == chore_id )
  38     return session.exec(cmd).first()
  39
  40 class Data():
  41     def __init__(self, db_file=None):
  42
  43         if not db_file:
  44             db_file = DB_FILE
  45
  46         self.file = db_file
  47         self.engine = self.get_engine()
  48
  49     def get_engine(self):
  50         return create_engine(f"sqlite:///{self.file}")
  51
  52     def setup(self):
  53         SQLModel.metadata.create_all(self.engine)
  54
  55
  56     @staticmethod
  57     def check_chore(session, chore_id) -> str:
  58         chore = get_chore(session, chore_id)
  59         if(not chore):
  60             return f'Chore not found: {chore_id}'
  61
  62         sql = text('''
  63         SELECT * FROM Record
  64         WHERE chore_id = :chore_id
  65         AND complete_date >= :today
  66         ORDER BY complete_date DESC
  67         LIMIT 1
  68        ''')
  69         params = {}
  70         params['chore_id'] = chore_id
  71         params['today'] = today()
  72         results = session.execute(sql, params).first()
  73
  74         logger.debug(results)
  75         if(results):
  76             return f'- [x] {chore}'
  77         else:
  78             return f'- [ ] {chore}'
  79
  80
  81     def check(self, chore_id):
  82         with Session(self.engine) as session:
  83             return Data.check_chore(session, chore_id)
  84
  85     def list(self):
  86         results = []
  87         with Session(self.engine) as session:
  88             cmd = select(Chore)
  89             chores = session.exec(cmd).all()
  90             if(not chores):
  91                 results.append('None Found')
  92
  93             for chore in chores:
  94                 results.append(Data.check_chore(session, chore.id))
  95         return results
  96
  97     def filter(self, term):
  98         results = []
  99         for item in self.list():
 100             if term in item:
 101                 results.append(item)
 102
 103         return results
 104
 105
 106
 107
 108     def add(self, name):
 109         try:
 110             with Session(self.engine) as session:
 111                 new_chore = Chore(name=name)
 112                 session.add(new_chore)
 113                 session.commit()
 114             return f'Created chore {name}.'
 115         except:
 116             return f'Unable to create {name}. May be duplicate.'
 117
 118
 119
 120     def done(self, chore_id):
 121         with Session(self.engine) as session:
 122             chore = get_chore(session, chore_id)
 123             if(not chore):
 124                 return f'Chore not found: {chore_id}'
 125
 126             try:
 127                 session.add(Record(chore_id=chore.id))
 128                 session.commit()
 129                 return f'Marked chore as done: {chore.name}.'
 130             except Exception as ex:
 131                 return f'Unable to mark chore as done: {ex}'
 132
 133
 134     def undo(self, chore_id):
 135         with Session(self.engine) as session:
 136             cmd = select(Chore).where(Chore.id == chore_id)
 137             chore = session.exec(cmd).first()
 138
 139             if(not chore):
 140                 return f'Not found: {chore_id}'
 141
 142             try:
 143                 cmd = select(Record).where(Record.chore_id == chore_id).where(Record.complete_date >= today())
 144                 results = session.exec(cmd).all()
 145                 for result in results:
 146                     session.delete(result)
 147                 session.commit()
 148                 return f'Marked chore as *not* done: {chore.name}.'
 149             except Exception as ex:
 150                 return f'Unable to mark chore as done: {ex}'
 151
 152     def remove(self, chore_id):
 153         try:
 154             with Session(self.engine) as session:
 155                 chore = get_chore(session, chore_id)
 156                 if(not chore):
 157                     return f'Chore not found: {chore_id}'
 158
 159                 sql = text('''
 160                 DELETE FROM Record
 161                 WHERE chore_id = :chore_id
 162                ''')
 163                 params = {}
 164                 params['chore_id'] = chore_id
 165                 results = session.execute(sql, params)
 166                 logger.debug(results)
 167                 sql = text('''
 168                 DELETE FROM Chore
 169                 WHERE id = :chore_id
 170                ''')
 171                 results = session.execute(sql, params)
 172                 logger.debug(results)
 173                 session.commit()
 174                 return f'Removed chore {chore_id}'
 175         except Exception as ex:
 176             return f'Unable to remove chore: {ex}'