fix: Provide separate URL for masks DB (#331)

* fix: Provide separate URL for masks DB * fmt * fmt
worldcoin · Sep 6, 2024 · 4d29249 · 4d29249
1 parent 6c91714
commit 4d29249
Show file tree

Hide file tree

Showing 4 changed files with 58 additions and 25 deletions.
diff --git a/iris-mpc-upgrade/src/bin/Readme.md b/iris-mpc-upgrade/src/bin/Readme.md
@@ -13,7 +13,7 @@ will bring up two old dbs on ports 6200 and 6201, and 3 new dbs on ports 6100,61
 ## Fill DBs with test data
 
 ```bash
-cargo run --release --bin seed-v1-dbs -- --db-urls postgres://postgres:postgres@localhost:6100/postgres --db-urls postgres://postgres:postgres@localhost:6101/postgres --num-elements 10000
+cargo run --release --bin seed-v1-dbs -- --shares-db-urls postgres://postgres:postgres@localhost:6100/shares --shares-db-urls postgres://postgres:postgres@localhost:6101/shares --masks-db-url postgres://postgres:postgres@localhost:6100/masks --num-elements 10000
 ```
 
 ## Upgrade for left eye
@@ -39,11 +39,11 @@ cargo run --release --bin upgrade-server -- --bind-addr 127.0.0.1:8002 --db-url
 Concurrently run:
 
 ```bash
-cargo run --release --bin upgrade-client -- --server1 127.0.0.1:8000 --server2 127.0.0.1:8001 --server3 127.0.0.1:8002 --db-start 0 --db-end 10000 --party-id 0 --eye left --db-url postgres://postgres:postgres@localhost:6100/postgres
+cargo run --release --bin upgrade-client -- --server1 127.0.0.1:8000 --server2 127.0.0.1:8001 --server3 127.0.0.1:8002 --db-start 0 --db-end 10000 --party-id 0 --eye left --shares-db-url postgres://postgres:postgres@localhost:6100/shares --masks-db-url postgres://postgres:postgres@localhost:6100/masks
 ```
 
 ```bash
-cargo run --release --bin upgrade-client -- --server1 127.0.0.1:8000 --server2 127.0.0.1:8001 --server3 127.0.0.1:8002 --db-start 0 --db-end 10000 --party-id 1 --eye left --db-url postgres://postgres:postgres@localhost:6101/postgres
+cargo run --release --bin upgrade-client -- --server1 127.0.0.1:8000 --server2 127.0.0.1:8001 --server3 127.0.0.1:8002 --db-start 0 --db-end 10000 --party-id 1 --eye left --shares-db-url postgres://postgres:postgres@localhost:6101/shares --masks-db-url postgres://postgres:postgres@localhost:6100/masks
 ```
 
 ## Upgrade for right eye
@@ -71,11 +71,11 @@ cargo run --release --bin upgrade-server -- --bind-addr 127.0.0.1:8002 --db-url
 Concurrently run:
 
 ```bash
-cargo run --release --bin upgrade-client -- --server1 127.0.0.1:8000 --server2 127.0.0.1:8001 --server3 127.0.0.1:8002 --db-start 0 --db-end 10000 --party-id 0 --eye right --db-url postgres://postgres:postgres@localhost:6100/postgres
+cargo run --release --bin upgrade-client -- --server1 127.0.0.1:8000 --server2 127.0.0.1:8001 --server3 127.0.0.1:8002 --db-start 0 --db-end 10000 --party-id 0 --eye right --shares-db-url postgres://postgres:postgres@localhost:6100/shares --masks-db-url postgres://postgres:postgres@localhost:6100/masks
 ```
 
 ```bash
-cargo run --release --bin upgrade-client -- --server1 127.0.0.1:8000 --server2 127.0.0.1:8001 --server3 127.0.0.1:8002 --db-start 0 --db-end 10000 --party-id 1 --eye right --db-url postgres://postgres:postgres@localhost:6101/postgres
+cargo run --release --bin upgrade-client -- --server1 127.0.0.1:8000 --server2 127.0.0.1:8001 --server3 127.0.0.1:8002 --db-start 0 --db-end 10000 --party-id 1 --eye right --shares-db-url postgres://postgres:postgres@localhost:6101/shares --masks-db-url postgres://postgres:postgres@localhost:6100/masks
 ```
 
 ## Check the upgrade was successful

diff --git a/iris-mpc-upgrade/src/bin/seed_v1_dbs.rs b/iris-mpc-upgrade/src/bin/seed_v1_dbs.rs
@@ -5,7 +5,11 @@ use rand::Rng;
 #[derive(Debug, Clone, Parser)]
 struct Args {
     #[clap(long)]
-    db_urls:      Vec<String>,
+    shares_db_urls: Vec<String>,
+
+    #[clap(long)]
+    masks_db_url: String,
+
     #[clap(long)]
     num_elements: u64,
 }
@@ -14,22 +18,35 @@ struct Args {
 async fn main() -> eyre::Result<()> {
     let args = Args::parse();
 
-    if args.db_urls.len() != 2 {
-        return Err(eyre::eyre!("Expect 2 db urls to be provided"));
+    if args.shares_db_urls.len() != 2 {
+        return Err(eyre::eyre!("Expect 2 shares db urls to be provided"));
     }
 
-    let db_config0 = DbConfig {
-        url:     args.db_urls[0].clone(),
+    if args.masks_db_url.is_empty() {
+        return Err(eyre::eyre!("Expect 1 masks db urls to be provided"));
+    }
+
+    let shares_db_config0 = DbConfig {
+        url:     args.shares_db_urls[0].clone(),
+        migrate: true,
+        create:  true,
+    };
+
+    let shares_db_config1 = DbConfig {
+        url:     args.shares_db_urls[1].clone(),
         migrate: true,
         create:  true,
     };
-    let db_config1 = DbConfig {
-        url:     args.db_urls[1].clone(),
+
+    let masks_db_config = DbConfig {
+        url:     args.masks_db_url.clone(),
         migrate: true,
         create:  true,
     };
-    let db0 = Db::new(&db_config0).await?;
-    let db1 = Db::new(&db_config1).await?;
+
+    let shares_db0 = Db::new(&shares_db_config0).await?;
+    let shares_db1 = Db::new(&shares_db_config1).await?;
+    let masks_db = Db::new(&masks_db_config).await?;
 
     let mut rng = rand::thread_rng();
 
@@ -49,10 +66,10 @@ async fn main() -> eyre::Result<()> {
         shares0.push((i, encoded[0]));
         shares1.push((i, encoded[1]));
     }
-    db0.insert_masks(&masks).await?;
-    db1.insert_masks(&masks).await?;
-    db0.insert_shares(&shares0).await?;
-    db1.insert_shares(&shares1).await?;
+    masks_db.insert_masks(&masks).await?;
+
+    shares_db0.insert_shares(&shares0).await?;
+    shares_db1.insert_shares(&shares1).await?;
 
     Ok(())
 }
diff --git a/iris-mpc-upgrade/src/bin/tcp_upgrade_client.rs b/iris-mpc-upgrade/src/bin/tcp_upgrade_client.rs
@@ -78,7 +78,10 @@ async fn main() -> eyre::Result<()> {
 
     // need this to maybe store the PG db or otherwise the borrow checker complains
     #[allow(unused_assignments)]
-    let mut maybe_db = None;
+    let mut maybe_shares_db = None;
+
+    #[allow(unused_assignments)]
+    let mut maybe_masks_db = None;
 
     #[allow(clippy::type_complexity)]
     let (mut shares_stream, mut mask_stream): (
@@ -97,13 +100,22 @@ async fn main() -> eyre::Result<()> {
             _ => unreachable!(),
         }
     } else {
-        maybe_db = Some(V1Database {
-            db: V1Db::new(&args.db_url).await?,
+        maybe_shares_db = Some(V1Database {
+            db: V1Db::new(&args.shares_db_url).await?,
+        });
+
+        maybe_masks_db = Some(V1Database {
+            db: V1Db::new(&args.masks_db_url).await?,
         });
 
         (
-            Box::pin(maybe_db.as_ref().unwrap().stream_shares(db_range.clone())?),
-            Box::pin(maybe_db.as_ref().unwrap().stream_masks(db_range)?),
+            Box::pin(
+                maybe_shares_db
+                    .as_ref()
+                    .unwrap()
+                    .stream_shares(db_range.clone())?,
+            ),
+            Box::pin(maybe_masks_db.as_ref().unwrap().stream_masks(db_range)?),
         )
     };
 

diff --git a/iris-mpc-upgrade/src/config.rs b/iris-mpc-upgrade/src/config.rs
@@ -83,7 +83,10 @@ pub struct UpgradeClientConfig {
     pub mock: bool,
 
     #[clap(long)]
-    pub db_url: String,
+    pub shares_db_url: String,
+
+    #[clap(long)]
+    pub masks_db_url: String,
 }
 
 fn resolve_host(hostname_port: &str) -> io::Result<SocketAddr> {
@@ -102,7 +105,8 @@ impl std::fmt::Debug for UpgradeClientConfig {
             .field("server1", &self.server1)
             .field("server2", &self.server2)
             .field("server3", &self.server3)
-            .field("db_url", &"<redacted>")
+            .field("shares_db_url", &"<redacted>")
+            .field("masks_db_url", &"<redacted>")
             .field("db_start", &self.db_start)
             .field("db_end", &self.db_end)
             .field("party_id", &self.party_id)