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