136 lines
3.6 KiB
Go
Executable File
136 lines
3.6 KiB
Go
Executable File
// 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 (
|
|
"context"
|
|
"fmt"
|
|
|
|
"github.com/dolthub/go-mysql-server/sql"
|
|
vitess "github.com/dolthub/vitess/go/vt/sqlparser"
|
|
|
|
"github.com/dolthub/doltgresql/server/types"
|
|
)
|
|
|
|
// Subscript represents a subscript expression, e.g. `a[1]`.
|
|
type Subscript struct {
|
|
Child sql.Expression
|
|
Index sql.Expression
|
|
}
|
|
|
|
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,
|
|
Index: index,
|
|
}
|
|
}
|
|
|
|
// 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(ctx *sql.Context) sql.Type {
|
|
dt, ok := s.Child.Type(ctx).(*types.DoltgresType)
|
|
if !ok {
|
|
panic(fmt.Sprintf("unexpected type %T for subscript", s.Child.Type(ctx)))
|
|
}
|
|
return dt.ArrayBaseType()
|
|
}
|
|
|
|
// IsNullable implements the sql.Expression interface.
|
|
func (s Subscript) IsNullable(ctx *sql.Context) 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)
|
|
}
|
|
}
|
|
|
|
// 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(ctx *sql.Context, 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
|
|
}
|
|
|
|
// WithResolvedChildren implements the vitess.Injectable interface.
|
|
func (s Subscript) WithResolvedChildren(ctx context.Context, 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
|
|
}
|