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")]
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, Deserialize, Default, Clone)]
30#[allow(dead_code)]
31pub struct QueryOptions {
32    #[serde(default)]
33    pub stream_rows: bool,
34    pub batch_rows: Option<usize>,
35    pub batch_bytes: Option<usize>,
36    pub statement_timeout_ms: Option<u64>,
37    pub lock_timeout_ms: Option<u64>,
38    pub read_only: Option<bool>,
39    pub inline_max_rows: Option<usize>,
40    pub inline_max_bytes: Option<usize>,
41}
42
43#[derive(Debug, Serialize)]
44#[serde(tag = "code")]
45pub enum Output {
46    #[serde(rename = "result")]
47    Result {
48        #[serde(skip_serializing_if = "Option::is_none")]
49        id: Option<String>,
50        #[serde(skip_serializing_if = "Option::is_none")]
51        session: Option<String>,
52        command_tag: String,
53        columns: Vec<ColumnInfo>,
54        rows: Vec<Value>,
55        row_count: usize,
56        trace: Trace,
57    },
58    #[serde(rename = "result_start")]
59    ResultStart {
60        id: String,
61        #[serde(skip_serializing_if = "Option::is_none")]
62        session: Option<String>,
63        columns: Vec<ColumnInfo>,
64    },
65    #[serde(rename = "result_rows")]
66    ResultRows {
67        id: String,
68        rows: Vec<Value>,
69        rows_batch_count: usize,
70    },
71    #[serde(rename = "result_end")]
72    ResultEnd {
73        id: String,
74        #[serde(skip_serializing_if = "Option::is_none")]
75        session: Option<String>,
76        command_tag: String,
77        trace: Trace,
78    },
79    #[serde(rename = "sql_error")]
80    SqlError {
81        #[serde(skip_serializing_if = "Option::is_none")]
82        id: Option<String>,
83        #[serde(skip_serializing_if = "Option::is_none")]
84        session: Option<String>,
85        sqlstate: String,
86        message: String,
87        #[serde(skip_serializing_if = "Option::is_none")]
88        detail: Option<String>,
89        #[serde(skip_serializing_if = "Option::is_none")]
90        hint: Option<String>,
91        #[serde(skip_serializing_if = "Option::is_none")]
92        position: Option<String>,
93        trace: Trace,
94    },
95    #[serde(rename = "error")]
96    Error {
97        #[serde(skip_serializing_if = "Option::is_none")]
98        id: Option<String>,
99        error_code: String,
100        error: String,
101        #[serde(skip_serializing_if = "Option::is_none")]
102        hint: Option<String>,
103        retryable: bool,
104        trace: Trace,
105    },
106    #[serde(rename = "dry_run")]
107    DryRun {
108        #[serde(skip_serializing_if = "Option::is_none")]
109        id: Option<String>,
110        sql: String,
111        params: Vec<String>,
112        #[serde(skip_serializing_if = "Option::is_none")]
113        session: Option<String>,
114        trace: Trace,
115    },
116    #[serde(rename = "config")]
117    Config(RuntimeConfig),
118    #[serde(rename = "pong")]
119    Pong { trace: PongTrace },
120    #[serde(rename = "close")]
121    Close { message: String, trace: CloseTrace },
122    #[serde(rename = "log")]
123    Log {
124        event: String,
125        #[serde(skip_serializing_if = "Option::is_none")]
126        request_id: Option<String>,
127        #[serde(skip_serializing_if = "Option::is_none")]
128        session: Option<String>,
129        #[serde(skip_serializing_if = "Option::is_none")]
130        error_code: Option<String>,
131        #[serde(skip_serializing_if = "Option::is_none")]
132        command_tag: Option<String>,
133        #[serde(skip_serializing_if = "Option::is_none")]
134        version: Option<String>,
135        #[serde(skip_serializing_if = "Option::is_none")]
136        argv: Option<Vec<String>>,
137        #[serde(skip_serializing_if = "Option::is_none")]
138        config: Option<Value>,
139        #[serde(skip_serializing_if = "Option::is_none")]
140        args: Option<Value>,
141        #[serde(skip_serializing_if = "Option::is_none")]
142        env: Option<Value>,
143        trace: Trace,
144    },
145}
146
147#[derive(Debug, Serialize, Clone)]
148pub struct ColumnInfo {
149    pub name: String,
150    #[serde(rename = "type")]
151    pub type_name: String,
152}
153
154#[derive(Debug, Serialize, Clone)]
155pub struct Trace {
156    pub duration_ms: u64,
157    #[serde(skip_serializing_if = "Option::is_none")]
158    pub row_count: Option<usize>,
159    #[serde(skip_serializing_if = "Option::is_none")]
160    pub payload_bytes: Option<usize>,
161}
162
163impl Trace {
164    pub fn only_duration(duration_ms: u64) -> Self {
165        Self {
166            duration_ms,
167            row_count: None,
168            payload_bytes: None,
169        }
170    }
171}
172
173#[derive(Debug, Serialize)]
174pub struct PongTrace {
175    pub uptime_s: u64,
176    pub requests_total: u64,
177    pub in_flight: usize,
178}
179
180#[derive(Debug, Serialize)]
181pub struct CloseTrace {
182    pub uptime_s: u64,
183    pub requests_total: u64,
184}
185
186#[derive(Debug, Serialize, Deserialize, Clone, Default)]
187pub struct SessionConfig {
188    #[serde(skip_serializing_if = "Option::is_none")]
189    pub dsn_secret: Option<String>,
190    #[serde(skip_serializing_if = "Option::is_none")]
191    pub conninfo_secret: Option<String>,
192    #[serde(skip_serializing_if = "Option::is_none")]
193    pub host: Option<String>,
194    #[serde(skip_serializing_if = "Option::is_none")]
195    pub port: Option<u16>,
196    #[serde(skip_serializing_if = "Option::is_none")]
197    pub user: Option<String>,
198    #[serde(skip_serializing_if = "Option::is_none")]
199    pub dbname: Option<String>,
200    #[serde(skip_serializing_if = "Option::is_none")]
201    pub password_secret: Option<String>,
202}
203
204#[derive(Debug, Serialize, Deserialize, Clone)]
205pub struct RuntimeConfig {
206    pub default_session: String,
207    #[serde(default)]
208    pub sessions: HashMap<String, SessionConfig>,
209    pub inline_max_rows: usize,
210    pub inline_max_bytes: usize,
211    pub statement_timeout_ms: u64,
212    pub lock_timeout_ms: u64,
213    #[serde(default)]
214    pub log: Vec<String>,
215}
216
217impl Default for RuntimeConfig {
218    fn default() -> Self {
219        let mut sessions = HashMap::new();
220        sessions.insert("default".to_string(), SessionConfig::default());
221        Self {
222            default_session: "default".to_string(),
223            sessions,
224            inline_max_rows: 1000,
225            inline_max_bytes: 1_048_576,
226            statement_timeout_ms: 30_000,
227            lock_timeout_ms: 5_000,
228            log: vec![],
229        }
230    }
231}
232
233#[derive(Debug, Deserialize, Default)]
234pub struct ConfigPatch {
235    pub default_session: Option<String>,
236    pub sessions: Option<HashMap<String, SessionConfigPatch>>,
237    pub inline_max_rows: Option<usize>,
238    pub inline_max_bytes: Option<usize>,
239    pub statement_timeout_ms: Option<u64>,
240    pub lock_timeout_ms: Option<u64>,
241    pub log: Option<Vec<String>>,
242}
243
244#[derive(Debug, Default)]
245pub enum PatchField<T> {
246    #[default]
247    Missing,
248    Null,
249    Value(T),
250}
251
252impl<'de, T> Deserialize<'de> for PatchField<T>
253where
254    T: Deserialize<'de>,
255{
256    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
257    where
258        D: serde::Deserializer<'de>,
259    {
260        let value = Option::<T>::deserialize(deserializer)?;
261        match value {
262            Some(value) => Ok(Self::Value(value)),
263            None => Ok(Self::Null),
264        }
265    }
266}
267
268impl<T> PatchField<T> {
269    pub fn into_update(self) -> Option<Option<T>> {
270        match self {
271            Self::Missing => None,
272            Self::Null => Some(None),
273            Self::Value(value) => Some(Some(value)),
274        }
275    }
276}
277
278#[derive(Debug, Deserialize, Default)]
279pub struct SessionConfigPatch {
280    #[serde(default)]
281    pub dsn_secret: PatchField<String>,
282    #[serde(default)]
283    pub conninfo_secret: PatchField<String>,
284    #[serde(default)]
285    pub host: PatchField<String>,
286    #[serde(default)]
287    pub port: PatchField<u16>,
288    #[serde(default)]
289    pub user: PatchField<String>,
290    #[serde(default)]
291    pub dbname: PatchField<String>,
292    #[serde(default)]
293    pub password_secret: PatchField<String>,
294}
295
296#[derive(Debug, Clone)]
297#[allow(dead_code)]
298pub struct ResolvedOptions {
299    pub stream_rows: bool,
300    pub batch_rows: usize,
301    pub batch_bytes: usize,
302    pub statement_timeout_ms: u64,
303    pub lock_timeout_ms: u64,
304    pub read_only: bool,
305    pub inline_max_rows: usize,
306    pub inline_max_bytes: usize,
307}
308
309#[cfg(test)]
310#[path = "../tests/support/unit_types.rs"]
311mod tests;
agent_first_psql/types.rs

agent_first_psql/
types.rs
