Refactoring out a schema override helper function

go/libraries/doltcore/sqle/database.go

@@ -242,44 +242,7 @@ func (db Database) GetTableInsensitive(ctx *sql.Context, tblName string) (sql.Ta
 		return nil, false, err
 	}
 
-	tbl, ok, err := db.getTableInsensitive(ctx, nil, ds, root, tblName)
-	if err != nil {
-		return nil, false, err
-	}
-
-	if !ok {
-		// If we didn't find the table in the data root... check if there is an overridden schema commit
-		schemaRoot, err := resolveOverriddenSchemaRoot(ctx, db)
-		if err != nil {
-			return nil, false, err
-		}
-
-		if schemaRoot != nil {
-			otherTable, _, ok, err := schemaRoot.GetTableInsensitive(ctx, tblName)
-			if err != nil {
-				return nil, false, err
-			}
-
-			if ok {
-				otherSchema, err := otherTable.GetSchema(ctx)
-				if err != nil {
-					return nil, false, err
-				}
-
-				otherSqlSchema, err := sqlutil.FromDoltSchema(db.Name(), tblName, otherSchema)
-				if err != nil {
-					return nil, false, err
-				}
-
-				emptyTable := plan.NewEmptyTableWithSchema(otherSqlSchema.Schema)
-				return emptyTable.(sql.Table), true, nil
-			}
-		}
-
-		return nil, false, nil
-	}
-
-	return tbl, true, nil
+	return db.getTableInsensitive(ctx, nil, ds, root, tblName)
 }
 
 // GetTableInsensitiveAsOf implements sql.VersionedDatabase
@@ -310,18 +273,24 @@ func (db Database) GetTableInsensitiveAsOf(ctx *sql.Context, tableName string, a
 		return table, ok, nil
 	}
 
-	versionableTable, ok := table.(dtables.VersionableTable)
-	if !ok {
+	switch t := table.(type) {
+	case dtables.VersionableTable:
+		versionedTable, err := t.LockedToRoot(ctx, root)
+		if err != nil {
+			return nil, false, err
+		}
+		return versionedTable, true, nil
+
+	case *plan.EmptyTable:
+		// getTableInsensitive returns *plan.EmptyTable if the table doesn't exist in the data root, but
+		// schemas have been locked to a commit where the table does exist. Since the table is empty,
+		// there's no need to lock it to a root.
+		return t, true, nil
+
+	default:
 		panic(fmt.Sprintf("unexpected table type %T", table))
 	}
 
-	versionedTable, err := versionableTable.LockedToRoot(ctx, root)
-
-	if err != nil {
-		return nil, false, err
-	}
-	return versionedTable, true, nil
-
 }
 
 func (db Database) getTableInsensitive(ctx *sql.Context, head *doltdb.Commit, ds *dsess.DoltSession, root *doltdb.RootValue, tblName string) (sql.Table, bool, error) {
@@ -508,7 +477,17 @@ func (db Database) getTableInsensitive(ctx *sql.Context, head *doltdb.Commit, ds
 	}
 
 	// TODO: this should reuse the root, not lookup the db state again
-	return db.getTable(ctx, root, tblName)
+	table, found, err := db.getTable(ctx, root, tblName)
+	if err != nil {
+		return nil, false, err
+	}
+	if found {
+		return table, found, err
+	}
+
+	// If the table wasn't found in the specified data root, check if there is an overridden
+	// schema commit that contains it and return an empty table if so.
+	return resolveOverriddenNonexistentTable(ctx, tblName, db)
 }
 
 // resolveAsOf resolves given expression to a commit, if one exists.
@@ -705,86 +684,6 @@ func (db Database) getTable(ctx *sql.Context, root *doltdb.RootValue, tableName
 	return table, true, nil
 }
 
-// TODO: Should this be in another type/file? Maybe in a schema-override.go file?
-// TODO: Godocs
-func overrideSchemaForTable(ctx *sql.Context, tableName string, tbl *doltdb.Table, overriddenSchemaRoot *doltdb.RootValue) error {
-	differentTable, _, ok, err := overriddenSchemaRoot.GetTableInsensitive(ctx, tableName)
-	if err != nil {
-		return fmt.Errorf("unable to find table at overridden schema root: %s", err.Error())
-	}
-	if !ok {
-		return fmt.Errorf("unable to find table at overridden schema root")
-	}
-	overriddenSchema, err := differentTable.GetSchema(ctx)
-	if err != nil {
-		return fmt.Errorf("unable to find table at overridden schema root: %s", err.Error())
-	}
-
-	tbl.OverrideSchema(overriddenSchema)
-	return nil
-}
-
-// TODO: Should this be in another type/file? Maybe in a schema-override.go file?
-// TODO: Godocs
-func getOverriddenSchemaValue(ctx *sql.Context) (*string, error) {
-	doltSession := dsess.DSessFromSess(ctx.Session)
-	varValue, err := doltSession.GetSessionVariable(ctx, dsess.DoltOverrideSchema)
-	if err != nil {
-		return nil, err
-	}
-
-	if varValue == nil {
-		return nil, nil
-	}
-
-	varString, ok := varValue.(string)
-	if !ok {
-		return nil, fmt.Errorf("value of %s is not a string", dsess.DoltOverrideSchema)
-	}
-
-	if varString == "" {
-		return nil, nil
-	}
-
-	return &varString, nil
-}
-
-// TODO: Should this be in another type/file? Maybe in a schema-override.go file?
-// TODO: Godocs
-func resolveOverriddenSchemaRoot(ctx *sql.Context, db Database) (*doltdb.RootValue, error) {
-	overriddenSchemaValue, err := getOverriddenSchemaValue(ctx)
-	if err != nil {
-		return nil, err
-	}
-
-	if overriddenSchemaValue == nil {
-		return nil, nil
-	}
-
-	commitSpec, err := doltdb.NewCommitSpec(*overriddenSchemaValue)
-	if err != nil {
-		return nil, fmt.Errorf("invalid commit spec specified in %s: %s", dsess.DoltOverrideSchema, err.Error())
-	}
-
-	doltSession := dsess.DSessFromSess(ctx.Session)
-	headRef, err := doltSession.CWBHeadRef(ctx, db.Name())
-	if err != nil {
-		return nil, fmt.Errorf("unable to retrieve current working branch head: " + err.Error())
-	}
-
-	commit, err := db.GetDoltDB().Resolve(ctx, commitSpec, headRef)
-	if err != nil {
-		return nil, fmt.Errorf("unable to resolve schema override value: " + err.Error())
-	}
-
-	rootValue, err := commit.GetRootValue(ctx)
-	if err != nil {
-		return nil, fmt.Errorf("unable to load root value for schema override commit: " + err.Error())
-	}
-
-	return rootValue, nil
-}
-
 // newDoltTable returns a sql.Table wrapping the given underlying dolt table
 func (db Database) newDoltTable(tableName string, sch schema.Schema, tbl *doltdb.Table) (sql.Table, error) {
 	readonlyTable, err := NewDoltTable(tableName, sch, tbl, db, db.editOpts)

go/libraries/doltcore/sqle/enginetest/dolt_queries_schema_override.go

@@ -418,6 +418,33 @@ var SchemaOverrideTests = []queries.ScriptTest{
 			},
 		},
 	},
+	{
+		Name: "Table Existence: Table exists in the pinned schema commit, NOT in AS OF data commit",
+		SetUpScript: []string{
+			"create table deletedTable (pk int primary key, c1 varchar(255));",
+			"insert into deletedTable values (1, 'one');",
+			"call dolt_commit('-Am', 'adding table deletedTable on main');",
+			"SET @commit1 = hashof('HEAD');",
+
+			"drop table deletedTable;",
+			"call dolt_commit('-Am', 'deleting table deletedTable on main');",
+			"SET @commit2 = hashof('HEAD');",
+		},
+		Assertions: []queries.ScriptTestAssertion{
+			{
+				Query:          "SELECT * from deletedTable;",
+				ExpectedErrStr: "table not found: deletedtable",
+			},
+			{
+				Query:    "SET @@dolt_override_schema=@commit1;",
+				Expected: []sql.Row{{}},
+			},
+			{
+				Query:    "SELECT * from deletedTable AS OF @commit2;",
+				Expected: []sql.Row{},
+			},
+		},
+	},
 
 	// INDEXES
 	{

go/libraries/doltcore/sqle/schema_override.go

@@ -0,0 +1,143 @@
+// 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 sqle
+
+import (
+	"fmt"
+	"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/sqlutil"
+	"github.com/dolthub/go-mysql-server/sql"
+	"github.com/dolthub/go-mysql-server/sql/plan"
+)
+
+// resolveOverriddenNonexistentTable check if there is an overridden schema commit set for this session, and if so
+// returns an empty table with that schema if |tblName| exists in the overridden schema commit. If no schema override
+// is set, this function returns a nil sql.Table and a false boolean return parameter.
+func resolveOverriddenNonexistentTable(ctx *sql.Context, tblName string, db Database) (sql.Table, bool, error) {
+	// Check to see if table schemas have been overridden
+	schemaRoot, err := resolveOverriddenSchemaRoot(ctx, db)
+	if err != nil {
+		return nil, false, err
+	}
+	if schemaRoot == nil {
+		return nil, false, nil
+	}
+
+	// If schema overrides are in place, see if the table exists in the overridden schema
+	t, _, ok, err := schemaRoot.GetTableInsensitive(ctx, tblName)
+	if err != nil {
+		return nil, false, err
+	}
+	if !ok {
+		return nil, false, nil
+	}
+
+	// Load the overridden schema and convert it to a sql.Schema
+	overriddenSchema, err := t.GetSchema(ctx)
+	if err != nil {
+		return nil, false, err
+	}
+	overriddenSqlSchema, err := sqlutil.FromDoltSchema(db.Name(), tblName, overriddenSchema)
+	if err != nil {
+		return nil, false, err
+	}
+
+	// Return an empty table with the overridden schema
+	emptyTable := plan.NewEmptyTableWithSchema(overriddenSqlSchema.Schema)
+	return emptyTable.(sql.Table), true, nil
+}
+
+// overrideSchemaForTable loads the schema from |overriddenSchemaRoot| for the table named |tableName| and sets the
+// override on |tbl|. If there are any problems loading the overridden schema, this function returns an error.
+func overrideSchemaForTable(ctx *sql.Context, tableName string, tbl *doltdb.Table, overriddenSchemaRoot *doltdb.RootValue) error {
+	differentTable, _, ok, err := overriddenSchemaRoot.GetTableInsensitive(ctx, tableName)
+	if err != nil {
+		return fmt.Errorf("unable to find table at overridden schema root: %s", err.Error())
+	}
+	if !ok {
+		return fmt.Errorf("unable to find table at overridden schema root")
+	}
+	overriddenSchema, err := differentTable.GetSchema(ctx)
+	if err != nil {
+		return fmt.Errorf("unable to find table at overridden schema root: %s", err.Error())
+	}
+
+	tbl.OverrideSchema(overriddenSchema)
+	return nil
+}
+
+// getOverriddenSchemaValue returns a pointer to the string value of the Dolt schema override session variable. If the
+// variable is not set (i.e. NULL or empty string) then this function returns nil so that callers can simply check for
+// nil.
+func getOverriddenSchemaValue(ctx *sql.Context) (*string, error) {
+	doltSession := dsess.DSessFromSess(ctx.Session)
+	varValue, err := doltSession.GetSessionVariable(ctx, dsess.DoltOverrideSchema)
+	if err != nil {
+		return nil, err
+	}
+
+	if varValue == nil {
+		return nil, nil
+	}
+
+	varString, ok := varValue.(string)
+	if !ok {
+		return nil, fmt.Errorf("value of %s is not a string", dsess.DoltOverrideSchema)
+	}
+
+	if varString == "" {
+		return nil, nil
+	}
+
+	return &varString, nil
+}
+
+// resolveOverriddenSchemaRoot loads the Dolt schema override session variable, resolves that commit reference, and
+// loads the RootValue for that commit. If the session variable is not set, this function returns nil. If there are
+// any problems resolving the commit or loading the root value, this function returns an error.
+func resolveOverriddenSchemaRoot(ctx *sql.Context, db Database) (*doltdb.RootValue, error) {
+	overriddenSchemaValue, err := getOverriddenSchemaValue(ctx)
+	if err != nil {
+		return nil, err
+	}
+
+	if overriddenSchemaValue == nil {
+		return nil, nil
+	}
+
+	commitSpec, err := doltdb.NewCommitSpec(*overriddenSchemaValue)
+	if err != nil {
+		return nil, fmt.Errorf("invalid commit spec specified in %s: %s", dsess.DoltOverrideSchema, err.Error())
+	}
+
+	doltSession := dsess.DSessFromSess(ctx.Session)
+	headRef, err := doltSession.CWBHeadRef(ctx, db.Name())
+	if err != nil {
+		return nil, fmt.Errorf("unable to retrieve current working branch head: " + err.Error())
+	}
+
+	commit, err := db.GetDoltDB().Resolve(ctx, commitSpec, headRef)
+	if err != nil {
+		return nil, fmt.Errorf("unable to resolve schema override value: " + err.Error())
+	}
+
+	rootValue, err := commit.GetRootValue(ctx)
+	if err != nil {
+		return nil, fmt.Errorf("unable to load root value for schema override commit: " + err.Error())
+	}
+
+	return rootValue, nil
+}