feat: Fix up batch export addon requirement (#27073)

Co-authored-by: github-actions <41898282+github-actions[bot]@users.noreply.github.com>

frontend/src/lib/components/PayGateMini/PayGateMini.tsx

@@ -1,4 +1,4 @@
-import { IconInfo, IconOpenSidebar } from '@posthog/icons'
+import { IconInfo, IconOpenSidebar, IconUnlock } from '@posthog/icons'
 import { LemonButton, Link, Tooltip } from '@posthog/lemon-ui'
 import clsx from 'clsx'
 import { useActions, useValues } from 'kea'
@@ -7,6 +7,7 @@ import { useEffect } from 'react'
 import { billingLogic } from 'scenes/billing/billingLogic'
 import { preflightLogic } from 'scenes/PreflightCheck/preflightLogic'
 import { getProductIcon } from 'scenes/products/Products'
+import { userLogic } from 'scenes/userLogic'
 
 import { AvailableFeature, BillingFeatureType, BillingProductV2AddonType, BillingProductV2Type } from '~/types'
 
@@ -41,10 +42,14 @@ export function PayGateMini({
     isGrandfathered,
     docsLink,
 }: PayGateMiniProps): JSX.Element | null {
-    const { productWithFeature, featureInfo, gateVariant } = useValues(payGateMiniLogic({ feature, currentUsage }))
+    const { productWithFeature, featureInfo, gateVariant, bypassPaywall } = useValues(
+        payGateMiniLogic({ feature, currentUsage })
+    )
+    const { setBypassPaywall } = useActions(payGateMiniLogic({ feature, currentUsage }))
     const { preflight, isCloudOrDev } = useValues(preflightLogic)
     const { billingLoading } = useValues(billingLogic)
     const { hideUpgradeModal } = useActions(upgradeModalLogic)
+    const { user } = useValues(userLogic)
 
     useEffect(() => {
         if (gateVariant) {
@@ -73,7 +78,7 @@ export function PayGateMini({
         return null // Don't show anything if paid features are explicitly disabled
     }
 
-    if (gateVariant && productWithFeature && featureInfo && !overrideShouldShowGate) {
+    if (gateVariant && productWithFeature && featureInfo && !overrideShouldShowGate && !bypassPaywall) {
         return (
             <PayGateContent
                 feature={feature}
@@ -96,6 +101,17 @@ export function PayGateMini({
                             Learn more <IconOpenSidebar className="ml-2" />
                         </LemonButton>
                     )}
+
+                    {user?.is_impersonated && (
+                        <LemonButton
+                            type="secondary"
+                            icon={<IconUnlock />}
+                            tooltip="Bypass this paywall - (UI only)"
+                            onClick={() => setBypassPaywall(true)}
+                        >
+                            Bypass paywall
+                        </LemonButton>
+                    )}
                 </div>
             </PayGateContent>
         )
@@ -142,7 +158,7 @@ function PayGateContent({
                 'PayGateMini rounded flex flex-col items-center p-4 text-center'
             )}
         >
-            <div className="flex text-4xl text-warning mb-2">
+            <div className="flex mb-2 text-4xl text-warning">
                 {getProductIcon(productWithFeature.name, featureInfo.icon_key)}
             </div>
             <h2>{featureInfo.name}</h2>
@@ -184,7 +200,7 @@ const renderUsageLimitMessage = (
                     </Tooltip>
                     .
                 </p>
-                <p className="border border-border bg-bg-3000 rounded p-4">
+                <p className="p-4 border rounded border-border bg-bg-3000">
                     <b>Your current plan limit:</b>{' '}
                     <span>
                         {featureAvailableOnOrg.limit} {featureAvailableOnOrg.unit}
@@ -198,7 +214,7 @@ const renderUsageLimitMessage = (
                                 <b>{featureInfoOnNextPlan?.limit} projects</b>.
                             </p>
                         )}
-                        <p className="italic text-xs text-muted mb-4">
+                        <p className="mb-4 text-xs italic text-muted">
                             Need unlimited projects? Check out the{' '}
                             <Link to="/organization/billing?products=platform_and_support" onClick={handleCtaClick}>
                                 Teams addon
@@ -244,9 +260,9 @@ const renderGateVariantMessage = (
 
 const GrandfatheredMessage = (): JSX.Element => {
     return (
-        <div className="flex gap-x-2 bg-bg-3000 rounded text-left mb-4">
-            <IconInfo className="text-muted text-2xl" />
-            <p className="text-muted mb-0">
+        <div className="flex mb-4 text-left rounded gap-x-2 bg-bg-3000">
+            <IconInfo className="text-2xl text-muted" />
+            <p className="mb-0 text-muted">
                 Your plan does not include this feature, but previously set settings may remain. Please upgrade your
                 plan to regain access.
             </p>

frontend/src/lib/components/PayGateMini/payGateMiniLogic.tsx

@@ -1,4 +1,4 @@
-import { actions, connect, kea, key, path, props, selectors } from 'kea'
+import { actions, connect, kea, key, path, props, reducers, selectors } from 'kea'
 import { billingLogic } from 'scenes/billing/billingLogic'
 import { preflightLogic } from 'scenes/PreflightCheck/preflightLogic'
 import { userLogic } from 'scenes/userLogic'
@@ -29,8 +29,17 @@ export const payGateMiniLogic = kea<payGateMiniLogicType>([
         ],
         actions: [],
     })),
+    reducers({
+        bypassPaywall: [
+            false,
+            {
+                setBypassPaywall: (_, { bypassPaywall }) => bypassPaywall,
+            },
+        ],
+    }),
     actions({
         setGateVariant: (gateVariant: GateVariantType) => ({ gateVariant }),
+        setBypassPaywall: (bypassPaywall: boolean) => ({ bypassPaywall }),
     }),
     selectors(({ values, props }) => ({
         productWithFeature: [

frontend/src/scenes/pipeline/destinations/NewDestinations.tsx

@@ -4,6 +4,7 @@ import { useActions, useValues } from 'kea'
 import { PayGateButton } from 'lib/components/PayGateMini/PayGateButton'
 import { PayGateMini } from 'lib/components/PayGateMini/PayGateMini'
 import { LemonTableLink } from 'lib/lemon-ui/LemonTable/LemonTableLink'
+import { userLogic } from 'scenes/userLogic'
 
 import { AvailableFeature, HogFunctionTypeType, PipelineStage } from '~/types'
 
@@ -31,6 +32,7 @@ export function DestinationOptionsTable({ types }: NewDestinationsProps): JSX.El
     const { loading, filteredDestinations, hiddenDestinations } = useValues(newDestinationsLogic({ types }))
     const { canEnableDestination } = useValues(pipelineAccessLogic)
     const { resetFilters } = useActions(destinationsFiltersLogic({ types }))
+    const { user } = useValues(userLogic)
 
     return (
         <>
@@ -75,15 +77,22 @@ export function DestinationOptionsTable({ types }: NewDestinationsProps): JSX.El
                                     type="primary"
                                     data-attr={`new-${PipelineStage.Destination}`}
                                     icon={<IconPlusSmall />}
-                                    // Preserve hash params to pass config in
                                     to={target.url}
-                                    fullWidth
                                 >
                                     Create
                                 </LemonButton>
                             ) : (
-                                <span className="whitespace-nowrap">
+                                <span className="flex items-center gap-2 whitespace-nowrap">
                                     <PayGateButton feature={AvailableFeature.DATA_PIPELINES} type="secondary" />
+                                    {/* Allow staff users to create destinations */}
+                                    {user?.is_impersonated && (
+                                        <LemonButton
+                                            type="primary"
+                                            icon={<IconPlusSmall />}
+                                            tooltip="Staff users can create destinations as an override"
+                                            to={target.url}
+                                        />
+                                    )}
                                 </span>
                             )
                         },

frontend/src/scenes/pipeline/pipelineBatchExportConfigurationLogic.tsx

@@ -215,7 +215,7 @@ export const pipelineBatchExportConfigurationLogic = kea<pipelineBatchExportConf
         setSavedConfiguration: (configuration: Record<string, any>) => ({ configuration }),
         setSelectedModel: (model: string) => ({ model }),
     }),
-    loaders(({ props, values, actions }) => ({
+    loaders(({ props, actions }) => ({
         batchExportConfig: [
             null as BatchExportConfiguration | null,
             {
@@ -226,13 +226,6 @@ export const pipelineBatchExportConfigurationLogic = kea<pipelineBatchExportConf
                     return null
                 },
                 updateBatchExportConfig: async (formdata) => {
-                    if (
-                        (!values.batchExportConfig || (values.batchExportConfig.paused && formdata.paused !== true)) &&
-                        !values.canEnableNewDestinations
-                    ) {
-                        lemonToast.error('Data pipelines add-on is required for enabling new destinations.')
-                        return null
-                    }
                     const { name, destination, interval, paused, created_at, start_at, end_at, model, ...config } =
                         formdata
                     const destinationObj = {

posthog/api/test/batch_exports/fixtures.py

@@ -0,0 +1,13 @@
+from posthog.api.test.test_organization import create_organization as create_organization_base
+from posthog.constants import AvailableFeature
+from posthog.models import Organization
+
+
+def create_organization(name: str, has_data_pipelines_feature: bool = True) -> Organization:
+    organization = create_organization_base(name)
+    if has_data_pipelines_feature:
+        organization.available_product_features = [
+            {"key": AvailableFeature.DATA_PIPELINES, "name": AvailableFeature.DATA_PIPELINES}
+        ]
+        organization.save()
+    return organization

posthog/api/test/batch_exports/test_backfill.py

@@ -10,7 +10,7 @@
     backfill_batch_export,
     create_batch_export_ok,
 )
-from posthog.api.test.test_organization import create_organization
+from posthog.api.test.batch_exports.fixtures import create_organization
 from posthog.api.test.test_team import create_team
 from posthog.api.test.test_user import create_user
 from posthog.temporal.common.client import sync_connect

posthog/api/test/batch_exports/test_create.py

@@ -9,8 +9,8 @@
 from rest_framework import status
 
 from posthog.api.test.batch_exports.conftest import describe_schedule, start_test_worker
+from posthog.api.test.batch_exports.fixtures import create_organization
 from posthog.api.test.batch_exports.operations import create_batch_export
-from posthog.api.test.test_organization import create_organization
 from posthog.api.test.test_team import create_team
 from posthog.api.test.test_user import create_user
 from posthog.batch_exports.models import BatchExport

posthog/api/test/batch_exports/test_delete.py

@@ -8,14 +8,14 @@
 from temporalio.service import RPCError
 
 from posthog.api.test.batch_exports.conftest import start_test_worker
+from posthog.api.test.batch_exports.fixtures import create_organization
 from posthog.api.test.batch_exports.operations import (
     backfill_batch_export_ok,
     create_batch_export_ok,
     delete_batch_export,
     delete_batch_export_ok,
     get_batch_export,
 )
-from posthog.api.test.test_organization import create_organization
 from posthog.api.test.test_team import create_team
 from posthog.api.test.test_user import create_user
 from posthog.temporal.common.client import sync_connect

posthog/api/test/batch_exports/test_get.py

@@ -7,7 +7,7 @@
     create_batch_export_ok,
     get_batch_export,
 )
-from posthog.api.test.test_organization import create_organization
+from posthog.api.test.batch_exports.fixtures import create_organization
 from posthog.api.test.test_team import create_team
 from posthog.api.test.test_user import create_user
 from posthog.temporal.common.client import sync_connect

posthog/api/test/batch_exports/test_list.py

@@ -6,7 +6,7 @@
     delete_batch_export_ok,
     list_batch_exports_ok,
 )
-from posthog.api.test.test_organization import create_organization
+from posthog.api.test.batch_exports.fixtures import create_organization
 from posthog.api.test.test_team import create_team
 from posthog.api.test.test_user import create_user
 

posthog/api/test/batch_exports/test_pause.py

@@ -14,7 +14,7 @@
     unpause_batch_export,
     unpause_batch_export_ok,
 )
-from posthog.api.test.test_organization import create_organization
+from posthog.api.test.batch_exports.fixtures import create_organization
 from posthog.api.test.test_team import create_team
 from posthog.api.test.test_user import create_user
 from posthog.batch_exports.service import batch_export_delete_schedule

posthog/api/test/batch_exports/test_runs.py

@@ -15,7 +15,7 @@
     get_batch_export_runs,
     get_batch_export_runs_ok,
 )
-from posthog.api.test.test_organization import create_organization
+from posthog.api.test.batch_exports.fixtures import create_organization
 from posthog.api.test.test_team import create_team
 from posthog.api.test.test_user import create_user
 from posthog.temporal.common.client import sync_connect

posthog/api/test/batch_exports/test_update.py

@@ -14,7 +14,7 @@
     patch_batch_export,
     put_batch_export,
 )
-from posthog.api.test.test_organization import create_organization
+from posthog.api.test.batch_exports.fixtures import create_organization
 from posthog.api.test.test_team import create_team
 from posthog.api.test.test_user import create_user
 from posthog.batch_exports.service import sync_batch_export

posthog/api/test/test_app_metrics.py

@@ -8,6 +8,7 @@
 from posthog.api.test.batch_exports.conftest import start_test_worker
 from posthog.api.test.batch_exports.operations import create_batch_export_ok
 from posthog.batch_exports.models import BatchExportRun
+from posthog.constants import AvailableFeature
 from posthog.models.activity_logging.activity_log import Detail, Trigger, log_activity
 from posthog.models.plugin import Plugin, PluginConfig
 from posthog.models.utils import UUIDT
@@ -27,6 +28,11 @@ def setUp(self):
         self.plugin = Plugin.objects.create(organization=self.organization)
         self.plugin_config = PluginConfig.objects.create(plugin=self.plugin, team=self.team, enabled=True, order=1)
 
+        self.organization.available_product_features = [
+            {"key": AvailableFeature.DATA_PIPELINES, "name": AvailableFeature.DATA_PIPELINES}
+        ]
+        self.organization.save()
+
     def test_retrieve(self):
         create_app_metric(
             team_id=self.team.pk,

posthog/api/test/test_team.py

@@ -458,7 +458,10 @@ def test_delete_bulky_postgres_data(self):
         def test_delete_batch_exports(self):
             self.organization_membership.level = OrganizationMembership.Level.ADMIN
             self.organization_membership.save()
-
+            self.organization.available_product_features = [
+                {"key": AvailableFeature.DATA_PIPELINES, "name": AvailableFeature.DATA_PIPELINES}
+            ]
+            self.organization.save()
             team: Team = Team.objects.create_with_data(initiating_user=self.user, organization=self.organization)
 
             destination_data = {
@@ -486,16 +489,16 @@ def test_delete_batch_exports(self):
                     json.dumps(batch_export_data),
                     content_type="application/json",
                 )
-                self.assertEqual(response.status_code, 201)
+                assert response.status_code == 201, response.json()
 
                 batch_export = response.json()
                 batch_export_id = batch_export["id"]
 
                 response = self.client.delete(f"/api/environments/{team.id}")
-                self.assertEqual(response.status_code, 204)
+                assert response.status_code == 204, response.json()
 
                 response = self.client.get(f"/api/environments/{team.id}/batch_exports/{batch_export_id}")
-                self.assertEqual(response.status_code, 404)
+                assert response.status_code == 404, response.json()
 
                 with self.assertRaises(RPCError):
                     describe_schedule(temporal, batch_export_id)

posthog/batch_exports/http.py

@@ -1,5 +1,6 @@
 import datetime as dt
 from typing import Any, TypedDict, cast
+from loginas.utils import is_impersonated_session
 
 import posthoganalytics
 import structlog
@@ -31,6 +32,7 @@
     sync_batch_export,
     unpause_batch_export,
 )
+from posthog.constants import AvailableFeature
 from posthog.hogql import ast, errors
 from posthog.hogql.hogql import HogQLContext
 from posthog.hogql.parser import parse_select
@@ -245,6 +247,20 @@ class Meta:
         ]
         read_only_fields = ["id", "team_id", "created_at", "last_updated_at", "latest_runs", "schema"]
 
+    def validate(self, attrs: dict) -> dict:
+        team = self.context["get_team"]()
+        attrs["team"] = team
+
+        has_addon = team.organization.is_feature_available(AvailableFeature.DATA_PIPELINES)
+
+        if not has_addon:
+            # Check if the user is impersonated - if so we allow changes as it could be an admin user fixing things
+
+            if not is_impersonated_session(self.context["request"]):
+                raise serializers.ValidationError("The Data Pipelines addon is required for batch exports.")
+
+        return attrs
+
     def create(self, validated_data: dict) -> BatchExport:
         """Create a BatchExport."""
         destination_data = validated_data.pop("destination")