Skip to content

Commit 8aa3c08

Browse files
authored
Merge pull request #1824 from eth3lbert/speedup-sqlite2pg
Optimize sqlite-to-postgres migration speed for large tables via temporary constraint suspension
2 parents 3b90189 + ef67816 commit 8aa3c08

File tree

1 file changed

+41
-0
lines changed

1 file changed

+41
-0
lines changed

database/src/bin/sqlite-to-postgres.rs

Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -554,6 +554,13 @@ async fn main() -> anyhow::Result<()> {
554554
// if everything succeeds. The transaction is not being used for isolation.
555555
let postgres_tx = postgres.transaction().await?;
556556

557+
// Disabling table triggers before copying data and then re-enabling them afterwards can
558+
// yield a significant speedup.
559+
// Note: Disabling or enabling internally generated constraint triggers requires superuser
560+
// privileges. [See PostgreSQL documentation](
561+
// https://www.postgresql.org/docs/current/sql-altertable.html#SQL-ALTERTABLE-DESC-DISABLE-ENABLE-TRIGGER)
562+
let tables = get_tables(&postgres_tx).await;
563+
disable_table_triggers(&postgres_tx, &tables).await;
557564
// Order matters to the extent necessary to satisfy foreign key constraints.
558565
copy::<Artifact>(&sqlite_tx, &postgres_tx).await;
559566
copy::<ArtifactCollectionDuration>(&sqlite_tx, &postgres_tx).await;
@@ -566,6 +573,7 @@ async fn main() -> anyhow::Result<()> {
566573
copy::<PullRequestBuild>(&sqlite_tx, &postgres_tx).await;
567574
copy::<RawSelfProfile>(&sqlite_tx, &postgres_tx).await;
568575
copy::<RustcCompilation>(&sqlite_tx, &postgres_tx).await;
576+
enable_table_triggers(&postgres_tx, &tables).await;
569577

570578
// This is overly paranoid, but don't commit the Postgres transaction until
571579
// the rollback of the SQLite transaction succeeds.
@@ -698,3 +706,36 @@ async fn copy<T: Table>(
698706
fn postgres_csv_writer<W: Write>(w: W) -> csv::Writer<W> {
699707
csv::WriterBuilder::new().has_headers(false).from_writer(w)
700708
}
709+
710+
async fn get_tables(postgres: &tokio_postgres::Transaction<'_>) -> Vec<String> {
711+
postgres
712+
.query(
713+
"SELECT tablename FROM pg_catalog.pg_tables WHERE schemaname = 'public'",
714+
&[],
715+
)
716+
.await
717+
.unwrap()
718+
.into_iter()
719+
.map(|row| row.get(0))
720+
.collect()
721+
}
722+
723+
async fn disable_table_triggers(postgres: &tokio_postgres::Transaction<'_>, tables: &[String]) {
724+
for table in tables {
725+
postgres
726+
.execute(&format!("ALTER TABLE {} DISABLE TRIGGER ALL", table), &[])
727+
.await
728+
.unwrap();
729+
}
730+
eprintln!("Disabled table triggers");
731+
}
732+
733+
async fn enable_table_triggers(postgres: &tokio_postgres::Transaction<'_>, tables: &[String]) {
734+
for table in tables {
735+
postgres
736+
.execute(&format!("ALTER TABLE {} ENABLE TRIGGER ALL", table), &[])
737+
.await
738+
.unwrap();
739+
}
740+
eprintln!("Enabled table triggers");
741+
}

0 commit comments

Comments
 (0)