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