| 1 |
|
| 2 |
|
| 3 |
|
| 4 |
|
| 5 |
|
| 6 |
|
| 7 |
|
| 8 |
|
| 9 |
|
| 10 |
|
| 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 |
|
| 19 |
|
| 20 |
const BATCH: usize = 500; |
| 21 |
|
| 22 |
|
| 23 |
struct Col { |
| 24 |
name: String, |
| 25 |
kind: Kind, |
| 26 |
|
| 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 |
|
| 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 |
|
| 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 |
|
| 98 |
|
| 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 |
|
| 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 |
} |