OXYGENHUB
# 4e2cad00a484 oxygenhub / crates / hub / src / db / copy.rs
9.1 KiB · Rust · last changed Postgres, step 4: admin copy-db moves a SQLite hub's records into Postgres 19 days ago source blame history raw
1
//! Copy a SQLite hub's records into PostgreSQL (#211), for moving a hub:
2
//! `oxygenhub admin copy-db --from <old data>/hub.db` on the new hub, built
3
//! with the `postgres` feature and migrated, before it first serves.
4
//!
5
//! Everything is copied in one transaction, ids kept, so a failed copy
6
//! leaves the database as it was. The tables go in the order the schema
7
//! made them, which is after every table they point at; a column pointing
8
//! at its own table (a fork's parent) is filled in once all its rows are
9
//! there. Then each id sequence is moved past the copied ids, the search
10
//! index is copied, and every table's count is checked against SQLite's.
11
12
use super::Pool;
13
use anyhow::{bail, Context, Result};
14
use sqlx::sqlite::{SqliteConnectOptions, SqliteRow};
15
use sqlx::{Connection, Row};
16
use std::path::Path;
17
18
/// Rows sent in one INSERT: well under Postgres's 65535 parameters for the
19
/// widest table.
20
const BATCH: usize = 500;
21
22
/// One column of a table here, and how to carry its values across.
23
struct Col {
24
    name: String,
25
    kind: Kind,
26
    /// Points at its own table: copied after the rows.
27
    self_ref: bool,
28
}
29
30
#[derive(Clone, Copy, PartialEq)]
31
enum Kind {
32
    Int,
33
    Bool,
34
    Real,
35
    Blob,
36
    Text,
37
}
38
39
enum Value {
40
    Int(Option<i64>),
41
    Bool(Option<bool>),
42
    Real(Option<f64>),
43
    Blob(Option<Vec<u8>>),
44
    Text(Option<String>),
45
}
46
47
/// What was copied: `(table, rows)`, in copy order.
48
pub async fn from_sqlite(pool: &Pool, path: &Path) -> Result<Vec<(String, i64)>> {
49
    let options = SqliteConnectOptions::new().filename(path).read_only(true);
50
    let mut lite = sqlx::SqliteConnection::connect_with(&options).await.with_context(|| format!("open {}", path.display()))?;
51
    up_to_date(&mut lite).await?;
52
    let tables: Vec<String> = sqlx::query_scalar(
53
        "SELECT CAST(relname AS TEXT) FROM pg_class WHERE relkind = 'r' AND relnamespace = CAST(current_schema() AS regnamespace) \
54
         AND relname NOT IN ('_sqlx_migrations', 'search') ORDER BY oid",
55
    )
56
    .fetch_all(pool)
57
    .await?;
58
    for t in &tables {
59
        let any: bool = sqlx::query_scalar(&format!("SELECT EXISTS (SELECT 1 FROM \"{t}\")")).fetch_one(pool).await?;
60
        if any {
61
            bail!("{t} already has rows here: copy into a new, empty database");
62
        }
63
    }
64
    let mut tx = pool.begin().await?;
65
    let mut copied = Vec::new();
66
    for t in &tables {
67
        let cols = columns(&mut tx, t).await?;
68
        let n = copy_table(&mut lite, &mut tx, t, &cols).await.with_context(|| format!("copy {t}"))?;
69
        copied.push((t.clone(), n));
70
    }
71
    // Ids carry on after the copied ones.
72
    let numbered: Vec<String> = sqlx::query_scalar(
73
        "SELECT CAST(t.relname AS TEXT) FROM pg_attribute a JOIN pg_class t ON t.oid = a.attrelid \
74
         WHERE t.relnamespace = CAST(current_schema() AS regnamespace) AND a.attname = 'id' AND a.attidentity <> ''",
75
    )
76
    .fetch_all(&mut *tx)
77
    .await?;
78
    for t in &numbered {
79
        sqlx::query(&format!("SELECT setval(pg_get_serial_sequence($1, 'id'), COALESCE((SELECT MAX(id) FROM \"{t}\"), 0) + 1, false)"))
80
            .bind(t)
81
            .execute(&mut *tx)
82
            .await?;
83
    }
84
    let search = copy_search(&mut lite, &mut tx).await.context("copy the search index")?;
85
    copied.push(("search".into(), search));
86
    for (t, n) in &copied {
87
        let there: i64 = sqlx::query_scalar(&format!("SELECT COUNT(*) FROM \"{t}\"")).fetch_one(&mut *tx).await?;
88
        if there != *n {
89
            bail!("{t}: {n} rows read but {there} here");
90
        }
91
    }
92
    tx.commit().await?;
93
    lite.close().await?;
94
    Ok(copied)
95
}
96
97
/// The SQLite hub must have every migration this hub knows, and no more:
98
/// the two schemas match only then.
99
async fn up_to_date(lite: &mut sqlx::SqliteConnection) -> Result<()> {
100
    let latest = sqlx::migrate!("./migrations").iter().map(|m| m.version).max().unwrap_or(0);
101
    let have: Option<i64> =
102
        sqlx::query_scalar("SELECT MAX(version) FROM _sqlx_migrations WHERE success = TRUE").fetch_one(&mut *lite).await?;
103
    match have {
104
        Some(v) if v == latest => Ok(()),
105
        Some(v) => bail!("the SQLite hub is at migration {v}, this one at {latest}: run the same version of the hub on both"),
106
        None => bail!("not a hub's database: it has no migrations"),
107
    }
108
}
109
110
async fn columns(tx: &mut sqlx::PgConnection, table: &str) -> Result<Vec<Col>> {
111
    let rows: Vec<(String, String, bool)> = sqlx::query_as(
112
        "SELECT CAST(a.attname AS TEXT), format_type(a.atttypid, a.atttypmod), \
113
           EXISTS (SELECT 1 FROM pg_constraint k WHERE k.conrelid = t.oid AND k.confrelid = t.oid AND k.contype = 'f' AND a.attnum = ANY (k.conkey)) \
114
         FROM pg_attribute a JOIN pg_class t ON t.oid = a.attrelid \
115
         WHERE t.relname = $1 AND t.relnamespace = CAST(current_schema() AS regnamespace) \
116
           AND a.attnum > 0 AND NOT a.attisdropped AND a.attgenerated = '' ORDER BY a.attnum",
117
    )
118
    .bind(table)
119
    .fetch_all(&mut *tx)
120
    .await?;
121
    Ok(rows
122
        .into_iter()
123
        .map(|(name, ty, self_ref)| {
124
            let kind = match ty.as_str() {
125
                "bigint" | "integer" | "smallint" => Kind::Int,
126
                "boolean" => Kind::Bool,
127
                "double precision" | "real" => Kind::Real,
128
                "bytea" => Kind::Blob,
129
                _ => Kind::Text,
130
            };
131
            Col { name, kind, self_ref }
132
        })
133
        .collect())
134
}
135
136
fn read(row: &SqliteRow, i: usize, kind: Kind) -> Result<Value> {
137
    Ok(match kind {
138
        Kind::Int => Value::Int(row.try_get(i)?),
139
        Kind::Bool => Value::Bool(row.try_get::<Option<i64>, _>(i)?.map(|v| v != 0)),
140
        Kind::Real => Value::Real(row.try_get(i)?),
141
        Kind::Blob => Value::Blob(row.try_get(i)?),
142
        Kind::Text => Value::Text(row.try_get(i)?),
143
    })
144
}
145
146
fn list(cols: &[&Col]) -> String {
147
    cols.iter().map(|c| format!("\"{}\"", c.name)).collect::<Vec<_>>().join(", ")
148
}
149
150
async fn copy_table(lite: &mut sqlx::SqliteConnection, tx: &mut sqlx::PgConnection, table: &str, cols: &[Col]) -> Result<i64> {
151
    use futures::TryStreamExt;
152
    let all: Vec<&Col> = cols.iter().collect();
153
    let now: Vec<&Col> = cols.iter().filter(|c| !c.self_ref).collect();
154
    let id_at = cols.iter().position(|c| c.name == "id");
155
    let sql = format!("SELECT {} FROM \"{table}\" ORDER BY rowid", list(&all));
156
    // Streamed, a batch at a time: a table (Flow's logs) can be large.
157
    let mut rows = sqlx::query(&sql).fetch(&mut *lite);
158
    let (mut n, mut batch, mut later) = (0i64, Vec::with_capacity(BATCH), Vec::new());
159
    loop {
160
        let row = rows.try_next().await?;
161
        if let Some(row) = &row {
162
            let mut values = Vec::with_capacity(now.len());
163
            let mut refs = Vec::new();
164
            for (i, c) in cols.iter().enumerate() {
165
                let v = read(row, i, c.kind).with_context(|| format!("read {table}.{}", c.name))?;
166
                if c.self_ref {
167
                    if matches!(v, Value::Int(Some(_))) {
168
                        refs.push((c.name.clone(), v));
169
                    }
170
                } else {
171
                    values.push(v);
172
                }
173
            }
174
            if !refs.is_empty() {
175
                let Some(i) = id_at else { bail!("{table} points at itself but has no id") };
176
                later.push((read(row, i, Kind::Int)?, refs));
177
            }
178
            batch.push(values);
179
            n += 1;
180
        }
181
        if batch.len() == BATCH || (row.is_none() && !batch.is_empty()) {
182
            let mut q = sqlx::QueryBuilder::<super::Db>::new(format!("INSERT INTO \"{table}\" ({}) ", list(&now)));
183
            q.push_values(batch.drain(..), |mut b, values| {
184
                for v in values {
185
                    match v {
186
                        Value::Int(v) => b.push_bind(v),
187
                        Value::Bool(v) => b.push_bind(v),
188
                        Value::Real(v) => b.push_bind(v),
189
                        Value::Blob(v) => b.push_bind(v),
190
                        Value::Text(v) => b.push_bind(v),
191
                    };
192
                }
193
            });
194
            q.build().execute(&mut *tx).await?;
195
        }
196
        if row.is_none() {
197
            break;
198
        }
199
    }
200
    for (id, refs) in later {
201
        let Value::Int(id) = id else { continue };
202
        for (name, v) in refs {
203
            if let Value::Int(v) = v {
204
                sqlx::query(&format!("UPDATE \"{table}\" SET \"{name}\" = $1 WHERE id = $2")).bind(v).bind(id).execute(&mut *tx).await?;
205
            }
206
        }
207
    }
208
    Ok(n)
209
}
210
211
async fn copy_search(lite: &mut sqlx::SqliteConnection, tx: &mut sqlx::PgConnection) -> Result<i64> {
212
    let rows: Vec<(String, i64, i64, i64, String, String)> =
213
        sqlx::query_as("SELECT kind, repo_id, target_id, number, title, body FROM search").fetch_all(&mut *lite).await?;
214
    for chunk in rows.chunks(BATCH) {
215
        let mut q = sqlx::QueryBuilder::<super::Db>::new("INSERT INTO search (kind, repo_id, target_id, number, title, body) ");
216
        q.push_values(chunk, |mut b, (kind, repo_id, target_id, number, title, body)| {
217
            b.push_bind(kind).push_bind(repo_id).push_bind(target_id).push_bind(number).push_bind(title).push_bind(body);
218
        });
219
        q.build().execute(&mut *tx).await?;
220
    }
221
    Ok(rows.len() as i64)
222
}