progress_store.py

from contextlib import closing

class ProgressStore:
    def _chunks(self, data, rows=10000):
        for i in range(0, len(data), rows):
            yield data[i:i+rows]
        
    def _execute(self, sql):
        with closing(self.connection.cursor()) as cursor:
            return [o for o in cursor.execute(sql)]
    
    def _commit(self):
        self.connection.commit()
        
    def __init__(self, db_location, bnumbers = None, table_name = 'progress'):
        self.db_location = db_location
        self.table_name = table_name
        
        self._bnumber_count = 0

        self.connection = sqlite3.connect(db_location)

        sql = (
            f"CREATE TABLE IF NOT EXISTS {self.table_name} "
            f"(bnumber text PRIMARY KEY, status text, notes text)"
        )
        
        self._execute(sql)
        self._commit()
        
        if (bnumbers):
            print(f"Initialising with reset!")
            self.reset(bnumbers)

            
    def count(self):
        sql = (
            f"SELECT COUNT(*) FROM {self.table_name}"
        )
        
        return self._execute(sql)[0][0]
    
    def reset(self, bnumbers):
        print(f"Resetting ProgressStore.")
        sql = f"DELETE FROM {self.table_name}"
        
        self._execute(sql)
        self._commit()  
        
        self._bnumber_count = len(bnumbers)
        
        chunks = self._chunks(bnumbers)
        
        print(f"Loading {len(bnumbers)} bnumbers in chunks.")

        for chunk in chunks:
            self._execute('BEGIN TRANSACTION')
            
            for bnumber in chunk:
                sql = (
                    f"INSERT INTO {self.table_name} "
                    f"(bnumber) VALUES ('{bnumber}')"
                )
                self.cursor.execute(sql)

            self._execute('COMMIT')

        print("Done resetting.")
        
    def close(self):
        self.connection.close()
    
    def get_all(self):
        sql = (
            f"SELECT * FROM {self.table_name}"
        )

        with closing(self.connection.cursor()) as cursor:
            result = cursor.execute(sql)
            next_one = result.fetchone()
        
            while next_one is not None:
                yield next_one
                next_one = result.fetchone()
    
    def get_status(self, status):
        sql = (
            f"SELECT * FROM {self.table_name} "
            f"WHERE status = '{status}'"
        )
        
        return self._execute(sql)
    
    def get(self, bnumber):
        sql = (
            f"SELECT * FROM {self.table_name} "
            f"WHERE bnumber = '{bnumber}'"
        )
        
        return self._execute(sql)
        
    def update(self, bnumber, status, notes = None, commit = True):        
        if(bnumber is None):
            raise Exception('Must include bnumber in status update!')
        
        if(status is None):
            raise Exception('Must include status in status update!')
        
        if(notes is None):
            notes = ""
        
        status = str(status).strip()
        
        sql = (
            f"UPDATE {self.table_name} SET "
            f"status = '{status}', notes = '{notes}' "
            f"WHERE bnumber = '{bnumber}'"
        )
                
        self._execute(sql)
        
        if commit:
            self._commit()
        
    def batch_update(self, status_updates):
        chunks = self._chunks(status_updates)
        
        print(f"Loading {len(status_updates)} status updates in chunks.")

        for chunk in chunks:
            self._execute('BEGIN TRANSACTION')
            
            for status_update in status_updates:
                self.update(
                    bnumber = status_update.get('bnumber'), 
                    status = status_update.get('status'), 
                    notes = status_update.get('notes'), 
                    commit = False
                )
                
            self._execute('COMMIT')
            
        print(f"Done loading status updates.")