chore: better tracing (generator.GenerateSchema, WriteSchema)

internal/lsp/handlers.go

@@ -282,7 +282,7 @@ func (s *Server) textDocHover(_ context.Context, r *jsonrpc2.Request) (*Hover, e
 	return hoverContents, nil
 }
 
-func (s *Server) textDocFormat(_ context.Context, r *jsonrpc2.Request) ([]lsp.TextEdit, error) {
+func (s *Server) textDocFormat(ctx context.Context, r *jsonrpc2.Request) ([]lsp.TextEdit, error) {
 	params, err := unmarshalParams[lsp.DocumentFormattingParams](r)
 	if err != nil {
 		return nil, err
@@ -295,7 +295,7 @@ func (s *Server) textDocFormat(_ context.Context, r *jsonrpc2.Request) ([]lsp.Te
 			return err
 		}
 
-		formattedSchema, _, err := generator.GenerateSchema(compiled.OrderedDefinitions)
+		formattedSchema, _, err := generator.GenerateSchema(ctx, compiled.OrderedDefinitions)
 		if err != nil {
 			return err
 		}

internal/services/shared/schema.go

@@ -4,6 +4,8 @@
 	"context"
 	"maps"
 
+	"go.opentelemetry.io/otel"
+
 	log "github.com/authzed/spicedb/internal/logging"
 	"github.com/authzed/spicedb/internal/namespace"
 	caveattypes "github.com/authzed/spicedb/pkg/caveats/types"
@@ -21,6 +23,8 @@
 	"github.com/authzed/spicedb/pkg/tuple"
 )
 
+var tracer = otel.Tracer("spicedb/internal/services/shared")
+
 // ValidatedSchemaChanges is a set of validated schema changes that can be applied to the datastore.
 type ValidatedSchemaChanges struct {
 	compiled             *compiler.CompiledSchema
@@ -36,6 +40,8 @@
 // ValidateSchemaChanges validates the schema found in the compiled schema and returns a
 // ValidatedSchemaChanges, if fully validated.
 func ValidateSchemaChanges(ctx context.Context, compiled *compiler.CompiledSchema, caveatTypeSet *caveattypes.TypeSet, additiveOnly bool, schemaText string) (*ValidatedSchemaChanges, error) {
+	ctx, span := tracer.Start(ctx, "schema.ValidateSchemaChanges")
+	defer span.End()
 	// 1) Validate the caveats defined.
 	newCaveatDefNames := mapz.NewSet[string]()
 	for _, caveatDef := range compiled.CaveatDefinitions {
@@ -94,6 +100,8 @@
 // ApplySchemaChanges applies schema changes found in the validated changes struct, via the specified
 // ReadWriteTransaction.
 func ApplySchemaChanges(ctx context.Context, rwt datalayer.ReadWriteTransaction, caveatTypeSet *caveattypes.TypeSet, validated *ValidatedSchemaChanges) (*AppliedSchemaChanges, error) {
+	ctx, span := tracer.Start(ctx, "schema.ApplySchemaChanges")
+	defer span.End()
 	sr, err := rwt.ReadSchema(ctx)
 	if err != nil {
 		return nil, err
@@ -132,6 +140,8 @@
 	existingCaveats []*core.CaveatDefinition,
 	existingObjectDefs []*core.NamespaceDefinition,
 ) (*AppliedSchemaChanges, error) {
+	ctx, span := tracer.Start(ctx, "schema.ApplySchemaChangesOverExisting")
+	defer span.End()
 	// Build a map of existing caveats to determine those being removed, if any.
 	existingCaveatDefMap := make(map[string]*core.CaveatDefinition, len(existingCaveats))
 	existingCaveatDefNames := mapz.NewSet[string]()

internal/services/v1/debug.go

@@ -60,7 +60,7 @@
 		defs = append(defs, ns.Definition)
 	}
 
-	schema, _, err := generator.GenerateSchema(defs)
+	schema, _, err := generator.GenerateSchema(ctx, defs)
 	if err != nil {
 		return "", err
 	}

internal/services/v1/schema.go

@@ -7,6 +7,7 @@
 
 	"buf.build/go/protovalidate"
 	grpcvalidate "github.com/grpc-ecosystem/go-grpc-middleware/v2/interceptors/protovalidate"
+	"go.opentelemetry.io/otel"
 
 	v1 "github.com/authzed/authzed-go/proto/authzed/api/v1"
 
@@ -29,6 +30,8 @@
 	"github.com/authzed/spicedb/pkg/zedtoken"
 )
 
+var tracer = otel.Tracer("spicedb/internal/services/v1/schema")
+
 type SchemaServerConfig struct {
 	// CaveatTypeSet is the set of caveat types that are allowed in the schema.
 	CaveatTypeSet *caveattypes.TypeSet
@@ -142,13 +145,16 @@
 	// the user must first compile them with `zed`
 	opts = append(opts, compiler.DisallowImportFlag())
 
+	_, span := tracer.Start(ctx, "compile")
 	compiled, err := compiler.Compile(compiler.InputSchema{
 		Source:       input.Source("schema"),
 		SchemaString: in.GetSchema(),
 	}, compiler.AllowUnprefixedObjectType(), opts...)
 	if err != nil {
+		span.End()
 		return nil, ss.rewriteError(ctx, err)
 	}
+	span.End()
 	log.Ctx(ctx).Trace().Int("objectDefinitions", len(compiled.ObjectDefinitions)).Int("caveatDefinitions", len(compiled.CaveatDefinitions)).Msg("compiled namespace definitions")
 
 	// Do as much validation as we can before talking to the datastore.

pkg/datastore/test/namespace.go

@@ -371,7 +371,7 @@ definition document {
 	testutil.RequireProtoEqual(t, caveatDef, readCaveatDef, "found changed caveat definition")
 
 	// Ensure the read namespace's string form matches the input as an extra check.
-	generated, _, err := generator.GenerateSchema([]compiler.SchemaDefinition{readCaveatDef, readNsDef})
+	generated, _, err := generator.GenerateSchema(ctx, []compiler.SchemaDefinition{readCaveatDef, readNsDef})
 	require.NoError(err)
 	require.Equal(schemaString, generated)
 }

pkg/development/wasm/operations.go

@@ -18,7 +18,7 @@ import (
 func runOperation(devContext *development.DevContext, operation *devinterface.Operation) (*devinterface.OperationResult, error) {
 	switch {
 	case operation.FormatSchemaParameters != nil:
-		formatted, _, err := generator.GenerateSchema(devContext.CompiledSchema.OrderedDefinitions)
+		formatted, _, err := generator.GenerateSchema(devContext.Ctx, devContext.CompiledSchema.OrderedDefinitions)
 		if err != nil {
 			return nil, err
 		}

pkg/schema/v2/builder_example_test.go

@@ -1,6 +1,7 @@
 package schema_test
 
 import (
+	"context"
 	"fmt"
 	"strings"
 
@@ -39,7 +40,7 @@ func ExampleSchemaBuilder() {
 		Build()
 
 	defs, caveats, _ := s.ToDefinitions()
-	schemaText, _, _ := generator.GenerateSchema(append(toSchemaDefinitions(defs), toSchemaDefinitions(caveats)...))
+	schemaText, _, _ := generator.GenerateSchema(context.Background(), append(toSchemaDefinitions(defs), toSchemaDefinitions(caveats)...))
 	fmt.Println(strings.TrimSpace(schemaText))
 }
 
@@ -67,7 +68,7 @@ func ExampleSchemaBuilder_incrementalUnion() {
 		Build()
 
 	defs, caveats, _ := s.ToDefinitions()
-	schemaText, _, _ := generator.GenerateSchema(append(toSchemaDefinitions(defs), toSchemaDefinitions(caveats)...))
+	schemaText, _, _ := generator.GenerateSchema(context.Background(), append(toSchemaDefinitions(defs), toSchemaDefinitions(caveats)...))
 	fmt.Println(strings.TrimSpace(schemaText))
 }
 
@@ -99,7 +100,7 @@ func ExampleSchemaBuilder_withArrow() {
 		Build()
 
 	defs, caveats, _ := s.ToDefinitions()
-	schemaText, _, _ := generator.GenerateSchema(append(toSchemaDefinitions(defs), toSchemaDefinitions(caveats)...))
+	schemaText, _, _ := generator.GenerateSchema(context.Background(), append(toSchemaDefinitions(defs), toSchemaDefinitions(caveats)...))
 	fmt.Println(strings.TrimSpace(schemaText))
 }
 
@@ -119,7 +120,7 @@ func ExampleSchemaBuilder_withCaveat() {
 		Build()
 
 	defs, caveats, _ := s.ToDefinitions()
-	schemaText, _, _ := generator.GenerateSchema(append(toSchemaDefinitions(defs), toSchemaDefinitions(caveats)...))
+	schemaText, _, _ := generator.GenerateSchema(context.Background(), append(toSchemaDefinitions(defs), toSchemaDefinitions(caveats)...))
 	fmt.Println(strings.TrimSpace(schemaText))
 }
 
@@ -148,7 +149,7 @@ func ExampleSchemaBuilder_withCaveatTypes() {
 		Build()
 
 	defs, caveats, _ := s.ToDefinitions()
-	schemaText, _, _ := generator.GenerateSchema(append(toSchemaDefinitions(defs), toSchemaDefinitions(caveats)...))
+	schemaText, _, _ := generator.GenerateSchema(context.Background(), append(toSchemaDefinitions(defs), toSchemaDefinitions(caveats)...))
 	fmt.Println(strings.TrimSpace(schemaText))
 }
 
@@ -172,7 +173,7 @@ func ExampleSchemaBuilder_withExclusion() {
 		Build()
 
 	defs, caveats, _ := s.ToDefinitions()
-	schemaText, _, _ := generator.GenerateSchema(append(toSchemaDefinitions(defs), toSchemaDefinitions(caveats)...))
+	schemaText, _, _ := generator.GenerateSchema(context.Background(), append(toSchemaDefinitions(defs), toSchemaDefinitions(caveats)...))
 	fmt.Println(strings.TrimSpace(schemaText))
 }
 
@@ -196,7 +197,7 @@ func ExampleSchemaBuilder_withIntersection() {
 		Build()
 
 	defs, caveats, _ := s.ToDefinitions()
-	schemaText, _, _ := generator.GenerateSchema(append(toSchemaDefinitions(defs), toSchemaDefinitions(caveats)...))
+	schemaText, _, _ := generator.GenerateSchema(context.Background(), append(toSchemaDefinitions(defs), toSchemaDefinitions(caveats)...))
 	fmt.Println(strings.TrimSpace(schemaText))
 }
 
@@ -211,7 +212,7 @@ func ExampleSchemaBuilder_withWildcard() {
 		Build()
 
 	defs, caveats, _ := s.ToDefinitions()
-	schemaText, _, _ := generator.GenerateSchema(append(toSchemaDefinitions(defs), toSchemaDefinitions(caveats)...))
+	schemaText, _, _ := generator.GenerateSchema(context.Background(), append(toSchemaDefinitions(defs), toSchemaDefinitions(caveats)...))
 	fmt.Println(strings.TrimSpace(schemaText))
 }
 
@@ -280,7 +281,7 @@ func ExampleSchemaBuilder_complex() {
 		Build()
 
 	defs, caveats, _ := s.ToDefinitions()
-	schemaText, _, _ := generator.GenerateSchema(append(toSchemaDefinitions(defs), toSchemaDefinitions(caveats)...))
+	schemaText, _, _ := generator.GenerateSchema(context.Background(), append(toSchemaDefinitions(defs), toSchemaDefinitions(caveats)...))
 	fmt.Println(strings.TrimSpace(schemaText))
 }
 
@@ -306,7 +307,7 @@ func ExampleSchemaBuilder_constructorPattern() {
 		Build()
 
 	defs, caveats, _ := s.ToDefinitions()
-	schemaText, _, _ := generator.GenerateSchema(append(toSchemaDefinitions(defs), toSchemaDefinitions(caveats)...))
+	schemaText, _, _ := generator.GenerateSchema(context.Background(), append(toSchemaDefinitions(defs), toSchemaDefinitions(caveats)...))
 	fmt.Println(strings.TrimSpace(schemaText))
 }
 
@@ -327,7 +328,7 @@ func ExampleSchemaBuilder_constructorWithIntersection() {
 		Build()
 
 	defs, caveats, _ := s.ToDefinitions()
-	schemaText, _, _ := generator.GenerateSchema(append(toSchemaDefinitions(defs), toSchemaDefinitions(caveats)...))
+	schemaText, _, _ := generator.GenerateSchema(context.Background(), append(toSchemaDefinitions(defs), toSchemaDefinitions(caveats)...))
 	fmt.Println(strings.TrimSpace(schemaText))
 }
 
@@ -348,7 +349,7 @@ func ExampleSchemaBuilder_constructorWithExclusion() {
 		Build()
 
 	defs, caveats, _ := s.ToDefinitions()
-	schemaText, _, _ := generator.GenerateSchema(append(toSchemaDefinitions(defs), toSchemaDefinitions(caveats)...))
+	schemaText, _, _ := generator.GenerateSchema(context.Background(), append(toSchemaDefinitions(defs), toSchemaDefinitions(caveats)...))
 	fmt.Println(strings.TrimSpace(schemaText))
 }
 
@@ -374,7 +375,7 @@ func ExampleSchemaBuilder_constructorWithArrow() {
 		Build()
 
 	defs, caveats, _ := s.ToDefinitions()
-	schemaText, _, _ := generator.GenerateSchema(append(toSchemaDefinitions(defs), toSchemaDefinitions(caveats)...))
+	schemaText, _, _ := generator.GenerateSchema(context.Background(), append(toSchemaDefinitions(defs), toSchemaDefinitions(caveats)...))
 	fmt.Println(strings.TrimSpace(schemaText))
 }
 
@@ -395,7 +396,7 @@ func ExampleSchemaBuilder_constructorWithIntersectionArrow() {
 		Build()
 
 	defs, caveats, _ := s.ToDefinitions()
-	schemaText, _, _ := generator.GenerateSchema(append(toSchemaDefinitions(defs), toSchemaDefinitions(caveats)...))
+	schemaText, _, _ := generator.GenerateSchema(context.Background(), append(toSchemaDefinitions(defs), toSchemaDefinitions(caveats)...))
 	fmt.Println(strings.TrimSpace(schemaText))
 }
 
@@ -418,7 +419,7 @@ func ExampleSchemaBuilder_constructorWithCaveat() {
 		Build()
 
 	defs, caveats, _ := s.ToDefinitions()
-	schemaText, _, _ := generator.GenerateSchema(append(toSchemaDefinitions(defs), toSchemaDefinitions(caveats)...))
+	schemaText, _, _ := generator.GenerateSchema(context.Background(), append(toSchemaDefinitions(defs), toSchemaDefinitions(caveats)...))
 	fmt.Println(strings.TrimSpace(schemaText))
 }
 
@@ -475,7 +476,7 @@ func ExampleSchemaBuilder_constructorComplexNested() {
 		Build()
 
 	defs, caveats, _ := s.ToDefinitions()
-	schemaText, _, _ := generator.GenerateSchema(append(toSchemaDefinitions(defs), toSchemaDefinitions(caveats)...))
+	schemaText, _, _ := generator.GenerateSchema(context.Background(), append(toSchemaDefinitions(defs), toSchemaDefinitions(caveats)...))
 	fmt.Println(strings.TrimSpace(schemaText))
 }
 
@@ -498,6 +499,6 @@ func ExampleSchemaBuilder_withIntersectionArrow() {
 		Build()
 
 	defs, caveats, _ := s.ToDefinitions()
-	schemaText, _, _ := generator.GenerateSchema(append(toSchemaDefinitions(defs), toSchemaDefinitions(caveats)...))
+	schemaText, _, _ := generator.GenerateSchema(context.Background(), append(toSchemaDefinitions(defs), toSchemaDefinitions(caveats)...))
 	fmt.Println(strings.TrimSpace(schemaText))
 }

pkg/schema/v2/flatten_test.go

@@ -327,7 +327,7 @@ definition folder {
 				schemaDefinitions = append(schemaDefinitions, caveat)
 			}
 
-			generatedSchema, _, err := generator.GenerateSchema(schemaDefinitions)
+			generatedSchema, _, err := generator.GenerateSchema(t.Context(), schemaDefinitions)
 			require.NoError(t, err)
 
 			// Verify the generated schema compiles
@@ -365,7 +365,7 @@ definition folder {
 				expectedSchemaDefinitions = append(expectedSchemaDefinitions, caveat)
 			}
 
-			expectedGenerated, _, err := generator.GenerateSchema(expectedSchemaDefinitions)
+			expectedGenerated, _, err := generator.GenerateSchema(t.Context(), expectedSchemaDefinitions)
 			require.NoError(t, err)
 
 			// Compare the two generated schemas
@@ -674,7 +674,7 @@ definition user {}`,
 				schemaDefinitions = append(schemaDefinitions, caveat)
 			}
 
-			generatedSchema, _, err := generator.GenerateSchema(schemaDefinitions)
+			generatedSchema, _, err := generator.GenerateSchema(t.Context(), schemaDefinitions)
 			require.NoError(t, err)
 
 			// Verify the generated schema compiles
@@ -712,7 +712,7 @@ definition user {}`,
 				expectedSchemaDefinitions = append(expectedSchemaDefinitions, caveat)
 			}
 
-			expectedGenerated, _, err := generator.GenerateSchema(expectedSchemaDefinitions)
+			expectedGenerated, _, err := generator.GenerateSchema(t.Context(), expectedSchemaDefinitions)
 			require.NoError(t, err)
 
 			// Compare the two generated schemas
@@ -1116,7 +1116,7 @@ definition user {}
 				schemaDefinitions = append(schemaDefinitions, caveat)
 			}
 
-			generatedSchema, _, err := generator.GenerateSchema(schemaDefinitions)
+			generatedSchema, _, err := generator.GenerateSchema(t.Context(), schemaDefinitions)
 			require.NoError(t, err)
 
 			// Verify the generated schema compiles
@@ -1154,7 +1154,7 @@ definition user {}
 				expectedSchemaDefinitions = append(expectedSchemaDefinitions, caveat)
 			}
 
-			expectedGenerated, _, err := generator.GenerateSchema(expectedSchemaDefinitions)
+			expectedGenerated, _, err := generator.GenerateSchema(t.Context(), expectedSchemaDefinitions)
 			require.NoError(t, err)
 
 			// Compare the two generated schemas

pkg/schema/v2/testing/generator_test.go

@@ -30,7 +30,7 @@ func TestExampleRunWithSchemaForTesting(t *testing.T) {
 			definitions = append(definitions, cd)
 		}
 
-		generated, _, err := generator.GenerateSchema(definitions)
+		generated, _, err := generator.GenerateSchema(t.Context(), definitions)
 		require.NoError(t, err)
 		t.Logf("Generated schema:\n%s", generated)
 

pkg/schemadsl/compiler/importer_test.go

@@ -82,7 +82,7 @@ func TestImporter(t *testing.T) {
 				compiler.SourceFolder(sourceFolder))
 			require.NoError(t, err)
 
-			generated, _, err := generator.GenerateSchema(compiled.OrderedDefinitions)
+			generated, _, err := generator.GenerateSchema(t.Context(), compiled.OrderedDefinitions)
 			require.NoError(t, err)
 
 			if os.Getenv("REGEN") == "true" {
@@ -107,7 +107,7 @@ func TestImporter(t *testing.T) {
 				compiler.SourceFS(fsys))
 			require.NoError(t, err)
 
-			generated, _, err := generator.GenerateSchema(compiled.OrderedDefinitions)
+			generated, _, err := generator.GenerateSchema(t.Context(), compiled.OrderedDefinitions)
 			require.NoError(t, err)
 
 			if os.Getenv("REGEN") == "true" {

pkg/schemadsl/generator/generator.go

@@ -29,8 +29,8 @@ 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(context.TODO(), definitions, caveattypes.Default.TypeSet)
+func GenerateSchema(ctx context.Context, definitions []compiler.SchemaDefinition) (string, bool, error) {
+	return GenerateSchemaWithCaveatTypeSet(ctx, definitions, caveattypes.Default.TypeSet)
 }
 
 // GenerateSchemaWithCaveatTypeSet generates a DSL view of the given schema.

pkg/schemadsl/generator/generator_test.go

@@ -459,7 +459,7 @@ definition user {
 			}, compiler.AllowUnprefixedObjectType())
 			require.NoError(err)
 
-			source, _, err := GenerateSchema(compiled.OrderedDefinitions)
+			source, _, err := GenerateSchema(t.Context(), compiled.OrderedDefinitions)
 			require.NoError(err)
 			require.Equal(test.expected, source)
 		})