fix: Update ALTER TABLE statement to include schema name in construct…

…_alter_table_statement function (paradedb#1481)
ericmand · Aug 7, 2024 · 8c567e9 · 8c567e9
1 parent 7dc87fd
commit 8c567e9
Show file tree

Hide file tree

Showing 2 changed files with 50 additions and 2 deletions.
diff --git a/pg_lakehouse/src/fdw/trigger.rs b/pg_lakehouse/src/fdw/trigger.rs
@@ -167,7 +167,7 @@ unsafe fn auto_create_schema_impl(fcinfo: pg_sys::FunctionCallInfo) -> Result<()
         .get("preserve_casing")
         .map_or(false, |s| s.eq_ignore_ascii_case("true"));
     let alter_table_statement =
-        construct_alter_table_statement(table_name, schema_rows, preserve_casing);
+        construct_alter_table_statement(schema_name, table_name, schema_rows, preserve_casing);
     Spi::run(alter_table_statement.as_str())?;
 
     Ok(())
@@ -231,6 +231,7 @@ fn duckdb_type_to_pg(column_name: &str, duckdb_type: &str) -> Result<String> {
 
 #[inline]
 fn construct_alter_table_statement(
+    schema_name: &str,
     table_name: &str,
     columns: Vec<(String, String)>,
     preserve_casing: bool,
@@ -251,7 +252,8 @@ fn construct_alter_table_statement(
         .collect();
 
     format!(
-        "ALTER TABLE {} {}",
+        "ALTER TABLE {}.{} {}",
+        spi::quote_identifier(schema_name),
         spi::quote_identifier(table_name),
         column_definitions.join(", ")
     )

diff --git a/pg_lakehouse/tests/table_config.rs b/pg_lakehouse/tests/table_config.rs
@@ -189,3 +189,49 @@ async fn test_preserve_casing_on_table_create(
 
     Ok(())
 }
+
+#[rstest]
+async fn test_table_with_custom_schema(mut conn: PgConnection, tempdir: TempDir) -> Result<()> {
+    let stored_batch = primitive_record_batch()?;
+    let parquet_path = tempdir.path().join("test_arrow_types.parquet");
+    let parquet_file = File::create(&parquet_path)?;
+
+    let mut writer = ArrowWriter::try_new(parquet_file, stored_batch.schema(), None).unwrap();
+    writer.write(&stored_batch)?;
+    writer.close()?;
+
+    primitive_setup_fdw_local_file_listing(parquet_path.as_path().to_str().unwrap(), "MyTable")
+        .execute(&mut conn);
+
+    // test quoted schema name
+    "CREATE SCHEMA \"MY_SCHEMA\"".to_string().execute(&mut conn);
+    match
+        format!( "CREATE FOREIGN TABLE \"MY_SCHEMA\".\"MyTable\" () SERVER parquet_server OPTIONS (files '{}', preserve_casing 'true')",
+                parquet_path.to_str().unwrap()
+        ).execute_result(&mut conn) {
+        Ok(_) => {}
+        Err(error) => {
+            panic!(
+                "should have successfully created table with custom schema \"MY_SCHEMA\".\"MyTable\": {}",
+                error
+            );
+        }
+    }
+
+    // test non-quoted schema name
+    "CREATE SCHEMA MY_SCHEMA".to_string().execute(&mut conn);
+    match
+        format!( "CREATE FOREIGN TABLE MY_SCHEMA.MyTable () SERVER parquet_server OPTIONS (files '{}', preserve_casing 'true')",
+                parquet_path.to_str().unwrap()
+        ).execute_result(&mut conn) {
+        Ok(_) => {}
+        Err(error) => {
+            panic!(
+                "should have successfully created table with custom schema MY_SCHEMA.MyTable: {}",
+                error
+            );
+        }
+    }
+
+    Ok(())
+}