pub(crate) use crate::application::json_input::{
json_bool_field, json_f32_field, json_string_field, json_usize_field,
};
pub(crate) use crate::application::{
AdminUseCases, CatalogUseCases, CreateEdgeInput, CreateEntityOutput, CreateNodeGraphLinkInput,
CreateNodeInput, CreateNodeTableLinkInput, CreateRowInput, CreateVectorInput,
DeleteEntityInput, EntityUseCases, ExecuteQueryInput, ExplainQueryInput, GraphCentralityInput,
GraphClusteringInput, GraphCommunitiesInput, GraphComponentsInput, GraphCyclesInput,
GraphHitsInput, GraphNeighborhoodInput, GraphPersonalizedPageRankInput, GraphShortestPathInput,
GraphTopologicalSortInput, GraphTraversalInput, GraphUseCases, InspectNativeArtifactInput,
NativeUseCases, PatchEntityInput, QueryUseCases, SearchHybridInput, SearchIvfInput,
SearchSimilarInput, SearchTextInput,
};
use std::collections::BTreeMap;
use std::pin::Pin;
use std::sync::Arc;
use std::time::{SystemTime, UNIX_EPOCH};
use crate::api::{RedDBOptions, RedDBResult};
use crate::auth::middleware::{check_permission, AuthResult};
use crate::auth::store::AuthStore;
use crate::auth::Role;
use crate::health::{HealthProvider, HealthState};
use crate::json::{
from_str as json_from_str, to_string as json_to_string, Map, Value as JsonValue,
};
use crate::runtime::{
RedDBRuntime, RuntimeFilter, RuntimeFilterValue, RuntimeGraphCentralityAlgorithm,
RuntimeGraphCentralityResult, RuntimeGraphClusteringResult, RuntimeGraphCommunityAlgorithm,
RuntimeGraphCommunityResult, RuntimeGraphComponentsMode, RuntimeGraphComponentsResult,
RuntimeGraphCyclesResult, RuntimeGraphDirection, RuntimeGraphHitsResult,
RuntimeGraphNeighborhoodResult, RuntimeGraphPathAlgorithm, RuntimeGraphPathResult,
RuntimeGraphPattern, RuntimeGraphProjection, RuntimeGraphTopologicalSortResult,
RuntimeGraphTraversalResult, RuntimeGraphTraversalStrategy, RuntimeIvfSearchResult,
RuntimeQueryResult, RuntimeQueryWeights, RuntimeStats, ScanPage,
};
use crate::storage::schema::Value;
use crate::storage::unified::devx::refs::{NodeRef, TableRef};
use crate::storage::unified::{Metadata, MetadataValue};
use crate::storage::{EntityData, EntityId, UnifiedEntity};
use tokio_stream::wrappers::TcpListenerStream;
use tonic::metadata::MetadataMap;
use tonic::{Request, Response, Status};
pub use reddb_grpc_proto as proto;
use proto::red_db_server::{RedDb, RedDbServer};
use proto::{
BatchQueryReply, BatchQueryRequest, BulkEntityReply, CollectionRequest, CollectionsReply,
DeleteEntityRequest, DeploymentProfileRequest, Empty, EntityReply, ExecutePreparedRequest,
ExportRequest, GraphProjectionUpsertRequest, HealthReply, IndexNameRequest, IndexToggleRequest,
JsonBulkCreateRequest, JsonCreateRequest, JsonPayloadRequest, KvWatchEvent, KvWatchRequest,
ManifestRequest, OperationReply, PayloadReply, PrepareQueryReply, PrepareQueryRequest,
QueryReply, QueryRequest, ScanEntity, ScanReply, ScanRequest, StatsReply, TopologyReply,
TopologyRequest, UpdateEntityRequest,
};
mod control_support;
mod entity_ops;
mod input_support;
pub(crate) mod scan_json;
use self::control_support::*;
use self::entity_ops::*;
use self::input_support::*;
use self::scan_json::*;
#[derive(Debug, Clone)]
pub struct GrpcServerOptions {
pub bind_addr: String,
pub tls: Option<GrpcTlsOptions>,
}
#[derive(Debug, Clone)]
pub struct GrpcTlsOptions {
pub cert_pem: Vec<u8>,
pub key_pem: Vec<u8>,
pub client_ca_pem: Option<Vec<u8>>,
}
impl GrpcTlsOptions {
pub fn to_tonic_config(
&self,
) -> Result<tonic::transport::ServerTlsConfig, Box<dyn std::error::Error>> {
let identity = tonic::transport::Identity::from_pem(&self.cert_pem, &self.key_pem);
let mut cfg = tonic::transport::ServerTlsConfig::new().identity(identity);
if let Some(ca_pem) = &self.client_ca_pem {
cfg = cfg.client_ca_root(tonic::transport::Certificate::from_pem(ca_pem));
}
Ok(cfg)
}
}
impl Default for GrpcServerOptions {
fn default() -> Self {
Self {
bind_addr: "127.0.0.1:5555".to_string(),
tls: None,
}
}
}
#[derive(Clone)]
pub struct RedDBGrpcServer {
runtime: RedDBRuntime,
options: GrpcServerOptions,
auth_store: Arc<AuthStore>,
oauth_validator: Option<Arc<crate::auth::OAuthValidator>>,
}
impl RedDBGrpcServer {
pub fn new(runtime: RedDBRuntime) -> Self {
let auth_config = crate::auth::AuthConfig::default();
let auth_store = Arc::new(AuthStore::new(auth_config));
Self::with_options(runtime, GrpcServerOptions::default(), auth_store)
}
pub fn from_database_options(
db_options: RedDBOptions,
options: GrpcServerOptions,
) -> RedDBResult<Self> {
let runtime = RedDBRuntime::with_options(db_options.clone())?;
let auth_store = if db_options.auth.vault_enabled {
let pager = runtime.db().store().pager().cloned().ok_or_else(|| {
crate::api::RedDBError::Internal(
"vault requires a paged database (persistent mode)".into(),
)
})?;
let store = AuthStore::with_vault(db_options.auth.clone(), pager, None)
.map_err(|e| crate::api::RedDBError::Internal(e.to_string()))?;
Arc::new(store)
} else {
Arc::new(AuthStore::new(db_options.auth.clone()))
};
auth_store.bootstrap_from_env();
Ok(Self::with_options(runtime, options, auth_store))
}
pub fn with_options(
runtime: RedDBRuntime,
options: GrpcServerOptions,
auth_store: Arc<AuthStore>,
) -> Self {
runtime.set_auth_store(Arc::clone(&auth_store));
Self {
runtime,
options,
auth_store,
oauth_validator: None,
}
}
pub fn with_oauth_validator(mut self, validator: Arc<crate::auth::OAuthValidator>) -> Self {
self.oauth_validator = Some(validator);
self
}
pub fn oauth_validator(&self) -> Option<&Arc<crate::auth::OAuthValidator>> {
self.oauth_validator.as_ref()
}
pub fn runtime(&self) -> &RedDBRuntime {
&self.runtime
}
pub fn options(&self) -> &GrpcServerOptions {
&self.options
}
pub fn auth_store(&self) -> &Arc<AuthStore> {
&self.auth_store
}
fn grpc_runtime(&self) -> GrpcRuntime {
GrpcRuntime {
runtime: self.runtime.clone(),
auth_store: self.auth_store.clone(),
prepared_registry: PreparedStatementRegistry::new(),
oauth_validator: self.oauth_validator.clone(),
}
}
pub async fn serve(&self) -> Result<(), Box<dyn std::error::Error>> {
let addr = self.options.bind_addr.parse()?;
let mut builder = tonic::transport::Server::builder();
if let Some(tls) = &self.options.tls {
log_grpc_tls_identity(tls);
builder = builder.tls_config(tls.to_tonic_config()?)?;
}
builder
.add_service(Self::configured_service(self.grpc_runtime()))
.serve(addr)
.await?;
Ok(())
}
pub async fn serve_on(
&self,
listener: std::net::TcpListener,
) -> Result<(), Box<dyn std::error::Error>> {
listener.set_nonblocking(true)?;
let listener = tokio::net::TcpListener::from_std(listener)?;
let incoming = TcpListenerStream::new(listener);
let mut builder = tonic::transport::Server::builder();
if let Some(tls) = &self.options.tls {
log_grpc_tls_identity(tls);
builder = builder.tls_config(tls.to_tonic_config()?)?;
}
builder
.add_service(Self::configured_service(self.grpc_runtime()))
.serve_with_incoming(incoming)
.await?;
Ok(())
}
fn configured_service(runtime: GrpcRuntime) -> RedDbServer<GrpcRuntime> {
use tonic::codec::CompressionEncoding;
RedDbServer::new(runtime)
.max_decoding_message_size(256 * 1024 * 1024)
.max_encoding_message_size(256 * 1024 * 1024)
.accept_compressed(CompressionEncoding::Zstd)
.accept_compressed(CompressionEncoding::Gzip)
.send_compressed(CompressionEncoding::Zstd)
}
}
struct GrpcPreparedStatement {
shape: std::sync::Arc<crate::storage::query::ast::QueryExpr>,
parameter_count: usize,
created_at: std::time::Instant,
}
struct PreparedStatementRegistry {
map: parking_lot::RwLock<std::collections::HashMap<u64, GrpcPreparedStatement>>,
next_id: std::sync::atomic::AtomicU64,
get_count: std::sync::atomic::AtomicU64,
}
impl PreparedStatementRegistry {
fn new() -> Arc<Self> {
Arc::new(Self {
map: parking_lot::RwLock::new(std::collections::HashMap::new()),
next_id: std::sync::atomic::AtomicU64::new(1),
get_count: std::sync::atomic::AtomicU64::new(0),
})
}
fn prepare(&self, shape: crate::storage::query::ast::QueryExpr, parameter_count: usize) -> u64 {
use std::sync::atomic::Ordering;
let id = self.next_id.fetch_add(1, Ordering::Relaxed);
let mut map = self.map.write();
self.evict_old_locked(&mut map);
map.insert(
id,
GrpcPreparedStatement {
shape: std::sync::Arc::new(shape),
parameter_count,
created_at: std::time::Instant::now(),
},
);
id
}
fn get_shape_and_count(
&self,
id: u64,
) -> Option<(std::sync::Arc<crate::storage::query::ast::QueryExpr>, usize)> {
let get_count = self
.get_count
.fetch_add(1, std::sync::atomic::Ordering::Relaxed)
+ 1;
if get_count.is_multiple_of(256) {
let mut map = self.map.write();
self.evict_old_locked(&mut map);
}
let map = self.map.read();
map.get(&id)
.map(|s| (std::sync::Arc::clone(&s.shape), s.parameter_count))
}
fn evict_old_locked(&self, map: &mut std::collections::HashMap<u64, GrpcPreparedStatement>) {
let threshold = std::time::Duration::from_secs(3600);
map.retain(|_, v| v.created_at.elapsed() < threshold);
}
}
#[derive(Clone)]
struct GrpcRuntime {
runtime: RedDBRuntime,
auth_store: Arc<AuthStore>,
prepared_registry: Arc<PreparedStatementRegistry>,
oauth_validator: Option<Arc<crate::auth::OAuthValidator>>,
}
impl GrpcRuntime {
fn admin_use_cases(&self) -> AdminUseCases<'_, RedDBRuntime> {
AdminUseCases::new(&self.runtime)
}
fn catalog_use_cases(&self) -> CatalogUseCases<'_, RedDBRuntime> {
CatalogUseCases::new(&self.runtime)
}
fn query_use_cases(&self) -> QueryUseCases<'_, RedDBRuntime> {
QueryUseCases::new(&self.runtime)
}
fn entity_use_cases(&self) -> EntityUseCases<'_, RedDBRuntime> {
EntityUseCases::new(&self.runtime)
}
fn graph_use_cases(&self) -> GraphUseCases<'_, RedDBRuntime> {
GraphUseCases::new(&self.runtime)
}
fn native_use_cases(&self) -> NativeUseCases<'_, RedDBRuntime> {
NativeUseCases::new(&self.runtime)
}
}
fn log_grpc_tls_identity(tls: &GrpcTlsOptions) {
use sha2::{Digest, Sha256};
let cert_fp = {
let mut h = Sha256::new();
h.update(&tls.cert_pem);
let digest = h.finalize();
let mut buf = String::with_capacity(64);
for b in digest.iter() {
buf.push_str(&format!("{b:02x}"));
}
buf
};
tracing::info!(
target: "reddb::security",
transport = "grpc",
cert_sha256 = %cert_fp,
mtls = tls.client_ca_pem.is_some(),
"gRPC TLS identity loaded"
);
}
include!("grpc/service_impl.rs");