Merge pull request #327 from briefercloud/sql-table-column-sort

introduce sorting by column to sql block result

apps/api/src/python/query/index.ts

@@ -7,6 +7,7 @@ import {
   RunQueryResult,
   SQLQueryConfiguration,
   SuccessRunQueryResult,
+  TableSort,
   jsonString,
 } from '@briefer/types'
 import { logger } from '../../logger.js'
@@ -346,10 +347,13 @@ export async function readDataframePage(
   queryId: string,
   dataframeName: string,
   page: number,
-  pageSize: number
+  pageSize: number,
+  sort: TableSort | null
 ): Promise<RunQueryResult | null> {
   const code = `import json
 
+sort_config = json.loads(${JSON.stringify(JSON.stringify(sort))})
+
 if not ("${dataframeName}" in globals()):
     import pandas as pd
     try:
@@ -360,9 +364,19 @@ if not ("${dataframeName}" in globals()):
 if "${dataframeName}" in globals():
     start = ${page * pageSize}
     end = (${page} + 1) * ${pageSize}
-    rows = json.loads(${dataframeName}.iloc[start:end].to_json(
-      orient="records", date_format="iso"
-    ))
+
+    df = ${dataframeName}
+    if sort_config:
+        try:
+            df = df.sort_values(by=sort_config["column"], ascending=sort_config["order"] == "asc")
+        except:
+            # try sorting as string
+            try:
+                df = df.sort_values(by=sort_config["column"], ascending=sort_config["order"] == "asc", key=lambda x: x.astype(str))
+            except:
+                pass
+
+    rows = json.loads(df.iloc[start:end].to_json(orient="records", date_format="iso"))
 
     # convert all values to string to make sure we preserve the python values
     # when displaying this data in the browser

apps/api/src/v1/workspaces/workspace/documents/document/queries/query/index.ts

@@ -5,6 +5,7 @@ import { Response, Request, Router } from 'express'
 import csvRouter from './csv.js'
 import { readDataframePage } from '../../../../../../../python/query/index.js'
 import { getJupyterManager } from '../../../../../../../jupyter/index.js'
+import { TableSort } from '@briefer/types'
 
 const queryRouter = Router({ mergeParams: true })
 
@@ -23,6 +24,8 @@ export async function getQueryHandler(req: Request, res: Response) {
         (a) => parseInt(z.string().parse(a), 10),
         z.number().nonnegative()
       ),
+      sortColumn: TableSort.shape.column.optional().nullable(),
+      sortOrder: TableSort.shape.order.optional().nullable(),
     })
     .safeParse(req.query)
 
@@ -44,7 +47,10 @@ export async function getQueryHandler(req: Request, res: Response) {
       queryId,
       data.dataframeName,
       data.page,
-      pageSize
+      pageSize,
+      data.sortColumn && data.sortOrder
+        ? { column: data.sortColumn, order: data.sortOrder }
+        : null
     )
     if (!result) {
       res.status(404).end()

apps/api/src/yjs/v2/executor/sql.ts

@@ -126,6 +126,7 @@ export class SQLExecutor implements ISQLExecutor {
       }
 
       block.setAttribute('result', null)
+      block.setAttribute('sort', null)
 
       let actualSource =
         (metadata.isSuggestion ? aiSuggestions : source)?.toJSON().trim() ?? ''

apps/web/src/components/v2Editor/customBlocks/sql/SQLResult.tsx

@@ -1,5 +1,3 @@
-import * as dfns from 'date-fns'
-import Ansi from '@cocalc/ansi-to-react'
 import { ExclamationTriangleIcon } from '@heroicons/react/24/outline'
 import PageButtons from '@/components/PageButtons'
 import Spin from '@/components/Spin'
@@ -9,11 +7,12 @@ import {
   RunQueryResult,
   SuccessRunQueryResult,
   SyntaxErrorRunQueryResult,
+  TableSort,
 } from '@briefer/types'
 import clsx from 'clsx'
 import { useCallback, useEffect, useMemo, useState } from 'react'
 import Table from './Table'
-import { fromPairs, splitEvery } from 'ramda'
+import { fromPairs, splitEvery, map } from 'ramda'
 import useResettableState from '@/hooks/useResettableState'
 import LargeSpinner from '@/components/LargeSpinner'
 import {
@@ -24,6 +23,7 @@ import {
 import { ArrowDownTrayIcon } from '@heroicons/react/24/solid'
 import { Tooltip } from '@/components/Tooltips'
 import { NEXT_PUBLIC_API_URL } from '@/utils/env'
+import qs from 'querystring'
 
 function formatMs(ms: number) {
   if (ms < 1000) {
@@ -50,6 +50,8 @@ interface Props {
   onFixWithAI: () => void
   dashboardMode: 'live' | 'editing' | 'none'
   canFixWithAI: boolean
+  sort: TableSort | null
+  onChangeSort: (sort: TableSort | null) => void
 }
 function SQLResult(props: Props) {
   switch (props.result.type) {
@@ -65,6 +67,8 @@ function SQLResult(props: Props) {
           toggleResultHidden={props.toggleResultHidden}
           blockId={props.blockId}
           dashboardMode={props.dashboardMode}
+          sort={props.sort}
+          onChangeSort={props.onChangeSort}
         />
       )
     case 'abort-error':
@@ -103,6 +107,8 @@ interface SQLSuccessProps {
   isResultHidden: boolean
   toggleResultHidden: () => void
   dashboardMode: 'live' | 'editing' | 'none'
+  sort: TableSort | null
+  onChangeSort: (sort: TableSort | null) => void
 }
 function SQLSuccess(props: SQLSuccessProps) {
   const [currentPageIndex, setCurrentPageIndex] = useState(0)
@@ -119,6 +125,10 @@ function SQLSuccess(props: SQLSuccessProps) {
     [rowsPerPage, props.result.rows]
   )
 
+  useEffect(() => {
+    setPages((pages) => map((page) => ({ ...page, status: 'loading' }), pages))
+  }, [props.sort])
+
   const currentRows = useMemo(() => {
     if (
       pages[currentPageIndex] &&
@@ -154,14 +164,21 @@ function SQLSuccess(props: SQLSuccessProps) {
       return
     }
 
+    const args: Record<string, string | number> = {
+      page: currentPageIndex,
+      pageSize: rowsPerPage,
+      dataframeName: props.dataframeName,
+    }
+
+    if (props.sort) {
+      args['sortColumn'] = props.sort.column
+      args['sortOrder'] = props.sort.order
+    }
+
     fetch(
       `${NEXT_PUBLIC_API_URL()}/v1/workspaces/${props.workspaceId}/documents/${
         props.documentId
-      }/queries/${
-        props.blockId
-      }?page=${currentPageIndex}&pageSize=${rowsPerPage}&dataframeName=${
-        props.dataframeName
-      }`,
+      }/queries/${props.blockId}?${qs.stringify(args)}`,
       {
         credentials: 'include',
       }
@@ -238,6 +255,7 @@ function SQLSuccess(props: SQLSuccessProps) {
     props.documentId,
     props.workspaceId,
     rowsPerPage,
+    props.sort,
   ])
 
   const prevPage = useCallback(() => {
@@ -365,6 +383,8 @@ function SQLSuccess(props: SQLSuccessProps) {
             rows={currentRows}
             columns={props.result.columns}
             isDashboard={props.dashboardMode !== 'none'}
+            sort={props.sort}
+            onChangeSort={props.onChangeSort}
           />
         )}
       </div>

apps/web/src/components/v2Editor/customBlocks/sql/Table.tsx

@@ -4,16 +4,42 @@ import {
   FlagIcon,
   HashtagIcon,
 } from '@heroicons/react/24/outline'
-import { DataFrameColumn, Json } from '@briefer/types'
+import {
+  DataFrameColumn,
+  exhaustiveCheck,
+  Json,
+  TableSort,
+} from '@briefer/types'
 import clsx from 'clsx'
 import ScrollBar from '@/components/ScrollBar'
+import { ArrowDownIcon, ArrowUpIcon } from '@heroicons/react/20/solid'
 
 interface Props {
   rows: Record<string, Json>[]
   columns: DataFrameColumn[]
   isDashboard: boolean
+  sort: TableSort | null
+  onChangeSort: (sort: TableSort | null) => void
 }
 function Table(props: Props) {
+  const onChangeSort = (column: string) => () => {
+    const currentOrder =
+      props.sort && props.sort.column === column ? props.sort.order : null
+    switch (currentOrder) {
+      case 'asc':
+        props.onChangeSort({ column, order: 'desc' })
+        break
+      case 'desc':
+        props.onChangeSort(null)
+        break
+      case null:
+        props.onChangeSort({ column, order: 'asc' })
+        break
+      default:
+        exhaustiveCheck(currentOrder)
+    }
+  }
+
   return (
     <ScrollBar
       className={clsx(
@@ -38,10 +64,24 @@ function Table(props: Props) {
                     'px-2 py-1.5 text-gray-500 whitespace-nowrap font-normal'
                   )}
                 >
-                  <div className="flex space-x-1 items-center">
-                    <Icon className="h-3 w-3 text-gray-400" />
-                    <span>{column.name}</span>
-                  </div>
+                  <button
+                    className="flex space-x-1 items-center w-full justify-between"
+                    onClick={onChangeSort(column.name.toString())}
+                  >
+                    <div className="flex items-center space-x-1">
+                      <Icon className="h-3 w-3 text-gray-400" />
+                      <span>{column.name}</span>
+                    </div>
+                    {props.sort && props.sort.column === column.name && (
+                      <div>
+                        {props.sort.order === 'asc' ? (
+                          <ArrowUpIcon className="h-3 w-3" />
+                        ) : (
+                          <ArrowDownIcon className="h-3 w-3" />
+                        )}
+                      </div>
+                    )}
+                  </button>
                 </th>
               )
             })}

apps/web/src/components/v2Editor/customBlocks/sql/index.tsx

@@ -59,7 +59,11 @@ import { SaveReusableComponentButton } from '@/components/ReusableComponents'
 import { useReusableComponents } from '@/hooks/useReusableComponents'
 import { CodeEditor } from '../../CodeEditor'
 import SQLQueryConfigurationButton from './SQLQueryConfigurationButton'
-import { exhaustiveCheck, SQLQueryConfiguration } from '@briefer/types'
+import {
+  exhaustiveCheck,
+  SQLQueryConfiguration,
+  TableSort,
+} from '@briefer/types'
 import { useBlockExecutions } from '@/hooks/useBlockExecution'
 import { head } from 'ramda'
 import { useAITasks } from '@/hooks/useAITasks'
@@ -130,6 +134,7 @@ function SQLBlock(props: Props) {
     dataSourceId,
     isFileDataSource,
     componentId,
+    sort,
   } = getSQLAttributes(props.block, props.blocks)
 
   const { startedAt: environmentStartedAt } = useEnvironmentStatus(
@@ -438,6 +443,13 @@ function SQLBlock(props: Props) {
     [props.block]
   )
 
+  const onChangeSort = useCallback(
+    (sort: TableSort | null) => {
+      props.block.setAttribute('sort', sort)
+    },
+    [props.block]
+  )
+
   if (props.dashboardMode !== 'none') {
     if (!result) {
       return (
@@ -465,6 +477,8 @@ function SQLBlock(props: Props) {
         onFixWithAI={onFixWithAI}
         dashboardMode={props.dashboardMode}
         canFixWithAI={hasOaiKey}
+        sort={sort}
+        onChangeSort={onChangeSort}
       />
     )
   }
@@ -714,6 +728,8 @@ function SQLBlock(props: Props) {
             onFixWithAI={onFixWithAI}
             dashboardMode={props.dashboardMode}
             canFixWithAI={hasOaiKey}
+            sort={sort}
+            onChangeSort={onChangeSort}
           />
         )}
       </div>

packages/editor/src/blocks/sql.ts

@@ -1,5 +1,9 @@
 import * as Y from 'yjs'
-import { RunQueryResult, SQLQueryConfiguration } from '@briefer/types'
+import {
+  RunQueryResult,
+  SQLQueryConfiguration,
+  TableSort,
+} from '@briefer/types'
 import {
   BlockType,
   BaseBlock,
@@ -33,6 +37,7 @@ export type SQLBlock = BaseBlock<BlockType.SQL> & {
   isEditWithAIPromptOpen: boolean
   aiSuggestions: Y.Text | null
   configuration: SQLQueryConfiguration | null
+  sort: TableSort | null
 
   // wether the block originated from a reusable component and the id of the component
   componentId: string | null
@@ -72,6 +77,7 @@ export const makeSQLBlock = (
     aiSuggestions: null,
     componentId: null,
     configuration: null,
+    sort: null,
   }
 
   for (const [key, value] of Object.entries(attrs)) {
@@ -111,6 +117,7 @@ export function getSQLAttributes(
     aiSuggestions: getSQLAISuggestions(block),
     componentId: getAttributeOr(block, 'componentId', null),
     configuration: getAttributeOr(block, 'configuration', null),
+    sort: getAttributeOr(block, 'sort', null),
   }
 }
 
@@ -153,6 +160,7 @@ export function duplicateSQLBlock(
         : null,
     componentId: options?.componentId ?? prevAttributes.componentId,
     configuration: clone(prevAttributes.configuration),
+    sort: clone(prevAttributes.sort),
   }
 
   const yBlock = new Y.XmlElement<SQLBlock>('block')

packages/types/src/index.ts

@@ -1238,3 +1238,10 @@ export type TutorialState = {
 export type FeatureFlags = {
   visualizationsV2: boolean
 }
+
+export const TableSort = z.object({
+  order: z.union([z.literal('asc'), z.literal('desc')]),
+  column: z.string(),
+})
+
+export type TableSort = z.infer<typeof TableSort>