improve user autocomplete for commands, add db user lookup

Author: dankmeme01

Cargo.lock

@@ -74,6 +74,7 @@ metrics-process = "2.4.3"
 aho-corasick = { version = "1.1.4", optional = true }
 async-watcher = { version = "0.4.0", optional = true }
 mimalloc = { version = "0.1.48", optional = true }
+fuzzy-matcher = { version = "0.3.7", optional = true }
 validator = { version = "0.20.0", features = ["derive"] }
 
 # Analytics
@@ -111,7 +112,7 @@ docker-alpine-build = ["all", "mimalloc"]
 word-filter = ["dep:aho-corasick", "dep:async-watcher"]
 featured-levels = ["dep:google-sheets4"]
 default = []
-discord = ["dep:poise", "dep:plotters", "dep:image"]
+discord = ["dep:poise", "dep:plotters", "dep:image", "dep:fuzzy-matcher"]
 quic = ["server-shared/quic"]
 websocket = ["server-shared/websocket"]
 mimalloc = ["dep:mimalloc"]

Cargo.toml

@@ -39,7 +39,7 @@ pub async fn punish(
     #[description = "Punishment type"]
     punishment_type: String,
 
-    #[autocomplete = "online_user_autocomplete"]
+    #[autocomplete = "online_and_db_user_autocomplete"]
     #[description = "Geometry Dash username or ID"]
     target_user: String,
 
@@ -88,7 +88,7 @@ pub async fn unpunish(
     #[autocomplete = "punish_autocomplete"]
     #[description = "Punishment type"]
     punishment_type: String,
-    #[autocomplete = "online_user_autocomplete"]
+    #[autocomplete = "online_and_db_user_autocomplete"]
     #[description = "Geometry Dash username or ID"]
     target_user: String,
 ) -> Result<(), BotError> {
@@ -300,7 +300,9 @@ pub async fn check_actions(
 #[poise::command(slash_command, guild_only = true)]
 pub async fn check_alts(
     ctx: Context<'_>,
-    #[description = "GD username or account ID of the target user"] user: String,
+    #[autocomplete = "db_user_autocomplete"]
+    #[description = "GD username or account ID of the target user"]
+    user: String,
 ) -> Result<(), BotError> {
     check_moderator(ctx).await?;
 

src/discord/commands/moderation.rs

@@ -1,5 +1,7 @@
-use std::{sync::Arc, time::Duration};
+use std::{cmp::Reverse, sync::Arc, time::Duration};
 
+use fuzzy_matcher::{FuzzyMatcher, skim::SkimMatcherV2};
+use itertools::Itertools;
 use poise::{
     CreateReply, ReplyHandle,
     serenity_prelude::{self as serenity, AutocompleteChoice},
@@ -8,7 +10,7 @@ use server_shared::qunet::server::Server;
 use thiserror::Error;
 
 use crate::{
-    core::handler::ConnectionHandler,
+    core::handler::{ClientStateHandle, ConnectionHandler},
     discord::{BotError, state::BotState},
     users::{ComputedRole, DbUser, UserPunishmentType, UsersModule},
 };
@@ -131,24 +133,75 @@ pub fn parse_duration_str(s: &str) -> Result<Duration, ParseDurationError> {
     Ok(Duration::from_secs(number * modifier))
 }
 
-pub async fn online_user_autocomplete(ctx: Context<'_>, partial: &str) -> Vec<AutocompleteChoice> {
+fn fuzzy_match(target: &str, candidate: &str) -> i64 {
+    let matcher = SkimMatcherV2::default();
+    matcher.fuzzy_match(target, candidate).unwrap_or(-1)
+}
+
+fn wrap_user_autocomplete<'a>(
+    query: &str,
+    iter: impl Iterator<Item = (&'a str, i32)>,
+) -> Vec<AutocompleteChoice> {
+    iter.sorted_by_key(|(username, id)| {
+        if let Ok(query_id) = query.parse::<i32>() {
+            if *id == query_id {
+                return Reverse(i64::MAX); // highest priority if ID matches
+            }
+        }
+
+        if username.eq_ignore_ascii_case(query) {
+            return Reverse(i64::MAX - 1); // second highest priority if username matches
+        }
+
+        // fuzzy match on username
+        Reverse(fuzzy_match(username, query))
+    })
+    .map(|(username, id)| AutocompleteChoice::new(username.to_owned(), id.to_string()))
+    .take(10)
+    .collect()
+}
+
+fn get_online_users_matching(ctx: Context<'_>, partial: &str) -> Vec<ClientStateHandle> {
     let server = ctx.data().server().unwrap();
-    let clients = server.handler().get_n_clients_matching(partial, 5);
+    let mut clients = server.handler().get_n_clients_matching(partial, 10);
+
+    if let Ok(query_id) = partial.parse::<i32>() {
+        if let Some(client) = server.handler().find_client(query_id) {
+            clients.push(client);
+        }
+    }
 
     clients
-        .into_iter()
-        .map(|c| AutocompleteChoice::new(c.username().to_string(), c.account_id().to_string()))
-        .collect()
 }
 
-// async fn db_user_autocomplete(ctx: Context<'_>, partial: &str) -> Vec<AutocompleteChoice> {
-//     let server = ctx.data().server().unwrap();
-//     let users = server.handler().module::<UsersModule>();
+async fn get_db_users_matching(ctx: Context<'_>, partial: &str) -> Vec<DbUser> {
+    let server = ctx.data().server().unwrap();
+    let users = server.handler().module::<UsersModule>();
 
-//     users.query_user();
+    users.query_matching_users(partial, 50).await.unwrap_or_default()
+}
 
-//     clients
-//         .into_iter()
-//         .map(|c| AutocompleteChoice::new(c.username().to_string(), c.account_id().to_string()))
-//         .collect()
-// }
+pub async fn online_user_autocomplete(ctx: Context<'_>, partial: &str) -> Vec<AutocompleteChoice> {
+    let clients = get_online_users_matching(ctx, partial);
+    wrap_user_autocomplete(partial, clients.iter().map(|c| (c.username(), c.account_id())))
+}
+
+pub async fn db_user_autocomplete(ctx: Context<'_>, partial: &str) -> Vec<AutocompleteChoice> {
+    let users = get_db_users_matching(ctx, partial).await;
+    wrap_user_autocomplete(partial, users.iter().map(|u| (u.username(), u.account_id)))
+}
+
+pub async fn online_and_db_user_autocomplete(
+    ctx: Context<'_>,
+    partial: &str,
+) -> Vec<AutocompleteChoice> {
+    let mut vec = Vec::new();
+
+    let first = get_online_users_matching(ctx, partial);
+    vec.extend(first.iter().map(|c| (c.username(), c.account_id())));
+
+    let second = get_db_users_matching(ctx, partial).await;
+    vec.extend(second.iter().map(|u| (u.username(), u.account_id)));
+
+    wrap_user_autocomplete(partial, vec.into_iter())
+}

src/discord/commands/util.rs

@@ -227,6 +227,34 @@ impl UsersDb {
         }
     }
 
+    pub async fn query_matching_users(
+        &self,
+        query: &str,
+        limit: usize,
+    ) -> DatabaseResult<Vec<DbUser>> {
+        let mut out = Vec::new();
+
+        // similar logic as query_user
+
+        if let Ok(id) = query.parse::<i32>() {
+            if let Some(user) = User::find_by_id(id).one(&self.conn).await? {
+                out.push(self.post_user_fetch(user).await?);
+            }
+        };
+
+        let users = User::find()
+            .filter(user::Column::Username.contains(query))
+            .limit(limit as u64)
+            .all(&self.conn)
+            .await?;
+
+        for user in users {
+            out.push(self.post_user_fetch(user).await?);
+        }
+
+        Ok(out)
+    }
+
     pub async fn query_user_with_role(&self, role_id: &str) -> DatabaseResult<Vec<DbUser>> {
         let users =
             User::find().filter(user::Column::Roles.contains(role_id)).all(&self.conn).await?;
@@ -894,7 +922,10 @@ impl DbUser {
     }
 
     pub fn username(&self) -> &str {
-        self.username.as_deref().unwrap_or("<unknown>")
+        match self.username.as_deref() {
+            Some(name) if !name.is_empty() => name,
+            _ => "Unknown",
+        }
     }
 }
 

src/users/database/mod.rs

@@ -320,6 +320,15 @@ impl UsersModule {
         self.db.query_user(query).await
     }
 
+    /// Queries users with a similar username
+    pub async fn query_matching_users(
+        &self,
+        query: &str,
+        max_users: usize,
+    ) -> DatabaseResult<Vec<DbUser>> {
+        self.db.query_matching_users(query, max_users).await
+    }
+
     /// Query a user by account ID or username, creating them if they don't exist
     /// If the user does not exist, this will fetch the data from GD servers
     pub async fn query_or_create_user(&self, query: &str) -> Result<Option<DbUser>, Error> {