feat: return timezone in HTTP handler. (#18936)

* feat: return timezone in HTTP handler.

* fix flaky test.

Author: youngsofun

src/query/service/src/servers/http/v1/http_query_handlers.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;
 
@@ -160,6 +161,11 @@ 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
+    // only include timezone for now
+    // only is_some in the initial response
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub settings: Option<BTreeMap<String, String>>,
 
     pub stats: QueryStats,
 
@@ -193,6 +199,7 @@ impl QueryResponse {
         }: 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)
@@ -257,6 +264,7 @@ impl QueryResponse {
             error: error.map(QueryError::from_error_code),
             has_result_set,
             result_timeout_secs: Some(result_timeout_secs),
+            settings,
         };
 
         match body_format {
@@ -389,6 +397,7 @@ async fn query_final_handler(
                     response,
                     true,
                     body_format,
+                    None,
                 ))
             }
             None => Err(query_id_not_found(&query_id, &ctx.node_id)),
@@ -519,6 +528,7 @@ async fn query_page_handler(
                     resp,
                     false,
                     body_format,
+                    None,
                 ))
             }
         }
@@ -577,6 +587,14 @@ 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);
@@ -612,10 +630,14 @@ 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)
-                        .into_response(),
+                Ok(QueryResponse::from_internal(
+                    query.id.to_string(),
+                    resp,
+                    false,
+                    body_format,
+                    driver_settings,
                 )
+                .into_response())
             }
         }
     };

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

@@ -128,6 +128,7 @@ impl HttpQueryRequest {
             error: Some(QueryError::from_error_code(err)),
             has_result_set: None,
             result_timeout_secs: None,
+            settings: None,
         })
     }
 
@@ -325,6 +326,7 @@ pub struct HttpSessionConf {
     pub role: Option<String>,
     #[serde(skip_serializing_if = "Option::is_none")]
     pub secondary_roles: Option<Vec<String>>,
+    // session level settings
     #[serde(skip_serializing_if = "Option::is_none")]
     pub settings: Option<BTreeMap<String, String>>,
     #[serde(skip_serializing_if = "Option::is_none")]
@@ -569,7 +571,7 @@ pub struct HttpQuery {
     // cache only
     pub(crate) result_timeout_secs: u64,
 
-    execute_state: Arc<Mutex<Executor>>,
+    pub(crate) execute_state: Arc<Mutex<Executor>>,
 
     // client states
     client_state: Arc<Mutex<ClientState>>,

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

@@ -7,7 +7,7 @@
 auth = ("root", "")
 STICKY_HEADER = "X-DATABEND-STICKY-NODE"
 
-max_threads = 32
+timezone = 'Asia/Shanghai'
 
 def patch_json(resp):
     if resp.get("stats", {}).get("running_time_ms"):
@@ -19,7 +19,7 @@ def do_query(query, timeout=10, pagination=None, port=8000, patch=True):
     session = {
         "settings": {
             "http_handler_result_timeout_secs": f"{timeout}",
-            "max_threads": f"{max_threads}"
+            "timezone": f"{timezone}"
         }
     }
 
@@ -118,7 +118,7 @@ def test_query_lifecycle_finalized(rows):
         "max_rows_in_buffer": max_rows_per_page,
     }
 
-    resp0 = do_query(f"select * from numbers({rows})", timeout=timeout, pagination=pagination)
+    resp0 = do_query(f"select number from numbers({rows}) order by number", timeout=timeout, pagination=pagination)
 
     query_id = resp0.get("id")
     node_id = resp0.get("node_id")
@@ -148,14 +148,15 @@ def test_query_lifecycle_finalized(rows):
            "final_uri": f"/v1/query/{query_id}/final",
            "next_uri": f"/v1/query/{query_id}/page/1",
            "kill_uri": f"/v1/query/{query_id}/kill",
+           'settings': {'timezone': 'Asia/Shanghai'},
            'session': {'catalog': 'default',
                        'database': 'default',
                        'internal': sessions_internal,
                        'need_keep_alive': False,
                        'need_sticky': False,
                        'role': 'account_admin',
                        'settings': {'http_handler_result_timeout_secs': f'{timeout}',
-                                    'max_threads': f"{max_threads}"},
+                                    'timezone': f"{timezone}"},
                        'txn_state': 'AutoCommit'}
            }
 
@@ -184,6 +185,7 @@ 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"