tracker/database.py

88 lines
2.6 KiB
Python

import sqlite3
import datetime
import pandas as pd
import os
DB_FILE = os.environ.get("DB_FILE", "work_log.db")
# Ensure the directory for the database file exists
db_dir = os.path.dirname(DB_FILE)
if db_dir:
os.makedirs(db_dir, exist_ok=True)
def get_connection():
conn = sqlite3.connect(DB_FILE)
return conn
def init_db():
conn = get_connection()
c = conn.cursor()
# Create table with new schema
c.execute('''
CREATE TABLE IF NOT EXISTS work_days (
date TEXT PRIMARY KEY,
hours REAL DEFAULT 8.0,
note TEXT,
type TEXT DEFAULT 'work',
course_name TEXT,
travel_costs INTEGER DEFAULT 0
)
''')
# Migration: Check if columns exist, if not add them
c.execute("PRAGMA table_info(work_days)")
columns = [info[1] for info in c.fetchall()]
if 'type' not in columns:
c.execute("ALTER TABLE work_days ADD COLUMN type TEXT DEFAULT 'work'")
if 'course_name' not in columns:
c.execute("ALTER TABLE work_days ADD COLUMN course_name TEXT")
if 'travel_costs' not in columns:
c.execute("ALTER TABLE work_days ADD COLUMN travel_costs INTEGER DEFAULT 0")
conn.commit()
conn.close()
def update_day(date_str, day_type="work", course_name=None, hours=8.0, travel_costs=False):
"""Updates or inserts a day. If course_name is None, it won't overwrite existing unless it's a new entry."""
conn = get_connection()
c = conn.cursor()
travel_val = 1 if travel_costs else 0
# Check if exists
c.execute("SELECT date FROM work_days WHERE date = ?", (date_str,))
data = c.fetchone()
if data:
# Update existing
c.execute("""
UPDATE work_days
SET type = ?, course_name = ?, hours = ?, travel_costs = ?
WHERE date = ?
""", (day_type, course_name, hours, travel_val, date_str))
else:
# Insert new
c.execute("""
INSERT INTO work_days (date, hours, type, course_name, travel_costs)
VALUES (?, ?, ?, ?, ?)
""", (date_str, hours, day_type, course_name, travel_val))
conn.commit()
conn.close()
def remove_day(date_str):
"""Removes a day from the database."""
conn = get_connection()
c = conn.cursor()
c.execute("DELETE FROM work_days WHERE date = ?", (date_str,))
conn.commit()
conn.close()
def get_all_days():
"""Returns a list of all work days."""
conn = get_connection()
df = pd.read_sql_query("SELECT * FROM work_days", conn)
conn.close()
return df