Stateless schema.

Author: khustup2

cpp/deeplake_pg/dl_catalog.cpp

@@ -29,6 +29,7 @@ constexpr const char* k_tables_name = "tables";
 constexpr const char* k_columns_name = "columns";
 constexpr const char* k_indexes_name = "indexes";
 constexpr const char* k_meta_name = "meta";
+constexpr const char* k_schemas_name = "schemas";
 constexpr const char* k_databases_name = "databases";
 
 // Shared (cluster-wide) path: {root}/__deeplake_catalog/{name}
@@ -184,6 +185,17 @@ deeplake_api::catalog_table_schema make_indexes_schema()
     return schema;
 }
 
+deeplake_api::catalog_table_schema make_schemas_schema()
+{
+    deeplake_api::catalog_table_schema schema;
+    schema.add("schema_name", deeplake_core::type::text(codecs::compression::null))
+        .add("owner", deeplake_core::type::text(codecs::compression::null))
+        .add("state", deeplake_core::type::text(codecs::compression::null))
+        .add("updated_at", deeplake_core::type::generic(nd::type::scalar(nd::dtype::int64)))
+        .set_primary_key("schema_name");
+    return schema;
+}
+
 deeplake_api::catalog_table_schema make_meta_schema()
 {
     deeplake_api::catalog_table_schema schema;
@@ -272,32 +284,35 @@ int64_t ensure_db_catalog(const std::string& root_path, const std::string& db_na
     const auto tables_path = join_db_path(root_path, db_name, k_tables_name);
     const auto columns_path = join_db_path(root_path, db_name, k_columns_name);
     const auto indexes_path = join_db_path(root_path, db_name, k_indexes_name);
+    const auto schemas_path = join_db_path(root_path, db_name, k_schemas_name);
     const auto meta_path = join_db_path(root_path, db_name, k_meta_name);
 
     try {
-        // Launch all 4 per-database catalog table creation in parallel
+        // Launch all 5 per-database catalog table creation in parallel
         icm::vector<async::promise<std::shared_ptr<deeplake_api::catalog_table>>> promises;
-        promises.reserve(4);
+        promises.reserve(5);
         promises.push_back(
             deeplake_api::open_or_create_catalog_table(tables_path, make_tables_schema(), icm::string_map<>(creds)));
         promises.push_back(
             deeplake_api::open_or_create_catalog_table(columns_path, make_columns_schema(), icm::string_map<>(creds)));
         promises.push_back(
             deeplake_api::open_or_create_catalog_table(indexes_path, make_indexes_schema(), icm::string_map<>(creds)));
+        promises.push_back(
+            deeplake_api::open_or_create_catalog_table(schemas_path, make_schemas_schema(), icm::string_map<>(creds)));
         promises.push_back(
             deeplake_api::open_or_create_catalog_table(meta_path, make_meta_schema(), icm::string_map<>(creds)));
 
         auto results = async::combine(std::move(promises)).get_future().get();
-        if (results.size() != 4) {
+        if (results.size() != 5) {
             elog(ERROR,
-                 "Failed to initialize per-db catalog at %s/%s: expected 4 catalog tables, got %zu",
+                 "Failed to initialize per-db catalog at %s/%s: expected 5 catalog tables, got %zu",
                  root_path.c_str(),
                  db_name.c_str(),
                  static_cast<size_t>(results.size()));
         }
 
-        // Initialize per-db meta table if empty (index 3 is meta)
-        auto& meta_table = results[3];
+        // Initialize per-db meta table if empty (index 4 is meta)
+        auto& meta_table = results[4];
         if (meta_table) {
             auto snapshot = meta_table->read().get_future().get();
             if (snapshot.row_count() == 0) {
@@ -447,6 +462,73 @@ std::vector<index_meta> load_indexes(const std::string&, const std::string&, icm
     return {};
 }
 
+std::vector<schema_meta> load_schemas(const std::string& root_path, const std::string& db_name, icm::string_map<> creds)
+{
+    std::vector<schema_meta> out;
+    try {
+        auto table = open_db_catalog_table(root_path, db_name, k_schemas_name, std::move(creds));
+        if (!table) {
+            return out;
+        }
+        auto snapshot = table->read().get_future().get();
+        if (snapshot.row_count() == 0) {
+            return out;
+        }
+
+        std::unordered_map<std::string, schema_meta> latest;
+        for (const auto& row : snapshot.rows()) {
+            auto schema_name_it = row.find("schema_name");
+            auto state_it = row.find("state");
+            if (schema_name_it == row.end() || state_it == row.end()) {
+                continue;
+            }
+
+            schema_meta meta;
+            meta.schema_name = deeplake_api::array_to_string(schema_name_it->second);
+            meta.state = deeplake_api::array_to_string(state_it->second);
+            auto owner_it = row.find("owner");
+            if (owner_it != row.end()) {
+                meta.owner = deeplake_api::array_to_string(owner_it->second);
+            }
+            auto updated_it = row.find("updated_at");
+            if (updated_it != row.end()) {
+                auto updated_vec = load_int64_vector(updated_it->second);
+                meta.updated_at = updated_vec.empty() ? 0 : updated_vec.front();
+            }
+
+            auto it = latest.find(meta.schema_name);
+            if (it == latest.end() || it->second.updated_at <= meta.updated_at) {
+                latest[meta.schema_name] = std::move(meta);
+            }
+        }
+
+        out.reserve(latest.size());
+        for (auto& [_, meta] : latest) {
+            if (meta.state == "ready") {
+                out.push_back(std::move(meta));
+            }
+        }
+        return out;
+    } catch (const std::exception& e) {
+        elog(DEBUG1, "Failed to load catalog schemas for db '%s': %s (may be old catalog)", db_name.c_str(), e.what());
+        return out;
+    } catch (...) {
+        elog(DEBUG1, "Failed to load catalog schemas for db '%s': unknown error (may be old catalog)", db_name.c_str());
+        return out;
+    }
+}
+
+void upsert_schema(const std::string& root_path, const std::string& db_name, icm::string_map<> creds, const schema_meta& meta)
+{
+    auto table = open_db_catalog_table(root_path, db_name, k_schemas_name, std::move(creds));
+    icm::string_map<nd::array> row;
+    row["schema_name"] = nd::adapt(meta.schema_name);
+    row["owner"] = nd::adapt(meta.owner);
+    row["state"] = nd::adapt(meta.state);
+    row["updated_at"] = nd::adapt(meta.updated_at == 0 ? now_ms() : meta.updated_at);
+    table->upsert(std::move(row)).get_future().get();
+}
+
 std::pair<std::vector<table_meta>, std::vector<column_meta>>
 load_tables_and_columns(const std::string& root_path, const std::string& db_name, icm::string_map<> creds)
 {

cpp/deeplake_pg/dl_catalog.hpp

@@ -40,6 +40,14 @@ struct index_meta
     int32_t order_type = 0;
 };
 
+struct schema_meta
+{
+    std::string schema_name;   // PK
+    std::string owner;
+    std::string state;         // "ready" or "dropping"
+    int64_t updated_at = 0;
+};
+
 struct database_meta
 {
     std::string db_name;       // PK
@@ -67,6 +75,10 @@ std::vector<index_meta> load_indexes(const std::string& root_path, const std::st
 std::pair<std::vector<table_meta>, std::vector<column_meta>>
 load_tables_and_columns(const std::string& root_path, const std::string& db_name, icm::string_map<> creds);
 
+// Per-database schema catalog
+std::vector<schema_meta> load_schemas(const std::string& root_path, const std::string& db_name, icm::string_map<> creds);
+void upsert_schema(const std::string& root_path, const std::string& db_name, icm::string_map<> creds, const schema_meta& meta);
+
 // Per-database upserts (write to {root}/{db_name}/__deeplake_catalog/)
 void upsert_table(const std::string& root_path, const std::string& db_name, icm::string_map<> creds, const table_meta& meta);
 void upsert_columns(const std::string& root_path, const std::string& db_name, icm::string_map<> creds, const std::vector<column_meta>& columns);

cpp/deeplake_pg/extension_init.cpp

@@ -9,7 +9,9 @@ extern "C" {
 #include <postgres.h>
 
 #include <catalog/namespace.h>
+#include <commands/dbcommands.h>
 #include <commands/defrem.h>
+#include <miscadmin.h>
 #include <commands/vacuum.h>
 #include <nodes/nodeFuncs.h>
 #include <optimizer/planner.h>
@@ -586,6 +588,35 @@ static void process_utility(PlannedStmt* pstmt,
                         }
                     }
                 }
+
+                // Mark schema as "dropping" in the S3 catalog
+                if (pg::stateless_enabled) {
+                    try {
+                        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()) {
+                            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_catalog::ensure_catalog(root_path, creds);
+                            pg::dl_catalog::ensure_db_catalog(root_path, db_name, creds);
+
+                            pg::dl_catalog::schema_meta s_meta;
+                            s_meta.schema_name = schema_name;
+                            s_meta.state = "dropping";
+                            pg::dl_catalog::upsert_schema(root_path, db_name, creds, s_meta);
+
+                            pg::dl_catalog::bump_db_catalog_version(root_path, db_name, pg::session_credentials::get_credentials());
+                            pg::dl_catalog::bump_catalog_version(root_path, pg::session_credentials::get_credentials());
+                        }
+                    } catch (const std::exception& e) {
+                        elog(WARNING, "pg_deeplake: failed to mark schema '%s' as dropping in catalog: %s", schema_name, e.what());
+                    }
+                }
             }
         } else if (stmt->removeType == OBJECT_DATABASE) {
             const char* query = "SELECT nspname, relname "
@@ -760,6 +791,40 @@ static void process_utility(PlannedStmt* pstmt,
         }
     }
 
+    // Post-hook: record CREATE SCHEMA in S3 catalog for multi-instance sync
+    if (IsA(pstmt->utilityStmt, CreateSchemaStmt) && pg::stateless_enabled) {
+        CreateSchemaStmt* schemastmt = (CreateSchemaStmt*)pstmt->utilityStmt;
+        try {
+            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() && schemastmt->schemaname != nullptr) {
+                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_catalog::ensure_catalog(root_path, creds);
+                pg::dl_catalog::ensure_db_catalog(root_path, db_name, creds);
+
+                pg::dl_catalog::schema_meta s_meta;
+                s_meta.schema_name = schemastmt->schemaname;
+                s_meta.state = "ready";
+                if (schemastmt->authrole != nullptr) {
+                    s_meta.owner = schemastmt->authrole->rolename;
+                }
+                pg::dl_catalog::upsert_schema(root_path, db_name, creds, s_meta);
+
+                pg::dl_catalog::bump_db_catalog_version(root_path, db_name, pg::session_credentials::get_credentials());
+                pg::dl_catalog::bump_catalog_version(root_path, pg::session_credentials::get_credentials());
+                elog(DEBUG1, "pg_deeplake: recorded CREATE SCHEMA '%s' in catalog", schemastmt->schemaname);
+            }
+        } catch (const std::exception& e) {
+            elog(DEBUG1, "pg_deeplake: failed to record CREATE SCHEMA in catalog: %s", e.what());
+        }
+    }
+
     // Post-process ALTER TABLE ADD COLUMN to add column to deeplake dataset
     if (IsA(pstmt->utilityStmt, AlterTableStmt)) {
         AlterTableStmt* stmt = (AlterTableStmt*)pstmt->utilityStmt;

cpp/deeplake_pg/sync_worker.cpp

@@ -309,6 +309,39 @@ void deeplake_sync_databases_from_catalog(const std::string& root_path, icm::str
     }
 }
 
+/**
+ * Sync schemas for a specific database from pre-loaded catalog data via libpq.
+ * Creates missing schemas in the target database.
+ */
+void deeplake_sync_schemas_for_db(const std::string& db_name,
+    const std::vector<pg::dl_catalog::schema_meta>& schemas)
+{
+    for (const auto& meta : schemas) {
+        if (meta.state == "dropping") {
+            continue;
+        }
+
+        // Skip system schemas
+        if (meta.schema_name == "public" || meta.schema_name == "pg_catalog" ||
+            meta.schema_name == "information_schema" ||
+            meta.schema_name.substr(0, 3) == "pg_") {
+            continue;
+        }
+
+        StringInfoData buf;
+        initStringInfo(&buf);
+        appendStringInfo(&buf, "CREATE SCHEMA IF NOT EXISTS %s",
+                         quote_identifier(meta.schema_name.c_str()));
+
+        if (execute_via_libpq(db_name.c_str(), buf.data)) {
+            elog(LOG, "pg_deeplake sync: created schema '%s' in database '%s'",
+                 meta.schema_name.c_str(), db_name.c_str());
+        }
+
+        pfree(buf.data);
+    }
+}
+
 /**
  * Sync tables for a specific database from pre-loaded catalog data via libpq.
  * Creates missing tables in the target database.
@@ -444,13 +477,19 @@ void sync_all_databases(
         return;
     }
 
-    // Step 5: For each changed database, load tables+columns and sync
+    // Step 5: For each changed database, load schemas first, then tables+columns and sync
     for (const auto& db_name : changed_dbs) {
         try {
+            // Sync schemas before tables so CREATE TABLE can find the target schema
+            auto schemas = pg::dl_catalog::load_schemas(root_path, db_name, creds);
+            if (!schemas.empty()) {
+                deeplake_sync_schemas_for_db(db_name, schemas);
+            }
+
             auto [tables, columns] = pg::dl_catalog::load_tables_and_columns(root_path, db_name, creds);
             deeplake_sync_tables_for_db(db_name, tables, columns);
-            elog(LOG, "pg_deeplake sync: synced %zu tables for database '%s'",
-                 tables.size(), db_name.c_str());
+            elog(LOG, "pg_deeplake sync: synced %zu schemas, %zu tables for database '%s'",
+                 schemas.size(), tables.size(), db_name.c_str());
         } catch (const std::exception& e) {
             elog(WARNING, "pg_deeplake sync: failed to sync database '%s': %s", db_name.c_str(), e.what());
         } catch (...) {

cpp/deeplake_pg/table_storage.cpp

@@ -289,6 +289,20 @@ void table_storage::save_table_metadata(const pg::table_data& table_data)
         }
         pg::dl_catalog::upsert_columns(root_dir, db_name, creds, columns);
 
+        // Belt-and-suspenders: ensure the schema is recorded even if
+        // the CREATE SCHEMA hook was missed (e.g., schema created before
+        // the extension was loaded, or via a different code path).
+        if (schema_name != "public") {
+            try {
+                pg::dl_catalog::schema_meta s_meta;
+                s_meta.schema_name = schema_name;
+                s_meta.state = "ready";
+                pg::dl_catalog::upsert_schema(root_dir, db_name, creds, s_meta);
+            } catch (...) {
+                elog(DEBUG1, "pg_deeplake: failed to upsert schema '%s' in catalog (non-fatal)", schema_name.c_str());
+            }
+        }
+
         pg::dl_catalog::bump_db_catalog_version(root_dir, db_name, session_credentials::get_credentials());
         pg::dl_catalog::bump_catalog_version(root_dir, session_credentials::get_credentials());
         catalog_version_ = pg::dl_catalog::get_db_catalog_version(root_dir, db_name, session_credentials::get_credentials());