Fixed add/rename column with custom schema and search path. (#3098)

khustup2 · web-flow · commit 5e830a920e2f · 2025-12-28T20:36:01.000-08:00
diff --git a/cpp/deeplake_pg/extension_init.cpp b/cpp/deeplake_pg/extension_init.cpp
@@ -611,8 +611,26 @@ static void process_utility(PlannedStmt* pstmt,
 
         /// Extract table name
         RangeVar* rel = stmt->relation;
-        const std::string schema_name = (rel->schemaname != nullptr ? rel->schemaname : "public");
-        const std::string table_name = schema_name + "." + rel->relname;
+
+        // Get the actual relation OID to resolve the proper schema name (handles custom schemas via search_path)
+        Oid rel_oid = RangeVarGetRelid(rel, NoLock, false);
+        Relation temp_rel = RelationIdGetRelation(rel_oid);
+        std::string table_name;
+
+        if (RelationIsValid(temp_rel)) {
+            Oid nspid = RelationGetNamespace(temp_rel);
+            char* nspname = get_namespace_name(nspid);
+            table_name = std::string(nspname ? nspname : "public") + "." + RelationGetRelationName(temp_rel);
+            if (nspname) {
+                pfree(nspname);
+            }
+            RelationClose(temp_rel);
+        } else {
+            // Fallback to RangeVar if relation is invalid (shouldn't happen)
+            const std::string schema_name = (rel->schemaname != nullptr ? rel->schemaname : "public");
+            table_name = schema_name + "." + rel->relname;
+        }
+
         auto* td = pg::table_storage::instance().get_table_data_if_exists(table_name);
         if (td != nullptr) {
             ListCell* lc = nullptr;
@@ -732,8 +750,26 @@ static void process_utility(PlannedStmt* pstmt,
     if (IsA(pstmt->utilityStmt, AlterTableStmt)) {
         AlterTableStmt* stmt = (AlterTableStmt*)pstmt->utilityStmt;
         RangeVar* rel = stmt->relation;
-        const std::string schema_name = (rel->schemaname != nullptr ? rel->schemaname : "public");
-        const std::string table_name = schema_name + "." + rel->relname;
+
+        // Get the actual relation OID to resolve the proper schema name (handles custom schemas via search_path)
+        Oid rel_oid = RangeVarGetRelid(rel, NoLock, false);
+        Relation temp_rel = RelationIdGetRelation(rel_oid);
+        std::string table_name;
+
+        if (RelationIsValid(temp_rel)) {
+            Oid nspid = RelationGetNamespace(temp_rel);
+            char* nspname = get_namespace_name(nspid);
+            table_name = std::string(nspname ? nspname : "public") + "." + RelationGetRelationName(temp_rel);
+            if (nspname) {
+                pfree(nspname);
+            }
+            RelationClose(temp_rel);
+        } else {
+            // Fallback to RangeVar if relation is invalid (shouldn't happen)
+            const std::string schema_name = (rel->schemaname != nullptr ? rel->schemaname : "public");
+            table_name = schema_name + "." + rel->relname;
+        }
+
         auto* td = pg::table_storage::instance().get_table_data_if_exists(table_name);
 
         if (td != nullptr) {
@@ -746,7 +782,6 @@ static void process_utility(PlannedStmt* pstmt,
                     const char* column_name = coldef->colname;
 
                     // Get the relation to query the new column's type from catalog
-                    Oid rel_oid = RangeVarGetRelid(rel, NoLock, false);
                     Relation relation = RelationIdGetRelation(rel_oid);
                     if (RelationIsValid(relation)) {
                         TupleDesc tupdesc = RelationGetDescr(relation);
@@ -981,8 +1016,26 @@ static void process_utility(PlannedStmt* pstmt,
         RenameStmt* stmt = (RenameStmt*)pstmt->utilityStmt;
         if (stmt->renameType == OBJECT_COLUMN && stmt->relation != nullptr) {
             RangeVar* rel = stmt->relation;
-            const std::string schema_name = (rel->schemaname != nullptr ? rel->schemaname : "public");
-            const std::string table_name = schema_name + "." + rel->relname;
+
+            // Get the actual relation OID to resolve the proper schema name (handles custom schemas via search_path)
+            Oid rel_oid = RangeVarGetRelid(rel, NoLock, false);
+            Relation temp_rel = RelationIdGetRelation(rel_oid);
+            std::string table_name;
+
+            if (RelationIsValid(temp_rel)) {
+                Oid nspid = RelationGetNamespace(temp_rel);
+                char* nspname = get_namespace_name(nspid);
+                table_name = std::string(nspname ? nspname : "public") + "." + RelationGetRelationName(temp_rel);
+                if (nspname) {
+                    pfree(nspname);
+                }
+                RelationClose(temp_rel);
+            } else {
+                // Fallback to RangeVar if relation is invalid (shouldn't happen)
+                const std::string schema_name = (rel->schemaname != nullptr ? rel->schemaname : "public");
+                table_name = schema_name + "." + rel->relname;
+            }
+
             auto* td = pg::table_storage::instance().get_table_data_if_exists(table_name);
 
             if (td != nullptr && stmt->subname != nullptr && stmt->newname != nullptr) {
@@ -1018,9 +1071,26 @@ static void process_utility(PlannedStmt* pstmt,
         CopyStmt* copy_stmt = (CopyStmt*)pstmt->utilityStmt;
         if (copy_stmt->relation) {
             // Build the qualified table name
-            const char* schema = copy_stmt->relation->schemaname ? copy_stmt->relation->schemaname : "public";
-            const char* table = copy_stmt->relation->relname;
-            std::string table_name = std::string(schema) + "." + table;
+            // Get the actual relation OID to resolve the proper schema name (handles custom schemas via search_path)
+            Oid rel_oid = RangeVarGetRelid(copy_stmt->relation, NoLock, false);
+            Relation temp_rel = RelationIdGetRelation(rel_oid);
+            std::string table_name;
+
+            if (RelationIsValid(temp_rel)) {
+                Oid nspid = RelationGetNamespace(temp_rel);
+                char* nspname = get_namespace_name(nspid);
+                table_name = std::string(nspname ? nspname : "public") + "." + RelationGetRelationName(temp_rel);
+                if (nspname) {
+                    pfree(nspname);
+                }
+                RelationClose(temp_rel);
+            } else {
+                // Fallback to RangeVar if relation is invalid (shouldn't happen)
+                const char* schema = copy_stmt->relation->schemaname ? copy_stmt->relation->schemaname : "public";
+                const char* table = copy_stmt->relation->relname;
+                table_name = std::string(schema) + "." + table;
+            }
+
             // If this is a deeplake table, flush/commit
             auto* td = pg::table_storage::instance().get_table_data_if_exists(table_name);
             if (td) {
diff --git a/postgres/tests/py_tests/test_add_column.py b/postgres/tests/py_tests/test_add_column.py
@@ -142,3 +142,78 @@ async def test_add_column_to_table(db_conn: asyncpg.Connection):
             await db_conn.execute("DROP TABLE IF EXISTS test_crash CASCADE")
         except:
             pass  # Connection may be dead after segfault
+
+
+@pytest.mark.asyncio
+async def test_add_column_with_custom_schema(db_conn: asyncpg.Connection):
+    """
+    Test that ALTER TABLE ADD COLUMN works correctly with custom schema.
+
+    Tests:
+    - Create a custom schema 'valod'
+    - Create table with id (INT4) and name (TEXT) columns using deeplake
+    - Insert one row with explicit id and empty string
+    - Verify data before adding column
+    - Add a new TEXT column 'surname' via ALTER TABLE
+    - Verify the column is added and existing data is preserved
+    - Verify SELECT returns correct columns after ALTER TABLE
+    """
+    assertions = Assertions(db_conn)
+
+    try:
+        # Create custom schema
+        await db_conn.execute("CREATE SCHEMA valod")
+
+        # Set search path to use the custom schema
+        await db_conn.execute("SET search_path TO 'valod'")
+
+        # Create table with id and name columns using deeplake storage
+        await db_conn.execute("""
+            CREATE TABLE a (
+                id INT4,
+                name TEXT
+            ) USING deeplake
+        """)
+
+        # Insert one row with explicit id and empty string
+        await db_conn.execute("INSERT INTO a VALUES (1, '')")
+
+        # Verify row was inserted
+        rows_before = await db_conn.fetch("SELECT * FROM a")
+        assert len(rows_before) == 1, f"Expected 1 row, got {len(rows_before)}"
+        assert rows_before[0]['id'] == 1, f"Expected id=1, got {rows_before[0]['id']}"
+        assert rows_before[0]['name'] == '', f"Expected empty string, got '{rows_before[0]['name']}'"
+
+        # Add new column 'surname' with TEXT type
+        await db_conn.execute("ALTER TABLE a ADD COLUMN surname TEXT")
+
+        # Verify column was added to PostgreSQL catalog
+        column_info = await db_conn.fetch("""
+            SELECT column_name, data_type
+            FROM information_schema.columns
+            WHERE table_schema = 'valod' AND table_name = 'a'
+            ORDER BY ordinal_position
+        """)
+        column_names = [col['column_name'] for col in column_info]
+        assert column_names == ['id', 'name', 'surname'], \
+            f"Expected ['id', 'name', 'surname'], got {column_names}"
+
+        # SELECT from table after adding column
+        rows_after = await db_conn.fetch("SELECT * FROM a")
+
+        # Verify existing data is preserved and new column exists
+        assert len(rows_after) == 1, f"Expected 1 row, got {len(rows_after)}"
+        assert rows_after[0]['id'] == 1, f"Expected id=1, got {rows_after[0]['id']}"
+        assert rows_after[0]['name'] == '', f"Expected empty string, got '{rows_after[0]['name']}'"
+        assert 'surname' in rows_after[0].keys(), "Expected 'surname' column to exist"
+        assert rows_after[0]['surname'] == '', \
+            f"Expected empty string for new column, got '{rows_after[0]['surname']}'"
+
+        print("✓ Test passed: ALTER TABLE ADD COLUMN works with custom schema")
+
+    finally:
+        # Cleanup
+        try:
+            await db_conn.execute("DROP SCHEMA IF EXISTS valod CASCADE")
+        except:
+            pass  # Connection may be dead after segfault
diff --git a/postgres/tests/py_tests/test_rename_column.py b/postgres/tests/py_tests/test_rename_column.py
@@ -148,3 +148,118 @@ async def test_rename_column_in_table(db_conn: asyncpg.Connection):
             await db_conn.execute("DROP TABLE IF EXISTS test_rename CASCADE")
         except:
             pass  # Connection may be dead after crash
+
+
+@pytest.mark.asyncio
+async def test_rename_column_with_custom_schema(db_conn: asyncpg.Connection):
+    """
+    Test that ALTER TABLE RENAME COLUMN works correctly with custom schema.
+
+    Tests:
+    - Create a custom schema 'test_schema'
+    - Create table with id (INT4), name (TEXT), and email (VARCHAR) columns using deeplake
+    - Insert one row with data
+    - Verify data before renaming column
+    - Rename the 'email' column to 'email_address' via ALTER TABLE
+    - Verify the column is renamed in both PostgreSQL catalog and deeplake dataset
+    - Verify existing data is preserved and accessible via new column name
+    - Verify old column name no longer exists
+    """
+    assertions = Assertions(db_conn)
+
+    try:
+        # Create custom schema
+        await db_conn.execute("CREATE SCHEMA test_schema")
+
+        # Set search path to use the custom schema
+        await db_conn.execute("SET search_path TO 'test_schema'")
+
+        # Create table with id, name, and email columns using deeplake storage
+        await db_conn.execute("""
+            CREATE TABLE test_rename_schema (
+                id INT4,
+                name TEXT,
+                email VARCHAR(255)
+            ) USING deeplake
+        """)
+
+        # Insert one row with explicit data
+        await db_conn.execute("""
+            INSERT INTO test_rename_schema VALUES (1, 'Alice', 'alice@example.com')
+        """)
+
+        # Verify row was inserted
+        rows_before = await db_conn.fetch("SELECT * FROM test_rename_schema")
+        assert len(rows_before) == 1, f"Expected 1 row, got {len(rows_before)}"
+        assert rows_before[0]['id'] == 1, f"Expected id=1, got {rows_before[0]['id']}"
+        assert rows_before[0]['name'] == 'Alice', f"Expected 'Alice', got '{rows_before[0]['name']}'"
+        assert rows_before[0]['email'] == 'alice@example.com', \
+            f"Expected 'alice@example.com', got '{rows_before[0]['email']}'"
+
+        # Rename email column to email_address
+        await db_conn.execute("""
+            ALTER TABLE test_rename_schema RENAME COLUMN email TO email_address
+        """)
+
+        # Verify column was renamed in PostgreSQL catalog
+        column_info = await db_conn.fetch("""
+            SELECT column_name, data_type, character_maximum_length
+            FROM information_schema.columns
+            WHERE table_schema = 'test_schema' AND table_name = 'test_rename_schema'
+            ORDER BY ordinal_position
+        """)
+        column_names = [col['column_name'] for col in column_info]
+
+        # Verify new column name exists
+        assert 'email_address' in column_names, \
+            f"Column 'email_address' should exist in catalog. Found: {column_names}"
+
+        # Verify old column name does not exist
+        assert 'email' not in column_names, \
+            f"Column 'email' should not exist anymore. Found: {column_names}"
+
+        # Verify expected column order
+        assert column_names == ['id', 'name', 'email_address'], \
+            f"Expected ['id', 'name', 'email_address'], got {column_names}"
+
+        # SELECT from table after renaming column
+        rows_after = await db_conn.fetch("SELECT * FROM test_rename_schema")
+
+        # Verify existing data is preserved and accessible via new column name
+        assert len(rows_after) == 1, f"Expected 1 row, got {len(rows_after)}"
+        assert rows_after[0]['id'] == 1, f"Expected id=1, got {rows_after[0]['id']}"
+        assert rows_after[0]['name'] == 'Alice', f"Expected 'Alice', got '{rows_after[0]['name']}'"
+        assert rows_after[0]['email_address'] == 'alice@example.com', \
+            f"Expected 'alice@example.com', got '{rows_after[0]['email_address']}'"
+
+        # Verify 'email_address' column is in the result
+        assert 'email_address' in rows_after[0].keys(), "Expected 'email_address' column to exist"
+
+        # Verify old column name is not accessible (should raise error)
+        try:
+            await db_conn.fetch("SELECT email FROM test_rename_schema")
+            assert False, "Should have raised error for non-existent column 'email'"
+        except asyncpg.UndefinedColumnError:
+            pass  # Expected error
+
+        # Insert another row using the new column name
+        await db_conn.execute("""
+            INSERT INTO test_rename_schema VALUES (2, 'Bob', 'bob@example.com')
+        """)
+
+        # Verify the new row was inserted successfully
+        all_rows = await db_conn.fetch("SELECT * FROM test_rename_schema ORDER BY id")
+        assert len(all_rows) == 2, f"Expected 2 rows, got {len(all_rows)}"
+        assert all_rows[1]['id'] == 2, f"Expected id=2, got {all_rows[1]['id']}"
+        assert all_rows[1]['name'] == 'Bob', f"Expected 'Bob', got '{all_rows[1]['name']}'"
+        assert all_rows[1]['email_address'] == 'bob@example.com', \
+            f"Expected 'bob@example.com', got '{all_rows[1]['email_address']}'"
+
+        print("✓ Test passed: ALTER TABLE RENAME COLUMN works with custom schema")
+
+    finally:
+        # Cleanup
+        try:
+            await db_conn.execute("DROP SCHEMA IF EXISTS test_schema CASCADE")
+        except:
+            pass  # Connection may be dead after crash
