Skip to main content

safe_migrate/
sync.rs

1// FILE: src/sync.rs
2
3use crate::ast::identifiers::ObjectId;
4use crate::db::cache::{DbCache, ForeignKeyCache, IndexCache};
5use crate::model::relation::{Persistence, RelationKind, RelationState};
6use anyhow::{Context, Result};
7use postgres::{Client, NoTls};
8use std::fs;
9use std::path::Path;
10
11pub fn sync_cache(out_path: &Path) -> Result<()> {
12    // Strict env-only credential enforcement
13    let db_url = std::env::var("DATABASE_URL")
14        .context("DATABASE_URL environment variable is required to sync database stats. Do not pass credentials via CLI flags or config files.")?;
15
16    // Destructive cache removal prevents corrupted reads on failures
17    if out_path.exists() {
18        fs::remove_file(out_path).context("Failed to remove old cache file before sync")?;
19    }
20
21    // Warn if connecting to a non-local host without TLS
22    let host = db_url
23        .split('@')
24        .nth(1)
25        .and_then(|h| h.split('/').next())
26        .unwrap_or("localhost");
27    if !host.starts_with("localhost")
28        && !host.starts_with("127.")
29        && !host.starts_with("/")
30        && host != "::1"
31    {
32        eprintln!(
33            "[WARN] Connecting to PostgreSQL at {} without TLS encryption.\n\
34             The database password will be sent in cleartext over the network.\n\
35             Use an SSH tunnel or a local connection for sensitive databases,\n\
36             or add native-tls support (see https://github.com/dsecurity49/safe-migrate).",
37            host
38        );
39    }
40
41    let mut client = Client::connect(&db_url, NoTls).context("Failed to connect to PostgreSQL")?;
42
43    let mut cache = DbCache::new();
44
45    // Query 1: Server Version
46    let version_row = client.query_one("SHOW server_version_num;", &[])?;
47    let version_str: String = version_row.get(0);
48    cache.pg_version_num = version_str.parse::<u32>().ok();
49
50    // Query 2: Relations + Staleness
51    let table_query = "
52        SELECT
53            n.nspname AS schema_name,
54            c.relname AS relation_name,
55            c.relkind AS relation_kind,
56            c.relpersistence AS persistence,
57            CASE WHEN c.reltuples < 0 THEN -1 ELSE c.reltuples::bigint END AS estimated_rows,
58            GREATEST(c.relpages::bigint, 0) AS relpages,
59            to_char(s.last_analyze, 'YYYY-MM-DD HH24:MI:SS') AS last_analyze,
60            to_char(s.last_autoanalyze, 'YYYY-MM-DD HH24:MI:SS') AS last_autoanalyze
61        FROM pg_class c
62        JOIN pg_namespace n ON n.oid = c.relnamespace
63        LEFT JOIN pg_stat_user_tables s ON s.relid = c.oid
64        WHERE c.relkind IN ('r', 'p', 'v', 'm')
65          AND n.nspname NOT IN ('pg_catalog', 'information_schema');
66    ";
67
68    for row in client.query(table_query, &[])? {
69        let schema_name: String = row.get("schema_name");
70        let relation_name: String = row.get("relation_name");
71        let relkind: i8 = row.get("relation_kind");
72        let persistence_char: i8 = row.get("persistence");
73        let raw_rows: i64 = row.get("estimated_rows");
74        let relpages: i64 = row.get("relpages");
75
76        let last_analyze: Option<String> = row.get("last_analyze");
77        let last_autoanalyze: Option<String> = row.get("last_autoanalyze");
78
79        let object_id = ObjectId::new(&schema_name, &relation_name);
80
81        let kind = match relkind as u8 {
82            b'v' => RelationKind::View,
83            b'm' => RelationKind::MaterializedView,
84            _ => RelationKind::Table,
85        };
86
87        let persistence = match persistence_char as u8 {
88            b't' => Persistence::Temporary,
89            b'u' => Persistence::Unlogged,
90            _ => Persistence::Permanent,
91        };
92
93        let estimated_rows = if raw_rows < 0 {
94            None
95        } else {
96            Some(raw_rows as u64)
97        };
98
99        let mut state = RelationState::new(
100            object_id.clone(),
101            ObjectId::new("public", "postgres"),
102            0,
103            estimated_rows,
104            kind,
105            persistence,
106            0,
107        );
108        state.relpages = Some(relpages as u64);
109        state.last_analyze = last_analyze;
110        state.last_autoanalyze = last_autoanalyze;
111
112        cache.insert_baseline(object_id, state);
113    }
114
115    // Query 3: Columns + Width
116    let col_query = "
117        SELECT
118            n.nspname AS schema_name,
119            c.relname AS relation_name,
120            a.attname AS column_name,
121            pg_catalog.format_type(a.atttypid, a.atttypmod) AS type_name,
122            a.attnotnull AS not_null,
123            s.avg_width AS avg_width,
124            pg_get_expr(ad.adbin, ad.adrelid) AS default_expr_text,
125            a.atttypmod AS type_modifier
126        FROM pg_attribute a
127        JOIN pg_class c ON a.attrelid = c.oid
128        JOIN pg_namespace n ON n.oid = c.relnamespace
129        LEFT JOIN pg_stats s ON s.schemaname = n.nspname AND s.tablename = c.relname AND s.attname = a.attname
130        LEFT JOIN pg_attrdef ad ON ad.adrelid = a.attrelid AND ad.adnum = a.attnum
131        WHERE a.attnum > 0 AND NOT a.attisdropped
132          AND c.relkind IN ('r', 'p', 'v', 'm')
133          AND n.nspname NOT IN ('pg_catalog', 'information_schema');
134    ";
135
136    for row in client.query(col_query, &[])? {
137        let schema_name: String = row.get("schema_name");
138        let relation_name: String = row.get("relation_name");
139        let column_name: String = row.get("column_name");
140        let type_name: String = row.get("type_name");
141        let not_null: bool = row.get("not_null");
142        let avg_width: Option<i32> = row.get("avg_width");
143        let default_expr_text: Option<String> = row.get("default_expr_text");
144        let type_modifier: Option<i32> = row.get("type_modifier");
145
146        let object_id = ObjectId::new(&schema_name, &relation_name);
147
148        if let Some(rel) = cache.relations.get_mut(&object_id) {
149            rel.columns.push(crate::model::column::Column {
150                name: column_name,
151                data_type: Some(type_name),
152                is_nullable: !not_null,
153                default: None,
154                avg_width,
155                default_expr_text,
156                type_modifier,
157            });
158        }
159    }
160
161    // Query 4: Triggers & Policies
162    let tp_query = "
163        SELECT 
164            n.nspname AS schema_name,
165            c.relname AS relation_name,
166            COALESCE(array_agg(DISTINCT t.tgname) FILTER (WHERE t.tgname IS NOT NULL AND t.tgisinternal = false), '{}') as triggers,
167            COALESCE(array_agg(DISTINCT p.polname) FILTER (WHERE p.polname IS NOT NULL), '{}') as policies
168        FROM pg_class c
169        JOIN pg_namespace n ON n.oid = c.relnamespace
170        LEFT JOIN pg_trigger t ON t.tgrelid = c.oid
171        LEFT JOIN pg_policy p ON p.polrelid = c.oid
172        WHERE c.relkind IN ('r', 'p', 'v', 'm') AND n.nspname NOT IN ('pg_catalog', 'information_schema')
173        GROUP BY n.nspname, c.relname;
174    ";
175
176    for row in client.query(tp_query, &[])? {
177        let schema_name: String = row.get("schema_name");
178        let relation_name: String = row.get("relation_name");
179        let triggers: Vec<String> = row.get("triggers");
180        let policies: Vec<String> = row.get("policies");
181
182        let object_id = ObjectId::new(&schema_name, &relation_name);
183
184        if let Some(rel) = cache.relations.get_mut(&object_id) {
185            rel.triggers.extend(triggers);
186            rel.policies.extend(policies);
187        }
188    }
189
190    // Query 5: Foreign Keys
191    let fk_query = "
192        SELECT 
193            c.conname AS constraint_name,
194            n1.nspname AS from_schema, t1.relname AS from_table,
195            n2.nspname AS to_schema, t2.relname AS to_table
196        FROM pg_constraint c
197        JOIN pg_class t1 ON t1.oid = c.conrelid
198        JOIN pg_namespace n1 ON n1.oid = t1.relnamespace
199        JOIN pg_class t2 ON t2.oid = c.confrelid
200        JOIN pg_namespace n2 ON n2.oid = t2.relnamespace
201        WHERE c.contype = 'f';
202    ";
203
204    for row in client.query(fk_query, &[])? {
205        let constraint_name: String = row.get("constraint_name");
206        let from_schema: String = row.get("from_schema");
207        let from_table: String = row.get("from_table");
208        let to_schema: String = row.get("to_schema");
209        let to_table: String = row.get("to_table");
210
211        cache.foreign_keys.push(ForeignKeyCache {
212            constraint_name,
213            from_table: ObjectId::new(&from_schema, &from_table),
214            to_table: ObjectId::new(&to_schema, &to_table),
215        });
216    }
217
218    // Query 6: Indexes
219    let idx_query = "
220        SELECT 
221            n_i.nspname AS index_schema, i.relname AS index_name,
222            n_t.nspname AS table_schema, t.relname AS table_name
223        FROM pg_index x
224        JOIN pg_class i ON i.oid = x.indexrelid
225        JOIN pg_namespace n_i ON n_i.oid = i.relnamespace
226        JOIN pg_class t ON t.oid = x.indrelid
227        JOIN pg_namespace n_t ON n_t.oid = t.relnamespace
228        WHERE x.indisvalid = true;
229    ";
230
231    for row in client.query(idx_query, &[])? {
232        let index_schema: String = row.get("index_schema");
233        let index_name: String = row.get("index_name");
234        let table_schema: String = row.get("table_schema");
235        let table_name: String = row.get("table_name");
236
237        cache.indexes.push(IndexCache {
238            index_id: ObjectId::new(&index_schema, &index_name),
239            table_id: ObjectId::new(&table_schema, &table_name),
240        });
241    }
242
243    // Atomic write via temp file
244    let json = serde_json::to_string_pretty(&cache)?;
245    let tmp_path = out_path.with_extension("tmp");
246
247    fs::write(&tmp_path, json).context("Failed to write temporary cache file")?;
248    fs::rename(&tmp_path, out_path).context("Failed to atomically rename cache file")?;
249
250    Ok(())
251}