refactor: setup database connection via prerun

cobble.go

@@ -1,21 +1,14 @@
 package main
 
 import (
+	"context"
 	"github.com/spf13/cobra"
 	"log"
 	"os"
 	"strconv"
 )
 
-var cli = &cobra.Command{
-	Use:   "cobble",
-	Short: "Cobble is a simple SQL migration utility.",
-}
-
-var up = &cobra.Command{
-	Use:   "up",
-	Short: "Applies migrations",
-	Run: func(cmd *cobra.Command, args []string) {
+func SetupDatabaseConnection(cmd *cobra.Command, args []string) {
 	dbPort, _ := strconv.Atoi(os.Getenv("DB_PORT"))
 
 	db := NewDatabase(os.Getenv("DB_HOST"), dbPort)
@@ -26,12 +19,26 @@ var up = &cobra.Command{
 		os.Exit(1)
 	}
 
+	cmd.SetContext(context.WithValue(cmd.Context(), "db", db))
+}
+
+var cli = &cobra.Command{
+	Use:   "cobble",
+	Short: "Cobble is a simple SQL migration utility.",
+}
+
+var up = &cobra.Command{
+	Use:    "up",
+	Short:  "Applies migrations",
+	PreRun: SetupDatabaseConnection,
+	Run: func(cmd *cobra.Command, args []string) {
+		db := cmd.Context().Value("db").(DB)
 		migrationRoot := args[0]
 		migrationGraph, _ := NewMigrationGraphFromDirectory(migrationRoot)
 		migrationHistory, _ := migrationGraph.GetLinearHistory()
 		for _, migration := range migrationHistory {
 			log.Printf("%s", migration.Name)
-			if err = migration.ApplyMigration(db); err != nil {
+			if err := migration.ApplyMigration(db); err != nil {
 				os.Exit(1)
 			}
 		}