chore: Migrate /superset/csv/<client_id> to API v1

superset-frontend/src/SqlLab/components/ResultSet/index.tsx

@@ -17,6 +17,7 @@
  * under the License.
  */
 import React, { useCallback, useEffect, useState } from 'react';
+import rison from 'rison';
 import { useDispatch } from 'react-redux';
 import ButtonGroup from 'src/components/ButtonGroup';
 import Alert from 'src/components/Alert';
@@ -219,6 +220,14 @@ const ResultSet = ({
     }
   };
 
+  const getExportCsvUrl = (clientId: string) => {
+    const params = rison.encode({
+      client_id: clientId,
+    });
+
+    return `/api/v1/sqllab/export/?q=${params}`;
+  };
+
   const renderControls = () => {
     if (search || visualize || csv) {
       let { data } = query.results;
@@ -257,7 +266,7 @@ const ResultSet = ({
               />
             )}
             {csv && (
-              <Button buttonSize="small" href={`/superset/csv/${query.id}`}>
+              <Button buttonSize="small" href={getExportCsvUrl(query.id)}>
                 <i className="fa fa-file-text-o" /> {t('Download to CSV')}
               </Button>
             )}

superset/sqllab/api.py

@@ -16,6 +16,7 @@
 # under the License.
 import logging
 from typing import Any, cast, Dict, Optional
+from urllib import parse
 
 import simplejson as json
 from flask import request
@@ -32,6 +33,7 @@
 from superset.sql_lab import get_sql_results
 from superset.sqllab.command_status import SqlJsonExecutionStatus
 from superset.sqllab.commands.execute import CommandResult, ExecuteSqlCommand
+from superset.sqllab.commands.export import SqlResultExportCommand
 from superset.sqllab.commands.results import SqlExecutionResultsCommand
 from superset.sqllab.exceptions import (
     QueryIsForbiddenToAccessException,
@@ -42,6 +44,7 @@
 from superset.sqllab.schemas import (
     ExecutePayloadSchema,
     QueryExecutionResponseSchema,
+    sql_lab_export_csv_schema,
     sql_lab_get_results_schema,
 )
 from superset.sqllab.sql_json_executer import (
@@ -53,7 +56,7 @@
 from superset.sqllab.validators import CanAccessQueryValidatorImpl
 from superset.superset_typing import FlaskResponse
 from superset.utils import core as utils
-from superset.views.base import json_success
+from superset.views.base import CsvResponse, generate_download_headers, json_success
 from superset.views.base_api import BaseSupersetApi, requires_json, statsd_metrics
 
 config = app.config
@@ -72,13 +75,81 @@ class SqlLabRestApi(BaseSupersetApi):
 
     apispec_parameter_schemas = {
         "sql_lab_get_results_schema": sql_lab_get_results_schema,
+        "sql_lab_export_csv_schema": sql_lab_export_csv_schema,
     }
     openapi_spec_tag = "SQL Lab"
     openapi_spec_component_schemas = (
         ExecutePayloadSchema,
         QueryExecutionResponseSchema,
     )
 
+    @expose("/export/")
+    @protect()
+    @statsd_metrics
+    @rison(sql_lab_export_csv_schema)
+    @event_logger.log_this_with_context(
+        action=lambda self, *args, **kwargs: f"{self.__class__.__name__}"
+        f".export_csv",
+        log_to_statsd=False,
+    )
+    def export_csv(self, **kwargs: Any) -> CsvResponse:
+        """Exports the SQL Query results to a CSV
+        ---
+        get:
+          summary: >-
+            Exports the SQL Query results to a CSV
+          parameters:
+          - in: query
+            name: q
+            content:
+              application/json:
+                schema:
+                  $ref: '#/components/schemas/sql_lab_export_csv_schema'
+          responses:
+            200:
+              description: SQL query results
+              content:
+                text/csv:
+                  schema:
+                    type: string
+            400:
+              $ref: '#/components/responses/400'
+            401:
+              $ref: '#/components/responses/401'
+            403:
+              $ref: '#/components/responses/403'
+            404:
+              $ref: '#/components/responses/404'
+            500:
+              $ref: '#/components/responses/500'
+        """
+        params = kwargs["rison"]
+        client_id = params.get("client_id")
+        result = SqlResultExportCommand(client_id=client_id).run()
+
+        query = result.get("query")
+        data = result.get("data")
+        row_count = result.get("row_count")
+
+        quoted_csv_name = parse.quote(query.name)
+        response = CsvResponse(
+            data, headers=generate_download_headers("csv", quoted_csv_name)
+        )
+        event_info = {
+            "event_type": "data_export",
+            "client_id": client_id,
+            "row_count": row_count,
+            "database": query.database.name,
+            "schema": query.schema,
+            "sql": query.sql,
+            "exported_format": "csv",
+        }
+        event_rep = repr(event_info)
+        logger.debug(
+            "CSV exported: %s", event_rep, extra={"superset_event": event_info}
+        )
+        return response
+
     @expose("/results/")
     @protect()
     @statsd_metrics

superset/sqllab/commands/export.py

@@ -0,0 +1,130 @@
+# 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.
+# pylint: disable=too-few-public-methods, too-many-arguments
+from __future__ import annotations
+
+import logging
+from typing import Any, cast, Dict
+
+import pandas as pd
+from flask_babel import gettext as __, lazy_gettext as _
+
+from superset import app, db, results_backend, results_backend_use_msgpack
+from superset.commands.base import BaseCommand
+from superset.errors import ErrorLevel, SupersetError, SupersetErrorType
+from superset.exceptions import SupersetErrorException, SupersetSecurityException
+from superset.models.sql_lab import Query
+from superset.sql_parse import ParsedQuery
+from superset.sqllab.limiting_factor import LimitingFactor
+from superset.utils import core as utils, csv
+from superset.utils.dates import now_as_float
+from superset.views.utils import _deserialize_results_payload
+
+config = app.config
+
+logger = logging.getLogger(__name__)
+
+
+class SqlResultExportCommand(BaseCommand):
+    _client_id: str
+    _query: Query
+
+    def __init__(
+        self,
+        client_id: str,
+    ) -> None:
+        self._client_id = client_id
+
+    def validate(self) -> None:
+        self._query = (
+            db.session.query(Query).filter_by(client_id=self._client_id).one_or_none()
+        )
+        if self._query is None:
+            raise SupersetErrorException(
+                SupersetError(
+                    message=__(
+                        "The query associated with these results could not be found. "
+                        "You need to re-run the original query."
+                    ),
+                    error_type=SupersetErrorType.RESULTS_BACKEND_ERROR,
+                    level=ErrorLevel.ERROR,
+                ),
+                status=404,
+            )
+
+        try:
+            self._query.raise_for_access()
+        except SupersetSecurityException:
+            raise SupersetErrorException(
+                SupersetError(
+                    message=__("Cannot access the query"),
+                    error_type=SupersetErrorType.QUERY_SECURITY_ACCESS_ERROR,
+                    level=ErrorLevel.ERROR,
+                ),
+                status=403,
+            )
+
+    def run(
+        self,
+    ) -> Dict[str, Any]:
+        self.validate()
+        blob = None
+        if results_backend and self._query.results_key:
+            logger.info(
+                "Fetching CSV from results backend [%s]", self._query.results_key
+            )
+            blob = results_backend.get(self._query.results_key)
+        if blob:
+            logger.info("Decompressing")
+            payload = utils.zlib_decompress(
+                blob, decode=not results_backend_use_msgpack
+            )
+            obj = _deserialize_results_payload(
+                payload, self._query, cast(bool, results_backend_use_msgpack)
+            )
+
+            df = pd.DataFrame(
+                data=obj["data"],
+                dtype=object,
+                columns=[c["name"] for c in obj["columns"]],
+            )
+
+            logger.info("Using pandas to convert to CSV")
+        else:
+            logger.info("Running a query to turn into CSV")
+            if self._query.select_sql:
+                sql = self._query.select_sql
+                limit = None
+            else:
+                sql = self._query.executed_sql
+                limit = ParsedQuery(sql).limit
+            if limit is not None and self._query.limiting_factor in {
+                LimitingFactor.QUERY,
+                LimitingFactor.DROPDOWN,
+                LimitingFactor.QUERY_AND_DROPDOWN,
+            }:
+                # remove extra row from `increased_limit`
+                limit -= 1
+            df = self._query.database.get_df(sql, self._query.schema)[:limit]
+
+        csv_data = csv.df_to_escaped_csv(df, index=False, **config["CSV_EXPORT"])
+
+        return {
+            "query": self._query,
+            "count": len(df.index),
+            "data": csv_data,
+        }

superset/sqllab/schemas.py

@@ -24,6 +24,14 @@
     "required": ["key"],
 }
 
+sql_lab_export_csv_schema = {
+    "type": "object",
+    "properties": {
+        "client_id": {"type": "string"},
+    },
+    "required": ["client_id"],
+}
+
 
 class ExecutePayloadSchema(Schema):
     database_id = fields.Integer(required=True)

superset/views/core.py

@@ -2392,6 +2392,7 @@ def _create_response_from_execution_context(  # pylint: disable=invalid-name, no
     @has_access
     @event_logger.log_this
     @expose("/csv/<client_id>")
+    @deprecated()
     def csv(self, client_id: str) -> FlaskResponse:  # pylint: disable=no-self-use
         """Download the query results as csv."""
         logger.info("Exporting CSV file [%s]", client_id)

tests/integration_tests/sql_lab/api_tests.py

@@ -19,14 +19,17 @@
 import datetime
 import json
 import random
+import csv
+import pandas as pd
+import io
 
 import pytest
 import prison
 from sqlalchemy.sql import func
 from unittest import mock
 
 from tests.integration_tests.test_app import app
-from superset import sql_lab
+from superset import db, sql_lab
 from superset.common.db_query_status import QueryStatus
 from superset.models.core import Database
 from superset.utils.database import get_example_database, get_main_database
@@ -176,3 +179,39 @@ def test_get_results_with_display_limit(self):
         self.assertEqual(result_limited, expected_limited)
 
         app.config["RESULTS_BACKEND_USE_MSGPACK"] = use_msgpack
+
+    @mock.patch("superset.models.sql_lab.Query.raise_for_access", lambda _: None)
+    @mock.patch("superset.models.core.Database.get_df")
+    def test_export_results(self, get_df_mock: mock.Mock) -> None:
+        self.login()
+
+        database = Database(
+            database_name="my_export_database", sqlalchemy_uri="sqlite://"
+        )
+        query_obj = Query(
+            client_id="test",
+            database=database,
+            tab_name="test_tab",
+            sql_editor_id="test_editor_id",
+            sql="select * from bar",
+            select_sql=None,
+            executed_sql="select * from bar limit 2",
+            limit=100,
+            select_as_cta=False,
+            rows=104,
+            error_message="none",
+            results_key="test_abc2",
+        )
+
+        db.session.add(database)
+        db.session.add(query_obj)
+
+        get_df_mock.return_value = pd.DataFrame({"foo": [1, 2, 3]})
+
+        arguments = {"client_id": "test"}
+        resp = self.get_resp(f"/api/v1/sqllab/export/?q={prison.dumps(arguments)}")
+        data = csv.reader(io.StringIO(resp))
+        expected_data = csv.reader(io.StringIO(f"foo\n1\n2"))
+
+        self.assertEqual(list(expected_data), list(data))
+        db.session.rollback()