snowflake-connector-python 5.0.0b3__tar.gz → 5.0.0b4__tar.gz

{snowflake_connector_python-5.0.0b3 → snowflake_connector_python-5.0.0b4}/Cargo.lock

@@ -4834,6 +4834,7 @@ dependencies = [
  "aws-lc-rs",
  "aws-sdk-s3",
  "aws-sdk-sts",
+ "aws-smithy-types",
  "axum",
  "base64 0.22.1",
  "brotli",
@@ -4854,6 +4855,8 @@ dependencies = [
  "hex",
  "html-escape",
  "http 1.3.1",
+ "http-body 1.0.1",
+ "http-body-util",
  "infer",
  "jwt",
  "keyring",
@@ -4869,6 +4872,7 @@ dependencies = [
  "opentelemetry-semantic-conventions",
  "opentelemetry_sdk",
  "os_info",
+ "pin-project-lite",
  "pprof",
  "prost 0.14.1",
  "proto_generator",
@@ -4881,6 +4885,7 @@ dependencies = [
  "rustls-native-certs 0.8.1",
  "rustls-pemfile 2.2.0",
  "rustls-webpki 0.102.8",
+ "scopeguard",
  "serde",
  "serde_json",
  "sf_core",

{snowflake_connector_python-5.0.0b3 → snowflake_connector_python-5.0.0b4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: snowflake-connector-python
-Version: 5.0.0b3
+Version: 5.0.0b4
 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.0b3 → snowflake_connector_python-5.0.0b4}/protobuf/database_driver_v1.proto

@@ -345,8 +345,36 @@ message ConnectionGetObjectsRequest {
   optional string column_name = 7;
 }
 
+// ConnectionGetObjects returns a result set (streaming, via result_set_get_stream)
+// containing a nested ADBC-shaped Arrow batch:
+//   catalog_name: utf8
+//   catalog_db_schemas: list<struct<
+//     db_schema_name: utf8,
+//     db_schema_tables: list<struct<
+//       table_name: utf8,
+//       table_type: utf8,
+//       table_columns: list (always empty in this impl),
+//       table_constraints: list (always empty in this impl)
+//     >>
+//   >>
+//
+// depth constants:
+//   0 = ALL (not used)
+//   1 = CATALOGS  -> populates only catalog_name
+//   2 = DB_SCHEMAS -> populates catalog_name + db_schema_name
+//   3 = TABLES    -> populates down to table_name/table_type
+//   4 = COLUMNS   -> deferred (not implemented yet)
+//
+// Filter semantics (search patterns, ADBC + backslash-escape extension):
+//   - NULL filter = no filter (account-wide or connection-context if
+//     CLIENT_METADATA_REQUEST_USE_CONNECTION_CTX is set)
+//   - empty string "" = matches nothing (returns empty result)
+//   - % = any sequence, _ = any single char
+//   - \%, \_, \\ = literal %, _, \ (escape extension, not in ADBC spec)
+//   - Matching is case-insensitive (Snowflake unquoted-identifier semantics)
 message ConnectionGetObjectsResponse {
-  bytes objects_data = 1;
+  ResultSetHandle result_set_handle = 1;
+  ResultSetDescriptor result_descriptor = 2;
 }
 
 message ConnectionGetTableSchemaRequest {
@@ -398,6 +426,14 @@ message ConnectionUseDatabaseRequest {
 message ConnectionUseDatabaseResponse {
 }
 
+message ConnectionUseSchemaRequest {
+  ConnectionHandle conn_handle = 1;
+  string schema = 2;
+}
+
+message ConnectionUseSchemaResponse {
+}
+
 message ConnectionSetSessionParametersRequest {
   ConnectionHandle conn_handle = 1;
   map<string, string> parameters = 2;
@@ -550,6 +586,7 @@ message PrepareResult {
   string query = 6;
   optional string sql_state = 7;
   bool array_bind_supported = 8;
+  repeated ColumnMetadata binds = 9;
 }
 
 message StatementPrepareRequest {
@@ -804,6 +841,10 @@ message ConfigGetPathsResponse {
 message TelemetrySendApiUsageRequest {
   ConnectionHandle conn_handle = 1;
   string api_method = 2;
+  // Names of the arguments the caller explicitly passed to the API method.
+  // Names only — never argument values, and parameters left at their default
+  // are omitted. Used to learn which optional fields callers actually exercise.
+  repeated string passed_arguments = 3;
 }
 
 message TelemetrySendWrapperErrorRequest {
@@ -840,6 +881,7 @@ service DatabaseDriver {
   rpc ConnectionRollback(ConnectionRollbackRequest) returns (ConnectionRollbackResponse);
   rpc ConnectionSetAutocommit(ConnectionSetAutocommitRequest) returns (ConnectionSetAutocommitResponse);
   rpc ConnectionUseDatabase(ConnectionUseDatabaseRequest) returns (ConnectionUseDatabaseResponse);
+  rpc ConnectionUseSchema(ConnectionUseSchemaRequest) returns (ConnectionUseSchemaResponse);
   // sets initial session parameters connection should be initialized with
   // call before opening the connection
   rpc ConnectionSetSessionParameters(ConnectionSetSessionParametersRequest) returns (ConnectionSetSessionParametersResponse);

{snowflake_connector_python-5.0.0b3 → snowflake_connector_python-5.0.0b4}/pyproject.toml

@@ -238,7 +238,11 @@ dependencies = [
     "protobuf>=6.32.1",
     "pytz",
     "tomlkit>=0.13",
-    "numpy>=2.1.0",
+    # Cap < 2.3 to match the mypy `python_version = "3.10"` target: numpy 2.3
+    # dropped 3.10 and its stubs use PEP-695 `type` statements mypy rejects
+    # under a 3.10 target. A real 3.10 install resolves numpy <= 2.2, so this is
+    # the numpy we should type-check against.
+    "numpy>=2.1.0,<2.3",
 ]
 
 [tool.hatch.envs.precommit.env-vars]

{snowflake_connector_python-5.0.0b3 → snowflake_connector_python-5.0.0b4}/sf_core/Cargo.toml

@@ -15,6 +15,7 @@ protobuf = ["prost", "proto_utils"]
 vendored-openssl = ["openssl/vendored"]
 auth_mfa_e2e = []
 auth_oauth_e2e = []
+auth_spcs_e2e = []
 gcs_e2e = []
 perf_timing = []
 test-utils = []
@@ -51,6 +52,10 @@ oauth2 = { version = "5", features = ["timing-resistant-secret-traits"] }
 axum = { version = "0.8", default-features = false, features = ["http1", "tokio", "query"] }
 webbrowser = "1"
 http = "1"
+http-body = "1"
+http-body-util = "0.1"
+pin-project-lite = "0.2"
+tempfile = "3.21.0"
 infer = "0.19.0"
 
 # Proto
@@ -62,6 +67,9 @@ aws-sdk-s3 = "1.103.0"
 aws-sdk-sts = "1.86.0"
 aws-config = "1.8.5"
 aws-credential-types = "1.2.5"
+# Enables `ByteStream::from_body_1_x` so the S3 PUT path can stream a custom
+# (lazily-encrypting) http-body without buffering ciphertext.
+aws-smithy-types = { version = "1.3", features = ["http-body-1-x"] }
 snafu = "0.8.7"
 error_trace = { path = "../error_trace" }
 fs2 = "0.4"
@@ -110,7 +118,6 @@ libc = "0.2"
 sf_core = { path = ".", default-features = false, features = ["test-utils"] }
 tokio = { version = "1.47.1", features = ["fs", "test-util"] }
 opentelemetry_sdk = { version = "0.30.0", default-features = false, features = ["experimental_metrics_custom_reader"] }
-tempfile = "3.21.0"
 arrow_deserialize_macro = { path = "tests/common/arrow_deserialize_macro" }
 bzip2 = "0.6.0"
 brotli = "8.0.2"
@@ -123,6 +130,7 @@ tokio-rustls = "0.26"
 tracing-test = { version = "0.2", features = ["no-env-filter"] }
 criterion = { version = "0.5", features = ["async_tokio"] }
 temp-env = { version = "0.3.6", features = ["async_closure"] }
+scopeguard = "1.2"
 
 [target.'cfg(not(target_os = "windows"))'.dev-dependencies]
 pprof = { version = "0.14", features = ["flamegraph", "criterion"] }
@@ -185,6 +193,11 @@ required-features = ["cli"]
 name = "param_defs_export"
 path = "src/bin/param_defs_export.rs"
 
+[[bin]]
+name = "spcs_probe"
+path = "src/bin/spcs_probe.rs"
+required-features = ["protobuf"]
+
 [[bench]]
 name = "prefetch_bench"
 harness = false

{snowflake_connector_python-5.0.0b3 → snowflake_connector_python-5.0.0b4}/sf_core/exports.def

@@ -4,6 +4,8 @@
 EXPORTS
     sf_core_init
     sf_core_api_call_proto
+    sf_core_api_call_proto_async
+    sf_core_api_cancel
     sf_core_free_buffer
     sf_core_perf_enabled
     sf_core_get_perf_data

{snowflake_connector_python-5.0.0b3 → snowflake_connector_python-5.0.0b4}/sf_core/src/apis/database_driver_v1/connection.rs

@@ -23,7 +23,7 @@ use super::validation::{
     validate_session_override_write,
 };
 use crate::config::ParamStore;
-use crate::config::connection_config::ConnectionConfig;
+use crate::config::connection_config::{ConnectionConfig, DiagnosticConfig};
 use crate::config::logout::LogoutConfig;
 use crate::config::param_registry::{
     DEFAULT_PUT_GET_MAX_ATTEMPTS, ParamKey, ParamScope, param_names,
@@ -34,6 +34,7 @@ use crate::config::rest_parameters::{
     resolve_log_query_parameters, resolve_log_query_text,
 };
 use crate::config::retry::RetryPolicy;
+use crate::diagnostic::DiagnosticRunner;
 use crate::handle_manager::Handle;
 use crate::rest::snowflake::{
     self, QueryExecutionMode, QueryInput, RestError, SessionTokens, SnowflakeResponseError,
@@ -42,6 +43,18 @@ use crate::rest::snowflake::{
 use crate::sensitive::SensitiveString;
 use crate::token_cache::TokenCache;
 
+/// Whether `execute_session_sql` should refresh the connection's local
+/// session-state cache (`session_parameters`, `final_session_names`) from
+/// the response after a successful query.
+#[derive(Copy, Clone, Debug)]
+enum SessionStateRefresh {
+    /// SQL may change session state (e.g. `ALTER SESSION SET`, `USE DATABASE`).
+    Apply,
+    /// SQL is non-stateful from the session's perspective (e.g. `COMMIT`,
+    /// `ROLLBACK`).
+    Skip,
+}
+
 impl DatabaseDriverV1 {
     /// Set autocommit on the given connection.
     ///
@@ -62,7 +75,8 @@ impl DatabaseDriverV1 {
                     } else {
                         "ALTER SESSION SET AUTOCOMMIT = FALSE"
                     };
-                    self.execute_session_sql(&conn, sql).await
+                    self.execute_session_sql(&conn, sql, SessionStateRefresh::Apply)
+                        .await
                 } else {
                     conn.init_session_parameters
                         .get_or_insert_with(HashMap::new)
@@ -92,8 +106,7 @@ impl DatabaseDriverV1 {
             }
             .fail();
         }
-        // TODO: ensure that this is correct escaping
-        let escaped = db.replace('"', "\"\"");
+        let escaped = escape_sql_identifier(db);
         let sql = format!("USE DATABASE \"{escaped}\"");
 
         match self.connections.get_obj(conn_handle) {
@@ -106,7 +119,89 @@ impl DatabaseDriverV1 {
                     }
                     .fail();
                 }
-                self.execute_session_sql(&conn, &sql).await
+                self.execute_session_sql(&conn, &sql, SessionStateRefresh::Apply)
+                    .await
+            }
+            None => InvalidArgumentSnafu {
+                argument: "Connection handle not found".to_string(),
+            }
+            .fail(),
+        }
+    }
+
+    /// Commit the open transaction on the given connection by executing `COMMIT`.
+    /// Must only be called after the connection is initialised (`is_post_connect()`).
+    pub async fn connection_commit(&self, conn_handle: Handle) -> Result<(), ApiError> {
+        self.execute_transaction_control(conn_handle, "COMMIT")
+            .await
+    }
+
+    /// Roll back the open transaction on the given connection by executing `ROLLBACK`.
+    /// Must only be called after the connection is initialised (`is_post_connect()`).
+    pub async fn connection_rollback(&self, conn_handle: Handle) -> Result<(), ApiError> {
+        self.execute_transaction_control(conn_handle, "ROLLBACK")
+            .await
+    }
+
+    /// Shared implementation for `connection_commit` / `connection_rollback`:
+    /// run the transaction-control statement on the connection's session.
+    async fn execute_transaction_control(
+        &self,
+        conn_handle: Handle,
+        sql: &str,
+    ) -> Result<(), ApiError> {
+        match self.connections.get_obj(conn_handle) {
+            Some(conn_ptr) => {
+                let conn = conn_ptr.lock().await;
+                if !conn.is_post_connect() {
+                    return InvalidArgumentSnafu {
+                        argument: format!("{sql} called before connection is open"),
+                    }
+                    .fail();
+                }
+                // COMMIT/ROLLBACK do not change session parameters or
+                // current database/schema/warehouse/role, so skip the
+                // session-state cache refresh.
+                self.execute_session_sql(&conn, sql, SessionStateRefresh::Skip)
+                    .await
+            }
+            None => InvalidArgumentSnafu {
+                argument: "Connection handle not found".to_string(),
+            }
+            .fail(),
+        }
+    }
+
+    /// Execute `USE SCHEMA "<name>"` or `USE SCHEMA "<db>"."<name>"` when a session database is set.
+    /// Schema and database names are escaped (internal `"` doubled).
+    /// Must only be called after the connection is initialised (`is_post_connect()`).
+    pub async fn connection_use_schema(
+        &self,
+        conn_handle: Handle,
+        schema: &str,
+    ) -> Result<(), ApiError> {
+        let schema = schema.trim();
+        if schema.is_empty() {
+            return InvalidArgumentSnafu {
+                argument: "schema name must not be empty".to_string(),
+            }
+            .fail();
+        }
+
+        match self.connections.get_obj(conn_handle) {
+            Some(conn_ptr) => {
+                let conn = conn_ptr.lock().await;
+                if !conn.is_post_connect() {
+                    return InvalidArgumentSnafu {
+                        argument: "connection_use_schema called before connection is open"
+                            .to_string(),
+                    }
+                    .fail();
+                }
+                let database = resolve_session_database(&conn)?;
+                let sql = build_use_schema_sql(database.as_deref(), schema);
+                self.execute_session_sql(&conn, &sql, SessionStateRefresh::Apply)
+                    .await
             }
             None => InvalidArgumentSnafu {
                 argument: "Connection handle not found".to_string(),
@@ -116,7 +211,20 @@ impl DatabaseDriverV1 {
     }
 
     /// Execute a session-scoped SQL command without creating a statement.
-    async fn execute_session_sql(&self, conn: &Connection, sql: &str) -> Result<(), ApiError> {
+    ///
+    /// `refresh` controls whether the local session-state cache
+    /// (`session_parameters`, `final_session_names`) is refreshed from the
+    /// response after a successful query. Pass `Apply` for SQL that may
+    /// change session state (e.g. `ALTER SESSION SET`, `USE DATABASE`); pass
+    /// `Skip` for SQL that is non-stateful from the session's perspective
+    /// (e.g. `COMMIT`, `ROLLBACK`) to avoid taking the cache write locks
+    /// for a no-op merge.
+    async fn execute_session_sql(
+        &self,
+        conn: &Connection,
+        sql: &str,
+        refresh: SessionStateRefresh,
+    ) -> Result<(), ApiError> {
         let query_input = QueryInput {
             sql: sql.to_string(),
             bindings: None,
@@ -150,7 +258,7 @@ impl DatabaseDriverV1 {
             }
         }?;
 
-        if response.success {
+        if response.success && matches!(refresh, SessionStateRefresh::Apply) {
             conn.update_session_params_cache(
                 sql,
                 response.data.parameters.as_ref(),
@@ -265,6 +373,24 @@ impl DatabaseDriverV1 {
                     read_spcs_token(self.fs_adapter().as_ref()),
                 );
 
+                // ---- Diagnostics: pre-connect -----------------------------------
+                let mut diag_runner = if let DiagnosticConfig::Enabled { .. } = config.diagnostic {
+                    let account = config.server.account.clone();
+                    // resolve_options() already ran derive_host_from_account, so host
+                    // is Some whenever account is set; the empty-string fallback is unreachable.
+                    let host_str = host.clone().unwrap_or_default();
+                    let diag_cfg = config.diagnostic.clone();
+                    tokio::task::spawn_blocking(move || {
+                        let mut runner = DiagnosticRunner::new(&account, &host_str, diag_cfg);
+                        runner.run_pre_connect();
+                        runner
+                    })
+                    .await
+                    .ok()
+                } else {
+                    None
+                };
+
                 let token_caching_requested = matches!(
                     &login_parameters.login_method,
                     LoginMethod::UserPasswordMfa {
@@ -304,8 +430,19 @@ impl DatabaseDriverV1 {
                     token_cache.map(|c| c as &dyn TokenCache),
                     Some(&self.prompt_locks),
                 )
-                .await
-                .context(LoginSnafu)?;
+                .await;
+
+                // ---- Diagnostics: post-connect ----------------------------------
+                if let Some(mut runner) = diag_runner.take() {
+                    tokio::task::spawn_blocking(move || {
+                        runner.run_post_connect(None);
+                        runner.write_report();
+                    })
+                    .await
+                    .ok();
+                }
+
+                let login_result = login_result.context(LoginSnafu)?;
 
                 // Initialize connection with session parameters from login response.
                 // The server returns system-level parameters but may not echo back
@@ -1465,6 +1602,34 @@ fn get_session_or_setting(
     resolved_or_seed_string(conn, setting_key)
 }
 
+// TODO: ensure that this is correct escaping
+fn escape_sql_identifier(name: &str) -> String {
+    name.replace('"', "\"\"")
+}
+
+fn build_use_schema_sql(database: Option<&str>, schema: &str) -> String {
+    let escaped_schema = escape_sql_identifier(schema);
+    match database.filter(|db| !db.is_empty()) {
+        Some(db) => {
+            let escaped_db = escape_sql_identifier(db);
+            format!("USE SCHEMA \"{escaped_db}\".\"{escaped_schema}\"")
+        }
+        None => format!("USE SCHEMA \"{escaped_schema}\""),
+    }
+}
+
+fn resolve_session_database(conn: &Connection) -> Result<Option<String>, ApiError> {
+    let final_names = conn
+        .final_session_names
+        .read()
+        .map_err(|_| ConnectionLockingSnafu {}.build())?;
+    Ok(final_names
+        .database
+        .clone()
+        .or_else(|| get_session_or_setting(conn, "DATABASE", param_names::DATABASE))
+        .filter(|db| !db.is_empty()))
+}
+
 impl DatabaseDriverV1 {
     /// Get connection information for the given connection handle
     pub async fn connection_get_info(
@@ -2187,6 +2352,30 @@ mod tests {
         assert_eq!(get_connection_seed_string(&conn, param_names::PORT), None);
     }
 
+    #[test]
+    fn build_use_schema_sql_without_database() {
+        assert_eq!(
+            build_use_schema_sql(None, "PUBLIC"),
+            "USE SCHEMA \"PUBLIC\""
+        );
+    }
+
+    #[test]
+    fn build_use_schema_sql_with_database() {
+        assert_eq!(
+            build_use_schema_sql(Some("TEST_DB"), "MY_SCHEMA"),
+            "USE SCHEMA \"TEST_DB\".\"MY_SCHEMA\""
+        );
+    }
+
+    #[test]
+    fn build_use_schema_sql_escapes_quotes() {
+        assert_eq!(
+            build_use_schema_sql(Some("DB\"1"), "SC\"2"),
+            "USE SCHEMA \"DB\"\"1\".\"SC\"\"2\""
+        );
+    }
+
     #[test]
     fn get_session_or_setting_prefers_session_parameter() {
         let conn =