ui: Show sampled query executions as sampled (#1465)

* ui: Show sampled query execution as sampled

* make metadata JSON column type

querybook/migrations/versions/f7b11b3e3a95_add_query_execution_metadata.py

@@ -0,0 +1,38 @@
+"""add query execution metadata
+
+Revision ID: f7b11b3e3a95
+Revises: 767cf0f573bd
+Create Date: 2024-06-18 19:33:42.080117
+
+"""
+
+from alembic import op
+import sqlalchemy as sa
+
+
+# revision identifiers, used by Alembic.
+revision = "f7b11b3e3a95"
+down_revision = "767cf0f573bd"
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.create_table(
+        "query_execution_metadata",
+        sa.Column("id", sa.Integer(), nullable=False),
+        sa.Column("query_execution_id", sa.Integer(), nullable=True),
+        sa.Column("execution_metadata", sa.JSON(), nullable=True),
+        sa.ForeignKeyConstraint(
+            ["query_execution_id"], ["query_execution.id"], ondelete="CASCADE"
+        ),
+        sa.PrimaryKeyConstraint("id"),
+    )
+    # ### end Alembic commands ###
+
+
+def downgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_table("query_execution_metadata")
+    # ### end Alembic commands ###

querybook/server/datasources/query_execution.py

@@ -61,7 +61,9 @@
 
 
 @register("/query_execution/", methods=["POST"])
-def create_query_execution(query, engine_id, data_cell_id=None, originator=None):
+def create_query_execution(
+    query, engine_id, metadata=None, data_cell_id=None, originator=None
+):
     with DBSession() as session:
         verify_query_engine_permission(engine_id, session=session)
 
@@ -70,6 +72,11 @@ def create_query_execution(query, engine_id, data_cell_id=None, originator=None)
             query=query, engine_id=engine_id, uid=uid, session=session
         )
 
+        if metadata:
+            logic.create_query_execution_metadata(
+                query_execution.id, metadata, session=session
+            )
+
         data_doc = None
         if data_cell_id:
             datadoc_logic.append_query_executions_to_data_cell(
@@ -258,6 +265,15 @@ def get_query_execution_error(query_execution_id):
         return logic.get_query_execution_error(query_execution_id, session=session)
 
 
+@register("/query_execution/<int:query_execution_id>/metadata/", methods=["GET"])
+def get_query_execution_metadata(query_execution_id):
+    verify_query_execution_permission(query_execution_id)
+    execution_metadata = logic.get_query_execution_metadata_by_execution_id(
+        query_execution_id
+    )
+    return execution_metadata.to_dict() if execution_metadata is not None else None
+
+
 @register(
     "/statement_execution/<int:statement_execution_id>/result/download/",
     methods=["GET"],

querybook/server/logic/query_execution.py

@@ -9,6 +9,7 @@
 from lib.logger import get_logger
 from models.query_execution import (
     QueryExecution,
+    QueryExecutionMetadata,
     StatementExecution,
     QueryExecutionNotification,
     QueryExecutionError,
@@ -170,6 +171,33 @@ def get_environments_by_execution_id(execution_id, session=None):
     )
 
 
+@with_session
+def create_query_execution_metadata(
+    query_execution_id,
+    metadata,
+    commit=True,
+    session=None,
+):
+    return QueryExecutionMetadata.create(
+        {
+            "query_execution_id": query_execution_id,
+            "execution_metadata": metadata,
+        },
+        commit=commit,
+        session=session,
+    )
+
+
+@with_session
+def get_query_execution_metadata_by_execution_id(
+    query_execution_id,
+    session=None,
+):
+    return QueryExecutionMetadata.get(
+        session=session, query_execution_id=query_execution_id
+    )
+
+
 """
     ----------------------------------------------------------------------------------------------------------
     STATEMENT EXECUTION

querybook/server/models/query_execution.py

@@ -188,6 +188,39 @@ def to_dict(self):
         }
 
 
+class QueryExecutionMetadata(CRUDMixin, Base):
+    """
+    Represents metadata for a query execution.
+
+    Attributes:
+        id: primary key
+        query_execution_id: foreign key of the related query execution.
+        execution_metadata: A JSON object containing metadata about the query execution.
+            - 'sample_rate': The rate at which the table was sampled for this query execution.
+    """
+
+    __tablename__ = "query_execution_metadata"
+
+    id = sql.Column(sql.Integer, primary_key=True)
+    query_execution_id = sql.Column(
+        sql.Integer, sql.ForeignKey("query_execution.id", ondelete="CASCADE")
+    )
+    execution_metadata = sql.Column(sql.JSON)
+
+    query_execution = relationship(
+        "QueryExecution",
+        uselist=False,
+        backref=backref("metadata", cascade="all, delete", passive_deletes=True),
+    )
+
+    def to_dict(self):
+        return {
+            "id": self.id,
+            "query_execution_id": self.query_execution_id,
+            "metadata": self.execution_metadata,
+        }
+
+
 class QueryExecutionViewer(CRUDMixin, Base):
     __tablename__ = "query_execution_viewer"
     __table_args__ = (

querybook/webapp/components/DataDocQueryCell/DataDocQueryCell.tsx

@@ -477,7 +477,10 @@ class DataDocQueryCellComponent extends React.PureComponent<IProps, IState> {
                     await this.props.createQueryExecution(
                         query,
                         engineId,
-                        this.props.cellId
+                        this.props.cellId,
+                        this.sampleRate > 0
+                            ? { sample_rate: this.sampleRate }
+                            : null
                     )
                 ).id;
 
@@ -548,7 +551,8 @@ class DataDocQueryCellComponent extends React.PureComponent<IProps, IState> {
             return this.props.createQueryExecution(
                 renderedQuery,
                 this.engineId,
-                this.props.cellId
+                this.props.cellId,
+                this.sampleRate > 0 ? { sample_rate: this.sampleRate } : null
             );
         }
     }
@@ -1092,8 +1096,9 @@ function mapDispatchToProps(dispatch: Dispatch) {
         createQueryExecution: (
             query: string,
             engineId: number,
-            cellId: number
-        ) => dispatch(createQueryExecution(query, engineId, cellId)),
+            cellId: number,
+            metadata: Record<string, string | number>
+        ) => dispatch(createQueryExecution(query, engineId, cellId, metadata)),
 
         setTableSidebarId: (id: number) => dispatch(setSidebarTableId(id)),
 

querybook/webapp/components/DataDocStatementExecution/DataDocStatementExecution.tsx

@@ -28,6 +28,8 @@ interface IProps {
     showStatementMeta: boolean;
 
     toggleStatementMeta: () => any;
+
+    queryExecutionId: number;
 }
 
 function useStatementResult(statementExecution: IStatementExecution) {
@@ -103,6 +105,8 @@ export const DataDocStatementExecution: React.FC<IProps> = ({
     showStatementLogs,
     showStatementMeta,
     toggleStatementMeta,
+
+    queryExecutionId,
 }) => {
     const [showInFullScreen, , toggleFullScreen] = useToggleState(false);
     const statementMeta = useStatementMeta(
@@ -188,6 +192,7 @@ export const DataDocStatementExecution: React.FC<IProps> = ({
                         resultLimit={resultLimit}
                         setResultLimit={setResultLimit}
                         isFetchingStatementResult={isFetchingStatementResult}
+                        queryExecutionId={queryExecutionId}
                     />
                 </>
             );
@@ -226,6 +231,7 @@ export const DataDocStatementExecution: React.FC<IProps> = ({
                     setResultLimit={setResultLimit}
                     resultLimit={resultLimit}
                     isFetchingStatementResult={isFetchingStatementResult}
+                    queryExecutionId={queryExecutionId}
                 />
             </Modal>
         ) : null;

querybook/webapp/components/DataDocStatementExecution/StatementResult.scss

@@ -23,6 +23,9 @@
             .warning-word {
                 color: var(--color-false);
             }
+            .accent-warning-text {
+                color: var(--color-accent-dark);
+            }
         }
 
         .view-type-switcher {

querybook/webapp/components/DataDocStatementExecution/StatementResult.tsx

@@ -6,16 +6,20 @@ 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 { getSelectStatementLimit } from 'lib/sql-helper/sql-limiter';
 import { stopPropagation } from 'lib/utils/noop';
 import { formatNumber } from 'lib/utils/number';
 import { IStoreState } from 'redux/store/types';
+import { QueryExecutionMetadataResource } from 'resource/queryExecution';
 import { TextButton } from 'ui/Button/Button';
+import { IconButton } from 'ui/Button/IconButton';
 import { InfoButton } from 'ui/Button/InfoButton';
 import { Checkbox } from 'ui/Checkbox/Checkbox';
 import { Icon } from 'ui/Icon/Icon';
@@ -39,6 +43,8 @@ interface IProps {
     isFetchingStatementResult: boolean;
     resultLimit: number;
     setResultLimit: (limit: number) => void;
+
+    queryExecutionId: number;
 }
 
 export const StatementResult: React.FC<IProps> = (props) => {
@@ -73,6 +79,8 @@ const StatementResultWithResult: React.FC<IProps> = ({
     isFetchingStatementResult,
     resultLimit,
     setResultLimit,
+
+    queryExecutionId,
 }) => {
     const { result_row_count: resultRowCount } = statementExecution;
     const { data: rawData } = statementResult;
@@ -104,6 +112,13 @@ const StatementResultWithResult: React.FC<IProps> = ({
     const actualRowMinusColCount = Math.max(data.length - 1, 0);
     const fetchedAllRows = resultRowMinusColCount === actualRowMinusColCount;
 
+    const { data: executionMetadata } = useResource(
+        React.useCallback(
+            () => QueryExecutionMetadataResource.get(queryExecutionId),
+            [queryExecutionId]
+        )
+    );
+
     const explorationButtons = [
         <ColumnToggleMenuButton
             columnNames={columnNames}
@@ -138,6 +153,7 @@ const StatementResultWithResult: React.FC<IProps> = ({
             resultLimit={resultLimit}
             setResultLimit={setResultLimit}
             isFetchingStatementResult={isFetchingStatementResult}
+            sampleRate={Number(executionMetadata?.metadata?.sample_rate) ?? 0}
         />
     );
     const visualizationDOM = data.length ? (
@@ -246,6 +262,8 @@ const FetchInfo: React.FC<{
     resultLimit: number;
     setResultLimit: (newLimit: number) => void;
     isFetchingStatementResult: boolean;
+
+    sampleRate: number;
 }> = ({
     statementQueryLimit,
     resultRowMinusColCount,
@@ -255,6 +273,8 @@ const FetchInfo: React.FC<{
     resultLimit,
     setResultLimit,
     isFetchingStatementResult,
+
+    sampleRate,
 }) => {
     const getFetchInfo = () => {
         if (isFetchingStatementResult) {
@@ -266,6 +286,29 @@ const FetchInfo: React.FC<{
             );
         }
 
+        const sampleUserGuideLink =
+            PublicConfig.table_sampling?.sample_user_guide_link ?? '';
+        const sampleRateText = (
+            <div className="flex-row ml4">
+                (Full Result,
+                <span className="accent-warning-text ml4">
+                    Sampled {sampleRate}%
+                </span>
+                {sampleUserGuideLink ? (
+                    <IconButton
+                        className="ml4"
+                        onClick={() => {
+                            window.open(sampleUserGuideLink, '_blank');
+                        }}
+                        icon="Info"
+                        size={16}
+                        noPadding
+                    />
+                ) : null}
+                )
+            </div>
+        );
+
         const limitReachedText =
             statementQueryLimit === resultRowMinusColCount ? (
                 <AccentText
@@ -283,7 +326,7 @@ const FetchInfo: React.FC<{
             return (
                 <span className="flex-row">
                     {formatNumber(actualRowMinusColCount, 'row')}
-                    (Full Result)
+                    {sampleRate > 0 ? sampleRateText : ' (Full Result)'}
                     {limitReachedText}
                 </span>
             );

querybook/webapp/components/QueryComposer/QueryComposer.tsx

@@ -524,7 +524,12 @@ const QueryComposer: React.FC = () => {
             engine.id,
             async (query, engineId) => {
                 const data = await dispatch(
-                    queryExecutionsAction.createQueryExecution(query, engineId)
+                    queryExecutionsAction.createQueryExecution(
+                        query,
+                        engineId,
+                        null,
+                        sampleRate > 0 ? { sample_rate: sampleRate } : null
+                    )
                 );
                 return data.id;
             }

querybook/webapp/components/QueryExecution/QueryExecution.tsx

@@ -147,6 +147,7 @@ export const QueryExecution: React.FC<IProps> = ({
                 showStatementMeta={showStatementMeta}
                 showStatementLogs={showStatementLogs}
                 toggleStatementMeta={toggleShowStatementMeta}
+                queryExecutionId={id}
             />
         ) : queryExecution.status <= QueryExecutionStatus.RUNNING ? (
             <Loading />

querybook/webapp/const/queryExecution.ts

@@ -45,6 +45,12 @@ export interface IQueryExecutionViewer {
     execution_id: number;
 }
 
+export interface IQueryExecutionMetadata {
+    id: number;
+    execution_id: number;
+    metadata: Record<string, string | number>;
+}
+
 export interface IQueryExecution {
     id: number;
     created_at: number;