From 408722dd9429cafd0402b20ce51a064177917b29 Mon Sep 17 00:00:00 2001
From: Philipp Sippl <philsippl@gmail.com>
Date: Thu, 26 Sep 2024 16:58:27 +0200
Subject: [PATCH] update the pg serial after all insertions (#460)

---
 iris-mpc-store/src/lib.rs                      | 11 +++++++++++
 iris-mpc-upgrade/src/bin/tcp_upgrade_server.rs |  9 ++++++++-
 iris-mpc-upgrade/src/lib.rs                    |  6 ++++++
 3 files changed, 25 insertions(+), 1 deletion(-)

diff --git a/iris-mpc-store/src/lib.rs b/iris-mpc-store/src/lib.rs
index 79af952f6..12df7daf4 100644
--- a/iris-mpc-store/src/lib.rs
+++ b/iris-mpc-store/src/lib.rs
@@ -304,6 +304,7 @@ DO UPDATE SET right_code = EXCLUDED.right_code, right_mask = EXCLUDED.right_mask
 
         Ok(())
     }
+
     pub async fn insert_results(
         &self,
         tx: &mut Transaction<'_, Postgres>,
@@ -321,6 +322,16 @@ DO UPDATE SET right_code = EXCLUDED.right_code, right_mask = EXCLUDED.right_mask
         Ok(())
     }
 
+    pub async fn update_iris_id_sequence(&self) -> Result<()> {
+        sqlx::query(
+            "SELECT setval(pg_get_serial_sequence('irises', 'id'), COALESCE(MAX(id), 0), false) \
+             FROM irises",
+        )
+        .execute(&self.pool)
+        .await?;
+        Ok(())
+    }
+
     pub async fn last_results(&self, count: usize) -> Result<Vec<String>> {
         let mut result_events: Vec<String> =
             sqlx::query_scalar("SELECT result_event FROM results ORDER BY id DESC LIMIT $1")
diff --git a/iris-mpc-upgrade/src/bin/tcp_upgrade_server.rs b/iris-mpc-upgrade/src/bin/tcp_upgrade_server.rs
index 46a416db8..7e7e80e85 100644
--- a/iris-mpc-upgrade/src/bin/tcp_upgrade_server.rs
+++ b/iris-mpc-upgrade/src/bin/tcp_upgrade_server.rs
@@ -61,7 +61,7 @@ async fn main() -> eyre::Result<()> {
     background_tasks.check_tasks();
     tracing::info!("Healthcheck server running on port 3000.");
 
-    let upgrader = IrisCodeUpgrader::new(args.party_id, sink);
+    let upgrader = IrisCodeUpgrader::new(args.party_id, sink.clone());
 
     // listen for incoming connections from clients
     let client_listener = tokio::net::TcpListener::bind(args.bind_addr).await?;
@@ -191,6 +191,9 @@ async fn main() -> eyre::Result<()> {
     }
     client_stream2.write_u8(FINAL_BATCH_SUCCESSFUL_ACK).await?;
     client_stream1.write_u8(FINAL_BATCH_SUCCESSFUL_ACK).await?;
+
+    sink.update_iris_id_sequence().await?;
+
     Ok(())
 }
 
@@ -227,4 +230,8 @@ impl NewIrisShareSink for IrisShareDbSink {
             }
         }
     }
+
+    async fn update_iris_id_sequence(&self) -> eyre::Result<()> {
+        self.store.update_iris_id_sequence().await
+    }
 }
diff --git a/iris-mpc-upgrade/src/lib.rs b/iris-mpc-upgrade/src/lib.rs
index 4de0f7ad9..e11c92961 100644
--- a/iris-mpc-upgrade/src/lib.rs
+++ b/iris-mpc-upgrade/src/lib.rs
@@ -41,6 +41,8 @@ pub trait NewIrisShareSink {
         code_share: &[u16; IRIS_CODE_LENGTH],
         mask_share: &[u16; MASK_CODE_LENGTH],
     ) -> eyre::Result<()>;
+
+    async fn update_iris_id_sequence(&self) -> eyre::Result<()>;
 }
 
 #[derive(Debug, Clone)]
@@ -80,6 +82,10 @@ impl NewIrisShareSink for IrisShareTestFileSink {
         file.flush()?;
         Ok(())
     }
+
+    async fn update_iris_id_sequence(&self) -> eyre::Result<()> {
+        Ok(())
+    }
 }
 
 #[derive(Clone)]