chore: import upstream snapshot with attribution

2026-07-13 12:32:25 +08:00
commit e014feafe1
2285 changed files with 1131979 additions and 0 deletions
@@ -0,0 +1,231 @@
+// Copyright 2023 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 ast
+
+import (
+	"fmt"
+	"strings"
+
+	"github.com/cockroachdb/errors"
+	vitess "github.com/dolthub/vitess/go/vt/sqlparser"
+
+	"github.com/dolthub/doltgresql/postgres/parser/sem/tree"
+	"github.com/dolthub/doltgresql/server/auth"
+)
+
+// nodeSelect handles *tree.Select nodes.
+func nodeSelect(ctx *Context, node *tree.Select) (vitess.SelectStatement, error) {
+	if node == nil {
+		return nil, nil
+	}
+	if node.Select == nil {
+		node.Select = &tree.ValuesClause{
+			Rows: []tree.Exprs{},
+		}
+	}
+	selectStmt, err := nodeSelectStatement(ctx, node.Select)
+	if err != nil {
+		return nil, err
+	}
+	orderBy, err := nodeOrderBy(ctx, node.OrderBy)
+	if err != nil {
+		return nil, err
+	}
+	with, err := nodeWith(ctx, node.With)
+	if err != nil {
+		return nil, err
+	}
+	limit, err := nodeLimit(ctx, node.Limit)
+	if err != nil {
+		return nil, err
+	}
+	_, err = nodeLockingClause(ctx, node.Locking)
+	if err != nil {
+		return nil, err
+	}
+
+	switch selectStmt := selectStmt.(type) {
+	case *vitess.ParenSelect:
+		return selectStmt, nil
+	case *vitess.Select:
+		selectStmt.OrderBy = orderBy
+		selectStmt.With = with
+		selectStmt.Limit = limit
+		return selectStmt, nil
+	case *vitess.SetOp:
+		selectStmt.OrderBy = orderBy
+		selectStmt.With = with
+		selectStmt.Limit = limit
+		return selectStmt, nil
+	default:
+		return nil, errors.Errorf("SELECT has encountered an unknown clause: `%T`", selectStmt)
+	}
+}
+
+// nodeSelectStatement handles tree.SelectStatement nodes.
+func nodeSelectStatement(ctx *Context, node tree.SelectStatement) (vitess.SelectStatement, error) {
+	if node == nil {
+		return nil, nil
+	}
+	ctx.Auth().PushAuthType(auth.AuthType_SELECT)
+	defer ctx.Auth().PopAuthType()
+
+	switch node := node.(type) {
+	case *tree.ParenSelect:
+		return nodeParenSelect(ctx, node)
+	case *tree.SelectClause:
+		return nodeSelectClause(ctx, node)
+	case *tree.UnionClause:
+		return nodeUnionClause(ctx, node)
+	case *tree.ValuesClause:
+		return nodeValuesClause(ctx, node)
+	default:
+		return nil, errors.Errorf("unknown type of SELECT statement: `%T`", node)
+	}
+}
+
+// nodeSelectExpr handles tree.SelectExpr nodes.
+func nodeSelectExpr(ctx *Context, node tree.SelectExpr) (vitess.SelectExpr, error) {
+	switch expr := node.Expr.(type) {
+	case *tree.AllColumnsSelector:
+		if expr.TableName.NumParts > 1 {
+			return nil, errors.Errorf("referencing items outside the schema or database is not yet supported")
+		}
+		return &vitess.StarExpr{
+			TableName: vitess.TableName{
+				Name: vitess.NewTableIdent(expr.TableName.Parts[0]),
+			},
+		}, nil
+	case tree.UnqualifiedStar:
+		return &vitess.StarExpr{}, nil
+	case *tree.UnresolvedName:
+		colName, err := unresolvedNameToColName(expr)
+		if err != nil {
+			return nil, err
+		}
+
+		if expr.Star {
+			return &vitess.StarExpr{
+				TableName: colName.Qualifier,
+			}, nil
+		}
+
+		// We don't set the InputExpression for ColName expressions. This matches the behavior in vitess's
+		// post-processing found in ast.go. Input expressions are load bearing for some parts of plan building
+		// so we need to match the behavior exactly.
+		return &vitess.AliasedExpr{
+			Expr: colName,
+			As:   vitess.NewColIdent(string(node.As)),
+		}, nil
+	default:
+		vitessExpr, err := nodeExpr(ctx, expr)
+		if err != nil {
+			return nil, err
+		}
+
+		if ce, ok := expr.(*tree.CastExpr); ok && node.As == "" {
+			hasConst := false
+			_, _ = tree.SimpleVisit(expr, func(visitingExpr tree.Expr) (recurse bool, newExpr tree.Expr, err error) {
+				switch visitingExpr.(type) {
+				case tree.Constant:
+					hasConst = true
+					return false, visitingExpr, nil
+				}
+				return true, visitingExpr, nil
+			})
+			if hasConst {
+				_, dt, err := nodeResolvableTypeReference(ctx, ce.Type, false)
+				if err != nil {
+					return nil, err
+				}
+				// constant value is not part of column name
+				// e.g. `1::INT2` should create column name as `int2`.
+				node.As = tree.UnrestrictedName(dt.Name())
+			} else {
+				// cast type is not part of column name
+				// e.g. `id::INT2` should create column name as `id`.
+				node.As = tree.UnrestrictedName(tree.AsString(ce.Expr))
+			}
+		}
+
+		return &vitess.AliasedExpr{
+			Expr:            vitessExpr,
+			As:              vitess.NewColIdent(string(node.As)),
+			InputExpression: inputExpressionForSelectExpr(node),
+		}, nil
+	}
+}
+
+// inputExpressionForSelectExpr returns the input expression for a tree.SelectExpr.
+// Postgres has specific handling for function calls that differs from the default printing behavior.
+func inputExpressionForSelectExpr(node tree.SelectExpr) string {
+	inputExpression := tree.AsStringWithFlags(&node, tree.FmtOmitFunctionArgs)
+	// To be consistent with vitess handling, InputExpression always gets its outer quotes trimmed
+	if strings.HasPrefix(inputExpression, "'") && strings.HasSuffix(inputExpression, "'") {
+		inputExpression = inputExpression[1 : len(inputExpression)-1]
+	}
+	return inputExpression
+}
+
+// nodeSelectExprs handles tree.SelectExprs nodes.
+func nodeSelectExprs(ctx *Context, node tree.SelectExprs) (vitess.SelectExprs, error) {
+	if len(node) == 0 {
+		return nil, nil
+	}
+	selectExprs := make(vitess.SelectExprs, len(node))
+	for i := range node {
+		var err error
+		// Unaliased string literals use Postgres's ?column? placeholder as the column name.
+		// We use a unique indexed sentinel rather than "?column?" directly to avoid GMS's alias
+		// scope map conflating multiple unnamed columns in the same SELECT clause; the handler
+		// translates the sentinel back to "?column?" on the wire.
+		if _, isStrVal := node[i].Expr.(*tree.StrVal); isStrVal && node[i].As == "" {
+			node[i].As = tree.UnrestrictedName(fmt.Sprintf("%s%d", UnknownColSentinelPrefix, i))
+		}
+		selectExprs[i], err = nodeSelectExpr(ctx, node[i])
+		if err != nil {
+			return nil, err
+		}
+	}
+	return selectExprs, nil
+}
+
+// nodeExprToSelectExpr handles tree.Expr nodes and returns the result as a vitess.SelectExpr.
+func nodeExprToSelectExpr(ctx *Context, node tree.Expr) (vitess.SelectExpr, error) {
+	if node == nil {
+		return nil, nil
+	}
+	return nodeSelectExpr(ctx, tree.SelectExpr{
+		Expr: node,
+	})
+}
+
+// nodeExprsToSelectExprs handles tree.Exprs nodes and returns the results as vitess.SelectExprs.
+func nodeExprsToSelectExprs(ctx *Context, node tree.Exprs) (vitess.SelectExprs, error) {
+	if len(node) == 0 {
+		return nil, nil
+	}
+	selectExprs := make(vitess.SelectExprs, len(node))
+	for i := range node {
+		var err error
+		selectExprs[i], err = nodeSelectExpr(ctx, tree.SelectExpr{
+			Expr: node[i],
+		})
+		if err != nil {
+			return nil, err
+		}
+	}
+	return selectExprs, nil
+}