diff --git a/CHANGELOG.md b/CHANGELOG.md index c08b7cf..7f798cd 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,8 @@ - changed - Refactor integration tests - Update to rust 1.77 + - Update dependencies + - Fix `sources::write_image` (invalid series dtype: expected `List`, got `binary`) ### v0.4.0 (2022-12-24) diff --git a/src/sources/image.rs b/src/sources/image.rs index 834d261..9fd0026 100644 --- a/src/sources/image.rs +++ b/src/sources/image.rs @@ -51,23 +51,15 @@ pub fn write_image( let image = df .column(image_column) .context("Can't find column for images")? - .list() + .binary() .context("Can't convert series to chunked array")? .get(i); println!("Save query result to file: {}", target_path.display()); - if let Some(image) = image { - let bytes = image - .u8() - .context("Can't convert series to chunked array")? - .into_iter() - .map(|byte| byte.expect("Can't convert series to bytes")) - .collect::>(); - + if let Some(bytes) = image { let mut file = File::create(target_path).context("Unable to create file")?; - file.write_all(bytes.as_slice()) - .context("Unable to write file.")?; + file.write_all(bytes).context("Unable to write file.")?; } } diff --git a/test_data/init.sql b/test_data/init.sql index 5d146ed..531244b 100644 --- a/test_data/init.sql +++ b/test_data/init.sql @@ -5,6 +5,10 @@ CREATE TABLE account ( email character varying NOT NULL ); +CREATE TABLE images (id serial PRIMARY KEY, image bytea); + COPY account(first_name, last_name, email) FROM - '/docker-entrypoint-initdb.d/contacts.csv' DELIMITER ',' CSV HEADER; \ No newline at end of file + '/docker-entrypoint-initdb.d/contacts.csv' DELIMITER ',' CSV HEADER; + +INSERT INTO images (image) VALUES (pg_read_binary_file('/docker-entrypoint-initdb.d/test.png')); diff --git a/tests/cmd/test_query.rs b/tests/cmd/test_query.rs index dd99c18..149bab2 100644 --- a/tests/cmd/test_query.rs +++ b/tests/cmd/test_query.rs @@ -7,7 +7,7 @@ */ use assert_cmd::Command; -use predicates::str; +use predicates::{boolean::PredicateBooleanExt, str}; use std::fs; use tempfile::tempdir; @@ -32,7 +32,7 @@ fn test_query_display() { } #[test] -fn test_query_save() { +fn test_query_save_csv() { let test_query = "select email, first_name, last_name from account"; let temp_dir = tempdir().unwrap(); let temp_path = temp_dir.path(); @@ -61,3 +61,46 @@ fn test_query_save() { assert!(dir_entry.is_some()); } } + +#[test] +fn test_query_save_png() { + let test_query = "select id, image from images"; + let temp_dir = tempdir().unwrap(); + let temp_path = temp_dir.path(); + assert!(temp_path.exists(), "Missing path: {}", temp_path.display()); + let save_dir = temp_path.to_str().unwrap(); + let file_type = "png"; + let image_column = "image"; + let image_name = "id"; + + println!( + "Execute 'pigeon query {test_query} --save --save-dir {save_dir} --file-type {file_type} --image-column {image_column} --image-name {image_name}'" + ); + let mut cmd = Command::cargo_bin("pigeon").unwrap(); + cmd.args([ + "query", + test_query, + "--save", + "--save-dir", + save_dir, + "--file-type", + file_type, + "--image-column", + image_column, + "--image-name", + image_name, + ]); + cmd.assert() + .success() + .stdout( + str::contains("Display query result: shape: (1, 2)").and(str::contains(format!( + "Save query result to file: {}/1.png", + temp_path.display() + ))), + ); + + let image_path = temp_path.join("1.png"); + let actual = fs::read(image_path).unwrap(); + let expected = fs::read("./test_data/test.png").unwrap(); + assert_eq!(actual, expected); +}