Fix: Updated to use optimized search columns when getting user counts (supabase#40107)

* Updated to use optimized search columns when getting user counts

* added unit tests for query builders

* minor

* Nit fix keywords URL param not getting loaded into search input field

* removed user footer

* updated tests for upstream

* updated integration tests

* updated tests to be int for paginated

* updated schema

* updated types

* updated type imports

* updated type imports

* Apply suggestion from @avallete

Co-authored-by: Andrew Valleteau <[email protected]>

* updated literal import

* refactor: use common escaping (supabase#40186)

refactor: use utils for sql escaping

---------

Co-authored-by: Joshen Lim <[email protected]>
Co-authored-by: Andrew Valleteau <[email protected]>

Author: 3 people

apps/studio/components/interfaces/Auth/Users/Users.constants.ts

@@ -1,8 +1,11 @@
 import { BASE_PATH } from 'lib/constants'
 import { PROVIDER_PHONE, PROVIDERS_SCHEMAS } from '../AuthProvidersFormValidation'
+import { OptimizedSearchColumns } from '@supabase/pg-meta/src/sql/studio/get-users-types'
 
 export type Filter = 'all' | 'verified' | 'unverified' | 'anonymous'
 
+export type SpecificFilterColumn = OptimizedSearchColumns | 'freeform'
+
 export const UUIDV4_LEFT_PREFIX_REGEX =
   /^(?:[0-9a-f]{1,8}|[0-9a-f]{8}-|[0-9a-f]{8}-[0-9a-f]{1,4}|[0-9a-f]{8}-[0-9a-f]{4}-|[0-9a-f]{8}-[0-9a-f]{4}-4[0-9a-f]{0,3}|[0-9a-f]{8}-[0-9a-f]{4}-4[0-9a-f]{3}-|[0-9a-f]{8}-[0-9a-f]{4}-4[0-9a-f]{3}-[89ab][0-9a-f]{0,3}|[0-9a-f]{8}-[0-9a-f]{4}-4[0-9a-f]{3}-[89ab][0-9a-f]{3}-|[0-9a-f]{8}-[0-9a-f]{4}-4[0-9a-f]{3}-[89ab][0-9a-f]{3}-[0-9a-f]{0,12})$/i
 

apps/studio/components/interfaces/Auth/Users/UsersFooter.tsx

@@ -2,19 +2,21 @@ import { THRESHOLD_COUNT } from '@supabase/pg-meta/src/sql/studio/get-count-esti
 import { HelpCircle, Loader2 } from 'lucide-react'
 import { useEffect, useState } from 'react'
 
+import type { Filter, SpecificFilterColumn } from './Users.constants'
+
+import { OptimizedSearchColumns } from '@supabase/pg-meta/src/sql/studio/get-users-types'
 import { useParams } from 'common'
 import { formatEstimatedCount } from 'components/grid/components/footer/pagination/Pagination.utils'
 import { useUsersCountQuery } from 'data/auth/users-count-query'
 import { useSelectedProjectQuery } from 'hooks/misc/useSelectedProject'
 import { Button, Tooltip, TooltipContent, TooltipTrigger } from 'ui'
 import ConfirmationModal from 'ui-patterns/Dialogs/ConfirmationModal'
-import { Filter } from './Users.constants'
 
 interface UsersFooterProps {
   filter: Filter
   filterKeywords: string
   selectedProviders: string[]
-  specificFilterColumn: string
+  specificFilterColumn: SpecificFilterColumn
 }
 
 export const UsersFooter = ({
@@ -42,6 +44,10 @@ export const UsersFooter = ({
       filter: filter === 'all' ? undefined : filter,
       providers: selectedProviders,
       forceExactCount,
+      // Use optimized search when filtering by specific column
+      ...(specificFilterColumn !== 'freeform'
+        ? { column: specificFilterColumn as OptimizedSearchColumns }
+        : { column: undefined }),
     },
     { keepPreviousData: true }
   )

apps/studio/components/interfaces/Auth/Users/UsersV2.tsx

@@ -5,6 +5,8 @@ import { UIEvent, useEffect, useMemo, useRef, useState } from 'react'
 import DataGrid, { Column, DataGridHandle, Row } from 'react-data-grid'
 import { toast } from 'sonner'
 
+import type { OptimizedSearchColumns } from '@supabase/pg-meta/src/sql/studio/get-users-types'
+import type { SpecificFilterColumn } from './Users.constants'
 import { LOCAL_STORAGE_KEYS, useParams } from 'common'
 import { useIsAPIDocsSidePanelEnabled } from 'components/interfaces/App/FeaturePreview/FeaturePreviewContext'
 import AlertError from 'components/ui/AlertError'
@@ -91,9 +93,9 @@ export const UsersV2 = () => {
     }
   }, [showEmailPhoneColumns])
 
-  const [specificFilterColumn, setSpecificFilterColumn] = useQueryState(
+  const [specificFilterColumn, setSpecificFilterColumn] = useQueryState<SpecificFilterColumn>(
     'filter',
-    parseAsStringEnum(['id', 'email', 'phone', 'freeform']).withDefault('id')
+    parseAsStringEnum<SpecificFilterColumn>(['id', 'email', 'phone', 'freeform']).withDefault('id')
   )
   const [filterUserType, setFilterUserType] = useQueryState(
     'userType',
@@ -138,7 +140,7 @@ export const UsersV2 = () => {
   )
 
   const [columns, setColumns] = useState<Column<any>[]>([])
-  const [search, setSearch] = useState('')
+  const [search, setSearch] = useState(filterKeywords)
   const [selectedUser, setSelectedUser] = useState<string>()
   const [selectedUsers, setSelectedUsers] = useState<Set<any>>(new Set([]))
   const [selectedUserToDelete, setSelectedUserToDelete] = useState<User>()
@@ -187,7 +189,7 @@ export const UsersV2 = () => {
       sort: sortColumn as 'id' | 'created_at' | 'email' | 'phone',
       order: sortOrder as 'asc' | 'desc',
       ...(specificFilterColumn !== 'freeform'
-        ? { column: specificFilterColumn }
+        ? { column: specificFilterColumn as OptimizedSearchColumns }
         : { column: undefined }),
     },
     {

apps/studio/data/auth/keys.ts

@@ -1,3 +1,5 @@
+import type { OptimizedSearchColumns } from '@supabase/pg-meta/src/sql/studio/get-users-types'
+
 export const authKeys = {
   users: (
     projectRef: string | undefined,
@@ -39,6 +41,7 @@ export const authKeys = {
       filter: string | undefined
       providers: string[] | undefined
       forceExactCount?: boolean
+      column?: OptimizedSearchColumns
     }
   ) =>
     ['projects', projectRef, 'users-count', ...(params ? [params].filter(Boolean) : [])] as const,

apps/studio/data/auth/users-count-query.ts

@@ -1,3 +1,4 @@
+import type { OptimizedSearchColumns } from '@supabase/pg-meta/src/sql/studio/get-users-types'
 import { getUsersCountSQL } from '@supabase/pg-meta/src/sql/studio/get-users-count'
 import { useQuery } from '@tanstack/react-query'
 
@@ -13,6 +14,9 @@ type UsersCountVariables = {
   filter?: Filter
   providers?: string[]
   forceExactCount?: boolean
+
+  /** If set, uses optimized prefix search for the specified column */
+  column?: OptimizedSearchColumns
 }
 
 export async function getUsersCount(
@@ -23,10 +27,11 @@ export async function getUsersCount(
     filter,
     providers,
     forceExactCount,
+    column,
   }: UsersCountVariables,
   signal?: AbortSignal
 ) {
-  const sql = getUsersCountSQL({ filter, keywords, providers, forceExactCount })
+  const sql = getUsersCountSQL({ filter, keywords, providers, forceExactCount, column })
 
   const { result } = await executeSql(
     {
@@ -63,6 +68,7 @@ export const useUsersCountQuery = <TData = UsersCountData>(
     filter,
     providers,
     forceExactCount,
+    column,
   }: UsersCountVariables,
   { enabled = true, ...options }: UseCustomQueryOptions<UsersCountData, UsersCountError, TData> = {}
 ) =>
@@ -72,6 +78,7 @@ export const useUsersCountQuery = <TData = UsersCountData>(
       filter,
       providers,
       forceExactCount,
+      column,
     }),
     queryFn: ({ signal }) =>
       getUsersCount(
@@ -82,6 +89,7 @@ export const useUsersCountQuery = <TData = UsersCountData>(
           filter,
           providers,
           forceExactCount,
+          column,
         },
         signal
       ),

apps/studio/data/auth/users-infinite-query.ts

@@ -1,6 +1,7 @@
 import { getPaginatedUsersSQL } from '@supabase/pg-meta/src/sql/studio/get-users-paginated'
 import { useInfiniteQuery } from '@tanstack/react-query'
 
+import { OptimizedSearchColumns } from '@supabase/pg-meta/src/sql/studio/get-users-types'
 import type { components } from 'data/api'
 import { executeSql, ExecuteSqlError } from 'data/sql/execute-sql-query'
 import { useSelectedProjectQuery } from 'hooks/misc/useSelectedProject'
@@ -20,8 +21,8 @@ type UsersVariables = {
   providers?: string[]
   sort?: 'id' | 'created_at' | 'email' | 'phone' | 'last_sign_in_at'
   order?: 'asc' | 'desc'
-
-  column?: 'id' | 'email' | 'phone'
+  /** If set, uses optimized prefix search for the specified column */
+  column?: OptimizedSearchColumns
   startAt?: string
 }
 

packages/pg-meta/src/sql/studio/get-users-common.ts

@@ -0,0 +1,51 @@
+export function prefixToUUID(prefix: string, max: boolean) {
+  const mapped = '00000000-0000-0000-0000-000000000000'
+    .split('')
+    .map((c, i) => (c === '-' ? c : prefix[i] ?? c))
+
+  if (prefix.length >= mapped.length) {
+    return mapped.join('')
+  }
+
+  if (prefix.length && prefix.length < 15) {
+    mapped[14] = '4'
+  }
+
+  if (prefix.length && prefix.length < 20) {
+    mapped[19] = max ? 'b' : '8'
+  }
+
+  if (max) {
+    for (let i = prefix.length; i < mapped.length; i += 1) {
+      if (mapped[i] === '0') {
+        mapped[i] = 'f'
+      }
+    }
+  }
+
+  return mapped.join('')
+}
+
+export function stringRange(prefix: string): [string, string | undefined] {
+  if (!prefix) {
+    return [prefix, undefined]
+  }
+
+  const lastCharCode = prefix.charCodeAt(prefix.length - 1)
+  const TILDE_CHAR_CODE = 126 // '~'
+  const Z_CHAR_CODE = 122 // 'z'
+
+  // 'z' (122): append '~' to avoid PostgreSQL collation issues with '{'
+  if (lastCharCode === Z_CHAR_CODE) {
+    return [prefix, prefix + '~']
+  }
+
+  // '~' (126) or beyond: append space since we can't increment further
+  if (lastCharCode >= TILDE_CHAR_CODE) {
+    return [prefix, prefix + ' ']
+  }
+
+  // All other characters: increment the last character
+  const upperBound = prefix.substring(0, prefix.length - 1) + String.fromCharCode(lastCharCode + 1)
+  return [prefix, upperBound]
+}

packages/pg-meta/src/sql/studio/get-users-count.ts

@@ -1,4 +1,7 @@
+import type { OptimizedSearchColumns } from './get-users-types'
 import { COUNT_ESTIMATE_SQL, THRESHOLD_COUNT } from './get-count-estimate'
+import { stringRange, prefixToUUID } from './get-users-common'
+import { literal } from '../../pg-format'
 
 export const USERS_COUNT_ESTIMATE_SQL = `select reltuples as estimate from pg_class where oid = 'auth.users'::regclass`
 
@@ -7,46 +10,74 @@ export const getUsersCountSQL = ({
   keywords,
   providers,
   forceExactCount = false,
+  column,
 }: {
   filter?: 'verified' | 'unverified' | 'anonymous'
   keywords?: string
   providers?: string[]
   forceExactCount?: boolean
+  /** If set, uses optimized prefix search for the specified column */
+  column?: OptimizedSearchColumns
 }) => {
   const hasValidKeywords = keywords && keywords !== ''
 
   const conditions: string[] = []
   const baseQueryCount = `select count(*) from auth.users`
   const baseQuerySelect = `select * from auth.users`
 
-  if (hasValidKeywords) {
-    // [Joshen] Escape single quotes properly
-    const formattedKeywords = keywords.replaceAll("'", "''")
-    conditions.push(
-      `id::text ilike '%${formattedKeywords}%' or email ilike '%${formattedKeywords}%' or phone ilike '%${formattedKeywords}%'`
-    )
-  }
-
-  if (filter === 'verified') {
-    conditions.push(`email_confirmed_at IS NOT NULL or phone_confirmed_at IS NOT NULL`)
-  } else if (filter === 'anonymous') {
-    conditions.push(`is_anonymous is true`)
-  } else if (filter === 'unverified') {
-    conditions.push(`email_confirmed_at IS NULL AND phone_confirmed_at IS NULL`)
-  }
-
-  if (providers && providers.length > 0) {
-    // [Joshen] This is arguarbly not fully optimized, but at the same time not commonly used
-    // JFYI in case we do eventually run into performance issues here when filtering for SAML provider
-    if (providers.includes('saml 2.0')) {
+  const optimizedSearchMode = column && hasValidKeywords
+  if (optimizedSearchMode) {
+    if (column === 'email') {
+      const range = stringRange(keywords)
+      const lowerBound = literal(range[0])
+      const upperBound = range[1] ? literal(range[1]) : null
       conditions.push(
-        `(select jsonb_agg(case when value ~ '^sso' then 'sso' else value end) from jsonb_array_elements_text((raw_app_meta_data ->> 'providers')::jsonb)) ?| array[${providers.map((p) => (p === 'saml 2.0' ? `'sso'` : `'${p}'`)).join(', ')}]`.trim()
+        `lower(email) >= ${lowerBound}${upperBound ? ` and lower(email) < ${upperBound}` : ''} and instance_id = '00000000-0000-0000-0000-000000000000'::uuid`
       )
-    } else {
+    } else if (column === 'phone') {
+      const range = stringRange(keywords)
+      const lowerBound = literal(range[0])
+      const upperBound = range[1] ? literal(range[1]) : null
+      conditions.push(`phone >= ${lowerBound}${upperBound ? ` and phone < ${upperBound}` : ''}`)
+    } else if (column === 'id') {
+      const lowerUUID = prefixToUUID(keywords, false)
+      const isMatchingUUIDValue = lowerUUID === keywords
+      if (isMatchingUUIDValue) {
+        conditions.push(`id = ${literal(keywords)}`)
+      } else {
+        const upperUUID = prefixToUUID(keywords, true)
+        conditions.push(`id >= ${literal(lowerUUID)} and id < ${literal(upperUUID)}`)
+      }
+    }
+  } else {
+    // Unified search mode - apply all filters
+    if (hasValidKeywords) {
+      const escapedFilterKeywords = literal(`%${keywords}%`)
       conditions.push(
-        `(raw_app_meta_data->>'providers')::jsonb ?| array[${providers.map((p) => `'${p}'`).join(', ')}]`
+        `id::text ilike ${escapedFilterKeywords} or email ilike ${escapedFilterKeywords} or phone ilike ${escapedFilterKeywords}`
       )
     }
+
+    if (filter === 'verified') {
+      conditions.push(`email_confirmed_at IS NOT NULL or phone_confirmed_at IS NOT NULL`)
+    } else if (filter === 'anonymous') {
+      conditions.push(`is_anonymous is true`)
+    } else if (filter === 'unverified') {
+      conditions.push(`email_confirmed_at IS NULL AND phone_confirmed_at IS NULL`)
+    }
+
+    if (providers && providers.length > 0) {
+      // [Joshen] This is arguarbly not fully optimized, but at the same time not commonly used
+      // JFYI in case we do eventually run into performance issues here when filtering for SAML provider
+      if (providers.includes('saml 2.0')) {
+        const mappedProviders = providers.map((p) => (p === 'saml 2.0' ? 'sso' : p))
+        conditions.push(
+          `(select jsonb_agg(case when value ~ '^sso' then 'sso' else value end) from jsonb_array_elements_text((raw_app_meta_data ->> 'providers')::jsonb)) ?| array[${literal(mappedProviders)}]`.trim()
+        )
+      } else {
+        conditions.push(`(raw_app_meta_data->>'providers')::jsonb ?| array[${literal(providers)}]`)
+      }
+    }
   }
 
   const combinedConditions = conditions.map((x) => `(${x})`).join(' and ')
@@ -58,16 +89,16 @@ export const getUsersCountSQL = ({
     const selectBaseSql = `${baseQuerySelect}${whereClause}`
     const countBaseSql = `${baseQueryCount}${whereClause}`
 
-    const escapedSelectSql = selectBaseSql.replaceAll("'", "''")
+    const escapedSelectSql = literal(selectBaseSql)
 
     const sql = `
 ${COUNT_ESTIMATE_SQL}
 
 with approximation as (${USERS_COUNT_ESTIMATE_SQL})
 select 
   case 
-    when estimate = -1 then (select pg_temp.count_estimate('${escapedSelectSql}'))::int
-    when estimate > ${THRESHOLD_COUNT} then ${conditions.length > 0 ? `(select pg_temp.count_estimate('${escapedSelectSql}'))::int` : 'estimate::int'}
+    when estimate = -1 then (select pg_temp.count_estimate(${escapedSelectSql}))::int
+    when estimate > ${THRESHOLD_COUNT} then ${conditions.length > 0 ? `(select pg_temp.count_estimate(${escapedSelectSql}))::int` : 'estimate::int'}
     else (${countBaseSql})
   end as count,
   estimate = -1 or estimate > ${THRESHOLD_COUNT} as is_estimate