From 1e4c50b70d6a45a46ca4b7f1a29f4e62c52c2cde Mon Sep 17 00:00:00 2001
From: Kay-Zee <kan@axiomzen.co>
Date: Thu, 6 Sep 2018 17:16:02 -0700
Subject: [PATCH] Add sql.DB to Postgres and Mysql structs to allow closing in
 Close()

---
 database/mysql/mysql.go       | 17 ++++++++++++-----
 database/postgres/postgres.go |  9 ++++++++-
 2 files changed, 20 insertions(+), 6 deletions(-)

diff --git a/database/mysql/mysql.go b/database/mysql/mysql.go
index 323edae..69ff063 100644
--- a/database/mysql/mysql.go
+++ b/database/mysql/mysql.go
@@ -31,10 +31,10 @@ func init() {
 var DefaultMigrationsTable = "schema_migrations"
 
 var (
-	ErrDatabaseDirty  = fmt.Errorf("database is dirty")
-	ErrNilConfig      = fmt.Errorf("no config")
-	ErrNoDatabaseName = fmt.Errorf("no database name")
-	ErrAppendPEM      = fmt.Errorf("failed to append PEM")
+	ErrDatabaseDirty    = fmt.Errorf("database is dirty")
+	ErrNilConfig        = fmt.Errorf("no config")
+	ErrNoDatabaseName   = fmt.Errorf("no database name")
+	ErrAppendPEM        = fmt.Errorf("failed to append PEM")
 	ErrTLSCertKeyConfig = fmt.Errorf("To use TLS client authentication, both x-tls-cert and x-tls-key must not be empty")
 )
 
@@ -47,6 +47,7 @@ type Mysql struct {
 	// mysql RELEASE_LOCK must be called from the same conn, so
 	// just do everything over a single conn anyway.
 	conn     *sql.Conn
+	db       *sql.DB
 	isLocked bool
 
 	config *Config
@@ -85,6 +86,7 @@ func WithInstance(instance *sql.DB, config *Config) (database.Driver, error) {
 
 	mx := &Mysql{
 		conn:   conn,
+		db:     instance,
 		config: config,
 	}
 
@@ -193,7 +195,12 @@ func (m *Mysql) Open(url string) (database.Driver, error) {
 }
 
 func (m *Mysql) Close() error {
-	return m.conn.Close()
+	connErr := m.conn.Close()
+	dbErr := m.db.Close()
+	if connErr != nil || dbErr != nil {
+		return fmt.Errorf("conn: %v, db: %v", connErr, dbErr)
+	}
+	return nil
 }
 
 func (m *Mysql) Lock() error {
diff --git a/database/postgres/postgres.go b/database/postgres/postgres.go
index 25e06b6..eac96ff 100644
--- a/database/postgres/postgres.go
+++ b/database/postgres/postgres.go
@@ -40,6 +40,7 @@ type Config struct {
 type Postgres struct {
 	// Locking and unlocking need to use the same connection
 	conn     *sql.Conn
+	db       *sql.DB
 	isLocked bool
 
 	// Open and WithInstance need to garantuee that config is never nil
@@ -79,6 +80,7 @@ func WithInstance(instance *sql.DB, config *Config) (database.Driver, error) {
 
 	px := &Postgres{
 		conn:   conn,
+		db:     instance,
 		config: config,
 	}
 
@@ -117,7 +119,12 @@ func (p *Postgres) Open(url string) (database.Driver, error) {
 }
 
 func (p *Postgres) Close() error {
-	return p.conn.Close()
+	connErr := p.conn.Close()
+	dbErr := p.db.Close()
+	if connErr != nil || dbErr != nil {
+		return fmt.Errorf("conn: %v, db: %v", connErr, dbErr)
+	}
+	return nil
 }
 
 // https://www.postgresql.org/docs/9.6/static/explicit-locking.html#ADVISORY-LOCKS