Move unshare_project to db module

Nathan Sobo created

Change summary

crates/collab/src/db.rs        | 83 +++++++++++++++++++++--------------
crates/collab/src/rpc.rs       | 13 +++-
crates/collab/src/rpc/store.rs | 66 ----------------------------
3 files changed, 59 insertions(+), 103 deletions(-)

Detailed changes

crates/collab/src/db.rs 🔗

@@ -1330,6 +1330,27 @@ where
         Ok(room)
     }
 
+    async fn get_guest_connection_ids(
+        &self,
+        project_id: ProjectId,
+        tx: &mut sqlx::Transaction<'_, D>,
+    ) -> Result<Vec<ConnectionId>> {
+        let mut guest_connection_ids = Vec::new();
+        let mut db_guest_connection_ids = sqlx::query_scalar::<_, i32>(
+            "
+            SELECT connection_id
+            FROM project_collaborators
+            WHERE project_id = $1 AND is_host = FALSE
+            ",
+        )
+        .bind(project_id)
+        .fetch(tx);
+        while let Some(connection_id) = db_guest_connection_ids.next().await {
+            guest_connection_ids.push(ConnectionId(connection_id? as u32));
+        }
+        Ok(guest_connection_ids)
+    }
+
     async fn get_room(
         &self,
         room_id: RoomId,
@@ -1539,6 +1560,31 @@ where
         .await
     }
 
+    pub async fn unshare_project(
+        &self,
+        project_id: ProjectId,
+        connection_id: ConnectionId,
+    ) -> Result<(proto::Room, Vec<ConnectionId>)> {
+        self.transact(|mut tx| async move {
+            let guest_connection_ids = self.get_guest_connection_ids(project_id, &mut tx).await?;
+            let room_id: RoomId = sqlx::query_scalar(
+                "
+                DELETE FROM projects
+                WHERE id = $1 AND host_connection_id = $2
+                RETURNING room_id
+                ",
+            )
+            .bind(project_id)
+            .bind(connection_id.0 as i32)
+            .fetch_one(&mut tx)
+            .await?;
+            let room = self.commit_room_transaction(room_id, tx).await?;
+
+            Ok((room, guest_connection_ids))
+        })
+        .await
+    }
+
     pub async fn update_project(
         &self,
         project_id: ProjectId,
@@ -1608,23 +1654,9 @@ where
             }
             query.execute(&mut tx).await?;
 
-            let mut guest_connection_ids = Vec::new();
-            {
-                let mut db_guest_connection_ids = sqlx::query_scalar::<_, i32>(
-                    "
-                    SELECT connection_id
-                    FROM project_collaborators
-                    WHERE project_id = $1 AND is_host = FALSE
-                    ",
-                )
-                .bind(project_id)
-                .fetch(&mut tx);
-                while let Some(connection_id) = db_guest_connection_ids.next().await {
-                    guest_connection_ids.push(ConnectionId(connection_id? as u32));
-                }
-            }
-
+            let guest_connection_ids = self.get_guest_connection_ids(project_id, &mut tx).await?;
             let room = self.commit_room_transaction(room_id, tx).await?;
+
             Ok((room, guest_connection_ids))
         })
         .await
@@ -2108,7 +2140,7 @@ where
             .execute(&mut tx)
             .await?;
 
-            if result.rows_affected() != 1 {
+            if result.rows_affected() == 0 {
                 Err(anyhow!("not a collaborator on this project"))?;
             }
 
@@ -2207,23 +2239,6 @@ where
         .await
     }
 
-    pub async fn unshare_project(&self, project_id: ProjectId) -> Result<()> {
-        todo!()
-        // test_support!(self, {
-        //     sqlx::query(
-        //         "
-        //         UPDATE projects
-        //         SET unregistered = TRUE
-        //         WHERE id = $1
-        //         ",
-        //     )
-        //     .bind(project_id)
-        //     .execute(&self.pool)
-        //     .await?;
-        //     Ok(())
-        // })
-    }
-
     // contacts
 
     pub async fn get_contacts(&self, user_id: UserId) -> Result<Vec<Contact>> {

crates/collab/src/rpc.rs 🔗

@@ -877,14 +877,19 @@ impl Server {
         message: Message<proto::UnshareProject>,
     ) -> Result<()> {
         let project_id = ProjectId::from_proto(message.payload.project_id);
-        let mut store = self.store().await;
-        let (room, project) = store.unshare_project(project_id, message.sender_connection_id)?;
+
+        let (room, guest_connection_ids) = self
+            .app_state
+            .db
+            .unshare_project(project_id, message.sender_connection_id)
+            .await?;
+
         broadcast(
             message.sender_connection_id,
-            project.guest_connection_ids(),
+            guest_connection_ids,
             |conn_id| self.peer.send(conn_id, message.payload.clone()),
         );
-        self.room_updated(room);
+        self.room_updated(&room);
 
         Ok(())
     }

crates/collab/src/rpc/store.rs 🔗

@@ -1,6 +1,6 @@
 use crate::db::{self, ProjectId, UserId};
 use anyhow::{anyhow, Result};
-use collections::{btree_map, BTreeMap, BTreeSet, HashMap, HashSet};
+use collections::{BTreeMap, BTreeSet, HashMap, HashSet};
 use rpc::{proto, ConnectionId};
 use serde::Serialize;
 use std::path::PathBuf;
@@ -72,14 +72,6 @@ pub struct Worktree {
 
 pub type ReplicaId = u16;
 
-pub struct LeftProject {
-    pub id: ProjectId,
-    pub host_user_id: UserId,
-    pub host_connection_id: ConnectionId,
-    pub connection_ids: Vec<ConnectionId>,
-    pub remove_collaborator: bool,
-}
-
 #[derive(Copy, Clone)]
 pub struct Metrics {
     pub connections: usize,
@@ -209,48 +201,6 @@ impl Store {
         &self.rooms
     }
 
-    pub fn unshare_project(
-        &mut self,
-        project_id: ProjectId,
-        connection_id: ConnectionId,
-    ) -> Result<(&proto::Room, Project)> {
-        match self.projects.entry(project_id) {
-            btree_map::Entry::Occupied(e) => {
-                if e.get().host_connection_id == connection_id {
-                    let project = e.remove();
-
-                    if let Some(host_connection) = self.connections.get_mut(&connection_id) {
-                        host_connection.projects.remove(&project_id);
-                    }
-
-                    for guest_connection in project.guests.keys() {
-                        if let Some(connection) = self.connections.get_mut(guest_connection) {
-                            connection.projects.remove(&project_id);
-                        }
-                    }
-
-                    let room = self
-                        .rooms
-                        .get_mut(&project.room_id)
-                        .ok_or_else(|| anyhow!("no such room"))?;
-                    let participant = room
-                        .participants
-                        .iter_mut()
-                        .find(|participant| participant.peer_id == connection_id.0)
-                        .ok_or_else(|| anyhow!("no such room"))?;
-                    participant
-                        .projects
-                        .retain(|project| project.id != project_id.to_proto());
-
-                    Ok((room, project))
-                } else {
-                    Err(anyhow!("no such project"))?
-                }
-            }
-            btree_map::Entry::Vacant(_) => Err(anyhow!("no such project"))?,
-        }
-    }
-
     #[cfg(test)]
     pub fn check_invariants(&self) {
         for (connection_id, connection) in &self.connections {
@@ -373,17 +323,3 @@ impl Store {
         }
     }
 }
-
-impl Project {
-    pub fn guest_connection_ids(&self) -> Vec<ConnectionId> {
-        self.guests.keys().copied().collect()
-    }
-
-    pub fn connection_ids(&self) -> Vec<ConnectionId> {
-        self.guests
-            .keys()
-            .copied()
-            .chain(Some(self.host_connection_id))
-            .collect()
-    }
-}