Skip to content

Commit

Permalink
Add pre_root for identities.
Browse files Browse the repository at this point in the history
  • Loading branch information
piohei committed Aug 6, 2024
1 parent 2c16474 commit 9a283f9
Show file tree
Hide file tree
Showing 8 changed files with 220 additions and 98 deletions.
7 changes: 7 additions & 0 deletions schemas/database/015_add_constraints_for_identities.down.sql
Original file line number Diff line number Diff line change
@@ -0,0 +1,7 @@
DROP UNIQUE INDEX idx_unique_insertion_leaf;
DROP UNIQUE INDEX idx_unique_deletion_leaf;

DROP TRIGGER validate_pre_root_trigger;
DROP FUNCTION validate_pre_root();

ALTER TABLE identities DROP COLUMN pre_root;
52 changes: 52 additions & 0 deletions schemas/database/015_add_constraints_for_identities.up.sql
Original file line number Diff line number Diff line change
@@ -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 tre 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;
125 changes: 80 additions & 45 deletions src/database/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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();
Expand Down Expand Up @@ -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?;
Expand All @@ -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?;
Expand All @@ -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?;
Expand Down Expand Up @@ -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?;
Expand Down Expand Up @@ -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");
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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?;
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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")?;

Expand All @@ -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();
Expand All @@ -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])
Expand Down Expand Up @@ -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);

Expand All @@ -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")?;

Expand Down Expand Up @@ -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?;

Expand Down Expand Up @@ -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(())
}

Expand Down
Loading

0 comments on commit 9a283f9

Please sign in to comment.