Function classify_sandbox_statement

mycli/packages/sql_utils.py:488–531

Classify a SQL statement for sandbox mode and extract the new password. Returns (statement_type, new_password) where statement_type is one of: - 'alter_user' — ALTER USER ... IDENTIFIED BY ... - 'set_password' — SET PASSWORD [FOR ...] = ... - 'quit' — quit, exit, \\q

(text: str)

def classify_sandbox_statement(text: str) -> tuple[str | None, str | None]:
    """Classify a SQL statement for sandbox mode and extract the new password.

    Returns (statement_type, new_password) where statement_type is one of:
    - 'alter_user'    — ALTER USER ... IDENTIFIED BY ...
    - 'set_password'  — SET PASSWORD [FOR ...] = ...
    - 'quit'          — quit, exit, \\q
    - None            — not allowed in sandbox mode
    """
    stripped = text.strip()
    if not stripped:
        return ('quit', None)

    try:
        tokens = list(sqlglot.tokenize(stripped, dialect='mysql'))
    except sqlglot.errors.TokenError:
        tokens = []

    if not tokens:
        return ('quit', None)

    types = [t.token_type for t in tokens]
    texts = [t.text.upper() for t in tokens]
    tt = sqlglot.tokens.TokenType

    # quit, exit
    if len(tokens) == 1 and types[0] == tt.VAR and texts[0] in ('QUIT', 'EXIT'):
        return ('quit', None)

    # \q
    if len(tokens) == 2 and types[0] in (tt.BACKSLASH, tt.SLASH) and texts[1] in ('Q', 'QUIT', 'EXIT'):
        return ('quit', None)

    # ALTER USER ...
    if len(tokens) >= 2 and types[0] == tt.ALTER and texts[1] == 'USER':
        pw = _find_password_after_by(tokens)
        return ('alter_user', pw)

    # SET PASSWORD ...
    if len(tokens) >= 2 and types[0] == tt.SET and texts[1] == 'PASSWORD':
        pw = _find_password_after_eq(tokens)
        return ('set_password', pw)

    return (None, None)


def _find_password_after_by(tokens: list[sqlglot.tokens.Token]) -> str | None: