fix: preserve loaded extensions when starting TP service (#754)

## What do these changes do?

Fix Python Database.serve() to prepare TP service without rebuilding
planner, preserving LOADed extensions. Add C++ regression coverage and
Python GDS TP E2E test.

## Related issue number

Fixes #750

---------

Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>

Author: zhanglei1949

include/neug/main/neug_db.h

@@ -276,6 +276,18 @@ class NeugDB {
    */
   void CloseAllConnection();
 
+  /**
+   * @brief Prepare an opened database for TP service without rebuilding
+   * planner.
+   *
+   * This closes local AP connections, persists and refreshes the live graph
+   * when needed, then rebuilds query runtime handles against the refreshed
+   * graph. The planner and its metadata registry are intentionally preserved so
+   * runtime extension registrations loaded in AP mode stay available in TP
+   * mode.
+   */
+  void PrepareForServing();
+
   inline const PropertyGraph& graph() const {
     return snapshot_store_->CurrentSnapshot();
   }
@@ -306,22 +318,25 @@ class NeugDB {
   void initAllocators(const std::string& allocator_dir);
   void openGraphAndIngestWals();
   void ingestWals(IWalParser& parser, PropertyGraph& graph);
+  void initPlanner();
+  void initQueryRuntime();
   void initPlannerAndQueryProcessor();
   std::shared_ptr<Checkpoint> consumeLiveGraphAndCommitCheckpoint(
       CheckpointSession& checkpoint_session);
-
   /**
-   * @brief Create a checkpoint after WAL recovery and keep the DB open.
-   *
-   * A recovery checkpoint publishes the recovered graph, reopens the graph from
-   * the published checkpoint so the live store owns checkpoint files, and rolls
-   * back to the previous checkpoint if reopening fails.
-   *
-   * A durable checkpoint is a transaction timeline reset boundary: it always
-   * compacts storage timestamps before dumping, and a successful checkpoint
-   * resets last_ts_ to 0. Must not be called while a NeugDBService is running.
+   * @brief Create a checkpoint and keep the DB open on the published graph.
+   *
+   * This publishes the current live graph and reopens it from the published
+   * checkpoint so the live store owns checkpoint files. If reopening fails, the
+   * published checkpoint is discarded and the checkpoint manager is restored to
+   * the previous checkpoint generation; callers should treat the failure as
+   * fatal. It is shared by recovery checkpoint and AP-to-TP service
+   * preparation. A durable checkpoint is a transaction timeline reset boundary:
+   * it always compacts storage timestamps before dumping, and a successful
+   * checkpoint resets last_ts_ to 0. Must not be called while a NeugDBService
+   * is running.
    */
-  void createCheckpointAfterRecovery();
+  void createCheckpointAndRefreshLiveGraph();
 
   /**
    * @brief Create a checkpoint while closing the DB.

src/main/neug_db.cc

@@ -145,7 +145,7 @@ bool NeugDB::Open(const NeugDBConfig& config) {
     if (last_ts_ > 0 && config.checkpoint_on_recovery &&
         config_.mode == DBMode::READ_WRITE) {
       LOG(INFO) << "Creating checkpoint after recovery at ts " << last_ts_;
-      createCheckpointAfterRecovery();
+      createCheckpointAndRefreshLiveGraph();
     }
     if (config_.mode == DBMode::READ_WRITE) {
       checkpoint_mgr_.CleanupRetiredCheckpoints();
@@ -212,6 +212,17 @@ void NeugDB::RemoveConnection(std::shared_ptr<Connection> conn) {
 
 void NeugDB::CloseAllConnection() { connection_manager_->Close(); }
 
+void NeugDB::PrepareForServing() {
+  if (IsClosed()) {
+    THROW_RUNTIME_ERROR("NeugDB instance is not ready for serving!");
+  }
+  CloseAllConnection();
+  if (config_.mode == DBMode::READ_WRITE) {
+    createCheckpointAndRefreshLiveGraph();
+  }
+  initQueryRuntime();
+}
+
 void NeugDB::preprocessConfig() {
   if (config_.max_thread_num < 0) {
     THROW_INVALID_ARGUMENT_EXCEPTION(
@@ -334,16 +345,20 @@ void NeugDB::ingestWals(IWalParser& parser, PropertyGraph& graph) {
   last_ts_ = parser.last_ts();
 }
 
-void NeugDB::initPlannerAndQueryProcessor() {
+void NeugDB::initPlanner() {
   if (config_.planner_kind == "gopt") {
-    // Gopt planner is the default planner, so we don't need to create it.
     planner_ = std::make_shared<GOptPlanner>();
   } else {
     THROW_INVALID_ARGUMENT_EXCEPTION("Invalid planner kind: " +
                                      config_.planner_kind);
   }
   LOG(INFO) << "Finish initializing planner";
+}
 
+void NeugDB::initQueryRuntime() {
+  if (!planner_) {
+    THROW_RUNTIME_ERROR("Planner is not initialized");
+  }
   global_query_cache_ = std::make_shared<execution::GlobalQueryCache>(planner_);
 
   query_processor_ = std::make_shared<QueryProcessor>(
@@ -354,6 +369,11 @@ void NeugDB::initPlannerAndQueryProcessor() {
       *snapshot_store_, planner_, query_processor_, config_);
 }
 
+void NeugDB::initPlannerAndQueryProcessor() {
+  initPlanner();
+  initQueryRuntime();
+}
+
 std::shared_ptr<Checkpoint> NeugDB::consumeLiveGraphAndCommitCheckpoint(
     CheckpointSession& checkpoint_session) {
   SnapshotGuard guard(*snapshot_store_);
@@ -365,7 +385,7 @@ std::shared_ptr<Checkpoint> NeugDB::consumeLiveGraphAndCommitCheckpoint(
   return published_checkpoint;
 }
 
-void NeugDB::createCheckpointAfterRecovery() {
+void NeugDB::createCheckpointAndRefreshLiveGraph() {
   std::lock_guard<std::mutex> lock(mutex_);
   auto previous_checkpoint = checkpoint_mgr_.CurrentCheckpoint();
   auto checkpoint_session = CheckpointSession::Begin(checkpoint_mgr_);
@@ -403,8 +423,8 @@ void NeugDB::createCheckpointAfterRecovery() {
     throw;
   }
 
-  // Replacing snapshot_store_ releases the consumed recovery graph before the
-  // retired checkpoint directory is removed.
+  // Replacing snapshot_store_ releases the consumed graph before the retired
+  // checkpoint directory is removed.
   previous_checkpoint.reset();
   checkpoint_mgr_.CleanupRetiredCheckpoints();
 

tests/unittest/test_connection.cc

@@ -16,13 +16,45 @@
 #include <gtest/gtest.h>
 #include <filesystem>
 
+#include "neug/compiler/extension/extension_api.h"
+#include "neug/compiler/function/neug_call_function.h"
+#include "neug/compiler/main/metadata_registry.h"
+#include "neug/compiler/transaction/transaction.h"
 #include "neug/main/connection.h"
 #include "neug/main/neug_db.h"
 #include "unittest/utils.h"
 
 namespace neug {
 
 namespace test {
+namespace {
+
+class PrepareForServingTestFunction : public function::NeugCallFunction {
+ public:
+  PrepareForServingTestFunction()
+      : NeugCallFunction(
+            "PREPARE_FOR_SERVING_TEST_EXTENSION", function::call_input_types{},
+            {{"name", ::neug::DataType(::neug::DataTypeId::kVarchar)}}) {}
+};
+
+struct PrepareForServingTestFunctionSet {
+  static constexpr const char* name = "PREPARE_FOR_SERVING_TEST_EXTENSION";
+  static function::function_set getFunctionSet() {
+    function::function_set function_set;
+    function_set.emplace_back(
+        std::make_unique<PrepareForServingTestFunction>());
+    return function_set;
+  }
+};
+
+bool HasPrepareForServingTestFunction() {
+  return main::MetadataRegistry::getCatalog()->containsFunction(
+      &transaction::DUMMY_TRANSACTION, PrepareForServingTestFunctionSet::name,
+      false);
+}
+
+}  // namespace
+
 class ConnectionTest : public ::testing::Test {
  protected:
   static constexpr const char* DB_DIR = "/tmp/connection_test";
@@ -183,6 +215,31 @@ TEST_F(ConnectionTest, TestExplicitReadAccessModeForCall) {
       << project_res.error().ToString();
 }
 
+TEST(ConnectionStandaloneTest, PrepareForServingPreservesLoadedExtensions) {
+  constexpr const char* db_dir = "/tmp/prepare_for_serving_test";
+  std::filesystem::remove_all(db_dir);
+
+  NeugDB db;
+  NeugDBConfig config;
+  config.data_dir = db_dir;
+  config.mode = DBMode::READ_WRITE;
+  config.checkpoint_on_close = false;
+  db.Open(config);
+
+  auto planner = db.GetPlanner();
+
+  extension::ExtensionAPI::registerFunction<PrepareForServingTestFunctionSet>(
+      catalog::CatalogEntryType::TABLE_FUNCTION_ENTRY);
+  ASSERT_TRUE(HasPrepareForServingTestFunction());
+
+  db.PrepareForServing();
+  EXPECT_EQ(planner, db.GetPlanner());
+  EXPECT_TRUE(HasPrepareForServingTestFunction());
+
+  db.Close();
+  std::filesystem::remove_all(db_dir);
+}
+
 // Test Parallel Execution
 TEST_F(ConnectionTest, TestParallelExecutionAtomicity) {
   NeugDB db;

tools/python_bind/neug/database.py

@@ -339,6 +339,12 @@ def serve(
             )
         except KeyboardInterrupt:
             self.stop_serving()
+            raise
+        except Exception:
+            self._serving = False
+            raise
+        if blocking:
+            self._serving = False
         return endpoint
 
     def stop_serving(self):

tools/python_bind/src/py_database.cc

@@ -92,6 +92,7 @@ std::string PyDatabase::serve(int port, const std::string& host,
                               int32_t thread_num, bool blocking,
                               bool auto_compaction) {
 #ifdef BUILD_HTTP_SERVER
+  std::lock_guard<std::recursive_mutex> lock(mtx_);
   if (!database) {
     THROW_RUNTIME_ERROR("Database is not initialized.");
   }
@@ -109,20 +110,8 @@ std::string PyDatabase::serve(int port, const std::string& host,
         ". Must be less than or equal to database max_thread_num: " +
         std::to_string(database->config().max_thread_num) + ".");
   }
-  /**
-   * Attention here: We utilize the NeugDBService to start the server, based on
-   * database. But we need to make some changes to the NeugDB to make it works
-   * well for service mode, where concurrent queries will be processed.
-   *
-   * But before all, we need to close the database, dump the data to disk, and
-   * then reload the database with VersionManager and multiple contexts. By
-   * doing this, we make sure all changes made during AP mode is persisted.
-   */
-
-  std::lock_guard<std::recursive_mutex> lock(mtx_);
 
-  database->Close();
-  database->Open(database->config());
+  database->PrepareForServing();
   neug::ServiceConfig config;
   config.query_port = port;
   config.host_str = host;
@@ -135,11 +124,17 @@ std::string PyDatabase::serve(int port, const std::string& host,
 #endif
 
   service_ = std::make_unique<neug::NeugDBService>(*database, config);
-  if (blocking) {
-    service_->run_and_wait_for_exit();
-    return "";
+  try {
+    if (blocking) {
+      service_->run_and_wait_for_exit();
+      service_.reset();
+      return "";
+    }
+    return service_->Start();
+  } catch (...) {
+    service_.reset();
+    throw;
   }
-  return service_->Start();
 #else
   THROW_RUNTIME_ERROR("HTTP server is not enabled in this build.");
 #endif

tools/python_bind/tests/test_gds.py

@@ -26,8 +26,10 @@
 sys.path.append(os.path.join(os.path.dirname(__file__), "../"))
 
 from conftest import ensure_result_cnt_gt_zero
+from conftest import wait_for_server_ready
 
 from neug.database import Database
+from neug.session import Session
 
 
 def _shown_projected_graph_names(conn):
@@ -132,6 +134,60 @@ def test_run_cdlp(tmp_path):
             assert isinstance(row[1], int), "label should be an integer"
 
 
+def test_loaded_gds_extension_available_after_starting_tp_service(
+    tmp_path, unused_tcp_port
+):
+    """LOAD gds in AP mode, then query a GDS procedure through TP service."""
+    db_dir = tmp_path / "gds_tp_service_db"
+    db = Database(db_path=str(db_dir), mode="w")
+    db.load_builtin_dataset("tinysnb")
+    conn = db.connect()
+    served = False
+    session = None
+    try:
+        conn.execute(
+            "CALL project_graph("
+            "'person_knows', "
+            "['person'], "
+            "{'[person, knows, person]': ''}"
+            ");"
+        )
+        conn.execute("LOAD gds;")
+        conn.close()
+
+        endpoint = db.serve(
+            port=unused_tcp_port,
+            host="localhost",
+            blocking=False,
+            auto_compaction=False,
+        )
+        served = True
+        wait_for_server_ready(endpoint)
+        session = Session.open(endpoint, timeout="10s")
+        rows = list(
+            session.execute(
+                """
+                CALL cdlp('person_knows', {concurrency: 1})
+                YIELD node, label
+                RETURN node.id, label;
+                """
+            )
+        )
+
+        assert len(rows) > 0, "TP cdlp must return at least one row"
+        for row in rows:
+            assert len(row) == 2, "each row should have (id, label)"
+            assert isinstance(row[1], int), "label should be an integer"
+    finally:
+        if session is not None:
+            session.close()
+        if conn.is_open:
+            conn.close()
+        if served:
+            db.stop_serving()
+        db.close()
+
+
 @contextmanager
 def tinysnb_simple_connection(tmp_path):
     """Open a writable DB with tinysnb loaded and a simple projected graph