Skip to content

Commit efff49d

Browse files
authored
feat: make set/show a hook implementation (#209)
* feat: make set/show a hook implementation * feat: implment extended query part * refactor: avoid duplicate testing code * chore: resolve lint issue
1 parent fd20234 commit efff49d

File tree

10 files changed

+488
-401
lines changed

10 files changed

+488
-401
lines changed

datafusion-postgres/src/client.rs

Lines changed: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,54 @@
1+
use pgwire::api::ClientInfo;
2+
3+
// Metadata keys for session-level settings
4+
const METADATA_STATEMENT_TIMEOUT: &str = "statement_timeout_ms";
5+
const METADATA_TIMEZONE: &str = "timezone";
6+
7+
/// Get statement timeout from client metadata
8+
pub fn get_statement_timeout<C>(client: &C) -> Option<std::time::Duration>
9+
where
10+
C: ClientInfo + ?Sized,
11+
{
12+
client
13+
.metadata()
14+
.get(METADATA_STATEMENT_TIMEOUT)
15+
.and_then(|s| s.parse::<u64>().ok())
16+
.map(std::time::Duration::from_millis)
17+
}
18+
19+
/// Set statement timeout in client metadata
20+
pub fn set_statement_timeout<C>(client: &mut C, timeout: Option<std::time::Duration>)
21+
where
22+
C: ClientInfo + ?Sized,
23+
{
24+
let metadata = client.metadata_mut();
25+
if let Some(duration) = timeout {
26+
metadata.insert(
27+
METADATA_STATEMENT_TIMEOUT.to_string(),
28+
duration.as_millis().to_string(),
29+
);
30+
} else {
31+
metadata.remove(METADATA_STATEMENT_TIMEOUT);
32+
}
33+
}
34+
35+
/// Get statement timeout from client metadata
36+
pub fn get_timezone<C>(client: &C) -> Option<&str>
37+
where
38+
C: ClientInfo + ?Sized,
39+
{
40+
client.metadata().get(METADATA_TIMEZONE).map(|s| s.as_str())
41+
}
42+
43+
/// Set statement timeout in client metadata
44+
pub fn set_timezone<C>(client: &mut C, timezone: Option<&str>)
45+
where
46+
C: ClientInfo + ?Sized,
47+
{
48+
let metadata = client.metadata_mut();
49+
if let Some(timezone) = timezone {
50+
metadata.insert(METADATA_TIMEZONE.to_string(), timezone.to_string());
51+
} else {
52+
metadata.remove(METADATA_TIMEZONE);
53+
}
54+
}

0 commit comments

Comments
 (0)