use super::*;use rpc::Notification;use util::ResultExt; impl Database { /// Initializes the different kinds of notifications by upserting records for them. pub async fn initialize_notification_kinds(&mut self) -> Result<()> { let all_kinds = Notification::all_variant_names(); let existing_kinds = notification_kind::Entity::find().all(&self.pool).await?; let kinds_to_create: Vec<_> = all_kinds .iter() .filter(|&kind| { !existing_kinds .iter() .any(|existing| existing.name == **kind) }) .map(|kind| notification_kind::ActiveModel { name: ActiveValue::Set(kind.to_string()), ..Default::default() }) .collect(); if !kinds_to_create.is_empty() { notification_kind::Entity::insert_many(kinds_to_create) .exec_without_returning(&self.pool) .await?; } let mut rows = notification_kind::Entity::find().stream(&self.pool).await?; while let Some(row) = rows.next().await { let row = row?; self.notification_kinds_by_name.insert(row.name, row.id); } for name in Notification::all_variant_names() { if let Some(id) = self.notification_kinds_by_name.get(*name).copied() { self.notification_kinds_by_id.insert(id, name); } } Ok(()) } /// Returns the notifications for the given recipient. pub async fn get_notifications( &self, recipient_id: UserId, limit: usize, before_id: Option<NotificationId>, ) -> Result<Vec<proto::Notification>> { self.transaction(|tx| async move { let mut result = Vec::new(); let mut condition = Condition::all().add(notification::Column::RecipientId.eq(recipient_id)); if let Some(before_id) = before_id { condition = condition.add(notification::Column::Id.lt(before_id)); } let mut rows = notification::Entity::find() .filter(condition) .order_by_desc(notification::Column::Id) .limit(limit as u64) .stream(&*tx) .await?; while let Some(row) = rows.next().await { let row = row?; if let Some(proto) = model_to_proto(self, row).log_err() { result.push(proto); } } result.reverse(); Ok(result) }) .await } /// Creates a notification. If `avoid_duplicates` is set to true, then avoid /// creating a new notification if the given recipient already has an /// unread notification with the given kind and entity id. pub async fn create_notification( &self, recipient_id: UserId, notification: Notification, avoid_duplicates: bool, tx: &DatabaseTransaction, ) -> Result<Option<(UserId, proto::Notification)>> { if avoid_duplicates && self .find_notification(recipient_id, ¬ification, tx) .await? .is_some() { return Ok(None); } let proto = notification.to_proto(); let kind = notification_kind_from_proto(self, &proto)?; let model = notification::ActiveModel { recipient_id: ActiveValue::Set(recipient_id), kind: ActiveValue::Set(kind), entity_id: ActiveValue::Set(proto.entity_id.map(|id| id as i32)), content: ActiveValue::Set(proto.content.clone()), ..Default::default() } .save(tx) .await?; Ok(Some(( recipient_id, proto::Notification { id: model.id.as_ref().to_proto(), kind: proto.kind, timestamp: model.created_at.as_ref().assume_utc().unix_timestamp() as u64, is_read: false, response: None, content: proto.content, entity_id: proto.entity_id, }, ))) } /// Remove an unread notification with the given recipient, kind and /// entity id. pub async fn remove_notification( &self, recipient_id: UserId, notification: Notification, tx: &DatabaseTransaction, ) -> Result<Option<NotificationId>> { let id = self .find_notification(recipient_id, ¬ification, tx) .await?; if let Some(id) = id { notification::Entity::delete_by_id(id).exec(tx).await?; } Ok(id) } /// Populate the response for the notification with the given kind and /// entity id. pub async fn mark_notification_as_read_with_response( &self, recipient_id: UserId, notification: &Notification, response: bool, tx: &DatabaseTransaction, ) -> Result<Option<(UserId, proto::Notification)>> { self.mark_notification_as_read_internal(recipient_id, notification, Some(response), tx) .await } /// Marks the given notification as read. pub async fn mark_notification_as_read( &self, recipient_id: UserId, notification: &Notification, tx: &DatabaseTransaction, ) -> Result<Option<(UserId, proto::Notification)>> { self.mark_notification_as_read_internal(recipient_id, notification, None, tx) .await } /// Marks the notification with the given ID as read. pub async fn mark_notification_as_read_by_id( &self, recipient_id: UserId, notification_id: NotificationId, ) -> Result<NotificationBatch> { self.transaction(|tx| async move { let row = notification::Entity::update(notification::ActiveModel { id: ActiveValue::Unchanged(notification_id), recipient_id: ActiveValue::Unchanged(recipient_id), is_read: ActiveValue::Set(true), ..Default::default() }) .exec(&*tx) .await?; Ok(model_to_proto(self, row) .map(|notification| (recipient_id, notification)) .into_iter() .collect()) }) .await } async fn mark_notification_as_read_internal( &self, recipient_id: UserId, notification: &Notification, response: Option<bool>, tx: &DatabaseTransaction, ) -> Result<Option<(UserId, proto::Notification)>> { if let Some(id) = self .find_notification(recipient_id, notification, tx) .await? { let row = notification::Entity::update(notification::ActiveModel { id: ActiveValue::Unchanged(id), recipient_id: ActiveValue::Unchanged(recipient_id), is_read: ActiveValue::Set(true), response: if let Some(response) = response { ActiveValue::Set(Some(response)) } else { ActiveValue::NotSet }, ..Default::default() }) .exec(tx) .await?; Ok(model_to_proto(self, row) .map(|notification| (recipient_id, notification)) .ok()) } else { Ok(None) } } /// Find an unread notification by its recipient, kind and entity id. async fn find_notification( &self, recipient_id: UserId, notification: &Notification, tx: &DatabaseTransaction, ) -> Result<Option<NotificationId>> { let proto = notification.to_proto(); let kind = notification_kind_from_proto(self, &proto)?; #[derive(Copy, Clone, Debug, EnumIter, DeriveColumn)] enum QueryIds { Id, } Ok(notification::Entity::find() .select_only() .column(notification::Column::Id) .filter( Condition::all() .add(notification::Column::RecipientId.eq(recipient_id)) .add(notification::Column::IsRead.eq(false)) .add(notification::Column::Kind.eq(kind)) .add(if proto.entity_id.is_some() { notification::Column::EntityId.eq(proto.entity_id) } else { notification::Column::EntityId.is_null() }), ) .into_values::<_, QueryIds>() .one(tx) .await?) }} pub fn model_to_proto(this: &Database, row: notification::Model) -> Result<proto::Notification> { let kind = this .notification_kinds_by_id .get(&row.kind) .ok_or_else(|| anyhow!("Unknown notification kind"))?; Ok(proto::Notification { id: row.id.to_proto(), kind: kind.to_string(), timestamp: row.created_at.assume_utc().unix_timestamp() as u64, is_read: row.is_read, response: row.response, content: row.content, entity_id: row.entity_id.map(|id| id as u64), })} fn notification_kind_from_proto( this: &Database, proto: &proto::Notification,) -> Result<NotificationKindId> { Ok(this .notification_kinds_by_name .get(&proto.kind) .copied() .ok_or_else(|| anyhow!("invalid notification kind {:?}", proto.kind))?)}
use super::*;
use rpc::Notification;
use util::ResultExt;
impl Database {
/// Initializes the different kinds of notifications by upserting records for them.
pub async fn initialize_notification_kinds(&mut self) -> Result<()> {
let all_kinds = Notification::all_variant_names();
let existing_kinds = notification_kind::Entity::find().all(&self.pool).await?;
let kinds_to_create: Vec<_> = all_kinds
.iter()
.filter(|&kind| {
!existing_kinds
.iter()
.any(|existing| existing.name == **kind)
})
.map(|kind| notification_kind::ActiveModel {
name: ActiveValue::Set(kind.to_string()),
..Default::default()
})
.collect();
if !kinds_to_create.is_empty() {
notification_kind::Entity::insert_many(kinds_to_create)
.exec_without_returning(&self.pool)
.await?;
}
let mut rows = notification_kind::Entity::find().stream(&self.pool).await?;
while let Some(row) = rows.next().await {
let row = row?;
self.notification_kinds_by_name.insert(row.name, row.id);
}
for name in Notification::all_variant_names() {
if let Some(id) = self.notification_kinds_by_name.get(*name).copied() {
self.notification_kinds_by_id.insert(id, name);
}
}
Ok(())
}
/// Returns the notifications for the given recipient.
pub async fn get_notifications(
&self,
recipient_id: UserId,
limit: usize,
before_id: Option<NotificationId>,
) -> Result<Vec<proto::Notification>> {
self.transaction(|tx| async move {
let mut result = Vec::new();
let mut condition =
Condition::all().add(notification::Column::RecipientId.eq(recipient_id));
if let Some(before_id) = before_id {
condition = condition.add(notification::Column::Id.lt(before_id));
}
let mut rows = notification::Entity::find()
.filter(condition)
.order_by_desc(notification::Column::Id)
.limit(limit as u64)
.stream(&*tx)
.await?;
while let Some(row) = rows.next().await {
let row = row?;
if let Some(proto) = model_to_proto(self, row).log_err() {
result.push(proto);
}
}
result.reverse();
Ok(result)
})
.await
}
/// Creates a notification. If `avoid_duplicates` is set to true, then avoid
/// creating a new notification if the given recipient already has an
/// unread notification with the given kind and entity id.
pub async fn create_notification(
&self,
recipient_id: UserId,
notification: Notification,
avoid_duplicates: bool,
tx: &DatabaseTransaction,
) -> Result<Option<(UserId, proto::Notification)>> {
if avoid_duplicates
&& self
.find_notification(recipient_id, ¬ification, tx)
.await?
.is_some()
{
return Ok(None);
}
let proto = notification.to_proto();
let kind = notification_kind_from_proto(self, &proto)?;
let model = notification::ActiveModel {
recipient_id: ActiveValue::Set(recipient_id),
kind: ActiveValue::Set(kind),
entity_id: ActiveValue::Set(proto.entity_id.map(|id| id as i32)),
content: ActiveValue::Set(proto.content.clone()),
..Default::default()
}
.save(tx)
.await?;
Ok(Some((
recipient_id,
proto::Notification {
id: model.id.as_ref().to_proto(),
kind: proto.kind,
timestamp: model.created_at.as_ref().assume_utc().unix_timestamp() as u64,
is_read: false,
response: None,
content: proto.content,
entity_id: proto.entity_id,
},
)))
}
/// Remove an unread notification with the given recipient, kind and
/// entity id.
pub async fn remove_notification(
&self,
recipient_id: UserId,
notification: Notification,
tx: &DatabaseTransaction,
) -> Result<Option<NotificationId>> {
let id = self
.find_notification(recipient_id, ¬ification, tx)
.await?;
if let Some(id) = id {
notification::Entity::delete_by_id(id).exec(tx).await?;
}
Ok(id)
}
/// Populate the response for the notification with the given kind and
/// entity id.
pub async fn mark_notification_as_read_with_response(
&self,
recipient_id: UserId,
notification: &Notification,
response: bool,
tx: &DatabaseTransaction,
) -> Result<Option<(UserId, proto::Notification)>> {
self.mark_notification_as_read_internal(recipient_id, notification, Some(response), tx)
.await
}
/// Marks the given notification as read.
pub async fn mark_notification_as_read(
&self,
recipient_id: UserId,
notification: &Notification,
tx: &DatabaseTransaction,
) -> Result<Option<(UserId, proto::Notification)>> {
self.mark_notification_as_read_internal(recipient_id, notification, None, tx)
.await
}
/// Marks the notification with the given ID as read.
pub async fn mark_notification_as_read_by_id(
&self,
recipient_id: UserId,
notification_id: NotificationId,
) -> Result<NotificationBatch> {
self.transaction(|tx| async move {
let row = notification::Entity::update(notification::ActiveModel {
id: ActiveValue::Unchanged(notification_id),
recipient_id: ActiveValue::Unchanged(recipient_id),
is_read: ActiveValue::Set(true),
..Default::default()
})
.exec(&*tx)
.await?;
Ok(model_to_proto(self, row)
.map(|notification| (recipient_id, notification))
.into_iter()
.collect())
})
.await
}
async fn mark_notification_as_read_internal(
&self,
recipient_id: UserId,
notification: &Notification,
response: Option<bool>,
tx: &DatabaseTransaction,
) -> Result<Option<(UserId, proto::Notification)>> {
if let Some(id) = self
.find_notification(recipient_id, notification, tx)
.await?
{
let row = notification::Entity::update(notification::ActiveModel {
id: ActiveValue::Unchanged(id),
recipient_id: ActiveValue::Unchanged(recipient_id),
is_read: ActiveValue::Set(true),
response: if let Some(response) = response {
ActiveValue::Set(Some(response))
} else {
ActiveValue::NotSet
},
..Default::default()
})
.exec(tx)
.await?;
Ok(model_to_proto(self, row)
.map(|notification| (recipient_id, notification))
.ok())
} else {
Ok(None)
}
}
/// Find an unread notification by its recipient, kind and entity id.
async fn find_notification(
&self,
recipient_id: UserId,
notification: &Notification,
tx: &DatabaseTransaction,
) -> Result<Option<NotificationId>> {
let proto = notification.to_proto();
let kind = notification_kind_from_proto(self, &proto)?;
#[derive(Copy, Clone, Debug, EnumIter, DeriveColumn)]
enum QueryIds {
Id,
}
Ok(notification::Entity::find()
.select_only()
.column(notification::Column::Id)
.filter(
Condition::all()
.add(notification::Column::RecipientId.eq(recipient_id))
.add(notification::Column::IsRead.eq(false))
.add(notification::Column::Kind.eq(kind))
.add(if proto.entity_id.is_some() {
notification::Column::EntityId.eq(proto.entity_id)
} else {
notification::Column::EntityId.is_null()
}),
)
.into_values::<_, QueryIds>()
.one(tx)
.await?)
}
}
pub fn model_to_proto(this: &Database, row: notification::Model) -> Result<proto::Notification> {
let kind = this
.notification_kinds_by_id
.get(&row.kind)
.ok_or_else(|| anyhow!("Unknown notification kind"))?;
Ok(proto::Notification {
id: row.id.to_proto(),
kind: kind.to_string(),
timestamp: row.created_at.assume_utc().unix_timestamp() as u64,
is_read: row.is_read,
response: row.response,
content: row.content,
entity_id: row.entity_id.map(|id| id as u64),
})
}
fn notification_kind_from_proto(
this: &Database,
proto: &proto::Notification,
) -> Result<NotificationKindId> {
Ok(this
.notification_kinds_by_name
.get(&proto.kind)
.copied()
.ok_or_else(|| anyhow!("invalid notification kind {:?}", proto.kind))?)
}