1
use clap::Parser;
2
use clap_complete::ArgValueCompleter;
3
use futures_util::SinkExt;
4
use tokio_stream::StreamExt;
5

            
6
use crate::{
7
    client::commands::{erroneous_server_response, print_authorization_owner_hint},
8
    core::{
9
        completion::mysql_user_completer,
10
        protocol::{
11
            ClientToServerMessageStream, LockUserError, Request, Response,
12
            print_lock_users_output_status, print_lock_users_output_status_json,
13
            request_validation::ValidationError,
14
        },
15
        types::MySQLUser,
16
    },
17
};
18

            
19
#[derive(Parser, Debug, Clone)]
20
pub struct LockUserArgs {
21
    /// The `MySQL` user(s) to loc
22
    #[arg(num_args = 1.., value_name = "USER_NAME")]
23
    #[cfg_attr(not(feature = "suid-sgid-mode"), arg(add = ArgValueCompleter::new(mysql_user_completer)))]
24
    username: Vec<MySQLUser>,
25

            
26
    /// Print the information as JSON
27
    #[arg(short, long)]
28
    json: bool,
29
}
30

            
31
pub async fn lock_users(
32
    args: LockUserArgs,
33
    mut server_connection: ClientToServerMessageStream,
34
) -> anyhow::Result<()> {
35
    if args.username.is_empty() {
36
        anyhow::bail!("No usernames provided");
37
    }
38

            
39
    let message = Request::LockUsers(args.username.clone());
40

            
41
    if let Err(err) = server_connection.send(message).await {
42
        server_connection.close().await.ok();
43
        anyhow::bail!(err);
44
    }
45

            
46
    let result = match server_connection.next().await {
47
        Some(Ok(Response::LockUsers(result))) => result,
48
        response => return erroneous_server_response(response),
49
    };
50

            
51
    if args.json {
52
        print_lock_users_output_status_json(&result);
53
    } else {
54
        print_lock_users_output_status(&result);
55

            
56
        if result.iter().any(|(_, res)| {
57
            matches!(
58
                res,
59
                Err(LockUserError::ValidationError(
60
                    ValidationError::AuthorizationError(_)
61
                ))
62
            )
63
        }) {
64
            print_authorization_owner_hint(&mut server_connection).await?;
65
        }
66
    }
67

            
68
    server_connection.send(Request::Exit).await?;
69

            
70
    if result.values().any(std::result::Result::is_err) {
71
        std::process::exit(1);
72
    }
73

            
74
    Ok(())
75
}