Function _parse_sql_duckdb

marimo/_sql/parse.py:96–174

Parse an SQL query using DuckDB. Returns parse result and unexpected errors. Note: - Only SELECT statements support json_serialize_sql - Invalid function names do not throw errors - Some syntax errors do not throw errors since they are not errors in the AST parser

(
    query: str,
)

def _parse_sql_duckdb(
    query: str,
) -> tuple[SqlParseResult | None, str | None]:
    """Parse an SQL query using DuckDB. Returns parse result and unexpected errors.

    Note:
    - Only SELECT statements support json_serialize_sql
    - Invalid function names do not throw errors
    - Some syntax errors do not throw errors since they are not errors in the AST parser
    """
    if not DependencyManager.duckdb.has():
        return None, "DuckDB not installed"

    import duckdb

    json_serialize_query = (
        JSON_SERIALIZE_QUERY
        if duckdb.__version__ >= "1.1.0"
        else JSON_SERIALIZE_LEGACY_QUERY
    )

    idx_to_offset_dict: dict[int, int] = {}
    try:
        query, idx_to_offset_dict = replace_brackets_with_quotes(query)
    except Exception as e:
        LOGGER.debug(f"Error sanitizing SQL query: {e}")

    relation = duckdb.execute(json_serialize_query, [query])
    fetch_result = relation.fetchone()
    if fetch_result is None:
        return None, "No result from DuckDB parse query"

    parse_response = fetch_result[0]
    parsed_error = msgspec.json.decode(parse_response, type=DuckDBParseError)

    if not parsed_error.error:
        return SqlParseResult(success=True, errors=[]), None

    if parsed_error.error_type == "not implemented":
        # This is a valid query, but not supported by DuckDB
        # Only SELECT statements support json_serialize_sql
        return SqlParseResult(success=True, errors=[]), None

    position = int(parsed_error.position or 0)
    subquery = query[:position]
    line_number = subquery.count("\n") + 1

    last_newline_idx = subquery.rfind("\n")
    column_number = position - last_newline_idx - 1

    # Adjust column_number to account for any added quotes from bracket replacements
    # SELECT {id} FRO users
    #             ^ error position should be here
    # SELECT {id} FRO users
    #               ^ user sees this
    # So in this case, we subtract the offset from the column position.
    # If the error is before the brackets, we don't need to add the offset because it just increases the string length
    # Column position will be the same as the user sees.