Skip to content

Commit

Permalink
fix: reserved column name in auto create schema
Browse files Browse the repository at this point in the history
  • Loading branch information
kysshsy committed Dec 13, 2024
1 parent ad3d5b6 commit 9bbc5e2
Show file tree
Hide file tree
Showing 3 changed files with 56 additions and 2 deletions.
6 changes: 5 additions & 1 deletion src/fdw/trigger.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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();

Expand Down
22 changes: 22 additions & 0 deletions tests/tests/fixtures/arrow.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
30 changes: 29 additions & 1 deletion tests/tests/table_config.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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};
Expand Down Expand Up @@ -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")
Expand Down

0 comments on commit 9bbc5e2

Please sign in to comment.