dwtt/DataStore.py

import sqlite3
import os
from datetime import datetime

HISTORY_DT_FMT = '%Y-%m-%d %H:%M:%S'

class Task:
    def __init__(self, taskid, name, longname, parentid):
        self.taskid = taskid
        self.name = name
        self.longname = longname
        self.parentid = parentid

class TaskRef:
    def __init__(self, taskid, time):
        self.taskid = taskid
        self.time = time

class DataStore:
    def __init__(self, dbfile):
        if not os.path.exists(dbfile):
            self.createdb(dbfile)
        self.conn = sqlite3.connect(dbfile)

    def __del__(self):
        self.conn.close()

    def createdb(self, dbfile):
        conn = sqlite3.connect(dbfile)
        c = conn.cursor()
        c.execute('''
CREATE TABLE tasks (
    id INTEGER PRIMARY KEY,
    name TEXT,
    longname TEXT,
    parentid INTEGER,
    FOREIGN KEY (parentid) REFERENCES tasks(id)
)''')
        c.execute('''
CREATE TABLE entries (
    date TEXT,
    taskid INTEGER,
    seconds INTEGER,
    PRIMARY KEY (date, taskid),
    FOREIGN KEY (taskid) REFERENCES tasks(id)
)''')
        c.execute('''
CREATE TABLE history (
    id INTEGER PRIMARY KEY,
    taskid INTEGER,
    datetime TEXT,
    FOREIGN KEY (taskid) REFERENCES tasks(id)
)''')
        conn.commit()
        c.close()
        conn.close()

    def getCurrentTask(self):
        c = self.conn.cursor()
        c.execute('''
SELECT taskid, datetime
FROM history
WHERE id = 0
''')
        ct = None
        for row in c:
            taskid = row[0]
            dt = datetime.strptime(row[1], HISTORY_DT_FMT)
            ct = TaskRef(taskid, dt)
        c.close()
        self.conn.commit()
        return ct

    def updateCurrentTask(self, ct):
        self.clearCurrentTask()
        c = self.conn.cursor()
        c.execute('''
INSERT INTO history
VALUES (0, ?, ?)
''', (ct.taskid, ct.time.strftime(HISTORY_DT_FMT)))
        c.close()
        self.conn.commit()

    def clearCurrentTask(self, ct):
        c = self.conn.cursor()
        c.execute('''
DELETE FROM history
WHERE id = 0
''')
        c.close()
        self.conn.commit()

    def updateTask(self, taskid, name, longname):
        c = self.conn.cursor()
        c.execute('''
UPDATE tasks
SET name = ?, longname = ?
WHERE id = ?
''', (name, longname, taskid))
        c.close()
        self.conn.commit()
        return True

    def getTaskByID(self, taskid):
        t = None
        c = self.conn.cursor()
        c.execute('''
SELECT name, longname, parentid
FROM tasks
WHERE id = ?
''', (taskid,))
        for row in c:
            t = Task(taskid, row[0], row[1], row[2])
        c.close()
        self.conn.commit()
        return t

    def getTaskByNameParent(self, name, parentid):
        t = None
        c = self.conn.cursor()
        name = name.strip()
        if parentid is None:
            c.execute('''
SELECT id, name, longname, parentid
FROM tasks
WHERE name = ? AND parentid IS NULL
''', (name,))
        else:
            c.execute('''
SELECT id, name, longname, parentid
FROM tasks
WHERE name = ? AND parentid = ?
''', (name, parentid))
        for row in c:
            t = Task(*row)
        c.close()
        self.conn.commit()
        return t

    def getTaskByName(self, name):
        parts = name.split(':')
        parentid = None
        task = None
        for p in parts:
            p = p.strip()
            task = self.getTaskByNameParent(p, parentid)
            if task is None:
                break
            parentid = task.taskid
        return task

    def getParentTaskByName(self, name):
        parts = name.split(':')
        parentname = ':'.join(parts[:-1])
        return self.getTaskByName(parentname)

    def createTask(self, name, longname, parentid):
        c = self.conn.cursor()
        if parentid is not None and parentid != '':
            c.execute('''
SELECT *
FROM tasks
WHERE id = ?
''', (parentid,))
            found = False
            for row in c:
                found = True
            if not found:
                return 0
        c.execute('''
SELECT MAX(id)
FROM tasks
''')
        nextid = 1
        for row in c:
            if row[0] is not None:
                nextid = row[0] + 1
        c.execute('''
INSERT INTO tasks
VALUES (?, ?, ?, ?)
''', (nextid, name, longname, parentid))
        c.close()
        self.conn.commit()
        return nextid

    def addTime(self, date, taskid, seconds):
        c = self.conn.cursor()
        exists = False
        oldseconds = 0
        c.execute('''
SELECT seconds
FROM entries
WHERE date = ? AND taskid = ?
''', (date, taskid))
        for row in c:
            if row[0] is not None:
                exists = True
                oldseconds = row[0]
        if exists:
            c.execute('''
UPDATE entries
SET seconds = ?
WHERE date = ? AND taskid = ?
''', (oldseconds + seconds, date, taskid))
        else:
            c.execute('''
INSERT INTO entries
VALUES(?, ?, ?)
''', (date, taskid, seconds))
        c.close()
        self.conn.commit()