Add budget-tracked toggles for projects/labels; fix SetLabelColor wiping the flag

4 files changed, 194 insertions, 8 deletions

diff --git a/internal/store/labels.go b/internal/store/labels.go
index 3c5f5eb..216b171 100644
--- a/internal/store/labels.go
+++ b/internal/store/labels.go
@@ -22,7 +22,7 @@ func (s *Store) SetTaskLabels(id string, labels []string) error {
 
 // GetLabelColors returns every label that has been assigned a color.
 func (s *Store) GetLabelColors() ([]models.LabelColor, error) {
-	rows, err := s.db.Query(`SELECT name, color FROM labels ORDER BY name ASC`)
+	rows, err := s.db.Query(`SELECT name, color, budget_tracked FROM labels ORDER BY name ASC`)
 	if err != nil {
 		return nil, err
 	}
@@ -31,7 +31,7 @@ func (s *Store) GetLabelColors() ([]models.LabelColor, error) {
 	var colors []models.LabelColor
 	for rows.Next() {
 		var c models.LabelColor
-		if err := rows.Scan(&c.Name, &c.Color); err != nil {
+		if err := rows.Scan(&c.Name, &c.Color, &c.BudgetTracked); err != nil {
 			return nil, err
 		}
 		colors = append(colors, c)
@@ -39,8 +39,43 @@ func (s *Store) GetLabelColors() ([]models.LabelColor, error) {
 	return colors, rows.Err()
 }
 
-// SetLabelColor assigns (or reassigns) a label's display color.
+// SetLabelColor assigns (or reassigns) a label's display color, preserving
+// any existing budget_tracked flag -- a plain INSERT OR REPLACE would
+// delete-and-reinsert the row, silently resetting budget_tracked to 0.
 func (s *Store) SetLabelColor(name, color string) error {
-	_, err := s.db.Exec(`INSERT OR REPLACE INTO labels (name, color) VALUES (?, ?)`, name, color)
+	_, err := s.db.Exec(`
+		INSERT INTO labels (name, color, budget_tracked) VALUES (?, ?, 0)
+		ON CONFLICT(name) DO UPDATE SET color = excluded.color
+	`, name, color)
 	return err
 }
+
+// SetLabelBudgetTracked marks a label as opted in (or out) of budget
+// tracking, preserving any existing color the same way SetLabelColor
+// preserves budget_tracked.
+func (s *Store) SetLabelBudgetTracked(name string, tracked bool) error {
+	_, err := s.db.Exec(`
+		INSERT INTO labels (name, color, budget_tracked) VALUES (?, '', ?)
+		ON CONFLICT(name) DO UPDATE SET budget_tracked = excluded.budget_tracked
+	`, name, tracked)
+	return err
+}
+
+// GetBudgetTrackedLabelNames returns the set of label names opted into budget tracking.
+func (s *Store) GetBudgetTrackedLabelNames() (map[string]bool, error) {
+	rows, err := s.db.Query(`SELECT name FROM labels WHERE budget_tracked = 1`)
+	if err != nil {
+		return nil, err
+	}
+	defer func() { _ = rows.Close() }()
+
+	names := make(map[string]bool)
+	for rows.Next() {
+		var name string
+		if err := rows.Scan(&name); err != nil {
+			return nil, err
+		}
+		names[name] = true
+	}
+	return names, rows.Err()
+}
diff --git a/internal/store/labels_test.go b/internal/store/labels_test.go
index 4e97084..79034c8 100644
--- a/internal/store/labels_test.go
+++ b/internal/store/labels_test.go
@@ -82,3 +82,63 @@ func TestSetLabelColor_SameName_Overwrites(t *testing.T) {
 		t.Fatalf("colors = %+v, want a single entry with the overwritten color", colors)
 	}
 }
+
+func TestSetLabelBudgetTracked_TogglesFlag(t *testing.T) {
+	s := newNativeTasksTestStore(t)
+
+	if err := s.SetLabelBudgetTracked("errands", true); err != nil {
+		t.Fatalf("SetLabelBudgetTracked: %v", err)
+	}
+	names, err := s.GetBudgetTrackedLabelNames()
+	if err != nil {
+		t.Fatal(err)
+	}
+	if !names["errands"] {
+		t.Error("expected 'errands' to be tracked")
+	}
+}
+
+func TestSetLabelColor_PreservesExistingBudgetTracked(t *testing.T) {
+	s := newNativeTasksTestStore(t)
+
+	if err := s.SetLabelBudgetTracked("errands", true); err != nil {
+		t.Fatal(err)
+	}
+	if err := s.SetLabelColor("errands", "#00ff00"); err != nil {
+		t.Fatalf("SetLabelColor: %v", err)
+	}
+
+	names, err := s.GetBudgetTrackedLabelNames()
+	if err != nil {
+		t.Fatal(err)
+	}
+	if !names["errands"] {
+		t.Error("expected budget_tracked to survive a later SetLabelColor call")
+	}
+	colors, err := s.GetLabelColors()
+	if err != nil {
+		t.Fatal(err)
+	}
+	if len(colors) != 1 || colors[0].Color != "#00ff00" {
+		t.Errorf("colors = %+v, want one entry with color #00ff00", colors)
+	}
+}
+
+func TestSetLabelBudgetTracked_PreservesExistingColor(t *testing.T) {
+	s := newNativeTasksTestStore(t)
+
+	if err := s.SetLabelColor("errands", "#00ff00"); err != nil {
+		t.Fatal(err)
+	}
+	if err := s.SetLabelBudgetTracked("errands", true); err != nil {
+		t.Fatal(err)
+	}
+
+	colors, err := s.GetLabelColors()
+	if err != nil {
+		t.Fatal(err)
+	}
+	if len(colors) != 1 || colors[0].Color != "#00ff00" {
+		t.Errorf("colors = %+v, want color to survive SetLabelBudgetTracked", colors)
+	}
+}
diff --git a/internal/store/projects.go b/internal/store/projects.go
index a05656f..a081529 100644
--- a/internal/store/projects.go
+++ b/internal/store/projects.go
@@ -18,7 +18,7 @@ func (s *Store) CreateProject(name, color string) (*models.Project, error) {
 // GetProjects returns all non-archived projects, alphabetically by name.
 func (s *Store) GetProjects() ([]models.Project, error) {
 	rows, err := s.db.Query(`
-		SELECT id, name, color, created_at, archived FROM projects WHERE archived = 0 ORDER BY name ASC
+		SELECT id, name, color, created_at, archived, budget_tracked FROM projects WHERE archived = 0 ORDER BY name ASC
 	`)
 	if err != nil {
 		return nil, err
@@ -28,7 +28,7 @@ func (s *Store) GetProjects() ([]models.Project, error) {
 	var projects []models.Project
 	for rows.Next() {
 		var p models.Project
-		if err := rows.Scan(&p.ID, &p.Name, &p.Color, &p.CreatedAt, &p.Archived); err != nil {
+		if err := rows.Scan(&p.ID, &p.Name, &p.Color, &p.CreatedAt, &p.Archived, &p.BudgetTracked); err != nil {
 			return nil, err
 		}
 		projects = append(projects, p)
@@ -40,8 +40,8 @@ func (s *Store) GetProjects() ([]models.Project, error) {
 func (s *Store) GetProjectByID(id string) (*models.Project, error) {
 	var p models.Project
 	err := s.db.QueryRow(`
-		SELECT id, name, color, created_at, archived FROM projects WHERE id = ?
-	`, id).Scan(&p.ID, &p.Name, &p.Color, &p.CreatedAt, &p.Archived)
+		SELECT id, name, color, created_at, archived, budget_tracked FROM projects WHERE id = ?
+	`, id).Scan(&p.ID, &p.Name, &p.Color, &p.CreatedAt, &p.Archived, &p.BudgetTracked)
 	if err == sql.ErrNoRows {
 		return nil, ErrNativeTaskNotFound
 	}
@@ -65,3 +65,32 @@ func (s *Store) SetTaskProject(id, projectID string) error {
 	}
 	return checkRowsAffected(result)
 }
+
+// SetProjectBudgetTracked marks a project as opted in (or out) of budget
+// tracking. Returns ErrNativeTaskNotFound if id doesn't match any row.
+func (s *Store) SetProjectBudgetTracked(id string, tracked bool) error {
+	result, err := s.db.Exec(`UPDATE projects SET budget_tracked = ? WHERE id = ?`, tracked, id)
+	if err != nil {
+		return err
+	}
+	return checkRowsAffected(result)
+}
+
+// GetBudgetTrackedProjectIDs returns the set of project IDs opted into budget tracking.
+func (s *Store) GetBudgetTrackedProjectIDs() (map[string]bool, error) {
+	rows, err := s.db.Query(`SELECT id FROM projects WHERE budget_tracked = 1`)
+	if err != nil {
+		return nil, err
+	}
+	defer func() { _ = rows.Close() }()
+
+	ids := make(map[string]bool)
+	for rows.Next() {
+		var id string
+		if err := rows.Scan(&id); err != nil {
+			return nil, err
+		}
+		ids[id] = true
+	}
+	return ids, rows.Err()
+}
diff --git a/internal/store/projects_test.go b/internal/store/projects_test.go
index beb1aa9..b39384a 100644
--- a/internal/store/projects_test.go
+++ b/internal/store/projects_test.go
@@ -84,3 +84,65 @@ func TestSetTaskProject_UnknownID_ReturnsErrNotFound(t *testing.T) {
 		t.Fatalf("expected ErrNativeTaskNotFound, got %v", err)
 	}
 }
+
+func TestSetProjectBudgetTracked_TogglesFlag(t *testing.T) {
+	s := newNativeTasksTestStore(t)
+	project, err := s.CreateProject("Sailing prep", "#3B82F6")
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if err := s.SetProjectBudgetTracked(project.ID, true); err != nil {
+		t.Fatalf("SetProjectBudgetTracked: %v", err)
+	}
+	got, err := s.GetProjectByID(project.ID)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if !got.BudgetTracked {
+		t.Error("expected BudgetTracked = true")
+	}
+
+	if err := s.SetProjectBudgetTracked(project.ID, false); err != nil {
+		t.Fatal(err)
+	}
+	got, err = s.GetProjectByID(project.ID)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if got.BudgetTracked {
+		t.Error("expected BudgetTracked = false after untoggling")
+	}
+}
+
+func TestSetProjectBudgetTracked_UnknownID_ReturnsErrNotFound(t *testing.T) {
+	s := newNativeTasksTestStore(t)
+	if err := s.SetProjectBudgetTracked("does-not-exist", true); !errors.Is(err, ErrNativeTaskNotFound) {
+		t.Errorf("err = %v, want ErrNativeTaskNotFound", err)
+	}
+}
+
+func TestGetBudgetTrackedProjectIDs_ReturnsOnlyTracked(t *testing.T) {
+	s := newNativeTasksTestStore(t)
+	tracked, err := s.CreateProject("Tracked", "#111111")
+	if err != nil {
+		t.Fatal(err)
+	}
+	if _, err := s.CreateProject("Untracked", "#222222"); err != nil {
+		t.Fatal(err)
+	}
+	if err := s.SetProjectBudgetTracked(tracked.ID, true); err != nil {
+		t.Fatal(err)
+	}
+
+	ids, err := s.GetBudgetTrackedProjectIDs()
+	if err != nil {
+		t.Fatalf("GetBudgetTrackedProjectIDs: %v", err)
+	}
+	if !ids[tracked.ID] {
+		t.Error("expected tracked project id present")
+	}
+	if len(ids) != 1 {
+		t.Errorf("len(ids) = %d, want 1", len(ids))
+	}
+}