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

omni_dev/snowflake/
client.rs
