1use crate::{Error, Result};
2use anyhow::{anyhow, Context};
3use async_trait::async_trait;
4use axum::http::StatusCode;
5use collections::HashMap;
6use futures::StreamExt;
7use serde::{Deserialize, Serialize};
8pub use sqlx::postgres::PgPoolOptions as DbOptions;
9use sqlx::{types::Uuid, FromRow, QueryBuilder};
10use std::{cmp, ops::Range, time::Duration};
11use time::{OffsetDateTime, PrimitiveDateTime};
12
13#[async_trait]
14pub trait Db: Send + Sync {
15 async fn create_user(
16 &self,
17 email_address: &str,
18 admin: bool,
19 params: NewUserParams,
20 ) -> Result<NewUserResult>;
21 async fn get_all_users(&self, page: u32, limit: u32) -> Result<Vec<User>>;
22 async fn fuzzy_search_users(&self, query: &str, limit: u32) -> Result<Vec<User>>;
23 async fn get_user_by_id(&self, id: UserId) -> Result<Option<User>>;
24 async fn get_user_metrics_id(&self, id: UserId) -> Result<String>;
25 async fn get_users_by_ids(&self, ids: Vec<UserId>) -> Result<Vec<User>>;
26 async fn get_users_with_no_invites(&self, invited_by_another_user: bool) -> Result<Vec<User>>;
27 async fn get_user_by_github_account(
28 &self,
29 github_login: &str,
30 github_user_id: Option<i32>,
31 ) -> Result<Option<User>>;
32 async fn set_user_is_admin(&self, id: UserId, is_admin: bool) -> Result<()>;
33 async fn set_user_connected_once(&self, id: UserId, connected_once: bool) -> Result<()>;
34 async fn destroy_user(&self, id: UserId) -> Result<()>;
35
36 async fn set_invite_count_for_user(&self, id: UserId, count: u32) -> Result<()>;
37 async fn get_invite_code_for_user(&self, id: UserId) -> Result<Option<(String, u32)>>;
38 async fn get_user_for_invite_code(&self, code: &str) -> Result<User>;
39 async fn create_invite_from_code(
40 &self,
41 code: &str,
42 email_address: &str,
43 device_id: Option<&str>,
44 ) -> Result<Invite>;
45
46 async fn create_signup(&self, signup: Signup) -> Result<()>;
47 async fn get_waitlist_summary(&self) -> Result<WaitlistSummary>;
48 async fn get_unsent_invites(&self, count: usize) -> Result<Vec<Invite>>;
49 async fn record_sent_invites(&self, invites: &[Invite]) -> Result<()>;
50 async fn create_user_from_invite(
51 &self,
52 invite: &Invite,
53 user: NewUserParams,
54 ) -> Result<NewUserResult>;
55
56 /// Registers a new project for the given user.
57 async fn register_project(&self, host_user_id: UserId) -> Result<ProjectId>;
58
59 /// Unregisters a project for the given project id.
60 async fn unregister_project(&self, project_id: ProjectId) -> Result<()>;
61
62 /// Update file counts by extension for the given project and worktree.
63 async fn update_worktree_extensions(
64 &self,
65 project_id: ProjectId,
66 worktree_id: u64,
67 extensions: HashMap<String, u32>,
68 ) -> Result<()>;
69
70 /// Get the file counts on the given project keyed by their worktree and extension.
71 async fn get_project_extensions(
72 &self,
73 project_id: ProjectId,
74 ) -> Result<HashMap<u64, HashMap<String, usize>>>;
75
76 /// Record which users have been active in which projects during
77 /// a given period of time.
78 async fn record_user_activity(
79 &self,
80 time_period: Range<OffsetDateTime>,
81 active_projects: &[(UserId, ProjectId)],
82 ) -> Result<()>;
83
84 /// Get the number of users who have been active in the given
85 /// time period for at least the given time duration.
86 async fn get_active_user_count(
87 &self,
88 time_period: Range<OffsetDateTime>,
89 min_duration: Duration,
90 only_collaborative: bool,
91 ) -> Result<usize>;
92
93 /// Get the users that have been most active during the given time period,
94 /// along with the amount of time they have been active in each project.
95 async fn get_top_users_activity_summary(
96 &self,
97 time_period: Range<OffsetDateTime>,
98 max_user_count: usize,
99 ) -> Result<Vec<UserActivitySummary>>;
100
101 /// Get the project activity for the given user and time period.
102 async fn get_user_activity_timeline(
103 &self,
104 time_period: Range<OffsetDateTime>,
105 user_id: UserId,
106 ) -> Result<Vec<UserActivityPeriod>>;
107
108 async fn get_contacts(&self, id: UserId) -> Result<Vec<Contact>>;
109 async fn has_contact(&self, user_id_a: UserId, user_id_b: UserId) -> Result<bool>;
110 async fn send_contact_request(&self, requester_id: UserId, responder_id: UserId) -> Result<()>;
111 async fn remove_contact(&self, requester_id: UserId, responder_id: UserId) -> Result<()>;
112 async fn dismiss_contact_notification(
113 &self,
114 responder_id: UserId,
115 requester_id: UserId,
116 ) -> Result<()>;
117 async fn respond_to_contact_request(
118 &self,
119 responder_id: UserId,
120 requester_id: UserId,
121 accept: bool,
122 ) -> Result<()>;
123
124 async fn create_access_token_hash(
125 &self,
126 user_id: UserId,
127 access_token_hash: &str,
128 max_access_token_count: usize,
129 ) -> Result<()>;
130 async fn get_access_token_hashes(&self, user_id: UserId) -> Result<Vec<String>>;
131
132 #[cfg(any(test, feature = "seed-support"))]
133 async fn find_org_by_slug(&self, slug: &str) -> Result<Option<Org>>;
134 #[cfg(any(test, feature = "seed-support"))]
135 async fn create_org(&self, name: &str, slug: &str) -> Result<OrgId>;
136 #[cfg(any(test, feature = "seed-support"))]
137 async fn add_org_member(&self, org_id: OrgId, user_id: UserId, is_admin: bool) -> Result<()>;
138 #[cfg(any(test, feature = "seed-support"))]
139 async fn create_org_channel(&self, org_id: OrgId, name: &str) -> Result<ChannelId>;
140 #[cfg(any(test, feature = "seed-support"))]
141
142 async fn get_org_channels(&self, org_id: OrgId) -> Result<Vec<Channel>>;
143 async fn get_accessible_channels(&self, user_id: UserId) -> Result<Vec<Channel>>;
144 async fn can_user_access_channel(&self, user_id: UserId, channel_id: ChannelId)
145 -> Result<bool>;
146
147 #[cfg(any(test, feature = "seed-support"))]
148 async fn add_channel_member(
149 &self,
150 channel_id: ChannelId,
151 user_id: UserId,
152 is_admin: bool,
153 ) -> Result<()>;
154 async fn create_channel_message(
155 &self,
156 channel_id: ChannelId,
157 sender_id: UserId,
158 body: &str,
159 timestamp: OffsetDateTime,
160 nonce: u128,
161 ) -> Result<MessageId>;
162 async fn get_channel_messages(
163 &self,
164 channel_id: ChannelId,
165 count: usize,
166 before_id: Option<MessageId>,
167 ) -> Result<Vec<ChannelMessage>>;
168
169 #[cfg(test)]
170 async fn teardown(&self, url: &str);
171
172 #[cfg(test)]
173 fn as_fake(&self) -> Option<&FakeDb>;
174}
175
176pub struct PostgresDb {
177 pool: sqlx::PgPool,
178}
179
180impl PostgresDb {
181 pub async fn new(url: &str, max_connections: u32) -> Result<Self> {
182 let pool = DbOptions::new()
183 .max_connections(max_connections)
184 .connect(url)
185 .await
186 .context("failed to connect to postgres database")?;
187 Ok(Self { pool })
188 }
189
190 pub fn fuzzy_like_string(string: &str) -> String {
191 let mut result = String::with_capacity(string.len() * 2 + 1);
192 for c in string.chars() {
193 if c.is_alphanumeric() {
194 result.push('%');
195 result.push(c);
196 }
197 }
198 result.push('%');
199 result
200 }
201}
202
203#[async_trait]
204impl Db for PostgresDb {
205 // users
206
207 async fn create_user(
208 &self,
209 email_address: &str,
210 admin: bool,
211 params: NewUserParams,
212 ) -> Result<NewUserResult> {
213 let query = "
214 INSERT INTO users (email_address, github_login, github_user_id, admin)
215 VALUES ($1, $2, $3, $4)
216 ON CONFLICT (github_login) DO UPDATE SET github_login = excluded.github_login
217 RETURNING id, metrics_id::text
218 ";
219 let (user_id, metrics_id): (UserId, String) = sqlx::query_as(query)
220 .bind(email_address)
221 .bind(params.github_login)
222 .bind(params.github_user_id)
223 .bind(admin)
224 .fetch_one(&self.pool)
225 .await?;
226 Ok(NewUserResult {
227 user_id,
228 metrics_id,
229 signup_device_id: None,
230 inviting_user_id: None,
231 })
232 }
233
234 async fn get_all_users(&self, page: u32, limit: u32) -> Result<Vec<User>> {
235 let query = "SELECT * FROM users ORDER BY github_login ASC LIMIT $1 OFFSET $2";
236 Ok(sqlx::query_as(query)
237 .bind(limit as i32)
238 .bind((page * limit) as i32)
239 .fetch_all(&self.pool)
240 .await?)
241 }
242
243 async fn fuzzy_search_users(&self, name_query: &str, limit: u32) -> Result<Vec<User>> {
244 let like_string = Self::fuzzy_like_string(name_query);
245 let query = "
246 SELECT users.*
247 FROM users
248 WHERE github_login ILIKE $1
249 ORDER BY github_login <-> $2
250 LIMIT $3
251 ";
252 Ok(sqlx::query_as(query)
253 .bind(like_string)
254 .bind(name_query)
255 .bind(limit as i32)
256 .fetch_all(&self.pool)
257 .await?)
258 }
259
260 async fn get_user_by_id(&self, id: UserId) -> Result<Option<User>> {
261 let users = self.get_users_by_ids(vec![id]).await?;
262 Ok(users.into_iter().next())
263 }
264
265 async fn get_user_metrics_id(&self, id: UserId) -> Result<String> {
266 let query = "
267 SELECT metrics_id::text
268 FROM users
269 WHERE id = $1
270 ";
271 Ok(sqlx::query_scalar(query)
272 .bind(id)
273 .fetch_one(&self.pool)
274 .await?)
275 }
276
277 async fn get_users_by_ids(&self, ids: Vec<UserId>) -> Result<Vec<User>> {
278 let ids = ids.into_iter().map(|id| id.0).collect::<Vec<_>>();
279 let query = "
280 SELECT users.*
281 FROM users
282 WHERE users.id = ANY ($1)
283 ";
284 Ok(sqlx::query_as(query)
285 .bind(&ids)
286 .fetch_all(&self.pool)
287 .await?)
288 }
289
290 async fn get_users_with_no_invites(&self, invited_by_another_user: bool) -> Result<Vec<User>> {
291 let query = format!(
292 "
293 SELECT users.*
294 FROM users
295 WHERE invite_count = 0
296 AND inviter_id IS{} NULL
297 ",
298 if invited_by_another_user { " NOT" } else { "" }
299 );
300
301 Ok(sqlx::query_as(&query).fetch_all(&self.pool).await?)
302 }
303
304 async fn get_user_by_github_account(
305 &self,
306 github_login: &str,
307 github_user_id: Option<i32>,
308 ) -> Result<Option<User>> {
309 if let Some(github_user_id) = github_user_id {
310 let mut user = sqlx::query_as::<_, User>(
311 "
312 UPDATE users
313 SET github_login = $1
314 WHERE github_user_id = $2
315 RETURNING *
316 ",
317 )
318 .bind(github_login)
319 .bind(github_user_id)
320 .fetch_optional(&self.pool)
321 .await?;
322
323 if user.is_none() {
324 user = sqlx::query_as::<_, User>(
325 "
326 UPDATE users
327 SET github_user_id = $1
328 WHERE github_login = $2
329 RETURNING *
330 ",
331 )
332 .bind(github_user_id)
333 .bind(github_login)
334 .fetch_optional(&self.pool)
335 .await?;
336 }
337
338 Ok(user)
339 } else {
340 Ok(sqlx::query_as(
341 "
342 SELECT * FROM users
343 WHERE github_login = $1
344 LIMIT 1
345 ",
346 )
347 .bind(github_login)
348 .fetch_optional(&self.pool)
349 .await?)
350 }
351 }
352
353 async fn set_user_is_admin(&self, id: UserId, is_admin: bool) -> Result<()> {
354 let query = "UPDATE users SET admin = $1 WHERE id = $2";
355 Ok(sqlx::query(query)
356 .bind(is_admin)
357 .bind(id.0)
358 .execute(&self.pool)
359 .await
360 .map(drop)?)
361 }
362
363 async fn set_user_connected_once(&self, id: UserId, connected_once: bool) -> Result<()> {
364 let query = "UPDATE users SET connected_once = $1 WHERE id = $2";
365 Ok(sqlx::query(query)
366 .bind(connected_once)
367 .bind(id.0)
368 .execute(&self.pool)
369 .await
370 .map(drop)?)
371 }
372
373 async fn destroy_user(&self, id: UserId) -> Result<()> {
374 let query = "DELETE FROM access_tokens WHERE user_id = $1;";
375 sqlx::query(query)
376 .bind(id.0)
377 .execute(&self.pool)
378 .await
379 .map(drop)?;
380 let query = "DELETE FROM users WHERE id = $1;";
381 Ok(sqlx::query(query)
382 .bind(id.0)
383 .execute(&self.pool)
384 .await
385 .map(drop)?)
386 }
387
388 // signups
389
390 async fn create_signup(&self, signup: Signup) -> Result<()> {
391 sqlx::query(
392 "
393 INSERT INTO signups
394 (
395 email_address,
396 email_confirmation_code,
397 email_confirmation_sent,
398 platform_linux,
399 platform_mac,
400 platform_windows,
401 platform_unknown,
402 editor_features,
403 programming_languages,
404 device_id
405 )
406 VALUES
407 ($1, $2, 'f', $3, $4, $5, 'f', $6, $7, $8)
408 RETURNING id
409 ",
410 )
411 .bind(&signup.email_address)
412 .bind(&random_email_confirmation_code())
413 .bind(&signup.platform_linux)
414 .bind(&signup.platform_mac)
415 .bind(&signup.platform_windows)
416 .bind(&signup.editor_features)
417 .bind(&signup.programming_languages)
418 .bind(&signup.device_id)
419 .execute(&self.pool)
420 .await?;
421 Ok(())
422 }
423
424 async fn get_waitlist_summary(&self) -> Result<WaitlistSummary> {
425 Ok(sqlx::query_as(
426 "
427 SELECT
428 COUNT(*) as count,
429 COALESCE(SUM(CASE WHEN platform_linux THEN 1 ELSE 0 END), 0) as linux_count,
430 COALESCE(SUM(CASE WHEN platform_mac THEN 1 ELSE 0 END), 0) as mac_count,
431 COALESCE(SUM(CASE WHEN platform_windows THEN 1 ELSE 0 END), 0) as windows_count,
432 COALESCE(SUM(CASE WHEN platform_unknown THEN 1 ELSE 0 END), 0) as unknown_count
433 FROM (
434 SELECT *
435 FROM signups
436 WHERE
437 NOT email_confirmation_sent
438 ) AS unsent
439 ",
440 )
441 .fetch_one(&self.pool)
442 .await?)
443 }
444
445 async fn get_unsent_invites(&self, count: usize) -> Result<Vec<Invite>> {
446 Ok(sqlx::query_as(
447 "
448 SELECT
449 email_address, email_confirmation_code
450 FROM signups
451 WHERE
452 NOT email_confirmation_sent AND
453 (platform_mac OR platform_unknown)
454 LIMIT $1
455 ",
456 )
457 .bind(count as i32)
458 .fetch_all(&self.pool)
459 .await?)
460 }
461
462 async fn record_sent_invites(&self, invites: &[Invite]) -> Result<()> {
463 sqlx::query(
464 "
465 UPDATE signups
466 SET email_confirmation_sent = 't'
467 WHERE email_address = ANY ($1)
468 ",
469 )
470 .bind(
471 &invites
472 .iter()
473 .map(|s| s.email_address.as_str())
474 .collect::<Vec<_>>(),
475 )
476 .execute(&self.pool)
477 .await?;
478 Ok(())
479 }
480
481 async fn create_user_from_invite(
482 &self,
483 invite: &Invite,
484 user: NewUserParams,
485 ) -> Result<NewUserResult> {
486 let mut tx = self.pool.begin().await?;
487
488 let (signup_id, existing_user_id, inviting_user_id, signup_device_id): (
489 i32,
490 Option<UserId>,
491 Option<UserId>,
492 Option<String>,
493 ) = sqlx::query_as(
494 "
495 SELECT id, user_id, inviting_user_id, device_id
496 FROM signups
497 WHERE
498 email_address = $1 AND
499 email_confirmation_code = $2
500 ",
501 )
502 .bind(&invite.email_address)
503 .bind(&invite.email_confirmation_code)
504 .fetch_optional(&mut tx)
505 .await?
506 .ok_or_else(|| Error::Http(StatusCode::NOT_FOUND, "no such invite".to_string()))?;
507
508 if existing_user_id.is_some() {
509 Err(Error::Http(
510 StatusCode::UNPROCESSABLE_ENTITY,
511 "invitation already redeemed".to_string(),
512 ))?;
513 }
514
515 let (user_id, metrics_id): (UserId, String) = sqlx::query_as(
516 "
517 INSERT INTO users
518 (email_address, github_login, github_user_id, admin, invite_count, invite_code)
519 VALUES
520 ($1, $2, $3, 'f', $4, $5)
521 RETURNING id, metrics_id::text
522 ",
523 )
524 .bind(&invite.email_address)
525 .bind(&user.github_login)
526 .bind(&user.github_user_id)
527 .bind(&user.invite_count)
528 .bind(random_invite_code())
529 .fetch_one(&mut tx)
530 .await?;
531
532 sqlx::query(
533 "
534 UPDATE signups
535 SET user_id = $1
536 WHERE id = $2
537 ",
538 )
539 .bind(&user_id)
540 .bind(&signup_id)
541 .execute(&mut tx)
542 .await?;
543
544 if let Some(inviting_user_id) = inviting_user_id {
545 let id: Option<UserId> = sqlx::query_scalar(
546 "
547 UPDATE users
548 SET invite_count = invite_count - 1
549 WHERE id = $1 AND invite_count > 0
550 RETURNING id
551 ",
552 )
553 .bind(&inviting_user_id)
554 .fetch_optional(&mut tx)
555 .await?;
556
557 if id.is_none() {
558 Err(Error::Http(
559 StatusCode::UNAUTHORIZED,
560 "no invites remaining".to_string(),
561 ))?;
562 }
563
564 sqlx::query(
565 "
566 INSERT INTO contacts
567 (user_id_a, user_id_b, a_to_b, should_notify, accepted)
568 VALUES
569 ($1, $2, 't', 't', 't')
570 ",
571 )
572 .bind(inviting_user_id)
573 .bind(user_id)
574 .execute(&mut tx)
575 .await?;
576 }
577
578 tx.commit().await?;
579 Ok(NewUserResult {
580 user_id,
581 metrics_id,
582 inviting_user_id,
583 signup_device_id,
584 })
585 }
586
587 // invite codes
588
589 async fn set_invite_count_for_user(&self, id: UserId, count: u32) -> Result<()> {
590 let mut tx = self.pool.begin().await?;
591 if count > 0 {
592 sqlx::query(
593 "
594 UPDATE users
595 SET invite_code = $1
596 WHERE id = $2 AND invite_code IS NULL
597 ",
598 )
599 .bind(random_invite_code())
600 .bind(id)
601 .execute(&mut tx)
602 .await?;
603 }
604
605 sqlx::query(
606 "
607 UPDATE users
608 SET invite_count = $1
609 WHERE id = $2
610 ",
611 )
612 .bind(count as i32)
613 .bind(id)
614 .execute(&mut tx)
615 .await?;
616 tx.commit().await?;
617 Ok(())
618 }
619
620 async fn get_invite_code_for_user(&self, id: UserId) -> Result<Option<(String, u32)>> {
621 let result: Option<(String, i32)> = sqlx::query_as(
622 "
623 SELECT invite_code, invite_count
624 FROM users
625 WHERE id = $1 AND invite_code IS NOT NULL
626 ",
627 )
628 .bind(id)
629 .fetch_optional(&self.pool)
630 .await?;
631 if let Some((code, count)) = result {
632 Ok(Some((code, count.try_into().map_err(anyhow::Error::new)?)))
633 } else {
634 Ok(None)
635 }
636 }
637
638 async fn get_user_for_invite_code(&self, code: &str) -> Result<User> {
639 sqlx::query_as(
640 "
641 SELECT *
642 FROM users
643 WHERE invite_code = $1
644 ",
645 )
646 .bind(code)
647 .fetch_optional(&self.pool)
648 .await?
649 .ok_or_else(|| {
650 Error::Http(
651 StatusCode::NOT_FOUND,
652 "that invite code does not exist".to_string(),
653 )
654 })
655 }
656
657 async fn create_invite_from_code(
658 &self,
659 code: &str,
660 email_address: &str,
661 device_id: Option<&str>,
662 ) -> Result<Invite> {
663 let mut tx = self.pool.begin().await?;
664
665 let existing_user: Option<UserId> = sqlx::query_scalar(
666 "
667 SELECT id
668 FROM users
669 WHERE email_address = $1
670 ",
671 )
672 .bind(email_address)
673 .fetch_optional(&mut tx)
674 .await?;
675 if existing_user.is_some() {
676 Err(anyhow!("email address is already in use"))?;
677 }
678
679 let row: Option<(UserId, i32)> = sqlx::query_as(
680 "
681 SELECT id, invite_count
682 FROM users
683 WHERE invite_code = $1
684 ",
685 )
686 .bind(code)
687 .fetch_optional(&mut tx)
688 .await?;
689
690 let (inviter_id, invite_count) = match row {
691 Some(row) => row,
692 None => Err(Error::Http(
693 StatusCode::NOT_FOUND,
694 "invite code not found".to_string(),
695 ))?,
696 };
697
698 if invite_count == 0 {
699 Err(Error::Http(
700 StatusCode::UNAUTHORIZED,
701 "no invites remaining".to_string(),
702 ))?;
703 }
704
705 let email_confirmation_code: String = sqlx::query_scalar(
706 "
707 INSERT INTO signups
708 (
709 email_address,
710 email_confirmation_code,
711 email_confirmation_sent,
712 inviting_user_id,
713 platform_linux,
714 platform_mac,
715 platform_windows,
716 platform_unknown,
717 device_id
718 )
719 VALUES
720 ($1, $2, 'f', $3, 'f', 'f', 'f', 't', $4)
721 ON CONFLICT (email_address)
722 DO UPDATE SET
723 inviting_user_id = excluded.inviting_user_id
724 RETURNING email_confirmation_code
725 ",
726 )
727 .bind(&email_address)
728 .bind(&random_email_confirmation_code())
729 .bind(&inviter_id)
730 .bind(&device_id)
731 .fetch_one(&mut tx)
732 .await?;
733
734 tx.commit().await?;
735
736 Ok(Invite {
737 email_address: email_address.into(),
738 email_confirmation_code,
739 })
740 }
741
742 // projects
743
744 async fn register_project(&self, host_user_id: UserId) -> Result<ProjectId> {
745 Ok(sqlx::query_scalar(
746 "
747 INSERT INTO projects(host_user_id)
748 VALUES ($1)
749 RETURNING id
750 ",
751 )
752 .bind(host_user_id)
753 .fetch_one(&self.pool)
754 .await
755 .map(ProjectId)?)
756 }
757
758 async fn unregister_project(&self, project_id: ProjectId) -> Result<()> {
759 sqlx::query(
760 "
761 UPDATE projects
762 SET unregistered = 't'
763 WHERE id = $1
764 ",
765 )
766 .bind(project_id)
767 .execute(&self.pool)
768 .await?;
769 Ok(())
770 }
771
772 async fn update_worktree_extensions(
773 &self,
774 project_id: ProjectId,
775 worktree_id: u64,
776 extensions: HashMap<String, u32>,
777 ) -> Result<()> {
778 if extensions.is_empty() {
779 return Ok(());
780 }
781
782 let mut query = QueryBuilder::new(
783 "INSERT INTO worktree_extensions (project_id, worktree_id, extension, count)",
784 );
785 query.push_values(extensions, |mut query, (extension, count)| {
786 query
787 .push_bind(project_id)
788 .push_bind(worktree_id as i32)
789 .push_bind(extension)
790 .push_bind(count as i32);
791 });
792 query.push(
793 "
794 ON CONFLICT (project_id, worktree_id, extension) DO UPDATE SET
795 count = excluded.count
796 ",
797 );
798 query.build().execute(&self.pool).await?;
799
800 Ok(())
801 }
802
803 async fn get_project_extensions(
804 &self,
805 project_id: ProjectId,
806 ) -> Result<HashMap<u64, HashMap<String, usize>>> {
807 #[derive(Clone, Debug, Default, FromRow, Serialize, PartialEq)]
808 struct WorktreeExtension {
809 worktree_id: i32,
810 extension: String,
811 count: i32,
812 }
813
814 let query = "
815 SELECT worktree_id, extension, count
816 FROM worktree_extensions
817 WHERE project_id = $1
818 ";
819 let counts = sqlx::query_as::<_, WorktreeExtension>(query)
820 .bind(&project_id)
821 .fetch_all(&self.pool)
822 .await?;
823
824 let mut extension_counts = HashMap::default();
825 for count in counts {
826 extension_counts
827 .entry(count.worktree_id as u64)
828 .or_insert_with(HashMap::default)
829 .insert(count.extension, count.count as usize);
830 }
831 Ok(extension_counts)
832 }
833
834 async fn record_user_activity(
835 &self,
836 time_period: Range<OffsetDateTime>,
837 projects: &[(UserId, ProjectId)],
838 ) -> Result<()> {
839 let query = "
840 INSERT INTO project_activity_periods
841 (ended_at, duration_millis, user_id, project_id)
842 VALUES
843 ($1, $2, $3, $4);
844 ";
845
846 let mut tx = self.pool.begin().await?;
847 let duration_millis =
848 ((time_period.end - time_period.start).as_seconds_f64() * 1000.0) as i32;
849 for (user_id, project_id) in projects {
850 sqlx::query(query)
851 .bind(time_period.end)
852 .bind(duration_millis)
853 .bind(user_id)
854 .bind(project_id)
855 .execute(&mut tx)
856 .await?;
857 }
858 tx.commit().await?;
859
860 Ok(())
861 }
862
863 async fn get_active_user_count(
864 &self,
865 time_period: Range<OffsetDateTime>,
866 min_duration: Duration,
867 only_collaborative: bool,
868 ) -> Result<usize> {
869 let mut with_clause = String::new();
870 with_clause.push_str("WITH\n");
871 with_clause.push_str(
872 "
873 project_durations AS (
874 SELECT user_id, project_id, SUM(duration_millis) AS project_duration
875 FROM project_activity_periods
876 WHERE $1 < ended_at AND ended_at <= $2
877 GROUP BY user_id, project_id
878 ),
879 ",
880 );
881 with_clause.push_str(
882 "
883 project_collaborators as (
884 SELECT project_id, COUNT(DISTINCT user_id) as max_collaborators
885 FROM project_durations
886 GROUP BY project_id
887 ),
888 ",
889 );
890
891 if only_collaborative {
892 with_clause.push_str(
893 "
894 user_durations AS (
895 SELECT user_id, SUM(project_duration) as total_duration
896 FROM project_durations, project_collaborators
897 WHERE
898 project_durations.project_id = project_collaborators.project_id AND
899 max_collaborators > 1
900 GROUP BY user_id
901 ORDER BY total_duration DESC
902 LIMIT $3
903 )
904 ",
905 );
906 } else {
907 with_clause.push_str(
908 "
909 user_durations AS (
910 SELECT user_id, SUM(project_duration) as total_duration
911 FROM project_durations
912 GROUP BY user_id
913 ORDER BY total_duration DESC
914 LIMIT $3
915 )
916 ",
917 );
918 }
919
920 let query = format!(
921 "
922 {with_clause}
923 SELECT count(user_durations.user_id)
924 FROM user_durations
925 WHERE user_durations.total_duration >= $3
926 "
927 );
928
929 let count: i64 = sqlx::query_scalar(&query)
930 .bind(time_period.start)
931 .bind(time_period.end)
932 .bind(min_duration.as_millis() as i64)
933 .fetch_one(&self.pool)
934 .await?;
935 Ok(count as usize)
936 }
937
938 async fn get_top_users_activity_summary(
939 &self,
940 time_period: Range<OffsetDateTime>,
941 max_user_count: usize,
942 ) -> Result<Vec<UserActivitySummary>> {
943 let query = "
944 WITH
945 project_durations AS (
946 SELECT user_id, project_id, SUM(duration_millis) AS project_duration
947 FROM project_activity_periods
948 WHERE $1 < ended_at AND ended_at <= $2
949 GROUP BY user_id, project_id
950 ),
951 user_durations AS (
952 SELECT user_id, SUM(project_duration) as total_duration
953 FROM project_durations
954 GROUP BY user_id
955 ORDER BY total_duration DESC
956 LIMIT $3
957 ),
958 project_collaborators as (
959 SELECT project_id, COUNT(DISTINCT user_id) as max_collaborators
960 FROM project_durations
961 GROUP BY project_id
962 )
963 SELECT user_durations.user_id, users.github_login, project_durations.project_id, project_duration, max_collaborators
964 FROM user_durations, project_durations, project_collaborators, users
965 WHERE
966 user_durations.user_id = project_durations.user_id AND
967 user_durations.user_id = users.id AND
968 project_durations.project_id = project_collaborators.project_id
969 ORDER BY total_duration DESC, user_id ASC, project_id ASC
970 ";
971
972 let mut rows = sqlx::query_as::<_, (UserId, String, ProjectId, i64, i64)>(query)
973 .bind(time_period.start)
974 .bind(time_period.end)
975 .bind(max_user_count as i32)
976 .fetch(&self.pool);
977
978 let mut result = Vec::<UserActivitySummary>::new();
979 while let Some(row) = rows.next().await {
980 let (user_id, github_login, project_id, duration_millis, project_collaborators) = row?;
981 let project_id = project_id;
982 let duration = Duration::from_millis(duration_millis as u64);
983 let project_activity = ProjectActivitySummary {
984 id: project_id,
985 duration,
986 max_collaborators: project_collaborators as usize,
987 };
988 if let Some(last_summary) = result.last_mut() {
989 if last_summary.id == user_id {
990 last_summary.project_activity.push(project_activity);
991 continue;
992 }
993 }
994 result.push(UserActivitySummary {
995 id: user_id,
996 project_activity: vec![project_activity],
997 github_login,
998 });
999 }
1000
1001 Ok(result)
1002 }
1003
1004 async fn get_user_activity_timeline(
1005 &self,
1006 time_period: Range<OffsetDateTime>,
1007 user_id: UserId,
1008 ) -> Result<Vec<UserActivityPeriod>> {
1009 const COALESCE_THRESHOLD: Duration = Duration::from_secs(30);
1010
1011 let query = "
1012 SELECT
1013 project_activity_periods.ended_at,
1014 project_activity_periods.duration_millis,
1015 project_activity_periods.project_id,
1016 worktree_extensions.extension,
1017 worktree_extensions.count
1018 FROM project_activity_periods
1019 LEFT OUTER JOIN
1020 worktree_extensions
1021 ON
1022 project_activity_periods.project_id = worktree_extensions.project_id
1023 WHERE
1024 project_activity_periods.user_id = $1 AND
1025 $2 < project_activity_periods.ended_at AND
1026 project_activity_periods.ended_at <= $3
1027 ORDER BY project_activity_periods.id ASC
1028 ";
1029
1030 let mut rows = sqlx::query_as::<
1031 _,
1032 (
1033 PrimitiveDateTime,
1034 i32,
1035 ProjectId,
1036 Option<String>,
1037 Option<i32>,
1038 ),
1039 >(query)
1040 .bind(user_id)
1041 .bind(time_period.start)
1042 .bind(time_period.end)
1043 .fetch(&self.pool);
1044
1045 let mut time_periods: HashMap<ProjectId, Vec<UserActivityPeriod>> = Default::default();
1046 while let Some(row) = rows.next().await {
1047 let (ended_at, duration_millis, project_id, extension, extension_count) = row?;
1048 let ended_at = ended_at.assume_utc();
1049 let duration = Duration::from_millis(duration_millis as u64);
1050 let started_at = ended_at - duration;
1051 let project_time_periods = time_periods.entry(project_id).or_default();
1052
1053 if let Some(prev_duration) = project_time_periods.last_mut() {
1054 if started_at <= prev_duration.end + COALESCE_THRESHOLD
1055 && ended_at >= prev_duration.start
1056 {
1057 prev_duration.end = cmp::max(prev_duration.end, ended_at);
1058 } else {
1059 project_time_periods.push(UserActivityPeriod {
1060 project_id,
1061 start: started_at,
1062 end: ended_at,
1063 extensions: Default::default(),
1064 });
1065 }
1066 } else {
1067 project_time_periods.push(UserActivityPeriod {
1068 project_id,
1069 start: started_at,
1070 end: ended_at,
1071 extensions: Default::default(),
1072 });
1073 }
1074
1075 if let Some((extension, extension_count)) = extension.zip(extension_count) {
1076 project_time_periods
1077 .last_mut()
1078 .unwrap()
1079 .extensions
1080 .insert(extension, extension_count as usize);
1081 }
1082 }
1083
1084 let mut durations = time_periods.into_values().flatten().collect::<Vec<_>>();
1085 durations.sort_unstable_by_key(|duration| duration.start);
1086 Ok(durations)
1087 }
1088
1089 // contacts
1090
1091 async fn get_contacts(&self, user_id: UserId) -> Result<Vec<Contact>> {
1092 let query = "
1093 SELECT user_id_a, user_id_b, a_to_b, accepted, should_notify
1094 FROM contacts
1095 WHERE user_id_a = $1 OR user_id_b = $1;
1096 ";
1097
1098 let mut rows = sqlx::query_as::<_, (UserId, UserId, bool, bool, bool)>(query)
1099 .bind(user_id)
1100 .fetch(&self.pool);
1101
1102 let mut contacts = Vec::new();
1103 while let Some(row) = rows.next().await {
1104 let (user_id_a, user_id_b, a_to_b, accepted, should_notify) = row?;
1105
1106 if user_id_a == user_id {
1107 if accepted {
1108 contacts.push(Contact::Accepted {
1109 user_id: user_id_b,
1110 should_notify: should_notify && a_to_b,
1111 });
1112 } else if a_to_b {
1113 contacts.push(Contact::Outgoing { user_id: user_id_b })
1114 } else {
1115 contacts.push(Contact::Incoming {
1116 user_id: user_id_b,
1117 should_notify,
1118 });
1119 }
1120 } else if accepted {
1121 contacts.push(Contact::Accepted {
1122 user_id: user_id_a,
1123 should_notify: should_notify && !a_to_b,
1124 });
1125 } else if a_to_b {
1126 contacts.push(Contact::Incoming {
1127 user_id: user_id_a,
1128 should_notify,
1129 });
1130 } else {
1131 contacts.push(Contact::Outgoing { user_id: user_id_a });
1132 }
1133 }
1134
1135 contacts.sort_unstable_by_key(|contact| contact.user_id());
1136
1137 Ok(contacts)
1138 }
1139
1140 async fn has_contact(&self, user_id_1: UserId, user_id_2: UserId) -> Result<bool> {
1141 let (id_a, id_b) = if user_id_1 < user_id_2 {
1142 (user_id_1, user_id_2)
1143 } else {
1144 (user_id_2, user_id_1)
1145 };
1146
1147 let query = "
1148 SELECT 1 FROM contacts
1149 WHERE user_id_a = $1 AND user_id_b = $2 AND accepted = 't'
1150 LIMIT 1
1151 ";
1152 Ok(sqlx::query_scalar::<_, i32>(query)
1153 .bind(id_a.0)
1154 .bind(id_b.0)
1155 .fetch_optional(&self.pool)
1156 .await?
1157 .is_some())
1158 }
1159
1160 async fn send_contact_request(&self, sender_id: UserId, receiver_id: UserId) -> Result<()> {
1161 let (id_a, id_b, a_to_b) = if sender_id < receiver_id {
1162 (sender_id, receiver_id, true)
1163 } else {
1164 (receiver_id, sender_id, false)
1165 };
1166 let query = "
1167 INSERT into contacts (user_id_a, user_id_b, a_to_b, accepted, should_notify)
1168 VALUES ($1, $2, $3, 'f', 't')
1169 ON CONFLICT (user_id_a, user_id_b) DO UPDATE
1170 SET
1171 accepted = 't',
1172 should_notify = 'f'
1173 WHERE
1174 NOT contacts.accepted AND
1175 ((contacts.a_to_b = excluded.a_to_b AND contacts.user_id_a = excluded.user_id_b) OR
1176 (contacts.a_to_b != excluded.a_to_b AND contacts.user_id_a = excluded.user_id_a));
1177 ";
1178 let result = sqlx::query(query)
1179 .bind(id_a.0)
1180 .bind(id_b.0)
1181 .bind(a_to_b)
1182 .execute(&self.pool)
1183 .await?;
1184
1185 if result.rows_affected() == 1 {
1186 Ok(())
1187 } else {
1188 Err(anyhow!("contact already requested"))?
1189 }
1190 }
1191
1192 async fn remove_contact(&self, requester_id: UserId, responder_id: UserId) -> Result<()> {
1193 let (id_a, id_b) = if responder_id < requester_id {
1194 (responder_id, requester_id)
1195 } else {
1196 (requester_id, responder_id)
1197 };
1198 let query = "
1199 DELETE FROM contacts
1200 WHERE user_id_a = $1 AND user_id_b = $2;
1201 ";
1202 let result = sqlx::query(query)
1203 .bind(id_a.0)
1204 .bind(id_b.0)
1205 .execute(&self.pool)
1206 .await?;
1207
1208 if result.rows_affected() == 1 {
1209 Ok(())
1210 } else {
1211 Err(anyhow!("no such contact"))?
1212 }
1213 }
1214
1215 async fn dismiss_contact_notification(
1216 &self,
1217 user_id: UserId,
1218 contact_user_id: UserId,
1219 ) -> Result<()> {
1220 let (id_a, id_b, a_to_b) = if user_id < contact_user_id {
1221 (user_id, contact_user_id, true)
1222 } else {
1223 (contact_user_id, user_id, false)
1224 };
1225
1226 let query = "
1227 UPDATE contacts
1228 SET should_notify = 'f'
1229 WHERE
1230 user_id_a = $1 AND user_id_b = $2 AND
1231 (
1232 (a_to_b = $3 AND accepted) OR
1233 (a_to_b != $3 AND NOT accepted)
1234 );
1235 ";
1236
1237 let result = sqlx::query(query)
1238 .bind(id_a.0)
1239 .bind(id_b.0)
1240 .bind(a_to_b)
1241 .execute(&self.pool)
1242 .await?;
1243
1244 if result.rows_affected() == 0 {
1245 Err(anyhow!("no such contact request"))?;
1246 }
1247
1248 Ok(())
1249 }
1250
1251 async fn respond_to_contact_request(
1252 &self,
1253 responder_id: UserId,
1254 requester_id: UserId,
1255 accept: bool,
1256 ) -> Result<()> {
1257 let (id_a, id_b, a_to_b) = if responder_id < requester_id {
1258 (responder_id, requester_id, false)
1259 } else {
1260 (requester_id, responder_id, true)
1261 };
1262 let result = if accept {
1263 let query = "
1264 UPDATE contacts
1265 SET accepted = 't', should_notify = 't'
1266 WHERE user_id_a = $1 AND user_id_b = $2 AND a_to_b = $3;
1267 ";
1268 sqlx::query(query)
1269 .bind(id_a.0)
1270 .bind(id_b.0)
1271 .bind(a_to_b)
1272 .execute(&self.pool)
1273 .await?
1274 } else {
1275 let query = "
1276 DELETE FROM contacts
1277 WHERE user_id_a = $1 AND user_id_b = $2 AND a_to_b = $3 AND NOT accepted;
1278 ";
1279 sqlx::query(query)
1280 .bind(id_a.0)
1281 .bind(id_b.0)
1282 .bind(a_to_b)
1283 .execute(&self.pool)
1284 .await?
1285 };
1286 if result.rows_affected() == 1 {
1287 Ok(())
1288 } else {
1289 Err(anyhow!("no such contact request"))?
1290 }
1291 }
1292
1293 // access tokens
1294
1295 async fn create_access_token_hash(
1296 &self,
1297 user_id: UserId,
1298 access_token_hash: &str,
1299 max_access_token_count: usize,
1300 ) -> Result<()> {
1301 let insert_query = "
1302 INSERT INTO access_tokens (user_id, hash)
1303 VALUES ($1, $2);
1304 ";
1305 let cleanup_query = "
1306 DELETE FROM access_tokens
1307 WHERE id IN (
1308 SELECT id from access_tokens
1309 WHERE user_id = $1
1310 ORDER BY id DESC
1311 OFFSET $3
1312 )
1313 ";
1314
1315 let mut tx = self.pool.begin().await?;
1316 sqlx::query(insert_query)
1317 .bind(user_id.0)
1318 .bind(access_token_hash)
1319 .execute(&mut tx)
1320 .await?;
1321 sqlx::query(cleanup_query)
1322 .bind(user_id.0)
1323 .bind(access_token_hash)
1324 .bind(max_access_token_count as i32)
1325 .execute(&mut tx)
1326 .await?;
1327 Ok(tx.commit().await?)
1328 }
1329
1330 async fn get_access_token_hashes(&self, user_id: UserId) -> Result<Vec<String>> {
1331 let query = "
1332 SELECT hash
1333 FROM access_tokens
1334 WHERE user_id = $1
1335 ORDER BY id DESC
1336 ";
1337 Ok(sqlx::query_scalar(query)
1338 .bind(user_id.0)
1339 .fetch_all(&self.pool)
1340 .await?)
1341 }
1342
1343 // orgs
1344
1345 #[allow(unused)] // Help rust-analyzer
1346 #[cfg(any(test, feature = "seed-support"))]
1347 async fn find_org_by_slug(&self, slug: &str) -> Result<Option<Org>> {
1348 let query = "
1349 SELECT *
1350 FROM orgs
1351 WHERE slug = $1
1352 ";
1353 Ok(sqlx::query_as(query)
1354 .bind(slug)
1355 .fetch_optional(&self.pool)
1356 .await?)
1357 }
1358
1359 #[cfg(any(test, feature = "seed-support"))]
1360 async fn create_org(&self, name: &str, slug: &str) -> Result<OrgId> {
1361 let query = "
1362 INSERT INTO orgs (name, slug)
1363 VALUES ($1, $2)
1364 RETURNING id
1365 ";
1366 Ok(sqlx::query_scalar(query)
1367 .bind(name)
1368 .bind(slug)
1369 .fetch_one(&self.pool)
1370 .await
1371 .map(OrgId)?)
1372 }
1373
1374 #[cfg(any(test, feature = "seed-support"))]
1375 async fn add_org_member(&self, org_id: OrgId, user_id: UserId, is_admin: bool) -> Result<()> {
1376 let query = "
1377 INSERT INTO org_memberships (org_id, user_id, admin)
1378 VALUES ($1, $2, $3)
1379 ON CONFLICT DO NOTHING
1380 ";
1381 Ok(sqlx::query(query)
1382 .bind(org_id.0)
1383 .bind(user_id.0)
1384 .bind(is_admin)
1385 .execute(&self.pool)
1386 .await
1387 .map(drop)?)
1388 }
1389
1390 // channels
1391
1392 #[cfg(any(test, feature = "seed-support"))]
1393 async fn create_org_channel(&self, org_id: OrgId, name: &str) -> Result<ChannelId> {
1394 let query = "
1395 INSERT INTO channels (owner_id, owner_is_user, name)
1396 VALUES ($1, false, $2)
1397 RETURNING id
1398 ";
1399 Ok(sqlx::query_scalar(query)
1400 .bind(org_id.0)
1401 .bind(name)
1402 .fetch_one(&self.pool)
1403 .await
1404 .map(ChannelId)?)
1405 }
1406
1407 #[allow(unused)] // Help rust-analyzer
1408 #[cfg(any(test, feature = "seed-support"))]
1409 async fn get_org_channels(&self, org_id: OrgId) -> Result<Vec<Channel>> {
1410 let query = "
1411 SELECT *
1412 FROM channels
1413 WHERE
1414 channels.owner_is_user = false AND
1415 channels.owner_id = $1
1416 ";
1417 Ok(sqlx::query_as(query)
1418 .bind(org_id.0)
1419 .fetch_all(&self.pool)
1420 .await?)
1421 }
1422
1423 async fn get_accessible_channels(&self, user_id: UserId) -> Result<Vec<Channel>> {
1424 let query = "
1425 SELECT
1426 channels.*
1427 FROM
1428 channel_memberships, channels
1429 WHERE
1430 channel_memberships.user_id = $1 AND
1431 channel_memberships.channel_id = channels.id
1432 ";
1433 Ok(sqlx::query_as(query)
1434 .bind(user_id.0)
1435 .fetch_all(&self.pool)
1436 .await?)
1437 }
1438
1439 async fn can_user_access_channel(
1440 &self,
1441 user_id: UserId,
1442 channel_id: ChannelId,
1443 ) -> Result<bool> {
1444 let query = "
1445 SELECT id
1446 FROM channel_memberships
1447 WHERE user_id = $1 AND channel_id = $2
1448 LIMIT 1
1449 ";
1450 Ok(sqlx::query_scalar::<_, i32>(query)
1451 .bind(user_id.0)
1452 .bind(channel_id.0)
1453 .fetch_optional(&self.pool)
1454 .await
1455 .map(|e| e.is_some())?)
1456 }
1457
1458 #[cfg(any(test, feature = "seed-support"))]
1459 async fn add_channel_member(
1460 &self,
1461 channel_id: ChannelId,
1462 user_id: UserId,
1463 is_admin: bool,
1464 ) -> Result<()> {
1465 let query = "
1466 INSERT INTO channel_memberships (channel_id, user_id, admin)
1467 VALUES ($1, $2, $3)
1468 ON CONFLICT DO NOTHING
1469 ";
1470 Ok(sqlx::query(query)
1471 .bind(channel_id.0)
1472 .bind(user_id.0)
1473 .bind(is_admin)
1474 .execute(&self.pool)
1475 .await
1476 .map(drop)?)
1477 }
1478
1479 // messages
1480
1481 async fn create_channel_message(
1482 &self,
1483 channel_id: ChannelId,
1484 sender_id: UserId,
1485 body: &str,
1486 timestamp: OffsetDateTime,
1487 nonce: u128,
1488 ) -> Result<MessageId> {
1489 let query = "
1490 INSERT INTO channel_messages (channel_id, sender_id, body, sent_at, nonce)
1491 VALUES ($1, $2, $3, $4, $5)
1492 ON CONFLICT (nonce) DO UPDATE SET nonce = excluded.nonce
1493 RETURNING id
1494 ";
1495 Ok(sqlx::query_scalar(query)
1496 .bind(channel_id.0)
1497 .bind(sender_id.0)
1498 .bind(body)
1499 .bind(timestamp)
1500 .bind(Uuid::from_u128(nonce))
1501 .fetch_one(&self.pool)
1502 .await
1503 .map(MessageId)?)
1504 }
1505
1506 async fn get_channel_messages(
1507 &self,
1508 channel_id: ChannelId,
1509 count: usize,
1510 before_id: Option<MessageId>,
1511 ) -> Result<Vec<ChannelMessage>> {
1512 let query = r#"
1513 SELECT * FROM (
1514 SELECT
1515 id, channel_id, sender_id, body, sent_at AT TIME ZONE 'UTC' as sent_at, nonce
1516 FROM
1517 channel_messages
1518 WHERE
1519 channel_id = $1 AND
1520 id < $2
1521 ORDER BY id DESC
1522 LIMIT $3
1523 ) as recent_messages
1524 ORDER BY id ASC
1525 "#;
1526 Ok(sqlx::query_as(query)
1527 .bind(channel_id.0)
1528 .bind(before_id.unwrap_or(MessageId::MAX))
1529 .bind(count as i64)
1530 .fetch_all(&self.pool)
1531 .await?)
1532 }
1533
1534 #[cfg(test)]
1535 async fn teardown(&self, url: &str) {
1536 use util::ResultExt;
1537
1538 let query = "
1539 SELECT pg_terminate_backend(pg_stat_activity.pid)
1540 FROM pg_stat_activity
1541 WHERE pg_stat_activity.datname = current_database() AND pid <> pg_backend_pid();
1542 ";
1543 sqlx::query(query).execute(&self.pool).await.log_err();
1544 self.pool.close().await;
1545 <sqlx::Postgres as sqlx::migrate::MigrateDatabase>::drop_database(url)
1546 .await
1547 .log_err();
1548 }
1549
1550 #[cfg(test)]
1551 fn as_fake(&self) -> Option<&FakeDb> {
1552 None
1553 }
1554}
1555
1556macro_rules! id_type {
1557 ($name:ident) => {
1558 #[derive(
1559 Clone,
1560 Copy,
1561 Debug,
1562 Default,
1563 PartialEq,
1564 Eq,
1565 PartialOrd,
1566 Ord,
1567 Hash,
1568 sqlx::Type,
1569 Serialize,
1570 Deserialize,
1571 )]
1572 #[sqlx(transparent)]
1573 #[serde(transparent)]
1574 pub struct $name(pub i32);
1575
1576 impl $name {
1577 #[allow(unused)]
1578 pub const MAX: Self = Self(i32::MAX);
1579
1580 #[allow(unused)]
1581 pub fn from_proto(value: u64) -> Self {
1582 Self(value as i32)
1583 }
1584
1585 #[allow(unused)]
1586 pub fn to_proto(self) -> u64 {
1587 self.0 as u64
1588 }
1589 }
1590
1591 impl std::fmt::Display for $name {
1592 fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
1593 self.0.fmt(f)
1594 }
1595 }
1596 };
1597}
1598
1599id_type!(UserId);
1600#[derive(Clone, Debug, Default, FromRow, Serialize, PartialEq)]
1601pub struct User {
1602 pub id: UserId,
1603 pub github_login: String,
1604 pub github_user_id: Option<i32>,
1605 pub email_address: Option<String>,
1606 pub admin: bool,
1607 pub invite_code: Option<String>,
1608 pub invite_count: i32,
1609 pub connected_once: bool,
1610}
1611
1612id_type!(ProjectId);
1613#[derive(Clone, Debug, Default, FromRow, Serialize, PartialEq)]
1614pub struct Project {
1615 pub id: ProjectId,
1616 pub host_user_id: UserId,
1617 pub unregistered: bool,
1618}
1619
1620#[derive(Clone, Debug, PartialEq, Serialize)]
1621pub struct UserActivitySummary {
1622 pub id: UserId,
1623 pub github_login: String,
1624 pub project_activity: Vec<ProjectActivitySummary>,
1625}
1626
1627#[derive(Clone, Debug, PartialEq, Serialize)]
1628pub struct ProjectActivitySummary {
1629 pub id: ProjectId,
1630 pub duration: Duration,
1631 pub max_collaborators: usize,
1632}
1633
1634#[derive(Clone, Debug, PartialEq, Serialize)]
1635pub struct UserActivityPeriod {
1636 pub project_id: ProjectId,
1637 #[serde(with = "time::serde::iso8601")]
1638 pub start: OffsetDateTime,
1639 #[serde(with = "time::serde::iso8601")]
1640 pub end: OffsetDateTime,
1641 pub extensions: HashMap<String, usize>,
1642}
1643
1644id_type!(OrgId);
1645#[derive(FromRow)]
1646pub struct Org {
1647 pub id: OrgId,
1648 pub name: String,
1649 pub slug: String,
1650}
1651
1652id_type!(ChannelId);
1653#[derive(Clone, Debug, FromRow, Serialize)]
1654pub struct Channel {
1655 pub id: ChannelId,
1656 pub name: String,
1657 pub owner_id: i32,
1658 pub owner_is_user: bool,
1659}
1660
1661id_type!(MessageId);
1662#[derive(Clone, Debug, FromRow)]
1663pub struct ChannelMessage {
1664 pub id: MessageId,
1665 pub channel_id: ChannelId,
1666 pub sender_id: UserId,
1667 pub body: String,
1668 pub sent_at: OffsetDateTime,
1669 pub nonce: Uuid,
1670}
1671
1672#[derive(Clone, Debug, PartialEq, Eq)]
1673pub enum Contact {
1674 Accepted {
1675 user_id: UserId,
1676 should_notify: bool,
1677 },
1678 Outgoing {
1679 user_id: UserId,
1680 },
1681 Incoming {
1682 user_id: UserId,
1683 should_notify: bool,
1684 },
1685}
1686
1687impl Contact {
1688 pub fn user_id(&self) -> UserId {
1689 match self {
1690 Contact::Accepted { user_id, .. } => *user_id,
1691 Contact::Outgoing { user_id } => *user_id,
1692 Contact::Incoming { user_id, .. } => *user_id,
1693 }
1694 }
1695}
1696
1697#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord)]
1698pub struct IncomingContactRequest {
1699 pub requester_id: UserId,
1700 pub should_notify: bool,
1701}
1702
1703#[derive(Clone, Deserialize)]
1704pub struct Signup {
1705 pub email_address: String,
1706 pub platform_mac: bool,
1707 pub platform_windows: bool,
1708 pub platform_linux: bool,
1709 pub editor_features: Vec<String>,
1710 pub programming_languages: Vec<String>,
1711 pub device_id: Option<String>,
1712}
1713
1714#[derive(Clone, Debug, PartialEq, Deserialize, Serialize, FromRow)]
1715pub struct WaitlistSummary {
1716 #[sqlx(default)]
1717 pub count: i64,
1718 #[sqlx(default)]
1719 pub linux_count: i64,
1720 #[sqlx(default)]
1721 pub mac_count: i64,
1722 #[sqlx(default)]
1723 pub windows_count: i64,
1724 #[sqlx(default)]
1725 pub unknown_count: i64,
1726}
1727
1728#[derive(FromRow, PartialEq, Debug, Serialize, Deserialize)]
1729pub struct Invite {
1730 pub email_address: String,
1731 pub email_confirmation_code: String,
1732}
1733
1734#[derive(Debug, Serialize, Deserialize)]
1735pub struct NewUserParams {
1736 pub github_login: String,
1737 pub github_user_id: i32,
1738 pub invite_count: i32,
1739}
1740
1741#[derive(Debug)]
1742pub struct NewUserResult {
1743 pub user_id: UserId,
1744 pub metrics_id: String,
1745 pub inviting_user_id: Option<UserId>,
1746 pub signup_device_id: Option<String>,
1747}
1748
1749fn random_invite_code() -> String {
1750 nanoid::nanoid!(16)
1751}
1752
1753fn random_email_confirmation_code() -> String {
1754 nanoid::nanoid!(64)
1755}
1756
1757#[cfg(test)]
1758pub use test::*;
1759
1760#[cfg(test)]
1761mod test {
1762 use super::*;
1763 use anyhow::anyhow;
1764 use collections::BTreeMap;
1765 use gpui::executor::Background;
1766 use lazy_static::lazy_static;
1767 use parking_lot::Mutex;
1768 use rand::prelude::*;
1769 use sqlx::{
1770 migrate::{MigrateDatabase, Migrator},
1771 Postgres,
1772 };
1773 use std::{path::Path, sync::Arc};
1774 use util::post_inc;
1775
1776 pub struct FakeDb {
1777 background: Arc<Background>,
1778 pub users: Mutex<BTreeMap<UserId, User>>,
1779 pub projects: Mutex<BTreeMap<ProjectId, Project>>,
1780 pub worktree_extensions: Mutex<BTreeMap<(ProjectId, u64, String), u32>>,
1781 pub orgs: Mutex<BTreeMap<OrgId, Org>>,
1782 pub org_memberships: Mutex<BTreeMap<(OrgId, UserId), bool>>,
1783 pub channels: Mutex<BTreeMap<ChannelId, Channel>>,
1784 pub channel_memberships: Mutex<BTreeMap<(ChannelId, UserId), bool>>,
1785 pub channel_messages: Mutex<BTreeMap<MessageId, ChannelMessage>>,
1786 pub contacts: Mutex<Vec<FakeContact>>,
1787 next_channel_message_id: Mutex<i32>,
1788 next_user_id: Mutex<i32>,
1789 next_org_id: Mutex<i32>,
1790 next_channel_id: Mutex<i32>,
1791 next_project_id: Mutex<i32>,
1792 }
1793
1794 #[derive(Debug)]
1795 pub struct FakeContact {
1796 pub requester_id: UserId,
1797 pub responder_id: UserId,
1798 pub accepted: bool,
1799 pub should_notify: bool,
1800 }
1801
1802 impl FakeDb {
1803 pub fn new(background: Arc<Background>) -> Self {
1804 Self {
1805 background,
1806 users: Default::default(),
1807 next_user_id: Mutex::new(0),
1808 projects: Default::default(),
1809 worktree_extensions: Default::default(),
1810 next_project_id: Mutex::new(1),
1811 orgs: Default::default(),
1812 next_org_id: Mutex::new(1),
1813 org_memberships: Default::default(),
1814 channels: Default::default(),
1815 next_channel_id: Mutex::new(1),
1816 channel_memberships: Default::default(),
1817 channel_messages: Default::default(),
1818 next_channel_message_id: Mutex::new(1),
1819 contacts: Default::default(),
1820 }
1821 }
1822 }
1823
1824 #[async_trait]
1825 impl Db for FakeDb {
1826 async fn create_user(
1827 &self,
1828 email_address: &str,
1829 admin: bool,
1830 params: NewUserParams,
1831 ) -> Result<NewUserResult> {
1832 self.background.simulate_random_delay().await;
1833
1834 let mut users = self.users.lock();
1835 let user_id = if let Some(user) = users
1836 .values()
1837 .find(|user| user.github_login == params.github_login)
1838 {
1839 user.id
1840 } else {
1841 let id = post_inc(&mut *self.next_user_id.lock());
1842 let user_id = UserId(id);
1843 users.insert(
1844 user_id,
1845 User {
1846 id: user_id,
1847 github_login: params.github_login,
1848 github_user_id: Some(params.github_user_id),
1849 email_address: Some(email_address.to_string()),
1850 admin,
1851 invite_code: None,
1852 invite_count: 0,
1853 connected_once: false,
1854 },
1855 );
1856 user_id
1857 };
1858 Ok(NewUserResult {
1859 user_id,
1860 metrics_id: "the-metrics-id".to_string(),
1861 inviting_user_id: None,
1862 signup_device_id: None,
1863 })
1864 }
1865
1866 async fn get_all_users(&self, _page: u32, _limit: u32) -> Result<Vec<User>> {
1867 unimplemented!()
1868 }
1869
1870 async fn fuzzy_search_users(&self, _: &str, _: u32) -> Result<Vec<User>> {
1871 unimplemented!()
1872 }
1873
1874 async fn get_user_by_id(&self, id: UserId) -> Result<Option<User>> {
1875 self.background.simulate_random_delay().await;
1876 Ok(self.get_users_by_ids(vec![id]).await?.into_iter().next())
1877 }
1878
1879 async fn get_user_metrics_id(&self, _id: UserId) -> Result<String> {
1880 Ok("the-metrics-id".to_string())
1881 }
1882
1883 async fn get_users_by_ids(&self, ids: Vec<UserId>) -> Result<Vec<User>> {
1884 self.background.simulate_random_delay().await;
1885 let users = self.users.lock();
1886 Ok(ids.iter().filter_map(|id| users.get(id).cloned()).collect())
1887 }
1888
1889 async fn get_users_with_no_invites(&self, _: bool) -> Result<Vec<User>> {
1890 unimplemented!()
1891 }
1892
1893 async fn get_user_by_github_account(
1894 &self,
1895 github_login: &str,
1896 github_user_id: Option<i32>,
1897 ) -> Result<Option<User>> {
1898 self.background.simulate_random_delay().await;
1899 if let Some(github_user_id) = github_user_id {
1900 for user in self.users.lock().values_mut() {
1901 if user.github_user_id == Some(github_user_id) {
1902 user.github_login = github_login.into();
1903 return Ok(Some(user.clone()));
1904 }
1905 if user.github_login == github_login {
1906 user.github_user_id = Some(github_user_id);
1907 return Ok(Some(user.clone()));
1908 }
1909 }
1910 Ok(None)
1911 } else {
1912 Ok(self
1913 .users
1914 .lock()
1915 .values()
1916 .find(|user| user.github_login == github_login)
1917 .cloned())
1918 }
1919 }
1920
1921 async fn set_user_is_admin(&self, _id: UserId, _is_admin: bool) -> Result<()> {
1922 unimplemented!()
1923 }
1924
1925 async fn set_user_connected_once(&self, id: UserId, connected_once: bool) -> Result<()> {
1926 self.background.simulate_random_delay().await;
1927 let mut users = self.users.lock();
1928 let mut user = users
1929 .get_mut(&id)
1930 .ok_or_else(|| anyhow!("user not found"))?;
1931 user.connected_once = connected_once;
1932 Ok(())
1933 }
1934
1935 async fn destroy_user(&self, _id: UserId) -> Result<()> {
1936 unimplemented!()
1937 }
1938
1939 // signups
1940
1941 async fn create_signup(&self, _signup: Signup) -> Result<()> {
1942 unimplemented!()
1943 }
1944
1945 async fn get_waitlist_summary(&self) -> Result<WaitlistSummary> {
1946 unimplemented!()
1947 }
1948
1949 async fn get_unsent_invites(&self, _count: usize) -> Result<Vec<Invite>> {
1950 unimplemented!()
1951 }
1952
1953 async fn record_sent_invites(&self, _invites: &[Invite]) -> Result<()> {
1954 unimplemented!()
1955 }
1956
1957 async fn create_user_from_invite(
1958 &self,
1959 _invite: &Invite,
1960 _user: NewUserParams,
1961 ) -> Result<NewUserResult> {
1962 unimplemented!()
1963 }
1964
1965 // invite codes
1966
1967 async fn set_invite_count_for_user(&self, _id: UserId, _count: u32) -> Result<()> {
1968 unimplemented!()
1969 }
1970
1971 async fn get_invite_code_for_user(&self, _id: UserId) -> Result<Option<(String, u32)>> {
1972 self.background.simulate_random_delay().await;
1973 Ok(None)
1974 }
1975
1976 async fn get_user_for_invite_code(&self, _code: &str) -> Result<User> {
1977 unimplemented!()
1978 }
1979
1980 async fn create_invite_from_code(
1981 &self,
1982 _code: &str,
1983 _email_address: &str,
1984 _device_id: Option<&str>,
1985 ) -> Result<Invite> {
1986 unimplemented!()
1987 }
1988
1989 // projects
1990
1991 async fn register_project(&self, host_user_id: UserId) -> Result<ProjectId> {
1992 self.background.simulate_random_delay().await;
1993 if !self.users.lock().contains_key(&host_user_id) {
1994 Err(anyhow!("no such user"))?;
1995 }
1996
1997 let project_id = ProjectId(post_inc(&mut *self.next_project_id.lock()));
1998 self.projects.lock().insert(
1999 project_id,
2000 Project {
2001 id: project_id,
2002 host_user_id,
2003 unregistered: false,
2004 },
2005 );
2006 Ok(project_id)
2007 }
2008
2009 async fn unregister_project(&self, project_id: ProjectId) -> Result<()> {
2010 self.background.simulate_random_delay().await;
2011 self.projects
2012 .lock()
2013 .get_mut(&project_id)
2014 .ok_or_else(|| anyhow!("no such project"))?
2015 .unregistered = true;
2016 Ok(())
2017 }
2018
2019 async fn update_worktree_extensions(
2020 &self,
2021 project_id: ProjectId,
2022 worktree_id: u64,
2023 extensions: HashMap<String, u32>,
2024 ) -> Result<()> {
2025 self.background.simulate_random_delay().await;
2026 if !self.projects.lock().contains_key(&project_id) {
2027 Err(anyhow!("no such project"))?;
2028 }
2029
2030 for (extension, count) in extensions {
2031 self.worktree_extensions
2032 .lock()
2033 .insert((project_id, worktree_id, extension), count);
2034 }
2035
2036 Ok(())
2037 }
2038
2039 async fn get_project_extensions(
2040 &self,
2041 _project_id: ProjectId,
2042 ) -> Result<HashMap<u64, HashMap<String, usize>>> {
2043 unimplemented!()
2044 }
2045
2046 async fn record_user_activity(
2047 &self,
2048 _time_period: Range<OffsetDateTime>,
2049 _active_projects: &[(UserId, ProjectId)],
2050 ) -> Result<()> {
2051 unimplemented!()
2052 }
2053
2054 async fn get_active_user_count(
2055 &self,
2056 _time_period: Range<OffsetDateTime>,
2057 _min_duration: Duration,
2058 _only_collaborative: bool,
2059 ) -> Result<usize> {
2060 unimplemented!()
2061 }
2062
2063 async fn get_top_users_activity_summary(
2064 &self,
2065 _time_period: Range<OffsetDateTime>,
2066 _limit: usize,
2067 ) -> Result<Vec<UserActivitySummary>> {
2068 unimplemented!()
2069 }
2070
2071 async fn get_user_activity_timeline(
2072 &self,
2073 _time_period: Range<OffsetDateTime>,
2074 _user_id: UserId,
2075 ) -> Result<Vec<UserActivityPeriod>> {
2076 unimplemented!()
2077 }
2078
2079 // contacts
2080
2081 async fn get_contacts(&self, id: UserId) -> Result<Vec<Contact>> {
2082 self.background.simulate_random_delay().await;
2083 let mut contacts = Vec::new();
2084
2085 for contact in self.contacts.lock().iter() {
2086 if contact.requester_id == id {
2087 if contact.accepted {
2088 contacts.push(Contact::Accepted {
2089 user_id: contact.responder_id,
2090 should_notify: contact.should_notify,
2091 });
2092 } else {
2093 contacts.push(Contact::Outgoing {
2094 user_id: contact.responder_id,
2095 });
2096 }
2097 } else if contact.responder_id == id {
2098 if contact.accepted {
2099 contacts.push(Contact::Accepted {
2100 user_id: contact.requester_id,
2101 should_notify: false,
2102 });
2103 } else {
2104 contacts.push(Contact::Incoming {
2105 user_id: contact.requester_id,
2106 should_notify: contact.should_notify,
2107 });
2108 }
2109 }
2110 }
2111
2112 contacts.sort_unstable_by_key(|contact| contact.user_id());
2113 Ok(contacts)
2114 }
2115
2116 async fn has_contact(&self, user_id_a: UserId, user_id_b: UserId) -> Result<bool> {
2117 self.background.simulate_random_delay().await;
2118 Ok(self.contacts.lock().iter().any(|contact| {
2119 contact.accepted
2120 && ((contact.requester_id == user_id_a && contact.responder_id == user_id_b)
2121 || (contact.requester_id == user_id_b && contact.responder_id == user_id_a))
2122 }))
2123 }
2124
2125 async fn send_contact_request(
2126 &self,
2127 requester_id: UserId,
2128 responder_id: UserId,
2129 ) -> Result<()> {
2130 self.background.simulate_random_delay().await;
2131 let mut contacts = self.contacts.lock();
2132 for contact in contacts.iter_mut() {
2133 if contact.requester_id == requester_id && contact.responder_id == responder_id {
2134 if contact.accepted {
2135 Err(anyhow!("contact already exists"))?;
2136 } else {
2137 Err(anyhow!("contact already requested"))?;
2138 }
2139 }
2140 if contact.responder_id == requester_id && contact.requester_id == responder_id {
2141 if contact.accepted {
2142 Err(anyhow!("contact already exists"))?;
2143 } else {
2144 contact.accepted = true;
2145 contact.should_notify = false;
2146 return Ok(());
2147 }
2148 }
2149 }
2150 contacts.push(FakeContact {
2151 requester_id,
2152 responder_id,
2153 accepted: false,
2154 should_notify: true,
2155 });
2156 Ok(())
2157 }
2158
2159 async fn remove_contact(&self, requester_id: UserId, responder_id: UserId) -> Result<()> {
2160 self.background.simulate_random_delay().await;
2161 self.contacts.lock().retain(|contact| {
2162 !(contact.requester_id == requester_id && contact.responder_id == responder_id)
2163 });
2164 Ok(())
2165 }
2166
2167 async fn dismiss_contact_notification(
2168 &self,
2169 user_id: UserId,
2170 contact_user_id: UserId,
2171 ) -> Result<()> {
2172 self.background.simulate_random_delay().await;
2173 let mut contacts = self.contacts.lock();
2174 for contact in contacts.iter_mut() {
2175 if contact.requester_id == contact_user_id
2176 && contact.responder_id == user_id
2177 && !contact.accepted
2178 {
2179 contact.should_notify = false;
2180 return Ok(());
2181 }
2182 if contact.requester_id == user_id
2183 && contact.responder_id == contact_user_id
2184 && contact.accepted
2185 {
2186 contact.should_notify = false;
2187 return Ok(());
2188 }
2189 }
2190 Err(anyhow!("no such notification"))?
2191 }
2192
2193 async fn respond_to_contact_request(
2194 &self,
2195 responder_id: UserId,
2196 requester_id: UserId,
2197 accept: bool,
2198 ) -> Result<()> {
2199 self.background.simulate_random_delay().await;
2200 let mut contacts = self.contacts.lock();
2201 for (ix, contact) in contacts.iter_mut().enumerate() {
2202 if contact.requester_id == requester_id && contact.responder_id == responder_id {
2203 if contact.accepted {
2204 Err(anyhow!("contact already confirmed"))?;
2205 }
2206 if accept {
2207 contact.accepted = true;
2208 contact.should_notify = true;
2209 } else {
2210 contacts.remove(ix);
2211 }
2212 return Ok(());
2213 }
2214 }
2215 Err(anyhow!("no such contact request"))?
2216 }
2217
2218 async fn create_access_token_hash(
2219 &self,
2220 _user_id: UserId,
2221 _access_token_hash: &str,
2222 _max_access_token_count: usize,
2223 ) -> Result<()> {
2224 unimplemented!()
2225 }
2226
2227 async fn get_access_token_hashes(&self, _user_id: UserId) -> Result<Vec<String>> {
2228 unimplemented!()
2229 }
2230
2231 async fn find_org_by_slug(&self, _slug: &str) -> Result<Option<Org>> {
2232 unimplemented!()
2233 }
2234
2235 async fn create_org(&self, name: &str, slug: &str) -> Result<OrgId> {
2236 self.background.simulate_random_delay().await;
2237 let mut orgs = self.orgs.lock();
2238 if orgs.values().any(|org| org.slug == slug) {
2239 Err(anyhow!("org already exists"))?
2240 } else {
2241 let org_id = OrgId(post_inc(&mut *self.next_org_id.lock()));
2242 orgs.insert(
2243 org_id,
2244 Org {
2245 id: org_id,
2246 name: name.to_string(),
2247 slug: slug.to_string(),
2248 },
2249 );
2250 Ok(org_id)
2251 }
2252 }
2253
2254 async fn add_org_member(
2255 &self,
2256 org_id: OrgId,
2257 user_id: UserId,
2258 is_admin: bool,
2259 ) -> Result<()> {
2260 self.background.simulate_random_delay().await;
2261 if !self.orgs.lock().contains_key(&org_id) {
2262 Err(anyhow!("org does not exist"))?;
2263 }
2264 if !self.users.lock().contains_key(&user_id) {
2265 Err(anyhow!("user does not exist"))?;
2266 }
2267
2268 self.org_memberships
2269 .lock()
2270 .entry((org_id, user_id))
2271 .or_insert(is_admin);
2272 Ok(())
2273 }
2274
2275 async fn create_org_channel(&self, org_id: OrgId, name: &str) -> Result<ChannelId> {
2276 self.background.simulate_random_delay().await;
2277 if !self.orgs.lock().contains_key(&org_id) {
2278 Err(anyhow!("org does not exist"))?;
2279 }
2280
2281 let mut channels = self.channels.lock();
2282 let channel_id = ChannelId(post_inc(&mut *self.next_channel_id.lock()));
2283 channels.insert(
2284 channel_id,
2285 Channel {
2286 id: channel_id,
2287 name: name.to_string(),
2288 owner_id: org_id.0,
2289 owner_is_user: false,
2290 },
2291 );
2292 Ok(channel_id)
2293 }
2294
2295 async fn get_org_channels(&self, org_id: OrgId) -> Result<Vec<Channel>> {
2296 self.background.simulate_random_delay().await;
2297 Ok(self
2298 .channels
2299 .lock()
2300 .values()
2301 .filter(|channel| !channel.owner_is_user && channel.owner_id == org_id.0)
2302 .cloned()
2303 .collect())
2304 }
2305
2306 async fn get_accessible_channels(&self, user_id: UserId) -> Result<Vec<Channel>> {
2307 self.background.simulate_random_delay().await;
2308 let channels = self.channels.lock();
2309 let memberships = self.channel_memberships.lock();
2310 Ok(channels
2311 .values()
2312 .filter(|channel| memberships.contains_key(&(channel.id, user_id)))
2313 .cloned()
2314 .collect())
2315 }
2316
2317 async fn can_user_access_channel(
2318 &self,
2319 user_id: UserId,
2320 channel_id: ChannelId,
2321 ) -> Result<bool> {
2322 self.background.simulate_random_delay().await;
2323 Ok(self
2324 .channel_memberships
2325 .lock()
2326 .contains_key(&(channel_id, user_id)))
2327 }
2328
2329 async fn add_channel_member(
2330 &self,
2331 channel_id: ChannelId,
2332 user_id: UserId,
2333 is_admin: bool,
2334 ) -> Result<()> {
2335 self.background.simulate_random_delay().await;
2336 if !self.channels.lock().contains_key(&channel_id) {
2337 Err(anyhow!("channel does not exist"))?;
2338 }
2339 if !self.users.lock().contains_key(&user_id) {
2340 Err(anyhow!("user does not exist"))?;
2341 }
2342
2343 self.channel_memberships
2344 .lock()
2345 .entry((channel_id, user_id))
2346 .or_insert(is_admin);
2347 Ok(())
2348 }
2349
2350 async fn create_channel_message(
2351 &self,
2352 channel_id: ChannelId,
2353 sender_id: UserId,
2354 body: &str,
2355 timestamp: OffsetDateTime,
2356 nonce: u128,
2357 ) -> Result<MessageId> {
2358 self.background.simulate_random_delay().await;
2359 if !self.channels.lock().contains_key(&channel_id) {
2360 Err(anyhow!("channel does not exist"))?;
2361 }
2362 if !self.users.lock().contains_key(&sender_id) {
2363 Err(anyhow!("user does not exist"))?;
2364 }
2365
2366 let mut messages = self.channel_messages.lock();
2367 if let Some(message) = messages
2368 .values()
2369 .find(|message| message.nonce.as_u128() == nonce)
2370 {
2371 Ok(message.id)
2372 } else {
2373 let message_id = MessageId(post_inc(&mut *self.next_channel_message_id.lock()));
2374 messages.insert(
2375 message_id,
2376 ChannelMessage {
2377 id: message_id,
2378 channel_id,
2379 sender_id,
2380 body: body.to_string(),
2381 sent_at: timestamp,
2382 nonce: Uuid::from_u128(nonce),
2383 },
2384 );
2385 Ok(message_id)
2386 }
2387 }
2388
2389 async fn get_channel_messages(
2390 &self,
2391 channel_id: ChannelId,
2392 count: usize,
2393 before_id: Option<MessageId>,
2394 ) -> Result<Vec<ChannelMessage>> {
2395 self.background.simulate_random_delay().await;
2396 let mut messages = self
2397 .channel_messages
2398 .lock()
2399 .values()
2400 .rev()
2401 .filter(|message| {
2402 message.channel_id == channel_id
2403 && message.id < before_id.unwrap_or(MessageId::MAX)
2404 })
2405 .take(count)
2406 .cloned()
2407 .collect::<Vec<_>>();
2408 messages.sort_unstable_by_key(|message| message.id);
2409 Ok(messages)
2410 }
2411
2412 async fn teardown(&self, _: &str) {}
2413
2414 #[cfg(test)]
2415 fn as_fake(&self) -> Option<&FakeDb> {
2416 Some(self)
2417 }
2418 }
2419
2420 pub struct TestDb {
2421 pub db: Option<Arc<dyn Db>>,
2422 pub url: String,
2423 }
2424
2425 impl TestDb {
2426 #[allow(clippy::await_holding_lock)]
2427 pub async fn postgres() -> Self {
2428 lazy_static! {
2429 static ref LOCK: Mutex<()> = Mutex::new(());
2430 }
2431
2432 let _guard = LOCK.lock();
2433 let mut rng = StdRng::from_entropy();
2434 let name = format!("zed-test-{}", rng.gen::<u128>());
2435 let url = format!("postgres://postgres@localhost/{}", name);
2436 let migrations_path = Path::new(concat!(env!("CARGO_MANIFEST_DIR"), "/migrations"));
2437 Postgres::create_database(&url)
2438 .await
2439 .expect("failed to create test db");
2440 let db = PostgresDb::new(&url, 5).await.unwrap();
2441 let migrator = Migrator::new(migrations_path).await.unwrap();
2442 migrator.run(&db.pool).await.unwrap();
2443 Self {
2444 db: Some(Arc::new(db)),
2445 url,
2446 }
2447 }
2448
2449 pub fn fake(background: Arc<Background>) -> Self {
2450 Self {
2451 db: Some(Arc::new(FakeDb::new(background))),
2452 url: Default::default(),
2453 }
2454 }
2455
2456 pub fn db(&self) -> &Arc<dyn Db> {
2457 self.db.as_ref().unwrap()
2458 }
2459 }
2460
2461 impl Drop for TestDb {
2462 fn drop(&mut self) {
2463 if let Some(db) = self.db.take() {
2464 futures::executor::block_on(db.teardown(&self.url));
2465 }
2466 }
2467 }
2468}