alexandrusoare commented on code in PR #35478:
URL: https://github.com/apache/superset/pull/35478#discussion_r2428789159
##########
superset-frontend/src/SqlLab/components/ResultSet/index.tsx:
##########
@@ -300,6 +318,32 @@ const ResultSet = ({
const getExportCsvUrl = (clientId: string) =>
`/api/v1/sqllab/export/${clientId}/`;
+ const getStreamingExportUrl = () => `/api/v1/sqllab/export_streaming/`;
+
Review Comment:
I don t think there s a need to save the url like this, for
`getExportCsvUrl` made sense because it changes depending on the clientId
##########
superset-frontend/src/SqlLab/components/ResultSet/index.tsx:
##########
@@ -222,11 +228,23 @@ const ResultSet = ({
const [searchText, setSearchText] = useState('');
const [cachedData, setCachedData] = useState<Record<string, unknown>[]>([]);
const [showSaveDatasetModal, setShowSaveDatasetModal] = useState(false);
+ const [showStreamingModal, setShowStreamingModal] = useState(false);
const history = useHistory();
const dispatch = useDispatch();
const logAction = useLogAction({ queryId, sqlEditorId: query.sqlEditorId });
+ // Streaming export hook
Review Comment:
```suggestion
```
##########
superset-frontend/src/SqlLab/components/ResultSet/index.tsx:
##########
@@ -300,6 +318,32 @@ const ResultSet = ({
const getExportCsvUrl = (clientId: string) =>
`/api/v1/sqllab/export/${clientId}/`;
+ const getStreamingExportUrl = () => `/api/v1/sqllab/export_streaming/`;
+
+ const handleCloseStreamingModal = () => {
+ setShowStreamingModal(false);
+ resetExport();
+ };
+
+ // Check if streaming export should be used
+ const shouldUseStreamingExport = () => {
Review Comment:
No need for these comments
##########
superset/charts/data/api.py:
##########
@@ -425,7 +441,25 @@ def _get_data_response(
except ChartDataQueryFailedError as exc:
return self.response_400(message=exc.message)
- return self._send_chart_response(result, form_data, datasource)
+ return self._send_chart_response(
+ result, form_data, datasource, filename, expected_rows
+ )
+
+ def _extract_export_params_from_request(self) -> tuple[str | None, int |
None]:
+ """Extract filename and expected_rows from request for streaming
exports."""
+ filename = request.form.get("filename")
+ if filename:
+ logger.info("📁 FRONTEND PROVIDED FILENAME: %s", filename)
+
+ expected_rows = None
+ if expected_rows_str := request.form.get("expected_rows"):
+ try:
+ expected_rows = int(expected_rows_str)
+ logger.info("📊 FRONTEND PROVIDED EXPECTED ROWS: %d",
expected_rows)
Review Comment:
We don t use emojis in loggers
##########
superset-frontend/src/dashboard/components/gridComponents/Chart/Chart.jsx:
##########
@@ -544,6 +616,21 @@ const Chart = props => {
emitCrossFilters={emitCrossFilters}
/>
</ChartWrapper>
+
+ {/* Streaming Export Modal */}
Review Comment:
No need for this comment
##########
superset-frontend/src/SqlLab/components/ResultSet/index.tsx:
##########
@@ -351,21 +395,47 @@ const ResultSet = ({
css={copyButtonStyles}
buttonSize="small"
buttonStyle="secondary"
- href={getExportCsvUrl(query.id)}
+ href={
+ !shouldUseStreamingExport()
+ ? getExportCsvUrl(query.id)
+ : undefined
+ }
Review Comment:
Can we try something like => href={!shouldUseStreamingExport() &&
getExportCsvUrl(query.id)}
##########
superset-frontend/src/SqlLab/components/ResultSet/index.tsx:
##########
@@ -351,21 +395,47 @@ const ResultSet = ({
css={copyButtonStyles}
buttonSize="small"
buttonStyle="secondary"
- href={getExportCsvUrl(query.id)}
+ href={
+ !shouldUseStreamingExport()
+ ? getExportCsvUrl(query.id)
+ : undefined
+ }
data-test="export-csv-button"
- onClick={() => {
- logAction(LOG_ACTIONS_SQLLAB_DOWNLOAD_CSV, {});
- if (
- limitingFactor === LimitingFactor.Dropdown &&
- limit === rowsCount
- ) {
- Modal.warning({
- title: t('Download is on the way'),
- content: t(
- 'Downloading %(rows)s rows based on the LIMIT
configuration. If you want the entire result set, you need to adjust the
LIMIT.',
- { rows: rowsCount.toLocaleString() },
- ),
+ onClick={e => {
+ const useStreaming = shouldUseStreamingExport();
+
+ if (useStreaming) {
+ e.preventDefault();
+ setShowStreamingModal(true);
+
+ const timestamp = new Date()
+ .toISOString()
+ .slice(0, 19)
+ .replace(/[-:]/g, '')
+ .replace('T', '_');
+ const filename = `sqllab_${query.id}_${timestamp}.csv`;
Review Comment:
Can we check if there s a util that already handles the creation of a
timestamp?
##########
superset/commands/sql_lab/streaming_export_command.py:
##########
@@ -0,0 +1,239 @@
+# 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.
+"""Command for streaming CSV exports of SQL Lab query results."""
+
+from __future__ import annotations
+
+import csv
+import io
+import logging
+import time
+from typing import Callable, Generator, TYPE_CHECKING
+
+from flask import current_app as app
+from flask_babel import gettext as __
+
+from superset import db
+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 SQLScript
+from superset.sqllab.limiting_factor import LimitingFactor
+
+if TYPE_CHECKING:
+ pass
Review Comment:
I noticed that we have two streaming_export_command.py, is there a way to
dry this up?
##########
superset-frontend/src/components/StreamingExportModal/useStreamingExport.test.ts:
##########
@@ -0,0 +1,128 @@
+/**
+ * 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.
+ */
+import { renderHook, act } from '@testing-library/react-hooks';
+import { useStreamingExport } from './useStreamingExport';
+import { ExportStatus } from './StreamingExportModal';
+
+// Mock SupersetClient
+jest.mock('@superset-ui/core', () => ({
+ ...jest.requireActual('@superset-ui/core'),
+ SupersetClient: {
+ getCSRFToken: jest.fn(() => Promise.resolve('mock-csrf-token')),
+ },
+}));
+
+// Mock URL APIs
Review Comment:
No need for comments here eiter
##########
superset-frontend/src/SqlLab/components/ResultSet/index.tsx:
##########
@@ -222,11 +228,23 @@ const ResultSet = ({
const [searchText, setSearchText] = useState('');
const [cachedData, setCachedData] = useState<Record<string, unknown>[]>([]);
const [showSaveDatasetModal, setShowSaveDatasetModal] = useState(false);
+ const [showStreamingModal, setShowStreamingModal] = useState(false);
const history = useHistory();
const dispatch = useDispatch();
const logAction = useLogAction({ queryId, sqlEditorId: query.sqlEditorId });
+ // Streaming export hook
+ const { progress, startExport, resetExport, retryExport } =
+ useStreamingExport({
+ onComplete: () => {
+ // Modal will show download button
Review Comment:
```suggestion
```
##########
superset/charts/data/api.py:
##########
@@ -462,3 +496,116 @@ def _create_query_context_from_form(
return ChartDataQueryContextSchema().load(form_data)
except KeyError as ex:
raise ValidationError("Request is incorrect") from ex
+
+ def _should_use_streaming(
+ self, result: dict[Any, Any], form_data: dict[str, Any] | None = None
+ ) -> bool:
+ """Determine if streaming should be used based on actual row count
threshold."""
+ from flask import current_app as app
+
+ query_context = result["query_context"]
+ result_format = query_context.result_format
+
+ # Only support CSV streaming currently
+ if result_format.lower() != "csv":
+ return False
+
+ # Get streaming threshold from config
+ threshold = app.config.get("CSV_STREAMING_ROW_THRESHOLD", 100000)
+
+ # Extract actual row count (same logic as frontend)
+ actual_row_count = None
+ viz_type = form_data.get("viz_type") if form_data else None
+
+ # For table viz, try to get actual row count from query results
+ if viz_type == "table" and result.get("queries"):
+ # Check if we have rowcount in the second query result (like
frontend does)
+ queries = result.get("queries", [])
+ if len(queries) > 1 and queries[1].get("data"):
+ data = queries[1]["data"]
+ if isinstance(data, list) and len(data) > 0:
+ actual_row_count = data[0].get("rowcount")
+
+ # Fallback to row_limit if actual count not available
Review Comment:
Let s get rid of the comments everywhere, and keep only those that are needed
##########
superset/charts/data/api.py:
##########
@@ -462,3 +496,116 @@ def _create_query_context_from_form(
return ChartDataQueryContextSchema().load(form_data)
except KeyError as ex:
raise ValidationError("Request is incorrect") from ex
+
+ def _should_use_streaming(
+ self, result: dict[Any, Any], form_data: dict[str, Any] | None = None
Review Comment:
is there something better than dict[Any, Any]?
##########
superset-frontend/src/components/StreamingExportModal/useStreamingExport.ts:
##########
@@ -0,0 +1,367 @@
+/**
+ * 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.
+ */
+import { useState, useCallback, useRef, useEffect } from 'react';
+import { SupersetClient } from '@superset-ui/core';
+import { ExportStatus, StreamingProgress } from './StreamingExportModal';
+
+interface UseStreamingExportOptions {
+ onComplete?: (downloadUrl: string, filename: string) => void;
+ onError?: (error: string) => void;
+}
+
+interface StreamingExportPayload {
+ [key: string]: unknown;
+}
Review Comment:
Is there a better type other than unknown?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
