Move client helpers to system-adapter-protocol

Author: Jeadie

Cargo.lock

@@ -10,5 +10,11 @@ repository.workspace = true
 serde = { workspace = true }
 serde_json = { workspace = true }
 uuid = { workspace = true, features = ["serde", "v4"] }
+tokio = { workspace = true, features = ["process", "io-util"] }
+reqwest = { workspace = true, optional = true }
+
+[features]
+default = ["client"]
+client = ["reqwest"]
 
 [dev-dependencies]

crates/system-adapter-protocol/Cargo.toml

@@ -0,0 +1,342 @@
+/*
+Copyright 2024-2025 The Spice.ai OSS Authors
+
+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
+
+     https://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.
+*/
+
+//! Client implementations for system adapter JSON-RPC communication.
+
+use crate::{methods, JsonRpcError, JsonRpcRequest, JsonRpcResponse};
+use serde::{de::DeserializeOwned, Serialize};
+use std::collections::HashMap;
+use tokio::{
+    io::{AsyncBufReadExt, AsyncWriteExt, BufReader},
+    process::{Child, ChildStdin, ChildStdout, Command},
+};
+
+/// Result type for client operations
+pub type Result<T> = std::result::Result<T, ClientError>;
+
+/// Errors that can occur during client operations
+#[derive(Debug)]
+pub enum ClientError {
+    /// JSON-RPC error returned by the server
+    JsonRpc(JsonRpcError),
+    /// I/O error during communication
+    Io(std::io::Error),
+    /// JSON serialization/deserialization error
+    Json(serde_json::Error),
+    /// HTTP transport error
+    #[cfg(feature = "client")]
+    Http(reqwest::Error),
+    /// Invalid response format
+    InvalidResponse(String),
+    /// Transport-specific error
+    Transport(String),
+}
+
+impl std::fmt::Display for ClientError {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            Self::JsonRpc(e) => write!(f, "JSON-RPC error: {}", e.message),
+            Self::Io(e) => write!(f, "I/O error: {e}"),
+            Self::Json(e) => write!(f, "JSON error: {e}"),
+            #[cfg(feature = "client")]
+            Self::Http(e) => write!(f, "HTTP error: {e}"),
+            Self::InvalidResponse(msg) => write!(f, "Invalid response: {msg}"),
+            Self::Transport(msg) => write!(f, "Transport error: {msg}"),
+        }
+    }
+}
+
+impl std::error::Error for ClientError {}
+
+impl From<std::io::Error> for ClientError {
+    fn from(e: std::io::Error) -> Self {
+        Self::Io(e)
+    }
+}
+
+impl From<serde_json::Error> for ClientError {
+    fn from(e: serde_json::Error) -> Self {
+        Self::Json(e)
+    }
+}
+
+#[cfg(feature = "client")]
+impl From<reqwest::Error> for ClientError {
+    fn from(e: reqwest::Error) -> Self {
+        Self::Http(e)
+    }
+}
+
+/// System adapter client for JSON-RPC communication
+pub enum Client {
+    /// Stdio transport - communicate via stdin/stdout with a child process
+    Stdio {
+        _child: Box<Child>,
+        stdin: ChildStdin,
+        stdout: BufReader<ChildStdout>,
+    },
+    /// HTTP transport - communicate via HTTP POST requests
+    #[cfg(feature = "client")]
+    Http {
+        client: reqwest::Client,
+        endpoint: String,
+    },
+}
+
+impl Client {
+    /// Create a client using stdio transport by spawning a command
+    pub fn stdio(
+        command: impl AsRef<str>,
+        args: Vec<String>,
+        env: HashMap<String, String>,
+    ) -> Result<Self> {
+        let command_str = command.as_ref();
+        let mut cmd = Command::new(command_str);
+
+        cmd.args(args);
+        for (key, value) in env {
+            cmd.env(key, value);
+        }
+
+        cmd.stdin(std::process::Stdio::piped())
+            .stdout(std::process::Stdio::piped())
+            .stderr(std::process::Stdio::inherit());
+
+        let mut child = cmd.spawn().map_err(|e| {
+            ClientError::Transport(format!(
+                "Failed to start stdio command '{command_str}': {e}"
+            ))
+        })?;
+
+        let stdin = child.stdin.take().ok_or_else(|| {
+            ClientError::Transport("Stdio child missing stdin".to_string())
+        })?;
+        let stdout = child.stdout.take().ok_or_else(|| {
+            ClientError::Transport("Stdio child missing stdout".to_string())
+        })?;
+
+        Ok(Self::Stdio {
+            _child: Box::new(child),
+            stdin,
+            stdout: BufReader::new(stdout),
+        })
+    }
+
+    /// Create a client using HTTP transport
+    #[cfg(feature = "client")]
+    pub fn http(endpoint: impl Into<String>) -> Self {
+        Self::Http {
+            client: reqwest::Client::new(),
+            endpoint: endpoint.into(),
+        }
+    }
+
+    /// Get the transport name
+    pub fn transport_name(&self) -> &'static str {
+        match self {
+            Self::Stdio { .. } => "stdio",
+            #[cfg(feature = "client")]
+            Self::Http { .. } => "http",
+        }
+    }
+
+    /// Query available RPC methods from the server
+    pub async fn rpc_methods(&mut self) -> Result<Vec<String>> {
+        let request = JsonRpcRequest::new(1, methods::RPC_METHODS, serde_json::json!({}));
+        let response: JsonRpcResponse<serde_json::Value> = self.call_typed(request).await?;
+
+        let methods = response
+            .result
+            .as_ref()
+            .and_then(|v| v.get("methods"))
+            .and_then(|v| v.as_array())
+            .ok_or_else(|| {
+                ClientError::InvalidResponse(
+                    "Response missing result.methods array".to_string(),
+                )
+            })?
+            .iter()
+            .filter_map(|v| v.as_str().map(ToString::to_string))
+            .collect();
+
+        Ok(methods)
+    }
+
+    /// Make a typed JSON-RPC call with request and response types
+    pub async fn call_typed<Req: Serialize, Resp: DeserializeOwned>(
+        &mut self,
+        request: JsonRpcRequest<Req>,
+    ) -> Result<JsonRpcResponse<Resp>> {
+        let request_value = serde_json::to_value(request)?;
+        let response_value = self.call_raw(request_value).await?;
+        let response: JsonRpcResponse<Resp> = serde_json::from_value(response_value)?;
+
+        if let Some(error) = response.error {
+            return Err(ClientError::JsonRpc(error));
+        }
+
+        Ok(response)
+    }
+
+    /// Make a raw JSON-RPC call with serde_json::Value
+    pub async fn call_raw(&mut self, request: serde_json::Value) -> Result<serde_json::Value> {
+        match self {
+            Self::Stdio {
+                _child: _,
+                stdin,
+                stdout,
+            } => {
+                let payload = serde_json::to_string(&request)?;
+                stdin.write_all(payload.as_bytes()).await?;
+                stdin.write_all(b"\n").await?;
+                stdin.flush().await?;
+
+                let mut line = String::new();
+                let read = stdout.read_line(&mut line).await?;
+                if read == 0 {
+                    return Err(ClientError::Transport(
+                        "Stdio process closed stdout before responding".to_string(),
+                    ));
+                }
+
+                let response: serde_json::Value = serde_json::from_str(line.trim_end())?;
+                if let Some(error) = response.get("error") {
+                    let error: JsonRpcError = serde_json::from_value(error.clone())?;
+                    return Err(ClientError::JsonRpc(error));
+                }
+                Ok(response)
+            }
+            #[cfg(feature = "client")]
+            Self::Http { client, endpoint } => {
+                let response = client
+                    .post(endpoint.as_str())
+                    .json(&request)
+                    .send()
+                    .await
+                    .map_err(|e| {
+                        ClientError::Transport(format!(
+                            "Failed to POST to {endpoint}: {e}"
+                        ))
+                    })?;
+
+                let status = response.status();
+                let value: serde_json::Value = response
+                    .json()
+                    .await
+                    .map_err(|e| {
+                        ClientError::Transport(format!(
+                            "Failed to parse response body (status {status}): {e}"
+                        ))
+                    })?;
+
+                if let Some(error) = value.get("error") {
+                    let error: JsonRpcError = serde_json::from_value(error.clone())?;
+                    return Err(ClientError::JsonRpc(error));
+                }
+                Ok(value)
+            }
+        }
+    }
+}
+
+/// Builder for creating a `Client` with various configuration options
+pub struct ClientBuilder {
+    transport: TransportConfig,
+}
+
+enum TransportConfig {
+    Stdio {
+        command: String,
+        args: Vec<String>,
+        env: HashMap<String, String>,
+    },
+    #[cfg(feature = "client")]
+    Http {
+        endpoint: String,
+    },
+}
+
+impl ClientBuilder {
+    /// Create a builder for stdio transport
+    pub fn stdio(command: impl Into<String>) -> Self {
+        Self {
+            transport: TransportConfig::Stdio {
+                command: command.into(),
+                args: Vec::new(),
+                env: HashMap::new(),
+            },
+        }
+    }
+
+    /// Create a builder for HTTP transport
+    #[cfg(feature = "client")]
+    pub fn http(endpoint: impl Into<String>) -> Self {
+        Self {
+            transport: TransportConfig::Http {
+                endpoint: endpoint.into(),
+            },
+        }
+    }
+
+    /// Add command-line arguments (stdio only)
+    pub fn with_args(mut self, args: Vec<String>) -> Self {
+        if let TransportConfig::Stdio { args: ref mut a, .. } = self.transport {
+            *a = args;
+        }
+        self
+    }
+
+    /// Add environment variables (stdio only)
+    pub fn with_env(mut self, env: HashMap<String, String>) -> Self {
+        if let TransportConfig::Stdio { env: ref mut e, .. } = self.transport {
+            *e = env;
+        }
+        self
+    }
+
+    /// Build the client
+    pub fn build(self) -> Result<Client> {
+        match self.transport {
+            TransportConfig::Stdio { command, args, env } => Client::stdio(command, args, env),
+            #[cfg(feature = "client")]
+            TransportConfig::Http { endpoint } => Ok(Client::http(endpoint)),
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_client_error_display() {
+        let error = ClientError::Transport("test error".to_string());
+        assert_eq!(format!("{error}"), "Transport error: test error");
+    }
+
+    #[test]
+    fn test_builder_stdio() {
+        let builder = ClientBuilder::stdio("python");
+        assert!(matches!(builder.transport, TransportConfig::Stdio { .. }));
+    }
+
+    #[cfg(feature = "client")]
+    #[test]
+    fn test_builder_http() {
+        let builder = ClientBuilder::http("http://localhost:8080");
+        assert!(matches!(builder.transport, TransportConfig::Http { .. }));
+    }
+}