initial support for querying database getters and iterators via admin cmd
Signed-off-by: strawberry <strawberry@puppygock.gay>
This commit is contained in:
strawberry
June
parent
f954cd2387
commit
e4a6a2325b
2 changed files with 222 additions and 6 deletions
|
|
@ -26,6 +26,7 @@ use serde_json::value::to_raw_value;
|
|||
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 appservice;
|
|||
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 @@ impl Service {
|
|||
}
|
||||
|
||||
// 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 @@ impl Service {
|
|||
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)
|
||||
|
|
|
|||
197
src/service/admin/query.rs
Normal file
197
src/service/admin/query.rs
Normal file
|
|
@ -0,0 +1,197 @@
|
|||
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
|
||||
),
|
||||
))
|
||||
},
|
||||
}
|
||||
}
|
||||
Loading…
Add table
Reference in a new issue