LevisNgigi commented on code in PR #35478: URL: https://github.com/apache/superset/pull/35478#discussion_r2419425306
########## superset-frontend/src/components/StreamingExportModal/StreamingExportModal.tsx: ########## @@ -0,0 +1,403 @@ +/** + * 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. + */ +/* eslint-disable theme-colors/no-literal-colors */ +import { styled, t } from '@superset-ui/core'; +import { Modal, Button, Typography, Progress } from 'antd'; +import { Icons } from '@superset-ui/core/components/Icons'; + +const { Text } = Typography; + +export enum ExportStatus { + STREAMING = 'streaming', + COMPLETED = 'completed', + ERROR = 'error', + CANCELLED = 'cancelled', +} + +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; + progress: StreamingProgress; +} + +const ModalContent = styled.div` + padding: ${({ theme }) => theme.sizeUnit * 4}px 0 + ${({ theme }) => theme.sizeUnit * 2}px; +`; + +const ProgressSection = styled.div` + margin: ${({ theme }) => theme.sizeUnit * 6}px 0; + position: relative; +`; + +const ProgressWrapper = styled.div` + display: flex; + align-items: center; + gap: ${({ theme }) => theme.sizeUnit * 3}px; +`; + +const SuccessIcon = styled(Icons.CheckCircleFilled)` + color: #52c41a; + font-size: 24px; + flex-shrink: 0; +`; + +const ErrorIconWrapper = styled.div` + display: flex; + align-items: center; + justify-content: center; + width: 16px; + height: 16px; + background-color: #ff4d4f; + border-radius: 50%; + flex-shrink: 0; +`; + +const ErrorIconStyled = styled(Icons.CloseOutlined)` + color: white; + font-size: 10px; +`; + +const ActionButtons = styled.div` + display: flex; + gap: ${({ theme }) => theme.sizeUnit * 2}px; + justify-content: flex-end; +`; + +const ProgressText = styled(Text)` + display: block; + text-align: center; + margin-top: ${({ theme }) => theme.sizeUnit * 4}px; +`; + +const ErrorText = styled(Text)` + display: block; + text-align: center; + margin-top: ${({ theme }) => theme.sizeUnit * 4}px; +`; + +const CancelButton = styled(Button)` + background-color: #f0fff8; + color: #1c997a; + border-color: #f0fff8; + + &:hover { + background-color: #f0fff8; + color: #1c997a; + border-color: #1c997a; + } + + &:focus { + background-color: #f0fff8; + color: #1c997a; + border-color: #1c997a; + } +`; + +const DownloadButton = styled(Button)` + &.ant-btn-primary { + background-color: #2ec196; + border-color: #2ec196; + color: #ffffff; + + &:hover:not(:disabled) { + background-color: #26a880; + border-color: #26a880; + color: #ffffff; + } + + &:focus:not(:disabled) { + background-color: #2ec196; + border-color: #2ec196; + color: #ffffff; + } + + &:disabled { + background-color: #f2f2f2; + border-color: #f2f2f2; + color: #b5b5b5; + } + } +`; + +interface ModalStateContentProps { + status: ExportStatus; + progress: StreamingProgress; + onCancel: () => void; + onRetry?: () => void; + onDownload: () => void; + getProgressPercentage: () => number; +} + +const ErrorContent = ({ + error, + onCancel, + onRetry, + getProgressPercentage, +}: { + error?: string; + onCancel: () => void; + onRetry?: () => void; + getProgressPercentage: () => number; +}) => ( + <ModalContent> + <ProgressSection> + <ProgressWrapper> + <Progress + percent={getProgressPercentage()} + status="exception" + showInfo={false} + style={{ flex: 1 }} + /> + <ErrorIconWrapper> + <ErrorIconStyled /> + </ErrorIconWrapper> + </ProgressWrapper> + <ErrorText type="danger">{error || t('Export failed')}</ErrorText> + </ProgressSection> + <ActionButtons> + <CancelButton onClick={onCancel}>{t('Close')}</CancelButton> + {onRetry && ( + <DownloadButton type="primary" onClick={onRetry}> + {t('Retry')} + </DownloadButton> + )} + </ActionButtons> + </ModalContent> +); + +const CancelledContent = ({ + getProgressPercentage, + onCancel, + onRetry, +}: { + getProgressPercentage: () => number; + onCancel: () => void; + onRetry?: () => void; +}) => ( + <ModalContent> + <ProgressSection> + <Progress + percent={getProgressPercentage()} + status="exception" + showInfo={false} + /> + <ProgressText>{t('Export cancelled')}</ProgressText> + </ProgressSection> + <ActionButtons> + <CancelButton onClick={onCancel}>{t('Close')}</CancelButton> + {onRetry && ( + <DownloadButton type="primary" onClick={onRetry}> Review Comment: Your DownloadButton already handles primary styling in line 128, might not need type="primary" here ########## superset-frontend/src/components/StreamingExportModal/useStreamingExport.ts: ########## @@ -0,0 +1,366 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +import { useState, useCallback, useRef, useEffect } from 'react'; +import { SupersetClient } from '@superset-ui/core'; +import { ExportStatus, StreamingProgress } from './StreamingExportModal'; + +interface UseStreamingExportOptions { + onComplete?: (downloadUrl: string, filename: string) => void; + onError?: (error: string) => void; +} + +interface StreamingExportPayload { + [key: string]: unknown; +} + +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, + _exportType: string, + expectedRows: number | undefined, + signal: AbortSignal, +): Promise<RequestInit> => { + const headers: Record<string, string> = { + 'Content-Type': 'application/x-www-form-urlencoded', + }; + + // Get CSRF token using SupersetClient + const csrfToken = await SupersetClient.getCSRFToken(); + if (csrfToken) { + headers['X-CSRFToken'] = csrfToken; + } + + // Build form data - if payload has client_id, it's SQL Lab export + // Otherwise it's a chart export with form_data + const formParams: Record<string, string> = { + filename, + 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<StreamingProgress>({ + rowsProcessed: 0, + totalRows: undefined, + totalSize: 0, + speed: 0, + mbPerSecond: 0, + elapsedTime: 0, + status: ExportStatus.STREAMING, + }); + const [retryCount, setRetryCount] = useState(0); + const abortControllerRef = useRef<AbortController | null>(null); + const lastExportParamsRef = useRef<StreamingExportParams | null>(null); + const currentBlobUrlRef = useRef<string | null>(null); + const isExportingRef = useRef(false); + + const updateProgress = useCallback((updates: Partial<StreamingProgress>) => { + setProgress(prev => ({ ...prev, ...updates })); + }, []); + + const executeExport = useCallback( + async (params: StreamingExportParams) => { + const { url, payload, filename, exportType, expectedRows } = params; + + abortControllerRef.current = new AbortController(); + + updateProgress({ + rowsProcessed: 0, + totalRows: expectedRows, + totalSize: 0, + speed: 0, + mbPerSecond: 0, + elapsedTime: 0, + status: ExportStatus.STREAMING, + filename, + }); + + try { + const defaultFilename = `export.${exportType}`; + const finalFilename = filename || defaultFilename; + + const fetchOptions = await createFetchRequest( + url, + payload, + finalFilename, + exportType, + expectedRows, + abortControllerRef.current.signal, + ); + + const response = await fetch(url, fetchOptions); + + if (!response.ok) { + throw new Error( + `Export failed: ${response.status} ${response.statusText}`, + ); + } + + if (!response.body) { + throw new Error('Response body is not available for streaming'); + } + + 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(); + Review Comment: Is there a way we could avoid disabling these ESLint rules here?Maybe we can refactor the loop logic ########## superset-frontend/src/components/StreamingExportModal/StreamingExportModal.tsx: ########## @@ -0,0 +1,403 @@ +/** + * 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. + */ +/* eslint-disable theme-colors/no-literal-colors */ +import { styled, t } from '@superset-ui/core'; +import { Modal, Button, Typography, Progress } from 'antd'; +import { Icons } from '@superset-ui/core/components/Icons'; + +const { Text } = Typography; + +export enum ExportStatus { + STREAMING = 'streaming', + COMPLETED = 'completed', + ERROR = 'error', + CANCELLED = 'cancelled', +} + +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; + progress: StreamingProgress; +} + +const ModalContent = styled.div` + padding: ${({ theme }) => theme.sizeUnit * 4}px 0 + ${({ theme }) => theme.sizeUnit * 2}px; +`; + +const ProgressSection = styled.div` + margin: ${({ theme }) => theme.sizeUnit * 6}px 0; + position: relative; +`; + +const ProgressWrapper = styled.div` + display: flex; + align-items: center; + gap: ${({ theme }) => theme.sizeUnit * 3}px; +`; + +const SuccessIcon = styled(Icons.CheckCircleFilled)` + color: #52c41a; + font-size: 24px; + flex-shrink: 0; +`; + +const ErrorIconWrapper = styled.div` + display: flex; + align-items: center; + justify-content: center; + width: 16px; + height: 16px; + background-color: #ff4d4f; + border-radius: 50%; + flex-shrink: 0; +`; + +const ErrorIconStyled = styled(Icons.CloseOutlined)` + color: white; + font-size: 10px; +`; + +const ActionButtons = styled.div` + display: flex; + gap: ${({ theme }) => theme.sizeUnit * 2}px; + justify-content: flex-end; +`; + +const ProgressText = styled(Text)` + display: block; + text-align: center; + margin-top: ${({ theme }) => theme.sizeUnit * 4}px; +`; + +const ErrorText = styled(Text)` + display: block; + text-align: center; + margin-top: ${({ theme }) => theme.sizeUnit * 4}px; +`; + +const CancelButton = styled(Button)` + background-color: #f0fff8; + color: #1c997a; + border-color: #f0fff8; + + &:hover { + background-color: #f0fff8; + color: #1c997a; + border-color: #1c997a; + } + + &:focus { + background-color: #f0fff8; + color: #1c997a; + border-color: #1c997a; + } +`; + +const DownloadButton = styled(Button)` + &.ant-btn-primary { + background-color: #2ec196; + border-color: #2ec196; + color: #ffffff; + + &:hover:not(:disabled) { + background-color: #26a880; + border-color: #26a880; + color: #ffffff; + } + + &:focus:not(:disabled) { + background-color: #2ec196; + border-color: #2ec196; + color: #ffffff; + } + + &:disabled { + background-color: #f2f2f2; + border-color: #f2f2f2; + color: #b5b5b5; + } + } +`; Review Comment: Would it be possible to use theme color tokens here instead of hardcoded hex values? ########## superset-frontend/src/components/StreamingExportModal/StreamingExportModal.tsx: ########## @@ -0,0 +1,403 @@ +/** + * 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. + */ +/* eslint-disable theme-colors/no-literal-colors */ +import { styled, t } from '@superset-ui/core'; +import { Modal, Button, Typography, Progress } from 'antd'; +import { Icons } from '@superset-ui/core/components/Icons'; + +const { Text } = Typography; + +export enum ExportStatus { + STREAMING = 'streaming', + COMPLETED = 'completed', + ERROR = 'error', + CANCELLED = 'cancelled', +} + +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; + progress: StreamingProgress; +} + +const ModalContent = styled.div` + padding: ${({ theme }) => theme.sizeUnit * 4}px 0 + ${({ theme }) => theme.sizeUnit * 2}px; +`; + +const ProgressSection = styled.div` + margin: ${({ theme }) => theme.sizeUnit * 6}px 0; + position: relative; +`; + +const ProgressWrapper = styled.div` + display: flex; + align-items: center; + gap: ${({ theme }) => theme.sizeUnit * 3}px; +`; + +const SuccessIcon = styled(Icons.CheckCircleFilled)` + color: #52c41a; + font-size: 24px; + flex-shrink: 0; +`; + +const ErrorIconWrapper = styled.div` + display: flex; + align-items: center; + justify-content: center; + width: 16px; + height: 16px; + background-color: #ff4d4f; + border-radius: 50%; + flex-shrink: 0; +`; + +const ErrorIconStyled = styled(Icons.CloseOutlined)` + color: white; + font-size: 10px; +`; + +const ActionButtons = styled.div` + display: flex; + gap: ${({ theme }) => theme.sizeUnit * 2}px; + justify-content: flex-end; +`; + +const ProgressText = styled(Text)` + display: block; + text-align: center; + margin-top: ${({ theme }) => theme.sizeUnit * 4}px; +`; + +const ErrorText = styled(Text)` + display: block; + text-align: center; + margin-top: ${({ theme }) => theme.sizeUnit * 4}px; +`; + +const CancelButton = styled(Button)` + background-color: #f0fff8; + color: #1c997a; + border-color: #f0fff8; + + &:hover { + background-color: #f0fff8; + color: #1c997a; + border-color: #1c997a; + } + + &:focus { + background-color: #f0fff8; + color: #1c997a; + border-color: #1c997a; + } +`; + +const DownloadButton = styled(Button)` + &.ant-btn-primary { + background-color: #2ec196; + border-color: #2ec196; + color: #ffffff; + + &:hover:not(:disabled) { + background-color: #26a880; + border-color: #26a880; + color: #ffffff; + } + + &:focus:not(:disabled) { + background-color: #2ec196; + border-color: #2ec196; + color: #ffffff; + } + + &:disabled { + background-color: #f2f2f2; + border-color: #f2f2f2; + color: #b5b5b5; + } + } +`; + +interface ModalStateContentProps { + status: ExportStatus; + progress: StreamingProgress; + onCancel: () => void; + onRetry?: () => void; + onDownload: () => void; + getProgressPercentage: () => number; +} + +const ErrorContent = ({ + error, + onCancel, + onRetry, + getProgressPercentage, +}: { + error?: string; + onCancel: () => void; + onRetry?: () => void; + getProgressPercentage: () => number; +}) => ( + <ModalContent> + <ProgressSection> + <ProgressWrapper> + <Progress + percent={getProgressPercentage()} + status="exception" + showInfo={false} + style={{ flex: 1 }} + /> + <ErrorIconWrapper> + <ErrorIconStyled /> + </ErrorIconWrapper> + </ProgressWrapper> + <ErrorText type="danger">{error || t('Export failed')}</ErrorText> + </ProgressSection> + <ActionButtons> + <CancelButton onClick={onCancel}>{t('Close')}</CancelButton> + {onRetry && ( + <DownloadButton type="primary" onClick={onRetry}> + {t('Retry')} + </DownloadButton> + )} + </ActionButtons> + </ModalContent> +); + +const CancelledContent = ({ + getProgressPercentage, + onCancel, + onRetry, +}: { + getProgressPercentage: () => number; + onCancel: () => void; + onRetry?: () => void; +}) => ( + <ModalContent> + <ProgressSection> + <Progress + percent={getProgressPercentage()} + status="exception" + showInfo={false} + /> + <ProgressText>{t('Export cancelled')}</ProgressText> + </ProgressSection> + <ActionButtons> + <CancelButton onClick={onCancel}>{t('Close')}</CancelButton> + {onRetry && ( + <DownloadButton type="primary" onClick={onRetry}> + {t('Retry')} + </DownloadButton> + )} + </ActionButtons> + </ModalContent> +); + +const CompletedContent = ({ + filename, + downloadUrl, + onCancel, + onDownload, +}: { + filename?: string; + downloadUrl?: string; + onCancel: () => void; + onDownload: () => void; +}) => ( + <ModalContent> + <ProgressSection> + <ProgressWrapper> + <Progress + percent={100} + status="success" + showInfo={false} + style={{ flex: 1 }} Review Comment: I think here as well ########## superset/commands/sql_lab/streaming_export_command.py: ########## @@ -0,0 +1,239 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""Command for streaming CSV exports of SQL Lab query results.""" + +from __future__ import annotations + +import csv +import io +import logging +import time +from typing import Callable, Generator, TYPE_CHECKING + +from flask import current_app as app +from flask_babel import gettext as __ + +from superset import db +from superset.commands.base import BaseCommand +from superset.errors import ErrorLevel, SupersetError, SupersetErrorType +from superset.exceptions import SupersetErrorException, SupersetSecurityException +from superset.models.sql_lab import Query +from superset.sql.parse import SQLScript +from superset.sqllab.limiting_factor import LimitingFactor + +if TYPE_CHECKING: + pass + +logger = logging.getLogger(__name__) + + +class StreamingSqlResultExportCommand(BaseCommand): + """ + Command to execute a streaming CSV export of SQL Lab query results. + + This command handles the business logic for: + - Fetching SQL Lab query results + - Generating CSV data in chunks + - Managing database connections + - Buffering data for efficient streaming + """ + + def __init__( + self, + client_id: str, + chunk_size: int = 1000, + ): + """ + Initialize the streaming export command. + + Args: + client_id: The SQL Lab query client ID + chunk_size: Number of rows to fetch per database query (default: 1000) + """ + self._client_id = client_id + self._chunk_size = chunk_size + self._current_app = app._get_current_object() + self._query: Query | None = None + + def validate(self) -> None: + """Validate permissions and query existence.""" + self._query = ( + db.session.query(Query).filter_by(client_id=self._client_id).one_or_none() + ) + if self._query is None: + raise SupersetErrorException( + SupersetError( + message=__( + "The query associated with these results could not be found. " + "You need to re-run the original query." + ), + error_type=SupersetErrorType.RESULTS_BACKEND_ERROR, + level=ErrorLevel.ERROR, + ), + status=404, + ) + + try: + self._query.raise_for_access() + except SupersetSecurityException as ex: + raise SupersetErrorException( + SupersetError( + message=__("Cannot access the query"), + error_type=SupersetErrorType.QUERY_SECURITY_ACCESS_ERROR, + level=ErrorLevel.ERROR, + ), + status=403, + ) from ex + + def run(self) -> Callable[[], Generator[str, None, None]]: # noqa: C901 + """ + Execute the streaming CSV export. + + Returns: + A callable that returns a generator yielding CSV data chunks as strings. + The callable is needed to maintain Flask app context during streaming. + """ + # Load all Query attributes while session is still active + # to avoid DetachedInstanceError + assert self._query is not None + + select_sql = self._query.select_sql + executed_sql = self._query.executed_sql + limiting_factor = self._query.limiting_factor + database = self._query.database + + # Get the SQL and limit + if select_sql: + sql = select_sql + limit = None + else: + sql = executed_sql + script = SQLScript(sql, database.db_engine_spec.engine) + # when a query has multiple statements only the last one returns data + limit = script.statements[-1].get_limit_value() + + if limit is not None and limiting_factor in { + LimitingFactor.QUERY, + LimitingFactor.DROPDOWN, + LimitingFactor.QUERY_AND_DROPDOWN, + }: + # remove extra row from `increased_limit` + limit -= 1 + + def csv_generator() -> Generator[str, None, None]: # noqa: C901 + """Generator that yields CSV data from SQL Lab query results.""" + with self._current_app.app_context(): + start_time = time.time() + total_bytes = 0 + + try: + # Create a new session to keep database object attached + with db.session() as session: + # Merge database to prevent DetachedInstanceError + merged_database = session.merge(database) + + # Execute query with streaming + with merged_database.get_sqla_engine() as engine: + connection = engine.connect() + + try: + from sqlalchemy import text + + result_proxy = connection.execution_options( + stream_results=True + ).execute(text(sql)) + + columns = list(result_proxy.keys()) + + # Use StringIO with csv.writer for proper escaping + buffer = io.StringIO() + csv_writer = csv.writer( + buffer, quoting=csv.QUOTE_MINIMAL + ) + + # Write CSV header + csv_writer.writerow(columns) + header_data = buffer.getvalue() + total_bytes += len(header_data.encode("utf-8")) + yield header_data + buffer.seek(0) + buffer.truncate() + + row_count = 0 + flush_threshold = 65536 # 64KB + + while True: + rows = result_proxy.fetchmany(self._chunk_size) + if not rows: + break + + for row in rows: + # Apply limit if specified + if limit is not None and row_count >= limit: + break + + csv_writer.writerow(row) + row_count += 1 + + # Check buffer size and flush if needed + current_size = buffer.tell() + if current_size >= flush_threshold: + data = buffer.getvalue() + data_bytes = len(data.encode("utf-8")) + total_bytes += data_bytes + yield data + buffer.seek(0) + buffer.truncate() + + # Break outer loop if limit reached + if limit is not None and row_count >= limit: + break + + # Flush remaining buffer + remaining_data = buffer.getvalue() + if remaining_data: + total_bytes += len(remaining_data.encode("utf-8")) + yield remaining_data + + # Log completion + total_time = time.time() - start_time + total_mb = total_bytes / (1024 * 1024) + logger.info( + "SQL Lab streaming CSV completed: %s rows, " + "%.1fMB in %.2fs", + f"{row_count:,}", + total_mb, + total_time, + ) + + finally: + connection.close() + + except Exception as e: + logger.error("Error in SQL Lab streaming CSV generator: %s", e) + import traceback + + logger.error("Traceback: %s", traceback.format_exc()) + + # Send error marker for frontend to detect + error_marker = ( + "__STREAM_ERROR__:Export failed. " + "Please try again in some time.\n" Review Comment: I saw Kasia say it should be "Export failed.Please try again". Should it apply here as well? ########## superset-frontend/src/components/StreamingExportModal/StreamingExportModal.tsx: ########## @@ -0,0 +1,403 @@ +/** + * 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. + */ +/* eslint-disable theme-colors/no-literal-colors */ +import { styled, t } from '@superset-ui/core'; +import { Modal, Button, Typography, Progress } from 'antd'; +import { Icons } from '@superset-ui/core/components/Icons'; + +const { Text } = Typography; + +export enum ExportStatus { + STREAMING = 'streaming', + COMPLETED = 'completed', + ERROR = 'error', + CANCELLED = 'cancelled', +} + +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; + progress: StreamingProgress; +} + +const ModalContent = styled.div` + padding: ${({ theme }) => theme.sizeUnit * 4}px 0 + ${({ theme }) => theme.sizeUnit * 2}px; +`; + +const ProgressSection = styled.div` + margin: ${({ theme }) => theme.sizeUnit * 6}px 0; + position: relative; +`; + +const ProgressWrapper = styled.div` + display: flex; + align-items: center; + gap: ${({ theme }) => theme.sizeUnit * 3}px; +`; + +const SuccessIcon = styled(Icons.CheckCircleFilled)` + color: #52c41a; + font-size: 24px; + flex-shrink: 0; +`; + +const ErrorIconWrapper = styled.div` + display: flex; + align-items: center; + justify-content: center; + width: 16px; + height: 16px; + background-color: #ff4d4f; + border-radius: 50%; + flex-shrink: 0; +`; + +const ErrorIconStyled = styled(Icons.CloseOutlined)` + color: white; + font-size: 10px; +`; + +const ActionButtons = styled.div` + display: flex; + gap: ${({ theme }) => theme.sizeUnit * 2}px; + justify-content: flex-end; +`; + +const ProgressText = styled(Text)` + display: block; + text-align: center; + margin-top: ${({ theme }) => theme.sizeUnit * 4}px; +`; + +const ErrorText = styled(Text)` + display: block; + text-align: center; + margin-top: ${({ theme }) => theme.sizeUnit * 4}px; +`; + +const CancelButton = styled(Button)` + background-color: #f0fff8; + color: #1c997a; + border-color: #f0fff8; + + &:hover { + background-color: #f0fff8; + color: #1c997a; + border-color: #1c997a; + } + + &:focus { + background-color: #f0fff8; + color: #1c997a; + border-color: #1c997a; + } +`; + +const DownloadButton = styled(Button)` + &.ant-btn-primary { + background-color: #2ec196; + border-color: #2ec196; + color: #ffffff; + + &:hover:not(:disabled) { + background-color: #26a880; + border-color: #26a880; + color: #ffffff; + } + + &:focus:not(:disabled) { + background-color: #2ec196; + border-color: #2ec196; + color: #ffffff; + } + + &:disabled { + background-color: #f2f2f2; + border-color: #f2f2f2; + color: #b5b5b5; + } + } +`; + +interface ModalStateContentProps { + status: ExportStatus; + progress: StreamingProgress; + onCancel: () => void; + onRetry?: () => void; + onDownload: () => void; + getProgressPercentage: () => number; +} + +const ErrorContent = ({ + error, + onCancel, + onRetry, + getProgressPercentage, +}: { + error?: string; + onCancel: () => void; + onRetry?: () => void; + getProgressPercentage: () => number; +}) => ( + <ModalContent> + <ProgressSection> + <ProgressWrapper> + <Progress + percent={getProgressPercentage()} + status="exception" + showInfo={false} + style={{ flex: 1 }} Review Comment: Consider replacing inline style={{ flex: 1 }} with a styled wrapper for consistency. ########## superset-frontend/src/components/StreamingExportModal/useStreamingExport.ts: ########## @@ -0,0 +1,366 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +import { useState, useCallback, useRef, useEffect } from 'react'; +import { SupersetClient } from '@superset-ui/core'; +import { ExportStatus, StreamingProgress } from './StreamingExportModal'; + +interface UseStreamingExportOptions { + onComplete?: (downloadUrl: string, filename: string) => void; + onError?: (error: string) => void; +} + +interface StreamingExportPayload { + [key: string]: unknown; +} + +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, + _exportType: string, + expectedRows: number | undefined, + signal: AbortSignal, +): Promise<RequestInit> => { + const headers: Record<string, string> = { + 'Content-Type': 'application/x-www-form-urlencoded', + }; + + // Get CSRF token using SupersetClient + const csrfToken = await SupersetClient.getCSRFToken(); + if (csrfToken) { + headers['X-CSRFToken'] = csrfToken; + } + + // Build form data - if payload has client_id, it's SQL Lab export + // Otherwise it's a chart export with form_data + const formParams: Record<string, string> = { + filename, + 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<StreamingProgress>({ + rowsProcessed: 0, + totalRows: undefined, + totalSize: 0, + speed: 0, + mbPerSecond: 0, + elapsedTime: 0, + status: ExportStatus.STREAMING, + }); + const [retryCount, setRetryCount] = useState(0); + const abortControllerRef = useRef<AbortController | null>(null); + const lastExportParamsRef = useRef<StreamingExportParams | null>(null); + const currentBlobUrlRef = useRef<string | null>(null); + const isExportingRef = useRef(false); + + const updateProgress = useCallback((updates: Partial<StreamingProgress>) => { + setProgress(prev => ({ ...prev, ...updates })); + }, []); + + const executeExport = useCallback( + async (params: StreamingExportParams) => { + const { url, payload, filename, exportType, expectedRows } = params; + + abortControllerRef.current = new AbortController(); + + updateProgress({ + rowsProcessed: 0, + totalRows: expectedRows, + totalSize: 0, + speed: 0, + mbPerSecond: 0, + elapsedTime: 0, + status: ExportStatus.STREAMING, + filename, + }); + + try { + const defaultFilename = `export.${exportType}`; + const finalFilename = filename || defaultFilename; + + const fetchOptions = await createFetchRequest( + url, + payload, + finalFilename, + exportType, + expectedRows, + abortControllerRef.current.signal, + ); + + const response = await fetch(url, fetchOptions); + + if (!response.ok) { + throw new Error( + `Export failed: ${response.status} ${response.statusText}`, + ); + } + + if (!response.body) { + throw new Error('Response body is not available for streaming'); + } + + 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, + }); + } + + // 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: finalFilename, + }); + + isExportingRef.current = false; + options.onComplete?.(downloadUrl, finalFilename); + } 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; Review Comment: This repeats itself a number of times "isExportingRef.current = false;" Maybe we could wrap it in a helper function then replace the different calls? -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: notifications-unsubscr...@superset.apache.org For queries about this service, please contact Infrastructure at: us...@infra.apache.org --------------------------------------------------------------------- To unsubscribe, e-mail: notifications-unsubscr...@superset.apache.org For additional commands, e-mail: notifications-h...@superset.apache.org
