Merge pull request #879 from dolthub/fulghum/nextval_regclass

Allow `nextval()` to take a regclass instance

Author: fulghum

server/functions/current_schema.go

@@ -18,6 +18,7 @@ import (
 	"github.com/dolthub/go-mysql-server/sql"
 
 	"github.com/dolthub/doltgresql/server/functions/framework"
+	"github.com/dolthub/doltgresql/server/settings"
 	pgtypes "github.com/dolthub/doltgresql/server/types"
 )
 
@@ -33,7 +34,7 @@ var current_schema = framework.Function0{
 	IsNonDeterministic: true,
 	Strict:             true,
 	Callable: func(ctx *sql.Context) (any, error) {
-		schemas, err := GetCurrentSchemas(ctx)
+		schemas, err := settings.GetCurrentSchemas(ctx)
 		if err != nil {
 			return nil, err
 		}

server/functions/current_schemas.go

@@ -15,12 +15,11 @@
 package functions
 
 import (
-	"strings"
-
 	"github.com/dolthub/go-mysql-server/sql"
 
 	"github.com/dolthub/doltgresql/postgres/parser/sessiondata"
 	"github.com/dolthub/doltgresql/server/functions/framework"
+	"github.com/dolthub/doltgresql/server/settings"
 	pgtypes "github.com/dolthub/doltgresql/server/types"
 )
 
@@ -41,7 +40,7 @@ var current_schemas = framework.Function1{
 		if val.(bool) {
 			schemas = append(schemas, sessiondata.PgCatalogName)
 		}
-		searchPaths, err := GetCurrentSchemas(ctx)
+		searchPaths, err := settings.GetCurrentSchemas(ctx)
 		if err != nil {
 			return nil, err
 		}
@@ -51,24 +50,3 @@ var current_schemas = framework.Function1{
 		return schemas, nil
 	},
 }
-
-// GetCurrentSchemas returns all the schemas in the search_path setting, with elements like "$user" excluded
-func GetCurrentSchemas(ctx *sql.Context) ([]string, error) {
-	searchPathVar, err := ctx.GetSessionVariable(ctx, "search_path")
-	if err != nil {
-		return nil, err
-	}
-
-	pathElems := strings.Split(searchPathVar.(string), ",")
-	var path []string
-
-	for _, schemaName := range pathElems {
-		schemaName = strings.Trim(schemaName, " ")
-		if schemaName == "\"$user\"" {
-			continue
-		}
-		path = append(path, schemaName)
-	}
-
-	return path, nil
-}

server/functions/nextval.go

@@ -25,9 +25,15 @@ import (
 // initNextVal registers the functions to the catalog.
 func initNextVal() {
 	framework.RegisterFunction(nextval_text)
+	framework.RegisterFunction(nextval_regclass)
 }
 
 // nextval_text represents the PostgreSQL function of the same name, taking the same parameters.
+//
+// TODO: Even though we can implicitly convert a text param to a regclass param, it's an expensive process
+// to convert it to a regclass, then convert the regclass back into the relation name, so we provide an overload
+// that takes a text param directly, in addition to the function form that takes a regclass. Once we can optimize
+// the regclass to text conversion, we can potentially remove this overload.
 var nextval_text = framework.Function1{
 	Name:               "nextval",
 	Return:             pgtypes.Int64,
@@ -47,3 +53,29 @@ var nextval_text = framework.Function1{
 		return collection.NextVal(schema, sequence)
 	},
 }
+
+// nextval_regclass represents the PostgreSQL function of the same name, taking the same parameters.
+var nextval_regclass = framework.Function1{
+	Name:               "nextval",
+	Return:             pgtypes.Int64,
+	Parameters:         [1]pgtypes.DoltgresType{pgtypes.Regclass},
+	IsNonDeterministic: true,
+	Strict:             true,
+	Callable: func(ctx *sql.Context, _ [2]pgtypes.DoltgresType, val any) (any, error) {
+		relationName, err := pgtypes.Regclass.IoOutput(ctx, val)
+		if err != nil {
+			return nil, err
+		}
+
+		schema, sequence, err := parseRelationName(ctx, relationName)
+		if err != nil {
+			return nil, err
+		}
+
+		collection, err := core.GetSequencesCollectionFromContext(ctx)
+		if err != nil {
+			return nil, err
+		}
+		return collection.NextVal(schema, sequence)
+	},
+}

server/settings/current_schemas.go

@@ -0,0 +1,56 @@
+// Copyright 2024 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 settings
+
+import (
+	"strings"
+
+	"github.com/dolthub/go-mysql-server/sql"
+)
+
+// GetCurrentSchemas returns all the schemas in the search_path setting, with elements like "$user" excluded
+func GetCurrentSchemas(ctx *sql.Context) ([]string, error) {
+	searchPathVar, err := ctx.GetSessionVariable(ctx, "search_path")
+	if err != nil {
+		return nil, err
+	}
+
+	pathElems := strings.Split(searchPathVar.(string), ",")
+	var path []string
+
+	for _, schemaName := range pathElems {
+		schemaName = strings.Trim(schemaName, " ")
+		if schemaName == "\"$user\"" {
+			continue
+		}
+		path = append(path, schemaName)
+	}
+
+	return path, nil
+}
+
+// GetCurrentSchemasAsMap returns the schemas from the search_path setting as a map for easy lookup. Any
+// elements like "$user" are excluded.
+func GetCurrentSchemasAsMap(ctx *sql.Context) (map[string]struct{}, error) {
+	schemas, err := GetCurrentSchemas(ctx)
+	if err != nil {
+		return nil, err
+	}
+	schemaMap := make(map[string]struct{}, len(schemas))
+	for _, schema := range schemas {
+		schemaMap[schema] = struct{}{}
+	}
+	return schemaMap, nil
+}

server/settings/doc.go

@@ -0,0 +1,19 @@
+// Copyright 2024 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 settings provides utility functions for working with settings, such as the search_path setting.
+//
+// This package is intended to be a leaf in the package dependency graph, and should not add dependencies to
+// other Doltgres packages.
+package settings

server/types/oid/regclass.go

@@ -20,6 +20,8 @@ import (
 
 	"github.com/dolthub/dolt/go/libraries/doltcore/sqle/resolve"
 	"github.com/dolthub/go-mysql-server/sql"
+
+	"github.com/dolthub/doltgresql/server/settings"
 )
 
 // regclass_IoInput is the implementation for IoInput that avoids circular dependencies by being declared in a separate
@@ -109,25 +111,56 @@ func regclass_IoInput(ctx *sql.Context, input string) (uint32, error) {
 // regclass_IoOutput is the implementation for IoOutput that avoids circular dependencies by being declared in a separate
 // package.
 func regclass_IoOutput(ctx *sql.Context, oid uint32) (string, error) {
+	// Find all the schemas on the search path. If a schema is on the search path, then it is not included in the
+	// output of relation name. If the relation's schema is not on the search path, then it is explicitly included.
+	schemasMap, err := settings.GetCurrentSchemasAsMap(ctx)
+	if err != nil {
+		return "", err
+	}
+
+	// The pg_catalog schema is always implicitly part of the search path
+	// https://www.postgresql.org/docs/current/ddl-schemas.html#DDL-SCHEMAS-CATALOG
+	schemasMap["pg_catalog"] = struct{}{}
+
 	output := strconv.FormatUint(uint64(oid), 10)
-	err := RunCallback(ctx, oid, Callbacks{
+	err = RunCallback(ctx, oid, Callbacks{
 		Index: func(ctx *sql.Context, schema ItemSchema, table ItemTable, index ItemIndex) (cont bool, err error) {
 			output = index.Item.ID()
 			if output == "PRIMARY" {
-				output = fmt.Sprintf("%s_pkey", index.Item.Table())
+				schemaName := schema.Item.SchemaName()
+				if _, ok := schemasMap[schemaName]; ok {
+					output = fmt.Sprintf("%s_pkey", index.Item.Table())
+				} else {
+					output = fmt.Sprintf("%s.%s_pkey", schemaName, index.Item.Table())
+				}
 			}
 			return false, nil
 		},
 		Sequence: func(ctx *sql.Context, schema ItemSchema, sequence ItemSequence) (cont bool, err error) {
-			output = sequence.Item.Name
+			schemaName := schema.Item.SchemaName()
+			if _, ok := schemasMap[schemaName]; ok {
+				output = sequence.Item.Name
+			} else {
+				output = fmt.Sprintf("%s.%s", schemaName, sequence.Item.Name)
+			}
 			return false, nil
 		},
 		Table: func(ctx *sql.Context, schema ItemSchema, table ItemTable) (cont bool, err error) {
-			output = table.Item.Name()
+			schemaName := schema.Item.SchemaName()
+			if _, ok := schemasMap[schemaName]; ok {
+				output = table.Item.Name()
+			} else {
+				output = fmt.Sprintf("%s.%s", schemaName, table.Item.Name())
+			}
 			return false, nil
 		},
 		View: func(ctx *sql.Context, schema ItemSchema, view ItemView) (cont bool, err error) {
-			output = view.Item.Name
+			schemaName := schema.Item.SchemaName()
+			if _, ok := schemasMap[schemaName]; ok {
+				output = view.Item.Name
+			} else {
+				output = fmt.Sprintf("%s.%s", schemaName, view.Item.Name)
+			}
 			return false, nil
 		},
 	})

server/types/text.go

@@ -106,6 +106,10 @@ func (b TextType) Convert(val any) (any, sql.ConvertInRange, error) {
 
 // Equals implements the DoltgresType interface.
 func (b TextType) Equals(otherType sql.Type) bool {
+	if _, ok := otherType.(TextType); !ok {
+		return false
+	}
+
 	if otherExtendedType, ok := otherType.(types.ExtendedType); ok {
 		return bytes.Equal(MustSerializeType(b), MustSerializeType(otherExtendedType))
 	}

testing/go/functions_test.go

@@ -202,6 +202,24 @@ func TestFunctionsOID(t *testing.T) {
 						{nil},
 					},
 				},
+				{
+					// When the relation is from a schema on the search path, it is not qualified with the schema name
+					Query: `SELECT to_regclass(('public.testing'::regclass)::text);`,
+					Expected: []sql.Row{
+						{"testing"},
+					},
+				},
+				{
+					// Clear out the current search_path setting to test fully qualified relation names
+					Query:    `SET search_path = '';`,
+					Expected: []sql.Row{},
+				},
+				{
+					Query: `SELECT to_regclass(('public.testing'::regclass)::text);`,
+					Expected: []sql.Row{
+						{"public.testing"},
+					},
+				},
 			},
 		},
 		{

testing/go/sequences_test.go

@@ -41,6 +41,14 @@ func TestSequences(t *testing.T) {
 					Query:    "SELECT nextval('test');",
 					Expected: []sql.Row{{3}},
 				},
+				{
+					Query:    "SELECT nextval('test'::regclass);",
+					Expected: []sql.Row{{4}},
+				},
+				{
+					Query:       "SELECT nextval('doesnotexist'::regclass);",
+					ExpectedErr: "does not exist",
+				},
 				{
 					Query:    "DROP SEQUENCE test;",
 					Expected: []sql.Row{},

testing/go/types_test.go

@@ -1798,6 +1798,7 @@ var typesTests = []ScriptTest{
 			`CREATE TABLE testing (pk INT primary key, v1 INT UNIQUE);`,
 			`CREATE TABLE "Testing2" (pk INT primary key, v1 INT);`,
 			`CREATE VIEW testview AS SELECT * FROM testing LIMIT 1;`,
+			`CREATE SEQUENCE seq1;`,
 		},
 		Assertions: []ScriptTestAssertion{
 			{
@@ -1838,6 +1839,12 @@ var typesTests = []ScriptTest{
 					{"testing"},
 				},
 			},
+			{
+				Query: `SELECT 'seq1'::regclass;`,
+				Expected: []sql.Row{
+					{"seq1"},
+				},
+			},
 			{
 				Query:       `SELECT 'Testing2'::regclass;`,
 				ExpectedErr: "does not exist",
@@ -1860,6 +1867,29 @@ var typesTests = []ScriptTest{
 					{"testing"},
 				},
 			},
+			{
+				// schema-qualified relation names are not returned if the schema is on the search path
+				Query: `SELECT 'public.testing'::regclass, 'public.seq1'::regclass, 'public.testview'::regclass, 'public.testing_pkey'::regclass;`,
+				Expected: []sql.Row{
+					{"testing", "seq1", "testview", "testing_pkey"},
+				},
+			},
+			{
+				// Clear out the current search_path setting to test schema-qualified relation names
+				Query:    `SET search_path = '';`,
+				Expected: []sql.Row{},
+			},
+			{
+				// Without 'public' on search_path, we get a does not exist error
+				Query:       `SELECT 'testing'::regclass;`,
+				ExpectedErr: "does not exist",
+			},
+			{
+				Query: `SELECT 'public.testing'::regclass, 'public.seq1'::regclass, 'public.testview'::regclass, 'public.testing_pkey'::regclass;`,
+				Expected: []sql.Row{
+					{"public.testing", "public.seq1", "public.testview", "public.testing_pkey"},
+				},
+			},
 		},
 	},
 	{