1use 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 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 if out_path.exists() {
18 fs::remove_file(out_path).context("Failed to remove old cache file before sync")?;
19 }
20
21 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 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 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 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 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 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 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 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}