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

Author: miparnisari

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) {
+	ctx, span := tracer.Start(ctx, "GenerateSchemaWithCaveatTypeSet")
+	defer span.End()
+
 	generated := make([]string, 0, len(definitions))
 	flags := mapz.NewSet[string]()