From 370948f8e3bd9476637d00e0cf4e4121bc27ef8c Mon Sep 17 00:00:00 2001
From: Krishna Gopal <kgopal@pinterest.com>
Date: Tue, 19 Dec 2023 12:39:18 -0600
Subject: [PATCH] feat: update column stats ui, add sort by usage to columns
 (#1389)

* feat: update column stats ui, add sort by usage to columns
---
 querybook/server/datasources/metastore.py     | 42 +++++-----
 querybook/server/logic/metastore.py           | 33 ++++++++
 .../DataTableStats/DataTableColumnStats.tsx   | 31 ++++---
 .../DataTableView/DataTableView.tsx           |  1 -
 .../DataTableColumnCard.tsx                   | 35 ++++----
 .../DataTableViewColumn.tsx                   | 81 +++++++++++++++----
 querybook/webapp/resource/table.ts            |  2 +
 7 files changed, 153 insertions(+), 72 deletions(-)

diff --git a/querybook/server/datasources/metastore.py b/querybook/server/datasources/metastore.py
index 706ee396d..7cde405c4 100644
--- a/querybook/server/datasources/metastore.py
+++ b/querybook/server/datasources/metastore.py
@@ -1,35 +1,29 @@
 from typing import Tuple, Union
-from flask_login import current_user
 
 from app.auth.permission import (
-    verify_query_engine_environment_permission,
-    verify_environment_permission,
-    verify_metastore_permission,
+    verify_data_column_permission,
     verify_data_schema_permission,
     verify_data_table_permission,
-    verify_data_column_permission,
+    verify_environment_permission,
+    verify_metastore_permission,
+    verify_query_engine_environment_permission,
 )
+from app.datasource import admin_only, api_assert, register, with_impression
 from app.db import DBSession
-from app.datasource import register, api_assert, with_impression, admin_only
 from app.flask_app import cache, limiter
+from const.datasources import RESOURCE_NOT_FOUND_STATUS_CODE
 from const.impression import ImpressionItemType
 from const.metastore import DataTableWarningSeverity, MetadataType
 from const.time import seconds_in_a_day
-from const.datasources import RESOURCE_NOT_FOUND_STATUS_CODE
+from flask_login import current_user
 from lib.lineage.utils import lineage
-from lib.metastore.utils import DataTableFinder
 from lib.metastore import get_metastore_loader
+from lib.metastore.utils import DataTableFinder
 from lib.query_analysis.samples import make_samples_query
 from lib.utils import mysql_cache
-from logic import metastore as logic
 from logic import admin as admin_logic
-from logic import data_element as data_element_logic
-from logic import tag as tag_logic
-from models.metastore import (
-    DataTableWarning,
-    DataTableStatistics,
-    DataTableColumnStatistics,
-)
+from logic import metastore as logic
+from models.metastore import DataTableStatistics, DataTableWarning
 from tasks.run_sample_query import run_sample_query
 
 
@@ -269,6 +263,13 @@ def get_columns_from_table(table_id):
         return logic.get_column_by_table_id(table_id, session=session)
 
 
+@register("/table/<int:table_id>/detailed_column/", methods=["GET"])
+def get_detailed_columns_from_table(table_id):
+    with DBSession() as session:
+        verify_data_table_permission(table_id, session=session)
+        return logic.get_detailed_columns_dict_by_table_id(table_id, session=session)
+
+
 @register("/table/<int:table_id>/raw_samples_query/", methods=["GET"])
 def get_table_samples_raw_query(
     table_id,
@@ -380,14 +381,7 @@ def get_column_by_table(table_id, column_name, with_table=False):
 def get_column(column_id, with_table=False):
     column = logic.get_column_by_id(column_id)
     verify_data_table_permission(column.table_id)
-    column_dict = column.to_dict(with_table)
-
-    column_dict["stats"] = DataTableColumnStatistics.get_all(column_id=column_id)
-    column_dict["tags"] = tag_logic.get_tags_by_column_id(column_id=column_id)
-    column_dict[
-        "data_element_association"
-    ] = data_element_logic.get_data_element_association_by_column_id(column_id)
-    return column_dict
+    return logic.get_detailed_column_dict(column, with_table=with_table)
 
 
 @register("/column/<int:column_id>/", methods=["PUT"])
diff --git a/querybook/server/logic/metastore.py b/querybook/server/logic/metastore.py
index 4f1546422..8f787f2be 100644
--- a/querybook/server/logic/metastore.py
+++ b/querybook/server/logic/metastore.py
@@ -5,6 +5,7 @@
 from const.metastore import DataOwner, DataTableWarningSeverity
 from lib.logger import get_logger
 from lib.sqlalchemy import update_model_fields
+from logic import data_element as data_element_logic
 from logic.user import get_user_by_name
 from models.admin import QueryEngineEnvironment
 from models.metastore import (
@@ -500,6 +501,38 @@ def get_column_by_table_id(table_id, session=None):
     )
 
 
+@with_session
+def get_detailed_column_dict(column: DataTableColumn, with_table=False, session=None):
+    from logic import tag as tag_logic
+
+    column_dict = column.to_dict(with_table)
+    column_dict["stats"] = DataTableColumnStatistics.get_all(
+        column_id=column.id, session=session
+    )
+    column_dict["tags"] = tag_logic.get_tags_by_column_id(
+        column_id=column.id, session=session
+    )
+    column_dict[
+        "data_element_association"
+    ] = data_element_logic.get_data_element_association_by_column_id(
+        column.id, session=session
+    )
+    return column_dict
+
+
+@with_session
+def get_detailed_columns_dict_by_table_id(table_id, session=None):
+    data_table_columns = (
+        session.query(DataTableColumn)
+        .filter(DataTableColumn.table_id == table_id)
+        .all()
+    )
+    columns_info = []
+    for col in data_table_columns:
+        columns_info.append(get_detailed_column_dict(col, session=session))
+    return columns_info
+
+
 @with_session
 def get_all_column_name_by_table_id(table_id, session=None):
     return (
diff --git a/querybook/webapp/components/DataTableStats/DataTableColumnStats.tsx b/querybook/webapp/components/DataTableStats/DataTableColumnStats.tsx
index 63cbeab79..9fd5739db 100644
--- a/querybook/webapp/components/DataTableStats/DataTableColumnStats.tsx
+++ b/querybook/webapp/components/DataTableStats/DataTableColumnStats.tsx
@@ -1,10 +1,9 @@
 import * as React from 'react';
 
 import { ITableColumnStats } from 'const/metastore';
-import { isNumeric } from 'lib/utils/number';
-import { KeyContentDisplay } from 'ui/KeyContentDisplay/KeyContentDisplay';
 
-import { TableStats } from './DataTableStatsCommon';
+import { Tag, TagGroup } from 'ui/Tag/Tag';
+import { KeyContentDisplay } from 'ui/KeyContentDisplay/KeyContentDisplay';
 
 interface IProps {
     stats: ITableColumnStats[];
@@ -13,15 +12,21 @@ interface IProps {
 export const DataTableColumnStats: React.FunctionComponent<IProps> = ({
     stats,
 }) => {
-    const statsDOM = (stats || []).map((tableColumnStat) => (
-        <KeyContentDisplay
-            key={tableColumnStat.id}
-            keyString={tableColumnStat.key}
-            rightAlign={isNumeric(tableColumnStat.value)}
-        >
-            <TableStats val={tableColumnStat.value} />
-        </KeyContentDisplay>
-    ));
+    const statsDOM = (stats || []).map((stat, i) => {
+        const formattedValue = Array.isArray(stat.value)
+            ? stat.value.join(', ')
+            : stat.value;
+        return (
+            <TagGroup key={stat.key}>
+                <Tag>{stat.key}</Tag>
+                <Tag highlighted>{formattedValue}</Tag>
+            </TagGroup>
+        );
+    });
 
-    return <div className="DataTableColumnStats">{statsDOM}</div>;
+    return (
+        <div className="DataTableColumnStats">
+            <KeyContentDisplay keyString="Stats">{statsDOM}</KeyContentDisplay>
+        </div>
+    );
 };
diff --git a/querybook/webapp/components/DataTableView/DataTableView.tsx b/querybook/webapp/components/DataTableView/DataTableView.tsx
index 0d471fd5b..40d33d810 100644
--- a/querybook/webapp/components/DataTableView/DataTableView.tsx
+++ b/querybook/webapp/components/DataTableView/DataTableView.tsx
@@ -281,7 +281,6 @@ export const DataTableView: React.FC<IDataTableViewProps> = ({ tableId }) => {
     const makeColumnsDOM = (numberOfRows = null) => (
         <DataTableViewColumn
             table={table}
-            tableColumns={tableColumns}
             numberOfRows={numberOfRows}
             updateDataColumnDescription={updateDataColumnDescription}
             onEditColumnDescriptionRedirect={getOnEditMetadata(
diff --git a/querybook/webapp/components/DataTableViewColumn/DataTableColumnCard.tsx b/querybook/webapp/components/DataTableViewColumn/DataTableColumnCard.tsx
index fad65a065..ebcd0db94 100644
--- a/querybook/webapp/components/DataTableViewColumn/DataTableColumnCard.tsx
+++ b/querybook/webapp/components/DataTableViewColumn/DataTableColumnCard.tsx
@@ -5,11 +5,10 @@ import { DataElement } from 'components/DataElement/DataElement';
 import { DataElementDescription } from 'components/DataElement/DataElementDescription';
 import { DataTableColumnStats } from 'components/DataTableStats/DataTableColumnStats';
 import { TableTag } from 'components/DataTableTags/DataTableTags';
-import { IDataColumn } from 'const/metastore';
-import { useResource } from 'hooks/useResource';
+import { IDetailedDataColumn } from 'const/metastore';
+import { convertRawToContentState } from 'lib/richtext/serialize';
 import { Nullable } from 'lib/typescript';
 import { parseType } from 'lib/utils/complex-types';
-import { TableColumnResource } from 'resource/table';
 import { Card } from 'ui/Card/Card';
 import { EditableTextField } from 'ui/EditableTextField/EditableTextField';
 import { KeyContentDisplay } from 'ui/KeyContentDisplay/KeyContentDisplay';
@@ -20,7 +19,7 @@ import { DataTableColumnCardNestedType } from './DataTableColumnCardNestedType';
 import './DataTableColumnCard.scss';
 
 interface IProps {
-    column: IDataColumn;
+    column: IDetailedDataColumn;
     onEditColumnDescriptionRedirect?: Nullable<() => Promise<void>>;
     updateDataColumnDescription: (
         columnId: number,
@@ -33,27 +32,25 @@ export const DataTableColumnCard: React.FunctionComponent<IProps> = ({
     onEditColumnDescriptionRedirect,
     updateDataColumnDescription,
 }) => {
-    const { data: detailedColumn } = useResource(
-        React.useCallback(() => TableColumnResource.get(column.id), [column.id])
-    );
     const parsedType = useMemo(() => parseType('', column.type), [column.type]);
 
-    const tagsDOM = (detailedColumn?.tags || []).map((tag) => (
+    const tagsDOM = (column?.tags || []).map((tag) => (
         <TableTag tag={tag} readonly={true} key={tag.id} mini={true} />
     ));
 
+    const columnDescription = convertRawToContentState(
+        column.description as string
+    );
     const descriptionContent = (
         <div>
-            {detailedColumn?.data_element_association &&
-                !(column.description as ContentState).hasText() && (
+            {column?.data_element_association &&
+                !columnDescription.hasText() && (
                     <DataElementDescription
-                        dataElementAssociation={
-                            detailedColumn.data_element_association
-                        }
+                        dataElementAssociation={column.data_element_association}
                     />
                 )}
             <EditableTextField
-                value={column.description as ContentState}
+                value={columnDescription}
                 onSave={updateDataColumnDescription.bind(null, column.id)}
                 placeholder="add column description"
                 onEditRedirect={onEditColumnDescriptionRedirect}
@@ -86,12 +83,10 @@ export const DataTableColumnCard: React.FunctionComponent<IProps> = ({
                             </div>
                         </KeyContentDisplay>
                     )}
-                    {detailedColumn?.data_element_association && (
+                    {column?.data_element_association && (
                         <KeyContentDisplay keyString="Data Element">
                             <DataElement
-                                association={
-                                    detailedColumn.data_element_association
-                                }
+                                association={column.data_element_association}
                             />
                         </KeyContentDisplay>
                     )}
@@ -103,7 +98,9 @@ export const DataTableColumnCard: React.FunctionComponent<IProps> = ({
                     <KeyContentDisplay keyString="Description">
                         {descriptionContent}
                     </KeyContentDisplay>
-                    <DataTableColumnStats stats={detailedColumn?.stats} />
+                    {!!column?.stats?.length && (
+                        <DataTableColumnStats stats={column.stats} />
+                    )}
                 </div>
             </Card>
         </div>
diff --git a/querybook/webapp/components/DataTableViewColumn/DataTableViewColumn.tsx b/querybook/webapp/components/DataTableViewColumn/DataTableViewColumn.tsx
index 202b0f679..2c77c7132 100644
--- a/querybook/webapp/components/DataTableViewColumn/DataTableViewColumn.tsx
+++ b/querybook/webapp/components/DataTableViewColumn/DataTableViewColumn.tsx
@@ -1,18 +1,18 @@
 import { ContentState } from 'draft-js';
 import React from 'react';
 
-import { IDataColumn, IDataTable } from 'const/metastore';
+import { IDataTable } from 'const/metastore';
+import { useResource } from 'hooks/useResource';
 import { Nullable } from 'lib/typescript';
+import { TableResource } from 'resource/table';
 import { Loading } from 'ui/Loading/Loading';
 import { SearchBar } from 'ui/SearchBar/SearchBar';
 import { OrderByButton } from 'ui/OrderByButton/OrderByButton';
-
 import { DataTableColumnCard } from './DataTableColumnCard';
 import './DataTableViewColumn.scss';
 
 export interface IDataTableViewColumnProps {
     table: IDataTable;
-    tableColumns: IDataColumn[];
     numberOfRows: number;
     updateDataColumnDescription: (
         columnId: number,
@@ -21,20 +21,50 @@ export interface IDataTableViewColumnProps {
     onEditColumnDescriptionRedirect?: Nullable<() => Promise<void>>;
 }
 
+type ColumnOrderBy = 'Default' | 'Aa' | 'Usage';
+
+const COLUMN_STATS_USAGE_KEY = 'usage';
+
 export const DataTableViewColumn: React.FunctionComponent<
     IDataTableViewColumnProps
 > = ({
     updateDataColumnDescription,
     table = null,
-    tableColumns = [],
     numberOfRows = null,
     onEditColumnDescriptionRedirect,
 }) => {
     const [filterString, setFilterString] = React.useState('');
-    const [orderColumnsBy, setOrdeColumnsBy] = React.useState(true);
-    const [orderBoardBy, setOrderBoardBy] = React.useState(false);
+    const [orderColumnsBy, setOrderColumnsBy] =
+        React.useState<ColumnOrderBy>('Aa');
+    const [orderColumnsByAsc, setOrderColumnsByAsc] = React.useState(false);
+    const { data: tableColumns } = useResource(
+        React.useCallback(
+            () => TableResource.getColumnDetails(table.id),
+            [table.id]
+        )
+    );
 
+    const usageByColumnId = React.useMemo(
+        () =>
+            tableColumns?.reduce((acc, column) => {
+                acc[column.id] =
+                    column.stats?.find(
+                        (stat) => stat.key === COLUMN_STATS_USAGE_KEY
+                    )?.value ?? 0;
+                return acc;
+            }, {}),
+        [tableColumns]
+    );
+    const canSortByUsage = React.useMemo(
+        () =>
+            !!usageByColumnId &&
+            Object.values(usageByColumnId).some((usage) => usage !== 0),
+        [usageByColumnId]
+    );
     const filteredColumns = React.useMemo(() => {
+        if (!tableColumns) {
+            return [];
+        }
         const filteredCols = tableColumns.filter((column) =>
             !!filterString
                 ? column.name.toLowerCase().includes(filterString.toLowerCase())
@@ -43,34 +73,55 @@ export const DataTableViewColumn: React.FunctionComponent<
         if (numberOfRows != null) {
             filteredCols.splice(numberOfRows);
         }
-        if (orderBoardBy) {
+        if (orderColumnsBy === 'Aa') {
             filteredCols.sort(
                 (a, b) =>
                     (a.name.toLowerCase() > b.name.toLowerCase() ? 1 : -1) *
-                    (orderColumnsBy ? 1 : -1)
+                    (orderColumnsByAsc ? 1 : -1)
+            );
+        } else if (orderColumnsBy === 'Usage') {
+            filteredCols.sort(
+                (a, b) =>
+                    (usageByColumnId[a.id] > usageByColumnId[b.id] ? 1 : -1) *
+                    (orderColumnsByAsc ? 1 : -1)
             );
         }
         return filteredCols;
     }, [
         tableColumns,
-        filterString,
         numberOfRows,
         orderColumnsBy,
-        orderBoardBy,
+        filterString,
+        orderColumnsByAsc,
+        usageByColumnId,
     ]);
 
+    const onOrderByFieldToggle = React.useCallback(() => {
+        setOrderColumnsBy((v) => {
+            if (v === 'Aa') {
+                v = canSortByUsage ? 'Usage' : 'Default';
+            } else if (v === 'Usage') {
+                v = 'Default';
+            } else {
+                // v = 'Default'
+                v = 'Aa';
+            }
+            return v;
+        });
+    }, [canSortByUsage]);
+
     if (!table || !tableColumns) {
         return <Loading />;
     }
 
     const sortButton = (
         <OrderByButton
-            asc={orderColumnsBy}
-            hideAscToggle={!orderBoardBy}
-            onAscToggle={() => setOrdeColumnsBy((v) => !v)}
+            asc={orderColumnsByAsc}
+            hideAscToggle={orderColumnsBy === 'Default'}
+            onAscToggle={() => setOrderColumnsByAsc((v) => !v)}
             orderByField="name"
-            orderByFieldSymbol={orderBoardBy ? 'Aa' : 'Default'}
-            onOrderByFieldToggle={() => setOrderBoardBy((v) => !v)}
+            orderByFieldSymbol={orderColumnsBy}
+            onOrderByFieldToggle={onOrderByFieldToggle}
         />
     );
 
diff --git a/querybook/webapp/resource/table.ts b/querybook/webapp/resource/table.ts
index 8d10bcd75..89073496e 100644
--- a/querybook/webapp/resource/table.ts
+++ b/querybook/webapp/resource/table.ts
@@ -131,6 +131,8 @@ export const TableResource = {
         ds.fetch<IDataTable>(`/table_name/${schemaName}/${tableName}/`, {
             metastore_id: metastoreId,
         }),
+    getColumnDetails: (tableId: number) =>
+        ds.fetch<IDetailedDataColumn[]>(`/table/${tableId}/detailed_column/`),
 
     getMetastoreLink: (tableId: number, metadataType: MetadataType) =>
         ds.fetch<string>(`/table/${tableId}/metastore_link/`, {