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 tenant: Option<String>,
367 role: Option<String>,
369 connect_timeout_secs: u64,
371 hello_timeout_secs: u64,
373 idle_timeout_secs: u64,
375}
376
377impl Client {
378 pub fn new(host: impl Into<String>, port: u16) -> Self {
398 Self {
399 transport: Transport::Quic,
400 host: host.into(),
401 port,
402 tls_enabled: true,
403 skip_verify: false,
404 page_size: 1000,
405 hello_name: "geode-rust".to_string(),
406 hello_ver: env!("CARGO_PKG_VERSION").to_string(),
407 conformance: "min".to_string(),
408 username: None,
409 password: None,
410 graph: None,
411 tenant: None,
412 role: None,
413 connect_timeout_secs: 10,
414 hello_timeout_secs: 15,
415 idle_timeout_secs: 30,
416 }
417 }
418
419 pub fn from_dsn(dsn_str: &str) -> Result<Self> {
467 let dsn = Dsn::parse(dsn_str)?;
468
469 Ok(Self {
470 transport: dsn.transport(),
471 host: dsn.host().to_string(),
472 port: dsn.port(),
473 tls_enabled: dsn.tls_enabled(),
474 skip_verify: dsn.skip_verify(),
475 page_size: dsn.page_size(),
476 hello_name: dsn.client_name().to_string(),
477 hello_ver: dsn.client_version().to_string(),
478 conformance: dsn.conformance().to_string(),
479 username: dsn.username().map(String::from),
480 password: dsn.password().map(|p| SecretString::from(p.to_string())),
481 graph: dsn.graph().map(String::from),
482 tenant: dsn.tenant().map(String::from),
483 role: dsn.role().map(String::from),
484 connect_timeout_secs: dsn.connect_timeout_secs().unwrap_or(10),
485 hello_timeout_secs: 15,
486 idle_timeout_secs: 30,
487 })
488 }
489
490 pub fn transport(&self) -> Transport {
492 self.transport
493 }
494
495 pub fn skip_verify(mut self, skip: bool) -> Self {
507 self.skip_verify = skip;
508 self
509 }
510
511 pub fn page_size(mut self, size: usize) -> Self {
520 self.page_size = size;
521 self
522 }
523
524 pub fn client_name(mut self, name: impl Into<String>) -> Self {
532 self.hello_name = name.into();
533 self
534 }
535
536 pub fn client_version(mut self, version: impl Into<String>) -> Self {
542 self.hello_ver = version.into();
543 self
544 }
545
546 pub fn conformance(mut self, level: impl Into<String>) -> Self {
552 self.conformance = level.into();
553 self
554 }
555
556 pub fn graph(mut self, graph: impl Into<String>) -> Self {
565 self.graph = Some(graph.into());
566 self
567 }
568
569 pub fn tenant(mut self, tenant: impl Into<String>) -> Self {
579 self.tenant = Some(tenant.into());
580 self
581 }
582
583 pub fn role(mut self, role: impl Into<String>) -> Self {
593 self.role = Some(role.into());
594 self
595 }
596
597 pub fn username(mut self, username: impl Into<String>) -> Self {
613 self.username = Some(username.into());
614 self
615 }
616
617 pub fn password(mut self, password: impl Into<String>) -> Self {
627 self.password = Some(SecretString::from(password.into()));
628 self
629 }
630
631 pub fn connect_timeout(mut self, seconds: u64) -> Self {
640 self.connect_timeout_secs = seconds.max(1);
641 self
642 }
643
644 pub fn hello_timeout(mut self, seconds: u64) -> Self {
653 self.hello_timeout_secs = seconds.max(1);
654 self
655 }
656
657 pub fn idle_timeout(mut self, seconds: u64) -> Self {
666 self.idle_timeout_secs = seconds.max(1);
667 self
668 }
669
670 pub fn validate(&self) -> Result<()> {
698 validate::hostname(&self.host)?;
700
701 validate::port(self.port)?;
703
704 validate::page_size(self.page_size)?;
706
707 Ok(())
708 }
709
710 pub async fn connect(&self) -> Result<Connection> {
741 self.validate()?;
743
744 let password_ref = self.password.as_ref().map(|s| s.expose_secret());
746
747 match self.transport {
748 Transport::Quic => {
749 Connection::new_quic(
750 &self.host,
751 self.port,
752 self.skip_verify,
753 self.page_size,
754 &self.hello_name,
755 &self.hello_ver,
756 &self.conformance,
757 self.username.as_deref(),
758 password_ref,
759 self.graph.as_deref(),
760 self.tenant.as_deref(),
761 self.role.as_deref(),
762 self.connect_timeout_secs,
763 self.hello_timeout_secs,
764 self.idle_timeout_secs,
765 )
766 .await
767 }
768 Transport::Grpc => {
769 #[cfg(feature = "grpc")]
770 {
771 Connection::new_grpc(
772 &self.host,
773 self.port,
774 self.tls_enabled,
775 self.skip_verify,
776 self.page_size,
777 self.username.as_deref(),
778 password_ref,
779 self.graph.as_deref(),
780 self.tenant.as_deref(),
781 self.role.as_deref(),
782 )
783 .await
784 }
785 #[cfg(not(feature = "grpc"))]
786 {
787 Err(Error::connection(
788 "gRPC transport requires the 'grpc' feature to be enabled",
789 ))
790 }
791 }
792 }
793 }
794}
795
796#[allow(dead_code)]
798enum ConnectionKind {
799 Quic {
801 conn: quinn::Connection,
802 send: quinn::SendStream,
803 recv: quinn::RecvStream,
804 buffer: Vec<u8>,
806 next_request_id: u64,
808 session_id: String,
810 },
811 #[cfg(feature = "grpc")]
813 Grpc {
814 client: Box<crate::grpc::GrpcClient>,
815 },
816}
817
818pub struct Connection {
863 kind: ConnectionKind,
864 #[allow(dead_code)]
866 page_size: usize,
867 in_transaction: bool,
869}
870
871impl Connection {
872 #[allow(clippy::too_many_arguments)]
874 async fn new_quic(
875 host: &str,
876 port: u16,
877 skip_verify: bool,
878 page_size: usize,
879 hello_name: &str,
880 hello_ver: &str,
881 conformance: &str,
882 username: Option<&str>,
883 password: Option<&str>,
884 graph: Option<&str>,
885 tenant: Option<&str>,
886 role: Option<&str>,
887 connect_timeout_secs: u64,
888 hello_timeout_secs: u64,
889 idle_timeout_secs: u64,
890 ) -> Result<Self> {
891 let mut last_err: Option<Error> = None;
892
893 for attempt in 1..=3 {
894 match Self::connect_quic_once(
895 host,
896 port,
897 skip_verify,
898 page_size,
899 hello_name,
900 hello_ver,
901 conformance,
902 username,
903 password,
904 graph,
905 tenant,
906 role,
907 connect_timeout_secs,
908 hello_timeout_secs,
909 idle_timeout_secs,
910 )
911 .await
912 {
913 Ok(conn) => return Ok(conn),
914 Err(e) => {
915 last_err = Some(e);
916 if attempt < 3 {
917 debug!("Connection attempt {} failed, retrying...", attempt);
918 tokio::time::sleep(Duration::from_millis(150)).await;
919 }
920 }
921 }
922 }
923
924 Err(last_err.unwrap_or_else(|| Error::connection("Failed to connect")))
925 }
926
927 #[cfg(feature = "grpc")]
929 #[allow(clippy::too_many_arguments)]
930 async fn new_grpc(
931 host: &str,
932 port: u16,
933 tls_enabled: bool,
934 skip_verify: bool,
935 page_size: usize,
936 username: Option<&str>,
937 password: Option<&str>,
938 graph: Option<&str>,
939 tenant: Option<&str>,
940 role: Option<&str>,
941 ) -> Result<Self> {
942 use crate::dsn::Dsn;
943
944 let tls_val = if tls_enabled { "1" } else { "0" };
946 let graph_suffix = graph
947 .map(|g| format!("&graph={}", urlencoding::encode(g)))
948 .unwrap_or_default();
949 let tenant_suffix = tenant
950 .map(|t| format!("&tenant={}", urlencoding::encode(t)))
951 .unwrap_or_default();
952 let role_suffix = role
953 .map(|r| format!("&role={}", urlencoding::encode(r)))
954 .unwrap_or_default();
955 let dsn_str = if let (Some(user), Some(pass)) = (username, password) {
956 format!(
957 "grpc://{}:{}@{}:{}?tls={}&insecure={}{}{}{}",
958 user,
959 pass,
960 host,
961 port,
962 tls_val,
963 skip_verify,
964 graph_suffix,
965 tenant_suffix,
966 role_suffix
967 )
968 } else {
969 format!(
970 "grpc://{}:{}?tls={}&insecure={}{}{}{}",
971 host, port, tls_val, skip_verify, graph_suffix, tenant_suffix, role_suffix
972 )
973 };
974
975 let dsn = Dsn::parse(&dsn_str)?;
976 let client = Box::new(crate::grpc::GrpcClient::connect(&dsn).await?);
977
978 Ok(Self {
979 kind: ConnectionKind::Grpc { client },
980 page_size,
981 in_transaction: false,
982 })
983 }
984
985 #[allow(clippy::too_many_arguments)]
986 async fn connect_quic_once(
987 host: &str,
988 port: u16,
989 skip_verify: bool,
990 page_size: usize,
991 hello_name: &str,
992 hello_ver: &str,
993 conformance: &str,
994 username: Option<&str>,
995 password: Option<&str>,
996 graph: Option<&str>,
997 tenant: Option<&str>,
998 role: Option<&str>,
999 connect_timeout_secs: u64,
1000 hello_timeout_secs: u64,
1001 idle_timeout_secs: u64,
1002 ) -> Result<Self> {
1003 debug!("Creating connection to {}:{}", host, port);
1004
1005 let _ = rustls::crypto::aws_lc_rs::default_provider().install_default();
1007
1008 let mut client_crypto = if skip_verify {
1010 warn!(
1013 "TLS certificate verification DISABLED - connection to {}:{} is vulnerable to MITM attacks. \
1014 Do NOT use skip_verify in production!",
1015 host, port
1016 );
1017 rustls::ClientConfig::builder_with_protocol_versions(&[&rustls::version::TLS13])
1018 .dangerous()
1019 .with_custom_certificate_verifier(Arc::new(SkipServerVerification))
1020 .with_no_client_auth()
1021 } else {
1022 let mut root_store = rustls::RootCertStore::empty();
1024
1025 let cert_result = rustls_native_certs::load_native_certs();
1026
1027 for err in &cert_result.errors {
1029 warn!("Error loading native certificate: {:?}", err);
1030 }
1031
1032 let mut certs_loaded = 0;
1033 let mut certs_failed = 0;
1034
1035 for cert in cert_result.certs {
1036 match root_store.add(cert) {
1037 Ok(()) => certs_loaded += 1,
1038 Err(_) => certs_failed += 1,
1039 }
1040 }
1041
1042 if certs_loaded == 0 {
1043 return Err(Error::tls(
1044 "No system root certificates found. TLS verification cannot proceed. \
1045 Either install system CA certificates or use skip_verify(true) for development only.",
1046 ));
1047 }
1048
1049 debug!(
1050 "Loaded {} system root certificates ({} failed to parse)",
1051 certs_loaded, certs_failed
1052 );
1053
1054 rustls::ClientConfig::builder_with_protocol_versions(&[&rustls::version::TLS13])
1055 .with_root_certificates(root_store)
1056 .with_no_client_auth()
1057 };
1058
1059 client_crypto.alpn_protocols = vec![GEODE_ALPN.to_vec()];
1061
1062 let mut client_config = ClientConfig::new(Arc::new(
1063 quinn::crypto::rustls::QuicClientConfig::try_from(client_crypto)
1064 .map_err(|e| Error::connection(format!("Failed to create QUIC config: {}", e)))?,
1065 ));
1066
1067 let mut transport = quinn::TransportConfig::default();
1069 let idle_timeout = Duration::from_secs(idle_timeout_secs.min(146_000 * 365 * 24 * 3600));
1072 transport.max_idle_timeout(Some(idle_timeout.try_into().map_err(|_| {
1073 Error::connection("Idle timeout value too large for QUIC protocol")
1074 })?));
1075 transport.keep_alive_interval(Some(Duration::from_secs(5)));
1076 client_config.transport_config(Arc::new(transport));
1077
1078 let mut endpoint = Endpoint::client(
1081 "0.0.0.0:0"
1082 .parse()
1083 .expect("0.0.0.0:0 is a valid socket address"),
1084 )
1085 .map_err(|e| Error::connection(format!("Failed to create endpoint: {}", e)))?;
1086 endpoint.set_default_client_config(client_config);
1087
1088 let mut resolved_addrs = format!("{}:{}", host, port)
1090 .to_socket_addrs()
1091 .map_err(|e| {
1092 Error::connection(format!(
1093 "Failed to resolve address {}:{} - {}",
1094 host, port, e
1095 ))
1096 })?;
1097
1098 let server_addr: SocketAddr = resolved_addrs
1099 .find(|addr| matches!(addr, SocketAddr::V4(_) | SocketAddr::V6(_)))
1100 .ok_or_else(|| Error::connection("Invalid address: could not resolve host"))?;
1101
1102 debug!("Connecting to {}", server_addr);
1103
1104 let server_name = if skip_verify {
1107 "localhost" } else {
1109 host
1110 };
1111
1112 trace!("Using server name for SNI: {}", server_name);
1113
1114 let conn = timeout(
1115 Duration::from_secs(connect_timeout_secs),
1116 endpoint
1117 .connect(server_addr, server_name)
1118 .map_err(|e| Error::connection(format!("Connection failed: {}", e)))?,
1119 )
1120 .await
1121 .map_err(|_| Error::connection("Connection timeout"))?
1122 .map_err(|e| Error::connection(format!("Failed to establish connection: {}", e)))?;
1123
1124 debug!("Connection established to {}:{}", host, port);
1125
1126 let (mut send, mut recv) = conn
1128 .open_bi()
1129 .await
1130 .map_err(|e| Error::connection(format!("Failed to open stream: {}", e)))?;
1131
1132 let hello_req = proto::HelloRequest {
1134 username: username.unwrap_or("").to_string(),
1135 password: password.unwrap_or("").to_string(),
1136 tenant_id: tenant.map(String::from),
1137 client_name: hello_name.to_string(),
1138 client_version: hello_ver.to_string(),
1139 wanted_conformance: conformance.to_string(),
1140 graph: graph.map(String::from),
1141 role: role.map(String::from),
1142 };
1143 let msg = proto::QuicClientMessage {
1144 msg: Some(proto::quic_client_message::Msg::Hello(hello_req)),
1145 };
1146 let data = proto::encode_with_length_prefix(&msg);
1147
1148 send.write_all(&data)
1149 .await
1150 .map_err(|e| Error::connection(format!("Failed to send HELLO: {}", e)))?;
1151
1152 let mut length_buf = [0u8; 4];
1154 timeout(
1155 Duration::from_secs(hello_timeout_secs),
1156 recv.read_exact(&mut length_buf),
1157 )
1158 .await
1159 .map_err(|_| Error::connection("HELLO response timeout"))?
1160 .map_err(|e| Error::connection(format!("Failed to read HELLO response length: {}", e)))?;
1161
1162 let msg_len = u32::from_be_bytes(length_buf) as usize;
1163
1164 if msg_len > MAX_PROTO_FRAME_BYTES {
1165 return Err(Error::limit(format!(
1166 "HELLO response frame size {} bytes exceeds max {} bytes",
1167 msg_len, MAX_PROTO_FRAME_BYTES
1168 )));
1169 }
1170
1171 let mut msg_buf = vec![0u8; msg_len];
1172 recv.read_exact(&mut msg_buf)
1173 .await
1174 .map_err(|e| Error::connection(format!("Failed to read HELLO response body: {}", e)))?;
1175
1176 let hello_response = proto::decode_quic_server_message(&msg_buf)?;
1177
1178 let session_id = match hello_response.msg {
1179 Some(proto::quic_server_message::Msg::Hello(ref hello_resp)) => {
1180 if !hello_resp.success {
1181 return Err(Error::connection(format!(
1182 "Authentication failed: {}",
1183 hello_resp.error_message
1184 )));
1185 }
1186 hello_resp.session_id.clone()
1187 }
1188 _ => {
1189 return Err(Error::connection("Expected HELLO response"));
1190 }
1191 };
1192
1193 debug!("HELLO handshake complete, session_id={}", session_id);
1194
1195 Ok(Self {
1196 kind: ConnectionKind::Quic {
1197 conn,
1198 send,
1199 recv,
1200 buffer: Vec::new(),
1201 next_request_id: 1,
1202 session_id,
1203 },
1204 page_size,
1205 in_transaction: false,
1206 })
1207 }
1208
1209 async fn send_proto_quic(
1211 send: &mut quinn::SendStream,
1212 msg: &proto::QuicClientMessage,
1213 ) -> Result<()> {
1214 let data = proto::encode_with_length_prefix(msg);
1215 send.write_all(&data)
1216 .await
1217 .map_err(|e| Error::connection(format!("Failed to send message: {}", e)))?;
1218 Ok(())
1219 }
1220
1221 async fn read_proto_quic(
1226 recv: &mut quinn::RecvStream,
1227 timeout_secs: u64,
1228 ) -> Result<proto::QuicServerMessage> {
1229 timeout(Duration::from_secs(timeout_secs), async {
1230 let mut length_buf = [0u8; 4];
1232 recv.read_exact(&mut length_buf)
1233 .await
1234 .map_err(|e| Error::connection(format!("Failed to read response length: {}", e)))?;
1235
1236 let msg_len = u32::from_be_bytes(length_buf) as usize;
1237
1238 if msg_len > MAX_PROTO_FRAME_BYTES {
1240 return Err(Error::limit(format!(
1241 "Frame size {} bytes exceeds max {} bytes",
1242 msg_len, MAX_PROTO_FRAME_BYTES
1243 )));
1244 }
1245
1246 let mut msg_buf = vec![0u8; msg_len];
1247 recv.read_exact(&mut msg_buf)
1248 .await
1249 .map_err(|e| Error::connection(format!("Failed to read response body: {}", e)))?;
1250
1251 proto::decode_quic_server_message(&msg_buf)
1252 })
1253 .await
1254 .map_err(|_| Error::timeout())?
1255 }
1256
1257 fn parse_proto_rows_static(
1259 proto_rows: &[proto::Row],
1260 columns: &[Column],
1261 ) -> Result<Vec<HashMap<String, Value>>> {
1262 let mut rows = Vec::with_capacity(proto_rows.len());
1263 for proto_row in proto_rows {
1264 let mut row = HashMap::with_capacity(columns.len());
1265 for (i, col) in columns.iter().enumerate() {
1266 let value = if i < proto_row.values.len() {
1267 crate::convert::proto_to_value(&proto_row.values[i])
1268 } else {
1269 Value::null()
1270 };
1271 row.insert(col.name.clone(), value);
1272 }
1273 rows.push(row);
1274 }
1275 Ok(rows)
1276 }
1277
1278 async fn send_begin_quic(
1280 send: &mut quinn::SendStream,
1281 recv: &mut quinn::RecvStream,
1282 session_id: &str,
1283 ) -> Result<()> {
1284 let msg = proto::QuicClientMessage {
1285 msg: Some(proto::quic_client_message::Msg::Begin(
1286 proto::BeginRequest {
1287 session_id: session_id.to_string(),
1288 ..Default::default()
1289 },
1290 )),
1291 };
1292 Self::send_proto_quic(send, &msg).await?;
1293
1294 let resp = Self::read_proto_quic(recv, 5).await?;
1295 if !matches!(resp.msg, Some(proto::quic_server_message::Msg::Begin(_))) {
1296 return Err(Error::protocol("Expected BEGIN response"));
1297 }
1298 Ok(())
1299 }
1300
1301 async fn send_commit_quic(
1303 send: &mut quinn::SendStream,
1304 recv: &mut quinn::RecvStream,
1305 session_id: &str,
1306 ) -> Result<()> {
1307 let msg = proto::QuicClientMessage {
1308 msg: Some(proto::quic_client_message::Msg::Commit(
1309 proto::CommitRequest {
1310 session_id: session_id.to_string(),
1311 },
1312 )),
1313 };
1314 Self::send_proto_quic(send, &msg).await?;
1315
1316 let resp = Self::read_proto_quic(recv, 5).await?;
1317 if !matches!(resp.msg, Some(proto::quic_server_message::Msg::Commit(_))) {
1318 return Err(Error::protocol("Expected COMMIT response"));
1319 }
1320 Ok(())
1321 }
1322
1323 async fn send_rollback_quic(
1325 send: &mut quinn::SendStream,
1326 recv: &mut quinn::RecvStream,
1327 session_id: &str,
1328 ) -> Result<()> {
1329 let msg = proto::QuicClientMessage {
1330 msg: Some(proto::quic_client_message::Msg::Rollback(
1331 proto::RollbackRequest {
1332 session_id: session_id.to_string(),
1333 },
1334 )),
1335 };
1336 Self::send_proto_quic(send, &msg).await?;
1337
1338 let resp = Self::read_proto_quic(recv, 5).await?;
1339 if !matches!(resp.msg, Some(proto::quic_server_message::Msg::Rollback(_))) {
1340 return Err(Error::protocol("Expected ROLLBACK response"));
1341 }
1342 Ok(())
1343 }
1344
1345 async fn send_savepoint_quic(
1347 send: &mut quinn::SendStream,
1348 recv: &mut quinn::RecvStream,
1349 session_id: &str,
1350 name: &str,
1351 ) -> Result<()> {
1352 let msg = proto::QuicClientMessage {
1353 msg: Some(proto::quic_client_message::Msg::Savepoint(
1354 proto::SavepointRequest {
1355 name: name.to_string(),
1356 session_id: session_id.to_string(),
1357 },
1358 )),
1359 };
1360 Self::send_proto_quic(send, &msg).await?;
1361
1362 let resp = Self::read_proto_quic(recv, 5).await?;
1363 if !matches!(
1364 resp.msg,
1365 Some(proto::quic_server_message::Msg::Savepoint(_))
1366 ) {
1367 return Err(Error::protocol("Expected SAVEPOINT response"));
1368 }
1369 Ok(())
1370 }
1371
1372 async fn send_rollback_to_quic(
1374 send: &mut quinn::SendStream,
1375 recv: &mut quinn::RecvStream,
1376 session_id: &str,
1377 name: &str,
1378 ) -> Result<()> {
1379 let msg = proto::QuicClientMessage {
1380 msg: Some(proto::quic_client_message::Msg::RollbackTo(
1381 proto::RollbackToRequest {
1382 name: name.to_string(),
1383 session_id: session_id.to_string(),
1384 },
1385 )),
1386 };
1387 Self::send_proto_quic(send, &msg).await?;
1388
1389 let resp = Self::read_proto_quic(recv, 5).await?;
1390 if !matches!(
1391 resp.msg,
1392 Some(proto::quic_server_message::Msg::RollbackTo(_))
1393 ) {
1394 return Err(Error::protocol("Expected ROLLBACK_TO response"));
1395 }
1396 Ok(())
1397 }
1398
1399 pub async fn query(&mut self, gql: &str) -> Result<(Page, Option<String>)> {
1429 self.query_with_params(gql, &HashMap::new()).await
1430 }
1431
1432 pub async fn query_with_params(
1472 &mut self,
1473 gql: &str,
1474 params: &HashMap<String, Value>,
1475 ) -> Result<(Page, Option<String>)> {
1476 validate::query(gql)?;
1477 for key in params.keys() {
1478 validate::param_name(key)?;
1479 }
1480 match &mut self.kind {
1481 ConnectionKind::Quic {
1482 send,
1483 recv,
1484 buffer,
1485 session_id,
1486 ..
1487 } => {
1488 Self::query_with_params_quic(
1489 send,
1490 recv,
1491 buffer,
1492 gql,
1493 params,
1494 session_id,
1495 self.page_size,
1496 )
1497 .await
1498 }
1499 #[cfg(feature = "grpc")]
1500 ConnectionKind::Grpc { client } => client.query_with_params(gql, params).await,
1501 }
1502 }
1503
1504 async fn query_with_params_quic(
1506 send: &mut quinn::SendStream,
1507 recv: &mut quinn::RecvStream,
1508 buffer: &mut Vec<u8>,
1509 gql: &str,
1510 params: &HashMap<String, Value>,
1511 session_id: &str,
1512 page_size: usize,
1513 ) -> Result<(Page, Option<String>)> {
1514 let (page, cursor) =
1515 Self::query_with_params_quic_inner(send, recv, buffer, 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: page_size as u32,
1548 session_id: session_id.to_string(),
1549 })),
1550 };
1551 Self::send_proto_quic(send, &pull_req).await?;
1552
1553 let resp = Self::read_proto_quic_buffered(recv, buffer, 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 Self::drain_execute_trailers_quic(recv, buffer).await?;
1598 return Ok((final_page, cursor));
1599 }
1600
1601 if page.final_page {
1602 Self::drain_execute_trailers_quic(recv, buffer).await?;
1603 }
1604 Ok((page, cursor))
1605 }
1606
1607 async fn query_with_params_quic_inner(
1609 send: &mut quinn::SendStream,
1610 recv: &mut quinn::RecvStream,
1611 buffer: &mut Vec<u8>,
1612 gql: &str,
1613 params: &HashMap<String, Value>,
1614 session_id: &str,
1615 ) -> Result<(Page, Option<String>)> {
1616 let params_proto: Vec<proto::Param> = params
1618 .iter()
1619 .map(|(k, v)| proto::Param {
1620 name: k.clone(),
1621 value: Some(v.to_proto_value()),
1622 })
1623 .collect();
1624
1625 let exec_req = proto::ExecuteRequest {
1627 session_id: session_id.to_string(),
1628 query: gql.to_string(),
1629 params: params_proto,
1630 };
1631 let msg = proto::QuicClientMessage {
1632 msg: Some(proto::quic_client_message::Msg::Execute(exec_req)),
1633 };
1634 Self::send_proto_quic(send, &msg)
1635 .await
1636 .map_err(|e| Error::query(format!("{}", e)))?;
1637
1638 let resp = Self::read_proto_quic_buffered(recv, buffer, 10).await?;
1640
1641 let exec_resp = match resp.msg {
1642 Some(proto::quic_server_message::Msg::Execute(e)) => e,
1643 _ => return Err(Error::protocol("Expected Execute response")),
1644 };
1645
1646 if let Some(proto::execution_response::Payload::Error(ref err)) = exec_resp.payload {
1648 let _ = Self::try_read_proto_quic_buffered(recv, buffer).await;
1650 return Err(Error::Query {
1651 code: err.code.clone(),
1652 message: err.message.clone(),
1653 });
1654 }
1655
1656 let columns: Vec<Column> = match exec_resp.payload {
1658 Some(proto::execution_response::Payload::Schema(ref s)) => s
1659 .columns
1660 .iter()
1661 .map(|c| Column {
1662 name: c.name.clone(),
1663 col_type: c.r#type.clone(),
1664 })
1665 .collect(),
1666 _ => Vec::new(),
1667 };
1668
1669 trace!("Schema columns: {:?}", columns);
1670
1671 while let Some(inline_resp) = Self::try_read_proto_quic_buffered(recv, buffer).await? {
1674 if let Some(proto::quic_server_message::Msg::Execute(inline_exec)) = inline_resp.msg {
1675 if let Some(proto::execution_response::Payload::Error(ref err)) =
1676 inline_exec.payload
1677 {
1678 return Err(Error::Query {
1679 code: err.code.clone(),
1680 message: err.message.clone(),
1681 });
1682 }
1683
1684 if let Some(proto::execution_response::Payload::Page(ref page_data)) =
1685 inline_exec.payload
1686 {
1687 let rows = Self::parse_proto_rows_static(&page_data.rows, &columns)?;
1688 let page = Page {
1689 columns,
1690 rows,
1691 ordered: page_data.ordered,
1692 order_keys: page_data.order_keys.clone(),
1693 final_page: page_data.r#final,
1694 };
1695 return Ok((page, None));
1696 }
1697 }
1698 }
1699
1700 if columns.is_empty() {
1702 return Ok((
1703 Page {
1704 columns,
1705 rows: Vec::new(),
1706 ordered: false,
1707 order_keys: Vec::new(),
1708 final_page: true,
1709 },
1710 None,
1711 ));
1712 }
1713
1714 if let Some(proto::execution_response::Payload::Page(ref page_data)) = exec_resp.payload {
1716 let rows = Self::parse_proto_rows_static(&page_data.rows, &columns)?;
1717 let page = Page {
1718 columns,
1719 rows,
1720 ordered: page_data.ordered,
1721 order_keys: page_data.order_keys.clone(),
1722 final_page: page_data.r#final,
1723 };
1724 return Ok((page, None));
1725 }
1726
1727 loop {
1729 let resp = Self::read_proto_quic_buffered(recv, buffer, 30).await?;
1730 if let Some(proto::quic_server_message::Msg::Execute(exec_resp)) = resp.msg {
1731 if let Some(proto::execution_response::Payload::Error(ref err)) = exec_resp.payload
1732 {
1733 return Err(Error::Query {
1734 code: err.code.clone(),
1735 message: err.message.clone(),
1736 });
1737 }
1738
1739 if let Some(proto::execution_response::Payload::Page(ref page_data)) =
1740 exec_resp.payload
1741 {
1742 let rows = Self::parse_proto_rows_static(&page_data.rows, &columns)?;
1743 let page = Page {
1744 columns,
1745 rows,
1746 ordered: page_data.ordered,
1747 order_keys: page_data.order_keys.clone(),
1748 final_page: page_data.r#final,
1749 };
1750 return Ok((page, None));
1751 }
1752 }
1753 }
1754 }
1755
1756 fn decode_buffered_quic_message(
1757 buffer: &mut Vec<u8>,
1758 ) -> Result<Option<proto::QuicServerMessage>> {
1759 if buffer.len() < 4 {
1760 return Ok(None);
1761 }
1762
1763 let msg_len = u32::from_be_bytes([buffer[0], buffer[1], buffer[2], buffer[3]]) as usize;
1764 if msg_len > MAX_PROTO_FRAME_BYTES {
1765 return Err(Error::limit(format!(
1766 "Frame size {} bytes exceeds max {} bytes",
1767 msg_len, MAX_PROTO_FRAME_BYTES
1768 )));
1769 }
1770
1771 let total_len = 4 + msg_len;
1772 if buffer.len() < total_len {
1773 return Ok(None);
1774 }
1775
1776 let msg = proto::decode_quic_server_message(&buffer[4..total_len])?;
1777 buffer.drain(..total_len);
1778 Ok(Some(msg))
1779 }
1780
1781 async fn read_proto_quic_buffered(
1782 recv: &mut quinn::RecvStream,
1783 buffer: &mut Vec<u8>,
1784 timeout_secs: u64,
1785 ) -> Result<proto::QuicServerMessage> {
1786 if let Some(msg) = Self::decode_buffered_quic_message(buffer)? {
1787 return Ok(msg);
1788 }
1789
1790 let deadline = Instant::now() + Duration::from_secs(timeout_secs);
1791 loop {
1792 let now = Instant::now();
1793 if now >= deadline {
1794 return Err(Error::timeout());
1795 }
1796
1797 let remaining = deadline.saturating_duration_since(now);
1798 let chunk = timeout(remaining, recv.read_chunk(MAX_PROTO_FRAME_BYTES, true))
1799 .await
1800 .map_err(|_| Error::timeout())?
1801 .map_err(|e| Error::connection(format!("Failed to read response: {}", e)))?
1802 .ok_or_else(|| Error::connection("Stream closed while reading response"))?;
1803 buffer.extend_from_slice(&chunk.bytes);
1804
1805 if let Some(msg) = Self::decode_buffered_quic_message(buffer)? {
1806 return Ok(msg);
1807 }
1808 }
1809 }
1810
1811 async fn try_read_proto_quic_buffered(
1812 recv: &mut quinn::RecvStream,
1813 buffer: &mut Vec<u8>,
1814 ) -> Result<Option<proto::QuicServerMessage>> {
1815 if let Some(msg) = Self::decode_buffered_quic_message(buffer)? {
1816 return Ok(Some(msg));
1817 }
1818
1819 let read_result = timeout(
1820 Duration::from_millis(500),
1821 recv.read_chunk(MAX_PROTO_FRAME_BYTES, true),
1822 )
1823 .await;
1824 let chunk = match read_result {
1825 Ok(Ok(Some(chunk))) => chunk,
1826 Ok(Ok(None)) => return Ok(None),
1827 Ok(Err(e)) => {
1828 return Err(Error::connection(format!("Failed to read response: {}", e)));
1829 }
1830 Err(_) => return Ok(None),
1831 };
1832
1833 buffer.extend_from_slice(&chunk.bytes);
1834 Self::decode_buffered_quic_message(buffer)
1835 }
1836
1837 async fn drain_execute_trailers_quic(
1838 recv: &mut quinn::RecvStream,
1839 buffer: &mut Vec<u8>,
1840 ) -> Result<()> {
1841 while let Some(resp) = Self::try_read_proto_quic_buffered(recv, buffer).await? {
1842 let Some(proto::quic_server_message::Msg::Execute(exec)) = resp.msg else {
1843 break;
1844 };
1845
1846 let is_trailer = matches!(
1847 exec.payload,
1848 Some(proto::execution_response::Payload::Metrics(_))
1849 | Some(proto::execution_response::Payload::Heartbeat(_))
1850 );
1851 if !is_trailer {
1852 break;
1853 }
1854 }
1855 Ok(())
1856 }
1857
1858 pub fn query_sync(
1860 &mut self,
1861 gql: &str,
1862 params: Option<HashMap<String, serde_json::Value>>,
1863 ) -> Result<Page> {
1864 let params_map = params.unwrap_or_default();
1865 let mut params_typed: HashMap<String, Value> = HashMap::new();
1866 for (k, v) in params_map {
1867 let typed_val = crate::types::Value::from_json(v)?;
1868 params_typed.insert(k, typed_val);
1869 }
1870
1871 match tokio::runtime::Handle::try_current() {
1872 Ok(handle) => {
1873 let (page, _cursor) =
1874 handle.block_on(self.query_with_params(gql, ¶ms_typed))?;
1875 Ok(page)
1876 }
1877 Err(_) => {
1878 let rt = tokio::runtime::Runtime::new()
1879 .map_err(|e| Error::query(format!("Failed to create runtime: {}", e)))?;
1880 let (page, _cursor) = rt.block_on(self.query_with_params(gql, ¶ms_typed))?;
1881 Ok(page)
1882 }
1883 }
1884 }
1885
1886 pub async fn begin(&mut self) -> Result<()> {
1911 let result = match &mut self.kind {
1912 ConnectionKind::Quic {
1913 send,
1914 recv,
1915 session_id,
1916 ..
1917 } => Self::send_begin_quic(send, recv, session_id).await,
1918 #[cfg(feature = "grpc")]
1919 ConnectionKind::Grpc { client } => client.begin().await,
1920 };
1921 if result.is_ok() {
1922 self.in_transaction = true;
1923 }
1924 result
1925 }
1926
1927 pub async fn commit(&mut self) -> Result<()> {
1950 let result = match &mut self.kind {
1951 ConnectionKind::Quic {
1952 send,
1953 recv,
1954 session_id,
1955 ..
1956 } => Self::send_commit_quic(send, recv, session_id).await,
1957 #[cfg(feature = "grpc")]
1958 ConnectionKind::Grpc { client } => client.commit().await,
1959 };
1960 if result.is_ok() {
1961 self.in_transaction = false;
1962 }
1963 result
1964 }
1965
1966 pub async fn rollback(&mut self) -> Result<()> {
1990 let result = match &mut self.kind {
1991 ConnectionKind::Quic {
1992 send,
1993 recv,
1994 session_id,
1995 ..
1996 } => Self::send_rollback_quic(send, recv, session_id).await,
1997 #[cfg(feature = "grpc")]
1998 ConnectionKind::Grpc { client } => client.rollback().await,
1999 };
2000 if result.is_ok() {
2001 self.in_transaction = false;
2002 }
2003 result
2004 }
2005
2006 pub async fn savepoint(&mut self, name: &str) -> Result<Savepoint> {
2040 match &mut self.kind {
2041 ConnectionKind::Quic {
2042 send,
2043 recv,
2044 session_id,
2045 ..
2046 } => Self::send_savepoint_quic(send, recv, session_id, name).await?,
2047 #[cfg(feature = "grpc")]
2048 ConnectionKind::Grpc { client } => client.savepoint(name).await?,
2049 }
2050 Ok(Savepoint {
2051 name: name.to_string(),
2052 })
2053 }
2054
2055 pub async fn rollback_to(&mut self, savepoint: &Savepoint) -> Result<()> {
2084 match &mut self.kind {
2085 ConnectionKind::Quic {
2086 send,
2087 recv,
2088 session_id,
2089 ..
2090 } => Self::send_rollback_to_quic(send, recv, session_id, &savepoint.name).await?,
2091 #[cfg(feature = "grpc")]
2092 ConnectionKind::Grpc { client } => client.rollback_to(&savepoint.name).await?,
2093 }
2094 Ok(())
2095 }
2096
2097 pub fn prepare(&self, query: &str) -> Result<PreparedStatement> {
2131 Ok(PreparedStatement::new(query))
2132 }
2133
2134 pub async fn explain(&mut self, gql: &str) -> Result<QueryPlan> {
2167 let explain_query = format!("EXPLAIN {}", gql);
2169 let (_page, _) = self.query(&explain_query).await?;
2170
2171 Ok(QueryPlan {
2174 operations: Vec::new(),
2175 estimated_rows: 0,
2176 raw: serde_json::json!({}),
2177 })
2178 }
2179
2180 pub async fn profile(&mut self, gql: &str) -> Result<QueryProfile> {
2211 let profile_query = format!("PROFILE {}", gql);
2213 let (page, _) = self.query(&profile_query).await?;
2214
2215 let plan = QueryPlan {
2217 operations: Vec::new(),
2218 estimated_rows: 0,
2219 raw: serde_json::json!({}),
2220 };
2221
2222 Ok(QueryProfile {
2223 plan,
2224 actual_rows: page.rows.len() as u64,
2225 execution_time_ms: 0.0,
2226 raw: serde_json::json!({}),
2227 })
2228 }
2229
2230 pub async fn batch(
2268 &mut self,
2269 queries: &[(&str, Option<&HashMap<String, Value>>)],
2270 ) -> Result<Vec<Page>> {
2271 let mut results = Vec::with_capacity(queries.len());
2272
2273 for (query, params) in queries {
2274 let (page, _) = match params {
2275 Some(p) => self.query_with_params(query, p).await?,
2276 None => self.query(query).await?,
2277 };
2278 results.push(page);
2279 }
2280
2281 Ok(results)
2282 }
2283
2284 #[allow(dead_code)]
2287 fn parse_plan_operations(result: &serde_json::Value) -> Vec<PlanOperation> {
2288 let mut operations = Vec::new();
2289
2290 if let Some(ops) = result.get("operations").and_then(|o| o.as_array()) {
2291 for op in ops {
2292 operations.push(Self::parse_single_operation(op));
2293 }
2294 } else if let Some(plan) = result.get("plan") {
2295 operations.push(Self::parse_single_operation(plan));
2297 }
2298
2299 operations
2300 }
2301
2302 #[allow(dead_code)]
2304 fn parse_single_operation(op: &serde_json::Value) -> PlanOperation {
2305 let op_type = op
2306 .get("type")
2307 .or_else(|| op.get("op_type"))
2308 .and_then(|t| t.as_str())
2309 .unwrap_or("Unknown")
2310 .to_string();
2311
2312 let description = op
2313 .get("description")
2314 .or_else(|| op.get("desc"))
2315 .and_then(|d| d.as_str())
2316 .unwrap_or("")
2317 .to_string();
2318
2319 let estimated_rows = op
2320 .get("estimated_rows")
2321 .or_else(|| op.get("rows"))
2322 .and_then(|r| r.as_u64());
2323
2324 let children = op
2325 .get("children")
2326 .and_then(|c| c.as_array())
2327 .map(|arr| arr.iter().map(Self::parse_single_operation).collect())
2328 .unwrap_or_default();
2329
2330 PlanOperation {
2331 op_type,
2332 description,
2333 estimated_rows,
2334 children,
2335 }
2336 }
2337
2338 pub fn close(&mut self) -> Result<()> {
2361 match &mut self.kind {
2362 ConnectionKind::Quic { conn, .. } => {
2363 conn.close(0u32.into(), b"client closing");
2367 Ok(())
2368 }
2369 #[cfg(feature = "grpc")]
2370 ConnectionKind::Grpc { client } => client.close(),
2371 }
2372 }
2373
2374 pub fn in_transaction(&self) -> bool {
2394 self.in_transaction
2395 }
2396
2397 pub fn is_healthy(&self) -> bool {
2398 match &self.kind {
2399 ConnectionKind::Quic { conn, .. } => {
2400 conn.close_reason().is_none()
2402 }
2403 #[cfg(feature = "grpc")]
2404 ConnectionKind::Grpc { .. } => {
2405 true
2407 }
2408 }
2409 }
2410}
2411
2412#[derive(Debug)]
2414struct SkipServerVerification;
2415
2416impl rustls::client::danger::ServerCertVerifier for SkipServerVerification {
2417 fn verify_server_cert(
2418 &self,
2419 _end_entity: &CertificateDer,
2420 _intermediates: &[CertificateDer],
2421 _server_name: &RustlsServerName,
2422 _ocsp_response: &[u8],
2423 _now: rustls::pki_types::UnixTime,
2424 ) -> std::result::Result<rustls::client::danger::ServerCertVerified, rustls::Error> {
2425 Ok(rustls::client::danger::ServerCertVerified::assertion())
2426 }
2427
2428 fn verify_tls12_signature(
2429 &self,
2430 _message: &[u8],
2431 _cert: &CertificateDer,
2432 _dss: &rustls::DigitallySignedStruct,
2433 ) -> std::result::Result<rustls::client::danger::HandshakeSignatureValid, rustls::Error> {
2434 Ok(rustls::client::danger::HandshakeSignatureValid::assertion())
2435 }
2436
2437 fn verify_tls13_signature(
2438 &self,
2439 _message: &[u8],
2440 _cert: &CertificateDer,
2441 _dss: &rustls::DigitallySignedStruct,
2442 ) -> std::result::Result<rustls::client::danger::HandshakeSignatureValid, rustls::Error> {
2443 Ok(rustls::client::danger::HandshakeSignatureValid::assertion())
2444 }
2445
2446 fn supported_verify_schemes(&self) -> Vec<rustls::SignatureScheme> {
2447 vec![
2448 rustls::SignatureScheme::RSA_PKCS1_SHA256,
2449 rustls::SignatureScheme::ECDSA_NISTP256_SHA256,
2450 rustls::SignatureScheme::ED25519,
2451 ]
2452 }
2453}
2454
2455#[cfg(test)]
2456mod tests {
2457 use super::*;
2458
2459 #[test]
2462 fn test_prepared_statement_new() {
2463 let stmt = PreparedStatement::new("MATCH (n:Person {id: $id}) RETURN n");
2464 assert_eq!(stmt.query(), "MATCH (n:Person {id: $id}) RETURN n");
2465 assert_eq!(stmt.param_names(), &["id"]);
2466 }
2467
2468 #[test]
2469 fn test_prepared_statement_multiple_params() {
2470 let stmt = PreparedStatement::new(
2471 "MATCH (p:Person {name: $name}) WHERE p.age > $min_age AND p.city = $city RETURN p",
2472 );
2473 assert!(stmt.query().contains("$name"));
2474 let names = stmt.param_names();
2475 assert_eq!(names.len(), 3);
2476 assert!(names.contains(&"name".to_string()));
2477 assert!(names.contains(&"min_age".to_string()));
2478 assert!(names.contains(&"city".to_string()));
2479 }
2480
2481 #[test]
2482 fn test_prepared_statement_no_params() {
2483 let stmt = PreparedStatement::new("MATCH (n) RETURN n LIMIT 10");
2484 assert!(stmt.param_names().is_empty());
2485 }
2486
2487 #[test]
2488 fn test_prepared_statement_duplicate_params() {
2489 let stmt =
2490 PreparedStatement::new("MATCH (a {id: $id})-[:KNOWS]->(b {id: $id}) RETURN a, b");
2491 assert_eq!(stmt.param_names(), &["id"]);
2493 }
2494
2495 #[test]
2496 fn test_prepared_statement_underscore_params() {
2497 let stmt = PreparedStatement::new("MATCH (n {user_id: $user_id}) RETURN n");
2498 assert_eq!(stmt.param_names(), &["user_id"]);
2499 }
2500
2501 #[test]
2502 fn test_prepared_statement_numeric_params() {
2503 let stmt = PreparedStatement::new("RETURN $param1, $param2, $param123");
2504 let names = stmt.param_names();
2505 assert_eq!(names.len(), 3);
2506 assert!(names.contains(&"param1".to_string()));
2507 assert!(names.contains(&"param2".to_string()));
2508 assert!(names.contains(&"param123".to_string()));
2509 }
2510
2511 #[test]
2514 fn test_plan_operation_struct() {
2515 let op = PlanOperation {
2516 op_type: "NodeScan".to_string(),
2517 description: "Scan Person nodes".to_string(),
2518 estimated_rows: Some(100),
2519 children: vec![],
2520 };
2521 assert_eq!(op.op_type, "NodeScan");
2522 assert_eq!(op.description, "Scan Person nodes");
2523 assert_eq!(op.estimated_rows, Some(100));
2524 assert!(op.children.is_empty());
2525 }
2526
2527 #[test]
2528 fn test_plan_operation_with_children() {
2529 let child = PlanOperation {
2530 op_type: "Filter".to_string(),
2531 description: "Filter by age".to_string(),
2532 estimated_rows: Some(50),
2533 children: vec![],
2534 };
2535 let parent = PlanOperation {
2536 op_type: "Projection".to_string(),
2537 description: "Project name, age".to_string(),
2538 estimated_rows: Some(50),
2539 children: vec![child],
2540 };
2541 assert_eq!(parent.children.len(), 1);
2542 assert_eq!(parent.children[0].op_type, "Filter");
2543 }
2544
2545 #[test]
2548 fn test_query_plan_struct() {
2549 let plan = QueryPlan {
2550 operations: vec![PlanOperation {
2551 op_type: "NodeScan".to_string(),
2552 description: "Full scan".to_string(),
2553 estimated_rows: Some(1000),
2554 children: vec![],
2555 }],
2556 estimated_rows: 1000,
2557 raw: serde_json::json!({"type": "plan"}),
2558 };
2559 assert_eq!(plan.operations.len(), 1);
2560 assert_eq!(plan.estimated_rows, 1000);
2561 }
2562
2563 #[test]
2566 fn test_query_profile_struct() {
2567 let plan = QueryPlan {
2568 operations: vec![],
2569 estimated_rows: 100,
2570 raw: serde_json::json!({}),
2571 };
2572 let profile = QueryProfile {
2573 plan,
2574 actual_rows: 95,
2575 execution_time_ms: 12.5,
2576 raw: serde_json::json!({"type": "profile"}),
2577 };
2578 assert_eq!(profile.actual_rows, 95);
2579 assert!((profile.execution_time_ms - 12.5).abs() < 0.001);
2580 }
2581
2582 #[test]
2585 fn test_page_struct() {
2586 let page = Page {
2587 columns: vec![Column {
2588 name: "x".to_string(),
2589 col_type: "INT".to_string(),
2590 }],
2591 rows: vec![],
2592 ordered: false,
2593 order_keys: vec![],
2594 final_page: true,
2595 };
2596 assert_eq!(page.columns.len(), 1);
2597 assert!(page.rows.is_empty());
2598 assert!(page.final_page);
2599 }
2600
2601 #[test]
2604 fn test_column_struct() {
2605 let col = Column {
2606 name: "age".to_string(),
2607 col_type: "INT".to_string(),
2608 };
2609 assert_eq!(col.name, "age");
2610 assert_eq!(col.col_type, "INT");
2611 }
2612
2613 #[test]
2616 fn test_savepoint_struct() {
2617 let sp = Savepoint {
2618 name: "before_update".to_string(),
2619 };
2620 assert_eq!(sp.name, "before_update");
2621 }
2622
2623 #[test]
2626 fn test_client_builder_defaults() {
2627 let _client = Client::new("localhost", 3141);
2628 }
2630
2631 #[test]
2632 fn test_client_builder_chain() {
2633 let _client = Client::new("example.com", 8443)
2634 .skip_verify(true)
2635 .page_size(500)
2636 .client_name("test-app")
2637 .client_version("2.0.0")
2638 .conformance("full");
2639 }
2641
2642 #[test]
2643 fn test_client_clone() {
2644 let client = Client::new("localhost", 3141).skip_verify(true);
2645 let _cloned = client.clone();
2646 }
2648
2649 #[test]
2652 fn test_parse_plan_operations_empty() {
2653 let result = serde_json::json!({});
2654 let ops = Connection::parse_plan_operations(&result);
2655 assert!(ops.is_empty());
2656 }
2657
2658 #[test]
2659 fn test_parse_plan_operations_array() {
2660 let result = serde_json::json!({
2661 "operations": [
2662 {"type": "NodeScan", "description": "Scan nodes", "estimated_rows": 100},
2663 {"type": "Filter", "description": "Apply filter", "estimated_rows": 50}
2664 ]
2665 });
2666 let ops = Connection::parse_plan_operations(&result);
2667 assert_eq!(ops.len(), 2);
2668 assert_eq!(ops[0].op_type, "NodeScan");
2669 assert_eq!(ops[1].op_type, "Filter");
2670 }
2671
2672 #[test]
2673 fn test_parse_plan_operations_single_plan() {
2674 let result = serde_json::json!({
2675 "plan": {"op_type": "FullScan", "desc": "Full table scan"}
2676 });
2677 let ops = Connection::parse_plan_operations(&result);
2678 assert_eq!(ops.len(), 1);
2679 assert_eq!(ops[0].op_type, "FullScan");
2680 assert_eq!(ops[0].description, "Full table scan");
2681 }
2682
2683 #[test]
2684 fn test_parse_single_operation() {
2685 let op_json = serde_json::json!({
2686 "type": "IndexScan",
2687 "description": "Use index on Person(name)",
2688 "estimated_rows": 25,
2689 "children": [
2690 {"type": "Filter", "description": "Filter results"}
2691 ]
2692 });
2693 let op = Connection::parse_single_operation(&op_json);
2694 assert_eq!(op.op_type, "IndexScan");
2695 assert_eq!(op.description, "Use index on Person(name)");
2696 assert_eq!(op.estimated_rows, Some(25));
2697 assert_eq!(op.children.len(), 1);
2698 assert_eq!(op.children[0].op_type, "Filter");
2699 }
2700
2701 #[test]
2702 fn test_parse_single_operation_minimal() {
2703 let op_json = serde_json::json!({});
2704 let op = Connection::parse_single_operation(&op_json);
2705 assert_eq!(op.op_type, "Unknown");
2706 assert_eq!(op.description, "");
2707 assert_eq!(op.estimated_rows, None);
2708 assert!(op.children.is_empty());
2709 }
2710
2711 #[test]
2712 fn test_parse_single_operation_alt_fields() {
2713 let op_json = serde_json::json!({
2714 "op_type": "Sort",
2715 "desc": "Sort by name ASC",
2716 "rows": 100
2717 });
2718 let op = Connection::parse_single_operation(&op_json);
2719 assert_eq!(op.op_type, "Sort");
2720 assert_eq!(op.description, "Sort by name ASC");
2721 assert_eq!(op.estimated_rows, Some(100));
2722 }
2723
2724 #[test]
2727 fn test_redact_dsn_url_with_password() {
2728 let dsn = "quic://admin:secret123@localhost:3141";
2729 let redacted = redact_dsn(dsn);
2730 assert!(redacted.contains("[REDACTED]"));
2731 assert!(!redacted.contains("secret123"));
2732 assert!(redacted.contains("admin"));
2733 assert!(redacted.contains("localhost"));
2734 }
2735
2736 #[test]
2737 fn test_redact_dsn_url_without_password() {
2738 let dsn = "quic://admin@localhost:3141";
2739 let redacted = redact_dsn(dsn);
2740 assert!(!redacted.contains("[REDACTED]"));
2741 assert!(redacted.contains("admin"));
2742 assert!(redacted.contains("localhost"));
2743 }
2744
2745 #[test]
2746 fn test_redact_dsn_url_no_auth() {
2747 let dsn = "quic://localhost:3141";
2748 let redacted = redact_dsn(dsn);
2749 assert_eq!(redacted, dsn);
2750 }
2751
2752 #[test]
2753 fn test_redact_dsn_query_param_password() {
2754 let dsn = "localhost:3141?username=admin&password=secret123";
2755 let redacted = redact_dsn(dsn);
2756 assert!(redacted.contains("[REDACTED]"));
2757 assert!(!redacted.contains("secret123"));
2758 assert!(redacted.contains("username=admin"));
2759 }
2760
2761 #[test]
2762 fn test_redact_dsn_query_param_pass() {
2763 let dsn = "localhost:3141?user=admin&pass=mysecret";
2764 let redacted = redact_dsn(dsn);
2765 assert!(redacted.contains("[REDACTED]"));
2766 assert!(!redacted.contains("mysecret"));
2767 }
2768
2769 #[test]
2770 fn test_redact_dsn_simple_no_password() {
2771 let dsn = "localhost:3141?insecure=true";
2772 let redacted = redact_dsn(dsn);
2773 assert_eq!(redacted, dsn);
2774 }
2775
2776 #[test]
2777 fn test_redact_dsn_url_with_query_and_password() {
2778 let dsn = "quic://user:pass@localhost:3141?insecure=true";
2779 let redacted = redact_dsn(dsn);
2780 assert!(redacted.contains("[REDACTED]"));
2781 assert!(!redacted.contains(":pass@"));
2782 assert!(redacted.contains("insecure=true"));
2783 }
2784
2785 #[test]
2788 fn test_client_validate_valid() {
2789 let client = Client::new("localhost", 3141);
2790 assert!(client.validate().is_ok());
2791 }
2792
2793 #[test]
2794 fn test_client_validate_valid_hostname() {
2795 let client = Client::new("geode.example.com", 3141);
2796 assert!(client.validate().is_ok());
2797 }
2798
2799 #[test]
2800 fn test_client_validate_valid_ipv4() {
2801 let client = Client::new("192.168.1.1", 8443);
2802 assert!(client.validate().is_ok());
2803 }
2804
2805 #[test]
2806 fn test_client_validate_invalid_hostname_hyphen_start() {
2807 let client = Client::new("-invalid", 3141);
2808 assert!(client.validate().is_err());
2809 }
2810
2811 #[test]
2812 fn test_client_validate_invalid_hostname_hyphen_end() {
2813 let client = Client::new("invalid-", 3141);
2814 assert!(client.validate().is_err());
2815 }
2816
2817 #[test]
2818 fn test_client_validate_invalid_port_zero() {
2819 let client = Client::new("localhost", 0);
2820 assert!(client.validate().is_err());
2821 }
2822
2823 #[test]
2824 fn test_client_validate_invalid_page_size_zero() {
2825 let client = Client::new("localhost", 3141).page_size(0);
2826 assert!(client.validate().is_err());
2827 }
2828
2829 #[test]
2830 fn test_client_validate_invalid_page_size_too_large() {
2831 let client = Client::new("localhost", 3141).page_size(200_000);
2832 assert!(client.validate().is_err());
2833 }
2834
2835 #[test]
2836 fn test_client_validate_with_all_options() {
2837 let client = Client::new("geode.example.com", 8443)
2838 .skip_verify(true)
2839 .page_size(500)
2840 .username("admin")
2841 .password("secret")
2842 .connect_timeout(15)
2843 .hello_timeout(10)
2844 .idle_timeout(60);
2845 assert!(client.validate().is_ok());
2846 }
2847
2848 #[test]
2850 fn test_client_extreme_timeout_values() {
2851 let _client = Client::new("localhost", 3141)
2853 .connect_timeout(u64::MAX)
2854 .hello_timeout(u64::MAX)
2855 .idle_timeout(u64::MAX);
2856 }
2858
2859 #[test]
2860 fn test_convert_edge_uses_type_field() {
2861 let edge = proto::EdgeValue {
2862 id: 100,
2863 from_id: 1,
2864 to_id: 2,
2865 label: "KNOWS".to_string(),
2866 properties: vec![],
2867 };
2868 let proto_val = proto::Value {
2869 kind: Some(proto::value::Kind::EdgeVal(edge)),
2870 };
2871 let val = crate::convert::proto_to_value(&proto_val);
2872 let obj = val.as_object().unwrap();
2873 assert_eq!(obj.get("type").unwrap().as_string().unwrap(), "KNOWS");
2874 assert!(
2875 obj.get("label").is_none(),
2876 "edge should not have 'label' field"
2877 );
2878 }
2879
2880 #[test]
2881 fn test_convert_edge_uses_start_end_node() {
2882 let edge = proto::EdgeValue {
2883 id: 100,
2884 from_id: 42,
2885 to_id: 99,
2886 label: "LIKES".to_string(),
2887 properties: vec![],
2888 };
2889 let proto_val = proto::Value {
2890 kind: Some(proto::value::Kind::EdgeVal(edge)),
2891 };
2892 let val = crate::convert::proto_to_value(&proto_val);
2893 let obj = val.as_object().unwrap();
2894 assert_eq!(obj.get("start_node").unwrap().as_int().unwrap(), 42);
2895 assert_eq!(obj.get("end_node").unwrap().as_int().unwrap(), 99);
2896 assert!(obj.get("from_id").is_none());
2897 assert!(obj.get("to_id").is_none());
2898 }
2899
2900 #[test]
2901 fn test_convert_edge_with_properties() {
2902 let edge = proto::EdgeValue {
2903 id: 100,
2904 from_id: 1,
2905 to_id: 2,
2906 label: "KNOWS".to_string(),
2907 properties: vec![proto::MapEntry {
2908 key: "since".to_string(),
2909 value: Some(proto::Value {
2910 kind: Some(proto::value::Kind::IntVal(proto::IntValue {
2911 value: 2020,
2912 kind: 1,
2913 })),
2914 }),
2915 }],
2916 };
2917 let proto_val = proto::Value {
2918 kind: Some(proto::value::Kind::EdgeVal(edge)),
2919 };
2920 let val = crate::convert::proto_to_value(&proto_val);
2921 let obj = val.as_object().unwrap();
2922 let props = obj.get("properties").unwrap().as_object().unwrap();
2923 assert_eq!(props.get("since").unwrap().as_int().unwrap(), 2020);
2924 }
2925
2926 #[test]
2927 fn test_convert_node_fields() {
2928 let node = proto::NodeValue {
2929 id: 42,
2930 labels: vec!["Person".to_string()],
2931 properties: vec![proto::MapEntry {
2932 key: "name".to_string(),
2933 value: Some(proto::Value {
2934 kind: Some(proto::value::Kind::StringVal(proto::StringValue {
2935 value: "Alice".to_string(),
2936 kind: 1,
2937 })),
2938 }),
2939 }],
2940 };
2941 let proto_val = proto::Value {
2942 kind: Some(proto::value::Kind::NodeVal(node)),
2943 };
2944 let val = crate::convert::proto_to_value(&proto_val);
2945 let obj = val.as_object().unwrap();
2946 assert_eq!(obj.get("id").unwrap().as_int().unwrap(), 42);
2947 let labels = obj.get("labels").unwrap().as_array().unwrap();
2948 assert_eq!(labels.len(), 1);
2949 let props = obj.get("properties").unwrap().as_object().unwrap();
2950 assert_eq!(props.get("name").unwrap().as_string().unwrap(), "Alice");
2951 }
2952
2953 fn null_proto_value() -> proto::Value {
2956 proto::Value {
2957 kind: Some(proto::value::Kind::NullVal(proto::NullValue {})),
2958 }
2959 }
2960
2961 fn node_proto_value(id: u64) -> proto::Value {
2962 proto::Value {
2963 kind: Some(proto::value::Kind::NodeVal(proto::NodeValue {
2964 id,
2965 labels: vec!["Person".to_string()],
2966 properties: vec![],
2967 })),
2968 }
2969 }
2970
2971 #[test]
2972 fn test_real_node_row_kept() {
2973 let columns = vec![Column {
2974 name: "n".to_string(),
2975 col_type: "NODE".to_string(),
2976 }];
2977 let proto_rows = vec![proto::Row {
2978 values: vec![node_proto_value(42)],
2979 }];
2980 let rows = Connection::parse_proto_rows_static(&proto_rows, &columns).unwrap();
2981 assert_eq!(rows.len(), 1, "real NODE row must be kept");
2982 let obj = rows[0].get("n").unwrap().as_object().unwrap();
2983 assert_eq!(obj.get("id").unwrap().as_int().unwrap(), 42);
2984 }
2985
2986 #[test]
2987 fn test_scalar_null_row_kept() {
2988 let columns = vec![Column {
2991 name: "x".to_string(),
2992 col_type: "STRING".to_string(),
2993 }];
2994 let proto_rows = vec![proto::Row {
2995 values: vec![null_proto_value()],
2996 }];
2997 let rows = Connection::parse_proto_rows_static(&proto_rows, &columns).unwrap();
2998 assert_eq!(rows.len(), 1, "scalar null row must be kept");
2999 assert!(rows[0].get("x").unwrap().is_null());
3000 }
3001}