1use indexmap::IndexMap;
10use serde::{Deserialize, Serialize};
11use smol_str::SmolStr;
12
13use super::{Attribute, Documentation, Ident, Span};
14
15#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
45pub struct ServerGroup {
46 pub name: Ident,
48 pub servers: IndexMap<SmolStr, Server>,
50 pub attributes: Vec<Attribute>,
52 pub documentation: Option<Documentation>,
54 pub span: Span,
56}
57
58impl ServerGroup {
59 pub fn new(name: Ident, span: Span) -> Self {
61 Self {
62 name,
63 servers: IndexMap::new(),
64 attributes: vec![],
65 documentation: None,
66 span,
67 }
68 }
69
70 pub fn add_server(&mut self, server: Server) {
72 self.servers.insert(server.name.name.clone(), server);
73 }
74
75 pub fn add_attribute(&mut self, attr: Attribute) {
77 self.attributes.push(attr);
78 }
79
80 pub fn set_documentation(&mut self, doc: Documentation) {
82 self.documentation = Some(doc);
83 }
84
85 pub fn strategy(&self) -> Option<ServerGroupStrategy> {
87 for attr in &self.attributes {
88 if attr.name.name == "strategy" {
89 if let Some(arg) = attr.args.first() {
90 let value_str = arg
91 .value
92 .as_string()
93 .map(|s| s.to_string())
94 .or_else(|| arg.value.as_ident().map(|s| s.to_string()))?;
95 return ServerGroupStrategy::parse(&value_str);
96 }
97 }
98 }
99 None
100 }
101
102 pub fn load_balance(&self) -> Option<LoadBalanceStrategy> {
104 for attr in &self.attributes {
105 if attr.name.name == "loadBalance" {
106 if let Some(arg) = attr.args.first() {
107 let value_str = arg
108 .value
109 .as_string()
110 .map(|s| s.to_string())
111 .or_else(|| arg.value.as_ident().map(|s| s.to_string()))?;
112 return LoadBalanceStrategy::parse(&value_str);
113 }
114 }
115 }
116 None
117 }
118
119 pub fn primary(&self) -> Option<&Server> {
121 self.servers
122 .values()
123 .find(|s| s.role() == Some(ServerRole::Primary))
124 }
125
126 pub fn replicas(&self) -> Vec<&Server> {
128 self.servers
129 .values()
130 .filter(|s| s.role() == Some(ServerRole::Replica))
131 .collect()
132 }
133
134 pub fn servers_in_region(&self, region: &str) -> Vec<&Server> {
136 self.servers
137 .values()
138 .filter(|s| s.region() == Some(region))
139 .collect()
140 }
141
142 pub fn failover_order(&self) -> Vec<&Server> {
144 let mut servers: Vec<_> = self.servers.values().collect();
145 servers.sort_by_key(|s| s.priority().unwrap_or(u32::MAX));
146 servers
147 }
148}
149
150#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
152pub struct Server {
153 pub name: Ident,
155 pub properties: IndexMap<SmolStr, ServerProperty>,
157 pub span: Span,
159}
160
161impl Server {
162 pub fn new(name: Ident, span: Span) -> Self {
164 Self {
165 name,
166 properties: IndexMap::new(),
167 span,
168 }
169 }
170
171 pub fn add_property(&mut self, prop: ServerProperty) {
173 self.properties.insert(prop.name.clone(), prop);
174 }
175
176 pub fn get_property(&self, name: &str) -> Option<&ServerPropertyValue> {
178 self.properties.get(name).map(|p| &p.value)
179 }
180
181 pub fn url(&self) -> Option<&str> {
183 match self.get_property("url")? {
184 ServerPropertyValue::String(s) => Some(s),
185 ServerPropertyValue::EnvVar(var) => Some(var),
186 _ => None,
187 }
188 }
189
190 pub fn role(&self) -> Option<ServerRole> {
192 match self.get_property("role")? {
193 ServerPropertyValue::String(s) | ServerPropertyValue::Identifier(s) => {
194 ServerRole::parse(s)
195 }
196 _ => None,
197 }
198 }
199
200 pub fn weight(&self) -> Option<u32> {
202 match self.get_property("weight")? {
203 ServerPropertyValue::Number(n) => Some(*n as u32),
204 _ => None,
205 }
206 }
207
208 pub fn region(&self) -> Option<&str> {
210 match self.get_property("region")? {
211 ServerPropertyValue::String(s) => Some(s),
212 _ => None,
213 }
214 }
215
216 pub fn priority(&self) -> Option<u32> {
218 match self.get_property("priority")? {
219 ServerPropertyValue::Number(n) => Some(*n as u32),
220 _ => None,
221 }
222 }
223
224 pub fn is_read_only(&self) -> bool {
226 match self.get_property("readOnly") {
227 Some(ServerPropertyValue::Boolean(b)) => *b,
228 _ => self.role() == Some(ServerRole::Replica),
229 }
230 }
231
232 pub fn max_connections(&self) -> Option<u32> {
234 match self.get_property("maxConnections")? {
235 ServerPropertyValue::Number(n) => Some(*n as u32),
236 _ => None,
237 }
238 }
239
240 pub fn health_check(&self) -> Option<&str> {
242 match self.get_property("healthCheck")? {
243 ServerPropertyValue::String(s) => Some(s),
244 _ => None,
245 }
246 }
247}
248
249#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
251pub struct ServerProperty {
252 pub name: SmolStr,
254 pub value: ServerPropertyValue,
256 pub span: Span,
258}
259
260impl ServerProperty {
261 pub fn new(name: impl Into<SmolStr>, value: ServerPropertyValue, span: Span) -> Self {
263 Self {
264 name: name.into(),
265 value,
266 span,
267 }
268 }
269}
270
271#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
273pub enum ServerPropertyValue {
274 String(String),
276 Number(f64),
278 Boolean(bool),
280 Identifier(String),
282 EnvVar(String),
284 Array(Vec<ServerPropertyValue>),
286}
287
288impl std::fmt::Display for ServerPropertyValue {
289 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
290 match self {
291 Self::String(s) => write!(f, "\"{}\"", s),
292 Self::Number(n) => write!(f, "{}", n),
293 Self::Boolean(b) => write!(f, "{}", b),
294 Self::Identifier(s) => write!(f, "{}", s),
295 Self::EnvVar(var) => write!(f, "env(\"{}\")", var),
296 Self::Array(arr) => {
297 write!(f, "[")?;
298 for (i, v) in arr.iter().enumerate() {
299 if i > 0 {
300 write!(f, ", ")?;
301 }
302 write!(f, "{}", v)?;
303 }
304 write!(f, "]")
305 }
306 }
307 }
308}
309
310#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
312pub enum ServerRole {
313 Primary,
315 Replica,
317 Analytics,
319 Archive,
321 Shard,
323}
324
325impl ServerRole {
326 pub fn parse(s: &str) -> Option<Self> {
328 match s.to_lowercase().as_str() {
329 "primary" | "master" | "writer" => Some(Self::Primary),
330 "replica" | "slave" | "reader" | "read" => Some(Self::Replica),
331 "analytics" | "reporting" | "olap" => Some(Self::Analytics),
332 "archive" | "historical" => Some(Self::Archive),
333 "shard" => Some(Self::Shard),
334 _ => None,
335 }
336 }
337
338 pub fn as_str(&self) -> &'static str {
340 match self {
341 Self::Primary => "primary",
342 Self::Replica => "replica",
343 Self::Analytics => "analytics",
344 Self::Archive => "archive",
345 Self::Shard => "shard",
346 }
347 }
348}
349
350#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
352pub enum ServerGroupStrategy {
353 ReadReplica,
355 Sharding,
357 MultiRegion,
359 HighAvailability,
361 Custom,
363}
364
365impl ServerGroupStrategy {
366 pub fn parse(s: &str) -> Option<Self> {
368 match s.to_lowercase().replace(['-', '_'], "").as_str() {
369 "readreplica" | "replication" => Some(Self::ReadReplica),
370 "sharding" | "shard" | "partition" => Some(Self::Sharding),
371 "multiregion" | "georeplica" | "geographic" => Some(Self::MultiRegion),
372 "highavailability" | "ha" | "failover" => Some(Self::HighAvailability),
373 "custom" => Some(Self::Custom),
374 _ => None,
375 }
376 }
377
378 pub fn as_str(&self) -> &'static str {
380 match self {
381 Self::ReadReplica => "ReadReplica",
382 Self::Sharding => "Sharding",
383 Self::MultiRegion => "MultiRegion",
384 Self::HighAvailability => "HighAvailability",
385 Self::Custom => "Custom",
386 }
387 }
388}
389
390#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
392pub enum LoadBalanceStrategy {
393 RoundRobin,
395 Random,
397 LeastConnections,
399 Weighted,
401 Nearest,
403 Sticky,
405}
406
407impl LoadBalanceStrategy {
408 pub fn parse(s: &str) -> Option<Self> {
410 match s.to_lowercase().replace(['-', '_'], "").as_str() {
411 "roundrobin" | "rr" => Some(Self::RoundRobin),
412 "random" | "rand" => Some(Self::Random),
413 "leastconnections" | "leastconn" | "least" => Some(Self::LeastConnections),
414 "weighted" | "weight" => Some(Self::Weighted),
415 "nearest" | "latency" | "geo" => Some(Self::Nearest),
416 "sticky" | "affinity" | "session" => Some(Self::Sticky),
417 _ => None,
418 }
419 }
420
421 pub fn as_str(&self) -> &'static str {
423 match self {
424 Self::RoundRobin => "RoundRobin",
425 Self::Random => "Random",
426 Self::LeastConnections => "LeastConnections",
427 Self::Weighted => "Weighted",
428 Self::Nearest => "Nearest",
429 Self::Sticky => "Sticky",
430 }
431 }
432}
433
434#[cfg(test)]
435mod tests {
436 use super::*;
437
438 #[test]
439 fn test_server_role_from_str() {
440 assert_eq!(ServerRole::parse("primary"), Some(ServerRole::Primary));
441 assert_eq!(ServerRole::parse("master"), Some(ServerRole::Primary));
442 assert_eq!(ServerRole::parse("replica"), Some(ServerRole::Replica));
443 assert_eq!(ServerRole::parse("slave"), Some(ServerRole::Replica));
444 assert_eq!(
445 ServerRole::parse("analytics"),
446 Some(ServerRole::Analytics)
447 );
448 assert_eq!(ServerRole::parse("shard"), Some(ServerRole::Shard));
449 assert_eq!(ServerRole::parse("invalid"), None);
450 }
451
452 #[test]
453 fn test_server_group_strategy_from_str() {
454 assert_eq!(
455 ServerGroupStrategy::parse("ReadReplica"),
456 Some(ServerGroupStrategy::ReadReplica)
457 );
458 assert_eq!(
459 ServerGroupStrategy::parse("sharding"),
460 Some(ServerGroupStrategy::Sharding)
461 );
462 assert_eq!(
463 ServerGroupStrategy::parse("multi-region"),
464 Some(ServerGroupStrategy::MultiRegion)
465 );
466 assert_eq!(
467 ServerGroupStrategy::parse("HA"),
468 Some(ServerGroupStrategy::HighAvailability)
469 );
470 }
471
472 #[test]
473 fn test_load_balance_strategy_from_str() {
474 assert_eq!(
475 LoadBalanceStrategy::parse("RoundRobin"),
476 Some(LoadBalanceStrategy::RoundRobin)
477 );
478 assert_eq!(
479 LoadBalanceStrategy::parse("rr"),
480 Some(LoadBalanceStrategy::RoundRobin)
481 );
482 assert_eq!(
483 LoadBalanceStrategy::parse("weighted"),
484 Some(LoadBalanceStrategy::Weighted)
485 );
486 assert_eq!(
487 LoadBalanceStrategy::parse("nearest"),
488 Some(LoadBalanceStrategy::Nearest)
489 );
490 }
491
492 #[test]
493 fn test_server_property_value_display() {
494 assert_eq!(
495 ServerPropertyValue::String("test".to_string()).to_string(),
496 "\"test\""
497 );
498 assert_eq!(ServerPropertyValue::Number(42.0).to_string(), "42");
499 assert_eq!(ServerPropertyValue::Boolean(true).to_string(), "true");
500 assert_eq!(
501 ServerPropertyValue::Identifier("primary".to_string()).to_string(),
502 "primary"
503 );
504 assert_eq!(
505 ServerPropertyValue::EnvVar("DATABASE_URL".to_string()).to_string(),
506 "env(\"DATABASE_URL\")"
507 );
508 }
509
510 fn test_span() -> Span {
511 Span::new(0, 0)
512 }
513
514 #[test]
515 fn test_server_group_primary_and_replicas() {
516 let mut group = ServerGroup::new(Ident::new("TestCluster", test_span()), test_span());
517
518 let mut primary = Server::new(Ident::new("primary", test_span()), test_span());
519 primary.add_property(ServerProperty::new(
520 "role",
521 ServerPropertyValue::Identifier("primary".to_string()),
522 test_span(),
523 ));
524 group.add_server(primary);
525
526 let mut replica1 = Server::new(Ident::new("replica1", test_span()), test_span());
527 replica1.add_property(ServerProperty::new(
528 "role",
529 ServerPropertyValue::Identifier("replica".to_string()),
530 test_span(),
531 ));
532 group.add_server(replica1);
533
534 let mut replica2 = Server::new(Ident::new("replica2", test_span()), test_span());
535 replica2.add_property(ServerProperty::new(
536 "role",
537 ServerPropertyValue::Identifier("replica".to_string()),
538 test_span(),
539 ));
540 group.add_server(replica2);
541
542 assert!(group.primary().is_some());
543 assert_eq!(group.primary().unwrap().name.name.as_str(), "primary");
544 assert_eq!(group.replicas().len(), 2);
545 }
546}