fix(segment): Posthog can identify after segment identifies a user, n…

…ot just during bootstrap. (#1373) This PR changes the behavior of the segment plugin, upon segment calling identify(), to : Set the USER_STATE to identified Set the distinct_id for the current user session Reload all feature flags for the new distinct_id
PostHog · Aug 26, 2024 · bbd08b7 · bbd08b7
1 parent ad550bf
commit bbd08b7
Show file tree

Hide file tree

Showing 4 changed files with 51 additions and 10 deletions.
diff --git a/playground/nextjs/pages/index.tsx b/playground/nextjs/pages/index.tsx
@@ -62,7 +62,9 @@ export default function Home() {
                     Set user properties
                 </button>
 
-                <button onClick={() => posthog?.reset()}>Reset</button>
+                <button onClick={() => posthog?.reset()} id="set-user-properties">
+                    Reset
+                </button>
             </div>
 
             {isClient && (

diff --git a/playground/segment/segment.html b/playground/segment/segment.html
@@ -9,7 +9,8 @@
                 analytics.load("GOEDfA21zZTtR7clsBuDvmBKAtAdZ6Np");
 
                 analytics.ready(() => {
-                    posthog.init('phc_mxO23D02CR3QU7V1u95ae3YEGNAFGuvyUlRpVNwX89B', {
+                    posthog.init('phc_lPwXKkvW2VtByfTsibDzFKAxn02rdVgMUiHWkZ5kW31', {
+                        api_host: "http://localhost:8000",
                         capture_pageview: false,
                         autocapture: false,
                         segment: window.analytics,

diff --git a/src/__tests__/segment.test.ts b/src/__tests__/segment.test.ts
@@ -10,10 +10,12 @@
 import { beforeEach, describe, expect, it, jest } from '@jest/globals'
 
 import { PostHog } from '../posthog-core'
+import { SegmentContext, SegmentPlugin } from '../extensions/segment-integration'
+import { USER_STATE } from '../constants'
 
 describe(`Segment integration`, () => {
     let segment: any
-    let segmentIntegration: any
+    let segmentIntegration: SegmentPlugin
     let posthogName: string
 
     jest.setTimeout(500)
@@ -28,7 +30,7 @@ describe(`Segment integration`, () => {
                 anonymousId: () => 'test-anonymous-id',
                 id: () => 'test-id',
             }),
-            register: (integration: any) => {
+            register: (integration: SegmentPlugin) => {
                 // IMPORTANT: the real register function returns a Promise. We
                 // want to do the same thing and have some way to verify that
                 // the integration is setup in time for the `loaded` callback.
@@ -108,4 +110,41 @@ describe(`Segment integration`, () => {
         expect(posthog.get_distinct_id()).toBe('test-id')
         expect(posthog.get_property('$device_id')).toBe('test-anonymous-id')
     })
+
+    it('should handle segment.identify after bootstrap', async () => {
+        segment.user = () => ({
+            anonymousId: () => 'test-anonymous-id',
+            id: () => '',
+        })
+
+        const posthog = await new Promise<PostHog>((resolve) => {
+            return new PostHog().init(
+                `test-token`,
+                {
+                    debug: true,
+                    persistence: `memory`,
+                    api_host: `https://test.com`,
+                    segment: segment,
+                    loaded: resolve,
+                },
+                posthogName
+            )
+        })
+
+        expect(posthog.get_distinct_id()).not.toEqual('test-id')
+        expect(posthog.persistence?.get_property(USER_STATE)).toEqual('anonymous')
+
+        if (segmentIntegration && segmentIntegration.identify) {
+            segmentIntegration.identify({
+                event: {
+                    event: '$identify',
+                    userId: 'distinguished user',
+                    anonymousId: 'anonymous segment user',
+                },
+            } as unknown as SegmentContext)
+
+            expect(posthog.get_distinct_id()).toEqual('distinguished user')
+            expect(posthog.persistence?.get_property(USER_STATE)).toEqual('identified')
+        }
+    })
 })
diff --git a/src/extensions/segment-integration.ts b/src/extensions/segment-integration.ts
@@ -34,7 +34,7 @@ export type SegmentAnalytics = {
 }
 
 // Loosely based on https://github.com/segmentio/analytics-next/blob/master/packages/core/src/plugins/index.ts
-interface SegmentContext {
+export interface SegmentContext {
     event: {
         event: string
         userId?: string
@@ -45,7 +45,7 @@ interface SegmentContext {
 
 type SegmentFunction = (ctx: SegmentContext) => Promise<SegmentContext> | SegmentContext
 
-interface SegmentPlugin {
+export interface SegmentPlugin {
     name: string
     version: string
     type: 'enrichment'
@@ -72,13 +72,12 @@ const createSegmentIntegration = (posthog: PostHog): SegmentPlugin => {
         }
         if (!ctx.event.userId && ctx.event.anonymousId !== posthog.get_distinct_id()) {
             // This is our only way of detecting that segment's analytics.reset() has been called so we also call it
+            logger.info('Segment integration does not have a userId set, resetting PostHog')
             posthog.reset()
         }
         if (ctx.event.userId && ctx.event.userId !== posthog.get_distinct_id()) {
-            posthog.register({
-                distinct_id: ctx.event.userId,
-            })
-            posthog.reloadFeatureFlags()
+            logger.info('Segment integration has a userId set, identifying with PostHog')
+            posthog.identify(ctx.event.userId)
         }
 
         const additionalProperties = posthog._calculate_event_properties(