Allow viewing ongoing queries and aborting them (#26887)

* Allow viewing ongoing queries and aborting them

* Fix

* Add refresh button + shift view running queries CTA

* Fix scrolling in SqlEditorMenu

* Small styling fix for table editor menu schema selector error state

apps/studio/components/interfaces/SQLEditor/OngoingQueriesPanel.tsx

@@ -0,0 +1,168 @@
+import dayjs from 'dayjs'
+import { RefreshCcw, RefreshCw, StopCircle } from 'lucide-react'
+import { useState } from 'react'
+import toast from 'react-hot-toast'
+
+import { useReadReplicasQuery } from 'data/read-replicas/replicas-query'
+import { useQueryAbortMutation } from 'data/sql/abort-query-mutation'
+import { useOngoingQueriesQuery } from 'data/sql/ongoing-queries-query'
+import { useSelectedProject } from 'hooks'
+import { useDatabaseSelectorStateSnapshot } from 'state/database-selector'
+import {
+  Button,
+  CodeBlock,
+  Sheet,
+  SheetContent,
+  SheetDescription,
+  SheetHeader,
+  SheetSection,
+  SheetTitle,
+  TooltipContent_Shadcn_,
+  TooltipTrigger_Shadcn_,
+  Tooltip_Shadcn_,
+  cn,
+} from 'ui'
+import ConfirmationModal from 'ui-patterns/Dialogs/ConfirmationModal'
+import AlertError from 'components/ui/AlertError'
+import { ResponseError } from 'types'
+
+interface OngoingQueriesPanel {
+  visible: boolean
+  onClose: () => void
+}
+
+export const OngoingQueriesPanel = ({ visible, onClose }: OngoingQueriesPanel) => {
+  const project = useSelectedProject()
+  const state = useDatabaseSelectorStateSnapshot()
+  const [selectedId, setSelectedId] = useState<number>()
+
+  const { data: databases } = useReadReplicasQuery({ projectRef: project?.ref })
+  const database = (databases ?? []).find((db) => db.identifier === state.selectedDatabaseId)
+
+  const {
+    data,
+    error,
+    isError,
+    isLoading: isLoadingOngoingQueries,
+    isFetching: isFetchingOngoingQueries,
+    refetch,
+  } = useOngoingQueriesQuery(
+    {
+      projectRef: project?.ref,
+      connectionString: database?.connectionString,
+    },
+    { staleTime: 5000 }
+  )
+  const queries = data ?? []
+
+  const { mutate: abortQuery, isLoading } = useQueryAbortMutation({
+    onSuccess: () => {
+      toast.success(`Successfully aborted query (ID: ${selectedId})`)
+      setSelectedId(undefined)
+    },
+  })
+
+  return (
+    <>
+      <Sheet open={visible} onOpenChange={() => onClose()}>
+        <SheetContent size="lg">
+          <SheetHeader>
+            <SheetTitle className="flex items-center gap-x-2">
+              Running queries on{' '}
+              {database?.identifier === project?.ref ? 'primary database' : 'read replica'}
+              <Button
+                type="default"
+                className="px-1.5"
+                loading={isLoadingOngoingQueries || isFetchingOngoingQueries}
+                icon={<RefreshCw />}
+                onClick={() => refetch()}
+              />
+            </SheetTitle>
+            <SheetDescription>
+              There {queries.length === 1 ? 'is' : 'are'}{' '}
+              <span className="text-foreground-light">{queries.length}</span> quer
+              {queries.length === 1 ? 'y' : 'ies'} currently running{' '}
+              {database?.identifier !== project?.ref ? `on replica ${database?.identifier}` : ''}
+            </SheetDescription>
+          </SheetHeader>
+          <div className="max-h-full h-full divide-y overflow-y-auto">
+            {isError && (
+              <div className="flex items-center justify-center h-full px-16">
+                <AlertError
+                  subject="Failed to retrieve ongoing queries"
+                  error={error as ResponseError}
+                />
+              </div>
+            )}
+            {queries.length === 0 && (
+              <div className="flex flex-col gap-y-2 items-center justify-center h-full text-foreground-light text-sm">
+                <span>No ongoing queries running on database</span>
+                <Button
+                  type="default"
+                  loading={isLoadingOngoingQueries || isFetchingOngoingQueries}
+                  icon={<RefreshCw />}
+                  onClick={() => refetch()}
+                >
+                  Refresh
+                </Button>
+              </div>
+            )}
+            {queries.map((query) => (
+              <SheetSection key={query.pid} className="flex justify-between gap-x-4">
+                <div className="flex flex-col gap-y-2 w-full">
+                  <CodeBlock
+                    hideLineNumbers
+                    value={query.query}
+                    language="sql"
+                    className={cn(
+                      'max-w-none max-h-52 w-full',
+                      '!bg-transparent !py-3 !px-3.5 prose dark:prose-dark',
+                      '[&>code]:m-0 [&>code>span]:flex [&>code>span]:flex-wrap'
+                    )}
+                  />
+                  <div className="flex items-center gap-x-2">
+                    <p className="text-foreground-light text-xs">PID: {query.pid}</p>
+                    <p className="text-foreground-light text-xs">•</p>
+                    <p className="text-foreground-light text-xs">
+                      Started since: {dayjs(query.query_start).format('DD MMM YYYY HH:mm (ZZ)')}
+                    </p>
+                  </div>
+                </div>
+
+                <Tooltip_Shadcn_>
+                  <TooltipTrigger_Shadcn_ asChild>
+                    <Button
+                      type="warning"
+                      className="px-1.5"
+                      icon={<StopCircle />}
+                      onClick={() => setSelectedId(query.pid)}
+                    />
+                  </TooltipTrigger_Shadcn_>
+                  <TooltipContent_Shadcn_ side="bottom">Abort query</TooltipContent_Shadcn_>
+                </Tooltip_Shadcn_>
+              </SheetSection>
+            ))}
+          </div>
+        </SheetContent>
+      </Sheet>
+
+      <ConfirmationModal
+        loading={isLoading}
+        variant="warning"
+        title={`Confirm to abort this query? (ID: ${selectedId})`}
+        visible={selectedId !== undefined}
+        onCancel={() => setSelectedId(undefined)}
+        onConfirm={() => {
+          if (selectedId !== undefined)
+            abortQuery({
+              pid: selectedId,
+              projectRef: project?.ref,
+              connectionString: database?.connectionString,
+            })
+        }}
+      >
+        <p className="text-sm">This will force the query to stop running.</p>
+      </ConfirmationModal>
+    </>
+  )
+}

apps/studio/components/interfaces/SQLEditor/SQLEditor.tsx

@@ -612,6 +612,7 @@ const SQLEditor = () => {
           executeQuery(true)
         }}
       />
+
       <div className="flex h-full">
         <ResizablePanelGroup
           className="h-full relative"
@@ -742,6 +743,7 @@ const SQLEditor = () => {
             )}
           </ResizablePanel>
         </ResizablePanelGroup>
+
         {isAiOpen && (
           <AiAssistantPanel
             messages={messages}

apps/studio/components/layouts/SQLEditorLayout/SQLEditorLayout.tsx

@@ -1,7 +1,8 @@
 import { withAuth } from 'hooks'
-import { ReactNode, useMemo } from 'react'
+import { ReactNode, useMemo, useState } from 'react'
 import { ProjectLayout } from '../'
-import SQLEditorMenu from './SQLEditorMenu'
+import { SQLEditorMenu } from './SQLEditorMenu'
+import { OngoingQueriesPanel } from 'components/interfaces/SQLEditor/OngoingQueriesPanel'
 
 export interface SQLEditorLayoutProps {
   title: string
@@ -9,7 +10,16 @@ export interface SQLEditorLayoutProps {
 }
 
 const SQLEditorLayout = ({ title, children }: SQLEditorLayoutProps) => {
-  const productMenu = useMemo(() => <SQLEditorMenu key="sql-editor-menu" />, [])
+  const [showOngoingQueries, setShowOngoingQueries] = useState(false)
+  const productMenu = useMemo(
+    () => (
+      <SQLEditorMenu
+        key="sql-editor-menu"
+        onViewOngoingQueries={() => setShowOngoingQueries(true)}
+      />
+    ),
+    []
+  )
 
   return (
     <ProjectLayout
@@ -20,6 +30,10 @@ const SQLEditorLayout = ({ title, children }: SQLEditorLayoutProps) => {
       resizableSidebar
     >
       {children}
+      <OngoingQueriesPanel
+        visible={showOngoingQueries}
+        onClose={() => setShowOngoingQueries(false)}
+      />
     </ProjectLayout>
   )
 }

apps/studio/components/layouts/SQLEditorLayout/SQLEditorMenu.tsx

@@ -32,7 +32,7 @@ import {
 import QueryItem from './QueryItem'
 import { selectItemsInRange } from './SQLEditorLayout.utils'
 
-const SideBarContent = () => {
+export const SQLEditorMenu = ({ onViewOngoingQueries }: { onViewOngoingQueries: () => void }) => {
   const { ref, id: activeId } = useParams()
   const router = useRouter()
   const { profile } = useProfile()
@@ -196,16 +196,19 @@ const SideBarContent = () => {
 
   return (
     <>
-      <div className="mt-6">
-        <div className="flex flex-col gap-4">
-          <Button
-            type="default"
-            className="justify-start mx-4"
-            onClick={() => handleNewQuery()}
-            icon={<Plus className="text-foreground-muted" strokeWidth={1} size={14} />}
-          >
-            New query
-          </Button>
+      <div className="h-full flex flex-col justify-between">
+        <div className="mt-4 mb-2 flex flex-col gap-4">
+          <div className="mx-4">
+            <Button
+              block
+              type="default"
+              className="justify-start"
+              onClick={() => handleNewQuery()}
+              icon={<Plus className="text-foreground-muted" strokeWidth={1} size={14} />}
+            >
+              New query
+            </Button>
+          </div>
 
           <div className="px-2">
             <InnerSideMenuItem
@@ -344,6 +347,12 @@ const SideBarContent = () => {
             </>
           )}
         </div>
+
+        <div className="p-4 border-t sticky bottom-0 bg-studio">
+          <Button block type="default" onClick={onViewOngoingQueries}>
+            View running queries
+          </Button>
+        </div>
       </div>
 
       <ConfirmationModal
@@ -369,5 +378,3 @@ const SideBarContent = () => {
     </>
   )
 }
-
-export default SideBarContent

apps/studio/components/layouts/TableEditorLayout/TableEditorMenu.tsx

@@ -1,5 +1,10 @@
 import * as Tooltip from '@radix-ui/react-tooltip'
 import { PermissionAction } from '@supabase/shared-types/out/constants'
+import { partition } from 'lodash'
+import { Plus } from 'lucide-react'
+import { useRouter } from 'next/router'
+import { useMemo, useState } from 'react'
+
 import { useParams } from 'common'
 import { ProtectedSchemaModal } from 'components/interfaces/Database/ProtectedSchemaWarning'
 import AlertError from 'components/ui/AlertError'
@@ -9,10 +14,6 @@ import { useSchemasQuery } from 'data/database/schemas-query'
 import { useEntityTypesQuery } from 'data/entity-types/entity-types-infinite-query'
 import { useCheckPermissions, useLocalStorage } from 'hooks'
 import { EXCLUDED_SCHEMAS } from 'lib/constants/schemas'
-import { partition } from 'lodash'
-import { Plus } from 'lucide-react'
-import { useRouter } from 'next/router'
-import { useMemo, useState } from 'react'
 import { useTableEditorStateSnapshot } from 'state/table-editor'
 import { AlertDescription_Shadcn_, AlertTitle_Shadcn_, Alert_Shadcn_, Button } from 'ui'
 import {
@@ -92,9 +93,9 @@ const TableEditorMenu = () => {
         className="pt-5 flex flex-col flex-grow gap-5 h-full"
         style={{ maxHeight: 'calc(100vh - 48px)' }}
       >
-        <div className="flex flex-col gap-1">
+        <div className="flex flex-col gap-y-1.5">
           <SchemaSelector
-            className="mx-4 h-7"
+            className="mx-4"
             selectedSchemaName={snap.selectedSchemaName}
             onSelectSchema={(name: string) => {
               setSearchText('')

apps/studio/components/ui/SchemaSelector.tsx

@@ -92,7 +92,6 @@ const SchemaSelector = ({
               size={size}
               disabled={disabled}
               type="default"
-              // className="w-full justify-start"
               className={`w-full [&>span]:w-full`}
               iconRight={
                 <ChevronsUpDown className="text-foreground-muted" strokeWidth={2} size={14} />

apps/studio/data/sql/abort-query-mutation.ts

@@ -0,0 +1,49 @@
+import { useMutation, UseMutationOptions, useQueryClient } from '@tanstack/react-query'
+import { toast } from 'react-hot-toast'
+
+import { executeSql } from 'data/sql/execute-sql-query'
+import type { ResponseError } from 'types'
+import { sqlKeys } from './keys'
+
+export type QueryAbortVariables = {
+  pid: number
+  projectRef?: string
+  connectionString?: string
+}
+
+export async function abortQuery({ pid, projectRef, connectionString }: QueryAbortVariables) {
+  const sql = /* SQL */ `select pg_terminate_backend(${pid})`
+  const { result } = await executeSql({ projectRef, connectionString, sql })
+  return result
+}
+
+type QueryAbortData = Awaited<ReturnType<typeof abortQuery>>
+
+export const useQueryAbortMutation = ({
+  onSuccess,
+  onError,
+  ...options
+}: Omit<
+  UseMutationOptions<QueryAbortData, ResponseError, QueryAbortVariables>,
+  'mutationFn'
+> = {}) => {
+  const queryClient = useQueryClient()
+  return useMutation<QueryAbortData, ResponseError, QueryAbortVariables>(
+    (vars) => abortQuery(vars),
+    {
+      async onSuccess(data, variables, context) {
+        const { projectRef } = variables
+        await queryClient.invalidateQueries(sqlKeys.query(projectRef, ['ongoing-queries']))
+        await onSuccess?.(data, variables, context)
+      },
+      async onError(data, variables, context) {
+        if (onError === undefined) {
+          toast.error(`Failed to abort query: ${data.message}`)
+        } else {
+          onError(data, variables, context)
+        }
+      },
+      ...options,
+    }
+  )
+}

apps/studio/data/sql/ongoing-queries-query.ts

@@ -0,0 +1,46 @@
+import { UseQueryOptions } from '@tanstack/react-query'
+import { ExecuteSqlData, ExecuteSqlError, useExecuteSqlQuery } from '../sql/execute-sql-query'
+
+type OngoingQuery = {
+  pid: number
+  query: string
+  query_start: string
+}
+
+export const getOngoingQueries = () => {
+  const sql = /* SQL */ `
+select pid, query, query_start from pg_stat_activity where state = 'active' and datname = 'postgres';
+`.trim()
+
+  return sql
+}
+
+export type OngoingQueriesVariables = {
+  projectRef?: string
+  connectionString?: string
+}
+
+export type OngoingQueriesData = OngoingQuery[]
+export type OngoingQueriesError = ExecuteSqlError
+
+export const useOngoingQueriesQuery = <TData extends OngoingQueriesData = OngoingQueriesData>(
+  { projectRef, connectionString }: OngoingQueriesVariables,
+  options: UseQueryOptions<ExecuteSqlData, OngoingQueriesError, TData> = {}
+) => {
+  return useExecuteSqlQuery(
+    {
+      projectRef,
+      connectionString,
+      sql: getOngoingQueries(),
+      queryKey: ['ongoing-queries'],
+    },
+    {
+      ...options,
+      select(data) {
+        return (data?.result ?? []).filter(
+          (x: OngoingQuery) => !x.query.startsWith(getOngoingQueries())
+        )
+      },
+    }
+  )
+}