1use std::sync::LazyLock;
4
5use anyhow::{bail, ensure, Context, Result};
6use serde_json::{json, Value};
7
8use crate::llm::actions::client_trait::{Client, ClientActionResult};
9use crate::llm::actions::protocol_trait::Protocol;
10use crate::llm::actions::{ActionDefinition, Parameter, ParameterDefinition, StartupExamples};
11use crate::protocol::log_template::LogTemplate;
12use crate::protocol::{ConnectContext, EventType};
13use crate::state::app_state::AppState;
14
15use super::wire;
16
17pub struct Db2ClientProtocol;
18
19impl Db2ClientProtocol {
20 pub fn new() -> Self {
21 Self
22 }
23}
24
25impl Default for Db2ClientProtocol {
26 fn default() -> Self {
27 Self::new()
28 }
29}
30
31fn p(name: &str, type_hint: &str, description: &str, required: bool) -> Parameter {
32 Parameter {
33 name: name.to_string(),
34 type_hint: type_hint.to_string(),
35 description: description.to_string(),
36 required,
37 }
38}
39
40fn execute_statement() -> ActionDefinition {
41 ActionDefinition {
42 name: "execute_statement".to_string(),
43 description: "Run one SQL statement that returns no rows (CREATE, INSERT, UPDATE, DELETE, DROP, ...) with EXCSQLIMM. The answer is db2_statement_result (SQLCODE, SQLSTATE, rows affected) or db2_error. Changes are not durable until commit unless the client was opened with autocommit.".to_string(),
44 parameters: vec![p("sql", "string", "The SQL statement, without a trailing semicolon", true)],
45 example: json!({"type": "execute_statement", "sql": "INSERT INTO notes (id, body) VALUES (1, 'hello')"}),
46 log_template: Some(LogTemplate::new().with_info("-> Db2 EXCSQLIMM {sql}")),
47 }
48}
49
50fn execute_query() -> ActionDefinition {
51 ActionDefinition {
52 name: "execute_query".to_string(),
53 description: format!("Run one SELECT (or VALUES) statement: prepare it, open a cursor and fetch rows. The answer is db2_query_result with the column descriptions and the rows as objects keyed by column name, or db2_error. At most max_rows rows are delivered (default {}, at most {}); truncated says whether more existed.", wire::DEFAULT_MAX_ROWS, wire::MAX_ROWS_LIMIT),
54 parameters: vec![
55 p("sql", "string", "The query, without a trailing semicolon", true),
56 p("max_rows", "number", "How many rows to deliver at most", false),
57 ],
58 example: json!({"type": "execute_query", "sql": "SELECT id, body FROM notes ORDER BY id", "max_rows": 50}),
59 log_template: Some(LogTemplate::new().with_info("-> Db2 PRPSQLSTT/OPNQRY {sql}")),
60 }
61}
62
63fn commit() -> ActionDefinition {
64 ActionDefinition {
65 name: "commit".to_string(),
66 description: "Commit the unit of work (RDBCMM). The answer is db2_transaction_result."
67 .to_string(),
68 parameters: vec![],
69 example: json!({"type": "commit"}),
70 log_template: Some(LogTemplate::new().with_info("-> Db2 RDBCMM")),
71 }
72}
73
74fn rollback() -> ActionDefinition {
75 ActionDefinition {
76 name: "rollback".to_string(),
77 description: "Roll back the unit of work (RDBRLLBCK), undoing every change since the last commit. The answer is db2_transaction_result.".to_string(),
78 parameters: vec![],
79 example: json!({"type": "rollback"}),
80 log_template: Some(LogTemplate::new().with_info("-> Db2 RDBRLLBCK")),
81 }
82}
83
84fn disconnect() -> ActionDefinition {
85 ActionDefinition {
86 name: "disconnect".to_string(),
87 description: "Close the connection. Uncommitted work is rolled back by the server."
88 .to_string(),
89 parameters: vec![],
90 example: json!({"type": "disconnect"}),
91 log_template: Some(LogTemplate::new().with_info("-> Db2 disconnect")),
92 }
93}
94
95pub fn actions() -> Vec<ActionDefinition> {
96 vec![
97 execute_statement(),
98 execute_query(),
99 commit(),
100 rollback(),
101 disconnect(),
102 ]
103}
104
105fn ev(id: &str, description: &str, params: Vec<Parameter>) -> EventType {
106 EventType::new(id, description, execute_query().example.clone())
107 .with_parameters(params)
108 .with_actions(actions())
109}
110
111pub static CONNECTED_EVENT: LazyLock<EventType> = LazyLock::new(|| {
112 ev(
113 "db2_connected",
114 "The client logged in (EXCSAT, ACCSEC, SECCHK) and opened the database (ACCRDB); nothing has run yet.",
115 vec![
116 p("remote_addr", "string", "The server's address", true),
117 p("database", "string", "The relational database name opened", true),
118 p("user", "string", "The user id that logged in", true),
119 p("server_class", "string", "The server's class name (SRVCLSNM), e.g. QDB2/LINUXX8664 or Apache Derby", true),
120 p("server_release", "string", "The server's product release level (SRVRLSLV)", true),
121 p("server_product_id", "string", "The product id the server gave when the database opened (PRDID)", true),
122 p("data_representation", "string", "The server's TYPDEFNAM (byte order and character set family)", true),
123 ],
124 )
125});
126
127pub static STATEMENT_RESULT_EVENT: LazyLock<EventType> = LazyLock::new(|| {
128 ev(
129 "db2_statement_result",
130 "A statement run with execute_statement succeeded (SQLCODE 0, or a positive warning).",
131 vec![
132 p(
133 "sql",
134 "string",
135 "The statement text exactly as it was sent with EXCSQLIMM",
136 true,
137 ),
138 p(
139 "sqlcode",
140 "number",
141 "0, or a positive warning code (100: no row matched)",
142 true,
143 ),
144 p(
145 "sqlstate",
146 "string",
147 "The SQLSTATE, blank on plain success",
148 true,
149 ),
150 p(
151 "row_count",
152 "number",
153 "Rows inserted, updated or deleted (SQLERRD3)",
154 true,
155 ),
156 p(
157 "message_tokens",
158 "array",
159 "Message tokens the server attached",
160 true,
161 ),
162 p(
163 "committed",
164 "boolean",
165 "True when autocommit committed it",
166 true,
167 ),
168 ],
169 )
170});
171
172pub static QUERY_RESULT_EVENT: LazyLock<EventType> = LazyLock::new(|| {
173 ev(
174 "db2_query_result",
175 "A query run with execute_query returned its rows.",
176 vec![
177 p("sql", "string", "The query text exactly as it was prepared and opened", true),
178 p("columns", "array", "One {name, type, nullable} per column, in order", true),
179 p("rows", "array", "One object per row keyed by column name. Integers and floats are numbers; DECIMAL is an exact string; DATE, TIME and TIMESTAMP are the server's strings; NULL is null; binary columns are {binary_length}", true),
180 p("row_count", "number", "How many rows are in rows (at most max_rows)", true),
181 p("truncated", "boolean", "True when more rows existed than max_rows", true),
182 p("warnings", "array", "Warning SQLCAs attached to rows", true),
183 ],
184 )
185});
186
187pub static TRANSACTION_RESULT_EVENT: LazyLock<EventType> = LazyLock::new(|| {
188 ev(
189 "db2_transaction_result",
190 "A commit or rollback finished.",
191 vec![
192 p("operation", "string", "commit or rollback", true),
193 p(
194 "outcome",
195 "string",
196 "committed or rolled_back, as the server reported the unit of work",
197 true,
198 ),
199 p(
200 "statements",
201 "array",
202 "The statements this client ran in that unit of work (the last 16)",
203 true,
204 ),
205 p(
206 "sqlcode",
207 "number",
208 "The SQLCODE the server returned with the commit or rollback (0 on success)",
209 true,
210 ),
211 ],
212 )
213});
214
215pub static ERROR_EVENT: LazyLock<EventType> = LazyLock::new(|| {
216 ev(
217 "db2_error",
218 "An action failed: the server refused the SQL (negative SQLCODE), answered with a DRDA error reply, or returned something this client cannot decode. The connection is still open.",
219 vec![
220 p("operation", "string", "execute_statement, prepare, open_query, fetch, commit or rollback", true),
221 p("sql", "string", "The statement, where there was one", false),
222 p("sqlcode", "number", "The negative SQLCODE, for an SQL error", false),
223 p("sqlstate", "string", "The SQLSTATE, for an SQL error", false),
224 p("message_tokens", "array", "The server's message tokens (Db2 sends tokens, not message text)", false),
225 p("reply_message", "string", "The DRDA reply message, for a protocol-level refusal", false),
226 p("message", "string", "What went wrong, in words", true),
227 p("rolled_back", "boolean", "True when the server ended the unit of work abnormally (deadlock, timeout)", false),
228 ],
229 )
230});
231
232pub fn request(v: &Value) -> Result<(String, Value)> {
234 let kind = v
235 .get("type")
236 .and_then(Value::as_str)
237 .context("an action needs a type")?;
238 match kind {
239 "execute_statement" | "execute_query" => {
240 let sql = v
241 .get("sql")
242 .and_then(Value::as_str)
243 .with_context(|| format!("{kind} needs sql"))?
244 .trim()
245 .trim_end_matches(';')
246 .trim();
247 ensure!(!sql.is_empty(), "{kind}: sql is empty");
248 ensure!(
249 sql.len() <= wire::MAX_SQL_BYTES,
250 "{kind}: the statement is {} bytes; at most {} are sent",
251 sql.len(),
252 wire::MAX_SQL_BYTES
253 );
254 let mut data = json!({ "sql": sql });
255 if kind == "execute_query" {
256 let max_rows = match v.get("max_rows") {
257 None | Some(Value::Null) => wire::DEFAULT_MAX_ROWS as u64,
258 Some(n) => n.as_u64().context("max_rows is a whole number")?,
259 };
260 ensure!(
261 (1..=wire::MAX_ROWS_LIMIT as u64).contains(&max_rows),
262 "max_rows is 1..={}",
263 wire::MAX_ROWS_LIMIT
264 );
265 data["max_rows"] = json!(max_rows);
266 }
267 Ok((kind.to_string(), data))
268 }
269 "commit" | "rollback" => Ok((kind.to_string(), json!({}))),
270 other => bail!("Unknown Db2 client action: {other}"),
271 }
272}
273
274impl Protocol for Db2ClientProtocol {
275 fn protocol_name(&self) -> &'static str {
276 "Db2"
277 }
278 fn stack_name(&self) -> &'static str {
279 "ETH>IP>TCP>DRDA>Db2"
280 }
281 fn keywords(&self) -> Vec<&'static str> {
282 vec![
283 "db2 client",
284 "drda client",
285 "connect to db2",
286 "ibm db2",
287 "apache derby",
288 "derby network server",
289 ]
290 }
291 fn description(&self) -> &'static str {
292 "Db2 client: speaks DRDA to IBM Db2 or Apache Derby, logs in, runs statements and queries, and commits or rolls back"
293 }
294 fn get_async_actions(&self, _: &AppState) -> Vec<ActionDefinition> {
295 actions()
296 }
297 fn get_sync_actions(&self) -> Vec<ActionDefinition> {
298 actions()
299 }
300 fn get_event_types(&self) -> Vec<EventType> {
301 vec![
302 CONNECTED_EVENT.clone(),
303 STATEMENT_RESULT_EVENT.clone(),
304 QUERY_RESULT_EVENT.clone(),
305 TRANSACTION_RESULT_EVENT.clone(),
306 ERROR_EVENT.clone(),
307 ]
308 }
309 fn get_startup_parameters(&self) -> Vec<ParameterDefinition> {
310 vec![
311 ParameterDefinition {
312 name: "database".into(),
313 type_hint: "string".into(),
314 description: "The relational database name (RDBNAM). Apache Derby reads connection attributes after it, e.g. sample;create=true".into(),
315 required: true,
316 example: json!("SAMPLE"),
317 default: None,
318 },
319 ParameterDefinition {
320 name: "username".into(),
321 type_hint: "string".into(),
322 description: "The user id sent in SECCHK (printable ASCII)".into(),
323 required: true,
324 example: json!("db2inst1"),
325 default: None,
326 },
327 ParameterDefinition {
328 name: "password".into(),
329 type_hint: "string".into(),
330 description: "The password, sent in the clear (security mechanism USRIDPWD). Without one the client logs in with the user id alone (USRIDONL)".into(),
331 required: false,
332 example: json!("secret"),
333 default: None,
334 },
335 ParameterDefinition {
336 name: "autocommit".into(),
337 type_hint: "boolean".into(),
338 description: "Commit after every successful statement and query".into(),
339 required: false,
340 example: json!(true),
341 default: Some(json!(super::DEFAULT_AUTOCOMMIT)),
342 },
343 ParameterDefinition {
344 name: "product_id".into(),
345 type_hint: "string".into(),
346 description: "The 8-character product id (PRDID, pppvvrrm) sent in ACCRDB and EXCSAT. Apache Derby admits only its own client's prefix, DNC; IBM Db2 accepts any".into(),
347 required: false,
348 example: json!("DNC10170"),
349 default: Some(json!(super::DEFAULT_PRODUCT_ID)),
350 },
351 ParameterDefinition {
352 name: "reply_timeout_secs".into(),
353 type_hint: "number".into(),
354 description: "Seconds to wait for a whole reply before the connection is given up (1..=3600). Statements waiting on a lock legitimately take a while".into(),
355 required: false,
356 example: json!(60),
357 default: Some(json!(super::REPLY_TIMEOUT.as_secs())),
358 },
359 ]
360 }
361 fn metadata(&self) -> crate::protocol::metadata::ProtocolMetadataV2 {
362 use crate::protocol::metadata::*;
363 ProtocolMetadataV2::builder()
364 .state(DevelopmentState::Experimental)
365 .privilege_requirement(PrivilegeRequirement::None)
366 .implementation("Hand-rolled DRDA requester (src/client/db2/wire.rs) at SQLAM 7: EXCSAT, ACCSEC, SECCHK (USRIDPWD or USRIDONL), ACCRDB declaring QTDSQLASC with UTF-8 TYPDEFOVR; EXCSQLIMM for statements; PRPSQLSTT (extended SQLDA) then OPNQRY/CNTQRY/CLSQRY for queries, with the QRYDSC row layout and QRYDTA rows decoded (integers, DECIMAL, floats, character, graphic, DATE/TIME/TIMESTAMP, BOOLEAN, NULL) honouring the server's byte order and CCSIDs; RDBCMM and RDBRLLBCK. DSS continuation, extended object lengths and rows split across query blocks are handled; every length is bounded")
367 .llm_control("Which statements and queries to run, how many rows to fetch, and when to commit or roll back; it reads SQLCODE, SQLSTATE, row counts, column descriptions and decoded rows")
368 .e2e_testing("tests/client/db2/real_server_test.rs drives Apache Derby 10.17.1.0's network server (a Java DRDA implementation sharing no code with NetGet) with the model mocked: create a table of ten types, insert, commit, select the rows back typed, insert a row built from them, commit, insert and roll back, and a bad statement's error SQLCA. Derby's own ij client reads the table back, a recording proxy feeds both directions to tshark's drda dissector, Derby's statement log shows what ran, and a refused login, the follow-up bound and an LLM failure are tested against it too. Fails rather than skips without java or the Derby jars")
369 .notes("Experimental: one independent server (Derby), not the two the client Beta bar asks for. IBM Db2 itself is not tested: its Community image is linux/amd64 only and needs a privileged container. No LOBs, XML, DECFLOAT, user-defined types or parameter markers; no encrypted security mechanisms (passwords go in the clear, as USRIDPWD sends them); DDM strings are EBCDIC CCSID 37 (UNICODEMGR is not negotiated); S/390 hexadecimal floats are refused. The product id defaults to DNC10170 because Derby refuses any other prefix")
370 .build()
371 }
372 fn example_prompt(&self) -> &'static str {
373 "Connect to the Db2 database SAMPLE at db2.example.com:50000 as db2inst1 and list the ten most recent orders"
374 }
375 fn get_startup_examples(&self) -> StartupExamples {
376 let llm = json!({
377 "type": "open_client",
378 "protocol": "Db2",
379 "remote_addr": "127.0.0.1:50000",
380 "startup_params": {"database": "SAMPLE", "username": "db2inst1", "password": "secret"},
381 "instruction": "List the tables in schema DB2INST1, then count the rows of each"
382 });
383 let mut static_example = llm.clone();
384 static_example["event_handlers"] = json!([
385 {"event_pattern": "db2_connected", "handler": {"type": "static", "actions": [
386 {"type": "execute_query", "sql": "SELECT TABNAME FROM SYSCAT.TABLES WHERE TABSCHEMA = 'DB2INST1'"}
387 ]}},
388 {"event_pattern": "*", "handler": {"type": "static", "actions": []}}
389 ]);
390 let mut scripted = static_example.clone();
391 scripted["event_handlers"][1] = json!({"event_pattern": "db2_query_result", "handler": {
392 "type": "script", "language": "python",
393 "code": "import json,sys\ne=json.load(sys.stdin)['event']\na=[{'type':'execute_query','sql':'SELECT COUNT(*) AS N FROM DB2INST1.'+r['TABNAME']} for r in e.get('rows',[]) if 'TABNAME' in r][:5]\nprint(json.dumps({'actions':a}))"
394 }});
395 StartupExamples::new(llm, scripted, static_example)
396 }
397 fn group_name(&self) -> &'static str {
398 "Database"
399 }
400}
401
402impl Client for Db2ClientProtocol {
403 fn connect(
404 &self,
405 ctx: ConnectContext,
406 ) -> std::pin::Pin<Box<dyn std::future::Future<Output = Result<std::net::SocketAddr>> + Send>>
407 {
408 Box::pin(super::connect(ctx))
409 }
410 fn execute_action(&self, v: Value) -> Result<ClientActionResult> {
411 if v.get("type").and_then(Value::as_str) == Some("disconnect") {
412 return Ok(ClientActionResult::Disconnect);
413 }
414 let (name, data) = request(&v)?;
415 Ok(ClientActionResult::Custom { name, data })
416 }
417}