You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@superset.apache.org by di...@apache.org on 2023/01/30 22:44:00 UTC
[superset] branch dm/migrate-csv-to-api-v1 created (now 81252926dc)
This is an automated email from the ASF dual-hosted git repository.
diegomedina24 pushed a change to branch dm/migrate-csv-to-api-v1
in repository https://gitbox.apache.org/repos/asf/superset.git
at 81252926dc chore: Migrate /superset/csv/<client_id> to API v1
This branch includes the following new commits:
new 81252926dc chore: Migrate /superset/csv/<client_id> to API v1
The 1 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
[superset] 01/01: chore: Migrate /superset/csv/ to API v1
Posted by di...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
diegomedina24 pushed a commit to branch dm/migrate-csv-to-api-v1
in repository https://gitbox.apache.org/repos/asf/superset.git
commit 81252926dcfa56473ece5091acf209752d7fb130
Author: Diego Medina <di...@gmail.com>
AuthorDate: Mon Jan 30 19:43:39 2023 -0300
chore: Migrate /superset/csv/<client_id> to API v1
---
docs/static/resources/openapi.json | 259 +++++++++++++++++----
.../src/SqlLab/components/ResultSet/index.tsx | 11 +-
superset/sqllab/api.py | 73 +++++-
superset/sqllab/commands/export.py | 133 +++++++++++
superset/sqllab/schemas.py | 8 +
superset/views/core.py | 1 +
tests/integration_tests/sql_lab/api_tests.py | 38 ++-
tests/integration_tests/sql_lab/commands_tests.py | 252 +++++++++++++++++++-
8 files changed, 710 insertions(+), 65 deletions(-)
diff --git a/docs/static/resources/openapi.json b/docs/static/resources/openapi.json
index 8077af91c1..c88379f6b6 100644
--- a/docs/static/resources/openapi.json
+++ b/docs/static/resources/openapi.json
@@ -1783,7 +1783,7 @@
"readOnly": true
},
"created_by": {
- "$ref": "#/components/schemas/ChartDataRestApi.get_list.User3"
+ "$ref": "#/components/schemas/ChartDataRestApi.get_list.User"
},
"created_on_delta_humanized": {
"readOnly": true
@@ -1830,7 +1830,7 @@
"type": "string"
},
"last_saved_by": {
- "$ref": "#/components/schemas/ChartDataRestApi.get_list.User"
+ "$ref": "#/components/schemas/ChartDataRestApi.get_list.User3"
},
"owners": {
"$ref": "#/components/schemas/ChartDataRestApi.get_list.User2"
@@ -2575,7 +2575,7 @@
"readOnly": true
},
"created_by": {
- "$ref": "#/components/schemas/ChartRestApi.get_list.User3"
+ "$ref": "#/components/schemas/ChartRestApi.get_list.User"
},
"created_on_delta_humanized": {
"readOnly": true
@@ -2622,7 +2622,7 @@
"type": "string"
},
"last_saved_by": {
- "$ref": "#/components/schemas/ChartRestApi.get_list.User"
+ "$ref": "#/components/schemas/ChartRestApi.get_list.User3"
},
"owners": {
"$ref": "#/components/schemas/ChartRestApi.get_list.User2"
@@ -3400,7 +3400,7 @@
"type": "string"
},
"changed_by": {
- "$ref": "#/components/schemas/DashboardRestApi.get_list.User"
+ "$ref": "#/components/schemas/DashboardRestApi.get_list.User1"
},
"changed_by_name": {
"readOnly": true
@@ -3415,7 +3415,7 @@
"readOnly": true
},
"created_by": {
- "$ref": "#/components/schemas/DashboardRestApi.get_list.User2"
+ "$ref": "#/components/schemas/DashboardRestApi.get_list.User"
},
"created_on_delta_humanized": {
"readOnly": true
@@ -3441,7 +3441,7 @@
"type": "string"
},
"owners": {
- "$ref": "#/components/schemas/DashboardRestApi.get_list.User1"
+ "$ref": "#/components/schemas/DashboardRestApi.get_list.User2"
},
"position_json": {
"nullable": true,
@@ -3500,25 +3500,16 @@
"last_name": {
"maxLength": 64,
"type": "string"
- },
- "username": {
- "maxLength": 64,
- "type": "string"
}
},
"required": [
"first_name",
- "last_name",
- "username"
+ "last_name"
],
"type": "object"
},
"DashboardRestApi.get_list.User1": {
"properties": {
- "email": {
- "maxLength": 64,
- "type": "string"
- },
"first_name": {
"maxLength": 64,
"type": "string"
@@ -3537,7 +3528,6 @@
}
},
"required": [
- "email",
"first_name",
"last_name",
"username"
@@ -3546,6 +3536,10 @@
},
"DashboardRestApi.get_list.User2": {
"properties": {
+ "email": {
+ "maxLength": 64,
+ "type": "string"
+ },
"first_name": {
"maxLength": 64,
"type": "string"
@@ -3557,11 +3551,17 @@
"last_name": {
"maxLength": 64,
"type": "string"
+ },
+ "username": {
+ "maxLength": 64,
+ "type": "string"
}
},
"required": [
+ "email",
"first_name",
- "last_name"
+ "last_name",
+ "username"
],
"type": "object"
},
@@ -4881,7 +4881,7 @@
"type": "integer"
},
"changed_by": {
- "$ref": "#/components/schemas/DatasetRestApi.get.User"
+ "$ref": "#/components/schemas/DatasetRestApi.get.User1"
},
"changed_on": {
"format": "date-time",
@@ -4895,7 +4895,7 @@
"$ref": "#/components/schemas/DatasetRestApi.get.TableColumn"
},
"created_by": {
- "$ref": "#/components/schemas/DatasetRestApi.get.User2"
+ "$ref": "#/components/schemas/DatasetRestApi.get.User"
},
"created_on": {
"format": "date-time",
@@ -4959,7 +4959,7 @@
"type": "integer"
},
"owners": {
- "$ref": "#/components/schemas/DatasetRestApi.get.User1"
+ "$ref": "#/components/schemas/DatasetRestApi.get.User2"
},
"schema": {
"maxLength": 255,
@@ -5173,23 +5173,14 @@
"maxLength": 64,
"type": "string"
},
- "id": {
- "format": "int32",
- "type": "integer"
- },
"last_name": {
"maxLength": 64,
"type": "string"
- },
- "username": {
- "maxLength": 64,
- "type": "string"
}
},
"required": [
"first_name",
- "last_name",
- "username"
+ "last_name"
],
"type": "object"
},
@@ -5199,21 +5190,30 @@
"maxLength": 64,
"type": "string"
},
+ "id": {
+ "format": "int32",
+ "type": "integer"
+ },
"last_name": {
"maxLength": 64,
"type": "string"
+ },
+ "username": {
+ "maxLength": 64,
+ "type": "string"
}
},
"required": [
"first_name",
- "last_name"
+ "last_name",
+ "username"
],
"type": "object"
},
"DatasetRestApi.get_list": {
"properties": {
"changed_by": {
- "$ref": "#/components/schemas/DatasetRestApi.get_list.User1"
+ "$ref": "#/components/schemas/DatasetRestApi.get_list.User"
},
"changed_by_name": {
"readOnly": true
@@ -5256,7 +5256,7 @@
"readOnly": true
},
"owners": {
- "$ref": "#/components/schemas/DatasetRestApi.get_list.User"
+ "$ref": "#/components/schemas/DatasetRestApi.get_list.User1"
},
"schema": {
"maxLength": 255,
@@ -5300,14 +5300,6 @@
"maxLength": 64,
"type": "string"
},
- "id": {
- "format": "int32",
- "type": "integer"
- },
- "last_name": {
- "maxLength": 64,
- "type": "string"
- },
"username": {
"maxLength": 64,
"type": "string"
@@ -5315,7 +5307,6 @@
},
"required": [
"first_name",
- "last_name",
"username"
],
"type": "object"
@@ -5326,6 +5317,14 @@
"maxLength": 64,
"type": "string"
},
+ "id": {
+ "format": "int32",
+ "type": "integer"
+ },
+ "last_name": {
+ "maxLength": 64,
+ "type": "string"
+ },
"username": {
"maxLength": 64,
"type": "string"
@@ -5333,6 +5332,7 @@
},
"required": [
"first_name",
+ "last_name",
"username"
],
"type": "object"
@@ -6933,7 +6933,7 @@
"type": "boolean"
},
"changed_by": {
- "$ref": "#/components/schemas/ReportScheduleRestApi.get_list.User"
+ "$ref": "#/components/schemas/ReportScheduleRestApi.get_list.User1"
},
"changed_on": {
"format": "date-time",
@@ -6949,7 +6949,7 @@
"type": "integer"
},
"created_by": {
- "$ref": "#/components/schemas/ReportScheduleRestApi.get_list.User2"
+ "$ref": "#/components/schemas/ReportScheduleRestApi.get_list.User"
},
"created_on": {
"format": "date-time",
@@ -6999,7 +6999,7 @@
"type": "string"
},
"owners": {
- "$ref": "#/components/schemas/ReportScheduleRestApi.get_list.User1"
+ "$ref": "#/components/schemas/ReportScheduleRestApi.get_list.User2"
},
"recipients": {
"$ref": "#/components/schemas/ReportScheduleRestApi.get_list.ReportRecipients"
@@ -7060,10 +7060,6 @@
"maxLength": 64,
"type": "string"
},
- "id": {
- "format": "int32",
- "type": "integer"
- },
"last_name": {
"maxLength": 64,
"type": "string"
@@ -7081,6 +7077,10 @@
"maxLength": 64,
"type": "string"
},
+ "id": {
+ "format": "int32",
+ "type": "integer"
+ },
"last_name": {
"maxLength": 64,
"type": "string"
@@ -9507,6 +9507,17 @@
},
"type": "object"
},
+ "sql_lab_export_csv_schema": {
+ "properties": {
+ "client_id": {
+ "type": "string"
+ }
+ },
+ "required": [
+ "client_id"
+ ],
+ "type": "object"
+ },
"sql_lab_get_results_schema": {
"properties": {
"key": {
@@ -16686,6 +16697,99 @@
]
}
},
+ "/api/v1/datasource/{datasource_type}/{datasource_id}/column/{column_name}/values/": {
+ "get": {
+ "parameters": [
+ {
+ "description": "The type of datasource",
+ "in": "path",
+ "name": "datasource_type",
+ "required": true,
+ "schema": {
+ "type": "string"
+ }
+ },
+ {
+ "description": "The id of the datasource",
+ "in": "path",
+ "name": "datasource_id",
+ "required": true,
+ "schema": {
+ "type": "integer"
+ }
+ },
+ {
+ "description": "The name of the column to get values for",
+ "in": "path",
+ "name": "column_name",
+ "required": true,
+ "schema": {
+ "type": "string"
+ }
+ }
+ ],
+ "responses": {
+ "200": {
+ "content": {
+ "application/json": {
+ "schema": {
+ "properties": {
+ "result": {
+ "items": {
+ "oneOf": [
+ {
+ "type": "string"
+ },
+ {
+ "type": "integer"
+ },
+ {
+ "type": "number"
+ },
+ {
+ "type": "boolean"
+ },
+ {
+ "type": "object"
+ }
+ ]
+ },
+ "type": "array"
+ }
+ },
+ "type": "object"
+ }
+ }
+ },
+ "description": "A List of distinct values for the column"
+ },
+ "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"
+ }
+ },
+ "security": [
+ {
+ "jwt": []
+ }
+ ],
+ "summary": "Get possible values for a datasource column",
+ "tags": [
+ "Datasources"
+ ]
+ }
+ },
"/api/v1/embedded_dashboard/{uuid}": {
"get": {
"description": "Get a report schedule log",
@@ -19799,6 +19903,59 @@
]
}
},
+ "/api/v1/sqllab/export/": {
+ "get": {
+ "parameters": [
+ {
+ "content": {
+ "application/json": {
+ "schema": {
+ "$ref": "#/components/schemas/sql_lab_export_csv_schema"
+ }
+ }
+ },
+ "in": "query",
+ "name": "q"
+ }
+ ],
+ "responses": {
+ "200": {
+ "content": {
+ "text/csv": {
+ "schema": {
+ "type": {}
+ }
+ }
+ },
+ "description": "SQL query results"
+ },
+ "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"
+ }
+ },
+ "security": [
+ {
+ "jwt": []
+ }
+ ],
+ "summary": "Exports the SQL Query results to a CSV",
+ "tags": [
+ "SQL Lab"
+ ]
+ }
+ },
"/api/v1/sqllab/results/": {
"get": {
"parameters": [
diff --git a/superset-frontend/src/SqlLab/components/ResultSet/index.tsx b/superset-frontend/src/SqlLab/components/ResultSet/index.tsx
index 81a4e47a11..47f8d4acdd 100644
--- a/superset-frontend/src/SqlLab/components/ResultSet/index.tsx
+++ b/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>
)}
diff --git a/superset/sqllab/api.py b/superset/sqllab/api.py
index 283c3ab638..75ac83c0c6 100644
--- a/superset/sqllab/api.py
+++ b/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.queries.dao import QueryDAO
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.query_render import SqlQueryRenderImpl
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.sqllab_execution_context import SqlJsonExecutionContext
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,6 +75,7 @@ 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 = (
@@ -79,6 +83,73 @@ class SqlLabRestApi(BaseSupersetApi):
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: {}
+ 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
diff --git a/superset/sqllab/commands/export.py b/superset/sqllab/commands/export.py
new file mode 100644
index 0000000000..483f8971f5
--- /dev/null
+++ b/superset/sqllab/commands/export.py
@@ -0,0 +1,133 @@
+# 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 superset.sql_parse import ParsedQuery
+from superset.sqllab.limiting_factor import LimitingFactor
+
+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.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,
+ }
diff --git a/superset/sqllab/schemas.py b/superset/sqllab/schemas.py
index f238fda5c9..428cdb89bb 100644
--- a/superset/sqllab/schemas.py
+++ b/superset/sqllab/schemas.py
@@ -24,6 +24,14 @@ sql_lab_get_results_schema = {
"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)
diff --git a/superset/views/core.py b/superset/views/core.py
index 8d632dcde2..283ea5df0d 100755
--- a/superset/views/core.py
+++ b/superset/views/core.py
@@ -2392,6 +2392,7 @@ class Superset(BaseSupersetView): # pylint: disable=too-many-public-methods
@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)
diff --git a/tests/integration_tests/sql_lab/api_tests.py b/tests/integration_tests/sql_lab/api_tests.py
index 4c2080ad4c..73506a38d2 100644
--- a/tests/integration_tests/sql_lab/api_tests.py
+++ b/tests/integration_tests/sql_lab/api_tests.py
@@ -19,6 +19,9 @@
import datetime
import json
import random
+import csv
+import pandas as pd
+import io
import pytest
import prison
@@ -26,7 +29,7 @@ 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,36 @@ class TestSqlLabApi(SupersetTestCase):
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_run_no_results_backend_executed_sql(self, get_df_mock: mock.Mock) -> None:
+ self.login()
+
+ database = Database(database_name="my_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))
diff --git a/tests/integration_tests/sql_lab/commands_tests.py b/tests/integration_tests/sql_lab/commands_tests.py
index 74c1fe7082..ea71ee6123 100644
--- a/tests/integration_tests/sql_lab/commands_tests.py
+++ b/tests/integration_tests/sql_lab/commands_tests.py
@@ -15,23 +15,253 @@
# specific language governing permissions and limitations
# under the License.
from unittest import mock, skip
-from unittest.mock import patch
+from unittest.mock import Mock, patch
+import pandas as pd
import pytest
from superset import db, sql_lab
from superset.common.db_query_status import QueryStatus
-from superset.errors import SupersetErrorType
-from superset.exceptions import SerializationError, SupersetErrorException
+from superset.errors import ErrorLevel, SupersetErrorType
+from superset.exceptions import (
+ SerializationError,
+ SupersetError,
+ SupersetErrorException,
+ SupersetSecurityException,
+)
from superset.models.core import Database
from superset.models.sql_lab import Query
-from superset.sqllab.commands import results
+from superset.sqllab.commands import export, results
+from superset.sqllab.limiting_factor import LimitingFactor
from superset.utils import core as utils
from tests.integration_tests.base_tests import SupersetTestCase
+class TestSqlResultExportCommand(SupersetTestCase):
+ def test_validation_query_not_found(self) -> None:
+ command = export.SqlResultExportCommand("test")
+
+ database = Database(database_name="my_database", sqlalchemy_uri="sqlite://")
+ query_obj = Query(
+ client_id="foo",
+ database=database,
+ tab_name="test_tab",
+ sql_editor_id="test_editor_id",
+ sql="select * from bar",
+ select_sql="select * from bar",
+ executed_sql="select * from bar",
+ limit=100,
+ select_as_cta=False,
+ rows=104,
+ error_message="none",
+ results_key="test_abc",
+ )
+
+ db.session.add(database)
+ db.session.add(query_obj)
+
+ with pytest.raises(SupersetErrorException) as ex_info:
+ command.run()
+ assert ex_info.value.error.error_type == SupersetErrorType.RESULTS_BACKEND_ERROR
+
+ def test_validation_invalid_access(self) -> None:
+ command = export.SqlResultExportCommand("test")
+
+ database = Database(database_name="my_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="select * from bar",
+ executed_sql="select * from bar",
+ limit=100,
+ select_as_cta=False,
+ rows=104,
+ error_message="none",
+ results_key="test_abc",
+ )
+
+ db.session.add(database)
+ db.session.add(query_obj)
+
+ with mock.patch(
+ "superset.security_manager.raise_for_access",
+ side_effect=SupersetSecurityException(
+ SupersetError(
+ "dummy",
+ SupersetErrorType.DATASOURCE_SECURITY_ACCESS_ERROR,
+ ErrorLevel.ERROR,
+ )
+ ),
+ ):
+ with pytest.raises(SupersetErrorException) as ex_info:
+ command.run()
+ assert (
+ ex_info.value.error.error_type
+ == SupersetErrorType.QUERY_SECURITY_ACCESS_ERROR
+ )
+
+ @patch("superset.models.sql_lab.Query.raise_for_access", lambda _: None)
+ @patch("superset.models.core.Database.get_df")
+ def test_run_no_results_backend_select_sql(self, get_df_mock: Mock) -> None:
+ command = export.SqlResultExportCommand("test")
+
+ database = Database(database_name="my_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="select * from bar",
+ executed_sql="select * from bar",
+ 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]})
+
+ result = command.run()
+
+ data = result.get("data")
+ count = result.get("count")
+ query = result.get("query")
+
+ assert data == "foo\n1\n2\n3\n"
+ assert count == 3
+ assert query.client_id == "test"
+
+ @patch("superset.models.sql_lab.Query.raise_for_access", lambda _: None)
+ @patch("superset.models.core.Database.get_df")
+ def test_run_no_results_backend_executed_sql(self, get_df_mock: Mock) -> None:
+ command = export.SqlResultExportCommand("test")
+
+ database = Database(database_name="my_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]})
+
+ result = command.run()
+
+ data = result.get("data")
+ count = result.get("count")
+ query = result.get("query")
+
+ assert data == "foo\n1\n2\n"
+ assert count == 2
+ assert query.client_id == "test"
+
+ @patch("superset.models.sql_lab.Query.raise_for_access", lambda _: None)
+ @patch("superset.models.core.Database.get_df")
+ def test_run_no_results_backend_executed_sql_limiting_factor(
+ self, get_df_mock: Mock
+ ) -> None:
+ command = export.SqlResultExportCommand("test")
+
+ database = Database(database_name="my_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",
+ limiting_factor=LimitingFactor.DROPDOWN,
+ )
+
+ db.session.add(database)
+ db.session.add(query_obj)
+
+ get_df_mock.return_value = pd.DataFrame({"foo": [1, 2, 3]})
+
+ result = command.run()
+
+ data = result.get("data")
+ count = result.get("count")
+ query = result.get("query")
+
+ assert data == "foo\n1\n"
+ assert count == 1
+ assert query.client_id == "test"
+
+ @patch("superset.models.sql_lab.Query.raise_for_access", lambda _: None)
+ @patch("superset.sqllab.commands.export.results_backend_use_msgpack", False)
+ def test_run_with_results_backend(self) -> None:
+ command = export.SqlResultExportCommand("test")
+
+ database = Database(database_name="my_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="select * from bar",
+ executed_sql="select * from bar",
+ limit=100,
+ select_as_cta=False,
+ rows=104,
+ error_message="none",
+ results_key="test_abc",
+ )
+
+ db.session.add(database)
+ db.session.add(query_obj)
+
+ data = [{"foo": i} for i in range(5)]
+ payload = {
+ "columns": [{"name": "foo"}],
+ "data": data,
+ }
+ serialized_payload = sql_lab._serialize_payload(payload, False)
+ compressed = utils.zlib_compress(serialized_payload)
+
+ export.results_backend = mock.Mock()
+ export.results_backend.get.return_value = compressed
+
+ result = command.run()
+
+ data = result.get("data")
+ count = result.get("count")
+ query = result.get("query")
+
+ assert data == "foo\n0\n1\n2\n3\n4\n"
+ assert count == 5
+ assert query.client_id == "test"
+
+
class TestSqlExecutionResultsCommand(SupersetTestCase):
- @mock.patch("superset.sqllab.commands.results.results_backend_use_msgpack", False)
+ @patch("superset.sqllab.commands.results.results_backend_use_msgpack", False)
def test_validation_no_results_backend(self) -> None:
results.results_backend = None
@@ -44,7 +274,7 @@ class TestSqlExecutionResultsCommand(SupersetTestCase):
== SupersetErrorType.RESULTS_BACKEND_NOT_CONFIGURED_ERROR
)
- @mock.patch("superset.sqllab.commands.results.results_backend_use_msgpack", False)
+ @patch("superset.sqllab.commands.results.results_backend_use_msgpack", False)
def test_validation_data_cannot_be_retrieved(self) -> None:
results.results_backend = mock.Mock()
results.results_backend.get.return_value = None
@@ -55,8 +285,8 @@ class TestSqlExecutionResultsCommand(SupersetTestCase):
command.run()
assert ex_info.value.error.error_type == SupersetErrorType.RESULTS_BACKEND_ERROR
- @mock.patch("superset.sqllab.commands.results.results_backend_use_msgpack", False)
- def test_validation_query_not_found(self) -> None:
+ @patch("superset.sqllab.commands.results.results_backend_use_msgpack", False)
+ def test_validation_data_not_found(self) -> None:
data = [{"col_0": i} for i in range(100)]
payload = {
"status": QueryStatus.SUCCESS,
@@ -75,8 +305,8 @@ class TestSqlExecutionResultsCommand(SupersetTestCase):
command.run()
assert ex_info.value.error.error_type == SupersetErrorType.RESULTS_BACKEND_ERROR
- @mock.patch("superset.sqllab.commands.results.results_backend_use_msgpack", False)
- def test_validation_query_not_found2(self) -> None:
+ @patch("superset.sqllab.commands.results.results_backend_use_msgpack", False)
+ def test_validation_query_not_found(self) -> None:
data = [{"col_0": i} for i in range(104)]
payload = {
"status": QueryStatus.SUCCESS,
@@ -120,7 +350,7 @@ class TestSqlExecutionResultsCommand(SupersetTestCase):
== SupersetErrorType.RESULTS_BACKEND_ERROR
)
- @mock.patch("superset.sqllab.commands.results.results_backend_use_msgpack", False)
+ @patch("superset.sqllab.commands.results.results_backend_use_msgpack", False)
def test_run_succeeds(self) -> None:
data = [{"col_0": i} for i in range(104)]
payload = {