From 4011f3cace18e4d87dc9d4f3f87dcab62a7ed2ef Mon Sep 17 00:00:00 2001
From: Jan Bader <jan@javil.eu>
Date: Thu, 2 Dec 2021 20:36:22 +0000
Subject: [PATCH] Also return sql.DB to be able to use goose

---
 cmd/budgeteer/main.go  |  4 ++--
 postgres/conn.go       | 15 +++++----------
 postgres/repository.go | 10 +++++++---
 3 files changed, 14 insertions(+), 15 deletions(-)

diff --git a/cmd/budgeteer/main.go b/cmd/budgeteer/main.go
index 7134a42..ca4129f 100644
--- a/cmd/budgeteer/main.go
+++ b/cmd/budgeteer/main.go
@@ -18,12 +18,12 @@ func main() {
 
 	bv := &bcrypt.Verifier{}
 
-	db, err := postgres.Connect(cfg.DatabaseHost, cfg.DatabaseUser, cfg.DatabasePassword, cfg.DatabaseName)
+	q, db, err := postgres.Connect(cfg.DatabaseHost, cfg.DatabaseUser, cfg.DatabasePassword, cfg.DatabaseName)
 	if err != nil {
 		log.Fatalf("Failed connecting to DB: %v", err)
 	}
 
-	us, err := postgres.NewRepository(db)
+	us, err := postgres.NewRepository(q, db)
 	if err != nil {
 		log.Fatalf("Failed building Repository: %v", err)
 	}
diff --git a/postgres/conn.go b/postgres/conn.go
index 84de83b..00d0c18 100644
--- a/postgres/conn.go
+++ b/postgres/conn.go
@@ -15,29 +15,24 @@ import (
 var migrations embed.FS
 
 // Connect to a database
-func Connect(server string, user string, password string, database string) (*Queries, error) {
+func Connect(server string, user string, password string, database string) (*Queries, *sql.DB, error) {
 	connString := fmt.Sprintf("postgres://%s:%s@%s/%s", user, password, server, database)
 	conn, err := sql.Open("pgx", connString)
 	if err != nil {
-		return nil, err
+		return nil, nil, err
 	}
 
 	goose.SetBaseFS(migrations)
 	if err = goose.Up(conn, "schema"); err != nil {
-		return nil, err
-	}
-
-	err = conn.Close()
-	if err != nil {
-		return nil, err
+		return nil, nil, err
 	}
 
 	connPG, err := pgx.Connect(context.Background(), connString)
 	if err != nil {
-		return nil, err
+		return nil, nil, err
 	}
 
-	return New(connPG), nil
+	return New(connPG), conn, nil
 }
 
 func (tx Transaction) GetAmount() float64 {
diff --git a/postgres/repository.go b/postgres/repository.go
index 2ea4048..ab33ab3 100644
--- a/postgres/repository.go
+++ b/postgres/repository.go
@@ -1,13 +1,17 @@
 package postgres
 
+import "database/sql"
+
 // Repository represents a PostgreSQL implementation of all ModelServices
 type Repository struct {
-	DB *Queries
+	DB       *Queries
+	LegacyDB *sql.DB
 }
 
-func NewRepository(queries *Queries) (*Repository, error) {
+func NewRepository(queries *Queries, db *sql.DB) (*Repository, error) {
 	repo := &Repository{
-		DB: queries,
+		DB:       queries,
+		LegacyDB: db,
 	}
 	return repo, nil
 }