1use std::collections::HashMap;
2use std::sync::Arc;
3
4use parking_lot::RwLock;
5
6use crate::wal::redo::UpdateSequenceRedo;
7use crate::wal::traits::WalWriter;
8use crate::wal::types::WalOpType;
9use crate::StorageError;
10
11use super::sequence::SequenceDef;
12
13pub trait SequenceStorage: Send + Sync {
15 fn load_all(&self) -> Result<Vec<SequenceDef>, StorageError>;
17 fn save(&self, def: &SequenceDef) -> Result<(), StorageError>;
19 fn delete(&self, name: &str) -> Result<(), StorageError>;
21 fn update_value(&self, name: &str, value: i64) -> Result<(), StorageError>;
23}
24
25pub struct SequenceManager {
30 sequences: RwLock<HashMap<String, Arc<SequenceDef>>>,
31 storage: Option<Arc<dyn SequenceStorage>>,
32 wal_writer: Option<Arc<parking_lot::Mutex<dyn WalWriter>>>,
33}
34
35impl SequenceManager {
36 pub fn new() -> Self {
38 Self {
39 sequences: RwLock::new(HashMap::new()),
40 storage: None,
41 wal_writer: None,
42 }
43 }
44
45 pub fn with_storage(storage: Arc<dyn SequenceStorage>) -> Self {
47 Self {
48 sequences: RwLock::new(HashMap::new()),
49 storage: Some(storage),
50 wal_writer: None,
51 }
52 }
53
54 pub fn with_wal(
56 storage: Arc<dyn SequenceStorage>,
57 wal_writer: Arc<parking_lot::Mutex<dyn WalWriter>>,
58 ) -> Self {
59 Self {
60 sequences: RwLock::new(HashMap::new()),
61 storage: Some(storage),
62 wal_writer: Some(wal_writer),
63 }
64 }
65
66 pub fn initialize(&self) -> Result<(), StorageError> {
68 if let Some(ref storage) = self.storage {
69 let defs = storage.load_all()?;
70 let mut map = self.sequences.write();
71 for def in defs {
72 let name = def.name.clone();
73 map.insert(name, Arc::new(def));
74 }
75 }
76 Ok(())
77 }
78
79 pub fn create_sequence(
81 &self,
82 name: String,
83 start: i64,
84 increment: i64,
85 min_value: i64,
86 max_value: i64,
87 cycle: bool,
88 ) -> Result<(), StorageError> {
89 {
90 let map = self.sequences.read();
91 if map.contains_key(&name) {
92 return Err(StorageError::db_error(format!(
93 "Sequence '{}' already exists",
94 name
95 )));
96 }
97 }
98
99 let def = SequenceDef::new(name.clone(), start, increment, min_value, max_value, cycle);
100
101 if let Some(ref storage) = self.storage {
102 storage.save(&def)?;
103 }
104
105 let mut map = self.sequences.write();
106 map.insert(name, Arc::new(def));
107 Ok(())
108 }
109
110 pub fn drop_sequence(&self, name: &str) -> Result<(), StorageError> {
112 {
113 let mut map = self.sequences.write();
114 if map.remove(name).is_none() {
115 return Err(StorageError::db_error(format!(
116 "Sequence '{}' does not exist",
117 name
118 )));
119 }
120 }
121
122 if let Some(ref storage) = self.storage {
123 storage.delete(name)?;
124 }
125
126 Ok(())
127 }
128
129 pub fn current_value(&self, name: &str) -> Result<i64, StorageError> {
131 let map = self.sequences.read();
132 let def = map
133 .get(name)
134 .ok_or_else(|| StorageError::db_error(format!("Sequence '{}' does not exist", name)))?;
135 Ok(def.current_value())
136 }
137
138 pub fn next_value(&self, name: &str) -> Result<i64, StorageError> {
140 let value = {
141 let map = self.sequences.read();
142 let def = map.get(name).ok_or_else(|| {
143 StorageError::db_error(format!("Sequence '{}' does not exist", name))
144 })?;
145 def.next_value()?
146 };
147
148 if let Some(ref storage) = self.storage {
149 storage.update_value(name, value)?;
150 }
151
152 if let Some(ref wal_writer) = self.wal_writer {
154 let redo = UpdateSequenceRedo {
155 space_id: 0,
156 table_name: name.to_string(),
157 next_value: value as u64,
158 };
159 let payload = postcard::to_allocvec(&redo)
160 .map_err(|e| StorageError::serialize_error(e.to_string()))?;
161 let mut writer = wal_writer.lock();
162 let ts = std::time::SystemTime::now()
163 .duration_since(std::time::UNIX_EPOCH)
164 .unwrap_or_default()
165 .as_millis() as u64;
166 writer
167 .append_entry(WalOpType::UpdateSequence, ts, &payload)
168 .map_err(|e| StorageError::db_error(format!("WAL write error: {}", e)))?;
169 }
170
171 Ok(value)
172 }
173
174 pub fn alter_sequence(
176 &self,
177 name: &str,
178 increment: Option<i64>,
179 min_value: Option<i64>,
180 max_value: Option<i64>,
181 cycle: Option<bool>,
182 ) -> Result<(), StorageError> {
183 let map = self.sequences.read();
184 let def = map
185 .get(name)
186 .ok_or_else(|| StorageError::db_error(format!("Sequence '{}' does not exist", name)))?;
187
188 let new_increment = increment.unwrap_or(def.increment);
189 let new_min = min_value.unwrap_or(def.min_value);
190 let new_max = max_value.unwrap_or(def.max_value);
191 let new_cycle = cycle.unwrap_or(def.cycle);
192 let current = def.current_value();
193
194 drop(map);
195
196 let new_def = SequenceDef::new(
197 name.to_string(),
198 current,
199 new_increment,
200 new_min,
201 new_max,
202 new_cycle,
203 );
204
205 if let Some(ref storage) = self.storage {
206 storage.save(&new_def)?;
207 }
208
209 let mut map = self.sequences.write();
210 map.insert(name.to_string(), Arc::new(new_def));
211 Ok(())
212 }
213
214 pub fn flush(&self) -> Result<(), StorageError> {
216 if let Some(ref storage) = self.storage {
217 let map = self.sequences.read();
218 for def in map.values() {
219 storage.save(def)?;
220 }
221 }
222 Ok(())
223 }
224
225 pub fn list_sequences(&self) -> Vec<String> {
227 let map = self.sequences.read();
228 map.keys().cloned().collect()
229 }
230
231 pub fn exists(&self, name: &str) -> bool {
233 let map = self.sequences.read();
234 map.contains_key(name)
235 }
236}
237
238impl Default for SequenceManager {
239 fn default() -> Self {
240 Self::new()
241 }
242}
243
244impl std::fmt::Debug for SequenceManager {
245 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
246 f.debug_struct("SequenceManager")
247 .field("count", &self.sequences.read().len())
248 .finish()
249 }
250}
251
252#[cfg(test)]
253mod tests {
254 use super::*;
255
256 struct MockStorage {
257 data: RwLock<HashMap<String, Vec<u8>>>,
258 }
259
260 impl MockStorage {
261 fn new() -> Self {
262 Self {
263 data: RwLock::new(HashMap::new()),
264 }
265 }
266 }
267
268 impl SequenceStorage for MockStorage {
269 fn load_all(&self) -> Result<Vec<SequenceDef>, StorageError> {
270 let data = self.data.read();
271 for _value in data.values() {
272 }
275 Ok(Vec::new())
276 }
277
278 fn save(&self, def: &SequenceDef) -> Result<(), StorageError> {
279 let key = format!("sequences:{}", def.name);
280 let value = postcard::to_allocvec(def)
281 .map_err(|e| StorageError::serialize_error(e.to_string()))?;
282 self.data.write().insert(key, value);
283 Ok(())
284 }
285
286 fn delete(&self, name: &str) -> Result<(), StorageError> {
287 let key = format!("sequences:{}", name);
288 self.data.write().remove(&key);
289 Ok(())
290 }
291
292 fn update_value(&self, name: &str, value: i64) -> Result<(), StorageError> {
293 let key = format!("sequences:{}", name);
294 let mut data = self.data.write();
295 if let Some(raw) = data.get(&key) {
296 let def: SequenceDef = postcard::from_bytes(raw)
297 .map_err(|e| StorageError::deserialize_error(e.to_string()))?;
298 def.set_value(value);
299 let new_raw = postcard::to_allocvec(&def)
300 .map_err(|e| StorageError::serialize_error(e.to_string()))?;
301 data.insert(key, new_raw);
302 }
303 Ok(())
304 }
305 }
306
307 #[test]
308 fn test_manager_create_drop() {
309 let manager = SequenceManager::new();
310 manager
311 .create_sequence("seq1".into(), 1, 1, 1, 100, false)
312 .unwrap();
313 assert!(manager.exists("seq1"));
314 assert_eq!(manager.current_value("seq1").unwrap(), 1);
315
316 manager.drop_sequence("seq1").unwrap();
317 assert!(!manager.exists("seq1"));
318 }
319
320 #[test]
321 fn test_manager_duplicate_create() {
322 let manager = SequenceManager::new();
323 manager
324 .create_sequence("seq1".into(), 1, 1, 1, 100, false)
325 .unwrap();
326 assert!(manager
327 .create_sequence("seq1".into(), 1, 1, 1, 100, false)
328 .is_err());
329 }
330
331 #[test]
332 fn test_manager_next_value() {
333 let manager = SequenceManager::new();
334 manager
335 .create_sequence("seq1".into(), 10, 5, 1, 1000, false)
336 .unwrap();
337 assert_eq!(manager.next_value("seq1").unwrap(), 15);
338 assert_eq!(manager.next_value("seq1").unwrap(), 20);
339 assert_eq!(manager.current_value("seq1").unwrap(), 20);
340 }
341
342 #[test]
343 fn test_manager_alter() {
344 let manager = SequenceManager::new();
345 manager
346 .create_sequence("seq1".into(), 1, 1, 1, 100, false)
347 .unwrap();
348 manager
349 .alter_sequence("seq1", Some(10), None, None, Some(true))
350 .unwrap();
351 assert_eq!(manager.next_value("seq1").unwrap(), 11);
352 }
353
354 #[test]
355 fn test_manager_list() {
356 let manager = SequenceManager::new();
357 manager
358 .create_sequence("a".into(), 1, 1, 1, 100, false)
359 .unwrap();
360 manager
361 .create_sequence("b".into(), 1, 1, 1, 100, false)
362 .unwrap();
363 let mut names = manager.list_sequences();
364 names.sort();
365 assert_eq!(names, vec!["a", "b"]);
366 }
367
368 #[test]
369 fn test_manager_with_storage() {
370 let storage = Arc::new(MockStorage::new());
371 let manager = SequenceManager::with_storage(storage);
372 manager
373 .create_sequence("seq1".into(), 1, 1, 1, 100, false)
374 .unwrap();
375 assert!(manager.exists("seq1"));
376 }
377}