From 53c9eba363c3f81cd542ee58b93c0d66e315ce04 Mon Sep 17 00:00:00 2001 From: amaannawab923 Date: Thu, 2 Oct 2025 21:50:24 +0530 Subject: [PATCH 01/57] fix(ag-grid-table): remove enterprise features to use community version --- .../plugin-chart-ag-grid-table/src/AgGridTable/index.tsx | 4 ---- 1 file changed, 4 deletions(-) diff --git a/superset-frontend/plugins/plugin-chart-ag-grid-table/src/AgGridTable/index.tsx b/superset-frontend/plugins/plugin-chart-ag-grid-table/src/AgGridTable/index.tsx index c294fb2f06d3..955624636600 100644 --- a/superset-frontend/plugins/plugin-chart-ag-grid-table/src/AgGridTable/index.tsx +++ b/superset-frontend/plugins/plugin-chart-ag-grid-table/src/AgGridTable/index.tsx @@ -131,10 +131,7 @@ const AgGridDataTable: FunctionComponent = memo( const defaultColDef = useMemo( () => ({ - flex: 1, filter: true, - enableRowGroup: true, - enableValue: true, sortable: true, resizable: true, minWidth: 100, @@ -312,7 +309,6 @@ const AgGridDataTable: FunctionComponent = memo( onCellClicked={handleCrossFilter} initialState={gridInitialState} suppressAggFuncInHeader - rowGroupPanelShow="always" enableCellTextSelection quickFilterText={serverPagination ? '' : quickFilterText} suppressMovableColumns={!allowRearrangeColumns} From 49cbbf49cfbf6ddbb18b15c69c52fa9032d0e16c Mon Sep 17 00:00:00 2001 From: amaannawab923 Date: Fri, 3 Oct 2025 12:14:52 +0530 Subject: [PATCH 02/57] feat: Add frontend streaming CSV export with progress modal Pulled from commit a4c85fa293333ce66f16f7da7d6dec7185c76496 Frontend Components: - StreamingExportModal: Modal with real-time progress tracking - Shows rows processed, file size, speed (rows/s, MB/s) - Progress bar with percentage completion - Cancel/retry functionality - Auto-download on completion - useStreamingExport hook: Manages streaming export lifecycle - Fetch API ReadableStream for chunked download - Real-time row counting from chunk data - Performance metrics (speed, throughput) - AbortController for cancellation - Blob creation for final download Integration: - Chart.jsx: Added streaming modal and export handling - exploreUtils: Modified exportChart() for streaming detection - Passes onStartStreamingExport callback for CSV - Extracts expectedRows for accurate progress - Fallback to legacy export for non-CSV User Experience: - Seamless CSV export with visual progress - Real-time feedback on large exports - No page navigation during export - Graceful error handling with retry option --- .../StreamingExportModal.tsx | 355 +++++++++++++ .../components/StreamingExportModal/index.ts | 21 + .../useStreamingExport.ts | 496 ++++++++++++++++++ .../components/gridComponents/Chart/Chart.jsx | 59 ++- .../src/explore/exploreUtils/index.js | 51 +- 5 files changed, 979 insertions(+), 3 deletions(-) create mode 100644 superset-frontend/src/components/StreamingExportModal/StreamingExportModal.tsx create mode 100644 superset-frontend/src/components/StreamingExportModal/index.ts create mode 100644 superset-frontend/src/components/StreamingExportModal/useStreamingExport.ts diff --git a/superset-frontend/src/components/StreamingExportModal/StreamingExportModal.tsx b/superset-frontend/src/components/StreamingExportModal/StreamingExportModal.tsx new file mode 100644 index 000000000000..6bbd5ad12827 --- /dev/null +++ b/superset-frontend/src/components/StreamingExportModal/StreamingExportModal.tsx @@ -0,0 +1,355 @@ +/** + * 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 React from 'react'; +import { styled, t } from '@superset-ui/core'; +import { Modal, Button, Typography, Space, Progress } from 'antd'; +import { + CheckCircleOutlined, + CloseCircleOutlined, + DownloadOutlined, + LoadingOutlined, + StopOutlined, +} from '@ant-design/icons'; + +const { Text, Title } = Typography; + +export interface StreamingProgress { + rowsProcessed: number; + totalRows?: number; // total expected rows for percentage calculation + totalSize: number; // in bytes + speed: number; // rows per second + mbPerSecond: number; + elapsedTime: number; // seconds + estimatedTimeRemaining?: number; // seconds + status: 'streaming' | 'completed' | 'error' | 'cancelled'; + downloadUrl?: string; + filename?: string; + error?: string; +} + +interface StreamingExportModalProps { + visible: boolean; + onCancel: () => void; + onRetry?: () => void; + progress: StreamingProgress; + exportType: 'csv' | 'xlsx'; +} + +const ModalContent = styled.div` + padding: ${({ theme }) => theme.sizeUnit * 4}px 0; + min-height: 200px; +`; + +const ProgressSection = styled.div` + margin: ${({ theme }) => theme.sizeUnit * 6}px 0; +`; + +const StatsGrid = styled.div` + display: grid; + grid-template-columns: 1fr 1fr; + gap: ${({ theme }) => theme.sizeUnit * 4}px; + margin: ${({ theme }) => theme.sizeUnit * 4}px 0; +`; + +const StatItem = styled.div` + text-align: center; + padding: ${({ theme }) => theme.sizeUnit * 3}px; + background: ${({ theme }) => theme.colorFillAlter}; + border-radius: ${({ theme }) => theme.borderRadius}px; +`; + +const ActionButtons = styled.div` + display: flex; + justify-content: center; + gap: ${({ theme }) => theme.sizeUnit * 3}px; + margin-top: ${({ theme }) => theme.sizeUnit * 6}px; +`; + +const StatusIcon = styled.div` + display: flex; + justify-content: center; + margin-bottom: ${({ theme }) => theme.sizeUnit * 4}px; + + .anticon { + font-size: 48px; + } +`; + +const formatFileSize = (bytes: number): string => { + if (bytes === 0) return '0 B'; + const k = 1024; + const sizes = ['B', 'KB', 'MB', 'GB']; + const i = Math.floor(Math.log(bytes) / Math.log(k)); + return `${parseFloat((bytes / Math.pow(k, i)).toFixed(1))} ${sizes[i]}`; +}; + +const formatTime = (seconds: number): string => { + if (seconds < 60) return `${Math.round(seconds)}s`; + const minutes = Math.floor(seconds / 60); + const remainingSeconds = Math.round(seconds % 60); + return `${minutes}m ${remainingSeconds}s`; +}; + +const formatNumber = (num: number): string => + new Intl.NumberFormat().format(num); + +const StreamingExportModal: React.FC = ({ + visible, + onCancel, + onRetry, + progress, + exportType, +}) => { + const { + rowsProcessed, + totalSize, + speed, + mbPerSecond, + elapsedTime, + estimatedTimeRemaining, + status, + downloadUrl, + filename, + error, + } = progress; + + const getProgressPercentage = (): number => { + if (status === 'completed') return 100; + + // Calculate actual percentage based on rows processed vs total expected rows + if (progress.totalRows && progress.totalRows > 0) { + // Use rows processed even if it's 0 - this allows for 0% at start + const percentage = Math.min( + 99, + (rowsProcessed / progress.totalRows) * 100, + ); + const rounded = Math.round(percentage); + + // ๐Ÿ” DEBUG: Log percentage calculation + console.log('๐ŸŽฏ FRONTEND MODAL PERCENTAGE:', { + rowsProcessed, + totalRows: progress.totalRows, + rawPercentage: percentage.toFixed(2), + roundedPercentage: rounded, + status, + }); + + return rounded; + } + + // Fallback: estimate based on time if we have estimatedTimeRemaining + if (estimatedTimeRemaining && elapsedTime > 0) { + const totalEstimatedTime = elapsedTime + estimatedTimeRemaining; + const percentage = Math.min( + 95, + Math.round((elapsedTime / totalEstimatedTime) * 100), + ); + + console.log('โฐ FRONTEND MODAL TIME-BASED:', { + elapsedTime, + estimatedTimeRemaining, + totalEstimatedTime, + percentage, + }); + + return percentage; + } + + // Default fallback for streaming status + const fallback = status === 'streaming' ? 10 : 0; + console.log('๐Ÿ”„ FRONTEND MODAL FALLBACK:', { + status, + fallback, + hasRows: rowsProcessed > 0, + hasTotalRows: !!progress.totalRows, + reasonForFallback: !progress.totalRows ? 'No totalRows' : 'Unknown', + }); + + return fallback; + }; + + const renderStatusIcon = () => { + switch (status) { + case 'streaming': + return ; + case 'completed': + return ; + case 'error': + return ; + case 'cancelled': + return ; + default: + return ; + } + }; + + const getTitle = () => t('CSV Export'); + + const handleDownload = () => { + if (downloadUrl && filename) { + const link = document.createElement('a'); + link.href = downloadUrl; + link.download = filename; + document.body.appendChild(link); + link.click(); + document.body.removeChild(link); + onCancel(); // Close modal after download + } + }; + + const renderContent = () => { + if (status === 'error') { + return ( + + {renderStatusIcon()} + + {error || t('An error occurred during export')} + + + {onRetry && ( + + )} + + + + ); + } + + if (status === 'cancelled') { + return ( + + {renderStatusIcon()} + + {t('Export was cancelled')} + + + {onRetry && ( + + )} + + + + ); + } + + if (status === 'completed') { + return ( + + {renderStatusIcon()} + + {t( + 'Successfully exported %s rows (%s)', + formatNumber(rowsProcessed), + formatFileSize(totalSize), + )} + + {filename && ( + + {filename} + + )} + + + + + + ); + } + + // Streaming status + return ( + + + `${Math.round(percent || 0)}%`} + /> + + {filename + ? t('Processing export for %s', filename) + : t( + 'Processing export for {dashboard_name}_{YYYY-MM-DD}_{HHMMSS}.csv', + )} + + + + + + + + + ); + }; + + return ( + + {renderContent()} + + ); +}; + +export default StreamingExportModal; diff --git a/superset-frontend/src/components/StreamingExportModal/index.ts b/superset-frontend/src/components/StreamingExportModal/index.ts new file mode 100644 index 000000000000..049a4d2af0f1 --- /dev/null +++ b/superset-frontend/src/components/StreamingExportModal/index.ts @@ -0,0 +1,21 @@ +/** + * 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. + */ +export { default as StreamingExportModal } from './StreamingExportModal'; +export type { StreamingProgress } from './StreamingExportModal'; +export { useStreamingExport } from './useStreamingExport'; \ No newline at end of file diff --git a/superset-frontend/src/components/StreamingExportModal/useStreamingExport.ts b/superset-frontend/src/components/StreamingExportModal/useStreamingExport.ts new file mode 100644 index 000000000000..b021f936602f --- /dev/null +++ b/superset-frontend/src/components/StreamingExportModal/useStreamingExport.ts @@ -0,0 +1,496 @@ +/** + * 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 } from 'react'; +import { SupersetClient } from '@superset-ui/core'; +import { StreamingProgress } from './StreamingExportModal'; + +interface UseStreamingExportOptions { + onProgress?: (progress: StreamingProgress) => void; + onComplete?: (downloadUrl: string, filename: string) => void; + onError?: (error: string) => void; +} + +interface StreamingExportParams { + url: string; + payload: any; + filename?: string; + exportType: 'csv' | 'xlsx'; + expectedRows?: number; // Total expected rows for progress calculation +} + +export const useStreamingExport = (options: UseStreamingExportOptions = {}) => { + const [progress, setProgress] = useState({ + rowsProcessed: 0, + totalRows: undefined, + totalSize: 0, + speed: 0, + mbPerSecond: 0, + elapsedTime: 0, + status: 'streaming', + }); + const [isExporting, setIsExporting] = useState(false); + const abortControllerRef = useRef(null); + const startTimeRef = useRef(0); + const lastUpdateRef = useRef(0); + const bytesHistoryRef = useRef>([]); + + const updateProgress = useCallback( + (updates: Partial) => { + setProgress(prev => { + const newProgress = { ...prev, ...updates }; + + // ๐Ÿ” DEBUG: Log every progress update + console.log('๐Ÿ“Š FRONTEND PROGRESS UPDATE:', { + rowsProcessed: newProgress.rowsProcessed, + totalRows: newProgress.totalRows, + percentage: newProgress.totalRows + ? Math.round( + (newProgress.rowsProcessed / newProgress.totalRows) * 100, + ) + : 'N/A', + totalSize: `${(newProgress.totalSize / 1024).toFixed(1)}KB`, + status: newProgress.status, + elapsedTime: `${newProgress.elapsedTime.toFixed(1)}s`, + }); + + options.onProgress?.(newProgress); + return newProgress; + }); + }, + [options], + ); + + const parseProgressFromHeaders = useCallback( + (headers: Headers, elapsedTime: number) => { + // Try to extract progress from response headers if backend provides them + const progressHeader = headers.get('X-Export-Progress'); + if (progressHeader) { + try { + const [rows, bytes, rate] = progressHeader.split(',').map(Number); + return { + rowsProcessed: rows || 0, + totalSize: bytes || 0, + speed: rate || 0, + mbPerSecond: bytes > 0 ? bytes / (1024 * 1024) / elapsedTime : 0, + elapsedTime, + }; + } catch (e) { + console.warn('Failed to parse progress header:', e); + } + } + return null; + }, + [], + ); + + const estimateProgressFromSize = useCallback( + (currentSize: number, elapsedTime: number, actualRows: number = 0) => { + const now = Date.now(); + + // Track bytes over time for smoother speed calculation + bytesHistoryRef.current.push({ time: now, bytes: currentSize }); + + // Keep only last 10 seconds of data for rolling average + const cutoffTime = now - 10000; + bytesHistoryRef.current = bytesHistoryRef.current.filter( + entry => entry.time > cutoffTime, + ); + + if (elapsedTime <= 0) { + return { + rowsProcessed: actualRows, + totalSize: currentSize, + speed: 0, + mbPerSecond: 0, + elapsedTime, + }; + } + + // Calculate speed based on recent data points for smoother updates + let mbPerSecond = 0; + if (bytesHistoryRef.current.length >= 2) { + const oldest = bytesHistoryRef.current[0]; + const newest = + bytesHistoryRef.current[bytesHistoryRef.current.length - 1]; + const timeDiff = (newest.time - oldest.time) / 1000; + const bytesDiff = newest.bytes - oldest.bytes; + + if (timeDiff > 0) { + mbPerSecond = bytesDiff / (1024 * 1024) / timeDiff; + } + } + + // If rolling average isn't available yet, use overall average + if (mbPerSecond === 0) { + mbPerSecond = currentSize / (1024 * 1024) / elapsedTime; + } + + // Use actual row count if available, otherwise estimate + let finalRowCount = actualRows; + if (actualRows === 0) { + // Fallback to size-based estimation only if no actual count + const avgBytesPerRow = + currentSize > 10000 + ? currentSize / Math.max(1, currentSize / 150) + : 200; + finalRowCount = Math.floor(currentSize / avgBytesPerRow); + } + + const rowsPerSecond = finalRowCount / elapsedTime; + + return { + rowsProcessed: finalRowCount, + totalSize: currentSize, + speed: rowsPerSecond, + mbPerSecond, + elapsedTime, + }; + }, + [], + ); + + const startExport = useCallback( + async ({ + url, + payload, + filename, + exportType, + expectedRows, + }: StreamingExportParams) => { + if (isExporting) { + console.warn('Export already in progress'); + return; + } + + setIsExporting(true); + abortControllerRef.current = new AbortController(); + startTimeRef.current = Date.now(); + lastUpdateRef.current = Date.now(); + bytesHistoryRef.current = []; // Reset bytes history for new export + + console.log('๐Ÿš€ FRONTEND STREAMING START:', { + url, + filename, + expectedRows, + exportType, + }); + + updateProgress({ + rowsProcessed: 0, + totalRows: expectedRows, + totalSize: 0, + speed: 0, + mbPerSecond: 0, + elapsedTime: 0, + status: 'streaming', + filename, + }); + + try { + // Initialize SupersetClient to ensure authentication and get CSRF token + await SupersetClient.init(); + const csrfToken = await SupersetClient.getCSRFToken(); + + // Use manual fetch for streaming while leveraging SupersetClient's authentication + const response = await fetch(url, { + method: 'POST', + headers: { + 'Content-Type': 'application/x-www-form-urlencoded', + ...(csrfToken && { 'X-CSRFToken': csrfToken }), + }, + body: new URLSearchParams({ + form_data: JSON.stringify(payload), + }), + signal: abortControllerRef.current.signal, + credentials: 'same-origin', + }); + + console.log('๐Ÿ“ก FRONTEND RESPONSE:', { + status: response.status, + statusText: response.statusText, + headers: Object.fromEntries(response.headers.entries()), + hasBody: !!response.body, + }); + + if (!response.ok) { + throw new Error( + `Export failed: ${response.status} ${response.statusText}`, + ); + } + + if (!response.body) { + throw new Error('Response body is not available for streaming'); + } + + console.log('๐Ÿ”ง STARTING STREAM READER'); + const reader = response.body.getReader(); + const chunks: Uint8Array[] = []; + let receivedLength = 0; + let actualRowCount = 0; + let accumulatedText = ''; + let chunkCount = 0; + + // More frequent progress updates for smoother UI + const progressInterval = setInterval(() => { + const now = Date.now(); + const elapsedTime = (now - startTimeRef.current) / 1000; + + // Always use the enhanced progress calculation with actual rows when available + const progressData = estimateProgressFromSize( + receivedLength, + elapsedTime, + actualRowCount, + ); + + // ๐Ÿ” DEBUG: Log interval-based progress updates + console.log('โฑ๏ธ FRONTEND INTERVAL UPDATE:', { + receivedLength, + actualRowCount, + elapsedTime: elapsedTime.toFixed(1), + progressData, + }); + + updateProgress(progressData); + }, 200); // Update every 200ms for smoother progress + + try { + // Read the streaming response + console.log('๐Ÿ”„ ENTERING STREAM READ LOOP'); + while (true) { + console.log('๐Ÿ”„ CALLING reader.read()...'); + const { done, value } = await reader.read(); + + console.log('๐Ÿ“ก READER RESULT:', { + done, + hasValue: !!value, + valueLength: value?.length, + }); + + if (done) { + console.log('โœ… STREAM READING COMPLETED'); + break; + } + + if (abortControllerRef.current?.signal.aborted) { + throw new Error('Export cancelled by user'); + } + + chunkCount++; + chunks.push(value); + const oldReceivedLength = receivedLength; + const oldRowCount = actualRowCount; + receivedLength += value.length; + + console.log('๐Ÿ“Š CHUNK RECEIVED:', { + chunkNumber: chunkCount, + chunkSize: value.length, + totalReceived: receivedLength, + }); + + // Convert chunk to text and count rows in real-time + try { + const chunkText = new TextDecoder().decode(value, { + stream: true, + }); + accumulatedText += chunkText; + + // Count complete lines (rows) in accumulated text + const lines = accumulatedText.split('\n'); + + // Keep the last incomplete line for next iteration + if (lines.length > 1) { + accumulatedText = lines[lines.length - 1]; // Keep incomplete line + + // Count completed lines (subtract 1 for header if this is first chunk) + const completedLines = lines.length - 1; + if (actualRowCount === 0 && completedLines > 0) { + // First chunk - subtract 1 for header row + actualRowCount = completedLines - 1; + } else { + // Subsequent chunks - add all completed lines + actualRowCount += completedLines; + } + } + + // ๐Ÿ” DEBUG: Log chunk processing and update progress immediately + if ( + actualRowCount !== oldRowCount || + receivedLength !== oldReceivedLength + ) { + console.log('๐Ÿ“ฆ FRONTEND CHUNK PROCESSED:', { + chunkSize: value.length, + totalBytes: receivedLength, + newRows: actualRowCount - oldRowCount, + totalRows: actualRowCount, + linesInChunk: lines.length - 1, + chunkTextPreview: `${chunkText.substring(0, 100)}...`, + }); + + // Immediately update progress with actual row count from chunk + const now = Date.now(); + const elapsedTime = (now - startTimeRef.current) / 1000; + const progressData = estimateProgressFromSize( + receivedLength, + elapsedTime, + actualRowCount, + ); + updateProgress(progressData); + } + } catch (decodeError) { + // If text decoding fails, fall back to byte-based estimation + console.warn( + 'โŒ Failed to decode chunk for row counting:', + decodeError, + ); + } + } + + clearInterval(progressInterval); + + // Create blob from chunks + const completeData = new Uint8Array(receivedLength); + let position = 0; + for (const chunk of chunks) { + completeData.set(chunk, position); + position += chunk.length; + } + + const mimeType = + exportType === 'csv' + ? 'text/csv;charset=utf-8' + : 'application/vnd.openxmlformats-officedocument.spreadsheetml.sheet'; + + const blob = new Blob([completeData], { type: mimeType }); + const downloadUrl = URL.createObjectURL(blob); + + const finalElapsedTime = (Date.now() - startTimeRef.current) / 1000; + + // Try to get accurate final count from server headers first + const headerProgress = parseProgressFromHeaders( + response.headers, + finalElapsedTime, + ); + + let finalProgress; + if (headerProgress && headerProgress.rowsProcessed > 0) { + // Use server-provided count if available + finalProgress = headerProgress; + } else { + // Use the real-time counted rows if available, otherwise parse final data + let finalRowCount = actualRowCount; + + if (finalRowCount === 0) { + // Fallback: parse the complete data for final count + try { + const text = new TextDecoder().decode(completeData); + const lines = text.split(/\r?\n/); + // Subtract 1 for header row, filter out empty lines + finalRowCount = + lines.filter(line => line.trim().length > 0).length - 1; + finalRowCount = Math.max(0, finalRowCount); + } catch (e) { + // Final fallback to size-based estimation + finalRowCount = Math.floor(receivedLength / 200); + } + } + + finalProgress = { + rowsProcessed: finalRowCount, + totalSize: receivedLength, + speed: finalRowCount / finalElapsedTime, + mbPerSecond: receivedLength / (1024 * 1024) / finalElapsedTime, + elapsedTime: finalElapsedTime, + }; + } + + updateProgress({ + ...finalProgress, + status: 'completed', + downloadUrl, + filename: filename || `export.${exportType}`, + }); + + options.onComplete?.(downloadUrl, filename || `export.${exportType}`); + } catch (streamError) { + clearInterval(progressInterval); + throw streamError; + } + } catch (error) { + const errorMessage = + error instanceof Error ? error.message : 'Unknown error occurred'; + + if ( + errorMessage.includes('cancelled') || + errorMessage.includes('aborted') + ) { + updateProgress({ + status: 'cancelled', + }); + } else { + updateProgress({ + status: 'error', + error: errorMessage, + }); + options.onError?.(errorMessage); + } + } finally { + setIsExporting(false); + abortControllerRef.current = null; + } + }, + [ + isExporting, + updateProgress, + parseProgressFromHeaders, + estimateProgressFromSize, + options, + ], + ); + + const cancelExport = useCallback(() => { + if (abortControllerRef.current) { + abortControllerRef.current.abort(); + updateProgress({ + status: 'cancelled', + }); + } + }, [updateProgress]); + + const resetExport = useCallback(() => { + setIsExporting(false); + abortControllerRef.current = null; + setProgress({ + rowsProcessed: 0, + totalRows: undefined, + totalSize: 0, + speed: 0, + mbPerSecond: 0, + elapsedTime: 0, + status: 'streaming', + }); + }, []); + + return { + progress, + isExporting, + startExport, + cancelExport, + resetExport, + }; +}; diff --git a/superset-frontend/src/dashboard/components/gridComponents/Chart/Chart.jsx b/superset-frontend/src/dashboard/components/gridComponents/Chart/Chart.jsx index d81a6ed0f253..21d612efa70c 100644 --- a/superset-frontend/src/dashboard/components/gridComponents/Chart/Chart.jsx +++ b/superset-frontend/src/dashboard/components/gridComponents/Chart/Chart.jsx @@ -27,6 +27,10 @@ import { useDispatch, useSelector } from 'react-redux'; import { exportChart, mountExploreUrl } from 'src/explore/exploreUtils'; import ChartContainer from 'src/components/Chart/ChartContainer'; +import { + StreamingExportModal, + useStreamingExport, +} from 'src/components/StreamingExportModal'; import { LOG_ACTIONS_CHANGE_DASHBOARD_FILTER, LOG_ACTIONS_EXPLORE_DASHBOARD_CHART, @@ -175,6 +179,20 @@ const Chart = props => { const [descriptionHeight, setDescriptionHeight] = useState(0); const [height, setHeight] = useState(props.height); const [width, setWidth] = useState(props.width); + + // Streaming export state + const [isStreamingModalVisible, setIsStreamingModalVisible] = useState(false); + const { progress, isExporting, startExport, cancelExport, resetExport } = + useStreamingExport({ + onComplete: (downloadUrl, filename) => { + boundActionCreators.addSuccessToast( + t('Export completed successfully: %s', filename), + ); + }, + onError: error => { + boundActionCreators.addDangerToast(t('Export failed: %s', error)); + }, + }); const history = useHistory(); const resize = useCallback( debounce(() => { @@ -378,12 +396,28 @@ const Chart = props => { slice_id: slice.slice_id, is_cached: isCached, }); + + const exportFormData = isFullCSV + ? { ...formData, row_limit: maxRows } + : formData; + const resultType = isPivot ? 'post_processed' : 'full'; + + // Handle streaming CSV exports for both regular and full data exports + const shouldUseStreaming = format === 'csv' && !isPivot; + exportChart({ - formData: isFullCSV ? { ...formData, row_limit: maxRows } : formData, - resultType: isPivot ? 'post_processed' : 'full', + formData: exportFormData, + resultType, resultFormat: format, force: true, ownState: dataMask[props.id]?.ownState, + onStartStreamingExport: shouldUseStreaming + ? exportParams => { + setIsStreamingModalVisible(true); + resetExport(); + startExport(exportParams); + } + : null, }); }, [ @@ -393,6 +427,8 @@ const Chart = props => { props.maxRows, dataMask[props.id]?.ownState, boundActionCreators.logEvent, + startExport, + resetExport, ], ); @@ -544,6 +580,25 @@ const Chart = props => { emitCrossFilters={emitCrossFilters} /> + + {/* Streaming Export Modal */} + { + if (isExporting) { + cancelExport(); + } + setIsStreamingModalVisible(false); + }} + onRetry={() => { + resetExport(); + // Note: Retry would need to store the last export parameters + // For now, just close the modal and let user retry manually + setIsStreamingModalVisible(false); + }} + progress={progress} + exportType="csv" + /> ); }; diff --git a/superset-frontend/src/explore/exploreUtils/index.js b/superset-frontend/src/explore/exploreUtils/index.js index d3bd62ef412a..2c3d7a392bab 100644 --- a/superset-frontend/src/explore/exploreUtils/index.js +++ b/superset-frontend/src/explore/exploreUtils/index.js @@ -248,6 +248,7 @@ export const exportChart = async ({ resultType = 'full', force = false, ownState = {}, + onStartStreamingExport = null, }) => { let url; let payload; @@ -272,7 +273,55 @@ export const exportChart = async ({ }); } - SupersetClient.postForm(url, { form_data: safeStringify(payload) }); + // Check if this should use streaming export for CSV + const shouldUseStreaming = + resultFormat === 'csv' && onStartStreamingExport && !useLegacyApi; + + if (shouldUseStreaming) { + // Use streaming export instead of opening new tab + const timestamp = new Date() + .toISOString() + .slice(0, 19) + .replace(/[-:]/g, '') + .replace('T', '_'); + const chartName = formData.slice_name || formData.viz_type || 'chart'; + const safeChartName = chartName.replace(/[^a-zA-Z0-9_-]/g, '_'); + const filename = `superset_${safeChartName}_${timestamp}.csv`; + + // Extract expected row count for progress calculation + // Try to get row limit from form data for accurate progress tracking + let expectedRows; + if (formData.row_limit && formData.row_limit > 0) { + expectedRows = formData.row_limit; + } else if ( + payload.queries && + payload.queries[0] && + payload.queries[0].row_limit + ) { + expectedRows = payload.queries[0].row_limit; + } else { + // Default fallback - estimate based on common chart sizes + expectedRows = 10000; // Conservative default for progress calculation + } + + console.log( + '๐ŸŽฏ EXPORT CHART: Setting expectedRows =', + expectedRows, + 'from formData.row_limit =', + formData.row_limit, + ); + + onStartStreamingExport({ + url, + payload, + filename, + exportType: 'csv', + expectedRows, // ๐ŸŽฏ This was missing! + }); + } else { + // Fallback to original behavior for non-streaming exports + SupersetClient.postForm(url, { form_data: safeStringify(payload) }); + } }; export const exploreChart = (formData, requestParams) => { From 2aaf39bd0d860df4384f4e28cb0782d44f825117 Mon Sep 17 00:00:00 2001 From: amaannawab923 Date: Fri, 3 Oct 2025 12:15:18 +0530 Subject: [PATCH 03/57] feat: Add testing mode for streaming CSV exports - Added configurable testing mode for UI demonstration - Testing mode: 10k row chunks with 3.5s delay between chunks - Production mode: 1k row chunks with no delay - Controlled by ENABLE_SLOW_STREAMING_TEST flag - Helps visualize progress modal during development --- superset/views/streaming.py | 402 ++++++++++++++++++++++++++++++++++++ 1 file changed, 402 insertions(+) create mode 100644 superset/views/streaming.py diff --git a/superset/views/streaming.py b/superset/views/streaming.py new file mode 100644 index 000000000000..67842d7c8213 --- /dev/null +++ b/superset/views/streaming.py @@ -0,0 +1,402 @@ +# 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. + +"""Streaming HTTP responses for large dataset exports.""" + +from __future__ import annotations + +import logging +import time +from datetime import datetime +from typing import Any, Generator, TYPE_CHECKING + +from flask import current_app as app, Response +from werkzeug.datastructures import Headers + +if TYPE_CHECKING: + from superset.common.query_context import QueryContext + +logger = logging.getLogger(__name__) + + +def create_streaming_csv_response_simple( + data_generator: Generator[str, None, None], + filename: str = "export.csv", + encoding: str = "utf-8", +) -> Response: + """ + Create a simple streaming CSV response using Flask's standard pattern. + + This follows the official Flask streaming documentation pattern. + """ + from flask import Response + + # Create response with proper headers + response = Response( + data_generator, + mimetype=f"text/csv; charset={encoding}", + headers={ + "Content-Disposition": f'attachment; filename="{filename}"', + "Cache-Control": "no-cache", + }, + ) + + logger.info("Created simple streaming CSV response for file: %s", filename) + return response + + +class StreamingExcelResponse(Response): + """ + Streaming response for Excel (XLSX) files. + + Note: Excel streaming is more complex than CSV due to the binary format. + This is a placeholder for future implementation. + """ + + def __init__( + self, + data_generator: Generator[bytes, None, None], + filename: str = "export.xlsx", + **kwargs: Any, + ) -> None: + headers = Headers() + headers.add( + "Content-Type", + "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet", + ) + headers.add("Content-Disposition", f'attachment; filename="{filename}"') + headers.add("Cache-Control", "no-cache, no-store, must-revalidate") + headers.add("Transfer-Encoding", "chunked") + + super().__init__( + response=data_generator, headers=headers, direct_passthrough=True, **kwargs + ) + + # Note: is_streamed is automatically True when response is a generator + + +def create_streaming_csv_response( + query_context: QueryContext, + filename: str | None = None, + chunk_size: int | None = None, + escape_formulas: bool = True, +) -> Response: + """ + Factory function to create a streaming CSV response using Flask's standard pattern. + + Args: + query_context: Superset query context + filename: Optional filename for download + chunk_size: Optional chunk size for processing + escape_formulas: Whether to escape formula injection + + Returns: + Flask Response configured for streaming CSV + """ + # Capture the Flask app instance in the current context + current_app = app._get_current_object() + + def csv_generator() -> Generator[str, None, None]: + """Generator that yields CSV data from database query.""" + # Use the captured app instance to create application context + with current_app.app_context(): + # Performance tracking + start_time = time.time() + total_bytes = 0 + + try: + logger.info("๐Ÿš€ STREAMING CSV: Starting streaming CSV generation") + logger.info( + "๐Ÿ“Š STREAMING CSV: Processing query with estimated large result set" + ) + + # Get the database connection and execute raw SQL query directly + from superset import db + from superset.connectors.sqla.models import SqlaTable + + # Get the datasource + datasource = query_context.datasource + + # Create a fresh session to avoid detached object issues + with db.session() as session: + # Refresh the datasource in the current session + if isinstance(datasource, SqlaTable): + datasource = session.merge(datasource) + + # Generate SQL query from query context + query_obj = query_context.queries[0] + + # Use the datasource to get the SQL query + sql_query = datasource.get_query_str(query_obj.to_dict()) + + query_start_time = time.time() + logger.info( + "โšก STREAMING CSV: Executing SQL query: %s...", + sql_query[:200], + ) + logger.info( + "โฑ๏ธ STREAMING CSV: Query execution started at %s", + time.strftime("%H:%M:%S"), + ) + + # Execute query directly with the database engine + # Use context manager for proper connection handling + with datasource.database.get_sqla_engine() as engine: + # Use server-side cursor for streaming + connection = engine.connect() + + try: + # Execute query with server-side cursor + from sqlalchemy import text + + result_proxy = connection.execution_options( + stream_results=True + ).execute(text(sql_query)) + + # Get column names + columns = list(result_proxy.keys()) + query_execution_time = time.time() - query_start_time + logger.info( + "๐Ÿ“‹ STREAMING CSV: Query columns (%d): %s", + len(columns), + columns, + ) + logger.info( + "โšก STREAMING CSV: Query executed in %.2fs, " + "starting data streaming...", + query_execution_time, + ) + + # Yield CSV header + header_row = ",".join(f'"{col}"' for col in columns) + "\n" + header_bytes = len(header_row.encode("utf-8")) + total_bytes += header_bytes + yield header_row + + # ๐Ÿงช TESTING CONFIGURATION - Enable slower streaming for UI testing + ENABLE_SLOW_STREAMING_TEST = ( + True # Set to False for production speed + ) + + if ENABLE_SLOW_STREAMING_TEST: + # Testing mode: 10k rows per chunk with 0.5s delay + chunk_size = 10000 + delay_between_chunks = 3.5 + logger.info( + "๐Ÿงช TESTING MODE: Using 10k row chunks with 0.5s delays" + ) + else: + # Production mode: 1k rows per chunk, no delay + chunk_size = 1000 + delay_between_chunks = 0 + + row_count = 0 + streaming_start_time = time.time() + last_progress_time = streaming_start_time + + while True: + # Fetch chunk of rows + rows = result_proxy.fetchmany(chunk_size) + if not rows: + break + + # Yield CSV rows + for row in rows: + csv_row = ",".join( + f'"{str(cell) if cell is not None else ""}"' + for cell in row + ) + csv_line = csv_row + "\n" + row_bytes = len(csv_line.encode("utf-8")) + total_bytes += row_bytes + yield csv_line + row_count += 1 + + # Performance logging every 10k rows or 5 seconds + current_time = time.time() + if ( + row_count % 10000 == 0 + or (current_time - last_progress_time) >= 5 + ): + elapsed = current_time - streaming_start_time + rows_per_sec = ( + row_count / elapsed if elapsed > 0 else 0 + ) + mb_streamed = total_bytes / (1024 * 1024) + mb_per_sec = ( + mb_streamed / elapsed if elapsed > 0 else 0 + ) + + logger.info( + "๐Ÿ“ˆ STREAMING CSV: %s rows streamed in %.1fs " + "(%.0f rows/s, %.1fMB, %.1fMB/s)", + f"{row_count:,}", + elapsed, + rows_per_sec, + mb_streamed, + mb_per_sec, + ) + last_progress_time = current_time + + # Apply testing delay for UI demonstration + if ( + ENABLE_SLOW_STREAMING_TEST + and delay_between_chunks > 0 + ): + time.sleep(delay_between_chunks) + + # Final performance summary + total_time = time.time() - start_time + streaming_time = time.time() - streaming_start_time + total_mb = total_bytes / (1024 * 1024) + + logger.info( + "โœ… STREAMING CSV: Completed streaming %s rows", + f"{row_count:,}", + ) + logger.info("๐Ÿ“Š STREAMING CSV PERFORMANCE:") + logger.info(" โ€ข Total Time: %.2fs", total_time) + logger.info(" โ€ข Query Time: %.2fs", query_execution_time) + logger.info(" โ€ข Streaming Time: %.2fs", streaming_time) + logger.info(" โ€ข Data Size: %.1fMB", total_mb) + logger.info( + " โ€ข Average Speed: %.0f rows/s, %.1fMB/s", + row_count / total_time, + total_mb / total_time, + ) + logger.info( + " โ€ข Memory Efficient: โœ… Constant memory usage " + "(~%d rows buffered)", + chunk_size, + ) + + finally: + connection.close() + + except Exception as e: + logger.error("Error in streaming CSV generator: %s", e) + import traceback + + logger.error("Traceback: %s", traceback.format_exc()) + + # Yield error info and fallback data + yield f"# Error occurred: {str(e)}\n" + yield "error,message\n" + yield f"CSV Export Error,{str(e)}\n" + + # Generate filename + if filename is None: + timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") + filename = f"superset_streaming_{timestamp}.csv" + + # Get encoding from the captured app + encoding = current_app.config.get("CSV_EXPORT", {}).get("encoding", "utf-8") + + logger.info("Creating Flask streaming CSV response: %s", filename) + + # Use simple Flask Response with generator (official pattern) + return create_streaming_csv_response_simple( + data_generator=csv_generator(), + filename=filename, + encoding=encoding, + ) + + +def should_use_streaming_response( + query_context: QueryContext, + result_format: str, +) -> bool: + """ + Determine if a streaming response should be used. + + Args: + query_context: The query context + result_format: The requested result format (csv, xlsx, etc.) + + Returns: + True if streaming should be used + """ + # Only stream for supported formats + if result_format.lower() not in ["csv"]: # Add "xlsx" when implemented + return False + + # Check if streaming is enabled in config + config = app.config.get("CSV_STREAMING", {}) + if not config.get("enabled", True): + return False + + # Simple row count estimation - always trigger streaming for CSV exports + # TODO: Implement proper row estimation logic based on query + estimated_rows = 50000 # Conservative estimate to trigger streaming + threshold = app.config.get("CSV_STREAMING_THRESHOLD", 10000) + + use_streaming = estimated_rows > threshold + + logger.info( + "Streaming decision: %d estimated rows, threshold %d, use_streaming=%s", + estimated_rows, + threshold, + use_streaming, + ) + + return use_streaming + + +class StreamingProgressTracker: + """ + Track progress of streaming exports for monitoring and user feedback. + + This could be extended to provide real-time progress updates via WebSocket + or SSE (Server-Sent Events) in the future. + """ + + def __init__(self, export_id: str) -> None: + self.export_id = export_id + self.start_time = datetime.now() + self.last_update = self.start_time + self.total_chunks = 0 + self.processed_rows = 0 + + def update_progress(self, chunks: int, rows: int) -> None: + """Update progress metrics.""" + self.total_chunks = chunks + self.processed_rows = rows + self.last_update = datetime.now() + + # Log progress every 100 chunks + if chunks % 100 == 0: + elapsed = (self.last_update - self.start_time).total_seconds() + rate = rows / elapsed if elapsed > 0 else 0 + logger.info( + "Export %s: %d rows, %d chunks, %.0f rows/sec", + self.export_id, + rows, + chunks, + rate, + ) + + def get_status(self) -> dict[str, Any]: + """Get current export status.""" + elapsed = (datetime.now() - self.start_time).total_seconds() + return { + "export_id": self.export_id, + "start_time": self.start_time.isoformat(), + "elapsed_seconds": elapsed, + "processed_rows": self.processed_rows, + "total_chunks": self.total_chunks, + "last_update": self.last_update.isoformat(), + } From d9d96a453420db51d48ab940f8b1d663bc136cd4 Mon Sep 17 00:00:00 2001 From: amaannawab923 Date: Fri, 3 Oct 2025 13:50:32 +0530 Subject: [PATCH 04/57] Streaming progress --- .../useStreamingExport.ts | 174 +++++---------- superset/charts/data/api.py | 165 ++++++++++++++- superset/views/streaming.py | 47 ++++- superset/views/streaming_progress.py | 199 ++++++++++++++++++ 4 files changed, 450 insertions(+), 135 deletions(-) create mode 100644 superset/views/streaming_progress.py diff --git a/superset-frontend/src/components/StreamingExportModal/useStreamingExport.ts b/superset-frontend/src/components/StreamingExportModal/useStreamingExport.ts index b021f936602f..23a2a79c5c37 100644 --- a/superset-frontend/src/components/StreamingExportModal/useStreamingExport.ts +++ b/superset-frontend/src/components/StreamingExportModal/useStreamingExport.ts @@ -202,6 +202,48 @@ export const useStreamingExport = (options: UseStreamingExportOptions = {}) => { filename, }); + // Start polling for progress using filename as export ID + const pollInterval = setInterval(async () => { + try { + const progressUrl = `/api/v1/chart/export/progress/${encodeURIComponent( + filename || `export.${exportType}`, + )}`; + console.log('๐Ÿ“Š POLLING PROGRESS:', progressUrl); + + const progressResponse = await SupersetClient.get({ + endpoint: progressUrl, + }); + + if (progressResponse.json) { + const serverProgress = progressResponse.json; + console.log('๐Ÿ“ˆ SERVER PROGRESS:', serverProgress); + + // Update progress from server response + updateProgress({ + rowsProcessed: serverProgress.rows_processed || 0, + totalRows: serverProgress.total_rows || expectedRows, + totalSize: serverProgress.bytes_processed || 0, + speed: serverProgress.speed_rows_per_sec || 0, + mbPerSecond: serverProgress.speed_mb_per_sec || 0, + elapsedTime: serverProgress.elapsed_time || 0, + status: serverProgress.status || 'streaming', + }); + + // Stop polling if completed or error + if ( + serverProgress.status === 'completed' || + serverProgress.status === 'error' + ) { + console.log('โœ… POLLING STOPPED: Export finished'); + clearInterval(pollInterval); + } + } + } catch (pollError) { + // Silently ignore polling errors - the main stream will handle failures + console.warn('โš ๏ธ Progress polling error (non-critical):', pollError); + } + }, 500); // Poll every 500ms as requested + try { // Initialize SupersetClient to ensure authentication and get CSRF token await SupersetClient.init(); @@ -216,6 +258,7 @@ export const useStreamingExport = (options: UseStreamingExportOptions = {}) => { }, body: new URLSearchParams({ form_data: JSON.stringify(payload), + filename: filename || `export.${exportType}`, // Pass filename to backend }), signal: abortControllerRef.current.signal, credentials: 'same-origin', @@ -242,32 +285,10 @@ export const useStreamingExport = (options: UseStreamingExportOptions = {}) => { const reader = response.body.getReader(); const chunks: Uint8Array[] = []; let receivedLength = 0; - let actualRowCount = 0; - let accumulatedText = ''; let chunkCount = 0; - // More frequent progress updates for smoother UI - const progressInterval = setInterval(() => { - const now = Date.now(); - const elapsedTime = (now - startTimeRef.current) / 1000; - - // Always use the enhanced progress calculation with actual rows when available - const progressData = estimateProgressFromSize( - receivedLength, - elapsedTime, - actualRowCount, - ); - - // ๐Ÿ” DEBUG: Log interval-based progress updates - console.log('โฑ๏ธ FRONTEND INTERVAL UPDATE:', { - receivedLength, - actualRowCount, - elapsedTime: elapsedTime.toFixed(1), - progressData, - }); - - updateProgress(progressData); - }, 200); // Update every 200ms for smoother progress + // Note: We rely on the polling interval (pollInterval) for progress updates + // No need for a separate progressInterval since the server tracks actual progress try { // Read the streaming response @@ -293,8 +314,6 @@ export const useStreamingExport = (options: UseStreamingExportOptions = {}) => { chunkCount++; chunks.push(value); - const oldReceivedLength = receivedLength; - const oldRowCount = actualRowCount; receivedLength += value.length; console.log('๐Ÿ“Š CHUNK RECEIVED:', { @@ -302,66 +321,9 @@ export const useStreamingExport = (options: UseStreamingExportOptions = {}) => { chunkSize: value.length, totalReceived: receivedLength, }); - - // Convert chunk to text and count rows in real-time - try { - const chunkText = new TextDecoder().decode(value, { - stream: true, - }); - accumulatedText += chunkText; - - // Count complete lines (rows) in accumulated text - const lines = accumulatedText.split('\n'); - - // Keep the last incomplete line for next iteration - if (lines.length > 1) { - accumulatedText = lines[lines.length - 1]; // Keep incomplete line - - // Count completed lines (subtract 1 for header if this is first chunk) - const completedLines = lines.length - 1; - if (actualRowCount === 0 && completedLines > 0) { - // First chunk - subtract 1 for header row - actualRowCount = completedLines - 1; - } else { - // Subsequent chunks - add all completed lines - actualRowCount += completedLines; - } - } - - // ๐Ÿ” DEBUG: Log chunk processing and update progress immediately - if ( - actualRowCount !== oldRowCount || - receivedLength !== oldReceivedLength - ) { - console.log('๐Ÿ“ฆ FRONTEND CHUNK PROCESSED:', { - chunkSize: value.length, - totalBytes: receivedLength, - newRows: actualRowCount - oldRowCount, - totalRows: actualRowCount, - linesInChunk: lines.length - 1, - chunkTextPreview: `${chunkText.substring(0, 100)}...`, - }); - - // Immediately update progress with actual row count from chunk - const now = Date.now(); - const elapsedTime = (now - startTimeRef.current) / 1000; - const progressData = estimateProgressFromSize( - receivedLength, - elapsedTime, - actualRowCount, - ); - updateProgress(progressData); - } - } catch (decodeError) { - // If text decoding fails, fall back to byte-based estimation - console.warn( - 'โŒ Failed to decode chunk for row counting:', - decodeError, - ); - } } - clearInterval(progressInterval); + clearInterval(pollInterval); // Stop polling when stream completes // Create blob from chunks const completeData = new Uint8Array(receivedLength); @@ -379,48 +341,8 @@ export const useStreamingExport = (options: UseStreamingExportOptions = {}) => { const blob = new Blob([completeData], { type: mimeType }); const downloadUrl = URL.createObjectURL(blob); - const finalElapsedTime = (Date.now() - startTimeRef.current) / 1000; - - // Try to get accurate final count from server headers first - const headerProgress = parseProgressFromHeaders( - response.headers, - finalElapsedTime, - ); - - let finalProgress; - if (headerProgress && headerProgress.rowsProcessed > 0) { - // Use server-provided count if available - finalProgress = headerProgress; - } else { - // Use the real-time counted rows if available, otherwise parse final data - let finalRowCount = actualRowCount; - - if (finalRowCount === 0) { - // Fallback: parse the complete data for final count - try { - const text = new TextDecoder().decode(completeData); - const lines = text.split(/\r?\n/); - // Subtract 1 for header row, filter out empty lines - finalRowCount = - lines.filter(line => line.trim().length > 0).length - 1; - finalRowCount = Math.max(0, finalRowCount); - } catch (e) { - // Final fallback to size-based estimation - finalRowCount = Math.floor(receivedLength / 200); - } - } - - finalProgress = { - rowsProcessed: finalRowCount, - totalSize: receivedLength, - speed: finalRowCount / finalElapsedTime, - mbPerSecond: receivedLength / (1024 * 1024) / finalElapsedTime, - elapsedTime: finalElapsedTime, - }; - } - + // Mark as completed - final progress will come from the last poll updateProgress({ - ...finalProgress, status: 'completed', downloadUrl, filename: filename || `export.${exportType}`, @@ -428,7 +350,7 @@ export const useStreamingExport = (options: UseStreamingExportOptions = {}) => { options.onComplete?.(downloadUrl, filename || `export.${exportType}`); } catch (streamError) { - clearInterval(progressInterval); + clearInterval(pollInterval); // Stop polling on error throw streamError; } } catch (error) { diff --git a/superset/charts/data/api.py b/superset/charts/data/api.py index 3591539a098a..0a89141ed1bd 100644 --- a/superset/charts/data/api.py +++ b/superset/charts/data/api.py @@ -62,7 +62,7 @@ class ChartDataRestApi(ChartRestApi): - include_route_methods = {"get_data", "data", "data_from_cache"} + include_route_methods = {"get_data", "data", "data_from_cache", "export_progress"} @expose("//data/", methods=("GET",)) @protect() @@ -257,8 +257,14 @@ def data(self) -> Response: return self._run_async(json_body, command) form_data = json_body.get("form_data") + + # Extract filename from request if provided (for streaming CSV) + filename = request.form.get("filename") + if filename: + logger.info("๐Ÿ“ FRONTEND PROVIDED FILENAME: %s", filename) + return self._get_data_response( - command, form_data=form_data, datasource=query_context.datasource + command, form_data=form_data, datasource=query_context.datasource, filename=filename ) @expose("/data/", methods=("GET",)) @@ -319,6 +325,109 @@ def data_from_cache(self, cache_key: str) -> Response: return self._get_data_response(command, True) + @expose("/export/progress/", methods=("GET",)) + @statsd_metrics + @event_logger.log_this_with_context( + action=lambda self, *args, **kwargs: f"{self.__class__.__name__}" + f".export_progress", + log_to_statsd=False, + ) + def export_progress(self, export_id: str) -> Response: + """ + Poll progress for a streaming CSV export. + --- + get: + summary: Get streaming export progress + description: >- + Returns the current progress of a streaming CSV export. + The export_id is the filename being exported. + This endpoint is polled by the frontend every 500ms. + parameters: + - in: path + schema: + type: string + name: export_id + description: The export ID (filename) + responses: + 200: + description: Export progress + content: + application/json: + schema: + type: object + properties: + export_id: + type: string + status: + type: string + enum: [streaming, completed, error] + rows_processed: + type: integer + total_rows: + type: integer + bytes_processed: + type: integer + elapsed_time: + type: number + percentage: + type: number + speed_rows_per_sec: + type: number + speed_mb_per_sec: + type: number + error_message: + type: string + 404: + $ref: '#/components/responses/404' + 401: + $ref: '#/components/responses/401' + """ + import logging + + logger = logging.getLogger(__name__) + + logger.info("๐Ÿ” POLLING API CALLED: export_id=%s", export_id) + + from superset.views.streaming_progress import progress_tracker + + logger.info("๐Ÿ“Š POLLING API: Getting progress for export_id=%s", export_id) + + # Get progress from tracker + progress = progress_tracker.get_progress(export_id) + + logger.info("๐Ÿ“ˆ POLLING API: Progress result=%s", progress) + + if progress is None: + logger.warning( + "โš ๏ธ POLLING API: Export not found: export_id=%s", export_id + ) + # Return a response indicating export not started yet or already completed + return self.response( + 200, + export_id=export_id, + status="not_found", + rows_processed=0, + total_rows=None, + bytes_processed=0, + elapsed_time=0, + percentage=None, + speed_rows_per_sec=0, + speed_mb_per_sec=0, + error_message="Export not found or not started yet", + ) + + logger.info( + "โœ… POLLING API: Returning progress: status=%s, rows=%s, bytes=%s", + progress.get("status"), + progress.get("rows_processed"), + progress.get("bytes_processed"), + ) + + # Cleanup old exports periodically + progress_tracker.cleanup_old_exports() + + return self.response(200, **progress) + def _run_async( self, form_data: dict[str, Any], command: ChartDataCommand ) -> Response: @@ -348,6 +457,7 @@ def _send_chart_response( # noqa: C901 result: dict[Any, Any], form_data: dict[str, Any] | None = None, datasource: BaseDatasource | Query | None = None, + filename: str | None = None, ) -> Response: result_type = result["query_context"].result_type result_format = result["query_context"].result_format @@ -368,6 +478,10 @@ def _send_chart_response( # noqa: C901 is_csv_format = result_format == ChartDataResultFormat.CSV + # Check if we should use streaming for large datasets + if is_csv_format and self._should_use_streaming(result, form_data): + return self._create_streaming_csv_response(result, form_data, filename=filename) + if len(result["queries"]) == 1: # return single query results data = result["queries"][0]["data"] @@ -417,6 +531,7 @@ def _get_data_response( force_cached: bool = False, form_data: dict[str, Any] | None = None, datasource: BaseDatasource | Query | None = None, + filename: str | None = None, ) -> Response: try: result = command.run(force_cached=force_cached) @@ -425,7 +540,7 @@ 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) # pylint: disable=invalid-name def _load_query_context_form_from_cache(self, cache_key: str) -> dict[str, Any]: @@ -462,3 +577,47 @@ 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 for this response.""" + from superset.views.streaming import should_use_streaming_response + + query_context = result["query_context"] + result_format = query_context.result_format + + return should_use_streaming_response(query_context, result_format) + + def _create_streaming_csv_response( + self, result: dict[Any, Any], form_data: dict[str, Any] | None = None, filename: str | None = None + ) -> Response: + """Create a streaming CSV response for large datasets.""" + from datetime import datetime + + from superset.views.streaming import create_streaming_csv_response + + query_context = result["query_context"] + + # Use filename from frontend if provided, otherwise generate one + if not filename: + timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") + chart_name = "export" + + if form_data and form_data.get("slice_name"): + chart_name = form_data["slice_name"] + elif form_data and form_data.get("viz_type"): + chart_name = form_data["viz_type"] + + # Sanitize chart name for filename + safe_chart_name = "".join( + c for c in chart_name if c.isalnum() or c in ("-", "_") + ) + filename = f"superset_{safe_chart_name}_{timestamp}.csv" + + logger.info("Creating streaming CSV response: %s (from frontend: %s)", filename, filename is not None) + + return create_streaming_csv_response( + query_context=query_context, + filename=filename, + ) diff --git a/superset/views/streaming.py b/superset/views/streaming.py index 67842d7c8213..fb024b36d369 100644 --- a/superset/views/streaming.py +++ b/superset/views/streaming.py @@ -21,12 +21,15 @@ import logging import time +import uuid from datetime import datetime from typing import Any, Generator, TYPE_CHECKING from flask import current_app as app, Response from werkzeug.datastructures import Headers +from superset.views.streaming_progress import progress_tracker + if TYPE_CHECKING: from superset.common.query_context import QueryContext @@ -107,6 +110,14 @@ def create_streaming_csv_response( Returns: Flask Response configured for streaming CSV """ + # Generate filename first if not provided + if filename is None: + timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") + filename = f"superset_streaming_{timestamp}.csv" + + # Use filename as export ID for progress tracking (frontend knows this) + export_id = filename + # Capture the Flask app instance in the current context current_app = app._get_current_object() @@ -120,10 +131,26 @@ def csv_generator() -> Generator[str, None, None]: try: logger.info("๐Ÿš€ STREAMING CSV: Starting streaming CSV generation") + logger.info("๐Ÿ“Š STREAMING CSV: Export ID: %s", export_id) logger.info( "๐Ÿ“Š STREAMING CSV: Processing query with estimated large result set" ) + # Estimate total rows for progress tracking + # TODO: Get actual row count from frontend instead of hardcoding + # For now, hardcoded to 148795 for testing + estimated_rows = 148795 + + logger.info( + "๐Ÿ“Š STREAMING CSV: Using hardcoded total_rows=%d for progress tracking", + estimated_rows, + ) + + # Initialize progress tracker + progress_tracker.create_export( + export_id=export_id, total_rows=estimated_rows + ) + # Get the database connection and execute raw SQL query directly from superset import db from superset.connectors.sqla.models import SqlaTable @@ -195,7 +222,7 @@ def csv_generator() -> Generator[str, None, None]: if ENABLE_SLOW_STREAMING_TEST: # Testing mode: 10k rows per chunk with 0.5s delay chunk_size = 10000 - delay_between_chunks = 3.5 + delay_between_chunks = 0.5 logger.info( "๐Ÿงช TESTING MODE: Using 10k row chunks with 0.5s delays" ) @@ -226,6 +253,13 @@ def csv_generator() -> Generator[str, None, None]: yield csv_line row_count += 1 + # Update progress tracker after each chunk + progress_tracker.update_progress( + export_id=export_id, + rows_processed=row_count, + bytes_processed=total_bytes, + ) + # Performance logging every 10k rows or 5 seconds current_time = time.time() if ( @@ -284,6 +318,9 @@ def csv_generator() -> Generator[str, None, None]: chunk_size, ) + # Mark export as completed in progress tracker + progress_tracker.complete_export(export_id) + finally: connection.close() @@ -293,16 +330,14 @@ def csv_generator() -> Generator[str, None, None]: logger.error("Traceback: %s", traceback.format_exc()) + # Mark export as failed in progress tracker + progress_tracker.fail_export(export_id, str(e)) + # Yield error info and fallback data yield f"# Error occurred: {str(e)}\n" yield "error,message\n" yield f"CSV Export Error,{str(e)}\n" - # Generate filename - if filename is None: - timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") - filename = f"superset_streaming_{timestamp}.csv" - # Get encoding from the captured app encoding = current_app.config.get("CSV_EXPORT", {}).get("encoding", "utf-8") diff --git a/superset/views/streaming_progress.py b/superset/views/streaming_progress.py new file mode 100644 index 000000000000..3bfdada6b6d5 --- /dev/null +++ b/superset/views/streaming_progress.py @@ -0,0 +1,199 @@ +# 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. + +"""Thread-safe progress tracking for streaming CSV exports.""" + +from __future__ import annotations + +import logging +import threading +import time +from dataclasses import dataclass +from typing import Any + +logger = logging.getLogger(__name__) + + +@dataclass +class ExportProgress: + """Progress information for a streaming export.""" + + export_id: str + status: str = "streaming" # streaming, completed, error + rows_processed: int = 0 + total_rows: int | None = None + bytes_processed: int = 0 + elapsed_time: float = 0.0 + start_time: float = 0.0 + error_message: str | None = None + + def to_dict(self) -> dict[str, Any]: + """Convert to dictionary for JSON response.""" + percentage = None + if self.total_rows and self.total_rows > 0: + percentage = round((self.rows_processed / self.total_rows) * 100, 1) + + # Calculate speeds + speed_rows_per_sec = 0.0 + speed_mb_per_sec = 0.0 + if self.elapsed_time > 0: + speed_rows_per_sec = self.rows_processed / self.elapsed_time + speed_mb_per_sec = (self.bytes_processed / self.elapsed_time) / ( + 1024 * 1024 + ) + + return { + "export_id": self.export_id, + "status": self.status, + "rows_processed": self.rows_processed, + "total_rows": self.total_rows, + "bytes_processed": self.bytes_processed, + "elapsed_time": round(self.elapsed_time, 2), + "percentage": percentage, + "speed_rows_per_sec": round(speed_rows_per_sec, 2), + "speed_mb_per_sec": round(speed_mb_per_sec, 3), + "error_message": self.error_message, + } + + +class StreamingProgressTracker: + """ + Thread-safe singleton for tracking streaming export progress. + + This allows the streaming generator to update progress while a separate + polling endpoint reads the current state without blocking the stream. + """ + + _instance: StreamingProgressTracker | None = None + _lock = threading.Lock() + + def __new__(cls) -> StreamingProgressTracker: + """Singleton pattern to ensure single instance.""" + if cls._instance is None: + with cls._lock: + if cls._instance is None: + cls._instance = super().__new__(cls) + cls._instance._initialized = False + return cls._instance + + def __init__(self) -> None: + """Initialize the progress tracker.""" + if self._initialized: + return + + self._progress_map: dict[str, ExportProgress] = {} + self._progress_lock = threading.Lock() + self._cleanup_interval = 300 # Cleanup after 5 minutes + self._last_cleanup = time.time() + self._initialized = True + + logger.info("โœ… StreamingProgressTracker initialized") + + def create_export( + self, export_id: str, total_rows: int | None = None + ) -> None: + """Create a new export progress entry.""" + with self._progress_lock: + self._progress_map[export_id] = ExportProgress( + export_id=export_id, + total_rows=total_rows, + start_time=time.time(), + ) + logger.info( + "๐Ÿ“Š Created progress tracker for export: %s (expected %s rows)", + export_id, + total_rows or "unknown", + ) + + def update_progress( + self, + export_id: str, + rows_processed: int, + bytes_processed: int, + ) -> None: + """Update progress for an export.""" + with self._progress_lock: + if export_id not in self._progress_map: + logger.warning("โš ๏ธ Export ID not found: %s", export_id) + return + + progress = self._progress_map[export_id] + progress.rows_processed = rows_processed + progress.bytes_processed = bytes_processed + progress.elapsed_time = time.time() - progress.start_time + + def complete_export(self, export_id: str) -> None: + """Mark export as completed.""" + with self._progress_lock: + if export_id in self._progress_map: + self._progress_map[export_id].status = "completed" + self._progress_map[export_id].elapsed_time = ( + time.time() - self._progress_map[export_id].start_time + ) + logger.info("โœ… Export completed: %s", export_id) + + def fail_export(self, export_id: str, error_message: str) -> None: + """Mark export as failed.""" + with self._progress_lock: + if export_id in self._progress_map: + self._progress_map[export_id].status = "error" + self._progress_map[export_id].error_message = error_message + self._progress_map[export_id].elapsed_time = ( + time.time() - self._progress_map[export_id].start_time + ) + logger.error("โŒ Export failed: %s - %s", export_id, error_message) + + def get_progress(self, export_id: str) -> dict[str, Any] | None: + """Get current progress for an export.""" + with self._progress_lock: + progress = self._progress_map.get(export_id) + if progress: + # Update elapsed time before returning + progress.elapsed_time = time.time() - progress.start_time + return progress.to_dict() + return None + + def cleanup_old_exports(self, max_age_seconds: int = 300) -> None: + """Remove old completed/failed exports to prevent memory leaks.""" + current_time = time.time() + + # Only cleanup periodically + if current_time - self._last_cleanup < self._cleanup_interval: + return + + with self._progress_lock: + to_remove = [] + for export_id, progress in self._progress_map.items(): + age = current_time - progress.start_time + if age > max_age_seconds and progress.status in ["completed", "error"]: + to_remove.append(export_id) + + for export_id in to_remove: + del self._progress_map[export_id] + + self._last_cleanup = current_time + + if to_remove: + logger.info( + "๐Ÿงน Cleaned up %d old exports, %d active remaining", + len(to_remove), + len(self._progress_map), + ) + + +# Global singleton instance +progress_tracker = StreamingProgressTracker() From b78895a30ddaa1c0d9cd982ecbe5939f4c5adf88 Mon Sep 17 00:00:00 2001 From: amaannawab923 Date: Fri, 3 Oct 2025 13:51:31 +0530 Subject: [PATCH 05/57] streaming --- superset/views/streaming.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/superset/views/streaming.py b/superset/views/streaming.py index fb024b36d369..c590cefd4ef6 100644 --- a/superset/views/streaming.py +++ b/superset/views/streaming.py @@ -164,10 +164,8 @@ def csv_generator() -> Generator[str, None, None]: if isinstance(datasource, SqlaTable): datasource = session.merge(datasource) - # Generate SQL query from query context query_obj = query_context.queries[0] - # Use the datasource to get the SQL query sql_query = datasource.get_query_str(query_obj.to_dict()) query_start_time = time.time() From 3d008616f0f172bc185e4ac01971f3916bc33209 Mon Sep 17 00:00:00 2001 From: amaannawab923 Date: Fri, 3 Oct 2025 20:45:02 +0530 Subject: [PATCH 06/57] use row limit for table from queries response --- .../useStreamingExport.ts | 1 + .../components/gridComponents/Chart/Chart.jsx | 19 +++++++++++- superset/charts/data/api.py | 24 +++++++++++---- superset/views/streaming.py | 30 ++++++++++++++----- 4 files changed, 60 insertions(+), 14 deletions(-) diff --git a/superset-frontend/src/components/StreamingExportModal/useStreamingExport.ts b/superset-frontend/src/components/StreamingExportModal/useStreamingExport.ts index 23a2a79c5c37..14bb778d6c73 100644 --- a/superset-frontend/src/components/StreamingExportModal/useStreamingExport.ts +++ b/superset-frontend/src/components/StreamingExportModal/useStreamingExport.ts @@ -259,6 +259,7 @@ export const useStreamingExport = (options: UseStreamingExportOptions = {}) => { body: new URLSearchParams({ form_data: JSON.stringify(payload), filename: filename || `export.${exportType}`, // Pass filename to backend + expected_rows: expectedRows?.toString() || '', // Pass expected row count to backend }), signal: abortControllerRef.current.signal, credentials: 'same-origin', diff --git a/superset-frontend/src/dashboard/components/gridComponents/Chart/Chart.jsx b/superset-frontend/src/dashboard/components/gridComponents/Chart/Chart.jsx index 21d612efa70c..f649cb8c00d5 100644 --- a/superset-frontend/src/dashboard/components/gridComponents/Chart/Chart.jsx +++ b/superset-frontend/src/dashboard/components/gridComponents/Chart/Chart.jsx @@ -402,6 +402,19 @@ const Chart = props => { : formData; const resultType = isPivot ? 'post_processed' : 'full'; + let actualRowCount; + const isTableViz = formData?.viz_type === 'table'; + + if ( + isTableViz && + queriesResponse?.length > 1 && + queriesResponse[1]?.data?.[0]?.rowcount + ) { + actualRowCount = queriesResponse[1].data[0].rowcount; + } else { + actualRowCount = exportFormData?.row_limit; + } + // Handle streaming CSV exports for both regular and full data exports const shouldUseStreaming = format === 'csv' && !isPivot; @@ -415,7 +428,10 @@ const Chart = props => { ? exportParams => { setIsStreamingModalVisible(true); resetExport(); - startExport(exportParams); + startExport({ + ...exportParams, + expectedRows: actualRowCount || exportParams.expectedRows, + }); } : null, }); @@ -427,6 +443,7 @@ const Chart = props => { props.maxRows, dataMask[props.id]?.ownState, boundActionCreators.logEvent, + queriesResponse, startExport, resetExport, ], diff --git a/superset/charts/data/api.py b/superset/charts/data/api.py index 0a89141ed1bd..7848de674e7f 100644 --- a/superset/charts/data/api.py +++ b/superset/charts/data/api.py @@ -263,8 +263,17 @@ def data(self) -> Response: if filename: logger.info("๐Ÿ“ FRONTEND PROVIDED FILENAME: %s", filename) + expected_rows = request.form.get("expected_rows") + if expected_rows: + try: + expected_rows = int(expected_rows) + logger.info("๐Ÿ“Š FRONTEND PROVIDED EXPECTED ROWS: %d", expected_rows) + except (ValueError, TypeError): + logger.warning("โš ๏ธ Invalid expected_rows value: %s", expected_rows) + expected_rows = None + return self._get_data_response( - command, form_data=form_data, datasource=query_context.datasource, filename=filename + command, form_data=form_data, datasource=query_context.datasource, filename=filename, expected_rows=expected_rows ) @expose("/data/", methods=("GET",)) @@ -458,6 +467,7 @@ def _send_chart_response( # noqa: C901 form_data: dict[str, Any] | None = None, datasource: BaseDatasource | Query | None = None, filename: str | None = None, + expected_rows: int | None = None, ) -> Response: result_type = result["query_context"].result_type result_format = result["query_context"].result_format @@ -479,8 +489,8 @@ def _send_chart_response( # noqa: C901 is_csv_format = result_format == ChartDataResultFormat.CSV # Check if we should use streaming for large datasets - if is_csv_format and self._should_use_streaming(result, form_data): - return self._create_streaming_csv_response(result, form_data, filename=filename) + if is_csv_format and True: + return self._create_streaming_csv_response(result, form_data, filename=filename, expected_rows=expected_rows) if len(result["queries"]) == 1: # return single query results @@ -532,6 +542,7 @@ def _get_data_response( form_data: dict[str, Any] | None = None, datasource: BaseDatasource | Query | None = None, filename: str | None = None, + expected_rows: int | None = None, ) -> Response: try: result = command.run(force_cached=force_cached) @@ -540,7 +551,7 @@ def _get_data_response( except ChartDataQueryFailedError as exc: return self.response_400(message=exc.message) - return self._send_chart_response(result, form_data, datasource, filename) + return self._send_chart_response(result, form_data, datasource, filename, expected_rows) # pylint: disable=invalid-name def _load_query_context_form_from_cache(self, cache_key: str) -> dict[str, Any]: @@ -590,7 +601,7 @@ def _should_use_streaming( return should_use_streaming_response(query_context, result_format) def _create_streaming_csv_response( - self, result: dict[Any, Any], form_data: dict[str, Any] | None = None, filename: str | None = None + self, result: dict[Any, Any], form_data: dict[str, Any] | None = None, filename: str | None = None, expected_rows: int | None = None ) -> Response: """Create a streaming CSV response for large datasets.""" from datetime import datetime @@ -616,8 +627,11 @@ def _create_streaming_csv_response( filename = f"superset_{safe_chart_name}_{timestamp}.csv" logger.info("Creating streaming CSV response: %s (from frontend: %s)", filename, filename is not None) + if expected_rows: + logger.info("๐Ÿ“Š Using expected_rows from frontend: %d", expected_rows) return create_streaming_csv_response( query_context=query_context, filename=filename, + expected_rows=expected_rows, ) diff --git a/superset/views/streaming.py b/superset/views/streaming.py index c590cefd4ef6..765a25e6ea54 100644 --- a/superset/views/streaming.py +++ b/superset/views/streaming.py @@ -97,6 +97,7 @@ def create_streaming_csv_response( filename: str | None = None, chunk_size: int | None = None, escape_formulas: bool = True, + expected_rows: int | None = None, ) -> Response: """ Factory function to create a streaming CSV response using Flask's standard pattern. @@ -136,15 +137,28 @@ def csv_generator() -> Generator[str, None, None]: "๐Ÿ“Š STREAMING CSV: Processing query with estimated large result set" ) - # Estimate total rows for progress tracking - # TODO: Get actual row count from frontend instead of hardcoding - # For now, hardcoded to 148795 for testing - estimated_rows = 148795 - logger.info( - "๐Ÿ“Š STREAMING CSV: Using hardcoded total_rows=%d for progress tracking", - estimated_rows, - ) + estimated_rows = expected_rows + + if estimated_rows: + logger.info( + "๐Ÿ“Š STREAMING CSV: Using expected_rows from frontend: %d", + estimated_rows, + ) + else: + form_data = query_context.form_data + estimated_rows = form_data.get('row_limit') if form_data else None + + if estimated_rows: + logger.info( + "๐Ÿ“Š STREAMING CSV: Using row_limit from form_data: %d", + estimated_rows, + ) + else: + logger.warning( + "โš ๏ธ STREAMING CSV: No expected_rows or row_limit available, progress percentage will be unavailable" + ) + estimated_rows = None # Initialize progress tracker progress_tracker.create_export( From 6b2b2399f528cf6006f7e40b35ff175be5f4f050 Mon Sep 17 00:00:00 2001 From: amaannawab923 Date: Sat, 4 Oct 2025 11:27:07 +0530 Subject: [PATCH 07/57] update: Streaming modal refactor --- .../StreamingExportModal.tsx | 503 ++++++++++-------- 1 file changed, 267 insertions(+), 236 deletions(-) diff --git a/superset-frontend/src/components/StreamingExportModal/StreamingExportModal.tsx b/superset-frontend/src/components/StreamingExportModal/StreamingExportModal.tsx index 6bbd5ad12827..b22032a38c37 100644 --- a/superset-frontend/src/components/StreamingExportModal/StreamingExportModal.tsx +++ b/superset-frontend/src/components/StreamingExportModal/StreamingExportModal.tsx @@ -18,26 +18,32 @@ */ import React from 'react'; import { styled, t } from '@superset-ui/core'; -import { Modal, Button, Typography, Space, Progress } from 'antd'; +import { Modal, Button, Typography, Progress, theme } from 'antd'; import { CheckCircleOutlined, CloseCircleOutlined, - DownloadOutlined, LoadingOutlined, StopOutlined, } from '@ant-design/icons'; -const { Text, Title } = Typography; +const { Text } = Typography; + +export enum ExportStatus { + STREAMING = 'streaming', + COMPLETED = 'completed', + ERROR = 'error', + CANCELLED = 'cancelled', +} export interface StreamingProgress { rowsProcessed: number; - totalRows?: number; // total expected rows for percentage calculation - totalSize: number; // in bytes - speed: number; // rows per second + totalRows?: number; + totalSize: number; + speed: number; mbPerSecond: number; - elapsedTime: number; // seconds - estimatedTimeRemaining?: number; // seconds - status: 'streaming' | 'completed' | 'error' | 'cancelled'; + elapsedTime: number; + estimatedTimeRemaining?: number; + status: ExportStatus; downloadUrl?: string; filename?: string; error?: string; @@ -52,31 +58,16 @@ interface StreamingExportModalProps { } const ModalContent = styled.div` - padding: ${({ theme }) => theme.sizeUnit * 4}px 0; - min-height: 200px; + padding: ${({ theme }) => theme.sizeUnit * 4}px 0 ${({ theme }) => theme.sizeUnit * 2}px; `; const ProgressSection = styled.div` margin: ${({ theme }) => theme.sizeUnit * 6}px 0; `; -const StatsGrid = styled.div` - display: grid; - grid-template-columns: 1fr 1fr; - gap: ${({ theme }) => theme.sizeUnit * 4}px; - margin: ${({ theme }) => theme.sizeUnit * 4}px 0; -`; - -const StatItem = styled.div` - text-align: center; - padding: ${({ theme }) => theme.sizeUnit * 3}px; - background: ${({ theme }) => theme.colorFillAlter}; - border-radius: ${({ theme }) => theme.borderRadius}px; -`; - const ActionButtons = styled.div` display: flex; - justify-content: center; + justify-content: flex-end; gap: ${({ theme }) => theme.sizeUnit * 3}px; margin-top: ${({ theme }) => theme.sizeUnit * 6}px; `; @@ -91,6 +82,51 @@ const StatusIcon = styled.div` } `; +const StyledIcon = styled.span<{ color: string }>` + color: ${({ color }) => color}; +`; + +const CenteredText = styled(Text)` + display: block; + text-align: center; + margin-bottom: ${({ theme }) => theme.sizeUnit * 4}px; +`; + +const ProgressText = styled(Text)` + display: block; + text-align: center; + margin-top: ${({ theme }) => theme.sizeUnit * 4}px; +`; + +const CancelButton = styled(Button)` + background-color: #f0fff8; + color: ${({ theme }) => theme.colorSuccess}; + border-color: transparent; + + &:hover { + background-color: #d9f7e8; + color: ${({ theme }) => theme.colorSuccess}; + border-color: transparent; + } +`; + +const DownloadButton = styled(Button)` + background-color: #2fc096; + color: white; + border-color: #2fc096; + + &:hover { + background-color: #26a77e; + border-color: #26a77e; + } + + &:disabled { + background-color: #f5f5f5; + color: rgba(0, 0, 0, 0.25); + border-color: #d9d9d9; + } +`; + const formatFileSize = (bytes: number): string => { if (bytes === 0) return '0 B'; const k = 1024; @@ -99,108 +135,195 @@ const formatFileSize = (bytes: number): string => { return `${parseFloat((bytes / Math.pow(k, i)).toFixed(1))} ${sizes[i]}`; }; -const formatTime = (seconds: number): string => { - if (seconds < 60) return `${Math.round(seconds)}s`; - const minutes = Math.floor(seconds / 60); - const remainingSeconds = Math.round(seconds % 60); - return `${minutes}m ${remainingSeconds}s`; -}; - const formatNumber = (num: number): string => new Intl.NumberFormat().format(num); +interface StatusIconProps { + status: ExportStatus; +} + +const StatusIconComponent: React.FC = ({ status }) => { + const { token } = theme.useToken(); + + const iconMap = { + [ExportStatus.STREAMING]: { Icon: LoadingOutlined, color: token.colorPrimary }, + [ExportStatus.COMPLETED]: { Icon: CheckCircleOutlined, color: token.colorSuccess }, + [ExportStatus.ERROR]: { Icon: CloseCircleOutlined, color: token.colorError }, + [ExportStatus.CANCELLED]: { Icon: StopOutlined, color: token.colorWarning }, + }; + + const { Icon, color } = iconMap[status] || iconMap[ExportStatus.STREAMING]; + return ( + + + + + + ); +}; + +interface ErrorContentProps { + error?: string; + onRetry?: () => void; + onCancel: () => void; +} + +const ErrorContent: React.FC = ({ + error, + onRetry, + onCancel, +}) => ( + + + + {error || t('An error occurred during export')} + + + {onRetry && ( + + )} + + + +); + +interface CancelledContentProps { + onRetry?: () => void; + onCancel: () => void; +} + +const CancelledContent: React.FC = ({ + onRetry, + onCancel, +}) => ( + + + {t('Export was cancelled')} + + {onRetry && ( + + )} + + + +); + +interface CompletedContentProps { + totalRows?: number; + rowsProcessed: number; + totalSize: number; + filename?: string; + downloadUrl?: string; + onDownload: () => void; + onCancel: () => void; +} + +const CompletedContent: React.FC = ({ + totalRows, + rowsProcessed, + totalSize, + filename, + downloadUrl, + onDownload, + onCancel, +}) => { + const { token } = theme.useToken(); + + return ( + + + + + {t('Export successful %s', filename || 'export.csv')} + + + + + + {t('Cancel')} + + + {t('Download')} + + + + ); +}; + +interface StreamingContentProps { + percentage: number; + filename?: string; + onCancel: () => void; +} + +const StreamingContent: React.FC = ({ + percentage, + filename, + onCancel, +}) => { + const { token } = theme.useToken(); + + return ( + + + `${Math.round(percent || 0)}%`} + /> + + {filename + ? t('Processing export for %s', filename) + : t('Processing export for {dashboard_name}_{YYYY-MM-DD}_{HHMMSS}.csv')} + + + + + + {t('Cancel')} + + + {t('Download')} + + + + ); +}; + const StreamingExportModal: React.FC = ({ visible, onCancel, onRetry, progress, - exportType, }) => { - const { - rowsProcessed, - totalSize, - speed, - mbPerSecond, - elapsedTime, - estimatedTimeRemaining, - status, - downloadUrl, - filename, - error, - } = progress; + const { totalRows, rowsProcessed, totalSize, status, downloadUrl, filename, error } = + progress; const getProgressPercentage = (): number => { - if (status === 'completed') return 100; - - // Calculate actual percentage based on rows processed vs total expected rows - if (progress.totalRows && progress.totalRows > 0) { - // Use rows processed even if it's 0 - this allows for 0% at start - const percentage = Math.min( - 99, - (rowsProcessed / progress.totalRows) * 100, - ); - const rounded = Math.round(percentage); - - // ๐Ÿ” DEBUG: Log percentage calculation - console.log('๐ŸŽฏ FRONTEND MODAL PERCENTAGE:', { - rowsProcessed, - totalRows: progress.totalRows, - rawPercentage: percentage.toFixed(2), - roundedPercentage: rounded, - status, - }); - - return rounded; - } - - // Fallback: estimate based on time if we have estimatedTimeRemaining - if (estimatedTimeRemaining && elapsedTime > 0) { - const totalEstimatedTime = elapsedTime + estimatedTimeRemaining; - const percentage = Math.min( - 95, - Math.round((elapsedTime / totalEstimatedTime) * 100), - ); + if (status === ExportStatus.COMPLETED) return 100; - console.log('โฐ FRONTEND MODAL TIME-BASED:', { - elapsedTime, - estimatedTimeRemaining, - totalEstimatedTime, - percentage, - }); - - return percentage; + if (totalRows && totalRows > 0) { + const percentage = Math.min(99, (rowsProcessed / totalRows) * 100); + return Math.round(percentage); } - // Default fallback for streaming status - const fallback = status === 'streaming' ? 10 : 0; - console.log('๐Ÿ”„ FRONTEND MODAL FALLBACK:', { - status, - fallback, - hasRows: rowsProcessed > 0, - hasTotalRows: !!progress.totalRows, - reasonForFallback: !progress.totalRows ? 'No totalRows' : 'Unknown', - }); - - return fallback; + return status === ExportStatus.STREAMING ? 10 : 0; }; - const renderStatusIcon = () => { - switch (status) { - case 'streaming': - return ; - case 'completed': - return ; - case 'error': - return ; - case 'cancelled': - return ; - default: - return ; - } - }; - - const getTitle = () => t('CSV Export'); - const handleDownload = () => { if (downloadUrl && filename) { const link = document.createElement('a'); @@ -209,145 +332,53 @@ const StreamingExportModal: React.FC = ({ document.body.appendChild(link); link.click(); document.body.removeChild(link); - onCancel(); // Close modal after download + onCancel(); } }; - const renderContent = () => { - if (status === 'error') { - return ( - - {renderStatusIcon()} - - {error || t('An error occurred during export')} - - - {onRetry && ( - - )} - - - + let content; + switch (status) { + case ExportStatus.ERROR: + content = ; + break; + case ExportStatus.CANCELLED: + content = ; + break; + case ExportStatus.COMPLETED: + content = ( + ); - } - - if (status === 'cancelled') { - return ( - - {renderStatusIcon()} - - {t('Export was cancelled')} - - - {onRetry && ( - - )} - - - + break; + default: + content = ( + ); - } - - if (status === 'completed') { - return ( - - {renderStatusIcon()} - - {t( - 'Successfully exported %s rows (%s)', - formatNumber(rowsProcessed), - formatFileSize(totalSize), - )} - - {filename && ( - - {filename} - - )} - - - - - - ); - } - - // Streaming status - return ( - - - `${Math.round(percent || 0)}%`} - /> - - {filename - ? t('Processing export for %s', filename) - : t( - 'Processing export for {dashboard_name}_{YYYY-MM-DD}_{HHMMSS}.csv', - )} - - - - - - - - - ); - }; + } return ( - {renderContent()} + {content} ); }; From 9e1625baf48cd39709633a3f2a30092680033b2f Mon Sep 17 00:00:00 2001 From: amaannawab923 Date: Sat, 4 Oct 2025 11:31:19 +0530 Subject: [PATCH 08/57] update: Refactor hook use streaming export --- .../useStreamingExport.ts | 215 ++---------------- 1 file changed, 21 insertions(+), 194 deletions(-) diff --git a/superset-frontend/src/components/StreamingExportModal/useStreamingExport.ts b/superset-frontend/src/components/StreamingExportModal/useStreamingExport.ts index 14bb778d6c73..446f574480e6 100644 --- a/superset-frontend/src/components/StreamingExportModal/useStreamingExport.ts +++ b/superset-frontend/src/components/StreamingExportModal/useStreamingExport.ts @@ -18,10 +18,9 @@ */ import { useState, useCallback, useRef } from 'react'; import { SupersetClient } from '@superset-ui/core'; -import { StreamingProgress } from './StreamingExportModal'; +import { ExportStatus, StreamingProgress } from './StreamingExportModal'; interface UseStreamingExportOptions { - onProgress?: (progress: StreamingProgress) => void; onComplete?: (downloadUrl: string, filename: string) => void; onError?: (error: string) => void; } @@ -31,7 +30,7 @@ interface StreamingExportParams { payload: any; filename?: string; exportType: 'csv' | 'xlsx'; - expectedRows?: number; // Total expected rows for progress calculation + expectedRows?: number; } export const useStreamingExport = (options: UseStreamingExportOptions = {}) => { @@ -42,125 +41,14 @@ export const useStreamingExport = (options: UseStreamingExportOptions = {}) => { speed: 0, mbPerSecond: 0, elapsedTime: 0, - status: 'streaming', + status: ExportStatus.STREAMING, }); const [isExporting, setIsExporting] = useState(false); const abortControllerRef = useRef(null); - const startTimeRef = useRef(0); - const lastUpdateRef = useRef(0); - const bytesHistoryRef = useRef>([]); const updateProgress = useCallback( (updates: Partial) => { - setProgress(prev => { - const newProgress = { ...prev, ...updates }; - - // ๐Ÿ” DEBUG: Log every progress update - console.log('๐Ÿ“Š FRONTEND PROGRESS UPDATE:', { - rowsProcessed: newProgress.rowsProcessed, - totalRows: newProgress.totalRows, - percentage: newProgress.totalRows - ? Math.round( - (newProgress.rowsProcessed / newProgress.totalRows) * 100, - ) - : 'N/A', - totalSize: `${(newProgress.totalSize / 1024).toFixed(1)}KB`, - status: newProgress.status, - elapsedTime: `${newProgress.elapsedTime.toFixed(1)}s`, - }); - - options.onProgress?.(newProgress); - return newProgress; - }); - }, - [options], - ); - - const parseProgressFromHeaders = useCallback( - (headers: Headers, elapsedTime: number) => { - // Try to extract progress from response headers if backend provides them - const progressHeader = headers.get('X-Export-Progress'); - if (progressHeader) { - try { - const [rows, bytes, rate] = progressHeader.split(',').map(Number); - return { - rowsProcessed: rows || 0, - totalSize: bytes || 0, - speed: rate || 0, - mbPerSecond: bytes > 0 ? bytes / (1024 * 1024) / elapsedTime : 0, - elapsedTime, - }; - } catch (e) { - console.warn('Failed to parse progress header:', e); - } - } - return null; - }, - [], - ); - - const estimateProgressFromSize = useCallback( - (currentSize: number, elapsedTime: number, actualRows: number = 0) => { - const now = Date.now(); - - // Track bytes over time for smoother speed calculation - bytesHistoryRef.current.push({ time: now, bytes: currentSize }); - - // Keep only last 10 seconds of data for rolling average - const cutoffTime = now - 10000; - bytesHistoryRef.current = bytesHistoryRef.current.filter( - entry => entry.time > cutoffTime, - ); - - if (elapsedTime <= 0) { - return { - rowsProcessed: actualRows, - totalSize: currentSize, - speed: 0, - mbPerSecond: 0, - elapsedTime, - }; - } - - // Calculate speed based on recent data points for smoother updates - let mbPerSecond = 0; - if (bytesHistoryRef.current.length >= 2) { - const oldest = bytesHistoryRef.current[0]; - const newest = - bytesHistoryRef.current[bytesHistoryRef.current.length - 1]; - const timeDiff = (newest.time - oldest.time) / 1000; - const bytesDiff = newest.bytes - oldest.bytes; - - if (timeDiff > 0) { - mbPerSecond = bytesDiff / (1024 * 1024) / timeDiff; - } - } - - // If rolling average isn't available yet, use overall average - if (mbPerSecond === 0) { - mbPerSecond = currentSize / (1024 * 1024) / elapsedTime; - } - - // Use actual row count if available, otherwise estimate - let finalRowCount = actualRows; - if (actualRows === 0) { - // Fallback to size-based estimation only if no actual count - const avgBytesPerRow = - currentSize > 10000 - ? currentSize / Math.max(1, currentSize / 150) - : 200; - finalRowCount = Math.floor(currentSize / avgBytesPerRow); - } - - const rowsPerSecond = finalRowCount / elapsedTime; - - return { - rowsProcessed: finalRowCount, - totalSize: currentSize, - speed: rowsPerSecond, - mbPerSecond, - elapsedTime, - }; + setProgress(prev => ({ ...prev, ...updates })); }, [], ); @@ -173,23 +61,10 @@ export const useStreamingExport = (options: UseStreamingExportOptions = {}) => { exportType, expectedRows, }: StreamingExportParams) => { - if (isExporting) { - console.warn('Export already in progress'); - return; - } + if (isExporting) return; setIsExporting(true); abortControllerRef.current = new AbortController(); - startTimeRef.current = Date.now(); - lastUpdateRef.current = Date.now(); - bytesHistoryRef.current = []; // Reset bytes history for new export - - console.log('๐Ÿš€ FRONTEND STREAMING START:', { - url, - filename, - expectedRows, - exportType, - }); updateProgress({ rowsProcessed: 0, @@ -198,17 +73,15 @@ export const useStreamingExport = (options: UseStreamingExportOptions = {}) => { speed: 0, mbPerSecond: 0, elapsedTime: 0, - status: 'streaming', + status: ExportStatus.STREAMING, filename, }); - // Start polling for progress using filename as export ID const pollInterval = setInterval(async () => { try { const progressUrl = `/api/v1/chart/export/progress/${encodeURIComponent( filename || `export.${exportType}`, )}`; - console.log('๐Ÿ“Š POLLING PROGRESS:', progressUrl); const progressResponse = await SupersetClient.get({ endpoint: progressUrl, @@ -216,9 +89,7 @@ export const useStreamingExport = (options: UseStreamingExportOptions = {}) => { if (progressResponse.json) { const serverProgress = progressResponse.json; - console.log('๐Ÿ“ˆ SERVER PROGRESS:', serverProgress); - // Update progress from server response updateProgress({ rowsProcessed: serverProgress.rows_processed || 0, totalRows: serverProgress.total_rows || expectedRows, @@ -226,30 +97,25 @@ export const useStreamingExport = (options: UseStreamingExportOptions = {}) => { speed: serverProgress.speed_rows_per_sec || 0, mbPerSecond: serverProgress.speed_mb_per_sec || 0, elapsedTime: serverProgress.elapsed_time || 0, - status: serverProgress.status || 'streaming', + status: serverProgress.status || ExportStatus.STREAMING, }); - // Stop polling if completed or error if ( serverProgress.status === 'completed' || serverProgress.status === 'error' ) { - console.log('โœ… POLLING STOPPED: Export finished'); clearInterval(pollInterval); } } - } catch (pollError) { - // Silently ignore polling errors - the main stream will handle failures - console.warn('โš ๏ธ Progress polling error (non-critical):', pollError); + } catch (error) { + // Ignore polling errors } - }, 500); // Poll every 500ms as requested + }, 500); try { - // Initialize SupersetClient to ensure authentication and get CSRF token await SupersetClient.init(); const csrfToken = await SupersetClient.getCSRFToken(); - // Use manual fetch for streaming while leveraging SupersetClient's authentication const response = await fetch(url, { method: 'POST', headers: { @@ -258,20 +124,13 @@ export const useStreamingExport = (options: UseStreamingExportOptions = {}) => { }, body: new URLSearchParams({ form_data: JSON.stringify(payload), - filename: filename || `export.${exportType}`, // Pass filename to backend - expected_rows: expectedRows?.toString() || '', // Pass expected row count to backend + filename: filename || `export.${exportType}`, + expected_rows: expectedRows?.toString() || '', }), signal: abortControllerRef.current.signal, credentials: 'same-origin', }); - console.log('๐Ÿ“ก FRONTEND RESPONSE:', { - status: response.status, - statusText: response.statusText, - headers: Object.fromEntries(response.headers.entries()), - hasBody: !!response.body, - }); - if (!response.ok) { throw new Error( `Export failed: ${response.status} ${response.statusText}`, @@ -282,51 +141,26 @@ export const useStreamingExport = (options: UseStreamingExportOptions = {}) => { throw new Error('Response body is not available for streaming'); } - console.log('๐Ÿ”ง STARTING STREAM READER'); const reader = response.body.getReader(); const chunks: Uint8Array[] = []; let receivedLength = 0; - let chunkCount = 0; - - // Note: We rely on the polling interval (pollInterval) for progress updates - // No need for a separate progressInterval since the server tracks actual progress try { - // Read the streaming response - console.log('๐Ÿ”„ ENTERING STREAM READ LOOP'); while (true) { - console.log('๐Ÿ”„ CALLING reader.read()...'); const { done, value } = await reader.read(); - console.log('๐Ÿ“ก READER RESULT:', { - done, - hasValue: !!value, - valueLength: value?.length, - }); - - if (done) { - console.log('โœ… STREAM READING COMPLETED'); - break; - } + if (done) break; if (abortControllerRef.current?.signal.aborted) { throw new Error('Export cancelled by user'); } - chunkCount++; chunks.push(value); receivedLength += value.length; - - console.log('๐Ÿ“Š CHUNK RECEIVED:', { - chunkNumber: chunkCount, - chunkSize: value.length, - totalReceived: receivedLength, - }); } - clearInterval(pollInterval); // Stop polling when stream completes + clearInterval(pollInterval); - // Create blob from chunks const completeData = new Uint8Array(receivedLength); let position = 0; for (const chunk of chunks) { @@ -342,16 +176,15 @@ export const useStreamingExport = (options: UseStreamingExportOptions = {}) => { const blob = new Blob([completeData], { type: mimeType }); const downloadUrl = URL.createObjectURL(blob); - // Mark as completed - final progress will come from the last poll updateProgress({ - status: 'completed', + status: ExportStatus.COMPLETED, downloadUrl, filename: filename || `export.${exportType}`, }); options.onComplete?.(downloadUrl, filename || `export.${exportType}`); } catch (streamError) { - clearInterval(pollInterval); // Stop polling on error + clearInterval(pollInterval); throw streamError; } } catch (error) { @@ -363,11 +196,11 @@ export const useStreamingExport = (options: UseStreamingExportOptions = {}) => { errorMessage.includes('aborted') ) { updateProgress({ - status: 'cancelled', + status: ExportStatus.CANCELLED, }); } else { updateProgress({ - status: 'error', + status: ExportStatus.ERROR, error: errorMessage, }); options.onError?.(errorMessage); @@ -377,20 +210,14 @@ export const useStreamingExport = (options: UseStreamingExportOptions = {}) => { abortControllerRef.current = null; } }, - [ - isExporting, - updateProgress, - parseProgressFromHeaders, - estimateProgressFromSize, - options, - ], + [isExporting, updateProgress, options], ); const cancelExport = useCallback(() => { if (abortControllerRef.current) { abortControllerRef.current.abort(); updateProgress({ - status: 'cancelled', + status: ExportStatus.CANCELLED, }); } }, [updateProgress]); @@ -405,7 +232,7 @@ export const useStreamingExport = (options: UseStreamingExportOptions = {}) => { speed: 0, mbPerSecond: 0, elapsedTime: 0, - status: 'streaming', + status: ExportStatus.STREAMING, }); }, []); From e2419ef622bcdaefcff3380fc88031f9635617a8 Mon Sep 17 00:00:00 2001 From: amaannawab923 Date: Sat, 4 Oct 2025 11:38:39 +0530 Subject: [PATCH 09/57] streaming export --- .../src/components/StreamingExportModal/useStreamingExport.ts | 3 --- 1 file changed, 3 deletions(-) diff --git a/superset-frontend/src/components/StreamingExportModal/useStreamingExport.ts b/superset-frontend/src/components/StreamingExportModal/useStreamingExport.ts index 446f574480e6..05fa05fe0926 100644 --- a/superset-frontend/src/components/StreamingExportModal/useStreamingExport.ts +++ b/superset-frontend/src/components/StreamingExportModal/useStreamingExport.ts @@ -113,14 +113,11 @@ export const useStreamingExport = (options: UseStreamingExportOptions = {}) => { }, 500); try { - await SupersetClient.init(); - const csrfToken = await SupersetClient.getCSRFToken(); const response = await fetch(url, { method: 'POST', headers: { 'Content-Type': 'application/x-www-form-urlencoded', - ...(csrfToken && { 'X-CSRFToken': csrfToken }), }, body: new URLSearchParams({ form_data: JSON.stringify(payload), From e67059b582652c4c00427b7a3c27b2f3fa3f405b Mon Sep 17 00:00:00 2001 From: amaannawab923 Date: Sat, 4 Oct 2025 11:44:01 +0530 Subject: [PATCH 10/57] update: correct toast messages --- .../components/gridComponents/Chart/Chart.jsx | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/superset-frontend/src/dashboard/components/gridComponents/Chart/Chart.jsx b/superset-frontend/src/dashboard/components/gridComponents/Chart/Chart.jsx index f649cb8c00d5..7efda7144fcb 100644 --- a/superset-frontend/src/dashboard/components/gridComponents/Chart/Chart.jsx +++ b/superset-frontend/src/dashboard/components/gridComponents/Chart/Chart.jsx @@ -80,8 +80,7 @@ const propTypes = { isInView: PropTypes.bool, }; -// we use state + shouldComponentUpdate() logic to prevent perf-wrecking -// resizing across all slices on a dashboard on every update + const RESIZE_TIMEOUT = 500; const DEFAULT_HEADER_HEIGHT = 22; @@ -180,17 +179,16 @@ const Chart = props => { const [height, setHeight] = useState(props.height); const [width, setWidth] = useState(props.width); - // Streaming export state const [isStreamingModalVisible, setIsStreamingModalVisible] = useState(false); const { progress, isExporting, startExport, cancelExport, resetExport } = useStreamingExport({ onComplete: (downloadUrl, filename) => { boundActionCreators.addSuccessToast( - t('Export completed successfully: %s', filename), + t('CSV file downloaded successfully'), ); }, - onError: error => { - boundActionCreators.addDangerToast(t('Export failed: %s', error)); + onError: () => { + boundActionCreators.addDangerToast(t('Export failed - please try again')); }, }); const history = useHistory(); From c309932c3cd87761fa770dbb4e0062c99d5d56a6 Mon Sep 17 00:00:00 2001 From: amaannawab923 Date: Mon, 6 Oct 2025 00:04:15 +0530 Subject: [PATCH 11/57] update: remove polling api --- .../useStreamingExport.ts | 59 ++---- superset-frontend/webpack.proxy-config.js | 19 +- superset/charts/data/api.py | 105 +-------- superset/views/streaming.py | 69 ++++-- superset/views/streaming_progress.py | 199 ------------------ 5 files changed, 85 insertions(+), 366 deletions(-) delete mode 100644 superset/views/streaming_progress.py diff --git a/superset-frontend/src/components/StreamingExportModal/useStreamingExport.ts b/superset-frontend/src/components/StreamingExportModal/useStreamingExport.ts index 05fa05fe0926..9c87afaeba20 100644 --- a/superset-frontend/src/components/StreamingExportModal/useStreamingExport.ts +++ b/superset-frontend/src/components/StreamingExportModal/useStreamingExport.ts @@ -17,7 +17,6 @@ * under the License. */ import { useState, useCallback, useRef } from 'react'; -import { SupersetClient } from '@superset-ui/core'; import { ExportStatus, StreamingProgress } from './StreamingExportModal'; interface UseStreamingExportOptions { @@ -77,43 +76,7 @@ export const useStreamingExport = (options: UseStreamingExportOptions = {}) => { filename, }); - const pollInterval = setInterval(async () => { - try { - const progressUrl = `/api/v1/chart/export/progress/${encodeURIComponent( - filename || `export.${exportType}`, - )}`; - - const progressResponse = await SupersetClient.get({ - endpoint: progressUrl, - }); - - if (progressResponse.json) { - const serverProgress = progressResponse.json; - - updateProgress({ - rowsProcessed: serverProgress.rows_processed || 0, - totalRows: serverProgress.total_rows || expectedRows, - totalSize: serverProgress.bytes_processed || 0, - speed: serverProgress.speed_rows_per_sec || 0, - mbPerSecond: serverProgress.speed_mb_per_sec || 0, - elapsedTime: serverProgress.elapsed_time || 0, - status: serverProgress.status || ExportStatus.STREAMING, - }); - - if ( - serverProgress.status === 'completed' || - serverProgress.status === 'error' - ) { - clearInterval(pollInterval); - } - } - } catch (error) { - // Ignore polling errors - } - }, 500); - try { - const response = await fetch(url, { method: 'POST', headers: { @@ -141,6 +104,8 @@ export const useStreamingExport = (options: UseStreamingExportOptions = {}) => { const reader = response.body.getReader(); const chunks: Uint8Array[] = []; let receivedLength = 0; + let rowsProcessed = 0; + const NEWLINE_BYTE = 10; // '\n' character code try { while (true) { @@ -154,9 +119,24 @@ export const useStreamingExport = (options: UseStreamingExportOptions = {}) => { chunks.push(value); receivedLength += value.length; - } - clearInterval(pollInterval); + // Count newlines directly in binary (faster than decoding + regex) + let newlineCount = 0; + for (let i = 0; i < value.length; i++) { + if (value[i] === NEWLINE_BYTE) { + newlineCount++; + } + } + rowsProcessed += newlineCount; + + // Update progress based on rows processed + updateProgress({ + status: ExportStatus.STREAMING, + rowsProcessed, + totalRows: expectedRows, + totalSize: receivedLength, + }); + } const completeData = new Uint8Array(receivedLength); let position = 0; @@ -181,7 +161,6 @@ export const useStreamingExport = (options: UseStreamingExportOptions = {}) => { options.onComplete?.(downloadUrl, filename || `export.${exportType}`); } catch (streamError) { - clearInterval(pollInterval); throw streamError; } } catch (error) { diff --git a/superset-frontend/webpack.proxy-config.js b/superset-frontend/webpack.proxy-config.js index 41835d88faf0..47b616ed7cd6 100644 --- a/superset-frontend/webpack.proxy-config.js +++ b/superset-frontend/webpack.proxy-config.js @@ -165,7 +165,24 @@ module.exports = newManifest => { if (isHTML(response)) { processHTML(proxyResponse, response); } else { - proxyResponse.pipe(response); + const isStreaming = proxyResponse.headers['x-superset-streaming'] === 'true'; + + if (isStreaming) { + proxyResponse.on('data', chunk => { + response.write(chunk); + if (response.flush) { + response.flush(); + } + }); + proxyResponse.on('end', () => { + response.end(); + }); + proxyResponse.on('error', () => { + response.end(); + }); + } else { + proxyResponse.pipe(response); + } } response.flushHeaders(); } catch (e) { diff --git a/superset/charts/data/api.py b/superset/charts/data/api.py index 7848de674e7f..5c76cd3b30c9 100644 --- a/superset/charts/data/api.py +++ b/superset/charts/data/api.py @@ -62,7 +62,7 @@ class ChartDataRestApi(ChartRestApi): - include_route_methods = {"get_data", "data", "data_from_cache", "export_progress"} + include_route_methods = {"get_data", "data", "data_from_cache"} @expose("//data/", methods=("GET",)) @protect() @@ -334,109 +334,6 @@ def data_from_cache(self, cache_key: str) -> Response: return self._get_data_response(command, True) - @expose("/export/progress/", methods=("GET",)) - @statsd_metrics - @event_logger.log_this_with_context( - action=lambda self, *args, **kwargs: f"{self.__class__.__name__}" - f".export_progress", - log_to_statsd=False, - ) - def export_progress(self, export_id: str) -> Response: - """ - Poll progress for a streaming CSV export. - --- - get: - summary: Get streaming export progress - description: >- - Returns the current progress of a streaming CSV export. - The export_id is the filename being exported. - This endpoint is polled by the frontend every 500ms. - parameters: - - in: path - schema: - type: string - name: export_id - description: The export ID (filename) - responses: - 200: - description: Export progress - content: - application/json: - schema: - type: object - properties: - export_id: - type: string - status: - type: string - enum: [streaming, completed, error] - rows_processed: - type: integer - total_rows: - type: integer - bytes_processed: - type: integer - elapsed_time: - type: number - percentage: - type: number - speed_rows_per_sec: - type: number - speed_mb_per_sec: - type: number - error_message: - type: string - 404: - $ref: '#/components/responses/404' - 401: - $ref: '#/components/responses/401' - """ - import logging - - logger = logging.getLogger(__name__) - - logger.info("๐Ÿ” POLLING API CALLED: export_id=%s", export_id) - - from superset.views.streaming_progress import progress_tracker - - logger.info("๐Ÿ“Š POLLING API: Getting progress for export_id=%s", export_id) - - # Get progress from tracker - progress = progress_tracker.get_progress(export_id) - - logger.info("๐Ÿ“ˆ POLLING API: Progress result=%s", progress) - - if progress is None: - logger.warning( - "โš ๏ธ POLLING API: Export not found: export_id=%s", export_id - ) - # Return a response indicating export not started yet or already completed - return self.response( - 200, - export_id=export_id, - status="not_found", - rows_processed=0, - total_rows=None, - bytes_processed=0, - elapsed_time=0, - percentage=None, - speed_rows_per_sec=0, - speed_mb_per_sec=0, - error_message="Export not found or not started yet", - ) - - logger.info( - "โœ… POLLING API: Returning progress: status=%s, rows=%s, bytes=%s", - progress.get("status"), - progress.get("rows_processed"), - progress.get("bytes_processed"), - ) - - # Cleanup old exports periodically - progress_tracker.cleanup_old_exports() - - return self.response(200, **progress) - def _run_async( self, form_data: dict[str, Any], command: ChartDataCommand ) -> Response: diff --git a/superset/views/streaming.py b/superset/views/streaming.py index 765a25e6ea54..c1b593cf5bd2 100644 --- a/superset/views/streaming.py +++ b/superset/views/streaming.py @@ -28,8 +28,6 @@ from flask import current_app as app, Response from werkzeug.datastructures import Headers -from superset.views.streaming_progress import progress_tracker - if TYPE_CHECKING: from superset.common.query_context import QueryContext @@ -48,16 +46,23 @@ def create_streaming_csv_response_simple( """ from flask import Response - # Create response with proper headers + # Create response with proper headers to disable buffering + # CRITICAL: Set direct_passthrough=False to ensure Flask actually iterates the generator response = Response( data_generator, mimetype=f"text/csv; charset={encoding}", headers={ "Content-Disposition": f'attachment; filename="{filename}"', "Cache-Control": "no-cache", + "X-Accel-Buffering": "no", # Disable nginx buffering + "X-Superset-Streaming": "true", # Identify streaming responses }, + direct_passthrough=False, # Flask must iterate generator, not pass to WSGI directly ) + # Force chunked transfer encoding (critical for streaming) + response.implicit_sequence_conversion = False + logger.info("Created simple streaming CSV response for file: %s", filename) return response @@ -160,11 +165,6 @@ def csv_generator() -> Generator[str, None, None]: ) estimated_rows = None - # Initialize progress tracker - progress_tracker.create_export( - export_id=export_id, total_rows=estimated_rows - ) - # Get the database connection and execute raw SQL query directly from superset import db from superset.connectors.sqla.models import SqlaTable @@ -247,13 +247,19 @@ def csv_generator() -> Generator[str, None, None]: streaming_start_time = time.time() last_progress_time = streaming_start_time + # Buffer to accumulate data before yielding (forces flush) + # WSGI servers need ~50KB minimum to start streaming + buffer = [] + buffer_size = 0 + FLUSH_THRESHOLD = 65536 # 64KB - exceeds WSGI buffering threshold + while True: # Fetch chunk of rows rows = result_proxy.fetchmany(chunk_size) if not rows: break - # Yield CSV rows + # Build CSV rows and accumulate in buffer for row in rows: csv_row = ",".join( f'"{str(cell) if cell is not None else ""}"' @@ -262,15 +268,34 @@ def csv_generator() -> Generator[str, None, None]: csv_line = csv_row + "\n" row_bytes = len(csv_line.encode("utf-8")) total_bytes += row_bytes - yield csv_line row_count += 1 - # Update progress tracker after each chunk - progress_tracker.update_progress( - export_id=export_id, - rows_processed=row_count, - bytes_processed=total_bytes, - ) + buffer.append(csv_line) + buffer_size += row_bytes + + # Yield when buffer exceeds threshold (forces immediate flush) + if buffer_size >= FLUSH_THRESHOLD: + chunk_data = "".join(buffer) + logger.info( + "๐Ÿ”ฅ YIELDING CHUNK: %d bytes at %s", + len(chunk_data), + time.strftime("%H:%M:%S") + ) + yield chunk_data + buffer = [] + buffer_size = 0 + + # Apply testing delay after yield for visible streaming + if ( + ENABLE_SLOW_STREAMING_TEST + and delay_between_chunks > 0 + ): + logger.info( + "โฑ๏ธ SLEEPING for %s seconds...", + delay_between_chunks + ) + time.sleep(delay_between_chunks) + logger.info("โฑ๏ธ WAKE UP - continuing...") # Performance logging every 10k rows or 5 seconds current_time = time.time() @@ -305,6 +330,12 @@ def csv_generator() -> Generator[str, None, None]: ): time.sleep(delay_between_chunks) + # Flush remaining buffer + if buffer: + yield "".join(buffer) + buffer = [] + buffer_size = 0 + # Final performance summary total_time = time.time() - start_time streaming_time = time.time() - streaming_start_time @@ -330,9 +361,6 @@ def csv_generator() -> Generator[str, None, None]: chunk_size, ) - # Mark export as completed in progress tracker - progress_tracker.complete_export(export_id) - finally: connection.close() @@ -342,9 +370,6 @@ def csv_generator() -> Generator[str, None, None]: logger.error("Traceback: %s", traceback.format_exc()) - # Mark export as failed in progress tracker - progress_tracker.fail_export(export_id, str(e)) - # Yield error info and fallback data yield f"# Error occurred: {str(e)}\n" yield "error,message\n" diff --git a/superset/views/streaming_progress.py b/superset/views/streaming_progress.py deleted file mode 100644 index 3bfdada6b6d5..000000000000 --- a/superset/views/streaming_progress.py +++ /dev/null @@ -1,199 +0,0 @@ -# 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. - -"""Thread-safe progress tracking for streaming CSV exports.""" - -from __future__ import annotations - -import logging -import threading -import time -from dataclasses import dataclass -from typing import Any - -logger = logging.getLogger(__name__) - - -@dataclass -class ExportProgress: - """Progress information for a streaming export.""" - - export_id: str - status: str = "streaming" # streaming, completed, error - rows_processed: int = 0 - total_rows: int | None = None - bytes_processed: int = 0 - elapsed_time: float = 0.0 - start_time: float = 0.0 - error_message: str | None = None - - def to_dict(self) -> dict[str, Any]: - """Convert to dictionary for JSON response.""" - percentage = None - if self.total_rows and self.total_rows > 0: - percentage = round((self.rows_processed / self.total_rows) * 100, 1) - - # Calculate speeds - speed_rows_per_sec = 0.0 - speed_mb_per_sec = 0.0 - if self.elapsed_time > 0: - speed_rows_per_sec = self.rows_processed / self.elapsed_time - speed_mb_per_sec = (self.bytes_processed / self.elapsed_time) / ( - 1024 * 1024 - ) - - return { - "export_id": self.export_id, - "status": self.status, - "rows_processed": self.rows_processed, - "total_rows": self.total_rows, - "bytes_processed": self.bytes_processed, - "elapsed_time": round(self.elapsed_time, 2), - "percentage": percentage, - "speed_rows_per_sec": round(speed_rows_per_sec, 2), - "speed_mb_per_sec": round(speed_mb_per_sec, 3), - "error_message": self.error_message, - } - - -class StreamingProgressTracker: - """ - Thread-safe singleton for tracking streaming export progress. - - This allows the streaming generator to update progress while a separate - polling endpoint reads the current state without blocking the stream. - """ - - _instance: StreamingProgressTracker | None = None - _lock = threading.Lock() - - def __new__(cls) -> StreamingProgressTracker: - """Singleton pattern to ensure single instance.""" - if cls._instance is None: - with cls._lock: - if cls._instance is None: - cls._instance = super().__new__(cls) - cls._instance._initialized = False - return cls._instance - - def __init__(self) -> None: - """Initialize the progress tracker.""" - if self._initialized: - return - - self._progress_map: dict[str, ExportProgress] = {} - self._progress_lock = threading.Lock() - self._cleanup_interval = 300 # Cleanup after 5 minutes - self._last_cleanup = time.time() - self._initialized = True - - logger.info("โœ… StreamingProgressTracker initialized") - - def create_export( - self, export_id: str, total_rows: int | None = None - ) -> None: - """Create a new export progress entry.""" - with self._progress_lock: - self._progress_map[export_id] = ExportProgress( - export_id=export_id, - total_rows=total_rows, - start_time=time.time(), - ) - logger.info( - "๐Ÿ“Š Created progress tracker for export: %s (expected %s rows)", - export_id, - total_rows or "unknown", - ) - - def update_progress( - self, - export_id: str, - rows_processed: int, - bytes_processed: int, - ) -> None: - """Update progress for an export.""" - with self._progress_lock: - if export_id not in self._progress_map: - logger.warning("โš ๏ธ Export ID not found: %s", export_id) - return - - progress = self._progress_map[export_id] - progress.rows_processed = rows_processed - progress.bytes_processed = bytes_processed - progress.elapsed_time = time.time() - progress.start_time - - def complete_export(self, export_id: str) -> None: - """Mark export as completed.""" - with self._progress_lock: - if export_id in self._progress_map: - self._progress_map[export_id].status = "completed" - self._progress_map[export_id].elapsed_time = ( - time.time() - self._progress_map[export_id].start_time - ) - logger.info("โœ… Export completed: %s", export_id) - - def fail_export(self, export_id: str, error_message: str) -> None: - """Mark export as failed.""" - with self._progress_lock: - if export_id in self._progress_map: - self._progress_map[export_id].status = "error" - self._progress_map[export_id].error_message = error_message - self._progress_map[export_id].elapsed_time = ( - time.time() - self._progress_map[export_id].start_time - ) - logger.error("โŒ Export failed: %s - %s", export_id, error_message) - - def get_progress(self, export_id: str) -> dict[str, Any] | None: - """Get current progress for an export.""" - with self._progress_lock: - progress = self._progress_map.get(export_id) - if progress: - # Update elapsed time before returning - progress.elapsed_time = time.time() - progress.start_time - return progress.to_dict() - return None - - def cleanup_old_exports(self, max_age_seconds: int = 300) -> None: - """Remove old completed/failed exports to prevent memory leaks.""" - current_time = time.time() - - # Only cleanup periodically - if current_time - self._last_cleanup < self._cleanup_interval: - return - - with self._progress_lock: - to_remove = [] - for export_id, progress in self._progress_map.items(): - age = current_time - progress.start_time - if age > max_age_seconds and progress.status in ["completed", "error"]: - to_remove.append(export_id) - - for export_id in to_remove: - del self._progress_map[export_id] - - self._last_cleanup = current_time - - if to_remove: - logger.info( - "๐Ÿงน Cleaned up %d old exports, %d active remaining", - len(to_remove), - len(self._progress_map), - ) - - -# Global singleton instance -progress_tracker = StreamingProgressTracker() From 04bb069a718aad902d95d552cdcfa60f4610ea50 Mon Sep 17 00:00:00 2001 From: amaannawab923 Date: Mon, 6 Oct 2025 00:21:13 +0530 Subject: [PATCH 12/57] Fix: streaming refactor of row estimation function --- .../components/gridComponents/Chart/Chart.jsx | 13 ++- superset/charts/data/api.py | 39 ++++++++- superset/config.py | 5 ++ superset/views/base.py | 1 + superset/views/streaming.py | 84 ------------------- 5 files changed, 51 insertions(+), 91 deletions(-) diff --git a/superset-frontend/src/dashboard/components/gridComponents/Chart/Chart.jsx b/superset-frontend/src/dashboard/components/gridComponents/Chart/Chart.jsx index 7efda7144fcb..f1bfcc426b5c 100644 --- a/superset-frontend/src/dashboard/components/gridComponents/Chart/Chart.jsx +++ b/superset-frontend/src/dashboard/components/gridComponents/Chart/Chart.jsx @@ -160,6 +160,9 @@ const Chart = props => { const maxRows = useSelector( state => state.dashboardInfo.common.conf.SQL_MAX_ROW, ); + const streamingThreshold = useSelector( + state => state.dashboardInfo.common.conf.CSV_STREAMING_ROW_THRESHOLD || 100000, + ); const datasource = useSelector( state => (chart && @@ -182,7 +185,7 @@ const Chart = props => { const [isStreamingModalVisible, setIsStreamingModalVisible] = useState(false); const { progress, isExporting, startExport, cancelExport, resetExport } = useStreamingExport({ - onComplete: (downloadUrl, filename) => { + onComplete: () => { boundActionCreators.addSuccessToast( t('CSV file downloaded successfully'), ); @@ -413,8 +416,11 @@ const Chart = props => { actualRowCount = exportFormData?.row_limit; } - // Handle streaming CSV exports for both regular and full data exports - const shouldUseStreaming = format === 'csv' && !isPivot; + // Handle streaming CSV exports based on row threshold + const shouldUseStreaming = + format === 'csv' && + !isPivot && + actualRowCount >= streamingThreshold; exportChart({ formData: exportFormData, @@ -444,6 +450,7 @@ const Chart = props => { queriesResponse, startExport, resetExport, + streamingThreshold, ], ); diff --git a/superset/charts/data/api.py b/superset/charts/data/api.py index 5c76cd3b30c9..b58a70e5fd47 100644 --- a/superset/charts/data/api.py +++ b/superset/charts/data/api.py @@ -386,7 +386,7 @@ def _send_chart_response( # noqa: C901 is_csv_format = result_format == ChartDataResultFormat.CSV # Check if we should use streaming for large datasets - if is_csv_format and True: + if is_csv_format and self._should_use_streaming(result, form_data): return self._create_streaming_csv_response(result, form_data, filename=filename, expected_rows=expected_rows) if len(result["queries"]) == 1: @@ -489,13 +489,44 @@ def _create_query_context_from_form( def _should_use_streaming( self, result: dict[Any, Any], form_data: dict[str, Any] | None = None ) -> bool: - """Determine if streaming should be used for this response.""" - from superset.views.streaming import should_use_streaming_response + """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 - return should_use_streaming_response(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 + if actual_row_count is None: + if form_data and "row_limit" in form_data: + actual_row_count = form_data.get("row_limit", 0) + elif query_context.form_data and "row_limit" in query_context.form_data: + actual_row_count = query_context.form_data.get("row_limit", 0) + + # Use streaming if row count meets or exceeds threshold + if actual_row_count is not None and actual_row_count >= threshold: + return True + + return False def _create_streaming_csv_response( self, result: dict[Any, Any], form_data: dict[str, Any] | None = None, filename: str | None = None, expected_rows: int | None = None diff --git a/superset/config.py b/superset/config.py index 170759740557..947720ea028d 100644 --- a/superset/config.py +++ b/superset/config.py @@ -1009,6 +1009,11 @@ class D3TimeFormat(TypedDict, total=False): # note: index option should not be overridden CSV_EXPORT = {"encoding": "utf-8-sig"} +# CSV Streaming: row threshold for using streaming CSV exports +# When row count >= this threshold, use streaming response instead of loading all data +# into memory. Streaming provides real-time progress and handles large datasets efficiently. +CSV_STREAMING_ROW_THRESHOLD = 100000 + # Excel Options: key/value pairs that will be passed as argument to DataFrame.to_excel # method. # note: index option should not be overridden diff --git a/superset/views/base.py b/superset/views/base.py index bf493628ad5d..47237eb34b40 100644 --- a/superset/views/base.py +++ b/superset/views/base.py @@ -121,6 +121,7 @@ "SYNC_DB_PERMISSIONS_IN_ASYNC_MODE", "TABLE_VIZ_MAX_ROW_SERVER", "MAPBOX_API_KEY", + "CSV_STREAMING_ROW_THRESHOLD", ) logger = logging.getLogger(__name__) diff --git a/superset/views/streaming.py b/superset/views/streaming.py index c1b593cf5bd2..ba2932c5b727 100644 --- a/superset/views/streaming.py +++ b/superset/views/streaming.py @@ -388,87 +388,3 @@ def csv_generator() -> Generator[str, None, None]: ) -def should_use_streaming_response( - query_context: QueryContext, - result_format: str, -) -> bool: - """ - Determine if a streaming response should be used. - - Args: - query_context: The query context - result_format: The requested result format (csv, xlsx, etc.) - - Returns: - True if streaming should be used - """ - # Only stream for supported formats - if result_format.lower() not in ["csv"]: # Add "xlsx" when implemented - return False - - # Check if streaming is enabled in config - config = app.config.get("CSV_STREAMING", {}) - if not config.get("enabled", True): - return False - - # Simple row count estimation - always trigger streaming for CSV exports - # TODO: Implement proper row estimation logic based on query - estimated_rows = 50000 # Conservative estimate to trigger streaming - threshold = app.config.get("CSV_STREAMING_THRESHOLD", 10000) - - use_streaming = estimated_rows > threshold - - logger.info( - "Streaming decision: %d estimated rows, threshold %d, use_streaming=%s", - estimated_rows, - threshold, - use_streaming, - ) - - return use_streaming - - -class StreamingProgressTracker: - """ - Track progress of streaming exports for monitoring and user feedback. - - This could be extended to provide real-time progress updates via WebSocket - or SSE (Server-Sent Events) in the future. - """ - - def __init__(self, export_id: str) -> None: - self.export_id = export_id - self.start_time = datetime.now() - self.last_update = self.start_time - self.total_chunks = 0 - self.processed_rows = 0 - - def update_progress(self, chunks: int, rows: int) -> None: - """Update progress metrics.""" - self.total_chunks = chunks - self.processed_rows = rows - self.last_update = datetime.now() - - # Log progress every 100 chunks - if chunks % 100 == 0: - elapsed = (self.last_update - self.start_time).total_seconds() - rate = rows / elapsed if elapsed > 0 else 0 - logger.info( - "Export %s: %d rows, %d chunks, %.0f rows/sec", - self.export_id, - rows, - chunks, - rate, - ) - - def get_status(self) -> dict[str, Any]: - """Get current export status.""" - elapsed = (datetime.now() - self.start_time).total_seconds() - return { - "export_id": self.export_id, - "start_time": self.start_time.isoformat(), - "elapsed_seconds": elapsed, - "processed_rows": self.processed_rows, - "total_chunks": self.total_chunks, - "last_update": self.last_update.isoformat(), - } From a71746c34b97bb39f8319a4117d5fefa15851e5f Mon Sep 17 00:00:00 2001 From: amaannawab923 Date: Mon, 6 Oct 2025 11:34:46 +0530 Subject: [PATCH 13/57] update: Refactor streaming.py --- superset/views/streaming.py | 219 +++--------------------------------- 1 file changed, 13 insertions(+), 206 deletions(-) diff --git a/superset/views/streaming.py b/superset/views/streaming.py index ba2932c5b727..ca0c9ce0d1a0 100644 --- a/superset/views/streaming.py +++ b/superset/views/streaming.py @@ -21,12 +21,10 @@ import logging import time -import uuid from datetime import datetime from typing import Any, Generator, TYPE_CHECKING from flask import current_app as app, Response -from werkzeug.datastructures import Headers if TYPE_CHECKING: from superset.common.query_context import QueryContext @@ -63,40 +61,9 @@ def create_streaming_csv_response_simple( # Force chunked transfer encoding (critical for streaming) response.implicit_sequence_conversion = False - logger.info("Created simple streaming CSV response for file: %s", filename) return response -class StreamingExcelResponse(Response): - """ - Streaming response for Excel (XLSX) files. - - Note: Excel streaming is more complex than CSV due to the binary format. - This is a placeholder for future implementation. - """ - - def __init__( - self, - data_generator: Generator[bytes, None, None], - filename: str = "export.xlsx", - **kwargs: Any, - ) -> None: - headers = Headers() - headers.add( - "Content-Type", - "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet", - ) - headers.add("Content-Disposition", f'attachment; filename="{filename}"') - headers.add("Cache-Control", "no-cache, no-store, must-revalidate") - headers.add("Transfer-Encoding", "chunked") - - super().__init__( - response=data_generator, headers=headers, direct_passthrough=True, **kwargs - ) - - # Note: is_streamed is automatically True when response is a generator - - def create_streaming_csv_response( query_context: QueryContext, filename: str | None = None, @@ -105,161 +72,74 @@ def create_streaming_csv_response( expected_rows: int | None = None, ) -> Response: """ - Factory function to create a streaming CSV response using Flask's standard pattern. + Factory function to create a streaming CSV response. Args: query_context: Superset query context filename: Optional filename for download - chunk_size: Optional chunk size for processing - escape_formulas: Whether to escape formula injection + chunk_size: Optional chunk size for processing (default: 1000) + escape_formulas: Whether to escape formula injection (not implemented) + expected_rows: Expected number of rows for progress tracking Returns: Flask Response configured for streaming CSV """ - # Generate filename first if not provided if filename is None: timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") filename = f"superset_streaming_{timestamp}.csv" - # Use filename as export ID for progress tracking (frontend knows this) - export_id = filename + if chunk_size is None: + chunk_size = 1000 - # Capture the Flask app instance in the current context current_app = app._get_current_object() def csv_generator() -> Generator[str, None, None]: """Generator that yields CSV data from database query.""" - # Use the captured app instance to create application context with current_app.app_context(): - # Performance tracking start_time = time.time() total_bytes = 0 try: - logger.info("๐Ÿš€ STREAMING CSV: Starting streaming CSV generation") - logger.info("๐Ÿ“Š STREAMING CSV: Export ID: %s", export_id) - logger.info( - "๐Ÿ“Š STREAMING CSV: Processing query with estimated large result set" - ) - - - estimated_rows = expected_rows - - if estimated_rows: - logger.info( - "๐Ÿ“Š STREAMING CSV: Using expected_rows from frontend: %d", - estimated_rows, - ) - else: - form_data = query_context.form_data - estimated_rows = form_data.get('row_limit') if form_data else None - - if estimated_rows: - logger.info( - "๐Ÿ“Š STREAMING CSV: Using row_limit from form_data: %d", - estimated_rows, - ) - else: - logger.warning( - "โš ๏ธ STREAMING CSV: No expected_rows or row_limit available, progress percentage will be unavailable" - ) - estimated_rows = None - - # Get the database connection and execute raw SQL query directly from superset import db from superset.connectors.sqla.models import SqlaTable # Get the datasource datasource = query_context.datasource - # Create a fresh session to avoid detached object issues with db.session() as session: - # Refresh the datasource in the current session if isinstance(datasource, SqlaTable): datasource = session.merge(datasource) query_obj = query_context.queries[0] - sql_query = datasource.get_query_str(query_obj.to_dict()) - query_start_time = time.time() - logger.info( - "โšก STREAMING CSV: Executing SQL query: %s...", - sql_query[:200], - ) - logger.info( - "โฑ๏ธ STREAMING CSV: Query execution started at %s", - time.strftime("%H:%M:%S"), - ) - - # Execute query directly with the database engine - # Use context manager for proper connection handling with datasource.database.get_sqla_engine() as engine: - # Use server-side cursor for streaming connection = engine.connect() try: - # Execute query with server-side cursor from sqlalchemy import text result_proxy = connection.execution_options( stream_results=True ).execute(text(sql_query)) - # Get column names columns = list(result_proxy.keys()) - query_execution_time = time.time() - query_start_time - logger.info( - "๐Ÿ“‹ STREAMING CSV: Query columns (%d): %s", - len(columns), - columns, - ) - logger.info( - "โšก STREAMING CSV: Query executed in %.2fs, " - "starting data streaming...", - query_execution_time, - ) # Yield CSV header header_row = ",".join(f'"{col}"' for col in columns) + "\n" - header_bytes = len(header_row.encode("utf-8")) - total_bytes += header_bytes + total_bytes += len(header_row.encode("utf-8")) yield header_row - # ๐Ÿงช TESTING CONFIGURATION - Enable slower streaming for UI testing - ENABLE_SLOW_STREAMING_TEST = ( - True # Set to False for production speed - ) - - if ENABLE_SLOW_STREAMING_TEST: - # Testing mode: 10k rows per chunk with 0.5s delay - chunk_size = 10000 - delay_between_chunks = 0.5 - logger.info( - "๐Ÿงช TESTING MODE: Using 10k row chunks with 0.5s delays" - ) - else: - # Production mode: 1k rows per chunk, no delay - chunk_size = 1000 - delay_between_chunks = 0 - row_count = 0 - streaming_start_time = time.time() - last_progress_time = streaming_start_time - - # Buffer to accumulate data before yielding (forces flush) - # WSGI servers need ~50KB minimum to start streaming buffer = [] buffer_size = 0 - FLUSH_THRESHOLD = 65536 # 64KB - exceeds WSGI buffering threshold + FLUSH_THRESHOLD = 65536 # 64KB while True: - # Fetch chunk of rows rows = result_proxy.fetchmany(chunk_size) if not rows: break - # Build CSV rows and accumulate in buffer for row in rows: csv_row = ",".join( f'"{str(cell) if cell is not None else ""}"' @@ -273,92 +153,23 @@ def csv_generator() -> Generator[str, None, None]: buffer.append(csv_line) buffer_size += row_bytes - # Yield when buffer exceeds threshold (forces immediate flush) if buffer_size >= FLUSH_THRESHOLD: - chunk_data = "".join(buffer) - logger.info( - "๐Ÿ”ฅ YIELDING CHUNK: %d bytes at %s", - len(chunk_data), - time.strftime("%H:%M:%S") - ) - yield chunk_data + yield "".join(buffer) buffer = [] buffer_size = 0 - # Apply testing delay after yield for visible streaming - if ( - ENABLE_SLOW_STREAMING_TEST - and delay_between_chunks > 0 - ): - logger.info( - "โฑ๏ธ SLEEPING for %s seconds...", - delay_between_chunks - ) - time.sleep(delay_between_chunks) - logger.info("โฑ๏ธ WAKE UP - continuing...") - - # Performance logging every 10k rows or 5 seconds - current_time = time.time() - if ( - row_count % 10000 == 0 - or (current_time - last_progress_time) >= 5 - ): - elapsed = current_time - streaming_start_time - rows_per_sec = ( - row_count / elapsed if elapsed > 0 else 0 - ) - mb_streamed = total_bytes / (1024 * 1024) - mb_per_sec = ( - mb_streamed / elapsed if elapsed > 0 else 0 - ) - - logger.info( - "๐Ÿ“ˆ STREAMING CSV: %s rows streamed in %.1fs " - "(%.0f rows/s, %.1fMB, %.1fMB/s)", - f"{row_count:,}", - elapsed, - rows_per_sec, - mb_streamed, - mb_per_sec, - ) - last_progress_time = current_time - - # Apply testing delay for UI demonstration - if ( - ENABLE_SLOW_STREAMING_TEST - and delay_between_chunks > 0 - ): - time.sleep(delay_between_chunks) - # Flush remaining buffer if buffer: yield "".join(buffer) - buffer = [] - buffer_size = 0 - # Final performance summary + # Log completion total_time = time.time() - start_time - streaming_time = time.time() - streaming_start_time total_mb = total_bytes / (1024 * 1024) - logger.info( - "โœ… STREAMING CSV: Completed streaming %s rows", + "Streaming CSV completed: %s rows, %.1fMB in %.2fs", f"{row_count:,}", - ) - logger.info("๐Ÿ“Š STREAMING CSV PERFORMANCE:") - logger.info(" โ€ข Total Time: %.2fs", total_time) - logger.info(" โ€ข Query Time: %.2fs", query_execution_time) - logger.info(" โ€ข Streaming Time: %.2fs", streaming_time) - logger.info(" โ€ข Data Size: %.1fMB", total_mb) - logger.info( - " โ€ข Average Speed: %.0f rows/s, %.1fMB/s", - row_count / total_time, - total_mb / total_time, - ) - logger.info( - " โ€ข Memory Efficient: โœ… Constant memory usage " - "(~%d rows buffered)", - chunk_size, + total_mb, + total_time, ) finally: @@ -375,12 +186,8 @@ def csv_generator() -> Generator[str, None, None]: yield "error,message\n" yield f"CSV Export Error,{str(e)}\n" - # Get encoding from the captured app encoding = current_app.config.get("CSV_EXPORT", {}).get("encoding", "utf-8") - logger.info("Creating Flask streaming CSV response: %s", filename) - - # Use simple Flask Response with generator (official pattern) return create_streaming_csv_response_simple( data_generator=csv_generator(), filename=filename, From 869f535e21d78293d5fd20de611c0e7566113e96 Mon Sep 17 00:00:00 2001 From: amaannawab923 Date: Mon, 6 Oct 2025 12:04:14 +0530 Subject: [PATCH 14/57] update: Streaming export command --- superset/charts/data/api.py | 37 +++- .../chart/data/streaming_export_command.py | 169 +++++++++++++++ superset/views/streaming.py | 197 ------------------ 3 files changed, 201 insertions(+), 202 deletions(-) create mode 100644 superset/commands/chart/data/streaming_export_command.py delete mode 100644 superset/views/streaming.py diff --git a/superset/charts/data/api.py b/superset/charts/data/api.py index b58a70e5fd47..92791b64f001 100644 --- a/superset/charts/data/api.py +++ b/superset/charts/data/api.py @@ -534,7 +534,11 @@ def _create_streaming_csv_response( """Create a streaming CSV response for large datasets.""" from datetime import datetime - from superset.views.streaming import create_streaming_csv_response + from flask import Response + + from superset.commands.chart.data.streaming_export_command import ( + StreamingCSVExportCommand, + ) query_context = result["query_context"] @@ -558,8 +562,31 @@ def _create_streaming_csv_response( if expected_rows: logger.info("๐Ÿ“Š Using expected_rows from frontend: %d", expected_rows) - return create_streaming_csv_response( - query_context=query_context, - filename=filename, - expected_rows=expected_rows, + # Execute streaming command + chunk_size = 1000 + command = StreamingCSVExportCommand(query_context, chunk_size) + command.validate() + + # Get the callable that returns the generator + csv_generator_callable = command.run() + + # Get encoding from config + encoding = app.config.get("CSV_EXPORT", {}).get("encoding", "utf-8") + + # Create response with streaming headers + response = Response( + csv_generator_callable(), # Call the callable to get generator + mimetype=f"text/csv; charset={encoding}", + headers={ + "Content-Disposition": f'attachment; filename="{filename}"', + "Cache-Control": "no-cache", + "X-Accel-Buffering": "no", # Disable nginx buffering + "X-Superset-Streaming": "true", # Identify streaming responses + }, + direct_passthrough=False, # Flask must iterate generator ) + + # Force chunked transfer encoding + response.implicit_sequence_conversion = False + + return response diff --git a/superset/commands/chart/data/streaming_export_command.py b/superset/commands/chart/data/streaming_export_command.py new file mode 100644 index 000000000000..c3874920c017 --- /dev/null +++ b/superset/commands/chart/data/streaming_export_command.py @@ -0,0 +1,169 @@ +# 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 large datasets.""" + +from __future__ import annotations + +import logging +import time +from typing import Callable, Generator, TYPE_CHECKING + +from flask import current_app as app + +from superset.commands.base import BaseCommand + +if TYPE_CHECKING: + from superset.common.query_context import QueryContext + +logger = logging.getLogger(__name__) + + +class StreamingCSVExportCommand(BaseCommand): + """ + Command to execute a streaming CSV export. + + This command handles the business logic for: + - Executing database queries with server-side cursors + - Generating CSV data in chunks + - Managing database connections + - Buffering data for efficient streaming + """ + + def __init__( + self, + query_context: QueryContext, + chunk_size: int = 1000, + ): + """ + Initialize the streaming export command. + + Args: + query_context: The query context containing datasource and query details + chunk_size: Number of rows to fetch per database query (default: 1000) + """ + self._query_context = query_context + self._chunk_size = chunk_size + self._current_app = app._get_current_object() + + def validate(self) -> None: + """Validate permissions and query context.""" + self._query_context.raise_for_access() + + def run(self) -> Callable[[], Generator[str, None, None]]: + """ + Execute the streaming CSV export. + + Returns: + A callable that returns a generator yielding CSV data chunks as strings. + The callable is needed to maintain Flask app context during streaming. + """ + + def csv_generator() -> Generator[str, None, None]: + """Generator that yields CSV data from database query.""" + with self._current_app.app_context(): + start_time = time.time() + total_bytes = 0 + + try: + from superset import db + from superset.connectors.sqla.models import SqlaTable + + datasource = self._query_context.datasource + + with db.session() as session: + if isinstance(datasource, SqlaTable): + datasource = session.merge(datasource) + + query_obj = self._query_context.queries[0] + sql_query = datasource.get_query_str(query_obj.to_dict()) + + with datasource.database.get_sqla_engine() as engine: + connection = engine.connect() + + try: + from sqlalchemy import text + + result_proxy = connection.execution_options( + stream_results=True + ).execute(text(sql_query)) + + columns = list(result_proxy.keys()) + + # Yield CSV header + header_row = ( + ",".join(f'"{col}"' for col in columns) + "\n" + ) + total_bytes += len(header_row.encode("utf-8")) + yield header_row + + row_count = 0 + buffer = [] + buffer_size = 0 + FLUSH_THRESHOLD = 65536 # 64KB + + while True: + rows = result_proxy.fetchmany(self._chunk_size) + if not rows: + break + + for row in rows: + csv_row = ",".join( + f'"{str(cell) if cell is not None else ""}"' + for cell in row + ) + csv_line = csv_row + "\n" + row_bytes = len(csv_line.encode("utf-8")) + total_bytes += row_bytes + row_count += 1 + + buffer.append(csv_line) + buffer_size += row_bytes + + if buffer_size >= FLUSH_THRESHOLD: + yield "".join(buffer) + buffer = [] + buffer_size = 0 + + # Flush remaining buffer + if buffer: + yield "".join(buffer) + + # Log completion + total_time = time.time() - start_time + total_mb = total_bytes / (1024 * 1024) + logger.info( + "Streaming CSV completed: %s rows, %.1fMB in %.2fs", + f"{row_count:,}", + total_mb, + total_time, + ) + + finally: + connection.close() + + except Exception as e: + logger.error("Error in streaming CSV generator: %s", e) + import traceback + + logger.error("Traceback: %s", traceback.format_exc()) + + # Yield error info and fallback data + yield f"# Error occurred: {str(e)}\n" + yield "error,message\n" + yield f"CSV Export Error,{str(e)}\n" + + return csv_generator diff --git a/superset/views/streaming.py b/superset/views/streaming.py deleted file mode 100644 index ca0c9ce0d1a0..000000000000 --- a/superset/views/streaming.py +++ /dev/null @@ -1,197 +0,0 @@ -# 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. - -"""Streaming HTTP responses for large dataset exports.""" - -from __future__ import annotations - -import logging -import time -from datetime import datetime -from typing import Any, Generator, TYPE_CHECKING - -from flask import current_app as app, Response - -if TYPE_CHECKING: - from superset.common.query_context import QueryContext - -logger = logging.getLogger(__name__) - - -def create_streaming_csv_response_simple( - data_generator: Generator[str, None, None], - filename: str = "export.csv", - encoding: str = "utf-8", -) -> Response: - """ - Create a simple streaming CSV response using Flask's standard pattern. - - This follows the official Flask streaming documentation pattern. - """ - from flask import Response - - # Create response with proper headers to disable buffering - # CRITICAL: Set direct_passthrough=False to ensure Flask actually iterates the generator - response = Response( - data_generator, - mimetype=f"text/csv; charset={encoding}", - headers={ - "Content-Disposition": f'attachment; filename="{filename}"', - "Cache-Control": "no-cache", - "X-Accel-Buffering": "no", # Disable nginx buffering - "X-Superset-Streaming": "true", # Identify streaming responses - }, - direct_passthrough=False, # Flask must iterate generator, not pass to WSGI directly - ) - - # Force chunked transfer encoding (critical for streaming) - response.implicit_sequence_conversion = False - - return response - - -def create_streaming_csv_response( - query_context: QueryContext, - filename: str | None = None, - chunk_size: int | None = None, - escape_formulas: bool = True, - expected_rows: int | None = None, -) -> Response: - """ - Factory function to create a streaming CSV response. - - Args: - query_context: Superset query context - filename: Optional filename for download - chunk_size: Optional chunk size for processing (default: 1000) - escape_formulas: Whether to escape formula injection (not implemented) - expected_rows: Expected number of rows for progress tracking - - Returns: - Flask Response configured for streaming CSV - """ - if filename is None: - timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") - filename = f"superset_streaming_{timestamp}.csv" - - if chunk_size is None: - chunk_size = 1000 - - current_app = app._get_current_object() - - def csv_generator() -> Generator[str, None, None]: - """Generator that yields CSV data from database query.""" - with current_app.app_context(): - start_time = time.time() - total_bytes = 0 - - try: - from superset import db - from superset.connectors.sqla.models import SqlaTable - - # Get the datasource - datasource = query_context.datasource - - with db.session() as session: - if isinstance(datasource, SqlaTable): - datasource = session.merge(datasource) - - query_obj = query_context.queries[0] - sql_query = datasource.get_query_str(query_obj.to_dict()) - - with datasource.database.get_sqla_engine() as engine: - connection = engine.connect() - - try: - from sqlalchemy import text - - result_proxy = connection.execution_options( - stream_results=True - ).execute(text(sql_query)) - - columns = list(result_proxy.keys()) - - # Yield CSV header - header_row = ",".join(f'"{col}"' for col in columns) + "\n" - total_bytes += len(header_row.encode("utf-8")) - yield header_row - - row_count = 0 - buffer = [] - buffer_size = 0 - FLUSH_THRESHOLD = 65536 # 64KB - - while True: - rows = result_proxy.fetchmany(chunk_size) - if not rows: - break - - for row in rows: - csv_row = ",".join( - f'"{str(cell) if cell is not None else ""}"' - for cell in row - ) - csv_line = csv_row + "\n" - row_bytes = len(csv_line.encode("utf-8")) - total_bytes += row_bytes - row_count += 1 - - buffer.append(csv_line) - buffer_size += row_bytes - - if buffer_size >= FLUSH_THRESHOLD: - yield "".join(buffer) - buffer = [] - buffer_size = 0 - - # Flush remaining buffer - if buffer: - yield "".join(buffer) - - # Log completion - total_time = time.time() - start_time - total_mb = total_bytes / (1024 * 1024) - logger.info( - "Streaming CSV completed: %s rows, %.1fMB in %.2fs", - f"{row_count:,}", - total_mb, - total_time, - ) - - finally: - connection.close() - - except Exception as e: - logger.error("Error in streaming CSV generator: %s", e) - import traceback - - logger.error("Traceback: %s", traceback.format_exc()) - - # Yield error info and fallback data - yield f"# Error occurred: {str(e)}\n" - yield "error,message\n" - yield f"CSV Export Error,{str(e)}\n" - - encoding = current_app.config.get("CSV_EXPORT", {}).get("encoding", "utf-8") - - return create_streaming_csv_response_simple( - data_generator=csv_generator(), - filename=filename, - encoding=encoding, - ) - - From 843a550910502f2f8eed3dae926616b48bdfb675 Mon Sep 17 00:00:00 2001 From: amaannawab923 Date: Mon, 6 Oct 2025 12:07:39 +0530 Subject: [PATCH 15/57] fix: remove unneccessary logic from explore utils --- .../src/explore/exploreUtils/index.js | 46 ++----------------- 1 file changed, 4 insertions(+), 42 deletions(-) diff --git a/superset-frontend/src/explore/exploreUtils/index.js b/superset-frontend/src/explore/exploreUtils/index.js index 2c3d7a392bab..a242ab988fab 100644 --- a/superset-frontend/src/explore/exploreUtils/index.js +++ b/superset-frontend/src/explore/exploreUtils/index.js @@ -273,53 +273,15 @@ export const exportChart = async ({ }); } - // Check if this should use streaming export for CSV - const shouldUseStreaming = - resultFormat === 'csv' && onStartStreamingExport && !useLegacyApi; - - if (shouldUseStreaming) { - // Use streaming export instead of opening new tab - const timestamp = new Date() - .toISOString() - .slice(0, 19) - .replace(/[-:]/g, '') - .replace('T', '_'); - const chartName = formData.slice_name || formData.viz_type || 'chart'; - const safeChartName = chartName.replace(/[^a-zA-Z0-9_-]/g, '_'); - const filename = `superset_${safeChartName}_${timestamp}.csv`; - - // Extract expected row count for progress calculation - // Try to get row limit from form data for accurate progress tracking - let expectedRows; - if (formData.row_limit && formData.row_limit > 0) { - expectedRows = formData.row_limit; - } else if ( - payload.queries && - payload.queries[0] && - payload.queries[0].row_limit - ) { - expectedRows = payload.queries[0].row_limit; - } else { - // Default fallback - estimate based on common chart sizes - expectedRows = 10000; // Conservative default for progress calculation - } - - console.log( - '๐ŸŽฏ EXPORT CHART: Setting expectedRows =', - expectedRows, - 'from formData.row_limit =', - formData.row_limit, - ); - + // Check if streaming export handler is provided (only available from dashboard) + if (onStartStreamingExport) { + // Streaming is handled by the caller (Chart.jsx in dashboard) onStartStreamingExport({ url, payload, - filename, - exportType: 'csv', - expectedRows, // ๐ŸŽฏ This was missing! }); } else { - // Fallback to original behavior for non-streaming exports + // Fallback to original behavior for explore view SupersetClient.postForm(url, { form_data: safeStringify(payload) }); } }; From c3ea06bb3d3b57d5f5cf68e8320831e6056108ff Mon Sep 17 00:00:00 2001 From: amaannawab923 Date: Mon, 6 Oct 2025 12:31:23 +0530 Subject: [PATCH 16/57] update: Generate file name --- .../components/gridComponents/Chart/Chart.jsx | 16 +++++++++++++++- .../src/explore/exploreUtils/index.js | 6 +++--- 2 files changed, 18 insertions(+), 4 deletions(-) diff --git a/superset-frontend/src/dashboard/components/gridComponents/Chart/Chart.jsx b/superset-frontend/src/dashboard/components/gridComponents/Chart/Chart.jsx index f1bfcc426b5c..15408a27b1cf 100644 --- a/superset-frontend/src/dashboard/components/gridComponents/Chart/Chart.jsx +++ b/superset-frontend/src/dashboard/components/gridComponents/Chart/Chart.jsx @@ -422,6 +422,19 @@ const Chart = props => { !isPivot && actualRowCount >= streamingThreshold; + // Generate filename for streaming exports + let filename; + if (shouldUseStreaming) { + const timestamp = new Date() + .toISOString() + .slice(0, 19) + .replace(/[-:]/g, '') + .replace('T', '_'); + const chartName = formData.slice_name || formData.viz_type || 'chart'; + const safeChartName = chartName.replace(/[^a-zA-Z0-9_-]/g, '_'); + filename = `superset_${safeChartName}_${timestamp}.csv`; + } + exportChart({ formData: exportFormData, resultType, @@ -434,7 +447,8 @@ const Chart = props => { resetExport(); startExport({ ...exportParams, - expectedRows: actualRowCount || exportParams.expectedRows, + filename, + expectedRows: actualRowCount, }); } : null, diff --git a/superset-frontend/src/explore/exploreUtils/index.js b/superset-frontend/src/explore/exploreUtils/index.js index a242ab988fab..396e02d2b557 100644 --- a/superset-frontend/src/explore/exploreUtils/index.js +++ b/superset-frontend/src/explore/exploreUtils/index.js @@ -273,15 +273,15 @@ export const exportChart = async ({ }); } - // Check if streaming export handler is provided (only available from dashboard) + // Check if streaming export handler is provided (from dashboard Chart.jsx) if (onStartStreamingExport) { - // Streaming is handled by the caller (Chart.jsx in dashboard) + // Streaming is handled by the caller - just pass URL and payload onStartStreamingExport({ url, payload, }); } else { - // Fallback to original behavior for explore view + // Fallback to original behavior for non-streaming exports SupersetClient.postForm(url, { form_data: safeStringify(payload) }); } }; From f320f424e28ec42d94f363f82feebb2f22b3e652 Mon Sep 17 00:00:00 2001 From: amaannawab923 Date: Tue, 7 Oct 2025 19:01:32 +0530 Subject: [PATCH 17/57] fix: ci cd fixes --- .../StreamingExportModal.tsx | 389 +++++------------- .../components/StreamingExportModal/index.ts | 2 +- .../useStreamingExport.ts | 97 +++-- .../components/gridComponents/Chart/Chart.jsx | 12 +- superset-frontend/webpack.proxy-config.js | 3 +- superset/charts/data/api.py | 57 ++- .../chart/data/streaming_export_command.py | 4 +- superset/config.py | 5 +- 8 files changed, 207 insertions(+), 362 deletions(-) diff --git a/superset-frontend/src/components/StreamingExportModal/StreamingExportModal.tsx b/superset-frontend/src/components/StreamingExportModal/StreamingExportModal.tsx index b22032a38c37..a084ae8b74f7 100644 --- a/superset-frontend/src/components/StreamingExportModal/StreamingExportModal.tsx +++ b/superset-frontend/src/components/StreamingExportModal/StreamingExportModal.tsx @@ -16,15 +16,8 @@ * specific language governing permissions and limitations * under the License. */ -import React from 'react'; import { styled, t } from '@superset-ui/core'; -import { Modal, Button, Typography, Progress, theme } from 'antd'; -import { - CheckCircleOutlined, - CloseCircleOutlined, - LoadingOutlined, - StopOutlined, -} from '@ant-design/icons'; +import { Modal, Button, Typography, Progress } from 'antd'; const { Text } = Typography; @@ -36,17 +29,16 @@ export enum ExportStatus { } export interface StreamingProgress { - rowsProcessed: number; totalRows?: number; + rowsProcessed: number; totalSize: number; - speed: number; - mbPerSecond: number; - elapsedTime: number; - estimatedTimeRemaining?: number; status: ExportStatus; downloadUrl?: string; - filename?: string; error?: string; + filename?: string; + speed?: number; + mbPerSecond?: number; + elapsedTime?: number; } interface StreamingExportModalProps { @@ -58,7 +50,8 @@ interface StreamingExportModalProps { } const ModalContent = styled.div` - padding: ${({ theme }) => theme.sizeUnit * 4}px 0 ${({ theme }) => theme.sizeUnit * 2}px; + padding: ${({ theme }) => theme.sizeUnit * 4}px 0 + ${({ theme }) => theme.sizeUnit * 2}px; `; const ProgressSection = styled.div` @@ -67,268 +60,51 @@ const ProgressSection = styled.div` const ActionButtons = styled.div` display: flex; + gap: ${({ theme }) => theme.sizeUnit * 2}px; justify-content: flex-end; - gap: ${({ theme }) => theme.sizeUnit * 3}px; - margin-top: ${({ theme }) => theme.sizeUnit * 6}px; -`; - -const StatusIcon = styled.div` - display: flex; - justify-content: center; - margin-bottom: ${({ theme }) => theme.sizeUnit * 4}px; - - .anticon { - font-size: 48px; - } -`; - -const StyledIcon = styled.span<{ color: string }>` - color: ${({ color }) => color}; `; -const CenteredText = styled(Text)` +const ProgressText = styled(Text)` display: block; text-align: center; - margin-bottom: ${({ theme }) => theme.sizeUnit * 4}px; + margin-top: ${({ theme }) => theme.sizeUnit * 4}px; `; -const ProgressText = styled(Text)` +const ErrorText = styled(Text)` display: block; text-align: center; margin-top: ${({ theme }) => theme.sizeUnit * 4}px; `; -const CancelButton = styled(Button)` - background-color: #f0fff8; - color: ${({ theme }) => theme.colorSuccess}; - border-color: transparent; - - &:hover { - background-color: #d9f7e8; - color: ${({ theme }) => theme.colorSuccess}; - border-color: transparent; - } -`; - -const DownloadButton = styled(Button)` - background-color: #2fc096; - color: white; - border-color: #2fc096; +const CancelButton = styled(Button)``; - &:hover { - background-color: #26a77e; - border-color: #26a77e; - } +const DownloadButton = styled(Button)``; - &:disabled { - background-color: #f5f5f5; - color: rgba(0, 0, 0, 0.25); - border-color: #d9d9d9; - } -`; - -const formatFileSize = (bytes: number): string => { - if (bytes === 0) return '0 B'; - const k = 1024; - const sizes = ['B', 'KB', 'MB', 'GB']; - const i = Math.floor(Math.log(bytes) / Math.log(k)); - return `${parseFloat((bytes / Math.pow(k, i)).toFixed(1))} ${sizes[i]}`; -}; - -const formatNumber = (num: number): string => - new Intl.NumberFormat().format(num); - -interface StatusIconProps { - status: ExportStatus; -} - -const StatusIconComponent: React.FC = ({ status }) => { - const { token } = theme.useToken(); - - const iconMap = { - [ExportStatus.STREAMING]: { Icon: LoadingOutlined, color: token.colorPrimary }, - [ExportStatus.COMPLETED]: { Icon: CheckCircleOutlined, color: token.colorSuccess }, - [ExportStatus.ERROR]: { Icon: CloseCircleOutlined, color: token.colorError }, - [ExportStatus.CANCELLED]: { Icon: StopOutlined, color: token.colorWarning }, - }; - - const { Icon, color } = iconMap[status] || iconMap[ExportStatus.STREAMING]; - return ( - - - - - - ); -}; - -interface ErrorContentProps { - error?: string; - onRetry?: () => void; - onCancel: () => void; -} - -const ErrorContent: React.FC = ({ - error, - onRetry, - onCancel, -}) => ( - - - - {error || t('An error occurred during export')} - - - {onRetry && ( - - )} - - - -); - -interface CancelledContentProps { - onRetry?: () => void; - onCancel: () => void; -} - -const CancelledContent: React.FC = ({ - onRetry, - onCancel, -}) => ( - - - {t('Export was cancelled')} - - {onRetry && ( - - )} - - - -); - -interface CompletedContentProps { - totalRows?: number; - rowsProcessed: number; - totalSize: number; - filename?: string; - downloadUrl?: string; - onDownload: () => void; - onCancel: () => void; -} - -const CompletedContent: React.FC = ({ - totalRows, - rowsProcessed, - totalSize, - filename, - downloadUrl, - onDownload, - onCancel, -}) => { - const { token } = theme.useToken(); - - return ( - - - - - {t('Export successful %s', filename || 'export.csv')} - - - - - - {t('Cancel')} - - - {t('Download')} - - - - ); -}; - -interface StreamingContentProps { - percentage: number; - filename?: string; - onCancel: () => void; -} - -const StreamingContent: React.FC = ({ - percentage, - filename, - onCancel, -}) => { - const { token } = theme.useToken(); - - return ( - - - `${Math.round(percent || 0)}%`} - /> - - {filename - ? t('Processing export for %s', filename) - : t('Processing export for {dashboard_name}_{YYYY-MM-DD}_{HHMMSS}.csv')} - - - - - - {t('Cancel')} - - - {t('Download')} - - - - ); -}; - -const StreamingExportModal: React.FC = ({ +const StreamingExportModal = ({ visible, onCancel, onRetry, progress, -}) => { - const { totalRows, rowsProcessed, totalSize, status, downloadUrl, filename, error } = - progress; +}: StreamingExportModalProps) => { + const { status, downloadUrl, filename, error } = progress; const getProgressPercentage = (): number => { if (status === ExportStatus.COMPLETED) return 100; - - if (totalRows && totalRows > 0) { - const percentage = Math.min(99, (rowsProcessed / totalRows) * 100); + if (progress.totalRows && progress.totalRows > 0) { + const percentage = Math.min( + 99, + (progress.rowsProcessed / progress.totalRows) * 100, + ); return Math.round(percentage); } - - return status === ExportStatus.STREAMING ? 10 : 0; + return 0; }; const handleDownload = () => { - if (downloadUrl && filename) { + if (downloadUrl) { const link = document.createElement('a'); link.href = downloadUrl; - link.download = filename; + link.download = filename || 'export.csv'; document.body.appendChild(link); link.click(); document.body.removeChild(link); @@ -337,46 +113,99 @@ const StreamingExportModal: React.FC = ({ }; let content; - switch (status) { - case ExportStatus.ERROR: - content = ; - break; - case ExportStatus.CANCELLED: - content = ; - break; - case ExportStatus.COMPLETED: - content = ( - - ); - break; - default: - content = ( - - ); + if (status === ExportStatus.ERROR) { + content = ( + + + + {error || t('Export failed')} + + + {t('Close')} + {onRetry && ( + + {t('Retry')} + + )} + + + ); + } else if (status === ExportStatus.CANCELLED) { + content = ( + + + + {t('Export cancelled')} + + + {t('Close')} + {onRetry && ( + + {t('Retry')} + + )} + + + ); + } else if (status === ExportStatus.COMPLETED) { + content = ( + + + + + {t('Export successful: %s', filename || 'export.csv')} + + + + {t('Close')} + + {t('Download')} + + + + ); + } else { + content = ( + + + `${Math.round(percent || 0)}%`} + /> + + {filename + ? t('Processing export for %s', filename) + : t('Processing export...')} + + + + {t('Cancel')} + + {t('Download')} + + + + ); } return ( {content} diff --git a/superset-frontend/src/components/StreamingExportModal/index.ts b/superset-frontend/src/components/StreamingExportModal/index.ts index 049a4d2af0f1..7c342b4f298e 100644 --- a/superset-frontend/src/components/StreamingExportModal/index.ts +++ b/superset-frontend/src/components/StreamingExportModal/index.ts @@ -18,4 +18,4 @@ */ export { default as StreamingExportModal } from './StreamingExportModal'; export type { StreamingProgress } from './StreamingExportModal'; -export { useStreamingExport } from './useStreamingExport'; \ No newline at end of file +export { useStreamingExport } from './useStreamingExport'; diff --git a/superset-frontend/src/components/StreamingExportModal/useStreamingExport.ts b/superset-frontend/src/components/StreamingExportModal/useStreamingExport.ts index 9c87afaeba20..59089a2c252c 100644 --- a/superset-frontend/src/components/StreamingExportModal/useStreamingExport.ts +++ b/superset-frontend/src/components/StreamingExportModal/useStreamingExport.ts @@ -45,12 +45,9 @@ export const useStreamingExport = (options: UseStreamingExportOptions = {}) => { const [isExporting, setIsExporting] = useState(false); const abortControllerRef = useRef(null); - const updateProgress = useCallback( - (updates: Partial) => { - setProgress(prev => ({ ...prev, ...updates })); - }, - [], - ); + const updateProgress = useCallback((updates: Partial) => { + setProgress(prev => ({ ...prev, ...updates })); + }, []); const startExport = useCallback( async ({ @@ -107,62 +104,60 @@ export const useStreamingExport = (options: UseStreamingExportOptions = {}) => { let rowsProcessed = 0; const NEWLINE_BYTE = 10; // '\n' character code - try { - while (true) { - const { done, value } = await reader.read(); + // eslint-disable-next-line no-constant-condition + while (true) { + // eslint-disable-next-line no-await-in-loop + const { done, value } = await reader.read(); - if (done) break; + if (done) break; - if (abortControllerRef.current?.signal.aborted) { - throw new Error('Export cancelled by user'); - } + if (abortControllerRef.current?.signal.aborted) { + throw new Error('Export cancelled by user'); + } - chunks.push(value); - receivedLength += value.length; + chunks.push(value); + receivedLength += value.length; - // Count newlines directly in binary (faster than decoding + regex) - let newlineCount = 0; - for (let i = 0; i < value.length; i++) { - if (value[i] === NEWLINE_BYTE) { - newlineCount++; - } + // Count newlines directly in binary (faster than decoding + regex) + let newlineCount = 0; + for (let i = 0; i < value.length; i += 1) { + if (value[i] === NEWLINE_BYTE) { + newlineCount += 1; } - rowsProcessed += newlineCount; - - // Update progress based on rows processed - updateProgress({ - status: ExportStatus.STREAMING, - rowsProcessed, - totalRows: expectedRows, - totalSize: receivedLength, - }); - } - - const completeData = new Uint8Array(receivedLength); - let position = 0; - for (const chunk of chunks) { - completeData.set(chunk, position); - position += chunk.length; } + rowsProcessed += newlineCount; - const mimeType = - exportType === 'csv' - ? 'text/csv;charset=utf-8' - : 'application/vnd.openxmlformats-officedocument.spreadsheetml.sheet'; - - const blob = new Blob([completeData], { type: mimeType }); - const downloadUrl = URL.createObjectURL(blob); - + // Update progress based on rows processed updateProgress({ - status: ExportStatus.COMPLETED, - downloadUrl, - filename: filename || `export.${exportType}`, + status: ExportStatus.STREAMING, + rowsProcessed, + totalRows: expectedRows, + totalSize: receivedLength, }); + } - options.onComplete?.(downloadUrl, filename || `export.${exportType}`); - } catch (streamError) { - throw streamError; + const completeData = new Uint8Array(receivedLength); + let position = 0; + for (const chunk of chunks) { + completeData.set(chunk, position); + position += chunk.length; } + + const mimeType = + exportType === 'csv' + ? 'text/csv;charset=utf-8' + : 'application/vnd.openxmlformats-officedocument.spreadsheetml.sheet'; + + const blob = new Blob([completeData], { type: mimeType }); + const downloadUrl = URL.createObjectURL(blob); + + updateProgress({ + status: ExportStatus.COMPLETED, + downloadUrl, + filename: filename || `export.${exportType}`, + }); + + options.onComplete?.(downloadUrl, filename || `export.${exportType}`); } catch (error) { const errorMessage = error instanceof Error ? error.message : 'Unknown error occurred'; diff --git a/superset-frontend/src/dashboard/components/gridComponents/Chart/Chart.jsx b/superset-frontend/src/dashboard/components/gridComponents/Chart/Chart.jsx index 15408a27b1cf..21e3d809246e 100644 --- a/superset-frontend/src/dashboard/components/gridComponents/Chart/Chart.jsx +++ b/superset-frontend/src/dashboard/components/gridComponents/Chart/Chart.jsx @@ -80,7 +80,6 @@ const propTypes = { isInView: PropTypes.bool, }; - const RESIZE_TIMEOUT = 500; const DEFAULT_HEADER_HEIGHT = 22; @@ -161,7 +160,8 @@ const Chart = props => { state => state.dashboardInfo.common.conf.SQL_MAX_ROW, ); const streamingThreshold = useSelector( - state => state.dashboardInfo.common.conf.CSV_STREAMING_ROW_THRESHOLD || 100000, + state => + state.dashboardInfo.common.conf.CSV_STREAMING_ROW_THRESHOLD || 100000, ); const datasource = useSelector( state => @@ -191,7 +191,9 @@ const Chart = props => { ); }, onError: () => { - boundActionCreators.addDangerToast(t('Export failed - please try again')); + boundActionCreators.addDangerToast( + t('Export failed - please try again'), + ); }, }); const history = useHistory(); @@ -418,9 +420,7 @@ const Chart = props => { // Handle streaming CSV exports based on row threshold const shouldUseStreaming = - format === 'csv' && - !isPivot && - actualRowCount >= streamingThreshold; + format === 'csv' && !isPivot && actualRowCount >= streamingThreshold; // Generate filename for streaming exports let filename; diff --git a/superset-frontend/webpack.proxy-config.js b/superset-frontend/webpack.proxy-config.js index 47b616ed7cd6..8fe035a1bb90 100644 --- a/superset-frontend/webpack.proxy-config.js +++ b/superset-frontend/webpack.proxy-config.js @@ -165,7 +165,8 @@ module.exports = newManifest => { if (isHTML(response)) { processHTML(proxyResponse, response); } else { - const isStreaming = proxyResponse.headers['x-superset-streaming'] === 'true'; + const isStreaming = + proxyResponse.headers['x-superset-streaming'] === 'true'; if (isStreaming) { proxyResponse.on('data', chunk => { diff --git a/superset/charts/data/api.py b/superset/charts/data/api.py index 92791b64f001..9df52b6f491d 100644 --- a/superset/charts/data/api.py +++ b/superset/charts/data/api.py @@ -257,23 +257,14 @@ def data(self) -> Response: return self._run_async(json_body, command) form_data = json_body.get("form_data") - - # Extract filename from request if provided (for streaming CSV) - filename = request.form.get("filename") - if filename: - logger.info("๐Ÿ“ FRONTEND PROVIDED FILENAME: %s", filename) - - expected_rows = request.form.get("expected_rows") - if expected_rows: - try: - expected_rows = int(expected_rows) - logger.info("๐Ÿ“Š FRONTEND PROVIDED EXPECTED ROWS: %d", expected_rows) - except (ValueError, TypeError): - logger.warning("โš ๏ธ Invalid expected_rows value: %s", expected_rows) - expected_rows = None + filename, expected_rows = self._extract_export_params_from_request() return self._get_data_response( - command, form_data=form_data, datasource=query_context.datasource, filename=filename, expected_rows=expected_rows + command, + form_data=form_data, + datasource=query_context.datasource, + filename=filename, + expected_rows=expected_rows, ) @expose("/data/", methods=("GET",)) @@ -387,7 +378,9 @@ def _send_chart_response( # noqa: C901 # Check if we should use streaming for large datasets if is_csv_format and self._should_use_streaming(result, form_data): - return self._create_streaming_csv_response(result, form_data, filename=filename, expected_rows=expected_rows) + return self._create_streaming_csv_response( + result, form_data, filename=filename, expected_rows=expected_rows + ) if len(result["queries"]) == 1: # return single query results @@ -448,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, filename, expected_rows) + 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) + except (ValueError, TypeError): + logger.warning("โš ๏ธ Invalid expected_rows value: %s", expected_rows_str) + + return filename, expected_rows # pylint: disable=invalid-name def _load_query_context_form_from_cache(self, cache_key: str) -> dict[str, Any]: @@ -529,7 +540,11 @@ def _should_use_streaming( return False def _create_streaming_csv_response( - self, result: dict[Any, Any], form_data: dict[str, Any] | None = None, filename: str | None = None, expected_rows: int | None = None + self, + result: dict[Any, Any], + form_data: dict[str, Any] | None = None, + filename: str | None = None, + expected_rows: int | None = None, ) -> Response: """Create a streaming CSV response for large datasets.""" from datetime import datetime @@ -558,7 +573,11 @@ def _create_streaming_csv_response( ) filename = f"superset_{safe_chart_name}_{timestamp}.csv" - logger.info("Creating streaming CSV response: %s (from frontend: %s)", filename, filename is not None) + logger.info( + "Creating streaming CSV response: %s (from frontend: %s)", + filename, + filename is not None, + ) if expected_rows: logger.info("๐Ÿ“Š Using expected_rows from frontend: %d", expected_rows) diff --git a/superset/commands/chart/data/streaming_export_command.py b/superset/commands/chart/data/streaming_export_command.py index c3874920c017..60493b597f7e 100644 --- a/superset/commands/chart/data/streaming_export_command.py +++ b/superset/commands/chart/data/streaming_export_command.py @@ -113,7 +113,7 @@ def csv_generator() -> Generator[str, None, None]: row_count = 0 buffer = [] buffer_size = 0 - FLUSH_THRESHOLD = 65536 # 64KB + flush_threshold = 65536 # 64KB while True: rows = result_proxy.fetchmany(self._chunk_size) @@ -133,7 +133,7 @@ def csv_generator() -> Generator[str, None, None]: buffer.append(csv_line) buffer_size += row_bytes - if buffer_size >= FLUSH_THRESHOLD: + if buffer_size >= flush_threshold: yield "".join(buffer) buffer = [] buffer_size = 0 diff --git a/superset/config.py b/superset/config.py index 947720ea028d..33bd243030bb 100644 --- a/superset/config.py +++ b/superset/config.py @@ -1010,8 +1010,9 @@ class D3TimeFormat(TypedDict, total=False): CSV_EXPORT = {"encoding": "utf-8-sig"} # CSV Streaming: row threshold for using streaming CSV exports -# When row count >= this threshold, use streaming response instead of loading all data -# into memory. Streaming provides real-time progress and handles large datasets efficiently. +# When row count >= this threshold, use streaming response instead of loading +# all data into memory. Streaming provides real-time progress and handles +# large datasets efficiently. CSV_STREAMING_ROW_THRESHOLD = 100000 # Excel Options: key/value pairs that will be passed as argument to DataFrame.to_excel From 130f9bacba1c52bc09c73757be9bbb7f0d35afe8 Mon Sep 17 00:00:00 2001 From: amaannawab923 Date: Tue, 7 Oct 2025 20:14:32 +0530 Subject: [PATCH 18/57] fix : ui fixes --- .../StreamingExportModal.tsx | 73 ++++++++++++++++++- 1 file changed, 69 insertions(+), 4 deletions(-) diff --git a/superset-frontend/src/components/StreamingExportModal/StreamingExportModal.tsx b/superset-frontend/src/components/StreamingExportModal/StreamingExportModal.tsx index a084ae8b74f7..8abb91dff418 100644 --- a/superset-frontend/src/components/StreamingExportModal/StreamingExportModal.tsx +++ b/superset-frontend/src/components/StreamingExportModal/StreamingExportModal.tsx @@ -16,8 +16,10 @@ * specific language governing permissions and limitations * under the License. */ +/* eslint-disable theme-colors/no-literal-colors */ import { styled, t } from '@superset-ui/core'; import { Modal, Button, Typography, Progress } from 'antd'; +import { Icons } from '@superset-ui/core/components/Icons'; const { Text } = Typography; @@ -56,6 +58,19 @@ const ModalContent = styled.div` const ProgressSection = styled.div` margin: ${({ theme }) => theme.sizeUnit * 6}px 0; + position: relative; +`; + +const ProgressWrapper = styled.div` + display: flex; + align-items: center; + gap: ${({ theme }) => theme.sizeUnit * 3}px; +`; + +const SuccessIcon = styled(Icons.CheckCircleFilled)` + color: #52c41a; + font-size: 24px; + flex-shrink: 0; `; const ActionButtons = styled.div` @@ -76,9 +91,49 @@ const ErrorText = styled(Text)` margin-top: ${({ theme }) => theme.sizeUnit * 4}px; `; -const CancelButton = styled(Button)``; +const CancelButton = styled(Button)` + background-color: #f0fff8; + color: #1c997a; + border-color: #f0fff8; + + &:hover { + background-color: #f0fff8; + color: #1c997a; + border-color: #1c997a; + } + + &:focus { + background-color: #f0fff8; + color: #1c997a; + border-color: #1c997a; + } +`; + +const DownloadButton = styled(Button)` + &.ant-btn-primary { + background-color: #2ec196; + border-color: #2ec196; + color: #ffffff; -const DownloadButton = styled(Button)``; + &:hover:not(:disabled) { + background-color: #26a880; + border-color: #26a880; + color: #ffffff; + } + + &:focus:not(:disabled) { + background-color: #2ec196; + border-color: #2ec196; + color: #ffffff; + } + + &:disabled { + background-color: #f2f2f2; + border-color: #f2f2f2; + color: #b5b5b5; + } + } +`; const StreamingExportModal = ({ visible, @@ -155,7 +210,15 @@ const StreamingExportModal = ({ content = ( - + + + + {t('Export successful: %s', filename || 'export.csv')} @@ -178,7 +241,8 @@ const StreamingExportModal = ({ `${Math.round(percent || 0)}%`} /> @@ -206,6 +270,7 @@ const StreamingExportModal = ({ footer={null} width={600} maskClosable={false} + centered > {content} From ec5ec0e3fa8f553aa2ce0eedf4f6e04701182a7d Mon Sep 17 00:00:00 2001 From: amaannawab923 Date: Tue, 7 Oct 2025 21:29:58 +0530 Subject: [PATCH 19/57] update: resolve comments --- .../StreamingExportModal.tsx | 287 ++++++++++++------ .../useStreamingExport.ts | 110 ++++--- .../chart/data/streaming_export_command.py | 51 ++-- 3 files changed, 289 insertions(+), 159 deletions(-) diff --git a/superset-frontend/src/components/StreamingExportModal/StreamingExportModal.tsx b/superset-frontend/src/components/StreamingExportModal/StreamingExportModal.tsx index 8abb91dff418..199e2c14b2ec 100644 --- a/superset-frontend/src/components/StreamingExportModal/StreamingExportModal.tsx +++ b/superset-frontend/src/components/StreamingExportModal/StreamingExportModal.tsx @@ -48,7 +48,6 @@ interface StreamingExportModalProps { onCancel: () => void; onRetry?: () => void; progress: StreamingProgress; - exportType: 'csv' | 'xlsx'; } const ModalContent = styled.div` @@ -135,13 +134,189 @@ const DownloadButton = styled(Button)` } `; +interface ModalStateContentProps { + status: ExportStatus; + progress: StreamingProgress; + onCancel: () => void; + onRetry?: () => void; + onDownload: () => void; + getProgressPercentage: () => number; +} + +const ErrorContent = ({ + error, + onCancel, + onRetry, +}: { + error?: string; + onCancel: () => void; + onRetry?: () => void; +}) => ( + + + + {error || t('Export failed')} + + + {t('Close')} + {onRetry && ( + + {t('Retry')} + + )} + + +); + +const CancelledContent = ({ + getProgressPercentage, + onCancel, + onRetry, +}: { + getProgressPercentage: () => number; + onCancel: () => void; + onRetry?: () => void; +}) => ( + + + + {t('Export cancelled')} + + + {t('Close')} + {onRetry && ( + + {t('Retry')} + + )} + + +); + +const CompletedContent = ({ + filename, + downloadUrl, + onCancel, + onDownload, +}: { + filename?: string; + downloadUrl?: string; + onCancel: () => void; + onDownload: () => void; +}) => ( + + + + + + + + {t('Export successful: %s', filename || 'export')} + + + + {t('Close')} + + {t('Download')} + + + +); + +const StreamingContent = ({ + filename, + getProgressPercentage, + onCancel, +}: { + filename?: string; + getProgressPercentage: () => number; + onCancel: () => void; +}) => ( + + + `${Math.round(percent || 0)}%`} + /> + + {filename + ? t('Processing export for %s', filename) + : t('Processing export...')} + + + + {t('Cancel')} + + {t('Download')} + + + +); + +const ModalStateContent = ({ + status, + progress, + onCancel, + onRetry, + onDownload, + getProgressPercentage, +}: ModalStateContentProps) => { + const { downloadUrl, filename, error } = progress; + + switch (status) { + case ExportStatus.ERROR: + return ; + case ExportStatus.CANCELLED: + return ( + + ); + case ExportStatus.COMPLETED: + return ( + + ); + default: + return ( + + ); + } +}; + const StreamingExportModal = ({ visible, onCancel, onRetry, progress, }: StreamingExportModalProps) => { - const { status, downloadUrl, filename, error } = progress; + const { status, downloadUrl, filename } = progress; const getProgressPercentage = (): number => { if (status === ExportStatus.COMPLETED) return 100; @@ -156,10 +331,10 @@ const StreamingExportModal = ({ }; const handleDownload = () => { - if (downloadUrl) { + if (downloadUrl && filename) { const link = document.createElement('a'); link.href = downloadUrl; - link.download = filename || 'export.csv'; + link.download = filename; document.body.appendChild(link); link.click(); document.body.removeChild(link); @@ -167,101 +342,6 @@ const StreamingExportModal = ({ } }; - let content; - if (status === ExportStatus.ERROR) { - content = ( - - - - {error || t('Export failed')} - - - {t('Close')} - {onRetry && ( - - {t('Retry')} - - )} - - - ); - } else if (status === ExportStatus.CANCELLED) { - content = ( - - - - {t('Export cancelled')} - - - {t('Close')} - {onRetry && ( - - {t('Retry')} - - )} - - - ); - } else if (status === ExportStatus.COMPLETED) { - content = ( - - - - - - - - {t('Export successful: %s', filename || 'export.csv')} - - - - {t('Close')} - - {t('Download')} - - - - ); - } else { - content = ( - - - `${Math.round(percent || 0)}%`} - /> - - {filename - ? t('Processing export for %s', filename) - : t('Processing export...')} - - - - {t('Cancel')} - - {t('Download')} - - - - ); - } - return ( - {content} + ); }; diff --git a/superset-frontend/src/components/StreamingExportModal/useStreamingExport.ts b/superset-frontend/src/components/StreamingExportModal/useStreamingExport.ts index 59089a2c252c..31716878acbf 100644 --- a/superset-frontend/src/components/StreamingExportModal/useStreamingExport.ts +++ b/superset-frontend/src/components/StreamingExportModal/useStreamingExport.ts @@ -24,14 +24,64 @@ interface UseStreamingExportOptions { onError?: (error: string) => void; } +interface StreamingExportPayload { + [key: string]: unknown; +} + interface StreamingExportParams { url: string; - payload: any; + payload: StreamingExportPayload; filename?: string; exportType: 'csv' | 'xlsx'; expectedRows?: number; } +const NEWLINE_BYTE = 10; // '\n' character code + +const createFetchRequest = ( + url: string, + payload: StreamingExportPayload, + filename: string, + exportType: string, + expectedRows: number | undefined, + signal: AbortSignal, +): RequestInit => ({ + method: 'POST', + headers: { + 'Content-Type': 'application/x-www-form-urlencoded', + }, + body: new URLSearchParams({ + form_data: JSON.stringify(payload), + filename, + expected_rows: expectedRows?.toString() || '', + }), + signal, + credentials: 'same-origin', +}); + +const countNewlines = (value: Uint8Array): number => + value.filter(byte => byte === NEWLINE_BYTE).length; + +const createBlob = ( + chunks: Uint8Array[], + receivedLength: number, + exportType: string, +): Blob => { + const completeData = new Uint8Array(receivedLength); + let position = 0; + for (const chunk of chunks) { + completeData.set(chunk, position); + position += chunk.length; + } + + const mimeType = + exportType === 'csv' + ? 'text/csv;charset=utf-8' + : 'application/vnd.openxmlformats-officedocument.spreadsheetml.sheet'; + + return new Blob([completeData], { type: mimeType }); +}; + export const useStreamingExport = (options: UseStreamingExportOptions = {}) => { const [progress, setProgress] = useState({ rowsProcessed: 0, @@ -74,19 +124,20 @@ export const useStreamingExport = (options: UseStreamingExportOptions = {}) => { }); try { - const response = await fetch(url, { - method: 'POST', - headers: { - 'Content-Type': 'application/x-www-form-urlencoded', - }, - body: new URLSearchParams({ - form_data: JSON.stringify(payload), - filename: filename || `export.${exportType}`, - expected_rows: expectedRows?.toString() || '', - }), - signal: abortControllerRef.current.signal, - credentials: 'same-origin', - }); + const defaultFilename = `export.${exportType}`; + const finalFilename = filename || defaultFilename; + + const response = await fetch( + url, + createFetchRequest( + url, + payload, + finalFilename, + exportType, + expectedRows, + abortControllerRef.current.signal, + ), + ); if (!response.ok) { throw new Error( @@ -102,7 +153,6 @@ export const useStreamingExport = (options: UseStreamingExportOptions = {}) => { const chunks: Uint8Array[] = []; let receivedLength = 0; let rowsProcessed = 0; - const NEWLINE_BYTE = 10; // '\n' character code // eslint-disable-next-line no-constant-condition while (true) { @@ -118,14 +168,10 @@ export const useStreamingExport = (options: UseStreamingExportOptions = {}) => { chunks.push(value); receivedLength += value.length; - // Count newlines directly in binary (faster than decoding + regex) - let newlineCount = 0; - for (let i = 0; i < value.length; i += 1) { - if (value[i] === NEWLINE_BYTE) { - newlineCount += 1; - } - } - rowsProcessed += newlineCount; + // Count newlines using filter (more efficient than loop) + // Note: This counts all newlines, including those within quoted CSV fields. + // For an exact row count, server should send row count in response headers. + rowsProcessed += countNewlines(value); // Update progress based on rows processed updateProgress({ @@ -136,28 +182,16 @@ export const useStreamingExport = (options: UseStreamingExportOptions = {}) => { }); } - const completeData = new Uint8Array(receivedLength); - let position = 0; - for (const chunk of chunks) { - completeData.set(chunk, position); - position += chunk.length; - } - - const mimeType = - exportType === 'csv' - ? 'text/csv;charset=utf-8' - : 'application/vnd.openxmlformats-officedocument.spreadsheetml.sheet'; - - const blob = new Blob([completeData], { type: mimeType }); + const blob = createBlob(chunks, receivedLength, exportType); const downloadUrl = URL.createObjectURL(blob); updateProgress({ status: ExportStatus.COMPLETED, downloadUrl, - filename: filename || `export.${exportType}`, + filename: finalFilename, }); - options.onComplete?.(downloadUrl, filename || `export.${exportType}`); + options.onComplete?.(downloadUrl, finalFilename); } catch (error) { const errorMessage = error instanceof Error ? error.message : 'Unknown error occurred'; diff --git a/superset/commands/chart/data/streaming_export_command.py b/superset/commands/chart/data/streaming_export_command.py index 60493b597f7e..9faaa99a32d5 100644 --- a/superset/commands/chart/data/streaming_export_command.py +++ b/superset/commands/chart/data/streaming_export_command.py @@ -18,6 +18,8 @@ from __future__ import annotations +import csv +import io import logging import time from typing import Callable, Generator, TYPE_CHECKING @@ -103,15 +105,21 @@ def csv_generator() -> Generator[str, None, None]: columns = list(result_proxy.keys()) - # Yield CSV header - header_row = ( - ",".join(f'"{col}"' for col in columns) + "\n" + # Use StringIO buffer with csv.writer for proper escaping + buffer = io.StringIO() + csv_writer = csv.writer( + buffer, quoting=csv.QUOTE_MINIMAL ) - total_bytes += len(header_row.encode("utf-8")) - yield header_row + + # Write CSV header + csv_writer.writerow(columns) + header_data = buffer.getvalue() + total_bytes += len(header_data.encode("utf-8")) + yield header_data + buffer.seek(0) + buffer.truncate() row_count = 0 - buffer = [] buffer_size = 0 flush_threshold = 65536 # 64KB @@ -121,26 +129,27 @@ def csv_generator() -> Generator[str, None, None]: break for row in rows: - csv_row = ",".join( - f'"{str(cell) if cell is not None else ""}"' - for cell in row - ) - csv_line = csv_row + "\n" - row_bytes = len(csv_line.encode("utf-8")) - total_bytes += row_bytes + csv_writer.writerow(row) row_count += 1 - buffer.append(csv_line) - buffer_size += row_bytes - - if buffer_size >= flush_threshold: - yield "".join(buffer) - buffer = [] + # Check buffer size and flush if needed + current_size = buffer.tell() + if current_size >= flush_threshold: + data = buffer.getvalue() + data_bytes = len(data.encode("utf-8")) + total_bytes += data_bytes + buffer_size = data_bytes + yield data + buffer.seek(0) + buffer.truncate() buffer_size = 0 + time.sleep(0.125) # Testing: delay between chunks # Flush remaining buffer - if buffer: - yield "".join(buffer) + remaining_data = buffer.getvalue() + if remaining_data: + total_bytes += len(remaining_data.encode("utf-8")) + yield remaining_data # Log completion total_time = time.time() - start_time From c03c2047a330a4748bfcc3439d5b7f8888fa86f7 Mon Sep 17 00:00:00 2001 From: amaannawab923 Date: Tue, 7 Oct 2025 21:31:02 +0530 Subject: [PATCH 20/57] fix: formatting --- .../StreamingExportModal/StreamingExportModal.tsx | 4 +++- superset/commands/chart/data/streaming_export_command.py | 6 +----- 2 files changed, 4 insertions(+), 6 deletions(-) diff --git a/superset-frontend/src/components/StreamingExportModal/StreamingExportModal.tsx b/superset-frontend/src/components/StreamingExportModal/StreamingExportModal.tsx index 199e2c14b2ec..0ede0b176c46 100644 --- a/superset-frontend/src/components/StreamingExportModal/StreamingExportModal.tsx +++ b/superset-frontend/src/components/StreamingExportModal/StreamingExportModal.tsx @@ -281,7 +281,9 @@ const ModalStateContent = ({ switch (status) { case ExportStatus.ERROR: - return ; + return ( + + ); case ExportStatus.CANCELLED: return ( Generator[str, None, None]: columns = list(result_proxy.keys()) - # Use StringIO buffer with csv.writer for proper escaping + # Use StringIO with csv.writer for proper escaping buffer = io.StringIO() csv_writer = csv.writer( buffer, quoting=csv.QUOTE_MINIMAL @@ -120,7 +120,6 @@ def csv_generator() -> Generator[str, None, None]: buffer.truncate() row_count = 0 - buffer_size = 0 flush_threshold = 65536 # 64KB while True: @@ -138,12 +137,9 @@ def csv_generator() -> Generator[str, None, None]: data = buffer.getvalue() data_bytes = len(data.encode("utf-8")) total_bytes += data_bytes - buffer_size = data_bytes yield data buffer.seek(0) buffer.truncate() - buffer_size = 0 - time.sleep(0.125) # Testing: delay between chunks # Flush remaining buffer remaining_data = buffer.getvalue() From a66bf3414b04b52ceea4199db2c641a5c94211c7 Mon Sep 17 00:00:00 2001 From: amaannawab923 Date: Tue, 7 Oct 2025 22:18:37 +0530 Subject: [PATCH 21/57] fix: filename format --- .../components/gridComponents/Chart/Chart.jsx | 15 ++++++--------- 1 file changed, 6 insertions(+), 9 deletions(-) diff --git a/superset-frontend/src/dashboard/components/gridComponents/Chart/Chart.jsx b/superset-frontend/src/dashboard/components/gridComponents/Chart/Chart.jsx index 21e3d809246e..95fba92d6506 100644 --- a/superset-frontend/src/dashboard/components/gridComponents/Chart/Chart.jsx +++ b/superset-frontend/src/dashboard/components/gridComponents/Chart/Chart.jsx @@ -421,18 +421,15 @@ const Chart = props => { // Handle streaming CSV exports based on row threshold const shouldUseStreaming = format === 'csv' && !isPivot && actualRowCount >= streamingThreshold; - - // Generate filename for streaming exports let filename; if (shouldUseStreaming) { - const timestamp = new Date() - .toISOString() - .slice(0, 19) - .replace(/[-:]/g, '') - .replace('T', '_'); - const chartName = formData.slice_name || formData.viz_type || 'chart'; + const now = new Date(); + const date = now.toISOString().slice(0, 10); + const time = now.toISOString().slice(11, 19).replace(/:/g, ''); + const timestamp = `_${date}_${time}`; + const chartName = slice.slice_name || formData.viz_type || 'chart'; const safeChartName = chartName.replace(/[^a-zA-Z0-9_-]/g, '_'); - filename = `superset_${safeChartName}_${timestamp}.csv`; + filename = `${safeChartName}${timestamp}.csv`; } exportChart({ From 0c3e7f5a4bc4c9ca40285af3d7f5fa8774f13f84 Mon Sep 17 00:00:00 2001 From: amaannawab923 Date: Wed, 8 Oct 2025 12:43:28 +0530 Subject: [PATCH 22/57] fix: sqllab streaming integration --- .../src/SqlLab/components/ResultSet/index.tsx | 194 ++++++++++---- .../useStreamingExport.ts | 69 +++-- .../sql_lab/streaming_export_command.py | 237 ++++++++++++++++++ superset/sqllab/api.py | 122 +++++++++ 4 files changed, 555 insertions(+), 67 deletions(-) create mode 100644 superset/commands/sql_lab/streaming_export_command.py diff --git a/superset-frontend/src/SqlLab/components/ResultSet/index.tsx b/superset-frontend/src/SqlLab/components/ResultSet/index.tsx index 643e4a712c7e..28a93af646bd 100644 --- a/superset-frontend/src/SqlLab/components/ResultSet/index.tsx +++ b/superset-frontend/src/SqlLab/components/ResultSet/index.tsx @@ -87,6 +87,8 @@ import { } from 'src/logger/LogUtils'; import { Icons } from '@superset-ui/core/components/Icons'; import { findPermission } from 'src/utils/findPermission'; +import { StreamingExportModal } from 'src/components/StreamingExportModal'; +import { useStreamingExport } from 'src/components/StreamingExportModal/useStreamingExport'; import ExploreCtasResultsButton from '../ExploreCtasResultsButton'; import ExploreResultsButton from '../ExploreResultsButton'; import HighlightedSql from '../HighlightedSql'; @@ -183,6 +185,10 @@ const ResultSet = ({ defaultQueryLimit, }: ResultSetProps) => { const user = useSelector(({ user }: SqlLabRootState) => user, shallowEqual); + const streamingThreshold = useSelector( + (state: SqlLabRootState) => + state.common?.conf?.CSV_STREAMING_ROW_THRESHOLD || 1000, + ); const query = useSelector( ({ sqlLab: { queries } }: SqlLabRootState) => pick(queries[queryId], [ @@ -222,11 +228,22 @@ const ResultSet = ({ const [searchText, setSearchText] = useState(''); const [cachedData, setCachedData] = useState[]>([]); 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 } = useStreamingExport({ + onComplete: () => { + // Modal will show download button + }, + onError: error => { + addDangerToast(t('Export failed: %s', error)); + }, + }); + const reRunQueryIfSessionTimeoutErrorOnMount = useCallback(() => { if ( query.errorMessage && @@ -300,6 +317,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 = () => { + const { rows, queryLimit, limitingFactor } = query; + const limit = queryLimit || query.results?.query?.limit; + const rowsCount = Math.min(rows || 0, query.results?.data?.length || 0); + + // Determine actual row count + let actualRowCount = rowsCount; + + // If not limited by dropdown/query, use the full row count + if (limitingFactor === LimitingFactor.NotLimited && rows) { + actualRowCount = rows; + } else if (limit) { + actualRowCount = Math.max(actualRowCount, limit); + } + + return actualRowCount >= streamingThreshold; + }; + const renderControls = () => { if (search || visualize || csv) { const { results, queryLimit, limitingFactor, rows } = query; @@ -351,21 +394,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`; + + startExport({ + url: getStreamingExportUrl(), + payload: { client_id: query.id }, + filename, + exportType: 'csv', + expectedRows: rows, }); + } else { + 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() }, + ), + }); + } } }} > @@ -714,42 +783,70 @@ const ResultSet = ({ )} + ); } if (data && data.length === 0) { - return ; + return ( + <> + + + + ); } } if (query.cached || (query.state === QueryState.Success && !query.results)) { if (query.isDataPreview) { return ( - + <> + + + ); } if (query.resultsKey) { return ( - + <> + + + ); } } @@ -764,15 +861,24 @@ const ResultSet = ({ const progressMsg = query?.extra?.progress ?? null; return ( - -
{!progressBar && }
- {/* show loading bar whenever progress bar is completed but needs time to render */} -
{query.progress === 100 && }
- -
{progressMsg && }
-
{query.progress !== 100 && progressBar}
- {trackingUrl &&
{trackingUrl}
} - + <> + +
{!progressBar && }
+ {/* show loading bar whenever progress bar is completed but needs time to render */} +
{query.progress === 100 && }
+ +
+ {progressMsg && } +
+
{query.progress !== 100 && progressBar}
+ {trackingUrl &&
{trackingUrl}
} + + + ); }; diff --git a/superset-frontend/src/components/StreamingExportModal/useStreamingExport.ts b/superset-frontend/src/components/StreamingExportModal/useStreamingExport.ts index 31716878acbf..1c26aa3fe1f4 100644 --- a/superset-frontend/src/components/StreamingExportModal/useStreamingExport.ts +++ b/superset-frontend/src/components/StreamingExportModal/useStreamingExport.ts @@ -17,6 +17,7 @@ * under the License. */ import { useState, useCallback, useRef } from 'react'; +import { SupersetClient } from '@superset-ui/core'; import { ExportStatus, StreamingProgress } from './StreamingExportModal'; interface UseStreamingExportOptions { @@ -38,26 +39,47 @@ interface StreamingExportParams { const NEWLINE_BYTE = 10; // '\n' character code -const createFetchRequest = ( - url: string, +const createFetchRequest = async ( + _url: string, payload: StreamingExportPayload, filename: string, - exportType: string, + _exportType: string, expectedRows: number | undefined, signal: AbortSignal, -): RequestInit => ({ - method: 'POST', - headers: { +): Promise => { + const headers: Record = { 'Content-Type': 'application/x-www-form-urlencoded', - }, - body: new URLSearchParams({ - form_data: JSON.stringify(payload), + }; + + // Get CSRF token using SupersetClient + const csrfToken = await SupersetClient.getCSRFToken(); + if (csrfToken) { + headers['X-CSRFToken'] = csrfToken; + } + + // Build form data - if payload has client_id, it's SQL Lab export + // Otherwise it's a chart export with form_data + const formParams: Record = { filename, expected_rows: expectedRows?.toString() || '', - }), - signal, - credentials: 'same-origin', -}); + }; + + if ('client_id' in payload) { + // SQL Lab export - pass client_id directly + formParams.client_id = String(payload.client_id); + } else { + // Chart export - wrap payload in form_data + formParams.form_data = JSON.stringify(payload); + } + + return { + method: 'POST', + headers, + body: new URLSearchParams(formParams), + signal, + credentials: 'same-origin', + }; +}; const countNewlines = (value: Uint8Array): number => value.filter(byte => byte === NEWLINE_BYTE).length; @@ -107,7 +129,9 @@ export const useStreamingExport = (options: UseStreamingExportOptions = {}) => { exportType, expectedRows, }: StreamingExportParams) => { - if (isExporting) return; + if (isExporting) { + return; + } setIsExporting(true); abortControllerRef.current = new AbortController(); @@ -127,18 +151,17 @@ export const useStreamingExport = (options: UseStreamingExportOptions = {}) => { const defaultFilename = `export.${exportType}`; const finalFilename = filename || defaultFilename; - const response = await fetch( + const fetchOptions = await createFetchRequest( url, - createFetchRequest( - url, - payload, - finalFilename, - exportType, - expectedRows, - abortControllerRef.current.signal, - ), + payload, + finalFilename, + exportType, + expectedRows, + abortControllerRef.current.signal, ); + const response = await fetch(url, fetchOptions); + if (!response.ok) { throw new Error( `Export failed: ${response.status} ${response.statusText}`, diff --git a/superset/commands/sql_lab/streaming_export_command.py b/superset/commands/sql_lab/streaming_export_command.py new file mode 100644 index 000000000000..886685fa3993 --- /dev/null +++ b/superset/commands/sql_lab/streaming_export_command.py @@ -0,0 +1,237 @@ +# 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 + +logger = logging.getLogger(__name__) + + +class StreamingSqlResultExportCommand(BaseCommand): + """ + Command to execute a streaming CSV export of SQL Lab query results. + + This command handles the business logic for: + - Fetching SQL Lab query results + - Generating CSV data in chunks + - Managing database connections + - Buffering data for efficient streaming + """ + + def __init__( + self, + client_id: str, + chunk_size: int = 1000, + ): + """ + Initialize the streaming export command. + + Args: + client_id: The SQL Lab query client ID + chunk_size: Number of rows to fetch per database query (default: 1000) + """ + self._client_id = client_id + self._chunk_size = chunk_size + self._current_app = app._get_current_object() + self._query: Query | None = None + + def validate(self) -> None: + """Validate permissions and query existence.""" + 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 as ex: + raise SupersetErrorException( + SupersetError( + message=__("Cannot access the query"), + error_type=SupersetErrorType.QUERY_SECURITY_ACCESS_ERROR, + level=ErrorLevel.ERROR, + ), + status=403, + ) from ex + + def run(self) -> Callable[[], Generator[str, None, None]]: # noqa: C901 + """ + Execute the streaming CSV export. + + Returns: + A callable that returns a generator yielding CSV data chunks as strings. + The callable is needed to maintain Flask app context during streaming. + """ + # Load all Query attributes while session is still active + # to avoid DetachedInstanceError + assert self._query is not None + + select_sql = self._query.select_sql + executed_sql = self._query.executed_sql + limiting_factor = self._query.limiting_factor + database = self._query.database + + # Get the SQL and limit + if select_sql: + sql = select_sql + limit = None + else: + sql = executed_sql + script = SQLScript(sql, database.db_engine_spec.engine) + # when a query has multiple statements only the last one returns data + limit = script.statements[-1].get_limit_value() + + if limit is not None and limiting_factor in { + LimitingFactor.QUERY, + LimitingFactor.DROPDOWN, + LimitingFactor.QUERY_AND_DROPDOWN, + }: + # remove extra row from `increased_limit` + limit -= 1 + + def csv_generator() -> Generator[str, None, None]: # noqa: C901 + """Generator that yields CSV data from SQL Lab query results.""" + with self._current_app.app_context(): + start_time = time.time() + total_bytes = 0 + + try: + # Create a new session to keep database object attached + with db.session() as session: + # Merge database to prevent DetachedInstanceError + merged_database = session.merge(database) + + # Execute query with streaming + with merged_database.get_sqla_engine() as engine: + connection = engine.connect() + + try: + from sqlalchemy import text + + result_proxy = connection.execution_options( + stream_results=True + ).execute(text(sql)) + + columns = list(result_proxy.keys()) + + # Use StringIO with csv.writer for proper escaping + buffer = io.StringIO() + csv_writer = csv.writer( + buffer, quoting=csv.QUOTE_MINIMAL + ) + + # Write CSV header + csv_writer.writerow(columns) + header_data = buffer.getvalue() + total_bytes += len(header_data.encode("utf-8")) + yield header_data + buffer.seek(0) + buffer.truncate() + + row_count = 0 + flush_threshold = 65536 # 64KB + + while True: + rows = result_proxy.fetchmany(self._chunk_size) + if not rows: + break + + for row in rows: + # Apply limit if specified + if limit is not None and row_count >= limit: + break + + csv_writer.writerow(row) + row_count += 1 + + # Check buffer size and flush if needed + current_size = buffer.tell() + if current_size >= flush_threshold: + data = buffer.getvalue() + data_bytes = len(data.encode("utf-8")) + total_bytes += data_bytes + yield data + buffer.seek(0) + buffer.truncate() + + # Break outer loop if limit reached + if limit is not None and row_count >= limit: + break + + # Flush remaining buffer + remaining_data = buffer.getvalue() + if remaining_data: + total_bytes += len(remaining_data.encode("utf-8")) + yield remaining_data + + # Log completion + total_time = time.time() - start_time + total_mb = total_bytes / (1024 * 1024) + logger.info( + "SQL Lab streaming CSV completed: %s rows, " + "%.1fMB in %.2fs", + f"{row_count:,}", + total_mb, + total_time, + ) + + finally: + connection.close() + + except Exception as e: + logger.error("Error in SQL Lab streaming CSV generator: %s", e) + import traceback + + logger.error("Traceback: %s", traceback.format_exc()) + + # Yield error info and fallback data + yield f"# Error occurred: {str(e)}\n" + yield "error,message\n" + yield f"CSV Export Error,{str(e)}\n" + + return csv_generator diff --git a/superset/sqllab/api.py b/superset/sqllab/api.py index 906dd72bcaf2..8f80499a2356 100644 --- a/superset/sqllab/api.py +++ b/superset/sqllab/api.py @@ -294,6 +294,128 @@ def export_csv(self, client_id: str) -> CsvResponse: ) return response + @expose("/export_streaming/", methods=("POST",)) + @statsd_metrics + @event_logger.log_this_with_context( + action=lambda self, + *args, + **kwargs: f"{self.__class__.__name__}.export_streaming_csv", + log_to_statsd=False, + ) + def export_streaming_csv(self) -> Response: + """Export SQL query results using streaming for large datasets. + --- + post: + summary: Export SQL query results to CSV with streaming + requestBody: + description: Export parameters + required: true + content: + application/x-www-form-urlencoded: + schema: + type: object + properties: + client_id: + type: string + description: The SQL query result identifier + filename: + type: string + description: Optional filename for the export + expected_rows: + type: integer + description: Optional expected row count for progress tracking + responses: + 200: + description: Streaming CSV export + 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' + """ + # Extract parameters from form data + client_id = request.form.get("client_id") + filename = request.form.get("filename") + + if not client_id: + return self.response_400(message="client_id is required") + + expected_rows = None + if expected_rows_str := request.form.get("expected_rows"): + try: + expected_rows = int(expected_rows_str) + except (ValueError, TypeError): + logger.warning("Invalid expected_rows value: %s", expected_rows_str) + + return self._create_streaming_csv_response(client_id, filename, expected_rows) + + def _create_streaming_csv_response( + self, + client_id: str, + filename: str | None = None, + expected_rows: int | None = None, + ) -> Response: + """Create a streaming CSV response for large SQL Lab result sets.""" + from datetime import datetime + + from superset.commands.sql_lab.streaming_export_command import ( + StreamingSqlResultExportCommand, + ) + + # Execute streaming command + chunk_size = 1000 + command = StreamingSqlResultExportCommand(client_id, chunk_size) + command.validate() + + # Generate filename if not provided + if not filename: + query = command._query + assert query is not None + timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") + safe_name = "".join( + c for c in (query.name or "query") if c.isalnum() or c in ("-", "_") + ) + filename = f"sqllab_{safe_name}_{timestamp}.csv" + + # Get the callable that returns the generator + csv_generator_callable = command.run() + + # Get encoding from config + encoding = app.config.get("CSV_EXPORT", {}).get("encoding", "utf-8") + + # Create response with streaming headers + response = Response( + csv_generator_callable(), # Call the callable to get generator + mimetype=f"text/csv; charset={encoding}", + headers={ + "Content-Disposition": f'attachment; filename="{filename}"', + "Cache-Control": "no-cache", + "X-Accel-Buffering": "no", # Disable nginx buffering + "X-Superset-Streaming": "true", # Identify streaming responses + }, + direct_passthrough=False, # Flask must iterate generator + ) + + # Force chunked transfer encoding + response.implicit_sequence_conversion = False + + logger.info( + "SQL Lab streaming CSV export started: client_id=%s, filename=%s", + client_id, + filename, + ) + + return response + @expose("/results/") @protect() @statsd_metrics From da789e9aff5bdaba44b5ad930ae97c816e76ef90 Mon Sep 17 00:00:00 2001 From: amaannawab923 Date: Wed, 8 Oct 2025 13:49:30 +0530 Subject: [PATCH 23/57] fix: permission for sqllab streaming --- superset/sqllab/api.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/superset/sqllab/api.py b/superset/sqllab/api.py index 8f80499a2356..e25ab4eb122e 100644 --- a/superset/sqllab/api.py +++ b/superset/sqllab/api.py @@ -295,6 +295,8 @@ def export_csv(self, client_id: str) -> CsvResponse: return response @expose("/export_streaming/", methods=("POST",)) + @protect() + @permission_name("read") @statsd_metrics @event_logger.log_this_with_context( action=lambda self, From bbb52dde1e647c3d4b94b502a2783beb257e4c7e Mon Sep 17 00:00:00 2001 From: amaannawab923 Date: Wed, 8 Oct 2025 16:36:01 +0530 Subject: [PATCH 24/57] fix: retry mechanism in case of error --- .../src/SqlLab/components/ResultSet/index.tsx | 21 ++-- .../StreamingExportModal.tsx | 38 +++++++- .../useStreamingExport.ts | 95 ++++++++++++++++--- .../chart/data/streaming_export_command.py | 14 +-- .../sql_lab/streaming_export_command.py | 10 +- 5 files changed, 143 insertions(+), 35 deletions(-) diff --git a/superset-frontend/src/SqlLab/components/ResultSet/index.tsx b/superset-frontend/src/SqlLab/components/ResultSet/index.tsx index 28a93af646bd..aac184baea3b 100644 --- a/superset-frontend/src/SqlLab/components/ResultSet/index.tsx +++ b/superset-frontend/src/SqlLab/components/ResultSet/index.tsx @@ -235,14 +235,15 @@ const ResultSet = ({ const logAction = useLogAction({ queryId, sqlEditorId: query.sqlEditorId }); // Streaming export hook - const { progress, startExport, resetExport } = useStreamingExport({ - onComplete: () => { - // Modal will show download button - }, - onError: error => { - addDangerToast(t('Export failed: %s', error)); - }, - }); + const { progress, startExport, resetExport, retryExport } = + useStreamingExport({ + onComplete: () => { + // Modal will show download button + }, + onError: error => { + addDangerToast(t('Export failed: %s', error)); + }, + }); const reRunQueryIfSessionTimeoutErrorOnMount = useCallback(() => { if ( @@ -786,6 +787,7 @@ const ResultSet = ({ @@ -798,6 +800,7 @@ const ResultSet = ({ @@ -826,6 +829,7 @@ const ResultSet = ({ @@ -844,6 +848,7 @@ const ResultSet = ({ diff --git a/superset-frontend/src/components/StreamingExportModal/StreamingExportModal.tsx b/superset-frontend/src/components/StreamingExportModal/StreamingExportModal.tsx index 0ede0b176c46..fc7a4531d63b 100644 --- a/superset-frontend/src/components/StreamingExportModal/StreamingExportModal.tsx +++ b/superset-frontend/src/components/StreamingExportModal/StreamingExportModal.tsx @@ -41,6 +41,7 @@ export interface StreamingProgress { speed?: number; mbPerSecond?: number; elapsedTime?: number; + retryCount?: number; } interface StreamingExportModalProps { @@ -72,6 +73,22 @@ const SuccessIcon = styled(Icons.CheckCircleFilled)` flex-shrink: 0; `; +const ErrorIconWrapper = styled.div` + display: flex; + align-items: center; + justify-content: center; + width: 16px; + height: 16px; + background-color: #ff4d4f; + border-radius: 50%; + flex-shrink: 0; +`; + +const ErrorIconStyled = styled(Icons.CloseOutlined)` + color: white; + font-size: 10px; +`; + const ActionButtons = styled.div` display: flex; gap: ${({ theme }) => theme.sizeUnit * 2}px; @@ -147,14 +164,26 @@ const ErrorContent = ({ error, onCancel, onRetry, + getProgressPercentage, }: { error?: string; onCancel: () => void; onRetry?: () => void; + getProgressPercentage: () => number; }) => ( - + + + + + + {error || t('Export failed')} @@ -282,7 +311,12 @@ const ModalStateContent = ({ switch (status) { case ExportStatus.ERROR: return ( - + ); case ExportStatus.CANCELLED: return ( diff --git a/superset-frontend/src/components/StreamingExportModal/useStreamingExport.ts b/superset-frontend/src/components/StreamingExportModal/useStreamingExport.ts index 1c26aa3fe1f4..650261fc8c73 100644 --- a/superset-frontend/src/components/StreamingExportModal/useStreamingExport.ts +++ b/superset-frontend/src/components/StreamingExportModal/useStreamingExport.ts @@ -115,25 +115,18 @@ export const useStreamingExport = (options: UseStreamingExportOptions = {}) => { status: ExportStatus.STREAMING, }); const [isExporting, setIsExporting] = useState(false); + const [retryCount, setRetryCount] = useState(0); const abortControllerRef = useRef(null); + const lastExportParamsRef = useRef(null); const updateProgress = useCallback((updates: Partial) => { setProgress(prev => ({ ...prev, ...updates })); }, []); - const startExport = useCallback( - async ({ - url, - payload, - filename, - exportType, - expectedRows, - }: StreamingExportParams) => { - if (isExporting) { - return; - } + const executeExport = useCallback( + async (params: StreamingExportParams) => { + const { url, payload, filename, exportType, expectedRows } = params; - setIsExporting(true); abortControllerRef.current = new AbortController(); updateProgress({ @@ -176,18 +169,46 @@ export const useStreamingExport = (options: UseStreamingExportOptions = {}) => { const chunks: Uint8Array[] = []; let receivedLength = 0; let rowsProcessed = 0; + let hasError = false; // eslint-disable-next-line no-constant-condition while (true) { // eslint-disable-next-line no-await-in-loop const { done, value } = await reader.read(); - if (done) break; + if (done) { + break; + } if (abortControllerRef.current?.signal.aborted) { throw new Error('Export cancelled by user'); } + // Check for error marker in the chunk + const textDecoder = new TextDecoder(); + const chunkText = textDecoder.decode(value); + + if (chunkText.includes('__STREAM_ERROR__')) { + const errorMatch = chunkText.match(/__STREAM_ERROR__:(.+)/); + const errorMsg = errorMatch + ? errorMatch[1].trim() + : 'Export failed. Please try again in some time.'; + + // Update progress to show error with current progress preserved + updateProgress({ + status: ExportStatus.ERROR, + error: errorMsg, + rowsProcessed, + totalRows: expectedRows, + totalSize: receivedLength, + }); + + setIsExporting(false); + options.onError?.(errorMsg); + hasError = true; + break; + } + chunks.push(value); receivedLength += value.length; @@ -205,6 +226,11 @@ export const useStreamingExport = (options: UseStreamingExportOptions = {}) => { }); } + // Check if we exited early due to error marker + if (hasError) { + return; + } + const blob = createBlob(chunks, receivedLength, exportType); const downloadUrl = URL.createObjectURL(blob); @@ -226,21 +252,58 @@ export const useStreamingExport = (options: UseStreamingExportOptions = {}) => { updateProgress({ status: ExportStatus.CANCELLED, }); + setIsExporting(false); } else { updateProgress({ status: ExportStatus.ERROR, error: errorMessage, }); options.onError?.(errorMessage); + setIsExporting(false); } } finally { - setIsExporting(false); abortControllerRef.current = null; } }, - [isExporting, updateProgress, options], + [updateProgress, options], + ); + + const startExport = useCallback( + async (params: StreamingExportParams) => { + if (isExporting) { + return; + } + + setIsExporting(true); + setRetryCount(0); + lastExportParamsRef.current = params; + + updateProgress({ + rowsProcessed: 0, + totalRows: params.expectedRows, + totalSize: 0, + speed: 0, + mbPerSecond: 0, + elapsedTime: 0, + status: ExportStatus.STREAMING, + filename: params.filename, + }); + + executeExport(params); + }, + [isExporting, updateProgress, executeExport], ); + const retryExport = useCallback(() => { + if (!lastExportParamsRef.current) { + return; + } + + setIsExporting(true); + setRetryCount(0); + executeExport(lastExportParamsRef.current); + }, [executeExport]); + const cancelExport = useCallback(() => { if (abortControllerRef.current) { abortControllerRef.current.abort(); @@ -267,8 +330,10 @@ export const useStreamingExport = (options: UseStreamingExportOptions = {}) => { return { progress, isExporting, + retryCount, startExport, cancelExport, resetExport, + retryExport, }; }; diff --git a/superset/commands/chart/data/streaming_export_command.py b/superset/commands/chart/data/streaming_export_command.py index 7cbcade773a7..153e4d9cd18a 100644 --- a/superset/commands/chart/data/streaming_export_command.py +++ b/superset/commands/chart/data/streaming_export_command.py @@ -65,7 +65,7 @@ def validate(self) -> None: """Validate permissions and query context.""" self._query_context.raise_for_access() - def run(self) -> Callable[[], Generator[str, None, None]]: + def run(self) -> Callable[[], Generator[str, None, None]]: # noqa: C901 """ Execute the streaming CSV export. @@ -74,7 +74,7 @@ def run(self) -> Callable[[], Generator[str, None, None]]: The callable is needed to maintain Flask app context during streaming. """ - def csv_generator() -> Generator[str, None, None]: + def csv_generator() -> Generator[str, None, None]: # noqa: C901 """Generator that yields CSV data from database query.""" with self._current_app.app_context(): start_time = time.time() @@ -166,9 +166,11 @@ def csv_generator() -> Generator[str, None, None]: logger.error("Traceback: %s", traceback.format_exc()) - # Yield error info and fallback data - yield f"# Error occurred: {str(e)}\n" - yield "error,message\n" - yield f"CSV Export Error,{str(e)}\n" + # Send error marker for frontend to detect + error_marker = ( + "__STREAM_ERROR__:Export failed. " + "Please try again in some time.\n" + ) + yield error_marker return csv_generator diff --git a/superset/commands/sql_lab/streaming_export_command.py b/superset/commands/sql_lab/streaming_export_command.py index 886685fa3993..6742f5dccc16 100644 --- a/superset/commands/sql_lab/streaming_export_command.py +++ b/superset/commands/sql_lab/streaming_export_command.py @@ -229,9 +229,11 @@ def csv_generator() -> Generator[str, None, None]: # noqa: C901 logger.error("Traceback: %s", traceback.format_exc()) - # Yield error info and fallback data - yield f"# Error occurred: {str(e)}\n" - yield "error,message\n" - yield f"CSV Export Error,{str(e)}\n" + # Send error marker for frontend to detect + error_marker = ( + "__STREAM_ERROR__:Export failed. " + "Please try again in some time.\n" + ) + yield error_marker return csv_generator From 957b7965bedb1b2ab40371c1e5a7b70e16b31f0f Mon Sep 17 00:00:00 2001 From: amaannawab923 Date: Wed, 8 Oct 2025 17:00:25 +0530 Subject: [PATCH 25/57] fix: retry mechanism for chart --- .../components/gridComponents/Chart/Chart.jsx | 37 +++++++++---------- 1 file changed, 18 insertions(+), 19 deletions(-) diff --git a/superset-frontend/src/dashboard/components/gridComponents/Chart/Chart.jsx b/superset-frontend/src/dashboard/components/gridComponents/Chart/Chart.jsx index 95fba92d6506..3a3b5f1afd88 100644 --- a/superset-frontend/src/dashboard/components/gridComponents/Chart/Chart.jsx +++ b/superset-frontend/src/dashboard/components/gridComponents/Chart/Chart.jsx @@ -183,19 +183,23 @@ const Chart = props => { const [width, setWidth] = useState(props.width); const [isStreamingModalVisible, setIsStreamingModalVisible] = useState(false); - const { progress, isExporting, startExport, cancelExport, resetExport } = - useStreamingExport({ - onComplete: () => { - boundActionCreators.addSuccessToast( - t('CSV file downloaded successfully'), - ); - }, - onError: () => { - boundActionCreators.addDangerToast( - t('Export failed - please try again'), - ); - }, - }); + const { + progress, + isExporting, + startExport, + cancelExport, + resetExport, + retryExport, + } = useStreamingExport({ + onComplete: () => { + boundActionCreators.addSuccessToast( + t('CSV file downloaded successfully'), + ); + }, + onError: () => { + boundActionCreators.addDangerToast(t('Export failed - please try again')); + }, + }); const history = useHistory(); const resize = useCallback( debounce(() => { @@ -623,12 +627,7 @@ const Chart = props => { } setIsStreamingModalVisible(false); }} - onRetry={() => { - resetExport(); - // Note: Retry would need to store the last export parameters - // For now, just close the modal and let user retry manually - setIsStreamingModalVisible(false); - }} + onRetry={retryExport} progress={progress} exportType="csv" /> From c6d18387399891b4236c68b9f07ebf26f524f05f Mon Sep 17 00:00:00 2001 From: amaannawab923 Date: Wed, 8 Oct 2025 17:04:41 +0530 Subject: [PATCH 26/57] update: Error message --- .../src/components/StreamingExportModal/useStreamingExport.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/superset-frontend/src/components/StreamingExportModal/useStreamingExport.ts b/superset-frontend/src/components/StreamingExportModal/useStreamingExport.ts index 650261fc8c73..7b31f90933e5 100644 --- a/superset-frontend/src/components/StreamingExportModal/useStreamingExport.ts +++ b/superset-frontend/src/components/StreamingExportModal/useStreamingExport.ts @@ -192,7 +192,7 @@ export const useStreamingExport = (options: UseStreamingExportOptions = {}) => { const errorMatch = chunkText.match(/__STREAM_ERROR__:(.+)/); const errorMsg = errorMatch ? errorMatch[1].trim() - : 'Export failed. Please try again in some time.'; + : 'Export failed. Please try again.'; // Update progress to show error with current progress preserved updateProgress({ From f9a72c8fe41cc71e45c92d3b58abacc0e486d273 Mon Sep 17 00:00:00 2001 From: amaannawab923 Date: Thu, 9 Oct 2025 18:35:59 +0530 Subject: [PATCH 27/57] fix: subsequent csv downloads failures fixed for charts --- .../useStreamingExport.ts | 49 ++++++++++++++----- .../components/gridComponents/Chart/Chart.jsx | 2 +- 2 files changed, 39 insertions(+), 12 deletions(-) diff --git a/superset-frontend/src/components/StreamingExportModal/useStreamingExport.ts b/superset-frontend/src/components/StreamingExportModal/useStreamingExport.ts index 7b31f90933e5..efd723fed6bc 100644 --- a/superset-frontend/src/components/StreamingExportModal/useStreamingExport.ts +++ b/superset-frontend/src/components/StreamingExportModal/useStreamingExport.ts @@ -16,7 +16,7 @@ * specific language governing permissions and limitations * under the License. */ -import { useState, useCallback, useRef } from 'react'; +import { useState, useCallback, useRef, useEffect } from 'react'; import { SupersetClient } from '@superset-ui/core'; import { ExportStatus, StreamingProgress } from './StreamingExportModal'; @@ -114,10 +114,11 @@ export const useStreamingExport = (options: UseStreamingExportOptions = {}) => { elapsedTime: 0, status: ExportStatus.STREAMING, }); - const [isExporting, setIsExporting] = useState(false); const [retryCount, setRetryCount] = useState(0); const abortControllerRef = useRef(null); const lastExportParamsRef = useRef(null); + const currentBlobUrlRef = useRef(null); + const isExportingRef = useRef(false); const updateProgress = useCallback((updates: Partial) => { setProgress(prev => ({ ...prev, ...updates })); @@ -203,7 +204,7 @@ export const useStreamingExport = (options: UseStreamingExportOptions = {}) => { totalSize: receivedLength, }); - setIsExporting(false); + isExportingRef.current = false; options.onError?.(errorMsg); hasError = true; break; @@ -232,7 +233,13 @@ export const useStreamingExport = (options: UseStreamingExportOptions = {}) => { } const blob = createBlob(chunks, receivedLength, exportType); + + if (currentBlobUrlRef.current) { + URL.revokeObjectURL(currentBlobUrlRef.current); + } + const downloadUrl = URL.createObjectURL(blob); + currentBlobUrlRef.current = downloadUrl; updateProgress({ status: ExportStatus.COMPLETED, @@ -240,6 +247,7 @@ export const useStreamingExport = (options: UseStreamingExportOptions = {}) => { filename: finalFilename, }); + isExportingRef.current = false; options.onComplete?.(downloadUrl, finalFilename); } catch (error) { const errorMessage = @@ -252,14 +260,14 @@ export const useStreamingExport = (options: UseStreamingExportOptions = {}) => { updateProgress({ status: ExportStatus.CANCELLED, }); - setIsExporting(false); + isExportingRef.current = false; } else { updateProgress({ status: ExportStatus.ERROR, error: errorMessage, }); options.onError?.(errorMessage); - setIsExporting(false); + isExportingRef.current = false; } } finally { abortControllerRef.current = null; @@ -270,11 +278,11 @@ export const useStreamingExport = (options: UseStreamingExportOptions = {}) => { const startExport = useCallback( async (params: StreamingExportParams) => { - if (isExporting) { + if (isExportingRef.current) { return; } - setIsExporting(true); + isExportingRef.current = true; setRetryCount(0); lastExportParamsRef.current = params; @@ -291,7 +299,7 @@ export const useStreamingExport = (options: UseStreamingExportOptions = {}) => { executeExport(params); }, - [isExporting, updateProgress, executeExport], + [updateProgress, executeExport], ); const retryExport = useCallback(() => { @@ -299,7 +307,11 @@ export const useStreamingExport = (options: UseStreamingExportOptions = {}) => { return; } - setIsExporting(true); + if (isExportingRef.current) { + return; + } + + isExportingRef.current = true; setRetryCount(0); executeExport(lastExportParamsRef.current); }, [executeExport]); @@ -314,7 +326,12 @@ export const useStreamingExport = (options: UseStreamingExportOptions = {}) => { }, [updateProgress]); const resetExport = useCallback(() => { - setIsExporting(false); + if (currentBlobUrlRef.current) { + URL.revokeObjectURL(currentBlobUrlRef.current); + currentBlobUrlRef.current = null; + } + + isExportingRef.current = false; abortControllerRef.current = null; setProgress({ rowsProcessed: 0, @@ -327,9 +344,19 @@ export const useStreamingExport = (options: UseStreamingExportOptions = {}) => { }); }, []); + // Cleanup blob URL on unmount to prevent memory leak + useEffect( + () => () => { + if (currentBlobUrlRef.current) { + URL.revokeObjectURL(currentBlobUrlRef.current); + } + }, + [], + ); + return { progress, - isExporting, + isExporting: isExportingRef.current, retryCount, startExport, cancelExport, diff --git a/superset-frontend/src/dashboard/components/gridComponents/Chart/Chart.jsx b/superset-frontend/src/dashboard/components/gridComponents/Chart/Chart.jsx index 3a3b5f1afd88..fe2b4a6396f7 100644 --- a/superset-frontend/src/dashboard/components/gridComponents/Chart/Chart.jsx +++ b/superset-frontend/src/dashboard/components/gridComponents/Chart/Chart.jsx @@ -445,7 +445,6 @@ const Chart = props => { onStartStreamingExport: shouldUseStreaming ? exportParams => { setIsStreamingModalVisible(true); - resetExport(); startExport({ ...exportParams, filename, @@ -626,6 +625,7 @@ const Chart = props => { cancelExport(); } setIsStreamingModalVisible(false); + resetExport(); }} onRetry={retryExport} progress={progress} From 1e97167c9d470f5b4c13379e04fe21c6dee4c619 Mon Sep 17 00:00:00 2001 From: amaannawab923 Date: Fri, 10 Oct 2025 01:36:48 +0530 Subject: [PATCH 28/57] Implemented Tests for Streaming Modal ui Component --- .../StreamingExportModal.test.tsx | 225 ++++++++++++++++++ 1 file changed, 225 insertions(+) create mode 100644 superset-frontend/src/components/StreamingExportModal/StreamingExportModal.test.tsx diff --git a/superset-frontend/src/components/StreamingExportModal/StreamingExportModal.test.tsx b/superset-frontend/src/components/StreamingExportModal/StreamingExportModal.test.tsx new file mode 100644 index 000000000000..eb288d23f301 --- /dev/null +++ b/superset-frontend/src/components/StreamingExportModal/StreamingExportModal.test.tsx @@ -0,0 +1,225 @@ +/** + * 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 { render, screen, userEvent } from 'spec/helpers/testing-library'; +import StreamingExportModal, { + ExportStatus, + StreamingProgress, +} from './StreamingExportModal'; + +const defaultProgress: StreamingProgress = { + rowsProcessed: 0, + totalRows: 1000, + totalSize: 0, + status: ExportStatus.STREAMING, + filename: 'test_export.csv', +}; + +const defaultProps = { + visible: true, + onCancel: jest.fn(), + onRetry: jest.fn(), + progress: defaultProgress, +}; + +beforeEach(() => { + jest.clearAllMocks(); + URL.revokeObjectURL = jest.fn(); + URL.createObjectURL = jest.fn(() => 'blob:mock-url'); +}); + +test('renders modal with streaming state', () => { + render(); + + expect(screen.getByText('Exporting Data')).toBeInTheDocument(); + expect(screen.getByText(/Processing export for test_export.csv/i)).toBeInTheDocument(); + expect(screen.getByRole('button', { name: 'Cancel' })).toBeInTheDocument(); + expect(screen.getByRole('button', { name: 'Download' })).toBeDisabled(); +}); + +test('shows progress percentage during streaming', () => { + const progress = { + ...defaultProgress, + rowsProcessed: 500, + totalRows: 1000, + status: ExportStatus.STREAMING, + }; + + render(); + + expect(screen.getByRole('progressbar')).toBeInTheDocument(); +}); + +test('shows completed state when export finishes', () => { + const progress = { + ...defaultProgress, + rowsProcessed: 1000, + totalRows: 1000, + status: ExportStatus.COMPLETED, + downloadUrl: 'blob:mock-url', + }; + + render(); + + expect(screen.getByText(/Export successful: test_export.csv/i)).toBeInTheDocument(); + expect(screen.getByRole('button', { name: 'Download' })).not.toBeDisabled(); +}); + +test('shows error state when export fails', () => { + const progress = { + ...defaultProgress, + status: ExportStatus.ERROR, + error: 'Database connection failed', + }; + + render(); + + expect(screen.getByText('Database connection failed')).toBeInTheDocument(); + expect(screen.getByRole('button', { name: 'Retry' })).toBeInTheDocument(); +}); + +test('shows cancelled state when export is cancelled', () => { + const progress = { + ...defaultProgress, + status: ExportStatus.CANCELLED, + }; + + render(); + + expect(screen.getByText('Export cancelled')).toBeInTheDocument(); + expect(screen.getAllByRole('button', { name: 'Close' })[0]).toBeInTheDocument(); +}); + +test('calls onCancel when cancel button is clicked during streaming', async () => { + const onCancel = jest.fn(); + render(); + + await userEvent.click(screen.getByRole('button', { name: 'Cancel' })); + expect(onCancel).toHaveBeenCalledTimes(1); +}); + +test('calls onRetry when retry button is clicked after error', async () => { + const onRetry = jest.fn(); + const progress = { + ...defaultProgress, + status: ExportStatus.ERROR, + error: 'Network error', + }; + + render( + , + ); + + await userEvent.click(screen.getByRole('button', { name: 'Retry' })); + expect(onRetry).toHaveBeenCalledTimes(1); +}); + +test('triggers download when download button is clicked', async () => { + const progress = { + ...defaultProgress, + rowsProcessed: 1000, + totalRows: 1000, + status: ExportStatus.COMPLETED, + downloadUrl: 'blob:mock-url', + filename: 'test_export.csv', + }; + + const onCancel = jest.fn(); + render( + , + ); + + const downloadButton = screen.getByRole('button', { name: 'Download' }); + expect(downloadButton).not.toBeDisabled(); + + await userEvent.click(downloadButton); + + expect(onCancel).toHaveBeenCalledTimes(1); +}); + +test('does not show download button when downloadUrl is missing', () => { + const progress = { + ...defaultProgress, + status: ExportStatus.COMPLETED, + }; + + render(); + + expect(screen.getByRole('button', { name: 'Download' })).toBeDisabled(); +}); + +test('progress bar shows correct percentage with decimal precision', () => { + const progress = { + ...defaultProgress, + rowsProcessed: 333, + totalRows: 1000, + status: ExportStatus.STREAMING, + }; + + render(); + + expect(screen.getByRole('progressbar')).toBeInTheDocument(); +}); + +test('shows generic processing message when filename is not provided', () => { + const progress = { + ...defaultProgress, + filename: undefined, + status: ExportStatus.STREAMING, + }; + + render(); + + expect(screen.getByText('Processing export...')).toBeInTheDocument(); +}); + +test('handles retry button visibility based on onRetry prop', () => { + const progress = { + ...defaultProgress, + status: ExportStatus.ERROR, + error: 'Test error', + }; + + const { rerender } = render( + , + ); + + expect(screen.queryByRole('button', { name: 'Retry' })).not.toBeInTheDocument(); + + rerender( + , + ); + + expect(screen.getByRole('button', { name: 'Retry' })).toBeInTheDocument(); +}); + +test('shows generic error message when error details are not provided', () => { + const progress = { + ...defaultProgress, + status: ExportStatus.ERROR, + }; + + render(); + + expect(screen.getByText('Export failed')).toBeInTheDocument(); +}); + From e6f65c4e6db3065bdb0a5dcb109c670329ecc5c2 Mon Sep 17 00:00:00 2001 From: amaannawab923 Date: Fri, 10 Oct 2025 01:37:44 +0530 Subject: [PATCH 29/57] Formatting fix --- .../StreamingExportModal.test.tsx | 35 ++++++++++++++----- 1 file changed, 27 insertions(+), 8 deletions(-) diff --git a/superset-frontend/src/components/StreamingExportModal/StreamingExportModal.test.tsx b/superset-frontend/src/components/StreamingExportModal/StreamingExportModal.test.tsx index eb288d23f301..56b647d6a396 100644 --- a/superset-frontend/src/components/StreamingExportModal/StreamingExportModal.test.tsx +++ b/superset-frontend/src/components/StreamingExportModal/StreamingExportModal.test.tsx @@ -47,7 +47,9 @@ test('renders modal with streaming state', () => { render(); expect(screen.getByText('Exporting Data')).toBeInTheDocument(); - expect(screen.getByText(/Processing export for test_export.csv/i)).toBeInTheDocument(); + expect( + screen.getByText(/Processing export for test_export.csv/i), + ).toBeInTheDocument(); expect(screen.getByRole('button', { name: 'Cancel' })).toBeInTheDocument(); expect(screen.getByRole('button', { name: 'Download' })).toBeDisabled(); }); @@ -76,7 +78,9 @@ test('shows completed state when export finishes', () => { render(); - expect(screen.getByText(/Export successful: test_export.csv/i)).toBeInTheDocument(); + expect( + screen.getByText(/Export successful: test_export.csv/i), + ).toBeInTheDocument(); expect(screen.getByRole('button', { name: 'Download' })).not.toBeDisabled(); }); @@ -102,7 +106,9 @@ test('shows cancelled state when export is cancelled', () => { render(); expect(screen.getByText('Export cancelled')).toBeInTheDocument(); - expect(screen.getAllByRole('button', { name: 'Close' })[0]).toBeInTheDocument(); + expect( + screen.getAllByRole('button', { name: 'Close' })[0], + ).toBeInTheDocument(); }); test('calls onCancel when cancel button is clicked during streaming', async () => { @@ -122,7 +128,11 @@ test('calls onRetry when retry button is clicked after error', async () => { }; render( - , + , ); await userEvent.click(screen.getByRole('button', { name: 'Retry' })); @@ -141,7 +151,11 @@ test('triggers download when download button is clicked', async () => { const onCancel = jest.fn(); render( - , + , ); const downloadButton = screen.getByRole('button', { name: 'Download' }); @@ -196,10 +210,16 @@ test('handles retry button visibility based on onRetry prop', () => { }; const { rerender } = render( - , + , ); - expect(screen.queryByRole('button', { name: 'Retry' })).not.toBeInTheDocument(); + expect( + screen.queryByRole('button', { name: 'Retry' }), + ).not.toBeInTheDocument(); rerender( { expect(screen.getByText('Export failed')).toBeInTheDocument(); }); - From 2371503a3dea67cdaa21027a12f7d2e640e8fd51 Mon Sep 17 00:00:00 2001 From: amaannawab923 Date: Fri, 10 Oct 2025 01:51:23 +0530 Subject: [PATCH 30/57] fix: eslint --- .../StreamingExportModal/StreamingExportModal.test.tsx | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/superset-frontend/src/components/StreamingExportModal/StreamingExportModal.test.tsx b/superset-frontend/src/components/StreamingExportModal/StreamingExportModal.test.tsx index 56b647d6a396..7663a8614e95 100644 --- a/superset-frontend/src/components/StreamingExportModal/StreamingExportModal.test.tsx +++ b/superset-frontend/src/components/StreamingExportModal/StreamingExportModal.test.tsx @@ -81,7 +81,7 @@ test('shows completed state when export finishes', () => { expect( screen.getByText(/Export successful: test_export.csv/i), ).toBeInTheDocument(); - expect(screen.getByRole('button', { name: 'Download' })).not.toBeDisabled(); + expect(screen.getByRole('button', { name: 'Download' })).toBeEnabled(); }); test('shows error state when export fails', () => { @@ -159,7 +159,7 @@ test('triggers download when download button is clicked', async () => { ); const downloadButton = screen.getByRole('button', { name: 'Download' }); - expect(downloadButton).not.toBeDisabled(); + expect(downloadButton).toBeEnabled(); await userEvent.click(downloadButton); From 3c2efa09f04bde18155c4b5fa57f076e2177aa27 Mon Sep 17 00:00:00 2001 From: amaannawab923 Date: Fri, 10 Oct 2025 13:16:51 +0530 Subject: [PATCH 31/57] update: hook unit tests --- .../useStreamingExport.test.ts | 130 ++++++++++++++++++ 1 file changed, 130 insertions(+) create mode 100644 superset-frontend/src/components/StreamingExportModal/useStreamingExport.test.ts diff --git a/superset-frontend/src/components/StreamingExportModal/useStreamingExport.test.ts b/superset-frontend/src/components/StreamingExportModal/useStreamingExport.test.ts new file mode 100644 index 000000000000..268da496ea74 --- /dev/null +++ b/superset-frontend/src/components/StreamingExportModal/useStreamingExport.test.ts @@ -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. + */ +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 +global.URL.createObjectURL = jest.fn(() => 'blob:mock-url'); +global.URL.revokeObjectURL = jest.fn(); + +// Mock fetch +global.fetch = jest.fn(); + +describe('useStreamingExport', () => { + beforeEach(() => { + jest.clearAllMocks(); + }); + + test('initializes with default progress state', () => { + const { result } = renderHook(() => useStreamingExport()); + + expect(result.current.progress).toEqual({ + rowsProcessed: 0, + totalRows: undefined, + totalSize: 0, + speed: 0, + mbPerSecond: 0, + elapsedTime: 0, + status: ExportStatus.STREAMING, + }); + }); + + test('provides startExport function', () => { + const { result } = renderHook(() => useStreamingExport()); + + expect(typeof result.current.startExport).toBe('function'); + }); + + test('provides resetExport function', () => { + const { result } = renderHook(() => useStreamingExport()); + + expect(typeof result.current.resetExport).toBe('function'); + }); + + test('provides retryExport function', () => { + const { result } = renderHook(() => useStreamingExport()); + + expect(typeof result.current.retryExport).toBe('function'); + }); + + test('provides cancelExport function', () => { + const { result } = renderHook(() => useStreamingExport()); + + expect(typeof result.current.cancelExport).toBe('function'); + }); + + test('resetExport resets progress to initial state', () => { + const { result } = renderHook(() => useStreamingExport()); + + act(() => { + result.current.resetExport(); + }); + + expect(result.current.progress.status).toBe(ExportStatus.STREAMING); + expect(result.current.progress.rowsProcessed).toBe(0); + expect(result.current.progress.totalSize).toBe(0); + }); + + test('accepts onComplete callback option', () => { + const onComplete = jest.fn(); + const { result } = renderHook(() => useStreamingExport({ onComplete })); + + expect(result.current).toBeDefined(); + }); + + test('accepts onError callback option', () => { + const onError = jest.fn(); + const { result } = renderHook(() => useStreamingExport({ onError })); + + expect(result.current).toBeDefined(); + }); + + test('progress includes all required fields', () => { + const { result } = renderHook(() => useStreamingExport()); + + expect(result.current.progress).toHaveProperty('rowsProcessed'); + expect(result.current.progress).toHaveProperty('totalRows'); + expect(result.current.progress).toHaveProperty('totalSize'); + expect(result.current.progress).toHaveProperty('status'); + expect(result.current.progress).toHaveProperty('speed'); + expect(result.current.progress).toHaveProperty('mbPerSecond'); + expect(result.current.progress).toHaveProperty('elapsedTime'); + }); + + test('cleans up on unmount', () => { + const revokeObjectURL = jest.fn(); + global.URL.revokeObjectURL = revokeObjectURL; + + const { unmount } = renderHook(() => useStreamingExport()); + + unmount(); + + // Cleanup should not throw errors + expect(true).toBe(true); + }); +}); From 096ba4e4b394fb9d440f2d111567409789f64b84 Mon Sep 17 00:00:00 2001 From: amaannawab923 Date: Fri, 10 Oct 2025 15:57:32 +0530 Subject: [PATCH 32/57] fix: eslint error --- .../useStreamingExport.test.ts | 136 +++++++++--------- 1 file changed, 67 insertions(+), 69 deletions(-) diff --git a/superset-frontend/src/components/StreamingExportModal/useStreamingExport.test.ts b/superset-frontend/src/components/StreamingExportModal/useStreamingExport.test.ts index 268da496ea74..583190d78741 100644 --- a/superset-frontend/src/components/StreamingExportModal/useStreamingExport.test.ts +++ b/superset-frontend/src/components/StreamingExportModal/useStreamingExport.test.ts @@ -35,96 +35,94 @@ global.URL.revokeObjectURL = jest.fn(); // Mock fetch global.fetch = jest.fn(); -describe('useStreamingExport', () => { - beforeEach(() => { - jest.clearAllMocks(); - }); +beforeEach(() => { + jest.clearAllMocks(); +}); - test('initializes with default progress state', () => { - const { result } = renderHook(() => useStreamingExport()); - - expect(result.current.progress).toEqual({ - rowsProcessed: 0, - totalRows: undefined, - totalSize: 0, - speed: 0, - mbPerSecond: 0, - elapsedTime: 0, - status: ExportStatus.STREAMING, - }); +test('useStreamingExport initializes with default progress state', () => { + const { result } = renderHook(() => useStreamingExport()); + + expect(result.current.progress).toEqual({ + rowsProcessed: 0, + totalRows: undefined, + totalSize: 0, + speed: 0, + mbPerSecond: 0, + elapsedTime: 0, + status: ExportStatus.STREAMING, }); +}); - test('provides startExport function', () => { - const { result } = renderHook(() => useStreamingExport()); +test('useStreamingExport provides startExport function', () => { + const { result } = renderHook(() => useStreamingExport()); - expect(typeof result.current.startExport).toBe('function'); - }); + expect(typeof result.current.startExport).toBe('function'); +}); - test('provides resetExport function', () => { - const { result } = renderHook(() => useStreamingExport()); +test('useStreamingExport provides resetExport function', () => { + const { result } = renderHook(() => useStreamingExport()); - expect(typeof result.current.resetExport).toBe('function'); - }); - - test('provides retryExport function', () => { - const { result } = renderHook(() => useStreamingExport()); + expect(typeof result.current.resetExport).toBe('function'); +}); - expect(typeof result.current.retryExport).toBe('function'); - }); +test('useStreamingExport provides retryExport function', () => { + const { result } = renderHook(() => useStreamingExport()); - test('provides cancelExport function', () => { - const { result } = renderHook(() => useStreamingExport()); + expect(typeof result.current.retryExport).toBe('function'); +}); - expect(typeof result.current.cancelExport).toBe('function'); - }); +test('useStreamingExport provides cancelExport function', () => { + const { result } = renderHook(() => useStreamingExport()); - test('resetExport resets progress to initial state', () => { - const { result } = renderHook(() => useStreamingExport()); + expect(typeof result.current.cancelExport).toBe('function'); +}); - act(() => { - result.current.resetExport(); - }); +test('useStreamingExport resetExport resets progress to initial state', () => { + const { result } = renderHook(() => useStreamingExport()); - expect(result.current.progress.status).toBe(ExportStatus.STREAMING); - expect(result.current.progress.rowsProcessed).toBe(0); - expect(result.current.progress.totalSize).toBe(0); + act(() => { + result.current.resetExport(); }); - test('accepts onComplete callback option', () => { - const onComplete = jest.fn(); - const { result } = renderHook(() => useStreamingExport({ onComplete })); + expect(result.current.progress.status).toBe(ExportStatus.STREAMING); + expect(result.current.progress.rowsProcessed).toBe(0); + expect(result.current.progress.totalSize).toBe(0); +}); - expect(result.current).toBeDefined(); - }); +test('useStreamingExport accepts onComplete callback option', () => { + const onComplete = jest.fn(); + const { result } = renderHook(() => useStreamingExport({ onComplete })); - test('accepts onError callback option', () => { - const onError = jest.fn(); - const { result } = renderHook(() => useStreamingExport({ onError })); + expect(result.current).toBeDefined(); +}); - expect(result.current).toBeDefined(); - }); +test('useStreamingExport accepts onError callback option', () => { + const onError = jest.fn(); + const { result } = renderHook(() => useStreamingExport({ onError })); - test('progress includes all required fields', () => { - const { result } = renderHook(() => useStreamingExport()); + expect(result.current).toBeDefined(); +}); - expect(result.current.progress).toHaveProperty('rowsProcessed'); - expect(result.current.progress).toHaveProperty('totalRows'); - expect(result.current.progress).toHaveProperty('totalSize'); - expect(result.current.progress).toHaveProperty('status'); - expect(result.current.progress).toHaveProperty('speed'); - expect(result.current.progress).toHaveProperty('mbPerSecond'); - expect(result.current.progress).toHaveProperty('elapsedTime'); - }); +test('useStreamingExport progress includes all required fields', () => { + const { result } = renderHook(() => useStreamingExport()); - test('cleans up on unmount', () => { - const revokeObjectURL = jest.fn(); - global.URL.revokeObjectURL = revokeObjectURL; + expect(result.current.progress).toHaveProperty('rowsProcessed'); + expect(result.current.progress).toHaveProperty('totalRows'); + expect(result.current.progress).toHaveProperty('totalSize'); + expect(result.current.progress).toHaveProperty('status'); + expect(result.current.progress).toHaveProperty('speed'); + expect(result.current.progress).toHaveProperty('mbPerSecond'); + expect(result.current.progress).toHaveProperty('elapsedTime'); +}); - const { unmount } = renderHook(() => useStreamingExport()); +test('useStreamingExport cleans up on unmount', () => { + const revokeObjectURL = jest.fn(); + global.URL.revokeObjectURL = revokeObjectURL; - unmount(); + const { unmount } = renderHook(() => useStreamingExport()); - // Cleanup should not throw errors - expect(true).toBe(true); - }); + unmount(); + + // Cleanup should not throw errors + expect(true).toBe(true); }); From 4b49a3228df67fafd994eeeb6ab0fc0f53c508c4 Mon Sep 17 00:00:00 2001 From: amaannawab923 Date: Sat, 11 Oct 2025 12:39:25 +0530 Subject: [PATCH 33/57] fix result format --- .../components/StreamingExportModal/useStreamingExport.ts | 8 +++++--- superset-frontend/src/explore/exploreUtils/index.js | 3 ++- 2 files changed, 7 insertions(+), 4 deletions(-) diff --git a/superset-frontend/src/components/StreamingExportModal/useStreamingExport.ts b/superset-frontend/src/components/StreamingExportModal/useStreamingExport.ts index efd723fed6bc..0f49e77bdb8a 100644 --- a/superset-frontend/src/components/StreamingExportModal/useStreamingExport.ts +++ b/superset-frontend/src/components/StreamingExportModal/useStreamingExport.ts @@ -127,6 +127,10 @@ export const useStreamingExport = (options: UseStreamingExportOptions = {}) => { const executeExport = useCallback( async (params: StreamingExportParams) => { const { url, payload, filename, exportType, expectedRows } = params; + if (isExportingRef.current) { + return; + } + isExportingRef.current = true; abortControllerRef.current = new AbortController(); @@ -153,7 +157,6 @@ export const useStreamingExport = (options: UseStreamingExportOptions = {}) => { expectedRows, abortControllerRef.current.signal, ); - const response = await fetch(url, fetchOptions); if (!response.ok) { @@ -241,6 +244,7 @@ export const useStreamingExport = (options: UseStreamingExportOptions = {}) => { const downloadUrl = URL.createObjectURL(blob); currentBlobUrlRef.current = downloadUrl; + updateProgress({ status: ExportStatus.COMPLETED, downloadUrl, @@ -282,7 +286,6 @@ export const useStreamingExport = (options: UseStreamingExportOptions = {}) => { return; } - isExportingRef.current = true; setRetryCount(0); lastExportParamsRef.current = params; @@ -311,7 +314,6 @@ export const useStreamingExport = (options: UseStreamingExportOptions = {}) => { return; } - isExportingRef.current = true; setRetryCount(0); executeExport(lastExportParamsRef.current); }, [executeExport]); diff --git a/superset-frontend/src/explore/exploreUtils/index.js b/superset-frontend/src/explore/exploreUtils/index.js index 396e02d2b557..1ebde5a3783b 100644 --- a/superset-frontend/src/explore/exploreUtils/index.js +++ b/superset-frontend/src/explore/exploreUtils/index.js @@ -275,10 +275,11 @@ export const exportChart = async ({ // Check if streaming export handler is provided (from dashboard Chart.jsx) if (onStartStreamingExport) { - // Streaming is handled by the caller - just pass URL and payload + // Streaming is handled by the caller - pass URL, payload, and export type onStartStreamingExport({ url, payload, + exportType: resultFormat, }); } else { // Fallback to original behavior for non-streaming exports From f27566c8d0d64bb9fa85701f893bc6507d4e9762 Mon Sep 17 00:00:00 2001 From: amaannawab923 Date: Tue, 14 Oct 2025 13:29:57 +0530 Subject: [PATCH 34/57] fix: eslint error --- .../src/components/StreamingExportModal/useStreamingExport.ts | 1 - 1 file changed, 1 deletion(-) diff --git a/superset-frontend/src/components/StreamingExportModal/useStreamingExport.ts b/superset-frontend/src/components/StreamingExportModal/useStreamingExport.ts index 0f49e77bdb8a..ea9cfc6a5f4e 100644 --- a/superset-frontend/src/components/StreamingExportModal/useStreamingExport.ts +++ b/superset-frontend/src/components/StreamingExportModal/useStreamingExport.ts @@ -244,7 +244,6 @@ export const useStreamingExport = (options: UseStreamingExportOptions = {}) => { const downloadUrl = URL.createObjectURL(blob); currentBlobUrlRef.current = downloadUrl; - updateProgress({ status: ExportStatus.COMPLETED, downloadUrl, From ca0a51e628f780e7b4e2ac2020ad9d8a52b70297 Mon Sep 17 00:00:00 2001 From: amaannawab923 Date: Tue, 14 Oct 2025 18:42:55 +0530 Subject: [PATCH 35/57] fix: comments resolution --- .../src/SqlLab/components/ResultSet/index.tsx | 20 +++++-------------- .../useStreamingExport.test.ts | 2 -- .../useStreamingExport.ts | 2 +- .../components/gridComponents/Chart/Chart.jsx | 1 - superset/charts/data/api.py | 8 ++++---- 5 files changed, 10 insertions(+), 23 deletions(-) diff --git a/superset-frontend/src/SqlLab/components/ResultSet/index.tsx b/superset-frontend/src/SqlLab/components/ResultSet/index.tsx index aac184baea3b..614575f5569d 100644 --- a/superset-frontend/src/SqlLab/components/ResultSet/index.tsx +++ b/superset-frontend/src/SqlLab/components/ResultSet/index.tsx @@ -234,12 +234,9 @@ const ResultSet = ({ 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 - }, + onComplete: () => {}, onError: error => { addDangerToast(t('Export failed: %s', error)); }, @@ -318,23 +315,18 @@ 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 = () => { const { rows, queryLimit, limitingFactor } = query; const limit = queryLimit || query.results?.query?.limit; const rowsCount = Math.min(rows || 0, query.results?.data?.length || 0); - // Determine actual row count let actualRowCount = rowsCount; - // If not limited by dropdown/query, use the full row count if (limitingFactor === LimitingFactor.NotLimited && rows) { actualRowCount = rows; } else if (limit) { @@ -395,11 +387,9 @@ const ResultSet = ({ css={copyButtonStyles} buttonSize="small" buttonStyle="secondary" - href={ - !shouldUseStreamingExport() - ? getExportCsvUrl(query.id) - : undefined - } + {...(!shouldUseStreamingExport() && { + href: getExportCsvUrl(query.id), + })} data-test="export-csv-button" onClick={e => { const useStreaming = shouldUseStreamingExport(); @@ -416,7 +406,7 @@ const ResultSet = ({ const filename = `sqllab_${query.id}_${timestamp}.csv`; startExport({ - url: getStreamingExportUrl(), + url: '/api/v1/sqllab/export_streaming/', payload: { client_id: query.id }, filename, exportType: 'csv', diff --git a/superset-frontend/src/components/StreamingExportModal/useStreamingExport.test.ts b/superset-frontend/src/components/StreamingExportModal/useStreamingExport.test.ts index 583190d78741..03241082dd1e 100644 --- a/superset-frontend/src/components/StreamingExportModal/useStreamingExport.test.ts +++ b/superset-frontend/src/components/StreamingExportModal/useStreamingExport.test.ts @@ -28,11 +28,9 @@ jest.mock('@superset-ui/core', () => ({ }, })); -// Mock URL APIs global.URL.createObjectURL = jest.fn(() => 'blob:mock-url'); global.URL.revokeObjectURL = jest.fn(); -// Mock fetch global.fetch = jest.fn(); beforeEach(() => { diff --git a/superset-frontend/src/components/StreamingExportModal/useStreamingExport.ts b/superset-frontend/src/components/StreamingExportModal/useStreamingExport.ts index ea9cfc6a5f4e..da31a7dbba9f 100644 --- a/superset-frontend/src/components/StreamingExportModal/useStreamingExport.ts +++ b/superset-frontend/src/components/StreamingExportModal/useStreamingExport.ts @@ -26,7 +26,7 @@ interface UseStreamingExportOptions { } interface StreamingExportPayload { - [key: string]: unknown; + [key: string]: any; } interface StreamingExportParams { diff --git a/superset-frontend/src/dashboard/components/gridComponents/Chart/Chart.jsx b/superset-frontend/src/dashboard/components/gridComponents/Chart/Chart.jsx index fe2b4a6396f7..f8aa6d6422e7 100644 --- a/superset-frontend/src/dashboard/components/gridComponents/Chart/Chart.jsx +++ b/superset-frontend/src/dashboard/components/gridComponents/Chart/Chart.jsx @@ -617,7 +617,6 @@ const Chart = props => { /> - {/* Streaming Export Modal */} { diff --git a/superset/charts/data/api.py b/superset/charts/data/api.py index 9df52b6f491d..199fc141095e 100644 --- a/superset/charts/data/api.py +++ b/superset/charts/data/api.py @@ -449,15 +449,15 @@ 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) + 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) + logger.info("FRONTEND PROVIDED EXPECTED ROWS: %d", expected_rows) except (ValueError, TypeError): - logger.warning("โš ๏ธ Invalid expected_rows value: %s", expected_rows_str) + logger.warning("Invalid expected_rows value: %s", expected_rows_str) return filename, expected_rows @@ -579,7 +579,7 @@ def _create_streaming_csv_response( filename is not None, ) if expected_rows: - logger.info("๐Ÿ“Š Using expected_rows from frontend: %d", expected_rows) + logger.info("Using expected_rows from frontend: %d", expected_rows) # Execute streaming command chunk_size = 1000 From edcdb2017df2db00950197c297c91159561ae455 Mon Sep 17 00:00:00 2001 From: amaannawab923 Date: Wed, 15 Oct 2025 19:08:58 +0530 Subject: [PATCH 36/57] Unit tests for streaming commands --- tests/unit_tests/commands/chart/__init__.py | 16 + .../chart/streaming_export_command_test.py | 270 +++++++++ tests/unit_tests/commands/sql_lab/__init__.py | 16 + .../sql_lab/streaming_export_command_test.py | 537 ++++++++++++++++++ 4 files changed, 839 insertions(+) create mode 100644 tests/unit_tests/commands/chart/__init__.py create mode 100644 tests/unit_tests/commands/chart/streaming_export_command_test.py create mode 100644 tests/unit_tests/commands/sql_lab/__init__.py create mode 100644 tests/unit_tests/commands/sql_lab/streaming_export_command_test.py diff --git a/tests/unit_tests/commands/chart/__init__.py b/tests/unit_tests/commands/chart/__init__.py new file mode 100644 index 000000000000..13a83393a912 --- /dev/null +++ b/tests/unit_tests/commands/chart/__init__.py @@ -0,0 +1,16 @@ +# 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. diff --git a/tests/unit_tests/commands/chart/streaming_export_command_test.py b/tests/unit_tests/commands/chart/streaming_export_command_test.py new file mode 100644 index 000000000000..a88f074f3d87 --- /dev/null +++ b/tests/unit_tests/commands/chart/streaming_export_command_test.py @@ -0,0 +1,270 @@ +# 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. +"""Unit tests for Chart Streaming CSV Export Command.""" + +import pytest +from pytest_mock import MockerFixture + +from superset.commands.chart.data.streaming_export_command import ( + StreamingCSVExportCommand, +) + + +def test_streaming_csv_export_command_init(mocker: MockerFixture) -> None: + """Test command initialization.""" + query_context = mocker.MagicMock() + command = StreamingCSVExportCommand(query_context, chunk_size=500) + + assert command._query_context == query_context + assert command._chunk_size == 500 + assert command._current_app is not None + + +def test_streaming_csv_export_command_default_chunk_size( + mocker: MockerFixture, +) -> None: + """Test command uses default chunk size.""" + query_context = mocker.MagicMock() + command = StreamingCSVExportCommand(query_context) + + assert command._chunk_size == 1000 + + +def test_validate_calls_raise_for_access(mocker: MockerFixture) -> None: + """Test validate method calls query context raise_for_access.""" + query_context = mocker.MagicMock() + command = StreamingCSVExportCommand(query_context) + + command.validate() + + query_context.raise_for_access.assert_called_once() + + +def test_validate_raises_exception_on_access_denied(mocker: MockerFixture) -> None: + """Test validate raises exception when access is denied.""" + query_context = mocker.MagicMock() + query_context.raise_for_access.side_effect = Exception("Access denied") + command = StreamingCSVExportCommand(query_context) + + with pytest.raises(Exception, match="Access denied"): + command.validate() + + +def test_csv_generation_with_small_dataset(mocker: MockerFixture) -> None: + """Test CSV generation with a small dataset.""" + mock_db = mocker.patch("superset.db") + mock_session = mocker.MagicMock() + mock_db.session.return_value.__enter__.return_value = mock_session + + query_context = mocker.MagicMock() + datasource = mocker.MagicMock() + datasource.get_query_str.return_value = "SELECT * FROM test" + query_context.datasource = datasource + query_context.queries = [mocker.MagicMock()] + mock_session.merge.return_value = datasource + + mock_result_proxy = mocker.MagicMock() + mock_result_proxy.keys.return_value = ["col1", "col2", "col3"] + mock_result_proxy.fetchmany.side_effect = [ + [ + ("row1_val1", "row1_val2", "row1_val3"), + ("row2_val1", "row2_val2", "row2_val3"), + ], + [("row3_val1", "row3_val2", "row3_val3")], + [], + ] + + mock_connection = mocker.MagicMock() + mock_connection.execution_options.return_value.execute.return_value = ( + mock_result_proxy + ) + mock_engine = mocker.MagicMock() + mock_engine.connect.return_value = mock_connection + datasource.database.get_sqla_engine.return_value.__enter__.return_value = ( + mock_engine + ) + + command = StreamingCSVExportCommand(query_context, chunk_size=2) + csv_generator_callable = command.run() + generator = csv_generator_callable() + + chunks = list(generator) + + csv_data = "".join(chunks) + lines = [line.strip() for line in csv_data.strip().split("\n")] + + assert len(lines) == 4 + assert lines[0] == "col1,col2,col3" + assert "row1_val1,row1_val2,row1_val3" in csv_data + assert "row2_val1,row2_val2,row2_val3" in csv_data + assert "row3_val1,row3_val2,row3_val3" in csv_data + + +def test_csv_generation_with_special_characters(mocker: MockerFixture) -> None: + """Test CSV generation properly escapes special characters.""" + mock_db = mocker.patch("superset.db") + mock_session = mocker.MagicMock() + mock_db.session.return_value.__enter__.return_value = mock_session + + query_context = mocker.MagicMock() + datasource = mocker.MagicMock() + datasource.get_query_str.return_value = "SELECT * FROM test" + query_context.datasource = datasource + query_context.queries = [mocker.MagicMock()] + mock_session.merge.return_value = datasource + + mock_result = mocker.MagicMock() + mock_result.keys.return_value = ["name", "description"] + mock_result.fetchmany.side_effect = [ + [("John, Jr.", 'Quote"Test'), ("Line\nBreak", "Comma,Value")], + [], + ] + + mock_connection = mocker.MagicMock() + mock_connection.execution_options.return_value.execute.return_value = mock_result + mock_engine = mocker.MagicMock() + mock_engine.connect.return_value = mock_connection + datasource.database.get_sqla_engine.return_value.__enter__.return_value = ( + mock_engine + ) + + command = StreamingCSVExportCommand(query_context, chunk_size=10) + csv_generator_callable = command.run() + generator = csv_generator_callable() + csv_data = "".join(generator) + + assert '"John, Jr."' in csv_data + assert '"Quote""Test"' in csv_data + assert "Line\nBreak" in csv_data + assert '"Comma,Value"' in csv_data + + +def test_streaming_with_null_values(mocker: MockerFixture) -> None: + """Test CSV generation handles NULL values correctly.""" + mock_db = mocker.patch("superset.db") + mock_session = mocker.MagicMock() + mock_db.session.return_value.__enter__.return_value = mock_session + + query_context = mocker.MagicMock() + datasource = mocker.MagicMock() + datasource.get_query_str.return_value = "SELECT * FROM test" + query_context.datasource = datasource + query_context.queries = [mocker.MagicMock()] + mock_session.merge.return_value = datasource + + mock_result = mocker.MagicMock() + mock_result.keys.return_value = ["col1", "col2", "col3"] + mock_result.fetchmany.side_effect = [ + [("value1", None, "value3"), (None, "value2", None)], + [], + ] + + mock_connection = mocker.MagicMock() + mock_connection.execution_options.return_value.execute.return_value = mock_result + mock_engine = mocker.MagicMock() + mock_engine.connect.return_value = mock_connection + datasource.database.get_sqla_engine.return_value.__enter__.return_value = ( + mock_engine + ) + + command = StreamingCSVExportCommand(query_context, chunk_size=10) + csv_generator_callable = command.run() + generator = csv_generator_callable() + csv_data = "".join(generator) + + lines = csv_data.strip().split("\n") + assert len(lines) == 3 + assert "value1,,value3" in csv_data + assert ",value2," in csv_data + + +def test_streaming_execution_options_enabled(mocker: MockerFixture) -> None: + """Test that streaming execution options are enabled.""" + mock_db = mocker.patch("superset.db") + mock_session = mocker.MagicMock() + mock_db.session.return_value.__enter__.return_value = mock_session + + query_context = mocker.MagicMock() + datasource = mocker.MagicMock() + datasource.get_query_str.return_value = "SELECT * FROM test" + query_context.datasource = datasource + query_context.queries = [mocker.MagicMock()] + mock_session.merge.return_value = datasource + + mock_result_proxy = mocker.MagicMock() + mock_result_proxy.keys.return_value = ["col1", "col2", "col3"] + mock_result_proxy.fetchmany.side_effect = [ + [ + ("row1_val1", "row1_val2", "row1_val3"), + ("row2_val1", "row2_val2", "row2_val3"), + ], + [("row3_val1", "row3_val2", "row3_val3")], + [], + ] + + mock_connection = mocker.MagicMock() + mock_execution_options = mocker.MagicMock() + mock_connection.execution_options.return_value = mock_execution_options + mock_execution_options.execute.return_value = mock_result_proxy + + mock_engine = mocker.MagicMock() + mock_engine.connect.return_value = mock_connection + datasource.database.get_sqla_engine.return_value.__enter__.return_value = ( + mock_engine + ) + + command = StreamingCSVExportCommand(query_context) + csv_generator_callable = command.run() + generator = csv_generator_callable() + list(generator) + + mock_connection.execution_options.assert_called_once_with(stream_results=True) + + +def test_empty_result_set(mocker: MockerFixture) -> None: + """Test CSV generation with empty result set.""" + mock_db = mocker.patch("superset.db") + mock_session = mocker.MagicMock() + mock_db.session.return_value.__enter__.return_value = mock_session + + query_context = mocker.MagicMock() + datasource = mocker.MagicMock() + datasource.get_query_str.return_value = "SELECT * FROM test" + query_context.datasource = datasource + query_context.queries = [mocker.MagicMock()] + mock_session.merge.return_value = datasource + + mock_result = mocker.MagicMock() + mock_result.keys.return_value = ["col1", "col2"] + mock_result.fetchmany.side_effect = [[]] + + mock_connection = mocker.MagicMock() + mock_connection.execution_options.return_value.execute.return_value = mock_result + mock_engine = mocker.MagicMock() + mock_engine.connect.return_value = mock_connection + datasource.database.get_sqla_engine.return_value.__enter__.return_value = ( + mock_engine + ) + + command = StreamingCSVExportCommand(query_context) + csv_generator_callable = command.run() + generator = csv_generator_callable() + csv_data = "".join(generator) + + lines = [line.strip() for line in csv_data.strip().split("\n")] + assert len(lines) == 1 + assert lines[0] == "col1,col2" diff --git a/tests/unit_tests/commands/sql_lab/__init__.py b/tests/unit_tests/commands/sql_lab/__init__.py new file mode 100644 index 000000000000..13a83393a912 --- /dev/null +++ b/tests/unit_tests/commands/sql_lab/__init__.py @@ -0,0 +1,16 @@ +# 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. diff --git a/tests/unit_tests/commands/sql_lab/streaming_export_command_test.py b/tests/unit_tests/commands/sql_lab/streaming_export_command_test.py new file mode 100644 index 000000000000..c0fc82ed04f8 --- /dev/null +++ b/tests/unit_tests/commands/sql_lab/streaming_export_command_test.py @@ -0,0 +1,537 @@ +# 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. +"""Unit tests for SQL Lab Streaming CSV Export Command.""" + +from unittest.mock import MagicMock, Mock, patch + +import pytest + +from superset.commands.sql_lab.streaming_export_command import ( + StreamingSqlResultExportCommand, +) +from superset.errors import SupersetErrorType +from superset.exceptions import SupersetErrorException, SupersetSecurityException +from superset.sqllab.limiting_factor import LimitingFactor + + +@pytest.fixture +def mock_query(): + """Create a mock SQL Lab query.""" + query = MagicMock() + query.client_id = "test_client_123" + query.select_sql = None + query.executed_sql = "SELECT * FROM test_table" + query.limiting_factor = LimitingFactor.NOT_LIMITED + query.database = MagicMock() + query.database.db_engine_spec = MagicMock() + query.database.db_engine_spec.engine = "postgresql" + query.raise_for_access = MagicMock() + return query + + +@pytest.fixture +def mock_result_proxy(): + """Create a mock database result proxy.""" + result = MagicMock() + result.keys.return_value = ["id", "name", "value"] + result.fetchmany.side_effect = [ + [(1, "test1", 100), (2, "test2", 200)], + [(3, "test3", 300)], + [], + ] + return result + + +def test_streaming_sql_result_export_command_init(): + """Test command initialization.""" + command = StreamingSqlResultExportCommand("client_123", chunk_size=500) + + assert command._client_id == "client_123" + assert command._chunk_size == 500 + assert command._query is None + assert command._current_app is not None + + +def test_streaming_sql_result_export_command_default_chunk_size(): + """Test command uses default chunk size.""" + command = StreamingSqlResultExportCommand("client_123") + + assert command._chunk_size == 1000 + + +@patch("superset.commands.sql_lab.streaming_export_command.db") +def test_validate_query_not_found(mock_db): + """Test validate raises exception when query is not found.""" + mock_query_result = mock_db.session.query.return_value.filter_by.return_value + mock_query_result.one_or_none.return_value = None + + command = StreamingSqlResultExportCommand("nonexistent_client") + + with pytest.raises(SupersetErrorException) as exc_info: + command.validate() + + assert exc_info.value.error.error_type == SupersetErrorType.RESULTS_BACKEND_ERROR + assert exc_info.value.status == 404 + + +@patch("superset.commands.sql_lab.streaming_export_command.db") +def test_validate_access_denied(mock_db, mock_query): + """Test validate raises exception when access is denied.""" + mock_query_result = mock_db.session.query.return_value.filter_by.return_value + mock_query_result.one_or_none.return_value = mock_query + mock_query.raise_for_access.side_effect = SupersetSecurityException( + Mock(message="Access denied") + ) + + command = StreamingSqlResultExportCommand("test_client_123") + + with pytest.raises(SupersetErrorException) as exc_info: + command.validate() + + assert ( + exc_info.value.error.error_type == SupersetErrorType.QUERY_SECURITY_ACCESS_ERROR + ) + assert exc_info.value.status == 403 + + +@patch("superset.commands.sql_lab.streaming_export_command.db") +def test_validate_success(mock_db, mock_query): + """Test successful validation.""" + mock_query_result = mock_db.session.query.return_value.filter_by.return_value + mock_query_result.one_or_none.return_value = mock_query + + command = StreamingSqlResultExportCommand("test_client_123") + command.validate() + + assert command._query == mock_query + mock_query.raise_for_access.assert_called_once() + + +@patch("superset.commands.sql_lab.streaming_export_command.db") +def test_csv_generation_with_select_sql(mock_db, mock_query, mock_result_proxy): + """Test CSV generation when query has select_sql.""" + mock_query.select_sql = "SELECT * FROM test WHERE id > 0" + mock_query.executed_sql = None + + mock_session = MagicMock() + mock_db.session.return_value.__enter__.return_value = mock_session + mock_session.merge.return_value = mock_query.database + + mock_connection = MagicMock() + mock_connection.execution_options.return_value.execute.return_value = ( + mock_result_proxy + ) + mock_engine = MagicMock() + mock_engine.connect.return_value = mock_connection + mock_query.database.get_sqla_engine.return_value.__enter__.return_value = ( + mock_engine + ) + + mock_query_result = mock_db.session.query.return_value.filter_by.return_value + mock_query_result.one_or_none.return_value = mock_query + command = StreamingSqlResultExportCommand("test_client_123", chunk_size=2) + command.validate() + + csv_generator_callable = command.run() + generator = csv_generator_callable() + chunks = list(generator) + + csv_data = "".join(chunks) + lines = [line.strip() for line in csv_data.strip().split("\n")] + + assert len(lines) == 4 + assert lines[0] == "id,name,value" + assert "1,test1,100" in csv_data + assert "2,test2,200" in csv_data + assert "3,test3,300" in csv_data + + +@patch("superset.commands.sql_lab.streaming_export_command.db") +@patch("superset.commands.sql_lab.streaming_export_command.SQLScript") +def test_csv_generation_with_executed_sql_and_limit( + mock_sqlscript, mock_db, mock_query, mock_result_proxy +): + """Test CSV generation with executed_sql and applies limit.""" + mock_query.select_sql = None + mock_query.executed_sql = "SELECT * FROM test LIMIT 2" + mock_query.limiting_factor = LimitingFactor.QUERY + + mock_statement = Mock() + mock_statement.get_limit_value.return_value = 3 + mock_script_instance = Mock() + mock_script_instance.statements = [mock_statement] + mock_sqlscript.return_value = mock_script_instance + + mock_session = MagicMock() + mock_db.session.return_value.__enter__.return_value = mock_session + mock_session.merge.return_value = mock_query.database + + mock_result = MagicMock() + mock_result.keys.return_value = ["id", "name"] + mock_result.fetchmany.side_effect = [ + [(1, "test1"), (2, "test2"), (3, "test3"), (4, "test4")], + [], + ] + + mock_connection = MagicMock() + mock_connection.execution_options.return_value.execute.return_value = mock_result + mock_engine = MagicMock() + mock_engine.connect.return_value = mock_connection + mock_query.database.get_sqla_engine.return_value.__enter__.return_value = ( + mock_engine + ) + + mock_query_result = mock_db.session.query.return_value.filter_by.return_value + mock_query_result.one_or_none.return_value = mock_query + command = StreamingSqlResultExportCommand("test_client_123", chunk_size=10) + command.validate() + + csv_generator_callable = command.run() + generator = csv_generator_callable() + csv_data = "".join(generator) + + lines = [line.strip() for line in csv_data.strip().split("\n")] + assert len(lines) == 3 # header + 2 rows (limit - 1) + + +@patch("superset.commands.sql_lab.streaming_export_command.db") +def test_csv_generation_with_special_characters(mock_db, mock_query): + """Test CSV generation properly escapes special characters.""" + mock_query.select_sql = "SELECT * FROM test" + + mock_result = MagicMock() + mock_result.keys.return_value = ["text", "description"] + mock_result.fetchmany.side_effect = [ + [('Text with "quotes"', "Line\nbreak"), ("Comma,value", "Tab\tchar")], + [], + ] + + mock_session = MagicMock() + mock_db.session.return_value.__enter__.return_value = mock_session + mock_session.merge.return_value = mock_query.database + + mock_connection = MagicMock() + mock_connection.execution_options.return_value.execute.return_value = mock_result + mock_engine = MagicMock() + mock_engine.connect.return_value = mock_connection + mock_query.database.get_sqla_engine.return_value.__enter__.return_value = ( + mock_engine + ) + + mock_query_result = mock_db.session.query.return_value.filter_by.return_value + mock_query_result.one_or_none.return_value = mock_query + command = StreamingSqlResultExportCommand("test_client_123") + command.validate() + + csv_generator_callable = command.run() + generator = csv_generator_callable() + csv_data = "".join(generator) + + assert '"Text with ""quotes"""' in csv_data # Quotes doubled + assert "Line\nbreak" in csv_data + assert '"Comma,value"' in csv_data + assert "Tab\tchar" in csv_data + + +@patch("superset.commands.sql_lab.streaming_export_command.db") +def test_limiting_factor_dropdown(mock_db, mock_query): + """Test limit adjustment for DROPDOWN limiting factor.""" + mock_query.select_sql = None + mock_query.executed_sql = "SELECT * FROM test LIMIT 101" + mock_query.limiting_factor = LimitingFactor.DROPDOWN + + with patch( + "superset.commands.sql_lab.streaming_export_command.SQLScript" + ) as mock_sqlscript: + mock_statement = Mock() + mock_statement.get_limit_value.return_value = 101 + mock_script_instance = Mock() + mock_script_instance.statements = [mock_statement] + mock_sqlscript.return_value = mock_script_instance + + mock_result = MagicMock() + mock_result.keys.return_value = ["id"] + mock_result.fetchmany.side_effect = [[(i,) for i in range(101)], []] + + mock_session = MagicMock() + mock_db.session.return_value.__enter__.return_value = mock_session + mock_session.merge.return_value = mock_query.database + + mock_connection = MagicMock() + mock_connection.execution_options.return_value.execute.return_value = ( + mock_result + ) + mock_engine = MagicMock() + mock_engine.connect.return_value = mock_connection + mock_query.database.get_sqla_engine.return_value.__enter__.return_value = ( + mock_engine + ) + + mock_query_result = mock_db.session.query.return_value.filter_by.return_value + mock_query_result.one_or_none.return_value = mock_query + command = StreamingSqlResultExportCommand("test_client_123", chunk_size=200) + command.validate() + + csv_generator_callable = command.run() + generator = csv_generator_callable() + csv_data = "".join(generator) + + lines = [line.strip() for line in csv_data.strip().split("\n")] + assert len(lines) == 101 + + +@patch("superset.commands.sql_lab.streaming_export_command.db") +def test_limiting_factor_query_and_dropdown(mock_db, mock_query): + """Test limit adjustment for QUERY_AND_DROPDOWN limiting factor.""" + mock_query.select_sql = None + mock_query.executed_sql = "SELECT * FROM test LIMIT 51" + mock_query.limiting_factor = LimitingFactor.QUERY_AND_DROPDOWN + + with patch( + "superset.commands.sql_lab.streaming_export_command.SQLScript" + ) as mock_sqlscript: + mock_statement = Mock() + mock_statement.get_limit_value.return_value = 51 + mock_script_instance = Mock() + mock_script_instance.statements = [mock_statement] + mock_sqlscript.return_value = mock_script_instance + + mock_result = MagicMock() + mock_result.keys.return_value = ["id"] + mock_result.fetchmany.side_effect = [[(i,) for i in range(51)], []] + + mock_session = MagicMock() + mock_db.session.return_value.__enter__.return_value = mock_session + mock_session.merge.return_value = mock_query.database + + mock_connection = MagicMock() + mock_connection.execution_options.return_value.execute.return_value = ( + mock_result + ) + mock_engine = MagicMock() + mock_engine.connect.return_value = mock_connection + mock_query.database.get_sqla_engine.return_value.__enter__.return_value = ( + mock_engine + ) + + mock_query_result = mock_db.session.query.return_value.filter_by.return_value + mock_query_result.one_or_none.return_value = mock_query + command = StreamingSqlResultExportCommand("test_client_123", chunk_size=100) + command.validate() + + csv_generator_callable = command.run() + generator = csv_generator_callable() + csv_data = "".join(generator) + + lines = [line.strip() for line in csv_data.strip().split("\n")] + assert len(lines) == 51 + + +@patch("superset.commands.sql_lab.streaming_export_command.db") +def test_empty_result_set(mock_db, mock_query): + """Test CSV generation with empty result set.""" + mock_query.select_sql = "SELECT * FROM empty_table" + + mock_result = MagicMock() + mock_result.keys.return_value = ["col1", "col2"] + mock_result.fetchmany.side_effect = [[]] + + mock_session = MagicMock() + mock_db.session.return_value.__enter__.return_value = mock_session + mock_session.merge.return_value = mock_query.database + + mock_connection = MagicMock() + mock_connection.execution_options.return_value.execute.return_value = mock_result + mock_engine = MagicMock() + mock_engine.connect.return_value = mock_connection + mock_query.database.get_sqla_engine.return_value.__enter__.return_value = ( + mock_engine + ) + + mock_query_result = mock_db.session.query.return_value.filter_by.return_value + mock_query_result.one_or_none.return_value = mock_query + command = StreamingSqlResultExportCommand("test_client_123") + command.validate() + + csv_generator_callable = command.run() + generator = csv_generator_callable() + csv_data = "".join(generator) + + lines = [line.strip() for line in csv_data.strip().split("\n")] + assert len(lines) == 1 + assert lines[0] == "col1,col2" + + +@patch("superset.commands.sql_lab.streaming_export_command.db") +def test_error_handling_yields_error_marker(mock_db, mock_query): + """Test that exceptions are caught and error marker is yielded.""" + mock_query.select_sql = "SELECT * FROM test" + + mock_session = MagicMock() + mock_db.session.return_value.__enter__.return_value = mock_session + mock_session.merge.side_effect = Exception("Database connection failed") + + mock_query_result = mock_db.session.query.return_value.filter_by.return_value + mock_query_result.one_or_none.return_value = mock_query + command = StreamingSqlResultExportCommand("test_client_123") + command.validate() + + csv_generator_callable = command.run() + generator = csv_generator_callable() + chunks = list(generator) + + error_output = "".join(chunks) + assert "__STREAM_ERROR__" in error_output + assert "Export failed" in error_output + + +@patch("superset.commands.sql_lab.streaming_export_command.db") +def test_connection_is_closed_after_streaming(mock_db, mock_query, mock_result_proxy): + """Test that database connection is properly closed.""" + mock_query.select_sql = "SELECT * FROM test" + + mock_session = MagicMock() + mock_db.session.return_value.__enter__.return_value = mock_session + mock_session.merge.return_value = mock_query.database + + mock_connection = MagicMock() + mock_connection.execution_options.return_value.execute.return_value = ( + mock_result_proxy + ) + mock_engine = MagicMock() + mock_engine.connect.return_value = mock_connection + mock_query.database.get_sqla_engine.return_value.__enter__.return_value = ( + mock_engine + ) + + mock_query_result = mock_db.session.query.return_value.filter_by.return_value + mock_query_result.one_or_none.return_value = mock_query + command = StreamingSqlResultExportCommand("test_client_123") + command.validate() + + csv_generator_callable = command.run() + generator = csv_generator_callable() + list(generator) + + mock_connection.close.assert_called_once() + + +@patch("superset.commands.sql_lab.streaming_export_command.db") +def test_streaming_execution_options_enabled(mock_db, mock_query, mock_result_proxy): + """Test that streaming execution options are enabled.""" + mock_query.select_sql = "SELECT * FROM test" + + mock_session = MagicMock() + mock_db.session.return_value.__enter__.return_value = mock_session + mock_session.merge.return_value = mock_query.database + + mock_connection = MagicMock() + mock_execution_options = Mock() + mock_connection.execution_options.return_value = mock_execution_options + mock_execution_options.execute.return_value = mock_result_proxy + + mock_engine = MagicMock() + mock_engine.connect.return_value = mock_connection + mock_query.database.get_sqla_engine.return_value.__enter__.return_value = ( + mock_engine + ) + + mock_query_result = mock_db.session.query.return_value.filter_by.return_value + mock_query_result.one_or_none.return_value = mock_query + command = StreamingSqlResultExportCommand("test_client_123") + command.validate() + + csv_generator_callable = command.run() + generator = csv_generator_callable() + list(generator) + + mock_connection.execution_options.assert_called_once_with(stream_results=True) + + +@patch("superset.commands.sql_lab.streaming_export_command.db") +@patch("superset.commands.sql_lab.streaming_export_command.logger") +def test_completion_logging(mock_logger, mock_db, mock_query, mock_result_proxy): + """Test that completion is logged with metrics.""" + mock_query.select_sql = "SELECT * FROM test" + + mock_session = MagicMock() + mock_db.session.return_value.__enter__.return_value = mock_session + mock_session.merge.return_value = mock_query.database + + mock_connection = MagicMock() + mock_connection.execution_options.return_value.execute.return_value = ( + mock_result_proxy + ) + mock_engine = MagicMock() + mock_engine.connect.return_value = mock_connection + mock_query.database.get_sqla_engine.return_value.__enter__.return_value = ( + mock_engine + ) + + mock_query_result = mock_db.session.query.return_value.filter_by.return_value + mock_query_result.one_or_none.return_value = mock_query + command = StreamingSqlResultExportCommand("test_client_123") + command.validate() + + csv_generator_callable = command.run() + generator = csv_generator_callable() + list(generator) + + assert mock_logger.info.called + log_message = str(mock_logger.info.call_args) + assert "SQL Lab streaming CSV completed" in log_message + assert "rows" in log_message + + +@patch("superset.commands.sql_lab.streaming_export_command.db") +def test_null_values_handling(mock_db, mock_query): + """Test CSV generation handles NULL values correctly.""" + mock_query.select_sql = "SELECT * FROM test" + + mock_result = MagicMock() + mock_result.keys.return_value = ["id", "name", "value"] + mock_result.fetchmany.side_effect = [ + [(1, None, 100), (2, "test", None), (None, None, None)], + [], + ] + + mock_session = MagicMock() + mock_db.session.return_value.__enter__.return_value = mock_session + mock_session.merge.return_value = mock_query.database + + mock_connection = MagicMock() + mock_connection.execution_options.return_value.execute.return_value = mock_result + mock_engine = MagicMock() + mock_engine.connect.return_value = mock_connection + mock_query.database.get_sqla_engine.return_value.__enter__.return_value = ( + mock_engine + ) + + mock_query_result = mock_db.session.query.return_value.filter_by.return_value + mock_query_result.one_or_none.return_value = mock_query + command = StreamingSqlResultExportCommand("test_client_123") + command.validate() + + csv_generator_callable = command.run() + generator = csv_generator_callable() + csv_data = "".join(generator) + + lines = [line.strip() for line in csv_data.strip().split("\n")] + assert len(lines) == 4 + assert "1,,100" in csv_data + assert "2,test," in csv_data + assert ",," in csv_data From c070c697bd38b8e888f9021447ce618973ccf41a Mon Sep 17 00:00:00 2001 From: amaannawab923 Date: Wed, 15 Oct 2025 19:59:25 +0530 Subject: [PATCH 37/57] Refactor & comments resolution --- superset-frontend/webpack.proxy-config.js | 2 +- superset/charts/data/api.py | 23 +- .../chart/data/streaming_export_command.py | 202 +++++++++------- .../sql_lab/streaming_export_command.py | 228 ++++++++++-------- superset/sqllab/api.py | 18 +- 5 files changed, 263 insertions(+), 210 deletions(-) diff --git a/superset-frontend/webpack.proxy-config.js b/superset-frontend/webpack.proxy-config.js index 8fe035a1bb90..5011814be569 100644 --- a/superset-frontend/webpack.proxy-config.js +++ b/superset-frontend/webpack.proxy-config.js @@ -166,7 +166,7 @@ module.exports = newManifest => { processHTML(proxyResponse, response); } else { const isStreaming = - proxyResponse.headers['x-superset-streaming'] === 'true'; + proxyResponse.headers['transfer-encoding'] === 'chunked'; if (isStreaming) { proxyResponse.on('data', chunk => { diff --git a/superset/charts/data/api.py b/superset/charts/data/api.py index 199fc141095e..4d8252fb5450 100644 --- a/superset/charts/data/api.py +++ b/superset/charts/data/api.py @@ -18,12 +18,14 @@ import contextlib import logging +from datetime import datetime from typing import Any, TYPE_CHECKING from flask import current_app as app, g, make_response, request, Response from flask_appbuilder.api import expose, protect from flask_babel import gettext as _ from marshmallow import ValidationError +from werkzeug.utils import secure_filename from superset import is_feature_enabled, security_manager from superset.async_events.async_query_manager import AsyncQueryTokenException @@ -35,6 +37,9 @@ CreateAsyncChartDataJobCommand, ) from superset.commands.chart.data.get_data_command import ChartDataCommand +from superset.commands.chart.data.streaming_export_command import ( + StreamingCSVExportCommand, +) from superset.commands.chart.exceptions import ( ChartDataCacheLoadError, ChartDataQueryFailedError, @@ -501,8 +506,6 @@ 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 @@ -514,7 +517,7 @@ def _should_use_streaming( threshold = app.config.get("CSV_STREAMING_ROW_THRESHOLD", 100000) # Extract actual row count (same logic as frontend) - actual_row_count = None + actual_row_count: int | None = 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 @@ -547,14 +550,6 @@ def _create_streaming_csv_response( expected_rows: int | None = None, ) -> Response: """Create a streaming CSV response for large datasets.""" - from datetime import datetime - - from flask import Response - - from superset.commands.chart.data.streaming_export_command import ( - StreamingCSVExportCommand, - ) - query_context = result["query_context"] # Use filename from frontend if provided, otherwise generate one @@ -568,10 +563,7 @@ def _create_streaming_csv_response( chart_name = form_data["viz_type"] # Sanitize chart name for filename - safe_chart_name = "".join( - c for c in chart_name if c.isalnum() or c in ("-", "_") - ) - filename = f"superset_{safe_chart_name}_{timestamp}.csv" + filename = secure_filename(f"superset_{chart_name}_{timestamp}.csv") logger.info( "Creating streaming CSV response: %s (from frontend: %s)", @@ -600,7 +592,6 @@ def _create_streaming_csv_response( "Content-Disposition": f'attachment; filename="{filename}"', "Cache-Control": "no-cache", "X-Accel-Buffering": "no", # Disable nginx buffering - "X-Superset-Streaming": "true", # Identify streaming responses }, direct_passthrough=False, # Flask must iterate generator ) diff --git a/superset/commands/chart/data/streaming_export_command.py b/superset/commands/chart/data/streaming_export_command.py index 153e4d9cd18a..425058de25bc 100644 --- a/superset/commands/chart/data/streaming_export_command.py +++ b/superset/commands/chart/data/streaming_export_command.py @@ -22,9 +22,10 @@ import io import logging import time -from typing import Callable, Generator, TYPE_CHECKING +from typing import Any, Callable, Generator, TYPE_CHECKING from flask import current_app as app +from sqlalchemy import text from superset.commands.base import BaseCommand @@ -65,7 +66,119 @@ def validate(self) -> None: """Validate permissions and query context.""" self._query_context.raise_for_access() - def run(self) -> Callable[[], Generator[str, None, None]]: # noqa: C901 + def _prepare_datasource(self, session: Any) -> Any: + """Prepare and merge datasource with session.""" + from superset.connectors.sqla.models import SqlaTable + + datasource = self._query_context.datasource + if isinstance(datasource, SqlaTable): + datasource = session.merge(datasource) + return datasource + + def _get_sql_query(self, datasource: Any) -> str: + """Generate SQL query from datasource and query context.""" + query_obj = self._query_context.queries[0] + return datasource.get_query_str(query_obj.to_dict()) + + def _write_csv_header( + self, columns: list[str], csv_writer: Any, buffer: io.StringIO + ) -> tuple[str, int]: + """Write CSV header and return header data with byte count.""" + csv_writer.writerow(columns) + header_data = buffer.getvalue() + total_bytes = len(header_data.encode("utf-8")) + buffer.seek(0) + buffer.truncate() + return header_data, total_bytes + + def _process_rows( + self, + result_proxy: Any, + csv_writer: Any, + buffer: io.StringIO, + ) -> Generator[tuple[str, int, int], None, None]: + """ + Process database rows and yield CSV data chunks. + + Yields tuples of (data_chunk, row_count, byte_count). + """ + row_count = 0 + flush_threshold = 65536 # 64KB + + while rows := result_proxy.fetchmany(self._chunk_size): + for row in rows: + csv_writer.writerow(row) + row_count += 1 + + # Check buffer size and flush if needed + current_size = buffer.tell() + if current_size >= flush_threshold: + data = buffer.getvalue() + data_bytes = len(data.encode("utf-8")) + yield data, row_count, data_bytes + buffer.seek(0) + buffer.truncate() + + # Flush remaining buffer + if remaining_data := buffer.getvalue(): + data_bytes = len(remaining_data.encode("utf-8")) + yield remaining_data, row_count, data_bytes + + def _execute_query_and_stream(self) -> Generator[str, None, None]: + """Execute query with streaming and yield CSV chunks.""" + from superset import db + + start_time = time.time() + total_bytes = 0 + + with db.session() as session: + datasource = self._prepare_datasource(session) + sql_query = self._get_sql_query(datasource) + + with datasource.database.get_sqla_engine() as engine: + connection = engine.connect() + + try: + result_proxy = connection.execution_options( + stream_results=True + ).execute(text(sql_query)) + + columns = list(result_proxy.keys()) + + # Use StringIO with csv.writer for proper escaping + buffer = io.StringIO() + csv_writer = csv.writer(buffer, quoting=csv.QUOTE_MINIMAL) + + # Write CSV header + header_data, header_bytes = self._write_csv_header( + columns, csv_writer, buffer + ) + total_bytes += header_bytes + yield header_data + + # Process rows and yield chunks + row_count = 0 + for data_chunk, rows_processed, chunk_bytes in self._process_rows( + result_proxy, csv_writer, buffer + ): + total_bytes += chunk_bytes + row_count = rows_processed + yield data_chunk + + # Log completion + total_time = time.time() - start_time + total_mb = total_bytes / (1024 * 1024) + logger.info( + "Streaming CSV completed: %s rows, %.1fMB in %.2fs", + f"{row_count:,}", + total_mb, + total_time, + ) + + finally: + connection.close() + + def run(self) -> Callable[[], Generator[str, None, None]]: """ Execute the streaming CSV export. @@ -74,92 +187,11 @@ def run(self) -> Callable[[], Generator[str, None, None]]: # noqa: C901 The callable is needed to maintain Flask app context during streaming. """ - def csv_generator() -> Generator[str, None, None]: # noqa: C901 + def csv_generator() -> Generator[str, None, None]: """Generator that yields CSV data from database query.""" with self._current_app.app_context(): - start_time = time.time() - total_bytes = 0 - try: - from superset import db - from superset.connectors.sqla.models import SqlaTable - - datasource = self._query_context.datasource - - with db.session() as session: - if isinstance(datasource, SqlaTable): - datasource = session.merge(datasource) - - query_obj = self._query_context.queries[0] - sql_query = datasource.get_query_str(query_obj.to_dict()) - - with datasource.database.get_sqla_engine() as engine: - connection = engine.connect() - - try: - from sqlalchemy import text - - result_proxy = connection.execution_options( - stream_results=True - ).execute(text(sql_query)) - - columns = list(result_proxy.keys()) - - # Use StringIO with csv.writer for proper escaping - buffer = io.StringIO() - csv_writer = csv.writer( - buffer, quoting=csv.QUOTE_MINIMAL - ) - - # Write CSV header - csv_writer.writerow(columns) - header_data = buffer.getvalue() - total_bytes += len(header_data.encode("utf-8")) - yield header_data - buffer.seek(0) - buffer.truncate() - - row_count = 0 - flush_threshold = 65536 # 64KB - - while True: - rows = result_proxy.fetchmany(self._chunk_size) - if not rows: - break - - for row in rows: - csv_writer.writerow(row) - row_count += 1 - - # Check buffer size and flush if needed - current_size = buffer.tell() - if current_size >= flush_threshold: - data = buffer.getvalue() - data_bytes = len(data.encode("utf-8")) - total_bytes += data_bytes - yield data - buffer.seek(0) - buffer.truncate() - - # Flush remaining buffer - remaining_data = buffer.getvalue() - if remaining_data: - total_bytes += len(remaining_data.encode("utf-8")) - yield remaining_data - - # Log completion - total_time = time.time() - start_time - total_mb = total_bytes / (1024 * 1024) - logger.info( - "Streaming CSV completed: %s rows, %.1fMB in %.2fs", - f"{row_count:,}", - total_mb, - total_time, - ) - - finally: - connection.close() - + yield from self._execute_query_and_stream() except Exception as e: logger.error("Error in streaming CSV generator: %s", e) import traceback diff --git a/superset/commands/sql_lab/streaming_export_command.py b/superset/commands/sql_lab/streaming_export_command.py index 6742f5dccc16..f253abc5e5bc 100644 --- a/superset/commands/sql_lab/streaming_export_command.py +++ b/superset/commands/sql_lab/streaming_export_command.py @@ -22,10 +22,11 @@ import io import logging import time -from typing import Callable, Generator, TYPE_CHECKING +from typing import Any, Callable, Generator, TYPE_CHECKING from flask import current_app as app from flask_babel import gettext as __ +from sqlalchemy import text from superset import db from superset.commands.base import BaseCommand @@ -99,16 +100,8 @@ def validate(self) -> None: status=403, ) from ex - def run(self) -> Callable[[], Generator[str, None, None]]: # noqa: C901 - """ - Execute the streaming CSV export. - - Returns: - A callable that returns a generator yielding CSV data chunks as strings. - The callable is needed to maintain Flask app context during streaming. - """ - # Load all Query attributes while session is still active - # to avoid DetachedInstanceError + def _get_sql_and_limit(self) -> tuple[str, int | None]: + """Get the SQL query and limit from the query object.""" assert self._query is not None select_sql = self._query.select_sql @@ -134,95 +127,136 @@ def run(self) -> Callable[[], Generator[str, None, None]]: # noqa: C901 # remove extra row from `increased_limit` limit -= 1 - def csv_generator() -> Generator[str, None, None]: # noqa: C901 - """Generator that yields CSV data from SQL Lab query results.""" - with self._current_app.app_context(): - start_time = time.time() - total_bytes = 0 + return sql, limit + + def _write_csv_header( + self, columns: list[str], csv_writer: Any, buffer: io.StringIO + ) -> tuple[str, int]: + """Write CSV header and return header data with byte count.""" + csv_writer.writerow(columns) + header_data = buffer.getvalue() + total_bytes = len(header_data.encode("utf-8")) + buffer.seek(0) + buffer.truncate() + return header_data, total_bytes + + def _process_rows( + self, + result_proxy: Any, + csv_writer: Any, + buffer: io.StringIO, + limit: int | None, + ) -> Generator[tuple[str, int, int], None, None]: + """ + Process database rows and yield CSV data chunks. + + Yields tuples of (data_chunk, row_count, byte_count). + """ + row_count = 0 + flush_threshold = 65536 # 64KB + + while rows := result_proxy.fetchmany(self._chunk_size): + for row in rows: + # Apply limit if specified + if limit is not None and row_count >= limit: + break + + csv_writer.writerow(row) + row_count += 1 + + # Check buffer size and flush if needed + current_size = buffer.tell() + if current_size >= flush_threshold: + data = buffer.getvalue() + data_bytes = len(data.encode("utf-8")) + yield data, row_count, data_bytes + buffer.seek(0) + buffer.truncate() + + # Break outer loop if limit reached + if limit is not None and row_count >= limit: + break + + # Flush remaining buffer + if remaining_data := buffer.getvalue(): + data_bytes = len(remaining_data.encode("utf-8")) + yield remaining_data, row_count, data_bytes + + def _execute_query_and_stream( + self, sql: str, database: Any + ) -> Generator[str, None, None]: + """Execute query with streaming and yield CSV chunks.""" + _, limit = self._get_sql_and_limit() + start_time = time.time() + total_bytes = 0 + + with db.session() as session: + # Merge database to prevent DetachedInstanceError + merged_database = session.merge(database) + + # Execute query with streaming + with merged_database.get_sqla_engine() as engine: + connection = engine.connect() try: - # Create a new session to keep database object attached - with db.session() as session: - # Merge database to prevent DetachedInstanceError - merged_database = session.merge(database) - - # Execute query with streaming - with merged_database.get_sqla_engine() as engine: - connection = engine.connect() - - try: - from sqlalchemy import text - - result_proxy = connection.execution_options( - stream_results=True - ).execute(text(sql)) - - columns = list(result_proxy.keys()) - - # Use StringIO with csv.writer for proper escaping - buffer = io.StringIO() - csv_writer = csv.writer( - buffer, quoting=csv.QUOTE_MINIMAL - ) - - # Write CSV header - csv_writer.writerow(columns) - header_data = buffer.getvalue() - total_bytes += len(header_data.encode("utf-8")) - yield header_data - buffer.seek(0) - buffer.truncate() - - row_count = 0 - flush_threshold = 65536 # 64KB - - while True: - rows = result_proxy.fetchmany(self._chunk_size) - if not rows: - break - - for row in rows: - # Apply limit if specified - if limit is not None and row_count >= limit: - break - - csv_writer.writerow(row) - row_count += 1 - - # Check buffer size and flush if needed - current_size = buffer.tell() - if current_size >= flush_threshold: - data = buffer.getvalue() - data_bytes = len(data.encode("utf-8")) - total_bytes += data_bytes - yield data - buffer.seek(0) - buffer.truncate() - - # Break outer loop if limit reached - if limit is not None and row_count >= limit: - break - - # Flush remaining buffer - remaining_data = buffer.getvalue() - if remaining_data: - total_bytes += len(remaining_data.encode("utf-8")) - yield remaining_data - - # Log completion - total_time = time.time() - start_time - total_mb = total_bytes / (1024 * 1024) - logger.info( - "SQL Lab streaming CSV completed: %s rows, " - "%.1fMB in %.2fs", - f"{row_count:,}", - total_mb, - total_time, - ) - - finally: - connection.close() + result_proxy = connection.execution_options( + stream_results=True + ).execute(text(sql)) + + columns = list(result_proxy.keys()) + + # Use StringIO with csv.writer for proper escaping + buffer = io.StringIO() + csv_writer = csv.writer(buffer, quoting=csv.QUOTE_MINIMAL) + # Write CSV header + header_data, header_bytes = self._write_csv_header( + columns, csv_writer, buffer + ) + total_bytes += header_bytes + yield header_data + + # Process rows and yield chunks + row_count = 0 + for data_chunk, rows_processed, chunk_bytes in self._process_rows( + result_proxy, csv_writer, buffer, limit + ): + total_bytes += chunk_bytes + row_count = rows_processed + yield data_chunk + + # Log completion + total_time = time.time() - start_time + total_mb = total_bytes / (1024 * 1024) + logger.info( + "SQL Lab streaming CSV completed: %s rows, %.1fMB in %.2fs", + f"{row_count:,}", + total_mb, + total_time, + ) + + finally: + connection.close() + + def run(self) -> Callable[[], Generator[str, None, None]]: + """ + Execute the streaming CSV export. + + Returns: + A callable that returns a generator yielding CSV data chunks as strings. + The callable is needed to maintain Flask app context during streaming. + """ + # Load all Query attributes while session is still active + # to avoid DetachedInstanceError + assert self._query is not None + sql, _ = self._get_sql_and_limit() + database = self._query.database + + def csv_generator() -> Generator[str, None, None]: + """Generator that yields CSV data from SQL Lab query results.""" + with self._current_app.app_context(): + try: + yield from self._execute_query_and_stream(sql, database) except Exception as e: logger.error("Error in SQL Lab streaming CSV generator: %s", e) import traceback diff --git a/superset/sqllab/api.py b/superset/sqllab/api.py index e25ab4eb122e..c376cd903ee5 100644 --- a/superset/sqllab/api.py +++ b/superset/sqllab/api.py @@ -15,6 +15,7 @@ # specific language governing permissions and limitations # under the License. import logging +from datetime import datetime from typing import Any, cast, Optional from urllib import parse @@ -23,12 +24,16 @@ from flask_appbuilder.api import expose, protect, rison, safe from flask_appbuilder.models.sqla.interface import SQLAInterface from marshmallow import ValidationError +from werkzeug.utils import secure_filename from superset import is_feature_enabled from superset.commands.sql_lab.estimate import QueryEstimationCommand from superset.commands.sql_lab.execute import CommandResult, ExecuteSqlCommand from superset.commands.sql_lab.export import SqlResultExportCommand from superset.commands.sql_lab.results import SqlExecutionResultsCommand +from superset.commands.sql_lab.streaming_export_command import ( + StreamingSqlResultExportCommand, +) from superset.constants import MODEL_API_RW_METHOD_PERMISSION_MAP from superset.daos.database import DatabaseDAO from superset.daos.query import QueryDAO @@ -367,12 +372,6 @@ def _create_streaming_csv_response( expected_rows: int | None = None, ) -> Response: """Create a streaming CSV response for large SQL Lab result sets.""" - from datetime import datetime - - from superset.commands.sql_lab.streaming_export_command import ( - StreamingSqlResultExportCommand, - ) - # Execute streaming command chunk_size = 1000 command = StreamingSqlResultExportCommand(client_id, chunk_size) @@ -383,10 +382,8 @@ def _create_streaming_csv_response( query = command._query assert query is not None timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") - safe_name = "".join( - c for c in (query.name or "query") if c.isalnum() or c in ("-", "_") - ) - filename = f"sqllab_{safe_name}_{timestamp}.csv" + query_name = query.name or "query" + filename = secure_filename(f"sqllab_{query_name}_{timestamp}.csv") # Get the callable that returns the generator csv_generator_callable = command.run() @@ -402,7 +399,6 @@ def _create_streaming_csv_response( "Content-Disposition": f'attachment; filename="{filename}"', "Cache-Control": "no-cache", "X-Accel-Buffering": "no", # Disable nginx buffering - "X-Superset-Streaming": "true", # Identify streaming responses }, direct_passthrough=False, # Flask must iterate generator ) From c494c79cbf4cb95f0ca1ff5d82285883c900d409 Mon Sep 17 00:00:00 2001 From: amaannawab923 Date: Wed, 15 Oct 2025 20:44:02 +0530 Subject: [PATCH 38/57] fix: refactor and drying up streaming export command with a unified base class --- .../chart/data/streaming_export_command.py | 183 +++------------ .../sql_lab/streaming_export_command.py | 212 ++++------------- .../commands/streaming_export/__init__.py | 16 ++ superset/commands/streaming_export/base.py | 219 ++++++++++++++++++ .../chart/streaming_export_command_test.py | 73 ++---- .../sql_lab/streaming_export_command_test.py | 90 +++---- 6 files changed, 367 insertions(+), 426 deletions(-) create mode 100644 superset/commands/streaming_export/__init__.py create mode 100644 superset/commands/streaming_export/base.py diff --git a/superset/commands/chart/data/streaming_export_command.py b/superset/commands/chart/data/streaming_export_command.py index 425058de25bc..3de46add1158 100644 --- a/superset/commands/chart/data/streaming_export_command.py +++ b/superset/commands/chart/data/streaming_export_command.py @@ -14,36 +14,26 @@ # 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 large datasets.""" +"""Command for streaming CSV exports of chart data.""" from __future__ import annotations -import csv -import io -import logging -import time -from typing import Any, Callable, Generator, TYPE_CHECKING +from typing import Any, TYPE_CHECKING -from flask import current_app as app -from sqlalchemy import text - -from superset.commands.base import BaseCommand +from superset.commands.streaming_export.base import BaseStreamingCSVExportCommand if TYPE_CHECKING: from superset.common.query_context import QueryContext -logger = logging.getLogger(__name__) - -class StreamingCSVExportCommand(BaseCommand): +class StreamingCSVExportCommand(BaseStreamingCSVExportCommand): """ - Command to execute a streaming CSV export. + Command to execute a streaming CSV export for chart data. - This command handles the business logic for: - - Executing database queries with server-side cursors - - Generating CSV data in chunks - - Managing database connections - - Buffering data for efficient streaming + This command handles chart-specific logic: + - QueryContext validation + - Datasource preparation and SQL generation + - No row limit (exports all chart data) """ def __init__( @@ -52,157 +42,46 @@ def __init__( chunk_size: int = 1000, ): """ - Initialize the streaming export command. + Initialize the chart streaming export command. Args: query_context: The query context containing datasource and query details chunk_size: Number of rows to fetch per database query (default: 1000) """ + super().__init__(chunk_size) self._query_context = query_context - self._chunk_size = chunk_size - self._current_app = app._get_current_object() def validate(self) -> None: """Validate permissions and query context.""" self._query_context.raise_for_access() - def _prepare_datasource(self, session: Any) -> Any: - """Prepare and merge datasource with session.""" - from superset.connectors.sqla.models import SqlaTable - - datasource = self._query_context.datasource - if isinstance(datasource, SqlaTable): - datasource = session.merge(datasource) - return datasource - - def _get_sql_query(self, datasource: Any) -> str: - """Generate SQL query from datasource and query context.""" - query_obj = self._query_context.queries[0] - return datasource.get_query_str(query_obj.to_dict()) - - def _write_csv_header( - self, columns: list[str], csv_writer: Any, buffer: io.StringIO - ) -> tuple[str, int]: - """Write CSV header and return header data with byte count.""" - csv_writer.writerow(columns) - header_data = buffer.getvalue() - total_bytes = len(header_data.encode("utf-8")) - buffer.seek(0) - buffer.truncate() - return header_data, total_bytes - - def _process_rows( - self, - result_proxy: Any, - csv_writer: Any, - buffer: io.StringIO, - ) -> Generator[tuple[str, int, int], None, None]: + def _get_sql_and_database(self) -> tuple[str, Any]: """ - Process database rows and yield CSV data chunks. + Get the SQL query and database for chart export. - Yields tuples of (data_chunk, row_count, byte_count). + Returns: + Tuple of (sql_query, database_object) """ - row_count = 0 - flush_threshold = 65536 # 64KB - - while rows := result_proxy.fetchmany(self._chunk_size): - for row in rows: - csv_writer.writerow(row) - row_count += 1 - - # Check buffer size and flush if needed - current_size = buffer.tell() - if current_size >= flush_threshold: - data = buffer.getvalue() - data_bytes = len(data.encode("utf-8")) - yield data, row_count, data_bytes - buffer.seek(0) - buffer.truncate() - - # Flush remaining buffer - if remaining_data := buffer.getvalue(): - data_bytes = len(remaining_data.encode("utf-8")) - yield remaining_data, row_count, data_bytes - - def _execute_query_and_stream(self) -> Generator[str, None, None]: - """Execute query with streaming and yield CSV chunks.""" from superset import db - - start_time = time.time() - total_bytes = 0 + from superset.connectors.sqla.models import SqlaTable with db.session() as session: - datasource = self._prepare_datasource(session) - sql_query = self._get_sql_query(datasource) - - with datasource.database.get_sqla_engine() as engine: - connection = engine.connect() - - try: - result_proxy = connection.execution_options( - stream_results=True - ).execute(text(sql_query)) - - columns = list(result_proxy.keys()) - - # Use StringIO with csv.writer for proper escaping - buffer = io.StringIO() - csv_writer = csv.writer(buffer, quoting=csv.QUOTE_MINIMAL) - - # Write CSV header - header_data, header_bytes = self._write_csv_header( - columns, csv_writer, buffer - ) - total_bytes += header_bytes - yield header_data - - # Process rows and yield chunks - row_count = 0 - for data_chunk, rows_processed, chunk_bytes in self._process_rows( - result_proxy, csv_writer, buffer - ): - total_bytes += chunk_bytes - row_count = rows_processed - yield data_chunk - - # Log completion - total_time = time.time() - start_time - total_mb = total_bytes / (1024 * 1024) - logger.info( - "Streaming CSV completed: %s rows, %.1fMB in %.2fs", - f"{row_count:,}", - total_mb, - total_time, - ) - - finally: - connection.close() - - def run(self) -> Callable[[], Generator[str, None, None]]: + # Prepare datasource + datasource = self._query_context.datasource + if isinstance(datasource, SqlaTable): + datasource = session.merge(datasource) + + # Generate SQL query + query_obj = self._query_context.queries[0] + sql_query = datasource.get_query_str(query_obj.to_dict()) + + return sql_query, datasource.database + + def _get_row_limit(self) -> int | None: """ - Execute the streaming CSV export. + Get the row limit for chart export. Returns: - A callable that returns a generator yielding CSV data chunks as strings. - The callable is needed to maintain Flask app context during streaming. + None (no limit for chart exports) """ - - def csv_generator() -> Generator[str, None, None]: - """Generator that yields CSV data from database query.""" - with self._current_app.app_context(): - try: - yield from self._execute_query_and_stream() - except Exception as e: - logger.error("Error in streaming CSV generator: %s", e) - import traceback - - logger.error("Traceback: %s", traceback.format_exc()) - - # Send error marker for frontend to detect - error_marker = ( - "__STREAM_ERROR__:Export failed. " - "Please try again in some time.\n" - ) - yield error_marker - - return csv_generator + return None diff --git a/superset/commands/sql_lab/streaming_export_command.py b/superset/commands/sql_lab/streaming_export_command.py index f253abc5e5bc..966bd9ab4215 100644 --- a/superset/commands/sql_lab/streaming_export_command.py +++ b/superset/commands/sql_lab/streaming_export_command.py @@ -18,39 +18,27 @@ from __future__ import annotations -import csv -import io -import logging -import time -from typing import Any, Callable, Generator, TYPE_CHECKING +from typing import Any -from flask import current_app as app from flask_babel import gettext as __ -from sqlalchemy import text from superset import db -from superset.commands.base import BaseCommand +from superset.commands.streaming_export.base import BaseStreamingCSVExportCommand 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 -logger = logging.getLogger(__name__) - - -class StreamingSqlResultExportCommand(BaseCommand): +class StreamingSqlResultExportCommand(BaseStreamingCSVExportCommand): """ Command to execute a streaming CSV export of SQL Lab query results. - This command handles the business logic for: - - Fetching SQL Lab query results - - Generating CSV data in chunks - - Managing database connections - - Buffering data for efficient streaming + This command handles SQL Lab-specific logic: + - Query validation and access control + - SQL parsing and limit extraction + - LimitingFactor-based row limit adjustment """ def __init__( @@ -59,15 +47,14 @@ def __init__( chunk_size: int = 1000, ): """ - Initialize the streaming export command. + Initialize the SQL Lab streaming export command. Args: client_id: The SQL Lab query client ID chunk_size: Number of rows to fetch per database query (default: 1000) """ + super().__init__(chunk_size) self._client_id = client_id - self._chunk_size = chunk_size - self._current_app = app._get_current_object() self._query: Query | None = None def validate(self) -> None: @@ -100,18 +87,45 @@ def validate(self) -> None: status=403, ) from ex - def _get_sql_and_limit(self) -> tuple[str, int | None]: - """Get the SQL query and limit from the query object.""" + def _get_sql_and_database(self) -> tuple[str, Any]: + """ + Get the SQL query and database for SQL Lab export. + + Returns: + Tuple of (sql_query, database_object) + """ assert self._query is not None select_sql = self._query.select_sql executed_sql = self._query.executed_sql - limiting_factor = self._query.limiting_factor database = self._query.database - # Get the SQL and limit + # Get the SQL query if select_sql: sql = select_sql + else: + sql = executed_sql + + return sql, database + + def _get_row_limit(self) -> int | None: + """ + Get the row limit for SQL Lab export. + + Handles SQL Lab's complex limit logic based on limiting_factor. + + Returns: + Adjusted row limit or None for unlimited + """ + assert self._query is not None + + select_sql = self._query.select_sql + executed_sql = self._query.executed_sql + limiting_factor = self._query.limiting_factor + database = self._query.database + + # Get limit from SQL + if select_sql: limit = None else: sql = executed_sql @@ -119,6 +133,7 @@ def _get_sql_and_limit(self) -> tuple[str, int | None]: # when a query has multiple statements only the last one returns data limit = script.statements[-1].get_limit_value() + # Adjust limit based on limiting factor if limit is not None and limiting_factor in { LimitingFactor.QUERY, LimitingFactor.DROPDOWN, @@ -127,147 +142,4 @@ def _get_sql_and_limit(self) -> tuple[str, int | None]: # remove extra row from `increased_limit` limit -= 1 - return sql, limit - - def _write_csv_header( - self, columns: list[str], csv_writer: Any, buffer: io.StringIO - ) -> tuple[str, int]: - """Write CSV header and return header data with byte count.""" - csv_writer.writerow(columns) - header_data = buffer.getvalue() - total_bytes = len(header_data.encode("utf-8")) - buffer.seek(0) - buffer.truncate() - return header_data, total_bytes - - def _process_rows( - self, - result_proxy: Any, - csv_writer: Any, - buffer: io.StringIO, - limit: int | None, - ) -> Generator[tuple[str, int, int], None, None]: - """ - Process database rows and yield CSV data chunks. - - Yields tuples of (data_chunk, row_count, byte_count). - """ - row_count = 0 - flush_threshold = 65536 # 64KB - - while rows := result_proxy.fetchmany(self._chunk_size): - for row in rows: - # Apply limit if specified - if limit is not None and row_count >= limit: - break - - csv_writer.writerow(row) - row_count += 1 - - # Check buffer size and flush if needed - current_size = buffer.tell() - if current_size >= flush_threshold: - data = buffer.getvalue() - data_bytes = len(data.encode("utf-8")) - yield data, row_count, data_bytes - buffer.seek(0) - buffer.truncate() - - # Break outer loop if limit reached - if limit is not None and row_count >= limit: - break - - # Flush remaining buffer - if remaining_data := buffer.getvalue(): - data_bytes = len(remaining_data.encode("utf-8")) - yield remaining_data, row_count, data_bytes - - def _execute_query_and_stream( - self, sql: str, database: Any - ) -> Generator[str, None, None]: - """Execute query with streaming and yield CSV chunks.""" - _, limit = self._get_sql_and_limit() - start_time = time.time() - total_bytes = 0 - - with db.session() as session: - # Merge database to prevent DetachedInstanceError - merged_database = session.merge(database) - - # Execute query with streaming - with merged_database.get_sqla_engine() as engine: - connection = engine.connect() - - try: - result_proxy = connection.execution_options( - stream_results=True - ).execute(text(sql)) - - columns = list(result_proxy.keys()) - - # Use StringIO with csv.writer for proper escaping - buffer = io.StringIO() - csv_writer = csv.writer(buffer, quoting=csv.QUOTE_MINIMAL) - - # Write CSV header - header_data, header_bytes = self._write_csv_header( - columns, csv_writer, buffer - ) - total_bytes += header_bytes - yield header_data - - # Process rows and yield chunks - row_count = 0 - for data_chunk, rows_processed, chunk_bytes in self._process_rows( - result_proxy, csv_writer, buffer, limit - ): - total_bytes += chunk_bytes - row_count = rows_processed - yield data_chunk - - # Log completion - total_time = time.time() - start_time - total_mb = total_bytes / (1024 * 1024) - logger.info( - "SQL Lab streaming CSV completed: %s rows, %.1fMB in %.2fs", - f"{row_count:,}", - total_mb, - total_time, - ) - - finally: - connection.close() - - def run(self) -> Callable[[], Generator[str, None, None]]: - """ - Execute the streaming CSV export. - - Returns: - A callable that returns a generator yielding CSV data chunks as strings. - The callable is needed to maintain Flask app context during streaming. - """ - # Load all Query attributes while session is still active - # to avoid DetachedInstanceError - assert self._query is not None - sql, _ = self._get_sql_and_limit() - database = self._query.database - - def csv_generator() -> Generator[str, None, None]: - """Generator that yields CSV data from SQL Lab query results.""" - with self._current_app.app_context(): - try: - yield from self._execute_query_and_stream(sql, database) - except Exception as e: - logger.error("Error in SQL Lab streaming CSV generator: %s", e) - import traceback - - logger.error("Traceback: %s", traceback.format_exc()) - - # Send error marker for frontend to detect - error_marker = ( - "__STREAM_ERROR__:Export failed. " - "Please try again in some time.\n" - ) - yield error_marker - - return csv_generator + return limit diff --git a/superset/commands/streaming_export/__init__.py b/superset/commands/streaming_export/__init__.py new file mode 100644 index 000000000000..13a83393a912 --- /dev/null +++ b/superset/commands/streaming_export/__init__.py @@ -0,0 +1,16 @@ +# 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. diff --git a/superset/commands/streaming_export/base.py b/superset/commands/streaming_export/base.py new file mode 100644 index 000000000000..243ba8148835 --- /dev/null +++ b/superset/commands/streaming_export/base.py @@ -0,0 +1,219 @@ +# 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. +"""Base command for streaming CSV exports.""" + +from __future__ import annotations + +import csv +import io +import logging +import time +from abc import abstractmethod +from typing import Any, Callable, Generator + +from flask import current_app as app +from sqlalchemy import text + +from superset import db +from superset.commands.base import BaseCommand + +logger = logging.getLogger(__name__) + + +class BaseStreamingCSVExportCommand(BaseCommand): + """ + Base class for streaming CSV export commands. + + Provides shared functionality for: + - Generating CSV data in chunks + - Managing database connections + - Buffering data for efficient streaming + - Error handling with user-friendly messages + + Subclasses must implement: + - _get_sql_and_database(): Return SQL query string and database object + - _get_row_limit(): Return optional row limit for the export + """ + + def __init__(self, chunk_size: int = 1000): + """ + Initialize the streaming export command. + + Args: + chunk_size: Number of rows to fetch per database query (default: 1000) + """ + self._chunk_size = chunk_size + self._current_app = app._get_current_object() + + @abstractmethod + def _get_sql_and_database(self) -> tuple[str, Any]: + """ + Get the SQL query and database for execution. + + Returns: + Tuple of (sql_query, database_object) + """ + + @abstractmethod + def _get_row_limit(self) -> int | None: + """ + Get the row limit for the export. + + Returns: + Row limit or None for unlimited + """ + + def _write_csv_header( + self, columns: list[str], csv_writer: Any, buffer: io.StringIO + ) -> tuple[str, int]: + """Write CSV header and return header data with byte count.""" + csv_writer.writerow(columns) + header_data = buffer.getvalue() + total_bytes = len(header_data.encode("utf-8")) + buffer.seek(0) + buffer.truncate() + return header_data, total_bytes + + def _process_rows( + self, + result_proxy: Any, + csv_writer: Any, + buffer: io.StringIO, + limit: int | None, + ) -> Generator[tuple[str, int, int], None, None]: + """ + Process database rows and yield CSV data chunks. + + Yields tuples of (data_chunk, row_count, byte_count). + """ + row_count = 0 + flush_threshold = 65536 # 64KB + + while rows := result_proxy.fetchmany(self._chunk_size): + for row in rows: + # Apply limit if specified + if limit is not None and row_count >= limit: + break + + csv_writer.writerow(row) + row_count += 1 + + # Check buffer size and flush if needed + current_size = buffer.tell() + if current_size >= flush_threshold: + data = buffer.getvalue() + data_bytes = len(data.encode("utf-8")) + yield data, row_count, data_bytes + buffer.seek(0) + buffer.truncate() + + # Break outer loop if limit reached + if limit is not None and row_count >= limit: + break + + # Flush remaining buffer + if remaining_data := buffer.getvalue(): + data_bytes = len(remaining_data.encode("utf-8")) + yield remaining_data, row_count, data_bytes + + def _execute_query_and_stream( + self, sql: str, database: Any, limit: int | None + ) -> Generator[str, None, None]: + """Execute query with streaming and yield CSV chunks.""" + start_time = time.time() + total_bytes = 0 + + with db.session() as session: + # Merge database to prevent DetachedInstanceError + merged_database = session.merge(database) + + # Execute query with streaming + with merged_database.get_sqla_engine() as engine: + connection = engine.connect() + + try: + result_proxy = connection.execution_options( + stream_results=True + ).execute(text(sql)) + + columns = list(result_proxy.keys()) + + # Use StringIO with csv.writer for proper escaping + buffer = io.StringIO() + csv_writer = csv.writer(buffer, quoting=csv.QUOTE_MINIMAL) + + # Write CSV header + header_data, header_bytes = self._write_csv_header( + columns, csv_writer, buffer + ) + total_bytes += header_bytes + yield header_data + + # Process rows and yield chunks + row_count = 0 + for data_chunk, rows_processed, chunk_bytes in self._process_rows( + result_proxy, csv_writer, buffer, limit + ): + total_bytes += chunk_bytes + row_count = rows_processed + yield data_chunk + + # Log completion + total_time = time.time() - start_time + total_mb = total_bytes / (1024 * 1024) + logger.info( + "Streaming CSV completed: %s rows, %.1fMB in %.2fs", + f"{row_count:,}", + total_mb, + total_time, + ) + + finally: + connection.close() + + def run(self) -> Callable[[], Generator[str, None, None]]: + """ + Execute the streaming CSV export. + + Returns: + A callable that returns a generator yielding CSV data chunks as strings. + The callable is needed to maintain Flask app context during streaming. + """ + # Load all needed data while session is still active + # to avoid DetachedInstanceError + sql, database = self._get_sql_and_database() + limit = self._get_row_limit() + + def csv_generator() -> Generator[str, None, None]: + """Generator that yields CSV data chunks.""" + with self._current_app.app_context(): + try: + yield from self._execute_query_and_stream(sql, database, limit) + except Exception as e: + logger.error("Error in streaming CSV generator: %s", e) + import traceback + + logger.error("Traceback: %s", traceback.format_exc()) + + # Send error marker for frontend to detect + error_marker = ( + "__STREAM_ERROR__:Export failed. " + "Please try again in some time.\n" + ) + yield error_marker + + return csv_generator diff --git a/tests/unit_tests/commands/chart/streaming_export_command_test.py b/tests/unit_tests/commands/chart/streaming_export_command_test.py index a88f074f3d87..e3fc6939753c 100644 --- a/tests/unit_tests/commands/chart/streaming_export_command_test.py +++ b/tests/unit_tests/commands/chart/streaming_export_command_test.py @@ -24,6 +24,24 @@ ) +def _setup_chart_mocks( + mocker: MockerFixture, sql: str = "SELECT * FROM test" +) -> tuple[MockerFixture, MockerFixture, MockerFixture]: + """Set up common mocks for chart streaming export tests.""" + mock_db = mocker.patch("superset.commands.streaming_export.base.db") + mock_session = mocker.MagicMock() + mock_db.session.return_value.__enter__.return_value = mock_session + + query_context = mocker.MagicMock() + datasource = mocker.MagicMock() + datasource.get_query_str.return_value = sql + query_context.datasource = datasource + query_context.queries = [mocker.MagicMock()] + mock_session.merge.return_value = datasource + + return mock_db, query_context, datasource + + def test_streaming_csv_export_command_init(mocker: MockerFixture) -> None: """Test command initialization.""" query_context = mocker.MagicMock() @@ -66,16 +84,7 @@ def test_validate_raises_exception_on_access_denied(mocker: MockerFixture) -> No def test_csv_generation_with_small_dataset(mocker: MockerFixture) -> None: """Test CSV generation with a small dataset.""" - mock_db = mocker.patch("superset.db") - mock_session = mocker.MagicMock() - mock_db.session.return_value.__enter__.return_value = mock_session - - query_context = mocker.MagicMock() - datasource = mocker.MagicMock() - datasource.get_query_str.return_value = "SELECT * FROM test" - query_context.datasource = datasource - query_context.queries = [mocker.MagicMock()] - mock_session.merge.return_value = datasource + mock_db, query_context, datasource = _setup_chart_mocks(mocker) mock_result_proxy = mocker.MagicMock() mock_result_proxy.keys.return_value = ["col1", "col2", "col3"] @@ -116,16 +125,7 @@ def test_csv_generation_with_small_dataset(mocker: MockerFixture) -> None: def test_csv_generation_with_special_characters(mocker: MockerFixture) -> None: """Test CSV generation properly escapes special characters.""" - mock_db = mocker.patch("superset.db") - mock_session = mocker.MagicMock() - mock_db.session.return_value.__enter__.return_value = mock_session - - query_context = mocker.MagicMock() - datasource = mocker.MagicMock() - datasource.get_query_str.return_value = "SELECT * FROM test" - query_context.datasource = datasource - query_context.queries = [mocker.MagicMock()] - mock_session.merge.return_value = datasource + mock_db, query_context, datasource = _setup_chart_mocks(mocker) mock_result = mocker.MagicMock() mock_result.keys.return_value = ["name", "description"] @@ -155,16 +155,7 @@ def test_csv_generation_with_special_characters(mocker: MockerFixture) -> None: def test_streaming_with_null_values(mocker: MockerFixture) -> None: """Test CSV generation handles NULL values correctly.""" - mock_db = mocker.patch("superset.db") - mock_session = mocker.MagicMock() - mock_db.session.return_value.__enter__.return_value = mock_session - - query_context = mocker.MagicMock() - datasource = mocker.MagicMock() - datasource.get_query_str.return_value = "SELECT * FROM test" - query_context.datasource = datasource - query_context.queries = [mocker.MagicMock()] - mock_session.merge.return_value = datasource + mock_db, query_context, datasource = _setup_chart_mocks(mocker) mock_result = mocker.MagicMock() mock_result.keys.return_value = ["col1", "col2", "col3"] @@ -194,16 +185,7 @@ def test_streaming_with_null_values(mocker: MockerFixture) -> None: def test_streaming_execution_options_enabled(mocker: MockerFixture) -> None: """Test that streaming execution options are enabled.""" - mock_db = mocker.patch("superset.db") - mock_session = mocker.MagicMock() - mock_db.session.return_value.__enter__.return_value = mock_session - - query_context = mocker.MagicMock() - datasource = mocker.MagicMock() - datasource.get_query_str.return_value = "SELECT * FROM test" - query_context.datasource = datasource - query_context.queries = [mocker.MagicMock()] - mock_session.merge.return_value = datasource + mock_db, query_context, datasource = _setup_chart_mocks(mocker) mock_result_proxy = mocker.MagicMock() mock_result_proxy.keys.return_value = ["col1", "col2", "col3"] @@ -237,16 +219,7 @@ def test_streaming_execution_options_enabled(mocker: MockerFixture) -> None: def test_empty_result_set(mocker: MockerFixture) -> None: """Test CSV generation with empty result set.""" - mock_db = mocker.patch("superset.db") - mock_session = mocker.MagicMock() - mock_db.session.return_value.__enter__.return_value = mock_session - - query_context = mocker.MagicMock() - datasource = mocker.MagicMock() - datasource.get_query_str.return_value = "SELECT * FROM test" - query_context.datasource = datasource - query_context.queries = [mocker.MagicMock()] - mock_session.merge.return_value = datasource + mock_db, query_context, datasource = _setup_chart_mocks(mocker) mock_result = mocker.MagicMock() mock_result.keys.return_value = ["col1", "col2"] diff --git a/tests/unit_tests/commands/sql_lab/streaming_export_command_test.py b/tests/unit_tests/commands/sql_lab/streaming_export_command_test.py index c0fc82ed04f8..b831e2da89cf 100644 --- a/tests/unit_tests/commands/sql_lab/streaming_export_command_test.py +++ b/tests/unit_tests/commands/sql_lab/streaming_export_command_test.py @@ -19,6 +19,7 @@ from unittest.mock import MagicMock, Mock, patch import pytest +from pytest_mock import MockerFixture from superset.commands.sql_lab.streaming_export_command import ( StreamingSqlResultExportCommand, @@ -28,6 +29,17 @@ from superset.sqllab.limiting_factor import LimitingFactor +def _setup_sqllab_mocks( + mocker: MockerFixture, mock_query: MagicMock +) -> tuple[MagicMock, MagicMock]: + """Set up common mocks for SQL Lab streaming export tests.""" + mock_db = mocker.patch("superset.commands.streaming_export.base.db") + mock_session = MagicMock() + mock_db.session.return_value.__enter__.return_value = mock_session + mock_session.merge.return_value = mock_query.database + return mock_db, mock_session + + @pytest.fixture def mock_query(): """Create a mock SQL Lab query.""" @@ -121,15 +133,12 @@ def test_validate_success(mock_db, mock_query): mock_query.raise_for_access.assert_called_once() -@patch("superset.commands.sql_lab.streaming_export_command.db") -def test_csv_generation_with_select_sql(mock_db, mock_query, mock_result_proxy): +def test_csv_generation_with_select_sql(mocker, mock_query, mock_result_proxy): """Test CSV generation when query has select_sql.""" mock_query.select_sql = "SELECT * FROM test WHERE id > 0" mock_query.executed_sql = None - mock_session = MagicMock() - mock_db.session.return_value.__enter__.return_value = mock_session - mock_session.merge.return_value = mock_query.database + mock_db, mock_session = _setup_sqllab_mocks(mocker, mock_query) mock_connection = MagicMock() mock_connection.execution_options.return_value.execute.return_value = ( @@ -160,10 +169,9 @@ def test_csv_generation_with_select_sql(mock_db, mock_query, mock_result_proxy): assert "3,test3,300" in csv_data -@patch("superset.commands.sql_lab.streaming_export_command.db") @patch("superset.commands.sql_lab.streaming_export_command.SQLScript") def test_csv_generation_with_executed_sql_and_limit( - mock_sqlscript, mock_db, mock_query, mock_result_proxy + mock_sqlscript, mocker, mock_query, mock_result_proxy ): """Test CSV generation with executed_sql and applies limit.""" mock_query.select_sql = None @@ -176,9 +184,7 @@ def test_csv_generation_with_executed_sql_and_limit( mock_script_instance.statements = [mock_statement] mock_sqlscript.return_value = mock_script_instance - mock_session = MagicMock() - mock_db.session.return_value.__enter__.return_value = mock_session - mock_session.merge.return_value = mock_query.database + mock_db, mock_session = _setup_sqllab_mocks(mocker, mock_query) mock_result = MagicMock() mock_result.keys.return_value = ["id", "name"] @@ -208,8 +214,7 @@ def test_csv_generation_with_executed_sql_and_limit( assert len(lines) == 3 # header + 2 rows (limit - 1) -@patch("superset.commands.sql_lab.streaming_export_command.db") -def test_csv_generation_with_special_characters(mock_db, mock_query): +def test_csv_generation_with_special_characters(mocker, mock_query): """Test CSV generation properly escapes special characters.""" mock_query.select_sql = "SELECT * FROM test" @@ -220,9 +225,7 @@ def test_csv_generation_with_special_characters(mock_db, mock_query): [], ] - mock_session = MagicMock() - mock_db.session.return_value.__enter__.return_value = mock_session - mock_session.merge.return_value = mock_query.database + mock_db, mock_session = _setup_sqllab_mocks(mocker, mock_query) mock_connection = MagicMock() mock_connection.execution_options.return_value.execute.return_value = mock_result @@ -247,8 +250,7 @@ def test_csv_generation_with_special_characters(mock_db, mock_query): assert "Tab\tchar" in csv_data -@patch("superset.commands.sql_lab.streaming_export_command.db") -def test_limiting_factor_dropdown(mock_db, mock_query): +def test_limiting_factor_dropdown(mocker, mock_query): """Test limit adjustment for DROPDOWN limiting factor.""" mock_query.select_sql = None mock_query.executed_sql = "SELECT * FROM test LIMIT 101" @@ -267,9 +269,7 @@ def test_limiting_factor_dropdown(mock_db, mock_query): mock_result.keys.return_value = ["id"] mock_result.fetchmany.side_effect = [[(i,) for i in range(101)], []] - mock_session = MagicMock() - mock_db.session.return_value.__enter__.return_value = mock_session - mock_session.merge.return_value = mock_query.database + mock_db, mock_session = _setup_sqllab_mocks(mocker, mock_query) mock_connection = MagicMock() mock_connection.execution_options.return_value.execute.return_value = ( @@ -294,8 +294,7 @@ def test_limiting_factor_dropdown(mock_db, mock_query): assert len(lines) == 101 -@patch("superset.commands.sql_lab.streaming_export_command.db") -def test_limiting_factor_query_and_dropdown(mock_db, mock_query): +def test_limiting_factor_query_and_dropdown(mocker, mock_query): """Test limit adjustment for QUERY_AND_DROPDOWN limiting factor.""" mock_query.select_sql = None mock_query.executed_sql = "SELECT * FROM test LIMIT 51" @@ -314,9 +313,7 @@ def test_limiting_factor_query_and_dropdown(mock_db, mock_query): mock_result.keys.return_value = ["id"] mock_result.fetchmany.side_effect = [[(i,) for i in range(51)], []] - mock_session = MagicMock() - mock_db.session.return_value.__enter__.return_value = mock_session - mock_session.merge.return_value = mock_query.database + mock_db, mock_session = _setup_sqllab_mocks(mocker, mock_query) mock_connection = MagicMock() mock_connection.execution_options.return_value.execute.return_value = ( @@ -341,8 +338,7 @@ def test_limiting_factor_query_and_dropdown(mock_db, mock_query): assert len(lines) == 51 -@patch("superset.commands.sql_lab.streaming_export_command.db") -def test_empty_result_set(mock_db, mock_query): +def test_empty_result_set(mocker, mock_query): """Test CSV generation with empty result set.""" mock_query.select_sql = "SELECT * FROM empty_table" @@ -350,9 +346,7 @@ def test_empty_result_set(mock_db, mock_query): mock_result.keys.return_value = ["col1", "col2"] mock_result.fetchmany.side_effect = [[]] - mock_session = MagicMock() - mock_db.session.return_value.__enter__.return_value = mock_session - mock_session.merge.return_value = mock_query.database + mock_db, mock_session = _setup_sqllab_mocks(mocker, mock_query) mock_connection = MagicMock() mock_connection.execution_options.return_value.execute.return_value = mock_result @@ -376,11 +370,11 @@ def test_empty_result_set(mock_db, mock_query): assert lines[0] == "col1,col2" -@patch("superset.commands.sql_lab.streaming_export_command.db") -def test_error_handling_yields_error_marker(mock_db, mock_query): +def test_error_handling_yields_error_marker(mocker, mock_query): """Test that exceptions are caught and error marker is yielded.""" mock_query.select_sql = "SELECT * FROM test" + mock_db = mocker.patch("superset.commands.streaming_export.base.db") mock_session = MagicMock() mock_db.session.return_value.__enter__.return_value = mock_session mock_session.merge.side_effect = Exception("Database connection failed") @@ -399,14 +393,11 @@ def test_error_handling_yields_error_marker(mock_db, mock_query): assert "Export failed" in error_output -@patch("superset.commands.sql_lab.streaming_export_command.db") -def test_connection_is_closed_after_streaming(mock_db, mock_query, mock_result_proxy): +def test_connection_is_closed_after_streaming(mocker, mock_query, mock_result_proxy): """Test that database connection is properly closed.""" mock_query.select_sql = "SELECT * FROM test" - mock_session = MagicMock() - mock_db.session.return_value.__enter__.return_value = mock_session - mock_session.merge.return_value = mock_query.database + mock_db, mock_session = _setup_sqllab_mocks(mocker, mock_query) mock_connection = MagicMock() mock_connection.execution_options.return_value.execute.return_value = ( @@ -430,14 +421,11 @@ def test_connection_is_closed_after_streaming(mock_db, mock_query, mock_result_p mock_connection.close.assert_called_once() -@patch("superset.commands.sql_lab.streaming_export_command.db") -def test_streaming_execution_options_enabled(mock_db, mock_query, mock_result_proxy): +def test_streaming_execution_options_enabled(mocker, mock_query, mock_result_proxy): """Test that streaming execution options are enabled.""" mock_query.select_sql = "SELECT * FROM test" - mock_session = MagicMock() - mock_db.session.return_value.__enter__.return_value = mock_session - mock_session.merge.return_value = mock_query.database + mock_db, mock_session = _setup_sqllab_mocks(mocker, mock_query) mock_connection = MagicMock() mock_execution_options = Mock() @@ -462,15 +450,12 @@ def test_streaming_execution_options_enabled(mock_db, mock_query, mock_result_pr mock_connection.execution_options.assert_called_once_with(stream_results=True) -@patch("superset.commands.sql_lab.streaming_export_command.db") -@patch("superset.commands.sql_lab.streaming_export_command.logger") -def test_completion_logging(mock_logger, mock_db, mock_query, mock_result_proxy): +@patch("superset.commands.streaming_export.base.logger") +def test_completion_logging(mock_logger, mocker, mock_query, mock_result_proxy): """Test that completion is logged with metrics.""" mock_query.select_sql = "SELECT * FROM test" - mock_session = MagicMock() - mock_db.session.return_value.__enter__.return_value = mock_session - mock_session.merge.return_value = mock_query.database + mock_db, mock_session = _setup_sqllab_mocks(mocker, mock_query) mock_connection = MagicMock() mock_connection.execution_options.return_value.execute.return_value = ( @@ -493,12 +478,11 @@ def test_completion_logging(mock_logger, mock_db, mock_query, mock_result_proxy) assert mock_logger.info.called log_message = str(mock_logger.info.call_args) - assert "SQL Lab streaming CSV completed" in log_message + assert "Streaming CSV completed" in log_message assert "rows" in log_message -@patch("superset.commands.sql_lab.streaming_export_command.db") -def test_null_values_handling(mock_db, mock_query): +def test_null_values_handling(mocker, mock_query): """Test CSV generation handles NULL values correctly.""" mock_query.select_sql = "SELECT * FROM test" @@ -509,9 +493,7 @@ def test_null_values_handling(mock_db, mock_query): [], ] - mock_session = MagicMock() - mock_db.session.return_value.__enter__.return_value = mock_session - mock_session.merge.return_value = mock_query.database + mock_db, mock_session = _setup_sqllab_mocks(mocker, mock_query) mock_connection = MagicMock() mock_connection.execution_options.return_value.execute.return_value = mock_result From bbf1b5eeaa87f6ab21d08ea20585301d521425d6 Mon Sep 17 00:00:00 2001 From: amaannawab923 Date: Wed, 15 Oct 2025 21:11:30 +0530 Subject: [PATCH 39/57] fix: use csv header instead of custom header from csv streaming api --- .../StreamingExportModal.tsx | 102 +++++++++--------- superset-frontend/webpack.proxy-config.js | 7 +- 2 files changed, 56 insertions(+), 53 deletions(-) diff --git a/superset-frontend/src/components/StreamingExportModal/StreamingExportModal.tsx b/superset-frontend/src/components/StreamingExportModal/StreamingExportModal.tsx index fc7a4531d63b..66d47b6164e2 100644 --- a/superset-frontend/src/components/StreamingExportModal/StreamingExportModal.tsx +++ b/superset-frontend/src/components/StreamingExportModal/StreamingExportModal.tsx @@ -16,8 +16,7 @@ * specific language governing permissions and limitations * under the License. */ -/* eslint-disable theme-colors/no-literal-colors */ -import { styled, t } from '@superset-ui/core'; +import { styled, t, useTheme } from '@superset-ui/core'; import { Modal, Button, Typography, Progress } from 'antd'; import { Icons } from '@superset-ui/core/components/Icons'; @@ -68,7 +67,7 @@ const ProgressWrapper = styled.div` `; const SuccessIcon = styled(Icons.CheckCircleFilled)` - color: #52c41a; + color: ${({ theme }) => theme.colorSuccess}; font-size: 24px; flex-shrink: 0; `; @@ -79,13 +78,13 @@ const ErrorIconWrapper = styled.div` justify-content: center; width: 16px; height: 16px; - background-color: #ff4d4f; + background-color: ${({ theme }) => theme.colorError}; border-radius: 50%; flex-shrink: 0; `; const ErrorIconStyled = styled(Icons.CloseOutlined)` - color: white; + color: ${({ theme }) => theme.colorWhite}; font-size: 10px; `; @@ -108,45 +107,45 @@ const ErrorText = styled(Text)` `; const CancelButton = styled(Button)` - background-color: #f0fff8; - color: #1c997a; - border-color: #f0fff8; + background-color: ${({ theme }) => theme.colorSuccessBg}; + color: ${({ theme }) => theme.colorSuccess}; + border-color: ${({ theme }) => theme.colorSuccessBg}; &:hover { - background-color: #f0fff8; - color: #1c997a; - border-color: #1c997a; + background-color: ${({ theme }) => theme.colorSuccessBg}; + color: ${({ theme }) => theme.colorSuccess}; + border-color: ${({ theme }) => theme.colorSuccess}; } &:focus { - background-color: #f0fff8; - color: #1c997a; - border-color: #1c997a; + background-color: ${({ theme }) => theme.colorSuccessBg}; + color: ${({ theme }) => theme.colorSuccess}; + border-color: ${({ theme }) => theme.colorSuccess}; } `; const DownloadButton = styled(Button)` &.ant-btn-primary { - background-color: #2ec196; - border-color: #2ec196; - color: #ffffff; + background-color: ${({ theme }) => theme.colorSuccess}; + border-color: ${({ theme }) => theme.colorSuccess}; + color: ${({ theme }) => theme.colorWhite}; &:hover:not(:disabled) { - background-color: #26a880; - border-color: #26a880; - color: #ffffff; + background-color: ${({ theme }) => theme.colorSuccessActive}; + border-color: ${({ theme }) => theme.colorSuccessActive}; + color: ${({ theme }) => theme.colorWhite}; } &:focus:not(:disabled) { - background-color: #2ec196; - border-color: #2ec196; - color: #ffffff; + background-color: ${({ theme }) => theme.colorSuccess}; + border-color: ${({ theme }) => theme.colorSuccess}; + color: ${({ theme }) => theme.colorWhite}; } &:disabled { - background-color: #f2f2f2; - border-color: #f2f2f2; - color: #b5b5b5; + background-color: ${({ theme }) => theme.colorBgContainerDisabled}; + border-color: ${({ theme }) => theme.colorBgContainerDisabled}; + color: ${({ theme }) => theme.colorTextDisabled}; } } `; @@ -273,30 +272,33 @@ const StreamingContent = ({ filename?: string; getProgressPercentage: () => number; onCancel: () => void; -}) => ( - - - `${Math.round(percent || 0)}%`} - /> - - {filename - ? t('Processing export for %s', filename) - : t('Processing export...')} - - - - {t('Cancel')} - - {t('Download')} - - - -); +}) => { + const theme = useTheme(); + return ( + + + `${Math.round(percent || 0)}%`} + /> + + {filename + ? t('Processing export for %s', filename) + : t('Processing export...')} + + + + {t('Cancel')} + + {t('Download')} + + + + ); +}; const ModalStateContent = ({ status, diff --git a/superset-frontend/webpack.proxy-config.js b/superset-frontend/webpack.proxy-config.js index 5011814be569..f15b79326960 100644 --- a/superset-frontend/webpack.proxy-config.js +++ b/superset-frontend/webpack.proxy-config.js @@ -165,10 +165,11 @@ module.exports = newManifest => { if (isHTML(response)) { processHTML(proxyResponse, response); } else { - const isStreaming = - proxyResponse.headers['transfer-encoding'] === 'chunked'; + const isCSV = (proxyResponse.headers['content-type'] || '').includes( + 'text/csv', + ); - if (isStreaming) { + if (isCSV) { proxyResponse.on('data', chunk => { response.write(chunk); if (response.flush) { From 62177919f6b56f86245c96fdc19b4eae23bfe3c3 Mon Sep 17 00:00:00 2001 From: amaannawab923 Date: Wed, 15 Oct 2025 21:48:32 +0530 Subject: [PATCH 40/57] fix: tests --- .../chart/data/streaming_export_command.py | 19 +++----- .../chart/streaming_export_command_test.py | 3 +- .../sql_lab/streaming_export_command_test.py | 43 ++++++++----------- 3 files changed, 25 insertions(+), 40 deletions(-) diff --git a/superset/commands/chart/data/streaming_export_command.py b/superset/commands/chart/data/streaming_export_command.py index 3de46add1158..b6ec3a36698a 100644 --- a/superset/commands/chart/data/streaming_export_command.py +++ b/superset/commands/chart/data/streaming_export_command.py @@ -62,20 +62,13 @@ def _get_sql_and_database(self) -> tuple[str, Any]: Returns: Tuple of (sql_query, database_object) """ - from superset import db - from superset.connectors.sqla.models import SqlaTable + # Get datasource and generate SQL query + # Note: datasource should already be attached to a session from query_context + datasource = self._query_context.datasource + query_obj = self._query_context.queries[0] + sql_query = datasource.get_query_str(query_obj.to_dict()) - with db.session() as session: - # Prepare datasource - datasource = self._query_context.datasource - if isinstance(datasource, SqlaTable): - datasource = session.merge(datasource) - - # Generate SQL query - query_obj = self._query_context.queries[0] - sql_query = datasource.get_query_str(query_obj.to_dict()) - - return sql_query, datasource.database + return sql_query, datasource.database def _get_row_limit(self) -> int | None: """ diff --git a/tests/unit_tests/commands/chart/streaming_export_command_test.py b/tests/unit_tests/commands/chart/streaming_export_command_test.py index e3fc6939753c..2ac26dde82a3 100644 --- a/tests/unit_tests/commands/chart/streaming_export_command_test.py +++ b/tests/unit_tests/commands/chart/streaming_export_command_test.py @@ -35,9 +35,10 @@ def _setup_chart_mocks( query_context = mocker.MagicMock() datasource = mocker.MagicMock() datasource.get_query_str.return_value = sql + datasource.database = mocker.MagicMock() query_context.datasource = datasource query_context.queries = [mocker.MagicMock()] - mock_session.merge.return_value = datasource + mock_session.merge.return_value = datasource.database return mock_db, query_context, datasource diff --git a/tests/unit_tests/commands/sql_lab/streaming_export_command_test.py b/tests/unit_tests/commands/sql_lab/streaming_export_command_test.py index b831e2da89cf..e4e2f9d8dd6f 100644 --- a/tests/unit_tests/commands/sql_lab/streaming_export_command_test.py +++ b/tests/unit_tests/commands/sql_lab/streaming_export_command_test.py @@ -33,11 +33,18 @@ def _setup_sqllab_mocks( mocker: MockerFixture, mock_query: MagicMock ) -> tuple[MagicMock, MagicMock]: """Set up common mocks for SQL Lab streaming export tests.""" - mock_db = mocker.patch("superset.commands.streaming_export.base.db") + mock_db_base = mocker.patch("superset.commands.streaming_export.base.db") mock_session = MagicMock() - mock_db.session.return_value.__enter__.return_value = mock_session + mock_db_base.session.return_value.__enter__.return_value = mock_session mock_session.merge.return_value = mock_query.database - return mock_db, mock_session + + mock_db_sqllab = mocker.patch( + "superset.commands.sql_lab.streaming_export_command.db" + ) + mock_query_result = mock_db_sqllab.session.query.return_value.filter_by.return_value + mock_query_result.one_or_none.return_value = mock_query + + return mock_db_base, mock_session @pytest.fixture @@ -150,8 +157,6 @@ def test_csv_generation_with_select_sql(mocker, mock_query, mock_result_proxy): mock_engine ) - mock_query_result = mock_db.session.query.return_value.filter_by.return_value - mock_query_result.one_or_none.return_value = mock_query command = StreamingSqlResultExportCommand("test_client_123", chunk_size=2) command.validate() @@ -201,8 +206,6 @@ def test_csv_generation_with_executed_sql_and_limit( mock_engine ) - mock_query_result = mock_db.session.query.return_value.filter_by.return_value - mock_query_result.one_or_none.return_value = mock_query command = StreamingSqlResultExportCommand("test_client_123", chunk_size=10) command.validate() @@ -235,8 +238,6 @@ def test_csv_generation_with_special_characters(mocker, mock_query): mock_engine ) - mock_query_result = mock_db.session.query.return_value.filter_by.return_value - mock_query_result.one_or_none.return_value = mock_query command = StreamingSqlResultExportCommand("test_client_123") command.validate() @@ -281,8 +282,6 @@ def test_limiting_factor_dropdown(mocker, mock_query): mock_engine ) - mock_query_result = mock_db.session.query.return_value.filter_by.return_value - mock_query_result.one_or_none.return_value = mock_query command = StreamingSqlResultExportCommand("test_client_123", chunk_size=200) command.validate() @@ -325,8 +324,6 @@ def test_limiting_factor_query_and_dropdown(mocker, mock_query): mock_engine ) - mock_query_result = mock_db.session.query.return_value.filter_by.return_value - mock_query_result.one_or_none.return_value = mock_query command = StreamingSqlResultExportCommand("test_client_123", chunk_size=100) command.validate() @@ -356,8 +353,6 @@ def test_empty_result_set(mocker, mock_query): mock_engine ) - mock_query_result = mock_db.session.query.return_value.filter_by.return_value - mock_query_result.one_or_none.return_value = mock_query command = StreamingSqlResultExportCommand("test_client_123") command.validate() @@ -374,13 +369,17 @@ def test_error_handling_yields_error_marker(mocker, mock_query): """Test that exceptions are caught and error marker is yielded.""" mock_query.select_sql = "SELECT * FROM test" - mock_db = mocker.patch("superset.commands.streaming_export.base.db") + mock_db_base = mocker.patch("superset.commands.streaming_export.base.db") mock_session = MagicMock() - mock_db.session.return_value.__enter__.return_value = mock_session + mock_db_base.session.return_value.__enter__.return_value = mock_session mock_session.merge.side_effect = Exception("Database connection failed") - mock_query_result = mock_db.session.query.return_value.filter_by.return_value + mock_db_sqllab = mocker.patch( + "superset.commands.sql_lab.streaming_export_command.db" + ) + mock_query_result = mock_db_sqllab.session.query.return_value.filter_by.return_value mock_query_result.one_or_none.return_value = mock_query + command = StreamingSqlResultExportCommand("test_client_123") command.validate() @@ -409,8 +408,6 @@ def test_connection_is_closed_after_streaming(mocker, mock_query, mock_result_pr mock_engine ) - mock_query_result = mock_db.session.query.return_value.filter_by.return_value - mock_query_result.one_or_none.return_value = mock_query command = StreamingSqlResultExportCommand("test_client_123") command.validate() @@ -438,8 +435,6 @@ def test_streaming_execution_options_enabled(mocker, mock_query, mock_result_pro mock_engine ) - mock_query_result = mock_db.session.query.return_value.filter_by.return_value - mock_query_result.one_or_none.return_value = mock_query command = StreamingSqlResultExportCommand("test_client_123") command.validate() @@ -467,8 +462,6 @@ def test_completion_logging(mock_logger, mocker, mock_query, mock_result_proxy): mock_engine ) - mock_query_result = mock_db.session.query.return_value.filter_by.return_value - mock_query_result.one_or_none.return_value = mock_query command = StreamingSqlResultExportCommand("test_client_123") command.validate() @@ -503,8 +496,6 @@ def test_null_values_handling(mocker, mock_query): mock_engine ) - mock_query_result = mock_db.session.query.return_value.filter_by.return_value - mock_query_result.one_or_none.return_value = mock_query command = StreamingSqlResultExportCommand("test_client_123") command.validate() From 7f55bcfd8e0b582d55a19a7af85487e7a3ad3cb9 Mon Sep 17 00:00:00 2001 From: amaannawab923 Date: Mon, 28 Oct 2024 23:30:00 +0530 Subject: [PATCH 41/57] Refactor Comments Resolution --- .../StreamingExportModal.tsx | 480 +++++++++--------- 1 file changed, 226 insertions(+), 254 deletions(-) diff --git a/superset-frontend/src/components/StreamingExportModal/StreamingExportModal.tsx b/superset-frontend/src/components/StreamingExportModal/StreamingExportModal.tsx index 66d47b6164e2..c945025897af 100644 --- a/superset-frontend/src/components/StreamingExportModal/StreamingExportModal.tsx +++ b/superset-frontend/src/components/StreamingExportModal/StreamingExportModal.tsx @@ -17,7 +17,12 @@ * under the License. */ import { styled, t, useTheme } from '@superset-ui/core'; -import { Modal, Button, Typography, Progress } from 'antd'; +import { + Modal, + Button, + Typography, + Progress, +} from '@superset-ui/core/components'; import { Icons } from '@superset-ui/core/components/Icons'; const { Text } = Typography; @@ -29,6 +34,9 @@ export enum ExportStatus { CANCELLED = 'cancelled', } +const MAX_PROGRESS_PERCENT = 99; +const COMPLETED_PERCENT = 100; + export interface StreamingProgress { totalRows?: number; rowsProcessed: number; @@ -51,105 +59,195 @@ interface StreamingExportModalProps { } const ModalContent = styled.div` - padding: ${({ theme }) => theme.sizeUnit * 4}px 0 - ${({ theme }) => theme.sizeUnit * 2}px; + ${({ theme }) => ` + padding: ${theme.sizeUnit * 4}px 0 ${theme.sizeUnit * 2}px; + `} `; const ProgressSection = styled.div` - margin: ${({ theme }) => theme.sizeUnit * 6}px 0; - position: relative; + ${({ theme }) => ` + margin: ${theme.sizeUnit * 6}px 0; + position: relative; + `} `; const ProgressWrapper = styled.div` - display: flex; - align-items: center; - gap: ${({ theme }) => theme.sizeUnit * 3}px; + ${({ theme }) => ` + display: flex; + align-items: center; + gap: ${theme.sizeUnit * 3}px; + `} +`; + +const StyledProgress = styled(Progress)` + flex: 1; `; const SuccessIcon = styled(Icons.CheckCircleFilled)` - color: ${({ theme }) => theme.colorSuccess}; - font-size: 24px; - flex-shrink: 0; + ${({ theme }) => ` + color: ${theme.colorSuccess}; + font-size: ${theme.sizeUnit * 6}px; + flex-shrink: 0; + `} `; const ErrorIconWrapper = styled.div` - display: flex; - align-items: center; - justify-content: center; - width: 16px; - height: 16px; - background-color: ${({ theme }) => theme.colorError}; - border-radius: 50%; - flex-shrink: 0; + ${({ theme }) => ` + display: flex; + align-items: center; + justify-content: center; + width: ${theme.sizeUnit * 4}px; + height: ${theme.sizeUnit * 4}px; + background-color: ${theme.colorError}; + border-radius: 50%; + flex-shrink: 0; + `} `; const ErrorIconStyled = styled(Icons.CloseOutlined)` - color: ${({ theme }) => theme.colorWhite}; - font-size: 10px; + ${({ theme }) => ` + color: ${theme.colorWhite}; + font-size: ${theme.sizeUnit * 2.5}px; + `} `; const ActionButtons = styled.div` - display: flex; - gap: ${({ theme }) => theme.sizeUnit * 2}px; - justify-content: flex-end; + ${({ theme }) => ` + display: flex; + gap: ${theme.sizeUnit * 2}px; + justify-content: flex-end; + `} `; -const ProgressText = styled(Text)` - display: block; - text-align: center; - margin-top: ${({ theme }) => theme.sizeUnit * 4}px; +const CenteredText = styled(Text)` + ${({ theme }) => ` + display: block; + text-align: center; + margin-top: ${theme.sizeUnit * 4}px; + `} `; -const ErrorText = styled(Text)` - display: block; - text-align: center; - margin-top: ${({ theme }) => theme.sizeUnit * 4}px; +const ErrorText = styled(CenteredText)` + ${({ theme }) => ` + color: ${theme.colorError}; + `} `; const CancelButton = styled(Button)` - background-color: ${({ theme }) => theme.colorSuccessBg}; - color: ${({ theme }) => theme.colorSuccess}; - border-color: ${({ theme }) => theme.colorSuccessBg}; - - &:hover { - background-color: ${({ theme }) => theme.colorSuccessBg}; - color: ${({ theme }) => theme.colorSuccess}; - border-color: ${({ theme }) => theme.colorSuccess}; - } + ${({ theme }) => ` + background-color: ${theme.colorSuccessBg}; + color: ${theme.colorSuccess}; + border-color: ${theme.colorSuccessBg}; + + &:hover { + background-color: ${theme.colorSuccessBg}; + color: ${theme.colorSuccess}; + border-color: ${theme.colorSuccess}; + } - &:focus { - background-color: ${({ theme }) => theme.colorSuccessBg}; - color: ${({ theme }) => theme.colorSuccess}; - border-color: ${({ theme }) => theme.colorSuccess}; - } + &:focus { + background-color: ${theme.colorSuccessBg}; + color: ${theme.colorSuccess}; + border-color: ${theme.colorSuccess}; + } + `} `; const DownloadButton = styled(Button)` - &.ant-btn-primary { - background-color: ${({ theme }) => theme.colorSuccess}; - border-color: ${({ theme }) => theme.colorSuccess}; - color: ${({ theme }) => theme.colorWhite}; + ${({ theme }) => ` + background-color: ${theme.colorSuccess}; + border-color: ${theme.colorSuccess}; + color: ${theme.colorWhite}; &:hover:not(:disabled) { - background-color: ${({ theme }) => theme.colorSuccessActive}; - border-color: ${({ theme }) => theme.colorSuccessActive}; - color: ${({ theme }) => theme.colorWhite}; + background-color: ${theme.colorSuccessActive}; + border-color: ${theme.colorSuccessActive}; + color: ${theme.colorWhite}; } &:focus:not(:disabled) { - background-color: ${({ theme }) => theme.colorSuccess}; - border-color: ${({ theme }) => theme.colorSuccess}; - color: ${({ theme }) => theme.colorWhite}; + background-color: ${theme.colorSuccess}; + border-color: ${theme.colorSuccess}; + color: ${theme.colorWhite}; } &:disabled { - background-color: ${({ theme }) => theme.colorBgContainerDisabled}; - border-color: ${({ theme }) => theme.colorBgContainerDisabled}; - color: ${({ theme }) => theme.colorTextDisabled}; + background-color: ${theme.colorBgContainerDisabled}; + border-color: ${theme.colorBgContainerDisabled}; + color: ${theme.colorTextDisabled}; } - } + `} `; +const triggerFileDownload = (url: string, filename: string) => { + const link = document.createElement('a'); + link.href = url; + link.download = filename; + document.body.appendChild(link); + link.click(); + document.body.removeChild(link); +}; + +const calculateProgressPercentage = ( + status: ExportStatus, + totalRows?: number, + rowsProcessed?: number, +): number => { + if (status === ExportStatus.COMPLETED) return COMPLETED_PERCENT; + + if (!totalRows || totalRows <= 0 || !rowsProcessed) return 0; + + const percentage = (rowsProcessed / totalRows) * 100; + return Math.round(Math.min(MAX_PROGRESS_PERCENT, percentage)); +}; + +const getProgressStatus = ( + status: ExportStatus, +): 'success' | 'exception' | 'normal' => { + switch (status) { + case ExportStatus.COMPLETED: + return 'success'; + case ExportStatus.ERROR: + case ExportStatus.CANCELLED: + return 'exception'; + case ExportStatus.STREAMING: + default: + return 'normal'; + } +}; + +const getMessageText = ( + status: ExportStatus, + filename?: string, + error?: string, +): string => { + switch (status) { + case ExportStatus.ERROR: + return error || t('Export failed'); + case ExportStatus.CANCELLED: + return t('Export cancelled'); + case ExportStatus.COMPLETED: + return t('Export successful: %s', filename || 'export'); + case ExportStatus.STREAMING: + default: + return filename + ? t('Processing export for %s', filename) + : t('Processing export...'); + } +}; + +const getButtonText = (status: ExportStatus): string => { + switch (status) { + case ExportStatus.ERROR: + case ExportStatus.CANCELLED: + case ExportStatus.COMPLETED: + return t('Close'); + case ExportStatus.STREAMING: + default: + return t('Cancel'); + } +}; + interface ModalStateContentProps { status: ExportStatus; progress: StreamingProgress; @@ -159,195 +257,79 @@ interface ModalStateContentProps { getProgressPercentage: () => number; } -const ErrorContent = ({ - error, - onCancel, - onRetry, - getProgressPercentage, -}: { - error?: string; - onCancel: () => void; - onRetry?: () => void; - getProgressPercentage: () => number; -}) => ( - - - - - - - - - {error || t('Export failed')} - - - {t('Close')} - {onRetry && ( - - {t('Retry')} - - )} - - -); - -const CancelledContent = ({ - getProgressPercentage, +const ModalStateContent = ({ + status, + progress, onCancel, onRetry, -}: { - getProgressPercentage: () => number; - onCancel: () => void; - onRetry?: () => void; -}) => ( - - - - {t('Export cancelled')} - - - {t('Close')} - {onRetry && ( - - {t('Retry')} - - )} - - -); - -const CompletedContent = ({ - filename, - downloadUrl, - onCancel, onDownload, -}: { - filename?: string; - downloadUrl?: string; - onCancel: () => void; - onDownload: () => void; -}) => ( - - - - - - - - {t('Export successful: %s', filename || 'export')} - - - - {t('Close')} - - {t('Download')} - - - -); - -const StreamingContent = ({ - filename, getProgressPercentage, - onCancel, -}: { - filename?: string; - getProgressPercentage: () => number; - onCancel: () => void; -}) => { +}: ModalStateContentProps) => { const theme = useTheme(); + const { downloadUrl, filename, error } = progress; + + const isError = status === ExportStatus.ERROR; + const isCancelled = status === ExportStatus.CANCELLED; + const isCompleted = status === ExportStatus.COMPLETED; + const isStreaming = status === ExportStatus.STREAMING; + + const hasIcon = isError || isCompleted; + const shouldShowRetry = (isError || isCancelled) && onRetry; + + const progressStatus = getProgressStatus(status); + const progressPercent = isCompleted ? 100 : getProgressPercentage(); + const messageText = getMessageText(status, filename, error); + const buttonText = getButtonText(status); + + const progressProps = { + percent: progressPercent, + status: progressStatus, + showInfo: isStreaming, + ...(isStreaming && { + strokeColor: theme.colorSuccess, + format: (percent?: number) => `${Math.round(percent || 0)}%`, + }), + }; + return ( - `${Math.round(percent || 0)}%`} - /> - - {filename - ? t('Processing export for %s', filename) - : t('Processing export...')} - + {hasIcon ? ( + + + {isError && ( + + + + )} + {isCompleted && } + + ) : ( + + )} + {isError ? ( + {messageText} + ) : ( + {messageText} + )} - {t('Cancel')} - - {t('Download')} - + {buttonText} + {shouldShowRetry ? ( + {t('Retry')} + ) : ( + + {t('Download')} + + )} ); }; -const ModalStateContent = ({ - status, - progress, - onCancel, - onRetry, - onDownload, - getProgressPercentage, -}: ModalStateContentProps) => { - const { downloadUrl, filename, error } = progress; - - switch (status) { - case ExportStatus.ERROR: - return ( - - ); - case ExportStatus.CANCELLED: - return ( - - ); - case ExportStatus.COMPLETED: - return ( - - ); - default: - return ( - - ); - } -}; - const StreamingExportModal = ({ visible, onCancel, @@ -356,26 +338,16 @@ const StreamingExportModal = ({ }: StreamingExportModalProps) => { const { status, downloadUrl, filename } = progress; - const getProgressPercentage = (): number => { - if (status === ExportStatus.COMPLETED) return 100; - if (progress.totalRows && progress.totalRows > 0) { - const percentage = Math.min( - 99, - (progress.rowsProcessed / progress.totalRows) * 100, - ); - return Math.round(percentage); - } - return 0; - }; + const getProgressPercentage = (): number => + calculateProgressPercentage( + status, + progress.totalRows, + progress.rowsProcessed, + ); const handleDownload = () => { if (downloadUrl && filename) { - const link = document.createElement('a'); - link.href = downloadUrl; - link.download = filename; - document.body.appendChild(link); - link.click(); - document.body.removeChild(link); + triggerFileDownload(downloadUrl, filename); onCancel(); } }; @@ -383,9 +355,9 @@ const StreamingExportModal = ({ return ( Date: Wed, 29 Oct 2025 12:57:15 +0530 Subject: [PATCH 42/57] fix: import for progress --- .../src/components/Progress/index.tsx | 24 +++++++++++++++++++ .../superset-ui-core/src/components/index.ts | 2 ++ 2 files changed, 26 insertions(+) create mode 100644 superset-frontend/packages/superset-ui-core/src/components/Progress/index.tsx diff --git a/superset-frontend/packages/superset-ui-core/src/components/Progress/index.tsx b/superset-frontend/packages/superset-ui-core/src/components/Progress/index.tsx new file mode 100644 index 000000000000..a2a41cb08474 --- /dev/null +++ b/superset-frontend/packages/superset-ui-core/src/components/Progress/index.tsx @@ -0,0 +1,24 @@ +/** + * 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 { Progress as AntdProgress } from 'antd'; +import type { ProgressProps as AntdProgressProps } from 'antd'; + +export type ProgressProps = AntdProgressProps; + +export const Progress = AntdProgress; diff --git a/superset-frontend/packages/superset-ui-core/src/components/index.ts b/superset-frontend/packages/superset-ui-core/src/components/index.ts index 55829c0b1c6e..47cb6c3a61f1 100644 --- a/superset-frontend/packages/superset-ui-core/src/components/index.ts +++ b/superset-frontend/packages/superset-ui-core/src/components/index.ts @@ -146,6 +146,8 @@ export { } from './ListViewCard'; export { Loading, type LoadingProps } from './Loading'; +export { Progress, type ProgressProps } from './Progress'; + export { Skeleton, type SkeletonProps } from './Skeleton'; export { Switch, type SwitchProps } from './Switch'; From 0bf8f390d905c3c0835e0e530d2eb09df95b98b7 Mon Sep 17 00:00:00 2001 From: amaannawab923 Date: Wed, 29 Oct 2025 13:28:53 +0530 Subject: [PATCH 43/57] move filename logic for sqllab to backend --- .../src/SqlLab/components/ResultSet/index.tsx | 8 ---- .../useStreamingExport.ts | 39 ++++++++++++------- superset/sqllab/api.py | 4 +- 3 files changed, 27 insertions(+), 24 deletions(-) diff --git a/superset-frontend/src/SqlLab/components/ResultSet/index.tsx b/superset-frontend/src/SqlLab/components/ResultSet/index.tsx index c3f91ba857d7..de207c33a4c1 100644 --- a/superset-frontend/src/SqlLab/components/ResultSet/index.tsx +++ b/superset-frontend/src/SqlLab/components/ResultSet/index.tsx @@ -423,17 +423,9 @@ const ResultSet = ({ e.preventDefault(); setShowStreamingModal(true); - const timestamp = new Date() - .toISOString() - .slice(0, 19) - .replace(/[-:]/g, '') - .replace('T', '_'); - const filename = `sqllab_${query.id}_${timestamp}.csv`; - startExport({ url: '/api/v1/sqllab/export_streaming/', payload: { client_id: query.id }, - filename, exportType: 'csv', expectedRows: rows, }); diff --git a/superset-frontend/src/components/StreamingExportModal/useStreamingExport.ts b/superset-frontend/src/components/StreamingExportModal/useStreamingExport.ts index da31a7dbba9f..fd0bf34cb2e3 100644 --- a/superset-frontend/src/components/StreamingExportModal/useStreamingExport.ts +++ b/superset-frontend/src/components/StreamingExportModal/useStreamingExport.ts @@ -42,7 +42,7 @@ const NEWLINE_BYTE = 10; // '\n' character code const createFetchRequest = async ( _url: string, payload: StreamingExportPayload, - filename: string, + filename: string | undefined, _exportType: string, expectedRows: number | undefined, signal: AbortSignal, @@ -57,12 +57,15 @@ const createFetchRequest = async ( headers['X-CSRFToken'] = csrfToken; } - // Build form data - if payload has client_id, it's SQL Lab export - // Otherwise it's a chart export with form_data - const formParams: Record = { - filename, - expected_rows: expectedRows?.toString() || '', - }; + const formParams: Record = {}; + + if (filename) { + formParams.filename = filename; + } + + if (expectedRows) { + formParams.expected_rows = expectedRows.toString(); + } if ('client_id' in payload) { // SQL Lab export - pass client_id directly @@ -146,13 +149,10 @@ export const useStreamingExport = (options: UseStreamingExportOptions = {}) => { }); try { - const defaultFilename = `export.${exportType}`; - const finalFilename = filename || defaultFilename; - const fetchOptions = await createFetchRequest( url, payload, - finalFilename, + filename, exportType, expectedRows, abortControllerRef.current.signal, @@ -169,6 +169,18 @@ export const useStreamingExport = (options: UseStreamingExportOptions = {}) => { throw new Error('Response body is not available for streaming'); } + const contentDisposition = response.headers.get('Content-Disposition'); + const defaultFilename = `export.${exportType}`; + let serverFilename = defaultFilename; + + if (contentDisposition) { + const filenameMatch = + contentDisposition.match(/filename="?([^"]+)"?/); + if (filenameMatch && filenameMatch[1]) { + serverFilename = filenameMatch[1]; + } + } + const reader = response.body.getReader(); const chunks: Uint8Array[] = []; let receivedLength = 0; @@ -227,6 +239,7 @@ export const useStreamingExport = (options: UseStreamingExportOptions = {}) => { rowsProcessed, totalRows: expectedRows, totalSize: receivedLength, + filename: serverFilename, }); } @@ -247,11 +260,11 @@ export const useStreamingExport = (options: UseStreamingExportOptions = {}) => { updateProgress({ status: ExportStatus.COMPLETED, downloadUrl, - filename: finalFilename, + filename: serverFilename, }); isExportingRef.current = false; - options.onComplete?.(downloadUrl, finalFilename); + options.onComplete?.(downloadUrl, serverFilename); } catch (error) { const errorMessage = error instanceof Error ? error.message : 'Unknown error occurred'; diff --git a/superset/sqllab/api.py b/superset/sqllab/api.py index c376cd903ee5..2b51e6a92c48 100644 --- a/superset/sqllab/api.py +++ b/superset/sqllab/api.py @@ -377,13 +377,11 @@ def _create_streaming_csv_response( command = StreamingSqlResultExportCommand(client_id, chunk_size) command.validate() - # Generate filename if not provided if not filename: query = command._query assert query is not None timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") - query_name = query.name or "query" - filename = secure_filename(f"sqllab_{query_name}_{timestamp}.csv") + filename = secure_filename(f"sqllab_{client_id}_{timestamp}.csv") # Get the callable that returns the generator csv_generator_callable = command.run() From 316421c100a8cd61c886f3c9fb76590fc5b5d9a0 Mon Sep 17 00:00:00 2001 From: amaannawab923 Date: Wed, 29 Oct 2025 15:48:44 +0530 Subject: [PATCH 44/57] fix: use constant default --- .../src/dashboard/components/gridComponents/Chart/Chart.jsx | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/superset-frontend/src/dashboard/components/gridComponents/Chart/Chart.jsx b/superset-frontend/src/dashboard/components/gridComponents/Chart/Chart.jsx index dbd64a5f3a70..8c5002cb78f6 100644 --- a/superset-frontend/src/dashboard/components/gridComponents/Chart/Chart.jsx +++ b/superset-frontend/src/dashboard/components/gridComponents/Chart/Chart.jsx @@ -82,6 +82,7 @@ const propTypes = { const RESIZE_TIMEOUT = 500; const DEFAULT_HEADER_HEIGHT = 22; +const DEFAULT_CSV_STREAMING_ROW_THRESHOLD = 100000; const ChartWrapper = styled.div` overflow: hidden; @@ -161,7 +162,8 @@ const Chart = props => { ); const streamingThreshold = useSelector( state => - state.dashboardInfo.common.conf.CSV_STREAMING_ROW_THRESHOLD || 100000, + state.dashboardInfo.common.conf.CSV_STREAMING_ROW_THRESHOLD || + DEFAULT_CSV_STREAMING_ROW_THRESHOLD, ); const datasource = useSelector( state => From 3fba39415219e27a1f7d603db481a8881fc88bf0 Mon Sep 17 00:00:00 2001 From: amaannawab923 Date: Fri, 7 Nov 2025 21:05:15 +0530 Subject: [PATCH 45/57] ci fix --- superset/charts/data/api.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/superset/charts/data/api.py b/superset/charts/data/api.py index 2124c1659c50..b4f6bfeb072a 100644 --- a/superset/charts/data/api.py +++ b/superset/charts/data/api.py @@ -19,7 +19,7 @@ import contextlib import logging from datetime import datetime -from typing import Any, TYPE_CHECKING +from typing import Any, Callable, TYPE_CHECKING from flask import current_app as app, g, make_response, request, Response from flask_appbuilder.api import expose, protect @@ -460,7 +460,7 @@ def _get_data_response( except ChartDataQueryFailedError as exc: return self.response_400(message=exc.message) - # Log is_cached if extra payload callback is provided + # Log is_cached if extra payload callback is provided if add_extra_log_payload and result and "queries" in result: is_cached_values = [query.get("is_cached") for query in result["queries"]] if len(is_cached_values) == 1: @@ -468,7 +468,6 @@ def _get_data_response( elif is_cached_values: add_extra_log_payload(is_cached=is_cached_values) - return self._send_chart_response( result, form_data, datasource, filename, expected_rows ) From dff377b0618ddd69f912ac7ed0ff0278581a7ac1 Mon Sep 17 00:00:00 2001 From: amaannawab923 Date: Mon, 10 Nov 2025 20:31:08 +0530 Subject: [PATCH 46/57] using sql row count --- .../src/dashboard/components/gridComponents/Chart/Chart.jsx | 2 ++ 1 file changed, 2 insertions(+) diff --git a/superset-frontend/src/dashboard/components/gridComponents/Chart/Chart.jsx b/superset-frontend/src/dashboard/components/gridComponents/Chart/Chart.jsx index 6da054266180..66e829e6a87a 100644 --- a/superset-frontend/src/dashboard/components/gridComponents/Chart/Chart.jsx +++ b/superset-frontend/src/dashboard/components/gridComponents/Chart/Chart.jsx @@ -443,6 +443,8 @@ const Chart = props => { queriesResponse[1]?.data?.[0]?.rowcount ) { actualRowCount = queriesResponse[1].data[0].rowcount; + } else if (queriesResponse?.[0]?.sql_rowcount != null) { + actualRowCount = queriesResponse[0].sql_rowcount; } else { actualRowCount = exportFormData?.row_limit; } From 9a4900a31a5f56679215f062e23c3d6a0ec21ee5 Mon Sep 17 00:00:00 2001 From: amaannawab923 Date: Tue, 11 Nov 2025 13:33:08 +0530 Subject: [PATCH 47/57] fix: Streaming works on explore page & Toast appears only after the end user clicks on Download cta --- .../StreamingExportModal.tsx | 3 + superset-frontend/src/constants.ts | 7 ++ .../components/gridComponents/Chart/Chart.jsx | 17 ++- .../components/ExploreChartHeader/index.jsx | 11 +- .../useExploreAdditionalActionsMenu/index.jsx | 116 ++++++++++++++++-- 5 files changed, 136 insertions(+), 18 deletions(-) diff --git a/superset-frontend/src/components/StreamingExportModal/StreamingExportModal.tsx b/superset-frontend/src/components/StreamingExportModal/StreamingExportModal.tsx index c945025897af..0d3edacd03a1 100644 --- a/superset-frontend/src/components/StreamingExportModal/StreamingExportModal.tsx +++ b/superset-frontend/src/components/StreamingExportModal/StreamingExportModal.tsx @@ -55,6 +55,7 @@ interface StreamingExportModalProps { visible: boolean; onCancel: () => void; onRetry?: () => void; + onDownload?: () => void; progress: StreamingProgress; } @@ -334,6 +335,7 @@ const StreamingExportModal = ({ visible, onCancel, onRetry, + onDownload, progress, }: StreamingExportModalProps) => { const { status, downloadUrl, filename } = progress; @@ -348,6 +350,7 @@ const StreamingExportModal = ({ const handleDownload = () => { if (downloadUrl && filename) { triggerFileDownload(downloadUrl, filename); + onDownload?.(); // Call onDownload callback if provided onCancel(); } }; diff --git a/superset-frontend/src/constants.ts b/superset-frontend/src/constants.ts index 4d8d4590fed4..f3b701b2ca28 100644 --- a/superset-frontend/src/constants.ts +++ b/superset-frontend/src/constants.ts @@ -189,3 +189,10 @@ export enum Actions { CREATE = 'create', UPDATE = 'update', } + +/** + * Default threshold for CSV streaming export. + * Exports with row counts >= this value will use streaming with progress tracking. + * Exports with row counts < this value will use traditional download. + */ +export const DEFAULT_CSV_STREAMING_ROW_THRESHOLD = 100000; diff --git a/superset-frontend/src/dashboard/components/gridComponents/Chart/Chart.jsx b/superset-frontend/src/dashboard/components/gridComponents/Chart/Chart.jsx index 66e829e6a87a..5f92aecddc02 100644 --- a/superset-frontend/src/dashboard/components/gridComponents/Chart/Chart.jsx +++ b/superset-frontend/src/dashboard/components/gridComponents/Chart/Chart.jsx @@ -39,7 +39,10 @@ import { LOG_ACTIONS_FORCE_REFRESH_CHART, } from 'src/logger/LogUtils'; import { postFormData } from 'src/explore/exploreUtils/formData'; -import { URL_PARAMS } from 'src/constants'; +import { + URL_PARAMS, + DEFAULT_CSV_STREAMING_ROW_THRESHOLD, +} from 'src/constants'; import { enforceSharedLabelsColorsArray } from 'src/utils/colorScheme'; import exportPivotExcel from 'src/utils/downloadAsPivotExcel'; import { @@ -87,7 +90,6 @@ const propTypes = { const RESIZE_TIMEOUT = 500; const DEFAULT_HEADER_HEIGHT = 22; -const DEFAULT_CSV_STREAMING_ROW_THRESHOLD = 100000; const ChartWrapper = styled.div` overflow: hidden; @@ -199,14 +201,18 @@ const Chart = props => { retryExport, } = useStreamingExport({ onComplete: () => { - boundActionCreators.addSuccessToast( - t('CSV file downloaded successfully'), - ); + // Don't show toast here - wait for user to click Download button }, onError: () => { boundActionCreators.addDangerToast(t('Export failed - please try again')); }, }); + + const handleDownloadComplete = useCallback(() => { + boundActionCreators.addSuccessToast( + t('CSV file downloaded successfully'), + ); + }, [boundActionCreators]); const history = useHistory(); const resize = useCallback( debounce(() => { @@ -686,6 +692,7 @@ const Chart = props => { resetExport(); }} onRetry={retryExport} + onDownload={handleDownloadComplete} progress={progress} exportType="csv" /> diff --git a/superset-frontend/src/explore/components/ExploreChartHeader/index.jsx b/superset-frontend/src/explore/components/ExploreChartHeader/index.jsx index 18b24bfc6d90..e1dd82579288 100644 --- a/superset-frontend/src/explore/components/ExploreChartHeader/index.jsx +++ b/superset-frontend/src/explore/components/ExploreChartHeader/index.jsx @@ -39,6 +39,7 @@ import ReportModal from 'src/features/reports/ReportModal'; import { deleteActiveReport } from 'src/features/reports/ReportModal/actions'; import { useUnsavedChangesPrompt } from 'src/hooks/useUnsavedChangesPrompt'; import { getChartFormDiffs } from 'src/utils/getChartFormDiffs'; +import { StreamingExportModal } from 'src/components/StreamingExportModal'; import { useExploreAdditionalActionsMenu } from '../useExploreAdditionalActionsMenu'; import { useExploreMetadataBar } from './useExploreMetadataBar'; @@ -172,7 +173,7 @@ export const ExploreChartHeader = ({ [redirectSQLLab, history], ); - const [menu, isDropdownVisible, setIsDropdownVisible] = + const [menu, isDropdownVisible, setIsDropdownVisible, streamingExportState] = useExploreAdditionalActionsMenu( latestQueryFormData, canDownload, @@ -345,6 +346,14 @@ export const ExploreChartHeader = ({ onConfirmNavigation={handleConfirmNavigation} handleSave={handleSaveAndCloseModal} /> + + ); }; diff --git a/superset-frontend/src/explore/components/useExploreAdditionalActionsMenu/index.jsx b/superset-frontend/src/explore/components/useExploreAdditionalActionsMenu/index.jsx index a6e91e9a6469..559b894be0b8 100644 --- a/superset-frontend/src/explore/components/useExploreAdditionalActionsMenu/index.jsx +++ b/superset-frontend/src/explore/components/useExploreAdditionalActionsMenu/index.jsx @@ -36,6 +36,7 @@ import { } from '@superset-ui/core/components'; import { Menu } from '@superset-ui/core/components/Menu'; import { useToasts } from 'src/components/MessageToasts/withToasts'; +import { DEFAULT_CSV_STREAMING_ROW_THRESHOLD } from 'src/constants'; import { exportChart, getChartKey } from 'src/explore/exploreUtils'; import downloadAsImage from 'src/utils/downloadAsImage'; import { getChartPermalink } from 'src/utils/urlUtils'; @@ -50,6 +51,7 @@ import { LOG_ACTIONS_CHART_DOWNLOAD_AS_XLS, } from 'src/logger/LogUtils'; import exportPivotExcel from 'src/utils/downloadAsPivotExcel'; +import { useStreamingExport } from 'src/components/StreamingExportModal'; import ViewQueryModal from '../controls/ViewQueryModal'; import EmbedCodeContent from '../EmbedCodeContent'; import { useDashboardsMenuItems } from './DashboardsSubMenu'; @@ -140,6 +142,38 @@ export const useExploreAdditionalActionsMenu = ( const chart = useSelector( state => state.charts?.[getChartKey(state.explore)], ); + const streamingThreshold = useSelector( + state => + state.common?.conf?.CSV_STREAMING_ROW_THRESHOLD || + DEFAULT_CSV_STREAMING_ROW_THRESHOLD, + ); + + // Streaming export state and handlers + const [isStreamingModalVisible, setIsStreamingModalVisible] = useState(false); + const { + progress, + isExporting, + startExport, + cancelExport, + resetExport, + retryExport, + } = useStreamingExport({ + onComplete: () => { + // Don't show toast here - wait for user to click Download button + }, + onError: () => { + addDangerToast(t('Export failed - please try again')); + }, + }); + + const handleCloseStreamingModal = useCallback(() => { + setIsStreamingModalVisible(false); + resetExport(); + }, [resetExport]); + + const handleDownloadComplete = useCallback(() => { + addSuccessToast(t('CSV file downloaded successfully')); + }, [addSuccessToast]); // Use the updated report menu items hook const reportMenuItem = useHeaderReportMenuItems({ @@ -170,18 +204,67 @@ export const useExploreAdditionalActionsMenu = ( } }, [addDangerToast, latestQueryFormData]); - const exportCSV = useCallback( - () => - canDownloadCSV - ? exportChart({ - formData: latestQueryFormData, - ownState, - resultType: 'full', - resultFormat: 'csv', - }) + const exportCSV = useCallback(() => { + if (!canDownloadCSV) return null; + + // Determine row count for streaming threshold check + let actualRowCount; + const isTableViz = latestQueryFormData?.viz_type === 'table'; + const queriesResponse = chart?.queriesResponse; + + if ( + isTableViz && + queriesResponse?.length > 1 && + queriesResponse[1]?.data?.[0]?.rowcount + ) { + actualRowCount = queriesResponse[1].data[0].rowcount; + } else if (queriesResponse?.[0]?.sql_rowcount != null) { + actualRowCount = queriesResponse[0].sql_rowcount; + } else { + actualRowCount = latestQueryFormData?.row_limit; + } + + // Check if streaming should be used + const shouldUseStreaming = + actualRowCount && actualRowCount >= streamingThreshold; + + let filename; + if (shouldUseStreaming) { + const now = new Date(); + const date = now.toISOString().slice(0, 10); + const time = now.toISOString().slice(11, 19).replace(/:/g, ''); + const timestamp = `_${date}_${time}`; + const chartName = + slice?.slice_name || latestQueryFormData.viz_type || 'chart'; + const safeChartName = chartName.replace(/[^a-zA-Z0-9_-]/g, '_'); + filename = `${safeChartName}${timestamp}.csv`; + } + + return exportChart({ + formData: latestQueryFormData, + ownState, + resultType: 'full', + resultFormat: 'csv', + onStartStreamingExport: shouldUseStreaming + ? exportParams => { + setIsStreamingModalVisible(true); + startExport({ + ...exportParams, + filename, + expectedRows: actualRowCount, + }); + } : null, - [canDownloadCSV, latestQueryFormData], - ); + }); + }, [ + canDownloadCSV, + latestQueryFormData, + ownState, + chart, + streamingThreshold, + slice, + startExport, + ]); const exportCSVPivoted = useCallback( () => @@ -545,5 +628,14 @@ export const useExploreAdditionalActionsMenu = ( theme.sizeUnit, ]); - return [menu, isDropdownVisible, setIsDropdownVisible]; + // Return streaming modal state and handlers for parent to render + const streamingExportState = { + isVisible: isStreamingModalVisible, + progress, + onCancel: handleCloseStreamingModal, + onRetry: retryExport, + onDownload: handleDownloadComplete, + }; + + return [menu, isDropdownVisible, setIsDropdownVisible, streamingExportState]; }; From 7c9bd3e4cfa6e9d44acb6fc68d4f5239bf839a45 Mon Sep 17 00:00:00 2001 From: amaannawab923 Date: Tue, 11 Nov 2025 14:31:57 +0530 Subject: [PATCH 48/57] ci fix --- .../dashboard/components/gridComponents/Chart/Chart.jsx | 9 ++------- 1 file changed, 2 insertions(+), 7 deletions(-) diff --git a/superset-frontend/src/dashboard/components/gridComponents/Chart/Chart.jsx b/superset-frontend/src/dashboard/components/gridComponents/Chart/Chart.jsx index 5f92aecddc02..e3e83d717e1f 100644 --- a/superset-frontend/src/dashboard/components/gridComponents/Chart/Chart.jsx +++ b/superset-frontend/src/dashboard/components/gridComponents/Chart/Chart.jsx @@ -39,10 +39,7 @@ import { LOG_ACTIONS_FORCE_REFRESH_CHART, } from 'src/logger/LogUtils'; import { postFormData } from 'src/explore/exploreUtils/formData'; -import { - URL_PARAMS, - DEFAULT_CSV_STREAMING_ROW_THRESHOLD, -} from 'src/constants'; +import { URL_PARAMS, DEFAULT_CSV_STREAMING_ROW_THRESHOLD } from 'src/constants'; import { enforceSharedLabelsColorsArray } from 'src/utils/colorScheme'; import exportPivotExcel from 'src/utils/downloadAsPivotExcel'; import { @@ -209,9 +206,7 @@ const Chart = props => { }); const handleDownloadComplete = useCallback(() => { - boundActionCreators.addSuccessToast( - t('CSV file downloaded successfully'), - ); + boundActionCreators.addSuccessToast(t('CSV file downloaded successfully')); }, [boundActionCreators]); const history = useHistory(); const resize = useCallback( From 16e0c82c5641ca51467dc8244492f3e0969f8d12 Mon Sep 17 00:00:00 2001 From: amaannawab923 Date: Tue, 11 Nov 2025 14:44:56 +0530 Subject: [PATCH 49/57] Rebase and update imports --- .../components/StreamingExportModal/StreamingExportModal.tsx | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/superset-frontend/src/components/StreamingExportModal/StreamingExportModal.tsx b/superset-frontend/src/components/StreamingExportModal/StreamingExportModal.tsx index 0d3edacd03a1..3c78dbfa9026 100644 --- a/superset-frontend/src/components/StreamingExportModal/StreamingExportModal.tsx +++ b/superset-frontend/src/components/StreamingExportModal/StreamingExportModal.tsx @@ -16,7 +16,8 @@ * specific language governing permissions and limitations * under the License. */ -import { styled, t, useTheme } from '@superset-ui/core'; +import { t } from '@superset-ui/core'; +import { styled, useTheme } from '@apache-superset/core/ui'; import { Modal, Button, From cfb4f49fe6567de41c108bc0c2c28540aabd0982 Mon Sep 17 00:00:00 2001 From: amaannawab923 Date: Tue, 11 Nov 2025 15:21:10 +0530 Subject: [PATCH 50/57] update: modal title as per figma --- .../components/StreamingExportModal/StreamingExportModal.tsx | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/superset-frontend/src/components/StreamingExportModal/StreamingExportModal.tsx b/superset-frontend/src/components/StreamingExportModal/StreamingExportModal.tsx index 3c78dbfa9026..161908735b77 100644 --- a/superset-frontend/src/components/StreamingExportModal/StreamingExportModal.tsx +++ b/superset-frontend/src/components/StreamingExportModal/StreamingExportModal.tsx @@ -358,7 +358,7 @@ const StreamingExportModal = ({ return ( Date: Tue, 11 Nov 2025 15:45:45 +0530 Subject: [PATCH 51/57] fix test --- .../StreamingExportModal/StreamingExportModal.test.tsx | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/superset-frontend/src/components/StreamingExportModal/StreamingExportModal.test.tsx b/superset-frontend/src/components/StreamingExportModal/StreamingExportModal.test.tsx index 7663a8614e95..30ac1379125b 100644 --- a/superset-frontend/src/components/StreamingExportModal/StreamingExportModal.test.tsx +++ b/superset-frontend/src/components/StreamingExportModal/StreamingExportModal.test.tsx @@ -46,7 +46,7 @@ beforeEach(() => { test('renders modal with streaming state', () => { render(); - expect(screen.getByText('Exporting Data')).toBeInTheDocument(); + expect(screen.getByText('CSV Export')).toBeInTheDocument(); expect( screen.getByText(/Processing export for test_export.csv/i), ).toBeInTheDocument(); From b5393a409e8c385b8bd6ff9b029caec7ec04a10a Mon Sep 17 00:00:00 2001 From: amaannawab923 Date: Tue, 11 Nov 2025 16:33:24 +0530 Subject: [PATCH 52/57] fix flaky test --- .../FiltersConfigModal/FiltersConfigModal.test.tsx | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/superset-frontend/src/dashboard/components/nativeFilters/FiltersConfigModal/FiltersConfigModal.test.tsx b/superset-frontend/src/dashboard/components/nativeFilters/FiltersConfigModal/FiltersConfigModal.test.tsx index 6db7f63edb96..653f6a32fcfb 100644 --- a/superset-frontend/src/dashboard/components/nativeFilters/FiltersConfigModal/FiltersConfigModal.test.tsx +++ b/superset-frontend/src/dashboard/components/nativeFilters/FiltersConfigModal/FiltersConfigModal.test.tsx @@ -456,7 +456,7 @@ test('deletes a filter', async () => { }), ), ); -}); +}, 30000); // Increase timeout to 30 seconds for slow async operations test('deletes a filter including dependencies', async () => { const nativeFilterState = [ From ed5c1e31ac1e5cedbe364b238b92fc58af9fa8cc Mon Sep 17 00:00:00 2001 From: amaannawab923 Date: Wed, 12 Nov 2025 12:00:02 +0530 Subject: [PATCH 53/57] data type fix --- superset/charts/data/api.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/superset/charts/data/api.py b/superset/charts/data/api.py index b4f6bfeb072a..c30669055ad1 100644 --- a/superset/charts/data/api.py +++ b/superset/charts/data/api.py @@ -549,14 +549,17 @@ def _should_use_streaming( 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") + rowcount = data[0].get("rowcount") + actual_row_count = int(rowcount) if rowcount else None # Fallback to row_limit if actual count not available if actual_row_count is None: if form_data and "row_limit" in form_data: - actual_row_count = form_data.get("row_limit", 0) + row_limit = form_data.get("row_limit", 0) + actual_row_count = int(row_limit) if row_limit else 0 elif query_context.form_data and "row_limit" in query_context.form_data: - actual_row_count = query_context.form_data.get("row_limit", 0) + row_limit = query_context.form_data.get("row_limit", 0) + actual_row_count = int(row_limit) if row_limit else 0 # Use streaming if row count meets or exceeds threshold if actual_row_count is not None and actual_row_count >= threshold: From c9af828e9b105ef619d1794bb34c3a0fa5616f82 Mon Sep 17 00:00:00 2001 From: amaannawab923 Date: Tue, 18 Nov 2025 16:21:44 +0530 Subject: [PATCH 54/57] Comments Resolution --- .../src/SqlLab/components/ResultSet/index.tsx | 3 ++- .../StreamingExportModal.tsx | 3 +-- .../components/gridComponents/Chart/Chart.jsx | 4 +--- superset/charts/data/api.py | 19 +++++-------------- .../sql_lab/streaming_export_command.py | 5 +---- superset/commands/streaming_export/base.py | 7 +------ superset/sqllab/api.py | 5 ++--- 7 files changed, 13 insertions(+), 33 deletions(-) diff --git a/superset-frontend/src/SqlLab/components/ResultSet/index.tsx b/superset-frontend/src/SqlLab/components/ResultSet/index.tsx index c9a039f0ee4e..21a10492d21b 100644 --- a/superset-frontend/src/SqlLab/components/ResultSet/index.tsx +++ b/superset-frontend/src/SqlLab/components/ResultSet/index.tsx @@ -237,7 +237,7 @@ const ResultSet = ({ const logAction = useLogAction({ queryId, sqlEditorId: query.sqlEditorId }); const { showConfirm, ConfirmModal } = useConfirmModal(); - const { progress, startExport, resetExport, retryExport } = + const { progress, startExport, resetExport, retryExport, cancelExport } = useStreamingExport({ onComplete: () => {}, onError: error => { @@ -318,6 +318,7 @@ const ResultSet = ({ ensureAppRoot(`/api/v1/sqllab/export/${clientId}/`); const handleCloseStreamingModal = () => { + cancelExport(); setShowStreamingModal(false); resetExport(); }; diff --git a/superset-frontend/src/components/StreamingExportModal/StreamingExportModal.tsx b/superset-frontend/src/components/StreamingExportModal/StreamingExportModal.tsx index 161908735b77..4e2c31a29fc1 100644 --- a/superset-frontend/src/components/StreamingExportModal/StreamingExportModal.tsx +++ b/superset-frontend/src/components/StreamingExportModal/StreamingExportModal.tsx @@ -35,7 +35,6 @@ export enum ExportStatus { CANCELLED = 'cancelled', } -const MAX_PROGRESS_PERCENT = 99; const COMPLETED_PERCENT = 100; export interface StreamingProgress { @@ -200,7 +199,7 @@ const calculateProgressPercentage = ( if (!totalRows || totalRows <= 0 || !rowsProcessed) return 0; const percentage = (rowsProcessed / totalRows) * 100; - return Math.round(Math.min(MAX_PROGRESS_PERCENT, percentage)); + return Math.floor(percentage); }; const getProgressStatus = ( diff --git a/superset-frontend/src/dashboard/components/gridComponents/Chart/Chart.jsx b/superset-frontend/src/dashboard/components/gridComponents/Chart/Chart.jsx index 5ab56e357c16..0538f9b3551b 100644 --- a/superset-frontend/src/dashboard/components/gridComponents/Chart/Chart.jsx +++ b/superset-frontend/src/dashboard/components/gridComponents/Chart/Chart.jsx @@ -681,9 +681,7 @@ const Chart = props => { { - if (isExporting) { - cancelExport(); - } + cancelExport(); setIsStreamingModalVisible(false); resetExport(); }} diff --git a/superset/charts/data/api.py b/superset/charts/data/api.py index b8e53df1a8f5..6166876a36b1 100644 --- a/superset/charts/data/api.py +++ b/superset/charts/data/api.py @@ -489,10 +489,7 @@ def _get_data_response( # Log is_cached if extra payload callback is provided if add_extra_log_payload and result and "queries" in result: is_cached_values = [query.get("is_cached") for query in result["queries"]] - if len(is_cached_values) == 1: - add_extra_log_payload(is_cached=is_cached_values[0]) - elif is_cached_values: - add_extra_log_payload(is_cached=is_cached_values) + add_extra_log_payload(is_cached=is_cached_values) return self._send_chart_response( result, form_data, datasource, filename, expected_rows @@ -588,10 +585,7 @@ def _should_use_streaming( actual_row_count = int(row_limit) if row_limit else 0 # Use streaming if row count meets or exceeds threshold - if actual_row_count is not None and actual_row_count >= threshold: - return True - - return False + return actual_row_count is not None and actual_row_count >= threshold def _create_streaming_csv_response( self, @@ -616,16 +610,13 @@ def _create_streaming_csv_response( # Sanitize chart name for filename filename = secure_filename(f"superset_{chart_name}_{timestamp}.csv") - logger.info( - "Creating streaming CSV response: %s (from frontend: %s)", - filename, - filename is not None, - ) + logger.info("Creating streaming CSV response: %s", filename) if expected_rows: logger.info("Using expected_rows from frontend: %d", expected_rows) # Execute streaming command - chunk_size = 1000 + # TODO: Make chunk size configurable via SUPERSET_CONFIG + chunk_size = 1024 command = StreamingCSVExportCommand(query_context, chunk_size) command.validate() diff --git a/superset/commands/sql_lab/streaming_export_command.py b/superset/commands/sql_lab/streaming_export_command.py index 966bd9ab4215..6b6585ac442f 100644 --- a/superset/commands/sql_lab/streaming_export_command.py +++ b/superset/commands/sql_lab/streaming_export_command.py @@ -101,10 +101,7 @@ def _get_sql_and_database(self) -> tuple[str, Any]: database = self._query.database # Get the SQL query - if select_sql: - sql = select_sql - else: - sql = executed_sql + sql = select_sql or executed_sql return sql, database diff --git a/superset/commands/streaming_export/base.py b/superset/commands/streaming_export/base.py index 243ba8148835..8b525a2c8226 100644 --- a/superset/commands/streaming_export/base.py +++ b/superset/commands/streaming_export/base.py @@ -143,9 +143,7 @@ def _execute_query_and_stream( # Execute query with streaming with merged_database.get_sqla_engine() as engine: - connection = engine.connect() - - try: + with engine.connect() as connection: result_proxy = connection.execution_options( stream_results=True ).execute(text(sql)) @@ -182,9 +180,6 @@ def _execute_query_and_stream( total_time, ) - finally: - connection.close() - def run(self) -> Callable[[], Generator[str, None, None]]: """ Execute the streaming CSV export. diff --git a/superset/sqllab/api.py b/superset/sqllab/api.py index 2b51e6a92c48..d3cd123346bb 100644 --- a/superset/sqllab/api.py +++ b/superset/sqllab/api.py @@ -373,13 +373,12 @@ def _create_streaming_csv_response( ) -> Response: """Create a streaming CSV response for large SQL Lab result sets.""" # Execute streaming command - chunk_size = 1000 + # TODO: Make chunk size configurable via SUPERSET_CONFIG + chunk_size = 1024 command = StreamingSqlResultExportCommand(client_id, chunk_size) command.validate() if not filename: - query = command._query - assert query is not None timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") filename = secure_filename(f"sqllab_{client_id}_{timestamp}.csv") From db69437dc6a8149df501a944c164f1543f774dd3 Mon Sep 17 00:00:00 2001 From: amaannawab923 Date: Wed, 19 Nov 2025 10:16:25 +0530 Subject: [PATCH 55/57] test fix --- .../charts/data/api_tests.py | 26 +++++++++---------- 1 file changed, 13 insertions(+), 13 deletions(-) diff --git a/tests/integration_tests/charts/data/api_tests.py b/tests/integration_tests/charts/data/api_tests.py index 30569e02807b..d4816cbbdf60 100644 --- a/tests/integration_tests/charts/data/api_tests.py +++ b/tests/integration_tests/charts/data/api_tests.py @@ -803,13 +803,13 @@ def test_chart_data_post_is_cached_in_event_logger(self, mock_event_logger): payload_with_force["force"] = True self.post_assert_metric(CHART_DATA_URI, payload_with_force, "data") - # Check that is_cached was logged as None (not from cache) + # Check that is_cached was logged as [None] (not from cache) call_kwargs = mock_event_logger.call_args[1] records = call_kwargs.get("records", []) assert len(records) > 0 - # is_cached should be None when force=True (bypasses cache) + # is_cached should be [None] when force=True (bypasses cache) assert "is_cached" in records[0] - assert records[0]["is_cached"] is None + assert records[0]["is_cached"] == [None] # Reset mock for second request mock_event_logger.reset_mock() @@ -819,12 +819,12 @@ def test_chart_data_post_is_cached_in_event_logger(self, mock_event_logger): payload_without_force["force"] = False self.post_assert_metric(CHART_DATA_URI, payload_without_force, "data") - # Check that is_cached was logged as True (from cache) + # Check that is_cached was logged as [True] (from cache) call_kwargs = mock_event_logger.call_args[1] records = call_kwargs.get("records", []) assert len(records) > 0 - # is_cached should be True when retrieved from cache - assert records[0]["is_cached"] is True + # is_cached should be [True] when retrieved from cache + assert records[0]["is_cached"] == [True] @with_feature_flags(GLOBAL_ASYNC_QUERIES=True) @pytest.mark.usefixtures("load_birth_names_dashboard_with_slices") @@ -1323,14 +1323,14 @@ def test_chart_data_is_cached_in_event_logger(self, mock_event_logger): # First request - should not be cached (force=true bypasses cache) self.get_assert_metric(f"api/v1/chart/{chart.id}/data/?force=true", "get_data") - # Check that is_cached was logged as None (not from cache) + # Check that is_cached was logged as [None] (not from cache) call_kwargs = mock_event_logger.call_args[1] records = call_kwargs.get("records", []) assert len(records) > 0 - # is_cached should be None when force=true (bypasses cache) - # The field should exist but be None + # is_cached should be [None] when force=true (bypasses cache) + # The field should exist but contain [None] assert "is_cached" in records[0] - assert records[0]["is_cached"] is None + assert records[0]["is_cached"] == [None] # Reset mock for second request mock_event_logger.reset_mock() @@ -1338,12 +1338,12 @@ def test_chart_data_is_cached_in_event_logger(self, mock_event_logger): # Second request - should be cached self.get_assert_metric(f"api/v1/chart/{chart.id}/data/", "get_data") - # Check that is_cached was logged as True (from cache) + # Check that is_cached was logged as [True] (from cache) call_kwargs = mock_event_logger.call_args[1] records = call_kwargs.get("records", []) assert len(records) > 0 - # is_cached should be True when retrieved from cache - assert records[0]["is_cached"] is True + # is_cached should be [True] when retrieved from cache + assert records[0]["is_cached"] == [True] @pytest.mark.usefixtures("load_birth_names_dashboard_with_slices") @with_feature_flags(GLOBAL_ASYNC_QUERIES=True) From 8090401bce03c6cc4fbba72274c9970d3aa7bf9e Mon Sep 17 00:00:00 2001 From: amaannawab923 Date: Wed, 19 Nov 2025 12:24:57 +0530 Subject: [PATCH 56/57] test fix --- .../chart/streaming_export_command_test.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/tests/unit_tests/commands/chart/streaming_export_command_test.py b/tests/unit_tests/commands/chart/streaming_export_command_test.py index 2ac26dde82a3..6096eaf5d7c1 100644 --- a/tests/unit_tests/commands/chart/streaming_export_command_test.py +++ b/tests/unit_tests/commands/chart/streaming_export_command_test.py @@ -102,6 +102,9 @@ def test_csv_generation_with_small_dataset(mocker: MockerFixture) -> None: mock_connection.execution_options.return_value.execute.return_value = ( mock_result_proxy ) + mock_connection.__enter__.return_value = mock_connection + mock_connection.__exit__.return_value = None + mock_engine = mocker.MagicMock() mock_engine.connect.return_value = mock_connection datasource.database.get_sqla_engine.return_value.__enter__.return_value = ( @@ -137,6 +140,9 @@ def test_csv_generation_with_special_characters(mocker: MockerFixture) -> None: mock_connection = mocker.MagicMock() mock_connection.execution_options.return_value.execute.return_value = mock_result + mock_connection.__enter__.return_value = mock_connection + mock_connection.__exit__.return_value = None + mock_engine = mocker.MagicMock() mock_engine.connect.return_value = mock_connection datasource.database.get_sqla_engine.return_value.__enter__.return_value = ( @@ -167,6 +173,9 @@ def test_streaming_with_null_values(mocker: MockerFixture) -> None: mock_connection = mocker.MagicMock() mock_connection.execution_options.return_value.execute.return_value = mock_result + mock_connection.__enter__.return_value = mock_connection + mock_connection.__exit__.return_value = None + mock_engine = mocker.MagicMock() mock_engine.connect.return_value = mock_connection datasource.database.get_sqla_engine.return_value.__enter__.return_value = ( @@ -203,6 +212,8 @@ def test_streaming_execution_options_enabled(mocker: MockerFixture) -> None: mock_execution_options = mocker.MagicMock() mock_connection.execution_options.return_value = mock_execution_options mock_execution_options.execute.return_value = mock_result_proxy + mock_connection.__enter__.return_value = mock_connection + mock_connection.__exit__.return_value = None mock_engine = mocker.MagicMock() mock_engine.connect.return_value = mock_connection @@ -228,6 +239,9 @@ def test_empty_result_set(mocker: MockerFixture) -> None: mock_connection = mocker.MagicMock() mock_connection.execution_options.return_value.execute.return_value = mock_result + mock_connection.__enter__.return_value = mock_connection + mock_connection.__exit__.return_value = None + mock_engine = mocker.MagicMock() mock_engine.connect.return_value = mock_connection datasource.database.get_sqla_engine.return_value.__enter__.return_value = ( From dfc23fbff6dfe7631681df88f91b7b3f6f02cb8b Mon Sep 17 00:00:00 2001 From: amaannawab923 Date: Wed, 19 Nov 2025 12:44:36 +0530 Subject: [PATCH 57/57] fix sqllab tests --- .../sql_lab/streaming_export_command_test.py | 32 ++++++++++++++++++- 1 file changed, 31 insertions(+), 1 deletion(-) diff --git a/tests/unit_tests/commands/sql_lab/streaming_export_command_test.py b/tests/unit_tests/commands/sql_lab/streaming_export_command_test.py index e4e2f9d8dd6f..5c7d4ac2d482 100644 --- a/tests/unit_tests/commands/sql_lab/streaming_export_command_test.py +++ b/tests/unit_tests/commands/sql_lab/streaming_export_command_test.py @@ -151,6 +151,9 @@ def test_csv_generation_with_select_sql(mocker, mock_query, mock_result_proxy): mock_connection.execution_options.return_value.execute.return_value = ( mock_result_proxy ) + mock_connection.__enter__.return_value = mock_connection + mock_connection.__exit__.return_value = None + mock_engine = MagicMock() mock_engine.connect.return_value = mock_connection mock_query.database.get_sqla_engine.return_value.__enter__.return_value = ( @@ -200,6 +203,9 @@ def test_csv_generation_with_executed_sql_and_limit( mock_connection = MagicMock() mock_connection.execution_options.return_value.execute.return_value = mock_result + mock_connection.__enter__.return_value = mock_connection + mock_connection.__exit__.return_value = None + mock_engine = MagicMock() mock_engine.connect.return_value = mock_connection mock_query.database.get_sqla_engine.return_value.__enter__.return_value = ( @@ -232,6 +238,9 @@ def test_csv_generation_with_special_characters(mocker, mock_query): mock_connection = MagicMock() mock_connection.execution_options.return_value.execute.return_value = mock_result + mock_connection.__enter__.return_value = mock_connection + mock_connection.__exit__.return_value = None + mock_engine = MagicMock() mock_engine.connect.return_value = mock_connection mock_query.database.get_sqla_engine.return_value.__enter__.return_value = ( @@ -276,6 +285,9 @@ def test_limiting_factor_dropdown(mocker, mock_query): mock_connection.execution_options.return_value.execute.return_value = ( mock_result ) + mock_connection.__enter__.return_value = mock_connection + mock_connection.__exit__.return_value = None + mock_engine = MagicMock() mock_engine.connect.return_value = mock_connection mock_query.database.get_sqla_engine.return_value.__enter__.return_value = ( @@ -318,6 +330,9 @@ def test_limiting_factor_query_and_dropdown(mocker, mock_query): mock_connection.execution_options.return_value.execute.return_value = ( mock_result ) + mock_connection.__enter__.return_value = mock_connection + mock_connection.__exit__.return_value = None + mock_engine = MagicMock() mock_engine.connect.return_value = mock_connection mock_query.database.get_sqla_engine.return_value.__enter__.return_value = ( @@ -347,6 +362,9 @@ def test_empty_result_set(mocker, mock_query): mock_connection = MagicMock() mock_connection.execution_options.return_value.execute.return_value = mock_result + mock_connection.__enter__.return_value = mock_connection + mock_connection.__exit__.return_value = None + mock_engine = MagicMock() mock_engine.connect.return_value = mock_connection mock_query.database.get_sqla_engine.return_value.__enter__.return_value = ( @@ -402,6 +420,9 @@ def test_connection_is_closed_after_streaming(mocker, mock_query, mock_result_pr mock_connection.execution_options.return_value.execute.return_value = ( mock_result_proxy ) + mock_connection.__enter__.return_value = mock_connection + mock_connection.__exit__.return_value = None + mock_engine = MagicMock() mock_engine.connect.return_value = mock_connection mock_query.database.get_sqla_engine.return_value.__enter__.return_value = ( @@ -415,7 +436,8 @@ def test_connection_is_closed_after_streaming(mocker, mock_query, mock_result_pr generator = csv_generator_callable() list(generator) - mock_connection.close.assert_called_once() + # With context managers, __exit__ is called to cleanup the connection + mock_connection.__exit__.assert_called_once() def test_streaming_execution_options_enabled(mocker, mock_query, mock_result_proxy): @@ -428,6 +450,8 @@ def test_streaming_execution_options_enabled(mocker, mock_query, mock_result_pro mock_execution_options = Mock() mock_connection.execution_options.return_value = mock_execution_options mock_execution_options.execute.return_value = mock_result_proxy + mock_connection.__enter__.return_value = mock_connection + mock_connection.__exit__.return_value = None mock_engine = MagicMock() mock_engine.connect.return_value = mock_connection @@ -456,6 +480,9 @@ def test_completion_logging(mock_logger, mocker, mock_query, mock_result_proxy): mock_connection.execution_options.return_value.execute.return_value = ( mock_result_proxy ) + mock_connection.__enter__.return_value = mock_connection + mock_connection.__exit__.return_value = None + mock_engine = MagicMock() mock_engine.connect.return_value = mock_connection mock_query.database.get_sqla_engine.return_value.__enter__.return_value = ( @@ -490,6 +517,9 @@ def test_null_values_handling(mocker, mock_query): mock_connection = MagicMock() mock_connection.execution_options.return_value.execute.return_value = mock_result + mock_connection.__enter__.return_value = mock_connection + mock_connection.__exit__.return_value = None + mock_engine = MagicMock() mock_engine.connect.return_value = mock_connection mock_query.database.get_sqla_engine.return_value.__enter__.return_value = (