1use log::{debug, trace, warn};
5use quinn::{ClientConfig, Endpoint};
6use rustls::pki_types::{CertificateDer, ServerName as RustlsServerName};
7use secrecy::{ExposeSecret, SecretString};
8use serde::{Deserialize, Serialize};
9use std::collections::HashMap;
10use std::net::{SocketAddr, ToSocketAddrs};
11use std::sync::Arc;
12use tokio::time::{Duration, Instant, timeout};
13
14use crate::dsn::{Dsn, Transport};
15use crate::error::{Error, Result};
16use crate::proto;
17use crate::types::Value;
18use crate::validate;
19
20const GEODE_ALPN: &[u8] = b"geode/1";
21const MAX_PROTO_FRAME_BYTES: usize = 8 * 1024 * 1024;
24const DEFAULT_MAX_ROWS: usize = 1_000_000;
26const DEFAULT_MAX_PAGES: usize = 10_000;
28
29pub fn redact_dsn(dsn: &str) -> String {
47 let mut result = dsn.to_string();
48
49 if let Some(scheme_end) = result.find("://") {
52 let after_scheme = scheme_end + 3;
53 if let Some(at_pos) = result[after_scheme..].find('@') {
54 let auth_section = &result[after_scheme..after_scheme + at_pos];
55 if let Some(colon_pos) = auth_section.find(':') {
56 let user = &auth_section[..colon_pos];
58 let rest_start = after_scheme + at_pos;
59 result = format!(
60 "{}{}:{}{}",
61 &result[..after_scheme],
62 user,
63 "[REDACTED]",
64 &result[rest_start..]
65 );
66 }
67 }
68 }
69
70 let patterns = ["password=", "pass="];
73 for pattern in patterns {
74 let lower = result.to_lowercase();
75 if let Some(start) = lower.find(pattern) {
76 let value_start = start + pattern.len();
77 let value_end = result[value_start..]
79 .find('&')
80 .map(|i| value_start + i)
81 .unwrap_or(result.len());
82
83 result = format!(
84 "{}[REDACTED]{}",
85 &result[..value_start],
86 &result[value_end..]
87 );
88 }
89 }
90
91 result
92}
93
94#[derive(Debug, Clone, Serialize, Deserialize)]
98pub struct Column {
99 pub name: String,
101 #[serde(rename = "type")]
103 pub col_type: String,
104}
105
106#[derive(Debug, Clone)]
127pub struct Page {
128 pub columns: Vec<Column>,
130 pub rows: Vec<HashMap<String, Value>>,
132 pub ordered: bool,
134 pub order_keys: Vec<String>,
136 pub final_page: bool,
138}
139
140#[derive(Debug, Clone)]
160pub struct Savepoint {
161 pub name: String,
163}
164
165#[derive(Debug, Clone)]
184pub struct PreparedStatement {
185 query: String,
187 param_names: Vec<String>,
189}
190
191impl PreparedStatement {
192 pub fn new(query: impl Into<String>) -> Self {
196 let query = query.into();
197 let param_names = Self::extract_param_names(&query);
198 Self { query, param_names }
199 }
200
201 fn extract_param_names(query: &str) -> Vec<String> {
203 let mut names = Vec::new();
204 let mut chars = query.chars().peekable();
205
206 while let Some(c) = chars.next() {
207 if c == '$' {
208 let mut name = String::new();
209 while let Some(&next) = chars.peek() {
210 if next.is_ascii_alphanumeric() || next == '_' {
211 name.push(chars.next().unwrap());
212 } else {
213 break;
214 }
215 }
216 if !name.is_empty() && !names.contains(&name) {
217 names.push(name);
218 }
219 }
220 }
221
222 names
223 }
224
225 pub fn query(&self) -> &str {
227 &self.query
228 }
229
230 pub fn param_names(&self) -> &[String] {
232 &self.param_names
233 }
234
235 pub async fn execute(
250 &self,
251 conn: &mut Connection,
252 params: &HashMap<String, crate::types::Value>,
253 ) -> crate::error::Result<(Page, Option<String>)> {
254 for name in &self.param_names {
256 if !params.contains_key(name) {
257 return Err(crate::error::Error::validation(format!(
258 "Missing required parameter: {}",
259 name
260 )));
261 }
262 }
263
264 conn.query_with_params(&self.query, params).await
265 }
266}
267
268#[derive(Debug, Clone)]
270pub struct PlanOperation {
271 pub op_type: String,
273 pub description: String,
275 pub estimated_rows: Option<u64>,
277 pub children: Vec<PlanOperation>,
279}
280
281#[derive(Debug, Clone)]
286pub struct QueryPlan {
287 pub operations: Vec<PlanOperation>,
289 pub estimated_rows: u64,
291 pub raw: serde_json::Value,
293}
294
295#[derive(Debug, Clone)]
299pub struct QueryProfile {
300 pub plan: QueryPlan,
302 pub actual_rows: u64,
304 pub execution_time_ms: f64,
306 pub raw: serde_json::Value,
308}
309
310#[derive(Clone)]
349pub struct Client {
350 transport: Transport,
351 host: String,
352 port: u16,
353 tls_enabled: bool,
354 skip_verify: bool,
355 page_size: usize,
356 hello_name: String,
357 hello_ver: String,
358 conformance: String,
359 username: Option<String>,
360 password: Option<SecretString>,
363 graph: Option<String>,
365 connect_timeout_secs: u64,
367 hello_timeout_secs: u64,
369 idle_timeout_secs: u64,
371}
372
373impl Client {
374 pub fn new(host: impl Into<String>, port: u16) -> Self {
394 Self {
395 transport: Transport::Quic,
396 host: host.into(),
397 port,
398 tls_enabled: true,
399 skip_verify: false,
400 page_size: 1000,
401 hello_name: "geode-rust".to_string(),
402 hello_ver: env!("CARGO_PKG_VERSION").to_string(),
403 conformance: "min".to_string(),
404 username: None,
405 password: None,
406 graph: None,
407 connect_timeout_secs: 10,
408 hello_timeout_secs: 15,
409 idle_timeout_secs: 30,
410 }
411 }
412
413 pub fn from_dsn(dsn_str: &str) -> Result<Self> {
461 let dsn = Dsn::parse(dsn_str)?;
462
463 Ok(Self {
464 transport: dsn.transport(),
465 host: dsn.host().to_string(),
466 port: dsn.port(),
467 tls_enabled: dsn.tls_enabled(),
468 skip_verify: dsn.skip_verify(),
469 page_size: dsn.page_size(),
470 hello_name: dsn.client_name().to_string(),
471 hello_ver: dsn.client_version().to_string(),
472 conformance: dsn.conformance().to_string(),
473 username: dsn.username().map(String::from),
474 password: dsn.password().map(|p| SecretString::from(p.to_string())),
475 graph: dsn.graph().map(String::from),
476 connect_timeout_secs: dsn.connect_timeout_secs().unwrap_or(10),
477 hello_timeout_secs: 15,
478 idle_timeout_secs: 30,
479 })
480 }
481
482 pub fn transport(&self) -> Transport {
484 self.transport
485 }
486
487 pub fn skip_verify(mut self, skip: bool) -> Self {
499 self.skip_verify = skip;
500 self
501 }
502
503 pub fn page_size(mut self, size: usize) -> Self {
512 self.page_size = size;
513 self
514 }
515
516 pub fn client_name(mut self, name: impl Into<String>) -> Self {
524 self.hello_name = name.into();
525 self
526 }
527
528 pub fn client_version(mut self, version: impl Into<String>) -> Self {
534 self.hello_ver = version.into();
535 self
536 }
537
538 pub fn conformance(mut self, level: impl Into<String>) -> Self {
544 self.conformance = level.into();
545 self
546 }
547
548 pub fn graph(mut self, graph: impl Into<String>) -> Self {
557 self.graph = Some(graph.into());
558 self
559 }
560
561 pub fn username(mut self, username: impl Into<String>) -> Self {
577 self.username = Some(username.into());
578 self
579 }
580
581 pub fn password(mut self, password: impl Into<String>) -> Self {
591 self.password = Some(SecretString::from(password.into()));
592 self
593 }
594
595 pub fn connect_timeout(mut self, seconds: u64) -> Self {
604 self.connect_timeout_secs = seconds.max(1);
605 self
606 }
607
608 pub fn hello_timeout(mut self, seconds: u64) -> Self {
617 self.hello_timeout_secs = seconds.max(1);
618 self
619 }
620
621 pub fn idle_timeout(mut self, seconds: u64) -> Self {
630 self.idle_timeout_secs = seconds.max(1);
631 self
632 }
633
634 pub fn validate(&self) -> Result<()> {
662 validate::hostname(&self.host)?;
664
665 validate::port(self.port)?;
667
668 validate::page_size(self.page_size)?;
670
671 Ok(())
672 }
673
674 pub async fn connect(&self) -> Result<Connection> {
705 self.validate()?;
707
708 let password_ref = self.password.as_ref().map(|s| s.expose_secret());
710
711 match self.transport {
712 Transport::Quic => {
713 Connection::new_quic(
714 &self.host,
715 self.port,
716 self.skip_verify,
717 self.page_size,
718 &self.hello_name,
719 &self.hello_ver,
720 &self.conformance,
721 self.username.as_deref(),
722 password_ref,
723 self.graph.as_deref(),
724 self.connect_timeout_secs,
725 self.hello_timeout_secs,
726 self.idle_timeout_secs,
727 )
728 .await
729 }
730 Transport::Grpc => {
731 #[cfg(feature = "grpc")]
732 {
733 Connection::new_grpc(
734 &self.host,
735 self.port,
736 self.tls_enabled,
737 self.skip_verify,
738 self.page_size,
739 self.username.as_deref(),
740 password_ref,
741 self.graph.as_deref(),
742 )
743 .await
744 }
745 #[cfg(not(feature = "grpc"))]
746 {
747 Err(Error::connection(
748 "gRPC transport requires the 'grpc' feature to be enabled",
749 ))
750 }
751 }
752 }
753 }
754}
755
756#[allow(dead_code)]
758enum ConnectionKind {
759 Quic {
761 conn: quinn::Connection,
762 send: quinn::SendStream,
763 recv: quinn::RecvStream,
764 buffer: Vec<u8>,
766 next_request_id: u64,
768 session_id: String,
770 },
771 #[cfg(feature = "grpc")]
773 Grpc {
774 client: Box<crate::grpc::GrpcClient>,
775 },
776}
777
778pub struct Connection {
823 kind: ConnectionKind,
824 #[allow(dead_code)]
826 page_size: usize,
827 in_transaction: bool,
829}
830
831impl Connection {
832 #[allow(clippy::too_many_arguments)]
834 async fn new_quic(
835 host: &str,
836 port: u16,
837 skip_verify: bool,
838 page_size: usize,
839 hello_name: &str,
840 hello_ver: &str,
841 conformance: &str,
842 username: Option<&str>,
843 password: Option<&str>,
844 graph: Option<&str>,
845 connect_timeout_secs: u64,
846 hello_timeout_secs: u64,
847 idle_timeout_secs: u64,
848 ) -> Result<Self> {
849 let mut last_err: Option<Error> = None;
850
851 for attempt in 1..=3 {
852 match Self::connect_quic_once(
853 host,
854 port,
855 skip_verify,
856 page_size,
857 hello_name,
858 hello_ver,
859 conformance,
860 username,
861 password,
862 graph,
863 connect_timeout_secs,
864 hello_timeout_secs,
865 idle_timeout_secs,
866 )
867 .await
868 {
869 Ok(conn) => return Ok(conn),
870 Err(e) => {
871 last_err = Some(e);
872 if attempt < 3 {
873 debug!("Connection attempt {} failed, retrying...", attempt);
874 tokio::time::sleep(Duration::from_millis(150)).await;
875 }
876 }
877 }
878 }
879
880 Err(last_err.unwrap_or_else(|| Error::connection("Failed to connect")))
881 }
882
883 #[cfg(feature = "grpc")]
885 #[allow(clippy::too_many_arguments)]
886 async fn new_grpc(
887 host: &str,
888 port: u16,
889 tls_enabled: bool,
890 skip_verify: bool,
891 page_size: usize,
892 username: Option<&str>,
893 password: Option<&str>,
894 graph: Option<&str>,
895 ) -> Result<Self> {
896 use crate::dsn::Dsn;
897
898 let tls_val = if tls_enabled { "1" } else { "0" };
900 let graph_suffix = graph
901 .map(|g| format!("&graph={}", urlencoding::encode(g)))
902 .unwrap_or_default();
903 let dsn_str = if let (Some(user), Some(pass)) = (username, password) {
904 format!(
905 "grpc://{}:{}@{}:{}?tls={}&insecure={}{}",
906 user, pass, host, port, tls_val, skip_verify, graph_suffix
907 )
908 } else {
909 format!(
910 "grpc://{}:{}?tls={}&insecure={}{}",
911 host, port, tls_val, skip_verify, graph_suffix
912 )
913 };
914
915 let dsn = Dsn::parse(&dsn_str)?;
916 let client = Box::new(crate::grpc::GrpcClient::connect(&dsn).await?);
917
918 Ok(Self {
919 kind: ConnectionKind::Grpc { client },
920 page_size,
921 in_transaction: false,
922 })
923 }
924
925 #[allow(clippy::too_many_arguments)]
926 async fn connect_quic_once(
927 host: &str,
928 port: u16,
929 skip_verify: bool,
930 page_size: usize,
931 hello_name: &str,
932 hello_ver: &str,
933 conformance: &str,
934 username: Option<&str>,
935 password: Option<&str>,
936 graph: Option<&str>,
937 connect_timeout_secs: u64,
938 hello_timeout_secs: u64,
939 idle_timeout_secs: u64,
940 ) -> Result<Self> {
941 debug!("Creating connection to {}:{}", host, port);
942
943 let _ = rustls::crypto::aws_lc_rs::default_provider().install_default();
945
946 let mut client_crypto = if skip_verify {
948 warn!(
951 "TLS certificate verification DISABLED - connection to {}:{} is vulnerable to MITM attacks. \
952 Do NOT use skip_verify in production!",
953 host, port
954 );
955 rustls::ClientConfig::builder_with_protocol_versions(&[&rustls::version::TLS13])
956 .dangerous()
957 .with_custom_certificate_verifier(Arc::new(SkipServerVerification))
958 .with_no_client_auth()
959 } else {
960 let mut root_store = rustls::RootCertStore::empty();
962
963 let cert_result = rustls_native_certs::load_native_certs();
964
965 for err in &cert_result.errors {
967 warn!("Error loading native certificate: {:?}", err);
968 }
969
970 let mut certs_loaded = 0;
971 let mut certs_failed = 0;
972
973 for cert in cert_result.certs {
974 match root_store.add(cert) {
975 Ok(()) => certs_loaded += 1,
976 Err(_) => certs_failed += 1,
977 }
978 }
979
980 if certs_loaded == 0 {
981 return Err(Error::tls(
982 "No system root certificates found. TLS verification cannot proceed. \
983 Either install system CA certificates or use skip_verify(true) for development only.",
984 ));
985 }
986
987 debug!(
988 "Loaded {} system root certificates ({} failed to parse)",
989 certs_loaded, certs_failed
990 );
991
992 rustls::ClientConfig::builder_with_protocol_versions(&[&rustls::version::TLS13])
993 .with_root_certificates(root_store)
994 .with_no_client_auth()
995 };
996
997 client_crypto.alpn_protocols = vec![GEODE_ALPN.to_vec()];
999
1000 let mut client_config = ClientConfig::new(Arc::new(
1001 quinn::crypto::rustls::QuicClientConfig::try_from(client_crypto)
1002 .map_err(|e| Error::connection(format!("Failed to create QUIC config: {}", e)))?,
1003 ));
1004
1005 let mut transport = quinn::TransportConfig::default();
1007 let idle_timeout = Duration::from_secs(idle_timeout_secs.min(146_000 * 365 * 24 * 3600));
1010 transport.max_idle_timeout(Some(idle_timeout.try_into().map_err(|_| {
1011 Error::connection("Idle timeout value too large for QUIC protocol")
1012 })?));
1013 transport.keep_alive_interval(Some(Duration::from_secs(5)));
1014 client_config.transport_config(Arc::new(transport));
1015
1016 let mut endpoint = Endpoint::client(
1019 "0.0.0.0:0"
1020 .parse()
1021 .expect("0.0.0.0:0 is a valid socket address"),
1022 )
1023 .map_err(|e| Error::connection(format!("Failed to create endpoint: {}", e)))?;
1024 endpoint.set_default_client_config(client_config);
1025
1026 let mut resolved_addrs = format!("{}:{}", host, port)
1028 .to_socket_addrs()
1029 .map_err(|e| {
1030 Error::connection(format!(
1031 "Failed to resolve address {}:{} - {}",
1032 host, port, e
1033 ))
1034 })?;
1035
1036 let server_addr: SocketAddr = resolved_addrs
1037 .find(|addr| matches!(addr, SocketAddr::V4(_) | SocketAddr::V6(_)))
1038 .ok_or_else(|| Error::connection("Invalid address: could not resolve host"))?;
1039
1040 debug!("Connecting to {}", server_addr);
1041
1042 let server_name = if skip_verify {
1045 "localhost" } else {
1047 host
1048 };
1049
1050 trace!("Using server name for SNI: {}", server_name);
1051
1052 let conn = timeout(
1053 Duration::from_secs(connect_timeout_secs),
1054 endpoint
1055 .connect(server_addr, server_name)
1056 .map_err(|e| Error::connection(format!("Connection failed: {}", e)))?,
1057 )
1058 .await
1059 .map_err(|_| Error::connection("Connection timeout"))?
1060 .map_err(|e| Error::connection(format!("Failed to establish connection: {}", e)))?;
1061
1062 debug!("Connection established to {}:{}", host, port);
1063
1064 let (mut send, mut recv) = conn
1066 .open_bi()
1067 .await
1068 .map_err(|e| Error::connection(format!("Failed to open stream: {}", e)))?;
1069
1070 let hello_req = proto::HelloRequest {
1072 username: username.unwrap_or("").to_string(),
1073 password: password.unwrap_or("").to_string(),
1074 tenant_id: None,
1075 client_name: hello_name.to_string(),
1076 client_version: hello_ver.to_string(),
1077 wanted_conformance: conformance.to_string(),
1078 graph: graph.map(String::from),
1079 };
1080 let msg = proto::QuicClientMessage {
1081 msg: Some(proto::quic_client_message::Msg::Hello(hello_req)),
1082 };
1083 let data = proto::encode_with_length_prefix(&msg);
1084
1085 send.write_all(&data)
1086 .await
1087 .map_err(|e| Error::connection(format!("Failed to send HELLO: {}", e)))?;
1088
1089 let mut length_buf = [0u8; 4];
1091 timeout(
1092 Duration::from_secs(hello_timeout_secs),
1093 recv.read_exact(&mut length_buf),
1094 )
1095 .await
1096 .map_err(|_| Error::connection("HELLO response timeout"))?
1097 .map_err(|e| Error::connection(format!("Failed to read HELLO response length: {}", e)))?;
1098
1099 let msg_len = u32::from_be_bytes(length_buf) as usize;
1100
1101 if msg_len > MAX_PROTO_FRAME_BYTES {
1102 return Err(Error::limit(format!(
1103 "HELLO response frame size {} bytes exceeds max {} bytes",
1104 msg_len, MAX_PROTO_FRAME_BYTES
1105 )));
1106 }
1107
1108 let mut msg_buf = vec![0u8; msg_len];
1109 recv.read_exact(&mut msg_buf)
1110 .await
1111 .map_err(|e| Error::connection(format!("Failed to read HELLO response body: {}", e)))?;
1112
1113 let hello_response = proto::decode_quic_server_message(&msg_buf)?;
1114
1115 let session_id = match hello_response.msg {
1116 Some(proto::quic_server_message::Msg::Hello(ref hello_resp)) => {
1117 if !hello_resp.success {
1118 return Err(Error::connection(format!(
1119 "Authentication failed: {}",
1120 hello_resp.error_message
1121 )));
1122 }
1123 hello_resp.session_id.clone()
1124 }
1125 _ => {
1126 return Err(Error::connection("Expected HELLO response"));
1127 }
1128 };
1129
1130 debug!("HELLO handshake complete, session_id={}", session_id);
1131
1132 Ok(Self {
1133 kind: ConnectionKind::Quic {
1134 conn,
1135 send,
1136 recv,
1137 buffer: Vec::new(),
1138 next_request_id: 1,
1139 session_id,
1140 },
1141 page_size,
1142 in_transaction: false,
1143 })
1144 }
1145
1146 async fn send_proto_quic(
1148 send: &mut quinn::SendStream,
1149 msg: &proto::QuicClientMessage,
1150 ) -> Result<()> {
1151 let data = proto::encode_with_length_prefix(msg);
1152 send.write_all(&data)
1153 .await
1154 .map_err(|e| Error::connection(format!("Failed to send message: {}", e)))?;
1155 Ok(())
1156 }
1157
1158 async fn read_proto_quic(
1163 recv: &mut quinn::RecvStream,
1164 timeout_secs: u64,
1165 ) -> Result<proto::QuicServerMessage> {
1166 timeout(Duration::from_secs(timeout_secs), async {
1167 let mut length_buf = [0u8; 4];
1169 recv.read_exact(&mut length_buf)
1170 .await
1171 .map_err(|e| Error::connection(format!("Failed to read response length: {}", e)))?;
1172
1173 let msg_len = u32::from_be_bytes(length_buf) as usize;
1174
1175 if msg_len > MAX_PROTO_FRAME_BYTES {
1177 return Err(Error::limit(format!(
1178 "Frame size {} bytes exceeds max {} bytes",
1179 msg_len, MAX_PROTO_FRAME_BYTES
1180 )));
1181 }
1182
1183 let mut msg_buf = vec![0u8; msg_len];
1184 recv.read_exact(&mut msg_buf)
1185 .await
1186 .map_err(|e| Error::connection(format!("Failed to read response body: {}", e)))?;
1187
1188 proto::decode_quic_server_message(&msg_buf)
1189 })
1190 .await
1191 .map_err(|_| Error::timeout())?
1192 }
1193
1194 fn parse_proto_rows_static(
1196 proto_rows: &[proto::Row],
1197 columns: &[Column],
1198 ) -> Result<Vec<HashMap<String, Value>>> {
1199 let mut rows = Vec::with_capacity(proto_rows.len());
1200 for proto_row in proto_rows {
1201 let mut row = HashMap::with_capacity(columns.len());
1202 for (i, col) in columns.iter().enumerate() {
1203 let value = if i < proto_row.values.len() {
1204 crate::convert::proto_to_value(&proto_row.values[i])
1205 } else {
1206 Value::null()
1207 };
1208 row.insert(col.name.clone(), value);
1209 }
1210 rows.push(row);
1211 }
1212 Ok(rows)
1213 }
1214
1215 async fn send_begin_quic(
1217 send: &mut quinn::SendStream,
1218 recv: &mut quinn::RecvStream,
1219 session_id: &str,
1220 ) -> Result<()> {
1221 let msg = proto::QuicClientMessage {
1222 msg: Some(proto::quic_client_message::Msg::Begin(
1223 proto::BeginRequest {
1224 session_id: session_id.to_string(),
1225 ..Default::default()
1226 },
1227 )),
1228 };
1229 Self::send_proto_quic(send, &msg).await?;
1230
1231 let resp = Self::read_proto_quic(recv, 5).await?;
1232 if !matches!(resp.msg, Some(proto::quic_server_message::Msg::Begin(_))) {
1233 return Err(Error::protocol("Expected BEGIN response"));
1234 }
1235 Ok(())
1236 }
1237
1238 async fn send_commit_quic(
1240 send: &mut quinn::SendStream,
1241 recv: &mut quinn::RecvStream,
1242 session_id: &str,
1243 ) -> Result<()> {
1244 let msg = proto::QuicClientMessage {
1245 msg: Some(proto::quic_client_message::Msg::Commit(
1246 proto::CommitRequest {
1247 session_id: session_id.to_string(),
1248 },
1249 )),
1250 };
1251 Self::send_proto_quic(send, &msg).await?;
1252
1253 let resp = Self::read_proto_quic(recv, 5).await?;
1254 if !matches!(resp.msg, Some(proto::quic_server_message::Msg::Commit(_))) {
1255 return Err(Error::protocol("Expected COMMIT response"));
1256 }
1257 Ok(())
1258 }
1259
1260 async fn send_rollback_quic(
1262 send: &mut quinn::SendStream,
1263 recv: &mut quinn::RecvStream,
1264 session_id: &str,
1265 ) -> Result<()> {
1266 let msg = proto::QuicClientMessage {
1267 msg: Some(proto::quic_client_message::Msg::Rollback(
1268 proto::RollbackRequest {
1269 session_id: session_id.to_string(),
1270 },
1271 )),
1272 };
1273 Self::send_proto_quic(send, &msg).await?;
1274
1275 let resp = Self::read_proto_quic(recv, 5).await?;
1276 if !matches!(resp.msg, Some(proto::quic_server_message::Msg::Rollback(_))) {
1277 return Err(Error::protocol("Expected ROLLBACK response"));
1278 }
1279 Ok(())
1280 }
1281
1282 async fn send_savepoint_quic(
1284 send: &mut quinn::SendStream,
1285 recv: &mut quinn::RecvStream,
1286 session_id: &str,
1287 name: &str,
1288 ) -> Result<()> {
1289 let msg = proto::QuicClientMessage {
1290 msg: Some(proto::quic_client_message::Msg::Savepoint(
1291 proto::SavepointRequest {
1292 name: name.to_string(),
1293 session_id: session_id.to_string(),
1294 },
1295 )),
1296 };
1297 Self::send_proto_quic(send, &msg).await?;
1298
1299 let resp = Self::read_proto_quic(recv, 5).await?;
1300 if !matches!(
1301 resp.msg,
1302 Some(proto::quic_server_message::Msg::Savepoint(_))
1303 ) {
1304 return Err(Error::protocol("Expected SAVEPOINT response"));
1305 }
1306 Ok(())
1307 }
1308
1309 async fn send_rollback_to_quic(
1311 send: &mut quinn::SendStream,
1312 recv: &mut quinn::RecvStream,
1313 session_id: &str,
1314 name: &str,
1315 ) -> Result<()> {
1316 let msg = proto::QuicClientMessage {
1317 msg: Some(proto::quic_client_message::Msg::RollbackTo(
1318 proto::RollbackToRequest {
1319 name: name.to_string(),
1320 session_id: session_id.to_string(),
1321 },
1322 )),
1323 };
1324 Self::send_proto_quic(send, &msg).await?;
1325
1326 let resp = Self::read_proto_quic(recv, 5).await?;
1327 if !matches!(
1328 resp.msg,
1329 Some(proto::quic_server_message::Msg::RollbackTo(_))
1330 ) {
1331 return Err(Error::protocol("Expected ROLLBACK_TO response"));
1332 }
1333 Ok(())
1334 }
1335
1336 pub async fn query(&mut self, gql: &str) -> Result<(Page, Option<String>)> {
1366 self.query_with_params(gql, &HashMap::new()).await
1367 }
1368
1369 pub async fn query_with_params(
1409 &mut self,
1410 gql: &str,
1411 params: &HashMap<String, Value>,
1412 ) -> Result<(Page, Option<String>)> {
1413 validate::query(gql)?;
1414 for key in params.keys() {
1415 validate::param_name(key)?;
1416 }
1417 match &mut self.kind {
1418 ConnectionKind::Quic {
1419 send,
1420 recv,
1421 buffer,
1422 session_id,
1423 ..
1424 } => {
1425 Self::query_with_params_quic(
1426 send,
1427 recv,
1428 buffer,
1429 gql,
1430 params,
1431 session_id,
1432 self.page_size,
1433 )
1434 .await
1435 }
1436 #[cfg(feature = "grpc")]
1437 ConnectionKind::Grpc { client } => client.query_with_params(gql, params).await,
1438 }
1439 }
1440
1441 async fn query_with_params_quic(
1443 send: &mut quinn::SendStream,
1444 recv: &mut quinn::RecvStream,
1445 buffer: &mut Vec<u8>,
1446 gql: &str,
1447 params: &HashMap<String, Value>,
1448 session_id: &str,
1449 page_size: usize,
1450 ) -> Result<(Page, Option<String>)> {
1451 let (page, cursor) =
1452 Self::query_with_params_quic_inner(send, recv, buffer, gql, params, session_id).await?;
1453
1454 if !page.final_page {
1456 let mut all_rows = page.rows;
1457 let columns = page.columns;
1458 let mut ordered = page.ordered;
1459 let mut order_keys = page.order_keys;
1460 let mut request_id: u64 = 0;
1461 let mut page_count: usize = 1; loop {
1464 if all_rows.len() > DEFAULT_MAX_ROWS {
1466 return Err(Error::limit(format!(
1467 "Total rows {} exceeds max {}",
1468 all_rows.len(),
1469 DEFAULT_MAX_ROWS
1470 )));
1471 }
1472 page_count += 1;
1473 if page_count > DEFAULT_MAX_PAGES {
1474 return Err(Error::limit(format!(
1475 "Page count {} exceeds max {}",
1476 page_count, DEFAULT_MAX_PAGES
1477 )));
1478 }
1479
1480 request_id += 1;
1481 let pull_req = proto::QuicClientMessage {
1482 msg: Some(proto::quic_client_message::Msg::Pull(proto::PullRequest {
1483 request_id,
1484 page_size: page_size as u32,
1485 session_id: session_id.to_string(),
1486 })),
1487 };
1488 Self::send_proto_quic(send, &pull_req).await?;
1489
1490 let resp = Self::read_proto_quic_buffered(recv, buffer, 30).await?;
1491
1492 let exec_resp = match &resp.msg {
1494 Some(proto::quic_server_message::Msg::Pull(pull)) => pull.response.as_ref(),
1495 Some(proto::quic_server_message::Msg::Execute(e)) => Some(e),
1496 _ => None,
1497 };
1498
1499 let exec_resp = match exec_resp {
1500 Some(e) => e,
1501 None => break,
1502 };
1503
1504 if let Some(proto::execution_response::Payload::Error(ref err)) = exec_resp.payload
1505 {
1506 return Err(Error::Query {
1507 code: err.code.clone(),
1508 message: err.message.clone(),
1509 });
1510 }
1511
1512 if let Some(proto::execution_response::Payload::Page(ref page_data)) =
1513 exec_resp.payload
1514 {
1515 let rows = Self::parse_proto_rows_static(&page_data.rows, &columns)?;
1516 all_rows.extend(rows);
1517 ordered = page_data.ordered;
1518 order_keys = page_data.order_keys.clone();
1519 if page_data.r#final {
1520 break;
1521 }
1522 } else {
1523 break;
1524 }
1525 }
1526
1527 let final_page = Page {
1528 columns,
1529 rows: all_rows,
1530 ordered,
1531 order_keys,
1532 final_page: true,
1533 };
1534 Self::drain_execute_trailers_quic(recv, buffer).await?;
1535 return Ok((final_page, cursor));
1536 }
1537
1538 if page.final_page {
1539 Self::drain_execute_trailers_quic(recv, buffer).await?;
1540 }
1541 Ok((page, cursor))
1542 }
1543
1544 async fn query_with_params_quic_inner(
1546 send: &mut quinn::SendStream,
1547 recv: &mut quinn::RecvStream,
1548 buffer: &mut Vec<u8>,
1549 gql: &str,
1550 params: &HashMap<String, Value>,
1551 session_id: &str,
1552 ) -> Result<(Page, Option<String>)> {
1553 let params_proto: Vec<proto::Param> = params
1555 .iter()
1556 .map(|(k, v)| proto::Param {
1557 name: k.clone(),
1558 value: Some(v.to_proto_value()),
1559 })
1560 .collect();
1561
1562 let exec_req = proto::ExecuteRequest {
1564 session_id: session_id.to_string(),
1565 query: gql.to_string(),
1566 params: params_proto,
1567 };
1568 let msg = proto::QuicClientMessage {
1569 msg: Some(proto::quic_client_message::Msg::Execute(exec_req)),
1570 };
1571 Self::send_proto_quic(send, &msg)
1572 .await
1573 .map_err(|e| Error::query(format!("{}", e)))?;
1574
1575 let resp = Self::read_proto_quic_buffered(recv, buffer, 10).await?;
1577
1578 let exec_resp = match resp.msg {
1579 Some(proto::quic_server_message::Msg::Execute(e)) => e,
1580 _ => return Err(Error::protocol("Expected Execute response")),
1581 };
1582
1583 if let Some(proto::execution_response::Payload::Error(ref err)) = exec_resp.payload {
1585 let _ = Self::try_read_proto_quic_buffered(recv, buffer).await;
1587 return Err(Error::Query {
1588 code: err.code.clone(),
1589 message: err.message.clone(),
1590 });
1591 }
1592
1593 let columns: Vec<Column> = match exec_resp.payload {
1595 Some(proto::execution_response::Payload::Schema(ref s)) => s
1596 .columns
1597 .iter()
1598 .map(|c| Column {
1599 name: c.name.clone(),
1600 col_type: c.r#type.clone(),
1601 })
1602 .collect(),
1603 _ => Vec::new(),
1604 };
1605
1606 trace!("Schema columns: {:?}", columns);
1607
1608 while let Some(inline_resp) = Self::try_read_proto_quic_buffered(recv, buffer).await? {
1611 if let Some(proto::quic_server_message::Msg::Execute(inline_exec)) = inline_resp.msg {
1612 if let Some(proto::execution_response::Payload::Error(ref err)) =
1613 inline_exec.payload
1614 {
1615 return Err(Error::Query {
1616 code: err.code.clone(),
1617 message: err.message.clone(),
1618 });
1619 }
1620
1621 if let Some(proto::execution_response::Payload::Page(ref page_data)) =
1622 inline_exec.payload
1623 {
1624 let rows = Self::parse_proto_rows_static(&page_data.rows, &columns)?;
1625 let page = Page {
1626 columns,
1627 rows,
1628 ordered: page_data.ordered,
1629 order_keys: page_data.order_keys.clone(),
1630 final_page: page_data.r#final,
1631 };
1632 return Ok((page, None));
1633 }
1634 }
1635 }
1636
1637 if columns.is_empty() {
1639 return Ok((
1640 Page {
1641 columns,
1642 rows: Vec::new(),
1643 ordered: false,
1644 order_keys: Vec::new(),
1645 final_page: true,
1646 },
1647 None,
1648 ));
1649 }
1650
1651 if let Some(proto::execution_response::Payload::Page(ref page_data)) = exec_resp.payload {
1653 let rows = Self::parse_proto_rows_static(&page_data.rows, &columns)?;
1654 let page = Page {
1655 columns,
1656 rows,
1657 ordered: page_data.ordered,
1658 order_keys: page_data.order_keys.clone(),
1659 final_page: page_data.r#final,
1660 };
1661 return Ok((page, None));
1662 }
1663
1664 loop {
1666 let resp = Self::read_proto_quic_buffered(recv, buffer, 30).await?;
1667 if let Some(proto::quic_server_message::Msg::Execute(exec_resp)) = resp.msg {
1668 if let Some(proto::execution_response::Payload::Error(ref err)) = exec_resp.payload
1669 {
1670 return Err(Error::Query {
1671 code: err.code.clone(),
1672 message: err.message.clone(),
1673 });
1674 }
1675
1676 if let Some(proto::execution_response::Payload::Page(ref page_data)) =
1677 exec_resp.payload
1678 {
1679 let rows = Self::parse_proto_rows_static(&page_data.rows, &columns)?;
1680 let page = Page {
1681 columns,
1682 rows,
1683 ordered: page_data.ordered,
1684 order_keys: page_data.order_keys.clone(),
1685 final_page: page_data.r#final,
1686 };
1687 return Ok((page, None));
1688 }
1689 }
1690 }
1691 }
1692
1693 fn decode_buffered_quic_message(
1694 buffer: &mut Vec<u8>,
1695 ) -> Result<Option<proto::QuicServerMessage>> {
1696 if buffer.len() < 4 {
1697 return Ok(None);
1698 }
1699
1700 let msg_len = u32::from_be_bytes([buffer[0], buffer[1], buffer[2], buffer[3]]) as usize;
1701 if msg_len > MAX_PROTO_FRAME_BYTES {
1702 return Err(Error::limit(format!(
1703 "Frame size {} bytes exceeds max {} bytes",
1704 msg_len, MAX_PROTO_FRAME_BYTES
1705 )));
1706 }
1707
1708 let total_len = 4 + msg_len;
1709 if buffer.len() < total_len {
1710 return Ok(None);
1711 }
1712
1713 let msg = proto::decode_quic_server_message(&buffer[4..total_len])?;
1714 buffer.drain(..total_len);
1715 Ok(Some(msg))
1716 }
1717
1718 async fn read_proto_quic_buffered(
1719 recv: &mut quinn::RecvStream,
1720 buffer: &mut Vec<u8>,
1721 timeout_secs: u64,
1722 ) -> Result<proto::QuicServerMessage> {
1723 if let Some(msg) = Self::decode_buffered_quic_message(buffer)? {
1724 return Ok(msg);
1725 }
1726
1727 let deadline = Instant::now() + Duration::from_secs(timeout_secs);
1728 loop {
1729 let now = Instant::now();
1730 if now >= deadline {
1731 return Err(Error::timeout());
1732 }
1733
1734 let remaining = deadline.saturating_duration_since(now);
1735 let chunk = timeout(remaining, recv.read_chunk(MAX_PROTO_FRAME_BYTES, true))
1736 .await
1737 .map_err(|_| Error::timeout())?
1738 .map_err(|e| Error::connection(format!("Failed to read response: {}", e)))?
1739 .ok_or_else(|| Error::connection("Stream closed while reading response"))?;
1740 buffer.extend_from_slice(&chunk.bytes);
1741
1742 if let Some(msg) = Self::decode_buffered_quic_message(buffer)? {
1743 return Ok(msg);
1744 }
1745 }
1746 }
1747
1748 async fn try_read_proto_quic_buffered(
1749 recv: &mut quinn::RecvStream,
1750 buffer: &mut Vec<u8>,
1751 ) -> Result<Option<proto::QuicServerMessage>> {
1752 if let Some(msg) = Self::decode_buffered_quic_message(buffer)? {
1753 return Ok(Some(msg));
1754 }
1755
1756 let read_result = timeout(
1757 Duration::from_millis(500),
1758 recv.read_chunk(MAX_PROTO_FRAME_BYTES, true),
1759 )
1760 .await;
1761 let chunk = match read_result {
1762 Ok(Ok(Some(chunk))) => chunk,
1763 Ok(Ok(None)) => return Ok(None),
1764 Ok(Err(e)) => {
1765 return Err(Error::connection(format!("Failed to read response: {}", e)));
1766 }
1767 Err(_) => return Ok(None),
1768 };
1769
1770 buffer.extend_from_slice(&chunk.bytes);
1771 Self::decode_buffered_quic_message(buffer)
1772 }
1773
1774 async fn drain_execute_trailers_quic(
1775 recv: &mut quinn::RecvStream,
1776 buffer: &mut Vec<u8>,
1777 ) -> Result<()> {
1778 while let Some(resp) = Self::try_read_proto_quic_buffered(recv, buffer).await? {
1779 let Some(proto::quic_server_message::Msg::Execute(exec)) = resp.msg else {
1780 break;
1781 };
1782
1783 let is_trailer = matches!(
1784 exec.payload,
1785 Some(proto::execution_response::Payload::Metrics(_))
1786 | Some(proto::execution_response::Payload::Heartbeat(_))
1787 );
1788 if !is_trailer {
1789 break;
1790 }
1791 }
1792 Ok(())
1793 }
1794
1795 pub fn query_sync(
1797 &mut self,
1798 gql: &str,
1799 params: Option<HashMap<String, serde_json::Value>>,
1800 ) -> Result<Page> {
1801 let params_map = params.unwrap_or_default();
1802 let mut params_typed: HashMap<String, Value> = HashMap::new();
1803 for (k, v) in params_map {
1804 let typed_val = crate::types::Value::from_json(v)?;
1805 params_typed.insert(k, typed_val);
1806 }
1807
1808 match tokio::runtime::Handle::try_current() {
1809 Ok(handle) => {
1810 let (page, _cursor) =
1811 handle.block_on(self.query_with_params(gql, ¶ms_typed))?;
1812 Ok(page)
1813 }
1814 Err(_) => {
1815 let rt = tokio::runtime::Runtime::new()
1816 .map_err(|e| Error::query(format!("Failed to create runtime: {}", e)))?;
1817 let (page, _cursor) = rt.block_on(self.query_with_params(gql, ¶ms_typed))?;
1818 Ok(page)
1819 }
1820 }
1821 }
1822
1823 pub async fn begin(&mut self) -> Result<()> {
1848 let result = match &mut self.kind {
1849 ConnectionKind::Quic {
1850 send,
1851 recv,
1852 session_id,
1853 ..
1854 } => Self::send_begin_quic(send, recv, session_id).await,
1855 #[cfg(feature = "grpc")]
1856 ConnectionKind::Grpc { client } => client.begin().await,
1857 };
1858 if result.is_ok() {
1859 self.in_transaction = true;
1860 }
1861 result
1862 }
1863
1864 pub async fn commit(&mut self) -> Result<()> {
1887 let result = match &mut self.kind {
1888 ConnectionKind::Quic {
1889 send,
1890 recv,
1891 session_id,
1892 ..
1893 } => Self::send_commit_quic(send, recv, session_id).await,
1894 #[cfg(feature = "grpc")]
1895 ConnectionKind::Grpc { client } => client.commit().await,
1896 };
1897 if result.is_ok() {
1898 self.in_transaction = false;
1899 }
1900 result
1901 }
1902
1903 pub async fn rollback(&mut self) -> Result<()> {
1927 let result = match &mut self.kind {
1928 ConnectionKind::Quic {
1929 send,
1930 recv,
1931 session_id,
1932 ..
1933 } => Self::send_rollback_quic(send, recv, session_id).await,
1934 #[cfg(feature = "grpc")]
1935 ConnectionKind::Grpc { client } => client.rollback().await,
1936 };
1937 if result.is_ok() {
1938 self.in_transaction = false;
1939 }
1940 result
1941 }
1942
1943 pub async fn savepoint(&mut self, name: &str) -> Result<Savepoint> {
1977 match &mut self.kind {
1978 ConnectionKind::Quic {
1979 send,
1980 recv,
1981 session_id,
1982 ..
1983 } => Self::send_savepoint_quic(send, recv, session_id, name).await?,
1984 #[cfg(feature = "grpc")]
1985 ConnectionKind::Grpc { client } => client.savepoint(name).await?,
1986 }
1987 Ok(Savepoint {
1988 name: name.to_string(),
1989 })
1990 }
1991
1992 pub async fn rollback_to(&mut self, savepoint: &Savepoint) -> Result<()> {
2021 match &mut self.kind {
2022 ConnectionKind::Quic {
2023 send,
2024 recv,
2025 session_id,
2026 ..
2027 } => Self::send_rollback_to_quic(send, recv, session_id, &savepoint.name).await?,
2028 #[cfg(feature = "grpc")]
2029 ConnectionKind::Grpc { client } => client.rollback_to(&savepoint.name).await?,
2030 }
2031 Ok(())
2032 }
2033
2034 pub fn prepare(&self, query: &str) -> Result<PreparedStatement> {
2068 Ok(PreparedStatement::new(query))
2069 }
2070
2071 pub async fn explain(&mut self, gql: &str) -> Result<QueryPlan> {
2104 let explain_query = format!("EXPLAIN {}", gql);
2106 let (_page, _) = self.query(&explain_query).await?;
2107
2108 Ok(QueryPlan {
2111 operations: Vec::new(),
2112 estimated_rows: 0,
2113 raw: serde_json::json!({}),
2114 })
2115 }
2116
2117 pub async fn profile(&mut self, gql: &str) -> Result<QueryProfile> {
2148 let profile_query = format!("PROFILE {}", gql);
2150 let (page, _) = self.query(&profile_query).await?;
2151
2152 let plan = QueryPlan {
2154 operations: Vec::new(),
2155 estimated_rows: 0,
2156 raw: serde_json::json!({}),
2157 };
2158
2159 Ok(QueryProfile {
2160 plan,
2161 actual_rows: page.rows.len() as u64,
2162 execution_time_ms: 0.0,
2163 raw: serde_json::json!({}),
2164 })
2165 }
2166
2167 pub async fn batch(
2205 &mut self,
2206 queries: &[(&str, Option<&HashMap<String, Value>>)],
2207 ) -> Result<Vec<Page>> {
2208 let mut results = Vec::with_capacity(queries.len());
2209
2210 for (query, params) in queries {
2211 let (page, _) = match params {
2212 Some(p) => self.query_with_params(query, p).await?,
2213 None => self.query(query).await?,
2214 };
2215 results.push(page);
2216 }
2217
2218 Ok(results)
2219 }
2220
2221 #[allow(dead_code)]
2224 fn parse_plan_operations(result: &serde_json::Value) -> Vec<PlanOperation> {
2225 let mut operations = Vec::new();
2226
2227 if let Some(ops) = result.get("operations").and_then(|o| o.as_array()) {
2228 for op in ops {
2229 operations.push(Self::parse_single_operation(op));
2230 }
2231 } else if let Some(plan) = result.get("plan") {
2232 operations.push(Self::parse_single_operation(plan));
2234 }
2235
2236 operations
2237 }
2238
2239 #[allow(dead_code)]
2241 fn parse_single_operation(op: &serde_json::Value) -> PlanOperation {
2242 let op_type = op
2243 .get("type")
2244 .or_else(|| op.get("op_type"))
2245 .and_then(|t| t.as_str())
2246 .unwrap_or("Unknown")
2247 .to_string();
2248
2249 let description = op
2250 .get("description")
2251 .or_else(|| op.get("desc"))
2252 .and_then(|d| d.as_str())
2253 .unwrap_or("")
2254 .to_string();
2255
2256 let estimated_rows = op
2257 .get("estimated_rows")
2258 .or_else(|| op.get("rows"))
2259 .and_then(|r| r.as_u64());
2260
2261 let children = op
2262 .get("children")
2263 .and_then(|c| c.as_array())
2264 .map(|arr| arr.iter().map(Self::parse_single_operation).collect())
2265 .unwrap_or_default();
2266
2267 PlanOperation {
2268 op_type,
2269 description,
2270 estimated_rows,
2271 children,
2272 }
2273 }
2274
2275 pub fn close(&mut self) -> Result<()> {
2298 match &mut self.kind {
2299 ConnectionKind::Quic { conn, .. } => {
2300 conn.close(0u32.into(), b"client closing");
2304 Ok(())
2305 }
2306 #[cfg(feature = "grpc")]
2307 ConnectionKind::Grpc { client } => client.close(),
2308 }
2309 }
2310
2311 pub fn in_transaction(&self) -> bool {
2331 self.in_transaction
2332 }
2333
2334 pub fn is_healthy(&self) -> bool {
2335 match &self.kind {
2336 ConnectionKind::Quic { conn, .. } => {
2337 conn.close_reason().is_none()
2339 }
2340 #[cfg(feature = "grpc")]
2341 ConnectionKind::Grpc { .. } => {
2342 true
2344 }
2345 }
2346 }
2347}
2348
2349#[derive(Debug)]
2351struct SkipServerVerification;
2352
2353impl rustls::client::danger::ServerCertVerifier for SkipServerVerification {
2354 fn verify_server_cert(
2355 &self,
2356 _end_entity: &CertificateDer,
2357 _intermediates: &[CertificateDer],
2358 _server_name: &RustlsServerName,
2359 _ocsp_response: &[u8],
2360 _now: rustls::pki_types::UnixTime,
2361 ) -> std::result::Result<rustls::client::danger::ServerCertVerified, rustls::Error> {
2362 Ok(rustls::client::danger::ServerCertVerified::assertion())
2363 }
2364
2365 fn verify_tls12_signature(
2366 &self,
2367 _message: &[u8],
2368 _cert: &CertificateDer,
2369 _dss: &rustls::DigitallySignedStruct,
2370 ) -> std::result::Result<rustls::client::danger::HandshakeSignatureValid, rustls::Error> {
2371 Ok(rustls::client::danger::HandshakeSignatureValid::assertion())
2372 }
2373
2374 fn verify_tls13_signature(
2375 &self,
2376 _message: &[u8],
2377 _cert: &CertificateDer,
2378 _dss: &rustls::DigitallySignedStruct,
2379 ) -> std::result::Result<rustls::client::danger::HandshakeSignatureValid, rustls::Error> {
2380 Ok(rustls::client::danger::HandshakeSignatureValid::assertion())
2381 }
2382
2383 fn supported_verify_schemes(&self) -> Vec<rustls::SignatureScheme> {
2384 vec![
2385 rustls::SignatureScheme::RSA_PKCS1_SHA256,
2386 rustls::SignatureScheme::ECDSA_NISTP256_SHA256,
2387 rustls::SignatureScheme::ED25519,
2388 ]
2389 }
2390}
2391
2392#[cfg(test)]
2393mod tests {
2394 use super::*;
2395
2396 #[test]
2399 fn test_prepared_statement_new() {
2400 let stmt = PreparedStatement::new("MATCH (n:Person {id: $id}) RETURN n");
2401 assert_eq!(stmt.query(), "MATCH (n:Person {id: $id}) RETURN n");
2402 assert_eq!(stmt.param_names(), &["id"]);
2403 }
2404
2405 #[test]
2406 fn test_prepared_statement_multiple_params() {
2407 let stmt = PreparedStatement::new(
2408 "MATCH (p:Person {name: $name}) WHERE p.age > $min_age AND p.city = $city RETURN p",
2409 );
2410 assert!(stmt.query().contains("$name"));
2411 let names = stmt.param_names();
2412 assert_eq!(names.len(), 3);
2413 assert!(names.contains(&"name".to_string()));
2414 assert!(names.contains(&"min_age".to_string()));
2415 assert!(names.contains(&"city".to_string()));
2416 }
2417
2418 #[test]
2419 fn test_prepared_statement_no_params() {
2420 let stmt = PreparedStatement::new("MATCH (n) RETURN n LIMIT 10");
2421 assert!(stmt.param_names().is_empty());
2422 }
2423
2424 #[test]
2425 fn test_prepared_statement_duplicate_params() {
2426 let stmt =
2427 PreparedStatement::new("MATCH (a {id: $id})-[:KNOWS]->(b {id: $id}) RETURN a, b");
2428 assert_eq!(stmt.param_names(), &["id"]);
2430 }
2431
2432 #[test]
2433 fn test_prepared_statement_underscore_params() {
2434 let stmt = PreparedStatement::new("MATCH (n {user_id: $user_id}) RETURN n");
2435 assert_eq!(stmt.param_names(), &["user_id"]);
2436 }
2437
2438 #[test]
2439 fn test_prepared_statement_numeric_params() {
2440 let stmt = PreparedStatement::new("RETURN $param1, $param2, $param123");
2441 let names = stmt.param_names();
2442 assert_eq!(names.len(), 3);
2443 assert!(names.contains(&"param1".to_string()));
2444 assert!(names.contains(&"param2".to_string()));
2445 assert!(names.contains(&"param123".to_string()));
2446 }
2447
2448 #[test]
2451 fn test_plan_operation_struct() {
2452 let op = PlanOperation {
2453 op_type: "NodeScan".to_string(),
2454 description: "Scan Person nodes".to_string(),
2455 estimated_rows: Some(100),
2456 children: vec![],
2457 };
2458 assert_eq!(op.op_type, "NodeScan");
2459 assert_eq!(op.description, "Scan Person nodes");
2460 assert_eq!(op.estimated_rows, Some(100));
2461 assert!(op.children.is_empty());
2462 }
2463
2464 #[test]
2465 fn test_plan_operation_with_children() {
2466 let child = PlanOperation {
2467 op_type: "Filter".to_string(),
2468 description: "Filter by age".to_string(),
2469 estimated_rows: Some(50),
2470 children: vec![],
2471 };
2472 let parent = PlanOperation {
2473 op_type: "Projection".to_string(),
2474 description: "Project name, age".to_string(),
2475 estimated_rows: Some(50),
2476 children: vec![child],
2477 };
2478 assert_eq!(parent.children.len(), 1);
2479 assert_eq!(parent.children[0].op_type, "Filter");
2480 }
2481
2482 #[test]
2485 fn test_query_plan_struct() {
2486 let plan = QueryPlan {
2487 operations: vec![PlanOperation {
2488 op_type: "NodeScan".to_string(),
2489 description: "Full scan".to_string(),
2490 estimated_rows: Some(1000),
2491 children: vec![],
2492 }],
2493 estimated_rows: 1000,
2494 raw: serde_json::json!({"type": "plan"}),
2495 };
2496 assert_eq!(plan.operations.len(), 1);
2497 assert_eq!(plan.estimated_rows, 1000);
2498 }
2499
2500 #[test]
2503 fn test_query_profile_struct() {
2504 let plan = QueryPlan {
2505 operations: vec![],
2506 estimated_rows: 100,
2507 raw: serde_json::json!({}),
2508 };
2509 let profile = QueryProfile {
2510 plan,
2511 actual_rows: 95,
2512 execution_time_ms: 12.5,
2513 raw: serde_json::json!({"type": "profile"}),
2514 };
2515 assert_eq!(profile.actual_rows, 95);
2516 assert!((profile.execution_time_ms - 12.5).abs() < 0.001);
2517 }
2518
2519 #[test]
2522 fn test_page_struct() {
2523 let page = Page {
2524 columns: vec![Column {
2525 name: "x".to_string(),
2526 col_type: "INT".to_string(),
2527 }],
2528 rows: vec![],
2529 ordered: false,
2530 order_keys: vec![],
2531 final_page: true,
2532 };
2533 assert_eq!(page.columns.len(), 1);
2534 assert!(page.rows.is_empty());
2535 assert!(page.final_page);
2536 }
2537
2538 #[test]
2541 fn test_column_struct() {
2542 let col = Column {
2543 name: "age".to_string(),
2544 col_type: "INT".to_string(),
2545 };
2546 assert_eq!(col.name, "age");
2547 assert_eq!(col.col_type, "INT");
2548 }
2549
2550 #[test]
2553 fn test_savepoint_struct() {
2554 let sp = Savepoint {
2555 name: "before_update".to_string(),
2556 };
2557 assert_eq!(sp.name, "before_update");
2558 }
2559
2560 #[test]
2563 fn test_client_builder_defaults() {
2564 let _client = Client::new("localhost", 3141);
2565 }
2567
2568 #[test]
2569 fn test_client_builder_chain() {
2570 let _client = Client::new("example.com", 8443)
2571 .skip_verify(true)
2572 .page_size(500)
2573 .client_name("test-app")
2574 .client_version("2.0.0")
2575 .conformance("full");
2576 }
2578
2579 #[test]
2580 fn test_client_clone() {
2581 let client = Client::new("localhost", 3141).skip_verify(true);
2582 let _cloned = client.clone();
2583 }
2585
2586 #[test]
2589 fn test_parse_plan_operations_empty() {
2590 let result = serde_json::json!({});
2591 let ops = Connection::parse_plan_operations(&result);
2592 assert!(ops.is_empty());
2593 }
2594
2595 #[test]
2596 fn test_parse_plan_operations_array() {
2597 let result = serde_json::json!({
2598 "operations": [
2599 {"type": "NodeScan", "description": "Scan nodes", "estimated_rows": 100},
2600 {"type": "Filter", "description": "Apply filter", "estimated_rows": 50}
2601 ]
2602 });
2603 let ops = Connection::parse_plan_operations(&result);
2604 assert_eq!(ops.len(), 2);
2605 assert_eq!(ops[0].op_type, "NodeScan");
2606 assert_eq!(ops[1].op_type, "Filter");
2607 }
2608
2609 #[test]
2610 fn test_parse_plan_operations_single_plan() {
2611 let result = serde_json::json!({
2612 "plan": {"op_type": "FullScan", "desc": "Full table scan"}
2613 });
2614 let ops = Connection::parse_plan_operations(&result);
2615 assert_eq!(ops.len(), 1);
2616 assert_eq!(ops[0].op_type, "FullScan");
2617 assert_eq!(ops[0].description, "Full table scan");
2618 }
2619
2620 #[test]
2621 fn test_parse_single_operation() {
2622 let op_json = serde_json::json!({
2623 "type": "IndexScan",
2624 "description": "Use index on Person(name)",
2625 "estimated_rows": 25,
2626 "children": [
2627 {"type": "Filter", "description": "Filter results"}
2628 ]
2629 });
2630 let op = Connection::parse_single_operation(&op_json);
2631 assert_eq!(op.op_type, "IndexScan");
2632 assert_eq!(op.description, "Use index on Person(name)");
2633 assert_eq!(op.estimated_rows, Some(25));
2634 assert_eq!(op.children.len(), 1);
2635 assert_eq!(op.children[0].op_type, "Filter");
2636 }
2637
2638 #[test]
2639 fn test_parse_single_operation_minimal() {
2640 let op_json = serde_json::json!({});
2641 let op = Connection::parse_single_operation(&op_json);
2642 assert_eq!(op.op_type, "Unknown");
2643 assert_eq!(op.description, "");
2644 assert_eq!(op.estimated_rows, None);
2645 assert!(op.children.is_empty());
2646 }
2647
2648 #[test]
2649 fn test_parse_single_operation_alt_fields() {
2650 let op_json = serde_json::json!({
2651 "op_type": "Sort",
2652 "desc": "Sort by name ASC",
2653 "rows": 100
2654 });
2655 let op = Connection::parse_single_operation(&op_json);
2656 assert_eq!(op.op_type, "Sort");
2657 assert_eq!(op.description, "Sort by name ASC");
2658 assert_eq!(op.estimated_rows, Some(100));
2659 }
2660
2661 #[test]
2664 fn test_redact_dsn_url_with_password() {
2665 let dsn = "quic://admin:secret123@localhost:3141";
2666 let redacted = redact_dsn(dsn);
2667 assert!(redacted.contains("[REDACTED]"));
2668 assert!(!redacted.contains("secret123"));
2669 assert!(redacted.contains("admin"));
2670 assert!(redacted.contains("localhost"));
2671 }
2672
2673 #[test]
2674 fn test_redact_dsn_url_without_password() {
2675 let dsn = "quic://admin@localhost:3141";
2676 let redacted = redact_dsn(dsn);
2677 assert!(!redacted.contains("[REDACTED]"));
2678 assert!(redacted.contains("admin"));
2679 assert!(redacted.contains("localhost"));
2680 }
2681
2682 #[test]
2683 fn test_redact_dsn_url_no_auth() {
2684 let dsn = "quic://localhost:3141";
2685 let redacted = redact_dsn(dsn);
2686 assert_eq!(redacted, dsn);
2687 }
2688
2689 #[test]
2690 fn test_redact_dsn_query_param_password() {
2691 let dsn = "localhost:3141?username=admin&password=secret123";
2692 let redacted = redact_dsn(dsn);
2693 assert!(redacted.contains("[REDACTED]"));
2694 assert!(!redacted.contains("secret123"));
2695 assert!(redacted.contains("username=admin"));
2696 }
2697
2698 #[test]
2699 fn test_redact_dsn_query_param_pass() {
2700 let dsn = "localhost:3141?user=admin&pass=mysecret";
2701 let redacted = redact_dsn(dsn);
2702 assert!(redacted.contains("[REDACTED]"));
2703 assert!(!redacted.contains("mysecret"));
2704 }
2705
2706 #[test]
2707 fn test_redact_dsn_simple_no_password() {
2708 let dsn = "localhost:3141?insecure=true";
2709 let redacted = redact_dsn(dsn);
2710 assert_eq!(redacted, dsn);
2711 }
2712
2713 #[test]
2714 fn test_redact_dsn_url_with_query_and_password() {
2715 let dsn = "quic://user:pass@localhost:3141?insecure=true";
2716 let redacted = redact_dsn(dsn);
2717 assert!(redacted.contains("[REDACTED]"));
2718 assert!(!redacted.contains(":pass@"));
2719 assert!(redacted.contains("insecure=true"));
2720 }
2721
2722 #[test]
2725 fn test_client_validate_valid() {
2726 let client = Client::new("localhost", 3141);
2727 assert!(client.validate().is_ok());
2728 }
2729
2730 #[test]
2731 fn test_client_validate_valid_hostname() {
2732 let client = Client::new("geode.example.com", 3141);
2733 assert!(client.validate().is_ok());
2734 }
2735
2736 #[test]
2737 fn test_client_validate_valid_ipv4() {
2738 let client = Client::new("192.168.1.1", 8443);
2739 assert!(client.validate().is_ok());
2740 }
2741
2742 #[test]
2743 fn test_client_validate_invalid_hostname_hyphen_start() {
2744 let client = Client::new("-invalid", 3141);
2745 assert!(client.validate().is_err());
2746 }
2747
2748 #[test]
2749 fn test_client_validate_invalid_hostname_hyphen_end() {
2750 let client = Client::new("invalid-", 3141);
2751 assert!(client.validate().is_err());
2752 }
2753
2754 #[test]
2755 fn test_client_validate_invalid_port_zero() {
2756 let client = Client::new("localhost", 0);
2757 assert!(client.validate().is_err());
2758 }
2759
2760 #[test]
2761 fn test_client_validate_invalid_page_size_zero() {
2762 let client = Client::new("localhost", 3141).page_size(0);
2763 assert!(client.validate().is_err());
2764 }
2765
2766 #[test]
2767 fn test_client_validate_invalid_page_size_too_large() {
2768 let client = Client::new("localhost", 3141).page_size(200_000);
2769 assert!(client.validate().is_err());
2770 }
2771
2772 #[test]
2773 fn test_client_validate_with_all_options() {
2774 let client = Client::new("geode.example.com", 8443)
2775 .skip_verify(true)
2776 .page_size(500)
2777 .username("admin")
2778 .password("secret")
2779 .connect_timeout(15)
2780 .hello_timeout(10)
2781 .idle_timeout(60);
2782 assert!(client.validate().is_ok());
2783 }
2784
2785 #[test]
2787 fn test_client_extreme_timeout_values() {
2788 let _client = Client::new("localhost", 3141)
2790 .connect_timeout(u64::MAX)
2791 .hello_timeout(u64::MAX)
2792 .idle_timeout(u64::MAX);
2793 }
2795
2796 #[test]
2797 fn test_convert_edge_uses_type_field() {
2798 let edge = proto::EdgeValue {
2799 id: 100,
2800 from_id: 1,
2801 to_id: 2,
2802 label: "KNOWS".to_string(),
2803 properties: vec![],
2804 };
2805 let proto_val = proto::Value {
2806 kind: Some(proto::value::Kind::EdgeVal(edge)),
2807 };
2808 let val = crate::convert::proto_to_value(&proto_val);
2809 let obj = val.as_object().unwrap();
2810 assert_eq!(obj.get("type").unwrap().as_string().unwrap(), "KNOWS");
2811 assert!(
2812 obj.get("label").is_none(),
2813 "edge should not have 'label' field"
2814 );
2815 }
2816
2817 #[test]
2818 fn test_convert_edge_uses_start_end_node() {
2819 let edge = proto::EdgeValue {
2820 id: 100,
2821 from_id: 42,
2822 to_id: 99,
2823 label: "LIKES".to_string(),
2824 properties: vec![],
2825 };
2826 let proto_val = proto::Value {
2827 kind: Some(proto::value::Kind::EdgeVal(edge)),
2828 };
2829 let val = crate::convert::proto_to_value(&proto_val);
2830 let obj = val.as_object().unwrap();
2831 assert_eq!(obj.get("start_node").unwrap().as_int().unwrap(), 42);
2832 assert_eq!(obj.get("end_node").unwrap().as_int().unwrap(), 99);
2833 assert!(obj.get("from_id").is_none());
2834 assert!(obj.get("to_id").is_none());
2835 }
2836
2837 #[test]
2838 fn test_convert_edge_with_properties() {
2839 let edge = proto::EdgeValue {
2840 id: 100,
2841 from_id: 1,
2842 to_id: 2,
2843 label: "KNOWS".to_string(),
2844 properties: vec![proto::MapEntry {
2845 key: "since".to_string(),
2846 value: Some(proto::Value {
2847 kind: Some(proto::value::Kind::IntVal(proto::IntValue {
2848 value: 2020,
2849 kind: 1,
2850 })),
2851 }),
2852 }],
2853 };
2854 let proto_val = proto::Value {
2855 kind: Some(proto::value::Kind::EdgeVal(edge)),
2856 };
2857 let val = crate::convert::proto_to_value(&proto_val);
2858 let obj = val.as_object().unwrap();
2859 let props = obj.get("properties").unwrap().as_object().unwrap();
2860 assert_eq!(props.get("since").unwrap().as_int().unwrap(), 2020);
2861 }
2862
2863 #[test]
2864 fn test_convert_node_fields() {
2865 let node = proto::NodeValue {
2866 id: 42,
2867 labels: vec!["Person".to_string()],
2868 properties: vec![proto::MapEntry {
2869 key: "name".to_string(),
2870 value: Some(proto::Value {
2871 kind: Some(proto::value::Kind::StringVal(proto::StringValue {
2872 value: "Alice".to_string(),
2873 kind: 1,
2874 })),
2875 }),
2876 }],
2877 };
2878 let proto_val = proto::Value {
2879 kind: Some(proto::value::Kind::NodeVal(node)),
2880 };
2881 let val = crate::convert::proto_to_value(&proto_val);
2882 let obj = val.as_object().unwrap();
2883 assert_eq!(obj.get("id").unwrap().as_int().unwrap(), 42);
2884 let labels = obj.get("labels").unwrap().as_array().unwrap();
2885 assert_eq!(labels.len(), 1);
2886 let props = obj.get("properties").unwrap().as_object().unwrap();
2887 assert_eq!(props.get("name").unwrap().as_string().unwrap(), "Alice");
2888 }
2889
2890 fn null_proto_value() -> proto::Value {
2893 proto::Value {
2894 kind: Some(proto::value::Kind::NullVal(proto::NullValue {})),
2895 }
2896 }
2897
2898 fn node_proto_value(id: u64) -> proto::Value {
2899 proto::Value {
2900 kind: Some(proto::value::Kind::NodeVal(proto::NodeValue {
2901 id,
2902 labels: vec!["Person".to_string()],
2903 properties: vec![],
2904 })),
2905 }
2906 }
2907
2908 #[test]
2909 fn test_real_node_row_kept() {
2910 let columns = vec![Column {
2911 name: "n".to_string(),
2912 col_type: "NODE".to_string(),
2913 }];
2914 let proto_rows = vec![proto::Row {
2915 values: vec![node_proto_value(42)],
2916 }];
2917 let rows = Connection::parse_proto_rows_static(&proto_rows, &columns).unwrap();
2918 assert_eq!(rows.len(), 1, "real NODE row must be kept");
2919 let obj = rows[0].get("n").unwrap().as_object().unwrap();
2920 assert_eq!(obj.get("id").unwrap().as_int().unwrap(), 42);
2921 }
2922
2923 #[test]
2924 fn test_scalar_null_row_kept() {
2925 let columns = vec![Column {
2928 name: "x".to_string(),
2929 col_type: "STRING".to_string(),
2930 }];
2931 let proto_rows = vec![proto::Row {
2932 values: vec![null_proto_value()],
2933 }];
2934 let rows = Connection::parse_proto_rows_static(&proto_rows, &columns).unwrap();
2935 assert_eq!(rows.len(), 1, "scalar null row must be kept");
2936 assert!(rows[0].get("x").unwrap().is_null());
2937 }
2938}