From 26d9674bb8a2545ca1572ac594a1533af178fad0 Mon Sep 17 00:00:00 2001
From: Thanatat Tamtan <acoshift@gmail.com>
Date: Wed, 21 Dec 2022 22:19:06 +0700
Subject: [PATCH] pgctx: add GetDB and GetDBKey

---
 pgctx/pgctx.go | 10 +++++++++-
 1 file changed, 9 insertions(+), 1 deletion(-)

diff --git a/pgctx/pgctx.go b/pgctx/pgctx.go
index 880e6d9..bb83116 100644
--- a/pgctx/pgctx.go
+++ b/pgctx/pgctx.go
@@ -51,6 +51,14 @@ func With(ctx context.Context, key any) context.Context {
 	return context.WithValue(ctx, ctxKeyDB{}, db)
 }
 
+func GetDB(ctx context.Context) DB {
+	return ctx.Value(ctxKeyDB{}).(DB)
+}
+
+func GetDBKey(ctx context.Context, key any) DB {
+	return ctx.Value(ctxKeyDB{key}).(DB)
+}
+
 type wrapTx struct {
 	*sql.Tx
 	onCommitted []func(ctx context.Context)
@@ -114,7 +122,7 @@ func Committed(ctx context.Context, f func(ctx context.Context)) {
 }
 
 type (
-	ctxKeyDB      struct{
+	ctxKeyDB struct {
 		key any
 	}
 	ctxKeyQueryer struct{}