feat: introduce feature_usage table to manage features

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/dao/factory.go

@@ -6,13 +6,24 @@ import (
 	"go.signoz.io/signoz/ee/query-service/dao/sqlite"
 )
 
+var db ModelDao
+
 func InitDao(engine, path string) (ModelDao, error) {
 
 	switch engine {
 	case "sqlite":
-		return sqlite.InitDB(path)
+		db, err:= sqlite.InitDB(path)
+		return db, err
 	default:
 		return nil, fmt.Errorf("qsdb type: %s is not supported in query service", engine)
 	}
 
 }
+
+func DB() ModelDao {
+	if db == nil {
+		// Should never reach here
+		panic("GetDB called before initialization")
+	}
+	return db
+}

ee/query-service/license/db.go

@@ -9,6 +9,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 +126,77 @@ 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 := r.CreateFeature(&feature)
+			if err != nil {
+				return err
+			}
+			continue
+		}
+		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.repo.InitFeatures(lm.activeFeatures)
+	if err != nil {
+		zap.S().Error("Couldn't activate features: ", err)
+	}
 	if !lm.validatorRunning {
 		// we want to make sure only one validator runs,
 		// we already have lock() so good to go
@@ -106,8 +111,8 @@ func (lm *Manager) SetActive(l *model.License) {
 }
 
 func setDefaultFeatures(lm *Manager) {
-	for k, v := range baseconstants.DEFAULT_FEATURE_SET {
-		lm.activeFeatures[k] = v
+	for _, feature := range baseconstants.DEFAULT_FEATURE_SET {
+		lm.activeFeatures = append(lm.activeFeatures, feature)
 	}
 }
 
@@ -123,8 +128,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.repo.InitFeatures(lm.activeFeatures)
+		if err != nil {
+			zap.S().Error("Couldn't initialize features: ", err)
+			return err
+		}
 	}
 
 	return nil
@@ -291,18 +301,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

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
 }