From 2bd5ca5a12b0366204f93cf45e3061da75ebfeab Mon Sep 17 00:00:00 2001 From: Zach Musgrave Date: Mon, 14 Apr 2025 14:27:39 -0700 Subject: [PATCH 1/5] Skipped test for select from function --- testing/go/functions_test.go | 34 ++++++++++++++++++++++++++++++++++ 1 file changed, 34 insertions(+) diff --git a/testing/go/functions_test.go b/testing/go/functions_test.go index eef6284020..d2741d268e 100644 --- a/testing/go/functions_test.go +++ b/testing/go/functions_test.go @@ -2310,5 +2310,39 @@ func TestSelectFromFunctions(t *testing.T) { }, }, }, + { + Name: "test select from dolt_ functions", + Skip: true, // need a way for single-row functions to declare a schema like table functions do, maybe just by modeling them as table functions in the first place + SetUpScript: []string{ + "CREATE TABLE test (pk INT primary key, v1 INT, v2 TEXT);", + "INSERT INTO test VALUES (1, 1, 'a'), (2, 2, 'b'), (3, 3, 'c'), (4, 4, 'd'), (5, 5, 'e');", + "call dolt_commit('-Am', 'first table');", + }, + Assertions: []ScriptTestAssertion{ + { + Query: `select * from dolt_branch('newBranch')`, + Expected: []sql.Row{{0}}, + }, + { + Query: `select status from dolt_checkout('newBranch')`, + Expected: []sql.Row{{0}}, + }, + { + Query: `insert into test values (6, 6, 'f')`, + }, + { + Query: `select length(commit_hash) > 0 from (select commit_hash from dolt_commit('-Am', 'added f') as result)`, + Expected: []sql.Row{{"t"}}, + }, + { + Query: "select dolt_checkout('main')", + Expected: []sql.Row{{0}}, + }, + { + Query: `select fast_forward, conflicts from dolt_merge('newBranch')`, + Expected: []sql.Row{{"t", 0}}, + }, + }, + }, }) } From 467a6d2685436dfea78817053b4af0518cafaa9e Mon Sep 17 00:00:00 2001 From: Zach Musgrave Date: Mon, 14 Apr 2025 15:20:48 -0700 Subject: [PATCH 2/5] Array subscript first pass --- server/ast/expr.go | 21 +++++- server/expression/subscript.go | 124 +++++++++++++++++++++++++++++++++ testing/go/expressions_test.go | 36 +++++++++- 3 files changed, 179 insertions(+), 2 deletions(-) create mode 100755 server/expression/subscript.go diff --git a/server/ast/expr.go b/server/ast/expr.go index 2899eccfaf..063c8b2d2d 100644 --- a/server/ast/expr.go +++ b/server/ast/expr.go @@ -607,7 +607,26 @@ func nodeExpr(ctx *Context, node tree.Expr) (vitess.Expr, error) { // TODO: figure out if I can delete this return nil, errors.Errorf("this should probably be deleted (internal error, IndexedVar)") case *tree.IndirectionExpr: - return nil, errors.Errorf("subscripts are not yet supported") + childExpr, err := nodeExpr(ctx, node.Expr) + if err != nil { + return nil, err + } + + if len(node.Indirection) > 1 { + return nil, errors.Errorf("multi dimensional array subscripts are not yet supported") + } else if node.Indirection[0].Slice { + return nil, errors.Errorf("slice subscripts are not yet supported") + } + + indexExpr, err := nodeExpr(ctx, node.Indirection[0].Begin) + if err != nil { + return nil, err + } + + return vitess.InjectedExpr{ + Expression: &pgexprs.Subscript{}, + Children: vitess.Exprs{childExpr, indexExpr}, + }, nil case *tree.IsNotNullExpr: expr, err := nodeExpr(ctx, node.Expr) if err != nil { diff --git a/server/expression/subscript.go b/server/expression/subscript.go new file mode 100755 index 0000000000..7d679c26a9 --- /dev/null +++ b/server/expression/subscript.go @@ -0,0 +1,124 @@ +// Copyright 2025 Dolthub, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package expression + +import ( + "fmt" + + "github.com/dolthub/doltgresql/server/types" + "github.com/dolthub/go-mysql-server/sql" + vitess "github.com/dolthub/vitess/go/vt/sqlparser" +) + +type Subscript struct { + Child sql.Expression + Index sql.Expression +} + +var _ vitess.Injectable = (*Subscript)(nil) +var _ sql.Expression = (*Subscript)(nil) + +func NewSubscript(child, index sql.Expression) *Subscript { + return &Subscript{ + Child: child, + Index: index, + } +} + +func (s Subscript) Resolved() bool { + return s.Child.Resolved() && s.Index.Resolved() +} + +func (s Subscript) String() string { + return fmt.Sprintf("%s[%s]", s.Child, s.Index) +} + +func (s Subscript) Type() sql.Type { + dt, ok := s.Child.Type().(*types.DoltgresType) + if !ok { + panic(fmt.Sprintf("unexpected type %T for subscript", s.Child.Type())) + } + return dt.ArrayBaseType() +} + +func (s Subscript) IsNullable() bool { + return true +} + +// Eval implements the sql.Expression interface. +func (s Subscript) Eval(ctx *sql.Context, row sql.Row) (interface{}, error) { + childVal, err := s.Child.Eval(ctx, row) + if err != nil { + return nil, err + } + if childVal == nil { + return nil, nil + } + + indexVal, err := s.Index.Eval(ctx, row) + if err != nil { + return nil, err + } + if indexVal == nil { + return nil, nil + } + + switch child := childVal.(type) { + case []interface{}: + index, ok := indexVal.(int32) + if !ok { + converted, _, err := types.Int32.Convert(ctx, indexVal) + if err != nil { + return nil, err + } + index = converted.(int32) + } + + // subscripts are 1-based + if index < 1 || int(index) > len(child) { + return nil, nil + } + return child[index-1], nil + default: + return nil, fmt.Errorf("unsupported type %T for subscript", child) + } +} + +func (s Subscript) Children() []sql.Expression { + return []sql.Expression{s.Child, s.Index} +} + +func (s Subscript) WithChildren(children ...sql.Expression) (sql.Expression, error) { + if len(children) != 2 { + return nil, fmt.Errorf("expected 2 children, got %d", len(children)) + } + return NewSubscript(children[0], children[1]), nil +} + +func (s Subscript) WithResolvedChildren(children []any) (any, error) { + if len(children) != 2 { + return nil, fmt.Errorf("expected 2 children, got %d", len(children)) + } + child, ok := children[0].(sql.Expression) + if !ok { + return nil, fmt.Errorf("expected child to be an expression but has type `%T`", children[0]) + } + index, ok := children[1].(sql.Expression) + if !ok { + return nil, fmt.Errorf("expected index to be an expression but has type `%T`", children[1]) + } + + return NewSubscript(child, index), nil +} diff --git a/testing/go/expressions_test.go b/testing/go/expressions_test.go index 894a08d59e..e605bf0007 100644 --- a/testing/go/expressions_test.go +++ b/testing/go/expressions_test.go @@ -21,7 +21,7 @@ import ( "github.com/dolthub/go-mysql-server/sql" ) -func TestExpressions(t *testing.T) { +func TestIn(t *testing.T) { RunScriptsWithoutNormalization(t, []ScriptTest{ anyTests("ANY"), anyTests("SOME"), @@ -292,3 +292,37 @@ func TestBinaryLogic(t *testing.T) { }, }) } + +func TestSubscript(t *testing.T) { + RunScripts(t, []ScriptTest{ + { + Name: "array literal", + Assertions: []ScriptTestAssertion{ + { + Query: `SELECT ARRAY[1, 2, 3][1];`, + Expected: []sql.Row{{1}}, + }, + { + Query: `SELECT (ARRAY[1, 2, 3])[3];`, + Expected: []sql.Row{{3}}, + }, + { + Query: `SELECT (ARRAY[1, 2, 3])[1+1];`, + Expected: []sql.Row{{2}}, + }, + { + Query: `SELECT ARRAY[1, 2, 3][0];`, + Expected: []sql.Row{{nil}}, + }, + { + Query: `SELECT ARRAY[1, 2, 3][4];`, + Expected: []sql.Row{{nil}}, + }, + { + Query: `SELECT ARRAY[1, 2, 3]['abc'];`, + ExpectedErr: "integer: unhandled type: string", + }, + }, + }, + }) +} From dabbb11ce7389c0c115e6fcc7c39eed6ee663a95 Mon Sep 17 00:00:00 2001 From: Zach Musgrave Date: Mon, 14 Apr 2025 16:20:19 -0700 Subject: [PATCH 3/5] more test --- testing/go/expressions_test.go | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/testing/go/expressions_test.go b/testing/go/expressions_test.go index e605bf0007..103cf35d8c 100644 --- a/testing/go/expressions_test.go +++ b/testing/go/expressions_test.go @@ -318,6 +318,14 @@ func TestSubscript(t *testing.T) { Query: `SELECT ARRAY[1, 2, 3][4];`, Expected: []sql.Row{{nil}}, }, + { + Query: `SELECT ARRAY[1, 2, 3][null];`, + Expected: []sql.Row{{nil}}, + }, + { + Query: `SELECT ARRAY[1, 2, 3][1:3];`, + ExpectedErr: "not yet supported", + }, { Query: `SELECT ARRAY[1, 2, 3]['abc'];`, ExpectedErr: "integer: unhandled type: string", From 73b2ce5fee8770c35ac7876721b6685797df9544 Mon Sep 17 00:00:00 2001 From: Zach Musgrave Date: Mon, 14 Apr 2025 16:28:59 -0700 Subject: [PATCH 4/5] formatting --- server/ast/expr.go | 2 +- server/expression/subscript.go | 3 ++- testing/go/expressions_test.go | 2 +- testing/go/functions_test.go | 10 +++++----- 4 files changed, 9 insertions(+), 8 deletions(-) diff --git a/server/ast/expr.go b/server/ast/expr.go index 063c8b2d2d..f4d8075877 100644 --- a/server/ast/expr.go +++ b/server/ast/expr.go @@ -617,7 +617,7 @@ func nodeExpr(ctx *Context, node tree.Expr) (vitess.Expr, error) { } else if node.Indirection[0].Slice { return nil, errors.Errorf("slice subscripts are not yet supported") } - + indexExpr, err := nodeExpr(ctx, node.Indirection[0].Begin) if err != nil { return nil, err diff --git a/server/expression/subscript.go b/server/expression/subscript.go index 7d679c26a9..185bad2131 100755 --- a/server/expression/subscript.go +++ b/server/expression/subscript.go @@ -17,9 +17,10 @@ package expression import ( "fmt" - "github.com/dolthub/doltgresql/server/types" "github.com/dolthub/go-mysql-server/sql" vitess "github.com/dolthub/vitess/go/vt/sqlparser" + + "github.com/dolthub/doltgresql/server/types" ) type Subscript struct { diff --git a/testing/go/expressions_test.go b/testing/go/expressions_test.go index 103cf35d8c..75fbfd9268 100644 --- a/testing/go/expressions_test.go +++ b/testing/go/expressions_test.go @@ -323,7 +323,7 @@ func TestSubscript(t *testing.T) { Expected: []sql.Row{{nil}}, }, { - Query: `SELECT ARRAY[1, 2, 3][1:3];`, + Query: `SELECT ARRAY[1, 2, 3][1:3];`, ExpectedErr: "not yet supported", }, { diff --git a/testing/go/functions_test.go b/testing/go/functions_test.go index d2741d268e..70aa685919 100644 --- a/testing/go/functions_test.go +++ b/testing/go/functions_test.go @@ -2320,26 +2320,26 @@ func TestSelectFromFunctions(t *testing.T) { }, Assertions: []ScriptTestAssertion{ { - Query: `select * from dolt_branch('newBranch')`, + Query: `select * from dolt_branch('newBranch')`, Expected: []sql.Row{{0}}, }, { - Query: `select status from dolt_checkout('newBranch')`, + Query: `select status from dolt_checkout('newBranch')`, Expected: []sql.Row{{0}}, }, { Query: `insert into test values (6, 6, 'f')`, }, { - Query: `select length(commit_hash) > 0 from (select commit_hash from dolt_commit('-Am', 'added f') as result)`, + Query: `select length(commit_hash) > 0 from (select commit_hash from dolt_commit('-Am', 'added f') as result)`, Expected: []sql.Row{{"t"}}, }, { - Query: "select dolt_checkout('main')", + Query: "select dolt_checkout('main')", Expected: []sql.Row{{0}}, }, { - Query: `select fast_forward, conflicts from dolt_merge('newBranch')`, + Query: `select fast_forward, conflicts from dolt_merge('newBranch')`, Expected: []sql.Row{{"t", 0}}, }, }, From 7819f7960701deba04ad2eb4ee6d990d96448a8d Mon Sep 17 00:00:00 2001 From: Zach Musgrave Date: Mon, 14 Apr 2025 17:10:17 -0700 Subject: [PATCH 5/5] More tests and comments --- server/expression/subscript.go | 9 +++++++++ testing/go/expressions_test.go | 33 +++++++++++++++++++++++++++++++++ 2 files changed, 42 insertions(+) diff --git a/server/expression/subscript.go b/server/expression/subscript.go index 185bad2131..08961018d8 100755 --- a/server/expression/subscript.go +++ b/server/expression/subscript.go @@ -23,6 +23,7 @@ import ( "github.com/dolthub/doltgresql/server/types" ) +// Subscript represents a subscript expression, e.g. `a[1]`. type Subscript struct { Child sql.Expression Index sql.Expression @@ -31,6 +32,7 @@ type Subscript struct { var _ vitess.Injectable = (*Subscript)(nil) var _ sql.Expression = (*Subscript)(nil) +// NewSubscript creates a new Subscript expression. func NewSubscript(child, index sql.Expression) *Subscript { return &Subscript{ Child: child, @@ -38,14 +40,17 @@ func NewSubscript(child, index sql.Expression) *Subscript { } } +// Resolved implements the sql.Expression interface. func (s Subscript) Resolved() bool { return s.Child.Resolved() && s.Index.Resolved() } +// String implements the sql.Expression interface. func (s Subscript) String() string { return fmt.Sprintf("%s[%s]", s.Child, s.Index) } +// Type implements the sql.Expression interface. func (s Subscript) Type() sql.Type { dt, ok := s.Child.Type().(*types.DoltgresType) if !ok { @@ -54,6 +59,7 @@ func (s Subscript) Type() sql.Type { return dt.ArrayBaseType() } +// IsNullable implements the sql.Expression interface. func (s Subscript) IsNullable() bool { return true } @@ -97,10 +103,12 @@ func (s Subscript) Eval(ctx *sql.Context, row sql.Row) (interface{}, error) { } } +// Children implements the sql.Expression interface. func (s Subscript) Children() []sql.Expression { return []sql.Expression{s.Child, s.Index} } +// WithChildren implements the sql.Expression interface. func (s Subscript) WithChildren(children ...sql.Expression) (sql.Expression, error) { if len(children) != 2 { return nil, fmt.Errorf("expected 2 children, got %d", len(children)) @@ -108,6 +116,7 @@ func (s Subscript) WithChildren(children ...sql.Expression) (sql.Expression, err return NewSubscript(children[0], children[1]), nil } +// WithResolvedChildren implements the vitess.Injectable interface. func (s Subscript) WithResolvedChildren(children []any) (any, error) { if len(children) != 2 { return nil, fmt.Errorf("expected 2 children, got %d", len(children)) diff --git a/testing/go/expressions_test.go b/testing/go/expressions_test.go index 75fbfd9268..2e2f50add5 100644 --- a/testing/go/expressions_test.go +++ b/testing/go/expressions_test.go @@ -293,6 +293,9 @@ func TestBinaryLogic(t *testing.T) { }) } +// Note that our parser is more forgiving of array subscripts than the actual Postgres parser. +// We can handle this: SELECT ARRAY[1, 2, 3][1] +// But postgres requires: SELECT (ARRAY[1, 2, 3])[1] func TestSubscript(t *testing.T) { RunScripts(t, []ScriptTest{ { @@ -322,6 +325,10 @@ func TestSubscript(t *testing.T) { Query: `SELECT ARRAY[1, 2, 3][null];`, Expected: []sql.Row{{nil}}, }, + { + Query: `SELECT ARRAY['a', 'b', 'c'][2];`, + Expected: []sql.Row{{"b"}}, + }, { Query: `SELECT ARRAY[1, 2, 3][1:3];`, ExpectedErr: "not yet supported", @@ -332,5 +339,31 @@ func TestSubscript(t *testing.T) { }, }, }, + { + Name: "array column", + SetUpScript: []string{ + `CREATE TABLE test (id INT, arr INT[]);`, + `INSERT INTO test VALUES (1, ARRAY[1, 2, 3]), (2, ARRAY[4, 5, 6]);`, + }, + Assertions: []ScriptTestAssertion{ + { + Query: `SELECT arr[2] FROM test order by 1;`, + Expected: []sql.Row{{2}, {5}}, + }, + }, + }, + { + Name: "array subquery", + SetUpScript: []string{ + "CREATE TABLE test (id INT);", + "INSERT INTO test VALUES (1), (2), (3);", + }, + Assertions: []ScriptTestAssertion{ + { + Query: `SELECT (array(select id from test order by 1))[2]`, + Expected: []sql.Row{{2}}, + }, + }, + }, }) }