ui: Add Query Nagging for Sampling (#1468)

* ui: Add Query Nagging for Sampling
pinterest · Jun 27, 2024 · 11caa51 · 11caa51
1 parent 7c210ce
commit 11caa51
Show file tree

Hide file tree

Showing 10 changed files with 304 additions and 143 deletions.
diff --git a/querybook/config/querybook_public_config.yaml b/querybook/config/querybook_public_config.yaml
@@ -36,3 +36,4 @@ table_sampling:
         - 0.1
     default_sample_rate: 0 # 0 means no sampling
     sample_user_guide_link: ''
+    sampling_tool_tip_delay: 10000 # delay duration (ms) of sampling tool tip
diff --git a/querybook/webapp/components/DataDocQueryCell/DataDocQueryCell.tsx b/querybook/webapp/components/DataDocQueryCell/DataDocQueryCell.tsx
@@ -459,6 +459,15 @@ class DataDocQueryCellComponent extends React.PureComponent<IProps, IState> {
         );
     }
 
+    @bind
+    public getQueryExecutionMetadata() {
+        const metadata = {};
+        if (this.sampleRate > 0) {
+            metadata['sample_rate'] = this.sampleRate;
+        }
+        return Object.keys(metadata).length === 0 ? null : metadata;
+    }
+
     @bind
     public async onRunButtonClick() {
         trackClick({
@@ -469,6 +478,7 @@ class DataDocQueryCellComponent extends React.PureComponent<IProps, IState> {
                 sampleRate: this.sampleRate,
             },
         });
+
         return runQuery(
             await this.getTransformedQuery(),
             this.engineId,
@@ -478,9 +488,7 @@ class DataDocQueryCellComponent extends React.PureComponent<IProps, IState> {
                         query,
                         engineId,
                         this.props.cellId,
-                        this.sampleRate > 0
-                            ? { sample_rate: this.sampleRate }
-                            : null
+                        this.getQueryExecutionMetadata()
                     )
                 ).id;
 
@@ -548,11 +556,14 @@ class DataDocQueryCellComponent extends React.PureComponent<IProps, IState> {
         );
 
         if (renderedQuery) {
+            const executionMetadata =
+                this.sampleRate > 0 ? { sample_rate: this.sampleRate } : null;
+
             return this.props.createQueryExecution(
                 renderedQuery,
                 this.engineId,
                 this.props.cellId,
-                this.sampleRate > 0 ? { sample_rate: this.sampleRate } : null
+                executionMetadata
             );
         }
     }
@@ -980,6 +991,8 @@ class DataDocQueryCellComponent extends React.PureComponent<IProps, IState> {
                 cellId={cellId}
                 isQueryCollapsed={this.queryCollapsed}
                 changeCellContext={isEditable ? this.handleChange : null}
+                onSamplingInfoClick={this.toggleShowTableSamplingInfoModal}
+                hasSamplingTables={this.hasSamplingTables}
             />
         );
     }

diff --git a/querybook/webapp/components/DataDocQueryExecutions/DataDocQueryExecutions.tsx b/querybook/webapp/components/DataDocQueryExecutions/DataDocQueryExecutions.tsx
@@ -24,136 +24,160 @@ interface IProps {
     isQueryCollapsed: boolean;
 
     changeCellContext?: (context: string, run?: boolean) => void;
+
+    onSamplingInfoClick?: () => void;
+    hasSamplingTables?: boolean;
 }
 
 export const DataDocQueryExecutions: React.FunctionComponent<IProps> =
-    React.memo(({ cellId, docId, changeCellContext, isQueryCollapsed }) => {
-        const { cellIdToExecutionId, onQueryCellSelectExecution } =
-            useContext(DataDocContext);
-
-        const queryExecutions =
-            useMakeSelector(
-                queryExecutionsSelectors.makeQueryExecutionsSelector,
-                cellId
-            ) ?? [];
-        const selectedExecutionIndex = useMemo(() => {
-            if (!(cellId in cellIdToExecutionId)) {
-                return 0;
-            }
-
-            const foundIndex = queryExecutions.findIndex(
-                (qe) => qe.id === cellIdToExecutionId[cellId]
-            );
-            // If not found, then return 0
-            return Math.max(0, foundIndex);
-        }, [cellIdToExecutionId[cellId], queryExecutions]);
-
-        const selectExecutionIndex = useCallback(
-            (index: number) => {
-                if (queryExecutions.length > index) {
-                    onQueryCellSelectExecution(
-                        cellId,
-                        queryExecutions[index].id
-                    );
+    React.memo(
+        ({
+            cellId,
+            docId,
+            changeCellContext,
+            isQueryCollapsed,
+            onSamplingInfoClick,
+            hasSamplingTables,
+        }) => {
+            const { cellIdToExecutionId, onQueryCellSelectExecution } =
+                useContext(DataDocContext);
+
+            const queryExecutions =
+                useMakeSelector(
+                    queryExecutionsSelectors.makeQueryExecutionsSelector,
+                    cellId
+                ) ?? [];
+            const selectedExecutionIndex = useMemo(() => {
+                if (!(cellId in cellIdToExecutionId)) {
+                    return 0;
                 }
-            },
-            [queryExecutions]
-        );
-
-        const dispatch = useDispatch();
-        useEffect(() => {
-            dispatch(queryExecutionsActions.fetchQueryExecutionsByCell(cellId));
-        }, [cellId]);
-
-        useEffect(() => {
-            selectExecutionIndex(0);
-        }, [queryExecutions.length]);
-
-        const [usedQueryParam, setUsedQueryParam] = useState(false);
-        useEffect(() => {
-            if (!usedQueryParam && queryExecutions.length) {
-                const queryParam = getQueryString();
-                if ('executionId' in queryParam && queryParam['executionId']) {
-                    const executionId = Number(queryParam['executionId']);
-                    for (const [
-                        index,
-                        queryExecution,
-                    ] of queryExecutions.entries()) {
-                        if (queryExecution.id === executionId) {
-                            selectExecutionIndex(index);
-                        }
+
+                const foundIndex = queryExecutions.findIndex(
+                    (qe) => qe.id === cellIdToExecutionId[cellId]
+                );
+                // If not found, then return 0
+                return Math.max(0, foundIndex);
+            }, [cellIdToExecutionId[cellId], queryExecutions]);
+
+            const selectExecutionIndex = useCallback(
+                (index: number) => {
+                    if (queryExecutions.length > index) {
+                        onQueryCellSelectExecution(
+                            cellId,
+                            queryExecutions[index].id
+                        );
                     }
-                }
-                setUsedQueryParam(true);
-            }
-        }, [queryExecutions, usedQueryParam]);
-
-        const handleQueryExecutionSelected = useCallback(
-            (qid: number) => {
-                const newIndex = queryExecutions.findIndex(
-                    ({ id }) => qid === id
+                },
+                [queryExecutions]
+            );
+
+            const dispatch = useDispatch();
+            useEffect(() => {
+                dispatch(
+                    queryExecutionsActions.fetchQueryExecutionsByCell(cellId)
                 );
-                if (newIndex >= 0) {
-                    selectExecutionIndex(newIndex);
+            }, [cellId]);
+
+            useEffect(() => {
+                selectExecutionIndex(0);
+            }, [queryExecutions.length]);
+
+            const [usedQueryParam, setUsedQueryParam] = useState(false);
+            useEffect(() => {
+                if (!usedQueryParam && queryExecutions.length) {
+                    const queryParam = getQueryString();
+                    if (
+                        'executionId' in queryParam &&
+                        queryParam['executionId']
+                    ) {
+                        const executionId = Number(queryParam['executionId']);
+                        for (const [
+                            index,
+                            queryExecution,
+                        ] of queryExecutions.entries()) {
+                            if (queryExecution.id === executionId) {
+                                selectExecutionIndex(index);
+                            }
+                        }
+                    }
+                    setUsedQueryParam(true);
                 }
-            },
-            [queryExecutions]
-        );
+            }, [queryExecutions, usedQueryParam]);
 
-        const generateExecutionsPickerDOM = () => {
-            const currentExecution = queryExecutions?.[selectedExecutionIndex];
-            if (!currentExecution) {
-                return null;
-            }
+            const handleQueryExecutionSelected = useCallback(
+                (qid: number) => {
+                    const newIndex = queryExecutions.findIndex(
+                        ({ id }) => qid === id
+                    );
+                    if (newIndex >= 0) {
+                        selectExecutionIndex(newIndex);
+                    }
+                },
+                [queryExecutions]
+            );
 
-            return (
-                <div className="execution-selector-section horizontal-space-between">
-                    <div className="flex-row">
-                        <QueryExecutionPicker
-                            queryExecutionId={currentExecution?.id}
-                            onSelection={handleQueryExecutionSelected}
-                            queryExecutions={queryExecutions}
-                        />
-                        <div className="ml8">
-                            <QueryExecutionDuration
+            const generateExecutionsPickerDOM = () => {
+                const currentExecution =
+                    queryExecutions?.[selectedExecutionIndex];
+                if (!currentExecution) {
+                    return null;
+                }
+
+                return (
+                    <div className="execution-selector-section horizontal-space-between">
+                        <div className="flex-row">
+                            <QueryExecutionPicker
+                                queryExecutionId={currentExecution?.id}
+                                onSelection={handleQueryExecutionSelected}
+                                queryExecutions={queryExecutions}
+                            />
+                            <div className="ml8">
+                                <QueryExecutionDuration
+                                    queryExecution={currentExecution}
+                                />
+                            </div>
+                        </div>
+                        <div className="execution-selector-bottom flex-row">
+                            <QueryExecutionBar
                                 queryExecution={currentExecution}
                             />
                         </div>
                     </div>
-                    <div className="execution-selector-bottom flex-row">
-                        <QueryExecutionBar queryExecution={currentExecution} />
-                    </div>
-                </div>
+                );
+            };
+
+            const selectedExecution = queryExecutions[selectedExecutionIndex];
+            const queryExecutionDOM = selectedExecution && (
+                <QueryExecution
+                    id={selectedExecution.id}
+                    docId={docId}
+                    key={selectedExecution.id}
+                    changeCellContext={changeCellContext}
+                    onSamplingInfoClick={onSamplingInfoClick}
+                    hasSamplingTables={hasSamplingTables}
+                />
             );
-        };
-        const selectedExecution = queryExecutions[selectedExecutionIndex];
-        const queryExecutionDOM = selectedExecution && (
-            <QueryExecution
-                id={selectedExecution.id}
-                docId={docId}
-                key={selectedExecution.id}
-                changeCellContext={changeCellContext}
-            />
-        );
-
-        const placeholderIfNoExecutionsDOM = isQueryCollapsed &&
-            !queryExecutionDOM && (
-                <div>
-                    Note: This query is collapsed, but it doesn't yet have an
-                    execution.
-                    <br />
-                    You can either run it, or edit it by clicking on the
-                    ellipsis button on the top right and then on "Show Query".
+
+            const placeholderIfNoExecutionsDOM = isQueryCollapsed &&
+                !queryExecutionDOM && (
+                    <div>
+                        Note: This query is collapsed, but it doesn't yet have
+                        an execution.
+                        <br />
+                        You can either run it, or edit it by clicking on the
+                        ellipsis button on the top right and then on "Show
+                        Query".
+                    </div>
+                );
+
+            return (
+                <div className="DataDocQueryExecutions">
+                    <StyledText size="xsmall">
+                        {generateExecutionsPickerDOM()}
+                        {queryExecutionDOM}
+                        {placeholderIfNoExecutionsDOM}
+                    </StyledText>
                 </div>
             );
-
-        return (
-            <div className="DataDocQueryExecutions">
-                <StyledText size="xsmall">
-                    {generateExecutionsPickerDOM()}
-                    {queryExecutionDOM}
-                    {placeholderIfNoExecutionsDOM}
-                </StyledText>
-            </div>
-        );
-    });
+        }
+    );