From fa0a69d7460bd4e1e55f559665ac0523fa25e52e Mon Sep 17 00:00:00 2001
From: Scott Weber <scott@beeper.com>
Date: Tue, 31 Jan 2023 14:11:01 -0500
Subject: [PATCH] Allow setting segment_user_id in config

---
 config/config.go    | 3 ++-
 example-config.yaml | 2 ++
 main.go             | 4 ++++
 segment.go          | 9 ++++++++-
 4 files changed, 16 insertions(+), 2 deletions(-)

diff --git a/config/config.go b/config/config.go
index 6398bd3..ea5463f 100644
--- a/config/config.go
+++ b/config/config.go
@@ -24,7 +24,8 @@ import (
 type Config struct {
 	*bridgeconfig.BaseConfig `yaml:",inline"`
 
-	SegmentKey string `yaml:"segment_key"`
+	SegmentKey    string `yaml:"segment_key"`
+	SegmentUserId string `yaml:"segment_user_id"`
 
 	Metrics struct {
 		Enabled bool   `yaml:"enabled"`
diff --git a/example-config.yaml b/example-config.yaml
index 4e2faaa..542978a 100644
--- a/example-config.yaml
+++ b/example-config.yaml
@@ -72,6 +72,8 @@ appservice:
 
 # Segment API key to track some events, like provisioning API login and encryption errors.
 segment_key: null
+# Optional user_id to use when sending Segment events. If null, defaults to using mxID.
+segment_user_id: null
 
 # Prometheus config.
 metrics:
diff --git a/main.go b/main.go
index 66d778b..23f4bb9 100644
--- a/main.go
+++ b/main.go
@@ -93,8 +93,12 @@ func (br *WABridge) Init() {
 
 	Segment.log = br.Log.Sub("Segment")
 	Segment.key = br.Config.SegmentKey
+	Segment.userID = br.Config.SegmentUserId
 	if Segment.IsEnabled() {
 		Segment.log.Infoln("Segment metrics are enabled")
+		if Segment.userID != "" {
+			Segment.log.Infoln("Overriding Segment user_id with %v", Segment.userID)
+		}
 	}
 
 	br.DB = database.New(br.Bridge.DB, br.Log.Sub("Database"))
diff --git a/segment.go b/segment.go
index c943cb2..16b6e53 100644
--- a/segment.go
+++ b/segment.go
@@ -30,6 +30,7 @@ const SegmentURL = "https://api.segment.io/v1/track"
 
 type SegmentClient struct {
 	key    string
+	userID string
 	log    log.Logger
 	client http.Client
 }
@@ -38,8 +39,14 @@ var Segment SegmentClient
 
 func (sc *SegmentClient) trackSync(userID id.UserID, event string, properties map[string]interface{}) error {
 	var buf bytes.Buffer
+	var segmentUserID string
+	if Segment.userID != "" {
+		segmentUserID = Segment.userID
+	} else {
+		segmentUserID = userID.String()
+	}
 	err := json.NewEncoder(&buf).Encode(map[string]interface{}{
-		"userId":     userID,
+		"userId":     segmentUserID,
 		"event":      event,
 		"properties": properties,
 	})