From 7ded7fecc23cf4e2a8d0c2cc236f055e4256fb68 Mon Sep 17 00:00:00 2001 From: zhang2014 Date: Sun, 15 Dec 2024 23:14:10 +0800 Subject: [PATCH 01/53] feat(cluster): support custom management cluster --- src/common/exception/src/exception_code.rs | 4 + src/query/management/src/lib.rs | 5 +- .../management/src/warehouse/cluster_mgr.rs | 229 ------- src/query/management/src/warehouse/mod.rs | 9 +- .../{cluster_api.rs => warehouse_api.rs} | 27 +- .../management/src/warehouse/warehouse_mgr.rs | 577 ++++++++++++++++++ src/query/management/tests/it/cluster.rs | 135 +++- src/query/management/tests/it/main.rs | 3 +- src/query/service/src/clusters/cluster.rs | 52 +- 9 files changed, 767 insertions(+), 274 deletions(-) delete mode 100644 src/query/management/src/warehouse/cluster_mgr.rs rename src/query/management/src/warehouse/{cluster_api.rs => warehouse_api.rs} (58%) create mode 100644 src/query/management/src/warehouse/warehouse_mgr.rs diff --git a/src/common/exception/src/exception_code.rs b/src/common/exception/src/exception_code.rs index ef4a30d220e6d..e5ac696f7f201 100644 --- a/src/common/exception/src/exception_code.rs +++ b/src/common/exception/src/exception_code.rs @@ -294,6 +294,10 @@ build_exceptions! { ClusterUnknownNode(2401), ClusterNodeAlreadyExists(2402), InvalidWarehouse(2403), + NoResourcesAvailable(2404), + WarehouseAlreadyExists(2405), + UnknownWarehouse(2406), + WarehouseOperateConflict(2407), // Stage error codes. UnknownStage(2501), diff --git a/src/query/management/src/lib.rs b/src/query/management/src/lib.rs index bed354004fd31..cc4a34a10bf99 100644 --- a/src/query/management/src/lib.rs +++ b/src/query/management/src/lib.rs @@ -50,5 +50,6 @@ pub use stage::StageApi; pub use stage::StageMgr; pub use user::UserApi; pub use user::UserMgr; -pub use warehouse::ClusterApi; -pub use warehouse::ClusterMgr; +pub use warehouse::SelectedNode; +pub use warehouse::WarehouseApi; +pub use warehouse::WarehouseMgr; diff --git a/src/query/management/src/warehouse/cluster_mgr.rs b/src/query/management/src/warehouse/cluster_mgr.rs deleted file mode 100644 index 8c9ce9e57c694..0000000000000 --- a/src/query/management/src/warehouse/cluster_mgr.rs +++ /dev/null @@ -1,229 +0,0 @@ -// Copyright 2021 Datafuse Labs -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -use std::time::Duration; - -use databend_common_base::base::escape_for_key; -use databend_common_base::base::unescape_for_key; -use databend_common_exception::ErrorCode; -use databend_common_exception::Result; -use databend_common_meta_kvapi::kvapi::KVApi; -use databend_common_meta_store::MetaStore; -use databend_common_meta_types::protobuf::SeqV; -use databend_common_meta_types::txn_op_response::Response; -use databend_common_meta_types::ConditionResult; -use databend_common_meta_types::MatchSeq; -use databend_common_meta_types::NodeInfo; -use databend_common_meta_types::TxnCondition; -use databend_common_meta_types::TxnOp; -use databend_common_meta_types::TxnOpResponse; -use databend_common_meta_types::TxnReply; -use databend_common_meta_types::TxnRequest; - -use crate::warehouse::ClusterApi; - -pub static CLUSTER_API_KEY_PREFIX: &str = "__fd_clusters_v5"; - -pub struct ClusterMgr { - metastore: MetaStore, - lift_time: Duration, - node_key_prefix: String, - meta_key_prefix: String, -} - -impl ClusterMgr { - pub fn create(metastore: MetaStore, tenant: &str, lift_time: Duration) -> Result { - if tenant.is_empty() { - return Err(ErrorCode::TenantIsEmpty( - "Tenant can not empty(while cluster mgr create)", - )); - } - - Ok(ClusterMgr { - metastore, - lift_time, - // all online node of tenant - node_key_prefix: format!( - "{}/{}/online_nodes", - CLUSTER_API_KEY_PREFIX, - escape_for_key(tenant)? - ), - // all computing cluster of tenant - meta_key_prefix: format!( - "{}/{}/online_clusters", - CLUSTER_API_KEY_PREFIX, - escape_for_key(tenant)? - ), - }) - } -} - -fn map_condition(k: &str, seq: MatchSeq) -> TxnCondition { - match seq { - MatchSeq::Any => TxnCondition::match_seq(k.to_owned(), ConditionResult::Ge, 0), - MatchSeq::GE(v) => TxnCondition::match_seq(k.to_owned(), ConditionResult::Ge, v), - MatchSeq::Exact(v) => TxnCondition::match_seq(k.to_owned(), ConditionResult::Eq, v), - } -} - -fn map_response(res: Option<&TxnOpResponse>) -> Option<&SeqV> { - res.and_then(|response| response.response.as_ref()) - .and_then(|response| match response { - Response::Put(v) => v.prev_value.as_ref(), - Response::Delete(v) => v.prev_value.as_ref(), - _ => unreachable!(), - }) -} - -impl ClusterMgr { - #[fastrace::trace] - #[async_backtrace::framed] - async fn upsert_node(&self, mut node: NodeInfo, seq: MatchSeq) -> Result { - let mut txn = TxnRequest::default(); - - let node_key = format!("{}/{}", self.node_key_prefix, escape_for_key(&node.id)?); - txn.if_then.push(TxnOp::put_with_ttl( - node_key.clone(), - serde_json::to_vec(&node)?, - Some(self.lift_time), - )); - - if !node.cluster_id.is_empty() && !node.warehouse_id.is_empty() { - let cluster_key = format!( - "{}/{}/{}/{}", - self.meta_key_prefix, - escape_for_key(&node.warehouse_id)?, - escape_for_key(&node.cluster_id)?, - escape_for_key(&node.id)? - ); - - txn.condition.push(map_condition(&cluster_key, seq)); - node.cluster_id = String::new(); - node.warehouse_id = String::new(); - txn.if_then.push(TxnOp::put_with_ttl( - cluster_key, - serde_json::to_vec(&node)?, - Some(self.lift_time), - )); - } - - Ok(self.metastore.transaction(txn).await?) - } -} - -#[async_trait::async_trait] -impl ClusterApi for ClusterMgr { - #[async_backtrace::framed] - #[fastrace::trace] - async fn add_node(&self, node: NodeInfo) -> Result<()> { - assert!(!node.cluster_id.is_empty()); - assert!(!node.warehouse_id.is_empty()); - - let res = self.upsert_node(node.clone(), MatchSeq::Exact(0)).await?; - - if res.success && map_response(res.responses.get(1)).is_none() { - return Ok(()); - } - - Err(ErrorCode::ClusterNodeAlreadyExists(format!( - "Node with ID '{}' already exists in the cluster.", - node.id - ))) - } - - async fn heartbeat(&self, node: &NodeInfo) -> Result<()> { - assert!(!node.cluster_id.is_empty()); - assert!(!node.warehouse_id.is_empty()); - - // Update or insert the node with GE(0). - let transition = self.upsert_node(node.clone(), MatchSeq::GE(0)).await?; - - match transition.success { - true => Ok(()), - false => Err(ErrorCode::MetaServiceError(format!( - "Unexpected None result returned when upsert heartbeat node {}", - node.id - ))), - } - } - - #[async_backtrace::framed] - #[fastrace::trace] - async fn get_nodes(&self, warehouse: &str, cluster: &str) -> Result> { - let cluster_key = format!( - "{}/{}/{}", - self.meta_key_prefix, - escape_for_key(warehouse)?, - escape_for_key(cluster)? - ); - - let values = self.metastore.prefix_list_kv(&cluster_key).await?; - - let mut nodes_info = Vec::with_capacity(values.len()); - for (node_key, value) in values { - let mut node_info = serde_json::from_slice::(&value.data)?; - - node_info.id = unescape_for_key(&node_key[cluster_key.len() + 1..])?; - node_info.cluster_id = cluster.to_string(); - node_info.warehouse_id = warehouse.to_string(); - nodes_info.push(node_info); - } - - Ok(nodes_info) - } - - #[async_backtrace::framed] - #[fastrace::trace] - async fn drop_node(&self, node_id: String) -> Result<()> { - let node_key = format!("{}/{}", self.node_key_prefix, escape_for_key(&node_id)?); - - if let Some(info) = self.metastore.get_kv(&node_key).await? { - let node_info: NodeInfo = serde_json::from_slice(&info.data)?; - - let mut txn = TxnRequest::default(); - - txn.if_then.push(TxnOp::delete(node_key)); - - if !node_info.cluster_id.is_empty() && !node_info.warehouse_id.is_empty() { - txn.if_then.push(TxnOp::delete(format!( - "{}/{}/{}/{}", - self.meta_key_prefix, - escape_for_key(&node_info.warehouse_id)?, - escape_for_key(&node_info.cluster_id)?, - escape_for_key(&node_info.id)? - ))); - } - - let res = self.metastore.transaction(txn).await?; - - if res.success - && map_response(res.responses.first()).is_some() - && map_response(res.responses.last()).is_some() - { - return Ok(()); - } - } - - Err(ErrorCode::ClusterUnknownNode(format!( - "Node with ID '{}' does not exist in the cluster.", - node_id - ))) - } - - #[async_backtrace::framed] - #[fastrace::trace] - async fn get_local_addr(&self) -> Result> { - Ok(self.metastore.get_local_addr().await?) - } -} diff --git a/src/query/management/src/warehouse/mod.rs b/src/query/management/src/warehouse/mod.rs index b7a58b2e551b8..206f779470b1c 100644 --- a/src/query/management/src/warehouse/mod.rs +++ b/src/query/management/src/warehouse/mod.rs @@ -12,8 +12,9 @@ // See the License for the specific language governing permissions and // limitations under the License. -mod cluster_api; -mod cluster_mgr; +mod warehouse_api; +mod warehouse_mgr; -pub use cluster_api::ClusterApi; -pub use cluster_mgr::ClusterMgr; +pub use warehouse_api::SelectedNode; +pub use warehouse_api::WarehouseApi; +pub use warehouse_mgr::WarehouseMgr; diff --git a/src/query/management/src/warehouse/cluster_api.rs b/src/query/management/src/warehouse/warehouse_api.rs similarity index 58% rename from src/query/management/src/warehouse/cluster_api.rs rename to src/query/management/src/warehouse/warehouse_api.rs index c9e2c6594cd51..64df7fc73109c 100644 --- a/src/query/management/src/warehouse/cluster_api.rs +++ b/src/query/management/src/warehouse/warehouse_api.rs @@ -15,14 +15,33 @@ use databend_common_exception::Result; use databend_common_meta_types::NodeInfo; +#[derive(serde::Serialize, serde::Deserialize, Clone)] +pub enum SelectedNode { + Random(Option), +} + +#[derive(serde::Serialize, serde::Deserialize)] +pub struct WarehouseInfo { + pub id: String, + pub status: String, + pub display_name: String, + pub clusters: Vec>, +} + /// Databend-query cluster management API #[async_trait::async_trait] -pub trait ClusterApi: Sync + Send { +pub trait WarehouseApi: Sync + Send { /// Add a new node. - async fn add_node(&self, node: NodeInfo) -> Result<()>; + async fn add_node(&self, node: NodeInfo) -> Result; /// Keep the tenant's cluster node alive. - async fn heartbeat(&self, node: &NodeInfo) -> Result<()>; + async fn heartbeat(&self, node: &mut NodeInfo, seq: u64) -> Result; + + async fn drop_warehouse(&self, warehouse: String) -> Result<()>; + + async fn create_warehouse(&self, warehouse: String, nodes: Vec) -> Result<()>; + + // async fn list_warehouses(&self) -> Result>; /// Get the tenant's cluster all nodes. async fn get_nodes(&self, warehouse: &str, cluster: &str) -> Result>; @@ -31,4 +50,6 @@ pub trait ClusterApi: Sync + Send { async fn drop_node(&self, node_id: String) -> Result<()>; async fn get_local_addr(&self) -> Result>; + + async fn get_node_info(&self, node_id: &str) -> Result; } diff --git a/src/query/management/src/warehouse/warehouse_mgr.rs b/src/query/management/src/warehouse/warehouse_mgr.rs new file mode 100644 index 0000000000000..e50b1d08fe4c6 --- /dev/null +++ b/src/query/management/src/warehouse/warehouse_mgr.rs @@ -0,0 +1,577 @@ +// Copyright 2021 Datafuse Labs +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::collections::HashSet; +use std::time::Duration; + +use databend_common_base::base::escape_for_key; +use databend_common_base::base::unescape_for_key; +use databend_common_base::base::GlobalUniqName; +use databend_common_exception::ErrorCode; +use databend_common_exception::Result; +use databend_common_meta_kvapi::kvapi::KVApi; +use databend_common_meta_store::MetaStore; +use databend_common_meta_types::protobuf::SeqV; +use databend_common_meta_types::txn_op_response::Response; +use databend_common_meta_types::ConditionResult; +use databend_common_meta_types::MatchSeq; +use databend_common_meta_types::NodeInfo; +use databend_common_meta_types::TxnCondition; +use databend_common_meta_types::TxnGetResponse; +use databend_common_meta_types::TxnOp; +use databend_common_meta_types::TxnOpResponse; +use databend_common_meta_types::TxnReply; +use databend_common_meta_types::TxnRequest; + +use crate::warehouse::warehouse_api::SelectedNode; +use crate::warehouse::warehouse_api::WarehouseInfo; +use crate::warehouse::WarehouseApi; + +pub static WAREHOUSE_API_KEY_PREFIX: &str = "__fd_clusters_v5"; +pub static WAREHOUSE_META_KEY_PREFIX: &str = "__fd_warehouses"; + +pub struct WarehouseMgr { + metastore: MetaStore, + lift_time: Duration, + node_key_prefix: String, + meta_key_prefix: String, + warehouse_key_prefix: String, +} + +impl WarehouseMgr { + pub fn create(metastore: MetaStore, tenant: &str, lift_time: Duration) -> Result { + if tenant.is_empty() { + return Err(ErrorCode::TenantIsEmpty( + "Tenant can not empty(while cluster mgr create)", + )); + } + + Ok(WarehouseMgr { + metastore, + lift_time, + // warehouses: + // all online node of tenant + node_key_prefix: format!( + "{}/{}/online_nodes", + WAREHOUSE_API_KEY_PREFIX, + escape_for_key(tenant)? + ), + // all computing cluster of tenant + meta_key_prefix: format!( + "{}/{}/online_clusters", + WAREHOUSE_API_KEY_PREFIX, + escape_for_key(tenant)? + ), + // all warehouse of tenant(required compatible with all versions) + warehouse_key_prefix: format!( + "{}/v1/{}", + WAREHOUSE_META_KEY_PREFIX, + escape_for_key(tenant)?, + ), + }) + } +} + +fn map_condition(k: &str, seq: MatchSeq) -> TxnCondition { + match seq { + MatchSeq::Any => TxnCondition::match_seq(k.to_owned(), ConditionResult::Ge, 0), + MatchSeq::GE(v) => TxnCondition::match_seq(k.to_owned(), ConditionResult::Ge, v), + MatchSeq::Exact(v) => TxnCondition::match_seq(k.to_owned(), ConditionResult::Eq, v), + } +} + +fn get_prev_value(res: Option<&TxnOpResponse>) -> Option<&SeqV> { + res.and_then(|response| response.response.as_ref()) + .and_then(|response| match response { + Response::Put(v) => v.prev_value.as_ref(), + Response::Delete(v) => v.prev_value.as_ref(), + _ => unreachable!(), + }) +} + +impl WarehouseMgr { + #[fastrace::trace] + #[async_backtrace::framed] + async fn upsert_node(&self, mut node: NodeInfo, seq: MatchSeq) -> Result { + let mut txn = TxnRequest::default(); + + let node_key = format!("{}/{}", self.node_key_prefix, escape_for_key(&node.id)?); + + txn.condition.push(map_condition(&node_key, seq)); + + txn.if_then.push(TxnOp::put_with_ttl( + node_key.clone(), + serde_json::to_vec(&node)?, + Some(self.lift_time), + )); + + if !node.cluster_id.is_empty() && !node.warehouse_id.is_empty() { + let cluster_key = format!( + "{}/{}/{}/{}", + self.meta_key_prefix, + escape_for_key(&node.warehouse_id)?, + escape_for_key(&node.cluster_id)?, + escape_for_key(&node.id)? + ); + + // txn.condition.push(map_condition(&cluster_key, seq)); + node.cluster_id = String::new(); + node.warehouse_id = String::new(); + txn.if_then.push(TxnOp::put_with_ttl( + cluster_key, + serde_json::to_vec(&node)?, + Some(self.lift_time), + )); + } + + txn.if_then.push(TxnOp::get(node_key.clone())); + txn.else_then.push(TxnOp::get(node_key.clone())); + Ok(self.metastore.transaction(txn).await?) + } + + async fn leave_cluster(&self, node_info: &mut NodeInfo, seq: u64) -> Result { + let mut cluster_id = String::new(); + let mut warehouse_id = String::new(); + + std::mem::swap(&mut node_info.cluster_id, &mut cluster_id); + std::mem::swap(&mut node_info.warehouse_id, &mut warehouse_id); + + let upsert_node = self.upsert_node(node_info.clone(), MatchSeq::Exact(seq)); + match upsert_node.await { + Err(err) => { + // rollback + std::mem::swap(&mut node_info.cluster_id, &mut cluster_id); + std::mem::swap(&mut node_info.warehouse_id, &mut warehouse_id); + Err(err) + } + Ok(response) if !response.success => { + // rollback + std::mem::swap(&mut node_info.cluster_id, &mut cluster_id); + std::mem::swap(&mut node_info.warehouse_id, &mut warehouse_id); + Ok(seq) + } + Ok(response) => match response.responses.last() { + Some(TxnOpResponse { + response: Some(Response::Get(TxnGetResponse { value: Some(v), .. })), + }) => Ok(v.seq), + _ => Err(ErrorCode::MetaServiceError("Meta insert failure.")), + }, + } + } + + async fn join_cluster(&self, node_info: &mut NodeInfo, seq: u64) -> Result { + let upsert_node = self.upsert_node(node_info.clone(), MatchSeq::Exact(seq)); + + match upsert_node.await { + Err(err) => { + // rollback + // std::mem::swap(&mut node_info.cluster_id, &mut cluster_id); + // std::mem::swap(&mut node_info.warehouse_id, &mut warehouse_id); + Err(err) + } + Ok(response) if !response.success => { + // rollback + // std::mem::swap(&mut node_info.cluster_id, &mut cluster_id); + // std::mem::swap(&mut node_info.warehouse_id, &mut warehouse_id); + Ok(seq) + } + Ok(response) => match response.responses.last() { + Some(TxnOpResponse { + response: Some(Response::Get(TxnGetResponse { value: Some(v), .. })), + }) => Ok(v.seq), + _ => Err(ErrorCode::MetaServiceError("Meta insert failure.")), + }, + } + } + + async fn resolve_conflicts(&self, reply: TxnReply, node: &mut NodeInfo) -> Result { + match reply.responses.first() { + None => self.leave_cluster(node, 0).await, + Some(TxnOpResponse { + response: Some(Response::Get(res)), + }) => match &res.value { + None => self.leave_cluster(node, 0).await, + Some(value) => { + let node_info = serde_json::from_slice::(&value.data)?; + + // Removed this node from the cluster in other nodes + if !node.cluster_id.is_empty() + && !node.warehouse_id.is_empty() + && node_info.cluster_id.is_empty() + && node_info.warehouse_id.is_empty() + { + return self.leave_cluster(node, value.seq).await; + } + + // Added this node to the cluster in other nodes + node.cluster_id = node_info.cluster_id; + node.warehouse_id = node_info.warehouse_id; + self.join_cluster(node, value.seq).await + } + }, + _ => Err(ErrorCode::Internal("Miss type while in meta response")), + } + } +} + +#[async_trait::async_trait] +impl WarehouseApi for WarehouseMgr { + #[async_backtrace::framed] + #[fastrace::trace] + async fn add_node(&self, node: NodeInfo) -> Result { + let res = self.upsert_node(node.clone(), MatchSeq::Exact(0)).await?; + + if res.success { + let Some(Response::Get(response)) = + res.responses.last().and_then(|x| x.response.as_ref()) + else { + return Err(ErrorCode::Internal("Unknown get response")); + }; + + let Some(node_info) = &response.value else { + return Err(ErrorCode::MetaServiceError("Add node info failure.")); + }; + + return Ok(node_info.seq); + } + + Err(ErrorCode::ClusterNodeAlreadyExists(format!( + "Node with ID '{}' already exists in the cluster.", + node.id + ))) + } + + async fn heartbeat(&self, node: &mut NodeInfo, seq: u64) -> Result { + assert!(!node.cluster_id.is_empty()); + assert!(!node.warehouse_id.is_empty()); + + let res = self.upsert_node(node.clone(), MatchSeq::Exact(seq)).await?; + + match res.success { + true => { + let Some(Response::Get(response)) = + res.responses.last().and_then(|x| x.response.as_ref()) + else { + return Err(ErrorCode::Internal("Unknown get response")); + }; + + let Some(node_info) = &response.value else { + return Err(ErrorCode::MetaServiceError("Add node info failure.")); + }; + + Ok(node_info.seq) + } + false => self.resolve_conflicts(res, node).await, + } + } + + async fn drop_warehouse(&self, warehouse: String) -> Result<()> { + let drop_key = format!("{}/{}", self.meta_key_prefix, escape_for_key(&warehouse)?); + + loop { + let mut delete_txn = TxnRequest::default(); + let values = self.metastore.prefix_list_kv(&drop_key).await?; + + let mut txn = TxnRequest::default(); + for (node_key, _value) in values { + txn.if_then.push(TxnOp::get(format!( + "{}/{}", + self.node_key_prefix, + &node_key[drop_key.len() + 1..] + ))); + + delete_txn.if_then.push(TxnOp::delete(node_key)); + } + + let warehouse_key = format!( + "{}/{}", + self.warehouse_key_prefix, + escape_for_key(&warehouse)? + ); + + txn.if_then.push(TxnOp::get(warehouse_key.clone())); + let mut fetch_reply = self.metastore.transaction(txn).await?; + + if let Some(response) = fetch_reply.responses.pop() { + let Some(Response::Get(response)) = response.response else { + return Err(ErrorCode::UnknownWarehouse(format!( + "Unknown warehouse or self managed warehouse {:?}", + warehouse + ))); + }; + + if response.key != warehouse_key { + return Err(ErrorCode::UnknownWarehouse(format!( + "Unknown warehouse or self managed warehouse {:?}", + warehouse + ))); + } + + let Some(value) = &response.value else { + return Err(ErrorCode::UnknownWarehouse(format!( + "Unknown warehouse or self managed warehouse {:?}", + warehouse + ))); + }; + + if fetch_reply.responses.len() != delete_txn.if_then.len() { + // TODO: maybe auto retry? + return Err(ErrorCode::WarehouseOperateConflict("Missing node info in online nodes list. Perhaps some nodes offline while in drop warehouse. You can try it again.")); + } + + for response in fetch_reply.responses { + let Some(Response::Get(response)) = response.response else { + // TODO: maybe auto retry? + return Err(ErrorCode::WarehouseOperateConflict("Missing node info in online nodes list. Perhaps some nodes offline while in drop warehouse. You can try it again.")); + }; + + let Some(value) = &response.value else { + // TODO: maybe auto retry? + return Err(ErrorCode::WarehouseOperateConflict("Missing node info in online nodes list. Perhaps some nodes offline while in drop warehouse. You can try it again.")); + }; + + let mut node_info = serde_json::from_slice::(&value.data)?; + + if node_info.warehouse_id != warehouse { + return Err(ErrorCode::WarehouseOperateConflict( + "Warehouse operate is conflict.", + )); + } + + delete_txn + .condition + .push(map_condition(&response.key, MatchSeq::Exact(value.seq))); + node_info.cluster_id = String::new(); + node_info.warehouse_id = String::new(); + delete_txn.if_then.push(TxnOp::put_with_ttl( + response.key, + serde_json::to_vec(&node_info)?, + Some(self.lift_time * 4), + )); + } + + delete_txn + .condition + .push(map_condition(&warehouse_key, MatchSeq::Exact(value.seq))); + delete_txn.if_then.push(TxnOp::delete(warehouse_key)); + + if !self.metastore.transaction(delete_txn).await?.success { + // seq is changed, will retry + continue; + } + + return Ok(()); + } + + return Err(ErrorCode::UnknownWarehouse(format!( + "Unknown warehouse or self managed warehouse {:?}", + warehouse + ))); + } + } + + async fn create_warehouse(&self, warehouse: String, nodes: Vec) -> Result<()> { + assert!(nodes.iter().all(|x| matches!(x, SelectedNode::Random(_)))); + + loop { + let mut selected_nodes = Vec::with_capacity(nodes.len()); + let mut allocated_ndoes = HashSet::with_capacity(nodes.len()); + + // get online nodes + let online_nodes = self.metastore.prefix_list_kv(&self.node_key_prefix).await?; + + for select_node in &nodes { + match select_node { + SelectedNode::Random(Some(_)) => { + return Err(ErrorCode::Unimplemented( + "Custom instance types are not supported.", + )); + } + // select random node + SelectedNode::Random(None) => { + for (_, v) in &online_nodes { + let mut node_info = serde_json::from_slice::(&v.data)?; + + if node_info.warehouse_id.is_empty() + && node_info.cluster_id.is_empty() + && !allocated_ndoes.contains(&node_info.id) + { + node_info.cluster_id = warehouse.clone(); + node_info.warehouse_id = warehouse.clone(); + allocated_ndoes.insert(node_info.id.clone()); + selected_nodes.push((v.seq, node_info)); + break; + } + } + } + } + } + + if selected_nodes.len() != nodes.len() { + return Err(ErrorCode::NoResourcesAvailable( + "Failed to create warehouse, reason: no resources available", + )); + } + + let mut txn = TxnRequest::default(); + + for (seq, mut node) in selected_nodes { + let node_key = format!("{}/{}", self.node_key_prefix, escape_for_key(&node.id)?); + + txn.condition + .push(map_condition(&node_key, MatchSeq::Exact(seq))); + txn.if_then.push(TxnOp::put_with_ttl( + node_key, + serde_json::to_vec(&node)?, + Some(self.lift_time * 4), + )); + + let cluster_key = format!( + "{}/{}/{}/{}", + self.meta_key_prefix, + escape_for_key(&node.warehouse_id)?, + escape_for_key(&node.cluster_id)?, + escape_for_key(&node.id)? + ); + + node.cluster_id = String::new(); + node.warehouse_id = String::new(); + txn.condition + .push(map_condition(&cluster_key, MatchSeq::Exact(0))); + txn.if_then.push(TxnOp::put_with_ttl( + cluster_key, + serde_json::to_vec(&node)?, + Some(self.lift_time * 4), + )); + } + + let warehouse_key = format!( + "{}/{}", + self.warehouse_key_prefix, + escape_for_key(&warehouse)? + ); + txn.condition + .push(map_condition(&warehouse_key, MatchSeq::Exact(0))); + txn.if_then.push(TxnOp::put( + warehouse_key.clone(), + serde_json::to_vec(&WarehouseInfo { + id: GlobalUniqName::unique(), + status: "Running".to_string(), + display_name: warehouse.clone(), + clusters: vec![nodes.clone()], + })?, + )); + txn.else_then.push(TxnOp::get(warehouse_key)); + + return match self.metastore.transaction(txn).await? { + res if res.success => Ok(()), + res => match res.responses.last() { + Some(TxnOpResponse { + response: Some(Response::Get(res)), + }) => { + if matches!(&res.value, Some(v) if v.seq != 0) { + return Err(ErrorCode::WarehouseAlreadyExists( + "Warehouse already exists", + )); + } + + // retry + continue; + } + _ => Err(ErrorCode::MetaServiceError( + "Missing type for meta response", + )), + }, + }; + } + } + + #[async_backtrace::framed] + #[fastrace::trace] + async fn get_nodes(&self, warehouse: &str, cluster: &str) -> Result> { + let cluster_key = format!( + "{}/{}/{}", + self.meta_key_prefix, + escape_for_key(warehouse)?, + escape_for_key(cluster)? + ); + + let values = self.metastore.prefix_list_kv(&cluster_key).await?; + + let mut nodes_info = Vec::with_capacity(values.len()); + for (node_key, value) in values { + let mut node_info = serde_json::from_slice::(&value.data)?; + + node_info.id = unescape_for_key(&node_key[cluster_key.len() + 1..])?; + node_info.cluster_id = cluster.to_string(); + node_info.warehouse_id = warehouse.to_string(); + nodes_info.push(node_info); + } + + Ok(nodes_info) + } + + #[async_backtrace::framed] + #[fastrace::trace] + async fn drop_node(&self, node_id: String) -> Result<()> { + let node_key = format!("{}/{}", self.node_key_prefix, escape_for_key(&node_id)?); + + if let Some(info) = self.metastore.get_kv(&node_key).await? { + let node_info: NodeInfo = serde_json::from_slice(&info.data)?; + + let mut txn = TxnRequest::default(); + + txn.if_then.push(TxnOp::delete(node_key)); + + if !node_info.cluster_id.is_empty() && !node_info.warehouse_id.is_empty() { + txn.if_then.push(TxnOp::delete(format!( + "{}/{}/{}/{}", + self.meta_key_prefix, + escape_for_key(&node_info.warehouse_id)?, + escape_for_key(&node_info.cluster_id)?, + escape_for_key(&node_info.id)? + ))); + } + + let res = self.metastore.transaction(txn).await?; + + if res.success + && get_prev_value(res.responses.first()).is_some() + && get_prev_value(res.responses.last()).is_some() + { + return Ok(()); + } + } + + Err(ErrorCode::ClusterUnknownNode(format!( + "Node with ID '{}' does not exist in the cluster.", + node_id + ))) + } + + #[async_backtrace::framed] + #[fastrace::trace] + async fn get_local_addr(&self) -> Result> { + Ok(self.metastore.get_local_addr().await?) + } + + async fn get_node_info(&self, node_id: &str) -> Result { + let node_key = format!("{}/{}", self.node_key_prefix, escape_for_key(&node_id)?); + let node_info = self.metastore.get_kv(&node_key).await?; + match node_info { + None => Err(ErrorCode::NotFoundClusterNode("")), + Some(v) => Ok(serde_json::from_slice(&v.data)?), + } + } +} diff --git a/src/query/management/tests/it/cluster.rs b/src/query/management/tests/it/cluster.rs index fda88e46fd904..6341f0fc24638 100644 --- a/src/query/management/tests/it/cluster.rs +++ b/src/query/management/tests/it/cluster.rs @@ -16,6 +16,7 @@ use std::sync::Arc; use std::time::Duration; use databend_common_base::base::tokio; +use databend_common_base::base::GlobalUniqName; use databend_common_exception::Result; use databend_common_management::*; use databend_common_meta_embedded::MemMeta; @@ -27,7 +28,7 @@ use databend_common_meta_types::NodeInfo; #[tokio::test(flavor = "multi_thread", worker_threads = 1)] async fn test_successfully_add_node() -> Result<()> { let now_ms = SeqV::<()>::now_ms(); - let (kv_api, cluster_api) = new_cluster_api().await?; + let (kv_api, cluster_api) = new_cluster_api(Duration::from_secs(60)).await?; let mut node_info = create_test_node_info(); cluster_api.add_node(node_info.clone()).await?; @@ -66,7 +67,7 @@ async fn test_successfully_add_node() -> Result<()> { #[tokio::test(flavor = "multi_thread", worker_threads = 1)] async fn test_already_exists_add_node() -> Result<()> { - let (_, cluster_api) = new_cluster_api().await?; + let (_, cluster_api) = new_cluster_api(Duration::from_secs(60)).await?; let node_info = create_test_node_info(); cluster_api.add_node(node_info.clone()).await?; @@ -81,7 +82,7 @@ async fn test_already_exists_add_node() -> Result<()> { #[tokio::test(flavor = "multi_thread", worker_threads = 1)] async fn test_successfully_get_nodes() -> Result<()> { - let (_, cluster_api) = new_cluster_api().await?; + let (_, cluster_api) = new_cluster_api(Duration::from_secs(60)).await?; let nodes = cluster_api .get_nodes("test-cluster-id", "test-cluster-id") @@ -100,7 +101,7 @@ async fn test_successfully_get_nodes() -> Result<()> { #[tokio::test(flavor = "multi_thread", worker_threads = 1)] async fn test_successfully_drop_node() -> Result<()> { - let (_, cluster_api) = new_cluster_api().await?; + let (_, cluster_api) = new_cluster_api(Duration::from_secs(60)).await?; let node_info = create_test_node_info(); cluster_api.add_node(node_info.clone()).await?; @@ -121,7 +122,7 @@ async fn test_successfully_drop_node() -> Result<()> { #[tokio::test(flavor = "multi_thread", worker_threads = 1)] async fn test_unknown_node_drop_node() -> Result<()> { - let (_, cluster_api) = new_cluster_api().await?; + let (_, cluster_api) = new_cluster_api(Duration::from_secs(60)).await?; match cluster_api.drop_node(String::from("UNKNOWN_ID")).await { Ok(_) => { /*panic!("Unknown node drop node must be return Err.")*/ } @@ -134,10 +135,10 @@ async fn test_unknown_node_drop_node() -> Result<()> { #[tokio::test(flavor = "multi_thread", worker_threads = 1)] async fn test_successfully_heartbeat_node() -> Result<()> { let now_ms = SeqV::<()>::now_ms(); - let (kv_api, cluster_api) = new_cluster_api().await?; + let (kv_api, cluster_api) = new_cluster_api(Duration::from_secs(60)).await?; - let node_info = create_test_node_info(); - cluster_api.add_node(node_info.clone()).await?; + let mut node_info = create_test_node_info(); + let res = cluster_api.add_node(node_info.clone()).await?; for key in [ "__fd_clusters_v5/test%2dtenant%2did/online_nodes/test_node", @@ -150,7 +151,7 @@ async fn test_successfully_heartbeat_node() -> Result<()> { } let now_ms = SeqV::<()>::now_ms(); - cluster_api.heartbeat(&node_info).await?; + cluster_api.heartbeat(&mut node_info, res).await?; for key in [ "__fd_clusters_v5/test%2dtenant%2did/online_nodes/test_node", @@ -163,6 +164,104 @@ async fn test_successfully_heartbeat_node() -> Result<()> { Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn test_successfully_create_warehouse() -> Result<()> { + let (_, cluster_mgr, nodes) = nodes(Duration::from_mins(30), 2).await?; + + let create_warehouse = cluster_mgr.create_warehouse("test_warehouse".to_string(), vec![ + SelectedNode::Random(None), + SelectedNode::Random(None), + ]); + + create_warehouse.await?; + + let get_warehouse_nodes = cluster_mgr.get_nodes("test_warehouse", "test_warehouse"); + + let warehouse_nodes = get_warehouse_nodes.await?; + + assert_eq!(warehouse_nodes.len(), 2); + assert_eq!(warehouse_nodes.last().map(|x| &x.id), nodes.last()); + assert_eq!(warehouse_nodes.first().map(|x| &x.id), nodes.first()); + + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn test_create_duplicated_warehouse() -> Result<()> { + let (_, cluster_mgr, _nodes) = nodes(Duration::from_mins(30), 2).await?; + + let create_warehouse = + cluster_mgr.create_warehouse("test_warehouse".to_string(), vec![SelectedNode::Random( + None, + )]); + + create_warehouse.await?; + + let create_warehouse = + cluster_mgr.create_warehouse("test_warehouse".to_string(), vec![SelectedNode::Random( + None, + )]); + + let res = create_warehouse.await; + assert!(res.is_err()); + assert_eq!(res.unwrap_err().code(), 2405); + + let create_warehouse = + cluster_mgr.create_warehouse("test_warehouse_2".to_string(), vec![SelectedNode::Random( + None, + )]); + + create_warehouse.await?; + + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn test_create_warehouse_with_no_resources() -> Result<()> { + let (_, cluster_mgr, _nodes) = nodes(Duration::from_mins(30), 2).await?; + + let create_warehouse = + cluster_mgr.create_warehouse("test_warehouse_1".to_string(), vec![SelectedNode::Random( + None, + )]); + + create_warehouse.await?; + + let create_warehouse = + cluster_mgr.create_warehouse("test_warehouse_2".to_string(), vec![SelectedNode::Random( + None, + )]); + + create_warehouse.await?; + + let create_warehouse = + cluster_mgr.create_warehouse("test_warehouse_3".to_string(), vec![SelectedNode::Random( + None, + )]); + + let res = create_warehouse.await; + + assert!(res.is_err()); + assert_eq!(res.unwrap_err().code(), 2404); + + Ok(()) +} + +fn empty_node(id: &str) -> NodeInfo { + NodeInfo { + id: id.to_string(), + secret: "".to_string(), + cpu_nums: 0, + version: 0, + http_address: "".to_string(), + flight_address: "".to_string(), + discovery_address: "".to_string(), + binary_version: "".to_string(), + cluster_id: "".to_string(), + warehouse_id: "".to_string(), + } +} + fn create_test_node_info() -> NodeInfo { NodeInfo { id: String::from("test_node"), @@ -178,9 +277,21 @@ fn create_test_node_info() -> NodeInfo { } } -async fn new_cluster_api() -> Result<(MetaStore, ClusterMgr)> { +async fn nodes(lift: Duration, size: usize) -> Result<(MetaStore, WarehouseMgr, Vec)> { + let (kv_api, cluster_manager) = new_cluster_api(lift).await?; + + let mut nodes = Vec::with_capacity(size); + for _index in 0..size { + let name = GlobalUniqName::unique(); + cluster_manager.add_node(empty_node(&name)).await?; + nodes.push(name); + } + + Ok((kv_api, cluster_manager, nodes)) +} + +async fn new_cluster_api(lift: Duration) -> Result<(MetaStore, WarehouseMgr)> { let test_api = MetaStore::L(Arc::new(MemMeta::default())); - let cluster_manager = - ClusterMgr::create(test_api.clone(), "test-tenant-id", Duration::from_secs(60))?; + let cluster_manager = WarehouseMgr::create(test_api.clone(), "test-tenant-id", lift)?; Ok((test_api, cluster_manager)) } diff --git a/src/query/management/tests/it/main.rs b/src/query/management/tests/it/main.rs index f3f2482be8ade..9f4f238e518d8 100644 --- a/src/query/management/tests/it/main.rs +++ b/src/query/management/tests/it/main.rs @@ -1,5 +1,5 @@ // Copyright 2021 Datafuse Labs -// + // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. // You may obtain a copy of the License at @@ -12,6 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. +#![feature(duration_constructors)] #![allow(clippy::uninlined_format_args)] mod cluster; diff --git a/src/query/service/src/clusters/cluster.rs b/src/query/service/src/clusters/cluster.rs index 8ab381200664d..f315c37fe7b2a 100644 --- a/src/query/service/src/clusters/cluster.rs +++ b/src/query/service/src/clusters/cluster.rs @@ -38,8 +38,8 @@ use databend_common_config::DATABEND_COMMIT_VERSION; use databend_common_exception::ErrorCode; use databend_common_exception::Result; use databend_common_grpc::ConnectionFactory; -use databend_common_management::ClusterApi; -use databend_common_management::ClusterMgr; +use databend_common_management::WarehouseApi; +use databend_common_management::WarehouseMgr; use databend_common_meta_store::MetaStore; use databend_common_meta_store::MetaStoreProvider; use databend_common_meta_types::NodeInfo; @@ -64,7 +64,7 @@ pub struct ClusterDiscovery { local_id: String, local_secret: String, heartbeat: Mutex, - api_provider: Arc, + api_provider: Arc, cluster_id: String, tenant_id: String, flight_address: String, @@ -235,11 +235,11 @@ impl ClusterDiscovery { fn create_provider( cfg: &InnerConfig, metastore: MetaStore, - ) -> Result<(Duration, Arc)> { + ) -> Result<(Duration, Arc)> { // TODO: generate if tenant or cluster id is empty let tenant_id = &cfg.query.tenant_id; let lift_time = Duration::from_secs(60); - let cluster_manager = ClusterMgr::create(metastore, tenant_id.tenant_name(), lift_time)?; + let cluster_manager = WarehouseMgr::create(metastore, tenant_id.tenant_name(), lift_time)?; Ok((lift_time, Arc::new(cluster_manager))) } @@ -441,15 +441,15 @@ impl ClusterDiscovery { node_info.warehouse_id = self.cluster_id.clone(); self.drop_invalid_nodes(&node_info).await?; match self.api_provider.add_node(node_info.clone()).await { - Ok(_) => self.start_heartbeat(node_info).await, + Ok(seq) => self.start_heartbeat(node_info, seq).await, Err(cause) => Err(cause.add_message_back("(while cluster api add_node).")), } } #[async_backtrace::framed] - async fn start_heartbeat(self: &Arc, node_info: NodeInfo) -> Result<()> { + async fn start_heartbeat(self: &Arc, node_info: NodeInfo, seq: u64) -> Result<()> { let mut heartbeat = self.heartbeat.lock().await; - heartbeat.start(node_info); + heartbeat.start(node_info, seq); Ok(()) } } @@ -458,7 +458,7 @@ struct ClusterHeartbeat { timeout: Duration, shutdown: Arc, shutdown_notify: Arc, - cluster_api: Arc, + cluster_api: Arc, shutdown_handler: Option>, cluster_id: String, tenant_id: String, @@ -467,7 +467,7 @@ struct ClusterHeartbeat { impl ClusterHeartbeat { pub fn create( timeout: Duration, - cluster_api: Arc, + cluster_api: Arc, cluster_id: String, tenant_id: String, ) -> ClusterHeartbeat { @@ -482,7 +482,7 @@ impl ClusterHeartbeat { } } - fn heartbeat_loop(&self, node: NodeInfo) -> impl Future + 'static { + fn heartbeat_loop(&self, mut node: NodeInfo, seq: u64) -> impl Future + 'static { let shutdown = self.shutdown.clone(); let shutdown_notify = self.shutdown_notify.clone(); let cluster_api = self.cluster_api.clone(); @@ -493,6 +493,7 @@ impl ClusterHeartbeat { async move { let mut shutdown_notified = Box::pin(shutdown_notify.notified()); + let mut match_seq = seq; while !shutdown.load(Ordering::Relaxed) { let mills = { let mut rng = thread_rng(); @@ -507,16 +508,21 @@ impl ClusterHeartbeat { } Either::Right((_, new_shutdown_notified)) => { shutdown_notified = new_shutdown_notified; - let heartbeat = cluster_api.heartbeat(&node); - if let Err(failure) = heartbeat.await { - metric_incr_cluster_heartbeat_count( - &node.id, - &node.flight_address, - &cluster_id, - &tenant_id, - "failure", - ); - error!("Cluster cluster api heartbeat failure: {:?}", failure); + let heartbeat = cluster_api.heartbeat(&mut node, match_seq); + match heartbeat.await { + Ok(new_match_seq) => { + match_seq = new_match_seq; + } + Err(failure) => { + metric_incr_cluster_heartbeat_count( + &node.id, + &node.flight_address, + &cluster_id, + &tenant_id, + "failure", + ); + error!("Cluster cluster api heartbeat failure: {:?}", failure); + } } } } @@ -528,9 +534,9 @@ impl ClusterHeartbeat { (duration / 3).as_millis()..=((duration / 3) * 2).as_millis() } - pub fn start(&mut self, node_info: NodeInfo) { + pub fn start(&mut self, node_info: NodeInfo, seq: u64) { self.shutdown_handler = Some(databend_common_base::runtime::spawn( - self.heartbeat_loop(node_info), + self.heartbeat_loop(node_info, seq), )); } From 566239cd77d67ab68b9e6cb3f97615d6f7564d21 Mon Sep 17 00:00:00 2001 From: zhang2014 Date: Tue, 17 Dec 2024 14:54:17 +0800 Subject: [PATCH 02/53] feat(cluster): support custom management cluster --- src/common/exception/src/exception_code.rs | 1 + src/meta/types/src/cluster.rs | 9 + src/meta/types/src/lib.rs | 1 + src/query/config/src/config.rs | 22 + src/query/management/src/lib.rs | 1 + src/query/management/src/warehouse/mod.rs | 1 + .../management/src/warehouse/warehouse_api.rs | 11 +- .../management/src/warehouse/warehouse_mgr.rs | 165 ++++++- src/query/management/tests/it/cluster.rs | 297 ------------- src/query/management/tests/it/main.rs | 2 +- src/query/management/tests/it/warehouse.rs | 409 ++++++++++++++++++ src/query/service/src/clusters/cluster.rs | 3 + 12 files changed, 601 insertions(+), 321 deletions(-) delete mode 100644 src/query/management/tests/it/cluster.rs create mode 100644 src/query/management/tests/it/warehouse.rs diff --git a/src/common/exception/src/exception_code.rs b/src/common/exception/src/exception_code.rs index e5ac696f7f201..7fdf7dc3289b2 100644 --- a/src/common/exception/src/exception_code.rs +++ b/src/common/exception/src/exception_code.rs @@ -298,6 +298,7 @@ build_exceptions! { WarehouseAlreadyExists(2405), UnknownWarehouse(2406), WarehouseOperateConflict(2407), + EmptyNodesForWarehouse(2408), // Stage error codes. UnknownStage(2501), diff --git a/src/meta/types/src/cluster.rs b/src/meta/types/src/cluster.rs index f0e25319f74ff..8bdae6f0750bd 100644 --- a/src/meta/types/src/cluster.rs +++ b/src/meta/types/src/cluster.rs @@ -70,6 +70,13 @@ impl fmt::Display for Node { } } +#[derive(Serialize, Deserialize, Clone, Debug, Eq, PartialEq, Ord, PartialOrd, Default)] +pub enum NodeType { + #[default] + SelfManaged, + SystemManaged, +} + /// Query node #[derive(Serialize, Deserialize, Clone, Debug, Eq, PartialEq, Ord, PartialOrd, Default)] #[serde(default)] @@ -82,6 +89,7 @@ pub struct NodeInfo { pub flight_address: String, pub discovery_address: String, pub binary_version: String, + pub node_type: NodeType, #[serde(skip_serializing_if = "String::is_empty")] pub cluster_id: String, @@ -110,6 +118,7 @@ impl NodeInfo { binary_version, cluster_id: "".to_string(), warehouse_id: "".to_string(), + node_type: NodeType::SystemManaged, } } diff --git a/src/meta/types/src/lib.rs b/src/meta/types/src/lib.rs index 458f7701dacfc..7c02068e24e66 100644 --- a/src/meta/types/src/lib.rs +++ b/src/meta/types/src/lib.rs @@ -63,6 +63,7 @@ pub use applied_state::AppliedState; pub use change::Change; pub use cluster::Node; pub use cluster::NodeInfo; +pub use cluster::NodeType; pub use endpoint::Endpoint; pub use errors::meta_api_errors::MetaAPIError; pub use errors::meta_api_errors::MetaDataError; diff --git a/src/query/config/src/config.rs b/src/query/config/src/config.rs index c693e97eadace..7ecdeaeb39693 100644 --- a/src/query/config/src/config.rs +++ b/src/query/config/src/config.rs @@ -152,6 +152,9 @@ pub struct Config { /// when converted from inner config, all catalog configurations will store in `catalogs` #[clap(skip)] pub catalogs: HashMap, + + #[clap(flatten)] + pub resources_manager: ResourceManagerConfig, } #[derive(Subcommand, Default, Debug, Eq, PartialEq, Clone, Serialize, Deserialize)] @@ -2952,6 +2955,25 @@ pub struct SpillConfig { pub spill_local_disk_max_bytes: u64, } +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize, Args, Default)] +#[serde(default, deny_unknown_fields)] +pub struct ResourceManagerConfig { + #[clap( + long = "cluster-type", + value_name = "VALUE", + default_value = "system-managed" + )] + #[serde(rename = "type")] + pub typ: String, + + #[clap(long, value_name = "VALUE")] + pub instance_type: Option, + + #[clap(long, value_name = "VALUE", default_value = "18446744073709551615")] + /// Allow space in bytes to spill to local disk. + pub spill_local_disk_max_bytes: u64, +} + mod cache_config_converters { use std::path::PathBuf; diff --git a/src/query/management/src/lib.rs b/src/query/management/src/lib.rs index cc4a34a10bf99..f66a003a11ba0 100644 --- a/src/query/management/src/lib.rs +++ b/src/query/management/src/lib.rs @@ -52,4 +52,5 @@ pub use user::UserApi; pub use user::UserMgr; pub use warehouse::SelectedNode; pub use warehouse::WarehouseApi; +pub use warehouse::WarehouseInfo; pub use warehouse::WarehouseMgr; diff --git a/src/query/management/src/warehouse/mod.rs b/src/query/management/src/warehouse/mod.rs index 206f779470b1c..0667ed3a7c87b 100644 --- a/src/query/management/src/warehouse/mod.rs +++ b/src/query/management/src/warehouse/mod.rs @@ -17,4 +17,5 @@ mod warehouse_mgr; pub use warehouse_api::SelectedNode; pub use warehouse_api::WarehouseApi; +pub use warehouse_api::WarehouseInfo; pub use warehouse_mgr::WarehouseMgr; diff --git a/src/query/management/src/warehouse/warehouse_api.rs b/src/query/management/src/warehouse/warehouse_api.rs index 64df7fc73109c..aaa2e5ef4d3a3 100644 --- a/src/query/management/src/warehouse/warehouse_api.rs +++ b/src/query/management/src/warehouse/warehouse_api.rs @@ -15,13 +15,18 @@ use databend_common_exception::Result; use databend_common_meta_types::NodeInfo; -#[derive(serde::Serialize, serde::Deserialize, Clone)] +#[derive(serde::Serialize, serde::Deserialize, Clone, Eq, PartialEq)] pub enum SelectedNode { Random(Option), } -#[derive(serde::Serialize, serde::Deserialize)] -pub struct WarehouseInfo { +#[derive(serde::Serialize, serde::Deserialize, Eq, PartialEq)] +pub enum WarehouseInfo { + SelfManaged, + SystemManaged(SystemManagedInfo), +} +#[derive(serde::Serialize, serde::Deserialize, Eq, PartialEq)] +pub struct SystemManagedInfo { pub id: String, pub status: String, pub display_name: String, diff --git a/src/query/management/src/warehouse/warehouse_mgr.rs b/src/query/management/src/warehouse/warehouse_mgr.rs index e50b1d08fe4c6..43e2446e344df 100644 --- a/src/query/management/src/warehouse/warehouse_mgr.rs +++ b/src/query/management/src/warehouse/warehouse_mgr.rs @@ -27,6 +27,7 @@ use databend_common_meta_types::txn_op_response::Response; use databend_common_meta_types::ConditionResult; use databend_common_meta_types::MatchSeq; use databend_common_meta_types::NodeInfo; +use databend_common_meta_types::NodeType; use databend_common_meta_types::TxnCondition; use databend_common_meta_types::TxnGetResponse; use databend_common_meta_types::TxnOp; @@ -35,6 +36,7 @@ use databend_common_meta_types::TxnReply; use databend_common_meta_types::TxnRequest; use crate::warehouse::warehouse_api::SelectedNode; +use crate::warehouse::warehouse_api::SystemManagedInfo; use crate::warehouse::warehouse_api::WarehouseInfo; use crate::warehouse::WarehouseApi; @@ -101,11 +103,115 @@ fn get_prev_value(res: Option<&TxnOpResponse>) -> Option<&SeqV> { } impl WarehouseMgr { - #[fastrace::trace] - #[async_backtrace::framed] - async fn upsert_node(&self, mut node: NodeInfo, seq: MatchSeq) -> Result { + fn cluster_key(&self, node: &NodeInfo) -> Result { + Ok(format!( + "{}/{}/{}/{}", + self.meta_key_prefix, + escape_for_key(&node.warehouse_id)?, + escape_for_key(&node.cluster_id)?, + escape_for_key(&node.id)? + )) + } + + async fn upsert_self_managed(&self, mut node: NodeInfo, seq: MatchSeq) -> Result { + if node.warehouse_id.is_empty() || node.cluster_id.is_empty() { + return Err(ErrorCode::InvalidWarehouse( + "The warehouse_id and cluster_id for self managed node must not be empty.", + )); + } + let mut txn = TxnRequest::default(); + let node_key = format!("{}/{}", self.node_key_prefix, escape_for_key(&node.id)?); + + txn.condition.push(map_condition(&node_key, seq)); + + txn.if_then.push(TxnOp::put_with_ttl( + node_key.clone(), + serde_json::to_vec(&node)?, + Some(self.lift_time), + )); + + let warehouse_node_key = self.cluster_key(&node)?; + let warehouse_info_key = format!( + "{}/{}", + self.warehouse_key_prefix, + escape_for_key(&node.warehouse_id)? + ); + + node.cluster_id = String::new(); + node.warehouse_id = String::new(); + txn.if_then.push(TxnOp::put_with_ttl( + warehouse_node_key, + serde_json::to_vec(&node)?, + Some(self.lift_time), + )); + + // upsert warehouse info if self-managed. + txn.if_then.push(TxnOp::put_with_ttl( + warehouse_info_key.clone(), + serde_json::to_vec(&WarehouseInfo::SelfManaged)?, + Some(self.lift_time), + )); + + txn.if_then.push(TxnOp::get(node_key.clone())); + txn.else_then.push(TxnOp::get(node_key.clone())); + + if seq != MatchSeq::Exact(0) { + return Ok(self.metastore.transaction(txn).await?); + } + + let mut exact_seq = 0; + let mut retry_count = 0; + + loop { + let mut warehouse_txn = txn.clone(); + + // insert if warehouse info is not exists or SelfManaged + warehouse_txn + .condition + .push(TxnCondition::eq_seq(warehouse_info_key.clone(), exact_seq)); + warehouse_txn + .else_then + .push(TxnOp::get(warehouse_info_key.clone())); + + return match self.metastore.transaction(warehouse_txn).await? { + mut response if !response.success => { + return match response.responses.pop().and_then(|x| x.response) { + Some(Response::Get(data)) => match data.value { + None => Ok(response), + Some(value) if value.seq == 0 => Ok(response), + Some(value) => match serde_json::from_slice(&value.data)? { + WarehouseInfo::SystemManaged(_) => { + Err(ErrorCode::WarehouseAlreadyExists("")) + } + WarehouseInfo::SelfManaged => match response.responses.first() { + // already exists node. + Some(TxnOpResponse { + response: + Some(Response::Get(TxnGetResponse { + value: Some(value), + .. + })), + }) if value.seq != 0 => Ok(response), + _ => { + log::info!("Self-managed warehouse has already been created by other nodes; attempt to join it. Retry count: {}", retry_count); + retry_count += 1; + exact_seq = value.seq; + continue; + } + }, + }, + }, + _ => Ok(response), + }; + } + response => Ok(response), + }; + } + } + async fn upsert_system_managed(&self, mut node: NodeInfo, seq: MatchSeq) -> Result { + let mut txn = TxnRequest::default(); let node_key = format!("{}/{}", self.node_key_prefix, escape_for_key(&node.id)?); txn.condition.push(map_condition(&node_key, seq)); @@ -116,6 +222,7 @@ impl WarehouseMgr { Some(self.lift_time), )); + // If the warehouse has already been assigned. if !node.cluster_id.is_empty() && !node.warehouse_id.is_empty() { let cluster_key = format!( "{}/{}/{}/{}", @@ -125,7 +232,6 @@ impl WarehouseMgr { escape_for_key(&node.id)? ); - // txn.condition.push(map_condition(&cluster_key, seq)); node.cluster_id = String::new(); node.warehouse_id = String::new(); txn.if_then.push(TxnOp::put_with_ttl( @@ -140,6 +246,15 @@ impl WarehouseMgr { Ok(self.metastore.transaction(txn).await?) } + #[fastrace::trace] + #[async_backtrace::framed] + async fn upsert_node(&self, node: NodeInfo, seq: MatchSeq) -> Result { + match node.node_type { + NodeType::SelfManaged => self.upsert_self_managed(node, seq).await, + NodeType::SystemManaged => self.upsert_system_managed(node, seq).await, + } + } + async fn leave_cluster(&self, node_info: &mut NodeInfo, seq: u64) -> Result { let mut cluster_id = String::new(); let mut warehouse_id = String::new(); @@ -163,8 +278,8 @@ impl WarehouseMgr { } Ok(response) => match response.responses.last() { Some(TxnOpResponse { - response: Some(Response::Get(TxnGetResponse { value: Some(v), .. })), - }) => Ok(v.seq), + response: Some(Response::Get(TxnGetResponse { value: Some(v), .. })), + }) => Ok(v.seq), _ => Err(ErrorCode::MetaServiceError("Meta insert failure.")), }, } @@ -188,8 +303,8 @@ impl WarehouseMgr { } Ok(response) => match response.responses.last() { Some(TxnOpResponse { - response: Some(Response::Get(TxnGetResponse { value: Some(v), .. })), - }) => Ok(v.seq), + response: Some(Response::Get(TxnGetResponse { value: Some(v), .. })), + }) => Ok(v.seq), _ => Err(ErrorCode::MetaServiceError("Meta insert failure.")), }, } @@ -199,8 +314,8 @@ impl WarehouseMgr { match reply.responses.first() { None => self.leave_cluster(node, 0).await, Some(TxnOpResponse { - response: Some(Response::Get(res)), - }) => match &res.value { + response: Some(Response::Get(res)), + }) => match &res.value { None => self.leave_cluster(node, 0).await, Some(value) => { let node_info = serde_json::from_slice::(&value.data)?; @@ -327,18 +442,18 @@ impl WarehouseApi for WarehouseMgr { if fetch_reply.responses.len() != delete_txn.if_then.len() { // TODO: maybe auto retry? - return Err(ErrorCode::WarehouseOperateConflict("Missing node info in online nodes list. Perhaps some nodes offline while in drop warehouse. You can try it again.")); + return Err(ErrorCode::WarehouseOperateConflict("Missing node info in online nodes list. It's possible that some nodes offline during the drop warehouse. You may try the operation again.")); } for response in fetch_reply.responses { let Some(Response::Get(response)) = response.response else { // TODO: maybe auto retry? - return Err(ErrorCode::WarehouseOperateConflict("Missing node info in online nodes list. Perhaps some nodes offline while in drop warehouse. You can try it again.")); + return Err(ErrorCode::WarehouseOperateConflict("Missing node info in online nodes list. It's possible that some nodes offline during the drop warehouse. You may try the operation again.")); }; let Some(value) = &response.value else { // TODO: maybe auto retry? - return Err(ErrorCode::WarehouseOperateConflict("Missing node info in online nodes list. Perhaps some nodes offline while in drop warehouse. You can try it again.")); + return Err(ErrorCode::WarehouseOperateConflict("Missing node info in online nodes list. It's possible that some nodes offline during the drop warehouse. You may try the operation again.")); }; let mut node_info = serde_json::from_slice::(&value.data)?; @@ -384,9 +499,19 @@ impl WarehouseApi for WarehouseMgr { async fn create_warehouse(&self, warehouse: String, nodes: Vec) -> Result<()> { assert!(nodes.iter().all(|x| matches!(x, SelectedNode::Random(_)))); + if warehouse.is_empty() { + return Err(ErrorCode::InvalidWarehouse("Warehouse name is empty.")); + } + + if nodes.is_empty() { + return Err(ErrorCode::EmptyNodesForWarehouse( + "Cannot create warehouse with empty nodes.", + )); + } + loop { let mut selected_nodes = Vec::with_capacity(nodes.len()); - let mut allocated_ndoes = HashSet::with_capacity(nodes.len()); + let mut allocated_nodes = HashSet::with_capacity(nodes.len()); // get online nodes let online_nodes = self.metastore.prefix_list_kv(&self.node_key_prefix).await?; @@ -405,11 +530,11 @@ impl WarehouseApi for WarehouseMgr { if node_info.warehouse_id.is_empty() && node_info.cluster_id.is_empty() - && !allocated_ndoes.contains(&node_info.id) + && !allocated_nodes.contains(&node_info.id) { node_info.cluster_id = warehouse.clone(); node_info.warehouse_id = warehouse.clone(); - allocated_ndoes.insert(node_info.id.clone()); + allocated_nodes.insert(node_info.id.clone()); selected_nodes.push((v.seq, node_info)); break; } @@ -465,12 +590,12 @@ impl WarehouseApi for WarehouseMgr { .push(map_condition(&warehouse_key, MatchSeq::Exact(0))); txn.if_then.push(TxnOp::put( warehouse_key.clone(), - serde_json::to_vec(&WarehouseInfo { + serde_json::to_vec(&WarehouseInfo::SystemManaged(SystemManagedInfo { id: GlobalUniqName::unique(), status: "Running".to_string(), display_name: warehouse.clone(), clusters: vec![nodes.clone()], - })?, + }))?, )); txn.else_then.push(TxnOp::get(warehouse_key)); @@ -478,8 +603,8 @@ impl WarehouseApi for WarehouseMgr { res if res.success => Ok(()), res => match res.responses.last() { Some(TxnOpResponse { - response: Some(Response::Get(res)), - }) => { + response: Some(Response::Get(res)), + }) => { if matches!(&res.value, Some(v) if v.seq != 0) { return Err(ErrorCode::WarehouseAlreadyExists( "Warehouse already exists", diff --git a/src/query/management/tests/it/cluster.rs b/src/query/management/tests/it/cluster.rs deleted file mode 100644 index 6341f0fc24638..0000000000000 --- a/src/query/management/tests/it/cluster.rs +++ /dev/null @@ -1,297 +0,0 @@ -// Copyright 2021 Datafuse Labs -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -use std::sync::Arc; -use std::time::Duration; - -use databend_common_base::base::tokio; -use databend_common_base::base::GlobalUniqName; -use databend_common_exception::Result; -use databend_common_management::*; -use databend_common_meta_embedded::MemMeta; -use databend_common_meta_kvapi::kvapi::KVApi; -use databend_common_meta_store::MetaStore; -use databend_common_meta_types::seq_value::SeqV; -use databend_common_meta_types::NodeInfo; - -#[tokio::test(flavor = "multi_thread", worker_threads = 1)] -async fn test_successfully_add_node() -> Result<()> { - let now_ms = SeqV::<()>::now_ms(); - let (kv_api, cluster_api) = new_cluster_api(Duration::from_secs(60)).await?; - - let mut node_info = create_test_node_info(); - cluster_api.add_node(node_info.clone()).await?; - let online_node = kv_api - .get_kv("__fd_clusters_v5/test%2dtenant%2did/online_nodes/test_node") - .await?; - - match online_node { - Some(SeqV { - seq: 1, - meta, - data: value, - }) => { - assert!(meta.unwrap().get_expire_at_ms().unwrap() - now_ms >= 59_000); - assert_eq!(value, serde_json::to_vec(&node_info)?); - } - catch => panic!("GetKVActionReply{:?}", catch), - } - - let online_cluster = kv_api.get_kv("__fd_clusters_v5/test%2dtenant%2did/online_clusters/test%2dcluster%2did/test%2dcluster%2did/test_node").await?; - - match online_cluster { - Some(SeqV { - meta, data: value, .. - }) => { - assert!(meta.unwrap().get_expire_at_ms().unwrap() - now_ms >= 59_000); - node_info.cluster_id = String::new(); - node_info.warehouse_id = String::new(); - assert_eq!(value, serde_json::to_vec(&node_info)?); - } - catch => panic!("GetKVActionReply{:?}", catch), - } - - Ok(()) -} - -#[tokio::test(flavor = "multi_thread", worker_threads = 1)] -async fn test_already_exists_add_node() -> Result<()> { - let (_, cluster_api) = new_cluster_api(Duration::from_secs(60)).await?; - - let node_info = create_test_node_info(); - cluster_api.add_node(node_info.clone()).await?; - - match cluster_api.add_node(node_info.clone()).await { - Ok(_) => panic!("Already exists add node must be return Err."), - Err(cause) => assert_eq!(cause.code(), 2402), - } - - Ok(()) -} - -#[tokio::test(flavor = "multi_thread", worker_threads = 1)] -async fn test_successfully_get_nodes() -> Result<()> { - let (_, cluster_api) = new_cluster_api(Duration::from_secs(60)).await?; - - let nodes = cluster_api - .get_nodes("test-cluster-id", "test-cluster-id") - .await?; - assert_eq!(nodes, vec![]); - - let node_info = create_test_node_info(); - cluster_api.add_node(node_info.clone()).await?; - - let nodes = cluster_api - .get_nodes("test-cluster-id", "test-cluster-id") - .await?; - assert_eq!(nodes, vec![node_info]); - Ok(()) -} - -#[tokio::test(flavor = "multi_thread", worker_threads = 1)] -async fn test_successfully_drop_node() -> Result<()> { - let (_, cluster_api) = new_cluster_api(Duration::from_secs(60)).await?; - - let node_info = create_test_node_info(); - cluster_api.add_node(node_info.clone()).await?; - - let nodes = cluster_api - .get_nodes("test-cluster-id", "test-cluster-id") - .await?; - assert_eq!(nodes, vec![node_info.clone()]); - - cluster_api.drop_node(node_info.id).await?; - - let nodes = cluster_api - .get_nodes("test-cluster-id", "test-cluster-id") - .await?; - assert_eq!(nodes, vec![]); - Ok(()) -} - -#[tokio::test(flavor = "multi_thread", worker_threads = 1)] -async fn test_unknown_node_drop_node() -> Result<()> { - let (_, cluster_api) = new_cluster_api(Duration::from_secs(60)).await?; - - match cluster_api.drop_node(String::from("UNKNOWN_ID")).await { - Ok(_) => { /*panic!("Unknown node drop node must be return Err.")*/ } - Err(cause) => assert_eq!(cause.code(), 2401), - } - - Ok(()) -} - -#[tokio::test(flavor = "multi_thread", worker_threads = 1)] -async fn test_successfully_heartbeat_node() -> Result<()> { - let now_ms = SeqV::<()>::now_ms(); - let (kv_api, cluster_api) = new_cluster_api(Duration::from_secs(60)).await?; - - let mut node_info = create_test_node_info(); - let res = cluster_api.add_node(node_info.clone()).await?; - - for key in [ - "__fd_clusters_v5/test%2dtenant%2did/online_nodes/test_node", - "__fd_clusters_v5/test%2dtenant%2did/online_clusters/test%2dcluster%2did/test%2dcluster%2did/test_node", - ] { - let value = kv_api.get_kv(key).await?; - let meta = value.unwrap().meta.unwrap(); - let expire_ms = meta.get_expire_at_ms().unwrap(); - assert!(expire_ms - now_ms >= 59_000); - } - - let now_ms = SeqV::<()>::now_ms(); - cluster_api.heartbeat(&mut node_info, res).await?; - - for key in [ - "__fd_clusters_v5/test%2dtenant%2did/online_nodes/test_node", - "__fd_clusters_v5/test%2dtenant%2did/online_clusters/test%2dcluster%2did/test%2dcluster%2did/test_node", - ] { - let value = kv_api.get_kv(key).await?; - assert!(value.unwrap().meta.unwrap().get_expire_at_ms().unwrap() - now_ms >= 59_000); - } - - Ok(()) -} - -#[tokio::test(flavor = "multi_thread", worker_threads = 1)] -async fn test_successfully_create_warehouse() -> Result<()> { - let (_, cluster_mgr, nodes) = nodes(Duration::from_mins(30), 2).await?; - - let create_warehouse = cluster_mgr.create_warehouse("test_warehouse".to_string(), vec![ - SelectedNode::Random(None), - SelectedNode::Random(None), - ]); - - create_warehouse.await?; - - let get_warehouse_nodes = cluster_mgr.get_nodes("test_warehouse", "test_warehouse"); - - let warehouse_nodes = get_warehouse_nodes.await?; - - assert_eq!(warehouse_nodes.len(), 2); - assert_eq!(warehouse_nodes.last().map(|x| &x.id), nodes.last()); - assert_eq!(warehouse_nodes.first().map(|x| &x.id), nodes.first()); - - Ok(()) -} - -#[tokio::test(flavor = "multi_thread", worker_threads = 1)] -async fn test_create_duplicated_warehouse() -> Result<()> { - let (_, cluster_mgr, _nodes) = nodes(Duration::from_mins(30), 2).await?; - - let create_warehouse = - cluster_mgr.create_warehouse("test_warehouse".to_string(), vec![SelectedNode::Random( - None, - )]); - - create_warehouse.await?; - - let create_warehouse = - cluster_mgr.create_warehouse("test_warehouse".to_string(), vec![SelectedNode::Random( - None, - )]); - - let res = create_warehouse.await; - assert!(res.is_err()); - assert_eq!(res.unwrap_err().code(), 2405); - - let create_warehouse = - cluster_mgr.create_warehouse("test_warehouse_2".to_string(), vec![SelectedNode::Random( - None, - )]); - - create_warehouse.await?; - - Ok(()) -} - -#[tokio::test(flavor = "multi_thread", worker_threads = 1)] -async fn test_create_warehouse_with_no_resources() -> Result<()> { - let (_, cluster_mgr, _nodes) = nodes(Duration::from_mins(30), 2).await?; - - let create_warehouse = - cluster_mgr.create_warehouse("test_warehouse_1".to_string(), vec![SelectedNode::Random( - None, - )]); - - create_warehouse.await?; - - let create_warehouse = - cluster_mgr.create_warehouse("test_warehouse_2".to_string(), vec![SelectedNode::Random( - None, - )]); - - create_warehouse.await?; - - let create_warehouse = - cluster_mgr.create_warehouse("test_warehouse_3".to_string(), vec![SelectedNode::Random( - None, - )]); - - let res = create_warehouse.await; - - assert!(res.is_err()); - assert_eq!(res.unwrap_err().code(), 2404); - - Ok(()) -} - -fn empty_node(id: &str) -> NodeInfo { - NodeInfo { - id: id.to_string(), - secret: "".to_string(), - cpu_nums: 0, - version: 0, - http_address: "".to_string(), - flight_address: "".to_string(), - discovery_address: "".to_string(), - binary_version: "".to_string(), - cluster_id: "".to_string(), - warehouse_id: "".to_string(), - } -} - -fn create_test_node_info() -> NodeInfo { - NodeInfo { - id: String::from("test_node"), - secret: "".to_string(), - cpu_nums: 0, - version: 0, - http_address: "ip3:port".to_string(), - flight_address: String::from("ip:port"), - discovery_address: "ip2:port".to_string(), - binary_version: "binary_version".to_string(), - cluster_id: "test-cluster-id".to_string(), - warehouse_id: "test-cluster-id".to_string(), - } -} - -async fn nodes(lift: Duration, size: usize) -> Result<(MetaStore, WarehouseMgr, Vec)> { - let (kv_api, cluster_manager) = new_cluster_api(lift).await?; - - let mut nodes = Vec::with_capacity(size); - for _index in 0..size { - let name = GlobalUniqName::unique(); - cluster_manager.add_node(empty_node(&name)).await?; - nodes.push(name); - } - - Ok((kv_api, cluster_manager, nodes)) -} - -async fn new_cluster_api(lift: Duration) -> Result<(MetaStore, WarehouseMgr)> { - let test_api = MetaStore::L(Arc::new(MemMeta::default())); - let cluster_manager = WarehouseMgr::create(test_api.clone(), "test-tenant-id", lift)?; - Ok((test_api, cluster_manager)) -} diff --git a/src/query/management/tests/it/main.rs b/src/query/management/tests/it/main.rs index 9f4f238e518d8..15a60a9b8bd1c 100644 --- a/src/query/management/tests/it/main.rs +++ b/src/query/management/tests/it/main.rs @@ -15,10 +15,10 @@ #![feature(duration_constructors)] #![allow(clippy::uninlined_format_args)] -mod cluster; mod quota; mod role; mod setting; mod stage; mod udf; mod user; +mod warehouse; diff --git a/src/query/management/tests/it/warehouse.rs b/src/query/management/tests/it/warehouse.rs new file mode 100644 index 0000000000000..0531a735d6205 --- /dev/null +++ b/src/query/management/tests/it/warehouse.rs @@ -0,0 +1,409 @@ +// Copyright 2021 Datafuse Labs +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::sync::Arc; +use std::time::Duration; + +use databend_common_base::base::tokio; +use databend_common_base::base::GlobalUniqName; +use databend_common_exception::Result; +use databend_common_management::*; +use databend_common_meta_embedded::MemMeta; +use databend_common_meta_kvapi::kvapi::KVApi; +use databend_common_meta_store::MetaStore; +use databend_common_meta_types::seq_value::SeqV; +use databend_common_meta_types::{MatchSeq, MatchSeqExt, NodeInfo}; +use databend_common_meta_types::NodeType; + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn test_empty_id_with_self_managed() -> Result<()> { + let (_kv, warehouse_manager, _nodes) = nodes(Duration::from_secs(60), 0).await?; + + let mut node = system_managed_node(&GlobalUniqName::unique()); + node.node_type = NodeType::SelfManaged; + node.warehouse_id = String::new(); + node.cluster_id = String::from("test_cluster_id"); + let res = warehouse_manager.add_node(node).await; + + assert!(res.is_err()); + assert_eq!(res.unwrap_err().code(), 2403); + + let mut node = system_managed_node(&GlobalUniqName::unique()); + node.node_type = NodeType::SelfManaged; + node.cluster_id = String::new(); + node.warehouse_id = String::from("test_cluster_id"); + let res = warehouse_manager.add_node(node).await; + + assert!(res.is_err()); + assert_eq!(res.unwrap_err().code(), 2403); + + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn test_successfully_add_self_managed_node() -> Result<()> { + let (kv, warehouse_manager, _nodes) = nodes(Duration::from_secs(60), 0).await?; + + let mut node_info_1 = self_managed_node("test_node_1"); + warehouse_manager.add_node(node_info_1.clone()).await?; + let node_key = "__fd_clusters_v5/test%2dtenant%2did/online_nodes/test_node_1"; + assert_key_value(&kv, node_key, serde_json::to_vec(&node_info_1)?).await; + + let warehouse_key = "__fd_clusters_v5/test%2dtenant%2did/online_clusters/test%2dcluster%2did/test%2dcluster%2did/test_node_1"; + + node_info_1.cluster_id = String::new(); + node_info_1.warehouse_id = String::new(); + + assert_key_value(&kv, warehouse_key, serde_json::to_vec(&node_info_1)?).await; + + let info_key = "__fd_warehouses/v1/test%2dtenant%2did/test%2dcluster%2did"; + assert_key_value( + &kv, + info_key, + serde_json::to_vec(&WarehouseInfo::SelfManaged)?, + ) + .await; + + let mut node_info_2 = self_managed_node("test_node_2"); + warehouse_manager.add_node(node_info_2.clone()).await?; + + let node_key = "__fd_clusters_v5/test%2dtenant%2did/online_nodes/test_node_2"; + assert_key_value(&kv, node_key, serde_json::to_vec(&node_info_2)?).await; + + let warehouse_key = "__fd_clusters_v5/test%2dtenant%2did/online_clusters/test%2dcluster%2did/test%2dcluster%2did/test_node_2"; + + node_info_2.cluster_id = String::new(); + node_info_2.warehouse_id = String::new(); + + assert_key_value(&kv, warehouse_key, serde_json::to_vec(&node_info_2)?).await; + + let info_key = "__fd_warehouses/v1/test%2dtenant%2did/test%2dcluster%2did"; + assert_key_value( + &kv, + info_key, + serde_json::to_vec(&WarehouseInfo::SelfManaged)?, + ) + .await; + + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn test_already_exists_add_self_managed_node() -> Result<()> { + let (kv, warehouse_manager, nodes) = nodes(Duration::from_secs(60), 1).await?; + + let node_info = self_managed_node("test_node_1"); + warehouse_manager.add_node(node_info.clone()).await?; + + let node_key = "__fd_clusters_v5/test%2dtenant%2did/online_nodes/test_node_1"; + assert_key_value(&kv, node_key, serde_json::to_vec(&node_info)?).await; + + // add already exists self-managed node and get failure + match warehouse_manager.add_node(node_info.clone()).await { + Ok(_) => panic!("Already exists add node must be return Err."), + Err(cause) => assert_eq!(cause.code(), 2402), + } + + // add already exists system-managed node and get failure + match warehouse_manager + .add_node(self_managed_node(&nodes[0])) + .await + { + Ok(_) => panic!("Already exists add node must be return Err."), + Err(cause) => assert_eq!(cause.code(), 2402), + } + + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn test_successfully_get_self_managed_nodes() -> Result<()> { + let (_kv, warehouse_manager, _nodes) = nodes(Duration::from_mins(60), 0).await?; + + let get_nodes = warehouse_manager.get_nodes("test-cluster-id", "test-cluster-id"); + + assert_eq!(get_nodes.await?, vec![]); + + let node_1 = self_managed_node("node_1"); + warehouse_manager.add_node(node_1.clone()).await?; + + let get_nodes = warehouse_manager.get_nodes("test-cluster-id", "test-cluster-id"); + + assert_eq!(get_nodes.await?, vec![node_1.clone()]); + + let node_2 = self_managed_node("node_2"); + warehouse_manager.add_node(node_2.clone()).await?; + + let get_nodes = warehouse_manager.get_nodes("test-cluster-id", "test-cluster-id"); + + assert_eq!(get_nodes.await?, vec![node_1, node_2]); + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn test_successfully_drop_self_managed_node() -> Result<()> { + let (_, warehouse_manager, _nodes) = nodes(Duration::from_mins(60), 0).await?; + + let node_info = self_managed_node("test_node"); + warehouse_manager.add_node(node_info.clone()).await?; + + let get_nodes = warehouse_manager.get_nodes("test-cluster-id", "test-cluster-id"); + + assert_eq!(get_nodes.await?, vec![node_info.clone()]); + + warehouse_manager.drop_node(node_info.id).await?; + + let get_nodes = warehouse_manager.get_nodes("test-cluster-id", "test-cluster-id"); + + assert_eq!(get_nodes.await?, vec![]); + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn test_unknown_node_drop_self_managed_node() -> Result<()> { + let (_, warehouse_manager, _nodes) = nodes(Duration::from_mins(60), 0).await?; + + match warehouse_manager.drop_node(String::from("UNKNOWN_ID")).await { + Ok(_) => { /*panic!("Unknown node drop node must be return Err.")*/ } + Err(cause) => assert_eq!(cause.code(), 2401), + } + + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn test_successfully_heartbeat_self_managed_node() -> Result<()> { + let (kv, warehouse_manager, _nodes) = nodes(Duration::from_mins(60), 0).await?; + + let mut node_info = self_managed_node("test_node"); + let seq = warehouse_manager.add_node(node_info.clone()).await?; + + let info_key = "__fd_clusters_v5/test%2dtenant%2did/online_nodes/test_node"; + assert_key_value(&kv, info_key, serde_json::to_vec(&node_info)?).await; + assert_key_expire(&kv, info_key, Duration::from_mins(50)).await; + + let warehouse_key = "__fd_clusters_v5/test%2dtenant%2did/online_clusters/test%2dcluster%2did/test%2dcluster%2did/test_node"; + let mut warehouse_node = node_info.clone(); + warehouse_node.cluster_id = String::new(); + warehouse_node.warehouse_id = String::new(); + assert_key_value(&kv, warehouse_key, serde_json::to_vec(&warehouse_node)?).await; + assert_key_expire(&kv, warehouse_key, Duration::from_mins(50)).await; + + let warehouse_info_key = "__fd_warehouses/v1/test%2dtenant%2did/test%2dcluster%2did"; + let info = serde_json::to_vec(&WarehouseInfo::SelfManaged)?; + assert_key_value(&kv, warehouse_info_key, info.clone()).await; + assert_key_expire(&kv, warehouse_info_key, Duration::from_mins(50)).await; + + warehouse_manager.heartbeat(&mut node_info, seq).await?; + assert_key_value(&kv, warehouse_info_key, info.clone()).await; + assert_key_value(&kv, info_key, serde_json::to_vec(&node_info)?).await; + assert_key_value(&kv, warehouse_key, serde_json::to_vec(&warehouse_node)?).await; + assert_key_seq(&kv, info_key, MatchSeq::GE(seq + 3)).await; + assert_key_seq(&kv, warehouse_key, MatchSeq::GE(seq + 3)).await; + assert_key_expire(&kv, info_key, Duration::from_mins(50)).await; + assert_key_expire(&kv, warehouse_key, Duration::from_mins(50)).await; + assert_key_expire(&kv, warehouse_info_key, Duration::from_mins(50)).await; + + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn test_successfully_create_system_managed_warehouse() -> Result<()> { + let (_, cluster_mgr, nodes) = nodes(Duration::from_mins(30), 2).await?; + + let create_warehouse = cluster_mgr.create_warehouse("test_warehouse".to_string(), vec![ + SelectedNode::Random(None), + SelectedNode::Random(None), + ]); + + create_warehouse.await?; + + let get_warehouse_nodes = cluster_mgr.get_nodes("test_warehouse", "test_warehouse"); + + let warehouse_nodes = get_warehouse_nodes.await?; + + assert_eq!(warehouse_nodes.len(), 2); + assert_eq!(warehouse_nodes.last().map(|x| &x.id), nodes.last()); + assert_eq!(warehouse_nodes.first().map(|x| &x.id), nodes.first()); + + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn test_create_duplicated_warehouse() -> Result<()> { + let (_, cluster_mgr, _nodes) = nodes(Duration::from_mins(30), 2).await?; + + let create_warehouse = + cluster_mgr.create_warehouse("test_warehouse".to_string(), vec![SelectedNode::Random( + None, + )]); + + create_warehouse.await?; + + let create_warehouse = + cluster_mgr.create_warehouse("test_warehouse".to_string(), vec![SelectedNode::Random( + None, + )]); + + let res = create_warehouse.await; + assert!(res.is_err()); + assert_eq!(res.unwrap_err().code(), 2405); + + let create_warehouse = + cluster_mgr.create_warehouse("test_warehouse_2".to_string(), vec![SelectedNode::Random( + None, + )]); + + create_warehouse.await?; + + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn test_create_warehouse_with_self_manage() -> Result<()> { + let (_, cluster_mgr, _nodes) = nodes(Duration::from_mins(30), 2).await?; + + // Self manage node online + let mut self_manage_node_1 = system_managed_node("self_manage_node_1"); + self_manage_node_1.cluster_id = String::from("test_warehouse"); + self_manage_node_1.warehouse_id = String::from("test_warehouse"); + cluster_mgr.add_node(self_manage_node_1).await?; + + let create_warehouse = + cluster_mgr.create_warehouse(String::from("test_warehouse"), vec![SelectedNode::Random( + None, + )]); + + let res = create_warehouse.await; + + assert!(res.is_err()); + assert_eq!(res.unwrap_err().code(), 2408); + + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn test_create_warehouse_with_no_resources() -> Result<()> { + let (_, cluster_mgr, _nodes) = nodes(Duration::from_mins(30), 2).await?; + + let create_warehouse = + cluster_mgr.create_warehouse("test_warehouse_1".to_string(), vec![SelectedNode::Random( + None, + )]); + + create_warehouse.await?; + + let create_warehouse = + cluster_mgr.create_warehouse("test_warehouse_2".to_string(), vec![SelectedNode::Random( + None, + )]); + + create_warehouse.await?; + + let create_warehouse = + cluster_mgr.create_warehouse("test_warehouse_3".to_string(), vec![SelectedNode::Random( + None, + )]); + + let res = create_warehouse.await; + + assert!(res.is_err()); + assert_eq!(res.unwrap_err().code(), 2404); + + Ok(()) +} + +fn system_managed_node(id: &str) -> NodeInfo { + NodeInfo { + id: id.to_string(), + secret: "".to_string(), + cpu_nums: 0, + version: 0, + http_address: "".to_string(), + flight_address: "".to_string(), + discovery_address: "".to_string(), + binary_version: "".to_string(), + node_type: NodeType::SystemManaged, + cluster_id: "".to_string(), + warehouse_id: "".to_string(), + } +} + +fn self_managed_node(node_id: &str) -> NodeInfo { + NodeInfo { + id: String::from(node_id), + secret: "".to_string(), + cpu_nums: 0, + version: 0, + http_address: "ip3:port".to_string(), + flight_address: String::from("ip:port"), + discovery_address: "ip2:port".to_string(), + binary_version: "binary_version".to_string(), + node_type: NodeType::SelfManaged, + cluster_id: "test-cluster-id".to_string(), + warehouse_id: "test-cluster-id".to_string(), + } +} + +async fn nodes(lift: Duration, size: usize) -> Result<(MetaStore, WarehouseMgr, Vec)> { + let (kv_api, cluster_manager) = new_cluster_api(lift).await?; + + let mut nodes = Vec::with_capacity(size); + for _index in 0..size { + let name = GlobalUniqName::unique(); + cluster_manager.add_node(system_managed_node(&name)).await?; + nodes.push(name); + } + + Ok((kv_api, cluster_manager, nodes)) +} + +async fn new_cluster_api(lift: Duration) -> Result<(MetaStore, WarehouseMgr)> { + let test_api = MetaStore::L(Arc::new(MemMeta::default())); + let cluster_manager = WarehouseMgr::create(test_api.clone(), "test-tenant-id", lift)?; + Ok((test_api, cluster_manager)) +} + +async fn assert_key_value(kv: &MetaStore, key: &str, value: Vec) { + let reply = kv.get_kv(key).await.unwrap(); + + match reply { + Some(SeqV { data, .. }) => { + assert_eq!(data, value); + } + catch => panic!("GetKVActionReply{:?}", catch), + } +} + +async fn assert_key_seq(kv: &MetaStore, key: &str, expect: MatchSeq) { + let reply = kv.get_kv(key).await.unwrap(); + + match reply { + Some(SeqV { seq, .. }) => { + assert!(expect.match_seq(seq).is_ok()); + } + catch => panic!("GetKVActionReply{:?}", catch), + } +} + +async fn assert_key_expire(kv: &MetaStore, key: &str, lift: Duration) { + let reply = kv.get_kv(key).await.unwrap(); + + match reply { + Some(SeqV { meta: Some(meta), .. }) => { + assert!(meta.get_expire_at_ms().unwrap() >= lift.as_millis() as u64); + } + catch => panic!("GetKVActionReply{:?}", catch), + } +} diff --git a/src/query/service/src/clusters/cluster.rs b/src/query/service/src/clusters/cluster.rs index f315c37fe7b2a..493e41b5682a9 100644 --- a/src/query/service/src/clusters/cluster.rs +++ b/src/query/service/src/clusters/cluster.rs @@ -43,6 +43,7 @@ use databend_common_management::WarehouseMgr; use databend_common_meta_store::MetaStore; use databend_common_meta_store::MetaStoreProvider; use databend_common_meta_types::NodeInfo; +use databend_common_meta_types::NodeType; use databend_common_metrics::cluster::*; use futures::future::select; use futures::future::Either; @@ -439,6 +440,8 @@ impl ClusterDiscovery { node_info.cluster_id = self.cluster_id.clone(); node_info.warehouse_id = self.cluster_id.clone(); + node_info.node_type = NodeType::SelfManaged; + self.drop_invalid_nodes(&node_info).await?; match self.api_provider.add_node(node_info.clone()).await { Ok(seq) => self.start_heartbeat(node_info, seq).await, From b4ab6823790b8ef325c0a0cf0e392aec1be84db2 Mon Sep 17 00:00:00 2001 From: zhang2014 Date: Tue, 17 Dec 2024 17:37:11 +0800 Subject: [PATCH 03/53] feat(cluster): support custom management cluster --- .../management/src/warehouse/warehouse_mgr.rs | 67 ++++--- src/query/management/tests/it/warehouse.rs | 167 ++++++++++++++++-- 2 files changed, 189 insertions(+), 45 deletions(-) diff --git a/src/query/management/src/warehouse/warehouse_mgr.rs b/src/query/management/src/warehouse/warehouse_mgr.rs index 43e2446e344df..cdeac8a69771d 100644 --- a/src/query/management/src/warehouse/warehouse_mgr.rs +++ b/src/query/management/src/warehouse/warehouse_mgr.rs @@ -12,7 +12,6 @@ // See the License for the specific language governing permissions and // limitations under the License. -use std::collections::HashSet; use std::time::Duration; use databend_common_base::base::escape_for_key; @@ -187,12 +186,12 @@ impl WarehouseMgr { WarehouseInfo::SelfManaged => match response.responses.first() { // already exists node. Some(TxnOpResponse { - response: - Some(Response::Get(TxnGetResponse { - value: Some(value), - .. - })), - }) if value.seq != 0 => Ok(response), + response: + Some(Response::Get(TxnGetResponse { + value: Some(value), + .. + })), + }) if value.seq != 0 => Ok(response), _ => { log::info!("Self-managed warehouse has already been created by other nodes; attempt to join it. Retry count: {}", retry_count); retry_count += 1; @@ -278,8 +277,8 @@ impl WarehouseMgr { } Ok(response) => match response.responses.last() { Some(TxnOpResponse { - response: Some(Response::Get(TxnGetResponse { value: Some(v), .. })), - }) => Ok(v.seq), + response: Some(Response::Get(TxnGetResponse { value: Some(v), .. })), + }) => Ok(v.seq), _ => Err(ErrorCode::MetaServiceError("Meta insert failure.")), }, } @@ -303,8 +302,8 @@ impl WarehouseMgr { } Ok(response) => match response.responses.last() { Some(TxnOpResponse { - response: Some(Response::Get(TxnGetResponse { value: Some(v), .. })), - }) => Ok(v.seq), + response: Some(Response::Get(TxnGetResponse { value: Some(v), .. })), + }) => Ok(v.seq), _ => Err(ErrorCode::MetaServiceError("Meta insert failure.")), }, } @@ -314,8 +313,8 @@ impl WarehouseMgr { match reply.responses.first() { None => self.leave_cluster(node, 0).await, Some(TxnOpResponse { - response: Some(Response::Get(res)), - }) => match &res.value { + response: Some(Response::Get(res)), + }) => match &res.value { None => self.leave_cluster(node, 0).await, Some(value) => { let node_info = serde_json::from_slice::(&value.data)?; @@ -511,39 +510,37 @@ impl WarehouseApi for WarehouseMgr { loop { let mut selected_nodes = Vec::with_capacity(nodes.len()); - let mut allocated_nodes = HashSet::with_capacity(nodes.len()); // get online nodes let online_nodes = self.metastore.prefix_list_kv(&self.node_key_prefix).await?; - for select_node in &nodes { - match select_node { - SelectedNode::Random(Some(_)) => { + let mut select_queue = nodes.clone(); + for (_, v) in &online_nodes { + match select_queue.last() { + None => { + break; + } + Some(SelectedNode::Random(Some(_))) => { return Err(ErrorCode::Unimplemented( "Custom instance types are not supported.", )); } - // select random node - SelectedNode::Random(None) => { - for (_, v) in &online_nodes { - let mut node_info = serde_json::from_slice::(&v.data)?; - - if node_info.warehouse_id.is_empty() - && node_info.cluster_id.is_empty() - && !allocated_nodes.contains(&node_info.id) - { - node_info.cluster_id = warehouse.clone(); - node_info.warehouse_id = warehouse.clone(); - allocated_nodes.insert(node_info.id.clone()); - selected_nodes.push((v.seq, node_info)); - break; - } + Some(SelectedNode::Random(None)) => { + // select random node + + let mut node_info = serde_json::from_slice::(&v.data)?; + + if node_info.warehouse_id.is_empty() && node_info.cluster_id.is_empty() { + node_info.cluster_id = warehouse.clone(); + node_info.warehouse_id = warehouse.clone(); + selected_nodes.push((v.seq, node_info)); + select_queue.pop(); } } } } - if selected_nodes.len() != nodes.len() { + if !select_queue.is_empty() { return Err(ErrorCode::NoResourcesAvailable( "Failed to create warehouse, reason: no resources available", )); @@ -603,8 +600,8 @@ impl WarehouseApi for WarehouseMgr { res if res.success => Ok(()), res => match res.responses.last() { Some(TxnOpResponse { - response: Some(Response::Get(res)), - }) => { + response: Some(Response::Get(res)), + }) => { if matches!(&res.value, Some(v) if v.seq != 0) { return Err(ErrorCode::WarehouseAlreadyExists( "Warehouse already exists", diff --git a/src/query/management/tests/it/warehouse.rs b/src/query/management/tests/it/warehouse.rs index 0531a735d6205..97404bd4166c0 100644 --- a/src/query/management/tests/it/warehouse.rs +++ b/src/query/management/tests/it/warehouse.rs @@ -23,7 +23,9 @@ use databend_common_meta_embedded::MemMeta; use databend_common_meta_kvapi::kvapi::KVApi; use databend_common_meta_store::MetaStore; use databend_common_meta_types::seq_value::SeqV; -use databend_common_meta_types::{MatchSeq, MatchSeqExt, NodeInfo}; +use databend_common_meta_types::MatchSeq; +use databend_common_meta_types::MatchSeqExt; +use databend_common_meta_types::NodeInfo; use databend_common_meta_types::NodeType; #[tokio::test(flavor = "multi_thread", worker_threads = 1)] @@ -73,7 +75,7 @@ async fn test_successfully_add_self_managed_node() -> Result<()> { info_key, serde_json::to_vec(&WarehouseInfo::SelfManaged)?, ) - .await; + .await; let mut node_info_2 = self_managed_node("test_node_2"); warehouse_manager.add_node(node_info_2.clone()).await?; @@ -94,7 +96,7 @@ async fn test_successfully_add_self_managed_node() -> Result<()> { info_key, serde_json::to_vec(&WarehouseInfo::SelfManaged)?, ) - .await; + .await; Ok(()) } @@ -174,7 +176,10 @@ async fn test_successfully_drop_self_managed_node() -> Result<()> { async fn test_unknown_node_drop_self_managed_node() -> Result<()> { let (_, warehouse_manager, _nodes) = nodes(Duration::from_mins(60), 0).await?; - match warehouse_manager.drop_node(String::from("UNKNOWN_ID")).await { + match warehouse_manager + .drop_node(String::from("UNKNOWN_ID")) + .await + { Ok(_) => { /*panic!("Unknown node drop node must be return Err.")*/ } Err(cause) => assert_eq!(cause.code(), 2401), } @@ -218,24 +223,152 @@ async fn test_successfully_heartbeat_self_managed_node() -> Result<()> { Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn test_empty_system_managed_warehouse() -> Result<()> { + let (_, cluster_mgr, _nodes) = nodes(Duration::from_mins(30), 2).await?; + + let create_warehouse = cluster_mgr.create_warehouse(String::new(), vec![]); + + assert_eq!(create_warehouse.await.unwrap_err().code(), 2403); + + let create_warehouse = cluster_mgr.create_warehouse(String::from("test"), vec![]); + + assert_eq!(create_warehouse.await.unwrap_err().code(), 2408); + + let create_warehouse = + cluster_mgr.create_warehouse(String::from("test"), vec![SelectedNode::Random(Some( + String::from("XLargeNode"), + ))]); + + assert_eq!(create_warehouse.await.unwrap_err().code(), 1002); + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 1)] async fn test_successfully_create_system_managed_warehouse() -> Result<()> { - let (_, cluster_mgr, nodes) = nodes(Duration::from_mins(30), 2).await?; + let (kv, warehouse_manager, nodes) = nodes(Duration::from_mins(30), 2).await?; - let create_warehouse = cluster_mgr.create_warehouse("test_warehouse".to_string(), vec![ + for node in &nodes { + let online_node = format!("__fd_clusters_v5/test%2dtenant%2did/online_nodes/{}", node); + assert_key_seq(&kv, &online_node, MatchSeq::GE(1)).await; + let warehouse_node = format!("__fd_clusters_v5/test%2dtenant%2did/online_clusters/test%2dcluster%2did/test%2dcluster%2did/{}", node); + assert_no_key(&kv, &warehouse_node).await; + } + + let create_warehouse = warehouse_manager.create_warehouse("test_warehouse".to_string(), vec![ SelectedNode::Random(None), SelectedNode::Random(None), ]); create_warehouse.await?; - let get_warehouse_nodes = cluster_mgr.get_nodes("test_warehouse", "test_warehouse"); + for node in &nodes { + let online_node = format!("__fd_clusters_v5/test%2dtenant%2did/online_nodes/{}", node); + assert_key_seq(&kv, &online_node, MatchSeq::GE(1)).await; + let warehouse_node = format!( + "__fd_clusters_v5/test%2dtenant%2did/online_clusters/test_warehouse/test_warehouse/{}", + node + ); + assert_key_seq(&kv, &warehouse_node, MatchSeq::GE(1)).await; + } + + let get_warehouse_nodes = warehouse_manager.get_nodes("test_warehouse", "test_warehouse"); let warehouse_nodes = get_warehouse_nodes.await?; assert_eq!(warehouse_nodes.len(), 2); - assert_eq!(warehouse_nodes.last().map(|x| &x.id), nodes.last()); - assert_eq!(warehouse_nodes.first().map(|x| &x.id), nodes.first()); + + for warehouse_node in &warehouse_nodes { + assert!(nodes.contains(&warehouse_node.id)); + assert_eq!(warehouse_node.cluster_id, "test_warehouse"); + assert_eq!(warehouse_node.warehouse_id, "test_warehouse"); + } + + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn test_create_system_managed_warehouse_with_offline_node() -> Result<()> { + let (_, warehouse_manager, mut nodes) = nodes(Duration::from_mins(30), 4).await?; + + // mock node offline + warehouse_manager.drop_node(nodes[0].to_string()).await?; + + let create_warehouse = warehouse_manager.create_warehouse("test_warehouse".to_string(), vec![ + SelectedNode::Random(None), + SelectedNode::Random(None), + SelectedNode::Random(None), + SelectedNode::Random(None), + ]); + + // no resources available + assert_eq!(create_warehouse.await.unwrap_err().code(), 2404); + + let create_warehouse = warehouse_manager.create_warehouse("test_warehouse".to_string(), vec![ + SelectedNode::Random(None), + SelectedNode::Random(None), + SelectedNode::Random(None), + ]); + + create_warehouse.await?; + + let get_warehouse_nodes = warehouse_manager.get_nodes("test_warehouse", "test_warehouse"); + + let warehouse_nodes = get_warehouse_nodes.await?; + + assert_eq!(warehouse_nodes.len(), 3); + + nodes.remove(0); + for warehouse_node in &warehouse_nodes { + assert!(nodes.contains(&warehouse_node.id)); + assert_eq!(warehouse_node.cluster_id, "test_warehouse"); + assert_eq!(warehouse_node.warehouse_id, "test_warehouse"); + } + + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn test_create_system_managed_warehouse_with_online_node() -> Result<()> { + let (_, warehouse_manager, mut nodes) = nodes(Duration::from_mins(30), 3).await?; + + let create_warehouse = warehouse_manager.create_warehouse("test_warehouse".to_string(), vec![ + SelectedNode::Random(None), + SelectedNode::Random(None), + SelectedNode::Random(None), + SelectedNode::Random(None), + ]); + + // no resources available + assert_eq!(create_warehouse.await.unwrap_err().code(), 2404); + + // mock node online + let new_node = GlobalUniqName::unique(); + warehouse_manager + .add_node(system_managed_node(&new_node)) + .await?; + + let create_warehouse = warehouse_manager.create_warehouse("test_warehouse".to_string(), vec![ + SelectedNode::Random(None), + SelectedNode::Random(None), + SelectedNode::Random(None), + SelectedNode::Random(None), + ]); + + create_warehouse.await?; + + let get_warehouse_nodes = warehouse_manager.get_nodes("test_warehouse", "test_warehouse"); + + let warehouse_nodes = get_warehouse_nodes.await?; + + assert_eq!(warehouse_nodes.len(), 4); + + nodes.push(new_node); + for warehouse_node in &warehouse_nodes { + assert!(nodes.contains(&warehouse_node.id)); + assert_eq!(warehouse_node.cluster_id, "test_warehouse"); + assert_eq!(warehouse_node.warehouse_id, "test_warehouse"); + } Ok(()) } @@ -375,6 +508,18 @@ async fn new_cluster_api(lift: Duration) -> Result<(MetaStore, WarehouseMgr)> { Ok((test_api, cluster_manager)) } +async fn assert_no_key(kv: &MetaStore, key: &str) { + let reply = kv.get_kv(key).await.unwrap(); + + match reply { + None => {} + Some(v) => match v.seq { + 0 => {} + _ => panic!("assert_no_key {}", key), + }, + } +} + async fn assert_key_value(kv: &MetaStore, key: &str, value: Vec) { let reply = kv.get_kv(key).await.unwrap(); @@ -401,7 +546,9 @@ async fn assert_key_expire(kv: &MetaStore, key: &str, lift: Duration) { let reply = kv.get_kv(key).await.unwrap(); match reply { - Some(SeqV { meta: Some(meta), .. }) => { + Some(SeqV { + meta: Some(meta), .. + }) => { assert!(meta.get_expire_at_ms().unwrap() >= lift.as_millis() as u64); } catch => panic!("GetKVActionReply{:?}", catch), From 3696d5746d4b79b41c904f686e6a8fe57a004f66 Mon Sep 17 00:00:00 2001 From: zhang2014 Date: Wed, 18 Dec 2024 00:12:53 +0800 Subject: [PATCH 04/53] feat(cluster): support custom management cluster --- .../management/src/warehouse/warehouse_api.rs | 4 +- .../management/src/warehouse/warehouse_mgr.rs | 18 ++++- src/query/management/tests/it/warehouse.rs | 73 ++++++++++++------- 3 files changed, 63 insertions(+), 32 deletions(-) diff --git a/src/query/management/src/warehouse/warehouse_api.rs b/src/query/management/src/warehouse/warehouse_api.rs index aaa2e5ef4d3a3..ccd062c18c1cb 100644 --- a/src/query/management/src/warehouse/warehouse_api.rs +++ b/src/query/management/src/warehouse/warehouse_api.rs @@ -12,6 +12,8 @@ // See the License for the specific language governing permissions and // limitations under the License. +use std::collections::HashMap; + use databend_common_exception::Result; use databend_common_meta_types::NodeInfo; @@ -30,7 +32,7 @@ pub struct SystemManagedInfo { pub id: String, pub status: String, pub display_name: String, - pub clusters: Vec>, + pub clusters: HashMap>, } /// Databend-query cluster management API diff --git a/src/query/management/src/warehouse/warehouse_mgr.rs b/src/query/management/src/warehouse/warehouse_mgr.rs index cdeac8a69771d..2e0936b0bfecc 100644 --- a/src/query/management/src/warehouse/warehouse_mgr.rs +++ b/src/query/management/src/warehouse/warehouse_mgr.rs @@ -12,6 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. +use std::collections::HashMap; use std::time::Duration; use databend_common_base::base::escape_for_key; @@ -181,7 +182,9 @@ impl WarehouseMgr { Some(value) if value.seq == 0 => Ok(response), Some(value) => match serde_json::from_slice(&value.data)? { WarehouseInfo::SystemManaged(_) => { - Err(ErrorCode::WarehouseAlreadyExists("")) + Err(ErrorCode::WarehouseAlreadyExists( + "Already exists same name system-managed warehouse.", + )) } WarehouseInfo::SelfManaged => match response.responses.first() { // already exists node. @@ -439,6 +442,14 @@ impl WarehouseApi for WarehouseMgr { ))); }; + if serde_json::from_slice::(&value.data)? + == WarehouseInfo::SelfManaged + { + return Err(ErrorCode::InvalidWarehouse( + "Cannot drop self-managed warehouse", + )); + } + if fetch_reply.responses.len() != delete_txn.if_then.len() { // TODO: maybe auto retry? return Err(ErrorCode::WarehouseOperateConflict("Missing node info in online nodes list. It's possible that some nodes offline during the drop warehouse. You may try the operation again.")); @@ -531,8 +542,8 @@ impl WarehouseApi for WarehouseMgr { let mut node_info = serde_json::from_slice::(&v.data)?; if node_info.warehouse_id.is_empty() && node_info.cluster_id.is_empty() { - node_info.cluster_id = warehouse.clone(); node_info.warehouse_id = warehouse.clone(); + node_info.cluster_id = String::from("default"); selected_nodes.push((v.seq, node_info)); select_queue.pop(); } @@ -583,6 +594,7 @@ impl WarehouseApi for WarehouseMgr { self.warehouse_key_prefix, escape_for_key(&warehouse)? ); + txn.condition .push(map_condition(&warehouse_key, MatchSeq::Exact(0))); txn.if_then.push(TxnOp::put( @@ -591,7 +603,7 @@ impl WarehouseApi for WarehouseMgr { id: GlobalUniqName::unique(), status: "Running".to_string(), display_name: warehouse.clone(), - clusters: vec![nodes.clone()], + clusters: HashMap::from([(String::from("default"), nodes.clone())]), }))?, )); txn.else_then.push(TxnOp::get(warehouse_key)); diff --git a/src/query/management/tests/it/warehouse.rs b/src/query/management/tests/it/warehouse.rs index 97404bd4166c0..26fcd2a13a828 100644 --- a/src/query/management/tests/it/warehouse.rs +++ b/src/query/management/tests/it/warehouse.rs @@ -187,6 +187,20 @@ async fn test_unknown_node_drop_self_managed_node() -> Result<()> { Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn test_drop_self_managed_warehouse() -> Result<()> { + let (_, warehouse_manager, _nodes) = nodes(Duration::from_mins(60), 0).await?; + + let node_info = self_managed_node("test_node"); + warehouse_manager.add_node(node_info.clone()).await?; + + let drop_warehouse = warehouse_manager.drop_warehouse(String::from("test-cluster-id")); + + assert_eq!(drop_warehouse.await.unwrap_err().code(), 2403); + + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 1)] async fn test_successfully_heartbeat_self_managed_node() -> Result<()> { let (kv, warehouse_manager, _nodes) = nodes(Duration::from_mins(60), 0).await?; @@ -251,7 +265,10 @@ async fn test_successfully_create_system_managed_warehouse() -> Result<()> { for node in &nodes { let online_node = format!("__fd_clusters_v5/test%2dtenant%2did/online_nodes/{}", node); assert_key_seq(&kv, &online_node, MatchSeq::GE(1)).await; - let warehouse_node = format!("__fd_clusters_v5/test%2dtenant%2did/online_clusters/test%2dcluster%2did/test%2dcluster%2did/{}", node); + let warehouse_node = format!( + "__fd_clusters_v5/test%2dtenant%2did/online_clusters/test%2dcluster%2did/default/{}", + node + ); assert_no_key(&kv, &warehouse_node).await; } @@ -266,13 +283,13 @@ async fn test_successfully_create_system_managed_warehouse() -> Result<()> { let online_node = format!("__fd_clusters_v5/test%2dtenant%2did/online_nodes/{}", node); assert_key_seq(&kv, &online_node, MatchSeq::GE(1)).await; let warehouse_node = format!( - "__fd_clusters_v5/test%2dtenant%2did/online_clusters/test_warehouse/test_warehouse/{}", + "__fd_clusters_v5/test%2dtenant%2did/online_clusters/test_warehouse/default/{}", node ); assert_key_seq(&kv, &warehouse_node, MatchSeq::GE(1)).await; } - let get_warehouse_nodes = warehouse_manager.get_nodes("test_warehouse", "test_warehouse"); + let get_warehouse_nodes = warehouse_manager.get_nodes("test_warehouse", "default"); let warehouse_nodes = get_warehouse_nodes.await?; @@ -280,7 +297,7 @@ async fn test_successfully_create_system_managed_warehouse() -> Result<()> { for warehouse_node in &warehouse_nodes { assert!(nodes.contains(&warehouse_node.id)); - assert_eq!(warehouse_node.cluster_id, "test_warehouse"); + assert_eq!(warehouse_node.cluster_id, "default"); assert_eq!(warehouse_node.warehouse_id, "test_warehouse"); } @@ -312,7 +329,7 @@ async fn test_create_system_managed_warehouse_with_offline_node() -> Result<()> create_warehouse.await?; - let get_warehouse_nodes = warehouse_manager.get_nodes("test_warehouse", "test_warehouse"); + let get_warehouse_nodes = warehouse_manager.get_nodes("test_warehouse", "default"); let warehouse_nodes = get_warehouse_nodes.await?; @@ -321,7 +338,7 @@ async fn test_create_system_managed_warehouse_with_offline_node() -> Result<()> nodes.remove(0); for warehouse_node in &warehouse_nodes { assert!(nodes.contains(&warehouse_node.id)); - assert_eq!(warehouse_node.cluster_id, "test_warehouse"); + assert_eq!(warehouse_node.cluster_id, "default"); assert_eq!(warehouse_node.warehouse_id, "test_warehouse"); } @@ -357,7 +374,7 @@ async fn test_create_system_managed_warehouse_with_online_node() -> Result<()> { create_warehouse.await?; - let get_warehouse_nodes = warehouse_manager.get_nodes("test_warehouse", "test_warehouse"); + let get_warehouse_nodes = warehouse_manager.get_nodes("test_warehouse", "default"); let warehouse_nodes = get_warehouse_nodes.await?; @@ -366,7 +383,7 @@ async fn test_create_system_managed_warehouse_with_online_node() -> Result<()> { nodes.push(new_node); for warehouse_node in &warehouse_nodes { assert!(nodes.contains(&warehouse_node.id)); - assert_eq!(warehouse_node.cluster_id, "test_warehouse"); + assert_eq!(warehouse_node.cluster_id, "default"); assert_eq!(warehouse_node.warehouse_id, "test_warehouse"); } @@ -375,26 +392,24 @@ async fn test_create_system_managed_warehouse_with_online_node() -> Result<()> { #[tokio::test(flavor = "multi_thread", worker_threads = 1)] async fn test_create_duplicated_warehouse() -> Result<()> { - let (_, cluster_mgr, _nodes) = nodes(Duration::from_mins(30), 2).await?; + let (_, warehouse_manager, _nodes) = nodes(Duration::from_mins(30), 2).await?; - let create_warehouse = - cluster_mgr.create_warehouse("test_warehouse".to_string(), vec![SelectedNode::Random( - None, - )]); + let create_warehouse = warehouse_manager.create_warehouse("test_warehouse".to_string(), vec![ + SelectedNode::Random(None), + ]); create_warehouse.await?; - let create_warehouse = - cluster_mgr.create_warehouse("test_warehouse".to_string(), vec![SelectedNode::Random( - None, - )]); + let create_warehouse = warehouse_manager.create_warehouse("test_warehouse".to_string(), vec![ + SelectedNode::Random(None), + ]); let res = create_warehouse.await; assert!(res.is_err()); assert_eq!(res.unwrap_err().code(), 2405); - let create_warehouse = - cluster_mgr.create_warehouse("test_warehouse_2".to_string(), vec![SelectedNode::Random( + let create_warehouse = warehouse_manager + .create_warehouse("test_warehouse_2".to_string(), vec![SelectedNode::Random( None, )]); @@ -405,23 +420,20 @@ async fn test_create_duplicated_warehouse() -> Result<()> { #[tokio::test(flavor = "multi_thread", worker_threads = 1)] async fn test_create_warehouse_with_self_manage() -> Result<()> { - let (_, cluster_mgr, _nodes) = nodes(Duration::from_mins(30), 2).await?; + let (_, warehouse_manager, _nodes) = nodes(Duration::from_mins(30), 2).await?; // Self manage node online - let mut self_manage_node_1 = system_managed_node("self_manage_node_1"); + let mut self_manage_node_1 = self_managed_node("self_manage_node_1"); self_manage_node_1.cluster_id = String::from("test_warehouse"); self_manage_node_1.warehouse_id = String::from("test_warehouse"); - cluster_mgr.add_node(self_manage_node_1).await?; + warehouse_manager.add_node(self_manage_node_1).await?; - let create_warehouse = - cluster_mgr.create_warehouse(String::from("test_warehouse"), vec![SelectedNode::Random( + let create_warehouse = warehouse_manager + .create_warehouse(String::from("test_warehouse"), vec![SelectedNode::Random( None, )]); - let res = create_warehouse.await; - - assert!(res.is_err()); - assert_eq!(res.unwrap_err().code(), 2408); + assert_eq!(create_warehouse.await.unwrap_err().code(), 2405); Ok(()) } @@ -457,6 +469,11 @@ async fn test_create_warehouse_with_no_resources() -> Result<()> { Ok(()) } +// #[tokio::test(flavor = "multi_thread", worker_threads = 1)] +// async fn test_create_warehouse_with_no_resources() -> Result<()> { +// +// } + fn system_managed_node(id: &str) -> NodeInfo { NodeInfo { id: id.to_string(), From 2c88215dbd692c8c252dcf95e7abcaa087f604ce Mon Sep 17 00:00:00 2001 From: zhang2014 Date: Wed, 18 Dec 2024 10:27:12 +0800 Subject: [PATCH 05/53] feat(cluster): support custom management cluster --- src/meta/types/tests/it/cluster.rs | 2 + src/query/config/src/config.rs | 4 +- .../management/src/warehouse/warehouse_mgr.rs | 32 +++--- src/query/management/tests/it/warehouse.rs | 102 ++++++++++++++++-- 4 files changed, 106 insertions(+), 34 deletions(-) diff --git a/src/meta/types/tests/it/cluster.rs b/src/meta/types/tests/it/cluster.rs index 6a169ff1ee697..49a814472a936 100644 --- a/src/meta/types/tests/it/cluster.rs +++ b/src/meta/types/tests/it/cluster.rs @@ -25,6 +25,7 @@ fn test_node_info_ip_port() -> anyhow::Result<()> { flight_address: "1.2.3.4:123".to_string(), discovery_address: "4.5.6.7:456".to_string(), binary_version: "v0.8-binary-version".to_string(), + node_type: Default::default(), cluster_id: "".to_string(), warehouse_id: "".to_string(), }; @@ -48,6 +49,7 @@ fn test_serde_node_info() { flight_address: "1.2.3.4:123".to_string(), discovery_address: "4.5.6.7:456".to_string(), binary_version: "v0.8-binary-version".to_string(), + node_type: Default::default(), cluster_id: String::new(), warehouse_id: String::new(), }; diff --git a/src/query/config/src/config.rs b/src/query/config/src/config.rs index 7ecdeaeb39693..0c1d4b707b903 100644 --- a/src/query/config/src/config.rs +++ b/src/query/config/src/config.rs @@ -152,9 +152,7 @@ pub struct Config { /// when converted from inner config, all catalog configurations will store in `catalogs` #[clap(skip)] pub catalogs: HashMap, - - #[clap(flatten)] - pub resources_manager: ResourceManagerConfig, + // #[clap(flatten)] } #[derive(Subcommand, Default, Debug, Eq, PartialEq, Clone, Serialize, Deserialize)] diff --git a/src/query/management/src/warehouse/warehouse_mgr.rs b/src/query/management/src/warehouse/warehouse_mgr.rs index 2e0936b0bfecc..e76aa4b0d29d1 100644 --- a/src/query/management/src/warehouse/warehouse_mgr.rs +++ b/src/query/management/src/warehouse/warehouse_mgr.rs @@ -22,7 +22,6 @@ use databend_common_exception::ErrorCode; use databend_common_exception::Result; use databend_common_meta_kvapi::kvapi::KVApi; use databend_common_meta_store::MetaStore; -use databend_common_meta_types::protobuf::SeqV; use databend_common_meta_types::txn_op_response::Response; use databend_common_meta_types::ConditionResult; use databend_common_meta_types::MatchSeq; @@ -93,15 +92,6 @@ fn map_condition(k: &str, seq: MatchSeq) -> TxnCondition { } } -fn get_prev_value(res: Option<&TxnOpResponse>) -> Option<&SeqV> { - res.and_then(|response| response.response.as_ref()) - .and_then(|response| match response { - Response::Put(v) => v.prev_value.as_ref(), - Response::Delete(v) => v.prev_value.as_ref(), - _ => unreachable!(), - }) -} - impl WarehouseMgr { fn cluster_key(&self, node: &NodeInfo) -> Result { Ok(format!( @@ -394,6 +384,10 @@ impl WarehouseApi for WarehouseMgr { } async fn drop_warehouse(&self, warehouse: String) -> Result<()> { + if warehouse.is_empty() { + return Err(ErrorCode::InvalidWarehouse("Warehouse name is empty.")); + } + let drop_key = format!("{}/{}", self.meta_key_prefix, escape_for_key(&warehouse)?); loop { @@ -402,11 +396,12 @@ impl WarehouseApi for WarehouseMgr { let mut txn = TxnRequest::default(); for (node_key, _value) in values { - txn.if_then.push(TxnOp::get(format!( - "{}/{}", - self.node_key_prefix, - &node_key[drop_key.len() + 1..] - ))); + let suffix = &node_key[drop_key.len() + 1..]; + + if let Some((_, node)) = suffix.split_once('/') { + txn.if_then + .push(TxnOp::get(format!("{}/{}", self.node_key_prefix, node))); + } delete_txn.if_then.push(TxnOp::delete(node_key)); } @@ -680,10 +675,7 @@ impl WarehouseApi for WarehouseMgr { let res = self.metastore.transaction(txn).await?; - if res.success - && get_prev_value(res.responses.first()).is_some() - && get_prev_value(res.responses.last()).is_some() - { + if res.success { return Ok(()); } } @@ -701,7 +693,7 @@ impl WarehouseApi for WarehouseMgr { } async fn get_node_info(&self, node_id: &str) -> Result { - let node_key = format!("{}/{}", self.node_key_prefix, escape_for_key(&node_id)?); + let node_key = format!("{}/{}", self.node_key_prefix, escape_for_key(node_id)?); let node_info = self.metastore.get_kv(&node_key).await?; match node_info { None => Err(ErrorCode::NotFoundClusterNode("")), diff --git a/src/query/management/tests/it/warehouse.rs b/src/query/management/tests/it/warehouse.rs index 26fcd2a13a828..bd374a1e5bcc8 100644 --- a/src/query/management/tests/it/warehouse.rs +++ b/src/query/management/tests/it/warehouse.rs @@ -440,24 +440,24 @@ async fn test_create_warehouse_with_self_manage() -> Result<()> { #[tokio::test(flavor = "multi_thread", worker_threads = 1)] async fn test_create_warehouse_with_no_resources() -> Result<()> { - let (_, cluster_mgr, _nodes) = nodes(Duration::from_mins(30), 2).await?; + let (_, warehouse_manager, _nodes) = nodes(Duration::from_mins(30), 2).await?; - let create_warehouse = - cluster_mgr.create_warehouse("test_warehouse_1".to_string(), vec![SelectedNode::Random( + let create_warehouse = warehouse_manager + .create_warehouse("test_warehouse_1".to_string(), vec![SelectedNode::Random( None, )]); create_warehouse.await?; - let create_warehouse = - cluster_mgr.create_warehouse("test_warehouse_2".to_string(), vec![SelectedNode::Random( + let create_warehouse = warehouse_manager + .create_warehouse("test_warehouse_2".to_string(), vec![SelectedNode::Random( None, )]); create_warehouse.await?; - let create_warehouse = - cluster_mgr.create_warehouse("test_warehouse_3".to_string(), vec![SelectedNode::Random( + let create_warehouse = warehouse_manager + .create_warehouse("test_warehouse_3".to_string(), vec![SelectedNode::Random( None, )]); @@ -469,10 +469,90 @@ async fn test_create_warehouse_with_no_resources() -> Result<()> { Ok(()) } -// #[tokio::test(flavor = "multi_thread", worker_threads = 1)] -// async fn test_create_warehouse_with_no_resources() -> Result<()> { -// -// } +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn test_drop_empty_warehouse() -> Result<()> { + let (_, warehouse_manager, _nodes) = nodes(Duration::from_mins(30), 2).await?; + let drop_warehouse = warehouse_manager.drop_warehouse(String::new()); + + assert_eq!(drop_warehouse.await.unwrap_err().code(), 2403); + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn test_drop_not_exists_warehouse() -> Result<()> { + let (_, warehouse_manager, _nodes) = nodes(Duration::from_mins(30), 2).await?; + let drop_warehouse = warehouse_manager.drop_warehouse(String::from("not_exists")); + + assert_eq!(drop_warehouse.await.unwrap_err().code(), 2406); + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn test_drop_system_managed_warehouse() -> Result<()> { + let (_, warehouse_manager, mut nodes) = nodes(Duration::from_mins(30), 2).await?; + let create_warehouse = + warehouse_manager.create_warehouse(String::from("test_warehouse"), vec![ + SelectedNode::Random(None), + SelectedNode::Random(None), + ]); + + create_warehouse.await?; + + let drop_warehouse = warehouse_manager.drop_warehouse(String::from("test_warehouse")); + drop_warehouse.await?; + + let create_warehouse = + warehouse_manager.create_warehouse(String::from("test_warehouse"), vec![ + SelectedNode::Random(None), + SelectedNode::Random(None), + ]); + + // create same name warehouse is successfully + create_warehouse.await?; + + // mock partial node offline + warehouse_manager.drop_node(nodes.remove(0)).await?; + + let drop_warehouse = warehouse_manager.drop_warehouse(String::from("test_warehouse")); + drop_warehouse.await?; + + // online node + let online_node_id = GlobalUniqName::unique(); + warehouse_manager + .add_node(system_managed_node(&online_node_id)) + .await?; + nodes.push(online_node_id); + let create_warehouse = + warehouse_manager.create_warehouse(String::from("test_warehouse"), vec![ + SelectedNode::Random(None), + SelectedNode::Random(None), + ]); + + // create same name warehouse is successfully + create_warehouse.await?; + + // mock all node offline + warehouse_manager.drop_node(nodes.remove(0)).await?; + warehouse_manager.drop_node(nodes.remove(0)).await?; + + let drop_warehouse = warehouse_manager.drop_warehouse(String::from("test_warehouse")); + drop_warehouse.await?; + + let online_node_id = GlobalUniqName::unique(); + warehouse_manager + .add_node(system_managed_node(&online_node_id)) + .await?; + nodes.push(online_node_id); + + // create same name warehouse is successfully + let create_warehouse = warehouse_manager + .create_warehouse(String::from("test_warehouse"), vec![SelectedNode::Random( + None, + )]); + create_warehouse.await?; + + Ok(()) +} fn system_managed_node(id: &str) -> NodeInfo { NodeInfo { From f1e51149c1a131e7866823195839075b537bd723 Mon Sep 17 00:00:00 2001 From: zhang2014 Date: Wed, 18 Dec 2024 15:44:48 +0800 Subject: [PATCH 06/53] feat(cluster): support custom management cluster --- src/query/management/src/lib.rs | 1 + src/query/management/src/warehouse/mod.rs | 1 + .../management/src/warehouse/warehouse_api.rs | 23 +- .../management/src/warehouse/warehouse_mgr.rs | 504 +++++++++++++----- src/query/management/tests/it/warehouse.rs | 297 ++++++++++- src/query/service/src/clusters/cluster.rs | 6 +- 6 files changed, 659 insertions(+), 173 deletions(-) diff --git a/src/query/management/src/lib.rs b/src/query/management/src/lib.rs index f66a003a11ba0..50fa86859a33c 100644 --- a/src/query/management/src/lib.rs +++ b/src/query/management/src/lib.rs @@ -51,6 +51,7 @@ pub use stage::StageMgr; pub use user::UserApi; pub use user::UserMgr; pub use warehouse::SelectedNode; +pub use warehouse::SystemManagedInfo; pub use warehouse::WarehouseApi; pub use warehouse::WarehouseInfo; pub use warehouse::WarehouseMgr; diff --git a/src/query/management/src/warehouse/mod.rs b/src/query/management/src/warehouse/mod.rs index 0667ed3a7c87b..271155ababfbe 100644 --- a/src/query/management/src/warehouse/mod.rs +++ b/src/query/management/src/warehouse/mod.rs @@ -16,6 +16,7 @@ mod warehouse_api; mod warehouse_mgr; pub use warehouse_api::SelectedNode; +pub use warehouse_api::SystemManagedInfo; pub use warehouse_api::WarehouseApi; pub use warehouse_api::WarehouseInfo; pub use warehouse_mgr::WarehouseMgr; diff --git a/src/query/management/src/warehouse/warehouse_api.rs b/src/query/management/src/warehouse/warehouse_api.rs index ccd062c18c1cb..7da8055e1823e 100644 --- a/src/query/management/src/warehouse/warehouse_api.rs +++ b/src/query/management/src/warehouse/warehouse_api.rs @@ -17,17 +17,18 @@ use std::collections::HashMap; use databend_common_exception::Result; use databend_common_meta_types::NodeInfo; -#[derive(serde::Serialize, serde::Deserialize, Clone, Eq, PartialEq)] +#[derive(serde::Serialize, serde::Deserialize, Clone, Eq, PartialEq, Debug)] pub enum SelectedNode { Random(Option), } -#[derive(serde::Serialize, serde::Deserialize, Eq, PartialEq)] +#[derive(serde::Serialize, serde::Deserialize, Eq, PartialEq, Debug)] pub enum WarehouseInfo { - SelfManaged, + SelfManaged(String), SystemManaged(SystemManagedInfo), } -#[derive(serde::Serialize, serde::Deserialize, Eq, PartialEq)] + +#[derive(serde::Serialize, serde::Deserialize, Eq, PartialEq, Debug)] pub struct SystemManagedInfo { pub id: String, pub status: String, @@ -38,8 +39,11 @@ pub struct SystemManagedInfo { /// Databend-query cluster management API #[async_trait::async_trait] pub trait WarehouseApi: Sync + Send { - /// Add a new node. - async fn add_node(&self, node: NodeInfo) -> Result; + /// Start a new node. + async fn start_node(&self, node: NodeInfo) -> Result; + + /// Shutdown the tenant's cluster one node by node.id. + async fn shutdown_node(&self, node_id: String) -> Result<()>; /// Keep the tenant's cluster node alive. async fn heartbeat(&self, node: &mut NodeInfo, seq: u64) -> Result; @@ -48,14 +52,13 @@ pub trait WarehouseApi: Sync + Send { async fn create_warehouse(&self, warehouse: String, nodes: Vec) -> Result<()>; - // async fn list_warehouses(&self) -> Result>; + async fn list_warehouses(&self) -> Result>; + + async fn rename_warehouse(&self, cur: String, to: String) -> Result<()>; /// Get the tenant's cluster all nodes. async fn get_nodes(&self, warehouse: &str, cluster: &str) -> Result>; - /// Drop the tenant's cluster one node by node.id. - async fn drop_node(&self, node_id: String) -> Result<()>; - async fn get_local_addr(&self) -> Result>; async fn get_node_info(&self, node_id: &str) -> Result; diff --git a/src/query/management/src/warehouse/warehouse_mgr.rs b/src/query/management/src/warehouse/warehouse_mgr.rs index e76aa4b0d29d1..8f5c7374ffd7f 100644 --- a/src/query/management/src/warehouse/warehouse_mgr.rs +++ b/src/query/management/src/warehouse/warehouse_mgr.rs @@ -92,7 +92,27 @@ fn map_condition(k: &str, seq: MatchSeq) -> TxnCondition { } } +struct ConsistentNodeInfo { + node_seq: u64, + cluster_seq: u64, + node_info: NodeInfo, +} + +struct ConsistentWarehouseInfo { + info_seq: u64, + warehouse_info: WarehouseInfo, + consistent_nodes: Vec, +} + impl WarehouseMgr { + fn node_key(&self, node: &NodeInfo) -> Result { + Ok(format!( + "{}/{}", + self.node_key_prefix, + escape_for_key(&node.id)? + )) + } + fn cluster_key(&self, node: &NodeInfo) -> Result { Ok(format!( "{}/{}/{}/{}", @@ -122,6 +142,7 @@ impl WarehouseMgr { )); let warehouse_node_key = self.cluster_key(&node)?; + let warehouse_info = WarehouseInfo::SelfManaged(node.warehouse_id.clone()); let warehouse_info_key = format!( "{}/{}", self.warehouse_key_prefix, @@ -137,9 +158,10 @@ impl WarehouseMgr { )); // upsert warehouse info if self-managed. + txn.if_then.push(TxnOp::put_with_ttl( warehouse_info_key.clone(), - serde_json::to_vec(&WarehouseInfo::SelfManaged)?, + serde_json::to_vec(&warehouse_info)?, Some(self.lift_time), )); @@ -176,7 +198,7 @@ impl WarehouseMgr { "Already exists same name system-managed warehouse.", )) } - WarehouseInfo::SelfManaged => match response.responses.first() { + WarehouseInfo::SelfManaged(_) => match response.responses.first() { // already exists node. Some(TxnOpResponse { response: @@ -330,13 +352,150 @@ impl WarehouseMgr { _ => Err(ErrorCode::Internal("Miss type while in meta response")), } } + + async fn consistent_warehouse_info(&self, id: &str) -> Result { + let warehouse_key = format!("{}/{}", self.warehouse_key_prefix, escape_for_key(id)?); + + let nodes_prefix = format!("{}/{}", self.meta_key_prefix, escape_for_key(id)?); + + 'retry: for _idx in 0..64 { + let Some(before_info) = self.metastore.get_kv(&warehouse_key).await? else { + return Err(ErrorCode::UnknownWarehouse(format!( + "Unknown warehouse or self managed warehouse {:?}", + id + ))); + }; + + let values = self.metastore.prefix_list_kv(&nodes_prefix).await?; + + let mut after_txn = TxnRequest::default(); + let mut cluster_node_seq = Vec::with_capacity(values.len()); + + for (node_key, value) in values { + let suffix = &node_key[nodes_prefix.len() + 1..]; + + if let Some((_cluster, node)) = suffix.split_once('/') { + let node_key = format!("{}/{}", self.node_key_prefix, node); + after_txn.if_then.push(TxnOp::get(node_key)); + cluster_node_seq.push(value.seq); + continue; + } + + return Err(ErrorCode::InvalidWarehouse(format!( + "Node key is invalid {:?}", + node_key + ))); + } + + let condition = map_condition(&warehouse_key, MatchSeq::Exact(before_info.seq)); + after_txn.condition.push(condition); + + match self.metastore.transaction(after_txn).await? { + response if response.success => { + let mut consistent_nodes = Vec::with_capacity(response.responses.len()); + for (idx, response) in response.responses.into_iter().enumerate() { + match response.response { + // TODO: maybe ignore none(not need retry) + Some(Response::Get(response)) => match response.value { + Some(value) => { + let node_info = + serde_json::from_slice::(&value.data)?; + + assert_eq!(node_info.warehouse_id, id); + assert!(!node_info.cluster_id.is_empty()); + + consistent_nodes.push(ConsistentNodeInfo { + node_seq: value.seq, + cluster_seq: cluster_node_seq[idx], + node_info, + }); + } + _ => { + continue 'retry; + } + }, + _ => { + continue 'retry; + } + } + } + + if consistent_nodes.len() == cluster_node_seq.len() { + return Ok(ConsistentWarehouseInfo { + info_seq: before_info.seq, + warehouse_info: serde_json::from_slice(&before_info.data)?, + consistent_nodes, + }); + } + } + _ => { + continue 'retry; + } + } + } + + Err(ErrorCode::Internal( + "Get consistent warehouse info failure(tried 64 times)", + )) + } + + async fn shutdown_self_managed_node(&self, node_info: &NodeInfo) -> Result<()> { + for _idx in 0..10 { + let consistent_info = self + .consistent_warehouse_info(&node_info.warehouse_id) + .await?; + + let mut txn = TxnRequest::default(); + + let node_key = self.node_key(node_info)?; + let cluster_key = self.cluster_key(node_info)?; + + if consistent_info.consistent_nodes.len() == 1 + && consistent_info.consistent_nodes[0].node_info.id == node_info.id + { + let warehouse_key = format!( + "{}/{}", + self.warehouse_key_prefix, + escape_for_key(&node_info.warehouse_id)? + ); + + txn.condition.push(map_condition( + &warehouse_key, + MatchSeq::Exact(consistent_info.info_seq), + )); + txn.if_then.push(TxnOp::delete(warehouse_key)); + + for consistent_node in consistent_info.consistent_nodes { + txn.condition.push(map_condition( + &node_key, + MatchSeq::Exact(consistent_node.node_seq), + )); + txn.condition.push(map_condition( + &cluster_key, + MatchSeq::Exact(consistent_node.cluster_seq), + )); + } + } + + txn.if_then.push(TxnOp::delete(node_key)); + txn.if_then.push(TxnOp::delete(cluster_key)); + + if self.metastore.transaction(txn).await?.success { + return Ok(()); + } + } + + Err(ErrorCode::WarehouseOperateConflict( + "Warehouse operate conflict(tried 10 times).", + )) + } } #[async_trait::async_trait] impl WarehouseApi for WarehouseMgr { #[async_backtrace::framed] #[fastrace::trace] - async fn add_node(&self, node: NodeInfo) -> Result { + async fn start_node(&self, node: NodeInfo) -> Result { let res = self.upsert_node(node.clone(), MatchSeq::Exact(0)).await?; if res.success { @@ -359,6 +518,48 @@ impl WarehouseApi for WarehouseMgr { ))) } + #[async_backtrace::framed] + #[fastrace::trace] + async fn shutdown_node(&self, node_id: String) -> Result<()> { + let node_key = format!("{}/{}", self.node_key_prefix, escape_for_key(&node_id)?); + + if let Some(info) = self.metastore.get_kv(&node_key).await? { + let node_info: NodeInfo = serde_json::from_slice(&info.data)?; + + return match node_info.node_type { + NodeType::SelfManaged => self.shutdown_self_managed_node(&node_info).await, + NodeType::SystemManaged => { + let mut txn = TxnRequest::default(); + + txn.if_then.push(TxnOp::delete(node_key)); + + if !node_info.cluster_id.is_empty() && !node_info.warehouse_id.is_empty() { + txn.if_then.push(TxnOp::delete(format!( + "{}/{}/{}/{}", + self.meta_key_prefix, + escape_for_key(&node_info.warehouse_id)?, + escape_for_key(&node_info.cluster_id)?, + escape_for_key(&node_info.id)? + ))); + } + + match self.metastore.transaction(txn).await?.success { + true => Ok(()), + false => Err(ErrorCode::ClusterUnknownNode(format!( + "Node with ID '{}' does not exist in the cluster.", + node_id + ))), + } + } + }; + } + + Err(ErrorCode::ClusterUnknownNode(format!( + "Node with ID '{}' does not exist in the cluster.", + node_id + ))) + } + async fn heartbeat(&self, node: &mut NodeInfo, seq: u64) -> Result { assert!(!node.cluster_id.is_empty()); assert!(!node.warehouse_id.is_empty()); @@ -388,117 +589,66 @@ impl WarehouseApi for WarehouseMgr { return Err(ErrorCode::InvalidWarehouse("Warehouse name is empty.")); } - let drop_key = format!("{}/{}", self.meta_key_prefix, escape_for_key(&warehouse)?); - - loop { - let mut delete_txn = TxnRequest::default(); - let values = self.metastore.prefix_list_kv(&drop_key).await?; + for _idx in 0..10 { + let consistent_info = self.consistent_warehouse_info(&warehouse).await?; - let mut txn = TxnRequest::default(); - for (node_key, _value) in values { - let suffix = &node_key[drop_key.len() + 1..]; - - if let Some((_, node)) = suffix.split_once('/') { - txn.if_then - .push(TxnOp::get(format!("{}/{}", self.node_key_prefix, node))); - } - - delete_txn.if_then.push(TxnOp::delete(node_key)); + if matches!( + consistent_info.warehouse_info, + WarehouseInfo::SelfManaged(_) + ) { + return Err(ErrorCode::InvalidWarehouse( + "Cannot drop self-managed warehouse", + )); } + let mut delete_txn = TxnRequest::default(); + let warehouse_key = format!( "{}/{}", self.warehouse_key_prefix, escape_for_key(&warehouse)? ); - txn.if_then.push(TxnOp::get(warehouse_key.clone())); - let mut fetch_reply = self.metastore.transaction(txn).await?; - - if let Some(response) = fetch_reply.responses.pop() { - let Some(Response::Get(response)) = response.response else { - return Err(ErrorCode::UnknownWarehouse(format!( - "Unknown warehouse or self managed warehouse {:?}", - warehouse - ))); - }; - - if response.key != warehouse_key { - return Err(ErrorCode::UnknownWarehouse(format!( - "Unknown warehouse or self managed warehouse {:?}", - warehouse - ))); - } - - let Some(value) = &response.value else { - return Err(ErrorCode::UnknownWarehouse(format!( - "Unknown warehouse or self managed warehouse {:?}", - warehouse - ))); - }; - - if serde_json::from_slice::(&value.data)? - == WarehouseInfo::SelfManaged - { - return Err(ErrorCode::InvalidWarehouse( - "Cannot drop self-managed warehouse", - )); - } - - if fetch_reply.responses.len() != delete_txn.if_then.len() { - // TODO: maybe auto retry? - return Err(ErrorCode::WarehouseOperateConflict("Missing node info in online nodes list. It's possible that some nodes offline during the drop warehouse. You may try the operation again.")); - } - - for response in fetch_reply.responses { - let Some(Response::Get(response)) = response.response else { - // TODO: maybe auto retry? - return Err(ErrorCode::WarehouseOperateConflict("Missing node info in online nodes list. It's possible that some nodes offline during the drop warehouse. You may try the operation again.")); - }; - - let Some(value) = &response.value else { - // TODO: maybe auto retry? - return Err(ErrorCode::WarehouseOperateConflict("Missing node info in online nodes list. It's possible that some nodes offline during the drop warehouse. You may try the operation again.")); - }; - - let mut node_info = serde_json::from_slice::(&value.data)?; - - if node_info.warehouse_id != warehouse { - return Err(ErrorCode::WarehouseOperateConflict( - "Warehouse operate is conflict.", - )); - } + delete_txn.condition.push(map_condition( + &warehouse_key, + MatchSeq::Exact(consistent_info.info_seq), + )); + delete_txn.if_then.push(TxnOp::delete(warehouse_key)); - delete_txn - .condition - .push(map_condition(&response.key, MatchSeq::Exact(value.seq))); - node_info.cluster_id = String::new(); - node_info.warehouse_id = String::new(); - delete_txn.if_then.push(TxnOp::put_with_ttl( - response.key, - serde_json::to_vec(&node_info)?, - Some(self.lift_time * 4), - )); - } + for mut consistent_node in consistent_info.consistent_nodes { + let node_key = self.node_key(&consistent_node.node_info)?; + let cluster_key = self.cluster_key(&consistent_node.node_info)?; - delete_txn - .condition - .push(map_condition(&warehouse_key, MatchSeq::Exact(value.seq))); - delete_txn.if_then.push(TxnOp::delete(warehouse_key)); + delete_txn.condition.push(map_condition( + &node_key, + MatchSeq::Exact(consistent_node.node_seq), + )); + delete_txn.condition.push(map_condition( + &cluster_key, + MatchSeq::Exact(consistent_node.cluster_seq), + )); - if !self.metastore.transaction(delete_txn).await?.success { - // seq is changed, will retry - continue; - } + delete_txn.if_then.push(TxnOp::delete(cluster_key)); + consistent_node.node_info.cluster_id = String::new(); + consistent_node.node_info.warehouse_id = String::new(); + delete_txn.if_then.push(TxnOp::put_with_ttl( + node_key, + serde_json::to_vec(&consistent_node.node_info)?, + Some(self.lift_time * 4), + )); + } - return Ok(()); + if !self.metastore.transaction(delete_txn).await?.success { + // seq is changed, will retry + continue; } - return Err(ErrorCode::UnknownWarehouse(format!( - "Unknown warehouse or self managed warehouse {:?}", - warehouse - ))); + return Ok(()); } + + Err(ErrorCode::WarehouseOperateConflict( + "Warehouse operate conflict(tried 10 times).", + )) } async fn create_warehouse(&self, warehouse: String, nodes: Vec) -> Result<()> { @@ -626,6 +776,129 @@ impl WarehouseApi for WarehouseMgr { } } + async fn list_warehouses(&self) -> Result> { + let values = self + .metastore + .prefix_list_kv(&self.warehouse_key_prefix) + .await?; + + let mut warehouses = Vec::with_capacity(values.len()); + for (_warehouse_key, value) in values { + warehouses.push(serde_json::from_slice::(&value.data)?); + } + + Ok(warehouses) + } + + async fn rename_warehouse(&self, current: String, to: String) -> Result<()> { + if current.is_empty() { + return Err(ErrorCode::InvalidWarehouse("Warehouse name is empty.")); + } + + if to.is_empty() { + return Err(ErrorCode::InvalidWarehouse("Warehouse name is empty.")); + } + + for _idx in 0..10 { + let mut consistent_info = self.consistent_warehouse_info(¤t).await?; + + consistent_info.warehouse_info = match consistent_info.warehouse_info { + WarehouseInfo::SelfManaged(_) => Err(ErrorCode::InvalidWarehouse( + "Cannot rename self-managed warehouse", + )), + WarehouseInfo::SystemManaged(mut info) => { + info.display_name = to.clone(); + Ok(WarehouseInfo::SystemManaged(info)) + } + }?; + + let mut rename_txn = TxnRequest::default(); + + let old_warehouse_key = format!( + "{}/{}", + self.warehouse_key_prefix, + escape_for_key(¤t)? + ); + + let new_warehouse_key = + format!("{}/{}", self.warehouse_key_prefix, escape_for_key(&to)?); + + rename_txn.condition.push(map_condition( + &old_warehouse_key, + MatchSeq::Exact(consistent_info.info_seq), + )); + + rename_txn + .condition + .push(map_condition(&new_warehouse_key, MatchSeq::Exact(0))); + rename_txn + .else_then + .push(TxnOp::get(new_warehouse_key.clone())); + + rename_txn.if_then.push(TxnOp::delete(old_warehouse_key)); + rename_txn.if_then.push(TxnOp::put_with_ttl( + new_warehouse_key, + serde_json::to_vec(&consistent_info.warehouse_info)?, + Some(self.lift_time * 4), + )); + + for mut consistent_node in consistent_info.consistent_nodes { + let node_key = self.node_key(&consistent_node.node_info)?; + let old_cluster_key = self.cluster_key(&consistent_node.node_info)?; + + consistent_node.node_info.warehouse_id = to.clone(); + + let new_cluster_key = self.cluster_key(&consistent_node.node_info)?; + + rename_txn.condition.push(map_condition( + &node_key, + MatchSeq::Exact(consistent_node.node_seq), + )); + rename_txn.condition.push(map_condition( + &old_cluster_key, + MatchSeq::Exact(consistent_node.cluster_seq), + )); + + rename_txn + .condition + .push(map_condition(&new_cluster_key, MatchSeq::Exact(0))); + + rename_txn.if_then.push(TxnOp::put_with_ttl( + node_key, + serde_json::to_vec(&consistent_node.node_info)?, + Some(self.lift_time * 4), + )); + consistent_node.node_info.cluster_id = String::new(); + consistent_node.node_info.warehouse_id = String::new(); + rename_txn.if_then.push(TxnOp::delete(old_cluster_key)); + rename_txn.if_then.push(TxnOp::put_with_ttl( + new_cluster_key.clone(), + serde_json::to_vec(&consistent_node.node_info)?, + Some(self.lift_time * 4), + )); + } + + return match self.metastore.transaction(rename_txn).await? { + response if response.success => Ok(()), + response => match response.responses.last() { + Some(TxnOpResponse { + response: Some(Response::Get(TxnGetResponse { value: Some(v), .. })), + }) if v.seq != 0 => Err(ErrorCode::WarehouseAlreadyExists(format!( + "Warehouse {} already exists.", + to + ))), + _ => { + continue; + } + }, + }; + } + + Err(ErrorCode::WarehouseOperateConflict( + "Warehouse operate conflict(tried 10 times while in rename warehouse).", + )) + } + #[async_backtrace::framed] #[fastrace::trace] async fn get_nodes(&self, warehouse: &str, cluster: &str) -> Result> { @@ -651,41 +924,6 @@ impl WarehouseApi for WarehouseMgr { Ok(nodes_info) } - #[async_backtrace::framed] - #[fastrace::trace] - async fn drop_node(&self, node_id: String) -> Result<()> { - let node_key = format!("{}/{}", self.node_key_prefix, escape_for_key(&node_id)?); - - if let Some(info) = self.metastore.get_kv(&node_key).await? { - let node_info: NodeInfo = serde_json::from_slice(&info.data)?; - - let mut txn = TxnRequest::default(); - - txn.if_then.push(TxnOp::delete(node_key)); - - if !node_info.cluster_id.is_empty() && !node_info.warehouse_id.is_empty() { - txn.if_then.push(TxnOp::delete(format!( - "{}/{}/{}/{}", - self.meta_key_prefix, - escape_for_key(&node_info.warehouse_id)?, - escape_for_key(&node_info.cluster_id)?, - escape_for_key(&node_info.id)? - ))); - } - - let res = self.metastore.transaction(txn).await?; - - if res.success { - return Ok(()); - } - } - - Err(ErrorCode::ClusterUnknownNode(format!( - "Node with ID '{}' does not exist in the cluster.", - node_id - ))) - } - #[async_backtrace::framed] #[fastrace::trace] async fn get_local_addr(&self) -> Result> { diff --git a/src/query/management/tests/it/warehouse.rs b/src/query/management/tests/it/warehouse.rs index bd374a1e5bcc8..a730c2fb702f9 100644 --- a/src/query/management/tests/it/warehouse.rs +++ b/src/query/management/tests/it/warehouse.rs @@ -12,6 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. +use std::collections::HashMap; use std::sync::Arc; use std::time::Duration; @@ -36,7 +37,7 @@ async fn test_empty_id_with_self_managed() -> Result<()> { node.node_type = NodeType::SelfManaged; node.warehouse_id = String::new(); node.cluster_id = String::from("test_cluster_id"); - let res = warehouse_manager.add_node(node).await; + let res = warehouse_manager.start_node(node).await; assert!(res.is_err()); assert_eq!(res.unwrap_err().code(), 2403); @@ -45,7 +46,7 @@ async fn test_empty_id_with_self_managed() -> Result<()> { node.node_type = NodeType::SelfManaged; node.cluster_id = String::new(); node.warehouse_id = String::from("test_cluster_id"); - let res = warehouse_manager.add_node(node).await; + let res = warehouse_manager.start_node(node).await; assert!(res.is_err()); assert_eq!(res.unwrap_err().code(), 2403); @@ -58,7 +59,7 @@ async fn test_successfully_add_self_managed_node() -> Result<()> { let (kv, warehouse_manager, _nodes) = nodes(Duration::from_secs(60), 0).await?; let mut node_info_1 = self_managed_node("test_node_1"); - warehouse_manager.add_node(node_info_1.clone()).await?; + warehouse_manager.start_node(node_info_1.clone()).await?; let node_key = "__fd_clusters_v5/test%2dtenant%2did/online_nodes/test_node_1"; assert_key_value(&kv, node_key, serde_json::to_vec(&node_info_1)?).await; @@ -73,12 +74,12 @@ async fn test_successfully_add_self_managed_node() -> Result<()> { assert_key_value( &kv, info_key, - serde_json::to_vec(&WarehouseInfo::SelfManaged)?, + serde_json::to_vec(&WarehouseInfo::SelfManaged(String::from("test-cluster-id")))?, ) .await; let mut node_info_2 = self_managed_node("test_node_2"); - warehouse_manager.add_node(node_info_2.clone()).await?; + warehouse_manager.start_node(node_info_2.clone()).await?; let node_key = "__fd_clusters_v5/test%2dtenant%2did/online_nodes/test_node_2"; assert_key_value(&kv, node_key, serde_json::to_vec(&node_info_2)?).await; @@ -94,7 +95,7 @@ async fn test_successfully_add_self_managed_node() -> Result<()> { assert_key_value( &kv, info_key, - serde_json::to_vec(&WarehouseInfo::SelfManaged)?, + serde_json::to_vec(&WarehouseInfo::SelfManaged(String::from("test-cluster-id")))?, ) .await; @@ -106,20 +107,20 @@ async fn test_already_exists_add_self_managed_node() -> Result<()> { let (kv, warehouse_manager, nodes) = nodes(Duration::from_secs(60), 1).await?; let node_info = self_managed_node("test_node_1"); - warehouse_manager.add_node(node_info.clone()).await?; + warehouse_manager.start_node(node_info.clone()).await?; let node_key = "__fd_clusters_v5/test%2dtenant%2did/online_nodes/test_node_1"; assert_key_value(&kv, node_key, serde_json::to_vec(&node_info)?).await; // add already exists self-managed node and get failure - match warehouse_manager.add_node(node_info.clone()).await { + match warehouse_manager.start_node(node_info.clone()).await { Ok(_) => panic!("Already exists add node must be return Err."), Err(cause) => assert_eq!(cause.code(), 2402), } // add already exists system-managed node and get failure match warehouse_manager - .add_node(self_managed_node(&nodes[0])) + .start_node(self_managed_node(&nodes[0])) .await { Ok(_) => panic!("Already exists add node must be return Err."), @@ -138,14 +139,14 @@ async fn test_successfully_get_self_managed_nodes() -> Result<()> { assert_eq!(get_nodes.await?, vec![]); let node_1 = self_managed_node("node_1"); - warehouse_manager.add_node(node_1.clone()).await?; + warehouse_manager.start_node(node_1.clone()).await?; let get_nodes = warehouse_manager.get_nodes("test-cluster-id", "test-cluster-id"); assert_eq!(get_nodes.await?, vec![node_1.clone()]); let node_2 = self_managed_node("node_2"); - warehouse_manager.add_node(node_2.clone()).await?; + warehouse_manager.start_node(node_2.clone()).await?; let get_nodes = warehouse_manager.get_nodes("test-cluster-id", "test-cluster-id"); @@ -158,13 +159,13 @@ async fn test_successfully_drop_self_managed_node() -> Result<()> { let (_, warehouse_manager, _nodes) = nodes(Duration::from_mins(60), 0).await?; let node_info = self_managed_node("test_node"); - warehouse_manager.add_node(node_info.clone()).await?; + warehouse_manager.start_node(node_info.clone()).await?; let get_nodes = warehouse_manager.get_nodes("test-cluster-id", "test-cluster-id"); assert_eq!(get_nodes.await?, vec![node_info.clone()]); - warehouse_manager.drop_node(node_info.id).await?; + warehouse_manager.shutdown_node(node_info.id).await?; let get_nodes = warehouse_manager.get_nodes("test-cluster-id", "test-cluster-id"); @@ -177,7 +178,7 @@ async fn test_unknown_node_drop_self_managed_node() -> Result<()> { let (_, warehouse_manager, _nodes) = nodes(Duration::from_mins(60), 0).await?; match warehouse_manager - .drop_node(String::from("UNKNOWN_ID")) + .shutdown_node(String::from("UNKNOWN_ID")) .await { Ok(_) => { /*panic!("Unknown node drop node must be return Err.")*/ } @@ -192,7 +193,7 @@ async fn test_drop_self_managed_warehouse() -> Result<()> { let (_, warehouse_manager, _nodes) = nodes(Duration::from_mins(60), 0).await?; let node_info = self_managed_node("test_node"); - warehouse_manager.add_node(node_info.clone()).await?; + warehouse_manager.start_node(node_info.clone()).await?; let drop_warehouse = warehouse_manager.drop_warehouse(String::from("test-cluster-id")); @@ -206,7 +207,7 @@ async fn test_successfully_heartbeat_self_managed_node() -> Result<()> { let (kv, warehouse_manager, _nodes) = nodes(Duration::from_mins(60), 0).await?; let mut node_info = self_managed_node("test_node"); - let seq = warehouse_manager.add_node(node_info.clone()).await?; + let seq = warehouse_manager.start_node(node_info.clone()).await?; let info_key = "__fd_clusters_v5/test%2dtenant%2did/online_nodes/test_node"; assert_key_value(&kv, info_key, serde_json::to_vec(&node_info)?).await; @@ -220,7 +221,7 @@ async fn test_successfully_heartbeat_self_managed_node() -> Result<()> { assert_key_expire(&kv, warehouse_key, Duration::from_mins(50)).await; let warehouse_info_key = "__fd_warehouses/v1/test%2dtenant%2did/test%2dcluster%2did"; - let info = serde_json::to_vec(&WarehouseInfo::SelfManaged)?; + let info = serde_json::to_vec(&WarehouseInfo::SelfManaged(String::from("test-cluster-id")))?; assert_key_value(&kv, warehouse_info_key, info.clone()).await; assert_key_expire(&kv, warehouse_info_key, Duration::from_mins(50)).await; @@ -309,7 +310,9 @@ async fn test_create_system_managed_warehouse_with_offline_node() -> Result<()> let (_, warehouse_manager, mut nodes) = nodes(Duration::from_mins(30), 4).await?; // mock node offline - warehouse_manager.drop_node(nodes[0].to_string()).await?; + warehouse_manager + .shutdown_node(nodes[0].to_string()) + .await?; let create_warehouse = warehouse_manager.create_warehouse("test_warehouse".to_string(), vec![ SelectedNode::Random(None), @@ -362,7 +365,7 @@ async fn test_create_system_managed_warehouse_with_online_node() -> Result<()> { // mock node online let new_node = GlobalUniqName::unique(); warehouse_manager - .add_node(system_managed_node(&new_node)) + .start_node(system_managed_node(&new_node)) .await?; let create_warehouse = warehouse_manager.create_warehouse("test_warehouse".to_string(), vec![ @@ -426,7 +429,9 @@ async fn test_create_warehouse_with_self_manage() -> Result<()> { let mut self_manage_node_1 = self_managed_node("self_manage_node_1"); self_manage_node_1.cluster_id = String::from("test_warehouse"); self_manage_node_1.warehouse_id = String::from("test_warehouse"); - warehouse_manager.add_node(self_manage_node_1).await?; + warehouse_manager + .start_node(self_manage_node_1.clone()) + .await?; let create_warehouse = warehouse_manager .create_warehouse(String::from("test_warehouse"), vec![SelectedNode::Random( @@ -435,7 +440,16 @@ async fn test_create_warehouse_with_self_manage() -> Result<()> { assert_eq!(create_warehouse.await.unwrap_err().code(), 2405); - Ok(()) + warehouse_manager + .shutdown_node(self_manage_node_1.id.clone()) + .await?; + + let create_warehouse = warehouse_manager + .create_warehouse(String::from("test_warehouse"), vec![SelectedNode::Random( + None, + )]); + + create_warehouse.await } #[tokio::test(flavor = "multi_thread", worker_threads = 1)] @@ -511,7 +525,7 @@ async fn test_drop_system_managed_warehouse() -> Result<()> { create_warehouse.await?; // mock partial node offline - warehouse_manager.drop_node(nodes.remove(0)).await?; + warehouse_manager.shutdown_node(nodes.remove(0)).await?; let drop_warehouse = warehouse_manager.drop_warehouse(String::from("test_warehouse")); drop_warehouse.await?; @@ -519,7 +533,7 @@ async fn test_drop_system_managed_warehouse() -> Result<()> { // online node let online_node_id = GlobalUniqName::unique(); warehouse_manager - .add_node(system_managed_node(&online_node_id)) + .start_node(system_managed_node(&online_node_id)) .await?; nodes.push(online_node_id); let create_warehouse = @@ -532,15 +546,15 @@ async fn test_drop_system_managed_warehouse() -> Result<()> { create_warehouse.await?; // mock all node offline - warehouse_manager.drop_node(nodes.remove(0)).await?; - warehouse_manager.drop_node(nodes.remove(0)).await?; + warehouse_manager.shutdown_node(nodes.remove(0)).await?; + warehouse_manager.shutdown_node(nodes.remove(0)).await?; let drop_warehouse = warehouse_manager.drop_warehouse(String::from("test_warehouse")); drop_warehouse.await?; let online_node_id = GlobalUniqName::unique(); warehouse_manager - .add_node(system_managed_node(&online_node_id)) + .start_node(system_managed_node(&online_node_id)) .await?; nodes.push(online_node_id); @@ -554,6 +568,233 @@ async fn test_drop_system_managed_warehouse() -> Result<()> { Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn test_list_warehouses() -> Result<()> { + let (_, warehouse_manager, _nodes) = nodes(Duration::from_mins(30), 10).await?; + + assert_eq!(warehouse_manager.list_warehouses().await?, vec![]); + + let self_managed_node_1 = GlobalUniqName::unique(); + warehouse_manager + .start_node(self_managed_node(&self_managed_node_1)) + .await?; + + let list_warehouse_1 = warehouse_manager.list_warehouses().await?; + assert_eq!(list_warehouse_1, vec![WarehouseInfo::SelfManaged( + String::from("test-cluster-id") + )]); + + let create_warehouse = warehouse_manager + .create_warehouse(String::from("test_warehouse_1"), vec![ + SelectedNode::Random(None), + ]); + + create_warehouse.await?; + + let list_warehouses_2 = warehouse_manager.list_warehouses().await?; + + assert_eq!(list_warehouses_2.len(), 2); + assert_eq!(list_warehouses_2[0], list_warehouse_1[0]); + let WarehouseInfo::SystemManaged(system_managed_info) = &list_warehouses_2[1] else { + unreachable!(); + }; + assert!(!system_managed_info.id.is_empty()); + assert_eq!(system_managed_info.status, "Running"); + assert_eq!(system_managed_info.display_name, "test_warehouse_1"); + assert_eq!( + system_managed_info.clusters, + HashMap::from([(String::from("default"), vec![SelectedNode::Random(None)])]) + ); + + let self_managed_node_2 = GlobalUniqName::unique(); + let mut self_managed_node = self_managed_node(&self_managed_node_2); + self_managed_node.warehouse_id = String::from("test_warehouse_2"); + warehouse_manager.start_node(self_managed_node).await?; + + let list_warehouses_3 = warehouse_manager.list_warehouses().await?; + + assert_eq!(list_warehouses_3.len(), 3); + assert_eq!(list_warehouses_3[0], list_warehouses_2[0]); + assert_eq!(list_warehouses_3[1], list_warehouses_2[1]); + assert_eq!( + list_warehouses_3[2], + WarehouseInfo::SelfManaged(String::from("test_warehouse_2")) + ); + + let create_warehouse = warehouse_manager + .create_warehouse(String::from("test_warehouse_3"), vec![ + SelectedNode::Random(None), + ]); + + create_warehouse.await?; + + let mut list_warehouses_4 = warehouse_manager.list_warehouses().await?; + + assert_eq!(list_warehouses_4.len(), 4); + assert_eq!(list_warehouses_4[0], list_warehouses_3[0]); + assert_eq!(list_warehouses_4[1], list_warehouses_3[1]); + assert_eq!(list_warehouses_4[2], list_warehouses_3[2]); + + let WarehouseInfo::SystemManaged(system_managed_info) = &list_warehouses_4[3] else { + unreachable!(); + }; + assert!(!system_managed_info.id.is_empty()); + assert_eq!(system_managed_info.status, "Running"); + assert_eq!(system_managed_info.display_name, "test_warehouse_3"); + assert_eq!( + system_managed_info.clusters, + HashMap::from([(String::from("default"), vec![SelectedNode::Random(None)])]) + ); + + warehouse_manager.shutdown_node(self_managed_node_1).await?; + list_warehouses_4.remove(0); + assert_eq!( + warehouse_manager.list_warehouses().await?, + list_warehouses_4 + ); + + warehouse_manager.shutdown_node(self_managed_node_2).await?; + list_warehouses_4.remove(1); + assert_eq!( + warehouse_manager.list_warehouses().await?, + list_warehouses_4 + ); + + warehouse_manager + .drop_warehouse(String::from("test_warehouse_1")) + .await?; + list_warehouses_4.remove(0); + assert_eq!( + warehouse_manager.list_warehouses().await?, + list_warehouses_4 + ); + + // keep show warehouse if all node offline + let nodes = warehouse_manager + .get_nodes("test_warehouse_3", "default") + .await?; + warehouse_manager.shutdown_node(nodes[0].id.clone()).await?; + assert_eq!( + warehouse_manager.list_warehouses().await?, + list_warehouses_4 + ); + warehouse_manager + .drop_warehouse(String::from("test_warehouse_3")) + .await?; + list_warehouses_4.remove(0); + assert_eq!( + warehouse_manager.list_warehouses().await?, + list_warehouses_4 + ); + + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn test_rename_warehouses() -> Result<()> { + let (kv, warehouse_manager, nodes) = nodes(Duration::from_mins(30), 1).await?; + + let self_managed_node_1 = GlobalUniqName::unique(); + warehouse_manager + .start_node(self_managed_node(&self_managed_node_1)) + .await?; + + let rename_warehouse = warehouse_manager.rename_warehouse( + String::from("test-cluster-id"), + String::from("test_warehouse"), + ); + + assert_eq!(rename_warehouse.await.unwrap_err().code(), 2403); + + warehouse_manager.shutdown_node(self_managed_node_1).await?; + + let warehouse_node_key = format!( + "__fd_clusters_v5/test%2dtenant%2did/online_clusters/test_warehouse/default/{}", + &nodes[0] + ); + + let create_warehouse = warehouse_manager + .create_warehouse(String::from("test_warehouse"), vec![SelectedNode::Random( + None, + )]); + + create_warehouse.await?; + + assert_key_seq(&kv, &warehouse_node_key, MatchSeq::GE(1)).await; + let list_warehouses = warehouse_manager.list_warehouses().await?; + + let WarehouseInfo::SystemManaged(system_managed_info) = &list_warehouses[0] else { + unreachable!(); + }; + + assert!(!system_managed_info.id.is_empty()); + assert_eq!(system_managed_info.status, "Running"); + assert_eq!(system_managed_info.display_name, "test_warehouse"); + assert_eq!( + system_managed_info.clusters, + HashMap::from([(String::from("default"), vec![SelectedNode::Random(None)])]) + ); + + let rename_warehouse = warehouse_manager.rename_warehouse( + String::from("test_warehouse"), + String::from("new_test_warehouse"), + ); + + rename_warehouse.await?; + + assert_no_key(&kv, &warehouse_node_key).await; + + let list_warehouses = warehouse_manager.list_warehouses().await?; + + let WarehouseInfo::SystemManaged(system_managed_info) = &list_warehouses[0] else { + unreachable!(); + }; + + assert!(!system_managed_info.id.is_empty()); + assert_eq!(system_managed_info.status, "Running"); + assert_eq!(system_managed_info.display_name, "new_test_warehouse"); + assert_eq!( + system_managed_info.clusters, + HashMap::from([(String::from("default"), vec![SelectedNode::Random(None)])]) + ); + + let system_managed_node_2 = GlobalUniqName::unique(); + warehouse_manager + .start_node(system_managed_node(&system_managed_node_2)) + .await?; + + let create_warehouse = warehouse_manager + .create_warehouse(String::from("test_warehouse"), vec![SelectedNode::Random( + None, + )]); + + create_warehouse.await?; + + let list_warehouses_2 = warehouse_manager.list_warehouses().await?; + + assert_eq!(list_warehouses_2[0], list_warehouses[0]); + let WarehouseInfo::SystemManaged(system_managed_info) = &list_warehouses_2[1] else { + unreachable!(); + }; + + assert!(!system_managed_info.id.is_empty()); + assert_eq!(system_managed_info.status, "Running"); + assert_eq!(system_managed_info.display_name, "test_warehouse"); + assert_eq!( + system_managed_info.clusters, + HashMap::from([(String::from("default"), vec![SelectedNode::Random(None)])]) + ); + + let rename_warehouse = warehouse_manager.rename_warehouse( + String::from("new_test_warehouse"), + String::from("test_warehouse"), + ); + + assert_eq!(rename_warehouse.await.unwrap_err().code(), 2405); + + Ok(()) +} + fn system_managed_node(id: &str) -> NodeInfo { NodeInfo { id: id.to_string(), @@ -592,7 +833,9 @@ async fn nodes(lift: Duration, size: usize) -> Result<(MetaStore, WarehouseMgr, let mut nodes = Vec::with_capacity(size); for _index in 0..size { let name = GlobalUniqName::unique(); - cluster_manager.add_node(system_managed_node(&name)).await?; + cluster_manager + .start_node(system_managed_node(&name)) + .await?; nodes.push(name); } diff --git a/src/query/service/src/clusters/cluster.rs b/src/query/service/src/clusters/cluster.rs index 493e41b5682a9..6a1f117c40fbb 100644 --- a/src/query/service/src/clusters/cluster.rs +++ b/src/query/service/src/clusters/cluster.rs @@ -346,7 +346,7 @@ impl ClusterDiscovery { // Restart in a very short time(< heartbeat timeout) after abnormal shutdown, Which will // lead to some invalid information if before_node.flight_address.eq(&node_info.flight_address) { - let drop_invalid_node = self.api_provider.drop_node(before_node.id); + let drop_invalid_node = self.api_provider.shutdown_node(before_node.id); if let Err(cause) = drop_invalid_node.await { warn!("Drop invalid node failure: {:?}", cause); } @@ -369,7 +369,7 @@ impl ClusterDiscovery { let mut mut_signal_pin = signal.as_mut(); let signal_future = Box::pin(mut_signal_pin.next()); - let drop_node = Box::pin(self.api_provider.drop_node(self.local_id.clone())); + let drop_node = Box::pin(self.api_provider.shutdown_node(self.local_id.clone())); match futures::future::select(drop_node, signal_future).await { Either::Left((drop_node_result, _)) => { if let Err(drop_node_failure) = drop_node_result { @@ -443,7 +443,7 @@ impl ClusterDiscovery { node_info.node_type = NodeType::SelfManaged; self.drop_invalid_nodes(&node_info).await?; - match self.api_provider.add_node(node_info.clone()).await { + match self.api_provider.start_node(node_info.clone()).await { Ok(seq) => self.start_heartbeat(node_info, seq).await, Err(cause) => Err(cause.add_message_back("(while cluster api add_node).")), } From eb69f108bce53d99c11b6f52871b68f56a7d681f Mon Sep 17 00:00:00 2001 From: zhang2014 Date: Wed, 18 Dec 2024 16:57:17 +0800 Subject: [PATCH 07/53] feat(cluster): support custom management cluster --- src/query/management/src/warehouse/warehouse_api.rs | 10 +++++++++- src/query/management/src/warehouse/warehouse_mgr.rs | 11 +++++++++++ 2 files changed, 20 insertions(+), 1 deletion(-) diff --git a/src/query/management/src/warehouse/warehouse_api.rs b/src/query/management/src/warehouse/warehouse_api.rs index 7da8055e1823e..848b0e7f9be5c 100644 --- a/src/query/management/src/warehouse/warehouse_api.rs +++ b/src/query/management/src/warehouse/warehouse_api.rs @@ -22,6 +22,8 @@ pub enum SelectedNode { Random(Option), } +pub type SelectedNodes = Vec; + #[derive(serde::Serialize, serde::Deserialize, Eq, PartialEq, Debug)] pub enum WarehouseInfo { SelfManaged(String), @@ -36,6 +38,8 @@ pub struct SystemManagedInfo { pub clusters: HashMap>, } +pub type RemoveNodes = Vec>; + /// Databend-query cluster management API #[async_trait::async_trait] pub trait WarehouseApi: Sync + Send { @@ -50,12 +54,16 @@ pub trait WarehouseApi: Sync + Send { async fn drop_warehouse(&self, warehouse: String) -> Result<()>; - async fn create_warehouse(&self, warehouse: String, nodes: Vec) -> Result<()>; + async fn create_warehouse(&self, warehouse: String, nodes: SelectedNodes) -> Result<()>; async fn list_warehouses(&self) -> Result>; async fn rename_warehouse(&self, cur: String, to: String) -> Result<()>; + async fn add_node(&self, warehouse: &str, cluster: &str, added: SelectedNodes) -> Result<()>; + + async fn remove_node(&self, warehouse: &str, cluster: &str, added: RemoveNodes) -> Result<()>; + /// Get the tenant's cluster all nodes. async fn get_nodes(&self, warehouse: &str, cluster: &str) -> Result>; diff --git a/src/query/management/src/warehouse/warehouse_mgr.rs b/src/query/management/src/warehouse/warehouse_mgr.rs index 8f5c7374ffd7f..c93d6b627275b 100644 --- a/src/query/management/src/warehouse/warehouse_mgr.rs +++ b/src/query/management/src/warehouse/warehouse_mgr.rs @@ -34,7 +34,9 @@ use databend_common_meta_types::TxnOpResponse; use databend_common_meta_types::TxnReply; use databend_common_meta_types::TxnRequest; +use crate::warehouse::warehouse_api::RemoveNodes; use crate::warehouse::warehouse_api::SelectedNode; +use crate::warehouse::warehouse_api::SelectedNodes; use crate::warehouse::warehouse_api::SystemManagedInfo; use crate::warehouse::warehouse_api::WarehouseInfo; use crate::warehouse::WarehouseApi; @@ -899,6 +901,15 @@ impl WarehouseApi for WarehouseMgr { )) } + async fn add_node(&self, warehouse: &str, cluster: &str, nodes: SelectedNodes) -> Result<()> { + // Move node to warehouse + todo!() + } + + async fn remove_node(&self, warehouse: &str, cluster: &str, nodes: RemoveNodes) -> Result<()> { + todo!() + } + #[async_backtrace::framed] #[fastrace::trace] async fn get_nodes(&self, warehouse: &str, cluster: &str) -> Result> { From 25fc50e2150f2bd688a8349b1c4914097f4cbf18 Mon Sep 17 00:00:00 2001 From: zhang2014 Date: Fri, 20 Dec 2024 10:26:23 +0800 Subject: [PATCH 08/53] feat(cluster): support custom management cluster --- src/common/exception/src/exception_code.rs | 2 + src/query/management/src/lib.rs | 1 + src/query/management/src/warehouse/mod.rs | 1 + .../management/src/warehouse/warehouse_api.rs | 34 +- .../management/src/warehouse/warehouse_mgr.rs | 406 +++++++++++++++++- src/query/management/tests/it/warehouse.rs | 30 +- 6 files changed, 447 insertions(+), 27 deletions(-) diff --git a/src/common/exception/src/exception_code.rs b/src/common/exception/src/exception_code.rs index 7fdf7dc3289b2..22131dec3912f 100644 --- a/src/common/exception/src/exception_code.rs +++ b/src/common/exception/src/exception_code.rs @@ -299,6 +299,8 @@ build_exceptions! { UnknownWarehouse(2406), WarehouseOperateConflict(2407), EmptyNodesForWarehouse(2408), + WarehouseClusterAlreadyExists(2409), + WarehouseClusterNotExists(2410), // Stage error codes. UnknownStage(2501), diff --git a/src/query/management/src/lib.rs b/src/query/management/src/lib.rs index 50fa86859a33c..06cfa0eb453fc 100644 --- a/src/query/management/src/lib.rs +++ b/src/query/management/src/lib.rs @@ -51,6 +51,7 @@ pub use stage::StageMgr; pub use user::UserApi; pub use user::UserMgr; pub use warehouse::SelectedNode; +pub use warehouse::SystemManagedCluster; pub use warehouse::SystemManagedInfo; pub use warehouse::WarehouseApi; pub use warehouse::WarehouseInfo; diff --git a/src/query/management/src/warehouse/mod.rs b/src/query/management/src/warehouse/mod.rs index 271155ababfbe..84e1f554e18b0 100644 --- a/src/query/management/src/warehouse/mod.rs +++ b/src/query/management/src/warehouse/mod.rs @@ -16,6 +16,7 @@ mod warehouse_api; mod warehouse_mgr; pub use warehouse_api::SelectedNode; +pub use warehouse_api::SystemManagedCluster; pub use warehouse_api::SystemManagedInfo; pub use warehouse_api::WarehouseApi; pub use warehouse_api::WarehouseInfo; diff --git a/src/query/management/src/warehouse/warehouse_api.rs b/src/query/management/src/warehouse/warehouse_api.rs index 848b0e7f9be5c..19b236fa29a2e 100644 --- a/src/query/management/src/warehouse/warehouse_api.rs +++ b/src/query/management/src/warehouse/warehouse_api.rs @@ -30,16 +30,19 @@ pub enum WarehouseInfo { SystemManaged(SystemManagedInfo), } +#[derive(serde::Serialize, serde::Deserialize, Eq, PartialEq, Debug)] +pub struct SystemManagedCluster { + pub nodes: Vec, +} + #[derive(serde::Serialize, serde::Deserialize, Eq, PartialEq, Debug)] pub struct SystemManagedInfo { pub id: String, pub status: String, pub display_name: String, - pub clusters: HashMap>, + pub clusters: HashMap, } -pub type RemoveNodes = Vec>; - /// Databend-query cluster management API #[async_trait::async_trait] pub trait WarehouseApi: Sync + Send { @@ -59,10 +62,27 @@ pub trait WarehouseApi: Sync + Send { async fn list_warehouses(&self) -> Result>; async fn rename_warehouse(&self, cur: String, to: String) -> Result<()>; - - async fn add_node(&self, warehouse: &str, cluster: &str, added: SelectedNodes) -> Result<()>; - - async fn remove_node(&self, warehouse: &str, cluster: &str, added: RemoveNodes) -> Result<()>; + // async fn list_warehouse_nodes(&self, warehouse: String) -> Result<()>; + // + async fn add_warehouse_cluster( + &self, + warehouse: String, + cluster: String, + nodes: SelectedNodes, + ) -> Result<()>; + + async fn drop_warehouse_cluster(&self, warehouse: String, cluster: String) -> Result<()>; + + async fn rename_warehouse_cluster( + &self, + warehouse: String, + cur: String, + to: String, + ) -> Result<()>; + + // async fn add_warehouse_cluster_node(&self, warehouse: &str, cluster: &str, added: SelectedNodes) -> Result<()>; + // + // async fn remove_warehouse_cluster_node(&self, warehouse: &str, cluster: &str, added: RemoveNodes) -> Result<()>; /// Get the tenant's cluster all nodes. async fn get_nodes(&self, warehouse: &str, cluster: &str) -> Result>; diff --git a/src/query/management/src/warehouse/warehouse_mgr.rs b/src/query/management/src/warehouse/warehouse_mgr.rs index c93d6b627275b..b700336653e8d 100644 --- a/src/query/management/src/warehouse/warehouse_mgr.rs +++ b/src/query/management/src/warehouse/warehouse_mgr.rs @@ -34,9 +34,9 @@ use databend_common_meta_types::TxnOpResponse; use databend_common_meta_types::TxnReply; use databend_common_meta_types::TxnRequest; -use crate::warehouse::warehouse_api::RemoveNodes; use crate::warehouse::warehouse_api::SelectedNode; use crate::warehouse::warehouse_api::SelectedNodes; +use crate::warehouse::warehouse_api::SystemManagedCluster; use crate::warehouse::warehouse_api::SystemManagedInfo; use crate::warehouse::warehouse_api::WarehouseInfo; use crate::warehouse::WarehouseApi; @@ -563,8 +563,10 @@ impl WarehouseApi for WarehouseMgr { } async fn heartbeat(&self, node: &mut NodeInfo, seq: u64) -> Result { - assert!(!node.cluster_id.is_empty()); - assert!(!node.warehouse_id.is_empty()); + if node.node_type == NodeType::SelfManaged { + assert!(!node.cluster_id.is_empty()); + assert!(!node.warehouse_id.is_empty()); + } let res = self.upsert_node(node.clone(), MatchSeq::Exact(seq)).await?; @@ -594,10 +596,7 @@ impl WarehouseApi for WarehouseMgr { for _idx in 0..10 { let consistent_info = self.consistent_warehouse_info(&warehouse).await?; - if matches!( - consistent_info.warehouse_info, - WarehouseInfo::SelfManaged(_) - ) { + if let WarehouseInfo::SelfManaged(_) = consistent_info.warehouse_info { return Err(ErrorCode::InvalidWarehouse( "Cannot drop self-managed warehouse", )); @@ -750,7 +749,9 @@ impl WarehouseApi for WarehouseMgr { id: GlobalUniqName::unique(), status: "Running".to_string(), display_name: warehouse.clone(), - clusters: HashMap::from([(String::from("default"), nodes.clone())]), + clusters: HashMap::from([(String::from("default"), SystemManagedCluster { + nodes: nodes.clone(), + })]), }))?, )); txn.else_then.push(TxnOp::get(warehouse_key)); @@ -838,10 +839,9 @@ impl WarehouseApi for WarehouseMgr { .push(TxnOp::get(new_warehouse_key.clone())); rename_txn.if_then.push(TxnOp::delete(old_warehouse_key)); - rename_txn.if_then.push(TxnOp::put_with_ttl( + rename_txn.if_then.push(TxnOp::put( new_warehouse_key, serde_json::to_vec(&consistent_info.warehouse_info)?, - Some(self.lift_time * 4), )); for mut consistent_node in consistent_info.consistent_nodes { @@ -901,15 +901,391 @@ impl WarehouseApi for WarehouseMgr { )) } - async fn add_node(&self, warehouse: &str, cluster: &str, nodes: SelectedNodes) -> Result<()> { - // Move node to warehouse - todo!() + async fn add_warehouse_cluster( + &self, + warehouse: String, + cluster: String, + nodes: SelectedNodes, + ) -> Result<()> { + if warehouse.is_empty() { + return Err(ErrorCode::InvalidWarehouse("Warehouse name is empty.")); + } + + if cluster.is_empty() { + return Err(ErrorCode::InvalidWarehouse( + "Warehouse cluster name is empty.", + )); + } + + if nodes.is_empty() { + return Err(ErrorCode::EmptyNodesForWarehouse( + "Cannot create warehouse cluster with empty nodes.", + )); + } + + for _idx in 0..10 { + let mut selected_nodes = Vec::with_capacity(nodes.len()); + + // get online nodes + let online_nodes = self.metastore.prefix_list_kv(&self.node_key_prefix).await?; + + let mut select_queue = nodes.clone(); + for (_, v) in &online_nodes { + match select_queue.last() { + None => { + break; + } + Some(SelectedNode::Random(Some(_))) => { + return Err(ErrorCode::Unimplemented( + "Custom instance types are not supported.", + )); + } + Some(SelectedNode::Random(None)) => { + // select random node + + let mut node_info = serde_json::from_slice::(&v.data)?; + + if node_info.warehouse_id.is_empty() && node_info.cluster_id.is_empty() { + node_info.warehouse_id = warehouse.clone(); + node_info.cluster_id = cluster.clone(); + selected_nodes.push((v.seq, node_info)); + select_queue.pop(); + } + } + } + } + + if !select_queue.is_empty() { + return Err(ErrorCode::NoResourcesAvailable( + "Failed to create warehouse cluster, reason: no resources available", + )); + } + + let mut create_cluster_txn = TxnRequest::default(); + + let mut consistent_info = self.consistent_warehouse_info(&warehouse).await?; + + consistent_info.warehouse_info = match consistent_info.warehouse_info { + WarehouseInfo::SelfManaged(_) => Err(ErrorCode::InvalidWarehouse(format!( + "Cannot add cluster for warehouse {:?}, because it's self-managed warehouse.", + warehouse + ))), + WarehouseInfo::SystemManaged(mut info) => { + match info.clusters.contains_key(&cluster) { + true => Err(ErrorCode::WarehouseClusterAlreadyExists(format!( + "Warehouse cluster {:?}.{:?} already exists", + warehouse, cluster + ))), + false => { + info.clusters.insert(cluster.clone(), SystemManagedCluster { + nodes: nodes.clone(), + }); + Ok(WarehouseInfo::SystemManaged(SystemManagedInfo { + id: info.id, + status: info.status, + display_name: info.display_name, + clusters: info.clusters, + })) + } + } + } + }?; + + let warehouse_key = format!( + "{}/{}", + self.warehouse_key_prefix, + escape_for_key(&warehouse)? + ); + + create_cluster_txn.condition.push(map_condition( + &warehouse_key, + MatchSeq::Exact(consistent_info.info_seq), + )); + + create_cluster_txn.if_then.push(TxnOp::put( + warehouse_key.clone(), + serde_json::to_vec(&consistent_info.warehouse_info)?, + )); + + // lock all cluster state + for consistent_node in consistent_info.consistent_nodes { + let node_key = self.node_key(&consistent_node.node_info)?; + let cluster_key = self.cluster_key(&consistent_node.node_info)?; + + create_cluster_txn.condition.push(map_condition( + &node_key, + MatchSeq::Exact(consistent_node.node_seq), + )); + create_cluster_txn.condition.push(map_condition( + &cluster_key, + MatchSeq::Exact(consistent_node.cluster_seq), + )); + } + + for (seq, mut node) in selected_nodes { + let node_key = self.node_key(&node)?; + let cluster_key = self.cluster_key(&node)?; + + create_cluster_txn + .condition + .push(map_condition(&node_key, MatchSeq::Exact(seq))); + create_cluster_txn.if_then.push(TxnOp::put_with_ttl( + node_key, + serde_json::to_vec(&node)?, + Some(self.lift_time * 4), + )); + + node.cluster_id = String::new(); + node.warehouse_id = String::new(); + create_cluster_txn + .condition + .push(map_condition(&cluster_key, MatchSeq::Exact(0))); + create_cluster_txn.if_then.push(TxnOp::put_with_ttl( + cluster_key, + serde_json::to_vec(&node)?, + Some(self.lift_time * 4), + )); + } + + return match self.metastore.transaction(create_cluster_txn).await? { + res if res.success => Ok(()), + _res => { + continue; + } + }; + } + + Err(ErrorCode::WarehouseOperateConflict( + "Warehouse operate conflict(tried 10 times while in create warehouse cluster).", + )) + } + + async fn drop_warehouse_cluster(&self, warehouse: String, cluster: String) -> Result<()> { + if warehouse.is_empty() { + return Err(ErrorCode::InvalidWarehouse("Warehouse name is empty.")); + } + + if cluster.is_empty() { + return Err(ErrorCode::InvalidWarehouse( + "Warehouse cluster name is empty.", + )); + } + + for _idx in 0..10 { + let mut drop_cluster_txn = TxnRequest::default(); + + let mut consistent_info = self.consistent_warehouse_info(&warehouse).await?; + + consistent_info.warehouse_info = match consistent_info.warehouse_info { + WarehouseInfo::SelfManaged(_) => Err(ErrorCode::InvalidWarehouse(format!( + "Cannot add cluster for warehouse {:?}, because it's self-managed warehouse.", + warehouse + ))), + WarehouseInfo::SystemManaged(mut info) => { + match info.clusters.contains_key(&cluster) { + false => Err(ErrorCode::WarehouseClusterNotExists(format!( + "Warehouse cluster {:?}.{:?} not exists", + warehouse, cluster + ))), + true => match info.clusters.len() == 1 { + true => Err(ErrorCode::EmptyNodesForWarehouse(format!( + "Warehouse {:?} only has {:?} one cluster, cannot drop it.", + warehouse, cluster + ))), + false => { + info.clusters.remove(&cluster); + Ok(WarehouseInfo::SystemManaged(SystemManagedInfo { + id: info.id, + status: info.status, + display_name: info.display_name, + clusters: info.clusters, + })) + } + }, + } + } + }?; + + let warehouse_key = format!( + "{}/{}", + self.warehouse_key_prefix, + escape_for_key(&warehouse)? + ); + + drop_cluster_txn.condition.push(map_condition( + &warehouse_key, + MatchSeq::Exact(consistent_info.info_seq), + )); + + drop_cluster_txn.if_then.push(TxnOp::put( + warehouse_key.clone(), + serde_json::to_vec(&consistent_info.warehouse_info)?, + )); + + // lock all cluster state + for mut consistent_node in consistent_info.consistent_nodes { + let node_key = self.node_key(&consistent_node.node_info)?; + let cluster_key = self.cluster_key(&consistent_node.node_info)?; + + drop_cluster_txn.condition.push(map_condition( + &node_key, + MatchSeq::Exact(consistent_node.node_seq), + )); + drop_cluster_txn.condition.push(map_condition( + &cluster_key, + MatchSeq::Exact(consistent_node.cluster_seq), + )); + + if consistent_node.node_info.cluster_id == cluster { + // Remove node + consistent_node.node_info.cluster_id = String::new(); + consistent_node.node_info.warehouse_id = String::new(); + + drop_cluster_txn.if_then.push(TxnOp::delete(cluster_key)); + drop_cluster_txn.if_then.push(TxnOp::put_with_ttl( + node_key, + serde_json::to_vec(&consistent_node.node_info)?, + Some(self.lift_time * 4), + )) + } + } + + return match self.metastore.transaction(drop_cluster_txn).await? { + res if res.success => Ok(()), + _ => { + continue; + } + }; + } + + Err(ErrorCode::WarehouseOperateConflict( + "Warehouse operate conflict(tried 10 times while in drop warehouse cluster).", + )) } - async fn remove_node(&self, warehouse: &str, cluster: &str, nodes: RemoveNodes) -> Result<()> { - todo!() + async fn rename_warehouse_cluster( + &self, + warehouse: String, + cur: String, + to: String, + ) -> Result<()> { + if warehouse.is_empty() { + return Err(ErrorCode::InvalidWarehouse("Warehouse name is empty.")); + } + + if cur.is_empty() { + return Err(ErrorCode::InvalidWarehouse( + "Warehouse cluster name is empty.", + )); + } + + if to.is_empty() { + return Err(ErrorCode::InvalidWarehouse( + "Warehouse cluster name is empty.", + )); + } + + for _idx in 0..10 { + let mut rename_cluster_txn = TxnRequest::default(); + + let mut consistent_info = self.consistent_warehouse_info(&warehouse).await?; + + consistent_info.warehouse_info = match consistent_info.warehouse_info { + WarehouseInfo::SelfManaged(_) => Err(ErrorCode::InvalidWarehouse(format!("Cannot rename cluster for warehouse {:?}, because it's self-managed warehouse.", warehouse))), + WarehouseInfo::SystemManaged(mut info) => match info.clusters.contains_key(&cur) { + false => Err(ErrorCode::WarehouseClusterNotExists(format!("Warehouse cluster {:?}.{:?} not exists", warehouse, cur))), + true => { + let cluster_info = info.clusters.remove(&cur); + info.clusters.insert(to.clone(), cluster_info.unwrap()); + Ok(WarehouseInfo::SystemManaged(SystemManagedInfo { + id: info.id, + status: info.status, + display_name: info.display_name, + clusters: info.clusters, + })) + } + } + }?; + + let warehouse_key = format!( + "{}/{}", + self.warehouse_key_prefix, + escape_for_key(&warehouse)? + ); + + rename_cluster_txn.condition.push(map_condition( + &warehouse_key, + MatchSeq::Exact(consistent_info.info_seq), + )); + + rename_cluster_txn.if_then.push(TxnOp::put( + warehouse_key.clone(), + serde_json::to_vec(&consistent_info.warehouse_info)?, + )); + + // lock all cluster state + for mut consistent_node in consistent_info.consistent_nodes { + let node_key = self.node_key(&consistent_node.node_info)?; + let old_cluster_key = self.cluster_key(&consistent_node.node_info)?; + + rename_cluster_txn.condition.push(map_condition( + &node_key, + MatchSeq::Exact(consistent_node.node_seq), + )); + rename_cluster_txn.condition.push(map_condition( + &old_cluster_key, + MatchSeq::Exact(consistent_node.cluster_seq), + )); + + if consistent_node.node_info.cluster_id == cur { + // rename node + consistent_node.node_info.cluster_id = to.clone(); + let new_cluster_key = self.cluster_key(&consistent_node.node_info)?; + + rename_cluster_txn.if_then.push(TxnOp::put_with_ttl( + node_key, + serde_json::to_vec(&consistent_node.node_info)?, + Some(self.lift_time * 4), + )); + + consistent_node.node_info.cluster_id = String::new(); + rename_cluster_txn + .condition + .push(map_condition(&new_cluster_key, MatchSeq::Exact(0))); + rename_cluster_txn + .if_then + .push(TxnOp::delete(old_cluster_key)); + rename_cluster_txn.if_then.push(TxnOp::put_with_ttl( + new_cluster_key, + serde_json::to_vec(&consistent_node.node_info)?, + Some(self.lift_time * 4), + )); + } + } + + return match self.metastore.transaction(rename_cluster_txn).await? { + res if res.success => Ok(()), + _ => { + continue; + } + }; + } + + Err(ErrorCode::WarehouseOperateConflict( + "Warehouse operate conflict(tried 10 times while in rename warehouse cluster).", + )) } + // async fn add_warehouse_cluster_node(&self, warehouse: &str, cluster: &str, nodes: SelectedNodes) -> Result<()> { + // // Move node to warehouse + // // self.consistent_warehouse_info() + // todo!() + // } + // + // async fn remove_warehouse_cluster_node(&self, warehouse: &str, cluster: &str, nodes: RemoveNodes) -> Result<()> { + // todo!() + // } + #[async_backtrace::framed] #[fastrace::trace] async fn get_nodes(&self, warehouse: &str, cluster: &str) -> Result> { diff --git a/src/query/management/tests/it/warehouse.rs b/src/query/management/tests/it/warehouse.rs index a730c2fb702f9..e436c78f0f33d 100644 --- a/src/query/management/tests/it/warehouse.rs +++ b/src/query/management/tests/it/warehouse.rs @@ -603,7 +603,9 @@ async fn test_list_warehouses() -> Result<()> { assert_eq!(system_managed_info.display_name, "test_warehouse_1"); assert_eq!( system_managed_info.clusters, - HashMap::from([(String::from("default"), vec![SelectedNode::Random(None)])]) + HashMap::from([(String::from("default"), SystemManagedCluster { + nodes: vec![SelectedNode::Random(None)] + })]) ); let self_managed_node_2 = GlobalUniqName::unique(); @@ -643,7 +645,9 @@ async fn test_list_warehouses() -> Result<()> { assert_eq!(system_managed_info.display_name, "test_warehouse_3"); assert_eq!( system_managed_info.clusters, - HashMap::from([(String::from("default"), vec![SelectedNode::Random(None)])]) + HashMap::from([(String::from("default"), SystemManagedCluster { + nodes: vec![SelectedNode::Random(None)] + })]) ); warehouse_manager.shutdown_node(self_managed_node_1).await?; @@ -690,6 +694,16 @@ async fn test_list_warehouses() -> Result<()> { Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn test_rename_not_exists_warehouses() -> Result<()> { + let (_, warehouse_manager, _) = nodes(Duration::from_mins(30), 1).await?; + let rename_warehouse = + warehouse_manager.rename_warehouse(String::from("test_warehouse"), String::from("aa")); + + assert_eq!(rename_warehouse.await.unwrap_err().code(), 2406); + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 1)] async fn test_rename_warehouses() -> Result<()> { let (kv, warehouse_manager, nodes) = nodes(Duration::from_mins(30), 1).await?; @@ -732,7 +746,9 @@ async fn test_rename_warehouses() -> Result<()> { assert_eq!(system_managed_info.display_name, "test_warehouse"); assert_eq!( system_managed_info.clusters, - HashMap::from([(String::from("default"), vec![SelectedNode::Random(None)])]) + HashMap::from([(String::from("default"), SystemManagedCluster { + nodes: vec![SelectedNode::Random(None)] + })]) ); let rename_warehouse = warehouse_manager.rename_warehouse( @@ -755,7 +771,9 @@ async fn test_rename_warehouses() -> Result<()> { assert_eq!(system_managed_info.display_name, "new_test_warehouse"); assert_eq!( system_managed_info.clusters, - HashMap::from([(String::from("default"), vec![SelectedNode::Random(None)])]) + HashMap::from([(String::from("default"), SystemManagedCluster { + nodes: vec![SelectedNode::Random(None)] + })]) ); let system_managed_node_2 = GlobalUniqName::unique(); @@ -782,7 +800,9 @@ async fn test_rename_warehouses() -> Result<()> { assert_eq!(system_managed_info.display_name, "test_warehouse"); assert_eq!( system_managed_info.clusters, - HashMap::from([(String::from("default"), vec![SelectedNode::Random(None)])]) + HashMap::from([(String::from("default"), SystemManagedCluster { + nodes: vec![SelectedNode::Random(None)] + })]) ); let rename_warehouse = warehouse_manager.rename_warehouse( From e3af0fb3c6e6af87c285ae91ecf46b3fc758f7f8 Mon Sep 17 00:00:00 2001 From: zhang2014 Date: Fri, 20 Dec 2024 12:51:51 +0800 Subject: [PATCH 09/53] feat(cluster): support custom management cluster --- .../management/src/warehouse/warehouse_api.rs | 21 +- .../management/src/warehouse/warehouse_mgr.rs | 302 +++++++++++++++++- 2 files changed, 309 insertions(+), 14 deletions(-) diff --git a/src/query/management/src/warehouse/warehouse_api.rs b/src/query/management/src/warehouse/warehouse_api.rs index 19b236fa29a2e..77acd9b855bea 100644 --- a/src/query/management/src/warehouse/warehouse_api.rs +++ b/src/query/management/src/warehouse/warehouse_api.rs @@ -62,8 +62,9 @@ pub trait WarehouseApi: Sync + Send { async fn list_warehouses(&self) -> Result>; async fn rename_warehouse(&self, cur: String, to: String) -> Result<()>; - // async fn list_warehouse_nodes(&self, warehouse: String) -> Result<()>; - // + + async fn list_warehouse_nodes(&self, warehouse: String) -> Result>; + async fn add_warehouse_cluster( &self, warehouse: String, @@ -80,9 +81,19 @@ pub trait WarehouseApi: Sync + Send { to: String, ) -> Result<()>; - // async fn add_warehouse_cluster_node(&self, warehouse: &str, cluster: &str, added: SelectedNodes) -> Result<()>; - // - // async fn remove_warehouse_cluster_node(&self, warehouse: &str, cluster: &str, added: RemoveNodes) -> Result<()>; + async fn add_warehouse_cluster_node( + &self, + warehouse: &str, + cluster: &str, + add_nodes: SelectedNodes, + ) -> Result<()>; + + async fn drop_warehouse_cluster_node( + &self, + warehouse: &str, + cluster: &str, + drop_nodes: Vec, + ) -> Result<()>; /// Get the tenant's cluster all nodes. async fn get_nodes(&self, warehouse: &str, cluster: &str) -> Result>; diff --git a/src/query/management/src/warehouse/warehouse_mgr.rs b/src/query/management/src/warehouse/warehouse_mgr.rs index b700336653e8d..16ee9ff3e670d 100644 --- a/src/query/management/src/warehouse/warehouse_mgr.rs +++ b/src/query/management/src/warehouse/warehouse_mgr.rs @@ -901,6 +901,20 @@ impl WarehouseApi for WarehouseMgr { )) } + async fn list_warehouse_nodes(&self, warehouse: String) -> Result> { + if warehouse.is_empty() { + return Err(ErrorCode::InvalidWarehouse("Warehouse name is empty.")); + } + + let consistent_info = self.consistent_warehouse_info(&warehouse).await?; + + Ok(consistent_info + .consistent_nodes + .into_iter() + .map(|x| x.node_info) + .collect()) + } + async fn add_warehouse_cluster( &self, warehouse: String, @@ -1276,15 +1290,285 @@ impl WarehouseApi for WarehouseMgr { )) } - // async fn add_warehouse_cluster_node(&self, warehouse: &str, cluster: &str, nodes: SelectedNodes) -> Result<()> { - // // Move node to warehouse - // // self.consistent_warehouse_info() - // todo!() - // } - // - // async fn remove_warehouse_cluster_node(&self, warehouse: &str, cluster: &str, nodes: RemoveNodes) -> Result<()> { - // todo!() - // } + async fn add_warehouse_cluster_node( + &self, + warehouse: &str, + cluster: &str, + nodes: SelectedNodes, + ) -> Result<()> { + if warehouse.is_empty() { + return Err(ErrorCode::InvalidWarehouse("Warehouse name is empty.")); + } + + if cluster.is_empty() { + return Err(ErrorCode::InvalidWarehouse( + "Warehouse cluster name is empty.", + )); + } + + if nodes.is_empty() { + return Err(ErrorCode::EmptyNodesForWarehouse( + "Cannot create warehouse cluster with empty nodes.", + )); + } + + for _idx in 0..10 { + let mut selected_nodes = Vec::with_capacity(nodes.len()); + + // get online nodes + let online_nodes = self.metastore.prefix_list_kv(&self.node_key_prefix).await?; + + let mut select_queue = nodes.clone(); + for (_, v) in &online_nodes { + match select_queue.last() { + None => { + break; + } + Some(SelectedNode::Random(Some(_))) => { + return Err(ErrorCode::Unimplemented( + "Custom instance types are not supported.", + )); + } + Some(SelectedNode::Random(None)) => { + // select random node + + let mut node_info = serde_json::from_slice::(&v.data)?; + + if node_info.warehouse_id.is_empty() && node_info.cluster_id.is_empty() { + node_info.cluster_id = cluster.to_string(); + node_info.warehouse_id = warehouse.to_string(); + selected_nodes.push((v.seq, node_info)); + select_queue.pop(); + } + } + } + } + + if !select_queue.is_empty() { + return Err(ErrorCode::NoResourcesAvailable( + "Failed to add warehouse cluster node, reason: no resources available", + )); + } + + let mut add_cluster_node_txn = TxnRequest::default(); + + let mut consistent_info = self.consistent_warehouse_info(warehouse).await?; + + consistent_info.warehouse_info = match consistent_info.warehouse_info { + WarehouseInfo::SelfManaged(_) => Err(ErrorCode::InvalidWarehouse(format!( + "Cannot add cluster node for warehouse {:?}, because it's self-managed warehouse.", + warehouse + ))), + WarehouseInfo::SystemManaged(mut info) => match info.clusters.get_mut(cluster) { + None => Err(ErrorCode::WarehouseClusterNotExists(format!( + "Warehouse cluster {:?}.{:?} not exists", + warehouse, cluster + ))), + Some(cluster_info) => { + cluster_info.nodes.extend(nodes.clone()); + Ok(WarehouseInfo::SystemManaged(SystemManagedInfo { + id: info.id, + status: info.status, + display_name: info.display_name, + clusters: info.clusters, + })) + } + }, + }?; + + let warehouse_key = format!( + "{}/{}", + self.warehouse_key_prefix, + escape_for_key(warehouse)? + ); + + add_cluster_node_txn.condition.push(map_condition( + &warehouse_key, + MatchSeq::Exact(consistent_info.info_seq), + )); + + add_cluster_node_txn.if_then.push(TxnOp::put( + warehouse_key.clone(), + serde_json::to_vec(&consistent_info.warehouse_info)?, + )); + + // lock all cluster state + for consistent_node in consistent_info.consistent_nodes { + let node_key = self.node_key(&consistent_node.node_info)?; + let cluster_key = self.cluster_key(&consistent_node.node_info)?; + + add_cluster_node_txn.condition.push(map_condition( + &node_key, + MatchSeq::Exact(consistent_node.node_seq), + )); + add_cluster_node_txn.condition.push(map_condition( + &cluster_key, + MatchSeq::Exact(consistent_node.cluster_seq), + )); + } + + for (seq, mut node) in selected_nodes { + let node_key = self.node_key(&node)?; + let cluster_key = self.cluster_key(&node)?; + + add_cluster_node_txn + .condition + .push(map_condition(&node_key, MatchSeq::Exact(seq))); + add_cluster_node_txn.if_then.push(TxnOp::put_with_ttl( + node_key, + serde_json::to_vec(&node)?, + Some(self.lift_time * 4), + )); + + node.cluster_id = String::new(); + node.warehouse_id = String::new(); + add_cluster_node_txn + .condition + .push(map_condition(&cluster_key, MatchSeq::Exact(0))); + add_cluster_node_txn.if_then.push(TxnOp::put_with_ttl( + cluster_key, + serde_json::to_vec(&node)?, + Some(self.lift_time * 4), + )); + } + + return match self.metastore.transaction(add_cluster_node_txn).await? { + res if res.success => Ok(()), + _res => { + continue; + } + }; + } + + Err(ErrorCode::WarehouseOperateConflict( + "Warehouse operate conflict(tried 10 times while in add warehouse cluster node).", + )) + } + + async fn drop_warehouse_cluster_node( + &self, + warehouse: &str, + cluster: &str, + nodes: Vec, + ) -> Result<()> { + if warehouse.is_empty() { + return Err(ErrorCode::InvalidWarehouse("Warehouse name is empty.")); + } + + if cluster.is_empty() { + return Err(ErrorCode::InvalidWarehouse( + "Warehouse cluster name is empty.", + )); + } + + for _idx in 0..10 { + let mut drop_cluster_node_txn = TxnRequest::default(); + + let mut consistent_info = self.consistent_warehouse_info(warehouse).await?; + + consistent_info.warehouse_info = match consistent_info.warehouse_info { + WarehouseInfo::SelfManaged(_) => Err(ErrorCode::InvalidWarehouse(format!( + "Cannot add cluster for warehouse {:?}, because it's self-managed warehouse.", + warehouse + ))), + WarehouseInfo::SystemManaged(mut info) => match info.clusters.get_mut(cluster) { + None => Err(ErrorCode::WarehouseClusterNotExists(format!( + "Warehouse cluster {:?}.{:?} not exists", + warehouse, cluster + ))), + Some(cluster) => match nodes.len() == cluster.nodes.len() { + true => Err(ErrorCode::EmptyNodesForWarehouse(format!( + "Warehouse cluster {:?}.{:?} only has {} nodes, cannot drop all.", + warehouse, + cluster, + nodes.len() + ))), + false => { + for remove_node in &nodes { + if consistent_info + .consistent_nodes + .iter() + .any(|x| &x.node_info.id == remove_node) + { + cluster.nodes.pop(); + continue; + } + + return Err(ErrorCode::ClusterUnknownNode(format!( + "Warehouse cluster {:?}.{:?} unknwon node {:?}", + warehouse, cluster, remove_node + ))); + } + + Ok(WarehouseInfo::SystemManaged(SystemManagedInfo { + id: info.id, + status: info.status, + display_name: info.display_name, + clusters: info.clusters, + })) + } + }, + }, + }?; + + let warehouse_key = format!( + "{}/{}", + self.warehouse_key_prefix, + escape_for_key(warehouse)? + ); + + drop_cluster_node_txn.condition.push(map_condition( + &warehouse_key, + MatchSeq::Exact(consistent_info.info_seq), + )); + + drop_cluster_node_txn.if_then.push(TxnOp::put( + warehouse_key.clone(), + serde_json::to_vec(&consistent_info.warehouse_info)?, + )); + + // lock all cluster state + for mut consistent_node in consistent_info.consistent_nodes { + let node_key = self.node_key(&consistent_node.node_info)?; + let cluster_key = self.cluster_key(&consistent_node.node_info)?; + + drop_cluster_node_txn.condition.push(map_condition( + &node_key, + MatchSeq::Exact(consistent_node.node_seq), + )); + drop_cluster_node_txn.condition.push(map_condition( + &cluster_key, + MatchSeq::Exact(consistent_node.cluster_seq), + )); + + if nodes.contains(&consistent_node.node_info.id) { + // Remove node + consistent_node.node_info.cluster_id = String::new(); + consistent_node.node_info.warehouse_id = String::new(); + + drop_cluster_node_txn + .if_then + .push(TxnOp::delete(cluster_key)); + drop_cluster_node_txn.if_then.push(TxnOp::put_with_ttl( + node_key, + serde_json::to_vec(&consistent_node.node_info)?, + Some(self.lift_time * 4), + )) + } + } + + return match self.metastore.transaction(drop_cluster_node_txn).await? { + res if res.success => Ok(()), + _ => { + continue; + } + }; + } + + Err(ErrorCode::WarehouseOperateConflict( + "Warehouse operate conflict(tried 10 times while in drop warehouse cluster node).", + )) + } #[async_backtrace::framed] #[fastrace::trace] From 4de94828dfccda79fe652a16a750a3cb48ce209d Mon Sep 17 00:00:00 2001 From: zhang2014 Date: Fri, 20 Dec 2024 12:59:03 +0800 Subject: [PATCH 10/53] feat(cluster): support custom management cluster --- .../management/src/warehouse/warehouse_api.rs | 6 ++++- .../management/src/warehouse/warehouse_mgr.rs | 6 ++++- src/query/management/tests/it/warehouse.rs | 26 ++++++++++++------- src/query/service/src/clusters/cluster.rs | 4 +-- 4 files changed, 29 insertions(+), 13 deletions(-) diff --git a/src/query/management/src/warehouse/warehouse_api.rs b/src/query/management/src/warehouse/warehouse_api.rs index 77acd9b855bea..63cafae9405d7 100644 --- a/src/query/management/src/warehouse/warehouse_api.rs +++ b/src/query/management/src/warehouse/warehouse_api.rs @@ -96,7 +96,11 @@ pub trait WarehouseApi: Sync + Send { ) -> Result<()>; /// Get the tenant's cluster all nodes. - async fn get_nodes(&self, warehouse: &str, cluster: &str) -> Result>; + async fn list_warehouse_cluster_nodes( + &self, + warehouse: &str, + cluster: &str, + ) -> Result>; async fn get_local_addr(&self) -> Result>; diff --git a/src/query/management/src/warehouse/warehouse_mgr.rs b/src/query/management/src/warehouse/warehouse_mgr.rs index 16ee9ff3e670d..9bb90b0741a9f 100644 --- a/src/query/management/src/warehouse/warehouse_mgr.rs +++ b/src/query/management/src/warehouse/warehouse_mgr.rs @@ -1572,7 +1572,11 @@ impl WarehouseApi for WarehouseMgr { #[async_backtrace::framed] #[fastrace::trace] - async fn get_nodes(&self, warehouse: &str, cluster: &str) -> Result> { + async fn list_warehouse_cluster_nodes( + &self, + warehouse: &str, + cluster: &str, + ) -> Result> { let cluster_key = format!( "{}/{}/{}", self.meta_key_prefix, diff --git a/src/query/management/tests/it/warehouse.rs b/src/query/management/tests/it/warehouse.rs index e436c78f0f33d..db57ad8aed857 100644 --- a/src/query/management/tests/it/warehouse.rs +++ b/src/query/management/tests/it/warehouse.rs @@ -134,21 +134,24 @@ async fn test_already_exists_add_self_managed_node() -> Result<()> { async fn test_successfully_get_self_managed_nodes() -> Result<()> { let (_kv, warehouse_manager, _nodes) = nodes(Duration::from_mins(60), 0).await?; - let get_nodes = warehouse_manager.get_nodes("test-cluster-id", "test-cluster-id"); + let get_nodes = + warehouse_manager.list_warehouse_cluster_nodes("test-cluster-id", "test-cluster-id"); assert_eq!(get_nodes.await?, vec![]); let node_1 = self_managed_node("node_1"); warehouse_manager.start_node(node_1.clone()).await?; - let get_nodes = warehouse_manager.get_nodes("test-cluster-id", "test-cluster-id"); + let get_nodes = + warehouse_manager.list_warehouse_cluster_nodes("test-cluster-id", "test-cluster-id"); assert_eq!(get_nodes.await?, vec![node_1.clone()]); let node_2 = self_managed_node("node_2"); warehouse_manager.start_node(node_2.clone()).await?; - let get_nodes = warehouse_manager.get_nodes("test-cluster-id", "test-cluster-id"); + let get_nodes = + warehouse_manager.list_warehouse_cluster_nodes("test-cluster-id", "test-cluster-id"); assert_eq!(get_nodes.await?, vec![node_1, node_2]); Ok(()) @@ -161,13 +164,15 @@ async fn test_successfully_drop_self_managed_node() -> Result<()> { let node_info = self_managed_node("test_node"); warehouse_manager.start_node(node_info.clone()).await?; - let get_nodes = warehouse_manager.get_nodes("test-cluster-id", "test-cluster-id"); + let get_nodes = + warehouse_manager.list_warehouse_cluster_nodes("test-cluster-id", "test-cluster-id"); assert_eq!(get_nodes.await?, vec![node_info.clone()]); warehouse_manager.shutdown_node(node_info.id).await?; - let get_nodes = warehouse_manager.get_nodes("test-cluster-id", "test-cluster-id"); + let get_nodes = + warehouse_manager.list_warehouse_cluster_nodes("test-cluster-id", "test-cluster-id"); assert_eq!(get_nodes.await?, vec![]); Ok(()) @@ -290,7 +295,8 @@ async fn test_successfully_create_system_managed_warehouse() -> Result<()> { assert_key_seq(&kv, &warehouse_node, MatchSeq::GE(1)).await; } - let get_warehouse_nodes = warehouse_manager.get_nodes("test_warehouse", "default"); + let get_warehouse_nodes = + warehouse_manager.list_warehouse_cluster_nodes("test_warehouse", "default"); let warehouse_nodes = get_warehouse_nodes.await?; @@ -332,7 +338,8 @@ async fn test_create_system_managed_warehouse_with_offline_node() -> Result<()> create_warehouse.await?; - let get_warehouse_nodes = warehouse_manager.get_nodes("test_warehouse", "default"); + let get_warehouse_nodes = + warehouse_manager.list_warehouse_cluster_nodes("test_warehouse", "default"); let warehouse_nodes = get_warehouse_nodes.await?; @@ -377,7 +384,8 @@ async fn test_create_system_managed_warehouse_with_online_node() -> Result<()> { create_warehouse.await?; - let get_warehouse_nodes = warehouse_manager.get_nodes("test_warehouse", "default"); + let get_warehouse_nodes = + warehouse_manager.list_warehouse_cluster_nodes("test_warehouse", "default"); let warehouse_nodes = get_warehouse_nodes.await?; @@ -675,7 +683,7 @@ async fn test_list_warehouses() -> Result<()> { // keep show warehouse if all node offline let nodes = warehouse_manager - .get_nodes("test_warehouse_3", "default") + .list_warehouse_cluster_nodes("test_warehouse_3", "default") .await?; warehouse_manager.shutdown_node(nodes[0].id.clone()).await?; assert_eq!( diff --git a/src/query/service/src/clusters/cluster.rs b/src/query/service/src/clusters/cluster.rs index 6a1f117c40fbb..c1719fe338267 100644 --- a/src/query/service/src/clusters/cluster.rs +++ b/src/query/service/src/clusters/cluster.rs @@ -249,7 +249,7 @@ impl ClusterDiscovery { pub async fn discover(&self, config: &InnerConfig) -> Result> { match self .api_provider - .get_nodes(&self.cluster_id, &self.cluster_id) + .list_warehouse_cluster_nodes(&self.cluster_id, &self.cluster_id) .await { Err(cause) => { @@ -326,7 +326,7 @@ impl ClusterDiscovery { async fn drop_invalid_nodes(self: &Arc, node_info: &NodeInfo) -> Result<()> { let current_nodes_info = match self .api_provider - .get_nodes(&node_info.warehouse_id, &node_info.cluster_id) + .list_warehouse_cluster_nodes(&node_info.warehouse_id, &node_info.cluster_id) .await { Ok(nodes) => nodes, From 1f12c1226634add901fa28184947fa8f2e9fc273 Mon Sep 17 00:00:00 2001 From: zhang2014 Date: Fri, 20 Dec 2024 13:00:27 +0800 Subject: [PATCH 11/53] feat(cluster): support custom management cluster --- src/query/management/src/lib.rs | 2 +- src/query/management/src/warehouse/mod.rs | 2 +- .../management/src/warehouse/warehouse_api.rs | 4 ++-- .../management/src/warehouse/warehouse_mgr.rs | 14 +++++++------- 4 files changed, 11 insertions(+), 11 deletions(-) diff --git a/src/query/management/src/lib.rs b/src/query/management/src/lib.rs index 06cfa0eb453fc..fff9251e23837 100644 --- a/src/query/management/src/lib.rs +++ b/src/query/management/src/lib.rs @@ -52,7 +52,7 @@ pub use user::UserApi; pub use user::UserMgr; pub use warehouse::SelectedNode; pub use warehouse::SystemManagedCluster; -pub use warehouse::SystemManagedInfo; +pub use warehouse::SystemManagedWarehouse; pub use warehouse::WarehouseApi; pub use warehouse::WarehouseInfo; pub use warehouse::WarehouseMgr; diff --git a/src/query/management/src/warehouse/mod.rs b/src/query/management/src/warehouse/mod.rs index 84e1f554e18b0..154a5b0c029cc 100644 --- a/src/query/management/src/warehouse/mod.rs +++ b/src/query/management/src/warehouse/mod.rs @@ -17,7 +17,7 @@ mod warehouse_mgr; pub use warehouse_api::SelectedNode; pub use warehouse_api::SystemManagedCluster; -pub use warehouse_api::SystemManagedInfo; +pub use warehouse_api::SystemManagedWarehouse; pub use warehouse_api::WarehouseApi; pub use warehouse_api::WarehouseInfo; pub use warehouse_mgr::WarehouseMgr; diff --git a/src/query/management/src/warehouse/warehouse_api.rs b/src/query/management/src/warehouse/warehouse_api.rs index 63cafae9405d7..c3d25a51f699a 100644 --- a/src/query/management/src/warehouse/warehouse_api.rs +++ b/src/query/management/src/warehouse/warehouse_api.rs @@ -27,7 +27,7 @@ pub type SelectedNodes = Vec; #[derive(serde::Serialize, serde::Deserialize, Eq, PartialEq, Debug)] pub enum WarehouseInfo { SelfManaged(String), - SystemManaged(SystemManagedInfo), + SystemManaged(SystemManagedWarehouse), } #[derive(serde::Serialize, serde::Deserialize, Eq, PartialEq, Debug)] @@ -36,7 +36,7 @@ pub struct SystemManagedCluster { } #[derive(serde::Serialize, serde::Deserialize, Eq, PartialEq, Debug)] -pub struct SystemManagedInfo { +pub struct SystemManagedWarehouse { pub id: String, pub status: String, pub display_name: String, diff --git a/src/query/management/src/warehouse/warehouse_mgr.rs b/src/query/management/src/warehouse/warehouse_mgr.rs index 9bb90b0741a9f..345f7d7d6327b 100644 --- a/src/query/management/src/warehouse/warehouse_mgr.rs +++ b/src/query/management/src/warehouse/warehouse_mgr.rs @@ -37,7 +37,7 @@ use databend_common_meta_types::TxnRequest; use crate::warehouse::warehouse_api::SelectedNode; use crate::warehouse::warehouse_api::SelectedNodes; use crate::warehouse::warehouse_api::SystemManagedCluster; -use crate::warehouse::warehouse_api::SystemManagedInfo; +use crate::warehouse::warehouse_api::SystemManagedWarehouse; use crate::warehouse::warehouse_api::WarehouseInfo; use crate::warehouse::WarehouseApi; @@ -745,7 +745,7 @@ impl WarehouseApi for WarehouseMgr { .push(map_condition(&warehouse_key, MatchSeq::Exact(0))); txn.if_then.push(TxnOp::put( warehouse_key.clone(), - serde_json::to_vec(&WarehouseInfo::SystemManaged(SystemManagedInfo { + serde_json::to_vec(&WarehouseInfo::SystemManaged(SystemManagedWarehouse { id: GlobalUniqName::unique(), status: "Running".to_string(), display_name: warehouse.clone(), @@ -994,7 +994,7 @@ impl WarehouseApi for WarehouseMgr { info.clusters.insert(cluster.clone(), SystemManagedCluster { nodes: nodes.clone(), }); - Ok(WarehouseInfo::SystemManaged(SystemManagedInfo { + Ok(WarehouseInfo::SystemManaged(SystemManagedWarehouse { id: info.id, status: info.status, display_name: info.display_name, @@ -1108,7 +1108,7 @@ impl WarehouseApi for WarehouseMgr { ))), false => { info.clusters.remove(&cluster); - Ok(WarehouseInfo::SystemManaged(SystemManagedInfo { + Ok(WarehouseInfo::SystemManaged(SystemManagedWarehouse { id: info.id, status: info.status, display_name: info.display_name, @@ -1211,7 +1211,7 @@ impl WarehouseApi for WarehouseMgr { true => { let cluster_info = info.clusters.remove(&cur); info.clusters.insert(to.clone(), cluster_info.unwrap()); - Ok(WarehouseInfo::SystemManaged(SystemManagedInfo { + Ok(WarehouseInfo::SystemManaged(SystemManagedWarehouse { id: info.id, status: info.status, display_name: info.display_name, @@ -1366,7 +1366,7 @@ impl WarehouseApi for WarehouseMgr { ))), Some(cluster_info) => { cluster_info.nodes.extend(nodes.clone()); - Ok(WarehouseInfo::SystemManaged(SystemManagedInfo { + Ok(WarehouseInfo::SystemManaged(SystemManagedWarehouse { id: info.id, status: info.status, display_name: info.display_name, @@ -1500,7 +1500,7 @@ impl WarehouseApi for WarehouseMgr { ))); } - Ok(WarehouseInfo::SystemManaged(SystemManagedInfo { + Ok(WarehouseInfo::SystemManaged(SystemManagedWarehouse { id: info.id, status: info.status, display_name: info.display_name, From eb12f849f5ca218a1de5376f315ac3332cdcfd8f Mon Sep 17 00:00:00 2001 From: zhang2014 Date: Fri, 20 Dec 2024 13:10:49 +0800 Subject: [PATCH 12/53] feat(cluster): support custom management cluster --- src/query/management/src/warehouse/warehouse_api.rs | 2 +- src/query/management/src/warehouse/warehouse_mgr.rs | 2 +- src/query/management/tests/it/warehouse.rs | 2 +- src/query/service/src/clusters/cluster.rs | 2 +- 4 files changed, 4 insertions(+), 4 deletions(-) diff --git a/src/query/management/src/warehouse/warehouse_api.rs b/src/query/management/src/warehouse/warehouse_api.rs index c3d25a51f699a..3e9b2a1809e59 100644 --- a/src/query/management/src/warehouse/warehouse_api.rs +++ b/src/query/management/src/warehouse/warehouse_api.rs @@ -53,7 +53,7 @@ pub trait WarehouseApi: Sync + Send { async fn shutdown_node(&self, node_id: String) -> Result<()>; /// Keep the tenant's cluster node alive. - async fn heartbeat(&self, node: &mut NodeInfo, seq: u64) -> Result; + async fn heartbeat_node(&self, node: &mut NodeInfo, seq: u64) -> Result; async fn drop_warehouse(&self, warehouse: String) -> Result<()>; diff --git a/src/query/management/src/warehouse/warehouse_mgr.rs b/src/query/management/src/warehouse/warehouse_mgr.rs index 345f7d7d6327b..fcc008aeb9f76 100644 --- a/src/query/management/src/warehouse/warehouse_mgr.rs +++ b/src/query/management/src/warehouse/warehouse_mgr.rs @@ -562,7 +562,7 @@ impl WarehouseApi for WarehouseMgr { ))) } - async fn heartbeat(&self, node: &mut NodeInfo, seq: u64) -> Result { + async fn heartbeat_node(&self, node: &mut NodeInfo, seq: u64) -> Result { if node.node_type == NodeType::SelfManaged { assert!(!node.cluster_id.is_empty()); assert!(!node.warehouse_id.is_empty()); diff --git a/src/query/management/tests/it/warehouse.rs b/src/query/management/tests/it/warehouse.rs index db57ad8aed857..1776ade64fd34 100644 --- a/src/query/management/tests/it/warehouse.rs +++ b/src/query/management/tests/it/warehouse.rs @@ -230,7 +230,7 @@ async fn test_successfully_heartbeat_self_managed_node() -> Result<()> { assert_key_value(&kv, warehouse_info_key, info.clone()).await; assert_key_expire(&kv, warehouse_info_key, Duration::from_mins(50)).await; - warehouse_manager.heartbeat(&mut node_info, seq).await?; + warehouse_manager.heartbeat_node(&mut node_info, seq).await?; assert_key_value(&kv, warehouse_info_key, info.clone()).await; assert_key_value(&kv, info_key, serde_json::to_vec(&node_info)?).await; assert_key_value(&kv, warehouse_key, serde_json::to_vec(&warehouse_node)?).await; diff --git a/src/query/service/src/clusters/cluster.rs b/src/query/service/src/clusters/cluster.rs index c1719fe338267..3b267cf0636a0 100644 --- a/src/query/service/src/clusters/cluster.rs +++ b/src/query/service/src/clusters/cluster.rs @@ -511,7 +511,7 @@ impl ClusterHeartbeat { } Either::Right((_, new_shutdown_notified)) => { shutdown_notified = new_shutdown_notified; - let heartbeat = cluster_api.heartbeat(&mut node, match_seq); + let heartbeat = cluster_api.heartbeat_node(&mut node, match_seq); match heartbeat.await { Ok(new_match_seq) => { match_seq = new_match_seq; From c3af3db165430ea8ac3f772fe93e54a3b844e508 Mon Sep 17 00:00:00 2001 From: zhang2014 Date: Fri, 20 Dec 2024 13:11:01 +0800 Subject: [PATCH 13/53] feat(cluster): support custom management cluster --- src/query/management/src/warehouse/warehouse_mgr.rs | 2 +- src/query/management/tests/it/warehouse.rs | 4 +++- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/src/query/management/src/warehouse/warehouse_mgr.rs b/src/query/management/src/warehouse/warehouse_mgr.rs index fcc008aeb9f76..f22b938e7d68b 100644 --- a/src/query/management/src/warehouse/warehouse_mgr.rs +++ b/src/query/management/src/warehouse/warehouse_mgr.rs @@ -1495,7 +1495,7 @@ impl WarehouseApi for WarehouseMgr { } return Err(ErrorCode::ClusterUnknownNode(format!( - "Warehouse cluster {:?}.{:?} unknwon node {:?}", + "Warehouse cluster {:?}.{:?} unknown node {:?}", warehouse, cluster, remove_node ))); } diff --git a/src/query/management/tests/it/warehouse.rs b/src/query/management/tests/it/warehouse.rs index 1776ade64fd34..54d82aa31fbf4 100644 --- a/src/query/management/tests/it/warehouse.rs +++ b/src/query/management/tests/it/warehouse.rs @@ -230,7 +230,9 @@ async fn test_successfully_heartbeat_self_managed_node() -> Result<()> { assert_key_value(&kv, warehouse_info_key, info.clone()).await; assert_key_expire(&kv, warehouse_info_key, Duration::from_mins(50)).await; - warehouse_manager.heartbeat_node(&mut node_info, seq).await?; + warehouse_manager + .heartbeat_node(&mut node_info, seq) + .await?; assert_key_value(&kv, warehouse_info_key, info.clone()).await; assert_key_value(&kv, info_key, serde_json::to_vec(&node_info)?).await; assert_key_value(&kv, warehouse_key, serde_json::to_vec(&warehouse_node)?).await; From 7df5abbf7c3dc960e71b863534617ac32dde17bf Mon Sep 17 00:00:00 2001 From: zhang2014 Date: Fri, 20 Dec 2024 20:17:07 +0800 Subject: [PATCH 14/53] feat(cluster): support custom management cluster --- .../management/src/warehouse/warehouse_api.rs | 4 + .../management/src/warehouse/warehouse_mgr.rs | 195 ++++++++++++++++++ 2 files changed, 199 insertions(+) diff --git a/src/query/management/src/warehouse/warehouse_api.rs b/src/query/management/src/warehouse/warehouse_api.rs index 3e9b2a1809e59..2dd2a7603bbdd 100644 --- a/src/query/management/src/warehouse/warehouse_api.rs +++ b/src/query/management/src/warehouse/warehouse_api.rs @@ -59,6 +59,10 @@ pub trait WarehouseApi: Sync + Send { async fn create_warehouse(&self, warehouse: String, nodes: SelectedNodes) -> Result<()>; + async fn resume_warehouse(&self, warehouse: String) -> Result<()>; + + async fn suspend_warehouse(&self, warehouse: String) -> Result<()>; + async fn list_warehouses(&self) -> Result>; async fn rename_warehouse(&self, cur: String, to: String) -> Result<()>; diff --git a/src/query/management/src/warehouse/warehouse_mgr.rs b/src/query/management/src/warehouse/warehouse_mgr.rs index f22b938e7d68b..61e12a517ab7f 100644 --- a/src/query/management/src/warehouse/warehouse_mgr.rs +++ b/src/query/management/src/warehouse/warehouse_mgr.rs @@ -779,6 +779,201 @@ impl WarehouseApi for WarehouseMgr { } } + async fn resume_warehouse(&self, warehouse: String) -> Result<()> { + if warehouse.is_empty() { + return Err(ErrorCode::InvalidWarehouse("Warehouse name is empty.")); + } + + for _idx in 0..10 { + let mut warehouse_info = self.consistent_warehouse_info(&warehouse).await?; + + let mut need_schedule_cluster = HashMap::new(); + warehouse_info.warehouse_info = match warehouse_info.warehouse_info { + WarehouseInfo::SelfManaged(_) => Err(ErrorCode::InvalidWarehouse( + "Cannot resume self-managed warehouse.", + )), + WarehouseInfo::SystemManaged(warehouse) => { + if warehouse.status != "SUSPEND" { + return Err(ErrorCode::InvalidWarehouse(format!( + "Cannot resume warehouse {:?}, because warehouse state is not suspend", + warehouse + ))); + } + + // TODO: support cluster resume? + need_schedule_cluster = warehouse.clusters.clone(); + Ok(WarehouseInfo::SystemManaged(SystemManagedWarehouse { + id: warehouse.id.clone(), + status: "RUNNING".to_string(), + display_name: warehouse.display_name, + clusters: warehouse.clusters, + })) + } + }?; + + // get online nodes + let online_nodes = self.metastore.prefix_list_kv(&self.node_key_prefix).await?; + + let mut unassign_online_nodes = Vec::with_capacity(online_nodes.len()); + + for (_key, v) in online_nodes { + let node_info = serde_json::from_slice::(&v.data)?; + + if node_info.warehouse_id.is_empty() && node_info.cluster_id.is_empty() { + assert_eq!(node_info.node_type, NodeType::SystemManaged); + unassign_online_nodes.push((v.seq, node_info)); + } + } + + let mut resume_txn = TxnRequest::default(); + for (cluster, info) in need_schedule_cluster { + for selected_node in info.nodes { + match selected_node { + SelectedNode::Random(Some(_)) => { + return Err(ErrorCode::Unimplemented( + "Custom instance types are not supported.", + )); + } + SelectedNode::Random(None) => match unassign_online_nodes.pop() { + None => { + return Err(ErrorCode::NoResourcesAvailable( + "Failed to create warehouse, reason: no resources available", + )); + } + Some((seq, mut node)) => { + node.cluster_id = cluster.clone(); + node.warehouse_id = warehouse.clone(); + + let node_key = self.node_key(&node)?; + let cluster_key = self.cluster_key(&node)?; + + resume_txn + .condition + .push(map_condition(&node_key, MatchSeq::Exact(seq))); + + resume_txn.if_then.push(TxnOp::put_with_ttl( + node_key, + serde_json::to_vec(&node)?, + Some(self.lift_time * 4), + )); + + node.cluster_id = String::new(); + node.warehouse_id = String::new(); + resume_txn + .condition + .push(map_condition(&cluster_key, MatchSeq::Exact(0))); + resume_txn.if_then.push(TxnOp::put_with_ttl( + cluster_key, + serde_json::to_vec(&node)?, + Some(self.lift_time * 4), + )); + } + }, + } + } + } + + let warehouse_key = format!( + "{}/{}", + self.warehouse_key_prefix, + escape_for_key(&warehouse)? + ); + + resume_txn.condition.push(map_condition( + &warehouse_key, + MatchSeq::Exact(warehouse_info.info_seq), + )); + resume_txn.if_then.push(TxnOp::put( + warehouse_key.clone(), + serde_json::to_vec(&warehouse_info.warehouse_info)?, + )); + + if self.metastore.transaction(resume_txn).await?.success { + return Ok(()); + } + } + + Err(ErrorCode::WarehouseOperateConflict( + "Warehouse operate conflict(tried 10 times while in resume warehouse).", + )) + } + + async fn suspend_warehouse(&self, warehouse: String) -> Result<()> { + if warehouse.is_empty() { + return Err(ErrorCode::InvalidWarehouse("Warehouse name is empty.")); + } + + for _idx in 0..10 { + let mut consistent_info = self.consistent_warehouse_info(&warehouse).await?; + + consistent_info.warehouse_info = match consistent_info.warehouse_info { + WarehouseInfo::SelfManaged(_) => Err(ErrorCode::InvalidWarehouse( + "Cannot suspend self-managed warehouse", + )), + WarehouseInfo::SystemManaged(warehouse) => { + if warehouse.status != "RUNNING" { + return Err(ErrorCode::InvalidWarehouse(format!("Cannot suspend warehouse {:?}, because warehouse state is not running.", warehouse))); + } + + Ok(WarehouseInfo::SystemManaged(SystemManagedWarehouse { + id: warehouse.id.clone(), + status: "RUNNING".to_string(), + display_name: warehouse.display_name, + clusters: warehouse.clusters, + })) + } + }?; + + let mut suspend_txn = TxnRequest::default(); + + let warehouse_key = format!( + "{}/{}", + self.warehouse_key_prefix, + escape_for_key(&warehouse)? + ); + + suspend_txn.condition.push(map_condition( + &warehouse_key, + MatchSeq::Exact(consistent_info.info_seq), + )); + suspend_txn.if_then.push(TxnOp::put( + warehouse_key, + serde_json::to_vec(&consistent_info.warehouse_info)?, + )); + + for mut consistent_node in consistent_info.consistent_nodes { + let node_key = self.node_key(&consistent_node.node_info)?; + let cluster_key = self.cluster_key(&consistent_node.node_info)?; + + suspend_txn.condition.push(map_condition( + &node_key, + MatchSeq::Exact(consistent_node.node_seq), + )); + suspend_txn.condition.push(map_condition( + &cluster_key, + MatchSeq::Exact(consistent_node.cluster_seq), + )); + + suspend_txn.if_then.push(TxnOp::delete(cluster_key)); + consistent_node.node_info.cluster_id = String::new(); + consistent_node.node_info.warehouse_id = String::new(); + suspend_txn.if_then.push(TxnOp::put_with_ttl( + node_key, + serde_json::to_vec(&consistent_node.node_info)?, + Some(self.lift_time * 4), + )); + } + + if self.metastore.transaction(suspend_txn).await?.success { + return Ok(()); + } + } + + Err(ErrorCode::WarehouseOperateConflict( + "Warehouse operate conflict(tried 10 times while in suspend warehouse).", + )) + } + async fn list_warehouses(&self) -> Result> { let values = self .metastore From 82cbba0150c2cd257afb3b5a84e63228d9b96a90 Mon Sep 17 00:00:00 2001 From: zhang2014 Date: Fri, 20 Dec 2024 20:22:06 +0800 Subject: [PATCH 15/53] feat(cluster): support custom management cluster --- src/query/management/src/warehouse/warehouse_api.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/query/management/src/warehouse/warehouse_api.rs b/src/query/management/src/warehouse/warehouse_api.rs index 2dd2a7603bbdd..7c72e83a020db 100644 --- a/src/query/management/src/warehouse/warehouse_api.rs +++ b/src/query/management/src/warehouse/warehouse_api.rs @@ -30,7 +30,7 @@ pub enum WarehouseInfo { SystemManaged(SystemManagedWarehouse), } -#[derive(serde::Serialize, serde::Deserialize, Eq, PartialEq, Debug)] +#[derive(serde::Serialize, serde::Deserialize, Eq, PartialEq, Debug, Clone)] pub struct SystemManagedCluster { pub nodes: Vec, } From 13656848ea6fda9458538ddf8f70dc59b884dd0c Mon Sep 17 00:00:00 2001 From: zhang2014 Date: Fri, 20 Dec 2024 22:37:36 +0800 Subject: [PATCH 16/53] feat(cluster): support custom management cluster --- src/query/ast/src/ast/statements/mod.rs | 2 + src/query/ast/src/ast/statements/statement.rs | 45 ++++++++++++++-- src/query/ast/src/ast/statements/warehouse.rs | 52 +++++++++++++++++++ src/query/ast/src/parser/statement.rs | 20 +++++++ src/query/ast/src/parser/token.rs | 2 + src/query/management/tests/it/warehouse.rs | 49 +++++++++++++++++ src/query/sql/src/planner/binder/binder.rs | 12 +++++ 7 files changed, 179 insertions(+), 3 deletions(-) create mode 100644 src/query/ast/src/ast/statements/warehouse.rs diff --git a/src/query/ast/src/ast/statements/mod.rs b/src/query/ast/src/ast/statements/mod.rs index b6fc9c0cf3ef3..da80d2136eafb 100644 --- a/src/query/ast/src/ast/statements/mod.rs +++ b/src/query/ast/src/ast/statements/mod.rs @@ -55,6 +55,7 @@ mod update; mod user; mod view; mod virtual_column; +mod warehouse; pub use call::*; pub use catalog::*; @@ -99,3 +100,4 @@ pub use update::*; pub use user::*; pub use view::*; pub use virtual_column::*; +pub use warehouse::*; diff --git a/src/query/ast/src/ast/statements/statement.rs b/src/query/ast/src/ast/statements/statement.rs index a9a7ca2885ac7..5fab0ea28b7a8 100644 --- a/src/query/ast/src/ast/statements/statement.rs +++ b/src/query/ast/src/ast/statements/statement.rs @@ -29,6 +29,7 @@ use crate::ast::statements::connection::CreateConnectionStmt; use crate::ast::statements::pipe::CreatePipeStmt; use crate::ast::statements::settings::Settings; use crate::ast::statements::task::CreateTaskStmt; +use crate::ast::statements::warehouse::ShowWarehousesStmt; use crate::ast::write_comma_separated_list; use crate::ast::CreateOption; use crate::ast::Identifier; @@ -128,6 +129,20 @@ pub enum Statement { catalog: Identifier, }, + // Warehouses + ShowWarehouses(ShowWarehousesStmt), + DropWarehouse(DropWarehouseStmt), + CreateWarehouse(CreateWarehouseStmt), + RenameWarehouse(RenameWarehouseStmt), + ResumeWarehouse(ResumeWarehouseStmt), + SuspendWarehouse(SuspendWarehouseStmt), + InspectWarehouse(InspectWarehouseStmt), + AddWarehouseCluster(AddWarehouseClusterStmt), + DropWarehouseCluster(DropWarehouseClusterStmt), + RenameWarehouseCluster(RenameWarehouseClusterStmt), + AddWarehouseClusterNode(AddWarehouseClusterNodeStmt), + DropWarehouseClusterNode(DropWarehouseClusterNodeStmt), + // Databases ShowDatabases(ShowDatabasesStmt), ShowDropDatabases(ShowDropDatabasesStmt), @@ -469,7 +484,9 @@ impl Statement { | Statement::ExecuteImmediate(..) | Statement::ShowProcedures { .. } | Statement::DescProcedure(..) - | Statement::CallProcedure(..) => true, + | Statement::CallProcedure(..) + | Statement::ShowWarehouses(..) + | Statement::InspectWarehouse(..) => true, Statement::CreateDatabase(..) | Statement::CreateTable(..) @@ -536,8 +553,17 @@ impl Statement { | Statement::AlterNotification(..) | Statement::DropNotification(..) | Statement::CreateProcedure(..) - | Statement::DropProcedure(..) => false, - + | Statement::DropProcedure(..) + | Statement::CreateWarehouse(..) + | Statement::DropWarehouse(..) + | Statement::RenameWarehouse(..) + | Statement::AddWarehouseCluster(..) + | Statement::DropWarehouseCluster(..) + | Statement::RenameWarehouseCluster(..) + | Statement::AddWarehouseClusterNode(..) + | Statement::DropWarehouseClusterNode(..) + | Statement::ResumeWarehouse(..) + | Statement::SuspendWarehouse(..) => false, Statement::StatementWithSettings { stmt, settings: _ } => { stmt.allowed_in_multi_statement() } @@ -948,6 +974,19 @@ impl Display for Statement { } Statement::System(stmt) => write!(f, "{stmt}")?, Statement::CallProcedure(stmt) => write!(f, "{stmt}")?, + + Statement::ShowWarehouses(_) => unimplemented!(), + Statement::DropWarehouse(_) => unimplemented!(), + Statement::CreateWarehouse(_) => unimplemented!(), + Statement::RenameWarehouse(_) => unimplemented!(), + Statement::ResumeWarehouse(_) => unimplemented!(), + Statement::SuspendWarehouse(_) => unimplemented!(), + Statement::InspectWarehouse(_) => unimplemented!(), + Statement::AddWarehouseCluster(_) => unimplemented!(), + Statement::DropWarehouseCluster(_) => unimplemented!(), + Statement::RenameWarehouseCluster(_) => unimplemented!(), + Statement::AddWarehouseClusterNode(_) => unimplemented!(), + Statement::DropWarehouseClusterNode(_) => unimplemented!(), } Ok(()) } diff --git a/src/query/ast/src/ast/statements/warehouse.rs b/src/query/ast/src/ast/statements/warehouse.rs new file mode 100644 index 0000000000000..700ade2a42e2f --- /dev/null +++ b/src/query/ast/src/ast/statements/warehouse.rs @@ -0,0 +1,52 @@ +// Copyright 2021 Datafuse Labs +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use derive_visitor::Drive; +use derive_visitor::DriveMut; + +#[derive(Debug, Clone, PartialEq, Drive, DriveMut)] +pub struct ShowWarehousesStmt {} + +#[derive(Debug, Clone, PartialEq, Drive, DriveMut)] +pub struct CreateWarehouseStmt {} + +#[derive(Debug, Clone, PartialEq, Drive, DriveMut)] +pub struct DropWarehouseStmt {} + +#[derive(Debug, Clone, PartialEq, Drive, DriveMut)] +pub struct SuspendWarehouseStmt {} + +#[derive(Debug, Clone, PartialEq, Drive, DriveMut)] +pub struct ResumeWarehouseStmt {} + +#[derive(Debug, Clone, PartialEq, Drive, DriveMut)] +pub struct RenameWarehouseStmt {} + +#[derive(Debug, Clone, PartialEq, Drive, DriveMut)] +pub struct InspectWarehouseStmt {} + +#[derive(Debug, Clone, PartialEq, Drive, DriveMut)] +pub struct AddWarehouseClusterStmt {} + +#[derive(Debug, Clone, PartialEq, Drive, DriveMut)] +pub struct DropWarehouseClusterStmt {} + +#[derive(Debug, Clone, PartialEq, Drive, DriveMut)] +pub struct RenameWarehouseClusterStmt {} + +#[derive(Debug, Clone, PartialEq, Drive, DriveMut)] +pub struct AddWarehouseClusterNodeStmt {} + +#[derive(Debug, Clone, PartialEq, Drive, DriveMut)] +pub struct DropWarehouseClusterNodeStmt {} diff --git a/src/query/ast/src/parser/statement.rs b/src/query/ast/src/parser/statement.rs index 8f5aa6d31829c..0224358196a09 100644 --- a/src/query/ast/src/parser/statement.rs +++ b/src/query/ast/src/parser/statement.rs @@ -523,6 +523,22 @@ pub fn statement_body(i: Input) -> IResult { |(_, _, catalog)| Statement::UseCatalog { catalog }, ); + let show_warehouses = map( + rule! { + SHOW ~ WAREHOUSES + }, + |(_, _)| Statement::ShowWarehouses(ShowWarehousesStmt {}), + ); + + // let create_warehouse = map( + // rule! { + // CREATE ~ WAREHOUSE ~ #ident + // }, + // |(_, _, name)| { + // Statement::ShowWarehouses(ShowWarehousesStmt {}) + // }, + // ); + let show_databases = map( rule! { SHOW ~ FULL? ~ ( DATABASES | SCHEMAS ) ~ ( ( FROM | IN ) ~ ^#ident )? ~ #show_limit? @@ -2292,6 +2308,10 @@ pub fn statement_body(i: Input) -> IResult { #use_catalog: "`USE CATALOG `" | #use_database : "`USE `" ), + // warehouse + rule!( + #show_warehouses: "`SHOW WAREHOUSES`" + ), // database rule!( #show_databases : "`SHOW [FULL] DATABASES [(FROM | IN) ] []`" diff --git a/src/query/ast/src/parser/token.rs b/src/query/ast/src/parser/token.rs index 82974cd017ce4..f36759b633bc6 100644 --- a/src/query/ast/src/parser/token.rs +++ b/src/query/ast/src/parser/token.rs @@ -1376,6 +1376,8 @@ pub enum TokenKind { SQL, #[token("SUNDAY", ignore(ascii_case))] SUNDAY, + #[token("WAREHOUSES", ignore(ascii_case))] + WAREHOUSES, } // Reference: https://www.postgresql.org/docs/current/sql-keywords-appendix.html diff --git a/src/query/management/tests/it/warehouse.rs b/src/query/management/tests/it/warehouse.rs index 54d82aa31fbf4..f9ef582bd3444 100644 --- a/src/query/management/tests/it/warehouse.rs +++ b/src/query/management/tests/it/warehouse.rs @@ -825,6 +825,55 @@ async fn test_rename_warehouses() -> Result<()> { Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn test_drop_warehouse_cluster_failure() -> Result<()> { + let (_, warehouse_manager, _nodes) = nodes(Duration::from_mins(30), 1).await?; + let drop_warehouse_cluster = warehouse_manager + .drop_warehouse_cluster(String::from("test_warehouse"), String::from("test_cluster")); + + assert_eq!(drop_warehouse_cluster.await.unwrap_err().code(), 2406); + + let create_warehouse = warehouse_manager + .create_warehouse(String::from("test_warehouse"), vec![SelectedNode::Random( + None, + )]); + + create_warehouse.await?; + + let drop_warehouse_cluster = + warehouse_manager.drop_warehouse_cluster(String::from(""), String::from("test_cluster")); + + assert_eq!(drop_warehouse_cluster.await.unwrap_err().code(), 2403); + + let drop_warehouse_cluster = + warehouse_manager.drop_warehouse_cluster(String::from("test_warehouse"), String::from("")); + + assert_eq!(drop_warehouse_cluster.await.unwrap_err().code(), 2403); + + let drop_warehouse_cluster = warehouse_manager + .drop_warehouse_cluster(String::from("test_warehouse"), String::from("test_cluster")); + + assert_eq!(drop_warehouse_cluster.await.unwrap_err().code(), 2410); + + let drop_warehouse_cluster = warehouse_manager + .drop_warehouse_cluster(String::from("test_warehouse"), String::from("default")); + + assert_eq!(drop_warehouse_cluster.await.unwrap_err().code(), 2408); + + warehouse_manager + .start_node(self_managed_node(&GlobalUniqName::unique())) + .await?; + + let drop_warehouse_cluster = warehouse_manager.drop_warehouse_cluster( + String::from("test-cluster-id"), + String::from("test-cluster-id"), + ); + + assert_eq!(drop_warehouse_cluster.await.unwrap_err().code(), 2403); + + Ok(()) +} + fn system_managed_node(id: &str) -> NodeInfo { NodeInfo { id: id.to_string(), diff --git a/src/query/sql/src/planner/binder/binder.rs b/src/query/sql/src/planner/binder/binder.rs index 9c1d566647f0f..8af8b2c4d0d77 100644 --- a/src/query/sql/src/planner/binder/binder.rs +++ b/src/query/sql/src/planner/binder/binder.rs @@ -645,6 +645,18 @@ impl<'a> Binder { return Err(ErrorCode::SyntaxException("CALL PROCEDURE, set enable_experimental_procedure=1")); } } + Statement::ShowWarehouses(_) => unimplemented!(), + Statement::DropWarehouse(_) => unimplemented!(), + Statement::CreateWarehouse(_) => unimplemented!(), + Statement::RenameWarehouse(_) => unimplemented!(), + Statement::ResumeWarehouse(_) => unimplemented!(), + Statement::SuspendWarehouse(_) => unimplemented!(), + Statement::InspectWarehouse(_) => unimplemented!(), + Statement::AddWarehouseCluster(_) => unimplemented!(), + Statement::DropWarehouseCluster(_) => unimplemented!(), + Statement::RenameWarehouseCluster(_) => unimplemented!(), + Statement::AddWarehouseClusterNode(_) => unimplemented!(), + Statement::DropWarehouseClusterNode(_) => unimplemented!(), }; match plan.kind() { From e092400d4487c5143ef9ce7212b27713f34ec2af Mon Sep 17 00:00:00 2001 From: zhang2014 Date: Sun, 22 Dec 2024 23:14:28 +0800 Subject: [PATCH 17/53] feat(cluster): support custom management cluster --- Cargo.lock | 14 ++ Cargo.toml | 2 + src/binaries/query/ee_main.rs | 2 +- src/binaries/query/entry.rs | 6 +- src/binaries/query/oss_main.rs | 2 +- src/binaries/tool/table_meta_inspector.rs | 2 +- src/query/ast/src/ast/statements/statement.rs | 24 ++-- src/query/ast/src/ast/statements/warehouse.rs | 130 +++++++++++++++-- src/query/ast/src/parser/statement.rs | 83 +++++++++-- src/query/ast/src/parser/token.rs | 2 + src/query/ee/Cargo.toml | 2 + src/query/ee/src/enterprise_services.rs | 2 + src/query/ee/src/lib.rs | 1 + src/query/ee/src/resource_management/mod.rs | 18 +++ .../resources_management_kubernetes.rs | 105 ++++++++++++++ .../resources_management_system.rs | 131 +++++++++++++++++ src/query/ee/src/test_kits/setup.rs | 2 +- .../resources_management/Cargo.toml | 23 +++ .../resources_management/src/lib.rs | 18 +++ .../src/resources_management.rs | 136 ++++++++++++++++++ src/query/service/Cargo.toml | 1 + src/query/service/src/global_services.rs | 11 +- .../interpreters/access/privilege_access.rs | 12 ++ .../interpreter_add_warehouse_cluster.rs | 54 +++++++ .../interpreter_add_warehouse_cluster_node.rs | 49 +++++++ .../interpreter_create_warehouses.rs | 54 +++++++ .../interpreter_drop_warehouse_cluster.rs | 55 +++++++ ...interpreter_drop_warehouse_cluster_node.rs | 49 +++++++ .../interpreter_drop_warehouses.rs | 55 +++++++ .../src/interpreters/interpreter_factory.rs | 49 +++++++ .../interpreter_inspect_warehouse.rs | 83 +++++++++++ .../interpreter_rename_warehouse.rs | 55 +++++++ .../interpreter_rename_warehouse_cluster.rs | 56 ++++++++ .../interpreter_resume_warehouse.rs | 52 +++++++ .../interpreter_show_warehouses.rs | 43 ++++++ .../interpreter_suspend_warehouse.rs | 52 +++++++ src/query/service/src/interpreters/mod.rs | 12 ++ src/query/service/src/local/mod.rs | 2 +- src/query/service/src/test_kits/fixture.rs | 2 +- src/query/sql/src/planner/binder/binder.rs | 97 +++++++------ src/query/sql/src/planner/binder/ddl/mod.rs | 1 + .../sql/src/planner/binder/ddl/warehouse.rs | 113 +++++++++++++++ .../sql/src/planner/format/display_plan.rs | 12 ++ src/query/sql/src/planner/plans/ddl/mod.rs | 2 + .../sql/src/planner/plans/ddl/warehouse.rs | 85 +++++++++++ src/query/sql/src/planner/plans/plan.rs | 26 ++++ 46 files changed, 1703 insertions(+), 84 deletions(-) create mode 100644 src/query/ee/src/resource_management/mod.rs create mode 100644 src/query/ee/src/resource_management/resources_management_kubernetes.rs create mode 100644 src/query/ee/src/resource_management/resources_management_system.rs create mode 100644 src/query/ee_features/resources_management/Cargo.toml create mode 100644 src/query/ee_features/resources_management/src/lib.rs create mode 100644 src/query/ee_features/resources_management/src/resources_management.rs create mode 100644 src/query/service/src/interpreters/interpreter_add_warehouse_cluster.rs create mode 100644 src/query/service/src/interpreters/interpreter_add_warehouse_cluster_node.rs create mode 100644 src/query/service/src/interpreters/interpreter_create_warehouses.rs create mode 100644 src/query/service/src/interpreters/interpreter_drop_warehouse_cluster.rs create mode 100644 src/query/service/src/interpreters/interpreter_drop_warehouse_cluster_node.rs create mode 100644 src/query/service/src/interpreters/interpreter_drop_warehouses.rs create mode 100644 src/query/service/src/interpreters/interpreter_inspect_warehouse.rs create mode 100644 src/query/service/src/interpreters/interpreter_rename_warehouse.rs create mode 100644 src/query/service/src/interpreters/interpreter_rename_warehouse_cluster.rs create mode 100644 src/query/service/src/interpreters/interpreter_resume_warehouse.rs create mode 100644 src/query/service/src/interpreters/interpreter_show_warehouses.rs create mode 100644 src/query/service/src/interpreters/interpreter_suspend_warehouse.rs create mode 100644 src/query/sql/src/planner/binder/ddl/warehouse.rs create mode 100644 src/query/sql/src/planner/plans/ddl/warehouse.rs diff --git a/Cargo.lock b/Cargo.lock index 5982a33dc2fce..87c5447c199ab 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4740,6 +4740,7 @@ dependencies = [ "databend-common-functions", "databend-common-io", "databend-common-license", + "databend-common-management", "databend-common-meta-api", "databend-common-meta-app", "databend-common-meta-store", @@ -4761,6 +4762,7 @@ dependencies = [ "databend-enterprise-data-mask-feature", "databend-enterprise-fail-safe", "databend-enterprise-inverted-index", + "databend-enterprise-resources-management", "databend-enterprise-storage-encryption", "databend-enterprise-storage-quota", "databend-enterprise-stream-handler", @@ -4782,6 +4784,17 @@ dependencies = [ "tempfile", ] +[[package]] +name = "databend-enterprise-resources-management" +version = "0.1.0" +dependencies = [ + "async-trait", + "databend-common-base", + "databend-common-exception", + "databend-common-management", + "databend-common-meta-types", +] + [[package]] name = "databend-enterprise-storage-encryption" version = "0.1.0" @@ -5019,6 +5032,7 @@ dependencies = [ "databend-enterprise-background-service", "databend-enterprise-data-mask-feature", "databend-enterprise-inverted-index", + "databend-enterprise-resources-management", "databend-enterprise-stream-handler", "databend-enterprise-vacuum-handler", "databend-enterprise-virtual-column", diff --git a/Cargo.toml b/Cargo.toml index 4eb9610c5c8e9..782e7a8ad2b4f 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -79,6 +79,7 @@ members = [ "src/query/ee_features/storage_quota", "src/query/ee_features/inverted_index", "src/query/ee_features/virtual_column", + "src/query/ee_features/resources_management", "src/query/service", "src/query/ee", "src/meta/api", @@ -179,6 +180,7 @@ databend-enterprise-fail-safe = { path = "src/query/ee_features/fail_safe" } databend-enterprise-inverted-index = { path = "src/query/ee_features/inverted_index" } databend-enterprise-meta = { path = "src/meta/ee" } databend-enterprise-query = { path = "src/query/ee" } +databend-enterprise-resources-management = { path = "src/query/ee_features/resources_management" } databend-enterprise-storage-encryption = { path = "src/query/ee_features/storage_encryption" } databend-enterprise-storage-quota = { path = "src/query/ee_features/storage_quota" } databend-enterprise-stream-handler = { path = "src/query/ee_features/stream_handler" } diff --git a/src/binaries/query/ee_main.rs b/src/binaries/query/ee_main.rs index b1d2782d61ea5..d85a8a217fb7f 100644 --- a/src/binaries/query/ee_main.rs +++ b/src/binaries/query/ee_main.rs @@ -69,7 +69,7 @@ pub async fn main_entrypoint() -> Result<(), MainError> { return Ok(()); } - init_services(&conf).await.with_context(make_error)?; + init_services(&conf, true).await.with_context(make_error)?; EnterpriseServices::init(conf.clone()) .await .with_context(make_error)?; diff --git a/src/binaries/query/entry.rs b/src/binaries/query/entry.rs index b48856e6dbf3a..8970c1bca48a1 100644 --- a/src/binaries/query/entry.rs +++ b/src/binaries/query/entry.rs @@ -69,7 +69,7 @@ pub async fn run_cmd(conf: &InnerConfig) -> Result { Ok(true) } -pub async fn init_services(conf: &InnerConfig) -> Result<(), MainError> { +pub async fn init_services(conf: &InnerConfig, ee_mode: bool) -> Result<(), MainError> { let make_error = || "failed to init services"; set_panic_hook(); @@ -92,7 +92,9 @@ pub async fn init_services(conf: &InnerConfig) -> Result<(), MainError> { .with_context(make_error); } // Make sure global services have been inited. - GlobalServices::init(conf).await.with_context(make_error) + GlobalServices::init(conf, ee_mode) + .await + .with_context(make_error) } async fn precheck_services(conf: &InnerConfig) -> Result<(), MainError> { diff --git a/src/binaries/query/oss_main.rs b/src/binaries/query/oss_main.rs index 92bc179e61a64..54b19666e29fb 100644 --- a/src/binaries/query/oss_main.rs +++ b/src/binaries/query/oss_main.rs @@ -70,7 +70,7 @@ async fn main_entrypoint() -> Result<(), MainError> { return Ok(()); } - init_services(&conf).await?; + init_services(&conf, false).await?; // init oss license manager OssLicenseManager::init(conf.query.tenant_id.tenant_name().to_string()) .with_context(make_error)?; diff --git a/src/binaries/tool/table_meta_inspector.rs b/src/binaries/tool/table_meta_inspector.rs index 4bf93b379953c..512ba900844d8 100644 --- a/src/binaries/tool/table_meta_inspector.rs +++ b/src/binaries/tool/table_meta_inspector.rs @@ -88,7 +88,7 @@ async fn parse_input_data(config: &InspectorConfig) -> Result> { builder = builder.collect(from_file(Toml, config_file)); let read_config = builder.build()?; let inner_config: InnerConfig = read_config.clone().try_into()?; - GlobalServices::init(&inner_config).await?; + GlobalServices::init(&inner_config, false).await?; let storage_config: StorageConfig = read_config.storage.try_into()?; init_operator(&storage_config.params)? } diff --git a/src/query/ast/src/ast/statements/statement.rs b/src/query/ast/src/ast/statements/statement.rs index 5fab0ea28b7a8..a5d5b89a88e24 100644 --- a/src/query/ast/src/ast/statements/statement.rs +++ b/src/query/ast/src/ast/statements/statement.rs @@ -975,18 +975,18 @@ impl Display for Statement { Statement::System(stmt) => write!(f, "{stmt}")?, Statement::CallProcedure(stmt) => write!(f, "{stmt}")?, - Statement::ShowWarehouses(_) => unimplemented!(), - Statement::DropWarehouse(_) => unimplemented!(), - Statement::CreateWarehouse(_) => unimplemented!(), - Statement::RenameWarehouse(_) => unimplemented!(), - Statement::ResumeWarehouse(_) => unimplemented!(), - Statement::SuspendWarehouse(_) => unimplemented!(), - Statement::InspectWarehouse(_) => unimplemented!(), - Statement::AddWarehouseCluster(_) => unimplemented!(), - Statement::DropWarehouseCluster(_) => unimplemented!(), - Statement::RenameWarehouseCluster(_) => unimplemented!(), - Statement::AddWarehouseClusterNode(_) => unimplemented!(), - Statement::DropWarehouseClusterNode(_) => unimplemented!(), + Statement::ShowWarehouses(stmt) => write!(f, "{stmt}")?, + Statement::DropWarehouse(stmt) => write!(f, "{stmt}")?, + Statement::CreateWarehouse(stmt) => write!(f, "{stmt}")?, + Statement::RenameWarehouse(stmt) => write!(f, "{stmt}")?, + Statement::ResumeWarehouse(stmt) => write!(f, "{stmt}")?, + Statement::SuspendWarehouse(stmt) => write!(f, "{stmt}")?, + Statement::InspectWarehouse(stmt) => write!(f, "{stmt}")?, + Statement::AddWarehouseCluster(stmt) => write!(f, "{stmt}")?, + Statement::DropWarehouseCluster(stmt) => write!(f, "{stmt}")?, + Statement::RenameWarehouseCluster(stmt) => write!(f, "{stmt}")?, + Statement::AddWarehouseClusterNode(stmt) => write!(f, "{stmt}")?, + Statement::DropWarehouseClusterNode(stmt) => write!(f, "{stmt}")?, } Ok(()) } diff --git a/src/query/ast/src/ast/statements/warehouse.rs b/src/query/ast/src/ast/statements/warehouse.rs index 700ade2a42e2f..968a06f6cd339 100644 --- a/src/query/ast/src/ast/statements/warehouse.rs +++ b/src/query/ast/src/ast/statements/warehouse.rs @@ -12,41 +12,153 @@ // See the License for the specific language governing permissions and // limitations under the License. +use std::fmt::Display; +use std::fmt::Formatter; + use derive_visitor::Drive; use derive_visitor::DriveMut; +use crate::ast::Identifier; + #[derive(Debug, Clone, PartialEq, Drive, DriveMut)] pub struct ShowWarehousesStmt {} +impl Display for ShowWarehousesStmt { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + write!(f, "SHOW WAREHOUSES") + } +} + #[derive(Debug, Clone, PartialEq, Drive, DriveMut)] -pub struct CreateWarehouseStmt {} +pub struct CreateWarehouseStmt { + pub warehouse: Identifier, +} + +impl Display for CreateWarehouseStmt { + fn fmt(&self, _f: &mut Formatter<'_>) -> std::fmt::Result { + todo!() + } +} #[derive(Debug, Clone, PartialEq, Drive, DriveMut)] -pub struct DropWarehouseStmt {} +pub struct DropWarehouseStmt { + pub warehouse: Identifier, +} + +impl Display for DropWarehouseStmt { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + write!(f, "DROP WAREHOUSE {}", self.warehouse) + } +} #[derive(Debug, Clone, PartialEq, Drive, DriveMut)] -pub struct SuspendWarehouseStmt {} +pub struct SuspendWarehouseStmt { + pub warehouse: Identifier, +} + +impl Display for SuspendWarehouseStmt { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + write!(f, "SUSPEND WAREHOUSE {}", self.warehouse) + } +} #[derive(Debug, Clone, PartialEq, Drive, DriveMut)] -pub struct ResumeWarehouseStmt {} +pub struct ResumeWarehouseStmt { + pub warehouse: Identifier, +} + +impl Display for ResumeWarehouseStmt { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + write!(f, "RESUME WAREHOUSE {}", self.warehouse) + } +} #[derive(Debug, Clone, PartialEq, Drive, DriveMut)] -pub struct RenameWarehouseStmt {} +pub struct RenameWarehouseStmt { + pub warehouse: Identifier, + pub new_warehouse: Identifier, +} + +impl Display for RenameWarehouseStmt { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + write!( + f, + "RENAME WAREHOUSE {} TO {}", + self.warehouse, self.new_warehouse + ) + } +} #[derive(Debug, Clone, PartialEq, Drive, DriveMut)] -pub struct InspectWarehouseStmt {} +pub struct InspectWarehouseStmt { + pub warehouse: Identifier, +} + +impl Display for InspectWarehouseStmt { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + write!(f, "INSPECT WAREHOUSE {}", self.warehouse) + } +} #[derive(Debug, Clone, PartialEq, Drive, DriveMut)] pub struct AddWarehouseClusterStmt {} +impl Display for AddWarehouseClusterStmt { + fn fmt(&self, _f: &mut Formatter<'_>) -> std::fmt::Result { + todo!() + } +} + #[derive(Debug, Clone, PartialEq, Drive, DriveMut)] -pub struct DropWarehouseClusterStmt {} +pub struct DropWarehouseClusterStmt { + pub cluster: Identifier, + pub warehouse: Identifier, +} + +impl Display for DropWarehouseClusterStmt { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + write!( + f, + "ALTER WAREHOUSE {} DROP CLUSTER {}", + self.warehouse, self.cluster + ) + } +} #[derive(Debug, Clone, PartialEq, Drive, DriveMut)] -pub struct RenameWarehouseClusterStmt {} +pub struct RenameWarehouseClusterStmt { + pub warehouse: Identifier, + pub cluster: Identifier, + pub new_cluster: Identifier, +} + +impl Display for RenameWarehouseClusterStmt { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + write!( + f, + "ALTER WAREHOUSE {} RENAME CLUSTER {} TO {}", + self.warehouse, self.cluster, self.new_cluster + ) + } +} #[derive(Debug, Clone, PartialEq, Drive, DriveMut)] pub struct AddWarehouseClusterNodeStmt {} +impl Display for AddWarehouseClusterNodeStmt { + fn fmt(&self, _f: &mut Formatter<'_>) -> std::fmt::Result { + todo!() + } +} + #[derive(Debug, Clone, PartialEq, Drive, DriveMut)] -pub struct DropWarehouseClusterNodeStmt {} +pub struct DropWarehouseClusterNodeStmt { + // warehouse:String, + // cluster: String, +} + +impl Display for DropWarehouseClusterNodeStmt { + fn fmt(&self, _f: &mut Formatter<'_>) -> std::fmt::Result { + todo!() + } +} diff --git a/src/query/ast/src/parser/statement.rs b/src/query/ast/src/parser/statement.rs index 0224358196a09..45244fe018ffd 100644 --- a/src/query/ast/src/parser/statement.rs +++ b/src/query/ast/src/parser/statement.rs @@ -530,14 +530,74 @@ pub fn statement_body(i: Input) -> IResult { |(_, _)| Statement::ShowWarehouses(ShowWarehousesStmt {}), ); - // let create_warehouse = map( - // rule! { - // CREATE ~ WAREHOUSE ~ #ident - // }, - // |(_, _, name)| { - // Statement::ShowWarehouses(ShowWarehousesStmt {}) - // }, - // ); + let _create_warehouse = map( + rule! { + CREATE ~ WAREHOUSE ~ #ident + }, + |(_, _, warehouse)| Statement::CreateWarehouse(CreateWarehouseStmt { warehouse }), + ); + + let drop_warehouse = map( + rule! { + DROP ~ WAREHOUSE ~ #ident + }, + |(_, _, warehouse)| Statement::DropWarehouse(DropWarehouseStmt { warehouse }), + ); + + let rename_warehouse = map( + rule! { + RENAME ~ WAREHOUSE ~ #ident ~ TO ~ #ident + }, + |(_, _, warehouse, _, new_warehouse)| { + Statement::RenameWarehouse(RenameWarehouseStmt { + warehouse, + new_warehouse, + }) + }, + ); + + let resume_warehouse = map( + rule! { + RESUME ~ WAREHOUSE ~ #ident + }, + |(_, _, warehouse)| Statement::ResumeWarehouse(ResumeWarehouseStmt { warehouse }), + ); + + let suspend_warehouse = map( + rule! { + SUSPEND ~ WAREHOUSE ~ #ident + }, + |(_, _, warehouse)| Statement::SuspendWarehouse(SuspendWarehouseStmt { warehouse }), + ); + + let inspect_warehouse = map( + rule! { + INSPECT ~ WAREHOUSE ~ #ident + }, + |(_, _, warehouse)| Statement::InspectWarehouse(InspectWarehouseStmt { warehouse }), + ); + + let drop_warehouse_cluster = map( + rule! { + ALTER ~ WAREHOUSE ~ #ident ~ DROP ~ CLUSTER ~ #ident + }, + |(_, _, warehouse, _, _, cluster)| { + Statement::DropWarehouseCluster(DropWarehouseClusterStmt { warehouse, cluster }) + }, + ); + + let rename_warehouse_cluster = map( + rule! { + ALTER ~ WAREHOUSE ~ #ident ~ RENAME ~ CLUSTER ~ #ident ~ TO ~ #ident + }, + |(_, _, warehouse, _, _, cluster, _, new_cluster)| { + Statement::RenameWarehouseCluster(RenameWarehouseClusterStmt { + warehouse, + cluster, + new_cluster, + }) + }, + ); let show_databases = map( rule! { @@ -2311,6 +2371,13 @@ pub fn statement_body(i: Input) -> IResult { // warehouse rule!( #show_warehouses: "`SHOW WAREHOUSES`" + | #drop_warehouse: "`DROP WAREHOUSE `" + | #rename_warehouse: "`RENAME WAREHOUSE TO `" + | #resume_warehouse: "`RESUME WAREHOUSE `" + | #suspend_warehouse: "`SUSPEND WAREHOUSE `" + | #inspect_warehouse: "`INSPECT WAREHOUSE `" + | #drop_warehouse_cluster: "`ALTER WAREHOUSE DROP CLUSTER `" + | #rename_warehouse_cluster: "`ALTER WAREHOUSE RENAME CLUSTER TO `" ), // database rule!( diff --git a/src/query/ast/src/parser/token.rs b/src/query/ast/src/parser/token.rs index f36759b633bc6..a4a5011a4c47e 100644 --- a/src/query/ast/src/parser/token.rs +++ b/src/query/ast/src/parser/token.rs @@ -1378,6 +1378,8 @@ pub enum TokenKind { SUNDAY, #[token("WAREHOUSES", ignore(ascii_case))] WAREHOUSES, + #[token("INSPECT", ignore(ascii_case))] + INSPECT, } // Reference: https://www.postgresql.org/docs/current/sql-keywords-appendix.html diff --git a/src/query/ee/Cargo.toml b/src/query/ee/Cargo.toml index def5aeb0bc62f..67a5fc9b55e28 100644 --- a/src/query/ee/Cargo.toml +++ b/src/query/ee/Cargo.toml @@ -27,6 +27,7 @@ databend-common-expression = { workspace = true } databend-common-functions = { workspace = true } databend-common-io = { workspace = true } databend-common-license = { workspace = true } +databend-common-management = { workspace = true } databend-common-meta-api = { workspace = true } databend-common-meta-app = { workspace = true } databend-common-meta-store = { workspace = true } @@ -48,6 +49,7 @@ databend-enterprise-background-service = { workspace = true } databend-enterprise-data-mask-feature = { workspace = true } databend-enterprise-fail-safe = { workspace = true } databend-enterprise-inverted-index = { workspace = true } +databend-enterprise-resources-management = { workspace = true } databend-enterprise-storage-encryption = { workspace = true } databend-enterprise-storage-quota = { workspace = true } databend-enterprise-stream-handler = { workspace = true } diff --git a/src/query/ee/src/enterprise_services.rs b/src/query/ee/src/enterprise_services.rs index 4d43c636a9b11..7433e65f253be 100644 --- a/src/query/ee/src/enterprise_services.rs +++ b/src/query/ee/src/enterprise_services.rs @@ -23,6 +23,7 @@ use crate::data_mask::RealDatamaskHandler; use crate::fail_safe::RealFailSafeHandler; use crate::inverted_index::RealInvertedIndexHandler; use crate::license::license_mgr::RealLicenseManager; +use crate::resource_management::SystemResourcesManagement; use crate::storage_encryption::RealStorageEncryptionHandler; use crate::storage_quota::RealStorageQuotaHandler; use crate::storages::fuse::operations::RealVacuumHandler; @@ -45,6 +46,7 @@ impl EnterpriseServices { RealInvertedIndexHandler::init()?; RealStorageQuotaHandler::init(&cfg)?; RealFailSafeHandler::init()?; + SystemResourcesManagement::init(&cfg).await?; Ok(()) } } diff --git a/src/query/ee/src/lib.rs b/src/query/ee/src/lib.rs index 7bf07a60abe3c..9ea8a9f26fbac 100644 --- a/src/query/ee/src/lib.rs +++ b/src/query/ee/src/lib.rs @@ -20,6 +20,7 @@ pub mod enterprise_services; pub mod fail_safe; pub mod inverted_index; pub mod license; +pub mod resource_management; pub mod storage_encryption; pub mod storage_quota; pub mod storages; diff --git a/src/query/ee/src/resource_management/mod.rs b/src/query/ee/src/resource_management/mod.rs new file mode 100644 index 0000000000000..c7fa58e84a8e9 --- /dev/null +++ b/src/query/ee/src/resource_management/mod.rs @@ -0,0 +1,18 @@ +// Copyright 2023 Databend Cloud +// +// Licensed under the Elastic License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// https://www.elastic.co/licensing/elastic-license +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +mod resources_management_kubernetes; +mod resources_management_system; + +pub use resources_management_system::SystemResourcesManagement; diff --git a/src/query/ee/src/resource_management/resources_management_kubernetes.rs b/src/query/ee/src/resource_management/resources_management_kubernetes.rs new file mode 100644 index 0000000000000..7d82318d55819 --- /dev/null +++ b/src/query/ee/src/resource_management/resources_management_kubernetes.rs @@ -0,0 +1,105 @@ +// Copyright 2023 Databend Cloud +// +// Licensed under the Elastic License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// https://www.elastic.co/licensing/elastic-license +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use databend_common_exception::ErrorCode; +use databend_common_exception::Result; +use databend_common_management::SelectedNode; +use databend_common_meta_types::NodeInfo; +use databend_enterprise_resources_management::ResourcesManagement; + +pub struct KubernetesResourcesManagement {} + +#[async_trait::async_trait] +impl ResourcesManagement for KubernetesResourcesManagement { + async fn create_warehouse(&self, _: String, _: Vec) -> Result<()> { + Err(ErrorCode::Unimplemented( + "Unimplemented kubernetes resources management", + )) + } + + async fn drop_warehouse(&self, _: String) -> Result<()> { + Err(ErrorCode::Unimplemented( + "Unimplemented kubernetes resources management", + )) + } + + async fn resume_warehouse(&self, _: String) -> Result<()> { + Err(ErrorCode::Unimplemented( + "Unimplemented kubernetes resources management", + )) + } + + async fn suspend_warehouse(&self, _: String) -> Result<()> { + Err(ErrorCode::Unimplemented( + "Unimplemented kubernetes resources management", + )) + } + + async fn rename_warehouse(&self, _: String, _: String) -> Result<()> { + Err(ErrorCode::Unimplemented( + "Unimplemented kubernetes resources management", + )) + } + + async fn inspect_warehouse(&self, _: String) -> Result> { + Err(ErrorCode::Unimplemented( + "Unimplemented kubernetes resources management", + )) + } + + async fn add_warehouse_cluster( + &self, + _: String, + _: String, + _: Vec, + ) -> Result<()> { + Err(ErrorCode::Unimplemented( + "Unimplemented kubernetes resources management", + )) + } + + async fn rename_warehouse_cluster(&self, _: String, _: String, _: String) -> Result<()> { + Err(ErrorCode::Unimplemented( + "Unimplemented kubernetes resources management", + )) + } + + async fn drop_warehouse_cluster(&self, _: String, _: String) -> Result<()> { + Err(ErrorCode::Unimplemented( + "Unimplemented kubernetes resources management", + )) + } + + async fn add_warehouse_cluster_node( + &self, + _: String, + _: String, + _: Vec, + ) -> Result<()> { + Err(ErrorCode::Unimplemented( + "Unimplemented kubernetes resources management", + )) + } + + async fn drop_warehouse_cluster_node( + &self, + _: String, + _: String, + _: Vec, + ) -> Result<()> { + Err(ErrorCode::Unimplemented( + "Unimplemented kubernetes resources management", + )) + } +} diff --git a/src/query/ee/src/resource_management/resources_management_system.rs b/src/query/ee/src/resource_management/resources_management_system.rs new file mode 100644 index 0000000000000..472e5fb5200e5 --- /dev/null +++ b/src/query/ee/src/resource_management/resources_management_system.rs @@ -0,0 +1,131 @@ +// Copyright 2023 Databend Cloud +// +// Licensed under the Elastic License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// https://www.elastic.co/licensing/elastic-license +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::sync::Arc; +use std::time::Duration; + +use databend_common_base::base::GlobalInstance; +use databend_common_config::InnerConfig; +use databend_common_exception::ErrorCode; +use databend_common_exception::Result; +use databend_common_management::SelectedNode; +use databend_common_management::WarehouseApi; +use databend_common_management::WarehouseMgr; +use databend_common_meta_store::MetaStoreProvider; +use databend_common_meta_types::NodeInfo; +use databend_enterprise_resources_management::ResourcesManagement; + +pub struct SystemResourcesManagement { + warehouse_manager: Arc, +} + +#[async_trait::async_trait] +impl ResourcesManagement for SystemResourcesManagement { + async fn create_warehouse(&self, name: String, nodes: Vec) -> Result<()> { + self.warehouse_manager.create_warehouse(name, nodes).await + } + + async fn drop_warehouse(&self, name: String) -> Result<()> { + self.warehouse_manager.drop_warehouse(name).await + } + + async fn resume_warehouse(&self, name: String) -> Result<()> { + self.warehouse_manager.resume_warehouse(name).await + } + + async fn suspend_warehouse(&self, name: String) -> Result<()> { + self.warehouse_manager.suspend_warehouse(name).await + } + + async fn rename_warehouse(&self, name: String, to: String) -> Result<()> { + self.warehouse_manager.rename_warehouse(name, to).await + } + + async fn inspect_warehouse(&self, name: String) -> Result> { + self.warehouse_manager.list_warehouse_nodes(name).await + } + + async fn add_warehouse_cluster( + &self, + name: String, + cluster: String, + nodes: Vec, + ) -> Result<()> { + self.warehouse_manager + .add_warehouse_cluster(name, cluster, nodes) + .await + } + + async fn rename_warehouse_cluster( + &self, + name: String, + cluster: String, + new_cluster: String, + ) -> Result<()> { + self.warehouse_manager + .rename_warehouse_cluster(name, cluster, new_cluster) + .await + } + + async fn drop_warehouse_cluster(&self, name: String, cluster: String) -> Result<()> { + self.warehouse_manager + .drop_warehouse_cluster(name, cluster) + .await + } + + async fn add_warehouse_cluster_node( + &self, + name: String, + cluster: String, + nodes: Vec, + ) -> Result<()> { + self.warehouse_manager + .add_warehouse_cluster_node(&name, &cluster, nodes) + .await + } + + async fn drop_warehouse_cluster_node( + &self, + name: String, + cluster: String, + nodes: Vec, + ) -> Result<()> { + self.warehouse_manager + .drop_warehouse_cluster_node(&name, &cluster, nodes) + .await + } +} + +impl SystemResourcesManagement { + pub async fn init(cfg: &InnerConfig) -> Result<()> { + let meta_api_provider = MetaStoreProvider::new(cfg.meta.to_meta_grpc_client_conf()); + match meta_api_provider.create_meta_store().await { + Err(cause) => { + Err(ErrorCode::from(cause).add_message_back("(while create resources management).")) + } + Ok(metastore) => { + let tenant_id = &cfg.query.tenant_id; + let lift_time = Duration::from_secs(60); + let warehouse_manager = + WarehouseMgr::create(metastore, tenant_id.tenant_name(), lift_time)?; + + let service: Arc = Arc::new(SystemResourcesManagement { + warehouse_manager: Arc::new(warehouse_manager), + }); + GlobalInstance::set(service); + Ok(()) + } + } + } +} diff --git a/src/query/ee/src/test_kits/setup.rs b/src/query/ee/src/test_kits/setup.rs index ba2b469d9aa0d..57b24f2e70b7e 100644 --- a/src/query/ee/src/test_kits/setup.rs +++ b/src/query/ee/src/test_kits/setup.rs @@ -34,7 +34,7 @@ impl TestFixture { databend_common_base::base::GlobalInstance::init_testing(&thread_name); } - GlobalServices::init_with(config).await?; + GlobalServices::init_with(config, false).await?; MockServices::init(config, public_key).await?; // Cluster register. diff --git a/src/query/ee_features/resources_management/Cargo.toml b/src/query/ee_features/resources_management/Cargo.toml new file mode 100644 index 0000000000000..493bdfb3a8a71 --- /dev/null +++ b/src/query/ee_features/resources_management/Cargo.toml @@ -0,0 +1,23 @@ +[package] +name = "databend-enterprise-resources-management" +version = { workspace = true } +authors = { workspace = true } +license = { workspace = true } +publish = { workspace = true } +edition = { workspace = true } + +[lib] +doctest = false +test = true + +[dependencies] +async-trait = { workspace = true } +databend-common-base = { workspace = true } +databend-common-exception = { workspace = true } +databend-common-management = { workspace = true } +databend-common-meta-types = { workspace = true } + +[build-dependencies] + +[lints] +workspace = true diff --git a/src/query/ee_features/resources_management/src/lib.rs b/src/query/ee_features/resources_management/src/lib.rs new file mode 100644 index 0000000000000..1624990934206 --- /dev/null +++ b/src/query/ee_features/resources_management/src/lib.rs @@ -0,0 +1,18 @@ +// Copyright 2021 Datafuse Labs +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +mod resources_management; + +pub use resources_management::DummyResourcesManagement; +pub use resources_management::ResourcesManagement; diff --git a/src/query/ee_features/resources_management/src/resources_management.rs b/src/query/ee_features/resources_management/src/resources_management.rs new file mode 100644 index 0000000000000..8d3e3b59ce36a --- /dev/null +++ b/src/query/ee_features/resources_management/src/resources_management.rs @@ -0,0 +1,136 @@ +// Copyright 2021 Datafuse Labs +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use databend_common_base::base::GlobalInstance; +use databend_common_exception::ErrorCode; +use databend_common_exception::Result; +use databend_common_management::SelectedNode; +use databend_common_meta_types::NodeInfo; + +#[async_trait::async_trait] +pub trait ResourcesManagement: Sync + Send + 'static { + async fn create_warehouse(&self, name: String, nodes: Vec) -> Result<()>; + + async fn drop_warehouse(&self, name: String) -> Result<()>; + + async fn resume_warehouse(&self, name: String) -> Result<()>; + + async fn suspend_warehouse(&self, name: String) -> Result<()>; + + async fn rename_warehouse(&self, name: String, to: String) -> Result<()>; + + async fn inspect_warehouse(&self, name: String) -> Result>; + + async fn add_warehouse_cluster( + &self, + name: String, + cluster: String, + nodes: Vec, + ) -> Result<()>; + + async fn rename_warehouse_cluster( + &self, + name: String, + cluster: String, + new_cluster: String, + ) -> Result<()>; + + async fn drop_warehouse_cluster(&self, name: String, cluster: String) -> Result<()>; + + async fn add_warehouse_cluster_node( + &self, + name: String, + cluster: String, + nodes: Vec, + ) -> Result<()>; + + async fn drop_warehouse_cluster_node( + &self, + name: String, + cluster: String, + nodes: Vec, + ) -> Result<()>; +} + +pub struct DummyResourcesManagement; + +#[async_trait::async_trait] +impl ResourcesManagement for DummyResourcesManagement { + async fn create_warehouse(&self, _: String, _: Vec) -> Result<()> { + Err(ErrorCode::Unimplemented("The use of this feature requires a Databend Enterprise Edition license. To unlock enterprise features, please contact Databend to obtain a license. Learn more at https://docs.databend.com/guides/overview/editions/dee/")) + } + + async fn drop_warehouse(&self, _: String) -> Result<()> { + Err(ErrorCode::Unimplemented("The use of this feature requires a Databend Enterprise Edition license. To unlock enterprise features, please contact Databend to obtain a license. Learn more at https://docs.databend.com/guides/overview/editions/dee/")) + } + + async fn resume_warehouse(&self, _: String) -> Result<()> { + Err(ErrorCode::Unimplemented("The use of this feature requires a Databend Enterprise Edition license. To unlock enterprise features, please contact Databend to obtain a license. Learn more at https://docs.databend.com/guides/overview/editions/dee/")) + } + + async fn suspend_warehouse(&self, _: String) -> Result<()> { + Err(ErrorCode::Unimplemented("The use of this feature requires a Databend Enterprise Edition license. To unlock enterprise features, please contact Databend to obtain a license. Learn more at https://docs.databend.com/guides/overview/editions/dee/")) + } + + async fn rename_warehouse(&self, _: String, _: String) -> Result<()> { + Err(ErrorCode::Unimplemented("The use of this feature requires a Databend Enterprise Edition license. To unlock enterprise features, please contact Databend to obtain a license. Learn more at https://docs.databend.com/guides/overview/editions/dee/")) + } + + async fn inspect_warehouse(&self, _: String) -> Result> { + Err(ErrorCode::Unimplemented("The use of this feature requires a Databend Enterprise Edition license. To unlock enterprise features, please contact Databend to obtain a license. Learn more at https://docs.databend.com/guides/overview/editions/dee/")) + } + + async fn add_warehouse_cluster( + &self, + _: String, + _: String, + _: Vec, + ) -> Result<()> { + Err(ErrorCode::Unimplemented("The use of this feature requires a Databend Enterprise Edition license. To unlock enterprise features, please contact Databend to obtain a license. Learn more at https://docs.databend.com/guides/overview/editions/dee/")) + } + + async fn rename_warehouse_cluster(&self, _: String, _: String, _: String) -> Result<()> { + Err(ErrorCode::Unimplemented("The use of this feature requires a Databend Enterprise Edition license. To unlock enterprise features, please contact Databend to obtain a license. Learn more at https://docs.databend.com/guides/overview/editions/dee/")) + } + + async fn drop_warehouse_cluster(&self, _: String, _: String) -> Result<()> { + Err(ErrorCode::Unimplemented("The use of this feature requires a Databend Enterprise Edition license. To unlock enterprise features, please contact Databend to obtain a license. Learn more at https://docs.databend.com/guides/overview/editions/dee/")) + } + + async fn add_warehouse_cluster_node( + &self, + _: String, + _: String, + _: Vec, + ) -> Result<()> { + Err(ErrorCode::Unimplemented("The use of this feature requires a Databend Enterprise Edition license. To unlock enterprise features, please contact Databend to obtain a license. Learn more at https://docs.databend.com/guides/overview/editions/dee/")) + } + + async fn drop_warehouse_cluster_node( + &self, + _: String, + _: String, + _: Vec, + ) -> Result<()> { + Err(ErrorCode::Unimplemented("The use of this feature requires a Databend Enterprise Edition license. To unlock enterprise features, please contact Databend to obtain a license. Learn more at https://docs.databend.com/guides/overview/editions/dee/")) + } +} + +impl DummyResourcesManagement { + pub fn init() -> Result<()> { + let instance: Box = Box::new(DummyResourcesManagement); + GlobalInstance::set(instance); + Ok(()) + } +} diff --git a/src/query/service/Cargo.toml b/src/query/service/Cargo.toml index 1ef9bdba27f68..95cb990c724f7 100644 --- a/src/query/service/Cargo.toml +++ b/src/query/service/Cargo.toml @@ -105,6 +105,7 @@ databend-enterprise-attach-table = { workspace = true } databend-enterprise-background-service = { workspace = true } databend-enterprise-data-mask-feature = { workspace = true } databend-enterprise-inverted-index = { workspace = true } +databend-enterprise-resources-management = { workspace = true } databend-enterprise-stream-handler = { workspace = true } databend-enterprise-vacuum-handler = { workspace = true } databend-enterprise-virtual-column = { workspace = true } diff --git a/src/query/service/src/global_services.rs b/src/query/service/src/global_services.rs index 263a7b5b86f3d..3f988ea4c1d06 100644 --- a/src/query/service/src/global_services.rs +++ b/src/query/service/src/global_services.rs @@ -35,6 +35,7 @@ use databend_common_tracing::GlobalLogger; use databend_common_users::builtin::BuiltIn; use databend_common_users::RoleCacheManager; use databend_common_users::UserApiProvider; +use databend_enterprise_resources_management::DummyResourcesManagement; use databend_storages_common_cache::CacheManager; use databend_storages_common_cache::TempDirManager; @@ -56,13 +57,13 @@ pub struct GlobalServices; impl GlobalServices { #[async_backtrace::framed] - pub async fn init(config: &InnerConfig) -> Result<()> { + pub async fn init(config: &InnerConfig, ee_mode: bool) -> Result<()> { GlobalInstance::init_production(); - GlobalServices::init_with(config).await + GlobalServices::init_with(config, ee_mode).await } #[async_backtrace::framed] - pub async fn init_with(config: &InnerConfig) -> Result<()> { + pub async fn init_with(config: &InnerConfig, ee_mode: bool) -> Result<()> { StackTrace::pre_load_symbol(); // app name format: node_id[0..7]@cluster_id @@ -163,6 +164,10 @@ impl GlobalServices { GlobalQueriesExecutor::init()?; } + if !ee_mode { + DummyResourcesManagement::init()?; + } + Ok(()) } } diff --git a/src/query/service/src/interpreters/access/privilege_access.rs b/src/query/service/src/interpreters/access/privilege_access.rs index d098c71609ef4..ac0da6c4b3a39 100644 --- a/src/query/service/src/interpreters/access/privilege_access.rs +++ b/src/query/service/src/interpreters/access/privilege_access.rs @@ -1255,6 +1255,18 @@ impl AccessChecker for PrivilegeAccess { } Plan::Commit => {} Plan::Abort => {} + Plan::ShowWarehouses => {} + Plan::DropWarehouse(_) => {} + Plan::ResumeWarehouse(_) => {} + Plan::SuspendWarehouse(_) => {} + Plan::RenameWarehouse(_) => {} + Plan::InspectWarehouse(_) => {} + Plan::DropWarehouseCluster(_) => {} + Plan::RenameWarehouseCluster(_) => {} + Plan::CreateWarehouse(_) => {} + Plan::AddWarehouseCluster(_) => {} + Plan::AddWarehouseClusterNode(_) => {} + Plan::DropWarehouseClusterNode(_) => {} } Ok(()) diff --git a/src/query/service/src/interpreters/interpreter_add_warehouse_cluster.rs b/src/query/service/src/interpreters/interpreter_add_warehouse_cluster.rs new file mode 100644 index 0000000000000..54958147491a5 --- /dev/null +++ b/src/query/service/src/interpreters/interpreter_add_warehouse_cluster.rs @@ -0,0 +1,54 @@ +// Copyright 2021 Datafuse Labs +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::sync::Arc; + +use databend_common_exception::Result; +use databend_common_sql::plans::AddWarehouseClusterPlan; + +use crate::interpreters::Interpreter; +use crate::pipelines::PipelineBuildResult; +use crate::sessions::QueryContext; + +pub struct AddWarehouseClusterInterpreter { + #[allow(dead_code)] + ctx: Arc, + #[allow(dead_code)] + plan: AddWarehouseClusterPlan, +} + +impl AddWarehouseClusterInterpreter { + pub fn try_create(ctx: Arc, plan: AddWarehouseClusterPlan) -> Result { + Ok(AddWarehouseClusterInterpreter { ctx, plan }) + } +} + +#[async_trait::async_trait] +impl Interpreter for AddWarehouseClusterInterpreter { + fn name(&self) -> &str { + "AddWarehouseClusterInterpreter" + } + + fn is_ddl(&self) -> bool { + true + } + + #[async_backtrace::framed] + async fn execute2(&self) -> Result { + // GlobalInstance::get::>() + // .drop_warehouse_cluster(self.plan.warehouse.clone(), self.plan.cluster.clone()) + // .await?; + Ok(PipelineBuildResult::create()) + } +} diff --git a/src/query/service/src/interpreters/interpreter_add_warehouse_cluster_node.rs b/src/query/service/src/interpreters/interpreter_add_warehouse_cluster_node.rs new file mode 100644 index 0000000000000..1566d7188ae04 --- /dev/null +++ b/src/query/service/src/interpreters/interpreter_add_warehouse_cluster_node.rs @@ -0,0 +1,49 @@ +// Copyright 2021 Datafuse Labs +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use databend_common_exception::Result; +use databend_common_sql::plans::AddWarehouseClusterNodePlan; + +use crate::interpreters::Interpreter; +use crate::pipelines::PipelineBuildResult; + +pub struct AddWarehouseClusterNodeInterpreter { + #[allow(dead_code)] + plan: AddWarehouseClusterNodePlan, +} + +impl AddWarehouseClusterNodeInterpreter { + pub fn try_create(plan: AddWarehouseClusterNodePlan) -> Result { + Ok(AddWarehouseClusterNodeInterpreter { plan }) + } +} + +#[async_trait::async_trait] +impl Interpreter for AddWarehouseClusterNodeInterpreter { + fn name(&self) -> &str { + "AddWarehouseClusterNodeInterpreter" + } + + fn is_ddl(&self) -> bool { + true + } + + #[async_backtrace::framed] + async fn execute2(&self) -> Result { + // GlobalInstance::get::>() + // .drop_warehouse_cluster(self.plan.warehouse.clone(), self.plan.cluster.clone()) + // .await?; + Ok(PipelineBuildResult::create()) + } +} diff --git a/src/query/service/src/interpreters/interpreter_create_warehouses.rs b/src/query/service/src/interpreters/interpreter_create_warehouses.rs new file mode 100644 index 0000000000000..609d32c83c844 --- /dev/null +++ b/src/query/service/src/interpreters/interpreter_create_warehouses.rs @@ -0,0 +1,54 @@ +// Copyright 2021 Datafuse Labs +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::sync::Arc; + +use databend_common_exception::Result; +use databend_common_sql::plans::CreateWarehousePlan; + +use crate::interpreters::Interpreter; +use crate::pipelines::PipelineBuildResult; +use crate::sessions::QueryContext; + +pub struct CreateWarehouseInterpreter { + #[allow(dead_code)] + ctx: Arc, + #[allow(dead_code)] + plan: CreateWarehousePlan, +} + +impl CreateWarehouseInterpreter { + pub fn try_create(ctx: Arc, plan: CreateWarehousePlan) -> Result { + Ok(CreateWarehouseInterpreter { ctx, plan }) + } +} + +#[async_trait::async_trait] +impl Interpreter for CreateWarehouseInterpreter { + fn name(&self) -> &str { + "CreateWarehouseInterpreter" + } + + fn is_ddl(&self) -> bool { + true + } + + #[async_backtrace::framed] + async fn execute2(&self) -> Result { + // GlobalInstance::get::>() + // .create_warehouse(self.plan.warehouse.clone()) + // .await?; + Ok(PipelineBuildResult::create()) + } +} diff --git a/src/query/service/src/interpreters/interpreter_drop_warehouse_cluster.rs b/src/query/service/src/interpreters/interpreter_drop_warehouse_cluster.rs new file mode 100644 index 0000000000000..107a544416661 --- /dev/null +++ b/src/query/service/src/interpreters/interpreter_drop_warehouse_cluster.rs @@ -0,0 +1,55 @@ +// Copyright 2021 Datafuse Labs +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::sync::Arc; + +use databend_common_base::base::GlobalInstance; +use databend_common_exception::Result; +use databend_common_sql::plans::DropWarehouseClusterPlan; +use databend_enterprise_resources_management::ResourcesManagement; + +use crate::interpreters::Interpreter; +use crate::pipelines::PipelineBuildResult; +use crate::sessions::QueryContext; + +pub struct DropWarehouseClusterInterpreter { + #[allow(dead_code)] + ctx: Arc, + plan: DropWarehouseClusterPlan, +} + +impl DropWarehouseClusterInterpreter { + pub fn try_create(ctx: Arc, plan: DropWarehouseClusterPlan) -> Result { + Ok(DropWarehouseClusterInterpreter { ctx, plan }) + } +} + +#[async_trait::async_trait] +impl Interpreter for DropWarehouseClusterInterpreter { + fn name(&self) -> &str { + "DropWarehouseClusterInterpreter" + } + + fn is_ddl(&self) -> bool { + true + } + + #[async_backtrace::framed] + async fn execute2(&self) -> Result { + GlobalInstance::get::>() + .drop_warehouse_cluster(self.plan.warehouse.clone(), self.plan.cluster.clone()) + .await?; + Ok(PipelineBuildResult::create()) + } +} diff --git a/src/query/service/src/interpreters/interpreter_drop_warehouse_cluster_node.rs b/src/query/service/src/interpreters/interpreter_drop_warehouse_cluster_node.rs new file mode 100644 index 0000000000000..e12120e8d8b5b --- /dev/null +++ b/src/query/service/src/interpreters/interpreter_drop_warehouse_cluster_node.rs @@ -0,0 +1,49 @@ +// Copyright 2021 Datafuse Labs +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use databend_common_exception::Result; +use databend_common_sql::plans::DropWarehouseClusterNodePlan; + +use crate::interpreters::Interpreter; +use crate::pipelines::PipelineBuildResult; + +pub struct DropWarehouseClusterNodeInterpreter { + #[allow(dead_code)] + plan: DropWarehouseClusterNodePlan, +} + +impl DropWarehouseClusterNodeInterpreter { + pub fn try_create(plan: DropWarehouseClusterNodePlan) -> Result { + Ok(DropWarehouseClusterNodeInterpreter { plan }) + } +} + +#[async_trait::async_trait] +impl Interpreter for DropWarehouseClusterNodeInterpreter { + fn name(&self) -> &str { + "DropWarehouseClusterNodeInterpreter" + } + + fn is_ddl(&self) -> bool { + true + } + + #[async_backtrace::framed] + async fn execute2(&self) -> Result { + // GlobalInstance::get::>() + // .drop_warehouse_cluster(self.plan.warehouse.clone(), self.plan.cluster.clone()) + // .await?; + Ok(PipelineBuildResult::create()) + } +} diff --git a/src/query/service/src/interpreters/interpreter_drop_warehouses.rs b/src/query/service/src/interpreters/interpreter_drop_warehouses.rs new file mode 100644 index 0000000000000..d077cb688429b --- /dev/null +++ b/src/query/service/src/interpreters/interpreter_drop_warehouses.rs @@ -0,0 +1,55 @@ +// Copyright 2021 Datafuse Labs +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::sync::Arc; + +use databend_common_base::base::GlobalInstance; +use databend_common_exception::Result; +use databend_common_sql::plans::DropWarehousePlan; +use databend_enterprise_resources_management::ResourcesManagement; + +use crate::interpreters::Interpreter; +use crate::pipelines::PipelineBuildResult; +use crate::sessions::QueryContext; + +pub struct DropWarehouseInterpreter { + #[allow(dead_code)] + ctx: Arc, + plan: DropWarehousePlan, +} + +impl DropWarehouseInterpreter { + pub fn try_create(ctx: Arc, plan: DropWarehousePlan) -> Result { + Ok(DropWarehouseInterpreter { ctx, plan }) + } +} + +#[async_trait::async_trait] +impl Interpreter for DropWarehouseInterpreter { + fn name(&self) -> &str { + "DropWarehouseInterpreter" + } + + fn is_ddl(&self) -> bool { + true + } + + #[async_backtrace::framed] + async fn execute2(&self) -> Result { + GlobalInstance::get::>() + .drop_warehouse(self.plan.warehouse.clone()) + .await?; + Ok(PipelineBuildResult::create()) + } +} diff --git a/src/query/service/src/interpreters/interpreter_factory.rs b/src/query/service/src/interpreters/interpreter_factory.rs index 4520200e047a4..af382e26faf6e 100644 --- a/src/query/service/src/interpreters/interpreter_factory.rs +++ b/src/query/service/src/interpreters/interpreter_factory.rs @@ -36,6 +36,8 @@ use super::interpreter_table_set_options::SetOptionsInterpreter; use super::interpreter_user_stage_drop::DropUserStageInterpreter; use super::*; use crate::interpreters::access::Accessor; +use crate::interpreters::interpreter_add_warehouse_cluster::AddWarehouseClusterInterpreter; +use crate::interpreters::interpreter_add_warehouse_cluster_node::AddWarehouseClusterNodeInterpreter; use crate::interpreters::interpreter_catalog_drop::DropCatalogInterpreter; use crate::interpreters::interpreter_connection_create::CreateConnectionInterpreter; use crate::interpreters::interpreter_connection_desc::DescConnectionInterpreter; @@ -43,9 +45,14 @@ use crate::interpreters::interpreter_connection_drop::DropConnectionInterpreter; use crate::interpreters::interpreter_connection_show::ShowConnectionsInterpreter; use crate::interpreters::interpreter_copy_into_location::CopyIntoLocationInterpreter; use crate::interpreters::interpreter_copy_into_table::CopyIntoTableInterpreter; +use crate::interpreters::interpreter_create_warehouses::CreateWarehouseInterpreter; +use crate::interpreters::interpreter_drop_warehouse_cluster::DropWarehouseClusterInterpreter; +use crate::interpreters::interpreter_drop_warehouse_cluster_node::DropWarehouseClusterNodeInterpreter; +use crate::interpreters::interpreter_drop_warehouses::DropWarehouseInterpreter; use crate::interpreters::interpreter_file_format_create::CreateFileFormatInterpreter; use crate::interpreters::interpreter_file_format_drop::DropFileFormatInterpreter; use crate::interpreters::interpreter_file_format_show::ShowFileFormatsInterpreter; +use crate::interpreters::interpreter_inspect_warehouse::InspectWarehouseInterpreter; use crate::interpreters::interpreter_notification_alter::AlterNotificationInterpreter; use crate::interpreters::interpreter_notification_create::CreateNotificationInterpreter; use crate::interpreters::interpreter_notification_desc::DescNotificationInterpreter; @@ -54,8 +61,13 @@ use crate::interpreters::interpreter_presign::PresignInterpreter; use crate::interpreters::interpreter_procedure_call::CallProcedureInterpreter; use crate::interpreters::interpreter_procedure_create::CreateProcedureInterpreter; use crate::interpreters::interpreter_procedure_drop::DropProcedureInterpreter; +use crate::interpreters::interpreter_rename_warehouse::RenameWarehouseInterpreter; +use crate::interpreters::interpreter_rename_warehouse_cluster::RenameWarehouseClusterInterpreter; +use crate::interpreters::interpreter_resume_warehouse::ResumeWarehouseInterpreter; use crate::interpreters::interpreter_role_show::ShowRolesInterpreter; use crate::interpreters::interpreter_set_priority::SetPriorityInterpreter; +use crate::interpreters::interpreter_show_warehouses::ShowWarehousesInterpreter; +use crate::interpreters::interpreter_suspend_warehouse::SuspendWarehouseInterpreter; use crate::interpreters::interpreter_system_action::SystemActionInterpreter; use crate::interpreters::interpreter_table_create::CreateTableInterpreter; use crate::interpreters::interpreter_table_revert::RevertTableInterpreter; @@ -631,6 +643,43 @@ impl InterpreterFactory { // ctx, // p.clone(), // )?)), + Plan::ShowWarehouses => Ok(Arc::new(ShowWarehousesInterpreter::try_create()?)), + Plan::CreateWarehouse(v) => Ok(Arc::new(CreateWarehouseInterpreter::try_create( + ctx, + *v.clone(), + )?)), + Plan::DropWarehouse(v) => Ok(Arc::new(DropWarehouseInterpreter::try_create( + ctx, + *v.clone(), + )?)), + Plan::ResumeWarehouse(v) => Ok(Arc::new(ResumeWarehouseInterpreter::try_create( + *v.clone(), + )?)), + Plan::SuspendWarehouse(v) => Ok(Arc::new(SuspendWarehouseInterpreter::try_create( + *v.clone(), + )?)), + Plan::RenameWarehouse(v) => Ok(Arc::new(RenameWarehouseInterpreter::try_create( + ctx, + *v.clone(), + )?)), + Plan::InspectWarehouse(v) => Ok(Arc::new(InspectWarehouseInterpreter::try_create( + *v.clone(), + )?)), + Plan::AddWarehouseCluster(v) => Ok(Arc::new( + AddWarehouseClusterInterpreter::try_create(ctx, *v.clone())?, + )), + Plan::DropWarehouseCluster(v) => Ok(Arc::new( + DropWarehouseClusterInterpreter::try_create(ctx, *v.clone())?, + )), + Plan::RenameWarehouseCluster(v) => Ok(Arc::new( + RenameWarehouseClusterInterpreter::try_create(*v.clone())?, + )), + Plan::AddWarehouseClusterNode(v) => Ok(Arc::new( + AddWarehouseClusterNodeInterpreter::try_create(*v.clone())?, + )), + Plan::DropWarehouseClusterNode(v) => Ok(Arc::new( + DropWarehouseClusterNodeInterpreter::try_create(*v.clone())?, + )), } } } diff --git a/src/query/service/src/interpreters/interpreter_inspect_warehouse.rs b/src/query/service/src/interpreters/interpreter_inspect_warehouse.rs new file mode 100644 index 0000000000000..15649f5950918 --- /dev/null +++ b/src/query/service/src/interpreters/interpreter_inspect_warehouse.rs @@ -0,0 +1,83 @@ +// Copyright 2021 Datafuse Labs +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::sync::Arc; + +use databend_common_base::base::GlobalInstance; +use databend_common_exception::Result; +use databend_common_expression::types::DataType; +use databend_common_expression::ColumnBuilder; +use databend_common_expression::DataBlock; +use databend_common_expression::Scalar; +use databend_common_meta_types::NodeType; +use databend_common_sql::plans::InspectWarehousePlan; +use databend_enterprise_resources_management::ResourcesManagement; + +use crate::interpreters::Interpreter; +use crate::pipelines::PipelineBuildResult; + +pub struct InspectWarehouseInterpreter { + plan: InspectWarehousePlan, +} + +impl InspectWarehouseInterpreter { + pub fn try_create(plan: InspectWarehousePlan) -> Result { + Ok(InspectWarehouseInterpreter { plan }) + } +} + +#[async_trait::async_trait] +impl Interpreter for InspectWarehouseInterpreter { + fn name(&self) -> &str { + "InspectWarehouseInterpreter" + } + + fn is_ddl(&self) -> bool { + false + } + + #[async_backtrace::framed] + async fn execute2(&self) -> Result { + let mut warehouse_nodes = GlobalInstance::get::>() + .inspect_warehouse(self.plan.warehouse.clone()) + .await?; + + warehouse_nodes.sort_by(|left, right| { + (&left.cluster_id, &left.id).cmp(&(&right.cluster_id, &right.id)) + }); + + let mut clusters = ColumnBuilder::with_capacity(&DataType::String, warehouse_nodes.len()); + let mut nodes_id = ColumnBuilder::with_capacity(&DataType::String, warehouse_nodes.len()); + let mut nodes_type = ColumnBuilder::with_capacity(&DataType::String, warehouse_nodes.len()); + // let mut nodes_pool = ColumnBuilder::with_capacity(&DataType::String, warehouse_nodes.len()); + + for warehouse_node in warehouse_nodes { + nodes_id.push(Scalar::String(warehouse_node.id.clone()).as_ref()); + clusters.push(Scalar::String(warehouse_node.cluster_id.clone()).as_ref()); + nodes_type.push( + Scalar::String(match warehouse_node.node_type { + NodeType::SelfManaged => String::from("SelfManaged"), + NodeType::SystemManaged => String::from("SystemManaged"), + }) + .as_ref(), + ); + } + + PipelineBuildResult::from_blocks(vec![DataBlock::new_from_columns(vec![ + clusters.build(), + nodes_id.build(), + nodes_type.build(), + ])]) + } +} diff --git a/src/query/service/src/interpreters/interpreter_rename_warehouse.rs b/src/query/service/src/interpreters/interpreter_rename_warehouse.rs new file mode 100644 index 0000000000000..1cf3536d16a57 --- /dev/null +++ b/src/query/service/src/interpreters/interpreter_rename_warehouse.rs @@ -0,0 +1,55 @@ +// Copyright 2021 Datafuse Labs +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::sync::Arc; + +use databend_common_base::base::GlobalInstance; +use databend_common_exception::Result; +use databend_common_sql::plans::RenameWarehousePlan; +use databend_enterprise_resources_management::ResourcesManagement; + +use crate::interpreters::Interpreter; +use crate::pipelines::PipelineBuildResult; +use crate::sessions::QueryContext; + +pub struct RenameWarehouseInterpreter { + #[allow(dead_code)] + ctx: Arc, + plan: RenameWarehousePlan, +} + +impl RenameWarehouseInterpreter { + pub fn try_create(ctx: Arc, plan: RenameWarehousePlan) -> Result { + Ok(RenameWarehouseInterpreter { ctx, plan }) + } +} + +#[async_trait::async_trait] +impl Interpreter for RenameWarehouseInterpreter { + fn name(&self) -> &str { + "RenameWarehouseInterpreter" + } + + fn is_ddl(&self) -> bool { + true + } + + #[async_backtrace::framed] + async fn execute2(&self) -> Result { + GlobalInstance::get::>() + .rename_warehouse(self.plan.warehouse.clone(), self.plan.new_warehouse.clone()) + .await?; + Ok(PipelineBuildResult::create()) + } +} diff --git a/src/query/service/src/interpreters/interpreter_rename_warehouse_cluster.rs b/src/query/service/src/interpreters/interpreter_rename_warehouse_cluster.rs new file mode 100644 index 0000000000000..9c6a91b840b3c --- /dev/null +++ b/src/query/service/src/interpreters/interpreter_rename_warehouse_cluster.rs @@ -0,0 +1,56 @@ +// Copyright 2021 Datafuse Labs +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::sync::Arc; + +use databend_common_base::base::GlobalInstance; +use databend_common_exception::Result; +use databend_common_sql::plans::RenameWarehouseClusterPlan; +use databend_enterprise_resources_management::ResourcesManagement; + +use crate::interpreters::Interpreter; +use crate::pipelines::PipelineBuildResult; + +pub struct RenameWarehouseClusterInterpreter { + plan: RenameWarehouseClusterPlan, +} + +impl RenameWarehouseClusterInterpreter { + pub fn try_create(plan: RenameWarehouseClusterPlan) -> Result { + Ok(RenameWarehouseClusterInterpreter { plan }) + } +} + +#[async_trait::async_trait] +impl Interpreter for RenameWarehouseClusterInterpreter { + fn name(&self) -> &str { + "RenameWarehouseClusterInterpreter" + } + + fn is_ddl(&self) -> bool { + true + } + + #[async_backtrace::framed] + async fn execute2(&self) -> Result { + GlobalInstance::get::>() + .rename_warehouse_cluster( + self.plan.warehouse.clone(), + self.plan.cluster.clone(), + self.plan.new_cluster.clone(), + ) + .await?; + Ok(PipelineBuildResult::create()) + } +} diff --git a/src/query/service/src/interpreters/interpreter_resume_warehouse.rs b/src/query/service/src/interpreters/interpreter_resume_warehouse.rs new file mode 100644 index 0000000000000..60419d6224a51 --- /dev/null +++ b/src/query/service/src/interpreters/interpreter_resume_warehouse.rs @@ -0,0 +1,52 @@ +// Copyright 2021 Datafuse Labs +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::sync::Arc; + +use databend_common_base::base::GlobalInstance; +use databend_common_exception::Result; +use databend_common_sql::plans::ResumeWarehousePlan; +use databend_enterprise_resources_management::ResourcesManagement; + +use crate::interpreters::Interpreter; +use crate::pipelines::PipelineBuildResult; + +pub struct ResumeWarehouseInterpreter { + plan: ResumeWarehousePlan, +} + +impl ResumeWarehouseInterpreter { + pub fn try_create(plan: ResumeWarehousePlan) -> Result { + Ok(ResumeWarehouseInterpreter { plan }) + } +} + +#[async_trait::async_trait] +impl Interpreter for ResumeWarehouseInterpreter { + fn name(&self) -> &str { + "ResumeWarehouseInterpreter" + } + + fn is_ddl(&self) -> bool { + true + } + + #[async_backtrace::framed] + async fn execute2(&self) -> Result { + GlobalInstance::get::>() + .resume_warehouse(self.plan.warehouse.clone()) + .await?; + Ok(PipelineBuildResult::create()) + } +} diff --git a/src/query/service/src/interpreters/interpreter_show_warehouses.rs b/src/query/service/src/interpreters/interpreter_show_warehouses.rs new file mode 100644 index 0000000000000..52f88b65914da --- /dev/null +++ b/src/query/service/src/interpreters/interpreter_show_warehouses.rs @@ -0,0 +1,43 @@ +// Copyright 2021 Datafuse Labs +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use databend_common_exception::Result; + +use crate::interpreters::Interpreter; +use crate::pipelines::PipelineBuildResult; + +pub struct ShowWarehousesInterpreter {} + +impl ShowWarehousesInterpreter { + pub fn try_create() -> Result { + Ok(ShowWarehousesInterpreter {}) + } +} + +#[async_trait::async_trait] +impl Interpreter for ShowWarehousesInterpreter { + fn name(&self) -> &str { + "ShowWarehousesInterpreter" + } + + fn is_ddl(&self) -> bool { + false + } + + #[async_backtrace::framed] + async fn execute2(&self) -> Result { + // WarehouseMgr::i + Ok(PipelineBuildResult::create()) + } +} diff --git a/src/query/service/src/interpreters/interpreter_suspend_warehouse.rs b/src/query/service/src/interpreters/interpreter_suspend_warehouse.rs new file mode 100644 index 0000000000000..eee6e2657426d --- /dev/null +++ b/src/query/service/src/interpreters/interpreter_suspend_warehouse.rs @@ -0,0 +1,52 @@ +// Copyright 2021 Datafuse Labs +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::sync::Arc; + +use databend_common_base::base::GlobalInstance; +use databend_common_exception::Result; +use databend_common_sql::plans::SuspendWarehousePlan; +use databend_enterprise_resources_management::ResourcesManagement; + +use crate::interpreters::Interpreter; +use crate::pipelines::PipelineBuildResult; + +pub struct SuspendWarehouseInterpreter { + plan: SuspendWarehousePlan, +} + +impl SuspendWarehouseInterpreter { + pub fn try_create(plan: SuspendWarehousePlan) -> Result { + Ok(SuspendWarehouseInterpreter { plan }) + } +} + +#[async_trait::async_trait] +impl Interpreter for SuspendWarehouseInterpreter { + fn name(&self) -> &str { + "SuspendWarehouseInterpreter" + } + + fn is_ddl(&self) -> bool { + true + } + + #[async_backtrace::framed] + async fn execute2(&self) -> Result { + GlobalInstance::get::>() + .suspend_warehouse(self.plan.warehouse.clone()) + .await?; + Ok(PipelineBuildResult::create()) + } +} diff --git a/src/query/service/src/interpreters/mod.rs b/src/query/service/src/interpreters/mod.rs index 8a31ac0dcafb8..ed4f2e1fdc31d 100644 --- a/src/query/service/src/interpreters/mod.rs +++ b/src/query/service/src/interpreters/mod.rs @@ -16,6 +16,8 @@ mod access; pub(crate) mod common; mod hook; mod interpreter; +mod interpreter_add_warehouse_cluster; +mod interpreter_add_warehouse_cluster_node; mod interpreter_catalog_create; mod interpreter_catalog_drop; mod interpreter_catalog_show_create; @@ -29,6 +31,7 @@ mod interpreter_connection_drop; mod interpreter_connection_show; mod interpreter_copy_into_location; mod interpreter_copy_into_table; +mod interpreter_create_warehouses; mod interpreter_data_mask_create; mod interpreter_data_mask_desc; mod interpreter_data_mask_drop; @@ -41,6 +44,9 @@ mod interpreter_dictionary_create; mod interpreter_dictionary_drop; mod interpreter_dictionary_rename; mod interpreter_dictionary_show_create; +mod interpreter_drop_warehouse_cluster; +mod interpreter_drop_warehouse_cluster_node; +mod interpreter_drop_warehouses; mod interpreter_execute_immediate; mod interpreter_explain; mod interpreter_factory; @@ -52,6 +58,7 @@ mod interpreter_index_drop; mod interpreter_index_refresh; mod interpreter_insert; mod interpreter_insert_multi_table; +mod interpreter_inspect_warehouse; mod interpreter_kill; mod interpreter_metrics; mod interpreter_mutation; @@ -77,7 +84,10 @@ mod interpreter_privilege_revoke; mod interpreter_procedure_call; mod interpreter_procedure_create; mod interpreter_procedure_drop; +mod interpreter_rename_warehouse; +mod interpreter_rename_warehouse_cluster; mod interpreter_replace; +mod interpreter_resume_warehouse; mod interpreter_role_create; mod interpreter_role_drop; mod interpreter_role_grant; @@ -90,8 +100,10 @@ mod interpreter_sequence_create; mod interpreter_sequence_drop; mod interpreter_set; mod interpreter_set_priority; +mod interpreter_show_warehouses; mod interpreter_stream_create; mod interpreter_stream_drop; +mod interpreter_suspend_warehouse; mod interpreter_system_action; mod interpreter_table_add_column; mod interpreter_table_analyze; diff --git a/src/query/service/src/local/mod.rs b/src/query/service/src/local/mod.rs index 33f5a9452f03a..0a64768e7e193 100644 --- a/src/query/service/src/local/mod.rs +++ b/src/query/service/src/local/mod.rs @@ -48,7 +48,7 @@ pub async fn query_local(query_sql: &str, output_format: &str) -> Result<()> { root: path.join("_data").to_str().unwrap().to_owned(), }); - GlobalServices::init(&conf).await?; + GlobalServices::init(&conf, false).await?; // init oss license manager OssLicenseManager::init(conf.query.tenant_id.tenant_name().to_string()).unwrap(); diff --git a/src/query/service/src/test_kits/fixture.rs b/src/query/service/src/test_kits/fixture.rs index e46472ab1754c..0417592029574 100644 --- a/src/query/service/src/test_kits/fixture.rs +++ b/src/query/service/src/test_kits/fixture.rs @@ -236,7 +236,7 @@ impl TestFixture { databend_common_base::base::GlobalInstance::init_testing(&thread_name); } - GlobalServices::init_with(config).await?; + GlobalServices::init_with(config, false).await?; OssLicenseManager::init(config.query.tenant_id.tenant_name().to_string())?; // Cluster register. diff --git a/src/query/sql/src/planner/binder/binder.rs b/src/query/sql/src/planner/binder/binder.rs index 8af8b2c4d0d77..5cd5d3f3fc8bc 100644 --- a/src/query/sql/src/planner/binder/binder.rs +++ b/src/query/sql/src/planner/binder/binder.rs @@ -204,7 +204,7 @@ impl<'a> Binder { self.bind_explain(bind_context, kind, options, query).await? } - Statement::ExplainAnalyze {partial, graphical, query } => { + Statement::ExplainAnalyze { partial, graphical, query } => { if let Statement::Explain { .. } | Statement::ExplainAnalyze { .. } = query.as_ref() { return Err(ErrorCode::SyntaxException("Invalid statement")); } @@ -254,7 +254,7 @@ impl<'a> Binder { Statement::ShowCreateCatalog(stmt) => self.bind_show_create_catalogs(stmt).await?, Statement::CreateCatalog(stmt) => self.bind_create_catalog(stmt).await?, Statement::DropCatalog(stmt) => self.bind_drop_catalog(stmt).await?, - Statement::UseCatalog {catalog} => { + Statement::UseCatalog { catalog } => { let catalog = normalize_identifier(catalog, &self.name_resolution_ctx).name; Plan::UseCatalog(Box::new(UseCatalogPlan { catalog, @@ -352,9 +352,10 @@ impl<'a> Binder { )); } Plan::CreateRole(Box::new(CreateRolePlan { - if_not_exists: *if_not_exists, - role_name: role_name.to_string(), - }))}, + if_not_exists: *if_not_exists, + role_name: role_name.to_string(), + })) + } Statement::DropRole { if_exists, role_name, @@ -386,9 +387,10 @@ impl<'a> Binder { )); } Plan::DropStage(Box::new(DropStagePlan { - if_exists: *if_exists, - name: stage_name.clone(), - }))}, + if_exists: *if_exists, + name: stage_name.clone(), + })) + } Statement::RemoveStage { location, pattern } => { self.bind_remove_stage(location, pattern).await? } @@ -453,7 +455,7 @@ impl<'a> Binder { Plan::CreateFileFormat(Box::new(CreateFileFormatPlan { create_option: create_option.clone().into(), name: name.clone(), - file_format_params: FileFormatParams::try_from_reader( FileFormatOptionsReader::from_ast(file_format_options), false)?, + file_format_params: FileFormatParams::try_from_reader(FileFormatOptionsReader::from_ast(file_format_options), false)?, })) } Statement::DropFileFormat { @@ -474,7 +476,7 @@ impl<'a> Binder { Statement::DescribeConnection(stmt) => Plan::DescConnection(Box::new(DescConnectionPlan { name: stmt.name.to_string(), })), - Statement::ShowConnections(_) => Plan::ShowConnections(Box::new(ShowConnectionsPlan{})), + Statement::ShowConnections(_) => Plan::ShowConnections(Box::new(ShowConnectionsPlan {})), // UDFs Statement::CreateUDF(stmt) => self.bind_create_udf(stmt).await?, @@ -493,7 +495,7 @@ impl<'a> Binder { .await? } - Statement::UnSetStmt{settings } => { + Statement::UnSetStmt { settings } => { let Settings { set_type, identifiers, .. } = settings; self.bind_unset(bind_context, *set_type, identifiers) .await? @@ -550,7 +552,7 @@ impl<'a> Binder { Statement::DescPasswordPolicy(stmt) => { self.bind_desc_password_policy(stmt).await? } - Statement::ShowPasswordPolicies{ show_options } => self.bind_show_password_policies(bind_context, show_options).await?, + Statement::ShowPasswordPolicies { show_options } => self.bind_show_password_policies(bind_context, show_options).await?, Statement::CreateTask(stmt) => { self.bind_create_task(stmt).await? } @@ -613,48 +615,55 @@ impl<'a> Binder { Statement::Commit => Plan::Commit, Statement::Abort => Plan::Abort, Statement::ExecuteImmediate(stmt) => self.bind_execute_immediate(stmt).await?, - Statement::SetPriority {priority, object_id} => { + Statement::SetPriority { priority, object_id } => { self.bind_set_priority(priority, object_id).await? - }, + } Statement::System(stmt) => self.bind_system(stmt).await?, - Statement::CreateProcedure(stmt) => { if self.ctx.get_settings().get_enable_experimental_procedure()? { - self.bind_create_procedure(stmt).await? - } else { - return Err(ErrorCode::SyntaxException("CREATE PROCEDURE, set enable_experimental_procedure=1")); - } - } - Statement::DropProcedure(stmt) => { if self.ctx.get_settings().get_enable_experimental_procedure()? { - self.bind_drop_procedure(stmt).await? - } else { - return Err(ErrorCode::SyntaxException("DROP PROCEDURE, set enable_experimental_procedure=1")); - } } - Statement::ShowProcedures { show_options } => { if self.ctx.get_settings().get_enable_experimental_procedure()? { - self.bind_show_procedures(bind_context, show_options).await? - } else { - return Err(ErrorCode::SyntaxException("SHOW PROCEDURES, set enable_experimental_procedure=1")); - } } - Statement::DescProcedure(stmt) => { if self.ctx.get_settings().get_enable_experimental_procedure()? { - self.bind_desc_procedure(stmt).await? - } else { - return Err(ErrorCode::SyntaxException("DESC PROCEDURE, set enable_experimental_procedure=1")); - } } + Statement::CreateProcedure(stmt) => { + if self.ctx.get_settings().get_enable_experimental_procedure()? { + self.bind_create_procedure(stmt).await? + } else { + return Err(ErrorCode::SyntaxException("CREATE PROCEDURE, set enable_experimental_procedure=1")); + } + } + Statement::DropProcedure(stmt) => { + if self.ctx.get_settings().get_enable_experimental_procedure()? { + self.bind_drop_procedure(stmt).await? + } else { + return Err(ErrorCode::SyntaxException("DROP PROCEDURE, set enable_experimental_procedure=1")); + } + } + Statement::ShowProcedures { show_options } => { + if self.ctx.get_settings().get_enable_experimental_procedure()? { + self.bind_show_procedures(bind_context, show_options).await? + } else { + return Err(ErrorCode::SyntaxException("SHOW PROCEDURES, set enable_experimental_procedure=1")); + } + } + Statement::DescProcedure(stmt) => { + if self.ctx.get_settings().get_enable_experimental_procedure()? { + self.bind_desc_procedure(stmt).await? + } else { + return Err(ErrorCode::SyntaxException("DESC PROCEDURE, set enable_experimental_procedure=1")); + } + } Statement::CallProcedure(stmt) => { if self.ctx.get_settings().get_enable_experimental_procedure()? { self.bind_call_procedure(bind_context, stmt).await? } else { return Err(ErrorCode::SyntaxException("CALL PROCEDURE, set enable_experimental_procedure=1")); } - } - Statement::ShowWarehouses(_) => unimplemented!(), - Statement::DropWarehouse(_) => unimplemented!(), + } + Statement::ShowWarehouses(v) => self.bind_show_warehouses(v)?, + Statement::DropWarehouse(v) => self.bind_drop_warehouse(v)?, Statement::CreateWarehouse(_) => unimplemented!(), - Statement::RenameWarehouse(_) => unimplemented!(), - Statement::ResumeWarehouse(_) => unimplemented!(), - Statement::SuspendWarehouse(_) => unimplemented!(), - Statement::InspectWarehouse(_) => unimplemented!(), + Statement::RenameWarehouse(v) => self.bind_rename_warehouse(v)?, + Statement::ResumeWarehouse(v) => self.bind_resume_warehouse(v)?, + Statement::SuspendWarehouse(v) => self.bind_suspend_warehouse(v)?, + Statement::InspectWarehouse(v) => self.bind_inspect_warehouse(v)?, Statement::AddWarehouseCluster(_) => unimplemented!(), - Statement::DropWarehouseCluster(_) => unimplemented!(), - Statement::RenameWarehouseCluster(_) => unimplemented!(), + Statement::DropWarehouseCluster(v) => self.bind_drop_warehouse_cluster(v)?, + Statement::RenameWarehouseCluster(v) => self.bind_rename_warehouse_cluster(v)?, Statement::AddWarehouseClusterNode(_) => unimplemented!(), Statement::DropWarehouseClusterNode(_) => unimplemented!(), }; diff --git a/src/query/sql/src/planner/binder/ddl/mod.rs b/src/query/sql/src/planner/binder/ddl/mod.rs index 9ef410d2db7e8..980f31eecb158 100644 --- a/src/query/sql/src/planner/binder/ddl/mod.rs +++ b/src/query/sql/src/planner/binder/ddl/mod.rs @@ -33,3 +33,4 @@ mod table; mod task; mod view; mod virtual_column; +mod warehouse; diff --git a/src/query/sql/src/planner/binder/ddl/warehouse.rs b/src/query/sql/src/planner/binder/ddl/warehouse.rs new file mode 100644 index 0000000000000..5efc5b04dd901 --- /dev/null +++ b/src/query/sql/src/planner/binder/ddl/warehouse.rs @@ -0,0 +1,113 @@ +// Copyright 2021 Datafuse Labs +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use databend_common_ast::ast::DropWarehouseClusterStmt; +use databend_common_ast::ast::DropWarehouseStmt; +use databend_common_ast::ast::InspectWarehouseStmt; +use databend_common_ast::ast::RenameWarehouseClusterStmt; +use databend_common_ast::ast::RenameWarehouseStmt; +use databend_common_ast::ast::ResumeWarehouseStmt; +use databend_common_ast::ast::ShowWarehousesStmt; +use databend_common_ast::ast::SuspendWarehouseStmt; +use databend_common_exception::Result; + +use crate::plans::DropWarehouseClusterPlan; +use crate::plans::DropWarehousePlan; +use crate::plans::InspectWarehousePlan; +use crate::plans::Plan; +use crate::plans::RenameWarehouseClusterPlan; +use crate::plans::RenameWarehousePlan; +use crate::plans::ResumeWarehousePlan; +use crate::plans::SuspendWarehousePlan; +use crate::Binder; + +impl Binder { + pub(in crate::planner::binder) fn bind_show_warehouses( + &mut self, + _stmt: &ShowWarehousesStmt, + ) -> Result { + Ok(Plan::ShowWarehouses) + } + + pub(in crate::planner::binder) fn bind_drop_warehouse( + &mut self, + stmt: &DropWarehouseStmt, + ) -> Result { + Ok(Plan::DropWarehouse(Box::new(DropWarehousePlan { + warehouse: stmt.warehouse.to_string(), + }))) + } + + pub(in crate::planner::binder) fn bind_resume_warehouse( + &mut self, + stmt: &ResumeWarehouseStmt, + ) -> Result { + Ok(Plan::ResumeWarehouse(Box::new(ResumeWarehousePlan { + warehouse: stmt.warehouse.to_string(), + }))) + } + + pub(in crate::planner::binder) fn bind_suspend_warehouse( + &mut self, + stmt: &SuspendWarehouseStmt, + ) -> Result { + Ok(Plan::SuspendWarehouse(Box::new(SuspendWarehousePlan { + warehouse: stmt.warehouse.to_string(), + }))) + } + + pub(in crate::planner::binder) fn bind_inspect_warehouse( + &mut self, + stmt: &InspectWarehouseStmt, + ) -> Result { + Ok(Plan::InspectWarehouse(Box::new(InspectWarehousePlan { + warehouse: stmt.warehouse.to_string(), + }))) + } + + pub(in crate::planner::binder) fn bind_rename_warehouse( + &mut self, + stmt: &RenameWarehouseStmt, + ) -> Result { + Ok(Plan::RenameWarehouse(Box::new(RenameWarehousePlan { + warehouse: stmt.warehouse.to_string(), + new_warehouse: stmt.new_warehouse.to_string(), + }))) + } + + pub(in crate::planner::binder) fn bind_drop_warehouse_cluster( + &mut self, + stmt: &DropWarehouseClusterStmt, + ) -> Result { + Ok(Plan::DropWarehouseCluster(Box::new( + DropWarehouseClusterPlan { + warehouse: stmt.warehouse.to_string(), + cluster: stmt.cluster.to_string(), + }, + ))) + } + + pub(in crate::planner::binder) fn bind_rename_warehouse_cluster( + &mut self, + stmt: &RenameWarehouseClusterStmt, + ) -> Result { + Ok(Plan::RenameWarehouseCluster(Box::new( + RenameWarehouseClusterPlan { + warehouse: stmt.warehouse.to_string(), + cluster: stmt.cluster.to_string(), + new_cluster: stmt.new_cluster.to_string(), + }, + ))) + } +} diff --git a/src/query/sql/src/planner/format/display_plan.rs b/src/query/sql/src/planner/format/display_plan.rs index cbb89f162eaef..3877b75d89b5a 100644 --- a/src/query/sql/src/planner/format/display_plan.rs +++ b/src/query/sql/src/planner/format/display_plan.rs @@ -213,6 +213,18 @@ impl Plan { Plan::DropDictionary(_) => Ok("DropDictionary".to_string()), Plan::ShowCreateDictionary(_) => Ok("ShowCreateDictionary".to_string()), Plan::RenameDictionary(_) => Ok("RenameDictionary".to_string()), + Plan::ShowWarehouses => Ok("ShowWarehouses".to_string()), + Plan::DropWarehouse(_) => Ok("DropWarehouse".to_string()), + Plan::ResumeWarehouse(_) => Ok("ResumeWarehouse".to_string()), + Plan::SuspendWarehouse(_) => Ok("SuspendWarehouse".to_string()), + Plan::RenameWarehouse(_) => Ok("RenameWarehouse".to_string()), + Plan::InspectWarehouse(_) => Ok("InspectWarehouse".to_string()), + Plan::DropWarehouseCluster(_) => Ok("DropWarehouseCluster".to_string()), + Plan::RenameWarehouseCluster(_) => Ok("RenameWarehouseCluster".to_string()), + Plan::CreateWarehouse(_) => Ok("CreateWarehouse".to_string()), + Plan::AddWarehouseCluster(_) => Ok("AddWarehouseCluster".to_string()), + Plan::AddWarehouseClusterNode(_) => Ok("AddWarehouseClusterNode".to_string()), + Plan::DropWarehouseClusterNode(_) => Ok("DropWarehouseClusterNode".to_string()), } } } diff --git a/src/query/sql/src/planner/plans/ddl/mod.rs b/src/query/sql/src/planner/plans/ddl/mod.rs index fe0b14ce21ad0..309bc5340b56f 100644 --- a/src/query/sql/src/planner/plans/ddl/mod.rs +++ b/src/query/sql/src/planner/plans/ddl/mod.rs @@ -30,6 +30,7 @@ mod task; mod udf; mod view; mod virtual_column; +mod warehouse; pub use account::*; pub use catalog::*; @@ -49,3 +50,4 @@ pub use task::*; pub use udf::*; pub use view::*; pub use virtual_column::*; +pub use warehouse::*; diff --git a/src/query/sql/src/planner/plans/ddl/warehouse.rs b/src/query/sql/src/planner/plans/ddl/warehouse.rs new file mode 100644 index 0000000000000..13154396f1574 --- /dev/null +++ b/src/query/sql/src/planner/plans/ddl/warehouse.rs @@ -0,0 +1,85 @@ +// Copyright 2021 Datafuse Labs +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use databend_common_expression::types::DataType; +use databend_common_expression::DataField; +use databend_common_expression::DataSchemaRef; +use databend_common_expression::DataSchemaRefExt; + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct CreateWarehousePlan { + pub warehouse: String, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct DropWarehousePlan { + pub warehouse: String, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct ResumeWarehousePlan { + pub warehouse: String, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct SuspendWarehousePlan { + pub warehouse: String, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct RenameWarehousePlan { + pub warehouse: String, + pub new_warehouse: String, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct InspectWarehousePlan { + pub warehouse: String, +} + +impl InspectWarehousePlan { + pub fn schema(&self) -> DataSchemaRef { + DataSchemaRefExt::create(vec![ + DataField::new("cluster", DataType::String), + DataField::new("node_name", DataType::String), + DataField::new("node_type", DataType::String), + ]) + } +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct AddWarehouseClusterPlan { + warehouse: String, + cluster: String, + // nodes:Vec +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct DropWarehouseClusterPlan { + pub warehouse: String, + pub cluster: String, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct RenameWarehouseClusterPlan { + pub warehouse: String, + pub cluster: String, + pub new_cluster: String, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct AddWarehouseClusterNodePlan {} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct DropWarehouseClusterNodePlan {} diff --git a/src/query/sql/src/planner/plans/plan.rs b/src/query/sql/src/planner/plans/plan.rs index dba3003b023be..2259c200b14aa 100644 --- a/src/query/sql/src/planner/plans/plan.rs +++ b/src/query/sql/src/planner/plans/plan.rs @@ -32,6 +32,8 @@ use crate::binder::ExplainConfig; use crate::optimizer::SExpr; use crate::plans::copy_into_location::CopyIntoLocationPlan; use crate::plans::AddTableColumnPlan; +use crate::plans::AddWarehouseClusterNodePlan; +use crate::plans::AddWarehouseClusterPlan; use crate::plans::AlterNetworkPolicyPlan; use crate::plans::AlterNotificationPlan; use crate::plans::AlterPasswordPolicyPlan; @@ -67,6 +69,7 @@ use crate::plans::CreateUDFPlan; use crate::plans::CreateUserPlan; use crate::plans::CreateViewPlan; use crate::plans::CreateVirtualColumnPlan; +use crate::plans::CreateWarehousePlan; use crate::plans::DescConnectionPlan; use crate::plans::DescDatamaskPolicyPlan; use crate::plans::DescNetworkPolicyPlan; @@ -99,6 +102,9 @@ use crate::plans::DropUDFPlan; use crate::plans::DropUserPlan; use crate::plans::DropViewPlan; use crate::plans::DropVirtualColumnPlan; +use crate::plans::DropWarehouseClusterNodePlan; +use crate::plans::DropWarehouseClusterPlan; +use crate::plans::DropWarehousePlan; use crate::plans::Exchange; use crate::plans::ExecuteImmediatePlan; use crate::plans::ExecuteTaskPlan; @@ -107,6 +113,7 @@ use crate::plans::GrantPrivilegePlan; use crate::plans::GrantRolePlan; use crate::plans::Insert; use crate::plans::InsertMultiTable; +use crate::plans::InspectWarehousePlan; use crate::plans::KillPlan; use crate::plans::ModifyTableColumnPlan; use crate::plans::ModifyTableCommentPlan; @@ -121,7 +128,10 @@ use crate::plans::RemoveStagePlan; use crate::plans::RenameDatabasePlan; use crate::plans::RenameTableColumnPlan; use crate::plans::RenameTablePlan; +use crate::plans::RenameWarehouseClusterPlan; +use crate::plans::RenameWarehousePlan; use crate::plans::Replace; +use crate::plans::ResumeWarehousePlan; use crate::plans::RevertTablePlan; use crate::plans::RevokePrivilegePlan; use crate::plans::RevokeRolePlan; @@ -138,6 +148,7 @@ use crate::plans::ShowFileFormatsPlan; use crate::plans::ShowNetworkPoliciesPlan; use crate::plans::ShowRolesPlan; use crate::plans::ShowTasksPlan; +use crate::plans::SuspendWarehousePlan; use crate::plans::SystemPlan; use crate::plans::TruncateTablePlan; use crate::plans::UndropDatabasePlan; @@ -191,6 +202,20 @@ pub enum Plan { DropCatalog(Box), UseCatalog(Box), + // Warehouses + ShowWarehouses, + CreateWarehouse(Box), + DropWarehouse(Box), + ResumeWarehouse(Box), + SuspendWarehouse(Box), + RenameWarehouse(Box), + InspectWarehouse(Box), + AddWarehouseCluster(Box), + DropWarehouseCluster(Box), + RenameWarehouseCluster(Box), + AddWarehouseClusterNode(Box), + DropWarehouseClusterNode(Box), + // Databases ShowCreateDatabase(Box), CreateDatabase(Box), @@ -495,6 +520,7 @@ impl Plan { Plan::InsertMultiTable(plan) => plan.schema(), Plan::DescUser(plan) => plan.schema(), Plan::Insert(plan) => plan.schema(), + Plan::InspectWarehouse(plan) => plan.schema(), _ => Arc::new(DataSchema::empty()), } From b98d6c7fd049cffaf0e4305568d7bcacf6b7ff2c Mon Sep 17 00:00:00 2001 From: zhang2014 Date: Sun, 22 Dec 2024 23:35:06 +0800 Subject: [PATCH 18/53] feat(cluster): support custom management cluster --- src/query/config/src/config.rs | 12 +----- .../resources_management_kubernetes.rs | 7 ++++ .../resources_management_system.rs | 5 +++ .../src/resources_management.rs | 7 ++++ .../interpreter_show_warehouses.rs | 39 ++++++++++++++++++- src/query/sql/src/planner/plans/plan.rs | 5 +++ 6 files changed, 63 insertions(+), 12 deletions(-) diff --git a/src/query/config/src/config.rs b/src/query/config/src/config.rs index 0c1d4b707b903..a0b90c75c258c 100644 --- a/src/query/config/src/config.rs +++ b/src/query/config/src/config.rs @@ -2955,21 +2955,13 @@ pub struct SpillConfig { #[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize, Args, Default)] #[serde(default, deny_unknown_fields)] -pub struct ResourceManagerConfig { - #[clap( - long = "cluster-type", - value_name = "VALUE", - default_value = "system-managed" - )] +pub struct ResourcesManagementConfig { + #[clap(long = "type", value_name = "VALUE", default_value = "system")] #[serde(rename = "type")] pub typ: String, #[clap(long, value_name = "VALUE")] pub instance_type: Option, - - #[clap(long, value_name = "VALUE", default_value = "18446744073709551615")] - /// Allow space in bytes to spill to local disk. - pub spill_local_disk_max_bytes: u64, } mod cache_config_converters { diff --git a/src/query/ee/src/resource_management/resources_management_kubernetes.rs b/src/query/ee/src/resource_management/resources_management_kubernetes.rs index 7d82318d55819..034a86c7b3b10 100644 --- a/src/query/ee/src/resource_management/resources_management_kubernetes.rs +++ b/src/query/ee/src/resource_management/resources_management_kubernetes.rs @@ -15,6 +15,7 @@ use databend_common_exception::ErrorCode; use databend_common_exception::Result; use databend_common_management::SelectedNode; +use databend_common_management::WarehouseInfo; use databend_common_meta_types::NodeInfo; use databend_enterprise_resources_management::ResourcesManagement; @@ -58,6 +59,12 @@ impl ResourcesManagement for KubernetesResourcesManagement { )) } + async fn list_warehouses(&self) -> Result> { + Err(ErrorCode::Unimplemented( + "Unimplemented kubernetes resources management", + )) + } + async fn add_warehouse_cluster( &self, _: String, diff --git a/src/query/ee/src/resource_management/resources_management_system.rs b/src/query/ee/src/resource_management/resources_management_system.rs index 472e5fb5200e5..e1451577edc07 100644 --- a/src/query/ee/src/resource_management/resources_management_system.rs +++ b/src/query/ee/src/resource_management/resources_management_system.rs @@ -21,6 +21,7 @@ use databend_common_exception::ErrorCode; use databend_common_exception::Result; use databend_common_management::SelectedNode; use databend_common_management::WarehouseApi; +use databend_common_management::WarehouseInfo; use databend_common_management::WarehouseMgr; use databend_common_meta_store::MetaStoreProvider; use databend_common_meta_types::NodeInfo; @@ -56,6 +57,10 @@ impl ResourcesManagement for SystemResourcesManagement { self.warehouse_manager.list_warehouse_nodes(name).await } + async fn list_warehouses(&self) -> Result> { + self.warehouse_manager.list_warehouses().await + } + async fn add_warehouse_cluster( &self, name: String, diff --git a/src/query/ee_features/resources_management/src/resources_management.rs b/src/query/ee_features/resources_management/src/resources_management.rs index 8d3e3b59ce36a..21261e6b5c3eb 100644 --- a/src/query/ee_features/resources_management/src/resources_management.rs +++ b/src/query/ee_features/resources_management/src/resources_management.rs @@ -16,6 +16,7 @@ use databend_common_base::base::GlobalInstance; use databend_common_exception::ErrorCode; use databend_common_exception::Result; use databend_common_management::SelectedNode; +use databend_common_management::WarehouseInfo; use databend_common_meta_types::NodeInfo; #[async_trait::async_trait] @@ -32,6 +33,8 @@ pub trait ResourcesManagement: Sync + Send + 'static { async fn inspect_warehouse(&self, name: String) -> Result>; + async fn list_warehouses(&self) -> Result>; + async fn add_warehouse_cluster( &self, name: String, @@ -91,6 +94,10 @@ impl ResourcesManagement for DummyResourcesManagement { Err(ErrorCode::Unimplemented("The use of this feature requires a Databend Enterprise Edition license. To unlock enterprise features, please contact Databend to obtain a license. Learn more at https://docs.databend.com/guides/overview/editions/dee/")) } + async fn list_warehouses(&self) -> Result> { + Err(ErrorCode::Unimplemented("The use of this feature requires a Databend Enterprise Edition license. To unlock enterprise features, please contact Databend to obtain a license. Learn more at https://docs.databend.com/guides/overview/editions/dee/")) + } + async fn add_warehouse_cluster( &self, _: String, diff --git a/src/query/service/src/interpreters/interpreter_show_warehouses.rs b/src/query/service/src/interpreters/interpreter_show_warehouses.rs index 52f88b65914da..4dbcd23e65c7d 100644 --- a/src/query/service/src/interpreters/interpreter_show_warehouses.rs +++ b/src/query/service/src/interpreters/interpreter_show_warehouses.rs @@ -12,7 +12,16 @@ // See the License for the specific language governing permissions and // limitations under the License. +use std::sync::Arc; + +use databend_common_base::base::GlobalInstance; use databend_common_exception::Result; +use databend_common_expression::types::DataType; +use databend_common_expression::ColumnBuilder; +use databend_common_expression::DataBlock; +use databend_common_expression::Scalar; +use databend_common_management::WarehouseInfo; +use databend_enterprise_resources_management::ResourcesManagement; use crate::interpreters::Interpreter; use crate::pipelines::PipelineBuildResult; @@ -37,7 +46,33 @@ impl Interpreter for ShowWarehousesInterpreter { #[async_backtrace::framed] async fn execute2(&self) -> Result { - // WarehouseMgr::i - Ok(PipelineBuildResult::create()) + let warehouses = GlobalInstance::get::>() + .list_warehouses() + .await?; + let mut warehouses_name = ColumnBuilder::with_capacity(&DataType::String, warehouses.len()); + let mut warehouses_type = ColumnBuilder::with_capacity(&DataType::String, warehouses.len()); + let mut warehouses_status = + ColumnBuilder::with_capacity(&DataType::String, warehouses.len()); + + for warehouse in warehouses { + match warehouse { + WarehouseInfo::SelfManaged(name) => { + warehouses_name.push(Scalar::String(name).as_ref()); + warehouses_type.push(Scalar::String(String::from("Self-Managed")).as_ref()); + warehouses_status.push(Scalar::String(String::from("Running")).as_ref()); + } + WarehouseInfo::SystemManaged(v) => { + warehouses_name.push(Scalar::String(v.display_name.clone()).as_ref()); + warehouses_type.push(Scalar::String(String::from("System-Managed")).as_ref()); + warehouses_status.push(Scalar::String(v.status.clone()).as_ref()); + } + } + } + + PipelineBuildResult::from_blocks(vec![DataBlock::new_from_columns(vec![ + warehouses_name.build(), + warehouses_type.build(), + warehouses_status.build(), + ])]) } } diff --git a/src/query/sql/src/planner/plans/plan.rs b/src/query/sql/src/planner/plans/plan.rs index 2259c200b14aa..94b7a85e028ff 100644 --- a/src/query/sql/src/planner/plans/plan.rs +++ b/src/query/sql/src/planner/plans/plan.rs @@ -521,6 +521,11 @@ impl Plan { Plan::DescUser(plan) => plan.schema(), Plan::Insert(plan) => plan.schema(), Plan::InspectWarehouse(plan) => plan.schema(), + Plan::ShowWarehouses => DataSchemaRefExt::create(vec![ + DataField::new("warehouse", DataType::String), + DataField::new("type", DataType::String), + DataField::new("status", DataType::String), + ]), _ => Arc::new(DataSchema::empty()), } From d723ef3289ace79a35f5a54aa70f3c285a0dbbd5 Mon Sep 17 00:00:00 2001 From: zhang2014 Date: Mon, 23 Dec 2024 12:35:01 +0800 Subject: [PATCH 19/53] feat(cluster): support custom management cluster --- src/query/ast/src/ast/statements/warehouse.rs | 95 +++++++++++++++++-- src/query/ast/src/parser/dynamic_table.rs | 4 +- src/query/ast/src/parser/statement.rs | 59 +++++++++++- src/query/ast/src/parser/token.rs | 2 + .../interpreter_create_warehouses.rs | 42 +++++++- .../sql/src/planner/binder/ddl/warehouse.rs | 29 ++++++ .../sql/src/planner/plans/ddl/warehouse.rs | 5 + 7 files changed, 219 insertions(+), 17 deletions(-) diff --git a/src/query/ast/src/ast/statements/warehouse.rs b/src/query/ast/src/ast/statements/warehouse.rs index 968a06f6cd339..474885666b228 100644 --- a/src/query/ast/src/ast/statements/warehouse.rs +++ b/src/query/ast/src/ast/statements/warehouse.rs @@ -12,6 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. +use std::collections::BTreeMap; use std::fmt::Display; use std::fmt::Formatter; @@ -32,11 +33,45 @@ impl Display for ShowWarehousesStmt { #[derive(Debug, Clone, PartialEq, Drive, DriveMut)] pub struct CreateWarehouseStmt { pub warehouse: Identifier, + pub node_list: Vec<(Option, u64)>, + pub options: BTreeMap, } impl Display for CreateWarehouseStmt { - fn fmt(&self, _f: &mut Formatter<'_>) -> std::fmt::Result { - todo!() + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + write!(f, "CREATE WAREHOUSE {}", self.warehouse)?; + + if !self.node_list.is_empty() { + write!(f, "(")?; + + for (idx, (resources_group, nodes)) in self.node_list.iter().enumerate() { + if idx != 0 { + write!(f, ",")?; + } + + write!(f, " ASSIGN {} NODES", nodes)?; + + if let Some(resources_group) = resources_group { + write!(f, " FROM '{}'", resources_group)?; + } + } + + write!(f, "(")?; + } + + if !self.options.is_empty() { + write!(f, " WITH ")?; + + for (idx, (key, value)) in self.options.iter().enumerate() { + if idx != 0 { + write!(f, ",")?; + } + + write!(f, " {} = '{}'", key, value)?; + } + } + + Ok(()) } } @@ -101,11 +136,52 @@ impl Display for InspectWarehouseStmt { } #[derive(Debug, Clone, PartialEq, Drive, DriveMut)] -pub struct AddWarehouseClusterStmt {} +pub struct AddWarehouseClusterStmt { + pub warehouse: Identifier, + pub cluster: Identifier, + pub node_list: Vec<(Option, u64)>, + pub options: BTreeMap, +} impl Display for AddWarehouseClusterStmt { - fn fmt(&self, _f: &mut Formatter<'_>) -> std::fmt::Result { - todo!() + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + write!( + f, + "ALTER WAREHOUSE {} ADD CLUSTER {}", + self.warehouse, self.cluster + )?; + + if !self.node_list.is_empty() { + write!(f, "(")?; + + for (idx, (resources_group, nodes)) in self.node_list.iter().enumerate() { + if idx != 0 { + write!(f, ",")?; + } + + write!(f, " ASSIGN {} NODES", nodes)?; + + if let Some(resources_group) = resources_group { + write!(f, " FROM '{}'", resources_group)?; + } + } + + write!(f, "(")?; + } + + if !self.options.is_empty() { + write!(f, " WITH ")?; + + for (idx, (key, value)) in self.options.iter().enumerate() { + if idx != 0 { + write!(f, ",")?; + } + + write!(f, " {} = '{}'", key, value)?; + } + } + + Ok(()) } } @@ -143,7 +219,12 @@ impl Display for RenameWarehouseClusterStmt { } #[derive(Debug, Clone, PartialEq, Drive, DriveMut)] -pub struct AddWarehouseClusterNodeStmt {} +pub struct AddWarehouseClusterNodeStmt { + pub warehouse: Identifier, + pub cluster: Identifier, + pub node_list: Vec<(Option, u64)>, + pub options: BTreeMap, +} impl Display for AddWarehouseClusterNodeStmt { fn fmt(&self, _f: &mut Formatter<'_>) -> std::fmt::Result { @@ -153,6 +234,8 @@ impl Display for AddWarehouseClusterNodeStmt { #[derive(Debug, Clone, PartialEq, Drive, DriveMut)] pub struct DropWarehouseClusterNodeStmt { + pub warehouse: Identifier, + pub cluster: Identifier, // warehouse:String, // cluster: String, } diff --git a/src/query/ast/src/parser/dynamic_table.rs b/src/query/ast/src/parser/dynamic_table.rs index b766bdfeec85b..b960a64c4ce30 100644 --- a/src/query/ast/src/parser/dynamic_table.rs +++ b/src/query/ast/src/parser/dynamic_table.rs @@ -38,7 +38,7 @@ use crate::parser::statement::cluster_type; use crate::parser::statement::create_table_source; use crate::parser::statement::parse_create_option; use crate::parser::statement::table_option; -use crate::parser::statement::warehouse_option; +use crate::parser::statement::task_warehouse_option; use crate::parser::token::TokenKind::*; use crate::parser::Input; @@ -145,7 +145,7 @@ fn dynamic_table_options( permutation(( target_lag, - warehouse_option, + task_warehouse_option, refresh_mode_opt, initialize_opt, ))(i) diff --git a/src/query/ast/src/parser/statement.rs b/src/query/ast/src/parser/statement.rs index 45244fe018ffd..b4c2eb6578134 100644 --- a/src/query/ast/src/parser/statement.rs +++ b/src/query/ast/src/parser/statement.rs @@ -125,7 +125,7 @@ pub fn statement_body(i: Input) -> IResult { rule! { CREATE ~ TASK ~ ( IF ~ ^NOT ~ ^EXISTS )? ~ #ident - ~ #warehouse_option + ~ #task_warehouse_option ~ ( SCHEDULE ~ "=" ~ #task_schedule_option )? ~ ( AFTER ~ #comma_separated_list0(literal_string) )? ~ ( WHEN ~ #expr )? @@ -530,11 +530,17 @@ pub fn statement_body(i: Input) -> IResult { |(_, _)| Statement::ShowWarehouses(ShowWarehousesStmt {}), ); - let _create_warehouse = map( + let create_warehouse = map( rule! { - CREATE ~ WAREHOUSE ~ #ident + CREATE ~ WAREHOUSE ~ #ident ~ ("(" ~ #assign_nodes_list ~ ")")? ~ (WITH ~ #warehouse_cluster_option)? + }, + |(_, _, warehouse, nodes, options)| { + Statement::CreateWarehouse(CreateWarehouseStmt { + warehouse, + node_list: nodes.map(|(_, nodes, _)| nodes).unwrap_or_else(Vec::new), + options: options.map(|(_, x)| x).unwrap_or_else(BTreeMap::new), + }) }, - |(_, _, warehouse)| Statement::CreateWarehouse(CreateWarehouseStmt { warehouse }), ); let drop_warehouse = map( @@ -577,6 +583,20 @@ pub fn statement_body(i: Input) -> IResult { |(_, _, warehouse)| Statement::InspectWarehouse(InspectWarehouseStmt { warehouse }), ); + let add_warehouse_cluster = map( + rule! { + ALTER ~ WAREHOUSE ~ #ident ~ ADD ~ CLUSTER ~ #ident ~ ("(" ~ #assign_nodes_list ~ ")")? ~ (WITH ~ #warehouse_cluster_option)? + }, + |(_, _, warehouse, _, _, cluster, nodes, options)| { + Statement::AddWarehouseCluster(AddWarehouseClusterStmt { + warehouse, + cluster, + node_list: nodes.map(|(_, nodes, _)| nodes).unwrap_or_else(Vec::new), + options: options.map(|(_, x)| x).unwrap_or_else(BTreeMap::new), + }) + }, + ); + let drop_warehouse_cluster = map( rule! { ALTER ~ WAREHOUSE ~ #ident ~ DROP ~ CLUSTER ~ #ident @@ -2371,11 +2391,13 @@ pub fn statement_body(i: Input) -> IResult { // warehouse rule!( #show_warehouses: "`SHOW WAREHOUSES`" + | #create_warehouse: "`CREATE WAREHOUSE [(ASSIGN NODES [FROM ] [, ...])] WITH [warehouse_size = ]`" | #drop_warehouse: "`DROP WAREHOUSE `" | #rename_warehouse: "`RENAME WAREHOUSE TO `" | #resume_warehouse: "`RESUME WAREHOUSE `" | #suspend_warehouse: "`SUSPEND WAREHOUSE `" | #inspect_warehouse: "`INSPECT WAREHOUSE `" + | #add_warehouse_cluster: "`ALTER WAREHOUSE ADD CLUSTER [(ASSIGN NODES [FROM ] [, ...])] WITH [cluster_size = ]`" | #drop_warehouse_cluster: "`ALTER WAREHOUSE DROP CLUSTER `" | #rename_warehouse_cluster: "`ALTER WAREHOUSE RENAME CLUSTER TO `" ), @@ -4026,7 +4048,7 @@ pub fn alter_pipe_option(i: Input) -> IResult { )(i) } -pub fn warehouse_option(i: Input) -> IResult { +pub fn task_warehouse_option(i: Input) -> IResult { alt((map( rule! { (WAREHOUSE ~ "=" ~ #literal_string)? @@ -4041,6 +4063,33 @@ pub fn warehouse_option(i: Input) -> IResult { ),))(i) } +pub fn assign_nodes_list(i: Input) -> IResult, u64)>> { + let nodes_list = map( + rule! { + ASSIGN ~ #literal_u64 ~ (FROM ~ #option_to_string)? + }, + |(_, node_size, resources_group)| (resources_group.map(|(_, x)| x), node_size), + ); + + map(comma_separated_list1(nodes_list), |opts| { + opts.into_iter().collect() + })(i) +} + +pub fn warehouse_cluster_option(i: Input) -> IResult> { + let option = map( + rule! { + #ident ~ "=" ~ #option_to_string + }, + |(k, _, v)| (k, v), + ); + map(comma_separated_list1(option), |opts| { + opts.into_iter() + .map(|(k, v)| (k.name.to_lowercase(), v.clone())) + .collect() + })(i) +} + pub fn task_schedule_option(i: Input) -> IResult { let interval = map( rule! { diff --git a/src/query/ast/src/parser/token.rs b/src/query/ast/src/parser/token.rs index a4a5011a4c47e..84bbcd2167e50 100644 --- a/src/query/ast/src/parser/token.rs +++ b/src/query/ast/src/parser/token.rs @@ -1380,6 +1380,8 @@ pub enum TokenKind { WAREHOUSES, #[token("INSPECT", ignore(ascii_case))] INSPECT, + #[token("ASSIGN", ignore(ascii_case))] + ASSIGN, } // Reference: https://www.postgresql.org/docs/current/sql-keywords-appendix.html diff --git a/src/query/service/src/interpreters/interpreter_create_warehouses.rs b/src/query/service/src/interpreters/interpreter_create_warehouses.rs index 609d32c83c844..e801b3b02ae05 100644 --- a/src/query/service/src/interpreters/interpreter_create_warehouses.rs +++ b/src/query/service/src/interpreters/interpreter_create_warehouses.rs @@ -14,8 +14,12 @@ use std::sync::Arc; +use databend_common_base::base::GlobalInstance; +use databend_common_exception::ErrorCode; use databend_common_exception::Result; +use databend_common_management::SelectedNode; use databend_common_sql::plans::CreateWarehousePlan; +use databend_enterprise_resources_management::ResourcesManagement; use crate::interpreters::Interpreter; use crate::pipelines::PipelineBuildResult; @@ -24,7 +28,6 @@ use crate::sessions::QueryContext; pub struct CreateWarehouseInterpreter { #[allow(dead_code)] ctx: Arc, - #[allow(dead_code)] plan: CreateWarehousePlan, } @@ -46,9 +49,40 @@ impl Interpreter for CreateWarehouseInterpreter { #[async_backtrace::framed] async fn execute2(&self) -> Result { - // GlobalInstance::get::>() - // .create_warehouse(self.plan.warehouse.clone()) - // .await?; + if let Some(warehouse_size) = self.plan.options.get("WAREHOUSE_SIZE") { + if !self.plan.nodes.is_empty() { + return Err(ErrorCode::InvalidArgument(format!(""))); + } + + let Ok(warehouse_size) = warehouse_size.parse::() else { + return Err(ErrorCode::InvalidArgument("")); + }; + + GlobalInstance::get::>() + .create_warehouse(self.plan.warehouse.clone(), vec![ + SelectedNode::Random(None); + warehouse_size + ]) + .await?; + + return Ok(PipelineBuildResult::create()); + } + + if self.plan.nodes.is_empty() { + return Err(ErrorCode::InvalidArgument("")); + } + + let mut selected_nodes = Vec::with_capacity(self.plan.nodes.len()); + for (group, nodes) in self.plan.nodes { + for _ in 0..nodes { + selected_nodes.push(SelectedNode::Random(group)); + } + } + + GlobalInstance::get::>() + .create_warehouse(self.plan.warehouse.clone(), selected_nodes) + .await?; + Ok(PipelineBuildResult::create()) } } diff --git a/src/query/sql/src/planner/binder/ddl/warehouse.rs b/src/query/sql/src/planner/binder/ddl/warehouse.rs index 5efc5b04dd901..b292e873154a5 100644 --- a/src/query/sql/src/planner/binder/ddl/warehouse.rs +++ b/src/query/sql/src/planner/binder/ddl/warehouse.rs @@ -12,6 +12,10 @@ // See the License for the specific language governing permissions and // limitations under the License. +use std::collections::hash_map::Entry; +use std::collections::HashMap; + +use databend_common_ast::ast::CreateWarehouseStmt; use databend_common_ast::ast::DropWarehouseClusterStmt; use databend_common_ast::ast::DropWarehouseStmt; use databend_common_ast::ast::InspectWarehouseStmt; @@ -22,6 +26,7 @@ use databend_common_ast::ast::ShowWarehousesStmt; use databend_common_ast::ast::SuspendWarehouseStmt; use databend_common_exception::Result; +use crate::plans::CreateWarehousePlan; use crate::plans::DropWarehouseClusterPlan; use crate::plans::DropWarehousePlan; use crate::plans::InspectWarehousePlan; @@ -40,6 +45,30 @@ impl Binder { Ok(Plan::ShowWarehouses) } + pub(in crate::planner::binder) fn bind_create_warehouse( + &mut self, + stmt: &CreateWarehouseStmt, + ) -> Result { + let mut nodes = HashMap::with_capacity(stmt.node_list.len()); + + for (group, nodes_num) in stmt.node_list { + match nodes.entry(group) { + Entry::Vacant(mut v) => { + v.insert(nodes_num); + } + Entry::Occupied(mut v) => { + *v.get_mut() += nodes_num; + } + } + } + + Ok(Plan::CreateWarehouse(Box::new(CreateWarehousePlan { + warehouse: stmt.warehouse.to_string(), + nodes: Default::default(), + options: stmt.options, + }))) + } + pub(in crate::planner::binder) fn bind_drop_warehouse( &mut self, stmt: &DropWarehouseStmt, diff --git a/src/query/sql/src/planner/plans/ddl/warehouse.rs b/src/query/sql/src/planner/plans/ddl/warehouse.rs index 13154396f1574..f1168f232b6ca 100644 --- a/src/query/sql/src/planner/plans/ddl/warehouse.rs +++ b/src/query/sql/src/planner/plans/ddl/warehouse.rs @@ -12,6 +12,9 @@ // See the License for the specific language governing permissions and // limitations under the License. +use std::collections::BTreeMap; +use std::collections::HashMap; + use databend_common_expression::types::DataType; use databend_common_expression::DataField; use databend_common_expression::DataSchemaRef; @@ -20,6 +23,8 @@ use databend_common_expression::DataSchemaRefExt; #[derive(Clone, Debug, PartialEq, Eq)] pub struct CreateWarehousePlan { pub warehouse: String, + pub nodes: HashMap, u64>, + pub options: BTreeMap, } #[derive(Clone, Debug, PartialEq, Eq)] From aa4bd74b93bfeafafac04ab01cfd68c5e3669a16 Mon Sep 17 00:00:00 2001 From: zhang2014 Date: Mon, 23 Dec 2024 17:58:48 +0800 Subject: [PATCH 20/53] feat(cluster): support custom management cluster --- src/meta/types/src/cluster.rs | 5 + src/meta/types/tests/it/cluster.rs | 4 + src/query/ast/src/ast/statements/statement.rs | 12 +- src/query/ast/src/ast/statements/warehouse.rs | 56 +- src/query/ast/src/parser/statement.rs | 56 ++ src/query/ast/src/parser/token.rs | 4 + .../resources_management_kubernetes.rs | 12 +- .../resources_management_system.rs | 15 +- .../src/resources_management.rs | 22 +- .../management/src/warehouse/warehouse_api.rs | 12 +- .../management/src/warehouse/warehouse_mgr.rs | 576 +++++++++++------- src/query/management/tests/it/warehouse.rs | 4 + .../interpreters/access/privilege_access.rs | 4 +- .../interpreter_add_warehouse_cluster.rs | 47 +- .../interpreter_add_warehouse_cluster_node.rs | 49 -- .../interpreter_assign_warehouse_nodes.rs | 66 ++ .../interpreter_create_warehouses.rs | 6 +- ...interpreter_drop_warehouse_cluster_node.rs | 6 +- .../src/interpreters/interpreter_factory.rs | 8 +- src/query/service/src/interpreters/mod.rs | 2 +- src/query/sql/src/planner/binder/binder.rs | 8 +- .../sql/src/planner/binder/ddl/warehouse.rs | 105 +++- .../sql/src/planner/format/display_plan.rs | 4 +- .../sql/src/planner/plans/ddl/warehouse.rs | 17 +- src/query/sql/src/planner/plans/plan.rs | 8 +- 25 files changed, 745 insertions(+), 363 deletions(-) delete mode 100644 src/query/service/src/interpreters/interpreter_add_warehouse_cluster_node.rs create mode 100644 src/query/service/src/interpreters/interpreter_assign_warehouse_nodes.rs diff --git a/src/meta/types/src/cluster.rs b/src/meta/types/src/cluster.rs index 8bdae6f0750bd..90d122038d062 100644 --- a/src/meta/types/src/cluster.rs +++ b/src/meta/types/src/cluster.rs @@ -90,11 +90,14 @@ pub struct NodeInfo { pub discovery_address: String, pub binary_version: String, pub node_type: NodeType, + pub resource_group: Option, #[serde(skip_serializing_if = "String::is_empty")] pub cluster_id: String, #[serde(skip_serializing_if = "String::is_empty")] pub warehouse_id: String, + + pub runtime_resource_group: Option, } impl NodeInfo { @@ -119,6 +122,8 @@ impl NodeInfo { cluster_id: "".to_string(), warehouse_id: "".to_string(), node_type: NodeType::SystemManaged, + resource_group: None, + runtime_resource_group: None, } } diff --git a/src/meta/types/tests/it/cluster.rs b/src/meta/types/tests/it/cluster.rs index 49a814472a936..f7a546c3aba31 100644 --- a/src/meta/types/tests/it/cluster.rs +++ b/src/meta/types/tests/it/cluster.rs @@ -26,8 +26,10 @@ fn test_node_info_ip_port() -> anyhow::Result<()> { discovery_address: "4.5.6.7:456".to_string(), binary_version: "v0.8-binary-version".to_string(), node_type: Default::default(), + resource_group: None, cluster_id: "".to_string(), warehouse_id: "".to_string(), + runtime_resource_group: None, }; let (ip, port) = n.ip_port()?; @@ -50,8 +52,10 @@ fn test_serde_node_info() { discovery_address: "4.5.6.7:456".to_string(), binary_version: "v0.8-binary-version".to_string(), node_type: Default::default(), + resource_group: None, cluster_id: String::new(), warehouse_id: String::new(), + runtime_resource_group: None, }; let json_str = serde_json::to_string(&info).unwrap(); diff --git a/src/query/ast/src/ast/statements/statement.rs b/src/query/ast/src/ast/statements/statement.rs index a5d5b89a88e24..62e528300efe0 100644 --- a/src/query/ast/src/ast/statements/statement.rs +++ b/src/query/ast/src/ast/statements/statement.rs @@ -140,8 +140,8 @@ pub enum Statement { AddWarehouseCluster(AddWarehouseClusterStmt), DropWarehouseCluster(DropWarehouseClusterStmt), RenameWarehouseCluster(RenameWarehouseClusterStmt), - AddWarehouseClusterNode(AddWarehouseClusterNodeStmt), - DropWarehouseClusterNode(DropWarehouseClusterNodeStmt), + AssignWarehouseNodes(AssignWarehouseNodesStmt), + UnassignWarehouseNodes(UnassignWarehouseNodesStmt), // Databases ShowDatabases(ShowDatabasesStmt), @@ -560,8 +560,8 @@ impl Statement { | Statement::AddWarehouseCluster(..) | Statement::DropWarehouseCluster(..) | Statement::RenameWarehouseCluster(..) - | Statement::AddWarehouseClusterNode(..) - | Statement::DropWarehouseClusterNode(..) + | Statement::AssignWarehouseNodes(..) + | Statement::UnassignWarehouseNodes(..) | Statement::ResumeWarehouse(..) | Statement::SuspendWarehouse(..) => false, Statement::StatementWithSettings { stmt, settings: _ } => { @@ -985,8 +985,8 @@ impl Display for Statement { Statement::AddWarehouseCluster(stmt) => write!(f, "{stmt}")?, Statement::DropWarehouseCluster(stmt) => write!(f, "{stmt}")?, Statement::RenameWarehouseCluster(stmt) => write!(f, "{stmt}")?, - Statement::AddWarehouseClusterNode(stmt) => write!(f, "{stmt}")?, - Statement::DropWarehouseClusterNode(stmt) => write!(f, "{stmt}")?, + Statement::AssignWarehouseNodes(stmt) => write!(f, "{stmt}")?, + Statement::UnassignWarehouseNodes(stmt) => write!(f, "{stmt}")?, } Ok(()) } diff --git a/src/query/ast/src/ast/statements/warehouse.rs b/src/query/ast/src/ast/statements/warehouse.rs index 474885666b228..10d36c2cb6f29 100644 --- a/src/query/ast/src/ast/statements/warehouse.rs +++ b/src/query/ast/src/ast/statements/warehouse.rs @@ -219,29 +219,57 @@ impl Display for RenameWarehouseClusterStmt { } #[derive(Debug, Clone, PartialEq, Drive, DriveMut)] -pub struct AddWarehouseClusterNodeStmt { +pub struct AssignWarehouseNodesStmt { pub warehouse: Identifier, - pub cluster: Identifier, - pub node_list: Vec<(Option, u64)>, - pub options: BTreeMap, + pub node_list: Vec<(Identifier, Option, u64)>, } -impl Display for AddWarehouseClusterNodeStmt { - fn fmt(&self, _f: &mut Formatter<'_>) -> std::fmt::Result { - todo!() +impl Display for AssignWarehouseNodesStmt { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + write!(f, "ALTER WAREHOUSE {} ASSIGN NODES (", self.warehouse)?; + + for (idx, (cluster, group, nodes)) in self.node_list.iter().enumerate() { + if idx != 0 { + write!(f, ", ")?; + } + + write!(f, "ASSIGN {} NODES", nodes)?; + + if let Some(group) = group { + write!(f, " FROM '{}'", group)?; + } + + write!(f, " FOR {}", cluster)?; + } + + Ok(()) } } #[derive(Debug, Clone, PartialEq, Drive, DriveMut)] -pub struct DropWarehouseClusterNodeStmt { +pub struct UnassignWarehouseNodesStmt { pub warehouse: Identifier, - pub cluster: Identifier, - // warehouse:String, - // cluster: String, + pub node_list: Vec<(Identifier, Option, u64)>, } -impl Display for DropWarehouseClusterNodeStmt { - fn fmt(&self, _f: &mut Formatter<'_>) -> std::fmt::Result { - todo!() +impl Display for UnassignWarehouseNodesStmt { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + write!(f, "ALTER WAREHOUSE {} UNASSIGN NODES (", self.warehouse)?; + + for (idx, (cluster, group, nodes)) in self.node_list.iter().enumerate() { + if idx != 0 { + write!(f, ", ")?; + } + + write!(f, "UNASSIGN {} NODES", nodes)?; + + if let Some(group) = group { + write!(f, " FROM '{}'", group)?; + } + + write!(f, " FOR {}", cluster)?; + } + + Ok(()) } } diff --git a/src/query/ast/src/parser/statement.rs b/src/query/ast/src/parser/statement.rs index b4c2eb6578134..46b100e798d2b 100644 --- a/src/query/ast/src/parser/statement.rs +++ b/src/query/ast/src/parser/statement.rs @@ -619,6 +619,30 @@ pub fn statement_body(i: Input) -> IResult { }, ); + let assign_warehouse_nodes = map( + rule! { + ALTER ~ WAREHOUSE ~ #ident ~ ASSIGN ~ NODES ~ "(" ~ #assign_warehouse_nodes_list ~ ")" + }, + |(_, _, warehouse, _, _, _, nodes, _)| { + Statement::AssignWarehouseNodes(AssignWarehouseNodesStmt { + warehouse, + node_list: nodes, + }) + }, + ); + + let unassign_warehouse_nodes = map( + rule! { + ALTER ~ WAREHOUSE ~ #ident ~ UNASSIGN ~ NODES ~ "(" ~ #unassign_warehouse_nodes_list ~ ")" + }, + |(_, _, warehouse, _, _, _, nodes, _)| { + Statement::UnassignWarehouseNodes(UnassignWarehouseNodesStmt { + warehouse, + node_list: nodes, + }) + }, + ); + let show_databases = map( rule! { SHOW ~ FULL? ~ ( DATABASES | SCHEMAS ) ~ ( ( FROM | IN ) ~ ^#ident )? ~ #show_limit? @@ -2400,6 +2424,8 @@ pub fn statement_body(i: Input) -> IResult { | #add_warehouse_cluster: "`ALTER WAREHOUSE ADD CLUSTER [(ASSIGN NODES [FROM ] [, ...])] WITH [cluster_size = ]`" | #drop_warehouse_cluster: "`ALTER WAREHOUSE DROP CLUSTER `" | #rename_warehouse_cluster: "`ALTER WAREHOUSE RENAME CLUSTER TO `" + | #assign_warehouse_nodes: "`ALTER WAREHOUSE ASSIGN NODES ( ASSIGN NODES [FROM ] FOR [, ...] )`" + | #unassign_warehouse_nodes: "`ALTER WAREHOUSE UNASSIGN NODES ( UNASSIGN NODES [FROM ] FOR [, ...] )`" ), // database rule!( @@ -4076,6 +4102,36 @@ pub fn assign_nodes_list(i: Input) -> IResult, u64)>> { })(i) } +pub fn assign_warehouse_nodes_list(i: Input) -> IResult, u64)>> { + let nodes_list = map( + rule! { + ASSIGN ~ #literal_u64 ~ NODES ~ (FROM ~ #option_to_string)? ~ FOR ~ #ident + }, + |(_, node_size, _, resources_group, _, cluster)| { + (cluster, resources_group.map(|(_, x)| x), node_size) + }, + ); + + map(comma_separated_list1(nodes_list), |opts| { + opts.into_iter().collect() + })(i) +} + +pub fn unassign_warehouse_nodes_list(i: Input) -> IResult, u64)>> { + let nodes_list = map( + rule! { + UNASSIGN ~ #literal_u64 ~ NODES ~ (FROM ~ #option_to_string)? ~ FOR ~ #ident + }, + |(_, node_size, _, resources_group, _, cluster)| { + (cluster, resources_group.map(|(_, x)| x), node_size) + }, + ); + + map(comma_separated_list1(nodes_list), |opts| { + opts.into_iter().collect() + })(i) +} + pub fn warehouse_cluster_option(i: Input) -> IResult> { let option = map( rule! { diff --git a/src/query/ast/src/parser/token.rs b/src/query/ast/src/parser/token.rs index 84bbcd2167e50..b9a266d711a98 100644 --- a/src/query/ast/src/parser/token.rs +++ b/src/query/ast/src/parser/token.rs @@ -1382,6 +1382,10 @@ pub enum TokenKind { INSPECT, #[token("ASSIGN", ignore(ascii_case))] ASSIGN, + #[token("NODES", ignore(ascii_case))] + NODES, + #[token("UNASSIGN", ignore(ascii_case))] + UNASSIGN, } // Reference: https://www.postgresql.org/docs/current/sql-keywords-appendix.html diff --git a/src/query/ee/src/resource_management/resources_management_kubernetes.rs b/src/query/ee/src/resource_management/resources_management_kubernetes.rs index 034a86c7b3b10..784050b2d16f4 100644 --- a/src/query/ee/src/resource_management/resources_management_kubernetes.rs +++ b/src/query/ee/src/resource_management/resources_management_kubernetes.rs @@ -12,6 +12,8 @@ // See the License for the specific language governing permissions and // limitations under the License. +use std::collections::HashMap; + use databend_common_exception::ErrorCode; use databend_common_exception::Result; use databend_common_management::SelectedNode; @@ -88,22 +90,20 @@ impl ResourcesManagement for KubernetesResourcesManagement { )) } - async fn add_warehouse_cluster_node( + async fn assign_warehouse_nodes( &self, _: String, - _: String, - _: Vec, + _: HashMap>, ) -> Result<()> { Err(ErrorCode::Unimplemented( "Unimplemented kubernetes resources management", )) } - async fn drop_warehouse_cluster_node( + async fn unassign_warehouse_nodes( &self, _: String, - _: String, - _: Vec, + _: HashMap>, ) -> Result<()> { Err(ErrorCode::Unimplemented( "Unimplemented kubernetes resources management", diff --git a/src/query/ee/src/resource_management/resources_management_system.rs b/src/query/ee/src/resource_management/resources_management_system.rs index e1451577edc07..5b7f412449416 100644 --- a/src/query/ee/src/resource_management/resources_management_system.rs +++ b/src/query/ee/src/resource_management/resources_management_system.rs @@ -12,6 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. +use std::collections::HashMap; use std::sync::Arc; use std::time::Duration; @@ -89,25 +90,23 @@ impl ResourcesManagement for SystemResourcesManagement { .await } - async fn add_warehouse_cluster_node( + async fn assign_warehouse_nodes( &self, name: String, - cluster: String, - nodes: Vec, + nodes: HashMap>, ) -> Result<()> { self.warehouse_manager - .add_warehouse_cluster_node(&name, &cluster, nodes) + .assign_warehouse_nodes(name, nodes) .await } - async fn drop_warehouse_cluster_node( + async fn unassign_warehouse_nodes( &self, name: String, - cluster: String, - nodes: Vec, + nodes: HashMap>, ) -> Result<()> { self.warehouse_manager - .drop_warehouse_cluster_node(&name, &cluster, nodes) + .unassign_warehouse_nodes(&name, nodes) .await } } diff --git a/src/query/ee_features/resources_management/src/resources_management.rs b/src/query/ee_features/resources_management/src/resources_management.rs index 21261e6b5c3eb..66ac7096b8bff 100644 --- a/src/query/ee_features/resources_management/src/resources_management.rs +++ b/src/query/ee_features/resources_management/src/resources_management.rs @@ -12,6 +12,8 @@ // See the License for the specific language governing permissions and // limitations under the License. +use std::collections::HashMap; + use databend_common_base::base::GlobalInstance; use databend_common_exception::ErrorCode; use databend_common_exception::Result; @@ -51,18 +53,16 @@ pub trait ResourcesManagement: Sync + Send + 'static { async fn drop_warehouse_cluster(&self, name: String, cluster: String) -> Result<()>; - async fn add_warehouse_cluster_node( + async fn assign_warehouse_nodes( &self, name: String, - cluster: String, - nodes: Vec, + nodes: HashMap>, ) -> Result<()>; - async fn drop_warehouse_cluster_node( + async fn unassign_warehouse_nodes( &self, name: String, - cluster: String, - nodes: Vec, + nodes: HashMap>, ) -> Result<()>; } @@ -115,20 +115,18 @@ impl ResourcesManagement for DummyResourcesManagement { Err(ErrorCode::Unimplemented("The use of this feature requires a Databend Enterprise Edition license. To unlock enterprise features, please contact Databend to obtain a license. Learn more at https://docs.databend.com/guides/overview/editions/dee/")) } - async fn add_warehouse_cluster_node( + async fn assign_warehouse_nodes( &self, _: String, - _: String, - _: Vec, + _: HashMap>, ) -> Result<()> { Err(ErrorCode::Unimplemented("The use of this feature requires a Databend Enterprise Edition license. To unlock enterprise features, please contact Databend to obtain a license. Learn more at https://docs.databend.com/guides/overview/editions/dee/")) } - async fn drop_warehouse_cluster_node( + async fn unassign_warehouse_nodes( &self, _: String, - _: String, - _: Vec, + _: HashMap>, ) -> Result<()> { Err(ErrorCode::Unimplemented("The use of this feature requires a Databend Enterprise Edition license. To unlock enterprise features, please contact Databend to obtain a license. Learn more at https://docs.databend.com/guides/overview/editions/dee/")) } diff --git a/src/query/management/src/warehouse/warehouse_api.rs b/src/query/management/src/warehouse/warehouse_api.rs index 7c72e83a020db..a0b388a5d9f96 100644 --- a/src/query/management/src/warehouse/warehouse_api.rs +++ b/src/query/management/src/warehouse/warehouse_api.rs @@ -85,18 +85,16 @@ pub trait WarehouseApi: Sync + Send { to: String, ) -> Result<()>; - async fn add_warehouse_cluster_node( + async fn assign_warehouse_nodes( &self, - warehouse: &str, - cluster: &str, - add_nodes: SelectedNodes, + name: String, + nodes: HashMap, ) -> Result<()>; - async fn drop_warehouse_cluster_node( + async fn unassign_warehouse_nodes( &self, warehouse: &str, - cluster: &str, - drop_nodes: Vec, + nodes: HashMap, ) -> Result<()>; /// Get the tenant's cluster all nodes. diff --git a/src/query/management/src/warehouse/warehouse_mgr.rs b/src/query/management/src/warehouse/warehouse_mgr.rs index 61e12a517ab7f..6646743777562 100644 --- a/src/query/management/src/warehouse/warehouse_mgr.rs +++ b/src/query/management/src/warehouse/warehouse_mgr.rs @@ -12,12 +12,14 @@ // See the License for the specific language governing permissions and // limitations under the License. +use std::collections::hash_map::Entry; use std::collections::HashMap; use std::time::Duration; use databend_common_base::base::escape_for_key; use databend_common_base::base::unescape_for_key; use databend_common_base::base::GlobalUniqName; +use databend_common_base::vec_ext::VecExt; use databend_common_exception::ErrorCode; use databend_common_exception::Result; use databend_common_meta_kvapi::kvapi::KVApi; @@ -491,6 +493,137 @@ impl WarehouseMgr { "Warehouse operate conflict(tried 10 times).", )) } + + async fn unassigned_nodes(&self) -> Result, Vec<(u64, NodeInfo)>>> { + let online_nodes = self.metastore.prefix_list_kv(&self.node_key_prefix).await?; + let mut group_nodes = HashMap::with_capacity(online_nodes.len()); + + for (_, seq_data) in online_nodes { + let node_info = serde_json::from_slice::(&seq_data.data)?; + + if node_info.cluster_id.is_empty() && node_info.warehouse_id.is_empty() { + match group_nodes.entry(node_info.resource_group.clone()) { + Entry::Vacant(v) => { + v.insert(vec![(seq_data.seq, node_info)]); + } + Entry::Occupied(mut v) => { + v.get_mut().push((seq_data.seq, node_info)); + } + } + } + } + + Ok(group_nodes) + } + + async fn pick_assign_warehouse_node( + &self, + warehouse: &str, + nodes: &HashMap, + ) -> Result>> { + let mut selected_nodes = HashMap::with_capacity(nodes.len()); + + let mut grouped_nodes = self.unassigned_nodes().await?; + + let mut after_assign_node = HashMap::new(); + + for (cluster, cluster_node_selector) in nodes { + let mut cluster_selected_nodes = Vec::with_capacity(cluster_node_selector.len()); + for node_selector in cluster_node_selector { + match node_selector { + SelectedNode::Random(None) => { + let Some(nodes_list) = grouped_nodes.get_mut(&None) else { + match after_assign_node.entry(cluster.clone()) { + Entry::Vacant(v) => { + v.insert(1); + } + Entry::Occupied(mut v) => { + *v.get_mut() += 1; + } + }; + + break; + }; + + let Some((seq, mut node)) = nodes_list.pop() else { + grouped_nodes.remove(&None); + match after_assign_node.entry(cluster.clone()) { + Entry::Vacant(v) => { + v.insert(1); + } + Entry::Occupied(mut v) => { + *v.get_mut() += 1; + } + }; + + break; + }; + + node.runtime_resource_group = None; + node.cluster_id = cluster.clone(); + node.warehouse_id = warehouse.to_string(); + cluster_selected_nodes.push((seq, node)); + } + SelectedNode::Random(Some(resource_group)) => { + let key = Some(resource_group.clone()); + let Some(nodes_list) = grouped_nodes.get_mut(&key) else { + return Err(ErrorCode::NoResourcesAvailable(format!( + "Failed to create warehouse, reason: no resources available for {} group", + resource_group + ))); + }; + + let Some((seq, mut node)) = nodes_list.pop() else { + grouped_nodes.remove(&key); + return Err(ErrorCode::NoResourcesAvailable(format!( + "Failed to create warehouse, reason: no resources available for {} group", + resource_group + ))); + }; + + node.cluster_id = cluster.clone(); + node.warehouse_id = warehouse.to_string(); + node.runtime_resource_group = Some(resource_group.clone()); + cluster_selected_nodes.push((seq, node)); + } + } + } + + selected_nodes.insert(cluster.clone(), cluster_selected_nodes); + } + + if !after_assign_node.is_empty() { + let mut remain_nodes = Vec::new(); + + let mut processed_data = true; + while processed_data { + processed_data = false; + for nodes in grouped_nodes.values_mut() { + if let Some((seq, node)) = nodes.pop() { + processed_data = true; + remain_nodes.push((seq, node)); + } + } + } + + for (cluster, remain_node) in after_assign_node { + for _idx in 0..remain_node { + let Some((seq, mut node)) = remain_nodes.pop() else { + return Err(ErrorCode::NoResourcesAvailable( + "Failed to create warehouse, reason: no resources available.", + )); + }; + + node.cluster_id = cluster.clone(); + node.warehouse_id = warehouse.to_string(); + node.runtime_resource_group = None; + selected_nodes.get_mut(&cluster).unwrap().push((seq, node)); + } + } + } + + Ok(selected_nodes) + } } #[async_trait::async_trait] @@ -652,9 +785,7 @@ impl WarehouseApi for WarehouseMgr { )) } - async fn create_warehouse(&self, warehouse: String, nodes: Vec) -> Result<()> { - assert!(nodes.iter().all(|x| matches!(x, SelectedNode::Random(_)))); - + async fn create_warehouse(&self, warehouse: String, nodes: SelectedNodes) -> Result<()> { if warehouse.is_empty() { return Err(ErrorCode::InvalidWarehouse("Warehouse name is empty.")); } @@ -665,43 +796,13 @@ impl WarehouseApi for WarehouseMgr { )); } - loop { - let mut selected_nodes = Vec::with_capacity(nodes.len()); - - // get online nodes - let online_nodes = self.metastore.prefix_list_kv(&self.node_key_prefix).await?; - - let mut select_queue = nodes.clone(); - for (_, v) in &online_nodes { - match select_queue.last() { - None => { - break; - } - Some(SelectedNode::Random(Some(_))) => { - return Err(ErrorCode::Unimplemented( - "Custom instance types are not supported.", - )); - } - Some(SelectedNode::Random(None)) => { - // select random node + let nodes_map = HashMap::from([(String::from("default"), nodes.clone())]); - let mut node_info = serde_json::from_slice::(&v.data)?; - - if node_info.warehouse_id.is_empty() && node_info.cluster_id.is_empty() { - node_info.warehouse_id = warehouse.clone(); - node_info.cluster_id = String::from("default"); - selected_nodes.push((v.seq, node_info)); - select_queue.pop(); - } - } - } - } - - if !select_queue.is_empty() { - return Err(ErrorCode::NoResourcesAvailable( - "Failed to create warehouse, reason: no resources available", - )); - } + loop { + let mut selected_nodes = self + .pick_assign_warehouse_node(&warehouse, &nodes_map) + .await?; + let selected_nodes = selected_nodes.remove("default").unwrap(); let mut txn = TxnRequest::default(); @@ -1132,43 +1233,13 @@ impl WarehouseApi for WarehouseMgr { )); } - for _idx in 0..10 { - let mut selected_nodes = Vec::with_capacity(nodes.len()); + let nodes_map = HashMap::from([(cluster.clone(), nodes.clone())]); - // get online nodes - let online_nodes = self.metastore.prefix_list_kv(&self.node_key_prefix).await?; - - let mut select_queue = nodes.clone(); - for (_, v) in &online_nodes { - match select_queue.last() { - None => { - break; - } - Some(SelectedNode::Random(Some(_))) => { - return Err(ErrorCode::Unimplemented( - "Custom instance types are not supported.", - )); - } - Some(SelectedNode::Random(None)) => { - // select random node - - let mut node_info = serde_json::from_slice::(&v.data)?; - - if node_info.warehouse_id.is_empty() && node_info.cluster_id.is_empty() { - node_info.warehouse_id = warehouse.clone(); - node_info.cluster_id = cluster.clone(); - selected_nodes.push((v.seq, node_info)); - select_queue.pop(); - } - } - } - } - - if !select_queue.is_empty() { - return Err(ErrorCode::NoResourcesAvailable( - "Failed to create warehouse cluster, reason: no resources available", - )); - } + for _idx in 0..10 { + let mut selected_nodes = self + .pick_assign_warehouse_node(&warehouse, &nodes_map) + .await?; + let selected_nodes = selected_nodes.remove(&cluster).unwrap(); let mut create_cluster_txn = TxnRequest::default(); @@ -1485,96 +1556,58 @@ impl WarehouseApi for WarehouseMgr { )) } - async fn add_warehouse_cluster_node( + async fn assign_warehouse_nodes( &self, - warehouse: &str, - cluster: &str, - nodes: SelectedNodes, + warehouse: String, + nodes: HashMap, ) -> Result<()> { if warehouse.is_empty() { return Err(ErrorCode::InvalidWarehouse("Warehouse name is empty.")); } - if cluster.is_empty() { - return Err(ErrorCode::InvalidWarehouse( - "Warehouse cluster name is empty.", - )); - } - if nodes.is_empty() { return Err(ErrorCode::EmptyNodesForWarehouse( - "Cannot create warehouse cluster with empty nodes.", + "Cannot assign warehouse nodes with empty nodes list.", )); } for _idx in 0..10 { - let mut selected_nodes = Vec::with_capacity(nodes.len()); - - // get online nodes - let online_nodes = self.metastore.prefix_list_kv(&self.node_key_prefix).await?; - - let mut select_queue = nodes.clone(); - for (_, v) in &online_nodes { - match select_queue.last() { - None => { - break; - } - Some(SelectedNode::Random(Some(_))) => { - return Err(ErrorCode::Unimplemented( - "Custom instance types are not supported.", - )); - } - Some(SelectedNode::Random(None)) => { - // select random node - - let mut node_info = serde_json::from_slice::(&v.data)?; - - if node_info.warehouse_id.is_empty() && node_info.cluster_id.is_empty() { - node_info.cluster_id = cluster.to_string(); - node_info.warehouse_id = warehouse.to_string(); - selected_nodes.push((v.seq, node_info)); - select_queue.pop(); - } - } - } - } - - if !select_queue.is_empty() { - return Err(ErrorCode::NoResourcesAvailable( - "Failed to add warehouse cluster node, reason: no resources available", - )); - } + let selected_nodes = self.pick_assign_warehouse_node(&warehouse, &nodes).await?; let mut add_cluster_node_txn = TxnRequest::default(); - let mut consistent_info = self.consistent_warehouse_info(warehouse).await?; + let mut consistent_info = self.consistent_warehouse_info(&warehouse).await?; consistent_info.warehouse_info = match consistent_info.warehouse_info { WarehouseInfo::SelfManaged(_) => Err(ErrorCode::InvalidWarehouse(format!( - "Cannot add cluster node for warehouse {:?}, because it's self-managed warehouse.", + "Cannot assign nodes for warehouse {:?}, because it's self-managed warehouse.", warehouse ))), - WarehouseInfo::SystemManaged(mut info) => match info.clusters.get_mut(cluster) { - None => Err(ErrorCode::WarehouseClusterNotExists(format!( - "Warehouse cluster {:?}.{:?} not exists", - warehouse, cluster - ))), - Some(cluster_info) => { + WarehouseInfo::SystemManaged(mut info) => { + for (cluster, nodes) in &nodes { + let Some(cluster_info) = info.clusters.get_mut(cluster) else { + return Err(ErrorCode::WarehouseClusterNotExists(format!( + "Warehouse cluster {:?}.{:?} not exists", + warehouse, cluster + ))); + }; + cluster_info.nodes.extend(nodes.clone()); - Ok(WarehouseInfo::SystemManaged(SystemManagedWarehouse { - id: info.id, - status: info.status, - display_name: info.display_name, - clusters: info.clusters, - })) } - }, + + Ok(WarehouseInfo::SystemManaged(SystemManagedWarehouse { + id: info.id, + status: info.status, + display_name: info.display_name, + clusters: info.clusters, + })) + } }?; let warehouse_key = format!( "{}/{}", self.warehouse_key_prefix, - escape_for_key(warehouse)? + escape_for_key(&warehouse)? ); add_cluster_node_txn.condition.push(map_condition( @@ -1602,29 +1635,31 @@ impl WarehouseApi for WarehouseMgr { )); } - for (seq, mut node) in selected_nodes { - let node_key = self.node_key(&node)?; - let cluster_key = self.cluster_key(&node)?; + for selected_nodes in selected_nodes.into_values() { + for (seq, mut node) in selected_nodes { + let node_key = self.node_key(&node)?; + let cluster_key = self.cluster_key(&node)?; - add_cluster_node_txn - .condition - .push(map_condition(&node_key, MatchSeq::Exact(seq))); - add_cluster_node_txn.if_then.push(TxnOp::put_with_ttl( - node_key, - serde_json::to_vec(&node)?, - Some(self.lift_time * 4), - )); + add_cluster_node_txn + .condition + .push(map_condition(&node_key, MatchSeq::Exact(seq))); + add_cluster_node_txn.if_then.push(TxnOp::put_with_ttl( + node_key, + serde_json::to_vec(&node)?, + Some(self.lift_time * 4), + )); - node.cluster_id = String::new(); - node.warehouse_id = String::new(); - add_cluster_node_txn - .condition - .push(map_condition(&cluster_key, MatchSeq::Exact(0))); - add_cluster_node_txn.if_then.push(TxnOp::put_with_ttl( - cluster_key, - serde_json::to_vec(&node)?, - Some(self.lift_time * 4), - )); + node.cluster_id = String::new(); + node.warehouse_id = String::new(); + add_cluster_node_txn + .condition + .push(map_condition(&cluster_key, MatchSeq::Exact(0))); + add_cluster_node_txn.if_then.push(TxnOp::put_with_ttl( + cluster_key, + serde_json::to_vec(&node)?, + Some(self.lift_time * 4), + )); + } } return match self.metastore.transaction(add_cluster_node_txn).await? { @@ -1636,74 +1671,64 @@ impl WarehouseApi for WarehouseMgr { } Err(ErrorCode::WarehouseOperateConflict( - "Warehouse operate conflict(tried 10 times while in add warehouse cluster node).", + "Warehouse operate conflict(tried 10 times while in assign warehouse nodes).", )) } - async fn drop_warehouse_cluster_node( + async fn unassign_warehouse_nodes( &self, warehouse: &str, - cluster: &str, - nodes: Vec, + nodes: HashMap, ) -> Result<()> { if warehouse.is_empty() { return Err(ErrorCode::InvalidWarehouse("Warehouse name is empty.")); } - if cluster.is_empty() { - return Err(ErrorCode::InvalidWarehouse( - "Warehouse cluster name is empty.", - )); - } - for _idx in 0..10 { + let mut nodes = nodes.clone(); let mut drop_cluster_node_txn = TxnRequest::default(); let mut consistent_info = self.consistent_warehouse_info(warehouse).await?; consistent_info.warehouse_info = match consistent_info.warehouse_info { WarehouseInfo::SelfManaged(_) => Err(ErrorCode::InvalidWarehouse(format!( - "Cannot add cluster for warehouse {:?}, because it's self-managed warehouse.", + "Cannot unassign nodes for warehouse {:?}, because it's self-managed warehouse.", warehouse ))), - WarehouseInfo::SystemManaged(mut info) => match info.clusters.get_mut(cluster) { - None => Err(ErrorCode::WarehouseClusterNotExists(format!( - "Warehouse cluster {:?}.{:?} not exists", - warehouse, cluster - ))), - Some(cluster) => match nodes.len() == cluster.nodes.len() { - true => Err(ErrorCode::EmptyNodesForWarehouse(format!( - "Warehouse cluster {:?}.{:?} only has {} nodes, cannot drop all.", - warehouse, - cluster, - nodes.len() - ))), - false => { - for remove_node in &nodes { - if consistent_info - .consistent_nodes - .iter() - .any(|x| &x.node_info.id == remove_node) - { - cluster.nodes.pop(); - continue; - } + WarehouseInfo::SystemManaged(mut info) => { + for (cluster, nodes) in &nodes { + let Some(cluster) = info.clusters.get_mut(cluster) else { + return Err(ErrorCode::WarehouseClusterNotExists(format!( + "Warehouse cluster {:?}.{:?} not exists", + warehouse, cluster + ))); + }; + + if cluster.nodes.len() == nodes.len() { + return Err(ErrorCode::EmptyNodesForWarehouse(format!( + "Cannot unassign all nodes for warehouse cluster {:?}.{:?}", + warehouse, cluster + ))); + } + for remove_node in nodes { + if cluster.nodes.remove_first(remove_node).is_none() { return Err(ErrorCode::ClusterUnknownNode(format!( "Warehouse cluster {:?}.{:?} unknown node {:?}", warehouse, cluster, remove_node ))); } - - Ok(WarehouseInfo::SystemManaged(SystemManagedWarehouse { - id: info.id, - status: info.status, - display_name: info.display_name, - clusters: info.clusters, - })) } - }, - }, + } + + + Ok(WarehouseInfo::SystemManaged(SystemManagedWarehouse { + id: info.id, + status: info.status, + display_name: info.display_name, + clusters: info.clusters, + })) + } }?; let warehouse_key = format!( @@ -1736,32 +1761,139 @@ impl WarehouseApi for WarehouseMgr { MatchSeq::Exact(consistent_node.cluster_seq), )); - if nodes.contains(&consistent_node.node_info.id) { - // Remove node - consistent_node.node_info.cluster_id = String::new(); - consistent_node.node_info.warehouse_id = String::new(); - - drop_cluster_node_txn - .if_then - .push(TxnOp::delete(cluster_key)); - drop_cluster_node_txn.if_then.push(TxnOp::put_with_ttl( - node_key, - serde_json::to_vec(&consistent_node.node_info)?, - Some(self.lift_time * 4), - )) + if let Some(v) = nodes.get_mut(&consistent_node.node_info.cluster_id) { + if let Some(remove_node) = v.pop() { + let SelectedNode::Random(resource_group) = remove_node; + if consistent_node.node_info.runtime_resource_group == resource_group { + consistent_node.node_info.cluster_id = String::new(); + consistent_node.node_info.warehouse_id = String::new(); + consistent_node.node_info.runtime_resource_group = None; + + drop_cluster_node_txn + .if_then + .push(TxnOp::delete(cluster_key)); + drop_cluster_node_txn.if_then.push(TxnOp::put_with_ttl( + node_key, + serde_json::to_vec(&consistent_node.node_info)?, + Some(self.lift_time * 4), + )) + } + } } } - - return match self.metastore.transaction(drop_cluster_node_txn).await? { - res if res.success => Ok(()), - _ => { - continue; - } - }; } + // // if cluster.is_empty() { + // // return Err(ErrorCode::InvalidWarehouse( + // // "Warehouse cluster name is empty.", + // // )); + // // } + // + // for _idx in 0..10 { + // let mut drop_cluster_node_txn = TxnRequest::default(); + // + // let mut consistent_info = self.consistent_warehouse_info(warehouse).await?; + // + // consistent_info.warehouse_info = match consistent_info.warehouse_info { + // WarehouseInfo::SelfManaged(_) => Err(ErrorCode::InvalidWarehouse(format!( + // "Cannot add cluster for warehouse {:?}, because it's self-managed warehouse.", + // warehouse + // ))), + // WarehouseInfo::SystemManaged(mut info) => match info.clusters.get_mut(cluster) { + // None => Err(ErrorCode::WarehouseClusterNotExists(format!( + // "Warehouse cluster {:?}.{:?} not exists", + // warehouse, cluster + // ))), + // Some(cluster) => match nodes.len() == cluster.nodes.len() { + // true => Err(ErrorCode::EmptyNodesForWarehouse(format!( + // "Warehouse cluster {:?}.{:?} only has {} nodes, cannot drop all.", + // warehouse, + // cluster, + // nodes.len() + // ))), + // false => { + // for remove_node in &nodes { + // if consistent_info + // .consistent_nodes + // .iter() + // .any(|x| &x.node_info.id == remove_node) + // { + // cluster.nodes.pop(); + // continue; + // } + // + // return Err(ErrorCode::ClusterUnknownNode(format!( + // "Warehouse cluster {:?}.{:?} unknown node {:?}", + // warehouse, cluster, remove_node + // ))); + // } + // + // Ok(WarehouseInfo::SystemManaged(SystemManagedWarehouse { + // id: info.id, + // status: info.status, + // display_name: info.display_name, + // clusters: info.clusters, + // })) + // } + // }, + // }, + // }?; + // + // let warehouse_key = format!( + // "{}/{}", + // self.warehouse_key_prefix, + // escape_for_key(warehouse)? + // ); + // + // drop_cluster_node_txn.condition.push(map_condition( + // &warehouse_key, + // MatchSeq::Exact(consistent_info.info_seq), + // )); + // + // drop_cluster_node_txn.if_then.push(TxnOp::put( + // warehouse_key.clone(), + // serde_json::to_vec(&consistent_info.warehouse_info)?, + // )); + // + // // lock all cluster state + // for mut consistent_node in consistent_info.consistent_nodes { + // let node_key = self.node_key(&consistent_node.node_info)?; + // let cluster_key = self.cluster_key(&consistent_node.node_info)?; + // + // drop_cluster_node_txn.condition.push(map_condition( + // &node_key, + // MatchSeq::Exact(consistent_node.node_seq), + // )); + // drop_cluster_node_txn.condition.push(map_condition( + // &cluster_key, + // MatchSeq::Exact(consistent_node.cluster_seq), + // )); + // + // if nodes.contains(&consistent_node.node_info.id) { + // // Remove node + // consistent_node.node_info.cluster_id = String::new(); + // consistent_node.node_info.warehouse_id = String::new(); + // + // drop_cluster_node_txn + // .if_then + // .push(TxnOp::delete(cluster_key)); + // drop_cluster_node_txn.if_then.push(TxnOp::put_with_ttl( + // node_key, + // serde_json::to_vec(&consistent_node.node_info)?, + // Some(self.lift_time * 4), + // )) + // } + // } + // + // return match self.metastore.transaction(drop_cluster_node_txn).await? { + // res if res.success => Ok(()), + // _ => { + // continue; + // } + // }; + // } Err(ErrorCode::WarehouseOperateConflict( - "Warehouse operate conflict(tried 10 times while in drop warehouse cluster node).", + "Warehouse operate conflict(tried 10 times while in unassign warehouse nodes).", )) } diff --git a/src/query/management/tests/it/warehouse.rs b/src/query/management/tests/it/warehouse.rs index f9ef582bd3444..4354072bf458e 100644 --- a/src/query/management/tests/it/warehouse.rs +++ b/src/query/management/tests/it/warehouse.rs @@ -885,8 +885,10 @@ fn system_managed_node(id: &str) -> NodeInfo { discovery_address: "".to_string(), binary_version: "".to_string(), node_type: NodeType::SystemManaged, + resource_group: None, cluster_id: "".to_string(), warehouse_id: "".to_string(), + runtime_resource_group: None, } } @@ -901,8 +903,10 @@ fn self_managed_node(node_id: &str) -> NodeInfo { discovery_address: "ip2:port".to_string(), binary_version: "binary_version".to_string(), node_type: NodeType::SelfManaged, + resource_group: None, cluster_id: "test-cluster-id".to_string(), warehouse_id: "test-cluster-id".to_string(), + runtime_resource_group: None, } } diff --git a/src/query/service/src/interpreters/access/privilege_access.rs b/src/query/service/src/interpreters/access/privilege_access.rs index ac0da6c4b3a39..490aa0e1c5912 100644 --- a/src/query/service/src/interpreters/access/privilege_access.rs +++ b/src/query/service/src/interpreters/access/privilege_access.rs @@ -1265,8 +1265,8 @@ impl AccessChecker for PrivilegeAccess { Plan::RenameWarehouseCluster(_) => {} Plan::CreateWarehouse(_) => {} Plan::AddWarehouseCluster(_) => {} - Plan::AddWarehouseClusterNode(_) => {} - Plan::DropWarehouseClusterNode(_) => {} + Plan::AssignWarehouseNodes(_) => {} + Plan::UnassignWarehouseNodes(_) => {} } Ok(()) diff --git a/src/query/service/src/interpreters/interpreter_add_warehouse_cluster.rs b/src/query/service/src/interpreters/interpreter_add_warehouse_cluster.rs index 54958147491a5..da0139a8314c3 100644 --- a/src/query/service/src/interpreters/interpreter_add_warehouse_cluster.rs +++ b/src/query/service/src/interpreters/interpreter_add_warehouse_cluster.rs @@ -14,8 +14,12 @@ use std::sync::Arc; +use databend_common_base::base::GlobalInstance; +use databend_common_exception::ErrorCode; use databend_common_exception::Result; +use databend_common_management::SelectedNode; use databend_common_sql::plans::AddWarehouseClusterPlan; +use databend_enterprise_resources_management::ResourcesManagement; use crate::interpreters::Interpreter; use crate::pipelines::PipelineBuildResult; @@ -24,7 +28,6 @@ use crate::sessions::QueryContext; pub struct AddWarehouseClusterInterpreter { #[allow(dead_code)] ctx: Arc, - #[allow(dead_code)] plan: AddWarehouseClusterPlan, } @@ -46,9 +49,45 @@ impl Interpreter for AddWarehouseClusterInterpreter { #[async_backtrace::framed] async fn execute2(&self) -> Result { - // GlobalInstance::get::>() - // .drop_warehouse_cluster(self.plan.warehouse.clone(), self.plan.cluster.clone()) - // .await?; + if let Some(cluster_size) = self.plan.options.get("CLUSTER_SIZE") { + if !self.plan.nodes.is_empty() { + return Err(ErrorCode::InvalidArgument(format!(""))); + } + + let Ok(cluster_size) = cluster_size.parse::() else { + return Err(ErrorCode::InvalidArgument("")); + }; + + GlobalInstance::get::>() + .add_warehouse_cluster( + self.plan.warehouse.clone(), + self.plan.cluster.clone(), + vec![SelectedNode::Random(None); cluster_size], + ) + .await?; + + return Ok(PipelineBuildResult::create()); + } + + if self.plan.nodes.is_empty() { + return Err(ErrorCode::InvalidArgument("")); + } + + let mut selected_nodes = Vec::with_capacity(self.plan.nodes.len()); + for (group, nodes) in &self.plan.nodes { + for _ in 0..*nodes { + selected_nodes.push(SelectedNode::Random(group.clone())); + } + } + + GlobalInstance::get::>() + .add_warehouse_cluster( + self.plan.warehouse.clone(), + self.plan.cluster.clone(), + selected_nodes, + ) + .await?; + Ok(PipelineBuildResult::create()) } } diff --git a/src/query/service/src/interpreters/interpreter_add_warehouse_cluster_node.rs b/src/query/service/src/interpreters/interpreter_add_warehouse_cluster_node.rs deleted file mode 100644 index 1566d7188ae04..0000000000000 --- a/src/query/service/src/interpreters/interpreter_add_warehouse_cluster_node.rs +++ /dev/null @@ -1,49 +0,0 @@ -// Copyright 2021 Datafuse Labs -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -use databend_common_exception::Result; -use databend_common_sql::plans::AddWarehouseClusterNodePlan; - -use crate::interpreters::Interpreter; -use crate::pipelines::PipelineBuildResult; - -pub struct AddWarehouseClusterNodeInterpreter { - #[allow(dead_code)] - plan: AddWarehouseClusterNodePlan, -} - -impl AddWarehouseClusterNodeInterpreter { - pub fn try_create(plan: AddWarehouseClusterNodePlan) -> Result { - Ok(AddWarehouseClusterNodeInterpreter { plan }) - } -} - -#[async_trait::async_trait] -impl Interpreter for AddWarehouseClusterNodeInterpreter { - fn name(&self) -> &str { - "AddWarehouseClusterNodeInterpreter" - } - - fn is_ddl(&self) -> bool { - true - } - - #[async_backtrace::framed] - async fn execute2(&self) -> Result { - // GlobalInstance::get::>() - // .drop_warehouse_cluster(self.plan.warehouse.clone(), self.plan.cluster.clone()) - // .await?; - Ok(PipelineBuildResult::create()) - } -} diff --git a/src/query/service/src/interpreters/interpreter_assign_warehouse_nodes.rs b/src/query/service/src/interpreters/interpreter_assign_warehouse_nodes.rs new file mode 100644 index 0000000000000..d23013aaa4179 --- /dev/null +++ b/src/query/service/src/interpreters/interpreter_assign_warehouse_nodes.rs @@ -0,0 +1,66 @@ +// Copyright 2021 Datafuse Labs +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::collections::HashMap; +use std::sync::Arc; + +use databend_common_base::base::GlobalInstance; +use databend_common_exception::Result; +use databend_common_management::SelectedNode; +use databend_common_sql::plans::AssignWarehouseNodesPlan; +use databend_enterprise_resources_management::ResourcesManagement; + +use crate::interpreters::Interpreter; +use crate::pipelines::PipelineBuildResult; + +pub struct AssignWarehouseNodesInterpreter { + plan: AssignWarehouseNodesPlan, +} + +impl AssignWarehouseNodesInterpreter { + pub fn try_create(plan: AssignWarehouseNodesPlan) -> Result { + Ok(AssignWarehouseNodesInterpreter { plan }) + } +} + +#[async_trait::async_trait] +impl Interpreter for AssignWarehouseNodesInterpreter { + fn name(&self) -> &str { + "AddWarehouseClusterNodeInterpreter" + } + + fn is_ddl(&self) -> bool { + true + } + + #[async_backtrace::framed] + async fn execute2(&self) -> Result { + let mut cluster_selected_nodes = HashMap::with_capacity(self.plan.assign_clusters.len()); + for (cluster, nodes_map) in &self.plan.assign_clusters { + let mut selected_nodes = Vec::with_capacity(nodes_map.len()); + for (group, nodes) in nodes_map { + for _ in 0..*nodes { + selected_nodes.push(SelectedNode::Random(group.clone())); + } + } + + cluster_selected_nodes.insert(cluster.clone(), selected_nodes); + } + + GlobalInstance::get::>() + .assign_warehouse_nodes(self.plan.warehouse.clone(), cluster_selected_nodes) + .await?; + Ok(PipelineBuildResult::create()) + } +} diff --git a/src/query/service/src/interpreters/interpreter_create_warehouses.rs b/src/query/service/src/interpreters/interpreter_create_warehouses.rs index e801b3b02ae05..a09c908d507bd 100644 --- a/src/query/service/src/interpreters/interpreter_create_warehouses.rs +++ b/src/query/service/src/interpreters/interpreter_create_warehouses.rs @@ -73,9 +73,9 @@ impl Interpreter for CreateWarehouseInterpreter { } let mut selected_nodes = Vec::with_capacity(self.plan.nodes.len()); - for (group, nodes) in self.plan.nodes { - for _ in 0..nodes { - selected_nodes.push(SelectedNode::Random(group)); + for (group, nodes) in &self.plan.nodes { + for _ in 0..*nodes { + selected_nodes.push(SelectedNode::Random(group.clone())); } } diff --git a/src/query/service/src/interpreters/interpreter_drop_warehouse_cluster_node.rs b/src/query/service/src/interpreters/interpreter_drop_warehouse_cluster_node.rs index e12120e8d8b5b..35d20accdc4c7 100644 --- a/src/query/service/src/interpreters/interpreter_drop_warehouse_cluster_node.rs +++ b/src/query/service/src/interpreters/interpreter_drop_warehouse_cluster_node.rs @@ -13,18 +13,18 @@ // limitations under the License. use databend_common_exception::Result; -use databend_common_sql::plans::DropWarehouseClusterNodePlan; +use databend_common_sql::plans::UnassignWarehouseNodesPlan; use crate::interpreters::Interpreter; use crate::pipelines::PipelineBuildResult; pub struct DropWarehouseClusterNodeInterpreter { #[allow(dead_code)] - plan: DropWarehouseClusterNodePlan, + plan: UnassignWarehouseNodesPlan, } impl DropWarehouseClusterNodeInterpreter { - pub fn try_create(plan: DropWarehouseClusterNodePlan) -> Result { + pub fn try_create(plan: UnassignWarehouseNodesPlan) -> Result { Ok(DropWarehouseClusterNodeInterpreter { plan }) } } diff --git a/src/query/service/src/interpreters/interpreter_factory.rs b/src/query/service/src/interpreters/interpreter_factory.rs index af382e26faf6e..f5581a8595788 100644 --- a/src/query/service/src/interpreters/interpreter_factory.rs +++ b/src/query/service/src/interpreters/interpreter_factory.rs @@ -37,7 +37,7 @@ use super::interpreter_user_stage_drop::DropUserStageInterpreter; use super::*; use crate::interpreters::access::Accessor; use crate::interpreters::interpreter_add_warehouse_cluster::AddWarehouseClusterInterpreter; -use crate::interpreters::interpreter_add_warehouse_cluster_node::AddWarehouseClusterNodeInterpreter; +use crate::interpreters::interpreter_assign_warehouse_nodes::AssignWarehouseNodesInterpreter; use crate::interpreters::interpreter_catalog_drop::DropCatalogInterpreter; use crate::interpreters::interpreter_connection_create::CreateConnectionInterpreter; use crate::interpreters::interpreter_connection_desc::DescConnectionInterpreter; @@ -674,10 +674,10 @@ impl InterpreterFactory { Plan::RenameWarehouseCluster(v) => Ok(Arc::new( RenameWarehouseClusterInterpreter::try_create(*v.clone())?, )), - Plan::AddWarehouseClusterNode(v) => Ok(Arc::new( - AddWarehouseClusterNodeInterpreter::try_create(*v.clone())?, + Plan::AssignWarehouseNodes(v) => Ok(Arc::new( + AssignWarehouseNodesInterpreter::try_create(*v.clone())?, )), - Plan::DropWarehouseClusterNode(v) => Ok(Arc::new( + Plan::UnassignWarehouseNodes(v) => Ok(Arc::new( DropWarehouseClusterNodeInterpreter::try_create(*v.clone())?, )), } diff --git a/src/query/service/src/interpreters/mod.rs b/src/query/service/src/interpreters/mod.rs index ed4f2e1fdc31d..b7b9ad71650e9 100644 --- a/src/query/service/src/interpreters/mod.rs +++ b/src/query/service/src/interpreters/mod.rs @@ -17,7 +17,7 @@ pub(crate) mod common; mod hook; mod interpreter; mod interpreter_add_warehouse_cluster; -mod interpreter_add_warehouse_cluster_node; +mod interpreter_assign_warehouse_nodes; mod interpreter_catalog_create; mod interpreter_catalog_drop; mod interpreter_catalog_show_create; diff --git a/src/query/sql/src/planner/binder/binder.rs b/src/query/sql/src/planner/binder/binder.rs index 5cd5d3f3fc8bc..9b0515ec668e1 100644 --- a/src/query/sql/src/planner/binder/binder.rs +++ b/src/query/sql/src/planner/binder/binder.rs @@ -656,16 +656,16 @@ impl<'a> Binder { } Statement::ShowWarehouses(v) => self.bind_show_warehouses(v)?, Statement::DropWarehouse(v) => self.bind_drop_warehouse(v)?, - Statement::CreateWarehouse(_) => unimplemented!(), + Statement::CreateWarehouse(v) => self.bind_create_warehouse(v)?, Statement::RenameWarehouse(v) => self.bind_rename_warehouse(v)?, Statement::ResumeWarehouse(v) => self.bind_resume_warehouse(v)?, Statement::SuspendWarehouse(v) => self.bind_suspend_warehouse(v)?, Statement::InspectWarehouse(v) => self.bind_inspect_warehouse(v)?, - Statement::AddWarehouseCluster(_) => unimplemented!(), + Statement::AddWarehouseCluster(v) => self.bind_add_warehouse_cluster(v)?, Statement::DropWarehouseCluster(v) => self.bind_drop_warehouse_cluster(v)?, Statement::RenameWarehouseCluster(v) => self.bind_rename_warehouse_cluster(v)?, - Statement::AddWarehouseClusterNode(_) => unimplemented!(), - Statement::DropWarehouseClusterNode(_) => unimplemented!(), + Statement::AssignWarehouseNodes(v) => self.bind_assign_warehouse_nodes(v)?, + Statement::UnassignWarehouseNodes(v) => self.bind_unassign_warehouse_nodes(v)?, }; match plan.kind() { diff --git a/src/query/sql/src/planner/binder/ddl/warehouse.rs b/src/query/sql/src/planner/binder/ddl/warehouse.rs index b292e873154a5..feb6f09f5d543 100644 --- a/src/query/sql/src/planner/binder/ddl/warehouse.rs +++ b/src/query/sql/src/planner/binder/ddl/warehouse.rs @@ -15,6 +15,8 @@ use std::collections::hash_map::Entry; use std::collections::HashMap; +use databend_common_ast::ast::AddWarehouseClusterStmt; +use databend_common_ast::ast::AssignWarehouseNodesStmt; use databend_common_ast::ast::CreateWarehouseStmt; use databend_common_ast::ast::DropWarehouseClusterStmt; use databend_common_ast::ast::DropWarehouseStmt; @@ -24,8 +26,11 @@ use databend_common_ast::ast::RenameWarehouseStmt; use databend_common_ast::ast::ResumeWarehouseStmt; use databend_common_ast::ast::ShowWarehousesStmt; use databend_common_ast::ast::SuspendWarehouseStmt; +use databend_common_ast::ast::UnassignWarehouseNodesStmt; use databend_common_exception::Result; +use crate::plans::AddWarehouseClusterPlan; +use crate::plans::AssignWarehouseNodesPlan; use crate::plans::CreateWarehousePlan; use crate::plans::DropWarehouseClusterPlan; use crate::plans::DropWarehousePlan; @@ -35,6 +40,7 @@ use crate::plans::RenameWarehouseClusterPlan; use crate::plans::RenameWarehousePlan; use crate::plans::ResumeWarehousePlan; use crate::plans::SuspendWarehousePlan; +use crate::plans::UnassignWarehouseNodesPlan; use crate::Binder; impl Binder { @@ -51,21 +57,21 @@ impl Binder { ) -> Result { let mut nodes = HashMap::with_capacity(stmt.node_list.len()); - for (group, nodes_num) in stmt.node_list { - match nodes.entry(group) { - Entry::Vacant(mut v) => { - v.insert(nodes_num); + for (group, nodes_num) in &stmt.node_list { + match nodes.entry(group.clone()) { + Entry::Vacant(v) => { + v.insert(*nodes_num); } Entry::Occupied(mut v) => { - *v.get_mut() += nodes_num; + *v.get_mut() += *nodes_num; } } } Ok(Plan::CreateWarehouse(Box::new(CreateWarehousePlan { + nodes, warehouse: stmt.warehouse.to_string(), - nodes: Default::default(), - options: stmt.options, + options: stmt.options.clone(), }))) } @@ -115,6 +121,33 @@ impl Binder { }))) } + pub(in crate::planner::binder) fn bind_add_warehouse_cluster( + &mut self, + stmt: &AddWarehouseClusterStmt, + ) -> Result { + let mut nodes = HashMap::with_capacity(stmt.node_list.len()); + + for (group, nodes_num) in &stmt.node_list { + match nodes.entry(group.clone()) { + Entry::Vacant(v) => { + v.insert(*nodes_num); + } + Entry::Occupied(mut v) => { + *v.get_mut() += *nodes_num; + } + } + } + + Ok(Plan::AddWarehouseCluster(Box::new( + AddWarehouseClusterPlan { + nodes, + options: stmt.options.clone(), + cluster: stmt.cluster.to_string(), + warehouse: stmt.warehouse.to_string(), + }, + ))) + } + pub(in crate::planner::binder) fn bind_drop_warehouse_cluster( &mut self, stmt: &DropWarehouseClusterStmt, @@ -139,4 +172,62 @@ impl Binder { }, ))) } + + pub(in crate::planner::binder) fn bind_assign_warehouse_nodes( + &mut self, + stmt: &AssignWarehouseNodesStmt, + ) -> Result { + let mut assign_clusters = HashMap::with_capacity(stmt.node_list.len()); + for (cluster, group, num) in &stmt.node_list { + match assign_clusters.entry(cluster.to_string()) { + Entry::Vacant(v) => { + v.insert(HashMap::from([(group.clone(), *num as usize)])); + } + Entry::Occupied(mut v) => match v.get_mut().entry(group.clone()) { + Entry::Vacant(v) => { + v.insert(*num as usize); + } + Entry::Occupied(mut v) => { + *v.get_mut() += *num as usize; + } + }, + } + } + + Ok(Plan::AssignWarehouseNodes(Box::new( + AssignWarehouseNodesPlan { + assign_clusters, + warehouse: stmt.warehouse.to_string(), + }, + ))) + } + + pub(in crate::planner::binder) fn bind_unassign_warehouse_nodes( + &mut self, + stmt: &UnassignWarehouseNodesStmt, + ) -> Result { + let mut unassign_clusters = HashMap::with_capacity(stmt.node_list.len()); + for (cluster, group, num) in &stmt.node_list { + match unassign_clusters.entry(cluster.to_string()) { + Entry::Vacant(v) => { + v.insert(HashMap::from([(group.clone(), *num as usize)])); + } + Entry::Occupied(mut v) => match v.get_mut().entry(group.clone()) { + Entry::Vacant(v) => { + v.insert(*num as usize); + } + Entry::Occupied(mut v) => { + *v.get_mut() += *num as usize; + } + }, + } + } + + Ok(Plan::UnassignWarehouseNodes(Box::new( + UnassignWarehouseNodesPlan { + unassign_clusters, + warehouse: stmt.warehouse.to_string(), + }, + ))) + } } diff --git a/src/query/sql/src/planner/format/display_plan.rs b/src/query/sql/src/planner/format/display_plan.rs index 3877b75d89b5a..cb451fa86a4e9 100644 --- a/src/query/sql/src/planner/format/display_plan.rs +++ b/src/query/sql/src/planner/format/display_plan.rs @@ -223,8 +223,8 @@ impl Plan { Plan::RenameWarehouseCluster(_) => Ok("RenameWarehouseCluster".to_string()), Plan::CreateWarehouse(_) => Ok("CreateWarehouse".to_string()), Plan::AddWarehouseCluster(_) => Ok("AddWarehouseCluster".to_string()), - Plan::AddWarehouseClusterNode(_) => Ok("AddWarehouseClusterNode".to_string()), - Plan::DropWarehouseClusterNode(_) => Ok("DropWarehouseClusterNode".to_string()), + Plan::AssignWarehouseNodes(_) => Ok("AddWarehouseClusterNode".to_string()), + Plan::UnassignWarehouseNodes(_) => Ok("DropWarehouseClusterNode".to_string()), } } } diff --git a/src/query/sql/src/planner/plans/ddl/warehouse.rs b/src/query/sql/src/planner/plans/ddl/warehouse.rs index f1168f232b6ca..5e028643d8acd 100644 --- a/src/query/sql/src/planner/plans/ddl/warehouse.rs +++ b/src/query/sql/src/planner/plans/ddl/warehouse.rs @@ -65,9 +65,10 @@ impl InspectWarehousePlan { #[derive(Clone, Debug, PartialEq, Eq)] pub struct AddWarehouseClusterPlan { - warehouse: String, - cluster: String, - // nodes:Vec + pub warehouse: String, + pub cluster: String, + pub nodes: HashMap, u64>, + pub options: BTreeMap, } #[derive(Clone, Debug, PartialEq, Eq)] @@ -84,7 +85,13 @@ pub struct RenameWarehouseClusterPlan { } #[derive(Clone, Debug, PartialEq, Eq)] -pub struct AddWarehouseClusterNodePlan {} +pub struct AssignWarehouseNodesPlan { + pub warehouse: String, + pub assign_clusters: HashMap, usize>>, +} #[derive(Clone, Debug, PartialEq, Eq)] -pub struct DropWarehouseClusterNodePlan {} +pub struct UnassignWarehouseNodesPlan { + pub warehouse: String, + pub unassign_clusters: HashMap, usize>>, +} diff --git a/src/query/sql/src/planner/plans/plan.rs b/src/query/sql/src/planner/plans/plan.rs index 94b7a85e028ff..8ac8e47006767 100644 --- a/src/query/sql/src/planner/plans/plan.rs +++ b/src/query/sql/src/planner/plans/plan.rs @@ -32,7 +32,6 @@ use crate::binder::ExplainConfig; use crate::optimizer::SExpr; use crate::plans::copy_into_location::CopyIntoLocationPlan; use crate::plans::AddTableColumnPlan; -use crate::plans::AddWarehouseClusterNodePlan; use crate::plans::AddWarehouseClusterPlan; use crate::plans::AlterNetworkPolicyPlan; use crate::plans::AlterNotificationPlan; @@ -44,6 +43,7 @@ use crate::plans::AlterUserPlan; use crate::plans::AlterViewPlan; use crate::plans::AlterVirtualColumnPlan; use crate::plans::AnalyzeTablePlan; +use crate::plans::AssignWarehouseNodesPlan; use crate::plans::CallProcedurePlan; use crate::plans::CopyIntoTableMode; use crate::plans::CopyIntoTablePlan; @@ -102,7 +102,6 @@ use crate::plans::DropUDFPlan; use crate::plans::DropUserPlan; use crate::plans::DropViewPlan; use crate::plans::DropVirtualColumnPlan; -use crate::plans::DropWarehouseClusterNodePlan; use crate::plans::DropWarehouseClusterPlan; use crate::plans::DropWarehousePlan; use crate::plans::Exchange; @@ -151,6 +150,7 @@ use crate::plans::ShowTasksPlan; use crate::plans::SuspendWarehousePlan; use crate::plans::SystemPlan; use crate::plans::TruncateTablePlan; +use crate::plans::UnassignWarehouseNodesPlan; use crate::plans::UndropDatabasePlan; use crate::plans::UndropTablePlan; use crate::plans::UnsetOptionsPlan; @@ -213,8 +213,8 @@ pub enum Plan { AddWarehouseCluster(Box), DropWarehouseCluster(Box), RenameWarehouseCluster(Box), - AddWarehouseClusterNode(Box), - DropWarehouseClusterNode(Box), + AssignWarehouseNodes(Box), + UnassignWarehouseNodes(Box), // Databases ShowCreateDatabase(Box), From da6155e669f2e52921b231bab3419472a67db5a0 Mon Sep 17 00:00:00 2001 From: zhang2014 Date: Tue, 24 Dec 2024 18:29:15 +0800 Subject: [PATCH 21/53] feat(cluster): support custom management cluster --- src/query/management/tests/it/warehouse.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/query/management/tests/it/warehouse.rs b/src/query/management/tests/it/warehouse.rs index 4354072bf458e..020ba9c954081 100644 --- a/src/query/management/tests/it/warehouse.rs +++ b/src/query/management/tests/it/warehouse.rs @@ -262,7 +262,7 @@ async fn test_empty_system_managed_warehouse() -> Result<()> { String::from("XLargeNode"), ))]); - assert_eq!(create_warehouse.await.unwrap_err().code(), 1002); + assert_eq!(create_warehouse.await.unwrap_err().code(), 2404); Ok(()) } From 63c0b9eef71a11322c533cbcb97ae40fffac2540 Mon Sep 17 00:00:00 2001 From: zhang2014 Date: Wed, 25 Dec 2024 21:49:56 +0800 Subject: [PATCH 22/53] feat(cluster): support custom management cluster --- .../databend-query-node-system-managed.toml | 96 +++++++++++++ .../deploy/databend-query-system-managed.sh | 126 ++++++++++++++++++ src/query/config/src/config.rs | 6 +- src/query/config/src/inner.rs | 3 + src/query/ee/src/resource_management/mod.rs | 1 + .../resources_management_self_managed.rs | 112 ++++++++++++++++ .../resources_management_system.rs | 47 ++++--- .../management/src/warehouse/warehouse_api.rs | 2 + .../management/src/warehouse/warehouse_mgr.rs | 11 ++ src/query/service/src/clusters/cluster.rs | 23 ++-- 10 files changed, 396 insertions(+), 31 deletions(-) create mode 100644 scripts/ci/deploy/config/databend-query-node-system-managed.toml create mode 100755 scripts/ci/deploy/databend-query-system-managed.sh create mode 100644 src/query/ee/src/resource_management/resources_management_self_managed.rs diff --git a/scripts/ci/deploy/config/databend-query-node-system-managed.toml b/scripts/ci/deploy/config/databend-query-node-system-managed.toml new file mode 100644 index 0000000000000..07374d0f845ee --- /dev/null +++ b/scripts/ci/deploy/config/databend-query-node-system-managed.toml @@ -0,0 +1,96 @@ +# Usage: +# databend-query -c databend_query_config_spec.toml + +[query] +max_active_sessions = 256 +shutdown_wait_timeout_ms = 5000 + +# For flight rpc. +flight_api_address = "0.0.0.0:flight_port" + +# Databend Query http address. +# For admin RESET API. +admin_api_address = "0.0.0.0:admin_api_port" + +# Databend Query metrics RESET API. +metric_api_address = "0.0.0.0:metric_api_port" + +# Databend Query MySQL Handler. +mysql_handler_host = "0.0.0.0" +mysql_handler_port = mysql_port + +# Databend Query ClickHouse Handler. +clickhouse_http_handler_host = "0.0.0.0" +clickhouse_http_handler_port = clickhouse_port + +# Databend Query HTTP Handler. +http_handler_host = "0.0.0.0" +http_handler_port = http_port + +# Databend Query FlightSQL Handler. +flight_sql_handler_host = "0.0.0.0" +flight_sql_handler_port = flight_sql_port + +tenant_id = "test_tenant" +cluster_id = "test_cluster" + +table_engine_memory_enabled = true +default_storage_format = 'parquet' +default_compression = 'zstd' + +[[query.users]] +name = "root" +auth_type = "no_password" + +[[query.users]] +name = "default" +auth_type = "no_password" + +# This for test +[[query.udfs]] +name = "ping" +definition = "CREATE FUNCTION ping(STRING) RETURNS STRING LANGUAGE python HANDLER = 'ping' ADDRESS = 'http://0.0.0.0:8815'" + +[query.resources_management] +type = "system" + +[log] + +[log.file] +level = "INFO" +format = "text" +dir = "./.databend/query_logs" +prefix_filter = "" + +[meta] +# It is a list of `grpc_api_advertise_host:` of databend-meta config +endpoints = ["0.0.0.0:9191"] +username = "root" +password = "root" +client_timeout_in_second = 60 +auto_sync_interval = 60 + +# Storage config. +[storage] +# fs | s3 | azblob | obs | oss +type = "fs" + +# Set a local folder to store your data. +# Comment out this block if you're NOT using local file system as storage. +[storage.fs] +data_path = "./.databend/stateless_test_data" + +# Cache config. +[cache] +# Type of storage to keep the table data cache +# +# available options: [none|disk] +# default is "none", which disable table data cache +# use "disk" to enabled disk cache +data_cache_storage = "none" + +[cache.disk] +# cache path +path = "./.databend/_cache" +# max bytes of cached data 20G +max_bytes = 21474836480 diff --git a/scripts/ci/deploy/databend-query-system-managed.sh b/scripts/ci/deploy/databend-query-system-managed.sh new file mode 100755 index 0000000000000..946286d1fa97e --- /dev/null +++ b/scripts/ci/deploy/databend-query-system-managed.sh @@ -0,0 +1,126 @@ +#!/bin/bash +# Copyright 2022 The Databend Authors. +# SPDX-License-Identifier: Apache-2.0. + +set -e + +SCRIPT_PATH="$(cd "$(dirname "$0")" >/dev/null 2>&1 && pwd)" +cd "$SCRIPT_PATH/../../.." || exit +BUILD_PROFILE=${BUILD_PROFILE:-debug} + +if [ $# -ne 1 ]; then + echo "Usage: $0 - Start number of databend-query with system-managed mode" + exit 1 +fi + +if ! [[ "$num" =~ ^[0-9]*$ ]]; then + echo "Error: Argument must be an integer." + exit 1 +fi + +# Caveat: has to kill query first. +# `query` tries to remove its liveness record from meta before shutting down. +# If meta is stopped, `query` will receive an error that hangs graceful +# shutdown. +killall databend-query || true +sleep 3 + +killall databend-meta || true +sleep 3 + +for bin in databend-query databend-meta; do + if test -n "$(pgrep $bin)"; then + echo "The $bin is not killed. force killing." + killall -9 $bin || true + fi +done + +# Wait for killed process to cleanup resources +sleep 1 + +echo 'Start Meta service HA cluster(3 nodes)...' + +mkdir -p ./.databend/ + +nohup ./target/${BUILD_PROFILE}/databend-meta -c scripts/ci/deploy/config/databend-meta-node-1.toml >./.databend/meta-1.out 2>&1 & +python3 scripts/ci/wait_tcp.py --timeout 30 --port 9191 + +# wait for cluster formation to complete. +sleep 1 + +nohup ./target/${BUILD_PROFILE}/databend-meta -c scripts/ci/deploy/config/databend-meta-node-2.toml >./.databend/meta-2.out 2>&1 & +python3 scripts/ci/wait_tcp.py --timeout 30 --port 28202 + +# wait for cluster formation to complete. +sleep 1 + +nohup ./target/${BUILD_PROFILE}/databend-meta -c scripts/ci/deploy/config/databend-meta-node-3.toml >./.databend/meta-3.out 2>&1 & +python3 scripts/ci/wait_tcp.py --timeout 30 --port 28302 + +# wait for cluster formation to complete. +sleep 1 + +num=$1 + +find_available_port() { + local base_port=20000 + local max_port=65535 + local attempts=10 + + for ((i=0; i/dev/null 2>&1; then + echo $port + return + fi + done + + echo "Unable to find an available port after $attempts attempts" >&2 + exit 1 +} + + +start_databend_query() { + local http_port=$1 + local mysql_port=$2 + local log_dir=$3 + system_managed_config="./scripts/ci/deploy/config/databend-query-node-system-managed.toml" + + temp_file=$(mktemp) + + if [ -f "$system_managed_config" ]; then + sed -e "s/flight_port/$(find_available_port)/g" \ + -e "s/admin_api_port/$(find_available_port)/g" \ + -e "s/metric_api_port/$(find_available_port)/g" \ + -e "s/mysql_port/${mysql_port}/g" \ + -e "s/clickhouse_port/$(find_available_port)/g" \ + -e "s/http_port/${http_port}/g" \ + -e "s/flight_sql_port/$(find_available_port)/g" \ + -e "s/query_logs/${log_dir}/g" \ + "$system_managed_config" > "$temp_file" + + if [ $? -eq 0 ]; then + echo "Start databend-query on port $http_port..." + nohup target/${BUILD_PROFILE}/databend-query -c $temp_file --internal-enable-sandbox-tenant & + + echo "Waiting on databend-query 10 seconds..." + python3 scripts/ci/wait_tcp.py --timeout 30 --port $http_port + else + echo "Error occurred during port replacement." + rm -f "$temp_file" + exit 1 + fi + else + echo "Error: system-managed config file is not exists." + exit 1 + fi +} + +start_databend_query 8000 3307 "logs_1" + +for (( i=1; i<$num; i++ )) +do + http_port=$(find_available_port) + mysql_port=$(find_available_port) + start_databend_query $http_port $mysql_port "logs_$(( i + 1 ))" +done diff --git a/src/query/config/src/config.rs b/src/query/config/src/config.rs index a0b90c75c258c..63d5eb488e281 100644 --- a/src/query/config/src/config.rs +++ b/src/query/config/src/config.rs @@ -152,7 +152,6 @@ pub struct Config { /// when converted from inner config, all catalog configurations will store in `catalogs` #[clap(skip)] pub catalogs: HashMap, - // #[clap(flatten)] } #[derive(Subcommand, Default, Debug, Eq, PartialEq, Clone, Serialize, Deserialize)] @@ -1690,6 +1689,9 @@ pub struct QueryConfig { #[clap(skip)] pub settings: HashMap, + + #[clap(skip)] + pub resources_management: Option, } impl Default for QueryConfig { @@ -1782,6 +1784,7 @@ impl TryInto for QueryConfig { .into_iter() .map(|(k, v)| (k, v.into())) .collect(), + resources_management: self.resources_management, }) } } @@ -1880,6 +1883,7 @@ impl From for QueryConfig { cloud_control_grpc_timeout: inner.cloud_control_grpc_timeout, max_cached_queries_profiles: inner.max_cached_queries_profiles, settings: HashMap::new(), + resources_management: None, } } } diff --git a/src/query/config/src/inner.rs b/src/query/config/src/inner.rs index 7f0df3db978ab..7a2b4c4ea31e0 100644 --- a/src/query/config/src/inner.rs +++ b/src/query/config/src/inner.rs @@ -36,6 +36,7 @@ use databend_common_tracing::Config as LogConfig; use super::config::Commands; use super::config::Config; +use super::config::ResourcesManagementConfig; use crate::background_config::InnerBackgroundConfig; use crate::BuiltInConfig; @@ -247,6 +248,7 @@ pub struct QueryConfig { pub cloud_control_grpc_timeout: u64, pub max_cached_queries_profiles: usize, pub settings: HashMap, + pub resources_management: Option, } impl Default for QueryConfig { @@ -323,6 +325,7 @@ impl Default for QueryConfig { data_retention_time_in_days_max: 90, max_cached_queries_profiles: 50, settings: HashMap::new(), + resources_management: None, } } } diff --git a/src/query/ee/src/resource_management/mod.rs b/src/query/ee/src/resource_management/mod.rs index c7fa58e84a8e9..0e57155906936 100644 --- a/src/query/ee/src/resource_management/mod.rs +++ b/src/query/ee/src/resource_management/mod.rs @@ -13,6 +13,7 @@ // limitations under the License. mod resources_management_kubernetes; +mod resources_management_self_managed; mod resources_management_system; pub use resources_management_system::SystemResourcesManagement; diff --git a/src/query/ee/src/resource_management/resources_management_self_managed.rs b/src/query/ee/src/resource_management/resources_management_self_managed.rs new file mode 100644 index 0000000000000..3eeb822c1d3b4 --- /dev/null +++ b/src/query/ee/src/resource_management/resources_management_self_managed.rs @@ -0,0 +1,112 @@ +// Copyright 2023 Databend Cloud +// +// Licensed under the Elastic License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// https://www.elastic.co/licensing/elastic-license +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::collections::HashMap; + +use databend_common_exception::ErrorCode; +use databend_common_exception::Result; +use databend_common_management::SelectedNode; +use databend_common_management::WarehouseInfo; +use databend_common_meta_types::NodeInfo; +use databend_enterprise_resources_management::ResourcesManagement; + +pub struct SelfManagedResourcesManagement {} + +#[async_trait::async_trait] +impl ResourcesManagement for SelfManagedResourcesManagement { + async fn create_warehouse(&self, _: String, _: Vec) -> Result<()> { + Err(ErrorCode::Unimplemented( + "Unimplemented create warehouse with self-managed resources management", + )) + } + + async fn drop_warehouse(&self, _: String) -> Result<()> { + Err(ErrorCode::Unimplemented( + "Unimplemented drop warehouse with self-managed resources management", + )) + } + + async fn resume_warehouse(&self, _: String) -> Result<()> { + Err(ErrorCode::Unimplemented( + "Unimplemented resume warehouse with self-managed resources management", + )) + } + + async fn suspend_warehouse(&self, _: String) -> Result<()> { + Err(ErrorCode::Unimplemented( + "Unimplemented suspend warehouse with self-managed resources management", + )) + } + + async fn rename_warehouse(&self, _: String, _: String) -> Result<()> { + Err(ErrorCode::Unimplemented( + "Unimplemented rename warehouse with self-managed resources management", + )) + } + + async fn inspect_warehouse(&self, _: String) -> Result> { + Err(ErrorCode::Unimplemented( + "Unimplemented inspect warehouse with self-managed resources management", + )) + } + + async fn list_warehouses(&self) -> Result> { + Err(ErrorCode::Unimplemented( + "Unimplemented list warehouses with self-managed resources management", + )) + } + + async fn add_warehouse_cluster( + &self, + _: String, + _: String, + _: Vec, + ) -> Result<()> { + Err(ErrorCode::Unimplemented( + "Unimplemented add warehouse cluster with self-managed resources management", + )) + } + + async fn rename_warehouse_cluster(&self, _: String, _: String, _: String) -> Result<()> { + Err(ErrorCode::Unimplemented( + "Unimplemented rename warehouse cluster with self-managed resources management", + )) + } + + async fn drop_warehouse_cluster(&self, _: String, _: String) -> Result<()> { + Err(ErrorCode::Unimplemented( + "Unimplemented drop warehouse cluster with self-managed resources management", + )) + } + + async fn assign_warehouse_nodes( + &self, + _: String, + _: HashMap>, + ) -> Result<()> { + Err(ErrorCode::Unimplemented( + "Unimplemented assign warehouse nodes with self-managed resources management", + )) + } + + async fn unassign_warehouse_nodes( + &self, + _: String, + _: HashMap>, + ) -> Result<()> { + Err(ErrorCode::Unimplemented( + "Unimplemented unassign warehouse nodes with self-managed resources management", + )) + } +} diff --git a/src/query/ee/src/resource_management/resources_management_system.rs b/src/query/ee/src/resource_management/resources_management_system.rs index 5b7f412449416..3df61f30d0db6 100644 --- a/src/query/ee/src/resource_management/resources_management_system.rs +++ b/src/query/ee/src/resource_management/resources_management_system.rs @@ -28,6 +28,8 @@ use databend_common_meta_store::MetaStoreProvider; use databend_common_meta_types::NodeInfo; use databend_enterprise_resources_management::ResourcesManagement; +use crate::resource_management::resources_management_self_managed::SelfManagedResourcesManagement; + pub struct SystemResourcesManagement { warehouse_manager: Arc, } @@ -113,23 +115,34 @@ impl ResourcesManagement for SystemResourcesManagement { impl SystemResourcesManagement { pub async fn init(cfg: &InnerConfig) -> Result<()> { - let meta_api_provider = MetaStoreProvider::new(cfg.meta.to_meta_grpc_client_conf()); - match meta_api_provider.create_meta_store().await { - Err(cause) => { - Err(ErrorCode::from(cause).add_message_back("(while create resources management).")) - } - Ok(metastore) => { - let tenant_id = &cfg.query.tenant_id; - let lift_time = Duration::from_secs(60); - let warehouse_manager = - WarehouseMgr::create(metastore, tenant_id.tenant_name(), lift_time)?; - - let service: Arc = Arc::new(SystemResourcesManagement { - warehouse_manager: Arc::new(warehouse_manager), - }); - GlobalInstance::set(service); - Ok(()) - } + if !cfg.query.cluster_id.is_empty() && cfg.query.resources_management.is_some() { + return Err(ErrorCode::InvalidConfig("Configuration error, cannot set both cluster_id and resources_management at the same time")); } + + let service: Arc = match &cfg.query.resources_management { + None => Arc::new(SelfManagedResourcesManagement {}), + Some(_resources_management) => { + let meta_api_provider = MetaStoreProvider::new(cfg.meta.to_meta_grpc_client_conf()); + match meta_api_provider.create_meta_store().await { + Err(cause) => { + return Err(ErrorCode::from(cause) + .add_message_back("(while create resources management).")) + } + Ok(metastore) => { + let tenant_id = &cfg.query.tenant_id; + let lift_time = Duration::from_secs(60); + let warehouse_manager = + WarehouseMgr::create(metastore, tenant_id.tenant_name(), lift_time)?; + + Arc::new(SystemResourcesManagement { + warehouse_manager: Arc::new(warehouse_manager), + }) + } + } + } + }; + + GlobalInstance::set(service); + Ok(()) } } diff --git a/src/query/management/src/warehouse/warehouse_api.rs b/src/query/management/src/warehouse/warehouse_api.rs index a0b388a5d9f96..df9567c5f6140 100644 --- a/src/query/management/src/warehouse/warehouse_api.rs +++ b/src/query/management/src/warehouse/warehouse_api.rs @@ -106,5 +106,7 @@ pub trait WarehouseApi: Sync + Send { async fn get_local_addr(&self) -> Result>; + async fn list_online_nodes(&self) -> Result>; + async fn get_node_info(&self, node_id: &str) -> Result; } diff --git a/src/query/management/src/warehouse/warehouse_mgr.rs b/src/query/management/src/warehouse/warehouse_mgr.rs index 6646743777562..a3dedc5f0aa5b 100644 --- a/src/query/management/src/warehouse/warehouse_mgr.rs +++ b/src/query/management/src/warehouse/warehouse_mgr.rs @@ -1932,6 +1932,17 @@ impl WarehouseApi for WarehouseMgr { Ok(self.metastore.get_local_addr().await?) } + async fn list_online_nodes(&self) -> Result> { + let reply = self.metastore.prefix_list_kv(&self.node_key_prefix).await?; + let mut online_nodes = Vec::with_capacity(reply.len()); + + for (_, seq_data) in reply { + online_nodes.push(serde_json::from_slice::(&seq_data.data)?); + } + + Ok(online_nodes) + } + async fn get_node_info(&self, node_id: &str) -> Result { let node_key = format!("{}/{}", self.node_key_prefix, escape_for_key(node_id)?); let node_info = self.metastore.get_kv(&node_key).await?; diff --git a/src/query/service/src/clusters/cluster.rs b/src/query/service/src/clusters/cluster.rs index 3b267cf0636a0..d6174185cde48 100644 --- a/src/query/service/src/clusters/cluster.rs +++ b/src/query/service/src/clusters/cluster.rs @@ -65,7 +65,7 @@ pub struct ClusterDiscovery { local_id: String, local_secret: String, heartbeat: Mutex, - api_provider: Arc, + warehouse_manager: Arc, cluster_id: String, tenant_id: String, flight_address: String, @@ -215,7 +215,7 @@ impl ClusterDiscovery { Ok(Arc::new(ClusterDiscovery { local_id: cfg.query.node_id.clone(), local_secret: cfg.query.node_secret.clone(), - api_provider: provider.clone(), + warehouse_manager: provider.clone(), heartbeat: Mutex::new(ClusterHeartbeat::create( lift_time, provider, @@ -248,7 +248,7 @@ impl ClusterDiscovery { #[async_backtrace::framed] pub async fn discover(&self, config: &InnerConfig) -> Result> { match self - .api_provider + .warehouse_manager .list_warehouse_cluster_nodes(&self.cluster_id, &self.cluster_id) .await { @@ -324,11 +324,7 @@ impl ClusterDiscovery { #[async_backtrace::framed] async fn drop_invalid_nodes(self: &Arc, node_info: &NodeInfo) -> Result<()> { - let current_nodes_info = match self - .api_provider - .list_warehouse_cluster_nodes(&node_info.warehouse_id, &node_info.cluster_id) - .await - { + let online_nodes = match self.warehouse_manager.list_online_nodes().await { Ok(nodes) => nodes, Err(cause) => { metric_incr_cluster_error_count( @@ -342,11 +338,11 @@ impl ClusterDiscovery { } }; - for before_node in current_nodes_info { + for before_node in online_nodes { // Restart in a very short time(< heartbeat timeout) after abnormal shutdown, Which will // lead to some invalid information if before_node.flight_address.eq(&node_info.flight_address) { - let drop_invalid_node = self.api_provider.shutdown_node(before_node.id); + let drop_invalid_node = self.warehouse_manager.shutdown_node(before_node.id); if let Err(cause) = drop_invalid_node.await { warn!("Drop invalid node failure: {:?}", cause); } @@ -369,7 +365,7 @@ impl ClusterDiscovery { let mut mut_signal_pin = signal.as_mut(); let signal_future = Box::pin(mut_signal_pin.next()); - let drop_node = Box::pin(self.api_provider.shutdown_node(self.local_id.clone())); + let drop_node = Box::pin(self.warehouse_manager.shutdown_node(self.local_id.clone())); match futures::future::select(drop_node, signal_future).await { Either::Left((drop_node_result, _)) => { if let Err(drop_node_failure) = drop_node_result { @@ -412,7 +408,7 @@ impl ClusterDiscovery { if let Ok(socket_addr) = SocketAddr::from_str(lookup_ip) { let ip_addr = socket_addr.ip(); if ip_addr.is_loopback() || ip_addr.is_unspecified() { - if let Some(local_addr) = self.api_provider.get_local_addr().await? { + if let Some(local_addr) = self.warehouse_manager.get_local_addr().await? { let local_socket_addr = SocketAddr::from_str(&local_addr)?; let new_addr = format!("{}:{}", local_socket_addr.ip(), socket_addr.port()); warn!( @@ -443,7 +439,8 @@ impl ClusterDiscovery { node_info.node_type = NodeType::SelfManaged; self.drop_invalid_nodes(&node_info).await?; - match self.api_provider.start_node(node_info.clone()).await { + + match self.warehouse_manager.start_node(node_info.clone()).await { Ok(seq) => self.start_heartbeat(node_info, seq).await, Err(cause) => Err(cause.add_message_back("(while cluster api add_node).")), } From ff0244d5653d79a6355f48e73878f9450f23a1c1 Mon Sep 17 00:00:00 2001 From: zhang2014 Date: Thu, 26 Dec 2024 16:07:19 +0800 Subject: [PATCH 23/53] feat(cluster): support custom management cluster --- Cargo.lock | 1 + .../databend-query-node-system-managed.toml | 3 +- src/query/catalog/src/cluster_info.rs | 2 + src/query/config/src/config.rs | 4 +- src/query/ee/src/enterprise_services.rs | 4 +- src/query/ee/src/resource_management/mod.rs | 55 ++++++++ .../resources_management_kubernetes.rs | 13 ++ .../resources_management_self_managed.rs | 24 ++++ .../resources_management_system.rs | 74 +++++------ .../resources_management/Cargo.toml | 1 + .../src/resources_management.rs | 12 ++ .../management/src/warehouse/warehouse_api.rs | 2 +- .../management/src/warehouse/warehouse_mgr.rs | 60 ++++++--- src/query/service/src/clusters/cluster.rs | 40 +++--- .../interpreter_add_warehouse_cluster.rs | 12 +- .../interpreter_create_warehouses.rs | 12 +- .../src/interpreters/interpreter_factory.rs | 125 +++++++++++------- .../flight/v1/packets/packet_publisher.rs | 10 +- src/query/service/src/test_kits/fixture.rs | 2 +- .../storages/system/src/clusters_table.rs | 10 +- 20 files changed, 321 insertions(+), 145 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 87c5447c199ab..5ddaf4f4ee3a7 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4790,6 +4790,7 @@ version = "0.1.0" dependencies = [ "async-trait", "databend-common-base", + "databend-common-config", "databend-common-exception", "databend-common-management", "databend-common-meta-types", diff --git a/scripts/ci/deploy/config/databend-query-node-system-managed.toml b/scripts/ci/deploy/config/databend-query-node-system-managed.toml index 07374d0f845ee..368f7f2ceab3a 100644 --- a/scripts/ci/deploy/config/databend-query-node-system-managed.toml +++ b/scripts/ci/deploy/config/databend-query-node-system-managed.toml @@ -32,7 +32,6 @@ flight_sql_handler_host = "0.0.0.0" flight_sql_handler_port = flight_sql_port tenant_id = "test_tenant" -cluster_id = "test_cluster" table_engine_memory_enabled = true default_storage_format = 'parquet' @@ -52,7 +51,7 @@ name = "ping" definition = "CREATE FUNCTION ping(STRING) RETURNS STRING LANGUAGE python HANDLER = 'ping' ADDRESS = 'http://0.0.0.0:8815'" [query.resources_management] -type = "system" +type = "system_managed" [log] diff --git a/src/query/catalog/src/cluster_info.rs b/src/query/catalog/src/cluster_info.rs index 20cba2e624845..0142742cc1933 100644 --- a/src/query/catalog/src/cluster_info.rs +++ b/src/query/catalog/src/cluster_info.rs @@ -18,7 +18,9 @@ use databend_common_meta_types::NodeInfo; #[derive(Debug, serde::Serialize, serde::Deserialize)] pub struct Cluster { + pub unassign: bool, pub local_id: String, + pub nodes: Vec>, } diff --git a/src/query/config/src/config.rs b/src/query/config/src/config.rs index 63d5eb488e281..1aab42befee29 100644 --- a/src/query/config/src/config.rs +++ b/src/query/config/src/config.rs @@ -2960,12 +2960,12 @@ pub struct SpillConfig { #[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize, Args, Default)] #[serde(default, deny_unknown_fields)] pub struct ResourcesManagementConfig { - #[clap(long = "type", value_name = "VALUE", default_value = "system")] + #[clap(long = "type", value_name = "VALUE", default_value = "self_managed")] #[serde(rename = "type")] pub typ: String, #[clap(long, value_name = "VALUE")] - pub instance_type: Option, + pub resource_group: Option, } mod cache_config_converters { diff --git a/src/query/ee/src/enterprise_services.rs b/src/query/ee/src/enterprise_services.rs index 7433e65f253be..1507ac44a43f0 100644 --- a/src/query/ee/src/enterprise_services.rs +++ b/src/query/ee/src/enterprise_services.rs @@ -23,7 +23,7 @@ use crate::data_mask::RealDatamaskHandler; use crate::fail_safe::RealFailSafeHandler; use crate::inverted_index::RealInvertedIndexHandler; use crate::license::license_mgr::RealLicenseManager; -use crate::resource_management::SystemResourcesManagement; +use crate::resource_management::init_resources_management; use crate::storage_encryption::RealStorageEncryptionHandler; use crate::storage_quota::RealStorageQuotaHandler; use crate::storages::fuse::operations::RealVacuumHandler; @@ -46,7 +46,7 @@ impl EnterpriseServices { RealInvertedIndexHandler::init()?; RealStorageQuotaHandler::init(&cfg)?; RealFailSafeHandler::init()?; - SystemResourcesManagement::init(&cfg).await?; + init_resources_management(&cfg).await?; Ok(()) } } diff --git a/src/query/ee/src/resource_management/mod.rs b/src/query/ee/src/resource_management/mod.rs index 0e57155906936..ee0a7877d15e3 100644 --- a/src/query/ee/src/resource_management/mod.rs +++ b/src/query/ee/src/resource_management/mod.rs @@ -16,4 +16,59 @@ mod resources_management_kubernetes; mod resources_management_self_managed; mod resources_management_system; +use std::sync::Arc; +use std::time::Duration; + +use databend_common_base::base::GlobalInstance; +use databend_common_config::InnerConfig; +use databend_common_exception::ErrorCode; +use databend_common_exception::Result; +use databend_common_management::WarehouseMgr; +use databend_common_meta_store::MetaStoreProvider; +use databend_enterprise_resources_management::ResourcesManagement; pub use resources_management_system::SystemResourcesManagement; + +use crate::resource_management::resources_management_kubernetes::KubernetesResourcesManagement; +use crate::resource_management::resources_management_self_managed::SelfManagedResourcesManagement; + +pub async fn init_resources_management(cfg: &InnerConfig) -> Result<()> { + let service: Arc = match &cfg.query.resources_management { + None => match cfg.query.cluster_id.is_empty() { + true => Err(ErrorCode::InvalidConfig( + "cluster_id is empty without resources management.", + )), + false => SelfManagedResourcesManagement::create(cfg), + }, + Some(resources_management) => { + match resources_management.typ.to_ascii_lowercase().as_str() { + "self_managed" => SelfManagedResourcesManagement::create(cfg), + "kubernetes_managed" => KubernetesResourcesManagement::create(), + "system_managed" => { + let meta_api_provider = + MetaStoreProvider::new(cfg.meta.to_meta_grpc_client_conf()); + match meta_api_provider.create_meta_store().await { + Err(cause) => Err(ErrorCode::from(cause) + .add_message_back("(while create resources management).")), + Ok(metastore) => { + let tenant_id = &cfg.query.tenant_id; + let lift_time = Duration::from_secs(60); + let warehouse_manager = WarehouseMgr::create( + metastore, + tenant_id.tenant_name(), + lift_time, + )?; + SystemResourcesManagement::create(Arc::new(warehouse_manager)) + } + } + } + _ => Err(ErrorCode::InvalidConfig(format!( + "unimplemented resources management {}", + resources_management.typ + ))), + } + } + }?; + + GlobalInstance::set(service); + Ok(()) +} diff --git a/src/query/ee/src/resource_management/resources_management_kubernetes.rs b/src/query/ee/src/resource_management/resources_management_kubernetes.rs index 784050b2d16f4..fb35e4e04d58a 100644 --- a/src/query/ee/src/resource_management/resources_management_kubernetes.rs +++ b/src/query/ee/src/resource_management/resources_management_kubernetes.rs @@ -13,6 +13,7 @@ // limitations under the License. use std::collections::HashMap; +use std::sync::Arc; use databend_common_exception::ErrorCode; use databend_common_exception::Result; @@ -23,8 +24,20 @@ use databend_enterprise_resources_management::ResourcesManagement; pub struct KubernetesResourcesManagement {} +impl KubernetesResourcesManagement { + pub fn create() -> Result> { + Ok(Arc::new(KubernetesResourcesManagement {})) + } +} + #[async_trait::async_trait] impl ResourcesManagement for KubernetesResourcesManagement { + async fn init_node(&self, _: &mut NodeInfo) -> Result<()> { + Err(ErrorCode::Unimplemented( + "Unimplemented kubernetes resources management", + )) + } + async fn create_warehouse(&self, _: String, _: Vec) -> Result<()> { Err(ErrorCode::Unimplemented( "Unimplemented kubernetes resources management", diff --git a/src/query/ee/src/resource_management/resources_management_self_managed.rs b/src/query/ee/src/resource_management/resources_management_self_managed.rs index 3eeb822c1d3b4..47a306b6c143c 100644 --- a/src/query/ee/src/resource_management/resources_management_self_managed.rs +++ b/src/query/ee/src/resource_management/resources_management_self_managed.rs @@ -13,18 +13,42 @@ // limitations under the License. use std::collections::HashMap; +use std::sync::Arc; +use databend_common_config::GlobalConfig; +use databend_common_config::InnerConfig; use databend_common_exception::ErrorCode; use databend_common_exception::Result; use databend_common_management::SelectedNode; use databend_common_management::WarehouseInfo; use databend_common_meta_types::NodeInfo; +use databend_common_meta_types::NodeType; use databend_enterprise_resources_management::ResourcesManagement; pub struct SelfManagedResourcesManagement {} +impl SelfManagedResourcesManagement { + pub fn create(cfg: &InnerConfig) -> Result> { + if cfg.query.cluster_id.is_empty() { + return Err(ErrorCode::InvalidConfig( + "cluster_id is empty with self-managed resources management", + )); + } + + Ok(Arc::new(SelfManagedResourcesManagement {})) + } +} + #[async_trait::async_trait] impl ResourcesManagement for SelfManagedResourcesManagement { + async fn init_node(&self, node: &mut NodeInfo) -> Result<()> { + let config = GlobalConfig::instance(); + node.cluster_id = config.query.cluster_id.clone(); + node.warehouse_id = config.query.cluster_id.clone(); + node.node_type = NodeType::SelfManaged; + Ok(()) + } + async fn create_warehouse(&self, _: String, _: Vec) -> Result<()> { Err(ErrorCode::Unimplemented( "Unimplemented create warehouse with self-managed resources management", diff --git a/src/query/ee/src/resource_management/resources_management_system.rs b/src/query/ee/src/resource_management/resources_management_system.rs index 3df61f30d0db6..66999cd83b027 100644 --- a/src/query/ee/src/resource_management/resources_management_system.rs +++ b/src/query/ee/src/resource_management/resources_management_system.rs @@ -14,28 +14,50 @@ use std::collections::HashMap; use std::sync::Arc; -use std::time::Duration; -use databend_common_base::base::GlobalInstance; -use databend_common_config::InnerConfig; -use databend_common_exception::ErrorCode; +use databend_common_config::GlobalConfig; use databend_common_exception::Result; use databend_common_management::SelectedNode; use databend_common_management::WarehouseApi; use databend_common_management::WarehouseInfo; -use databend_common_management::WarehouseMgr; -use databend_common_meta_store::MetaStoreProvider; use databend_common_meta_types::NodeInfo; +use databend_common_meta_types::NodeType; use databend_enterprise_resources_management::ResourcesManagement; -use crate::resource_management::resources_management_self_managed::SelfManagedResourcesManagement; - pub struct SystemResourcesManagement { - warehouse_manager: Arc, + pub warehouse_manager: Arc, +} + +impl SystemResourcesManagement { + pub fn create(warehouse_mgr: Arc) -> Result> { + Ok(Arc::new(SystemResourcesManagement { + warehouse_manager: warehouse_mgr, + })) + } } #[async_trait::async_trait] impl ResourcesManagement for SystemResourcesManagement { + async fn init_node(&self, node: &mut NodeInfo) -> Result<()> { + let config = GlobalConfig::instance(); + assert!(config.query.cluster_id.is_empty()); + assert!(config.query.resources_management.is_some()); + + if let Some(resources_management) = &config.query.resources_management { + assert_eq!( + &resources_management.typ.to_ascii_lowercase(), + "system_managed" + ); + + node.cluster_id = String::new(); + node.warehouse_id = String::new(); + node.node_type = NodeType::SystemManaged; + node.resource_group = resources_management.resource_group.clone(); + } + + Ok(()) + } + async fn create_warehouse(&self, name: String, nodes: Vec) -> Result<()> { self.warehouse_manager.create_warehouse(name, nodes).await } @@ -112,37 +134,3 @@ impl ResourcesManagement for SystemResourcesManagement { .await } } - -impl SystemResourcesManagement { - pub async fn init(cfg: &InnerConfig) -> Result<()> { - if !cfg.query.cluster_id.is_empty() && cfg.query.resources_management.is_some() { - return Err(ErrorCode::InvalidConfig("Configuration error, cannot set both cluster_id and resources_management at the same time")); - } - - let service: Arc = match &cfg.query.resources_management { - None => Arc::new(SelfManagedResourcesManagement {}), - Some(_resources_management) => { - let meta_api_provider = MetaStoreProvider::new(cfg.meta.to_meta_grpc_client_conf()); - match meta_api_provider.create_meta_store().await { - Err(cause) => { - return Err(ErrorCode::from(cause) - .add_message_back("(while create resources management).")) - } - Ok(metastore) => { - let tenant_id = &cfg.query.tenant_id; - let lift_time = Duration::from_secs(60); - let warehouse_manager = - WarehouseMgr::create(metastore, tenant_id.tenant_name(), lift_time)?; - - Arc::new(SystemResourcesManagement { - warehouse_manager: Arc::new(warehouse_manager), - }) - } - } - } - }; - - GlobalInstance::set(service); - Ok(()) - } -} diff --git a/src/query/ee_features/resources_management/Cargo.toml b/src/query/ee_features/resources_management/Cargo.toml index 493bdfb3a8a71..428f76c3e3b33 100644 --- a/src/query/ee_features/resources_management/Cargo.toml +++ b/src/query/ee_features/resources_management/Cargo.toml @@ -13,6 +13,7 @@ test = true [dependencies] async-trait = { workspace = true } databend-common-base = { workspace = true } +databend-common-config = { workspace = true } databend-common-exception = { workspace = true } databend-common-management = { workspace = true } databend-common-meta-types = { workspace = true } diff --git a/src/query/ee_features/resources_management/src/resources_management.rs b/src/query/ee_features/resources_management/src/resources_management.rs index 66ac7096b8bff..32da1e0742059 100644 --- a/src/query/ee_features/resources_management/src/resources_management.rs +++ b/src/query/ee_features/resources_management/src/resources_management.rs @@ -15,14 +15,18 @@ use std::collections::HashMap; use databend_common_base::base::GlobalInstance; +use databend_common_config::GlobalConfig; use databend_common_exception::ErrorCode; use databend_common_exception::Result; use databend_common_management::SelectedNode; use databend_common_management::WarehouseInfo; use databend_common_meta_types::NodeInfo; +use databend_common_meta_types::NodeType; #[async_trait::async_trait] pub trait ResourcesManagement: Sync + Send + 'static { + async fn init_node(&self, node: &mut NodeInfo) -> Result<()>; + async fn create_warehouse(&self, name: String, nodes: Vec) -> Result<()>; async fn drop_warehouse(&self, name: String) -> Result<()>; @@ -70,6 +74,14 @@ pub struct DummyResourcesManagement; #[async_trait::async_trait] impl ResourcesManagement for DummyResourcesManagement { + async fn init_node(&self, node: &mut NodeInfo) -> Result<()> { + let config = GlobalConfig::instance(); + node.cluster_id = config.query.cluster_id.clone(); + node.warehouse_id = config.query.cluster_id.clone(); + node.node_type = NodeType::SelfManaged; + Ok(()) + } + async fn create_warehouse(&self, _: String, _: Vec) -> Result<()> { Err(ErrorCode::Unimplemented("The use of this feature requires a Databend Enterprise Edition license. To unlock enterprise features, please contact Databend to obtain a license. Learn more at https://docs.databend.com/guides/overview/editions/dee/")) } diff --git a/src/query/management/src/warehouse/warehouse_api.rs b/src/query/management/src/warehouse/warehouse_api.rs index df9567c5f6140..ab93c0cb03b9c 100644 --- a/src/query/management/src/warehouse/warehouse_api.rs +++ b/src/query/management/src/warehouse/warehouse_api.rs @@ -108,5 +108,5 @@ pub trait WarehouseApi: Sync + Send { async fn list_online_nodes(&self) -> Result>; - async fn get_node_info(&self, node_id: &str) -> Result; + async fn discover(&self, node_id: &str) -> Result<(bool, Vec)>; } diff --git a/src/query/management/src/warehouse/warehouse_mgr.rs b/src/query/management/src/warehouse/warehouse_mgr.rs index a3dedc5f0aa5b..57cf9781c8016 100644 --- a/src/query/management/src/warehouse/warehouse_mgr.rs +++ b/src/query/management/src/warehouse/warehouse_mgr.rs @@ -360,7 +360,7 @@ impl WarehouseMgr { async fn consistent_warehouse_info(&self, id: &str) -> Result { let warehouse_key = format!("{}/{}", self.warehouse_key_prefix, escape_for_key(id)?); - let nodes_prefix = format!("{}/{}", self.meta_key_prefix, escape_for_key(id)?); + let nodes_prefix = format!("{}/{}/", self.meta_key_prefix, escape_for_key(id)?); 'retry: for _idx in 0..64 { let Some(before_info) = self.metastore.get_kv(&warehouse_key).await? else { @@ -376,7 +376,7 @@ impl WarehouseMgr { let mut cluster_node_seq = Vec::with_capacity(values.len()); for (node_key, value) in values { - let suffix = &node_key[nodes_prefix.len() + 1..]; + let suffix = &node_key[nodes_prefix.len()..]; if let Some((_cluster, node)) = suffix.split_once('/') { let node_key = format!("{}/{}", self.node_key_prefix, node); @@ -894,7 +894,7 @@ impl WarehouseApi for WarehouseMgr { "Cannot resume self-managed warehouse.", )), WarehouseInfo::SystemManaged(warehouse) => { - if warehouse.status != "SUSPEND" { + if warehouse.status.to_uppercase() != "SUSPENDED" { return Err(ErrorCode::InvalidWarehouse(format!( "Cannot resume warehouse {:?}, because warehouse state is not suspend", warehouse @@ -905,7 +905,7 @@ impl WarehouseApi for WarehouseMgr { need_schedule_cluster = warehouse.clusters.clone(); Ok(WarehouseInfo::SystemManaged(SystemManagedWarehouse { id: warehouse.id.clone(), - status: "RUNNING".to_string(), + status: "Running".to_string(), display_name: warehouse.display_name, clusters: warehouse.clusters, })) @@ -1012,13 +1012,13 @@ impl WarehouseApi for WarehouseMgr { "Cannot suspend self-managed warehouse", )), WarehouseInfo::SystemManaged(warehouse) => { - if warehouse.status != "RUNNING" { - return Err(ErrorCode::InvalidWarehouse(format!("Cannot suspend warehouse {:?}, because warehouse state is not running.", warehouse))); + if warehouse.status.to_uppercase() != "RUNNING" { + return Err(ErrorCode::InvalidWarehouse(format!("Cannot suspend warehouse {:?}, because warehouse state is not running.", warehouse.display_name))); } Ok(WarehouseInfo::SystemManaged(SystemManagedWarehouse { id: warehouse.id.clone(), - status: "RUNNING".to_string(), + status: "Suspended".to_string(), display_name: warehouse.display_name, clusters: warehouse.clusters, })) @@ -1904,20 +1904,20 @@ impl WarehouseApi for WarehouseMgr { warehouse: &str, cluster: &str, ) -> Result> { - let cluster_key = format!( - "{}/{}/{}", + let cluster_prefix = format!( + "{}/{}/{}/", self.meta_key_prefix, escape_for_key(warehouse)?, escape_for_key(cluster)? ); - let values = self.metastore.prefix_list_kv(&cluster_key).await?; + let values = self.metastore.prefix_list_kv(&cluster_prefix).await?; let mut nodes_info = Vec::with_capacity(values.len()); for (node_key, value) in values { let mut node_info = serde_json::from_slice::(&value.data)?; - node_info.id = unescape_for_key(&node_key[cluster_key.len() + 1..])?; + node_info.id = unescape_for_key(&node_key[cluster_prefix.len()..])?; node_info.cluster_id = cluster.to_string(); node_info.warehouse_id = warehouse.to_string(); nodes_info.push(node_info); @@ -1943,12 +1943,40 @@ impl WarehouseApi for WarehouseMgr { Ok(online_nodes) } - async fn get_node_info(&self, node_id: &str) -> Result { + async fn discover(&self, node_id: &str) -> Result<(bool, Vec)> { let node_key = format!("{}/{}", self.node_key_prefix, escape_for_key(node_id)?); - let node_info = self.metastore.get_kv(&node_key).await?; - match node_info { - None => Err(ErrorCode::NotFoundClusterNode("")), - Some(v) => Ok(serde_json::from_slice(&v.data)?), + let Some(seq) = self.metastore.get_kv(&node_key).await? else { + return Err(ErrorCode::NotFoundClusterNode(format!( + "Node {} is offline, Please restart this node.", + node_id + ))); + }; + + let self_info = serde_json::from_slice::(&seq.data)?; + + if self_info.warehouse_id.is_empty() && self_info.cluster_id.is_empty() { + return Ok((false, vec![self_info])); } + + let cluster_prefix = format!( + "{}/{}/{}/", + self.meta_key_prefix, + escape_for_key(&self_info.warehouse_id)?, + escape_for_key(&self_info.cluster_id)? + ); + + let values = self.metastore.prefix_list_kv(&cluster_prefix).await?; + + let mut cluster_nodes_info = Vec::with_capacity(values.len()); + for (node_key, value) in values { + let mut cluster_node = serde_json::from_slice::(&value.data)?; + + cluster_node.id = unescape_for_key(&node_key[cluster_prefix.len()..])?; + cluster_node.cluster_id = self_info.cluster_id.clone(); + cluster_node.warehouse_id = self_info.warehouse_id.clone(); + cluster_nodes_info.push(cluster_node); + } + + Ok((true, cluster_nodes_info)) } } diff --git a/src/query/service/src/clusters/cluster.rs b/src/query/service/src/clusters/cluster.rs index d6174185cde48..38e774a273b32 100644 --- a/src/query/service/src/clusters/cluster.rs +++ b/src/query/service/src/clusters/cluster.rs @@ -43,8 +43,8 @@ use databend_common_management::WarehouseMgr; use databend_common_meta_store::MetaStore; use databend_common_meta_store::MetaStoreProvider; use databend_common_meta_types::NodeInfo; -use databend_common_meta_types::NodeType; use databend_common_metrics::cluster::*; +use databend_enterprise_resources_management::ResourcesManagement; use futures::future::select; use futures::future::Either; use futures::Future; @@ -75,7 +75,7 @@ pub struct ClusterDiscovery { // avoid leak FlightClient to common-xxx #[async_trait::async_trait] pub trait ClusterHelper { - fn create(nodes: Vec>, local_id: String) -> Arc; + fn create(unassign: bool, nodes: Vec>, local_id: String) -> Arc; fn empty() -> Arc; fn is_empty(&self) -> bool; fn is_local(&self, node: &NodeInfo) -> bool; @@ -93,12 +93,17 @@ pub trait ClusterHelper { #[async_trait::async_trait] impl ClusterHelper for Cluster { - fn create(nodes: Vec>, local_id: String) -> Arc { - Arc::new(Cluster { local_id, nodes }) + fn create(unassign: bool, nodes: Vec>, local_id: String) -> Arc { + Arc::new(Cluster { + unassign, + local_id, + nodes, + }) } fn empty() -> Arc { Arc::new(Cluster { + unassign: false, local_id: String::from(""), nodes: Vec::new(), }) @@ -247,11 +252,16 @@ impl ClusterDiscovery { #[async_backtrace::framed] pub async fn discover(&self, config: &InnerConfig) -> Result> { - match self - .warehouse_manager - .list_warehouse_cluster_nodes(&self.cluster_id, &self.cluster_id) - .await - { + let nodes = match config.query.cluster_id.is_empty() { + true => self.warehouse_manager.discover(&config.query.node_id).await, + false => self + .warehouse_manager + .list_warehouse_cluster_nodes(&self.cluster_id, &self.cluster_id) + .await + .map(|x| (true, x)), + }; + + match nodes { Err(cause) => { metric_incr_cluster_error_count( &self.local_id, @@ -262,7 +272,7 @@ impl ClusterDiscovery { ); Err(cause.add_message_back("(while cluster api get_nodes).")) } - Ok(cluster_nodes) => { + Ok((has_cluster, cluster_nodes)) => { let mut res = Vec::with_capacity(cluster_nodes.len()); for node in &cluster_nodes { if node.id != self.local_id { @@ -289,8 +299,9 @@ impl ClusterDiscovery { &self.flight_address, cluster_nodes.len() as f64, ); - let res = Cluster::create(res, self.local_id.clone()); - *self.cached_cluster.write() = Some(res.clone()); + + let res = Cluster::create(!has_cluster, res, self.local_id.clone()); + // *self.cached_cluster.write() = Some(res.clone()); Ok(res) } } @@ -434,9 +445,8 @@ impl ClusterDiscovery { DATABEND_COMMIT_VERSION.to_string(), ); - node_info.cluster_id = self.cluster_id.clone(); - node_info.warehouse_id = self.cluster_id.clone(); - node_info.node_type = NodeType::SelfManaged; + let resources_management = GlobalInstance::get::>(); + resources_management.init_node(&mut node_info).await?; self.drop_invalid_nodes(&node_info).await?; diff --git a/src/query/service/src/interpreters/interpreter_add_warehouse_cluster.rs b/src/query/service/src/interpreters/interpreter_add_warehouse_cluster.rs index da0139a8314c3..160dc6ea1e023 100644 --- a/src/query/service/src/interpreters/interpreter_add_warehouse_cluster.rs +++ b/src/query/service/src/interpreters/interpreter_add_warehouse_cluster.rs @@ -49,13 +49,17 @@ impl Interpreter for AddWarehouseClusterInterpreter { #[async_backtrace::framed] async fn execute2(&self) -> Result { - if let Some(cluster_size) = self.plan.options.get("CLUSTER_SIZE") { + if let Some(cluster_size) = self.plan.options.get("cluster_size") { if !self.plan.nodes.is_empty() { - return Err(ErrorCode::InvalidArgument(format!(""))); + return Err(ErrorCode::InvalidArgument( + "CLUSTER_SIZE option and node list exists in one query.", + )); } let Ok(cluster_size) = cluster_size.parse::() else { - return Err(ErrorCode::InvalidArgument("")); + return Err(ErrorCode::InvalidArgument( + "CLUSTER_SIZE must be a ", + )); }; GlobalInstance::get::>() @@ -70,7 +74,7 @@ impl Interpreter for AddWarehouseClusterInterpreter { } if self.plan.nodes.is_empty() { - return Err(ErrorCode::InvalidArgument("")); + return Err(ErrorCode::InvalidArgument("Cluster nodes list is empty")); } let mut selected_nodes = Vec::with_capacity(self.plan.nodes.len()); diff --git a/src/query/service/src/interpreters/interpreter_create_warehouses.rs b/src/query/service/src/interpreters/interpreter_create_warehouses.rs index a09c908d507bd..e58ad4ff86cd8 100644 --- a/src/query/service/src/interpreters/interpreter_create_warehouses.rs +++ b/src/query/service/src/interpreters/interpreter_create_warehouses.rs @@ -49,13 +49,17 @@ impl Interpreter for CreateWarehouseInterpreter { #[async_backtrace::framed] async fn execute2(&self) -> Result { - if let Some(warehouse_size) = self.plan.options.get("WAREHOUSE_SIZE") { + if let Some(warehouse_size) = self.plan.options.get("warehouse_size") { if !self.plan.nodes.is_empty() { - return Err(ErrorCode::InvalidArgument(format!(""))); + return Err(ErrorCode::InvalidArgument( + "WAREHOUSE_SIZE option and node list exists in one query.", + )); } let Ok(warehouse_size) = warehouse_size.parse::() else { - return Err(ErrorCode::InvalidArgument("")); + return Err(ErrorCode::InvalidArgument( + "WAREHOUSE_SIZE must be a ", + )); }; GlobalInstance::get::>() @@ -69,7 +73,7 @@ impl Interpreter for CreateWarehouseInterpreter { } if self.plan.nodes.is_empty() { - return Err(ErrorCode::InvalidArgument("")); + return Err(ErrorCode::InvalidArgument("Warehouse nodes list is empty")); } let mut selected_nodes = Vec::with_capacity(self.plan.nodes.len()); diff --git a/src/query/service/src/interpreters/interpreter_factory.rs b/src/query/service/src/interpreters/interpreter_factory.rs index f5581a8595788..a615fa6acbc9e 100644 --- a/src/query/service/src/interpreters/interpreter_factory.rs +++ b/src/query/service/src/interpreters/interpreter_factory.rs @@ -16,6 +16,7 @@ use std::sync::Arc; use databend_common_ast::ast::ExplainKind; use databend_common_catalog::lock::LockTableOption; +use databend_common_catalog::table_context::TableContext; use databend_common_exception::ErrorCode; use databend_common_exception::Result; use databend_common_sql::binder::ExplainConfig; @@ -101,17 +102,80 @@ impl InterpreterFactory { pub async fn get(ctx: Arc, plan: &Plan) -> Result { // Check the access permission. let access_checker = Accessor::create(ctx.clone()); - access_checker - .check(plan) - .await - .map_err(|e| match e.code() { - ErrorCode::PERMISSION_DENIED => { - error!("Access.denied(v2): {:?}", e); - e + access_checker.check(plan).await.inspect_err(|e| { + if e.code() == ErrorCode::PERMISSION_DENIED { + error!("Access.denied(v2): {:?}", e); + } + })?; + + Self::get_warehouses_interpreter(ctx, plan, Self::get_inner) + } + + pub fn get_warehouses_interpreter( + ctx: Arc, + plan: &Plan, + other: impl FnOnce(Arc, &Plan) -> Result, + ) -> Result { + match plan { + Plan::ShowWarehouses => Ok(Arc::new(ShowWarehousesInterpreter::try_create()?)), + Plan::CreateWarehouse(v) => Ok(Arc::new(CreateWarehouseInterpreter::try_create( + ctx.clone(), + *v.clone(), + )?)), + Plan::DropWarehouse(v) => Ok(Arc::new(DropWarehouseInterpreter::try_create( + ctx.clone(), + *v.clone(), + )?)), + Plan::ResumeWarehouse(v) => Ok(Arc::new(ResumeWarehouseInterpreter::try_create( + *v.clone(), + )?)), + Plan::SuspendWarehouse(v) => Ok(Arc::new(SuspendWarehouseInterpreter::try_create( + *v.clone(), + )?)), + Plan::RenameWarehouse(v) => Ok(Arc::new(RenameWarehouseInterpreter::try_create( + ctx.clone(), + *v.clone(), + )?)), + Plan::InspectWarehouse(v) => Ok(Arc::new(InspectWarehouseInterpreter::try_create( + *v.clone(), + )?)), + Plan::AddWarehouseCluster(v) => Ok(Arc::new( + AddWarehouseClusterInterpreter::try_create(ctx.clone(), *v.clone())?, + )), + Plan::DropWarehouseCluster(v) => Ok(Arc::new( + DropWarehouseClusterInterpreter::try_create(ctx.clone(), *v.clone())?, + )), + Plan::RenameWarehouseCluster(v) => Ok(Arc::new( + RenameWarehouseClusterInterpreter::try_create(*v.clone())?, + )), + Plan::AssignWarehouseNodes(v) => Ok(Arc::new( + AssignWarehouseNodesInterpreter::try_create(*v.clone())?, + )), + Plan::UnassignWarehouseNodes(v) => Ok(Arc::new( + DropWarehouseClusterNodeInterpreter::try_create(*v.clone())?, + )), + Plan::Query { metadata, .. } => { + let read_guard = metadata.read(); + for table in read_guard.tables() { + let db = table.database(); + // System library access is secure + if db != "system" && db != "information_schema" && ctx.get_cluster().unassign { + return Err(ErrorCode::InvalidWarehouse(format!( + "Node {} has not been assigned warehouse, please assign to warehouse before executing SQL on this node.", + ctx.get_cluster().local_id + ))); + } } - _ => e, - })?; - Self::get_inner(ctx, plan) + other(ctx, plan) + } + _ => match ctx.get_cluster().unassign { + true => Err(ErrorCode::InvalidWarehouse(format!( + "Node {} has not been assigned warehouse, please assign to warehouse before executing SQL on this node.", + ctx.get_cluster().local_id + ))), + false => other(ctx, plan), + }, + } } pub fn get_inner(ctx: Arc, plan: &Plan) -> Result { @@ -643,43 +707,10 @@ impl InterpreterFactory { // ctx, // p.clone(), // )?)), - Plan::ShowWarehouses => Ok(Arc::new(ShowWarehousesInterpreter::try_create()?)), - Plan::CreateWarehouse(v) => Ok(Arc::new(CreateWarehouseInterpreter::try_create( - ctx, - *v.clone(), - )?)), - Plan::DropWarehouse(v) => Ok(Arc::new(DropWarehouseInterpreter::try_create( - ctx, - *v.clone(), - )?)), - Plan::ResumeWarehouse(v) => Ok(Arc::new(ResumeWarehouseInterpreter::try_create( - *v.clone(), - )?)), - Plan::SuspendWarehouse(v) => Ok(Arc::new(SuspendWarehouseInterpreter::try_create( - *v.clone(), - )?)), - Plan::RenameWarehouse(v) => Ok(Arc::new(RenameWarehouseInterpreter::try_create( - ctx, - *v.clone(), - )?)), - Plan::InspectWarehouse(v) => Ok(Arc::new(InspectWarehouseInterpreter::try_create( - *v.clone(), - )?)), - Plan::AddWarehouseCluster(v) => Ok(Arc::new( - AddWarehouseClusterInterpreter::try_create(ctx, *v.clone())?, - )), - Plan::DropWarehouseCluster(v) => Ok(Arc::new( - DropWarehouseClusterInterpreter::try_create(ctx, *v.clone())?, - )), - Plan::RenameWarehouseCluster(v) => Ok(Arc::new( - RenameWarehouseClusterInterpreter::try_create(*v.clone())?, - )), - Plan::AssignWarehouseNodes(v) => Ok(Arc::new( - AssignWarehouseNodesInterpreter::try_create(*v.clone())?, - )), - Plan::UnassignWarehouseNodes(v) => Ok(Arc::new( - DropWarehouseClusterNodeInterpreter::try_create(*v.clone())?, - )), + _ => Err(ErrorCode::Unimplemented(format!( + "Unimplemented interpreter for {:?} plan", + plan + ))), } } } diff --git a/src/query/service/src/servers/flight/v1/packets/packet_publisher.rs b/src/query/service/src/servers/flight/v1/packets/packet_publisher.rs index 12bb478627759..753c1eb0c2d1d 100644 --- a/src/query/service/src/servers/flight/v1/packets/packet_publisher.rs +++ b/src/query/service/src/servers/flight/v1/packets/packet_publisher.rs @@ -93,11 +93,10 @@ impl DataflowDiagramBuilder { let source = self .nodes .get(source) - .ok_or_else(|| ErrorCode::NotFoundClusterNode(""))?; - let destination = self - .nodes - .get(destination) - .ok_or_else(|| ErrorCode::NotFoundClusterNode(""))?; + .ok_or_else(|| ErrorCode::NotFoundClusterNode(format!("not found {}", source)))?; + let destination = self.nodes.get(destination).ok_or_else(|| { + ErrorCode::NotFoundClusterNode(format!("not found {}", destination)) + })?; self.graph .add_edge(*source, *destination, Edge::Fragment(v)); @@ -166,6 +165,7 @@ impl QueryEnv { )?; let query_ctx = session.create_query_context_with_cluster(Arc::new(Cluster { + unassign: self.cluster.unassign, nodes: self.cluster.nodes.clone(), local_id: GlobalConfig::instance().query.node_id.clone(), }))?; diff --git a/src/query/service/src/test_kits/fixture.rs b/src/query/service/src/test_kits/fixture.rs index 0417592029574..2e46303dae24c 100644 --- a/src/query/service/src/test_kits/fixture.rs +++ b/src/query/service/src/test_kits/fixture.rs @@ -272,7 +272,7 @@ impl TestFixture { let dummy_query_context = QueryContext::create_from_shared(QueryContextShared::try_create( self.default_session.clone(), - Cluster::create(nodes, local_id), + Cluster::create(false, nodes, local_id), )?); dummy_query_context.get_settings().set_max_threads(8)?; diff --git a/src/query/storages/system/src/clusters_table.rs b/src/query/storages/system/src/clusters_table.rs index abe4a13baba19..ba0e038ffd9aa 100644 --- a/src/query/storages/system/src/clusters_table.rs +++ b/src/query/storages/system/src/clusters_table.rs @@ -16,7 +16,6 @@ use std::sync::Arc; use databend_common_catalog::table::Table; use databend_common_catalog::table_context::TableContext; -use databend_common_config::GlobalConfig; use databend_common_exception::Result; use databend_common_expression::types::number::NumberScalar; use databend_common_expression::types::DataType; @@ -57,12 +56,17 @@ impl SyncSystemTable for ClustersTable { ); let mut versions = ColumnBuilder::with_capacity(&DataType::String, cluster_nodes.len()); - let cluster_id = GlobalConfig::instance().query.cluster_id.clone(); for cluster_node in &cluster_nodes { let (ip, port) = cluster_node.ip_port()?; names.push(Scalar::String(cluster_node.id.clone()).as_ref()); - clusters.push(Scalar::String(cluster_id.clone()).as_ref()); + clusters.push( + Scalar::String(format!( + "{}/{}", + cluster_node.warehouse_id, cluster_node.cluster_id + )) + .as_ref(), + ); addresses.push(Scalar::String(ip).as_ref()); addresses_port.push(Scalar::Number(NumberScalar::UInt16(port)).as_ref()); versions.push(Scalar::String(cluster_node.binary_version.clone()).as_ref()); From b598f15ac18e7131a81b6539cd30317e23620c80 Mon Sep 17 00:00:00 2001 From: zhang2014 Date: Thu, 26 Dec 2024 20:57:53 +0800 Subject: [PATCH 24/53] feat(cluster): support custom management cluster --- .../management/src/warehouse/warehouse_api.rs | 2 + .../management/src/warehouse/warehouse_mgr.rs | 8 +++ src/query/service/src/clusters/cluster.rs | 56 ++++++++++--------- .../src/servers/http/middleware/session.rs | 4 +- 4 files changed, 43 insertions(+), 27 deletions(-) diff --git a/src/query/management/src/warehouse/warehouse_api.rs b/src/query/management/src/warehouse/warehouse_api.rs index ab93c0cb03b9c..69a9e2aefd2e0 100644 --- a/src/query/management/src/warehouse/warehouse_api.rs +++ b/src/query/management/src/warehouse/warehouse_api.rs @@ -109,4 +109,6 @@ pub trait WarehouseApi: Sync + Send { async fn list_online_nodes(&self) -> Result>; async fn discover(&self, node_id: &str) -> Result<(bool, Vec)>; + + async fn get_node_info(&self, node_id: &str) -> Result>; } diff --git a/src/query/management/src/warehouse/warehouse_mgr.rs b/src/query/management/src/warehouse/warehouse_mgr.rs index 57cf9781c8016..24ea64b351dec 100644 --- a/src/query/management/src/warehouse/warehouse_mgr.rs +++ b/src/query/management/src/warehouse/warehouse_mgr.rs @@ -1979,4 +1979,12 @@ impl WarehouseApi for WarehouseMgr { Ok((true, cluster_nodes_info)) } + + async fn get_node_info(&self, node_id: &str) -> Result> { + let node_key = format!("{}/{}", self.node_key_prefix, escape_for_key(node_id)?); + match self.metastore.get_kv(&node_key).await? { + None => Ok(None), + Some(seq) => Ok(Some(serde_json::from_slice(&seq.data)?)), + } + } } diff --git a/src/query/service/src/clusters/cluster.rs b/src/query/service/src/clusters/cluster.rs index 47cd817d3ed95..fc14fef64da4a 100644 --- a/src/query/service/src/clusters/cluster.rs +++ b/src/query/service/src/clusters/cluster.rs @@ -31,6 +31,9 @@ use databend_common_base::base::DummySignalStream; use databend_common_base::base::GlobalInstance; use databend_common_base::base::SignalStream; use databend_common_base::base::SignalType; +use databend_common_cache::Cache; +use databend_common_cache::LruCache; +use databend_common_cache::MemSized; pub use databend_common_catalog::cluster_info::Cluster; use databend_common_config::GlobalConfig; use databend_common_config::InnerConfig; @@ -52,7 +55,6 @@ use futures::StreamExt; use log::error; use log::info; use log::warn; -use parking_lot::RwLock; use rand::thread_rng; use rand::Rng; use serde::Deserialize; @@ -69,7 +71,7 @@ pub struct ClusterDiscovery { cluster_id: String, tenant_id: String, flight_address: String, - cached_cluster: RwLock>>, + lru_cache: parking_lot::Mutex>, } // avoid leak FlightClient to common-xxx @@ -240,7 +242,7 @@ impl ClusterDiscovery { cluster_id: cfg.query.cluster_id.clone(), tenant_id: cfg.query.tenant_id.tenant_name().to_string(), flight_address: cfg.query.flight_api_address.clone(), - cached_cluster: Default::default(), + lru_cache: parking_lot::Mutex::new(LruCache::with_items_capacity(100)), })) } @@ -311,35 +313,28 @@ impl ClusterDiscovery { ); let res = Cluster::create(!has_cluster, res, self.local_id.clone()); - // *self.cached_cluster.write() = Some(res.clone()); Ok(res) } } } - fn cached_cluster(self: &Arc) -> Option> { - (*self.cached_cluster.read()).clone() - } - - pub async fn find_node_by_id( - self: Arc, - id: &str, - config: &InnerConfig, - ) -> Result>> { - let (mut cluster, mut is_cached) = if let Some(cluster) = self.cached_cluster() { - (cluster, true) - } else { - (self.discover(config).await?, false) - }; - while is_cached { - for node in cluster.get_nodes() { - if node.id == id { - return Ok(Some(node.clone())); - } + pub async fn find_node_by_id(self: Arc, id: &str) -> Result>> { + { + let mut lru_cache = self.lru_cache.lock(); + if let Some(node_info) = lru_cache.get(id) { + return Ok(Some(node_info.node.clone())); } - cluster = self.discover(config).await?; - is_cached = false; } + + if let Some(node_info) = self.warehouse_manager.get_node_info(id).await? { + let cache_object = Arc::new(node_info); + let mut lru_cache = self.lru_cache.lock(); + lru_cache.insert(id.to_string(), CachedNode { + node: cache_object.clone(), + }); + return Ok(Some(cache_object)); + } + Ok(None) } @@ -601,3 +596,14 @@ pub struct FlightParams { pub(crate) retry_times: u64, pub(crate) retry_interval: u64, } + +#[derive(Clone)] +pub struct CachedNode { + pub node: Arc, +} + +impl MemSized for CachedNode { + fn mem_bytes(&self) -> usize { + 0 + } +} diff --git a/src/query/service/src/servers/http/middleware/session.rs b/src/query/service/src/servers/http/middleware/session.rs index 47834d27f6593..9bbadc60e49ce 100644 --- a/src/query/service/src/servers/http/middleware/session.rs +++ b/src/query/service/src/servers/http/middleware/session.rs @@ -517,9 +517,8 @@ impl Endpoint for HTTPSessionEndpoint { .to_string(); let local_id = GlobalConfig::instance().query.node_id.clone(); if local_id != sticky_node_id { - let config = GlobalConfig::instance(); return if let Some(node) = ClusterDiscovery::instance() - .find_node_by_id(&sticky_node_id, &config) + .find_node_by_id(&sticky_node_id) .await .map_err(HttpErrorCode::server_error)? { @@ -537,6 +536,7 @@ impl Endpoint for HTTPSessionEndpoint { }; } } + let method = req.method().clone(); let uri = req.uri().clone(); From bebdd21e2bc84b8d841334f5d2f8f7ddd1104b3e Mon Sep 17 00:00:00 2001 From: zhang2014 Date: Fri, 27 Dec 2024 09:49:07 +0800 Subject: [PATCH 25/53] feat(cluster): support custom management cluster --- .../resources_management/src/resources_management.rs | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/query/ee_features/resources_management/src/resources_management.rs b/src/query/ee_features/resources_management/src/resources_management.rs index 32da1e0742059..b617aa8b16823 100644 --- a/src/query/ee_features/resources_management/src/resources_management.rs +++ b/src/query/ee_features/resources_management/src/resources_management.rs @@ -13,6 +13,7 @@ // limitations under the License. use std::collections::HashMap; +use std::sync::Arc; use databend_common_base::base::GlobalInstance; use databend_common_config::GlobalConfig; @@ -146,7 +147,7 @@ impl ResourcesManagement for DummyResourcesManagement { impl DummyResourcesManagement { pub fn init() -> Result<()> { - let instance: Box = Box::new(DummyResourcesManagement); + let instance: Arc = Arc::new(DummyResourcesManagement); GlobalInstance::set(instance); Ok(()) } From 704fbf6d1aa16f1fbe0db1db45733267d686575d Mon Sep 17 00:00:00 2001 From: zhang2014 Date: Fri, 27 Dec 2024 14:32:55 +0800 Subject: [PATCH 26/53] feat(cluster): support custom management cluster --- .../databend-query-node-system-managed.toml | 1 + .../deploy/databend-query-system-managed.sh | 25 +++-- src/query/ast/src/ast/statements/statement.rs | 3 + src/query/ast/src/ast/statements/warehouse.rs | 15 ++- src/query/ast/src/parser/statement.rs | 12 ++- src/query/ast/src/parser/token.rs | 2 + .../resources_management_kubernetes.rs | 6 ++ .../resources_management_self_managed.rs | 6 ++ .../resources_management_system.rs | 4 + .../src/resources_management.rs | 6 ++ .../management/src/warehouse/warehouse_mgr.rs | 17 +++- .../interpreters/access/privilege_access.rs | 1 + .../src/interpreters/interpreter_factory.rs | 2 + .../interpreter_show_online_nodes.rs | 91 +++++++++++++++++++ src/query/service/src/interpreters/mod.rs | 1 + src/query/sql/src/planner/binder/binder.rs | 1 + .../sql/src/planner/binder/ddl/warehouse.rs | 8 ++ .../sql/src/planner/format/display_plan.rs | 1 + src/query/sql/src/planner/plans/plan.rs | 9 ++ 19 files changed, 198 insertions(+), 13 deletions(-) create mode 100644 src/query/service/src/interpreters/interpreter_show_online_nodes.rs diff --git a/scripts/ci/deploy/config/databend-query-node-system-managed.toml b/scripts/ci/deploy/config/databend-query-node-system-managed.toml index 368f7f2ceab3a..b39671a402325 100644 --- a/scripts/ci/deploy/config/databend-query-node-system-managed.toml +++ b/scripts/ci/deploy/config/databend-query-node-system-managed.toml @@ -52,6 +52,7 @@ definition = "CREATE FUNCTION ping(STRING) RETURNS STRING LANGUAGE python HANDLE [query.resources_management] type = "system_managed" +resource_group [log] diff --git a/scripts/ci/deploy/databend-query-system-managed.sh b/scripts/ci/deploy/databend-query-system-managed.sh index 946286d1fa97e..0510df4d959c8 100755 --- a/scripts/ci/deploy/databend-query-system-managed.sh +++ b/scripts/ci/deploy/databend-query-system-managed.sh @@ -8,9 +8,15 @@ SCRIPT_PATH="$(cd "$(dirname "$0")" >/dev/null 2>&1 && pwd)" cd "$SCRIPT_PATH/../../.." || exit BUILD_PROFILE=${BUILD_PROFILE:-debug} -if [ $# -ne 1 ]; then - echo "Usage: $0 - Start number of databend-query with system-managed mode" - exit 1 +if [ $# -eq 1 ]; then + num=$1 + resources_group="" +elif [ $# -eq 2 ]; then + num=$1 + resources_group=$2 +else + echo "Usage: $0 - Start number of databend-query with system-managed mode" + exit 1 fi if ! [[ "$num" =~ ^[0-9]*$ ]]; then @@ -60,8 +66,6 @@ python3 scripts/ci/wait_tcp.py --timeout 30 --port 28302 # wait for cluster formation to complete. sleep 1 -num=$1 - find_available_port() { local base_port=20000 local max_port=65535 @@ -84,6 +88,7 @@ start_databend_query() { local http_port=$1 local mysql_port=$2 local log_dir=$3 + local resources_group=$4 system_managed_config="./scripts/ci/deploy/config/databend-query-node-system-managed.toml" temp_file=$(mktemp) @@ -97,6 +102,7 @@ start_databend_query() { -e "s/http_port/${http_port}/g" \ -e "s/flight_sql_port/$(find_available_port)/g" \ -e "s/query_logs/${log_dir}/g" \ + -e "s/resource_group/resource_group=\"${resources_group}\"/g" \ "$system_managed_config" > "$temp_file" if [ $? -eq 0 ]; then @@ -116,11 +122,14 @@ start_databend_query() { fi } -start_databend_query 8000 3307 "logs_1" +if ! lsof -i :8000 >/dev/null 2>&1; then + start_databend_query 8000 3307 "logs_1" $resources_group + num=$(( num - 1 )) +fi -for (( i=1; i<$num; i++ )) +for (( i=0; i<$num; i++ )) do http_port=$(find_available_port) mysql_port=$(find_available_port) - start_databend_query $http_port $mysql_port "logs_$(( i + 1 ))" + start_databend_query $http_port $mysql_port "logs_$http_port" $resources_group done diff --git a/src/query/ast/src/ast/statements/statement.rs b/src/query/ast/src/ast/statements/statement.rs index 62e528300efe0..093c3da855418 100644 --- a/src/query/ast/src/ast/statements/statement.rs +++ b/src/query/ast/src/ast/statements/statement.rs @@ -130,6 +130,7 @@ pub enum Statement { }, // Warehouses + ShowOnlineNodes(ShowOnlineNodesStmt), ShowWarehouses(ShowWarehousesStmt), DropWarehouse(DropWarehouseStmt), CreateWarehouse(CreateWarehouseStmt), @@ -486,6 +487,7 @@ impl Statement { | Statement::DescProcedure(..) | Statement::CallProcedure(..) | Statement::ShowWarehouses(..) + | Statement::ShowOnlineNodes(..) | Statement::InspectWarehouse(..) => true, Statement::CreateDatabase(..) @@ -975,6 +977,7 @@ impl Display for Statement { Statement::System(stmt) => write!(f, "{stmt}")?, Statement::CallProcedure(stmt) => write!(f, "{stmt}")?, + Statement::ShowOnlineNodes(stmt) => write!(f, "{stmt}")?, Statement::ShowWarehouses(stmt) => write!(f, "{stmt}")?, Statement::DropWarehouse(stmt) => write!(f, "{stmt}")?, Statement::CreateWarehouse(stmt) => write!(f, "{stmt}")?, diff --git a/src/query/ast/src/ast/statements/warehouse.rs b/src/query/ast/src/ast/statements/warehouse.rs index 10d36c2cb6f29..4a48d10575b8a 100644 --- a/src/query/ast/src/ast/statements/warehouse.rs +++ b/src/query/ast/src/ast/statements/warehouse.rs @@ -21,6 +21,15 @@ use derive_visitor::DriveMut; use crate::ast::Identifier; +#[derive(Debug, Clone, PartialEq, Drive, DriveMut)] +pub struct ShowOnlineNodesStmt {} + +impl Display for ShowOnlineNodesStmt { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + write!(f, "SHOW ONLINE NODES") + } +} + #[derive(Debug, Clone, PartialEq, Drive, DriveMut)] pub struct ShowWarehousesStmt {} @@ -56,7 +65,7 @@ impl Display for CreateWarehouseStmt { } } - write!(f, "(")?; + write!(f, ")")?; } if !self.options.is_empty() { @@ -242,6 +251,8 @@ impl Display for AssignWarehouseNodesStmt { write!(f, " FOR {}", cluster)?; } + write!(f, ")")?; + Ok(()) } } @@ -270,6 +281,8 @@ impl Display for UnassignWarehouseNodesStmt { write!(f, " FOR {}", cluster)?; } + write!(f, ")")?; + Ok(()) } } diff --git a/src/query/ast/src/parser/statement.rs b/src/query/ast/src/parser/statement.rs index d38bbf796c75a..18369de413771 100644 --- a/src/query/ast/src/parser/statement.rs +++ b/src/query/ast/src/parser/statement.rs @@ -523,6 +523,13 @@ pub fn statement_body(i: Input) -> IResult { |(_, _, catalog)| Statement::UseCatalog { catalog }, ); + let show_online_nodes = map( + rule! { + SHOW ~ ONLINE ~ NODES + }, + |(_, _, _)| Statement::ShowOnlineNodes(ShowOnlineNodesStmt {}), + ); + let show_warehouses = map( rule! { SHOW ~ WAREHOUSES @@ -2415,6 +2422,7 @@ pub fn statement_body(i: Input) -> IResult { // warehouse rule!( #show_warehouses: "`SHOW WAREHOUSES`" + | #show_online_nodes: "`SHOW ONLINE NODES`" | #create_warehouse: "`CREATE WAREHOUSE [(ASSIGN NODES [FROM ] [, ...])] WITH [warehouse_size = ]`" | #drop_warehouse: "`DROP WAREHOUSE `" | #rename_warehouse: "`RENAME WAREHOUSE TO `" @@ -4092,9 +4100,9 @@ pub fn task_warehouse_option(i: Input) -> IResult { pub fn assign_nodes_list(i: Input) -> IResult, u64)>> { let nodes_list = map( rule! { - ASSIGN ~ #literal_u64 ~ (FROM ~ #option_to_string)? + ASSIGN ~ #literal_u64 ~ NODES ~ (FROM ~ #option_to_string)? }, - |(_, node_size, resources_group)| (resources_group.map(|(_, x)| x), node_size), + |(_, node_size, _, resources_group)| (resources_group.map(|(_, x)| x), node_size), ); map(comma_separated_list1(nodes_list), |opts| { diff --git a/src/query/ast/src/parser/token.rs b/src/query/ast/src/parser/token.rs index 018195899fb92..3cd2bcfbf195a 100644 --- a/src/query/ast/src/parser/token.rs +++ b/src/query/ast/src/parser/token.rs @@ -1388,6 +1388,8 @@ pub enum TokenKind { NODES, #[token("UNASSIGN", ignore(ascii_case))] UNASSIGN, + #[token("ONLINE", ignore(ascii_case))] + ONLINE, } // Reference: https://www.postgresql.org/docs/current/sql-keywords-appendix.html diff --git a/src/query/ee/src/resource_management/resources_management_kubernetes.rs b/src/query/ee/src/resource_management/resources_management_kubernetes.rs index fb35e4e04d58a..ea362f1346ae4 100644 --- a/src/query/ee/src/resource_management/resources_management_kubernetes.rs +++ b/src/query/ee/src/resource_management/resources_management_kubernetes.rs @@ -122,4 +122,10 @@ impl ResourcesManagement for KubernetesResourcesManagement { "Unimplemented kubernetes resources management", )) } + + async fn list_online_nodes(&self) -> Result> { + Err(ErrorCode::Unimplemented( + "Unimplemented kubernetes resources management", + )) + } } diff --git a/src/query/ee/src/resource_management/resources_management_self_managed.rs b/src/query/ee/src/resource_management/resources_management_self_managed.rs index 47a306b6c143c..fdd97fd6c14ff 100644 --- a/src/query/ee/src/resource_management/resources_management_self_managed.rs +++ b/src/query/ee/src/resource_management/resources_management_self_managed.rs @@ -133,4 +133,10 @@ impl ResourcesManagement for SelfManagedResourcesManagement { "Unimplemented unassign warehouse nodes with self-managed resources management", )) } + + async fn list_online_nodes(&self) -> Result> { + Err(ErrorCode::Unimplemented( + "Unimplemented list online nodes with self-managed resources management", + )) + } } diff --git a/src/query/ee/src/resource_management/resources_management_system.rs b/src/query/ee/src/resource_management/resources_management_system.rs index 66999cd83b027..ebf5cbedfb2ec 100644 --- a/src/query/ee/src/resource_management/resources_management_system.rs +++ b/src/query/ee/src/resource_management/resources_management_system.rs @@ -133,4 +133,8 @@ impl ResourcesManagement for SystemResourcesManagement { .unassign_warehouse_nodes(&name, nodes) .await } + + async fn list_online_nodes(&self) -> Result> { + self.warehouse_manager.list_online_nodes().await + } } diff --git a/src/query/ee_features/resources_management/src/resources_management.rs b/src/query/ee_features/resources_management/src/resources_management.rs index b617aa8b16823..6a010f5591692 100644 --- a/src/query/ee_features/resources_management/src/resources_management.rs +++ b/src/query/ee_features/resources_management/src/resources_management.rs @@ -69,6 +69,8 @@ pub trait ResourcesManagement: Sync + Send + 'static { name: String, nodes: HashMap>, ) -> Result<()>; + + async fn list_online_nodes(&self) -> Result>; } pub struct DummyResourcesManagement; @@ -143,6 +145,10 @@ impl ResourcesManagement for DummyResourcesManagement { ) -> Result<()> { Err(ErrorCode::Unimplemented("The use of this feature requires a Databend Enterprise Edition license. To unlock enterprise features, please contact Databend to obtain a license. Learn more at https://docs.databend.com/guides/overview/editions/dee/")) } + + async fn list_online_nodes(&self) -> Result> { + Err(ErrorCode::Unimplemented("The use of this feature requires a Databend Enterprise Edition license. To unlock enterprise features, please contact Databend to obtain a license. Learn more at https://docs.databend.com/guides/overview/editions/dee/")) + } } impl DummyResourcesManagement { diff --git a/src/query/management/src/warehouse/warehouse_mgr.rs b/src/query/management/src/warehouse/warehouse_mgr.rs index 24ea64b351dec..29fd004dc4af7 100644 --- a/src/query/management/src/warehouse/warehouse_mgr.rs +++ b/src/query/management/src/warehouse/warehouse_mgr.rs @@ -276,9 +276,14 @@ impl WarehouseMgr { async fn leave_cluster(&self, node_info: &mut NodeInfo, seq: u64) -> Result { let mut cluster_id = String::new(); let mut warehouse_id = String::new(); + let mut runtime_resource_group = None; std::mem::swap(&mut node_info.cluster_id, &mut cluster_id); std::mem::swap(&mut node_info.warehouse_id, &mut warehouse_id); + std::mem::swap( + &mut node_info.runtime_resource_group, + &mut runtime_resource_group, + ); let upsert_node = self.upsert_node(node_info.clone(), MatchSeq::Exact(seq)); match upsert_node.await { @@ -286,12 +291,20 @@ impl WarehouseMgr { // rollback std::mem::swap(&mut node_info.cluster_id, &mut cluster_id); std::mem::swap(&mut node_info.warehouse_id, &mut warehouse_id); + std::mem::swap( + &mut node_info.runtime_resource_group, + &mut runtime_resource_group, + ); Err(err) } Ok(response) if !response.success => { // rollback std::mem::swap(&mut node_info.cluster_id, &mut cluster_id); std::mem::swap(&mut node_info.warehouse_id, &mut warehouse_id); + std::mem::swap( + &mut node_info.runtime_resource_group, + &mut runtime_resource_group, + ); Ok(seq) } Ok(response) => match response.responses.last() { @@ -542,7 +555,7 @@ impl WarehouseMgr { } }; - break; + continue; }; let Some((seq, mut node)) = nodes_list.pop() else { @@ -556,7 +569,7 @@ impl WarehouseMgr { } }; - break; + continue; }; node.runtime_resource_group = None; diff --git a/src/query/service/src/interpreters/access/privilege_access.rs b/src/query/service/src/interpreters/access/privilege_access.rs index 490aa0e1c5912..9fd2967dd2825 100644 --- a/src/query/service/src/interpreters/access/privilege_access.rs +++ b/src/query/service/src/interpreters/access/privilege_access.rs @@ -1256,6 +1256,7 @@ impl AccessChecker for PrivilegeAccess { Plan::Commit => {} Plan::Abort => {} Plan::ShowWarehouses => {} + Plan::ShowOnlineNodes => {} Plan::DropWarehouse(_) => {} Plan::ResumeWarehouse(_) => {} Plan::SuspendWarehouse(_) => {} diff --git a/src/query/service/src/interpreters/interpreter_factory.rs b/src/query/service/src/interpreters/interpreter_factory.rs index a615fa6acbc9e..7610e4acb44da 100644 --- a/src/query/service/src/interpreters/interpreter_factory.rs +++ b/src/query/service/src/interpreters/interpreter_factory.rs @@ -67,6 +67,7 @@ use crate::interpreters::interpreter_rename_warehouse_cluster::RenameWarehouseCl use crate::interpreters::interpreter_resume_warehouse::ResumeWarehouseInterpreter; use crate::interpreters::interpreter_role_show::ShowRolesInterpreter; use crate::interpreters::interpreter_set_priority::SetPriorityInterpreter; +use crate::interpreters::interpreter_show_online_nodes::ShowOnlineNodesInterpreter; use crate::interpreters::interpreter_show_warehouses::ShowWarehousesInterpreter; use crate::interpreters::interpreter_suspend_warehouse::SuspendWarehouseInterpreter; use crate::interpreters::interpreter_system_action::SystemActionInterpreter; @@ -118,6 +119,7 @@ impl InterpreterFactory { ) -> Result { match plan { Plan::ShowWarehouses => Ok(Arc::new(ShowWarehousesInterpreter::try_create()?)), + Plan::ShowOnlineNodes => Ok(Arc::new(ShowOnlineNodesInterpreter::try_create()?)), Plan::CreateWarehouse(v) => Ok(Arc::new(CreateWarehouseInterpreter::try_create( ctx.clone(), *v.clone(), diff --git a/src/query/service/src/interpreters/interpreter_show_online_nodes.rs b/src/query/service/src/interpreters/interpreter_show_online_nodes.rs new file mode 100644 index 0000000000000..674b2cbcff4e4 --- /dev/null +++ b/src/query/service/src/interpreters/interpreter_show_online_nodes.rs @@ -0,0 +1,91 @@ +// Copyright 2021 Datafuse Labs +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::sync::Arc; + +use databend_common_base::base::GlobalInstance; +use databend_common_exception::Result; +use databend_common_expression::types::DataType; +use databend_common_expression::ColumnBuilder; +use databend_common_expression::DataBlock; +use databend_common_expression::Scalar; +use databend_common_meta_types::NodeType; +use databend_enterprise_resources_management::ResourcesManagement; + +use crate::interpreters::Interpreter; +use crate::pipelines::PipelineBuildResult; + +pub struct ShowOnlineNodesInterpreter {} + +impl ShowOnlineNodesInterpreter { + pub fn try_create() -> Result { + Ok(ShowOnlineNodesInterpreter {}) + } +} + +#[async_trait::async_trait] +impl Interpreter for ShowOnlineNodesInterpreter { + fn name(&self) -> &str { + "ShowWarehousesInterpreter" + } + + fn is_ddl(&self) -> bool { + false + } + + #[async_backtrace::framed] + async fn execute2(&self) -> Result { + let online_nodes = GlobalInstance::get::>() + .list_online_nodes() + .await?; + + let mut nodes_id = ColumnBuilder::with_capacity(&DataType::String, online_nodes.len()); + let mut nodes_type = ColumnBuilder::with_capacity(&DataType::String, online_nodes.len()); + let mut nodes_resource_group = + ColumnBuilder::with_capacity(&DataType::String, online_nodes.len()); + let mut nodes_warehouse = + ColumnBuilder::with_capacity(&DataType::String, online_nodes.len()); + let mut nodes_cluster = ColumnBuilder::with_capacity(&DataType::String, online_nodes.len()); + let mut nodes_version = ColumnBuilder::with_capacity(&DataType::String, online_nodes.len()); + + for node in online_nodes { + let node_type = match node.node_type { + NodeType::SelfManaged => String::from("SelfManaged"), + NodeType::SystemManaged => String::from("SystemManaged"), + }; + + let binary_version = match node.binary_version.split_once('(') { + None => node.binary_version, + Some((left, _right)) => left.to_string(), + }; + + nodes_id.push(Scalar::String(node.id).as_ref()); + nodes_type.push(Scalar::String(node_type).as_ref()); + nodes_resource_group + .push(Scalar::String(node.resource_group.clone().unwrap_or_default()).as_ref()); + nodes_warehouse.push(Scalar::String(node.warehouse_id).as_ref()); + nodes_cluster.push(Scalar::String(node.cluster_id).as_ref()); + nodes_version.push(Scalar::String(binary_version).as_ref()); + } + + PipelineBuildResult::from_blocks(vec![DataBlock::new_from_columns(vec![ + nodes_id.build(), + nodes_type.build(), + nodes_resource_group.build(), + nodes_warehouse.build(), + nodes_cluster.build(), + nodes_version.build(), + ])]) + } +} diff --git a/src/query/service/src/interpreters/mod.rs b/src/query/service/src/interpreters/mod.rs index b7b9ad71650e9..9cb860f828b02 100644 --- a/src/query/service/src/interpreters/mod.rs +++ b/src/query/service/src/interpreters/mod.rs @@ -100,6 +100,7 @@ mod interpreter_sequence_create; mod interpreter_sequence_drop; mod interpreter_set; mod interpreter_set_priority; +mod interpreter_show_online_nodes; mod interpreter_show_warehouses; mod interpreter_stream_create; mod interpreter_stream_drop; diff --git a/src/query/sql/src/planner/binder/binder.rs b/src/query/sql/src/planner/binder/binder.rs index 5fdb319688f94..976d8b1d2edf3 100644 --- a/src/query/sql/src/planner/binder/binder.rs +++ b/src/query/sql/src/planner/binder/binder.rs @@ -656,6 +656,7 @@ impl<'a> Binder { return Err(ErrorCode::SyntaxException("CALL PROCEDURE, set enable_experimental_procedure=1")); } } + Statement::ShowOnlineNodes(v) => self.bind_show_online_nodes(v)?, Statement::ShowWarehouses(v) => self.bind_show_warehouses(v)?, Statement::DropWarehouse(v) => self.bind_drop_warehouse(v)?, Statement::CreateWarehouse(v) => self.bind_create_warehouse(v)?, diff --git a/src/query/sql/src/planner/binder/ddl/warehouse.rs b/src/query/sql/src/planner/binder/ddl/warehouse.rs index feb6f09f5d543..9a36ca6f18359 100644 --- a/src/query/sql/src/planner/binder/ddl/warehouse.rs +++ b/src/query/sql/src/planner/binder/ddl/warehouse.rs @@ -24,6 +24,7 @@ use databend_common_ast::ast::InspectWarehouseStmt; use databend_common_ast::ast::RenameWarehouseClusterStmt; use databend_common_ast::ast::RenameWarehouseStmt; use databend_common_ast::ast::ResumeWarehouseStmt; +use databend_common_ast::ast::ShowOnlineNodesStmt; use databend_common_ast::ast::ShowWarehousesStmt; use databend_common_ast::ast::SuspendWarehouseStmt; use databend_common_ast::ast::UnassignWarehouseNodesStmt; @@ -44,6 +45,13 @@ use crate::plans::UnassignWarehouseNodesPlan; use crate::Binder; impl Binder { + pub(in crate::planner::binder) fn bind_show_online_nodes( + &mut self, + _stmt: &ShowOnlineNodesStmt, + ) -> Result { + Ok(Plan::ShowOnlineNodes) + } + pub(in crate::planner::binder) fn bind_show_warehouses( &mut self, _stmt: &ShowWarehousesStmt, diff --git a/src/query/sql/src/planner/format/display_plan.rs b/src/query/sql/src/planner/format/display_plan.rs index cb451fa86a4e9..a88171e5548a6 100644 --- a/src/query/sql/src/planner/format/display_plan.rs +++ b/src/query/sql/src/planner/format/display_plan.rs @@ -214,6 +214,7 @@ impl Plan { Plan::ShowCreateDictionary(_) => Ok("ShowCreateDictionary".to_string()), Plan::RenameDictionary(_) => Ok("RenameDictionary".to_string()), Plan::ShowWarehouses => Ok("ShowWarehouses".to_string()), + Plan::ShowOnlineNodes => Ok("ShowOnlineNodes".to_string()), Plan::DropWarehouse(_) => Ok("DropWarehouse".to_string()), Plan::ResumeWarehouse(_) => Ok("ResumeWarehouse".to_string()), Plan::SuspendWarehouse(_) => Ok("SuspendWarehouse".to_string()), diff --git a/src/query/sql/src/planner/plans/plan.rs b/src/query/sql/src/planner/plans/plan.rs index 96162ee648eb2..1523b7b281251 100644 --- a/src/query/sql/src/planner/plans/plan.rs +++ b/src/query/sql/src/planner/plans/plan.rs @@ -203,6 +203,7 @@ pub enum Plan { UseCatalog(Box), // Warehouses + ShowOnlineNodes, ShowWarehouses, CreateWarehouse(Box), DropWarehouse(Box), @@ -526,6 +527,14 @@ impl Plan { DataField::new("type", DataType::String), DataField::new("status", DataType::String), ]), + Plan::ShowOnlineNodes => DataSchemaRefExt::create(vec![ + DataField::new("id", DataType::String), + DataField::new("type", DataType::String), + DataField::new("resource_group", DataType::String), + DataField::new("warehouse", DataType::String), + DataField::new("cluster", DataType::String), + DataField::new("version", DataType::String), + ]), _ => Arc::new(DataSchema::empty()), } From f6c40252f27c906a0bb4b56860042c457f856189 Mon Sep 17 00:00:00 2001 From: zhang2014 Date: Sun, 29 Dec 2024 14:17:51 +0800 Subject: [PATCH 27/53] feat(cluster): support custom management cluster --- .../databend-query-node-system-managed.toml | 2 +- .../deploy/databend-query-system-managed.sh | 12 +-- src/common/license/src/license.rs | 3 + src/meta/types/src/cluster.rs | 8 +- src/meta/types/tests/it/cluster.rs | 8 +- src/query/ast/src/ast/statements/warehouse.rs | 12 +-- src/query/ast/src/parser/statement.rs | 18 ++-- .../ast/tests/it/testdata/stmt-error.txt | 6 +- src/query/config/src/config.rs | 2 +- src/query/ee/src/resource_management/mod.rs | 5 +- .../resources_management_system.rs | 2 +- .../management/src/warehouse/warehouse_mgr.rs | 60 ++++++-------- src/query/management/tests/it/warehouse.rs | 8 +- .../interpreter_add_warehouse_cluster.rs | 15 +++- .../interpreter_assign_warehouse_nodes.rs | 21 ++++- .../interpreter_create_warehouses.rs | 14 ++++ .../interpreter_drop_warehouse_cluster.rs | 15 ++++ ...interpreter_drop_warehouse_cluster_node.rs | 49 ----------- .../interpreter_drop_warehouses.rs | 15 ++++ .../src/interpreters/interpreter_factory.rs | 15 ++-- .../interpreter_inspect_warehouse.rs | 12 ++- .../interpreter_rename_warehouse.rs | 15 ++++ .../interpreter_rename_warehouse_cluster.rs | 21 ++++- .../interpreter_resume_warehouse.rs | 21 ++++- .../interpreter_show_online_nodes.rs | 23 +++-- .../interpreter_show_warehouses.rs | 15 +++- .../interpreter_suspend_warehouse.rs | 21 ++++- .../interpreter_unassign_warehouse_nodes.rs | 83 +++++++++++++++++++ src/query/service/src/interpreters/mod.rs | 2 +- src/query/service/src/interpreters/util.rs | 20 +++++ src/query/service/tests/it/clusters.rs | 2 + .../storages/testdata/configs_table_basic.txt | 1 + .../sql/src/planner/plans/ddl/warehouse.rs | 20 ++--- src/query/sql/src/planner/plans/plan.rs | 2 +- 34 files changed, 384 insertions(+), 164 deletions(-) delete mode 100644 src/query/service/src/interpreters/interpreter_drop_warehouse_cluster_node.rs create mode 100644 src/query/service/src/interpreters/interpreter_unassign_warehouse_nodes.rs diff --git a/scripts/ci/deploy/config/databend-query-node-system-managed.toml b/scripts/ci/deploy/config/databend-query-node-system-managed.toml index b39671a402325..b9937839479e1 100644 --- a/scripts/ci/deploy/config/databend-query-node-system-managed.toml +++ b/scripts/ci/deploy/config/databend-query-node-system-managed.toml @@ -52,7 +52,7 @@ definition = "CREATE FUNCTION ping(STRING) RETURNS STRING LANGUAGE python HANDLE [query.resources_management] type = "system_managed" -resource_group +node_group [log] diff --git a/scripts/ci/deploy/databend-query-system-managed.sh b/scripts/ci/deploy/databend-query-system-managed.sh index 0510df4d959c8..4701daaa58c06 100755 --- a/scripts/ci/deploy/databend-query-system-managed.sh +++ b/scripts/ci/deploy/databend-query-system-managed.sh @@ -10,10 +10,10 @@ BUILD_PROFILE=${BUILD_PROFILE:-debug} if [ $# -eq 1 ]; then num=$1 - resources_group="" + node_group="" elif [ $# -eq 2 ]; then num=$1 - resources_group=$2 + node_group=$2 else echo "Usage: $0 - Start number of databend-query with system-managed mode" exit 1 @@ -88,7 +88,7 @@ start_databend_query() { local http_port=$1 local mysql_port=$2 local log_dir=$3 - local resources_group=$4 + local node_group=$4 system_managed_config="./scripts/ci/deploy/config/databend-query-node-system-managed.toml" temp_file=$(mktemp) @@ -102,7 +102,7 @@ start_databend_query() { -e "s/http_port/${http_port}/g" \ -e "s/flight_sql_port/$(find_available_port)/g" \ -e "s/query_logs/${log_dir}/g" \ - -e "s/resource_group/resource_group=\"${resources_group}\"/g" \ + -e "s/node_group/node_group=\"${node_group}\"/g" \ "$system_managed_config" > "$temp_file" if [ $? -eq 0 ]; then @@ -123,7 +123,7 @@ start_databend_query() { } if ! lsof -i :8000 >/dev/null 2>&1; then - start_databend_query 8000 3307 "logs_1" $resources_group + start_databend_query 8000 3307 "logs_1" $node_group num=$(( num - 1 )) fi @@ -131,5 +131,5 @@ for (( i=0; i<$num; i++ )) do http_port=$(find_available_port) mysql_port=$(find_available_port) - start_databend_query $http_port $mysql_port "logs_$http_port" $resources_group + start_databend_query $http_port $mysql_port "logs_$http_port" $node_group done diff --git a/src/common/license/src/license.rs b/src/common/license/src/license.rs index b2d9769e2ffca..219a903498acb 100644 --- a/src/common/license/src/license.rs +++ b/src/common/license/src/license.rs @@ -73,6 +73,8 @@ pub enum Feature { StorageQuota(StorageQuota), #[serde(alias = "amend_table", alias = "AMEND_TABLE")] AmendTable, + #[serde(alias = "system_management", alias = "SYSTEM_MANAGEMENT")] + SystemManagement, #[serde(other)] Unknown, } @@ -119,6 +121,7 @@ impl fmt::Display for Feature { write!(f, ")") } Feature::AmendTable => write!(f, "amend_table"), + Feature::SystemManagement => write!(f, "system_management"), Feature::Unknown => write!(f, "unknown"), } } diff --git a/src/meta/types/src/cluster.rs b/src/meta/types/src/cluster.rs index 90d122038d062..1162c648e0525 100644 --- a/src/meta/types/src/cluster.rs +++ b/src/meta/types/src/cluster.rs @@ -90,14 +90,14 @@ pub struct NodeInfo { pub discovery_address: String, pub binary_version: String, pub node_type: NodeType, - pub resource_group: Option, + pub node_group: Option, #[serde(skip_serializing_if = "String::is_empty")] pub cluster_id: String, #[serde(skip_serializing_if = "String::is_empty")] pub warehouse_id: String, - pub runtime_resource_group: Option, + pub runtime_node_group: Option, } impl NodeInfo { @@ -122,8 +122,8 @@ impl NodeInfo { cluster_id: "".to_string(), warehouse_id: "".to_string(), node_type: NodeType::SystemManaged, - resource_group: None, - runtime_resource_group: None, + node_group: None, + runtime_node_group: None, } } diff --git a/src/meta/types/tests/it/cluster.rs b/src/meta/types/tests/it/cluster.rs index f7a546c3aba31..2ecf109e408ce 100644 --- a/src/meta/types/tests/it/cluster.rs +++ b/src/meta/types/tests/it/cluster.rs @@ -26,10 +26,10 @@ fn test_node_info_ip_port() -> anyhow::Result<()> { discovery_address: "4.5.6.7:456".to_string(), binary_version: "v0.8-binary-version".to_string(), node_type: Default::default(), - resource_group: None, + node_group: None, cluster_id: "".to_string(), warehouse_id: "".to_string(), - runtime_resource_group: None, + runtime_node_group: None, }; let (ip, port) = n.ip_port()?; @@ -52,10 +52,10 @@ fn test_serde_node_info() { discovery_address: "4.5.6.7:456".to_string(), binary_version: "v0.8-binary-version".to_string(), node_type: Default::default(), - resource_group: None, + node_group: None, cluster_id: String::new(), warehouse_id: String::new(), - runtime_resource_group: None, + runtime_node_group: None, }; let json_str = serde_json::to_string(&info).unwrap(); diff --git a/src/query/ast/src/ast/statements/warehouse.rs b/src/query/ast/src/ast/statements/warehouse.rs index 4a48d10575b8a..af812a9abd3ac 100644 --- a/src/query/ast/src/ast/statements/warehouse.rs +++ b/src/query/ast/src/ast/statements/warehouse.rs @@ -53,15 +53,15 @@ impl Display for CreateWarehouseStmt { if !self.node_list.is_empty() { write!(f, "(")?; - for (idx, (resources_group, nodes)) in self.node_list.iter().enumerate() { + for (idx, (node_group, nodes)) in self.node_list.iter().enumerate() { if idx != 0 { write!(f, ",")?; } write!(f, " ASSIGN {} NODES", nodes)?; - if let Some(resources_group) = resources_group { - write!(f, " FROM '{}'", resources_group)?; + if let Some(node_group) = node_group { + write!(f, " FROM '{}'", node_group)?; } } @@ -163,15 +163,15 @@ impl Display for AddWarehouseClusterStmt { if !self.node_list.is_empty() { write!(f, "(")?; - for (idx, (resources_group, nodes)) in self.node_list.iter().enumerate() { + for (idx, (node_group, nodes)) in self.node_list.iter().enumerate() { if idx != 0 { write!(f, ",")?; } write!(f, " ASSIGN {} NODES", nodes)?; - if let Some(resources_group) = resources_group { - write!(f, " FROM '{}'", resources_group)?; + if let Some(node_group) = node_group { + write!(f, " FROM '{}'", node_group)?; } } diff --git a/src/query/ast/src/parser/statement.rs b/src/query/ast/src/parser/statement.rs index 18369de413771..2aa9da2ac37db 100644 --- a/src/query/ast/src/parser/statement.rs +++ b/src/query/ast/src/parser/statement.rs @@ -2423,17 +2423,17 @@ pub fn statement_body(i: Input) -> IResult { rule!( #show_warehouses: "`SHOW WAREHOUSES`" | #show_online_nodes: "`SHOW ONLINE NODES`" - | #create_warehouse: "`CREATE WAREHOUSE [(ASSIGN NODES [FROM ] [, ...])] WITH [warehouse_size = ]`" + | #create_warehouse: "`CREATE WAREHOUSE [(ASSIGN NODES [FROM ] [, ...])] WITH [warehouse_size = ]`" | #drop_warehouse: "`DROP WAREHOUSE `" | #rename_warehouse: "`RENAME WAREHOUSE TO `" | #resume_warehouse: "`RESUME WAREHOUSE `" | #suspend_warehouse: "`SUSPEND WAREHOUSE `" | #inspect_warehouse: "`INSPECT WAREHOUSE `" - | #add_warehouse_cluster: "`ALTER WAREHOUSE ADD CLUSTER [(ASSIGN NODES [FROM ] [, ...])] WITH [cluster_size = ]`" + | #add_warehouse_cluster: "`ALTER WAREHOUSE ADD CLUSTER [(ASSIGN NODES [FROM ] [, ...])] WITH [cluster_size = ]`" | #drop_warehouse_cluster: "`ALTER WAREHOUSE DROP CLUSTER `" | #rename_warehouse_cluster: "`ALTER WAREHOUSE RENAME CLUSTER TO `" - | #assign_warehouse_nodes: "`ALTER WAREHOUSE ASSIGN NODES ( ASSIGN NODES [FROM ] FOR [, ...] )`" - | #unassign_warehouse_nodes: "`ALTER WAREHOUSE UNASSIGN NODES ( UNASSIGN NODES [FROM ] FOR [, ...] )`" + | #assign_warehouse_nodes: "`ALTER WAREHOUSE ASSIGN NODES ( ASSIGN NODES [FROM ] FOR [, ...] )`" + | #unassign_warehouse_nodes: "`ALTER WAREHOUSE UNASSIGN NODES ( UNASSIGN NODES [FROM ] FOR [, ...] )`" ), // database rule!( @@ -4102,7 +4102,7 @@ pub fn assign_nodes_list(i: Input) -> IResult, u64)>> { rule! { ASSIGN ~ #literal_u64 ~ NODES ~ (FROM ~ #option_to_string)? }, - |(_, node_size, _, resources_group)| (resources_group.map(|(_, x)| x), node_size), + |(_, node_size, _, node_group)| (node_group.map(|(_, x)| x), node_size), ); map(comma_separated_list1(nodes_list), |opts| { @@ -4115,8 +4115,8 @@ pub fn assign_warehouse_nodes_list(i: Input) -> IResult IResult SQL:1:6 | 1 | drop a - | ^ unexpected `a`, expecting `TASK`, `TABLE`, `MASKING`, `CATALOG`, `DATABASE`, `PASSWORD`, `AGGREGATING`, `SCHEMA`, `NETWORK`, `USER`, `ROLE`, `FUNCTION`, `DICTIONARY`, `VIEW`, `INVERTED`, `VIRTUAL`, `SEQUENCE`, `STAGE`, `FILE`, `STREAM`, `PIPE`, `NOTIFICATION`, `CONNECTION`, or `PROCEDURE` + | ^ unexpected `a`, expecting `TASK`, `TABLE`, `MASKING`, `CATALOG`, `DATABASE`, `PASSWORD`, `WAREHOUSE`, `AGGREGATING`, `SCHEMA`, `NETWORK`, `USER`, `ROLE`, `FUNCTION`, `DICTIONARY`, `VIEW`, `INVERTED`, `VIRTUAL`, `SEQUENCE`, `STAGE`, `FILE`, `STREAM`, `PIPE`, `NOTIFICATION`, `CONNECTION`, or `PROCEDURE` ---------- Input ---------- @@ -198,7 +198,7 @@ error: --> SQL:1:6 | 1 | drop usar if exists 'test-j'; - | ^^^^ unexpected `usar`, expecting `USER`, `STREAM`, `STAGE`, `PASSWORD`, `AGGREGATING`, `ROLE`, `TABLE`, `SCHEMA`, `NETWORK`, `VIRTUAL`, `CATALOG`, `DATABASE`, `FUNCTION`, `INVERTED`, `PROCEDURE`, `TASK`, `NOTIFICATION`, `MASKING`, `SEQUENCE`, `DICTIONARY`, `VIEW`, `FILE`, `PIPE`, or `CONNECTION` + | ^^^^ unexpected `usar`, expecting `USER`, `STREAM`, `STAGE`, `PASSWORD`, `WAREHOUSE`, `AGGREGATING`, `ROLE`, `TABLE`, `SCHEMA`, `NETWORK`, `VIRTUAL`, `CATALOG`, `DATABASE`, `FUNCTION`, `INVERTED`, `PROCEDURE`, `TASK`, `NOTIFICATION`, `MASKING`, `SEQUENCE`, `DICTIONARY`, `VIEW`, `FILE`, `PIPE`, or `CONNECTION` ---------- Input ---------- @@ -238,7 +238,7 @@ error: --> SQL:1:6 | 1 | SHOW GRANT FOR ROLE 'role1'; - | ^^^^^ unexpected `GRANT`, expecting `GRANTS`, `CREATE`, `NETWORK`, `VIRTUAL`, `CATALOGS`, `STREAMS`, `FUNCTIONS`, `DATABASES`, `CONNECTIONS`, `TABLE_FUNCTIONS`, `DROP`, `ROLES`, `TABLE`, `TASKS`, `INDEXES`, `COLUMNS`, `PASSWORD`, `PROCEDURES`, `PROCESSLIST`, `STAGES`, `TABLES`, `DICTIONARIES`, `ENGINES`, `METRICS`, `SETTINGS`, `VARIABLES`, `LOCKS`, `SCHEMAS`, `USERS`, `USER`, `FIELDS`, `VIEWS`, `FILE`, or `FULL` + | ^^^^^ unexpected `GRANT`, expecting `GRANTS`, `CREATE`, `NETWORK`, `VIRTUAL`, `CATALOGS`, `STREAMS`, `FUNCTIONS`, `DATABASES`, `CONNECTIONS`, `TABLE_FUNCTIONS`, `DROP`, `ROLES`, `TABLE`, `TASKS`, `ONLINE`, `INDEXES`, `COLUMNS`, `PASSWORD`, `PROCEDURES`, `PROCESSLIST`, `STAGES`, `TABLES`, `DICTIONARIES`, `ENGINES`, `METRICS`, `SETTINGS`, `VARIABLES`, `WAREHOUSES`, `LOCKS`, `SCHEMAS`, `USERS`, `USER`, `FIELDS`, `VIEWS`, `FILE`, or `FULL` ---------- Input ---------- diff --git a/src/query/config/src/config.rs b/src/query/config/src/config.rs index ebab9f5daba3f..88781b68644f8 100644 --- a/src/query/config/src/config.rs +++ b/src/query/config/src/config.rs @@ -2993,7 +2993,7 @@ pub struct ResourcesManagementConfig { pub typ: String, #[clap(long, value_name = "VALUE")] - pub resource_group: Option, + pub node_group: Option, } mod cache_config_converters { diff --git a/src/query/ee/src/resource_management/mod.rs b/src/query/ee/src/resource_management/mod.rs index ee0a7877d15e3..efd48e0a84fe6 100644 --- a/src/query/ee/src/resource_management/mod.rs +++ b/src/query/ee/src/resource_management/mod.rs @@ -26,11 +26,10 @@ use databend_common_exception::Result; use databend_common_management::WarehouseMgr; use databend_common_meta_store::MetaStoreProvider; use databend_enterprise_resources_management::ResourcesManagement; +pub use resources_management_kubernetes::KubernetesResourcesManagement; +pub use resources_management_self_managed::SelfManagedResourcesManagement; pub use resources_management_system::SystemResourcesManagement; -use crate::resource_management::resources_management_kubernetes::KubernetesResourcesManagement; -use crate::resource_management::resources_management_self_managed::SelfManagedResourcesManagement; - pub async fn init_resources_management(cfg: &InnerConfig) -> Result<()> { let service: Arc = match &cfg.query.resources_management { None => match cfg.query.cluster_id.is_empty() { diff --git a/src/query/ee/src/resource_management/resources_management_system.rs b/src/query/ee/src/resource_management/resources_management_system.rs index ebf5cbedfb2ec..2b929a8f024b0 100644 --- a/src/query/ee/src/resource_management/resources_management_system.rs +++ b/src/query/ee/src/resource_management/resources_management_system.rs @@ -52,7 +52,7 @@ impl ResourcesManagement for SystemResourcesManagement { node.cluster_id = String::new(); node.warehouse_id = String::new(); node.node_type = NodeType::SystemManaged; - node.resource_group = resources_management.resource_group.clone(); + node.node_group = resources_management.node_group.clone(); } Ok(()) diff --git a/src/query/management/src/warehouse/warehouse_mgr.rs b/src/query/management/src/warehouse/warehouse_mgr.rs index 29fd004dc4af7..5d7825985a876 100644 --- a/src/query/management/src/warehouse/warehouse_mgr.rs +++ b/src/query/management/src/warehouse/warehouse_mgr.rs @@ -276,14 +276,11 @@ impl WarehouseMgr { async fn leave_cluster(&self, node_info: &mut NodeInfo, seq: u64) -> Result { let mut cluster_id = String::new(); let mut warehouse_id = String::new(); - let mut runtime_resource_group = None; + let mut runtime_node_group = None; std::mem::swap(&mut node_info.cluster_id, &mut cluster_id); std::mem::swap(&mut node_info.warehouse_id, &mut warehouse_id); - std::mem::swap( - &mut node_info.runtime_resource_group, - &mut runtime_resource_group, - ); + std::mem::swap(&mut node_info.runtime_node_group, &mut runtime_node_group); let upsert_node = self.upsert_node(node_info.clone(), MatchSeq::Exact(seq)); match upsert_node.await { @@ -291,20 +288,14 @@ impl WarehouseMgr { // rollback std::mem::swap(&mut node_info.cluster_id, &mut cluster_id); std::mem::swap(&mut node_info.warehouse_id, &mut warehouse_id); - std::mem::swap( - &mut node_info.runtime_resource_group, - &mut runtime_resource_group, - ); + std::mem::swap(&mut node_info.runtime_node_group, &mut runtime_node_group); Err(err) } Ok(response) if !response.success => { // rollback std::mem::swap(&mut node_info.cluster_id, &mut cluster_id); std::mem::swap(&mut node_info.warehouse_id, &mut warehouse_id); - std::mem::swap( - &mut node_info.runtime_resource_group, - &mut runtime_resource_group, - ); + std::mem::swap(&mut node_info.runtime_node_group, &mut runtime_node_group); Ok(seq) } Ok(response) => match response.responses.last() { @@ -515,7 +506,7 @@ impl WarehouseMgr { let node_info = serde_json::from_slice::(&seq_data.data)?; if node_info.cluster_id.is_empty() && node_info.warehouse_id.is_empty() { - match group_nodes.entry(node_info.resource_group.clone()) { + match group_nodes.entry(node_info.node_group.clone()) { Entry::Vacant(v) => { v.insert(vec![(seq_data.seq, node_info)]); } @@ -572,17 +563,17 @@ impl WarehouseMgr { continue; }; - node.runtime_resource_group = None; + node.runtime_node_group = None; node.cluster_id = cluster.clone(); node.warehouse_id = warehouse.to_string(); cluster_selected_nodes.push((seq, node)); } - SelectedNode::Random(Some(resource_group)) => { - let key = Some(resource_group.clone()); + SelectedNode::Random(Some(node_group)) => { + let key = Some(node_group.clone()); let Some(nodes_list) = grouped_nodes.get_mut(&key) else { return Err(ErrorCode::NoResourcesAvailable(format!( "Failed to create warehouse, reason: no resources available for {} group", - resource_group + node_group ))); }; @@ -590,13 +581,13 @@ impl WarehouseMgr { grouped_nodes.remove(&key); return Err(ErrorCode::NoResourcesAvailable(format!( "Failed to create warehouse, reason: no resources available for {} group", - resource_group + node_group ))); }; node.cluster_id = cluster.clone(); node.warehouse_id = warehouse.to_string(); - node.runtime_resource_group = Some(resource_group.clone()); + node.runtime_node_group = Some(node_group.clone()); cluster_selected_nodes.push((seq, node)); } } @@ -629,7 +620,7 @@ impl WarehouseMgr { node.cluster_id = cluster.clone(); node.warehouse_id = warehouse.to_string(); - node.runtime_resource_group = None; + node.runtime_node_group = None; selected_nodes.get_mut(&cluster).unwrap().push((seq, node)); } } @@ -1487,15 +1478,18 @@ impl WarehouseApi for WarehouseMgr { WarehouseInfo::SelfManaged(_) => Err(ErrorCode::InvalidWarehouse(format!("Cannot rename cluster for warehouse {:?}, because it's self-managed warehouse.", warehouse))), WarehouseInfo::SystemManaged(mut info) => match info.clusters.contains_key(&cur) { false => Err(ErrorCode::WarehouseClusterNotExists(format!("Warehouse cluster {:?}.{:?} not exists", warehouse, cur))), - true => { - let cluster_info = info.clusters.remove(&cur); - info.clusters.insert(to.clone(), cluster_info.unwrap()); - Ok(WarehouseInfo::SystemManaged(SystemManagedWarehouse { - id: info.id, - status: info.status, - display_name: info.display_name, - clusters: info.clusters, - })) + true => match info.clusters.contains_key(&to) { + true => Err(ErrorCode::WarehouseClusterAlreadyExists(format!("Warehouse cluster {:?}.{:?} already exists", warehouse, to))), + false => { + let cluster_info = info.clusters.remove(&cur); + info.clusters.insert(to.clone(), cluster_info.unwrap()); + Ok(WarehouseInfo::SystemManaged(SystemManagedWarehouse { + id: info.id, + status: info.status, + display_name: info.display_name, + clusters: info.clusters, + })) + } } } }?; @@ -1776,11 +1770,11 @@ impl WarehouseApi for WarehouseMgr { if let Some(v) = nodes.get_mut(&consistent_node.node_info.cluster_id) { if let Some(remove_node) = v.pop() { - let SelectedNode::Random(resource_group) = remove_node; - if consistent_node.node_info.runtime_resource_group == resource_group { + let SelectedNode::Random(node_group) = remove_node; + if consistent_node.node_info.runtime_node_group == node_group { consistent_node.node_info.cluster_id = String::new(); consistent_node.node_info.warehouse_id = String::new(); - consistent_node.node_info.runtime_resource_group = None; + consistent_node.node_info.runtime_node_group = None; drop_cluster_node_txn .if_then diff --git a/src/query/management/tests/it/warehouse.rs b/src/query/management/tests/it/warehouse.rs index 020ba9c954081..a7c897e2a0c9a 100644 --- a/src/query/management/tests/it/warehouse.rs +++ b/src/query/management/tests/it/warehouse.rs @@ -885,10 +885,10 @@ fn system_managed_node(id: &str) -> NodeInfo { discovery_address: "".to_string(), binary_version: "".to_string(), node_type: NodeType::SystemManaged, - resource_group: None, + node_group: None, cluster_id: "".to_string(), warehouse_id: "".to_string(), - runtime_resource_group: None, + runtime_node_group: None, } } @@ -903,10 +903,10 @@ fn self_managed_node(node_id: &str) -> NodeInfo { discovery_address: "ip2:port".to_string(), binary_version: "binary_version".to_string(), node_type: NodeType::SelfManaged, - resource_group: None, + node_group: None, cluster_id: "test-cluster-id".to_string(), warehouse_id: "test-cluster-id".to_string(), - runtime_resource_group: None, + runtime_node_group: None, } } diff --git a/src/query/service/src/interpreters/interpreter_add_warehouse_cluster.rs b/src/query/service/src/interpreters/interpreter_add_warehouse_cluster.rs index 160dc6ea1e023..47586f4b96d0d 100644 --- a/src/query/service/src/interpreters/interpreter_add_warehouse_cluster.rs +++ b/src/query/service/src/interpreters/interpreter_add_warehouse_cluster.rs @@ -15,18 +15,21 @@ use std::sync::Arc; use databend_common_base::base::GlobalInstance; +use databend_common_catalog::table_context::TableContext; use databend_common_exception::ErrorCode; use databend_common_exception::Result; +use databend_common_license::license::Feature; +use databend_common_license::license_manager::LicenseManagerSwitch; use databend_common_management::SelectedNode; use databend_common_sql::plans::AddWarehouseClusterPlan; use databend_enterprise_resources_management::ResourcesManagement; +use crate::interpreters::util::AuditElement; use crate::interpreters::Interpreter; use crate::pipelines::PipelineBuildResult; use crate::sessions::QueryContext; pub struct AddWarehouseClusterInterpreter { - #[allow(dead_code)] ctx: Arc, plan: AddWarehouseClusterPlan, } @@ -49,6 +52,9 @@ impl Interpreter for AddWarehouseClusterInterpreter { #[async_backtrace::framed] async fn execute2(&self) -> Result { + LicenseManagerSwitch::instance() + .check_enterprise_enabled(self.ctx.get_license_key(), Feature::SystemManagement)?; + if let Some(cluster_size) = self.plan.options.get("cluster_size") { if !self.plan.nodes.is_empty() { return Err(ErrorCode::InvalidArgument( @@ -92,6 +98,13 @@ impl Interpreter for AddWarehouseClusterInterpreter { ) .await?; + let user_info = self.ctx.get_current_user()?; + log::info!( + target: "databend::log::audit", + "{}", + serde_json::to_string(&AuditElement::create(&user_info, "alter_warehouse_add_cluster", &self.plan))? + ); + Ok(PipelineBuildResult::create()) } } diff --git a/src/query/service/src/interpreters/interpreter_assign_warehouse_nodes.rs b/src/query/service/src/interpreters/interpreter_assign_warehouse_nodes.rs index d23013aaa4179..5b3a8621989f2 100644 --- a/src/query/service/src/interpreters/interpreter_assign_warehouse_nodes.rs +++ b/src/query/service/src/interpreters/interpreter_assign_warehouse_nodes.rs @@ -16,21 +16,27 @@ use std::collections::HashMap; use std::sync::Arc; use databend_common_base::base::GlobalInstance; +use databend_common_catalog::table_context::TableContext; use databend_common_exception::Result; +use databend_common_license::license::Feature; +use databend_common_license::license_manager::LicenseManagerSwitch; use databend_common_management::SelectedNode; use databend_common_sql::plans::AssignWarehouseNodesPlan; use databend_enterprise_resources_management::ResourcesManagement; +use crate::interpreters::util::AuditElement; use crate::interpreters::Interpreter; use crate::pipelines::PipelineBuildResult; +use crate::sessions::QueryContext; pub struct AssignWarehouseNodesInterpreter { + ctx: Arc, plan: AssignWarehouseNodesPlan, } impl AssignWarehouseNodesInterpreter { - pub fn try_create(plan: AssignWarehouseNodesPlan) -> Result { - Ok(AssignWarehouseNodesInterpreter { plan }) + pub fn try_create(ctx: Arc, plan: AssignWarehouseNodesPlan) -> Result { + Ok(AssignWarehouseNodesInterpreter { ctx, plan }) } } @@ -46,6 +52,9 @@ impl Interpreter for AssignWarehouseNodesInterpreter { #[async_backtrace::framed] async fn execute2(&self) -> Result { + LicenseManagerSwitch::instance() + .check_enterprise_enabled(self.ctx.get_license_key(), Feature::SystemManagement)?; + let mut cluster_selected_nodes = HashMap::with_capacity(self.plan.assign_clusters.len()); for (cluster, nodes_map) in &self.plan.assign_clusters { let mut selected_nodes = Vec::with_capacity(nodes_map.len()); @@ -61,6 +70,14 @@ impl Interpreter for AssignWarehouseNodesInterpreter { GlobalInstance::get::>() .assign_warehouse_nodes(self.plan.warehouse.clone(), cluster_selected_nodes) .await?; + + let user_info = self.ctx.get_current_user()?; + log::info!( + target: "databend::log::audit", + "{}", + serde_json::to_string(&AuditElement::create(&user_info, "alter_warehouse_assign_nodes", &self.plan))? + ); + Ok(PipelineBuildResult::create()) } } diff --git a/src/query/service/src/interpreters/interpreter_create_warehouses.rs b/src/query/service/src/interpreters/interpreter_create_warehouses.rs index e58ad4ff86cd8..aafa1554c29bb 100644 --- a/src/query/service/src/interpreters/interpreter_create_warehouses.rs +++ b/src/query/service/src/interpreters/interpreter_create_warehouses.rs @@ -15,12 +15,16 @@ use std::sync::Arc; use databend_common_base::base::GlobalInstance; +use databend_common_catalog::table_context::TableContext; use databend_common_exception::ErrorCode; use databend_common_exception::Result; +use databend_common_license::license::Feature; +use databend_common_license::license_manager::LicenseManagerSwitch; use databend_common_management::SelectedNode; use databend_common_sql::plans::CreateWarehousePlan; use databend_enterprise_resources_management::ResourcesManagement; +use crate::interpreters::util::AuditElement; use crate::interpreters::Interpreter; use crate::pipelines::PipelineBuildResult; use crate::sessions::QueryContext; @@ -49,6 +53,9 @@ impl Interpreter for CreateWarehouseInterpreter { #[async_backtrace::framed] async fn execute2(&self) -> Result { + LicenseManagerSwitch::instance() + .check_enterprise_enabled(self.ctx.get_license_key(), Feature::SystemManagement)?; + if let Some(warehouse_size) = self.plan.options.get("warehouse_size") { if !self.plan.nodes.is_empty() { return Err(ErrorCode::InvalidArgument( @@ -87,6 +94,13 @@ impl Interpreter for CreateWarehouseInterpreter { .create_warehouse(self.plan.warehouse.clone(), selected_nodes) .await?; + let user_info = self.ctx.get_current_user()?; + log::info!( + target: "databend::log::audit", + "{}", + serde_json::to_string(&AuditElement::create(&user_info, "create_warehouse", &self.plan))? + ); + Ok(PipelineBuildResult::create()) } } diff --git a/src/query/service/src/interpreters/interpreter_drop_warehouse_cluster.rs b/src/query/service/src/interpreters/interpreter_drop_warehouse_cluster.rs index 107a544416661..1a5424adf515b 100644 --- a/src/query/service/src/interpreters/interpreter_drop_warehouse_cluster.rs +++ b/src/query/service/src/interpreters/interpreter_drop_warehouse_cluster.rs @@ -15,10 +15,14 @@ use std::sync::Arc; use databend_common_base::base::GlobalInstance; +use databend_common_catalog::table_context::TableContext; use databend_common_exception::Result; +use databend_common_license::license::Feature; +use databend_common_license::license_manager::LicenseManagerSwitch; use databend_common_sql::plans::DropWarehouseClusterPlan; use databend_enterprise_resources_management::ResourcesManagement; +use crate::interpreters::util::AuditElement; use crate::interpreters::Interpreter; use crate::pipelines::PipelineBuildResult; use crate::sessions::QueryContext; @@ -47,9 +51,20 @@ impl Interpreter for DropWarehouseClusterInterpreter { #[async_backtrace::framed] async fn execute2(&self) -> Result { + LicenseManagerSwitch::instance() + .check_enterprise_enabled(self.ctx.get_license_key(), Feature::SystemManagement)?; + GlobalInstance::get::>() .drop_warehouse_cluster(self.plan.warehouse.clone(), self.plan.cluster.clone()) .await?; + + let user_info = self.ctx.get_current_user()?; + log::info!( + target: "databend::log::audit", + "{}", + serde_json::to_string(&AuditElement::create(&user_info, "alter_warehouse_drop_cluster", &self.plan))? + ); + Ok(PipelineBuildResult::create()) } } diff --git a/src/query/service/src/interpreters/interpreter_drop_warehouse_cluster_node.rs b/src/query/service/src/interpreters/interpreter_drop_warehouse_cluster_node.rs deleted file mode 100644 index 35d20accdc4c7..0000000000000 --- a/src/query/service/src/interpreters/interpreter_drop_warehouse_cluster_node.rs +++ /dev/null @@ -1,49 +0,0 @@ -// Copyright 2021 Datafuse Labs -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -use databend_common_exception::Result; -use databend_common_sql::plans::UnassignWarehouseNodesPlan; - -use crate::interpreters::Interpreter; -use crate::pipelines::PipelineBuildResult; - -pub struct DropWarehouseClusterNodeInterpreter { - #[allow(dead_code)] - plan: UnassignWarehouseNodesPlan, -} - -impl DropWarehouseClusterNodeInterpreter { - pub fn try_create(plan: UnassignWarehouseNodesPlan) -> Result { - Ok(DropWarehouseClusterNodeInterpreter { plan }) - } -} - -#[async_trait::async_trait] -impl Interpreter for DropWarehouseClusterNodeInterpreter { - fn name(&self) -> &str { - "DropWarehouseClusterNodeInterpreter" - } - - fn is_ddl(&self) -> bool { - true - } - - #[async_backtrace::framed] - async fn execute2(&self) -> Result { - // GlobalInstance::get::>() - // .drop_warehouse_cluster(self.plan.warehouse.clone(), self.plan.cluster.clone()) - // .await?; - Ok(PipelineBuildResult::create()) - } -} diff --git a/src/query/service/src/interpreters/interpreter_drop_warehouses.rs b/src/query/service/src/interpreters/interpreter_drop_warehouses.rs index d077cb688429b..46276c39750ee 100644 --- a/src/query/service/src/interpreters/interpreter_drop_warehouses.rs +++ b/src/query/service/src/interpreters/interpreter_drop_warehouses.rs @@ -15,10 +15,14 @@ use std::sync::Arc; use databend_common_base::base::GlobalInstance; +use databend_common_catalog::table_context::TableContext; use databend_common_exception::Result; +use databend_common_license::license::Feature; +use databend_common_license::license_manager::LicenseManagerSwitch; use databend_common_sql::plans::DropWarehousePlan; use databend_enterprise_resources_management::ResourcesManagement; +use crate::interpreters::util::AuditElement; use crate::interpreters::Interpreter; use crate::pipelines::PipelineBuildResult; use crate::sessions::QueryContext; @@ -47,9 +51,20 @@ impl Interpreter for DropWarehouseInterpreter { #[async_backtrace::framed] async fn execute2(&self) -> Result { + LicenseManagerSwitch::instance() + .check_enterprise_enabled(self.ctx.get_license_key(), Feature::SystemManagement)?; + GlobalInstance::get::>() .drop_warehouse(self.plan.warehouse.clone()) .await?; + + let user_info = self.ctx.get_current_user()?; + log::info!( + target: "databend::log::audit", + "{}", + serde_json::to_string(&AuditElement::create(&user_info, "drop_warehouse", &self.plan))? + ); + Ok(PipelineBuildResult::create()) } } diff --git a/src/query/service/src/interpreters/interpreter_factory.rs b/src/query/service/src/interpreters/interpreter_factory.rs index 7610e4acb44da..d767a51771e0d 100644 --- a/src/query/service/src/interpreters/interpreter_factory.rs +++ b/src/query/service/src/interpreters/interpreter_factory.rs @@ -48,7 +48,6 @@ use crate::interpreters::interpreter_copy_into_location::CopyIntoLocationInterpr use crate::interpreters::interpreter_copy_into_table::CopyIntoTableInterpreter; use crate::interpreters::interpreter_create_warehouses::CreateWarehouseInterpreter; use crate::interpreters::interpreter_drop_warehouse_cluster::DropWarehouseClusterInterpreter; -use crate::interpreters::interpreter_drop_warehouse_cluster_node::DropWarehouseClusterNodeInterpreter; use crate::interpreters::interpreter_drop_warehouses::DropWarehouseInterpreter; use crate::interpreters::interpreter_file_format_create::CreateFileFormatInterpreter; use crate::interpreters::interpreter_file_format_drop::DropFileFormatInterpreter; @@ -83,6 +82,7 @@ use crate::interpreters::interpreter_tasks_show::ShowTasksInterpreter; use crate::interpreters::interpreter_txn_abort::AbortInterpreter; use crate::interpreters::interpreter_txn_begin::BeginInterpreter; use crate::interpreters::interpreter_txn_commit::CommitInterpreter; +use crate::interpreters::interpreter_unassign_warehouse_nodes::UnassignWarehouseNodesInterpreter; use crate::interpreters::interpreter_view_describe::DescribeViewInterpreter; use crate::interpreters::AlterUserInterpreter; use crate::interpreters::CreateStreamInterpreter; @@ -118,8 +118,8 @@ impl InterpreterFactory { other: impl FnOnce(Arc, &Plan) -> Result, ) -> Result { match plan { - Plan::ShowWarehouses => Ok(Arc::new(ShowWarehousesInterpreter::try_create()?)), - Plan::ShowOnlineNodes => Ok(Arc::new(ShowOnlineNodesInterpreter::try_create()?)), + Plan::ShowWarehouses => Ok(Arc::new(ShowWarehousesInterpreter::try_create(ctx.clone())?)), + Plan::ShowOnlineNodes => Ok(Arc::new(ShowOnlineNodesInterpreter::try_create(ctx.clone())?)), Plan::CreateWarehouse(v) => Ok(Arc::new(CreateWarehouseInterpreter::try_create( ctx.clone(), *v.clone(), @@ -129,9 +129,11 @@ impl InterpreterFactory { *v.clone(), )?)), Plan::ResumeWarehouse(v) => Ok(Arc::new(ResumeWarehouseInterpreter::try_create( + ctx.clone(), *v.clone(), )?)), Plan::SuspendWarehouse(v) => Ok(Arc::new(SuspendWarehouseInterpreter::try_create( + ctx.clone(), *v.clone(), )?)), Plan::RenameWarehouse(v) => Ok(Arc::new(RenameWarehouseInterpreter::try_create( @@ -139,6 +141,7 @@ impl InterpreterFactory { *v.clone(), )?)), Plan::InspectWarehouse(v) => Ok(Arc::new(InspectWarehouseInterpreter::try_create( + ctx.clone(), *v.clone(), )?)), Plan::AddWarehouseCluster(v) => Ok(Arc::new( @@ -148,13 +151,13 @@ impl InterpreterFactory { DropWarehouseClusterInterpreter::try_create(ctx.clone(), *v.clone())?, )), Plan::RenameWarehouseCluster(v) => Ok(Arc::new( - RenameWarehouseClusterInterpreter::try_create(*v.clone())?, + RenameWarehouseClusterInterpreter::try_create(ctx.clone(), *v.clone())?, )), Plan::AssignWarehouseNodes(v) => Ok(Arc::new( - AssignWarehouseNodesInterpreter::try_create(*v.clone())?, + AssignWarehouseNodesInterpreter::try_create(ctx.clone(), *v.clone())?, )), Plan::UnassignWarehouseNodes(v) => Ok(Arc::new( - DropWarehouseClusterNodeInterpreter::try_create(*v.clone())?, + UnassignWarehouseNodesInterpreter::try_create(ctx.clone(), *v.clone())?, )), Plan::Query { metadata, .. } => { let read_guard = metadata.read(); diff --git a/src/query/service/src/interpreters/interpreter_inspect_warehouse.rs b/src/query/service/src/interpreters/interpreter_inspect_warehouse.rs index 15649f5950918..0ed2b5243d421 100644 --- a/src/query/service/src/interpreters/interpreter_inspect_warehouse.rs +++ b/src/query/service/src/interpreters/interpreter_inspect_warehouse.rs @@ -15,25 +15,30 @@ use std::sync::Arc; use databend_common_base::base::GlobalInstance; +use databend_common_catalog::table_context::TableContext; use databend_common_exception::Result; use databend_common_expression::types::DataType; use databend_common_expression::ColumnBuilder; use databend_common_expression::DataBlock; use databend_common_expression::Scalar; +use databend_common_license::license::Feature; +use databend_common_license::license_manager::LicenseManagerSwitch; use databend_common_meta_types::NodeType; use databend_common_sql::plans::InspectWarehousePlan; use databend_enterprise_resources_management::ResourcesManagement; use crate::interpreters::Interpreter; use crate::pipelines::PipelineBuildResult; +use crate::sessions::QueryContext; pub struct InspectWarehouseInterpreter { + ctx: Arc, plan: InspectWarehousePlan, } impl InspectWarehouseInterpreter { - pub fn try_create(plan: InspectWarehousePlan) -> Result { - Ok(InspectWarehouseInterpreter { plan }) + pub fn try_create(ctx: Arc, plan: InspectWarehousePlan) -> Result { + Ok(InspectWarehouseInterpreter { ctx, plan }) } } @@ -49,6 +54,9 @@ impl Interpreter for InspectWarehouseInterpreter { #[async_backtrace::framed] async fn execute2(&self) -> Result { + LicenseManagerSwitch::instance() + .check_enterprise_enabled(self.ctx.get_license_key(), Feature::SystemManagement)?; + let mut warehouse_nodes = GlobalInstance::get::>() .inspect_warehouse(self.plan.warehouse.clone()) .await?; diff --git a/src/query/service/src/interpreters/interpreter_rename_warehouse.rs b/src/query/service/src/interpreters/interpreter_rename_warehouse.rs index 1cf3536d16a57..b044c82bb624f 100644 --- a/src/query/service/src/interpreters/interpreter_rename_warehouse.rs +++ b/src/query/service/src/interpreters/interpreter_rename_warehouse.rs @@ -15,10 +15,14 @@ use std::sync::Arc; use databend_common_base::base::GlobalInstance; +use databend_common_catalog::table_context::TableContext; use databend_common_exception::Result; +use databend_common_license::license::Feature; +use databend_common_license::license_manager::LicenseManagerSwitch; use databend_common_sql::plans::RenameWarehousePlan; use databend_enterprise_resources_management::ResourcesManagement; +use crate::interpreters::util::AuditElement; use crate::interpreters::Interpreter; use crate::pipelines::PipelineBuildResult; use crate::sessions::QueryContext; @@ -47,9 +51,20 @@ impl Interpreter for RenameWarehouseInterpreter { #[async_backtrace::framed] async fn execute2(&self) -> Result { + LicenseManagerSwitch::instance() + .check_enterprise_enabled(self.ctx.get_license_key(), Feature::SystemManagement)?; + GlobalInstance::get::>() .rename_warehouse(self.plan.warehouse.clone(), self.plan.new_warehouse.clone()) .await?; + + let user_info = self.ctx.get_current_user()?; + log::info!( + target: "databend::log::audit", + "{}", + serde_json::to_string(&AuditElement::create(&user_info, "rename_warehouse", &self.plan))? + ); + Ok(PipelineBuildResult::create()) } } diff --git a/src/query/service/src/interpreters/interpreter_rename_warehouse_cluster.rs b/src/query/service/src/interpreters/interpreter_rename_warehouse_cluster.rs index 9c6a91b840b3c..5724e2826857c 100644 --- a/src/query/service/src/interpreters/interpreter_rename_warehouse_cluster.rs +++ b/src/query/service/src/interpreters/interpreter_rename_warehouse_cluster.rs @@ -15,20 +15,26 @@ use std::sync::Arc; use databend_common_base::base::GlobalInstance; +use databend_common_catalog::table_context::TableContext; use databend_common_exception::Result; +use databend_common_license::license::Feature; +use databend_common_license::license_manager::LicenseManagerSwitch; use databend_common_sql::plans::RenameWarehouseClusterPlan; use databend_enterprise_resources_management::ResourcesManagement; +use crate::interpreters::util::AuditElement; use crate::interpreters::Interpreter; use crate::pipelines::PipelineBuildResult; +use crate::sessions::QueryContext; pub struct RenameWarehouseClusterInterpreter { + ctx: Arc, plan: RenameWarehouseClusterPlan, } impl RenameWarehouseClusterInterpreter { - pub fn try_create(plan: RenameWarehouseClusterPlan) -> Result { - Ok(RenameWarehouseClusterInterpreter { plan }) + pub fn try_create(ctx: Arc, plan: RenameWarehouseClusterPlan) -> Result { + Ok(RenameWarehouseClusterInterpreter { ctx, plan }) } } @@ -44,6 +50,9 @@ impl Interpreter for RenameWarehouseClusterInterpreter { #[async_backtrace::framed] async fn execute2(&self) -> Result { + LicenseManagerSwitch::instance() + .check_enterprise_enabled(self.ctx.get_license_key(), Feature::SystemManagement)?; + GlobalInstance::get::>() .rename_warehouse_cluster( self.plan.warehouse.clone(), @@ -51,6 +60,14 @@ impl Interpreter for RenameWarehouseClusterInterpreter { self.plan.new_cluster.clone(), ) .await?; + + let user_info = self.ctx.get_current_user()?; + log::info!( + target: "databend::log::audit", + "{}", + serde_json::to_string(&AuditElement::create(&user_info, "alter_warehouse_rename_cluster", &self.plan))? + ); + Ok(PipelineBuildResult::create()) } } diff --git a/src/query/service/src/interpreters/interpreter_resume_warehouse.rs b/src/query/service/src/interpreters/interpreter_resume_warehouse.rs index 60419d6224a51..f763908ba3f87 100644 --- a/src/query/service/src/interpreters/interpreter_resume_warehouse.rs +++ b/src/query/service/src/interpreters/interpreter_resume_warehouse.rs @@ -15,20 +15,26 @@ use std::sync::Arc; use databend_common_base::base::GlobalInstance; +use databend_common_catalog::table_context::TableContext; use databend_common_exception::Result; +use databend_common_license::license::Feature; +use databend_common_license::license_manager::LicenseManagerSwitch; use databend_common_sql::plans::ResumeWarehousePlan; use databend_enterprise_resources_management::ResourcesManagement; +use crate::interpreters::util::AuditElement; use crate::interpreters::Interpreter; use crate::pipelines::PipelineBuildResult; +use crate::sessions::QueryContext; pub struct ResumeWarehouseInterpreter { + ctx: Arc, plan: ResumeWarehousePlan, } impl ResumeWarehouseInterpreter { - pub fn try_create(plan: ResumeWarehousePlan) -> Result { - Ok(ResumeWarehouseInterpreter { plan }) + pub fn try_create(ctx: Arc, plan: ResumeWarehousePlan) -> Result { + Ok(ResumeWarehouseInterpreter { ctx, plan }) } } @@ -44,9 +50,20 @@ impl Interpreter for ResumeWarehouseInterpreter { #[async_backtrace::framed] async fn execute2(&self) -> Result { + LicenseManagerSwitch::instance() + .check_enterprise_enabled(self.ctx.get_license_key(), Feature::SystemManagement)?; + GlobalInstance::get::>() .resume_warehouse(self.plan.warehouse.clone()) .await?; + + let user_info = self.ctx.get_current_user()?; + log::info!( + target: "databend::log::audit", + "{}", + serde_json::to_string(&AuditElement::create(&user_info, "resume_warehouse", &self.plan))? + ); + Ok(PipelineBuildResult::create()) } } diff --git a/src/query/service/src/interpreters/interpreter_show_online_nodes.rs b/src/query/service/src/interpreters/interpreter_show_online_nodes.rs index 674b2cbcff4e4..21f1c6ca8d950 100644 --- a/src/query/service/src/interpreters/interpreter_show_online_nodes.rs +++ b/src/query/service/src/interpreters/interpreter_show_online_nodes.rs @@ -15,22 +15,28 @@ use std::sync::Arc; use databend_common_base::base::GlobalInstance; +use databend_common_catalog::table_context::TableContext; use databend_common_exception::Result; use databend_common_expression::types::DataType; use databend_common_expression::ColumnBuilder; use databend_common_expression::DataBlock; use databend_common_expression::Scalar; +use databend_common_license::license::Feature; +use databend_common_license::license_manager::LicenseManagerSwitch; use databend_common_meta_types::NodeType; use databend_enterprise_resources_management::ResourcesManagement; use crate::interpreters::Interpreter; use crate::pipelines::PipelineBuildResult; +use crate::sessions::QueryContext; -pub struct ShowOnlineNodesInterpreter {} +pub struct ShowOnlineNodesInterpreter { + ctx: Arc, +} impl ShowOnlineNodesInterpreter { - pub fn try_create() -> Result { - Ok(ShowOnlineNodesInterpreter {}) + pub fn try_create(ctx: Arc) -> Result { + Ok(ShowOnlineNodesInterpreter { ctx }) } } @@ -46,14 +52,16 @@ impl Interpreter for ShowOnlineNodesInterpreter { #[async_backtrace::framed] async fn execute2(&self) -> Result { + LicenseManagerSwitch::instance() + .check_enterprise_enabled(self.ctx.get_license_key(), Feature::SystemManagement)?; + let online_nodes = GlobalInstance::get::>() .list_online_nodes() .await?; let mut nodes_id = ColumnBuilder::with_capacity(&DataType::String, online_nodes.len()); let mut nodes_type = ColumnBuilder::with_capacity(&DataType::String, online_nodes.len()); - let mut nodes_resource_group = - ColumnBuilder::with_capacity(&DataType::String, online_nodes.len()); + let mut nodes_group = ColumnBuilder::with_capacity(&DataType::String, online_nodes.len()); let mut nodes_warehouse = ColumnBuilder::with_capacity(&DataType::String, online_nodes.len()); let mut nodes_cluster = ColumnBuilder::with_capacity(&DataType::String, online_nodes.len()); @@ -72,8 +80,7 @@ impl Interpreter for ShowOnlineNodesInterpreter { nodes_id.push(Scalar::String(node.id).as_ref()); nodes_type.push(Scalar::String(node_type).as_ref()); - nodes_resource_group - .push(Scalar::String(node.resource_group.clone().unwrap_or_default()).as_ref()); + nodes_group.push(Scalar::String(node.node_group.clone().unwrap_or_default()).as_ref()); nodes_warehouse.push(Scalar::String(node.warehouse_id).as_ref()); nodes_cluster.push(Scalar::String(node.cluster_id).as_ref()); nodes_version.push(Scalar::String(binary_version).as_ref()); @@ -82,7 +89,7 @@ impl Interpreter for ShowOnlineNodesInterpreter { PipelineBuildResult::from_blocks(vec![DataBlock::new_from_columns(vec![ nodes_id.build(), nodes_type.build(), - nodes_resource_group.build(), + nodes_group.build(), nodes_warehouse.build(), nodes_cluster.build(), nodes_version.build(), diff --git a/src/query/service/src/interpreters/interpreter_show_warehouses.rs b/src/query/service/src/interpreters/interpreter_show_warehouses.rs index 4dbcd23e65c7d..883318df7f930 100644 --- a/src/query/service/src/interpreters/interpreter_show_warehouses.rs +++ b/src/query/service/src/interpreters/interpreter_show_warehouses.rs @@ -15,22 +15,28 @@ use std::sync::Arc; use databend_common_base::base::GlobalInstance; +use databend_common_catalog::table_context::TableContext; use databend_common_exception::Result; use databend_common_expression::types::DataType; use databend_common_expression::ColumnBuilder; use databend_common_expression::DataBlock; use databend_common_expression::Scalar; +use databend_common_license::license::Feature; +use databend_common_license::license_manager::LicenseManagerSwitch; use databend_common_management::WarehouseInfo; use databend_enterprise_resources_management::ResourcesManagement; use crate::interpreters::Interpreter; use crate::pipelines::PipelineBuildResult; +use crate::sessions::QueryContext; -pub struct ShowWarehousesInterpreter {} +pub struct ShowWarehousesInterpreter { + ctx: Arc, +} impl ShowWarehousesInterpreter { - pub fn try_create() -> Result { - Ok(ShowWarehousesInterpreter {}) + pub fn try_create(ctx: Arc) -> Result { + Ok(ShowWarehousesInterpreter { ctx }) } } @@ -46,6 +52,9 @@ impl Interpreter for ShowWarehousesInterpreter { #[async_backtrace::framed] async fn execute2(&self) -> Result { + LicenseManagerSwitch::instance() + .check_enterprise_enabled(self.ctx.get_license_key(), Feature::SystemManagement)?; + let warehouses = GlobalInstance::get::>() .list_warehouses() .await?; diff --git a/src/query/service/src/interpreters/interpreter_suspend_warehouse.rs b/src/query/service/src/interpreters/interpreter_suspend_warehouse.rs index eee6e2657426d..31d2691cbc877 100644 --- a/src/query/service/src/interpreters/interpreter_suspend_warehouse.rs +++ b/src/query/service/src/interpreters/interpreter_suspend_warehouse.rs @@ -15,20 +15,26 @@ use std::sync::Arc; use databend_common_base::base::GlobalInstance; +use databend_common_catalog::table_context::TableContext; use databend_common_exception::Result; +use databend_common_license::license::Feature; +use databend_common_license::license_manager::LicenseManagerSwitch; use databend_common_sql::plans::SuspendWarehousePlan; use databend_enterprise_resources_management::ResourcesManagement; +use crate::interpreters::util::AuditElement; use crate::interpreters::Interpreter; use crate::pipelines::PipelineBuildResult; +use crate::sessions::QueryContext; pub struct SuspendWarehouseInterpreter { + ctx: Arc, plan: SuspendWarehousePlan, } impl SuspendWarehouseInterpreter { - pub fn try_create(plan: SuspendWarehousePlan) -> Result { - Ok(SuspendWarehouseInterpreter { plan }) + pub fn try_create(ctx: Arc, plan: SuspendWarehousePlan) -> Result { + Ok(SuspendWarehouseInterpreter { ctx, plan }) } } @@ -44,9 +50,20 @@ impl Interpreter for SuspendWarehouseInterpreter { #[async_backtrace::framed] async fn execute2(&self) -> Result { + LicenseManagerSwitch::instance() + .check_enterprise_enabled(self.ctx.get_license_key(), Feature::SystemManagement)?; + GlobalInstance::get::>() .suspend_warehouse(self.plan.warehouse.clone()) .await?; + + let user_info = self.ctx.get_current_user()?; + log::info!( + target: "databend::log::audit", + "{}", + serde_json::to_string(&AuditElement::create(&user_info, "suspend_warehouse", &self.plan))? + ); + Ok(PipelineBuildResult::create()) } } diff --git a/src/query/service/src/interpreters/interpreter_unassign_warehouse_nodes.rs b/src/query/service/src/interpreters/interpreter_unassign_warehouse_nodes.rs new file mode 100644 index 0000000000000..4f5b5d6cee0d3 --- /dev/null +++ b/src/query/service/src/interpreters/interpreter_unassign_warehouse_nodes.rs @@ -0,0 +1,83 @@ +// Copyright 2021 Datafuse Labs +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::collections::HashMap; +use std::sync::Arc; + +use databend_common_base::base::GlobalInstance; +use databend_common_catalog::table_context::TableContext; +use databend_common_exception::Result; +use databend_common_license::license::Feature; +use databend_common_license::license_manager::LicenseManagerSwitch; +use databend_common_management::SelectedNode; +use databend_common_sql::plans::UnassignWarehouseNodesPlan; +use databend_enterprise_resources_management::ResourcesManagement; + +use crate::interpreters::util::AuditElement; +use crate::interpreters::Interpreter; +use crate::pipelines::PipelineBuildResult; +use crate::sessions::QueryContext; + +pub struct UnassignWarehouseNodesInterpreter { + ctx: Arc, + plan: UnassignWarehouseNodesPlan, +} + +impl UnassignWarehouseNodesInterpreter { + pub fn try_create(ctx: Arc, plan: UnassignWarehouseNodesPlan) -> Result { + Ok(UnassignWarehouseNodesInterpreter { ctx, plan }) + } +} + +#[async_trait::async_trait] +impl Interpreter for UnassignWarehouseNodesInterpreter { + fn name(&self) -> &str { + "DropWarehouseClusterNodeInterpreter" + } + + fn is_ddl(&self) -> bool { + true + } + + #[async_backtrace::framed] + async fn execute2(&self) -> Result { + LicenseManagerSwitch::instance() + .check_enterprise_enabled(self.ctx.get_license_key(), Feature::SystemManagement)?; + + let mut cluster_selected_nodes = HashMap::with_capacity(self.plan.unassign_clusters.len()); + for (cluster, nodes_map) in &self.plan.unassign_clusters { + let mut selected_nodes = Vec::with_capacity(nodes_map.len()); + for (group, nodes) in nodes_map { + for _ in 0..*nodes { + selected_nodes.push(SelectedNode::Random(group.clone())); + } + } + + cluster_selected_nodes.insert(cluster.clone(), selected_nodes); + } + + GlobalInstance::get::>() + .unassign_warehouse_nodes(self.plan.warehouse.clone(), cluster_selected_nodes) + .await?; + + let user_info = self.ctx.get_current_user()?; + log::info!( + target: "databend::log::audit", + "{}", + serde_json::to_string(&AuditElement::create(&user_info, "alter_warehouse_unassign_nodes", &self.plan))? + ); + + Ok(PipelineBuildResult::create()) + } +} diff --git a/src/query/service/src/interpreters/mod.rs b/src/query/service/src/interpreters/mod.rs index 9cb860f828b02..d978c2cee591f 100644 --- a/src/query/service/src/interpreters/mod.rs +++ b/src/query/service/src/interpreters/mod.rs @@ -45,7 +45,6 @@ mod interpreter_dictionary_drop; mod interpreter_dictionary_rename; mod interpreter_dictionary_show_create; mod interpreter_drop_warehouse_cluster; -mod interpreter_drop_warehouse_cluster_node; mod interpreter_drop_warehouses; mod interpreter_execute_immediate; mod interpreter_explain; @@ -137,6 +136,7 @@ mod interpreter_tasks_show; mod interpreter_txn_abort; mod interpreter_txn_begin; mod interpreter_txn_commit; +mod interpreter_unassign_warehouse_nodes; mod interpreter_unset; mod interpreter_use_database; mod interpreter_user_alter; diff --git a/src/query/service/src/interpreters/util.rs b/src/query/service/src/interpreters/util.rs index ba71b1fc0c401..21ecc616b7109 100644 --- a/src/query/service/src/interpreters/util.rs +++ b/src/query/service/src/interpreters/util.rs @@ -24,6 +24,7 @@ use databend_common_expression::DataBlock; use databend_common_expression::DataSchemaRef; use databend_common_expression::Scalar; use databend_common_expression::TableSchemaRef; +use databend_common_meta_app::principal::UserInfo; use databend_common_script::ir::ColumnAccess; use databend_common_script::Client; use databend_common_sql::Planner; @@ -185,3 +186,22 @@ impl Client for ScriptClient { } } } + +#[derive(serde::Serialize)] +pub struct AuditElement<'a, T: serde::Serialize> { + user: String, + hostname: String, + action_type: &'static str, + action: &'a T, +} + +impl<'a, T: serde::Serialize> AuditElement<'a, T> { + pub fn create(user: &UserInfo, action_type: &'static str, action: &'a T) -> Self { + AuditElement { + action, + action_type, + user: user.name.clone(), + hostname: user.hostname.clone(), + } + } +} diff --git a/src/query/service/tests/it/clusters.rs b/src/query/service/tests/it/clusters.rs index 8bb4ca7e14257..96b31b3e5899c 100644 --- a/src/query/service/tests/it/clusters.rs +++ b/src/query/service/tests/it/clusters.rs @@ -36,6 +36,8 @@ async fn test_single_cluster_discovery() -> Result<()> { #[tokio::test(flavor = "current_thread")] async fn test_remove_invalid_nodes() -> Result<()> { + let _guard = TestFixture::setup().await?; + let config_1 = ConfigBuilder::create() .query_flight_address("invalid_address_1") .build(); diff --git a/src/query/service/tests/it/storages/testdata/configs_table_basic.txt b/src/query/service/tests/it/storages/testdata/configs_table_basic.txt index 20762066b467a..d5a4dd3c75516 100644 --- a/src/query/service/tests/it/storages/testdata/configs_table_basic.txt +++ b/src/query/service/tests/it/storages/testdata/configs_table_basic.txt @@ -128,6 +128,7 @@ DB.Table: 'system'.'configs', Table: configs-table_id:1, ver:0, Engine: SystemCo | 'query' | 'openai_api_version' | '' | '' | | 'query' | 'parquet_fast_read_bytes' | 'null' | '' | | 'query' | 'quota' | 'null' | '' | +| 'query' | 'resources_management' | 'null' | '' | | 'query' | 'rpc_client_timeout_secs' | '0' | '' | | 'query' | 'rpc_tls_query_server_root_ca_cert' | '' | '' | | 'query' | 'rpc_tls_query_service_domain_name' | 'localhost' | '' | diff --git a/src/query/sql/src/planner/plans/ddl/warehouse.rs b/src/query/sql/src/planner/plans/ddl/warehouse.rs index 5e028643d8acd..c66fe220c652d 100644 --- a/src/query/sql/src/planner/plans/ddl/warehouse.rs +++ b/src/query/sql/src/planner/plans/ddl/warehouse.rs @@ -20,29 +20,29 @@ use databend_common_expression::DataField; use databend_common_expression::DataSchemaRef; use databend_common_expression::DataSchemaRefExt; -#[derive(Clone, Debug, PartialEq, Eq)] +#[derive(Clone, Debug, PartialEq, Eq, serde::Serialize)] pub struct CreateWarehousePlan { pub warehouse: String, pub nodes: HashMap, u64>, pub options: BTreeMap, } -#[derive(Clone, Debug, PartialEq, Eq)] +#[derive(Clone, Debug, PartialEq, Eq, serde::Serialize)] pub struct DropWarehousePlan { pub warehouse: String, } -#[derive(Clone, Debug, PartialEq, Eq)] +#[derive(Clone, Debug, PartialEq, Eq, serde::Serialize)] pub struct ResumeWarehousePlan { pub warehouse: String, } -#[derive(Clone, Debug, PartialEq, Eq)] +#[derive(Clone, Debug, PartialEq, Eq, serde::Serialize)] pub struct SuspendWarehousePlan { pub warehouse: String, } -#[derive(Clone, Debug, PartialEq, Eq)] +#[derive(Clone, Debug, PartialEq, Eq, serde::Serialize)] pub struct RenameWarehousePlan { pub warehouse: String, pub new_warehouse: String, @@ -63,7 +63,7 @@ impl InspectWarehousePlan { } } -#[derive(Clone, Debug, PartialEq, Eq)] +#[derive(Clone, Debug, PartialEq, Eq, serde::Serialize)] pub struct AddWarehouseClusterPlan { pub warehouse: String, pub cluster: String, @@ -71,26 +71,26 @@ pub struct AddWarehouseClusterPlan { pub options: BTreeMap, } -#[derive(Clone, Debug, PartialEq, Eq)] +#[derive(Clone, Debug, PartialEq, Eq, serde::Serialize)] pub struct DropWarehouseClusterPlan { pub warehouse: String, pub cluster: String, } -#[derive(Clone, Debug, PartialEq, Eq)] +#[derive(Clone, Debug, PartialEq, Eq, serde::Serialize)] pub struct RenameWarehouseClusterPlan { pub warehouse: String, pub cluster: String, pub new_cluster: String, } -#[derive(Clone, Debug, PartialEq, Eq)] +#[derive(Clone, Debug, PartialEq, Eq, serde::Serialize)] pub struct AssignWarehouseNodesPlan { pub warehouse: String, pub assign_clusters: HashMap, usize>>, } -#[derive(Clone, Debug, PartialEq, Eq)] +#[derive(Clone, Debug, PartialEq, Eq, serde::Serialize)] pub struct UnassignWarehouseNodesPlan { pub warehouse: String, pub unassign_clusters: HashMap, usize>>, diff --git a/src/query/sql/src/planner/plans/plan.rs b/src/query/sql/src/planner/plans/plan.rs index 1523b7b281251..a189fe78224a0 100644 --- a/src/query/sql/src/planner/plans/plan.rs +++ b/src/query/sql/src/planner/plans/plan.rs @@ -530,7 +530,7 @@ impl Plan { Plan::ShowOnlineNodes => DataSchemaRefExt::create(vec![ DataField::new("id", DataType::String), DataField::new("type", DataType::String), - DataField::new("resource_group", DataType::String), + DataField::new("node_group", DataType::String), DataField::new("warehouse", DataType::String), DataField::new("cluster", DataType::String), DataField::new("version", DataType::String), From 9c97163748b5b68b6aca4c0d676b853d2797757b Mon Sep 17 00:00:00 2001 From: zhang2014 Date: Sun, 29 Dec 2024 19:12:52 +0800 Subject: [PATCH 28/53] feat(cluster): support custom management cluster --- .../management/src/warehouse/warehouse_mgr.rs | 131 +----------------- 1 file changed, 5 insertions(+), 126 deletions(-) diff --git a/src/query/management/src/warehouse/warehouse_mgr.rs b/src/query/management/src/warehouse/warehouse_mgr.rs index 5d7825985a876..96ded24cc2087 100644 --- a/src/query/management/src/warehouse/warehouse_mgr.rs +++ b/src/query/management/src/warehouse/warehouse_mgr.rs @@ -135,7 +135,7 @@ impl WarehouseMgr { } let mut txn = TxnRequest::default(); - let node_key = format!("{}/{}", self.node_key_prefix, escape_for_key(&node.id)?); + let node_key = self.node_key(&node)?; txn.condition.push(map_condition(&node_key, seq)); @@ -230,7 +230,7 @@ impl WarehouseMgr { async fn upsert_system_managed(&self, mut node: NodeInfo, seq: MatchSeq) -> Result { let mut txn = TxnRequest::default(); - let node_key = format!("{}/{}", self.node_key_prefix, escape_for_key(&node.id)?); + let node_key = self.node_key(&node)?; txn.condition.push(map_condition(&node_key, seq)); @@ -242,13 +242,7 @@ impl WarehouseMgr { // If the warehouse has already been assigned. if !node.cluster_id.is_empty() && !node.warehouse_id.is_empty() { - let cluster_key = format!( - "{}/{}/{}/{}", - self.meta_key_prefix, - escape_for_key(&node.warehouse_id)?, - escape_for_key(&node.cluster_id)?, - escape_for_key(&node.id)? - ); + let cluster_key = self.cluster_key(&node)?; node.cluster_id = String::new(); node.warehouse_id = String::new(); @@ -811,7 +805,7 @@ impl WarehouseApi for WarehouseMgr { let mut txn = TxnRequest::default(); for (seq, mut node) in selected_nodes { - let node_key = format!("{}/{}", self.node_key_prefix, escape_for_key(&node.id)?); + let node_key = self.node_key(&node)?; txn.condition .push(map_condition(&node_key, MatchSeq::Exact(seq))); @@ -821,13 +815,7 @@ impl WarehouseApi for WarehouseMgr { Some(self.lift_time * 4), )); - let cluster_key = format!( - "{}/{}/{}/{}", - self.meta_key_prefix, - escape_for_key(&node.warehouse_id)?, - escape_for_key(&node.cluster_id)?, - escape_for_key(&node.id)? - ); + let cluster_key = self.cluster_key(&node)?; node.cluster_id = String::new(); node.warehouse_id = String::new(); @@ -1789,115 +1777,6 @@ impl WarehouseApi for WarehouseMgr { } } } - // // if cluster.is_empty() { - // // return Err(ErrorCode::InvalidWarehouse( - // // "Warehouse cluster name is empty.", - // // )); - // // } - // - // for _idx in 0..10 { - // let mut drop_cluster_node_txn = TxnRequest::default(); - // - // let mut consistent_info = self.consistent_warehouse_info(warehouse).await?; - // - // consistent_info.warehouse_info = match consistent_info.warehouse_info { - // WarehouseInfo::SelfManaged(_) => Err(ErrorCode::InvalidWarehouse(format!( - // "Cannot add cluster for warehouse {:?}, because it's self-managed warehouse.", - // warehouse - // ))), - // WarehouseInfo::SystemManaged(mut info) => match info.clusters.get_mut(cluster) { - // None => Err(ErrorCode::WarehouseClusterNotExists(format!( - // "Warehouse cluster {:?}.{:?} not exists", - // warehouse, cluster - // ))), - // Some(cluster) => match nodes.len() == cluster.nodes.len() { - // true => Err(ErrorCode::EmptyNodesForWarehouse(format!( - // "Warehouse cluster {:?}.{:?} only has {} nodes, cannot drop all.", - // warehouse, - // cluster, - // nodes.len() - // ))), - // false => { - // for remove_node in &nodes { - // if consistent_info - // .consistent_nodes - // .iter() - // .any(|x| &x.node_info.id == remove_node) - // { - // cluster.nodes.pop(); - // continue; - // } - // - // return Err(ErrorCode::ClusterUnknownNode(format!( - // "Warehouse cluster {:?}.{:?} unknown node {:?}", - // warehouse, cluster, remove_node - // ))); - // } - // - // Ok(WarehouseInfo::SystemManaged(SystemManagedWarehouse { - // id: info.id, - // status: info.status, - // display_name: info.display_name, - // clusters: info.clusters, - // })) - // } - // }, - // }, - // }?; - // - // let warehouse_key = format!( - // "{}/{}", - // self.warehouse_key_prefix, - // escape_for_key(warehouse)? - // ); - // - // drop_cluster_node_txn.condition.push(map_condition( - // &warehouse_key, - // MatchSeq::Exact(consistent_info.info_seq), - // )); - // - // drop_cluster_node_txn.if_then.push(TxnOp::put( - // warehouse_key.clone(), - // serde_json::to_vec(&consistent_info.warehouse_info)?, - // )); - // - // // lock all cluster state - // for mut consistent_node in consistent_info.consistent_nodes { - // let node_key = self.node_key(&consistent_node.node_info)?; - // let cluster_key = self.cluster_key(&consistent_node.node_info)?; - // - // drop_cluster_node_txn.condition.push(map_condition( - // &node_key, - // MatchSeq::Exact(consistent_node.node_seq), - // )); - // drop_cluster_node_txn.condition.push(map_condition( - // &cluster_key, - // MatchSeq::Exact(consistent_node.cluster_seq), - // )); - // - // if nodes.contains(&consistent_node.node_info.id) { - // // Remove node - // consistent_node.node_info.cluster_id = String::new(); - // consistent_node.node_info.warehouse_id = String::new(); - // - // drop_cluster_node_txn - // .if_then - // .push(TxnOp::delete(cluster_key)); - // drop_cluster_node_txn.if_then.push(TxnOp::put_with_ttl( - // node_key, - // serde_json::to_vec(&consistent_node.node_info)?, - // Some(self.lift_time * 4), - // )) - // } - // } - // - // return match self.metastore.transaction(drop_cluster_node_txn).await? { - // res if res.success => Ok(()), - // _ => { - // continue; - // } - // }; - // } Err(ErrorCode::WarehouseOperateConflict( "Warehouse operate conflict(tried 10 times while in unassign warehouse nodes).", From d5f6795a4d8030b449eebc8003d36c187e029774 Mon Sep 17 00:00:00 2001 From: zhang2014 Date: Sun, 29 Dec 2024 19:49:36 +0800 Subject: [PATCH 29/53] feat(cluster): support custom management cluster --- src/meta/types/src/cluster.rs | 1 + .../management/src/warehouse/warehouse_api.rs | 17 +- .../management/src/warehouse/warehouse_mgr.rs | 300 ++++++++---------- 3 files changed, 139 insertions(+), 179 deletions(-) diff --git a/src/meta/types/src/cluster.rs b/src/meta/types/src/cluster.rs index 1162c648e0525..78b0a9fdf6f1b 100644 --- a/src/meta/types/src/cluster.rs +++ b/src/meta/types/src/cluster.rs @@ -97,6 +97,7 @@ pub struct NodeInfo { #[serde(skip_serializing_if = "String::is_empty")] pub warehouse_id: String, + #[serde(skip_serializing_if = "Option::is_none")] pub runtime_node_group: Option, } diff --git a/src/query/management/src/warehouse/warehouse_api.rs b/src/query/management/src/warehouse/warehouse_api.rs index 69a9e2aefd2e0..c4ec4df8434e6 100644 --- a/src/query/management/src/warehouse/warehouse_api.rs +++ b/src/query/management/src/warehouse/warehouse_api.rs @@ -17,16 +17,17 @@ use std::collections::HashMap; use databend_common_exception::Result; use databend_common_meta_types::NodeInfo; +// Used to describe the nodes in warehouse, excluding the details of the nodes. #[derive(serde::Serialize, serde::Deserialize, Clone, Eq, PartialEq, Debug)] pub enum SelectedNode { - Random(Option), + // Randomly select a node from the tenant's list of online nodes. The selected node's group must match it if node group is specified. + Random(Option /* node group */), } -pub type SelectedNodes = Vec; - +// Used to describe information about the warehouse, including both self-managed and system-managed warehouses to ensure compatibility. #[derive(serde::Serialize, serde::Deserialize, Eq, PartialEq, Debug)] pub enum WarehouseInfo { - SelfManaged(String), + SelfManaged(String /* cluster_id of self-managed cluster */), SystemManaged(SystemManagedWarehouse), } @@ -57,7 +58,7 @@ pub trait WarehouseApi: Sync + Send { async fn drop_warehouse(&self, warehouse: String) -> Result<()>; - async fn create_warehouse(&self, warehouse: String, nodes: SelectedNodes) -> Result<()>; + async fn create_warehouse(&self, warehouse: String, nodes: Vec) -> Result<()>; async fn resume_warehouse(&self, warehouse: String) -> Result<()>; @@ -73,7 +74,7 @@ pub trait WarehouseApi: Sync + Send { &self, warehouse: String, cluster: String, - nodes: SelectedNodes, + nodes: Vec, ) -> Result<()>; async fn drop_warehouse_cluster(&self, warehouse: String, cluster: String) -> Result<()>; @@ -88,13 +89,13 @@ pub trait WarehouseApi: Sync + Send { async fn assign_warehouse_nodes( &self, name: String, - nodes: HashMap, + nodes: HashMap>, ) -> Result<()>; async fn unassign_warehouse_nodes( &self, warehouse: &str, - nodes: HashMap, + nodes: HashMap>, ) -> Result<()>; /// Get the tenant's cluster all nodes. diff --git a/src/query/management/src/warehouse/warehouse_mgr.rs b/src/query/management/src/warehouse/warehouse_mgr.rs index 96ded24cc2087..504a34a74b4ce 100644 --- a/src/query/management/src/warehouse/warehouse_mgr.rs +++ b/src/query/management/src/warehouse/warehouse_mgr.rs @@ -37,7 +37,6 @@ use databend_common_meta_types::TxnReply; use databend_common_meta_types::TxnRequest; use crate::warehouse::warehouse_api::SelectedNode; -use crate::warehouse::warehouse_api::SelectedNodes; use crate::warehouse::warehouse_api::SystemManagedCluster; use crate::warehouse::warehouse_api::SystemManagedWarehouse; use crate::warehouse::warehouse_api::WarehouseInfo; @@ -65,20 +64,19 @@ impl WarehouseMgr { Ok(WarehouseMgr { metastore, lift_time, - // warehouses: - // all online node of tenant + // Prefix for all online nodes of the tenant node_key_prefix: format!( "{}/{}/online_nodes", WAREHOUSE_API_KEY_PREFIX, escape_for_key(tenant)? ), - // all computing cluster of tenant + // Prefix for all online computing clusters of the tenant meta_key_prefix: format!( "{}/{}/online_clusters", WAREHOUSE_API_KEY_PREFIX, escape_for_key(tenant)? ), - // all warehouse of tenant(required compatible with all versions) + // Prefix for all warehouses of the tenant (must ensure compatibility across all versions) warehouse_key_prefix: format!( "{}/v1/{}", WAREHOUSE_META_KEY_PREFIX, @@ -117,7 +115,7 @@ impl WarehouseMgr { )) } - fn cluster_key(&self, node: &NodeInfo) -> Result { + fn cluster_node_key(&self, node: &NodeInfo) -> Result { Ok(format!( "{}/{}/{}/{}", self.meta_key_prefix, @@ -127,6 +125,22 @@ impl WarehouseMgr { )) } + fn warehouse_info_key(&self, warehouse: &str) -> Result { + Ok(format!( + "{}/{}", + self.warehouse_key_prefix, + escape_for_key(warehouse)? + )) + } + + // Unload the warehouse and cluster from the node. + // 1. Used when a node is removed from the cluster. + // 2. For cluster_node_key: node_info, since its warehouse and cluster are already encoded in the key, we do not need to write the warehouse and cluster into its value again. + fn unload_warehouse_info(&self, node: &mut NodeInfo) { + node.cluster_id = String::new(); + node.warehouse_id = String::new(); + } + async fn upsert_self_managed(&self, mut node: NodeInfo, seq: MatchSeq) -> Result { if node.warehouse_id.is_empty() || node.cluster_id.is_empty() { return Err(ErrorCode::InvalidWarehouse( @@ -145,18 +159,13 @@ impl WarehouseMgr { Some(self.lift_time), )); - let warehouse_node_key = self.cluster_key(&node)?; + let cluster_node_key = self.cluster_node_key(&node)?; let warehouse_info = WarehouseInfo::SelfManaged(node.warehouse_id.clone()); - let warehouse_info_key = format!( - "{}/{}", - self.warehouse_key_prefix, - escape_for_key(&node.warehouse_id)? - ); + let warehouse_info_key = self.warehouse_info_key(&node.warehouse_id)?; - node.cluster_id = String::new(); - node.warehouse_id = String::new(); + self.unload_warehouse_info(&mut node); txn.if_then.push(TxnOp::put_with_ttl( - warehouse_node_key, + cluster_node_key, serde_json::to_vec(&node)?, Some(self.lift_time), )); @@ -242,12 +251,11 @@ impl WarehouseMgr { // If the warehouse has already been assigned. if !node.cluster_id.is_empty() && !node.warehouse_id.is_empty() { - let cluster_key = self.cluster_key(&node)?; + let cluster_node_key = self.cluster_node_key(&node)?; - node.cluster_id = String::new(); - node.warehouse_id = String::new(); + self.unload_warehouse_info(&mut node); txn.if_then.push(TxnOp::put_with_ttl( - cluster_key, + cluster_node_key, serde_json::to_vec(&node)?, Some(self.lift_time), )); @@ -356,12 +364,12 @@ impl WarehouseMgr { } async fn consistent_warehouse_info(&self, id: &str) -> Result { - let warehouse_key = format!("{}/{}", self.warehouse_key_prefix, escape_for_key(id)?); + let warehouse_info_key = self.warehouse_info_key(id)?; let nodes_prefix = format!("{}/{}/", self.meta_key_prefix, escape_for_key(id)?); 'retry: for _idx in 0..64 { - let Some(before_info) = self.metastore.get_kv(&warehouse_key).await? else { + let Some(before_info) = self.metastore.get_kv(&warehouse_info_key).await? else { return Err(ErrorCode::UnknownWarehouse(format!( "Unknown warehouse or self managed warehouse {:?}", id @@ -389,7 +397,7 @@ impl WarehouseMgr { ))); } - let condition = map_condition(&warehouse_key, MatchSeq::Exact(before_info.seq)); + let condition = map_condition(&warehouse_info_key, MatchSeq::Exact(before_info.seq)); after_txn.condition.push(condition); match self.metastore.transaction(after_txn).await? { @@ -450,22 +458,18 @@ impl WarehouseMgr { let mut txn = TxnRequest::default(); let node_key = self.node_key(node_info)?; - let cluster_key = self.cluster_key(node_info)?; + let cluster_node_key = self.cluster_node_key(node_info)?; if consistent_info.consistent_nodes.len() == 1 && consistent_info.consistent_nodes[0].node_info.id == node_info.id { - let warehouse_key = format!( - "{}/{}", - self.warehouse_key_prefix, - escape_for_key(&node_info.warehouse_id)? - ); + let warehouse_info_key = self.warehouse_info_key(&node_info.warehouse_id)?; txn.condition.push(map_condition( - &warehouse_key, + &warehouse_info_key, MatchSeq::Exact(consistent_info.info_seq), )); - txn.if_then.push(TxnOp::delete(warehouse_key)); + txn.if_then.push(TxnOp::delete(warehouse_info_key)); for consistent_node in consistent_info.consistent_nodes { txn.condition.push(map_condition( @@ -473,14 +477,14 @@ impl WarehouseMgr { MatchSeq::Exact(consistent_node.node_seq), )); txn.condition.push(map_condition( - &cluster_key, + &cluster_node_key, MatchSeq::Exact(consistent_node.cluster_seq), )); } } txn.if_then.push(TxnOp::delete(node_key)); - txn.if_then.push(TxnOp::delete(cluster_key)); + txn.if_then.push(TxnOp::delete(cluster_node_key)); if self.metastore.transaction(txn).await?.success { return Ok(()); @@ -517,7 +521,7 @@ impl WarehouseMgr { async fn pick_assign_warehouse_node( &self, warehouse: &str, - nodes: &HashMap, + nodes: &HashMap>, ) -> Result>> { let mut selected_nodes = HashMap::with_capacity(nodes.len()); @@ -735,34 +739,29 @@ impl WarehouseApi for WarehouseMgr { let mut delete_txn = TxnRequest::default(); - let warehouse_key = format!( - "{}/{}", - self.warehouse_key_prefix, - escape_for_key(&warehouse)? - ); + let warehouse_info_key = self.warehouse_info_key(&warehouse)?; delete_txn.condition.push(map_condition( - &warehouse_key, + &warehouse_info_key, MatchSeq::Exact(consistent_info.info_seq), )); - delete_txn.if_then.push(TxnOp::delete(warehouse_key)); + delete_txn.if_then.push(TxnOp::delete(warehouse_info_key)); for mut consistent_node in consistent_info.consistent_nodes { let node_key = self.node_key(&consistent_node.node_info)?; - let cluster_key = self.cluster_key(&consistent_node.node_info)?; + let cluster_node_key = self.cluster_node_key(&consistent_node.node_info)?; delete_txn.condition.push(map_condition( &node_key, MatchSeq::Exact(consistent_node.node_seq), )); delete_txn.condition.push(map_condition( - &cluster_key, + &cluster_node_key, MatchSeq::Exact(consistent_node.cluster_seq), )); - delete_txn.if_then.push(TxnOp::delete(cluster_key)); - consistent_node.node_info.cluster_id = String::new(); - consistent_node.node_info.warehouse_id = String::new(); + delete_txn.if_then.push(TxnOp::delete(cluster_node_key)); + self.unload_warehouse_info(&mut consistent_node.node_info); delete_txn.if_then.push(TxnOp::put_with_ttl( node_key, serde_json::to_vec(&consistent_node.node_info)?, @@ -783,7 +782,7 @@ impl WarehouseApi for WarehouseMgr { )) } - async fn create_warehouse(&self, warehouse: String, nodes: SelectedNodes) -> Result<()> { + async fn create_warehouse(&self, warehouse: String, nodes: Vec) -> Result<()> { if warehouse.is_empty() { return Err(ErrorCode::InvalidWarehouse("Warehouse name is empty.")); } @@ -815,29 +814,24 @@ impl WarehouseApi for WarehouseMgr { Some(self.lift_time * 4), )); - let cluster_key = self.cluster_key(&node)?; + let cluster_node_key = self.cluster_node_key(&node)?; - node.cluster_id = String::new(); - node.warehouse_id = String::new(); + self.unload_warehouse_info(&mut node); txn.condition - .push(map_condition(&cluster_key, MatchSeq::Exact(0))); + .push(map_condition(&cluster_node_key, MatchSeq::Exact(0))); txn.if_then.push(TxnOp::put_with_ttl( - cluster_key, + cluster_node_key, serde_json::to_vec(&node)?, Some(self.lift_time * 4), )); } - let warehouse_key = format!( - "{}/{}", - self.warehouse_key_prefix, - escape_for_key(&warehouse)? - ); + let warehouse_info_key = self.warehouse_info_key(&warehouse)?; txn.condition - .push(map_condition(&warehouse_key, MatchSeq::Exact(0))); + .push(map_condition(&warehouse_info_key, MatchSeq::Exact(0))); txn.if_then.push(TxnOp::put( - warehouse_key.clone(), + warehouse_info_key.clone(), serde_json::to_vec(&WarehouseInfo::SystemManaged(SystemManagedWarehouse { id: GlobalUniqName::unique(), status: "Running".to_string(), @@ -847,7 +841,7 @@ impl WarehouseApi for WarehouseMgr { })]), }))?, )); - txn.else_then.push(TxnOp::get(warehouse_key)); + txn.else_then.push(TxnOp::get(warehouse_info_key)); return match self.metastore.transaction(txn).await? { res if res.success => Ok(()), @@ -938,7 +932,7 @@ impl WarehouseApi for WarehouseMgr { node.warehouse_id = warehouse.clone(); let node_key = self.node_key(&node)?; - let cluster_key = self.cluster_key(&node)?; + let cluster_node_key = self.cluster_node_key(&node)?; resume_txn .condition @@ -950,13 +944,12 @@ impl WarehouseApi for WarehouseMgr { Some(self.lift_time * 4), )); - node.cluster_id = String::new(); - node.warehouse_id = String::new(); + self.unload_warehouse_info(&mut node); resume_txn .condition - .push(map_condition(&cluster_key, MatchSeq::Exact(0))); + .push(map_condition(&cluster_node_key, MatchSeq::Exact(0))); resume_txn.if_then.push(TxnOp::put_with_ttl( - cluster_key, + cluster_node_key, serde_json::to_vec(&node)?, Some(self.lift_time * 4), )); @@ -966,18 +959,14 @@ impl WarehouseApi for WarehouseMgr { } } - let warehouse_key = format!( - "{}/{}", - self.warehouse_key_prefix, - escape_for_key(&warehouse)? - ); + let warehouse_info_key = self.warehouse_info_key(&warehouse)?; resume_txn.condition.push(map_condition( - &warehouse_key, + &warehouse_info_key, MatchSeq::Exact(warehouse_info.info_seq), )); resume_txn.if_then.push(TxnOp::put( - warehouse_key.clone(), + warehouse_info_key.clone(), serde_json::to_vec(&warehouse_info.warehouse_info)?, )); @@ -1019,37 +1008,32 @@ impl WarehouseApi for WarehouseMgr { let mut suspend_txn = TxnRequest::default(); - let warehouse_key = format!( - "{}/{}", - self.warehouse_key_prefix, - escape_for_key(&warehouse)? - ); + let warehouse_info_key = self.warehouse_info_key(&warehouse)?; suspend_txn.condition.push(map_condition( - &warehouse_key, + &warehouse_info_key, MatchSeq::Exact(consistent_info.info_seq), )); suspend_txn.if_then.push(TxnOp::put( - warehouse_key, + warehouse_info_key, serde_json::to_vec(&consistent_info.warehouse_info)?, )); for mut consistent_node in consistent_info.consistent_nodes { let node_key = self.node_key(&consistent_node.node_info)?; - let cluster_key = self.cluster_key(&consistent_node.node_info)?; + let cluster_node_key = self.cluster_node_key(&consistent_node.node_info)?; suspend_txn.condition.push(map_condition( &node_key, MatchSeq::Exact(consistent_node.node_seq), )); suspend_txn.condition.push(map_condition( - &cluster_key, + &cluster_node_key, MatchSeq::Exact(consistent_node.cluster_seq), )); - suspend_txn.if_then.push(TxnOp::delete(cluster_key)); - consistent_node.node_info.cluster_id = String::new(); - consistent_node.node_info.warehouse_id = String::new(); + suspend_txn.if_then.push(TxnOp::delete(cluster_node_key)); + self.unload_warehouse_info(&mut consistent_node.node_info); suspend_txn.if_then.push(TxnOp::put_with_ttl( node_key, serde_json::to_vec(&consistent_node.node_info)?, @@ -1105,64 +1089,60 @@ impl WarehouseApi for WarehouseMgr { let mut rename_txn = TxnRequest::default(); - let old_warehouse_key = format!( - "{}/{}", - self.warehouse_key_prefix, - escape_for_key(¤t)? - ); + let old_warehouse_info_key = self.warehouse_info_key(¤t)?; - let new_warehouse_key = - format!("{}/{}", self.warehouse_key_prefix, escape_for_key(&to)?); + let new_warehouse_info_key = self.warehouse_info_key(&to)?; rename_txn.condition.push(map_condition( - &old_warehouse_key, + &old_warehouse_info_key, MatchSeq::Exact(consistent_info.info_seq), )); rename_txn .condition - .push(map_condition(&new_warehouse_key, MatchSeq::Exact(0))); + .push(map_condition(&new_warehouse_info_key, MatchSeq::Exact(0))); rename_txn .else_then - .push(TxnOp::get(new_warehouse_key.clone())); + .push(TxnOp::get(new_warehouse_info_key.clone())); - rename_txn.if_then.push(TxnOp::delete(old_warehouse_key)); + rename_txn + .if_then + .push(TxnOp::delete(old_warehouse_info_key)); rename_txn.if_then.push(TxnOp::put( - new_warehouse_key, + new_warehouse_info_key, serde_json::to_vec(&consistent_info.warehouse_info)?, )); for mut consistent_node in consistent_info.consistent_nodes { let node_key = self.node_key(&consistent_node.node_info)?; - let old_cluster_key = self.cluster_key(&consistent_node.node_info)?; + let old_cluster_node_key = self.cluster_node_key(&consistent_node.node_info)?; consistent_node.node_info.warehouse_id = to.clone(); - let new_cluster_key = self.cluster_key(&consistent_node.node_info)?; + let new_cluster_node_key = self.cluster_node_key(&consistent_node.node_info)?; rename_txn.condition.push(map_condition( &node_key, MatchSeq::Exact(consistent_node.node_seq), )); rename_txn.condition.push(map_condition( - &old_cluster_key, + &old_cluster_node_key, MatchSeq::Exact(consistent_node.cluster_seq), )); rename_txn .condition - .push(map_condition(&new_cluster_key, MatchSeq::Exact(0))); + .push(map_condition(&new_cluster_node_key, MatchSeq::Exact(0))); rename_txn.if_then.push(TxnOp::put_with_ttl( node_key, serde_json::to_vec(&consistent_node.node_info)?, Some(self.lift_time * 4), )); - consistent_node.node_info.cluster_id = String::new(); - consistent_node.node_info.warehouse_id = String::new(); - rename_txn.if_then.push(TxnOp::delete(old_cluster_key)); + self.unload_warehouse_info(&mut consistent_node.node_info); + rename_txn.if_then.push(TxnOp::delete(old_cluster_node_key)); rename_txn.if_then.push(TxnOp::put_with_ttl( - new_cluster_key.clone(), + new_cluster_node_key.clone(), serde_json::to_vec(&consistent_node.node_info)?, Some(self.lift_time * 4), )); @@ -1207,7 +1187,7 @@ impl WarehouseApi for WarehouseMgr { &self, warehouse: String, cluster: String, - nodes: SelectedNodes, + nodes: Vec, ) -> Result<()> { if warehouse.is_empty() { return Err(ErrorCode::InvalidWarehouse("Warehouse name is empty.")); @@ -1263,40 +1243,36 @@ impl WarehouseApi for WarehouseMgr { } }?; - let warehouse_key = format!( - "{}/{}", - self.warehouse_key_prefix, - escape_for_key(&warehouse)? - ); + let warehouse_info_key = self.warehouse_info_key(&warehouse)?; create_cluster_txn.condition.push(map_condition( - &warehouse_key, + &warehouse_info_key, MatchSeq::Exact(consistent_info.info_seq), )); create_cluster_txn.if_then.push(TxnOp::put( - warehouse_key.clone(), + warehouse_info_key.clone(), serde_json::to_vec(&consistent_info.warehouse_info)?, )); // lock all cluster state for consistent_node in consistent_info.consistent_nodes { let node_key = self.node_key(&consistent_node.node_info)?; - let cluster_key = self.cluster_key(&consistent_node.node_info)?; + let cluster_node_key = self.cluster_node_key(&consistent_node.node_info)?; create_cluster_txn.condition.push(map_condition( &node_key, MatchSeq::Exact(consistent_node.node_seq), )); create_cluster_txn.condition.push(map_condition( - &cluster_key, + &cluster_node_key, MatchSeq::Exact(consistent_node.cluster_seq), )); } for (seq, mut node) in selected_nodes { let node_key = self.node_key(&node)?; - let cluster_key = self.cluster_key(&node)?; + let cluster_node_key = self.cluster_node_key(&node)?; create_cluster_txn .condition @@ -1307,13 +1283,12 @@ impl WarehouseApi for WarehouseMgr { Some(self.lift_time * 4), )); - node.cluster_id = String::new(); - node.warehouse_id = String::new(); + self.unload_warehouse_info(&mut node); create_cluster_txn .condition - .push(map_condition(&cluster_key, MatchSeq::Exact(0))); + .push(map_condition(&cluster_node_key, MatchSeq::Exact(0))); create_cluster_txn.if_then.push(TxnOp::put_with_ttl( - cluster_key, + cluster_node_key, serde_json::to_vec(&node)?, Some(self.lift_time * 4), )); @@ -1378,42 +1353,39 @@ impl WarehouseApi for WarehouseMgr { } }?; - let warehouse_key = format!( - "{}/{}", - self.warehouse_key_prefix, - escape_for_key(&warehouse)? - ); + let warehouse_info_key = self.warehouse_info_key(&warehouse)?; drop_cluster_txn.condition.push(map_condition( - &warehouse_key, + &warehouse_info_key, MatchSeq::Exact(consistent_info.info_seq), )); drop_cluster_txn.if_then.push(TxnOp::put( - warehouse_key.clone(), + warehouse_info_key.clone(), serde_json::to_vec(&consistent_info.warehouse_info)?, )); // lock all cluster state for mut consistent_node in consistent_info.consistent_nodes { let node_key = self.node_key(&consistent_node.node_info)?; - let cluster_key = self.cluster_key(&consistent_node.node_info)?; + let cluster_node_key = self.cluster_node_key(&consistent_node.node_info)?; drop_cluster_txn.condition.push(map_condition( &node_key, MatchSeq::Exact(consistent_node.node_seq), )); drop_cluster_txn.condition.push(map_condition( - &cluster_key, + &cluster_node_key, MatchSeq::Exact(consistent_node.cluster_seq), )); if consistent_node.node_info.cluster_id == cluster { // Remove node - consistent_node.node_info.cluster_id = String::new(); - consistent_node.node_info.warehouse_id = String::new(); + self.unload_warehouse_info(&mut consistent_node.node_info); - drop_cluster_txn.if_then.push(TxnOp::delete(cluster_key)); + drop_cluster_txn + .if_then + .push(TxnOp::delete(cluster_node_key)); drop_cluster_txn.if_then.push(TxnOp::put_with_ttl( node_key, serde_json::to_vec(&consistent_node.node_info)?, @@ -1482,40 +1454,36 @@ impl WarehouseApi for WarehouseMgr { } }?; - let warehouse_key = format!( - "{}/{}", - self.warehouse_key_prefix, - escape_for_key(&warehouse)? - ); + let warehouse_info_key = self.warehouse_info_key(&warehouse)?; rename_cluster_txn.condition.push(map_condition( - &warehouse_key, + &warehouse_info_key, MatchSeq::Exact(consistent_info.info_seq), )); rename_cluster_txn.if_then.push(TxnOp::put( - warehouse_key.clone(), + warehouse_info_key.clone(), serde_json::to_vec(&consistent_info.warehouse_info)?, )); // lock all cluster state for mut consistent_node in consistent_info.consistent_nodes { let node_key = self.node_key(&consistent_node.node_info)?; - let old_cluster_key = self.cluster_key(&consistent_node.node_info)?; + let old_cluster_node_key = self.cluster_node_key(&consistent_node.node_info)?; rename_cluster_txn.condition.push(map_condition( &node_key, MatchSeq::Exact(consistent_node.node_seq), )); rename_cluster_txn.condition.push(map_condition( - &old_cluster_key, + &old_cluster_node_key, MatchSeq::Exact(consistent_node.cluster_seq), )); if consistent_node.node_info.cluster_id == cur { // rename node consistent_node.node_info.cluster_id = to.clone(); - let new_cluster_key = self.cluster_key(&consistent_node.node_info)?; + let new_cluster_node_key = self.cluster_node_key(&consistent_node.node_info)?; rename_cluster_txn.if_then.push(TxnOp::put_with_ttl( node_key, @@ -1526,12 +1494,12 @@ impl WarehouseApi for WarehouseMgr { consistent_node.node_info.cluster_id = String::new(); rename_cluster_txn .condition - .push(map_condition(&new_cluster_key, MatchSeq::Exact(0))); + .push(map_condition(&new_cluster_node_key, MatchSeq::Exact(0))); rename_cluster_txn .if_then - .push(TxnOp::delete(old_cluster_key)); + .push(TxnOp::delete(old_cluster_node_key)); rename_cluster_txn.if_then.push(TxnOp::put_with_ttl( - new_cluster_key, + new_cluster_node_key, serde_json::to_vec(&consistent_node.node_info)?, Some(self.lift_time * 4), )); @@ -1554,7 +1522,7 @@ impl WarehouseApi for WarehouseMgr { async fn assign_warehouse_nodes( &self, warehouse: String, - nodes: HashMap, + nodes: HashMap>, ) -> Result<()> { if warehouse.is_empty() { return Err(ErrorCode::InvalidWarehouse("Warehouse name is empty.")); @@ -1599,33 +1567,29 @@ impl WarehouseApi for WarehouseMgr { } }?; - let warehouse_key = format!( - "{}/{}", - self.warehouse_key_prefix, - escape_for_key(&warehouse)? - ); + let warehouse_info_key = self.warehouse_info_key(&warehouse)?; add_cluster_node_txn.condition.push(map_condition( - &warehouse_key, + &warehouse_info_key, MatchSeq::Exact(consistent_info.info_seq), )); add_cluster_node_txn.if_then.push(TxnOp::put( - warehouse_key.clone(), + warehouse_info_key.clone(), serde_json::to_vec(&consistent_info.warehouse_info)?, )); // lock all cluster state for consistent_node in consistent_info.consistent_nodes { let node_key = self.node_key(&consistent_node.node_info)?; - let cluster_key = self.cluster_key(&consistent_node.node_info)?; + let cluster_node_key = self.cluster_node_key(&consistent_node.node_info)?; add_cluster_node_txn.condition.push(map_condition( &node_key, MatchSeq::Exact(consistent_node.node_seq), )); add_cluster_node_txn.condition.push(map_condition( - &cluster_key, + &cluster_node_key, MatchSeq::Exact(consistent_node.cluster_seq), )); } @@ -1633,7 +1597,7 @@ impl WarehouseApi for WarehouseMgr { for selected_nodes in selected_nodes.into_values() { for (seq, mut node) in selected_nodes { let node_key = self.node_key(&node)?; - let cluster_key = self.cluster_key(&node)?; + let cluster_node_key = self.cluster_node_key(&node)?; add_cluster_node_txn .condition @@ -1644,13 +1608,12 @@ impl WarehouseApi for WarehouseMgr { Some(self.lift_time * 4), )); - node.cluster_id = String::new(); - node.warehouse_id = String::new(); + self.unload_warehouse_info(&mut node); add_cluster_node_txn .condition - .push(map_condition(&cluster_key, MatchSeq::Exact(0))); + .push(map_condition(&cluster_node_key, MatchSeq::Exact(0))); add_cluster_node_txn.if_then.push(TxnOp::put_with_ttl( - cluster_key, + cluster_node_key, serde_json::to_vec(&node)?, Some(self.lift_time * 4), )); @@ -1673,7 +1636,7 @@ impl WarehouseApi for WarehouseMgr { async fn unassign_warehouse_nodes( &self, warehouse: &str, - nodes: HashMap, + nodes: HashMap>, ) -> Result<()> { if warehouse.is_empty() { return Err(ErrorCode::InvalidWarehouse("Warehouse name is empty.")); @@ -1726,33 +1689,29 @@ impl WarehouseApi for WarehouseMgr { } }?; - let warehouse_key = format!( - "{}/{}", - self.warehouse_key_prefix, - escape_for_key(warehouse)? - ); + let warehouse_info_key = self.warehouse_info_key(warehouse)?; drop_cluster_node_txn.condition.push(map_condition( - &warehouse_key, + &warehouse_info_key, MatchSeq::Exact(consistent_info.info_seq), )); drop_cluster_node_txn.if_then.push(TxnOp::put( - warehouse_key.clone(), + warehouse_info_key.clone(), serde_json::to_vec(&consistent_info.warehouse_info)?, )); // lock all cluster state for mut consistent_node in consistent_info.consistent_nodes { let node_key = self.node_key(&consistent_node.node_info)?; - let cluster_key = self.cluster_key(&consistent_node.node_info)?; + let cluster_node_key = self.cluster_node_key(&consistent_node.node_info)?; drop_cluster_node_txn.condition.push(map_condition( &node_key, MatchSeq::Exact(consistent_node.node_seq), )); drop_cluster_node_txn.condition.push(map_condition( - &cluster_key, + &cluster_node_key, MatchSeq::Exact(consistent_node.cluster_seq), )); @@ -1760,13 +1719,12 @@ impl WarehouseApi for WarehouseMgr { if let Some(remove_node) = v.pop() { let SelectedNode::Random(node_group) = remove_node; if consistent_node.node_info.runtime_node_group == node_group { - consistent_node.node_info.cluster_id = String::new(); - consistent_node.node_info.warehouse_id = String::new(); + self.unload_warehouse_info(&mut consistent_node.node_info); consistent_node.node_info.runtime_node_group = None; drop_cluster_node_txn .if_then - .push(TxnOp::delete(cluster_key)); + .push(TxnOp::delete(cluster_node_key)); drop_cluster_node_txn.if_then.push(TxnOp::put_with_ttl( node_key, serde_json::to_vec(&consistent_node.node_info)?, From d680947aa94fcc58918e7bcd85d7de871e09ec8c Mon Sep 17 00:00:00 2001 From: zhang2014 Date: Sun, 29 Dec 2024 19:52:42 +0800 Subject: [PATCH 30/53] feat(cluster): support custom management cluster --- src/query/management/src/warehouse/warehouse_mgr.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/query/management/src/warehouse/warehouse_mgr.rs b/src/query/management/src/warehouse/warehouse_mgr.rs index 504a34a74b4ce..70c690d172e2f 100644 --- a/src/query/management/src/warehouse/warehouse_mgr.rs +++ b/src/query/management/src/warehouse/warehouse_mgr.rs @@ -66,13 +66,13 @@ impl WarehouseMgr { lift_time, // Prefix for all online nodes of the tenant node_key_prefix: format!( - "{}/{}/online_nodes", + "{}/{}/online_nodes_v2", WAREHOUSE_API_KEY_PREFIX, escape_for_key(tenant)? ), // Prefix for all online computing clusters of the tenant meta_key_prefix: format!( - "{}/{}/online_clusters", + "{}/{}/online_clusters_v2", WAREHOUSE_API_KEY_PREFIX, escape_for_key(tenant)? ), From b1d432310edfc4407d1f17a7c4ebcc6ff99d9fc6 Mon Sep 17 00:00:00 2001 From: zhang2014 Date: Sun, 29 Dec 2024 21:00:09 +0800 Subject: [PATCH 31/53] feat(cluster): support custom management cluster --- .../management/src/warehouse/warehouse_mgr.rs | 6 ++--- src/query/management/tests/it/warehouse.rs | 24 +++++++++---------- 2 files changed, 15 insertions(+), 15 deletions(-) diff --git a/src/query/management/src/warehouse/warehouse_mgr.rs b/src/query/management/src/warehouse/warehouse_mgr.rs index 70c690d172e2f..cd416f1edecab 100644 --- a/src/query/management/src/warehouse/warehouse_mgr.rs +++ b/src/query/management/src/warehouse/warehouse_mgr.rs @@ -42,7 +42,7 @@ use crate::warehouse::warehouse_api::SystemManagedWarehouse; use crate::warehouse::warehouse_api::WarehouseInfo; use crate::warehouse::WarehouseApi; -pub static WAREHOUSE_API_KEY_PREFIX: &str = "__fd_clusters_v5"; +pub static WAREHOUSE_API_KEY_PREFIX: &str = "__fd_clusters_v6"; pub static WAREHOUSE_META_KEY_PREFIX: &str = "__fd_warehouses"; pub struct WarehouseMgr { @@ -66,13 +66,13 @@ impl WarehouseMgr { lift_time, // Prefix for all online nodes of the tenant node_key_prefix: format!( - "{}/{}/online_nodes_v2", + "{}/{}/online_nodes", WAREHOUSE_API_KEY_PREFIX, escape_for_key(tenant)? ), // Prefix for all online computing clusters of the tenant meta_key_prefix: format!( - "{}/{}/online_clusters_v2", + "{}/{}/online_clusters", WAREHOUSE_API_KEY_PREFIX, escape_for_key(tenant)? ), diff --git a/src/query/management/tests/it/warehouse.rs b/src/query/management/tests/it/warehouse.rs index a7c897e2a0c9a..ea5dd7ff3eee6 100644 --- a/src/query/management/tests/it/warehouse.rs +++ b/src/query/management/tests/it/warehouse.rs @@ -60,10 +60,10 @@ async fn test_successfully_add_self_managed_node() -> Result<()> { let mut node_info_1 = self_managed_node("test_node_1"); warehouse_manager.start_node(node_info_1.clone()).await?; - let node_key = "__fd_clusters_v5/test%2dtenant%2did/online_nodes/test_node_1"; + let node_key = "__fd_clusters_v6/test%2dtenant%2did/online_nodes/test_node_1"; assert_key_value(&kv, node_key, serde_json::to_vec(&node_info_1)?).await; - let warehouse_key = "__fd_clusters_v5/test%2dtenant%2did/online_clusters/test%2dcluster%2did/test%2dcluster%2did/test_node_1"; + let warehouse_key = "__fd_clusters_v6/test%2dtenant%2did/online_clusters/test%2dcluster%2did/test%2dcluster%2did/test_node_1"; node_info_1.cluster_id = String::new(); node_info_1.warehouse_id = String::new(); @@ -81,10 +81,10 @@ async fn test_successfully_add_self_managed_node() -> Result<()> { let mut node_info_2 = self_managed_node("test_node_2"); warehouse_manager.start_node(node_info_2.clone()).await?; - let node_key = "__fd_clusters_v5/test%2dtenant%2did/online_nodes/test_node_2"; + let node_key = "__fd_clusters_v6/test%2dtenant%2did/online_nodes/test_node_2"; assert_key_value(&kv, node_key, serde_json::to_vec(&node_info_2)?).await; - let warehouse_key = "__fd_clusters_v5/test%2dtenant%2did/online_clusters/test%2dcluster%2did/test%2dcluster%2did/test_node_2"; + let warehouse_key = "__fd_clusters_v6/test%2dtenant%2did/online_clusters/test%2dcluster%2did/test%2dcluster%2did/test_node_2"; node_info_2.cluster_id = String::new(); node_info_2.warehouse_id = String::new(); @@ -109,7 +109,7 @@ async fn test_already_exists_add_self_managed_node() -> Result<()> { let node_info = self_managed_node("test_node_1"); warehouse_manager.start_node(node_info.clone()).await?; - let node_key = "__fd_clusters_v5/test%2dtenant%2did/online_nodes/test_node_1"; + let node_key = "__fd_clusters_v6/test%2dtenant%2did/online_nodes/test_node_1"; assert_key_value(&kv, node_key, serde_json::to_vec(&node_info)?).await; // add already exists self-managed node and get failure @@ -214,11 +214,11 @@ async fn test_successfully_heartbeat_self_managed_node() -> Result<()> { let mut node_info = self_managed_node("test_node"); let seq = warehouse_manager.start_node(node_info.clone()).await?; - let info_key = "__fd_clusters_v5/test%2dtenant%2did/online_nodes/test_node"; + let info_key = "__fd_clusters_v6/test%2dtenant%2did/online_nodes/test_node"; assert_key_value(&kv, info_key, serde_json::to_vec(&node_info)?).await; assert_key_expire(&kv, info_key, Duration::from_mins(50)).await; - let warehouse_key = "__fd_clusters_v5/test%2dtenant%2did/online_clusters/test%2dcluster%2did/test%2dcluster%2did/test_node"; + let warehouse_key = "__fd_clusters_v6/test%2dtenant%2did/online_clusters/test%2dcluster%2did/test%2dcluster%2did/test_node"; let mut warehouse_node = node_info.clone(); warehouse_node.cluster_id = String::new(); warehouse_node.warehouse_id = String::new(); @@ -271,10 +271,10 @@ async fn test_successfully_create_system_managed_warehouse() -> Result<()> { let (kv, warehouse_manager, nodes) = nodes(Duration::from_mins(30), 2).await?; for node in &nodes { - let online_node = format!("__fd_clusters_v5/test%2dtenant%2did/online_nodes/{}", node); + let online_node = format!("__fd_clusters_v6/test%2dtenant%2did/online_nodes/{}", node); assert_key_seq(&kv, &online_node, MatchSeq::GE(1)).await; let warehouse_node = format!( - "__fd_clusters_v5/test%2dtenant%2did/online_clusters/test%2dcluster%2did/default/{}", + "__fd_clusters_v6/test%2dtenant%2did/online_clusters/test%2dcluster%2did/default/{}", node ); assert_no_key(&kv, &warehouse_node).await; @@ -288,10 +288,10 @@ async fn test_successfully_create_system_managed_warehouse() -> Result<()> { create_warehouse.await?; for node in &nodes { - let online_node = format!("__fd_clusters_v5/test%2dtenant%2did/online_nodes/{}", node); + let online_node = format!("__fd_clusters_v6/test%2dtenant%2did/online_nodes/{}", node); assert_key_seq(&kv, &online_node, MatchSeq::GE(1)).await; let warehouse_node = format!( - "__fd_clusters_v5/test%2dtenant%2did/online_clusters/test_warehouse/default/{}", + "__fd_clusters_v6/test%2dtenant%2did/online_clusters/test_warehouse/default/{}", node ); assert_key_seq(&kv, &warehouse_node, MatchSeq::GE(1)).await; @@ -733,7 +733,7 @@ async fn test_rename_warehouses() -> Result<()> { warehouse_manager.shutdown_node(self_managed_node_1).await?; let warehouse_node_key = format!( - "__fd_clusters_v5/test%2dtenant%2did/online_clusters/test_warehouse/default/{}", + "__fd_clusters_v6/test%2dtenant%2did/online_clusters/test_warehouse/default/{}", &nodes[0] ); From 885ebf4b9de69ffc45b6e99e6e1f9de6d1f7862a Mon Sep 17 00:00:00 2001 From: zhang2014 Date: Mon, 30 Dec 2024 10:27:59 +0800 Subject: [PATCH 32/53] feat(cluster): support custom management cluster --- .../management/src/warehouse/warehouse_mgr.rs | 33 ++++++++++++------- 1 file changed, 22 insertions(+), 11 deletions(-) diff --git a/src/query/management/src/warehouse/warehouse_mgr.rs b/src/query/management/src/warehouse/warehouse_mgr.rs index cd416f1edecab..fffffe5ed0d07 100644 --- a/src/query/management/src/warehouse/warehouse_mgr.rs +++ b/src/query/management/src/warehouse/warehouse_mgr.rs @@ -45,12 +45,23 @@ use crate::warehouse::WarehouseApi; pub static WAREHOUSE_API_KEY_PREFIX: &str = "__fd_clusters_v6"; pub static WAREHOUSE_META_KEY_PREFIX: &str = "__fd_warehouses"; +// example: +// __fd_clusters_v6/test_tenant/online_nodes +// |- /RUV9DQArNnP4Hej4A74f07: NodeInfo { id: "RUV9DQArNnP4Hej4A74f07", secret: "", cpu_nums: 0, version: 0, http_address: "", flight_address: "", discovery_address: "", binary_version: "", node_type: SystemManaged, node_group: Some("test"), cluster_id: "", warehouse_id: "", runtime_node_group: Some("test") } +// |- /9a9DU1KufVmSEIDSPnFLZ: NodeInfo { id: "9a9DU1KufVmSEIDSPnFLZ", secret: "", cpu_nums: 0, version: 0, http_address: "", flight_address: "", discovery_address: "", binary_version: "", node_type: SystemManaged, node_group: None, cluster_id: “test_cluster”, warehouse_id: “test_warehouse”, runtime_node_group: None } +// |- /5jTQoMBGJb4TLHeD4pjoa: NodeInfo { id: "5jTQoMBGJb4TLHeD4pjoa", secret: "", cpu_nums: 0, version: 0, http_address: "", flight_address: "", discovery_address: "", binary_version: "", node_type: SystemManaged, node_group: None, cluster_id: “test_cluster”, warehouse_id: “test_warehouse”, runtime_node_group: None } +// +// __fd_clusters_v6/test_tenant/online_clusters +// |- /test_warehouse/test_cluster/9a9DU1KufVmSEIDSPnFLZ: NodeInfo { id: "9a9DU1KufVmSEIDSPnFLZ", secret: "", cpu_nums: 0, version: 0, http_address: "", flight_address: "", discovery_address: "", binary_version: "", node_type: SystemManaged, node_group: None, cluster_id: "", warehouse_id: "", runtime_node_group: None } +// |- /test_warehouse/test_cluster/5jTQoMBGJb4TLHeD4pjoa: NodeInfo { id: "5jTQoMBGJb4TLHeD4pjoa", secret: "", cpu_nums: 0, version: 0, http_address: "", flight_address: "", discovery_address: "", binary_version: "", node_type: SystemManaged, node_group: None, cluster_id: "", warehouse_id: "", runtime_node_group: None } +// __fd_warehouses/v1/test_tenant +// |- /test_warehouse: SystemManaged(SystemManagedWarehouse { id: "udcQQnniiBBFmfP9jziSy4", status: "Running", display_name: "test_warehouse", clusters: {"test_cluster": SystemManagedCluster { nodes: [Random(Some("test")), Random(None)] }} }) pub struct WarehouseMgr { metastore: MetaStore, lift_time: Duration, node_key_prefix: String, - meta_key_prefix: String, - warehouse_key_prefix: String, + cluster_node_key_prefix: String, + warehouse_info_key_prefix: String, } impl WarehouseMgr { @@ -71,13 +82,13 @@ impl WarehouseMgr { escape_for_key(tenant)? ), // Prefix for all online computing clusters of the tenant - meta_key_prefix: format!( + cluster_node_key_prefix: format!( "{}/{}/online_clusters", WAREHOUSE_API_KEY_PREFIX, escape_for_key(tenant)? ), // Prefix for all warehouses of the tenant (must ensure compatibility across all versions) - warehouse_key_prefix: format!( + warehouse_info_key_prefix: format!( "{}/v1/{}", WAREHOUSE_META_KEY_PREFIX, escape_for_key(tenant)?, @@ -118,7 +129,7 @@ impl WarehouseMgr { fn cluster_node_key(&self, node: &NodeInfo) -> Result { Ok(format!( "{}/{}/{}/{}", - self.meta_key_prefix, + self.cluster_node_key_prefix, escape_for_key(&node.warehouse_id)?, escape_for_key(&node.cluster_id)?, escape_for_key(&node.id)? @@ -128,7 +139,7 @@ impl WarehouseMgr { fn warehouse_info_key(&self, warehouse: &str) -> Result { Ok(format!( "{}/{}", - self.warehouse_key_prefix, + self.warehouse_info_key_prefix, escape_for_key(warehouse)? )) } @@ -366,7 +377,7 @@ impl WarehouseMgr { async fn consistent_warehouse_info(&self, id: &str) -> Result { let warehouse_info_key = self.warehouse_info_key(id)?; - let nodes_prefix = format!("{}/{}/", self.meta_key_prefix, escape_for_key(id)?); + let nodes_prefix = format!("{}/{}/", self.cluster_node_key_prefix, escape_for_key(id)?); 'retry: for _idx in 0..64 { let Some(before_info) = self.metastore.get_kv(&warehouse_info_key).await? else { @@ -673,7 +684,7 @@ impl WarehouseApi for WarehouseMgr { if !node_info.cluster_id.is_empty() && !node_info.warehouse_id.is_empty() { txn.if_then.push(TxnOp::delete(format!( "{}/{}/{}/{}", - self.meta_key_prefix, + self.cluster_node_key_prefix, escape_for_key(&node_info.warehouse_id)?, escape_for_key(&node_info.cluster_id)?, escape_for_key(&node_info.id)? @@ -1054,7 +1065,7 @@ impl WarehouseApi for WarehouseMgr { async fn list_warehouses(&self) -> Result> { let values = self .metastore - .prefix_list_kv(&self.warehouse_key_prefix) + .prefix_list_kv(&self.warehouse_info_key_prefix) .await?; let mut warehouses = Vec::with_capacity(values.len()); @@ -1750,7 +1761,7 @@ impl WarehouseApi for WarehouseMgr { ) -> Result> { let cluster_prefix = format!( "{}/{}/{}/", - self.meta_key_prefix, + self.cluster_node_key_prefix, escape_for_key(warehouse)?, escape_for_key(cluster)? ); @@ -1804,7 +1815,7 @@ impl WarehouseApi for WarehouseMgr { let cluster_prefix = format!( "{}/{}/{}/", - self.meta_key_prefix, + self.cluster_node_key_prefix, escape_for_key(&self_info.warehouse_id)?, escape_for_key(&self_info.cluster_id)? ); From 81d8b317ea6525ebe1bdf484df60b9bcb5e41f3a Mon Sep 17 00:00:00 2001 From: zhang2014 Date: Mon, 30 Dec 2024 10:42:13 +0800 Subject: [PATCH 33/53] feat(cluster): support custom management cluster --- src/query/management/src/warehouse/warehouse_mgr.rs | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/src/query/management/src/warehouse/warehouse_mgr.rs b/src/query/management/src/warehouse/warehouse_mgr.rs index fffffe5ed0d07..cfb1d192b9de5 100644 --- a/src/query/management/src/warehouse/warehouse_mgr.rs +++ b/src/query/management/src/warehouse/warehouse_mgr.rs @@ -47,12 +47,12 @@ pub static WAREHOUSE_META_KEY_PREFIX: &str = "__fd_warehouses"; // example: // __fd_clusters_v6/test_tenant/online_nodes -// |- /RUV9DQArNnP4Hej4A74f07: NodeInfo { id: "RUV9DQArNnP4Hej4A74f07", secret: "", cpu_nums: 0, version: 0, http_address: "", flight_address: "", discovery_address: "", binary_version: "", node_type: SystemManaged, node_group: Some("test"), cluster_id: "", warehouse_id: "", runtime_node_group: Some("test") } -// |- /9a9DU1KufVmSEIDSPnFLZ: NodeInfo { id: "9a9DU1KufVmSEIDSPnFLZ", secret: "", cpu_nums: 0, version: 0, http_address: "", flight_address: "", discovery_address: "", binary_version: "", node_type: SystemManaged, node_group: None, cluster_id: “test_cluster”, warehouse_id: “test_warehouse”, runtime_node_group: None } -// |- /5jTQoMBGJb4TLHeD4pjoa: NodeInfo { id: "5jTQoMBGJb4TLHeD4pjoa", secret: "", cpu_nums: 0, version: 0, http_address: "", flight_address: "", discovery_address: "", binary_version: "", node_type: SystemManaged, node_group: None, cluster_id: “test_cluster”, warehouse_id: “test_warehouse”, runtime_node_group: None } +// |- /RUV9DQArNnP4Hej4A74f07: NodeInfo { id: "RUV9DQArNnP4Hej4A74f07", secret: "", cpu_nums: 0, version: 0, http_address: "", flight_address: "", discovery_address: "", binary_version: "", node_type: SystemManaged, node_group: None, cluster_id: "", warehouse_id: "", runtime_node_group: None } +// |- /9a9DU1KufVmSEIDSOnFLZ: NodeInfo { id: "9a9DU1KufVmSEIDSOnFLZ", secret: "", cpu_nums: 0, version: 0, http_address: "", flight_address: "", discovery_address: "", binary_version: "", node_type: SystemManaged, node_group: Some("test"), cluster_id: "test_cluster", warehouse_id: "test_warehouse", runtime_node_group: Some("test") } +// |- /5jTQoMBGJb4TLHeD4pjoa: NodeInfo { id: "5jTQoMBGJb4TLHeD4pjoa", secret: "", cpu_nums: 0, version: 0, http_address: "", flight_address: "", discovery_address: "", binary_version: "", node_type: SystemManaged, node_group: None, cluster_id: "test_cluster", warehouse_id: "test_warehouse", runtime_node_group: None } // // __fd_clusters_v6/test_tenant/online_clusters -// |- /test_warehouse/test_cluster/9a9DU1KufVmSEIDSPnFLZ: NodeInfo { id: "9a9DU1KufVmSEIDSPnFLZ", secret: "", cpu_nums: 0, version: 0, http_address: "", flight_address: "", discovery_address: "", binary_version: "", node_type: SystemManaged, node_group: None, cluster_id: "", warehouse_id: "", runtime_node_group: None } +// |- /test_warehouse/test_cluster/9a9DU1KufVmSEIDSOnFLZ: NodeInfo { id: "9a9DU1KufVmSEIDSOnFLZ", secret: "", cpu_nums: 0, version: 0, http_address: "", flight_address: "", discovery_address: "", binary_version: "", node_type: SystemManaged, node_group: None, cluster_id: "", warehouse_id: "", runtime_node_group: None } // |- /test_warehouse/test_cluster/5jTQoMBGJb4TLHeD4pjoa: NodeInfo { id: "5jTQoMBGJb4TLHeD4pjoa", secret: "", cpu_nums: 0, version: 0, http_address: "", flight_address: "", discovery_address: "", binary_version: "", node_type: SystemManaged, node_group: None, cluster_id: "", warehouse_id: "", runtime_node_group: None } // __fd_warehouses/v1/test_tenant // |- /test_warehouse: SystemManaged(SystemManagedWarehouse { id: "udcQQnniiBBFmfP9jziSy4", status: "Running", display_name: "test_warehouse", clusters: {"test_cluster": SystemManagedCluster { nodes: [Random(Some("test")), Random(None)] }} }) From 630c1a99a18cf7bc25c94bef6cf11fd478656096 Mon Sep 17 00:00:00 2001 From: zhang2014 Date: Mon, 30 Dec 2024 13:24:57 +0800 Subject: [PATCH 34/53] feat(cluster): support custom management cluster --- src/query/management/src/warehouse/warehouse_mgr.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/query/management/src/warehouse/warehouse_mgr.rs b/src/query/management/src/warehouse/warehouse_mgr.rs index cfb1d192b9de5..32b516e5c4e69 100644 --- a/src/query/management/src/warehouse/warehouse_mgr.rs +++ b/src/query/management/src/warehouse/warehouse_mgr.rs @@ -215,7 +215,7 @@ impl WarehouseMgr { return match response.responses.pop().and_then(|x| x.response) { Some(Response::Get(data)) => match data.value { None => Ok(response), - Some(value) if value.seq == 0 => Ok(response), + Some(value) if value.seq == exact_seq => Ok(response), Some(value) => match serde_json::from_slice(&value.data)? { WarehouseInfo::SystemManaged(_) => { Err(ErrorCode::WarehouseAlreadyExists( From 577a8890236d9867231d2249f80963aad77fa0c5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=BC=A0=E7=82=8E=E6=B3=BC?= Date: Mon, 30 Dec 2024 17:26:06 +0800 Subject: [PATCH 35/53] refactor: improve and clean up `warehouse_mgr::upsert_self_managed()` ### Improvements: - Added retry mechanism with a fallback in the retry loop. - Return an error if an unexpected response is received when `TxnGetResponse` is expected. - Refined quit-retry condition: now only triggered when the seq of `NodeInfo` changes. ### Refactoring: - Simplified and decoupled nested branching for better readability and maintainability. - Consolidated related logic, e.g., building `txn if_then` operations in a single place. - Differentiated `NodeInfo` with and without warehouse-related information. ### Documentation: - Added details explaining behavioral differences between insert and update modes. --- src/meta/types/src/match_seq.rs | 8 + src/meta/types/src/proto_ext/txn_ext.rs | 16 ++ .../management/src/warehouse/warehouse_api.rs | 26 ++- .../management/src/warehouse/warehouse_mgr.rs | 190 +++++++++++------- 4 files changed, 165 insertions(+), 75 deletions(-) diff --git a/src/meta/types/src/match_seq.rs b/src/meta/types/src/match_seq.rs index ea514378e0632..90caa7876b4ca 100644 --- a/src/meta/types/src/match_seq.rs +++ b/src/meta/types/src/match_seq.rs @@ -18,6 +18,7 @@ use std::fmt::Formatter; use serde::Deserialize; use serde::Serialize; +use crate::protobuf as pb; use crate::seq_value::SeqV; use crate::ConflictSeq; @@ -68,6 +69,13 @@ impl MatchSeqExt<&Option>> for MatchSeq { } } +impl MatchSeqExt<&Option> for MatchSeq { + fn match_seq(&self, sv: &Option) -> Result<(), ConflictSeq> { + let seq = sv.as_ref().map_or(0, |sv| sv.seq); + self.match_seq(seq) + } +} + impl MatchSeqExt<&SeqV> for MatchSeq { fn match_seq(&self, sv: &SeqV) -> Result<(), ConflictSeq> { let seq = sv.seq; diff --git a/src/meta/types/src/proto_ext/txn_ext.rs b/src/meta/types/src/proto_ext/txn_ext.rs index 80d1598b33aed..5efc4a9e30b53 100644 --- a/src/meta/types/src/proto_ext/txn_ext.rs +++ b/src/meta/types/src/proto_ext/txn_ext.rs @@ -343,6 +343,22 @@ impl pb::TxnOpResponse { })), } } + + /// Consumes and returns the response as a `Get` response if it is one. + pub fn into_get(self) -> Option { + match self.response { + Some(pb::txn_op_response::Response::Get(resp)) => Some(resp), + _ => None, + } + } + + /// Returns the response as a `Get` response if it is one. + pub fn as_get(&self) -> Option<&pb::TxnGetResponse> { + match &self.response { + Some(pb::txn_op_response::Response::Get(resp)) => Some(resp), + _ => None, + } + } } impl pb::TxnGetResponse { diff --git a/src/query/management/src/warehouse/warehouse_api.rs b/src/query/management/src/warehouse/warehouse_api.rs index c4ec4df8434e6..93dc31ac14fc6 100644 --- a/src/query/management/src/warehouse/warehouse_api.rs +++ b/src/query/management/src/warehouse/warehouse_api.rs @@ -17,17 +17,33 @@ use std::collections::HashMap; use databend_common_exception::Result; use databend_common_meta_types::NodeInfo; -// Used to describe the nodes in warehouse, excluding the details of the nodes. +/// Databend-query cluster ID. +/// +/// A cluster is a collection of databend nodes for computation. +pub type ClusterId = String; + +/// Name of a node group. +/// +/// It is used to filter nodes in a warehouse when creating a cluster. +pub type NodeGroupName = String; + +/// Specifies how to select nodes in a warehouse without exposing node details. #[derive(serde::Serialize, serde::Deserialize, Clone, Eq, PartialEq, Debug)] pub enum SelectedNode { - // Randomly select a node from the tenant's list of online nodes. The selected node's group must match it if node group is specified. - Random(Option /* node group */), + /// Select a random online node from the tenant's node list. + /// If `NodeGroupName` is specified, only nodes with this group name will be selected. + Random(Option), } -// Used to describe information about the warehouse, including both self-managed and system-managed warehouses to ensure compatibility. +/// Contains warehouse metadata that applies to both self-managed and system-managed warehouses. +/// +/// This enum allows uniform handling of different warehouse types while maintaining their +/// distinct properties. #[derive(serde::Serialize, serde::Deserialize, Eq, PartialEq, Debug)] pub enum WarehouseInfo { - SelfManaged(String /* cluster_id of self-managed cluster */), + /// A warehouse managed by the tenant, i.e., the cluster id is statically configured in config file. + SelfManaged(ClusterId), + /// A warehouse managed by the system, i.e., clusters in it can be dynamically assigned and updated. SystemManaged(SystemManagedWarehouse), } diff --git a/src/query/management/src/warehouse/warehouse_mgr.rs b/src/query/management/src/warehouse/warehouse_mgr.rs index 32b516e5c4e69..b3e9ea357809d 100644 --- a/src/query/management/src/warehouse/warehouse_mgr.rs +++ b/src/query/management/src/warehouse/warehouse_mgr.rs @@ -17,6 +17,7 @@ use std::collections::HashMap; use std::time::Duration; use databend_common_base::base::escape_for_key; +use databend_common_base::base::tokio; use databend_common_base::base::unescape_for_key; use databend_common_base::base::GlobalUniqName; use databend_common_base::vec_ext::VecExt; @@ -24,9 +25,13 @@ use databend_common_exception::ErrorCode; use databend_common_exception::Result; use databend_common_meta_kvapi::kvapi::KVApi; use databend_common_meta_store::MetaStore; +use databend_common_meta_types::anyerror::AnyError; use databend_common_meta_types::txn_op_response::Response; use databend_common_meta_types::ConditionResult; +use databend_common_meta_types::InvalidReply; use databend_common_meta_types::MatchSeq; +use databend_common_meta_types::MatchSeqExt; +use databend_common_meta_types::MetaError; use databend_common_meta_types::NodeInfo; use databend_common_meta_types::NodeType; use databend_common_meta_types::TxnCondition; @@ -35,6 +40,8 @@ use databend_common_meta_types::TxnOp; use databend_common_meta_types::TxnOpResponse; use databend_common_meta_types::TxnReply; use databend_common_meta_types::TxnRequest; +use log::error; +use log::info; use crate::warehouse::warehouse_api::SelectedNode; use crate::warehouse::warehouse_api::SystemManagedCluster; @@ -159,93 +166,125 @@ impl WarehouseMgr { )); } - let mut txn = TxnRequest::default(); - let node_key = self.node_key(&node)?; - - txn.condition.push(map_condition(&node_key, seq)); - - txn.if_then.push(TxnOp::put_with_ttl( - node_key.clone(), - serde_json::to_vec(&node)?, - Some(self.lift_time), - )); + // It's a many-to-one map from `node` to `warehouse`. + // + // In insert mode, there may not be any node registered in the self-managed warehouse, + // thus we need to asset either the warehouse does not exist, + // or the warehouse is a self-managed warehouse. + // Otherwise(warehouse is found a system-managed), return an error immediately. + // + // In update mode(`!insert`), we already assert the seq of node_key, + // and `node` and `warehouse` are always updated together, + // so we can be sure that the warehouse is always a self-managed + // and is safe to update in this transaction. + // If the node changed(node seq does not match), + // just ignore, and no retry is required. + let is_insert = seq == MatchSeq::Exact(0); + + let node_with_wh = node.clone(); + self.unload_warehouse_info(&mut node); + let node_without_wh = node; - let cluster_node_key = self.cluster_node_key(&node)?; - let warehouse_info = WarehouseInfo::SelfManaged(node.warehouse_id.clone()); - let warehouse_info_key = self.warehouse_info_key(&node.warehouse_id)?; + let node_key = self.node_key(&node_with_wh)?; + let cluster_node_key = self.cluster_node_key(&node_with_wh)?; + let warehouse_key = self.warehouse_info_key(&node_with_wh.warehouse_id)?; + let warehouse = WarehouseInfo::SelfManaged(node_with_wh.warehouse_id.clone()); - self.unload_warehouse_info(&mut node); - txn.if_then.push(TxnOp::put_with_ttl( - cluster_node_key, - serde_json::to_vec(&node)?, - Some(self.lift_time), - )); + let mut txn = TxnRequest::default(); - // upsert warehouse info if self-managed. + txn.condition.push(map_condition(&node_key, seq)); - txn.if_then.push(TxnOp::put_with_ttl( - warehouse_info_key.clone(), - serde_json::to_vec(&warehouse_info)?, - Some(self.lift_time), - )); + let ttl = Some(self.lift_time); + txn.if_then = vec![ + TxnOp::put_with_ttl(&node_key, serde_json::to_vec(&node_with_wh)?, ttl), + TxnOp::put_with_ttl( + &cluster_node_key, + serde_json::to_vec(&node_without_wh)?, + ttl, + ), + TxnOp::put_with_ttl(&warehouse_key, serde_json::to_vec(&warehouse)?, ttl), + TxnOp::get(&node_key), + ]; - txn.if_then.push(TxnOp::get(node_key.clone())); - txn.else_then.push(TxnOp::get(node_key.clone())); + txn.else_then.push(TxnOp::get(&node_key)); - if seq != MatchSeq::Exact(0) { + if !is_insert { return Ok(self.metastore.transaction(txn).await?); } - let mut exact_seq = 0; - let mut retry_count = 0; + let mut wh_seq = 0; + let max_retry = 20; - loop { + for retry_count in 0..max_retry { let mut warehouse_txn = txn.clone(); // insert if warehouse info is not exists or SelfManaged warehouse_txn .condition - .push(TxnCondition::eq_seq(warehouse_info_key.clone(), exact_seq)); - warehouse_txn - .else_then - .push(TxnOp::get(warehouse_info_key.clone())); - - return match self.metastore.transaction(warehouse_txn).await? { - mut response if !response.success => { - return match response.responses.pop().and_then(|x| x.response) { - Some(Response::Get(data)) => match data.value { - None => Ok(response), - Some(value) if value.seq == exact_seq => Ok(response), - Some(value) => match serde_json::from_slice(&value.data)? { - WarehouseInfo::SystemManaged(_) => { - Err(ErrorCode::WarehouseAlreadyExists( - "Already exists same name system-managed warehouse.", - )) - } - WarehouseInfo::SelfManaged(_) => match response.responses.first() { - // already exists node. - Some(TxnOpResponse { - response: - Some(Response::Get(TxnGetResponse { - value: Some(value), - .. - })), - }) if value.seq != 0 => Ok(response), - _ => { - log::info!("Self-managed warehouse has already been created by other nodes; attempt to join it. Retry count: {}", retry_count); - retry_count += 1; - exact_seq = value.seq; - continue; - } - }, - }, - }, - _ => Ok(response), - }; + .push(TxnCondition::eq_seq(&warehouse_key, wh_seq)); + warehouse_txn.else_then.push(TxnOp::get(&warehouse_key)); + + let mut txn_reply = self.metastore.transaction(warehouse_txn).await?; + + if txn_reply.success { + return Ok(txn_reply); + } + + // txn_reply includes: 1. get node response 2. get warehouse response + + // 1. Check node state + { + let first = txn_reply.responses.first(); + + let node_get_resp = first.and_then(|r| r.as_get()); + let node_get_resp = node_get_resp.ok_or_else(|| invalid_get_resp(first))?; + + if seq.match_seq(&node_get_resp.value).is_err() { + // Node changed, no more retry is needed. + return Ok(txn_reply); } - response => Ok(response), - }; + } + + // 2. Check warehouse state + { + // The last response is not meant to be returned. + let last = txn_reply.responses.pop(); + + let wh_get_resp = last.as_ref().and_then(|r| r.as_get()); + let wh_get_resp = wh_get_resp.ok_or_else(|| invalid_get_resp(last.as_ref()))?; + + if let Some(wh_seqv) = &wh_get_resp.value { + let wh: WarehouseInfo = serde_json::from_slice(&wh_seqv.data)?; + + match wh { + WarehouseInfo::SystemManaged(_) => { + return Err(ErrorCode::WarehouseAlreadyExists( + "Already exists same name system-managed warehouse.", + )) + } + WarehouseInfo::SelfManaged(_) => { + // Safe to retry with the updated seq of warehouse + } + }; + + info!( + "Self-managed warehouse has already been created by other nodes;\ + attempt to join it. Retry count: {}", + retry_count + ); + }; + + // Retry with the updated seq of warehouse + wh_seq = wh_get_resp.value.as_ref().map_or(0, |v| v.seq); + } + + // upon retry, fallback a little while. + tokio::time::sleep(Duration::from_millis(30 * retry_count)).await; } + + Err(ErrorCode::WarehouseOperateConflict(format!( + "Warehouse operate conflict(tried {max_retry} times)." + ))) } async fn upsert_system_managed(&self, mut node: NodeInfo, seq: MatchSeq) -> Result { @@ -1843,3 +1882,14 @@ impl WarehouseApi for WarehouseMgr { } } } + +/// Build an error indicating that databend-meta responded with an unexpected response, +/// while expecting a TxnGetResponse. +fn invalid_get_resp(resp: Option<&TxnOpResponse>) -> MetaError { + let invalid = InvalidReply::new( + "Invalid response while upsert self managed node", + &AnyError::error("Expect TxnGetResponse"), + ); + error!("{} got: {:?}", invalid, resp); + MetaError::from(invalid) +} From 456f42f925f8296b221391d4352c95f063b6897f Mon Sep 17 00:00:00 2001 From: zhang2014 Date: Thu, 2 Jan 2025 10:51:31 +0800 Subject: [PATCH 36/53] feat(cluster): clean code --- src/meta/types/src/cluster.rs | 45 +++++ src/meta/types/src/proto_ext/txn_ext.rs | 6 +- .../management/src/warehouse/warehouse_mgr.rs | 155 ++++++++---------- 3 files changed, 115 insertions(+), 91 deletions(-) diff --git a/src/meta/types/src/cluster.rs b/src/meta/types/src/cluster.rs index 78b0a9fdf6f1b..0aa10cdc32060 100644 --- a/src/meta/types/src/cluster.rs +++ b/src/meta/types/src/cluster.rs @@ -133,6 +133,51 @@ impl NodeInfo { Ok((addr.ip().to_string(), addr.port())) } + + pub fn assigned_warehouse(&self) -> bool { + !self.warehouse_id.is_empty() && !self.cluster_id.is_empty() + } + + pub fn leave_warehouse(&self) -> NodeInfo { + NodeInfo { + id: self.id.clone(), + secret: self.secret.clone(), + cpu_nums: self.cpu_nums, + version: self.version.clone(), + http_address: self.http_address.clone(), + flight_address: self.flight_address.clone(), + discovery_address: self.discovery_address.clone(), + binary_version: self.binary_version.clone(), + node_type: self.node_type.clone(), + node_group: self.node_group.clone(), + cluster_id: String::new(), + warehouse_id: String::new(), + runtime_node_group: None, + } + } + + pub fn join_warehouse( + &self, + warehouse: String, + cluster: String, + group: Option, + ) -> NodeInfo { + NodeInfo { + id: self.id.clone(), + secret: self.secret.clone(), + cpu_nums: self.cpu_nums, + version: self.version.clone(), + http_address: self.http_address.clone(), + flight_address: self.flight_address.clone(), + discovery_address: self.discovery_address.clone(), + binary_version: self.binary_version.clone(), + node_type: self.node_type.clone(), + node_group: self.node_group.clone(), + cluster_id: cluster, + warehouse_id: warehouse, + runtime_node_group: group, + } + } } #[cfg(test)] diff --git a/src/meta/types/src/proto_ext/txn_ext.rs b/src/meta/types/src/proto_ext/txn_ext.rs index 5efc4a9e30b53..2ca1992a3f3a5 100644 --- a/src/meta/types/src/proto_ext/txn_ext.rs +++ b/src/meta/types/src/proto_ext/txn_ext.rs @@ -352,8 +352,12 @@ impl pb::TxnOpResponse { } } + pub fn as_get(&self) -> &pb::TxnGetResponse { + self.try_as_get().unwrap() + } + /// Returns the response as a `Get` response if it is one. - pub fn as_get(&self) -> Option<&pb::TxnGetResponse> { + pub fn try_as_get(&self) -> Option<&pb::TxnGetResponse> { match &self.response { Some(pb::txn_op_response::Response::Get(resp)) => Some(resp), _ => None, diff --git a/src/query/management/src/warehouse/warehouse_mgr.rs b/src/query/management/src/warehouse/warehouse_mgr.rs index b3e9ea357809d..afbc83c8c873a 100644 --- a/src/query/management/src/warehouse/warehouse_mgr.rs +++ b/src/query/management/src/warehouse/warehouse_mgr.rs @@ -236,7 +236,7 @@ impl WarehouseMgr { { let first = txn_reply.responses.first(); - let node_get_resp = first.and_then(|r| r.as_get()); + let node_get_resp = first.map(|r| r.as_get()); let node_get_resp = node_get_resp.ok_or_else(|| invalid_get_resp(first))?; if seq.match_seq(&node_get_resp.value).is_err() { @@ -250,7 +250,7 @@ impl WarehouseMgr { // The last response is not meant to be returned. let last = txn_reply.responses.pop(); - let wh_get_resp = last.as_ref().and_then(|r| r.as_get()); + let wh_get_resp = last.as_ref().and_then(|r| r.try_as_get()); let wh_get_resp = wh_get_resp.ok_or_else(|| invalid_get_resp(last.as_ref()))?; if let Some(wh_seqv) = &wh_get_resp.value { @@ -326,61 +326,54 @@ impl WarehouseMgr { } async fn leave_cluster(&self, node_info: &mut NodeInfo, seq: u64) -> Result { - let mut cluster_id = String::new(); - let mut warehouse_id = String::new(); - let mut runtime_node_group = None; - - std::mem::swap(&mut node_info.cluster_id, &mut cluster_id); - std::mem::swap(&mut node_info.warehouse_id, &mut warehouse_id); - std::mem::swap(&mut node_info.runtime_node_group, &mut runtime_node_group); - - let upsert_node = self.upsert_node(node_info.clone(), MatchSeq::Exact(seq)); - match upsert_node.await { - Err(err) => { - // rollback - std::mem::swap(&mut node_info.cluster_id, &mut cluster_id); - std::mem::swap(&mut node_info.warehouse_id, &mut warehouse_id); - std::mem::swap(&mut node_info.runtime_node_group, &mut runtime_node_group); - Err(err) - } - Ok(response) if !response.success => { - // rollback - std::mem::swap(&mut node_info.cluster_id, &mut cluster_id); - std::mem::swap(&mut node_info.warehouse_id, &mut warehouse_id); - std::mem::swap(&mut node_info.runtime_node_group, &mut runtime_node_group); - Ok(seq) - } - Ok(response) => match response.responses.last() { - Some(TxnOpResponse { - response: Some(Response::Get(TxnGetResponse { value: Some(v), .. })), - }) => Ok(v.seq), - _ => Err(ErrorCode::MetaServiceError("Meta insert failure.")), - }, + let leave_node = node_info.leave_warehouse(); + + let response = self + .upsert_node(leave_node.clone(), MatchSeq::Exact(seq)) + .await?; + + if !response.success { + return Err(ErrorCode::WarehouseOperateConflict( + "Node's status changed while in leave cluster.", + )); + } + + *node_info = leave_node; + + match response.responses.last() { + Some(TxnOpResponse { + response: + Some(Response::Get(TxnGetResponse { + value: Some(seq_v), .. + })), + }) => Ok(seq_v.seq), + _ => Err(ErrorCode::MetaServiceError( + "Invalid response while leave cluster, expect get response", + )), } } - async fn join_cluster(&self, node_info: &mut NodeInfo, seq: u64) -> Result { - let upsert_node = self.upsert_node(node_info.clone(), MatchSeq::Exact(seq)); + async fn join_cluster(&self, node_info: NodeInfo, seq: u64) -> Result { + let response = self + .upsert_node(node_info.clone(), MatchSeq::Exact(seq)) + .await?; - match upsert_node.await { - Err(err) => { - // rollback - // std::mem::swap(&mut node_info.cluster_id, &mut cluster_id); - // std::mem::swap(&mut node_info.warehouse_id, &mut warehouse_id); - Err(err) - } - Ok(response) if !response.success => { - // rollback - // std::mem::swap(&mut node_info.cluster_id, &mut cluster_id); - // std::mem::swap(&mut node_info.warehouse_id, &mut warehouse_id); - Ok(seq) - } - Ok(response) => match response.responses.last() { - Some(TxnOpResponse { - response: Some(Response::Get(TxnGetResponse { value: Some(v), .. })), - }) => Ok(v.seq), - _ => Err(ErrorCode::MetaServiceError("Meta insert failure.")), - }, + if !response.success { + return Err(ErrorCode::WarehouseOperateConflict( + "Node's status changed while in join cluster.", + )); + } + + match response.responses.last() { + Some(TxnOpResponse { + response: + Some(Response::Get(TxnGetResponse { + value: Some(seq_v), .. + })), + }) => Ok(seq_v.seq), + _ => Err(ErrorCode::MetaServiceError( + "Invalid response while join cluster, expect get response", + )), } } @@ -392,21 +385,18 @@ impl WarehouseMgr { }) => match &res.value { None => self.leave_cluster(node, 0).await, Some(value) => { - let node_info = serde_json::from_slice::(&value.data)?; + // node info from meta service + let node_from_meta = serde_json::from_slice::(&value.data)?; // Removed this node from the cluster in other nodes - if !node.cluster_id.is_empty() - && !node.warehouse_id.is_empty() - && node_info.cluster_id.is_empty() - && node_info.warehouse_id.is_empty() - { + if node.assigned_warehouse() && !node_from_meta.assigned_warehouse() { return self.leave_cluster(node, value.seq).await; } - // Added this node to the cluster in other nodes - node.cluster_id = node_info.cluster_id; - node.warehouse_id = node_info.warehouse_id; - self.join_cluster(node, value.seq).await + // Move this node to new cluster + let seq = self.join_cluster(node_from_meta.clone(), value.seq).await?; + *node = node_from_meta; + Ok(seq) } }, _ => Err(ErrorCode::Internal("Miss type while in meta response")), @@ -683,11 +673,13 @@ impl WarehouseApi for WarehouseMgr { #[async_backtrace::framed] #[fastrace::trace] async fn start_node(&self, node: NodeInfo) -> Result { - let res = self.upsert_node(node.clone(), MatchSeq::Exact(0)).await?; + let upsert_reply = self.upsert_node(node.clone(), MatchSeq::Exact(0)).await?; - if res.success { - let Some(Response::Get(response)) = - res.responses.last().and_then(|x| x.response.as_ref()) + if upsert_reply.success { + let Some(Response::Get(response)) = upsert_reply + .responses + .last() + .and_then(|x| x.response.as_ref()) else { return Err(ErrorCode::Internal("Unknown get response")); }; @@ -1839,6 +1831,7 @@ impl WarehouseApi for WarehouseMgr { async fn discover(&self, node_id: &str) -> Result<(bool, Vec)> { let node_key = format!("{}/{}", self.node_key_prefix, escape_for_key(node_id)?); + let Some(seq) = self.metastore.get_kv(&node_key).await? else { return Err(ErrorCode::NotFoundClusterNode(format!( "Node {} is offline, Please restart this node.", @@ -1846,32 +1839,14 @@ impl WarehouseApi for WarehouseMgr { ))); }; - let self_info = serde_json::from_slice::(&seq.data)?; - - if self_info.warehouse_id.is_empty() && self_info.cluster_id.is_empty() { - return Ok((false, vec![self_info])); - } - - let cluster_prefix = format!( - "{}/{}/{}/", - self.cluster_node_key_prefix, - escape_for_key(&self_info.warehouse_id)?, - escape_for_key(&self_info.cluster_id)? - ); - - let values = self.metastore.prefix_list_kv(&cluster_prefix).await?; - - let mut cluster_nodes_info = Vec::with_capacity(values.len()); - for (node_key, value) in values { - let mut cluster_node = serde_json::from_slice::(&value.data)?; + let node = serde_json::from_slice::(&seq.data)?; - cluster_node.id = unescape_for_key(&node_key[cluster_prefix.len()..])?; - cluster_node.cluster_id = self_info.cluster_id.clone(); - cluster_node.warehouse_id = self_info.warehouse_id.clone(); - cluster_nodes_info.push(cluster_node); + if node.warehouse_id.is_empty() && node.cluster_id.is_empty() { + return Ok((false, vec![node])); } - Ok((true, cluster_nodes_info)) + let list_nodes = self.list_warehouse_cluster_nodes(&node.warehouse_id, &node.cluster_id); + Ok((true, list_nodes.await?)) } async fn get_node_info(&self, node_id: &str) -> Result> { From d6f784ec4b6566ffcd0fd0d7573dd404f2ea87d6 Mon Sep 17 00:00:00 2001 From: zhang2014 Date: Thu, 2 Jan 2025 15:40:59 +0800 Subject: [PATCH 37/53] feat(cluster): clean code --- src/meta/types/src/cluster.rs | 25 +- .../management/src/warehouse/warehouse_mgr.rs | 235 +++++++++++------- 2 files changed, 152 insertions(+), 108 deletions(-) diff --git a/src/meta/types/src/cluster.rs b/src/meta/types/src/cluster.rs index 0aa10cdc32060..a5e8f6fcd4965 100644 --- a/src/meta/types/src/cluster.rs +++ b/src/meta/types/src/cluster.rs @@ -143,7 +143,7 @@ impl NodeInfo { id: self.id.clone(), secret: self.secret.clone(), cpu_nums: self.cpu_nums, - version: self.version.clone(), + version: self.version, http_address: self.http_address.clone(), flight_address: self.flight_address.clone(), discovery_address: self.discovery_address.clone(), @@ -155,29 +155,6 @@ impl NodeInfo { runtime_node_group: None, } } - - pub fn join_warehouse( - &self, - warehouse: String, - cluster: String, - group: Option, - ) -> NodeInfo { - NodeInfo { - id: self.id.clone(), - secret: self.secret.clone(), - cpu_nums: self.cpu_nums, - version: self.version.clone(), - http_address: self.http_address.clone(), - flight_address: self.flight_address.clone(), - discovery_address: self.discovery_address.clone(), - binary_version: self.binary_version.clone(), - node_type: self.node_type.clone(), - node_group: self.node_group.clone(), - cluster_id: cluster, - warehouse_id: warehouse, - runtime_node_group: group, - } - } } #[cfg(test)] diff --git a/src/query/management/src/warehouse/warehouse_mgr.rs b/src/query/management/src/warehouse/warehouse_mgr.rs index afbc83c8c873a..94bb53de0f94e 100644 --- a/src/query/management/src/warehouse/warehouse_mgr.rs +++ b/src/query/management/src/warehouse/warehouse_mgr.rs @@ -159,28 +159,13 @@ impl WarehouseMgr { node.warehouse_id = String::new(); } - async fn upsert_self_managed(&self, mut node: NodeInfo, seq: MatchSeq) -> Result { + async fn start_self_managed(&self, node: NodeInfo) -> Result { if node.warehouse_id.is_empty() || node.cluster_id.is_empty() { return Err(ErrorCode::InvalidWarehouse( "The warehouse_id and cluster_id for self managed node must not be empty.", )); } - // It's a many-to-one map from `node` to `warehouse`. - // - // In insert mode, there may not be any node registered in the self-managed warehouse, - // thus we need to asset either the warehouse does not exist, - // or the warehouse is a self-managed warehouse. - // Otherwise(warehouse is found a system-managed), return an error immediately. - // - // In update mode(`!insert`), we already assert the seq of node_key, - // and `node` and `warehouse` are always updated together, - // so we can be sure that the warehouse is always a self-managed - // and is safe to update in this transaction. - // If the node changed(node seq does not match), - // just ignore, and no retry is required. - let is_insert = seq == MatchSeq::Exact(0); - let node_with_wh = node.clone(); self.unload_warehouse_info(&mut node); let node_without_wh = node; @@ -192,7 +177,8 @@ impl WarehouseMgr { let mut txn = TxnRequest::default(); - txn.condition.push(map_condition(&node_key, seq)); + txn.condition + .push(map_condition(&node_key, MatchSeq::Exact(0))); let ttl = Some(self.lift_time); txn.if_then = vec![ @@ -206,11 +192,7 @@ impl WarehouseMgr { TxnOp::get(&node_key), ]; - txn.else_then.push(TxnOp::get(&node_key)); - - if !is_insert { - return Ok(self.metastore.transaction(txn).await?); - } + txn.else_then = vec![TxnOp::get(&node_key), TxnOp::get(&warehouse_key)]; let mut wh_seq = 0; let max_retry = 20; @@ -222,7 +204,6 @@ impl WarehouseMgr { warehouse_txn .condition .push(TxnCondition::eq_seq(&warehouse_key, wh_seq)); - warehouse_txn.else_then.push(TxnOp::get(&warehouse_key)); let mut txn_reply = self.metastore.transaction(warehouse_txn).await?; @@ -239,7 +220,8 @@ impl WarehouseMgr { let node_get_resp = first.map(|r| r.as_get()); let node_get_resp = node_get_resp.ok_or_else(|| invalid_get_resp(first))?; - if seq.match_seq(&node_get_resp.value).is_err() { + let exact_seq = MatchSeq::Exact(0); + if exact_seq.match_seq(&node_get_resp.value).is_err() { // Node changed, no more retry is needed. return Ok(txn_reply); } @@ -287,12 +269,71 @@ impl WarehouseMgr { ))) } - async fn upsert_system_managed(&self, mut node: NodeInfo, seq: MatchSeq) -> Result { + async fn heartbeat_self_managed(&self, node: NodeInfo, seq: MatchSeq) -> Result { + let node_with_wh = node.clone(); + self.unload_warehouse_info(&mut node); + let node_without_wh = node; + + let node_key = self.node_key(&node_with_wh)?; + let cluster_node_key = self.cluster_node_key(&node_with_wh)?; + let warehouse_key = self.warehouse_info_key(&node_with_wh.warehouse_id)?; + let warehouse = WarehouseInfo::SelfManaged(node_with_wh.warehouse_id.clone()); + let mut txn = TxnRequest::default(); - let node_key = self.node_key(&node)?; txn.condition.push(map_condition(&node_key, seq)); + let ttl = Some(self.lift_time); + txn.if_then = vec![ + TxnOp::put_with_ttl(&node_key, serde_json::to_vec(&node_with_wh)?, ttl), + TxnOp::put_with_ttl( + &cluster_node_key, + serde_json::to_vec(&node_without_wh)?, + ttl, + ), + TxnOp::put_with_ttl(&warehouse_key, serde_json::to_vec(&warehouse)?, ttl), + TxnOp::get(&node_key), + ]; + + txn.else_then = vec![TxnOp::get(&node_key), TxnOp::get(&warehouse_key)]; + + let txn_reply = self.metastore.transaction(txn).await?; + + if !txn_reply.success { + // txn_reply includes: 1. get node response 2. get warehouse response + + // Check warehouse state + { + // The last response is not meant to be returned. + let last = txn_reply.responses.pop(); + + let wh_get_resp = last.as_ref().and_then(|r| r.try_as_get()); + let wh_get_resp = wh_get_resp.ok_or_else(|| invalid_get_resp(last.as_ref()))?; + + if let Some(wh_seqv) = &wh_get_resp.value { + let wh: WarehouseInfo = serde_json::from_slice(&wh_seqv.data)?; + + if let WarehouseInfo::SystemManaged(_) = wh { + // an unknown error caused all nodes of the warehouse to fail the heartbeat with the meta service. If someone creates a warehouse with the same name during this period. + return Err(ErrorCode::WarehouseAlreadyExists( + "Already exists same name system-managed warehouse.", + )); + } + }; + } + } + + Ok(txn_reply) + } + + async fn start_system_managed(&self, mut node: NodeInfo) -> Result { + // TODO: find attach warehouse + let mut txn = TxnRequest::default(); + let node_key = self.node_key(&node)?; + + txn.condition + .push(map_condition(&node_key, MatchSeq::Exact(0))); + txn.if_then.push(TxnOp::put_with_ttl( node_key.clone(), serde_json::to_vec(&node)?, @@ -316,23 +357,43 @@ impl WarehouseMgr { Ok(self.metastore.transaction(txn).await?) } - #[fastrace::trace] - #[async_backtrace::framed] - async fn upsert_node(&self, node: NodeInfo, seq: MatchSeq) -> Result { - match node.node_type { - NodeType::SelfManaged => self.upsert_self_managed(node, seq).await, - NodeType::SystemManaged => self.upsert_system_managed(node, seq).await, + async fn heartbeat_system_managed(&self, node: NodeInfo, seq: MatchSeq) -> Result { + let mut txn = TxnRequest::default(); + let node_key = self.node_key(&node)?; + + txn.condition.push(map_condition(&node_key, seq)); + + txn.if_then.push(TxnOp::put_with_ttl( + node_key.clone(), + serde_json::to_vec(&node)?, + Some(self.lift_time), + )); + + // If the warehouse has already been assigned. + if !node.cluster_id.is_empty() && !node.warehouse_id.is_empty() { + let cluster_node_key = self.cluster_node_key(&node)?; + + self.unload_warehouse_info(&mut node); + txn.if_then.push(TxnOp::put_with_ttl( + cluster_node_key, + serde_json::to_vec(&node)?, + Some(self.lift_time), + )); } + + txn.if_then.push(TxnOp::get(node_key.clone())); + txn.else_then.push(TxnOp::get(node_key.clone())); + Ok(self.metastore.transaction(txn).await?) } - async fn leave_cluster(&self, node_info: &mut NodeInfo, seq: u64) -> Result { + async fn leave_warehouse(&self, node_info: &mut NodeInfo, seq: u64) -> Result { let leave_node = node_info.leave_warehouse(); - let response = self - .upsert_node(leave_node.clone(), MatchSeq::Exact(seq)) + let reply = self + .heartbeat_system_managed(leave_node, MatchSeq::Exact(seq)) .await?; - if !response.success { + if !reply.success { return Err(ErrorCode::WarehouseOperateConflict( "Node's status changed while in leave cluster.", )); @@ -340,7 +401,7 @@ impl WarehouseMgr { *node_info = leave_node; - match response.responses.last() { + match reply.responses.last() { Some(TxnOpResponse { response: Some(Response::Get(TxnGetResponse { @@ -353,18 +414,18 @@ impl WarehouseMgr { } } - async fn join_cluster(&self, node_info: NodeInfo, seq: u64) -> Result { - let response = self - .upsert_node(node_info.clone(), MatchSeq::Exact(seq)) + async fn join_warehouse(&self, node_info: NodeInfo, seq: u64) -> Result { + let reply = self + .heartbeat_system_managed(node_info.clone(), MatchSeq::Exact(seq)) .await?; - if !response.success { + if !reply.success { return Err(ErrorCode::WarehouseOperateConflict( "Node's status changed while in join cluster.", )); } - match response.responses.last() { + match reply.responses.last() { Some(TxnOpResponse { response: Some(Response::Get(TxnGetResponse { @@ -379,22 +440,24 @@ impl WarehouseMgr { async fn resolve_conflicts(&self, reply: TxnReply, node: &mut NodeInfo) -> Result { match reply.responses.first() { - None => self.leave_cluster(node, 0).await, + None => self.leave_warehouse(node, 0).await, Some(TxnOpResponse { response: Some(Response::Get(res)), }) => match &res.value { - None => self.leave_cluster(node, 0).await, + None => self.leave_warehouse(node, 0).await, Some(value) => { // node info from meta service let node_from_meta = serde_json::from_slice::(&value.data)?; // Removed this node from the cluster in other nodes if node.assigned_warehouse() && !node_from_meta.assigned_warehouse() { - return self.leave_cluster(node, value.seq).await; + return self.leave_warehouse(node, value.seq).await; } - // Move this node to new cluster - let seq = self.join_cluster(node_from_meta.clone(), value.seq).await?; + // Move this node to new warehouse and cluster + let seq = self + .join_warehouse(node_from_meta.clone(), value.seq) + .await?; *node = node_from_meta; Ok(seq) } @@ -673,28 +736,27 @@ impl WarehouseApi for WarehouseMgr { #[async_backtrace::framed] #[fastrace::trace] async fn start_node(&self, node: NodeInfo) -> Result { - let upsert_reply = self.upsert_node(node.clone(), MatchSeq::Exact(0)).await?; - - if upsert_reply.success { - let Some(Response::Get(response)) = upsert_reply - .responses - .last() - .and_then(|x| x.response.as_ref()) - else { - return Err(ErrorCode::Internal("Unknown get response")); - }; - - let Some(node_info) = &response.value else { - return Err(ErrorCode::MetaServiceError("Add node info failure.")); - }; + let start_reply = match node.node_type { + NodeType::SelfManaged => self.start_self_managed(node).await?, + NodeType::SystemManaged => self.start_system_managed(node).await?, + }; - return Ok(node_info.seq); + if !start_reply.success { + return Err(ErrorCode::ClusterNodeAlreadyExists(format!( + "Node with ID '{}' already exists in the cluster.", + node.id + ))); } - Err(ErrorCode::ClusterNodeAlreadyExists(format!( - "Node with ID '{}' already exists in the cluster.", - node.id - ))) + match start_reply.responses.last() { + Some(TxnOpResponse { + response: + Some(Response::Get(TxnGetResponse { + value: Some(seq_v), .. + })), + }) => Ok(seq_v.seq), + _ => Err(ErrorCode::MetaServiceError("Add node info failure.")), + } } #[async_backtrace::framed] @@ -740,28 +802,33 @@ impl WarehouseApi for WarehouseMgr { } async fn heartbeat_node(&self, node: &mut NodeInfo, seq: u64) -> Result { - if node.node_type == NodeType::SelfManaged { - assert!(!node.cluster_id.is_empty()); - assert!(!node.warehouse_id.is_empty()); - } + let exact = MatchSeq::Exact(seq); + let heartbeat_reply = match node.node_type { + NodeType::SelfManaged => { + assert!(!node.cluster_id.is_empty()); + assert!(!node.warehouse_id.is_empty()); - let res = self.upsert_node(node.clone(), MatchSeq::Exact(seq)).await?; - - match res.success { - true => { - let Some(Response::Get(response)) = - res.responses.last().and_then(|x| x.response.as_ref()) - else { - return Err(ErrorCode::Internal("Unknown get response")); - }; + self.heartbeat_self_managed(node.clone(), exact).await? + } + NodeType::SystemManaged => { + let reply = self.heartbeat_system_managed(node.clone(), exact).await?; - let Some(node_info) = &response.value else { - return Err(ErrorCode::MetaServiceError("Add node info failure.")); - }; + if !reply.success { + return self.resolve_conflicts(reply, node).await; + } - Ok(node_info.seq) + reply } - false => self.resolve_conflicts(res, node).await, + }; + + match heartbeat_reply.responses.last() { + Some(TxnOpResponse { + response: + Some(Response::Get(TxnGetResponse { + value: Some(seq_v), .. + })), + }) => Ok(seq_v.seq), + _ => Err(ErrorCode::MetaServiceError("Heartbeat node info failure.")), } } From 8402d2815e2392b0b96d8ca1fb5e726a9f9b8fa8 Mon Sep 17 00:00:00 2001 From: zhang2014 Date: Thu, 2 Jan 2025 16:07:07 +0800 Subject: [PATCH 38/53] feat(cluster): clean code --- .../management/src/warehouse/warehouse_mgr.rs | 66 +++++++++---------- 1 file changed, 31 insertions(+), 35 deletions(-) diff --git a/src/query/management/src/warehouse/warehouse_mgr.rs b/src/query/management/src/warehouse/warehouse_mgr.rs index 94bb53de0f94e..806d671aa4903 100644 --- a/src/query/management/src/warehouse/warehouse_mgr.rs +++ b/src/query/management/src/warehouse/warehouse_mgr.rs @@ -159,7 +159,7 @@ impl WarehouseMgr { node.warehouse_id = String::new(); } - async fn start_self_managed(&self, node: NodeInfo) -> Result { + async fn start_self_managed(&self, mut node: NodeInfo) -> Result { if node.warehouse_id.is_empty() || node.cluster_id.is_empty() { return Err(ErrorCode::InvalidWarehouse( "The warehouse_id and cluster_id for self managed node must not be empty.", @@ -269,7 +269,7 @@ impl WarehouseMgr { ))) } - async fn heartbeat_self_managed(&self, node: NodeInfo, seq: MatchSeq) -> Result { + async fn heartbeat_self_managed(&self, mut node: NodeInfo, seq: u64) -> Result { let node_with_wh = node.clone(); self.unload_warehouse_info(&mut node); let node_without_wh = node; @@ -281,7 +281,8 @@ impl WarehouseMgr { let mut txn = TxnRequest::default(); - txn.condition.push(map_condition(&node_key, seq)); + txn.condition + .push(map_condition(&node_key, MatchSeq::Exact(seq))); let ttl = Some(self.lift_time); txn.if_then = vec![ @@ -297,7 +298,7 @@ impl WarehouseMgr { txn.else_then = vec![TxnOp::get(&node_key), TxnOp::get(&warehouse_key)]; - let txn_reply = self.metastore.transaction(txn).await?; + let mut txn_reply = self.metastore.transaction(txn).await?; if !txn_reply.success { // txn_reply includes: 1. get node response 2. get warehouse response @@ -357,11 +358,12 @@ impl WarehouseMgr { Ok(self.metastore.transaction(txn).await?) } - async fn heartbeat_system_managed(&self, node: NodeInfo, seq: MatchSeq) -> Result { + async fn heartbeat_system_managed(&self, mut node: NodeInfo, seq: u64) -> Result { let mut txn = TxnRequest::default(); let node_key = self.node_key(&node)?; - txn.condition.push(map_condition(&node_key, seq)); + txn.condition + .push(map_condition(&node_key, MatchSeq::Exact(seq))); txn.if_then.push(TxnOp::put_with_ttl( node_key.clone(), @@ -390,7 +392,7 @@ impl WarehouseMgr { let leave_node = node_info.leave_warehouse(); let reply = self - .heartbeat_system_managed(leave_node, MatchSeq::Exact(seq)) + .heartbeat_system_managed(leave_node.clone(), seq) .await?; if !reply.success { @@ -415,9 +417,7 @@ impl WarehouseMgr { } async fn join_warehouse(&self, node_info: NodeInfo, seq: u64) -> Result { - let reply = self - .heartbeat_system_managed(node_info.clone(), MatchSeq::Exact(seq)) - .await?; + let reply = self.heartbeat_system_managed(node_info, seq).await?; if !reply.success { return Err(ErrorCode::WarehouseOperateConflict( @@ -439,30 +439,26 @@ impl WarehouseMgr { } async fn resolve_conflicts(&self, reply: TxnReply, node: &mut NodeInfo) -> Result { - match reply.responses.first() { + // system-managed heartbeat_node reply include 'get node response' at last response + let get_node_response = reply.responses.last(); + let get_node_response = get_node_response.and_then(TxnOpResponse::try_as_get); + + match get_node_response { None => self.leave_warehouse(node, 0).await, - Some(TxnOpResponse { - response: Some(Response::Get(res)), - }) => match &res.value { - None => self.leave_warehouse(node, 0).await, - Some(value) => { - // node info from meta service - let node_from_meta = serde_json::from_slice::(&value.data)?; - - // Removed this node from the cluster in other nodes - if node.assigned_warehouse() && !node_from_meta.assigned_warehouse() { - return self.leave_warehouse(node, value.seq).await; - } + Some(TxnGetResponse { value: None, .. }) => self.leave_warehouse(node, 0).await, + Some(TxnGetResponse { value: Some(v), .. }) => { + let node_from_meta = serde_json::from_slice::(&v.data)?; - // Move this node to new warehouse and cluster - let seq = self - .join_warehouse(node_from_meta.clone(), value.seq) - .await?; - *node = node_from_meta; - Ok(seq) + // Removed this node from the cluster in other nodes + if node.assigned_warehouse() && !node_from_meta.assigned_warehouse() { + return self.leave_warehouse(node, v.seq).await; } - }, - _ => Err(ErrorCode::Internal("Miss type while in meta response")), + + // Move this node to new warehouse and cluster + let seq = self.join_warehouse(node_from_meta.clone(), v.seq).await?; + *node = node_from_meta; + Ok(seq) + } } } @@ -737,8 +733,8 @@ impl WarehouseApi for WarehouseMgr { #[fastrace::trace] async fn start_node(&self, node: NodeInfo) -> Result { let start_reply = match node.node_type { - NodeType::SelfManaged => self.start_self_managed(node).await?, - NodeType::SystemManaged => self.start_system_managed(node).await?, + NodeType::SelfManaged => self.start_self_managed(node.clone()).await?, + NodeType::SystemManaged => self.start_system_managed(node.clone()).await?, }; if !start_reply.success { @@ -801,8 +797,8 @@ impl WarehouseApi for WarehouseMgr { ))) } - async fn heartbeat_node(&self, node: &mut NodeInfo, seq: u64) -> Result { - let exact = MatchSeq::Exact(seq); + async fn heartbeat_node(&self, node: &mut NodeInfo, exact: u64) -> Result { + // let exact = MatchSeq::Exact(seq); let heartbeat_reply = match node.node_type { NodeType::SelfManaged => { assert!(!node.cluster_id.is_empty()); From 00252102c56fbd7da6ca546f616f46622f2cf3b1 Mon Sep 17 00:00:00 2001 From: zhang2014 Date: Thu, 2 Jan 2025 16:27:54 +0800 Subject: [PATCH 39/53] feat(cluster): clean code --- .../management/src/warehouse/warehouse_mgr.rs | 118 ++++++++---------- 1 file changed, 53 insertions(+), 65 deletions(-) diff --git a/src/query/management/src/warehouse/warehouse_mgr.rs b/src/query/management/src/warehouse/warehouse_mgr.rs index 806d671aa4903..7c62cb4683c64 100644 --- a/src/query/management/src/warehouse/warehouse_mgr.rs +++ b/src/query/management/src/warehouse/warehouse_mgr.rs @@ -403,13 +403,11 @@ impl WarehouseMgr { *node_info = leave_node; - match reply.responses.last() { - Some(TxnOpResponse { - response: - Some(Response::Get(TxnGetResponse { - value: Some(seq_v), .. - })), - }) => Ok(seq_v.seq), + let get_node_response = reply.responses.last(); + let get_node_response = get_node_response.and_then(TxnOpResponse::try_as_get); + + match get_node_response { + Some(TxnGetResponse { value: Some(v), .. }) => Ok(v.seq), _ => Err(ErrorCode::MetaServiceError( "Invalid response while leave cluster, expect get response", )), @@ -425,13 +423,11 @@ impl WarehouseMgr { )); } - match reply.responses.last() { - Some(TxnOpResponse { - response: - Some(Response::Get(TxnGetResponse { - value: Some(seq_v), .. - })), - }) => Ok(seq_v.seq), + let get_node_response = reply.responses.last(); + let get_node_response = get_node_response.and_then(TxnOpResponse::try_as_get); + + match get_node_response { + Some(TxnGetResponse { value: Some(v), .. }) => Ok(v.seq), _ => Err(ErrorCode::MetaServiceError( "Invalid response while join cluster, expect get response", )), @@ -463,11 +459,13 @@ impl WarehouseMgr { } async fn consistent_warehouse_info(&self, id: &str) -> Result { + // Obtain a consistent snapshot of the warehouse nodes using 3 RPC calls + // TODO: meta does not support using list within transactions, which will be optimized in later version let warehouse_info_key = self.warehouse_info_key(id)?; let nodes_prefix = format!("{}/{}/", self.cluster_node_key_prefix, escape_for_key(id)?); - 'retry: for _idx in 0..64 { + 'get_warehouse_snapshot: for _idx in 0..64 { let Some(before_info) = self.metastore.get_kv(&warehouse_info_key).await? else { return Err(ErrorCode::UnknownWarehouse(format!( "Unknown warehouse or self managed warehouse {:?}", @@ -483,64 +481,54 @@ impl WarehouseMgr { for (node_key, value) in values { let suffix = &node_key[nodes_prefix.len()..]; - if let Some((_cluster, node)) = suffix.split_once('/') { - let node_key = format!("{}/{}", self.node_key_prefix, node); - after_txn.if_then.push(TxnOp::get(node_key)); - cluster_node_seq.push(value.seq); - continue; - } + let Some((_cluster, node)) = suffix.split_once('/') else { + return Err(ErrorCode::InvalidWarehouse(format!( + "Node key is invalid {:?}", + node_key + ))); + }; - return Err(ErrorCode::InvalidWarehouse(format!( - "Node key is invalid {:?}", - node_key - ))); + let node_key = format!("{}/{}", self.node_key_prefix, node); + after_txn.if_then.push(TxnOp::get(node_key)); + cluster_node_seq.push(value.seq); } let condition = map_condition(&warehouse_info_key, MatchSeq::Exact(before_info.seq)); after_txn.condition.push(condition); - match self.metastore.transaction(after_txn).await? { - response if response.success => { - let mut consistent_nodes = Vec::with_capacity(response.responses.len()); - for (idx, response) in response.responses.into_iter().enumerate() { - match response.response { - // TODO: maybe ignore none(not need retry) - Some(Response::Get(response)) => match response.value { - Some(value) => { - let node_info = - serde_json::from_slice::(&value.data)?; - - assert_eq!(node_info.warehouse_id, id); - assert!(!node_info.cluster_id.is_empty()); - - consistent_nodes.push(ConsistentNodeInfo { - node_seq: value.seq, - cluster_seq: cluster_node_seq[idx], - node_info, - }); - } - _ => { - continue 'retry; - } - }, - _ => { - continue 'retry; - } - } - } + let txn_reply = self.metastore.transaction(after_txn).await?; - if consistent_nodes.len() == cluster_node_seq.len() { - return Ok(ConsistentWarehouseInfo { - info_seq: before_info.seq, - warehouse_info: serde_json::from_slice(&before_info.data)?, - consistent_nodes, - }); - } - } - _ => { - continue 'retry; - } + if !txn_reply.success { + // The snapshot version status has changed; need to retry obtaining the snapshot. + continue 'get_warehouse_snapshot; + } + + let mut consistent_nodes = Vec::with_capacity(txn_reply.responses.len()); + for (idx, response) in txn_reply.responses.into_iter().enumerate() { + let get_node_info = response.as_get(); + + let Some(seq_v) = get_node_info.value.as_ref() else { + // The node went offline during the get snapshot. + continue 'get_warehouse_snapshot; + }; + + let node_info = serde_json::from_slice::(&seq_v.data)?; + + assert_eq!(node_info.warehouse_id, id); + assert!(!node_info.cluster_id.is_empty()); + + consistent_nodes.push(ConsistentNodeInfo { + node_seq: seq_v.seq, + cluster_seq: cluster_node_seq[idx], + node_info, + }); } + + return Ok(ConsistentWarehouseInfo { + info_seq: before_info.seq, + warehouse_info: serde_json::from_slice(&before_info.data)?, + consistent_nodes, + }); } Err(ErrorCode::Internal( From 19d996c548e1ccdfc28c670d6f31ca0b64d9176f Mon Sep 17 00:00:00 2001 From: zhang2014 Date: Thu, 2 Jan 2025 18:04:50 +0800 Subject: [PATCH 40/53] feat(cluster): clean code --- .../management/src/warehouse/warehouse_mgr.rs | 242 +++++++++--------- 1 file changed, 121 insertions(+), 121 deletions(-) diff --git a/src/query/management/src/warehouse/warehouse_mgr.rs b/src/query/management/src/warehouse/warehouse_mgr.rs index 7c62cb4683c64..a1ca1891ea80d 100644 --- a/src/query/management/src/warehouse/warehouse_mgr.rs +++ b/src/query/management/src/warehouse/warehouse_mgr.rs @@ -112,16 +112,16 @@ fn map_condition(k: &str, seq: MatchSeq) -> TxnCondition { } } -struct ConsistentNodeInfo { +struct NodeInfoSnapshot { node_seq: u64, cluster_seq: u64, node_info: NodeInfo, } -struct ConsistentWarehouseInfo { +struct WarehouseSnapshot { info_seq: u64, warehouse_info: WarehouseInfo, - consistent_nodes: Vec, + snapshot_nodes: Vec, } impl WarehouseMgr { @@ -458,7 +458,7 @@ impl WarehouseMgr { } } - async fn consistent_warehouse_info(&self, id: &str) -> Result { + async fn warehouse_snapshot(&self, id: &str) -> Result { // Obtain a consistent snapshot of the warehouse nodes using 3 RPC calls // TODO: meta does not support using list within transactions, which will be optimized in later version let warehouse_info_key = self.warehouse_info_key(id)?; @@ -503,7 +503,7 @@ impl WarehouseMgr { continue 'get_warehouse_snapshot; } - let mut consistent_nodes = Vec::with_capacity(txn_reply.responses.len()); + let mut snapshot_nodes = Vec::with_capacity(txn_reply.responses.len()); for (idx, response) in txn_reply.responses.into_iter().enumerate() { let get_node_info = response.as_get(); @@ -517,17 +517,17 @@ impl WarehouseMgr { assert_eq!(node_info.warehouse_id, id); assert!(!node_info.cluster_id.is_empty()); - consistent_nodes.push(ConsistentNodeInfo { + snapshot_nodes.push(NodeInfoSnapshot { node_seq: seq_v.seq, cluster_seq: cluster_node_seq[idx], node_info, }); } - return Ok(ConsistentWarehouseInfo { + return Ok(WarehouseSnapshot { info_seq: before_info.seq, warehouse_info: serde_json::from_slice(&before_info.data)?, - consistent_nodes, + snapshot_nodes, }); } @@ -538,34 +538,34 @@ impl WarehouseMgr { async fn shutdown_self_managed_node(&self, node_info: &NodeInfo) -> Result<()> { for _idx in 0..10 { - let consistent_info = self - .consistent_warehouse_info(&node_info.warehouse_id) - .await?; + let warehouse_snapshot = self.warehouse_snapshot(&node_info.warehouse_id).await?; let mut txn = TxnRequest::default(); let node_key = self.node_key(node_info)?; let cluster_node_key = self.cluster_node_key(node_info)?; - if consistent_info.consistent_nodes.len() == 1 - && consistent_info.consistent_nodes[0].node_info.id == node_info.id + // Attempt to delete the warehouse info if the last node is shutting down. + if warehouse_snapshot.snapshot_nodes.len() == 1 + && warehouse_snapshot.snapshot_nodes[0].node_info.id == node_info.id { let warehouse_info_key = self.warehouse_info_key(&node_info.warehouse_id)?; txn.condition.push(map_condition( &warehouse_info_key, - MatchSeq::Exact(consistent_info.info_seq), + MatchSeq::Exact(warehouse_snapshot.info_seq), )); txn.if_then.push(TxnOp::delete(warehouse_info_key)); - for consistent_node in consistent_info.consistent_nodes { + for node_info_snapshot in warehouse_snapshot.snapshot_nodes { txn.condition.push(map_condition( &node_key, - MatchSeq::Exact(consistent_node.node_seq), + MatchSeq::Exact(node_info_snapshot.node_seq), )); + txn.condition.push(map_condition( &cluster_node_key, - MatchSeq::Exact(consistent_node.cluster_seq), + MatchSeq::Exact(node_info_snapshot.cluster_seq), )); } } @@ -608,18 +608,18 @@ impl WarehouseMgr { async fn pick_assign_warehouse_node( &self, warehouse: &str, - nodes: &HashMap>, + warehouse_nodes_matcher: &HashMap>, ) -> Result>> { - let mut selected_nodes = HashMap::with_capacity(nodes.len()); + let mut selected_nodes = HashMap::with_capacity(warehouse_nodes_matcher.len()); let mut grouped_nodes = self.unassigned_nodes().await?; let mut after_assign_node = HashMap::new(); - for (cluster, cluster_node_selector) in nodes { - let mut cluster_selected_nodes = Vec::with_capacity(cluster_node_selector.len()); - for node_selector in cluster_node_selector { - match node_selector { + for (cluster, cluster_node_matcher) in warehouse_nodes_matcher { + let mut cluster_selected_nodes = Vec::with_capacity(cluster_node_matcher.len()); + for selected_node in cluster_node_matcher { + match selected_node { SelectedNode::Random(None) => { let Some(nodes_list) = grouped_nodes.get_mut(&None) else { match after_assign_node.entry(cluster.clone()) { @@ -822,9 +822,9 @@ impl WarehouseApi for WarehouseMgr { } for _idx in 0..10 { - let consistent_info = self.consistent_warehouse_info(&warehouse).await?; + let warehouse_snapshot = self.warehouse_snapshot(&warehouse).await?; - if let WarehouseInfo::SelfManaged(_) = consistent_info.warehouse_info { + if let WarehouseInfo::SelfManaged(_) = warehouse_snapshot.warehouse_info { return Err(ErrorCode::InvalidWarehouse( "Cannot drop self-managed warehouse", )); @@ -836,28 +836,28 @@ impl WarehouseApi for WarehouseMgr { delete_txn.condition.push(map_condition( &warehouse_info_key, - MatchSeq::Exact(consistent_info.info_seq), + MatchSeq::Exact(warehouse_snapshot.info_seq), )); delete_txn.if_then.push(TxnOp::delete(warehouse_info_key)); - for mut consistent_node in consistent_info.consistent_nodes { - let node_key = self.node_key(&consistent_node.node_info)?; - let cluster_node_key = self.cluster_node_key(&consistent_node.node_info)?; + for mut node_snapshot in warehouse_snapshot.snapshot_nodes { + let node_key = self.node_key(&node_snapshot.node_info)?; + let cluster_node_key = self.cluster_node_key(&node_snapshot.node_info)?; delete_txn.condition.push(map_condition( &node_key, - MatchSeq::Exact(consistent_node.node_seq), + MatchSeq::Exact(node_snapshot.node_seq), )); delete_txn.condition.push(map_condition( &cluster_node_key, - MatchSeq::Exact(consistent_node.cluster_seq), + MatchSeq::Exact(node_snapshot.cluster_seq), )); delete_txn.if_then.push(TxnOp::delete(cluster_node_key)); - self.unload_warehouse_info(&mut consistent_node.node_info); + self.unload_warehouse_info(&mut node_snapshot.node_info); delete_txn.if_then.push(TxnOp::put_with_ttl( node_key, - serde_json::to_vec(&consistent_node.node_info)?, + serde_json::to_vec(&node_snapshot.node_info)?, Some(self.lift_time * 4), )); } @@ -965,7 +965,7 @@ impl WarehouseApi for WarehouseMgr { } for _idx in 0..10 { - let mut warehouse_info = self.consistent_warehouse_info(&warehouse).await?; + let mut warehouse_info = self.warehouse_snapshot(&warehouse).await?; let mut need_schedule_cluster = HashMap::new(); warehouse_info.warehouse_info = match warehouse_info.warehouse_info { @@ -1079,9 +1079,9 @@ impl WarehouseApi for WarehouseMgr { } for _idx in 0..10 { - let mut consistent_info = self.consistent_warehouse_info(&warehouse).await?; + let mut warehouse_snapshot = self.warehouse_snapshot(&warehouse).await?; - consistent_info.warehouse_info = match consistent_info.warehouse_info { + warehouse_snapshot.warehouse_info = match warehouse_snapshot.warehouse_info { WarehouseInfo::SelfManaged(_) => Err(ErrorCode::InvalidWarehouse( "Cannot suspend self-managed warehouse", )), @@ -1105,31 +1105,31 @@ impl WarehouseApi for WarehouseMgr { suspend_txn.condition.push(map_condition( &warehouse_info_key, - MatchSeq::Exact(consistent_info.info_seq), + MatchSeq::Exact(warehouse_snapshot.info_seq), )); suspend_txn.if_then.push(TxnOp::put( warehouse_info_key, - serde_json::to_vec(&consistent_info.warehouse_info)?, + serde_json::to_vec(&warehouse_snapshot.warehouse_info)?, )); - for mut consistent_node in consistent_info.consistent_nodes { - let node_key = self.node_key(&consistent_node.node_info)?; - let cluster_node_key = self.cluster_node_key(&consistent_node.node_info)?; + for mut node_snapshot in warehouse_snapshot.snapshot_nodes { + let node_key = self.node_key(&node_snapshot.node_info)?; + let cluster_node_key = self.cluster_node_key(&node_snapshot.node_info)?; suspend_txn.condition.push(map_condition( &node_key, - MatchSeq::Exact(consistent_node.node_seq), + MatchSeq::Exact(node_snapshot.node_seq), )); suspend_txn.condition.push(map_condition( &cluster_node_key, - MatchSeq::Exact(consistent_node.cluster_seq), + MatchSeq::Exact(node_snapshot.cluster_seq), )); suspend_txn.if_then.push(TxnOp::delete(cluster_node_key)); - self.unload_warehouse_info(&mut consistent_node.node_info); + self.unload_warehouse_info(&mut node_snapshot.node_info); suspend_txn.if_then.push(TxnOp::put_with_ttl( node_key, - serde_json::to_vec(&consistent_node.node_info)?, + serde_json::to_vec(&node_snapshot.node_info)?, Some(self.lift_time * 4), )); } @@ -1168,9 +1168,9 @@ impl WarehouseApi for WarehouseMgr { } for _idx in 0..10 { - let mut consistent_info = self.consistent_warehouse_info(¤t).await?; + let mut warehouse_snapshot = self.warehouse_snapshot(¤t).await?; - consistent_info.warehouse_info = match consistent_info.warehouse_info { + warehouse_snapshot.warehouse_info = match warehouse_snapshot.warehouse_info { WarehouseInfo::SelfManaged(_) => Err(ErrorCode::InvalidWarehouse( "Cannot rename self-managed warehouse", )), @@ -1188,7 +1188,7 @@ impl WarehouseApi for WarehouseMgr { rename_txn.condition.push(map_condition( &old_warehouse_info_key, - MatchSeq::Exact(consistent_info.info_seq), + MatchSeq::Exact(warehouse_snapshot.info_seq), )); rename_txn @@ -1203,24 +1203,24 @@ impl WarehouseApi for WarehouseMgr { .push(TxnOp::delete(old_warehouse_info_key)); rename_txn.if_then.push(TxnOp::put( new_warehouse_info_key, - serde_json::to_vec(&consistent_info.warehouse_info)?, + serde_json::to_vec(&warehouse_snapshot.warehouse_info)?, )); - for mut consistent_node in consistent_info.consistent_nodes { - let node_key = self.node_key(&consistent_node.node_info)?; - let old_cluster_node_key = self.cluster_node_key(&consistent_node.node_info)?; + for mut node_snapshot in warehouse_snapshot.snapshot_nodes { + let node_key = self.node_key(&node_snapshot.node_info)?; + let old_cluster_node_key = self.cluster_node_key(&node_snapshot.node_info)?; - consistent_node.node_info.warehouse_id = to.clone(); + node_snapshot.node_info.warehouse_id = to.clone(); - let new_cluster_node_key = self.cluster_node_key(&consistent_node.node_info)?; + let new_cluster_node_key = self.cluster_node_key(&node_snapshot.node_info)?; rename_txn.condition.push(map_condition( &node_key, - MatchSeq::Exact(consistent_node.node_seq), + MatchSeq::Exact(node_snapshot.node_seq), )); rename_txn.condition.push(map_condition( &old_cluster_node_key, - MatchSeq::Exact(consistent_node.cluster_seq), + MatchSeq::Exact(node_snapshot.cluster_seq), )); rename_txn @@ -1229,14 +1229,14 @@ impl WarehouseApi for WarehouseMgr { rename_txn.if_then.push(TxnOp::put_with_ttl( node_key, - serde_json::to_vec(&consistent_node.node_info)?, + serde_json::to_vec(&node_snapshot.node_info)?, Some(self.lift_time * 4), )); - self.unload_warehouse_info(&mut consistent_node.node_info); + self.unload_warehouse_info(&mut node_snapshot.node_info); rename_txn.if_then.push(TxnOp::delete(old_cluster_node_key)); rename_txn.if_then.push(TxnOp::put_with_ttl( new_cluster_node_key.clone(), - serde_json::to_vec(&consistent_node.node_info)?, + serde_json::to_vec(&node_snapshot.node_info)?, Some(self.lift_time * 4), )); } @@ -1267,10 +1267,10 @@ impl WarehouseApi for WarehouseMgr { return Err(ErrorCode::InvalidWarehouse("Warehouse name is empty.")); } - let consistent_info = self.consistent_warehouse_info(&warehouse).await?; + let warehouse_snapshot = self.warehouse_snapshot(&warehouse).await?; - Ok(consistent_info - .consistent_nodes + Ok(warehouse_snapshot + .snapshot_nodes .into_iter() .map(|x| x.node_info) .collect()) @@ -1308,9 +1308,9 @@ impl WarehouseApi for WarehouseMgr { let mut create_cluster_txn = TxnRequest::default(); - let mut consistent_info = self.consistent_warehouse_info(&warehouse).await?; + let mut warehouse_snapshot = self.warehouse_snapshot(&warehouse).await?; - consistent_info.warehouse_info = match consistent_info.warehouse_info { + warehouse_snapshot.warehouse_info = match warehouse_snapshot.warehouse_info { WarehouseInfo::SelfManaged(_) => Err(ErrorCode::InvalidWarehouse(format!( "Cannot add cluster for warehouse {:?}, because it's self-managed warehouse.", warehouse @@ -1340,26 +1340,26 @@ impl WarehouseApi for WarehouseMgr { create_cluster_txn.condition.push(map_condition( &warehouse_info_key, - MatchSeq::Exact(consistent_info.info_seq), + MatchSeq::Exact(warehouse_snapshot.info_seq), )); create_cluster_txn.if_then.push(TxnOp::put( warehouse_info_key.clone(), - serde_json::to_vec(&consistent_info.warehouse_info)?, + serde_json::to_vec(&warehouse_snapshot.warehouse_info)?, )); // lock all cluster state - for consistent_node in consistent_info.consistent_nodes { - let node_key = self.node_key(&consistent_node.node_info)?; - let cluster_node_key = self.cluster_node_key(&consistent_node.node_info)?; + for node_snapshot in warehouse_snapshot.snapshot_nodes { + let node_key = self.node_key(&node_snapshot.node_info)?; + let cluster_node_key = self.cluster_node_key(&node_snapshot.node_info)?; create_cluster_txn.condition.push(map_condition( &node_key, - MatchSeq::Exact(consistent_node.node_seq), + MatchSeq::Exact(node_snapshot.node_seq), )); create_cluster_txn.condition.push(map_condition( &cluster_node_key, - MatchSeq::Exact(consistent_node.cluster_seq), + MatchSeq::Exact(node_snapshot.cluster_seq), )); } @@ -1414,9 +1414,9 @@ impl WarehouseApi for WarehouseMgr { for _idx in 0..10 { let mut drop_cluster_txn = TxnRequest::default(); - let mut consistent_info = self.consistent_warehouse_info(&warehouse).await?; + let mut warehouse_snapshot = self.warehouse_snapshot(&warehouse).await?; - consistent_info.warehouse_info = match consistent_info.warehouse_info { + warehouse_snapshot.warehouse_info = match warehouse_snapshot.warehouse_info { WarehouseInfo::SelfManaged(_) => Err(ErrorCode::InvalidWarehouse(format!( "Cannot add cluster for warehouse {:?}, because it's self-managed warehouse.", warehouse @@ -1450,38 +1450,38 @@ impl WarehouseApi for WarehouseMgr { drop_cluster_txn.condition.push(map_condition( &warehouse_info_key, - MatchSeq::Exact(consistent_info.info_seq), + MatchSeq::Exact(warehouse_snapshot.info_seq), )); drop_cluster_txn.if_then.push(TxnOp::put( warehouse_info_key.clone(), - serde_json::to_vec(&consistent_info.warehouse_info)?, + serde_json::to_vec(&warehouse_snapshot.warehouse_info)?, )); // lock all cluster state - for mut consistent_node in consistent_info.consistent_nodes { - let node_key = self.node_key(&consistent_node.node_info)?; - let cluster_node_key = self.cluster_node_key(&consistent_node.node_info)?; + for mut node_snapshot in warehouse_snapshot.snapshot_nodes { + let node_key = self.node_key(&node_snapshot.node_info)?; + let cluster_node_key = self.cluster_node_key(&node_snapshot.node_info)?; drop_cluster_txn.condition.push(map_condition( &node_key, - MatchSeq::Exact(consistent_node.node_seq), + MatchSeq::Exact(node_snapshot.node_seq), )); drop_cluster_txn.condition.push(map_condition( &cluster_node_key, - MatchSeq::Exact(consistent_node.cluster_seq), + MatchSeq::Exact(node_snapshot.cluster_seq), )); - if consistent_node.node_info.cluster_id == cluster { + if node_snapshot.node_info.cluster_id == cluster { // Remove node - self.unload_warehouse_info(&mut consistent_node.node_info); + self.unload_warehouse_info(&mut node_snapshot.node_info); drop_cluster_txn .if_then .push(TxnOp::delete(cluster_node_key)); drop_cluster_txn.if_then.push(TxnOp::put_with_ttl( node_key, - serde_json::to_vec(&consistent_node.node_info)?, + serde_json::to_vec(&node_snapshot.node_info)?, Some(self.lift_time * 4), )) } @@ -1525,9 +1525,9 @@ impl WarehouseApi for WarehouseMgr { for _idx in 0..10 { let mut rename_cluster_txn = TxnRequest::default(); - let mut consistent_info = self.consistent_warehouse_info(&warehouse).await?; + let mut warehouse_snapshot = self.warehouse_snapshot(&warehouse).await?; - consistent_info.warehouse_info = match consistent_info.warehouse_info { + warehouse_snapshot.warehouse_info = match warehouse_snapshot.warehouse_info { WarehouseInfo::SelfManaged(_) => Err(ErrorCode::InvalidWarehouse(format!("Cannot rename cluster for warehouse {:?}, because it's self-managed warehouse.", warehouse))), WarehouseInfo::SystemManaged(mut info) => match info.clusters.contains_key(&cur) { false => Err(ErrorCode::WarehouseClusterNotExists(format!("Warehouse cluster {:?}.{:?} not exists", warehouse, cur))), @@ -1551,40 +1551,40 @@ impl WarehouseApi for WarehouseMgr { rename_cluster_txn.condition.push(map_condition( &warehouse_info_key, - MatchSeq::Exact(consistent_info.info_seq), + MatchSeq::Exact(warehouse_snapshot.info_seq), )); rename_cluster_txn.if_then.push(TxnOp::put( warehouse_info_key.clone(), - serde_json::to_vec(&consistent_info.warehouse_info)?, + serde_json::to_vec(&warehouse_snapshot.warehouse_info)?, )); // lock all cluster state - for mut consistent_node in consistent_info.consistent_nodes { - let node_key = self.node_key(&consistent_node.node_info)?; - let old_cluster_node_key = self.cluster_node_key(&consistent_node.node_info)?; + for mut node_snapshot in warehouse_snapshot.snapshot_nodes { + let node_key = self.node_key(&node_snapshot.node_info)?; + let old_cluster_node_key = self.cluster_node_key(&node_snapshot.node_info)?; rename_cluster_txn.condition.push(map_condition( &node_key, - MatchSeq::Exact(consistent_node.node_seq), + MatchSeq::Exact(node_snapshot.node_seq), )); rename_cluster_txn.condition.push(map_condition( &old_cluster_node_key, - MatchSeq::Exact(consistent_node.cluster_seq), + MatchSeq::Exact(node_snapshot.cluster_seq), )); - if consistent_node.node_info.cluster_id == cur { + if node_snapshot.node_info.cluster_id == cur { // rename node - consistent_node.node_info.cluster_id = to.clone(); - let new_cluster_node_key = self.cluster_node_key(&consistent_node.node_info)?; + node_snapshot.node_info.cluster_id = to.clone(); + let new_cluster_node_key = self.cluster_node_key(&node_snapshot.node_info)?; rename_cluster_txn.if_then.push(TxnOp::put_with_ttl( node_key, - serde_json::to_vec(&consistent_node.node_info)?, + serde_json::to_vec(&node_snapshot.node_info)?, Some(self.lift_time * 4), )); - consistent_node.node_info.cluster_id = String::new(); + node_snapshot.node_info.cluster_id = String::new(); rename_cluster_txn .condition .push(map_condition(&new_cluster_node_key, MatchSeq::Exact(0))); @@ -1593,7 +1593,7 @@ impl WarehouseApi for WarehouseMgr { .push(TxnOp::delete(old_cluster_node_key)); rename_cluster_txn.if_then.push(TxnOp::put_with_ttl( new_cluster_node_key, - serde_json::to_vec(&consistent_node.node_info)?, + serde_json::to_vec(&node_snapshot.node_info)?, Some(self.lift_time * 4), )); } @@ -1632,9 +1632,9 @@ impl WarehouseApi for WarehouseMgr { let mut add_cluster_node_txn = TxnRequest::default(); - let mut consistent_info = self.consistent_warehouse_info(&warehouse).await?; + let mut warehouse_snapshot = self.warehouse_snapshot(&warehouse).await?; - consistent_info.warehouse_info = match consistent_info.warehouse_info { + warehouse_snapshot.warehouse_info = match warehouse_snapshot.warehouse_info { WarehouseInfo::SelfManaged(_) => Err(ErrorCode::InvalidWarehouse(format!( "Cannot assign nodes for warehouse {:?}, because it's self-managed warehouse.", warehouse @@ -1664,26 +1664,26 @@ impl WarehouseApi for WarehouseMgr { add_cluster_node_txn.condition.push(map_condition( &warehouse_info_key, - MatchSeq::Exact(consistent_info.info_seq), + MatchSeq::Exact(warehouse_snapshot.info_seq), )); add_cluster_node_txn.if_then.push(TxnOp::put( warehouse_info_key.clone(), - serde_json::to_vec(&consistent_info.warehouse_info)?, + serde_json::to_vec(&warehouse_snapshot.warehouse_info)?, )); // lock all cluster state - for consistent_node in consistent_info.consistent_nodes { - let node_key = self.node_key(&consistent_node.node_info)?; - let cluster_node_key = self.cluster_node_key(&consistent_node.node_info)?; + for node_snapshot in warehouse_snapshot.snapshot_nodes { + let node_key = self.node_key(&node_snapshot.node_info)?; + let cluster_node_key = self.cluster_node_key(&node_snapshot.node_info)?; add_cluster_node_txn.condition.push(map_condition( &node_key, - MatchSeq::Exact(consistent_node.node_seq), + MatchSeq::Exact(node_snapshot.node_seq), )); add_cluster_node_txn.condition.push(map_condition( &cluster_node_key, - MatchSeq::Exact(consistent_node.cluster_seq), + MatchSeq::Exact(node_snapshot.cluster_seq), )); } @@ -1739,9 +1739,9 @@ impl WarehouseApi for WarehouseMgr { let mut nodes = nodes.clone(); let mut drop_cluster_node_txn = TxnRequest::default(); - let mut consistent_info = self.consistent_warehouse_info(warehouse).await?; + let mut warehouse_snapshot = self.warehouse_snapshot(warehouse).await?; - consistent_info.warehouse_info = match consistent_info.warehouse_info { + warehouse_snapshot.warehouse_info = match warehouse_snapshot.warehouse_info { WarehouseInfo::SelfManaged(_) => Err(ErrorCode::InvalidWarehouse(format!( "Cannot unassign nodes for warehouse {:?}, because it's self-managed warehouse.", warehouse @@ -1786,41 +1786,41 @@ impl WarehouseApi for WarehouseMgr { drop_cluster_node_txn.condition.push(map_condition( &warehouse_info_key, - MatchSeq::Exact(consistent_info.info_seq), + MatchSeq::Exact(warehouse_snapshot.info_seq), )); drop_cluster_node_txn.if_then.push(TxnOp::put( warehouse_info_key.clone(), - serde_json::to_vec(&consistent_info.warehouse_info)?, + serde_json::to_vec(&warehouse_snapshot.warehouse_info)?, )); // lock all cluster state - for mut consistent_node in consistent_info.consistent_nodes { - let node_key = self.node_key(&consistent_node.node_info)?; - let cluster_node_key = self.cluster_node_key(&consistent_node.node_info)?; + for mut node_snapshot in warehouse_snapshot.snapshot_nodes { + let node_key = self.node_key(&node_snapshot.node_info)?; + let cluster_node_key = self.cluster_node_key(&node_snapshot.node_info)?; drop_cluster_node_txn.condition.push(map_condition( &node_key, - MatchSeq::Exact(consistent_node.node_seq), + MatchSeq::Exact(node_snapshot.node_seq), )); drop_cluster_node_txn.condition.push(map_condition( &cluster_node_key, - MatchSeq::Exact(consistent_node.cluster_seq), + MatchSeq::Exact(node_snapshot.cluster_seq), )); - if let Some(v) = nodes.get_mut(&consistent_node.node_info.cluster_id) { + if let Some(v) = nodes.get_mut(&node_snapshot.node_info.cluster_id) { if let Some(remove_node) = v.pop() { let SelectedNode::Random(node_group) = remove_node; - if consistent_node.node_info.runtime_node_group == node_group { - self.unload_warehouse_info(&mut consistent_node.node_info); - consistent_node.node_info.runtime_node_group = None; + if node_snapshot.node_info.runtime_node_group == node_group { + self.unload_warehouse_info(&mut node_snapshot.node_info); + node_snapshot.node_info.runtime_node_group = None; drop_cluster_node_txn .if_then .push(TxnOp::delete(cluster_node_key)); drop_cluster_node_txn.if_then.push(TxnOp::put_with_ttl( node_key, - serde_json::to_vec(&consistent_node.node_info)?, + serde_json::to_vec(&node_snapshot.node_info)?, Some(self.lift_time * 4), )) } From 3387b1d0060c647c75141ef16b2b2d7b9c252369 Mon Sep 17 00:00:00 2001 From: zhang2014 Date: Thu, 2 Jan 2025 20:04:42 +0800 Subject: [PATCH 41/53] feat(cluster): clean code --- .../management/src/warehouse/warehouse_mgr.rs | 63 ++++++++++--------- 1 file changed, 33 insertions(+), 30 deletions(-) diff --git a/src/query/management/src/warehouse/warehouse_mgr.rs b/src/query/management/src/warehouse/warehouse_mgr.rs index a1ca1891ea80d..a203bc116d7df 100644 --- a/src/query/management/src/warehouse/warehouse_mgr.rs +++ b/src/query/management/src/warehouse/warehouse_mgr.rs @@ -102,6 +102,32 @@ impl WarehouseMgr { ), }) } + + async fn shutdown_system_managed(&self, node_info: &NodeInfo) -> Result<()> { + let node_key = self.node_key(&node_info)?; + + let mut txn = TxnRequest::default(); + + txn.if_then.push(TxnOp::delete(node_key)); + + if node_info.assigned_warehouse() { + txn.if_then.push(TxnOp::delete(format!( + "{}/{}/{}/{}", + self.cluster_node_key_prefix, + escape_for_key(&node_info.warehouse_id)?, + escape_for_key(&node_info.cluster_id)?, + escape_for_key(&node_info.id)? + ))); + } + + match self.metastore.transaction(txn).await?.success { + true => Ok(()), + false => Err(ErrorCode::ClusterUnknownNode(format!( + "Node with ID '{}' does not exist in the cluster.", + node_info.id + ))), + } + } } fn map_condition(k: &str, seq: MatchSeq) -> TxnCondition { @@ -585,13 +611,13 @@ impl WarehouseMgr { async fn unassigned_nodes(&self) -> Result, Vec<(u64, NodeInfo)>>> { let online_nodes = self.metastore.prefix_list_kv(&self.node_key_prefix).await?; - let mut group_nodes = HashMap::with_capacity(online_nodes.len()); + let mut unassigned_nodes = HashMap::with_capacity(online_nodes.len()); for (_, seq_data) in online_nodes { let node_info = serde_json::from_slice::(&seq_data.data)?; - if node_info.cluster_id.is_empty() && node_info.warehouse_id.is_empty() { - match group_nodes.entry(node_info.node_group.clone()) { + if !node_info.assigned_warehouse() { + match unassigned_nodes.entry(node_info.node_group.clone()) { Entry::Vacant(v) => { v.insert(vec![(seq_data.seq, node_info)]); } @@ -602,7 +628,7 @@ impl WarehouseMgr { } } - Ok(group_nodes) + Ok(unassigned_nodes) } async fn pick_assign_warehouse_node( @@ -753,29 +779,7 @@ impl WarehouseApi for WarehouseMgr { return match node_info.node_type { NodeType::SelfManaged => self.shutdown_self_managed_node(&node_info).await, - NodeType::SystemManaged => { - let mut txn = TxnRequest::default(); - - txn.if_then.push(TxnOp::delete(node_key)); - - if !node_info.cluster_id.is_empty() && !node_info.warehouse_id.is_empty() { - txn.if_then.push(TxnOp::delete(format!( - "{}/{}/{}/{}", - self.cluster_node_key_prefix, - escape_for_key(&node_info.warehouse_id)?, - escape_for_key(&node_info.cluster_id)?, - escape_for_key(&node_info.id)? - ))); - } - - match self.metastore.transaction(txn).await?.success { - true => Ok(()), - false => Err(ErrorCode::ClusterUnknownNode(format!( - "Node with ID '{}' does not exist in the cluster.", - node_id - ))), - } - } + NodeType::SystemManaged => self.shutdown_system_managed(&node_info).await, }; } @@ -786,7 +790,6 @@ impl WarehouseApi for WarehouseMgr { } async fn heartbeat_node(&self, node: &mut NodeInfo, exact: u64) -> Result { - // let exact = MatchSeq::Exact(seq); let heartbeat_reply = match node.node_type { NodeType::SelfManaged => { assert!(!node.cluster_id.is_empty()); @@ -999,7 +1002,7 @@ impl WarehouseApi for WarehouseMgr { for (_key, v) in online_nodes { let node_info = serde_json::from_slice::(&v.data)?; - if node_info.warehouse_id.is_empty() && node_info.cluster_id.is_empty() { + if !node_info.assigned_warehouse() { assert_eq!(node_info.node_type, NodeType::SystemManaged); unassign_online_nodes.push((v.seq, node_info)); } @@ -1892,7 +1895,7 @@ impl WarehouseApi for WarehouseMgr { let node = serde_json::from_slice::(&seq.data)?; - if node.warehouse_id.is_empty() && node.cluster_id.is_empty() { + if !node.assigned_warehouse() { return Ok((false, vec![node])); } From 7a7f715b9ac690dae0b95373b00b0362027e7d09 Mon Sep 17 00:00:00 2001 From: zhang2014 Date: Fri, 3 Jan 2025 20:15:19 +0800 Subject: [PATCH 42/53] feat(cluster): clean code --- src/query/ast/src/ast/statements/statement.rs | 3 + src/query/ast/src/ast/statements/warehouse.rs | 11 ++++ src/query/ast/src/parser/statement.rs | 8 +++ .../management/src/warehouse/warehouse_mgr.rs | 2 +- .../interpreters/access/privilege_access.rs | 1 + .../src/interpreters/interpreter_factory.rs | 2 + .../interpreters/interpreter_use_warehouse.rs | 62 +++++++++++++++++++ src/query/service/src/interpreters/mod.rs | 1 + src/query/sql/src/planner/binder/binder.rs | 1 + .../sql/src/planner/binder/ddl/warehouse.rs | 11 ++++ .../sql/src/planner/format/display_plan.rs | 1 + .../sql/src/planner/plans/ddl/warehouse.rs | 5 ++ src/query/sql/src/planner/plans/plan.rs | 2 + 13 files changed, 109 insertions(+), 1 deletion(-) create mode 100644 src/query/service/src/interpreters/interpreter_use_warehouse.rs diff --git a/src/query/ast/src/ast/statements/statement.rs b/src/query/ast/src/ast/statements/statement.rs index 093c3da855418..651d9532f1095 100644 --- a/src/query/ast/src/ast/statements/statement.rs +++ b/src/query/ast/src/ast/statements/statement.rs @@ -131,6 +131,7 @@ pub enum Statement { // Warehouses ShowOnlineNodes(ShowOnlineNodesStmt), + UseWarehouse(UseWarehouseStmt), ShowWarehouses(ShowWarehousesStmt), DropWarehouse(DropWarehouseStmt), CreateWarehouse(CreateWarehouseStmt), @@ -557,6 +558,7 @@ impl Statement { | Statement::CreateProcedure(..) | Statement::DropProcedure(..) | Statement::CreateWarehouse(..) + | Statement::UseWarehouse(..) | Statement::DropWarehouse(..) | Statement::RenameWarehouse(..) | Statement::AddWarehouseCluster(..) @@ -979,6 +981,7 @@ impl Display for Statement { Statement::ShowOnlineNodes(stmt) => write!(f, "{stmt}")?, Statement::ShowWarehouses(stmt) => write!(f, "{stmt}")?, + Statement::UseWarehouse(stmt) => write!(f, "{stmt}")?, Statement::DropWarehouse(stmt) => write!(f, "{stmt}")?, Statement::CreateWarehouse(stmt) => write!(f, "{stmt}")?, Statement::RenameWarehouse(stmt) => write!(f, "{stmt}")?, diff --git a/src/query/ast/src/ast/statements/warehouse.rs b/src/query/ast/src/ast/statements/warehouse.rs index af812a9abd3ac..d0200b8ad60d6 100644 --- a/src/query/ast/src/ast/statements/warehouse.rs +++ b/src/query/ast/src/ast/statements/warehouse.rs @@ -39,6 +39,17 @@ impl Display for ShowWarehousesStmt { } } +#[derive(Debug, Clone, PartialEq, Drive, DriveMut)] +pub struct UseWarehouseStmt { + pub warehouse: Identifier, +} + +impl Display for UseWarehouseStmt { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + write!(f, "USE WAREHOUSE {}", self.warehouse) + } +} + #[derive(Debug, Clone, PartialEq, Drive, DriveMut)] pub struct CreateWarehouseStmt { pub warehouse: Identifier, diff --git a/src/query/ast/src/parser/statement.rs b/src/query/ast/src/parser/statement.rs index 39c1a8977de60..fb8145245afcd 100644 --- a/src/query/ast/src/parser/statement.rs +++ b/src/query/ast/src/parser/statement.rs @@ -537,6 +537,13 @@ pub fn statement_body(i: Input) -> IResult { |(_, _)| Statement::ShowWarehouses(ShowWarehousesStmt {}), ); + let use_warehouse = map( + rule! { + USE ~ WAREHOUSE ~ #ident + }, + |(_, _, warehouse)| Statement::UseWarehouse(UseWarehouseStmt { warehouse }), + ); + let create_warehouse = map( rule! { CREATE ~ WAREHOUSE ~ #ident ~ ("(" ~ #assign_nodes_list ~ ")")? ~ (WITH ~ #warehouse_cluster_option)? @@ -2423,6 +2430,7 @@ pub fn statement_body(i: Input) -> IResult { rule!( #show_warehouses: "`SHOW WAREHOUSES`" | #show_online_nodes: "`SHOW ONLINE NODES`" + | #use_warehouse: "`USE WAREHOUSE `" | #create_warehouse: "`CREATE WAREHOUSE [(ASSIGN NODES [FROM ] [, ...])] WITH [warehouse_size = ]`" | #drop_warehouse: "`DROP WAREHOUSE `" | #rename_warehouse: "`RENAME WAREHOUSE TO `" diff --git a/src/query/management/src/warehouse/warehouse_mgr.rs b/src/query/management/src/warehouse/warehouse_mgr.rs index a203bc116d7df..09c503f13dbdd 100644 --- a/src/query/management/src/warehouse/warehouse_mgr.rs +++ b/src/query/management/src/warehouse/warehouse_mgr.rs @@ -104,7 +104,7 @@ impl WarehouseMgr { } async fn shutdown_system_managed(&self, node_info: &NodeInfo) -> Result<()> { - let node_key = self.node_key(&node_info)?; + let node_key = self.node_key(node_info)?; let mut txn = TxnRequest::default(); diff --git a/src/query/service/src/interpreters/access/privilege_access.rs b/src/query/service/src/interpreters/access/privilege_access.rs index 9fd2967dd2825..3f0a0edc15dcf 100644 --- a/src/query/service/src/interpreters/access/privilege_access.rs +++ b/src/query/service/src/interpreters/access/privilege_access.rs @@ -1264,6 +1264,7 @@ impl AccessChecker for PrivilegeAccess { Plan::InspectWarehouse(_) => {} Plan::DropWarehouseCluster(_) => {} Plan::RenameWarehouseCluster(_) => {} + Plan::UseWarehouse(_) => {} Plan::CreateWarehouse(_) => {} Plan::AddWarehouseCluster(_) => {} Plan::AssignWarehouseNodes(_) => {} diff --git a/src/query/service/src/interpreters/interpreter_factory.rs b/src/query/service/src/interpreters/interpreter_factory.rs index d767a51771e0d..19f36d0bc65f8 100644 --- a/src/query/service/src/interpreters/interpreter_factory.rs +++ b/src/query/service/src/interpreters/interpreter_factory.rs @@ -83,6 +83,7 @@ use crate::interpreters::interpreter_txn_abort::AbortInterpreter; use crate::interpreters::interpreter_txn_begin::BeginInterpreter; use crate::interpreters::interpreter_txn_commit::CommitInterpreter; use crate::interpreters::interpreter_unassign_warehouse_nodes::UnassignWarehouseNodesInterpreter; +use crate::interpreters::interpreter_use_warehouse::UseWarehouseInterpreter; use crate::interpreters::interpreter_view_describe::DescribeViewInterpreter; use crate::interpreters::AlterUserInterpreter; use crate::interpreters::CreateStreamInterpreter; @@ -120,6 +121,7 @@ impl InterpreterFactory { match plan { Plan::ShowWarehouses => Ok(Arc::new(ShowWarehousesInterpreter::try_create(ctx.clone())?)), Plan::ShowOnlineNodes => Ok(Arc::new(ShowOnlineNodesInterpreter::try_create(ctx.clone())?)), + Plan::UseWarehouse(v) => Ok(Arc::new(UseWarehouseInterpreter::try_create(ctx.clone(), *v.clone())?)), Plan::CreateWarehouse(v) => Ok(Arc::new(CreateWarehouseInterpreter::try_create( ctx.clone(), *v.clone(), diff --git a/src/query/service/src/interpreters/interpreter_use_warehouse.rs b/src/query/service/src/interpreters/interpreter_use_warehouse.rs new file mode 100644 index 0000000000000..c6fb637d8140a --- /dev/null +++ b/src/query/service/src/interpreters/interpreter_use_warehouse.rs @@ -0,0 +1,62 @@ +// Copyright 2021 Datafuse Labs +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::sync::Arc; + +use databend_common_base::base::GlobalInstance; +use databend_common_catalog::table_context::TableContext; +use databend_common_exception::Result; +use databend_common_license::license::Feature; +use databend_common_license::license_manager::LicenseManagerSwitch; +use databend_common_sql::plans::UseWarehousePlan; +use databend_enterprise_resources_management::ResourcesManagement; + +use crate::interpreters::Interpreter; +use crate::pipelines::PipelineBuildResult; +use crate::sessions::QueryContext; + +pub struct UseWarehouseInterpreter { + ctx: Arc, + plan: UseWarehousePlan, +} + +impl UseWarehouseInterpreter { + pub fn try_create(ctx: Arc, plan: UseWarehousePlan) -> Result { + Ok(UseWarehouseInterpreter { ctx, plan }) + } +} + +#[async_trait::async_trait] +impl Interpreter for UseWarehouseInterpreter { + fn name(&self) -> &str { + "UseWarehouseInterpreter" + } + + fn is_ddl(&self) -> bool { + false + } + + #[async_backtrace::framed] + async fn execute2(&self) -> Result { + LicenseManagerSwitch::instance() + .check_enterprise_enabled(self.ctx.get_license_key(), Feature::SystemManagement)?; + + // check warehouse exists + let _nodes = GlobalInstance::get::>() + .inspect_warehouse(self.plan.warehouse.clone()) + .await?; + + Ok(PipelineBuildResult::create()) + } +} diff --git a/src/query/service/src/interpreters/mod.rs b/src/query/service/src/interpreters/mod.rs index d978c2cee591f..d9028b0c00952 100644 --- a/src/query/service/src/interpreters/mod.rs +++ b/src/query/service/src/interpreters/mod.rs @@ -139,6 +139,7 @@ mod interpreter_txn_commit; mod interpreter_unassign_warehouse_nodes; mod interpreter_unset; mod interpreter_use_database; +mod interpreter_use_warehouse; mod interpreter_user_alter; mod interpreter_user_create; mod interpreter_user_desc; diff --git a/src/query/sql/src/planner/binder/binder.rs b/src/query/sql/src/planner/binder/binder.rs index 976d8b1d2edf3..ede8956c0f3fd 100644 --- a/src/query/sql/src/planner/binder/binder.rs +++ b/src/query/sql/src/planner/binder/binder.rs @@ -658,6 +658,7 @@ impl<'a> Binder { } Statement::ShowOnlineNodes(v) => self.bind_show_online_nodes(v)?, Statement::ShowWarehouses(v) => self.bind_show_warehouses(v)?, + Statement::UseWarehouse(v) => self.bind_use_warehouse(v)?, Statement::DropWarehouse(v) => self.bind_drop_warehouse(v)?, Statement::CreateWarehouse(v) => self.bind_create_warehouse(v)?, Statement::RenameWarehouse(v) => self.bind_rename_warehouse(v)?, diff --git a/src/query/sql/src/planner/binder/ddl/warehouse.rs b/src/query/sql/src/planner/binder/ddl/warehouse.rs index 9a36ca6f18359..fa5bd5e242421 100644 --- a/src/query/sql/src/planner/binder/ddl/warehouse.rs +++ b/src/query/sql/src/planner/binder/ddl/warehouse.rs @@ -28,6 +28,7 @@ use databend_common_ast::ast::ShowOnlineNodesStmt; use databend_common_ast::ast::ShowWarehousesStmt; use databend_common_ast::ast::SuspendWarehouseStmt; use databend_common_ast::ast::UnassignWarehouseNodesStmt; +use databend_common_ast::ast::UseWarehouseStmt; use databend_common_exception::Result; use crate::plans::AddWarehouseClusterPlan; @@ -42,6 +43,7 @@ use crate::plans::RenameWarehousePlan; use crate::plans::ResumeWarehousePlan; use crate::plans::SuspendWarehousePlan; use crate::plans::UnassignWarehouseNodesPlan; +use crate::plans::UseWarehousePlan; use crate::Binder; impl Binder { @@ -59,6 +61,15 @@ impl Binder { Ok(Plan::ShowWarehouses) } + pub(in crate::planner::binder) fn bind_use_warehouse( + &mut self, + stmt: &UseWarehouseStmt, + ) -> Result { + Ok(Plan::UseWarehouse(Box::new(UseWarehousePlan { + warehouse: stmt.warehouse.to_string(), + }))) + } + pub(in crate::planner::binder) fn bind_create_warehouse( &mut self, stmt: &CreateWarehouseStmt, diff --git a/src/query/sql/src/planner/format/display_plan.rs b/src/query/sql/src/planner/format/display_plan.rs index a88171e5548a6..b80cfe416fe74 100644 --- a/src/query/sql/src/planner/format/display_plan.rs +++ b/src/query/sql/src/planner/format/display_plan.rs @@ -223,6 +223,7 @@ impl Plan { Plan::DropWarehouseCluster(_) => Ok("DropWarehouseCluster".to_string()), Plan::RenameWarehouseCluster(_) => Ok("RenameWarehouseCluster".to_string()), Plan::CreateWarehouse(_) => Ok("CreateWarehouse".to_string()), + Plan::UseWarehouse(_) => Ok("UseWarehouse".to_string()), Plan::AddWarehouseCluster(_) => Ok("AddWarehouseCluster".to_string()), Plan::AssignWarehouseNodes(_) => Ok("AddWarehouseClusterNode".to_string()), Plan::UnassignWarehouseNodes(_) => Ok("DropWarehouseClusterNode".to_string()), diff --git a/src/query/sql/src/planner/plans/ddl/warehouse.rs b/src/query/sql/src/planner/plans/ddl/warehouse.rs index c66fe220c652d..07c73af657cc4 100644 --- a/src/query/sql/src/planner/plans/ddl/warehouse.rs +++ b/src/query/sql/src/planner/plans/ddl/warehouse.rs @@ -20,6 +20,11 @@ use databend_common_expression::DataField; use databend_common_expression::DataSchemaRef; use databend_common_expression::DataSchemaRefExt; +#[derive(Clone, Debug, PartialEq, Eq, serde::Serialize)] +pub struct UseWarehousePlan { + pub warehouse: String, +} + #[derive(Clone, Debug, PartialEq, Eq, serde::Serialize)] pub struct CreateWarehousePlan { pub warehouse: String, diff --git a/src/query/sql/src/planner/plans/plan.rs b/src/query/sql/src/planner/plans/plan.rs index a189fe78224a0..e3abfd95793f2 100644 --- a/src/query/sql/src/planner/plans/plan.rs +++ b/src/query/sql/src/planner/plans/plan.rs @@ -157,6 +157,7 @@ use crate::plans::UnsetOptionsPlan; use crate::plans::UnsetPlan; use crate::plans::UseCatalogPlan; use crate::plans::UseDatabasePlan; +use crate::plans::UseWarehousePlan; use crate::plans::VacuumDropTablePlan; use crate::plans::VacuumTablePlan; use crate::plans::VacuumTemporaryFilesPlan; @@ -205,6 +206,7 @@ pub enum Plan { // Warehouses ShowOnlineNodes, ShowWarehouses, + UseWarehouse(Box), CreateWarehouse(Box), DropWarehouse(Box), ResumeWarehouse(Box), From 84923b2fd7c3cd4d4b8e65baf447bd28afd798ae Mon Sep 17 00:00:00 2001 From: zhang2014 Date: Fri, 3 Jan 2025 22:22:31 +0800 Subject: [PATCH 43/53] feat(cluster): clean code --- src/query/ast/src/parser/statement.rs | 2 +- .../service/src/interpreters/interpreter_use_warehouse.rs | 6 ++++++ src/query/settings/src/settings_default.rs | 7 +++++++ src/query/settings/src/settings_getter_setter.rs | 5 +++++ 4 files changed, 19 insertions(+), 1 deletion(-) diff --git a/src/query/ast/src/parser/statement.rs b/src/query/ast/src/parser/statement.rs index fb8145245afcd..6c2629bdafc2b 100644 --- a/src/query/ast/src/parser/statement.rs +++ b/src/query/ast/src/parser/statement.rs @@ -2424,13 +2424,13 @@ pub fn statement_body(i: Input) -> IResult { // use rule!( #use_catalog: "`USE CATALOG `" + | #use_warehouse: "`USE WAREHOUSE `" | #use_database : "`USE `" ), // warehouse rule!( #show_warehouses: "`SHOW WAREHOUSES`" | #show_online_nodes: "`SHOW ONLINE NODES`" - | #use_warehouse: "`USE WAREHOUSE `" | #create_warehouse: "`CREATE WAREHOUSE [(ASSIGN NODES [FROM ] [, ...])] WITH [warehouse_size = ]`" | #drop_warehouse: "`DROP WAREHOUSE `" | #rename_warehouse: "`RENAME WAREHOUSE TO `" diff --git a/src/query/service/src/interpreters/interpreter_use_warehouse.rs b/src/query/service/src/interpreters/interpreter_use_warehouse.rs index c6fb637d8140a..df89a0d59b960 100644 --- a/src/query/service/src/interpreters/interpreter_use_warehouse.rs +++ b/src/query/service/src/interpreters/interpreter_use_warehouse.rs @@ -57,6 +57,12 @@ impl Interpreter for UseWarehouseInterpreter { .inspect_warehouse(self.plan.warehouse.clone()) .await?; + unsafe { + self.ctx + .get_session_settings() + .set_warehouse(self.plan.warehouse.clone())?; + } + Ok(PipelineBuildResult::create()) } } diff --git a/src/query/settings/src/settings_default.rs b/src/query/settings/src/settings_default.rs index b895d7ffd65d7..6fc84d47fc340 100644 --- a/src/query/settings/src/settings_default.rs +++ b/src/query/settings/src/settings_default.rs @@ -1151,6 +1151,13 @@ impl DefaultSettings { scope: SettingScope::Both, range: Some(SettingRange::Numeric(0..=u64::MAX)), }), + ("warehouse", DefaultSettingValue { + value: UserSettingValue::String("".to_string()), + desc: "Please use the statement to set the warehouse, this settings is only used to synchronize the warehouse status between the client and the server.", + mode: SettingMode::Read, + scope: SettingScope::Session, + range: None, + }), ]); Ok(Arc::new(DefaultSettings { diff --git a/src/query/settings/src/settings_getter_setter.rs b/src/query/settings/src/settings_getter_setter.rs index 21d12dbb56fcb..803af09caf731 100644 --- a/src/query/settings/src/settings_getter_setter.rs +++ b/src/query/settings/src/settings_getter_setter.rs @@ -848,4 +848,9 @@ impl Settings { let v = self.try_get_u64("stream_consume_batch_size_hint")?; Ok(if v == 0 { None } else { Some(v) }) } + + /// # Safety + pub unsafe fn set_warehouse(&self, warehouse: String) -> Result<()> { + self.unchecked_set_setting(String::from("warehouse"), warehouse) + } } From 8bd09f8d97fffcf96bc2c05b50ea4363fa22fd9e Mon Sep 17 00:00:00 2001 From: zhang2014 Date: Sun, 5 Jan 2025 01:29:18 +0800 Subject: [PATCH 44/53] feat(cluster): clean code --- src/common/base/src/headers.rs | 1 + src/query/service/src/clusters/cluster.rs | 43 ++++++++ .../src/servers/http/middleware/session.rs | 103 +++++++++++++----- src/query/settings/src/settings_default.rs | 2 +- 4 files changed, 118 insertions(+), 31 deletions(-) diff --git a/src/common/base/src/headers.rs b/src/common/base/src/headers.rs index aaaafbbecbbaf..06dc158f9bb72 100644 --- a/src/common/base/src/headers.rs +++ b/src/common/base/src/headers.rs @@ -25,6 +25,7 @@ pub const HEADER_QUERY_STATE: &str = "X-DATABEND-QUERY-STATE"; pub const HEADER_QUERY_PAGE_ROWS: &str = "X-DATABEND-QUERY-PAGE-ROWS"; pub const HEADER_VERSION: &str = "X-DATABEND-VERSION"; pub const HEADER_STICKY: &str = "X-DATABEND-STICKY-NODE"; +pub const HEADER_WAREHOUSE: &str = "X-DATABEND-WAREHOUSE"; pub const HEADER_SIGNATURE: &str = "X-DATABEND-SIGNATURE"; pub const HEADER_AUTH_METHOD: &str = "X-DATABEND-AUTH-METHOD"; diff --git a/src/query/service/src/clusters/cluster.rs b/src/query/service/src/clusters/cluster.rs index e6be6ca216797..994f8b3d0ce38 100644 --- a/src/query/service/src/clusters/cluster.rs +++ b/src/query/service/src/clusters/cluster.rs @@ -12,6 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. +use std::collections::hash_map::Entry; use std::collections::HashMap; use std::net::SocketAddr; use std::ops::RangeInclusive; @@ -325,6 +326,48 @@ impl ClusterDiscovery { } } + pub async fn find_node_by_warehouse( + self: Arc, + warehouse: &str, + ) -> Result>> { + let nodes = self + .warehouse_manager + .list_warehouse_nodes(warehouse.to_string()) + .await?; + + let mut warehouse_clusters_nodes = Vec::new(); + let mut warehouse_clusters_nodes_index = HashMap::new(); + + for node in nodes { + match warehouse_clusters_nodes_index + .entry((node.version.to_string(), node.cluster_id.clone())) + { + Entry::Vacant(v) => { + v.insert(warehouse_clusters_nodes.len()); + warehouse_clusters_nodes.push(vec![Arc::new(node)]); + } + Entry::Occupied(v) => { + warehouse_clusters_nodes[*v.get()].push(Arc::new(node)); + } + }; + } + + if warehouse_clusters_nodes.is_empty() { + return Ok(None); + } + + let system_time = std::time::SystemTime::now(); + let system_timestamp = system_time + .duration_since(std::time::UNIX_EPOCH) + .expect("expect time"); + + let nanos = system_timestamp.as_nanos(); + let cluster_idx = (nanos % warehouse_clusters_nodes_index.len() as u128) as usize; + let pick_cluster_nodes = &warehouse_clusters_nodes[cluster_idx]; + let nodes_idx = (nanos % pick_cluster_nodes.len() as u128) as usize; + Ok(Some(pick_cluster_nodes[nodes_idx].clone())) + } + pub async fn find_node_by_id(self: Arc, id: &str) -> Result>> { { let mut lru_cache = self.lru_cache.lock(); diff --git a/src/query/service/src/servers/http/middleware/session.rs b/src/query/service/src/servers/http/middleware/session.rs index 9bbadc60e49ce..0c4b991d5b3d2 100644 --- a/src/query/service/src/servers/http/middleware/session.rs +++ b/src/query/service/src/servers/http/middleware/session.rs @@ -21,6 +21,7 @@ use databend_common_base::headers::HEADER_QUERY_ID; use databend_common_base::headers::HEADER_STICKY; use databend_common_base::headers::HEADER_TENANT; use databend_common_base::headers::HEADER_VERSION; +use databend_common_base::headers::HEADER_WAREHOUSE; use databend_common_base::runtime::ThreadTracker; use databend_common_config::GlobalConfig; use databend_common_config::DATABEND_SEMVER; @@ -504,36 +505,78 @@ impl Endpoint for HTTPSessionEndpoint { async fn call(&self, mut req: Request) -> PoemResult { let headers = req.headers().clone(); - if self.endpoint_kind.may_need_sticky() - && let Some(sticky_node_id) = headers.get(HEADER_STICKY) - { - let sticky_node_id = sticky_node_id - .to_str() - .map_err(|e| { - HttpErrorCode::bad_request(ErrorCode::BadArguments(format!( - "Invalid Header ({HEADER_STICKY}: {sticky_node_id:?}): {e}" - ))) - })? - .to_string(); - let local_id = GlobalConfig::instance().query.node_id.clone(); - if local_id != sticky_node_id { - return if let Some(node) = ClusterDiscovery::instance() - .find_node_by_id(&sticky_node_id) - .await - .map_err(HttpErrorCode::server_error)? - { - log::info!( - "forwarding /v1{} from {local_id} to {sticky_node_id}", - req.uri() - ); - forward_request(req, node).await - } else { - let msg = format!("sticky_node_id '{sticky_node_id}' not found in cluster",); - warn!("{}", msg); - Err(Error::from(HttpErrorCode::bad_request( - ErrorCode::BadArguments(msg), - ))) - }; + if self.endpoint_kind.may_need_sticky() { + if let Some(sticky_node_id) = headers.get(HEADER_STICKY) { + let sticky_node_id = sticky_node_id + .to_str() + .map_err(|e| { + HttpErrorCode::bad_request(ErrorCode::BadArguments(format!( + "Invalid Header ({HEADER_STICKY}: {sticky_node_id:?}): {e}" + ))) + })? + .to_string(); + let local_id = GlobalConfig::instance().query.node_id.clone(); + if local_id != sticky_node_id { + return if let Some(node) = ClusterDiscovery::instance() + .find_node_by_id(&sticky_node_id) + .await + .map_err(HttpErrorCode::server_error)? + { + log::info!( + "forwarding /v1{} from {local_id} to {sticky_node_id}", + req.uri() + ); + forward_request(req, node).await + } else { + let msg = + format!("sticky_node_id '{sticky_node_id}' not found in cluster",); + warn!("{}", msg); + Err(Error::from(HttpErrorCode::bad_request( + ErrorCode::BadArguments(msg), + ))) + }; + } + } else if let Some(warehouse) = headers.get(HEADER_WAREHOUSE) { + req.headers_mut().remove(HEADER_WAREHOUSE); + + let warehouse = warehouse + .to_str() + .map_err(|e| { + HttpErrorCode::bad_request(ErrorCode::BadArguments(format!( + "Invalid Header ({HEADER_WAREHOUSE}: {warehouse:?}): {e}" + ))) + })? + .to_string(); + + let cluster_discovery = ClusterDiscovery::instance(); + + let forward_node = cluster_discovery.find_node_by_warehouse(&warehouse).await; + + match forward_node { + Err(error) => { + return Err(HttpErrorCode::server_error(error).into()); + } + Ok(None) => { + let msg = format!("Not find the '{}' warehouse; it is possible that all nodes of the warehouse have gone offline. Please exit the client and reconnect, or use `use warehouse `", warehouse); + warn!("{}", msg); + return Err(Error::from(HttpErrorCode::bad_request( + ErrorCode::UnknownWarehouse(msg), + ))); + } + Ok(Some(node)) => { + let local_id = GlobalConfig::instance().query.node_id.clone(); + if node.id != local_id { + log::info!( + "forwarding /v1{} from {} to warehouse {}({})", + req.uri(), + local_id, + warehouse, + node.id + ); + return forward_request(req, node).await; + } + } + } } } diff --git a/src/query/settings/src/settings_default.rs b/src/query/settings/src/settings_default.rs index b2ad94e9acd6c..602990689a963 100644 --- a/src/query/settings/src/settings_default.rs +++ b/src/query/settings/src/settings_default.rs @@ -1154,7 +1154,7 @@ impl DefaultSettings { ("warehouse", DefaultSettingValue { value: UserSettingValue::String("".to_string()), desc: "Please use the statement to set the warehouse, this settings is only used to synchronize the warehouse status between the client and the server.", - mode: SettingMode::Read, + mode: SettingMode::Both, scope: SettingScope::Session, range: None, }), From 8c45fe91c80c3678463ef8f85811e6023f020672 Mon Sep 17 00:00:00 2001 From: zhang2014 Date: Sun, 5 Jan 2025 20:41:33 +0800 Subject: [PATCH 45/53] feat(cluster): clean code --- .../management/src/warehouse/warehouse_api.rs | 2 +- .../management/src/warehouse/warehouse_mgr.rs | 77 +++++++++++-------- src/query/service/src/clusters/cluster.rs | 19 ++--- src/query/service/src/test_kits/fixture.rs | 2 +- 4 files changed, 55 insertions(+), 45 deletions(-) diff --git a/src/query/management/src/warehouse/warehouse_api.rs b/src/query/management/src/warehouse/warehouse_api.rs index 93dc31ac14fc6..090e003063f97 100644 --- a/src/query/management/src/warehouse/warehouse_api.rs +++ b/src/query/management/src/warehouse/warehouse_api.rs @@ -125,7 +125,7 @@ pub trait WarehouseApi: Sync + Send { async fn list_online_nodes(&self) -> Result>; - async fn discover(&self, node_id: &str) -> Result<(bool, Vec)>; + async fn discover(&self, node_id: &str) -> Result>; async fn get_node_info(&self, node_id: &str) -> Result>; } diff --git a/src/query/management/src/warehouse/warehouse_mgr.rs b/src/query/management/src/warehouse/warehouse_mgr.rs index 09c503f13dbdd..ff3e2e97bd2da 100644 --- a/src/query/management/src/warehouse/warehouse_mgr.rs +++ b/src/query/management/src/warehouse/warehouse_mgr.rs @@ -21,6 +21,7 @@ use databend_common_base::base::tokio; use databend_common_base::base::unescape_for_key; use databend_common_base::base::GlobalUniqName; use databend_common_base::vec_ext::VecExt; +use databend_common_base::version::DATABEND_COMMIT_VERSION; use databend_common_exception::ErrorCode; use databend_common_exception::Result; use databend_common_meta_kvapi::kvapi::KVApi; @@ -49,6 +50,8 @@ use crate::warehouse::warehouse_api::SystemManagedWarehouse; use crate::warehouse::warehouse_api::WarehouseInfo; use crate::warehouse::WarehouseApi; +static DEFAULT_CLUSTER_ID: &str = "default"; + pub static WAREHOUSE_API_KEY_PREFIX: &str = "__fd_clusters_v6"; pub static WAREHOUSE_META_KEY_PREFIX: &str = "__fd_warehouses"; @@ -330,24 +333,22 @@ impl WarehouseMgr { // txn_reply includes: 1. get node response 2. get warehouse response // Check warehouse state - { - // The last response is not meant to be returned. - let last = txn_reply.responses.pop(); + // The last response is not meant to be returned. + let last = txn_reply.responses.pop(); - let wh_get_resp = last.as_ref().and_then(|r| r.try_as_get()); - let wh_get_resp = wh_get_resp.ok_or_else(|| invalid_get_resp(last.as_ref()))?; + let wh_get_resp = last.as_ref().and_then(|r| r.try_as_get()); + let wh_get_resp = wh_get_resp.ok_or_else(|| invalid_get_resp(last.as_ref()))?; - if let Some(wh_seqv) = &wh_get_resp.value { - let wh: WarehouseInfo = serde_json::from_slice(&wh_seqv.data)?; + if let Some(wh_seqv) = &wh_get_resp.value { + let wh: WarehouseInfo = serde_json::from_slice(&wh_seqv.data)?; - if let WarehouseInfo::SystemManaged(_) = wh { - // an unknown error caused all nodes of the warehouse to fail the heartbeat with the meta service. If someone creates a warehouse with the same name during this period. - return Err(ErrorCode::WarehouseAlreadyExists( - "Already exists same name system-managed warehouse.", - )); - } - }; - } + if let WarehouseInfo::SystemManaged(_) = wh { + // an unknown error caused all nodes of the warehouse to fail the heartbeat with the meta service. If someone creates a warehouse with the same name during this period. + return Err(ErrorCode::WarehouseAlreadyExists( + "Already exists same name system-managed warehouse.", + )); + } + }; } Ok(txn_reply) @@ -634,15 +635,15 @@ impl WarehouseMgr { async fn pick_assign_warehouse_node( &self, warehouse: &str, - warehouse_nodes_matcher: &HashMap>, + nodes_matcher: &HashMap>, ) -> Result>> { - let mut selected_nodes = HashMap::with_capacity(warehouse_nodes_matcher.len()); + let mut selected_nodes = HashMap::with_capacity(nodes_matcher.len()); let mut grouped_nodes = self.unassigned_nodes().await?; let mut after_assign_node = HashMap::new(); - for (cluster, cluster_node_matcher) in warehouse_nodes_matcher { + for (cluster, cluster_node_matcher) in nodes_matcher { let mut cluster_selected_nodes = Vec::with_capacity(cluster_node_matcher.len()); for selected_node in cluster_node_matcher { match selected_node { @@ -843,7 +844,7 @@ impl WarehouseApi for WarehouseMgr { )); delete_txn.if_then.push(TxnOp::delete(warehouse_info_key)); - for mut node_snapshot in warehouse_snapshot.snapshot_nodes { + for node_snapshot in warehouse_snapshot.snapshot_nodes { let node_key = self.node_key(&node_snapshot.node_info)?; let cluster_node_key = self.cluster_node_key(&node_snapshot.node_info)?; @@ -857,10 +858,10 @@ impl WarehouseApi for WarehouseMgr { )); delete_txn.if_then.push(TxnOp::delete(cluster_node_key)); - self.unload_warehouse_info(&mut node_snapshot.node_info); + let leave_node = node_snapshot.node_info.leave_warehouse(); delete_txn.if_then.push(TxnOp::put_with_ttl( node_key, - serde_json::to_vec(&node_snapshot.node_info)?, + serde_json::to_vec(&leave_node)?, Some(self.lift_time * 4), )); } @@ -889,13 +890,13 @@ impl WarehouseApi for WarehouseMgr { )); } - let nodes_map = HashMap::from([(String::from("default"), nodes.clone())]); + let nodes_map = HashMap::from([(String::from(DEFAULT_CLUSTER_ID), nodes.clone())]); loop { let mut selected_nodes = self .pick_assign_warehouse_node(&warehouse, &nodes_map) .await?; - let selected_nodes = selected_nodes.remove("default").unwrap(); + let selected_nodes = selected_nodes.remove(DEFAULT_CLUSTER_ID).unwrap(); let mut txn = TxnRequest::default(); @@ -932,9 +933,12 @@ impl WarehouseApi for WarehouseMgr { id: GlobalUniqName::unique(), status: "Running".to_string(), display_name: warehouse.clone(), - clusters: HashMap::from([(String::from("default"), SystemManagedCluster { - nodes: nodes.clone(), - })]), + clusters: HashMap::from([( + String::from(DEFAULT_CLUSTER_ID), + SystemManagedCluster { + nodes: nodes.clone(), + }, + )]), }))?, )); txn.else_then.push(TxnOp::get(warehouse_info_key)); @@ -1462,7 +1466,7 @@ impl WarehouseApi for WarehouseMgr { )); // lock all cluster state - for mut node_snapshot in warehouse_snapshot.snapshot_nodes { + for node_snapshot in warehouse_snapshot.snapshot_nodes { let node_key = self.node_key(&node_snapshot.node_info)?; let cluster_node_key = self.cluster_node_key(&node_snapshot.node_info)?; @@ -1476,15 +1480,14 @@ impl WarehouseApi for WarehouseMgr { )); if node_snapshot.node_info.cluster_id == cluster { - // Remove node - self.unload_warehouse_info(&mut node_snapshot.node_info); + let leave_node = node_snapshot.node_info.leave_warehouse(); drop_cluster_txn .if_then .push(TxnOp::delete(cluster_node_key)); drop_cluster_txn.if_then.push(TxnOp::put_with_ttl( node_key, - serde_json::to_vec(&node_snapshot.node_info)?, + serde_json::to_vec(&leave_node)?, Some(self.lift_time * 4), )) } @@ -1883,7 +1886,7 @@ impl WarehouseApi for WarehouseMgr { Ok(online_nodes) } - async fn discover(&self, node_id: &str) -> Result<(bool, Vec)> { + async fn discover(&self, node_id: &str) -> Result> { let node_key = format!("{}/{}", self.node_key_prefix, escape_for_key(node_id)?); let Some(seq) = self.metastore.get_kv(&node_key).await? else { @@ -1896,11 +1899,17 @@ impl WarehouseApi for WarehouseMgr { let node = serde_json::from_slice::(&seq.data)?; if !node.assigned_warehouse() { - return Ok((false, vec![node])); + return Ok(vec![node]); } - let list_nodes = self.list_warehouse_cluster_nodes(&node.warehouse_id, &node.cluster_id); - Ok((true, list_nodes.await?)) + let expect_version = DATABEND_COMMIT_VERSION.to_string(); + + Ok(self + .list_warehouse_cluster_nodes(&node.warehouse_id, &node.cluster_id) + .await? + .into_iter() + .filter(|x| x.binary_version == expect_version) + .collect::>()) } async fn get_node_info(&self, node_id: &str) -> Result> { diff --git a/src/query/service/src/clusters/cluster.rs b/src/query/service/src/clusters/cluster.rs index 994f8b3d0ce38..8509e4f2438a2 100644 --- a/src/query/service/src/clusters/cluster.rs +++ b/src/query/service/src/clusters/cluster.rs @@ -78,7 +78,7 @@ pub struct ClusterDiscovery { // avoid leak FlightClient to common-xxx #[async_trait::async_trait] pub trait ClusterHelper { - fn create(unassign: bool, nodes: Vec>, local_id: String) -> Arc; + fn create(nodes: Vec>, local_id: String) -> Arc; fn empty() -> Arc; fn is_empty(&self) -> bool; fn is_local(&self, node: &NodeInfo) -> bool; @@ -98,7 +98,8 @@ pub trait ClusterHelper { #[async_trait::async_trait] impl ClusterHelper for Cluster { - fn create(unassign: bool, nodes: Vec>, local_id: String) -> Arc { + fn create(nodes: Vec>, local_id: String) -> Arc { + let unassign = nodes.iter().all(|node| !node.assigned_warehouse()); Arc::new(Cluster { unassign, local_id, @@ -274,11 +275,11 @@ impl ClusterDiscovery { pub async fn discover(&self, config: &InnerConfig) -> Result> { let nodes = match config.query.cluster_id.is_empty() { true => self.warehouse_manager.discover(&config.query.node_id).await, - false => self - .warehouse_manager - .list_warehouse_cluster_nodes(&self.cluster_id, &self.cluster_id) - .await - .map(|x| (true, x)), + false => { + self.warehouse_manager + .list_warehouse_cluster_nodes(&self.cluster_id, &self.cluster_id) + .await + } }; match nodes { @@ -292,7 +293,7 @@ impl ClusterDiscovery { ); Err(cause.add_message_back("(while cluster api get_nodes).")) } - Ok((has_cluster, cluster_nodes)) => { + Ok(cluster_nodes) => { let mut res = Vec::with_capacity(cluster_nodes.len()); for node in &cluster_nodes { if node.id != self.local_id { @@ -320,7 +321,7 @@ impl ClusterDiscovery { cluster_nodes.len() as f64, ); - let res = Cluster::create(!has_cluster, res, self.local_id.clone()); + let res = Cluster::create(res, self.local_id.clone()); Ok(res) } } diff --git a/src/query/service/src/test_kits/fixture.rs b/src/query/service/src/test_kits/fixture.rs index 858858673e09f..8cbfe09776d96 100644 --- a/src/query/service/src/test_kits/fixture.rs +++ b/src/query/service/src/test_kits/fixture.rs @@ -274,7 +274,7 @@ impl TestFixture { let dummy_query_context = QueryContext::create_from_shared(QueryContextShared::try_create( self.default_session.clone(), - Cluster::create(false, nodes, local_id), + Cluster::create(nodes, local_id), )?); dummy_query_context.get_settings().set_max_threads(8)?; From 30d07cb05bcb24675e28a720a86d998f54c45a5b Mon Sep 17 00:00:00 2001 From: zhang2014 Date: Mon, 6 Jan 2025 01:03:07 +0800 Subject: [PATCH 46/53] feat(cluster): clean code --- .../management/src/warehouse/warehouse_api.rs | 2 +- .../management/src/warehouse/warehouse_mgr.rs | 224 ++++++++++++++++-- src/query/service/src/clusters/cluster.rs | 4 +- 3 files changed, 204 insertions(+), 26 deletions(-) diff --git a/src/query/management/src/warehouse/warehouse_api.rs b/src/query/management/src/warehouse/warehouse_api.rs index 090e003063f97..dffee08fe64b9 100644 --- a/src/query/management/src/warehouse/warehouse_api.rs +++ b/src/query/management/src/warehouse/warehouse_api.rs @@ -64,7 +64,7 @@ pub struct SystemManagedWarehouse { #[async_trait::async_trait] pub trait WarehouseApi: Sync + Send { /// Start a new node. - async fn start_node(&self, node: NodeInfo) -> Result; + async fn start_node(&self, node: NodeInfo) -> Result<(u64, NodeInfo)>; /// Shutdown the tenant's cluster one node by node.id. async fn shutdown_node(&self, node_id: String) -> Result<()>; diff --git a/src/query/management/src/warehouse/warehouse_mgr.rs b/src/query/management/src/warehouse/warehouse_mgr.rs index ff3e2e97bd2da..3168cc4ca54f0 100644 --- a/src/query/management/src/warehouse/warehouse_mgr.rs +++ b/src/query/management/src/warehouse/warehouse_mgr.rs @@ -55,6 +55,8 @@ static DEFAULT_CLUSTER_ID: &str = "default"; pub static WAREHOUSE_API_KEY_PREFIX: &str = "__fd_clusters_v6"; pub static WAREHOUSE_META_KEY_PREFIX: &str = "__fd_warehouses"; +type LostNodes = Vec; + // example: // __fd_clusters_v6/test_tenant/online_nodes // |- /RUV9DQArNnP4Hej4A74f07: NodeInfo { id: "RUV9DQArNnP4Hej4A74f07", secret: "", cpu_nums: 0, version: 0, http_address: "", flight_address: "", discovery_address: "", binary_version: "", node_type: SystemManaged, node_group: None, cluster_id: "", warehouse_id: "", runtime_node_group: None } @@ -354,34 +356,210 @@ impl WarehouseMgr { Ok(txn_reply) } + // return list of unhealthy warehouse clusters. + // return type: HashMap<(WarehouseId, ClusterId), LostNodes> + async fn unhealthy_warehouses(&self) -> Result> { + let online_nodes = self.list_online_nodes().await?; + let mut cluster_online_nodes = HashMap::with_capacity(online_nodes.len()); + + for node in online_nodes { + match cluster_online_nodes.entry((node.warehouse_id.clone(), node.cluster_id.clone())) { + Entry::Vacant(v) => { + v.insert(vec![node]); + } + Entry::Occupied(mut v) => { + v.get_mut().push(node); + } + } + } + + let list_reply = self + .metastore + .prefix_list_kv(&self.warehouse_info_key_prefix) + .await?; + + let mut unhealthy_warehouses = HashMap::with_capacity(list_reply.len()); + + for (_key, seq_v) in list_reply { + let wh: WarehouseInfo = serde_json::from_slice(&seq_v.data)?; + + // skip if self-managed warehouse + let WarehouseInfo::SystemManaged(wh) = wh else { + continue; + }; + + for (cluster_id, cluster) in wh.clusters { + let mut lost_nodes = cluster.nodes; + + let key = (wh.display_name.clone(), cluster_id.clone()); + if let Some(online_nodes) = cluster_online_nodes.remove(&key) { + for online_node in online_nodes { + lost_nodes.remove_first(&SelectedNode::Random( + online_node.runtime_node_group.clone(), + )); + } + } + + if !lost_nodes.is_empty() { + unhealthy_warehouses + .insert((wh.display_name.clone(), cluster_id.clone()), lost_nodes); + } + } + } + + Ok(unhealthy_warehouses) + } + + async fn recovery_system_managed_warehouse(&self, node: &mut NodeInfo) -> Result { + for retry_count in 0..20 { + let unhealthy_warehouses = self.unhealthy_warehouses().await?; + + if unhealthy_warehouses.is_empty() { + // All warehouses are healthy. + return Ok(false); + } + + let mut runtime_node_group = None; + let mut warehouse_and_cluster = None; + let match_group = SelectedNode::Random(node.node_group.clone()); + for ((warehouse, cluster), lost_nodes) in unhealthy_warehouses { + if lost_nodes.iter().any(|x| x == &match_group) { + runtime_node_group = node.node_group.clone(); + warehouse_and_cluster = Some((warehouse, cluster)); + break; + } else if warehouse_and_cluster.is_none() { + if !lost_nodes.iter().any(|x| x == &SelectedNode::Random(None)) { + continue; + } + + runtime_node_group = None; + warehouse_and_cluster = Some((warehouse, cluster)); + } + } + + let Some((warehouse, cluster)) = warehouse_and_cluster else { + // No warehouse needs to fixed + return Ok(false); + }; + + let Ok(warehouse_snapshot) = self.warehouse_snapshot(&warehouse).await else { + // The warehouse may have been dropped. + log::warn!("The warehouse {} may have been dropped.", warehouse); + continue; + }; + + let WarehouseInfo::SystemManaged(wh) = warehouse_snapshot.warehouse_info else { + // The warehouse may have been dropped and started self-managed warehouse with same name. + log::warn!("The warehouse({}) may have been dropped and started self-managed warehouse with same name.", warehouse); + continue; + }; + + let Some(recovery_cluster) = wh.clusters.get(&cluster) else { + // The cluster may have been dropped + log::warn!( + "The cluster({}/{}) may have been dropped", + warehouse, + cluster + ); + continue; + }; + + // Check the node list again; if conflicts occur, skip this correction. + { + let mut lost_nodes = recovery_cluster.nodes.clone(); + + for node_snapshot in warehouse_snapshot.snapshot_nodes { + if node_snapshot.node_info.cluster_id != cluster { + continue; + } + + lost_nodes.remove_first(&SelectedNode::Random( + node_snapshot.node_info.runtime_node_group.clone(), + )); + } + + if !lost_nodes + .iter() + .any(|x| x == &SelectedNode::Random(runtime_node_group.clone())) + { + continue; + } + } + + let mut txn = TxnRequest::default(); + + let mut node_with_wh = node.clone(); + + node_with_wh.cluster_id = cluster.clone(); + node_with_wh.warehouse_id = warehouse.clone(); + node.runtime_node_group = runtime_node_group.clone(); + node_with_wh.runtime_node_group = runtime_node_group; + + let node_key = self.node_key(&node_with_wh)?; + let cluster_node_key = self.cluster_node_key(&node_with_wh)?; + let warehouse_info_key = self.warehouse_info_key(&warehouse)?; + + txn.condition = vec![ + map_condition(&node_key, MatchSeq::Exact(0)), + map_condition(&cluster_node_key, MatchSeq::Exact(0)), + map_condition( + &warehouse_info_key, + MatchSeq::Exact(warehouse_snapshot.info_seq), + ), + ]; + + txn.if_then = vec![ + TxnOp::put_with_ttl( + node_key, + serde_json::to_vec(&node_with_wh)?, + Some(self.lift_time), + ), + TxnOp::put_with_ttl( + cluster_node_key, + serde_json::to_vec(&node)?, + Some(self.lift_time), + ), + TxnOp::put( + warehouse_info_key, + serde_json::to_vec(&WarehouseInfo::SystemManaged(wh))?, + ), + ]; + + if self.metastore.transaction(txn).await?.success { + return Ok(true); + } + + // upon retry, fallback a little while. + tokio::time::sleep(Duration::from_millis(30 * retry_count)).await; + } + + Err(ErrorCode::WarehouseOperateConflict( + "Warehouse operate conflict(tried 20 times).".to_string(), + )) + } + async fn start_system_managed(&self, mut node: NodeInfo) -> Result { - // TODO: find attach warehouse let mut txn = TxnRequest::default(); let node_key = self.node_key(&node)?; - txn.condition - .push(map_condition(&node_key, MatchSeq::Exact(0))); - - txn.if_then.push(TxnOp::put_with_ttl( - node_key.clone(), - serde_json::to_vec(&node)?, - Some(self.lift_time), - )); + if !self.recovery_system_managed_warehouse(&mut node).await? { + txn.else_then.push(TxnOp::get(node_key.clone())); + txn.condition + .push(map_condition(&node_key, MatchSeq::Exact(0))); - // If the warehouse has already been assigned. - if !node.cluster_id.is_empty() && !node.warehouse_id.is_empty() { - let cluster_node_key = self.cluster_node_key(&node)?; + txn.if_then = vec![ + TxnOp::put_with_ttl( + node_key.clone(), + serde_json::to_vec(&node)?, + Some(self.lift_time), + ), + TxnOp::get(node_key.clone()), + ]; - self.unload_warehouse_info(&mut node); - txn.if_then.push(TxnOp::put_with_ttl( - cluster_node_key, - serde_json::to_vec(&node)?, - Some(self.lift_time), - )); + return Ok(self.metastore.transaction(txn).await?); } - txn.if_then.push(TxnOp::get(node_key.clone())); - txn.else_then.push(TxnOp::get(node_key.clone())); + txn.if_then = vec![TxnOp::get(node_key.clone())]; Ok(self.metastore.transaction(txn).await?) } @@ -399,7 +577,7 @@ impl WarehouseMgr { )); // If the warehouse has already been assigned. - if !node.cluster_id.is_empty() && !node.warehouse_id.is_empty() { + if node.assigned_warehouse() { let cluster_node_key = self.cluster_node_key(&node)?; self.unload_warehouse_info(&mut node); @@ -746,7 +924,7 @@ impl WarehouseMgr { impl WarehouseApi for WarehouseMgr { #[async_backtrace::framed] #[fastrace::trace] - async fn start_node(&self, node: NodeInfo) -> Result { + async fn start_node(&self, node: NodeInfo) -> Result<(u64, NodeInfo)> { let start_reply = match node.node_type { NodeType::SelfManaged => self.start_self_managed(node.clone()).await?, NodeType::SystemManaged => self.start_system_managed(node.clone()).await?, @@ -765,7 +943,7 @@ impl WarehouseApi for WarehouseMgr { Some(Response::Get(TxnGetResponse { value: Some(seq_v), .. })), - }) => Ok(seq_v.seq), + }) => Ok((seq_v.seq, serde_json::from_slice(&seq_v.data)?)), _ => Err(ErrorCode::MetaServiceError("Add node info failure.")), } } diff --git a/src/query/service/src/clusters/cluster.rs b/src/query/service/src/clusters/cluster.rs index 8509e4f2438a2..023dd7fca1404 100644 --- a/src/query/service/src/clusters/cluster.rs +++ b/src/query/service/src/clusters/cluster.rs @@ -506,8 +506,8 @@ impl ClusterDiscovery { self.drop_invalid_nodes(&node_info).await?; - match self.warehouse_manager.start_node(node_info.clone()).await { - Ok(seq) => self.start_heartbeat(node_info, seq).await, + match self.warehouse_manager.start_node(node_info).await { + Ok((seq, node_info)) => self.start_heartbeat(node_info, seq).await, Err(cause) => Err(cause.add_message_back("(while cluster api add_node).")), } } From 5da3a5e74757317f2e3b3302b77e755ccd037b47 Mon Sep 17 00:00:00 2001 From: zhang2014 Date: Mon, 6 Jan 2025 01:12:03 +0800 Subject: [PATCH 47/53] feat(cluster): clean code --- src/query/management/tests/it/warehouse.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/query/management/tests/it/warehouse.rs b/src/query/management/tests/it/warehouse.rs index ea5dd7ff3eee6..6b3b68bfcc878 100644 --- a/src/query/management/tests/it/warehouse.rs +++ b/src/query/management/tests/it/warehouse.rs @@ -212,7 +212,7 @@ async fn test_successfully_heartbeat_self_managed_node() -> Result<()> { let (kv, warehouse_manager, _nodes) = nodes(Duration::from_mins(60), 0).await?; let mut node_info = self_managed_node("test_node"); - let seq = warehouse_manager.start_node(node_info.clone()).await?; + let (seq, _node) = warehouse_manager.start_node(node_info.clone()).await?; let info_key = "__fd_clusters_v6/test%2dtenant%2did/online_nodes/test_node"; assert_key_value(&kv, info_key, serde_json::to_vec(&node_info)?).await; From 724477c8f3ba13c20541a0b085c79b8b6b215eaf Mon Sep 17 00:00:00 2001 From: zhang2014 Date: Mon, 6 Jan 2025 10:08:18 +0800 Subject: [PATCH 48/53] feat(cluster): clean code --- .../management/src/warehouse/warehouse_api.rs | 2 +- .../management/src/warehouse/warehouse_mgr.rs | 41 +++++++++---------- src/query/management/tests/it/warehouse.rs | 20 ++++----- .../interpreter_show_warehouses.rs | 2 +- 4 files changed, 32 insertions(+), 33 deletions(-) diff --git a/src/query/management/src/warehouse/warehouse_api.rs b/src/query/management/src/warehouse/warehouse_api.rs index dffee08fe64b9..f6f96189f6416 100644 --- a/src/query/management/src/warehouse/warehouse_api.rs +++ b/src/query/management/src/warehouse/warehouse_api.rs @@ -55,8 +55,8 @@ pub struct SystemManagedCluster { #[derive(serde::Serialize, serde::Deserialize, Eq, PartialEq, Debug)] pub struct SystemManagedWarehouse { pub id: String, + pub role_id: String, pub status: String, - pub display_name: String, pub clusters: HashMap, } diff --git a/src/query/management/src/warehouse/warehouse_mgr.rs b/src/query/management/src/warehouse/warehouse_mgr.rs index 3168cc4ca54f0..223c99c8403c7 100644 --- a/src/query/management/src/warehouse/warehouse_mgr.rs +++ b/src/query/management/src/warehouse/warehouse_mgr.rs @@ -391,7 +391,7 @@ impl WarehouseMgr { for (cluster_id, cluster) in wh.clusters { let mut lost_nodes = cluster.nodes; - let key = (wh.display_name.clone(), cluster_id.clone()); + let key = (wh.id.clone(), cluster_id.clone()); if let Some(online_nodes) = cluster_online_nodes.remove(&key) { for online_node in online_nodes { lost_nodes.remove_first(&SelectedNode::Random( @@ -401,8 +401,7 @@ impl WarehouseMgr { } if !lost_nodes.is_empty() { - unhealthy_warehouses - .insert((wh.display_name.clone(), cluster_id.clone()), lost_nodes); + unhealthy_warehouses.insert((wh.id.clone(), cluster_id.clone()), lost_nodes); } } } @@ -1108,9 +1107,9 @@ impl WarehouseApi for WarehouseMgr { txn.if_then.push(TxnOp::put( warehouse_info_key.clone(), serde_json::to_vec(&WarehouseInfo::SystemManaged(SystemManagedWarehouse { - id: GlobalUniqName::unique(), + role_id: GlobalUniqName::unique(), status: "Running".to_string(), - display_name: warehouse.clone(), + id: warehouse.clone(), clusters: HashMap::from([( String::from(DEFAULT_CLUSTER_ID), SystemManagedCluster { @@ -1168,9 +1167,9 @@ impl WarehouseApi for WarehouseMgr { // TODO: support cluster resume? need_schedule_cluster = warehouse.clusters.clone(); Ok(WarehouseInfo::SystemManaged(SystemManagedWarehouse { - id: warehouse.id.clone(), + role_id: warehouse.role_id.clone(), status: "Running".to_string(), - display_name: warehouse.display_name, + id: warehouse.id, clusters: warehouse.clusters, })) } @@ -1272,13 +1271,13 @@ impl WarehouseApi for WarehouseMgr { )), WarehouseInfo::SystemManaged(warehouse) => { if warehouse.status.to_uppercase() != "RUNNING" { - return Err(ErrorCode::InvalidWarehouse(format!("Cannot suspend warehouse {:?}, because warehouse state is not running.", warehouse.display_name))); + return Err(ErrorCode::InvalidWarehouse(format!("Cannot suspend warehouse {:?}, because warehouse state is not running.", warehouse.id))); } Ok(WarehouseInfo::SystemManaged(SystemManagedWarehouse { - id: warehouse.id.clone(), + role_id: warehouse.role_id.clone(), status: "Suspended".to_string(), - display_name: warehouse.display_name, + id: warehouse.id, clusters: warehouse.clusters, })) } @@ -1360,7 +1359,7 @@ impl WarehouseApi for WarehouseMgr { "Cannot rename self-managed warehouse", )), WarehouseInfo::SystemManaged(mut info) => { - info.display_name = to.clone(); + info.id = to.clone(); Ok(WarehouseInfo::SystemManaged(info)) } }?; @@ -1511,9 +1510,9 @@ impl WarehouseApi for WarehouseMgr { nodes: nodes.clone(), }); Ok(WarehouseInfo::SystemManaged(SystemManagedWarehouse { - id: info.id, + role_id: info.role_id, status: info.status, - display_name: info.display_name, + id: info.id, clusters: info.clusters, })) } @@ -1620,9 +1619,9 @@ impl WarehouseApi for WarehouseMgr { false => { info.clusters.remove(&cluster); Ok(WarehouseInfo::SystemManaged(SystemManagedWarehouse { - id: info.id, + role_id: info.role_id, status: info.status, - display_name: info.display_name, + id: info.id, clusters: info.clusters, })) } @@ -1721,9 +1720,9 @@ impl WarehouseApi for WarehouseMgr { let cluster_info = info.clusters.remove(&cur); info.clusters.insert(to.clone(), cluster_info.unwrap()); Ok(WarehouseInfo::SystemManaged(SystemManagedWarehouse { - id: info.id, + role_id: info.role_id, status: info.status, - display_name: info.display_name, + id: info.id, clusters: info.clusters, })) } @@ -1836,9 +1835,9 @@ impl WarehouseApi for WarehouseMgr { } Ok(WarehouseInfo::SystemManaged(SystemManagedWarehouse { - id: info.id, + role_id: info.role_id, status: info.status, - display_name: info.display_name, + id: info.id, clusters: info.clusters, })) } @@ -1958,9 +1957,9 @@ impl WarehouseApi for WarehouseMgr { Ok(WarehouseInfo::SystemManaged(SystemManagedWarehouse { - id: info.id, + role_id: info.role_id, status: info.status, - display_name: info.display_name, + id: info.id, clusters: info.clusters, })) } diff --git a/src/query/management/tests/it/warehouse.rs b/src/query/management/tests/it/warehouse.rs index 6b3b68bfcc878..ad1223fc91ea1 100644 --- a/src/query/management/tests/it/warehouse.rs +++ b/src/query/management/tests/it/warehouse.rs @@ -608,9 +608,9 @@ async fn test_list_warehouses() -> Result<()> { let WarehouseInfo::SystemManaged(system_managed_info) = &list_warehouses_2[1] else { unreachable!(); }; - assert!(!system_managed_info.id.is_empty()); + assert!(!system_managed_info.role_id.is_empty()); assert_eq!(system_managed_info.status, "Running"); - assert_eq!(system_managed_info.display_name, "test_warehouse_1"); + assert_eq!(system_managed_info.id, "test_warehouse_1"); assert_eq!( system_managed_info.clusters, HashMap::from([(String::from("default"), SystemManagedCluster { @@ -650,9 +650,9 @@ async fn test_list_warehouses() -> Result<()> { let WarehouseInfo::SystemManaged(system_managed_info) = &list_warehouses_4[3] else { unreachable!(); }; - assert!(!system_managed_info.id.is_empty()); + assert!(!system_managed_info.role_id.is_empty()); assert_eq!(system_managed_info.status, "Running"); - assert_eq!(system_managed_info.display_name, "test_warehouse_3"); + assert_eq!(system_managed_info.id, "test_warehouse_3"); assert_eq!( system_managed_info.clusters, HashMap::from([(String::from("default"), SystemManagedCluster { @@ -751,9 +751,9 @@ async fn test_rename_warehouses() -> Result<()> { unreachable!(); }; - assert!(!system_managed_info.id.is_empty()); + assert!(!system_managed_info.role_id.is_empty()); assert_eq!(system_managed_info.status, "Running"); - assert_eq!(system_managed_info.display_name, "test_warehouse"); + assert_eq!(system_managed_info.id, "test_warehouse"); assert_eq!( system_managed_info.clusters, HashMap::from([(String::from("default"), SystemManagedCluster { @@ -776,9 +776,9 @@ async fn test_rename_warehouses() -> Result<()> { unreachable!(); }; - assert!(!system_managed_info.id.is_empty()); + assert!(!system_managed_info.role_id.is_empty()); assert_eq!(system_managed_info.status, "Running"); - assert_eq!(system_managed_info.display_name, "new_test_warehouse"); + assert_eq!(system_managed_info.id, "new_test_warehouse"); assert_eq!( system_managed_info.clusters, HashMap::from([(String::from("default"), SystemManagedCluster { @@ -805,9 +805,9 @@ async fn test_rename_warehouses() -> Result<()> { unreachable!(); }; - assert!(!system_managed_info.id.is_empty()); + assert!(!system_managed_info.role_id.is_empty()); assert_eq!(system_managed_info.status, "Running"); - assert_eq!(system_managed_info.display_name, "test_warehouse"); + assert_eq!(system_managed_info.id, "test_warehouse"); assert_eq!( system_managed_info.clusters, HashMap::from([(String::from("default"), SystemManagedCluster { diff --git a/src/query/service/src/interpreters/interpreter_show_warehouses.rs b/src/query/service/src/interpreters/interpreter_show_warehouses.rs index 883318df7f930..be44d4993c717 100644 --- a/src/query/service/src/interpreters/interpreter_show_warehouses.rs +++ b/src/query/service/src/interpreters/interpreter_show_warehouses.rs @@ -71,7 +71,7 @@ impl Interpreter for ShowWarehousesInterpreter { warehouses_status.push(Scalar::String(String::from("Running")).as_ref()); } WarehouseInfo::SystemManaged(v) => { - warehouses_name.push(Scalar::String(v.display_name.clone()).as_ref()); + warehouses_name.push(Scalar::String(v.id.clone()).as_ref()); warehouses_type.push(Scalar::String(String::from("System-Managed")).as_ref()); warehouses_status.push(Scalar::String(v.status.clone()).as_ref()); } From 9d5506ad064b958600aeda2ae9ef496bb9e17db5 Mon Sep 17 00:00:00 2001 From: zhang2014 Date: Mon, 6 Jan 2025 13:01:28 +0800 Subject: [PATCH 49/53] feat(cluster): clean code --- .../management/src/warehouse/warehouse_mgr.rs | 15 +++++---------- 1 file changed, 5 insertions(+), 10 deletions(-) diff --git a/src/query/management/src/warehouse/warehouse_mgr.rs b/src/query/management/src/warehouse/warehouse_mgr.rs index 223c99c8403c7..0fee9ebb4d7b2 100644 --- a/src/query/management/src/warehouse/warehouse_mgr.rs +++ b/src/query/management/src/warehouse/warehouse_mgr.rs @@ -116,13 +116,8 @@ impl WarehouseMgr { txn.if_then.push(TxnOp::delete(node_key)); if node_info.assigned_warehouse() { - txn.if_then.push(TxnOp::delete(format!( - "{}/{}/{}/{}", - self.cluster_node_key_prefix, - escape_for_key(&node_info.warehouse_id)?, - escape_for_key(&node_info.cluster_id)?, - escape_for_key(&node_info.id)? - ))); + let cluster_node_key = self.cluster_node_key(node_info)?; + txn.if_then.push(TxnOp::delete(cluster_node_key)); } match self.metastore.transaction(txn).await?.success { @@ -2024,20 +2019,20 @@ impl WarehouseApi for WarehouseMgr { warehouse: &str, cluster: &str, ) -> Result> { - let cluster_prefix = format!( + let cluster_nodes_prefix = format!( "{}/{}/{}/", self.cluster_node_key_prefix, escape_for_key(warehouse)?, escape_for_key(cluster)? ); - let values = self.metastore.prefix_list_kv(&cluster_prefix).await?; + let values = self.metastore.prefix_list_kv(&cluster_nodes_prefix).await?; let mut nodes_info = Vec::with_capacity(values.len()); for (node_key, value) in values { let mut node_info = serde_json::from_slice::(&value.data)?; - node_info.id = unescape_for_key(&node_key[cluster_prefix.len()..])?; + node_info.id = unescape_for_key(&node_key[cluster_nodes_prefix.len()..])?; node_info.cluster_id = cluster.to_string(); node_info.warehouse_id = warehouse.to_string(); nodes_info.push(node_info); From 2e956e859a19131f3dfd5000116a10fc6f0733db Mon Sep 17 00:00:00 2001 From: zhang2014 Date: Mon, 6 Jan 2025 14:19:30 +0800 Subject: [PATCH 50/53] feat(cluster): add concurrent unit test --- src/query/management/tests/it/warehouse.rs | 179 +++++++++++++++++++++ 1 file changed, 179 insertions(+) diff --git a/src/query/management/tests/it/warehouse.rs b/src/query/management/tests/it/warehouse.rs index ad1223fc91ea1..d221d5f08a3c0 100644 --- a/src/query/management/tests/it/warehouse.rs +++ b/src/query/management/tests/it/warehouse.rs @@ -17,7 +17,10 @@ use std::sync::Arc; use std::time::Duration; use databend_common_base::base::tokio; +use databend_common_base::base::tokio::sync::Barrier; use databend_common_base::base::GlobalUniqName; +use databend_common_base::runtime::Runtime; +use databend_common_base::runtime::TrySpawn; use databend_common_exception::Result; use databend_common_management::*; use databend_common_meta_embedded::MemMeta; @@ -403,6 +406,44 @@ async fn test_create_system_managed_warehouse_with_online_node() -> Result<()> { Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn test_concurrent_create_warehouse() -> Result<()> { + let (_, warehouse_manager, _nodes) = nodes(Duration::from_mins(30), 9).await?; + + let barrier = Arc::new(Barrier::new(10)); + let warehouse_manager = Arc::new(warehouse_manager); + + let mut runtimes = Vec::with_capacity(10); + let mut join_handler = Vec::with_capacity(10); + for idx in 0..10 { + let runtime = Arc::new(Runtime::with_worker_threads(2, None)?); + + runtimes.push(runtime.clone()); + + join_handler.push(runtime.spawn({ + let barrier = barrier.clone(); + let warehouse_manager = warehouse_manager.clone(); + async move { + let _ = barrier.wait().await; + + let create_warehouse = warehouse_manager.create_warehouse( + format!("warehouse_{}", idx), + vec![SelectedNode::Random(None); 1], + ); + + create_warehouse.await.is_ok() + } + })); + } + + let create_res = futures::future::try_join_all(join_handler).await?; + + assert_eq!(create_res.len(), 10); + assert_eq!(create_res.iter().filter(|x| **x).count(), 9); + + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 1)] async fn test_create_duplicated_warehouse() -> Result<()> { let (_, warehouse_manager, _nodes) = nodes(Duration::from_mins(30), 2).await?; @@ -493,6 +534,144 @@ async fn test_create_warehouse_with_no_resources() -> Result<()> { Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn test_recovery_create_warehouse() -> Result<()> { + let (_, warehouse_manager, nodes) = nodes(Duration::from_mins(30), 2).await?; + + let create_warehouse = warehouse_manager.create_warehouse( + String::from("test_warehouse"), + vec![SelectedNode::Random(None); 2], + ); + let _ = create_warehouse.await?; + + let list_warehouse_nodes = + warehouse_manager.list_warehouse_nodes(String::from("test_warehouse")); + + assert_eq!(list_warehouse_nodes.await?.len(), 2); + + let shutdown_node = warehouse_manager.shutdown_node(nodes[0].clone()); + shutdown_node.await?; + + let shutdown_node = warehouse_manager.shutdown_node(nodes[1].clone()); + shutdown_node.await?; + + let list_warehouse_nodes = + warehouse_manager.list_warehouse_nodes(String::from("test_warehouse")); + + assert_eq!(list_warehouse_nodes.await?.len(), 0); + + let node_1 = GlobalUniqName::unique(); + let start_node_1 = warehouse_manager.start_node(system_managed_node(&node_1)); + assert!(start_node_1.await.is_ok()); + + let list_warehouse_nodes = + warehouse_manager.list_warehouse_nodes(String::from("test_warehouse")); + + let nodes = list_warehouse_nodes + .await? + .into_iter() + .map(|x| x.id) + .collect::>(); + assert_eq!(nodes.len(), 1); + assert!(nodes.contains(&node_1)); + + let node_2 = GlobalUniqName::unique(); + let mut node_info_2 = system_managed_node(&node_2); + node_info_2.node_group = Some(String::from("test_group")); + let start_node_2 = warehouse_manager.start_node(node_info_2); + assert!(start_node_2.await.is_ok()); + + let list_warehouse_nodes = + warehouse_manager.list_warehouse_nodes(String::from("test_warehouse")); + + let nodes = list_warehouse_nodes + .await? + .into_iter() + .map(|x| x.id) + .collect::>(); + assert_eq!(nodes.len(), 2); + assert!(nodes.contains(&node_1)); + assert!(nodes.contains(&node_2)); + + // warehouse is fixed + let node_3 = GlobalUniqName::unique(); + let start_node_3 = warehouse_manager.start_node(system_managed_node(&node_3)); + assert!(start_node_3.await.is_ok()); + + let list_warehouse_nodes = + warehouse_manager.list_warehouse_nodes(String::from("test_warehouse")); + + let nodes = list_warehouse_nodes + .await? + .into_iter() + .map(|x| x.id) + .collect::>(); + assert_eq!(nodes.len(), 2); + assert!(nodes.contains(&node_1)); + assert!(nodes.contains(&node_2)); + + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn test_concurrent_recovery_create_warehouse() -> Result<()> { + let (_, warehouse_manager, nodes) = nodes(Duration::from_mins(30), 2).await?; + + let create_warehouse = warehouse_manager.create_warehouse( + String::from("test_warehouse"), + vec![SelectedNode::Random(None); 2], + ); + let _ = create_warehouse.await?; + + let list_warehouse_nodes = + warehouse_manager.list_warehouse_nodes(String::from("test_warehouse")); + + assert_eq!(list_warehouse_nodes.await?.len(), 2); + + let shutdown_node = warehouse_manager.shutdown_node(nodes[0].clone()); + shutdown_node.await?; + + let shutdown_node = warehouse_manager.shutdown_node(nodes[1].clone()); + shutdown_node.await?; + + let barrier = Arc::new(Barrier::new(10)); + let warehouse_manager = Arc::new(warehouse_manager); + + let mut runtimes = Vec::with_capacity(10); + let mut join_handler = Vec::with_capacity(10); + for _idx in 0..10 { + let runtime = Arc::new(Runtime::with_worker_threads(2, None)?); + + runtimes.push(runtime.clone()); + + join_handler.push(runtime.spawn({ + let barrier = barrier.clone(); + let warehouse_manager = warehouse_manager.clone(); + async move { + let _ = barrier.wait().await; + + let node_id = GlobalUniqName::unique(); + let start_node = warehouse_manager.start_node(system_managed_node(&node_id)); + + let (_, node_info) = start_node.await.unwrap(); + node_info.id + } + })); + } + + let start_res = futures::future::try_join_all(join_handler).await?; + + assert_eq!(start_res.len(), 10); + let list_warehouse_nodes = + warehouse_manager.list_warehouse_nodes(String::from("test_warehouse")); + assert_eq!(list_warehouse_nodes.await?.len(), 2); + + let list_online_nodes = warehouse_manager.list_online_nodes(); + assert_eq!(list_online_nodes.await?.len(), 10); + + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 1)] async fn test_drop_empty_warehouse() -> Result<()> { let (_, warehouse_manager, _nodes) = nodes(Duration::from_mins(30), 2).await?; From bd993204e2bd943c32366f1b2f4d95b673a0e9cf Mon Sep 17 00:00:00 2001 From: zhang2014 Date: Mon, 6 Jan 2025 14:38:27 +0800 Subject: [PATCH 51/53] feat(cluster): add concurrent unit test --- src/query/management/tests/it/warehouse.rs | 134 ++++++++++++++++++++- 1 file changed, 132 insertions(+), 2 deletions(-) diff --git a/src/query/management/tests/it/warehouse.rs b/src/query/management/tests/it/warehouse.rs index d221d5f08a3c0..823ead83c29ac 100644 --- a/src/query/management/tests/it/warehouse.rs +++ b/src/query/management/tests/it/warehouse.rs @@ -542,7 +542,7 @@ async fn test_recovery_create_warehouse() -> Result<()> { String::from("test_warehouse"), vec![SelectedNode::Random(None); 2], ); - let _ = create_warehouse.await?; + create_warehouse.await?; let list_warehouse_nodes = warehouse_manager.list_warehouse_nodes(String::from("test_warehouse")); @@ -621,7 +621,7 @@ async fn test_concurrent_recovery_create_warehouse() -> Result<()> { String::from("test_warehouse"), vec![SelectedNode::Random(None); 2], ); - let _ = create_warehouse.await?; + create_warehouse.await?; let list_warehouse_nodes = warehouse_manager.list_warehouse_nodes(String::from("test_warehouse")); @@ -1004,6 +1004,136 @@ async fn test_rename_warehouses() -> Result<()> { Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn test_concurrent_rename_to_warehouse() -> Result<()> { + let (_, warehouse_manager, _nodes) = nodes(Duration::from_mins(30), 9).await?; + + let create_warehouse = warehouse_manager.create_warehouse( + String::from("test_warehouse"), + vec![SelectedNode::Random(None); 9], + ); + create_warehouse.await?; + + let barrier = Arc::new(Barrier::new(10)); + let warehouse_manager = Arc::new(warehouse_manager); + + let mut runtimes = Vec::with_capacity(10); + let mut join_handler = Vec::with_capacity(10); + for idx in 0..10 { + let runtime = Arc::new(Runtime::with_worker_threads(2, None)?); + + runtimes.push(runtime.clone()); + + join_handler.push(runtime.spawn({ + let barrier = barrier.clone(); + let warehouse_manager = warehouse_manager.clone(); + async move { + let _ = barrier.wait().await; + + let rename_warehouse = warehouse_manager.rename_warehouse( + String::from("test_warehouse"), + format!("test_warehouse_{}", idx), + ); + + rename_warehouse.await.is_ok() + } + })); + } + + let rename_res = futures::future::try_join_all(join_handler).await?; + + assert_eq!(rename_res.len(), 10); + assert_eq!(rename_res.iter().filter(|x| **x).count(), 1); + + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn test_concurrent_rename_from_warehouse() -> Result<()> { + let (_, warehouse_manager, _nodes) = nodes(Duration::from_mins(30), 10).await?; + + let barrier = Arc::new(Barrier::new(10)); + let warehouse_manager = Arc::new(warehouse_manager); + + let mut runtimes = Vec::with_capacity(10); + let mut join_handler = Vec::with_capacity(10); + for idx in 0..10 { + let runtime = Arc::new(Runtime::with_worker_threads(2, None)?); + + runtimes.push(runtime.clone()); + + join_handler.push(runtime.spawn({ + let barrier = barrier.clone(); + let warehouse_manager = warehouse_manager.clone(); + async move { + let create_warehouse = warehouse_manager.create_warehouse( + format!("test_warehouse_{}", idx), + vec![SelectedNode::Random(None); 1], + ); + create_warehouse.await.unwrap(); + + let _ = barrier.wait().await; + + let rename_warehouse = warehouse_manager.rename_warehouse( + format!("test_warehouse_{}", idx), + String::from("test_warehouse"), + ); + + rename_warehouse.await.is_ok() + } + })); + } + + let rename_res = futures::future::try_join_all(join_handler).await?; + + assert_eq!(rename_res.len(), 10); + assert_eq!(rename_res.iter().filter(|x| **x).count(), 1); + + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn test_concurrent_drop_warehouse() -> Result<()> { + let (_, warehouse_manager, _nodes) = nodes(Duration::from_mins(30), 9).await?; + + let create_warehouse = warehouse_manager.create_warehouse( + String::from("test_warehouse"), + vec![SelectedNode::Random(None); 9], + ); + create_warehouse.await?; + + let barrier = Arc::new(Barrier::new(10)); + let warehouse_manager = Arc::new(warehouse_manager); + + let mut runtimes = Vec::with_capacity(10); + let mut join_handler = Vec::with_capacity(10); + for _idx in 0..10 { + let runtime = Arc::new(Runtime::with_worker_threads(2, None)?); + + runtimes.push(runtime.clone()); + + join_handler.push(runtime.spawn({ + let barrier = barrier.clone(); + let warehouse_manager = warehouse_manager.clone(); + async move { + let _ = barrier.wait().await; + + let drop_warehouse = + warehouse_manager.drop_warehouse(String::from("test_warehouse")); + + drop_warehouse.await.is_ok() + } + })); + } + + let drop_res = futures::future::try_join_all(join_handler).await?; + + assert_eq!(drop_res.len(), 10); + assert_eq!(drop_res.iter().filter(|x| **x).count(), 1); + + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 1)] async fn test_drop_warehouse_cluster_failure() -> Result<()> { let (_, warehouse_manager, _nodes) = nodes(Duration::from_mins(30), 1).await?; From 66587720066f36fd80bf8d313bcb3d11d25de42a Mon Sep 17 00:00:00 2001 From: zhang2014 Date: Mon, 6 Jan 2025 14:46:51 +0800 Subject: [PATCH 52/53] feat(cluster): clean code --- .../management/src/warehouse/warehouse_mgr.rs | 34 +++++++++++++++++-- 1 file changed, 32 insertions(+), 2 deletions(-) diff --git a/src/query/management/src/warehouse/warehouse_mgr.rs b/src/query/management/src/warehouse/warehouse_mgr.rs index 0fee9ebb4d7b2..b2a8b76fdeb05 100644 --- a/src/query/management/src/warehouse/warehouse_mgr.rs +++ b/src/query/management/src/warehouse/warehouse_mgr.rs @@ -557,6 +557,36 @@ impl WarehouseMgr { Ok(self.metastore.transaction(txn).await?) } + async fn update_system_managed_node(&self, mut node: NodeInfo, seq: u64) -> Result { + let mut txn = TxnRequest::default(); + let node_key = self.node_key(&node)?; + + txn.condition + .push(map_condition(&node_key, MatchSeq::Exact(seq))); + + txn.if_then.push(TxnOp::put_with_ttl( + node_key.clone(), + serde_json::to_vec(&node)?, + Some(self.lift_time), + )); + + // If the warehouse has already been assigned. + if node.assigned_warehouse() { + let cluster_node_key = self.cluster_node_key(&node)?; + + self.unload_warehouse_info(&mut node); + txn.if_then.push(TxnOp::put_with_ttl( + cluster_node_key, + serde_json::to_vec(&node)?, + Some(self.lift_time), + )); + } + + txn.if_then.push(TxnOp::get(node_key.clone())); + txn.else_then.push(TxnOp::get(node_key.clone())); + Ok(self.metastore.transaction(txn).await?) + } + async fn heartbeat_system_managed(&self, mut node: NodeInfo, seq: u64) -> Result { let mut txn = TxnRequest::default(); let node_key = self.node_key(&node)?; @@ -591,7 +621,7 @@ impl WarehouseMgr { let leave_node = node_info.leave_warehouse(); let reply = self - .heartbeat_system_managed(leave_node.clone(), seq) + .update_system_managed_node(leave_node.clone(), seq) .await?; if !reply.success { @@ -614,7 +644,7 @@ impl WarehouseMgr { } async fn join_warehouse(&self, node_info: NodeInfo, seq: u64) -> Result { - let reply = self.heartbeat_system_managed(node_info, seq).await?; + let reply = self.update_system_managed_node(node_info, seq).await?; if !reply.success { return Err(ErrorCode::WarehouseOperateConflict( From 5928a0d49ad2519fd9fe7acd1151a060d47d9f23 Mon Sep 17 00:00:00 2001 From: zhang2014 Date: Mon, 6 Jan 2025 15:19:51 +0800 Subject: [PATCH 53/53] feat(cluster): clean code --- src/meta/types/src/cluster.rs | 21 ++++++ .../management/src/warehouse/warehouse_api.rs | 3 +- .../management/src/warehouse/warehouse_mgr.rs | 72 +++++++++---------- src/query/management/tests/it/warehouse.rs | 12 ++-- src/query/service/src/clusters/cluster.rs | 8 ++- 5 files changed, 69 insertions(+), 47 deletions(-) diff --git a/src/meta/types/src/cluster.rs b/src/meta/types/src/cluster.rs index a5e8f6fcd4965..eff7dd5bbb013 100644 --- a/src/meta/types/src/cluster.rs +++ b/src/meta/types/src/cluster.rs @@ -138,6 +138,27 @@ impl NodeInfo { !self.warehouse_id.is_empty() && !self.cluster_id.is_empty() } + // Unload the warehouse and cluster from the node. + // 1. Used when a node is removed from the cluster. + // 2. For cluster_node_key: node_info, since its warehouse and cluster are already encoded in the key, we do not need to write the warehouse and cluster into its value again. + pub fn unload_warehouse_info(&self) -> NodeInfo { + NodeInfo { + id: self.id.clone(), + secret: self.secret.clone(), + cpu_nums: self.cpu_nums, + version: self.version, + http_address: self.http_address.clone(), + flight_address: self.flight_address.clone(), + discovery_address: self.discovery_address.clone(), + binary_version: self.binary_version.clone(), + node_type: self.node_type.clone(), + node_group: self.node_group.clone(), + cluster_id: String::new(), + warehouse_id: String::new(), + runtime_node_group: self.runtime_node_group.clone(), + } + } + pub fn leave_warehouse(&self) -> NodeInfo { NodeInfo { id: self.id.clone(), diff --git a/src/query/management/src/warehouse/warehouse_api.rs b/src/query/management/src/warehouse/warehouse_api.rs index f6f96189f6416..aa1afecaaee9a 100644 --- a/src/query/management/src/warehouse/warehouse_api.rs +++ b/src/query/management/src/warehouse/warehouse_api.rs @@ -16,6 +16,7 @@ use std::collections::HashMap; use databend_common_exception::Result; use databend_common_meta_types::NodeInfo; +use databend_common_meta_types::SeqV; /// Databend-query cluster ID. /// @@ -64,7 +65,7 @@ pub struct SystemManagedWarehouse { #[async_trait::async_trait] pub trait WarehouseApi: Sync + Send { /// Start a new node. - async fn start_node(&self, node: NodeInfo) -> Result<(u64, NodeInfo)>; + async fn start_node(&self, node: NodeInfo) -> Result>; /// Shutdown the tenant's cluster one node by node.id. async fn shutdown_node(&self, node_id: String) -> Result<()>; diff --git a/src/query/management/src/warehouse/warehouse_mgr.rs b/src/query/management/src/warehouse/warehouse_mgr.rs index b2a8b76fdeb05..6b93ead41bc23 100644 --- a/src/query/management/src/warehouse/warehouse_mgr.rs +++ b/src/query/management/src/warehouse/warehouse_mgr.rs @@ -35,6 +35,7 @@ use databend_common_meta_types::MatchSeqExt; use databend_common_meta_types::MetaError; use databend_common_meta_types::NodeInfo; use databend_common_meta_types::NodeType; +use databend_common_meta_types::SeqV; use databend_common_meta_types::TxnCondition; use databend_common_meta_types::TxnGetResponse; use databend_common_meta_types::TxnOp; @@ -177,24 +178,14 @@ impl WarehouseMgr { )) } - // Unload the warehouse and cluster from the node. - // 1. Used when a node is removed from the cluster. - // 2. For cluster_node_key: node_info, since its warehouse and cluster are already encoded in the key, we do not need to write the warehouse and cluster into its value again. - fn unload_warehouse_info(&self, node: &mut NodeInfo) { - node.cluster_id = String::new(); - node.warehouse_id = String::new(); - } - - async fn start_self_managed(&self, mut node: NodeInfo) -> Result { - if node.warehouse_id.is_empty() || node.cluster_id.is_empty() { + async fn start_self_managed(&self, node_with_wh: NodeInfo) -> Result { + if node_with_wh.warehouse_id.is_empty() || node_with_wh.cluster_id.is_empty() { return Err(ErrorCode::InvalidWarehouse( "The warehouse_id and cluster_id for self managed node must not be empty.", )); } - let node_with_wh = node.clone(); - self.unload_warehouse_info(&mut node); - let node_without_wh = node; + let node_without_wh = node_with_wh.unload_warehouse_info(); let node_key = self.node_key(&node_with_wh)?; let cluster_node_key = self.cluster_node_key(&node_with_wh)?; @@ -295,10 +286,9 @@ impl WarehouseMgr { ))) } - async fn heartbeat_self_managed(&self, mut node: NodeInfo, seq: u64) -> Result { + async fn heartbeat_self_managed(&self, node: NodeInfo, seq: u64) -> Result { let node_with_wh = node.clone(); - self.unload_warehouse_info(&mut node); - let node_without_wh = node; + let node_without_wh = node.unload_warehouse_info(); let node_key = self.node_key(&node_with_wh)?; let cluster_node_key = self.cluster_node_key(&node_with_wh)?; @@ -557,7 +547,7 @@ impl WarehouseMgr { Ok(self.metastore.transaction(txn).await?) } - async fn update_system_managed_node(&self, mut node: NodeInfo, seq: u64) -> Result { + async fn update_system_managed_node(&self, node: NodeInfo, seq: u64) -> Result { let mut txn = TxnRequest::default(); let node_key = self.node_key(&node)?; @@ -574,7 +564,7 @@ impl WarehouseMgr { if node.assigned_warehouse() { let cluster_node_key = self.cluster_node_key(&node)?; - self.unload_warehouse_info(&mut node); + let node = node.unload_warehouse_info(); txn.if_then.push(TxnOp::put_with_ttl( cluster_node_key, serde_json::to_vec(&node)?, @@ -587,7 +577,7 @@ impl WarehouseMgr { Ok(self.metastore.transaction(txn).await?) } - async fn heartbeat_system_managed(&self, mut node: NodeInfo, seq: u64) -> Result { + async fn heartbeat_system_managed(&self, node: NodeInfo, seq: u64) -> Result { let mut txn = TxnRequest::default(); let node_key = self.node_key(&node)?; @@ -604,7 +594,7 @@ impl WarehouseMgr { if node.assigned_warehouse() { let cluster_node_key = self.cluster_node_key(&node)?; - self.unload_warehouse_info(&mut node); + let node = node.unload_warehouse_info(); txn.if_then.push(TxnOp::put_with_ttl( cluster_node_key, serde_json::to_vec(&node)?, @@ -718,7 +708,10 @@ impl WarehouseMgr { }; let node_key = format!("{}/{}", self.node_key_prefix, node); + // Fetch the seq of node under `node_prefix/**` after_txn.if_then.push(TxnOp::get(node_key)); + // While we always assert the seq of `node_key_prefix/**` nodes during updates to ensure a consistent snapshot view, + // the seq of `cluster_node_key_prefix/**` nodes is also asserted as a defensive measure. cluster_node_seq.push(value.seq); } @@ -948,7 +941,7 @@ impl WarehouseMgr { impl WarehouseApi for WarehouseMgr { #[async_backtrace::framed] #[fastrace::trace] - async fn start_node(&self, node: NodeInfo) -> Result<(u64, NodeInfo)> { + async fn start_node(&self, node: NodeInfo) -> Result> { let start_reply = match node.node_type { NodeType::SelfManaged => self.start_self_managed(node.clone()).await?, NodeType::SystemManaged => self.start_system_managed(node.clone()).await?, @@ -967,7 +960,7 @@ impl WarehouseApi for WarehouseMgr { Some(Response::Get(TxnGetResponse { value: Some(seq_v), .. })), - }) => Ok((seq_v.seq, serde_json::from_slice(&seq_v.data)?)), + }) => Ok(SeqV::new(seq_v.seq, serde_json::from_slice(&seq_v.data)?)), _ => Err(ErrorCode::MetaServiceError("Add node info failure.")), } } @@ -1102,7 +1095,7 @@ impl WarehouseApi for WarehouseMgr { let mut txn = TxnRequest::default(); - for (seq, mut node) in selected_nodes { + for (seq, node) in selected_nodes { let node_key = self.node_key(&node)?; txn.condition @@ -1115,7 +1108,7 @@ impl WarehouseApi for WarehouseMgr { let cluster_node_key = self.cluster_node_key(&node)?; - self.unload_warehouse_info(&mut node); + let node = node.unload_warehouse_info(); txn.condition .push(map_condition(&cluster_node_key, MatchSeq::Exact(0))); txn.if_then.push(TxnOp::put_with_ttl( @@ -1246,7 +1239,7 @@ impl WarehouseApi for WarehouseMgr { Some(self.lift_time * 4), )); - self.unload_warehouse_info(&mut node); + let node = node.unload_warehouse_info(); resume_txn .condition .push(map_condition(&cluster_node_key, MatchSeq::Exact(0))); @@ -1321,7 +1314,7 @@ impl WarehouseApi for WarehouseMgr { serde_json::to_vec(&warehouse_snapshot.warehouse_info)?, )); - for mut node_snapshot in warehouse_snapshot.snapshot_nodes { + for node_snapshot in warehouse_snapshot.snapshot_nodes { let node_key = self.node_key(&node_snapshot.node_info)?; let cluster_node_key = self.cluster_node_key(&node_snapshot.node_info)?; @@ -1335,10 +1328,10 @@ impl WarehouseApi for WarehouseMgr { )); suspend_txn.if_then.push(TxnOp::delete(cluster_node_key)); - self.unload_warehouse_info(&mut node_snapshot.node_info); + let node = node_snapshot.node_info.unload_warehouse_info(); suspend_txn.if_then.push(TxnOp::put_with_ttl( node_key, - serde_json::to_vec(&node_snapshot.node_info)?, + serde_json::to_vec(&node)?, Some(self.lift_time * 4), )); } @@ -1441,11 +1434,12 @@ impl WarehouseApi for WarehouseMgr { serde_json::to_vec(&node_snapshot.node_info)?, Some(self.lift_time * 4), )); - self.unload_warehouse_info(&mut node_snapshot.node_info); + + let node = node_snapshot.node_info.unload_warehouse_info(); rename_txn.if_then.push(TxnOp::delete(old_cluster_node_key)); rename_txn.if_then.push(TxnOp::put_with_ttl( new_cluster_node_key.clone(), - serde_json::to_vec(&node_snapshot.node_info)?, + serde_json::to_vec(&node)?, Some(self.lift_time * 4), )); } @@ -1572,7 +1566,7 @@ impl WarehouseApi for WarehouseMgr { )); } - for (seq, mut node) in selected_nodes { + for (seq, node) in selected_nodes { let node_key = self.node_key(&node)?; let cluster_node_key = self.cluster_node_key(&node)?; @@ -1585,7 +1579,7 @@ impl WarehouseApi for WarehouseMgr { Some(self.lift_time * 4), )); - self.unload_warehouse_info(&mut node); + let node = node.unload_warehouse_info(); create_cluster_txn .condition .push(map_condition(&cluster_node_key, MatchSeq::Exact(0))); @@ -1896,7 +1890,7 @@ impl WarehouseApi for WarehouseMgr { } for selected_nodes in selected_nodes.into_values() { - for (seq, mut node) in selected_nodes { + for (seq, node) in selected_nodes { let node_key = self.node_key(&node)?; let cluster_node_key = self.cluster_node_key(&node)?; @@ -1909,7 +1903,7 @@ impl WarehouseApi for WarehouseMgr { Some(self.lift_time * 4), )); - self.unload_warehouse_info(&mut node); + let node = node.unload_warehouse_info(); add_cluster_node_txn .condition .push(map_condition(&cluster_node_key, MatchSeq::Exact(0))); @@ -2003,7 +1997,7 @@ impl WarehouseApi for WarehouseMgr { )); // lock all cluster state - for mut node_snapshot in warehouse_snapshot.snapshot_nodes { + for node_snapshot in warehouse_snapshot.snapshot_nodes { let node_key = self.node_key(&node_snapshot.node_info)?; let cluster_node_key = self.cluster_node_key(&node_snapshot.node_info)?; @@ -2020,15 +2014,14 @@ impl WarehouseApi for WarehouseMgr { if let Some(remove_node) = v.pop() { let SelectedNode::Random(node_group) = remove_node; if node_snapshot.node_info.runtime_node_group == node_group { - self.unload_warehouse_info(&mut node_snapshot.node_info); - node_snapshot.node_info.runtime_node_group = None; + let node = node_snapshot.node_info.leave_warehouse(); drop_cluster_node_txn .if_then .push(TxnOp::delete(cluster_node_key)); drop_cluster_node_txn.if_then.push(TxnOp::put_with_ttl( node_key, - serde_json::to_vec(&node_snapshot.node_info)?, + serde_json::to_vec(&node)?, Some(self.lift_time * 4), )) } @@ -2088,6 +2081,9 @@ impl WarehouseApi for WarehouseMgr { Ok(online_nodes) } + // Discovers and returns all the nodes that are in the same cluster with the given node, including the provided one. + // + // If the given node is not in a cluster yet, just return itself. async fn discover(&self, node_id: &str) -> Result> { let node_key = format!("{}/{}", self.node_key_prefix, escape_for_key(node_id)?); diff --git a/src/query/management/tests/it/warehouse.rs b/src/query/management/tests/it/warehouse.rs index 823ead83c29ac..af87241bb101d 100644 --- a/src/query/management/tests/it/warehouse.rs +++ b/src/query/management/tests/it/warehouse.rs @@ -215,7 +215,7 @@ async fn test_successfully_heartbeat_self_managed_node() -> Result<()> { let (kv, warehouse_manager, _nodes) = nodes(Duration::from_mins(60), 0).await?; let mut node_info = self_managed_node("test_node"); - let (seq, _node) = warehouse_manager.start_node(node_info.clone()).await?; + let node = warehouse_manager.start_node(node_info.clone()).await?; let info_key = "__fd_clusters_v6/test%2dtenant%2did/online_nodes/test_node"; assert_key_value(&kv, info_key, serde_json::to_vec(&node_info)?).await; @@ -234,13 +234,13 @@ async fn test_successfully_heartbeat_self_managed_node() -> Result<()> { assert_key_expire(&kv, warehouse_info_key, Duration::from_mins(50)).await; warehouse_manager - .heartbeat_node(&mut node_info, seq) + .heartbeat_node(&mut node_info, node.seq) .await?; assert_key_value(&kv, warehouse_info_key, info.clone()).await; assert_key_value(&kv, info_key, serde_json::to_vec(&node_info)?).await; assert_key_value(&kv, warehouse_key, serde_json::to_vec(&warehouse_node)?).await; - assert_key_seq(&kv, info_key, MatchSeq::GE(seq + 3)).await; - assert_key_seq(&kv, warehouse_key, MatchSeq::GE(seq + 3)).await; + assert_key_seq(&kv, info_key, MatchSeq::GE(node.seq + 3)).await; + assert_key_seq(&kv, warehouse_key, MatchSeq::GE(node.seq + 3)).await; assert_key_expire(&kv, info_key, Duration::from_mins(50)).await; assert_key_expire(&kv, warehouse_key, Duration::from_mins(50)).await; assert_key_expire(&kv, warehouse_info_key, Duration::from_mins(50)).await; @@ -653,8 +653,8 @@ async fn test_concurrent_recovery_create_warehouse() -> Result<()> { let node_id = GlobalUniqName::unique(); let start_node = warehouse_manager.start_node(system_managed_node(&node_id)); - let (_, node_info) = start_node.await.unwrap(); - node_info.id + let seq_node = start_node.await.unwrap(); + seq_node.id.clone() } })); } diff --git a/src/query/service/src/clusters/cluster.rs b/src/query/service/src/clusters/cluster.rs index 023dd7fca1404..563dfa98fb092 100644 --- a/src/query/service/src/clusters/cluster.rs +++ b/src/query/service/src/clusters/cluster.rs @@ -47,6 +47,8 @@ use databend_common_management::WarehouseMgr; use databend_common_meta_store::MetaStore; use databend_common_meta_store::MetaStoreProvider; use databend_common_meta_types::NodeInfo; +use databend_common_meta_types::SeqV; +use databend_common_meta_types::SeqValue; use databend_common_metrics::cluster::*; use databend_enterprise_resources_management::ResourcesManagement; use futures::future::select; @@ -507,14 +509,16 @@ impl ClusterDiscovery { self.drop_invalid_nodes(&node_info).await?; match self.warehouse_manager.start_node(node_info).await { - Ok((seq, node_info)) => self.start_heartbeat(node_info, seq).await, + Ok(seq_node) => self.start_heartbeat(seq_node).await, Err(cause) => Err(cause.add_message_back("(while cluster api add_node).")), } } #[async_backtrace::framed] - async fn start_heartbeat(self: &Arc, node_info: NodeInfo, seq: u64) -> Result<()> { + async fn start_heartbeat(self: &Arc, seq_node: SeqV) -> Result<()> { let mut heartbeat = self.heartbeat.lock().await; + let seq = seq_node.seq; + let node_info = seq_node.into_value().unwrap(); heartbeat.start(node_info, seq); Ok(()) }