chore: import upstream snapshot with attribution

core/context.go

@@ -0,0 +1,672 @@
+// 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 core
+
+import (
+	"maps"
+	"slices"
+
+	"github.com/cockroachdb/errors"
+	"github.com/dolthub/dolt/go/libraries/doltcore/doltdb"
+	"github.com/dolthub/dolt/go/libraries/doltcore/sqle/dsess"
+	"github.com/dolthub/dolt/go/libraries/doltcore/sqle/resolve"
+	"github.com/dolthub/go-mysql-server/sql"
+
+	"github.com/dolthub/doltgresql/core/casts"
+	"github.com/dolthub/doltgresql/core/extensions"
+	"github.com/dolthub/doltgresql/core/functions"
+	"github.com/dolthub/doltgresql/core/procedures"
+	"github.com/dolthub/doltgresql/core/rootobject/objinterface"
+	"github.com/dolthub/doltgresql/core/sequences"
+	"github.com/dolthub/doltgresql/core/triggers"
+	"github.com/dolthub/doltgresql/core/typecollection"
+)
+
+// contextValues contains a set of cached data passed alongside the context. This data is considered temporary
+// and may be refreshed at any point, including during the middle of a query. Callers should not assume that
+// data stored in contextValues is persisted, and other types of data should not be added to contextValues.
+type contextValues struct {
+	seqs  map[string]*sequences.Collection
+	types map[string]*typecollection.TypeCollection
+	funcs map[string]*functions.Collection
+	procs map[string]*procedures.Collection
+	trigs map[string]*triggers.Collection
+	exts  map[string]*extensions.Collection
+	casts map[string]*casts.Collection
+
+	pgCatalogCache any
+	runner         sql.StatementRunner
+
+	// cache the dateOutputFormat, this is refreshed on SET
+	dateOutputFormat string
+}
+
+// getContextValues accesses the contextValues in the given context. If the context does not have a contextValues, then
+// it creates one and adds it to the context.
+func getContextValues(ctx *sql.Context) (*contextValues, error) {
+	if ctx == nil {
+		return nil, errors.New("context is nil")
+	}
+	sess := dsess.DSessFromSess(ctx.Session)
+	if sess.DoltgresSessObj == nil {
+		cv := &contextValues{}
+		sess.DoltgresSessObj = cv
+		return cv, nil
+	}
+	cv, ok := sess.DoltgresSessObj.(*contextValues)
+	if !ok {
+		return nil, errors.Errorf("context contains an unknown values struct of type: %T", sess.DoltgresSessObj)
+	}
+	return cv, nil
+}
+
+// ClearContextValues clears all context values. This is primarily for operations that are directly called from Dolt, as
+// Dolt does not have the Doltgres concept of context values. Care must be taken to ensure that intermediate state
+// written to the context values are not overwritten.
+func ClearContextValues(ctx *sql.Context) {
+	sess := dsess.DSessFromSess(ctx.Session)
+	if sess.DoltgresSessObj != nil {
+		// We want to persist the runner between clears since it's static
+		var runner sql.StatementRunner
+		if cv, ok := sess.DoltgresSessObj.(*contextValues); ok {
+			runner = cv.runner
+		}
+		sess.DoltgresSessObj = &contextValues{
+			runner: runner,
+		}
+	}
+}
+
+// GetRootFromContext returns the working session's root from the context, along with the session.
+func GetRootFromContext(ctx *sql.Context) (*dsess.DoltSession, *RootValue, error) {
+	return getRootFromContextForDatabase(ctx, "")
+}
+
+// getRootFromContextForDatabase returns the working session's root from the context for a specific database, along with the session.
+func getRootFromContextForDatabase(ctx *sql.Context, database string) (*dsess.DoltSession, *RootValue, error) {
+	session := dsess.DSessFromSess(ctx.Session)
+
+	if len(database) == 0 {
+		database = ctx.GetCurrentDatabase()
+	}
+	state, ok, err := session.LookupDbState(ctx, database)
+	if err != nil {
+		return nil, nil, err
+	}
+	if !ok {
+		return nil, nil, sql.ErrDatabaseNotFound.New(database)
+	}
+	return session, state.WorkingRoot().(*RootValue), nil
+}
+
+// IsContextValid returns whether the context is valid for use with any of the functions in the package. If this is not
+// false, then there's a high likelihood that the context is being used in a temporary scenario (such as setting up the
+// database, etc.).
+func IsContextValid(ctx *sql.Context) bool {
+	if ctx == nil {
+		return false
+	}
+	_, ok := ctx.Session.(*dsess.DoltSession)
+	return ok
+}
+
+// GetPgCatalogCache returns a cache of data for pg_catalog tables. This function should only be used by
+// pg_catalog table handlers. The catalog cache instance stores generated pg_catalog table data so that
+// it only has to generated table data once per query.
+//
+// TODO: The return type here is currently any, to avoid a package import cycle. This could be cleaned up by
+// introducing a new interface type, in a package that does not depend on core or pgcatalog packages.
+func GetPgCatalogCache(ctx *sql.Context) (any, error) {
+	cv, err := getContextValues(ctx)
+	if err != nil {
+		return nil, err
+	}
+	return cv.pgCatalogCache, nil
+}
+
+// SetPgCatalogCache sets |pgCatalogCache| as the catalog cache instance for this session.
+//
+// TODO: The input type here is currently any, to avoid a package import cycle. This could be cleaned up by
+// introducing a new interface type, in a package that does not depend on core or pgcatalog packages.
+func SetPgCatalogCache(ctx *sql.Context, pgCatalogCache any) error {
+	cv, err := getContextValues(ctx)
+	if err != nil {
+		return err
+	}
+	cv.pgCatalogCache = pgCatalogCache
+	return nil
+}
+
+// SetRunnerOnContext sets the given runner within the context values.
+func SetRunnerOnContext(ctx *sql.Context, runner sql.StatementRunner) error {
+	if runner == nil {
+		return nil
+	}
+	cv, err := getContextValues(ctx)
+	if err != nil {
+		return err
+	}
+	cv.runner = runner
+	return nil
+}
+
+// GetRunnerFromContext returns the sql.StatementRunner from within the context.
+func GetRunnerFromContext(ctx *sql.Context) (sql.StatementRunner, error) {
+	cv, err := getContextValues(ctx)
+	if err != nil {
+		return nil, err
+	}
+	return cv.runner, nil
+}
+
+// GetDoltTableFromContext returns the Dolt table from the context. Returns nil if no table was found.
+func GetDoltTableFromContext(ctx *sql.Context, tableName doltdb.TableName) (*doltdb.Table, error) {
+	_, root, err := GetRootFromContext(ctx)
+	if err != nil {
+		return nil, err
+	}
+
+	var table *doltdb.Table
+	if tableName.Schema == "" {
+		_, table, _, err = resolve.Table(ctx, root, tableName.Name)
+		if err != nil {
+			return nil, err
+		}
+	} else {
+		table, _, err = root.GetTable(ctx, tableName)
+	}
+
+	if err != nil {
+		return nil, err
+	}
+
+	return table, nil
+}
+
+// GetSqlDatabaseFromContext returns the database from the context. Uses the context's current database if an empty
+// string is provided. Returns nil if the database was not found.
+func GetSqlDatabaseFromContext(ctx *sql.Context, database string) (sql.Database, error) {
+	session := dsess.DSessFromSess(ctx.Session)
+	if len(database) == 0 {
+		database = ctx.GetCurrentDatabase()
+	}
+	db, err := session.Provider().Database(ctx, database)
+	if err != nil {
+		if sql.ErrDatabaseNotFound.Is(err) {
+			return nil, nil
+		}
+		return nil, err
+	}
+	return db, nil
+}
+
+// GetSqlTableFromContext returns the table from the context. Uses the context's current database if an empty database
+// name is provided. Returns nil if no table was found.
+func GetSqlTableFromContext(ctx *sql.Context, databaseName string, tableName doltdb.TableName) (sql.Table, error) {
+	db, err := GetSqlDatabaseFromContext(ctx, databaseName)
+	if err != nil || db == nil {
+		return nil, err
+	}
+	schemaDb, ok := db.(sql.SchemaDatabase)
+	if !ok {
+		// Fairly confident that Dolt only has database implementations that inherit sql.SchemaDatabase, so only GMS
+		// databases may fail here (like information_schema). In this scenario, we expect that no schema will be passed,
+		// so we'll special-case it here.
+		if len(tableName.Schema) == 0 {
+			tbl, ok, err := db.GetTableInsensitive(ctx, tableName.Name)
+			if err != nil || !ok {
+				return nil, err
+			}
+			return tbl, nil
+		}
+		return nil, nil
+	}
+
+	var searchPath []string
+	if len(tableName.Schema) == 0 {
+		// If a schema was not provided, then we'll use the search path
+		searchPath, err = SearchPath(ctx)
+		if err != nil {
+			return nil, err
+		}
+	} else {
+		// A specific schema is given, so we'll only use that one for the search path
+		searchPath = []string{tableName.Schema}
+	}
+
+	for _, schema := range searchPath {
+		db, ok, err = schemaDb.GetSchema(ctx, schema)
+		if err != nil {
+			return nil, err
+		}
+		if !ok {
+			continue
+		}
+		tbl, ok, err := db.GetTableInsensitive(ctx, tableName.Name)
+		if err != nil {
+			return nil, err
+		}
+		if !ok {
+			continue
+		}
+		return tbl, nil
+	}
+	return nil, nil
+}
+
+// SearchPath returns the effective schema search path for the current session
+func SearchPath(ctx *sql.Context) ([]string, error) {
+	path, err := resolve.SearchPath(ctx)
+	if err != nil {
+		return nil, err
+	}
+
+	// pg_catalog is *always* implicitly part of the search path as the first element, unless it's specifically
+	// included later. This allows users to override built-in names with user-defined names, but they have to
+	// opt in to that behavior.
+	hasPgCatalog := false
+	for _, schema := range path {
+		if schema == "pg_catalog" {
+			hasPgCatalog = true
+			break
+		}
+	}
+
+	if !hasPgCatalog {
+		path = append([]string{"pg_catalog"}, path...)
+	}
+	return path, nil
+}
+
+// GetExtensionsCollectionFromContext returns the extensions collection from the given context. Will always return a
+// collection if no error is returned.
+func GetExtensionsCollectionFromContext(ctx *sql.Context, database string) (*extensions.Collection, error) {
+	cv, err := getContextValues(ctx)
+	if err != nil {
+		return nil, err
+	}
+	if cv.exts == nil {
+		cv.exts = make(map[string]*extensions.Collection)
+	}
+	if len(database) == 0 {
+		database = ctx.GetCurrentDatabase()
+	}
+	if cv.exts[database] == nil {
+		_, root, err := getRootFromContextForDatabase(ctx, database)
+		if err != nil {
+			return nil, err
+		}
+		cv.exts[database], err = extensions.LoadExtensions(ctx, root)
+		if err != nil {
+			return nil, err
+		}
+	}
+	return cv.exts[database], nil
+}
+
+// GetFunctionsCollectionFromContext returns the functions collection from the given context. Will always return a
+// collection if no error is returned.
+func GetFunctionsCollectionFromContext(ctx *sql.Context, database string) (*functions.Collection, error) {
+	cv, err := getContextValues(ctx)
+	if err != nil {
+		return nil, err
+	}
+	if cv.funcs == nil {
+		cv.funcs = make(map[string]*functions.Collection)
+	}
+	if len(database) == 0 {
+		database = ctx.GetCurrentDatabase()
+	}
+	if cv.funcs[database] == nil {
+		_, root, err := getRootFromContextForDatabase(ctx, database)
+		if err != nil {
+			return nil, err
+		}
+		cv.funcs[database], err = functions.LoadFunctions(ctx, root)
+		if err != nil {
+			return nil, err
+		}
+	}
+	return cv.funcs[database], nil
+}
+
+// GetProceduresCollectionFromContext returns the procedures collection from the given context. Will always return a
+// collection if no error is returned.
+func GetProceduresCollectionFromContext(ctx *sql.Context, database string) (*procedures.Collection, error) {
+	cv, err := getContextValues(ctx)
+	if err != nil {
+		return nil, err
+	}
+	if cv.procs == nil {
+		cv.procs = make(map[string]*procedures.Collection)
+	}
+	if len(database) == 0 {
+		database = ctx.GetCurrentDatabase()
+	}
+	if cv.procs[database] == nil {
+		_, root, err := getRootFromContextForDatabase(ctx, database)
+		if err != nil {
+			return nil, err
+		}
+		cv.procs[database], err = procedures.LoadProcedures(ctx, root)
+		if err != nil {
+			return nil, err
+		}
+	}
+	return cv.procs[database], nil
+}
+
+// GetSequencesCollectionFromContext returns the given sequence collection from the context for the database
+// named. If no database is provided, the context's current database is used.
+// Will always return a collection if no error is returned.
+func GetSequencesCollectionFromContext(ctx *sql.Context, database string) (*sequences.Collection, error) {
+	cv, err := getContextValues(ctx)
+	if err != nil {
+		return nil, err
+	}
+	if cv.seqs == nil {
+		cv.seqs = make(map[string]*sequences.Collection)
+	}
+	if len(database) == 0 {
+		database = ctx.GetCurrentDatabase()
+	}
+	if cv.seqs[database] == nil {
+		_, root, err := getRootFromContextForDatabase(ctx, database)
+		if err != nil {
+			return nil, err
+		}
+		cv.seqs[database], err = sequences.LoadSequences(ctx, root)
+		if err != nil {
+			return nil, err
+		}
+	}
+	return cv.seqs[database], nil
+}
+
+// GetTriggersCollectionFromContext returns the triggers collection from the given context. Will always return a
+// collection if no error is returned.
+func GetTriggersCollectionFromContext(ctx *sql.Context, database string) (*triggers.Collection, error) {
+	cv, err := getContextValues(ctx)
+	if err != nil {
+		return nil, err
+	}
+	if cv.trigs == nil {
+		cv.trigs = make(map[string]*triggers.Collection)
+	}
+	if len(database) == 0 {
+		database = ctx.GetCurrentDatabase()
+	}
+	if cv.trigs[database] == nil {
+		_, root, err := getRootFromContextForDatabase(ctx, database)
+		if err != nil {
+			return nil, err
+		}
+		cv.trigs[database], err = triggers.LoadTriggers(ctx, root)
+		if err != nil {
+			return nil, err
+		}
+	}
+	return cv.trigs[database], nil
+}
+
+// GetTypesCollectionFromContext returns the given type collection from the context.
+// Will always return a collection if no error is returned.
+func GetTypesCollectionFromContext(ctx *sql.Context, database string) (*typecollection.TypeCollection, error) {
+	cv, err := getContextValues(ctx)
+	if err != nil {
+		return nil, err
+	}
+	if cv.types == nil {
+		cv.types = make(map[string]*typecollection.TypeCollection)
+	}
+	if len(database) == 0 {
+		database = ctx.GetCurrentDatabase()
+	}
+	if cv.types[database] == nil {
+		_, root, err := getRootFromContextForDatabase(ctx, database)
+		if err != nil {
+			return nil, err
+		}
+		cv.types[database], err = typecollection.LoadTypes(ctx, root)
+		if err != nil {
+			return nil, err
+		}
+	}
+	return cv.types[database], nil
+}
+
+// GetCastsCollectionFromContext returns the given casts collection from the context.
+// Will always return a collection if no error is returned.
+func GetCastsCollectionFromContext(ctx *sql.Context, database string) (*casts.Collection, error) {
+	cv, err := getContextValues(ctx)
+	if err != nil {
+		return nil, err
+	}
+	if cv.casts == nil {
+		cv.casts = make(map[string]*casts.Collection)
+	}
+	if len(database) == 0 {
+		database = ctx.GetCurrentDatabase()
+	}
+	if cv.casts[database] == nil {
+		_, root, err := getRootFromContextForDatabase(ctx, database)
+		if err != nil {
+			return nil, err
+		}
+		cv.casts[database], err = casts.LoadCasts(ctx, root)
+		if err != nil {
+			return nil, err
+		}
+	}
+	return cv.casts[database], nil
+}
+
+// CloseContextRootFinalizer finalizes any changes persisted within the context by writing them to the working root.
+// This should ONLY be called by the ContextRootFinalizer node.
+func CloseContextRootFinalizer(ctx *sql.Context) error {
+	sess := dsess.DSessFromSess(ctx.Session)
+	if sess.DoltgresSessObj == nil {
+		return nil
+	}
+	cv, ok := sess.DoltgresSessObj.(*contextValues)
+	if !ok {
+		return nil
+	}
+
+	// We need to update the root for all databases used by this context. This logic parallels what happens during
+	// transaction commit in the dolt/sqle layer, where we check each branch state to see if it's dirty
+	for _, db := range databasesInContext(ctx, cv) {
+		err := updateSessionRootForDatabase(ctx, db, cv)
+		if err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
+// GetDateStyleOutputFormat returns the cached DateOutputFormat
+func GetDateStyleOutputFormat(ctx *sql.Context) (string, error) {
+	cv, err := getContextValues(ctx)
+	if err != nil {
+		return "", err
+	}
+	return cv.dateOutputFormat, nil
+}
+
+// SetDateStyleOutputFormat cached the provided dateOutputFormat
+func SetDateStyleOutputFormat(ctx *sql.Context, dateOutputFormat string) error {
+	cv, err := getContextValues(ctx)
+	if err != nil {
+		return err
+	}
+	cv.dateOutputFormat = dateOutputFormat
+	return nil
+}
+
+// updateSessionRootForDatabase updates the root for all changes made to root object collections within the context
+// values.
+func updateSessionRootForDatabase(ctx *sql.Context, db string, cv *contextValues) error {
+	session, root, err := getRootFromContextForDatabase(ctx, db)
+	if err != nil {
+		return err
+	}
+
+	newRoot := root
+	if cv.seqs != nil && cv.seqs[db] != nil {
+		retRoot, err := cv.seqs[db].UpdateRoot(ctx, newRoot)
+		if err != nil {
+			return err
+		}
+		newRoot = retRoot.(*RootValue)
+		delete(cv.seqs, db)
+	}
+
+	if cv.funcs != nil && cv.funcs[db] != nil && cv.funcs[db].DiffersFrom(ctx, root) {
+		retRoot, err := cv.funcs[db].UpdateRoot(ctx, newRoot)
+		if err != nil {
+			return err
+		}
+		newRoot = retRoot.(*RootValue)
+		delete(cv.funcs, db)
+	}
+
+	if cv.procs != nil && cv.procs[db] != nil && cv.procs[db].DiffersFrom(ctx, root) {
+		retRoot, err := cv.procs[db].UpdateRoot(ctx, newRoot)
+		if err != nil {
+			return err
+		}
+		newRoot = retRoot.(*RootValue)
+		delete(cv.procs, db)
+	}
+
+	if cv.trigs != nil && cv.trigs[db] != nil && cv.trigs[db].DiffersFrom(ctx, root) {
+		retRoot, err := cv.trigs[db].UpdateRoot(ctx, newRoot)
+		if err != nil {
+			return err
+		}
+		newRoot = retRoot.(*RootValue)
+		delete(cv.trigs, db)
+	}
+
+	if cv.exts != nil && cv.exts[db] != nil && cv.exts[db].DiffersFrom(ctx, root) {
+		retRoot, err := cv.exts[db].UpdateRoot(ctx, newRoot)
+		if err != nil {
+			return err
+		}
+		newRoot = retRoot.(*RootValue)
+		delete(cv.exts, db)
+	}
+
+	if cv.types != nil && cv.types[db] != nil {
+		retRoot, err := cv.types[db].UpdateRoot(ctx, newRoot)
+		if err != nil {
+			return err
+		}
+		newRoot = retRoot.(*RootValue)
+		delete(cv.types, db)
+	}
+
+	if cv.casts != nil && cv.casts[db] != nil && cv.casts[db].DiffersFrom(ctx, root) {
+		retRoot, err := cv.casts[db].UpdateRoot(ctx, newRoot)
+		if err != nil {
+			return err
+		}
+		newRoot = retRoot.(*RootValue)
+		delete(cv.casts, db)
+	}
+
+	// Setting the session working root doesn't do a check to see if anything actually changed or not before marking that
+	// branch state dirty, and dolt only allows a single dirty working set per commit. So it's important here to only
+	// update the session root if something actually changed for that db.
+	if err, rootChanged := rootValueChanged(newRoot, root); rootChanged {
+		if err = session.SetWorkingRoot(ctx, db, newRoot); err != nil {
+			// TODO: We need a way to see if the session has a writeable working root
+			// (new interface method on session probably), and avoid setting it if so
+			if errors.Is(err, doltdb.ErrOperationNotSupportedInDetachedHead) {
+				return nil
+			}
+			return err
+		}
+	} else if err != nil {
+		return err
+	}
+
+	return nil
+}
+
+// rootValueChanged returns whether the new root value is different from the old one
+func rootValueChanged(newRoot *RootValue, root *RootValue) (error, bool) {
+	if newRoot == root {
+		return nil, false
+	}
+
+	newHash, err := newRoot.HashOf()
+	if err != nil {
+		return err, false
+	}
+
+	oldHash, err := root.HashOf()
+	if err != nil {
+		return err, false
+	}
+
+	if newHash == oldHash {
+		return nil, false
+	}
+
+	return nil, true
+}
+
+// databasesInContext returns all databases found within the context values.
+func databasesInContext(ctx *sql.Context, cv *contextValues) []string {
+	dbs := make(map[string]struct{})
+	if cv.seqs != nil {
+		for db := range cv.seqs {
+			dbs[db] = struct{}{}
+		}
+	}
+	currentDb := ctx.GetCurrentDatabase()
+	if len(currentDb) > 0 {
+		dbs[currentDb] = struct{}{}
+	}
+
+	return slices.Sorted(maps.Keys(dbs))
+}
+
+// clear removes the collection from the cache.
+func (cv *contextValues) clear(objID objinterface.RootObjectID) {
+	switch objID {
+	case objinterface.RootObjectID_None:
+		// Nothing to cache with this
+	case objinterface.RootObjectID_Sequences:
+		cv.seqs = nil
+	case objinterface.RootObjectID_Types:
+		cv.types = nil
+	case objinterface.RootObjectID_Functions:
+		cv.funcs = nil
+	case objinterface.RootObjectID_Triggers:
+		cv.trigs = nil
+	case objinterface.RootObjectID_Extensions:
+		// We don't cache these
+	case objinterface.RootObjectID_Conflicts:
+		// We don't cache these
+	case objinterface.RootObjectID_Procedures:
+		cv.procs = nil
+	case objinterface.RootObjectID_Casts:
+		cv.casts = nil
+	default:
+		panic("unhandled context clear object ID")
+	}
+}