use super::*; use rpc::Notification; use sea_orm::{SelectColumns, TryInsertResult}; use time::OffsetDateTime; use util::ResultExt; impl Database { /// Inserts a record representing a user joining the chat for a given channel. pub async fn join_channel_chat( &self, channel_id: ChannelId, connection_id: ConnectionId, user_id: UserId, ) -> Result<()> { self.transaction(|tx| async move { let channel = self.get_channel_internal(channel_id, &tx).await?; self.check_user_is_channel_participant(&channel, user_id, &tx) .await?; channel_chat_participant::ActiveModel { id: ActiveValue::NotSet, channel_id: ActiveValue::Set(channel_id), user_id: ActiveValue::Set(user_id), connection_id: ActiveValue::Set(connection_id.id as i32), connection_server_id: ActiveValue::Set(ServerId(connection_id.owner_id as i32)), } .insert(&*tx) .await?; Ok(()) }) .await } /// Removes `channel_chat_participant` records associated with the given connection ID. pub async fn channel_chat_connection_lost( &self, connection_id: ConnectionId, tx: &DatabaseTransaction, ) -> Result<()> { channel_chat_participant::Entity::delete_many() .filter( Condition::all() .add( channel_chat_participant::Column::ConnectionServerId .eq(connection_id.owner_id), ) .add(channel_chat_participant::Column::ConnectionId.eq(connection_id.id)), ) .exec(tx) .await?; Ok(()) } /// Removes `channel_chat_participant` records associated with the given user ID so they /// will no longer get chat notifications. pub async fn leave_channel_chat( &self, channel_id: ChannelId, connection_id: ConnectionId, _user_id: UserId, ) -> Result<()> { self.transaction(|tx| async move { channel_chat_participant::Entity::delete_many() .filter( Condition::all() .add( channel_chat_participant::Column::ConnectionServerId .eq(connection_id.owner_id), ) .add(channel_chat_participant::Column::ConnectionId.eq(connection_id.id)) .add(channel_chat_participant::Column::ChannelId.eq(channel_id)), ) .exec(&*tx) .await?; Ok(()) }) .await } /// Retrieves the messages in the specified channel. /// /// Use `before_message_id` to paginate through the channel's messages. pub async fn get_channel_messages( &self, channel_id: ChannelId, user_id: UserId, count: usize, before_message_id: Option, ) -> Result> { self.transaction(|tx| async move { let channel = self.get_channel_internal(channel_id, &tx).await?; self.check_user_is_channel_participant(&channel, user_id, &tx) .await?; let mut condition = Condition::all().add(channel_message::Column::ChannelId.eq(channel_id)); if let Some(before_message_id) = before_message_id { condition = condition.add(channel_message::Column::Id.lt(before_message_id)); } let rows = channel_message::Entity::find() .filter(condition) .order_by_desc(channel_message::Column::Id) .limit(count as u64) .all(&*tx) .await?; self.load_channel_messages(rows, &tx).await }) .await } /// Returns the channel messages with the given IDs. pub async fn get_channel_messages_by_id( &self, user_id: UserId, message_ids: &[MessageId], ) -> Result> { self.transaction(|tx| async move { let rows = channel_message::Entity::find() .filter(channel_message::Column::Id.is_in(message_ids.iter().copied())) .order_by_desc(channel_message::Column::Id) .all(&*tx) .await?; let mut channels = HashMap::::default(); for row in &rows { channels.insert( row.channel_id, self.get_channel_internal(row.channel_id, &tx).await?, ); } for (_, channel) in channels { self.check_user_is_channel_participant(&channel, user_id, &tx) .await?; } let messages = self.load_channel_messages(rows, &tx).await?; Ok(messages) }) .await } async fn load_channel_messages( &self, rows: Vec, tx: &DatabaseTransaction, ) -> Result> { let mut messages = rows .into_iter() .map(|row| { let nonce = row.nonce.as_u64_pair(); proto::ChannelMessage { id: row.id.to_proto(), sender_id: row.sender_id.to_proto(), body: row.body, timestamp: row.sent_at.assume_utc().unix_timestamp() as u64, mentions: vec![], nonce: Some(proto::Nonce { upper_half: nonce.0, lower_half: nonce.1, }), reply_to_message_id: row.reply_to_message_id.map(|id| id.to_proto()), edited_at: row .edited_at .map(|t| t.assume_utc().unix_timestamp() as u64), } }) .collect::>(); messages.reverse(); let mut mentions = channel_message_mention::Entity::find() .filter(channel_message_mention::Column::MessageId.is_in(messages.iter().map(|m| m.id))) .order_by_asc(channel_message_mention::Column::MessageId) .order_by_asc(channel_message_mention::Column::StartOffset) .stream(tx) .await?; let mut message_ix = 0; while let Some(mention) = mentions.next().await { let mention = mention?; let message_id = mention.message_id.to_proto(); while let Some(message) = messages.get_mut(message_ix) { if message.id < message_id { message_ix += 1; } else { if message.id == message_id { message.mentions.push(proto::ChatMention { range: Some(proto::Range { start: mention.start_offset as u64, end: mention.end_offset as u64, }), user_id: mention.user_id.to_proto(), }); } break; } } } Ok(messages) } fn format_mentions_to_entities( &self, message_id: MessageId, body: &str, mentions: &[proto::ChatMention], ) -> Result> { Ok(mentions .iter() .filter_map(|mention| { let range = mention.range.as_ref()?; if !body.is_char_boundary(range.start as usize) || !body.is_char_boundary(range.end as usize) { return None; } Some(channel_message_mention::ActiveModel { message_id: ActiveValue::Set(message_id), start_offset: ActiveValue::Set(range.start as i32), end_offset: ActiveValue::Set(range.end as i32), user_id: ActiveValue::Set(UserId::from_proto(mention.user_id)), }) }) .collect::>()) } /// Creates a new channel message. pub async fn create_channel_message( &self, channel_id: ChannelId, user_id: UserId, body: &str, mentions: &[proto::ChatMention], timestamp: OffsetDateTime, nonce: u128, reply_to_message_id: Option, ) -> Result { self.transaction(|tx| async move { let channel = self.get_channel_internal(channel_id, &tx).await?; self.check_user_is_channel_participant(&channel, user_id, &tx) .await?; let mut rows = channel_chat_participant::Entity::find() .filter(channel_chat_participant::Column::ChannelId.eq(channel_id)) .stream(&*tx) .await?; let mut is_participant = false; let mut participant_connection_ids = HashSet::default(); let mut participant_user_ids = Vec::new(); while let Some(row) = rows.next().await { let row = row?; if row.user_id == user_id { is_participant = true; } participant_user_ids.push(row.user_id); participant_connection_ids.insert(row.connection()); } drop(rows); if !is_participant { Err(anyhow!("not a chat participant"))?; } let timestamp = timestamp.to_offset(time::UtcOffset::UTC); let timestamp = time::PrimitiveDateTime::new(timestamp.date(), timestamp.time()); let result = channel_message::Entity::insert(channel_message::ActiveModel { channel_id: ActiveValue::Set(channel_id), sender_id: ActiveValue::Set(user_id), body: ActiveValue::Set(body.to_string()), sent_at: ActiveValue::Set(timestamp), nonce: ActiveValue::Set(Uuid::from_u128(nonce)), id: ActiveValue::NotSet, reply_to_message_id: ActiveValue::Set(reply_to_message_id), edited_at: ActiveValue::NotSet, }) .on_conflict( OnConflict::columns([ channel_message::Column::SenderId, channel_message::Column::Nonce, ]) .do_nothing() .to_owned(), ) .do_nothing() .exec(&*tx) .await?; let message_id; let mut notifications = Vec::new(); match result { TryInsertResult::Inserted(result) => { message_id = result.last_insert_id; let mentioned_user_ids = mentions.iter().map(|m| m.user_id).collect::>(); let mentions = self.format_mentions_to_entities(message_id, body, mentions)?; if !mentions.is_empty() { channel_message_mention::Entity::insert_many(mentions) .exec(&*tx) .await?; } for mentioned_user in mentioned_user_ids { notifications.extend( self.create_notification( UserId::from_proto(mentioned_user), rpc::Notification::ChannelMessageMention { message_id: message_id.to_proto(), sender_id: user_id.to_proto(), channel_id: channel_id.to_proto(), }, false, &tx, ) .await?, ); } self.observe_channel_message_internal(channel_id, user_id, message_id, &tx) .await?; } _ => { message_id = channel_message::Entity::find() .filter(channel_message::Column::Nonce.eq(Uuid::from_u128(nonce))) .one(&*tx) .await? .ok_or_else(|| anyhow!("failed to insert message"))? .id; } } Ok(CreatedChannelMessage { message_id, participant_connection_ids, notifications, }) }) .await } pub async fn observe_channel_message( &self, channel_id: ChannelId, user_id: UserId, message_id: MessageId, ) -> Result { self.transaction(|tx| async move { self.observe_channel_message_internal(channel_id, user_id, message_id, &tx) .await?; let mut batch = NotificationBatch::default(); batch.extend( self.mark_notification_as_read( user_id, &Notification::ChannelMessageMention { message_id: message_id.to_proto(), sender_id: Default::default(), channel_id: Default::default(), }, &tx, ) .await?, ); Ok(batch) }) .await } async fn observe_channel_message_internal( &self, channel_id: ChannelId, user_id: UserId, message_id: MessageId, tx: &DatabaseTransaction, ) -> Result<()> { observed_channel_messages::Entity::insert(observed_channel_messages::ActiveModel { user_id: ActiveValue::Set(user_id), channel_id: ActiveValue::Set(channel_id), channel_message_id: ActiveValue::Set(message_id), }) .on_conflict( OnConflict::columns([ observed_channel_messages::Column::ChannelId, observed_channel_messages::Column::UserId, ]) .update_column(observed_channel_messages::Column::ChannelMessageId) .action_cond_where(observed_channel_messages::Column::ChannelMessageId.lt(message_id)) .to_owned(), ) // TODO: Try to upgrade SeaORM so we don't have to do this hack around their bug .exec_without_returning(tx) .await?; Ok(()) } pub async fn observed_channel_messages( &self, channel_ids: &[ChannelId], user_id: UserId, tx: &DatabaseTransaction, ) -> Result> { let rows = observed_channel_messages::Entity::find() .filter(observed_channel_messages::Column::UserId.eq(user_id)) .filter( observed_channel_messages::Column::ChannelId .is_in(channel_ids.iter().map(|id| id.0)), ) .all(tx) .await?; Ok(rows .into_iter() .map(|message| proto::ChannelMessageId { channel_id: message.channel_id.to_proto(), message_id: message.channel_message_id.to_proto(), }) .collect()) } pub async fn latest_channel_messages( &self, channel_ids: &[ChannelId], tx: &DatabaseTransaction, ) -> Result> { let mut values = String::new(); for id in channel_ids { if !values.is_empty() { values.push_str(", "); } write!(&mut values, "({})", id).unwrap(); } if values.is_empty() { return Ok(Vec::default()); } let sql = format!( r#" SELECT * FROM ( SELECT *, row_number() OVER ( PARTITION BY channel_id ORDER BY id DESC ) as row_number FROM channel_messages WHERE channel_id in ({values}) ) AS messages WHERE row_number = 1 "#, ); let stmt = Statement::from_string(self.pool.get_database_backend(), sql); let mut last_messages = channel_message::Model::find_by_statement(stmt) .stream(tx) .await?; let mut results = Vec::new(); while let Some(result) = last_messages.next().await { let message = result?; results.push(proto::ChannelMessageId { channel_id: message.channel_id.to_proto(), message_id: message.id.to_proto(), }); } Ok(results) } fn get_notification_kind_id_by_name(&self, notification_kind: &str) -> Option { self.notification_kinds_by_id .iter() .find(|(_, kind)| **kind == notification_kind) .map(|kind| kind.0 .0) } /// Removes the channel message with the given ID. pub async fn remove_channel_message( &self, channel_id: ChannelId, message_id: MessageId, user_id: UserId, ) -> Result<(Vec, Vec)> { self.transaction(|tx| async move { let mut rows = channel_chat_participant::Entity::find() .filter(channel_chat_participant::Column::ChannelId.eq(channel_id)) .stream(&*tx) .await?; let mut is_participant = false; let mut participant_connection_ids = Vec::new(); while let Some(row) = rows.next().await { let row = row?; if row.user_id == user_id { is_participant = true; } participant_connection_ids.push(row.connection()); } drop(rows); if !is_participant { Err(anyhow!("not a chat participant"))?; } let result = channel_message::Entity::delete_by_id(message_id) .filter(channel_message::Column::SenderId.eq(user_id)) .exec(&*tx) .await?; if result.rows_affected == 0 { let channel = self.get_channel_internal(channel_id, &tx).await?; if self .check_user_is_channel_admin(&channel, user_id, &tx) .await .is_ok() { let result = channel_message::Entity::delete_by_id(message_id) .exec(&*tx) .await?; if result.rows_affected == 0 { Err(anyhow!("no such message"))?; } } else { Err(anyhow!("operation could not be completed"))?; } } let notification_kind_id = self.get_notification_kind_id_by_name("ChannelMessageMention"); let existing_notifications = notification::Entity::find() .filter(notification::Column::EntityId.eq(message_id)) .filter(notification::Column::Kind.eq(notification_kind_id)) .select_column(notification::Column::Id) .all(&*tx) .await?; let existing_notification_ids = existing_notifications .into_iter() .map(|notification| notification.id) .collect(); // remove all the mention notifications for this message notification::Entity::delete_many() .filter(notification::Column::EntityId.eq(message_id)) .filter(notification::Column::Kind.eq(notification_kind_id)) .exec(&*tx) .await?; Ok((participant_connection_ids, existing_notification_ids)) }) .await } /// Updates the channel message with the given ID, body and timestamp(edited_at). pub async fn update_channel_message( &self, channel_id: ChannelId, message_id: MessageId, user_id: UserId, body: &str, mentions: &[proto::ChatMention], edited_at: OffsetDateTime, ) -> Result { self.transaction(|tx| async move { let channel = self.get_channel_internal(channel_id, &tx).await?; self.check_user_is_channel_participant(&channel, user_id, &tx) .await?; let mut rows = channel_chat_participant::Entity::find() .filter(channel_chat_participant::Column::ChannelId.eq(channel_id)) .stream(&*tx) .await?; let mut is_participant = false; let mut participant_connection_ids = Vec::new(); let mut participant_user_ids = Vec::new(); while let Some(row) = rows.next().await { let row = row?; if row.user_id == user_id { is_participant = true; } participant_user_ids.push(row.user_id); participant_connection_ids.push(row.connection()); } drop(rows); if !is_participant { Err(anyhow!("not a chat participant"))?; } let channel_message = channel_message::Entity::find_by_id(message_id) .filter(channel_message::Column::SenderId.eq(user_id)) .one(&*tx) .await?; let Some(channel_message) = channel_message else { Err(anyhow!("Channel message not found"))? }; let edited_at = edited_at.to_offset(time::UtcOffset::UTC); let edited_at = time::PrimitiveDateTime::new(edited_at.date(), edited_at.time()); let updated_message = channel_message::ActiveModel { body: ActiveValue::Set(body.to_string()), edited_at: ActiveValue::Set(Some(edited_at)), reply_to_message_id: ActiveValue::Unchanged(channel_message.reply_to_message_id), id: ActiveValue::Unchanged(message_id), channel_id: ActiveValue::Unchanged(channel_id), sender_id: ActiveValue::Unchanged(user_id), sent_at: ActiveValue::Unchanged(channel_message.sent_at), nonce: ActiveValue::Unchanged(channel_message.nonce), }; let result = channel_message::Entity::update_many() .set(updated_message) .filter(channel_message::Column::Id.eq(message_id)) .filter(channel_message::Column::SenderId.eq(user_id)) .exec(&*tx) .await?; if result.rows_affected == 0 { return Err(anyhow!( "Attempted to edit a message (id: {message_id}) which does not exist anymore." ))?; } // we have to fetch the old mentions, // so we don't send a notification when the message has been edited that you are mentioned in let old_mentions = channel_message_mention::Entity::find() .filter(channel_message_mention::Column::MessageId.eq(message_id)) .all(&*tx) .await?; // remove all existing mentions channel_message_mention::Entity::delete_many() .filter(channel_message_mention::Column::MessageId.eq(message_id)) .exec(&*tx) .await?; let new_mentions = self.format_mentions_to_entities(message_id, body, mentions)?; if !new_mentions.is_empty() { // insert new mentions channel_message_mention::Entity::insert_many(new_mentions) .exec(&*tx) .await?; } let mut update_mention_user_ids = HashSet::default(); let mut new_mention_user_ids = mentions.iter().map(|m| m.user_id).collect::>(); // Filter out users that were mentioned before for mention in &old_mentions { if new_mention_user_ids.contains(&mention.user_id.to_proto()) { update_mention_user_ids.insert(mention.user_id.to_proto()); } new_mention_user_ids.remove(&mention.user_id.to_proto()); } let notification_kind_id = self.get_notification_kind_id_by_name("ChannelMessageMention"); let existing_notifications = notification::Entity::find() .filter(notification::Column::EntityId.eq(message_id)) .filter(notification::Column::Kind.eq(notification_kind_id)) .all(&*tx) .await?; // determine which notifications should be updated or deleted let mut deleted_notification_ids = HashSet::default(); let mut updated_mention_notifications = Vec::new(); for notification in existing_notifications { if update_mention_user_ids.contains(¬ification.recipient_id.to_proto()) { if let Some(notification) = self::notifications::model_to_proto(self, notification).log_err() { updated_mention_notifications.push(notification); } } else { deleted_notification_ids.insert(notification.id); } } let mut notifications = Vec::new(); for mentioned_user in new_mention_user_ids { notifications.extend( self.create_notification( UserId::from_proto(mentioned_user), rpc::Notification::ChannelMessageMention { message_id: message_id.to_proto(), sender_id: user_id.to_proto(), channel_id: channel_id.to_proto(), }, false, &tx, ) .await?, ); } Ok(UpdatedChannelMessage { message_id, participant_connection_ids, notifications, reply_to_message_id: channel_message.reply_to_message_id, timestamp: channel_message.sent_at, deleted_mention_notification_ids: deleted_notification_ids .into_iter() .collect::>(), updated_mention_notifications, }) }) .await } }