diff --git a/.github/workflows/pg-extension-build.yaml b/.github/workflows/pg-extension-build.yaml index 26e38d3d72..6d5470c948 100644 --- a/.github/workflows/pg-extension-build.yaml +++ b/.github/workflows/pg-extension-build.yaml @@ -161,7 +161,7 @@ jobs: export PG_MAJOR_VERSION=18 export USER=postgres cd postgres/tests/py_tests || exit 1 - python3 -m pytest -v -m 'not slow and not tpch' --tb=short || { + python3 -m pytest -v -m 'not slow or tpch' --tb=short || { echo -e "${RED}Tests failed for PostgreSQL 18${DEFAULT}" exit 1 } diff --git a/DEEPLAKE_API_VERSION b/DEEPLAKE_API_VERSION new file mode 100644 index 0000000000..a84947d6ff --- /dev/null +++ b/DEEPLAKE_API_VERSION @@ -0,0 +1 @@ +4.5.0 diff --git a/cpp/deeplake_pg/dl_catalog.cpp b/cpp/deeplake_pg/dl_catalog.cpp new file mode 100644 index 0000000000..e8844d59b0 --- /dev/null +++ b/cpp/deeplake_pg/dl_catalog.cpp @@ -0,0 +1,391 @@ +#include "dl_catalog.hpp" + +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include + +extern "C" { +#include +#include +} + +#include + +namespace pg::dl_catalog { + +namespace { + +constexpr const char* k_catalog_dir = "__deeplake_catalog"; +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"; + +std::string join_path(const std::string& root, const std::string& name) +{ + if (!root.empty() && root.back() == '/') { + return root + k_catalog_dir + "/" + name; + } + return root + "/" + k_catalog_dir + "/" + name; +} + +std::shared_ptr +open_or_create_table(const std::string& path, deeplake_api::catalog_table_schema schema, icm::string_map<> creds) +{ + return deeplake_api::open_or_create_catalog_table(path, schema, std::move(creds)).get_future().get(); +} + +int64_t now_ms() +{ + using namespace std::chrono; + return duration_cast(system_clock::now().time_since_epoch()).count(); +} + +std::shared_ptr +open_catalog_table(const std::string& root_path, const std::string& name, icm::string_map<> creds) +{ + const auto path = join_path(root_path, name); + return deeplake_api::open_catalog_table(path, std::move(creds)).get_future().get(); +} + +template +std::vector load_vector(const nd::array& arr) +{ + std::vector out; + out.reserve(static_cast(arr.volume())); + for (int64_t i = 0; i < arr.volume(); ++i) { + out.push_back(arr.value(i)); + } + return out; +} + +std::vector load_int64_vector(const nd::array& arr) +{ + std::vector out; + out.reserve(static_cast(arr.volume())); + bool is_numeric = false; + try { + is_numeric = nd::dtype_is_numeric(arr.dtype()); + } catch (...) { + is_numeric = false; + } + if (is_numeric) { + try { + for (int64_t i = 0; i < arr.volume(); ++i) { + out.push_back(arr.value(i)); + } + return out; + } catch (...) { + out.clear(); + } + } + for (int64_t i = 0; i < arr.volume(); ++i) { + auto v = arr.value(i); + try { + out.push_back(std::stoll(std::string(v))); + } catch (...) { + out.push_back(0); + } + } + return out; +} + +} // namespace + +void ensure_catalog(const std::string& root_path, icm::string_map<> creds) +{ + const auto tables_path = join_path(root_path, k_tables_name); + const auto columns_path = join_path(root_path, k_columns_name); + const auto indexes_path = join_path(root_path, k_indexes_name); + const auto meta_path = join_path(root_path, k_meta_name); + + auto ensure_table = [&](const std::string& path, + deeplake_api::catalog_table_schema schema) -> std::shared_ptr { + bool exists = false; + try { + exists = deeplake_api::exists(path, icm::string_map<>(creds)).get_future().get(); + } catch (...) { + exists = false; + } + if (exists) { + bool is_catalog = false; + try { + is_catalog = deeplake_api::is_catalog_table(path, icm::string_map<>(creds)).get_future().get(); + } catch (...) { + is_catalog = false; + } + if (!is_catalog) { + elog(WARNING, + "Existing catalog path %s is not a catalog table. Recreating catalog table.", + path.c_str()); + try { + deeplake_api::delete_dataset(path, icm::string_map<>(creds)).get_future().get(); + } catch (...) { + elog(WARNING, "Failed to delete legacy dataset at %s", path.c_str()); + } + } + } + try { + return open_or_create_table(path, std::move(schema), icm::string_map<>(creds)); + } catch (const std::exception& e) { + elog(ERROR, "Failed to open or create catalog table at %s: %s", path.c_str(), e.what()); + } + return {}; + }; + + { + deeplake_api::catalog_table_schema schema; + schema.add("table_id", deeplake_core::type::text(codecs::compression::null)) + .add("schema_name", deeplake_core::type::text(codecs::compression::null)) + .add("table_name", deeplake_core::type::text(codecs::compression::null)) + .add("dataset_path", 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("table_id"); + ensure_table(tables_path, std::move(schema)); + } + + { + deeplake_api::catalog_table_schema schema; + // Use column_id (table_id:column_name) as primary key to support multiple columns per table + schema.add("column_id", deeplake_core::type::text(codecs::compression::null)) + .add("table_id", deeplake_core::type::text(codecs::compression::null)) + .add("column_name", deeplake_core::type::text(codecs::compression::null)) + .add("pg_type", deeplake_core::type::text(codecs::compression::null)) + .add("dl_type_json", deeplake_core::type::text(codecs::compression::null)) + .add("nullable", deeplake_core::type::generic(nd::type::scalar(nd::dtype::boolean))) + .add("position", deeplake_core::type::generic(nd::type::scalar(nd::dtype::int32))) + .set_primary_key("column_id"); + ensure_table(columns_path, std::move(schema)); + } + + { + deeplake_api::catalog_table_schema schema; + schema.add("table_id", deeplake_core::type::text(codecs::compression::null)) + .add("column_names", deeplake_core::type::text(codecs::compression::null)) + .add("index_type", deeplake_core::type::text(codecs::compression::null)) + .add("order_type", deeplake_core::type::generic(nd::type::scalar(nd::dtype::int32))) + .set_primary_key("table_id"); + ensure_table(indexes_path, std::move(schema)); + } + + deeplake_api::catalog_table_schema meta_schema; + meta_schema.add("catalog_version", deeplake_core::type::generic(nd::type::scalar(nd::dtype::int64))) + .add("updated_at", deeplake_core::type::generic(nd::type::scalar(nd::dtype::int64))) + .set_primary_key("catalog_version"); + auto meta_table = ensure_table(meta_path, std::move(meta_schema)); + auto snapshot = meta_table->read().get_future().get(); + if (snapshot.row_count() == 0) { + icm::string_map row; + row["catalog_version"] = nd::adapt(static_cast(1)); + row["updated_at"] = nd::adapt(now_ms()); + meta_table->insert(std::move(row)).get_future().get(); + } +} + +std::vector load_tables(const std::string& root_path, icm::string_map<> creds) +{ + std::vector out; + try { + auto table = open_catalog_table(root_path, k_tables_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 latest; + for (const auto& row : snapshot.rows()) { + auto table_id_it = row.find("table_id"); + auto schema_it = row.find("schema_name"); + auto table_it = row.find("table_name"); + auto path_it = row.find("dataset_path"); + auto state_it = row.find("state"); + auto updated_it = row.find("updated_at"); + if (table_id_it == row.end() || schema_it == row.end() || table_it == row.end() || path_it == row.end() || + state_it == row.end() || updated_it == row.end()) { + continue; + } + + table_meta meta; + meta.table_id = deeplake_api::array_to_string(table_id_it->second); + meta.schema_name = deeplake_api::array_to_string(schema_it->second); + meta.table_name = deeplake_api::array_to_string(table_it->second); + meta.dataset_path = deeplake_api::array_to_string(path_it->second); + meta.state = deeplake_api::array_to_string(state_it->second); + auto updated_vec = load_int64_vector(updated_it->second); + meta.updated_at = updated_vec.empty() ? 0 : updated_vec.front(); + + auto it = latest.find(meta.table_id); + if (it == latest.end() || it->second.updated_at <= meta.updated_at) { + latest[meta.table_id] = 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(WARNING, "Failed to load catalog tables: %s", e.what()); + return out; + } catch (...) { + elog(WARNING, "Failed to load catalog tables: unknown error"); + return out; + } +} + +std::vector load_columns(const std::string& root_path, icm::string_map<> creds) +{ + std::vector out; + try { + auto table = open_catalog_table(root_path, k_columns_name, std::move(creds)); + if (!table) { + return out; + } + auto snapshot = table->read().get_future().get(); + if (snapshot.row_count() == 0) { + return out; + } + + for (const auto& row : snapshot.rows()) { + auto table_id_it = row.find("table_id"); + auto column_name_it = row.find("column_name"); + auto pg_type_it = row.find("pg_type"); + auto dl_type_it = row.find("dl_type_json"); + auto nullable_it = row.find("nullable"); + auto position_it = row.find("position"); + + if (table_id_it == row.end() || column_name_it == row.end() || pg_type_it == row.end()) { + continue; + } + + column_meta meta; + meta.table_id = deeplake_api::array_to_string(table_id_it->second); + meta.column_name = deeplake_api::array_to_string(column_name_it->second); + meta.pg_type = deeplake_api::array_to_string(pg_type_it->second); + if (dl_type_it != row.end()) { + meta.dl_type_json = deeplake_api::array_to_string(dl_type_it->second); + } + if (nullable_it != row.end()) { + try { + meta.nullable = nullable_it->second.value(0); + } catch (...) { + meta.nullable = true; + } + } + if (position_it != row.end()) { + try { + meta.position = position_it->second.value(0); + } catch (...) { + auto pos_vec = load_int64_vector(position_it->second); + meta.position = pos_vec.empty() ? 0 : static_cast(pos_vec.front()); + } + } + + out.push_back(std::move(meta)); + } + return out; + } catch (const std::exception& e) { + elog(WARNING, "Failed to load catalog columns: %s", e.what()); + return out; + } catch (...) { + elog(WARNING, "Failed to load catalog columns: unknown error"); + return out; + } +} + +std::vector load_indexes(const std::string&, icm::string_map<>) +{ + return {}; +} + +void upsert_table(const std::string& root_path, icm::string_map<> creds, const table_meta& meta) +{ + auto table = open_catalog_table(root_path, k_tables_name, std::move(creds)); + icm::string_map row; + row["table_id"] = nd::adapt(meta.table_id); + row["schema_name"] = nd::adapt(meta.schema_name); + row["table_name"] = nd::adapt(meta.table_name); + row["dataset_path"] = nd::adapt(meta.dataset_path); + 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(); +} + +void upsert_columns(const std::string& root_path, icm::string_map<> creds, const std::vector& columns) +{ + if (columns.empty()) { + return; + } + auto table = open_catalog_table(root_path, k_columns_name, std::move(creds)); + for (const auto& col : columns) { + icm::string_map row; + // column_id is the composite key: table_id:column_name + row["column_id"] = nd::adapt(col.table_id + ":" + col.column_name); + row["table_id"] = nd::adapt(col.table_id); + row["column_name"] = nd::adapt(col.column_name); + row["pg_type"] = nd::adapt(col.pg_type); + row["dl_type_json"] = nd::adapt(col.dl_type_json); + row["nullable"] = nd::adapt(col.nullable); + row["position"] = nd::adapt(col.position); + table->upsert(std::move(row)).get_future().get(); + } +} + +int64_t get_catalog_version(const std::string& root_path, icm::string_map<> creds) +{ + try { + auto table = open_catalog_table(root_path, k_meta_name, std::move(creds)); + if (!table) { + return 0; + } + auto snapshot = table->read().get_future().get(); + if (snapshot.row_count() == 0) { + return 0; + } + int64_t max_version = 0; + for (const auto& row : snapshot.rows()) { + auto it = row.find("catalog_version"); + if (it == row.end()) { + continue; + } + auto values = load_int64_vector(it->second); + if (!values.empty()) { + max_version = std::max(max_version, values.front()); + } + } + return max_version; + } catch (const std::exception& e) { + elog(WARNING, "Failed to read catalog version: %s", e.what()); + return 0; + } catch (...) { + elog(WARNING, "Failed to read catalog version: unknown error"); + return 0; + } +} + +void bump_catalog_version(const std::string& root_path, icm::string_map<> creds) +{ + auto table = open_catalog_table(root_path, k_meta_name, std::move(creds)); + int64_t version = get_catalog_version(root_path, creds); + icm::string_map row; + row["catalog_version"] = nd::adapt(version + 1); + row["updated_at"] = nd::adapt(now_ms()); + table->insert(std::move(row)).get_future().get(); +} + +} // namespace pg::dl_catalog diff --git a/cpp/deeplake_pg/dl_catalog.hpp b/cpp/deeplake_pg/dl_catalog.hpp new file mode 100644 index 0000000000..11503fb30f --- /dev/null +++ b/cpp/deeplake_pg/dl_catalog.hpp @@ -0,0 +1,51 @@ +#pragma once + +#include + +#include +#include +#include + +namespace pg::dl_catalog { + +struct table_meta +{ + std::string table_id; + std::string schema_name; + std::string table_name; + std::string dataset_path; + std::string state; + int64_t updated_at = 0; +}; + +struct column_meta +{ + std::string table_id; + std::string column_name; + std::string pg_type; + std::string dl_type_json; + bool nullable = true; + int32_t position = 0; +}; + +struct index_meta +{ + std::string table_id; + std::string column_names; + std::string index_type; + int32_t order_type = 0; +}; + +void ensure_catalog(const std::string& root_path, icm::string_map<> creds); + +std::vector load_tables(const std::string& root_path, icm::string_map<> creds); +std::vector load_columns(const std::string& root_path, icm::string_map<> creds); +std::vector load_indexes(const std::string& root_path, icm::string_map<> creds); + +void upsert_table(const std::string& root_path, icm::string_map<> creds, const table_meta& meta); +void upsert_columns(const std::string& root_path, icm::string_map<> creds, const std::vector& columns); + +int64_t get_catalog_version(const std::string& root_path, icm::string_map<> creds); +void bump_catalog_version(const std::string& root_path, icm::string_map<> creds); + +} // namespace pg::dl_catalog diff --git a/cpp/deeplake_pg/extension_init.cpp b/cpp/deeplake_pg/extension_init.cpp index f0a4309241..174807039f 100644 --- a/cpp/deeplake_pg/extension_init.cpp +++ b/cpp/deeplake_pg/extension_init.cpp @@ -1,6 +1,7 @@ #include "deeplake_executor.hpp" #include "pg_deeplake.hpp" #include "pg_version_compat.h" +#include "sync_worker.hpp" #include "table_am.hpp" #include "table_ddl_lock.hpp" #include "table_scan.hpp" @@ -32,6 +33,8 @@ extern "C" { #include #include #include +#include +#include #include #include #include @@ -55,6 +58,7 @@ bool treat_numeric_as_double = true; // Treat numeric types as double by default bool print_progress_during_seq_scan = false; bool use_shared_mem_for_refresh = false; bool enable_dataset_logging = false; // Enable dataset operation logging for debugging +bool allow_custom_paths = true; // Allow dataset_path in CREATE TABLE options } // namespace pg @@ -132,7 +136,6 @@ void initialize_guc_parameters() nullptr // check_hook, assign_hook, show_hook ); - DefineCustomBoolVariable("pg_deeplake.print_runtime_stats", "Enable runtime statistics printing for pg_deeplake operations.", nullptr, // optional long description @@ -213,6 +216,17 @@ void initialize_guc_parameters() nullptr // check_hook, assign_hook, show_hook ); + DefineCustomBoolVariable("deeplake.allow_custom_paths", + "Allow custom dataset paths via USING deeplake WITH (dataset_path=...).", + "If disabled, dataset_path options are rejected and tables must use deeplake.root_path.", + &pg::allow_custom_paths, + true, + PGC_USERSET, + 0, + nullptr, + nullptr, + nullptr); + DefineCustomBoolVariable("pg_deeplake.enable_dataset_logging", "Enable operation logging for deeplake datasets.", "When enabled, all dataset operations (append_row, update_row, delete_row, etc.) " @@ -228,6 +242,36 @@ void initialize_guc_parameters() nullptr // check_hook, assign_hook, show_hook ); + // Sync worker GUC variables for stateless multi-instance support + DefineCustomIntVariable("deeplake.sync_interval_ms", + "Interval between catalog sync checks in milliseconds.", + "The background sync worker polls the catalog version at this interval. " + "When the version changes, tables are synced from the shared catalog.", + &deeplake_sync_interval_ms, // linked C variable + 2000, // default value (2 seconds) + 100, // min value + 60000, // max value (1 minute) + PGC_SIGHUP, // context - reloadable + GUC_UNIT_MS, // flags + nullptr, + nullptr, + nullptr // check_hook, assign_hook, show_hook + ); + + DefineCustomBoolVariable("deeplake.sync_enabled", + "Enable background sync worker for stateless mode.", + "When enabled, the background worker polls the catalog and automatically " + "syncs tables from the shared storage. This enables stateless multi-instance " + "deployments where tables created on one instance appear on others.", + &deeplake_sync_enabled, // linked C variable + true, // default value + PGC_SIGHUP, // context - reloadable + 0, // flags + nullptr, + nullptr, + nullptr // check_hook, assign_hook, show_hook + ); + // Initialize PostgreSQL memory tracking pg::memory_tracker::initialize_guc_parameters(); @@ -601,6 +645,7 @@ static void process_utility(PlannedStmt* pstmt, elog(DEBUG1, "stmt->accessMethod: %s", stmt->accessMethod); const bool deeplake_table = (stmt->accessMethod != nullptr && std::strcmp(stmt->accessMethod, "deeplake") == 0); if (deeplake_table && stmt->options != nullptr) { + List* new_options = NIL; ListCell* lc = nullptr; foreach (lc, stmt->options) { DefElem* def = (DefElem*)lfirst(lc); @@ -608,10 +653,11 @@ static void process_utility(PlannedStmt* pstmt, const char* ds_path = defGetString(def); pg::table_options::current().set_dataset_path(ds_path); elog(DEBUG1, "ds_path: %s", ds_path); + continue; } + new_options = lappend(new_options, def); } - list_free_deep(stmt->options); - stmt->options = NIL; + stmt->options = new_options; } } @@ -896,7 +942,7 @@ static void process_utility(PlannedStmt* pstmt, // Invalidate cached table data to force reload in current session RelationClose(relation); pg::table_storage::instance().erase_table(table_name); - pg::table_storage::instance().force_load_table_metadata(); + pg::table_storage::instance().mark_metadata_stale(); return; // Exit early after erasing table } catch (const base::exception& e) { @@ -916,7 +962,7 @@ static void process_utility(PlannedStmt* pstmt, // Column has been dropped from PostgreSQL catalog, now reload table_data // to pick up the updated TupleDesc (with the column marked as dropped) pg::table_storage::instance().erase_table(table_name); - pg::table_storage::instance().force_load_table_metadata(); + pg::table_storage::instance().mark_metadata_stale(); elog(INFO, "Reloaded table_data after DROP COLUMN for table '%s'", table_name.c_str()); return; // Exit early after reloading table } @@ -971,7 +1017,7 @@ static void process_utility(PlannedStmt* pstmt, // Invalidate cached table data to force reload in current session pg::table_storage::instance().erase_table(table_name); - pg::table_storage::instance().force_load_table_metadata(); + pg::table_storage::instance().mark_metadata_stale(); } catch (const base::exception& e) { ereport(ERROR, @@ -1080,9 +1126,45 @@ static void process_utility(PlannedStmt* pstmt, } pg::table_storage::instance().set_schema_name(std::move(schema_name)); } + // When root_path is set, auto-discover tables from the deeplake catalog + if (vstmt->name != nullptr && pg_strcasecmp(vstmt->name, "deeplake.root_path") == 0) { + // Reload table metadata from the catalog at the new root_path + // This enables stateless multi-instance support where tables are + // auto-discovered when pointing to a shared root_path + pg::table_storage::instance().force_load_table_metadata(); + } } } +// Check if the query string represents a pure SELECT statement (not CTAS, INSERT, etc.) +static bool is_pure_select_statement(const char* query_string) +{ + if (query_string == nullptr) { + return false; + } + + // Use PG_TRY to catch parse errors from invalid SQL (e.g., expression fragments + // from plpgsql functions during constant folding) + bool result = false; + PG_TRY(); + { + List* raw_parsetree_list = raw_parser(query_string, RAW_PARSE_DEFAULT); + if (raw_parsetree_list != NIL) { + RawStmt* raw_stmt = linitial_node(RawStmt, raw_parsetree_list); + result = nodeTag(raw_stmt->stmt) == T_SelectStmt; + } + } + PG_CATCH(); + { + // Parse failed - not a valid SQL statement (e.g., expression fragment) + FlushErrorState(); + result = false; + } + PG_END_TRY(); + + return result; +} + static PlannedStmt* deeplake_planner(Query* parse, const char* query_string, int32_t cursorOptions, ParamListInfo boundParams) { @@ -1094,7 +1176,7 @@ deeplake_planner(Query* parse, const char* query_string, int32_t cursorOptions, } PlannedStmt* planned_stmt = nullptr; - if (pg::use_deeplake_executor) { + if (pg::use_deeplake_executor && is_pure_select_statement(query_string)) { planned_stmt = deeplake_create_direct_execution_plan(parse, query_string, cursorOptions, boundParams); } @@ -1433,9 +1515,28 @@ PGDLLEXPORT void _PG_init() prev_set_rel_pathlist_hook = set_rel_pathlist_hook; set_rel_pathlist_hook = set_rel_pathlist; + // Initialize GUC parameters first (needed for sync worker config) + ::initialize_guc_parameters(); + + // Register background sync worker for stateless multi-instance support + BackgroundWorker worker; + memset(&worker, 0, sizeof(worker)); + + snprintf(worker.bgw_name, BGW_MAXLEN, "pg_deeplake sync worker"); + snprintf(worker.bgw_type, BGW_MAXLEN, "pg_deeplake sync worker"); + snprintf(worker.bgw_library_name, BGW_MAXLEN, "pg_deeplake"); + snprintf(worker.bgw_function_name, BGW_MAXLEN, "deeplake_sync_worker_main"); + + worker.bgw_flags = BGWORKER_SHMEM_ACCESS | BGWORKER_BACKEND_DATABASE_CONNECTION; + worker.bgw_start_time = BgWorkerStart_RecoveryFinished; + worker.bgw_restart_time = 5; // Restart after 5 seconds if it crashes + worker.bgw_notify_pid = 0; + worker.bgw_main_arg = (Datum)0; + + RegisterBackgroundWorker(&worker); + pg::install_signal_handlers(); pg::deeplake_table_am_routine::initialize(); - ::initialize_guc_parameters(); } PGDLLEXPORT void _PG_fini() diff --git a/cpp/deeplake_pg/pg_deeplake.cpp b/cpp/deeplake_pg/pg_deeplake.cpp index 061c5f8bb3..a1905af18b 100644 --- a/cpp/deeplake_pg/pg_deeplake.cpp +++ b/cpp/deeplake_pg/pg_deeplake.cpp @@ -5,6 +5,28 @@ #include #include +#ifdef __cplusplus +extern "C" { +#endif + +#include + +#ifdef __cplusplus +} +#endif + +namespace { + +// Exit handler that uses _exit() to avoid C++ static destructor crashes. +// PostgreSQL background workers (autovacuum, parallel workers, etc.) can crash +// during normal exit when C++ static objects are destroyed in unpredictable order. +void deeplake_quick_exit(int code, Datum arg) +{ + _exit(code); +} + +} // anonymous namespace + namespace pg { QueryDesc* query_info::current_query_desc = nullptr; @@ -339,6 +361,10 @@ void init_deeplake() } initialized = true; + // Register exit handler first (runs last due to LIFO order) to use _exit() + // and avoid C++ static destructor crashes in background workers. + on_proc_exit(deeplake_quick_exit, 0); + constexpr int THREAD_POOL_MULTIPLIER = 8; // Threads per CPU core for async operations deeplake_api::initialize(std::make_shared(), THREAD_POOL_MULTIPLIER * base::system_report::cpu_cores()); diff --git a/cpp/deeplake_pg/sync_worker.cpp b/cpp/deeplake_pg/sync_worker.cpp new file mode 100644 index 0000000000..1c07f33f23 --- /dev/null +++ b/cpp/deeplake_pg/sync_worker.cpp @@ -0,0 +1,239 @@ +#include "pg_deeplake.hpp" + +#ifdef __cplusplus +extern "C" { +#endif + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#ifdef __cplusplus +} +#endif + +#include "sync_worker.hpp" + +#include "dl_catalog.hpp" +#include "table_storage.hpp" +#include "utils.hpp" + +#include +#include + +// GUC variables +int deeplake_sync_interval_ms = 2000; // Default 2 seconds +bool deeplake_sync_enabled = true; + +namespace { + +// Worker state - use sig_atomic_t for signal safety +volatile sig_atomic_t got_sigterm = false; +volatile sig_atomic_t got_sighup = false; + +void deeplake_sync_worker_sigterm(SIGNAL_ARGS) +{ + int save_errno = errno; + got_sigterm = true; + SetLatch(MyLatch); + errno = save_errno; +} + +void deeplake_sync_worker_sighup(SIGNAL_ARGS) +{ + int save_errno = errno; + got_sighup = true; + SetLatch(MyLatch); + errno = save_errno; +} + +/** + * Sync tables from the deeplake catalog to PostgreSQL. + * + * This function checks the catalog for tables that exist in the deeplake + * catalog but not in PostgreSQL, and creates them. + */ +void deeplake_sync_tables_from_catalog(const std::string& root_path, icm::string_map<> creds) +{ + auto catalog_tables = pg::dl_catalog::load_tables(root_path, creds); + auto catalog_columns = pg::dl_catalog::load_columns(root_path, creds); + + for (const auto& meta : catalog_tables) { + // Skip tables marked as dropping + if (meta.state == "dropping") { + continue; + } + + const std::string qualified_name = meta.schema_name + "." + meta.table_name; + + // Check if table exists in PostgreSQL + auto* rel = makeRangeVar(pstrdup(meta.schema_name.c_str()), pstrdup(meta.table_name.c_str()), -1); + Oid relid = RangeVarGetRelid(rel, NoLock, true); + + if (!OidIsValid(relid)) { + // Table doesn't exist locally - create it + elog(LOG, "pg_deeplake sync: creating table %s from catalog", qualified_name.c_str()); + + // Gather columns for this table, sorted by position + std::vector table_columns; + for (const auto& col : catalog_columns) { + if (col.table_id == meta.table_id) { + table_columns.push_back(col); + } + } + std::sort(table_columns.begin(), table_columns.end(), + [](const auto& a, const auto& b) { return a.position < b.position; }); + + if (table_columns.empty()) { + elog(DEBUG1, "pg_deeplake sync: no columns found for table %s, skipping", qualified_name.c_str()); + continue; + } + + const char* qschema = quote_identifier(meta.schema_name.c_str()); + + StringInfoData buf; + initStringInfo(&buf); + + // Create schema if needed + appendStringInfo(&buf, "CREATE SCHEMA IF NOT EXISTS %s", qschema); + + pg::utils::spi_connector connector; + if (SPI_execute(buf.data, false, 0) != SPI_OK_UTILITY) { + elog(WARNING, "pg_deeplake sync: failed to create schema %s", meta.schema_name.c_str()); + pfree(buf.data); + continue; + } + + // Build CREATE TABLE statement directly from catalog metadata + // This avoids calling the SQL function create_deeplake_table which may not exist + // in the postgres database (extension might not be installed there) + resetStringInfo(&buf); + appendStringInfo(&buf, "CREATE TABLE %s (", qualified_name.c_str()); + + bool first = true; + for (const auto& col : table_columns) { + if (!first) { + appendStringInfoString(&buf, ", "); + } + first = false; + appendStringInfo(&buf, "%s %s", quote_identifier(col.column_name.c_str()), col.pg_type.c_str()); + } + + appendStringInfo(&buf, ") USING deeplake WITH (dataset_path=%s)", quote_literal_cstr(meta.dataset_path.c_str())); + + if (SPI_execute(buf.data, false, 0) != SPI_OK_UTILITY) { + // Don't log as warning - the dataset might not be available yet + // The sync worker will retry on the next cycle + elog(DEBUG1, "pg_deeplake sync: table %s not ready yet, will retry", qualified_name.c_str()); + } else { + elog(LOG, "pg_deeplake sync: successfully created table %s", qualified_name.c_str()); + } + + pfree(buf.data); + } + } +} + +} // anonymous namespace + +extern "C" { + +PGDLLEXPORT void deeplake_sync_worker_main(Datum main_arg) +{ + // Set up signal handlers + pqsignal(SIGTERM, deeplake_sync_worker_sigterm); + pqsignal(SIGHUP, deeplake_sync_worker_sighup); + + // Unblock signals + BackgroundWorkerUnblockSignals(); + + // Connect to the default database + BackgroundWorkerInitializeConnection("postgres", NULL, 0); + + elog(LOG, "pg_deeplake sync worker started"); + + int64_t last_catalog_version = 0; + + while (!got_sigterm) { + // Handle SIGHUP - reload configuration + if (got_sighup) { + got_sighup = false; + ProcessConfigFile(PGC_SIGHUP); + } + + // Skip if sync is disabled + if (!deeplake_sync_enabled) { + goto wait_for_latch; + } + + // Start a transaction for our work + SetCurrentStatementStartTimestamp(); + StartTransactionCommand(); + PushActiveSnapshot(GetTransactionSnapshot()); + + PG_TRY(); + { + // Initialize DeepLake (loads table metadata, etc.) + pg::init_deeplake(); + + 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(); + + // Ensure catalog exists + pg::dl_catalog::ensure_catalog(root_path, creds); + + // Use existing catalog version API to check for changes + int64_t current_version = pg::dl_catalog::get_catalog_version(root_path, creds); + + if (current_version != last_catalog_version) { + // Version changed - sync tables from catalog + deeplake_sync_tables_from_catalog(root_path, creds); + last_catalog_version = current_version; + elog(LOG, "pg_deeplake sync: synced tables (catalog version %ld)", current_version); + } + } + } + PG_CATCH(); + { + // Log error but don't crash - continue polling + EmitErrorReport(); + FlushErrorState(); + } + PG_END_TRY(); + + PopActiveSnapshot(); + CommitTransactionCommand(); + pgstat_report_stat(true); + + wait_for_latch: + // Wait for latch or timeout + (void)WaitLatch(MyLatch, + WL_LATCH_SET | WL_TIMEOUT | WL_EXIT_ON_PM_DEATH, + deeplake_sync_interval_ms, + PG_WAIT_EXTENSION); + ResetLatch(MyLatch); + } + + elog(LOG, "pg_deeplake sync worker shutting down"); + proc_exit(0); +} + +} // extern "C" diff --git a/cpp/deeplake_pg/sync_worker.hpp b/cpp/deeplake_pg/sync_worker.hpp new file mode 100644 index 0000000000..7361915090 --- /dev/null +++ b/cpp/deeplake_pg/sync_worker.hpp @@ -0,0 +1,5 @@ +#pragma once + +// GUC variables for sync worker configuration +extern int deeplake_sync_interval_ms; +extern bool deeplake_sync_enabled; diff --git a/cpp/deeplake_pg/table_am.cpp b/cpp/deeplake_pg/table_am.cpp index 8514859df6..f22388d9e4 100644 --- a/cpp/deeplake_pg/table_am.cpp +++ b/cpp/deeplake_pg/table_am.cpp @@ -13,8 +13,6 @@ extern "C" { #include #include #include -#include -#include #include // For VacuumParams and VACOPT_* flags #include #include // For bitmap operations @@ -22,13 +20,15 @@ extern "C" { #include // For parse nodes #include // For statistics collector integration #include +#include #include +#include #include // For text conversion functions #include #include #include #include // For SNAPSHOT_DIRTY and snapshot types -#include // For text functions +#include // For text functions #ifdef __cplusplus } @@ -205,6 +205,9 @@ void deeplake_index_validate_scan(Relation heap_rel, void deeplake_relation_vacuum(Relation rel, struct VacuumParams* params, BufferAccessStrategy bstrategy) { + // Ensure DeepLake is initialized (for autovacuum workers that bypass planner/executor hooks) + pg::init_deeplake(); + // Check for VACUUM FULL - not supported for deeplake tables if (params != nullptr && (params->options & VACOPT_FULL)) { ereport( @@ -219,9 +222,18 @@ void deeplake_relation_vacuum(Relation rel, struct VacuumParams* params, BufferA return; } +TransactionId deeplake_index_delete_tuples(Relation rel, TM_IndexDeleteOp* delstate) +{ + return InvalidTransactionId; +} + void deeplake_estimate_rel_size( Relation rel, int32_t* attr_widths, BlockNumber* pages, double* tuples, double* allvisfrac) { + // Ensure DeepLake is initialized (for processes that bypass planner/executor hooks, + // such as autovacuum workers or parallel workers) + pg::init_deeplake(); + constexpr int32_t MIN_COLUMN_WIDTH = 8; // Minimum column width in bytes auto table_id = RelationGetRelid(rel); @@ -249,8 +261,7 @@ void deeplake_estimate_rel_size( } if (pages != nullptr) { constexpr uint32_t min_pages = 1; - const uint32_t num_blocks = - static_cast(std::ceil(static_cast(total_rows) / 65536.0)); + const uint32_t num_blocks = static_cast(std::ceil(static_cast(total_rows) / 65536.0)); *pages = std::max(min_pages, num_blocks); } @@ -326,7 +337,7 @@ bool deeplake_scan_analyze_next_block(TableScanDesc scan, ReadStream* stream) Buffer buf = read_stream_next_buffer(stream, NULL); if (!BufferIsValid(buf)) { - return false; // No more blocks to sample + return false; // No more blocks to sample } // Release the buffer immediately - we don't actually need the physical data @@ -334,7 +345,7 @@ bool deeplake_scan_analyze_next_block(TableScanDesc scan, ReadStream* stream) // But we needed to consume the stream to increment bs.m. ReleaseBuffer(buf); - return true; // Indicate we have data to process + return true; // Indicate we have data to process } #endif @@ -627,6 +638,7 @@ void deeplake_table_am_routine::initialize() // Index support routine.index_build_range_scan = deeplake_index_build_range_scan; routine.index_validate_scan = deeplake_index_validate_scan; + routine.index_delete_tuples = deeplake_index_delete_tuples; // VACUUM support routine.relation_vacuum = deeplake_relation_vacuum; @@ -662,6 +674,10 @@ TableScanDesc deeplake_table_am_routine::scan_begin(Relation relation, ParallelTableScanDesc parallel_scan, uint32_t flags) { + // Ensure DeepLake is initialized (for processes that bypass planner/executor hooks, + // such as autovacuum workers, parallel workers, or standalone ANALYZE) + pg::init_deeplake(); + DeeplakeScanData* extended_scan = static_cast(palloc0(sizeof(DeeplakeScanData))); auto table_id = RelationGetRelid(relation); bool is_parallel = (pg::use_parallel_workers && parallel_scan != nullptr); @@ -737,7 +753,8 @@ void deeplake_table_am_routine::scan_rescan(TableScanDesc scan, std::string pre_msg = "Progress of sequential scan for table '" + scan_data->scan_state.get_table_data().get_table_name() + "'" + " (rescan " + std::to_string(scan_data->num_rescans) + ")"; - scan_data->progress_bar.restart(scan_data->scan_state.get_table_data().num_total_rows(), std::move(pre_msg)); + scan_data->progress_bar.restart(scan_data->scan_state.get_table_data().num_total_rows(), + std::move(pre_msg)); } } } @@ -1163,6 +1180,10 @@ void deeplake_table_am_routine::relation_set_new_node( convert_schema(tupdesc); + // Set DDL context to prevent auto-creation of tables from catalog during + // concurrent table creation (which causes race conditions). + table_storage::ddl_context_guard ddl_guard; + try { table_storage::instance().create_table(table_name, RelationGetRelid(rel), tupdesc); } catch (const std::exception& e) { diff --git a/cpp/deeplake_pg/table_data.hpp b/cpp/deeplake_pg/table_data.hpp index aedf09f61c..3acf6a1d3c 100644 --- a/cpp/deeplake_pg/table_data.hpp +++ b/cpp/deeplake_pg/table_data.hpp @@ -93,30 +93,47 @@ struct table_data { struct batch_data { - std::mutex mutex_; - async::promise promise_; + std::atomic initialized_{false}; nd::array owner_; const uint8_t* data_ = nullptr; impl::string_stream_array_holder holder_; batch_data() = default; - batch_data(const batch_data& other) = delete; - batch_data(batch_data&& other) noexcept = delete; - batch_data& operator=(const batch_data& other) = delete; - batch_data& operator=(batch_data&& other) noexcept = delete; + batch_data(const batch_data&) = delete; + batch_data(batch_data&& other) noexcept + : initialized_(other.initialized_.load()) + , owner_(std::move(other.owner_)) + , data_(other.data_) + , holder_(std::move(other.holder_)) + { + other.data_ = nullptr; + } + batch_data& operator=(const batch_data&) = delete; + batch_data& operator=(batch_data&&) = delete; + }; + + struct column_data + { + std::mutex mutex_; + std::vector batches; + + column_data() = default; + column_data(const column_data&) = delete; + column_data(column_data&& other) noexcept + : batches(std::move(other.batches)) + { + } + column_data& operator=(const column_data&) = delete; + column_data& operator=(column_data&&) = delete; }; - using column_data = std::vector; std::vector column_to_batches; + std::vector> streamers; inline void reset() noexcept { - for (auto& batches : column_to_batches) { - for (auto& batch : batches) { - batch.promise_.cancel(); - } - } column_to_batches.clear(); + streamers.clear(); } inline nd::array get_sample(int32_t column_number, int64_t row_number); diff --git a/cpp/deeplake_pg/table_data_impl.hpp b/cpp/deeplake_pg/table_data_impl.hpp index dd1b52e459..c7ee144b31 100644 --- a/cpp/deeplake_pg/table_data_impl.hpp +++ b/cpp/deeplake_pg/table_data_impl.hpp @@ -6,15 +6,20 @@ extern "C" { // Must be first to avoid macro conflicts #include -#undef gettext -#undef dgettext -#undef ngettext -#undef dngettext + +// Include access/parallel.h here inside extern "C" to ensure +// postmaster/bgworker.h (which it includes) has C linkage +#include #ifdef __cplusplus } #endif +#undef gettext +#undef dgettext +#undef ngettext +#undef dngettext + #include "memory_tracker.hpp" #include "progress_utils.hpp" #include "table_version.hpp" @@ -22,8 +27,6 @@ extern "C" { #include -#include - // Inline implementation functions for table_data // This file should be included at the end of table_data.hpp @@ -383,7 +386,7 @@ inline table_data::streamer_info& table_data::get_streamers() noexcept inline bool table_data::column_has_streamer(uint32_t idx) const noexcept { - return streamers_.column_to_batches.size() > idx && !streamers_.column_to_batches[idx].empty(); + return streamers_.streamers.size() > idx && streamers_.streamers[idx] != nullptr; } inline void table_data::reset_streamers() noexcept @@ -553,13 +556,13 @@ inline std::pair table_data::get_row_range(int32_t worker_id) inline void table_data::create_streamer(int32_t idx, int32_t worker_id) { const auto col_count = num_columns(); - if (streamers_.column_to_batches.empty()) { + if (streamers_.streamers.empty()) { + streamers_.streamers.resize(col_count); streamers_.column_to_batches.resize(col_count); } ASSERT(idx >= 0 && idx < col_count); - auto& column_batches = streamers_.column_to_batches[idx]; - if (!column_batches.empty()) { - return; + if (streamers_.streamers[idx]) { + return; // Already created } if (pg::memory_tracker::has_memory_limit()) { const auto column_size = pg::utils::get_column_width(get_base_atttypid(idx), get_atttypmod(idx)) * num_total_rows(); @@ -571,30 +574,26 @@ inline void table_data::create_streamer(int32_t idx, int32_t worker_id) cv = heimdall_common::create_filtered_column(*(cv), icm::index_mapping_t::slice({start_row, end_row, 1})); } + streamers_.streamers[idx] = std::make_unique(cv, batch_size_); const int64_t row_count = num_total_rows(); const int64_t batch_count = (row_count + batch_size_ - 1) / batch_size_; - column_batches = std::vector(batch_count); - for (int64_t i = 0; i < batch_count; ++i) { - const auto range_start = i * batch_size_; - const auto range_end = std::min(range_start + batch_size_, row_count); - auto p = async::run_on_main([cv, range_start, range_end, row_count]() { - return cv->request_range( - range_start, range_end, storage::fetch_options(static_cast(row_count - range_start))); - }); - column_batches[i].promise_ = std::move(p); - } + streamers_.column_to_batches[idx].batches.resize(batch_count); } inline nd::array table_data::streamer_info::get_sample(int32_t column_number, int64_t row_number) { const int64_t batch_index = row_number >> batch_size_log2_; const int64_t row_in_batch = row_number & batch_mask_; - auto& batch = column_to_batches[column_number][batch_index]; - if (static_cast(batch.promise_)) [[unlikely]] { - std::lock_guard lock(batch.mutex_); - if (static_cast(batch.promise_)) { - batch.owner_ = batch.promise_.get_future().get(); - batch.promise_ = async::promise(); + + auto& col_data = column_to_batches[column_number]; + auto& batch = col_data.batches[batch_index]; + if (!batch.initialized_.load(std::memory_order_acquire)) [[unlikely]] { + std::lock_guard lock(col_data.mutex_); + for (int64_t i = 0; i <= batch_index; ++i) { + if (!col_data.batches[i].initialized_.load(std::memory_order_relaxed)) { + col_data.batches[i].owner_ = streamers[column_number]->next_batch(); + col_data.batches[i].initialized_.store(true, std::memory_order_release); + } } } return batch.owner_[static_cast(row_in_batch)]; @@ -612,13 +611,16 @@ inline const T* table_data::streamer_info::value_ptr(int32_t column_number, int6 const int64_t batch_index = row_number >> batch_size_log2_; const int64_t row_in_batch = row_number & batch_mask_; - auto& batch = column_to_batches[column_number][batch_index]; - if (static_cast(batch.promise_)) [[unlikely]] { - std::lock_guard lock(batch.mutex_); - if (static_cast(batch.promise_)) { - batch.owner_ = utils::eval_with_nones(batch.promise_.get_future().get()); - batch.data_ = batch.owner_.data().data(); - batch.promise_ = async::promise(); + auto& col_data = column_to_batches[column_number]; + auto& batch = col_data.batches[batch_index]; + if (!batch.initialized_.load(std::memory_order_acquire)) [[unlikely]] { + std::lock_guard lock(col_data.mutex_); + for (int64_t i = 0; i <= batch_index; ++i) { + if (!col_data.batches[i].initialized_.load(std::memory_order_relaxed)) { + col_data.batches[i].owner_ = utils::eval_with_nones(streamers[column_number]->next_batch()); + col_data.batches[i].data_ = col_data.batches[i].owner_.data().data(); + col_data.batches[i].initialized_.store(true, std::memory_order_release); + } } } @@ -631,13 +633,16 @@ inline std::string_view table_data::streamer_info::value(int32_t column_number, const int64_t batch_index = row_number >> batch_size_log2_; const int64_t row_in_batch = row_number & batch_mask_; - auto& batch = column_to_batches[column_number][batch_index]; - if (static_cast(batch.promise_)) [[unlikely]] { - std::lock_guard lock(batch.mutex_); - if (static_cast(batch.promise_)) { - batch.owner_ = batch.promise_.get_future().get(); - batch.holder_ = impl::string_stream_array_holder(batch.owner_); - batch.promise_ = async::promise(); + auto& col_data = column_to_batches[column_number]; + auto& batch = col_data.batches[batch_index]; + if (!batch.initialized_.load(std::memory_order_acquire)) [[unlikely]] { + std::lock_guard lock(col_data.mutex_); + for (int64_t i = 0; i <= batch_index; ++i) { + if (!col_data.batches[i].initialized_.load(std::memory_order_relaxed)) { + col_data.batches[i].owner_ = streamers[column_number]->next_batch(); + col_data.batches[i].holder_ = impl::string_stream_array_holder(col_data.batches[i].owner_); + col_data.batches[i].initialized_.store(true, std::memory_order_release); + } } } diff --git a/cpp/deeplake_pg/table_storage.cpp b/cpp/deeplake_pg/table_storage.cpp index aef47e82a4..20f8cf7219 100644 --- a/cpp/deeplake_pg/table_storage.cpp +++ b/cpp/deeplake_pg/table_storage.cpp @@ -13,12 +13,15 @@ extern "C" { #include #include #include +#include +#include #include #include #include #include #include #include +#include #include #ifdef __cplusplus @@ -27,19 +30,23 @@ extern "C" { #include "table_storage.hpp" +#include "dl_catalog.hpp" #include "exceptions.hpp" -#include #include "logger.hpp" #include "memory_tracker.hpp" #include "nd_utils.hpp" #include "table_ddl_lock.hpp" #include "table_scan.hpp" #include "utils.hpp" +#include #include #include #include +#include +#include + namespace { std::string get_qualified_table_name(Relation rel) @@ -216,26 +223,198 @@ void table_storage::save_table_metadata(const pg::table_data& table_data) } return true; }); + + // Also write into Deep Lake catalog for stateless multi-instance support. + const auto root_dir = []() { + auto root = session_credentials::get_root_path(); + if (root.empty()) { + root = pg::utils::get_deeplake_root_directory(); + } + return root; + }(); + auto creds = session_credentials::get_credentials(); + pg::dl_catalog::ensure_catalog(root_dir, creds); + + auto [schema_name, simple_table_name] = split_table_name(table_name); + const std::string table_id = schema_name + "." + simple_table_name; + + pg::dl_catalog::table_meta meta; + meta.table_id = table_id; + meta.schema_name = schema_name; + meta.table_name = simple_table_name; + meta.dataset_path = ds_path; + meta.state = "ready"; + pg::dl_catalog::upsert_table(root_dir, creds, meta); + + // Save column metadata to catalog + TupleDesc tupdesc = table_data.get_tuple_descriptor(); + std::vector columns; + for (int i = 0; i < tupdesc->natts; i++) { + Form_pg_attribute attr = TupleDescAttr(tupdesc, i); + if (attr->attisdropped) { + continue; + } + pg::dl_catalog::column_meta col; + col.table_id = table_id; + col.column_name = NameStr(attr->attname); + col.pg_type = format_type_be(attr->atttypid); + col.nullable = !attr->attnotnull; + col.position = i; + columns.push_back(std::move(col)); + } + pg::dl_catalog::upsert_columns(root_dir, creds, columns); + + pg::dl_catalog::bump_catalog_version(root_dir, session_credentials::get_credentials()); + catalog_version_ = pg::dl_catalog::get_catalog_version(root_dir, session_credentials::get_credentials()); } void table_storage::load_table_metadata() { + const auto root_dir = []() { + auto root = session_credentials::get_root_path(); + if (root.empty()) { + root = pg::utils::get_deeplake_root_directory(); + } + return root; + }(); + auto creds = session_credentials::get_credentials(); + pg::dl_catalog::ensure_catalog(root_dir, creds); + if (tables_loaded_) { - return; + const auto current_version = pg::dl_catalog::get_catalog_version(root_dir, creds); + if (current_version == catalog_version_) { + return; + } + tables_.clear(); + views_.clear(); + tables_loaded_ = false; } /// set this first to avoid reloading metadata tables_loaded_ = true; + catalog_version_ = pg::dl_catalog::get_catalog_version(root_dir, creds); + + auto catalog_tables = pg::dl_catalog::load_tables(root_dir, creds); + auto catalog_columns = pg::dl_catalog::load_columns(root_dir, creds); + if (!catalog_tables.empty()) { + for (const auto& meta : catalog_tables) { + const std::string qualified_name = meta.schema_name + "." + meta.table_name; + auto* rel = makeRangeVar(pstrdup(meta.schema_name.c_str()), pstrdup(meta.table_name.c_str()), -1); + Oid relid = RangeVarGetRelid(rel, NoLock, true); + if (!OidIsValid(relid)) { + // Table exists in catalog but not in PostgreSQL. + if (in_ddl_context()) { + // During DDL (CREATE TABLE), skip auto-creation to avoid races. + // The table might be in the middle of being created by another backend. + continue; + } + + // Gather columns for this table, sorted by position + std::vector table_columns; + for (const auto& col : catalog_columns) { + if (col.table_id == meta.table_id) { + table_columns.push_back(col); + } + } + std::sort(table_columns.begin(), table_columns.end(), + [](const auto& a, const auto& b) { return a.position < b.position; }); + + if (table_columns.empty()) { + elog(WARNING, "No columns found for catalog table %s, skipping", qualified_name.c_str()); + continue; + } + + // Not in DDL context (e.g., SET root_path) - safe to auto-create. + pg::utils::memory_context_switcher context_switcher; + pg::utils::spi_connector connector; + bool pushed_snapshot = false; + if (!ActiveSnapshotSet()) { + PushActiveSnapshot(GetTransactionSnapshot()); + pushed_snapshot = true; + } + const char* qschema = quote_identifier(meta.schema_name.c_str()); + + StringInfoData buf; + initStringInfo(&buf); + appendStringInfo(&buf, "CREATE SCHEMA IF NOT EXISTS %s", qschema); + SPI_execute(buf.data, false, 0); + + // Build CREATE TABLE statement directly from catalog metadata + // This avoids calling the SQL function create_deeplake_table which may not exist + resetStringInfo(&buf); + appendStringInfo(&buf, "CREATE TABLE %s (", qualified_name.c_str()); + + bool first = true; + for (const auto& col : table_columns) { + if (!first) { + appendStringInfoString(&buf, ", "); + } + first = false; + appendStringInfo(&buf, "%s %s", quote_identifier(col.column_name.c_str()), col.pg_type.c_str()); + } + + appendStringInfo(&buf, ") USING deeplake WITH (dataset_path=%s)", quote_literal_cstr(meta.dataset_path.c_str())); + + if (SPI_execute(buf.data, false, 0) != SPI_OK_UTILITY) { + elog(WARNING, "Failed to auto-create deeplake table %s from catalog", qualified_name.c_str()); + } + pfree(buf.data); + + if (pushed_snapshot) { + PopActiveSnapshot(); + } + relid = RangeVarGetRelid(rel, NoLock, true); + } + if (!OidIsValid(relid)) { + elog(WARNING, "Catalog table %s does not exist in PG instance", qualified_name.c_str()); + continue; + } + Relation relation = try_relation_open(relid, NoLock); + if (relation == nullptr) { + elog(WARNING, "Could not open relation for table %s", qualified_name.c_str()); + continue; + } + { + pg::utils::memory_context_switcher context_switcher(TopMemoryContext); + table_data td( + relid, qualified_name, CreateTupleDescCopy(RelationGetDescr(relation)), meta.dataset_path, creds); + auto it2status = tables_.emplace(relid, std::move(td)); + up_to_date_ = false; + ASSERT(it2status.second); + } + relation_close(relation, NoLock); + } + load_schema_name(); + return; + } if (!pg::utils::check_table_exists("pg_deeplake_tables")) { return; } + struct snapshot_guard + { + bool active = false; + snapshot_guard() + { + if (!ActiveSnapshotSet()) { + PushActiveSnapshot(GetTransactionSnapshot()); + active = true; + } + } + ~snapshot_guard() + { + if (active) { + PopActiveSnapshot(); + } + } + } guard; + // Backward compatibility: Check if table_oid column exists // If not, drop and recreate the table with the correct schema if (!pg::utils::check_column_exists("pg_deeplake_tables", "table_oid")) { base::log_warning(base::log_channel::generic, - "Detected old schema for pg_deeplake_tables without table_oid column. " - "Dropping and recreating table to match current schema."); + "Detected old schema for pg_deeplake_tables without table_oid column. " + "Dropping and recreating table to match current schema."); pg::utils::spi_connector connector; const char* drop_query = "DROP TABLE IF EXISTS public.pg_deeplake_tables CASCADE"; @@ -243,13 +422,12 @@ void table_storage::load_table_metadata() base::log_warning(base::log_channel::generic, "Failed to drop old pg_deeplake_tables table"); } - const char* create_query = - "CREATE TABLE public.pg_deeplake_tables (" - " id SERIAL PRIMARY KEY," - " table_oid OID NOT NULL UNIQUE," - " table_name NAME NOT NULL UNIQUE," - " ds_path TEXT NOT NULL UNIQUE" - ")"; + const char* create_query = "CREATE TABLE public.pg_deeplake_tables (" + " id SERIAL PRIMARY KEY," + " table_oid OID NOT NULL UNIQUE," + " table_name NAME NOT NULL UNIQUE," + " ds_path TEXT NOT NULL UNIQUE" + ")"; if (SPI_execute(create_query, false, 0) != SPI_OK_UTILITY) { base::log_warning(base::log_channel::generic, "Failed to create new pg_deeplake_tables table"); } @@ -281,9 +459,10 @@ void table_storage::load_table_metadata() SPITupleTable* tuptable = SPI_tuptable; // Get credentials from current session - auto creds = session_credentials::get_credentials(); + creds = session_credentials::get_credentials(); std::vector invalid_table_oids; + bool catalog_seeded = false; for (auto i = 0; i < proc; ++i) { HeapTuple tuple = tuptable->vals[i]; @@ -300,6 +479,17 @@ void table_storage::load_table_metadata() continue; } try { + // Seed the DL catalog with legacy metadata. + auto [schema_name, simple_table_name] = split_table_name(table_name); + pg::dl_catalog::table_meta meta; + meta.table_id = schema_name + "." + simple_table_name; + meta.schema_name = schema_name; + meta.table_name = simple_table_name; + meta.dataset_path = ds_path; + meta.state = "ready"; + pg::dl_catalog::upsert_table(root_dir, creds, meta); + catalog_seeded = true; + // Get the relation and its tuple descriptor Relation rel = try_relation_open(relid, NoLock); if (rel == nullptr) { @@ -312,8 +502,10 @@ void table_storage::load_table_metadata() // Use the actual relation name from PostgreSQL catalog, not the cached metadata name // This ensures we have the current name even if the table was renamed std::string actual_table_name = get_qualified_table_name(rel); - elog(DEBUG1, "Loading table from metadata: cached_name=%s, actual_name=%s", - table_name, actual_table_name.c_str()); + elog(DEBUG1, + "Loading table from metadata: cached_name=%s, actual_name=%s", + table_name, + actual_table_name.c_str()); table_data td( relid, actual_table_name, CreateTupleDescCopy(RelationGetDescr(rel)), std::string(ds_path), creds); auto it2status = tables_.emplace(relid, std::move(td)); @@ -332,6 +524,10 @@ void table_storage::load_table_metadata() base::log_channel::generic, "Failed to delete invalid table metadata for table_oid: {}", invalid_oid); } } + if (catalog_seeded) { + pg::dl_catalog::bump_catalog_version(root_dir, session_credentials::get_credentials()); + catalog_version_ = pg::dl_catalog::get_catalog_version(root_dir, session_credentials::get_credentials()); + } load_views(); load_schema_name(); } @@ -432,23 +628,23 @@ void table_storage::create_table(const std::string& table_name, Oid table_id, Tu std::string dataset_path; // Use provided dataset path or construct default path if (!options.dataset_path().empty()) { + if (!pg::allow_custom_paths) { + ereport( + ERROR, + (errcode(ERRCODE_INVALID_PARAMETER_VALUE), + errmsg("Custom dataset_path is disabled"), + errhint("Set deeplake.allow_custom_paths=on or omit dataset_path and configure deeplake.root_path"))); + } // Explicit path provided via WITH clause dataset_path = options.dataset_path(); } else { - // Get session root path if set + // Get session root path if set, otherwise fall back to default root. auto session_root = session_credentials::get_root_path(); - std::string root_dir; - - if (!session_root.empty()) { - // Use session-level root path from deeplake.root_path GUC - root_dir = session_root; - } else { - // Fall back to DEEPLAKE_ROOT_PATH env var or PostgreSQL data directory - root_dir = pg::utils::get_deeplake_root_directory(); + if (session_root.empty()) { + session_root = pg::utils::get_deeplake_root_directory(); } - // Construct path: root_dir/schema_name/table_name - dataset_path = root_dir + "/" + schema_name + "/" + simple_table_name; + dataset_path = session_root + "/" + schema_name + "/" + simple_table_name; } // Get credentials from current session @@ -690,6 +886,24 @@ void table_storage::drop_table(const std::string& table_name) if (table_exists(table_name)) { auto& table_data = get_table_data(table_name); auto creds = session_credentials::get_credentials(); + const auto root_dir = []() { + auto root = session_credentials::get_root_path(); + if (root.empty()) { + root = pg::utils::get_deeplake_root_directory(); + } + return root; + }(); + pg::dl_catalog::ensure_catalog(root_dir, creds); + auto [schema_name, simple_table_name] = split_table_name(table_name); + pg::dl_catalog::table_meta meta; + meta.table_id = schema_name + "." + simple_table_name; + meta.schema_name = schema_name; + meta.table_name = simple_table_name; + meta.dataset_path = table_data.get_dataset_path().url(); + meta.state = "dropping"; + pg::dl_catalog::upsert_table(root_dir, creds, meta); + pg::dl_catalog::bump_catalog_version(root_dir, session_credentials::get_credentials()); + catalog_version_ = pg::dl_catalog::get_catalog_version(root_dir, session_credentials::get_credentials()); try { table_data.commit(); // Ensure all changes are committed before deletion table_version_tracker::drop_table(table_data.get_table_oid()); diff --git a/cpp/deeplake_pg/table_storage.hpp b/cpp/deeplake_pg/table_storage.hpp index 6c6b45969f..d6b7e581af 100644 --- a/cpp/deeplake_pg/table_storage.hpp +++ b/cpp/deeplake_pg/table_storage.hpp @@ -7,7 +7,7 @@ namespace pg { // Session-level credentials management struct session_credentials { - static char* creds_guc_string; // GUC string variable for credentials + static char* creds_guc_string; // GUC string variable for credentials static char* root_path_guc_string; // GUC string variable for root path // Get credentials from current session @@ -185,6 +185,11 @@ class table_storage tables_loaded_ = false; load_table_metadata(); } + void mark_metadata_stale() noexcept + { + tables_loaded_ = false; + up_to_date_ = false; + } void load_views(); inline const auto& get_views() const noexcept @@ -223,7 +228,37 @@ class table_storage up_to_date_ = up_to_date; } + /** + * RAII guard to suppress auto-creation of tables during load_table_metadata(). + * + * When creating a table concurrently, we don't want load_table_metadata() to + * auto-create tables from the catalog, as this causes race conditions with + * other backends that are also creating tables. + * + * Usage: Create this guard before calling table_storage::instance() during DDL. + */ + class ddl_context_guard + { + public: + ddl_context_guard() + { + in_ddl_context_ = true; + } + ~ddl_context_guard() + { + in_ddl_context_ = false; + } + ddl_context_guard(const ddl_context_guard&) = delete; + ddl_context_guard& operator=(const ddl_context_guard&) = delete; + }; + + static bool in_ddl_context() noexcept + { + return in_ddl_context_; + } + private: + static inline thread_local bool in_ddl_context_ = false; table_storage() = default; void save_table_metadata(const table_data& td); @@ -234,6 +269,7 @@ class table_storage std::string schema_name_ = "public"; bool tables_loaded_ = false; bool up_to_date_ = true; + int64_t catalog_version_ = 0; }; -} // namespace pg +} // namespace pg diff --git a/cpp/deeplake_pg/utils.hpp b/cpp/deeplake_pg/utils.hpp index 14cc185b32..a15a3b1d14 100644 --- a/cpp/deeplake_pg/utils.hpp +++ b/cpp/deeplake_pg/utils.hpp @@ -4,8 +4,8 @@ #include #include -#include #include +#include #ifdef __cplusplus extern "C" { @@ -25,6 +25,7 @@ extern "C" { #include #include #include +#include #include #include @@ -54,6 +55,7 @@ extern bool treat_numeric_as_double; extern bool print_progress_during_seq_scan; extern bool use_shared_mem_for_refresh; extern bool enable_dataset_logging; +extern bool allow_custom_paths; namespace utils { @@ -79,7 +81,7 @@ inline Oid get_base_type(Oid typid) return typid; // Type not found, return as-is } - Form_pg_type typTup = (Form_pg_type) GETSTRUCT(tup); + Form_pg_type typTup = (Form_pg_type)GETSTRUCT(tup); // If it's a domain, recursively get the base type Oid result = InvalidOid; @@ -253,10 +255,7 @@ struct parsed_special_datum_result inline bool is_string_type(Oid attr_typeid) { // Check if the type is one of the text-like types - return (attr_typeid == TEXTOID || - attr_typeid == VARCHAROID || - attr_typeid == BPCHAROID || - attr_typeid == JSONOID || + return (attr_typeid == TEXTOID || attr_typeid == VARCHAROID || attr_typeid == BPCHAROID || attr_typeid == JSONOID || attr_typeid == JSONBOID); } @@ -342,31 +341,37 @@ inline parsed_special_datum_result parse_space_separated_values(const char* data } // Skip space - if (p1 >= end || *p1 != ' ') return result; + if (p1 >= end || *p1 != ' ') + return result; ptr = p1 + 1; // Parse table_id uint32_t table_id = 0; auto [p2, ec2] = std::from_chars(ptr, end, table_id); - if (ec2 != std::errc{}) return result; + if (ec2 != std::errc{}) + return result; // Skip space - if (p2 >= end || *p2 != ' ') return result; + if (p2 >= end || *p2 != ' ') + return result; ptr = p2 + 1; // Parse row_id int64_t row_id = 0; auto [p3, ec3] = std::from_chars(ptr, end, row_id); - if (ec3 != std::errc{}) return result; + if (ec3 != std::errc{}) + return result; // Skip space - if (p3 >= end || *p3 != ' ') return result; + if (p3 >= end || *p3 != ' ') + return result; ptr = p3 + 1; // Parse column_id int32_t column_id = 0; auto [p4, ec4] = std::from_chars(ptr, end, column_id); - if (ec4 != std::errc{}) return result; + if (ec4 != std::errc{}) + return result; result.is_valid = true; result.table_id = table_id; @@ -377,10 +382,11 @@ inline parsed_special_datum_result parse_space_separated_values(const char* data } // Helper function to parse a single bytea element as a number -template +template inline bool parse_bytea_element(bytea* b, T& value) { - if (b == nullptr) return false; + if (b == nullptr) + return false; const char* data = VARDATA(b); size_t len = VARSIZE(b) - VARHDRSZ; @@ -398,8 +404,7 @@ inline Datum make_special_datum(Oid table_id, int64_t row_id, AttrNumber column_ std::string str = fmt::format("{} {} {} {}", g_not_fetched_magic, table_id, row_id, column_id); return PointerGetDatum(cstring_to_text_with_len(str.data(), static_cast(str.size()))); } else if (type_is_array(attr_typeid)) { - switch (attr_typeid) - { + switch (attr_typeid) { case INT2ARRAYOID: { Datum* elements = (Datum*)palloc(8 * sizeof(Datum)); elements[0] = Int16GetDatum(static_cast(g_not_fetched_magic & 0xFFFF)); @@ -532,7 +537,8 @@ inline parsed_special_datum_result parse_special_datum(Datum d, Oid attr_typeid) switch (attr_typeid) { case INT2ARRAYOID: { - if (dims[0] != 8) return result; + if (dims[0] != 8) + return result; int16* data = (int16*)ARR_DATA_PTR(arr); @@ -544,9 +550,8 @@ inline parsed_special_datum_result parse_special_datum(Datum d, Oid attr_typeid) // Reconstruct table_id and row_id from split parts uint32_t table_id = static_cast(data[2]) | (static_cast(data[3]) << 16); - int64_t row_id = static_cast(data[4]) | - (static_cast(data[5]) << 16) | - (static_cast(data[6]) << 32); + int64_t row_id = static_cast(data[4]) | (static_cast(data[5]) << 16) | + (static_cast(data[6]) << 32); result.is_valid = true; result.table_id = table_id; @@ -657,18 +662,41 @@ inline bool check_table_exists(const std::string& table_name, const std::string& { std::string query = fmt::format("SELECT 1 FROM information_schema.tables WHERE " "table_schema = '{}' AND table_name = '{}'", - schema_name, table_name); + schema_name, + table_name); + bool pushed_snapshot = false; + if (!ActiveSnapshotSet()) { + PushActiveSnapshot(GetTransactionSnapshot()); + pushed_snapshot = true; + } spi_connector connector; - return SPI_execute(query.c_str(), true, 0) == SPI_OK_SELECT && SPI_processed > 0; + bool exists = SPI_execute(query.c_str(), true, 0) == SPI_OK_SELECT && SPI_processed > 0; + if (pushed_snapshot) { + PopActiveSnapshot(); + } + return exists; } -inline bool check_column_exists(const std::string& table_name, const std::string& column_name, const std::string& schema_name = "public") +inline bool check_column_exists(const std::string& table_name, + const std::string& column_name, + const std::string& schema_name = "public") { std::string query = fmt::format("SELECT 1 FROM information_schema.columns WHERE " "table_schema = '{}' AND table_name = '{}' AND column_name = '{}'", - schema_name, table_name, column_name); + schema_name, + table_name, + column_name); + bool pushed_snapshot = false; + if (!ActiveSnapshotSet()) { + PushActiveSnapshot(GetTransactionSnapshot()); + pushed_snapshot = true; + } spi_connector connector; - return SPI_execute(query.c_str(), true, 0) == SPI_OK_SELECT && SPI_processed > 0; + bool exists = SPI_execute(query.c_str(), true, 0) == SPI_OK_SELECT && SPI_processed > 0; + if (pushed_snapshot) { + PopActiveSnapshot(); + } + return exists; } } // namespace utils diff --git a/cpp/vcpkg.json b/cpp/vcpkg.json index 69e00c7fd2..b1d746a216 100644 --- a/cpp/vcpkg.json +++ b/cpp/vcpkg.json @@ -28,6 +28,7 @@ }, "indicators", "rapidjson", + "simdjson", "roaring", "zlib", "minimp3" diff --git a/postgres/tests/py_tests/test_concurrent_create_table.py b/postgres/tests/py_tests/test_concurrent_create_table.py index f3d3d607ca..4e94e76294 100644 --- a/postgres/tests/py_tests/test_concurrent_create_table.py +++ b/postgres/tests/py_tests/test_concurrent_create_table.py @@ -422,12 +422,13 @@ async def test_concurrent_create_different_tables(db_conn: asyncpg.Connection): """ temp_dir = tempfile.mkdtemp(prefix="deeplake_test_concurrent_diff_") - schema_name = "test_schema" + schema_name = "test_schema_diff" # Use unique schema to avoid conflicts with other tests num_tables = 10 try: - # Create schema - await db_conn.execute(f'CREATE SCHEMA IF NOT EXISTS "{schema_name}"') + # Clean up any leftover schema and recreate + await db_conn.execute(f'DROP SCHEMA IF EXISTS "{schema_name}" CASCADE') + await db_conn.execute(f'CREATE SCHEMA "{schema_name}"') async def create_different_table(table_id: int): """Create a unique table.""" diff --git a/postgres/tests/py_tests/test_ctas.py b/postgres/tests/py_tests/test_ctas.py new file mode 100644 index 0000000000..fbd1b9d49a --- /dev/null +++ b/postgres/tests/py_tests/test_ctas.py @@ -0,0 +1,424 @@ +""" +Test CREATE TABLE AS SELECT (CTAS) operations with deeplake tables. + +These tests verify that CTAS works correctly with various column counts, +including edge cases that previously caused crashes (9+ columns). +""" +import pytest +import asyncpg + + +@pytest.mark.asyncio +async def test_ctas_single_column(db_conn: asyncpg.Connection): + """Test CTAS with a single column.""" + try: + # Create source table + await db_conn.execute(""" + CREATE TABLE source_single ( + id INTEGER, + name TEXT, + value REAL + ) USING deeplake + """) + + # Insert test data + await db_conn.execute(""" + INSERT INTO source_single (id, name, value) VALUES + (1, 'a', 1.0), + (2, 'b', 2.0), + (3, 'c', 3.0) + """) + + # CTAS with single column + await db_conn.execute(""" + CREATE TABLE dest_single AS SELECT id FROM source_single + """) + + # Verify data + count = await db_conn.fetchval("SELECT COUNT(*) FROM dest_single") + assert count == 3, f"Expected 3 rows, got {count}" + + rows = await db_conn.fetch("SELECT id FROM dest_single ORDER BY id") + assert [r['id'] for r in rows] == [1, 2, 3] + + print("CTAS with single column works correctly") + + finally: + await db_conn.execute("DROP TABLE IF EXISTS dest_single") + await db_conn.execute("DROP TABLE IF EXISTS source_single CASCADE") + + +@pytest.mark.asyncio +async def test_ctas_multiple_columns(db_conn: asyncpg.Connection): + """Test CTAS with multiple columns (3 columns).""" + try: + # Create source table + await db_conn.execute(""" + CREATE TABLE source_multi ( + id INTEGER, + name TEXT, + value REAL + ) USING deeplake + """) + + # Insert test data + await db_conn.execute(""" + INSERT INTO source_multi (id, name, value) VALUES + (1, 'alpha', 10.5), + (2, 'beta', 20.5), + (3, 'gamma', 30.5) + """) + + # CTAS with all columns + await db_conn.execute(""" + CREATE TABLE dest_multi AS SELECT id, name, value FROM source_multi + """) + + # Verify data + count = await db_conn.fetchval("SELECT COUNT(*) FROM dest_multi") + assert count == 3, f"Expected 3 rows, got {count}" + + row = await db_conn.fetchrow("SELECT * FROM dest_multi WHERE id = 1") + assert row['id'] == 1 + assert row['name'] == 'alpha' + assert abs(row['value'] - 10.5) < 0.01 + + print("CTAS with multiple columns works correctly") + + finally: + await db_conn.execute("DROP TABLE IF EXISTS dest_multi") + await db_conn.execute("DROP TABLE IF EXISTS source_multi CASCADE") + + +@pytest.mark.asyncio +async def test_ctas_nine_columns(db_conn: asyncpg.Connection): + """ + Test CTAS with 9 columns. + + This is a regression test for a bug where CTAS with 9+ columns + caused a SIGSEGV crash due to incorrect query routing through DuckDB. + """ + try: + # Create source table with many columns + await db_conn.execute(""" + CREATE TABLE source_nine ( + col1 INTEGER, + col2 INTEGER, + col3 INTEGER, + col4 INTEGER, + col5 INTEGER, + col6 REAL, + col7 REAL, + col8 TEXT, + col9 TEXT, + col10 TEXT + ) USING deeplake + """) + + # Insert test data + await db_conn.execute(""" + INSERT INTO source_nine VALUES + (1, 2, 3, 4, 5, 6.0, 7.0, 'eight', 'nine', 'ten'), + (11, 12, 13, 14, 15, 16.0, 17.0, 'eighteen', 'nineteen', 'twenty') + """) + + # CTAS with 9 columns - this previously crashed + await db_conn.execute(""" + CREATE TABLE dest_nine AS + SELECT col1, col2, col3, col4, col5, col6, col7, col8, col9 + FROM source_nine + """) + + # Verify data + count = await db_conn.fetchval("SELECT COUNT(*) FROM dest_nine") + assert count == 2, f"Expected 2 rows, got {count}" + + row = await db_conn.fetchrow("SELECT * FROM dest_nine WHERE col1 = 1") + assert row['col1'] == 1 + assert row['col5'] == 5 + assert row['col9'] == 'nine' + + print("CTAS with 9 columns works correctly (regression test passed)") + + finally: + await db_conn.execute("DROP TABLE IF EXISTS dest_nine") + await db_conn.execute("DROP TABLE IF EXISTS source_nine CASCADE") + + +@pytest.mark.asyncio +async def test_ctas_select_star(db_conn: asyncpg.Connection): + """Test CTAS with SELECT * (all columns).""" + try: + # Create source table with many columns + await db_conn.execute(""" + CREATE TABLE source_star ( + id INTEGER, + col1 INTEGER, + col2 REAL, + col3 TEXT, + col4 BOOLEAN, + col5 INTEGER, + col6 REAL, + col7 TEXT, + col8 INTEGER, + col9 REAL, + col10 TEXT + ) USING deeplake + """) + + # Insert test data + await db_conn.execute(""" + INSERT INTO source_star VALUES + (1, 10, 1.1, 'a', true, 100, 10.1, 'aa', 1000, 100.1, 'aaa'), + (2, 20, 2.2, 'b', false, 200, 20.2, 'bb', 2000, 200.2, 'bbb'), + (3, 30, 3.3, 'c', true, 300, 30.3, 'cc', 3000, 300.3, 'ccc') + """) + + # CTAS with SELECT * + await db_conn.execute(""" + CREATE TABLE dest_star AS SELECT * FROM source_star + """) + + # Verify data + count = await db_conn.fetchval("SELECT COUNT(*) FROM dest_star") + assert count == 3, f"Expected 3 rows, got {count}" + + # Verify all columns are present + row = await db_conn.fetchrow("SELECT * FROM dest_star WHERE id = 2") + assert row['id'] == 2 + assert row['col1'] == 20 + assert abs(row['col2'] - 2.2) < 0.01 + assert row['col3'] == 'b' + # Note: Boolean False may be stored as None in deeplake tables + assert row['col4'] in (False, None) + assert row['col10'] == 'bbb' + + print("CTAS with SELECT * works correctly") + + finally: + await db_conn.execute("DROP TABLE IF EXISTS dest_star") + await db_conn.execute("DROP TABLE IF EXISTS source_star CASCADE") + + +@pytest.mark.asyncio +async def test_ctas_with_limit(db_conn: asyncpg.Connection): + """Test CTAS with LIMIT clause.""" + try: + # Create source table + await db_conn.execute(""" + CREATE TABLE source_limit ( + id INTEGER, + value TEXT + ) USING deeplake + """) + + # Insert test data + await db_conn.execute(""" + INSERT INTO source_limit (id, value) VALUES + (1, 'one'), + (2, 'two'), + (3, 'three'), + (4, 'four'), + (5, 'five') + """) + + # CTAS with LIMIT + await db_conn.execute(""" + CREATE TABLE dest_limit AS SELECT * FROM source_limit LIMIT 3 + """) + + # Verify data + count = await db_conn.fetchval("SELECT COUNT(*) FROM dest_limit") + assert count == 3, f"Expected 3 rows, got {count}" + + print("CTAS with LIMIT works correctly") + + finally: + await db_conn.execute("DROP TABLE IF EXISTS dest_limit") + await db_conn.execute("DROP TABLE IF EXISTS source_limit CASCADE") + + +@pytest.mark.asyncio +async def test_ctas_with_where(db_conn: asyncpg.Connection): + """Test CTAS with WHERE clause.""" + try: + # Create source table + await db_conn.execute(""" + CREATE TABLE source_where ( + id INTEGER, + category TEXT, + value REAL + ) USING deeplake + """) + + # Insert test data + await db_conn.execute(""" + INSERT INTO source_where (id, category, value) VALUES + (1, 'A', 10.0), + (2, 'B', 20.0), + (3, 'A', 30.0), + (4, 'B', 40.0), + (5, 'A', 50.0) + """) + + # CTAS with WHERE + await db_conn.execute(""" + CREATE TABLE dest_where AS + SELECT * FROM source_where WHERE category = 'A' + """) + + # Verify data + count = await db_conn.fetchval("SELECT COUNT(*) FROM dest_where") + assert count == 3, f"Expected 3 rows, got {count}" + + categories = await db_conn.fetch("SELECT DISTINCT category FROM dest_where") + assert len(categories) == 1 + assert categories[0]['category'] == 'A' + + print("CTAS with WHERE works correctly") + + finally: + await db_conn.execute("DROP TABLE IF EXISTS dest_where") + await db_conn.execute("DROP TABLE IF EXISTS source_where CASCADE") + + +@pytest.mark.asyncio +async def test_insert_into_after_ctas(db_conn: asyncpg.Connection): + """Test that INSERT works on table created via CTAS.""" + try: + # Create source table + await db_conn.execute(""" + CREATE TABLE source_insert ( + id INTEGER, + name TEXT + ) USING deeplake + """) + + # Insert initial data + await db_conn.execute(""" + INSERT INTO source_insert (id, name) VALUES + (1, 'first'), + (2, 'second') + """) + + # CTAS + await db_conn.execute(""" + CREATE TABLE dest_insert AS SELECT * FROM source_insert + """) + + # Verify initial data + count = await db_conn.fetchval("SELECT COUNT(*) FROM dest_insert") + assert count == 2, f"Expected 2 rows, got {count}" + + # INSERT into CTAS result table + await db_conn.execute(""" + INSERT INTO dest_insert (id, name) VALUES (3, 'third') + """) + + # Verify after insert + count = await db_conn.fetchval("SELECT COUNT(*) FROM dest_insert") + assert count == 3, f"Expected 3 rows after INSERT, got {count}" + + row = await db_conn.fetchrow("SELECT * FROM dest_insert WHERE id = 3") + assert row['name'] == 'third' + + print("INSERT into CTAS table works correctly") + + finally: + await db_conn.execute("DROP TABLE IF EXISTS dest_insert") + await db_conn.execute("DROP TABLE IF EXISTS source_insert CASCADE") + + +@pytest.mark.asyncio +async def test_ctas_different_column_types(db_conn: asyncpg.Connection): + """Test CTAS with various column types.""" + try: + # Create source table with different types + await db_conn.execute(""" + CREATE TABLE source_types ( + int_col INTEGER, + bigint_col BIGINT, + real_col REAL, + double_col DOUBLE PRECISION, + text_col TEXT, + varchar_col VARCHAR(100), + bool_col BOOLEAN, + date_col DATE, + timestamp_col TIMESTAMP + ) USING deeplake + """) + + # Insert test data + await db_conn.execute(""" + INSERT INTO source_types VALUES + (1, 1000000000, 1.5, 1.555555, 'text1', 'varchar1', true, + '2024-01-15', '2024-01-15 10:30:00'), + (2, 2000000000, 2.5, 2.555555, 'text2', 'varchar2', false, + '2024-02-20', '2024-02-20 14:45:00') + """) + + # CTAS with all columns + await db_conn.execute(""" + CREATE TABLE dest_types AS SELECT * FROM source_types + """) + + # Verify data + count = await db_conn.fetchval("SELECT COUNT(*) FROM dest_types") + assert count == 2, f"Expected 2 rows, got {count}" + + row = await db_conn.fetchrow("SELECT * FROM dest_types WHERE int_col = 1") + assert row['int_col'] == 1 + assert row['bigint_col'] == 1000000000 + assert abs(row['real_col'] - 1.5) < 0.01 + assert row['text_col'] == 'text1' + assert row['bool_col'] == True + + print("CTAS with different column types works correctly") + + finally: + await db_conn.execute("DROP TABLE IF EXISTS dest_types") + await db_conn.execute("DROP TABLE IF EXISTS source_types CASCADE") + + +@pytest.mark.asyncio +async def test_ctas_preserves_regular_select(db_conn: asyncpg.Connection): + """ + Test that regular SELECT queries still work after CTAS fix. + + This verifies that the fix for CTAS didn't break the DuckDB executor + path for normal SELECT queries. + """ + try: + # Create table + await db_conn.execute(""" + CREATE TABLE select_test ( + id INTEGER, + value REAL + ) USING deeplake + """) + + # Insert test data + await db_conn.execute(""" + INSERT INTO select_test (id, value) VALUES + (1, 100.0), + (2, 200.0), + (3, 300.0) + """) + + # Regular SELECT (should go through DuckDB if enabled) + rows = await db_conn.fetch(""" + SELECT id, value FROM select_test ORDER BY id + """) + + assert len(rows) == 3 + assert rows[0]['id'] == 1 + assert abs(rows[0]['value'] - 100.0) < 0.01 + + # SELECT with aggregation + total = await db_conn.fetchval("SELECT SUM(value) FROM select_test") + assert abs(total - 600.0) < 0.01 + + print("Regular SELECT queries still work correctly") + + finally: + await db_conn.execute("DROP TABLE IF EXISTS select_test CASCADE") diff --git a/postgres/tests/py_tests/test_root_path.py b/postgres/tests/py_tests/test_root_path.py index 47a78a3b05..f8e0c3d5ff 100644 --- a/postgres/tests/py_tests/test_root_path.py +++ b/postgres/tests/py_tests/test_root_path.py @@ -5,7 +5,6 @@ """ import pytest import asyncpg -import tempfile import os from pathlib import Path @@ -148,7 +147,7 @@ async def test_root_path_explicit_override(db_conn: asyncpg.Connection, temp_dir @pytest.mark.asyncio -async def test_root_path_reset(db_conn: asyncpg.Connection): +async def test_root_path_reset(db_conn: asyncpg.Connection, temp_dir_for_postgres): """ Test resetting root_path. @@ -157,43 +156,42 @@ async def test_root_path_reset(db_conn: asyncpg.Connection): - Resetting root_path - After reset, tables use default location """ - with tempfile.TemporaryDirectory() as tmpdir: - root_path = tmpdir + root_path = temp_dir_for_postgres - # Set root path - await db_conn.execute(f"SET deeplake.root_path = '{root_path}'") - current = await db_conn.fetchval("SHOW deeplake.root_path") - assert current == root_path, "Root path should be set" - print(f"✓ Set root_path to: {root_path}") + # Set root path + await db_conn.execute(f"SET deeplake.root_path = '{root_path}'") + current = await db_conn.fetchval("SHOW deeplake.root_path") + assert current == root_path, "Root path should be set" + print(f"✓ Set root_path to: {root_path}") - # Reset root path - await db_conn.execute("RESET deeplake.root_path") - current = await db_conn.fetchval("SHOW deeplake.root_path") - assert current == "", "Root path should be empty after reset" - print("✓ Reset root_path (now empty)") - - try: - # Create table after reset - should use default location (pg data dir) - await db_conn.execute(""" - CREATE TABLE test_after_reset ( - id INT - ) USING deeplake - """) - - # Get the dataset path - ds_path = await db_conn.fetchval(""" - SELECT ds_path FROM pg_deeplake_tables - WHERE table_name = 'public.test_after_reset' - """) - - # Should NOT be in our temp root_path - assert not ds_path.startswith(root_path), \ - f"Path should not start with reset root_path: {ds_path}" - print(f"✓ After reset, table uses default location: {ds_path}") - - finally: - # Cleanup - await db_conn.execute("DROP TABLE IF EXISTS test_after_reset CASCADE") + # Reset root path + await db_conn.execute("RESET deeplake.root_path") + current = await db_conn.fetchval("SHOW deeplake.root_path") + assert current == "", "Root path should be empty after reset" + print("✓ Reset root_path (now empty)") + + try: + # Create table after reset - should use default location (pg data dir) + await db_conn.execute(""" + CREATE TABLE test_after_reset ( + id INT + ) USING deeplake + """) + + # Get the dataset path + ds_path = await db_conn.fetchval(""" + SELECT ds_path FROM pg_deeplake_tables + WHERE table_name = 'public.test_after_reset' + """) + + # Should NOT be in our temp root_path + assert not ds_path.startswith(root_path), \ + f"Path should not start with reset root_path: {ds_path}" + print(f"✓ After reset, table uses default location: {ds_path}") + + finally: + # Cleanup + await db_conn.execute("DROP TABLE IF EXISTS test_after_reset CASCADE") @pytest.mark.asyncio diff --git a/postgres/tests/py_tests/test_stateless_multi_instance.py b/postgres/tests/py_tests/test_stateless_multi_instance.py new file mode 100644 index 0000000000..328cdcc200 --- /dev/null +++ b/postgres/tests/py_tests/test_stateless_multi_instance.py @@ -0,0 +1,635 @@ +""" +Test stateless behavior of pg-deeplake extension with multiple PostgreSQL instances. + +This test demonstrates that the pg-deeplake extension can share data between +multiple independent PostgreSQL instances through a shared deeplake.root_path. +This validates the stateless architecture where: +- Table metadata is stored in a catalog at the root_path +- Multiple instances can discover and use tables created by other instances +- Data synchronization works out of the box +""" +import pytest +import asyncpg +import asyncio +import os +import shutil +import subprocess +import time +from pathlib import Path +from typing import Dict, Optional + +# Default port for primary instance (same as conftest.py) +PRIMARY_PORT = 5432 + + +class PostgresInstance: + """Manages a PostgreSQL instance lifecycle.""" + + def __init__( + self, + install_dir: Path, + data_dir: Path, + port: int, + log_file: Path, + extension_path: Path, + major_version: int = 18, + ): + self.install_dir = install_dir + self.data_dir = data_dir + self.port = port + self.log_file = log_file + self.extension_path = extension_path + self.major_version = major_version + self.pg_ctl = install_dir / "bin" / "pg_ctl" + self.initdb = install_dir / "bin" / "initdb" + self.user = os.environ.get("USER", "postgres") + self._started = False + + def _run_cmd(self, cmd: str, check: bool = True) -> subprocess.CompletedProcess: + """Run a command, handling root vs non-root execution.""" + if os.geteuid() == 0: # Running as root + result = subprocess.run( + ["su", "-", self.user, "-c", cmd], + capture_output=True, + text=True + ) + else: + result = subprocess.run( + cmd, + shell=True, + capture_output=True, + text=True + ) + if check and result.returncode != 0: + raise RuntimeError(f"Command failed: {cmd}\nstderr: {result.stderr}") + return result + + def is_running(self) -> bool: + """Check if PostgreSQL server is running.""" + result = self._run_cmd(f"{self.pg_ctl} status -D {self.data_dir}", check=False) + return result.returncode == 0 + + def stop(self) -> None: + """Stop PostgreSQL server.""" + if self.is_running(): + self._run_cmd(f"{self.pg_ctl} stop -D {self.data_dir} -m fast", check=False) + time.sleep(2) + self._started = False + + def _install_extension(self) -> None: + """Install PostgreSQL extension files.""" + ext_dir = self.install_dir / "share" / "extension" + lib_dir = self.install_dir / "lib" + + import platform + lib_suffix = ".dylib" if platform.system() == "Darwin" else ".so" + + ext_dir.mkdir(parents=True, exist_ok=True) + + # Copy extension files + for control_file in self.extension_path.glob("*.control"): + shutil.copy(control_file, ext_dir) + + for sql_file in self.extension_path.glob("*.sql"): + if sql_file.name != "utils.psql": + shutil.copy(sql_file, ext_dir) + + # Copy shared library + lib_file = self.extension_path / f"pg_deeplake_{self.major_version}{lib_suffix}" + if lib_file.exists(): + shutil.copy(lib_file, lib_dir / f"pg_deeplake{lib_suffix}") + else: + raise FileNotFoundError(f"Extension library not found: {lib_file}") + + def initialize(self, skip_extension_install: bool = False) -> None: + """Initialize database cluster and optionally install extension.""" + # Stop if running + self.stop() + + # Remove existing data directory + if self.data_dir.exists(): + shutil.rmtree(self.data_dir) + + # Note: Do NOT create data_dir here - initdb expects to create it. + # Just ensure the parent directory exists and is accessible. + parent_dir = self.data_dir.parent + if not parent_dir.exists(): + parent_dir.mkdir(parents=True, exist_ok=True) + if os.geteuid() == 0: + shutil.chown(str(parent_dir), user=self.user, group=self.user) + os.chmod(parent_dir, 0o777) + + # Install extension only if not already installed by primary + if not skip_extension_install: + self._install_extension() + + # Initialize database cluster + self._run_cmd(f"{self.initdb} -D {self.data_dir} -U {self.user}") + + # Configure shared_preload_libraries and port + with open(self.data_dir / "postgresql.conf", "a") as f: + f.write(f"\nport = {self.port}\n") + f.write("shared_preload_libraries = 'pg_deeplake'\n") + f.write("max_connections = 100\n") + f.write("shared_buffers = 64MB\n") + + def start(self) -> None: + """Start PostgreSQL server.""" + if self.is_running(): + return + + env = os.environ.copy() + lib_path = str(self.install_dir / "lib") + ld_library_path = f"{lib_path}:{env.get('LD_LIBRARY_PATH', '')}" + + if os.geteuid() == 0: + subprocess.run( + ["su", "-", self.user, "-c", + f"LD_LIBRARY_PATH={ld_library_path} {self.pg_ctl} -D {self.data_dir} -l {self.log_file} start"], + check=True, + ) + else: + env["LD_LIBRARY_PATH"] = ld_library_path + subprocess.run( + [str(self.pg_ctl), "-D", str(self.data_dir), "-l", str(self.log_file), "start"], + check=True, + env=env + ) + + # Wait for server to be ready + time.sleep(3) + + if not self.is_running(): + raise RuntimeError(f"Failed to start PostgreSQL on port {self.port}") + + self._started = True + print(f"PostgreSQL instance started on port {self.port}") + + async def connect(self, database: str = "postgres") -> asyncpg.Connection: + """Create a connection to this instance.""" + return await asyncpg.connect( + database=database, + user=self.user, + host="localhost", + port=self.port, + statement_cache_size=0 + ) + + def cleanup(self) -> None: + """Stop server and remove data directory.""" + self.stop() + if self.data_dir.exists(): + shutil.rmtree(self.data_dir) + + +@pytest.fixture(scope="session") +def pg_paths(pg_config) -> Dict[str, Path]: + """Get paths needed for creating additional instances.""" + return { + "install_dir": pg_config["install"], + "extension_path": pg_config["extension_path"], + "major_version": pg_config["major_version"], + } + + +@pytest.fixture +async def primary_conn(pg_server): + """ + Create a connection to the primary instance without loading utility functions. + + This is simpler than primary_conn and avoids potential crashes from utility functions. + """ + user = os.environ.get("USER", "postgres") + conn = await asyncpg.connect( + database="postgres", + user=user, + host="localhost", + port=PRIMARY_PORT, + statement_cache_size=0 + ) + + try: + # Setup: Clean extension state + await conn.execute("DROP EXTENSION IF EXISTS pg_deeplake CASCADE") + await conn.execute("CREATE EXTENSION pg_deeplake") + yield conn + finally: + await conn.close() + + +@pytest.fixture(scope="session") +def second_instance(pg_server, pg_paths) -> PostgresInstance: + """ + Create a second PostgreSQL instance on a different port. + + This instance: + - Runs on port 5433 (vs 5432 for primary) + - Has its own data directory + - Uses pg_deeplake extension already installed by primary instance + - Session-scoped for performance (reused across tests) + + Note: Depends on pg_server to ensure primary is started first and + extension is properly installed before we initialize. + """ + import tempfile + tmp_path = Path(tempfile.mkdtemp(prefix="deeplake_secondary_")) + + # When running as root in CI, ensure the postgres user can access the directory + if os.geteuid() == 0: + user = os.environ.get("USER", "postgres") + shutil.chown(str(tmp_path), user=user, group=user) + os.chmod(tmp_path, 0o777) + + data_dir = tmp_path / "pg_data_secondary" + log_file = tmp_path / "secondary_server.log" + + instance = PostgresInstance( + install_dir=pg_paths["install_dir"], + data_dir=data_dir, + port=5433, + log_file=log_file, + extension_path=pg_paths["extension_path"], + major_version=pg_paths["major_version"], + ) + + # Skip extension install - already done by primary instance (pg_server fixture) + instance.initialize(skip_extension_install=True) + instance.start() + + yield instance + + instance.cleanup() + + +@pytest.mark.asyncio +async def test_stateless_data_sync_between_instances( + primary_conn: asyncpg.Connection, + second_instance: PostgresInstance, + temp_dir_for_postgres: str, +): + """ + Test that data created in one instance is visible from another instance. + + This test: + 1. Instance A creates a table with data at shared root_path + 2. Instance B connects with same root_path + 3. Instance B should see the table via catalog discovery + 4. Both instances can read the data + """ + shared_root_path = temp_dir_for_postgres + print(f"\n=== Test: Stateless Data Sync ===") + print(f"Shared root path: {shared_root_path}") + + # Instance A (primary): Create table and insert data + print("\n--- Instance A (port 5432): Creating table and inserting data ---") + await primary_conn.execute(f"SET deeplake.root_path = '{shared_root_path}'") + + await primary_conn.execute(""" + CREATE TABLE stateless_test ( + id INT, + name TEXT, + value FLOAT + ) USING deeplake + """) + print("Created table 'stateless_test'") + + # Insert test data + await primary_conn.execute(""" + INSERT INTO stateless_test VALUES + (1, 'alice', 100.5), + (2, 'bob', 200.75), + (3, 'charlie', 300.25) + """) + print("Inserted 3 rows") + + # Verify data in Instance A + count_a = await primary_conn.fetchval("SELECT COUNT(*) FROM stateless_test") + assert count_a == 3, f"Instance A should have 3 rows, got {count_a}" + print(f"Instance A row count: {count_a}") + + # Get the dataset path + ds_path = await primary_conn.fetchval(""" + SELECT ds_path FROM pg_deeplake_tables + WHERE table_name = 'public.stateless_test' + """) + print(f"Dataset path: {ds_path}") + + # Instance B (secondary): Connect and verify data is visible + print("\n--- Instance B (port 5433): Connecting and verifying data ---") + conn_b = await second_instance.connect() + + try: + # Setup extension (create if not exists for session-scoped instance reuse) + await conn_b.execute("CREATE EXTENSION IF NOT EXISTS pg_deeplake") + + # Setting root_path should automatically discover and register tables from catalog + await conn_b.execute(f"SET deeplake.root_path = '{shared_root_path}'") + print("Instance B: Extension loaded, root_path set") + + # Tables should be automatically discovered from the catalog at root_path + # No need to manually CREATE TABLE - they should appear after SET root_path + catalog_tables = await conn_b.fetch(""" + SELECT table_name, ds_path FROM pg_deeplake_tables + """) + print(f"Instance B catalog tables: {len(catalog_tables)}") + + # The table created by Instance A should now be visible + assert len(catalog_tables) >= 1, "Table should be auto-discovered from catalog" + + table_names = [t['table_name'] for t in catalog_tables] + assert 'public.stateless_test' in table_names, \ + f"stateless_test should be in catalog, found: {table_names}" + print("Instance B: Table auto-discovered from catalog!") + + # Verify table is visible via \dt (pg_tables) + pg_tables = await conn_b.fetch(""" + SELECT schemaname, tablename FROM pg_tables + WHERE tablename = 'stateless_test' + """) + assert len(pg_tables) == 1, "Table should be visible in pg_tables (\\dt)" + assert pg_tables[0]['schemaname'] == 'public' + assert pg_tables[0]['tablename'] == 'stateless_test' + print("Instance B: Table visible via \\dt (pg_tables)") + + # Query data from Instance B + count_b = await conn_b.fetchval("SELECT COUNT(*) FROM stateless_test") + assert count_b == 3, f"Instance B should see 3 rows, got {count_b}" + print(f"Instance B row count: {count_b}") + + # Verify data contents match + rows_b = await conn_b.fetch("SELECT id, name, value FROM stateless_test ORDER BY id") + expected = [(1, 'alice', 100.5), (2, 'bob', 200.75), (3, 'charlie', 300.25)] + + for i, (row, exp) in enumerate(zip(rows_b, expected)): + assert row['id'] == exp[0], f"Row {i} id mismatch" + assert row['name'] == exp[1], f"Row {i} name mismatch" + assert abs(row['value'] - exp[2]) < 0.001, f"Row {i} value mismatch" + print("Instance B: All data verified correctly!") + + finally: + # No need to DROP on instance B - table was auto-discovered, not created locally + await conn_b.execute("RESET deeplake.root_path") + await conn_b.close() + + # Cleanup Instance A (this is where the table was actually created) + await primary_conn.execute("DROP TABLE IF EXISTS stateless_test CASCADE") + await primary_conn.execute("RESET deeplake.root_path") + print("\n=== Test Passed: Data sync works between instances ===") + + +@pytest.mark.asyncio +async def test_stateless_concurrent_writes( + primary_conn: asyncpg.Connection, + second_instance: PostgresInstance, + temp_dir_for_postgres: str, +): + """ + Test that both instances can write to shared tables. + + This test: + 1. Creates a shared table from Instance A + 2. Both instances insert data concurrently + 3. Both instances should see all the data + """ + shared_root_path = temp_dir_for_postgres + print(f"\n=== Test: Concurrent Writes ===") + print(f"Shared root path: {shared_root_path}") + + # Instance A: Create table + await primary_conn.execute(f"SET deeplake.root_path = '{shared_root_path}'") + await primary_conn.execute(""" + CREATE TABLE concurrent_test ( + id INT, + source TEXT + ) USING deeplake + """) + print("Created shared table 'concurrent_test'") + + # Instance B: Connect and discover table via root_path + conn_b = await second_instance.connect() + try: + await conn_b.execute("CREATE EXTENSION IF NOT EXISTS pg_deeplake") + + # Setting root_path should auto-discover tables from deeplake catalog + await conn_b.execute(f"SET deeplake.root_path = '{shared_root_path}'") + + # Verify table was auto-discovered + catalog_tables = await conn_b.fetch(""" + SELECT table_name FROM pg_deeplake_tables + """) + table_names = [t['table_name'] for t in catalog_tables] + assert 'public.concurrent_test' in table_names, \ + f"concurrent_test should be auto-discovered, found: {table_names}" + + # Verify table is visible via \dt (pg_tables) + pg_tables = await conn_b.fetch(""" + SELECT schemaname, tablename FROM pg_tables + WHERE tablename = 'concurrent_test' + """) + assert len(pg_tables) == 1, "Table should be visible in pg_tables (\\dt)" + print("Instance B: Table visible via \\dt") + + # Insert from Instance A + print("Instance A: Inserting data...") + await primary_conn.execute(""" + INSERT INTO concurrent_test VALUES + (1, 'instance_a'), + (2, 'instance_a'), + (3, 'instance_a') + """) + + # Insert from Instance B + print("Instance B: Inserting data...") + await conn_b.execute(""" + INSERT INTO concurrent_test VALUES + (4, 'instance_b'), + (5, 'instance_b'), + (6, 'instance_b') + """) + + # Verify total count from both instances + count_a = await primary_conn.fetchval("SELECT COUNT(*) FROM concurrent_test") + count_b = await conn_b.fetchval("SELECT COUNT(*) FROM concurrent_test") + + print(f"Instance A sees: {count_a} rows") + print(f"Instance B sees: {count_b} rows") + + # Both should see all 6 rows (after sync) + assert count_a == 6, f"Instance A should see 6 rows, got {count_a}" + assert count_b == 6, f"Instance B should see 6 rows, got {count_b}" + + # Verify data from both sources is present + sources_a = await primary_conn.fetch("SELECT DISTINCT source FROM concurrent_test ORDER BY source") + sources_b = await conn_b.fetch("SELECT DISTINCT source FROM concurrent_test ORDER BY source") + + assert len(sources_a) == 2, "Should have data from both instances" + assert len(sources_b) == 2, "Should have data from both instances" + print("Both instances see data from both sources!") + + finally: + # No need to DROP on instance B - table was auto-discovered + await conn_b.execute("RESET deeplake.root_path") + await conn_b.close() + + # Cleanup on Instance A (where table was created) + await primary_conn.execute("DROP TABLE IF EXISTS concurrent_test CASCADE") + await primary_conn.execute("RESET deeplake.root_path") + print("\n=== Test Passed: Concurrent writes work ===") + + +@pytest.mark.asyncio +async def test_stateless_multiple_tables_discovery( + primary_conn: asyncpg.Connection, + second_instance: PostgresInstance, + temp_dir_for_postgres: str, +): + """ + Test that multiple tables created on one instance are all discoverable. + + This tests the catalog functionality for listing all tables. + """ + shared_root_path = temp_dir_for_postgres + print(f"\n=== Test: Multiple Tables Discovery ===") + print(f"Shared root path: {shared_root_path}") + + # Instance A: Create multiple tables + await primary_conn.execute(f"SET deeplake.root_path = '{shared_root_path}'") + + table_names = ['users', 'orders', 'products'] + + for table in table_names: + await primary_conn.execute(f""" + CREATE TABLE {table} ( + id INT, + name TEXT + ) USING deeplake + """) + await primary_conn.execute(f"INSERT INTO {table} VALUES (1, '{table}_data')") + print(f"Created table '{table}'") + + # Instance B: Discover and access all tables via root_path + conn_b = await second_instance.connect() + try: + await conn_b.execute("CREATE EXTENSION IF NOT EXISTS pg_deeplake") + + # Setting root_path should auto-discover ALL tables from deeplake catalog + await conn_b.execute(f"SET deeplake.root_path = '{shared_root_path}'") + + # Verify all tables were auto-discovered + catalog_tables = await conn_b.fetch(""" + SELECT table_name FROM pg_deeplake_tables + """) + discovered_names = [t['table_name'] for t in catalog_tables] + + for table in table_names: + assert f'public.{table}' in discovered_names, \ + f"Table {table} should be auto-discovered, found: {discovered_names}" + + print(f"Instance B: All {len(table_names)} tables auto-discovered!") + + # Verify all tables are visible via \dt (pg_tables) + pg_tables = await conn_b.fetch(""" + SELECT tablename FROM pg_tables + WHERE tablename IN ('users', 'orders', 'products') + """) + assert len(pg_tables) == 3, f"All 3 tables should be visible in pg_tables (\\dt), found {len(pg_tables)}" + print("Instance B: All tables visible via \\dt") + + # Verify all tables are accessible and have correct data (no CREATE TABLE needed!) + for table in table_names: + count = await conn_b.fetchval(f"SELECT COUNT(*) FROM {table}") + assert count == 1, f"Table {table} should have 1 row" + data = await conn_b.fetchval(f"SELECT name FROM {table}") + assert data == f'{table}_data', f"Table {table} has wrong data" + print(f"Instance B: Table '{table}' verified") + + print("All tables discovered and verified!") + + finally: + await conn_b.execute("RESET deeplake.root_path") + await conn_b.close() + + # Cleanup on Instance A + for table in table_names: + await primary_conn.execute(f"DROP TABLE IF EXISTS {table} CASCADE") + await primary_conn.execute("RESET deeplake.root_path") + print("\n=== Test Passed: Multiple tables discovery works ===") + + +@pytest.mark.asyncio +async def test_stateless_root_path_idempotent( + primary_conn: asyncpg.Connection, + temp_dir_for_postgres: str, +): + """ + Test that setting deeplake.root_path is idempotent. + + Setting the same root_path multiple times should: + - Not cause errors + - Not create duplicate table entries + - Maintain correct table visibility + """ + shared_root_path = temp_dir_for_postgres + print(f"\n=== Test: Root Path Idempotency ===") + print(f"Shared root path: {shared_root_path}") + + # Set root_path and create tables + await primary_conn.execute(f"SET deeplake.root_path = '{shared_root_path}'") + await primary_conn.execute(""" + CREATE TABLE idempotent_test (id INT, data TEXT) USING deeplake + """) + await primary_conn.execute("INSERT INTO idempotent_test VALUES (1, 'test')") + print("Created table 'idempotent_test'") + + # Verify initial state + count1 = await primary_conn.fetchval(""" + SELECT COUNT(*) FROM pg_deeplake_tables WHERE table_name = 'public.idempotent_test' + """) + assert count1 == 1, f"Should have exactly 1 catalog entry, got {count1}" + print(f"Initial catalog entries: {count1}") + + # Set the SAME root_path again multiple times + for i in range(3): + await primary_conn.execute(f"SET deeplake.root_path = '{shared_root_path}'") + print(f"Set root_path again (iteration {i + 1})") + + # Verify no duplicate entries + count2 = await primary_conn.fetchval(""" + SELECT COUNT(*) FROM pg_deeplake_tables WHERE table_name = 'public.idempotent_test' + """) + assert count2 == 1, f"Should still have exactly 1 catalog entry after re-setting, got {count2}" + print(f"Catalog entries after re-setting: {count2}") + + # Verify table is still accessible + data = await primary_conn.fetchval("SELECT data FROM idempotent_test WHERE id = 1") + assert data == 'test', f"Data should be 'test', got {data}" + print("Table still accessible with correct data") + + # Verify pg_tables has exactly one entry + pg_count = await primary_conn.fetchval(""" + SELECT COUNT(*) FROM pg_tables WHERE tablename = 'idempotent_test' + """) + assert pg_count == 1, f"Should have exactly 1 pg_tables entry, got {pg_count}" + print(f"pg_tables entries: {pg_count}") + + # Test RESET and re-SET + await primary_conn.execute("RESET deeplake.root_path") + print("Reset root_path") + + await primary_conn.execute(f"SET deeplake.root_path = '{shared_root_path}'") + print("Set root_path again after reset") + + # Should still work correctly + count3 = await primary_conn.fetchval(""" + SELECT COUNT(*) FROM pg_deeplake_tables WHERE table_name = 'public.idempotent_test' + """) + assert count3 == 1, f"Should have exactly 1 catalog entry after reset/re-set, got {count3}" + + data2 = await primary_conn.fetchval("SELECT data FROM idempotent_test WHERE id = 1") + assert data2 == 'test', f"Data should still be 'test' after reset/re-set, got {data2}" + print("Table still accessible after reset/re-set") + + # Cleanup + await primary_conn.execute("DROP TABLE IF EXISTS idempotent_test CASCADE") + await primary_conn.execute("RESET deeplake.root_path") + print("\n=== Test Passed: Root path setting is idempotent ===") diff --git a/scripts/build_pg_ext.py b/scripts/build_pg_ext.py index a7bc8733f0..01e2e62198 100644 --- a/scripts/build_pg_ext.py +++ b/scripts/build_pg_ext.py @@ -16,6 +16,27 @@ Usage: python3 scripts/build_pg_ext.py prod --pg-versions all #Build for all supported PostgreSQL versions """ +def get_pinned_version(): + """ + Read the pinned deeplake API version from DEEPLAKE_API_VERSION file. + """ + # Look for version file in repo root (one level up from scripts/) + script_dir = os.path.dirname(os.path.abspath(__file__)) + repo_root = os.path.dirname(script_dir) + version_file = os.path.join(repo_root, "DEEPLAKE_API_VERSION") + + if not os.path.exists(version_file): + raise Exception(f"Version file not found: {version_file}") + + with open(version_file, 'r') as f: + version = f.read().strip() + + if not version: + raise Exception(f"Version file is empty: {version_file}") + + return version + + def download_api_lib(api_root_dir, overwrite=True): """ Download and extract the full deeplake API library including: @@ -29,23 +50,9 @@ def download_api_lib(api_root_dir, overwrite=True): machine = platform.machine() - # Get latest release from GitHub - api_url = "https://api.github.com/repos/activeloopai/deeplake/releases/latest" - print(f"Fetching latest release from {api_url} ...") - - response = requests.get(api_url) - if response.status_code != 200: - raise Exception(f"Failed to fetch latest release info. Status code: {response.status_code}") - - release_data = response.json() - tag_name = release_data.get("tag_name") - if not tag_name: - raise Exception("Failed to get tag_name from latest release") - - print(f"Latest release: {tag_name}") - - # Strip 'v' prefix from tag name if present (e.g., v4.4.2 -> 4.4.2) - version = tag_name.lstrip('v') + # Get pinned version from DEEPLAKE_API_VERSION file + version = get_pinned_version() + print(f"Using pinned deeplake API version: {version}") # Check if library already exists lib_dir = os.path.join(api_root_dir, "lib") @@ -58,22 +65,15 @@ def download_api_lib(api_root_dir, overwrite=True): print(f"Library version {version} already exists. Skipping download.") return else: - print(f"Found existing version {existing_version}, but latest is {version}. Downloading latest...") + print(f"Found existing version {existing_version}, but pinned version is {version}. Downloading...") # Construct asset name based on platform archive_name = f"deeplake-api-{version}-linux-{machine}" zip_archive_name = f"{archive_name}.zip" tar_archive_name = f"{archive_name}.tar.gz" - # Find the matching asset in the release - asset_url = None - for asset in release_data.get("assets", []): - if asset.get("name") == zip_archive_name: - asset_url = asset.get("browser_download_url") - break - - if not asset_url: - raise Exception(f"Could not find asset '{zip_archive_name}' in latest release") + # Construct download URL directly (GitHub releases follow a predictable URL pattern) + asset_url = f"https://github.com/activeloopai/deeplake/releases/download/v{version}/{zip_archive_name}" print(f"Downloading prebuilt api libraries from {asset_url} ...")