diff --git a/querybook/webapp/components/DataDocQueryCell/DataDocQueryCell.tsx b/querybook/webapp/components/DataDocQueryCell/DataDocQueryCell.tsx index 649056de5..98b1dea72 100644 --- a/querybook/webapp/components/DataDocQueryCell/DataDocQueryCell.tsx +++ b/querybook/webapp/components/DataDocQueryCell/DataDocQueryCell.tsx @@ -993,6 +993,7 @@ class DataDocQueryCellComponent extends React.PureComponent { changeCellContext={isEditable ? this.handleChange : null} onSamplingInfoClick={this.toggleShowTableSamplingInfoModal} hasSamplingTables={this.hasSamplingTables} + sampleRate={this.sampleRate} /> ); } diff --git a/querybook/webapp/components/DataDocQueryExecutions/DataDocQueryExecutions.tsx b/querybook/webapp/components/DataDocQueryExecutions/DataDocQueryExecutions.tsx index b7fc780c4..ad23efc07 100644 --- a/querybook/webapp/components/DataDocQueryExecutions/DataDocQueryExecutions.tsx +++ b/querybook/webapp/components/DataDocQueryExecutions/DataDocQueryExecutions.tsx @@ -27,6 +27,7 @@ interface IProps { onSamplingInfoClick?: () => void; hasSamplingTables?: boolean; + sampleRate: number; } export const DataDocQueryExecutions: React.FunctionComponent = @@ -38,6 +39,7 @@ export const DataDocQueryExecutions: React.FunctionComponent = isQueryCollapsed, onSamplingInfoClick, hasSamplingTables, + sampleRate, }) => { const { cellIdToExecutionId, onQueryCellSelectExecution } = useContext(DataDocContext); @@ -155,6 +157,7 @@ export const DataDocQueryExecutions: React.FunctionComponent = changeCellContext={changeCellContext} onSamplingInfoClick={onSamplingInfoClick} hasSamplingTables={hasSamplingTables} + sampleRate={sampleRate} /> ); diff --git a/querybook/webapp/components/QueryComposer/QueryComposer.tsx b/querybook/webapp/components/QueryComposer/QueryComposer.tsx index f16b2638a..fdc2ef971 100644 --- a/querybook/webapp/components/QueryComposer/QueryComposer.tsx +++ b/querybook/webapp/components/QueryComposer/QueryComposer.tsx @@ -163,7 +163,11 @@ const useRowLimit = (dispatch: Dispatch, environmentId: number) => { return { rowLimit, setRowLimit }; }; -const useTableSampleRate = (dispatch: Dispatch, environmentId: number) => { +const useTableSampleRate = ( + dispatch: Dispatch, + environmentId: number, + samplingTables: Record +) => { const sampleRate = useSelector( (state: IStoreState) => state.adhocQuery[environmentId]?.sampleRate ?? 0 ); @@ -178,7 +182,19 @@ const useTableSampleRate = (dispatch: Dispatch, environmentId: number) => { [dispatch, environmentId] ); - return { sampleRate, setSampleRate }; + const getSampleRate = useCallback( + () => (Object.keys(samplingTables).length > 0 ? sampleRate : -1), + [sampleRate, samplingTables] + ); + + const getSamplingTables = useCallback(() => { + Object.keys(samplingTables).forEach((tableName) => { + samplingTables[tableName].sample_rate = getSampleRate(); + }); + return samplingTables; + }, [getSampleRate, samplingTables]); + + return { setSampleRate, getSampleRate, getSamplingTables }; }; const useTemplatedVariables = (dispatch: Dispatch, environmentId: number) => { @@ -395,14 +411,6 @@ function useTranspileQuery( }; } -function getQueryExecutionMetadata(sampleRate: number) { - const metadata = {}; - if (sampleRate > 0) { - metadata['sample_rate'] = sampleRate; - } - return Object.keys(metadata).length === 0 ? null : metadata; -} - const QueryComposer: React.FC = () => { useTrackView(ComponentType.ADHOC_QUERY); useBrowserTitle('Adhoc Query'); @@ -421,12 +429,9 @@ const QueryComposer: React.FC = () => { environmentId ); const { rowLimit, setRowLimit } = useRowLimit(dispatch, environmentId); - const { sampleRate, setSampleRate } = useTableSampleRate( - dispatch, - environmentId - ); const [samplingTables, setSamplingTables] = useState({}); - + const { setSampleRate, getSampleRate, getSamplingTables } = + useTableSampleRate(dispatch, environmentId, samplingTables); const [resultsCollapsed, setResultsCollapsed] = useState(false); const { searchAndReplaceProps, searchAndReplaceRef } = @@ -507,9 +512,22 @@ const QueryComposer: React.FC = () => { const triggerSurvey = useSurveyTrigger(); - const queryExecutionMetadata = getQueryExecutionMetadata(sampleRate); + const getQueryExecutionMetadata = useCallback(() => { + const metadata = {}; + + const sampleRate = getSampleRate(); + if (sampleRate > 0) { + metadata['sample_rate'] = sampleRate; + } + + return Object.keys(metadata).length === 0 ? null : metadata; + }, [getSampleRate]); const handleRunQuery = React.useCallback(async () => { + const sampleRate = getSampleRate(); + const samplingTables = getSamplingTables(); + const queryExecutionMetadata = getQueryExecutionMetadata(); + trackClick({ component: ComponentType.ADHOC_QUERY, element: ElementType.RUN_QUERY_BUTTON, @@ -520,6 +538,7 @@ const QueryComposer: React.FC = () => { }); // Throttle to prevent double run await sleep(250); + const transformedQuery = await transformQuery( getCurrentSelectedQuery(), engine.language, @@ -553,16 +572,16 @@ const QueryComposer: React.FC = () => { setResultsCollapsed(false); } }, [ + getSampleRate, + getSamplingTables, + getQueryExecutionMetadata, hasLintErrors, - sampleRate, getCurrentSelectedQuery, engine, templatedVariables, rowLimit, - samplingTables, triggerSurvey, dispatch, - queryExecutionMetadata, setExecutionId, ]); @@ -596,7 +615,6 @@ const QueryComposer: React.FC = () => { if (table?.custom_properties?.sampling) { samplingTables[tableName] = { sampled_table: table.custom_properties?.sampled_table, - sample_rate: sampleRate, }; } }); @@ -666,6 +684,7 @@ const QueryComposer: React.FC = () => { hasSamplingTables={ Object.keys(samplingTables).length > 0 } + sampleRate={getSampleRate()} /> @@ -692,7 +711,7 @@ const QueryComposer: React.FC = () => { setShowTableSamplingInfoModal(false)} /> ); @@ -732,7 +751,7 @@ const QueryComposer: React.FC = () => { rowLimit={rowLimit} onRowLimitChange={setRowLimit} hasSamplingTables={Object.keys(samplingTables).length > 0} - sampleRate={sampleRate} + sampleRate={getSampleRate()} onSampleRateChange={setSampleRate} onTableSamplingInfoClick={() => setShowTableSamplingInfoModal(true) diff --git a/querybook/webapp/components/QueryComposer/QueryComposerExecution.tsx b/querybook/webapp/components/QueryComposer/QueryComposerExecution.tsx index f553ebe16..1a604cb42 100644 --- a/querybook/webapp/components/QueryComposer/QueryComposerExecution.tsx +++ b/querybook/webapp/components/QueryComposer/QueryComposerExecution.tsx @@ -17,12 +17,14 @@ interface IProps { id: number; onSamplingInfoClick?: () => void; hasSamplingTables?: boolean; + sampleRate: number; } export const QueryComposerExecution: React.FunctionComponent = ({ id, onSamplingInfoClick, hasSamplingTables, + sampleRate, }) => { const execution = useSelector( (state: IStoreState) => state.queryExecutions.queryExecutionById[id] @@ -61,6 +63,7 @@ export const QueryComposerExecution: React.FunctionComponent = ({ id={id} onSamplingInfoClick={onSamplingInfoClick} hasSamplingTables={hasSamplingTables} + sampleRate={sampleRate} /> ); diff --git a/querybook/webapp/components/QueryExecution/QueryExecution.tsx b/querybook/webapp/components/QueryExecution/QueryExecution.tsx index 470a513a6..758087c80 100644 --- a/querybook/webapp/components/QueryExecution/QueryExecution.tsx +++ b/querybook/webapp/components/QueryExecution/QueryExecution.tsx @@ -33,6 +33,7 @@ interface IProps { onSamplingInfoClick?: () => void; hasSamplingTables?: boolean; + sampleRate?: number; } function useQueryExecutionReduxState(queryId: number) { @@ -92,6 +93,7 @@ export const QueryExecution: React.FC = ({ onSamplingInfoClick, hasSamplingTables, + sampleRate, }) => { const isEditable = useSelector((state: IStoreState) => canCurrentUserEditSelector(state, docId) @@ -144,6 +146,7 @@ export const QueryExecution: React.FC = ({ queryExecution={queryExecution} onSamplingInfoClick={onSamplingInfoClick} hasSamplingTables={hasSamplingTables} + sampleRate={sampleRate} /> ); diff --git a/querybook/webapp/components/QueryExecution/SamplingToolTip.tsx b/querybook/webapp/components/QueryExecution/SamplingToolTip.tsx index 6c23245a4..fadaee97c 100644 --- a/querybook/webapp/components/QueryExecution/SamplingToolTip.tsx +++ b/querybook/webapp/components/QueryExecution/SamplingToolTip.tsx @@ -12,23 +12,28 @@ interface SamplingTooltipProps { queryExecution: IQueryExecution; onSamplingInfoClick?: () => void; hasSamplingTables?: boolean; + sampleRate: number; } export const SamplingTooltip: React.FC = ({ queryExecution: { status, id }, onSamplingInfoClick, hasSamplingTables, + sampleRate, }) => { const { enabled, sampling_tool_tip_delay: delay } = PublicConfig.table_sampling; - const queryCanBeSampled = enabled && hasSamplingTables; - const [showSamplingTip, setShowSamplingTip] = useState(false); // If query run longer than 10 seconds, we show a suggestion to users they can sample their query useEffect(() => { - if (queryCanBeSampled && status <= QueryExecutionStatus.RUNNING) { + if ( + enabled && + hasSamplingTables && + sampleRate <= 0 && + status === QueryExecutionStatus.RUNNING + ) { const timer = setTimeout(() => { setShowSamplingTip(true); }, delay); @@ -38,7 +43,7 @@ export const SamplingTooltip: React.FC = ({ setShowSamplingTip(false); }; } - }, [status, id, queryCanBeSampled, delay]); + }, [id, delay, enabled, hasSamplingTables, sampleRate, status]); const samplingTipDOM = showSamplingTip && (