1use serde::{Deserialize, Serialize};
11
12#[derive(
14 Debug,
15 Clone,
16 Serialize,
17 Deserialize,
18 rkyv::Archive,
19 rkyv::Serialize,
20 rkyv::Deserialize,
21 zerompk::ToMessagePack,
22 zerompk::FromMessagePack,
23)]
24pub struct ShapeDefinition {
25 pub shape_id: String,
27 pub tenant_id: u32,
29 pub shape_type: ShapeType,
31 pub description: String,
33 #[serde(default)]
36 pub field_filter: Vec<String>,
37}
38
39#[derive(
41 Debug,
42 Clone,
43 Serialize,
44 Deserialize,
45 rkyv::Archive,
46 rkyv::Serialize,
47 rkyv::Deserialize,
48 zerompk::ToMessagePack,
49 zerompk::FromMessagePack,
50)]
51pub enum ShapeType {
52 Document {
56 collection: String,
57 predicate: Vec<u8>,
59 },
60
61 Graph {
66 root_nodes: Vec<String>,
67 max_depth: usize,
68 edge_label: Option<String>,
69 },
70
71 Vector {
75 collection: String,
76 field_name: Option<String>,
77 },
78}
79
80impl ShapeDefinition {
81 pub fn could_match(&self, collection: &str, _doc_id: &str) -> bool {
86 match &self.shape_type {
87 ShapeType::Document {
88 collection: shape_coll,
89 ..
90 } => shape_coll == collection,
91 ShapeType::Graph { root_nodes, .. } => {
92 !root_nodes.is_empty()
94 }
95 ShapeType::Vector {
96 collection: shape_coll,
97 ..
98 } => shape_coll == collection,
99 }
100 }
101
102 pub fn collection(&self) -> Option<&str> {
104 match &self.shape_type {
105 ShapeType::Document { collection, .. } => Some(collection),
106 ShapeType::Vector { collection, .. } => Some(collection),
107 ShapeType::Graph { .. } => None,
108 }
109 }
110}
111
112#[cfg(test)]
113mod tests {
114 use super::*;
115
116 #[test]
117 fn document_shape_matches_collection() {
118 let shape = ShapeDefinition {
119 shape_id: "s1".into(),
120 tenant_id: 1,
121 shape_type: ShapeType::Document {
122 collection: "orders".into(),
123 predicate: Vec::new(),
124 },
125 description: "all orders".into(),
126 field_filter: vec![],
127 };
128
129 assert!(shape.could_match("orders", "o1"));
130 assert!(!shape.could_match("users", "u1"));
131 assert_eq!(shape.collection(), Some("orders"));
132 }
133
134 #[test]
135 fn graph_shape() {
136 let shape = ShapeDefinition {
137 shape_id: "g1".into(),
138 tenant_id: 1,
139 shape_type: ShapeType::Graph {
140 root_nodes: vec!["alice".into()],
141 max_depth: 2,
142 edge_label: Some("KNOWS".into()),
143 },
144 description: "alice's network".into(),
145 field_filter: vec![],
146 };
147
148 assert!(shape.could_match("any_collection", "any_doc"));
149 assert_eq!(shape.collection(), None);
150 }
151
152 #[test]
153 fn vector_shape() {
154 let shape = ShapeDefinition {
155 shape_id: "v1".into(),
156 tenant_id: 1,
157 shape_type: ShapeType::Vector {
158 collection: "embeddings".into(),
159 field_name: Some("title".into()),
160 },
161 description: "title embeddings".into(),
162 field_filter: vec![],
163 };
164
165 assert!(shape.could_match("embeddings", "e1"));
166 assert!(!shape.could_match("other", "e1"));
167 }
168
169 #[test]
170 fn msgpack_roundtrip() {
171 let shape = ShapeDefinition {
172 shape_id: "test".into(),
173 tenant_id: 5,
174 shape_type: ShapeType::Document {
175 collection: "users".into(),
176 predicate: vec![1, 2, 3],
177 },
178 description: "test shape".into(),
179 field_filter: vec![],
180 };
181 let bytes = zerompk::to_msgpack_vec(&shape).unwrap();
182 let decoded: ShapeDefinition = zerompk::from_msgpack(&bytes).unwrap();
183 assert_eq!(decoded.shape_id, "test");
184 assert_eq!(decoded.tenant_id, 5);
185 assert!(matches!(decoded.shape_type, ShapeType::Document { .. }));
186 }
187}