Support for Dedicated Pooler in Connection Pooling (#33817)

* Init

* Initial set up for hooking up supavisor and pgbouncer

* Hook up pgbouncer status check after swapping pooler type

* Add check for nano compute for switching to pg bouncer

* Add check for ipv4 addon

* Remove expect error tag

* Add badge to select options for pooler types

* Remove statement mode

* Resolve undefined problem with react hook form

* Fix

* Update UI texts from PgBouncer to Dedicated Pooler

* Feex

* FEEX

* Fix

* Small update to UI

* Smol update

apps/studio/components/interfaces/DiskManagement/fields/ComputeSizeField.tsx

@@ -5,10 +5,12 @@ import { UseFormReturn } from 'react-hook-form'
 import { components } from 'api-types'
 import { useParams } from 'common'
 import { useProjectContext } from 'components/layouts/ProjectLayout/ProjectContext'
+import { DocsButton } from 'components/ui/DocsButton'
 import { useOrgSubscriptionQuery } from 'data/subscriptions/org-subscription-query'
 import { useProjectAddonsQuery } from 'data/subscriptions/project-addons-query'
 import { useSelectedOrganization } from 'hooks/misc/useSelectedOrganization'
 import { getCloudProviderArchitecture } from 'lib/cloudprovider-utils'
+import { InstanceSpecs } from 'lib/constants'
 import { cn, FormField_Shadcn_, RadioGroupCard, RadioGroupCardItem, Skeleton } from 'ui'
 import { ComputeBadge } from 'ui-patterns'
 import { FormItemLayout } from 'ui-patterns/form/FormItemLayout/FormItemLayout'
@@ -22,8 +24,6 @@ import {
 import { BillingChangeBadge } from '../ui/BillingChangeBadge'
 import FormMessage from '../ui/FormMessage'
 import { NoticeBar } from '../ui/NoticeBar'
-import { InstanceSpecs } from 'lib/constants'
-import { DocsButton } from 'components/ui/DocsButton'
 
 /**
  * to do: this could be a type from api-types
@@ -196,7 +196,6 @@ export function ComputeSizeField({ form, disabled }: ComputeSizeFieldProps) {
                         lockedOption && 'opacity-50'
                       )}
                       disabled={disabled || lockedOption}
-                      // @ts-expect-error
                       label={
                         <>
                           {showUpgradeBadge && compute.identifier === 'ci_micro' && (

apps/studio/components/interfaces/DiskManagement/ui/DiskManagement.constants.tsx

@@ -57,11 +57,6 @@ export const DISK_LIMITS = {
   },
 }
 
-export const DISK_TYPE_LABELS = {
-  [DiskType.GP3]: 'General Purpose SSD (gp3)',
-  [DiskType.IO2]: 'Provisioned IOPS SSD (io2)',
-}
-
 interface PlanDetails {
   includedDiskGB: { gp3: number; io2: number }
 }

apps/studio/components/interfaces/Settings/Database/ConnectionPooling/ConnectionPooling.constants.ts

@@ -1,5 +1,7 @@
 // https://supabase.com/docs/guides/platform/performance#optimizing-the-number-of-connections
 // https://github.com/supabase/infrastructure/blob/develop/worker/src/lib/constants.ts#L544-L596
+// https://github.com/supabase/supabase-admin-api/blob/master/optimizations/pgbouncer.go
+// [Joshen] This matches for both Supavisor and PgBouncer
 
 export const POOLING_OPTIMIZATIONS = {
   ci_nano: {

apps/studio/components/ui/Forms/Form.constants.ts

@@ -1,19 +1,5 @@
 import * as z from 'zod'
 
-// This validator validates a string to be a positive integer or if it's an empty string, transforms it to a null
-export const StringToPositiveNumber = z.union([
-  // parse the value if it's a number
-  z.number().positive().int(),
-  // parse the value if it's a non-empty string
-  z.string().min(1).pipe(z.coerce.number().positive().int()),
-  // transform a non-empty string into a null value
-  z
-    .string()
-    .max(0, 'The field accepts only a number')
-    .transform((v) => null),
-  z.null(),
-])
-
 export const StringNumberOrNull = z
   .string()
   .transform((v) => (v === '' ? null : v))
@@ -22,3 +8,10 @@ export const StringNumberOrNull = z
     message: 'Invalid number',
   })
   .transform((value) => (value === null ? null : Number(value)))
+
+/**
+ * [Joshen] After wrangling with RHF I think this is the easiest way to handle nullable number fields
+ * - Declare the field normally as you would in the zod form schema (e.g field: z.number().nullable())
+ * - In the InputField, add a form.register call `{...form.register('field_name', { setValueAs: setValueAsNullableNumber })}`
+ */
+export const setValueAsNullableNumber = (v: any) => (v === '' || v === null ? null : parseInt(v))

apps/studio/components/ui/Panel.tsx

@@ -36,11 +36,7 @@ function Panel(props: PropsWithChildren<PanelProps>) {
         </div>
       )}
       {props.children}
-      {props.footer && (
-        <div className="bg-surface-100 border-t border-default">
-          <div className="flex h-12 items-center px-4 md:px-6">{props.footer}</div>
-        </div>
-      )}
+      {props.footer && <Footer>{props.footer}</Footer>}
     </div>
   )
 
@@ -55,6 +51,14 @@ function Content({ children, className }: { children: ReactNode; className?: str
   return <div className={cn('px-4 md:px-6 py-4', className)}>{children}</div>
 }
 
+function Footer({ children }: { children: ReactNode; className?: string }) {
+  return (
+    <div className="bg-surface-100 border-t border-default">
+      <div className="flex h-12 items-center px-4 md:px-6">{children}</div>
+    </div>
+  )
+}
+
 const PanelNotice = forwardRef<
   HTMLDivElement,
   {
@@ -139,5 +143,6 @@ const PanelNotice = forwardRef<
 PanelNotice.displayName = 'PanelNotice'
 
 Panel.Content = Content
+Panel.Footer = Footer
 Panel.Notice = PanelNotice
 export default Panel

apps/studio/data/database/pgbouncer-config-query.ts

@@ -4,12 +4,12 @@ import { get, handleError } from 'data/fetchers'
 import { ResponseError } from 'types'
 import { databaseKeys } from './keys'
 
-export type ProjectPgbouncerConfigVariables = {
+export type PgbouncerConfigVariables = {
   projectRef?: string
 }
 
-export async function getProjectPgbouncerConfig(
-  { projectRef }: ProjectPgbouncerConfigVariables,
+export async function getPgbouncerConfig(
+  { projectRef }: PgbouncerConfigVariables,
   signal?: AbortSignal
 ) {
   if (!projectRef) throw new Error('projectRef is required')
@@ -22,19 +22,19 @@ export async function getProjectPgbouncerConfig(
   return data
 }
 
-export type ProjectPgbouncerConfigData = Awaited<ReturnType<typeof getProjectPgbouncerConfig>>
-export type ProjectPgbouncerConfigError = ResponseError
+export type PgbouncerConfigData = Awaited<ReturnType<typeof getPgbouncerConfig>>
+export type PgbouncerConfigError = ResponseError
 
-export const useProjectPgbouncerConfigQuery = <TData = ProjectPgbouncerConfigData>(
-  { projectRef }: ProjectPgbouncerConfigVariables,
+export const usePgbouncerConfigQuery = <TData = PgbouncerConfigData>(
+  { projectRef }: PgbouncerConfigVariables,
   {
     enabled = true,
     ...options
-  }: UseQueryOptions<ProjectPgbouncerConfigData, ProjectPgbouncerConfigError, TData> = {}
+  }: UseQueryOptions<PgbouncerConfigData, PgbouncerConfigError, TData> = {}
 ) =>
-  useQuery<ProjectPgbouncerConfigData, ProjectPgbouncerConfigError, TData>(
+  useQuery<PgbouncerConfigData, PgbouncerConfigError, TData>(
     databaseKeys.pgbouncerConfig(projectRef),
-    ({ signal }) => getProjectPgbouncerConfig({ projectRef }, signal),
+    ({ signal }) => getPgbouncerConfig({ projectRef }, signal),
     {
       enabled: enabled && typeof projectRef !== 'undefined',
       ...options,

apps/studio/data/database/pgbouncer-config-update-mutation.ts

@@ -0,0 +1,73 @@
+import { useMutation, UseMutationOptions, useQueryClient } from '@tanstack/react-query'
+import { toast } from 'sonner'
+
+import type { components } from 'data/api'
+import { handleError, patch } from 'data/fetchers'
+import type { ResponseError } from 'types'
+import { databaseKeys } from './keys'
+
+export type PgbouncerConfigurationUpdateVariables = {
+  ref: string
+} & components['schemas']['UpdatePgbouncerConfigBody']
+
+export async function updatePgbouncerConfiguration({
+  ref,
+  pool_mode,
+  default_pool_size,
+  pgbouncer_enabled,
+  ignore_startup_parameters,
+  max_client_conn,
+}: PgbouncerConfigurationUpdateVariables) {
+  if (!ref) return console.error('Project ref is required')
+
+  const { data, error } = await patch('/platform/projects/{ref}/config/pgbouncer', {
+    params: { path: { ref } },
+    body: {
+      default_pool_size,
+      pool_mode,
+      pgbouncer_enabled,
+      ignore_startup_parameters,
+      max_client_conn,
+    },
+  })
+
+  if (error) handleError(error)
+  return data
+}
+
+type PgbouncerConfigurationUpdateData = Awaited<ReturnType<typeof updatePgbouncerConfiguration>>
+
+export const usePgbouncerConfigurationUpdateMutation = ({
+  onSuccess,
+  onError,
+  ...options
+}: Omit<
+  UseMutationOptions<
+    PgbouncerConfigurationUpdateData,
+    ResponseError,
+    PgbouncerConfigurationUpdateVariables
+  >,
+  'mutationFn'
+> = {}) => {
+  const queryClient = useQueryClient()
+
+  return useMutation<
+    PgbouncerConfigurationUpdateData,
+    ResponseError,
+    PgbouncerConfigurationUpdateVariables
+  >((vars) => updatePgbouncerConfiguration(vars), {
+    async onSuccess(data, variables, context) {
+      const { ref } = variables
+      await queryClient.invalidateQueries(databaseKeys.pgbouncerConfig(ref))
+      await onSuccess?.(data, variables, context)
+    },
+    async onError(data, variables, context) {
+      if (onError === undefined) {
+        toast.error(`Failed to update PgBouncer configuration: ${data.message}`)
+      } else {
+        onError(data, variables, context)
+      }
+    },
+    ...options,
+  })
+}

apps/studio/data/database/pgbouncer-status-query.ts

@@ -4,12 +4,12 @@ import { get, handleError } from 'data/fetchers'
 import { ResponseError } from 'types'
 import { databaseKeys } from './keys'
 
-export type ProjectPgbouncerStatusVariables = {
+export type PgbouncerStatusVariables = {
   projectRef?: string
 }
 
-export async function getProjectPgbouncerStatus(
-  { projectRef }: ProjectPgbouncerStatusVariables,
+export async function getPgbouncerStatus(
+  { projectRef }: PgbouncerStatusVariables,
   signal?: AbortSignal
 ) {
   if (!projectRef) throw new Error('projectRef is required')
@@ -22,19 +22,19 @@ export async function getProjectPgbouncerStatus(
   return data
 }
 
-export type ProjectPgbouncerStatusData = Awaited<ReturnType<typeof getProjectPgbouncerStatus>>
-export type ProjectPgbouncerStatusError = ResponseError
+export type PgbouncerStatusData = Awaited<ReturnType<typeof getPgbouncerStatus>>
+export type PgbouncerStatusError = ResponseError
 
-export const useProjectPgbouncerStatusQuery = <TData = ProjectPgbouncerStatusData>(
-  { projectRef }: ProjectPgbouncerStatusVariables,
+export const usePgbouncerStatusQuery = <TData = PgbouncerStatusData>(
+  { projectRef }: PgbouncerStatusVariables,
   {
     enabled = true,
     ...options
-  }: UseQueryOptions<ProjectPgbouncerStatusData, ProjectPgbouncerStatusError, TData> = {}
+  }: UseQueryOptions<PgbouncerStatusData, PgbouncerStatusError, TData> = {}
 ) =>
-  useQuery<ProjectPgbouncerStatusData, ProjectPgbouncerStatusError, TData>(
+  useQuery<PgbouncerStatusData, PgbouncerStatusError, TData>(
     databaseKeys.pgbouncerStatus(projectRef),
-    ({ signal }) => getProjectPgbouncerStatus({ projectRef }, signal),
+    ({ signal }) => getPgbouncerStatus({ projectRef }, signal),
     {
       enabled: enabled && typeof projectRef !== 'undefined',
       ...options,

apps/studio/data/database/pooling-configuration-query.ts

@@ -31,6 +31,9 @@ export async function getPoolingConfiguration(
 export type PoolingConfigurationData = Awaited<ReturnType<typeof getPoolingConfiguration>>
 export type PoolingConfigurationError = ResponseError
 
+/**
+ * @deprecated use useSupavisorConfigurationQuery isntead
+ */
 export const usePoolingConfigurationQuery = <TData = PoolingConfigurationData>(
   { projectRef }: PoolingConfigurationVariables,
   {
@@ -46,3 +49,22 @@ export const usePoolingConfigurationQuery = <TData = PoolingConfigurationData>(
       ...options,
     }
   )
+
+/**
+ * Just a duplicate of usePoolingConfigurationQuery until we move everything over
+ */
+export const useSupavisorConfigurationQuery = <TData = PoolingConfigurationData>(
+  { projectRef }: PoolingConfigurationVariables,
+  {
+    enabled = true,
+    ...options
+  }: UseQueryOptions<PoolingConfigurationData, PoolingConfigurationError, TData> = {}
+) =>
+  useQuery<PoolingConfigurationData, PoolingConfigurationError, TData>(
+    databaseKeys.poolingConfiguration(projectRef),
+    ({ signal }) => getPoolingConfiguration({ projectRef }, signal),
+    {
+      enabled: enabled && typeof projectRef !== 'undefined',
+      ...options,
+    }
+  )

apps/studio/data/database/pooling-configuration-update-mutation.ts

@@ -57,7 +57,42 @@ export const usePoolingConfigurationUpdateMutation = ({
     },
     async onError(data, variables, context) {
       if (onError === undefined) {
-        toast.error(`Failed to update pooling configuration: ${data.message}`)
+        toast.error(`Failed to update PgBouncer configuration: ${data.message}`)
+      } else {
+        onError(data, variables, context)
+      }
+    },
+    ...options,
+  })
+}
+
+export const useSupavisorConfigurationUpdateMutation = ({
+  onSuccess,
+  onError,
+  ...options
+}: Omit<
+  UseMutationOptions<
+    PoolingConfigurationUpdateData,
+    ResponseError,
+    PoolingConfigurationUpdateVariables
+  >,
+  'mutationFn'
+> = {}) => {
+  const queryClient = useQueryClient()
+
+  return useMutation<
+    PoolingConfigurationUpdateData,
+    ResponseError,
+    PoolingConfigurationUpdateVariables
+  >((vars) => updatePoolingConfiguration(vars), {
+    async onSuccess(data, variables, context) {
+      const { ref } = variables
+      await queryClient.invalidateQueries(databaseKeys.poolingConfiguration(ref))
+      await onSuccess?.(data, variables, context)
+    },
+    async onError(data, variables, context) {
+      if (onError === undefined) {
+        toast.error(`Failed to update Supavisor configuration: ${data.message}`)
       } else {
         onError(data, variables, context)
       }

packages/ui/src/components/radio-group-card.tsx

@@ -16,7 +16,7 @@ RadioGroupCard.displayName = RadioGroupPrimitive.Root.displayName
 
 interface RadioGroupCardItemProps {
   image?: React.ReactNode
-  label: string
+  label: string | React.ReactNode
   showIndicator?: boolean
 }