Added WithContext() and FromContext()

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
+}