Fix SPI stack leak, error logging, and search_path during DDL WAL replay

Author: khustup2

cpp/deeplake_pg/extension_init.cpp

@@ -9,7 +9,10 @@ extern "C" {
 #include <postgres.h>
 
 #include <catalog/namespace.h>
+#include <commands/dbcommands.h>
 #include <commands/defrem.h>
+#include <commands/extension.h>
+#include <miscadmin.h>
 #include <commands/vacuum.h>
 #include <nodes/nodeFuncs.h>
 #include <optimizer/planner.h>
@@ -25,7 +28,7 @@ extern "C" {
 
 #include "column_statistics.hpp"
 #include "deeplake_executor.hpp"
-#include "dl_catalog.hpp"
+#include "dl_wal.hpp"
 #include "pg_deeplake.hpp"
 #include "pg_version_compat.h"
 #include "sync_worker.hpp"
@@ -42,6 +45,7 @@ extern "C" {
 #include <climits>
 #include <cmath>
 #include <cstdint>
+#include <unistd.h>
 #include <map>
 #include <memory>
 #include <numeric>
@@ -70,6 +74,58 @@ bool stateless_enabled = false;      // Enable stateless catalog sync across ins
 
 namespace {
 
+bool is_ddl_log_suppressed()
+{
+    if (pg::table_storage::is_catalog_only_create()) {
+        return true;
+    }
+    if (creating_extension) {
+        return true;
+    }
+    const char* app_name = GetConfigOption("application_name", true, false);
+    return app_name && strcmp(app_name, "pg_deeplake_sync") == 0;
+}
+
+void append_to_ddl_log_if_needed(const char* command_tag, const char* object_identity, const char* query_string)
+{
+    if (!pg::stateless_enabled || is_ddl_log_suppressed()) {
+        return;
+    }
+    if (query_string == nullptr || query_string[0] == '\0') {
+        return;
+    }
+
+    auto root_path = pg::session_credentials::get_root_path();
+    if (root_path.empty()) {
+        root_path = pg::utils::get_deeplake_root_directory();
+    }
+    if (root_path.empty()) {
+        return;
+    }
+
+    try {
+        auto creds = pg::session_credentials::get_credentials();
+        const char* dbname = get_database_name(MyDatabaseId);
+        std::string db_name = dbname ? dbname : "postgres";
+        if (dbname) {
+            pfree(const_cast<char*>(dbname));
+        }
+
+        pg::dl_wal::ddl_log_entry entry;
+        entry.seq = pg::dl_wal::next_ddl_seq();
+        entry.origin_instance_id = pg::dl_wal::local_instance_id();
+        const char* current_search_path = GetConfigOption("search_path", true, false);
+        entry.search_path = current_search_path != nullptr ? current_search_path : "";
+        entry.command_tag = command_tag != nullptr ? command_tag : "";
+        entry.object_identity = object_identity != nullptr ? object_identity : "";
+        entry.ddl_sql = query_string;
+
+        pg::dl_wal::append_ddl_log(root_path, db_name, creds, entry);
+    } catch (const std::exception& e) {
+        elog(WARNING, "pg_deeplake: failed to append DDL to WAL log: %s", e.what());
+    }
+}
+
 bool is_count_star(TargetEntry* node)
 {
     if (node == nullptr || node->expr == nullptr || !IsA(node->expr, Aggref)) {
@@ -691,11 +747,11 @@ static void process_utility(PlannedStmt* pstmt,
             }
             if (!root_path.empty()) {
                 auto creds = pg::session_credentials::get_credentials();
-                pg::dl_catalog::database_meta db_meta;
+                pg::dl_wal::ensure_catalog(root_path, creds);
+                pg::dl_wal::database_meta db_meta;
                 db_meta.db_name = dbstmt->dbname;
                 db_meta.state = "dropping";
-                pg::dl_catalog::upsert_database(root_path, creds, db_meta);
-                pg::dl_catalog::bump_catalog_version(root_path, creds);
+                pg::dl_wal::upsert_database(root_path, creds, db_meta);
                 elog(LOG, "pg_deeplake: marked database '%s' as dropping in catalog", dbstmt->dbname);
             }
         } catch (const std::exception& e) {
@@ -727,7 +783,8 @@ static void process_utility(PlannedStmt* pstmt,
                 }
                 if (!root_path.empty()) {
                     auto creds = pg::session_credentials::get_credentials();
-                    pg::dl_catalog::database_meta db_meta;
+                    pg::dl_wal::ensure_catalog(root_path, creds);
+                    pg::dl_wal::database_meta db_meta;
                     db_meta.db_name = dbstmt->dbname;
                     db_meta.state = "ready";
 
@@ -748,8 +805,7 @@ static void process_utility(PlannedStmt* pstmt,
                         }
                     }
 
-                    pg::dl_catalog::upsert_database(root_path, creds, db_meta);
-                    pg::dl_catalog::bump_catalog_version(root_path, creds);
+                    pg::dl_wal::upsert_database(root_path, creds, db_meta);
                     elog(DEBUG1, "pg_deeplake: recorded CREATE DATABASE '%s' in catalog", dbstmt->dbname);
                 }
             } catch (const std::exception& e) {
@@ -758,6 +814,14 @@ static void process_utility(PlannedStmt* pstmt,
         }
     }
 
+    // Post-hook: record CREATE SCHEMA in DDL WAL log
+    if (IsA(pstmt->utilityStmt, CreateSchemaStmt)) {
+        CreateSchemaStmt* schemastmt = (CreateSchemaStmt*)pstmt->utilityStmt;
+        if (schemastmt->schemaname != nullptr) {
+            append_to_ddl_log_if_needed("CREATE SCHEMA", schemastmt->schemaname, queryString);
+        }
+    }
+
     // Post-process ALTER TABLE ADD COLUMN to add column to deeplake dataset
     if (IsA(pstmt->utilityStmt, AlterTableStmt)) {
         AlterTableStmt* stmt = (AlterTableStmt*)pstmt->utilityStmt;
@@ -1256,6 +1320,59 @@ static void process_utility(PlannedStmt* pstmt,
             }
         }
     }
+
+    if (nodeTag(pstmt->utilityStmt) == T_DropStmt) {
+        DropStmt* stmt = (DropStmt*)pstmt->utilityStmt;
+        if (stmt->removeType == OBJECT_SCHEMA) {
+            append_to_ddl_log_if_needed("DROP SCHEMA", nullptr, queryString);
+        } else if (stmt->removeType == OBJECT_TABLE) {
+            append_to_ddl_log_if_needed("DROP TABLE", nullptr, queryString);
+        } else if (stmt->removeType == OBJECT_INDEX) {
+            append_to_ddl_log_if_needed("DROP INDEX", nullptr, queryString);
+        } else if (stmt->removeType == OBJECT_VIEW) {
+            append_to_ddl_log_if_needed("DROP VIEW", nullptr, queryString);
+        }
+    }
+
+    if (IsA(pstmt->utilityStmt, CreateStmt)) {
+        CreateStmt* stmt = (CreateStmt*)pstmt->utilityStmt;
+        if (stmt->accessMethod != nullptr && std::strcmp(stmt->accessMethod, "deeplake") == 0) {
+            const char* schema_name = stmt->relation->schemaname ? stmt->relation->schemaname : "public";
+            std::string object_id = std::string(schema_name) + "." + stmt->relation->relname;
+            append_to_ddl_log_if_needed("CREATE TABLE", object_id.c_str(), queryString);
+        }
+    }
+
+    if (IsA(pstmt->utilityStmt, AlterTableStmt)) {
+        if (queryString != nullptr && strncasecmp(queryString, "ALTER TABLE", 11) == 0) {
+            AlterTableStmt* stmt = (AlterTableStmt*)pstmt->utilityStmt;
+            const char* schema_name = stmt->relation->schemaname ? stmt->relation->schemaname : "public";
+            std::string object_id = std::string(schema_name) + "." + stmt->relation->relname;
+            append_to_ddl_log_if_needed("ALTER TABLE", object_id.c_str(), queryString);
+        }
+    }
+
+    if (IsA(pstmt->utilityStmt, IndexStmt)) {
+        if (queryString && strncasecmp(queryString, "CREATE", 6) == 0 &&
+            strncasecmp(queryString, "CREATE TABLE", 12) != 0) {
+            IndexStmt* stmt = (IndexStmt*)pstmt->utilityStmt;
+            const char* schema_name = stmt->relation->schemaname ? stmt->relation->schemaname : "public";
+            std::string object_id = std::string(schema_name) + "." + stmt->relation->relname;
+            append_to_ddl_log_if_needed("CREATE INDEX", object_id.c_str(), queryString);
+        }
+    }
+
+    if (IsA(pstmt->utilityStmt, ViewStmt)) {
+        ViewStmt* stmt = (ViewStmt*)pstmt->utilityStmt;
+        const char* schema_name = stmt->view->schemaname ? stmt->view->schemaname : "public";
+        std::string object_id = std::string(schema_name) + "." + stmt->view->relname;
+        append_to_ddl_log_if_needed("CREATE VIEW", object_id.c_str(), queryString);
+    }
+
+    if (IsA(pstmt->utilityStmt, RenameStmt)) {
+        append_to_ddl_log_if_needed("RENAME", nullptr, queryString);
+    }
+
     if (IsA(pstmt->utilityStmt, VariableSetStmt)) {
         VariableSetStmt* vstmt = (VariableSetStmt*)pstmt->utilityStmt;
         if (vstmt->name != nullptr && pg_strcasecmp(vstmt->name, "search_path") == 0) {
@@ -1533,26 +1650,6 @@ static void executor_end(QueryDesc* query_desc)
     }
 
     if (pg::query_info::is_in_executor_context(query_desc)) {
-        if (query_desc->operation == CMD_INSERT || query_desc->operation == CMD_UPDATE ||
-            query_desc->operation == CMD_DELETE || query_desc->operation == CMD_UTILITY) {
-            // Use PG_TRY/CATCH to handle errors during flush without cascading aborts
-            PG_TRY();
-            {
-                if (!pg::table_storage::instance().flush_all()) {
-                    pg::table_storage::instance().rollback_all();
-                    ereport(ERROR, (errcode(ERRCODE_INTERNAL_ERROR), errmsg("Failed to flush table storage")));
-                }
-            }
-            PG_CATCH();
-            {
-                // Error occurred during flush - rollback and suppress to prevent cascade
-                // This prevents "Deeplake does not support transaction aborts" cascade
-                pg::table_storage::instance().rollback_all();
-                // Don't re-throw - let the transaction abort naturally
-                FlushErrorState();
-            }
-            PG_END_TRY();
-        }
         pg::query_info::pop_context(query_desc);
         pg::table_storage::instance().reset_requested_columns();
     }

cpp/deeplake_pg/pg_deeplake.cpp

@@ -1,6 +1,8 @@
 #include "pg_deeplake.hpp"
+#include "dl_wal.hpp"
 #include "logger.hpp"
 #include "table_storage.hpp"
+#include "utils.hpp"
 
 #include <deeplake_api/deeplake_api.hpp>
 #include <deeplake_core/deeplake_index_type.hpp>
@@ -9,6 +11,9 @@
 extern "C" {
 #endif
 
+#include <access/xact.h>
+#include <commands/dbcommands.h>
+#include <miscadmin.h>
 #include <storage/ipc.h>
 
 #ifdef __cplusplus
@@ -260,6 +265,8 @@ void save_index_metadata(Oid oid)
     if (SPI_execute(buf.data, false, 0) != SPI_OK_INSERT) {
         ereport(ERROR, (errcode(ERRCODE_INTERNAL_ERROR), errmsg("Failed to save metadata")));
     }
+
+    // Cross-instance propagation is driven by DDL WAL logging in ProcessUtility.
 }
 
 void load_index_metadata()
@@ -353,6 +360,23 @@ void deeplake_xact_callback(XactEvent event, void *arg)
     }
 }
 
+void deeplake_subxact_callback(SubXactEvent event,
+                               SubTransactionId my_subid,
+                               SubTransactionId parent_subid,
+                               void* arg)
+{
+    switch (event) {
+    case SUBXACT_EVENT_ABORT_SUB:
+        pg::table_storage::instance().rollback_subxact(my_subid);
+        break;
+    case SUBXACT_EVENT_COMMIT_SUB:
+        pg::table_storage::instance().commit_subxact(my_subid, parent_subid);
+        break;
+    default:
+        break;
+    }
+}
+
 void init_deeplake()
 {
     static bool initialized = false;
@@ -377,6 +401,7 @@ void init_deeplake()
     pg::table_storage::instance(); /// initialize table storage
 
     RegisterXactCallback(deeplake_xact_callback, nullptr);
+    RegisterSubXactCallback(deeplake_subxact_callback, nullptr);
 }
 
 } // namespace pg