1use crate::llm::actions::{
2 protocol_trait::{ActionResult, Protocol, Server},
3 ActionDefinition, Parameter, ParameterDefinition,
4};
5use crate::protocol::log_template::LogTemplate;
6use crate::protocol::{EventType, SpawnContext};
7use crate::state::app_state::AppState;
8use anyhow::{bail, ensure, Result};
9use serde_json::{json, Value};
10use std::sync::LazyLock;
11
12#[derive(Default)]
13pub struct ClickhouseProtocol;
14impl ClickhouseProtocol {
15 pub fn new() -> Self {
16 Self
17 }
18}
19
20pub fn parameter(name: &str, kind: &str, description: &str, required: bool) -> Parameter {
21 Parameter {
22 name: name.into(),
23 type_hint: kind.into(),
24 description: description.into(),
25 required,
26 }
27}
28
29pub fn action(
30 name: &str,
31 description: &str,
32 parameters: Vec<Parameter>,
33 example: Value,
34) -> ActionDefinition {
35 let log_template = match name {
36 "clickhouse_result" => LogTemplate::new()
37 .with_info("-> ClickHouse result {preview(columns,80)} {preview(rows,80)}"),
38 "clickhouse_exception" => {
39 LogTemplate::new().with_info("-> ClickHouse exception {code}: {message}")
40 }
41 "clickhouse_query" => {
42 LogTemplate::new().with_info("-> ClickHouse query {preview(query,100)}")
43 }
44 _ => LogTemplate::new().with_info(format!("-> ClickHouse {name}")),
45 };
46 ActionDefinition {
47 name: name.into(),
48 description: description.into(),
49 parameters,
50 example,
51 log_template: Some(log_template),
52 }
53}
54
55pub fn columns_param() -> Parameter {
56 parameter(
57 "columns",
58 "array",
59 "[{name, type}] with types UInt8..UInt64, Int8..Int64, Float32, Float64, Bool, String, Date, DateTime or Nullable(...) of those",
60 true,
61 )
62}
63
64fn result_action() -> ActionDefinition {
65 action(
66 "clickhouse_result",
67 "Answer a query with a result set. Rust sends the header block, the rows, progress and end of stream.",
68 vec![
69 columns_param(),
70 parameter("rows", "array", "Rows as arrays in column order; Date as YYYY-MM-DD, DateTime as YYYY-MM-DD hh:mm:ss, null for Nullable", true),
71 ],
72 json!({"type":"clickhouse_result","columns":[{"name":"x","type":"UInt8"},{"name":"s","type":"String"}],"rows":[[1,"a"],[2,"b"]]}),
73 )
74}
75
76fn ok_action() -> ActionDefinition {
77 action(
78 "clickhouse_ok",
79 "Finish a statement that returns no rows (CREATE, DROP, SET, an accepted INSERT's data).",
80 vec![],
81 json!({"type":"clickhouse_ok"}),
82 )
83}
84
85fn insert_action() -> ActionDefinition {
86 action(
87 "clickhouse_insert",
88 "Accept an INSERT and describe the table's columns; the client then sends its rows, which arrive as clickhouse_insert_data.",
89 vec![columns_param()],
90 json!({"type":"clickhouse_insert","columns":[{"name":"id","type":"UInt32"},{"name":"name","type":"String"}]}),
91 )
92}
93
94fn exception_action() -> ActionDefinition {
95 action(
96 "clickhouse_exception",
97 "Fail the query with a ClickHouse exception the client shows as Code: N. DB::Exception: message.",
98 vec![
99 parameter("code", "number", "Error code, e.g. 60 UNKNOWN_TABLE, 62 SYNTAX_ERROR, 497 ACCESS_DENIED; default 1002", false),
100 parameter("message", "string", "Exception message", true),
101 ],
102 json!({"type":"clickhouse_exception","code":60,"message":"Table default.missing does not exist"}),
103 )
104}
105
106pub static QUERY_EVENT: LazyLock<EventType> = LazyLock::new(|| {
107 EventType::new(
108 "clickhouse_query",
109 "A client ran a SQL query; answer with rows, no rows, acceptance of an INSERT, or an exception.",
110 result_action().example.clone(),
111 )
112 .with_parameters(vec![
113 parameter("query", "string", "The SQL statement text", true),
114 parameter("query_id", "string", "Query id the client set, often empty", true),
115 parameter("database", "string", "Default database from the client's hello", true),
116 parameter("user", "string", "User the client logged in as", true),
117 parameter("settings", "object", "Settings the client sent with the query", true),
118 parameter("remote_addr", "string", "Client address and port", true),
119 ])
120 .with_actions(vec![result_action(), ok_action(), insert_action(), exception_action()])
121});
122
123pub static INSERT_DATA_EVENT: LazyLock<EventType> = LazyLock::new(|| {
124 EventType::new(
125 "clickhouse_insert_data",
126 "The rows of an INSERT you accepted have arrived; confirm or fail it.",
127 ok_action().example.clone(),
128 )
129 .with_parameters(vec![
130 parameter("query", "string", "The INSERT statement", true),
131 parameter(
132 "columns",
133 "array",
134 "[{name, type}] as the client sent them",
135 true,
136 ),
137 parameter("rows", "array", "The inserted rows, in column order", true),
138 parameter("user", "string", "User the client logged in as", true),
139 ])
140 .with_actions(vec![ok_action(), exception_action()])
141});
142
143impl Protocol for ClickhouseProtocol {
144 fn protocol_name(&self) -> &'static str {
145 "ClickHouse"
146 }
147 fn stack_name(&self) -> &'static str {
148 "ETH>IP>TCP>ClickHouse"
149 }
150 fn description(&self) -> &'static str {
151 "ClickHouse native TCP server whose query results and inserts the handler decides"
152 }
153 fn keywords(&self) -> Vec<&'static str> {
154 vec![
155 "clickhouse",
156 "clickhouse native",
157 "olap",
158 "columnar database",
159 "sql",
160 ]
161 }
162 fn get_async_actions(&self, _: &AppState) -> Vec<ActionDefinition> {
163 vec![]
164 }
165 fn get_sync_actions(&self) -> Vec<ActionDefinition> {
166 vec![
167 result_action(),
168 ok_action(),
169 insert_action(),
170 exception_action(),
171 ]
172 }
173 fn get_event_types(&self) -> Vec<EventType> {
174 vec![QUERY_EVENT.clone(), INSERT_DATA_EVENT.clone()]
175 }
176 fn get_startup_parameters(&self) -> Vec<ParameterDefinition> {
177 vec![
178 ParameterDefinition {
179 name: "user".into(),
180 type_hint: "string".into(),
181 description: "With password: the only user allowed to log in, checked in Rust; without both, every login is accepted".into(),
182 required: false,
183 example: json!("default"),
184 default: None,
185 },
186 ParameterDefinition {
187 name: "password".into(),
188 type_hint: "string".into(),
189 description: "With user: the password checked in Rust".into(),
190 required: false,
191 example: json!("secret"),
192 default: None,
193 },
194 ParameterDefinition {
195 name: "idle_timeout_secs".into(),
196 type_hint: "number".into(),
197 description: "Seconds a connection may stay silent between packets (1..=86400)".into(),
198 required: false,
199 example: json!(600),
200 default: Some(json!(super::IDLE_TIMEOUT.as_secs())),
201 },
202 ]
203 }
204 fn metadata(&self) -> crate::protocol::metadata::ProtocolMetadataV2 {
205 use crate::protocol::metadata::*;
206 ProtocolMetadataV2::builder()
207 .state(DevelopmentState::Experimental)
208 .privilege_requirement(PrivilegeRequirement::None)
209 .well_known_port(9000)
210 .implementation("ClickHouse native TCP at protocol revision 54429, hand-written: hello, ping, query with settings, external-table terminator, Data blocks with typed columns, LZ4 compressed frames with CityHash128 checksums, INSERT header and row blocks, progress, exceptions, end of stream")
211 .llm_control("Each query's result (columns and rows), acceptance of DDL, the table shape for an INSERT and whether its rows are accepted, and exceptions")
212 .e2e_testing("tests/server/clickhouse: raw packets and bounds; the official clickhouse-client 24.8 (with and without compression) and Python clickhouse-driver as independent clients")
213 .notes("Column types: UInt8..UInt64, Int8..Int64, Float32/64, Bool, String, Date, DateTime and Nullable of those; no arrays, decimals, LowCardinality or external tables. opensrv-clickhouse was not used: it allocates from wire lengths with no bound. Strings are capped at 1 MiB, a block at 16 MiB, 1000 columns and 100000 rows; each connection may be silent idle_timeout_secs. A handler failure answers an exception (code 1002, or 202 when the backend is saturated) with a generic message, never a fabricated result.")
214 .request_only("Every packet answers the client's hello, ping or query")
215 .answers_on_failure()
216 .max_inbound_bytes(super::wire::MAX_BLOCK_BYTES)
217 .build()
218 }
219 fn example_prompt(&self) -> &'static str {
220 "ClickHouse server on port 9000 with a web_events table the handler answers queries about"
221 }
222 fn get_startup_examples(&self) -> crate::llm::actions::StartupExamples {
223 let llm = json!({"type":"open_server","base_stack":"clickhouse","port":9000,"instruction":"Pretend to hold a web_events table (ts DateTime, url String, hits UInt32) and answer queries about it"});
224 let mut scripted = llm.clone();
225 scripted["event_handlers"] = json!([{"event_pattern":"clickhouse_query","handler":{"type":"script","language":"python","code":"import json,sys\nq=json.load(sys.stdin)['event']['query'].strip().rstrip(';')\na={'type':'clickhouse_result','columns':[{'name':'version()','type':'String'}],'rows':[['24.8.0']]} if q.lower()=='select version()' else {'type':'clickhouse_exception','code':62,'message':'Unsupported query'}\nprint(json.dumps({'actions':[a]}))"}}]);
226 let mut static_example = llm.clone();
227 static_example["event_handlers"] = json!([{"event_pattern":"clickhouse_query","handler":{"type":"static","actions":[{"type":"clickhouse_result","columns":[{"name":"x","type":"UInt8"}],"rows":[[1]]}]}}]);
228 crate::llm::actions::StartupExamples::new(llm, scripted, static_example)
229 }
230 fn group_name(&self) -> &'static str {
231 "Database"
232 }
233}
234
235impl Server for ClickhouseProtocol {
236 fn spawn(
237 &self,
238 ctx: SpawnContext,
239 ) -> std::pin::Pin<Box<dyn std::future::Future<Output = Result<std::net::SocketAddr>> + Send>>
240 {
241 Box::pin(super::spawn(ctx))
242 }
243 fn execute_action(&self, v: Value) -> Result<ActionResult> {
244 let name = v["type"].as_str().unwrap_or_default().to_string();
245 match name.as_str() {
246 "clickhouse_result" => {
247 super::wire::Block::from_json(&v["columns"], &v["rows"])?;
248 }
249 "clickhouse_insert" => {
250 let block = super::wire::Block::from_json(&v["columns"], &Value::Null)?;
251 ensure!(
252 !block.columns.is_empty(),
253 "an INSERT needs at least one column"
254 );
255 }
256 "clickhouse_exception" => {
257 ensure!(v["message"].is_string(), "message must be a string");
258 ensure!(
259 v.get("code").is_none_or(|c| c
260 .as_i64()
261 .is_some_and(|c| (0..=i64::from(i32::MAX)).contains(&c))),
262 "code must be a non-negative integer"
263 );
264 }
265 "clickhouse_ok" => {}
266 _ => bail!("Unknown ClickHouse server action"),
267 }
268 Ok(ActionResult::Custom { name, data: v })
269 }
270}