diff --git a/ee/query-service/app/server.go b/ee/query-service/app/server.go index a5c7c1db22..ea0b0344ad 100644 --- a/ee/query-service/app/server.go +++ b/ee/query-service/app/server.go @@ -193,7 +193,7 @@ func NewServer(serverOptions *ServerOptions) (*Server, error) { } // 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 { return nil, err } diff --git a/ee/query-service/model/usage.go b/ee/query-service/model/usage.go index d9129531dc..3cedb0532e 100644 --- a/ee/query-service/model/usage.go +++ b/ee/query-service/model/usage.go @@ -20,6 +20,8 @@ type Usage struct { TimeStamp time.Time `json:"timestamp"` Count int64 `json:"count"` Size int64 `json:"size"` + OrgName string `json:"orgName"` + TenantId string `json:"tenantId"` } type UsageDB struct { diff --git a/ee/query-service/usage/manager.go b/ee/query-service/usage/manager.go index 335eae8143..28d4e3020e 100644 --- a/ee/query-service/usage/manager.go +++ b/ee/query-service/usage/manager.go @@ -4,6 +4,8 @@ import ( "context" "encoding/json" "fmt" + "os" + "regexp" "strings" "sync/atomic" "time" @@ -11,10 +13,10 @@ import ( "github.com/ClickHouse/clickhouse-go/v2" "github.com/go-co-op/gocron" "github.com/google/uuid" - "github.com/jmoiron/sqlx" "go.uber.org/zap" + "go.signoz.io/signoz/ee/query-service/dao" 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/model" @@ -38,15 +40,29 @@ type Manager struct { licenseRepo *license.Repo 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.*):`) + hostNameRegexMatches := hostNameRegex.FindStringSubmatch(os.Getenv("ClickHouseUrl")) + + tenantID := "" + if len(hostNameRegexMatches) == 2 { + tenantID = hostNameRegexMatches[1] + tenantID = strings.TrimRight(tenantID, "-clickhouse") + } m := &Manager{ // repository: repo, clickhouseConn: clickhouseConn, licenseRepo: licenseRepo, 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 } @@ -123,6 +139,19 @@ func (lm *Manager) UploadUsage() { 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{} for _, usage := range usages { usageDataBytes, err := encryption.Decrypt([]byte(usage.ExporterID[:32]), []byte(usage.Data)) @@ -142,6 +171,8 @@ func (lm *Manager) UploadUsage() { usageData.ExporterID = usage.ExporterID usageData.Type = usage.Type usageData.Tenant = usage.Tenant + usageData.OrgName = orgName + usageData.TenantId = lm.tenantID usagesPayload = append(usagesPayload, usageData) }