Function classify_sql_statement

marimo/_ast/sql_utils.py:22–63

Identifies whether a SQL statement is a DDL, DML, or DQL statement.

(
    sql_statement: str, dialect: SQLGLOT_DIALECTS | None = None
)

def classify_sql_statement(
    sql_statement: str, dialect: SQLGLOT_DIALECTS | None = None
) -> SQL_TYPE | Literal["unknown"]:
    """
    Identifies whether a SQL statement is a DDL, DML, or DQL statement.
    """
    DependencyManager.sqlglot.require(why="SQL parsing")

    from sqlglot import exp, parse
    from sqlglot.errors import ParseError

    sql_statement = sql_statement.strip().lower()
    try:
        with _loggers.suppress_warnings_logs("sqlglot"):
            expression_list = parse(sql_statement, dialect=dialect)
    except ParseError as e:
        log_sql_error(
            LOGGER.debug,
            message="Failed to parse SQL statement for classification.",
            exception=e,
            rule_code="MF005",
            node=None,
            sql_content=sql_statement,
        )
        return "unknown"

    for expression in expression_list:
        if expression is None:
            continue

        if bool(
            expression.find(
                exp.Create, exp.Drop, exp.Alter, exp.Attach, exp.Detach
            )
        ):
            return "DDL"
        elif bool(expression.find(exp.Insert, exp.Update, exp.Delete)):
            return "DML"
        else:
            return "DQL"

    return "unknown"