feat: Support for selecting batch export model (#22914)

* refactor: Update metrics to fetch counts at request time * fix: Move import to method * fix: Add function * fix: Typing fixes * feat: Custom schemas for batch exports * feat: Add django api support for defining model * feat: Frontend support for model field * fix: Clean-up * fix: Add model to story * fix: Add missing migration * fix: Make new field nullable * Update UI snapshots for `chromium` (1) * Update UI snapshots for `chromium` (1) * Update UI snapshots for `chromium` (1) * Update UI snapshots for `chromium` (1) * Update UI snapshots for `chromium` (1) * Update UI snapshots for `chromium` (1) * fix: Bump migration number --------- Co-authored-by: github-actions <41898282+github-actions[bot]@users.noreply.github.com>
PostHog · Jun 18, 2024 · 52eb1bf · 52eb1bf
1 parent 1ed340d
commit 52eb1bf
Show file tree

Hide file tree

Showing 16 changed files with 387 additions and 12 deletions.
diff --git a/frontend/src/lib/constants.tsx b/frontend/src/lib/constants.tsx
@@ -155,6 +155,7 @@ export const FEATURE_FLAGS = {
     WEB_ANALYTICS: 'web-analytics', // owner @robbie-c #team-web-analytics
     WEB_ANALYTICS_SAMPLING: 'web-analytics-sampling', // owner @robbie-c #team-web-analytics
     HIGH_FREQUENCY_BATCH_EXPORTS: 'high-frequency-batch-exports', // owner: @tomasfarias
+    PERSON_BATCH_EXPORTS: 'person-batch-exports', // owner: @tomasfarias
     // owner: #team-replay, only to be enabled for PostHog team testing
     EXCEPTION_AUTOCAPTURE: 'exception-autocapture',
     DATA_WAREHOUSE: 'data-warehouse', // owner: @EDsCODE

diff --git a/frontend/src/queries/schema.json b/frontend/src/queries/schema.json
@@ -2400,6 +2400,29 @@
             "required": ["key", "operator", "type"],
             "type": "object"
         },
+        "DatabaseSchemaBatchExportTable": {
+            "additionalProperties": false,
+            "properties": {
+                "fields": {
+                    "additionalProperties": {
+                        "$ref": "#/definitions/DatabaseSchemaField"
+                    },
+                    "type": "object"
+                },
+                "id": {
+                    "type": "string"
+                },
+                "name": {
+                    "type": "string"
+                },
+                "type": {
+                    "const": "batch_export",
+                    "type": "string"
+                }
+            },
+            "required": ["fields", "id", "name", "type"],
+            "type": "object"
+        },
         "DatabaseSchemaDataWarehouseTable": {
             "additionalProperties": false,
             "properties": {
@@ -2587,6 +2610,9 @@
                 },
                 {
                     "$ref": "#/definitions/DatabaseSchemaViewTable"
+                },
+                {
+                    "$ref": "#/definitions/DatabaseSchemaBatchExportTable"
                 }
             ]
         },
@@ -2606,7 +2632,7 @@
                     "type": "string"
                 },
                 "type": {
-                    "enum": ["posthog", "data_warehouse", "view"],
+                    "enum": ["posthog", "data_warehouse", "view", "batch_export"],
                     "type": "string"
                 }
             },

diff --git a/frontend/src/queries/schema.ts b/frontend/src/queries/schema.ts
@@ -1448,7 +1448,7 @@ export interface DatabaseSchemaField {
 }
 
 export interface DatabaseSchemaTableCommon {
-    type: 'posthog' | 'data_warehouse' | 'view'
+    type: 'posthog' | 'data_warehouse' | 'view' | 'batch_export'
     id: string
     name: string
     fields: Record<string, DatabaseSchemaField>
@@ -1471,10 +1471,15 @@ export interface DatabaseSchemaDataWarehouseTable extends DatabaseSchemaTableCom
     source?: DatabaseSchemaSource
 }
 
+export interface DatabaseSchemaBatchExportTable extends DatabaseSchemaTableCommon {
+    type: 'batch_export'
+}
+
 export type DatabaseSchemaTable =
     | DatabaseSchemaPostHogTable
     | DatabaseSchemaDataWarehouseTable
     | DatabaseSchemaViewTable
+    | DatabaseSchemaBatchExportTable
 
 export interface DatabaseSchemaQueryResponse {
     tables: Record<string, DatabaseSchemaTable>

diff --git a/frontend/src/scenes/batch_exports/BatchExportEditForm.tsx b/frontend/src/scenes/batch_exports/BatchExportEditForm.tsx
@@ -107,7 +107,16 @@ export function BatchExportGeneralEditFields({
                         <LemonInput placeholder="Name your workflow for future reference" />
                     </LemonField>
                 )}
-
+                {featureFlags[FEATURE_FLAGS.PERSON_BATCH_EXPORTS] && (
+                    <LemonField name="model" label="Model" info="A model defines the data that will be exported.">
+                        <LemonSelect
+                            options={[
+                                { value: 'events', label: 'Events' },
+                                { value: 'persons', label: 'Persons' },
+                            ]}
+                        />
+                    </LemonField>
+                )}
                 <div className="flex gap-2 items-start flex-wrap">
                     <LemonField
                         name="interval"

diff --git a/frontend/src/scenes/batch_exports/BatchExports.stories.tsx b/frontend/src/scenes/batch_exports/BatchExports.stories.tsx
@@ -27,6 +27,7 @@ export default {
                 1: {
                     id: '1',
                     team_id: 1,
+                    model: 'events',
                     name: 'My S3 Exporter',
                     destination: {
                         type: 'S3',

diff --git a/frontend/src/scenes/batch_exports/batchExportEditLogic.ts b/frontend/src/scenes/batch_exports/batchExportEditLogic.ts
@@ -52,6 +52,7 @@ export const batchExportFormFields = (
         name: !name && !isPipeline ? 'Please enter a name' : '', // In pipeline UI the name is in the top bar
         destination: !destination ? 'Please select a destination' : '',
         interval: !interval ? 'Please select a frequency' : '',
+        model: !name && !isPipeline ? 'Please choose a model' : '',
         paused: '',
         start_at: '',
         end_at: '',
@@ -156,9 +157,10 @@ export const batchExportsEditLogic = kea<batchExportsEditLogicType>([
         batchExportConfigForm: {
             defaults: {
                 name: '',
+                model: 'events',
             } as BatchExportConfigurationForm,
             errors: (form) => batchExportFormFields(props.id === 'new', form),
-            submit: async ({ name, destination, interval, start_at, end_at, paused, ...config }) => {
+            submit: async ({ name, destination, interval, start_at, end_at, paused, model, ...config }) => {
                 const destinationObject: BatchExportService =
                     destination === 'Postgres'
                         ? ({
@@ -194,6 +196,7 @@ export const batchExportsEditLogic = kea<batchExportsEditLogicType>([
                     paused,
                     name,
                     interval,
+                    model,
                     start_at: start_at?.toISOString() ?? null,
                     end_at: end_at?.toISOString() ?? null,
                     destination: destinationObject,

diff --git a/frontend/src/scenes/pipeline/PipelineBatchExportConfiguration.tsx b/frontend/src/scenes/pipeline/PipelineBatchExportConfiguration.tsx
@@ -1,14 +1,17 @@
-import { LemonSwitch } from '@posthog/lemon-ui'
+import { LemonSelect, LemonSwitch } from '@posthog/lemon-ui'
 import { useActions, useValues } from 'kea'
 import { Form } from 'kea-forms'
 import { NotFound } from 'lib/components/NotFound'
 import { PageHeader } from 'lib/components/PageHeader'
+import { FEATURE_FLAGS } from 'lib/constants'
 import { LemonButton } from 'lib/lemon-ui/LemonButton'
 import { LemonField } from 'lib/lemon-ui/LemonField'
 import { LemonInput } from 'lib/lemon-ui/LemonInput'
 import { SpinnerOverlay } from 'lib/lemon-ui/Spinner'
+import { featureFlagLogic } from 'lib/logic/featureFlagLogic'
 import { BatchExportGeneralEditFields, BatchExportsEditFields } from 'scenes/batch_exports/BatchExportEditForm'
 import { BatchExportConfigurationForm } from 'scenes/batch_exports/batchExportEditLogic'
+import { DatabaseTable } from 'scenes/data-management/database/DatabaseTable'
 
 import { BATCH_EXPORT_SERVICE_NAMES, BatchExportService } from '~/types'
 
@@ -22,13 +25,16 @@ export function PipelineBatchExportConfiguration({ service, id }: { service?: st
     const {
         isNew,
         configuration,
+        tables,
         savedConfiguration,
         isConfigurationSubmitting,
         batchExportConfigLoading,
         configurationChanged,
         batchExportConfig,
+        selectedModel,
     } = useValues(logic)
-    const { resetConfiguration, submitConfiguration } = useActions(logic)
+    const { resetConfiguration, submitConfiguration, setSelectedModel } = useActions(logic)
+    const { featureFlags } = useValues(featureFlagLogic)
 
     if (service && !BATCH_EXPORT_SERVICE_NAMES.includes(service)) {
         return <NotFound object={`batch export service ${service}`} />
@@ -105,6 +111,30 @@ export function PipelineBatchExportConfiguration({ service, id }: { service?: st
                             >
                                 <LemonInput type="text" />
                             </LemonField>
+
+                            {featureFlags[FEATURE_FLAGS.PERSON_BATCH_EXPORTS] && (
+                                <>
+                                    <LemonField
+                                        name="model"
+                                        label="Model"
+                                        info="A model defines the data that will be exported."
+                                    >
+                                        <LemonSelect
+                                            options={tables.map((table) => ({ value: table.name, label: table.id }))}
+                                            value={selectedModel}
+                                            onSelect={(newValue) => {
+                                                setSelectedModel(newValue)
+                                            }}
+                                        />
+                                    </LemonField>
+
+                                    <DatabaseTable
+                                        table={selectedModel ? selectedModel : 'events'}
+                                        tables={tables}
+                                        inEditSchemaMode={false}
+                                    />
+                                </>
+                            )}
                         </div>
                         <div className="border bg-bg-light p-3 rounded flex-2 min-w-100">
                             <BatchExportConfigurationFields