Skip to main content

omni_dev/snowflake/client/
config.rs

1//! Configuration for the clean-room Snowflake client.
2
3use std::net::{IpAddr, Ipv4Addr};
4use std::time::Duration;
5
6/// Default per-request HTTP timeout for Snowflake REST calls.
7///
8/// Generous so the query-request long-poll (the server holds it open until
9/// results are ready or its synchronous window elapses, then returns an "in
10/// progress" code) is not cut short; long-running queries are bounded by
11/// `query_timeout` instead.
12pub const DEFAULT_HTTP_TIMEOUT: Duration = Duration::from_secs(120);
13
14/// Default overall deadline for one query, including async-result polling.
15pub const DEFAULT_QUERY_TIMEOUT: Duration = Duration::from_secs(3600);
16
17/// How to open the SSO URL during external-browser authentication.
18#[derive(Clone, Debug, Default)]
19pub enum BrowserLaunch {
20    /// Open with the OS default handler (`open` / `xdg-open` / `start`).
21    #[default]
22    Auto,
23    /// Run a custom command; `{url}` (or a trailing arg) receives the SSO URL.
24    /// Use this to target a specific Chrome profile in a new window, e.g.
25    /// `Google Chrome --profile-directory=Profile 1 --new-window {url}`.
26    Command(Vec<String>),
27    /// Do not open a browser; the SSO URL is logged for manual opening.
28    Manual,
29}
30
31/// External-browser SSO settings.
32#[derive(Clone, Debug)]
33pub struct BrowserConfig {
34    /// How to open the SSO URL.
35    pub launch: BrowserLaunch,
36    /// Bind address for the localhost callback listener.
37    pub callback_addr: IpAddr,
38    /// Bind port for the callback listener (`0` = OS-assigned ephemeral port).
39    pub callback_port: u16,
40}
41
42impl Default for BrowserConfig {
43    fn default() -> Self {
44        Self {
45            launch: BrowserLaunch::Auto,
46            callback_addr: IpAddr::V4(Ipv4Addr::LOCALHOST),
47            callback_port: 0,
48        }
49    }
50}
51
52/// The authentication method used to establish a session.
53#[derive(Clone, Debug)]
54pub enum AuthMethod {
55    /// External-browser SSO (the only method implemented).
56    ExternalBrowser(BrowserConfig),
57}
58
59/// Everything needed to authenticate and run queries against one account/user.
60#[derive(Clone, Debug)]
61pub struct SnowflakeClientConfig {
62    /// Snowflake account identifier.
63    pub account: String,
64    /// Login user.
65    pub user: String,
66    /// Authentication method.
67    pub auth: AuthMethod,
68    /// Default warehouse applied at login.
69    pub warehouse: Option<String>,
70    /// Default role applied at login.
71    pub role: Option<String>,
72    /// Default database applied at login.
73    pub database: Option<String>,
74    /// Default schema applied at login.
75    pub schema: Option<String>,
76    /// Override the API host (defaults to `<account>.snowflakecomputing.com`).
77    pub host: Option<String>,
78    /// Per-request HTTP timeout for REST calls (not the SSO callback wait).
79    pub http_timeout: Duration,
80    /// Overall deadline for one query (submit + async-result polling).
81    pub query_timeout: Duration,
82}
83
84impl SnowflakeClientConfig {
85    /// Builds a config for external-browser SSO with default browser settings.
86    #[must_use]
87    pub fn external_browser(account: impl Into<String>, user: impl Into<String>) -> Self {
88        Self {
89            account: account.into(),
90            user: user.into(),
91            auth: AuthMethod::ExternalBrowser(BrowserConfig::default()),
92            warehouse: None,
93            role: None,
94            database: None,
95            schema: None,
96            host: None,
97            http_timeout: DEFAULT_HTTP_TIMEOUT,
98            query_timeout: DEFAULT_QUERY_TIMEOUT,
99        }
100    }
101
102    /// The API host for this account.
103    ///
104    /// Uses the `host` override when set; otherwise derives
105    /// `<account>.snowflakecomputing.com` from the account identifier — lowercased
106    /// and with underscores mapped to dashes (Snowflake's URL rule), while region/
107    /// cloud dot-segments (e.g. `acct.us-east-1.aws`) are preserved.
108    #[must_use]
109    pub fn api_host(&self) -> String {
110        if let Some(host) = &self.host {
111            return host.clone();
112        }
113        let account = self.account.trim().to_lowercase().replace('_', "-");
114        format!("{account}.snowflakecomputing.com")
115    }
116}
117
118#[cfg(test)]
119#[allow(clippy::unwrap_used, clippy::expect_used)]
120mod tests {
121    use super::*;
122
123    #[test]
124    fn api_host_maps_account_identifiers() {
125        let host =
126            |account: &str| SnowflakeClientConfig::external_browser(account, "me").api_host();
127        // Lowercased.
128        assert_eq!(host("MyAcct"), "myacct.snowflakecomputing.com");
129        // Underscores → dashes (Snowflake URL rule).
130        assert_eq!(host("my_org-acct"), "my-org-acct.snowflakecomputing.com");
131        // Region/cloud dot-segments preserved.
132        assert_eq!(
133            host("xy12345.us-east-1.aws"),
134            "xy12345.us-east-1.aws.snowflakecomputing.com"
135        );
136    }
137
138    #[test]
139    fn api_host_override_wins() {
140        let cfg = SnowflakeClientConfig {
141            host: Some("acct.privatelink.snowflakecomputing.com".to_string()),
142            ..SnowflakeClientConfig::external_browser("MyAcct", "me")
143        };
144        assert_eq!(cfg.api_host(), "acct.privatelink.snowflakecomputing.com");
145    }
146
147    #[test]
148    fn default_http_timeout_is_set() {
149        let cfg = SnowflakeClientConfig::external_browser("a", "b");
150        assert_eq!(cfg.http_timeout, DEFAULT_HTTP_TIMEOUT);
151    }
152}