mirror of
https://github.com/apache/superset.git
synced 2026-08-01 19:42:29 +00:00
Compare commits
4 Commits
fix-screen
...
fix-sql-la
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
7d5e8ad785 | ||
|
|
22c305f758 | ||
|
|
6929d032b8 | ||
|
|
f6c574edd8 |
@@ -62,6 +62,7 @@ export function ErrorMessageWithStackTrace({
|
||||
fallback,
|
||||
compact,
|
||||
closable = true,
|
||||
errorMitigationFunction,
|
||||
}: Props) {
|
||||
// Check if a custom error message component was registered for this message
|
||||
if (error) {
|
||||
@@ -77,6 +78,7 @@ export function ErrorMessageWithStackTrace({
|
||||
error={error}
|
||||
source={source}
|
||||
subtitle={subtitle}
|
||||
errorMitigationFunction={errorMitigationFunction}
|
||||
/>
|
||||
);
|
||||
}
|
||||
|
||||
@@ -20,7 +20,7 @@
|
||||
import * as reduxHooks from 'react-redux';
|
||||
import { Provider } from 'react-redux';
|
||||
import { createStore, Store } from 'redux';
|
||||
import { render, waitFor } from 'spec/helpers/testing-library';
|
||||
import { act, render, waitFor } from 'spec/helpers/testing-library';
|
||||
import { ErrorLevel, ErrorSource, ErrorTypeEnum } from '@superset-ui/core';
|
||||
import { reRunQuery } from 'src/SqlLab/actions/sqlLab';
|
||||
import { triggerQuery } from 'src/components/Chart/chartAction';
|
||||
@@ -166,15 +166,55 @@ describe('OAuth2RedirectMessage Component', () => {
|
||||
render(setup());
|
||||
|
||||
simulateBroadcastMessage({ tabId: 'tabId' });
|
||||
simulateStorageMessage({ tabId: 'tabId' });
|
||||
|
||||
await waitFor(() => {
|
||||
expect(reRunQuery).toHaveBeenCalledWith({ sql: 'SELECT * FROM table' });
|
||||
});
|
||||
expect(reRunQuery).toHaveBeenCalledTimes(1);
|
||||
});
|
||||
|
||||
test('dispatches reRunQuery action when storage event has matching tab ID', async () => {
|
||||
render(setup());
|
||||
|
||||
simulateStorageMessage({ tabId: 'tabId' });
|
||||
|
||||
await waitFor(() => {
|
||||
expect(reRunQuery).toHaveBeenCalledWith({ sql: 'SELECT * FROM table' });
|
||||
});
|
||||
});
|
||||
|
||||
test('dispatches reRunQuery action when storage event has matching tab ID', async () => {
|
||||
render(setup());
|
||||
test('waits for the SQL Lab query before consuming the completion', async () => {
|
||||
const initialState = {
|
||||
sqlLab: {
|
||||
queries: {},
|
||||
queryEditors: [{ id: 'editor-id', latestQueryId: 'query-id' }],
|
||||
tabHistory: ['editor-id'],
|
||||
},
|
||||
explore: { slice: null },
|
||||
charts: {},
|
||||
dashboardInfo: {},
|
||||
};
|
||||
const delayedQueryStore = createStore(
|
||||
(state: typeof initialState = initialState, action) =>
|
||||
action.type === 'load-query'
|
||||
? {
|
||||
...state,
|
||||
sqlLab: {
|
||||
...state.sqlLab,
|
||||
queries: { 'query-id': { sql: 'SELECT * FROM table' } },
|
||||
},
|
||||
}
|
||||
: state,
|
||||
);
|
||||
render(setup({}, delayedQueryStore));
|
||||
|
||||
simulateBroadcastMessage({ tabId: 'tabId' });
|
||||
expect(reRunQuery).not.toHaveBeenCalled();
|
||||
|
||||
act(() => {
|
||||
delayedQueryStore.dispatch({ type: 'load-query' });
|
||||
});
|
||||
simulateStorageMessage({ tabId: 'tabId' });
|
||||
|
||||
await waitFor(() => {
|
||||
@@ -234,4 +274,22 @@ describe('OAuth2RedirectMessage Component', () => {
|
||||
]);
|
||||
});
|
||||
});
|
||||
|
||||
test('runs scoped mitigation once instead of CRUD invalidation', async () => {
|
||||
const errorMitigationFunction = jest.fn();
|
||||
render(
|
||||
setup({
|
||||
source: 'crud' as ErrorSource,
|
||||
errorMitigationFunction,
|
||||
}),
|
||||
);
|
||||
|
||||
simulateBroadcastMessage({ tabId: 'tabId' });
|
||||
simulateStorageMessage({ tabId: 'tabId' });
|
||||
|
||||
await waitFor(() => {
|
||||
expect(errorMitigationFunction).toHaveBeenCalledTimes(1);
|
||||
});
|
||||
expect(api.util.invalidateTags).not.toHaveBeenCalled();
|
||||
});
|
||||
});
|
||||
|
||||
@@ -16,7 +16,7 @@
|
||||
* specific language governing permissions and limitations
|
||||
* under the License.
|
||||
*/
|
||||
import { useEffect } from 'react';
|
||||
import { useEffect, useRef } from 'react';
|
||||
|
||||
import { useDispatch, useSelector } from 'react-redux';
|
||||
import { QueryEditor, SqlLabRootState } from 'src/SqlLab/types';
|
||||
@@ -58,15 +58,16 @@ interface OAuth2RedirectExtra {
|
||||
*
|
||||
* After the token has been stored, the opened tab will broadcast a message to the
|
||||
* original tab and close itself. This component, running on the original tab, listens
|
||||
* on a same-origin BroadcastChannel and re-runs the query for the user once it
|
||||
* receives the success message — be it in SQL Lab, Explore, or a dashboard. Both tabs
|
||||
* share a "tab ID" (a UUID generated by the backend) which is echoed back through the
|
||||
* channel so the original tab only reacts to its own OAuth2 flow.
|
||||
* for same-origin BroadcastChannel and storage notifications and re-runs the query
|
||||
* for the user once it receives the success message — be it in SQL Lab, Explore, or
|
||||
* a dashboard. Both tabs share a "tab ID" (a UUID generated by the backend) which is
|
||||
* echoed back so the original tab only reacts to its own OAuth2 flow.
|
||||
*/
|
||||
export function OAuth2RedirectMessage({
|
||||
error,
|
||||
source,
|
||||
closable,
|
||||
errorMitigationFunction,
|
||||
}: ErrorMessageComponentProps<OAuth2RedirectExtra>) {
|
||||
const { extra, level } = error;
|
||||
|
||||
@@ -103,13 +104,17 @@ export function OAuth2RedirectMessage({
|
||||
);
|
||||
|
||||
const dispatch = useDispatch();
|
||||
const lastHandledTabIdRef = useRef<string>();
|
||||
|
||||
useEffect(() => {
|
||||
const handleOAuthComplete = (tabId?: string) => {
|
||||
if (tabId !== extra.tab_id) {
|
||||
if (tabId !== extra.tab_id || tabId === lastHandledTabIdRef.current) {
|
||||
return;
|
||||
}
|
||||
if (source === 'sqllab' && query) {
|
||||
|
||||
if (errorMitigationFunction) {
|
||||
errorMitigationFunction();
|
||||
} else if (source === 'sqllab' && query) {
|
||||
dispatch(reRunQuery(query));
|
||||
} else if (source === 'explore') {
|
||||
dispatch(triggerQuery(true, chartId));
|
||||
@@ -123,7 +128,11 @@ export function OAuth2RedirectMessage({
|
||||
'Tables',
|
||||
]),
|
||||
);
|
||||
} else {
|
||||
return;
|
||||
}
|
||||
|
||||
lastHandledTabIdRef.current = tabId;
|
||||
};
|
||||
|
||||
const channel =
|
||||
@@ -156,7 +165,16 @@ export function OAuth2RedirectMessage({
|
||||
window.removeEventListener('storage', handleStorage);
|
||||
channel?.close();
|
||||
};
|
||||
}, [source, extra.tab_id, dispatch, query, chartId, chartList, dashboardId]);
|
||||
}, [
|
||||
source,
|
||||
extra.tab_id,
|
||||
dispatch,
|
||||
query,
|
||||
chartId,
|
||||
chartList,
|
||||
dashboardId,
|
||||
errorMitigationFunction,
|
||||
]);
|
||||
|
||||
const body = (
|
||||
<p>
|
||||
|
||||
@@ -27,6 +27,7 @@ export type ErrorMessageComponentProps<ExtraType = Record<string, any> | null> =
|
||||
subtitle?: ReactNode;
|
||||
compact?: boolean;
|
||||
closable?: boolean;
|
||||
errorMitigationFunction?: () => void;
|
||||
};
|
||||
|
||||
export type ErrorMessageComponent = ComponentType<ErrorMessageComponentProps>;
|
||||
|
||||
@@ -35,6 +35,5 @@ export const Basic: StoryFn<typeof DatasetPanel> = args => (
|
||||
Basic.args = {
|
||||
tableName: 'example_table',
|
||||
loading: false,
|
||||
hasError: false,
|
||||
columnList: exampleColumns,
|
||||
};
|
||||
|
||||
@@ -77,7 +77,6 @@ test('View Dataset opens a single-prefixed URL under a subdirectory deployment',
|
||||
render(
|
||||
<DatasetPanel
|
||||
tableName="example_table"
|
||||
hasError={false}
|
||||
columnList={exampleColumns}
|
||||
loading={false}
|
||||
datasets={datasetWith(`${APP_ROOT}/explore/?datasource=1__table`)}
|
||||
@@ -97,7 +96,6 @@ test('View Dataset passes an external explore_url through unprefixed', async ()
|
||||
render(
|
||||
<DatasetPanel
|
||||
tableName="example_table"
|
||||
hasError={false}
|
||||
columnList={exampleColumns}
|
||||
loading={false}
|
||||
datasets={datasetWith('https://external.example.com/custom-endpoint')}
|
||||
|
||||
@@ -17,10 +17,12 @@
|
||||
* under the License.
|
||||
*/
|
||||
import { render, screen } from 'spec/helpers/testing-library';
|
||||
import { ErrorTypeEnum } from '@superset-ui/core';
|
||||
import DatasetPanel, {
|
||||
REFRESHING,
|
||||
tableColumnDefinition,
|
||||
COLUMN_TITLE,
|
||||
ERROR_TITLE,
|
||||
} from 'src/features/datasets/AddDataset/DatasetPanel/DatasetPanel';
|
||||
import { exampleColumns, exampleDataset } from './fixtures';
|
||||
import { ITableColumn } from './types';
|
||||
@@ -31,8 +33,6 @@ import {
|
||||
SELECT_TABLE_TITLE,
|
||||
NO_COLUMNS_TITLE,
|
||||
NO_COLUMNS_DESCRIPTION,
|
||||
ERROR_TITLE,
|
||||
ERROR_DESCRIPTION,
|
||||
} from './MessageContent';
|
||||
|
||||
jest.mock(
|
||||
@@ -47,7 +47,7 @@ jest.mock(
|
||||
// eslint-disable-next-line no-restricted-globals -- TODO: Migrate from describe blocks
|
||||
describe('DatasetPanel', () => {
|
||||
test('renders a blank state DatasetPanel', () => {
|
||||
render(<DatasetPanel hasError={false} columnList={[]} loading={false} />, {
|
||||
render(<DatasetPanel columnList={[]} loading={false} />, {
|
||||
useRouter: true,
|
||||
});
|
||||
|
||||
@@ -70,17 +70,9 @@ describe('DatasetPanel', () => {
|
||||
});
|
||||
|
||||
test('renders a no columns screen', () => {
|
||||
render(
|
||||
<DatasetPanel
|
||||
tableName="Name"
|
||||
hasError={false}
|
||||
columnList={[]}
|
||||
loading={false}
|
||||
/>,
|
||||
{
|
||||
useRouter: true,
|
||||
},
|
||||
);
|
||||
render(<DatasetPanel tableName="Name" columnList={[]} loading={false} />, {
|
||||
useRouter: true,
|
||||
});
|
||||
|
||||
const blankDatasetImg = screen.getByRole('img', { name: /empty/i });
|
||||
expect(blankDatasetImg).toBeVisible();
|
||||
@@ -91,17 +83,9 @@ describe('DatasetPanel', () => {
|
||||
});
|
||||
|
||||
test('renders a loading screen', () => {
|
||||
render(
|
||||
<DatasetPanel
|
||||
tableName="Name"
|
||||
hasError={false}
|
||||
columnList={[]}
|
||||
loading
|
||||
/>,
|
||||
{
|
||||
useRouter: true,
|
||||
},
|
||||
);
|
||||
render(<DatasetPanel tableName="Name" columnList={[]} loading />, {
|
||||
useRouter: true,
|
||||
});
|
||||
|
||||
const loadingIndicator = screen.getByTestId('loading-indicator');
|
||||
expect(loadingIndicator).toBeVisible();
|
||||
@@ -113,7 +97,12 @@ describe('DatasetPanel', () => {
|
||||
render(
|
||||
<DatasetPanel
|
||||
tableName="Name"
|
||||
hasError
|
||||
error={{
|
||||
error_type: ErrorTypeEnum.GENERIC_BACKEND_ERROR,
|
||||
extra: null,
|
||||
level: 'error',
|
||||
message: 'Structured backend failure',
|
||||
}}
|
||||
columnList={[]}
|
||||
loading={false}
|
||||
/>,
|
||||
@@ -124,8 +113,9 @@ describe('DatasetPanel', () => {
|
||||
|
||||
const errorTitle = screen.getByText(ERROR_TITLE);
|
||||
expect(errorTitle).toBeVisible();
|
||||
const errorDescription = screen.getByText(ERROR_DESCRIPTION);
|
||||
const errorDescription = screen.getByText('Structured backend failure');
|
||||
expect(errorDescription).toBeVisible();
|
||||
expect(screen.getByTitle('Name')).toHaveStyle({ position: 'relative' });
|
||||
});
|
||||
|
||||
test('renders a table with columns displayed', async () => {
|
||||
@@ -133,7 +123,6 @@ describe('DatasetPanel', () => {
|
||||
render(
|
||||
<DatasetPanel
|
||||
tableName={tableName}
|
||||
hasError={false}
|
||||
columnList={exampleColumns}
|
||||
loading={false}
|
||||
/>,
|
||||
@@ -159,7 +148,6 @@ describe('DatasetPanel', () => {
|
||||
render(
|
||||
<DatasetPanel
|
||||
tableName="example_table"
|
||||
hasError={false}
|
||||
columnList={exampleColumns}
|
||||
loading={false}
|
||||
datasets={exampleDataset}
|
||||
|
||||
@@ -21,11 +21,13 @@ import { Alert } from '@apache-superset/core/components';
|
||||
import { css, styled } from '@apache-superset/core/theme';
|
||||
import { Icons } from '@superset-ui/core/components/Icons';
|
||||
import { Loading } from '@superset-ui/core/components';
|
||||
import type { SupersetError } from '@superset-ui/core';
|
||||
import Table, {
|
||||
ColumnsType,
|
||||
TableSize,
|
||||
} from '@superset-ui/core/components/Table';
|
||||
import { DatasetObject } from 'src/features/datasets/AddDataset/types';
|
||||
import { ErrorMessageWithStackTrace } from 'src/components';
|
||||
import { openInNewTab, stripAppRoot } from 'src/utils/navigationUtils';
|
||||
import { ITableColumn } from './types';
|
||||
import MessageContent from './MessageContent';
|
||||
@@ -146,6 +148,10 @@ const TableScrollContainer = styled.div`
|
||||
right: 0;
|
||||
`;
|
||||
|
||||
const ErrorContainer = styled.div`
|
||||
padding: 0 ${({ theme }) => theme.sizeUnit * 6}px;
|
||||
`;
|
||||
|
||||
const StyledAlert = styled(Alert)`
|
||||
${({ theme }) => `
|
||||
border: 1px solid ${theme.colorInfoText};
|
||||
@@ -167,6 +173,7 @@ const StyledAlert = styled(Alert)`
|
||||
|
||||
export const REFRESHING = t('Refreshing columns');
|
||||
export const COLUMN_TITLE = t('Table columns');
|
||||
export const ERROR_TITLE = t('An Error Occurred');
|
||||
|
||||
const pageSizeOptions = ['5', '10', '15', '25'];
|
||||
const DEFAULT_PAGE_SIZE = 25;
|
||||
@@ -201,9 +208,13 @@ export interface IDatasetPanelProps {
|
||||
*/
|
||||
columnList: ITableColumn[];
|
||||
/**
|
||||
* Boolean indicating if there is an error state
|
||||
* Error returned while loading the table metadata
|
||||
*/
|
||||
hasError: boolean;
|
||||
error?: SupersetError;
|
||||
/**
|
||||
* Function used to retry loading the table metadata after error mitigation
|
||||
*/
|
||||
errorMitigationFunction?: () => void;
|
||||
/**
|
||||
* Boolean indicating if the component is in a loading state
|
||||
*/
|
||||
@@ -256,11 +267,11 @@ const DatasetPanel = ({
|
||||
tableName,
|
||||
columnList,
|
||||
loading,
|
||||
hasError,
|
||||
error,
|
||||
errorMitigationFunction,
|
||||
datasets,
|
||||
}: IDatasetPanelProps) => {
|
||||
const hasColumns = Boolean(columnList?.length > 0);
|
||||
const datasetNames = datasets?.map(dataset => dataset.table_name);
|
||||
const hasColumns = columnList.length > 0;
|
||||
const tableWithDataset = datasets?.find(
|
||||
dataset => dataset.table_name === tableName,
|
||||
);
|
||||
@@ -278,7 +289,19 @@ const DatasetPanel = ({
|
||||
);
|
||||
}
|
||||
if (!loading) {
|
||||
if (!loading && tableName && hasColumns && !hasError) {
|
||||
if (error) {
|
||||
component = (
|
||||
<ErrorContainer>
|
||||
<ErrorMessageWithStackTrace
|
||||
error={error}
|
||||
errorMitigationFunction={errorMitigationFunction}
|
||||
source="crud"
|
||||
subtitle={error.message}
|
||||
title={ERROR_TITLE}
|
||||
/>
|
||||
</ErrorContainer>
|
||||
);
|
||||
} else if (tableName && hasColumns) {
|
||||
component = (
|
||||
<>
|
||||
<StyledTitle title={COLUMN_TITLE}>{COLUMN_TITLE}</StyledTitle>
|
||||
@@ -312,13 +335,7 @@ const DatasetPanel = ({
|
||||
</>
|
||||
);
|
||||
} else {
|
||||
component = (
|
||||
<MessageContent
|
||||
hasColumns={hasColumns}
|
||||
hasError={hasError}
|
||||
tableName={tableName}
|
||||
/>
|
||||
);
|
||||
component = <MessageContent tableName={tableName} />;
|
||||
}
|
||||
}
|
||||
|
||||
@@ -326,11 +343,12 @@ const DatasetPanel = ({
|
||||
<>
|
||||
{tableName && (
|
||||
<>
|
||||
{datasetNames?.includes(tableName) &&
|
||||
renderExistingDatasetAlert(tableWithDataset)}
|
||||
{tableWithDataset && renderExistingDatasetAlert(tableWithDataset)}
|
||||
<StyledHeader
|
||||
position={
|
||||
!loading && hasColumns ? EPosition.RELATIVE : EPosition.ABSOLUTE
|
||||
!loading && (hasColumns || error)
|
||||
? EPosition.RELATIVE
|
||||
: EPosition.ABSOLUTE
|
||||
}
|
||||
title={tableName || ''}
|
||||
>
|
||||
|
||||
@@ -16,8 +16,14 @@
|
||||
* specific language governing permissions and limitations
|
||||
* under the License.
|
||||
*/
|
||||
import { render, waitFor } from 'spec/helpers/testing-library';
|
||||
import { SupersetClient } from '@superset-ui/core';
|
||||
import { act, render, screen, waitFor } from 'spec/helpers/testing-library';
|
||||
import { ErrorTypeEnum, SupersetClient } from '@superset-ui/core';
|
||||
import type { SupersetClientResponse } from '@superset-ui/core';
|
||||
import {
|
||||
DatabaseErrorMessage,
|
||||
getErrorMessageComponentRegistry,
|
||||
OAuth2RedirectMessage,
|
||||
} from 'src/components/ErrorMessage';
|
||||
import DatasetPanelWrapper from 'src/features/datasets/AddDataset/DatasetPanel';
|
||||
|
||||
jest.mock(
|
||||
@@ -29,17 +35,29 @@ jest.mock(
|
||||
),
|
||||
);
|
||||
|
||||
const errorMessageRegistry = getErrorMessageComponentRegistry();
|
||||
|
||||
afterEach(() => {
|
||||
errorMessageRegistry.remove(ErrorTypeEnum.GENERIC_BACKEND_ERROR);
|
||||
errorMessageRegistry.remove(ErrorTypeEnum.OAUTH2_REDIRECT);
|
||||
jest.restoreAllMocks();
|
||||
});
|
||||
|
||||
const tableMetadataResponse = (
|
||||
name: string,
|
||||
columnName: string,
|
||||
): SupersetClientResponse => ({
|
||||
response: new Response(),
|
||||
json: {
|
||||
name,
|
||||
columns: [{ name: columnName, type: 'INTEGER', longType: 'INTEGER' }],
|
||||
},
|
||||
});
|
||||
|
||||
test('fetches table metadata for schema-less database without schema', async () => {
|
||||
const getSpy = jest.spyOn(SupersetClient, 'get').mockResolvedValue({
|
||||
json: {
|
||||
name: 'my_table',
|
||||
columns: [{ name: 'id', type: 'INTEGER', longType: 'INTEGER' }],
|
||||
},
|
||||
} as any);
|
||||
const getSpy = jest
|
||||
.spyOn(SupersetClient, 'get')
|
||||
.mockResolvedValue(tableMetadataResponse('my_table', 'id'));
|
||||
|
||||
render(
|
||||
<DatasetPanelWrapper
|
||||
@@ -58,3 +76,99 @@ test('fetches table metadata for schema-less database without schema', async ()
|
||||
);
|
||||
});
|
||||
});
|
||||
|
||||
test('renders a fallback message for an unstructured metadata error', async () => {
|
||||
jest.spyOn(SupersetClient, 'get').mockRejectedValue({
|
||||
response: new Response('{}', {
|
||||
status: 500,
|
||||
headers: { 'Content-Type': 'application/json' },
|
||||
}),
|
||||
});
|
||||
errorMessageRegistry.registerValue(
|
||||
ErrorTypeEnum.GENERIC_BACKEND_ERROR,
|
||||
DatabaseErrorMessage,
|
||||
);
|
||||
|
||||
render(
|
||||
<DatasetPanelWrapper
|
||||
tableName="broken_table"
|
||||
dbId={1}
|
||||
database={{ supports_schemas: false }}
|
||||
/>,
|
||||
{ useRouter: true },
|
||||
);
|
||||
|
||||
expect(
|
||||
await screen.findByText('Unable to load columns for the selected table.'),
|
||||
).toBeVisible();
|
||||
});
|
||||
|
||||
test('retries only table metadata after matching OAuth completion', async () => {
|
||||
const oauthError = {
|
||||
error_type: ErrorTypeEnum.OAUTH2_REDIRECT,
|
||||
message: 'OAuth authorization is required.',
|
||||
extra: {
|
||||
url: 'https://example.com/authorize',
|
||||
tab_id: 'dataset-oauth-tab',
|
||||
},
|
||||
level: 'warning',
|
||||
};
|
||||
const getSpy = jest
|
||||
.spyOn(SupersetClient, 'get')
|
||||
.mockRejectedValueOnce({
|
||||
response: new Response(JSON.stringify({ errors: [oauthError] }), {
|
||||
status: 403,
|
||||
headers: { 'Content-Type': 'application/json' },
|
||||
}),
|
||||
})
|
||||
.mockResolvedValueOnce(tableMetadataResponse('oauth_table', 'oauth_id'));
|
||||
|
||||
errorMessageRegistry.registerValue(
|
||||
ErrorTypeEnum.OAUTH2_REDIRECT,
|
||||
OAuth2RedirectMessage,
|
||||
);
|
||||
|
||||
render(
|
||||
<DatasetPanelWrapper
|
||||
tableName="oauth_table"
|
||||
dbId={1}
|
||||
database={{ supports_schemas: false }}
|
||||
/>,
|
||||
{
|
||||
initialState: {
|
||||
charts: {},
|
||||
dashboardInfo: {},
|
||||
explore: {},
|
||||
sqlLab: {
|
||||
queries: {},
|
||||
queryEditors: [],
|
||||
tabHistory: [],
|
||||
},
|
||||
},
|
||||
useRedux: true,
|
||||
useRouter: true,
|
||||
},
|
||||
);
|
||||
|
||||
const authorizationLink = await screen.findByRole('link', {
|
||||
name: /provide authorization/i,
|
||||
});
|
||||
expect(authorizationLink).toHaveAttribute(
|
||||
'href',
|
||||
'https://example.com/authorize',
|
||||
);
|
||||
expect(getSpy).toHaveBeenCalledTimes(1);
|
||||
|
||||
act(() => {
|
||||
window.dispatchEvent(
|
||||
new StorageEvent('storage', {
|
||||
key: 'oauth2_auth_complete',
|
||||
newValue: JSON.stringify({ tabId: 'dataset-oauth-tab' }),
|
||||
}),
|
||||
);
|
||||
});
|
||||
|
||||
expect(await screen.findByText('oauth_id')).toBeVisible();
|
||||
expect(getSpy).toHaveBeenCalledTimes(2);
|
||||
expect(getSpy.mock.calls[1]).toEqual(getSpy.mock.calls[0]);
|
||||
});
|
||||
|
||||
@@ -65,27 +65,17 @@ export const NO_COLUMNS_TITLE = t('No table columns');
|
||||
export const NO_COLUMNS_DESCRIPTION = t(
|
||||
'This database table does not contain any data. Please select a different table.',
|
||||
);
|
||||
export const ERROR_TITLE = t('An Error Occurred');
|
||||
export const ERROR_DESCRIPTION = t(
|
||||
'Unable to load columns for the selected table. Please select a different table.',
|
||||
);
|
||||
|
||||
interface MessageContentProps {
|
||||
hasError: boolean;
|
||||
tableName?: string | null;
|
||||
hasColumns: boolean;
|
||||
}
|
||||
|
||||
export const MessageContent = (props: MessageContentProps) => {
|
||||
const { hasError, tableName, hasColumns } = props;
|
||||
let currentImage: string | undefined = 'empty-dataset.svg';
|
||||
const { tableName } = props;
|
||||
let currentImage = 'empty-dataset.svg';
|
||||
let currentTitle = SELECT_TABLE_TITLE;
|
||||
let currentDescription = renderEmptyDescription();
|
||||
if (hasError) {
|
||||
currentTitle = ERROR_TITLE;
|
||||
currentDescription = <>{ERROR_DESCRIPTION}</>;
|
||||
currentImage = undefined;
|
||||
} else if (tableName && !hasColumns) {
|
||||
if (tableName) {
|
||||
currentImage = 'no-columns.svg';
|
||||
currentTitle = NO_COLUMNS_TITLE;
|
||||
currentDescription = <>{NO_COLUMNS_DESCRIPTION}</>;
|
||||
|
||||
@@ -16,9 +16,14 @@
|
||||
* specific language governing permissions and limitations
|
||||
* under the License.
|
||||
*/
|
||||
import { useEffect, useState, useRef } from 'react';
|
||||
import { useCallback, useEffect, useRef, useState } from 'react';
|
||||
import { t } from '@apache-superset/core/translation';
|
||||
import { SupersetClient } from '@superset-ui/core';
|
||||
import {
|
||||
ErrorTypeEnum,
|
||||
getClientErrorObject,
|
||||
SupersetClient,
|
||||
} from '@superset-ui/core';
|
||||
import type { SupersetError } from '@superset-ui/core';
|
||||
import { logging } from '@apache-superset/core/utils';
|
||||
import { DatasetObject } from 'src/features/datasets/AddDataset/types';
|
||||
import { addDangerToast } from 'src/components/MessageToasts/actions';
|
||||
@@ -30,7 +35,7 @@ import { ITableColumn, IDatabaseTable, isIDatabaseTable } from './types';
|
||||
/**
|
||||
* Interface for the getTableMetadata API call
|
||||
*/
|
||||
interface IColumnProps {
|
||||
interface TableMetadataRequest {
|
||||
/**
|
||||
* Unique id of the database
|
||||
*/
|
||||
@@ -43,6 +48,10 @@ interface IColumnProps {
|
||||
* Name of the schema (optional for databases that don't support schemas)
|
||||
*/
|
||||
schema?: string | null;
|
||||
/**
|
||||
* Name of the catalog (optional for databases that don't support catalogs)
|
||||
*/
|
||||
catalog?: string | null;
|
||||
}
|
||||
|
||||
export interface IDatasetPanelWrapperProps {
|
||||
@@ -63,7 +72,7 @@ export interface IDatasetPanelWrapperProps {
|
||||
* The selected database object (used to check engine capabilities)
|
||||
*/
|
||||
database?: Partial<DatabaseObject> | null;
|
||||
setHasColumns?: Function;
|
||||
setHasColumns?: (hasColumns: boolean) => void;
|
||||
datasets?: DatasetObject[] | undefined;
|
||||
}
|
||||
|
||||
@@ -78,74 +87,131 @@ const DatasetPanelWrapper = ({
|
||||
}: IDatasetPanelWrapperProps) => {
|
||||
const [columnList, setColumnList] = useState<ITableColumn[]>([]);
|
||||
const [loading, setLoading] = useState(false);
|
||||
const [hasError, setHasError] = useState(false);
|
||||
const tableNameRef = useRef(tableName);
|
||||
const [error, setError] = useState<SupersetError>();
|
||||
const requestIdRef = useRef(0);
|
||||
const currentRequestRef = useRef<TableMetadataRequest>();
|
||||
const supportsSchemas = database?.supports_schemas;
|
||||
|
||||
const getTableMetadata = async (props: IColumnProps) => {
|
||||
const { dbId, tableName, schema } = props;
|
||||
setLoading(true);
|
||||
setHasColumns?.(false);
|
||||
const path = `/api/v1/database/${dbId}/table_metadata/${toQueryString({
|
||||
name: tableName,
|
||||
catalog,
|
||||
schema,
|
||||
})}`;
|
||||
try {
|
||||
const response = await SupersetClient.get({
|
||||
endpoint: path,
|
||||
});
|
||||
const getTableMetadata = useCallback(
|
||||
async (props: TableMetadataRequest) => {
|
||||
const { dbId, tableName, catalog, schema } = props;
|
||||
requestIdRef.current += 1;
|
||||
const requestId = requestIdRef.current;
|
||||
setLoading(true);
|
||||
setColumnList([]);
|
||||
setError(undefined);
|
||||
setHasColumns?.(false);
|
||||
const path = `/api/v1/database/${dbId}/table_metadata/${toQueryString({
|
||||
name: tableName,
|
||||
catalog,
|
||||
schema,
|
||||
})}`;
|
||||
try {
|
||||
const response = await SupersetClient.get({
|
||||
endpoint: path,
|
||||
});
|
||||
|
||||
if (isIDatabaseTable(response?.json)) {
|
||||
const table: IDatabaseTable = response.json as IDatabaseTable;
|
||||
/**
|
||||
* The user is able to click other table columns while the http call for last selected table column is made
|
||||
* This check ensures we process the response that matches the last selected table name and ignore the others
|
||||
*/
|
||||
if (table.name === tableNameRef.current) {
|
||||
if (requestId !== requestIdRef.current) {
|
||||
return;
|
||||
}
|
||||
|
||||
const table = isIDatabaseTable(response?.json)
|
||||
? (response.json as IDatabaseTable)
|
||||
: undefined;
|
||||
if (table?.name === tableName) {
|
||||
setColumnList(table.columns);
|
||||
setHasColumns?.(table.columns.length > 0);
|
||||
setHasError(false);
|
||||
setError(undefined);
|
||||
} else {
|
||||
const message = t(
|
||||
'The API response from %s does not match the IDatabaseTable interface.',
|
||||
path,
|
||||
);
|
||||
setColumnList([]);
|
||||
setHasColumns?.(false);
|
||||
setError({
|
||||
error_type: ErrorTypeEnum.GENERIC_BACKEND_ERROR,
|
||||
extra: null,
|
||||
level: 'error',
|
||||
message,
|
||||
});
|
||||
addDangerToast(message);
|
||||
logging.error(message);
|
||||
}
|
||||
} else {
|
||||
setColumnList([]);
|
||||
setHasColumns?.(false);
|
||||
setHasError(true);
|
||||
addDangerToast(
|
||||
t(
|
||||
'The API response from %s does not match the IDatabaseTable interface.',
|
||||
path,
|
||||
),
|
||||
);
|
||||
logging.error(
|
||||
t(
|
||||
'The API response from %s does not match the IDatabaseTable interface.',
|
||||
path,
|
||||
),
|
||||
} catch (caughtError) {
|
||||
const clientError = await getClientErrorObject(
|
||||
caughtError as Parameters<typeof getClientErrorObject>[0],
|
||||
);
|
||||
|
||||
if (requestId === requestIdRef.current) {
|
||||
const parsedError = clientError.errors?.[0] ?? {
|
||||
error_type: ErrorTypeEnum.GENERIC_BACKEND_ERROR,
|
||||
extra: null,
|
||||
level: 'error' as const,
|
||||
message:
|
||||
clientError.error ||
|
||||
clientError.message ||
|
||||
clientError.statusText ||
|
||||
t('Unable to load columns for the selected table.'),
|
||||
};
|
||||
|
||||
setColumnList([]);
|
||||
setHasColumns?.(false);
|
||||
setError(parsedError);
|
||||
}
|
||||
} finally {
|
||||
if (requestId === requestIdRef.current) {
|
||||
setLoading(false);
|
||||
}
|
||||
}
|
||||
} catch (error) {
|
||||
setColumnList([]);
|
||||
setHasColumns?.(false);
|
||||
setHasError(true);
|
||||
} finally {
|
||||
setLoading(false);
|
||||
},
|
||||
[setHasColumns],
|
||||
);
|
||||
|
||||
const retryGetTableMetadata = useCallback(() => {
|
||||
if (currentRequestRef.current) {
|
||||
getTableMetadata(currentRequestRef.current);
|
||||
}
|
||||
};
|
||||
}, [getTableMetadata]);
|
||||
|
||||
useEffect(() => {
|
||||
tableNameRef.current = tableName;
|
||||
const schemaRequired = database?.supports_schemas !== false;
|
||||
const schemaRequired = supportsSchemas !== false;
|
||||
if (tableName && dbId && (schema || !schemaRequired)) {
|
||||
getTableMetadata({ tableName, dbId, schema: schema || undefined });
|
||||
const request = {
|
||||
tableName,
|
||||
dbId,
|
||||
catalog,
|
||||
schema: schema || undefined,
|
||||
};
|
||||
currentRequestRef.current = request;
|
||||
getTableMetadata(request);
|
||||
} else if (currentRequestRef.current) {
|
||||
currentRequestRef.current = undefined;
|
||||
requestIdRef.current += 1;
|
||||
setColumnList([]);
|
||||
setError(undefined);
|
||||
setHasColumns?.(false);
|
||||
setLoading(false);
|
||||
}
|
||||
// getTableMetadata is a const and should not be in dependency array
|
||||
// eslint-disable-next-line react-hooks/exhaustive-deps
|
||||
}, [tableName, dbId, schema, database]);
|
||||
|
||||
return () => {
|
||||
requestIdRef.current += 1;
|
||||
};
|
||||
}, [
|
||||
tableName,
|
||||
dbId,
|
||||
catalog,
|
||||
schema,
|
||||
supportsSchemas,
|
||||
getTableMetadata,
|
||||
setHasColumns,
|
||||
]);
|
||||
|
||||
return (
|
||||
<DatasetPanel
|
||||
columnList={columnList}
|
||||
hasError={hasError}
|
||||
error={error}
|
||||
errorMitigationFunction={retryGetTableMetadata}
|
||||
loading={loading}
|
||||
tableName={tableName}
|
||||
datasets={datasets}
|
||||
|
||||
@@ -138,7 +138,6 @@ describe('DatasetLayout', () => {
|
||||
<DatasetPanelComponent
|
||||
tableName="large_table"
|
||||
columnList={manyColumns}
|
||||
hasError={false}
|
||||
loading={false}
|
||||
/>
|
||||
}
|
||||
|
||||
@@ -161,6 +161,9 @@ def get_query(query_id: int) -> Query:
|
||||
try:
|
||||
return db.session.query(Query).filter_by(id=query_id).one()
|
||||
except Exception as ex:
|
||||
# roll back so a poisoned session (e.g. PendingRollbackError after a
|
||||
# failed flush) doesn't fail every subsequent backoff retry identically
|
||||
db.session.rollback()
|
||||
raise SqlLabException("Failed at getting query") from ex
|
||||
|
||||
|
||||
|
||||
@@ -87,6 +87,29 @@ def resolve_screenshot_task_budget_seconds(
|
||||
return None
|
||||
|
||||
|
||||
# Fallback wall-clock budget, in seconds, for the entire tiled-screenshot
|
||||
# operation (element lookup plus all per-tile readiness/animation waits
|
||||
# combined), used when resolve_screenshot_task_budget_seconds() returns None
|
||||
# (no Celery task context -- e.g. synchronous thumbnail generation -- or no
|
||||
# usable task limit). The non-tiled readiness path treats None as "keep the
|
||||
# configured SCREENSHOT_LOAD_WAIT" because it makes exactly one bounded wait;
|
||||
# the tiled path cannot, because its per-tile waits accumulate: with N tiles,
|
||||
# an uncapped load_wait allows N * load_wait of total wall-clock time, so the
|
||||
# operation still needs one fixed total ceiling. Sized against the longest
|
||||
# Celery hard task_time_limit observed in production for report execution
|
||||
# (1740s), minus the same 300s cleanup margin the runtime derivation reserves
|
||||
# for combining tiles, building the PDF, and delivering the notification.
|
||||
TILED_SCREENSHOT_TOTAL_WAIT_BUDGET_SECONDS = 1440 # 1740s limit - 300s margin
|
||||
|
||||
|
||||
class ScreenshotTaskBudgetExceededError(RuntimeError):
|
||||
"""Raised when no safe task budget remains before screenshot capture."""
|
||||
|
||||
|
||||
class TiledScreenshotBudgetExceededError(ScreenshotTaskBudgetExceededError):
|
||||
"""Raised when the tiled-screenshot time budget runs out mid-capture."""
|
||||
|
||||
|
||||
try:
|
||||
from playwright.sync_api import TimeoutError as PlaywrightTimeout
|
||||
except ImportError:
|
||||
@@ -251,7 +274,7 @@ def combine_screenshot_tiles(screenshot_tiles: list[bytes]) -> bytes:
|
||||
return screenshot_tiles[0]
|
||||
|
||||
|
||||
def take_tiled_screenshot(
|
||||
def take_tiled_screenshot( # noqa: C901
|
||||
page: "Page",
|
||||
element_name: str,
|
||||
tile_height: int,
|
||||
@@ -274,6 +297,12 @@ def take_tiled_screenshot(
|
||||
|
||||
Returns:
|
||||
Combined screenshot bytes or None if failed
|
||||
|
||||
Raises:
|
||||
TiledScreenshotBudgetExceededError: If the total time budget for the
|
||||
tiled-screenshot operation runs out before every tile has been
|
||||
verifiably captured. Callers must treat this as a hard failure
|
||||
rather than fall back to an unchecked/partial screenshot.
|
||||
"""
|
||||
context_suffix = f" [{log_context}]" if log_context else ""
|
||||
# Set right before re-raising the per-tile readiness timeout below, and
|
||||
@@ -286,6 +315,15 @@ def take_tiled_screenshot(
|
||||
# match `except PlaywrightTimeout` and incorrectly propagate instead of
|
||||
# degrading to `None` like every other unexpected error in this function.
|
||||
readiness_timeout = False
|
||||
# Cap the whole tiled operation against the running Celery task's own
|
||||
# time limit, using the same runtime derivation as the non-tiled
|
||||
# readiness wait (#42253/#42427). Unlike that path, a None budget does
|
||||
# not mean "keep the configured timeout": per-tile waits accumulate, so
|
||||
# the operation falls back to a fixed total ceiling instead.
|
||||
wait_budget_seconds = resolve_screenshot_task_budget_seconds(log_context)
|
||||
if wait_budget_seconds is None:
|
||||
wait_budget_seconds = float(TILED_SCREENSHOT_TOTAL_WAIT_BUDGET_SECONDS)
|
||||
start_time = time.monotonic()
|
||||
try:
|
||||
# Get the target element
|
||||
element = page.locator(f".{element_name}")
|
||||
@@ -320,9 +358,44 @@ def take_tiled_screenshot(
|
||||
num_tiles = max(1, (dashboard_height + tile_height - 1) // tile_height)
|
||||
logger.info("Taking %s screenshot tiles", num_tiles)
|
||||
|
||||
screenshot_tiles = []
|
||||
screenshot_tiles: list[bytes] = []
|
||||
|
||||
def _raise_if_budget_exhausted(elapsed: float, remaining_budget: float) -> None:
|
||||
if remaining_budget > 0:
|
||||
return
|
||||
# A customer-side chart-loading issue (a slow/hung dashboard),
|
||||
# not a Superset system fault, so this is a WARNING rather
|
||||
# than an ERROR -- consistent with #38130/#38441, which
|
||||
# deliberately downgraded screenshot timeout logs the same way.
|
||||
logger.warning(
|
||||
"Tiled screenshot time budget exhausted on tile %s/%s: "
|
||||
"%s/%s tiles captured so far, %.1fs elapsed of a %.1fs "
|
||||
"budget. Aborting instead of capturing remaining tiles "
|
||||
"unchecked.%s",
|
||||
i + 1,
|
||||
num_tiles,
|
||||
len(screenshot_tiles),
|
||||
num_tiles,
|
||||
elapsed,
|
||||
wait_budget_seconds,
|
||||
context_suffix,
|
||||
)
|
||||
raise TiledScreenshotBudgetExceededError(
|
||||
f"Tiled screenshot budget of "
|
||||
f"{wait_budget_seconds:.1f}s exhausted "
|
||||
f"after {len(screenshot_tiles)}/{num_tiles} tiles"
|
||||
)
|
||||
|
||||
for i in range(num_tiles):
|
||||
# Check the time budget before starting this tile's readiness wait.
|
||||
# If it's already exhausted, we can no longer verify this (or any
|
||||
# later) tile is actually ready to capture -- fail loudly instead
|
||||
# of silently snapshotting a spinner or blank chart, or running
|
||||
# past the Celery task time limit and getting SIGKILLed.
|
||||
elapsed = time.monotonic() - start_time
|
||||
remaining_budget = wait_budget_seconds - elapsed
|
||||
_raise_if_budget_exhausted(elapsed, remaining_budget)
|
||||
|
||||
# Calculate scroll position to show this tile's content
|
||||
scroll_y = dashboard_top + (i * tile_height)
|
||||
|
||||
@@ -332,17 +405,31 @@ def take_tiled_screenshot(
|
||||
)
|
||||
# Wait for scroll to settle and content to load
|
||||
page.wait_for_timeout(SCROLL_SETTLE_TIMEOUT_MS)
|
||||
|
||||
# Recompute the remaining budget after the scroll-settle sleep --
|
||||
# which itself consumes real wall-clock time -- rather than
|
||||
# reusing the value from before it, so the readiness-check
|
||||
# timeout below is capped against a fresh number instead of a
|
||||
# stale one that would let each tile overrun the budget by up
|
||||
# to one settle interval.
|
||||
tile_wait_start = time.monotonic()
|
||||
elapsed = tile_wait_start - start_time
|
||||
remaining_budget = wait_budget_seconds - elapsed
|
||||
_raise_if_budget_exhausted(elapsed, remaining_budget)
|
||||
|
||||
# Wait for every chart holder visible in the current viewport to reach
|
||||
# a terminal state (rendered chart or error/empty state). Only check
|
||||
# a terminal state (rendered chart or error/empty state), capped at
|
||||
# whatever remains of the total time budget so a slow dashboard
|
||||
# degrades gracefully instead of exceeding it. Only check
|
||||
# viewport-visible chart holders to avoid blocking on virtualization
|
||||
# placeholders rendered for off-screen charts. A holder that hasn't
|
||||
# mounted anything yet does not satisfy this check -- unlike checking
|
||||
# for the absence of `.loading`, which passes vacuously in that case.
|
||||
tile_wait_start = time.monotonic()
|
||||
tile_load_wait = min(load_wait, remaining_budget)
|
||||
try:
|
||||
page.wait_for_function(
|
||||
CHART_HOLDERS_READY_JS,
|
||||
timeout=load_wait * 1000,
|
||||
timeout=tile_load_wait * 1000,
|
||||
)
|
||||
except PlaywrightTimeout:
|
||||
elapsed = time.monotonic() - tile_wait_start
|
||||
@@ -354,14 +441,21 @@ def take_tiled_screenshot(
|
||||
# made the same call for the other screenshot timeout paths.
|
||||
logger.warning(
|
||||
"Timed out after %.2fs waiting for %s chart container(s) to "
|
||||
"become ready on tile %s/%s (load_wait=%ss)%s; unready chart "
|
||||
"holders (chart id, state): %s. Aborting tiled screenshot "
|
||||
"rather than capturing a blank or partially-loaded tile.",
|
||||
"become ready on tile %s/%s (waited %.1fs of a %ss requested "
|
||||
"load_wait; %.1fs elapsed of a %.1fs total budget; %s/%s "
|
||||
"tiles captured so far)%s; unready chart holders (chart id, "
|
||||
"state): %s. Aborting tiled screenshot rather than capturing "
|
||||
"a blank or partially-loaded tile.",
|
||||
elapsed,
|
||||
len(unready_chart_holders),
|
||||
i + 1,
|
||||
num_tiles,
|
||||
tile_load_wait,
|
||||
load_wait,
|
||||
time.monotonic() - start_time,
|
||||
wait_budget_seconds,
|
||||
len(screenshot_tiles),
|
||||
num_tiles,
|
||||
context_suffix,
|
||||
unready_chart_holders,
|
||||
)
|
||||
@@ -377,12 +471,36 @@ def take_tiled_screenshot(
|
||||
load_wait,
|
||||
context_suffix,
|
||||
)
|
||||
readiness_wait_elapsed = time.monotonic() - tile_wait_start
|
||||
|
||||
# Wait for chart animations (e.g. ECharts) to finish after spinner clears.
|
||||
# The global animation wait before tiling only covers the first tile;
|
||||
# subsequent tiles need their own wait after data loads.
|
||||
# subsequent tiles need their own wait after data loads. Capped at
|
||||
# whatever remains of the budget; unlike the readiness wait above this
|
||||
# is cosmetic settling, not a readiness check, so we simply skip it
|
||||
# (rather than raise) once the budget runs out.
|
||||
animation_wait_elapsed = 0.0
|
||||
if animation_wait > 0:
|
||||
page.wait_for_timeout(animation_wait * 1000)
|
||||
elapsed = time.monotonic() - start_time
|
||||
remaining_budget = wait_budget_seconds - elapsed
|
||||
tile_animation_wait = max(0, min(animation_wait, remaining_budget))
|
||||
if tile_animation_wait > 0:
|
||||
animation_wait_start = time.monotonic()
|
||||
page.wait_for_timeout(tile_animation_wait * 1000)
|
||||
animation_wait_elapsed = time.monotonic() - animation_wait_start
|
||||
|
||||
# Per-tile timing breakdown so slow dashboards can be profiled from
|
||||
# logs alone. DEBUG rather than INFO: this fires once per tile, and
|
||||
# large dashboards can have dozens of tiles per report run.
|
||||
logger.debug(
|
||||
"Tile %s/%s timing: %.2fs waiting for chart readiness, "
|
||||
"%.2fs waiting for animations.%s",
|
||||
i + 1,
|
||||
num_tiles,
|
||||
readiness_wait_elapsed,
|
||||
animation_wait_elapsed,
|
||||
context_suffix,
|
||||
)
|
||||
|
||||
# Calculate what portion of the element we want to capture for this tile
|
||||
tile_start_in_element = i * tile_height
|
||||
@@ -431,6 +549,12 @@ def take_tiled_screenshot(
|
||||
|
||||
return combined_screenshot
|
||||
|
||||
except TiledScreenshotBudgetExceededError:
|
||||
# Budget exhaustion must fail cleanly, not be swallowed into the
|
||||
# generic `return None` degradation below -- the raise carries the
|
||||
# budget diagnostics to the caller, which fails the capture loudly
|
||||
# (#42273) instead of receiving an anonymous empty result.
|
||||
raise
|
||||
except Exception as e:
|
||||
if readiness_timeout:
|
||||
# Let the per-tile readiness timeout propagate so the caller
|
||||
|
||||
@@ -85,6 +85,26 @@ class ScreenshotCachePayloadType(TypedDict):
|
||||
status: str
|
||||
|
||||
|
||||
# Magic bytes for a cheap image sanity check. This is intentionally not a full
|
||||
# decode: it's meant to catch 0-byte/corrupt/blank payloads before they're
|
||||
# cached or served, not to validate the image is renderable.
|
||||
PNG_MAGIC_BYTES = b"\x89PNG\r\n\x1a\n"
|
||||
JPEG_MAGIC_BYTES = b"\xff\xd8\xff"
|
||||
|
||||
|
||||
def validate_screenshot_image(image: bytes | None) -> str | None:
|
||||
"""Cheaply validate screenshot bytes before they're cached or served.
|
||||
|
||||
:return: None if the bytes look like a usable image, otherwise a short
|
||||
reason ("empty" or "undecodable") suitable for logging.
|
||||
"""
|
||||
if not image:
|
||||
return "empty"
|
||||
if not image.startswith((PNG_MAGIC_BYTES, JPEG_MAGIC_BYTES)):
|
||||
return "undecodable"
|
||||
return None
|
||||
|
||||
|
||||
class ScreenshotCachePayload:
|
||||
def __init__(
|
||||
self,
|
||||
@@ -147,6 +167,13 @@ class ScreenshotCachePayload:
|
||||
def get_status(self) -> str:
|
||||
return self.status.value
|
||||
|
||||
def get_invalid_image_reason(self) -> str | None:
|
||||
"""Reason this payload's image should not be served/cached, or None if
|
||||
it passes validation (or it isn't claiming a successful screenshot)."""
|
||||
if self.status != StatusValues.UPDATED:
|
||||
return None
|
||||
return validate_screenshot_image(self._image)
|
||||
|
||||
def is_error_cache_ttl_expired(self) -> bool:
|
||||
error_cache_ttl = app.config["THUMBNAIL_ERROR_CACHE_TTL"]
|
||||
return (
|
||||
@@ -263,6 +290,14 @@ class BaseScreenshot:
|
||||
elif isinstance(payload, dict):
|
||||
payload = cast(ScreenshotCachePayloadType, payload)
|
||||
payload = ScreenshotCachePayload.from_dict(payload)
|
||||
if invalid_reason := payload.get_invalid_image_reason():
|
||||
logger.warning(
|
||||
"Rejecting cached screenshot for %s: %s image payload; "
|
||||
"treating as a cache miss",
|
||||
cache_key,
|
||||
invalid_reason,
|
||||
)
|
||||
return None
|
||||
return payload
|
||||
logger.info("Failed at getting from cache: %s", cache_key)
|
||||
return None
|
||||
@@ -331,15 +366,28 @@ class BaseScreenshot:
|
||||
image = None
|
||||
|
||||
# Cache the result (success or error) to avoid immediate retries
|
||||
if image:
|
||||
invalid_reason = validate_screenshot_image(image)
|
||||
# `image and` is redundant at runtime (validate_screenshot_image
|
||||
# only returns None for truthy, well-formed bytes) but mypy can't
|
||||
# infer that image is non-None from invalid_reason being None
|
||||
# across the function-call boundary, so it's kept for narrowing.
|
||||
if image and invalid_reason is None:
|
||||
with event_logger.log_context(
|
||||
f"screenshot.cache.{self.thumbnail_type}"
|
||||
):
|
||||
cache_payload.update(image)
|
||||
elif cache_payload.status != StatusValues.ERROR:
|
||||
# Only call error() if not already set — avoids overwriting
|
||||
# the timestamp recorded when the actual failure occurred above.
|
||||
cache_payload.error()
|
||||
else:
|
||||
if invalid_reason:
|
||||
logger.warning(
|
||||
"Not caching screenshot result for %s: %s image payload",
|
||||
cache_key,
|
||||
invalid_reason,
|
||||
)
|
||||
if cache_payload.status != StatusValues.ERROR:
|
||||
# Only call error() if not already set — avoids overwriting
|
||||
# the timestamp recorded when the actual failure occurred
|
||||
# above.
|
||||
cache_payload.error()
|
||||
|
||||
logger.info("Caching thumbnail: %s", cache_key)
|
||||
self.cache.set(cache_key, cache_payload.to_dict())
|
||||
|
||||
@@ -47,6 +47,7 @@ from superset.utils.screenshot_utils import (
|
||||
CHART_HOLDERS_READY_JS,
|
||||
FIND_CHART_HOLDER_STATES_JS,
|
||||
resolve_screenshot_task_budget_seconds,
|
||||
ScreenshotTaskBudgetExceededError,
|
||||
take_tiled_screenshot,
|
||||
)
|
||||
|
||||
@@ -61,10 +62,6 @@ PLAYWRIGHT_INSTALL_MESSAGE = (
|
||||
)
|
||||
|
||||
|
||||
class ScreenshotTaskBudgetExceededError(RuntimeError):
|
||||
"""Raised when no safe task budget remains before screenshot capture."""
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from typing import Any
|
||||
|
||||
|
||||
@@ -35,6 +35,7 @@ from superset.sql.parse import SQLStatement, Table
|
||||
from superset.sql_lab import (
|
||||
execute_query,
|
||||
execute_sql_statements,
|
||||
get_query,
|
||||
get_sql_results,
|
||||
)
|
||||
from superset.utils.rls import apply_rls, get_predicates_for_table
|
||||
@@ -71,6 +72,34 @@ def test_execute_query(mocker: MockerFixture, app: None) -> None:
|
||||
SupersetResultSet.assert_called_with([(42,)], cursor.description, db_engine_spec)
|
||||
|
||||
|
||||
def test_get_query_rolls_back_session_before_retrying(
|
||||
mocker: MockerFixture, app: SupersetApp
|
||||
) -> None:
|
||||
"""
|
||||
A broken transaction (e.g. `PendingRollbackError` following a failed flush)
|
||||
leaves the session unusable until `session.rollback()` is called, so without
|
||||
it every `backoff` retry would reuse the same poisoned session and fail
|
||||
identically. `get_query` must roll back on failure so each retry gets a
|
||||
clean session and has a real chance to succeed.
|
||||
"""
|
||||
# avoid actually sleeping through the `backoff` decorator's retry interval
|
||||
mocker.patch("backoff._sync.time.sleep")
|
||||
|
||||
expected_query = mocker.MagicMock()
|
||||
mock_one = mocker.patch("superset.sql_lab.db.session.query")
|
||||
mock_one.return_value.filter_by.return_value.one.side_effect = [
|
||||
Exception("session is broken"),
|
||||
expected_query,
|
||||
]
|
||||
mock_rollback = mocker.patch("superset.sql_lab.db.session.rollback")
|
||||
|
||||
result = get_query(query_id=1)
|
||||
|
||||
assert result is expected_query
|
||||
assert mock_one.return_value.filter_by.return_value.one.call_count == 2
|
||||
mock_rollback.assert_called_once()
|
||||
|
||||
|
||||
@with_config(
|
||||
{
|
||||
"SQLLAB_PAYLOAD_MAX_MB": 50,
|
||||
|
||||
@@ -33,6 +33,10 @@ from superset.utils.screenshots import (
|
||||
|
||||
BASE_SCREENSHOT_PATH = "superset.utils.screenshots.BaseScreenshot"
|
||||
|
||||
# A minimal valid PNG header, used wherever a test needs bytes that pass
|
||||
# ScreenshotCachePayload's image validation.
|
||||
FAKE_PNG_BYTES = b"\x89PNG\r\n\x1a\n" + b"fake-png-body"
|
||||
|
||||
|
||||
class MockCache:
|
||||
"""A class to manage screenshot cache."""
|
||||
@@ -92,7 +96,7 @@ def test_get_cache_key(app_context, screenshot_obj):
|
||||
def test_get_from_cache_key(mocker: MockerFixture, screenshot_obj):
|
||||
"""get_from_cache_key should always return a ScreenshotCachePayload Object"""
|
||||
# backwards compatibility test for retrieving plain bytes
|
||||
fake_bytes = b"fake_screenshot_data"
|
||||
fake_bytes = FAKE_PNG_BYTES
|
||||
BaseScreenshot.cache = MockCache()
|
||||
BaseScreenshot.cache.set("key", fake_bytes)
|
||||
cache_payload = screenshot_obj.get_from_cache_key("key")
|
||||
@@ -108,10 +112,10 @@ class TestComputeAndCache:
|
||||
BASE_SCREENSHOT_PATH + ".get_from_cache_key", return_value=None
|
||||
)
|
||||
get_screenshot = mocker.patch(
|
||||
BASE_SCREENSHOT_PATH + ".get_screenshot", return_value=b"new_image_data"
|
||||
BASE_SCREENSHOT_PATH + ".get_screenshot", return_value=FAKE_PNG_BYTES
|
||||
)
|
||||
resize_image = mocker.patch(
|
||||
BASE_SCREENSHOT_PATH + ".resize_image", return_value=b"resized_image_data"
|
||||
BASE_SCREENSHOT_PATH + ".resize_image", return_value=FAKE_PNG_BYTES
|
||||
)
|
||||
BaseScreenshot.cache = MockCache()
|
||||
return {
|
||||
|
||||
@@ -37,6 +37,10 @@ from superset.utils.screenshots import (
|
||||
BASE_SCREENSHOT_PATH = "superset.utils.screenshots.BaseScreenshot"
|
||||
DISTRIBUTED_LOCK_PATH = "superset.utils.screenshots.DistributedLock"
|
||||
|
||||
# A minimal valid PNG header, used wherever a test needs bytes that pass
|
||||
# ScreenshotCachePayload's image validation.
|
||||
FAKE_PNG_BYTES = b"\x89PNG\r\n\x1a\n" + b"fake-png-body"
|
||||
|
||||
|
||||
class MockCache:
|
||||
"""A class to manage screenshot cache for testing."""
|
||||
@@ -83,11 +87,11 @@ class TestCacheOnlyOnSuccess:
|
||||
mocker.patch(DISTRIBUTED_LOCK_PATH)
|
||||
mocker.patch(BASE_SCREENSHOT_PATH + ".get_from_cache_key", return_value=None)
|
||||
get_screenshot = mocker.patch(
|
||||
BASE_SCREENSHOT_PATH + ".get_screenshot", return_value=b"image_data"
|
||||
BASE_SCREENSHOT_PATH + ".get_screenshot", return_value=FAKE_PNG_BYTES
|
||||
)
|
||||
# Mock resize_image to avoid PIL errors with fake image data
|
||||
mocker.patch(
|
||||
BASE_SCREENSHOT_PATH + ".resize_image", return_value=b"resized_image_data"
|
||||
BASE_SCREENSHOT_PATH + ".resize_image", return_value=FAKE_PNG_BYTES
|
||||
)
|
||||
BaseScreenshot.cache = MockCache()
|
||||
return get_screenshot
|
||||
@@ -161,13 +165,15 @@ class TestCacheOnlyOnSuccess:
|
||||
screenshot_obj: BaseScreenshot,
|
||||
mock_user: MagicMock,
|
||||
) -> None:
|
||||
"""Empty bytes from get_screenshot must set ERROR, not leave COMPUTING."""
|
||||
"""Empty bytes from get_screenshot must set ERROR, not leave COMPUTING,
|
||||
and must log a WARNING that includes the cache key."""
|
||||
mocker.patch(DISTRIBUTED_LOCK_PATH)
|
||||
mocker.patch(BASE_SCREENSHOT_PATH + ".get_from_cache_key", return_value=None)
|
||||
mocker.patch(
|
||||
BASE_SCREENSHOT_PATH + ".get_screenshot",
|
||||
return_value=b"",
|
||||
)
|
||||
mock_logger = mocker.patch("superset.utils.screenshots.logger")
|
||||
BaseScreenshot.cache = MockCache()
|
||||
|
||||
screenshot_obj.compute_and_cache(user=mock_user, force=True)
|
||||
@@ -177,6 +183,43 @@ class TestCacheOnlyOnSuccess:
|
||||
assert cached_value is not None
|
||||
assert cached_value["status"] == "Error"
|
||||
assert cached_value.get("image") is None
|
||||
assert any(
|
||||
cache_key in call.args and "empty" in call.args
|
||||
for call in mock_logger.warning.call_args_list
|
||||
)
|
||||
|
||||
def test_cache_error_status_when_screenshot_returns_garbage_bytes(
|
||||
self,
|
||||
mocker: MockerFixture,
|
||||
screenshot_obj: BaseScreenshot,
|
||||
mock_user: MagicMock,
|
||||
) -> None:
|
||||
"""Non-empty bytes without a valid image header must set ERROR, not be
|
||||
cached as a success, and must log a WARNING that includes the cache key."""
|
||||
mocker.patch(DISTRIBUTED_LOCK_PATH)
|
||||
mocker.patch(BASE_SCREENSHOT_PATH + ".get_from_cache_key", return_value=None)
|
||||
mocker.patch(
|
||||
BASE_SCREENSHOT_PATH + ".get_screenshot",
|
||||
return_value=b"this-is-not-a-real-image",
|
||||
)
|
||||
mocker.patch(
|
||||
BASE_SCREENSHOT_PATH + ".resize_image",
|
||||
return_value=b"this-is-not-a-real-image",
|
||||
)
|
||||
mock_logger = mocker.patch("superset.utils.screenshots.logger")
|
||||
BaseScreenshot.cache = MockCache()
|
||||
|
||||
screenshot_obj.compute_and_cache(user=mock_user, force=True)
|
||||
|
||||
cache_key = screenshot_obj.get_cache_key()
|
||||
cached_value = BaseScreenshot.cache.get(cache_key)
|
||||
assert cached_value is not None
|
||||
assert cached_value["status"] == "Error"
|
||||
assert cached_value.get("image") is None
|
||||
assert any(
|
||||
cache_key in call.args and "undecodable" in call.args
|
||||
for call in mock_logger.warning.call_args_list
|
||||
)
|
||||
|
||||
def test_computing_status_written_to_cache_early(
|
||||
self,
|
||||
@@ -197,14 +240,14 @@ class TestCacheOnlyOnSuccess:
|
||||
"Cache should be set to COMPUTING before screenshot starts"
|
||||
)
|
||||
assert cached_value["status"] == "Computing"
|
||||
return b"image_data"
|
||||
return FAKE_PNG_BYTES
|
||||
|
||||
mocker.patch(
|
||||
BASE_SCREENSHOT_PATH + ".get_screenshot",
|
||||
side_effect=check_cache_during_screenshot,
|
||||
)
|
||||
mocker.patch(
|
||||
BASE_SCREENSHOT_PATH + ".resize_image", return_value=b"resized_image_data"
|
||||
BASE_SCREENSHOT_PATH + ".resize_image", return_value=FAKE_PNG_BYTES
|
||||
)
|
||||
|
||||
screenshot_obj.compute_and_cache(user=mock_user, force=True)
|
||||
@@ -429,11 +472,11 @@ class TestIntegrationCacheBugFix:
|
||||
BaseScreenshot.cache.set(cache_key, stale_payload.to_dict())
|
||||
|
||||
mocker.patch(
|
||||
BASE_SCREENSHOT_PATH + ".get_screenshot", return_value=b"recovered_image"
|
||||
BASE_SCREENSHOT_PATH + ".get_screenshot", return_value=FAKE_PNG_BYTES
|
||||
)
|
||||
# Mock resize to avoid PIL errors
|
||||
mocker.patch(
|
||||
BASE_SCREENSHOT_PATH + ".resize_image", return_value=b"resized_image"
|
||||
BASE_SCREENSHOT_PATH + ".resize_image", return_value=FAKE_PNG_BYTES
|
||||
)
|
||||
|
||||
# Should trigger task because COMPUTING is stale
|
||||
@@ -482,3 +525,72 @@ class TestIntegrationCacheBugFix:
|
||||
|
||||
assert payload._image == old_image
|
||||
assert payload.status == StatusValues.COMPUTING
|
||||
|
||||
|
||||
class TestReadSideImageValidation:
|
||||
"""A cached payload that claims a successful screenshot (status UPDATED)
|
||||
but carries invalid image bytes must be served as a cache miss, not
|
||||
returned to the caller — this is what the dashboard/chart screenshot
|
||||
endpoints call to fetch bytes to serve."""
|
||||
|
||||
def test_zero_byte_image_is_treated_as_cache_miss(
|
||||
self, mocker: MockerFixture, screenshot_obj: BaseScreenshot
|
||||
) -> None:
|
||||
mock_logger = mocker.patch("superset.utils.screenshots.logger")
|
||||
BaseScreenshot.cache = MockCache()
|
||||
cache_key = screenshot_obj.get_cache_key()
|
||||
stale_payload = ScreenshotCachePayload(image=b"", status=StatusValues.UPDATED)
|
||||
BaseScreenshot.cache.set(cache_key, stale_payload.to_dict())
|
||||
|
||||
result = screenshot_obj.get_from_cache_key(cache_key)
|
||||
|
||||
assert result is None
|
||||
assert any(
|
||||
cache_key in call.args and "empty" in call.args
|
||||
for call in mock_logger.warning.call_args_list
|
||||
)
|
||||
|
||||
def test_garbage_bytes_image_is_treated_as_cache_miss(
|
||||
self, mocker: MockerFixture, screenshot_obj: BaseScreenshot
|
||||
) -> None:
|
||||
mock_logger = mocker.patch("superset.utils.screenshots.logger")
|
||||
BaseScreenshot.cache = MockCache()
|
||||
cache_key = screenshot_obj.get_cache_key()
|
||||
garbage_payload = ScreenshotCachePayload(image=b"not-an-image-at-all")
|
||||
BaseScreenshot.cache.set(cache_key, garbage_payload.to_dict())
|
||||
|
||||
result = screenshot_obj.get_from_cache_key(cache_key)
|
||||
|
||||
assert result is None
|
||||
assert any(
|
||||
cache_key in call.args and "undecodable" in call.args
|
||||
for call in mock_logger.warning.call_args_list
|
||||
)
|
||||
|
||||
def test_valid_image_is_served_normally(
|
||||
self, screenshot_obj: BaseScreenshot
|
||||
) -> None:
|
||||
BaseScreenshot.cache = MockCache()
|
||||
cache_key = screenshot_obj.get_cache_key()
|
||||
valid_payload = ScreenshotCachePayload(image=FAKE_PNG_BYTES)
|
||||
BaseScreenshot.cache.set(cache_key, valid_payload.to_dict())
|
||||
|
||||
result = screenshot_obj.get_from_cache_key(cache_key)
|
||||
|
||||
assert result is not None
|
||||
assert result.get_image().read() == FAKE_PNG_BYTES
|
||||
|
||||
def test_pending_status_with_no_image_is_not_rejected(
|
||||
self, screenshot_obj: BaseScreenshot
|
||||
) -> None:
|
||||
"""Non-UPDATED statuses (e.g. PENDING/COMPUTING) aren't claiming a
|
||||
successful screenshot, so they should be returned as-is."""
|
||||
BaseScreenshot.cache = MockCache()
|
||||
cache_key = screenshot_obj.get_cache_key()
|
||||
pending_payload = ScreenshotCachePayload(status=StatusValues.PENDING)
|
||||
BaseScreenshot.cache.set(cache_key, pending_payload.to_dict())
|
||||
|
||||
result = screenshot_obj.get_from_cache_key(cache_key)
|
||||
|
||||
assert result is not None
|
||||
assert result.status == StatusValues.PENDING
|
||||
|
||||
@@ -25,8 +25,11 @@ from superset.utils.screenshot_utils import (
|
||||
combine_screenshot_tiles,
|
||||
resolve_screenshot_task_budget_seconds,
|
||||
SCREENSHOT_TASK_BUDGET_MAX_MARGIN_SECONDS,
|
||||
ScreenshotTaskBudgetExceededError,
|
||||
SCROLL_SETTLE_TIMEOUT_MS,
|
||||
take_tiled_screenshot,
|
||||
TILED_SCREENSHOT_TOTAL_WAIT_BUDGET_SECONDS,
|
||||
TiledScreenshotBudgetExceededError,
|
||||
)
|
||||
|
||||
|
||||
@@ -453,12 +456,17 @@ class TestTakeTiledScreenshot:
|
||||
assert warning_args[2] == 1 # count of unready chart containers
|
||||
assert warning_args[3] == 1 # tile index
|
||||
assert warning_args[4] == 3 # total tiles
|
||||
assert warning_args[5] == 30 # load_wait
|
||||
assert warning_args[6] == "" # no log_context passed
|
||||
assert warning_args[5] == 30 # tile_load_wait (uncapped: budget remains)
|
||||
assert warning_args[6] == 30 # requested load_wait
|
||||
assert isinstance(warning_args[7], float) # total elapsed vs budget
|
||||
assert warning_args[8] == 1440 # total budget (fixed fallback)
|
||||
assert warning_args[9] == 0 # tiles captured so far
|
||||
assert warning_args[10] == 3 # total tiles
|
||||
assert warning_args[11] == "" # no log_context passed
|
||||
# Diagnostic payload identifies chart id AND the state it's stuck in
|
||||
# (spinner mounted vs nothing mounted vs waiting-on-database) so a
|
||||
# slow query can be told apart from the virtualization race.
|
||||
assert warning_args[7] == [{"chartId": "42", "state": "waiting_on_database"}]
|
||||
assert warning_args[12] == [{"chartId": "42", "state": "waiting_on_database"}]
|
||||
|
||||
def test_timeout_warning_includes_log_context(self, mock_page):
|
||||
"""The log context (e.g. report execution id) is threaded through for
|
||||
@@ -484,7 +492,7 @@ class TestTakeTiledScreenshot:
|
||||
)
|
||||
|
||||
warning_args = mock_logger.warning.call_args[0]
|
||||
assert warning_args[6] == " [execution_id=abc-123]"
|
||||
assert warning_args[11] == " [execution_id=abc-123]"
|
||||
|
||||
def test_chart_holder_with_nothing_mounted_blocks_wait(self, mock_page):
|
||||
"""Regression test for the vacuous-pass race (PR #39895).
|
||||
@@ -646,3 +654,311 @@ class TestTakeTiledScreenshot:
|
||||
|
||||
sig = inspect.signature(take_tiled_screenshot)
|
||||
assert sig.parameters["animation_wait"].default == 0
|
||||
|
||||
|
||||
class TestTileWaitBudget:
|
||||
"""The tiled operation's cumulative per-tile waits are capped by one
|
||||
wall-clock budget derived from the running Celery task's own time limit
|
||||
(resolve_screenshot_task_budget_seconds), falling back to a fixed total
|
||||
ceiling outside Celery because per-tile waits accumulate."""
|
||||
|
||||
@pytest.fixture
|
||||
def mock_page(self):
|
||||
"""Create a mock Playwright page object for a 3-tile (5000px) dashboard."""
|
||||
page = MagicMock()
|
||||
element = MagicMock()
|
||||
page.locator.return_value = element
|
||||
page.evaluate.return_value = {
|
||||
"height": 5000,
|
||||
"top": 100,
|
||||
"left": 50,
|
||||
"width": 800,
|
||||
}
|
||||
page.screenshot.return_value = b"fake_screenshot_data"
|
||||
return page
|
||||
|
||||
class _FakeClock:
|
||||
"""Stateful monotonic() stand-in the test advances explicitly.
|
||||
|
||||
Robust to how many times the code under test samples the clock per
|
||||
tile (budget check, per-tile wait timing, animation budget) -- only
|
||||
explicit advances move time forward.
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.now = 0.0
|
||||
|
||||
def __call__(self) -> float:
|
||||
return self.now
|
||||
|
||||
def test_budget_error_is_task_budget_error_subclass(self):
|
||||
"""Callers can catch the whole budget-error family with the base
|
||||
ScreenshotTaskBudgetExceededError type."""
|
||||
assert issubclass(
|
||||
TiledScreenshotBudgetExceededError, ScreenshotTaskBudgetExceededError
|
||||
)
|
||||
|
||||
def test_per_tile_wait_shrinks_as_budget_depletes(self, mock_page, monkeypatch):
|
||||
"""Each tile's readiness-wait timeout is capped at the remaining budget."""
|
||||
monkeypatch.setattr(
|
||||
"superset.utils.screenshot_utils.TILED_SCREENSHOT_TOTAL_WAIT_BUDGET_SECONDS", # noqa: E501
|
||||
1000,
|
||||
)
|
||||
clock = self._FakeClock()
|
||||
# Simulate slow tiles: the readiness wait itself consumes wall time,
|
||||
# so each subsequent tile sees less remaining budget.
|
||||
wait_durations = iter([950, 40, 5])
|
||||
|
||||
def slow_wait(*args, **kwargs):
|
||||
clock.now += next(wait_durations)
|
||||
|
||||
mock_page.wait_for_function.side_effect = slow_wait
|
||||
|
||||
with patch("superset.utils.screenshot_utils.current_task", None):
|
||||
with patch("superset.utils.screenshot_utils.time.monotonic", new=clock):
|
||||
with patch("superset.utils.screenshot_utils.combine_screenshot_tiles"):
|
||||
result = take_tiled_screenshot(
|
||||
mock_page, "dashboard", tile_height=2000, load_wait=100
|
||||
)
|
||||
|
||||
assert result is not None
|
||||
timeouts = [
|
||||
call[1]["timeout"] for call in mock_page.wait_for_function.call_args_list
|
||||
]
|
||||
# remaining budget at each tile's wait: 1000, 50, 10 seconds
|
||||
# -> capped timeouts shrink
|
||||
assert timeouts == [100 * 1000, 50 * 1000, 10 * 1000]
|
||||
assert timeouts == sorted(timeouts, reverse=True)
|
||||
|
||||
def test_readiness_wait_uses_budget_recomputed_after_scroll_settle(
|
||||
self, mock_page, monkeypatch
|
||||
):
|
||||
"""The readiness-wait timeout must be capped using the budget
|
||||
recomputed *after* the scroll-settle sleep, not the stale value from
|
||||
before it -- otherwise each tile could overrun the total budget by up
|
||||
to one settle interval."""
|
||||
monkeypatch.setattr(
|
||||
"superset.utils.screenshot_utils.TILED_SCREENSHOT_TOTAL_WAIT_BUDGET_SECONDS", # noqa: E501
|
||||
1000,
|
||||
)
|
||||
# A single-tile dashboard to keep the scenario simple.
|
||||
mock_page.evaluate.return_value = {
|
||||
"height": 1000,
|
||||
"top": 100,
|
||||
"left": 50,
|
||||
"width": 800,
|
||||
}
|
||||
clock = self._FakeClock()
|
||||
# The scroll-settle sleep itself consumes 950s of wall-clock time,
|
||||
# leaving only 50s of the 1000s budget by the time the readiness
|
||||
# wait is capped.
|
||||
mock_page.wait_for_timeout.side_effect = lambda *args, **kwargs: setattr(
|
||||
clock, "now", clock.now + 950
|
||||
)
|
||||
|
||||
with patch("superset.utils.screenshot_utils.current_task", None):
|
||||
with patch("superset.utils.screenshot_utils.time.monotonic", new=clock):
|
||||
with patch("superset.utils.screenshot_utils.combine_screenshot_tiles"):
|
||||
take_tiled_screenshot(
|
||||
mock_page, "dashboard", tile_height=2000, load_wait=999
|
||||
)
|
||||
|
||||
timeout = mock_page.wait_for_function.call_args_list[0][1]["timeout"]
|
||||
# Must reflect the post-settle remaining budget (50s), not the
|
||||
# stale pre-settle value (1000s, which would have let load_wait's
|
||||
# full 999s through uncapped).
|
||||
assert timeout == 50 * 1000
|
||||
|
||||
def test_budget_exhausted_raises_and_stops_capturing(self, mock_page, monkeypatch):
|
||||
"""Exhausting the budget aborts cleanly instead of capturing unchecked."""
|
||||
monkeypatch.setattr(
|
||||
"superset.utils.screenshot_utils.TILED_SCREENSHOT_TOTAL_WAIT_BUDGET_SECONDS", # noqa: E501
|
||||
1000,
|
||||
)
|
||||
clock = self._FakeClock()
|
||||
# Tile 0's readiness wait consumes the whole budget; tile 1's budget
|
||||
# check then sees remaining <= 0 and raises before capturing.
|
||||
mock_page.wait_for_function.side_effect = lambda *args, **kwargs: setattr(
|
||||
clock, "now", 1000.0
|
||||
)
|
||||
|
||||
with patch("superset.utils.screenshot_utils.current_task", None):
|
||||
with patch("superset.utils.screenshot_utils.time.monotonic", new=clock):
|
||||
with patch(
|
||||
"superset.utils.screenshot_utils.combine_screenshot_tiles"
|
||||
) as mock_combine:
|
||||
with patch("superset.utils.screenshot_utils.logger") as mock_logger:
|
||||
with pytest.raises(TiledScreenshotBudgetExceededError):
|
||||
take_tiled_screenshot(
|
||||
mock_page, "dashboard", tile_height=2000, load_wait=100
|
||||
)
|
||||
|
||||
# Only the first tile was captured before the budget ran out.
|
||||
assert mock_page.screenshot.call_count == 1
|
||||
# Tiles were never combined -- the function raised before that point.
|
||||
mock_combine.assert_not_called()
|
||||
|
||||
# Budget exhaustion is a customer chart-loading issue, not a Superset
|
||||
# system fault, so it must log at WARNING (not ERROR) -- consistent
|
||||
# with the #38130/#38441 precedent for screenshot timeout logging.
|
||||
assert mock_logger.error.call_count == 0
|
||||
mock_logger.warning.assert_called_once()
|
||||
warning_args = mock_logger.warning.call_args[0]
|
||||
assert "budget exhausted" in warning_args[0]
|
||||
# tile index, tiles total, tiles captured, tiles total,
|
||||
# elapsed seconds, budget seconds, log-context suffix
|
||||
assert warning_args[1] == 2
|
||||
assert warning_args[2] == 3
|
||||
assert warning_args[3] == 1
|
||||
assert warning_args[4] == 3
|
||||
assert warning_args[5] == 1000
|
||||
assert warning_args[6] == 1000
|
||||
assert warning_args[7] == ""
|
||||
|
||||
def test_budget_exhausted_warning_includes_log_context(
|
||||
self, mock_page, monkeypatch
|
||||
):
|
||||
"""log_context (e.g. report execution id) is appended to the warning."""
|
||||
monkeypatch.setattr(
|
||||
"superset.utils.screenshot_utils.TILED_SCREENSHOT_TOTAL_WAIT_BUDGET_SECONDS", # noqa: E501
|
||||
1000,
|
||||
)
|
||||
clock = self._FakeClock()
|
||||
# Tile 0's readiness wait consumes the whole budget; tile 1's budget
|
||||
# check then sees remaining <= 0 and raises.
|
||||
mock_page.wait_for_function.side_effect = lambda *args, **kwargs: setattr(
|
||||
clock, "now", 1000.0
|
||||
)
|
||||
with patch("superset.utils.screenshot_utils.current_task", None):
|
||||
with patch("superset.utils.screenshot_utils.time.monotonic", new=clock):
|
||||
with patch("superset.utils.screenshot_utils.combine_screenshot_tiles"):
|
||||
with patch("superset.utils.screenshot_utils.logger") as mock_logger:
|
||||
with pytest.raises(TiledScreenshotBudgetExceededError):
|
||||
take_tiled_screenshot(
|
||||
mock_page,
|
||||
"dashboard",
|
||||
tile_height=2000,
|
||||
load_wait=100,
|
||||
log_context="execution_id=abc-123",
|
||||
)
|
||||
|
||||
warning_args = mock_logger.warning.call_args[0]
|
||||
assert warning_args[-1] == " [execution_id=abc-123]"
|
||||
|
||||
def test_budget_exhausted_before_first_tile_raises_without_capture(
|
||||
self, mock_page, monkeypatch
|
||||
):
|
||||
"""No budget floor: a budget already exhausted by setup (element
|
||||
lookup/dimension probing) raises before the first tile is captured,
|
||||
matching the non-tiled path's raise-before-capture semantics."""
|
||||
monkeypatch.setattr(
|
||||
"superset.utils.screenshot_utils.TILED_SCREENSHOT_TOTAL_WAIT_BUDGET_SECONDS", # noqa: E501
|
||||
1000,
|
||||
)
|
||||
clock = self._FakeClock()
|
||||
# The dashboard-dimension evaluate() itself consumes the whole budget.
|
||||
original_return = {"height": 5000, "top": 100, "left": 50, "width": 800}
|
||||
|
||||
def slow_evaluate(*args, **kwargs):
|
||||
clock.now = 1000.0
|
||||
return original_return
|
||||
|
||||
mock_page.evaluate.side_effect = slow_evaluate
|
||||
|
||||
with patch("superset.utils.screenshot_utils.current_task", None):
|
||||
with patch("superset.utils.screenshot_utils.time.monotonic", new=clock):
|
||||
with patch(
|
||||
"superset.utils.screenshot_utils.combine_screenshot_tiles"
|
||||
) as mock_combine:
|
||||
with pytest.raises(TiledScreenshotBudgetExceededError):
|
||||
take_tiled_screenshot(
|
||||
mock_page, "dashboard", tile_height=2000, load_wait=100
|
||||
)
|
||||
|
||||
mock_page.screenshot.assert_not_called()
|
||||
mock_combine.assert_not_called()
|
||||
|
||||
def test_no_celery_context_uses_fixed_total_fallback(self, mock_page):
|
||||
"""Outside Celery the helper returns None; the tiled path must fall
|
||||
back to the fixed total ceiling rather than running uncapped, because
|
||||
per-tile waits accumulate across tiles."""
|
||||
clock = self._FakeClock()
|
||||
with patch("superset.utils.screenshot_utils.current_task", None):
|
||||
with patch("superset.utils.screenshot_utils.time.monotonic", new=clock):
|
||||
with patch("superset.utils.screenshot_utils.combine_screenshot_tiles"):
|
||||
take_tiled_screenshot(
|
||||
mock_page,
|
||||
"dashboard",
|
||||
tile_height=2000,
|
||||
load_wait=10_000, # deliberately above the fallback
|
||||
)
|
||||
|
||||
first_timeout = mock_page.wait_for_function.call_args_list[0][1]["timeout"]
|
||||
assert first_timeout == TILED_SCREENSHOT_TOTAL_WAIT_BUDGET_SECONDS * 1000
|
||||
|
||||
def test_derived_task_budget_caps_tile_wait(self, mock_page):
|
||||
"""Inside Celery, the tiled path caps waits using the same
|
||||
task-derived budget as the non-tiled path (helper reuse, #42427)."""
|
||||
task = MagicMock()
|
||||
task.request.timelimit = (120, None) # (hard, soft): 120s hard limit
|
||||
|
||||
clock = self._FakeClock()
|
||||
with patch("superset.utils.screenshot_utils.current_task", task):
|
||||
with patch("superset.utils.screenshot_utils.time.monotonic", new=clock):
|
||||
with patch("superset.utils.screenshot_utils.combine_screenshot_tiles"):
|
||||
take_tiled_screenshot(
|
||||
mock_page, "dashboard", tile_height=2000, load_wait=200
|
||||
)
|
||||
|
||||
# margin = min(300, 120 * 0.2) = 24; budget = 120 - 24 = 96
|
||||
first_timeout = mock_page.wait_for_function.call_args_list[0][1]["timeout"]
|
||||
assert first_timeout == 96 * 1000
|
||||
assert first_timeout < 200 * 1000
|
||||
|
||||
def test_fast_dashboard_matches_default_behavior(self, mock_page):
|
||||
"""Well under budget, waits are not capped and behavior is unchanged."""
|
||||
with patch("superset.utils.screenshot_utils.current_task", None):
|
||||
with patch("superset.utils.screenshot_utils.combine_screenshot_tiles"):
|
||||
result = take_tiled_screenshot(
|
||||
mock_page,
|
||||
"dashboard",
|
||||
tile_height=2000,
|
||||
load_wait=30,
|
||||
animation_wait=5,
|
||||
)
|
||||
|
||||
assert result is not None
|
||||
assert mock_page.screenshot.call_count == 3
|
||||
|
||||
for call in mock_page.wait_for_function.call_args_list:
|
||||
assert call[1]["timeout"] == 30 * 1000
|
||||
|
||||
animation_calls = [
|
||||
call
|
||||
for call in mock_page.wait_for_timeout.call_args_list
|
||||
if call[0][0] == 5 * 1000
|
||||
]
|
||||
assert len(animation_calls) == 3
|
||||
|
||||
def test_per_tile_timing_debug_line_logged(self, mock_page):
|
||||
"""Each tile logs a DEBUG timing breakdown (readiness wait, animation
|
||||
wait) so slow dashboards can be profiled from logs alone."""
|
||||
with patch("superset.utils.screenshot_utils.current_task", None):
|
||||
with patch("superset.utils.screenshot_utils.logger") as mock_logger:
|
||||
with patch("superset.utils.screenshot_utils.combine_screenshot_tiles"):
|
||||
take_tiled_screenshot(
|
||||
mock_page,
|
||||
"dashboard",
|
||||
tile_height=2000,
|
||||
log_context="cache_key=xyz",
|
||||
)
|
||||
|
||||
timing_calls = [
|
||||
call for call in mock_logger.debug.call_args_list if "timing" in call[0][0]
|
||||
]
|
||||
assert len(timing_calls) == 3
|
||||
for i, call in enumerate(timing_calls):
|
||||
args = call[0]
|
||||
assert args[1] == i + 1 # tile index
|
||||
assert args[2] == 3 # total tiles
|
||||
assert args[-1] == " [cache_key=xyz]"
|
||||
|
||||
Reference in New Issue
Block a user