FocusFlowAI / storage.py
avaliev's picture
Upload 4 files
c6ad652 verified
"""
Task Manager with SQLite backend for CRUD operations.
"""
import sqlite3
import json
from typing import List, Dict, Optional
from datetime import datetime
import os
class TaskManager:
"""Manages tasks with SQLite persistence."""
# Strict status enum
VALID_STATUSES = {"Todo", "In Progress", "Done"}
def __init__(self, db_path: str = "focusflow.db", use_memory: bool = False):
"""
Initialize the task manager.
Args:
db_path: Path to SQLite database file
use_memory: If True, use in-memory list instead of SQLite (for Demo/HF Spaces)
"""
self.db_path = db_path
self.use_memory = use_memory
self.memory_tasks = [] # List of dicts for in-memory storage
self.memory_counter = 0 # Auto-increment ID for in-memory
if not self.use_memory:
self._init_db()
else:
print("ℹ️ TaskManager initialized in IN-MEMORY mode (non-persistent)")
def _init_db(self):
"""Create the tasks table if it doesn't exist."""
try:
conn = sqlite3.connect(self.db_path)
cursor = conn.cursor()
cursor.execute("""
CREATE TABLE IF NOT EXISTS tasks (
id INTEGER PRIMARY KEY AUTOINCREMENT,
title TEXT NOT NULL,
description TEXT,
status TEXT DEFAULT 'Todo',
estimated_duration TEXT,
position INTEGER,
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
)
""")
conn.commit()
conn.close()
except Exception as e:
print(f"⚠️ Database initialization failed: {e}. Falling back to in-memory mode.")
self.use_memory = True
self.memory_tasks = []
self.memory_counter = 0
def add_task(self, title: str, description: str = "",
estimated_duration: str = "", status: str = "Todo") -> int:
"""Add a new task and return its ID."""
# Validate status
if status not in self.VALID_STATUSES:
status = "Todo"
if self.use_memory:
self.memory_counter += 1
# Calculate position
max_pos = 0
if self.memory_tasks:
max_pos = max(t.get('position', 0) for t in self.memory_tasks)
position = max_pos + 1
new_task = {
"id": self.memory_counter,
"title": title,
"description": description,
"status": status,
"estimated_duration": estimated_duration,
"position": position,
"created_at": datetime.now().isoformat(),
"updated_at": datetime.now().isoformat()
}
self.memory_tasks.append(new_task)
return self.memory_counter
conn = sqlite3.connect(self.db_path)
cursor = conn.cursor()
# Get max position
cursor.execute("SELECT MAX(position) FROM tasks")
result = cursor.fetchone()
max_pos = result[0] if result and result[0] is not None else 0
position = max_pos + 1
cursor.execute("""
INSERT INTO tasks (title, description, status, estimated_duration, position)
VALUES (?, ?, ?, ?, ?)
""", (title, description, status, estimated_duration, position))
task_id = cursor.lastrowid or 0
conn.commit()
conn.close()
return task_id
def get_all_tasks(self) -> List[Dict]:
"""Get all tasks ordered by position."""
if self.use_memory:
return sorted(self.memory_tasks, key=lambda x: x.get('position', 0))
conn = sqlite3.connect(self.db_path)
conn.row_factory = sqlite3.Row
cursor = conn.cursor()
cursor.execute("""
SELECT id, title, description, status, estimated_duration, position
FROM tasks ORDER BY position
""")
tasks = [dict(row) for row in cursor.fetchall()]
conn.close()
return tasks
def get_task(self, task_id: int) -> Optional[Dict]:
"""Get a specific task by ID."""
if self.use_memory:
for task in self.memory_tasks:
if task['id'] == task_id:
return task.copy()
return None
conn = sqlite3.connect(self.db_path)
conn.row_factory = sqlite3.Row
cursor = conn.cursor()
cursor.execute("""
SELECT id, title, description, status, estimated_duration, position
FROM tasks WHERE id = ?
""", (task_id,))
row = cursor.fetchone()
conn.close()
return dict(row) if row else None
def update_task(self, task_id: int, **kwargs):
"""Update a task's fields with validation."""
# Validate status if provided
if 'status' in kwargs and kwargs['status'] not in self.VALID_STATUSES:
raise ValueError(f"Invalid status. Must be one of: {', '.join(self.VALID_STATUSES)}")
allowed_fields = ['title', 'description', 'status', 'estimated_duration', 'position']
if self.use_memory:
for task in self.memory_tasks:
if task['id'] == task_id:
for key, value in kwargs.items():
if key in allowed_fields:
task[key] = value
task['updated_at'] = datetime.now().isoformat()
return
return
conn = sqlite3.connect(self.db_path)
cursor = conn.cursor()
updates = []
values = []
for key, value in kwargs.items():
if key in allowed_fields:
updates.append(f"{key} = ?")
values.append(value)
if updates:
values.append(task_id)
query = f"UPDATE tasks SET {', '.join(updates)}, updated_at = CURRENT_TIMESTAMP WHERE id = ?"
cursor.execute(query, values)
conn.commit()
conn.close()
def delete_task(self, task_id: int):
"""Delete a task by ID."""
if self.use_memory:
self.memory_tasks = [t for t in self.memory_tasks if t['id'] != task_id]
return
conn = sqlite3.connect(self.db_path)
cursor = conn.cursor()
cursor.execute("DELETE FROM tasks WHERE id = ?", (task_id,))
conn.commit()
conn.close()
def reorder_tasks(self, task_ids: List[int]):
"""Reorder tasks based on new order."""
if self.use_memory:
# Create a map for O(1) lookup
task_map = {t['id']: t for t in self.memory_tasks}
# Update positions based on the input list order
for position, task_id in enumerate(task_ids, start=1):
if task_id in task_map:
task_map[task_id]['position'] = position
return
conn = sqlite3.connect(self.db_path)
cursor = conn.cursor()
for position, task_id in enumerate(task_ids, start=1):
cursor.execute("UPDATE tasks SET position = ? WHERE id = ?", (position, task_id))
conn.commit()
conn.close()
def get_active_task(self) -> Optional[Dict]:
"""Get the task marked as 'In Progress'."""
if self.use_memory:
# Filter for In Progress tasks and sort by position
active_tasks = [t for t in self.memory_tasks if t['status'] == 'In Progress']
if not active_tasks:
return None
# Return the first one (lowest position)
return sorted(active_tasks, key=lambda x: x.get('position', 0))[0].copy()
conn = sqlite3.connect(self.db_path)
conn.row_factory = sqlite3.Row
cursor = conn.cursor()
cursor.execute("""
SELECT id, title, description, status, estimated_duration
FROM tasks WHERE status = 'In Progress'
ORDER BY position LIMIT 1
""")
row = cursor.fetchone()
conn.close()
return dict(row) if row else None
def set_active_task(self, task_id: int) -> bool:
"""Set a task as 'In Progress' and ensure only one task has this status.
Returns True if successful, False otherwise."""
if self.use_memory:
# Check if task exists
target_task = None
for t in self.memory_tasks:
if t['id'] == task_id:
target_task = t
break
if not target_task:
return False
if target_task['status'] == 'Done':
return False
# Reset other In Progress tasks
for t in self.memory_tasks:
if t['status'] == 'In Progress':
t['status'] = 'Todo'
t['updated_at'] = datetime.now().isoformat()
# Set target task
target_task['status'] = 'In Progress'
target_task['updated_at'] = datetime.now().isoformat()
return True
conn = sqlite3.connect(self.db_path)
cursor = conn.cursor()
# Check if task exists and is not already Done
cursor.execute("SELECT status FROM tasks WHERE id = ?", (task_id,))
result = cursor.fetchone()
if not result:
conn.close()
return False
current_status = result[0]
if current_status == "Done":
conn.close()
return False
# Enforce single "In Progress" rule: set all current "In Progress" tasks back to "Todo"
cursor.execute("""
UPDATE tasks SET status = 'Todo', updated_at = CURRENT_TIMESTAMP
WHERE status = 'In Progress'
""")
# Set the selected task as 'In Progress'
cursor.execute("""
UPDATE tasks SET status = 'In Progress', updated_at = CURRENT_TIMESTAMP
WHERE id = ?
""", (task_id,))
conn.commit()
conn.close()
return True
def clear_all_tasks(self):
"""Clear all tasks from the database."""
if self.use_memory:
self.memory_tasks = []
return
conn = sqlite3.connect(self.db_path)
cursor = conn.cursor()
cursor.execute("DELETE FROM tasks")
conn.commit()
conn.close()