Skip to content

Commit 572a819

Browse files
aykut-bozkurtaykut-bozkurt
committed
add more tests
1 parent 2688cb0 commit 572a819

File tree

5 files changed

+235
-54
lines changed

5 files changed

+235
-54
lines changed

‎src/arrow_parquet/parquet_reader.rs‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ use arrow::array::RecordBatch;
77
use arrow_cast::{cast_with_options, CastOptions};
88
use arrow_schema::SchemaRef;
99
use futures::StreamExt;
10+
use glob::Pattern;
1011
use parquet::arrow::async_reader::{ParquetObjectReader, ParquetRecordBatchStream};
1112
use pgrx::{
1213
check_for_interrupts,
@@ -207,9 +208,9 @@ impl ParquetReaderContext {
207208
SingleParquetReader::try_new(uri_info, match_by, tupledesc_schema.clone(), &attributes)
208209
.map(|reader| vec![reader])
209210
.unwrap_or_else(|e| {
210-
// if uri contains any pattern, try to create readers from the pattern uri
211+
// if uri contains a valid pattern, try to create readers from the pattern uri
211212
// otherwise, panic with the original error
212-
if !uri_info.path.as_ref().contains('*') {
213+
if Pattern::try_from(uri_info).is_err() {
213214
panic!("{e}");
214215
}
215216

‎src/parquet_udfs/list.rs‎

Lines changed: 16 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -22,27 +22,19 @@ mod parquet {
2222
}
2323
}
2424

25-
fn error_if_list_http_store(uri_info: &ParsedUriInfo) {
26-
if uri_info.uri.scheme() == "http" || uri_info.uri.scheme() == "https" {
27-
panic!("list operation on http(s) object stores is not supported");
28-
}
29-
}
30-
3125
pub(crate) fn list_uri(uri_info: &ParsedUriInfo) -> Vec<(String, i64)> {
3226
ensure_access_privilege_to_uri(uri_info, true);
3327

34-
error_if_list_http_store(uri_info);
35-
3628
let base_uri = object_store_base_uri(&uri_info.uri);
3729

38-
let copy_from = true;
39-
let (parquet_object_store, location) = get_or_create_object_store(uri_info, copy_from);
40-
4130
// build the pattern before we start the stream to bail out early
42-
let pattern = Pattern::new(location.as_ref()).unwrap_or_else(|e| {
31+
let pattern = Pattern::try_from(uri_info).unwrap_or_else(|e| {
4332
panic!("{}", e);
4433
});
4534

35+
let copy_from = true;
36+
let (parquet_object_store, location) = get_or_create_object_store(uri_info, copy_from);
37+
4638
// prefix is the part of the location that doesn't contain any patterns
4739
let prefix = location
4840
.parts()
@@ -90,3 +82,15 @@ pub(crate) fn list_uri(uri_info: &ParsedUriInfo) -> Vec<(String, i64)> {
9082
})
9183
.collect::<Vec<_>>()
9284
}
85+
86+
impl TryFrom<&ParsedUriInfo> for Pattern {
87+
type Error = String;
88+
89+
fn try_from(uri_info: &ParsedUriInfo) -> Result<Self, Self::Error> {
90+
if uri_info.uri.scheme() == "http" || uri_info.uri.scheme() == "https" {
91+
return Err("list operation on http(s) object stores is not supported".into());
92+
}
93+
94+
Self::new(uri_info.path.as_ref()).map_err(|e| e.to_string())
95+
}
96+
}

‎src/pgrx_tests/copy_pattern.rs‎

Lines changed: 41 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -2,16 +2,22 @@
22
mod tests {
33
use pgrx::{pg_test, Spi};
44

5+
use crate::pgrx_tests::common::FileCleanup;
6+
57
#[pg_test]
68
#[should_panic(expected = "EOF: file size of 2 is less than footer")]
7-
fn test_pattern_invalid_parquet_file() {
9+
fn test_non_parquet_path_in_pattern() {
10+
let filename1 = "/tmp/pg_parquet_test/dummy1.parquet";
11+
let filename2 = "/tmp/pg_parquet_test/dummy2.csv";
12+
let _file1 = FileCleanup::new(filename1);
13+
let _file2 = FileCleanup::new(filename2);
14+
815
let copy_to_parquet =
9-
"copy (select 1 as a) to '/tmp/pg_parquet_test/dummy1.parquet' with (format parquet);";
10-
Spi::run(copy_to_parquet).unwrap();
16+
format!("copy (select 1 as a) to '{filename1}' with (format parquet);");
17+
Spi::run(&copy_to_parquet).unwrap();
1118

12-
let copy_to_csv =
13-
"copy (select 1 as a) to '/tmp/pg_parquet_test/dummy2.csv' with (format csv);";
14-
Spi::run(copy_to_csv).unwrap();
19+
let copy_to_csv = format!("copy (select 1 as a) to '{filename2}' with (format csv);");
20+
Spi::run(&copy_to_csv).unwrap();
1521

1622
let create_table = "create table test_table(a int);";
1723
Spi::run(create_table).unwrap();
@@ -25,8 +31,10 @@ mod tests {
2531
}
2632

2733
#[pg_test]
28-
fn test_pattern_with_special_parquet_file_name() {
34+
fn test_path_with_special_parquet_file_name() {
2935
let filename = "/tmp/pg_parquet_test/du**mm*y1.parquet";
36+
let _file = FileCleanup::new(filename);
37+
3038
let copy_to_parquet =
3139
format!("copy (select 1 as a) to '{filename}' with (format parquet);");
3240
Spi::run(&copy_to_parquet).unwrap();
@@ -41,4 +49,30 @@ mod tests {
4149
let result = Spi::get_one::<i64>(count_query).unwrap().unwrap();
4250
assert_eq!(result, 1);
4351
}
52+
53+
#[pg_test]
54+
#[should_panic(expected = "failed to get object store metadata")]
55+
fn test_path_with_special_chars() {
56+
let filename = "/tmp/pg_parquet_test/du**\\mm*y1.parquet";
57+
let _file = FileCleanup::new(filename);
58+
59+
let create_table = "create table test_table(a int);";
60+
Spi::run(create_table).unwrap();
61+
62+
let copy_from_parquet = format!("COPY test_table FROM '{filename}' WITH (format parquet);");
63+
Spi::run(copy_from_parquet.as_str()).unwrap();
64+
}
65+
66+
#[pg_test]
67+
#[should_panic(expected = "no files found that match the pattern")]
68+
fn test_with_nonexistent_pattern_path() {
69+
let filename = "/tmp/pg_parquet_test/**";
70+
let _file = FileCleanup::new(filename);
71+
72+
let create_table = "create table test_table(a int);";
73+
Spi::run(create_table).unwrap();
74+
75+
let copy_from_parquet = format!("COPY test_table FROM '{filename}' WITH (format parquet);");
76+
Spi::run(copy_from_parquet.as_str()).unwrap();
77+
}
4478
}

‎src/pgrx_tests/object_store.rs‎

Lines changed: 163 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -91,6 +91,49 @@ mod tests {
9191
test_table.assert_expected_and_result_rows();
9292
}
9393

94+
#[pg_test]
95+
fn test_s3_uri_with_special_chars() {
96+
object_store_cache_clear();
97+
98+
let test_bucket_name: String =
99+
std::env::var("AWS_S3_TEST_BUCKET").expect("AWS_S3_TEST_BUCKET not found");
100+
101+
let s3_uri = format!("s3://{test_bucket_name}/so\\mek*ey/**testme.parquet");
102+
103+
let copy_to_command = format!(
104+
"COPY (SELECT a FROM generate_series(1,10) a) TO '{s3_uri}' WITH (format parquet);"
105+
);
106+
Spi::run(&copy_to_command).unwrap();
107+
108+
let create_table = "create table test_table(a int);";
109+
Spi::run(create_table).unwrap();
110+
111+
let copy_from_command = format!("COPY test_table FROM '{s3_uri}' WITH (format parquet);");
112+
Spi::run(copy_from_command.as_str()).unwrap();
113+
114+
let count_query = "select count(*) from test_table;";
115+
let result = Spi::get_one::<i64>(count_query).unwrap().unwrap();
116+
assert_eq!(result, 10);
117+
}
118+
119+
#[pg_test]
120+
#[should_panic(expected = "no files found that match the pattern")]
121+
fn test_s3_with_nonexistent_pattern_uri() {
122+
object_store_cache_clear();
123+
124+
let test_bucket_name: String =
125+
std::env::var("AWS_S3_TEST_BUCKET").expect("AWS_S3_TEST_BUCKET not found");
126+
127+
let s3_uri_pattern = format!("s3://{test_bucket_name}/notexists/**");
128+
129+
let create_table = "create table test_table(a int);";
130+
Spi::run(create_table).unwrap();
131+
132+
let copy_from_command =
133+
format!("COPY test_table FROM '{s3_uri_pattern}' WITH (format parquet);");
134+
Spi::run(copy_from_command.as_str()).unwrap();
135+
}
136+
94137
#[pg_test]
95138
fn test_s3_from_config_file() {
96139
object_store_cache_clear();
@@ -370,22 +413,6 @@ mod tests {
370413
.unwrap();
371414
}
372415

373-
#[pg_test]
374-
#[should_panic(expected = "no files found that match the pattern")]
375-
fn test_s3_empty_pattern() {
376-
object_store_cache_clear();
377-
378-
let test_bucket_name: String =
379-
std::env::var("AWS_S3_TEST_BUCKET").expect("AWS_S3_TEST_BUCKET not found");
380-
381-
let create_table = "create table test_table(id int);";
382-
Spi::run(create_table).unwrap();
383-
384-
let s3_uri_pattern = format!("s3://{test_bucket_name}/dummy*.parquet");
385-
let copy_from_command = format!("COPY test_table FROM '{}';", s3_uri_pattern);
386-
Spi::run(copy_from_command.as_str()).unwrap();
387-
}
388-
389416
#[pg_test]
390417
#[should_panic(expected = "column count mismatch")]
391418
fn test_s3_pattern_schema_mismatch() {
@@ -491,6 +518,52 @@ mod tests {
491518
test_table.assert_expected_and_result_rows();
492519
}
493520

521+
#[pg_test]
522+
#[cfg(not(rhel8))]
523+
fn test_azure_uri_with_special_chars() {
524+
object_store_cache_clear();
525+
526+
let test_container_name: String = std::env::var("AZURE_TEST_CONTAINER_NAME")
527+
.expect("AZURE_TEST_CONTAINER_NAME not found");
528+
529+
let azure_uri = format!("az://{test_container_name}/so\\mek*ey/**testme.parquet");
530+
531+
let copy_to_command = format!(
532+
"COPY (SELECT a FROM generate_series(1,10) a) TO '{azure_uri}' WITH (format parquet);"
533+
);
534+
Spi::run(&copy_to_command).unwrap();
535+
536+
let create_table = "create table test_table(a int);";
537+
Spi::run(create_table).unwrap();
538+
539+
let copy_from_command =
540+
format!("COPY test_table FROM '{azure_uri}' WITH (format parquet);");
541+
Spi::run(copy_from_command.as_str()).unwrap();
542+
543+
let count_query = "select count(*) from test_table;";
544+
let result = Spi::get_one::<i64>(count_query).unwrap().unwrap();
545+
assert_eq!(result, 10);
546+
}
547+
548+
#[pg_test]
549+
#[cfg(not(rhel8))]
550+
#[should_panic(expected = "no files found that match the pattern")]
551+
fn test_azure_with_nonexistent_pattern_uri() {
552+
object_store_cache_clear();
553+
554+
let test_container_name: String = std::env::var("AZURE_TEST_CONTAINER_NAME")
555+
.expect("AZURE_TEST_CONTAINER_NAME not found");
556+
557+
let azure_uri_pattern = format!("az://{test_container_name}/notexists/**");
558+
559+
let create_table = "create table test_table(a int);";
560+
Spi::run(create_table).unwrap();
561+
562+
let copy_from_command =
563+
format!("COPY test_table FROM '{azure_uri_pattern}' WITH (format parquet);");
564+
Spi::run(copy_from_command.as_str()).unwrap();
565+
}
566+
494567
#[pg_test]
495568
#[cfg(not(rhel8))]
496569
fn test_azure_from_config_file() {
@@ -767,32 +840,46 @@ mod tests {
767840
}
768841

769842
#[pg_test]
770-
#[should_panic(expected = "list operation on http(s) object stores is not supported")]
771-
fn test_http_uri_glob_pattern() {
843+
fn test_http_uri_with_special_chars() {
772844
object_store_cache_clear();
773845

774846
let http_endpoint: String =
775847
std::env::var("HTTP_ENDPOINT").expect("HTTP_ENDPOINT not found");
776848

777-
let http_uri =
778-
format!("{http_endpoint}/test_http_uri_glob_pattern/some/pg_parquet_test.parquet");
779-
780-
let http_uri_pattern = format!("{http_endpoint}/test_http_uri_glob_pattern/**/*.parquet");
849+
let http_uri = format!("{http_endpoint}/so\\mek*ey/**testme.parquet");
781850

782-
let mut copy_options = HashMap::new();
783-
copy_options.insert(
784-
"file_size_bytes".to_string(),
785-
CopyOptionValue::StringOption("1MB".to_string()),
851+
let copy_to_command = format!(
852+
"COPY (SELECT a FROM generate_series(1,10) a) TO '{http_uri}' WITH (format parquet);"
786853
);
854+
Spi::run(&copy_to_command).unwrap();
787855

788-
let test_table = TestTable::<i32>::new("int4".into())
789-
.with_uri(http_uri)
790-
.with_uri_pattern(http_uri_pattern)
791-
.with_order_by_col("a".into())
792-
.with_copy_to_options(copy_options);
856+
let create_table = "create table test_table(a int);";
857+
Spi::run(create_table).unwrap();
793858

794-
test_table.insert("INSERT INTO test_expected select i from generate_series(1, 1000000) i;");
795-
test_table.assert_expected_and_result_rows();
859+
let copy_from_command = format!("COPY test_table FROM '{http_uri}' WITH (format parquet);");
860+
Spi::run(copy_from_command.as_str()).unwrap();
861+
862+
let count_query = "select count(*) from test_table;";
863+
let result = Spi::get_one::<i64>(count_query).unwrap().unwrap();
864+
assert_eq!(result, 10);
865+
}
866+
867+
#[pg_test]
868+
#[should_panic(expected = "failed to get object store metadata for uri")]
869+
fn test_http_with_nonexistent_pattern_uri() {
870+
object_store_cache_clear();
871+
872+
let http_endpoint: String =
873+
std::env::var("HTTP_ENDPOINT").expect("HTTP_ENDPOINT not found");
874+
875+
let http_uri_pattern = format!("{http_endpoint}/notexists/**");
876+
877+
let create_table = "create table test_table(a int);";
878+
Spi::run(create_table).unwrap();
879+
880+
let copy_from_command =
881+
format!("COPY test_table FROM '{http_uri_pattern}' WITH (format parquet);");
882+
Spi::run(copy_from_command.as_str()).unwrap();
796883
}
797884

798885
#[pg_test]
@@ -839,6 +926,49 @@ mod tests {
839926
test_table.assert_expected_and_result_rows();
840927
}
841928

929+
#[pg_test]
930+
fn test_gcs_uri_with_special_chars() {
931+
object_store_cache_clear();
932+
933+
let test_bucket_name: String =
934+
std::env::var("GOOGLE_TEST_BUCKET").expect("GOOGLE_TEST_BUCKET not found");
935+
936+
let gcs_uri = format!("gs://{test_bucket_name}/so\\mek*ey/**testme.parquet");
937+
938+
let copy_to_command = format!(
939+
"COPY (SELECT a FROM generate_series(1,10) a) TO '{gcs_uri}' WITH (format parquet);"
940+
);
941+
Spi::run(&copy_to_command).unwrap();
942+
943+
let create_table = "create table test_table(a int);";
944+
Spi::run(create_table).unwrap();
945+
946+
let copy_from_command = format!("COPY test_table FROM '{gcs_uri}' WITH (format parquet);");
947+
Spi::run(copy_from_command.as_str()).unwrap();
948+
949+
let count_query = "select count(*) from test_table;";
950+
let result = Spi::get_one::<i64>(count_query).unwrap().unwrap();
951+
assert_eq!(result, 10);
952+
}
953+
954+
#[pg_test]
955+
#[should_panic(expected = "no files found that match the pattern")]
956+
fn test_gcs_with_nonexistent_pattern_uri() {
957+
object_store_cache_clear();
958+
959+
let test_bucket_name: String =
960+
std::env::var("GOOGLE_TEST_BUCKET").expect("GOOGLE_TEST_BUCKET not found");
961+
962+
let gcs_uri_pattern = format!("gs://{test_bucket_name}/notexists/**");
963+
964+
let create_table = "create table test_table(a int);";
965+
Spi::run(create_table).unwrap();
966+
967+
let copy_from_command =
968+
format!("COPY test_table FROM '{gcs_uri_pattern}' WITH (format parquet);");
969+
Spi::run(copy_from_command.as_str()).unwrap();
970+
}
971+
842972
#[pg_test]
843973
#[should_panic(expected = "404 Not Found")]
844974
fn test_gcs_write_wrong_bucket() {

‎src/pgrx_tests/udfs.rs‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1082,4 +1082,16 @@ mod tests {
10821082
vec![]
10831083
);
10841084
}
1085+
1086+
#[pg_test]
1087+
#[should_panic(expected = "list operation on http(s) object stores is not supported")]
1088+
fn test_http_list_not_allowed() {
1089+
let http_endpoint: String =
1090+
std::env::var("HTTP_ENDPOINT").expect("HTTP_ENDPOINT not found");
1091+
1092+
let http_uri_pattern = format!("{http_endpoint}/**");
1093+
1094+
let parquet_list_command = format!("select * from parquet.list('{http_uri_pattern}');");
1095+
Spi::run(&parquet_list_command).unwrap();
1096+
}
10851097
}

0 commit comments

Comments
 (0)