1pub(crate) use crate::application::json_input::{
2 json_bool_field, json_f32_field, json_string_field, json_usize_field,
3};
4pub(crate) use crate::application::{
5 AdminUseCases, CatalogUseCases, CreateEdgeInput, CreateEntityOutput, CreateNodeGraphLinkInput,
6 CreateNodeInput, CreateNodeTableLinkInput, CreateRowInput, CreateVectorInput,
7 DeleteEntityInput, EntityUseCases, ExecuteQueryInput, ExplainQueryInput, GraphCentralityInput,
8 GraphClusteringInput, GraphCommunitiesInput, GraphComponentsInput, GraphCyclesInput,
9 GraphHitsInput, GraphNeighborhoodInput, GraphPersonalizedPageRankInput, GraphShortestPathInput,
10 GraphTopologicalSortInput, GraphTraversalInput, GraphUseCases, InspectNativeArtifactInput,
11 NativeUseCases, PatchEntityInput, QueryUseCases, SearchHybridInput, SearchIvfInput,
12 SearchSimilarInput, SearchTextInput,
13};
14use std::collections::BTreeMap;
15use std::pin::Pin;
16use std::sync::Arc;
17use std::time::{SystemTime, UNIX_EPOCH};
18
19use crate::api::{RedDBOptions, RedDBResult};
20use crate::auth::middleware::{check_permission, AuthResult, AuthSource};
21use crate::auth::store::AuthStore;
22use crate::auth::Role;
23use crate::health::{HealthProvider, HealthState};
24use crate::json::{
25 from_str as json_from_str, to_string as json_to_string, Map, Value as JsonValue,
26};
27use crate::runtime::{
28 RedDBRuntime, RuntimeFilter, RuntimeFilterValue, RuntimeGraphCentralityAlgorithm,
29 RuntimeGraphCentralityResult, RuntimeGraphClusteringResult, RuntimeGraphCommunityAlgorithm,
30 RuntimeGraphCommunityResult, RuntimeGraphComponentsMode, RuntimeGraphComponentsResult,
31 RuntimeGraphCyclesResult, RuntimeGraphDirection, RuntimeGraphHitsResult,
32 RuntimeGraphNeighborhoodResult, RuntimeGraphPathAlgorithm, RuntimeGraphPathResult,
33 RuntimeGraphPattern, RuntimeGraphProjection, RuntimeGraphTopologicalSortResult,
34 RuntimeGraphTraversalResult, RuntimeGraphTraversalStrategy, RuntimeIvfSearchResult,
35 RuntimeQueryResult, RuntimeQueryWeights, RuntimeStats, ScanPage,
36};
37use crate::storage::schema::Value;
38use crate::storage::unified::devx::refs::{NodeRef, TableRef};
39use crate::storage::unified::{Metadata, MetadataValue};
40use crate::storage::{EntityData, EntityId, UnifiedEntity};
41use tokio_stream::wrappers::TcpListenerStream;
42use tonic::metadata::MetadataMap;
43use tonic::{Request, Response, Status};
44
45pub use reddb_grpc_proto as proto;
51
52use proto::red_db_server::{RedDb, RedDbServer};
53use proto::{
54 ask_stream_event, AskAnswerToken, AskReply, AskRequest, AskSources, AskStreamEvent,
55 BatchInsertChunk, BatchInsertReply, BatchQueryReply, BatchQueryRequest, BulkEntityReply,
56 Citation, CollectionRequest, CollectionsReply, DeleteEntityRequest, DeploymentProfileRequest,
57 Empty, EntityReply, ExecutePreparedRequest, ExportRequest, GraphProjectionUpsertRequest,
58 HealthReply, IndexNameRequest, IndexToggleRequest, JsonBulkCreateRequest, JsonCreateRequest,
59 JsonPayloadRequest, KvWatchEvent, KvWatchRequest, ManifestRequest, OperationReply,
60 PayloadReply, PrepareQueryReply, PrepareQueryRequest, QueryReply, QueryRequest, QueryValue,
61 ScanEntity, ScanReply, ScanRequest, StatsReply, TopologyReply, TopologyRequest,
62 UpdateEntityRequest, Validation, ValidationItem,
63};
64
65mod control_support;
66mod entity_ops;
67mod input_support;
68pub(crate) mod scan_json;
69
70use self::control_support::*;
71use self::entity_ops::*;
72use self::input_support::*;
73use self::scan_json::*;
74
75#[derive(Debug, Clone)]
76pub struct GrpcServerOptions {
77 pub bind_addr: String,
78 pub tls: Option<GrpcTlsOptions>,
83}
84
85#[derive(Debug, Clone)]
91pub struct GrpcTlsOptions {
92 pub cert_pem: Vec<u8>,
94 pub key_pem: Vec<u8>,
96 pub client_ca_pem: Option<Vec<u8>>,
101}
102
103impl GrpcTlsOptions {
104 pub fn to_tonic_config(
108 &self,
109 ) -> Result<tonic::transport::ServerTlsConfig, Box<dyn std::error::Error>> {
110 let identity = tonic::transport::Identity::from_pem(&self.cert_pem, &self.key_pem);
111 let mut cfg = tonic::transport::ServerTlsConfig::new().identity(identity);
112 if let Some(ca_pem) = &self.client_ca_pem {
113 cfg = cfg.client_ca_root(tonic::transport::Certificate::from_pem(ca_pem));
114 }
115 Ok(cfg)
116 }
117}
118
119impl Default for GrpcServerOptions {
120 fn default() -> Self {
121 Self {
122 bind_addr: "127.0.0.1:5555".to_string(),
123 tls: None,
124 }
125 }
126}
127
128#[derive(Clone)]
129pub struct RedDBGrpcServer {
130 runtime: RedDBRuntime,
131 options: GrpcServerOptions,
132 auth_store: Arc<AuthStore>,
133 oauth_validator: Option<Arc<crate::auth::OAuthValidator>>,
139}
140
141impl RedDBGrpcServer {
142 pub fn new(runtime: RedDBRuntime) -> Self {
143 let auth_config = crate::auth::AuthConfig::default();
144 let auth_store = Arc::new(AuthStore::new(auth_config));
145 Self::with_options(runtime, GrpcServerOptions::default(), auth_store)
146 }
147
148 pub fn from_database_options(
149 db_options: RedDBOptions,
150 options: GrpcServerOptions,
151 ) -> RedDBResult<Self> {
152 let runtime = RedDBRuntime::with_options(db_options.clone())?;
154
155 let auth_store = if db_options.auth.vault_enabled {
156 let pager = runtime.db().store().pager().cloned().ok_or_else(|| {
160 crate::api::RedDBError::Internal(
161 "vault requires a paged database (persistent mode)".into(),
162 )
163 })?;
164 let store = AuthStore::with_vault(db_options.auth.clone(), pager, None)
165 .map_err(|e| crate::api::RedDBError::Internal(e.to_string()))?;
166 Arc::new(store)
167 } else {
168 Arc::new(AuthStore::new(db_options.auth.clone()))
169 };
170 auth_store.bootstrap_from_env();
171 Ok(Self::with_options(runtime, options, auth_store))
172 }
173
174 pub fn with_options(
175 runtime: RedDBRuntime,
176 options: GrpcServerOptions,
177 auth_store: Arc<AuthStore>,
178 ) -> Self {
179 runtime.set_auth_store(Arc::clone(&auth_store));
182 Self {
183 runtime,
184 options,
185 auth_store,
186 oauth_validator: None,
187 }
188 }
189
190 pub fn with_oauth_validator(mut self, validator: Arc<crate::auth::OAuthValidator>) -> Self {
196 self.oauth_validator = Some(validator);
197 self
198 }
199
200 pub fn oauth_validator(&self) -> Option<&Arc<crate::auth::OAuthValidator>> {
202 self.oauth_validator.as_ref()
203 }
204
205 pub fn runtime(&self) -> &RedDBRuntime {
206 &self.runtime
207 }
208
209 pub fn options(&self) -> &GrpcServerOptions {
210 &self.options
211 }
212
213 pub fn auth_store(&self) -> &Arc<AuthStore> {
214 &self.auth_store
215 }
216
217 fn grpc_runtime(&self) -> GrpcRuntime {
218 GrpcRuntime {
219 runtime: self.runtime.clone(),
220 auth_store: self.auth_store.clone(),
221 prepared_registry: PreparedStatementRegistry::new(),
222 oauth_validator: self.oauth_validator.clone(),
223 }
224 }
225
226 pub async fn serve(&self) -> Result<(), Box<dyn std::error::Error>> {
227 let addr = self.options.bind_addr.parse()?;
228 let mut builder = tonic::transport::Server::builder();
229 if let Some(tls) = &self.options.tls {
230 log_grpc_tls_identity(tls);
233 builder = builder.tls_config(tls.to_tonic_config()?)?;
234 }
235 builder
236 .add_service(Self::configured_service(self.grpc_runtime()))
237 .serve(addr)
238 .await?;
239 Ok(())
240 }
241
242 pub async fn serve_on(
243 &self,
244 listener: std::net::TcpListener,
245 ) -> Result<(), Box<dyn std::error::Error>> {
246 listener.set_nonblocking(true)?;
247 let listener = tokio::net::TcpListener::from_std(listener)?;
248 let incoming = TcpListenerStream::new(listener);
249 let mut builder = tonic::transport::Server::builder();
250 if let Some(tls) = &self.options.tls {
251 log_grpc_tls_identity(tls);
252 builder = builder.tls_config(tls.to_tonic_config()?)?;
253 }
254 builder
255 .add_service(Self::configured_service(self.grpc_runtime()))
256 .serve_with_incoming(incoming)
257 .await?;
258 Ok(())
259 }
260
261 pub(crate) async fn serve_router_demux(
268 &self,
269 rx: tokio::sync::mpsc::Receiver<tokio::net::TcpStream>,
270 ) -> Result<(), Box<dyn std::error::Error>> {
271 use tokio_stream::StreamExt;
272 let incoming = tokio_stream::wrappers::ReceiverStream::new(rx).map(Ok::<_, std::io::Error>);
273 let mut builder = tonic::transport::Server::builder();
274 if let Some(tls) = &self.options.tls {
275 log_grpc_tls_identity(tls);
276 builder = builder.tls_config(tls.to_tonic_config()?)?;
277 }
278 builder
279 .add_service(Self::configured_service(self.grpc_runtime()))
280 .serve_with_incoming(incoming)
281 .await?;
282 Ok(())
283 }
284
285 fn configured_service(runtime: GrpcRuntime) -> RedDbServer<GrpcRuntime> {
286 use tonic::codec::CompressionEncoding;
290 RedDbServer::new(runtime)
291 .max_decoding_message_size(256 * 1024 * 1024)
292 .max_encoding_message_size(256 * 1024 * 1024)
293 .accept_compressed(CompressionEncoding::Zstd)
294 .accept_compressed(CompressionEncoding::Gzip)
295 .send_compressed(CompressionEncoding::Zstd)
296 }
297}
298
299struct GrpcPreparedStatement {
301 shape: std::sync::Arc<crate::storage::query::ast::QueryExpr>,
302 parameter_count: usize,
303 created_at: std::time::Instant,
304}
305
306struct PreparedStatementRegistry {
309 map: parking_lot::RwLock<std::collections::HashMap<u64, GrpcPreparedStatement>>,
311 next_id: std::sync::atomic::AtomicU64,
312 get_count: std::sync::atomic::AtomicU64,
313}
314
315impl PreparedStatementRegistry {
316 fn new() -> Arc<Self> {
317 Arc::new(Self {
318 map: parking_lot::RwLock::new(std::collections::HashMap::new()),
319 next_id: std::sync::atomic::AtomicU64::new(1),
320 get_count: std::sync::atomic::AtomicU64::new(0),
321 })
322 }
323
324 fn prepare(&self, shape: crate::storage::query::ast::QueryExpr, parameter_count: usize) -> u64 {
325 use std::sync::atomic::Ordering;
326 let id = self.next_id.fetch_add(1, Ordering::Relaxed);
327 let mut map = self.map.write();
328 self.evict_old_locked(&mut map);
329 map.insert(
330 id,
331 GrpcPreparedStatement {
332 shape: std::sync::Arc::new(shape),
334 parameter_count,
335 created_at: std::time::Instant::now(),
336 },
337 );
338 id
339 }
340
341 fn get_shape_and_count(
342 &self,
343 id: u64,
344 ) -> Option<(std::sync::Arc<crate::storage::query::ast::QueryExpr>, usize)> {
345 let get_count = self
348 .get_count
349 .fetch_add(1, std::sync::atomic::Ordering::Relaxed)
350 + 1;
351 if get_count.is_multiple_of(256) {
352 let mut map = self.map.write();
353 self.evict_old_locked(&mut map);
354 }
355 let map = self.map.read();
356 map.get(&id)
357 .map(|s| (std::sync::Arc::clone(&s.shape), s.parameter_count))
358 }
359
360 fn evict_old_locked(&self, map: &mut std::collections::HashMap<u64, GrpcPreparedStatement>) {
361 let threshold = std::time::Duration::from_secs(3600);
362 map.retain(|_, v| v.created_at.elapsed() < threshold);
363 }
364}
365
366#[derive(Clone)]
367struct GrpcRuntime {
368 runtime: RedDBRuntime,
369 auth_store: Arc<AuthStore>,
370 prepared_registry: Arc<PreparedStatementRegistry>,
371 oauth_validator: Option<Arc<crate::auth::OAuthValidator>>,
375}
376
377impl GrpcRuntime {
378 fn admin_use_cases(&self) -> AdminUseCases<'_, RedDBRuntime> {
379 AdminUseCases::new(&self.runtime)
380 }
381
382 fn catalog_use_cases(&self) -> CatalogUseCases<'_, RedDBRuntime> {
383 CatalogUseCases::new(&self.runtime)
384 }
385
386 fn query_use_cases(&self) -> QueryUseCases<'_, RedDBRuntime> {
387 QueryUseCases::new(&self.runtime)
388 }
389
390 fn entity_use_cases(&self) -> EntityUseCases<'_, RedDBRuntime> {
391 EntityUseCases::new(&self.runtime)
392 }
393
394 fn graph_use_cases(&self) -> GraphUseCases<'_, RedDBRuntime> {
395 GraphUseCases::new(&self.runtime)
396 }
397
398 fn native_use_cases(&self) -> NativeUseCases<'_, RedDBRuntime> {
399 NativeUseCases::new(&self.runtime)
400 }
401}
402
403fn grpc_query_value_to_schema_value(value: QueryValue) -> Result<Value, Status> {
404 use proto::query_value::Kind;
405
406 match value
407 .kind
408 .ok_or_else(|| Status::invalid_argument("missing query param value"))?
409 {
410 Kind::NullValue(_) => Ok(Value::Null),
411 Kind::BoolValue(value) => Ok(Value::Boolean(value)),
412 Kind::IntValue(value) => Ok(Value::Integer(value)),
413 Kind::FloatValue(value) => Ok(Value::Float(value)),
414 Kind::TextValue(value) => Ok(Value::Text(std::sync::Arc::from(value))),
415 Kind::BytesValue(value) => Ok(Value::Blob(value)),
416 Kind::VectorValue(value) => Ok(Value::Vector(value.values)),
417 Kind::JsonValue(value) => {
418 let parsed = json_from_str::<JsonValue>(&value)
419 .map_err(|e| Status::invalid_argument(format!("json param parse error: {e}")))?;
420 let encoded = json_to_string(&parsed)
421 .map_err(|e| Status::invalid_argument(format!("json param encode error: {e}")))?;
422 Ok(Value::Json(encoded.into_bytes()))
423 }
424 Kind::TimestampValue(value) => Ok(Value::Timestamp(value)),
425 Kind::UuidValue(value) => {
426 let bytes: [u8; 16] = value.try_into().map_err(|value: Vec<u8>| {
427 Status::invalid_argument(format!(
428 "uuid param must be 16 bytes, got {}",
429 value.len()
430 ))
431 })?;
432 Ok(Value::Uuid(bytes))
433 }
434 }
435}
436
437fn execute_grpc_query_with_optional_params(
438 runtime: &RedDBRuntime,
439 query: String,
440 params: Vec<QueryValue>,
441) -> Result<RuntimeQueryResult, Status> {
442 if query.trim().is_empty() {
443 return Err(Status::invalid_argument("query field cannot be empty"));
444 }
445
446 if params.is_empty() {
447 let result = runtime.execute_query(&query).map_err(to_status)?;
448 enforce_grpc_commit_policy_after_query_result(runtime, &result)?;
449 return Ok(result);
450 }
451
452 let binds = params
453 .into_iter()
454 .map(grpc_query_value_to_schema_value)
455 .collect::<Result<Vec<_>, _>>()?;
456 let parsed = crate::storage::query::modes::parse_multi(&query)
457 .map_err(|e| Status::invalid_argument(format!("parse error: {e}")))?;
458 let bound = crate::storage::query::user_params::bind(&parsed, &binds)
459 .map_err(|e| Status::invalid_argument(format!("bind error: {e}")))?;
460 let result = runtime.execute_query_expr(bound).map_err(to_status)?;
461 enforce_grpc_commit_policy_after_query_result(runtime, &result)?;
462 Ok(result)
463}
464
465fn enforce_grpc_commit_policy_after_query_result(
466 runtime: &RedDBRuntime,
467 result: &RuntimeQueryResult,
468) -> Result<(), Status> {
469 let is_mutation = matches!(result.statement_type, "insert" | "update" | "delete");
470 if !is_mutation {
471 return Ok(());
472 }
473 let post_lsn = runtime.cdc_current_lsn();
474 runtime
475 .enforce_commit_policy(post_lsn)
476 .map(|_| ())
477 .map_err(|err| Status::deadline_exceeded(err.to_string()))
478}
479
480#[cfg(test)]
481mod grpc_query_value_tests {
482 use super::*;
483 use proto::query_value::Kind;
484
485 #[test]
486 fn grpc_query_value_maps_to_schema_value_variants() {
487 let cases = vec![
488 (
489 QueryValue {
490 kind: Some(Kind::NullValue(proto::QueryNull {})),
491 },
492 Value::Null,
493 ),
494 (
495 QueryValue {
496 kind: Some(Kind::BoolValue(true)),
497 },
498 Value::Boolean(true),
499 ),
500 (
501 QueryValue {
502 kind: Some(Kind::IntValue(42)),
503 },
504 Value::Integer(42),
505 ),
506 (
507 QueryValue {
508 kind: Some(Kind::FloatValue(1.5)),
509 },
510 Value::Float(1.5),
511 ),
512 (
513 QueryValue {
514 kind: Some(Kind::BytesValue(vec![0, 1, 2])),
515 },
516 Value::Blob(vec![0, 1, 2]),
517 ),
518 (
519 QueryValue {
520 kind: Some(Kind::VectorValue(proto::QueryVector {
521 values: vec![0.25, 0.5],
522 })),
523 },
524 Value::Vector(vec![0.25, 0.5]),
525 ),
526 (
527 QueryValue {
528 kind: Some(Kind::TimestampValue(1_779_999_000)),
529 },
530 Value::Timestamp(1_779_999_000),
531 ),
532 (
533 QueryValue {
534 kind: Some(Kind::UuidValue(vec![0x11; 16])),
535 },
536 Value::Uuid([0x11; 16]),
537 ),
538 ];
539
540 for (input, expected) in cases {
541 assert_eq!(grpc_query_value_to_schema_value(input).unwrap(), expected);
542 }
543
544 assert_eq!(
545 grpc_query_value_to_schema_value(QueryValue {
546 kind: Some(Kind::TextValue("alice".into())),
547 })
548 .unwrap(),
549 Value::Text(std::sync::Arc::from("alice"))
550 );
551 assert_eq!(
552 grpc_query_value_to_schema_value(QueryValue {
553 kind: Some(Kind::JsonValue("{\"role\":\"admin\"}".into())),
554 })
555 .unwrap(),
556 Value::Json(b"{\"role\":\"admin\"}".to_vec())
557 );
558 }
559
560 #[test]
561 fn grpc_query_value_rejects_missing_kind_and_bad_uuid() {
562 assert!(grpc_query_value_to_schema_value(QueryValue { kind: None }).is_err());
563 assert!(grpc_query_value_to_schema_value(QueryValue {
564 kind: Some(Kind::UuidValue(vec![0; 15])),
565 })
566 .is_err());
567 }
568
569 #[test]
570 fn grpc_query_rejects_empty_query_before_runtime_parse() {
571 let runtime =
572 RedDBRuntime::with_options(crate::api::RedDBOptions::in_memory()).expect("runtime");
573 let err = execute_grpc_query_with_optional_params(&runtime, " ".to_string(), Vec::new())
574 .expect_err("empty query should fail");
575
576 assert_eq!(err.code(), tonic::Code::InvalidArgument);
577 assert_eq!(err.message(), "query field cannot be empty");
578 }
579
580 #[test]
581 fn grpc_query_params_are_bound_before_execution() {
582 let runtime =
583 RedDBRuntime::with_options(crate::api::RedDBOptions::in_memory()).expect("runtime");
584 seed_grpc_param_table(&runtime);
585
586 let result = execute_grpc_query_with_optional_params(
587 &runtime,
588 "SELECT id, name FROM p WHERE id = $1 AND name = $2".to_string(),
589 grpc_param_values(),
590 )
591 .expect("parameterized query");
592
593 assert_eq!(result.result.records.len(), 1);
594 }
595
596 #[test]
597 fn grpc_query_enforces_ack_n_commit_policy_fail_closed() {
598 let _env_lock = env_lock().lock().expect("env lock");
599 let _env = EnvGuard::set(&[
600 ("RED_PRIMARY_COMMIT_POLICY", "ack_n=1"),
601 ("RED_REPLICATION_ACK_TIMEOUT_MS", "20"),
602 ("RED_COMMIT_FAIL_ON_TIMEOUT", "true"),
603 ]);
604 let data_path = temp_data_path("grpc_ack_n_timeout");
605 cleanup(&data_path);
606
607 let runtime = RedDBRuntime::with_options(
608 crate::api::RedDBOptions::persistent(&data_path)
609 .with_replication(crate::replication::ReplicationConfig::primary()),
610 )
611 .expect("runtime");
612
613 let err = execute_grpc_query_with_optional_params(
614 &runtime,
615 "INSERT INTO grpc_ack_items (id, name) VALUES (1, 'alpha')".to_string(),
616 Vec::new(),
617 )
618 .expect_err("ack_n without replica ack must fail closed");
619
620 assert_eq!(err.code(), tonic::Code::DeadlineExceeded);
621 assert!(
622 err.message().contains("commit policy timed out")
623 && err.message().contains("RED_COMMIT_FAIL_ON_TIMEOUT"),
624 "error should identify commit policy timeout, got {err:?}"
625 );
626 assert!(
627 runtime.cdc_current_lsn() > 0,
628 "local mutation should advance CDC before gRPC response fails"
629 );
630
631 cleanup(&data_path);
632 }
633
634 #[tokio::test]
635 async fn grpc_query_rpc_binds_query_request_params() {
636 let runtime =
637 RedDBRuntime::with_options(crate::api::RedDBOptions::in_memory()).expect("runtime");
638 seed_grpc_param_table(&runtime);
639 let service = GrpcRuntime {
640 runtime,
641 auth_store: Arc::new(AuthStore::new(crate::auth::AuthConfig::default())),
642 prepared_registry: PreparedStatementRegistry::new(),
643 oauth_validator: None,
644 };
645
646 let reply = RedDb::query(
647 &service,
648 Request::new(QueryRequest {
649 query: "SELECT id, name FROM p WHERE id = $1 AND name = $2".to_string(),
650 entity_types: Vec::new(),
651 capabilities: Vec::new(),
652 params: grpc_param_values(),
653 }),
654 )
655 .await
656 .expect("query rpc")
657 .into_inner();
658
659 assert_eq!(reply.record_count, 1);
660 assert!(reply.result_json.contains("Alice"), "{}", reply.result_json);
661 assert!(!reply.result_json.contains("Bob"), "{}", reply.result_json);
662 }
663
664 fn seed_grpc_param_table(runtime: &RedDBRuntime) {
665 runtime
666 .execute_query("CREATE TABLE p (id INTEGER, name TEXT)")
667 .expect("create table");
668 runtime
669 .execute_query("INSERT INTO p (id, name) VALUES (1, 'Alice')")
670 .expect("insert alice");
671 runtime
672 .execute_query("INSERT INTO p (id, name) VALUES (2, 'Bob')")
673 .expect("insert bob");
674 }
675
676 fn grpc_param_values() -> Vec<QueryValue> {
677 vec![
678 QueryValue {
679 kind: Some(Kind::IntValue(1)),
680 },
681 QueryValue {
682 kind: Some(Kind::TextValue("Alice".to_string())),
683 },
684 ]
685 }
686
687 fn env_lock() -> &'static std::sync::Mutex<()> {
688 static LOCK: std::sync::OnceLock<std::sync::Mutex<()>> = std::sync::OnceLock::new();
689 LOCK.get_or_init(|| std::sync::Mutex::new(()))
690 }
691
692 struct EnvGuard {
693 previous: Vec<(&'static str, Option<String>)>,
694 }
695
696 impl EnvGuard {
697 fn set(vars: &[(&'static str, &'static str)]) -> Self {
698 let previous = vars
699 .iter()
700 .map(|(key, _)| (*key, std::env::var(key).ok()))
701 .collect();
702 for (key, value) in vars {
703 std::env::set_var(key, value);
704 }
705 Self { previous }
706 }
707 }
708
709 impl Drop for EnvGuard {
710 fn drop(&mut self) {
711 for (key, value) in self.previous.iter().rev() {
712 match value {
713 Some(value) => std::env::set_var(key, value),
714 None => std::env::remove_var(key),
715 }
716 }
717 }
718 }
719
720 fn temp_data_path(name: &str) -> std::path::PathBuf {
721 let suffix = SystemTime::now()
722 .duration_since(UNIX_EPOCH)
723 .unwrap()
724 .as_nanos();
725 std::env::temp_dir().join(format!("reddb_{name}_{suffix}.rdb"))
726 }
727
728 fn cleanup(data_path: &std::path::Path) {
729 let _ = std::fs::remove_file(data_path);
730 let _ = std::fs::remove_file(
731 crate::replication::primary::PrimaryReplication::slot_path_for(data_path),
732 );
733 let _ = std::fs::remove_file(crate::replication::primary::LogicalWalSpool::path_for(
734 data_path,
735 ));
736 let _ = std::fs::remove_dir_all(
737 crate::replication::primary::PrimaryReplication::primary_replica_root_for(data_path),
738 );
739 reddb_file::cleanup_rebootstrap_artifacts(data_path);
740 }
741}
742
743#[cfg(test)]
744mod grpc_ask_query_reply_tests {
745 use super::*;
746 use crate::storage::query::modes::QueryMode;
747 use crate::storage::query::unified::{UnifiedRecord, UnifiedResult};
748 use crate::storage::schema::Value as SchemaValue;
749
750 fn ask_runtime_result() -> RuntimeQueryResult {
751 let mut result = UnifiedResult::with_columns(vec![
752 "answer".into(),
753 "provider".into(),
754 "model".into(),
755 "mode".into(),
756 "retry_count".into(),
757 "prompt_tokens".into(),
758 "completion_tokens".into(),
759 "sources_flat".into(),
760 "citations".into(),
761 "validation".into(),
762 ]);
763 let mut record = UnifiedRecord::new();
764 record.set("answer", SchemaValue::text("Deploy failed [^1]."));
765 record.set("provider", SchemaValue::text("openai"));
766 record.set("model", SchemaValue::text("gpt-4o-mini"));
767 record.set("mode", SchemaValue::text("strict"));
768 record.set("retry_count", SchemaValue::Integer(0));
769 record.set("prompt_tokens", SchemaValue::Integer(11));
770 record.set("completion_tokens", SchemaValue::Integer(7));
771 record.set(
772 "sources_flat",
773 SchemaValue::Json(
774 br#"[{"urn":"urn:reddb:row:deployments:1","kind":"row","collection":"deployments","id":"1"}]"#.to_vec(),
775 ),
776 );
777 record.set(
778 "citations",
779 SchemaValue::Json(br#"[{"marker":1,"urn":"urn:reddb:row:deployments:1"}]"#.to_vec()),
780 );
781 record.set(
782 "validation",
783 SchemaValue::Json(br#"{"ok":true,"warnings":[],"errors":[]}"#.to_vec()),
784 );
785 result.push(record);
786
787 RuntimeQueryResult {
788 query: "ASK 'why did deploy fail?'".to_string(),
789 mode: QueryMode::Sql,
790 statement: "ask",
791 engine: "runtime-ai",
792 result,
793 affected_rows: 0,
794 statement_type: "select",
795 bookmark: None,
796 }
797 }
798
799 #[test]
800 fn query_reply_ask_result_json_uses_full_canonical_schema() {
801 let reply = query_reply(ask_runtime_result(), &None, &None);
802 let json: crate::json::Value =
803 crate::json::from_str(&reply.result_json).expect("valid ask json");
804
805 assert_eq!(
806 json.get("answer").and_then(crate::json::Value::as_str),
807 Some("Deploy failed [^1].")
808 );
809 assert_eq!(
810 json.get("cache_hit").and_then(crate::json::Value::as_bool),
811 Some(false)
812 );
813 assert_eq!(
814 json.get("cost_usd").and_then(crate::json::Value::as_f64),
815 Some(0.0)
816 );
817 assert_eq!(
818 json.get("mode").and_then(crate::json::Value::as_str),
819 Some("strict")
820 );
821 assert_eq!(
822 json.get("retry_count").and_then(crate::json::Value::as_u64),
823 Some(0)
824 );
825 assert!(
826 json.get("records").is_none(),
827 "ASK must not be row-wrapped: {}",
828 reply.result_json
829 );
830 assert!(
831 json.get("sources_flat")
832 .and_then(crate::json::Value::as_array)
833 .is_some_and(|sources| sources.len() == 1
834 && sources[0]
835 .get("payload")
836 .and_then(crate::json::Value::as_str)
837 .is_some()),
838 "sources_flat must be parsed with payload fallback: {}",
839 reply.result_json
840 );
841 assert!(
842 json.get("citations")
843 .and_then(crate::json::Value::as_array)
844 .is_some_and(|citations| citations.len() == 1),
845 "citations must be parsed: {}",
846 reply.result_json
847 );
848 assert_eq!(
849 json.get("validation")
850 .and_then(|v| v.get("ok"))
851 .and_then(crate::json::Value::as_bool),
852 Some(true)
853 );
854 }
855
856 #[test]
857 fn query_reply_non_ask_answer_column_keeps_row_shape() {
858 let mut result = UnifiedResult::with_columns(vec!["answer".into()]);
859 let mut record = UnifiedRecord::new();
860 record.set("answer", SchemaValue::text("plain select"));
861 result.push(record);
862
863 let reply = query_reply(
864 RuntimeQueryResult {
865 query: "SELECT 'plain select' AS answer".to_string(),
866 mode: QueryMode::Sql,
867 statement: "select",
868 engine: "runtime-sql",
869 result,
870 affected_rows: 0,
871 statement_type: "select",
872 bookmark: None,
873 },
874 &None,
875 &None,
876 );
877 let json: crate::json::Value =
878 crate::json::from_str(&reply.result_json).expect("valid query json");
879
880 assert!(
881 json.get("records").is_some(),
882 "non-ASK must stay row-wrapped"
883 );
884 assert!(
885 json.get("answer").is_none(),
886 "non-ASK must not use ASK envelope"
887 );
888 }
889}
890
891fn log_grpc_tls_identity(tls: &GrpcTlsOptions) {
894 use sha2::{Digest, Sha256};
895 let cert_fp = {
896 let mut h = Sha256::new();
897 h.update(&tls.cert_pem);
898 let digest = h.finalize();
899 let mut buf = String::with_capacity(64);
902 for b in digest.iter() {
903 buf.push_str(&format!("{b:02x}"));
904 }
905 buf
906 };
907 tracing::info!(
908 target: "reddb::security",
909 transport = "grpc",
910 cert_sha256 = %cert_fp,
911 mtls = tls.client_ca_pem.is_some(),
912 "gRPC TLS identity loaded"
913 );
914}
915
916include!("grpc/service_impl.rs");