From 7896a453984e3a7fc271d23614d780f1021994c1 Mon Sep 17 00:00:00 2001
From: Sebastian Tobie <sebastian@sebastian-tobie.de>
Date: Sun, 28 Feb 2021 14:01:31 +0100
Subject: [PATCH] the contextkey now refers now to the function for getting the
 database.

---
 middleware/db/db.go | 24 ++++++++++++++++++------
 1 file changed, 18 insertions(+), 6 deletions(-)

diff --git a/middleware/db/db.go b/middleware/db/db.go
index 6c125b0..ff8b228 100644
--- a/middleware/db/db.go
+++ b/middleware/db/db.go
@@ -17,6 +17,10 @@ import (
 // ContextKey is the key that is used in a gin.Context to get the Middleware
 const ContextKey string = "db"
 
+// ConnGet is an function that returns an connection instance.
+type ConnGet func(string) *pgxpool.Conn
+
+var _ ConnGet = NewMiddleware().GetConn
 var _ middleware.Middleware = &Middleware{}
 
 // Middleware return a handler that sets the db into the context of every request.
@@ -68,14 +72,11 @@ func (m *Middleware) AddDB(name, uri string) (err error) {
 	return
 }
 
-// GetDB returns a pgx.Conn Object or nil if the name is not found.
-func (m *Middleware) GetDB(name string) *pgxpool.Conn {
-	m.lock.Lock()
-	defer m.lock.Unlock()
+func (m *Middleware) getconn(name string) *pgxpool.Conn {
 	if db, found := m.databases[name]; found {
 		conn, err := db.Acquire(context.TODO())
 		if err != nil {
-			log.Error().Err(err).Msg("Could not get a connection from the pool")
+			log.Error().Err(err).Msgf("Could not get the connection from the pool %s", name)
 			return nil
 		}
 		return conn
@@ -83,9 +84,20 @@ func (m *Middleware) GetDB(name string) *pgxpool.Conn {
 	return nil
 }
 
+// GetConn returns a connection from the specified database or if not found one of the default database.
+func (m *Middleware) GetConn(name string) *pgxpool.Conn {
+	m.lock.Lock()
+	defer m.lock.Unlock()
+	if conn := m.getconn(name); conn != nil {
+		return conn
+	}
+	conn := m.getconn("default")
+	return conn
+}
+
 // Gin is the Entrypoint for Gin.
 func (m *Middleware) Gin(c *gin.Context) {
-	c.Set(ContextKey, m)
+	c.Set(ContextKey, m.GetConn)
 }
 
 // Teardown closes the DBConnection