Skip to main content

omni_dev/snowflake/client/
mod.rs

1//! A clean-room Snowflake client: external-browser SSO + the v1 query endpoint,
2//! with session-token renewal and keep-alive heartbeats.
3//!
4//! Implemented from Snowflake's documented REST protocol — no third-party
5//! connector. Unlike `snowflake-connector-rs`, it captures the **master token**
6//! at login, so the session token can be renewed
7//! ([`SnowflakeSession::renew`]) and kept alive
8//! ([`SnowflakeSession::heartbeat`]) — letting the daemon authenticate once and
9//! stay live indefinitely instead of dying when the ~1h session token lapses.
10//!
11//! Verification status: the offline-decodable paths (result parsing, row→JSON,
12//! request shaping, callback parsing) are unit-tested; the live
13//! SSO/login/query/renew paths follow the documented protocol but need a real
14//! account to verify. The Arrow result format and result sets that span external
15//! chunks are not yet supported (a clear [`Error::Unsupported`] is returned).
16
17pub mod config;
18pub mod error;
19pub mod row;
20
21mod auth;
22mod session;
23mod transport;
24
25use std::sync::Arc;
26
27pub use config::{AuthMethod, BrowserConfig, BrowserLaunch, SnowflakeClientConfig};
28pub use error::{Error, Result};
29pub use row::{rows_to_payload, value_to_json, Column, Row};
30pub use session::SnowflakeSession;
31
32use transport::Transport;
33
34/// A clean-room Snowflake client bound to one account/user configuration.
35pub struct SnowflakeClient {
36    transport: Arc<Transport>,
37    config: SnowflakeClientConfig,
38}
39
40impl SnowflakeClient {
41    /// Builds a client. Cheap — resolves the API host; does not authenticate.
42    ///
43    /// # Errors
44    ///
45    /// [`Error::Protocol`] if the resolved host is invalid, or a transport build
46    /// failure.
47    pub fn new(config: SnowflakeClientConfig) -> Result<Self> {
48        let transport = Arc::new(Transport::new(&config.api_host(), config.http_timeout)?);
49        Ok(Self { transport, config })
50    }
51
52    /// The configuration this client was built with.
53    #[must_use]
54    pub fn config(&self) -> &SnowflakeClientConfig {
55        &self.config
56    }
57
58    /// Authenticates and returns a live session. For external-browser SSO this
59    /// opens a browser once.
60    ///
61    /// # Errors
62    ///
63    /// [`Error::Auth`] when the SSO flow fails, or a transport/server error.
64    pub async fn create_session(&self) -> Result<SnowflakeSession> {
65        let AuthMethod::ExternalBrowser(browser) = &self.config.auth;
66        let tokens = auth::external_browser_login(&self.transport, &self.config, browser).await?;
67        Ok(SnowflakeSession::new(
68            Arc::clone(&self.transport),
69            tokens,
70            self.config.query_timeout,
71        ))
72    }
73}
74
75/// Builds a session whose transport targets `base_url` with placeholder tokens,
76/// so the engine's query orchestration can be exercised offline against a mock
77/// server without going through the (live-only) browser SSO flow.
78#[cfg(test)]
79#[allow(clippy::unwrap_used, clippy::expect_used)]
80pub(crate) fn test_session(base_url: &str, query_timeout: std::time::Duration) -> SnowflakeSession {
81    use std::time::Duration;
82    let url = url::Url::parse(base_url).unwrap().join("/").unwrap();
83    let transport =
84        Arc::new(transport::Transport::with_base_url(url, Duration::from_secs(5)).unwrap());
85    SnowflakeSession::new(
86        transport,
87        session::LoginTokens {
88            session_token: "test-sess".to_string(),
89            master_token: "test-mast".to_string(),
90            session_validity_secs: 3600,
91            master_validity_secs: 14_400,
92        },
93        query_timeout,
94    )
95}
96
97#[cfg(test)]
98#[allow(clippy::unwrap_used, clippy::expect_used)]
99mod tests {
100    use super::*;
101
102    #[test]
103    fn new_builds_a_client_and_exposes_its_config() {
104        let client = SnowflakeClient::new(SnowflakeClientConfig::external_browser("MyAcct", "me"))
105            .expect("a valid account resolves to a valid host");
106        assert_eq!(client.config().account, "MyAcct");
107        assert_eq!(client.config().user, "me");
108        assert_eq!(client.config().api_host(), "myacct.snowflakecomputing.com");
109    }
110
111    #[test]
112    fn new_rejects_an_invalid_host_override() {
113        let config = SnowflakeClientConfig {
114            // An unterminated IPv6 literal can never parse as a URL host.
115            host: Some("[".to_string()),
116            ..SnowflakeClientConfig::external_browser("acct", "me")
117        };
118        assert!(matches!(
119            SnowflakeClient::new(config),
120            Err(Error::Protocol(_))
121        ));
122    }
123}