feat(persons): save persons as a static cohort (#25310)

Co-authored-by: github-actions <41898282+github-actions[bot]@users.noreply.github.com>
PostHog · Oct 1, 2024 · efdc026 · efdc026
1 parent 38bbb90
commit efdc026
Show file tree

Hide file tree

Showing 3 changed files with 66 additions and 5 deletions.
diff --git a/frontend/src/lib/components/ExportButton/exportsLogic.ts b/frontend/src/lib/components/ExportButton/exportsLogic.ts
@@ -1,14 +1,18 @@
 import { actions, connect, kea, listeners, path, reducers } from 'kea'
 import { loaders } from 'kea-loaders'
+import { router } from 'kea-router'
 import api from 'lib/api'
 import { downloadBlob, downloadExportedAsset, TriggerExportProps } from 'lib/components/ExportButton/exporter'
 import { dayjs } from 'lib/dayjs'
 import { lemonToast } from 'lib/lemon-ui/LemonToast'
 import { delay } from 'lib/utils'
 import posthog from 'posthog-js'
+import { urls } from 'scenes/urls'
 
 import { sidePanelStateLogic } from '~/layout/navigation-3000/sidepanel/sidePanelStateLogic'
-import { ExportContext, ExportedAssetType, ExporterFormat, LocalExportContext, SidePanelTab } from '~/types'
+import { cohortsModel } from '~/models/cohortsModel'
+import { DataNode } from '~/queries/schema'
+import { CohortType, ExportContext, ExportedAssetType, ExporterFormat, LocalExportContext, SidePanelTab } from '~/types'
 
 import type { exportsLogicType } from './exportsLogicType'
 
@@ -29,6 +33,7 @@ export const exportsLogic = kea<exportsLogicType>([
         pollExportStatus: (exportedAsset: ExportedAssetType) => ({ exportedAsset }),
         addFresh: (exportedAsset: ExportedAssetType) => ({ exportedAsset }),
         removeFresh: (exportedAsset: ExportedAssetType) => ({ exportedAsset }),
+        createStaticCohort: (name: string, query: DataNode) => ({ query, name }),
     }),
 
     connect({
@@ -134,6 +139,29 @@ export const exportsLogic = kea<exportsLogicType>([
                 error: 'Export failed!',
             })
         },
+        createStaticCohort: async ({ query, name }) => {
+            try {
+                const toastId = 'toast-' + Math.random()
+                lemonToast.info('Saving cohort...', { toastId, autoClose: false })
+                const cohort: CohortType = await api.create('api/cohort', {
+                    is_static: true,
+                    name: name || 'Query cohort',
+                    query: query,
+                })
+                cohortsModel.actions.cohortCreated(cohort)
+                await delay(500) // just in case the toast is too fast
+                lemonToast.dismiss(toastId)
+                lemonToast.success('Cohort saved', {
+                    toastId: `${toastId}-success`,
+                    button: {
+                        label: 'View cohort',
+                        action: () => router.actions.push(urls.cohort(cohort.id)),
+                    },
+                })
+            } catch (e) {
+                lemonToast.error('Cohort save failed')
+            }
+        },
     })),
 
     loaders(() => ({

diff --git a/frontend/src/models/cohortsModel.ts b/frontend/src/models/cohortsModel.ts
@@ -89,10 +89,10 @@ function processCohortCriteria(criteria: AnyCohortCriteriaType): AnyCohortCriter
 
 export const cohortsModel = kea<cohortsModelType>([
     path(['models', 'cohortsModel']),
-    connect({
+    connect(() => ({
         values: [teamLogic, ['currentTeam']],
         actions: [exportsLogic, ['startExport']],
-    }),
+    })),
     actions(() => ({
         setPollTimeout: (pollTimeout: number | null) => ({ pollTimeout }),
         updateCohort: (cohort: CohortType) => ({ cohort }),

diff --git a/frontend/src/queries/nodes/DataTable/DataTableExport.tsx b/frontend/src/queries/nodes/DataTable/DataTableExport.tsx
@@ -1,8 +1,9 @@
 import { IconDownload } from '@posthog/icons'
-import { LemonButton, LemonMenu, lemonToast } from '@posthog/lemon-ui'
+import { LemonButton, LemonDialog, LemonInput, LemonMenu, lemonToast } from '@posthog/lemon-ui'
 import { useActions, useValues } from 'kea'
 import { TriggerExportProps } from 'lib/components/ExportButton/exporter'
 import { exportsLogic } from 'lib/components/ExportButton/exportsLogic'
+import { LemonField } from 'lib/lemon-ui/LemonField'
 import { copyToClipboard } from 'lib/utils/copyToClipboard'
 import Papa from 'papaparse'
 import { asDisplay } from 'scenes/persons/person-utils'
@@ -195,7 +196,7 @@ interface DataTableExportProps {
 
 export function DataTableExport({ query }: DataTableExportProps): JSX.Element | null {
     const { dataTableRows, columnsInResponse, columnsInQuery, queryWithDefaults } = useValues(dataTableLogic)
-    const { startExport } = useActions(exportsLogic)
+    const { startExport, createStaticCohort } = useActions(exportsLogic)
 
     const source: DataNode = query.source
     const filterCount =
@@ -205,6 +206,7 @@ export function DataTableExport({ query }: DataTableExportProps): JSX.Element |
     const canExportAllColumns =
         (isEventsQuery(source) && source.select.includes('*')) || isPersonsNode(source) || isActorsQuery(source)
     const showExportClipboardButtons = isPersonsNode(source) || isEventsQuery(source) || isHogQLQuery(source)
+    const canSaveAsCohort = isActorsQuery(source)
 
     return (
         <LemonMenu
@@ -270,6 +272,37 @@ export function DataTableExport({ query }: DataTableExportProps): JSX.Element |
                         },
                     ],
                 },
+                canSaveAsCohort && {
+                    label: 'Save as cohort',
+                    items: [
+                        {
+                            label: 'Save as static cohort',
+                            onClick: () => {
+                                LemonDialog.openForm({
+                                    title: 'Save as static cohort',
+                                    description: 'This will create a cohort with the current results of the query.',
+                                    initialValues: {
+                                        name: '',
+                                    },
+                                    content: (
+                                        <LemonField name="name">
+                                            <LemonInput
+                                                type="text"
+                                                data-attr="insight-name"
+                                                placeholder="Please enter a name for the cohort"
+                                                autoFocus
+                                            />
+                                        </LemonField>
+                                    ),
+                                    errors: {
+                                        name: (name) => (!name ? 'You must enter a name' : undefined),
+                                    },
+                                    onSubmit: async ({ name }) => createStaticCohort(name, source),
+                                })
+                            },
+                        },
+                    ],
+                },
             ].filter(Boolean)}
         >
             <LemonButton type="secondary" icon={<IconDownload />} data-attr="data-table-export-menu">