fix: return timezone when set in query level. (#18952)

* ci: drop table at end of test

* ci: support ttc_dev for developing drivers.

* fix: return timezone when set in query level.

Author: youngsofun

src/query/service/src/servers/http/v1/http_query_handlers.rs

@@ -161,9 +161,9 @@ pub struct QueryResponse {
     pub data: Arc<BlocksSerializer>,
     pub affect: Option<QueryAffect>,
     pub result_timeout_secs: Option<u64>,
-    // settings also used by driver, may be set in session or global
+    // settings also used by driver, may be set in query/session/global level
     // only include timezone for now
-    // only is_some in the initial response
+    // only available after binding
     #[serde(skip_serializing_if = "Option::is_none")]
     pub settings: Option<BTreeMap<String, String>>,
 
@@ -195,11 +195,11 @@ impl QueryResponse {
                     affect,
                     error,
                     warnings,
+                    driver_settings,
                 },
         }: HttpQueryResponseInternal,
         is_final: bool,
         body_format: BodyFormat,
-        settings: Option<BTreeMap<String, String>>,
     ) -> Response {
         let (data, next_uri) = if is_final {
             (Arc::new(BlocksSerializer::empty()), None)
@@ -264,7 +264,7 @@ impl QueryResponse {
             error: error.map(QueryError::from_error_code),
             has_result_set,
             result_timeout_secs: Some(result_timeout_secs),
-            settings,
+            settings: driver_settings,
         };
 
         match body_format {
@@ -397,7 +397,6 @@ async fn query_final_handler(
                     response,
                     true,
                     body_format,
-                    None,
                 ))
             }
             None => Err(query_id_not_found(&query_id, &ctx.node_id)),
@@ -528,7 +527,6 @@ async fn query_page_handler(
                     resp,
                     false,
                     body_format,
-                    None,
                 ))
             }
         }
@@ -587,14 +585,6 @@ pub(crate) async fn query_handler(
                 Ok(req.fail_to_start_sql(err).into_response())
             }
             Ok(mut query) => {
-                let settings = query.execute_state.lock().get_session_state().settings;
-                let tz = settings.get_timezone().unwrap();
-                let driver_settings = if tz == "UTC" {
-                    None
-                } else {
-                    Some(BTreeMap::from_iter([("timezone".to_string(), tz)]))
-                };
-
                 if let Err(err) = query.start_query(sql.clone()).await {
                     let err = err.display_with_sql(&sql);
                     error!("[HTTP-QUERY] Failed to start SQL query, error: {:?}", err);
@@ -630,14 +620,10 @@ pub(crate) async fn query_handler(
                 query
                     .update_expire_time(false, resp.is_data_drained())
                     .await;
-                Ok(QueryResponse::from_internal(
-                    query.id.to_string(),
-                    resp,
-                    false,
-                    body_format,
-                    driver_settings,
+                Ok(
+                    QueryResponse::from_internal(query.id.to_string(), resp, false, body_format)
+                        .into_response(),
                 )
-                .into_response())
             }
         }
     };

src/query/service/src/servers/http/v1/query/execute_state.rs

@@ -12,6 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use std::collections::BTreeMap;
 use std::collections::HashMap;
 use std::sync::Arc;
 use std::time::SystemTime;
@@ -129,13 +130,15 @@ pub struct ExecuteRunning {
     // mainly used to get progress for now
     pub(crate) ctx: Arc<QueryContext>,
     schema: DataSchemaRef,
+    driver_settings: Option<BTreeMap<String, String>>,
     has_result_set: bool,
     #[allow(dead_code)]
     queue_guard: AcquireQueueGuard,
 }
 
 pub struct ExecuteStopped {
     pub schema: DataSchemaRef,
+    pub driver_settings: Option<BTreeMap<String, String>>,
     pub has_result_set: Option<bool>,
     pub stats: Progresses,
     pub affect: Option<QueryAffect>,
@@ -197,11 +200,18 @@ impl Executor {
             Stopped(f) => f.schema.clone(),
         };
 
+        let driver_settings = match &self.state {
+            Starting(_) => None,
+            Running(r) => r.driver_settings.clone(),
+            Stopped(f) => f.driver_settings.clone(),
+        };
+
         ResponseState {
             running_time_ms: self.get_query_duration_ms(),
             progresses: self.get_progress(),
             state: exe_state,
             error: err,
+            driver_settings,
             warnings: self.get_warnings(),
             affect: self.get_affect(),
             schema,
@@ -309,6 +319,7 @@ impl Executor {
                 ExecuteStopped {
                     stats: Default::default(),
                     schema: Default::default(),
+                    driver_settings: None,
                     has_result_set: None,
                     reason: reason.clone(),
                     session_state: ExecutorSessionState::new(s.ctx.get_current_session()),
@@ -331,6 +342,7 @@ impl Executor {
                 ExecuteStopped {
                     stats: Progresses::from_context(&r.ctx),
                     schema: r.schema.clone(),
+                    driver_settings: r.driver_settings.clone(),
                     has_result_set: Some(r.has_result_set),
                     reason: reason.clone(),
                     session_state: ExecutorSessionState::new(r.ctx.get_current_session()),
@@ -389,12 +401,15 @@ impl ExecuteState {
         } else {
             Default::default()
         };
+        let tz = ctx.get_settings().get_timezone().with_context(make_error)?;
+        let driver_settings = Some(BTreeMap::from_iter([("timezone".to_string(), tz)]));
 
         let running_state = ExecuteRunning {
             session,
             ctx: ctx.clone(),
             queue_guard,
             schema,
+            driver_settings,
             has_result_set,
         };
         info!("[HTTP-QUERY] Query state changed to Running");

src/query/service/src/servers/http/v1/query/http_query.rs

@@ -491,6 +491,7 @@ pub struct StageAttachmentConf {
 pub struct ResponseState {
     pub has_result_set: Option<bool>,
     pub schema: DataSchemaRef,
+    pub driver_settings: Option<BTreeMap<String, String>>,
     pub running_time_ms: i64,
     pub progresses: Progresses,
     pub state: ExecuteStateKind,
@@ -847,6 +848,7 @@ impl HttpQuery {
                     let state = ExecuteStopped {
                         stats: Progresses::default(),
                         schema: Default::default(),
+                        driver_settings: None,
                         has_result_set: None,
                         reason: Err(e.clone()),
                         session_state: ExecutorSessionState::new(

tests/nox/suites/1_stateful/09_http_handler/test_09_0001_json_response.py

@@ -44,3 +44,22 @@ def test_json_response_errors():
     assert response3.json() == json.loads(
         '{"error": {"code": 404, "message": "not found"}}'
     )
+
+def query_page_0(sql):
+    data = {
+        'sql': sql,
+        "pagination": { "wait_time_secs": 5}
+    }
+    data = json.dumps(data)
+    response = requests.post(
+        query_url,
+        auth=auth,
+        headers={"Content-Type": "application/json"},
+        data=data,
+    )
+    return response.json()
+
+def test_timezone():
+    timezone = 'Asia/Shanghai'
+    r = query_page_0(f"settings (timezone='{timezone}') select 1")
+    assert r.get('settings', {}).get('timezone') == timezone

tests/nox/suites/1_stateful/09_http_handler/test_09_0014_query_lifecycle.py

@@ -191,7 +191,6 @@ def test_query_lifecycle_finalized(rows):
 
     # not return session since nothing changed
     exp["session"] = None
-    del exp["settings"] # only in the first resp
     exp["state"] = "Succeeded"
     if rows == 8:
         exp["next_uri"] = f"/v1/query/{query_id}/page/2"

tests/sqllogictests/src/main.rs

@@ -57,7 +57,6 @@ mod util;
 const HANDLER_MYSQL: &str = "mysql";
 const HANDLER_HTTP: &str = "http";
 const HANDLER_HYBRID: &str = "hybrid";
-const HANDLER_TTC_RUST: &str = "ttc_rust";
 const TTC_PORT_START: u16 = 9902;
 
 use std::sync::LazyLock;
@@ -136,10 +135,11 @@ pub async fn main() -> Result<()> {
             HANDLER_HYBRID => {
                 run_hybrid_client(args.clone(), &mut containers).await?;
             }
-            HANDLER_TTC_RUST => {}
             handler if handler.starts_with("ttc") => {
-                let image = format!("ghcr.io/databendlabs/{handler}:latest");
-                run_ttc_container(&image, TTC_PORT_START, args.port, &mut containers).await?;
+                if handler != "ttc_dev" {
+                    let image = format!("ghcr.io/databendlabs/{handler}:latest");
+                    run_ttc_container(&image, TTC_PORT_START, args.port, &mut containers).await?;
+                }
                 run_ttc_client(
                     args.clone(),
                     ClientType::Ttc(handler.to_string(), TTC_PORT_START),

tests/sqllogictests/suites/base/05_ddl/05_0037_ddl_dictionary.test

@@ -145,3 +145,6 @@ dd2 CREATE DICTIONARY dd2 ( c1 INT NOT NULL, c2 VARCHAR NOT NULL ) PRIMARY KEY c
 
 statement ok
 DROP DATABASE db1
+
+statement ok
+DROP DICTIONARY IF EXISTS default.dd2