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 queries;
7mod relationship;
8mod routine;
9mod table;
10
11pub use relationship::{Cardinality, Junction, Relationship, RelationshipsMap};
12pub use routine::{FuncVolatility, RetType, Routine, RoutineMap, RoutineParam};
13pub use table::{Column, ColumnMap, Table, TablesMap};
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(
84        &self,
85        qi: &QualifiedIdentifier,
86        schema: &str,
87    ) -> Option<&Vec<Relationship>> {
88        self.relationships.get(&(qi.clone(), schema.to_string()))
89    }
90
91    /// Get a routine by qualified identifier.
92    pub fn get_routines(&self, qi: &QualifiedIdentifier) -> Option<&Vec<Routine>> {
93        self.routines.get(qi)
94    }
95
96    /// Check if a timezone is valid.
97    pub fn is_valid_timezone(&self, tz: &str) -> bool {
98        self.timezones.contains(tz)
99    }
100
101    /// Get a summary of the cached schema.
102    pub fn summary(&self) -> String {
103        format!(
104            "SchemaCache: {} tables, {} relationship sets, {} routines, PG {}",
105            self.tables.len(),
106            self.relationships.len(),
107            self.routines.len(),
108            self.pg_version
109        )
110    }
111
112    /// Find a relationship between two tables by name.
113    pub fn find_relationship(
114        &self,
115        from: &QualifiedIdentifier,
116        to_name: &str,
117        schema: &str,
118    ) -> Option<&Relationship> {
119        self.get_relationships(from, schema)?
120            .iter()
121            .find(|r| match r {
122                Relationship::ForeignKey { foreign_table, .. } => foreign_table.name == to_name,
123                Relationship::Computed { foreign_table, .. } => foreign_table.name == to_name,
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    /// Create a schema cache reference from a static cache.
139    pub fn from_static(cache: SchemaCache) -> Self {
140        Self(Arc::new(tokio::sync::RwLock::new(Some(cache))))
141    }
142
143    /// Load or reload the schema cache.
144    pub async fn load(&self, pool: &PgPool, schemas: &[String]) -> Result<()> {
145        let cache = SchemaCache::load(pool, schemas).await?;
146        let mut guard = self.0.write().await;
147        *guard = Some(cache);
148        Ok(())
149    }
150
151    /// Get a read reference to the schema cache.
152    pub async fn get(&self) -> Result<tokio::sync::RwLockReadGuard<'_, Option<SchemaCache>>> {
153        let guard = self.0.read().await;
154        if guard.is_none() {
155            return Err(Error::SchemaCacheNotLoaded);
156        }
157        Ok(guard)
158    }
159
160    /// Check if the cache is loaded.
161    pub async fn is_loaded(&self) -> bool {
162        self.0.read().await.is_some()
163    }
164}
165
166impl Default for SchemaCacheRef {
167    fn default() -> Self {
168        Self::new()
169    }
170}
postrust_core/schema_cache/mod.rs

postrust_core/schema_cache/
mod.rs
