From a241b02ebfa3921834cb814678815ee15e5d0266 Mon Sep 17 00:00:00 2001
From: Ali <bincyber@users.noreply.github.com>
Date: Sat, 13 Mar 2021 17:33:47 -0500
Subject: [PATCH] Added WithContext() and FromContext()

---
 context.go | 27 +++++++++++++++++++++++++++
 1 file changed, 27 insertions(+)
 create mode 100644 context.go

diff --git a/context.go b/context.go
new file mode 100644
index 00000000..c7d76e05
--- /dev/null
+++ b/context.go
@@ -0,0 +1,27 @@
+package gorm
+
+import (
+	"context"
+)
+
+// contextKeyType is an unexported type so that the context key never
+// collides with any other context keys.
+type contextKeyType struct{}
+
+// contextKey is the key used for the context to store the DB object.
+var contextKey = contextKeyType{}
+
+// WithContext inserts a DB into the context and is retrievable using FromContext().
+func WithContext(ctx context.Context, db *DB) context.Context {
+	return context.WithValue(ctx, contextKey, db)
+}
+
+// FromContext extracts a DB from the context. An error is returned if
+// the context does not contain a DB object.
+func FromContext(ctx context.Context) (*DB, error) {
+	db, _ := ctx.Value(contextKey).(*DB)
+	if db == nil {
+		return nil, ErrDBNotFoundInContext
+	}
+	return db, nil
+}