supabase/apps/studio/components/interfaces/Settings/Database/ConnectionPooling/ConnectionPooling.tsx

import { zodResolver } from '@hookform/resolvers/zod'
import { PermissionAction } from '@supabase/shared-types/out/constants'
import { useParams } from 'common'
import { capitalize } from 'lodash'
import Link from 'next/link'
import { Fragment, useEffect } from 'react'
import { SubmitHandler, useForm, useWatch } from 'react-hook-form'
import { toast } from 'sonner'
import {
  Alert,
  AlertDescription,
  AlertTitle,
  Badge,
  Button,
  Form,
  FormControl,
  FormField,
  FormInputGroupInput,
  InputGroup,
  InputGroupAddon,
  InputGroupText,
  Separator,
} from 'ui'
import { Admonition } from 'ui-patterns/Admonition'
import { FormItemLayout } from 'ui-patterns/form/FormItemLayout/FormItemLayout'
import {
  PageSection,
  PageSectionAside,
  PageSectionContent,
  PageSectionMeta,
  PageSectionSummary,
  PageSectionTitle,
} from 'ui-patterns/PageSection'
import { ShimmeringLoader } from 'ui-patterns/ShimmeringLoader'
import z from 'zod'

import { POOLING_OPTIMIZATIONS } from './ConnectionPooling.constants'
import { AlertError } from '@/components/ui/AlertError'
import { DocsButton } from '@/components/ui/DocsButton'
import { FormActions } from '@/components/ui/Forms/FormActions'
import { HighAvailabilityDisabledSectionNotice } from '@/components/ui/HighAvailability/HighAvailabilityDisabledSectionNotice'
import { InlineLink } from '@/components/ui/InlineLink'
import Panel from '@/components/ui/Panel'
import { useMaxConnectionsQuery } from '@/data/database/max-connections-query'
import { usePgbouncerConfigQuery } from '@/data/database/pgbouncer-config-query'
import { usePgbouncerConfigurationUpdateMutation } from '@/data/database/pgbouncer-config-update-mutation'
import { useProjectAddonsQuery } from '@/data/subscriptions/project-addons-query'
import { useCheckEntitlements } from '@/hooks/misc/useCheckEntitlements'
import { useAsyncCheckPermissions } from '@/hooks/misc/useCheckPermissions'
import { useHighAvailability } from '@/hooks/misc/useHighAvailability'
import { useSelectedProjectQuery } from '@/hooks/misc/useSelectedProject'
import { DOCS_URL } from '@/lib/constants'
import { preprocessEmptyNumberInput } from '@/lib/forms/zod-number-input'

const formId = 'pooling-configuration-form'
const HIGH_AVAILABILITY_MAX_CLIENT_CONNECTIONS = 100_000
const HA_DISABLED_TITLE =
  'Connection pooling settings are managed automatically on High Availability projects'

const PoolingConfigurationFormSchema = z.object({
  default_pool_size: preprocessEmptyNumberInput(z.coerce.number().optional()),
  max_client_conn: preprocessEmptyNumberInput(z.coerce.number().optional()),
})

/**
 * [Joshen] PgBouncer configuration will be the main endpoint for GET and PATCH of pooling config
 */
export const ConnectionPooling = () => {
  const { ref: projectRef } = useParams()
  const { data: project } = useSelectedProjectQuery()
  const { isHighAvailability, isPending: isHighAvailabilityPending } = useHighAvailability()
  const canLoadPoolingConfig = !isHighAvailability && !isHighAvailabilityPending
  const { can: canUpdateConnectionPoolingConfiguration } = useAsyncCheckPermissions(
    PermissionAction.UPDATE,
    'projects',
    { resource: { project_id: project?.id } }
  )

  const {
    data: pgbouncerConfig,
    error: pgbouncerConfigError,
    isPending: isLoadingPgbouncerConfig,
    isError: isErrorPgbouncerConfig,
    isSuccess: isSuccessPgbouncerConfig,
  } = usePgbouncerConfigQuery({ projectRef }, { enabled: canLoadPoolingConfig })

  const { hasAccess: hasDedicatedPooler } = useCheckEntitlements('dedicated_pooler')
  const disablePoolModeSelection = !hasDedicatedPooler

  const { data: maxConnData } = useMaxConnectionsQuery(
    {
      projectRef: project?.ref,
      connectionString: project?.connectionString,
    },
    { enabled: canLoadPoolingConfig }
  )
  const { data: addons, isSuccess: isSuccessAddons } = useProjectAddonsQuery({ projectRef })

  const { mutate: updatePoolerConfig, isPending: isUpdatingPoolerConfig } =
    usePgbouncerConfigurationUpdateMutation()

  const hasIpv4Addon = !!addons?.selected_addons.find((addon) => addon.type === 'ipv4')
  const computeInstance = addons?.selected_addons.find((addon) => addon.type === 'compute_instance')
  const computeSize =
    computeInstance?.variant.name ?? capitalize(project?.infra_compute_size) ?? 'Nano'
  const poolingOptimizations =
    POOLING_OPTIMIZATIONS[
      (computeInstance?.variant.identifier as keyof typeof POOLING_OPTIMIZATIONS) ??
        (project?.infra_compute_size === 'nano' ? 'ci_nano' : 'ci_micro')
    ]
  const defaultPoolSize = poolingOptimizations.poolSize ?? 15
  const defaultMaxClientConn = poolingOptimizations.maxClientConn ?? 200

  const form = useForm<z.infer<typeof PoolingConfigurationFormSchema>>({
    resolver: zodResolver(PoolingConfigurationFormSchema),
    defaultValues: {
      default_pool_size: undefined,
      max_client_conn: undefined,
    },
  })
  const default_pool_size = useWatch({ control: form.control, name: 'default_pool_size' })
  const connectionPoolingUnavailable = pgbouncerConfig?.pool_mode === null
  const ignoreStartupParameters = pgbouncerConfig?.ignore_startup_parameters

  const onSubmit: SubmitHandler<z.infer<typeof PoolingConfigurationFormSchema>> = async (data) => {
    const { default_pool_size } = data

    if (!projectRef || isHighAvailability) return

    updatePoolerConfig(
      {
        ref: projectRef,
        default_pool_size: default_pool_size === null ? undefined : default_pool_size,
        ignore_startup_parameters: ignoreStartupParameters ?? '',
      },
      {
        onSuccess: (data) => {
          toast.success(`Successfully updated pooler configuration`)
          if (data) {
            form.reset({
              default_pool_size: data.default_pool_size,
            })
          }
        },
      }
    )
  }

  const resetForm = () => {
    form.reset({
      default_pool_size: pgbouncerConfig?.default_pool_size ?? defaultPoolSize,
      max_client_conn: pgbouncerConfig?.max_client_conn ?? defaultMaxClientConn,
    })
  }

  useEffect(() => {
    if (isSuccessPgbouncerConfig) resetForm()
  }, [isSuccessPgbouncerConfig])

  return (
    <PageSection id="connection-pooler">
      <PageSectionMeta>
        <PageSectionSummary>
          <PageSectionTitle>Connection pooling</PageSectionTitle>
        </PageSectionSummary>
        <PageSectionAside>
          <DocsButton
            href={`${DOCS_URL}/guides/database/connecting-to-postgres/pooling-and-limits#how-connection-pooling-works`}
          />
        </PageSectionAside>
      </PageSectionMeta>
      <PageSectionContent className="space-y-4">
        {isHighAvailability && (
          <HighAvailabilityDisabledSectionNotice
            title={HA_DISABLED_TITLE}
            description={
              <>
                High Availability projects run one pooler per Postgres pod, each supporting up to{' '}
                {HIGH_AVAILABILITY_MAX_CLIENT_CONNECTIONS.toLocaleString()} active or passive
                connections, with pooling behavior selected automatically.{' '}
                <InlineLink href="https://multigres.com/blog/pooling-without-choosing-a-mode">
                  Learn more
                </InlineLink>
                .
              </>
            }
          />
        )}

        {isSuccessAddons && !isHighAvailability && !disablePoolModeSelection && !hasIpv4Addon && (
          <Admonition
            type="default"
            layout="responsive"
            title="Dedicated pooler uses IPv6 by default"
            description="Connections from IPv4-only networks require enabling the IPv4 add-on on your project instance."
            actions={
              <Button variant="default" asChild>
                <Link href={`/project/${projectRef}/settings/addons?panel=ipv4`}>
                  Enable IPv4 add-on
                </Link>
              </Button>
            }
          />
        )}

        <Panel
          noMargin
          footer={
            isHighAvailability ? undefined : (
              <FormActions
                form={formId}
                isSubmitting={isUpdatingPoolerConfig}
                hasChanges={form.formState.isDirty}
                handleReset={() => resetForm()}
                helper={
                  !canUpdateConnectionPoolingConfiguration
                    ? 'You need additional permissions to update connection pooling settings'
                    : undefined
                }
              />
            )
          }
        >
          <Panel.Content>
            {!isHighAvailability && isLoadingPgbouncerConfig && (
              <div className="flex flex-col gap-y-4">
                {Array.from({ length: 4 }).map((_, i) => (
                  <Fragment key={`loader-${i}`}>
                    <div className="grid gap-2 items-center md:grid md:grid-cols-12 md:gap-x-4 w-full">
                      <ShimmeringLoader className="h-4 w-1/3 col-span-4" delayIndex={i} />
                      <ShimmeringLoader className="h-8 w-full col-span-8" delayIndex={i} />
                    </div>
                    <Separator />
                  </Fragment>
                ))}

                <ShimmeringLoader className="h-8 w-full" />
              </div>
            )}
            {!isHighAvailability && isErrorPgbouncerConfig && (
              <AlertError
                error={pgbouncerConfigError}
                subject="Failed to retrieve connection pooler configuration"
              />
            )}
            {!isHighAvailability && connectionPoolingUnavailable && (
              <Admonition
                type="default"
                title="Unable to retrieve pooling configuration"
                description="Please start a new project to enable this feature"
              />
            )}
            {(isHighAvailability ||
              (isSuccessPgbouncerConfig && !connectionPoolingUnavailable)) && (
              <>
                <div className="flex flex-row gap-2 justify-between w-full">
                  <div className="flex flex-col text-sm">
                    <h5 className="text-foreground font-normal">Connection poolers</h5>
                    <p className="text-foreground-lighter">
                      {isHighAvailability
                        ? 'One pooler runs for each Postgres pod in the cluster.'
                        : 'Configuration is shared across all connection poolers.'}
                    </p>
                  </div>
                  <div className="flex flex-row gap-1 items-center">
                    {isHighAvailability ? (
                      <Badge>High Availability</Badge>
                    ) : (
                      <>
                        <Badge>Shared</Badge>
                        {!disablePoolModeSelection && <Badge>Dedicated</Badge>}
                      </>
                    )}
                  </div>
                </div>
                <Separator className="bg-border -mx-6 w-[calc(100%+3rem)] my-4" />
                <Form {...form}>
                  <form
                    id={formId}
                    className="flex flex-col gap-y-4 w-full"
                    onSubmit={form.handleSubmit(onSubmit)}
                  >
                    <FormField
                      control={form.control}
                      name="default_pool_size"
                      render={({ field }) => (
                        <FormItemLayout
                          layout="flex-row-reverse"
                          label="Connection pool size"
                          description={
                            isHighAvailability ? (
                              <p>
                                Pool size is managed automatically for each Postgres pod and cannot
                                be changed.
                              </p>
                            ) : (
                              <p>
                                The maximum number of connections made to the underlying Postgres
                                cluster, per user+db combination. Pool size has a default of{' '}
                                {defaultPoolSize} based on your compute size of {computeSize}.
                              </p>
                            )
                          }
                          className="[&>div]:md:w-1/2 [&>div]:xl:w-2/5 [&>div>div]:w-full"
                        >
                          <FormControl>
                            <InputGroup>
                              <FormInputGroupInput
                                {...field}
                                disabled={isHighAvailability}
                                type="number"
                                className="w-full"
                                value={isHighAvailability ? '' : (field.value ?? '')}
                                placeholder={
                                  isHighAvailability
                                    ? 'Managed automatically'
                                    : defaultPoolSize.toString()
                                }
                                onChange={(event) =>
                                  field.onChange(
                                    isNaN(event.target.valueAsNumber)
                                      ? null
                                      : event.target.valueAsNumber
                                  )
                                }
                              />
                              {!isHighAvailability && (
                                <InputGroupAddon align="inline-end">
                                  <InputGroupText>connections</InputGroupText>
                                </InputGroupAddon>
                              )}
                            </InputGroup>
                          </FormControl>
                          {!isHighAvailability &&
                            !!maxConnData &&
                            (default_pool_size ?? 15) > maxConnData.maxConnections * 0.8 && (
                              <Alert variant="warning" className="mt-2">
                                <AlertTitle className="text-foreground">
                                  Pool size is greater than 80% of the max connections (
                                  {maxConnData.maxConnections}) on your database
                                </AlertTitle>
                                <AlertDescription>
                                  This may result in instability and unreliability with your
                                  database connections.
                                </AlertDescription>
                              </Alert>
                            )}
                        </FormItemLayout>
                      )}
                    />

                    <Separator className="bg-border -mx-6 w-[calc(100%+3rem)]" />

                    <FormField
                      control={form.control}
                      disabled
                      name="max_client_conn"
                      render={({ field }) => (
                        <FormItemLayout
                          layout="flex-row-reverse"
                          label="Max client connections"
                          className="[&>div]:md:w-1/2 [&>div]:xl:w-2/5 [&>div>div]:w-full"
                          description={
                            isHighAvailability ? (
                              <p>
                                Each pooler can support up to{' '}
                                {HIGH_AVAILABILITY_MAX_CLIENT_CONNECTIONS.toLocaleString()} active
                                or passive client connections. This value is managed automatically
                                and cannot be changed.
                              </p>
                            ) : (
                              <p>
                                The maximum number of concurrent client connections allowed. This
                                value is fixed at {defaultMaxClientConn} based on your compute size
                                of {computeSize} and cannot be changed.{' '}
                                <InlineLink
                                  href={`${DOCS_URL}/guides/database/connection-management#configuring-supavisors-pool-size`}
                                >
                                  Learn more
                                </InlineLink>
                              </p>
                            )
                          }
                        >
                          <FormControl>
                            <InputGroup>
                              <FormInputGroupInput
                                {...field}
                                type="number"
                                className="w-full"
                                value={
                                  isHighAvailability
                                    ? HIGH_AVAILABILITY_MAX_CLIENT_CONNECTIONS
                                    : (pgbouncerConfig?.max_client_conn ?? '')
                                }
                                placeholder={
                                  isHighAvailability
                                    ? HIGH_AVAILABILITY_MAX_CLIENT_CONNECTIONS.toString()
                                    : defaultMaxClientConn.toString()
                                }
                                onChange={(event) =>
                                  field.onChange(
                                    isNaN(event.target.valueAsNumber)
                                      ? null
                                      : event.target.valueAsNumber
                                  )
                                }
                              />
                              <InputGroupAddon align="inline-end">
                                <InputGroupText>clients</InputGroupText>
                              </InputGroupAddon>
                            </InputGroup>
                          </FormControl>
                        </FormItemLayout>
                      )}
                    />
                  </form>
                </Form>
              </>
            )}
          </Panel.Content>
        </Panel>
      </PageSectionContent>
    </PageSection>
  )
}