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

@@ -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)
 }

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

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
+}

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 = append(lm.activeFeatures, baseconstants.DEFAULT_FEATURE_SET...)
-		lm.activeFeatures[k] = v
-	}
 }
 
 // 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 {
+	feature, err := lm.repo.GetFeature(featureKey)
-		if value {
+	if err != nil {
-			return nil
+		return err
-		}
+	}
-		return basemodel.ErrFeatureUnavailable{Key: featureKey}
+	if feature.Active {
+		return nil
 	}
 	return basemodel.ErrFeatureUnavailable{Key: featureKey}
 }
 
 // GetFeatureFlags returns current active features
-func (lm *Manager) GetFeatureFlags() basemodel.FeatureSet {
+func (lm *Manager) GetFeatureFlags() (basemodel.FeatureSet, error) {
-	return lm.activeFeatures
+	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

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
 }

ee/query-service/model/plans.go

@@ -11,21 +11,143 @@ const Enterprise = "ENTERPRISE_PLAN"
 const DisableUpsell = "DISABLE_UPSELL"
 
 var BasicPlan = basemodel.FeatureSet{
-	Basic:         true,
+	basemodel.Feature{
-	SSO:           false,
+		Name:       SSO,
-	DisableUpsell: false,
+		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,
+	basemodel.Feature{
-	SSO:                             true,
+		Name:       SSO,
-	basemodel.SmartTraceDetail:      true,
+		Active:     true,
-	basemodel.CustomMetricsFunction: 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,
+	basemodel.Feature{
-	SSO:                             true,
+		Name:       SSO,
-	basemodel.SmartTraceDetail:      true,
+		Active:     true,
-	basemodel.CustomMetricsFunction: 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:      "",
+	},
 }

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
+}

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)
 }

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)
 }
 

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

pkg/query-service/constants/constants.go

@@ -71,8 +71,19 @@ func IsTimestampSortFeatureEnabled() bool {
 }
 
 var DEFAULT_FEATURE_SET = model.FeatureSet{
-	DurationSort:  IsDurationSortFeatureEnabled(),
+	model.Feature{
-	TimestampSort: IsTimestampSortFeatureEnabled(),
+		Name:       DurationSort,
+		Active:     IsDurationSortFeatureEnabled(),
+		Usage:      0,
+		UsageLimit: -1,
+		Route:      "",
+	}, model.Feature{
+		Name:       TimestampSort,
+		Active:     IsTimestampSortFeatureEnabled(),
+		Usage:      0,
+		UsageLimit: -1,
+		Route:      "",
+	},
 }
 
 func GetContextTimeout() time.Duration {

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{
+	fM := &FeatureManager{}
-		activeFeatures: constants.DEFAULT_FEATURE_SET,
-	}
 	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 {
+	feature, err := fm.GetFeatureFlag(featureKey)
-			return nil
+	if err != nil {
-		}
+		return err
-		return model.ErrFeatureUnavailable{Key: featureKey}
 	}
+
+	if feature.Active {
+		return nil
+	}
+
 	return model.ErrFeatureUnavailable{Key: featureKey}
 }
 
-// GetFeatureFlags returns current active features
+// GetFeatureFlags returns current features
-func (fm *FeatureManager) GetFeatureFlags() model.FeatureSet {
+func (fm *FeatureManager) GetFeatureFlags() (model.FeatureSet, error) {
-	return fm.activeFeatures
+	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}
+}

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
 }

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"
-
+const OSS = "OSS"
-var BasicPlan = FeatureSet{
+const QueryBuilderPanels = "QUERY_BUILDER_PANELS"
-	Basic: true,
+const QueryBuilderAlerts = "QUERY_BUILDER_ALERTS"
-}

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{},
+		tasks:        map[string]Task{},
-		rules:    map[string]Rule{},
+		rules:        map[string]Rule{},
-		notifier: notifier,
+		notifier:     notifier,
-		ruleDB:   db,
+		ruleDB:       db,
-		opts:     o,
+		opts:         o,
-		block:    make(chan struct{}),
+		block:        make(chan struct{}),
-		logger:   o.Logger,
+		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(&currentRule.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(&currentRule.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(&currentRule.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
+//   - get the latest rule from db
-//  - over write the patch attributes received in input (ruleStr)
+//   - over write the patch attributes received in input (ruleStr)
-//  - re-deploy or undeploy task as necessary
+//   - re-deploy or undeploy task as necessary
-//  - update the patched rule in the DB
+//   - update the patched rule in the DB
 func (m *Manager) PatchRule(ruleStr string, ruleId string) (*GettableRule, error) {
 
 	if ruleId == "" {