ui: Disable tip for sampled queries (#1471)

* ui: disable tip for sampled queries
pinterest · Jul 8, 2024 · dd50b5d · dd50b5d
1 parent 1092268
commit dd50b5d
Show file tree

Hide file tree

Showing 6 changed files with 60 additions and 26 deletions.
diff --git a/querybook/webapp/components/DataDocQueryCell/DataDocQueryCell.tsx b/querybook/webapp/components/DataDocQueryCell/DataDocQueryCell.tsx
@@ -993,6 +993,7 @@ class DataDocQueryCellComponent extends React.PureComponent<IProps, IState> {
                 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
@@ -27,6 +27,7 @@ interface IProps {
 
     onSamplingInfoClick?: () => void;
     hasSamplingTables?: boolean;
+    sampleRate: number;
 }
 
 export const DataDocQueryExecutions: React.FunctionComponent<IProps> =
@@ -38,6 +39,7 @@ export const DataDocQueryExecutions: React.FunctionComponent<IProps> =
             isQueryCollapsed,
             onSamplingInfoClick,
             hasSamplingTables,
+            sampleRate,
         }) => {
             const { cellIdToExecutionId, onQueryCellSelectExecution } =
                 useContext(DataDocContext);
@@ -155,6 +157,7 @@ export const DataDocQueryExecutions: React.FunctionComponent<IProps> =
                     changeCellContext={changeCellContext}
                     onSamplingInfoClick={onSamplingInfoClick}
                     hasSamplingTables={hasSamplingTables}
+                    sampleRate={sampleRate}
                 />
             );
 

diff --git 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<string, any>
+) => {
     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()}
                     />
                 </div>
             </Resizable>
@@ -692,7 +711,7 @@ const QueryComposer: React.FC = () => {
         <DataDocTableSamplingInfo
             query={getCurrentSelectedQuery()}
             language={engine.language}
-            samplingTables={samplingTables}
+            samplingTables={getSamplingTables()}
             onHide={() => 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
@@ -17,12 +17,14 @@ interface IProps {
     id: number;
     onSamplingInfoClick?: () => void;
     hasSamplingTables?: boolean;
+    sampleRate: number;
 }
 
 export const QueryComposerExecution: React.FunctionComponent<IProps> = ({
     id,
     onSamplingInfoClick,
     hasSamplingTables,
+    sampleRate,
 }) => {
     const execution = useSelector(
         (state: IStoreState) => state.queryExecutions.queryExecutionById[id]
@@ -61,6 +63,7 @@ export const QueryComposerExecution: React.FunctionComponent<IProps> = ({
                 id={id}
                 onSamplingInfoClick={onSamplingInfoClick}
                 hasSamplingTables={hasSamplingTables}
+                sampleRate={sampleRate}
             />
         </div>
     );

diff --git 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<IProps> = ({
 
     onSamplingInfoClick,
     hasSamplingTables,
+    sampleRate,
 }) => {
     const isEditable = useSelector((state: IStoreState) =>
         canCurrentUserEditSelector(state, docId)
@@ -144,6 +146,7 @@ export const QueryExecution: React.FC<IProps> = ({
                 queryExecution={queryExecution}
                 onSamplingInfoClick={onSamplingInfoClick}
                 hasSamplingTables={hasSamplingTables}
+                sampleRate={sampleRate}
             />
         );
 

diff --git 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<SamplingTooltipProps> = ({
     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<SamplingTooltipProps> = ({
                 setShowSamplingTip(false);
             };
         }
-    }, [status, id, queryCanBeSampled, delay]);
+    }, [id, delay, enabled, hasSamplingTables, sampleRate, status]);
 
     const samplingTipDOM = showSamplingTip && (
         <Message size="small" type="warning" className="SamplingToolTip">