Fix write image (#9)

* Update changelog

* Add test

* Add test table

* Fix test

* Fix write image

* Improve test

* Update changelog

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)
 

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::<Vec<_>>();
-
+        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.")?;
         }
     }
 

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;
+    '/docker-entrypoint-initdb.d/contacts.csv' DELIMITER ',' CSV HEADER;
+
+INSERT INTO images (image) VALUES (pg_read_binary_file('/docker-entrypoint-initdb.d/test.png'));

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);
+}