add admin-shell subcommand

Author: MarinPostma

libsql-server/src/admin_shell.rs

@@ -2,3 +2,233 @@ use std::fmt::Display;
 use std::pin::Pin;
 
 use bytes::Bytes;
+use dialoguer::BasicHistory;
+use rusqlite::types::ValueRef;
+use tokio_stream::{Stream, StreamExt as _};
+use tonic::metadata::BinaryMetadataValue;
+
+use crate::connection::Connection as _;
+use crate::database::Connection;
+use crate::namespace::{NamespaceName, NamespaceStore};
+
+use self::rpc::admin_shell_service_server::{AdminShellService, AdminShellServiceServer};
+use self::rpc::response::Resp;
+use self::rpc::Null;
+
+mod rpc {
+    #![allow(clippy::all)]
+    include!("generated/admin_shell.rs");
+}
+
+pub(crate) fn make_svc(namespace_store: NamespaceStore) -> AdminShellServiceServer<AdminShell> {
+    let admin_shell = AdminShell::new(namespace_store);
+    rpc::admin_shell_service_server::AdminShellServiceServer::new(admin_shell)
+}
+
+pub(super) struct AdminShell {
+    namespace_store: NamespaceStore,
+}
+
+impl AdminShell {
+    fn new(namespace_store: NamespaceStore) -> Self {
+        Self { namespace_store }
+    }
+
+    async fn with_namespace(
+        &self,
+        ns: Bytes,
+        queries: impl Stream<Item = Result<rpc::Query, tonic::Status>>,
+    ) -> anyhow::Result<impl Stream<Item = Result<rpc::Response, tonic::Status>>> {
+        let namespace = NamespaceName::from_bytes(ns).unwrap();
+        let connection_maker = self
+            .namespace_store
+            .with(namespace, |ns| ns.db.connection_maker())
+            .await?;
+        let connection = connection_maker.create().await?;
+        Ok(run_shell(connection, queries))
+    }
+}
+
+fn run_shell(
+    conn: Connection,
+    queries: impl Stream<Item = Result<rpc::Query, tonic::Status>>,
+) -> impl Stream<Item = Result<rpc::Response, tonic::Status>> {
+    async_stream::stream! {
+        tokio::pin!(queries);
+        while let Some(q) = queries.next().await {
+            let Ok(q) = q else { break };
+            let res = tokio::task::block_in_place(|| {
+                conn.with_raw(move |conn| {
+                    run_one(conn, q.query)
+                })
+            });
+
+            yield res
+        }
+    }
+}
+
+fn run_one(conn: &mut rusqlite::Connection, q: String) -> Result<rpc::Response, tonic::Status> {
+    match try_run_one(conn, q) {
+        Ok(resp) => Ok(resp),
+        Err(e) => Ok(rpc::Response {
+            resp: Some(Resp::Error(rpc::Error {
+                error: e.to_string(),
+            })),
+        }),
+    }
+}
+
+fn try_run_one(conn: &mut rusqlite::Connection, q: String) -> anyhow::Result<rpc::Response> {
+    let mut stmt = conn.prepare(&q)?;
+    let col_count = stmt.column_count();
+    let mut rows = stmt.query(())?;
+    let mut out_rows = Vec::new();
+    while let Some(row) = rows.next()? {
+        let mut out_row = Vec::with_capacity(col_count);
+        for i in 0..col_count {
+            let rpc_value = match row.get_ref(i).unwrap() {
+                ValueRef::Null => rpc::value::Value::Null(Null {}),
+                ValueRef::Integer(i) => rpc::value::Value::Integer(i),
+                ValueRef::Real(x) => rpc::value::Value::Real(x),
+                ValueRef::Text(s) => rpc::value::Value::Text(String::from_utf8(s.to_vec())?),
+                ValueRef::Blob(b) => rpc::value::Value::Blob(b.to_vec()),
+            };
+            out_row.push(rpc::Value {
+                value: Some(rpc_value),
+            });
+        }
+        out_rows.push(rpc::Row { values: out_row });
+    }
+
+    Ok(rpc::Response {
+        resp: Some(Resp::Rows(rpc::Rows { rows: out_rows })),
+    })
+}
+
+#[async_trait::async_trait]
+impl AdminShellService for AdminShell {
+    type ShellStream = Pin<Box<dyn Stream<Item = Result<rpc::Response, tonic::Status>> + Send>>;
+
+    async fn shell(
+        &self,
+        request: tonic::Request<tonic::Streaming<rpc::Query>>,
+    ) -> std::result::Result<tonic::Response<Self::ShellStream>, tonic::Status> {
+        let Some(namespace) = request.metadata().get_bin("x-namespace-bin") else {
+            return Err(tonic::Status::new(
+                tonic::Code::InvalidArgument,
+                "missing namespace",
+            ));
+        };
+        let Ok(ns_bytes) = namespace.to_bytes() else {
+            return Err(tonic::Status::new(
+                tonic::Code::InvalidArgument,
+                "bad namespace encoding",
+            ));
+        };
+
+        match self.with_namespace(ns_bytes, request.into_inner()).await {
+            Ok(s) => Ok(tonic::Response::new(Box::pin(s))),
+            Err(e) => Err(tonic::Status::new(
+                tonic::Code::FailedPrecondition,
+                e.to_string(),
+            )),
+        }
+    }
+}
+
+pub struct AdminShellClient {
+    remote_url: String,
+}
+
+impl AdminShellClient {
+    pub fn new(remote_url: String) -> Self {
+        Self { remote_url }
+    }
+
+    pub async fn run_namespace(&self, namespace: &str) -> anyhow::Result<()> {
+        let namespace = NamespaceName::from_string(namespace.to_string())?;
+        let mut client = rpc::admin_shell_service_client::AdminShellServiceClient::connect(
+            self.remote_url.clone(),
+        )
+        .await?;
+        let (sender, receiver) = tokio::sync::mpsc::channel(1);
+        let req_stream = tokio_stream::wrappers::ReceiverStream::new(receiver);
+
+        let mut req = tonic::Request::new(req_stream);
+        req.metadata_mut().insert_bin(
+            "x-namespace-bin",
+            BinaryMetadataValue::from_bytes(namespace.as_slice()),
+        );
+        let mut resp_stream = client.shell(req).await?.into_inner();
+
+        let mut history = BasicHistory::new();
+        loop {
+            // this is blocking, but the shell runs in it's own process with no other tasks, so
+            // that's ok
+            let prompt = dialoguer::Input::<String>::new()
+                .with_prompt("> ")
+                .history_with(&mut history)
+                .interact_text();
+
+            match prompt {
+                Ok(query) => {
+                    let q = rpc::Query { query };
+                    sender.send(q).await?;
+                    match resp_stream.next().await {
+                        Some(Ok(rpc::Response {
+                            resp: Some(rpc::response::Resp::Rows(rows)),
+                        })) => {
+                            println!("{}", RowsFormatter(rows));
+                        }
+                        Some(Ok(rpc::Response {
+                            resp: Some(rpc::response::Resp::Error(rpc::Error { error })),
+                        })) => {
+                            println!("query error: {error}");
+                        }
+                        Some(Err(e)) => {
+                            println!("rpc error: {}", e.message());
+                            break;
+                        }
+                        _ => break,
+                    }
+                }
+                Err(e) => println!("error: {e}"),
+            }
+        }
+
+        Ok(())
+    }
+}
+
+struct RowsFormatter(rpc::Rows);
+
+impl Display for RowsFormatter {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        for row in self.0.rows.iter() {
+            let mut is_first = true;
+            for value in row.values.iter() {
+                if !is_first {
+                    f.write_str(", ")?;
+                }
+                is_first = false;
+
+                match value.value.as_ref().unwrap() {
+                    rpc::value::Value::Null(_) => f.write_str("null")?,
+                    rpc::value::Value::Real(x) => write!(f, "{x}")?,
+                    rpc::value::Value::Integer(i) => write!(f, "{i}")?,
+                    rpc::value::Value::Text(s) => f.write_str(&s)?,
+                    rpc::value::Value::Blob(b) => {
+                        for x in b {
+                            write!(f, "{x:0x}")?
+                        }
+                    }
+                }
+            }
+
+            writeln!(f)?;
+        }
+
+        Ok(())
+    }
+}

libsql-server/src/main.rs

@@ -301,16 +301,18 @@ struct Cli {
         default_value = "8"
     )]
     sync_conccurency: usize,
+
+    #[clap(subcommand)]
+    subcommand: Option<UtilsSubcommands>,
 }
 
 #[derive(clap::Subcommand, Debug)]
 enum UtilsSubcommands {
-    Dump {
-        #[clap(long)]
-        /// Path at which to write the dump
-        path: Option<PathBuf>,
+    AdminShell {
+        #[clap(long, default_value = "http://127.0.0.1:9090")]
+        admin_api_url: String,
         #[clap(long)]
-        namespace: String,
+        namespace: Option<String>,
     },
 }
 
@@ -710,6 +712,25 @@ async fn build_server(config: &Cli) -> anyhow::Result<Server> {
 
 #[tokio::main]
 async fn main() -> Result<()> {
+    let args = Cli::parse();
+
+    if let Some(ref subcommand) = args.subcommand {
+        match subcommand {
+            UtilsSubcommands::AdminShell {
+                admin_api_url,
+                namespace,
+            } => {
+                let client =
+                    libsql_server::admin_shell::AdminShellClient::new(admin_api_url.clone());
+                if let Some(ns) = namespace {
+                    client.run_namespace(ns).await?;
+                }
+            }
+        }
+
+        return Ok(());
+    }
+
     if std::env::var("RUST_LOG").is_err() {
         std::env::set_var("RUST_LOG", "info");
     }
@@ -730,8 +751,6 @@ async fn main() -> Result<()> {
         )
         .init();
 
-    let args = Cli::parse();
-
     args.print_welcome_message();
     let server = build_server(&args).await?;
     server.start().await?;