googlesql: Spanner schemas, arrays, and a function seed from Spanner

internal/compiler/parse_core.go

@@ -201,5 +201,13 @@ func coreParamColumn(p core.Parameter, params *named.ParamSet) *Column {
 		col.IsSqlcSlice = param.IsSqlcSlice()
 		col.IsNamedParam = true
 	}
+	// sqlc.slice() already says the argument is a list of values, so a
+	// slice the analyzer typed as an array, as it does "= ANY(sqlc.slice(x))",
+	// is a slice of the array's elements.
+	if col.IsSqlcSlice && col.TypeExpr.IsArray() {
+		col.TypeExpr = col.TypeExpr.Element()
+		col.IsArray = false
+		describeType(col, col.TypeExpr)
+	}
 	return col
 }

internal/core/analyzer/expr.go

@@ -79,6 +79,9 @@ func (a *analyzer) typeExpr(n ast.Node) (exprType, error) {
 	case *ast.A_ArrayExpr:
 		return a.typeArrayExpr(e)
 
+	case *ast.A_Indirection:
+		return a.typeIndirection(e)
+
 	case *ast.SubLink:
 		return a.typeSubLink(e)
 
@@ -366,7 +369,17 @@ func (a *analyzer) typeQuantifiedExpr(e *ast.A_Expr) (exprType, error) {
 	if err != nil {
 		return exprType{}, err
 	}
-	if err := a.typeOperands(e.Rexpr, leftT); err != nil {
+	// "x = ANY(SELECT ...)" compares x against the subquery's column, but
+	// "x = ANY(array)" against the elements of an array, so a placeholder
+	// there is an array of x's type.
+	right := leftT
+	if _, sub := e.Rexpr.(*ast.SubLink); !sub {
+		right = a.arrayOf(leftT)
+		right.sourceClassOID = leftT.sourceClassOID
+		right.sourceAttributeOID = leftT.sourceAttributeOID
+		right.sourceTableAlias = leftT.sourceTableAlias
+	}
+	if err := a.typeOperands(e.Rexpr, right); err != nil {
 		return exprType{}, err
 	}
 	return a.boolType(false)
@@ -524,6 +537,51 @@ func (a *analyzer) typeArrayExpr(e *ast.A_ArrayExpr) (exprType, error) {
 	return a.lookupType(core.Array(element.WithNullable(false))), nil
 }
 
+// typeIndirection types an array subscript, arr[i], as the array's element,
+// and a slice, arr[i:j], as the array itself. An element may be NULL, and a
+// subscript past the end of the array yields NULL in the dialects that do not
+// raise an error for it. The subscripts are integers.
+func (a *analyzer) typeIndirection(e *ast.A_Indirection) (exprType, error) {
+	t, err := a.typeExpr(e.Arg)
+	if err != nil {
+		return exprType{}, err
+	}
+	intT := exprType{}
+	if oid, err := a.cat.ConstTypeOID(core.ConstInteger); err == nil {
+		intT.typeOID = oid
+	}
+	for _, item := range listItems(e.Indirection) {
+		idx, ok := item.(*ast.A_Indices)
+		if !ok {
+			return exprType{}, fmt.Errorf("typeExpr: unsupported indirection %T", item)
+		}
+		for _, bound := range []ast.Node{idx.Lidx, idx.Uidx} {
+			if err := a.typeOperands(bound, intT); err != nil {
+				return exprType{}, err
+			}
+		}
+		if idx.IsSlice {
+			continue
+		}
+		element := a.exprOf(t).Element()
+		if element == nil {
+			return exprType{nullable: true}, nil
+		}
+		t = a.lookupType(element)
+		t.nullable = true
+	}
+	return t, nil
+}
+
+// arrayOf is the type of an array of t's type.
+func (a *analyzer) arrayOf(t exprType) exprType {
+	element := a.exprOf(t)
+	if element == nil {
+		return exprType{}
+	}
+	return a.lookupType(core.Array(element.WithNullable(false)))
+}
+
 // lookupType is the type an expression refers to: its row when the catalog
 // holds one, its family's row when it holds only that, and the expression
 // alone when it holds neither.
@@ -619,14 +677,21 @@ func (a *analyzer) typeSubLink(e *ast.SubLink) (exprType, error) {
 		return exprType{}, err
 	}
 	switch e.SubLinkType {
-	case ast.EXPR_SUBLINK, ast.ARRAY_SUBLINK:
+	case ast.EXPR_SUBLINK:
 		if len(cols) == 0 {
 			return exprType{nullable: true}, nil
 		}
 		// A subquery that matches no row yields NULL.
 		t := columnExprType(cols[0])
 		t.nullable = true
 		return t, nil
+	case ast.ARRAY_SUBLINK:
+		// ARRAY(SELECT ...) collects the column into an array, which is
+		// empty rather than NULL when no row matches.
+		if len(cols) == 0 {
+			return exprType{}, nil
+		}
+		return a.arrayOf(columnExprType(cols[0])), nil
 	default:
 		return a.boolType(false)
 	}
@@ -831,12 +896,20 @@ func (a *analyzer) typeFuncCall(f *ast.FuncCall) (exprType, error) {
 	}
 	p := a.pickOverload(overloads, argOIDs)
 	// An argument that is a bare placeholder takes the parameter's type,
-	// unless the parameter is polymorphic and says nothing.
+	// unless the parameter is polymorphic, when it takes what the other
+	// arguments say: the element of the array passed for an anyarray, or
+	// an array of the value passed for an anyelement, as
+	// array_includes(anyarray, anyelement) relates them.
 	for i, arg := range args {
 		if i >= len(p.ArgTypes) {
 			break
 		}
 		if a.isPolymorphicOID(p.ArgTypes[i]) {
+			if t, ok := a.polymorphicArg(p, i, argTypes); ok {
+				if err := a.typeOperands(arg, t); err != nil {
+					return exprType{}, err
+				}
+			}
 			continue
 		}
 		if err := a.typeOperands(arg, exprType{typeOID: p.ArgTypes[i]}); err != nil {
@@ -916,11 +989,61 @@ func (a *analyzer) returnType(p core.ProcOverload, argTypes []exprType) exprType
 		return exprType{}
 	}
 	if isPolymorphic(name) && (argTypes[0].typeOID != 0 || argTypes[0].expr != nil) {
-		return exprType{typeOID: argTypes[0].typeOID, expr: argTypes[0].expr}
+		arg := exprType{typeOID: argTypes[0].typeOID, expr: argTypes[0].expr}
+		// An element taken from an array argument, unnest(anyarray)
+		// returning anyelement, is the array's element; an array built from
+		// an element argument is an array of it.
+		paramName := ""
+		if len(p.ArgTypes) > 0 {
+			paramName, _ = a.cat.TypeName(p.ArgTypes[0])
+		}
+		switch {
+		case isPolymorphicArray(paramName) && !isPolymorphicArray(name):
+			if element := a.exprOf(arg).Element(); element != nil {
+				return a.lookupType(element)
+			}
+		case !isPolymorphicArray(paramName) && isPolymorphicArray(name) && isPolymorphic(paramName):
+			return a.arrayOf(arg)
+		}
+		return arg
 	}
 	return exprType{typeOID: p.ReturnTypeOID}
 }
 
+// polymorphicArg is the type the other arguments of a call give its i-th,
+// polymorphic, parameter: the element of an array passed for an anyarray
+// when the parameter is an element, or an array of an element passed for
+// an anyelement when the parameter is an array.
+func (a *analyzer) polymorphicArg(p core.ProcOverload, i int, argTypes []exprType) (exprType, bool) {
+	name, err := a.cat.TypeName(p.ArgTypes[i])
+	if err != nil || !isPolymorphic(name) {
+		return exprType{}, false
+	}
+	for j, other := range p.ArgTypes {
+		if j == i || j >= len(argTypes) || (argTypes[j].typeOID == 0 && argTypes[j].expr == nil) {
+			continue
+		}
+		otherName, err := a.cat.TypeName(other)
+		if err != nil || !isPolymorphic(otherName) {
+			continue
+		}
+		switch {
+		case isPolymorphicArray(otherName) && !isPolymorphicArray(name):
+			if element := a.exprOf(argTypes[j]).Element(); element != nil {
+				return a.lookupType(element), true
+			}
+		case !isPolymorphicArray(otherName) && isPolymorphicArray(name):
+			return a.arrayOf(argTypes[j]), true
+		}
+	}
+	return exprType{}, false
+}
+
+// isPolymorphicArray reports whether a polymorphic type stands for an array.
+func isPolymorphicArray(typeName string) bool {
+	return typeName == "anyarray" || typeName == "anycompatiblearray"
+}
+
 // argIndex reads a seed's "$n" pseudo-type as the zero-based index of the
 // argument whose type it stands for.
 func argIndex(typeName string) (int, bool) {

internal/core/analyzer/scope.go

@@ -100,19 +100,28 @@ func (a *analyzer) bindRangeFunction(rf *ast.RangeFunction) (scopeRel, error) {
 	}
 
 	// The arguments are typed against the scope built so far, which is what
-	// gives a placeholder passed to the function its type.
-	if _, err := a.typeFuncCall(call); err != nil {
+	// gives a placeholder passed to the function its type, and what
+	// resolves a polymorphic result: unnest of an array of strings yields
+	// strings.
+	t, err := a.typeFuncCall(call)
+	if err != nil {
 		return scopeRel{}, err
 	}
 
-	p := overloads[0]
-	rel := scopeRel{
-		alias: name,
-		cols:  []core.ClassColumn{{Name: name, TypeOID: p.ReturnTypeOID, NotNull: !p.ReturnNullable}},
-	}
+	// The relation's alias names its single column too, as "unnest(tags)
+	// AS tag" does.
+	colName := name
+	rel := scopeRel{alias: name}
 	if rf.Alias != nil && rf.Alias.Aliasname != nil && *rf.Alias.Aliasname != "" {
 		rel.alias = *rf.Alias.Aliasname
+		colName = rel.alias
+	}
+	col := core.ClassColumn{Name: colName, TypeOID: t.typeOID, Type: t.expr, NotNull: !t.nullable}
+	if t.typeOID == 0 && t.expr == nil {
+		p := overloads[0]
+		col.TypeOID, col.NotNull = p.ReturnTypeOID, !p.ReturnNullable
 	}
+	rel.cols = []core.ClassColumn{col}
 	return rel, nil
 }
 

internal/core/seed/seed.go

@@ -320,12 +320,24 @@ func apply(cat *core.Catalog, fsys fs.FS, settings Settings) error {
 	if err := b.rules(); err != nil {
 		return err
 	}
-	if err := b.categoryOperators(); err != nil {
+	// An operator the dialect lists replaces the one its rules would
+	// generate for the same operands, as GoogleSQL's INT64 / INT64 is a
+	// FLOAT64, so the list is read before the rules apply.
+	var operators []Operator
+	if err := stream(fsys, OperatorsFile, func(op Operator) error {
+		operators = append(operators, op)
+		return nil
+	}); err != nil {
 		return err
 	}
-	if err := stream(fsys, OperatorsFile, b.addOperator); err != nil {
+	if err := b.categoryOperators(operators); err != nil {
 		return err
 	}
+	for i, op := range operators {
+		if err := b.addOperator(op); err != nil {
+			return fmt.Errorf("seed: %s: record %d: %w", OperatorsFile, i+1, err)
+		}
+	}
 	if err := stream(fsys, CastsFile, b.addCast); err != nil {
 		return err
 	}
@@ -713,8 +725,12 @@ func (b *builder) rules() error {
 	return nil
 }
 
-func (b *builder) categoryOperators() error {
+func (b *builder) categoryOperators(listed []Operator) error {
 	s := b.settings
+	overridden := map[[3]string]bool{}
+	for _, op := range listed {
+		overridden[[3]string{op.Name, strings.ToLower(op.Left), strings.ToLower(op.Right)}] = true
+	}
 	boolOID, ok := b.oids[strings.ToLower(s.Bool)]
 	if len(s.Comparison) > 0 && !ok {
 		return fmt.Errorf("seed %s: comparison operators need a bool type", s.Dialect)
@@ -729,6 +745,9 @@ func (b *builder) categoryOperators() error {
 		}
 		if strings.Contains(s.ArithmeticCategories, t.category) {
 			for _, op := range s.Arithmetic {
+				if overridden[[3]string{op, strings.ToLower(t.name), strings.ToLower(t.name)}] {
+					continue
+				}
 				if err := b.operator(op, t.name, t.name, t.name, b.oids[t.name]); err != nil {
 					return err
 				}

internal/endtoend/testdata/analyze_arrays/googlesql/exec.json

@@ -0,0 +1,5 @@
+{
+  "command": "analyze",
+  "args": ["--dialect", "googlesql", "--schema", "schema.sql", "query.sql"],
+  "contexts": ["base"]
+}

internal/endtoend/testdata/analyze_arrays/googlesql/query.sql

@@ -0,0 +1,8 @@
+-- name: AlbumIDs :one
+SELECT singer_id, ARRAY(SELECT album_id FROM albums WHERE albums.singer_id = singers.singer_id) AS album_ids
+FROM singers WHERE singer_id = @singer_id;
+
+-- name: Literals :one
+SELECT [1, 2, 3] AS ints, ARRAY<STRING>['a'] AS strs, DATE '2024-01-02' AS d,
+  TIMESTAMP '2024-01-02 03:04:05+00' AS ts, NUMERIC '1.5' AS n, JSON '{"a": 1}' AS j,
+  b'abc' AS raw, INTERVAL 3 DAY AS iv;

internal/endtoend/testdata/analyze_arrays/googlesql/schema.sql

@@ -0,0 +1,13 @@
+CREATE TABLE singers (
+  singer_id INT64 NOT NULL,
+  name STRING(MAX) NOT NULL,
+  tags ARRAY<STRING(MAX)>,
+  scores ARRAY<INT64> NOT NULL,
+) PRIMARY KEY (singer_id);
+
+CREATE TABLE albums (
+  singer_id INT64 NOT NULL,
+  album_id INT64 NOT NULL,
+  title STRING(MAX),
+) PRIMARY KEY (singer_id, album_id),
+  INTERLEAVE IN PARENT singers ON DELETE CASCADE;