Implement per-event enrichment (#351)

* Example of singular enrichment

* Added enrichment for other events.

Author: bsneed

Sources/Segment/Analytics.swift

@@ -93,11 +93,11 @@ public class Analytics {
         Self.removeActiveWriteKey(configuration.values.writeKey)
     }
 
-    internal func process<E: RawEvent>(incomingEvent: E) {
+    internal func process<E: RawEvent>(incomingEvent: E, enrichments: [EnrichmentClosure]? = nil) {
         guard enabled == true else { return }
         let event = incomingEvent.applyRawEventData(store: store)
 
-        _ = timeline.process(incomingEvent: event)
+        _ = timeline.process(incomingEvent: event, enrichments: enrichments)
 
         let flushPolicies = configuration.values.flushPolicies
         for policy in flushPolicies {

Sources/Segment/Errors.swift

@@ -27,6 +27,8 @@ public enum AnalyticsError: Error {
     case jsonUnknown(Error)
 
     case pluginError(Error)
+    
+    case enrichmentError(String)
 }
 
 extension Analytics {

Sources/Segment/Events.swift

@@ -224,3 +224,232 @@ extension Analytics {
         process(incomingEvent: event)
     }
 }
+
+// MARK: - Enrichment event signatures
+
+extension Analytics {
+    // Tracks an event performed by a user, including some additional event properties.
+    /// - Parameters:
+    ///   - name: Name of the action, e.g., 'Purchased a T-Shirt'
+    ///   - properties: Properties specific to the named event. For example, an event with
+    ///     the name 'Purchased a Shirt' might have properties like revenue or size.
+    ///   - enrichments: Enrichments to be applied to this specific event only, or `nil` for none.
+    public func track<P: Codable>(name: String, properties: P?, enrichments: [EnrichmentClosure]?) {
+        do {
+            if let properties = properties {
+                let jsonProperties = try JSON(with: properties)
+                let event = TrackEvent(event: name, properties: jsonProperties)
+                process(incomingEvent: event, enrichments: enrichments)
+            } else {
+                let event = TrackEvent(event: name, properties: nil)
+                process(incomingEvent: event, enrichments: enrichments)
+            }
+        } catch {
+            reportInternalError(error, fatal: true)
+        }
+    }
+    
+    /// Tracks an event performed by a user.
+    /// - Parameters:
+    ///   - name: Name of the action, e.g., 'Purchased a T-Shirt'
+    ///   - enrichments: Enrichments to be applied to this specific event only, or `nil` for none.
+    public func track(name: String, enrichments: [EnrichmentClosure]?) {
+        track(name: name, properties: nil as TrackEvent?, enrichments: enrichments)
+    }
+    
+    /// Tracks an event performed by a user, including some additional event properties.
+    /// - Parameters:
+    ///   - name: Name of the action, e.g., 'Purchased a T-Shirt'
+    ///   - properties: A dictionary or properties specific to the named event.
+    ///     For example, an event with the name 'Purchased a Shirt' might have properties
+    ///     like revenue or size.
+    ///   - enrichments: Enrichments to be applied to this specific event only, or `nil` for none.
+    public func track(name: String, properties: [String: Any]?, enrichments: [EnrichmentClosure]?) {
+        var props: JSON? = nil
+        if let properties = properties {
+            do {
+                props = try JSON(properties)
+            } catch {
+                reportInternalError(error, fatal: true)
+            }
+        }
+        let event = TrackEvent(event: name, properties: props)
+        process(incomingEvent: event, enrichments: enrichments)
+    }
+
+    /// Associate a user with their unique ID and record traits about them.
+    /// - Parameters:
+    ///   - userId: A database ID for this user. If you don't have a userId
+    ///     but want to record traits, just pass traits into the event and they will be associated
+    ///     with the anonymousId of that user.  In the case when user logs out, make sure to
+    ///     call ``reset()`` to clear the user's identity info. For more information on how we
+    ///     generate the UUID and Apple's policies on IDs, see
+    ///      https://segment.io/libraries/ios#ids
+    ///  - traits: A dictionary of traits you know about the user. Things like: email, name, plan, etc.
+    ///  - enrichments: Enrichments to be applied to this specific event only, or `nil` for none.
+    public func identify<T: Codable>(userId: String, traits: T?, enrichments: [EnrichmentClosure]?) {
+        do {
+            if let traits = traits {
+                let jsonTraits = try JSON(with: traits)
+                store.dispatch(action: UserInfo.SetUserIdAndTraitsAction(userId: userId, traits: jsonTraits))
+                let event = IdentifyEvent(userId: userId, traits: jsonTraits)
+                process(incomingEvent: event, enrichments: enrichments)
+            } else {
+                store.dispatch(action: UserInfo.SetUserIdAndTraitsAction(userId: userId, traits: nil))
+                let event = IdentifyEvent(userId: userId, traits: nil)
+                process(incomingEvent: event, enrichments: enrichments)
+            }
+        } catch {
+            reportInternalError(error, fatal: true)
+        }
+    }
+    
+    /// Associate a user with their unique ID and record traits about them.
+    /// - Parameters:
+    ///   - traits: A dictionary of traits you know about the user. Things like: email, name, plan, etc.
+    ///   - enrichments: Enrichments to be applied to this specific event only, or `nil` for none.
+    public func identify<T: Codable>(traits: T, enrichments: [EnrichmentClosure]?) {
+        do {
+            let jsonTraits = try JSON(with: traits)
+            store.dispatch(action: UserInfo.SetTraitsAction(traits: jsonTraits))
+            let event = IdentifyEvent(traits: jsonTraits)
+            process(incomingEvent: event, enrichments: enrichments)
+        } catch {
+            reportInternalError(error, fatal: true)
+        }
+    }
+
+    /// Associate a user with their unique ID and record traits about them.
+    /// - Parameters:
+    ///   - userId: A database ID for this user.
+    ///     For more information on how we generate the UUID and Apple's policies on IDs, see
+    ///     https://segment.io/libraries/ios#ids
+    ///   - enrichments: Enrichments to be applied to this specific event only, or `nil` for none.
+    /// In the case when user logs out, make sure to call ``reset()`` to clear user's identity info.
+    public func identify(userId: String, enrichments: [EnrichmentClosure]?) {
+        let event = IdentifyEvent(userId: userId, traits: nil)
+        store.dispatch(action: UserInfo.SetUserIdAction(userId: userId))
+        process(incomingEvent: event, enrichments: enrichments)
+    }
+
+    /// Associate a user with their unique ID and record traits about them.
+    /// - Parameters:
+    ///   - userId: A database ID for this user. If you don't have a userId
+    ///     but want to record traits, just pass traits into the event and they will be associated
+    ///     with the anonymousId of that user.  In the case when user logs out, make sure to
+    ///     call ``reset()`` to clear the user's identity info. For more information on how we
+    ///     generate the UUID and Apple's policies on IDs, see
+    ///      https://segment.io/libraries/ios#ids
+    ///   - traits: A dictionary of traits you know about the user. Things like: email, name, plan, etc.
+    ///   - enrichments: Enrichments to be applied to this specific event only, or `nil` for none.
+    /// In the case when user logs out, make sure to call ``reset()`` to clear user's identity info.
+    public func identify(userId: String, traits: [String: Any]? = nil, enrichments: [EnrichmentClosure]?) {
+        do {
+            if let traits = traits {
+                let traits = try JSON(traits as Any)
+                store.dispatch(action: UserInfo.SetUserIdAndTraitsAction(userId: userId, traits: traits))
+                let event = IdentifyEvent(userId: userId, traits: traits)
+                process(incomingEvent: event, enrichments: enrichments)
+            } else {
+                store.dispatch(action: UserInfo.SetUserIdAndTraitsAction(userId: userId, traits: nil))
+                let event = IdentifyEvent(userId: userId, traits: nil)
+                process(incomingEvent: event, enrichments: enrichments)
+            }
+        } catch {
+            reportInternalError(error, fatal: true)
+        }
+    }
+    
+    /// Track a screen change with a title, category and other properties.
+    /// - Parameters:
+    ///   - screenTitle: The title of the screen being tracked.
+    ///   - category: A category to the type of screen if it applies.
+    ///   - properties: Any extra metadata associated with the screen. e.g. method of access, size, etc.
+    ///   - enrichments: Enrichments to be applied to this specific event only, or `nil` for none.
+    public func screen<P: Codable>(title: String, category: String? = nil, properties: P?, enrichments: [EnrichmentClosure]?) {
+        do {
+            if let properties = properties {
+                let jsonProperties = try JSON(with: properties)
+                let event = ScreenEvent(title: title, category: category, properties: jsonProperties)
+                process(incomingEvent: event, enrichments: enrichments)
+            } else {
+                let event = ScreenEvent(title: title, category: category)
+                process(incomingEvent: event, enrichments: enrichments)
+            }
+        } catch {
+            reportInternalError(error, fatal: true)
+        }
+    }
+    
+    /// Track a screen change with a title, category and other properties.
+    /// - Parameters:
+    ///   - screenTitle: The title of the screen being tracked.
+    ///   - category: A category to the type of screen if it applies.
+    ///   - enrichments: Enrichments to be applied to this specific event only, or `nil` for none.
+    public func screen(title: String, category: String? = nil, enrichments: [EnrichmentClosure]?) {
+        screen(title: title, category: category, properties: nil as ScreenEvent?, enrichments: enrichments)
+    }
+    
+    /// Track a screen change with a title, category and other properties.
+    /// - Parameters:
+    ///   - screenTitle: The title of the screen being tracked.
+    ///   - category: A category to the type of screen if it applies.
+    ///   - properties: Any extra metadata associated with the screen. e.g. method of access, size, etc.
+    ///   - enrichments: Enrichments to be applied to this specific event only, or `nil` for none.
+    public func screen(title: String, category: String? = nil, properties: [String: Any]?, enrichments: [EnrichmentClosure]?) {
+        // if properties is nil, this is the event that'll get used.
+        var event = ScreenEvent(title: title, category: category, properties: nil)
+        // if we have properties, get a new one rolling.
+        if let properties = properties {
+            do {
+                let jsonProperties = try JSON(properties)
+                event = ScreenEvent(title: title, category: category, properties: jsonProperties)
+            } catch {
+                reportInternalError(error, fatal: true)
+            }
+        }
+        process(incomingEvent: event, enrichments: enrichments)
+    }
+    
+    public func group<T: Codable>(groupId: String, traits: T?, enrichments: [EnrichmentClosure]?) {
+        do {
+            if let traits = traits {
+                let jsonTraits = try JSON(with: traits)
+                let event = GroupEvent(groupId: groupId, traits: jsonTraits)
+                process(incomingEvent: event)
+            } else {
+                let event = GroupEvent(groupId: groupId)
+                process(incomingEvent: event)
+            }
+        } catch {
+            reportInternalError(error, fatal: true)
+        }
+    }
+    
+    public func group(groupId: String, enrichments: [EnrichmentClosure]?) {
+        group(groupId: groupId, traits: nil as GroupEvent?, enrichments: enrichments)
+    }
+    
+    /// Associate a user with a group such as a company, organization, project, etc.
+    /// - Parameters:
+    ///   - groupId: A unique identifier for the group identification in your system.
+    ///   - traits: Traits of the group you may be interested in such as email, phone or name.
+    public func group(groupId: String, traits: [String: Any]?, enrichments: [EnrichmentClosure]?) {
+        var event = GroupEvent(groupId: groupId)
+        if let traits = traits {
+            do {
+                let jsonTraits = try JSON(traits)
+                event = GroupEvent(groupId: groupId, traits: jsonTraits)
+            } catch {
+                reportInternalError(error, fatal: true)
+            }
+        }
+        process(incomingEvent: event, enrichments: enrichments)
+    }
+    
+    public func alias(newId: String, enrichments: [EnrichmentClosure]?) {
+        let event = AliasEvent(newId: newId, previousId: self.userId)
+        store.dispatch(action: UserInfo.SetUserIdAction(userId: newId))
+        process(incomingEvent: event, enrichments: enrichments)
+    }
+}

Sources/Segment/Plugins/Context.swift

@@ -160,4 +160,11 @@ public class Context: PlatformPlugin {
         // other stuff?? ...
     }
 
+    public static func insertOrigin(event: RawEvent?, data: [String: Any]) -> RawEvent? {
+        guard var working = event else { return event }
+        if let newContext = try? working.context?.add(value: data, forKey: "__eventOrigin") {
+            working.context = newContext
+        }
+        return working
+    }
 }

Sources/Segment/Timeline.swift

@@ -25,11 +25,21 @@ public class Timeline {
     }
 
     @discardableResult
-    internal func process<E: RawEvent>(incomingEvent: E) -> E? {
+    internal func process<E: RawEvent>(incomingEvent: E, enrichments: [EnrichmentClosure]? = nil) -> E? {
         // apply .before and .enrichment types first ...
         let beforeResult = applyPlugins(type: .before, event: incomingEvent)
         // .enrichment here is akin to source middleware in the old analytics-ios.
-        let enrichmentResult = applyPlugins(type: .enrichment, event: beforeResult)
+        var enrichmentResult = applyPlugins(type: .enrichment, event: beforeResult)
+        
+        if let enrichments {
+            for closure in enrichments {
+                if let result = closure(enrichmentResult) as? E {
+                    enrichmentResult = result
+                } else {
+                    Analytics.reportInternalError(AnalyticsError.enrichmentError("The given enrichment attempted to change the event type!"))
+                }
+            }
+        }
 
         // once the event enters a destination, we don't want
         // to know about changes that happen there. those changes

Tests/Segment-Tests/Analytics_Tests.swift

@@ -958,6 +958,24 @@ final class Analytics_Tests: XCTestCase {
         XCTAssertEqual(outputReader.lastEvent?.anonymousId, "blah-111")
         XCTAssertEqual(anonIdGenerator.currentId, "blah-111")
         XCTAssertEqual(outputReader.lastEvent?.anonymousId, anonIdGenerator.currentId)
+    }
+    
+    func testSingularEnrichment() throws {
+        let analytics = Analytics(configuration: Configuration(writeKey: "test"))
+        let outputReader = OutputReaderPlugin()
+        analytics.add(plugin: outputReader)
+
+        waitUntilStarted(analytics: analytics)
 
+        let addEventOrigin: EnrichmentClosure = { event in
+            return Context.insertOrigin(event: event, data: [
+                "type": "mobile"
+            ])
+        }
+
+        analytics.track(name: "enrichment check", enrichments: [addEventOrigin])
+
+        let trackEvent: TrackEvent? = outputReader.lastEvent as? TrackEvent
+        XCTAssertEqual(trackEvent?.context?.value(forKeyPath: "__eventOrigin.type"), "mobile")
     }
 }