Invalidate all sessions' db state for drop database

go/cmd/dolt/commands/sqlserver/server.go

@@ -206,6 +206,7 @@ func Serve(
 
 	if errors.Is(startError, server.UnixSocketInUseError) {
 		lgr.Warn("unix socket set up failed: file already in use: ", serverConf.Socket)
+		startError = nil
 	} else if startError != nil {
 		cli.PrintErr(startError)
 		return

go/libraries/doltcore/sqle/database_provider.go

@@ -33,6 +33,7 @@ import (
 	"github.com/dolthub/dolt/go/libraries/doltcore/sqle/dfunctions"
 	"github.com/dolthub/dolt/go/libraries/doltcore/sqle/dprocedures"
 	"github.com/dolthub/dolt/go/libraries/doltcore/sqle/dsess"
+	"github.com/dolthub/dolt/go/libraries/doltcore/sqlserver"
 	"github.com/dolthub/dolt/go/libraries/doltcore/table/editor"
 	"github.com/dolthub/dolt/go/libraries/utils/filesys"
 	"github.com/dolthub/dolt/go/store/types"
@@ -109,6 +110,13 @@ func NewDoltDatabaseProviderWithDatabases(defaultBranch string, fs filesys.Files
 		externalProcedures.Register(esp)
 	}
 
+	// If the specified |fs| is an in mem file system, default to using the InMemDoltDB dbFactoryUrl so that all
+	// databases are created with the same file system type.
+	dbFactoryUrl := doltdb.LocalDirDoltDB
+	if _, ok := fs.(*filesys.InMemFS); ok {
+		dbFactoryUrl = doltdb.InMemDoltDB
+	}
+
 	return DoltDatabaseProvider{
 		dbLocations:        dbLocations,
 		databases:          dbs,
@@ -117,7 +125,7 @@ func NewDoltDatabaseProviderWithDatabases(defaultBranch string, fs filesys.Files
 		mu:                 &sync.RWMutex{},
 		fs:                 fs,
 		defaultBranch:      defaultBranch,
-		dbFactoryUrl:       doltdb.LocalDirDoltDB,
+		dbFactoryUrl:       dbFactoryUrl,
 		InitDatabaseHook:   ConfigureReplicationDatabaseHook,
 		isStandby:          new(bool),
 	}, nil
@@ -698,6 +706,35 @@ func (p DoltDatabaseProvider) DropDatabase(ctx *sql.Context, name string) error
 	}
 
 	delete(p.databases, dbKey)
+
+	return p.invalidateDbStateInAllSessions(ctx, name)
+}
+
+// invalidateDbStateInAllSessions removes the db state for this database from every session. This is necessary when a
+// database is dropped, so that other sessions don't use stale db state.
+func (p DoltDatabaseProvider) invalidateDbStateInAllSessions(ctx *sql.Context, name string) error {
+	runningServer := sqlserver.GetRunningServer()
+	if runningServer != nil {
+		sessionManager := runningServer.SessionManager()
+		err := sessionManager.Iter(func(session sql.Session) (bool, error) {
+			dsess, ok := session.(*dsess.DoltSession)
+			if !ok {
+				return false, fmt.Errorf("unexpected session type: %T", session)
+			}
+
+			// We need to invalidate this database state for EVERY session, even if other sessions aren't actively
+			// using this database, since they could still reference it with a db-qualified table name.
+			err := dsess.RemoveDbState(ctx, name)
+			if err != nil {
+				return true, err
+			}
+			return false, nil
+		})
+		if err != nil {
+			return err
+		}
+	}
+
 	return nil
 }
 

go/libraries/doltcore/sqle/dsess/session.go

@@ -136,7 +136,9 @@ func DSessFromSess(sess sql.Session) *DoltSession {
 // LookupDbState returns the session state for the database named
 func (d *DoltSession) lookupDbState(ctx *sql.Context, dbName string) (*DatabaseSessionState, bool, error) {
 	dbName = strings.ToLower(dbName)
+	d.mu.Lock()
 	dbState, ok := d.dbStates[dbName]
+	d.mu.Unlock()
 	if ok {
 		return dbState, ok, nil
 	}
@@ -165,12 +167,14 @@ func (d *DoltSession) lookupDbState(ctx *sql.Context, dbName string) (*DatabaseS
 		}
 	}
 
-	// If we got this far, we have a valid inital database state, so add it to the session for future reuse
+	// If we got this far, we have a valid initial database state, so add it to the session for future reuse
 	if err = d.AddDB(ctx, init); err != nil {
 		return nil, ok, err
 	}
 
+	d.mu.Lock()
 	dbState, ok = d.dbStates[dbName]
+	d.mu.Unlock()
 	if !ok {
 		return nil, false, sql.ErrDatabaseNotFound.New(dbName)
 	}
@@ -190,6 +194,14 @@ func (d *DoltSession) LookupDbState(ctx *sql.Context, dbName string) (*DatabaseS
 	return s, ok, nil
 }
 
+// RemoveDbState invalidates any cached db state in this session, for example, if a database is dropped.
+func (d *DoltSession) RemoveDbState(_ *sql.Context, dbName string) error {
+	d.mu.Lock()
+	defer d.mu.Unlock()
+	delete(d.dbStates, strings.ToLower(dbName))
+	return nil
+}
+
 // Flush flushes all changes sitting in edit sessions to the session root for the database named. This normally
 // happens automatically as part of statement execution, and is only necessary when the session is manually batched (as
 // for bulk SQL import)
@@ -1063,6 +1075,9 @@ func (d *DoltSession) setHeadRefSessionVar(ctx *sql.Context, db, value string) e
 }
 
 func (d *DoltSession) setForeignKeyChecksSessionVar(ctx *sql.Context, key string, value interface{}) error {
+	d.mu.Lock()
+	defer d.mu.Unlock()
+
 	convertedVal, err := sqltypes.Int64.Convert(value)
 	if err != nil {
 		return err
@@ -1091,7 +1106,9 @@ func (d *DoltSession) setForeignKeyChecksSessionVar(ctx *sql.Context, key string
 }
 
 // HasDB returns true if |sess| is tracking state for this database.
-func (d *DoltSession) HasDB(ctx *sql.Context, dbName string) bool {
+func (d *DoltSession) HasDB(_ *sql.Context, dbName string) bool {
+	d.mu.Lock()
+	defer d.mu.Unlock()
 	_, ok := d.dbStates[strings.ToLower(dbName)]
 	return ok
 }
@@ -1106,7 +1123,9 @@ func (d *DoltSession) AddDB(ctx *sql.Context, dbState InitialDbState) error {
 	DefineSystemVariablesForDB(db.Name())
 
 	sessionState := NewEmptyDatabaseSessionState()
+	d.mu.Lock()
 	d.dbStates[strings.ToLower(db.Name())] = sessionState
+	d.mu.Unlock()
 	sessionState.dbName = db.Name()
 	sessionState.db = db
 

go/libraries/doltcore/sqle/enginetest/dolt_server_test.go

@@ -301,12 +301,133 @@ var DoltBranchMultiSessionScriptTests = []queries.ScriptTest{
 	},
 }
 
+// DropDatabaseMultiSessionScriptTests test that when dropping a database, other sessions are properly updated
+// and don't get left with old state that causes incorrect results.
+// Note: this test needs to be run against a real Dolt sql-server, and not just with our transaction test scripts,
+// because the transaction tests currently have a different behavior for session management and don't emulate prod.
+var DropDatabaseMultiSessionScriptTests = []queries.ScriptTest{
+	{
+		Name: "Test multi-session behavior for dropping databases",
+		SetUpScript: []string{
+			"create database db01;",
+			"create table db01.t01 (pk int primary key);",
+			"insert into db01.t01 values (101), (202), (303);",
+		},
+		Assertions: []queries.ScriptTestAssertion{
+			{
+				Query:    "/* client a */ use db01;",
+				Expected: []sql.Row{},
+			},
+			{
+				Query:    "/* client b */ use db01;",
+				Expected: []sql.Row{},
+			},
+			{
+				Query:    "/* client a */ show tables;",
+				Expected: []sql.Row{{"t01"}},
+			},
+			{
+				Query:    "/* client b */ show tables;",
+				Expected: []sql.Row{{"t01"}},
+			},
+			{
+				Query:    "/* client a */ drop database db01;",
+				Expected: []sql.Row{},
+			},
+			{
+				// TODO: This test runner doesn't currently support asserting against null values
+				Query:    "/* client a */ select database() is NULL;",
+				Expected: []sql.Row{{1}},
+			},
+			{
+				Query:    "/* client a */ show databases like 'db01';",
+				Expected: []sql.Row{},
+			},
+			{
+				Query:    "/* client a */ create database db01;",
+				Expected: []sql.Row{},
+			},
+			{
+				Query:    "/* client b */ select database();",
+				Expected: []sql.Row{{"db01"}},
+			},
+			{
+				Query:    "/* client b */ show tables;",
+				Expected: []sql.Row{},
+			},
+		},
+	},
+	{
+		Name: "Test multi-session behavior for dropping databases with a revision db",
+		SetUpScript: []string{
+			"create database db01;",
+			"use db01;",
+			"create table db01.t01 (pk int primary key);",
+			"insert into db01.t01 values (101), (202), (303);",
+			"call dolt_commit('-Am', 'commit on main');",
+			"call dolt_checkout('-b', 'branch1');",
+			"insert into db01.t01 values (1001), (2002), (3003);",
+			"call dolt_commit('-Am', 'commit on branch1');",
+		},
+		Assertions: []queries.ScriptTestAssertion{
+			{
+				Query:    "/* client a */ use db01;",
+				Expected: []sql.Row{},
+			},
+			{
+				Query:    "/* client b */ use `db01/branch1`;",
+				Expected: []sql.Row{},
+			},
+			{
+				Query:    "/* client a */ show tables;",
+				Expected: []sql.Row{{"t01"}},
+			},
+			{
+				Query:    "/* client b */ show tables;",
+				Expected: []sql.Row{{"t01"}},
+			},
+			{
+				Query:    "/* client a */ drop database db01;",
+				Expected: []sql.Row{},
+			},
+			{
+				// TODO: This test runner doesn't currently support asserting against null values
+				Query:    "/* client a */ select database() is NULL;",
+				Expected: []sql.Row{{1}},
+			},
+			{
+				Query:    "/* client a */ show databases like 'db01';",
+				Expected: []sql.Row{},
+			},
+			{
+				Query:    "/* client a */ create database db01;",
+				Expected: []sql.Row{},
+			},
+			{
+				// At this point, this is an invalid revision database, and any queries against it will fail.
+				Query:    "/* client b */ select database();",
+				Expected: []sql.Row{{"db01/branch1"}},
+			},
+			{
+				Query:          "/* client b */ show tables;",
+				ExpectedErrStr: "Error 1105: database not found: db01/branch1",
+			},
+		},
+	},
+}
+
 // TestDoltMultiSessionBehavior runs tests that exercise multi-session logic on a running SQL server. Statements
 // are sent through the server, from out of process, instead of directly to the in-process engine API.
 func TestDoltMultiSessionBehavior(t *testing.T) {
 	testMultiSessionScriptTests(t, DoltBranchMultiSessionScriptTests)
 }
 
+// TestDropDatabaseMultiSessionBehavior tests that dropping a database from one session correctly updates state
+// in other sessions.
+func TestDropDatabaseMultiSessionBehavior(t *testing.T) {
+	testMultiSessionScriptTests(t, DropDatabaseMultiSessionScriptTests)
+}
+
 func testMultiSessionScriptTests(t *testing.T, tests []queries.ScriptTest) {
 	for _, test := range tests {
 		t.Run(test.Name, func(t *testing.T) {