From 9bbc5e24e63ee5df18331343e7947c0fd85584f3 Mon Sep 17 00:00:00 2001
From: kysshsy <kysshsy@gmail.com>
Date: Fri, 13 Dec 2024 23:01:14 +0800
Subject: [PATCH] fix: reserved column name in auto create schema

---
 src/fdw/trigger.rs            |  6 +++++-
 tests/tests/fixtures/arrow.rs | 22 ++++++++++++++++++++++
 tests/tests/table_config.rs   | 30 +++++++++++++++++++++++++++++-
 3 files changed, 56 insertions(+), 2 deletions(-)

diff --git a/src/fdw/trigger.rs b/src/fdw/trigger.rs
index 7ac2f334..d035ee99 100644
--- a/src/fdw/trigger.rs
+++ b/src/fdw/trigger.rs
@@ -281,7 +281,11 @@ fn construct_alter_table_statement(
                 column_name.to_string()
             };
 
-            format!("ADD COLUMN {} {}", column_name, pg_type)
+            format!(
+                "ADD COLUMN {} {}",
+                spi::quote_identifier(column_name),
+                pg_type
+            )
         })
         .collect();
 
diff --git a/tests/tests/fixtures/arrow.rs b/tests/tests/fixtures/arrow.rs
index 56578502..0c004d68 100644
--- a/tests/tests/fixtures/arrow.rs
+++ b/tests/tests/fixtures/arrow.rs
@@ -244,6 +244,28 @@ pub fn primitive_record_batch_single() -> Result<RecordBatch> {
     )?)
 }
 
+pub fn reserved_column_record_batch() -> Result<RecordBatch> {
+    // INTEGER authorization is a reserved column name
+    let fields = vec![
+        Field::new("INTEGER", DataType::Int32, false),
+        Field::new("authorization", DataType::Utf8, false),
+    ];
+
+    let schema = Arc::new(Schema::new(fields));
+
+    Ok(RecordBatch::try_new(
+        schema,
+        vec![
+            Arc::new(Int32Array::from(vec![1, 2, 3])),
+            Arc::new(StringArray::from(vec![
+                Some("auth_1"),
+                Some("auth_2"),
+                Some("auth_3"),
+            ])),
+        ],
+    )?)
+}
+
 pub fn primitive_create_foreign_data_wrapper(
     wrapper: &str,
     handler: &str,
diff --git a/tests/tests/table_config.rs b/tests/tests/table_config.rs
index 531944c4..d01077c7 100644
--- a/tests/tests/table_config.rs
+++ b/tests/tests/table_config.rs
@@ -19,7 +19,8 @@ mod fixtures;
 
 use crate::fixtures::arrow::{
     primitive_record_batch, primitive_setup_fdw_local_file_listing, record_batch_with_casing,
-    setup_local_file_listing_with_casing,
+    reserved_column_record_batch, setup_local_file_listing_with_casing,
+    setup_parquet_wrapper_and_server,
 };
 use crate::fixtures::db::Query;
 use crate::fixtures::{conn, tempdir};
@@ -91,6 +92,33 @@ async fn test_reserved_table_name(mut conn: PgConnection, tempdir: TempDir) -> R
     Ok(())
 }
 
+#[rstest]
+fn test_reserved_column_name(mut conn: PgConnection, tempdir: TempDir) -> Result<()> {
+    let stored_batch = reserved_column_record_batch()?;
+    let parquet_path = tempdir.path().join("reserved_column_table.parquet");
+    let parquet_file = File::create(&parquet_path).unwrap();
+
+    let mut writer = ArrowWriter::try_new(parquet_file, stored_batch.schema(), None).unwrap();
+    writer.write(&stored_batch)?;
+    writer.close()?;
+
+    setup_parquet_wrapper_and_server().execute(&mut conn);
+
+    match format!(
+        "CREATE FOREIGN TABLE reserved_table_name () SERVER parquet_server OPTIONS (files '{}')",
+        parquet_path.to_str().unwrap()
+    )
+    .execute_result(&mut conn)
+    {
+        Ok(_) => {}
+        Err(e) => {
+            panic!("fail to create table with reserved column name: {}", e)
+        }
+    }
+
+    Ok(())
+}
+
 #[rstest]
 async fn test_invalid_file(mut conn: PgConnection) -> Result<()> {
     match primitive_setup_fdw_local_file_listing("invalid_file.parquet", "primitive")