Add checkpoint configuration parameters to the API bindings (#4739)

royi-luo · web-flow · commit 0a520d4daa55 · 2025-01-20T20:09:09.000-05:00
diff --git a/src_cpp/include/py_database.h b/src_cpp/include/py_database.h
@@ -18,7 +18,8 @@ class PyDatabase {
     static uint64_t getStorageVersion();
 
     explicit PyDatabase(const std::string& databasePath, uint64_t bufferPoolSize,
-        uint64_t maxNumThreads, bool compression, bool readOnly, uint64_t maxDBSize);
+        uint64_t maxNumThreads, bool compression, bool readOnly, uint64_t maxDBSize,
+        bool autoCheckpoint, int64_t checkpointThreshold);
 
     ~PyDatabase();
 
diff --git a/src_cpp/py_database.cpp b/src_cpp/py_database.cpp
@@ -11,10 +11,12 @@ using namespace kuzu::common;
 
 void PyDatabase::initialize(py::handle& m) {
     py::class_<PyDatabase>(m, "Database")
-        .def(py::init<const std::string&, uint64_t, uint64_t, bool, bool, uint64_t>(),
+        .def(
+            py::init<const std::string&, uint64_t, uint64_t, bool, bool, uint64_t, bool, int64_t>(),
             py::arg("database_path"), py::arg("buffer_pool_size") = 0,
             py::arg("max_num_threads") = 0, py::arg("compression") = true,
-            py::arg("read_only") = false, py::arg("max_db_size") = (uint64_t)1 << 43)
+            py::arg("read_only") = false, py::arg("max_db_size") = (uint64_t)1 << 43,
+            py::arg("auto_checkpoint") = true, py::arg("checkpoint_threshold") = -1)
         .def("scan_node_table_as_int64", &PyDatabase::scanNodeTable<std::int64_t>,
             py::arg("table_name"), py::arg("prop_name"), py::arg("indices"), py::arg("np_array"),
             py::arg("num_threads"))
@@ -44,9 +46,13 @@ uint64_t PyDatabase::getStorageVersion() {
 }
 
 PyDatabase::PyDatabase(const std::string& databasePath, uint64_t bufferPoolSize,
-    uint64_t maxNumThreads, bool compression, bool readOnly, uint64_t maxDBSize) {
-    auto systemConfig =
-        SystemConfig(bufferPoolSize, maxNumThreads, compression, readOnly, maxDBSize);
+    uint64_t maxNumThreads, bool compression, bool readOnly, uint64_t maxDBSize,
+    bool autoCheckpoint, int64_t checkpointThreshold) {
+    auto systemConfig = SystemConfig(bufferPoolSize, maxNumThreads, compression, readOnly,
+        maxDBSize, autoCheckpoint);
+    if (checkpointThreshold >= 0) {
+        systemConfig.checkpointThreshold = static_cast<uint64_t>(checkpointThreshold);
+    }
     database = std::make_unique<Database>(databasePath, systemConfig);
     database->addTableFunction(kuzu::PandasScanFunction::name,
         kuzu::PandasScanFunction::getFunctionSet());
diff --git a/src_py/database.py b/src_py/database.py
@@ -35,6 +35,8 @@ def __init__(
         lazy_init: bool = False,
         read_only: bool = False,
         max_db_size: int = (1 << 43),
+        auto_checkpoint: bool = True,
+        checkpoint_threshold: int = -1,
     ):
         """
         Parameters
@@ -73,6 +75,14 @@ def __init__(
              a better solution later. The value is default to 1 << 43 (8TB) under 64-bit
              environment and 1GB under 32-bit one.
 
+        auto_checkpoint: bool
+            If true, the database will automatically checkpoint when the size of
+            the WAL file exceeds the checkpoint threshold.
+
+        checkpoint_threshold: int
+            The threshold of the WAL file size in bytes. When the size of the
+            WAL file exceeds this threshold, the database will checkpoint if autoCheckpoint is true.
+
         """
         if database_path is None:
             database_path = ":memory:"
@@ -85,6 +95,8 @@ def __init__(
         self.compression = compression
         self.read_only = read_only
         self.max_db_size = max_db_size
+        self.auto_checkpoint = auto_checkpoint
+        self.checkpoint_threshold = checkpoint_threshold
         self.is_closed = False
 
         self._database: Any = None  # (type: _kuzu.Database from pybind11)
@@ -147,6 +159,8 @@ def init_database(self) -> None:
                 self.compression,
                 self.read_only,
                 self.max_db_size,
+                self.auto_checkpoint,
+                self.checkpoint_threshold,
             )
 
     def get_torch_geometric_remote_backend(
diff --git a/test/test_database.py b/test/test_database.py
@@ -117,3 +117,25 @@ def test_in_mem_database_no_db_path() -> None:
         conn.execute("CREATE (:person {name: 'Bob', age: 40});")
         with conn.execute("MATCH (p:person) RETURN p.*") as result:
             assert result.get_num_tuples() == 2
+
+
+def test_database_auto_checkpoint_config(tmp_path: Path) -> None:
+    with kuzu.Database(database_path=tmp_path, auto_checkpoint=False) as db:
+        assert not db.is_closed
+        assert db._database is not None
+
+        conn = kuzu.Connection(db)
+        with conn.execute("CALL current_setting('auto_checkpoint') RETURN *") as result:
+            assert result.get_num_tuples() == 1
+            assert result.get_next()[0] == "False"
+
+
+def test_database_checkpoint_threshold_config(tmp_path: Path) -> None:
+    with kuzu.Database(database_path=tmp_path, checkpoint_threshold=1234) as db:
+        assert not db.is_closed
+        assert db._database is not None
+
+        conn = kuzu.Connection(db)
+        with conn.execute("CALL current_setting('checkpoint_threshold') RETURN *") as result:
+            assert result.get_num_tuples() == 1
+            assert result.get_next()[0] == "1234"
