Implement a combined builder pattern for relationship SQL construction

This moves the behavior out of Spanner datastore and into a common lib where possible

internal/datastore/common/relationships.go

@@ -17,26 +17,6 @@ import (
 
 const errUnableToQueryRels = "unable to query relationships: %w"
 
-// StaticValueOrAddColumnForSelect adds a column to the list of columns to select if the value
-// is not static, otherwise it sets the value to the static value.
-func StaticValueOrAddColumnForSelect(colsToSelect []any, queryInfo QueryInfo, colName string, field *string) []any {
-	if queryInfo.Schema.ColumnOptimization == ColumnOptimizationOptionNone {
-		// If column optimization is disabled, always add the column to the list of columns to select.
-		colsToSelect = append(colsToSelect, field)
-		return colsToSelect
-	}
-
-	// If the value is static, set the field to it and return.
-	if found, ok := queryInfo.FilteringValues[colName]; ok && found.SingleValue != nil {
-		*field = *found.SingleValue
-		return colsToSelect
-	}
-
-	// Otherwise, add the column to the list of columns to select, as the value is not static.
-	colsToSelect = append(colsToSelect, field)
-	return colsToSelect
-}
-
 // Querier is an interface for querying the database.
 type Querier[R Rows] interface {
 	QueryFunc(ctx context.Context, f func(context.Context, R) error, sql string, args ...any) error
@@ -60,10 +40,14 @@ type closeRows interface {
 }
 
 // QueryRelationships queries relationships for the given query and transaction.
-func QueryRelationships[R Rows, C ~map[string]any](ctx context.Context, queryInfo QueryInfo, sqlStatement string, args []any, span trace.Span, tx Querier[R], withIntegrity bool) (datastore.RelationshipIterator, error) {
+func QueryRelationships[R Rows, C ~map[string]any](ctx context.Context, builder RelationshipsQueryBuilder, span trace.Span, tx Querier[R]) (datastore.RelationshipIterator, error) {
 	defer span.End()
 
-	colsToSelect := make([]any, 0, 8)
+	sqlString, args, err := builder.SelectSQL()
+	if err != nil {
+		return nil, fmt.Errorf(errUnableToQueryRels, err)
+	}
+
 	var resourceObjectType string
 	var resourceObjectID string
 	var resourceRelation string
@@ -78,26 +62,9 @@ func QueryRelationships[R Rows, C ~map[string]any](ctx context.Context, queryInf
 	var integrityHash []byte
 	var timestamp time.Time
 
-	colsToSelect = StaticValueOrAddColumnForSelect(colsToSelect, queryInfo, queryInfo.Schema.ColNamespace, &resourceObjectType)
-	colsToSelect = StaticValueOrAddColumnForSelect(colsToSelect, queryInfo, queryInfo.Schema.ColObjectID, &resourceObjectID)
-	colsToSelect = StaticValueOrAddColumnForSelect(colsToSelect, queryInfo, queryInfo.Schema.ColRelation, &resourceRelation)
-	colsToSelect = StaticValueOrAddColumnForSelect(colsToSelect, queryInfo, queryInfo.Schema.ColUsersetNamespace, &subjectObjectType)
-	colsToSelect = StaticValueOrAddColumnForSelect(colsToSelect, queryInfo, queryInfo.Schema.ColUsersetObjectID, &subjectObjectID)
-	colsToSelect = StaticValueOrAddColumnForSelect(colsToSelect, queryInfo, queryInfo.Schema.ColUsersetRelation, &subjectRelation)
-
-	if !queryInfo.SkipCaveats || queryInfo.Schema.ColumnOptimization == ColumnOptimizationOptionNone {
-		colsToSelect = append(colsToSelect, &caveatName, &caveatCtx)
-	}
-
-	colsToSelect = append(colsToSelect, &expiration)
-
-	if withIntegrity {
-		colsToSelect = append(colsToSelect, &integrityKeyID, &integrityHash, &timestamp)
-	}
-
-	if len(colsToSelect) == 0 {
-		var unused int
-		colsToSelect = append(colsToSelect, &unused)
+	colsToSelect, err := ColumnsToSelect(builder, &resourceObjectType, &resourceObjectID, &resourceRelation, &subjectObjectType, &subjectObjectID, &subjectRelation, &caveatName, &caveatCtx, &expiration, &integrityKeyID, &integrityHash, &timestamp)
+	if err != nil {
+		return nil, fmt.Errorf(errUnableToQueryRels, err)
 	}
 
 	return func(yield func(tuple.Relationship, error) bool) {
@@ -117,7 +84,7 @@ func QueryRelationships[R Rows, C ~map[string]any](ctx context.Context, queryInf
 				}
 
 				var caveat *corev1.ContextualizedCaveat
-				if !queryInfo.SkipCaveats || queryInfo.Schema.ColumnOptimization == ColumnOptimizationOptionNone {
+				if !builder.SkipCaveats || builder.Schema.ColumnOptimization == ColumnOptimizationOptionNone {
 					if caveatName.Valid {
 						var err error
 						caveat, err = ContextualizedCaveatFrom(caveatName.String, caveatCtx)
@@ -171,7 +138,7 @@ func QueryRelationships[R Rows, C ~map[string]any](ctx context.Context, queryInf
 
 			span.AddEvent("Rels loaded", trace.WithAttributes(attribute.Int("relCount", relCount)))
 			return nil
-		}, sqlStatement, args...)
+		}, sqlString, args...)
 		if err != nil {
 			if !yield(tuple.Relationship{}, err) {
 				return

internal/datastore/common/schema.go

@@ -0,0 +1,122 @@
+package common
+
+import (
+	sq "github.com/Masterminds/squirrel"
+
+	"github.com/authzed/spicedb/pkg/spiceerrors"
+)
+
+// SchemaInformation holds the schema information from the SQL datastore implementation.
+//
+//go:generate go run github.com/ecordell/optgen -output schema_options.go . SchemaInformation
+type SchemaInformation struct {
+	RelationshipTableName string `debugmap:"visible"`
+
+	ColNamespace        string `debugmap:"visible"`
+	ColObjectID         string `debugmap:"visible"`
+	ColRelation         string `debugmap:"visible"`
+	ColUsersetNamespace string `debugmap:"visible"`
+	ColUsersetObjectID  string `debugmap:"visible"`
+	ColUsersetRelation  string `debugmap:"visible"`
+	ColCaveatName       string `debugmap:"visible"`
+	ColCaveatContext    string `debugmap:"visible"`
+	ColExpiration       string `debugmap:"visible"`
+
+	ColIntegrityKeyID     string `debugmap:"visible"`
+	ColIntegrityHash      string `debugmap:"visible"`
+	ColIntegrityTimestamp string `debugmap:"visible"`
+
+	// PaginationFilterType is the type of pagination filter to use for this schema.
+	PaginationFilterType PaginationFilterType `debugmap:"visible"`
+
+	// PlaceholderFormat is the format of placeholders to use for this schema.
+	PlaceholderFormat sq.PlaceholderFormat `debugmap:"visible"`
+
+	// NowFunction is the function to use to get the current time in the datastore.
+	NowFunction string `debugmap:"visible"`
+
+	// ColumnOptimization is the optimization to use for columns in the schema, if any.
+	ColumnOptimization ColumnOptimizationOption `debugmap:"visible"`
+
+	// WithIntegrityColumns is a flag to indicate if the schema has integrity columns.
+	WithIntegrityColumns bool `debugmap:"visible"`
+}
+
+func (si SchemaInformation) debugValidate() {
+	spiceerrors.DebugAssert(func() bool {
+		si.mustValidate()
+		return true
+	}, "SchemaInformation failed to validate")
+}
+
+func (si SchemaInformation) mustValidate() {
+	if si.RelationshipTableName == "" {
+		panic("RelationshipTableName is required")
+	}
+
+	if si.ColNamespace == "" {
+		panic("ColNamespace is required")
+	}
+
+	if si.ColObjectID == "" {
+		panic("ColObjectID is required")
+	}
+
+	if si.ColRelation == "" {
+		panic("ColRelation is required")
+	}
+
+	if si.ColUsersetNamespace == "" {
+		panic("ColUsersetNamespace is required")
+	}
+
+	if si.ColUsersetObjectID == "" {
+		panic("ColUsersetObjectID is required")
+	}
+
+	if si.ColUsersetRelation == "" {
+		panic("ColUsersetRelation is required")
+	}
+
+	if si.ColCaveatName == "" {
+		panic("ColCaveatName is required")
+	}
+
+	if si.ColCaveatContext == "" {
+		panic("ColCaveatContext is required")
+	}
+
+	if si.ColExpiration == "" {
+		panic("ColExpiration is required")
+	}
+
+	if si.WithIntegrityColumns {
+		if si.ColIntegrityKeyID == "" {
+			panic("ColIntegrityKeyID is required")
+		}
+
+		if si.ColIntegrityHash == "" {
+			panic("ColIntegrityHash is required")
+		}
+
+		if si.ColIntegrityTimestamp == "" {
+			panic("ColIntegrityTimestamp is required")
+		}
+	}
+
+	if si.NowFunction == "" {
+		panic("NowFunction is required")
+	}
+
+	if si.ColumnOptimization == ColumnOptimizationOptionUnknown {
+		panic("ColumnOptimization is required")
+	}
+
+	if si.PaginationFilterType == PaginationFilterTypeUnknown {
+		panic("PaginationFilterType is required")
+	}
+
+	if si.PlaceholderFormat == nil {
+		panic("PlaceholderFormat is required")
+	}
+}