Skip to main content

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    #[serde(rename = "session_info")]
28    SessionInfo {
29        #[serde(default)]
30        id: Option<String>,
31        #[serde(default)]
32        session: Option<String>,
33    },
34}
35
36#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
37pub enum Permission {
38    #[serde(rename = "read")]
39    Read,
40    #[serde(rename = "write")]
41    Write,
42    #[serde(rename = "ssh-read")]
43    SshRead,
44    #[serde(rename = "ssh-write")]
45    SshWrite,
46    #[serde(rename = "container-read")]
47    ContainerRead,
48    #[serde(rename = "container-write")]
49    ContainerWrite,
50}
51
52impl Permission {
53    pub fn as_str(self) -> &'static str {
54        match self {
55            Self::Read => "read",
56            Self::Write => "write",
57            Self::SshRead => "ssh-read",
58            Self::SshWrite => "ssh-write",
59            Self::ContainerRead => "container-read",
60            Self::ContainerWrite => "container-write",
61        }
62    }
63
64    pub fn is_read_only(self) -> bool {
65        matches!(self, Self::Read | Self::SshRead | Self::ContainerRead)
66    }
67
68    pub fn allows_ssh(self) -> bool {
69        matches!(self, Self::SshRead | Self::SshWrite)
70    }
71
72    pub fn allows_container(self) -> bool {
73        matches!(self, Self::ContainerRead | Self::ContainerWrite)
74    }
75}
76
77impl std::str::FromStr for Permission {
78    type Err = String;
79
80    fn from_str(value: &str) -> Result<Self, Self::Err> {
81        match value {
82            "read" => Ok(Self::Read),
83            "write" => Ok(Self::Write),
84            "ssh-read" => Ok(Self::SshRead),
85            "ssh-write" => Ok(Self::SshWrite),
86            "container-read" => Ok(Self::ContainerRead),
87            "container-write" => Ok(Self::ContainerWrite),
88            _ => Err(format!(
89                "invalid permission `{value}`; expected read, write, ssh-read, ssh-write, container-read, or container-write"
90            )),
91        }
92    }
93}
94
95#[derive(Debug, Deserialize, Default, Clone)]
96#[serde(deny_unknown_fields)]
97#[allow(dead_code)]
98pub struct QueryOptions {
99    #[serde(default)]
100    pub stream_rows: bool,
101    pub batch_rows: Option<usize>,
102    pub batch_bytes: Option<usize>,
103    pub statement_timeout_ms: Option<u64>,
104    pub lock_timeout_ms: Option<u64>,
105    pub permission: Option<Permission>,
106    pub inline_max_rows: Option<usize>,
107    pub inline_max_bytes: Option<usize>,
108}
109
110#[derive(Debug, Serialize)]
111#[serde(tag = "code")]
112pub enum Output {
113    #[serde(rename = "result")]
114    Result {
115        #[serde(skip_serializing_if = "Option::is_none")]
116        id: Option<String>,
117        #[serde(skip_serializing_if = "Option::is_none")]
118        session: Option<String>,
119        command_tag: String,
120        columns: Vec<ColumnInfo>,
121        rows: Vec<Value>,
122        row_count: usize,
123        trace: Trace,
124    },
125    #[serde(rename = "result_start")]
126    ResultStart {
127        id: String,
128        #[serde(skip_serializing_if = "Option::is_none")]
129        session: Option<String>,
130        columns: Vec<ColumnInfo>,
131    },
132    #[serde(rename = "result_rows")]
133    ResultRows {
134        id: String,
135        rows: Vec<Value>,
136        rows_batch_count: usize,
137    },
138    #[serde(rename = "result_end")]
139    ResultEnd {
140        id: String,
141        #[serde(skip_serializing_if = "Option::is_none")]
142        session: Option<String>,
143        command_tag: String,
144        trace: Trace,
145    },
146    #[serde(rename = "sql_error")]
147    SqlError {
148        #[serde(skip_serializing_if = "Option::is_none")]
149        id: Option<String>,
150        #[serde(skip_serializing_if = "Option::is_none")]
151        session: Option<String>,
152        sqlstate: String,
153        message: String,
154        #[serde(skip_serializing_if = "Option::is_none")]
155        detail: Option<String>,
156        #[serde(skip_serializing_if = "Option::is_none")]
157        hint: Option<String>,
158        #[serde(skip_serializing_if = "Option::is_none")]
159        position: Option<String>,
160        trace: Trace,
161    },
162    #[serde(rename = "error")]
163    Error {
164        #[serde(skip_serializing_if = "Option::is_none")]
165        id: Option<String>,
166        error_code: String,
167        error: String,
168        #[serde(skip_serializing_if = "Option::is_none")]
169        hint: Option<String>,
170        retryable: bool,
171        trace: Trace,
172    },
173    #[serde(rename = "error")]
174    ConnectError {
175        #[serde(skip_serializing_if = "Option::is_none")]
176        id: Option<String>,
177        error_code: String,
178        error: String,
179        #[serde(skip_serializing_if = "Option::is_none")]
180        sqlstate: Option<String>,
181        #[serde(skip_serializing_if = "Option::is_none")]
182        message: Option<String>,
183        #[serde(skip_serializing_if = "Option::is_none")]
184        detail: Option<String>,
185        #[serde(skip_serializing_if = "Option::is_none")]
186        hint: Option<String>,
187        retryable: bool,
188        trace: Trace,
189    },
190    #[serde(rename = "dry_run")]
191    DryRun {
192        #[serde(skip_serializing_if = "Option::is_none")]
193        id: Option<String>,
194        sql: String,
195        params: Vec<String>,
196        #[serde(skip_serializing_if = "Option::is_none")]
197        session: Option<String>,
198        trace: Trace,
199    },
200    #[serde(rename = "config")]
201    Config(RuntimeConfig),
202    #[serde(rename = "pong")]
203    Pong { trace: PongTrace },
204    #[serde(rename = "close")]
205    Close { message: String, trace: CloseTrace },
206    #[serde(rename = "session_info")]
207    SessionInfo {
208        #[serde(skip_serializing_if = "Option::is_none")]
209        id: Option<String>,
210        session: String,
211        transport_kind: String,
212        permission_default: String,
213        stream_rows_default: bool,
214        batch_rows: usize,
215        batch_bytes: usize,
216        inline_max_rows: usize,
217        inline_max_bytes: usize,
218        statement_timeout_ms: u64,
219        lock_timeout_ms: u64,
220        trace: Trace,
221    },
222    #[serde(rename = "log")]
223    Log {
224        event: String,
225        #[serde(skip_serializing_if = "Option::is_none")]
226        request_id: Option<String>,
227        #[serde(skip_serializing_if = "Option::is_none")]
228        session: Option<String>,
229        #[serde(skip_serializing_if = "Option::is_none")]
230        error_code: Option<String>,
231        #[serde(skip_serializing_if = "Option::is_none")]
232        command_tag: Option<String>,
233        #[serde(skip_serializing_if = "Option::is_none")]
234        version: Option<String>,
235        #[serde(skip_serializing_if = "Option::is_none")]
236        config: Option<Value>,
237        #[serde(skip_serializing_if = "Option::is_none")]
238        args: Option<Value>,
239        #[serde(skip_serializing_if = "Option::is_none")]
240        env: Option<Value>,
241        #[serde(skip_serializing_if = "Option::is_none")]
242        chain: Option<String>,
243        trace: Trace,
244    },
245}
246
247#[derive(Debug, Serialize, Clone)]
248pub struct ColumnInfo {
249    pub name: String,
250    #[serde(rename = "type")]
251    pub type_name: String,
252}
253
254#[derive(Debug, Serialize, Clone)]
255pub struct Trace {
256    pub duration_ms: u64,
257    #[serde(skip_serializing_if = "Option::is_none")]
258    pub row_count: Option<usize>,
259    #[serde(skip_serializing_if = "Option::is_none")]
260    pub payload_bytes: Option<usize>,
261}
262
263impl Trace {
264    pub fn only_duration(duration_ms: u64) -> Self {
265        Self {
266            duration_ms,
267            row_count: None,
268            payload_bytes: None,
269        }
270    }
271}
272
273#[derive(Debug, Serialize)]
274pub struct PongTrace {
275    pub uptime_s: u64,
276    pub requests_total: u64,
277    pub in_flight: usize,
278}
279
280#[derive(Debug, Serialize)]
281pub struct CloseTrace {
282    pub uptime_s: u64,
283    pub requests_total: u64,
284}
285
286#[derive(Debug, Serialize, Clone, Default)]
287pub struct SessionConfig {
288    #[serde(skip_serializing_if = "Option::is_none")]
289    pub dsn_secret: Option<String>,
290    #[serde(skip_serializing_if = "Option::is_none")]
291    pub conninfo_secret: Option<String>,
292    #[serde(skip_serializing_if = "Option::is_none")]
293    pub host: Option<String>,
294    #[serde(skip_serializing_if = "Option::is_none")]
295    pub port: Option<u16>,
296    #[serde(skip_serializing_if = "Option::is_none")]
297    pub user: Option<String>,
298    #[serde(skip_serializing_if = "Option::is_none")]
299    pub dbname: Option<String>,
300    #[serde(skip_serializing_if = "Option::is_none")]
301    pub password_secret: Option<String>,
302    #[serde(flatten)]
303    pub ssh: SshConfig,
304    #[serde(flatten)]
305    pub container: ContainerConfig,
306}
307
308#[derive(Debug, Serialize, Deserialize, Clone, Default)]
309pub struct SshConfig {
310    #[serde(rename = "ssh", skip_serializing_if = "Option::is_none")]
311    pub destination: Option<String>,
312    #[serde(rename = "ssh_options", default, skip_serializing_if = "Vec::is_empty")]
313    pub options: Vec<String>,
314    #[serde(rename = "ssh_local_host", skip_serializing_if = "Option::is_none")]
315    pub local_host: Option<String>,
316    #[serde(rename = "ssh_local_port", skip_serializing_if = "Option::is_none")]
317    pub local_port: Option<u16>,
318    #[serde(rename = "ssh_remote_socket", skip_serializing_if = "Option::is_none")]
319    pub remote_socket: Option<String>,
320    #[serde(rename = "ssh_sudo_user", skip_serializing_if = "Option::is_none")]
321    pub sudo_user: Option<String>,
322}
323
324impl SshConfig {
325    pub fn has_transport_fields(&self) -> bool {
326        self.destination.is_some()
327            || !self.options.is_empty()
328            || self.local_host.is_some()
329            || self.local_port.is_some()
330            || self.remote_socket.is_some()
331            || self.sudo_user.is_some()
332    }
333
334    pub fn has_tunnel_or_bridge_options(&self) -> bool {
335        self.local_host.is_some()
336            || self.local_port.is_some()
337            || self.remote_socket.is_some()
338            || self.sudo_user.is_some()
339    }
340}
341
342#[derive(Debug, Serialize, Deserialize, Clone, Default)]
343pub struct ContainerConfig {
344    #[serde(rename = "container", skip_serializing_if = "Option::is_none")]
345    pub target: Option<String>,
346    #[serde(rename = "container_driver", skip_serializing_if = "Option::is_none")]
347    pub driver: Option<String>,
348    #[serde(rename = "container_runtime", skip_serializing_if = "Option::is_none")]
349    pub runtime: Option<String>,
350    #[serde(rename = "container_user", skip_serializing_if = "Option::is_none")]
351    pub user: Option<String>,
352    #[serde(
353        rename = "container_namespace",
354        skip_serializing_if = "Option::is_none"
355    )]
356    pub namespace: Option<String>,
357    #[serde(rename = "container_context", skip_serializing_if = "Option::is_none")]
358    pub context: Option<String>,
359    #[serde(
360        rename = "container_compose_files",
361        default,
362        skip_serializing_if = "Vec::is_empty"
363    )]
364    pub compose_files: Vec<String>,
365    #[serde(
366        rename = "container_compose_project",
367        skip_serializing_if = "Option::is_none"
368    )]
369    pub compose_project: Option<String>,
370    #[serde(
371        rename = "container_pod_container",
372        skip_serializing_if = "Option::is_none"
373    )]
374    pub pod_container: Option<String>,
375}
376
377impl ContainerConfig {
378    pub fn has_transport_fields(&self) -> bool {
379        self.target.is_some()
380            || self.driver.is_some()
381            || self.runtime.is_some()
382            || self.user.is_some()
383            || self.namespace.is_some()
384            || self.context.is_some()
385            || !self.compose_files.is_empty()
386            || self.compose_project.is_some()
387            || self.pod_container.is_some()
388    }
389}
390
391#[derive(Debug, Deserialize, Default)]
392#[serde(deny_unknown_fields)]
393struct SessionConfigFlat {
394    #[serde(default)]
395    dsn_secret: Option<String>,
396    #[serde(default)]
397    conninfo_secret: Option<String>,
398    #[serde(default)]
399    host: Option<String>,
400    #[serde(default)]
401    port: Option<u16>,
402    #[serde(default)]
403    user: Option<String>,
404    #[serde(default)]
405    dbname: Option<String>,
406    #[serde(default)]
407    password_secret: Option<String>,
408    #[serde(default)]
409    ssh: Option<String>,
410    #[serde(default)]
411    ssh_options: Vec<String>,
412    #[serde(default)]
413    ssh_local_host: Option<String>,
414    #[serde(default)]
415    ssh_local_port: Option<u16>,
416    #[serde(default)]
417    ssh_remote_socket: Option<String>,
418    #[serde(default)]
419    ssh_sudo_user: Option<String>,
420    #[serde(default)]
421    container: Option<String>,
422    #[serde(default)]
423    container_driver: Option<String>,
424    #[serde(default)]
425    container_runtime: Option<String>,
426    #[serde(default)]
427    container_user: Option<String>,
428    #[serde(default)]
429    container_namespace: Option<String>,
430    #[serde(default)]
431    container_context: Option<String>,
432    #[serde(default)]
433    container_compose_files: Vec<String>,
434    #[serde(default)]
435    container_compose_project: Option<String>,
436    #[serde(default)]
437    container_pod_container: Option<String>,
438}
439
440impl From<SessionConfigFlat> for SessionConfig {
441    fn from(flat: SessionConfigFlat) -> Self {
442        Self {
443            dsn_secret: flat.dsn_secret,
444            conninfo_secret: flat.conninfo_secret,
445            host: flat.host,
446            port: flat.port,
447            user: flat.user,
448            dbname: flat.dbname,
449            password_secret: flat.password_secret,
450            ssh: SshConfig {
451                destination: flat.ssh,
452                options: flat.ssh_options,
453                local_host: flat.ssh_local_host,
454                local_port: flat.ssh_local_port,
455                remote_socket: flat.ssh_remote_socket,
456                sudo_user: flat.ssh_sudo_user,
457            },
458            container: ContainerConfig {
459                target: flat.container,
460                driver: flat.container_driver,
461                runtime: flat.container_runtime,
462                user: flat.container_user,
463                namespace: flat.container_namespace,
464                context: flat.container_context,
465                compose_files: flat.container_compose_files,
466                compose_project: flat.container_compose_project,
467                pod_container: flat.container_pod_container,
468            },
469        }
470    }
471}
472
473impl<'de> Deserialize<'de> for SessionConfig {
474    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
475    where
476        D: serde::Deserializer<'de>,
477    {
478        SessionConfigFlat::deserialize(deserializer).map(Self::from)
479    }
480}
481
482impl SessionConfig {
483    pub fn uses_ssh_transport(&self) -> bool {
484        self.ssh.has_transport_fields()
485    }
486
487    pub fn uses_container_transport(&self) -> bool {
488        self.container.has_transport_fields()
489    }
490
491    pub fn transport_kind(&self) -> Result<TransportKind, String> {
492        let uses_ssh = self.uses_ssh_transport();
493        let uses_container = self.uses_container_transport();
494        match (uses_ssh, uses_container) {
495            (false, false) => Ok(TransportKind::Direct),
496            (true, false) => Ok(TransportKind::Ssh),
497            (false, true) => Ok(TransportKind::Container),
498            // --ssh + --container means "run container exec on that remote host".
499            // The PostgreSQL connection still crosses the container boundary.
500            (true, true) => Ok(TransportKind::Container),
501        }
502    }
503}
504
505#[derive(Debug, Clone, Copy, PartialEq, Eq)]
506pub enum TransportKind {
507    Direct,
508    Ssh,
509    Container,
510}
511
512impl TransportKind {
513    pub fn as_str(self) -> &'static str {
514        match self {
515            Self::Direct => "direct",
516            Self::Ssh => "ssh",
517            Self::Container => "container",
518        }
519    }
520}
521
522#[derive(Debug, Serialize, Deserialize, Clone)]
523pub struct RuntimeConfig {
524    pub default_session: String,
525    #[serde(default)]
526    pub sessions: HashMap<String, SessionConfig>,
527    pub inline_max_rows: usize,
528    pub inline_max_bytes: usize,
529    pub statement_timeout_ms: u64,
530    pub lock_timeout_ms: u64,
531    #[serde(default)]
532    pub log: Vec<String>,
533}
534
535impl Default for RuntimeConfig {
536    fn default() -> Self {
537        let mut sessions = HashMap::new();
538        sessions.insert("default".to_string(), SessionConfig::default());
539        Self {
540            default_session: "default".to_string(),
541            sessions,
542            inline_max_rows: 1000,
543            inline_max_bytes: 1_048_576,
544            statement_timeout_ms: 30_000,
545            lock_timeout_ms: 5_000,
546            log: vec![],
547        }
548    }
549}
550
551#[derive(Debug, Deserialize, Default)]
552#[serde(deny_unknown_fields)]
553pub struct ConfigPatch {
554    pub default_session: Option<String>,
555    pub sessions: Option<HashMap<String, SessionConfigPatch>>,
556    pub inline_max_rows: Option<usize>,
557    pub inline_max_bytes: Option<usize>,
558    pub statement_timeout_ms: Option<u64>,
559    pub lock_timeout_ms: Option<u64>,
560    pub log: Option<Vec<String>>,
561}
562
563#[derive(Debug, Default)]
564pub enum PatchField<T> {
565    #[default]
566    Missing,
567    Null,
568    Value(T),
569}
570
571impl<'de, T> Deserialize<'de> for PatchField<T>
572where
573    T: Deserialize<'de>,
574{
575    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
576    where
577        D: serde::Deserializer<'de>,
578    {
579        let value = Option::<T>::deserialize(deserializer)?;
580        match value {
581            Some(value) => Ok(Self::Value(value)),
582            None => Ok(Self::Null),
583        }
584    }
585}
586
587impl<T> PatchField<T> {
588    pub fn into_update(self) -> Option<Option<T>> {
589        match self {
590            Self::Missing => None,
591            Self::Null => Some(None),
592            Self::Value(value) => Some(Some(value)),
593        }
594    }
595}
596
597#[derive(Debug, Default)]
598pub struct SessionConfigPatch {
599    pub dsn_secret: PatchField<String>,
600    pub conninfo_secret: PatchField<String>,
601    pub host: PatchField<String>,
602    pub port: PatchField<u16>,
603    pub user: PatchField<String>,
604    pub dbname: PatchField<String>,
605    pub password_secret: PatchField<String>,
606    pub ssh: SshConfigPatch,
607    pub container: ContainerConfigPatch,
608}
609
610#[derive(Debug, Default)]
611pub struct SshConfigPatch {
612    pub destination: PatchField<String>,
613    pub options: PatchField<Vec<String>>,
614    pub local_host: PatchField<String>,
615    pub local_port: PatchField<u16>,
616    pub remote_socket: PatchField<String>,
617    pub sudo_user: PatchField<String>,
618}
619
620#[derive(Debug, Default)]
621pub struct ContainerConfigPatch {
622    pub target: PatchField<String>,
623    pub driver: PatchField<String>,
624    pub runtime: PatchField<String>,
625    pub user: PatchField<String>,
626    pub namespace: PatchField<String>,
627    pub context: PatchField<String>,
628    pub compose_files: PatchField<Vec<String>>,
629    pub compose_project: PatchField<String>,
630    pub pod_container: PatchField<String>,
631}
632
633#[derive(Debug, Deserialize, Default)]
634#[serde(deny_unknown_fields)]
635struct SessionConfigPatchFlat {
636    #[serde(default)]
637    dsn_secret: PatchField<String>,
638    #[serde(default)]
639    conninfo_secret: PatchField<String>,
640    #[serde(default)]
641    host: PatchField<String>,
642    #[serde(default)]
643    port: PatchField<u16>,
644    #[serde(default)]
645    user: PatchField<String>,
646    #[serde(default)]
647    dbname: PatchField<String>,
648    #[serde(default)]
649    password_secret: PatchField<String>,
650    #[serde(default)]
651    ssh: PatchField<String>,
652    #[serde(default)]
653    ssh_options: PatchField<Vec<String>>,
654    #[serde(default)]
655    ssh_local_host: PatchField<String>,
656    #[serde(default)]
657    ssh_local_port: PatchField<u16>,
658    #[serde(default)]
659    ssh_remote_socket: PatchField<String>,
660    #[serde(default)]
661    ssh_sudo_user: PatchField<String>,
662    #[serde(default)]
663    container: PatchField<String>,
664    #[serde(default)]
665    container_driver: PatchField<String>,
666    #[serde(default)]
667    container_runtime: PatchField<String>,
668    #[serde(default)]
669    container_user: PatchField<String>,
670    #[serde(default)]
671    container_namespace: PatchField<String>,
672    #[serde(default)]
673    container_context: PatchField<String>,
674    #[serde(default)]
675    container_compose_files: PatchField<Vec<String>>,
676    #[serde(default)]
677    container_compose_project: PatchField<String>,
678    #[serde(default)]
679    container_pod_container: PatchField<String>,
680}
681
682impl From<SessionConfigPatchFlat> for SessionConfigPatch {
683    fn from(flat: SessionConfigPatchFlat) -> Self {
684        Self {
685            dsn_secret: flat.dsn_secret,
686            conninfo_secret: flat.conninfo_secret,
687            host: flat.host,
688            port: flat.port,
689            user: flat.user,
690            dbname: flat.dbname,
691            password_secret: flat.password_secret,
692            ssh: SshConfigPatch {
693                destination: flat.ssh,
694                options: flat.ssh_options,
695                local_host: flat.ssh_local_host,
696                local_port: flat.ssh_local_port,
697                remote_socket: flat.ssh_remote_socket,
698                sudo_user: flat.ssh_sudo_user,
699            },
700            container: ContainerConfigPatch {
701                target: flat.container,
702                driver: flat.container_driver,
703                runtime: flat.container_runtime,
704                user: flat.container_user,
705                namespace: flat.container_namespace,
706                context: flat.container_context,
707                compose_files: flat.container_compose_files,
708                compose_project: flat.container_compose_project,
709                pod_container: flat.container_pod_container,
710            },
711        }
712    }
713}
714
715impl<'de> Deserialize<'de> for SessionConfigPatch {
716    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
717    where
718        D: serde::Deserializer<'de>,
719    {
720        SessionConfigPatchFlat::deserialize(deserializer).map(Self::from)
721    }
722}
723
724#[derive(Debug, Clone)]
725#[allow(dead_code)]
726pub struct ResolvedOptions {
727    pub stream_rows: bool,
728    pub batch_rows: usize,
729    pub batch_bytes: usize,
730    pub statement_timeout_ms: u64,
731    pub lock_timeout_ms: u64,
732    pub read_only: bool,
733    pub inline_max_rows: usize,
734    pub inline_max_bytes: usize,
735}
736
737#[cfg(test)]
738#[path = "../tests/support/unit_types.rs"]
739mod tests;