1#![allow(clippy::doc_markdown)]
5
6use alloc::collections::BTreeMap;
22use alloc::string::{String, ToString};
23use alloc::vec::Vec;
24
25use spg_storage::{ColumnSchema, DataType, Row, Value};
26
27use crate::{Engine, EngineError, QueryResult};
28
29use spg_sql::ast::{CreatePublicationStatement, PublicationScope};
30
31const SCOPE_ALL_TABLES: u8 = 0;
35const SCOPE_FOR_TABLES: u8 = 1;
36const SCOPE_ALL_TABLES_EXCEPT: u8 = 2;
37
38#[derive(Debug, Clone, PartialEq, Eq, Default)]
39pub struct Publications {
40 inner: BTreeMap<String, PublicationScope>,
43}
44
45#[derive(Debug, PartialEq, Eq)]
46pub enum PublicationError {
47 DuplicateName(String),
48 Corrupt(String),
52}
53
54impl Publications {
55 pub fn new() -> Self {
56 Self::default()
57 }
58
59 pub fn len(&self) -> usize {
60 self.inner.len()
61 }
62
63 pub fn is_empty(&self) -> bool {
64 self.inner.is_empty()
65 }
66
67 pub fn contains(&self, name: &str) -> bool {
68 self.inner.contains_key(name)
69 }
70
71 pub fn get(&self, name: &str) -> Option<&PublicationScope> {
76 self.inner.get(name)
77 }
78
79 pub fn iter(&self) -> impl Iterator<Item = (&String, &PublicationScope)> {
82 self.inner.iter()
83 }
84
85 pub fn create(
89 &mut self,
90 name: String,
91 scope: PublicationScope,
92 ) -> Result<(), PublicationError> {
93 if self.inner.contains_key(&name) {
94 return Err(PublicationError::DuplicateName(name));
95 }
96 self.inner.insert(name, scope);
97 Ok(())
98 }
99
100 pub fn drop(&mut self, name: &str) -> bool {
105 self.inner.remove(name).is_some()
106 }
107
108 pub fn serialize(&self) -> Vec<u8> {
120 let mut out = Vec::with_capacity(2 + self.inner.len() * 16);
121 let n = u16::try_from(self.inner.len()).expect("≤ 65,535 publications per cluster");
122 out.extend_from_slice(&n.to_le_bytes());
123 for (name, scope) in &self.inner {
124 write_str(&mut out, name);
125 match scope {
126 PublicationScope::AllTables => out.push(SCOPE_ALL_TABLES),
127 PublicationScope::ForTables(ts) => {
128 out.push(SCOPE_FOR_TABLES);
129 write_table_list(&mut out, ts);
130 }
131 PublicationScope::AllTablesExcept(ts) => {
132 out.push(SCOPE_ALL_TABLES_EXCEPT);
133 write_table_list(&mut out, ts);
134 }
135 PublicationScope::TablesInSchema(_) => unreachable!(),
138 }
139 }
140 out
141 }
142
143 pub fn deserialize(buf: &[u8]) -> Result<Self, PublicationError> {
144 let mut p = 0usize;
145 let n = read_u16(buf, &mut p)? as usize;
146 let mut inner = BTreeMap::new();
147 for _ in 0..n {
148 let name = read_str(buf, &mut p)?;
149 let tag = read_u8(buf, &mut p)?;
150 let scope = match tag {
151 SCOPE_ALL_TABLES => PublicationScope::AllTables,
152 SCOPE_FOR_TABLES => PublicationScope::ForTables(read_table_list(buf, &mut p)?),
153 SCOPE_ALL_TABLES_EXCEPT => {
154 PublicationScope::AllTablesExcept(read_table_list(buf, &mut p)?)
155 }
156 other => {
157 return Err(PublicationError::Corrupt(alloc::format!(
158 "unknown publication scope tag {other:#x}"
159 )));
160 }
161 };
162 if inner.insert(name.clone(), scope).is_some() {
163 return Err(PublicationError::Corrupt(alloc::format!(
164 "duplicate publication name {name:?} in serialised payload"
165 )));
166 }
167 }
168 if p != buf.len() {
169 return Err(PublicationError::Corrupt(alloc::format!(
170 "trailing bytes in publications payload: read {p}, len {}",
171 buf.len()
172 )));
173 }
174 Ok(Self { inner })
175 }
176}
177
178fn write_str(out: &mut Vec<u8>, s: &str) {
179 let n = u16::try_from(s.len()).expect("publication / table name fits in u16");
180 out.extend_from_slice(&n.to_le_bytes());
181 out.extend_from_slice(s.as_bytes());
182}
183
184fn write_table_list(out: &mut Vec<u8>, ts: &[String]) {
185 let n = u16::try_from(ts.len()).expect("≤ 65,535 tables per publication");
186 out.extend_from_slice(&n.to_le_bytes());
187 for t in ts {
188 write_str(out, t);
189 }
190}
191
192fn read_u8(buf: &[u8], p: &mut usize) -> Result<u8, PublicationError> {
193 let v = buf
194 .get(*p)
195 .copied()
196 .ok_or_else(|| PublicationError::Corrupt("short read (u8)".to_string()))?;
197 *p += 1;
198 Ok(v)
199}
200
201fn read_u16(buf: &[u8], p: &mut usize) -> Result<u16, PublicationError> {
202 let slice = buf
203 .get(*p..*p + 2)
204 .ok_or_else(|| PublicationError::Corrupt("short read (u16)".to_string()))?;
205 let arr: [u8; 2] = slice
206 .try_into()
207 .map_err(|_| PublicationError::Corrupt("u16 slice".to_string()))?;
208 *p += 2;
209 Ok(u16::from_le_bytes(arr))
210}
211
212fn read_str(buf: &[u8], p: &mut usize) -> Result<String, PublicationError> {
213 let n = read_u16(buf, p)? as usize;
214 let slice = buf
215 .get(*p..*p + n)
216 .ok_or_else(|| PublicationError::Corrupt(alloc::format!("short read (str, {n} bytes)")))?;
217 *p += n;
218 core::str::from_utf8(slice)
219 .map(ToString::to_string)
220 .map_err(|e| PublicationError::Corrupt(alloc::format!("non-UTF-8 str: {e}")))
221}
222
223fn read_table_list(buf: &[u8], p: &mut usize) -> Result<Vec<String>, PublicationError> {
224 let n = read_u16(buf, p)? as usize;
225 let mut out = Vec::with_capacity(n);
226 for _ in 0..n {
227 out.push(read_str(buf, p)?);
228 }
229 Ok(out)
230}
231
232impl Engine {
233 pub(crate) fn exec_show_publications(&self) -> QueryResult {
245 let columns = alloc::vec![
246 ColumnSchema::new("name", DataType::Text, false),
247 ColumnSchema::new("scope", DataType::Text, false),
248 ColumnSchema::new("table_count", DataType::Int, true),
249 ];
250 let rows: Vec<Row<'static>> = self
251 .publications
252 .iter()
253 .map(|(name, scope)| {
254 let (scope_str, count_val) = match scope {
255 spg_sql::ast::PublicationScope::AllTables => {
256 ("FOR ALL TABLES".to_string(), Value::Null)
257 }
258 spg_sql::ast::PublicationScope::ForTables(ts) => (
259 alloc::format!("FOR TABLE {}", ts.join(", ")),
260 Value::Int(i32::try_from(ts.len()).unwrap_or(i32::MAX)),
261 ),
262 spg_sql::ast::PublicationScope::AllTablesExcept(ts) => (
263 alloc::format!("FOR ALL TABLES EXCEPT {}", ts.join(", ")),
264 Value::Int(i32::try_from(ts.len()).unwrap_or(i32::MAX)),
265 ),
266 spg_sql::ast::PublicationScope::TablesInSchema(_) => unreachable!(),
268 };
269 Row::new(alloc::vec![
270 Value::text(name.clone()),
271 Value::text(scope_str),
272 count_val,
273 ])
274 })
275 .collect();
276 QueryResult::Rows { columns, rows }
277 }
278
279 pub(crate) fn exec_create_publication(
287 &mut self,
288 s: CreatePublicationStatement,
289 ) -> Result<QueryResult, EngineError> {
290 if let PublicationScope::ForTables(ts) | PublicationScope::AllTablesExcept(ts) = &s.scope {
303 for t in ts {
304 if self.active_catalog().get(t).is_none() {
305 return Err(EngineError::Unsupported(alloc::format!(
306 "relation \"{t}\" does not exist"
307 )));
308 }
309 }
310 }
311 let scope = match s.scope {
312 PublicationScope::TablesInSchema(schema) => {
313 if schema.eq_ignore_ascii_case("public") {
314 PublicationScope::AllTables
315 } else {
316 return Err(EngineError::Unsupported(alloc::format!(
317 "schema \"{schema}\" does not exist"
318 )));
319 }
320 }
321 other => other,
322 };
323 self.publications
324 .create(s.name, scope)
325 .map_err(|e| EngineError::Unsupported(alloc::format!("CREATE PUBLICATION: {e:?}")))?;
326 Ok(QueryResult::CommandOk {
327 affected: 1,
328 modified_catalog: true,
329 })
330 }
331
332 pub(crate) fn exec_drop_publication(
339 &mut self,
340 name: &str,
341 if_exists: bool,
342 ) -> Result<QueryResult, EngineError> {
343 let removed = self.publications.drop(name);
344 if !removed && !if_exists {
345 return Err(EngineError::Unsupported(alloc::format!(
346 "publication \"{name}\" does not exist"
347 )));
348 }
349 Ok(QueryResult::CommandOk {
350 affected: usize::from(removed),
351 modified_catalog: removed,
352 })
353 }
354
355 pub const fn publications(&self) -> &Publications {
360 &self.publications
361 }
362}
363
364#[cfg(test)]
365mod tests {
366 use super::*;
367
368 #[test]
369 fn empty_roundtrips() {
370 let p = Publications::new();
371 let bytes = p.serialize();
372 let p2 = Publications::deserialize(&bytes).unwrap();
373 assert_eq!(p, p2);
374 }
375
376 #[test]
377 fn single_all_tables_roundtrips() {
378 let mut p = Publications::new();
379 p.create("pub_a".into(), PublicationScope::AllTables)
380 .unwrap();
381 let bytes = p.serialize();
382 let p2 = Publications::deserialize(&bytes).unwrap();
383 assert_eq!(p, p2);
384 assert!(p2.contains("pub_a"));
385 assert_eq!(p2.len(), 1);
386 }
387
388 #[test]
389 fn duplicate_create_errors() {
390 let mut p = Publications::new();
391 p.create("pub_a".into(), PublicationScope::AllTables)
392 .unwrap();
393 let err = p
394 .create("pub_a".into(), PublicationScope::AllTables)
395 .unwrap_err();
396 assert_eq!(err, PublicationError::DuplicateName("pub_a".into()));
397 }
398
399 #[test]
400 fn drop_present_returns_true_drop_absent_false() {
401 let mut p = Publications::new();
402 p.create("pub_a".into(), PublicationScope::AllTables)
403 .unwrap();
404 assert!(p.drop("pub_a"));
405 assert!(!p.drop("pub_a"));
406 assert!(!p.drop("never_existed"));
407 }
408
409 #[test]
413 fn for_tables_scope_roundtrips() {
414 let mut p = Publications::new();
415 p.create(
416 "p_pick".into(),
417 PublicationScope::ForTables(alloc::vec!["t1".into(), "t2".into()]),
418 )
419 .unwrap();
420 let bytes = p.serialize();
421 let p2 = Publications::deserialize(&bytes).unwrap();
422 assert_eq!(p, p2);
423 }
424
425 #[test]
426 fn all_tables_except_scope_roundtrips() {
427 let mut p = Publications::new();
428 p.create(
429 "p_neg".into(),
430 PublicationScope::AllTablesExcept(alloc::vec!["t3".into()]),
431 )
432 .unwrap();
433 let bytes = p.serialize();
434 let p2 = Publications::deserialize(&bytes).unwrap();
435 assert_eq!(p, p2);
436 }
437
438 #[test]
439 fn corrupt_tag_errors() {
440 let mut buf = Vec::new();
442 buf.extend_from_slice(&1u16.to_le_bytes()); buf.extend_from_slice(&3u16.to_le_bytes()); buf.extend_from_slice(b"bad");
445 buf.push(0xFF); let err = Publications::deserialize(&buf).unwrap_err();
447 assert!(matches!(err, PublicationError::Corrupt(_)));
448 }
449
450 #[test]
451 fn trailing_bytes_errors() {
452 let mut p = Publications::new();
453 p.create("pub_a".into(), PublicationScope::AllTables)
454 .unwrap();
455 let mut bytes = p.serialize();
456 bytes.push(0xCC);
457 let err = Publications::deserialize(&bytes).unwrap_err();
458 assert!(matches!(err, PublicationError::Corrupt(_)));
459 }
460
461 #[test]
462 fn deterministic_order_independent_of_insert_sequence() {
463 let mut p1 = Publications::new();
465 p1.create("z".into(), PublicationScope::AllTables).unwrap();
466 p1.create("a".into(), PublicationScope::AllTables).unwrap();
467 let mut p2 = Publications::new();
468 p2.create("a".into(), PublicationScope::AllTables).unwrap();
469 p2.create("z".into(), PublicationScope::AllTables).unwrap();
470 assert_eq!(p1.serialize(), p2.serialize());
471 }
472}