1use crate::llm::actions::{
2 client_trait::{Client, ClientActionResult},
3 protocol_trait::Protocol,
4 ActionDefinition, ParameterDefinition,
5};
6use crate::protocol::{ConnectContext, EventType};
7use crate::server::clickhouse::{
8 actions::{action, parameter},
9 wire,
10};
11use crate::state::app_state::AppState;
12use anyhow::{bail, ensure, Result};
13use serde_json::{json, Value};
14use std::sync::LazyLock;
15
16#[derive(Default)]
17pub struct ClickhouseClientProtocol;
18impl ClickhouseClientProtocol {
19 pub fn new() -> Self {
20 Self
21 }
22}
23
24fn query_action() -> ActionDefinition {
25 action(
26 "clickhouse_query",
27 "Run a SQL statement that is not an INSERT; its rows arrive as clickhouse_result.",
28 vec![parameter("query", "string", "The SQL statement text", true)],
29 json!({"type":"clickhouse_query","query":"SELECT name, engine FROM system.tables LIMIT 5"}),
30 )
31}
32
33fn insert_action() -> ActionDefinition {
34 action(
35 "clickhouse_insert",
36 "Insert rows: the server describes the columns, Rust encodes the rows to match, and the outcome arrives as clickhouse_result.",
37 vec![
38 parameter("query", "string", "The INSERT statement without data, e.g. INSERT INTO t (id, name) VALUES", true),
39 parameter("rows", "array", "Rows as arrays in the table's column order", true),
40 ],
41 json!({"type":"clickhouse_insert","query":"INSERT INTO events (id, name) VALUES","rows":[[1,"signup"]]}),
42 )
43}
44
45fn disconnect_action() -> ActionDefinition {
46 action(
47 "disconnect",
48 "Close the ClickHouse connection",
49 vec![],
50 json!({"type":"disconnect"}),
51 )
52}
53
54fn actions() -> Vec<ActionDefinition> {
55 vec![query_action(), insert_action(), disconnect_action()]
56}
57
58pub static CONNECTED_EVENT: LazyLock<EventType> = LazyLock::new(|| {
59 EventType::new(
60 "clickhouse_connected",
61 "Logged in to the ClickHouse server.",
62 query_action().example.clone(),
63 )
64 .with_parameters(vec![
65 parameter("server_name", "string", "Server name from its hello", true),
66 parameter(
67 "version",
68 "string",
69 "Server version major.minor.patch",
70 true,
71 ),
72 parameter("revision", "number", "Protocol revision in use", true),
73 parameter("timezone", "string", "Server timezone, when sent", false),
74 ])
75 .with_actions(actions())
76});
77
78pub static RESULT_EVENT: LazyLock<EventType> = LazyLock::new(|| {
79 EventType::new(
80 "clickhouse_result",
81 "The outcome of one query or insert.",
82 query_action().example.clone(),
83 )
84 .with_parameters(vec![
85 parameter("query", "string", "The statement this answers", true),
86 parameter(
87 "ok",
88 "boolean",
89 "False when the server raised an exception",
90 true,
91 ),
92 parameter("columns", "array", "[{name, type}] of the result", false),
93 parameter("rows", "array", "Result rows, in column order", false),
94 parameter(
95 "exception",
96 "object",
97 "{code, message} when ok is false",
98 false,
99 ),
100 parameter("rows_written", "number", "For an insert: rows sent", false),
101 ])
102 .with_actions(actions())
103});
104
105pub fn check(v: &Value) -> Result<()> {
106 let query = v["query"]
107 .as_str()
108 .ok_or_else(|| anyhow::anyhow!("query must be a string"))?;
109 ensure!(
110 !query.trim().is_empty() && query.len() <= wire::MAX_STRING,
111 "query must be 1 byte to 1 MiB"
112 );
113 match v["type"].as_str() {
114 Some("clickhouse_query") => {}
115 Some("clickhouse_insert") => {
116 let rows = v["rows"]
117 .as_array()
118 .ok_or_else(|| anyhow::anyhow!("rows must be an array of arrays"))?;
119 ensure!(
120 !rows.is_empty() && rows.len() <= wire::MAX_ROWS,
121 "1 to {} rows",
122 wire::MAX_ROWS
123 );
124 ensure!(
125 rows.iter().all(Value::is_array),
126 "each row must be an array"
127 );
128 }
129 _ => bail!("Unknown ClickHouse client action"),
130 }
131 Ok(())
132}
133
134impl Protocol for ClickhouseClientProtocol {
135 fn protocol_name(&self) -> &'static str {
136 "ClickHouse"
137 }
138 fn stack_name(&self) -> &'static str {
139 "ETH>IP>TCP>ClickHouse"
140 }
141 fn keywords(&self) -> Vec<&'static str> {
142 vec!["clickhouse", "clickhouse native", "olap", "sql"]
143 }
144 fn description(&self) -> &'static str {
145 "ClickHouse native TCP client: queries and inserts"
146 }
147 fn get_async_actions(&self, _: &AppState) -> Vec<ActionDefinition> {
148 actions()
149 }
150 fn get_sync_actions(&self) -> Vec<ActionDefinition> {
151 vec![]
152 }
153 fn get_event_types(&self) -> Vec<EventType> {
154 vec![CONNECTED_EVENT.clone(), RESULT_EVENT.clone()]
155 }
156 fn get_startup_parameters(&self) -> Vec<ParameterDefinition> {
157 vec![
158 ParameterDefinition {
159 name: "user".into(),
160 type_hint: "string".into(),
161 description: "User to log in as".into(),
162 required: false,
163 example: json!("default"),
164 default: Some(json!(super::DEFAULT_USER)),
165 },
166 ParameterDefinition {
167 name: "password".into(),
168 type_hint: "string".into(),
169 description: "Password for that user".into(),
170 required: false,
171 example: json!("secret"),
172 default: None,
173 },
174 ParameterDefinition {
175 name: "database".into(),
176 type_hint: "string".into(),
177 description: "Default database; empty for the server's own default".into(),
178 required: false,
179 example: json!("analytics"),
180 default: None,
181 },
182 ]
183 }
184 fn metadata(&self) -> crate::protocol::metadata::ProtocolMetadataV2 {
185 use crate::protocol::metadata::*;
186 ProtocolMetadataV2::builder()
187 .state(DevelopmentState::Experimental)
188 .privilege_requirement(PrivilegeRequirement::None)
189 .well_known_port(9000)
190 .implementation("ClickHouse native TCP at protocol revision 54429: hello, uncompressed queries with the external-tables terminator, Data, Progress, ProfileInfo, Totals, Extremes, Log and TableColumns packets, exceptions, and the INSERT header/data exchange")
191 .llm_control("Which statements to run and rows to insert, and what to do with each result")
192 .e2e_testing("tests/client/clickhouse: NetGet's own server; the official ClickHouse server 24.8 as the independent peer")
193 .notes("Result columns are decoded for integers, floats, Bool, String, Date, DateTime and Nullable of those; a result with another type ends the session, because its bytes cannot be skipped without knowing their layout. Strings are capped at 1 MiB, blocks at 16 MiB, 100000 rows. Use clickhouse_insert, not clickhouse_query, for INSERT.")
194 .max_inbound_bytes(wire::MAX_BLOCK_BYTES)
195 .build()
196 }
197 fn example_prompt(&self) -> &'static str {
198 "Connect to ClickHouse at 127.0.0.1:9000 and list the tables in the default database"
199 }
200 fn get_startup_examples(&self) -> crate::llm::actions::StartupExamples {
201 let llm = json!({"type":"open_client","protocol":"clickhouse","remote_addr":"127.0.0.1:9000","instruction":"List the tables and count rows in each"});
202 let mut static_example = llm.clone();
203 static_example["event_handlers"] = json!([
204 {"event_pattern":"clickhouse_connected","handler":{"type":"static","actions":[{"type":"clickhouse_query","query":"SELECT version()"}]}},
205 {"event_pattern":"clickhouse_result","handler":{"type":"static","actions":[{"type":"disconnect"}]}}
206 ]);
207 let mut scripted = static_example.clone();
208 scripted["event_handlers"][1]["handler"] = json!({"type":"script","language":"python","code":"import json,sys\ne=json.load(sys.stdin)['event']\nprint(json.dumps({'actions':[{'type':'clickhouse_query','query':'SHOW TABLES'}] if e['query']=='SELECT version()' else [{'type':'disconnect'}]}))"});
209 crate::llm::actions::StartupExamples::new(llm, scripted, static_example)
210 }
211 fn group_name(&self) -> &'static str {
212 "Database"
213 }
214}
215
216impl Client for ClickhouseClientProtocol {
217 fn connect(
218 &self,
219 ctx: ConnectContext,
220 ) -> std::pin::Pin<Box<dyn std::future::Future<Output = Result<std::net::SocketAddr>> + Send>>
221 {
222 Box::pin(super::connect(ctx))
223 }
224 fn execute_action(&self, v: Value) -> Result<ClientActionResult> {
225 if v["type"] == "disconnect" {
226 return Ok(ClientActionResult::Disconnect);
227 }
228 check(&v)?;
229 Ok(ClientActionResult::Custom {
230 name: v["type"].as_str().unwrap_or_default().to_string(),
231 data: v,
232 })
233 }
234}