Implement invite codes using sea-orm

Antonio Scandurra created

Change summary

crates/collab/src/db2.rs        | 220 +++++++++++++++++++
crates/collab/src/db2/signup.rs |  33 ++
crates/collab/src/db2/tests.rs  | 386 +++++++++++++++++-----------------
3 files changed, 446 insertions(+), 193 deletions(-)

Detailed changes

crates/collab/src/db2.rs 🔗

@@ -4,6 +4,7 @@ mod project;
 mod project_collaborator;
 mod room;
 mod room_participant;
+mod signup;
 #[cfg(test)]
 mod tests;
 mod user;
@@ -14,6 +15,7 @@ use anyhow::anyhow;
 use collections::HashMap;
 use dashmap::DashMap;
 use futures::StreamExt;
+use hyper::StatusCode;
 use rpc::{proto, ConnectionId};
 use sea_orm::{
     entity::prelude::*, ConnectOptions, DatabaseConnection, DatabaseTransaction, DbErr,
@@ -34,6 +36,7 @@ use std::{future::Future, marker::PhantomData, rc::Rc, sync::Arc};
 use tokio::sync::{Mutex, OwnedMutexGuard};
 
 pub use contact::Contact;
+pub use signup::Invite;
 pub use user::Model as User;
 
 pub struct Database {
@@ -523,6 +526,222 @@ impl Database {
         .await
     }
 
+    // invite codes
+
+    pub async fn create_invite_from_code(
+        &self,
+        code: &str,
+        email_address: &str,
+        device_id: Option<&str>,
+    ) -> Result<Invite> {
+        self.transact(|tx| async move {
+            let existing_user = user::Entity::find()
+                .filter(user::Column::EmailAddress.eq(email_address))
+                .one(&tx)
+                .await?;
+
+            if existing_user.is_some() {
+                Err(anyhow!("email address is already in use"))?;
+            }
+
+            let inviter = match user::Entity::find()
+                .filter(user::Column::InviteCode.eq(code))
+                .one(&tx)
+                .await?
+            {
+                Some(inviter) => inviter,
+                None => {
+                    return Err(Error::Http(
+                        StatusCode::NOT_FOUND,
+                        "invite code not found".to_string(),
+                    ))?
+                }
+            };
+
+            if inviter.invite_count == 0 {
+                Err(Error::Http(
+                    StatusCode::UNAUTHORIZED,
+                    "no invites remaining".to_string(),
+                ))?;
+            }
+
+            let signup = signup::Entity::insert(signup::ActiveModel {
+                email_address: ActiveValue::set(email_address.into()),
+                email_confirmation_code: ActiveValue::set(random_email_confirmation_code()),
+                email_confirmation_sent: ActiveValue::set(false),
+                inviting_user_id: ActiveValue::set(Some(inviter.id)),
+                platform_linux: ActiveValue::set(false),
+                platform_mac: ActiveValue::set(false),
+                platform_windows: ActiveValue::set(false),
+                platform_unknown: ActiveValue::set(true),
+                device_id: ActiveValue::set(device_id.map(|device_id| device_id.into())),
+                ..Default::default()
+            })
+            .on_conflict(
+                OnConflict::column(signup::Column::EmailAddress)
+                    .update_column(signup::Column::InvitingUserId)
+                    .to_owned(),
+            )
+            .exec_with_returning(&tx)
+            .await?;
+            tx.commit().await?;
+
+            Ok(Invite {
+                email_address: signup.email_address,
+                email_confirmation_code: signup.email_confirmation_code,
+            })
+        })
+        .await
+    }
+
+    pub async fn create_user_from_invite(
+        &self,
+        invite: &Invite,
+        user: NewUserParams,
+    ) -> Result<Option<NewUserResult>> {
+        self.transact(|tx| async {
+            let tx = tx;
+            let signup = signup::Entity::find()
+                .filter(
+                    signup::Column::EmailAddress
+                        .eq(invite.email_address.as_str())
+                        .and(
+                            signup::Column::EmailConfirmationCode
+                                .eq(invite.email_confirmation_code.as_str()),
+                        ),
+                )
+                .one(&tx)
+                .await?
+                .ok_or_else(|| Error::Http(StatusCode::NOT_FOUND, "no such invite".to_string()))?;
+
+            if signup.user_id.is_some() {
+                return Ok(None);
+            }
+
+            let user = user::Entity::insert(user::ActiveModel {
+                email_address: ActiveValue::set(Some(invite.email_address.clone())),
+                github_login: ActiveValue::set(user.github_login.clone()),
+                github_user_id: ActiveValue::set(Some(user.github_user_id)),
+                admin: ActiveValue::set(false),
+                invite_count: ActiveValue::set(user.invite_count),
+                invite_code: ActiveValue::set(Some(random_invite_code())),
+                metrics_id: ActiveValue::set(Uuid::new_v4()),
+                ..Default::default()
+            })
+            .on_conflict(
+                OnConflict::column(user::Column::GithubLogin)
+                    .update_columns([
+                        user::Column::EmailAddress,
+                        user::Column::GithubUserId,
+                        user::Column::Admin,
+                    ])
+                    .to_owned(),
+            )
+            .exec_with_returning(&tx)
+            .await?;
+
+            let mut signup = signup.into_active_model();
+            signup.user_id = ActiveValue::set(Some(user.id));
+            let signup = signup.update(&tx).await?;
+
+            if let Some(inviting_user_id) = signup.inviting_user_id {
+                let result = user::Entity::update_many()
+                    .filter(
+                        user::Column::Id
+                            .eq(inviting_user_id)
+                            .and(user::Column::InviteCount.gt(0)),
+                    )
+                    .col_expr(
+                        user::Column::InviteCount,
+                        Expr::col(user::Column::InviteCount).sub(1),
+                    )
+                    .exec(&tx)
+                    .await?;
+
+                if result.rows_affected == 0 {
+                    Err(Error::Http(
+                        StatusCode::UNAUTHORIZED,
+                        "no invites remaining".to_string(),
+                    ))?;
+                }
+
+                contact::Entity::insert(contact::ActiveModel {
+                    user_id_a: ActiveValue::set(inviting_user_id),
+                    user_id_b: ActiveValue::set(user.id),
+                    a_to_b: ActiveValue::set(true),
+                    should_notify: ActiveValue::set(true),
+                    accepted: ActiveValue::set(true),
+                    ..Default::default()
+                })
+                .on_conflict(OnConflict::new().do_nothing().to_owned())
+                .exec_without_returning(&tx)
+                .await?;
+            }
+
+            tx.commit().await?;
+            Ok(Some(NewUserResult {
+                user_id: user.id,
+                metrics_id: user.metrics_id.to_string(),
+                inviting_user_id: signup.inviting_user_id,
+                signup_device_id: signup.device_id,
+            }))
+        })
+        .await
+    }
+
+    pub async fn set_invite_count_for_user(&self, id: UserId, count: u32) -> Result<()> {
+        self.transact(|tx| async move {
+            if count > 0 {
+                user::Entity::update_many()
+                    .filter(
+                        user::Column::Id
+                            .eq(id)
+                            .and(user::Column::InviteCode.is_null()),
+                    )
+                    .col_expr(user::Column::InviteCode, random_invite_code().into())
+                    .exec(&tx)
+                    .await?;
+            }
+
+            user::Entity::update_many()
+                .filter(user::Column::Id.eq(id))
+                .col_expr(user::Column::InviteCount, count.into())
+                .exec(&tx)
+                .await?;
+            tx.commit().await?;
+            Ok(())
+        })
+        .await
+    }
+
+    pub async fn get_invite_code_for_user(&self, id: UserId) -> Result<Option<(String, u32)>> {
+        self.transact(|tx| async move {
+            match user::Entity::find_by_id(id).one(&tx).await? {
+                Some(user) if user.invite_code.is_some() => {
+                    Ok(Some((user.invite_code.unwrap(), user.invite_count as u32)))
+                }
+                _ => Ok(None),
+            }
+        })
+        .await
+    }
+
+    pub async fn get_user_for_invite_code(&self, code: &str) -> Result<User> {
+        self.transact(|tx| async move {
+            user::Entity::find()
+                .filter(user::Column::InviteCode.eq(code))
+                .one(&tx)
+                .await?
+                .ok_or_else(|| {
+                    Error::Http(
+                        StatusCode::NOT_FOUND,
+                        "that invite code does not exist".to_string(),
+                    )
+                })
+        })
+        .await
+    }
+
     // projects
 
     pub async fn share_project(
@@ -966,6 +1185,7 @@ id_type!(RoomId);
 id_type!(RoomParticipantId);
 id_type!(ProjectId);
 id_type!(ProjectCollaboratorId);
+id_type!(SignupId);
 id_type!(WorktreeId);
 
 #[cfg(test)]

crates/collab/src/db2/signup.rs 🔗

@@ -0,0 +1,33 @@
+use super::{SignupId, UserId};
+use sea_orm::entity::prelude::*;
+
+#[derive(Clone, Debug, PartialEq, Eq, DeriveEntityModel)]
+#[sea_orm(table_name = "signups")]
+pub struct Model {
+    #[sea_orm(primary_key)]
+    pub id: SignupId,
+    pub email_address: String,
+    pub email_confirmation_code: String,
+    pub email_confirmation_sent: bool,
+    pub created_at: DateTime,
+    pub device_id: Option<String>,
+    pub user_id: Option<UserId>,
+    pub inviting_user_id: Option<UserId>,
+    pub platform_mac: bool,
+    pub platform_linux: bool,
+    pub platform_windows: bool,
+    pub platform_unknown: bool,
+    pub editor_features: Option<String>,
+    pub programming_languages: Option<String>,
+}
+
+#[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)]
+pub enum Relation {}
+
+impl ActiveModelBehavior for ActiveModel {}
+
+#[derive(Debug)]
+pub struct Invite {
+    pub email_address: String,
+    pub email_confirmation_code: String,
+}

crates/collab/src/db2/tests.rs 🔗

@@ -457,210 +457,210 @@ async fn test_fuzzy_search_users() {
     }
 }
 
-// #[gpui::test]
-// async fn test_invite_codes() {
-//     let test_db = PostgresTestDb::new(build_background_executor());
-//     let db = test_db.db();
+#[gpui::test]
+async fn test_invite_codes() {
+    let test_db = TestDb::postgres(build_background_executor());
+    let db = test_db.db();
 
-//     let NewUserResult { user_id: user1, .. } = db
-//         .create_user(
-//             "user1@example.com",
-//             false,
-//             NewUserParams {
-//                 github_login: "user1".into(),
-//                 github_user_id: 0,
-//                 invite_count: 0,
-//             },
-//         )
-//         .await
-//         .unwrap();
+    let NewUserResult { user_id: user1, .. } = db
+        .create_user(
+            "user1@example.com",
+            false,
+            NewUserParams {
+                github_login: "user1".into(),
+                github_user_id: 0,
+                invite_count: 0,
+            },
+        )
+        .await
+        .unwrap();
 
-//     // Initially, user 1 has no invite code
-//     assert_eq!(db.get_invite_code_for_user(user1).await.unwrap(), None);
+    // Initially, user 1 has no invite code
+    assert_eq!(db.get_invite_code_for_user(user1).await.unwrap(), None);
 
-//     // Setting invite count to 0 when no code is assigned does not assign a new code
-//     db.set_invite_count_for_user(user1, 0).await.unwrap();
-//     assert!(db.get_invite_code_for_user(user1).await.unwrap().is_none());
+    // Setting invite count to 0 when no code is assigned does not assign a new code
+    db.set_invite_count_for_user(user1, 0).await.unwrap();
+    assert!(db.get_invite_code_for_user(user1).await.unwrap().is_none());
 
-//     // User 1 creates an invite code that can be used twice.
-//     db.set_invite_count_for_user(user1, 2).await.unwrap();
-//     let (invite_code, invite_count) = db.get_invite_code_for_user(user1).await.unwrap().unwrap();
-//     assert_eq!(invite_count, 2);
+    // User 1 creates an invite code that can be used twice.
+    db.set_invite_count_for_user(user1, 2).await.unwrap();
+    let (invite_code, invite_count) = db.get_invite_code_for_user(user1).await.unwrap().unwrap();
+    assert_eq!(invite_count, 2);
 
-//     // User 2 redeems the invite code and becomes a contact of user 1.
-//     let user2_invite = db
-//         .create_invite_from_code(&invite_code, "user2@example.com", Some("user-2-device-id"))
-//         .await
-//         .unwrap();
-//     let NewUserResult {
-//         user_id: user2,
-//         inviting_user_id,
-//         signup_device_id,
-//         metrics_id,
-//     } = db
-//         .create_user_from_invite(
-//             &user2_invite,
-//             NewUserParams {
-//                 github_login: "user2".into(),
-//                 github_user_id: 2,
-//                 invite_count: 7,
-//             },
-//         )
-//         .await
-//         .unwrap()
-//         .unwrap();
-//     let (_, invite_count) = db.get_invite_code_for_user(user1).await.unwrap().unwrap();
-//     assert_eq!(invite_count, 1);
-//     assert_eq!(inviting_user_id, Some(user1));
-//     assert_eq!(signup_device_id.unwrap(), "user-2-device-id");
-//     assert_eq!(db.get_user_metrics_id(user2).await.unwrap(), metrics_id);
-//     assert_eq!(
-//         db.get_contacts(user1).await.unwrap(),
-//         [Contact::Accepted {
-//             user_id: user2,
-//             should_notify: true,
-//             busy: false,
-//         }]
-//     );
-//     assert_eq!(
-//         db.get_contacts(user2).await.unwrap(),
-//         [Contact::Accepted {
-//             user_id: user1,
-//             should_notify: false,
-//             busy: false,
-//         }]
-//     );
-//     assert_eq!(
-//         db.get_invite_code_for_user(user2).await.unwrap().unwrap().1,
-//         7
-//     );
+    // User 2 redeems the invite code and becomes a contact of user 1.
+    let user2_invite = db
+        .create_invite_from_code(&invite_code, "user2@example.com", Some("user-2-device-id"))
+        .await
+        .unwrap();
+    let NewUserResult {
+        user_id: user2,
+        inviting_user_id,
+        signup_device_id,
+        metrics_id,
+    } = db
+        .create_user_from_invite(
+            &user2_invite,
+            NewUserParams {
+                github_login: "user2".into(),
+                github_user_id: 2,
+                invite_count: 7,
+            },
+        )
+        .await
+        .unwrap()
+        .unwrap();
+    let (_, invite_count) = db.get_invite_code_for_user(user1).await.unwrap().unwrap();
+    assert_eq!(invite_count, 1);
+    assert_eq!(inviting_user_id, Some(user1));
+    assert_eq!(signup_device_id.unwrap(), "user-2-device-id");
+    assert_eq!(db.get_user_metrics_id(user2).await.unwrap(), metrics_id);
+    assert_eq!(
+        db.get_contacts(user1).await.unwrap(),
+        [Contact::Accepted {
+            user_id: user2,
+            should_notify: true,
+            busy: false,
+        }]
+    );
+    assert_eq!(
+        db.get_contacts(user2).await.unwrap(),
+        [Contact::Accepted {
+            user_id: user1,
+            should_notify: false,
+            busy: false,
+        }]
+    );
+    assert_eq!(
+        db.get_invite_code_for_user(user2).await.unwrap().unwrap().1,
+        7
+    );
 
-//     // User 3 redeems the invite code and becomes a contact of user 1.
-//     let user3_invite = db
-//         .create_invite_from_code(&invite_code, "user3@example.com", None)
-//         .await
-//         .unwrap();
-//     let NewUserResult {
-//         user_id: user3,
-//         inviting_user_id,
-//         signup_device_id,
-//         ..
-//     } = db
-//         .create_user_from_invite(
-//             &user3_invite,
-//             NewUserParams {
-//                 github_login: "user-3".into(),
-//                 github_user_id: 3,
-//                 invite_count: 3,
-//             },
-//         )
-//         .await
-//         .unwrap()
-//         .unwrap();
-//     let (_, invite_count) = db.get_invite_code_for_user(user1).await.unwrap().unwrap();
-//     assert_eq!(invite_count, 0);
-//     assert_eq!(inviting_user_id, Some(user1));
-//     assert!(signup_device_id.is_none());
-//     assert_eq!(
-//         db.get_contacts(user1).await.unwrap(),
-//         [
-//             Contact::Accepted {
-//                 user_id: user2,
-//                 should_notify: true,
-//                 busy: false,
-//             },
-//             Contact::Accepted {
-//                 user_id: user3,
-//                 should_notify: true,
-//                 busy: false,
-//             }
-//         ]
-//     );
-//     assert_eq!(
-//         db.get_contacts(user3).await.unwrap(),
-//         [Contact::Accepted {
-//             user_id: user1,
-//             should_notify: false,
-//             busy: false,
-//         }]
-//     );
-//     assert_eq!(
-//         db.get_invite_code_for_user(user3).await.unwrap().unwrap().1,
-//         3
-//     );
+    // User 3 redeems the invite code and becomes a contact of user 1.
+    let user3_invite = db
+        .create_invite_from_code(&invite_code, "user3@example.com", None)
+        .await
+        .unwrap();
+    let NewUserResult {
+        user_id: user3,
+        inviting_user_id,
+        signup_device_id,
+        ..
+    } = db
+        .create_user_from_invite(
+            &user3_invite,
+            NewUserParams {
+                github_login: "user-3".into(),
+                github_user_id: 3,
+                invite_count: 3,
+            },
+        )
+        .await
+        .unwrap()
+        .unwrap();
+    let (_, invite_count) = db.get_invite_code_for_user(user1).await.unwrap().unwrap();
+    assert_eq!(invite_count, 0);
+    assert_eq!(inviting_user_id, Some(user1));
+    assert!(signup_device_id.is_none());
+    assert_eq!(
+        db.get_contacts(user1).await.unwrap(),
+        [
+            Contact::Accepted {
+                user_id: user2,
+                should_notify: true,
+                busy: false,
+            },
+            Contact::Accepted {
+                user_id: user3,
+                should_notify: true,
+                busy: false,
+            }
+        ]
+    );
+    assert_eq!(
+        db.get_contacts(user3).await.unwrap(),
+        [Contact::Accepted {
+            user_id: user1,
+            should_notify: false,
+            busy: false,
+        }]
+    );
+    assert_eq!(
+        db.get_invite_code_for_user(user3).await.unwrap().unwrap().1,
+        3
+    );
 
-//     // Trying to reedem the code for the third time results in an error.
-//     db.create_invite_from_code(&invite_code, "user4@example.com", Some("user-4-device-id"))
-//         .await
-//         .unwrap_err();
+    // Trying to reedem the code for the third time results in an error.
+    db.create_invite_from_code(&invite_code, "user4@example.com", Some("user-4-device-id"))
+        .await
+        .unwrap_err();
 
-//     // Invite count can be updated after the code has been created.
-//     db.set_invite_count_for_user(user1, 2).await.unwrap();
-//     let (latest_code, invite_count) = db.get_invite_code_for_user(user1).await.unwrap().unwrap();
-//     assert_eq!(latest_code, invite_code); // Invite code doesn't change when we increment above 0
-//     assert_eq!(invite_count, 2);
+    // Invite count can be updated after the code has been created.
+    db.set_invite_count_for_user(user1, 2).await.unwrap();
+    let (latest_code, invite_count) = db.get_invite_code_for_user(user1).await.unwrap().unwrap();
+    assert_eq!(latest_code, invite_code); // Invite code doesn't change when we increment above 0
+    assert_eq!(invite_count, 2);
 
-//     // User 4 can now redeem the invite code and becomes a contact of user 1.
-//     let user4_invite = db
-//         .create_invite_from_code(&invite_code, "user4@example.com", Some("user-4-device-id"))
-//         .await
-//         .unwrap();
-//     let user4 = db
-//         .create_user_from_invite(
-//             &user4_invite,
-//             NewUserParams {
-//                 github_login: "user-4".into(),
-//                 github_user_id: 4,
-//                 invite_count: 5,
-//             },
-//         )
-//         .await
-//         .unwrap()
-//         .unwrap()
-//         .user_id;
+    // User 4 can now redeem the invite code and becomes a contact of user 1.
+    let user4_invite = db
+        .create_invite_from_code(&invite_code, "user4@example.com", Some("user-4-device-id"))
+        .await
+        .unwrap();
+    let user4 = db
+        .create_user_from_invite(
+            &user4_invite,
+            NewUserParams {
+                github_login: "user-4".into(),
+                github_user_id: 4,
+                invite_count: 5,
+            },
+        )
+        .await
+        .unwrap()
+        .unwrap()
+        .user_id;
 
-//     let (_, invite_count) = db.get_invite_code_for_user(user1).await.unwrap().unwrap();
-//     assert_eq!(invite_count, 1);
-//     assert_eq!(
-//         db.get_contacts(user1).await.unwrap(),
-//         [
-//             Contact::Accepted {
-//                 user_id: user2,
-//                 should_notify: true,
-//                 busy: false,
-//             },
-//             Contact::Accepted {
-//                 user_id: user3,
-//                 should_notify: true,
-//                 busy: false,
-//             },
-//             Contact::Accepted {
-//                 user_id: user4,
-//                 should_notify: true,
-//                 busy: false,
-//             }
-//         ]
-//     );
-//     assert_eq!(
-//         db.get_contacts(user4).await.unwrap(),
-//         [Contact::Accepted {
-//             user_id: user1,
-//             should_notify: false,
-//             busy: false,
-//         }]
-//     );
-//     assert_eq!(
-//         db.get_invite_code_for_user(user4).await.unwrap().unwrap().1,
-//         5
-//     );
+    let (_, invite_count) = db.get_invite_code_for_user(user1).await.unwrap().unwrap();
+    assert_eq!(invite_count, 1);
+    assert_eq!(
+        db.get_contacts(user1).await.unwrap(),
+        [
+            Contact::Accepted {
+                user_id: user2,
+                should_notify: true,
+                busy: false,
+            },
+            Contact::Accepted {
+                user_id: user3,
+                should_notify: true,
+                busy: false,
+            },
+            Contact::Accepted {
+                user_id: user4,
+                should_notify: true,
+                busy: false,
+            }
+        ]
+    );
+    assert_eq!(
+        db.get_contacts(user4).await.unwrap(),
+        [Contact::Accepted {
+            user_id: user1,
+            should_notify: false,
+            busy: false,
+        }]
+    );
+    assert_eq!(
+        db.get_invite_code_for_user(user4).await.unwrap().unwrap().1,
+        5
+    );
 
-//     // An existing user cannot redeem invite codes.
-//     db.create_invite_from_code(&invite_code, "user2@example.com", Some("user-2-device-id"))
-//         .await
-//         .unwrap_err();
-//     let (_, invite_count) = db.get_invite_code_for_user(user1).await.unwrap().unwrap();
-//     assert_eq!(invite_count, 1);
-// }
+    // An existing user cannot redeem invite codes.
+    db.create_invite_from_code(&invite_code, "user2@example.com", Some("user-2-device-id"))
+        .await
+        .unwrap_err();
+    let (_, invite_count) = db.get_invite_code_for_user(user1).await.unwrap().unwrap();
+    assert_eq!(invite_count, 1);
+}
 
 // #[gpui::test]
 // async fn test_signups() {