From 33927ec5404488f674456448b4fcac1114b2f46d Mon Sep 17 00:00:00 2001
From: Aaron Son <aaron@dolthub.com>
Date: Fri, 15 Dec 2023 14:39:48 -0800
Subject: [PATCH] go: sqle: DoltDatabaseProvider: If we encounter an error
 while creating a database, try to clean up after ourselves so we do not leave
 partially initialized database state around.

---
 .../doltcore/sqle/database_provider.go        | 11 +++++--
 .../sqle/enginetest/dolt_engine_test.go       | 30 +++++++++++++++++++
 2 files changed, 39 insertions(+), 2 deletions(-)

diff --git a/go/libraries/doltcore/sqle/database_provider.go b/go/libraries/doltcore/sqle/database_provider.go
index 3ce7d9f4ca..865a6b6418 100644
--- a/go/libraries/doltcore/sqle/database_provider.go
+++ b/go/libraries/doltcore/sqle/database_provider.go
@@ -362,7 +362,7 @@ func (p *DoltDatabaseProvider) CreateDatabase(ctx *sql.Context, name string) err
 	return p.CreateCollatedDatabase(ctx, name, sql.Collation_Default)
 }
 
-func (p *DoltDatabaseProvider) CreateCollatedDatabase(ctx *sql.Context, name string, collation sql.CollationID) error {
+func (p *DoltDatabaseProvider) CreateCollatedDatabase(ctx *sql.Context, name string, collation sql.CollationID) (err error) {
 	p.mu.Lock()
 	defer p.mu.Unlock()
 
@@ -373,10 +373,17 @@ func (p *DoltDatabaseProvider) CreateCollatedDatabase(ctx *sql.Context, name str
 		return fmt.Errorf("Cannot create DB, file exists at %s", name)
 	}
 
-	err := p.fs.MkDirs(name)
+	err = p.fs.MkDirs(name)
 	if err != nil {
 		return err
 	}
+	defer func() {
+		// We do not want to leave this directory behind if we do not
+		// successfully create the database.
+		if err != nil {
+			p.fs.Delete(name, true /* force / recursive */)
+		}
+	}()
 
 	newFs, err := p.fs.WithWorkingDir(name)
 	if err != nil {
diff --git a/go/libraries/doltcore/sqle/enginetest/dolt_engine_test.go b/go/libraries/doltcore/sqle/enginetest/dolt_engine_test.go
index 7005429dfb..7ce0ef8d26 100644
--- a/go/libraries/doltcore/sqle/enginetest/dolt_engine_test.go
+++ b/go/libraries/doltcore/sqle/enginetest/dolt_engine_test.go
@@ -2970,6 +2970,36 @@ func TestThreeWayMergeWithSchemaChangeScriptsPrepared(t *testing.T) {
 	})
 }
 
+// If CREATE DATABASE has an error within the DatabaseProvider, it should not
+// leave behind intermediate filesystem state.
+func TestCreateDatabaseErrorCleansUp(t *testing.T) {
+	dh := newDoltHarness(t)
+	require.NotNil(t, dh)
+	e, err := dh.NewEngine(t)
+	require.NoError(t, err)
+	require.NotNil(t, e)
+
+	dh.provider.(*sqle.DoltDatabaseProvider).InitDatabaseHook = func(_ *sql.Context, _ *sqle.DoltDatabaseProvider, name string, _ *env.DoltEnv) error {
+		if name == "cannot_create" {
+			return fmt.Errorf("there was an error initializing this database. abort!")
+		}
+		return nil
+	}
+
+	err = dh.provider.CreateDatabase(enginetest.NewContext(dh), "can_create")
+	require.NoError(t, err)
+
+	err = dh.provider.CreateDatabase(enginetest.NewContext(dh), "cannot_create")
+	require.Error(t, err)
+
+	fs := dh.multiRepoEnv.FileSystem()
+	exists, _ := fs.Exists("cannot_create")
+	require.False(t, exists)
+	exists, isDir := fs.Exists("can_create")
+	require.True(t, exists)
+	require.True(t, isDir)
+}
+
 // runMergeScriptTestsInBothDirections creates a new test run, named |name|, and runs the specified merge |tests|
 // in both directions (right to left merge, and left to right merge). If
 // |runAsPrepared| is true then the test scripts will be run using the prepared