From 1e0b34367bf71471222095e56356b19aeb226513 Mon Sep 17 00:00:00 2001
From: strawberry <strawberry@puppygock.gay>
Date: Sun, 21 Apr 2024 19:37:52 -0400
Subject: [PATCH] add users query command, initial fsck admin command

Signed-off-by: strawberry <strawberry@puppygock.gay>
---
 src/service/admin/fsck.rs               | 19 ------------------
 src/service/admin/fsck/fsck_commands.rs | 26 +++++++++++++++++++++++++
 src/service/admin/fsck/mod.rs           | 19 ++++++++++++++++++
 src/service/admin/mod.rs                |  6 ++++++
 src/service/admin/query/mod.rs          | 17 ++++++++++++++--
 src/service/admin/query/users.rs        | 25 ++++++++++++++++++++++++
 6 files changed, 91 insertions(+), 21 deletions(-)
 delete mode 100644 src/service/admin/fsck.rs
 create mode 100644 src/service/admin/fsck/fsck_commands.rs
 create mode 100644 src/service/admin/fsck/mod.rs
 create mode 100644 src/service/admin/query/users.rs

diff --git a/src/service/admin/fsck.rs b/src/service/admin/fsck.rs
deleted file mode 100644
index 9e9b64a1..00000000
--- a/src/service/admin/fsck.rs
+++ /dev/null
@@ -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!()
-		},
-	}
-}
diff --git a/src/service/admin/fsck/fsck_commands.rs b/src/service/admin/fsck/fsck_commands.rs
new file mode 100644
index 00000000..46a04c3c
--- /dev/null
+++ b/src/service/admin/fsck/fsck_commands.rs
@@ -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()))
+}
diff --git a/src/service/admin/fsck/mod.rs b/src/service/admin/fsck/mod.rs
new file mode 100644
index 00000000..e618fdc5
--- /dev/null
+++ b/src/service/admin/fsck/mod.rs
@@ -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?,
+	})
+}
diff --git a/src/service/admin/mod.rs b/src/service/admin/mod.rs
index 0f361a4c..f1a63652 100644
--- a/src/service/admin/mod.rs
+++ b/src/service/admin/mod.rs
@@ -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)
diff --git a/src/service/admin/query/mod.rs b/src/service/admin/query/mod.rs
index f1473a46..a1d7c53a 100644
--- a/src/service/admin/query/mod.rs
+++ b/src/service/admin/query/mod.rs
@@ -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?,
 	})
 }
diff --git a/src/service/admin/query/users.rs b/src/service/admin/query/users.rs
new file mode 100644
index 00000000..818ff6d6
--- /dev/null
+++ b/src/service/admin/query/users.rs
@@ -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
+				),
+			))
+		},
+	}
+}