messages.rs

  1use super::*;
  2use rpc::Notification;
  3use sea_orm::TryInsertResult;
  4use time::OffsetDateTime;
  5
  6impl Database {
  7    /// Inserts a record representing a user joining the chat for a given channel.
  8    pub async fn join_channel_chat(
  9        &self,
 10        channel_id: ChannelId,
 11        connection_id: ConnectionId,
 12        user_id: UserId,
 13    ) -> Result<()> {
 14        self.transaction(|tx| async move {
 15            let channel = self.get_channel_internal(channel_id, &tx).await?;
 16            self.check_user_is_channel_participant(&channel, user_id, &tx)
 17                .await?;
 18            channel_chat_participant::ActiveModel {
 19                id: ActiveValue::NotSet,
 20                channel_id: ActiveValue::Set(channel_id),
 21                user_id: ActiveValue::Set(user_id),
 22                connection_id: ActiveValue::Set(connection_id.id as i32),
 23                connection_server_id: ActiveValue::Set(ServerId(connection_id.owner_id as i32)),
 24            }
 25            .insert(&*tx)
 26            .await?;
 27            Ok(())
 28        })
 29        .await
 30    }
 31
 32    /// Removes `channel_chat_participant` records associated with the given connection ID.
 33    pub async fn channel_chat_connection_lost(
 34        &self,
 35        connection_id: ConnectionId,
 36        tx: &DatabaseTransaction,
 37    ) -> Result<()> {
 38        channel_chat_participant::Entity::delete_many()
 39            .filter(
 40                Condition::all()
 41                    .add(
 42                        channel_chat_participant::Column::ConnectionServerId
 43                            .eq(connection_id.owner_id),
 44                    )
 45                    .add(channel_chat_participant::Column::ConnectionId.eq(connection_id.id)),
 46            )
 47            .exec(tx)
 48            .await?;
 49        Ok(())
 50    }
 51
 52    /// Removes `channel_chat_participant` records associated with the given user ID so they
 53    /// will no longer get chat notifications.
 54    pub async fn leave_channel_chat(
 55        &self,
 56        channel_id: ChannelId,
 57        connection_id: ConnectionId,
 58        _user_id: UserId,
 59    ) -> Result<()> {
 60        self.transaction(|tx| async move {
 61            channel_chat_participant::Entity::delete_many()
 62                .filter(
 63                    Condition::all()
 64                        .add(
 65                            channel_chat_participant::Column::ConnectionServerId
 66                                .eq(connection_id.owner_id),
 67                        )
 68                        .add(channel_chat_participant::Column::ConnectionId.eq(connection_id.id))
 69                        .add(channel_chat_participant::Column::ChannelId.eq(channel_id)),
 70                )
 71                .exec(&*tx)
 72                .await?;
 73
 74            Ok(())
 75        })
 76        .await
 77    }
 78
 79    /// Retrieves the messages in the specified channel.
 80    ///
 81    /// Use `before_message_id` to paginate through the channel's messages.
 82    pub async fn get_channel_messages(
 83        &self,
 84        channel_id: ChannelId,
 85        user_id: UserId,
 86        count: usize,
 87        before_message_id: Option<MessageId>,
 88    ) -> Result<Vec<proto::ChannelMessage>> {
 89        self.transaction(|tx| async move {
 90            let channel = self.get_channel_internal(channel_id, &tx).await?;
 91            self.check_user_is_channel_participant(&channel, user_id, &tx)
 92                .await?;
 93
 94            let mut condition =
 95                Condition::all().add(channel_message::Column::ChannelId.eq(channel_id));
 96
 97            if let Some(before_message_id) = before_message_id {
 98                condition = condition.add(channel_message::Column::Id.lt(before_message_id));
 99            }
100
101            let rows = channel_message::Entity::find()
102                .filter(condition)
103                .order_by_desc(channel_message::Column::Id)
104                .limit(count as u64)
105                .all(&*tx)
106                .await?;
107
108            self.load_channel_messages(rows, &tx).await
109        })
110        .await
111    }
112
113    /// Returns the channel messages with the given IDs.
114    pub async fn get_channel_messages_by_id(
115        &self,
116        user_id: UserId,
117        message_ids: &[MessageId],
118    ) -> Result<Vec<proto::ChannelMessage>> {
119        self.transaction(|tx| async move {
120            let rows = channel_message::Entity::find()
121                .filter(channel_message::Column::Id.is_in(message_ids.iter().copied()))
122                .order_by_desc(channel_message::Column::Id)
123                .all(&*tx)
124                .await?;
125
126            let mut channels = HashMap::<ChannelId, channel::Model>::default();
127            for row in &rows {
128                channels.insert(
129                    row.channel_id,
130                    self.get_channel_internal(row.channel_id, &tx).await?,
131                );
132            }
133
134            for (_, channel) in channels {
135                self.check_user_is_channel_participant(&channel, user_id, &tx)
136                    .await?;
137            }
138
139            let messages = self.load_channel_messages(rows, &tx).await?;
140            Ok(messages)
141        })
142        .await
143    }
144
145    async fn load_channel_messages(
146        &self,
147        rows: Vec<channel_message::Model>,
148        tx: &DatabaseTransaction,
149    ) -> Result<Vec<proto::ChannelMessage>> {
150        let mut messages = rows
151            .into_iter()
152            .map(|row| {
153                let nonce = row.nonce.as_u64_pair();
154                proto::ChannelMessage {
155                    id: row.id.to_proto(),
156                    sender_id: row.sender_id.to_proto(),
157                    body: row.body,
158                    timestamp: row.sent_at.assume_utc().unix_timestamp() as u64,
159                    mentions: vec![],
160                    nonce: Some(proto::Nonce {
161                        upper_half: nonce.0,
162                        lower_half: nonce.1,
163                    }),
164                    reply_to_message_id: row.reply_to_message_id.map(|id| id.to_proto()),
165                }
166            })
167            .collect::<Vec<_>>();
168        messages.reverse();
169
170        let mut mentions = channel_message_mention::Entity::find()
171            .filter(channel_message_mention::Column::MessageId.is_in(messages.iter().map(|m| m.id)))
172            .order_by_asc(channel_message_mention::Column::MessageId)
173            .order_by_asc(channel_message_mention::Column::StartOffset)
174            .stream(tx)
175            .await?;
176
177        let mut message_ix = 0;
178        while let Some(mention) = mentions.next().await {
179            let mention = mention?;
180            let message_id = mention.message_id.to_proto();
181            while let Some(message) = messages.get_mut(message_ix) {
182                if message.id < message_id {
183                    message_ix += 1;
184                } else {
185                    if message.id == message_id {
186                        message.mentions.push(proto::ChatMention {
187                            range: Some(proto::Range {
188                                start: mention.start_offset as u64,
189                                end: mention.end_offset as u64,
190                            }),
191                            user_id: mention.user_id.to_proto(),
192                        });
193                    }
194                    break;
195                }
196            }
197        }
198
199        Ok(messages)
200    }
201
202    /// Creates a new channel message.
203    #[allow(clippy::too_many_arguments)]
204    pub async fn create_channel_message(
205        &self,
206        channel_id: ChannelId,
207        user_id: UserId,
208        body: &str,
209        mentions: &[proto::ChatMention],
210        timestamp: OffsetDateTime,
211        nonce: u128,
212        reply_to_message_id: Option<MessageId>,
213    ) -> Result<CreatedChannelMessage> {
214        self.transaction(|tx| async move {
215            let channel = self.get_channel_internal(channel_id, &tx).await?;
216            self.check_user_is_channel_participant(&channel, user_id, &tx)
217                .await?;
218
219            let mut rows = channel_chat_participant::Entity::find()
220                .filter(channel_chat_participant::Column::ChannelId.eq(channel_id))
221                .stream(&*tx)
222                .await?;
223
224            let mut is_participant = false;
225            let mut participant_connection_ids = Vec::new();
226            let mut participant_user_ids = Vec::new();
227            while let Some(row) = rows.next().await {
228                let row = row?;
229                if row.user_id == user_id {
230                    is_participant = true;
231                }
232                participant_user_ids.push(row.user_id);
233                participant_connection_ids.push(row.connection());
234            }
235            drop(rows);
236
237            if !is_participant {
238                Err(anyhow!("not a chat participant"))?;
239            }
240
241            let timestamp = timestamp.to_offset(time::UtcOffset::UTC);
242            let timestamp = time::PrimitiveDateTime::new(timestamp.date(), timestamp.time());
243
244            let result = channel_message::Entity::insert(channel_message::ActiveModel {
245                channel_id: ActiveValue::Set(channel_id),
246                sender_id: ActiveValue::Set(user_id),
247                body: ActiveValue::Set(body.to_string()),
248                sent_at: ActiveValue::Set(timestamp),
249                nonce: ActiveValue::Set(Uuid::from_u128(nonce)),
250                id: ActiveValue::NotSet,
251                reply_to_message_id: ActiveValue::Set(reply_to_message_id),
252            })
253            .on_conflict(
254                OnConflict::columns([
255                    channel_message::Column::SenderId,
256                    channel_message::Column::Nonce,
257                ])
258                .do_nothing()
259                .to_owned(),
260            )
261            .do_nothing()
262            .exec(&*tx)
263            .await?;
264
265            let message_id;
266            let mut notifications = Vec::new();
267            match result {
268                TryInsertResult::Inserted(result) => {
269                    message_id = result.last_insert_id;
270                    let mentioned_user_ids =
271                        mentions.iter().map(|m| m.user_id).collect::<HashSet<_>>();
272
273                    let mentions = mentions
274                        .iter()
275                        .filter_map(|mention| {
276                            let range = mention.range.as_ref()?;
277                            if !body.is_char_boundary(range.start as usize)
278                                || !body.is_char_boundary(range.end as usize)
279                            {
280                                return None;
281                            }
282                            Some(channel_message_mention::ActiveModel {
283                                message_id: ActiveValue::Set(message_id),
284                                start_offset: ActiveValue::Set(range.start as i32),
285                                end_offset: ActiveValue::Set(range.end as i32),
286                                user_id: ActiveValue::Set(UserId::from_proto(mention.user_id)),
287                            })
288                        })
289                        .collect::<Vec<_>>();
290                    if !mentions.is_empty() {
291                        channel_message_mention::Entity::insert_many(mentions)
292                            .exec(&*tx)
293                            .await?;
294                    }
295
296                    for mentioned_user in mentioned_user_ids {
297                        notifications.extend(
298                            self.create_notification(
299                                UserId::from_proto(mentioned_user),
300                                rpc::Notification::ChannelMessageMention {
301                                    message_id: message_id.to_proto(),
302                                    sender_id: user_id.to_proto(),
303                                    channel_id: channel_id.to_proto(),
304                                },
305                                false,
306                                &tx,
307                            )
308                            .await?,
309                        );
310                    }
311
312                    self.observe_channel_message_internal(channel_id, user_id, message_id, &tx)
313                        .await?;
314                }
315                _ => {
316                    message_id = channel_message::Entity::find()
317                        .filter(channel_message::Column::Nonce.eq(Uuid::from_u128(nonce)))
318                        .one(&*tx)
319                        .await?
320                        .ok_or_else(|| anyhow!("failed to insert message"))?
321                        .id;
322                }
323            }
324
325            let mut channel_members = self.get_channel_participants(&channel, &tx).await?;
326            channel_members.retain(|member| !participant_user_ids.contains(member));
327
328            Ok(CreatedChannelMessage {
329                message_id,
330                participant_connection_ids,
331                channel_members,
332                notifications,
333            })
334        })
335        .await
336    }
337
338    pub async fn observe_channel_message(
339        &self,
340        channel_id: ChannelId,
341        user_id: UserId,
342        message_id: MessageId,
343    ) -> Result<NotificationBatch> {
344        self.transaction(|tx| async move {
345            self.observe_channel_message_internal(channel_id, user_id, message_id, &tx)
346                .await?;
347            let mut batch = NotificationBatch::default();
348            batch.extend(
349                self.mark_notification_as_read(
350                    user_id,
351                    &Notification::ChannelMessageMention {
352                        message_id: message_id.to_proto(),
353                        sender_id: Default::default(),
354                        channel_id: Default::default(),
355                    },
356                    &tx,
357                )
358                .await?,
359            );
360            Ok(batch)
361        })
362        .await
363    }
364
365    async fn observe_channel_message_internal(
366        &self,
367        channel_id: ChannelId,
368        user_id: UserId,
369        message_id: MessageId,
370        tx: &DatabaseTransaction,
371    ) -> Result<()> {
372        observed_channel_messages::Entity::insert(observed_channel_messages::ActiveModel {
373            user_id: ActiveValue::Set(user_id),
374            channel_id: ActiveValue::Set(channel_id),
375            channel_message_id: ActiveValue::Set(message_id),
376        })
377        .on_conflict(
378            OnConflict::columns([
379                observed_channel_messages::Column::ChannelId,
380                observed_channel_messages::Column::UserId,
381            ])
382            .update_column(observed_channel_messages::Column::ChannelMessageId)
383            .action_cond_where(observed_channel_messages::Column::ChannelMessageId.lt(message_id))
384            .to_owned(),
385        )
386        // TODO: Try to upgrade SeaORM so we don't have to do this hack around their bug
387        .exec_without_returning(tx)
388        .await?;
389        Ok(())
390    }
391
392    pub async fn observed_channel_messages(
393        &self,
394        channel_ids: &[ChannelId],
395        user_id: UserId,
396        tx: &DatabaseTransaction,
397    ) -> Result<Vec<proto::ChannelMessageId>> {
398        let rows = observed_channel_messages::Entity::find()
399            .filter(observed_channel_messages::Column::UserId.eq(user_id))
400            .filter(
401                observed_channel_messages::Column::ChannelId
402                    .is_in(channel_ids.iter().map(|id| id.0)),
403            )
404            .all(tx)
405            .await?;
406
407        Ok(rows
408            .into_iter()
409            .map(|message| proto::ChannelMessageId {
410                channel_id: message.channel_id.to_proto(),
411                message_id: message.channel_message_id.to_proto(),
412            })
413            .collect())
414    }
415
416    pub async fn latest_channel_messages(
417        &self,
418        channel_ids: &[ChannelId],
419        tx: &DatabaseTransaction,
420    ) -> Result<Vec<proto::ChannelMessageId>> {
421        let mut values = String::new();
422        for id in channel_ids {
423            if !values.is_empty() {
424                values.push_str(", ");
425            }
426            write!(&mut values, "({})", id).unwrap();
427        }
428
429        if values.is_empty() {
430            return Ok(Vec::default());
431        }
432
433        let sql = format!(
434            r#"
435            SELECT
436                *
437            FROM (
438                SELECT
439                    *,
440                    row_number() OVER (
441                        PARTITION BY channel_id
442                        ORDER BY id DESC
443                    ) as row_number
444                FROM channel_messages
445                WHERE
446                    channel_id in ({values})
447            ) AS messages
448            WHERE
449                row_number = 1
450            "#,
451        );
452
453        let stmt = Statement::from_string(self.pool.get_database_backend(), sql);
454        let mut last_messages = channel_message::Model::find_by_statement(stmt)
455            .stream(tx)
456            .await?;
457
458        let mut results = Vec::new();
459        while let Some(result) = last_messages.next().await {
460            let message = result?;
461            results.push(proto::ChannelMessageId {
462                channel_id: message.channel_id.to_proto(),
463                message_id: message.id.to_proto(),
464            });
465        }
466
467        Ok(results)
468    }
469
470    /// Removes the channel message with the given ID.
471    pub async fn remove_channel_message(
472        &self,
473        channel_id: ChannelId,
474        message_id: MessageId,
475        user_id: UserId,
476    ) -> Result<Vec<ConnectionId>> {
477        self.transaction(|tx| async move {
478            let mut rows = channel_chat_participant::Entity::find()
479                .filter(channel_chat_participant::Column::ChannelId.eq(channel_id))
480                .stream(&*tx)
481                .await?;
482
483            let mut is_participant = false;
484            let mut participant_connection_ids = Vec::new();
485            while let Some(row) = rows.next().await {
486                let row = row?;
487                if row.user_id == user_id {
488                    is_participant = true;
489                }
490                participant_connection_ids.push(row.connection());
491            }
492            drop(rows);
493
494            if !is_participant {
495                Err(anyhow!("not a chat participant"))?;
496            }
497
498            let result = channel_message::Entity::delete_by_id(message_id)
499                .filter(channel_message::Column::SenderId.eq(user_id))
500                .exec(&*tx)
501                .await?;
502
503            if result.rows_affected == 0 {
504                let channel = self.get_channel_internal(channel_id, &tx).await?;
505                if self
506                    .check_user_is_channel_admin(&channel, user_id, &tx)
507                    .await
508                    .is_ok()
509                {
510                    let result = channel_message::Entity::delete_by_id(message_id)
511                        .exec(&*tx)
512                        .await?;
513                    if result.rows_affected == 0 {
514                        Err(anyhow!("no such message"))?;
515                    }
516                } else {
517                    Err(anyhow!("operation could not be completed"))?;
518                }
519            }
520
521            Ok(participant_connection_ids)
522        })
523        .await
524    }
525}