1mod ids;
2mod queries;
3mod tables;
4#[cfg(test)]
5pub mod tests;
6
7use crate::{executor::Executor, Error, Result};
8use anyhow::anyhow;
9use collections::{BTreeMap, HashMap, HashSet};
10use dashmap::DashMap;
11use futures::StreamExt;
12use rand::{prelude::StdRng, Rng, SeedableRng};
13use rpc::{
14 proto::{self},
15 ConnectionId,
16};
17use sea_orm::{
18 entity::prelude::*,
19 sea_query::{Alias, Expr, OnConflict},
20 ActiveValue, Condition, ConnectionTrait, DatabaseConnection, DatabaseTransaction, DbErr,
21 FromQueryResult, IntoActiveModel, IsolationLevel, JoinType, QueryOrder, QuerySelect, Statement,
22 TransactionTrait,
23};
24use serde::{ser::Error as _, Deserialize, Serialize, Serializer};
25use sqlx::{
26 migrate::{Migrate, Migration, MigrationSource},
27 Connection,
28};
29use std::{
30 fmt::Write as _,
31 future::Future,
32 marker::PhantomData,
33 ops::{Deref, DerefMut},
34 path::Path,
35 rc::Rc,
36 sync::Arc,
37 time::Duration,
38};
39use time::{format_description::well_known::iso8601, PrimitiveDateTime};
40use tokio::sync::{Mutex, OwnedMutexGuard};
41
42#[cfg(test)]
43pub use tests::TestDb;
44
45pub use ids::*;
46pub use queries::contributors::ContributorSelector;
47pub use sea_orm::ConnectOptions;
48pub use tables::user::Model as User;
49pub use tables::*;
50
51/// Database gives you a handle that lets you access the database.
52/// It handles pooling internally.
53pub struct Database {
54 options: ConnectOptions,
55 pool: DatabaseConnection,
56 rooms: DashMap<RoomId, Arc<Mutex<()>>>,
57 rng: Mutex<StdRng>,
58 executor: Executor,
59 notification_kinds_by_id: HashMap<NotificationKindId, &'static str>,
60 notification_kinds_by_name: HashMap<String, NotificationKindId>,
61 #[cfg(test)]
62 runtime: Option<tokio::runtime::Runtime>,
63}
64
65// The `Database` type has so many methods that its impl blocks are split into
66// separate files in the `queries` folder.
67impl Database {
68 /// Connects to the database with the given options
69 pub async fn new(options: ConnectOptions, executor: Executor) -> Result<Self> {
70 sqlx::any::install_default_drivers();
71 Ok(Self {
72 options: options.clone(),
73 pool: sea_orm::Database::connect(options).await?,
74 rooms: DashMap::with_capacity(16384),
75 rng: Mutex::new(StdRng::seed_from_u64(0)),
76 notification_kinds_by_id: HashMap::default(),
77 notification_kinds_by_name: HashMap::default(),
78 executor,
79 #[cfg(test)]
80 runtime: None,
81 })
82 }
83
84 #[cfg(test)]
85 pub fn reset(&self) {
86 self.rooms.clear();
87 }
88
89 /// Runs the database migrations.
90 pub async fn migrate(
91 &self,
92 migrations_path: &Path,
93 ignore_checksum_mismatch: bool,
94 ) -> anyhow::Result<Vec<(Migration, Duration)>> {
95 let migrations = MigrationSource::resolve(migrations_path)
96 .await
97 .map_err(|err| anyhow!("failed to load migrations: {err:?}"))?;
98
99 let mut connection = sqlx::AnyConnection::connect(self.options.get_url()).await?;
100
101 connection.ensure_migrations_table().await?;
102 let applied_migrations: HashMap<_, _> = connection
103 .list_applied_migrations()
104 .await?
105 .into_iter()
106 .map(|m| (m.version, m))
107 .collect();
108
109 let mut new_migrations = Vec::new();
110 for migration in migrations {
111 match applied_migrations.get(&migration.version) {
112 Some(applied_migration) => {
113 if migration.checksum != applied_migration.checksum && !ignore_checksum_mismatch
114 {
115 Err(anyhow!(
116 "checksum mismatch for applied migration {}",
117 migration.description
118 ))?;
119 }
120 }
121 None => {
122 let elapsed = connection.apply(&migration).await?;
123 new_migrations.push((migration, elapsed));
124 }
125 }
126 }
127
128 Ok(new_migrations)
129 }
130
131 /// Initializes static data that resides in the database by upserting it.
132 pub async fn initialize_static_data(&mut self) -> Result<()> {
133 self.initialize_notification_kinds().await?;
134 Ok(())
135 }
136
137 /// Transaction runs things in a transaction. If you want to call other methods
138 /// and pass the transaction around you need to reborrow the transaction at each
139 /// call site with: `&*tx`.
140 pub async fn transaction<F, Fut, T>(&self, f: F) -> Result<T>
141 where
142 F: Send + Fn(TransactionHandle) -> Fut,
143 Fut: Send + Future<Output = Result<T>>,
144 {
145 let body = async {
146 let mut i = 0;
147 loop {
148 let (tx, result) = self.with_transaction(&f).await?;
149 match result {
150 Ok(result) => match tx.commit().await.map_err(Into::into) {
151 Ok(()) => return Ok(result),
152 Err(error) => {
153 if !self.retry_on_serialization_error(&error, i).await {
154 return Err(error);
155 }
156 }
157 },
158 Err(error) => {
159 tx.rollback().await?;
160 if !self.retry_on_serialization_error(&error, i).await {
161 return Err(error);
162 }
163 }
164 }
165 i += 1;
166 }
167 };
168
169 self.run(body).await
170 }
171
172 pub async fn weak_transaction<F, Fut, T>(&self, f: F) -> Result<T>
173 where
174 F: Send + Fn(TransactionHandle) -> Fut,
175 Fut: Send + Future<Output = Result<T>>,
176 {
177 let body = async {
178 let (tx, result) = self.with_weak_transaction(&f).await?;
179 match result {
180 Ok(result) => match tx.commit().await.map_err(Into::into) {
181 Ok(()) => return Ok(result),
182 Err(error) => {
183 return Err(error);
184 }
185 },
186 Err(error) => {
187 tx.rollback().await?;
188 return Err(error);
189 }
190 }
191 };
192
193 self.run(body).await
194 }
195
196 /// The same as room_transaction, but if you need to only optionally return a Room.
197 async fn optional_room_transaction<F, Fut, T>(&self, f: F) -> Result<Option<RoomGuard<T>>>
198 where
199 F: Send + Fn(TransactionHandle) -> Fut,
200 Fut: Send + Future<Output = Result<Option<(RoomId, T)>>>,
201 {
202 let body = async {
203 let mut i = 0;
204 loop {
205 let (tx, result) = self.with_transaction(&f).await?;
206 match result {
207 Ok(Some((room_id, data))) => {
208 let lock = self.rooms.entry(room_id).or_default().clone();
209 let _guard = lock.lock_owned().await;
210 match tx.commit().await.map_err(Into::into) {
211 Ok(()) => {
212 return Ok(Some(RoomGuard {
213 data,
214 _guard,
215 _not_send: PhantomData,
216 }));
217 }
218 Err(error) => {
219 if !self.retry_on_serialization_error(&error, i).await {
220 return Err(error);
221 }
222 }
223 }
224 }
225 Ok(None) => match tx.commit().await.map_err(Into::into) {
226 Ok(()) => return Ok(None),
227 Err(error) => {
228 if !self.retry_on_serialization_error(&error, i).await {
229 return Err(error);
230 }
231 }
232 },
233 Err(error) => {
234 tx.rollback().await?;
235 if !self.retry_on_serialization_error(&error, i).await {
236 return Err(error);
237 }
238 }
239 }
240 i += 1;
241 }
242 };
243
244 self.run(body).await
245 }
246
247 /// room_transaction runs the block in a transaction. It returns a RoomGuard, that keeps
248 /// the database locked until it is dropped. This ensures that updates sent to clients are
249 /// properly serialized with respect to database changes.
250 async fn room_transaction<F, Fut, T>(&self, room_id: RoomId, f: F) -> Result<RoomGuard<T>>
251 where
252 F: Send + Fn(TransactionHandle) -> Fut,
253 Fut: Send + Future<Output = Result<T>>,
254 {
255 let body = async {
256 let mut i = 0;
257 loop {
258 let lock = self.rooms.entry(room_id).or_default().clone();
259 let _guard = lock.lock_owned().await;
260 let (tx, result) = self.with_transaction(&f).await?;
261 match result {
262 Ok(data) => match tx.commit().await.map_err(Into::into) {
263 Ok(()) => {
264 return Ok(RoomGuard {
265 data,
266 _guard,
267 _not_send: PhantomData,
268 });
269 }
270 Err(error) => {
271 if !self.retry_on_serialization_error(&error, i).await {
272 return Err(error);
273 }
274 }
275 },
276 Err(error) => {
277 tx.rollback().await?;
278 if !self.retry_on_serialization_error(&error, i).await {
279 return Err(error);
280 }
281 }
282 }
283 i += 1;
284 }
285 };
286
287 self.run(body).await
288 }
289
290 async fn with_transaction<F, Fut, T>(&self, f: &F) -> Result<(DatabaseTransaction, Result<T>)>
291 where
292 F: Send + Fn(TransactionHandle) -> Fut,
293 Fut: Send + Future<Output = Result<T>>,
294 {
295 let tx = self
296 .pool
297 .begin_with_config(Some(IsolationLevel::Serializable), None)
298 .await?;
299
300 let mut tx = Arc::new(Some(tx));
301 let result = f(TransactionHandle(tx.clone())).await;
302 let Some(tx) = Arc::get_mut(&mut tx).and_then(|tx| tx.take()) else {
303 return Err(anyhow!(
304 "couldn't complete transaction because it's still in use"
305 ))?;
306 };
307
308 Ok((tx, result))
309 }
310
311 async fn with_weak_transaction<F, Fut, T>(
312 &self,
313 f: &F,
314 ) -> Result<(DatabaseTransaction, Result<T>)>
315 where
316 F: Send + Fn(TransactionHandle) -> Fut,
317 Fut: Send + Future<Output = Result<T>>,
318 {
319 let tx = self
320 .pool
321 .begin_with_config(Some(IsolationLevel::ReadCommitted), None)
322 .await?;
323
324 let mut tx = Arc::new(Some(tx));
325 let result = f(TransactionHandle(tx.clone())).await;
326 let Some(tx) = Arc::get_mut(&mut tx).and_then(|tx| tx.take()) else {
327 return Err(anyhow!(
328 "couldn't complete transaction because it's still in use"
329 ))?;
330 };
331
332 Ok((tx, result))
333 }
334
335 async fn run<F, T>(&self, future: F) -> Result<T>
336 where
337 F: Future<Output = Result<T>>,
338 {
339 #[cfg(test)]
340 {
341 if let Executor::Deterministic(executor) = &self.executor {
342 executor.simulate_random_delay().await;
343 }
344
345 self.runtime.as_ref().unwrap().block_on(future)
346 }
347
348 #[cfg(not(test))]
349 {
350 future.await
351 }
352 }
353
354 async fn retry_on_serialization_error(&self, error: &Error, prev_attempt_count: usize) -> bool {
355 // If the error is due to a failure to serialize concurrent transactions, then retry
356 // this transaction after a delay. With each subsequent retry, double the delay duration.
357 // Also vary the delay randomly in order to ensure different database connections retry
358 // at different times.
359 const SLEEPS: [f32; 10] = [10., 20., 40., 80., 160., 320., 640., 1280., 2560., 5120.];
360 if is_serialization_error(error) && prev_attempt_count < SLEEPS.len() {
361 let base_delay = SLEEPS[prev_attempt_count];
362 let randomized_delay = base_delay * self.rng.lock().await.gen_range(0.5..=2.0);
363 log::info!(
364 "retrying transaction after serialization error. delay: {} ms.",
365 randomized_delay
366 );
367 self.executor
368 .sleep(Duration::from_millis(randomized_delay as u64))
369 .await;
370 true
371 } else {
372 false
373 }
374 }
375}
376
377fn is_serialization_error(error: &Error) -> bool {
378 const SERIALIZATION_FAILURE_CODE: &str = "40001";
379 match error {
380 Error::Database(
381 DbErr::Exec(sea_orm::RuntimeErr::SqlxError(error))
382 | DbErr::Query(sea_orm::RuntimeErr::SqlxError(error)),
383 ) if error
384 .as_database_error()
385 .and_then(|error| error.code())
386 .as_deref()
387 == Some(SERIALIZATION_FAILURE_CODE) =>
388 {
389 true
390 }
391 _ => false,
392 }
393}
394
395/// A handle to a [`DatabaseTransaction`].
396pub struct TransactionHandle(Arc<Option<DatabaseTransaction>>);
397
398impl Deref for TransactionHandle {
399 type Target = DatabaseTransaction;
400
401 fn deref(&self) -> &Self::Target {
402 self.0.as_ref().as_ref().unwrap()
403 }
404}
405
406/// [`RoomGuard`] keeps a database transaction alive until it is dropped.
407/// so that updates to rooms are serialized.
408pub struct RoomGuard<T> {
409 data: T,
410 _guard: OwnedMutexGuard<()>,
411 _not_send: PhantomData<Rc<()>>,
412}
413
414impl<T> Deref for RoomGuard<T> {
415 type Target = T;
416
417 fn deref(&self) -> &T {
418 &self.data
419 }
420}
421
422impl<T> DerefMut for RoomGuard<T> {
423 fn deref_mut(&mut self) -> &mut T {
424 &mut self.data
425 }
426}
427
428impl<T> RoomGuard<T> {
429 /// Returns the inner value of the guard.
430 pub fn into_inner(self) -> T {
431 self.data
432 }
433}
434
435#[derive(Clone, Debug, PartialEq, Eq)]
436pub enum Contact {
437 Accepted { user_id: UserId, busy: bool },
438 Outgoing { user_id: UserId },
439 Incoming { user_id: UserId },
440}
441
442impl Contact {
443 pub fn user_id(&self) -> UserId {
444 match self {
445 Contact::Accepted { user_id, .. } => *user_id,
446 Contact::Outgoing { user_id } => *user_id,
447 Contact::Incoming { user_id, .. } => *user_id,
448 }
449 }
450}
451
452pub type NotificationBatch = Vec<(UserId, proto::Notification)>;
453
454pub struct CreatedChannelMessage {
455 pub message_id: MessageId,
456 pub participant_connection_ids: Vec<ConnectionId>,
457 pub channel_members: Vec<UserId>,
458 pub notifications: NotificationBatch,
459}
460
461pub struct UpdatedChannelMessage {
462 pub message_id: MessageId,
463 pub participant_connection_ids: Vec<ConnectionId>,
464 pub notifications: NotificationBatch,
465 pub reply_to_message_id: Option<MessageId>,
466 pub timestamp: PrimitiveDateTime,
467}
468
469#[derive(Clone, Debug, PartialEq, Eq, FromQueryResult, Serialize, Deserialize)]
470pub struct Invite {
471 pub email_address: String,
472 pub email_confirmation_code: String,
473}
474
475#[derive(Clone, Debug, Deserialize)]
476pub struct NewSignup {
477 pub email_address: String,
478 pub platform_mac: bool,
479 pub platform_windows: bool,
480 pub platform_linux: bool,
481 pub editor_features: Vec<String>,
482 pub programming_languages: Vec<String>,
483 pub device_id: Option<String>,
484 pub added_to_mailing_list: bool,
485 pub created_at: Option<DateTime>,
486}
487
488#[derive(Clone, Debug, PartialEq, Deserialize, Serialize, FromQueryResult)]
489pub struct WaitlistSummary {
490 pub count: i64,
491 pub linux_count: i64,
492 pub mac_count: i64,
493 pub windows_count: i64,
494 pub unknown_count: i64,
495}
496
497/// The parameters to create a new user.
498#[derive(Debug, Serialize, Deserialize)]
499pub struct NewUserParams {
500 pub github_login: String,
501 pub github_user_id: i32,
502}
503
504/// The result of creating a new user.
505#[derive(Debug)]
506pub struct NewUserResult {
507 pub user_id: UserId,
508 pub metrics_id: String,
509 pub inviting_user_id: Option<UserId>,
510 pub signup_device_id: Option<String>,
511}
512
513/// The result of updating a channel membership.
514#[derive(Debug)]
515pub struct MembershipUpdated {
516 pub channel_id: ChannelId,
517 pub new_channels: ChannelsForUser,
518 pub removed_channels: Vec<ChannelId>,
519}
520
521/// The result of setting a member's role.
522#[derive(Debug)]
523pub enum SetMemberRoleResult {
524 InviteUpdated(Channel),
525 MembershipUpdated(MembershipUpdated),
526}
527
528/// The result of inviting a member to a channel.
529#[derive(Debug)]
530pub struct InviteMemberResult {
531 pub channel: Channel,
532 pub notifications: NotificationBatch,
533}
534
535#[derive(Debug)]
536pub struct RespondToChannelInvite {
537 pub membership_update: Option<MembershipUpdated>,
538 pub notifications: NotificationBatch,
539}
540
541#[derive(Debug)]
542pub struct RemoveChannelMemberResult {
543 pub membership_update: MembershipUpdated,
544 pub notification_id: Option<NotificationId>,
545}
546
547#[derive(Debug, PartialEq, Eq, Hash)]
548pub struct Channel {
549 pub id: ChannelId,
550 pub name: String,
551 pub visibility: ChannelVisibility,
552 /// parent_path is the channel ids from the root to this one (not including this one)
553 pub parent_path: Vec<ChannelId>,
554}
555
556impl Channel {
557 pub fn from_model(value: channel::Model) -> Self {
558 Channel {
559 id: value.id,
560 visibility: value.visibility,
561 name: value.clone().name,
562 parent_path: value.ancestors().collect(),
563 }
564 }
565
566 pub fn to_proto(&self) -> proto::Channel {
567 proto::Channel {
568 id: self.id.to_proto(),
569 name: self.name.clone(),
570 visibility: self.visibility.into(),
571 parent_path: self.parent_path.iter().map(|c| c.to_proto()).collect(),
572 }
573 }
574}
575
576#[derive(Debug, PartialEq, Eq, Hash)]
577pub struct ChannelMember {
578 pub role: ChannelRole,
579 pub user_id: UserId,
580 pub kind: proto::channel_member::Kind,
581}
582
583impl ChannelMember {
584 pub fn to_proto(&self) -> proto::ChannelMember {
585 proto::ChannelMember {
586 role: self.role.into(),
587 user_id: self.user_id.to_proto(),
588 kind: self.kind.into(),
589 }
590 }
591}
592
593#[derive(Debug, PartialEq)]
594pub struct ChannelsForUser {
595 pub channels: Vec<Channel>,
596 pub channel_memberships: Vec<channel_member::Model>,
597 pub channel_participants: HashMap<ChannelId, Vec<UserId>>,
598 pub hosted_projects: Vec<proto::HostedProject>,
599
600 pub observed_buffer_versions: Vec<proto::ChannelBufferVersion>,
601 pub observed_channel_messages: Vec<proto::ChannelMessageId>,
602 pub latest_buffer_versions: Vec<proto::ChannelBufferVersion>,
603 pub latest_channel_messages: Vec<proto::ChannelMessageId>,
604}
605
606#[derive(Debug)]
607pub struct RejoinedChannelBuffer {
608 pub buffer: proto::RejoinedChannelBuffer,
609 pub old_connection_id: ConnectionId,
610}
611
612#[derive(Clone)]
613pub struct JoinRoom {
614 pub room: proto::Room,
615 pub channel: Option<channel::Model>,
616}
617
618pub struct RejoinedRoom {
619 pub room: proto::Room,
620 pub rejoined_projects: Vec<RejoinedProject>,
621 pub reshared_projects: Vec<ResharedProject>,
622 pub channel: Option<channel::Model>,
623}
624
625pub struct ResharedProject {
626 pub id: ProjectId,
627 pub old_connection_id: ConnectionId,
628 pub collaborators: Vec<ProjectCollaborator>,
629 pub worktrees: Vec<proto::WorktreeMetadata>,
630}
631
632pub struct RejoinedProject {
633 pub id: ProjectId,
634 pub old_connection_id: ConnectionId,
635 pub collaborators: Vec<ProjectCollaborator>,
636 pub worktrees: Vec<RejoinedWorktree>,
637 pub language_servers: Vec<proto::LanguageServer>,
638}
639
640#[derive(Debug)]
641pub struct RejoinedWorktree {
642 pub id: u64,
643 pub abs_path: String,
644 pub root_name: String,
645 pub visible: bool,
646 pub updated_entries: Vec<proto::Entry>,
647 pub removed_entries: Vec<u64>,
648 pub updated_repositories: Vec<proto::RepositoryEntry>,
649 pub removed_repositories: Vec<u64>,
650 pub diagnostic_summaries: Vec<proto::DiagnosticSummary>,
651 pub settings_files: Vec<WorktreeSettingsFile>,
652 pub scan_id: u64,
653 pub completed_scan_id: u64,
654}
655
656pub struct LeftRoom {
657 pub room: proto::Room,
658 pub channel: Option<channel::Model>,
659 pub left_projects: HashMap<ProjectId, LeftProject>,
660 pub canceled_calls_to_user_ids: Vec<UserId>,
661 pub deleted: bool,
662}
663
664pub struct RefreshedRoom {
665 pub room: proto::Room,
666 pub channel: Option<channel::Model>,
667 pub stale_participant_user_ids: Vec<UserId>,
668 pub canceled_calls_to_user_ids: Vec<UserId>,
669}
670
671pub struct RefreshedChannelBuffer {
672 pub connection_ids: Vec<ConnectionId>,
673 pub collaborators: Vec<proto::Collaborator>,
674}
675
676pub struct Project {
677 pub id: ProjectId,
678 pub role: ChannelRole,
679 pub collaborators: Vec<ProjectCollaborator>,
680 pub worktrees: BTreeMap<u64, Worktree>,
681 pub language_servers: Vec<proto::LanguageServer>,
682}
683
684pub struct ProjectCollaborator {
685 pub connection_id: ConnectionId,
686 pub user_id: UserId,
687 pub replica_id: ReplicaId,
688 pub is_host: bool,
689}
690
691impl ProjectCollaborator {
692 pub fn to_proto(&self) -> proto::Collaborator {
693 proto::Collaborator {
694 peer_id: Some(self.connection_id.into()),
695 replica_id: self.replica_id.0 as u32,
696 user_id: self.user_id.to_proto(),
697 }
698 }
699}
700
701#[derive(Debug)]
702pub struct LeftProject {
703 pub id: ProjectId,
704 pub host_user_id: Option<UserId>,
705 pub host_connection_id: Option<ConnectionId>,
706 pub connection_ids: Vec<ConnectionId>,
707}
708
709pub struct Worktree {
710 pub id: u64,
711 pub abs_path: String,
712 pub root_name: String,
713 pub visible: bool,
714 pub entries: Vec<proto::Entry>,
715 pub repository_entries: BTreeMap<u64, proto::RepositoryEntry>,
716 pub diagnostic_summaries: Vec<proto::DiagnosticSummary>,
717 pub settings_files: Vec<WorktreeSettingsFile>,
718 pub scan_id: u64,
719 pub completed_scan_id: u64,
720}
721
722#[derive(Debug)]
723pub struct WorktreeSettingsFile {
724 pub path: String,
725 pub content: String,
726}
727
728pub struct NewExtensionVersion {
729 pub name: String,
730 pub version: semver::Version,
731 pub description: String,
732 pub authors: Vec<String>,
733 pub repository: String,
734 pub published_at: PrimitiveDateTime,
735}
736
737#[derive(Debug, Serialize, PartialEq)]
738pub struct ExtensionMetadata {
739 pub id: String,
740 pub name: String,
741 pub version: String,
742 pub authors: Vec<String>,
743 pub description: String,
744 pub repository: String,
745 #[serde(serialize_with = "serialize_iso8601")]
746 pub published_at: PrimitiveDateTime,
747 pub download_count: u64,
748}
749
750pub fn serialize_iso8601<S: Serializer>(
751 datetime: &PrimitiveDateTime,
752 serializer: S,
753) -> Result<S::Ok, S::Error> {
754 const SERDE_CONFIG: iso8601::EncodedConfig = iso8601::Config::DEFAULT
755 .set_year_is_six_digits(false)
756 .set_time_precision(iso8601::TimePrecision::Second {
757 decimal_digits: None,
758 })
759 .encode();
760
761 datetime
762 .assume_utc()
763 .format(&time::format_description::well_known::Iso8601::<SERDE_CONFIG>)
764 .map_err(S::Error::custom)?
765 .serialize(serializer)
766}