use bytes::{BufMut, Bytes, BytesMut};
use chrono::{DateTime, TimeZone, Utc};
use database::pool::{postgres, sqlite, ConnectionManager};
use futures_util::sink::SinkExt;
use hashbrown::HashMap;
use serde::{Serialize, Serializer};
use std::io::Write;
use std::time::Instant;
const NULL_STRING: &str = "\\N";
trait Table {
fn name() -> &'static str;
fn sqlite_attributes() -> &'static str;
fn postgres_attributes() -> &'static str;
fn postgres_generated_id_attribute() -> Option<&'static str>;
fn write_postgres_csv_row<W: Write>(writer: &mut csv::Writer<W>, row: &rusqlite::Row);
}
struct Artifact;
#[derive(Serialize)]
struct ArtifactRow<'a> {
id: i32,
name: &'a str,
date: Nullable<DateTime<Utc>>,
typ: &'a str,
}
impl Table for Artifact {
fn name() -> &'static str {
"artifact"
}
fn sqlite_attributes() -> &'static str {
"id, name, date, type"
}
fn postgres_attributes() -> &'static str {
"id, name, date, type"
}
fn postgres_generated_id_attribute() -> Option<&'static str> {
Some("id")
}
fn write_postgres_csv_row<W: Write>(writer: &mut csv::Writer<W>, row: &rusqlite::Row) {
let date: Option<i64> = row.get(2).unwrap();
writer
.serialize(ArtifactRow {
id: row.get(0).unwrap(),
name: row.get_ref(1).unwrap().as_str().unwrap(),
date: Nullable(date.map(|seconds| Utc.timestamp_opt(seconds, 0).unwrap())),
typ: row.get_ref(3).unwrap().as_str().unwrap(),
})
.unwrap();
}
}
struct Benchmark;
#[derive(Serialize)]
struct BenchmarkRow<'a> {
name: &'a str,
stabilized: Nullable<bool>,
category: &'a str,
}
impl Table for Benchmark {
fn name() -> &'static str {
"benchmark"
}
fn sqlite_attributes() -> &'static str {
"name, stabilized, category"
}
fn postgres_attributes() -> &'static str {
"name, stabilized, category"
}
fn postgres_generated_id_attribute() -> Option<&'static str> {
None
}
fn write_postgres_csv_row<W: Write>(writer: &mut csv::Writer<W>, row: &rusqlite::Row) {
writer
.serialize(BenchmarkRow {
name: row.get_ref(0).unwrap().as_str().unwrap(),
stabilized: row.get(1).unwrap(),
category: row.get_ref(2).unwrap().as_str().unwrap(),
})
.unwrap();
}
}
struct Collection;
#[derive(Serialize)]
struct CollectionRow<'a> {
id: i32,
perf_commit: Nullable<&'a str>,
}
impl Table for Collection {
fn name() -> &'static str {
"collection"
}
fn sqlite_attributes() -> &'static str {
"id, perf_commit"
}
fn postgres_attributes() -> &'static str {
"id, perf_commit"
}
fn postgres_generated_id_attribute() -> Option<&'static str> {
Some("id")
}
fn write_postgres_csv_row<W: Write>(writer: &mut csv::Writer<W>, row: &rusqlite::Row) {
writer
.serialize(CollectionRow {
id: row.get(0).unwrap(),
perf_commit: row.get_ref(1).unwrap().try_into().unwrap(),
})
.unwrap();
}
}
struct Error;
#[derive(Serialize)]
struct ErrorRow<'a> {
id: i32,
aid: i32,
context: &'a str,
message: Nullable<&'a str>,
}
impl Table for Error {
fn name() -> &'static str {
"error"
}
fn sqlite_attributes() -> &'static str {
"id, aid, context, message, job_id"
}
fn postgres_attributes() -> &'static str {
"id, aid, context, message, job_id"
}
fn postgres_generated_id_attribute() -> Option<&'static str> {
None
}
fn write_postgres_csv_row<W: Write>(writer: &mut csv::Writer<W>, row: &rusqlite::Row) {
writer
.serialize(ErrorRow {
id: row.get(0).unwrap(),
aid: row.get(1).unwrap(),
context: row.get_ref(2).unwrap().as_str().unwrap(),
message: row.get_ref(3).unwrap().try_into().unwrap(),
})
.unwrap();
}
}
struct Pstat;
#[derive(Serialize)]
struct PstatRow {
series: i32,
aid: i32,
cid: i32,
value: f64,
}
impl Table for Pstat {
fn name() -> &'static str {
"pstat"
}
fn sqlite_attributes() -> &'static str {
"series, aid, cid, value"
}
fn postgres_attributes() -> &'static str {
"series, aid, cid, value"
}
fn postgres_generated_id_attribute() -> Option<&'static str> {
None
}
fn write_postgres_csv_row<W: Write>(writer: &mut csv::Writer<W>, row: &rusqlite::Row) {
writer
.serialize(PstatRow {
series: row.get(0).unwrap(),
aid: row.get(1).unwrap(),
cid: row.get(2).unwrap(),
value: row.get(3).unwrap(),
})
.unwrap();
}
}
struct PstatSeries;
#[derive(Serialize)]
struct PstatSeriesRow<'a> {
id: i32,
krate: &'a str,
profile: &'a str,
scenario: &'a str,
backend: &'a str,
target: &'a str,
metric: &'a str,
}
impl Table for PstatSeries {
fn name() -> &'static str {
"pstat_series"
}
fn sqlite_attributes() -> &'static str {
"id, crate, profile, scenario, backend, target, metric"
}
fn postgres_attributes() -> &'static str {
"id, crate, profile, scenario, backend, target, metric"
}
fn postgres_generated_id_attribute() -> Option<&'static str> {
Some("id")
}
fn write_postgres_csv_row<W: Write>(writer: &mut csv::Writer<W>, row: &rusqlite::Row) {
writer
.serialize(PstatSeriesRow {
id: row.get(0).unwrap(),
krate: row.get_ref(1).unwrap().as_str().unwrap(),
profile: row.get_ref(2).unwrap().as_str().unwrap(),
scenario: row.get_ref(3).unwrap().as_str().unwrap(),
backend: row.get_ref(4).unwrap().as_str().unwrap(),
target: row.get_ref(5).unwrap().as_str().unwrap(),
metric: row.get_ref(6).unwrap().as_str().unwrap(),
})
.unwrap();
}
}
struct RawSelfProfile;
#[derive(Serialize)]
struct RawSelfProfileRow<'a> {
aid: i32,
cid: i32,
krate: &'a str,
profile: &'a str,
cache: &'a str,
}
impl Table for RawSelfProfile {
fn name() -> &'static str {
"raw_self_profile"
}
fn sqlite_attributes() -> &'static str {
"aid, cid, crate, profile, cache"
}
fn postgres_attributes() -> &'static str {
"aid, cid, crate, profile, cache"
}
fn postgres_generated_id_attribute() -> Option<&'static str> {
None
}
fn write_postgres_csv_row<W: Write>(writer: &mut csv::Writer<W>, row: &rusqlite::Row) {
writer
.serialize(RawSelfProfileRow {
aid: row.get(0).unwrap(),
cid: row.get(1).unwrap(),
krate: row.get_ref(2).unwrap().as_str().unwrap(),
profile: row.get_ref(3).unwrap().as_str().unwrap(),
cache: row.get_ref(4).unwrap().as_str().unwrap(),
})
.unwrap();
}
}
struct RustcCompilation;
#[derive(Serialize)]
struct RustcCompilationRow<'a> {
aid: i32,
cid: i32,
krate: &'a str,
duration: i64,
}
impl Table for RustcCompilation {
fn name() -> &'static str {
"rustc_compilation"
}
fn sqlite_attributes() -> &'static str {
"aid, cid, crate, duration"
}
fn postgres_attributes() -> &'static str {
"aid, cid, crate, duration"
}
fn postgres_generated_id_attribute() -> Option<&'static str> {
None
}
fn write_postgres_csv_row<W: Write>(writer: &mut csv::Writer<W>, row: &rusqlite::Row) {
writer
.serialize(RustcCompilationRow {
aid: row.get(0).unwrap(),
cid: row.get(1).unwrap(),
krate: row.get_ref(2).unwrap().as_str().unwrap(),
duration: row.get(3).unwrap(),
})
.unwrap();
}
}
struct RuntimePstat;
#[derive(Serialize)]
struct RuntimePstatRow {
series: i32,
aid: i32,
cid: i32,
value: f64,
}
impl Table for RuntimePstat {
fn name() -> &'static str {
"runtime_pstat"
}
fn sqlite_attributes() -> &'static str {
"series, aid, cid, value"
}
fn postgres_attributes() -> &'static str {
"series, aid, cid, value"
}
fn postgres_generated_id_attribute() -> Option<&'static str> {
None
}
fn write_postgres_csv_row<W: Write>(writer: &mut csv::Writer<W>, row: &rusqlite::Row) {
writer
.serialize(RuntimePstatRow {
series: row.get(0).unwrap(),
aid: row.get(1).unwrap(),
cid: row.get(2).unwrap(),
value: row.get(3).unwrap(),
})
.unwrap();
}
}
struct RuntimePstatSeries;
#[derive(Serialize)]
struct RuntimePstatSeriesRow<'a> {
id: i32,
benchmark: &'a str,
target: &'a str,
metric: &'a str,
}
impl Table for RuntimePstatSeries {
fn name() -> &'static str {
"runtime_pstat_series"
}
fn sqlite_attributes() -> &'static str {
"id, benchmark, target, metric"
}
fn postgres_attributes() -> &'static str {
"id, benchmark, target, metric"
}
fn postgres_generated_id_attribute() -> Option<&'static str> {
Some("id")
}
fn write_postgres_csv_row<W: Write>(writer: &mut csv::Writer<W>, row: &rusqlite::Row) {
writer
.serialize(RuntimePstatSeriesRow {
id: row.get(0).unwrap(),
benchmark: row.get_ref(1).unwrap().as_str().unwrap(),
target: row.get_ref(2).unwrap().as_str().unwrap(),
metric: row.get_ref(3).unwrap().as_str().unwrap(),
})
.unwrap();
}
}
struct ArtifactSize;
#[derive(Serialize)]
struct ArtifactSizeRow<'a> {
aid: i32,
component: &'a str,
size: i32,
}
impl Table for ArtifactSize {
fn name() -> &'static str {
"artifact_size"
}
fn sqlite_attributes() -> &'static str {
"aid, component, size"
}
fn postgres_attributes() -> &'static str {
"aid, component, size"
}
fn postgres_generated_id_attribute() -> Option<&'static str> {
None
}
fn write_postgres_csv_row<W: Write>(writer: &mut csv::Writer<W>, row: &rusqlite::Row) {
writer
.serialize(ArtifactSizeRow {
aid: row.get(0).unwrap(),
component: row.get_ref(1).unwrap().as_str().unwrap(),
size: row.get(2).unwrap(),
})
.unwrap();
}
}
struct Nullable<T>(Option<T>);
impl<T: Serialize> Serialize for Nullable<T> {
fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
where
S: Serializer,
{
match self.0 {
Some(ref t) => t.serialize(serializer),
None => NULL_STRING.serialize(serializer),
}
}
}
impl<T: rusqlite::types::FromSql> rusqlite::types::FromSql for Nullable<T> {
fn column_result(value: rusqlite::types::ValueRef<'_>) -> rusqlite::types::FromSqlResult<Self> {
Ok(Nullable(rusqlite::types::FromSql::column_result(value)?))
}
}
impl<'a> TryFrom<rusqlite::types::ValueRef<'a>> for Nullable<&'a str> {
type Error = rusqlite::types::FromSqlError;
fn try_from(value: rusqlite::types::ValueRef<'a>) -> Result<Self, Self::Error> {
use rusqlite::types::ValueRef;
match value {
ValueRef::Null => Ok(Nullable(None)),
ValueRef::Text(_) => Ok(Nullable(Some(value.as_str()?))),
_ => Err(rusqlite::types::FromSqlError::InvalidType),
}
}
}
#[tokio::main]
async fn main() -> anyhow::Result<()> {
env_logger::init();
let matches = clap::Command::new("sqlite-to-postgres")
.about("Exports a rustc-perf SQLite database to a Postgres database")
.version(clap::crate_version!())
.arg(
clap::Arg::new("sqlite-db")
.required(true)
.value_name("SQLITE_DB")
.help("SQLite database file"),
)
.arg(
clap::Arg::new("postgres-db")
.required(true)
.value_name("POSTGRES_DB")
.help(
"Postgres database connection string, \
e.g. postgres://user:password@localhost:5432",
),
)
.get_matches();
let postgres = matches.get_one::<String>("postgres-db").unwrap();
let sqlite = matches.get_one::<String>("sqlite-db").unwrap();
let mut sqlite = sqlite::Sqlite::new(sqlite.into())
.open()
.await
.into_inner()
.unwrap();
let mut postgres: tokio_postgres::Client =
postgres::Postgres::new(postgres.into()).open().await.into();
let sqlite_tx = sqlite.transaction().unwrap();
let postgres_tx = postgres.transaction().await?;
let tables = get_tables(&postgres_tx).await;
disable_table_triggers(&postgres_tx, &tables).await;
copy::<Artifact>(&sqlite_tx, &postgres_tx).await;
copy::<Benchmark>(&sqlite_tx, &postgres_tx).await;
copy::<Collection>(&sqlite_tx, &postgres_tx).await;
copy::<Error>(&sqlite_tx, &postgres_tx).await;
copy::<PstatSeries>(&sqlite_tx, &postgres_tx).await;
copy::<Pstat>(&sqlite_tx, &postgres_tx).await;
copy::<RawSelfProfile>(&sqlite_tx, &postgres_tx).await;
copy::<RustcCompilation>(&sqlite_tx, &postgres_tx).await;
copy::<RuntimePstatSeries>(&sqlite_tx, &postgres_tx).await;
copy::<RuntimePstat>(&sqlite_tx, &postgres_tx).await;
copy::<ArtifactSize>(&sqlite_tx, &postgres_tx).await;
enable_table_triggers(&postgres_tx, &tables).await;
sqlite_tx.rollback().unwrap();
postgres_tx.commit().await?;
Ok(())
}
async fn copy<T: Table>(
sqlite: &rusqlite::Transaction<'_>,
postgres: &tokio_postgres::Transaction<'_>,
) {
let table = T::name();
let postgres_columns = postgres
.query(
r#"SELECT column_name FROM information_schema.columns
WHERE table_schema = 'public' AND table_name = $1
ORDER BY ordinal_position"#,
&[&table],
)
.await
.unwrap()
.into_iter()
.map(|row| row.get(0))
.collect::<Vec<String>>();
let attributes = mapping_pg_columns_to_attributes::<T>(&postgres_columns);
let copy = postgres
.prepare(&format!(
r#"copy {table} ({attributes}) from stdin (encoding utf8, format csv, null '{NULL_STRING}')"#,
))
.await
.unwrap();
let copy_in_sink = postgres.copy_in::<_, Bytes>(©).await.unwrap();
tokio::pin!(copy_in_sink);
let mut csv_writer = postgres_csv_writer(BytesMut::new().writer());
let mut select = sqlite
.prepare(&format!("select {} from {}", T::sqlite_attributes(), table))
.unwrap();
let start = Instant::now();
let mut rows = select.query([]).unwrap();
let mut count = 0;
const ROWS_PER_SEND: usize = 1024;
while let Some(result) = rows.next().transpose() {
let row = result.unwrap();
T::write_postgres_csv_row(&mut csv_writer, row);
count += 1;
if count % ROWS_PER_SEND == 0 {
let mut bytes_writer = csv_writer.into_inner().unwrap();
let bytes = bytes_writer.get_mut().split().freeze();
copy_in_sink.send(bytes).await.unwrap();
csv_writer = postgres_csv_writer(bytes_writer);
}
}
if count % ROWS_PER_SEND != 0 {
let bytes = csv_writer.into_inner().unwrap().into_inner().freeze();
copy_in_sink.send(bytes).await.unwrap();
}
copy_in_sink.close().await.unwrap();
if count > 0 {
if let Some(generated_id_attr) = T::postgres_generated_id_attribute() {
postgres
.execute(
&format!(
"select setval(
pg_get_serial_sequence($1, $2),
coalesce(max({generated_id_attr}) + 1, 1), false)
from {table}"
) as &str,
&[&table, &generated_id_attr],
)
.await
.unwrap();
}
}
let elapsed = start.elapsed();
eprintln!(
"Copied {} rows from {} table in {:?} ({:.0} rows/second)",
count,
table,
elapsed,
count as f64 / elapsed.as_secs_f64()
);
}
fn postgres_csv_writer<W: Write>(w: W) -> csv::Writer<W> {
csv::WriterBuilder::new().has_headers(false).from_writer(w)
}
fn mapping_pg_columns_to_attributes<T: Table>(postgres_columns: &[impl AsRef<str>]) -> String {
let sl = T::sqlite_attributes()
.split(',')
.map(str::trim)
.collect::<Vec<_>>();
let pg = T::postgres_attributes()
.split(',')
.map(str::trim)
.collect::<Vec<_>>();
assert_eq!(
sl.len(),
pg.len(),
"The number of attributes in SQLite and Postgres is mismatched."
);
let map = pg
.iter()
.enumerate()
.map(|(i, p)| (p, i))
.collect::<HashMap<_, _>>();
let mut out_attrs = vec![""; pg.len()];
for col in postgres_columns.iter().map(AsRef::as_ref) {
let idx = map.get(&col).unwrap_or_else(|| {
panic!(
"Failed to find a corresponding attribute for column {} in table {}.",
col,
T::name()
)
});
out_attrs[*idx] = col;
}
out_attrs.join(", ")
}
async fn get_tables(postgres: &tokio_postgres::Transaction<'_>) -> Vec<String> {
postgres
.query(
"SELECT tablename FROM pg_catalog.pg_tables WHERE schemaname = 'public'",
&[],
)
.await
.unwrap()
.into_iter()
.map(|row| row.get(0))
.collect()
}
async fn disable_table_triggers(postgres: &tokio_postgres::Transaction<'_>, tables: &[String]) {
for table in tables {
postgres
.execute(&format!("ALTER TABLE {table} DISABLE TRIGGER ALL"), &[])
.await
.unwrap();
}
eprintln!("Disabled table triggers");
}
async fn enable_table_triggers(postgres: &tokio_postgres::Transaction<'_>, tables: &[String]) {
for table in tables {
postgres
.execute(&format!("ALTER TABLE {table} ENABLE TRIGGER ALL"), &[])
.await
.unwrap();
}
eprintln!("Enabled table triggers");
}