From 0db2784d6b0403e78d815d8119ff3661337726ec Mon Sep 17 00:00:00 2001
From: Vishal Sharma <makeavish786@gmail.com>
Date: Fri, 6 Sep 2024 14:46:18 +0530
Subject: [PATCH] =?UTF-8?q?chore:=20calculate=20user=20count=20dynamically?=
 =?UTF-8?q?=20and=20set=20user=20role=20in=20identity=E2=80=A6=20(#5870)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

* chore: calculate user count dynamically and set user role in identity event

* chore: move to callbacks

---------

Co-authored-by: Srikanth Chekuri <srikanth.chekuri92@gmail.com>
---
 pkg/query-service/dao/sqlite/connection.go |  4 ++-
 pkg/query-service/dao/sqlite/rbac.go       | 16 +++++++++
 pkg/query-service/telemetry/telemetry.go   | 38 ++++++++++++++++------
 3 files changed, 47 insertions(+), 11 deletions(-)

diff --git a/pkg/query-service/dao/sqlite/connection.go b/pkg/query-service/dao/sqlite/connection.go
index d7e5ad5de9..a4373d5ecd 100644
--- a/pkg/query-service/dao/sqlite/connection.go
+++ b/pkg/query-service/dao/sqlite/connection.go
@@ -103,6 +103,9 @@ func InitDB(dataSourceName string) (*ModelDaoSqlite, error) {
 		return nil, err
 	}
 
+	telemetry.GetInstance().SetUserCountCallback(mds.GetUserCount)
+	telemetry.GetInstance().SetUserRoleCallback(mds.GetUserRole)
+
 	return mds, nil
 }
 
@@ -140,7 +143,6 @@ func (mds *ModelDaoSqlite) initializeOrgPreferences(ctx context.Context) error {
 
 	users, _ := mds.GetUsers(ctx)
 	countUsers := len(users)
-	telemetry.GetInstance().SetCountUsers(int8(countUsers))
 	if countUsers > 0 {
 		telemetry.GetInstance().SetCompanyDomain(users[countUsers-1].Email)
 		telemetry.GetInstance().SetUserEmail(users[countUsers-1].Email)
diff --git a/pkg/query-service/dao/sqlite/rbac.go b/pkg/query-service/dao/sqlite/rbac.go
index aba9beb065..bb594ac463 100644
--- a/pkg/query-service/dao/sqlite/rbac.go
+++ b/pkg/query-service/dao/sqlite/rbac.go
@@ -612,3 +612,19 @@ func (mds *ModelDaoSqlite) PrecheckLogin(ctx context.Context, email, sourceUrl s
 
 	return resp, nil
 }
+
+func (mds *ModelDaoSqlite) GetUserRole(ctx context.Context, groupId string) (string, error) {
+	role, err := mds.GetGroup(ctx, groupId)
+	if err != nil || role == nil {
+		return "", err
+	}
+	return role.Name, nil
+}
+
+func (mds *ModelDaoSqlite) GetUserCount(ctx context.Context) (int, error) {
+	users, err := mds.GetUsers(ctx)
+	if err != nil {
+		return 0, err
+	}
+	return len(users), nil
+}
diff --git a/pkg/query-service/telemetry/telemetry.go b/pkg/query-service/telemetry/telemetry.go
index c916135f4e..88f3a09542 100644
--- a/pkg/query-service/telemetry/telemetry.go
+++ b/pkg/query-service/telemetry/telemetry.go
@@ -176,16 +176,25 @@ type Telemetry struct {
 	rateLimits    map[string]int8
 	activeUser    map[string]int8
 	patTokenUser  bool
-	countUsers    int8
 	mutex         sync.RWMutex
 
 	alertsInfoCallback func(ctx context.Context) (*model.AlertsInfo, error)
+	userCountCallback  func(ctx context.Context) (int, error)
+	userRoleCallback   func(ctx context.Context, groupId string) (string, error)
 }
 
 func (a *Telemetry) SetAlertsInfoCallback(callback func(ctx context.Context) (*model.AlertsInfo, error)) {
 	a.alertsInfoCallback = callback
 }
 
+func (a *Telemetry) SetUserCountCallback(callback func(ctx context.Context) (int, error)) {
+	a.userCountCallback = callback
+}
+
+func (a *Telemetry) SetUserRoleCallback(callback func(ctx context.Context, groupId string) (string, error)) {
+	a.userRoleCallback = callback
+}
+
 func createTelemetry() {
 	// Do not do anything in CI (not even resolving the outbound IP address)
 	if testing.Testing() {
@@ -259,6 +268,8 @@ func createTelemetry() {
 		metricsTTL, _ := telemetry.reader.GetTTL(ctx, &model.GetTTLParams{Type: constants.MetricsTTL})
 		logsTTL, _ := telemetry.reader.GetTTL(ctx, &model.GetTTLParams{Type: constants.LogsTTL})
 
+		userCount, _ := telemetry.userCountCallback(ctx)
+
 		data := map[string]interface{}{
 			"totalSpans":                            totalSpans,
 			"spansInLastHeartBeatInterval":          spansInLastHeartBeatInterval,
@@ -266,7 +277,7 @@ func createTelemetry() {
 			"getSamplesInfoInLastHeartBeatInterval": getSamplesInfoInLastHeartBeatInterval,
 			"totalLogs":                             totalLogs,
 			"getLogsInfoInLastHeartBeatInterval":    getLogsInfoInLastHeartBeatInterval,
-			"countUsers":                            telemetry.countUsers,
+			"countUsers":                            userCount,
 			"metricsTTLStatus":                      metricsTTL.Status,
 			"tracesTTLStatus":                       traceTTL.Status,
 			"logsTTLStatus":                         logsTTL.Status,
@@ -450,11 +461,22 @@ func (a *Telemetry) IdentifyUser(user *model.User) {
 	if !a.isTelemetryEnabled() || a.isTelemetryAnonymous() {
 		return
 	}
+	// extract user group from user.groupId
+	role, _ := a.userRoleCallback(context.Background(), user.GroupId)
+
 	if a.saasOperator != nil {
-		a.saasOperator.Enqueue(analytics.Identify{
-			UserId: a.userEmail,
-			Traits: analytics.NewTraits().SetName(user.Name).SetEmail(user.Email),
-		})
+		if role != "" {
+			a.saasOperator.Enqueue(analytics.Identify{
+				UserId: a.userEmail,
+				Traits: analytics.NewTraits().SetName(user.Name).SetEmail(user.Email).Set("role", role),
+			})
+		} else {
+			a.saasOperator.Enqueue(analytics.Identify{
+				UserId: a.userEmail,
+				Traits: analytics.NewTraits().SetName(user.Name).SetEmail(user.Email),
+			})
+		}
+
 		a.saasOperator.Enqueue(analytics.Group{
 			UserId:  a.userEmail,
 			GroupId: a.getCompanyDomain(),
@@ -474,10 +496,6 @@ func (a *Telemetry) IdentifyUser(user *model.User) {
 	})
 }
 
-func (a *Telemetry) SetCountUsers(countUsers int8) {
-	a.countUsers = countUsers
-}
-
 func (a *Telemetry) SetUserEmail(email string) {
 	a.userEmail = email
 }