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

     1  package datastore
     2  
     3  import (
     4  	"context"
     5  
     6  	"go.opentelemetry.io/otel/trace"
     7  
     8  	log "github.com/authzed/spicedb/internal/logging"
     9  	"github.com/authzed/spicedb/pkg/datastore"
    10  	"github.com/authzed/spicedb/pkg/datastore/options"
    11  	"github.com/authzed/spicedb/pkg/datastore/test"
    12  	core "github.com/authzed/spicedb/pkg/proto/core/v1"
    13  )
    14  
    15  // SeparateContextWithTracing is a utility method which allows for severing the
    16  // context between grpc and the datastore to prevent context cancellation from
    17  // killing database connections that should otherwise go back to the connection
    18  // pool.
    19  func SeparateContextWithTracing(ctx context.Context) context.Context {
    20  	span := trace.SpanFromContext(ctx)
    21  	ctxWithObservability := trace.ContextWithSpan(context.Background(), span)
    22  
    23  	loggerFromContext := log.Ctx(ctx)
    24  	if loggerFromContext != nil {
    25  		ctxWithObservability = loggerFromContext.WithContext(ctxWithObservability)
    26  	}
    27  
    28  	return ctxWithObservability
    29  }
    30  
    31  // NewSeparatingContextDatastoreProxy severs any timeouts in the context being
    32  // passed to the datastore and only retains tracing metadata.
    33  //
    34  // This is useful for datastores that do not want to close connections when a
    35  // cancel or deadline occurs.
    36  func NewSeparatingContextDatastoreProxy(d datastore.Datastore) datastore.Datastore {
    37  	return &ctxProxy{d}
    38  }
    39  
    40  type ctxProxy struct{ delegate datastore.Datastore }
    41  
    42  func (p *ctxProxy) ReadWriteTx(
    43  	ctx context.Context,
    44  	f datastore.TxUserFunc,
    45  	opts ...options.RWTOptionsOption,
    46  ) (datastore.Revision, error) {
    47  	return p.delegate.ReadWriteTx(ctx, f, opts...)
    48  }
    49  
    50  func (p *ctxProxy) OptimizedRevision(ctx context.Context) (datastore.Revision, error) {
    51  	return p.delegate.OptimizedRevision(SeparateContextWithTracing(ctx))
    52  }
    53  
    54  func (p *ctxProxy) CheckRevision(ctx context.Context, revision datastore.Revision) error {
    55  	return p.delegate.CheckRevision(SeparateContextWithTracing(ctx), revision)
    56  }
    57  
    58  func (p *ctxProxy) HeadRevision(ctx context.Context) (datastore.Revision, error) {
    59  	return p.delegate.HeadRevision(SeparateContextWithTracing(ctx))
    60  }
    61  
    62  func (p *ctxProxy) RevisionFromString(serialized string) (datastore.Revision, error) {
    63  	return p.delegate.RevisionFromString(serialized)
    64  }
    65  
    66  func (p *ctxProxy) Watch(ctx context.Context, afterRevision datastore.Revision, options datastore.WatchOptions) (<-chan *datastore.RevisionChanges, <-chan error) {
    67  	return p.delegate.Watch(ctx, afterRevision, options)
    68  }
    69  
    70  func (p *ctxProxy) Features(ctx context.Context) (*datastore.Features, error) {
    71  	return p.delegate.Features(SeparateContextWithTracing(ctx))
    72  }
    73  
    74  func (p *ctxProxy) Statistics(ctx context.Context) (datastore.Stats, error) {
    75  	return p.delegate.Statistics(SeparateContextWithTracing(ctx))
    76  }
    77  
    78  func (p *ctxProxy) ReadyState(ctx context.Context) (datastore.ReadyState, error) {
    79  	return p.delegate.ReadyState(SeparateContextWithTracing(ctx))
    80  }
    81  
    82  func (p *ctxProxy) Close() error { return p.delegate.Close() }
    83  
    84  func (p *ctxProxy) SnapshotReader(rev datastore.Revision) datastore.Reader {
    85  	delegateReader := p.delegate.SnapshotReader(rev)
    86  	return &ctxReader{delegateReader}
    87  }
    88  
    89  func (p *ctxProxy) Unwrap() datastore.Datastore {
    90  	return p.delegate
    91  }
    92  
    93  // Implement the TestableDatastore interface
    94  func (p *ctxProxy) ExampleRetryableError() error {
    95  	return p.delegate.(test.TestableDatastore).ExampleRetryableError()
    96  }
    97  
    98  type ctxReader struct{ delegate datastore.Reader }
    99  
   100  func (r *ctxReader) ReadCaveatByName(ctx context.Context, name string) (*core.CaveatDefinition, datastore.Revision, error) {
   101  	return r.delegate.ReadCaveatByName(SeparateContextWithTracing(ctx), name)
   102  }
   103  
   104  func (r *ctxReader) ListAllCaveats(ctx context.Context) ([]datastore.RevisionedCaveat, error) {
   105  	return r.delegate.ListAllCaveats(SeparateContextWithTracing(ctx))
   106  }
   107  
   108  func (r *ctxReader) LookupCaveatsWithNames(ctx context.Context, caveatNames []string) ([]datastore.RevisionedCaveat, error) {
   109  	return r.delegate.LookupCaveatsWithNames(SeparateContextWithTracing(ctx), caveatNames)
   110  }
   111  
   112  func (r *ctxReader) ListAllNamespaces(ctx context.Context) ([]datastore.RevisionedNamespace, error) {
   113  	return r.delegate.ListAllNamespaces(SeparateContextWithTracing(ctx))
   114  }
   115  
   116  func (r *ctxReader) LookupNamespacesWithNames(ctx context.Context, nsNames []string) ([]datastore.RevisionedNamespace, error) {
   117  	return r.delegate.LookupNamespacesWithNames(SeparateContextWithTracing(ctx), nsNames)
   118  }
   119  
   120  func (r *ctxReader) ReadNamespaceByName(ctx context.Context, nsName string) (*core.NamespaceDefinition, datastore.Revision, error) {
   121  	return r.delegate.ReadNamespaceByName(SeparateContextWithTracing(ctx), nsName)
   122  }
   123  
   124  func (r *ctxReader) QueryRelationships(ctx context.Context, filter datastore.RelationshipsFilter, options ...options.QueryOptionsOption) (datastore.RelationshipIterator, error) {
   125  	return r.delegate.QueryRelationships(SeparateContextWithTracing(ctx), filter, options...)
   126  }
   127  
   128  func (r *ctxReader) ReverseQueryRelationships(ctx context.Context, subjectsFilter datastore.SubjectsFilter, options ...options.ReverseQueryOptionsOption) (datastore.RelationshipIterator, error) {
   129  	return r.delegate.ReverseQueryRelationships(SeparateContextWithTracing(ctx), subjectsFilter, options...)
   130  }
   131  
   132  var (
   133  	_ datastore.Datastore = (*ctxProxy)(nil)
   134  	_ datastore.Reader    = (*ctxReader)(nil)
   135  )