github.com/authzed/spicedb@v1.32.1-0.20240520085336-ebda56537386/internal/middleware/datastore/datastore.go (about)

     1  package datastore
     2  
     3  import (
     4  	"context"
     5  
     6  	middleware "github.com/grpc-ecosystem/go-grpc-middleware/v2"
     7  	"google.golang.org/grpc"
     8  
     9  	"github.com/authzed/spicedb/pkg/datastore"
    10  )
    11  
    12  type ctxKeyType struct{}
    13  
    14  var datastoreKey ctxKeyType = struct{}{}
    15  
    16  type datastoreHandle struct {
    17  	datastore datastore.Datastore
    18  }
    19  
    20  // ContextWithHandle adds a placeholder to a context that will later be
    21  // filled by the datastore
    22  func ContextWithHandle(ctx context.Context) context.Context {
    23  	return context.WithValue(ctx, datastoreKey, &datastoreHandle{})
    24  }
    25  
    26  // FromContext reads the selected datastore out of a context.Context
    27  // and returns nil if it does not exist.
    28  func FromContext(ctx context.Context) datastore.Datastore {
    29  	if c := ctx.Value(datastoreKey); c != nil {
    30  		handle := c.(*datastoreHandle)
    31  		return handle.datastore
    32  	}
    33  	return nil
    34  }
    35  
    36  // MustFromContext reads the selected datastore out of a context.Context and panics if it does not exist
    37  func MustFromContext(ctx context.Context) datastore.Datastore {
    38  	datastore := FromContext(ctx)
    39  	if datastore == nil {
    40  		panic("datastore middleware did not inject datastore")
    41  	}
    42  
    43  	return datastore
    44  }
    45  
    46  // SetInContext adds a datastore to the given context
    47  func SetInContext(ctx context.Context, datastore datastore.Datastore) error {
    48  	handle := ctx.Value(datastoreKey)
    49  	if handle == nil {
    50  		return nil
    51  	}
    52  	handle.(*datastoreHandle).datastore = datastore
    53  	return nil
    54  }
    55  
    56  // ContextWithDatastore adds the handle and datastore in one step
    57  func ContextWithDatastore(ctx context.Context, datastore datastore.Datastore) context.Context {
    58  	return context.WithValue(ctx, datastoreKey, &datastoreHandle{datastore: datastore})
    59  }
    60  
    61  // UnaryServerInterceptor returns a new unary server interceptor that adds the
    62  // datastore to the context
    63  func UnaryServerInterceptor(datastore datastore.Datastore) grpc.UnaryServerInterceptor {
    64  	return func(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (interface{}, error) {
    65  		newCtx := ContextWithHandle(ctx)
    66  		if err := SetInContext(newCtx, datastore); err != nil {
    67  			return nil, err
    68  		}
    69  
    70  		return handler(newCtx, req)
    71  	}
    72  }
    73  
    74  // StreamServerInterceptor returns a new stream server interceptor that adds the
    75  // datastore to the context
    76  func StreamServerInterceptor(datastore datastore.Datastore) grpc.StreamServerInterceptor {
    77  	return func(srv interface{}, stream grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
    78  		wrapped := middleware.WrapServerStream(stream)
    79  		wrapped.WrappedContext = ContextWithHandle(wrapped.WrappedContext)
    80  		if err := SetInContext(wrapped.WrappedContext, datastore); err != nil {
    81  			return err
    82  		}
    83  		return handler(srv, wrapped)
    84  	}
    85  }