unitycatalog_storage_proxy/cli/
mod.rs

1//! CLI surface for the standalone `storage-proxy` binary.
2//!
3//! Two subcommands, mirroring the deployable half of the `server` crate's CLI
4//! (the proxy has no database, so there is no `migrate`):
5//!   - `serve` — run the byte-proxy (see [`serve::serve`]). Flags overlay the
6//!     config file, CLI flags taking precedence. A config file is optional when
7//!     the required upstream URL is supplied via `--upstream-url`.
8//!   - `healthcheck` — probe the configured `/health` endpoint and exit 0
9//!     (healthy) or non-zero (unhealthy). This is the probe the distroless
10//!     image's Docker `HEALTHCHECK` runs, since distroless has no shell/`curl`.
11
12pub mod config;
13pub mod serve;
14
15use std::time::Duration;
16
17use clap::{Args, Parser, Subcommand};
18
19use crate::cli::config::{AuthConfig, Config, DEFAULT_HOST, DEFAULT_PORT, UpstreamConfig};
20
21/// `storage-proxy` — standalone Unity Catalog storage byte-proxy.
22#[derive(Debug, Parser)]
23#[command(name = "storage-proxy", version, about, long_about = None)]
24pub struct Cli {
25    #[command(subcommand)]
26    pub command: Command,
27}
28
29#[derive(Debug, Subcommand)]
30pub enum Command {
31    /// Run the byte-proxy.
32    Serve(ServeArgs),
33    /// Probe the configured `/health` endpoint; exit 0 if healthy, non-zero otherwise.
34    Healthcheck(HealthcheckArgs),
35}
36
37/// Arguments for `serve`. Every flag is optional and, when present, overlays the
38/// value loaded from the config file (highest precedence). With no config file,
39/// `--upstream-url` is required.
40#[derive(Debug, Default, Clone, Args)]
41pub struct ServeArgs {
42    /// Config file path (YAML). Also read from `STORAGE_PROXY_CONFIG`.
43    #[arg(short, long, env = "STORAGE_PROXY_CONFIG", value_name = "PATH")]
44    pub config: Option<String>,
45
46    /// Host/interface to bind. Overrides config; default 0.0.0.0 (all interfaces).
47    #[arg(long)]
48    pub host: Option<String>,
49
50    /// TCP port to listen on. Overrides config; default 8080.
51    #[arg(long, short)]
52    pub port: Option<u16>,
53
54    /// Upstream Unity Catalog base URL to resolve + vend through, e.g.
55    /// `http://uc:8080/api/2.1/unity-catalog/`. Overrides config; required when
56    /// no config file is supplied.
57    #[arg(long, value_name = "URL")]
58    pub upstream_url: Option<String>,
59
60    /// Bearer token for the upstream UC. Overrides config. Absent = contact the
61    /// upstream unauthenticated. Prefer a config-file `{ env: ... }` reference in
62    /// production so the token is not visible in the process arguments.
63    #[arg(long, value_name = "TOKEN")]
64    pub upstream_token: Option<String>,
65}
66
67impl ServeArgs {
68    /// Resolve the effective [`Config`]: load the file when given, else start
69    /// from the CLI-supplied upstream, then overlay the remaining flags.
70    pub fn resolve_config(&self) -> Result<Config, String> {
71        let mut cfg = match &self.config {
72            Some(path) => Config::load(path)?,
73            None => {
74                let base_url = self.upstream_url.clone().ok_or_else(|| {
75                    "no config file and no --upstream-url: the upstream Unity Catalog URL is \
76                     required"
77                        .to_string()
78                })?;
79                Config {
80                    host: None,
81                    port: None,
82                    base_path: None,
83                    upstream: UpstreamConfig {
84                        base_url,
85                        token: None,
86                    },
87                    auth: AuthConfig::default(),
88                }
89            }
90        };
91        self.overlay(&mut cfg);
92        Ok(cfg)
93    }
94
95    /// Overlay the provided flags onto `cfg` (highest precedence).
96    fn overlay(&self, cfg: &mut Config) {
97        if let Some(host) = &self.host {
98            cfg.host = Some(host.clone());
99        }
100        if let Some(port) = self.port {
101            cfg.port = Some(port);
102        }
103        if let Some(url) = &self.upstream_url {
104            cfg.upstream.base_url = url.clone();
105        }
106        if let Some(token) = &self.upstream_token {
107            cfg.upstream.token = Some(config::ConfigValue::Value(token.clone()));
108        }
109    }
110}
111
112/// Arguments for `healthcheck`. Shares the config/host/port resolution inputs so
113/// the probe targets the same address the server binds.
114#[derive(Debug, Default, Clone, Args)]
115pub struct HealthcheckArgs {
116    /// Config file path used to resolve the probe target. Also
117    /// `STORAGE_PROXY_CONFIG`.
118    #[arg(short, long, env = "STORAGE_PROXY_CONFIG", value_name = "PATH")]
119    pub config: Option<String>,
120
121    /// Host to connect to (overrides config). A wildcard bind host maps to loopback.
122    #[arg(long)]
123    pub host: Option<String>,
124
125    /// Port to connect to (overrides config).
126    #[arg(long, short)]
127    pub port: Option<u16>,
128
129    /// Probe the full health URL directly, bypassing config load entirely. Use
130    /// when no config file is available to the probe process (e.g. the Docker
131    /// `HEALTHCHECK`, which relies on the default host/port).
132    #[arg(long, value_name = "URL")]
133    pub url: Option<String>,
134
135    /// Probe timeout in seconds.
136    #[arg(long, default_value_t = 3)]
137    pub timeout_secs: u64,
138}
139
140impl HealthcheckArgs {
141    /// The URL to GET: `--url` verbatim if given, else assembled from the loaded
142    /// config (or defaults) with `--host`/`--port` overlaid.
143    fn target_url(&self) -> Result<String, String> {
144        if let Some(url) = &self.url {
145            return Ok(url.clone());
146        }
147        // Without a config file the probe targets the default bind address —
148        // which is exactly what the container's HEALTHCHECK relies on.
149        let host = self
150            .host
151            .clone()
152            .or_else(|| {
153                self.config
154                    .as_ref()
155                    .and_then(|p| Config::load(p).ok())
156                    .and_then(|c| c.host.clone())
157            })
158            .unwrap_or_else(|| DEFAULT_HOST.to_string());
159        let host = match host.as_str() {
160            "0.0.0.0" | "" | "::" => "127.0.0.1".to_string(),
161            _ => host,
162        };
163        let port = self
164            .port
165            .or_else(|| {
166                self.config
167                    .as_ref()
168                    .and_then(|p| Config::load(p).ok())
169                    .and_then(|c| c.port)
170            })
171            .unwrap_or(DEFAULT_PORT);
172        Ok(format!("http://{host}:{port}/health"))
173    }
174}
175
176/// Run the health probe. `Ok(())` iff the endpoint returns a 2xx with body `OK`;
177/// the caller maps any `Err` to a non-zero exit.
178pub fn run_healthcheck(args: &HealthcheckArgs) -> Result<(), String> {
179    let url = args.target_url()?;
180    let client = reqwest::blocking::Client::builder()
181        .timeout(Duration::from_secs(args.timeout_secs))
182        .build()
183        .map_err(|e| e.to_string())?;
184    let resp = client.get(&url).send().map_err(|e| e.to_string())?;
185    let status = resp.status();
186    if !status.is_success() {
187        return Err(format!("health endpoint returned {status}"));
188    }
189    let body = resp.text().map_err(|e| e.to_string())?;
190    if body.trim() != "OK" {
191        return Err(format!("unexpected health body {body:?}"));
192    }
193    Ok(())
194}
195
196#[cfg(test)]
197mod tests {
198    use super::*;
199    use clap::CommandFactory;
200
201    #[test]
202    fn cli_arg_tree_is_valid() {
203        Cli::command().debug_assert();
204    }
205
206    #[test]
207    fn serve_parses_with_upstream_url() {
208        let cli = Cli::try_parse_from([
209            "storage-proxy",
210            "serve",
211            "--upstream-url",
212            "http://uc/",
213            "--port",
214            "9000",
215        ])
216        .unwrap();
217        let Command::Serve(args) = cli.command else {
218            panic!("expected serve");
219        };
220        assert_eq!(args.upstream_url.as_deref(), Some("http://uc/"));
221        assert_eq!(args.port, Some(9000));
222    }
223
224    #[test]
225    fn serve_config_requires_upstream_without_file() {
226        let args = ServeArgs::default();
227        assert!(args.resolve_config().is_err());
228    }
229
230    #[test]
231    fn serve_flags_overlay_and_build_config() {
232        let args = ServeArgs {
233            upstream_url: Some("http://uc/".into()),
234            upstream_token: Some("secret".into()),
235            host: Some("127.0.0.1".into()),
236            port: Some(9000),
237            ..ServeArgs::default()
238        };
239        let cfg = args.resolve_config().unwrap();
240        assert_eq!(cfg.upstream.base_url, "http://uc/");
241        assert_eq!(
242            cfg.upstream.token.as_ref().unwrap().value().as_deref(),
243            Some("secret")
244        );
245        assert_eq!(cfg.resolved_host(), "127.0.0.1");
246        assert_eq!(cfg.resolved_port(), 9000);
247    }
248
249    #[test]
250    fn healthcheck_parses() {
251        let cli = Cli::try_parse_from(["storage-proxy", "healthcheck", "--port", "9000"]).unwrap();
252        let Command::Healthcheck(args) = cli.command else {
253            panic!("expected healthcheck");
254        };
255        assert_eq!(args.port, Some(9000));
256        assert_eq!(args.timeout_secs, 3);
257    }
258
259    #[test]
260    fn healthcheck_url_overrides_config() {
261        let args = HealthcheckArgs {
262            url: Some("http://example.test:1234/health".into()),
263            ..HealthcheckArgs::default()
264        };
265        assert_eq!(
266            args.target_url().unwrap(),
267            "http://example.test:1234/health"
268        );
269    }
270
271    #[test]
272    fn healthcheck_defaults_to_loopback_default_port() {
273        let args = HealthcheckArgs::default();
274        assert_eq!(args.target_url().unwrap(), "http://127.0.0.1:8080/health");
275    }
276
277    #[test]
278    fn no_subcommand_is_an_error() {
279        assert!(Cli::try_parse_from(["storage-proxy"]).is_err());
280    }
281}
unitycatalog_storage_proxy/cli/mod.rs

unitycatalog_storage_proxy/cli/
mod.rs
