From 4c85d77d0eecb14cbf5e571161ec2566df97146e Mon Sep 17 00:00:00 2001
From: violog <51th.apprent1ce.f0rce@gmail.com>
Date: Thu, 22 Feb 2024 15:35:06 +0200
Subject: [PATCH] Fix DB session cloning on each new request

---
 internal/service/handlers/middleware.go | 28 +++++++++++++++++++++++++
 internal/service/router.go              |  6 +-----
 2 files changed, 29 insertions(+), 5 deletions(-)

diff --git a/internal/service/handlers/middleware.go b/internal/service/handlers/middleware.go
index f87b5aa..b13c721 100644
--- a/internal/service/handlers/middleware.go
+++ b/internal/service/handlers/middleware.go
@@ -1,11 +1,14 @@
 package handlers
 
 import (
+	"context"
 	"net/http"
 
 	"github.com/rarimo/auth-svc/pkg/auth"
+	"github.com/rarimo/rarime-points-svc/internal/data/pg"
 	"gitlab.com/distributed_lab/ape"
 	"gitlab.com/distributed_lab/ape/problems"
+	"gitlab.com/distributed_lab/kit/pgdb"
 	"gitlab.com/distributed_lab/logan/v3"
 )
 
@@ -29,3 +32,28 @@ func AuthMiddleware(auth *auth.Client, log *logan.Entry) func(http.Handler) http
 		})
 	}
 }
+
+type ctxExtender func(context.Context) context.Context
+
+// DBCloneMiddleware is designed to clone DB session on each request. You must
+// put all new DB handlers here instead of ape.CtxMiddleware, unless you have a
+// reason to do otherwise.
+func DBCloneMiddleware(db *pgdb.DB) func(http.Handler) http.Handler {
+	return func(next http.Handler) http.Handler {
+		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			clone := db.Clone()
+			ctx := r.Context()
+
+			extenders := []ctxExtender{
+				CtxEventsQ(pg.NewEvents(clone)),
+				CtxBalancesQ(pg.NewBalances(clone)),
+				CtxWithdrawalsQ(pg.NewWithdrawals(clone)),
+			}
+
+			for _, extender := range extenders {
+				ctx = extender(ctx)
+			}
+			next.ServeHTTP(w, r.WithContext(ctx))
+		})
+	}
+}
diff --git a/internal/service/router.go b/internal/service/router.go
index 8c3a8bc..2e47bfc 100644
--- a/internal/service/router.go
+++ b/internal/service/router.go
@@ -5,27 +5,23 @@ import (
 
 	"github.com/go-chi/chi"
 	"github.com/rarimo/rarime-points-svc/internal/config"
-	"github.com/rarimo/rarime-points-svc/internal/data/pg"
 	"github.com/rarimo/rarime-points-svc/internal/service/handlers"
 	"gitlab.com/distributed_lab/ape"
 )
 
 func Run(ctx context.Context, cfg config.Config) {
 	r := chi.NewRouter()
-	db := cfg.DB().Clone()
 
 	r.Use(
 		ape.RecoverMiddleware(cfg.Log()),
 		ape.LoganMiddleware(cfg.Log()),
 		ape.CtxMiddleware(
 			handlers.CtxLog(cfg.Log()),
-			handlers.CtxEventsQ(pg.NewEvents(db)),
-			handlers.CtxBalancesQ(pg.NewBalances(db)),
-			handlers.CtxWithdrawalsQ(pg.NewWithdrawals(db)),
 			handlers.CtxEventTypes(cfg.EventTypes()),
 			handlers.CtxBroadcaster(cfg.Broadcaster()),
 			handlers.CtxPointPrice(cfg.PointPrice()),
 		),
+		handlers.DBCloneMiddleware(cfg.DB()),
 	)
 	r.Route("/integrations/rarime-points-svc/v1", func(r chi.Router) {
 		r.Route("/balances", func(r chi.Router) {