Skip to main content

agent_first_psql/
config.rs

1use crate::types::*;
2use agent_first_data::cli_parse_log_filters;
3
4pub const VERSION: &str = env!("CARGO_PKG_VERSION");
5
6impl RuntimeConfig {
7    pub fn apply_update(&mut self, patch: ConfigPatch) {
8        if let Some(v) = patch.default_session {
9            self.default_session = v;
10        }
11        if let Some(v) = patch.inline_max_rows {
12            self.inline_max_rows = v;
13        }
14        if let Some(v) = patch.inline_max_bytes {
15            self.inline_max_bytes = v;
16        }
17        if let Some(v) = patch.statement_timeout_ms {
18            self.statement_timeout_ms = v;
19        }
20        if let Some(v) = patch.lock_timeout_ms {
21            self.lock_timeout_ms = v;
22        }
23        if let Some(v) = patch.log {
24            self.log = cli_parse_log_filters(&v);
25        }
26        if let Some(sessions) = patch.sessions {
27            for (name, s) in sessions {
28                let entry = self.sessions.entry(name).or_default();
29                if let Some(v) = s.dsn_secret.into_update() {
30                    entry.dsn_secret = v;
31                }
32                if let Some(v) = s.conninfo_secret.into_update() {
33                    entry.conninfo_secret = v;
34                }
35                if let Some(v) = s.host.into_update() {
36                    entry.host = v;
37                }
38                if let Some(v) = s.port.into_update() {
39                    entry.port = v;
40                }
41                if let Some(v) = s.user.into_update() {
42                    entry.user = v;
43                }
44                if let Some(v) = s.dbname.into_update() {
45                    entry.dbname = v;
46                }
47                if let Some(v) = s.password_secret.into_update() {
48                    entry.password_secret = v;
49                }
50                if let Some(v) = s.ssh.destination.into_update() {
51                    entry.ssh.destination = v;
52                }
53                if let Some(v) = s.ssh.via.into_update() {
54                    entry.ssh.via = v.unwrap_or_default();
55                }
56                if let Some(v) = s.ssh.options.into_update() {
57                    entry.ssh.options = v.unwrap_or_default();
58                }
59                if let Some(v) = s.ssh.local_host.into_update() {
60                    entry.ssh.local_host = v;
61                }
62                if let Some(v) = s.ssh.local_port.into_update() {
63                    entry.ssh.local_port = v;
64                }
65                if let Some(v) = s.ssh.remote_socket.into_update() {
66                    entry.ssh.remote_socket = v;
67                }
68                if let Some(v) = s.ssh.sudo_user.into_update() {
69                    entry.ssh.sudo_user = v;
70                }
71                if let Some(v) = s.container.target.into_update() {
72                    entry.container.target = v;
73                }
74                if let Some(v) = s.container.driver.into_update() {
75                    entry.container.driver = v;
76                }
77                if let Some(v) = s.container.runtime.into_update() {
78                    entry.container.runtime = v;
79                }
80                if let Some(v) = s.container.user.into_update() {
81                    entry.container.user = v;
82                }
83                if let Some(v) = s.container.namespace.into_update() {
84                    entry.container.namespace = v;
85                }
86                if let Some(v) = s.container.context.into_update() {
87                    entry.container.context = v;
88                }
89                if let Some(v) = s.container.compose_files.into_update() {
90                    entry.container.compose_files = v.unwrap_or_default();
91                }
92                if let Some(v) = s.container.compose_project.into_update() {
93                    entry.container.compose_project = v;
94                }
95                if let Some(v) = s.container.pod_container.into_update() {
96                    entry.container.pod_container = v;
97                }
98            }
99        }
100        if !self.sessions.contains_key(&self.default_session) {
101            self.sessions
102                .insert(self.default_session.clone(), SessionConfig::default());
103        }
104    }
105
106    #[allow(dead_code)]
107    pub fn resolve_options(&self, q: &QueryOptions) -> ResolvedOptions {
108        self.resolve_options_with_permission(q, q.permission.unwrap_or(Permission::Read))
109    }
110
111    pub fn resolve_options_for_session(
112        &self,
113        q: &QueryOptions,
114        session: &SessionConfig,
115    ) -> Result<ResolvedOptions, String> {
116        let transport = session.transport_kind()?;
117        let permission = q.permission.unwrap_or(match transport {
118            TransportKind::Direct => Permission::Read,
119            TransportKind::Ssh => Permission::SshRead,
120            TransportKind::Container => Permission::ContainerRead,
121        });
122        match transport {
123            TransportKind::Direct if permission.allows_ssh() => {
124                return Err(format!(
125                    "permission `{}` requires SSH transport; use `read` or `write` for direct connections",
126                    permission.as_str()
127                ));
128            }
129            TransportKind::Direct if permission.allows_container() => {
130                return Err(format!(
131                    "permission `{}` requires container transport; use `read` or `write` for direct connections",
132                    permission.as_str()
133                ));
134            }
135            TransportKind::Ssh if !permission.allows_ssh() => {
136                return Err(format!(
137                    "permission `{}` does not allow SSH transport; use `ssh-read` or `ssh-write`",
138                    permission.as_str()
139                ));
140            }
141            TransportKind::Container if !permission.allows_container() => {
142                return Err(format!(
143                    "permission `{}` does not allow container transport; use `container-read` or `container-write`",
144                    permission.as_str()
145                ));
146            }
147            _ => {}
148        }
149        Ok(self.resolve_options_with_permission(q, permission))
150    }
151
152    fn resolve_options_with_permission(
153        &self,
154        q: &QueryOptions,
155        permission: Permission,
156    ) -> ResolvedOptions {
157        ResolvedOptions {
158            stream_rows: q.stream_rows,
159            batch_rows: q.batch_rows.unwrap_or(1000).max(1),
160            batch_bytes: q.batch_bytes.unwrap_or(262_144).max(1024),
161            statement_timeout_ms: q.statement_timeout_ms.unwrap_or(self.statement_timeout_ms),
162            lock_timeout_ms: q.lock_timeout_ms.unwrap_or(self.lock_timeout_ms),
163            read_only: permission.is_read_only(),
164            inline_max_rows: q.inline_max_rows.unwrap_or(self.inline_max_rows),
165            inline_max_bytes: q.inline_max_bytes.unwrap_or(self.inline_max_bytes),
166        }
167    }
168}
169
170pub fn sessions_to_invalidate(patch: &ConfigPatch) -> Vec<String> {
171    let mut sessions: Vec<String> = vec![];
172    if let Some(default_session) = patch.default_session.as_ref() {
173        sessions.push(default_session.clone());
174    }
175    if let Some(update_sessions) = patch.sessions.as_ref() {
176        sessions.extend(update_sessions.keys().cloned());
177    }
178    sessions.sort();
179    sessions.dedup();
180    sessions
181}
182
183#[cfg(test)]
184#[path = "../tests/support/unit_config.rs"]
185mod tests;