agent_first_psql/
types.rs

1use serde::{Deserialize, Serialize};
2use serde_json::Value;
3use std::collections::HashMap;
4
5#[derive(Debug, Deserialize)]
6#[serde(tag = "code", deny_unknown_fields)]
7pub enum Input {
8    #[serde(rename = "query")]
9    Query {
10        id: String,
11        #[serde(default)]
12        session: Option<String>,
13        sql: String,
14        #[serde(default)]
15        params: Vec<Value>,
16        #[serde(default)]
17        options: QueryOptions,
18    },
19    #[serde(rename = "config")]
20    Config(ConfigPatch),
21    #[serde(rename = "cancel")]
22    Cancel { id: String },
23    #[serde(rename = "ping")]
24    Ping,
25    #[serde(rename = "close")]
26    Close,
27}
28
29#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
30pub enum Permission {
31    #[serde(rename = "read")]
32    Read,
33    #[serde(rename = "write")]
34    Write,
35    #[serde(rename = "ssh-read")]
36    SshRead,
37    #[serde(rename = "ssh-write")]
38    SshWrite,
39}
40
41impl Permission {
42    pub fn as_str(self) -> &'static str {
43        match self {
44            Self::Read => "read",
45            Self::Write => "write",
46            Self::SshRead => "ssh-read",
47            Self::SshWrite => "ssh-write",
48        }
49    }
50
51    pub fn is_read_only(self) -> bool {
52        matches!(self, Self::Read | Self::SshRead)
53    }
54
55    pub fn allows_ssh(self) -> bool {
56        matches!(self, Self::SshRead | Self::SshWrite)
57    }
58}
59
60impl std::str::FromStr for Permission {
61    type Err = String;
62
63    fn from_str(value: &str) -> Result<Self, Self::Err> {
64        match value {
65            "read" => Ok(Self::Read),
66            "write" => Ok(Self::Write),
67            "ssh-read" => Ok(Self::SshRead),
68            "ssh-write" => Ok(Self::SshWrite),
69            _ => Err(format!(
70                "invalid permission `{value}`; expected read, write, ssh-read, or ssh-write"
71            )),
72        }
73    }
74}
75
76#[derive(Debug, Deserialize, Default, Clone)]
77#[serde(deny_unknown_fields)]
78#[allow(dead_code)]
79pub struct QueryOptions {
80    #[serde(default)]
81    pub stream_rows: bool,
82    pub batch_rows: Option<usize>,
83    pub batch_bytes: Option<usize>,
84    pub statement_timeout_ms: Option<u64>,
85    pub lock_timeout_ms: Option<u64>,
86    pub permission: Option<Permission>,
87    pub inline_max_rows: Option<usize>,
88    pub inline_max_bytes: Option<usize>,
89}
90
91#[derive(Debug, Serialize)]
92#[serde(tag = "code")]
93pub enum Output {
94    #[serde(rename = "result")]
95    Result {
96        #[serde(skip_serializing_if = "Option::is_none")]
97        id: Option<String>,
98        #[serde(skip_serializing_if = "Option::is_none")]
99        session: Option<String>,
100        command_tag: String,
101        columns: Vec<ColumnInfo>,
102        rows: Vec<Value>,
103        row_count: usize,
104        trace: Trace,
105    },
106    #[serde(rename = "result_start")]
107    ResultStart {
108        id: String,
109        #[serde(skip_serializing_if = "Option::is_none")]
110        session: Option<String>,
111        columns: Vec<ColumnInfo>,
112    },
113    #[serde(rename = "result_rows")]
114    ResultRows {
115        id: String,
116        rows: Vec<Value>,
117        rows_batch_count: usize,
118    },
119    #[serde(rename = "result_end")]
120    ResultEnd {
121        id: String,
122        #[serde(skip_serializing_if = "Option::is_none")]
123        session: Option<String>,
124        command_tag: String,
125        trace: Trace,
126    },
127    #[serde(rename = "sql_error")]
128    SqlError {
129        #[serde(skip_serializing_if = "Option::is_none")]
130        id: Option<String>,
131        #[serde(skip_serializing_if = "Option::is_none")]
132        session: Option<String>,
133        sqlstate: String,
134        message: String,
135        #[serde(skip_serializing_if = "Option::is_none")]
136        detail: Option<String>,
137        #[serde(skip_serializing_if = "Option::is_none")]
138        hint: Option<String>,
139        #[serde(skip_serializing_if = "Option::is_none")]
140        position: Option<String>,
141        trace: Trace,
142    },
143    #[serde(rename = "error")]
144    Error {
145        #[serde(skip_serializing_if = "Option::is_none")]
146        id: Option<String>,
147        error_code: String,
148        error: String,
149        #[serde(skip_serializing_if = "Option::is_none")]
150        hint: Option<String>,
151        retryable: bool,
152        trace: Trace,
153    },
154    #[serde(rename = "dry_run")]
155    DryRun {
156        #[serde(skip_serializing_if = "Option::is_none")]
157        id: Option<String>,
158        sql: String,
159        params: Vec<String>,
160        #[serde(skip_serializing_if = "Option::is_none")]
161        session: Option<String>,
162        trace: Trace,
163    },
164    #[serde(rename = "config")]
165    Config(RuntimeConfig),
166    #[serde(rename = "pong")]
167    Pong { trace: PongTrace },
168    #[serde(rename = "close")]
169    Close { message: String, trace: CloseTrace },
170    #[serde(rename = "log")]
171    Log {
172        event: String,
173        #[serde(skip_serializing_if = "Option::is_none")]
174        request_id: Option<String>,
175        #[serde(skip_serializing_if = "Option::is_none")]
176        session: Option<String>,
177        #[serde(skip_serializing_if = "Option::is_none")]
178        error_code: Option<String>,
179        #[serde(skip_serializing_if = "Option::is_none")]
180        command_tag: Option<String>,
181        #[serde(skip_serializing_if = "Option::is_none")]
182        version: Option<String>,
183        #[serde(skip_serializing_if = "Option::is_none")]
184        config: Option<Value>,
185        #[serde(skip_serializing_if = "Option::is_none")]
186        args: Option<Value>,
187        #[serde(skip_serializing_if = "Option::is_none")]
188        env: Option<Value>,
189        trace: Trace,
190    },
191}
192
193#[derive(Debug, Serialize, Clone)]
194pub struct ColumnInfo {
195    pub name: String,
196    #[serde(rename = "type")]
197    pub type_name: String,
198}
199
200#[derive(Debug, Serialize, Clone)]
201pub struct Trace {
202    pub duration_ms: u64,
203    #[serde(skip_serializing_if = "Option::is_none")]
204    pub row_count: Option<usize>,
205    #[serde(skip_serializing_if = "Option::is_none")]
206    pub payload_bytes: Option<usize>,
207}
208
209impl Trace {
210    pub fn only_duration(duration_ms: u64) -> Self {
211        Self {
212            duration_ms,
213            row_count: None,
214            payload_bytes: None,
215        }
216    }
217}
218
219#[derive(Debug, Serialize)]
220pub struct PongTrace {
221    pub uptime_s: u64,
222    pub requests_total: u64,
223    pub in_flight: usize,
224}
225
226#[derive(Debug, Serialize)]
227pub struct CloseTrace {
228    pub uptime_s: u64,
229    pub requests_total: u64,
230}
231
232#[derive(Debug, Serialize, Deserialize, Clone, Default)]
233#[serde(deny_unknown_fields)]
234pub struct SessionConfig {
235    #[serde(skip_serializing_if = "Option::is_none")]
236    pub dsn_secret: Option<String>,
237    #[serde(skip_serializing_if = "Option::is_none")]
238    pub conninfo_secret: Option<String>,
239    #[serde(skip_serializing_if = "Option::is_none")]
240    pub host: Option<String>,
241    #[serde(skip_serializing_if = "Option::is_none")]
242    pub port: Option<u16>,
243    #[serde(skip_serializing_if = "Option::is_none")]
244    pub user: Option<String>,
245    #[serde(skip_serializing_if = "Option::is_none")]
246    pub dbname: Option<String>,
247    #[serde(skip_serializing_if = "Option::is_none")]
248    pub password_secret: Option<String>,
249    #[serde(skip_serializing_if = "Option::is_none")]
250    pub ssh: Option<String>,
251    #[serde(default, skip_serializing_if = "Vec::is_empty")]
252    pub ssh_options: Vec<String>,
253    #[serde(skip_serializing_if = "Option::is_none")]
254    pub ssh_local_host: Option<String>,
255    #[serde(skip_serializing_if = "Option::is_none")]
256    pub ssh_local_port: Option<u16>,
257    #[serde(skip_serializing_if = "Option::is_none")]
258    pub ssh_remote_socket: Option<String>,
259    #[serde(skip_serializing_if = "Option::is_none")]
260    pub ssh_sudo_user: Option<String>,
261}
262
263impl SessionConfig {
264    pub fn uses_ssh_transport(&self) -> bool {
265        self.ssh.is_some()
266            || !self.ssh_options.is_empty()
267            || self.ssh_local_host.is_some()
268            || self.ssh_local_port.is_some()
269            || self.ssh_remote_socket.is_some()
270            || self.ssh_sudo_user.is_some()
271    }
272}
273
274#[derive(Debug, Serialize, Deserialize, Clone)]
275pub struct RuntimeConfig {
276    pub default_session: String,
277    #[serde(default)]
278    pub sessions: HashMap<String, SessionConfig>,
279    pub inline_max_rows: usize,
280    pub inline_max_bytes: usize,
281    pub statement_timeout_ms: u64,
282    pub lock_timeout_ms: u64,
283    #[serde(default)]
284    pub log: Vec<String>,
285}
286
287impl Default for RuntimeConfig {
288    fn default() -> Self {
289        let mut sessions = HashMap::new();
290        sessions.insert("default".to_string(), SessionConfig::default());
291        Self {
292            default_session: "default".to_string(),
293            sessions,
294            inline_max_rows: 1000,
295            inline_max_bytes: 1_048_576,
296            statement_timeout_ms: 30_000,
297            lock_timeout_ms: 5_000,
298            log: vec![],
299        }
300    }
301}
302
303#[derive(Debug, Deserialize, Default)]
304#[serde(deny_unknown_fields)]
305pub struct ConfigPatch {
306    pub default_session: Option<String>,
307    pub sessions: Option<HashMap<String, SessionConfigPatch>>,
308    pub inline_max_rows: Option<usize>,
309    pub inline_max_bytes: Option<usize>,
310    pub statement_timeout_ms: Option<u64>,
311    pub lock_timeout_ms: Option<u64>,
312    pub log: Option<Vec<String>>,
313}
314
315#[derive(Debug, Default)]
316pub enum PatchField<T> {
317    #[default]
318    Missing,
319    Null,
320    Value(T),
321}
322
323impl<'de, T> Deserialize<'de> for PatchField<T>
324where
325    T: Deserialize<'de>,
326{
327    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
328    where
329        D: serde::Deserializer<'de>,
330    {
331        let value = Option::<T>::deserialize(deserializer)?;
332        match value {
333            Some(value) => Ok(Self::Value(value)),
334            None => Ok(Self::Null),
335        }
336    }
337}
338
339impl<T> PatchField<T> {
340    pub fn into_update(self) -> Option<Option<T>> {
341        match self {
342            Self::Missing => None,
343            Self::Null => Some(None),
344            Self::Value(value) => Some(Some(value)),
345        }
346    }
347}
348
349#[derive(Debug, Deserialize, Default)]
350#[serde(deny_unknown_fields)]
351pub struct SessionConfigPatch {
352    #[serde(default)]
353    pub dsn_secret: PatchField<String>,
354    #[serde(default)]
355    pub conninfo_secret: PatchField<String>,
356    #[serde(default)]
357    pub host: PatchField<String>,
358    #[serde(default)]
359    pub port: PatchField<u16>,
360    #[serde(default)]
361    pub user: PatchField<String>,
362    #[serde(default)]
363    pub dbname: PatchField<String>,
364    #[serde(default)]
365    pub password_secret: PatchField<String>,
366    #[serde(default)]
367    pub ssh: PatchField<String>,
368    #[serde(default)]
369    pub ssh_options: PatchField<Vec<String>>,
370    #[serde(default)]
371    pub ssh_local_host: PatchField<String>,
372    #[serde(default)]
373    pub ssh_local_port: PatchField<u16>,
374    #[serde(default)]
375    pub ssh_remote_socket: PatchField<String>,
376    #[serde(default)]
377    pub ssh_sudo_user: PatchField<String>,
378}
379
380#[derive(Debug, Clone)]
381#[allow(dead_code)]
382pub struct ResolvedOptions {
383    pub stream_rows: bool,
384    pub batch_rows: usize,
385    pub batch_bytes: usize,
386    pub statement_timeout_ms: u64,
387    pub lock_timeout_ms: u64,
388    pub read_only: bool,
389    pub inline_max_rows: usize,
390    pub inline_max_bytes: usize,
391}
392
393#[cfg(test)]
394#[path = "../tests/support/unit_types.rs"]
395mod tests;
agent_first_psql/types.rs

agent_first_psql/
types.rs
