|
| 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