From ae0cc16de519312643c8069c193300e8f776338e Mon Sep 17 00:00:00 2001
From: Siyuan Huang <73871299+kysshsy@users.noreply.github.com>
Date: Sun, 15 Dec 2024 09:22:24 +0800
Subject: [PATCH] test: add a test of creating a foreign table with reserved
 column name (#182)

* fix: reserved column name in auto create schema

* fix: remove code change and fix test
---
 tests/tests/fixtures/arrow.rs | 40 +++++++++++++++++++++++++++++++++++
 tests/tests/table_config.rs   | 30 +++++++++++++++++++++++++-
 2 files changed, 69 insertions(+), 1 deletion(-)

diff --git a/tests/tests/fixtures/arrow.rs b/tests/tests/fixtures/arrow.rs
index 56578502..f7caa58c 100644
--- a/tests/tests/fixtures/arrow.rs
+++ b/tests/tests/fixtures/arrow.rs
@@ -244,6 +244,46 @@ pub fn primitive_record_batch_single() -> Result<RecordBatch> {
     )?)
 }
 
+pub fn reserved_column_record_batch() -> Result<RecordBatch> {
+    // authorization is a reserved column name
+    let fields = vec![
+        Field::new("id", DataType::Utf8, false),
+        Field::new("operation_id", DataType::Utf8, false),
+        Field::new("order_id", DataType::Utf8, false),
+        Field::new("dealer_id", DataType::Utf8, false),
+        Field::new("dealer_name", DataType::Utf8, false),
+        Field::new("authorization", DataType::Utf8, false),
+        Field::new("on_day", DataType::Boolean, false),
+    ];
+
+    let schema = Arc::new(Schema::new(fields));
+
+    Ok(RecordBatch::try_new(
+        schema,
+        vec![
+            Arc::new(StringArray::from(vec![Some("1"), Some("2"), Some("3")])),
+            Arc::new(StringArray::from(vec![Some("1"), Some("2"), Some("3")])),
+            Arc::new(StringArray::from(vec![Some("1"), Some("2"), Some("3")])),
+            Arc::new(StringArray::from(vec![Some("1"), Some("2"), Some("3")])),
+            Arc::new(StringArray::from(vec![
+                Some("Jason"),
+                Some("Alice"),
+                Some("Bob"),
+            ])),
+            Arc::new(StringArray::from(vec![
+                Some("auth_1"),
+                Some("auth_2"),
+                Some("auth_3"),
+            ])),
+            Arc::new(BooleanArray::from(vec![
+                Some(true),
+                Some(false),
+                Some(true),
+            ])),
+        ],
+    )?)
+}
+
 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..3bd247da 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 '{}', preserve_casing 'true')",
+        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")