chore: add tenantId and orgName in usage (#4399)

* feat: add tenantId and orgName in usage

* fix: update regex

* fix: if else logic updated
This commit is contained in:
Nityananda Gohain 2024-01-20 01:24:09 +05:30 committed by GitHub
parent a8d70206ab
commit f3fdd2dd6c
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 36 additions and 3 deletions

View File

@ -193,7 +193,7 @@ func NewServer(serverOptions *ServerOptions) (*Server, error) {
} }
// start the usagemanager // start the usagemanager
usageManager, err := usage.New("sqlite", localDB, lm.GetRepo(), reader.GetConn()) usageManager, err := usage.New("sqlite", modelDao, lm.GetRepo(), reader.GetConn())
if err != nil { if err != nil {
return nil, err return nil, err
} }

View File

@ -20,6 +20,8 @@ type Usage struct {
TimeStamp time.Time `json:"timestamp"` TimeStamp time.Time `json:"timestamp"`
Count int64 `json:"count"` Count int64 `json:"count"`
Size int64 `json:"size"` Size int64 `json:"size"`
OrgName string `json:"orgName"`
TenantId string `json:"tenantId"`
} }
type UsageDB struct { type UsageDB struct {

View File

@ -4,6 +4,8 @@ import (
"context" "context"
"encoding/json" "encoding/json"
"fmt" "fmt"
"os"
"regexp"
"strings" "strings"
"sync/atomic" "sync/atomic"
"time" "time"
@ -11,10 +13,10 @@ import (
"github.com/ClickHouse/clickhouse-go/v2" "github.com/ClickHouse/clickhouse-go/v2"
"github.com/go-co-op/gocron" "github.com/go-co-op/gocron"
"github.com/google/uuid" "github.com/google/uuid"
"github.com/jmoiron/sqlx"
"go.uber.org/zap" "go.uber.org/zap"
"go.signoz.io/signoz/ee/query-service/dao"
licenseserver "go.signoz.io/signoz/ee/query-service/integrations/signozio" licenseserver "go.signoz.io/signoz/ee/query-service/integrations/signozio"
"go.signoz.io/signoz/ee/query-service/license" "go.signoz.io/signoz/ee/query-service/license"
"go.signoz.io/signoz/ee/query-service/model" "go.signoz.io/signoz/ee/query-service/model"
@ -38,15 +40,29 @@ type Manager struct {
licenseRepo *license.Repo licenseRepo *license.Repo
scheduler *gocron.Scheduler scheduler *gocron.Scheduler
modelDao dao.ModelDao
tenantID string
} }
func New(dbType string, db *sqlx.DB, licenseRepo *license.Repo, clickhouseConn clickhouse.Conn) (*Manager, error) { func New(dbType string, modelDao dao.ModelDao, licenseRepo *license.Repo, clickhouseConn clickhouse.Conn) (*Manager, error) {
hostNameRegex := regexp.MustCompile(`tcp://(?P<hostname>.*):`)
hostNameRegexMatches := hostNameRegex.FindStringSubmatch(os.Getenv("ClickHouseUrl"))
tenantID := ""
if len(hostNameRegexMatches) == 2 {
tenantID = hostNameRegexMatches[1]
tenantID = strings.TrimRight(tenantID, "-clickhouse")
}
m := &Manager{ m := &Manager{
// repository: repo, // repository: repo,
clickhouseConn: clickhouseConn, clickhouseConn: clickhouseConn,
licenseRepo: licenseRepo, licenseRepo: licenseRepo,
scheduler: gocron.NewScheduler(time.UTC).Every(1).Day().At("00:00"), // send usage every at 00:00 UTC scheduler: gocron.NewScheduler(time.UTC).Every(1).Day().At("00:00"), // send usage every at 00:00 UTC
modelDao: modelDao,
tenantID: tenantID,
} }
return m, nil return m, nil
} }
@ -123,6 +139,19 @@ func (lm *Manager) UploadUsage() {
zap.S().Info("uploading usage data") zap.S().Info("uploading usage data")
// Try to get the org name
orgName := ""
orgNames, err := lm.modelDao.GetOrgs(ctx)
if err != nil {
zap.S().Errorf("failed to get org data: %v", zap.Error(err))
} else {
if len(orgNames) != 1 {
zap.S().Errorf("expected one org but got %d orgs", len(orgNames))
} else {
orgName = orgNames[0].Name
}
}
usagesPayload := []model.Usage{} usagesPayload := []model.Usage{}
for _, usage := range usages { for _, usage := range usages {
usageDataBytes, err := encryption.Decrypt([]byte(usage.ExporterID[:32]), []byte(usage.Data)) usageDataBytes, err := encryption.Decrypt([]byte(usage.ExporterID[:32]), []byte(usage.Data))
@ -142,6 +171,8 @@ func (lm *Manager) UploadUsage() {
usageData.ExporterID = usage.ExporterID usageData.ExporterID = usage.ExporterID
usageData.Type = usage.Type usageData.Type = usage.Type
usageData.Tenant = usage.Tenant usageData.Tenant = usage.Tenant
usageData.OrgName = orgName
usageData.TenantId = lm.tenantID
usagesPayload = append(usagesPayload, usageData) usagesPayload = append(usagesPayload, usageData)
} }