add users query command, initial fsck admin command

Signed-off-by: strawberry <strawberry@puppygock.gay>
This commit is contained in:
strawberry 2024-04-21 19:37:52 -04:00 committed by June
parent affd063df6
commit 1e0b34367b
6 changed files with 91 additions and 21 deletions

View file

@ -1,19 +0,0 @@
use clap::Subcommand;
use ruma::events::room::message::RoomMessageEventContent;
use crate::Result;
#[cfg_attr(test, derive(Debug))]
#[derive(Subcommand)]
pub(crate) enum FsckCommand {
Register,
}
#[allow(dead_code)]
pub(crate) async fn fsck(command: FsckCommand, _body: Vec<&str>) -> Result<RoomMessageEventContent> {
match command {
FsckCommand::Register => {
todo!()
},
}
}

View file

@ -0,0 +1,26 @@
use ruma::events::room::message::RoomMessageEventContent;
use crate::{services, Result};
/// Uses the iterator in `src/database/key_value/users.rs` to iterator over
/// every user in our database (remote and local). Reports total count, any
/// errors if there were any, etc
pub(super) async fn check_all_users(_body: Vec<&str>) -> Result<RoomMessageEventContent> {
let timer = tokio::time::Instant::now();
let results = services().users.db.iter();
let query_time = timer.elapsed();
let users = results.collect::<Vec<_>>();
let total = users.len();
let err_count = users.iter().filter(|user| user.is_err()).count();
let ok_count = users.iter().filter(|user| user.is_ok()).count();
let message = format!(
"Database query completed in {query_time:?}:\n\n```\nTotal entries: {:?}\nFailure/Invalid user count: \
{:?}\nSuccess/Valid user count: {:?}```",
total, err_count, ok_count
);
Ok(RoomMessageEventContent::notice_html(message, String::new()))
}

View file

@ -0,0 +1,19 @@
use clap::Subcommand;
use ruma::events::room::message::RoomMessageEventContent;
use self::fsck_commands::check_all_users;
use crate::Result;
pub(crate) mod fsck_commands;
#[cfg_attr(test, derive(Debug))]
#[derive(Subcommand)]
pub(crate) enum FsckCommand {
CheckAllUsers,
}
pub(crate) async fn process(command: FsckCommand, body: Vec<&str>) -> Result<RoomMessageEventContent> {
Ok(match command {
FsckCommand::CheckAllUsers => check_all_users(body).await?,
})
}

View file

@ -26,6 +26,7 @@ use serde_json::value::to_raw_value;
use tokio::sync::Mutex;
use tracing::{error, warn};
use self::fsck::FsckCommand;
use super::pdu::PduBuilder;
use crate::{
service::admin::{
@ -82,6 +83,10 @@ enum AdminCommand {
#[command(subcommand)]
/// - Query all the database getters and iterators
Query(QueryCommand),
#[command(subcommand)]
/// - Query all the database getters and iterators
Fsck(FsckCommand),
}
#[derive(Debug)]
@ -283,6 +288,7 @@ impl Service {
AdminCommand::Server(command) => server::process(command, body).await?,
AdminCommand::Debug(command) => debug::process(command, body).await?,
AdminCommand::Query(command) => query::process(command, body).await?,
AdminCommand::Fsck(command) => fsck::process(command, body).await?,
};
Ok(reply_message_content)

View file

@ -4,6 +4,7 @@ pub(crate) mod globals;
pub(crate) mod presence;
pub(crate) mod room_alias;
pub(crate) mod sending;
pub(crate) mod users;
use clap::Subcommand;
use ruma::{
@ -13,7 +14,7 @@ use ruma::{
use self::{
account_data::account_data, appservice::appservice, globals::globals, presence::presence, room_alias::room_alias,
sending::sending,
sending::sending, users::users,
};
use crate::Result;
@ -41,9 +42,13 @@ pub(crate) enum QueryCommand {
#[command(subcommand)]
Globals(Globals),
/// - globals.rs iterators and getters
/// - sending.rs iterators and getters
#[command(subcommand)]
Sending(Sending),
/// - users.rs iterators and getters
#[command(subcommand)]
Users(Users),
}
#[cfg_attr(test, derive(Debug))]
@ -149,6 +154,13 @@ pub(crate) enum Sending {
},
}
#[cfg_attr(test, derive(Debug))]
#[derive(Subcommand)]
/// All the getters and iterators from src/database/key_value/users.rs
pub(crate) enum Users {
Iter,
}
/// Processes admin query commands
pub(crate) async fn process(command: QueryCommand, _body: Vec<&str>) -> Result<RoomMessageEventContent> {
Ok(match command {
@ -158,5 +170,6 @@ pub(crate) async fn process(command: QueryCommand, _body: Vec<&str>) -> Result<R
QueryCommand::RoomAlias(command) => room_alias(command).await?,
QueryCommand::Globals(command) => globals(command).await?,
QueryCommand::Sending(command) => sending(command).await?,
QueryCommand::Users(command) => users(command).await?,
})
}

View file

@ -0,0 +1,25 @@
use ruma::events::room::message::RoomMessageEventContent;
use super::Users;
use crate::{services, Result};
/// All the getters and iterators in key_value/users.rs
pub(super) async fn users(subcommand: Users) -> Result<RoomMessageEventContent> {
match subcommand {
Users::Iter => {
let timer = tokio::time::Instant::now();
let results = services().users.db.iter();
let query_time = timer.elapsed();
let users = results.collect::<Vec<_>>();
Ok(RoomMessageEventContent::text_html(
format!("Query completed in {query_time:?}:\n\n```\n{:?}```", users),
format!(
"<p>Query completed in {query_time:?}:</p>\n<pre><code>{:?}\n</code></pre>",
users
),
))
},
}
}