more function support

server/ast/aliased_table_expr.go

@@ -15,8 +15,9 @@
 package ast
 
 import (
-	"github.com/cockroachdb/errors"
+	"strings"
 
+	"github.com/cockroachdb/errors"
 	vitess "github.com/dolthub/vitess/go/vt/sqlparser"
 
 	"github.com/dolthub/doltgresql/postgres/parser/sem/tree"
@@ -26,7 +27,9 @@ import (
 // nodeAliasedTableExpr handles *tree.AliasedTableExpr nodes.
 func nodeAliasedTableExpr(ctx *Context, node *tree.AliasedTableExpr) (*vitess.AliasedTableExpr, error) {
 	if node.Ordinality {
-		return nil, errors.Errorf("ordinality is not yet supported")
+		if _, ok := node.Expr.(*tree.RowsFromExpr); !ok {
+			return nil, errors.Errorf("WITH ORDINALITY is only supported for functions")
+		}
 	}
 	if node.IndexFlags != nil {
 		return nil, errors.Errorf("index flags are not yet supported")
@@ -101,16 +104,51 @@ func nodeAliasedTableExpr(ctx *Context, node *tree.AliasedTableExpr) (*vitess.Al
 		}
 		aliasExpr = subquery
 	case *tree.RowsFromExpr:
-		tableExpr, err := nodeTableExpr(ctx, expr)
-		if err != nil {
-			return nil, err
-		}
+		var selectStmt vitess.SelectStatement
+		if node.Ordinality {
+			// WITH ORDINALITY appends a bigint column numbering the function's result rows, named
+			// "ordinality" unless renamed by a column alias list. The numbering projection has to
+			// live one level above the function's expansion, so we expand the function in the
+			// select list of a wrapped subquery.
+			items, err := nodeExprs(ctx, expr.Items)
+			if err != nil {
+				return nil, err
+			}
+			innerExprs := make(vitess.SelectExprs, len(items))
+			for i := range items {
+				innerExprs[i] = &vitess.AliasedExpr{Expr: items[i]}
+			}
+			selectStmt = &vitess.Select{
+				SelectExprs: vitess.SelectExprs{
+					&vitess.StarExpr{},
+					&vitess.AliasedExpr{
+						Expr: &vitess.FuncExpr{
+							Name: vitess.NewColIdent("row_number"),
+							Over: &vitess.Over{},
+						},
+						As: vitess.NewColIdent("ordinality"),
+					},
+				},
+				From: vitess.TableExprs{
+					&vitess.AliasedTableExpr{
+						Expr: &vitess.Subquery{Select: &vitess.Select{SelectExprs: innerExprs}},
+						As:   vitess.NewTableIdent("with_ordinality"),
+					},
+				},
+			}
+		} else {
+			tableExpr, err := nodeTableExpr(ctx, expr)
+			if err != nil {
+				return nil, err
+			}
 
-		// TODO: this should be represented as a table function more directly
-		subquery := &vitess.Subquery{
-			Select: &vitess.Select{
+			// TODO: this should be represented as a table function more directly
+			selectStmt = &vitess.Select{
 				From: vitess.TableExprs{tableExpr},
-			},
+			}
+		}
+		subquery := &vitess.Subquery{
+			Select: selectStmt,
 		}
 
 		if len(node.As.Cols) > 0 {
@@ -125,6 +163,17 @@ func nodeAliasedTableExpr(ctx *Context, node *tree.AliasedTableExpr) (*vitess.Al
 		return nil, errors.Errorf("unhandled table expression: `%T`", expr)
 	}
 	alias := string(node.As.Alias)
+	if alias == "" && node.Ordinality {
+		// A derived table needs an alias; the implicit alias of a function called in FROM is the
+		// function's name, matching Postgres
+		alias = "with_ordinality"
+		if rf, ok := node.Expr.(*tree.RowsFromExpr); ok && len(rf.Items) == 1 {
+			if fe, ok := rf.Items[0].(*tree.FuncExpr); ok {
+				nameParts := strings.Split(fe.Func.String(), ".")
+				alias = strings.ToLower(nameParts[len(nameParts)-1])
+			}
+		}
+	}
 
 	var asOf *vitess.AsOf
 	if node.AsOf != nil {

server/connection_handler.go

@@ -493,11 +493,6 @@ func (h *ConnectionHandler) handleMessage(msg pgproto3.Message) (stop, endOfMess
 // expected as part of this query, in which case the server will send a READY FOR QUERY message back to the client so
 // that it can send its next query.
 func (h *ConnectionHandler) handleQuery(message *pgproto3.Query) (endOfMessages bool, err error) {
-	handled, err := h.handledPSQLCommands(message.String)
-	if handled || err != nil {
-		return true, err
-	}
-
 	queries, err := h.convertQuery(message.String)
 	if err != nil {
 		if printErrorStackTraces {
@@ -510,6 +505,7 @@ func (h *ConnectionHandler) handleQuery(message *pgproto3.Query) (endOfMessages
 	delete(h.preparedStatements, "")
 	delete(h.portals, "")
 
+	var handled bool
 	if len(queries) == 1 {
 		// empty query special case
 		if queries[0].AST == nil {
@@ -1379,25 +1375,6 @@ func (h *ConnectionHandler) sendDescribeResponse(fields []pgproto3.FieldDescript
 	}
 }
 
-// handledPSQLCommands handles the special PSQL commands, such as \l and \dt.
-func (h *ConnectionHandler) handledPSQLCommands(statement string) (bool, error) {
-	statement = strings.ToLower(statement)
-	// Command: \d table_name
-	if strings.HasPrefix(statement, "select c.oid,\n  n.nspname,\n  c.relname\nfrom pg_catalog.pg_class c\n     left join pg_catalog.pg_namespace n on n.oid = c.relnamespace\nwhere c.relname operator(pg_catalog.~) '^(") && strings.HasSuffix(statement, ")$' collate pg_catalog.default\n  and pg_catalog.pg_table_is_visible(c.oid)\norder by 2, 3;") {
-		// There are >at least< 15 separate statements sent for this command, which is far too much to validate and
-		// implement, so we'll just return an error for now
-		return true, errors.Errorf("PSQL command not yet supported")
-	}
-	// Command: \df
-	if statement == "select n.nspname as \"schema\",\n  p.proname as \"name\",\n  pg_catalog.pg_get_function_result(p.oid) as \"result data type\",\n  pg_catalog.pg_get_function_arguments(p.oid) as \"argument data types\",\n case p.prokind\n  when 'a' then 'agg'\n  when 'w' then 'window'\n  when 'p' then 'proc'\n  else 'func'\n end as \"type\"\nfrom pg_catalog.pg_proc p\n     left join pg_catalog.pg_namespace n on n.oid = p.pronamespace\nwhere pg_catalog.pg_function_is_visible(p.oid)\n      and n.nspname <> 'pg_catalog'\n      and n.nspname <> 'information_schema'\norder by 1, 2, 4;" {
-		return true, h.query(ConvertedQuery{
-			String:       `SELECT '' AS "Schema", '' AS "Name", '' AS "Result data type", '' AS "Argument data types", '' AS "Type" LIMIT 0;`,
-			StatementTag: "SELECT",
-		})
-	}
-	return false, nil
-}
-
 // endOfMessages should be called from HandleConnection or a function within HandleConnection. This represents the end
 // of the message slice, which may occur naturally (all relevant response messages have been sent) or on error. Once
 // endOfMessages has been called, no further messages should be sent, and the connection loop should wait for the next

server/functions/binary/equal.go

@@ -499,7 +499,12 @@ var oideq = framework.Function2{
 		// This method doesn't use DoltgresType.Compare because it's on the critical path for many tooling queries that
 		// examine the pg_catalog tables.
 		val1id, val2id := val1.(id.Id), val2.(id.Id)
-		return val1id == val2id, nil
+		if val1id == val2id {
+			return true, nil
+		}
+		// Different internal IDs can still map to the same OID: an OID given to us by a client resolves to a raw
+		// numeric ID unless its assignment is already cached, and a raw OID of 0 is a distinct value from id.Null.
+		return id.Cache().ToOID(val1id) == id.Cache().ToOID(val2id), nil
 	},
 }
 

server/functions/init.go

@@ -169,6 +169,7 @@ func Init() {
 	initPgOpclassIsVisible()
 	initPgOperatorIsVisible()
 	initPgOpfamilyIsVisible()
+	initPgPartitionAncestors()
 	initPgPostmasterStartTime()
 	initPgRelationIsPublishable()
 	initPgRelationSize()

server/functions/oid.go

@@ -114,6 +114,8 @@ var btoidcmp = framework.Function2{
 	Parameters: [2]*pgtypes.DoltgresType{pgtypes.Oid, pgtypes.Oid},
 	Strict:     true,
 	Callable: func(ctx *sql.Context, _ [3]*pgtypes.DoltgresType, val1, val2 any) (any, error) {
-		return int32(cmp.Compare(val1.(id.Id), val2.(id.Id))), nil
+		// Postgres compares OIDs numerically, and different internal IDs can map to the same OID (e.g. a raw
+		// numeric OID given by a client and the cached ID it refers to)
+		return int32(cmp.Compare(id.Cache().ToOID(val1.(id.Id)), id.Cache().ToOID(val2.(id.Id)))), nil
 	},
 }

server/functions/pg_function_is_visible.go

@@ -17,8 +17,8 @@ package functions
 import (
 	"github.com/dolthub/go-mysql-server/sql"
 
+	"github.com/dolthub/doltgresql/core"
 	"github.com/dolthub/doltgresql/core/id"
-
 	"github.com/dolthub/doltgresql/server/functions/framework"
 	pgtypes "github.com/dolthub/doltgresql/server/types"
 )
@@ -36,7 +36,45 @@ var pg_function_is_visible_oid = framework.Function1{
 	IsNonDeterministic: true,
 	Strict:             true,
 	Callable: func(ctx *sql.Context, _ [2]*pgtypes.DoltgresType, val any) (any, error) {
-		// TODO: Functions are not contained within a schema for now, so will be true if function is found
-		return id.Cache().Exists(val.(id.Id)), nil
+		oidVal := val.(id.Id)
+		if oidVal.Section() != id.Section_Function && oidVal.Section() != id.Section_Procedure {
+			return false, nil
+		}
+		// TODO: Postgres additionally checks that the function isn't shadowed by one with the same
+		// name and argument types in an earlier schema of the search path.
+		paths, err := core.SearchPath(ctx)
+		if err != nil {
+			return false, err
+		}
+
+		// This is either id.Function or id.Procedure, so the first segment is the schema name
+		schemaName := oidVal.Segment(0)
+		inPath := false
+		for _, path := range paths {
+			if path == schemaName {
+				inPath = true
+				break
+			}
+		}
+		if !inPath {
+			return false, nil
+		}
+		// Built-in functions are registered in the OID cache at startup, and don't appear in the
+		// function and procedure collections that RunCallback iterates over.
+		if id.Cache().Exists(oidVal) {
+			return true, nil
+		}
+		isVisible := false
+		err = RunCallback(ctx, oidVal, Callbacks{
+			Function: func(ctx *sql.Context, schema ItemSchema, function ItemFunction) (cont bool, err error) {
+				isVisible = true
+				return false, nil
+			},
+			Procedure: func(ctx *sql.Context, schema ItemSchema, procedure ItemProcedure) (cont bool, err error) {
+				isVisible = true
+				return false, nil
+			},
+		})
+		return isVisible, err
 	},
 }

server/functions/pg_get_indexdef.go

@@ -18,7 +18,6 @@ import (
 	"fmt"
 	"strings"
 
-	"github.com/cockroachdb/errors"
 	"github.com/dolthub/go-mysql-server/sql"
 	"github.com/dolthub/go-mysql-server/sql/plan"
 
@@ -59,12 +58,29 @@ var pg_get_indexdef_oid = framework.Function1{
 // buildIndexDef generates a CREATE INDEX DDL statement for the given index.
 func buildIndexDef(ctx *sql.Context, index sql.Index, table sql.Table, schemaName string) string {
 	name := index.ID()
+	if name == "PRIMARY" {
+		// Primary key indexes are displayed with their postgres-convention name, matching pg_class
+		name = fmt.Sprintf("%s_pkey", index.Table())
+	}
 	using := strings.ToLower(index.IndexType())
 	unique := ""
 	if index.IsUnique() {
 		unique = " UNIQUE"
 	}
 
+	colsStr := strings.Join(indexColumnExprs(ctx, index, table), ", ")
+
+	def := fmt.Sprintf("CREATE%s INDEX %s ON %s.%s USING %s (%s)", unique, name, schemaName, index.Table(), using, colsStr)
+	if pi, ok := index.(sql.PartialIndex); ok && pi.Predicate() != "" {
+		def += " WHERE (" + pi.Predicate() + ")"
+	}
+	return def
+}
+
+// indexColumnExprs returns the rendered text of each column of the given index, in index column
+// order. Plain columns render as the bare column name, functional expressions as their original
+// SQL text.
+func indexColumnExprs(ctx *sql.Context, index sql.Index, table sql.Table) []string {
 	cols := make([]string, len(index.Expressions()))
 	for i, expr := range index.Expressions() {
 		if exprText, ok := RenderHiddenIndexColumnExpr(plan.GetColumnFromIndexExpr(ctx, expr, table)); ok {
@@ -79,13 +95,7 @@ func buildIndexDef(ctx *sql.Context, index sql.Index, table sql.Table, schemaNam
 			cols[i] = expr
 		}
 	}
-	colsStr := strings.Join(cols, ", ")
-
-	def := fmt.Sprintf("CREATE%s INDEX %s ON %s.%s USING %s (%s)", unique, name, schemaName, index.Table(), using, colsStr)
-	if pi, ok := index.(sql.PartialIndex); ok && pi.Predicate() != "" {
-		def += " WHERE (" + pi.Predicate() + ")"
-	}
-	return def
+	return cols
 }
 
 // RenderHiddenIndexColumnExpr returns the original SQL text of the functional expression backing
@@ -112,23 +122,27 @@ var pg_get_indexdef_oid_integer_bool = framework.Function3{
 	Callable: func(ctx *sql.Context, _ [4]*pgtypes.DoltgresType, val1, val2, val3 any) (any, error) {
 		oidVal := val1.(id.Id)
 		colNo := val2.(int32)
-		pretty := val3.(bool)
-		if pretty {
-			return "", errors.Errorf("pretty printing is not yet supported")
-		}
+		// The pretty flag only affects the formatting of expressions, which we don't reproduce, so
+		// we return the same text either way.
+		result := ""
 		err := RunCallback(ctx, oidVal, Callbacks{
 			Index: func(ctx *sql.Context, schema ItemSchema, table ItemTable, index ItemIndex) (cont bool, err error) {
-				exprs := index.Item.Expressions()
-				if int(colNo) >= len(exprs) {
-					return false, errors.Errorf("column not found")
+				if colNo == 0 {
+					result = buildIndexDef(ctx, index.Item, table.Item, schema.Item.SchemaName())
+					return false, nil
+				}
+				// A non-zero column number selects just that column's definition, or an empty
+				// string if the index has no such column.
+				cols := indexColumnExprs(ctx, index.Item, table.Item)
+				if colNo >= 1 && int(colNo) <= len(cols) {
+					result = cols[colNo-1]
 				}
-				// TODO: make `create index` statement
 				return false, nil
 			},
 		})
 		if err != nil {
 			return "", err
 		}
-		return "", nil
+		return result, nil
 	},
 }

server/functions/pg_get_triggerdef.go

@@ -15,10 +15,12 @@
 package functions
 
 import (
-	"github.com/cockroachdb/errors"
+	"strings"
 
 	"github.com/dolthub/go-mysql-server/sql"
 
+	"github.com/dolthub/doltgresql/core"
+	"github.com/dolthub/doltgresql/core/id"
 	"github.com/dolthub/doltgresql/server/functions/framework"
 	pgtypes "github.com/dolthub/doltgresql/server/types"
 )
@@ -37,8 +39,7 @@ var pg_get_triggerdef_oid = framework.Function1{
 	IsNonDeterministic: true,
 	Strict:             true,
 	Callable: func(ctx *sql.Context, _ [2]*pgtypes.DoltgresType, val any) (any, error) {
-		// TODO: triggers are not supported yet
-		return "", nil
+		return getTriggerDef(ctx, val.(id.Id))
 	},
 }
 
@@ -50,11 +51,29 @@ var pg_get_triggerdef_oid_bool = framework.Function2{
 	IsNonDeterministic: true,
 	Strict:             true,
 	Callable: func(ctx *sql.Context, _ [3]*pgtypes.DoltgresType, val1, val2 any) (any, error) {
-		pretty := val2.(bool)
-		if pretty {
-			return "", errors.Errorf("pretty printing is not yet supported")
-		}
-		// TODO: triggers are not supported yet
-		return "", nil
+		// The pretty flag only affects the formatting of expressions, which we don't reproduce, so
+		// we return the same text either way.
+		return getTriggerDef(ctx, val1.(id.Id))
 	},
 }
+
+// getTriggerDef returns the definition of the trigger for the given OID, or an empty string if the OID doesn't refer
+// to an existing trigger. Postgres reconstructs the statement from the catalog, whereas we return the statement that
+// created the trigger.
+func getTriggerDef(ctx *sql.Context, oidVal id.Id) (string, error) {
+	if oidVal.Section() != id.Section_Trigger {
+		return "", nil
+	}
+	collection, err := core.GetTriggersCollectionFromContext(ctx, ctx.GetCurrentDatabase())
+	if err != nil {
+		return "", err
+	}
+	trigger, err := collection.GetTrigger(ctx, id.Trigger(oidVal))
+	if err != nil {
+		return "", err
+	}
+	if !trigger.ID.IsValid() {
+		return "", nil
+	}
+	return strings.TrimSuffix(strings.TrimSpace(trigger.Definition), ";"), nil
+}