postrust_core/schema_cache/
mod.rs

1//! PostgreSQL schema introspection and caching.
2//!
3//! This module provides functionality to discover and cache database schema
4//! metadata including tables, columns, relationships, and functions.
5
6mod table;
7mod relationship;
8mod routine;
9mod queries;
10
11pub use table::{Table, Column, ColumnMap, TablesMap};
12pub use relationship::{Relationship, Cardinality, Junction, RelationshipsMap};
13pub use routine::{Routine, RoutineParam, RetType, FuncVolatility, RoutineMap};
14
15use crate::api_request::QualifiedIdentifier;
16use crate::error::{Error, Result};
17use sqlx::PgPool;
18use std::collections::HashSet;
19use std::sync::Arc;
20use tracing::info;
21
22/// Cached PostgreSQL schema metadata.
23#[derive(Clone, Debug)]
24pub struct SchemaCache {
25    /// Tables and views by qualified identifier.
26    pub tables: TablesMap,
27    /// Relationships between tables.
28    pub relationships: RelationshipsMap,
29    /// Stored functions/procedures.
30    pub routines: RoutineMap,
31    /// Valid timezone names.
32    pub timezones: HashSet<String>,
33    /// PostgreSQL version.
34    pub pg_version: i32,
35}
36
37impl SchemaCache {
38    /// Load schema cache from the database.
39    pub async fn load(pool: &PgPool, schemas: &[String]) -> Result<Self> {
40        info!("Loading schema cache for schemas: {:?}", schemas);
41
42        // Get PostgreSQL version
43        let pg_version = queries::get_pg_version(pool).await?;
44        info!("PostgreSQL version: {}", pg_version);
45
46        // Load tables and columns
47        let tables = queries::load_tables(pool, schemas).await?;
48        info!("Loaded {} tables/views", tables.len());
49
50        // Load relationships
51        let relationships = queries::load_relationships(pool, schemas).await?;
52        info!("Loaded {} relationship sets", relationships.len());
53
54        // Load routines
55        let routines = queries::load_routines(pool, schemas).await?;
56        info!("Loaded {} routines", routines.len());
57
58        // Load timezone names
59        let timezones = queries::load_timezones(pool).await?;
60        info!("Loaded {} timezones", timezones.len());
61
62        Ok(Self {
63            tables,
64            relationships,
65            routines,
66            timezones,
67            pg_version,
68        })
69    }
70
71    /// Get a table by qualified identifier.
72    pub fn get_table(&self, qi: &QualifiedIdentifier) -> Option<&Table> {
73        self.tables.get(qi)
74    }
75
76    /// Get a table, returning an error if not found.
77    pub fn require_table(&self, qi: &QualifiedIdentifier) -> Result<&Table> {
78        self.get_table(qi)
79            .ok_or_else(|| Error::TableNotFound(qi.to_string()))
80    }
81
82    /// Get relationships for a table.
83    pub fn get_relationships(&self, qi: &QualifiedIdentifier, schema: &str) -> Option<&Vec<Relationship>> {
84        self.relationships.get(&(qi.clone(), schema.to_string()))
85    }
86
87    /// Get a routine by qualified identifier.
88    pub fn get_routines(&self, qi: &QualifiedIdentifier) -> Option<&Vec<Routine>> {
89        self.routines.get(qi)
90    }
91
92    /// Check if a timezone is valid.
93    pub fn is_valid_timezone(&self, tz: &str) -> bool {
94        self.timezones.contains(tz)
95    }
96
97    /// Get a summary of the cached schema.
98    pub fn summary(&self) -> String {
99        format!(
100            "SchemaCache: {} tables, {} relationship sets, {} routines, PG {}",
101            self.tables.len(),
102            self.relationships.len(),
103            self.routines.len(),
104            self.pg_version
105        )
106    }
107
108    /// Find a relationship between two tables by name.
109    pub fn find_relationship(
110        &self,
111        from: &QualifiedIdentifier,
112        to_name: &str,
113        schema: &str,
114    ) -> Option<&Relationship> {
115        self.get_relationships(from, schema)?
116            .iter()
117            .find(|r| match r {
118                Relationship::ForeignKey { foreign_table, .. } => {
119                    foreign_table.name == to_name
120                }
121                Relationship::Computed { foreign_table, .. } => {
122                    foreign_table.name == to_name
123                }
124            })
125    }
126}
127
128/// Thread-safe schema cache wrapper.
129#[derive(Clone)]
130pub struct SchemaCacheRef(Arc<tokio::sync::RwLock<Option<SchemaCache>>>);
131
132impl SchemaCacheRef {
133    /// Create a new empty schema cache reference.
134    pub fn new() -> Self {
135        Self(Arc::new(tokio::sync::RwLock::new(None)))
136    }
137
138    /// Load or reload the schema cache.
139    pub async fn load(&self, pool: &PgPool, schemas: &[String]) -> Result<()> {
140        let cache = SchemaCache::load(pool, schemas).await?;
141        let mut guard = self.0.write().await;
142        *guard = Some(cache);
143        Ok(())
144    }
145
146    /// Get a read reference to the schema cache.
147    pub async fn get(&self) -> Result<tokio::sync::RwLockReadGuard<'_, Option<SchemaCache>>> {
148        let guard = self.0.read().await;
149        if guard.is_none() {
150            return Err(Error::SchemaCacheNotLoaded);
151        }
152        Ok(guard)
153    }
154
155    /// Check if the cache is loaded.
156    pub async fn is_loaded(&self) -> bool {
157        self.0.read().await.is_some()
158    }
159}
160
161impl Default for SchemaCacheRef {
162    fn default() -> Self {
163        Self::new()
164    }
165}
postrust_core/schema_cache/mod.rs

postrust_core/schema_cache/
mod.rs
