handle multiple database instantiations for import caching

mxwli · mxwli · commit 085d3a9d4987 · 2024-03-11T11:47:31.000-04:00
diff --git a/src_cpp/include/py_database.h b/src_cpp/include/py_database.h
@@ -2,6 +2,7 @@
 
 #include "main/kuzu.h"
 #include "main/storage_driver.h"
+#include "cached_import/py_cached_import.h"
 #include "pybind_include.h" // IWYU pragma: keep (used for py:: namespace)
 #define PYBIND11_DETAILED_ERROR_MESSAGES
 using namespace kuzu::main;
diff --git a/src_cpp/kuzu_binding.cpp b/src_cpp/kuzu_binding.cpp
@@ -1,3 +1,4 @@
+#include "include/cached_import/py_cached_import.h"
 #include "include/py_connection.h"
 #include "include/py_database.h"
 #include "include/py_prepared_statement.h"
@@ -8,6 +9,10 @@ void bind(py::module& m) {
     PyConnection::initialize(m);
     PyPreparedStatement::initialize(m);
     PyQueryResult::initialize(m);
+    auto cleanImportCache = []() {
+        kuzu::importCache.reset();
+    };
+    m.add_object("_clean_import_cache", py::capsule(cleanImportCache));
 }
 
 PYBIND11_MODULE(_kuzu, m) {
diff --git a/src_cpp/py_database.cpp b/src_cpp/py_database.cpp
@@ -48,8 +48,14 @@ PyDatabase::PyDatabase(const std::string& databasePath, uint64_t bufferPoolSize,
     database = std::make_unique<Database>(databasePath, systemConfig);
     database->addBuiltInFunction(READ_PANDAS_FUNC_NAME, kuzu::PandasScanFunction::getFunctionSet());
     storageDriver = std::make_unique<kuzu::main::StorageDriver>(database.get());
+    py::gil_scoped_acquire acquire;
+    if (kuzu::importCache.get() == nullptr) {
+        kuzu::importCache = std::make_shared<kuzu::PythonCachedImport>();
+    }
 }
 
+PyDatabase::~PyDatabase() {}
+
 template<class T>
 void PyDatabase::scanNodeTable(const std::string& tableName, const std::string& propName,
     const py::array_t<uint64_t>& indices, py::array_t<T>& result, int numThreads) {
