From 0f0bdd980e6c5cfef2e0510319ba857c75e7b348 Mon Sep 17 00:00:00 2001 From: Royce Lam Date: Tue, 13 Oct 2026 22:42:26 +0800 Subject: [PATCH 1/2] Tests: Snowflake and Databricks error / lifecycle / edge paths via wiremock Signed-off-by: Royce Lam --- .../src/query_engine/databricks.rs | 110 ++++++++++++++++++ .../src/query_engine/snowflake.rs | 74 ++++++++++++ 2 files changed, 184 insertions(+) diff --git a/crates/utopia-server/src/query_engine/databricks.rs b/crates/utopia-server/src/query_engine/databricks.rs index 47e8dd800..56d86ee81 100644 --- a/crates/utopia-server/src/query_engine/databricks.rs +++ b/crates/utopia-server/src/query_engine/databricks.rs @@ -847,4 +847,114 @@ mod tests { assert!(err.contains("schema read returned 1 rows"), "{err}"); assert!(err.contains("99"), "{err}"); } + + /// 401/403 一类：token 失效或 PAT 撤销。`error_for_status()` 会把非 2xx + /// 抛成 reqwest 错误——表现上等价「连接失败 / 鉴权失败」。把这层向上层送 + /// 调用者时仍可分辨：「reqwest」还是业务 message 的边界 + #[tokio::test] + async fn an_unauthorized_http_status_is_surfaced_as_a_connection_error() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/api/2.0/sql/statements")) + .respond_with(ResponseTemplate::new(401).set_body_string("Unauthorized")) + .mount(&server) + .await; + let err = DatabricksEngine::new(conn(&server)) + .execute("SELECT 1") + .await + .unwrap_err() + .to_string(); + // reqwest 的 error_for_status 把 status 包成 builder 错误；我们要确保 + // 调用者看到的是 HTTP 401 字样，不是「Snowflake: …」的错形态 + assert!(err.contains("401"), "{err}"); + } + + /// 5xx 一类（warehouse suspended / cluster down）走同一条 error_for_status + /// 路径——表面上仍是 reqwest 错误，需要可识别的 status code + #[tokio::test] + async fn a_5xx_response_is_surfaced_with_status_code() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/api/2.0/sql/statements")) + .respond_with(ResponseTemplate::new(503).set_body_string( + "Service Unavailable — warehouse suspended", + )) + .mount(&server) + .await; + let err = DatabricksEngine::new(conn(&server)) + .execute("SELECT 1") + .await + .unwrap_err() + .to_string(); + assert!(err.contains("503"), "{err}"); + } + + /// 用户主动取消（cancel button 或 stmt cancellation）是 Databricks 的合法 + /// 状态：`CANCELED`。业务错误带 message 后应能让上层识别这是「取消」 + /// 而不是「失败」，所以 message 应完整透传 + #[tokio::test] + async fn a_canceled_statement_carries_a_distinct_message() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/api/2.0/sql/statements")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "statement_id": "s3", + "status": { "state": "CANCELED", "error": { "error_code": "CANCELED", "message": "Statement was canceled by user." } } + }))) + .mount(&server) + .await; + let err = DatabricksEngine::new(conn(&server)) + .execute("SELECT 1") + .await + .unwrap_err() + .to_string(); + assert!(err.contains("CANCELED"), "{err}"); + assert!(err.contains("canceled by user"), "{err}"); + } + + /// PENDING 响应里若缺失 `statement_id`，现有逻辑会以「Databricks returned + /// no statement_id」抛出——把这条兜底钉住，免得一改 polling 就把它吞掉 + #[tokio::test] + async fn pending_state_without_statement_id_is_surfaced() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/api/2.0/sql/statements")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "status": { "state": "PENDING" } + }))) + .mount(&server) + .await; + let err = DatabricksEngine::new(conn(&server)) + .execute("SELECT 1") + .await + .unwrap_err() + .to_string(); + assert!( + err.contains("no statement_id"), + "PENDING 但没 statement_id 应被显式提示，避免 polling 失效时不报错：{err}" + ); + } + + /// `SUCCEEDED` 但 `result.data_array` 是空（0 行查询 / DDL 走过 execute 路径）， + /// rows 取 []、columns 也取空向量——`rows_to_json_lines` 仍要返回空集、不崩 + #[tokio::test] + async fn a_succeeded_statement_with_no_rows_returns_empty() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/api/2.0/sql/statements")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "statement_id": "s4", + "status": { "state": "SUCCEEDED" }, + "manifest": { "schema": { "columns": [] } }, + "result": { "data_array": [] } + }))) + .mount(&server) + .await; + let out = DatabricksEngine::new(conn(&server)) + .execute("SELECT 1 WHERE 1=0") + .await + .unwrap(); + assert!(out.rows.is_empty(), "{out:?}"); + assert!(!out.truncated); + } } diff --git a/crates/utopia-server/src/query_engine/snowflake.rs b/crates/utopia-server/src/query_engine/snowflake.rs index b9052313a..dfbe0c865 100644 --- a/crates/utopia-server/src/query_engine/snowflake.rs +++ b/crates/utopia-server/src/query_engine/snowflake.rs @@ -491,4 +491,78 @@ mod tests { assert!(err.contains(expected_error), "{err}"); } } + + /// 401/403 一类（auth、权限）也是 status 路径——message 字段照样送出。 + /// 与 422 的区别仅在状态码，调用者拿到的错误形态应保持一致（都是 + /// `Snowflake {status}: {message}`），便于上层按错误形态分类 + #[tokio::test] + async fn auth_failure_body_is_surfaced() { + for status in [401u16, 403, 503] { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/api/v2/statements")) + .respond_with(ResponseTemplate::new(status).set_body_json(json!({ + "code": "390100", + "message": format!("HTTP {status}: not authorized — token expired or missing role"), + }))) + .mount(&server) + .await; + let err = SnowflakeEngine::new(conn(&server)) + .execute("SELECT 1") + .await + .unwrap_err() + .to_string(); + assert!(err.contains(&status.to_string()), "{err}"); + assert!(err.contains("token expired"), "{err}"); + } + } + + /// 5xx 一类（如 Snowflake 内部 504 / 503 + GW 错配）时 body 不是 JSON—— + /// 现有逻辑会尝试 JSON.parse 然后退回原文，所以「消息文案」仍能送达 + /// 调用者。把这层兜底路径钉住，免得被人「美化」掉 + #[tokio::test] + async fn non_json_error_body_falls_back_to_plain_text() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/api/v2/statements")) + .respond_with(ResponseTemplate::new(500).set_body_string( + "Internal error: warehouse suspended\n", + )) + .mount(&server) + .await; + let err = SnowflakeEngine::new(conn(&server)) + .execute("SELECT 1") + .await + .unwrap_err() + .to_string(); + assert!(err.contains("500"), "{err}"); + assert!( + err.contains("warehouse suspended"), + "非 JSON body 应原样透传，给上层一份能落库的描述：{err}" + ); + } + + /// 2xx body 里若带非「成功」code（如 query canceled 取消信号），必须以业务 + /// 错误形态抛出，不让屏蔽在「看似 OK」的响应里——这是 Snowflake 的双层错误 + /// 模型，应当与 4xx/5xx 共用同一条 message 通道 + #[tokio::test] + async fn business_error_code_in_a_2xx_body_is_bubbled_up() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/api/v2/statements")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "code": "000605", + "message": "Statement was canceled by user.", + // 没有 resultSetMetaData，与状态一致——成功响应必有 meta 或 data + }))) + .mount(&server) + .await; + let err = SnowflakeEngine::new(conn(&server)) + .execute("SELECT 1") + .await + .unwrap_err() + .to_string(); + assert!(err.contains("000605"), "{err}"); + assert!(err.contains("canceled"), "{err}"); + } } From 28b0468bc45f2a054ae4e6fc534269e58af709a9 Mon Sep 17 00:00:00 2001 From: Wayland Yang Date: Wed, 14 Oct 2026 16:37:24 +0800 Subject: [PATCH 2/2] Format the Snowflake and Databricks error-path tests Co-Authored-By: Claude Opus 5.5 Signed-off-by: Wayland Yang --- crates/utopia-server/src/query_engine/databricks.rs | 7 ++++--- crates/utopia-server/src/query_engine/snowflake.rs | 6 +++--- 2 files changed, 7 insertions(+), 6 deletions(-) diff --git a/crates/utopia-server/src/query_engine/databricks.rs b/crates/utopia-server/src/query_engine/databricks.rs index 56d86ee81..f18276137 100644 --- a/crates/utopia-server/src/query_engine/databricks.rs +++ b/crates/utopia-server/src/query_engine/databricks.rs @@ -876,9 +876,10 @@ mod tests { let server = MockServer::start().await; Mock::given(method("POST")) .and(path("/api/2.0/sql/statements")) - .respond_with(ResponseTemplate::new(503).set_body_string( - "Service Unavailable — warehouse suspended", - )) + .respond_with( + ResponseTemplate::new(503) + .set_body_string("Service Unavailable — warehouse suspended"), + ) .mount(&server) .await; let err = DatabricksEngine::new(conn(&server)) diff --git a/crates/utopia-server/src/query_engine/snowflake.rs b/crates/utopia-server/src/query_engine/snowflake.rs index dfbe0c865..dce0e707c 100644 --- a/crates/utopia-server/src/query_engine/snowflake.rs +++ b/crates/utopia-server/src/query_engine/snowflake.rs @@ -525,9 +525,9 @@ mod tests { let server = MockServer::start().await; Mock::given(method("POST")) .and(path("/api/v2/statements")) - .respond_with(ResponseTemplate::new(500).set_body_string( - "Internal error: warehouse suspended\n", - )) + .respond_with( + ResponseTemplate::new(500).set_body_string("Internal error: warehouse suspended\n"), + ) .mount(&server) .await; let err = SnowflakeEngine::new(conn(&server))