28 lines
763 B
Go
28 lines
763 B
Go
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
|
|
}
|