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