Skip to main content

prax_schema/ast/
schema.rs

1//! Top-level schema definition.
2
3use indexmap::IndexMap;
4use serde::{Deserialize, Serialize};
5use smol_str::SmolStr;
6
7use super::procedure::Procedure;
8use super::{
9    CompositeType, Datasource, Enum, Generator, Model, Policy, Relation, ServerGroup, View,
10};
11
12/// A complete Prax schema.
13#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
14pub struct Schema {
15    /// Datasource configuration (database connection and extensions).
16    pub datasource: Option<Datasource>,
17    /// Generator configurations.
18    pub generators: IndexMap<SmolStr, Generator>,
19    /// All models in the schema.
20    pub models: IndexMap<SmolStr, Model>,
21    /// All enums in the schema.
22    pub enums: IndexMap<SmolStr, Enum>,
23    /// All composite types in the schema.
24    pub types: IndexMap<SmolStr, CompositeType>,
25    /// All views in the schema.
26    pub views: IndexMap<SmolStr, View>,
27    /// Server groups for multi-server configurations.
28    pub server_groups: IndexMap<SmolStr, ServerGroup>,
29    /// PostgreSQL Row-Level Security policies.
30    pub policies: Vec<Policy>,
31    /// Stored procedures and functions.
32    pub procedures: IndexMap<SmolStr, Procedure>,
33    /// Raw SQL definitions.
34    pub raw_sql: Vec<RawSql>,
35    /// Resolved relations (populated after validation).
36    pub relations: Vec<Relation>,
37}
38
39impl Schema {
40    /// Create a new empty schema.
41    pub fn new() -> Self {
42        Self::default()
43    }
44
45    /// Set the datasource configuration.
46    pub fn set_datasource(&mut self, datasource: Datasource) {
47        self.datasource = Some(datasource);
48    }
49
50    /// Get the datasource configuration.
51    pub fn datasource(&self) -> Option<&Datasource> {
52        self.datasource.as_ref()
53    }
54
55    /// Check if the schema has vector extension enabled.
56    pub fn has_vector_support(&self) -> bool {
57        self.datasource
58            .as_ref()
59            .is_some_and(|ds| ds.has_vector_support())
60    }
61
62    /// Get all required PostgreSQL extensions from the datasource.
63    pub fn required_extensions(&self) -> Vec<&super::PostgresExtension> {
64        self.datasource
65            .as_ref()
66            .map(|ds| ds.extensions.iter().collect())
67            .unwrap_or_default()
68    }
69
70    /// Add a generator to the schema.
71    pub fn add_generator(&mut self, generator: Generator) {
72        self.generators.insert(generator.name.clone(), generator);
73    }
74
75    /// Get a generator by name.
76    pub fn get_generator(&self, name: &str) -> Option<&Generator> {
77        self.generators.get(name)
78    }
79
80    /// Get all enabled generators.
81    pub fn enabled_generators(&self) -> Vec<&Generator> {
82        self.generators
83            .values()
84            .filter(|g| g.is_enabled())
85            .collect()
86    }
87
88    /// Add a model to the schema.
89    pub fn add_model(&mut self, model: Model) {
90        self.models.insert(model.name.name.clone(), model);
91    }
92
93    /// Add an enum to the schema.
94    pub fn add_enum(&mut self, e: Enum) {
95        self.enums.insert(e.name.name.clone(), e);
96    }
97
98    /// Add a composite type to the schema.
99    pub fn add_type(&mut self, t: CompositeType) {
100        self.types.insert(t.name.name.clone(), t);
101    }
102
103    /// Add a view to the schema.
104    pub fn add_view(&mut self, v: View) {
105        self.views.insert(v.name.name.clone(), v);
106    }
107
108    /// Add a server group to the schema.
109    pub fn add_server_group(&mut self, sg: ServerGroup) {
110        self.server_groups.insert(sg.name.name.clone(), sg);
111    }
112
113    /// Add a PostgreSQL Row-Level Security policy.
114    pub fn add_policy(&mut self, policy: Policy) {
115        self.policies.push(policy);
116    }
117
118    /// Add a raw SQL definition.
119    pub fn add_raw_sql(&mut self, sql: RawSql) {
120        self.raw_sql.push(sql);
121    }
122
123    /// Get a model by name.
124    pub fn get_model(&self, name: &str) -> Option<&Model> {
125        self.models.get(name)
126    }
127
128    /// Get a mutable model by name.
129    pub fn get_model_mut(&mut self, name: &str) -> Option<&mut Model> {
130        self.models.get_mut(name)
131    }
132
133    /// Get an enum by name.
134    pub fn get_enum(&self, name: &str) -> Option<&Enum> {
135        self.enums.get(name)
136    }
137
138    /// Get a composite type by name.
139    pub fn get_type(&self, name: &str) -> Option<&CompositeType> {
140        self.types.get(name)
141    }
142
143    /// Get a view by name.
144    pub fn get_view(&self, name: &str) -> Option<&View> {
145        self.views.get(name)
146    }
147
148    /// Get a server group by name.
149    pub fn get_server_group(&self, name: &str) -> Option<&ServerGroup> {
150        self.server_groups.get(name)
151    }
152
153    /// Get all server group names.
154    pub fn server_group_names(&self) -> impl Iterator<Item = &str> {
155        self.server_groups.keys().map(|s| s.as_str())
156    }
157
158    /// Get a policy by name.
159    pub fn get_policy(&self, name: &str) -> Option<&Policy> {
160        self.policies.iter().find(|p| p.name() == name)
161    }
162
163    /// Get all policies for a specific model/table.
164    pub fn policies_for(&self, model: &str) -> Vec<&Policy> {
165        self.policies
166            .iter()
167            .filter(|p| p.table() == model)
168            .collect()
169    }
170
171    /// Check if a model has Row-Level Security policies.
172    pub fn has_policies(&self, model: &str) -> bool {
173        self.policies.iter().any(|p| p.table() == model)
174    }
175
176    /// Get all policy names.
177    pub fn policy_names(&self) -> impl Iterator<Item = &str> {
178        self.policies.iter().map(|p| p.name())
179    }
180
181    /// Check if a type name exists (model, enum, type, or view).
182    pub fn type_exists(&self, name: &str) -> bool {
183        self.models.contains_key(name)
184            || self.enums.contains_key(name)
185            || self.types.contains_key(name)
186            || self.views.contains_key(name)
187    }
188
189    /// Get all model names.
190    pub fn model_names(&self) -> impl Iterator<Item = &str> {
191        self.models.keys().map(|s| s.as_str())
192    }
193
194    /// Get all enum names.
195    pub fn enum_names(&self) -> impl Iterator<Item = &str> {
196        self.enums.keys().map(|s| s.as_str())
197    }
198
199    /// Get relations for a specific model.
200    pub fn relations_for(&self, model: &str) -> Vec<&Relation> {
201        self.relations
202            .iter()
203            .filter(|r| r.from_model == model || r.to_model == model)
204            .collect()
205    }
206
207    /// Get relations originating from a specific model.
208    pub fn relations_from(&self, model: &str) -> Vec<&Relation> {
209        self.relations
210            .iter()
211            .filter(|r| r.from_model == model)
212            .collect()
213    }
214
215    /// Merge another schema into this one (deprecated: use [`Schema::try_merge`]).
216    ///
217    /// Silently overwrites duplicates. Kept for backward compatibility; will be
218    /// removed in a future release.
219    #[deprecated(since = "0.9.8", note = "use try_merge for collision-aware merging")]
220    pub fn merge(&mut self, other: Schema) {
221        self.models.extend(other.models);
222        self.enums.extend(other.enums);
223        self.types.extend(other.types);
224        self.views.extend(other.views);
225        self.server_groups.extend(other.server_groups);
226        self.policies.extend(other.policies);
227        self.raw_sql.extend(other.raw_sql);
228    }
229
230    /// Merge `other` into `self`, returning every collision found rather than
231    /// silently overwriting.
232    ///
233    /// Items whose `source_id` is `None` (built outside the loader, e.g. in
234    /// tests) are reported with `SourceId(u32::MAX)`. In production usage the
235    /// loader stamps every item via [`crate::loader::stamp_source`] first.
236    pub fn try_merge(&mut self, other: Schema) -> Result<(), Vec<crate::loader::MergeConflict>> {
237        use crate::loader::MergeConflict;
238        use crate::loader::merge::loc;
239
240        let mut conflicts: Vec<MergeConflict> = Vec::new();
241
242        for (name, m) in other.models {
243            if let Some(existing) = self.models.get(&name) {
244                conflicts.push(MergeConflict::DuplicateModel {
245                    name: name.clone(),
246                    existing: loc(existing.source_id, existing.span),
247                    incoming: loc(m.source_id, m.span),
248                });
249            } else {
250                self.models.insert(name, m);
251            }
252        }
253
254        for (name, e) in other.enums {
255            if let Some(existing) = self.enums.get(&name) {
256                conflicts.push(MergeConflict::DuplicateEnum {
257                    name: name.clone(),
258                    existing: loc(existing.source_id, existing.span),
259                    incoming: loc(e.source_id, e.span),
260                });
261            } else {
262                self.enums.insert(name, e);
263            }
264        }
265
266        for (name, t) in other.types {
267            if let Some(existing) = self.types.get(&name) {
268                conflicts.push(MergeConflict::DuplicateType {
269                    name: name.clone(),
270                    existing: loc(existing.source_id, existing.span),
271                    incoming: loc(t.source_id, t.span),
272                });
273            } else {
274                self.types.insert(name, t);
275            }
276        }
277
278        for (name, v) in other.views {
279            if let Some(existing) = self.views.get(&name) {
280                conflicts.push(MergeConflict::DuplicateView {
281                    name: name.clone(),
282                    existing: loc(existing.source_id, existing.span),
283                    incoming: loc(v.source_id, v.span),
284                });
285            } else {
286                self.views.insert(name, v);
287            }
288        }
289
290        for (name, p) in other.procedures {
291            if let Some(existing) = self.procedures.get(&name) {
292                conflicts.push(MergeConflict::DuplicateProcedure {
293                    name: name.clone(),
294                    existing: loc(existing.source_id, existing.span),
295                    incoming: loc(p.source_id, p.span),
296                });
297            } else {
298                self.procedures.insert(name, p);
299            }
300        }
301
302        for (name, sg) in other.server_groups {
303            if let Some(existing) = self.server_groups.get(&name) {
304                conflicts.push(MergeConflict::DuplicateServerGroup {
305                    name: name.clone(),
306                    existing: loc(existing.source_id, existing.span),
307                    incoming: loc(sg.source_id, sg.span),
308                });
309            } else {
310                self.server_groups.insert(name, sg);
311            }
312        }
313
314        for (name, g) in other.generators {
315            if let Some(existing) = self.generators.get(&name) {
316                conflicts.push(MergeConflict::DuplicateGenerator {
317                    name: name.clone(),
318                    existing: loc(existing.source_id, existing.span),
319                    incoming: loc(g.source_id, g.span),
320                });
321            } else {
322                self.generators.insert(name, g);
323            }
324        }
325
326        for p in other.policies {
327            if let Some(existing) = self.policies.iter().find(|x| x.name() == p.name()) {
328                conflicts.push(MergeConflict::DuplicatePolicy {
329                    name: SmolStr::new(p.name()),
330                    existing: loc(existing.source_id, existing.span),
331                    incoming: loc(p.source_id, p.span),
332                });
333            } else {
334                self.policies.push(p);
335            }
336        }
337
338        for r in other.raw_sql {
339            if let Some(existing) = self.raw_sql.iter().find(|x| x.name == r.name) {
340                conflicts.push(MergeConflict::DuplicateRawSql {
341                    name: r.name.clone(),
342                    existing: loc(existing.source_id, super::Span::new(0, 0)),
343                    incoming: loc(r.source_id, super::Span::new(0, 0)),
344                });
345            } else {
346                self.raw_sql.push(r);
347            }
348        }
349
350        match (&self.datasource, other.datasource) {
351            (Some(existing), Some(incoming)) => {
352                conflicts.push(MergeConflict::MultipleDatasource {
353                    existing: loc(existing.source_id, existing.span),
354                    incoming: loc(incoming.source_id, incoming.span),
355                });
356            }
357            (None, Some(incoming)) => self.datasource = Some(incoming),
358            (_, None) => {}
359        }
360
361        if conflicts.is_empty() {
362            Ok(())
363        } else {
364            Err(conflicts)
365        }
366    }
367}
368
369/// A raw SQL definition.
370#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
371pub struct RawSql {
372    /// Name/identifier for the SQL (e.g., view name).
373    pub name: SmolStr,
374    /// The raw SQL content.
375    pub sql: String,
376    /// Source file this raw SQL was parsed from (None for single-file path).
377    #[serde(default, skip_serializing_if = "Option::is_none")]
378    pub source_id: Option<crate::loader::SourceId>,
379}
380
381impl RawSql {
382    /// Create a new raw SQL definition.
383    pub fn new(name: impl Into<SmolStr>, sql: impl Into<String>) -> Self {
384        Self {
385            name: name.into(),
386            sql: sql.into(),
387            source_id: None,
388        }
389    }
390}
391
392/// Schema statistics for debugging/info.
393#[derive(Debug, Clone, Default)]
394pub struct SchemaStats {
395    /// Number of models.
396    pub model_count: usize,
397    /// Number of enums.
398    pub enum_count: usize,
399    /// Number of composite types.
400    pub type_count: usize,
401    /// Number of views.
402    pub view_count: usize,
403    /// Number of server groups.
404    pub server_group_count: usize,
405    /// Number of procedures.
406    pub procedure_count: usize,
407    /// Number of RLS policies.
408    pub policy_count: usize,
409    /// Total number of fields across all models.
410    pub field_count: usize,
411    /// Number of relations.
412    pub relation_count: usize,
413}
414
415impl Schema {
416    /// Get statistics about the schema.
417    pub fn stats(&self) -> SchemaStats {
418        SchemaStats {
419            model_count: self.models.len(),
420            enum_count: self.enums.len(),
421            type_count: self.types.len(),
422            view_count: self.views.len(),
423            server_group_count: self.server_groups.len(),
424            procedure_count: self.procedures.len(),
425            policy_count: self.policies.len(),
426            field_count: self.models.values().map(|m| m.fields.len()).sum(),
427            relation_count: self.relations.len(),
428        }
429    }
430}
431
432impl std::fmt::Display for Schema {
433    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
434        let stats = self.stats();
435        write!(
436            f,
437            "Schema({} models, {} enums, {} types, {} views, {} procedures, {} server groups, {} policies, {} fields, {} relations)",
438            stats.model_count,
439            stats.enum_count,
440            stats.type_count,
441            stats.view_count,
442            stats.procedure_count,
443            stats.server_group_count,
444            stats.policy_count,
445            stats.field_count,
446            stats.relation_count
447        )
448    }
449}
450
451#[cfg(test)]
452mod tests {
453    use super::*;
454    use crate::ast::{
455        Attribute, EnumVariant, Field, FieldType, Ident, Policy, RelationType, ScalarType, Span,
456        TypeModifier,
457    };
458
459    fn make_span() -> Span {
460        Span::new(0, 10)
461    }
462
463    fn make_ident(name: &str) -> Ident {
464        Ident::new(name, make_span())
465    }
466
467    fn make_model(name: &str) -> Model {
468        let mut model = Model::new(make_ident(name), make_span());
469        let id_field = make_id_field();
470        model.add_field(id_field);
471        model
472    }
473
474    fn make_id_field() -> Field {
475        let mut field = Field::new(
476            make_ident("id"),
477            FieldType::Scalar(ScalarType::Int),
478            TypeModifier::Required,
479            vec![],
480            make_span(),
481        );
482        field
483            .attributes
484            .push(Attribute::simple(make_ident("id"), make_span()));
485        field
486    }
487
488    fn make_field(name: &str, field_type: FieldType) -> Field {
489        Field::new(
490            make_ident(name),
491            field_type,
492            TypeModifier::Required,
493            vec![],
494            make_span(),
495        )
496    }
497
498    fn make_enum(name: &str, variants: &[&str]) -> Enum {
499        let mut e = Enum::new(make_ident(name), make_span());
500        for v in variants {
501            e.add_variant(EnumVariant::new(make_ident(v), make_span()));
502        }
503        e
504    }
505
506    // ==================== Schema Tests ====================
507
508    #[test]
509    fn test_schema_new() {
510        let schema = Schema::new();
511        assert!(schema.models.is_empty());
512        assert!(schema.enums.is_empty());
513        assert!(schema.types.is_empty());
514        assert!(schema.views.is_empty());
515        assert!(schema.policies.is_empty());
516        assert!(schema.raw_sql.is_empty());
517        assert!(schema.relations.is_empty());
518    }
519
520    #[test]
521    fn test_schema_default() {
522        let schema = Schema::default();
523        assert!(schema.models.is_empty());
524    }
525
526    #[test]
527    fn test_schema_add_model() {
528        let mut schema = Schema::new();
529        let model = make_model("User");
530
531        schema.add_model(model);
532
533        assert_eq!(schema.models.len(), 1);
534        assert!(schema.models.contains_key("User"));
535    }
536
537    #[test]
538    fn test_schema_add_multiple_models() {
539        let mut schema = Schema::new();
540        schema.add_model(make_model("User"));
541        schema.add_model(make_model("Post"));
542        schema.add_model(make_model("Comment"));
543
544        assert_eq!(schema.models.len(), 3);
545    }
546
547    #[test]
548    fn test_schema_add_enum() {
549        let mut schema = Schema::new();
550        let e = make_enum("Role", &["User", "Admin"]);
551
552        schema.add_enum(e);
553
554        assert_eq!(schema.enums.len(), 1);
555        assert!(schema.enums.contains_key("Role"));
556    }
557
558    #[test]
559    fn test_schema_add_type() {
560        let mut schema = Schema::new();
561        let ct = CompositeType::new(make_ident("Address"), make_span());
562
563        schema.add_type(ct);
564
565        assert_eq!(schema.types.len(), 1);
566        assert!(schema.types.contains_key("Address"));
567    }
568
569    #[test]
570    fn test_schema_add_view() {
571        let mut schema = Schema::new();
572        let view = View::new(make_ident("UserStats"), make_span());
573
574        schema.add_view(view);
575
576        assert_eq!(schema.views.len(), 1);
577        assert!(schema.views.contains_key("UserStats"));
578    }
579
580    #[test]
581    fn test_schema_add_raw_sql() {
582        let mut schema = Schema::new();
583        let sql = RawSql::new("migration_1", "CREATE TABLE test ();");
584
585        schema.add_raw_sql(sql);
586
587        assert_eq!(schema.raw_sql.len(), 1);
588    }
589
590    #[test]
591    fn test_schema_get_model() {
592        let mut schema = Schema::new();
593        schema.add_model(make_model("User"));
594
595        let model = schema.get_model("User");
596        assert!(model.is_some());
597        assert_eq!(model.unwrap().name(), "User");
598
599        assert!(schema.get_model("NonExistent").is_none());
600    }
601
602    #[test]
603    fn test_schema_get_model_mut() {
604        let mut schema = Schema::new();
605        schema.add_model(make_model("User"));
606
607        let model = schema.get_model_mut("User");
608        assert!(model.is_some());
609
610        // Modify the model
611        let model = model.unwrap();
612        model.add_field(make_field("email", FieldType::Scalar(ScalarType::String)));
613
614        // Verify modification persisted
615        assert_eq!(schema.get_model("User").unwrap().fields.len(), 2);
616    }
617
618    #[test]
619    fn test_schema_get_enum() {
620        let mut schema = Schema::new();
621        schema.add_enum(make_enum("Role", &["User", "Admin"]));
622
623        let e = schema.get_enum("Role");
624        assert!(e.is_some());
625        assert_eq!(e.unwrap().name(), "Role");
626
627        assert!(schema.get_enum("NonExistent").is_none());
628    }
629
630    #[test]
631    fn test_schema_get_type() {
632        let mut schema = Schema::new();
633        schema.add_type(CompositeType::new(make_ident("Address"), make_span()));
634
635        let ct = schema.get_type("Address");
636        assert!(ct.is_some());
637
638        assert!(schema.get_type("NonExistent").is_none());
639    }
640
641    #[test]
642    fn test_schema_get_view() {
643        let mut schema = Schema::new();
644        schema.add_view(View::new(make_ident("Stats"), make_span()));
645
646        let v = schema.get_view("Stats");
647        assert!(v.is_some());
648
649        assert!(schema.get_view("NonExistent").is_none());
650    }
651
652    #[test]
653    fn test_schema_type_exists() {
654        let mut schema = Schema::new();
655        schema.add_model(make_model("User"));
656        schema.add_enum(make_enum("Role", &["User"]));
657        schema.add_type(CompositeType::new(make_ident("Address"), make_span()));
658        schema.add_view(View::new(make_ident("Stats"), make_span()));
659
660        assert!(schema.type_exists("User")); // model
661        assert!(schema.type_exists("Role")); // enum
662        assert!(schema.type_exists("Address")); // type
663        assert!(schema.type_exists("Stats")); // view
664        assert!(!schema.type_exists("NonExistent"));
665    }
666
667    #[test]
668    fn test_schema_model_names() {
669        let mut schema = Schema::new();
670        schema.add_model(make_model("User"));
671        schema.add_model(make_model("Post"));
672
673        let names: Vec<_> = schema.model_names().collect();
674        assert_eq!(names.len(), 2);
675        assert!(names.contains(&"User"));
676        assert!(names.contains(&"Post"));
677    }
678
679    #[test]
680    fn test_schema_enum_names() {
681        let mut schema = Schema::new();
682        schema.add_enum(make_enum("Role", &["User"]));
683        schema.add_enum(make_enum("Status", &["Active"]));
684
685        let names: Vec<_> = schema.enum_names().collect();
686        assert_eq!(names.len(), 2);
687        assert!(names.contains(&"Role"));
688        assert!(names.contains(&"Status"));
689    }
690
691    #[test]
692    fn test_schema_relations_for() {
693        let mut schema = Schema::new();
694        schema.relations.push(Relation::new(
695            "Post",
696            "author",
697            "User",
698            RelationType::ManyToOne,
699        ));
700        schema.relations.push(Relation::new(
701            "Comment",
702            "user",
703            "User",
704            RelationType::ManyToOne,
705        ));
706        schema.relations.push(Relation::new(
707            "Post",
708            "tags",
709            "Tag",
710            RelationType::ManyToMany,
711        ));
712
713        let user_relations = schema.relations_for("User");
714        assert_eq!(user_relations.len(), 2);
715
716        let post_relations = schema.relations_for("Post");
717        assert_eq!(post_relations.len(), 2);
718
719        let tag_relations = schema.relations_for("Tag");
720        assert_eq!(tag_relations.len(), 1);
721    }
722
723    #[test]
724    fn test_schema_relations_from() {
725        let mut schema = Schema::new();
726        schema.relations.push(Relation::new(
727            "Post",
728            "author",
729            "User",
730            RelationType::ManyToOne,
731        ));
732        schema.relations.push(Relation::new(
733            "Post",
734            "tags",
735            "Tag",
736            RelationType::ManyToMany,
737        ));
738        schema.relations.push(Relation::new(
739            "User",
740            "posts",
741            "Post",
742            RelationType::OneToMany,
743        ));
744
745        let post_relations = schema.relations_from("Post");
746        assert_eq!(post_relations.len(), 2);
747
748        let user_relations = schema.relations_from("User");
749        assert_eq!(user_relations.len(), 1);
750
751        let tag_relations = schema.relations_from("Tag");
752        assert_eq!(tag_relations.len(), 0);
753    }
754
755    #[test]
756    #[allow(deprecated)]
757    fn test_schema_merge() {
758        let mut schema1 = Schema::new();
759        schema1.add_model(make_model("User"));
760        schema1.add_enum(make_enum("Role", &["User"]));
761
762        let mut schema2 = Schema::new();
763        schema2.add_model(make_model("Post"));
764        schema2.add_enum(make_enum("Status", &["Active"]));
765        schema2.add_raw_sql(RawSql::new("init", "-- init"));
766
767        schema1.merge(schema2);
768
769        assert_eq!(schema1.models.len(), 2);
770        assert_eq!(schema1.enums.len(), 2);
771        assert_eq!(schema1.raw_sql.len(), 1);
772    }
773
774    #[test]
775    fn test_schema_stats() {
776        let mut schema = Schema::new();
777
778        let mut user = make_model("User");
779        user.add_field(make_field("email", FieldType::Scalar(ScalarType::String)));
780        user.add_field(make_field("name", FieldType::Scalar(ScalarType::String)));
781        schema.add_model(user);
782
783        let mut post = make_model("Post");
784        post.add_field(make_field("title", FieldType::Scalar(ScalarType::String)));
785        schema.add_model(post);
786
787        schema.add_enum(make_enum("Role", &["User", "Admin"]));
788        schema.add_type(CompositeType::new(make_ident("Address"), make_span()));
789        schema.add_view(View::new(make_ident("Stats"), make_span()));
790        schema.relations.push(Relation::new(
791            "Post",
792            "author",
793            "User",
794            RelationType::ManyToOne,
795        ));
796
797        let stats = schema.stats();
798        assert_eq!(stats.model_count, 2);
799        assert_eq!(stats.enum_count, 1);
800        assert_eq!(stats.type_count, 1);
801        assert_eq!(stats.view_count, 1);
802        assert_eq!(stats.field_count, 5); // 3 in User + 2 in Post
803        assert_eq!(stats.relation_count, 1);
804    }
805
806    #[test]
807    fn test_schema_display() {
808        let mut schema = Schema::new();
809        schema.add_model(make_model("User"));
810        schema.add_enum(make_enum("Role", &["User"]));
811
812        let display = format!("{}", schema);
813        assert!(display.contains("1 models"));
814        assert!(display.contains("1 enums"));
815        assert!(display.contains("0 policies"));
816    }
817
818    #[test]
819    fn test_schema_equality() {
820        let schema1 = Schema::new();
821        let schema2 = Schema::new();
822        assert_eq!(schema1, schema2);
823    }
824
825    #[test]
826    fn test_schema_clone() {
827        let mut schema = Schema::new();
828        schema.add_model(make_model("User"));
829
830        let cloned = schema.clone();
831        assert_eq!(cloned.models.len(), 1);
832    }
833
834    // ==================== RawSql Tests ====================
835
836    #[test]
837    fn test_raw_sql_new() {
838        let sql = RawSql::new("create_users", "CREATE TABLE users ();");
839
840        assert_eq!(sql.name.as_str(), "create_users");
841        assert_eq!(sql.sql, "CREATE TABLE users ();");
842    }
843
844    #[test]
845    fn test_raw_sql_from_strings() {
846        let name = String::from("migration");
847        let content = String::from("ALTER TABLE users ADD COLUMN age INT;");
848        let sql = RawSql::new(name, content);
849
850        assert_eq!(sql.name.as_str(), "migration");
851    }
852
853    #[test]
854    fn test_raw_sql_equality() {
855        let sql1 = RawSql::new("test", "SELECT 1;");
856        let sql2 = RawSql::new("test", "SELECT 1;");
857        let sql3 = RawSql::new("test", "SELECT 2;");
858
859        assert_eq!(sql1, sql2);
860        assert_ne!(sql1, sql3);
861    }
862
863    #[test]
864    fn test_raw_sql_clone() {
865        let sql = RawSql::new("test", "SELECT 1;");
866        let cloned = sql.clone();
867        assert_eq!(sql, cloned);
868    }
869
870    // ==================== SchemaStats Tests ====================
871
872    #[test]
873    fn test_schema_stats_default() {
874        let stats = SchemaStats::default();
875        assert_eq!(stats.model_count, 0);
876        assert_eq!(stats.enum_count, 0);
877        assert_eq!(stats.type_count, 0);
878        assert_eq!(stats.view_count, 0);
879        assert_eq!(stats.policy_count, 0);
880        assert_eq!(stats.field_count, 0);
881        assert_eq!(stats.relation_count, 0);
882    }
883
884    #[test]
885    fn test_schema_stats_debug() {
886        let stats = SchemaStats::default();
887        let debug = format!("{:?}", stats);
888        assert!(debug.contains("SchemaStats"));
889    }
890
891    #[test]
892    fn test_schema_stats_clone() {
893        let stats = SchemaStats {
894            model_count: 5,
895            enum_count: 2,
896            type_count: 1,
897            view_count: 3,
898            procedure_count: 0,
899            server_group_count: 2,
900            policy_count: 4,
901            field_count: 25,
902            relation_count: 10,
903        };
904        let cloned = stats.clone();
905        assert_eq!(cloned.model_count, 5);
906        assert_eq!(cloned.field_count, 25);
907        assert_eq!(cloned.policy_count, 4);
908    }
909
910    // ==================== Policy Schema Tests ====================
911
912    #[test]
913    fn test_schema_add_policy() {
914        let mut schema = Schema::new();
915        let policy = Policy::new(make_ident("read_own"), make_ident("User"), make_span());
916
917        schema.add_policy(policy);
918
919        assert_eq!(schema.policies.len(), 1);
920    }
921
922    #[test]
923    fn test_schema_get_policy() {
924        let mut schema = Schema::new();
925        schema.add_policy(Policy::new(
926            make_ident("read_own"),
927            make_ident("User"),
928            make_span(),
929        ));
930
931        let policy = schema.get_policy("read_own");
932        assert!(policy.is_some());
933        assert_eq!(policy.unwrap().name(), "read_own");
934
935        assert!(schema.get_policy("nonexistent").is_none());
936    }
937
938    #[test]
939    fn test_schema_policies_for_model() {
940        let mut schema = Schema::new();
941        schema.add_policy(Policy::new(
942            make_ident("user_read"),
943            make_ident("User"),
944            make_span(),
945        ));
946        schema.add_policy(Policy::new(
947            make_ident("user_write"),
948            make_ident("User"),
949            make_span(),
950        ));
951        schema.add_policy(Policy::new(
952            make_ident("post_read"),
953            make_ident("Post"),
954            make_span(),
955        ));
956
957        let user_policies = schema.policies_for("User");
958        assert_eq!(user_policies.len(), 2);
959
960        let post_policies = schema.policies_for("Post");
961        assert_eq!(post_policies.len(), 1);
962
963        let comment_policies = schema.policies_for("Comment");
964        assert!(comment_policies.is_empty());
965    }
966
967    #[test]
968    fn test_schema_has_policies() {
969        let mut schema = Schema::new();
970        schema.add_policy(Policy::new(
971            make_ident("test"),
972            make_ident("User"),
973            make_span(),
974        ));
975
976        assert!(schema.has_policies("User"));
977        assert!(!schema.has_policies("Post"));
978    }
979
980    #[test]
981    fn test_schema_policy_names() {
982        let mut schema = Schema::new();
983        schema.add_policy(Policy::new(
984            make_ident("policy1"),
985            make_ident("User"),
986            make_span(),
987        ));
988        schema.add_policy(Policy::new(
989            make_ident("policy2"),
990            make_ident("Post"),
991            make_span(),
992        ));
993
994        let names: Vec<_> = schema.policy_names().collect();
995        assert_eq!(names.len(), 2);
996        assert!(names.contains(&"policy1"));
997        assert!(names.contains(&"policy2"));
998    }
999
1000    #[test]
1001    #[allow(deprecated)]
1002    fn test_schema_merge_with_policies() {
1003        let mut schema1 = Schema::new();
1004        schema1.add_policy(Policy::new(
1005            make_ident("policy1"),
1006            make_ident("User"),
1007            make_span(),
1008        ));
1009
1010        let mut schema2 = Schema::new();
1011        schema2.add_policy(Policy::new(
1012            make_ident("policy2"),
1013            make_ident("Post"),
1014            make_span(),
1015        ));
1016
1017        schema1.merge(schema2);
1018
1019        assert_eq!(schema1.policies.len(), 2);
1020    }
1021
1022    #[test]
1023    fn test_schema_stats_with_policies() {
1024        let mut schema = Schema::new();
1025        schema.add_model(make_model("User"));
1026        schema.add_policy(Policy::new(
1027            make_ident("policy1"),
1028            make_ident("User"),
1029            make_span(),
1030        ));
1031        schema.add_policy(Policy::new(
1032            make_ident("policy2"),
1033            make_ident("User"),
1034            make_span(),
1035        ));
1036
1037        let stats = schema.stats();
1038        assert_eq!(stats.model_count, 1);
1039        assert_eq!(stats.policy_count, 2);
1040    }
1041}