Allow foreign keys on nonlocal tables

go/cmd/dolt/commands/cvcmds/verify_constraints.go

@@ -29,6 +29,7 @@ import (
 	"github.com/dolthub/dolt/go/libraries/doltcore/doltdb"
 	"github.com/dolthub/dolt/go/libraries/doltcore/env"
 	"github.com/dolthub/dolt/go/libraries/doltcore/merge"
+	"github.com/dolthub/dolt/go/libraries/doltcore/sqle/dsess"
 	"github.com/dolthub/dolt/go/libraries/utils/argparser"
 	"github.com/dolthub/dolt/go/store/types"
 )
@@ -111,30 +112,37 @@ func (cmd VerifyConstraintsCmd) Exec(ctx context.Context, commandStr string, arg
 		return commands.HandleVErrAndExitCode(errhand.BuildDError("Unable to get head commit hash.").AddCause(err).Build(), nil)
 	}
 
-	endRoot, tablesWithViolations, err := merge.AddForeignKeyViolations(ctx, working, comparingRoot, tableSet, h)
+	eng, dbName, err := engine.NewSqlEngineForEnv(ctx, dEnv)
+	if err != nil {
+		return commands.HandleVErrAndExitCode(errhand.BuildDError("Failed to build sql engine.").AddCause(err).Build(), nil)
+	}
+	sqlCtx, err := eng.NewLocalContext(ctx)
+	if err != nil {
+		return commands.HandleVErrAndExitCode(errhand.BuildDError("Failed to build sql context.").AddCause(err).Build(), nil)
+	}
+	defer sql.SessionEnd(sqlCtx.Session)
+	sql.SessionCommandBegin(sqlCtx.Session)
+	defer sql.SessionCommandEnd(sqlCtx.Session)
+	sqlCtx.SetCurrentDatabase(dbName)
+
+	tableResolver, err := dsess.GetTableResolver(sqlCtx)
+	if err != nil {
+		cli.PrintErrln(errhand.VerboseErrorFromError(err))
+		return 1
+	}
+
+	endRoot, tablesWithViolations, err := merge.AddForeignKeyViolations(sqlCtx, tableResolver, working, comparingRoot, tableSet, h)
 	if err != nil {
 		return commands.HandleVErrAndExitCode(errhand.BuildDError("Unable to process constraint violations.").AddCause(err).Build(), nil)
 	}
 
-	err = dEnv.UpdateWorkingRoot(ctx, endRoot)
+	err = dEnv.UpdateWorkingRoot(sqlCtx, endRoot)
 	if err != nil {
 		return commands.HandleVErrAndExitCode(errhand.BuildDError("Unable to update working root.").AddCause(err).Build(), nil)
 	}
 
 	if tablesWithViolations.Size() > 0 {
 		cli.PrintErrln("All constraints are not satisfied.")
-		eng, dbName, err := engine.NewSqlEngineForEnv(ctx, dEnv)
-		if err != nil {
-			return commands.HandleVErrAndExitCode(errhand.BuildDError("Failed to build sql engine.").AddCause(err).Build(), nil)
-		}
-		sqlCtx, err := eng.NewLocalContext(ctx)
-		if err != nil {
-			return commands.HandleVErrAndExitCode(errhand.BuildDError("Failed to build sql context.").AddCause(err).Build(), nil)
-		}
-		defer sql.SessionEnd(sqlCtx.Session)
-		sql.SessionCommandBegin(sqlCtx.Session)
-		defer sql.SessionCommandEnd(sqlCtx.Session)
-		sqlCtx.SetCurrentDatabase(dbName)
 
 		for _, tableName := range tablesWithViolations.AsSortedSlice() {
 			tbl, ok, err := endRoot.GetTable(sqlCtx, tableName)

go/cmd/dolt/commands/schcmds/copy-tags.go

@@ -18,6 +18,8 @@ import (
 	"context"
 	"fmt"
 
+	"github.com/dolthub/go-mysql-server/sql"
+
 	"github.com/dolthub/dolt/go/cmd/dolt/cli"
 	"github.com/dolthub/dolt/go/cmd/dolt/commands"
 	"github.com/dolthub/dolt/go/cmd/dolt/errhand"
@@ -27,6 +29,7 @@ import (
 	"github.com/dolthub/dolt/go/libraries/doltcore/env/actions"
 	"github.com/dolthub/dolt/go/libraries/doltcore/ref"
 	"github.com/dolthub/dolt/go/libraries/doltcore/schema"
+	"github.com/dolthub/dolt/go/libraries/doltcore/sqle/dsess"
 	"github.com/dolthub/dolt/go/libraries/utils/argparser"
 	"github.com/dolthub/dolt/go/libraries/utils/config"
 	"github.com/dolthub/dolt/go/store/datas"
@@ -85,13 +88,26 @@ func (cmd CopyTagsCmd) ArgParser() *argparser.ArgParser {
 }
 
 // Exec implements the cli.Command interface.
-func (cmd CopyTagsCmd) Exec(ctx context.Context, commandStr string, args []string, dEnv *env.DoltEnv, _ cli.CliContext) int {
+func (cmd CopyTagsCmd) Exec(ctx context.Context, commandStr string, args []string, dEnv *env.DoltEnv, cliCtx cli.CliContext) int {
 	fromBranchName, fromRoots, toRoots, usage, err := cmd.validateArgs(ctx, commandStr, args, dEnv)
 	if err != nil {
 		verr := errhand.BuildDError("error validating arguments")
 		return commands.HandleVErrAndExitCode(verr.AddCause(err).Build(), usage)
 	}
 
+	queryist, err := cliCtx.QueryEngine(ctx)
+	if err != nil {
+		cli.PrintErrln(errhand.VerboseErrorFromError(err))
+		return 1
+	}
+	sqlCtx := queryist.Context
+
+	tableResolver, err := dsess.GetTableResolver(sqlCtx)
+	if err != nil {
+		cli.PrintErrln(errhand.VerboseErrorFromError(err))
+		return 1
+	}
+
 	// Load all the tags from fromBranch and toBranch
 	fromBranchTags, err := getAllTagsForRoots(ctx, fromRoots.Head)
 	if err != nil {
@@ -158,7 +174,7 @@ func (cmd CopyTagsCmd) Exec(ctx context.Context, commandStr string, args []strin
 	}
 
 	if tagsSynced > 0 {
-		if err = doltCommitUpdatedTags(ctx, dEnv, workingRoot, fromBranchName); err != nil {
+		if err = doltCommitUpdatedTags(sqlCtx, tableResolver, dEnv, workingRoot, fromBranchName); err != nil {
 			vErr := errhand.BuildDError("failed to commit column tag updates").AddCause(err).Build()
 			return commands.HandleVErrAndExitCode(vErr, usage)
 		}
@@ -240,7 +256,7 @@ func validateDestinationBranch(ctx context.Context, toRoots *doltdb.Roots) error
 
 // doltCommitUpdatedTags commits tag changes in |workingRoot| for the specified DoltEnv, |dEnv|. The commit message uses
 // |fromBranchName| to document the source of the tag changes.
-func doltCommitUpdatedTags(ctx context.Context, dEnv *env.DoltEnv, workingRoot doltdb.RootValue, fromBranchName string) (err error) {
+func doltCommitUpdatedTags(ctx *sql.Context, tableResolver doltdb.TableResolver, dEnv *env.DoltEnv, workingRoot doltdb.RootValue, fromBranchName string) (err error) {
 	if err = dEnv.UpdateWorkingRoot(ctx, workingRoot); err != nil {
 		return err
 	}
@@ -271,7 +287,7 @@ func doltCommitUpdatedTags(ctx context.Context, dEnv *env.DoltEnv, workingRoot d
 	}
 
 	doltDB := dEnv.DoltDB(ctx)
-	pendingCommit, err := actions.GetCommitStaged(ctx, roots, workingSet, nil, doltDB, actions.CommitStagedProps{
+	pendingCommit, err := actions.GetCommitStaged(ctx, tableResolver, roots, workingSet, nil, doltDB, actions.CommitStagedProps{
 		Name:    name,
 		Email:   email,
 		Message: "Syncing column tags from " + fromBranchName + " branch",

go/cmd/dolt/dolt.go

@@ -46,7 +46,6 @@ import (
 	"github.com/dolthub/dolt/go/cmd/dolt/commands/cvcmds"
 	"github.com/dolthub/dolt/go/cmd/dolt/commands/docscmds"
 	"github.com/dolthub/dolt/go/cmd/dolt/commands/indexcmds"
-	"github.com/dolthub/dolt/go/cmd/dolt/commands/schcmds"
 	"github.com/dolthub/dolt/go/cmd/dolt/commands/sqlserver"
 	"github.com/dolthub/dolt/go/cmd/dolt/doltcmd"
 	"github.com/dolthub/dolt/go/cmd/dolt/doltversion"
@@ -74,7 +73,6 @@ var commandsWithoutCliCtx = []cli.Command{
 	commands.BackupCmd{},
 	commands.LoginCmd{},
 	credcmds.Commands,
-	schcmds.Commands,
 	cvcmds.Commands,
 	commands.SendMetricsCmd{},
 	commands.MigrateCmd{},

go/libraries/doltcore/cherry_pick/cherry_pick.go

@@ -249,6 +249,11 @@ func cherryPick(ctx *sql.Context, dSess *dsess.DoltSession, roots doltdb.Roots,
 		return nil, "", nil, err
 	}
 
+	tableResolver, err := dsess.GetTableResolver(ctx)
+	if err != nil {
+		return nil, "", nil, err
+	}
+
 	doltDB, ok := dSess.GetDoltDB(ctx, dbName)
 	if !ok {
 		return nil, "", nil, fmt.Errorf("failed to get doltDB")
@@ -333,7 +338,7 @@ func cherryPick(ctx *sql.Context, dSess *dsess.DoltSession, roots doltdb.Roots,
 		IsCherryPick:        true,
 		KeepSchemaConflicts: false,
 	}
-	result, err := merge.MergeRoots(ctx, roots.Working, cherryRoot, parentRoot, cherryCommit, parentCommit, dbState.EditOpts(), mo)
+	result, err := merge.MergeRoots(ctx, tableResolver, roots.Working, cherryRoot, parentRoot, cherryCommit, parentCommit, dbState.EditOpts(), mo)
 	if err != nil {
 		return result, "", nil, err
 	}

go/libraries/doltcore/doltdb/doltdb.go

@@ -76,6 +76,35 @@ var InMemDoltDB = "mem://"
 var ErrNoRootValAtHash = errors.New("there is no dolt root value at that hash")
 var ErrCannotDeleteLastBranch = errors.New("cannot delete the last branch")
 
+// TableResolver allows the user of a DoltDB to configure how table names are resolved on roots.
+// This is useful because the user-backed system table dolt_nonlocal_tables allows table names to resolve to
+// tables on other refs, but sqle.Database is necessary to resolve those refs.
+type TableResolver interface {
+	ResolveTableInsensitive(ctx *sql.Context, root RootValue, tblName TableName) (trueTableName TableName, table *Table, found bool, err error)
+	ResolveTable(ctx *sql.Context, root RootValue, tblName TableName) (table *Table, found bool, err error)
+}
+
+type SimpleTableResolver struct{}
+
+var _ TableResolver = SimpleTableResolver{}
+
+func (t SimpleTableResolver) ResolveTableInsensitive(ctx *sql.Context, root RootValue, tblName TableName) (trueTableName TableName, table *Table, exists bool, err error) {
+	trueTableNameString, exists, err := root.ResolveTableName(ctx, tblName)
+	if err != nil || !exists {
+		return TableName{}, nil, false, err
+	}
+	trueTableName = TableName{
+		Name:   trueTableNameString,
+		Schema: tblName.Schema,
+	}
+	table, exists, err = root.GetTable(ctx, trueTableName)
+	return trueTableName, table, exists, err
+}
+
+func (t SimpleTableResolver) ResolveTable(ctx *sql.Context, root RootValue, tblName TableName) (table *Table, exists bool, err error) {
+	return root.GetTable(ctx, tblName)
+}
+
 // DoltDB wraps access to the underlying noms database and hides some of the details of the underlying storage.
 type DoltDB struct {
 	db  hooksDatabase

go/libraries/doltcore/doltdb/root_val.go

@@ -1191,55 +1191,47 @@ func (root *rootValue) PutForeignKeyCollection(ctx context.Context, fkc *Foreign
 // ValidateForeignKeysOnSchemas ensures that all foreign keys' tables are present, removing any foreign keys where the declared
 // table is missing, and returning an error if a key is in an invalid state or a referenced table is missing. Does not
 // check any tables' row data.
-func ValidateForeignKeysOnSchemas(ctx context.Context, root RootValue) (RootValue, error) {
+func ValidateForeignKeysOnSchemas(ctx *sql.Context, tableResolver TableResolver, root RootValue) (RootValue, error) {
 	fkCollection, err := root.GetForeignKeyCollection(ctx)
 	if err != nil {
 		return nil, err
 	}
 
-	allTablesSlice, err := UnionTableNames(ctx, root)
-	if err != nil {
-		return nil, err
-	}
 	allTablesSet := make(map[TableName]schema.Schema)
-	var rootObjNamesMap map[TableName]struct{}
-	for _, tableName := range allTablesSlice {
-		tbl, ok, err := root.GetTable(ctx, tableName)
+	getTableSchema := func(tableName TableName) (schema.Schema, bool, error) {
+		if tableSch, ok := allTablesSet[tableName]; ok {
+			return tableSch, true, nil
+		}
+		tbl, ok, err := tableResolver.ResolveTable(ctx, root, tableName)
 		if err != nil {
-			return nil, err
+			return nil, false, err
 		}
 		if !ok {
-			if rootObjNamesMap == nil {
-				rootObjNames, err := root.FilterRootObjectNames(ctx, allTablesSlice)
-				if err != nil {
-					return nil, err
-				}
-				rootObjNamesMap = make(map[TableName]struct{})
-				for _, rootObjName := range rootObjNames {
-					rootObjNamesMap[rootObjName] = struct{}{}
-				}
-			}
-			if _, ok = rootObjNamesMap[tableName]; ok {
-				continue
-			}
-			return nil, fmt.Errorf("found table `%s` in staging but could not load for foreign key check", tableName)
+			return nil, false, nil
 		}
 		tblSch, err := tbl.GetSchema(ctx)
 		if err != nil {
-			return nil, err
+			return nil, false, err
 		}
 		allTablesSet[tableName] = tblSch
+		return tblSch, true, nil
 	}
 
 	// some of these checks are sanity checks and should never happen
 	allForeignKeys := fkCollection.AllKeys()
 	for _, foreignKey := range allForeignKeys {
-		tblSch, existsInRoot := allTablesSet[foreignKey.TableName]
+		tblSch, existsInRoot, err := getTableSchema(foreignKey.TableName)
+		if err != nil {
+			return nil, err
+		}
 		if existsInRoot {
 			if err := foreignKey.ValidateTableSchema(tblSch); err != nil {
 				return nil, err
 			}
-			parentSch, existsInRoot := allTablesSet[foreignKey.ReferencedTableName]
+			parentSch, existsInRoot, err := getTableSchema(foreignKey.ReferencedTableName)
+			if err != nil {
+				return nil, err
+			}
 			if !existsInRoot {
 				return nil, fmt.Errorf("foreign key `%s` requires the referenced table `%s`", foreignKey.Name, foreignKey.ReferencedTableName)
 			}

go/libraries/doltcore/dtestutils/testcommands/multienv.go

@@ -20,19 +20,20 @@ import (
 	"os"
 	"path/filepath"
 
-	cmd "github.com/dolthub/dolt/go/cmd/dolt/commands"
-	"github.com/dolthub/dolt/go/libraries/doltcore/sqle/dprocedures"
-	"github.com/dolthub/dolt/go/libraries/utils/config"
+	"github.com/dolthub/go-mysql-server/sql"
 
 	"github.com/dolthub/dolt/go/cmd/dolt/cli"
+	cmd "github.com/dolthub/dolt/go/cmd/dolt/commands"
 	"github.com/dolthub/dolt/go/libraries/doltcore/doltdb"
 	"github.com/dolthub/dolt/go/libraries/doltcore/doltdb/durable"
 	"github.com/dolthub/dolt/go/libraries/doltcore/dtestutils"
 	"github.com/dolthub/dolt/go/libraries/doltcore/env"
 	"github.com/dolthub/dolt/go/libraries/doltcore/env/actions"
 	"github.com/dolthub/dolt/go/libraries/doltcore/row"
 	"github.com/dolthub/dolt/go/libraries/doltcore/schema"
+	"github.com/dolthub/dolt/go/libraries/doltcore/sqle/dprocedures"
 	"github.com/dolthub/dolt/go/libraries/doltcore/table"
+	"github.com/dolthub/dolt/go/libraries/utils/config"
 	"github.com/dolthub/dolt/go/libraries/utils/filesys"
 	"github.com/dolthub/dolt/go/store/datas"
 	"github.com/dolthub/dolt/go/store/types"
@@ -240,7 +241,7 @@ func (mr *MultiRepoTestSetup) GetDB(dbName string) *doltdb.DoltDB {
 }
 
 func (mr *MultiRepoTestSetup) CommitWithWorkingSet(dbName string) *doltdb.Commit {
-	ctx := context.Background()
+	ctx := sql.NewEmptyContext()
 	dEnv := mr.envs[dbName]
 	ws, err := dEnv.WorkingSet(ctx)
 	if err != nil {
@@ -262,7 +263,7 @@ func (mr *MultiRepoTestSetup) CommitWithWorkingSet(dbName string) *doltdb.Commit
 	if err != nil {
 		panic("couldn't get roots: " + err.Error())
 	}
-	pendingCommit, err := actions.GetCommitStaged(ctx, roots, ws, mergeParentCommits, dEnv.DbData(ctx).Ddb, actions.CommitStagedProps{
+	pendingCommit, err := actions.GetCommitStaged(ctx, doltdb.SimpleTableResolver{}, roots, ws, mergeParentCommits, dEnv.DbData(ctx).Ddb, actions.CommitStagedProps{
 		Message:    "auto commit",
 		Date:       t,
 		AllowEmpty: true,

go/libraries/doltcore/env/actions/commit.go

@@ -15,9 +15,10 @@
 package actions
 
 import (
-	"context"
 	"time"
 
+	"github.com/dolthub/go-mysql-server/sql"
+
 	"github.com/dolthub/dolt/go/libraries/doltcore/diff"
 	"github.com/dolthub/dolt/go/libraries/doltcore/doltdb"
 	"github.com/dolthub/dolt/go/store/datas"
@@ -36,7 +37,8 @@ type CommitStagedProps struct {
 
 // GetCommitStaged returns a new pending commit with the roots and commit properties given.
 func GetCommitStaged(
-	ctx context.Context,
+	ctx *sql.Context,
+	tableResolver doltdb.TableResolver,
 	roots doltdb.Roots,
 	ws *doltdb.WorkingSet,
 	mergeParents []*doltdb.Commit,
@@ -99,7 +101,7 @@ func GetCommitStaged(
 			}
 		}
 
-		roots.Staged, err = doltdb.ValidateForeignKeysOnSchemas(ctx, roots.Staged)
+		roots.Staged, err = doltdb.ValidateForeignKeysOnSchemas(ctx, tableResolver, roots.Staged)
 		if err != nil {
 			return nil, err
 		}

go/libraries/doltcore/merge/data_merge_test.go

@@ -129,7 +129,7 @@ func testDataMergeHelper(t *testing.T, tests []dataMergeTest, flipSides bool) {
 					var eo editor.Options
 					eo = eo.WithDeaf(editor.NewInMemDeaf(a.VRW()))
 					// attempt merge before skipping to assert no panics
-					result, err := merge.MergeRoots(sql.NewContext(ctx), l, r, a, rootish{r}, rootish{a}, eo, mo)
+					result, err := merge.MergeRoots(sql.NewContext(ctx), doltdb.SimpleTableResolver{}, l, r, a, rootish{r}, rootish{a}, eo, mo)
 
 					if data.dataConflict {
 						// TODO: Test the conflict error message more deeply