Resolves rust-lang#530 - add_package logic, make database calls in for loops

Author: syphar

src/db/add_package.rs

@@ -1,4 +1,5 @@
 use std::{
+    collections::{HashMap, HashSet},
     fs,
     io::{BufRead, BufReader},
     path::Path,
@@ -186,15 +187,16 @@ pub(crate) fn add_build_into_database(
 }
 
 fn initialize_package_in_database(conn: &mut Client, pkg: &MetadataPackage) -> Result<i32> {
-    let mut rows = conn.query("SELECT id FROM crates WHERE name = $1", &[&pkg.name])?;
-    // insert crate into database if it is not exists
-    if rows.is_empty() {
-        rows = conn.query(
-            "INSERT INTO crates (name) VALUES ($1) RETURNING id",
+    Ok(conn
+        .query_one(
+            "INSERT INTO crates (name) 
+            VALUES ($1) 
+            ON CONFLICT (name) DO UPDATE 
+                SET name = EXCLUDED.name
+            RETURNING id",
             &[&pkg.name],
-        )?;
-    }
-    Ok(rows[0].get(0))
+        )?
+        .get(0))
 }
 
 /// Convert dependencies into Vec<(String, String, String)>
@@ -316,28 +318,41 @@ fn add_keywords_into_database(
     pkg: &MetadataPackage,
     release_id: i32,
 ) -> Result<()> {
-    for keyword in &pkg.keywords {
-        let slug = slugify(&keyword);
-        let keyword_id: i32 = {
-            let rows = conn.query("SELECT id FROM keywords WHERE slug = $1", &[&slug])?;
-            if !rows.is_empty() {
-                rows[0].get(0)
-            } else {
-                conn.query(
-                    "INSERT INTO keywords (name, slug) VALUES ($1, $2) RETURNING id",
-                    &[&keyword, &slug],
-                )?[0]
-                    .get(0)
-            }
-        };
-
-        // add releationship
-        let _ = conn.query(
-            "INSERT INTO keyword_rels (rid, kid) VALUES ($1, $2)",
-            &[&release_id, &keyword_id],
-        );
+    let wanted_keywords: HashMap<String, String> = pkg
+        .keywords
+        .iter()
+        .map(|kw| (slugify(&kw), kw.clone()))
+        .collect();
+
+    let existing_keyword_slugs: HashSet<String> = conn
+        .query(
+            "SELECT slug FROM keywords WHERE slug = ANY($1)",
+            &[&wanted_keywords.keys().collect::<Vec<_>>()],
+        )?
+        .iter()
+        .map(|row| row.get(0))
+        .collect();
+
+    // we create new keywords one-by-one, since most of the time we already have them,
+    // and because support for multi-record inserts is a mess without adding a new
+    // library
+    let insert_keyword_query = conn.prepare("INSERT INTO keywords (name, slug) VALUES ($1, $2)")?;
+    for (slug, name) in wanted_keywords
+        .iter()
+        .filter(|(k, _)| !(existing_keyword_slugs.contains(*k)))
+        .collect::<Vec<(_, _)>>()
+    {
+        conn.query(&insert_keyword_query, &[&name, &slug])?;
     }
 
+    conn.query(
+        "INSERT INTO keyword_rels (rid, kid) 
+        SELECT $1 as rid, id as kid 
+        FROM keywords 
+        WHERE slug = ANY($2)",
+        &[&release_id, &wanted_keywords.keys().collect::<Vec<_>>()],
+    )?;
+
     Ok(())
 }
 
@@ -347,7 +362,9 @@ pub fn update_crate_data_in_database(
     registry_data: &CrateData,
 ) -> Result<()> {
     info!("Updating crate data for {}", name);
-    let crate_id = conn.query("SELECT id FROM crates WHERE crates.name = $1", &[&name])?[0].get(0);
+    let crate_id = conn
+        .query_one("SELECT id FROM crates WHERE crates.name = $1", &[&name])?
+        .get(0);
 
     update_owners_in_database(conn, &registry_data.owners, crate_id)?;
 
@@ -360,64 +377,51 @@ fn update_owners_in_database(
     owners: &[CrateOwner],
     crate_id: i32,
 ) -> Result<()> {
-    let rows = conn.query(
-        "
-            SELECT login
-            FROM owners
-            INNER JOIN owner_rels
-                ON owner_rels.oid = owners.id
-            WHERE owner_rels.cid = $1
-        ",
-        &[&crate_id],
+    // Update any existing owner data since it is mutable and could have changed since last
+    // time we pulled it
+    let owner_upsert = conn.prepare(
+        "INSERT INTO owners (login, avatar, name, email)
+        VALUES ($1, $2, $3, $4)
+        ON CONFLICT (login) DO UPDATE
+            SET
+                avatar = EXCLUDED.avatar,
+                name = EXCLUDED.name,
+                email = EXCLUDED.email",
     )?;
-    let existing_owners = rows.into_iter().map(|row| -> String { row.get(0) });
-
     for owner in owners {
-        debug!("Updating owner data for {}: {:?}", owner.login, owner);
-
-        // Update any existing owner data since it is mutable and could have changed since last
-        // time we pulled it
-        let owner_id: i32 = {
-            conn.query(
-                "
-                    INSERT INTO owners (login, avatar, name, email)
-                    VALUES ($1, $2, $3, $4)
-                    ON CONFLICT (login) DO UPDATE
-                        SET
-                            avatar = $2,
-                            name = $3,
-                            email = $4
-                    RETURNING id
-                ",
-                &[&owner.login, &owner.avatar, &owner.name, &owner.email],
-            )?[0]
-                .get(0)
-        };
-
-        // add relationship
         conn.query(
-            "INSERT INTO owner_rels (cid, oid) VALUES ($1, $2) ON CONFLICT DO NOTHING",
-            &[&crate_id, &owner_id],
+            &owner_upsert,
+            &[&owner.login, &owner.avatar, &owner.name, &owner.email],
         )?;
     }
 
-    let to_remove =
-        existing_owners.filter(|login| !owners.iter().any(|owner| &owner.login == login));
-
-    for login in to_remove {
-        debug!("Removing owner relationship {}", login);
-        // remove relationship
-        conn.query(
-            "
-                DELETE FROM owner_rels
-                USING owners
-                WHERE owner_rels.cid = $1
-                    AND owner_rels.oid = owners.id
-                    AND owners.login = $2
-            ",
-            &[&crate_id, &login],
-        )?;
-    }
+    let updated_oids: Vec<i32> = conn
+        .query(
+            "INSERT INTO owner_rels (cid, oid)
+             SELECT $1,id
+             FROM owners 
+             WHERE login = ANY($2)
+             ON CONFLICT (cid,oid) 
+             DO UPDATE -- we need this so the existing/updated records end 
+                       -- up being in the returned OIDs
+                 SET oid=excluded.oid 
+             RETURNING oid",
+            &[
+                &crate_id,
+                &owners.iter().map(|o| o.login.clone()).collect::<Vec<_>>(),
+            ],
+        )?
+        .iter()
+        .map(|row| row.get(0))
+        .collect();
+
+    conn.query(
+        "DELETE FROM owner_rels
+         WHERE 
+            cid = $1 AND 
+            NOT (oid = ANY($2))",
+        &[&crate_id, &updated_oids],
+    )?;
 
     Ok(())
 }