Improve handling of context

postgres/budgetservice.go

@@ -6,32 +6,15 @@ import (
 	"github.com/google/uuid"
 )
 
-// Budget returns a budget for a given id.
-func (s *Repository) Budget(id uuid.UUID) (*Budget, error) {
-	budget, err := s.DB.GetBudget(context.Background(), id)
-	if err != nil {
-		return nil, err
-	}
-	return &budget, nil
-}
-
-func (s *Repository) BudgetsForUser(id uuid.UUID) ([]Budget, error) {
-	budgets, err := s.DB.GetBudgetsForUser(context.Background(), id)
-	if err != nil {
-		return nil, err
-	}
-	return budgets, nil
-
-}
-
-func (s *Repository) NewBudget(name string, userID uuid.UUID) (*Budget, error) {
-	budget, err := s.DB.CreateBudget(context.Background(), name)
+// NewBudget creates a budget and adds it to the current user
+func (s *Repository) NewBudget(context context.Context, name string, userID uuid.UUID) (*Budget, error) {
+	budget, err := s.DB.CreateBudget(context, name)
 	if err != nil {
 		return nil, err
 	}
 
 	ub := LinkBudgetToUserParams{UserID: userID, BudgetID: budget.ID}
-	_, err = s.DB.LinkBudgetToUser(context.Background(), ub)
+	_, err = s.DB.LinkBudgetToUser(context, ub)
 	if err != nil {
 		return nil, err
 	}

postgres/ynab-import.go

@@ -22,29 +22,29 @@ type YNABImport struct {
 	budgetID       uuid.UUID
 }
 
-func NewYNABImport(q *Queries, budgetID uuid.UUID) (*YNABImport, error) {
-	accounts, err := q.GetAccounts(context.Background(), budgetID)
+func NewYNABImport(context context.Context, q *Queries, budgetID uuid.UUID) (*YNABImport, error) {
+	accounts, err := q.GetAccounts(context, budgetID)
 	if err != nil {
 		return nil, err
 	}
 
-	payees, err := q.GetPayees(context.Background(), budgetID)
+	payees, err := q.GetPayees(context, budgetID)
 	if err != nil {
 		return nil, err
 	}
 
-	categories, err := q.GetCategories(context.Background(), budgetID)
+	categories, err := q.GetCategories(context, budgetID)
 	if err != nil {
 		return nil, err
 	}
 
-	categoryGroups, err := q.GetCategoryGroups(context.Background(), budgetID)
+	categoryGroups, err := q.GetCategoryGroups(context, budgetID)
 	if err != nil {
 		return nil, err
 	}
 
 	return &YNABImport{
-		Context:        context.Background(),
+		Context:        context,
 		accounts:       accounts,
 		payees:         payees,
 		categories:     categories,