initial support for querying database getters and iterators via admin cmd

Signed-off-by: strawberry <strawberry@puppygock.gay>

src/service/admin/mod.rs

@@ -26,6 +26,7 @@
 use tokio::sync::Mutex;
 use tracing::{error, warn};
 
+use self::query::QueryCommand;
 use super::pdu::PduBuilder;
 use crate::{
 	service::admin::{
@@ -39,6 +40,7 @@
 pub(crate) mod debug;
 pub(crate) mod federation;
 pub(crate) mod media;
+pub(crate) mod query;
 pub(crate) mod room;
 pub(crate) mod room_alias;
 pub(crate) mod room_directory;
@@ -77,11 +79,12 @@ enum AdminCommand {
 	Media(MediaCommand),
 
 	#[command(subcommand)]
-	// TODO: should i split out debug commands to a separate thing? the
-	// debug commands seem like they could fit in the other categories fine
-	// this is more like a "miscellaneous" category than a debug one
 	/// - Commands for debugging things
 	Debug(DebugCommand),
+
+	#[command(subcommand)]
+	/// - Query all the database getters and iterators
+	Query(QueryCommand),
 }
 
 #[derive(Debug)]
@@ -249,10 +252,25 @@ fn parse_admin_command(&self, command_line: &str) -> Result<AdminCommand, String
 		}
 
 		// Backwards compatibility with `register_appservice`-style commands
-		let command_with_dashes;
+		let command_with_dashes_argv1;
 		if argv.len() > 1 && argv[1].contains('_') {
-			command_with_dashes = argv[1].replace('_', "-");
-			argv[1] = &command_with_dashes;
+			command_with_dashes_argv1 = argv[1].replace('_', "-");
+			argv[1] = &command_with_dashes_argv1;
+		}
+
+		// Backwards compatibility with `register_appservice`-style commands
+		let command_with_dashes_argv2;
+		if argv.len() > 2 && argv[2].contains('_') {
+			command_with_dashes_argv2 = argv[2].replace('_', "-");
+			argv[2] = &command_with_dashes_argv2;
+		}
+
+		// if the user is using the `query` command (argv[1]), replace the database
+		// function/table calls with underscores to match the codebase
+		let command_with_dashes_argv3;
+		if argv.len() > 3 && argv[1].eq("query") {
+			command_with_dashes_argv3 = argv[3].replace('_', "-");
+			argv[3] = &command_with_dashes_argv3;
 		}
 
 		AdminCommand::try_parse_from(argv).map_err(|error| error.to_string())
@@ -267,6 +285,7 @@ async fn process_admin_command(&self, command: AdminCommand, body: Vec<&str>) ->
 			AdminCommand::Federation(command) => federation::process(command, body).await?,
 			AdminCommand::Server(command) => server::process(command, body).await?,
 			AdminCommand::Debug(command) => debug::process(command, body).await?,
+			AdminCommand::Query(command) => query::process(command, body).await?,
 		};
 
 		Ok(reply_message_content)

src/service/admin/query.rs

+use clap::Subcommand;
+use ruma::{
+	events::{room::message::RoomMessageEventContent, RoomAccountDataEventType},
+	RoomId, UserId,
+};
+
+use crate::{services, Result};
+
+#[cfg_attr(test, derive(Debug))]
+#[derive(Subcommand)]
+/// Query tables from database
+pub(crate) enum QueryCommand {
+	/// - account_data.rs iterators and getters
+	#[command(subcommand)]
+	AccountData(AccountData),
+
+	/// - appservice.rs iterators and getters
+	#[command(subcommand)]
+	Appservice(Appservice),
+
+	/// - presence.rs iterators and getters
+	#[command(subcommand)]
+	Presence(Presence),
+}
+
+#[cfg_attr(test, derive(Debug))]
+#[derive(Subcommand)]
+/// All the getters and iterators from src/database/key_value/account_data.rs
+/// via services()
+pub(crate) enum AccountData {
+	/// - Returns all changes to the account data that happened after `since`.
+	ChangesSince {
+		/// Full user ID
+		user_id: Box<UserId>,
+		/// UNIX timestamp since (u64)
+		since: u64,
+		/// Optional room ID of the account data
+		room_id: Option<Box<RoomId>>,
+	},
+
+	/// - Searches the account data for a specific kind.
+	Get {
+		/// Full user ID
+		user_id: Box<UserId>,
+		/// Account data event type
+		kind: RoomAccountDataEventType,
+		/// Optional room ID of the account data
+		room_id: Option<Box<RoomId>>,
+	},
+}
+
+#[cfg_attr(test, derive(Debug))]
+#[derive(Subcommand)]
+/// All the getters and iterators from src/database/key_value/appservice.rs via
+/// services()
+pub(crate) enum Appservice {
+	/// - Gets the appservice registration info/details from the ID as a string
+	GetRegistration {
+		/// Appservice registration ID
+		appservice_id: Box<str>,
+	},
+}
+
+#[cfg_attr(test, derive(Debug))]
+#[derive(Subcommand)]
+/// All the getters and iterators from src/database/key_value/presence.rs via
+/// services()
+pub(crate) enum Presence {
+	/// - Returns the latest presence event for the given user.
+	GetPresence {
+		/// Full user ID
+		user_id: Box<UserId>,
+	},
+
+	/// - Returns the most recent presence updates that happened after the event
+	///   with id `since`.
+	PresenceSince {
+		/// UNIX timestamp since (u64)
+		since: u64,
+	},
+}
+
+/// Processes admin command
+#[allow(non_snake_case)]
+pub(crate) async fn process(command: QueryCommand, _body: Vec<&str>) -> Result<RoomMessageEventContent> {
+	match command {
+		QueryCommand::AccountData(AccountData) => account_data(AccountData).await,
+		QueryCommand::Appservice(Appservice) => appservice(Appservice).await,
+		QueryCommand::Presence(Presence) => presence(Presence).await,
+	}
+}
+
+/// All the getters and iterators in key_value/account_data.rs via services()
+async fn account_data(subcommand: AccountData) -> Result<RoomMessageEventContent> {
+	match subcommand {
+		AccountData::ChangesSince {
+			user_id,
+			since,
+			room_id,
+		} => {
+			let timer = tokio::time::Instant::now();
+			let results = services()
+				.account_data
+				.changes_since(room_id.as_deref(), &user_id, since)?;
+			let query_time = timer.elapsed();
+
+			Ok(RoomMessageEventContent::text_html(
+				format!("Query completed in {query_time:?}:\n\n```\n{:?}```", results),
+				format!(
+					"<p>Query completed in {query_time:?}:</p>\n<pre><code>{:?}\n</code></pre>",
+					results
+				),
+			))
+		},
+		AccountData::Get {
+			user_id,
+			kind,
+			room_id,
+		} => {
+			let timer = tokio::time::Instant::now();
+			let results = services()
+				.account_data
+				.get(room_id.as_deref(), &user_id, kind)?;
+			let query_time = timer.elapsed();
+
+			Ok(RoomMessageEventContent::text_html(
+				format!("Query completed in {query_time:?}:\n\n```\n{:?}```", results),
+				format!(
+					"<p>Query completed in {query_time:?}:</p>\n<pre><code>{:?}\n</code></pre>",
+					results
+				),
+			))
+		},
+	}
+}
+
+/// All the getters and iterators in key_value/appservice.rs via services()
+async fn appservice(subcommand: Appservice) -> Result<RoomMessageEventContent> {
+	match subcommand {
+		Appservice::GetRegistration {
+			appservice_id,
+		} => {
+			let timer = tokio::time::Instant::now();
+			let results = services()
+				.appservice
+				.get_registration(appservice_id.as_ref())
+				.await;
+			let query_time = timer.elapsed();
+
+			Ok(RoomMessageEventContent::text_html(
+				format!("Query completed in {query_time:?}:\n\n```\n{:?}```", results),
+				format!(
+					"<p>Query completed in {query_time:?}:</p>\n<pre><code>{:?}\n</code></pre>",
+					results
+				),
+			))
+		},
+	}
+}
+
+/// All the getters and iterators in key_value/appservice.rs via services()
+async fn presence(subcommand: Presence) -> Result<RoomMessageEventContent> {
+	match subcommand {
+		Presence::GetPresence {
+			user_id,
+		} => {
+			let timer = tokio::time::Instant::now();
+			let results = services().presence.get_presence(&user_id)?;
+			let query_time = timer.elapsed();
+
+			Ok(RoomMessageEventContent::text_html(
+				format!("Query completed in {query_time:?}:\n\n```\n{:?}```", results),
+				format!(
+					"<p>Query completed in {query_time:?}:</p>\n<pre><code>{:?}\n</code></pre>",
+					results
+				),
+			))
+		},
+		Presence::PresenceSince {
+			since,
+		} => {
+			let timer = tokio::time::Instant::now();
+			let results = services().presence.presence_since(since);
+			let query_time = timer.elapsed();
+
+			let presence_since: Vec<(_, _, _)> = results.collect();
+
+			Ok(RoomMessageEventContent::text_html(
+				format!("Query completed in {query_time:?}:\n\n```\n{:?}```", presence_since),
+				format!(
+					"<p>Query completed in {query_time:?}:</p>\n<pre><code>{:?}\n</code></pre>",
+					presence_since
+				),
+			))
+		},
+	}
+}