Refactor

superset-frontend/plugins/plugin-chart-ag-grid-table/src/AgGridTable/components/CustomHeader.tsx

@@ -26,11 +26,16 @@ import {
   ArrowUpOutlined,
   PlusOutlined,
 } from '@ant-design/icons';
+import { ColDef } from 'ag-grid-community';
 import CustomPopover from './CustomPopover';
-import { CustomColDef } from '..';
 import FilterIcon from './Filter';
 import KebabMenu from './KebabMenu';
-import { CustomHeaderParams, SortState, UserProvidedColDef } from '../../types';
+import {
+  CustomColDef,
+  CustomHeaderParams,
+  SortState,
+  UserProvidedColDef,
+} from '../../types';
 
 // Styled Components
 const Container = styled.div`
@@ -212,14 +217,16 @@ const CustomHeader: React.FC<CustomHeaderParams> = ({
     setIsMenuVisible(!isMenuVisible);
   };
 
+  const isTimeComparisonColumn = (col: ColDef<any, any>) =>
+    (col as UserProvidedColDef).timeComparisonKey === timeComparisonKey &&
+    !(col as UserProvidedColDef).isMain;
+
   const handleToggleComparison = (event: React.MouseEvent) => {
     event.stopPropagation();
 
     const allColumns = api.getColumnDefs();
-    const timeComparisonColumns = allColumns?.filter(
-      col =>
-        (col as UserProvidedColDef).timeComparisonKey === timeComparisonKey &&
-        !(col as UserProvidedColDef).isMain,
+    const timeComparisonColumns = allColumns?.filter(col =>
+      isTimeComparisonColumn(col),
     );
     const timeComparsionColIds = timeComparisonColumns?.map(
       item => (item as UserProvidedColDef).field || '',
@@ -300,7 +307,7 @@ const CustomHeader: React.FC<CustomHeaderParams> = ({
           {FilterIcon}
         </FilterIconWrapper>
       </CustomPopover>
-      {!isPercentMetric && (
+      {!isPercentMetric && !isTimeComparison && (
         <CustomPopover
           content={menuContent}
           isOpen={isMenuVisible}

superset-frontend/plugins/plugin-chart-ag-grid-table/src/AgGridTable/components/TimeComparisonVisibility.tsx

@@ -0,0 +1,127 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+/* eslint-disable import/no-extraneous-dependencies */
+import { useState } from 'react';
+import { Dropdown, Menu } from 'antd';
+import { TableOutlined, DownOutlined, CheckOutlined } from '@ant-design/icons';
+import { css, useTheme, t } from '@superset-ui/core';
+
+interface ComparisonColumn {
+  key: string;
+  label: string;
+}
+
+interface TimeComparisonVisibilityProps {
+  comparisonColumns: ComparisonColumn[];
+  selectedComparisonColumns: string[];
+  onSelectionChange: (selectedColumns: string[]) => void;
+}
+
+const TimeComparisonVisibility: React.FC<TimeComparisonVisibilityProps> = ({
+  comparisonColumns,
+  selectedComparisonColumns,
+  onSelectionChange,
+}) => {
+  const [showComparisonDropdown, setShowComparisonDropdown] = useState(false);
+  const theme = useTheme();
+
+  const allKey = comparisonColumns[0].key;
+
+  const handleOnClick = (data: any) => {
+    const { key } = data;
+    // Toggle 'All' key selection
+    if (key === allKey) {
+      onSelectionChange([allKey]);
+    } else if (selectedComparisonColumns.includes(allKey)) {
+      onSelectionChange([key]);
+    } else {
+      // Toggle selection for other keys
+      onSelectionChange(
+        selectedComparisonColumns.includes(key)
+          ? selectedComparisonColumns.filter((k: string) => k !== key) // Deselect if already selected
+          : [...selectedComparisonColumns, key],
+      ); // Select if not already selected
+    }
+  };
+
+  const handleOnBlur = () => {
+    if (selectedComparisonColumns.length === 3) {
+      onSelectionChange([comparisonColumns[0].key]);
+    }
+  };
+
+  return (
+    <Dropdown
+      placement="bottomRight"
+      visible={showComparisonDropdown}
+      onVisibleChange={(flag: boolean) => {
+        setShowComparisonDropdown(flag);
+      }}
+      overlay={
+        <Menu
+          multiple
+          onClick={handleOnClick}
+          onBlur={handleOnBlur}
+          selectedKeys={selectedComparisonColumns}
+        >
+          <div
+            css={css`
+              max-width: 242px;
+              padding: 0 ${theme.gridUnit * 2}px;
+              color: ${theme.colors.grayscale.base};
+              font-size: ${theme.typography.sizes.s}px;
+            `}
+          >
+            {t(
+              'Select columns that will be displayed in the table. You can multiselect columns.',
+            )}
+          </div>
+          {comparisonColumns.map((column: ComparisonColumn) => (
+            <Menu.Item key={column.key}>
+              <span
+                css={css`
+                  color: ${theme.colors.grayscale.dark2};
+                `}
+              >
+                {column.label}
+              </span>
+              <span
+                css={css`
+                  float: right;
+                  font-size: ${theme.typography.sizes.s}px;
+                `}
+              >
+                {selectedComparisonColumns.includes(column.key) && (
+                  <CheckOutlined />
+                )}
+              </span>
+            </Menu.Item>
+          ))}
+        </Menu>
+      }
+      trigger={['click']}
+    >
+      <span>
+        <TableOutlined /> <DownOutlined />
+      </span>
+    </Dropdown>
+  );
+};
+
+export default TimeComparisonVisibility;

superset-frontend/plugins/plugin-chart-ag-grid-table/src/AgGridTable/index.tsx

@@ -33,9 +33,7 @@ import {
   type ColDef,
   ModuleRegistry,
   GridReadyEvent,
-  GridApi,
   GridState,
-  SortModelItem,
   CellClickedEvent,
 } from 'ag-grid-community';
 import './styles/ag-grid.css';
@@ -53,17 +51,11 @@ import { SearchOutlined } from '@ant-design/icons';
 import { debounce, isEqual } from 'lodash';
 import Pagination from './components/Pagination';
 import SearchSelectDropdown from './components/SearchSelectDropdown';
-import { SearchOption, SortByItem } from '../types';
-import getInitialSortState from '../utils/getInitialSortState';
-
-export interface CustomColDef extends ColDef {
-  customMeta?: {
-    isMetric?: boolean;
-    isPercentMetric?: boolean;
-  };
-}
+import { CustomColDef, SearchOption, SortByItem } from '../types';
+import getInitialSortState, { shouldSort } from '../utils/getInitialSortState';
+import { PAGE_SIZE_OPTIONS } from '../consts';
 
-export interface Props {
+export interface AgGridTableProps {
   gridTheme?: string;
   isDarkMode?: boolean;
   gridHeight?: number | null;
@@ -171,13 +163,10 @@ const StyledContainer = styled.div`
 
 const isSearchFocused = new Map<string, boolean>();
 
-const AgGridDataTable: FunctionComponent<Props> = memo(
+const AgGridDataTable: FunctionComponent<AgGridTableProps> = memo(
   ({
-    gridTheme = 'ag-theme-quartz',
-    isDarkMode = false,
     gridHeight = null,
     data = [],
-    onGridReady,
     colDefsFromProps,
     includeSearch,
     allowRearrangeColumns,
@@ -203,7 +192,6 @@ const AgGridDataTable: FunctionComponent<Props> = memo(
     showTotals,
   }) => {
     const gridRef = useRef<AgGridReact>(null);
-    const gridApiRef = useRef<GridApi | null>(null);
     const inputRef = useRef<HTMLInputElement>(null);
     const searchId = `search-${id}`;
     const gridInitialState: GridState = {
@@ -227,39 +215,15 @@ const AgGridDataTable: FunctionComponent<Props> = memo(
       [],
     );
 
-    // Handle grid ready event
-    const handleGridReady = useCallback(
-      (params: GridReadyEvent) => {
-        gridApiRef.current = params.api;
-        if (onGridReady) {
-          onGridReady(params);
-        }
-
-        // Optional: Auto-size columns on initial load
-        params.api.sizeColumnsToFit();
-      },
-      [onGridReady],
-    );
-
     // Memoize container style
     const containerStyle = useMemo(
       () => ({
-        height: gridHeight
-          ? includeSearch
-            ? `${gridHeight - 16}px`
-            : `${gridHeight}px`
-          : '100%',
+        height: gridHeight ? `${gridHeight}px` : '100%',
         width: '100%',
       }),
       [gridHeight],
     );
 
-    // Memoize grid class name
-    const gridClassName = useMemo(
-      () => `ag-theme-quartz${isDarkMode ? '-dark' : ''}`,
-      [isDarkMode],
-    );
-
     const [quickFilterText, setQuickFilterText] = useState<string>();
     const [searchValue, setSearchValue] = useState(
       serverPaginationData?.searchText || '',
@@ -310,32 +274,6 @@ const AgGridDataTable: FunctionComponent<Props> = memo(
       [serverPagination, debouncedSearch, searchId],
     );
 
-    const shouldSort = ({
-      colId,
-      sortDir,
-      percentMetrics,
-      serverPagination,
-      gridInitialState,
-    }: {
-      colId: string;
-      sortDir: string;
-      percentMetrics: string[];
-      serverPagination: boolean;
-      gridInitialState: GridState;
-    }) => {
-      if (percentMetrics.includes(colId)) return false;
-      if (!serverPagination) return false;
-
-      const initialSort: Partial<SortModelItem> =
-        gridInitialState?.sort?.sortModel?.[0] || {};
-      const { colId: initialColId = '', sort: initialSortDir = '' } =
-        initialSort;
-
-      if (initialColId === colId && initialSortDir === sortDir) return false;
-
-      return true;
-    };
-
     const handleColumnHeaderClick = useCallback(
       params => {
         const colId = params?.column?.colId;
@@ -392,7 +330,7 @@ const AgGridDataTable: FunctionComponent<Props> = memo(
 
     return (
       <StyledContainer>
-        <div className={gridClassName} style={containerStyle}>
+        <div className="ag-theme-quartz" style={containerStyle}>
           <div className="dropdown-controls-container">
             {renderTimeComparisonDropdown && (
               <div className="time-comparison-dropdown">
@@ -446,12 +384,11 @@ const AgGridDataTable: FunctionComponent<Props> = memo(
             groupDefaultExpanded={-1}
             rowGroupPanelShow="always"
             enableCellTextSelection
-            onGridReady={handleGridReady}
             quickFilterText={serverPagination ? '' : quickFilterText}
             suppressMovableColumns={!allowRearrangeColumns}
             pagination={pagination}
             paginationPageSize={pageSize}
-            paginationPageSizeSelector={[10, 20, 50, 100, 200]}
+            paginationPageSizeSelector={PAGE_SIZE_OPTIONS}
             suppressDragLeaveHidesColumns
             pinnedBottomRowData={showTotals ? [cleanedTotals] : undefined}
             context={{