From c32b8638a4f5bb580a85316f1ae37e27cd41d9bd Mon Sep 17 00:00:00 2001 From: Vishal Sharma Date: Wed, 17 May 2023 16:10:43 +0530 Subject: [PATCH] feat: introduce feature_usage table to manage features (#2661) * feat: introduce feature_usage table to manage features * feat: introduce limit on QB alerts and dashboards --- ee/query-service/app/api/featureFlags.go | 6 +- ee/query-service/app/server.go | 9 +- ee/query-service/license/db.go | 78 ++++++++++ ee/query-service/license/manager.go | 43 ++++-- ee/query-service/license/sqlite/init.go | 15 ++ ee/query-service/model/plans.go | 144 ++++++++++++++++-- pkg/query-service/app/dashboards/model.go | 125 ++++++++++++++- pkg/query-service/app/dashboards/provision.go | 9 +- pkg/query-service/app/http_handler.go | 16 +- pkg/query-service/app/server.go | 6 +- pkg/query-service/constants/constants.go | 15 +- pkg/query-service/featureManager/manager.go | 56 +++++-- pkg/query-service/interfaces/featureLookup.go | 5 +- pkg/query-service/model/featureSet.go | 17 ++- pkg/query-service/rules/manager.go | 144 ++++++++++++++++-- 15 files changed, 613 insertions(+), 75 deletions(-) diff --git a/ee/query-service/app/api/featureFlags.go b/ee/query-service/app/api/featureFlags.go index 9c979d17ba..63b36d45c4 100644 --- a/ee/query-service/app/api/featureFlags.go +++ b/ee/query-service/app/api/featureFlags.go @@ -5,6 +5,10 @@ import ( ) func (ah *APIHandler) getFeatureFlags(w http.ResponseWriter, r *http.Request) { - featureSet := ah.FF().GetFeatureFlags() + featureSet, err := ah.FF().GetFeatureFlags() + if err != nil { + ah.HandleError(w, err, http.StatusInternalServerError) + return + } ah.Respond(w, featureSet) } diff --git a/ee/query-service/app/server.go b/ee/query-service/app/server.go index 30daa090b6..315a211f9f 100644 --- a/ee/query-service/app/server.go +++ b/ee/query-service/app/server.go @@ -22,6 +22,8 @@ import ( "go.signoz.io/signoz/ee/query-service/app/db" "go.signoz.io/signoz/ee/query-service/dao" "go.signoz.io/signoz/ee/query-service/interfaces" + baseInterface "go.signoz.io/signoz/pkg/query-service/interfaces" + licensepkg "go.signoz.io/signoz/ee/query-service/license" "go.signoz.io/signoz/ee/query-service/usage" @@ -126,7 +128,8 @@ func NewServer(serverOptions *ServerOptions) (*Server, error) { serverOptions.RuleRepoURL, localDB, reader, - serverOptions.DisableRules) + serverOptions.DisableRules, + lm) if err != nil { return nil, err @@ -544,7 +547,8 @@ func makeRulesManager( ruleRepoURL string, db *sqlx.DB, ch baseint.Reader, - disableRules bool) (*rules.Manager, error) { + disableRules bool, + fm baseInterface.FeatureLookup) (*rules.Manager, error) { // create engine pqle, err := pqle.FromConfigPath(promConfigPath) @@ -571,6 +575,7 @@ func makeRulesManager( Context: context.Background(), Logger: nil, DisableRules: disableRules, + FeatureFlags: fm, } // create Manager diff --git a/ee/query-service/license/db.go b/ee/query-service/license/db.go index a82f0377e2..8d2f7065ff 100644 --- a/ee/query-service/license/db.go +++ b/ee/query-service/license/db.go @@ -2,6 +2,7 @@ package license import ( "context" + "database/sql" "fmt" "time" @@ -9,6 +10,7 @@ import ( "go.signoz.io/signoz/ee/query-service/license/sqlite" "go.signoz.io/signoz/ee/query-service/model" + basemodel "go.signoz.io/signoz/pkg/query-service/model" "go.uber.org/zap" ) @@ -125,3 +127,79 @@ func (r *Repo) UpdatePlanDetails(ctx context.Context, return nil } + +func (r *Repo) CreateFeature(req *basemodel.Feature) *basemodel.ApiError { + + _, err := r.db.Exec( + `INSERT INTO feature_status (name, active, usage, usage_limit, route) + VALUES (?, ?, ?, ?, ?);`, + req.Name, req.Active, req.Usage, req.UsageLimit, req.Route) + if err != nil { + return &basemodel.ApiError{Typ: basemodel.ErrorInternal, Err: err} + } + return nil +} + +func (r *Repo) GetFeature(featureName string) (basemodel.Feature, error) { + + var feature basemodel.Feature + + err := r.db.Get(&feature, + `SELECT * FROM feature_status WHERE name = ?;`, featureName) + if err != nil { + return feature, err + } + if feature.Name == "" { + return feature, basemodel.ErrFeatureUnavailable{Key: featureName} + } + return feature, nil +} + +func (r *Repo) GetAllFeatures() ([]basemodel.Feature, error) { + + var feature []basemodel.Feature + + err := r.db.Select(&feature, + `SELECT * FROM feature_status;`) + if err != nil { + return feature, err + } + + return feature, nil +} + +func (r *Repo) UpdateFeature(req basemodel.Feature) error { + + _, err := r.db.Exec( + `UPDATE feature_status SET active = ?, usage = ?, usage_limit = ?, route = ? WHERE name = ?;`, + req.Active, req.Usage, req.UsageLimit, req.Route, req.Name) + if err != nil { + return err + } + return nil +} + +func (r *Repo) InitFeatures(req basemodel.FeatureSet) error { + // get a feature by name, if it doesn't exist, create it. If it does exist, update it. + for _, feature := range req { + currentFeature, err := r.GetFeature(feature.Name) + if err != nil && err == sql.ErrNoRows { + err := r.CreateFeature(&feature) + if err != nil { + return err + } + continue + } else if err != nil { + return err + } + feature.Usage = currentFeature.Usage + if feature.Usage >= feature.UsageLimit && feature.UsageLimit != -1 { + feature.Active = false + } + err = r.UpdateFeature(feature) + if err != nil { + return err + } + } + return nil +} diff --git a/ee/query-service/license/manager.go b/ee/query-service/license/manager.go index a3e9ba0771..7a1be5118f 100644 --- a/ee/query-service/license/manager.go +++ b/ee/query-service/license/manager.go @@ -96,6 +96,11 @@ func (lm *Manager) SetActive(l *model.License) { lm.activeFeatures = l.FeatureSet // set default features setDefaultFeatures(lm) + + err := lm.InitFeatures(lm.activeFeatures) + if err != nil { + zap.S().Panicf("Couldn't activate features: %v", err) + } if !lm.validatorRunning { // we want to make sure only one validator runs, // we already have lock() so good to go @@ -106,9 +111,7 @@ func (lm *Manager) SetActive(l *model.License) { } func setDefaultFeatures(lm *Manager) { - for k, v := range baseconstants.DEFAULT_FEATURE_SET { - lm.activeFeatures[k] = v - } + lm.activeFeatures = append(lm.activeFeatures, baseconstants.DEFAULT_FEATURE_SET...) } // LoadActiveLicense loads the most recent active license @@ -123,8 +126,13 @@ func (lm *Manager) LoadActiveLicense() error { } else { zap.S().Info("No active license found, defaulting to basic plan") // if no active license is found, we default to basic(free) plan with all default features - lm.activeFeatures = basemodel.BasicPlan + lm.activeFeatures = model.BasicPlan setDefaultFeatures(lm) + err := lm.InitFeatures(lm.activeFeatures) + if err != nil { + zap.S().Error("Couldn't initialize features: ", err) + return err + } } return nil @@ -291,18 +299,31 @@ func (lm *Manager) Activate(ctx context.Context, key string) (licenseResponse *m // CheckFeature will be internally used by backend routines // for feature gating func (lm *Manager) CheckFeature(featureKey string) error { - if value, ok := lm.activeFeatures[featureKey]; ok { - if value { - return nil - } - return basemodel.ErrFeatureUnavailable{Key: featureKey} + feature, err := lm.repo.GetFeature(featureKey) + if err != nil { + return err + } + if feature.Active { + return nil } return basemodel.ErrFeatureUnavailable{Key: featureKey} } // GetFeatureFlags returns current active features -func (lm *Manager) GetFeatureFlags() basemodel.FeatureSet { - return lm.activeFeatures +func (lm *Manager) GetFeatureFlags() (basemodel.FeatureSet, error) { + return lm.repo.GetAllFeatures() +} + +func (lm *Manager) InitFeatures(features basemodel.FeatureSet) error { + return lm.repo.InitFeatures(features) +} + +func (lm *Manager) UpdateFeatureFlag(feature basemodel.Feature) error { + return lm.repo.UpdateFeature(feature) +} + +func (lm *Manager) GetFeatureFlag(key string) (basemodel.Feature, error) { + return lm.repo.GetFeature(key) } // GetRepo return the license repo diff --git a/ee/query-service/license/sqlite/init.go b/ee/query-service/license/sqlite/init.go index a03153659c..e500ddb4aa 100644 --- a/ee/query-service/license/sqlite/init.go +++ b/ee/query-service/license/sqlite/init.go @@ -2,6 +2,7 @@ package sqlite import ( "fmt" + "github.com/jmoiron/sqlx" ) @@ -33,5 +34,19 @@ func InitDB(db *sqlx.DB) error { if err != nil { return fmt.Errorf("Error in creating licenses table: %s", err.Error()) } + + table_schema = `CREATE TABLE IF NOT EXISTS feature_status ( + name TEXT PRIMARY KEY, + active bool, + usage INTEGER DEFAULT 0, + usage_limit INTEGER DEFAULT 0, + route TEXT + );` + + _, err = db.Exec(table_schema) + if err != nil { + return fmt.Errorf("Error in creating feature_status table: %s", err.Error()) + } + return nil } diff --git a/ee/query-service/model/plans.go b/ee/query-service/model/plans.go index c42712f693..52ebd5c5b5 100644 --- a/ee/query-service/model/plans.go +++ b/ee/query-service/model/plans.go @@ -11,21 +11,143 @@ const Enterprise = "ENTERPRISE_PLAN" const DisableUpsell = "DISABLE_UPSELL" var BasicPlan = basemodel.FeatureSet{ - Basic: true, - SSO: false, - DisableUpsell: false, + basemodel.Feature{ + Name: SSO, + Active: false, + Usage: 0, + UsageLimit: -1, + Route: "", + }, + basemodel.Feature{ + Name: basemodel.OSS, + Active: false, + Usage: 0, + UsageLimit: -1, + Route: "", + }, + basemodel.Feature{ + Name: DisableUpsell, + Active: false, + Usage: 0, + UsageLimit: -1, + Route: "", + }, + basemodel.Feature{ + Name: basemodel.SmartTraceDetail, + Active: false, + Usage: 0, + UsageLimit: -1, + Route: "", + }, + basemodel.Feature{ + Name: basemodel.CustomMetricsFunction, + Active: false, + Usage: 0, + UsageLimit: -1, + Route: "", + }, + basemodel.Feature{ + Name: basemodel.QueryBuilderPanels, + Active: true, + Usage: 0, + UsageLimit: 5, + Route: "", + }, + basemodel.Feature{ + Name: basemodel.QueryBuilderAlerts, + Active: true, + Usage: 0, + UsageLimit: 5, + Route: "", + }, } var ProPlan = basemodel.FeatureSet{ - Pro: true, - SSO: true, - basemodel.SmartTraceDetail: true, - basemodel.CustomMetricsFunction: true, + basemodel.Feature{ + Name: SSO, + Active: true, + Usage: 0, + UsageLimit: -1, + Route: "", + }, + basemodel.Feature{ + Name: basemodel.OSS, + Active: false, + Usage: 0, + UsageLimit: -1, + Route: "", + }, + basemodel.Feature{ + Name: basemodel.SmartTraceDetail, + Active: true, + Usage: 0, + UsageLimit: -1, + Route: "", + }, + basemodel.Feature{ + Name: basemodel.CustomMetricsFunction, + Active: true, + Usage: 0, + UsageLimit: -1, + Route: "", + }, + basemodel.Feature{ + Name: basemodel.QueryBuilderPanels, + Active: true, + Usage: 0, + UsageLimit: -1, + Route: "", + }, + basemodel.Feature{ + Name: basemodel.QueryBuilderAlerts, + Active: true, + Usage: 0, + UsageLimit: -1, + Route: "", + }, } var EnterprisePlan = basemodel.FeatureSet{ - Enterprise: true, - SSO: true, - basemodel.SmartTraceDetail: true, - basemodel.CustomMetricsFunction: true, + basemodel.Feature{ + Name: SSO, + Active: true, + Usage: 0, + UsageLimit: -1, + Route: "", + }, + basemodel.Feature{ + Name: basemodel.OSS, + Active: false, + Usage: 0, + UsageLimit: -1, + Route: "", + }, + basemodel.Feature{ + Name: basemodel.SmartTraceDetail, + Active: true, + Usage: 0, + UsageLimit: -1, + Route: "", + }, + basemodel.Feature{ + Name: basemodel.CustomMetricsFunction, + Active: true, + Usage: 0, + UsageLimit: -1, + Route: "", + }, + basemodel.Feature{ + Name: basemodel.QueryBuilderPanels, + Active: true, + Usage: 0, + UsageLimit: -1, + Route: "", + }, + basemodel.Feature{ + Name: basemodel.QueryBuilderAlerts, + Active: true, + Usage: 0, + UsageLimit: -1, + Route: "", + }, } diff --git a/pkg/query-service/app/dashboards/model.go b/pkg/query-service/app/dashboards/model.go index 56adf9aae7..2edf0624ef 100644 --- a/pkg/query-service/app/dashboards/model.go +++ b/pkg/query-service/app/dashboards/model.go @@ -14,6 +14,7 @@ import ( "github.com/gosimple/slug" "github.com/jmoiron/sqlx" "github.com/mitchellh/mapstructure" + "go.signoz.io/signoz/pkg/query-service/interfaces" "go.signoz.io/signoz/pkg/query-service/model" "go.uber.org/zap" ) @@ -131,7 +132,7 @@ func (c *Data) Scan(src interface{}) error { } // CreateDashboard creates a new dashboard -func CreateDashboard(data map[string]interface{}) (*Dashboard, *model.ApiError) { +func CreateDashboard(data map[string]interface{}, fm interfaces.FeatureLookup) (*Dashboard, *model.ApiError) { dash := &Dashboard{ Data: data, } @@ -146,6 +147,13 @@ func CreateDashboard(data map[string]interface{}) (*Dashboard, *model.ApiError) return nil, &model.ApiError{Typ: model.ErrorExec, Err: err} } + if countTraceAndLogsPanel(data) > 0 { + fErr := checkFeatureUsage(fm, countTraceAndLogsPanel(data)) + if fErr != nil { + return nil, fErr + } + } + // db.Prepare("Insert into dashboards where") result, err := db.Exec("INSERT INTO dashboards (uuid, created_at, updated_at, data) VALUES ($1, $2, $3, $4)", dash.Uuid, dash.CreatedAt, dash.UpdatedAt, map_data) @@ -160,6 +168,11 @@ func CreateDashboard(data map[string]interface{}) (*Dashboard, *model.ApiError) } dash.Id = int(lastInsertId) + traceAndLogsPanelUsage := countTraceAndLogsPanel(data) + if traceAndLogsPanelUsage > 0 { + updateFeatureUsage(fm, traceAndLogsPanelUsage) + } + return dash, nil } @@ -176,7 +189,13 @@ func GetDashboards() ([]Dashboard, *model.ApiError) { return dashboards, nil } -func DeleteDashboard(uuid string) *model.ApiError { +func DeleteDashboard(uuid string, fm interfaces.FeatureLookup) *model.ApiError { + + dashboard, dErr := GetDashboard(uuid) + if dErr != nil { + zap.S().Errorf("Error in getting dashboard: ", uuid, dErr) + return dErr + } query := fmt.Sprintf("DELETE FROM dashboards WHERE uuid='%s';", uuid) @@ -194,6 +213,11 @@ func DeleteDashboard(uuid string) *model.ApiError { return &model.ApiError{Typ: model.ErrorNotFound, Err: fmt.Errorf("no dashboard found with uuid: %s", uuid)} } + traceAndLogsPanelUsage := countTraceAndLogsPanel(dashboard.Data) + if traceAndLogsPanelUsage > 0 { + updateFeatureUsage(fm, -traceAndLogsPanelUsage) + } + return nil } @@ -210,7 +234,7 @@ func GetDashboard(uuid string) (*Dashboard, *model.ApiError) { return &dashboard, nil } -func UpdateDashboard(uuid string, data map[string]interface{}) (*Dashboard, *model.ApiError) { +func UpdateDashboard(uuid string, data map[string]interface{}, fm interfaces.FeatureLookup) (*Dashboard, *model.ApiError) { map_data, err := json.Marshal(data) if err != nil { @@ -223,6 +247,16 @@ func UpdateDashboard(uuid string, data map[string]interface{}) (*Dashboard, *mod return nil, apiErr } + // check if the count of trace and logs QB panel has changed, if yes, then check feature flag count + existingCount := countTraceAndLogsPanel(dashboard.Data) + newCount := countTraceAndLogsPanel(data) + if newCount > existingCount { + err := checkFeatureUsage(fm, newCount-existingCount) + if err != nil { + return nil, err + } + } + dashboard.UpdatedAt = time.Now() dashboard.Data = data @@ -233,10 +267,58 @@ func UpdateDashboard(uuid string, data map[string]interface{}) (*Dashboard, *mod zap.S().Errorf("Error in inserting dashboard data: ", data, err) return nil, &model.ApiError{Typ: model.ErrorExec, Err: err} } - + if existingCount != newCount { + // if the count of trace and logs panel has changed, we need to update feature flag count as well + updateFeatureUsage(fm, newCount-existingCount) + } return dashboard, nil } +func updateFeatureUsage(fm interfaces.FeatureLookup, usage int64) *model.ApiError { + feature, err := fm.GetFeatureFlag(model.QueryBuilderPanels) + if err != nil { + switch err.(type) { + case model.ErrFeatureUnavailable: + zap.S().Errorf("feature unavailable", zap.String("featureKey", model.QueryBuilderPanels), zap.Error(err)) + return model.BadRequest(err) + default: + zap.S().Errorf("feature check failed", zap.String("featureKey", model.QueryBuilderPanels), zap.Error(err)) + return model.BadRequest(err) + } + } + feature.Usage += usage + if feature.Usage >= feature.UsageLimit { + feature.Active = false + } + if feature.Usage < feature.UsageLimit { + feature.Active = true + } + err = fm.UpdateFeatureFlag(feature) + if err != nil { + return model.BadRequest(err) + } + + return nil +} + +func checkFeatureUsage(fm interfaces.FeatureLookup, usage int64) *model.ApiError { + feature, err := fm.GetFeatureFlag(model.QueryBuilderPanels) + if err != nil { + switch err.(type) { + case model.ErrFeatureUnavailable: + zap.S().Errorf("feature unavailable", zap.String("featureKey", model.QueryBuilderPanels), zap.Error(err)) + return model.BadRequest(err) + default: + zap.S().Errorf("feature check failed", zap.String("featureKey", model.QueryBuilderPanels), zap.Error(err)) + return model.BadRequest(err) + } + } + if feature.UsageLimit-(feature.Usage+usage) < 0 { + return model.BadRequest(fmt.Errorf("feature usage exceeded")) + } + return nil +} + // UpdateSlug updates the slug func (d *Dashboard) UpdateSlug() { var title string @@ -505,3 +587,38 @@ func TransformGrafanaJSONToSignoz(grafanaJSON model.GrafanaJSON) model.Dashboard } return toReturn } + +func countTraceAndLogsPanel(data map[string]interface{}) int64 { + count := int64(0) + if data != nil && data["widgets"] != nil { + widgets, ok := data["widgets"].(interface{}) + if ok { + data, ok := widgets.([]interface{}) + if ok { + for _, widget := range data { + sData, ok := widget.(map[string]interface{}) + if ok && sData["query"] != nil { + query, ok := sData["query"].(interface{}).(map[string]interface{}) + if ok && query["queryType"] == "builder" && query["builder"] != nil { + builderData, ok := query["builder"].(interface{}).(map[string]interface{}) + if ok && builderData["queryData"] != nil { + builderQueryData, ok := builderData["queryData"].([]interface{}) + if ok { + for _, queryData := range builderQueryData { + data, ok := queryData.(map[string]interface{}) + if ok { + if data["dataSource"] == "traces" || data["dataSource"] == "logs" { + count++ + } + } + } + } + } + } + } + } + } + } + } + return count +} diff --git a/pkg/query-service/app/dashboards/provision.go b/pkg/query-service/app/dashboards/provision.go index bc0bc39059..d8869f048b 100644 --- a/pkg/query-service/app/dashboards/provision.go +++ b/pkg/query-service/app/dashboards/provision.go @@ -6,10 +6,11 @@ import ( "os" "go.signoz.io/signoz/pkg/query-service/constants" + "go.signoz.io/signoz/pkg/query-service/interfaces" "go.uber.org/zap" ) -func readCurrentDir(dir string) error { +func readCurrentDir(dir string, fm interfaces.FeatureLookup) error { file, err := os.Open(dir) if err != nil { zap.S().Errorf("failed opening directory: %s", err) @@ -43,7 +44,7 @@ func readCurrentDir(dir string) error { continue } - _, apiErr = CreateDashboard(data) + _, apiErr = CreateDashboard(data, fm) if apiErr != nil { zap.S().Errorf("Creating Dashboards: Error in file: %s\t%s", filename, apiErr.Err) continue @@ -53,7 +54,7 @@ func readCurrentDir(dir string) error { return nil } -func LoadDashboardFiles() error { +func LoadDashboardFiles(fm interfaces.FeatureLookup) error { dashboardsPath := constants.GetOrDefaultEnv("DASHBOARDS_PATH", "./config/dashboards") - return readCurrentDir(dashboardsPath) + return readCurrentDir(dashboardsPath, fm) } diff --git a/pkg/query-service/app/http_handler.go b/pkg/query-service/app/http_handler.go index fb3af10b3e..b1ecc25aad 100644 --- a/pkg/query-service/app/http_handler.go +++ b/pkg/query-service/app/http_handler.go @@ -116,7 +116,7 @@ func NewAPIHandler(opts APIHandlerOpts) (*APIHandler, error) { aH.ready = aH.testReady - dashboards.LoadDashboardFiles() + dashboards.LoadDashboardFiles(aH.featureFlags) // if errReadingDashboards != nil { // return nil, errReadingDashboards // } @@ -723,7 +723,7 @@ func (aH *APIHandler) getDashboards(w http.ResponseWriter, r *http.Request) { func (aH *APIHandler) deleteDashboard(w http.ResponseWriter, r *http.Request) { uuid := mux.Vars(r)["uuid"] - err := dashboards.DeleteDashboard(uuid) + err := dashboards.DeleteDashboard(uuid, aH.featureFlags) if err != nil { RespondError(w, err, nil) @@ -829,7 +829,7 @@ func (aH *APIHandler) updateDashboard(w http.ResponseWriter, r *http.Request) { return } - dashboard, apiError := dashboards.UpdateDashboard(uuid, postData) + dashboard, apiError := dashboards.UpdateDashboard(uuid, postData, aH.featureFlags) if apiError != nil { RespondError(w, apiError, nil) return @@ -863,7 +863,7 @@ func (aH *APIHandler) saveAndReturn(w http.ResponseWriter, signozDashboard model toSave["widgets"] = signozDashboard.Widgets toSave["variables"] = signozDashboard.Variables - dashboard, apiError := dashboards.CreateDashboard(toSave) + dashboard, apiError := dashboards.CreateDashboard(toSave, aH.featureFlags) if apiError != nil { RespondError(w, apiError, nil) return @@ -904,7 +904,7 @@ func (aH *APIHandler) createDashboards(w http.ResponseWriter, r *http.Request) { return } - dash, apiErr := dashboards.CreateDashboard(postData) + dash, apiErr := dashboards.CreateDashboard(postData, aH.featureFlags) if apiErr != nil { RespondError(w, apiErr, nil) @@ -1613,7 +1613,11 @@ func (aH *APIHandler) getVersion(w http.ResponseWriter, r *http.Request) { } func (aH *APIHandler) getFeatureFlags(w http.ResponseWriter, r *http.Request) { - featureSet := aH.FF().GetFeatureFlags() + featureSet, err := aH.FF().GetFeatureFlags() + if err != nil { + aH.HandleError(w, err, http.StatusInternalServerError) + return + } aH.Respond(w, featureSet) } diff --git a/pkg/query-service/app/server.go b/pkg/query-service/app/server.go index 73b1baed1e..78101ea406 100644 --- a/pkg/query-service/app/server.go +++ b/pkg/query-service/app/server.go @@ -107,7 +107,7 @@ func NewServer(serverOptions *ServerOptions) (*Server, error) { } <-readerReady - rm, err := makeRulesManager(serverOptions.PromConfigPath, constants.GetAlertManagerApiPrefix(), serverOptions.RuleRepoURL, localDB, reader, serverOptions.DisableRules) + rm, err := makeRulesManager(serverOptions.PromConfigPath, constants.GetAlertManagerApiPrefix(), serverOptions.RuleRepoURL, localDB, reader, serverOptions.DisableRules, fm) if err != nil { return nil, err } @@ -491,7 +491,8 @@ func makeRulesManager( ruleRepoURL string, db *sqlx.DB, ch interfaces.Reader, - disableRules bool) (*rules.Manager, error) { + disableRules bool, + fm interfaces.FeatureLookup) (*rules.Manager, error) { // create engine pqle, err := pqle.FromReader(ch) @@ -518,6 +519,7 @@ func makeRulesManager( Context: context.Background(), Logger: nil, DisableRules: disableRules, + FeatureFlags: fm, } // create Manager diff --git a/pkg/query-service/constants/constants.go b/pkg/query-service/constants/constants.go index 140d76c67a..db1555f1a3 100644 --- a/pkg/query-service/constants/constants.go +++ b/pkg/query-service/constants/constants.go @@ -71,8 +71,19 @@ func IsTimestampSortFeatureEnabled() bool { } var DEFAULT_FEATURE_SET = model.FeatureSet{ - DurationSort: IsDurationSortFeatureEnabled(), - TimestampSort: IsTimestampSortFeatureEnabled(), + model.Feature{ + Name: DurationSort, + Active: IsDurationSortFeatureEnabled(), + Usage: 0, + UsageLimit: -1, + Route: "", + }, model.Feature{ + Name: TimestampSort, + Active: IsTimestampSortFeatureEnabled(), + Usage: 0, + UsageLimit: -1, + Route: "", + }, } func GetContextTimeout() time.Duration { diff --git a/pkg/query-service/featureManager/manager.go b/pkg/query-service/featureManager/manager.go index 1c8c953982..15175b1882 100644 --- a/pkg/query-service/featureManager/manager.go +++ b/pkg/query-service/featureManager/manager.go @@ -3,32 +3,64 @@ package featureManager import ( "go.signoz.io/signoz/pkg/query-service/constants" "go.signoz.io/signoz/pkg/query-service/model" + "go.uber.org/zap" ) type FeatureManager struct { - activeFeatures model.FeatureSet } func StartManager() *FeatureManager { - fM := &FeatureManager{ - activeFeatures: constants.DEFAULT_FEATURE_SET, - } + fM := &FeatureManager{} return fM } // CheckFeature will be internally used by backend routines // for feature gating func (fm *FeatureManager) CheckFeature(featureKey string) error { - if value, ok := fm.activeFeatures[featureKey]; ok { - if value { - return nil - } - return model.ErrFeatureUnavailable{Key: featureKey} + + feature, err := fm.GetFeatureFlag(featureKey) + if err != nil { + return err } + + if feature.Active { + return nil + } + return model.ErrFeatureUnavailable{Key: featureKey} } -// GetFeatureFlags returns current active features -func (fm *FeatureManager) GetFeatureFlags() model.FeatureSet { - return fm.activeFeatures +// GetFeatureFlags returns current features +func (fm *FeatureManager) GetFeatureFlags() (model.FeatureSet, error) { + features := append(constants.DEFAULT_FEATURE_SET, model.Feature{ + Name: model.OSS, + Active: true, + Usage: 0, + UsageLimit: -1, + Route: "", + }) + return features, nil } + +func (fm *FeatureManager) InitFeatures(req model.FeatureSet) error { + zap.S().Error("InitFeatures not implemented in OSS") + return nil +} + +func (fm *FeatureManager) UpdateFeatureFlag(req model.Feature) error { + zap.S().Error("UpdateFeatureFlag not implemented in OSS") + return nil +} + +func (fm *FeatureManager) GetFeatureFlag(key string) (model.Feature, error) { + features, err := fm.GetFeatureFlags() + if err != nil { + return model.Feature{}, err + } + for _, feature := range features { + if feature.Name == key { + return feature, nil + } + } + return model.Feature{}, model.ErrFeatureUnavailable{Key: key} +} \ No newline at end of file diff --git a/pkg/query-service/interfaces/featureLookup.go b/pkg/query-service/interfaces/featureLookup.go index 0f7ab49a03..32c7041cf0 100644 --- a/pkg/query-service/interfaces/featureLookup.go +++ b/pkg/query-service/interfaces/featureLookup.go @@ -6,5 +6,8 @@ import ( type FeatureLookup interface { CheckFeature(f string) error - GetFeatureFlags() model.FeatureSet + GetFeatureFlags() (model.FeatureSet, error) + GetFeatureFlag(f string) (model.Feature, error) + UpdateFeatureFlag(features model.Feature) error + InitFeatures(features model.FeatureSet) error } diff --git a/pkg/query-service/model/featureSet.go b/pkg/query-service/model/featureSet.go index bba153e861..7f4cfb645e 100644 --- a/pkg/query-service/model/featureSet.go +++ b/pkg/query-service/model/featureSet.go @@ -1,11 +1,16 @@ package model -type FeatureSet map[string]bool +type FeatureSet []Feature +type Feature struct { + Name string `db:"name" json:"name"` + Active bool `db:"active" json:"active"` + Usage int64 `db:"usage" json:"usage"` + UsageLimit int64 `db:"usage_limit" json:"usage_limit"` + Route string `db:"route" json:"route"` +} -const Basic = "BASIC_PLAN" const SmartTraceDetail = "SMART_TRACE_DETAIL" const CustomMetricsFunction = "CUSTOM_METRICS_FUNCTION" - -var BasicPlan = FeatureSet{ - Basic: true, -} +const OSS = "OSS" +const QueryBuilderPanels = "QUERY_BUILDER_PANELS" +const QueryBuilderAlerts = "QUERY_BUILDER_ALERTS" \ No newline at end of file diff --git a/pkg/query-service/rules/manager.go b/pkg/query-service/rules/manager.go index 2437355daa..2c5bd926a1 100644 --- a/pkg/query-service/rules/manager.go +++ b/pkg/query-service/rules/manager.go @@ -21,7 +21,9 @@ import ( // opentracing "github.com/opentracing/opentracing-go" am "go.signoz.io/signoz/pkg/query-service/integrations/alertManager" + "go.signoz.io/signoz/pkg/query-service/interfaces" "go.signoz.io/signoz/pkg/query-service/model" + v3 "go.signoz.io/signoz/pkg/query-service/model/v3" "go.signoz.io/signoz/pkg/query-service/utils/labels" ) @@ -59,6 +61,7 @@ type ManagerOptions struct { Logger log.Logger ResendDelay time.Duration DisableRules bool + FeatureFlags interfaces.FeatureLookup } // The Manager manages recording and alerting rules. @@ -77,6 +80,8 @@ type Manager struct { // pause all rule tasks pause bool logger log.Logger + + featureFlags interfaces.FeatureLookup } func defaultOptions(o *ManagerOptions) *ManagerOptions { @@ -109,13 +114,14 @@ func NewManager(o *ManagerOptions) (*Manager, error) { db := newRuleDB(o.DBConn) m := &Manager{ - tasks: map[string]Task{}, - rules: map[string]Rule{}, - notifier: notifier, - ruleDB: db, - opts: o, - block: make(chan struct{}), - logger: o.Logger, + tasks: map[string]Task{}, + rules: map[string]Rule{}, + notifier: notifier, + ruleDB: db, + opts: o, + block: make(chan struct{}), + logger: o.Logger, + featureFlags: o.FeatureFlags, } return m, nil } @@ -221,6 +227,20 @@ func (m *Manager) EditRule(ruleStr string, id string) error { parsedRule, errs := ParsePostableRule([]byte(ruleStr)) + currentRule, err := m.GetRule(id) + if err != nil { + zap.S().Errorf("msg: ", "failed to get the rule from rule db", "\t ruleid: ", id) + return err + } + + if !checkIfTraceOrLogQB(¤tRule.PostableRule) { + // check if the new rule uses any feature that is not enabled + err = m.checkFeatureUsage(parsedRule) + if err != nil { + return err + } + } + if len(errs) > 0 { zap.S().Errorf("failed to parse rules:", errs) // just one rule is being parsed so expect just one error @@ -233,7 +253,24 @@ func (m *Manager) EditRule(ruleStr string, id string) error { } if !m.opts.DisableRules { - return m.syncRuleStateWithTask(taskName, parsedRule) + err = m.syncRuleStateWithTask(taskName, parsedRule) + if err != nil { + return err + } + } + + // update feature usage if the current rule is not a trace or log query builder + if !checkIfTraceOrLogQB(¤tRule.PostableRule) { + err = m.updateFeatureUsage(parsedRule, 1) + if err != nil { + zap.S().Errorf("error updating feature usage: %v", err) + } + // update feature usage if the new rule is not a trace or log query builder and the current rule is + } else if !checkIfTraceOrLogQB(parsedRule) { + err = m.updateFeatureUsage(¤tRule.PostableRule, -1) + if err != nil { + zap.S().Errorf("error updating feature usage: %v", err) + } } return nil @@ -285,6 +322,13 @@ func (m *Manager) DeleteRule(id string) error { return fmt.Errorf("delete rule received an rule id in invalid format, must be a number") } + // update feature usage + rule, err := m.GetRule(id) + if err != nil { + zap.S().Errorf("msg: ", "failed to get the rule from rule db", "\t ruleid: ", id) + return err + } + taskName := prepareTaskName(int64(idInt)) if !m.opts.DisableRules { m.deleteTask(taskName) @@ -295,6 +339,11 @@ func (m *Manager) DeleteRule(id string) error { return err } + err = m.updateFeatureUsage(&rule.PostableRule, -1) + if err != nil { + zap.S().Errorf("error updating feature usage: %v", err) + } + return nil } @@ -319,6 +368,12 @@ func (m *Manager) deleteTask(taskName string) { func (m *Manager) CreateRule(ruleStr string) error { parsedRule, errs := ParsePostableRule([]byte(ruleStr)) + // check if the rule uses any feature that is not enabled + err := m.checkFeatureUsage(parsedRule) + if err != nil { + return err + } + if len(errs) > 0 { zap.S().Errorf("failed to parse rules:", errs) // just one rule is being parsed so expect just one error @@ -335,7 +390,70 @@ func (m *Manager) CreateRule(ruleStr string) error { return err } } - return tx.Commit() + err = tx.Commit() + if err != nil { + return err + } + + // update feature usage + err = m.updateFeatureUsage(parsedRule, 1) + if err != nil { + zap.S().Errorf("error updating feature usage: %v", err) + } + return nil +} + +func (m *Manager) updateFeatureUsage(parsedRule *PostableRule, usage int64) error { + isTraceOrLogQB := checkIfTraceOrLogQB(parsedRule) + if isTraceOrLogQB { + feature, err := m.featureFlags.GetFeatureFlag(model.QueryBuilderAlerts) + if err != nil { + return err + } + feature.Usage += usage + if feature.Usage == feature.UsageLimit { + feature.Active = false + } + if feature.Usage < feature.UsageLimit { + feature.Active = true + } + err = m.featureFlags.UpdateFeatureFlag(feature) + if err != nil { + return err + } + } + return nil +} + +func (m *Manager) checkFeatureUsage(parsedRule *PostableRule) error { + isTraceOrLogQB := checkIfTraceOrLogQB(parsedRule) + if isTraceOrLogQB { + err := m.featureFlags.CheckFeature(model.QueryBuilderAlerts) + if err != nil { + switch err.(type) { + case model.ErrFeatureUnavailable: + zap.S().Errorf("feature unavailable", zap.String("featureKey", model.QueryBuilderAlerts), zap.Error(err)) + return model.BadRequest(err) + default: + zap.S().Errorf("feature check failed", zap.String("featureKey", model.QueryBuilderAlerts), zap.Error(err)) + return model.BadRequest(err) + } + } + } + return nil +} + +func checkIfTraceOrLogQB(parsedRule *PostableRule) bool { + if parsedRule != nil { + if parsedRule.RuleCondition.QueryType() == v3.QueryTypeBuilder { + for _, query := range parsedRule.RuleCondition.CompositeQuery.BuilderQueries { + if query.DataSource == v3.DataSourceTraces || query.DataSource == v3.DataSourceLogs { + return true + } + } + } + } + return false } func (m *Manager) addTask(rule *PostableRule, taskName string) error { @@ -623,10 +741,10 @@ func (m *Manager) syncRuleStateWithTask(taskName string, rule *PostableRule) err // PatchRule supports attribute level changes to the rule definition unlike // EditRule, which updates entire rule definition in the DB. // the process: -// - get the latest rule from db -// - over write the patch attributes received in input (ruleStr) -// - re-deploy or undeploy task as necessary -// - update the patched rule in the DB +// - get the latest rule from db +// - over write the patch attributes received in input (ruleStr) +// - re-deploy or undeploy task as necessary +// - update the patched rule in the DB func (m *Manager) PatchRule(ruleStr string, ruleId string) (*GettableRule, error) { if ruleId == "" {