chore: fix regression from #2912 - tracing in SchemaText() (#2947)

Author: miparnisari

CHANGELOG.md

@@ -7,6 +7,9 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/).
 ### Changed
 - Updated CI so that Postgres tests run against v18 which is GA and not against v13 which is EOL (https://github.com/authzed/spicedb/pull/2926)
 
+### Fixed
+- Regression introduced in 1.49.2: missing spans in ReadSchema calls (https://github.com/authzed/spicedb/pull/2947)
+
 ## [1.49.2] - 2026-03-02
 ### Added
 - feat(query planner): add recursive direction strategies, and fix IS BFS (https://github.com/authzed/spicedb/pull/2891)

internal/middleware/memoryprotection/memory_protection.go

@@ -9,13 +9,16 @@ import (
 	middleware "github.com/grpc-ecosystem/go-grpc-middleware/v2"
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/prometheus/client_golang/prometheus/promauto"
+	"go.opentelemetry.io/otel"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
 
 	log "github.com/authzed/spicedb/internal/logging"
 )
 
+var tracer = otel.Tracer("spicedb/internal/middleware/memory_protection")
+
 // RequestsProcessed tracks requests that were processed by this middleware.
 var RequestsProcessed = promauto.NewCounterVec(prometheus.CounterOpts{
 	Namespace: "spicedb",
@@ -44,7 +47,7 @@ func New(usageProvider MemoryUsageProvider, name string) *MemoryProtectionMiddle
 // UnaryServerInterceptor returns a unary server interceptor that rejects incoming requests is memory usage is too high
 func (am *MemoryProtectionMiddleware) UnaryServerInterceptor() grpc.UnaryServerInterceptor {
 	return func(ctx context.Context, req any, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (any, error) {
-		if err := am.checkAdmission(info.FullMethod); err != nil {
+		if err := am.checkAdmission(ctx, info.FullMethod); err != nil {
 			return nil, err
 		}
 
@@ -55,7 +58,7 @@ func (am *MemoryProtectionMiddleware) UnaryServerInterceptor() grpc.UnaryServerI
 // StreamServerInterceptor returns a stream server interceptor that rejects incoming requests is memory usage is too high
 func (am *MemoryProtectionMiddleware) StreamServerInterceptor() grpc.StreamServerInterceptor {
 	return func(srv any, stream grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
-		if err := am.checkAdmission(info.FullMethod); err != nil {
+		if err := am.checkAdmission(stream.Context(), info.FullMethod); err != nil {
 			return err
 		}
 
@@ -65,7 +68,10 @@ func (am *MemoryProtectionMiddleware) StreamServerInterceptor() grpc.StreamServe
 }
 
 // checkAdmission returns an error if the request should be denied because memory usage is too high.
-func (am *MemoryProtectionMiddleware) checkAdmission(method string) error {
+func (am *MemoryProtectionMiddleware) checkAdmission(ctx context.Context, method string) error {
+	_, span := tracer.Start(ctx, "checkMemoryUsage")
+	defer span.End()
+
 	accept := true
 	defer func() {
 		am.recordMetric(method, accept)

internal/middleware/memoryprotection/memory_protection_test.go

@@ -48,7 +48,7 @@ func TestNew(t *testing.T) {
 			am := New(tt.inputProvider, "name")
 			require.NotNil(t, am)
 
-			err := am.checkAdmission("some_method")
+			err := am.checkAdmission(t.Context(), "some_method")
 			if tt.expectReqLetThrough {
 				require.NoError(t, err) // if the middleware is off, every request is let through
 			} else {

internal/middleware/perfinsights/perfinsights.go

@@ -11,6 +11,7 @@ import (
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/prometheus/client_golang/prometheus/promauto"
 	"github.com/rs/zerolog/log"
+	"go.opentelemetry.io/otel"
 	"go.opentelemetry.io/otel/trace"
 	"google.golang.org/grpc"
 
@@ -69,6 +70,8 @@ var APIShapeLatency = promauto.NewHistogramVec(prometheus.HistogramOpts{
 	NativeHistogramBucketFactor: 1.1,
 }, append([]string{"api_kind"}, allLabels...))
 
+var tracer = otel.Tracer("spicedb/internal/middleware")
+
 // ShapeBuilder is a function that returns a slice of strings representing the shape of the API call.
 // This is used to report the shape of the API call to Prometheus.
 type ShapeBuilder func() APIShapeLabels
@@ -79,6 +82,8 @@ func ObserveShapeLatency(ctx context.Context, methodName string, shape APIShapeL
 }
 
 func observeShapeLatency(ctx context.Context, metric *prometheus.HistogramVec, methodName string, shape APIShapeLabels, duration time.Duration) {
+	ctx, span := tracer.Start(ctx, "perfInsights.observeShapeLatency")
+	defer span.End()
 	labels := buildLabels(methodName, shape)
 	if len(labels) == 0 {
 		log.Warn().Str("method", methodName).

internal/middleware/usagemetrics/usagemetrics.go

@@ -9,6 +9,7 @@ import (
 	"github.com/grpc-ecosystem/go-grpc-middleware/v2/interceptors"
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/prometheus/client_golang/prometheus/promauto"
+	"go.opentelemetry.io/otel"
 	"google.golang.org/grpc"
 
 	"github.com/authzed/authzed-go/pkg/responsemeta"
@@ -34,6 +35,8 @@ var (
 		Help:      "Histogram of cluster dispatches performed by the instance.",
 		Buckets:   []float64{1, 5, 10, 25, 50, 100, 250},
 	}, DispatchedCountLabels)
+
+	tracer = otel.Tracer("spicedb/internal/middleware")
 )
 
 type reporter struct{}
@@ -52,6 +55,9 @@ type serverReporter struct {
 
 // PostCall is invoked after all PostMsgSend operations.
 func (r *serverReporter) PostCall(_ error, _ time.Duration) {
+	_, span := tracer.Start(r.ctx, "usagemetrics.PostCall")
+	defer span.End()
+
 	responseMeta := FromContext(r.ctx)
 	if responseMeta == nil {
 		responseMeta = &dispatch.ResponseMeta{}

internal/services/shared/schema_test.go

@@ -778,7 +778,7 @@ func TestApplySchemaChangesOverExisting(t *testing.T) {
 
 				sr, err := rwt.ReadSchema()
 				require.NoError(err)
-				schemaText, err := sr.SchemaText()
+				schemaText, err := sr.SchemaText(t.Context())
 				require.NoError(err)
 				require.Equal(tc.expectedSchema, schemaText)
 				return nil

internal/services/v1/schema.go

@@ -96,7 +96,7 @@ func (ss *schemaServer) ReadSchema(ctx context.Context, _ *v1.ReadSchemaRequest)
 		return nil, ss.rewriteError(ctx, err)
 	}
 
-	schemaText, err := sr.SchemaText()
+	schemaText, err := sr.SchemaText(ctx)
 	if err != nil {
 		return nil, ss.rewriteError(ctx, err)
 	}

pkg/datalayer/datalayer.go

@@ -51,7 +51,7 @@ type DataLayer interface {
 // SchemaReader groups schema read methods, accessed via RevisionedReader.ReadSchema().
 type SchemaReader interface {
 	// SchemaText returns the full schema text.
-	SchemaText() (string, error)
+	SchemaText(ctx context.Context) (string, error)
 
 	// LookupTypeDefByName looks up a type definition by name.
 	LookupTypeDefByName(ctx context.Context, name string) (datastore.RevisionedTypeDefinition, bool, error)

pkg/datalayer/schema_adapter.go

@@ -30,9 +30,7 @@ type legacySchemaReaderAdapter struct {
 
 // SchemaText returns the schema text at the current revision by reading all namespaces and caveats
 // and generating the schema text from them.
-func (l *legacySchemaReaderAdapter) SchemaText() (string, error) {
-	ctx := context.Background()
-
+func (l *legacySchemaReaderAdapter) SchemaText(ctx context.Context) (string, error) {
 	// Read all namespaces
 	namespaces, err := l.ListAllTypeDefinitions(ctx)
 	if err != nil {
@@ -61,7 +59,7 @@ func (l *legacySchemaReaderAdapter) SchemaText() (string, error) {
 	}
 
 	// Generate schema text with proper use directives
-	schemaText, _, err := generator.GenerateSchemaWithCaveatTypeSet(definitions, caveatTypeSet)
+	schemaText, _, err := generator.GenerateSchemaWithCaveatTypeSet(ctx, definitions, caveatTypeSet)
 	if err != nil {
 		return "", fmt.Errorf("failed to generate schema: %w", err)
 	}

pkg/schemadsl/generator/generator.go

@@ -2,12 +2,15 @@ package generator
 
 import (
 	"bufio"
+	"context"
 	"fmt"
 	"maps"
 	"slices"
 	"sort"
 	"strings"
 
+	"go.opentelemetry.io/otel"
+
 	"github.com/authzed/spicedb/pkg/caveats"
 	caveattypes "github.com/authzed/spicedb/pkg/caveats/types"
 	"github.com/authzed/spicedb/pkg/genutil/mapz"
@@ -18,18 +21,23 @@ import (
 	"github.com/authzed/spicedb/pkg/spiceerrors"
 )
 
+var tracer = otel.Tracer("spicedb/pkg/schemadsl/generator")
+
 // Ellipsis is the relation name for terminal subjects.
 const Ellipsis = "..."
 
 // MaxSingleLineCommentLength sets the maximum length for a comment to made single line.
 const MaxSingleLineCommentLength = 70 // 80 - the comment parts and some padding
 
 func GenerateSchema(definitions []compiler.SchemaDefinition) (string, bool, error) {
-	return GenerateSchemaWithCaveatTypeSet(definitions, caveattypes.Default.TypeSet)
+	return GenerateSchemaWithCaveatTypeSet(context.TODO(), definitions, caveattypes.Default.TypeSet)
 }
 
 // GenerateSchemaWithCaveatTypeSet generates a DSL view of the given schema.
-func GenerateSchemaWithCaveatTypeSet(definitions []compiler.SchemaDefinition, caveatTypeSet *caveattypes.TypeSet) (string, bool, error) {
+func GenerateSchemaWithCaveatTypeSet(ctx context.Context, definitions []compiler.SchemaDefinition, caveatTypeSet *caveattypes.TypeSet) (string, bool, error) {
+	_, span := tracer.Start(ctx, "GenerateSchemaWithCaveatTypeSet")
+	defer span.End()
+
 	generated := make([]string, 0, len(definitions))
 	flags := mapz.NewSet[string]()