mirror of
https://git.mirrors.martin98.com/https://github.com/SigNoz/signoz
synced 2025-07-29 23:31:59 +08:00
118 lines
3.3 KiB
Go
118 lines
3.3 KiB
Go
package sqlite
|
|
|
|
import (
|
|
"context"
|
|
|
|
"github.com/SigNoz/signoz/pkg/query-service/constants"
|
|
"github.com/SigNoz/signoz/pkg/query-service/telemetry"
|
|
"github.com/SigNoz/signoz/pkg/sqlstore"
|
|
"github.com/SigNoz/signoz/pkg/types"
|
|
"github.com/pkg/errors"
|
|
"github.com/uptrace/bun"
|
|
"go.uber.org/zap"
|
|
)
|
|
|
|
type ModelDaoSqlite struct {
|
|
bundb *bun.DB
|
|
}
|
|
|
|
// InitDB sets up setting up the connection pool global variable.
|
|
func InitDB(sqlStore sqlstore.SQLStore) (*ModelDaoSqlite, error) {
|
|
mds := &ModelDaoSqlite{bundb: sqlStore.BunDB()}
|
|
|
|
ctx := context.Background()
|
|
if err := mds.initializeOrgPreferences(ctx); err != nil {
|
|
return nil, err
|
|
}
|
|
if err := mds.initializeRBAC(ctx); err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
telemetry.GetInstance().SetUserCountCallback(mds.GetUserCount)
|
|
telemetry.GetInstance().SetUserRoleCallback(mds.GetUserRole)
|
|
telemetry.GetInstance().SetGetUsersCallback(mds.GetUsers)
|
|
|
|
return mds, nil
|
|
}
|
|
|
|
// DB returns database connection
|
|
func (mds *ModelDaoSqlite) DB() *bun.DB {
|
|
return mds.bundb
|
|
}
|
|
|
|
// initializeOrgPreferences initializes in-memory telemetry settings. It is planned to have
|
|
// multiple orgs in the system. In case of multiple orgs, there will be separate instance
|
|
// of in-memory telemetry for each of the org, having their own settings. As of now, we only
|
|
// have one org so this method relies on the settings of this org to initialize the telemetry etc.
|
|
// TODO(Ahsan): Make it multi-tenant when we move to a system with multiple orgs.
|
|
func (mds *ModelDaoSqlite) initializeOrgPreferences(ctx context.Context) error {
|
|
|
|
// set anonymous setting as default in case of any failures to fetch UserPreference in below section
|
|
telemetry.GetInstance().SetTelemetryAnonymous(constants.DEFAULT_TELEMETRY_ANONYMOUS)
|
|
|
|
orgs, apiError := mds.GetOrgs(ctx)
|
|
if apiError != nil {
|
|
return apiError.Err
|
|
}
|
|
|
|
if len(orgs) > 1 {
|
|
return errors.Errorf("Found %d organizations, expected one or none.", len(orgs))
|
|
}
|
|
|
|
var org types.Organization
|
|
if len(orgs) == 1 {
|
|
org = orgs[0]
|
|
}
|
|
|
|
// set telemetry fields from userPreferences
|
|
telemetry.GetInstance().SetDistinctId(org.ID)
|
|
|
|
users, _ := mds.GetUsers(ctx)
|
|
countUsers := len(users)
|
|
if countUsers > 0 {
|
|
telemetry.GetInstance().SetCompanyDomain(users[countUsers-1].Email)
|
|
telemetry.GetInstance().SetUserEmail(users[countUsers-1].Email)
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
// initializeRBAC creates the ADMIN, EDITOR and VIEWER groups if they are not present.
|
|
func (mds *ModelDaoSqlite) initializeRBAC(ctx context.Context) error {
|
|
f := func(groupName string) error {
|
|
_, err := mds.createGroupIfNotPresent(ctx, groupName)
|
|
return errors.Wrap(err, "Failed to create group")
|
|
}
|
|
|
|
if err := f(constants.AdminGroup); err != nil {
|
|
return err
|
|
}
|
|
if err := f(constants.EditorGroup); err != nil {
|
|
return err
|
|
}
|
|
if err := f(constants.ViewerGroup); err != nil {
|
|
return err
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
func (mds *ModelDaoSqlite) createGroupIfNotPresent(ctx context.Context,
|
|
name string) (*types.Group, error) {
|
|
|
|
group, err := mds.GetGroupByName(ctx, name)
|
|
if err != nil {
|
|
return nil, errors.Wrap(err.Err, "Failed to query for root group")
|
|
}
|
|
if group != nil {
|
|
return group, nil
|
|
}
|
|
|
|
zap.L().Debug("group is not found, creating it", zap.String("group_name", name))
|
|
group, cErr := mds.CreateGroup(ctx, &types.Group{Name: name})
|
|
if cErr != nil {
|
|
return nil, cErr.Err
|
|
}
|
|
return group, nil
|
|
}
|