feat: add table sample rate to user settings

pinterest · Nov 21, 2024 · 5c3c505 · 5c3c505
1 parent 7dd47c0
commit 5c3c505
Show file tree

Hide file tree

Showing 14 changed files with 117 additions and 65 deletions.
diff --git a/package.json b/package.json
@@ -1,6 +1,6 @@
 {
     "name": "querybook",
-    "version": "3.37.0",
+    "version": "3.37.1",
     "description": "A Big Data Webapp",
     "private": true,
     "scripts": {

diff --git a/querybook/config/user_setting.yaml b/querybook/config/user_setting.yaml
@@ -50,6 +50,12 @@ show_full_view:
         - disabled
     helper: Instead of modal, show full view when opening table/execution/snippet
 
+table_sample_rate:
+    default: ''
+    tab: general
+    options: []
+    helper: The sample rate will be pre-selected when a table supports sampling.
+
 editor_font_size:
     default: medium
     tab: editor

diff --git a/querybook/webapp/components/DataDocQueryCell/DataDocQueryCell.tsx b/querybook/webapp/components/DataDocQueryCell/DataDocQueryCell.tsx
@@ -24,7 +24,6 @@ import { QuerySnippetInsertionModal } from 'components/QuerySnippetInsertionModa
 import { TemplatedQueryView } from 'components/TemplateQueryView/TemplatedQueryView';
 import { TranspileQueryModal } from 'components/TranspileQueryModal/TranspileQueryModal';
 import { UDFForm } from 'components/UDFForm/UDFForm';
-import PublicConfig from 'config/querybook_public_config.yaml';
 import { ComponentType, ElementType } from 'const/analytics';
 import {
     IDataQueryCellMeta,
@@ -37,6 +36,7 @@ import { SurveySurfaceType } from 'const/survey';
 import { triggerSurvey } from 'hooks/ui/useSurveyTrigger';
 import { trackClick } from 'lib/analytics';
 import CodeMirror from 'lib/codemirror';
+import { isAIFeatureEnabled } from 'lib/public-config';
 import { getQueryAsExplain } from 'lib/sql-helper/sql-lexer';
 import { DEFAULT_ROW_LIMIT } from 'lib/sql-helper/sql-limiter';
 import { getPossibleTranspilers } from 'lib/templated-query/transpile';
@@ -65,8 +65,6 @@ import { ErrorQueryCell } from './ErrorQueryCell';
 
 import './DataDocQueryCell.scss';
 
-const AIAssistantConfig = PublicConfig.ai_assistant;
-
 const ON_CHANGE_DEBOUNCE_MS = 500;
 const FORMAT_QUERY_SHORTCUT = getShortcutSymbols(
     KeyMap.queryEditor.formatQuery.key
@@ -221,8 +219,11 @@ class DataDocQueryCellComponent extends React.PureComponent<IProps, IState> {
     }
 
     public get sampleRate() {
-        // -1 for tables don't support sampling, 0 for default sample rate (which means disable sampling)
-        return this.hasSamplingTables ? this.state.meta.sample_rate ?? 0 : -1;
+        // -1 for tables don't support sampling
+        const sampleRate = this.hasSamplingTables
+            ? this.state.meta.sample_rate
+            : -1;
+        return sampleRate;
     }
 
     @decorate(memoizeOne)
@@ -809,7 +810,7 @@ class DataDocQueryCellComponent extends React.PureComponent<IProps, IState> {
                         {this.getAdditionalDropDownButtonDOM()}
                     </div>
                 </div>
-                {AIAssistantConfig.enabled && isEditable && (
+                {isAIFeatureEnabled() && isEditable && (
                     <AICommandBar
                         query={query}
                         queryEngine={queryEngineById[this.engineId]}

diff --git a/querybook/webapp/components/DataDocStatementExecution/StatementResult.tsx b/querybook/webapp/components/DataDocStatementExecution/StatementResult.tsx
@@ -6,13 +6,13 @@ import {
     IStatementResultTableHandles,
     StatementResultTable,
 } from 'components/StatementResultTable/StatementResultTable';
-import PublicConfig from 'config/querybook_public_config.yaml';
 import { IStatementExecution, IStatementResult } from 'const/queryExecution';
 import { StatementExecutionResultSizes } from 'const/queryResultLimit';
 import { MIN_COLUMN_TO_SHOW_FILTER } from 'const/uiConfig';
 import { useImmer } from 'hooks/useImmer';
 import { useResource } from 'hooks/useResource';
 import { useToggleState } from 'hooks/useToggleState';
+import { TABLE_SAMPLING_CONFIG } from 'lib/public-config';
 import { getSelectStatementLimit } from 'lib/sql-helper/sql-limiter';
 import { stopPropagation } from 'lib/utils/noop';
 import { formatNumber } from 'lib/utils/number';
@@ -287,7 +287,7 @@ const FetchInfo: React.FC<{
         }
 
         const sampleUserGuideLink =
-            PublicConfig.table_sampling?.sample_user_guide_link ?? '';
+            TABLE_SAMPLING_CONFIG.sample_user_guide_link ?? '';
         const sampleRateText = (
             <div className="flex-row ml4">
                 (Full Result,

diff --git a/querybook/webapp/components/DataDocTableSamplingInfo/DataDocTableSamplingInfo.tsx b/querybook/webapp/components/DataDocTableSamplingInfo/DataDocTableSamplingInfo.tsx
@@ -1,9 +1,9 @@
 import React from 'react';
 
 import { QueryComparison } from 'components/TranspileQueryModal/QueryComparison';
-import PublicConfig from 'config/querybook_public_config.yaml';
 import { ISamplingTables } from 'const/datadoc';
 import { useResource } from 'hooks/useResource';
+import { TABLE_SAMPLING_CONFIG } from 'lib/public-config';
 import { formatError } from 'lib/utils/error';
 import { QueryTransformResource } from 'resource/queryTransform';
 import { Link } from 'ui/Link/Link';
@@ -27,7 +27,7 @@ export const DataDocTableSamplingInfo: React.FC<IProps> = ({
     onHide,
 }) => {
     const sampleUserGuideLink =
-        PublicConfig.table_sampling?.sample_user_guide_link ?? '';
+        TABLE_SAMPLING_CONFIG.sample_user_guide_link ?? '';
 
     const {
         data: sampledQuery,

diff --git a/querybook/webapp/components/QueryCellTitle/QueryCellTitle.tsx b/querybook/webapp/components/QueryCellTitle/QueryCellTitle.tsx
@@ -1,17 +1,15 @@
 import React, { useCallback, useEffect, useState } from 'react';
 
-import PublicConfig from 'config/querybook_public_config.yaml';
 import { AICommandType } from 'const/aiAssistant';
 import { ComponentType, ElementType } from 'const/analytics';
 import { useAISocket } from 'hooks/useAISocket';
 import { trackClick } from 'lib/analytics';
+import { isAIFeatureEnabled } from 'lib/public-config';
 import { IconButton } from 'ui/Button/IconButton';
 import { ResizableTextArea } from 'ui/ResizableTextArea/ResizableTextArea';
 
 import './QueryCellTitle.scss';
 
-const AIAssistantConfig = PublicConfig.ai_assistant;
-
 interface IQueryCellTitleProps {
     cellId: number;
     value: string;
@@ -30,9 +28,7 @@ export const QueryCellTitle: React.FC<IQueryCellTitleProps> = ({
     forceSaveQuery,
 }) => {
     const titleGenerationEnabled =
-        AIAssistantConfig.enabled &&
-        AIAssistantConfig.query_title_generation.enabled &&
-        query;
+        isAIFeatureEnabled('query_title_generation') && query;
     const [title, setTitle] = useState<string>('');
 
     const socket = useAISocket(AICommandType.SQL_TITLE, ({ data }) => {

diff --git a/querybook/webapp/components/QueryComposer/QueryComposer.tsx b/querybook/webapp/components/QueryComposer/QueryComposer.tsx
@@ -167,7 +167,7 @@ const useTableSampleRate = (
     samplingTables: Record<string, any>
 ) => {
     const sampleRate = useSelector(
-        (state: IStoreState) => state.adhocQuery[environmentId]?.sampleRate ?? 0
+        (state: IStoreState) => state.adhocQuery[environmentId]?.sampleRate
     );
     const setSampleRate = useCallback(
         (newSampleRate: number) =>

diff --git a/querybook/webapp/components/QueryExecution/QueryError.tsx b/querybook/webapp/components/QueryExecution/QueryError.tsx
@@ -3,14 +3,14 @@ import { useDispatch, useSelector } from 'react-redux';
 
 import { AutoFixButton } from 'components/AIAssistant/AutoFixButton';
 import { ErrorSuggestion } from 'components/DataDocStatementExecution/ErrorSuggestion';
-import PublicConfig from 'config/querybook_public_config.yaml';
 import { IQueryEngine } from 'const/queryEngine';
 import {
     IQueryError,
     IQueryExecution,
     IStatementExecution,
     QueryExecutionErrorType,
 } from 'const/queryExecution';
+import { isAIFeatureEnabled } from 'lib/public-config';
 import {
     getQueryLinePosition,
     IToken,
@@ -31,8 +31,6 @@ import { ExecutedQueryCell } from './ExecutedQueryCell';
 
 import './QueryError.scss';
 
-const AIAssistantConfig = PublicConfig.ai_assistant;
-
 interface IProps {
     queryEngine: IQueryEngine;
     queryError: IQueryError;
@@ -186,15 +184,13 @@ export const QueryError: React.FunctionComponent<IProps> = ({
                 <Icon name="AlertOctagon" size={20} className="mr8" />
                 {errorTitle}
             </div>
-            {!readonly &&
-                AIAssistantConfig.enabled &&
-                AIAssistantConfig.query_auto_fix.enabled && (
-                    <AutoFixButton
-                        query={queryExecution.query}
-                        queryExecutionId={queryExecution.id}
-                        onUpdateQuery={changeCellContext}
-                    />
-                )}
+            {!readonly && isAIFeatureEnabled('query_auto_fix') && (
+                <AutoFixButton
+                    query={queryExecution.query}
+                    queryExecutionId={queryExecution.id}
+                    onUpdateQuery={changeCellContext}
+                />
+            )}
         </div>
     );
 

diff --git a/querybook/webapp/components/QueryExecution/SamplingToolTip.tsx b/querybook/webapp/components/QueryExecution/SamplingToolTip.tsx
@@ -1,8 +1,8 @@
 import React, { useEffect, useState } from 'react';
 
 import { SamplingInfoButton } from 'components/QueryRunButton/QueryRunButton';
-import PublicConfig from 'config/querybook_public_config.yaml';
 import { IQueryExecution, QueryExecutionStatus } from 'const/queryExecution';
+import { TABLE_SAMPLING_CONFIG } from 'lib/public-config';
 import { Message } from 'ui/Message/Message';
 import { AccentText } from 'ui/StyledText/StyledText';
 
@@ -21,8 +21,7 @@ export const SamplingTooltip: React.FC<SamplingTooltipProps> = ({
     hasSamplingTables,
     sampleRate,
 }) => {
-    const { enabled, sampling_tool_tip_delay: delay } =
-        PublicConfig.table_sampling;
+    const { enabled, sampling_tool_tip_delay: delay } = TABLE_SAMPLING_CONFIG;
 
     const [showSamplingTip, setShowSamplingTip] = useState(false);
 

diff --git a/querybook/webapp/components/QueryRunButton/QueryRunButton.tsx b/querybook/webapp/components/QueryRunButton/QueryRunButton.tsx
@@ -3,11 +3,14 @@ import * as React from 'react';
 import { useState } from 'react';
 import { useSelector } from 'react-redux';
 
-import PublicConfig from 'config/querybook_public_config.yaml';
 import { IQueryEngine, QueryEngineStatus } from 'const/queryEngine';
 import { queryEngineStatusToIconStatus } from 'const/queryStatusIcon';
 import { TooltipDirection } from 'const/tooltip';
 import { MIN_ENGINE_TO_SHOW_FILTER } from 'const/uiConfig';
+import {
+    getTableSamplingRateOptions,
+    TABLE_SAMPLING_CONFIG,
+} from 'lib/public-config';
 import {
     ALLOW_UNLIMITED_QUERY,
     DEFAULT_ROW_LIMIT,
@@ -17,6 +20,7 @@ import { getShortcutSymbols, KeyMap } from 'lib/utils/keyboard';
 import { stopPropagation } from 'lib/utils/noop';
 import { formatNumber } from 'lib/utils/number';
 import { queryEngineStatusByIdEnvSelector } from 'redux/queryEngine/selector';
+import { IStoreState } from 'redux/store/types';
 import { AsyncButton, IAsyncButtonHandles } from 'ui/AsyncButton/AsyncButton';
 import { IconButton } from 'ui/Button/IconButton';
 import { Dropdown } from 'ui/Dropdown/Dropdown';
@@ -287,29 +291,30 @@ const QueryLimitSelector: React.FC<{
     );
 };
 
-const TABLE_SAMPLING_CONFIG = PublicConfig.table_sampling ?? {
-    enabled: false,
-    sample_rates: [],
-    default_sample_rate: 0,
-};
-const sampleRateOptions = [{ label: 'none', value: 0 }].concat(
-    TABLE_SAMPLING_CONFIG.sample_rates.map((value) => ({
-        label: value + '%',
-        value,
-    }))
-);
-const DEFAULT_SAMPLE_RATE = TABLE_SAMPLING_CONFIG.default_sample_rate;
 const TableSamplingSelector: React.FC<{
-    sampleRate: number;
+    sampleRate: number | undefined;
     setSampleRate: (sampleRate: number) => void;
     tooltipPos: TooltipDirection;
     onTableSamplingInfoClick: () => void;
 }> = ({ sampleRate, setSampleRate, tooltipPos, onTableSamplingInfoClick }) => {
+    const sampleRateOptions = React.useMemo(getTableSamplingRateOptions, []);
+    const userDefaultTableSampleRate = useSelector(
+        (state: IStoreState) => state.user.computedSettings['table_sample_rate']
+    );
+
     React.useEffect(() => {
-        if (!sampleRateOptions.some((option) => option.value === sampleRate)) {
-            setSampleRate(DEFAULT_SAMPLE_RATE);
+        if (
+            sampleRate === undefined ||
+            !sampleRateOptions.some((option) => option.value === sampleRate)
+        ) {
+            setSampleRate(parseFloat(userDefaultTableSampleRate));
         }
-    }, [sampleRate, setSampleRate]);
+    }, [
+        sampleRate,
+        setSampleRate,
+        sampleRateOptions,
+        userDefaultTableSampleRate,
+    ]);
 
     const selectedSampleRateText = React.useMemo(() => {
         if (sampleRate > 0) {

diff --git a/querybook/webapp/components/Search/SearchOverview.tsx b/querybook/webapp/components/Search/SearchOverview.tsx
@@ -6,7 +6,6 @@ import CreatableSelect from 'react-select/creatable';
 
 import { UserAvatar } from 'components/UserBadge/UserAvatar';
 import { UserSelect } from 'components/UserSelect/UserSelect';
-import PublicConfig from 'config/querybook_public_config.yaml';
 import { ComponentType, ElementType } from 'const/analytics';
 import {
     IBoardPreview,
@@ -19,6 +18,7 @@ import { useShallowSelector } from 'hooks/redux/useShallowSelector';
 import { useSurveyTrigger } from 'hooks/ui/useSurveyTrigger';
 import { useTrackView } from 'hooks/useTrackView';
 import { trackClick, trackView } from 'lib/analytics';
+import { isAIFeatureEnabled } from 'lib/public-config';
 import { titleize } from 'lib/utils';
 import { getCurrentEnv } from 'lib/utils/query-string';
 import {
@@ -69,8 +69,6 @@ import { TableSelect } from './TableSelect';
 
 import './SearchOverview.scss';
 
-const AIAssistantConfig = PublicConfig.ai_assistant;
-
 const userReactSelectStyle = makeReactSelectStyle(
     true,
     miniAsyncReactSelectStyles
@@ -314,8 +312,7 @@ export const SearchOverview: React.FC<ISearchOverviewProps> = ({
                     autoFocus
                 />
                 {searchType === SearchType.Table &&
-                    AIAssistantConfig.enabled &&
-                    AIAssistantConfig.table_vector_search.enabled && (
+                    isAIFeatureEnabled('table_vector_search') && (
                         <div className="mt8 flex-row">
                             <AccentText weight="bold" className="ml8 mr12">
                                 Natural Language Search

diff --git a/querybook/webapp/components/UserSettingsMenu/UserSettingsMenu.tsx b/querybook/webapp/components/UserSettingsMenu/UserSettingsMenu.tsx
@@ -2,8 +2,12 @@ import React, { useMemo } from 'react';
 import { useDispatch, useSelector } from 'react-redux';
 
 import { UserSettingsTab } from 'components/EnvironmentAppRouter/modalRoute/UserSettingsMenuRoute';
-import PublicConfig from 'config/querybook_public_config.yaml';
 import userSettingConfig from 'config/user_setting.yaml';
+import {
+    getTableSamplingRateOptions,
+    isAIFeatureEnabled,
+    TABLE_SAMPLING_CONFIG,
+} from 'lib/public-config';
 import { titleize } from 'lib/utils';
 import { availableEnvironmentsSelector } from 'redux/environment/selector';
 import { notificationServiceSelector } from 'redux/notificationService/selector';
@@ -14,8 +18,6 @@ import { makeSelectOptions, Select } from 'ui/Select/Select';
 
 import './UserSettingsMenu.scss';
 
-const AIAssistantConfig = PublicConfig.ai_assistant;
-
 export const UserSettingsMenu: React.FC<{ tab: UserSettingsTab }> = ({
     tab,
 }) => {
@@ -41,9 +43,7 @@ export const UserSettingsMenu: React.FC<{ tab: UserSettingsTab }> = ({
             Object.entries(userSettingConfig).filter(([key, value]) => {
                 if (key === 'sql_complete') {
                     return (
-                        AIAssistantConfig.enabled &&
-                        AIAssistantConfig.sql_complete.enabled &&
-                        value.tab === tab
+                        isAIFeatureEnabled('sql_complete') && value.tab === tab
                     );
                 }
                 return value.tab === tab;
@@ -81,6 +81,9 @@ export const UserSettingsMenu: React.FC<{ tab: UserSettingsTab }> = ({
                 return makeSelectOptions(
                     notifiers.map((notifier) => notifier.name)
                 );
+            } else if (key === 'table_sample_rate') {
+                const options = getTableSamplingRateOptions();
+                return makeSelectOptions(options);
             }
             return makeSelectOptions(userSettingConfig[key].options);
         },
@@ -100,9 +103,18 @@ export const UserSettingsMenu: React.FC<{ tab: UserSettingsTab }> = ({
         [userSettingByKey, setUserSettings, getRawKey]
     );
 
+    const getValueByKey = (key: string) => {
+        let defaultValue = userSettingConfig[key].default;
+
+        if (key === 'table_sample_rate') {
+            defaultValue = TABLE_SAMPLING_CONFIG.default_sample_rate.toString();
+        }
+
+        return userSettingByKey[getRawKey(key)] ?? defaultValue;
+    };
+
     const makeFieldByKey = (key: string) => {
-        const value =
-            userSettingByKey[getRawKey(key)] ?? userSettingConfig[key].default;
+        const value = getValueByKey(key);
 
         const formField = (
             <>