assay_engine/migrate/
schema.rs

1//! Reading the shape of both stores.
2//!
3//! Nothing here hard-codes a table list. The source's tables come from
4//! each ATTACHed database's `sqlite_master`, the target's columns and
5//! foreign keys from the Postgres catalog, so a module that adds a table
6//! is carried without touching this file.
7
8use anyhow::{Context, Result, bail};
9use std::collections::{HashMap, HashSet};
10
11/// A table to copy, named the same on both sides.
12#[derive(Debug, Clone, PartialEq, Eq, Hash)]
13pub struct Table {
14    pub schema: String,
15    pub name: String,
16}
17
18impl std::fmt::Display for Table {
19    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
20        write!(f, "{}.{}", self.schema, self.name)
21    }
22}
23
24/// One target column: its name and the Postgres type to bind for.
25#[derive(Debug, Clone)]
26pub struct Column {
27    pub name: String,
28    pub udt: String,
29}
30
31/// Every table in the SQLite store, in the order its module was attached.
32pub async fn source_tables(pool: &sqlx::SqlitePool, modules: &[&str]) -> Result<Vec<Table>> {
33    let mut tables = Vec::new();
34    for module in modules {
35        let sql = format!(
36            "SELECT name FROM {module}.sqlite_master
37             WHERE type = 'table' AND name NOT LIKE 'sqlite_%'
38             ORDER BY name"
39        );
40        let names: Vec<(String,)> = sqlx::query_as(&sql)
41            .fetch_all(pool)
42            .await
43            .with_context(|| format!("list tables in sqlite database {module}"))?;
44        tables.extend(names.into_iter().map(|(name,)| Table {
45            schema: (*module).to_string(),
46            name,
47        }));
48    }
49    Ok(tables)
50}
51
52/// Column names of a SQLite table, in declaration order. The second
53/// argument to `pragma_table_info` selects the ATTACHed database.
54pub async fn source_columns(pool: &sqlx::SqlitePool, table: &Table) -> Result<Vec<String>> {
55    let rows: Vec<(String,)> = sqlx::query_as("SELECT name FROM pragma_table_info(?, ?)")
56        .bind(&table.name)
57        .bind(&table.schema)
58        .fetch_all(pool)
59        .await
60        .with_context(|| format!("read columns of {table}"))?;
61    Ok(rows.into_iter().map(|(n,)| n).collect())
62}
63
64/// Columns of a Postgres table, in ordinal order. Empty when the table
65/// does not exist.
66pub async fn target_columns(pool: &sqlx::PgPool, table: &Table) -> Result<Vec<Column>> {
67    let rows: Vec<(String, String)> = sqlx::query_as(
68        "SELECT column_name, udt_name FROM information_schema.columns
69         WHERE table_schema = $1 AND table_name = $2
70         ORDER BY ordinal_position",
71    )
72    .bind(&table.schema)
73    .bind(&table.name)
74    .fetch_all(pool)
75    .await
76    .with_context(|| format!("read postgres columns of {table}"))?;
77    Ok(rows
78        .into_iter()
79        .map(|(name, udt)| Column { name, udt })
80        .collect())
81}
82
83/// Tables in the target's engine schemas that already hold rows.
84///
85/// Run before any DDL: a target the engine has never touched has no such
86/// schemas at all, and that is the only state a migration may write into.
87pub async fn non_empty_tables(pool: &sqlx::PgPool, schemas: &[&str]) -> Result<Vec<(Table, i64)>> {
88    let owned: Vec<String> = schemas.iter().map(|s| (*s).to_string()).collect();
89    let existing: Vec<(String, String)> = sqlx::query_as(
90        "SELECT table_schema, table_name FROM information_schema.tables
91         WHERE table_schema = ANY($1) AND table_type = 'BASE TABLE'
92         ORDER BY table_schema, table_name",
93    )
94    .bind(&owned)
95    .fetch_all(pool)
96    .await
97    .context("list existing tables in the target")?;
98
99    let mut occupied = Vec::new();
100    for (schema, name) in existing {
101        let table = Table { schema, name };
102        let count = count_rows_pg(pool, &table).await?;
103        if count > 0 {
104            occupied.push((table, count));
105        }
106    }
107    Ok(occupied)
108}
109
110/// Schema-qualified, quoted table reference. Identical syntax on both
111/// backends because the SQLite store ATTACHes one database per schema.
112pub fn quoted(table: &Table) -> String {
113    format!(r#""{}"."{}""#, table.schema, table.name)
114}
115
116fn count_sql(table: &Table) -> String {
117    format!("SELECT COUNT(*) FROM {}", quoted(table))
118}
119
120pub async fn count_rows_pg(pool: &sqlx::PgPool, table: &Table) -> Result<i64> {
121    let count = sqlx::query_scalar(&count_sql(table)).fetch_one(pool).await;
122    count.with_context(|| format!("count rows of {table}"))
123}
124
125pub async fn count_rows_sqlite(pool: &sqlx::SqlitePool, table: &Table) -> Result<i64> {
126    let count = sqlx::query_scalar(&count_sql(table)).fetch_one(pool).await;
127    count.with_context(|| format!("count rows of {table}"))
128}
129
130/// Order `tables` so that every table follows the tables it references.
131///
132/// Insert order has to respect foreign keys, and the graph is read from
133/// the target catalog rather than declared here so a new reference does
134/// not need a matching edit.
135pub async fn insert_order(pool: &sqlx::PgPool, tables: &[Table]) -> Result<Vec<Table>> {
136    let schemas: Vec<String> = tables
137        .iter()
138        .map(|t| t.schema.clone())
139        .collect::<HashSet<_>>()
140        .into_iter()
141        .collect();
142    let edges: Vec<(String, String, String, String)> = sqlx::query_as(
143        "SELECT ns.nspname, cl.relname, fns.nspname, fcl.relname
144         FROM pg_constraint c
145         JOIN pg_class cl ON cl.oid = c.conrelid
146         JOIN pg_namespace ns ON ns.oid = cl.relnamespace
147         JOIN pg_class fcl ON fcl.oid = c.confrelid
148         JOIN pg_namespace fns ON fns.oid = fcl.relnamespace
149         WHERE c.contype = 'f' AND ns.nspname = ANY($1)",
150    )
151    .bind(&schemas)
152    .fetch_all(pool)
153    .await
154    .context("read foreign keys from the target")?;
155
156    let present: HashSet<Table> = tables.iter().cloned().collect();
157    let mut blockers: HashMap<Table, HashSet<Table>> =
158        tables.iter().map(|t| (t.clone(), HashSet::new())).collect();
159    for (schema, name, ref_schema, ref_name) in edges {
160        let child = Table { schema, name };
161        let parent = Table {
162            schema: ref_schema,
163            name: ref_name,
164        };
165        if child == parent || !present.contains(&child) || !present.contains(&parent) {
166            continue;
167        }
168        blockers.entry(child).or_default().insert(parent);
169    }
170
171    let mut ordered: Vec<Table> = Vec::with_capacity(tables.len());
172    let mut placed: HashSet<Table> = HashSet::new();
173    while ordered.len() < tables.len() {
174        let ready: Vec<Table> = tables
175            .iter()
176            .filter(|t| !placed.contains(*t))
177            .filter(|t| blockers[*t].iter().all(|p| placed.contains(p)))
178            .cloned()
179            .collect();
180        if ready.is_empty() {
181            let stuck: Vec<String> = tables
182                .iter()
183                .filter(|t| !placed.contains(*t))
184                .map(|t| t.to_string())
185                .collect();
186            bail!("foreign keys form a cycle across {}", stuck.join(", "));
187        }
188        for table in ready {
189            placed.insert(table.clone());
190            ordered.push(table);
191        }
192    }
193    Ok(ordered)
194}
195
196/// Re-point every sequence at the ids that were just copied in.
197///
198/// Copying preserves ids, so a sequence still sitting at 1 would hand
199/// the next insert an id the migrated rows already own.
200pub async fn resync_sequences(pool: &sqlx::PgPool, tables: &[Table]) -> Result<usize> {
201    let mut resynced = 0;
202    for table in tables {
203        for column in target_columns(pool, table).await? {
204            let qualified = format!(r#""{}"."{}""#, table.schema, table.name);
205            let sequence: Option<String> = sqlx::query_scalar("SELECT pg_get_serial_sequence($1, $2)")
206                .bind(&qualified)
207                .bind(&column.name)
208                .fetch_one(pool)
209                .await
210                .with_context(|| format!("resolve sequence for {table}.{}", column.name))?;
211            let Some(sequence) = sequence else { continue };
212            let sql = format!(
213                r#"SELECT setval('{sequence}', COALESCE((SELECT MAX("{}") FROM {qualified}), 0) + 1, false)"#,
214                column.name
215            );
216            sqlx::query(&sql)
217                .execute(pool)
218                .await
219                .with_context(|| format!("resync sequence {sequence}"))?;
220            resynced += 1;
221        }
222    }
223    Ok(resynced)
224}
assay_engine/migrate/schema.rs

assay_engine/migrate/
schema.rs
