feat(cohorts): enable cohort pagination

frontend/src/layout/navigation-3000/sidebars/cohorts.ts

@@ -94,7 +94,7 @@ export const cohortsSidebarLogic = kea<cohortsSidebarLogicType>([
                 if (searchTerm) {
                     return fuse.search(searchTerm).map((result) => [result.item, result.matches as FuseSearchMatch[]])
                 }
-                return cohorts.map((cohort) => [cohort, null])
+                return cohorts.results.map((cohort) => [cohort, null])
             },
         ],
     })),

frontend/src/lib/api.ts

@@ -1488,16 +1488,25 @@ const api = {
         async duplicate(cohortId: CohortType['id']): Promise<CohortType> {
             return await new ApiRequest().cohortsDuplicate(cohortId).get()
         },
-        async list(): Promise<PaginatedResponse<CohortType>> {
-            // TODO: Remove hard limit and paginate cohorts
-            return await new ApiRequest().cohorts().withQueryString('limit=600').get()
+        async list(): Promise<CohortType[]> {
+            // TODO: Deprecate this in favor of listPaginated
+            return (await this.listPaginated({ limit: 600 })).results
         },
         determineDeleteEndpoint(): string {
             return new ApiRequest().cohorts().assembleEndpointUrl()
         },
         determineListUrl(cohortId: number | 'new', params: PersonListParams): string {
             return `/api/cohort/${cohortId}/persons?${toParams(params)}`
         },
+        async listPaginated(
+            params: {
+                limit?: number
+                offset?: number
+                search?: string
+            } = {}
+        ): Promise<CountedPaginatedResponse<CohortType>> {
+            return await new ApiRequest().cohorts().withQueryString(toParams(params)).get()
+        },
     },
 
     dashboards: {

frontend/src/lib/components/InsightLegend/InsightLegendRow.tsx

@@ -48,7 +48,7 @@ export function InsightLegendRow({ rowIndex, item }: InsightLegendRowProps): JSX
     const formattedBreakdownValue = formatBreakdownLabel(
         item.breakdown_value,
         breakdownFilter,
-        cohorts,
+        cohorts.results,
         formatPropertyValueForDisplay
     )
 

frontend/src/lib/components/PropertyFilters/components/taxonomicPropertyFilterLogic.ts

@@ -68,7 +68,8 @@ export const taxonomicPropertyFilterLogic = kea<taxonomicPropertyFilterLogicType
         ],
         selectedCohortName: [
             (s) => [s.filter, cohortsModel.selectors.cohorts],
-            (filter, cohorts) => (filter?.type === 'cohort' ? cohorts.find((c) => c.id === filter?.value)?.name : null),
+            (filter, cohorts) =>
+                filter?.type === 'cohort' ? cohorts.results.find((c) => c.id === filter?.value)?.name : null,
         ],
         activeTaxonomicGroup: [
             (s) => [s.filter, s.taxonomicGroups],

frontend/src/lib/components/TaxonomicFilter/infiniteListLogic.ts

@@ -247,7 +247,8 @@ export const infiniteListLogic = kea<infiniteListLogicType>([
 
                     if (group?.logic && group?.value) {
                         let items = group.logic.selectors[group.value]?.(state)
-                        if (group?.value === 'featureFlags' && items.results) {
+                        // Handle paginated responses for cohorts
+                        if (items?.results) {
                             items = items.results
                         }
                         // TRICKY: Feature flags don't support dynamic behavioral cohorts,

frontend/src/lib/components/TaxonomicFilter/taxonomicFilterLogic.tsx

@@ -24,7 +24,6 @@ import { ReplayTaxonomicFilters } from 'scenes/session-recordings/filters/Replay
 import { teamLogic } from 'scenes/teamLogic'
 
 import { actionsModel } from '~/models/actionsModel'
-import { cohortsModel } from '~/models/cohortsModel'
 import { dashboardsModel } from '~/models/dashboardsModel'
 import { groupPropertiesModel } from '~/models/groupPropertiesModel'
 import { groupsModel } from '~/models/groupsModel'
@@ -355,7 +354,7 @@ export const taxonomicFilterLogic = kea<taxonomicFilterLogicType>([
                         name: 'Cohorts',
                         searchPlaceholder: 'cohorts',
                         type: TaxonomicFilterGroupType.Cohorts,
-                        logic: cohortsModel,
+                        endpoint: combineUrl(`api/projects/${projectId}/cohorts/`).url,
                         value: 'cohorts',
                         getName: (cohort: CohortType) => cohort.name || `Cohort ${cohort.id}`,
                         getValue: (cohort: CohortType) => cohort.id,
@@ -368,7 +367,7 @@ export const taxonomicFilterLogic = kea<taxonomicFilterLogicType>([
                         name: 'Cohorts',
                         searchPlaceholder: 'cohorts',
                         type: TaxonomicFilterGroupType.CohortsWithAllUsers,
-                        logic: cohortsModel,
+                        endpoint: combineUrl(`api/projects/${projectId}/cohorts/`).url,
                         value: 'cohortsWithAllUsers',
                         getName: (cohort: CohortType) => cohort.name || `Cohort ${cohort.id}`,
                         getValue: (cohort: CohortType) => cohort.id,

frontend/src/models/cohortsModel.ts

@@ -1,14 +1,16 @@
-import Fuse from 'fuse.js'
-import { actions, afterMount, beforeUnmount, connect, kea, listeners, path, reducers, selectors } from 'kea'
+import { actions, beforeUnmount, connect, events, kea, listeners, path, reducers, selectors } from 'kea'
 import { loaders } from 'kea-loaders'
-import api from 'lib/api'
+import { actionToUrl, router, urlToAction } from 'kea-router'
+import api, { CountedPaginatedResponse } from 'lib/api'
 import { exportsLogic } from 'lib/components/ExportButton/exportsLogic'
+import { PaginationManual } from 'lib/lemon-ui/PaginationControl'
 import { deleteWithUndo } from 'lib/utils/deleteWithUndo'
 import { permanentlyMount } from 'lib/utils/kea-logic-builders'
 import { COHORT_EVENT_TYPES_WITH_EXPLICIT_DATETIME } from 'scenes/cohorts/CohortFilters/constants'
 import { BehavioralFilterKey } from 'scenes/cohorts/CohortFilters/types'
 import { personsLogic } from 'scenes/persons/personsLogic'
-import { isAuthenticatedTeam, teamLogic } from 'scenes/teamLogic'
+import { personsManagementSceneLogic } from 'scenes/persons-management/personsManagementSceneLogic'
+import { teamLogic } from 'scenes/teamLogic'
 import { urls } from 'scenes/urls'
 
 import {
@@ -24,6 +26,18 @@ import type { cohortsModelType } from './cohortsModelType'
 
 const POLL_TIMEOUT = 5000
 
+export const COHORTS_PER_PAGE = 100
+
+export interface CohortFilters {
+    search?: string
+    page?: number
+}
+
+export const DEFAULT_COHORT_FILTERS: CohortFilters = {
+    search: undefined,
+    page: 1,
+}
+
 export function processCohort(cohort: CohortType): CohortType {
     return {
         ...cohort,
@@ -91,23 +105,28 @@ export const cohortsModel = kea<cohortsModelType>([
     path(['models', 'cohortsModel']),
     connect(() => ({
         values: [teamLogic, ['currentTeam']],
-        actions: [exportsLogic, ['startExport']],
+        actions: [exportsLogic, ['startExport'], personsManagementSceneLogic, ['setTabKey']],
     })),
     actions(() => ({
         setPollTimeout: (pollTimeout: number | null) => ({ pollTimeout }),
         updateCohort: (cohort: CohortType) => ({ cohort }),
         deleteCohort: (cohort: Partial<CohortType>) => ({ cohort }),
         cohortCreated: (cohort: CohortType) => ({ cohort }),
         exportCohortPersons: (id: CohortType['id'], columns?: string[]) => ({ id, columns }),
+        setCohortFilters: (filters: Partial<CohortFilters>) => ({ filters }),
     })),
-    loaders(() => ({
+    loaders(({ values }) => ({
         cohorts: {
-            __default: [] as CohortType[],
+            __default: { count: 0, results: [] } as CountedPaginatedResponse<CohortType>,
             loadCohorts: async () => {
-                // TRICKY in tests this was returning undefined without calling list
-                const response = await api.cohorts.list()
-                personsLogic.findMounted({ syncWithUrl: true })?.actions.loadCohorts() // To ensure sync on person page
-                return response?.results?.map((cohort) => processCohort(cohort)) || []
+                const response = await api.cohorts.listPaginated({
+                    ...values.paramsFromFilters,
+                })
+                personsLogic.findMounted({ syncWithUrl: true })?.actions.loadCohorts()
+                return {
+                    count: response.count,
+                    results: response.results.map((cohort) => processCohort(cohort)),
+                }
             },
         },
     })),
@@ -123,45 +142,72 @@ export const cohortsModel = kea<cohortsModelType>([
                 if (!cohort) {
                     return state
                 }
-                return [...state].map((existingCohort) => (existingCohort.id === cohort.id ? cohort : existingCohort))
+                return {
+                    ...state,
+                    results: state.results.map((existingCohort) =>
+                        existingCohort.id === cohort.id ? cohort : existingCohort
+                    ),
+                }
             },
             cohortCreated: (state, { cohort }) => {
                 if (!cohort) {
                     return state
                 }
-                return [cohort, ...state]
+                return {
+                    ...state,
+                    results: [cohort, ...state.results],
+                }
             },
             deleteCohort: (state, { cohort }) => {
                 if (!cohort.id) {
                     return state
                 }
-                return [...state].filter((c) => c.id !== cohort.id)
+                return {
+                    ...state,
+                    results: state.results.filter((c) => c.id !== cohort.id),
+                }
             },
         },
+        cohortFilters: [
+            DEFAULT_COHORT_FILTERS,
+            {
+                setCohortFilters: (state, { filters }) => {
+                    return { ...state, ...filters }
+                },
+            },
+        ],
     }),
     selectors({
-        cohortsWithAllUsers: [(s) => [s.cohorts], (cohorts) => [{ id: 'all', name: 'All Users*' }, ...cohorts]],
+        cohortsWithAllUsers: [(s) => [s.cohorts], (cohorts) => [{ id: 'all', name: 'All Users*' }, ...cohorts.results]],
         cohortsById: [
             (s) => [s.cohorts],
             (cohorts): Partial<Record<string | number, CohortType>> =>
-                Object.fromEntries(cohorts.map((cohort) => [cohort.id, cohort])),
+                Object.fromEntries(cohorts.results.map((cohort) => [cohort.id, cohort])),
         ],
-
-        cohortsSearch: [
-            (s) => [s.cohorts],
-            (cohorts): ((term: string) => CohortType[]) => {
-                const fuse = new Fuse<CohortType>(cohorts ?? [], {
-                    keys: ['name'],
-                    threshold: 0.3,
-                })
-
-                return (term) => fuse.search(term).map((result) => result.item)
+        count: [(selectors) => [selectors.cohorts], (cohorts) => cohorts.count],
+        paramsFromFilters: [
+            (s) => [s.cohortFilters],
+            (filters: CohortFilters) => ({
+                ...filters,
+                limit: COHORTS_PER_PAGE,
+                offset: filters.page ? (filters.page - 1) * COHORTS_PER_PAGE : 0,
+            }),
+        ],
+        pagination: [
+            (s) => [s.cohortFilters, s.count],
+            (filters, count): PaginationManual => {
+                return {
+                    controlled: true,
+                    pageSize: COHORTS_PER_PAGE,
+                    currentPage: filters.page || 1,
+                    entryCount: count,
+                }
             },
         ],
     }),
     listeners(({ actions }) => ({
-        loadCohortsSuccess: async ({ cohorts }: { cohorts: CohortType[] }) => {
-            const is_calculating = cohorts.filter((cohort) => cohort.is_calculating).length > 0
+        loadCohortsSuccess: async ({ cohorts }: { cohorts: CountedPaginatedResponse<CohortType> }) => {
+            const is_calculating = cohorts.results.filter((cohort) => cohort.is_calculating).length > 0
             if (!is_calculating || !window.location.pathname.includes(urls.cohorts())) {
                 return
             }
@@ -172,7 +218,8 @@ export const cohortsModel = kea<cohortsModelType>([
                 export_format: ExporterFormat.CSV,
                 export_context: {
                     path: `/api/cohort/${id}/persons`,
-                },
+                    columns,
+                } as { path: string; columns?: string[] },
             }
             if (columns && columns.length > 0) {
                 exportCommand.export_context['columns'] = columns
@@ -186,16 +233,51 @@ export const cohortsModel = kea<cohortsModelType>([
                 callback: actions.loadCohorts,
             })
         },
+        setCohortFilters: async () => {
+            if (!router.values.location.pathname.includes(urls.cohorts())) {
+                return
+            }
+            actions.loadCohorts()
+        },
+    })),
+    actionToUrl(({ values }) => ({
+        setCohortFilters: () => {
+            const searchParams: Record<string, any> = {
+                ...values.cohortFilters,
+            }
+
+            // Only include non-default values in URL
+            Object.keys(searchParams).forEach((key) => {
+                if (
+                    searchParams[key] === undefined ||
+                    searchParams[key] === DEFAULT_COHORT_FILTERS[key as keyof CohortFilters]
+                ) {
+                    delete searchParams[key]
+                }
+            })
+
+            return [router.values.location.pathname, searchParams, router.values.hashParams, { replace: true }]
+        },
+    })),
+    urlToAction(({ actions }) => ({
+        [urls.cohorts()]: (_, searchParams) => {
+            const { page, search } = searchParams
+            const filtersFromUrl: Partial<CohortFilters> = {
+                search,
+            }
+
+            filtersFromUrl.page = page !== undefined ? parseInt(page) : undefined
+
+            actions.setCohortFilters({ ...DEFAULT_COHORT_FILTERS, ...filtersFromUrl })
+        },
     })),
     beforeUnmount(({ values }) => {
         clearTimeout(values.pollTimeout || undefined)
     }),
-
-    afterMount(({ actions, values }) => {
-        if (isAuthenticatedTeam(values.currentTeam)) {
-            // Don't load on shared insights/dashboards
+    events(({ actions }) => ({
+        afterMount: () => {
             actions.loadCohorts()
-        }
-    }),
+        },
+    })),
     permanentlyMount(),
 ])

frontend/src/queries/nodes/InsightViz/ResultCustomizationsModal.tsx

@@ -108,7 +108,7 @@ function TrendsInfo({ dataset, resultCustomizationBy }: TrendsInfoProps): JSX.El
                         {formatBreakdownLabel(
                             dataset.breakdown_value,
                             breakdownFilter,
-                            cohorts,
+                            cohorts.results,
                             formatPropertyValueForDisplay
                         )}
                     </b>

frontend/src/scenes/cohorts/Cohorts.tsx

@@ -21,10 +21,10 @@ import { CohortType, ProductKey } from '~/types'
 import { cohortsModel } from '../../models/cohortsModel'
 
 export function Cohorts(): JSX.Element {
-    const { cohorts, cohortsSearch, cohortsLoading } = useValues(cohortsModel)
-    const { deleteCohort, exportCohortPersons } = useActions(cohortsModel)
+    const { cohorts, cohortsLoading, pagination, cohortFilters } = useValues(cohortsModel)
+    const { deleteCohort, exportCohortPersons, setCohortFilters } = useActions(cohortsModel)
     const { searchParams } = useValues(router)
-    const [searchTerm, setSearchTerm] = useState<string>('')
+    const [searchTerm, setSearchTerm] = useState(cohortFilters.search || '')
 
     const columns: LemonTableColumns<CohortType> = [
         {
@@ -142,7 +142,7 @@ export function Cohorts(): JSX.Element {
                 productKey={ProductKey.COHORTS}
                 thingName="cohort"
                 description="Use cohorts to group people together, such as users who used your app in the last week, or people who viewed the signup page but didn’t convert."
-                isEmpty={cohorts?.length == 0 && !cohortsLoading}
+                isEmpty={cohorts.count == 0 && !cohortsLoading && !searchTerm}
                 docsURL="https://posthog.com/docs/data/cohorts"
                 action={() => router.actions.push(urls.cohort('new'))}
                 customHog={ListHog}
@@ -152,16 +152,19 @@ export function Cohorts(): JSX.Element {
                 <LemonInput
                     type="search"
                     placeholder="Search for cohorts"
-                    onChange={setSearchTerm}
+                    onChange={(search) => {
+                        setSearchTerm(search)
+                        setCohortFilters({ search: search || undefined, page: 1 })
+                    }}
                     value={searchTerm}
                 />
             </div>
             <LemonTable
                 columns={columns}
                 loading={cohortsLoading}
                 rowKey="id"
-                pagination={{ pageSize: 100 }}
-                dataSource={searchTerm ? cohortsSearch(searchTerm) : cohorts ?? []}
+                pagination={pagination}
+                dataSource={cohorts.results}
                 nouns={['cohort', 'cohorts']}
                 data-attr="cohorts-table"
             />

frontend/src/scenes/feature-flags/FeatureFlagProjects.tsx

@@ -138,7 +138,7 @@ function FeatureFlagCopySection(): JSX.Element {
     const { currentTeam } = useValues(teamLogic)
     const { cohorts } = useValues(cohortsModel)
 
-    const hasStaticCohort = checkHasStaticCohort(featureFlag, cohorts)
+    const hasStaticCohort = checkHasStaticCohort(featureFlag, cohorts.results)
     const hasMultipleProjects = (currentOrganization?.teams?.length ?? 0) > 1
 
     return hasMultipleProjects && featureFlag.can_edit ? (

frontend/src/scenes/funnels/FunnelTooltip.tsx

@@ -58,7 +58,7 @@ export function FunnelTooltip({
                     {formatBreakdownLabel(
                         series.breakdown_value,
                         breakdownFilter,
-                        cohorts,
+                        cohorts.results,
                         formatPropertyValueForDisplay
                     )}
                 </strong>