Skip to content

Commit

Permalink
chore: Migrate /superset/csv/<client_id> to API v1
Browse files Browse the repository at this point in the history
  • Loading branch information
diegomedina248 committed Jan 31, 2023
1 parent b94052e commit 54191d5
Show file tree
Hide file tree
Showing 8 changed files with 729 additions and 76 deletions.
271 changes: 214 additions & 57 deletions docs/static/resources/openapi.json

Large diffs are not rendered by default.

11 changes: 10 additions & 1 deletion superset-frontend/src/SqlLab/components/ResultSet/index.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -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';
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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>
)}
Expand Down
73 changes: 72 additions & 1 deletion superset/sqllab/api.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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,
Expand All @@ -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 (
Expand All @@ -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
Expand All @@ -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
Expand Down
130 changes: 130 additions & 0 deletions superset/sqllab/commands/export.py
Original file line number Diff line number Diff line change
@@ -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,
}
8 changes: 8 additions & 0 deletions superset/sqllab/schemas.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
1 change: 1 addition & 0 deletions superset/views/core.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
41 changes: 40 additions & 1 deletion tests/integration_tests/sql_lab/api_tests.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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_abc",
)

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()
Loading

0 comments on commit 54191d5

Please sign in to comment.