Optimize sqlite-to-postgres migration speed for large tables via temporary constraint suspension

Author: eth3lbert

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

@@ -37,6 +37,11 @@ trait Table {
     fn write_postgres_csv_row<W: Write>(writer: &mut csv::Writer<W>, row: &rusqlite::Row);
 }
 
+trait TableConstraint {
+    /// Table constraints
+    fn constraints() -> &'static str;
+}
+
 struct Artifact;
 
 #[derive(Serialize)]
@@ -282,6 +287,19 @@ impl Table for Pstat {
     }
 }
 
+impl TableConstraint for Pstat {
+    // These constraints can be inspected using `\d` in psql. However, only the primary key needs
+    // slight modification to comply with the execution statement.
+    fn constraints() -> &'static str {
+        r#"
+        "pstat_pkey" PRIMARY KEY (series, aid, cid);
+        "pstat_aid_fkey" FOREIGN KEY (aid) REFERENCES artifact(id) ON UPDATE CASCADE ON DELETE CASCADE;
+        "pstat_cid_fkey" FOREIGN KEY (cid) REFERENCES collection(id) ON UPDATE CASCADE ON DELETE CASCADE;
+        "pstat_series_fkey" FOREIGN KEY (series) REFERENCES pstat_series(id) ON UPDATE CASCADE ON DELETE CASCADE;
+        "#
+    }
+}
+
 struct PstatSeries;
 
 #[derive(Serialize)]
@@ -410,6 +428,19 @@ impl Table for RawSelfProfile {
     }
 }
 
+impl TableConstraint for RawSelfProfile {
+    // These constraints can be inspected using `\d` in psql. However, only the primary key needs
+    // slight modification to comply with the execution statement.
+    fn constraints() -> &'static str {
+        r#"
+        "raw_self_profile_pkey" PRIMARY KEY (aid, cid, crate, profile, cache);
+        "raw_self_profile_aid_fkey" FOREIGN KEY (aid) REFERENCES artifact(id) ON UPDATE CASCADE ON DELETE CASCADE;
+        "raw_self_profile_cid_fkey" FOREIGN KEY (cid) REFERENCES collection(id) ON UPDATE CASCADE ON DELETE CASCADE;
+        "raw_self_profile_crate_fkey" FOREIGN KEY (crate) REFERENCES benchmark(name) ON UPDATE CASCADE ON DELETE CASCADE;
+        "#
+    }
+}
+
 struct RustcCompilation;
 
 #[derive(Serialize)]
@@ -562,9 +593,19 @@ async fn main() -> anyhow::Result<()> {
     copy::<CollectorProgress>(&sqlite_tx, &postgres_tx).await;
     copy::<Error>(&sqlite_tx, &postgres_tx).await;
     copy::<PstatSeries>(&sqlite_tx, &postgres_tx).await;
+
+    // Dropping constraints of large tables before copying data and then adding them back.
+    // This approach should yield a significant speedup.
+    drop_constraints::<Pstat>(&postgres_tx).await;
     copy::<Pstat>(&sqlite_tx, &postgres_tx).await;
+    add_constraints::<Pstat>(&postgres_tx).await;
+
     copy::<PullRequestBuild>(&sqlite_tx, &postgres_tx).await;
+
+    drop_constraints::<RawSelfProfile>(&postgres_tx).await;
     copy::<RawSelfProfile>(&sqlite_tx, &postgres_tx).await;
+    add_constraints::<RawSelfProfile>(&postgres_tx).await;
+
     copy::<RustcCompilation>(&sqlite_tx, &postgres_tx).await;
 
     // This is overly paranoid, but don't commit the Postgres transaction until
@@ -698,3 +739,44 @@ async fn copy<T: Table>(
 fn postgres_csv_writer<W: Write>(w: W) -> csv::Writer<W> {
     csv::WriterBuilder::new().has_headers(false).from_writer(w)
 }
+
+async fn drop_constraints<T: Table + TableConstraint>(postgres: &tokio_postgres::Transaction<'_>) {
+    let table = T::name();
+    let constraints = T::constraints();
+
+    let start = Instant::now();
+    for constraint in constraints.split(';').map(str::trim) {
+        if !constraint.is_empty() {
+            let name = constraint.split_whitespace().next().unwrap();
+            postgres
+                .execute(
+                    &format!("ALTER TABLE {} DROP CONSTRAINT {}", table, name),
+                    &[],
+                )
+                .await
+                .unwrap();
+        }
+    }
+    let elapsed = start.elapsed();
+    eprintln!("Drop constraints for table {} in {:?}", table, elapsed);
+}
+
+async fn add_constraints<T: Table + TableConstraint>(postgres: &tokio_postgres::Transaction<'_>) {
+    let table = T::name();
+    let constraints = T::constraints();
+
+    let start = Instant::now();
+    for constraint in constraints.split(';').map(str::trim) {
+        if !constraint.is_empty() {
+            postgres
+                .execute(
+                    &format!("ALTER TABLE {} ADD CONSTRAINT {}", table, constraint),
+                    &[],
+                )
+                .await
+                .unwrap();
+        }
+    }
+    let elapsed = start.elapsed();
+    eprintln!("Add constraints for table {} in {:?}", table, elapsed);
+}