refactor: allow retrieving DB instance from context

Author: Vilsol

db/postgres/postgres.go

@@ -119,6 +119,11 @@ func DeleteForced(ctx context.Context, object interface{}) {
 
 func DBCtx(ctx context.Context) *gorm.DB {
 	if ctx != nil {
+		dbCtx := DBFromContext(ctx)
+		if dbCtx != nil {
+			return dbCtx
+		}
+
 		return db.WithContext(ctx)
 	}
 

db/postgres/utils.go

@@ -0,0 +1,24 @@
+package postgres
+
+import (
+	"context"
+	"gorm.io/gorm"
+)
+
+type (
+	ContextDB struct{}
+)
+
+func ContextWithDB(ctx context.Context, db *gorm.DB) context.Context {
+	return context.WithValue(ctx, ContextDB{}, db)
+}
+
+func DBFromContext(ctx context.Context) *gorm.DB {
+	value := ctx.Value(ContextDB{})
+
+	if value == nil {
+		return nil
+	}
+
+	return value.(*gorm.DB)
+}

db/statistics.go

@@ -38,16 +38,16 @@ func RunAsyncStatisticLoop(ctx context.Context) {
 				}
 			}
 
-			updateTx := postgres.DBCtx(ctx).Begin()
-
 			for entityType, entityValue := range resultMap {
 				for action, actionValue := range entityValue {
 					for entityID, count := range actionValue {
+						updateTx := postgres.DBCtx(ctx).Begin()
+						ctxWithTx := postgres.ContextWithDB(ctx, updateTx)
 						switch entityType {
 						case "mod":
 							switch action {
 							case "view":
-								mod := postgres.GetModByID(ctx, entityID)
+								mod := postgres.GetModByID(ctxWithTx, entityID)
 								if mod != nil {
 									currentHotness := mod.Hotness
 									if currentHotness > 4 {
@@ -60,7 +60,7 @@ func RunAsyncStatisticLoop(ctx context.Context) {
 						case "version":
 							switch action {
 							case "download":
-								version := postgres.GetVersion(ctx, entityID)
+								version := postgres.GetVersion(ctxWithTx, entityID)
 								if version != nil {
 									currentHotness := version.Hotness
 									if currentHotness > 4 {
@@ -71,13 +71,11 @@ func RunAsyncStatisticLoop(ctx context.Context) {
 								}
 							}
 						}
+						updateTx.Commit()
 					}
 				}
 			}
 
-			updateTx.Commit()
-			updateTx = postgres.DBCtx(ctx).Begin()
-
 			type Result struct {
 				ModID     string
 				Hotness   uint
@@ -89,7 +87,9 @@ func RunAsyncStatisticLoop(ctx context.Context) {
 			postgres.DBCtx(ctx).Raw("SELECT mod_id, SUM(hotness) AS hotness, SUM(downloads) AS downloads FROM versions GROUP BY mod_id").Scan(&resultRows)
 
 			for _, row := range resultRows {
-				mod := postgres.GetModByID(ctx, row.ModID)
+				updateTx := postgres.DBCtx(ctx).Begin()
+				ctxWithTx := postgres.ContextWithDB(ctx, updateTx)
+				mod := postgres.GetModByID(ctxWithTx, row.ModID)
 				if mod != nil {
 					currentPopularity := mod.Popularity
 					if currentPopularity > 4 {
@@ -101,10 +101,9 @@ func RunAsyncStatisticLoop(ctx context.Context) {
 						Downloads:  row.Downloads,
 					})
 				}
+				updateTx.Commit()
 			}
 
-			updateTx.Commit()
-
 			log.Ctx(ctx).Info().Msgf("Statistics Updated! Took %s", time.Since(start).String())
 			time.Sleep(time.Minute)
 		}