chore(cubestore): Small refactoring of import process (#7189)

Author: waralexrom

rust/cubestore/cubestore/src/config/mod.rs

@@ -15,7 +15,7 @@ use crate::config::injection::{DIService, Injector};
 use crate::config::processing_loop::ProcessingLoop;
 use crate::http::HttpServer;
 use crate::import::limits::ConcurrencyLimits;
-use crate::import::{ImportService, ImportServiceImpl};
+use crate::import::{ImportService, ImportServiceImpl, LocationsValidator, LocationsValidatorImpl};
 use crate::metastore::{
     BaseRocksStoreFs, MetaStore, MetaStoreRpcClient, RocksMetaStore, RocksStoreConfig,
 };
@@ -1960,6 +1960,12 @@ impl Config {
             })
             .await;
 
+        self.injector
+            .register_typed::<dyn LocationsValidator, _, _, _>(async move |_| {
+                LocationsValidatorImpl::new()
+            })
+            .await;
+
         self.injector
             .register_typed::<dyn ImportService, _, _, _>(async move |i| {
                 ImportServiceImpl::new(
@@ -1969,6 +1975,7 @@ impl Config {
                     i.get_service_typed().await,
                     i.get_service_typed().await,
                     i.get_service_typed().await,
+                    i.get_service_typed().await,
                 )
             })
             .await;

rust/cubestore/cubestore/src/import/mod.rs

@@ -469,6 +469,28 @@ impl<R: AsyncBufRead> Stream for CsvLineStream<R> {
     }
 }
 
+#[async_trait]
+pub trait LocationsValidator: DIService + Send + Sync {
+    async fn validate(&self, locations: &Vec<String>) -> Result<(), CubeError>;
+}
+
+pub struct LocationsValidatorImpl;
+
+#[async_trait]
+impl LocationsValidator for LocationsValidatorImpl {
+    async fn validate(&self, _locations: &Vec<String>) -> Result<(), CubeError> {
+        Ok(())
+    }
+}
+
+impl LocationsValidatorImpl {
+    pub fn new() -> Arc<Self> {
+        Arc::new(Self {})
+    }
+}
+
+crate::di_service!(LocationsValidatorImpl, [LocationsValidator]);
+
 #[automock]
 #[async_trait]
 pub trait ImportService: DIService + Send + Sync {
@@ -482,6 +504,7 @@ pub trait ImportService: DIService + Send + Sync {
     async fn validate_table_location(&self, table_id: u64, location: &str)
         -> Result<(), CubeError>;
     async fn estimate_location_row_count(&self, location: &str) -> Result<u64, CubeError>;
+    async fn validate_locations_size(&self, locations: &Vec<String>) -> Result<(), CubeError>;
 }
 
 crate::di_service!(MockImportService, [ImportService]);
@@ -493,6 +516,7 @@ pub struct ImportServiceImpl {
     remote_fs: Arc<dyn RemoteFs>,
     config_obj: Arc<dyn ConfigObj>,
     limits: Arc<ConcurrencyLimits>,
+    validator: Arc<dyn LocationsValidator>,
 }
 
 crate::di_service!(ImportServiceImpl, [ImportService]);
@@ -505,6 +529,7 @@ impl ImportServiceImpl {
         remote_fs: Arc<dyn RemoteFs>,
         config_obj: Arc<dyn ConfigObj>,
         limits: Arc<ConcurrencyLimits>,
+        validator: Arc<dyn LocationsValidator>,
     ) -> Arc<ImportServiceImpl> {
         Arc::new(ImportServiceImpl {
             meta_store,
@@ -513,6 +538,7 @@ impl ImportServiceImpl {
             remote_fs,
             config_obj,
             limits,
+            validator,
         })
     }
 
@@ -618,23 +644,19 @@ impl ImportServiceImpl {
     }
 
     async fn download_temp_file(&self, location: &str) -> Result<File, CubeError> {
-        let to_download = ImportServiceImpl::temp_uploads_path(location);
+        let to_download = LocationHelper::temp_uploads_path(location);
         // TODO check file size
         let local_file = self.remote_fs.download_file(&to_download, None).await?;
         Ok(File::open(local_file.clone())
             .await
             .map_err(|e| CubeError::internal(format!("Open temp_file {}: {}", local_file, e)))?)
     }
 
-    fn temp_uploads_path(location: &str) -> String {
-        location.replace("temp://", "temp-uploads/")
-    }
-
     async fn drop_temp_uploads(&self, location: &str) -> Result<(), CubeError> {
         // TODO There also should be a process which collects orphaned uploads due to failed imports
         if location.starts_with("temp://") {
             self.remote_fs
-                .delete_file(&ImportServiceImpl::temp_uploads_path(location))
+                .delete_file(&LocationHelper::temp_uploads_path(location))
                 .await?;
         }
         Ok(())
@@ -810,28 +832,57 @@ impl ImportService for ImportServiceImpl {
     }
 
     async fn estimate_location_row_count(&self, location: &str) -> Result<u64, CubeError> {
-        if location.starts_with("http") {
+        let file_size =
+            LocationHelper::location_file_size(location, self.remote_fs.clone()).await?;
+        Ok(ImportServiceImpl::estimate_rows(location, file_size))
+    }
+
+    async fn validate_locations_size(&self, locations: &Vec<String>) -> Result<(), CubeError> {
+        self.validator.validate(locations).await
+    }
+}
+
+pub struct LocationHelper;
+
+impl LocationHelper {
+    pub async fn location_file_size(
+        location: &str,
+        remote_fs: Arc<dyn RemoteFs>,
+    ) -> Result<Option<u64>, CubeError> {
+        let res = if location.starts_with("http") {
             let client = reqwest::Client::new();
             let res = client.head(location).send().await?;
             let length = res.headers().get(reqwest::header::CONTENT_LENGTH);
 
-            let size = if let Some(length) = length {
+            if let Some(length) = length {
                 Some(length.to_str()?.parse::<u64>()?)
             } else {
                 None
-            };
-            Ok(ImportServiceImpl::estimate_rows(location, size))
+            }
         } else if location.starts_with("temp://") {
-            // TODO do the actual estimation
-            Ok(ImportServiceImpl::estimate_rows(location, None))
+            let remote_path = Self::temp_uploads_path(location);
+            match remote_fs.list_with_metadata(&remote_path).await {
+                Ok(list) => {
+                    let list_res = list.iter().next().ok_or(CubeError::internal(format!(
+                        "Location {} can't be listed in remote_fs",
+                        location
+                    )));
+                    match list_res {
+                        Ok(file) => Ok(Some(file.file_size)),
+                        Err(e) => Err(e),
+                    }
+                }
+                Err(e) => Err(e),
+            }?
         } else if location.starts_with("stream://") {
-            Ok(ImportServiceImpl::estimate_rows(location, None))
+            None
         } else {
-            Ok(ImportServiceImpl::estimate_rows(
-                location,
-                Some(tokio::fs::metadata(location).await?.len()),
-            ))
-        }
+            Some(tokio::fs::metadata(location).await?.len())
+        };
+        Ok(res)
+    }
+    pub fn temp_uploads_path(location: &str) -> String {
+        location.replace("temp://", "temp-uploads/")
     }
 }
 

rust/cubestore/cubestore/src/sql/mod.rs

@@ -363,6 +363,11 @@ impl SqlServiceImpl {
         }
 
         let listener = self.cluster.job_result_listener();
+        if let Some(locations) = locations.as_ref() {
+            self.import_service
+                .validate_locations_size(locations)
+                .await?;
+        }
 
         let partition_split_threshold = if let Some(locations) = locations.as_ref() {
             let size = join_all(