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 add(self, name):
  98         try:
  99             with Session(self.engine) as session:
 100                 new_chore = Chore(name=name)
 101                 session.add(new_chore)
 102                 session.commit()
 103             return f'Created chore {name}.'
 104         except:
 105             return f'Unable to create {name}. May be duplicate.'
 106
 107
 108
 109     def done(self, chore_id):
 110         with Session(self.engine) as session:
 111             chore = get_chore(session, chore_id)
 112             if(not chore):
 113                 return f'Chore not found: {chore_id}'
 114
 115             try:
 116                 session.add(Record(chore_id=chore.id))
 117                 session.commit()
 118                 return f'Marked chore as done: {chore.name}.'
 119             except Exception as ex:
 120                 return f'Unable to mark chore as done: {ex}'
 121
 122
 123     def undo(self, chore_id):
 124         with Session(self.engine) as session:
 125             cmd = select(Chore).where(Chore.id == chore_id)
 126             chore = session.exec(cmd).first()
 127
 128             if(not chore):
 129                 return f'Not found: {chore_id}'
 130
 131             try:
 132                 cmd = select(Record).where(Record.chore_id == chore_id).where(Record.complete_date >= today())
 133                 results = session.exec(cmd).all()
 134                 for result in results:
 135                     session.delete(result)
 136                 session.commit()
 137                 return f'Marked chore as *not* done: {chore.name}.'
 138             except Exception as ex:
 139                 return f'Unable to mark chore as done: {ex}'
 140
 141     def remove(self, chore_id):
 142         try:
 143             with Session(self.engine) as session:
 144                 chore = get_chore(session, chore_id)
 145                 if(not chore):
 146                     return f'Chore not found: {chore_id}'
 147
 148                 sql = text('''
 149                 DELETE FROM Record
 150                 WHERE chore_id = :chore_id
 151                ''')
 152                 params = {}
 153                 params['chore_id'] = chore_id
 154                 results = session.execute(sql, params)
 155                 logger.debug(results)
 156                 sql = text('''
 157                 DELETE FROM Chore
 158                 WHERE id = :chore_id
 159                ''')
 160                 results = session.execute(sql, params)
 161                 logger.debug(results)
 162                 session.commit()
 163                 return f'Removed chore {chore_id}'
 164         except Exception as ex:
 165             return f'Unable to remove chore: {ex}'