Refactor pybind ownership around shared handles

Author: adsharma

src_cpp/include/py_connection.h

@@ -1,9 +1,11 @@
 #pragma once
 
+#include <memory>
 #include <unordered_map>
 
 #include "main/storage_driver.h"
 #include "py_database.h"
+#include "py_handle_state.h"
 #include "py_prepared_statement.h"
 #include "py_query_result.h"
 
@@ -29,8 +31,7 @@ class PyConnection {
         const py::dict& params);
 
     std::unique_ptr<PyQueryResult> query(const std::string& statement);
-    std::unique_ptr<PyQueryResult> queryAsArrow(const std::string& statement,
-        int64_t chunkSize);
+    std::unique_ptr<PyQueryResult> queryAsArrow(const std::string& statement, int64_t chunkSize);
 
     void setMaxNumThreadForExec(uint64_t numThreads);
 
@@ -65,12 +66,10 @@ class PyConnection {
         const LogicalType& type);
 
 private:
-    void checkOpen() const;
+    PyConnectionState& refState() const;
 
-    std::unique_ptr<StorageDriver> storageDriver;
-    std::unique_ptr<Connection> conn;
-    std::unordered_map<std::string, py::object> arrowTableRefs;
+    std::shared_ptr<PyConnectionState> state;
 
     static std::unique_ptr<PyQueryResult> checkAndWrapQueryResult(
-        std::unique_ptr<QueryResult>& queryResult);
+        std::unique_ptr<QueryResult>& queryResult, std::shared_ptr<PyConnectionState> state);
 };

src_cpp/include/py_database.h

@@ -1,7 +1,10 @@
 #pragma once
 
+#include <memory>
+
 #include "main/lbug.h"
 #include "main/storage_driver.h"
+#include "py_handle_state.h"
 #include "pybind_include.h" // IWYU pragma: keep (used for py:: namespace)
 #define PYBIND11_DETAILED_ERROR_MESSAGES
 using namespace lbug::main;
@@ -30,6 +33,5 @@ class PyDatabase {
         const py::array_t<uint64_t>& indices, py::array_t<T>& result, int numThreads);
 
 private:
-    std::unique_ptr<Database> database;
-    std::unique_ptr<StorageDriver> storageDriver;
+    std::shared_ptr<PyDatabaseState> state;
 };

src_cpp/include/py_handle_state.h

@@ -0,0 +1,94 @@
+#pragma once
+
+#include <memory>
+#include <string>
+#include <unordered_map>
+
+#include "common/exception/runtime.h"
+#include "main/lbug.h"
+#include "main/prepared_statement.h"
+#include "main/storage_driver.h"
+#include "pybind_include.h"
+
+struct PyDatabaseState {
+    std::unique_ptr<lbug::main::Database> database;
+    std::unique_ptr<lbug::main::StorageDriver> storageDriver;
+
+    ~PyDatabaseState() { closeNative(); }
+
+    void closeNative() {
+        storageDriver.reset();
+        database.reset();
+    }
+
+    lbug::main::Database& ref() const {
+        if (database == nullptr) {
+            throw lbug::common::RuntimeException("Database is closed.");
+        }
+        return *database;
+    }
+
+    lbug::main::StorageDriver& storage() const {
+        if (storageDriver == nullptr) {
+            throw lbug::common::RuntimeException("Database is closed.");
+        }
+        return *storageDriver;
+    }
+};
+
+struct PyConnectionState {
+    std::shared_ptr<PyDatabaseState> database;
+    std::unique_ptr<lbug::main::StorageDriver> storageDriver;
+    std::unique_ptr<lbug::main::Connection> conn;
+    std::unordered_map<std::string, py::object> arrowTableRefs;
+
+    ~PyConnectionState() { closeNative(); }
+
+    void closeNative() {
+        arrowTableRefs.clear();
+        conn.reset();
+        storageDriver.reset();
+        database.reset();
+    }
+
+    lbug::main::Connection& ref() const {
+        if (conn == nullptr) {
+            throw lbug::common::RuntimeException("Connection is closed.");
+        }
+        return *conn;
+    }
+
+    lbug::main::StorageDriver& storage() const {
+        if (storageDriver == nullptr) {
+            throw lbug::common::RuntimeException("Connection is closed.");
+        }
+        return *storageDriver;
+    }
+};
+
+struct PyPreparedStatementState {
+    std::shared_ptr<PyConnectionState> connection;
+    std::unique_ptr<lbug::main::PreparedStatement> preparedStatement;
+
+    lbug::main::PreparedStatement& ref() const {
+        if (preparedStatement == nullptr) {
+            throw lbug::common::RuntimeException("Prepared statement is closed.");
+        }
+        return *preparedStatement;
+    }
+};
+
+struct PyQueryResultState {
+    std::shared_ptr<PyConnectionState> connection;
+    std::shared_ptr<PyQueryResultState> parent;
+    std::unique_ptr<lbug::main::QueryResult> owned;
+    lbug::main::QueryResult* borrowed = nullptr;
+
+    lbug::main::QueryResult& ref() const {
+        auto* result = owned != nullptr ? owned.get() : borrowed;
+        if (result == nullptr) {
+            throw lbug::common::RuntimeException("Query result is closed.");
+        }
+        return *result;
+    }
+};

src_cpp/include/py_prepared_statement.h

@@ -2,6 +2,7 @@
 
 #include "main/lbug.h"
 #include "main/prepared_statement.h"
+#include "py_handle_state.h"
 #include "pybind_include.h"
 
 using namespace lbug::main;
@@ -17,5 +18,5 @@ class PyPreparedStatement {
     bool isSuccess() const;
 
 private:
-    std::unique_ptr<PreparedStatement> preparedStatement;
+    std::shared_ptr<PyPreparedStatementState> state;
 };

src_cpp/include/py_query_result.h

@@ -6,6 +6,7 @@
 #include "arrow_array.h"
 #include "common/arrow/arrow.h"
 #include "main/lbug.h"
+#include "py_handle_state.h"
 #include "pybind_include.h"
 
 using namespace lbug::main;
@@ -54,7 +55,7 @@ class PyQueryResult {
     size_t getNumTuples();
 
 private:
-    void checkOpen() const;
+    PyQueryResultState& refState() const;
 
     static py::dict convertNodeIdToPyDict(const lbug::common::nodeID_t& nodeId);
 
@@ -65,6 +66,5 @@ class PyQueryResult {
         const std::vector<std::string>& names, std::int64_t chunkSize, bool fallbackExtensionTypes);
 
 private:
-    QueryResult* queryResult = nullptr;
-    bool isOwned = false;
+    std::shared_ptr<PyQueryResultState> state;
 };