1 files changed, 39 insertions, 4 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()
+}