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 CHORE_DB = os.environ.get('CHORE_DB')
  12
  13 DATE_FORMAT = '%Y-%m-%d'
  14
  15 def today():
  16     return dt.now().strftime(DATE_FORMAT)
  17
  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 = CHORE_DB):
  42         self.file = db_file
  43         self.engine = self.get_engine()
  44
  45     def get_engine(self):
  46         return create_engine(f"sqlite:///{self.file}")
  47
  48     def setup(self):
  49         SQLModel.metadata.create_all(self.engine)
  50
  51
  52     @staticmethod
  53     def check_chore(session, chore_id) -> str:
  54         chore = get_chore(session, chore_id)
  55         if(not chore):
  56             return f'Chore not found: {chore_id}'
  57
  58         sql = text('''
  59         SELECT * FROM Record
  60         WHERE chore_id = :chore_id
  61         AND complete_date >= :today
  62         ORDER BY complete_date DESC
  63         LIMIT 1
  64        ''')
  65         params = {}
  66         params['chore_id'] = chore_id
  67         params['today'] = today()
  68         results = session.execute(sql, params).first()
  69
  70         logger.debug(results)
  71         if(results):
  72             return f'- [x] {chore}'
  73         else:
  74             return f'- [ ] {chore}'
  75
  76
  77     def check(self, chore_id):
  78         with Session(self.engine) as session:
  79             return Data.check_chore(session, chore_id)
  80
  81     def list(self):
  82         results = []
  83         with Session(self.engine) as session:
  84             cmd = select(Chore)
  85             chores = session.exec(cmd).all()
  86             if(not chores):
  87                 retuls.append('None Found')
  88
  89             for chore in chores:
  90                 results.append(Data.check_chore(session, chore.id))
  91         return results
  92
  93     def add(self, name):
  94         try:
  95             with Session(self.engine) as session:
  96                 new_chore = Chore(name=name)
  97                 session.add(new_chore)
  98                 session.commit()
  99             return f'Created chore {name}.'
 100         except:
 101             return f'Unable to create {name}. May be duplicate.'
 102
 103
 104
 105     def done(self, chore_id):
 106         with Session(self.engine) as session:
 107             chore = get_chore(session, chore_id)
 108             if(not chore):
 109                 return f'Chore not found: {chore_id}'
 110
 111             try:
 112                 session.add(Record(chore_id=chore.id))
 113                 session.commit()
 114                 return f'Marked chore as done: {chore.name}.'
 115             except Exception as ex:
 116                 return f'Unable to mark chore as done: {ex}'
 117
 118
 119     def undo(self, chore_id):
 120         with Session(self.engine) as session:
 121             cmd = select(Chore).where(Chore.id == chore_id)
 122             chore = session.exec(cmd).first()
 123
 124             if(not chore):
 125                 return f'Not found: {chore_id}'
 126
 127             try:
 128                 cmd = select(Record).where(Record.chore_id == chore_id).where(Record.complete_date >= today())
 129                 results = session.exec(cmd).all()
 130                 for result in results:
 131                     session.delete(result)
 132                 session.commit()
 133                 return f'Marked chore as *not* done: {chore.name}.'
 134             except Exception as ex:
 135                 return f'Unable to mark chore as done: {ex}'
 136
 137     def remove(self, chore_id):
 138         try:
 139             with Session(self.engine) as session:
 140                 chore = get_chore(session, chore_id)
 141                 if(not chore):
 142                     return f'Chore not found: {chore_id}'
 143
 144                 sql = text('''
 145                 DELETE FROM Record
 146                 WHERE chore_id = :chore_id
 147                ''')
 148                 params = {}
 149                 params['chore_id'] = chore_id
 150                 results = session.execute(sql, params)
 151                 logger.debug(results)
 152                 sql = text('''
 153                 DELETE FROM Chore
 154                 WHERE id = :chore_id
 155                ''')
 156                 results = session.execute(sql, params)
 157                 logger.debug(results)
 158                 session.commit()
 159                 return f'Removed chore {chore_id}'
 160         except Exception as ex:
 161             return f'Unable to remove chore: {ex}'