snowflake-connector-python 5.0.0b5__tar.gz → 5.0.0b6__tar.gz

{snowflake_connector_python-5.0.0b5 → snowflake_connector_python-5.0.0b6}/Cargo.lock

@@ -3128,6 +3128,7 @@ dependencies = [
  "napi-sys",
  "nohash-hasher",
  "rustc-hash",
+ "tokio",
 ]
 
 [[package]]
@@ -3209,6 +3210,9 @@ dependencies = [
  "napi-build",
  "napi-derive",
  "sf_core",
+ "tokio",
+ "tracing",
+ "tracing-subscriber",
 ]
 
 [[package]]

{snowflake_connector_python-5.0.0b5 → snowflake_connector_python-5.0.0b6}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: snowflake-connector-python
-Version: 5.0.0b5
+Version: 5.0.0b6
 Summary: Snowflake DB driver for Python (Private Preview)
 Project-URL: Homepage, https://github.com/snowflakedb/universal-driver
 Project-URL: Bug Reports, https://github.com/snowflakedb/universal-driver/issues

{snowflake_connector_python-5.0.0b5 → snowflake_connector_python-5.0.0b6}/protobuf/database_driver_v1.proto

@@ -351,6 +351,11 @@ message ConnectionGetInfoResponse {
   optional string warehouse = 11;
   optional string master_token = 12;
   optional string user_agent = 13;
+  optional string proxy_host = 14;
+  optional int64 proxy_port = 15;
+  optional string proxy_user = 16;
+  optional string proxy_password = 17;
+  optional string no_proxy = 18;
 }
 
 message ConnectionGetObjectsRequest {
@@ -563,6 +568,39 @@ message ConnectionAbortQueryResponse {
   bool success = 1;
 }
 
+// Streaming PUT: upload an in-memory byte buffer using a PUT SQL statement.
+// Backs JDBC `uploadStream` and Python `cursor.execute(file_stream=…)`.
+message ConnectionUploadStreamRequest {
+  ConnectionHandle conn_handle = 1;
+  // The PUT SQL statement. Must begin with PUT and reference a `file://`
+  // path whose basename is used as the destination filename in the stage.
+  string sql = 2;
+  // Raw bytes to upload. The caller has already drained the source stream.
+  bytes data = 3;
+}
+
+// Mirrors `ResultSetResponse`: a PUT-shaped result set (handle + descriptor),
+// so the caller can read the standard 8-column PUT rowset.
+message ConnectionUploadStreamResponse {
+  ResultSetHandle result_set_handle = 1;
+  ResultSetDescriptor result_descriptor = 2;
+}
+
+// Streaming GET: download a file from a stage and return its bytes. Backs
+// JDBC `downloadStream`.
+message ConnectionDownloadStreamRequest {
+  ConnectionHandle conn_handle = 1;
+  string stage_name = 2;
+  string source_filename = 3;
+  // When true, the core gunzips the downloaded bytes before returning
+  // (opt-in; the bytes are returned verbatim otherwise).
+  bool decompress = 4;
+}
+
+message ConnectionDownloadStreamResponse {
+  bytes data = 1;
+}
+
 // Statement service requests and responses
 
 message StatementNewRequest {
@@ -696,9 +734,9 @@ message ResultChunk {
 
 
 message DatabaseFetchChunkRequest {
-  DatabaseHandle db_handle = 1;
-  ResultChunk chunk = 2;
-  // Required when chunk.format == CHUNK_FORMAT_JSON (forward from ResultSetGetChunksResponse.columns); ignored for CHUNK_FORMAT_ARROW_IPC.
+  ConnectionHandle conn_handle = 1;
+  repeated ResultChunk chunks = 2;
+  // Required for JSON remote chunks (forward from ResultSetGetChunksResponse.columns). Ignored otherwise.
   repeated ColumnMetadata columns = 3;
 }
 
@@ -722,7 +760,7 @@ message ResultSetGetChunksRequest {
 
 message ResultSetGetChunksResponse {
   repeated ResultChunk chunks = 1;
-  // Column metadata. Required for CHUNK_FORMAT_JSON (forward on DatabaseFetchChunkRequest); ignored for CHUNK_FORMAT_ARROW_IPC.
+  // Column metadata. Required when any chunk is CHUNK_FORMAT_JSON; ignored for Arrow-only results.
   repeated ColumnMetadata columns = 2;
 }
 
@@ -922,6 +960,9 @@ service DatabaseDriver {
   rpc ConnectionGetQueryStatus(ConnectionGetQueryStatusRequest) returns (ConnectionGetQueryStatusResponse);
   // Abort a running query by Snowflake Query ID
   rpc ConnectionAbortQuery(ConnectionAbortQueryRequest) returns (ConnectionAbortQueryResponse);
+  // Streaming PUT (upload bytes) / GET (download to bytes) from a stage
+  rpc ConnectionUploadStream(ConnectionUploadStreamRequest) returns (ConnectionUploadStreamResponse);
+  rpc ConnectionDownloadStream(ConnectionDownloadStreamRequest) returns (ConnectionDownloadStreamResponse);
   // Generic HTTP request through sf_core's TLS/OCSP/CRL client
   rpc ConnectionSendHttp(ConnectionSendHttpRequest) returns (ConnectionSendHttpResponse);
   // Token request (ISSUE/RENEW) executed with master token in sf_core

snowflake_connector_python-5.0.0b6/sf_core/CHANGELOG.md

@@ -0,0 +1,11 @@
+# Changelog
+
+## Upcoming Release
+
+Bug fixes:
+
+- Fixed an issue where cached OAuth tokens could be incorrectly shared across different Snowflake accounts or roles that used the same identity provider, and where tokens stored by one driver could not be read by another. The token cache key now includes the IdP URL, Snowflake account URL, username, role, and token type, all normalized and hashed uniformly across drivers (`SnowflakeTokenCache.v2.<sha256>`). Existing v1 cache entries are orphaned; the driver re-authenticates transparently on the next connection and writes a v2 entry. (snowflakedb/universal-driver#TBD)
+
+Test improvements:
+
+- `SnowflakeTestClient` now automatically releases statement handles and result set handles via Drop, eliminating manual `release_statement()` and `result_set_release()` calls in tests and preventing resource leaks when tests panic.

{snowflake_connector_python-5.0.0b5 → snowflake_connector_python-5.0.0b6}/sf_core/Cargo.toml

@@ -139,6 +139,7 @@ arrow_deserialize_macro = { path = "tests/common/arrow_deserialize_macro" }
 bzip2 = "0.6.0"
 brotli = "8.0.2"
 zstd = "0.13.3"
+flate2 = "1.1.2"
 test-case = "3.3.1"
 # Dev/test reqwest: also rustls-only so test builds don't re-introduce
 # native-tls (feature unification would otherwise pull `default-tls` back in).

{snowflake_connector_python-5.0.0b5 → snowflake_connector_python-5.0.0b6}/sf_core/src/apis/database_driver_v1/connection.rs

@@ -25,13 +25,14 @@ use super::validation::{
 use crate::config::ParamStore;
 use crate::config::connection_config::{ConnectionConfig, DiagnosticConfig};
 use crate::config::logout::LogoutConfig;
-use crate::config::param_registry::{ParamKey, ParamScope, param_names};
+use crate::config::param_registry::{ParamKey, param_names};
 use crate::config::resolver;
 use crate::config::rest_parameters::{
     ClientInfo, LoginMethod, LoginParameters, QueryParameters, resolve_log_max_query_length,
     resolve_log_query_parameters, resolve_log_query_text,
 };
 use crate::config::retry::RetryPolicy;
+use crate::config::settings::Settings;
 use crate::diagnostic::DiagnosticRunner;
 use crate::handle_manager::Handle;
 use crate::rest::snowflake::{
@@ -39,7 +40,7 @@ use crate::rest::snowflake::{
     heartbeat, snowflake_query_with_client,
 };
 use crate::sensitive::SensitiveString;
-use crate::token_cache::TokenCache;
+use crate::tls::config::ProxyConfig;
 
 /// Whether `execute_session_sql` should refresh the connection's local
 /// session-state cache (`session_parameters`, `final_session_names`) from
@@ -298,16 +299,16 @@ impl DatabaseDriverV1 {
                     let init_params = conn.init_session_parameters.clone();
                     let resolved_snapshot = resolved.clone();
 
-                    // Forward unrecognized settings as session parameters so
-                    // drivers can set arbitrary Snowflake session params
-                    // via regular connection options.
-                    let mut unknown_settings = collect_unknown_settings(&conn.connection_seed);
-                    // `CLIENT_SESSION_KEEP_ALIVE` and
-                    // `CLIENT_SESSION_KEEP_ALIVE_HEARTBEAT_FREQUENCY` are registered
-                    // params so they don't show up as "unknown"; mirror them into
-                    // login session parameters to match the Python connector.
+                    // Session parameters to send in the login request. Seeded
+                    // with unrecognized connection options (so drivers can set
+                    // arbitrary Snowflake session params via regular connection
+                    // options), then augmented below with registered session
+                    // params that belong in SESSION_PARAMETERS
+                    // (CLIENT_SESSION_KEEP_ALIVE, its heartbeat frequency,
+                    // CLIENT_PREFETCH_THREADS, QUERY_TAG) to match the Python connector.
+                    let mut login_session_params = collect_unknown_settings(&conn.connection_seed);
                     if let Some(v) = resolved.get_bool(param_names::CLIENT_SESSION_KEEP_ALIVE) {
-                        unknown_settings.insert(
+                        login_session_params.insert(
                             param_names::CLIENT_SESSION_KEEP_ALIVE.as_str().to_string(),
                             v.to_string(),
                         );
@@ -315,22 +316,34 @@ impl DatabaseDriverV1 {
                     if let Some(v) =
                         resolved.get_int(param_names::CLIENT_SESSION_KEEP_ALIVE_HEARTBEAT_FREQUENCY)
                     {
-                        unknown_settings.insert(
+                        login_session_params.insert(
                             param_names::CLIENT_SESSION_KEEP_ALIVE_HEARTBEAT_FREQUENCY
                                 .as_str()
                                 .to_string(),
                             v.to_string(),
                         );
                     }
+                    if let Some(v) = resolved.get_int(param_names::CLIENT_PREFETCH_THREADS) {
+                        login_session_params.insert(
+                            param_names::CLIENT_PREFETCH_THREADS.as_str().to_string(),
+                            v.to_string(),
+                        );
+                    }
                     if resolved
                         .get_bool(param_names::VALIDATE_DEFAULT_PARAMETERS)
                         .unwrap_or(false)
                     {
-                        unknown_settings.insert(
+                        login_session_params.insert(
                             "CLIENT_VALIDATE_DEFAULT_PARAMETERS".to_string(),
                             "true".to_string(),
                         );
                     }
+                    // QUERY_TAG is a registered session parameter (also
+                    // statement-overridable), so it isn't "unknown"; mirror a
+                    // connection-level value into the login session parameters.
+                    if let Some(v) = resolved.get_string(param_names::QUERY_TAG) {
+                        login_session_params.insert("QUERY_TAG".to_string(), v);
+                    }
                     let init_params = match init_params {
                         Some(explicit) => {
                             // Normalize explicit keys to uppercase so precedence
@@ -340,12 +353,12 @@ impl DatabaseDriverV1 {
                                 .map(|(k, v)| (k.to_uppercase(), v))
                                 .collect();
                             // Explicit session params take precedence
-                            for (k, v) in unknown_settings {
+                            for (k, v) in login_session_params {
                                 merged.entry(k).or_insert(v);
                             }
                             Some(merged)
                         }
-                        None if !unknown_settings.is_empty() => Some(unknown_settings),
+                        None if !login_session_params.is_empty() => Some(login_session_params),
                         None => None,
                     };
 
@@ -439,7 +452,7 @@ impl DatabaseDriverV1 {
                     &http_client,
                     &login_parameters,
                     init_params.as_ref(),
-                    token_cache.map(|c| c as &dyn TokenCache),
+                    token_cache,
                     Some(&self.prompt_locks),
                     &retry_policy,
                 );
@@ -557,13 +570,21 @@ impl DatabaseDriverV1 {
                             session_token: conn.tokens.clone(),
                         });
 
-                        // unwrap is safe: telemetry_enabled is only true when
-                        // telemetry_sessions() is Some.
-                        self.telemetry_sessions()
-                            .unwrap()
-                            .write()
-                            .unwrap_or_else(|e| e.into_inner())
-                            .insert(session_id, exporter_session);
+                        if let Some(sessions) = self.telemetry_sessions() {
+                            sessions
+                                .write()
+                                .unwrap_or_else(|e| e.into_inner())
+                                .insert(session_id, exporter_session);
+                        } else {
+                            // `telemetry_enabled` implies `telemetry_sessions()`
+                            // was Some (checked above); reaching None here means
+                            // the registry was cleared concurrently.
+                            tracing::error!(
+                                "telemetry session registry unexpectedly absent after \
+                                 telemetry_enabled check; session {session_id} not \
+                                 registered for telemetry"
+                            );
+                        }
 
                         let env_info = conn
                             .wrapper_identity
@@ -1082,6 +1103,14 @@ impl Connection {
         crate::rest::snowflake::query_response::read_use_s3_regional_url_session_param(&params)
     }
 
+    /// Reads `unsafe_file_write` from the connection seed. Default `false`
+    /// (owner-only 0o600 permissions on GET downloads on Unix).
+    pub(crate) fn unsafe_file_write(&self) -> bool {
+        self.connection_seed
+            .get_bool(param_names::UNSAFE_FILE_WRITE)
+            .unwrap_or(false)
+    }
+
     /// The resolved TLS config for this connection, read from the established
     /// [`ClientInfo`]. Falls back to the default `TlsConfig` before login
     /// (when `client_info` is unset). Used by the storage (S3/GCS/Azure) HTTP
@@ -1313,7 +1342,7 @@ impl RefreshContext {
     pub async fn from_arc(conn: &Arc<Mutex<Connection>>) -> Result<Self, ApiError> {
         let guard = conn.lock().await;
         if guard.is_closed.load(Ordering::SeqCst) {
-            return Err(ConnectionClosedSnafu {}.build());
+            return ConnectionClosedSnafu {}.fail();
         }
         Self::new(&guard)
     }
@@ -1503,10 +1532,9 @@ impl RefreshContext {
         F: Fn(SensitiveString) -> Fut,
         Fut: Future<Output = Result<T, RestError>>,
     {
-        use snafu::IntoError;
         crate::refresh::execute_with_refresh(self, |token| {
             let fut = f(token);
-            async move { fut.await.map_err(|e| QuerySnafu.into_error(e)) }
+            async move { fut.await.context(QuerySnafu) }
         })
         .await
     }
@@ -1695,6 +1723,16 @@ pub struct ConnectionInfo {
     pub master_token: Option<SensitiveString>,
     /// The User-Agent string built by the core for this client
     pub user_agent: Option<String>,
+    /// The configured HTTP proxy hostname, if any
+    pub proxy_host: Option<String>,
+    /// The configured HTTP proxy port, if any
+    pub proxy_port: Option<i64>,
+    /// The configured HTTP proxy username for Basic auth, if any
+    pub proxy_user: Option<String>,
+    /// The configured HTTP proxy password for Basic auth (redacted in Debug output)
+    pub proxy_password: Option<SensitiveString>,
+    /// Comma-separated list of hosts that bypass the proxy, if configured
+    pub no_proxy: Option<String>,
 }
 
 fn setting_as_display_string(setting: &Setting) -> Option<String> {
@@ -1848,6 +1886,18 @@ impl DatabaseDriverV1 {
                     .as_ref()
                     .map(crate::rest::snowflake::user_agent);
 
+                // Rebuild the effective proxy config from the resolved login
+                // snapshot (falling back to the raw seed before login) so the
+                // legacy `PROXY` URL form is merged with the individual
+                // `proxy_*` fields exactly as the HTTP client sees them. Reuses
+                // the single parser rather than re-reading fields piecemeal.
+                let proxy_settings: &dyn Settings = conn
+                    .resolved_connect
+                    .as_ref()
+                    .map(|s| s as &dyn Settings)
+                    .unwrap_or(&conn.connection_seed);
+                let proxy = ProxyConfig::from_settings(proxy_settings);
+
                 Ok(ConnectionInfo {
                     host,
                     port,
@@ -1862,6 +1912,11 @@ impl DatabaseDriverV1 {
                     warehouse,
                     master_token,
                     user_agent,
+                    proxy_host: proxy.host,
+                    proxy_port: proxy.port,
+                    proxy_user: proxy.user,
+                    proxy_password: proxy.password,
+                    no_proxy: proxy.no_proxy,
                 })
             }
             None => InvalidArgumentSnafu {
@@ -2023,7 +2078,7 @@ impl DatabaseDriverV1 {
 
                 let (canonical, def) = canonicalize_setting_key(&key);
                 if let Some(d) = def
-                    && d.scope == ParamScope::Session
+                    && d.is_session_scoped()
                 {
                     if let Some(s) = conn
                         .session_overrides
@@ -2908,6 +2963,85 @@ mod tests {
         assert_eq!(info.warehouse, None);
         assert!(info.session_token.is_none());
         assert_eq!(info.session_id, None);
+        assert_eq!(info.proxy_host, None);
+        assert_eq!(info.proxy_port, None);
+        assert_eq!(info.proxy_user, None);
+        assert!(info.proxy_password.is_none());
+        assert_eq!(info.no_proxy, None);
+
+        ds.connection_release(handle).unwrap();
+    }
+
+    #[tokio::test]
+    async fn connection_get_info_returns_proxy_settings() {
+        let ds = DatabaseDriverV1::new();
+        let handle = ds.connection_new();
+        ds.connection_set_option(
+            handle,
+            "proxy_host".into(),
+            Setting::String("proxy.example.com".into()),
+        )
+        .await
+        .unwrap();
+        ds.connection_set_option(handle, "proxy_port".into(), Setting::Int(8080))
+            .await
+            .unwrap();
+        ds.connection_set_option(handle, "proxy_user".into(), Setting::String("puser".into()))
+            .await
+            .unwrap();
+        ds.connection_set_option(
+            handle,
+            "proxy_password".into(),
+            Setting::String("ppass".into()),
+        )
+        .await
+        .unwrap();
+        ds.connection_set_option(
+            handle,
+            "no_proxy".into(),
+            Setting::String("localhost,127.0.0.1".into()),
+        )
+        .await
+        .unwrap();
+
+        let info = ds.connection_get_info(handle).await.unwrap();
+
+        assert_eq!(info.proxy_host, Some("proxy.example.com".into()));
+        assert_eq!(info.proxy_port, Some(8080));
+        assert_eq!(info.proxy_user, Some("puser".into()));
+        assert_eq!(
+            info.proxy_password.as_ref().map(|p| p.reveal().as_str()),
+            Some("ppass")
+        );
+        assert_eq!(info.no_proxy, Some("localhost,127.0.0.1".into()));
+
+        ds.connection_release(handle).unwrap();
+    }
+
+    /// The legacy ODBC `PROXY` URL form must be parsed and surfaced through
+    /// `connection_get_info` exactly as the HTTP client sees it, since the
+    /// info aggregation reuses `ProxyConfig::from_settings`.
+    #[tokio::test]
+    async fn connection_get_info_parses_legacy_proxy_url() {
+        let ds = DatabaseDriverV1::new();
+        let handle = ds.connection_new();
+        ds.connection_set_option(
+            handle,
+            "proxy".into(),
+            Setting::String("http://puser:ppass@proxy.example.com:8080".into()),
+        )
+        .await
+        .unwrap();
+
+        let info = ds.connection_get_info(handle).await.unwrap();
+
+        assert_eq!(info.proxy_host, Some("proxy.example.com".into()));
+        assert_eq!(info.proxy_port, Some(8080));
+        assert_eq!(info.proxy_user, Some("puser".into()));
+        assert_eq!(
+            info.proxy_password.as_ref().map(|p| p.reveal().as_str()),
+            Some("ppass")
+        );
 
         ds.connection_release(handle).unwrap();
     }

snowflake_connector_python-5.0.0b6/sf_core/src/apis/database_driver_v1/database.rs

@@ -0,0 +1,126 @@
+use std::collections::HashMap;
+use tokio::sync::Mutex;
+
+use super::error::*;
+use super::global_state::DatabaseDriverV1;
+use super::validation::{ValidationIssue, resolve_and_apply_options};
+use crate::chunks::{ChunkFormatKind, FetchChunkInput, PrefetchConfig, fetch_chunks_reader};
+use crate::config::ParamStore;
+use crate::config::settings::Setting;
+use crate::handle_manager::Handle;
+use crate::query_types::RowType;
+use crate::tls;
+use arrow::ffi_stream::FFI_ArrowArrayStream;
+use snafu::ResultExt;
+
+impl DatabaseDriverV1 {
+    pub fn database_new(&self) -> Handle {
+        self.databases.add_handle(Mutex::new(Database::new()))
+    }
+
+    pub async fn database_set_options(
+        &self,
+        db_handle: Handle,
+        options: HashMap<String, Setting>,
+    ) -> Result<Vec<ValidationIssue>, ApiError> {
+        match self.databases.get_obj(db_handle) {
+            Some(db_ptr) => {
+                let mut db = db_ptr.lock().await;
+                resolve_and_apply_options(&mut db.settings, options)
+            }
+            None => InvalidArgumentSnafu {
+                argument: "Database handle not found".to_string(),
+            }
+            .fail(),
+        }
+    }
+
+    pub fn database_init(&self, db_handle: Handle) -> Result<(), ApiError> {
+        match self.databases.get_obj(db_handle) {
+            Some(_db_ptr) => Ok(()),
+            None => InvalidArgumentSnafu {
+                argument: "Database handle not found".to_string(),
+            }
+            .fail(),
+        }
+    }
+
+    pub fn database_release(&self, db_handle: Handle) -> Result<(), ApiError> {
+        match self.databases.delete_handle(db_handle) {
+            true => Ok(()),
+            false => InvalidArgumentSnafu {
+                argument: "Failed to release database handle".to_string(),
+            }
+            .fail(),
+        }
+    }
+
+    pub async fn database_fetch_chunk(
+        &self,
+        conn_handle: Option<Handle>,
+        chunks: Vec<FetchChunkInput>,
+        chunk_format: ChunkFormatKind,
+        nullable_flags: &[bool],
+        row_types: Vec<RowType>,
+    ) -> Result<Box<FFI_ArrowArrayStream>, ApiError> {
+        let (client, prefetch_config) = match conn_handle {
+            Some(conn_handle) => {
+                let conn_ptr = self.connections.get_obj(conn_handle).ok_or_else(|| {
+                    InvalidArgumentSnafu {
+                        argument: "Connection handle not found".to_string(),
+                    }
+                    .build()
+                })?;
+                let conn = conn_ptr.lock().await;
+                let client = conn
+                    .http_client
+                    .clone()
+                    .ok_or_else(|| ConnectionNotInitializedSnafu.build())?;
+                let session_params = conn.session_parameters.read().await;
+                let prefetch_config = PrefetchConfig::from_session_params(&session_params);
+                (client, prefetch_config)
+            }
+            None => {
+                // TODO(SNOW-3801967): when no conn handle, we fall back to a fresh TLS client and
+                //  default prefetch config, but we could restore proper configuration if it's
+                //  serialized with the result chunk
+                let client = tls::create_tls_client_with_config(
+                    tls::TlsConfig::default(),
+                    self.crl_worker.clone(),
+                )
+                .context(TlsClientCreationSnafu)?;
+                (client, PrefetchConfig::default())
+            }
+        };
+        let reader = fetch_chunks_reader(
+            chunks,
+            chunk_format,
+            row_types,
+            nullable_flags,
+            client,
+            &prefetch_config,
+        )
+        .await
+        .context(ChunkFetchSnafu)?;
+
+        Ok(Box::new(FFI_ArrowArrayStream::new(reader)))
+    }
+}
+
+pub struct Database {
+    pub(crate) settings: ParamStore,
+}
+
+impl Default for Database {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl Database {
+    pub fn new() -> Self {
+        Database {
+            settings: ParamStore::new(),
+        }
+    }
+}

{snowflake_connector_python-5.0.0b5 → snowflake_connector_python-5.0.0b6}/sf_core/src/apis/database_driver_v1/error.rs

@@ -159,6 +159,12 @@ pub enum ApiError {
         #[snafu(implicit)]
         location: Location,
     },
+    #[snafu(display("Background chunk-decode task failed to join"))]
+    BlockingTaskJoin {
+        source: tokio::task::JoinError,
+        #[snafu(implicit)]
+        location: Location,
+    },
     #[snafu(display("Failed to encode inline JSON rowset as Arrow IPC"))]
     InlineJsonEncode {
         #[snafu(implicit)]

{snowflake_connector_python-5.0.0b5 → snowflake_connector_python-5.0.0b6}/sf_core/src/apis/database_driver_v1/global_state.rs

@@ -14,7 +14,7 @@ use crate::logging::LogManager;
 use crate::rest::snowflake::prompt_lock::PromptLockMap;
 use crate::telemetry::platform_detection::{DetectionConfig, detect_platforms};
 use crate::telemetry::snowflake_exporter::SessionRegistry;
-use crate::token_cache::{KeyringTokenCache, TokenCacheError};
+use crate::token_cache::{KeyringTokenCache, TokenCache, TokenCacheError};
 
 /// Which shape the PUT/GET result set should take.
 #[derive(Debug, Clone, Default, PartialEq, Eq)]
@@ -95,12 +95,13 @@ pub struct DatabaseDriverV1 {
     pub(super) connections: HandleManager<Mutex<Connection>>,
     pub(super) statements: HandleManager<Mutex<Statement>>,
     pub(super) results: HandleManager<Mutex<ResultSet>>,
-    token_cache: once_cell::sync::OnceCell<KeyringTokenCache>,
+    token_cache: once_cell::sync::OnceCell<Arc<dyn TokenCache>>,
     fs: Arc<dyn FsAdapter>,
     platforms: tokio::sync::OnceCell<Vec<String>>,
     log_manager: Option<LogManager>,
     pub(super) wrapper_presets: WrapperPresets,
-    /// Process-global per-`(host, user, token-type)` prompt locks.
+    /// Process-global per-[`crate::token_cache::CacheKey`] prompt locks
+    /// (scoped by idp, snowflake, username, role, and token_type).
     /// Shared across all connections on this driver instance.
     pub(crate) prompt_locks: Arc<PromptLockMap>,
     /// Lazy CRL worker shared across all connections on this driver instance.
@@ -186,8 +187,12 @@ impl DatabaseDriverV1 {
         }
     }
 
-    pub fn token_cache(&self) -> Result<&KeyringTokenCache, TokenCacheError> {
-        self.token_cache.get_or_try_init(KeyringTokenCache::new)
+    pub fn token_cache(&self) -> Result<Arc<dyn TokenCache>, TokenCacheError> {
+        self.token_cache
+            .get_or_try_init(|| {
+                KeyringTokenCache::new().map(|c| Arc::new(c) as Arc<dyn TokenCache>)
+            })
+            .map(Arc::clone)
     }
 
     pub fn fs_adapter(&self) -> Arc<dyn FsAdapter> {
@@ -244,7 +249,7 @@ mod tests {
         let first = driver.token_cache().expect("first call failed");
         let second = driver.token_cache().expect("second call failed");
         assert!(
-            std::ptr::eq(first, second),
+            Arc::ptr_eq(&first, &second),
             "token_cache() should return the same instance on repeated calls"
         );
     }

{snowflake_connector_python-5.0.0b5 → snowflake_connector_python-5.0.0b6}/sf_core/src/apis/database_driver_v1/logout.rs

@@ -138,12 +138,11 @@ fn validate_config(config: &LogoutConfig) -> Result<(), ApiError> {
     if let Some(timeout) = config.logout_request_timeout
         && timeout.is_zero()
     {
-        return Err(InvalidArgumentSnafu {
+        return InvalidArgumentSnafu {
             argument:
                 "logout_request_timeout: 0s. Zero timeout means immediate failure. Must be positive."
                     .to_string(),
-        }
-        .build());
+        }.fail();
     }
     Ok(())
 }

{snowflake_connector_python-5.0.0b5 → snowflake_connector_python-5.0.0b6}/sf_core/src/apis/database_driver_v1/mod.rs

@@ -20,13 +20,14 @@ mod query;
 pub(crate) mod result_set;
 pub mod spcs_token;
 pub(crate) mod statement;
+mod stream_transfer;
 pub(crate) mod validation;
 
+pub use crate::chunks::FetchChunkInput;
 pub use crate::config::settings::Setting;
 pub use crate::handle_manager::Handle;
 pub use async_query_registry::AsyncQueryRegistry;
 pub use connection::{Connection, ConnectionInfo, RefreshContext, with_valid_session};
-pub use database::FetchChunkInput;
 pub use error::ApiError;
 pub use get_objects::{
     ColumnDescriptor, DEPTH_CATALOGS, DEPTH_COLUMNS, DEPTH_DB_SCHEMAS, DEPTH_TABLES,