Repository: apache/superset
File: superset/common/query_context_processor.py (Python)

Code before:
            include_index = not isinstance(df.index, pd.RangeIndex)
            columns = list(df.columns)
            verbose_map = self._qc_datasource.data.get("verbose_map", {})
            if verbose_map:
                df.columns = [verbose_map.get(column, column) for column in columns]

            result = None
            if self._query_context.result_format == ChartDataResultFormat.CSV:
                result = csv.df_to_escaped_csv(
                    df, index=include_index, **config["CSV_EXPORT"]
                )
            elif self._query_context.result_format == ChartDataResultFormat.XLSX:
                excel.apply_column_types(df, coltypes)
                result = excel.df_to_excel(df, **config["EXCEL_EXPORT"])
            return result or ""

        return df.to_dict(orient="records")

    def get_payload(  # noqa: C901
        self,
        cache_query_context: bool | None = False,
        force_cached: bool = False,
    ) -> dict[str, Any]:
        """Returns the query results with both metadata and data"""
        # Handle empty queries case
        if not self._query_context.queries:
            return {"queries": []}

        totals_data = None
        main_query = self._query_context.queries[0]
        totals_query_index = None
        totals_result = None
        for i, query_obj in enumerate(self._query_context.queries):
            if (
                query_obj.columns == []
                and not query_obj.post_processing
                and query_obj.metrics == main_query.metrics
            ):
                totals_query_index = i
                totals_result = get_query_results(
                    query_obj.result_type or self._query_context.result_type,
                    self._query_context,
                    query_obj,
                    force_cached,
                )

                if (
                    totals_result
                    and isinstance(totals_result, dict)
                    and "data" in totals_result
                    and isinstance(totals_result["data"], list)

Review comment (refactor):
Looking at the `buildQuery` this is something to be computed by `query_object -> exe_post_processing`. I don't  think we should bring these into the context processor which is a more abstract concept. Meaning, these changes would be introducing an implementation that's very specific for your use case into something that's broadly used across the app.

Is there a way to include your new post processor without affecting the context? how is the current one working without having to alter it?

Code after:
            include_index = not isinstance(df.index, pd.RangeIndex)
            columns = list(df.columns)
            verbose_map = self._qc_datasource.data.get("verbose_map", {})
            if verbose_map:
                df.columns = [verbose_map.get(column, column) for column in columns]

            result = None
            if self._query_context.result_format == ChartDataResultFormat.CSV:
                result = csv.df_to_escaped_csv(
                    df, index=include_index, **config["CSV_EXPORT"]
                )
            elif self._query_context.result_format == ChartDataResultFormat.XLSX:
                excel.apply_column_types(df, coltypes)
                result = excel.df_to_excel(df, **config["EXCEL_EXPORT"])
            return result or ""

        return df.to_dict(orient="records")

    def ensure_totals_available(self) -> None:
        queries_needing_totals = []
        totals_queries = []

        for i, query in enumerate(self._query_context.queries):
            needs_totals = any(
                pp.get("operation") == "contribution"
                for pp in getattr(query, "post_processing", []) or []
            )

            if needs_totals:
                queries_needing_totals.append(i)

            is_totals_query = (
                not query.columns and query.metrics and not query.post_processing
            )
            if is_totals_query:
                totals_queries.append(i)

        if not queries_needing_totals or not totals_queries:
            return

        totals_idx = totals_queries[0]
        totals_query = self._query_context.queries[totals_idx]

        totals_query.row_limit = None

        result = self._query_context.get_query_result(totals_query)
        df = result.df

        totals = {
            col: df[col].sum() for col in df.columns if df[col].dtype.kind in "biufc"
        }