Add actions to list nodes and retrieve node metrics

crates/modelardb_embedded/Cargo.toml

@@ -35,6 +35,7 @@ modelardb_auth.workspace = true
 modelardb_compression.workspace = true
 modelardb_storage.workspace = true
 modelardb_types.workspace = true
+prost.workspace = true
 tokio.workspace = true
 tonic.workspace = true
 

crates/modelardb_embedded/src/error.rs

@@ -29,6 +29,7 @@ use deltalake::{DeltaTableError, ObjectStoreError};
 use modelardb_compression::error::ModelarDbCompressionError;
 use modelardb_storage::error::ModelarDbStorageError;
 use modelardb_types::error::ModelarDbTypesError;
+use prost::DecodeError;
 use tonic::Status as TonicStatusError;
 use tonic::transport::Error as TonicTransportError;
 
@@ -58,6 +59,8 @@ pub enum ModelarDbEmbeddedError {
     ObjectStore(ObjectStoreError),
     /// Error returned by Apache Parquet.
     Parquet(ParquetError),
+    /// Error returned by Prost when decoding a message that is not valid.
+    ProstDecode(DecodeError),
     /// Status returned by Tonic.
     TonicStatus(Box<TonicStatusError>),
     /// Error returned by Tonic.
@@ -83,6 +86,7 @@ impl Display for ModelarDbEmbeddedError {
             Self::ModelarDbTypes(reason) => write!(f, "ModelarDB Types Error: {reason}"),
             Self::ObjectStore(reason) => write!(f, "Object Store Error: {reason}"),
             Self::Parquet(reason) => write!(f, "Parquet Error: {reason}"),
+            Self::ProstDecode(reason) => write!(f, "Prost Decode Error: {reason}"),
             Self::TonicStatus(reason) => write!(f, "Tonic Status Error: {reason}"),
             Self::TonicTransport(reason) => write!(f, "Tonic Transport Error: {reason}"),
             Self::Unimplemented(reason) => write!(f, "Unimplemented Error: {reason}"),
@@ -104,6 +108,7 @@ impl Error for ModelarDbEmbeddedError {
             Self::ModelarDbTypes(reason) => Some(reason),
             Self::ObjectStore(reason) => Some(reason),
             Self::Parquet(reason) => Some(reason),
+            Self::ProstDecode(reason) => Some(reason),
             Self::TonicStatus(reason) => Some(reason),
             Self::TonicTransport(reason) => Some(reason),
             Self::Unimplemented(_reason) => None,
@@ -166,6 +171,12 @@ impl From<ParquetError> for ModelarDbEmbeddedError {
     }
 }
 
+impl From<DecodeError> for ModelarDbEmbeddedError {
+    fn from(error: DecodeError) -> Self {
+        Self::ProstDecode(error)
+    }
+}
+
 impl From<TonicStatusError> for ModelarDbEmbeddedError {
     fn from(error: TonicStatusError) -> Self {
         Self::TonicStatus(Box::new(error))

crates/modelardb_embedded/src/operations/client.rs

@@ -28,16 +28,20 @@ use arrow::record_batch::RecordBatch;
 use arrow_flight::decode::FlightRecordBatchStream;
 use arrow_flight::encode::FlightDataEncoderBuilder;
 use arrow_flight::flight_service_client::FlightServiceClient;
-use arrow_flight::{Action, Criteria, FlightDescriptor, Ticket};
+use arrow_flight::{Action, Criteria, FlightDescriptor, Result as FlightResult, Ticket};
 use async_trait::async_trait;
 use datafusion::error::DataFusionError;
 use datafusion::execution::RecordBatchStream;
 use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
 use futures::{StreamExt, TryStreamExt, stream};
 use modelardb_auth::BearerInterceptor;
+use modelardb_types::flight::protocol;
+use modelardb_types::types::Node;
+use prost::Message;
+use prost::bytes::Bytes;
 use tonic::codegen::InterceptedService;
 use tonic::transport::{Channel, Endpoint};
-use tonic::{Request, Status};
+use tonic::{Request, Status, Streaming};
 
 use crate::error::{ModelarDbEmbeddedError, Result};
 use crate::operations::{
@@ -65,6 +69,134 @@ impl Client {
 
         Ok(Client { flight_client })
     }
+
+    /// Returns the current configuration of the node. If the configuration could not be retrieved,
+    /// [`ModelarDbEmbeddedError`] is returned.
+    pub async fn configuration(&mut self) -> Result<protocol::Configuration> {
+        let bytes = self.retrieve_action_bytes("GetConfiguration").await?;
+
+        Ok(protocol::Configuration::decode(bytes)?)
+    }
+
+    /// Updates `setting` in the node configuration to `new_value`. If the setting could not be
+    /// updated, [`ModelarDbEmbeddedError`] is returned.
+    pub async fn update_configuration(
+        &mut self,
+        setting: protocol::update_configuration::Setting,
+        new_value: Option<u64>,
+    ) -> Result<()> {
+        let update_configuration = protocol::UpdateConfiguration {
+            setting: setting as i32,
+            new_value,
+        };
+
+        self.send_action("UpdateConfiguration", update_configuration.encode_to_vec())
+            .await?;
+
+        Ok(())
+    }
+
+    /// Flushes all data in memory to disk. If the data could not be flushed,
+    /// [`ModelarDbEmbeddedError`] is returned.
+    pub async fn flush_memory(&mut self) -> Result<()> {
+        self.send_action("FlushMemory", vec![]).await?;
+
+        Ok(())
+    }
+
+    /// Flushes all data in memory to disk and then transfers all compressed data to the remote
+    /// object store. If the data could not be flushed, [`ModelarDbEmbeddedError`] is returned.
+    pub async fn flush_node(&mut self) -> Result<()> {
+        self.send_action("FlushNode", vec![]).await?;
+
+        Ok(())
+    }
+
+    /// Flushes all data to disk, transfers it to the remote object store, removes the node from the
+    /// cluster if necessary, and kills the node process. Since the process is killed, a
+    /// conventional response cannot be returned, so a dropped connection is not treated as an
+    /// error. If the data could not be flushed before the node was killed,
+    /// [`ModelarDbEmbeddedError`] is returned.
+    pub async fn kill_node(&mut self) -> Result<()> {
+        // The node exits while handling this action, so the response stream is dropped.
+        let _ = self.send_action("KillNode", vec![]).await;
+
+        Ok(())
+    }
+
+    /// Returns the nodes that are currently part of the cluster. A single node returns only itself.
+    /// If the nodes could not be retrieved, [`ModelarDbEmbeddedError`] is returned.
+    pub async fn list_nodes(&mut self) -> Result<Vec<Node>> {
+        let bytes = self.retrieve_action_bytes("ListNodes").await?;
+
+        Ok(modelardb_types::flight::deserialize_and_extract_cluster_nodes(&bytes)?)
+    }
+
+    /// Returns the current resource usage metrics of the node. If the metrics could not be
+    /// retrieved, [`ModelarDbEmbeddedError`] is returned.
+    pub async fn node_metrics(&mut self) -> Result<protocol::NodeMetrics> {
+        let bytes = self.retrieve_action_bytes("NodeMetrics").await?;
+
+        Ok(protocol::NodeMetrics::decode(bytes)?)
+    }
+
+    /// Sends the action with the type `action_type` and an empty body to the node and returns the
+    /// body of the response. If the action could not be performed, [`ModelarDbEmbeddedError`] is
+    /// returned.
+    async fn retrieve_action_bytes(&mut self, action_type: &str) -> Result<Bytes> {
+        let mut response = self.send_action(action_type, vec![]).await?;
+
+        let message = response.message().await?.ok_or_else(|| {
+            ModelarDbEmbeddedError::from(Status::internal(format!(
+                "Action '{action_type}' did not return a response message."
+            )))
+        })?;
+
+        Ok(message.body)
+    }
+
+    /// Sends the action with the type `action_type` and `body` to the node and returns the response
+    /// stream. If the action could not be performed, [`ModelarDbEmbeddedError`] is returned.
+    async fn send_action(
+        &mut self,
+        action_type: &str,
+        body: Vec<u8>,
+    ) -> Result<Streaming<FlightResult>> {
+        let action = Action {
+            r#type: action_type.to_owned(),
+            body: body.into(),
+        };
+
+        let response = self.flight_client.do_action(Request::new(action)).await?;
+
+        Ok(response.into_inner())
+    }
+
+    /// Returns the URL of the cloud node that the node assigns to execute the SQL in `sql`. If the
+    /// node is not running in a cluster, or a cloud node could not be assigned,
+    /// [`ModelarDbEmbeddedError`] is returned.
+    pub async fn cloud_query_node(&mut self, sql: &str) -> Result<String> {
+        let flight_descriptor = FlightDescriptor::new_cmd(sql.to_owned());
+        let flight_info = self
+            .flight_client
+            .get_flight_info(Request::new(flight_descriptor))
+            .await?
+            .into_inner();
+
+        let endpoint = flight_info.endpoint.into_iter().next().ok_or_else(|| {
+            ModelarDbEmbeddedError::InvalidArgument(
+                "The node did not return an endpoint for the query.".to_owned(),
+            )
+        })?;
+
+        let location = endpoint.location.into_iter().next().ok_or_else(|| {
+            ModelarDbEmbeddedError::InvalidArgument(
+                "The endpoint did not return a cloud node location for the query.".to_owned(),
+            )
+        })?;
+
+        Ok(location.uri)
+    }
 }
 
 #[async_trait]
@@ -76,21 +208,9 @@ impl Operations for Client {
 
     /// Returns the type of the ModelarDB node that the client is connected to.
     async fn modelardb_type(&mut self) -> Result<ModelarDBType> {
-        // Retrieve the node type from the ModelarDB node.
-        let action = Action {
-            r#type: "NodeType".to_owned(),
-            body: vec![].into(),
-        };
+        let bytes = self.retrieve_action_bytes("NodeType").await?;
 
-        let response = self.flight_client.do_action(Request::new(action)).await?;
-
-        let message = response
-            .into_inner()
-            .message()
-            .await?
-            .expect("Flight message should exist.");
-
-        ModelarDBType::from_str(str::from_utf8(&message.body)?)
+        ModelarDBType::from_str(str::from_utf8(&bytes)?)
     }
 
     /// Creates a table with the name in `table_name` and the information in `table_type`. If the
@@ -118,12 +238,7 @@ impl Operations for Client {
             }
         };
 
-        let action = Action {
-            r#type: "CreateTable".to_owned(),
-            body: protobuf_bytes.into(),
-        };
-
-        self.flight_client.do_action(action).await?;
+        self.send_action("CreateTable", protobuf_bytes).await?;
 
         Ok(())
     }

crates/modelardb_server/Cargo.toml

@@ -46,6 +46,7 @@ rand.workspace = true
 serde.workspace = true
 snmalloc-rs = { workspace = true, features = ["build_cc"] }
 sqlparser.workspace = true
+sysinfo.workspace = true
 tokio = { workspace = true, features = ["rt-multi-thread", "signal"] }
 tokio-stream.workspace = true
 toml.workspace = true

crates/modelardb_server/src/cluster.rs

@@ -36,6 +36,26 @@ use tonic::transport::Endpoint;
 use crate::context::Context;
 use crate::error::{ModelarDbServerError, Result};
 
+/// The different possible modes that a ModelarDB server can be deployed in, assigned when the
+/// server is started.
+#[derive(Clone)]
+pub(crate) enum ClusterMode {
+    SingleNode(Node),
+    MultiNode(Cluster),
+}
+
+impl ClusterMode {
+    /// Return all nodes in the deployment. A single node returns only itself, while a node in a
+    /// cluster returns every node in the cluster. If the nodes could not be retrieved, return
+    /// [`ModelarDbServerError`].
+    pub(crate) async fn nodes(&self) -> Result<Vec<Node>> {
+        match self {
+            ClusterMode::SingleNode(node) => Ok(vec![node.clone()]),
+            ClusterMode::MultiNode(cluster) => cluster.nodes().await,
+        }
+    }
+}
+
 /// Stores the node that represents the local system and allows for performing operations that need
 /// to be applied to every peer node in the cluster.
 #[derive(Clone)]
@@ -357,10 +377,19 @@ impl Cluster {
         Ok(url.to_owned())
     }
 
+    /// Return all nodes currently in the cluster. If the nodes could not be retrieved, return
+    /// [`ModelarDbServerError`].
+    pub(crate) async fn nodes(&self) -> Result<Vec<Node>> {
+        self.remote_data_folder
+            .nodes()
+            .await
+            .map_err(|error| error.into())
+    }
+
     /// Return all nodes in the cluster except the node that was saved when the [`Cluster`] was
     /// created. If the nodes could not be retrieved, return [`ModelarDbServerError`].
     async fn peer_nodes(&self) -> Result<Vec<Node>> {
-        let nodes = self.remote_data_folder.nodes().await?;
+        let nodes = self.nodes().await?;
 
         Ok(nodes
             .into_iter()
@@ -485,8 +514,8 @@ mod test {
     use modelardb_types::types::{ArrowTimestamp, ArrowValue, ErrorBound, ServerMode};
     use tempfile::TempDir;
 
+    use crate::ServerArgs;
     use crate::data_folders::DataFolders;
-    use crate::{ClusterMode, ServerArgs};
 
     // Tests for Cluster.
     #[tokio::test]
@@ -721,7 +750,7 @@ mod test {
                     Some(cluster.remote_data_folder.clone()),
                     local_data_folder,
                 ),
-                ClusterMode::MultiNode(Box::new(cluster)),
+                ClusterMode::MultiNode(cluster),
                 &ServerArgs::parse_from(["modelardbd", "edge", "data", "s3://bucket"]),
             )
             .await

crates/modelardb_server/src/configuration.rs

@@ -34,9 +34,10 @@ use prost::Message;
 use serde::{Deserialize, Serialize};
 use tokio::sync::RwLock;
 
+use crate::ServerArgs;
+use crate::cluster::ClusterMode;
 use crate::error::{ModelarDbServerError, Result};
 use crate::storage::StorageEngine;
-use crate::{ClusterMode, ServerArgs};
 
 const CONFIGURATION_FILE_NAME: &str = "modelardbd.toml";
 
@@ -602,9 +603,10 @@ mod tests {
             .await
             .unwrap();
 
+        let node = Node::new("edge".to_owned(), ServerMode::Edge);
         let result = ConfigurationManager::try_new(
             local_data_folder,
-            ClusterMode::SingleNode,
+            ClusterMode::SingleNode(node),
             &default_args(),
         )
         .await;
@@ -625,9 +627,10 @@ mod tests {
         let path = temp_dir.path().join(CONFIGURATION_FILE_NAME);
         std::fs::write(path, "invalid_toml").unwrap();
 
+        let node = Node::new("edge".to_owned(), ServerMode::Edge);
         let result = ConfigurationManager::try_new(
             local_data_folder,
-            ClusterMode::SingleNode,
+            ClusterMode::SingleNode(node),
             &default_args(),
         )
         .await;
@@ -1020,7 +1023,7 @@ mod tests {
         let configuration_manager = Arc::new(RwLock::new(
             ConfigurationManager::try_new(
                 local_data_folder.clone(),
-                ClusterMode::MultiNode(Box::new(cluster)),
+                ClusterMode::MultiNode(cluster),
                 &default_args(),
             )
             .await