Merge pull request #781 from worldcoin/piohei/add_pre_root_for_identi…

…ties

Add pre_root for identities.

schemas/database/015_add_constraints_for_identities.down.sql

@@ -0,0 +1,7 @@
+DROP INDEX idx_unique_insertion_leaf;
+DROP INDEX idx_unique_deletion_leaf;
+
+DROP TRIGGER validate_pre_root_trigger ON identities;
+DROP FUNCTION validate_pre_root();
+
+ALTER TABLE identities DROP COLUMN pre_root;

schemas/database/015_add_constraints_for_identities.up.sql

@@ -0,0 +1,52 @@
+CREATE UNIQUE INDEX idx_unique_insertion_leaf on identities(leaf_index) WHERE commitment != E'\\x0000000000000000000000000000000000000000000000000000000000000000';
+CREATE UNIQUE INDEX idx_unique_deletion_leaf on identities(leaf_index) WHERE commitment = E'\\x0000000000000000000000000000000000000000000000000000000000000000';
+
+-- Add the new 'prev_root' column
+ALTER TABLE identities ADD COLUMN pre_root BYTEA;
+
+-- This constraint ensures that we have consistent database and changes to the tree are done in a valid sequence.
+CREATE OR REPLACE FUNCTION validate_pre_root() returns trigger as $$
+    DECLARE
+        last_id identities.id%type;
+        last_root identities.root%type;
+    BEGIN
+        SELECT id, root
+        INTO last_id, last_root
+        FROM identities
+        ORDER BY id DESC
+        LIMIT 1;
+
+        -- When last_id is NULL that means there are no records in identities table. The first prev_root can
+        -- be a value not referencing previous root in database.
+        IF last_id IS NULL THEN RETURN NEW;
+        END IF;
+
+        IF NEW.pre_root IS NULL THEN RAISE EXCEPTION 'Sent pre_root (%) can be null only for first record in table.', NEW.pre_root;
+        END IF;
+
+        IF (last_root != NEW.pre_root) THEN RAISE EXCEPTION 'Sent pre_root (%) is different than last root (%) in database.', NEW.pre_root, last_root;
+        END IF;
+
+        RETURN NEW;
+    END;
+$$ language plpgsql;
+
+CREATE TRIGGER validate_pre_root_trigger BEFORE INSERT ON identities FOR EACH ROW EXECUTE PROCEDURE validate_pre_root();
+
+-- Below function took around 10 minutes for 10 million records.
+DO
+$do$
+DECLARE
+    prev_root identities.pre_root%type := NULL;
+    identity identities%rowtype;
+BEGIN
+    FOR identity IN SELECT * FROM identities ORDER BY id ASC
+    LOOP
+        IF identity.pre_root IS NULL THEN UPDATE identities SET pre_root = prev_root WHERE id = identity.id;
+        END IF;
+        prev_root = identity.root;
+    END LOOP;
+END
+$do$;
+
+CREATE UNIQUE INDEX idx_unique_pre_root on identities(pre_root) WHERE pre_root IS NOT NULL;

src/database/mod.rs

@@ -158,6 +158,7 @@ mod test {
     use ethers::types::U256;
     use postgres_docker_utils::DockerContainer;
     use ruint::Uint;
+    use semaphore::poseidon_tree::LazyPoseidonTree;
     use semaphore::Field;
     use testcontainers::clients::Cli;
 
@@ -284,12 +285,15 @@ mod test {
         let (db, _db_container) = setup_db(&docker).await?;
 
         let zero: Hash = U256::zero().into();
+        let initial_root = LazyPoseidonTree::new(4, zero).root();
         let zero_root: Hash = U256::from_dec_str("6789")?.into();
         let root: Hash = U256::from_dec_str("54321")?.into();
         let commitment: Hash = U256::from_dec_str("12345")?.into();
 
-        db.insert_pending_identity(0, &commitment, &root).await?;
-        db.insert_pending_identity(0, &zero, &zero_root).await?;
+        db.insert_pending_identity(0, &commitment, &root, &initial_root)
+            .await?;
+        db.insert_pending_identity(0, &zero, &zero_root, &root)
+            .await?;
 
         let leaf_index = db
             .get_identity_leaf_index(&commitment)
@@ -557,23 +561,6 @@ mod test {
         Ok(())
     }
 
-    #[tokio::test]
-    async fn test_update_insertion_timestamp() -> anyhow::Result<()> {
-        let docker = Cli::default();
-        let (db, _db_container) = setup_db(&docker).await?;
-
-        let insertion_timestamp = Utc::now();
-
-        db.update_latest_insertion_timestamp(insertion_timestamp)
-            .await?;
-
-        let latest_insertion_timestamp = db.get_latest_insertion_timestamp().await?.unwrap();
-
-        assert!(latest_insertion_timestamp.timestamp() - insertion_timestamp.timestamp() <= 1);
-
-        Ok(())
-    }
-
     #[tokio::test]
     async fn test_insert_deletion() -> anyhow::Result<()> {
         let docker = Cli::default();
@@ -637,14 +624,15 @@ mod test {
         let docker = Cli::default();
         let (db, _db_container) = setup_db(&docker).await?;
 
+        let initial_root = LazyPoseidonTree::new(4, Hash::ZERO).root();
         let identities = mock_identities(1);
         let roots = mock_roots(1);
 
         let next_leaf_index = db.get_next_leaf_index().await?;
 
         assert_eq!(next_leaf_index, 0, "Db should contain not leaf indexes");
 
-        db.insert_pending_identity(0, &identities[0], &roots[0])
+        db.insert_pending_identity(0, &identities[0], &roots[0], &initial_root)
             .await?;
 
         let next_leaf_index = db.get_next_leaf_index().await?;
@@ -658,13 +646,16 @@ mod test {
         let docker = Cli::default();
         let (db, _db_container) = setup_db(&docker).await?;
 
+        let initial_root = LazyPoseidonTree::new(4, Hash::ZERO).root();
         let identities = mock_identities(5);
         let roots = mock_roots(5);
 
+        let mut pre_root = &initial_root;
         for i in 0..5 {
-            db.insert_pending_identity(i, &identities[i], &roots[i])
+            db.insert_pending_identity(i, &identities[i], &roots[i], pre_root)
                 .await
                 .context("Inserting identity")?;
+            pre_root = &roots[i];
         }
 
         db.mark_root_as_processed_tx(&roots[2]).await?;
@@ -688,13 +679,16 @@ mod test {
         let docker = Cli::default();
         let (db, _db_container) = setup_db(&docker).await?;
 
+        let initial_root = LazyPoseidonTree::new(4, Hash::ZERO).root();
         let identities = mock_identities(5);
         let roots = mock_roots(5);
 
+        let mut pre_root = &initial_root;
         for i in 0..5 {
-            db.insert_pending_identity(i, &identities[i], &roots[i])
+            db.insert_pending_identity(i, &identities[i], &roots[i], pre_root)
                 .await
                 .context("Inserting identity")?;
+            pre_root = &roots[i];
         }
 
         db.mark_root_as_processed_tx(&roots[2]).await?;
@@ -731,13 +725,16 @@ mod test {
         let docker = Cli::default();
         let (db, _db_container) = setup_db(&docker).await?;
 
+        let initial_root = LazyPoseidonTree::new(4, Hash::ZERO).root();
         let identities = mock_identities(5);
         let roots = mock_roots(5);
 
+        let mut pre_root = &initial_root;
         for i in 0..5 {
-            db.insert_pending_identity(i, &identities[i], &roots[i])
+            db.insert_pending_identity(i, &identities[i], &roots[i], pre_root)
                 .await
                 .context("Inserting identity")?;
+            pre_root = &roots[i];
         }
 
         db.mark_root_as_mined_tx(&roots[2]).await?;
@@ -776,13 +773,16 @@ mod test {
 
         let num_identities = 6;
 
+        let initial_root = LazyPoseidonTree::new(4, Hash::ZERO).root();
         let identities = mock_identities(num_identities);
         let roots = mock_roots(num_identities);
 
+        let mut pre_root = &initial_root;
         for i in 0..num_identities {
-            db.insert_pending_identity(i, &identities[i], &roots[i])
+            db.insert_pending_identity(i, &identities[i], &roots[i], pre_root)
                 .await
                 .context("Inserting identity")?;
+            pre_root = &roots[i];
         }
 
         println!("Marking roots up to 2nd as processed");
@@ -818,13 +818,16 @@ mod test {
         let docker = Cli::default();
         let (db, _db_container) = setup_db(&docker).await?;
 
+        let initial_root = LazyPoseidonTree::new(4, Hash::ZERO).root();
         let identities = mock_identities(5);
         let roots = mock_roots(5);
 
+        let mut pre_root = &initial_root;
         for i in 0..5 {
-            db.insert_pending_identity(i, &identities[i], &roots[i])
+            db.insert_pending_identity(i, &identities[i], &roots[i], pre_root)
                 .await
                 .context("Inserting identity")?;
+            pre_root = &roots[i];
         }
 
         // root[2] is somehow erroneously marked as mined
@@ -862,14 +865,15 @@ mod test {
         let docker = Cli::default();
         let (db, _db_container) = setup_db(&docker).await?;
 
+        let initial_root = LazyPoseidonTree::new(4, Hash::ZERO).root();
         let identities = mock_identities(5);
         let roots = mock_roots(5);
 
         let root = db.get_root_state(&roots[0]).await?;
 
         assert!(root.is_none(), "Root should not exist");
 
-        db.insert_pending_identity(0, &identities[0], &roots[0])
+        db.insert_pending_identity(0, &identities[0], &roots[0], &initial_root)
             .await?;
 
         let root = db
@@ -920,14 +924,17 @@ mod test {
         let docker = Cli::default();
         let (db, _db_container) = setup_db(&docker).await?;
 
+        let initial_root = LazyPoseidonTree::new(4, Hash::ZERO).root();
         let identities = mock_identities(5);
 
         let roots = mock_roots(7);
 
+        let mut pre_root = &initial_root;
         for i in 0..5 {
-            db.insert_pending_identity(i, &identities[i], &roots[i])
+            db.insert_pending_identity(i, &identities[i], &roots[i], pre_root)
                 .await
                 .context("Inserting identity")?;
+            pre_root = &roots[i];
         }
 
         db.mark_root_as_processed_tx(&roots[2]).await?;
@@ -959,20 +966,23 @@ mod test {
         let docker = Cli::default();
         let (db, _db_container) = setup_db(&docker).await?;
 
+        let initial_root = LazyPoseidonTree::new(4, Hash::ZERO).root();
         let identities = mock_identities(5);
 
         let roots = mock_roots(5);
         let zero_roots = mock_zero_roots(5);
 
+        let mut pre_root = &initial_root;
         for i in 0..5 {
-            db.insert_pending_identity(i, &identities[i], &roots[i])
+            db.insert_pending_identity(i, &identities[i], &roots[i], pre_root)
                 .await
                 .context("Inserting identity")?;
+            pre_root = &roots[i];
         }
 
-        db.insert_pending_identity(0, &Hash::ZERO, &zero_roots[0])
+        db.insert_pending_identity(0, &Hash::ZERO, &zero_roots[0], &roots[4])
             .await?;
-        db.insert_pending_identity(3, &Hash::ZERO, &zero_roots[3])
+        db.insert_pending_identity(3, &Hash::ZERO, &zero_roots[3], &zero_roots[0])
             .await?;
 
         let pending_tree_updates = db
@@ -1014,10 +1024,11 @@ mod test {
         let docker = Cli::default();
         let (db, _db_container) = setup_db(&docker).await?;
 
+        let initial_root = LazyPoseidonTree::new(4, Hash::ZERO).root();
         let identities = mock_identities(5);
         let roots = mock_roots(5);
 
-        db.insert_pending_identity(0, &identities[0], &roots[0])
+        db.insert_pending_identity(0, &identities[0], &roots[0], &initial_root)
             .await
             .context("Inserting identity 1")?;
 
@@ -1034,9 +1045,9 @@ mod test {
 
         // Inserting a new pending root sets invalidation time for the
         // previous root
-        db.insert_pending_identity(1, &identities[1], &roots[1])
+        db.insert_pending_identity(1, &identities[1], &roots[1], &roots[0])
             .await?;
-        db.insert_pending_identity(2, &identities[2], &roots[2])
+        db.insert_pending_identity(2, &identities[2], &roots[2], &roots[1])
             .await?;
 
         let root_1_inserted_at = Utc::now();
@@ -1056,7 +1067,7 @@ mod test {
         assert_same_time!(root_item_1.pending_valid_as_of, root_1_inserted_at);
 
         // Test mined roots
-        db.insert_pending_identity(3, &identities[3], &roots[3])
+        db.insert_pending_identity(3, &identities[3], &roots[3], &roots[2])
             .await?;
 
         db.mark_root_as_processed_tx(&roots[0])
@@ -1091,6 +1102,7 @@ mod test {
         let docker = Cli::default();
         let (db, _db_container) = setup_db(&docker).await?;
 
+        let initial_root = LazyPoseidonTree::new(4, Hash::ZERO).root();
         let identities = mock_identities(2);
         let roots = mock_roots(1);
 
@@ -1106,7 +1118,7 @@ mod test {
         assert!(db.identity_exists(identities[0]).await?);
 
         // When there's only processed identity
-        db.insert_pending_identity(0, &identities[1], &roots[0])
+        db.insert_pending_identity(0, &identities[1], &roots[0], &initial_root)
             .await
             .context("Inserting identity")?;
 
@@ -1148,7 +1160,35 @@ mod test {
     }
 
     #[tokio::test]
-    async fn test_latest_deletion_root() -> anyhow::Result<()> {
+    async fn test_latest_insertion() -> anyhow::Result<()> {
+        let docker = Cli::default();
+        let (db, _db_container) = setup_db(&docker).await?;
+
+        // Update with initial timestamp
+        let initial_timestamp = chrono::Utc::now();
+        db.update_latest_insertion(initial_timestamp)
+            .await
+            .context("Inserting initial root")?;
+
+        // Assert values
+        let initial_entry = db.get_latest_insertion().await?;
+        assert!(initial_entry.timestamp.timestamp() - initial_timestamp.timestamp() <= 1);
+
+        // Update with a new timestamp
+        let new_timestamp = chrono::Utc::now();
+        db.update_latest_insertion(new_timestamp)
+            .await
+            .context("Updating with new root")?;
+
+        // Assert values
+        let new_entry = db.get_latest_insertion().await?;
+        assert!((new_entry.timestamp.timestamp() - new_timestamp.timestamp()) <= 1);
+
+        Ok(())
+    }
+
+    #[tokio::test]
+    async fn test_latest_deletion() -> anyhow::Result<()> {
         let docker = Cli::default();
         let (db, _db_container) = setup_db(&docker).await?;
 
@@ -1179,25 +1219,20 @@ mod test {
     async fn can_not_insert_same_root_multiple_times() -> anyhow::Result<()> {
         let docker = Cli::default();
         let (db, _db_container) = setup_db(&docker).await?;
+
+        let initial_root = LazyPoseidonTree::new(4, Hash::ZERO).root();
         let identities = mock_identities(2);
         let roots = mock_roots(2);
 
-        db.insert_pending_identity(0, &identities[0], &roots[0])
+        db.insert_pending_identity(0, &identities[0], &roots[0], &initial_root)
             .await?;
 
         let res = db
-            .insert_pending_identity(1, &identities[1], &roots[0])
+            .insert_pending_identity(1, &identities[1], &roots[0], &roots[0])
             .await;
 
         assert!(res.is_err(), "Inserting duplicate root should fail");
 
-        let root_state = db
-            .get_root_state(&roots[0])
-            .await?
-            .context("Missing root")?;
-
-        assert_eq!(root_state.status, ProcessedStatus::Pending);
-
         Ok(())
     }