From b5ce7ad06050adb433b7274095fd743c583c66d7 Mon Sep 17 00:00:00 2001
From: Heiru Wu <heiru.wu@instill.tech>
Date: Tue, 23 Jan 2024 14:06:50 +0800
Subject: [PATCH] refactor(usage): add org data in usage report

---
 go.mod             |  4 ++--
 go.sum             |  8 ++++----
 pkg/usage/usage.go | 22 ++++++++++++++++++++--
 3 files changed, 26 insertions(+), 8 deletions(-)

diff --git a/go.mod b/go.mod
index 202ca9c..8fc31b3 100644
--- a/go.mod
+++ b/go.mod
@@ -15,7 +15,7 @@ require (
 	github.com/grpc-ecosystem/grpc-gateway/v2 v2.16.0
 	github.com/iancoleman/strcase v0.2.0
 	github.com/influxdata/influxdb-client-go/v2 v2.12.3
-	github.com/instill-ai/protogen-go v0.3.3-alpha.0.20231210131526-67e990838339
+	github.com/instill-ai/protogen-go v0.3.3-alpha.0.20240122214718-7d090df83765
 	github.com/instill-ai/usage-client v0.2.4-alpha.0.20231206035716-4c05f872df97
 	github.com/instill-ai/x v0.3.0-alpha
 	github.com/knadh/koanf v1.4.4
@@ -39,7 +39,7 @@ require (
 	golang.org/x/net v0.17.0
 	google.golang.org/genproto/googleapis/api v0.0.0-20230530153820-e85fd2cbaebc
 	google.golang.org/genproto/googleapis/rpc v0.0.0-20230530153820-e85fd2cbaebc
-	google.golang.org/grpc v1.56.0
+	google.golang.org/grpc v1.56.3
 	google.golang.org/protobuf v1.30.0
 	gorm.io/datatypes v1.2.0
 	gorm.io/driver/postgres v1.5.0
diff --git a/go.sum b/go.sum
index 60f30e4..708a2ac 100644
--- a/go.sum
+++ b/go.sum
@@ -1107,8 +1107,8 @@ github.com/influxdata/line-protocol/v2 v2.0.0-20210312151457-c52fdecb625a/go.mod
 github.com/influxdata/line-protocol/v2 v2.1.0/go.mod h1:QKw43hdUBg3GTk2iC3iyCxksNj7PX9aUSeYOYE/ceHY=
 github.com/influxdata/line-protocol/v2 v2.2.1 h1:EAPkqJ9Km4uAxtMRgUubJyqAr6zgWM0dznKMLRauQRE=
 github.com/influxdata/line-protocol/v2 v2.2.1/go.mod h1:DmB3Cnh+3oxmG6LOBIxce4oaL4CPj3OmMPgvauXh+tM=
-github.com/instill-ai/protogen-go v0.3.3-alpha.0.20231210131526-67e990838339 h1:Q48Mm+0i6gL4ZYMiHPddMfBQaslk83y3jmPg9T1T7IQ=
-github.com/instill-ai/protogen-go v0.3.3-alpha.0.20231210131526-67e990838339/go.mod h1:q/YL5TZXD9nvmJ7Rih4gY3/B2HT2+GiFdxeZp9D+yE4=
+github.com/instill-ai/protogen-go v0.3.3-alpha.0.20240122214718-7d090df83765 h1:nFS0byEkAMloOQq9ULBKF3WDvTv5cfvv12BrRgaGwcU=
+github.com/instill-ai/protogen-go v0.3.3-alpha.0.20240122214718-7d090df83765/go.mod h1:jhEL0SauySMoPLVvx105DWyThju9sYTbsXIySVCArmM=
 github.com/instill-ai/usage-client v0.2.4-alpha.0.20231206035716-4c05f872df97 h1:WycXqzJP1ihjJwrkxlNP2TQc1DSUxUtfl/PtCpLBa3Y=
 github.com/instill-ai/usage-client v0.2.4-alpha.0.20231206035716-4c05f872df97/go.mod h1:Da8RdKakfxy1iNdvI5FSTcL1lSDtda+b9jOgOEEO68E=
 github.com/instill-ai/x v0.3.0-alpha h1:z9fedROOG2dVHhswBfVwU/hzHuq8/JKSUON7inF+FH8=
@@ -2465,8 +2465,8 @@ google.golang.org/grpc v1.50.0/go.mod h1:ZgQEeidpAuNRZ8iRrlBKXZQP1ghovWIVhdJRyCD
 google.golang.org/grpc v1.50.1/go.mod h1:ZgQEeidpAuNRZ8iRrlBKXZQP1ghovWIVhdJRyCDK+GI=
 google.golang.org/grpc v1.51.0/go.mod h1:wgNDFcnuBGmxLKI/qn4T+m5BtEBYXJPvibbUPsAIPww=
 google.golang.org/grpc v1.52.3/go.mod h1:pu6fVzoFb+NBYNAvQL08ic+lvB2IojljRYuun5vorUY=
-google.golang.org/grpc v1.56.0 h1:+y7Bs8rtMd07LeXmL3NxcTLn7mUkbKZqEpPhMNkwJEE=
-google.golang.org/grpc v1.56.0/go.mod h1:I9bI3vqKfayGqPUAwGdOSu7kt6oIJLixfffKrpXqQ9s=
+google.golang.org/grpc v1.56.3 h1:8I4C0Yq1EjstUzUJzpcRVbuYA2mODtEmpWiQoN/b2nc=
+google.golang.org/grpc v1.56.3/go.mod h1:I9bI3vqKfayGqPUAwGdOSu7kt6oIJLixfffKrpXqQ9s=
 google.golang.org/grpc/cmd/protoc-gen-go-grpc v1.1.0/go.mod h1:6Kw0yEErY5E/yWrBtf03jp27GLLJujG4z/JK95pnjjw=
 google.golang.org/protobuf v0.0.0-20200109180630-ec00e32a8dfd/go.mod h1:DFci5gLYBciE7Vtevhsrf46CRTquxDuWsQurQQe4oz8=
 google.golang.org/protobuf v0.0.0-20200221191635-4d8936d0db64/go.mod h1:kwYJMbMJ01Woi6D6+Kah6886xMZcty6N08ah7+eCXa0=
diff --git a/pkg/usage/usage.go b/pkg/usage/usage.go
index 0f5daf7..b571af3 100644
--- a/pkg/usage/usage.go
+++ b/pkg/usage/usage.go
@@ -85,13 +85,31 @@ func (u *usage) RetrieveUsageData() interface{} {
 		}
 	}
 
-	logger.Debug(fmt.Sprintf("[mgmt-backend] usage data length: %v", len(allUsers)))
+	allOrgs := []*mgmtPB.Organization{}
+	pageToken = ""
+	for {
+		orgs, _, token, err := u.service.ListOrganizationsAdmin(ctx, 100, pageToken, filtering.Filter{})
+		if err != nil {
+			logger.Error(fmt.Sprintf("%s", err))
+			break
+		}
+
+		pageToken = token
+		allOrgs = append(allOrgs, orgs...)
+		if token == "" {
+			break
+		}
+	}
+
+	logger.Debug(fmt.Sprintf("[mgmt-backend] user usage data length: %v", len(allUsers)))
+	logger.Debug(fmt.Sprintf("[mgmt-backend] org usage data length: %v", len(allOrgs)))
 
 	logger.Debug("[mgmt-backend] send usage data...")
 
 	return &usagePB.SessionReport_MgmtUsageData{
 		MgmtUsageData: &usagePB.MgmtUsageData{
-			Usages: allUsers,
+			UserUsages: allUsers,
+			OrgUsages:  allOrgs,
 		},
 	}
 }