feat: http handler return geometry_output_format with data. (#18963)

Author: youngsofun

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

@@ -12,7 +12,6 @@
 // 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;
 
@@ -142,6 +141,14 @@ impl QueryResponseField {
     }
 }
 
+// settings also used by driver, may be set in query/session/global level
+// only available after binding
+#[derive(Serialize, Debug, Clone)]
+pub struct ResultFormatSettings {
+    pub timezone: String,
+    pub geometry_output_format: String,
+}
+
 #[derive(Serialize, Debug, Clone)]
 pub struct QueryResponse {
     pub id: String,
@@ -161,11 +168,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 query/session/global level
-    // only include timezone for now
-    // only available after binding
+
     #[serde(skip_serializing_if = "Option::is_none")]
-    pub settings: Option<BTreeMap<String, String>>,
+    pub settings: Option<ResultFormatSettings>,
 
     pub stats: QueryStats,
 
@@ -195,7 +200,7 @@ impl QueryResponse {
                     affect,
                     error,
                     warnings,
-                    driver_settings,
+                    result_format_settings: driver_settings,
                 },
         }: HttpQueryResponseInternal,
         is_final: bool,

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

@@ -12,7 +12,6 @@
 // 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;
@@ -45,6 +44,7 @@ use crate::interpreters::interpreter_plan_sql;
 use crate::interpreters::Interpreter;
 use crate::interpreters::InterpreterFactory;
 use crate::interpreters::InterpreterQueryLog;
+use crate::servers::http::v1::http_query_handlers::ResultFormatSettings;
 use crate::sessions::AcquireQueueGuard;
 use crate::sessions::QueryAffect;
 use crate::sessions::QueryContext;
@@ -130,15 +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>>,
+    result_format_settings: Option<ResultFormatSettings>,
     has_result_set: bool,
     #[allow(dead_code)]
     queue_guard: AcquireQueueGuard,
 }
 
 pub struct ExecuteStopped {
     pub schema: DataSchemaRef,
-    pub driver_settings: Option<BTreeMap<String, String>>,
+    pub result_format_settings: Option<ResultFormatSettings>,
     pub has_result_set: Option<bool>,
     pub stats: Progresses,
     pub affect: Option<QueryAffect>,
@@ -200,18 +200,18 @@ impl Executor {
             Stopped(f) => f.schema.clone(),
         };
 
-        let driver_settings = match &self.state {
+        let result_format_settings = match &self.state {
             Starting(_) => None,
-            Running(r) => r.driver_settings.clone(),
-            Stopped(f) => f.driver_settings.clone(),
+            Running(r) => r.result_format_settings.clone(),
+            Stopped(f) => f.result_format_settings.clone(),
         };
 
         ResponseState {
             running_time_ms: self.get_query_duration_ms(),
             progresses: self.get_progress(),
             state: exe_state,
             error: err,
-            driver_settings,
+            result_format_settings,
             warnings: self.get_warnings(),
             affect: self.get_affect(),
             schema,
@@ -319,7 +319,7 @@ impl Executor {
                 ExecuteStopped {
                     stats: Default::default(),
                     schema: Default::default(),
-                    driver_settings: None,
+                    result_format_settings: None,
                     has_result_set: None,
                     reason: reason.clone(),
                     session_state: ExecutorSessionState::new(s.ctx.get_current_session()),
@@ -342,7 +342,7 @@ impl Executor {
                 ExecuteStopped {
                     stats: Progresses::from_context(&r.ctx),
                     schema: r.schema.clone(),
-                    driver_settings: r.driver_settings.clone(),
+                    result_format_settings: r.result_format_settings.clone(),
                     has_result_set: Some(r.has_result_set),
                     reason: reason.clone(),
                     session_state: ExecutorSessionState::new(r.ctx.get_current_session()),
@@ -401,15 +401,23 @@ 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 settings = ctx.get_settings();
+        let timezone = settings.get_timezone().with_context(make_error)?;
+        let geometry_output_format = settings
+            .get_geometry_output_format()
+            .with_context(make_error)?
+            .to_string();
+        let result_format_settings = Some(ResultFormatSettings {
+            timezone,
+            geometry_output_format,
+        });
 
         let running_state = ExecuteRunning {
             session,
             ctx: ctx.clone(),
             queue_guard,
             schema,
-            driver_settings,
+            result_format_settings,
             has_result_set,
         };
         info!("[HTTP-QUERY] Query state changed to Running");

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

@@ -67,6 +67,7 @@ use super::PageManager;
 use super::ResponseData;
 use super::Wait;
 use crate::servers::http::error::QueryError;
+use crate::servers::http::v1::http_query_handlers::ResultFormatSettings;
 use crate::servers::http::v1::ClientSessionManager;
 use crate::servers::http::v1::HttpQueryManager;
 use crate::servers::http::v1::QueryResponse;
@@ -491,7 +492,7 @@ pub struct StageAttachmentConf {
 pub struct ResponseState {
     pub has_result_set: Option<bool>,
     pub schema: DataSchemaRef,
-    pub driver_settings: Option<BTreeMap<String, String>>,
+    pub result_format_settings: Option<ResultFormatSettings>,
     pub running_time_ms: i64,
     pub progresses: Progresses,
     pub state: ExecuteStateKind,
@@ -848,7 +849,7 @@ impl HttpQuery {
                     let state = ExecuteStopped {
                         stats: Progresses::default(),
                         schema: Default::default(),
-                        driver_settings: None,
+                        result_format_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,22 +44,3 @@ 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

@@ -8,6 +8,7 @@
 STICKY_HEADER = "X-DATABEND-STICKY-NODE"
 
 timezone = 'Asia/Shanghai'
+geometry_output_format = 'EWKB'
 
 scan_progress = "scan_progress"
 
@@ -25,7 +26,8 @@ def do_query(query, timeout=10, pagination=None, port=8000, patch=True):
     session = {
         "settings": {
             "http_handler_result_timeout_secs": f"{timeout}",
-            "timezone": f"{timezone}"
+            "timezone": f"{timezone}",
+            "geometry_output_format": f"{geometry_output_format}"
         }
     }
 
@@ -154,15 +156,17 @@ 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'},
+           'settings': {'timezone': 'Asia/Shanghai',
+                        "geometry_output_format": f"{geometry_output_format}" },
            '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}',
-                                    'timezone': f"{timezone}"},
+                                    'timezone': f"{timezone}",
+                                    "geometry_output_format": f"{geometry_output_format}" },
                        'txn_state': 'AutoCommit'}
            }
 

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

@@ -69,3 +69,28 @@ def drain_reader(reader):
 
     rows.sort()
     assert rows == [x for x in range(97)]
+
+def query_page_0(sql):
+    query_url = "http://localhost:8000/v1/query"
+
+    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_result_format_settings():
+    timezone = 'Asia/Shanghai'
+    r = query_page_0(f"settings (timezone='{timezone}') select 1")
+    assert r.get('settings', {}).get('timezone') == timezone
+
+    geometry_output_format = 'EWKB'
+    r = query_page_0(f"settings (geometry_output_format='{geometry_output_format}') select 1")
+    assert r.get('settings', {}).get('geometry_output_format') == geometry_output_format