feat(#540): add dataset-size to response

The dataset-response now also contains the number of rows
and columns of the dataset.

Signed-off-by: Tobias Anker <tobias.anker@kitsunemimi.moe>

Author: kitsudaiki

.github/workflows/build_test.yml

@@ -619,6 +619,7 @@ jobs:
           mkdir -p test_dir
           cp sdk_api_test.py test_dir/
           cp test_values.py test_dir/
+          cp csv_test.csv test_dir/
           cd test_dir
           /usr/bin/python3 sdk_api_test.py
 

src/hanami/src/api/http_endpoints/dataset/create_dataset_v1_0.rs

@@ -30,6 +30,7 @@ use crate::database::dataset_table;
 use crate::config;
 
 use hanami_dataset::converter::{load_mnist_images, load_csv_file};
+use hanami_dataset::dataset_io::read_data_set_file;
 use hanami_common::error::HanamiError;
 
 use super::dataset_structs::DatasetResp;
@@ -231,10 +232,19 @@ pub async fn upload_binary(mut payload: Multipart, path: Path<(String, String)>,
         }
     }
 
+    let file_handle = match read_data_set_file(&target_filepath) {
+        Ok(file_handle) => file_handle,
+        Err(_) => {
+            return Err(ErrorResponse::InternalError("".to_string()));
+        }
+    };
+
     // create response
     let resp = DatasetResp {
         uuid: dataset_uuid.clone(),
         name: dataset.name.clone(),
+        number_of_rows: file_handle.get_number_of_rows(),
+        number_of_columns: file_handle.header.columns.len() as u64,
         created_by: dataset.created_by.clone(),
         created_at: dataset.created_at.clone(),
         updated_by: dataset.updated_by.clone(),

src/hanami/src/api/http_endpoints/dataset/dataset_structs.rs

@@ -27,6 +27,8 @@ pub struct DatasetCreateReq {
 pub struct DatasetResp {
     pub uuid: Uuid,
     pub name: String,
+    pub number_of_rows: u64,
+    pub number_of_columns: u64,
     pub created_at: String,
     pub created_by: String,
     pub updated_at: String,

src/hanami/src/api/http_endpoints/dataset/get_dataset_v1_0.rs

@@ -16,11 +16,14 @@ use actix_web::web::Json;
 use actix_web::web::Path;
 use apistos::api_operation;
 use uuid::Uuid;
+use std::path::PathBuf;
 
 use crate::api::errors::ErrorResponse;
 use crate::api::user_context::UserContext;
 use crate::database::dataset_table;
+
 use hanami_common::enums;
+use hanami_dataset::dataset_io::read_data_set_file;
 
 use super::dataset_structs::DatasetResp;
 
@@ -34,19 +37,8 @@ use super::dataset_structs::DatasetResp;
     error_code = 500
 )]
 pub async fn get_dataset(dataset_uuid: Path<Uuid>, context: UserContext) -> Result<Json<DatasetResp>, ErrorResponse> {
-    match dataset_table::get_dataset(&dataset_uuid, &context) {
-        Ok(dataset) => {
-            let resp = DatasetResp {
-                uuid: dataset_uuid.clone(),
-                name: dataset.name.clone(),
-                created_by: dataset.created_by.clone(),
-                created_at: dataset.created_at.clone(),
-                updated_by: dataset.updated_by.clone(),
-                updated_at: dataset.updated_at.clone(),
-            };
-        
-            return Ok(Json(resp));
-        },
+    let dataset_data = match dataset_table::get_dataset(&dataset_uuid, &context) {
+        Ok(dataset_data) => dataset_data,
         Err(enums::DbError::InternalError) => {
             return Err(ErrorResponse::InternalError("".to_string()));
         },
@@ -55,4 +47,24 @@ pub async fn get_dataset(dataset_uuid: Path<Uuid>, context: UserContext) -> Resu
             return Err(ErrorResponse::NotFound(msg));
         }
     };
+
+    let file_handle = match read_data_set_file(&PathBuf::from(dataset_data.file_path)) {
+        Ok(file_handle) => file_handle,
+        Err(_) => {
+            return Err(ErrorResponse::InternalError("".to_string()));
+        }
+    };
+
+    let resp = DatasetResp {
+        uuid: dataset_uuid.clone(),
+        name: dataset_data.name.clone(),
+        number_of_rows: file_handle.get_number_of_rows(),
+        number_of_columns: file_handle.header.columns.len() as u64,
+        created_by: dataset_data.created_by.clone(),
+        created_at: dataset_data.created_at.clone(),
+        updated_by: dataset_data.updated_by.clone(),
+        updated_at: dataset_data.updated_at.clone(),
+    };
+
+    return Ok(Json(resp));
 }

testing/python_sdk_api/sdk_api_test.py

@@ -137,21 +137,29 @@ def test_dataset():
 
     result = dataset.upload_mnist_files(
         token, address, train_dataset_name, train_inputs, train_labels, False)
-    dataset_uuid = result["uuid"]
+    mnist_dataset_uuid = result["uuid"]
 
     dataset.list_datasets(token, address, False)
-    dataset.get_dataset(token, address, dataset_uuid, False)
+    mnist_dataset = dataset.get_dataset(token, address, mnist_dataset_uuid, False)
+    assert mnist_dataset["number_of_rows"] == 60000
+    assert mnist_dataset["number_of_columns"] == 2
 
-    dataset.upload_csv_files(
+    result = dataset.upload_csv_files(
         token, address, "csv_test", "./csv_test.csv", False)
+    csv_dataset_uuid = result["uuid"]
+
+    csv_dataset = dataset.get_dataset(token, address, csv_dataset_uuid, False)
+    assert csv_dataset["number_of_rows"] == 3
+    assert csv_dataset["number_of_columns"] == 3
 
     try:
         dataset.get_dataset(token, address, " 569003fd-bf24-410b-8678-28f141877ac9", False)
     except hanami_exceptions.NotFoundException:
         pass
-    dataset.delete_dataset(token, address, dataset_uuid, False)
+    dataset.delete_dataset(token, address, mnist_dataset_uuid, False)
+    dataset.delete_dataset(token, address, csv_dataset_uuid, False)
     try:
-        dataset.delete_dataset(token, address, dataset_uuid, False)
+        dataset.delete_dataset(token, address, mnist_dataset_uuid, False)
     except hanami_exceptions.NotFoundException:
         pass