diff --git a/ee/http/middleware/pat.go b/ee/http/middleware/pat.go index a890f7c79c..cb7e8e0917 100644 --- a/ee/http/middleware/pat.go +++ b/ee/http/middleware/pat.go @@ -4,20 +4,20 @@ import ( "net/http" "time" + "github.com/SigNoz/signoz/pkg/sqlstore" "github.com/SigNoz/signoz/pkg/types" "github.com/SigNoz/signoz/pkg/types/authtypes" - "github.com/uptrace/bun" "go.uber.org/zap" ) type Pat struct { - db *bun.DB + store sqlstore.SQLStore uuid *authtypes.UUID headers []string } -func NewPat(db *bun.DB, headers []string) *Pat { - return &Pat{db: db, uuid: authtypes.NewUUID(), headers: headers} +func NewPat(store sqlstore.SQLStore, headers []string) *Pat { + return &Pat{store: store, uuid: authtypes.NewUUID(), headers: headers} } func (p *Pat) Wrap(next http.Handler) http.Handler { @@ -41,7 +41,7 @@ func (p *Pat) Wrap(next http.Handler) http.Handler { return } - err = p.db.NewSelect().Model(&pat).Where("token = ?", patToken).Scan(r.Context()) + err = p.store.BunDB().NewSelect().Model(&pat).Where("token = ?", patToken).Scan(r.Context()) if err != nil { next.ServeHTTP(w, r) return @@ -54,7 +54,7 @@ func (p *Pat) Wrap(next http.Handler) http.Handler { // get user from db user := types.User{} - err = p.db.NewSelect().Model(&user).Where("id = ?", pat.UserID).Scan(r.Context()) + err = p.store.BunDB().NewSelect().Model(&user).Where("id = ?", pat.UserID).Scan(r.Context()) if err != nil { next.ServeHTTP(w, r) return @@ -74,7 +74,7 @@ func (p *Pat) Wrap(next http.Handler) http.Handler { next.ServeHTTP(w, r) pat.LastUsed = time.Now().Unix() - _, err = p.db.NewUpdate().Model(&pat).Column("last_used").Where("token = ?", patToken).Where("revoked = false").Exec(r.Context()) + _, err = p.store.BunDB().NewUpdate().Model(&pat).Column("last_used").Where("token = ?", patToken).Where("revoked = false").Exec(r.Context()) if err != nil { zap.L().Error("Failed to update PAT last used in db, err: %v", zap.Error(err)) } diff --git a/ee/query-service/app/server.go b/ee/query-service/app/server.go index bc44bfb29b..cc0fb251ba 100644 --- a/ee/query-service/app/server.go +++ b/ee/query-service/app/server.go @@ -112,7 +112,7 @@ func NewServer(serverOptions *ServerOptions) (*Server, error) { return nil, err } - if err := baseexplorer.InitWithDSN(serverOptions.SigNoz.SQLStore.BunDB()); err != nil { + if err := baseexplorer.InitWithDSN(serverOptions.SigNoz.SQLStore); err != nil { return nil, err } @@ -120,7 +120,7 @@ func NewServer(serverOptions *ServerOptions) (*Server, error) { return nil, err } - if err := dashboards.InitDB(serverOptions.SigNoz.SQLStore.BunDB()); err != nil { + if err := dashboards.InitDB(serverOptions.SigNoz.SQLStore); err != nil { return nil, err } @@ -130,7 +130,7 @@ func NewServer(serverOptions *ServerOptions) (*Server, error) { } // initiate license manager - lm, err := licensepkg.StartManager(serverOptions.SigNoz.SQLStore.SQLxDB(), serverOptions.SigNoz.SQLStore.BunDB()) + lm, err := licensepkg.StartManager(serverOptions.SigNoz.SQLStore.SQLxDB(), serverOptions.SigNoz.SQLStore) if err != nil { return nil, err } @@ -317,7 +317,7 @@ func (s *Server) createPrivateServer(apiHandler *api.APIHandler) (*http.Server, r := baseapp.NewRouter() r.Use(middleware.NewAuth(zap.L(), s.serverOptions.Jwt, []string{"Authorization", "Sec-WebSocket-Protocol"}).Wrap) - r.Use(eemiddleware.NewPat(s.serverOptions.SigNoz.SQLStore.BunDB(), []string{"SIGNOZ-API-KEY"}).Wrap) + r.Use(eemiddleware.NewPat(s.serverOptions.SigNoz.SQLStore, []string{"SIGNOZ-API-KEY"}).Wrap) r.Use(middleware.NewTimeout(zap.L(), s.serverOptions.Config.APIServer.Timeout.ExcludedRoutes, s.serverOptions.Config.APIServer.Timeout.Default, @@ -365,7 +365,7 @@ func (s *Server) createPublicServer(apiHandler *api.APIHandler, web web.Web) (*h am := baseapp.NewAuthMiddleware(getUserFromRequest) r.Use(middleware.NewAuth(zap.L(), s.serverOptions.Jwt, []string{"Authorization", "Sec-WebSocket-Protocol"}).Wrap) - r.Use(eemiddleware.NewPat(s.serverOptions.SigNoz.SQLStore.BunDB(), []string{"SIGNOZ-API-KEY"}).Wrap) + r.Use(eemiddleware.NewPat(s.serverOptions.SigNoz.SQLStore, []string{"SIGNOZ-API-KEY"}).Wrap) r.Use(middleware.NewTimeout(zap.L(), s.serverOptions.Config.APIServer.Timeout.ExcludedRoutes, s.serverOptions.Config.APIServer.Timeout.Default, diff --git a/ee/query-service/license/db.go b/ee/query-service/license/db.go index 0f30ec78db..c241ad8766 100644 --- a/ee/query-service/license/db.go +++ b/ee/query-service/license/db.go @@ -9,10 +9,10 @@ import ( "github.com/jmoiron/sqlx" "github.com/mattn/go-sqlite3" - "github.com/uptrace/bun" "github.com/SigNoz/signoz/ee/query-service/model" basemodel "github.com/SigNoz/signoz/pkg/query-service/model" + "github.com/SigNoz/signoz/pkg/sqlstore" "github.com/SigNoz/signoz/pkg/types" "go.uber.org/zap" ) @@ -20,14 +20,14 @@ import ( // Repo is license repo. stores license keys in a secured DB type Repo struct { db *sqlx.DB - bundb *bun.DB + store sqlstore.SQLStore } // NewLicenseRepo initiates a new license repo -func NewLicenseRepo(db *sqlx.DB, bundb *bun.DB) Repo { +func NewLicenseRepo(db *sqlx.DB, store sqlstore.SQLStore) Repo { return Repo{ db: db, - bundb: bundb, + store: store, } } @@ -171,7 +171,7 @@ func (r *Repo) UpdateLicenseV3(ctx context.Context, l *model.LicenseV3) error { func (r *Repo) CreateFeature(req *types.FeatureStatus) *basemodel.ApiError { - _, err := r.bundb.NewInsert(). + _, err := r.store.BunDB().NewInsert(). Model(req). Exec(context.Background()) if err != nil { @@ -183,7 +183,7 @@ func (r *Repo) CreateFeature(req *types.FeatureStatus) *basemodel.ApiError { func (r *Repo) GetFeature(featureName string) (types.FeatureStatus, error) { var feature types.FeatureStatus - err := r.bundb.NewSelect(). + err := r.store.BunDB().NewSelect(). Model(&feature). Where("name = ?", featureName). Scan(context.Background()) @@ -212,7 +212,7 @@ func (r *Repo) GetAllFeatures() ([]basemodel.Feature, error) { func (r *Repo) UpdateFeature(req types.FeatureStatus) error { - _, err := r.bundb.NewUpdate(). + _, err := r.store.BunDB().NewUpdate(). Model(&req). Where("name = ?", req.Name). Exec(context.Background()) diff --git a/ee/query-service/license/manager.go b/ee/query-service/license/manager.go index 900a5dd75e..c67dcf7a25 100644 --- a/ee/query-service/license/manager.go +++ b/ee/query-service/license/manager.go @@ -7,11 +7,11 @@ import ( "github.com/jmoiron/sqlx" "github.com/pkg/errors" - "github.com/uptrace/bun" "sync" baseconstants "github.com/SigNoz/signoz/pkg/query-service/constants" + "github.com/SigNoz/signoz/pkg/sqlstore" "github.com/SigNoz/signoz/pkg/types" "github.com/SigNoz/signoz/pkg/types/authtypes" @@ -45,12 +45,12 @@ type Manager struct { activeFeatures basemodel.FeatureSet } -func StartManager(db *sqlx.DB, bundb *bun.DB, features ...basemodel.Feature) (*Manager, error) { +func StartManager(db *sqlx.DB, store sqlstore.SQLStore, features ...basemodel.Feature) (*Manager, error) { if LM != nil { return LM, nil } - repo := NewLicenseRepo(db, bundb) + repo := NewLicenseRepo(db, store) m := &Manager{ repo: &repo, } diff --git a/pkg/query-service/app/dashboards/model.go b/pkg/query-service/app/dashboards/model.go index c90eb0690e..8f05b44f18 100644 --- a/pkg/query-service/app/dashboards/model.go +++ b/pkg/query-service/app/dashboards/model.go @@ -11,16 +11,16 @@ import ( "github.com/SigNoz/signoz/pkg/query-service/interfaces" "github.com/SigNoz/signoz/pkg/query-service/model" + "github.com/SigNoz/signoz/pkg/sqlstore" "github.com/SigNoz/signoz/pkg/types" "github.com/google/uuid" - "github.com/uptrace/bun" "github.com/SigNoz/signoz/pkg/query-service/telemetry" "go.uber.org/zap" ) // This time the global variable is unexported. -var db *bun.DB +var store sqlstore.SQLStore // User for mapping job,instance from grafana var ( @@ -33,8 +33,8 @@ var ( ) // InitDB sets up setting up the connection pool global variable. -func InitDB(inputDB *bun.DB) error { - db = inputDB +func InitDB(sqlStore sqlstore.SQLStore) error { + store = sqlStore telemetry.GetInstance().SetDashboardsInfoCallback(GetDashboardsInfo) return nil @@ -57,7 +57,7 @@ func CreateDashboard(ctx context.Context, orgID string, email string, data map[s dash.UUID = data["uuid"].(string) } - err := db.NewInsert().Model(dash).Returning("id").Scan(ctx, &dash.ID) + err := store.BunDB().NewInsert().Model(dash).Returning("id").Scan(ctx, &dash.ID) if err != nil { zap.L().Error("Error in inserting dashboard data: ", zap.Any("dashboard", dash), zap.Error(err)) return nil, &model.ApiError{Typ: model.ErrorExec, Err: err} @@ -69,7 +69,7 @@ func CreateDashboard(ctx context.Context, orgID string, email string, data map[s func GetDashboards(ctx context.Context, orgID string) ([]types.Dashboard, *model.ApiError) { dashboards := []types.Dashboard{} - err := db.NewSelect().Model(&dashboards).Where("org_id = ?", orgID).Scan(ctx) + err := store.BunDB().NewSelect().Model(&dashboards).Where("org_id = ?", orgID).Scan(ctx) if err != nil { return nil, &model.ApiError{Typ: model.ErrorExec, Err: err} } @@ -89,7 +89,7 @@ func DeleteDashboard(ctx context.Context, orgID, uuid string, fm interfaces.Feat return model.BadRequest(fmt.Errorf("dashboard is locked, please unlock the dashboard to be able to delete it")) } - result, err := db.NewDelete().Model(&types.Dashboard{}).Where("org_id = ?", orgID).Where("uuid = ?", uuid).Exec(ctx) + result, err := store.BunDB().NewDelete().Model(&types.Dashboard{}).Where("org_id = ?", orgID).Where("uuid = ?", uuid).Exec(ctx) if err != nil { return &model.ApiError{Typ: model.ErrorExec, Err: err} } @@ -108,7 +108,7 @@ func DeleteDashboard(ctx context.Context, orgID, uuid string, fm interfaces.Feat func GetDashboard(ctx context.Context, orgID, uuid string) (*types.Dashboard, *model.ApiError) { dashboard := types.Dashboard{} - err := db.NewSelect().Model(&dashboard).Where("org_id = ?", orgID).Where("uuid = ?", uuid).Scan(ctx) + err := store.BunDB().NewSelect().Model(&dashboard).Where("org_id = ?", orgID).Where("uuid = ?", uuid).Scan(ctx) if err != nil { return nil, &model.ApiError{Typ: model.ErrorNotFound, Err: fmt.Errorf("no dashboard found with uuid: %s", uuid)} } @@ -148,7 +148,7 @@ func UpdateDashboard(ctx context.Context, orgID, userEmail, uuid string, data ma dashboard.UpdatedBy = userEmail dashboard.Data = data - _, err = db.NewUpdate().Model(dashboard).Set("updated_at = ?", dashboard.UpdatedAt).Set("updated_by = ?", userEmail).Set("data = ?", mapData).Where("uuid = ?", dashboard.UUID).Exec(ctx) + _, err = store.BunDB().NewUpdate().Model(dashboard).Set("updated_at = ?", dashboard.UpdatedAt).Set("updated_by = ?", userEmail).Set("data = ?", mapData).Where("uuid = ?", dashboard.UUID).Exec(ctx) if err != nil { zap.L().Error("Error in inserting dashboard data", zap.Any("data", data), zap.Error(err)) @@ -170,7 +170,7 @@ func LockUnlockDashboard(ctx context.Context, orgID, uuid string, lock bool) *mo lockValue = 0 } - _, err := db.NewUpdate().Model(dashboard).Set("locked = ?", lockValue).Where("org_id = ?", orgID).Where("uuid = ?", uuid).Exec(ctx) + _, err := store.BunDB().NewUpdate().Model(dashboard).Set("locked = ?", lockValue).Where("org_id = ?", orgID).Where("uuid = ?", uuid).Exec(ctx) if err != nil { zap.L().Error("Error in updating dashboard", zap.String("uuid", uuid), zap.Error(err)) return &model.ApiError{Typ: model.ErrorExec, Err: err} @@ -242,7 +242,7 @@ func GetDashboardsInfo(ctx context.Context) (*model.DashboardsInfo, error) { dashboardsInfo := model.DashboardsInfo{} // fetch dashboards from dashboard db dashboards := []types.Dashboard{} - err := db.NewSelect().Model(&dashboards).Scan(ctx) + err := store.BunDB().NewSelect().Model(&dashboards).Scan(ctx) if err != nil { zap.L().Error("Error in processing sql query", zap.Error(err)) return &dashboardsInfo, err @@ -451,7 +451,7 @@ func countPanelsInDashboard(inputData map[string]interface{}) model.DashboardsIn func GetDashboardsWithMetricNames(ctx context.Context, orgID string, metricNames []string) (map[string][]map[string]string, *model.ApiError) { dashboards := []types.Dashboard{} - err := db.NewSelect().Model(&dashboards).Where("org_id = ?", orgID).Scan(ctx) + err := store.BunDB().NewSelect().Model(&dashboards).Where("org_id = ?", orgID).Scan(ctx) if err != nil { zap.L().Error("Error in getting dashboards", zap.Error(err)) return nil, &model.ApiError{Typ: model.ErrorExec, Err: err} diff --git a/pkg/query-service/app/explorer/db.go b/pkg/query-service/app/explorer/db.go index c397034059..1feb543a91 100644 --- a/pkg/query-service/app/explorer/db.go +++ b/pkg/query-service/app/explorer/db.go @@ -11,30 +11,30 @@ import ( "github.com/SigNoz/signoz/pkg/query-service/model" v3 "github.com/SigNoz/signoz/pkg/query-service/model/v3" "github.com/SigNoz/signoz/pkg/query-service/telemetry" + "github.com/SigNoz/signoz/pkg/sqlstore" "github.com/SigNoz/signoz/pkg/types" "github.com/SigNoz/signoz/pkg/types/authtypes" "github.com/google/uuid" - "github.com/uptrace/bun" "go.uber.org/zap" ) -var db *bun.DB +var store sqlstore.SQLStore // InitWithDSN sets up setting up the connection pool global variable. -func InitWithDSN(inputDB *bun.DB) error { - db = inputDB +func InitWithDSN(sqlStore sqlstore.SQLStore) error { + store = sqlStore telemetry.GetInstance().SetSavedViewsInfoCallback(GetSavedViewsInfo) return nil } -func InitWithDB(bunDB *bun.DB) { - db = bunDB +func InitWithDB(sqlStore sqlstore.SQLStore) { + store = sqlStore } func GetViews(ctx context.Context, orgID string) ([]*v3.SavedView, error) { var views []types.SavedView - err := db.NewSelect().Model(&views).Where("org_id = ?", orgID).Scan(ctx) + err := store.BunDB().NewSelect().Model(&views).Where("org_id = ?", orgID).Scan(ctx) if err != nil { return nil, fmt.Errorf("error in getting saved views: %s", err.Error()) } @@ -67,9 +67,9 @@ func GetViewsForFilters(ctx context.Context, orgID string, sourcePage string, na var views []types.SavedView var err error if len(category) == 0 { - err = db.NewSelect().Model(&views).Where("org_id = ? AND source_page = ? AND name LIKE ?", orgID, sourcePage, "%"+name+"%").Scan(ctx) + err = store.BunDB().NewSelect().Model(&views).Where("org_id = ? AND source_page = ? AND name LIKE ?", orgID, sourcePage, "%"+name+"%").Scan(ctx) } else { - err = db.NewSelect().Model(&views).Where("org_id = ? AND source_page = ? AND category LIKE ? AND name LIKE ?", orgID, sourcePage, "%"+category+"%", "%"+name+"%").Scan(ctx) + err = store.BunDB().NewSelect().Model(&views).Where("org_id = ? AND source_page = ? AND category LIKE ? AND name LIKE ?", orgID, sourcePage, "%"+category+"%", "%"+name+"%").Scan(ctx) } if err != nil { return nil, fmt.Errorf("error in getting saved views: %s", err.Error()) @@ -139,7 +139,7 @@ func CreateView(ctx context.Context, orgID string, view v3.SavedView) (string, e ExtraData: view.ExtraData, } - _, err = db.NewInsert().Model(&dbView).Exec(ctx) + _, err = store.BunDB().NewInsert().Model(&dbView).Exec(ctx) if err != nil { return "", fmt.Errorf("error in creating saved view: %s", err.Error()) } @@ -148,7 +148,7 @@ func CreateView(ctx context.Context, orgID string, view v3.SavedView) (string, e func GetView(ctx context.Context, orgID string, uuid_ string) (*v3.SavedView, error) { var view types.SavedView - err := db.NewSelect().Model(&view).Where("org_id = ? AND uuid = ?", orgID, uuid_).Scan(ctx) + err := store.BunDB().NewSelect().Model(&view).Where("org_id = ? AND uuid = ?", orgID, uuid_).Scan(ctx) if err != nil { return nil, fmt.Errorf("error in getting saved view: %s", err.Error()) } @@ -187,7 +187,7 @@ func UpdateView(ctx context.Context, orgID string, uuid_ string, view v3.SavedVi updatedAt := time.Now() updatedBy := claims.Email - _, err = db.NewUpdate(). + _, err = store.BunDB().NewUpdate(). Model(&types.SavedView{}). Set("updated_at = ?, updated_by = ?, name = ?, category = ?, source_page = ?, tags = ?, data = ?, extra_data = ?", updatedAt, updatedBy, view.Name, view.Category, view.SourcePage, strings.Join(view.Tags, ","), data, view.ExtraData). @@ -201,7 +201,7 @@ func UpdateView(ctx context.Context, orgID string, uuid_ string, view v3.SavedVi } func DeleteView(ctx context.Context, orgID string, uuid_ string) error { - _, err := db.NewDelete(). + _, err := store.BunDB().NewDelete(). Model(&types.SavedView{}). Where("uuid = ?", uuid_). Where("org_id = ?", orgID). @@ -216,7 +216,7 @@ func GetSavedViewsInfo(ctx context.Context) (*model.SavedViewsInfo, error) { savedViewsInfo := model.SavedViewsInfo{} // get single org ID from db var orgIDs []string - err := db.NewSelect().Model((*types.Organization)(nil)).Column("id").Scan(ctx, &orgIDs) + err := store.BunDB().NewSelect().Model((*types.Organization)(nil)).Column("id").Scan(ctx, &orgIDs) if err != nil { return nil, fmt.Errorf("error in getting org IDs: %s", err.Error()) } diff --git a/pkg/query-service/app/server.go b/pkg/query-service/app/server.go index 6f2ac1d8e3..706ea955b7 100644 --- a/pkg/query-service/app/server.go +++ b/pkg/query-service/app/server.go @@ -101,11 +101,11 @@ func NewServer(serverOptions *ServerOptions) (*Server, error) { return nil, err } - if err := dashboards.InitDB(serverOptions.SigNoz.SQLStore.BunDB()); err != nil { + if err := dashboards.InitDB(serverOptions.SigNoz.SQLStore); err != nil { return nil, err } - if err := explorer.InitWithDSN(serverOptions.SigNoz.SQLStore.BunDB()); err != nil { + if err := explorer.InitWithDSN(serverOptions.SigNoz.SQLStore); err != nil { return nil, err } diff --git a/pkg/query-service/utils/testutils.go b/pkg/query-service/utils/testutils.go index 29a5e547c2..ce1ae7c850 100644 --- a/pkg/query-service/utils/testutils.go +++ b/pkg/query-service/utils/testutils.go @@ -72,7 +72,7 @@ func NewQueryServiceDBForTests(t *testing.T) sqlstore.SQLStore { if err != nil { t.Fatalf("could not initialize dao: %v", err) } - _ = dashboards.InitDB(sqlStore.BunDB()) + _ = dashboards.InitDB(sqlStore) return sqlStore }