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 jwt;
25mod session;
26mod transport;
27
28use std::sync::Arc;
29
30pub use config::{AuthMethod, BrowserConfig, BrowserLaunch, KeyPairConfig, SnowflakeClientConfig};
31pub use error::{Error, Result};
32pub use row::{rows_to_payload, value_to_json, Column, Row};
33pub use session::SnowflakeSession;
34
35use transport::Transport;
36
37/// A clean-room Snowflake client bound to one account/user configuration.
38pub struct SnowflakeClient {
39    transport: Arc<Transport>,
40    config: SnowflakeClientConfig,
41}
42
43impl SnowflakeClient {
44    /// Builds a client. Cheap — resolves the API host; does not authenticate.
45    ///
46    /// # Errors
47    ///
48    /// [`Error::Protocol`] if the resolved host is invalid, or a transport build
49    /// failure.
50    pub fn new(config: SnowflakeClientConfig) -> Result<Self> {
51        let transport = Arc::new(Transport::new(&config.api_host(), config.http_timeout)?);
52        Ok(Self { transport, config })
53    }
54
55    /// The configuration this client was built with.
56    #[must_use]
57    pub fn config(&self) -> &SnowflakeClientConfig {
58        &self.config
59    }
60
61    /// Authenticates and returns a live session. For external-browser SSO this
62    /// opens a browser once; PAT and key-pair JWT are non-interactive.
63    ///
64    /// # Errors
65    ///
66    /// [`Error::Auth`] when the auth flow fails, or a transport/server error.
67    pub async fn create_session(&self) -> Result<SnowflakeSession> {
68        let tokens = match &self.config.auth {
69            AuthMethod::ExternalBrowser(browser) => {
70                auth::external_browser_login(&self.transport, &self.config, browser).await?
71            }
72            AuthMethod::ProgrammaticAccessToken { token } => {
73                auth::pat_login(&self.transport, &self.config, token).await?
74            }
75            AuthMethod::KeyPairJwt(keypair) => {
76                auth::keypair_jwt_login(&self.transport, &self.config, keypair).await?
77            }
78        };
79        Ok(SnowflakeSession::new(
80            Arc::clone(&self.transport),
81            tokens,
82            self.config.query_timeout,
83        ))
84    }
85}
86
87/// Builds a session whose transport targets `base_url` with placeholder tokens,
88/// so the engine's query orchestration can be exercised offline against a mock
89/// server without going through the (live-only) browser SSO flow.
90#[cfg(test)]
91#[allow(clippy::unwrap_used, clippy::expect_used)]
92pub(crate) fn test_session(base_url: &str, query_timeout: std::time::Duration) -> SnowflakeSession {
93    use std::time::Duration;
94    let url = url::Url::parse(base_url).unwrap().join("/").unwrap();
95    let transport =
96        Arc::new(transport::Transport::with_base_url(url, Duration::from_secs(5)).unwrap());
97    SnowflakeSession::new(
98        transport,
99        session::LoginTokens {
100            session_token: "test-sess".into(),
101            master_token: "test-mast".into(),
102            session_validity_secs: 3600,
103            master_validity_secs: 14_400,
104        },
105        query_timeout,
106    )
107}
108
109/// Builds a client whose transport targets `base_url` with the given auth
110/// method, so the non-interactive login flows (PAT / key-pair JWT) can be
111/// exercised offline against a mock `session/v1/login-request`.
112#[cfg(test)]
113#[allow(clippy::unwrap_used, clippy::expect_used)]
114pub(crate) fn test_client(base_url: &str, auth: AuthMethod) -> SnowflakeClient {
115    use std::time::Duration;
116    let url = url::Url::parse(base_url).unwrap().join("/").unwrap();
117    let transport =
118        Arc::new(transport::Transport::with_base_url(url, Duration::from_secs(5)).unwrap());
119    let mut config = SnowflakeClientConfig::external_browser("TESTACCT", "tester");
120    config.auth = auth;
121    SnowflakeClient { transport, config }
122}
123
124#[cfg(test)]
125#[allow(clippy::unwrap_used, clippy::expect_used)]
126mod tests {
127    use super::*;
128
129    #[test]
130    fn new_builds_a_client_and_exposes_its_config() {
131        let client = SnowflakeClient::new(SnowflakeClientConfig::external_browser("MyAcct", "me"))
132            .expect("a valid account resolves to a valid host");
133        assert_eq!(client.config().account, "MyAcct");
134        assert_eq!(client.config().user, "me");
135        assert_eq!(client.config().api_host(), "myacct.snowflakecomputing.com");
136    }
137
138    #[test]
139    fn new_rejects_an_invalid_host_override() {
140        let config = SnowflakeClientConfig {
141            // An unterminated IPv6 literal can never parse as a URL host.
142            host: Some("[".to_string()),
143            ..SnowflakeClientConfig::external_browser("acct", "me")
144        };
145        assert!(matches!(
146            SnowflakeClient::new(config),
147            Err(Error::Protocol(_))
148        ));
149    }
150
151    #[tokio::test]
152    async fn pat_create_session_sends_the_token_and_captures_login_tokens() {
153        use wiremock::matchers::{method, path};
154        use wiremock::{Mock, MockServer, ResponseTemplate};
155
156        let server = MockServer::start().await;
157        Mock::given(method("POST"))
158            .and(path("/session/v1/login-request"))
159            .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
160                "success": true,
161                "data": {
162                    "token": "sess-tok",
163                    "masterToken": "mast-tok",
164                    "validityInSeconds": 3600,
165                    "masterValidityInSeconds": 14_400,
166                },
167            })))
168            .expect(1)
169            .mount(&server)
170            .await;
171
172        let client = test_client(
173            &server.uri(),
174            AuthMethod::ProgrammaticAccessToken {
175                token: "pat-xyz".into(),
176            },
177        );
178        // A live session means the login response was parsed into tokens.
179        client
180            .create_session()
181            .await
182            .expect("PAT login should succeed against the mock");
183
184        // The request must present the PAT as TOKEN under the PAT authenticator.
185        let requests = server.received_requests().await.expect("recorded requests");
186        let body: serde_json::Value = serde_json::from_slice(&requests[0].body).unwrap();
187        assert_eq!(body["data"]["AUTHENTICATOR"], "PROGRAMMATIC_ACCESS_TOKEN");
188        assert_eq!(body["data"]["TOKEN"], "pat-xyz");
189        assert_eq!(body["data"]["LOGIN_NAME"], "tester");
190        assert_eq!(body["data"]["ACCOUNT_NAME"], "TESTACCT");
191    }
192
193    #[tokio::test]
194    async fn keypair_jwt_create_session_sends_a_jwt_and_captures_login_tokens() {
195        use aws_lc_rs::encoding::AsDer as _;
196        use base64::Engine as _;
197        use wiremock::matchers::{method, path};
198        use wiremock::{Mock, MockServer, ResponseTemplate};
199
200        let server = MockServer::start().await;
201        Mock::given(method("POST"))
202            .and(path("/session/v1/login-request"))
203            .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
204                "success": true,
205                "data": {
206                    "token": "sess-tok",
207                    "masterToken": "mast-tok",
208                    "validityInSeconds": 3600,
209                    "masterValidityInSeconds": 14_400,
210                },
211            })))
212            .expect(1)
213            .mount(&server)
214            .await;
215
216        // A throwaway RSA key so `build_jwt` has something to sign with.
217        let key_pair = aws_lc_rs::rsa::KeyPair::generate(aws_lc_rs::rsa::KeySize::Rsa2048).unwrap();
218        let der = key_pair.as_der().unwrap();
219        let pem = format!(
220            "-----BEGIN PRIVATE KEY-----\n{}\n-----END PRIVATE KEY-----\n",
221            base64::engine::general_purpose::STANDARD.encode(der.as_ref())
222        );
223
224        let client = test_client(
225            &server.uri(),
226            AuthMethod::KeyPairJwt(KeyPairConfig {
227                private_key_pem: pem.into(),
228            }),
229        );
230        client
231            .create_session()
232            .await
233            .expect("key-pair JWT login should succeed against the mock");
234
235        let requests = server.received_requests().await.expect("recorded requests");
236        let body: serde_json::Value = serde_json::from_slice(&requests[0].body).unwrap();
237        assert_eq!(body["data"]["AUTHENTICATOR"], "SNOWFLAKE_JWT");
238        let jwt = body["data"]["TOKEN"].as_str().expect("a JWT string");
239        assert_eq!(jwt.split('.').count(), 3, "TOKEN must be a compact JWS");
240    }
241}
omni_dev/snowflake/client.rs

omni_dev/snowflake/
client.rs
