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 64c0a6db1d71..12a0504ce55e 100644 --- a/superset-frontend/packages/superset-ui-core/src/components/index.ts +++ b/superset-frontend/packages/superset-ui-core/src/components/index.ts @@ -145,6 +145,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'; diff --git a/superset-frontend/src/SqlLab/components/ResultSet/index.tsx b/superset-frontend/src/SqlLab/components/ResultSet/index.tsx index 3297aa7b9467..21a10492d21b 100644 --- a/superset-frontend/src/SqlLab/components/ResultSet/index.tsx +++ b/superset-frontend/src/SqlLab/components/ResultSet/index.tsx @@ -83,6 +83,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 { ensureAppRoot } from 'src/utils/pathUtils'; import { useConfirmModal } from 'src/hooks/useConfirmModal'; import ExploreCtasResultsButton from '../ExploreCtasResultsButton'; @@ -184,6 +186,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], [ @@ -224,12 +230,21 @@ 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 }); const { showConfirm, ConfirmModal } = useConfirmModal(); + const { progress, startExport, resetExport, retryExport, cancelExport } = + useStreamingExport({ + onComplete: () => {}, + onError: error => { + addDangerToast(t('Export failed: %s', error)); + }, + }); + const reRunQueryIfSessionTimeoutErrorOnMount = useCallback(() => { if ( query.errorMessage && @@ -302,6 +317,28 @@ const ResultSet = ({ const getExportCsvUrl = (clientId: string) => ensureAppRoot(`/api/v1/sqllab/export/${clientId}/`); + const handleCloseStreamingModal = () => { + cancelExport(); + setShowStreamingModal(false); + resetExport(); + }; + + 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); + + let actualRowCount = rowsCount; + + 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 { limitingFactor, queryLimit, results, rows } = query; @@ -372,9 +409,27 @@ const ResultSet = ({ { + const useStreaming = shouldUseStreamingExport(); + + if (useStreaming) { + e.preventDefault(); + setShowStreamingModal(true); + + startExport({ + url: '/api/v1/sqllab/export_streaming/', + payload: { client_id: query.id }, + exportType: 'csv', + expectedRows: rows, + }); + } else { + handleDownloadCsv(e); + } + }} > {t('Download to CSV')} @@ -723,43 +778,75 @@ const ResultSet = ({ + {ConfirmModal} ); } if (data && data.length === 0) { - return ; + return ( + <> + + + + ); } } if (query.cached || (query.state === QueryState.Success && !query.results)) { if (query.isDataPreview) { return ( - + <> + + + ); } if (query.resultsKey) { return ( - + <> + + + ); } } @@ -774,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/StreamingExportModal.test.tsx b/superset-frontend/src/components/StreamingExportModal/StreamingExportModal.test.tsx new file mode 100644 index 000000000000..30ac1379125b --- /dev/null +++ b/superset-frontend/src/components/StreamingExportModal/StreamingExportModal.test.tsx @@ -0,0 +1,244 @@ +/** + * 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('CSV Export')).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' })).toBeEnabled(); +}); + +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).toBeEnabled(); + + 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(); +}); diff --git a/superset-frontend/src/components/StreamingExportModal/StreamingExportModal.tsx b/superset-frontend/src/components/StreamingExportModal/StreamingExportModal.tsx new file mode 100644 index 000000000000..4e2c31a29fc1 --- /dev/null +++ b/superset-frontend/src/components/StreamingExportModal/StreamingExportModal.tsx @@ -0,0 +1,380 @@ +/** + * 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 { t } from '@superset-ui/core'; +import { styled, useTheme } from '@apache-superset/core/ui'; +import { + Modal, + Button, + Typography, + Progress, +} from '@superset-ui/core/components'; +import { Icons } from '@superset-ui/core/components/Icons'; + +const { Text } = Typography; + +export enum ExportStatus { + STREAMING = 'streaming', + COMPLETED = 'completed', + ERROR = 'error', + CANCELLED = 'cancelled', +} + +const COMPLETED_PERCENT = 100; + +export interface StreamingProgress { + totalRows?: number; + rowsProcessed: number; + totalSize: number; + status: ExportStatus; + downloadUrl?: string; + error?: string; + filename?: string; + speed?: number; + mbPerSecond?: number; + elapsedTime?: number; + retryCount?: number; +} + +interface StreamingExportModalProps { + visible: boolean; + onCancel: () => void; + onRetry?: () => void; + onDownload?: () => void; + progress: StreamingProgress; +} + +const ModalContent = styled.div` + ${({ theme }) => ` + padding: ${theme.sizeUnit * 4}px 0 ${theme.sizeUnit * 2}px; + `} +`; + +const ProgressSection = styled.div` + ${({ theme }) => ` + margin: ${theme.sizeUnit * 6}px 0; + position: relative; + `} +`; + +const ProgressWrapper = styled.div` + ${({ theme }) => ` + display: flex; + align-items: center; + gap: ${theme.sizeUnit * 3}px; + `} +`; + +const StyledProgress = styled(Progress)` + flex: 1; +`; + +const SuccessIcon = styled(Icons.CheckCircleFilled)` + ${({ theme }) => ` + color: ${theme.colorSuccess}; + font-size: ${theme.sizeUnit * 6}px; + flex-shrink: 0; + `} +`; + +const ErrorIconWrapper = styled.div` + ${({ 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)` + ${({ theme }) => ` + color: ${theme.colorWhite}; + font-size: ${theme.sizeUnit * 2.5}px; + `} +`; + +const ActionButtons = styled.div` + ${({ theme }) => ` + display: flex; + gap: ${theme.sizeUnit * 2}px; + justify-content: flex-end; + `} +`; + +const CenteredText = styled(Text)` + ${({ theme }) => ` + display: block; + text-align: center; + margin-top: ${theme.sizeUnit * 4}px; + `} +`; + +const ErrorText = styled(CenteredText)` + ${({ theme }) => ` + color: ${theme.colorError}; + `} +`; + +const CancelButton = styled(Button)` + ${({ 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.colorSuccessBg}; + color: ${theme.colorSuccess}; + border-color: ${theme.colorSuccess}; + } + `} +`; + +const DownloadButton = styled(Button)` + ${({ theme }) => ` + background-color: ${theme.colorSuccess}; + border-color: ${theme.colorSuccess}; + color: ${theme.colorWhite}; + + &:hover:not(:disabled) { + background-color: ${theme.colorSuccessActive}; + border-color: ${theme.colorSuccessActive}; + color: ${theme.colorWhite}; + } + + &:focus:not(:disabled) { + background-color: ${theme.colorSuccess}; + border-color: ${theme.colorSuccess}; + color: ${theme.colorWhite}; + } + + &:disabled { + 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.floor(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; + onCancel: () => void; + onRetry?: () => void; + onDownload: () => void; + getProgressPercentage: () => number; +} + +const ModalStateContent = ({ + status, + progress, + onCancel, + onRetry, + onDownload, + getProgressPercentage, +}: 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 ( + + + {hasIcon ? ( + + + {isError && ( + + + + )} + {isCompleted && } + + ) : ( + + )} + {isError ? ( + {messageText} + ) : ( + {messageText} + )} + + + {buttonText} + {shouldShowRetry ? ( + {t('Retry')} + ) : ( + + {t('Download')} + + )} + + + ); +}; + +const StreamingExportModal = ({ + visible, + onCancel, + onRetry, + onDownload, + progress, +}: StreamingExportModalProps) => { + const { status, downloadUrl, filename } = progress; + + const getProgressPercentage = (): number => + calculateProgressPercentage( + status, + progress.totalRows, + progress.rowsProcessed, + ); + + const handleDownload = () => { + if (downloadUrl && filename) { + triggerFileDownload(downloadUrl, filename); + onDownload?.(); // Call onDownload callback if provided + onCancel(); + } + }; + + return ( + + + + ); +}; + +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..7c342b4f298e --- /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'; 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..03241082dd1e --- /dev/null +++ b/superset-frontend/src/components/StreamingExportModal/useStreamingExport.test.ts @@ -0,0 +1,126 @@ +/** + * 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')), + }, +})); + +global.URL.createObjectURL = jest.fn(() => 'blob:mock-url'); +global.URL.revokeObjectURL = jest.fn(); + +global.fetch = jest.fn(); + +beforeEach(() => { + jest.clearAllMocks(); +}); + +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('useStreamingExport provides startExport function', () => { + const { result } = renderHook(() => useStreamingExport()); + + expect(typeof result.current.startExport).toBe('function'); +}); + +test('useStreamingExport provides resetExport function', () => { + const { result } = renderHook(() => useStreamingExport()); + + expect(typeof result.current.resetExport).toBe('function'); +}); + +test('useStreamingExport provides retryExport function', () => { + const { result } = renderHook(() => useStreamingExport()); + + expect(typeof result.current.retryExport).toBe('function'); +}); + +test('useStreamingExport provides cancelExport function', () => { + const { result } = renderHook(() => useStreamingExport()); + + expect(typeof result.current.cancelExport).toBe('function'); +}); + +test('useStreamingExport 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('useStreamingExport accepts onComplete callback option', () => { + const onComplete = jest.fn(); + const { result } = renderHook(() => useStreamingExport({ onComplete })); + + expect(result.current).toBeDefined(); +}); + +test('useStreamingExport accepts onError callback option', () => { + const onError = jest.fn(); + const { result } = renderHook(() => useStreamingExport({ onError })); + + expect(result.current).toBeDefined(); +}); + +test('useStreamingExport 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('useStreamingExport 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); +}); diff --git a/superset-frontend/src/components/StreamingExportModal/useStreamingExport.ts b/superset-frontend/src/components/StreamingExportModal/useStreamingExport.ts new file mode 100644 index 000000000000..fd0bf34cb2e3 --- /dev/null +++ b/superset-frontend/src/components/StreamingExportModal/useStreamingExport.ts @@ -0,0 +1,380 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +import { useState, useCallback, useRef, useEffect } from 'react'; +import { SupersetClient } from '@superset-ui/core'; +import { ExportStatus, StreamingProgress } from './StreamingExportModal'; + +interface UseStreamingExportOptions { + onComplete?: (downloadUrl: string, filename: string) => void; + onError?: (error: string) => void; +} + +interface StreamingExportPayload { + [key: string]: any; +} + +interface StreamingExportParams { + url: string; + payload: StreamingExportPayload; + filename?: string; + exportType: 'csv' | 'xlsx'; + expectedRows?: number; +} + +const NEWLINE_BYTE = 10; // '\n' character code + +const createFetchRequest = async ( + _url: string, + payload: StreamingExportPayload, + filename: string | undefined, + _exportType: string, + expectedRows: number | undefined, + signal: AbortSignal, +): Promise => { + const headers: Record = { + 'Content-Type': 'application/x-www-form-urlencoded', + }; + + // Get CSRF token using SupersetClient + const csrfToken = await SupersetClient.getCSRFToken(); + if (csrfToken) { + headers['X-CSRFToken'] = csrfToken; + } + + 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 + 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; + +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, + totalRows: undefined, + totalSize: 0, + speed: 0, + mbPerSecond: 0, + elapsedTime: 0, + status: ExportStatus.STREAMING, + }); + 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 })); + }, []); + + const executeExport = useCallback( + async (params: StreamingExportParams) => { + const { url, payload, filename, exportType, expectedRows } = params; + if (isExportingRef.current) { + return; + } + isExportingRef.current = true; + + abortControllerRef.current = new AbortController(); + + updateProgress({ + rowsProcessed: 0, + totalRows: expectedRows, + totalSize: 0, + speed: 0, + mbPerSecond: 0, + elapsedTime: 0, + status: ExportStatus.STREAMING, + filename, + }); + + try { + const fetchOptions = await createFetchRequest( + url, + payload, + filename, + exportType, + expectedRows, + abortControllerRef.current.signal, + ); + const response = await fetch(url, fetchOptions); + + 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'); + } + + 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; + 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 (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.'; + + // Update progress to show error with current progress preserved + updateProgress({ + status: ExportStatus.ERROR, + error: errorMsg, + rowsProcessed, + totalRows: expectedRows, + totalSize: receivedLength, + }); + + isExportingRef.current = false; + options.onError?.(errorMsg); + hasError = true; + break; + } + + chunks.push(value); + receivedLength += value.length; + + // 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({ + status: ExportStatus.STREAMING, + rowsProcessed, + totalRows: expectedRows, + totalSize: receivedLength, + filename: serverFilename, + }); + } + + // Check if we exited early due to error marker + if (hasError) { + return; + } + + 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, + downloadUrl, + filename: serverFilename, + }); + + isExportingRef.current = false; + options.onComplete?.(downloadUrl, serverFilename); + } catch (error) { + const errorMessage = + error instanceof Error ? error.message : 'Unknown error occurred'; + + if ( + errorMessage.includes('cancelled') || + errorMessage.includes('aborted') + ) { + updateProgress({ + status: ExportStatus.CANCELLED, + }); + isExportingRef.current = false; + } else { + updateProgress({ + status: ExportStatus.ERROR, + error: errorMessage, + }); + options.onError?.(errorMessage); + isExportingRef.current = false; + } + } finally { + abortControllerRef.current = null; + } + }, + [updateProgress, options], + ); + + const startExport = useCallback( + async (params: StreamingExportParams) => { + if (isExportingRef.current) { + return; + } + + 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); + }, + [updateProgress, executeExport], + ); + + const retryExport = useCallback(() => { + if (!lastExportParamsRef.current) { + return; + } + + if (isExportingRef.current) { + return; + } + + setRetryCount(0); + executeExport(lastExportParamsRef.current); + }, [executeExport]); + + const cancelExport = useCallback(() => { + if (abortControllerRef.current) { + abortControllerRef.current.abort(); + updateProgress({ + status: ExportStatus.CANCELLED, + }); + } + }, [updateProgress]); + + const resetExport = useCallback(() => { + if (currentBlobUrlRef.current) { + URL.revokeObjectURL(currentBlobUrlRef.current); + currentBlobUrlRef.current = null; + } + + isExportingRef.current = false; + abortControllerRef.current = null; + setProgress({ + rowsProcessed: 0, + totalRows: undefined, + totalSize: 0, + speed: 0, + mbPerSecond: 0, + elapsedTime: 0, + status: ExportStatus.STREAMING, + }); + }, []); + + // Cleanup blob URL on unmount to prevent memory leak + useEffect( + () => () => { + if (currentBlobUrlRef.current) { + URL.revokeObjectURL(currentBlobUrlRef.current); + } + }, + [], + ); + + return { + progress, + isExporting: isExportingRef.current, + retryCount, + startExport, + cancelExport, + resetExport, + retryExport, + }; +}; diff --git a/superset-frontend/src/constants.ts b/superset-frontend/src/constants.ts index 8bbd14e93c1e..fddd72288a0e 100644 --- a/superset-frontend/src/constants.ts +++ b/superset-frontend/src/constants.ts @@ -194,3 +194,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 4fb02867d862..f8307626103e 100644 --- a/superset-frontend/src/dashboard/components/gridComponents/Chart/Chart.jsx +++ b/superset-frontend/src/dashboard/components/gridComponents/Chart/Chart.jsx @@ -28,6 +28,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, @@ -36,7 +40,7 @@ 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 { @@ -82,8 +86,6 @@ 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; @@ -164,6 +166,11 @@ 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 || + DEFAULT_CSV_STREAMING_ROW_THRESHOLD, + ); const datasource = useSelector( state => (chart && @@ -182,6 +189,27 @@ const Chart = props => { const [descriptionHeight, setDescriptionHeight] = useState(0); const [height, setHeight] = useState(props.height); const [width, setWidth] = useState(props.width); + + 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: () => { + 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(() => { @@ -425,6 +453,39 @@ const Chart = props => { is_cached: isCached, }); + const exportFormData = isFullCSV + ? { ...formData, row_limit: maxRows } + : 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 if (queriesResponse?.[0]?.sql_rowcount != null) { + actualRowCount = queriesResponse[0].sql_rowcount; + } else { + actualRowCount = exportFormData?.row_limit; + } + + // Handle streaming CSV exports based on row threshold + const shouldUseStreaming = + format === 'csv' && !isPivot && 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 || formData.viz_type || 'chart'; + const safeChartName = chartName.replace(/[^a-zA-Z0-9_-]/g, '_'); + filename = `${safeChartName}${timestamp}.csv`; + } let ownState = dataMask[props.id]?.ownState || {}; // Convert chart-specific state to backend format using registered converter @@ -440,11 +501,21 @@ const Chart = props => { } exportChart({ - formData: isFullCSV ? { ...formData, row_limit: maxRows } : formData, - resultType: isPivot ? 'post_processed' : 'full', + formData: exportFormData, + resultType, resultFormat: format, force: true, ownState, + onStartStreamingExport: shouldUseStreaming + ? exportParams => { + setIsStreamingModalVisible(true); + startExport({ + ...exportParams, + filename, + expectedRows: actualRowCount, + }); + } + : null, }); }, [ @@ -457,6 +528,10 @@ const Chart = props => { chartState, props.id, boundActionCreators.logEvent, + queriesResponse, + startExport, + resetExport, + streamingThreshold, ], ); @@ -609,6 +684,19 @@ const Chart = props => { onChartStateChange={handleChartStateChange} /> + + { + cancelExport(); + setIsStreamingModalVisible(false); + resetExport(); + }} + onRetry={retryExport} + onDownload={handleDownloadComplete} + progress={progress} + exportType="csv" + /> ); }; 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 = [ diff --git a/superset-frontend/src/explore/components/ExploreChartHeader/index.jsx b/superset-frontend/src/explore/components/ExploreChartHeader/index.jsx index 2bd18b35ac2d..147148c567c2 100644 --- a/superset-frontend/src/explore/components/ExploreChartHeader/index.jsx +++ b/superset-frontend/src/explore/components/ExploreChartHeader/index.jsx @@ -40,6 +40,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'; @@ -173,7 +174,7 @@ export const ExploreChartHeader = ({ [redirectSQLLab, history], ); - const [menu, isDropdownVisible, setIsDropdownVisible] = + const [menu, isDropdownVisible, setIsDropdownVisible, streamingExportState] = useExploreAdditionalActionsMenu( latestQueryFormData, canDownload, @@ -346,6 +347,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 56762424cb90..b85817f698e8 100644 --- a/superset-frontend/src/explore/components/useExploreAdditionalActionsMenu/index.jsx +++ b/superset-frontend/src/explore/components/useExploreAdditionalActionsMenu/index.jsx @@ -29,6 +29,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'; @@ -43,6 +44,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'; @@ -133,6 +135,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({ @@ -163,18 +197,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( () => @@ -538,5 +621,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]; }; diff --git a/superset-frontend/src/explore/exploreUtils/index.js b/superset-frontend/src/explore/exploreUtils/index.js index d3bd62ef412a..1ebde5a3783b 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,18 @@ export const exportChart = async ({ }); } - SupersetClient.postForm(url, { form_data: safeStringify(payload) }); + // Check if streaming export handler is provided (from dashboard Chart.jsx) + if (onStartStreamingExport) { + // 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 + SupersetClient.postForm(url, { form_data: safeStringify(payload) }); + } }; export const exploreChart = (formData, requestParams) => { diff --git a/superset-frontend/webpack.proxy-config.js b/superset-frontend/webpack.proxy-config.js index 41835d88faf0..f15b79326960 100644 --- a/superset-frontend/webpack.proxy-config.js +++ b/superset-frontend/webpack.proxy-config.js @@ -165,7 +165,26 @@ module.exports = newManifest => { if (isHTML(response)) { processHTML(proxyResponse, response); } else { - proxyResponse.pipe(response); + const isCSV = (proxyResponse.headers['content-type'] || '').includes( + 'text/csv', + ); + + if (isCSV) { + 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 08d3b0b97c18..6166876a36b1 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, Callable, 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, @@ -268,11 +273,15 @@ def data( # noqa: C901 return self._run_async(json_body, command, add_extra_log_payload) form_data = json_body.get("form_data") + filename, expected_rows = self._extract_export_params_from_request() + return self._get_data_response( - command=command, + command, form_data=form_data, datasource=query_context.datasource, add_extra_log_payload=add_extra_log_payload, + filename=filename, + expected_rows=expected_rows, ) @expose("/data/", methods=("GET",)) @@ -369,6 +378,8 @@ 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, + expected_rows: int | None = None, ) -> Response: result_type = result["query_context"].result_type result_format = result["query_context"].result_format @@ -389,6 +400,12 @@ 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, expected_rows=expected_rows + ) + if len(result["queries"]) == 1: # return single query results data = result["queries"][0]["data"] @@ -457,6 +474,8 @@ def _get_data_response( force_cached: bool = False, form_data: dict[str, Any] | None = None, datasource: BaseDatasource | Query | None = None, + filename: str | None = None, + expected_rows: int | None = None, add_extra_log_payload: Callable[..., None] | None = None, ) -> Response: """Get data response and optionally log is_cached information.""" @@ -467,10 +486,30 @@ def _get_data_response( except ChartDataQueryFailedError as exc: return self.response_400(message=exc.message) - # Log is_cached if extra payload callback is provided - self._log_is_cached(result, add_extra_log_payload) + # 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"]] + add_extra_log_payload(is_cached=is_cached_values) + + 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) - return self._send_chart_response(result, form_data, datasource) + 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]: @@ -507,3 +546,99 @@ def _create_query_context_from_form( return ChartDataQueryContextSchema().load(form_data) except KeyError as ex: raise ValidationError("Request is incorrect") from ex + + def _should_use_streaming( + self, result: dict[Any, Any], form_data: dict[str, Any] | None = None + ) -> bool: + """Determine if streaming should be used based on actual row count threshold.""" + query_context = result["query_context"] + result_format = query_context.result_format + + # Only support CSV streaming currently + if result_format.lower() != "csv": + return False + + # Get streaming threshold from config + threshold = app.config.get("CSV_STREAMING_ROW_THRESHOLD", 100000) + + # Extract actual row count (same logic as frontend) + actual_row_count: 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 + 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: + 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: + 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: + 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 + return actual_row_count is not None and actual_row_count >= threshold + + 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, + ) -> Response: + """Create a streaming CSV response for large datasets.""" + 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 + filename = secure_filename(f"superset_{chart_name}_{timestamp}.csv") + + logger.info("Creating streaming CSV response: %s", filename) + if expected_rows: + logger.info("Using expected_rows from frontend: %d", expected_rows) + + # Execute streaming command + # TODO: Make chunk size configurable via SUPERSET_CONFIG + chunk_size = 1024 + 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 + }, + 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..b6ec3a36698a --- /dev/null +++ b/superset/commands/chart/data/streaming_export_command.py @@ -0,0 +1,80 @@ +# 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 chart data.""" + +from __future__ import annotations + +from typing import Any, TYPE_CHECKING + +from superset.commands.streaming_export.base import BaseStreamingCSVExportCommand + +if TYPE_CHECKING: + from superset.common.query_context import QueryContext + + +class StreamingCSVExportCommand(BaseStreamingCSVExportCommand): + """ + Command to execute a streaming CSV export for chart data. + + This command handles chart-specific logic: + - QueryContext validation + - Datasource preparation and SQL generation + - No row limit (exports all chart data) + """ + + def __init__( + self, + query_context: QueryContext, + chunk_size: int = 1000, + ): + """ + 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 + + def validate(self) -> None: + """Validate permissions and query context.""" + self._query_context.raise_for_access() + + def _get_sql_and_database(self) -> tuple[str, Any]: + """ + Get the SQL query and database for chart export. + + Returns: + Tuple of (sql_query, database_object) + """ + # 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()) + + return sql_query, datasource.database + + def _get_row_limit(self) -> int | None: + """ + Get the row limit for chart export. + + Returns: + None (no limit for chart exports) + """ + return None 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..6b6585ac442f --- /dev/null +++ b/superset/commands/sql_lab/streaming_export_command.py @@ -0,0 +1,142 @@ +# 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 + +from typing import Any + +from flask_babel import gettext as __ + +from superset import db +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 + + +class StreamingSqlResultExportCommand(BaseStreamingCSVExportCommand): + """ + Command to execute a streaming CSV export of SQL Lab query results. + + This command handles SQL Lab-specific logic: + - Query validation and access control + - SQL parsing and limit extraction + - LimitingFactor-based row limit adjustment + """ + + def __init__( + self, + client_id: str, + chunk_size: int = 1000, + ): + """ + 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._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 _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 + database = self._query.database + + # Get the SQL query + sql = select_sql or 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 + 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() + + # Adjust limit based on limiting factor + 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 + + 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..8b525a2c8226 --- /dev/null +++ b/superset/commands/streaming_export/base.py @@ -0,0 +1,214 @@ +# 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: + with engine.connect() as connection: + 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, + ) + + 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/superset/config.py b/superset/config.py index 4912fc4d954e..93ec7091ff6c 100644 --- a/superset/config.py +++ b/superset/config.py @@ -1009,6 +1009,12 @@ 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/sqllab/api.py b/superset/sqllab/api.py index 906dd72bcaf2..d3cd123346bb 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 @@ -294,6 +299,118 @@ 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, + *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.""" + # Execute streaming command + # TODO: Make chunk size configurable via SUPERSET_CONFIG + chunk_size = 1024 + command = StreamingSqlResultExportCommand(client_id, chunk_size) + command.validate() + + if not filename: + timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") + filename = secure_filename(f"sqllab_{client_id}_{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 + }, + 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 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/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) 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..6096eaf5d7c1 --- /dev/null +++ b/tests/unit_tests/commands/chart/streaming_export_command_test.py @@ -0,0 +1,258 @@ +# 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 _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 + datasource.database = mocker.MagicMock() + query_context.datasource = datasource + query_context.queries = [mocker.MagicMock()] + mock_session.merge.return_value = datasource.database + + return mock_db, query_context, datasource + + +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, query_context, datasource = _setup_chart_mocks(mocker) + + 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_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 = ( + 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, query_context, datasource = _setup_chart_mocks(mocker) + + 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_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 = ( + 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, query_context, datasource = _setup_chart_mocks(mocker) + + 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_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 = ( + 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, query_context, datasource = _setup_chart_mocks(mocker) + + 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_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 = ( + 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, query_context, datasource = _setup_chart_mocks(mocker) + + 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_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 = ( + 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..5c7d4ac2d482 --- /dev/null +++ b/tests/unit_tests/commands/sql_lab/streaming_export_command_test.py @@ -0,0 +1,540 @@ +# 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 pytest_mock import MockerFixture + +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 + + +def _setup_sqllab_mocks( + mocker: MockerFixture, mock_query: MagicMock +) -> tuple[MagicMock, MagicMock]: + """Set up common mocks for SQL Lab streaming export tests.""" + mock_db_base = mocker.patch("superset.commands.streaming_export.base.db") + mock_session = MagicMock() + mock_db_base.session.return_value.__enter__.return_value = mock_session + mock_session.merge.return_value = mock_query.database + + 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 +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() + + +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_db, mock_session = _setup_sqllab_mocks(mocker, mock_query) + + mock_connection = MagicMock() + 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 = ( + mock_engine + ) + + 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.SQLScript") +def test_csv_generation_with_executed_sql_and_limit( + mock_sqlscript, mocker, 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_db, mock_session = _setup_sqllab_mocks(mocker, mock_query) + + 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_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 = ( + mock_engine + ) + + 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) + + +def test_csv_generation_with_special_characters(mocker, 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_db, mock_session = _setup_sqllab_mocks(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 = ( + mock_engine + ) + + 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 + + +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" + 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_db, mock_session = _setup_sqllab_mocks(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 = ( + mock_engine + ) + + 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 + + +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" + 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_db, mock_session = _setup_sqllab_mocks(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 = ( + mock_engine + ) + + 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 + + +def test_empty_result_set(mocker, 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_db, mock_session = _setup_sqllab_mocks(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 = ( + mock_engine + ) + + 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" + + +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_base = mocker.patch("superset.commands.streaming_export.base.db") + mock_session = MagicMock() + mock_db_base.session.return_value.__enter__.return_value = mock_session + mock_session.merge.side_effect = Exception("Database connection failed") + + 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() + + 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 + + +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_db, mock_session = _setup_sqllab_mocks(mocker, mock_query) + + mock_connection = MagicMock() + 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 = ( + mock_engine + ) + + command = StreamingSqlResultExportCommand("test_client_123") + command.validate() + + csv_generator_callable = command.run() + generator = csv_generator_callable() + list(generator) + + # 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): + """Test that streaming execution options are enabled.""" + mock_query.select_sql = "SELECT * FROM test" + + mock_db, mock_session = _setup_sqllab_mocks(mocker, mock_query) + + 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_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 = ( + mock_engine + ) + + 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.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_db, mock_session = _setup_sqllab_mocks(mocker, mock_query) + + mock_connection = MagicMock() + 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 = ( + mock_engine + ) + + 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 "Streaming CSV completed" in log_message + assert "rows" in log_message + + +def test_null_values_handling(mocker, 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_db, mock_session = _setup_sqllab_mocks(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 = ( + mock_engine + ) + + 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