address comment

pinterest · Jul 8, 2024 · 4558e08 · 4558e08
1 parent 208f6d0
commit 4558e08
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 20 deletions.
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 { sampleRate, setSampleRate, getSampleRate, getSamplingTables };
 };
 
 const useTemplatedVariables = (dispatch: Dispatch, environmentId: number) => {
@@ -413,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 { sampleRate, setSampleRate, getSampleRate, getSamplingTables } =
+        useTableSampleRate(dispatch, environmentId, samplingTables);
     const [resultsCollapsed, setResultsCollapsed] = useState(false);
 
     const { searchAndReplaceProps, searchAndReplaceRef } =
@@ -499,18 +512,6 @@ const QueryComposer: React.FC = () => {
 
     const triggerSurvey = useSurveyTrigger();
 
-    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]);
-
     const getQueryExecutionMetadata = useCallback(() => {
         const metadata = {};
 

diff --git a/querybook/webapp/components/QueryExecution/SamplingToolTip.tsx b/querybook/webapp/components/QueryExecution/SamplingToolTip.tsx
@@ -32,7 +32,7 @@ export const SamplingTooltip: React.FC<SamplingTooltipProps> = ({
             enabled &&
             hasSamplingTables &&
             sampleRate <= 0 &&
-            status <= QueryExecutionStatus.RUNNING
+            status === QueryExecutionStatus.RUNNING
         ) {
             const timer = setTimeout(() => {
                 setShowSamplingTip(true);