Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
76 changes: 0 additions & 76 deletions cpp/deeplake_pg/extension_init.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -46,7 +46,6 @@ namespace pg {
bool use_parallel_workers = false;
bool use_deeplake_executor = true;
bool explain_query_before_execute = false;
bool ignore_primary_keys = true;
bool print_runtime_stats = false;
bool support_json_index = false;
bool is_filter_pushdown_enabled = true;
Expand Down Expand Up @@ -133,17 +132,6 @@ void initialize_guc_parameters()
nullptr // check_hook, assign_hook, show_hook
);

DefineCustomBoolVariable("pg_deeplake.ignore_primary_keys",
"If set to true, PRIMARY KEY constraints will be ignored during table creation.",
nullptr, // optional long description
&pg::ignore_primary_keys, // linked C variable
true, // default value
PGC_USERSET, // context (USERSET, SUSET, etc.)
0, // flags
nullptr,
nullptr,
nullptr // check_hook, assign_hook, show_hook
);

DefineCustomBoolVariable("pg_deeplake.print_runtime_stats",
"Enable runtime statistics printing for pg_deeplake operations.",
Expand Down Expand Up @@ -625,70 +613,6 @@ static void process_utility(PlannedStmt* pstmt,
list_free_deep(stmt->options);
stmt->options = NIL;
}
// Remove PRIMARY KEY constraints
if (pg::ignore_primary_keys && deeplake_table && stmt->tableElts != nullptr) {
List* new_table_elts = NIL;
ListCell* lc = nullptr;
bool has_primary_key = false;
// Get table name from the CreateStmt
std::string table_name = stmt->relation ? stmt->relation->relname : "";
std::map<std::string, std::set<std::string>> primary_keys;
foreach (lc, stmt->tableElts) {
Node* element = (Node*)lfirst(lc);
// Handle table-level PRIMARY KEY constraints
if (IsA(element, Constraint)) {
Constraint* constraint = (Constraint*)element;
if (constraint->contype == CONSTR_PRIMARY) {
has_primary_key = true;
// Extract primary key column names
if (constraint->keys != nullptr) {
ListCell* key_lc = nullptr;
foreach (key_lc, constraint->keys) {
std::string col_name = strVal(lfirst(key_lc));
elog(DEBUG1,
"Removing table-level PRIMARY KEY constraint on table '%s' for column: %s",
table_name.c_str(),
col_name.c_str());
primary_keys[table_name].insert(col_name);
}
}
continue; // Skip adding this constraint
}
}

// Handle column-level PRIMARY KEY constraints (e.g., "col INT PRIMARY KEY")
if (IsA(element, ColumnDef)) {
ColumnDef* coldef = (ColumnDef*)element;
if (coldef->constraints != nullptr) {
List* new_constraints = NIL;
ListCell* const_lc = nullptr;

foreach (const_lc, coldef->constraints) {
Constraint* constraint = (Constraint*)lfirst(const_lc);
if (constraint->contype == CONSTR_PRIMARY) {
has_primary_key = true;
elog(DEBUG1,
"Removing column-level PRIMARY KEY constraint on table '%s' for column: %s",
table_name.c_str(),
coldef->colname);
primary_keys[table_name].insert(coldef->colname);
continue;
}
new_constraints = lappend(new_constraints, constraint);
}

// Update column's constraints list
coldef->constraints = new_constraints;
}
}

new_table_elts = lappend(new_table_elts, element);
}
if (has_primary_key) {
stmt->tableElts = new_table_elts;
pg::table_storage::instance().set_primary_keys(std::move(primary_keys));
}
}
}

std::optional<pg::utils::parallel_workers_switcher> switcher;
Expand Down
19 changes: 17 additions & 2 deletions cpp/deeplake_pg/table_am.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,7 @@ extern "C" {
#include <utils/lsyscache.h>
#include <utils/rel.h>
#include <utils/relcache.h>
#include <utils/snapshot.h> // For SNAPSHOT_DIRTY and snapshot types
#include <utils/varlena.h> // For text functions

#ifdef __cplusplus
Expand Down Expand Up @@ -966,13 +967,27 @@ bool deeplake_table_am_routine::index_fetch_tuple(struct IndexFetchTableData* sc

pg::utils::memory_context_switcher context_switcher(idx_scan->memory_context);
idx_scan->scan_state.set_current_position(utils::tid_to_row_number(tid));

if (!idx_scan->scan_state.get_next_tuple(slot)) {
*all_dead = true;
if (all_dead != nullptr) {
*all_dead = true;
}
return false;
}

// For SNAPSHOT_DIRTY (used by btree unique checking),
// we need to indicate that no in-progress transaction is affecting this tuple.
// This prevents PostgreSQL from trying to look up transaction status in pg_subtrans.
// Deeplake tuples are always immediately visible (no MVCC).
if (snapshot != nullptr && snapshot->snapshot_type == SNAPSHOT_DIRTY) {
snapshot->xmin = InvalidTransactionId;
snapshot->xmax = InvalidTransactionId;
}

*call_again = false;
*all_dead = false;
if (all_dead != nullptr) {
*all_dead = false;
}
return true;
}

Expand Down
1 change: 0 additions & 1 deletion cpp/deeplake_pg/table_data.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -77,7 +77,6 @@ struct table_data
inline void clear_delete_rows() noexcept;
inline void add_update_row(int64_t row_id, icm::string_map<nd::array> update_row);
inline void clear_update_rows() noexcept;
inline void set_primary_keys(const std::set<std::string>& primary_keys);
inline Oid get_table_oid() const noexcept;
inline bool flush();

Expand Down
26 changes: 0 additions & 26 deletions cpp/deeplake_pg/table_data_impl.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -530,32 +530,6 @@ inline bool table_data::flush_updates()
return true;
}

inline void table_data::set_primary_keys(const std::set<std::string>& primary_keys)
{
get_dataset()->set_indexing_mode(deeplake::indexing_mode::always);
bool index_created = false;
for (const auto& column_name : primary_keys) {
auto& column = get_dataset()->get_column(column_name);
auto index_holder = column.index_holder();
if (index_holder != nullptr) {
continue;
}
auto column_type_kind = column.type().kind();
if (column_type_kind == deeplake_core::type_kind::generic && !column.type().data_type().is_array() &&
nd::dtype_is_numeric(column.type().data_type().get_dtype())) {
column.create_index(deeplake_core::index_type(
deeplake_core::numeric_index_type(deeplake_core::deeplake_index_type::type::inverted_index)));
index_created = true;
elog(DEBUG1, "Created numeric index on table '%s' column '%s'", table_name_.c_str(), column_name.c_str());
} else {
elog(DEBUG1, "Column %s is not supported yet for indexing", column_name.c_str());
}
}
if (index_created) {
commit();
}
}

inline Oid table_data::get_table_oid() const noexcept
{
return table_oid_;
Expand Down
13 changes: 4 additions & 9 deletions cpp/deeplake_pg/table_storage.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -677,14 +677,6 @@ void table_storage::create_table(const std::string& table_name, Oid table_id, Tu
ereport(ERROR, (errcode(ERRCODE_INTERNAL_ERROR), errmsg("%s", e.what())));
}

if (auto it = primary_keys_.find(simple_table_name); it != primary_keys_.end() && !it->second.empty()) {
try {
td.set_primary_keys(it->second);
} catch (const base::exception& e) {
elog(WARNING, "Failed to set primary keys for table %s: %s", table_name.c_str(), e.what());
}
}

tables_.emplace(table_id, std::move(td));
up_to_date_ = false;
}
Expand Down Expand Up @@ -792,7 +784,10 @@ bool table_storage::fetch_tuple(Oid table_id, ItemPointer tid, TupleTableSlot* s
ExecClearTuple(slot);

const auto row_number = utils::tid_to_row_number(tid);
if (row_number >= table_data.num_rows()) {
// Use num_total_rows() to include uncommitted rows in the current transaction.
// This is necessary for AFTER triggers (like FK checks) that need to see
// rows inserted earlier in the same transaction.
if (row_number >= table_data.num_total_rows()) {
return false;
}
Datum* values = slot->tts_values;
Expand Down
6 changes: 0 additions & 6 deletions cpp/deeplake_pg/table_storage.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -223,11 +223,6 @@ class table_storage
up_to_date_ = up_to_date;
}

inline void set_primary_keys(std::map<std::string, std::set<std::string>>&& primary_keys) noexcept
{
primary_keys_ = std::move(primary_keys);
}

private:
table_storage() = default;

Expand All @@ -236,7 +231,6 @@ class table_storage

std::unordered_map<Oid, table_data> tables_;
std::unordered_map<Oid, std::pair<std::string, std::string>> views_;
std::map<std::string, std::set<std::string>> primary_keys_;
std::string schema_name_ = "public";
bool tables_loaded_ = false;
bool up_to_date_ = true;
Expand Down
1 change: 0 additions & 1 deletion cpp/deeplake_pg/utils.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -46,7 +46,6 @@ inline static constexpr const char* index_type_option_name = "index_type";
extern bool use_parallel_workers;
extern bool use_deeplake_executor;
extern bool explain_query_before_execute;
extern bool ignore_primary_keys;
extern bool print_runtime_stats;
extern bool is_filter_pushdown_enabled;
extern int32_t max_streamable_column_width;
Expand Down
69 changes: 69 additions & 0 deletions postgres/tests/py_tests/test_constraint_enforcement.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,69 @@
"""
Tests for constraint enforcement in pg_deeplake.

These tests verify that PRIMARY KEY, UNIQUE, and FOREIGN KEY constraints
work correctly with deeplake tables.
"""
import pytest
import asyncpg


@pytest.mark.asyncio
async def test_primary_key_rejects_duplicates(db_conn: asyncpg.Connection):
"""PRIMARY KEY should reject duplicate values."""
try:
await db_conn.execute("DROP TABLE IF EXISTS pk_test CASCADE")
await db_conn.execute("""
CREATE TABLE pk_test (id INT PRIMARY KEY, name TEXT) USING deeplake
""")

await db_conn.execute("INSERT INTO pk_test VALUES (1, 'alice')")

with pytest.raises(asyncpg.UniqueViolationError):
await db_conn.execute("INSERT INTO pk_test VALUES (1, 'bob')")

finally:
await db_conn.execute("DROP TABLE IF EXISTS pk_test CASCADE")


@pytest.mark.asyncio
async def test_unique_constraint_rejects_duplicates(db_conn: asyncpg.Connection):
"""UNIQUE constraint should reject duplicate values."""
try:
await db_conn.execute("DROP TABLE IF EXISTS unique_test CASCADE")
await db_conn.execute("""
CREATE TABLE unique_test (id INT PRIMARY KEY, email TEXT UNIQUE) USING deeplake
""")

await db_conn.execute("INSERT INTO unique_test VALUES (1, 'alice@test.com')")

with pytest.raises(asyncpg.UniqueViolationError):
await db_conn.execute("INSERT INTO unique_test VALUES (2, 'alice@test.com')")

finally:
await db_conn.execute("DROP TABLE IF EXISTS unique_test CASCADE")


@pytest.mark.asyncio
async def test_foreign_key_insert(db_conn: asyncpg.Connection):
"""INSERT into child table with FK should trigger parent lookup."""
try:
await db_conn.execute("DROP TABLE IF EXISTS fk_child CASCADE")
await db_conn.execute("DROP TABLE IF EXISTS fk_parent CASCADE")

await db_conn.execute("""
CREATE TABLE fk_parent (id INT PRIMARY KEY) USING deeplake
""")
await db_conn.execute("""
CREATE TABLE fk_child (
id INT PRIMARY KEY,
parent_id INT REFERENCES fk_parent(id)
) USING deeplake
""")

await db_conn.execute("INSERT INTO fk_parent VALUES (1)")
await db_conn.execute("INSERT INTO fk_child VALUES (1, 1)")

finally:
await db_conn.execute("DROP TABLE IF EXISTS fk_child CASCADE")
await db_conn.execute("DROP TABLE IF EXISTS fk_parent CASCADE")
Loading
Loading