Source code for sqlspec.builder._explain

"""EXPLAIN statement builder.

Provides a fluent interface for building EXPLAIN statements with
dialect-aware SQL generation.
"""

from typing import TYPE_CHECKING, Any

from mypy_extensions import trait
from sqlglot import Dialect, exp
from typing_extensions import Self

from sqlspec.core import SQL, StatementConfig
from sqlspec.core.explain import ORACLE_EXPLAIN_PREFIX, ORACLE_MANAGED_EXPLAIN_META_KEY, ExplainFormat, ExplainOptions
from sqlspec.exceptions import SQLBuilderError
from sqlspec.utils.type_guards import has_expression_and_sql, has_parameter_builder, is_expression

if TYPE_CHECKING:
    from sqlglot.dialects.dialect import DialectType

    from sqlspec.protocols import SQLBuilderProtocol


__all__ = (
    "Explain",
    "ExplainMixin",
    "build_bigquery_explain",
    "build_duckdb_explain",
    "build_explain_sql",
    "build_generic_explain",
    "build_mysql_explain",
    "build_oracle_explain",
    "build_postgres_explain",
    "build_sqlite_explain",
    "normalize_dialect_name",
)


POSTGRES_DIALECTS = frozenset({"postgres", "postgresql", "redshift"})
MYSQL_DIALECTS = frozenset({"mysql", "mariadb"})
SQLITE_DIALECTS = frozenset({"sqlite"})
DUCKDB_DIALECTS = frozenset({"duckdb"})
ORACLE_DIALECTS = frozenset({"oracle"})
BIGQUERY_DIALECTS = frozenset({"bigquery"})
SPANNER_DIALECTS = frozenset({"spanner"})


def normalize_dialect_name(dialect: "DialectType | None") -> str | None:
    """Normalize dialect to lowercase string.

    Args:
        dialect: Dialect type, string, or None

    Returns:
        Lowercase string representation of dialect or None
    """
    if dialect is None:
        return None
    if isinstance(dialect, str):
        return dialect.lower()
    if isinstance(dialect, type) and issubclass(dialect, Dialect):
        return dialect.__name__.lower()
    return type(dialect).__name__.lower()


def build_postgres_explain(statement_sql: str, options: "ExplainOptions") -> str:
    """Build PostgreSQL EXPLAIN statement.

    PostgreSQL uses the syntax: EXPLAIN (OPTIONS) statement

    Args:
        statement_sql: The SQL statement to explain
        options: ExplainOptions configuration

    Returns:
        Complete EXPLAIN SQL string
    """
    option_parts: list[str] = []

    if options.analyze:
        option_parts.append("ANALYZE")
    if options.verbose:
        option_parts.append("VERBOSE")
    if options.costs is not None:
        option_parts.append(f"COSTS {'TRUE' if options.costs else 'FALSE'}")
    if options.buffers is not None:
        option_parts.append(f"BUFFERS {'TRUE' if options.buffers else 'FALSE'}")
    if options.timing is not None:
        option_parts.append(f"TIMING {'TRUE' if options.timing else 'FALSE'}")
    if options.summary is not None:
        option_parts.append(f"SUMMARY {'TRUE' if options.summary else 'FALSE'}")
    if options.memory is not None:
        option_parts.append(f"MEMORY {'TRUE' if options.memory else 'FALSE'}")
    if options.settings is not None:
        option_parts.append(f"SETTINGS {'TRUE' if options.settings else 'FALSE'}")
    if options.wal is not None:
        option_parts.append(f"WAL {'TRUE' if options.wal else 'FALSE'}")
    if options.generic_plan is not None:
        option_parts.append(f"GENERIC_PLAN {'TRUE' if options.generic_plan else 'FALSE'}")
    if options.format is not None:
        option_parts.append(f"FORMAT {options.format.value.upper()}")

    if option_parts:
        options_str = ", ".join(option_parts)
        return f"EXPLAIN ({options_str}) {statement_sql}"
    return f"EXPLAIN {statement_sql}"


def build_mysql_explain(statement_sql: str, options: "ExplainOptions") -> str:
    """Build MySQL EXPLAIN statement.

    MySQL uses:
        - EXPLAIN [FORMAT = TRADITIONAL|JSON|TREE] statement
        - EXPLAIN ANALYZE statement (always TREE format)

    Args:
        statement_sql: The SQL statement to explain
        options: ExplainOptions configuration

    Returns:
        Complete EXPLAIN SQL string
    """
    if options.analyze:
        return f"EXPLAIN ANALYZE {statement_sql}"

    if options.format is not None:
        format_map = {
            ExplainFormat.JSON: "JSON",
            ExplainFormat.TREE: "TREE",
            ExplainFormat.TRADITIONAL: "TRADITIONAL",
            ExplainFormat.TEXT: "TRADITIONAL",
        }
        fmt = format_map.get(options.format, "TRADITIONAL")
        return f"EXPLAIN FORMAT = {fmt} {statement_sql}"

    return f"EXPLAIN {statement_sql}"


def build_sqlite_explain(statement_sql: str, options: "ExplainOptions") -> str:
    """Build SQLite EXPLAIN statement.

    SQLite only supports EXPLAIN QUERY PLAN (no additional options).
    Raw EXPLAIN returns virtual machine opcodes which is rarely useful.

    Args:
        statement_sql: The SQL statement to explain
        options: ExplainOptions configuration (mostly ignored for SQLite)

    Returns:
        Complete EXPLAIN SQL string
    """
    return f"EXPLAIN QUERY PLAN {statement_sql}"


def build_duckdb_explain(statement_sql: str, options: "ExplainOptions") -> str:
    """Build DuckDB EXPLAIN statement.

    DuckDB supports:
        - EXPLAIN statement
        - EXPLAIN ANALYZE statement
        - EXPLAIN (FORMAT JSON) statement

    Args:
        statement_sql: The SQL statement to explain
        options: ExplainOptions configuration

    Returns:
        Complete EXPLAIN SQL string
    """
    if options.analyze:
        return f"EXPLAIN ANALYZE {statement_sql}"

    if options.format == ExplainFormat.JSON:
        return f"EXPLAIN (FORMAT JSON) {statement_sql}"

    return f"EXPLAIN {statement_sql}"


def build_oracle_explain(statement_sql: str, options: "ExplainOptions") -> str:
    """Build Oracle EXPLAIN statement.

    Oracle requires a two-step process:
        1. EXPLAIN PLAN FOR statement
        2. SELECT * FROM TABLE(DBMS_XPLAN.DISPLAY())

    This function returns only the first step. The driver must handle
    executing both statements.

    Args:
        statement_sql: The SQL statement to explain
        options: ExplainOptions configuration (mostly ignored for Oracle)

    Returns:
        EXPLAIN PLAN FOR SQL string
    """
    return f"{ORACLE_EXPLAIN_PREFIX}{statement_sql}"


def build_bigquery_explain(statement_sql: str, options: "ExplainOptions") -> str:
    """Build BigQuery EXPLAIN statement.

    BigQuery supports:
        - EXPLAIN statement
        - EXPLAIN ANALYZE statement (incurs query execution costs!)

    Args:
        statement_sql: The SQL statement to explain
        options: ExplainOptions configuration

    Returns:
        Complete EXPLAIN SQL string
    """
    if options.analyze:
        return f"EXPLAIN ANALYZE {statement_sql}"
    return f"EXPLAIN {statement_sql}"


def build_generic_explain(statement_sql: str, options: "ExplainOptions") -> str:
    """Build generic EXPLAIN statement for unknown dialects.

    Args:
        statement_sql: The SQL statement to explain
        options: ExplainOptions configuration

    Returns:
        Complete EXPLAIN SQL string
    """
    if options.analyze:
        return f"EXPLAIN ANALYZE {statement_sql}"
    return f"EXPLAIN {statement_sql}"


def build_explain_sql(statement_sql: str, options: "ExplainOptions", dialect: "DialectType | None" = None) -> str:
    """Build dialect-specific EXPLAIN SQL.

    Args:
        statement_sql: The SQL statement to explain
        options: ExplainOptions configuration
        dialect: Target SQL dialect

    Returns:
        Complete EXPLAIN SQL string for the target dialect
    """
    dialect_name = normalize_dialect_name(dialect)

    if dialect_name in POSTGRES_DIALECTS:
        return build_postgres_explain(statement_sql, options)
    if dialect_name in MYSQL_DIALECTS:
        return build_mysql_explain(statement_sql, options)
    if dialect_name in SQLITE_DIALECTS:
        return build_sqlite_explain(statement_sql, options)
    if dialect_name in DUCKDB_DIALECTS:
        return build_duckdb_explain(statement_sql, options)
    if dialect_name in ORACLE_DIALECTS:
        return build_oracle_explain(statement_sql, options)
    if dialect_name in BIGQUERY_DIALECTS:
        return build_bigquery_explain(statement_sql, options)
    if dialect_name in SPANNER_DIALECTS:
        return build_generic_explain(statement_sql, options)

    return build_generic_explain(statement_sql, options)


class Explain:
    """Builder for EXPLAIN statements with dialect-aware rendering.

    Provides a fluent API for constructing EXPLAIN statements with
    various options that are translated to dialect-specific syntax.
    """

    __slots__ = ("_dialect", "_options", "_parameters", "_source_config", "_statement_sql")

    def __init__(
        self,
        statement: "str | exp.Expr | SQL | SQLBuilderProtocol",
        dialect: "DialectType | None" = None,
        options: "ExplainOptions | None" = None,
    ) -> None:
        """Initialize ExplainBuilder.

        Args:
            statement: SQL statement to explain (string, expression, SQL object, or builder)
            dialect: Target SQL dialect
            options: Initial ExplainOptions (or None for defaults)
        """
        self._dialect = dialect
        self._options = options if options is not None else ExplainOptions()
        self._parameters: dict[str, Any] = {}
        self._source_config: StatementConfig | None = None

        self._statement_sql = self._resolve_statement_sql(statement)
def analyze(self, enabled: bool = True) -> Self: """Enable ANALYZE option (execute statement for real statistics). Args: enabled: Whether to enable ANALYZE Returns: Self for method chaining """ self._options = self._options.copy(analyze=enabled) return self