jmug/stoatchat
1
0
Fork 1
Code Issues Packages Projects Releases Wiki Activity

chore(monorepo): delta, january, quark

Browse Source
This commit is contained in:
Paul Makles
2022-06-02 00:04:22 +01:00
parent 5d8432e267
commit 2ce610e1e7
232 changed files with 18094 additions and 554 deletions
Download Patch File Download Diff File Expand all files Collapse all files

26
crates/quark/src/impl/mongo/admin/migrations.rs Normal file
View File

@@ -0,0 +1,26 @@
use crate::{AbstractMigrations, Result};
use super::super::MongoDb;
mod init;
mod scripts;
#[async_trait]
impl AbstractMigrations for MongoDb {
async fn migrate_database(&self) -> Result<()> {
info!("Migrating the database.");
let list = self
.list_database_names(None, None)
.await
.expect("Failed to fetch database names.");
if list.iter().any(|x| x == "revolt") {
scripts::migrate_database(self).await;
} else {
init::create_database(self).await;
}
Ok(())
}
}

187
crates/quark/src/impl/mongo/admin/migrations/init.rs Normal file
View File

@@ -0,0 +1,187 @@
use crate::r#impl::mongo::MongoDb;
use super::scripts::LATEST_REVISION;
use log::info;
use mongodb::bson::doc;
use mongodb::options::CreateCollectionOptions;
pub async fn create_database(db: &MongoDb) {
info!("Creating database.");
let db = db.db();
db.create_collection("accounts", None)
.await
.expect("Failed to create accounts collection.");
db.create_collection("users", None)
.await
.expect("Failed to create users collection.");
db.create_collection("channels", None)
.await
.expect("Failed to create channels collection.");
db.create_collection("messages", None)
.await
.expect("Failed to create messages collection.");
db.create_collection("servers", None)
.await
.expect("Failed to create servers collection.");
db.create_collection("server_members", None)
.await
.expect("Failed to create server_members collection.");
db.create_collection("server_bans", None)
.await
.expect("Failed to create server_bans collection.");
db.create_collection("channel_invites", None)
.await
.expect("Failed to create channel_invites collection.");
db.create_collection("channel_unreads", None)
.await
.expect("Failed to create channel_unreads collection.");
db.create_collection("migrations", None)
.await
.expect("Failed to create migrations collection.");
db.create_collection("attachments", None)
.await
.expect("Failed to create attachments collection.");
db.create_collection("user_settings", None)
.await
.expect("Failed to create user_settings collection.");
db.create_collection("bots", None)
.await
.expect("Failed to create bots collection.");
db.create_collection(
"pubsub",
CreateCollectionOptions::builder()
.capped(true)
.size(1_000_000)
.build(),
)
.await
.expect("Failed to create pubsub collection.");
db.run_command(
doc! {
"createIndexes": "users",
"indexes": [
{
"key": {
"username": 1_i32
},
"name": "username",
"unique": true,
"collation": {
"locale": "en",
"strength": 2_i32
}
}
]
},
None,
)
.await
.expect("Failed to create username index.");
db.run_command(
doc! {
"createIndexes": "messages",
"indexes": [
{
"key": {
"content": "text"
},
"name": "content"
},
{
"key": {
"channel": 1_i32
},
"name": "channel"
},
{
"key": {
"channel": 1_i32,
"_id": 1_i32
},
"name": "channel_id_compound"
}
]
},
None,
)
.await
.expect("Failed to create message index.");
db.run_command(
doc! {
"createIndexes": "channel_unreads",
"indexes": [
{
"key": {
"_id.channel": 1_i32,
"_id.user": 1_i32,
},
"name": "compound_id"
},
{
"key": {
"_id.user": 1_i32,
},
"name": "user_id"
}
]
},
None,
)
.await
.expect("Failed to create channel_unreads index.");
db.run_command(
doc! {
"createIndexes": "server_members",
"indexes": [
{
"key": {
"_id.server": 1_i32,
"_id.user": 1_i32,
},
"name": "compound_id"
},
{
"key": {
"_id.user": 1_i32,
},
"name": "user_id"
}
]
},
None,
)
.await
.expect("Failed to create server_members index.");
db.collection("migrations")
.insert_one(
doc! {
"_id": 0_i32,
"revision": LATEST_REVISION
},
None,
)
.await
.expect("Failed to save migration info.");
info!("Created database.");
}

610
crates/quark/src/impl/mongo/admin/migrations/scripts.rs Normal file
View File

@@ -0,0 +1,610 @@
use std::time::Duration;
use bson::Bson;
use futures::StreamExt;
use log::info;
use mongodb::{
bson::{doc, from_bson, from_document, to_document, Document},
options::FindOptions,
};
use serde::{Deserialize, Serialize};
use crate::{r#impl::mongo::MongoDb, Permission, DEFAULT_PERMISSION_SERVER};
#[derive(Serialize, Deserialize)]
struct MigrationInfo {
_id: i32,
revision: i32,
}
pub const LATEST_REVISION: i32 = 15;
pub async fn migrate_database(db: &MongoDb) {
let migrations = db.col::<Document>("migrations");
let data = migrations
.find_one(None, None)
.await
.expect("Failed to fetch migration data.");
if let Some(doc) = data {
let info: MigrationInfo =
from_document(doc).expect("Failed to read migration information.");
let revision = run_migrations(db, info.revision).await;
migrations
.update_one(
doc! {
"_id": info._id
},
doc! {
"$set": {
"revision": revision
}
},
None,
)
.await
.expect("Failed to commit migration information.");
info!("Migration complete. Currently at revision {}.", revision);
} else {
panic!("Database was configured incorrectly, possibly because initalization failed.")
}
}
pub async fn run_migrations(db: &MongoDb, revision: i32) -> i32 {
info!("Starting database migration.");
if revision <= 0 {
info!("Running migration [revision 0]: Test migration system.");
}
if revision <= 1 {
info!("Running migration [revision 1 / 2021-04-24]: Migrate to Autumn v1.0.0.");
let messages = db.col::<Document>("messages");
let attachments = db.col::<Document>("attachments");
messages
.update_many(
doc! { "attachment": { "$exists": 1_i32 } },
doc! { "$set": { "attachment.tag": "attachments", "attachment.size": 0_i32 } },
None,
)
.await
.expect("Failed to update messages.");
attachments
.update_many(
doc! {},
doc! { "$set": { "tag": "attachments", "size": 0_i32 } },
None,
)
.await
.expect("Failed to update attachments.");
}
if revision <= 2 {
info!("Running migration [revision 2 / 2021-05-08]: Add servers collection.");
db.db()
.create_collection("servers", None)
.await
.expect("Failed to create servers collection.");
}
if revision <= 3 {
info!("Running migration [revision 3 / 2021-05-25]: Support multiple file uploads, add channel_unreads and user_settings.");
let messages = db.col::<Document>("messages");
let mut cursor = messages
.find(
doc! {
"attachment": {
"$exists": 1_i32
}
},
FindOptions::builder()
.projection(doc! {
"_id": 1_i32,
"attachments": [ "$attachment" ]
})
.build(),
)
.await
.expect("Failed to fetch messages.");
while let Some(result) = cursor.next().await {
let doc = result.unwrap();
let id = doc.get_str("_id").unwrap();
let attachments = doc.get_array("attachments").unwrap();
messages
.update_one(
doc! { "_id": id },
doc! { "$unset": { "attachment": 1_i32 }, "$set": { "attachments": attachments } },
None,
)
.await
.unwrap();
}
db.db()
.create_collection("channel_unreads", None)
.await
.expect("Failed to create channel_unreads collection.");
db.db()
.create_collection("user_settings", None)
.await
.expect("Failed to create user_settings collection.");
}
if revision <= 4 {
info!("Running migration [revision 4 / 2021-06-01]: Add more server collections.");
db.db()
.create_collection("server_members", None)
.await
.expect("Failed to create server_members collection.");
db.db()
.create_collection("server_bans", None)
.await
.expect("Failed to create server_bans collection.");
db.db()
.create_collection("channel_invites", None)
.await
.expect("Failed to create channel_invites collection.");
}
if revision <= 5 {
info!("Running migration [revision 5 / 2021-06-26]: Add permissions.");
#[derive(Serialize)]
struct Server {
pub default_permissions: (i32, i32),
}
let server = Server {
default_permissions: (0_i32, 0_i32),
};
db.col::<Document>("servers")
.update_many(
doc! {},
doc! {
"$set": to_document(&server).unwrap()
},
None,
)
.await
.expect("Failed to migrate servers.");
}
if revision <= 6 {
info!("Running migration [revision 6 / 2021-07-09]: Add message text index.");
db.db()
.run_command(
doc! {
"createIndexes": "messages",
"indexes": [
{
"key": {
"content": "text"
},
"name": "content"
}
]
},
None,
)
.await
.expect("Failed to create message index.");
}
if revision <= 7 {
info!("Running migration [revision 7 / 2021-08-11]: Add message text index.");
db.db()
.create_collection("bots", None)
.await
.expect("Failed to create bots collection.");
}
if revision <= 8 {
info!("Running migration [revision 8 / 2021-09-10]: Update to rAuth version 1.");
db.db()
.run_command(
doc! {
"dropIndexes": "accounts",
"index": ["email", "email_normalised"]
},
None,
)
.await
.expect("Failed to delete legacy account indexes.");
let col = db.col::<Document>("sessions");
let mut cursor = db
.col::<Document>("accounts")
.find(doc! {}, None)
.await
.unwrap();
while let Some(doc) = cursor.next().await {
if let Ok(account) = doc {
let id = account.get_str("_id").unwrap();
if let Some(sessions) = account.get("sessions") {
#[derive(Deserialize)]
struct Session {
id: String,
token: String,
friendly_name: String,
subscription: Option<Document>,
}
let sessions = from_bson::<Vec<Session>>(sessions.clone()).unwrap();
for session in sessions {
info!("Converting session {} to new format.", &session.id);
let mut doc = doc! {
"_id": session.id,
"token": session.token,
"user_id": id,
"name": session.friendly_name,
};
if let Some(sub) = session.subscription {
doc.insert("subscription", sub);
}
col.insert_one(doc, None).await.ok();
}
} else {
info!("Account doesn't have any sessions!");
}
}
}
db.col::<Document>("accounts")
.update_many(
doc! {},
doc! {
"$unset": {
"sessions": 1_i32,
},
"$set": {
"mfa": {
"recovery_codes": []
}
}
},
None,
)
.await
.unwrap();
}
if revision <= 9 {
info!("Running migration [revision 9 / 2021-09-14]: Switch from last_message to last_message_id.");
let mut cursor = db
.col::<Document>("channels")
.find(doc! {}, None)
.await
.unwrap();
while let Some(doc) = cursor.next().await {
if let Ok(channel) = doc {
let channel_id = channel.get_str("_id").unwrap();
if let Some(last_message) = channel.get("last_message") {
#[derive(Serialize, Deserialize, Debug, Clone)]
pub struct Obj {
#[serde(rename = "_id")]
id: String,
}
#[derive(Serialize, Deserialize, Debug, Clone)]
#[serde(untagged)]
pub enum LastMessage {
Obj(Obj),
Id(String),
}
let lm = from_bson::<LastMessage>(last_message.clone()).unwrap();
let id = match lm {
LastMessage::Obj(Obj { id }) => id,
LastMessage::Id(id) => id,
};
info!("Converting session {} to new format.", &channel_id);
db.col::<Document>("channels")
.update_one(
doc! {
"_id": channel_id
},
doc! {
"$set": {
"last_message_id": id
},
"$unset": {
"last_message": 1_i32,
}
},
None,
)
.await
.unwrap();
} else {
info!("{} has no last_message.", &channel_id);
}
}
}
}
if revision <= 10 {
info!("Running migration [revision 10 / 2021-11-01]: Remove nonce values on channels and servers.");
db.col::<Document>("servers")
.update_many(
doc! {},
doc! {
"$unset": {
"nonce": 1_i32,
}
},
None,
)
.await
.unwrap();
db.col::<Document>("channels")
.update_many(
doc! {},
doc! {
"$unset": {
"nonce": 1_i32,
}
},
None,
)
.await
.unwrap();
}
if revision <= 11 {
info!("Running migration [revision 11 / 2021-11-14]: Add indexes to database.");
db.db()
.run_command(
doc! {
"createIndexes": "messages",
"indexes": [
{
"key": {
"channel": 1_i32
},
"name": "channel"
}
]
},
None,
)
.await
.expect("Failed to create message index.");
db.db()
.run_command(
doc! {
"createIndexes": "channel_unreads",
"indexes": [
{
"key": {
"_id.channel": 1_i32,
"_id.user": 1_i32,
},
"name": "compound_id"
},
{
"key": {
"_id.user": 1_i32,
},
"name": "user_id"
}
]
},
None,
)
.await
.expect("Failed to create channel_unreads index.");
db.db()
.run_command(
doc! {
"createIndexes": "server_members",
"indexes": [
{
"key": {
"_id.server": 1_i32,
"_id.user": 1_i32,
},
"name": "compound_id"
},
{
"key": {
"_id.user": 1_i32,
},
"name": "user_id"
}
]
},
None,
)
.await
.expect("Failed to create server_members index.");
}
if revision <= 12 {
info!("Running migration [revision 12 / 2021-11-21]: Add indexes to database.");
db.db()
.run_command(
doc! {
"createIndexes": "messages",
"indexes": [
{
"key": {
"channel": 1_i32,
"_id": 1_i32
},
"name": "channel_id_compound"
}
]
},
None,
)
.await
.expect("Failed to create message index.");
}
if revision <= 13 {
info!("Running migration [revision 13 / 22-02-2022]: Wipe legacy permission values.");
warn!("This is a destructive operation and will wipe existing permission data (excl. defaults for SendMessage).");
warn!("Taking a backup is advised.");
warn!("Continuing in 10 seconds...");
async_std::task::sleep(Duration::from_secs(10)).await;
let servers = db.col::<Document>("servers");
let mut cursor = servers.find(doc! {}, None).await.unwrap();
while let Some(Ok(mut document)) = cursor.next().await {
let id = document.get_str("_id").unwrap().to_string();
info!("Updating server {id}");
let mut update = doc! {};
// Try to pluck channel permission SendMessage (0x2)
// Structure of default_permissions used to be [server, channel]
let has_send = document
.get_array("default_permissions")
.map(|x| {
x.get(1)
.map(|x| x.as_i32().map(|x| (x as u32 & 0x2) == 0x2))
})
.ok()
.flatten()
.flatten()
.unwrap_or_default();
update.insert(
"default_permissions",
(*DEFAULT_PERMISSION_SERVER
// Remove Send Message permission if it wasn't originally granted
^ (if has_send {
0
} else {
Permission::SendMessage as u64
})) as i64,
);
if let Some(Bson::Document(mut roles)) = document.remove("roles") {
for role in roles.keys().cloned().collect::<Vec<String>>() {
if let Some(Bson::Document(role)) = roles.get_mut(role) {
role.insert(
"permissions",
doc! {
"a": 0_i64,
"d": 0_i64,
},
);
}
}
update.insert("roles", roles);
}
servers
.update_one(doc! { "_id": id }, doc! { "$set": update }, None)
.await
.unwrap();
}
let channels = db.col::<Document>("channels");
let mut cursor = channels.find(doc! {}, None).await.unwrap();
while let Some(Ok(document)) = cursor.next().await {
let id = document.get_str("_id").unwrap().to_string();
info!("Updating channel {id}");
let mut unset = doc! {
"permissions": 1_i32,
"role_permissions": 1_i32,
};
// Try to pluck channel permission SendMessage (0x2)
let has_send = document
.get_i32("default_permissions")
.map(|x| (x as u32 & 0x2) == 0x2)
.unwrap_or(true);
if has_send {
// Let parent permissions fall through.
unset.insert("default_permissions", 1_i32);
}
let mut update = doc! {
"$unset": unset
};
if !has_send {
// Block send message permission.
update.insert(
"$set",
doc! {
"default_permissions": {
"a": 0_i64,
"d": Permission::SendMessage as i64
}
},
);
}
channels
.update_one(doc! { "_id": id }, update, None)
.await
.unwrap();
}
}
if revision <= 14 {
info!("Running migration [revision 14 / 21-04-2022]: Split content into content and system fields.");
db.col::<Document>("messages")
.update_many(
doc! {
"content": {
"$type": "object"
}
},
doc! {
"$rename": {
"content": "system"
}
},
None,
)
.await
.unwrap();
}
// Need to migrate fields on attachments, change `user_id`, `object_id`, etc to `parent`.
// Reminder to update LATEST_REVISION when adding new migrations.
LATEST_REVISION
}

125
crates/quark/src/impl/mongo/autumn/attachment.rs Normal file
View File

@@ -0,0 +1,125 @@
use bson::Document;
use crate::models::attachment::File;
use crate::{AbstractAttachment, Error, Result};
use super::super::MongoDb;
static COL: &str = "attachments";
impl MongoDb {
pub async fn delete_many_attachments(&self, projection: Document) -> Result<()> {
self.col::<Document>(COL)
.update_many(
projection,
doc! {
"$set": {
"deleted": true
}
},
None,
)
.await
.map(|_| ())
.map_err(|_| Error::DatabaseError {
operation: "update_many",
with: "attachment",
})
}
}
#[async_trait]
impl AbstractAttachment for MongoDb {
async fn find_and_use_attachment(
&self,
attachment_id: &str,
tag: &str,
parent_type: &str,
parent_id: &str,
) -> Result<File> {
let key = format!("{}_id", parent_type);
match self
.find_one::<File>(
COL,
doc! {
"_id": attachment_id,
"tag": tag,
&key: {
"$exists": false
}
},
)
.await
{
Ok(file) => {
self.col::<Document>(COL)
.update_one(
doc! {
"_id": &file.id
},
doc! {
"$set": {
key: parent_id
}
},
None,
)
.await
.map_err(|_| Error::DatabaseError {
operation: "update_one",
with: "attachment",
})?;
Ok(file)
}
Err(Error::NotFound) => Err(Error::UnknownAttachment),
Err(error) => Err(error),
}
}
async fn insert_attachment(&self, attachment: &File) -> Result<()> {
self.insert_one(COL, attachment).await.map(|_| ())
}
async fn mark_attachment_as_reported(&self, id: &str) -> Result<()> {
self.col::<Document>(COL)
.update_one(
doc! {
"_id": id
},
doc! {
"$set": {
"reported": true
}
},
None,
)
.await
.map(|_| ())
.map_err(|_| Error::DatabaseError {
operation: "update_one",
with: "attachment",
})
}
async fn mark_attachment_as_deleted(&self, id: &str) -> Result<()> {
self.col::<Document>(COL)
.update_one(
doc! {
"_id": id
},
doc! {
"$set": {
"deleted": true
}
},
None,
)
.await
.map(|_| ())
.map_err(|_| Error::DatabaseError {
operation: "update_one",
with: "attachment",
})
}
}

323
crates/quark/src/impl/mongo/channels/channel.rs Normal file
View File

@@ -0,0 +1,323 @@
use std::collections::HashSet;
use bson::{Bson, Document};
use crate::models::channel::{Channel, FieldsChannel, PartialChannel};
use crate::r#impl::mongo::IntoDocumentPath;
use crate::{AbstractChannel, AbstractServer, Error, OverrideField, Result};
use super::super::MongoDb;
static COL: &str = "channels";
impl MongoDb {
pub async fn delete_associated_channel_objects(&self, id: Bson) -> Result<()> {
// Delete all invites to these channels.
self.col::<Document>("channel_invites")
.delete_many(
doc! {
"channel": &id
},
None,
)
.await
.map_err(|_| Error::DatabaseError {
operation: "delete_many",
with: "channel_invites",
})?;
// Delete unread message objects on channels.
self.col::<Document>("channel_unreads")
.delete_many(
doc! {
"_id.channel": &id
},
None,
)
.await
.map_err(|_| Error::DatabaseError {
operation: "delete_many",
with: "channel_unreads",
})
.map(|_| ())
// update many attachments with parent id
}
}
#[async_trait]
impl AbstractChannel for MongoDb {
async fn fetch_channel(&self, id: &str) -> Result<Channel> {
self.find_one_by_id(COL, id).await
}
async fn fetch_channels<'a>(&self, ids: &'a [String]) -> Result<Vec<Channel>> {
self.find(
COL,
doc! {
"_id": {
"$in": ids
}
},
)
.await
}
async fn insert_channel(&self, channel: &Channel) -> Result<()> {
self.insert_one(COL, channel).await.map(|_| ())
}
async fn update_channel(
&self,
id: &str,
channel: &PartialChannel,
remove: Vec<FieldsChannel>,
) -> Result<()> {
self.update_one_by_id(
COL,
id,
channel,
remove.iter().map(|x| x as &dyn IntoDocumentPath).collect(),
None,
)
.await
.map(|_| ())
}
async fn delete_channel(&self, channel: &Channel) -> Result<()> {
let id = channel.id().to_string();
let server_id = match channel {
Channel::TextChannel { server, .. } | Channel::VoiceChannel { server, .. } => {
Some(server)
}
_ => None,
};
// Delete invites and unreads.
self.delete_associated_channel_objects(Bson::String(id.to_string()))
.await?;
// Delete messages.
self.delete_bulk_messages(doc! {
"channel": &id
})
.await?;
// Remove from server object.
if let Some(server) = server_id {
let server = self.fetch_server(server).await?;
let mut update = doc! {
"$pull": {
"channels": &id
}
};
if let Some(sys) = &server.system_messages {
let mut unset = doc! {};
if let Some(cid) = &sys.user_joined {
if &id == cid {
unset.insert("system_messages.user_joined", 1_i32);
}
}
if let Some(cid) = &sys.user_left {
if &id == cid {
unset.insert("system_messages.user_left", 1_i32);
}
}
if let Some(cid) = &sys.user_kicked {
if &id == cid {
unset.insert("system_messages.user_kicked", 1_i32);
}
}
if let Some(cid) = &sys.user_banned {
if &id == cid {
unset.insert("system_messages.user_banned", 1_i32);
}
}
if !unset.is_empty() {
update.insert("$unset", unset);
}
}
self.col::<Document>("servers")
.update_one(
doc! {
"_id": server.id
},
update,
None,
)
.await
.map_err(|_| Error::DatabaseError {
operation: "update_one",
with: "servers",
})?;
}
// Delete associated attachments
self.delete_many_attachments(doc! {
"object_id": &id
})
.await?;
// Delete the channel itself
self.delete_one_by_id(COL, &id).await.map(|_| ())
}
async fn find_direct_messages(&self, user_id: &str) -> Result<Vec<Channel>> {
self.find(
COL,
doc! {
"$or": [
{
"$or": [
{
"channel_type": "DirectMessage"
},
{
"channel_type": "Group"
}
],
"recipients": user_id
},
{
"channel_type": "SavedMessages",
"user": user_id
}
]
},
)
.await
}
async fn find_saved_messages_channel(&self, user_id: &str) -> Result<Channel> {
self.find_one(
COL,
doc! {
"channel_type": "SavedMessages",
"user": user_id
},
)
.await
}
async fn find_direct_message_channel(&self, user_a: &str, user_b: &str) -> Result<Channel> {
self.find_one(
COL,
if user_a == user_b {
doc! {
"channel_type": "SavedMessages",
"user": user_a
}
} else {
doc! {
"channel_type": "DirectMessage",
"recipients": {
"$all": [ user_a, user_b ]
}
}
},
)
.await
}
async fn add_user_to_group(&self, channel: &str, user: &str) -> Result<()> {
self.col::<Document>(COL)
.update_one(
doc! {
"_id": channel
},
doc! {
"$push": {
"recipients": user
}
},
None,
)
.await
.map(|_| ())
.map_err(|_| Error::DatabaseError {
operation: "update_one",
with: "channel",
})
}
async fn remove_user_from_group(&self, channel: &str, user: &str) -> Result<()> {
self.col::<Document>(COL)
.update_one(
doc! {
"_id": channel
},
doc! {
"$pull": {
"recipients": user
}
},
None,
)
.await
.map(|_| ())
.map_err(|_| Error::DatabaseError {
operation: "update_one",
with: "channel",
})
}
async fn set_channel_role_permission(
&self,
channel: &str,
role: &str,
permissions: OverrideField,
) -> Result<()> {
self.col::<Document>(COL)
.update_one(
doc! { "_id": channel },
doc! {
"$set": {
"role_permissions.".to_owned() + role: permissions
}
},
None,
)
.await
.map(|_| ())
.map_err(|_| Error::DatabaseError {
operation: "update_one",
with: "channel",
})
}
async fn check_channels_exist(&self, channels: &HashSet<String>) -> Result<bool> {
let count = channels.len() as u64;
self.col::<Document>(COL)
.count_documents(
doc! {
"_id": {
"$in": channels.iter().cloned().collect::<Vec<String>>()
}
},
None,
)
.await
.map(|x| x == count)
.map_err(|_| Error::DatabaseError {
operation: "count_documents",
with: "channel",
})
}
}
impl IntoDocumentPath for FieldsChannel {
fn as_path(&self) -> Option<&'static str> {
Some(match self {
FieldsChannel::DefaultPermissions => "default_permissions",
FieldsChannel::Description => "description",
FieldsChannel::Icon => "icon",
})
}
}

31
crates/quark/src/impl/mongo/channels/channel_invite.rs Normal file
View File

@@ -0,0 +1,31 @@
use crate::models::Invite;
use crate::{AbstractChannelInvite, Result};
use super::super::MongoDb;
static COL: &str = "channel_invites";
#[async_trait]
impl AbstractChannelInvite for MongoDb {
async fn fetch_invite(&self, code: &str) -> Result<Invite> {
self.find_one_by_id(COL, code).await
}
async fn insert_invite(&self, invite: &Invite) -> Result<()> {
self.insert_one(COL, invite).await.map(|_| ())
}
async fn delete_invite(&self, code: &str) -> Result<()> {
self.delete_one_by_id(COL, code).await.map(|_| ())
}
async fn fetch_invites_for_server(&self, server: &str) -> Result<Vec<Invite>> {
self.find(
COL,
doc! {
"server": server
},
)
.await
}
}

104
crates/quark/src/impl/mongo/channels/channel_unread.rs Normal file
View File

@@ -0,0 +1,104 @@
use bson::Document;
use mongodb::options::UpdateOptions;
use ulid::Ulid;
use crate::models::channel_unread::ChannelUnread;
use crate::{AbstractChannelUnread, Error, Result};
use super::super::MongoDb;
static COL: &str = "channel_unreads";
#[async_trait]
impl AbstractChannelUnread for MongoDb {
async fn acknowledge_message(&self, channel: &str, user: &str, message: &str) -> Result<()> {
self.col::<Document>(COL)
.update_one(
doc! {
"_id.channel": channel,
"_id.user": user,
},
doc! {
"$unset": {
"mentions": 1_i32
},
"$set": {
"last_id": message
}
},
UpdateOptions::builder().upsert(true).build(),
)
.await
.map(|_| ())
.map_err(|_| Error::DatabaseError {
operation: "update_one",
with: "channel_unread",
})
}
async fn acknowledge_channels(&self, user: &str, channels: &[String]) -> Result<()> {
self.col::<Document>(COL)
.update_one(
doc! {
"_id.channel": {
"$in": channels
},
"_id.user": user,
},
doc! {
"$unset": {
"mentions": 1_i32
},
"$set": {
"last_id": Ulid::new().to_string()
}
},
UpdateOptions::builder().upsert(true).build(),
)
.await
.map(|_| ())
.map_err(|_| Error::DatabaseError {
operation: "update",
with: "channel_unread",
})
}
async fn add_mention_to_unread<'a>(
&self,
channel: &str,
user: &str,
ids: &[String],
) -> Result<()> {
self.col::<Document>(COL)
.update_one(
doc! {
"_id.channel": channel,
"_id.user": user,
},
doc! {
"$push": {
"mentions": {
"$each": ids
}
}
},
UpdateOptions::builder().upsert(true).build(),
)
.await
.map(|_| ())
.map_err(|_| Error::DatabaseError {
operation: "update_one",
with: "channel_unread",
})
}
async fn fetch_unreads(&self, user: &str) -> Result<Vec<ChannelUnread>> {
self.find(
COL,
doc! {
"_id.user": user
},
)
.await
}
}

285
crates/quark/src/impl/mongo/channels/message.rs Normal file
View File

@@ -0,0 +1,285 @@
use bson::{to_bson, Document};
use futures::try_join;
use mongodb::options::FindOptions;
use crate::models::message::{AppendMessage, Message, MessageSort, PartialMessage};
use crate::r#impl::mongo::DocumentId;
use crate::{AbstractMessage, Error, Result};
use super::super::MongoDb;
static COL: &str = "messages";
impl MongoDb {
pub async fn delete_bulk_messages(&self, projection: Document) -> Result<()> {
let mut for_attachments = projection.clone();
for_attachments.insert(
"attachment",
doc! {
"$exists": 1_i32
},
);
// Check if there are any attachments we need to delete.
let message_ids_with_attachments = self
.find_with_options::<_, DocumentId>(
COL,
for_attachments,
FindOptions::builder()
.projection(doc! { "_id": 1_i32 })
.build(),
)
.await?
.into_iter()
.map(|x| x.id)
.collect::<Vec<String>>();
// If we found any, mark them as deleted.
if !message_ids_with_attachments.is_empty() {
self.col::<Document>("attachments")
.update_many(
doc! {
"message_id": {
"$in": message_ids_with_attachments
}
},
doc! {
"$set": {
"deleted": true
}
},
None,
)
.await
.map_err(|_| Error::DatabaseError {
operation: "update_many",
with: "attachments",
})?;
}
// And then delete said messages.
self.col::<Document>(COL)
.delete_many(projection, None)
.await
.map(|_| ())
.map_err(|_| Error::DatabaseError {
operation: "delete_many",
with: "messages",
})
}
}
#[async_trait]
impl AbstractMessage for MongoDb {
async fn fetch_message(&self, id: &str) -> Result<Message> {
self.find_one_by_id(COL, id).await
}
async fn insert_message(&self, message: &Message) -> Result<()> {
self.insert_one(COL, message).await.map(|_| ())
}
async fn update_message(&self, id: &str, message: &PartialMessage) -> Result<()> {
self.update_one_by_id(COL, id, message, vec![], None)
.await
.map(|_| ())
}
async fn append_message(&self, id: &str, append: &AppendMessage) -> Result<()> {
let mut query = doc! {};
if let Some(embeds) = &append.embeds {
if !embeds.is_empty() {
query.insert(
"$push",
doc! {
"embeds": {
"$each": to_bson(embeds)
.map_err(|_| Error::DatabaseError {
operation: "to_bson",
with: "embeds"
})?
}
},
);
}
}
if query.is_empty() {
return Ok(());
}
self.col::<Document>(COL)
.update_one(
doc! {
"_id": id
},
query,
None,
)
.await
.map_err(|_| Error::DatabaseError {
operation: "update_one",
with: "message",
})
.map(|_| ())
}
async fn delete_message(&self, id: &str) -> Result<()> {
self.delete_bulk_messages(doc! {
"_id": id
})
.await
}
async fn delete_messages(&self, channel: &str, ids: Vec<String>) -> Result<()> {
self.delete_bulk_messages(doc! {
"channel": channel,
"_id": {
"$in": ids
}
})
.await
}
async fn fetch_messages(
&self,
channel: &str,
limit: Option<i64>,
before: Option<String>,
after: Option<String>,
sort: Option<MessageSort>,
nearby: Option<String>,
) -> Result<Vec<Message>> {
let limit = limit.unwrap_or(50);
Ok(if let Some(nearby) = nearby {
let (a, b) = try_join!(
self.find_with_options::<_, Message>(
COL,
doc! {
"channel": channel,
"_id": {
"$gte": &nearby
}
},
FindOptions::builder()
.limit(limit / 2 + 1)
.sort(doc! {
"_id": 1_i32
})
.build(),
),
self.find_with_options::<_, Message>(
COL,
doc! {
"channel": channel,
"_id": {
"$lt": &nearby
}
},
FindOptions::builder()
.limit(limit / 2)
.sort(doc! {
"_id": -1_i32
})
.build(),
)
)?;
[a, b].concat()
} else {
let mut query = doc! { "channel": channel };
if let Some(before) = before {
query.insert("_id", doc! { "$lt": before });
}
if let Some(after) = after {
query.insert("_id", doc! { "$gt": after });
}
let sort: i32 = if let MessageSort::Latest = sort.unwrap_or(MessageSort::Latest) {
-1
} else {
1
};
self.find_with_options::<_, Message>(
COL,
query,
FindOptions::builder()
.limit(limit)
.sort(doc! {
"_id": sort
})
.build(),
)
.await?
})
}
async fn search_messages(
&self,
channel: &str,
query: &str,
limit: Option<i64>,
before: Option<String>,
after: Option<String>,
sort: MessageSort,
) -> Result<Vec<Message>> {
let limit = limit.unwrap_or(50);
let mut filter = doc! {
"channel": channel,
"$text": {
"$search": query
}
};
if let Some(doc) = match (before, after) {
(Some(before), Some(after)) => Some(doc! {
"lt": before,
"gt": after
}),
(Some(before), _) => Some(doc! {
"lt": before
}),
(_, Some(after)) => Some(doc! {
"gt": after
}),
_ => None,
} {
filter.insert("_id", doc);
}
self.find_with_options(
COL,
filter,
FindOptions::builder()
.projection(if let MessageSort::Relevance = &sort {
doc! {
"score": {
"$meta": "textScore"
}
}
} else {
doc! {}
})
.limit(limit)
.sort(match &sort {
MessageSort::Relevance => doc! {
"score": {
"$meta": "textScore"
}
},
MessageSort::Latest => doc! {
"_id": -1_i32
},
MessageSort::Oldest => doc! {
"_id": 1_i32
},
})
.build(),
)
.await
}
}

250
crates/quark/src/impl/mongo/mod.rs Normal file
View File

@@ -0,0 +1,250 @@
use std::ops::Deref;
use bson::{to_document, Document};
use futures::StreamExt;
use mongodb::{
options::{FindOneOptions, FindOptions},
results::{DeleteResult, InsertOneResult, UpdateResult},
};
use rocket::serde::DeserializeOwned;
use serde::{Deserialize, Serialize};
use crate::{util::manipulation::prefix_keys, AbstractDatabase, Error, Result};
pub mod admin {
pub mod migrations;
}
pub mod autumn {
pub mod attachment;
}
pub mod channels {
pub mod channel;
pub mod channel_invite;
pub mod channel_unread;
pub mod message;
}
pub mod servers {
pub mod server;
pub mod server_ban;
pub mod server_member;
}
pub mod users {
pub mod bot;
pub mod user;
pub mod user_settings;
}
#[derive(Debug, Clone)]
pub struct MongoDb(pub mongodb::Client);
impl AbstractDatabase for MongoDb {}
impl Deref for MongoDb {
type Target = mongodb::Client;
fn deref(&self) -> &Self::Target {
&self.0
}
}
impl MongoDb {
pub fn db(&self) -> mongodb::Database {
self.database("revolt")
}
pub fn col<T>(&self, collection: &str) -> mongodb::Collection<T> {
self.db().collection(collection)
}
async fn insert_one<T: Serialize>(
&self,
collection: &'static str,
document: T,
) -> Result<InsertOneResult> {
self.col::<T>(collection)
.insert_one(document, None)
.await
.map_err(|_| Error::DatabaseError {
operation: "insert_one",
with: collection,
})
}
async fn find_with_options<O, T: DeserializeOwned + Unpin + Send + Sync>(
&self,
collection: &'static str,
projection: Document,
options: O,
) -> Result<Vec<T>>
where
O: Into<Option<FindOptions>>,
{
Ok(self
.col::<T>(collection)
.find(projection, options)
.await
.map_err(|_| Error::DatabaseError {
operation: "find",
with: collection,
})?
.filter_map(|s| async { s.ok() })
.collect::<Vec<T>>()
.await)
}
async fn find<T: DeserializeOwned + Unpin + Send + Sync>(
&self,
collection: &'static str,
projection: Document,
) -> Result<Vec<T>> {
self.find_with_options(collection, projection, None).await
}
async fn find_one_with_options<O, T: DeserializeOwned + Unpin + Send + Sync>(
&self,
collection: &'static str,
projection: Document,
options: O,
) -> Result<T>
where
O: Into<Option<FindOneOptions>>,
{
self.col::<T>(collection)
.find_one(projection, options)
.await
.map_err(|_| Error::DatabaseError {
operation: "find_one",
with: collection,
})?
.ok_or(Error::NotFound)
}
async fn find_one<T: DeserializeOwned + Unpin + Send + Sync>(
&self,
collection: &'static str,
projection: Document,
) -> Result<T> {
self.find_one_with_options(collection, projection, None)
.await
}
async fn find_one_by_id<T: DeserializeOwned + Unpin + Send + Sync>(
&self,
collection: &'static str,
id: &str,
) -> Result<T> {
self.find_one(
collection,
doc! {
"_id": id
},
)
.await
}
async fn update_one<P, T: Serialize>(
&self,
collection: &'static str,
projection: Document,
partial: T,
remove: Vec<&dyn IntoDocumentPath>,
prefix: P,
) -> Result<UpdateResult>
where
P: Into<Option<String>>,
{
let prefix = prefix.into();
let mut unset = doc! {};
for field in remove {
if let Some(path) = field.as_path() {
if let Some(prefix) = &prefix {
unset.insert(prefix.to_owned() + path, 1_i32);
} else {
unset.insert(path, 1_i32);
}
}
}
let query = doc! {
"$unset": unset,
"$set": if let Some(prefix) = &prefix {
to_document(&prefix_keys(&partial, prefix))
} else {
to_document(&partial)
}.map_err(|_| Error::DatabaseError {
operation: "to_document",
with: collection
})?
};
self.col::<Document>(collection)
.update_one(projection, query, None)
.await
.map_err(|_| Error::DatabaseError {
operation: "update_one",
with: collection,
})
}
async fn update_one_by_id<P, T: Serialize>(
&self,
collection: &'static str,
id: &str,
partial: T,
remove: Vec<&dyn IntoDocumentPath>,
prefix: P,
) -> Result<UpdateResult>
where
P: Into<Option<String>>,
{
self.update_one(
collection,
doc! {
"_id": id
},
partial,
remove,
prefix,
)
.await
}
async fn delete_one(
&self,
collection: &'static str,
projection: Document,
) -> Result<DeleteResult> {
self.col::<Document>(collection)
.delete_one(projection, None)
.await
.map_err(|_| Error::DatabaseError {
operation: "delete_one",
with: collection,
})
}
async fn delete_one_by_id(&self, collection: &'static str, id: &str) -> Result<DeleteResult> {
self.delete_one(
collection,
doc! {
"_id": id
},
)
.await
}
}
#[derive(Deserialize)]
pub struct DocumentId {
#[serde(rename = "_id")]
pub id: String,
}
pub trait IntoDocumentPath: Send + Sync {
fn as_path(&self) -> Option<&'static str>;
}

228
crates/quark/src/impl/mongo/servers/server.rs Normal file
View File

@@ -0,0 +1,228 @@
use bson::{to_document, Bson, Document};
use crate::models::server::{FieldsRole, FieldsServer, PartialRole, PartialServer, Role, Server};
use crate::r#impl::mongo::IntoDocumentPath;
use crate::{AbstractServer, Database, Error, Result};
use super::super::MongoDb;
static COL: &str = "servers";
impl MongoDb {
pub async fn delete_associated_server_objects(&self, server: &Server) -> Result<()> {
// Check if there are any attachments we need to delete.
self.delete_bulk_messages(doc! {
"channel": {
"$in": &server.channels
}
})
.await?;
// Delete all channels.
self.col::<Document>("channels")
.delete_many(
doc! {
"server": &server.id
},
None,
)
.await
.map_err(|_| Error::DatabaseError {
operation: "delete_many",
with: "channels",
})?;
// Delete any associated objects, e.g. unreads and invites.
self.delete_associated_channel_objects(Bson::Document(doc! { "$in": &server.channels }))
.await?;
// Delete members and bans.
for with in &["server_members", "server_bans"] {
self.col::<Document>(with)
.delete_many(
doc! {
"_id.server": &server.id
},
None,
)
.await
.map_err(|_| Error::DatabaseError {
operation: "delete_many",
with,
})?;
}
// Update many attachments with parent id.
self.delete_many_attachments(doc! {
"object_id": &server.id
})
.await?;
Ok(())
}
}
#[async_trait]
impl AbstractServer for MongoDb {
async fn fetch_server(&self, id: &str) -> Result<Server> {
self.find_one_by_id(COL, id).await
}
async fn fetch_servers<'a>(&self, ids: &'a [String]) -> Result<Vec<Server>> {
self.find(
COL,
doc! {
"_id": {
"$in": ids
}
},
)
.await
}
async fn insert_server(&self, server: &Server) -> Result<()> {
self.insert_one(COL, server).await.map(|_| ())
}
async fn update_server(
&self,
id: &str,
server: &PartialServer,
remove: Vec<FieldsServer>,
) -> Result<()> {
self.update_one_by_id(
COL,
id,
server,
remove.iter().map(|x| x as &dyn IntoDocumentPath).collect(),
None,
)
.await
.map(|_| ())
}
async fn delete_server(&self, server: &Server) -> Result<()> {
self.delete_associated_server_objects(server).await?;
self.delete_one_by_id(COL, &server.id).await.map(|_| ())
}
async fn insert_role(&self, server_id: &str, role_id: &str, role: &Role) -> Result<()> {
self.col::<Database>(COL)
.update_one(
doc! {
"_id": server_id
},
doc! {
"$set": {
"roles.".to_owned() + role_id: to_document(role)
.map_err(|_| Error::DatabaseError {
operation: "to_document",
with: "role"
})?
}
},
None,
)
.await
.map(|_| ())
.map_err(|_| Error::DatabaseError {
operation: "update_one",
with: "server",
})
}
async fn update_role(
&self,
server_id: &str,
role_id: &str,
role: &PartialRole,
remove: Vec<FieldsRole>,
) -> Result<()> {
self.update_one_by_id(
COL,
server_id,
role,
remove.iter().map(|x| x as &dyn IntoDocumentPath).collect(),
"roles.".to_owned() + role_id + ".",
)
.await
.map(|_| ())
}
async fn delete_role(&self, server_id: &str, role_id: &str) -> Result<()> {
self.col::<Document>("server_members")
.update_many(
doc! {
"_id.server": server_id
},
doc! {
"$pull": {
"roles": &role_id
}
},
None,
)
.await
.map_err(|_| Error::DatabaseError {
operation: "update_many",
with: "server_members",
})?;
self.col::<Document>("channels")
.update_one(
doc! {
"server": server_id
},
doc! {
"$unset": {
"role_permissions.".to_owned() + role_id: 1_i32
}
},
None,
)
.await
.map_err(|_| Error::DatabaseError {
operation: "update_one",
with: "channels",
})?;
self.col::<Document>("servers")
.update_one(
doc! {
"_id": server_id
},
doc! {
"$unset": {
"roles.".to_owned() + role_id: 1_i32
}
},
None,
)
.await
.map_err(|_| Error::DatabaseError {
operation: "update_one",
with: "servers",
})
.map(|_| ())
}
}
impl IntoDocumentPath for FieldsServer {
fn as_path(&self) -> Option<&'static str> {
Some(match self {
FieldsServer::Banner => "banner",
FieldsServer::Categories => "categories",
FieldsServer::Description => "description",
FieldsServer::Icon => "icon",
FieldsServer::SystemMessages => "system_messages",
})
}
}
impl IntoDocumentPath for FieldsRole {
fn as_path(&self) -> Option<&'static str> {
Some(match self {
FieldsRole::Colour => "colour",
})
}
}

47
crates/quark/src/impl/mongo/servers/server_ban.rs Normal file
View File

@@ -0,0 +1,47 @@
use crate::models::server_member::MemberCompositeKey;
use crate::models::ServerBan;
use crate::{AbstractServerBan, Result};
use super::super::MongoDb;
static COL: &str = "server_bans";
#[async_trait]
impl AbstractServerBan for MongoDb {
async fn fetch_ban(&self, server: &str, user: &str) -> Result<ServerBan> {
self.find_one(
COL,
doc! {
"_id.server": server,
"_id.user": user
},
)
.await
}
async fn fetch_bans(&self, server: &str) -> Result<Vec<ServerBan>> {
self.find(
COL,
doc! {
"_id.server": server
},
)
.await
}
async fn insert_ban(&self, ban: &ServerBan) -> Result<()> {
self.insert_one(COL, ban).await.map(|_| ())
}
async fn delete_ban(&self, id: &MemberCompositeKey) -> Result<()> {
self.delete_one(
COL,
doc! {
"_id.server": &id.server,
"_id.user": &id.user
},
)
.await
.map(|_| ())
}
}

134
crates/quark/src/impl/mongo/servers/server_member.rs Normal file
View File

@@ -0,0 +1,134 @@
use bson::Document;
use crate::models::server_member::{FieldsMember, Member, MemberCompositeKey, PartialMember};
use crate::r#impl::mongo::IntoDocumentPath;
use crate::{AbstractServerMember, Error, Result};
use super::super::MongoDb;
static COL: &str = "server_members";
#[async_trait]
impl AbstractServerMember for MongoDb {
async fn fetch_member(&self, server: &str, user: &str) -> Result<Member> {
self.find_one(
COL,
doc! {
"_id.server": server,
"_id.user": user
},
)
.await
}
async fn insert_member(&self, member: &Member) -> Result<()> {
self.insert_one(COL, member).await.map(|_| ())
}
async fn update_member(
&self,
id: &MemberCompositeKey,
member: &PartialMember,
remove: Vec<FieldsMember>,
) -> Result<()> {
self.update_one(
COL,
doc! {
"_id.server": &id.server,
"_id.user": &id.user
},
member,
remove.iter().map(|x| x as &dyn IntoDocumentPath).collect(),
None,
)
.await
.map(|_| ())
}
async fn delete_member(&self, id: &MemberCompositeKey) -> Result<()> {
self.delete_one(
COL,
doc! {
"_id.server": &id.server,
"_id.user": &id.user
},
)
.await
.map(|_| ())
}
async fn fetch_all_members<'a>(&self, server: &str) -> Result<Vec<Member>> {
self.find(
COL,
doc! {
"_id.server": server
},
)
.await
}
async fn fetch_all_memberships<'a>(&self, user: &str) -> Result<Vec<Member>> {
self.find(
COL,
doc! {
"_id.user": user
},
)
.await
}
async fn fetch_members<'a>(&self, server: &str, ids: &'a [String]) -> Result<Vec<Member>> {
self.find(
COL,
doc! {
"_id.server": server,
"_id.user": {
"$in": ids
}
},
)
.await
}
async fn fetch_member_count(&self, server: &str) -> Result<usize> {
self.col::<Document>(COL)
.count_documents(
doc! {
"_id.server": server
},
None,
)
.await
.map(|c| c as usize)
.map_err(|_| Error::DatabaseError {
operation: "count_documents",
with: "server_members",
})
}
async fn fetch_server_count(&self, user: &str) -> Result<usize> {
self.col::<Document>(COL)
.count_documents(
doc! {
"_id.user": user
},
None,
)
.await
.map(|c| c as usize)
.map_err(|_| Error::DatabaseError {
operation: "count_documents",
with: "server_members",
})
}
}
impl IntoDocumentPath for FieldsMember {
fn as_path(&self) -> Option<&'static str> {
Some(match self {
FieldsMember::Avatar => "avatar",
FieldsMember::Nickname => "nickname",
FieldsMember::Roles => "roles",
})
}
}

68
crates/quark/src/impl/mongo/users/bot.rs Normal file
View File

@@ -0,0 +1,68 @@
use crate::models::bot::{Bot, FieldsBot, PartialBot};
use crate::r#impl::mongo::IntoDocumentPath;
use crate::{AbstractBot, Result};
use super::super::MongoDb;
static COL: &str = "bots";
#[async_trait]
impl AbstractBot for MongoDb {
async fn fetch_bot(&self, id: &str) -> Result<Bot> {
self.find_one_by_id(COL, id).await
}
async fn fetch_bot_by_token(&self, token: &str) -> Result<Bot> {
self.find_one(
COL,
doc! {
"token": token
},
)
.await
}
async fn insert_bot(&self, bot: &Bot) -> Result<()> {
self.insert_one(COL, &bot).await.map(|_| ())
}
async fn update_bot(&self, id: &str, bot: &PartialBot, remove: Vec<FieldsBot>) -> Result<()> {
self.update_one_by_id(
COL,
id,
bot,
remove.iter().map(|x| x as &dyn IntoDocumentPath).collect(),
None,
)
.await
.map(|_| ())
}
async fn delete_bot(&self, id: &str) -> Result<()> {
self.delete_one_by_id(COL, id).await.map(|_| ())
}
async fn fetch_bots_by_user(&self, user_id: &str) -> Result<Vec<Bot>> {
self.find(
COL,
doc! {
"owner": user_id
},
)
.await
}
async fn get_number_of_bots_by_user(&self, user_id: &str) -> Result<usize> {
// ! FIXME: move this to generic?
self.fetch_bots_by_user(user_id).await.map(|x| x.len())
}
}
impl IntoDocumentPath for FieldsBot {
fn as_path(&self) -> Option<&'static str> {
match self {
FieldsBot::InteractionsURL => Some("interactions_url"),
FieldsBot::Token => None,
}
}
}

319
crates/quark/src/impl/mongo/users/user.rs Normal file
View File

@@ -0,0 +1,319 @@
use bson::Document;
use futures::StreamExt;
use mongodb::options::{Collation, CollationStrength, FindOneOptions, FindOptions};
use crate::models::user::{FieldsUser, PartialUser, RelationshipStatus, User};
use crate::r#impl::mongo::IntoDocumentPath;
use crate::{AbstractUser, Error, Result};
use super::super::MongoDb;
lazy_static! {
static ref FIND_USERNAME_OPTIONS: FindOneOptions = FindOneOptions::builder()
.collation(
Collation::builder()
.locale("en")
.strength(CollationStrength::Secondary)
.build()
)
.build();
}
static COL: &str = "users";
#[async_trait]
impl AbstractUser for MongoDb {
async fn fetch_user(&self, id: &str) -> Result<User> {
self.find_one_by_id(COL, id).await
}
async fn fetch_user_by_username(&self, username: &str) -> Result<User> {
self.find_one_with_options(
COL,
doc! {
"username": username
},
FIND_USERNAME_OPTIONS.clone(),
)
.await
}
async fn fetch_user_by_token(&self, token: &str) -> Result<User> {
let session = self
.col::<Document>("sessions")
.find_one(
doc! {
"token": token
},
None,
)
.await
.map_err(|_| Error::DatabaseError {
operation: "find_one",
with: "sessions",
})?
.ok_or(Error::InvalidSession)?;
self.fetch_user(session.get_str("user_id").unwrap()).await
}
async fn insert_user(&self, user: &User) -> Result<()> {
self.insert_one(COL, user).await.map(|_| ())
}
async fn update_user(
&self,
id: &str,
user: &PartialUser,
remove: Vec<FieldsUser>,
) -> Result<()> {
self.update_one_by_id(
COL,
id,
user,
remove.iter().map(|x| x as &dyn IntoDocumentPath).collect(),
None,
)
.await
.map(|_| ())
}
async fn delete_user(&self, id: &str) -> Result<()> {
self.delete_one_by_id(COL, id).await.map(|_| ())
}
async fn fetch_users<'a>(&self, ids: &'a [String]) -> Result<Vec<User>> {
let mut cursor = self
.col::<User>(COL)
.find(
doc! {
"_id": {
"$in": ids
}
},
None,
)
.await
.map_err(|_| Error::DatabaseError {
operation: "find",
with: "users",
})?;
let mut users = vec![];
while let Some(Ok(user)) = cursor.next().await {
users.push(user);
}
Ok(users)
}
async fn is_username_taken(&self, username: &str) -> Result<bool> {
// ! FIXME: move this up to generic
match self.fetch_user_by_username(username).await {
Ok(_) => Ok(true),
Err(Error::NotFound) => Ok(false),
Err(error) => Err(error),
}
}
async fn fetch_mutual_user_ids(&self, user_a: &str, user_b: &str) -> Result<Vec<String>> {
Ok(self
.col::<Document>(COL)
.find(
doc! {
"$and": [
{ "relations": { "$elemMatch": { "_id": &user_a, "status": "Friend" } } },
{ "relations": { "$elemMatch": { "_id": &user_b, "status": "Friend" } } }
]
},
FindOptions::builder().projection(doc! { "_id": 1 }).build(),
)
.await
.map_err(|_| Error::DatabaseError {
operation: "find",
with: "users",
})?
.filter_map(|s| async { s.ok() })
.collect::<Vec<Document>>()
.await
.into_iter()
.filter_map(|x| x.get_str("_id").ok().map(|x| x.to_string()))
.collect::<Vec<String>>())
}
async fn fetch_mutual_channel_ids(&self, user_a: &str, user_b: &str) -> Result<Vec<String>> {
Ok(self
.col::<Document>("channels")
.find(
doc! {
"channel_type": {
"$in": ["Group", "DirectMessage"]
},
"recipients": {
"$all": [ user_a, user_b ]
}
},
FindOptions::builder().projection(doc! { "_id": 1 }).build(),
)
.await
.map_err(|_| Error::DatabaseError {
operation: "find",
with: "channels",
})?
.filter_map(|s| async { s.ok() })
.collect::<Vec<Document>>()
.await
.into_iter()
.filter_map(|x| x.get_str("_id").ok().map(|x| x.to_string()))
.collect::<Vec<String>>())
}
async fn fetch_mutual_server_ids(&self, user_a: &str, user_b: &str) -> Result<Vec<String>> {
Ok(self
.col::<Document>("server_members")
.aggregate(
vec![
doc! {
"$match": {
"_id.user": user_a
}
},
doc! {
"$lookup": {
"from": "server_members",
"as": "members",
"let": {
"server": "$_id.server"
},
"pipeline": [
{
"$match": {
"$expr": {
"$and": [
{ "$eq": [ "$_id.user", user_b ] },
{ "$eq": [ "$_id.server", "$$server" ] }
]
}
}
}
]
}
},
doc! {
"$match": {
"members": {
"$size": 1_i32
}
}
},
doc! {
"$project": {
"_id": "$_id.server"
}
},
],
None,
)
.await
.map_err(|_| Error::DatabaseError {
operation: "aggregate",
with: "server_members",
})?
.filter_map(|s| async { s.ok() })
.collect::<Vec<Document>>()
.await
.into_iter()
.filter_map(|x| x.get_str("_id").ok().map(|i| i.to_string()))
.collect::<Vec<String>>())
}
async fn set_relationship(
&self,
user_id: &str,
target_id: &str,
relationship: &RelationshipStatus,
) -> Result<()> {
if let RelationshipStatus::None = relationship {
return self.pull_relationship(user_id, target_id).await;
}
self.col::<Document>(COL)
.update_one(
doc! {
"_id": user_id
},
vec![doc! {
"$set": {
"relations": {
"$concatArrays": [
{
"$ifNull": [
{
"$filter": {
"input": "$relations",
"cond": {
"$ne": [
"$$this._id",
target_id
]
}
}
},
[]
]
},
[
{
"_id": target_id,
"status": format!("{:?}", relationship)
}
]
]
}
}
}],
None,
)
.await
.map(|_| ())
.map_err(|_| Error::DatabaseError {
operation: "update_one",
with: "user",
})
}
async fn pull_relationship(&self, user_id: &str, target_id: &str) -> Result<()> {
self.col::<Document>(COL)
.update_one(
doc! {
"_id": user_id
},
doc! {
"$pull": {
"relations": {
"_id": target_id
}
}
},
None,
)
.await
.map(|_| ())
.map_err(|_| Error::DatabaseError {
operation: "update_one",
with: "user",
})
}
}
impl IntoDocumentPath for FieldsUser {
fn as_path(&self) -> Option<&'static str> {
Some(match self {
FieldsUser::Avatar => "avatar",
FieldsUser::ProfileBackground => "profile.background",
FieldsUser::ProfileContent => "profile.content",
FieldsUser::StatusPresence => "status.presence",
FieldsUser::StatusText => "status.text",
})
}
}

62
crates/quark/src/impl/mongo/users/user_settings.rs Normal file
View File

@@ -0,0 +1,62 @@
use bson::{to_bson, Document};
use mongodb::options::{FindOneOptions, UpdateOptions};
use crate::models::UserSettings;
use crate::{AbstractUserSettings, Error, Result};
use super::super::MongoDb;
static COL: &str = "user_settings";
#[async_trait]
impl AbstractUserSettings for MongoDb {
async fn fetch_user_settings(&'_ self, id: &str, filter: &'_ [String]) -> Result<UserSettings> {
let mut projection = doc! {
"_id": 0,
};
for key in filter {
projection.insert(key, 1);
}
self.find_one_with_options(
COL,
doc! {
"_id": id
},
FindOneOptions::builder().projection(projection).build(),
)
.await
}
async fn set_user_settings(&self, id: &str, settings: &UserSettings) -> Result<()> {
let mut set = doc! {};
for (key, data) in settings {
set.insert(
key,
vec![to_bson(&data.0).unwrap(), to_bson(&data.1).unwrap()],
);
}
self.col::<Document>(COL)
.update_one(
doc! {
"_id": id
},
doc! {
"$set": set
},
UpdateOptions::builder().upsert(true).build(),
)
.await
.map(|_| ())
.map_err(|_| Error::DatabaseError {
operation: "update_one",
with: "user_settings",
})
}
async fn delete_user_settings(&self, id: &str) -> Result<()> {
self.delete_one_by_id(COL, id).await.map(|_| ())
}
}