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