From 3a2a830a390c009f6816638be35ba6bc61492f85 Mon Sep 17 00:00:00 2001
From: Jackson Newhouse <jnewhouse@influxdata.com>
Date: Thu, 30 Jan 2025 16:01:41 -0800
Subject: [PATCH] fix: persist create database action. (#25944)

---
 influxdb3/tests/server/cli.rs    | 39 +++++++++++++++++++++++++++++++-
 influxdb3_catalog/src/catalog.rs | 22 ++++++++----------
 2 files changed, 48 insertions(+), 13 deletions(-)

diff --git a/influxdb3/tests/server/cli.rs b/influxdb3/tests/server/cli.rs
index 9abf552c22..c5aca13bcb 100644
--- a/influxdb3/tests/server/cli.rs
+++ b/influxdb3/tests/server/cli.rs
@@ -10,7 +10,6 @@ use std::{
     thread,
 };
 use test_helpers::tempfile::NamedTempFile;
-#[cfg(feature = "system-py")]
 use test_helpers::tempfile::TempDir;
 use test_helpers::{assert_contains, assert_not_contains};
 
@@ -848,6 +847,44 @@ async fn test_triggers_are_started() {
     }
 }
 
+#[test_log::test(tokio::test)]
+async fn test_database_create_persists() {
+    // create tmp dir for object store
+    let tmp_file = TempDir::new().unwrap();
+    let tmp_dir = tmp_file.path().to_str().unwrap();
+
+    let mut server = TestServer::configure()
+        .with_object_store_dir(tmp_dir)
+        .spawn()
+        .await;
+
+    let server_addr = server.client_addr();
+    let db_name = "foo";
+    let result = run_with_confirmation(&["create", "database", db_name, "--host", &server_addr]);
+    debug!(result = ?result, "create database");
+    assert_contains!(&result, "Database \"foo\" created successfully");
+
+    // restart the server
+    server.kill();
+
+    server = TestServer::configure()
+        .with_object_store_dir(tmp_dir)
+        .spawn()
+        .await;
+
+    let server_addr = server.client_addr();
+
+    let result = run_with_confirmation(&["show", "databases", "--host", &server_addr]);
+    assert_eq!(
+        r#"+---------------+
+| iox::database |
++---------------+
+| foo           |
++---------------+"#,
+        result
+    );
+}
+
 #[test_log::test(tokio::test)]
 async fn test_trigger_enable() {
     let plugin_file = create_plugin_file(
diff --git a/influxdb3_catalog/src/catalog.rs b/influxdb3_catalog/src/catalog.rs
index 6759c9def7..d8ba64a479 100644
--- a/influxdb3_catalog/src/catalog.rs
+++ b/influxdb3_catalog/src/catalog.rs
@@ -579,16 +579,6 @@ impl DatabaseSchema {
             catalog_batch.database_id,
             Arc::clone(&catalog_batch.database_name),
         );
-
-        // We need to special case when we create a DB via the commandline as
-        // this will only contain one op to create a database. If we don't
-        // startup of the database will fail
-        if catalog_batch.ops.len() == 1 {
-            if let CatalogOp::CreateDatabase(_) = catalog_batch.ops[0] {
-                return Ok(db_schema);
-            }
-        }
-
         let new_db = DatabaseSchema::new_if_updated_from_batch(&db_schema, catalog_batch)?
             .expect("database must be new");
         Ok(new_db)
@@ -702,10 +692,18 @@ trait UpdateDatabaseSchema {
 impl UpdateDatabaseSchema for CatalogOp {
     fn update_schema<'a>(
         &self,
-        schema: Cow<'a, DatabaseSchema>,
+        mut schema: Cow<'a, DatabaseSchema>,
     ) -> Result<Cow<'a, DatabaseSchema>> {
         match &self {
-            CatalogOp::CreateDatabase(_) => Ok(schema),
+            CatalogOp::CreateDatabase(create_database) => {
+                if create_database.database_id != schema.id
+                    || create_database.database_name != schema.name
+                {
+                    warn!("Create database call received by a mismatched DatabaseSchema. This should not be possible.")
+                }
+                schema.to_mut();
+                Ok(schema)
+            }
             CatalogOp::CreateTable(create_table) => create_table.update_schema(schema),
             CatalogOp::AddFields(field_additions) => field_additions.update_schema(schema),
             CatalogOp::CreateDistinctCache(distinct_cache_definition) => {