supabase/apps/studio/data/database/database-size-query.ts

import { getDatabaseSizeSql } from '@supabase/pg-meta'
import { useQuery } from '@tanstack/react-query'

import { databaseKeys } from './keys'
import { executeSql } from '@/data/sql/execute-sql-mutation'
import { ResponseError, UseCustomQueryOptions } from '@/types'

export type DatabaseSizeVariables = {
  projectRef?: string
  connectionString?: string | null
}

export async function getDatabaseSize(
  { projectRef, connectionString }: DatabaseSizeVariables,
  signal?: AbortSignal
) {
  const sql = getDatabaseSizeSql()

  const { result } = await executeSql(
    {
      projectRef,
      connectionString,
      sql,
      queryKey: ['database-size'],
    },
    signal
  )

  const dbSize = result?.[0]?.db_size
  if (typeof dbSize !== 'number') {
    throw new Error('Error fetching dbSize')
  }

  return dbSize
}

export type DatabaseSizeData = Awaited<ReturnType<typeof getDatabaseSize>>
export type DatabaseSizeError = ResponseError

export const useDatabaseSizeQuery = <TData = DatabaseSizeData>(
  { projectRef, connectionString }: DatabaseSizeVariables,
  {
    enabled = true,
    ...options
  }: UseCustomQueryOptions<DatabaseSizeData, DatabaseSizeError, TData> = {}
) =>
  useQuery<DatabaseSizeData, DatabaseSizeError, TData>({
    queryKey: databaseKeys.databaseSize(projectRef),
    queryFn: ({ signal }) => getDatabaseSize({ projectRef, connectionString }, signal),
    enabled: enabled && typeof projectRef !== 'undefined',
    ...options,
  })