Support specifying max DB size in rust (#2963)

kuzudb · Feb 28, 2024 · e09b3ec · e09b3ec
1 parent e056583
commit e09b3ec
Show file tree

Hide file tree

Showing 4 changed files with 35 additions and 2 deletions.
diff --git a/tools/rust_api/include/kuzu_rs.h b/tools/rust_api/include/kuzu_rs.h
@@ -80,7 +80,8 @@ std::unique_ptr<std::vector<kuzu::common::LogicalType>> logical_type_get_struct_
 
 /* Database */
 std::unique_ptr<kuzu::main::Database> new_database(std::string_view databasePath,
-    uint64_t bufferPoolSize, uint64_t maxNumThreads, bool enableCompression, bool readOnly);
+    uint64_t bufferPoolSize, uint64_t maxNumThreads, bool enableCompression, bool readOnly,
+    uint64_t maxDBSize);
 
 void database_set_logging_level(kuzu::main::Database& database, const std::string& level);
 

diff --git a/tools/rust_api/src/database.rs b/tools/rust_api/src/database.rs
@@ -36,7 +36,11 @@ pub struct SystemConfig {
     /// When true, new columns will be compressed if possible
     /// Defaults to true
     enable_compression: bool,
+    /// Whether or not the database can be written to from this instance
+    /// Defaults to not read-only
     read_only: bool,
+    /// The maximum size of the database. Defaults to 8TB
+    max_db_size: u64,
 }
 
 impl Default for SystemConfig {
@@ -46,6 +50,8 @@ impl Default for SystemConfig {
             max_num_threads: 0,
             enable_compression: true,
             read_only: false,
+            // This is a little weird, but it's a temporary interface
+            max_db_size: u32::MAX as u64,
         }
     }
 }
@@ -67,6 +73,10 @@ impl SystemConfig {
         self.read_only = read_only;
         self
     }
+    pub fn max_db_size(mut self, max_db_size: u64) -> Self {
+        self.max_db_size = max_db_size;
+        self
+    }
 }
 
 impl Database {
@@ -86,6 +96,7 @@ impl Database {
                 config.max_num_threads,
                 config.enable_compression,
                 config.read_only,
+                config.max_db_size,
             )?),
         })
     }
@@ -172,4 +183,21 @@ mod tests {
         );
         Ok(())
     }
+
+    #[test]
+    fn test_database_max_size() -> Result<()> {
+        let temp_dir = tempfile::tempdir()?;
+        {
+            // Max DB size of 4GB should be fine
+            Database::new(
+                temp_dir.path(),
+                SystemConfig::default().max_db_size(1 << 32),
+            )?;
+        }
+        {
+            Database::new(temp_dir.path(), SystemConfig::default().max_db_size(0))
+                .expect_err("0 is not a valid max DB size");
+        }
+        Ok(())
+    }
 }
diff --git a/tools/rust_api/src/ffi.rs b/tools/rust_api/src/ffi.rs
@@ -125,6 +125,7 @@ pub(crate) mod ffi {
             maxNumThreads: u64,
             enableCompression: bool,
             readOnly: bool,
+            maxDBSize: u64,
         ) -> Result<UniquePtr<Database>>;
 
         fn database_set_logging_level(database: Pin<&mut Database>, level: &CxxString);

diff --git a/tools/rust_api/src/kuzu_rs.cpp b/tools/rust_api/src/kuzu_rs.cpp
@@ -65,7 +65,7 @@ std::unique_ptr<std::vector<kuzu::common::LogicalType>> logical_type_get_struct_
 }
 
 std::unique_ptr<Database> new_database(std::string_view databasePath, uint64_t bufferPoolSize,
-    uint64_t maxNumThreads, bool enableCompression, bool readOnly) {
+    uint64_t maxNumThreads, bool enableCompression, bool readOnly, uint64_t maxDBSize) {
     auto systemConfig = SystemConfig();
     if (bufferPoolSize > 0) {
         systemConfig.bufferPoolSize = bufferPoolSize;
@@ -75,6 +75,9 @@ std::unique_ptr<Database> new_database(std::string_view databasePath, uint64_t b
     }
     systemConfig.readOnly = readOnly;
     systemConfig.enableCompression = enableCompression;
+    if (maxDBSize != -1u) {
+        systemConfig.maxDBSize = maxDBSize;
+    }
     return std::make_unique<Database>(databasePath, systemConfig);
 }