From cf330041466b288b105d1786fa4fc514299b64b9 Mon Sep 17 00:00:00 2001 From: Travis Gilbert <1travisgilbert@gmail.com> Date: Mon, 8 Jun 2026 03:02:20 -0400 Subject: [PATCH] feat(plugins): add geometry plugin API --- Cargo.lock | 220 +++++ crates/rustyred-core/Cargo.toml | 4 + crates/rustyred-core/src/executor.rs | 31 + crates/rustyred-core/src/fulltext.rs | 30 +- crates/rustyred-core/src/geometry.rs | 899 +++++++++++++++++++++ crates/rustyred-core/src/lib.rs | 13 + crates/rustyred-core/src/plugin.rs | 387 +++++++++ crates/rustyred-core/src/spatial.rs | 29 +- crates/rustyred-mcp/Cargo.toml | 2 + crates/rustyred-mcp/src/lib.rs | 194 +++++ crates/rustyred-server/Cargo.toml | 5 + crates/rustyred-server/src/grpc/service.rs | 335 +++++++- crates/rustyred-server/src/openapi.rs | 181 +++++ crates/rustyred-server/src/router.rs | 377 ++++++++- crates/rustyred-server/src/state.rs | 233 +++++- vendor/proto/rustyred/v1/rustyred.proto | 11 + 16 files changed, 2901 insertions(+), 50 deletions(-) create mode 100644 crates/rustyred-core/src/geometry.rs create mode 100644 crates/rustyred-core/src/plugin.rs diff --git a/Cargo.lock b/Cargo.lock index 8dfa038..004b881 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -561,6 +561,16 @@ dependencies = [ "dtoa", ] +[[package]] +name = "earcutr" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "79127ed59a85d7687c409e9978547cffb7dc79675355ed22da6b66fd5f6ead01" +dependencies = [ + "itertools 0.11.0", + "num-traits", +] + [[package]] name = "ed25519" version = "2.2.3" @@ -661,6 +671,12 @@ dependencies = [ "libc", ] +[[package]] +name = "float_next_after" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8bf7cc16383c4b8d58b9905a8509f02926ce3058053c056376248d958c9df1e8" + [[package]] name = "fnv" version = "1.0.7" @@ -757,6 +773,71 @@ dependencies = [ "version_check", ] +[[package]] +name = "geo" +version = "0.32.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3f3901269ec6d4f6068d3f09e5f02f995bd076398dcd1dfec407cd230b02d11b" +dependencies = [ + "earcutr", + "float_next_after", + "geo-types", + "geographiclib-rs", + "i_overlay", + "log", + "num-traits", + "rand 0.8.6", + "robust", + "rstar", + "sif-itree", + "spade", +] + +[[package]] +name = "geo-traits" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2e7c353d12a704ccfab1ba8bfb1a7fe6cb18b665bf89d37f4f7890edcd260206" +dependencies = [ + "geo-types", +] + +[[package]] +name = "geo-types" +version = "0.7.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "94776032c45f950d30a13af6113c2ad5625316c9abfbccee4dd5a6695f8fe0f5" +dependencies = [ + "approx", + "num-traits", + "rayon", + "rstar", + "serde", +] + +[[package]] +name = "geographiclib-rs" +version = "0.2.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c5a7f08910fd98737a6eda7568e7c5e645093e073328eeef49758cfe8b0489c7" +dependencies = [ + "libm", +] + +[[package]] +name = "geozero" +version = "0.15.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "db5eda63aa99ac06160fd53328ed75c34f14e3196d3f56a3649e247ed796e54b" +dependencies = [ + "geo-types", + "log", + "scroll", + "serde_json", + "thiserror 2.0.18", + "wkt", +] + [[package]] name = "getrandom" version = "0.2.17" @@ -835,6 +916,15 @@ version = "0.1.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3b42eb4efef1f96510ae1a33b2682562a677d504641e9903a77bf5c666b9013e" +[[package]] +name = "hash32" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "47d60b12902ba28e2730cd37e95b8c9223af2808df9e902d4df49588d1470606" +dependencies = [ + "byteorder", +] + [[package]] name = "hashbrown" version = "0.12.3" @@ -852,6 +942,17 @@ dependencies = [ "foldhash 0.1.5", ] +[[package]] +name = "hashbrown" +version = "0.16.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "841d1cc9bed7f9236f321df977030373f4a4163ae1a7dbfe1a51a2c1a51d9100" +dependencies = [ + "allocator-api2", + "equivalent", + "foldhash 0.2.0", +] + [[package]] name = "hashbrown" version = "0.17.1" @@ -863,6 +964,16 @@ dependencies = [ "foldhash 0.2.0", ] +[[package]] +name = "heapless" +version = "0.8.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0bfb9eb618601c89945a70e254898da93b13be0388091d42117462b265bb3fad" +dependencies = [ + "hash32", + "stable_deref_trait", +] + [[package]] name = "heck" version = "0.5.0" @@ -1006,6 +1117,49 @@ dependencies = [ "tracing", ] +[[package]] +name = "i_float" +version = "1.15.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "010025c2c532c8d82e42d0b8bb5184afa449fa6f06c709ea9adcb16c49ae405b" +dependencies = [ + "libm", +] + +[[package]] +name = "i_key_sort" +version = "0.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9190f86706ca38ac8add223b2aed8b1330002b5cdbbce28fb58b10914d38fc27" + +[[package]] +name = "i_overlay" +version = "4.0.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "413183068e6e0289e18d7d0a1f661b81546e6918d5453a44570b9ab30cbed1b3" +dependencies = [ + "i_float", + "i_key_sort", + "i_shape", + "i_tree", + "rayon", +] + +[[package]] +name = "i_shape" +version = "1.14.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ea154b742f7d43dae2897fcd5ead86bc7b5eefcedd305a7ebf9f69d44d61082" +dependencies = [ + "i_float", +] + +[[package]] +name = "i_tree" +version = "0.16.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "35e6d558e6d4c7b82bc51d9c771e7a927862a161a7d87bf2b0541450e0e20915" + [[package]] name = "icu_collections" version = "2.2.0" @@ -1168,6 +1322,15 @@ version = "2.12.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d98f6fed1fde3f8c21bc40a1abb88dd75e67924f9cffc3ef95607bad8017f8e2" +[[package]] +name = "itertools" +version = "0.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b1c173a5686ce8bfa551b3563d0c2170bf24ca44da99c7ca4bfdab5418c3fe57" +dependencies = [ + "either", +] + [[package]] name = "itertools" version = "0.12.1" @@ -2062,6 +2225,23 @@ dependencies = [ "windows-sys 0.52.0", ] +[[package]] +name = "robust" +version = "1.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4e27ee8bb91ca0adcf0ecb116293afa12d393f9c2b9b9cd54d33e8078fe19839" + +[[package]] +name = "rstar" +version = "0.12.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "421400d13ccfd26dfa5858199c30a5d76f9c54e0dba7575273025b43c5175dbb" +dependencies = [ + "heapless", + "num-traits", + "smallvec", +] + [[package]] name = "rust-stemmers" version = "1.2.0" @@ -2180,6 +2360,9 @@ name = "rustyred-core" version = "0.9.1" dependencies = [ "fs2", + "geo", + "geo-types", + "geozero", "h3o", "instant-distance", "redis", @@ -2279,6 +2462,12 @@ version = "1.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49" +[[package]] +name = "scroll" +version = "0.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c1257cd4248b4132760d6524d6dda4e053bc648c9070b960929bf50cfb1e7add" + [[package]] name = "selectors" version = "0.37.0" @@ -2411,6 +2600,12 @@ version = "1.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0fda2ff0d084019ba4d7c6f371c95d8fd75ce3524c3cb8fb653a3023f6323e64" +[[package]] +name = "sif-itree" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d7f45b8998ced5134fb1d75732c77842a3e888f19c1ff98481822e8fbfbf930b" + [[package]] name = "signal-hook-registry" version = "1.4.8" @@ -2477,6 +2672,18 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "spade" +version = "2.15.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9699399fd9349b00b184f5635b074f9ec93afffef30c853f8c875b32c0f8c7fa" +dependencies = [ + "hashbrown 0.16.1", + "num-traits", + "robust", + "smallvec", +] + [[package]] name = "spki" version = "0.7.3" @@ -3487,6 +3694,19 @@ dependencies = [ "wasmparser", ] +[[package]] +name = "wkt" +version = "0.14.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "efb2b923ccc882312e559ffaa832a055ba9d1ac0cc8e86b3e25453247e4b81d7" +dependencies = [ + "geo-traits", + "geo-types", + "log", + "num-traits", + "thiserror 1.0.69", +] + [[package]] name = "writeable" version = "0.6.3" diff --git a/crates/rustyred-core/Cargo.toml b/crates/rustyred-core/Cargo.toml index 2848122..c3f30d8 100644 --- a/crates/rustyred-core/Cargo.toml +++ b/crates/rustyred-core/Cargo.toml @@ -18,9 +18,13 @@ tantivy = ["dep:tantivy"] # §P8-A pa8.1+pa8.2: S2 cells as alternative spatial backend behind # RUSTY_RED_SPATIAL_BACKEND=s2. Default builds keep H3. s2 = ["dep:s2"] +geometry = ["dep:geo", "dep:geo-types", "dep:geozero"] [dependencies] fs2 = "0.4" +geo = { version = "0.32.0", optional = true } +geo-types = { version = "0.7.19", optional = true } +geozero = { version = "0.15.1", default-features = false, features = ["with-geo", "with-wkb", "with-wkt"], optional = true } h3o = "0.7" instant-distance = "0.6.1" redis = { version = "0.27", optional = true } diff --git a/crates/rustyred-core/src/executor.rs b/crates/rustyred-core/src/executor.rs index 02f2bf1..b3d469a 100644 --- a/crates/rustyred-core/src/executor.rs +++ b/crates/rustyred-core/src/executor.rs @@ -523,6 +523,13 @@ impl RustyredExecutor for InMemoryRustyredExecutor { fn execute_request(&mut self, request: RustyredRequest) -> RustyredResponse { let command_name = request.command.clone(); + if let Some(operation) = crate::plugin::builtin_plugin_registry().operation(&command_name) { + let context = crate::plugin::PluginOperationContext { + command: operation.command, + state_hash: self.state_hash(), + }; + return (operation.handler)(context, request.args); + } match RustyredCommand::from_name(&request.command) { Ok(command) => self.execute(command, request.args).unwrap_or_else(|error| { RustyredResponse::err(command_name, error, self.state_hash()) @@ -551,6 +558,17 @@ impl RustyredExecutor for StoreBackedRustyredExecutor { fn execute_request(&mut self, request: RustyredRequest) -> RustyredResponse { let command_name = request.command.clone(); + if let Some(operation) = crate::plugin::builtin_plugin_registry().operation(&command_name) { + let context = crate::plugin::PluginOperationContext { + command: operation.command, + state_hash: self.state_hash(), + }; + let response = (operation.handler)(context, request.args); + if response.ok { + self.persist(); + } + return response; + } match RustyredCommand::from_name(&request.command) { Ok(command) => self.execute(command, request.args).unwrap_or_else(|error| { RustyredResponse::err(command_name, error, self.state_hash()) @@ -907,4 +925,17 @@ mod tests { Some("missing_graph_endpoint") ); } + + #[test] + fn plugin_operation_round_trips_through_json_executor() { + let mut executor = InMemoryRustyredExecutor::new(); + let raw = executor.execute_json( + r#"{"command":"RUSTYRED.PLUGIN.ECHO","args":{"message":"hello plugin"}}"#, + ); + let parsed: serde_json::Value = serde_json::from_str(&raw).unwrap(); + + assert_eq!(parsed["ok"], true); + assert_eq!(parsed["command"], "RUSTYRED.PLUGIN.ECHO"); + assert_eq!(parsed["payload"]["args"]["message"], "hello plugin"); + } } diff --git a/crates/rustyred-core/src/fulltext.rs b/crates/rustyred-core/src/fulltext.rs index b524251..544ff08 100644 --- a/crates/rustyred-core/src/fulltext.rs +++ b/crates/rustyred-core/src/fulltext.rs @@ -347,30 +347,14 @@ pub fn make_fulltext_backend_from_value( designation: FullTextDesignation, raw: &str, ) -> Result, FullTextBackendError> { - match raw.trim().to_ascii_lowercase().as_str() { - "" | FULLTEXT_BACKEND_HAND_ROLLED | "hand-rolled" | "bm25" => { - Ok(Box::new(FullTextIndex::for_designation(designation))) - } - FULLTEXT_BACKEND_TANTIVY => { - #[cfg(feature = "tantivy")] - { - Ok(Box::new( - crate::fulltext_tantivy::TantivyFullTextBackend::new(designation) - .map_err(FullTextBackendError::TantivyInit)?, - )) - } - #[cfg(not(feature = "tantivy"))] - { - let _ = designation; - Err(FullTextBackendError::UnknownBackend( - "tantivy backend requires building with --features tantivy".to_string(), - )) - } - } - other => Err(FullTextBackendError::UnknownBackend(format!( - "unknown {RUSTY_RED_FULLTEXT_BACKEND_ENV} value: {other}" - ))), + let registry = crate::plugin::builtin_plugin_registry(); + if let Some(backend) = registry.fulltext_backend(raw) { + return (backend.constructor)(designation); } + Err(FullTextBackendError::UnknownBackend(format!( + "unknown {RUSTY_RED_FULLTEXT_BACKEND_ENV} value: {}", + raw.trim().to_ascii_lowercase() + ))) } #[derive(Debug, Clone)] diff --git a/crates/rustyred-core/src/geometry.rs b/crates/rustyred-core/src/geometry.rs new file mode 100644 index 0000000..8fabda4 --- /dev/null +++ b/crates/rustyred-core/src/geometry.rs @@ -0,0 +1,899 @@ +use std::collections::{BTreeSet, HashMap}; + +use geo::{BoundingRect, Contains, Intersects, Within}; +use geo_types::{Geometry, Point}; +use geozero::wkb::Wkb; +use geozero::wkt::Wkt; +use geozero::{CoordDimensions, ToGeo, ToWkb, ToWkt}; +use serde::{Deserialize, Serialize}; +use serde_json::{Number, Value}; + +use crate::plugin::{PluginCapability, PluginCapabilityKind, RustyRedPlugin}; +use crate::spatial::SpatialDesignation; + +#[cfg(feature = "s2")] +use s2::cellid::CellID; +#[cfg(feature = "s2")] +use s2::latlng::LatLng as S2LatLng; +#[cfg(feature = "s2")] +use s2::rect::Rect as S2Rect; +#[cfg(feature = "s2")] +use s2::region::RegionCoverer; + +#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)] +#[serde(rename_all = "snake_case")] +pub enum GeometryEncoding { + Point, + Wkb, + Wkt, + Subgraph, +} + +#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)] +pub struct GeometryDesignation { + pub label: String, + pub property: String, + pub encoding: GeometryEncoding, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub lat_property: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub lon_property: Option, + pub resolution: u8, +} + +impl GeometryDesignation { + pub fn point( + label: impl Into, + lat_property: impl Into, + lon_property: impl Into, + resolution: u8, + ) -> Self { + let lat_property = lat_property.into(); + let lon_property = lon_property.into(); + Self { + label: label.into(), + property: format!("{lat_property},{lon_property}"), + encoding: GeometryEncoding::Point, + lat_property: Some(lat_property), + lon_property: Some(lon_property), + resolution, + } + } + + pub fn property( + label: impl Into, + property: impl Into, + encoding: GeometryEncoding, + resolution: u8, + ) -> Self { + Self { + label: label.into(), + property: property.into(), + encoding, + lat_property: None, + lon_property: None, + resolution, + } + } +} + +impl From for GeometryDesignation { + fn from(value: SpatialDesignation) -> Self { + Self::point( + value.label, + value.lat_property, + value.lon_property, + value.resolution, + ) + } +} + +#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)] +pub struct GeometryEncoderDescriptor { + pub name: String, + pub encoding: GeometryEncoding, +} + +pub trait GeometryEncoder: Send + Sync + std::fmt::Debug { + fn name(&self) -> &'static str; + fn encoding(&self) -> GeometryEncoding; + fn encode( + &self, + geom: &Geometry, + props: &mut Value, + designation: &GeometryDesignation, + ) -> Result<(), GeometryError>; + fn decode( + &self, + props: &Value, + designation: &GeometryDesignation, + ) -> Result, GeometryError>; +} + +#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)] +pub enum GeometryError { + MissingProperty(String), + InvalidProperty(String), + UnsupportedEncoding(String), + UnsupportedGeometry(String), + Codec(String), + Cover(String), +} + +impl GeometryError { + pub fn code(&self) -> &'static str { + match self { + Self::MissingProperty(_) => "missing_geometry_property", + Self::InvalidProperty(_) => "invalid_geometry_property", + Self::UnsupportedEncoding(_) => "unsupported_geometry_encoding", + Self::UnsupportedGeometry(_) => "unsupported_geometry", + Self::Codec(_) => "geometry_codec_error", + Self::Cover(_) => "geometry_cover_error", + } + } + + pub fn message(&self) -> String { + match self { + Self::MissingProperty(message) + | Self::InvalidProperty(message) + | Self::UnsupportedEncoding(message) + | Self::UnsupportedGeometry(message) + | Self::Codec(message) + | Self::Cover(message) => message.clone(), + } + } +} + +impl std::fmt::Display for GeometryError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_str(&self.message()) + } +} + +impl std::error::Error for GeometryError {} + +#[derive(Clone, Debug)] +pub struct PointEncoder; + +impl GeometryEncoder for PointEncoder { + fn name(&self) -> &'static str { + "point" + } + + fn encoding(&self) -> GeometryEncoding { + GeometryEncoding::Point + } + + fn encode( + &self, + geom: &Geometry, + props: &mut Value, + designation: &GeometryDesignation, + ) -> Result<(), GeometryError> { + let point = match geom { + Geometry::Point(point) => point, + other => { + return Err(GeometryError::UnsupportedGeometry(format!( + "PointEncoder cannot encode {other:?}" + ))) + } + }; + let lat_property = designation_lat_property(designation)?; + let lon_property = designation_lon_property(designation)?; + let object = props_object_mut(props)?; + object.insert( + lat_property.to_string(), + Number::from_f64(point.y()) + .map(Value::Number) + .ok_or_else(|| GeometryError::InvalidProperty("lat is non-finite".to_string()))?, + ); + object.insert( + lon_property.to_string(), + Number::from_f64(point.x()) + .map(Value::Number) + .ok_or_else(|| GeometryError::InvalidProperty("lon is non-finite".to_string()))?, + ); + Ok(()) + } + + fn decode( + &self, + props: &Value, + designation: &GeometryDesignation, + ) -> Result, GeometryError> { + let lat_property = designation_lat_property(designation)?; + let lon_property = designation_lon_property(designation)?; + let lat = props + .get(lat_property) + .and_then(Value::as_f64) + .ok_or_else(|| GeometryError::MissingProperty(lat_property.to_string()))?; + let lon = props + .get(lon_property) + .and_then(Value::as_f64) + .ok_or_else(|| GeometryError::MissingProperty(lon_property.to_string()))?; + Ok(Geometry::Point(Point::new(lon, lat))) + } +} + +#[derive(Clone, Debug)] +pub struct WkbEncoder; + +impl GeometryEncoder for WkbEncoder { + fn name(&self) -> &'static str { + "wkb" + } + + fn encoding(&self) -> GeometryEncoding { + GeometryEncoding::Wkb + } + + fn encode( + &self, + geom: &Geometry, + props: &mut Value, + designation: &GeometryDesignation, + ) -> Result<(), GeometryError> { + let bytes = geom + .to_wkb(CoordDimensions::xy()) + .map_err(|err| GeometryError::Codec(err.to_string()))?; + let object = props_object_mut(props)?; + object.insert( + designation.property.clone(), + Value::Array( + bytes + .into_iter() + .map(|byte| Value::Number(Number::from(byte))) + .collect(), + ), + ); + Ok(()) + } + + fn decode( + &self, + props: &Value, + designation: &GeometryDesignation, + ) -> Result, GeometryError> { + let bytes = json_bytes( + props + .get(&designation.property) + .ok_or_else(|| GeometryError::MissingProperty(designation.property.clone()))?, + )?; + Wkb(bytes.as_slice()) + .to_geo() + .map_err(|err| GeometryError::Codec(err.to_string())) + } +} + +#[derive(Clone, Debug)] +pub struct WktEncoder; + +impl GeometryEncoder for WktEncoder { + fn name(&self) -> &'static str { + "wkt" + } + + fn encoding(&self) -> GeometryEncoding { + GeometryEncoding::Wkt + } + + fn encode( + &self, + geom: &Geometry, + props: &mut Value, + designation: &GeometryDesignation, + ) -> Result<(), GeometryError> { + let wkt = geom + .to_wkt() + .map_err(|err| GeometryError::Codec(err.to_string()))?; + let object = props_object_mut(props)?; + object.insert(designation.property.clone(), Value::String(wkt)); + Ok(()) + } + + fn decode( + &self, + props: &Value, + designation: &GeometryDesignation, + ) -> Result, GeometryError> { + let wkt = props + .get(&designation.property) + .and_then(Value::as_str) + .ok_or_else(|| GeometryError::MissingProperty(designation.property.clone()))?; + Wkt(wkt) + .to_geo() + .map_err(|err| GeometryError::Codec(err.to_string())) + } +} + +pub fn encoder_for_encoding( + encoding: &GeometryEncoding, +) -> Result, GeometryError> { + match encoding { + GeometryEncoding::Point => Ok(Box::new(PointEncoder)), + GeometryEncoding::Wkb => Ok(Box::new(WkbEncoder)), + GeometryEncoding::Wkt => Ok(Box::new(WktEncoder)), + GeometryEncoding::Subgraph => Err(GeometryError::UnsupportedEncoding( + "subgraph geometry encoding is declared for consumers and is not a core scalar encoder" + .to_string(), + )), + } +} + +pub fn geometry_encoder_descriptors() -> Vec { + vec![ + GeometryEncoderDescriptor { + name: "point".to_string(), + encoding: GeometryEncoding::Point, + }, + GeometryEncoderDescriptor { + name: "wkb".to_string(), + encoding: GeometryEncoding::Wkb, + }, + GeometryEncoderDescriptor { + name: "wkt".to_string(), + encoding: GeometryEncoding::Wkt, + }, + ] +} + +pub fn decode_geometry_value( + encoding: GeometryEncoding, + value: &Value, +) -> Result, GeometryError> { + let designation = GeometryDesignation::property("__query", "geometry", encoding.clone(), 0); + let encoder = encoder_for_encoding(&encoding)?; + let props = serde_json::json!({ "geometry": value.clone() }); + encoder.decode(&props, &designation) +} + +#[derive(Debug)] +pub struct GeometryIndex { + designation: GeometryDesignation, + encoder: Box, + geometries: HashMap>, + #[cfg(feature = "s2")] + cells: HashMap>, + #[cfg(feature = "s2")] + node_to_cells: HashMap>, +} + +impl GeometryIndex { + pub fn new(designation: GeometryDesignation) -> Result { + let encoder = encoder_for_encoding(&designation.encoding)?; + Ok(Self { + designation, + encoder, + geometries: HashMap::new(), + #[cfg(feature = "s2")] + cells: HashMap::new(), + #[cfg(feature = "s2")] + node_to_cells: HashMap::new(), + }) + } + + pub fn designation(&self) -> &GeometryDesignation { + &self.designation + } + + pub fn node_count(&self) -> usize { + self.geometries.len() + } + + pub fn decode_properties(&self, props: &Value) -> Result, GeometryError> { + self.encoder.decode(props, &self.designation) + } + + pub fn encode_geometry( + &self, + geom: &Geometry, + props: &mut Value, + ) -> Result<(), GeometryError> { + self.encoder.encode(geom, props, &self.designation) + } + + pub fn upsert_from_properties( + &mut self, + node_id: &str, + props: &Value, + ) -> Result<(), GeometryError> { + let geom = self.decode_properties(props)?; + self.upsert_geometry(node_id, geom) + } + + pub fn upsert_geometry( + &mut self, + node_id: &str, + geom: Geometry, + ) -> Result<(), GeometryError> { + #[cfg(feature = "s2")] + { + self.remove(node_id); + let cover = cover_geometry(&geom, self.designation.resolution)?; + for cell in &cover { + self.cells + .entry(*cell) + .or_default() + .push(node_id.to_string()); + } + self.node_to_cells + .insert(node_id.to_string(), cover.into_iter().collect()); + self.geometries.insert(node_id.to_string(), geom); + return Ok(()); + } + #[cfg(not(feature = "s2"))] + { + let _ = node_id; + let _ = geom; + Err(GeometryError::Cover( + "geometry indexing requires building with --features s2".to_string(), + )) + } + } + + pub fn remove(&mut self, node_id: &str) { + #[cfg(feature = "s2")] + { + if let Some(old_cells) = self.node_to_cells.remove(node_id) { + for cell in old_cells { + if let Some(nodes) = self.cells.get_mut(&cell) { + nodes.retain(|id| id != node_id); + } + } + } + } + self.geometries.remove(node_id); + } + + pub fn contains_point(&self, lat: f64, lon: f64) -> Result, GeometryError> { + let point = Point::new(lon, lat); + let candidates = self.candidates_for_point(lat, lon)?; + Ok(self.refine_candidates(candidates, |geom| match geom { + Geometry::Polygon(_) | Geometry::MultiPolygon(_) => geom.contains(&point), + _ => false, + })) + } + + pub fn intersects(&self, query: &Geometry) -> Result, GeometryError> { + let candidates = self.candidates_for_geometry(query)?; + Ok(self.refine_candidates(candidates, |geom| geom.intersects(query))) + } + + pub fn intersects_value( + &self, + encoding: GeometryEncoding, + value: &Value, + ) -> Result, GeometryError> { + let query = decode_geometry_value(encoding, value)?; + self.intersects(&query) + } + + pub fn within(&self, query: &Geometry) -> Result, GeometryError> { + let candidates = self.candidates_for_geometry(query)?; + Ok(self.refine_candidates(candidates, |geom| geom.is_within(query))) + } + + pub fn within_value( + &self, + encoding: GeometryEncoding, + value: &Value, + ) -> Result, GeometryError> { + let query = decode_geometry_value(encoding, value)?; + self.within(&query) + } + + fn refine_candidates( + &self, + candidates: BTreeSet, + predicate: impl Fn(&Geometry) -> bool, + ) -> Vec { + let mut out: Vec = candidates + .into_iter() + .filter(|node_id| { + self.geometries + .get(node_id) + .is_some_and(|geom| predicate(geom)) + }) + .collect(); + out.sort(); + out + } + + fn candidates_for_point(&self, lat: f64, lon: f64) -> Result, GeometryError> { + #[cfg(feature = "s2")] + { + let cell = cell_for_lat_lon(lat, lon, self.designation.resolution)?; + return Ok(self.candidates_for_cell(cell)); + } + #[cfg(not(feature = "s2"))] + { + let _ = (lat, lon); + Err(GeometryError::Cover( + "geometry candidates require building with --features s2".to_string(), + )) + } + } + + fn candidates_for_geometry( + &self, + geom: &Geometry, + ) -> Result, GeometryError> { + #[cfg(feature = "s2")] + { + let mut out = BTreeSet::new(); + for cell in cover_geometry(geom, self.designation.resolution)? { + out.extend(self.candidates_for_cell(cell)); + } + return Ok(out); + } + #[cfg(not(feature = "s2"))] + { + let _ = geom; + Err(GeometryError::Cover( + "geometry candidates require building with --features s2".to_string(), + )) + } + } + + #[cfg(feature = "s2")] + fn candidates_for_cell(&self, cell: u64) -> BTreeSet { + let query_cell = CellID(cell); + self.cells + .iter() + .filter(|(stored_cell, _)| CellID(**stored_cell).intersects(&query_cell)) + .flat_map(|(_, nodes)| nodes.iter().cloned()) + .collect() + } + + #[cfg(feature = "s2")] + pub fn cells_for_node(&self, node_id: &str) -> Vec { + self.node_to_cells.get(node_id).cloned().unwrap_or_default() + } + + #[cfg(feature = "s2")] + pub fn cell_contains_node(&self, cell: u64, node_id: &str) -> bool { + self.cells + .get(&cell) + .is_some_and(|nodes| nodes.iter().any(|id| id == node_id)) + } +} + +#[derive(Clone, Debug)] +pub struct GeometryPlugin; + +impl RustyRedPlugin for GeometryPlugin { + fn name(&self) -> &'static str { + "rustyred.geometry" + } + + fn capabilities(&self) -> Vec { + let mut capabilities = vec![ + PluginCapability { + kind: PluginCapabilityKind::Designation, + name: "geometry.designation".to_string(), + }, + PluginCapability { + kind: PluginCapabilityKind::Index, + name: "geometry.s2_cover".to_string(), + }, + PluginCapability { + kind: PluginCapabilityKind::Hook, + name: "geometry.node_upsert_index".to_string(), + }, + ]; + capabilities.extend( + geometry_encoder_descriptors() + .into_iter() + .map(|descriptor| PluginCapability { + kind: PluginCapabilityKind::Encoder, + name: format!("geometry.encoder.{}", descriptor.name), + }), + ); + capabilities + } +} + +fn props_object_mut( + props: &mut Value, +) -> Result<&mut serde_json::Map, GeometryError> { + if props.is_null() { + *props = Value::Object(Default::default()); + } + props + .as_object_mut() + .ok_or_else(|| GeometryError::InvalidProperty("properties must be a JSON object".into())) +} + +fn designation_lat_property(designation: &GeometryDesignation) -> Result<&str, GeometryError> { + designation + .lat_property + .as_deref() + .ok_or_else(|| GeometryError::InvalidProperty("lat_property is required".to_string())) +} + +fn designation_lon_property(designation: &GeometryDesignation) -> Result<&str, GeometryError> { + designation + .lon_property + .as_deref() + .ok_or_else(|| GeometryError::InvalidProperty("lon_property is required".to_string())) +} + +fn json_bytes(value: &Value) -> Result, GeometryError> { + let bytes = value + .as_array() + .ok_or_else(|| GeometryError::InvalidProperty("WKB property must be a byte array".into()))? + .iter() + .map(|value| { + value + .as_u64() + .and_then(|byte| u8::try_from(byte).ok()) + .ok_or_else(|| { + GeometryError::InvalidProperty( + "WKB byte array must contain integers in 0..=255".to_string(), + ) + }) + }) + .collect::, _>>()?; + if bytes.is_empty() { + return Err(GeometryError::InvalidProperty( + "WKB property cannot be empty".to_string(), + )); + } + Ok(bytes) +} + +#[cfg(feature = "s2")] +const MAX_S2_LEVEL: u8 = 30; +#[cfg(feature = "s2")] +const MAX_S2_COVER_CELLS: usize = 1024; + +#[cfg(feature = "s2")] +fn level_from_resolution(resolution: u8) -> u8 { + resolution.saturating_mul(2).clamp(1, MAX_S2_LEVEL) +} + +#[cfg(feature = "s2")] +fn cell_for_lat_lon(lat: f64, lon: f64, resolution: u8) -> Result { + validate_coord(lat, lon)?; + let level = level_from_resolution(resolution); + let cell = CellID::from(S2LatLng::from_degrees(lat, lon)).parent(u64::from(level)); + Ok(cell.0) +} + +#[cfg(feature = "s2")] +fn cover_geometry(geom: &Geometry, resolution: u8) -> Result, GeometryError> { + let level = level_from_resolution(resolution); + match geom { + Geometry::Point(point) => { + let mut cells = BTreeSet::new(); + cells.insert(cell_for_lat_lon(point.y(), point.x(), resolution)?); + Ok(cells) + } + _ => { + let rect = geom.bounding_rect().ok_or_else(|| { + GeometryError::Cover("cannot cover geometry without a bounding rectangle".into()) + })?; + let min_lat = rect.min().y; + let max_lat = rect.max().y; + let min_lon = rect.min().x; + let max_lon = rect.max().x; + validate_coord(min_lat, min_lon)?; + validate_coord(max_lat, max_lon)?; + + let coverer = RegionCoverer { + min_level: 1, + max_level: level, + level_mod: 1, + max_cells: MAX_S2_COVER_CELLS, + }; + let s2_rect = S2Rect::from_degrees(min_lat, min_lon, max_lat, max_lon); + Ok(coverer + .covering(&s2_rect) + .0 + .into_iter() + .map(|cell| cell.0) + .collect()) + } + } +} + +#[cfg(feature = "s2")] +fn validate_coord(lat: f64, lon: f64) -> Result<(), GeometryError> { + if !lat.is_finite() || !lon.is_finite() { + return Err(GeometryError::InvalidProperty(format!( + "non-finite coordinate ({lat}, {lon})" + ))); + } + if !(-90.0..=90.0).contains(&lat) || !(-180.0..=180.0).contains(&lon) { + return Err(GeometryError::InvalidProperty(format!( + "out-of-range coordinate ({lat}, {lon})" + ))); + } + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + use geo_types::{coord, LineString, Polygon}; + use serde_json::json; + + fn square() -> Geometry { + Geometry::Polygon(Polygon::new( + LineString::from(vec![ + coord! { x: 0.0, y: 0.0 }, + coord! { x: 4.0, y: 0.0 }, + coord! { x: 4.0, y: 4.0 }, + coord! { x: 0.0, y: 4.0 }, + coord! { x: 0.0, y: 0.0 }, + ]), + Vec::new(), + )) + } + + #[test] + fn wkb_encoder_round_trips_polygon_and_linestring() { + let encoder = WkbEncoder; + let designation = GeometryDesignation::property("Parcel", "geom", GeometryEncoding::Wkb, 7); + let geometries = vec![ + square(), + Geometry::LineString(LineString::from(vec![ + coord! { x: -1.0, y: 1.0 }, + coord! { x: 5.0, y: 1.0 }, + ])), + ]; + + for geom in geometries { + let mut props = json!({}); + encoder.encode(&geom, &mut props, &designation).unwrap(); + let decoded = encoder.decode(&props, &designation).unwrap(); + assert_eq!(decoded, geom); + } + } + + #[test] + fn wkt_encoder_round_trips_polygon() { + let encoder = WktEncoder; + let designation = GeometryDesignation::property("Parcel", "geom", GeometryEncoding::Wkt, 7); + let mut props = json!({}); + encoder.encode(&square(), &mut props, &designation).unwrap(); + + let decoded = encoder.decode(&props, &designation).unwrap(); + assert_eq!(decoded, square()); + } + + #[test] + fn point_encoder_matches_lat_lon_properties() { + let encoder = PointEncoder; + let designation = GeometryDesignation::point("Place", "lat", "lon", 7); + let mut props = json!({}); + let point = Geometry::Point(Point::new(-83.6875, 43.0125)); + + encoder.encode(&point, &mut props, &designation).unwrap(); + assert_eq!(props["lat"], 43.0125); + assert_eq!(props["lon"], -83.6875); + assert_eq!(encoder.decode(&props, &designation).unwrap(), point); + } + + #[test] + fn geometry_plugin_enumerates_encoders() { + let plugin = GeometryPlugin; + let capabilities = plugin.capabilities(); + + assert!(capabilities + .iter() + .any(|capability| capability.name == "geometry.encoder.point")); + assert!(capabilities + .iter() + .any(|capability| capability.name == "geometry.encoder.wkb")); + assert!(capabilities + .iter() + .any(|capability| capability.name == "geometry.encoder.wkt")); + } + + #[cfg(feature = "s2")] + #[test] + fn geometry_index_stores_and_retrieves_polygon_from_properties() { + let designation = GeometryDesignation::property("Parcel", "geom", GeometryEncoding::Wkb, 7); + let mut index = GeometryIndex::new(designation).unwrap(); + let polygon = square(); + let mut props = json!({}); + index.encode_geometry(&polygon, &mut props).unwrap(); + + index.upsert_from_properties("parcel:1", &props).unwrap(); + + assert_eq!(index.node_count(), 1); + assert_eq!(index.decode_properties(&props).unwrap(), polygon); + assert!(!index.cells_for_node("parcel:1").is_empty()); + } + + #[cfg(feature = "s2")] + #[test] + fn point_geometry_indexes_to_one_cell() { + let designation = GeometryDesignation::point("Place", "lat", "lon", 7); + let mut index = GeometryIndex::new(designation).unwrap(); + index + .upsert_geometry("point:1", Geometry::Point(Point::new(-83.7, 43.0))) + .unwrap(); + + assert_eq!(index.cells_for_node("point:1").len(), 1); + assert_eq!( + index.contains_point(43.0, -83.7).unwrap(), + Vec::::new() + ); + } + + #[cfg(feature = "s2")] + #[test] + fn geometry_index_moves_between_covers_without_stale_cells() { + let designation = GeometryDesignation::property("Parcel", "geom", GeometryEncoding::Wkb, 7); + let mut index = GeometryIndex::new(designation).unwrap(); + let first = square(); + index.upsert_geometry("parcel:1", first).unwrap(); + let old_cells = index.cells_for_node("parcel:1"); + let shifted = Geometry::Polygon(Polygon::new( + LineString::from(vec![ + coord! { x: 20.0, y: 20.0 }, + coord! { x: 24.0, y: 20.0 }, + coord! { x: 24.0, y: 24.0 }, + coord! { x: 20.0, y: 24.0 }, + coord! { x: 20.0, y: 20.0 }, + ]), + Vec::new(), + )); + + index.upsert_geometry("parcel:1", shifted).unwrap(); + + assert_ne!(old_cells, index.cells_for_node("parcel:1")); + assert!(old_cells + .into_iter() + .all(|cell| !index.cell_contains_node(cell, "parcel:1"))); + } + + #[cfg(feature = "s2")] + #[test] + fn topology_predicates_refine_broad_phase_candidates() { + let designation = GeometryDesignation::property("Parcel", "geom", GeometryEncoding::Wkb, 7); + let mut index = GeometryIndex::new(designation).unwrap(); + index.upsert_geometry("parcel:1", square()).unwrap(); + index + .upsert_geometry( + "parcel:2", + Geometry::Polygon(Polygon::new( + LineString::from(vec![ + coord! { x: 10.0, y: 10.0 }, + coord! { x: 14.0, y: 10.0 }, + coord! { x: 14.0, y: 14.0 }, + coord! { x: 10.0, y: 14.0 }, + coord! { x: 10.0, y: 10.0 }, + ]), + Vec::new(), + )), + ) + .unwrap(); + + assert_eq!( + index.contains_point(2.0, 2.0).unwrap(), + vec!["parcel:1".to_string()] + ); + assert!(index.contains_point(8.0, 8.0).unwrap().is_empty()); + + let crossing = Geometry::LineString(LineString::from(vec![ + coord! { x: -1.0, y: 1.0 }, + coord! { x: 5.0, y: 1.0 }, + ])); + assert_eq!(index.intersects(&crossing).unwrap(), vec!["parcel:1"]); + + let envelope = Geometry::Polygon(Polygon::new( + LineString::from(vec![ + coord! { x: -1.0, y: -1.0 }, + coord! { x: 5.0, y: -1.0 }, + coord! { x: 5.0, y: 5.0 }, + coord! { x: -1.0, y: 5.0 }, + coord! { x: -1.0, y: -1.0 }, + ]), + Vec::new(), + )); + assert_eq!(index.within(&envelope).unwrap(), vec!["parcel:1"]); + } +} diff --git a/crates/rustyred-core/src/lib.rs b/crates/rustyred-core/src/lib.rs index 62b6114..3d023f4 100644 --- a/crates/rustyred-core/src/lib.rs +++ b/crates/rustyred-core/src/lib.rs @@ -15,9 +15,12 @@ pub mod executor; pub mod fulltext; #[cfg(feature = "tantivy")] pub mod fulltext_tantivy; +#[cfg(feature = "geometry")] +pub mod geometry; pub mod graph; pub mod graph_store; pub mod instant_kg; +pub mod plugin; pub mod spatial; #[cfg(feature = "s2")] pub mod spatial_s2; @@ -32,6 +35,11 @@ pub use fulltext::{ make_fulltext_backend, make_fulltext_backend_from_value, FullTextBackend, FullTextBackendError, FullTextDesignation, FullTextIndex, RUSTY_RED_FULLTEXT_BACKEND_ENV, }; +#[cfg(feature = "geometry")] +pub use geometry::{ + GeometryDesignation, GeometryEncoder, GeometryEncoding, GeometryError, GeometryIndex, + GeometryPlugin, PointEncoder, WkbEncoder, WktEncoder, +}; #[allow(deprecated)] pub use graph::louvain_communities; pub use graph::{ @@ -55,6 +63,11 @@ pub use instant_kg::{ InstantKgStatus, PprResult, SearchResult, SessionDelta, INSTANT_KG_DEFAULT_ENCODER_VERSION, INSTANT_KG_DEFAULT_INGEST_VERSION, INSTANT_KG_PROTOCOL_VERSION, }; +pub use plugin::{ + builtin_plugin_registry, FullTextBackendRegistration, PluginCapability, PluginCapabilityKind, + PluginOperationContext, PluginOperationRegistration, PluginRegistry, RustyRedPlugin, + SpatialBackendRegistration, +}; pub use spatial::{ make_spatial_backend, make_spatial_backend_from_value, SpatialBackend, SpatialDesignation, SpatialError, SpatialIndex, RUSTY_RED_SPATIAL_BACKEND_ENV, diff --git a/crates/rustyred-core/src/plugin.rs b/crates/rustyred-core/src/plugin.rs new file mode 100644 index 0000000..b6dccff --- /dev/null +++ b/crates/rustyred-core/src/plugin.rs @@ -0,0 +1,387 @@ +use std::collections::BTreeMap; +use std::sync::Arc; + +use serde::{Deserialize, Serialize}; +use serde_json::Value; + +use crate::commands::RustyredResponse; +use crate::fulltext::{ + FullTextBackend, FullTextBackendError, FullTextDesignation, FullTextIndex, + FULLTEXT_BACKEND_HAND_ROLLED, FULLTEXT_BACKEND_TANTIVY, +}; +use crate::spatial::{ + SpatialBackend, SpatialDesignation, SpatialError, SpatialIndex, SPATIAL_BACKEND_H3, + SPATIAL_BACKEND_S2, +}; + +pub type SpatialBackendFactory = + fn(SpatialDesignation) -> Result, SpatialError>; +pub type FullTextBackendFactory = + fn(FullTextDesignation) -> Result, FullTextBackendError>; +pub type PluginOperationHandler = fn(PluginOperationContext<'_>, Value) -> RustyredResponse; + +#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)] +pub enum PluginCapabilityKind { + Designation, + Encoder, + Index, + Operation, + Hook, +} + +#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)] +pub struct PluginCapability { + pub kind: PluginCapabilityKind, + pub name: String, +} + +#[derive(Clone)] +pub struct SpatialBackendRegistration { + pub name: &'static str, + pub aliases: &'static [&'static str], + pub constructor: SpatialBackendFactory, +} + +impl std::fmt::Debug for SpatialBackendRegistration { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("SpatialBackendRegistration") + .field("name", &self.name) + .field("aliases", &self.aliases) + .finish_non_exhaustive() + } +} + +#[derive(Clone)] +pub struct FullTextBackendRegistration { + pub name: &'static str, + pub aliases: &'static [&'static str], + pub constructor: FullTextBackendFactory, +} + +impl std::fmt::Debug for FullTextBackendRegistration { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("FullTextBackendRegistration") + .field("name", &self.name) + .field("aliases", &self.aliases) + .finish_non_exhaustive() + } +} + +#[derive(Clone)] +pub struct PluginOperationRegistration { + pub command: &'static str, + pub summary: &'static str, + pub handler: PluginOperationHandler, +} + +impl std::fmt::Debug for PluginOperationRegistration { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("PluginOperationRegistration") + .field("command", &self.command) + .field("summary", &self.summary) + .finish_non_exhaustive() + } +} + +#[derive(Clone, Debug)] +pub struct PluginOperationContext<'a> { + pub command: &'a str, + pub state_hash: String, +} + +pub trait RustyRedPlugin: Send + Sync + std::fmt::Debug { + fn name(&self) -> &'static str; + + fn capabilities(&self) -> Vec { + Vec::new() + } + + fn spatial_backends(&self) -> Vec { + Vec::new() + } + + fn fulltext_backends(&self) -> Vec { + Vec::new() + } + + fn operations(&self) -> Vec { + Vec::new() + } +} + +#[derive(Clone, Debug, Default)] +pub struct PluginRegistry { + plugins: Vec>, + spatial_backends: BTreeMap, + fulltext_backends: BTreeMap, + operations: BTreeMap, +} + +impl PluginRegistry { + pub fn new() -> Self { + Self::default() + } + + pub fn with_builtin_plugins() -> Self { + let mut registry = Self::new(); + registry.register(CoreSpatialPlugin); + registry.register(CoreFullTextPlugin); + registry.register(CoreOperationsPlugin); + #[cfg(feature = "geometry")] + registry.register(crate::geometry::GeometryPlugin); + registry + } + + pub fn register(&mut self, plugin: impl RustyRedPlugin + 'static) { + let plugin = Arc::new(plugin); + for backend in plugin.spatial_backends() { + self.insert_spatial_backend(backend); + } + for backend in plugin.fulltext_backends() { + self.insert_fulltext_backend(backend); + } + for operation in plugin.operations() { + self.operations + .insert(normalize_command(operation.command), operation); + } + self.plugins.push(plugin); + } + + pub fn plugins(&self) -> Vec<&dyn RustyRedPlugin> { + self.plugins + .iter() + .map(|plugin| plugin.as_ref() as &dyn RustyRedPlugin) + .collect() + } + + pub fn capabilities(&self) -> Vec { + self.plugins + .iter() + .flat_map(|plugin| plugin.capabilities()) + .collect() + } + + pub fn spatial_backend(&self, raw: &str) -> Option<&SpatialBackendRegistration> { + let key = normalize_backend_name(raw); + self.spatial_backends.get(&key) + } + + pub fn fulltext_backend(&self, raw: &str) -> Option<&FullTextBackendRegistration> { + let key = normalize_backend_name(raw); + self.fulltext_backends.get(&key) + } + + pub fn operation(&self, command: &str) -> Option<&PluginOperationRegistration> { + self.operations.get(&normalize_command(command)) + } + + pub fn operations(&self) -> Vec<&PluginOperationRegistration> { + self.operations.values().collect() + } + + fn insert_spatial_backend(&mut self, backend: SpatialBackendRegistration) { + self.spatial_backends + .insert(normalize_backend_name(backend.name), backend.clone()); + for alias in backend.aliases { + self.spatial_backends + .insert(normalize_backend_name(alias), backend.clone()); + } + } + + fn insert_fulltext_backend(&mut self, backend: FullTextBackendRegistration) { + self.fulltext_backends + .insert(normalize_backend_name(backend.name), backend.clone()); + for alias in backend.aliases { + self.fulltext_backends + .insert(normalize_backend_name(alias), backend.clone()); + } + } +} + +pub fn builtin_plugin_registry() -> PluginRegistry { + PluginRegistry::with_builtin_plugins() +} + +fn normalize_backend_name(raw: &str) -> String { + raw.trim().to_ascii_lowercase() +} + +fn normalize_command(command: &str) -> String { + command.trim().to_ascii_uppercase() +} + +#[derive(Clone, Debug)] +pub struct CoreSpatialPlugin; + +impl RustyRedPlugin for CoreSpatialPlugin { + fn name(&self) -> &'static str { + "rustyred.core.spatial" + } + + fn capabilities(&self) -> Vec { + vec![PluginCapability { + kind: PluginCapabilityKind::Index, + name: "spatial".to_string(), + }] + } + + fn spatial_backends(&self) -> Vec { + vec![ + SpatialBackendRegistration { + name: SPATIAL_BACKEND_H3, + aliases: &["", "hand_rolled", "hand-rolled"], + constructor: |designation| Ok(Box::new(SpatialIndex::for_designation(designation))), + }, + SpatialBackendRegistration { + name: SPATIAL_BACKEND_S2, + aliases: &[], + constructor: s2_spatial_backend, + }, + ] + } +} + +#[cfg(feature = "s2")] +fn s2_spatial_backend( + designation: SpatialDesignation, +) -> Result, SpatialError> { + Ok(Box::new(crate::spatial_s2::S2SpatialBackend::new( + designation, + ))) +} + +#[cfg(not(feature = "s2"))] +fn s2_spatial_backend( + _designation: SpatialDesignation, +) -> Result, SpatialError> { + Err(SpatialError::UnknownBackend( + "s2 backend requires building with --features s2".to_string(), + )) +} + +#[derive(Clone, Debug)] +pub struct CoreFullTextPlugin; + +impl RustyRedPlugin for CoreFullTextPlugin { + fn name(&self) -> &'static str { + "rustyred.core.fulltext" + } + + fn capabilities(&self) -> Vec { + vec![PluginCapability { + kind: PluginCapabilityKind::Index, + name: "fulltext".to_string(), + }] + } + + fn fulltext_backends(&self) -> Vec { + vec![ + FullTextBackendRegistration { + name: FULLTEXT_BACKEND_HAND_ROLLED, + aliases: &["", "hand-rolled", "bm25"], + constructor: |designation| { + Ok(Box::new(FullTextIndex::for_designation(designation))) + }, + }, + FullTextBackendRegistration { + name: FULLTEXT_BACKEND_TANTIVY, + aliases: &[], + constructor: tantivy_fulltext_backend, + }, + ] + } +} + +#[cfg(feature = "tantivy")] +fn tantivy_fulltext_backend( + designation: FullTextDesignation, +) -> Result, FullTextBackendError> { + Ok(Box::new( + crate::fulltext_tantivy::TantivyFullTextBackend::new(designation) + .map_err(FullTextBackendError::TantivyInit)?, + )) +} + +#[cfg(not(feature = "tantivy"))] +fn tantivy_fulltext_backend( + _designation: FullTextDesignation, +) -> Result, FullTextBackendError> { + Err(FullTextBackendError::UnknownBackend( + "tantivy backend requires building with --features tantivy".to_string(), + )) +} + +#[derive(Clone, Debug)] +pub struct CoreOperationsPlugin; + +impl RustyRedPlugin for CoreOperationsPlugin { + fn name(&self) -> &'static str { + "rustyred.core.operations" + } + + fn capabilities(&self) -> Vec { + vec![PluginCapability { + kind: PluginCapabilityKind::Operation, + name: "plugin.echo".to_string(), + }] + } + + fn operations(&self) -> Vec { + vec![PluginOperationRegistration { + command: "RUSTYRED.PLUGIN.ECHO", + summary: "Echo a JSON payload through the plugin operation registry.", + handler: |context, args| { + RustyredResponse::ok( + context.command, + "ok", + serde_json::json!({ + "operation": context.command, + "args": args, + }), + context.state_hash, + ) + }, + }] + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[derive(Debug)] + struct NoopPlugin; + + impl RustyRedPlugin for NoopPlugin { + fn name(&self) -> &'static str { + "test.noop" + } + } + + #[test] + fn noop_plugin_registers_and_is_enumerable() { + let mut registry = PluginRegistry::new(); + registry.register(NoopPlugin); + + let names: Vec<&str> = registry + .plugins() + .iter() + .map(|plugin| plugin.name()) + .collect(); + assert_eq!(names, vec!["test.noop"]); + } + + #[test] + fn builtin_registry_enumerates_core_plugins() { + let registry = builtin_plugin_registry(); + let names: Vec<&str> = registry + .plugins() + .iter() + .map(|plugin| plugin.name()) + .collect(); + + assert!(names.contains(&"rustyred.core.spatial")); + assert!(names.contains(&"rustyred.core.fulltext")); + assert!(registry.operation("RUSTYRED.PLUGIN.ECHO").is_some()); + } +} diff --git a/crates/rustyred-core/src/spatial.rs b/crates/rustyred-core/src/spatial.rs index 23ec254..2c44674 100644 --- a/crates/rustyred-core/src/spatial.rs +++ b/crates/rustyred-core/src/spatial.rs @@ -248,29 +248,14 @@ pub fn make_spatial_backend_from_value( designation: SpatialDesignation, raw: &str, ) -> Result, SpatialError> { - match raw.trim().to_ascii_lowercase().as_str() { - "" | SPATIAL_BACKEND_H3 | "hand_rolled" | "hand-rolled" => { - Ok(Box::new(SpatialIndex::for_designation(designation))) - } - SPATIAL_BACKEND_S2 => { - #[cfg(feature = "s2")] - { - Ok(Box::new(crate::spatial_s2::S2SpatialBackend::new( - designation, - ))) - } - #[cfg(not(feature = "s2"))] - { - let _ = designation; - Err(SpatialError::UnknownBackend( - "s2 backend requires building with --features s2".to_string(), - )) - } - } - other => Err(SpatialError::UnknownBackend(format!( - "unknown {RUSTY_RED_SPATIAL_BACKEND_ENV} value: {other}" - ))), + let registry = crate::plugin::builtin_plugin_registry(); + if let Some(backend) = registry.spatial_backend(raw) { + return (backend.constructor)(designation); } + Err(SpatialError::UnknownBackend(format!( + "unknown {RUSTY_RED_SPATIAL_BACKEND_ENV} value: {}", + raw.trim().to_ascii_lowercase() + ))) } #[cfg(test)] diff --git a/crates/rustyred-mcp/Cargo.toml b/crates/rustyred-mcp/Cargo.toml index d404820..74f7f68 100644 --- a/crates/rustyred-mcp/Cargo.toml +++ b/crates/rustyred-mcp/Cargo.toml @@ -10,6 +10,8 @@ publish = false [features] default = [] redis-store = ["rustyred-core/redis-store"] +geometry = ["rustyred-core/geometry"] +s2 = ["rustyred-core/s2"] [dependencies] serde = { version = "1.0", features = ["derive"] } diff --git a/crates/rustyred-mcp/src/lib.rs b/crates/rustyred-mcp/src/lib.rs index b9dfe90..7b4f8ab 100644 --- a/crates/rustyred-mcp/src/lib.rs +++ b/crates/rustyred-mcp/src/lib.rs @@ -169,6 +169,54 @@ pub trait McpGraphBackend { "spatial bbox search is not supported by this MCP backend", )) } + fn designate_geometry_property( + &mut self, + _label: &str, + _property: &str, + _encoding: &str, + _resolution: u8, + ) -> GraphStoreResult<()> { + Err(GraphStoreError::new( + "unsupported_operation", + "geometry designation is not supported by this MCP backend", + )) + } + fn spatial_contains_point( + &self, + _label: &str, + _property: &str, + _lat: f64, + _lon: f64, + ) -> GraphStoreResult> { + Err(GraphStoreError::new( + "unsupported_operation", + "geometry contains search is not supported by this MCP backend", + )) + } + fn spatial_intersects_geometry( + &self, + _label: &str, + _property: &str, + _encoding: &str, + _geometry: &Value, + ) -> GraphStoreResult> { + Err(GraphStoreError::new( + "unsupported_operation", + "geometry intersects search is not supported by this MCP backend", + )) + } + fn spatial_within_geometry( + &self, + _label: &str, + _property: &str, + _encoding: &str, + _geometry: &Value, + ) -> GraphStoreResult> { + Err(GraphStoreError::new( + "unsupported_operation", + "geometry within search is not supported by this MCP backend", + )) + } fn epistemic_neighbors( &self, node_id: &str, @@ -830,6 +878,79 @@ fn call_tool( } }) } + "rustyred.spatial.designate_geometry" | "rustyred.graph.spatial.designate_geometry" => { + if let Some(error) = require_write_tool(config, context, name) { + return Ok(error); + } + let label = required_str(&arguments, "label", name)?; + let property = required_str(&arguments, "property", name)?; + let encoding = arguments + .get("encoding") + .and_then(Value::as_str) + .unwrap_or("wkb"); + let resolution = arguments + .get("resolution") + .and_then(Value::as_u64) + .unwrap_or(9) + .min(u8::MAX as u64) as u8; + backend.designate_geometry_property(label, property, encoding, resolution)?; + json!({ + "tenant": tenant, + "designated": { + "label": label, + "property": property, + "encoding": encoding, + "resolution": resolution + } + }) + } + "rustyred.spatial.contains" | "rustyred.graph.spatial.contains" => { + let label = required_str(&arguments, "label", name)?; + let property = required_str(&arguments, "property", name)?; + let lat = required_f64(&arguments, "lat", name)?; + let lon = required_f64(&arguments, "lon", name)?; + let node_ids = backend.spatial_contains_point(label, property, lat, lon)?; + json!({ + "tenant": tenant, + "node_ids": node_ids, + "stats": { "returned": node_ids.len() } + }) + } + "rustyred.spatial.intersects" | "rustyred.graph.spatial.intersects" => { + let label = required_str(&arguments, "label", name)?; + let property = required_str(&arguments, "property", name)?; + let encoding = arguments + .get("encoding") + .and_then(Value::as_str) + .unwrap_or("wkt"); + let geometry = arguments.get("geometry").ok_or_else(|| { + McpError::invalid_params("rustyred.spatial.intersects requires geometry") + })?; + let node_ids = + backend.spatial_intersects_geometry(label, property, encoding, geometry)?; + json!({ + "tenant": tenant, + "node_ids": node_ids, + "stats": { "returned": node_ids.len() } + }) + } + "rustyred.spatial.within" | "rustyred.graph.spatial.within" => { + let label = required_str(&arguments, "label", name)?; + let property = required_str(&arguments, "property", name)?; + let encoding = arguments + .get("encoding") + .and_then(Value::as_str) + .unwrap_or("wkt"); + let geometry = arguments.get("geometry").ok_or_else(|| { + McpError::invalid_params("rustyred.spatial.within requires geometry") + })?; + let node_ids = backend.spatial_within_geometry(label, property, encoding, geometry)?; + json!({ + "tenant": tenant, + "node_ids": node_ids, + "stats": { "returned": node_ids.len() } + }) + } "rustyred.bulk.nodes" | "rustyred.graph.bulk.nodes" => { if let Some(error) = require_write_tool(config, context, name) { return Ok(error); @@ -2181,6 +2302,51 @@ fn tool_definitions(config: &McpServerConfig) -> Vec { "required": ["label", "lat_property", "lon_property", "min_lat", "min_lon", "max_lat", "max_lon"] }), )); + tools.push(tool( + "rustyred.spatial.contains", + "Search designated geometry properties for polygons containing a point.", + json!({ + "type": "object", + "properties": { + "tenant": { "type": "string" }, + "label": { "type": "string" }, + "property": { "type": "string" }, + "lat": { "type": "number" }, + "lon": { "type": "number" } + }, + "required": ["label", "property", "lat", "lon"] + }), + )); + tools.push(tool( + "rustyred.spatial.intersects", + "Search designated geometry properties that intersect a query geometry.", + json!({ + "type": "object", + "properties": { + "tenant": { "type": "string" }, + "label": { "type": "string" }, + "property": { "type": "string" }, + "encoding": { "type": "string", "enum": ["wkb", "wkt"], "default": "wkt" }, + "geometry": {} + }, + "required": ["label", "property", "geometry"] + }), + )); + tools.push(tool( + "rustyred.spatial.within", + "Search designated geometry properties contained within a query geometry.", + json!({ + "type": "object", + "properties": { + "tenant": { "type": "string" }, + "label": { "type": "string" }, + "property": { "type": "string" }, + "encoding": { "type": "string", "enum": ["wkb", "wkt"], "default": "wkt" }, + "geometry": {} + }, + "required": ["label", "property", "geometry"] + }), + )); tools.push(tool( "rustyred.vector.search", "Run a pure vector similarity search using HNSW indexes. Returns top-k nearest nodes.", @@ -2263,6 +2429,21 @@ fn tool_definitions(config: &McpServerConfig) -> Vec { "required": ["label", "lat_property", "lon_property"] }), )); + tools.push(tool_write( + "rustyred.spatial.designate_geometry", + "Designate a node geometry property for geometry predicate search.", + json!({ + "type": "object", + "properties": { + "tenant": { "type": "string" }, + "label": { "type": "string" }, + "property": { "type": "string" }, + "encoding": { "type": "string", "enum": ["point", "wkb", "wkt", "subgraph"], "default": "wkb" }, + "resolution": { "type": "integer", "default": 9 } + }, + "required": ["label", "property"] + }), + )); tools.push(tool_write( "rustyred.bulk.nodes", "Bulk upsert node records from a JSON array.", @@ -2901,6 +3082,15 @@ mod tests { assert!(tools .iter() .any(|tool| tool["name"] == "rustyred.spatial.radius")); + assert!(tools + .iter() + .any(|tool| tool["name"] == "rustyred.spatial.contains")); + assert!(tools + .iter() + .any(|tool| tool["name"] == "rustyred.spatial.intersects")); + assert!(tools + .iter() + .any(|tool| tool["name"] == "rustyred.spatial.within")); assert!(tools .iter() .any(|tool| tool["name"] == "rustyred.graph.version.compile")); @@ -3272,6 +3462,9 @@ mod tests { assert!(tools .iter() .any(|tool| tool["name"] == "rustyred.spatial.designate")); + assert!(tools + .iter() + .any(|tool| tool["name"] == "rustyred.spatial.designate_geometry")); } #[test] @@ -3282,6 +3475,7 @@ mod tests { for tool_name in [ "rustyred.fulltext.designate", "rustyred.spatial.designate", + "rustyred.spatial.designate_geometry", "rustyred.bulk.nodes", "rustyred.bulk.edges", "rustyred.vector.designate", diff --git a/crates/rustyred-server/Cargo.toml b/crates/rustyred-server/Cargo.toml index ce5fe15..4bb04e7 100644 --- a/crates/rustyred-server/Cargo.toml +++ b/crates/rustyred-server/Cargo.toml @@ -31,6 +31,11 @@ tower-http = { version = "0.6", features = ["cors", "trace"] } tracing = "0.1" tracing-subscriber = { version = "0.3", features = ["env-filter", "fmt"] } +[features] +default = [] +geometry = ["rustyred-core/geometry", "rustyred-mcp/geometry"] +s2 = ["rustyred-core/s2", "rustyred-mcp/s2"] + [build-dependencies] tonic-build = "0.12" diff --git a/crates/rustyred-server/src/grpc/service.rs b/crates/rustyred-server/src/grpc/service.rs index 389ad52..035efa5 100644 --- a/crates/rustyred-server/src/grpc/service.rs +++ b/crates/rustyred-server/src/grpc/service.rs @@ -9,10 +9,14 @@ use std::collections::{BTreeMap, HashMap, HashSet}; use std::time::Instant; use http::{HeaderMap, HeaderValue, StatusCode}; +use rustyred_core::executor::StoreBackedRustyredExecutor; +#[cfg(feature = "geometry")] +use rustyred_core::GeometryEncoding; use rustyred_core::{ - connected_components, label_propagation_communities, pagerank, personalized_pagerank, - Direction, EdgeRecord, EpistemicType, GraphMutation, GraphMutationBatch, GraphStoreError, - NeighborHit, NeighborQuery, NodeQuery, NodeRecord, + builtin_plugin_registry, connected_components, label_propagation_communities, pagerank, + personalized_pagerank, Direction, EdgeRecord, EpistemicType, GraphMutation, GraphMutationBatch, + GraphStoreError, InMemoryRustyredExecutor, NeighborHit, NeighborQuery, NodeQuery, NodeRecord, + RustyredExecutor, RustyredRequest, }; use serde_json::{json, Number, Value}; use tonic::{Request, Response, Status}; @@ -20,6 +24,7 @@ use tonic::{Request, Response, Status}; use super::proto; use super::proto::graph_database_server::GraphDatabase; use crate::auth::require_scope; +use crate::config::StorageMode; use crate::graph_cache::{GraphCacheInvalidateBody, GraphCacheLookupBody, GraphCachePutBody}; use crate::query_surface::{ execute_cypher_query, explain_cypher_query, parse_tx_cypher_mutations, PublicCypherBody, @@ -319,6 +324,193 @@ fn spatial_property_pair(property_name: &str) -> Result<(String, String), Status )) } +fn normalize_plugin_operation(operation: &str) -> String { + operation.trim().to_ascii_lowercase().replace('-', "_") +} + +fn plugin_operation_requires_write(operation: &str) -> bool { + matches!( + normalize_plugin_operation(operation).as_str(), + "designate_geometry" | "spatial_designate_geometry" | "rustyred.spatial.designate_geometry" + ) +} + +#[cfg(feature = "geometry")] +fn required_plugin_arg_str<'a>( + args: &'a Value, + key: &str, + operation: &str, +) -> Result<&'a str, Status> { + args.get(key) + .and_then(Value::as_str) + .filter(|value| !value.trim().is_empty()) + .ok_or_else(|| Status::invalid_argument(format!("{operation} requires string arg {key}"))) +} + +#[cfg(feature = "geometry")] +fn required_plugin_arg_f64(args: &Value, key: &str, operation: &str) -> Result { + args.get(key) + .and_then(Value::as_f64) + .ok_or_else(|| Status::invalid_argument(format!("{operation} requires numeric arg {key}"))) +} + +#[cfg(feature = "geometry")] +fn parse_geometry_encoding_for_grpc(raw: &str) -> Result { + match raw.trim().to_ascii_lowercase().as_str() { + "point" => Ok(GeometryEncoding::Point), + "wkb" => Ok(GeometryEncoding::Wkb), + "wkt" => Ok(GeometryEncoding::Wkt), + "subgraph" => Ok(GeometryEncoding::Subgraph), + other => Err(Status::invalid_argument(format!( + "unsupported geometry encoding: {other}" + ))), + } +} + +#[cfg(feature = "geometry")] +fn execute_geometry_plugin_operation( + state: &AppState, + tenant_id: &str, + operation: &str, + args: &Value, +) -> Result, Status> { + let normalized = normalize_plugin_operation(operation); + match normalized.as_str() { + "designate_geometry" + | "spatial_designate_geometry" + | "rustyred.spatial.designate_geometry" => { + let label = required_plugin_arg_str(args, "label", operation)?; + let property = required_plugin_arg_str(args, "property", operation)?; + let encoding = args + .get("encoding") + .and_then(Value::as_str) + .unwrap_or("wkb"); + let resolution = args + .get("resolution") + .and_then(Value::as_u64) + .unwrap_or(9) + .min(u8::MAX as u64) as u8; + state + .designate_geometry_property( + tenant_id, + label, + property, + parse_geometry_encoding_for_grpc(encoding)?, + resolution, + ) + .map_err(|err| state_status(operation, err))?; + Ok(Some(json!({ + "ok": true, + "tenant": tenant_id, + "operation": normalized, + "designated": { + "label": label, + "property": property, + "encoding": encoding, + "resolution": resolution + } + }))) + } + "contains" | "spatial_contains" | "rustyred.spatial.contains" => { + let label = required_plugin_arg_str(args, "label", operation)?; + let property = required_plugin_arg_str(args, "property", operation)?; + let lat = required_plugin_arg_f64(args, "lat", operation)?; + let lon = required_plugin_arg_f64(args, "lon", operation)?; + let node_ids = state + .geometry_contains_point(tenant_id, label, property, lat, lon) + .map_err(|err| state_status(operation, err))?; + Ok(Some(json!({ + "ok": true, + "tenant": tenant_id, + "operation": normalized, + "count": node_ids.len(), + "node_ids": node_ids + }))) + } + "intersects" | "spatial_intersects" | "rustyred.spatial.intersects" => { + let label = required_plugin_arg_str(args, "label", operation)?; + let property = required_plugin_arg_str(args, "property", operation)?; + let encoding = args + .get("encoding") + .and_then(Value::as_str) + .unwrap_or("wkt"); + let geometry = args.get("geometry").ok_or_else(|| { + Status::invalid_argument(format!("{operation} requires geometry")) + })?; + let node_ids = state + .geometry_intersects( + tenant_id, + label, + property, + parse_geometry_encoding_for_grpc(encoding)?, + geometry, + ) + .map_err(|err| state_status(operation, err))?; + Ok(Some(json!({ + "ok": true, + "tenant": tenant_id, + "operation": normalized, + "count": node_ids.len(), + "node_ids": node_ids + }))) + } + "within" | "spatial_within" | "rustyred.spatial.within" => { + let label = required_plugin_arg_str(args, "label", operation)?; + let property = required_plugin_arg_str(args, "property", operation)?; + let encoding = args + .get("encoding") + .and_then(Value::as_str) + .unwrap_or("wkt"); + let geometry = args.get("geometry").ok_or_else(|| { + Status::invalid_argument(format!("{operation} requires geometry")) + })?; + let node_ids = state + .geometry_within( + tenant_id, + label, + property, + parse_geometry_encoding_for_grpc(encoding)?, + geometry, + ) + .map_err(|err| state_status(operation, err))?; + Ok(Some(json!({ + "ok": true, + "tenant": tenant_id, + "operation": normalized, + "count": node_ids.len(), + "node_ids": node_ids + }))) + } + _ => Ok(None), + } +} + +#[cfg(not(feature = "geometry"))] +fn execute_geometry_plugin_operation( + _state: &AppState, + _tenant_id: &str, + operation: &str, + _args: &Value, +) -> Result, Status> { + match normalize_plugin_operation(operation).as_str() { + "designate_geometry" + | "spatial_designate_geometry" + | "rustyred.spatial.designate_geometry" + | "contains" + | "spatial_contains" + | "rustyred.spatial.contains" + | "intersects" + | "spatial_intersects" + | "rustyred.spatial.intersects" + | "within" + | "spatial_within" + | "rustyred.spatial.within" => Err(Status::unimplemented( + "geometry plugin operations require building rustyred-server with --features geometry,s2", + )), + _ => Ok(None), + } +} + fn epistemic_types_from_proto(values: Vec) -> Result>, Status> { let mut types = Vec::new(); for value in values { @@ -691,6 +883,9 @@ impl GraphDatabase for GraphDatabaseService { .maybe_index_node_fulltext(&request.tenant_id, &node); self.state .maybe_index_node_spatially(&request.tenant_id, &node); + #[cfg(feature = "geometry")] + self.state + .maybe_index_node_geometry(&request.tenant_id, &node); } else { self.state .append_graph_transaction_mutations( @@ -890,6 +1085,9 @@ impl GraphDatabase for GraphDatabaseService { .maybe_index_node_fulltext(&request.tenant_id, node); self.state .maybe_index_node_spatially(&request.tenant_id, node); + #[cfg(feature = "geometry")] + self.state + .maybe_index_node_geometry(&request.tenant_id, node); } transaction.graph_version } @@ -1243,6 +1441,67 @@ impl GraphDatabase for GraphDatabaseService { })) } + async fn execute_plugin_operation( + &self, + request: Request, + ) -> Result, Status> { + let scope = if plugin_operation_requires_write(&request.get_ref().operation) { + "graph:write" + } else { + "graph:read" + }; + self.require_request_scope(&request, scope)?; + let request = request.into_inner(); + let args = if request.args_json.trim().is_empty() { + json!({}) + } else { + serde_json::from_str::(&request.args_json).map_err(|err| { + Status::invalid_argument(format!("args_json is invalid JSON: {err}")) + })? + }; + + if let Some(result) = execute_geometry_plugin_operation( + &self.state, + &request.tenant_id, + &request.operation, + &args, + )? { + return Ok(Response::new(proto::PluginOperationResponse { + result_json: result.to_string(), + })); + } + + if builtin_plugin_registry() + .operation(&request.operation) + .is_none() + { + return Err(Status::invalid_argument(format!( + "unknown plugin operation: {}", + request.operation + ))); + } + + let response = match self.state.tenant_store(&request.tenant_id) { + Ok(store) => { + let mut executor = StoreBackedRustyredExecutor::new(store); + executor.execute_request(RustyredRequest::new(request.operation, args)) + } + Err(_error) if self.state.config.storage_mode != StorageMode::Redis => { + let mut executor = InMemoryRustyredExecutor::new(); + executor.execute_request(RustyredRequest::new(request.operation, args)) + } + Err(error) => return Err(state_status("execute_plugin_operation", error)), + }; + let result_json = serde_json::to_string(&response).map_err(|err| { + Status::internal(format!( + "plugin operation response serialization failed: {err}" + )) + })?; + Ok(Response::new(proto::PluginOperationResponse { + result_json, + })) + } + // ==================================================================== // Epistemic traversal // ==================================================================== @@ -1771,6 +2030,76 @@ mod tests { assert_eq!(neighbors.neighbors[0].score, 0.82); } + #[cfg(feature = "geometry")] + #[tokio::test] + async fn grpc_plugin_operation_executes_geometry_and_registry_operations() { + let service = service_with_graph(); + + let designated = service + .execute_plugin_operation(Request::new(proto::PluginOperationRequest { + tenant_id: "smoke".to_string(), + operation: "rustyred.spatial.designate_geometry".to_string(), + args_json: json!({ + "label": "Parcel", + "property": "geom", + "encoding": "wkt", + "resolution": 9 + }) + .to_string(), + })) + .await + .unwrap() + .into_inner(); + let designated_payload: Value = serde_json::from_str(&designated.result_json).unwrap(); + assert_eq!(designated_payload["ok"], true); + + service + .upsert_node(Request::new(proto::UpsertNodeRequest { + tenant_id: "smoke".to_string(), + node: Some(proto::Node { + id: "parcel:grpc".to_string(), + labels: vec!["Parcel".to_string()], + properties: Some(properties_to_proto(&json!({ + "geom": "POLYGON((0 0,4 0,4 4,0 4,0 0))" + }))), + }), + transaction_id: String::new(), + })) + .await + .unwrap(); + + let contains = service + .execute_plugin_operation(Request::new(proto::PluginOperationRequest { + tenant_id: "smoke".to_string(), + operation: "rustyred.spatial.contains".to_string(), + args_json: json!({ + "label": "Parcel", + "property": "geom", + "lat": 2.0, + "lon": 2.0 + }) + .to_string(), + })) + .await + .unwrap() + .into_inner(); + let contains_payload: Value = serde_json::from_str(&contains.result_json).unwrap(); + assert_eq!(contains_payload["node_ids"], json!(["parcel:grpc"])); + + let echo = service + .execute_plugin_operation(Request::new(proto::PluginOperationRequest { + tenant_id: "smoke".to_string(), + operation: "RUSTYRED.PLUGIN.ECHO".to_string(), + args_json: json!({ "ping": "pong" }).to_string(), + })) + .await + .unwrap() + .into_inner(); + let echo_payload: Value = serde_json::from_str(&echo.result_json).unwrap(); + assert_eq!(echo_payload["ok"], true); + assert_eq!(echo_payload["payload"]["args"], json!({ "ping": "pong" })); + } + #[tokio::test] async fn grpc_write_query_diagnostics_and_cache_surfaces_work() { let service = service_with_graph(); diff --git a/crates/rustyred-server/src/openapi.rs b/crates/rustyred-server/src/openapi.rs index abd077f..184a440 100644 --- a/crates/rustyred-server/src/openapi.rs +++ b/crates/rustyred-server/src/openapi.rs @@ -1469,6 +1469,98 @@ pub async fn openapi(State(state): State) -> Json { } } }, + "/v1/tenants/{tenant_id}/graph/spatial/designate_geometry": { + "post": { + "tags": ["graph"], + "summary": "Designate a geometry index", + "description": "Registers a node label and property holding geometry (point, WKB, or WKT) for the S2-cover geometry index. Served by the spatial plugin-operation route. Upserts of matching nodes refresh the cover.", + "parameters": [tenant_parameter.clone()], + "requestBody": { + "required": true, + "content": { + "application/json": { + "schema": { "$ref": "#/components/schemas/SpatialGeometryDesignateRequest" } + } + } + }, + "responses": { + "200": { "$ref": "#/components/responses/GeometryDesignationResponse" }, + "400": { "$ref": "#/components/responses/GraphStoreError" }, + "401": { "$ref": "#/components/responses/Unauthorized" }, + "403": { "$ref": "#/components/responses/Forbidden" }, + "503": { "$ref": "#/components/responses/StoreUnavailable" } + } + } + }, + "/v1/tenants/{tenant_id}/graph/spatial/within": { + "post": { + "tags": ["graph"], + "summary": "Run a geometry within query", + "description": "Returns node ids whose designated geometry lies within the query geometry. Two-phase: S2 cover broad phase, then an exact geo predicate narrow phase.", + "parameters": [tenant_parameter.clone()], + "requestBody": { + "required": true, + "content": { + "application/json": { + "schema": { "$ref": "#/components/schemas/SpatialGeometryQueryRequest" } + } + } + }, + "responses": { + "200": { "$ref": "#/components/responses/GeometryIdsResponse" }, + "400": { "$ref": "#/components/responses/GraphStoreError" }, + "401": { "$ref": "#/components/responses/Unauthorized" }, + "403": { "$ref": "#/components/responses/Forbidden" }, + "503": { "$ref": "#/components/responses/StoreUnavailable" } + } + } + }, + "/v1/tenants/{tenant_id}/graph/spatial/intersects": { + "post": { + "tags": ["graph"], + "summary": "Run a geometry intersects query", + "description": "Returns node ids whose designated geometry intersects the query geometry. Two-phase: S2 cover broad phase, then an exact geo predicate narrow phase.", + "parameters": [tenant_parameter.clone()], + "requestBody": { + "required": true, + "content": { + "application/json": { + "schema": { "$ref": "#/components/schemas/SpatialGeometryQueryRequest" } + } + } + }, + "responses": { + "200": { "$ref": "#/components/responses/GeometryIdsResponse" }, + "400": { "$ref": "#/components/responses/GraphStoreError" }, + "401": { "$ref": "#/components/responses/Unauthorized" }, + "403": { "$ref": "#/components/responses/Forbidden" }, + "503": { "$ref": "#/components/responses/StoreUnavailable" } + } + } + }, + "/v1/tenants/{tenant_id}/graph/spatial/contains": { + "post": { + "tags": ["graph"], + "summary": "Run a geometry contains-point query", + "description": "Returns node ids whose designated polygon or multipolygon geometry contains the given lat/lon point. Two-phase: S2 cell lookup broad phase, then an exact geo contains narrow phase.", + "parameters": [tenant_parameter.clone()], + "requestBody": { + "required": true, + "content": { + "application/json": { + "schema": { "$ref": "#/components/schemas/SpatialContainsRequest" } + } + } + }, + "responses": { + "200": { "$ref": "#/components/responses/GeometryIdsResponse" }, + "400": { "$ref": "#/components/responses/GraphStoreError" }, + "401": { "$ref": "#/components/responses/Unauthorized" }, + "403": { "$ref": "#/components/responses/Forbidden" }, + "503": { "$ref": "#/components/responses/StoreUnavailable" } + } + } + }, "/v1/tenants/{tenant_id}/graph/fulltext/designate": { "post": { "tags": ["graph"], @@ -1837,6 +1929,22 @@ pub async fn openapi(State(state): State) -> Json { } } }, + "GeometryIdsResponse": { + "description": "Geometry topology query result (within, intersects, contains).", + "content": { + "application/json": { + "schema": { "$ref": "#/components/schemas/GeometryIdsResponseBody" } + } + } + }, + "GeometryDesignationResponse": { + "description": "Geometry index designation acknowledgement.", + "content": { + "application/json": { + "schema": { "$ref": "#/components/schemas/GeometryDesignationResponseBody" } + } + } + }, "BulkIngestResponse": { "description": "Bulk ingest report.", "content": { @@ -2725,6 +2833,72 @@ pub async fn openapi(State(state): State) -> Json { }, "additionalProperties": false }, + "SpatialGeometryDesignateRequest": { + "type": "object", + "required": ["label", "property"], + "properties": { + "label": { "type": "string" }, + "property": { "type": "string" }, + "encoding": { "type": "string", "enum": ["point", "wkb", "wkt", "subgraph"], "default": "wkb" }, + "resolution": { "type": "integer", "minimum": 0, "maximum": 15, "default": 8 } + }, + "additionalProperties": false + }, + "SpatialContainsRequest": { + "type": "object", + "required": ["label", "property", "lat", "lon"], + "properties": { + "label": { "type": "string" }, + "property": { "type": "string" }, + "lat": { "type": "number" }, + "lon": { "type": "number" } + }, + "additionalProperties": false + }, + "SpatialGeometryQueryRequest": { + "type": "object", + "required": ["label", "property", "geometry"], + "properties": { + "label": { "type": "string" }, + "property": { "type": "string" }, + "geometry": { "description": "Query geometry whose concrete shape depends on encoding: a WKT string, a WKB byte array, or a point value." }, + "encoding": { "type": "string", "enum": ["point", "wkb", "wkt", "subgraph"], "default": "wkt" } + }, + "additionalProperties": false + }, + "GeometryIdsResponseBody": { + "type": "object", + "required": ["ok", "tenant", "operation", "count", "node_ids"], + "properties": { + "ok": { "type": "boolean" }, + "tenant": { "type": "string" }, + "operation": { "type": "string" }, + "count": { "type": "integer", "minimum": 0 }, + "node_ids": { "type": "array", "items": { "type": "string" } } + }, + "additionalProperties": false + }, + "GeometryDesignationResponseBody": { + "type": "object", + "required": ["ok", "tenant", "operation", "designated"], + "properties": { + "ok": { "type": "boolean" }, + "tenant": { "type": "string" }, + "operation": { "type": "string" }, + "designated": { + "type": "object", + "required": ["label", "property", "encoding", "resolution"], + "properties": { + "label": { "type": "string" }, + "property": { "type": "string" }, + "encoding": { "type": "string" }, + "resolution": { "type": "integer" } + }, + "additionalProperties": false + } + }, + "additionalProperties": false + }, "FullTextDesignateRequest": { "type": "object", "required": ["label", "property"], @@ -3456,6 +3630,13 @@ mod tests { ("/v1/tenants/{tenant_id}/graph/spatial/designate", "post"), ("/v1/tenants/{tenant_id}/graph/spatial/radius", "post"), ("/v1/tenants/{tenant_id}/graph/spatial/bbox", "post"), + ( + "/v1/tenants/{tenant_id}/graph/spatial/designate_geometry", + "post", + ), + ("/v1/tenants/{tenant_id}/graph/spatial/within", "post"), + ("/v1/tenants/{tenant_id}/graph/spatial/intersects", "post"), + ("/v1/tenants/{tenant_id}/graph/spatial/contains", "post"), ("/v1/tenants/{tenant_id}/graph/fulltext/designate", "post"), ("/v1/tenants/{tenant_id}/graph/fulltext/search", "post"), ("/v1/tenants/{tenant_id}/graph/bulk/nodes", "post"), diff --git a/crates/rustyred-server/src/router.rs b/crates/rustyred-server/src/router.rs index f241b2c..1935d8f 100644 --- a/crates/rustyred-server/src/router.rs +++ b/crates/rustyred-server/src/router.rs @@ -423,6 +423,10 @@ pub fn build_router(state: AppState) -> Router { "/v1/tenants/:tenant_id/graph/spatial/bbox", post(graph_spatial_bbox), ) + .route( + "/v1/tenants/:tenant_id/graph/spatial/:operation", + post(graph_spatial_plugin_operation), + ) .route( "/v1/tenants/:tenant_id/graph/fulltext/designate", post(graph_fulltext_designate), @@ -848,6 +852,8 @@ fn index_committed_batch( if let rustyred_core::GraphMutation::NodeUpsert(node) = mutation { state.maybe_index_node_spatially(tenant_id, node); state.maybe_index_node_fulltext(tenant_id, node); + #[cfg(feature = "geometry")] + state.maybe_index_node_geometry(tenant_id, node); } } } @@ -2169,6 +2175,8 @@ async fn graph_node_upsert( state.observability.record_mutation(); state.maybe_index_node_spatially(&tenant_id, &index_clone); state.maybe_index_node_fulltext(&tenant_id, &index_clone); + #[cfg(feature = "geometry")] + state.maybe_index_node_geometry(&tenant_id, &index_clone); Json(json!({ "ok": true, "node": result })).into_response() } Err(error) => { @@ -2639,7 +2647,27 @@ fn execute_tenant_graph_command( ) } }; - execute_graph_store_command(&mut store, command_name, args) + let node_for_hooks = matches!( + RustyredCommand::from_name(command_name), + Ok(RustyredCommand::GraphNodeUpsert) + ) + .then(|| { + serde_json::from_value::(args.clone()) + .ok() + .map(NodeWriteBody::into_record) + }) + .flatten(); + let response = execute_graph_store_command(&mut store, command_name, args); + if response.ok { + if let Some(node) = node_for_hooks { + state.observability.record_mutation(); + state.maybe_index_node_spatially(tenant_id, &node); + state.maybe_index_node_fulltext(tenant_id, &node); + #[cfg(feature = "geometry")] + state.maybe_index_node_geometry(tenant_id, &node); + } + } + response } fn execute_tenant_cache_command( @@ -3754,6 +3782,8 @@ fn flush_node_batch( state.observability.record_mutation(); state.maybe_index_node_spatially(tenant_id, node); state.maybe_index_node_fulltext(tenant_id, node); + #[cfg(feature = "geometry")] + state.maybe_index_node_geometry(tenant_id, node); } } Err(_) => { @@ -3768,6 +3798,8 @@ fn flush_node_batch( state.observability.record_mutation(); state.maybe_index_node_spatially(tenant_id, &node); state.maybe_index_node_fulltext(tenant_id, &node); + #[cfg(feature = "geometry")] + state.maybe_index_node_geometry(tenant_id, &node); } Err(err) => { *failed += 1; @@ -4256,6 +4288,245 @@ async fn graph_spatial_bbox( } } +#[cfg(feature = "geometry")] +#[derive(Debug, Deserialize)] +struct SpatialGeometryDesignateBody { + label: String, + property: String, + #[serde(default = "default_geometry_encoding")] + encoding: String, + #[serde(default = "default_h3_resolution")] + resolution: u8, +} + +#[cfg(feature = "geometry")] +#[derive(Debug, Deserialize)] +struct SpatialContainsBody { + label: String, + property: String, + lat: f64, + lon: f64, +} + +#[cfg(feature = "geometry")] +#[derive(Debug, Deserialize)] +struct SpatialGeometryQueryBody { + label: String, + property: String, + geometry: Value, + #[serde(default = "default_query_geometry_encoding")] + encoding: String, +} + +#[cfg(feature = "geometry")] +fn default_geometry_encoding() -> String { + "wkb".to_string() +} + +#[cfg(feature = "geometry")] +fn default_query_geometry_encoding() -> String { + "wkt".to_string() +} + +async fn graph_spatial_plugin_operation( + State(state): State, + Path((tenant_id, operation)): Path<(String, String)>, + headers: HeaderMap, + Json(payload): Json, +) -> impl IntoResponse { + #[cfg(feature = "geometry")] + { + return graph_spatial_plugin_operation_enabled( + state, tenant_id, operation, headers, payload, + ); + } + #[cfg(not(feature = "geometry"))] + { + let _ = (state, tenant_id, operation, headers, payload); + ( + StatusCode::NOT_IMPLEMENTED, + Json(json!({ + "error": "geometry_feature_disabled", + "message": "geometry operations require building rustyred-server with --features geometry,s2" + })), + ) + .into_response() + } +} + +#[cfg(feature = "geometry")] +fn graph_spatial_plugin_operation_enabled( + state: AppState, + tenant_id: String, + operation: String, + headers: HeaderMap, + payload: Value, +) -> axum::response::Response { + let operation = operation.trim().to_ascii_lowercase().replace('-', "_"); + let scope = if matches!( + operation.as_str(), + "designate_geometry" | "spatial_designate_geometry" + ) { + "graph:write" + } else { + "graph:read" + }; + if let Err(status) = require_scope( + &headers, + &state.config.api_tokens, + scope, + state.config.require_auth, + ) { + return status.into_response(); + } + state.observability.record_spatial_search(); + match operation.as_str() { + "designate_geometry" | "spatial_designate_geometry" => { + let body = match serde_json::from_value::(payload) { + Ok(body) => body, + Err(error) => return invalid_body_response(error), + }; + let encoding = match parse_geometry_encoding(&body.encoding) { + Ok(encoding) => encoding, + Err(response) => return response, + }; + match state.designate_geometry_property( + &tenant_id, + &body.label, + &body.property, + encoding, + body.resolution, + ) { + Ok(()) => Json(json!({ + "ok": true, + "tenant": tenant_id, + "operation": operation, + "designated": { + "label": body.label, + "property": body.property, + "encoding": body.encoding, + "resolution": body.resolution, + } + })) + .into_response(), + Err(error) => store_unavailable_response(error), + } + } + "contains" | "spatial_contains" => { + let body = match serde_json::from_value::(payload) { + Ok(body) => body, + Err(error) => return invalid_body_response(error), + }; + match state.geometry_contains_point( + &tenant_id, + &body.label, + &body.property, + body.lat, + body.lon, + ) { + Ok(ids) => geometry_ids_response(&tenant_id, &operation, ids), + Err(error) => store_unavailable_response(error), + } + } + "intersects" | "spatial_intersects" => { + let body = match serde_json::from_value::(payload) { + Ok(body) => body, + Err(error) => return invalid_body_response(error), + }; + let encoding = match parse_geometry_encoding(&body.encoding) { + Ok(encoding) => encoding, + Err(response) => return response, + }; + match state.geometry_intersects( + &tenant_id, + &body.label, + &body.property, + encoding, + &body.geometry, + ) { + Ok(ids) => geometry_ids_response(&tenant_id, &operation, ids), + Err(error) => store_unavailable_response(error), + } + } + "within" | "spatial_within" => { + let body = match serde_json::from_value::(payload) { + Ok(body) => body, + Err(error) => return invalid_body_response(error), + }; + let encoding = match parse_geometry_encoding(&body.encoding) { + Ok(encoding) => encoding, + Err(response) => return response, + }; + match state.geometry_within( + &tenant_id, + &body.label, + &body.property, + encoding, + &body.geometry, + ) { + Ok(ids) => geometry_ids_response(&tenant_id, &operation, ids), + Err(error) => store_unavailable_response(error), + } + } + _ => ( + StatusCode::NOT_FOUND, + Json(json!({ + "error": "unknown_spatial_operation", + "message": format!("unknown spatial plugin operation: {operation}") + })), + ) + .into_response(), + } +} + +#[cfg(feature = "geometry")] +fn parse_geometry_encoding( + raw: &str, +) -> Result { + match raw.trim().to_ascii_lowercase().as_str() { + "point" => Ok(rustyred_core::GeometryEncoding::Point), + "wkb" => Ok(rustyred_core::GeometryEncoding::Wkb), + "wkt" => Ok(rustyred_core::GeometryEncoding::Wkt), + "subgraph" => Ok(rustyred_core::GeometryEncoding::Subgraph), + other => Err(( + StatusCode::BAD_REQUEST, + Json(json!({ + "error": "invalid_geometry_encoding", + "message": format!("unsupported geometry encoding: {other}") + })), + ) + .into_response()), + } +} + +#[cfg(feature = "geometry")] +fn geometry_ids_response( + tenant_id: &str, + operation: &str, + ids: Vec, +) -> axum::response::Response { + Json(json!({ + "ok": true, + "tenant": tenant_id, + "operation": operation, + "count": ids.len(), + "node_ids": ids, + })) + .into_response() +} + +#[cfg(feature = "geometry")] +fn invalid_body_response(error: serde_json::Error) -> axum::response::Response { + ( + StatusCode::BAD_REQUEST, + Json(json!({ + "error": "invalid_request_body", + "message": error.to_string(), + })), + ) + .into_response() +} + async fn graph_algorithm_communities( State(state): State, Path(tenant_id): Path, @@ -4906,6 +5177,110 @@ mod tests { .starts_with("content_snapshot:")); } + #[cfg(feature = "geometry")] + #[tokio::test] + async fn spatial_geometry_plugin_routes_use_node_upsert_hooks_and_preserve_bbox() { + let state = memory_product_state(); + let tenant_id = "geometry-tenant"; + + let designate = super::graph_spatial_plugin_operation( + State(state.clone()), + Path((tenant_id.to_string(), "designate_geometry".to_string())), + HeaderMap::new(), + Json(json!({ + "label": "Parcel", + "property": "geom", + "encoding": "wkt", + "resolution": 9 + })), + ) + .await + .into_response(); + assert_eq!(designate.status(), StatusCode::OK); + + let upsert = execute_tenant_command( + &state, + tenant_id, + "RUSTYRED.GRAPH.NODE.UPSERT", + json!({ + "id": "parcel-1", + "labels": ["Parcel"], + "properties": { + "geom": "POLYGON((0 0,4 0,4 4,0 4,0 0))", + "lat": 2.0, + "lon": 2.0 + } + }), + ); + let upsert_payload = response_payload_json(upsert).await; + assert_eq!(upsert_payload["ok"], true); + + for (operation, payload) in [ + ( + "contains", + json!({ + "label": "Parcel", + "property": "geom", + "lat": 2.0, + "lon": 2.0 + }), + ), + ( + "intersects", + json!({ + "label": "Parcel", + "property": "geom", + "encoding": "wkt", + "geometry": "LINESTRING(2 -1,2 5)" + }), + ), + ( + "within", + json!({ + "label": "Parcel", + "property": "geom", + "encoding": "wkt", + "geometry": "POLYGON((-1 -1,5 -1,5 5,-1 5,-1 -1))" + }), + ), + ] { + let response = super::graph_spatial_plugin_operation( + State(state.clone()), + Path((tenant_id.to_string(), operation.to_string())), + HeaderMap::new(), + Json(payload), + ) + .await + .into_response(); + assert_eq!(response.status(), StatusCode::OK); + let payload = response_payload_json(response).await; + assert_eq!(payload["node_ids"], json!(["parcel-1"])); + } + + state + .designate_spatial_property(tenant_id, "Parcel", "lat", "lon", 9) + .unwrap(); + let bbox = super::graph_spatial_bbox( + State(state), + Path(tenant_id.to_string()), + HeaderMap::new(), + Json(super::SpatialBboxBody { + label: "Parcel".to_string(), + lat_property: "lat".to_string(), + lon_property: "lon".to_string(), + min_lat: 1.0, + min_lon: 1.0, + max_lat: 3.0, + max_lon: 3.0, + }), + ) + .await + .into_response(); + assert_eq!(bbox.status(), StatusCode::OK); + let bbox_payload = response_payload_json(bbox).await; + assert_eq!(bbox_payload["node_ids"], json!(["parcel-1"])); + } + #[tokio::test] async fn crawl_route_rejects_empty_seed_lists_before_fetching() { let response = crawl_submit( diff --git a/crates/rustyred-server/src/state.rs b/crates/rustyred-server/src/state.rs index 6914298..d2b172b 100644 --- a/crates/rustyred-server/src/state.rs +++ b/crates/rustyred-server/src/state.rs @@ -15,8 +15,12 @@ use rustyred_core::{ NeighborQuery, NodeQuery, NodeRecord, RedCoreGraphStore, RedCoreOptions, RedisGraphStore, SpatialBackend, SpatialDesignation, VectorDesignation, VerifyReport, }; +#[cfg(feature = "geometry")] +use rustyred_core::{GeometryDesignation, GeometryEncoding, GeometryIndex}; use rustyred_mcp::{McpError, McpGraphBackend, McpGraphProvider, McpServerConfig}; use serde_json::json; +#[cfg(feature = "geometry")] +use serde_json::Value; use crate::config::{Config, StorageMode}; use crate::graph_cache::GraphCacheTenant; @@ -40,6 +44,9 @@ type SpatialIndexes = BTreeMap>>; +#[cfg(feature = "geometry")] +type GeometryIndexes = BTreeMap>; + #[derive(Clone)] pub struct AppState { pub config: Arc, @@ -50,6 +57,8 @@ pub struct AppState { next_graph_txn_id: Arc, spatial_indexes: Arc>, fulltext_indexes: Arc>, + #[cfg(feature = "geometry")] + geometry_indexes: Arc>, } impl AppState { @@ -68,6 +77,8 @@ impl AppState { next_graph_txn_id: Arc::new(AtomicU64::new(1)), spatial_indexes: Arc::new(Mutex::new(BTreeMap::new())), fulltext_indexes: Arc::new(Mutex::new(BTreeMap::new())), + #[cfg(feature = "geometry")] + geometry_indexes: Arc::new(Mutex::new(BTreeMap::new())), } } @@ -329,6 +340,125 @@ impl AppState { Ok(index.bbox_search(min_lat, min_lon, max_lat, max_lon)) } + #[cfg(feature = "geometry")] + pub fn designate_geometry_property( + &self, + tenant_id: &str, + label: &str, + property: &str, + encoding: GeometryEncoding, + resolution: u8, + ) -> Result<(), StoreAccessError> { + if !(0..=15).contains(&resolution) { + return Err(StoreAccessError::unsupported(format!( + "geometry resolution {resolution} is outside 0..=15" + ))); + } + let store = self.tenant_graph_store(tenant_id)?; + let designation = GeometryDesignation::property( + label.to_string(), + property.to_string(), + encoding, + resolution, + ); + let mut index = GeometryIndex::new(designation) + .map_err(|err| StoreAccessError::unsupported(err.message()))?; + let nodes = store + .query_nodes(NodeQuery { + label: Some(label.to_string()), + ..NodeQuery::default() + }) + .map_err(StoreAccessError::from)?; + for node in nodes { + let _ = index.upsert_from_properties(&node.id, &node.properties); + } + let mut indexes = self + .geometry_indexes + .lock() + .map_err(|_| StoreAccessError::internal("geometry index lock poisoned"))?; + indexes + .entry(tenant_id.to_string()) + .or_default() + .insert((label.to_string(), property.to_string()), index); + Ok(()) + } + + #[cfg(feature = "geometry")] + pub fn maybe_index_node_geometry(&self, tenant_id: &str, node: &NodeRecord) { + let Ok(mut indexes) = self.geometry_indexes.lock() else { + return; + }; + let Some(tenant_map) = indexes.get_mut(tenant_id) else { + return; + }; + for ((label, _property), index) in tenant_map.iter_mut() { + if !node.labels.iter().any(|node_label| node_label == label) { + continue; + } + match index.upsert_from_properties(&node.id, &node.properties) { + Ok(()) => {} + Err(_) => index.remove(&node.id), + } + } + } + + #[cfg(feature = "geometry")] + pub fn geometry_contains_point( + &self, + tenant_id: &str, + label: &str, + property: &str, + lat: f64, + lon: f64, + ) -> Result, StoreAccessError> { + let indexes = self + .geometry_indexes + .lock() + .map_err(|_| StoreAccessError::internal("geometry index lock poisoned"))?; + let index = geometry_index(&indexes, tenant_id, label, property)?; + index + .contains_point(lat, lon) + .map_err(|err| StoreAccessError::unsupported(err.message())) + } + + #[cfg(feature = "geometry")] + pub fn geometry_intersects( + &self, + tenant_id: &str, + label: &str, + property: &str, + encoding: GeometryEncoding, + geometry: &Value, + ) -> Result, StoreAccessError> { + let indexes = self + .geometry_indexes + .lock() + .map_err(|_| StoreAccessError::internal("geometry index lock poisoned"))?; + let index = geometry_index(&indexes, tenant_id, label, property)?; + index + .intersects_value(encoding, geometry) + .map_err(|err| StoreAccessError::unsupported(err.message())) + } + + #[cfg(feature = "geometry")] + pub fn geometry_within( + &self, + tenant_id: &str, + label: &str, + property: &str, + encoding: GeometryEncoding, + geometry: &Value, + ) -> Result, StoreAccessError> { + let indexes = self + .geometry_indexes + .lock() + .map_err(|_| StoreAccessError::internal("geometry index lock poisoned"))?; + let index = geometry_index(&indexes, tenant_id, label, property)?; + index + .within_value(encoding, geometry) + .map_err(|err| StoreAccessError::unsupported(err.message())) + } + pub fn begin_graph_transaction(&self, tenant_id: &str) -> Result { self.purge_expired_graph_transactions()?; let store = match self.tenant_graph_store(tenant_id)? { @@ -428,9 +558,18 @@ impl AppState { "graph transaction snapshot conflict", )); } + let committed_mutations = context.mutations.clone(); let transaction = store - .commit_batch(context.mutations) + .commit_batch(committed_mutations.clone()) .map_err(StoreAccessError::from)?; + for mutation in &committed_mutations.mutations { + if let GraphMutation::NodeUpsert(node) = mutation { + self.maybe_index_node_spatially(tenant_id, node); + self.maybe_index_node_fulltext(tenant_id, node); + #[cfg(feature = "geometry")] + self.maybe_index_node_geometry(tenant_id, node); + } + } let mut transactions = self .graph_transactions .lock() @@ -984,6 +1123,27 @@ impl RedCoreTenantExecutor { } } +#[cfg(feature = "geometry")] +fn geometry_index<'a>( + indexes: &'a GeometryIndexes, + tenant_id: &str, + label: &str, + property: &str, +) -> Result<&'a GeometryIndex, StoreAccessError> { + let Some(tenant_map) = indexes.get(tenant_id) else { + return Err(StoreAccessError::unsupported( + "no geometry designations for this tenant", + )); + }; + tenant_map + .get(&(label.to_string(), property.to_string())) + .ok_or_else(|| { + StoreAccessError::unsupported( + "geometry designation not found; call /spatial/designate_geometry first", + ) + }) +} + fn now_millis() -> u64 { SystemTime::now() .duration_since(UNIX_EPOCH) @@ -1511,6 +1671,8 @@ impl McpGraphBackend for ProductMcpBackend { self.state .maybe_index_node_spatially(&self.tenant_id, &node); self.state.maybe_index_node_fulltext(&self.tenant_id, &node); + #[cfg(feature = "geometry")] + self.state.maybe_index_node_geometry(&self.tenant_id, &node); Ok(()) } @@ -1668,6 +1830,61 @@ impl McpGraphBackend for ProductMcpBackend { .map_err(|error| GraphStoreError::new(error.code, error.message)) } + #[cfg(feature = "geometry")] + fn designate_geometry_property( + &mut self, + label: &str, + property: &str, + encoding: &str, + resolution: u8, + ) -> GraphStoreResult<()> { + let encoding = parse_geometry_encoding_for_mcp(encoding)?; + self.state + .designate_geometry_property(&self.tenant_id, label, property, encoding, resolution) + .map_err(|error| GraphStoreError::new(error.code, error.message)) + } + + #[cfg(feature = "geometry")] + fn spatial_contains_point( + &self, + label: &str, + property: &str, + lat: f64, + lon: f64, + ) -> GraphStoreResult> { + self.state + .geometry_contains_point(&self.tenant_id, label, property, lat, lon) + .map_err(|error| GraphStoreError::new(error.code, error.message)) + } + + #[cfg(feature = "geometry")] + fn spatial_intersects_geometry( + &self, + label: &str, + property: &str, + encoding: &str, + geometry: &Value, + ) -> GraphStoreResult> { + let encoding = parse_geometry_encoding_for_mcp(encoding)?; + self.state + .geometry_intersects(&self.tenant_id, label, property, encoding, geometry) + .map_err(|error| GraphStoreError::new(error.code, error.message)) + } + + #[cfg(feature = "geometry")] + fn spatial_within_geometry( + &self, + label: &str, + property: &str, + encoding: &str, + geometry: &Value, + ) -> GraphStoreResult> { + let encoding = parse_geometry_encoding_for_mcp(encoding)?; + self.state + .geometry_within(&self.tenant_id, label, property, encoding, geometry) + .map_err(|error| GraphStoreError::new(error.code, error.message)) + } + fn epistemic_neighbors( &self, node_id: &str, @@ -1680,6 +1897,20 @@ impl McpGraphBackend for ProductMcpBackend { } } +#[cfg(feature = "geometry")] +fn parse_geometry_encoding_for_mcp(raw: &str) -> GraphStoreResult { + match raw.trim().to_ascii_lowercase().as_str() { + "point" => Ok(GeometryEncoding::Point), + "wkb" => Ok(GeometryEncoding::Wkb), + "wkt" => Ok(GeometryEncoding::Wkt), + "subgraph" => Ok(GeometryEncoding::Subgraph), + other => Err(GraphStoreError::new( + "invalid_geometry_encoding", + format!("unsupported geometry encoding: {other}"), + )), + } +} + impl McpGraphProvider for AppState { type Backend = ProductMcpBackend; diff --git a/vendor/proto/rustyred/v1/rustyred.proto b/vendor/proto/rustyred/v1/rustyred.proto index 445202e..8d2c459 100644 --- a/vendor/proto/rustyred/v1/rustyred.proto +++ b/vendor/proto/rustyred/v1/rustyred.proto @@ -68,6 +68,7 @@ service GraphDatabase { rpc SpatialRadius(SpatialRadiusRequest) returns (SpatialResponse); rpc SpatialBoundingBox(SpatialBboxRequest) returns (SpatialResponse); rpc DesignateSpatialProperty(DesignateSpatialRequest) returns (DesignateAck); + rpc ExecutePluginOperation(PluginOperationRequest) returns (PluginOperationResponse); // ---- Epistemic traversal ------------------------------------------------- rpc EpistemicNeighbors(EpistemicNeighborsRequest) returns (EpistemicNeighborsResponse); @@ -435,6 +436,16 @@ message DesignateSpatialRequest { string property_name = 3; } +message PluginOperationRequest { + string tenant_id = 1; + string operation = 2; + string args_json = 3; +} + +message PluginOperationResponse { + string result_json = 1; +} + // ============================================================================ // Epistemic traversal