prax_schema/ast/
server_group.rs

1//! Server group definitions for multi-server database configurations.
2//!
3//! Server groups allow organizing multiple database servers for:
4//! - Read replicas (load balancing reads)
5//! - Sharding (horizontal scaling)
6//! - Multi-region deployment
7//! - Failover and high availability
8
9use indexmap::IndexMap;
10use serde::{Deserialize, Serialize};
11use smol_str::SmolStr;
12
13use super::{Attribute, Documentation, Ident, Span};
14
15/// A server group containing multiple database servers.
16///
17/// # Example
18/// ```prax
19/// serverGroup MainCluster {
20///     @@strategy(ReadReplica)
21///     @@loadBalance(RoundRobin)
22///
23///     server primary {
24///         url = env("PRIMARY_DATABASE_URL")
25///         role = "primary"
26///         weight = 1
27///     }
28///
29///     server replica1 {
30///         url = env("REPLICA1_DATABASE_URL")
31///         role = "replica"
32///         weight = 2
33///         region = "us-east-1"
34///     }
35///
36///     server replica2 {
37///         url = env("REPLICA2_DATABASE_URL")
38///         role = "replica"
39///         weight = 2
40///         region = "us-west-2"
41///     }
42/// }
43/// ```
44#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
45pub struct ServerGroup {
46    /// Group name.
47    pub name: Ident,
48    /// Servers in this group.
49    pub servers: IndexMap<SmolStr, Server>,
50    /// Group-level attributes (strategy, load balancing, etc.).
51    pub attributes: Vec<Attribute>,
52    /// Documentation comment.
53    pub documentation: Option<Documentation>,
54    /// Source location.
55    pub span: Span,
56}
57
58impl ServerGroup {
59    /// Create a new server group.
60    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    /// Add a server to the group.
71    pub fn add_server(&mut self, server: Server) {
72        self.servers.insert(server.name.name.clone(), server);
73    }
74
75    /// Add an attribute to the group.
76    pub fn add_attribute(&mut self, attr: Attribute) {
77        self.attributes.push(attr);
78    }
79
80    /// Set documentation.
81    pub fn set_documentation(&mut self, doc: Documentation) {
82        self.documentation = Some(doc);
83    }
84
85    /// Get the group strategy (e.g., ReadReplica, Sharding, MultiRegion).
86    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    /// Get the load balancing strategy.
103    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    /// Get the primary server.
120    pub fn primary(&self) -> Option<&Server> {
121        self.servers
122            .values()
123            .find(|s| s.role() == Some(ServerRole::Primary))
124    }
125
126    /// Get all replica servers.
127    pub fn replicas(&self) -> Vec<&Server> {
128        self.servers
129            .values()
130            .filter(|s| s.role() == Some(ServerRole::Replica))
131            .collect()
132    }
133
134    /// Get servers by region.
135    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    /// Get the failover order (sorted by priority).
143    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/// An individual server within a server group.
151#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
152pub struct Server {
153    /// Server name.
154    pub name: Ident,
155    /// Server properties (url, role, weight, etc.).
156    pub properties: IndexMap<SmolStr, ServerProperty>,
157    /// Source location.
158    pub span: Span,
159}
160
161impl Server {
162    /// Create a new server.
163    pub fn new(name: Ident, span: Span) -> Self {
164        Self {
165            name,
166            properties: IndexMap::new(),
167            span,
168        }
169    }
170
171    /// Add a property to the server.
172    pub fn add_property(&mut self, prop: ServerProperty) {
173        self.properties.insert(prop.name.clone(), prop);
174    }
175
176    /// Get a property value by name.
177    pub fn get_property(&self, name: &str) -> Option<&ServerPropertyValue> {
178        self.properties.get(name).map(|p| &p.value)
179    }
180
181    /// Get the server URL.
182    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    /// Get the server role.
191    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    /// Get the server weight (for load balancing).
201    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    /// Get the server region.
209    pub fn region(&self) -> Option<&str> {
210        match self.get_property("region")? {
211            ServerPropertyValue::String(s) => Some(s),
212            _ => None,
213        }
214    }
215
216    /// Get the server priority (for failover).
217    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    /// Check if this is a read-only server.
225    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    /// Get the maximum connections for this server.
233    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    /// Get the health check endpoint.
241    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/// A server property key-value pair.
250#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
251pub struct ServerProperty {
252    /// Property name.
253    pub name: SmolStr,
254    /// Property value.
255    pub value: ServerPropertyValue,
256    /// Source location.
257    pub span: Span,
258}
259
260impl ServerProperty {
261    /// Create a new server property.
262    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/// Server property value types.
272#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
273pub enum ServerPropertyValue {
274    /// String value.
275    String(String),
276    /// Number value.
277    Number(f64),
278    /// Boolean value.
279    Boolean(bool),
280    /// Identifier (enum-like value).
281    Identifier(String),
282    /// Environment variable reference.
283    EnvVar(String),
284    /// Array of values.
285    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/// Server role within a group.
311#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
312pub enum ServerRole {
313    /// Primary/master server (handles writes).
314    Primary,
315    /// Replica/slave server (handles reads).
316    Replica,
317    /// Analytics server (for reporting queries).
318    Analytics,
319    /// Archive server (for historical data).
320    Archive,
321    /// Shard server (for horizontal partitioning).
322    Shard,
323}
324
325impl ServerRole {
326    /// Parse role from string.
327    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    /// Get the role as a string.
339    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/// Server group strategy.
351#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
352pub enum ServerGroupStrategy {
353    /// Read replica configuration (primary + replicas).
354    ReadReplica,
355    /// Sharding configuration (horizontal partitioning).
356    Sharding,
357    /// Multi-region deployment.
358    MultiRegion,
359    /// High availability with automatic failover.
360    HighAvailability,
361    /// Custom strategy.
362    Custom,
363}
364
365impl ServerGroupStrategy {
366    /// Parse strategy from string.
367    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    /// Get the strategy as a string.
379    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/// Load balancing strategy for distributing queries.
391#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
392pub enum LoadBalanceStrategy {
393    /// Round-robin distribution.
394    RoundRobin,
395    /// Random selection.
396    Random,
397    /// Least connections.
398    LeastConnections,
399    /// Weighted distribution based on server weights.
400    Weighted,
401    /// Route to nearest (by latency or region).
402    Nearest,
403    /// Sticky sessions (same client to same server).
404    Sticky,
405}
406
407impl LoadBalanceStrategy {
408    /// Parse strategy from string.
409    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    /// Get the strategy as a string.
422    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!(ServerRole::parse("analytics"), Some(ServerRole::Analytics));
445        assert_eq!(ServerRole::parse("shard"), Some(ServerRole::Shard));
446        assert_eq!(ServerRole::parse("invalid"), None);
447    }
448
449    #[test]
450    fn test_server_group_strategy_from_str() {
451        assert_eq!(
452            ServerGroupStrategy::parse("ReadReplica"),
453            Some(ServerGroupStrategy::ReadReplica)
454        );
455        assert_eq!(
456            ServerGroupStrategy::parse("sharding"),
457            Some(ServerGroupStrategy::Sharding)
458        );
459        assert_eq!(
460            ServerGroupStrategy::parse("multi-region"),
461            Some(ServerGroupStrategy::MultiRegion)
462        );
463        assert_eq!(
464            ServerGroupStrategy::parse("HA"),
465            Some(ServerGroupStrategy::HighAvailability)
466        );
467    }
468
469    #[test]
470    fn test_load_balance_strategy_from_str() {
471        assert_eq!(
472            LoadBalanceStrategy::parse("RoundRobin"),
473            Some(LoadBalanceStrategy::RoundRobin)
474        );
475        assert_eq!(
476            LoadBalanceStrategy::parse("rr"),
477            Some(LoadBalanceStrategy::RoundRobin)
478        );
479        assert_eq!(
480            LoadBalanceStrategy::parse("weighted"),
481            Some(LoadBalanceStrategy::Weighted)
482        );
483        assert_eq!(
484            LoadBalanceStrategy::parse("nearest"),
485            Some(LoadBalanceStrategy::Nearest)
486        );
487    }
488
489    #[test]
490    fn test_server_property_value_display() {
491        assert_eq!(
492            ServerPropertyValue::String("test".to_string()).to_string(),
493            "\"test\""
494        );
495        assert_eq!(ServerPropertyValue::Number(42.0).to_string(), "42");
496        assert_eq!(ServerPropertyValue::Boolean(true).to_string(), "true");
497        assert_eq!(
498            ServerPropertyValue::Identifier("primary".to_string()).to_string(),
499            "primary"
500        );
501        assert_eq!(
502            ServerPropertyValue::EnvVar("DATABASE_URL".to_string()).to_string(),
503            "env(\"DATABASE_URL\")"
504        );
505    }
506
507    fn test_span() -> Span {
508        Span::new(0, 0)
509    }
510
511    #[test]
512    fn test_server_group_primary_and_replicas() {
513        let mut group = ServerGroup::new(Ident::new("TestCluster", test_span()), test_span());
514
515        let mut primary = Server::new(Ident::new("primary", test_span()), test_span());
516        primary.add_property(ServerProperty::new(
517            "role",
518            ServerPropertyValue::Identifier("primary".to_string()),
519            test_span(),
520        ));
521        group.add_server(primary);
522
523        let mut replica1 = Server::new(Ident::new("replica1", test_span()), test_span());
524        replica1.add_property(ServerProperty::new(
525            "role",
526            ServerPropertyValue::Identifier("replica".to_string()),
527            test_span(),
528        ));
529        group.add_server(replica1);
530
531        let mut replica2 = Server::new(Ident::new("replica2", test_span()), test_span());
532        replica2.add_property(ServerProperty::new(
533            "role",
534            ServerPropertyValue::Identifier("replica".to_string()),
535            test_span(),
536        ));
537        group.add_server(replica2);
538
539        assert!(group.primary().is_some());
540        assert_eq!(group.primary().unwrap().name.name.as_str(), "primary");
541        assert_eq!(group.replicas().len(), 2);
542    }
543}