superset2/superset/commands/sql_lab/estimate.py

# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.
from __future__ import annotations

import logging
from typing import Any, TypedDict

from flask import current_app as app
from flask_babel import gettext as __

from superset import db, is_feature_enabled, security_manager
from superset.commands.base import BaseCommand
from superset.errors import ErrorLevel, SupersetError, SupersetErrorType
from superset.exceptions import (
    SupersetDisallowedSQLFunctionException,
    SupersetDisallowedSQLTableException,
    SupersetDMLNotAllowedException,
    SupersetErrorException,
    SupersetTimeoutException,
)
from superset.jinja_context import get_template_processor
from superset.models.core import Database
from superset.sql.parse import SQLScript
from superset.utils import core as utils
from superset.utils.rls import apply_rls

logger = logging.getLogger(__name__)


class EstimateQueryCostType(TypedDict):
    database_id: int
    sql: str
    template_params: dict[str, Any]
    catalog: str | None
    schema: str | None


class QueryEstimationCommand(BaseCommand):
    _database_id: int
    _sql: str
    _template_params: dict[str, Any]
    _schema: str
    _database: Database
    _catalog: str | None

    def __init__(self, params: EstimateQueryCostType) -> None:
        self._database_id = params["database_id"]
        self._sql = params.get("sql", "")
        self._template_params = params.get("template_params", {})
        self._schema = params.get("schema") or ""
        self._catalog = params.get("catalog")

    def validate(self) -> None:
        self._database = db.session.query(Database).get(self._database_id)
        if not self._database:
            raise SupersetErrorException(
                SupersetError(
                    message=__("The database could not be found"),
                    error_type=SupersetErrorType.RESULTS_BACKEND_ERROR,
                    level=ErrorLevel.ERROR,
                ),
                status=404,
            )
        security_manager.raise_for_access(database=self._database)

    def _apply_sql_security(self, sql: str) -> str:
        """Run the disallowed-function/table, DML and RLS controls against the
        SQL to be estimated, mirroring ``sql_lab.execute_sql_statements``.

        Returns the SQL with RLS predicates injected (when ``RLS_IN_SQLLAB`` is
        enabled), so the cost estimate reflects the same constrained query the
        user would actually be allowed to run.
        """
        db_engine_spec = self._database.db_engine_spec
        parsed_script = SQLScript(sql, engine=db_engine_spec.engine)

        disallowed_functions = app.config["DISALLOWED_SQL_FUNCTIONS"].get(
            db_engine_spec.engine,
            set(),
        )
        if disallowed_functions and parsed_script.check_functions_present(
            disallowed_functions
        ):
            raise SupersetDisallowedSQLFunctionException(disallowed_functions)

        disallowed_tables = app.config["DISALLOWED_SQL_TABLES"].get(
            db_engine_spec.engine,
            set(),
        )
        rls_enabled = is_feature_enabled("RLS_IN_SQLLAB")

        # Resolve the effective per-query schema once, the same way the execution
        # path does (``sql_lab.execute_sql_statements``), but only when a control
        # below actually needs it. Going through ``get_default_schema_for_query``
        # rather than the static ``get_default_schema`` runs engine-specific
        # per-query security gates too — e.g. ``PostgresEngineSpec`` rejects a
        # query that sets ``search_path`` — and resolves unqualified references to
        # the schema the engine uses at runtime, so both the denylist check and
        # RLS injection match the execution path exactly.
        catalog: str | None = None
        effective_schema = ""
        if disallowed_tables or rls_enabled:
            catalog = self._catalog or self._database.get_default_catalog()
            resolved_schema = self._database.resolve_query_default_schema(
                self._sql, self._schema, catalog, self._template_params
            )
            # An explicit schema still wins for matching/RLS targeting; otherwise
            # fall back to the runtime-resolved default.
            effective_schema = self._schema or resolved_schema or ""

        if disallowed_tables:
            # Honors schema-qualified denylist entries (e.g.
            # ``information_schema.tables``) and reports only the tables
            # actually referenced by the query.
            found_tables = parsed_script.get_disallowed_tables(
                disallowed_tables, effective_schema
            )
            if found_tables:
                raise SupersetDisallowedSQLTableException(found_tables)

        if parsed_script.has_mutation() and not self._database.allow_dml:
            raise SupersetDMLNotAllowedException()

        if rls_enabled:
            for statement in parsed_script.statements:
                apply_rls(self._database, catalog, effective_schema, statement)
            return parsed_script.format()

        return sql

    def run(
        self,
    ) -> list[dict[str, Any]]:
        self.validate()

        sql = self._sql
        if self._template_params:
            template_processor = get_template_processor(self._database)
            sql = template_processor.process_template(sql, **self._template_params)

        # Apply the same SQL security controls used by the execution path
        # (sql_lab.execute_sql_statements) so cost estimation cannot be used to
        # probe disallowed functions/tables, bypass the DML guard, or confirm
        # the existence of rows hidden by row-level security.
        sql = self._apply_sql_security(sql)

        timeout = app.config["SQLLAB_QUERY_COST_ESTIMATE_TIMEOUT"]
        timeout_msg = f"The estimation exceeded the {timeout} seconds timeout."
        try:
            with utils.timeout(seconds=timeout, error_message=timeout_msg):
                cost = self._database.db_engine_spec.estimate_query_cost(
                    self._database,
                    self._catalog,
                    self._schema,
                    sql,
                    utils.QuerySource.SQL_LAB,
                )
        except SupersetTimeoutException as ex:
            logger.exception(ex)
            raise SupersetErrorException(
                SupersetError(
                    message=__(
                        "The query estimation was killed after %(sqllab_timeout)s "
                        "seconds. It might be too complex, or the database might be "
                        "under heavy load.",
                        sqllab_timeout=app.config["SQLLAB_QUERY_COST_ESTIMATE_TIMEOUT"],
                    ),
                    error_type=SupersetErrorType.SQLLAB_TIMEOUT_ERROR,
                    level=ErrorLevel.ERROR,
                ),
                status=500,
            ) from ex

        spec = self._database.db_engine_spec
        query_cost_formatters: dict[str, Any] = app.config[
            "QUERY_COST_FORMATTERS_BY_ENGINE"
        ]
        query_cost_formatter = query_cost_formatters.get(
            spec.engine, spec.query_cost_formatter
        )
        cost = query_cost_formatter(cost)
        return cost