feat(query): support direct the spill data to a different bucket (#17159)

* add storage.spill bucket

* fix

* fix

* fix

* remove strict

* vacuum_temporary_files

* SpillConfig.storage

* test

* fix

* fix

* test

* update

* fix

* fix

* update test

---------

Co-authored-by: sundyli <[email protected]>

Author: forsaken628

.github/actions/test_sqllogic_standalone_linux_minio/action.yml

@@ -45,6 +45,7 @@ runs:
         export AWS_EC2_METADATA_DISABLED=true
 
         aws --endpoint-url http://127.0.0.1:9900/ s3 mb s3://testbucket
+        aws --endpoint-url http://127.0.0.1:9900/ s3 mb s3://spillbucket
 
     - name: Run sqllogic Tests with Standalone mode
       if: inputs.storage-format == 'all' || inputs.storage-format == 'parquet'

scripts/ci/ci-run-sqllogic-tests-minio.sh

@@ -18,6 +18,20 @@ export STORAGE_S3_ACCESS_KEY_ID=minioadmin
 export STORAGE_S3_SECRET_ACCESS_KEY=minioadmin
 export STORAGE_ALLOW_INSECURE=true
 
+export SPILL_SPILL_LOCAL_DISK_PATH=''
+config="[spill.storage]
+type = \"s3\"
+
+[spill.storage.s3]
+bucket = \"spillbucket\"
+root = \"admin\"
+endpoint_url = \"http://127.0.0.1:9900\"
+access_key_id = \"minioadmin\"
+secret_access_key = \"minioadmin\"
+allow_insecure = true"
+
+echo "$config" >> ./scripts/ci/deploy/config/databend-query-node-1.toml
+
 echo "Starting standalone DatabendQuery and DatabendMeta"
 ./scripts/ci/deploy/databend-query-standalone.sh
 

src/common/storage/src/lib.rs

@@ -40,6 +40,7 @@ pub use http_client::StorageHttpClient;
 
 mod operator;
 pub use operator::build_operator;
+pub use operator::check_operator;
 pub use operator::init_operator;
 pub use operator::DataOperator;
 

src/common/storage/src/operator.rs

@@ -407,7 +407,9 @@ impl RetryInterceptor for DatabendRetryInterceptor {
 #[derive(Clone, Debug)]
 pub struct DataOperator {
     operator: Operator,
+    spill_operator: Option<Operator>,
     params: StorageParams,
+    spill_params: Option<StorageParams>,
 }
 
 impl DataOperator {
@@ -416,60 +418,69 @@ impl DataOperator {
         self.operator.clone()
     }
 
+    pub fn spill_operator(&self) -> Operator {
+        match &self.spill_operator {
+            Some(op) => op.clone(),
+            None => self.operator.clone(),
+        }
+    }
+
+    pub fn spill_params(&self) -> Option<&StorageParams> {
+        self.spill_params.as_ref()
+    }
+
     pub fn params(&self) -> StorageParams {
         self.params.clone()
     }
 
     #[async_backtrace::framed]
-    pub async fn init(conf: &StorageConfig) -> databend_common_exception::Result<()> {
-        GlobalInstance::set(Self::try_create(&conf.params).await?);
+    pub async fn init(
+        conf: &StorageConfig,
+        spill_params: Option<StorageParams>,
+    ) -> databend_common_exception::Result<()> {
+        GlobalInstance::set(Self::try_create(conf, spill_params).await?);
 
         Ok(())
     }
 
     /// Create a new data operator without check.
-    pub fn try_new(sp: &StorageParams) -> databend_common_exception::Result<DataOperator> {
-        let operator = init_operator(sp)?;
+    pub fn try_new(
+        conf: &StorageConfig,
+        spill_params: Option<StorageParams>,
+    ) -> databend_common_exception::Result<DataOperator> {
+        let operator = init_operator(&conf.params)?;
+        let spill_operator = spill_params.as_ref().map(init_operator).transpose()?;
 
         Ok(DataOperator {
             operator,
-            params: sp.clone(),
+            params: conf.params.clone(),
+            spill_operator,
+            spill_params,
         })
     }
 
     #[async_backtrace::framed]
-    pub async fn try_create(sp: &StorageParams) -> databend_common_exception::Result<DataOperator> {
-        let sp = sp.clone();
-
-        let operator = init_operator(&sp)?;
-
-        // OpenDAL will send a real request to underlying storage to check whether it works or not.
-        // If this check failed, it's highly possible that the users have configured it wrongly.
-        //
-        // Make sure the check is called inside GlobalIORuntime to prevent
-        // IO hang on reuse connection.
-        let op = operator.clone();
-        if let Err(cause) = GlobalIORuntime::instance()
-            .spawn(async move {
-                let res = op.stat("/").await;
-                match res {
-                    Ok(_) => Ok(()),
-                    Err(e) if e.kind() == opendal::ErrorKind::NotFound => Ok(()),
-                    Err(e) => Err(e),
-                }
-            })
-            .await
-            .expect("join must succeed")
-        {
-            return Err(ErrorCode::StorageUnavailable(format!(
-                "current configured storage is not available: config: {:?}, cause: {cause}",
-                sp
-            )));
-        }
+    pub async fn try_create(
+        conf: &StorageConfig,
+        spill_params: Option<StorageParams>,
+    ) -> databend_common_exception::Result<DataOperator> {
+        let operator = init_operator(&conf.params)?;
+        check_operator(&operator, &conf.params).await?;
+
+        let spill_operator = match &spill_params {
+            Some(params) => {
+                let op = init_operator(params)?;
+                check_operator(&op, params).await?;
+                Some(op)
+            }
+            None => None,
+        };
 
         Ok(DataOperator {
             operator,
-            params: sp.clone(),
+            params: conf.params.clone(),
+            spill_operator,
+            spill_params,
         })
     }
 
@@ -485,3 +496,33 @@ impl DataOperator {
         GlobalInstance::get()
     }
 }
+
+pub async fn check_operator(
+    operator: &Operator,
+    params: &StorageParams,
+) -> databend_common_exception::Result<()> {
+    // OpenDAL will send a real request to underlying storage to check whether it works or not.
+    // If this check failed, it's highly possible that the users have configured it wrongly.
+    //
+    // Make sure the check is called inside GlobalIORuntime to prevent
+    // IO hang on reuse connection.
+    let op = operator.clone();
+
+    GlobalIORuntime::instance()
+        .spawn(async move {
+            let res = op.stat("/").await;
+            match res {
+                Ok(_) => Ok(()),
+                Err(e) if e.kind() == opendal::ErrorKind::NotFound => Ok(()),
+                Err(e) => Err(e),
+            }
+        })
+        .await
+        .expect("join must succeed")
+        .map_err(|cause| {
+            ErrorCode::StorageUnavailable(format!(
+                "current configured storage is not available: config: {:?}, cause: {cause}",
+                params
+            ))
+        })
+}

src/query/config/src/config.rs

@@ -2972,6 +2972,10 @@ pub struct SpillConfig {
     #[clap(long, value_name = "VALUE", default_value = "18446744073709551615")]
     /// Allow space in bytes to spill to local disk.
     pub spill_local_disk_max_bytes: u64,
+
+    // TODO: We need to fix StorageConfig so that it supports environment variables and command line injections.
+    #[clap(skip)]
+    pub storage: Option<StorageConfig>,
 }
 
 impl Default for SpillConfig {
@@ -2981,8 +2985,6 @@ impl Default for SpillConfig {
 }
 
 mod cache_config_converters {
-    use std::path::PathBuf;
-
     use log::warn;
 
     use super::*;
@@ -3024,7 +3026,7 @@ mod cache_config_converters {
                 storage,
                 catalog,
                 cache,
-                mut spill,
+                spill,
                 background,
                 catalogs: input_catalogs,
                 ..
@@ -3044,18 +3046,7 @@ mod cache_config_converters {
                 catalogs.insert(CATALOG_HIVE.to_string(), catalog);
             }
 
-            // Trick for cloud, perhaps we should introduce a new configuration for the local writeable root.
-            if cache.disk_cache_config.path != inner::DiskCacheConfig::default().path
-                && spill.spill_local_disk_path == inner::SpillConfig::default().path
-            {
-                spill.spill_local_disk_path = PathBuf::from(&cache.disk_cache_config.path)
-                    .join("temp/_query_spill")
-                    .into_os_string()
-                    .into_string()
-                    .map_err(|s| {
-                        ErrorCode::Internal(format!("failed to convert os string to string: {s:?}"))
-                    })?
-            };
+            let spill = convert_local_spill_config(spill, &cache.disk_cache_config)?;
 
             Ok(InnerConfig {
                 subcommand,
@@ -3066,7 +3057,7 @@ mod cache_config_converters {
                 storage: storage.try_into()?,
                 catalogs,
                 cache: cache.try_into()?,
-                spill: spill.try_into()?,
+                spill,
                 background: background.try_into()?,
             })
         }
@@ -3129,37 +3120,51 @@ mod cache_config_converters {
         }
     }
 
-    impl TryFrom<SpillConfig> for inner::SpillConfig {
-        type Error = ErrorCode;
+    fn convert_local_spill_config(
+        spill: SpillConfig,
+        cache: &DiskCacheConfig,
+    ) -> Result<inner::SpillConfig> {
+        // Trick for cloud, perhaps we should introduce a new configuration for the local writeable root.
+        let local_writeable_root = if cache.path != DiskCacheConfig::default().path
+            && spill.spill_local_disk_path.is_empty()
+        {
+            Some(cache.path.clone())
+        } else {
+            None
+        };
 
-        fn try_from(value: SpillConfig) -> std::result::Result<Self, Self::Error> {
-            let SpillConfig {
-                spill_local_disk_path,
-                spill_local_disk_reserved_space_percentage: reserved,
-                spill_local_disk_max_bytes,
-            } = value;
-            if !reserved.is_normal()
-                || reserved.is_sign_negative()
-                || reserved > OrderedFloat(100.0)
-            {
-                Err(ErrorCode::InvalidArgument(format!(
-                    "invalid spill_local_disk_reserved_space_percentage: {reserved}"
-                )))?;
-            }
-            Ok(Self {
-                path: spill_local_disk_path,
-                reserved_disk_ratio: reserved / 100.0,
-                global_bytes_limit: spill_local_disk_max_bytes,
+        let storage_params = spill
+            .storage
+            .map(|storage| {
+                let storage: InnerStorageConfig = storage.try_into()?;
+                Ok::<_, ErrorCode>(storage.params)
             })
-        }
+            .transpose()?;
+
+        Ok(inner::SpillConfig {
+            local_writeable_root,
+            path: spill.spill_local_disk_path,
+            reserved_disk_ratio: spill.spill_local_disk_reserved_space_percentage / 100.0,
+            global_bytes_limit: spill.spill_local_disk_max_bytes,
+            storage_params,
+        })
     }
 
     impl From<inner::SpillConfig> for SpillConfig {
         fn from(value: inner::SpillConfig) -> Self {
+            let storage = value.storage_params.map(|params| {
+                InnerStorageConfig {
+                    params,
+                    ..Default::default()
+                }
+                .into()
+            });
+
             Self {
                 spill_local_disk_path: value.path,
                 spill_local_disk_reserved_space_percentage: value.reserved_disk_ratio * 100.0,
                 spill_local_disk_max_bytes: value.global_bytes_limit,
+                storage,
             }
         }
     }

src/query/config/src/inner.rs

@@ -17,6 +17,7 @@ use std::fmt;
 use std::fmt::Debug;
 use std::fmt::Display;
 use std::fmt::Formatter;
+use std::path::PathBuf;
 use std::str::FromStr;
 use std::time::Duration;
 
@@ -28,6 +29,7 @@ use databend_common_exception::Result;
 use databend_common_grpc::RpcClientConf;
 use databend_common_grpc::RpcClientTlsConfig;
 use databend_common_meta_app::principal::UserSettingValue;
+use databend_common_meta_app::storage::StorageParams;
 use databend_common_meta_app::tenant::Tenant;
 use databend_common_meta_app::tenant::TenantQuota;
 use databend_common_storage::StorageConfig;
@@ -717,22 +719,55 @@ impl Default for CacheConfig {
 
 #[derive(Clone, Debug, PartialEq, Eq)]
 pub struct SpillConfig {
-    /// Path of spill to local disk. disable if it's empty.
-    pub path: String,
+    pub(crate) local_writeable_root: Option<String>,
+    pub(crate) path: String,
 
     /// Ratio of the reserve of the disk space.
     pub reserved_disk_ratio: OrderedFloat<f64>,
 
     /// Allow bytes use of disk space.
     pub global_bytes_limit: u64,
+
+    pub storage_params: Option<StorageParams>,
+}
+
+impl SpillConfig {
+    /// Path of spill to local disk.
+    pub fn local_path(&self) -> Option<PathBuf> {
+        if self.global_bytes_limit == 0 {
+            return None;
+        }
+
+        if !self.path.is_empty() {
+            return Some(self.path.clone().into());
+        }
+
+        if let Some(root) = &self.local_writeable_root {
+            return Some(PathBuf::from(root).join("temp/_query_spill"));
+        }
+
+        None
+    }
+
+    pub fn new_for_test(path: String, reserved_disk_ratio: f64, global_bytes_limit: u64) -> Self {
+        Self {
+            local_writeable_root: None,
+            path,
+            reserved_disk_ratio: OrderedFloat(reserved_disk_ratio),
+            global_bytes_limit,
+            storage_params: None,
+        }
+    }
 }
 
 impl Default for SpillConfig {
     fn default() -> Self {
         Self {
+            local_writeable_root: None,
             path: "".to_string(),
             reserved_disk_ratio: OrderedFloat(0.3),
             global_bytes_limit: u64::MAX,
+            storage_params: None,
         }
     }
 }