github.com/openfga/openfga@v1.5.4-rc1/pkg/storage/storagewrappers/context.go (about)

     1  package storagewrappers
     2  
     3  import (
     4  	"context"
     5  
     6  	openfgav1 "github.com/openfga/api/proto/openfga/v1"
     7  	"go.opentelemetry.io/otel/trace"
     8  
     9  	"github.com/openfga/openfga/pkg/storage"
    10  )
    11  
    12  // ContextTracerWrapper is a wrapper for a datastore that introduces a new context to the underlying datastore methods.
    13  // Its purpose is to prevent the closure of the underlying database connection in case the original context is cancelled,
    14  // such as when a client cancels the context. This ensures that ongoing queries are allowed to complete even if the
    15  // original context is cancelled, helping to avoid unnecessary database connection churn.
    16  type ContextTracerWrapper struct {
    17  	storage.OpenFGADatastore
    18  }
    19  
    20  var _ storage.OpenFGADatastore = (*ContextTracerWrapper)(nil)
    21  
    22  // NewContextWrapper creates a new instance of [ContextTracerWrapper], wrapping the specified datastore. It is crucial
    23  // for [ContextTracerWrapper] to be the first wrapper around the datastore for traces to function correctly.
    24  func NewContextWrapper(inner storage.OpenFGADatastore) *ContextTracerWrapper {
    25  	return &ContextTracerWrapper{inner}
    26  }
    27  
    28  // queryContext generates a new context that is independent of the provided
    29  // context and its timeout with the exception of the trace context.
    30  func queryContext(ctx context.Context) context.Context {
    31  	span := trace.SpanFromContext(ctx)
    32  	return trace.ContextWithSpan(context.Background(), span)
    33  }
    34  
    35  // Close ensures proper cleanup and closure of resources associated with the OpenFGADatastore.
    36  func (c *ContextTracerWrapper) Close() {
    37  	c.OpenFGADatastore.Close()
    38  }
    39  
    40  // Read see [storage.RelationshipTupleReader.ReadUserTuple].
    41  func (c *ContextTracerWrapper) Read(ctx context.Context, store string, tupleKey *openfgav1.TupleKey) (storage.TupleIterator, error) {
    42  	queryCtx := queryContext(ctx)
    43  
    44  	return c.OpenFGADatastore.Read(queryCtx, store, tupleKey)
    45  }
    46  
    47  // ReadPage see [storage.RelationshipTupleReader.ReadPage].
    48  func (c *ContextTracerWrapper) ReadPage(ctx context.Context, store string, tupleKey *openfgav1.TupleKey, opts storage.PaginationOptions) ([]*openfgav1.Tuple, []byte, error) {
    49  	queryCtx := queryContext(ctx)
    50  
    51  	return c.OpenFGADatastore.ReadPage(queryCtx, store, tupleKey, opts)
    52  }
    53  
    54  // ReadUserTuple see [storage.RelationshipTupleReader].ReadUserTuple.
    55  func (c *ContextTracerWrapper) ReadUserTuple(ctx context.Context, store string, tupleKey *openfgav1.TupleKey) (*openfgav1.Tuple, error) {
    56  	queryCtx := queryContext(ctx)
    57  
    58  	return c.OpenFGADatastore.ReadUserTuple(queryCtx, store, tupleKey)
    59  }
    60  
    61  // ReadUsersetTuples see [storage.RelationshipTupleReader].ReadUsersetTuples.
    62  func (c *ContextTracerWrapper) ReadUsersetTuples(ctx context.Context, store string, filter storage.ReadUsersetTuplesFilter) (storage.TupleIterator, error) {
    63  	queryCtx := queryContext(ctx)
    64  
    65  	return c.OpenFGADatastore.ReadUsersetTuples(queryCtx, store, filter)
    66  }
    67  
    68  // ReadStartingWithUser see [storage.RelationshipTupleReader].ReadStartingWithUser.
    69  func (c *ContextTracerWrapper) ReadStartingWithUser(ctx context.Context, store string, opts storage.ReadStartingWithUserFilter) (storage.TupleIterator, error) {
    70  	queryCtx := queryContext(ctx)
    71  
    72  	return c.OpenFGADatastore.ReadStartingWithUser(queryCtx, store, opts)
    73  }