Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
17 changes: 11 additions & 6 deletions src/query/service/src/servers/http/v1/http_query_handlers.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand Down Expand Up @@ -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,
Expand All @@ -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,

Expand Down Expand Up @@ -195,7 +200,7 @@ impl QueryResponse {
affect,
error,
warnings,
driver_settings,
result_format_settings: driver_settings,
},
}: HttpQueryResponseInternal,
is_final: bool,
Expand Down
32 changes: 20 additions & 12 deletions src/query/service/src/servers/http/v1/query/execute_state.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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>,
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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()),
Expand All @@ -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()),
Expand Down Expand Up @@ -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");
Expand Down
5 changes: 3 additions & 2 deletions src/query/service/src/servers/http/v1/query/http_query.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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(
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@
STICKY_HEADER = "X-DATABEND-STICKY-NODE"

timezone = 'Asia/Shanghai'
geometry_output_format = 'EWKB'

scan_progress = "scan_progress"

Expand All @@ -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}"
}
}

Expand Down Expand Up @@ -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'}
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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
Loading