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, 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
29#[allow(clippy::too_many_arguments)]
33fn build_hello_request(
34 username: Option<&str>,
35 password: Option<&str>,
36 hello_name: &str,
37 hello_ver: &str,
38 conformance: &str,
39 graph: Option<&str>,
40 tenant_id: Option<&str>,
41 role: Option<&str>,
42) -> proto::HelloRequest {
43 proto::HelloRequest {
44 username: username.unwrap_or("").to_string(),
45 password: password.unwrap_or("").to_string(),
46 tenant_id: tenant_id.filter(|t| !t.is_empty()).map(String::from),
47 client_name: hello_name.to_string(),
48 client_version: hello_ver.to_string(),
49 wanted_conformance: conformance.to_string(),
50 graph: graph.filter(|g| !g.is_empty()).map(String::from),
51 role: role.filter(|r| !r.is_empty()).map(String::from),
52 }
53}
54
55#[allow(dead_code)] fn redact_dsn(dsn: &str) -> String {
59 let mut result = dsn.to_string();
60
61 if let Some(scheme_end) = result.find("://") {
64 let after_scheme = scheme_end + 3;
65 if let Some(at_pos) = result[after_scheme..].find('@') {
66 let auth_section = &result[after_scheme..after_scheme + at_pos];
67 if let Some(colon_pos) = auth_section.find(':') {
68 let user = &auth_section[..colon_pos];
70 let rest_start = after_scheme + at_pos;
71 result = format!(
72 "{}{}:{}{}",
73 &result[..after_scheme],
74 user,
75 "[REDACTED]",
76 &result[rest_start..]
77 );
78 }
79 }
80 }
81
82 let patterns = ["password=", "pass="];
85 for pattern in patterns {
86 let lower = result.to_lowercase();
87 if let Some(start) = lower.find(pattern) {
88 let value_start = start + pattern.len();
89 let value_end = result[value_start..]
91 .find('&')
92 .map(|i| value_start + i)
93 .unwrap_or(result.len());
94
95 result = format!(
96 "{}[REDACTED]{}",
97 &result[..value_start],
98 &result[value_end..]
99 );
100 }
101 }
102
103 result
104}
105
106#[derive(Debug, Clone, Serialize, Deserialize)]
110pub struct Column {
111 pub name: String,
113 #[serde(rename = "type")]
115 pub col_type: String,
116}
117
118#[derive(Debug, Clone)]
139pub struct Page {
140 pub columns: Vec<Column>,
142 pub rows: Vec<HashMap<String, Value>>,
144 pub ordered: bool,
146 pub order_keys: Vec<String>,
148 pub final_page: bool,
150}
151
152#[derive(Debug, Clone)]
172pub struct Savepoint {
173 pub name: String,
175}
176
177#[derive(Debug, Clone)]
196pub struct PreparedStatement {
197 query: String,
199 param_names: Vec<String>,
201}
202
203impl PreparedStatement {
204 pub fn new(query: impl Into<String>) -> Self {
208 let query = query.into();
209 let param_names = Self::extract_param_names(&query);
210 Self { query, param_names }
211 }
212
213 fn extract_param_names(query: &str) -> Vec<String> {
215 let mut names = Vec::new();
216 let mut chars = query.chars().peekable();
217
218 while let Some(c) = chars.next() {
219 if c == '$' {
220 let mut name = String::new();
221 while let Some(&next) = chars.peek() {
222 if next.is_ascii_alphanumeric() || next == '_' {
223 name.push(chars.next().unwrap());
224 } else {
225 break;
226 }
227 }
228 if !name.is_empty() && !names.contains(&name) {
229 names.push(name);
230 }
231 }
232 }
233
234 names
235 }
236
237 pub fn query(&self) -> &str {
239 &self.query
240 }
241
242 pub fn param_names(&self) -> &[String] {
244 &self.param_names
245 }
246
247 pub async fn execute(
262 &self,
263 conn: &mut Connection,
264 params: &HashMap<String, crate::types::Value>,
265 ) -> crate::error::Result<(Page, Option<String>)> {
266 for name in &self.param_names {
268 if !params.contains_key(name) {
269 return Err(crate::error::Error::validation(format!(
270 "Missing required parameter: {}",
271 name
272 )));
273 }
274 }
275
276 conn.query_with_params(&self.query, params).await
277 }
278}
279
280#[derive(Debug, Clone)]
282pub struct PlanOperation {
283 pub op_type: String,
285 pub description: String,
287 pub estimated_rows: Option<u64>,
289 pub children: Vec<PlanOperation>,
291}
292
293#[derive(Debug, Clone)]
298pub struct QueryPlan {
299 pub operations: Vec<PlanOperation>,
301 pub estimated_rows: u64,
303 pub raw: serde_json::Value,
305}
306
307#[derive(Debug, Clone)]
311pub struct QueryProfile {
312 pub plan: QueryPlan,
314 pub actual_rows: u64,
316 pub execution_time_ms: f64,
318 pub raw: serde_json::Value,
320}
321
322#[derive(Clone)]
361pub struct Client {
362 transport: Transport,
363 host: String,
364 port: u16,
365 tls_enabled: bool,
366 skip_verify: bool,
367 page_size: usize,
368 hello_name: String,
369 hello_ver: String,
370 conformance: String,
371 username: Option<String>,
372 password: Option<SecretString>,
375 graph: Option<String>,
377 tenant_id: Option<String>,
379 role: Option<String>,
381 connect_timeout_secs: u64,
383 hello_timeout_secs: u64,
385 idle_timeout_secs: u64,
387}
388
389impl Client {
390 pub fn new(host: impl Into<String>, port: u16) -> Self {
410 Self {
411 transport: Transport::Quic,
412 host: host.into(),
413 port,
414 tls_enabled: true,
415 skip_verify: false,
416 page_size: 1000,
417 hello_name: "geode-rust".to_string(),
418 hello_ver: env!("CARGO_PKG_VERSION").to_string(),
419 conformance: "min".to_string(),
420 username: None,
421 password: None,
422 graph: None,
423 tenant_id: None,
424 role: None,
425 connect_timeout_secs: 10,
426 hello_timeout_secs: 5,
427 idle_timeout_secs: 30,
428 }
429 }
430
431 pub fn from_dsn(dsn_str: &str) -> Result<Self> {
479 let dsn = Dsn::parse(dsn_str)?;
480
481 Ok(Self {
482 transport: dsn.transport(),
483 host: dsn.host().to_string(),
484 port: dsn.port(),
485 tls_enabled: dsn.tls_enabled(),
486 skip_verify: dsn.skip_verify(),
487 page_size: dsn.page_size(),
488 hello_name: dsn.client_name().to_string(),
489 hello_ver: dsn.client_version().to_string(),
490 conformance: dsn.conformance().to_string(),
491 username: dsn.username().map(String::from),
492 password: dsn.password().map(|p| SecretString::from(p.to_string())),
493 graph: dsn.graph().map(String::from),
494 tenant_id: dsn.tenant_id().map(String::from),
495 role: dsn.role().map(String::from),
496 connect_timeout_secs: dsn.connect_timeout_secs().unwrap_or(10),
497 hello_timeout_secs: 5,
498 idle_timeout_secs: 30,
499 })
500 }
501
502 pub fn transport(&self) -> Transport {
504 self.transport
505 }
506
507 pub fn skip_verify(mut self, skip: bool) -> Self {
519 self.skip_verify = skip;
520 self
521 }
522
523 pub fn page_size(mut self, size: usize) -> Self {
532 self.page_size = size;
533 self
534 }
535
536 pub fn client_name(mut self, name: impl Into<String>) -> Self {
544 self.hello_name = name.into();
545 self
546 }
547
548 pub fn client_version(mut self, version: impl Into<String>) -> Self {
554 self.hello_ver = version.into();
555 self
556 }
557
558 pub fn conformance(mut self, level: impl Into<String>) -> Self {
564 self.conformance = level.into();
565 self
566 }
567
568 pub fn graph(mut self, graph: impl Into<String>) -> Self {
577 self.graph = Some(graph.into());
578 self
579 }
580
581 pub fn tenant_id(mut self, tenant_id: impl Into<String>) -> Self {
583 self.tenant_id = Some(tenant_id.into());
584 self
585 }
586
587 pub fn role(mut self, role: impl Into<String>) -> Self {
589 self.role = Some(role.into());
590 self
591 }
592
593 pub fn username(mut self, username: impl Into<String>) -> Self {
609 self.username = Some(username.into());
610 self
611 }
612
613 pub fn password(mut self, password: impl Into<String>) -> Self {
623 self.password = Some(SecretString::from(password.into()));
624 self
625 }
626
627 pub fn connect_timeout(mut self, seconds: u64) -> Self {
636 self.connect_timeout_secs = seconds.max(1);
637 self
638 }
639
640 pub fn hello_timeout(mut self, seconds: u64) -> Self {
649 self.hello_timeout_secs = seconds.max(1);
650 self
651 }
652
653 pub fn idle_timeout(mut self, seconds: u64) -> Self {
662 self.idle_timeout_secs = seconds.max(1);
663 self
664 }
665
666 pub fn validate(&self) -> Result<()> {
694 validate::hostname(&self.host)?;
696
697 validate::port(self.port)?;
699
700 validate::page_size(self.page_size)?;
702
703 Ok(())
704 }
705
706 pub async fn connect(&self) -> Result<Connection> {
737 self.validate()?;
739
740 let password_ref = self.password.as_ref().map(|s| s.expose_secret());
742
743 match self.transport {
744 Transport::Quic => {
745 Connection::new_quic(
746 &self.host,
747 self.port,
748 self.skip_verify,
749 self.page_size,
750 &self.hello_name,
751 &self.hello_ver,
752 &self.conformance,
753 self.username.as_deref(),
754 password_ref,
755 self.graph.as_deref(),
756 self.tenant_id.as_deref(),
757 self.role.as_deref(),
758 self.connect_timeout_secs,
759 self.hello_timeout_secs,
760 self.idle_timeout_secs,
761 )
762 .await
763 }
764 Transport::Grpc => {
765 #[cfg(feature = "grpc")]
766 {
767 Connection::new_grpc(
768 &self.host,
769 self.port,
770 self.tls_enabled,
771 self.skip_verify,
772 self.page_size,
773 self.username.as_deref(),
774 password_ref,
775 self.graph.as_deref(),
776 )
777 .await
778 }
779 #[cfg(not(feature = "grpc"))]
780 {
781 Err(Error::connection(
782 "gRPC transport requires the 'grpc' feature to be enabled",
783 ))
784 }
785 }
786 }
787 }
788}
789
790#[allow(dead_code)]
792enum ConnectionKind {
793 Quic {
795 conn: quinn::Connection,
796 send: quinn::SendStream,
797 recv: quinn::RecvStream,
798 buffer: Vec<u8>,
800 next_request_id: u64,
802 session_id: String,
804 },
805 #[cfg(feature = "grpc")]
807 Grpc { client: crate::grpc::GrpcClient },
808}
809
810pub struct Connection {
855 kind: ConnectionKind,
856 #[allow(dead_code)]
858 page_size: usize,
859 in_transaction: bool,
861}
862
863impl Connection {
864 #[allow(clippy::too_many_arguments)]
866 async fn new_quic(
867 host: &str,
868 port: u16,
869 skip_verify: bool,
870 page_size: usize,
871 hello_name: &str,
872 hello_ver: &str,
873 conformance: &str,
874 username: Option<&str>,
875 password: Option<&str>,
876 graph: Option<&str>,
877 tenant_id: Option<&str>,
878 role: Option<&str>,
879 connect_timeout_secs: u64,
880 hello_timeout_secs: u64,
881 idle_timeout_secs: u64,
882 ) -> Result<Self> {
883 let mut last_err: Option<Error> = None;
884
885 for attempt in 1..=3 {
886 match Self::connect_quic_once(
887 host,
888 port,
889 skip_verify,
890 page_size,
891 hello_name,
892 hello_ver,
893 conformance,
894 username,
895 password,
896 graph,
897 tenant_id,
898 role,
899 connect_timeout_secs,
900 hello_timeout_secs,
901 idle_timeout_secs,
902 )
903 .await
904 {
905 Ok(conn) => return Ok(conn),
906 Err(e) => {
907 last_err = Some(e);
908 if attempt < 3 {
909 debug!("Connection attempt {} failed, retrying...", attempt);
910 tokio::time::sleep(Duration::from_millis(150)).await;
911 }
912 }
913 }
914 }
915
916 Err(last_err.unwrap_or_else(|| Error::connection("Failed to connect")))
917 }
918
919 #[cfg(feature = "grpc")]
921 #[allow(clippy::too_many_arguments)]
922 async fn new_grpc(
923 host: &str,
924 port: u16,
925 tls_enabled: bool,
926 skip_verify: bool,
927 page_size: usize,
928 username: Option<&str>,
929 password: Option<&str>,
930 graph: Option<&str>,
931 ) -> Result<Self> {
932 use crate::dsn::Dsn;
933
934 let tls_val = if tls_enabled { "1" } else { "0" };
936 let graph_suffix = graph
937 .map(|g| format!("&graph={}", urlencoding::encode(g)))
938 .unwrap_or_default();
939 let dsn_str = if let (Some(user), Some(pass)) = (username, password) {
940 format!(
941 "grpc://{}:{}@{}:{}?tls={}&insecure={}{}",
942 user, pass, host, port, tls_val, skip_verify, graph_suffix
943 )
944 } else {
945 format!(
946 "grpc://{}:{}?tls={}&insecure={}{}",
947 host, port, tls_val, skip_verify, graph_suffix
948 )
949 };
950
951 let dsn = Dsn::parse(&dsn_str)?;
952 let client = crate::grpc::GrpcClient::connect(&dsn).await?;
953
954 Ok(Self {
955 kind: ConnectionKind::Grpc { client },
956 page_size,
957 in_transaction: false,
958 })
959 }
960
961 #[allow(clippy::too_many_arguments)]
962 async fn connect_quic_once(
963 host: &str,
964 port: u16,
965 skip_verify: bool,
966 page_size: usize,
967 hello_name: &str,
968 hello_ver: &str,
969 conformance: &str,
970 username: Option<&str>,
971 password: Option<&str>,
972 graph: Option<&str>,
973 tenant_id: Option<&str>,
974 role: Option<&str>,
975 connect_timeout_secs: u64,
976 hello_timeout_secs: u64,
977 idle_timeout_secs: u64,
978 ) -> Result<Self> {
979 debug!("Creating connection to {}:{}", host, port);
980
981 let _ = rustls::crypto::aws_lc_rs::default_provider().install_default();
983
984 let mut client_crypto = if skip_verify {
986 warn!(
989 "TLS certificate verification DISABLED - connection to {}:{} is vulnerable to MITM attacks. \
990 Do NOT use skip_verify in production!",
991 host, port
992 );
993 rustls::ClientConfig::builder_with_protocol_versions(&[&rustls::version::TLS13])
994 .dangerous()
995 .with_custom_certificate_verifier(Arc::new(SkipServerVerification))
996 .with_no_client_auth()
997 } else {
998 let mut root_store = rustls::RootCertStore::empty();
1000
1001 let cert_result = rustls_native_certs::load_native_certs();
1002
1003 for err in &cert_result.errors {
1005 warn!("Error loading native certificate: {:?}", err);
1006 }
1007
1008 let mut certs_loaded = 0;
1009 let mut certs_failed = 0;
1010
1011 for cert in cert_result.certs {
1012 match root_store.add(cert) {
1013 Ok(()) => certs_loaded += 1,
1014 Err(_) => certs_failed += 1,
1015 }
1016 }
1017
1018 if certs_loaded == 0 {
1019 return Err(Error::tls(
1020 "No system root certificates found. TLS verification cannot proceed. \
1021 Either install system CA certificates or use skip_verify(true) for development only.",
1022 ));
1023 }
1024
1025 debug!(
1026 "Loaded {} system root certificates ({} failed to parse)",
1027 certs_loaded, certs_failed
1028 );
1029
1030 rustls::ClientConfig::builder_with_protocol_versions(&[&rustls::version::TLS13])
1031 .with_root_certificates(root_store)
1032 .with_no_client_auth()
1033 };
1034
1035 client_crypto.alpn_protocols = vec![GEODE_ALPN.to_vec()];
1037
1038 let mut client_config = ClientConfig::new(Arc::new(
1039 quinn::crypto::rustls::QuicClientConfig::try_from(client_crypto)
1040 .map_err(|e| Error::connection(format!("Failed to create QUIC config: {}", e)))?,
1041 ));
1042
1043 let mut transport = quinn::TransportConfig::default();
1045 let idle_timeout = Duration::from_secs(idle_timeout_secs.min(146_000 * 365 * 24 * 3600));
1048 transport.max_idle_timeout(Some(idle_timeout.try_into().map_err(|_| {
1049 Error::connection("Idle timeout value too large for QUIC protocol")
1050 })?));
1051 transport.keep_alive_interval(Some(Duration::from_secs(5)));
1052 client_config.transport_config(Arc::new(transport));
1053
1054 let mut endpoint = Endpoint::client(
1057 "0.0.0.0:0"
1058 .parse()
1059 .expect("0.0.0.0:0 is a valid socket address"),
1060 )
1061 .map_err(|e| Error::connection(format!("Failed to create endpoint: {}", e)))?;
1062 endpoint.set_default_client_config(client_config);
1063
1064 let mut resolved_addrs = format!("{}:{}", host, port)
1066 .to_socket_addrs()
1067 .map_err(|e| {
1068 Error::connection(format!(
1069 "Failed to resolve address {}:{} - {}",
1070 host, port, e
1071 ))
1072 })?;
1073
1074 let server_addr: SocketAddr = resolved_addrs
1075 .find(|addr| matches!(addr, SocketAddr::V4(_) | SocketAddr::V6(_)))
1076 .ok_or_else(|| Error::connection("Invalid address: could not resolve host"))?;
1077
1078 debug!("Connecting to {}", server_addr);
1079
1080 let server_name = if skip_verify {
1083 "localhost" } else {
1085 host
1086 };
1087
1088 trace!("Using server name for SNI: {}", server_name);
1089
1090 let conn = timeout(
1091 Duration::from_secs(connect_timeout_secs),
1092 endpoint
1093 .connect(server_addr, server_name)
1094 .map_err(|e| Error::connection(format!("Connection failed: {}", e)))?,
1095 )
1096 .await
1097 .map_err(|_| Error::connection("Connection timeout"))?
1098 .map_err(|e| Error::connection(format!("Failed to establish connection: {}", e)))?;
1099
1100 debug!("Connection established to {}:{}", host, port);
1101
1102 let (mut send, mut recv) = conn
1104 .open_bi()
1105 .await
1106 .map_err(|e| Error::connection(format!("Failed to open stream: {}", e)))?;
1107
1108 let hello_req = build_hello_request(
1110 username,
1111 password,
1112 hello_name,
1113 hello_ver,
1114 conformance,
1115 graph,
1116 tenant_id,
1117 role,
1118 );
1119 let msg = proto::QuicClientMessage {
1120 msg: Some(proto::quic_client_message::Msg::Hello(hello_req)),
1121 };
1122 let data = proto::encode_with_length_prefix(&msg);
1123
1124 send.write_all(&data)
1125 .await
1126 .map_err(|e| Error::connection(format!("Failed to send HELLO: {}", e)))?;
1127
1128 let mut length_buf = [0u8; 4];
1130 timeout(
1131 Duration::from_secs(hello_timeout_secs),
1132 recv.read_exact(&mut length_buf),
1133 )
1134 .await
1135 .map_err(|_| Error::connection("HELLO response timeout"))?
1136 .map_err(|e| Error::connection(format!("Failed to read HELLO response length: {}", e)))?;
1137
1138 let msg_len = u32::from_be_bytes(length_buf) as usize;
1139
1140 if msg_len > MAX_PROTO_FRAME_BYTES {
1141 return Err(Error::limit(format!(
1142 "HELLO response frame size {} bytes exceeds max {} bytes",
1143 msg_len, MAX_PROTO_FRAME_BYTES
1144 )));
1145 }
1146
1147 let mut msg_buf = vec![0u8; msg_len];
1148 recv.read_exact(&mut msg_buf)
1149 .await
1150 .map_err(|e| Error::connection(format!("Failed to read HELLO response body: {}", e)))?;
1151
1152 let hello_response = proto::decode_quic_server_message(&msg_buf)?;
1153
1154 let session_id = match hello_response.msg {
1155 Some(proto::quic_server_message::Msg::Hello(ref hello_resp)) => {
1156 if !hello_resp.success {
1157 return Err(Error::connection(format!(
1158 "Authentication failed: {}",
1159 hello_resp.error_message
1160 )));
1161 }
1162 hello_resp.session_id.clone()
1163 }
1164 _ => {
1165 return Err(Error::connection("Expected HELLO response"));
1166 }
1167 };
1168
1169 debug!("HELLO handshake complete, session_id={}", session_id);
1170
1171 Ok(Self {
1172 kind: ConnectionKind::Quic {
1173 conn,
1174 send,
1175 recv,
1176 buffer: Vec::new(),
1177 next_request_id: 1,
1178 session_id,
1179 },
1180 page_size,
1181 in_transaction: false,
1182 })
1183 }
1184
1185 async fn send_proto_quic(
1187 send: &mut quinn::SendStream,
1188 msg: &proto::QuicClientMessage,
1189 ) -> Result<()> {
1190 let data = proto::encode_with_length_prefix(msg);
1191 send.write_all(&data)
1192 .await
1193 .map_err(|e| Error::connection(format!("Failed to send message: {}", e)))?;
1194 Ok(())
1195 }
1196
1197 async fn read_proto_quic(
1202 recv: &mut quinn::RecvStream,
1203 timeout_secs: u64,
1204 ) -> Result<proto::QuicServerMessage> {
1205 timeout(Duration::from_secs(timeout_secs), async {
1206 let mut length_buf = [0u8; 4];
1208 recv.read_exact(&mut length_buf)
1209 .await
1210 .map_err(|e| Error::connection(format!("Failed to read response length: {}", e)))?;
1211
1212 let msg_len = u32::from_be_bytes(length_buf) as usize;
1213
1214 if msg_len > MAX_PROTO_FRAME_BYTES {
1216 return Err(Error::limit(format!(
1217 "Frame size {} bytes exceeds max {} bytes",
1218 msg_len, MAX_PROTO_FRAME_BYTES
1219 )));
1220 }
1221
1222 let mut msg_buf = vec![0u8; msg_len];
1223 recv.read_exact(&mut msg_buf)
1224 .await
1225 .map_err(|e| Error::connection(format!("Failed to read response body: {}", e)))?;
1226
1227 proto::decode_quic_server_message(&msg_buf)
1228 })
1229 .await
1230 .map_err(|_| Error::timeout())?
1231 }
1232
1233 async fn try_read_proto_quic(
1238 recv: &mut quinn::RecvStream,
1239 ) -> Result<Option<proto::QuicServerMessage>> {
1240 let read_result = timeout(Duration::from_millis(500), async {
1241 let mut length_buf = [0u8; 4];
1242 recv.read_exact(&mut length_buf)
1243 .await
1244 .map_err(|e| Error::connection(format!("Failed to read response: {}", e)))?;
1245
1246 let msg_len = u32::from_be_bytes(length_buf) as usize;
1247
1248 if msg_len > MAX_PROTO_FRAME_BYTES {
1249 return Err(Error::limit(format!(
1250 "Frame size {} bytes exceeds max {} bytes",
1251 msg_len, MAX_PROTO_FRAME_BYTES
1252 )));
1253 }
1254
1255 let mut msg_buf = vec![0u8; msg_len];
1256 recv.read_exact(&mut msg_buf)
1257 .await
1258 .map_err(|e| Error::connection(format!("Failed to read response body: {}", e)))?;
1259
1260 proto::decode_quic_server_message(&msg_buf)
1261 })
1262 .await;
1263
1264 match read_result {
1265 Ok(Ok(msg)) => Ok(Some(msg)),
1266 Ok(Err(e)) => Err(e),
1267 Err(_) => Ok(None), }
1269 }
1270
1271 fn parse_proto_rows_static(
1273 proto_rows: &[proto::Row],
1274 columns: &[Column],
1275 ) -> Result<Vec<HashMap<String, Value>>> {
1276 let mut rows = Vec::with_capacity(proto_rows.len());
1277 for proto_row in proto_rows {
1278 let mut row = HashMap::with_capacity(columns.len());
1279 for (i, col) in columns.iter().enumerate() {
1280 let value = if i < proto_row.values.len() {
1281 crate::convert::proto_to_value(&proto_row.values[i])
1282 } else {
1283 Value::null()
1284 };
1285 row.insert(col.name.clone(), value);
1286 }
1287 rows.push(row);
1288 }
1289 Ok(rows)
1290 }
1291
1292 async fn send_begin_quic(
1294 send: &mut quinn::SendStream,
1295 recv: &mut quinn::RecvStream,
1296 session_id: &str,
1297 ) -> Result<()> {
1298 let msg = proto::QuicClientMessage {
1299 msg: Some(proto::quic_client_message::Msg::Begin(
1300 proto::BeginRequest {
1301 session_id: session_id.to_string(),
1302 ..Default::default()
1303 },
1304 )),
1305 };
1306 Self::send_proto_quic(send, &msg).await?;
1307
1308 let resp = Self::read_proto_quic(recv, 5).await?;
1309 if !matches!(resp.msg, Some(proto::quic_server_message::Msg::Begin(_))) {
1310 return Err(Error::protocol("Expected BEGIN response"));
1311 }
1312 Ok(())
1313 }
1314
1315 async fn send_commit_quic(
1317 send: &mut quinn::SendStream,
1318 recv: &mut quinn::RecvStream,
1319 session_id: &str,
1320 ) -> Result<()> {
1321 let msg = proto::QuicClientMessage {
1322 msg: Some(proto::quic_client_message::Msg::Commit(
1323 proto::CommitRequest {
1324 session_id: session_id.to_string(),
1325 },
1326 )),
1327 };
1328 Self::send_proto_quic(send, &msg).await?;
1329
1330 let resp = Self::read_proto_quic(recv, 5).await?;
1331 if !matches!(resp.msg, Some(proto::quic_server_message::Msg::Commit(_))) {
1332 return Err(Error::protocol("Expected COMMIT response"));
1333 }
1334 Ok(())
1335 }
1336
1337 async fn send_rollback_quic(
1339 send: &mut quinn::SendStream,
1340 recv: &mut quinn::RecvStream,
1341 session_id: &str,
1342 ) -> Result<()> {
1343 let msg = proto::QuicClientMessage {
1344 msg: Some(proto::quic_client_message::Msg::Rollback(
1345 proto::RollbackRequest {
1346 session_id: session_id.to_string(),
1347 },
1348 )),
1349 };
1350 Self::send_proto_quic(send, &msg).await?;
1351
1352 let resp = Self::read_proto_quic(recv, 5).await?;
1353 if !matches!(resp.msg, Some(proto::quic_server_message::Msg::Rollback(_))) {
1354 return Err(Error::protocol("Expected ROLLBACK response"));
1355 }
1356 Ok(())
1357 }
1358
1359 async fn send_savepoint_quic(
1361 send: &mut quinn::SendStream,
1362 recv: &mut quinn::RecvStream,
1363 session_id: &str,
1364 name: &str,
1365 ) -> Result<()> {
1366 let msg = proto::QuicClientMessage {
1367 msg: Some(proto::quic_client_message::Msg::Savepoint(
1368 proto::SavepointRequest {
1369 name: name.to_string(),
1370 session_id: session_id.to_string(),
1371 },
1372 )),
1373 };
1374 Self::send_proto_quic(send, &msg).await?;
1375
1376 let resp = Self::read_proto_quic(recv, 5).await?;
1377 if !matches!(
1378 resp.msg,
1379 Some(proto::quic_server_message::Msg::Savepoint(_))
1380 ) {
1381 return Err(Error::protocol("Expected SAVEPOINT response"));
1382 }
1383 Ok(())
1384 }
1385
1386 async fn send_rollback_to_quic(
1388 send: &mut quinn::SendStream,
1389 recv: &mut quinn::RecvStream,
1390 session_id: &str,
1391 name: &str,
1392 ) -> Result<()> {
1393 let msg = proto::QuicClientMessage {
1394 msg: Some(proto::quic_client_message::Msg::RollbackTo(
1395 proto::RollbackToRequest {
1396 name: name.to_string(),
1397 session_id: session_id.to_string(),
1398 },
1399 )),
1400 };
1401 Self::send_proto_quic(send, &msg).await?;
1402
1403 let resp = Self::read_proto_quic(recv, 5).await?;
1404 if !matches!(
1405 resp.msg,
1406 Some(proto::quic_server_message::Msg::RollbackTo(_))
1407 ) {
1408 return Err(Error::protocol("Expected ROLLBACK_TO response"));
1409 }
1410 Ok(())
1411 }
1412
1413 pub async fn query(&mut self, gql: &str) -> Result<(Page, Option<String>)> {
1443 self.query_with_params(gql, &HashMap::new()).await
1444 }
1445
1446 pub async fn query_with_params(
1486 &mut self,
1487 gql: &str,
1488 params: &HashMap<String, Value>,
1489 ) -> Result<(Page, Option<String>)> {
1490 validate::query(gql)?;
1491 for key in params.keys() {
1492 validate::param_name(key)?;
1493 }
1494 match &mut self.kind {
1495 ConnectionKind::Quic {
1496 send,
1497 recv,
1498 session_id,
1499 ..
1500 } => Self::query_with_params_quic(send, recv, gql, params, session_id).await,
1501 #[cfg(feature = "grpc")]
1502 ConnectionKind::Grpc { client } => client.query_with_params(gql, params).await,
1503 }
1504 }
1505
1506 async fn query_with_params_quic(
1508 send: &mut quinn::SendStream,
1509 recv: &mut quinn::RecvStream,
1510 gql: &str,
1511 params: &HashMap<String, Value>,
1512 session_id: &str,
1513 ) -> Result<(Page, Option<String>)> {
1514 let (page, cursor) =
1515 Self::query_with_params_quic_inner(send, recv, gql, params, session_id).await?;
1516
1517 if !page.final_page {
1519 let mut all_rows = page.rows;
1520 let columns = page.columns;
1521 let mut ordered = page.ordered;
1522 let mut order_keys = page.order_keys;
1523 let mut request_id: u64 = 0;
1524 let mut page_count: usize = 1; loop {
1527 if all_rows.len() > DEFAULT_MAX_ROWS {
1529 return Err(Error::limit(format!(
1530 "Total rows {} exceeds max {}",
1531 all_rows.len(),
1532 DEFAULT_MAX_ROWS
1533 )));
1534 }
1535 page_count += 1;
1536 if page_count > DEFAULT_MAX_PAGES {
1537 return Err(Error::limit(format!(
1538 "Page count {} exceeds max {}",
1539 page_count, DEFAULT_MAX_PAGES
1540 )));
1541 }
1542
1543 request_id += 1;
1544 let pull_req = proto::QuicClientMessage {
1545 msg: Some(proto::quic_client_message::Msg::Pull(proto::PullRequest {
1546 request_id,
1547 page_size: 1000,
1548 session_id: String::new(),
1549 })),
1550 };
1551 Self::send_proto_quic(send, &pull_req).await?;
1552
1553 let resp = Self::read_proto_quic(recv, 30).await?;
1554
1555 let exec_resp = match &resp.msg {
1557 Some(proto::quic_server_message::Msg::Pull(pull)) => pull.response.as_ref(),
1558 Some(proto::quic_server_message::Msg::Execute(e)) => Some(e),
1559 _ => None,
1560 };
1561
1562 let exec_resp = match exec_resp {
1563 Some(e) => e,
1564 None => break,
1565 };
1566
1567 if let Some(proto::execution_response::Payload::Error(ref err)) = exec_resp.payload
1568 {
1569 return Err(Error::Query {
1570 code: err.code.clone(),
1571 message: err.message.clone(),
1572 });
1573 }
1574
1575 if let Some(proto::execution_response::Payload::Page(ref page_data)) =
1576 exec_resp.payload
1577 {
1578 let rows = Self::parse_proto_rows_static(&page_data.rows, &columns)?;
1579 all_rows.extend(rows);
1580 ordered = page_data.ordered;
1581 order_keys = page_data.order_keys.clone();
1582 if page_data.r#final {
1583 break;
1584 }
1585 } else {
1586 break;
1587 }
1588 }
1589
1590 let final_page = Page {
1591 columns,
1592 rows: all_rows,
1593 ordered,
1594 order_keys,
1595 final_page: true,
1596 };
1597 return Ok((final_page, cursor));
1598 }
1599
1600 Ok((page, cursor))
1601 }
1602
1603 async fn query_with_params_quic_inner(
1605 send: &mut quinn::SendStream,
1606 recv: &mut quinn::RecvStream,
1607 gql: &str,
1608 params: &HashMap<String, Value>,
1609 session_id: &str,
1610 ) -> Result<(Page, Option<String>)> {
1611 let params_proto: Vec<proto::Param> = params
1613 .iter()
1614 .map(|(k, v)| proto::Param {
1615 name: k.clone(),
1616 value: Some(v.to_proto_value()),
1617 })
1618 .collect();
1619
1620 let exec_req = proto::ExecuteRequest {
1622 session_id: session_id.to_string(),
1623 query: gql.to_string(),
1624 params: params_proto,
1625 };
1626 let msg = proto::QuicClientMessage {
1627 msg: Some(proto::quic_client_message::Msg::Execute(exec_req)),
1628 };
1629 Self::send_proto_quic(send, &msg)
1630 .await
1631 .map_err(|e| Error::query(format!("{}", e)))?;
1632
1633 let resp = Self::read_proto_quic(recv, 10).await?;
1635
1636 let exec_resp = match resp.msg {
1637 Some(proto::quic_server_message::Msg::Execute(e)) => e,
1638 _ => return Err(Error::protocol("Expected Execute response")),
1639 };
1640
1641 if let Some(proto::execution_response::Payload::Error(ref err)) = exec_resp.payload {
1643 let _ = Self::try_read_proto_quic(recv).await;
1645 return Err(Error::Query {
1646 code: err.code.clone(),
1647 message: err.message.clone(),
1648 });
1649 }
1650
1651 let columns: Vec<Column> = match exec_resp.payload {
1653 Some(proto::execution_response::Payload::Schema(ref s)) => s
1654 .columns
1655 .iter()
1656 .map(|c| Column {
1657 name: c.name.clone(),
1658 col_type: c.r#type.clone(),
1659 })
1660 .collect(),
1661 _ => Vec::new(),
1662 };
1663
1664 trace!("Schema columns: {:?}", columns);
1665
1666 if let Some(inline_resp) = Self::try_read_proto_quic(recv).await? {
1668 if let Some(proto::quic_server_message::Msg::Execute(inline_exec)) = inline_resp.msg {
1669 if let Some(proto::execution_response::Payload::Error(ref err)) =
1670 inline_exec.payload
1671 {
1672 return Err(Error::Query {
1673 code: err.code.clone(),
1674 message: err.message.clone(),
1675 });
1676 }
1677
1678 if let Some(proto::execution_response::Payload::Page(ref page_data)) =
1679 inline_exec.payload
1680 {
1681 let rows = Self::parse_proto_rows_static(&page_data.rows, &columns)?;
1682 let page = Page {
1683 columns,
1684 rows,
1685 ordered: page_data.ordered,
1686 order_keys: page_data.order_keys.clone(),
1687 final_page: page_data.r#final,
1688 };
1689 return Ok((page, None));
1690 }
1691
1692 let page = Page {
1694 columns,
1695 rows: Vec::new(),
1696 ordered: false,
1697 order_keys: Vec::new(),
1698 final_page: true,
1699 };
1700 return Ok((page, None));
1701 }
1702 }
1703
1704 if let Some(proto::execution_response::Payload::Page(ref page_data)) = exec_resp.payload {
1706 let rows = Self::parse_proto_rows_static(&page_data.rows, &columns)?;
1707 let page = Page {
1708 columns,
1709 rows,
1710 ordered: page_data.ordered,
1711 order_keys: page_data.order_keys.clone(),
1712 final_page: page_data.r#final,
1713 };
1714 return Ok((page, None));
1715 }
1716
1717 let resp = Self::read_proto_quic(recv, 30).await?;
1719 if let Some(proto::quic_server_message::Msg::Execute(exec_resp)) = resp.msg {
1720 if let Some(proto::execution_response::Payload::Error(ref err)) = exec_resp.payload {
1721 return Err(Error::Query {
1722 code: err.code.clone(),
1723 message: err.message.clone(),
1724 });
1725 }
1726
1727 if let Some(proto::execution_response::Payload::Page(ref page_data)) = exec_resp.payload
1728 {
1729 let rows = Self::parse_proto_rows_static(&page_data.rows, &columns)?;
1730 let page = Page {
1731 columns,
1732 rows,
1733 ordered: page_data.ordered,
1734 order_keys: page_data.order_keys.clone(),
1735 final_page: page_data.r#final,
1736 };
1737 return Ok((page, None));
1738 }
1739 }
1740
1741 let page = Page {
1743 columns,
1744 rows: Vec::new(),
1745 ordered: false,
1746 order_keys: Vec::new(),
1747 final_page: true,
1748 };
1749
1750 Ok((page, None))
1751 }
1752
1753 pub fn query_sync(
1755 &mut self,
1756 gql: &str,
1757 params: Option<HashMap<String, serde_json::Value>>,
1758 ) -> Result<Page> {
1759 let params_map = params.unwrap_or_default();
1760 let mut params_typed: HashMap<String, Value> = HashMap::new();
1761 for (k, v) in params_map {
1762 let typed_val = crate::types::Value::from_json(v)?;
1763 params_typed.insert(k, typed_val);
1764 }
1765
1766 match tokio::runtime::Handle::try_current() {
1767 Ok(handle) => {
1768 let (page, _cursor) =
1769 handle.block_on(self.query_with_params(gql, ¶ms_typed))?;
1770 Ok(page)
1771 }
1772 Err(_) => {
1773 let rt = tokio::runtime::Runtime::new()
1774 .map_err(|e| Error::query(format!("Failed to create runtime: {}", e)))?;
1775 let (page, _cursor) = rt.block_on(self.query_with_params(gql, ¶ms_typed))?;
1776 Ok(page)
1777 }
1778 }
1779 }
1780
1781 pub async fn begin(&mut self) -> Result<()> {
1806 let result = match &mut self.kind {
1807 ConnectionKind::Quic {
1808 send,
1809 recv,
1810 session_id,
1811 ..
1812 } => Self::send_begin_quic(send, recv, session_id).await,
1813 #[cfg(feature = "grpc")]
1814 ConnectionKind::Grpc { client } => client.begin().await,
1815 };
1816 if result.is_ok() {
1817 self.in_transaction = true;
1818 }
1819 result
1820 }
1821
1822 pub async fn commit(&mut self) -> Result<()> {
1845 let result = match &mut self.kind {
1846 ConnectionKind::Quic {
1847 send,
1848 recv,
1849 session_id,
1850 ..
1851 } => Self::send_commit_quic(send, recv, session_id).await,
1852 #[cfg(feature = "grpc")]
1853 ConnectionKind::Grpc { client } => client.commit().await,
1854 };
1855 if result.is_ok() {
1856 self.in_transaction = false;
1857 }
1858 result
1859 }
1860
1861 pub async fn rollback(&mut self) -> Result<()> {
1885 let result = match &mut self.kind {
1886 ConnectionKind::Quic {
1887 send,
1888 recv,
1889 session_id,
1890 ..
1891 } => Self::send_rollback_quic(send, recv, session_id).await,
1892 #[cfg(feature = "grpc")]
1893 ConnectionKind::Grpc { client } => client.rollback().await,
1894 };
1895 if result.is_ok() {
1896 self.in_transaction = false;
1897 }
1898 result
1899 }
1900
1901 pub async fn savepoint(&mut self, name: &str) -> Result<Savepoint> {
1935 match &mut self.kind {
1936 ConnectionKind::Quic {
1937 send,
1938 recv,
1939 session_id,
1940 ..
1941 } => Self::send_savepoint_quic(send, recv, session_id, name).await?,
1942 #[cfg(feature = "grpc")]
1943 ConnectionKind::Grpc { client } => client.savepoint(name).await?,
1944 }
1945 Ok(Savepoint {
1946 name: name.to_string(),
1947 })
1948 }
1949
1950 pub async fn rollback_to(&mut self, savepoint: &Savepoint) -> Result<()> {
1979 match &mut self.kind {
1980 ConnectionKind::Quic {
1981 send,
1982 recv,
1983 session_id,
1984 ..
1985 } => Self::send_rollback_to_quic(send, recv, session_id, &savepoint.name).await?,
1986 #[cfg(feature = "grpc")]
1987 ConnectionKind::Grpc { client } => client.rollback_to(&savepoint.name).await?,
1988 }
1989 Ok(())
1990 }
1991
1992 pub fn prepare(&self, query: &str) -> Result<PreparedStatement> {
2026 Ok(PreparedStatement::new(query))
2027 }
2028
2029 pub async fn explain(&mut self, gql: &str) -> Result<QueryPlan> {
2062 let explain_query = format!("EXPLAIN {}", gql);
2064 let (_page, _) = self.query(&explain_query).await?;
2065
2066 Ok(QueryPlan {
2069 operations: Vec::new(),
2070 estimated_rows: 0,
2071 raw: serde_json::json!({}),
2072 })
2073 }
2074
2075 pub async fn profile(&mut self, gql: &str) -> Result<QueryProfile> {
2106 let profile_query = format!("PROFILE {}", gql);
2108 let (page, _) = self.query(&profile_query).await?;
2109
2110 let plan = QueryPlan {
2112 operations: Vec::new(),
2113 estimated_rows: 0,
2114 raw: serde_json::json!({}),
2115 };
2116
2117 Ok(QueryProfile {
2118 plan,
2119 actual_rows: page.rows.len() as u64,
2120 execution_time_ms: 0.0,
2121 raw: serde_json::json!({}),
2122 })
2123 }
2124
2125 pub async fn batch(
2163 &mut self,
2164 queries: &[(&str, Option<&HashMap<String, Value>>)],
2165 ) -> Result<Vec<Page>> {
2166 let mut results = Vec::with_capacity(queries.len());
2167
2168 for (query, params) in queries {
2169 let (page, _) = match params {
2170 Some(p) => self.query_with_params(query, p).await?,
2171 None => self.query(query).await?,
2172 };
2173 results.push(page);
2174 }
2175
2176 Ok(results)
2177 }
2178
2179 #[allow(dead_code)]
2182 fn parse_plan_operations(result: &serde_json::Value) -> Vec<PlanOperation> {
2183 let mut operations = Vec::new();
2184
2185 if let Some(ops) = result.get("operations").and_then(|o| o.as_array()) {
2186 for op in ops {
2187 operations.push(Self::parse_single_operation(op));
2188 }
2189 } else if let Some(plan) = result.get("plan") {
2190 operations.push(Self::parse_single_operation(plan));
2192 }
2193
2194 operations
2195 }
2196
2197 #[allow(dead_code)]
2199 fn parse_single_operation(op: &serde_json::Value) -> PlanOperation {
2200 let op_type = op
2201 .get("type")
2202 .or_else(|| op.get("op_type"))
2203 .and_then(|t| t.as_str())
2204 .unwrap_or("Unknown")
2205 .to_string();
2206
2207 let description = op
2208 .get("description")
2209 .or_else(|| op.get("desc"))
2210 .and_then(|d| d.as_str())
2211 .unwrap_or("")
2212 .to_string();
2213
2214 let estimated_rows = op
2215 .get("estimated_rows")
2216 .or_else(|| op.get("rows"))
2217 .and_then(|r| r.as_u64());
2218
2219 let children = op
2220 .get("children")
2221 .and_then(|c| c.as_array())
2222 .map(|arr| arr.iter().map(Self::parse_single_operation).collect())
2223 .unwrap_or_default();
2224
2225 PlanOperation {
2226 op_type,
2227 description,
2228 estimated_rows,
2229 children,
2230 }
2231 }
2232
2233 pub fn close(&mut self) -> Result<()> {
2256 match &mut self.kind {
2257 ConnectionKind::Quic { conn, .. } => {
2258 conn.close(0u32.into(), b"client closing");
2262 Ok(())
2263 }
2264 #[cfg(feature = "grpc")]
2265 ConnectionKind::Grpc { client } => client.close(),
2266 }
2267 }
2268
2269 pub fn in_transaction(&self) -> bool {
2289 self.in_transaction
2290 }
2291
2292 pub fn is_healthy(&self) -> bool {
2293 match &self.kind {
2294 ConnectionKind::Quic { conn, .. } => {
2295 conn.close_reason().is_none()
2297 }
2298 #[cfg(feature = "grpc")]
2299 ConnectionKind::Grpc { .. } => {
2300 true
2302 }
2303 }
2304 }
2305}
2306
2307#[derive(Debug)]
2309struct SkipServerVerification;
2310
2311impl rustls::client::danger::ServerCertVerifier for SkipServerVerification {
2312 fn verify_server_cert(
2313 &self,
2314 _end_entity: &CertificateDer,
2315 _intermediates: &[CertificateDer],
2316 _server_name: &RustlsServerName,
2317 _ocsp_response: &[u8],
2318 _now: rustls::pki_types::UnixTime,
2319 ) -> std::result::Result<rustls::client::danger::ServerCertVerified, rustls::Error> {
2320 Ok(rustls::client::danger::ServerCertVerified::assertion())
2321 }
2322
2323 fn verify_tls12_signature(
2324 &self,
2325 _message: &[u8],
2326 _cert: &CertificateDer,
2327 _dss: &rustls::DigitallySignedStruct,
2328 ) -> std::result::Result<rustls::client::danger::HandshakeSignatureValid, rustls::Error> {
2329 Ok(rustls::client::danger::HandshakeSignatureValid::assertion())
2330 }
2331
2332 fn verify_tls13_signature(
2333 &self,
2334 _message: &[u8],
2335 _cert: &CertificateDer,
2336 _dss: &rustls::DigitallySignedStruct,
2337 ) -> std::result::Result<rustls::client::danger::HandshakeSignatureValid, rustls::Error> {
2338 Ok(rustls::client::danger::HandshakeSignatureValid::assertion())
2339 }
2340
2341 fn supported_verify_schemes(&self) -> Vec<rustls::SignatureScheme> {
2342 vec![
2343 rustls::SignatureScheme::RSA_PKCS1_SHA256,
2344 rustls::SignatureScheme::ECDSA_NISTP256_SHA256,
2345 rustls::SignatureScheme::ED25519,
2346 ]
2347 }
2348}
2349
2350#[cfg(test)]
2351mod tests {
2352 use super::*;
2353
2354 #[test]
2357 fn test_prepared_statement_new() {
2358 let stmt = PreparedStatement::new("MATCH (n:Person {id: $id}) RETURN n");
2359 assert_eq!(stmt.query(), "MATCH (n:Person {id: $id}) RETURN n");
2360 assert_eq!(stmt.param_names(), &["id"]);
2361 }
2362
2363 #[test]
2364 fn test_prepared_statement_multiple_params() {
2365 let stmt = PreparedStatement::new(
2366 "MATCH (p:Person {name: $name}) WHERE p.age > $min_age AND p.city = $city RETURN p",
2367 );
2368 assert!(stmt.query().contains("$name"));
2369 let names = stmt.param_names();
2370 assert_eq!(names.len(), 3);
2371 assert!(names.contains(&"name".to_string()));
2372 assert!(names.contains(&"min_age".to_string()));
2373 assert!(names.contains(&"city".to_string()));
2374 }
2375
2376 #[test]
2377 fn test_prepared_statement_no_params() {
2378 let stmt = PreparedStatement::new("MATCH (n) RETURN n LIMIT 10");
2379 assert!(stmt.param_names().is_empty());
2380 }
2381
2382 #[test]
2383 fn test_prepared_statement_duplicate_params() {
2384 let stmt =
2385 PreparedStatement::new("MATCH (a {id: $id})-[:KNOWS]->(b {id: $id}) RETURN a, b");
2386 assert_eq!(stmt.param_names(), &["id"]);
2388 }
2389
2390 #[test]
2391 fn test_prepared_statement_underscore_params() {
2392 let stmt = PreparedStatement::new("MATCH (n {user_id: $user_id}) RETURN n");
2393 assert_eq!(stmt.param_names(), &["user_id"]);
2394 }
2395
2396 #[test]
2397 fn test_prepared_statement_numeric_params() {
2398 let stmt = PreparedStatement::new("RETURN $param1, $param2, $param123");
2399 let names = stmt.param_names();
2400 assert_eq!(names.len(), 3);
2401 assert!(names.contains(&"param1".to_string()));
2402 assert!(names.contains(&"param2".to_string()));
2403 assert!(names.contains(&"param123".to_string()));
2404 }
2405
2406 #[test]
2409 fn test_plan_operation_struct() {
2410 let op = PlanOperation {
2411 op_type: "NodeScan".to_string(),
2412 description: "Scan Person nodes".to_string(),
2413 estimated_rows: Some(100),
2414 children: vec![],
2415 };
2416 assert_eq!(op.op_type, "NodeScan");
2417 assert_eq!(op.description, "Scan Person nodes");
2418 assert_eq!(op.estimated_rows, Some(100));
2419 assert!(op.children.is_empty());
2420 }
2421
2422 #[test]
2423 fn test_plan_operation_with_children() {
2424 let child = PlanOperation {
2425 op_type: "Filter".to_string(),
2426 description: "Filter by age".to_string(),
2427 estimated_rows: Some(50),
2428 children: vec![],
2429 };
2430 let parent = PlanOperation {
2431 op_type: "Projection".to_string(),
2432 description: "Project name, age".to_string(),
2433 estimated_rows: Some(50),
2434 children: vec![child],
2435 };
2436 assert_eq!(parent.children.len(), 1);
2437 assert_eq!(parent.children[0].op_type, "Filter");
2438 }
2439
2440 #[test]
2443 fn test_query_plan_struct() {
2444 let plan = QueryPlan {
2445 operations: vec![PlanOperation {
2446 op_type: "NodeScan".to_string(),
2447 description: "Full scan".to_string(),
2448 estimated_rows: Some(1000),
2449 children: vec![],
2450 }],
2451 estimated_rows: 1000,
2452 raw: serde_json::json!({"type": "plan"}),
2453 };
2454 assert_eq!(plan.operations.len(), 1);
2455 assert_eq!(plan.estimated_rows, 1000);
2456 }
2457
2458 #[test]
2461 fn test_query_profile_struct() {
2462 let plan = QueryPlan {
2463 operations: vec![],
2464 estimated_rows: 100,
2465 raw: serde_json::json!({}),
2466 };
2467 let profile = QueryProfile {
2468 plan,
2469 actual_rows: 95,
2470 execution_time_ms: 12.5,
2471 raw: serde_json::json!({"type": "profile"}),
2472 };
2473 assert_eq!(profile.actual_rows, 95);
2474 assert!((profile.execution_time_ms - 12.5).abs() < 0.001);
2475 }
2476
2477 #[test]
2480 fn test_page_struct() {
2481 let page = Page {
2482 columns: vec![Column {
2483 name: "x".to_string(),
2484 col_type: "INT".to_string(),
2485 }],
2486 rows: vec![],
2487 ordered: false,
2488 order_keys: vec![],
2489 final_page: true,
2490 };
2491 assert_eq!(page.columns.len(), 1);
2492 assert!(page.rows.is_empty());
2493 assert!(page.final_page);
2494 }
2495
2496 #[test]
2499 fn test_column_struct() {
2500 let col = Column {
2501 name: "age".to_string(),
2502 col_type: "INT".to_string(),
2503 };
2504 assert_eq!(col.name, "age");
2505 assert_eq!(col.col_type, "INT");
2506 }
2507
2508 #[test]
2511 fn test_savepoint_struct() {
2512 let sp = Savepoint {
2513 name: "before_update".to_string(),
2514 };
2515 assert_eq!(sp.name, "before_update");
2516 }
2517
2518 #[test]
2521 fn test_client_builder_defaults() {
2522 let _client = Client::new("localhost", 3141);
2523 }
2525
2526 #[test]
2527 fn test_client_builder_chain() {
2528 let _client = Client::new("example.com", 8443)
2529 .skip_verify(true)
2530 .page_size(500)
2531 .client_name("test-app")
2532 .client_version("2.0.0")
2533 .conformance("full");
2534 }
2536
2537 #[test]
2538 fn test_client_clone() {
2539 let client = Client::new("localhost", 3141).skip_verify(true);
2540 let _cloned = client.clone();
2541 }
2543
2544 #[test]
2547 fn test_parse_plan_operations_empty() {
2548 let result = serde_json::json!({});
2549 let ops = Connection::parse_plan_operations(&result);
2550 assert!(ops.is_empty());
2551 }
2552
2553 #[test]
2554 fn test_parse_plan_operations_array() {
2555 let result = serde_json::json!({
2556 "operations": [
2557 {"type": "NodeScan", "description": "Scan nodes", "estimated_rows": 100},
2558 {"type": "Filter", "description": "Apply filter", "estimated_rows": 50}
2559 ]
2560 });
2561 let ops = Connection::parse_plan_operations(&result);
2562 assert_eq!(ops.len(), 2);
2563 assert_eq!(ops[0].op_type, "NodeScan");
2564 assert_eq!(ops[1].op_type, "Filter");
2565 }
2566
2567 #[test]
2568 fn test_parse_plan_operations_single_plan() {
2569 let result = serde_json::json!({
2570 "plan": {"op_type": "FullScan", "desc": "Full table scan"}
2571 });
2572 let ops = Connection::parse_plan_operations(&result);
2573 assert_eq!(ops.len(), 1);
2574 assert_eq!(ops[0].op_type, "FullScan");
2575 assert_eq!(ops[0].description, "Full table scan");
2576 }
2577
2578 #[test]
2579 fn test_parse_single_operation() {
2580 let op_json = serde_json::json!({
2581 "type": "IndexScan",
2582 "description": "Use index on Person(name)",
2583 "estimated_rows": 25,
2584 "children": [
2585 {"type": "Filter", "description": "Filter results"}
2586 ]
2587 });
2588 let op = Connection::parse_single_operation(&op_json);
2589 assert_eq!(op.op_type, "IndexScan");
2590 assert_eq!(op.description, "Use index on Person(name)");
2591 assert_eq!(op.estimated_rows, Some(25));
2592 assert_eq!(op.children.len(), 1);
2593 assert_eq!(op.children[0].op_type, "Filter");
2594 }
2595
2596 #[test]
2597 fn test_parse_single_operation_minimal() {
2598 let op_json = serde_json::json!({});
2599 let op = Connection::parse_single_operation(&op_json);
2600 assert_eq!(op.op_type, "Unknown");
2601 assert_eq!(op.description, "");
2602 assert_eq!(op.estimated_rows, None);
2603 assert!(op.children.is_empty());
2604 }
2605
2606 #[test]
2607 fn test_parse_single_operation_alt_fields() {
2608 let op_json = serde_json::json!({
2609 "op_type": "Sort",
2610 "desc": "Sort by name ASC",
2611 "rows": 100
2612 });
2613 let op = Connection::parse_single_operation(&op_json);
2614 assert_eq!(op.op_type, "Sort");
2615 assert_eq!(op.description, "Sort by name ASC");
2616 assert_eq!(op.estimated_rows, Some(100));
2617 }
2618
2619 #[test]
2622 fn test_redact_dsn_url_with_password() {
2623 let dsn = "quic://admin:secret123@localhost:3141";
2624 let redacted = redact_dsn(dsn);
2625 assert!(redacted.contains("[REDACTED]"));
2626 assert!(!redacted.contains("secret123"));
2627 assert!(redacted.contains("admin"));
2628 assert!(redacted.contains("localhost"));
2629 }
2630
2631 #[test]
2632 fn test_redact_dsn_url_without_password() {
2633 let dsn = "quic://admin@localhost:3141";
2634 let redacted = redact_dsn(dsn);
2635 assert!(!redacted.contains("[REDACTED]"));
2636 assert!(redacted.contains("admin"));
2637 assert!(redacted.contains("localhost"));
2638 }
2639
2640 #[test]
2641 fn test_redact_dsn_url_no_auth() {
2642 let dsn = "quic://localhost:3141";
2643 let redacted = redact_dsn(dsn);
2644 assert_eq!(redacted, dsn);
2645 }
2646
2647 #[test]
2648 fn test_redact_dsn_query_param_password() {
2649 let dsn = "localhost:3141?username=admin&password=secret123";
2650 let redacted = redact_dsn(dsn);
2651 assert!(redacted.contains("[REDACTED]"));
2652 assert!(!redacted.contains("secret123"));
2653 assert!(redacted.contains("username=admin"));
2654 }
2655
2656 #[test]
2657 fn test_redact_dsn_query_param_pass() {
2658 let dsn = "localhost:3141?user=admin&pass=mysecret";
2659 let redacted = redact_dsn(dsn);
2660 assert!(redacted.contains("[REDACTED]"));
2661 assert!(!redacted.contains("mysecret"));
2662 }
2663
2664 #[test]
2665 fn test_redact_dsn_simple_no_password() {
2666 let dsn = "localhost:3141?insecure=true";
2667 let redacted = redact_dsn(dsn);
2668 assert_eq!(redacted, dsn);
2669 }
2670
2671 #[test]
2672 fn test_redact_dsn_url_with_query_and_password() {
2673 let dsn = "quic://user:pass@localhost:3141?insecure=true";
2674 let redacted = redact_dsn(dsn);
2675 assert!(redacted.contains("[REDACTED]"));
2676 assert!(!redacted.contains(":pass@"));
2677 assert!(redacted.contains("insecure=true"));
2678 }
2679
2680 #[test]
2683 fn test_client_validate_valid() {
2684 let client = Client::new("localhost", 3141);
2685 assert!(client.validate().is_ok());
2686 }
2687
2688 #[test]
2689 fn test_client_validate_valid_hostname() {
2690 let client = Client::new("geode.example.com", 3141);
2691 assert!(client.validate().is_ok());
2692 }
2693
2694 #[test]
2695 fn test_client_validate_valid_ipv4() {
2696 let client = Client::new("192.168.1.1", 8443);
2697 assert!(client.validate().is_ok());
2698 }
2699
2700 #[test]
2701 fn test_client_validate_invalid_hostname_hyphen_start() {
2702 let client = Client::new("-invalid", 3141);
2703 assert!(client.validate().is_err());
2704 }
2705
2706 #[test]
2707 fn test_client_validate_invalid_hostname_hyphen_end() {
2708 let client = Client::new("invalid-", 3141);
2709 assert!(client.validate().is_err());
2710 }
2711
2712 #[test]
2713 fn test_client_validate_invalid_port_zero() {
2714 let client = Client::new("localhost", 0);
2715 assert!(client.validate().is_err());
2716 }
2717
2718 #[test]
2719 fn test_client_validate_invalid_page_size_zero() {
2720 let client = Client::new("localhost", 3141).page_size(0);
2721 assert!(client.validate().is_err());
2722 }
2723
2724 #[test]
2725 fn test_client_validate_invalid_page_size_too_large() {
2726 let client = Client::new("localhost", 3141).page_size(200_000);
2727 assert!(client.validate().is_err());
2728 }
2729
2730 #[test]
2731 fn test_client_validate_with_all_options() {
2732 let client = Client::new("geode.example.com", 8443)
2733 .skip_verify(true)
2734 .page_size(500)
2735 .username("admin")
2736 .password("secret")
2737 .connect_timeout(15)
2738 .hello_timeout(10)
2739 .idle_timeout(60);
2740 assert!(client.validate().is_ok());
2741 }
2742
2743 #[test]
2744 fn test_build_hello_request_includes_tenant_and_role() {
2745 use prost::Message;
2746 let req = build_hello_request(
2747 Some("admin"),
2748 Some("secret"),
2749 "geode-rust",
2750 "0.1",
2751 "min",
2752 None,
2753 Some("acme"),
2754 Some("analyst"),
2755 );
2756 assert_eq!(req.tenant_id, Some("acme".to_string()));
2757 assert_eq!(req.role, Some("analyst".to_string()));
2758 let encoded = req.encode_to_vec();
2759 let decoded = proto::HelloRequest::decode(encoded.as_slice()).unwrap();
2760 assert_eq!(decoded.tenant_id.as_deref(), Some("acme"));
2761 assert_eq!(decoded.role.as_deref(), Some("analyst"));
2762 }
2763
2764 #[test]
2765 fn test_build_hello_request_omits_empty_optionals() {
2766 let req = build_hello_request(
2767 Some("u"),
2768 Some("p"),
2769 "n",
2770 "v",
2771 "min",
2772 Some(""),
2773 Some(""),
2774 Some(""),
2775 );
2776 assert_eq!(req.graph, None);
2777 assert_eq!(req.tenant_id, None);
2778 assert_eq!(req.role, None);
2779 }
2780
2781 #[test]
2782 fn test_client_tenant_role_builders() {
2783 let _client = Client::new("localhost", 3141)
2784 .tenant_id("acme")
2785 .role("analyst");
2786 }
2787
2788 #[test]
2790 fn test_client_extreme_timeout_values() {
2791 let _client = Client::new("localhost", 3141)
2793 .connect_timeout(u64::MAX)
2794 .hello_timeout(u64::MAX)
2795 .idle_timeout(u64::MAX);
2796 }
2798
2799 #[test]
2800 fn test_convert_edge_uses_type_field() {
2801 let edge = proto::EdgeValue {
2802 id: 100,
2803 from_id: 1,
2804 to_id: 2,
2805 label: "KNOWS".to_string(),
2806 properties: vec![],
2807 };
2808 let val = crate::convert::proto_to_value(&proto::Value {
2809 kind: Some(proto::value::Kind::EdgeVal(edge)),
2810 });
2811 let e = val.as_edge().unwrap();
2812 assert_eq!(e.edge_type, "KNOWS");
2813 }
2814
2815 #[test]
2816 fn test_convert_edge_uses_start_end_node() {
2817 let edge = proto::EdgeValue {
2818 id: 100,
2819 from_id: 42,
2820 to_id: 99,
2821 label: "LIKES".to_string(),
2822 properties: vec![],
2823 };
2824 let val = crate::convert::proto_to_value(&proto::Value {
2825 kind: Some(proto::value::Kind::EdgeVal(edge)),
2826 });
2827 let e = val.as_edge().unwrap();
2828 assert_eq!(e.start_node, 42);
2829 assert_eq!(e.end_node, 99);
2830 }
2831
2832 #[test]
2833 fn test_convert_edge_with_properties() {
2834 let edge = proto::EdgeValue {
2835 id: 100,
2836 from_id: 1,
2837 to_id: 2,
2838 label: "KNOWS".to_string(),
2839 properties: vec![proto::MapEntry {
2840 key: "since".to_string(),
2841 value: Some(proto::Value {
2842 kind: Some(proto::value::Kind::IntVal(proto::IntValue {
2843 value: 2020,
2844 kind: 1,
2845 })),
2846 }),
2847 }],
2848 };
2849 let val = crate::convert::proto_to_value(&proto::Value {
2850 kind: Some(proto::value::Kind::EdgeVal(edge)),
2851 });
2852 let e = val.as_edge().unwrap();
2853 assert_eq!(e.properties.get("since").unwrap().as_int().unwrap(), 2020);
2854 }
2855
2856 #[test]
2857 fn test_convert_node_fields() {
2858 let node = proto::NodeValue {
2859 id: 42,
2860 labels: vec!["Person".to_string()],
2861 properties: vec![proto::MapEntry {
2862 key: "name".to_string(),
2863 value: Some(proto::Value {
2864 kind: Some(proto::value::Kind::StringVal(proto::StringValue {
2865 value: "Alice".to_string(),
2866 kind: 1,
2867 })),
2868 }),
2869 }],
2870 };
2871 let val = crate::convert::proto_to_value(&proto::Value {
2872 kind: Some(proto::value::Kind::NodeVal(node)),
2873 });
2874 let n = val.as_node().unwrap();
2875 assert_eq!(n.id, 42);
2876 assert_eq!(n.labels.len(), 1);
2877 assert_eq!(
2878 n.properties.get("name").unwrap().as_string().unwrap(),
2879 "Alice"
2880 );
2881 }
2882}