Tests: Snowflake and Databricks error and lifecycle paths (from #1218)

crates/utopia-server/src/query_engine/databricks.rs

@@ -847,4 +847,115 @@ 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);
+    }
 }

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}");
+    }
 }