mirror of
https://git.mirrors.martin98.com/https://github.com/SigNoz/signoz
synced 2025-06-04 11:25:52 +08:00
346 lines
9.2 KiB
Go
346 lines
9.2 KiB
Go
package sqlmigration
|
|
|
|
import (
|
|
"context"
|
|
"database/sql"
|
|
"database/sql/driver"
|
|
"encoding/json"
|
|
"strconv"
|
|
"time"
|
|
|
|
"github.com/SigNoz/signoz/pkg/factory"
|
|
"github.com/SigNoz/signoz/pkg/sqlstore"
|
|
"github.com/SigNoz/signoz/pkg/types"
|
|
ruletypes "github.com/SigNoz/signoz/pkg/types/ruletypes"
|
|
"github.com/SigNoz/signoz/pkg/valuer"
|
|
"github.com/uptrace/bun"
|
|
"github.com/uptrace/bun/migrate"
|
|
)
|
|
|
|
type updateRules struct {
|
|
store sqlstore.SQLStore
|
|
}
|
|
|
|
type AlertIds []string
|
|
|
|
func (a *AlertIds) Scan(src interface{}) error {
|
|
if data, ok := src.([]byte); ok {
|
|
return json.Unmarshal(data, a)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (a *AlertIds) Value() (driver.Value, error) {
|
|
return json.Marshal(a)
|
|
}
|
|
|
|
type existingRule struct {
|
|
bun.BaseModel `bun:"table:rules"`
|
|
ID int `bun:"id,pk,autoincrement"`
|
|
CreatedAt time.Time `bun:"created_at,type:datetime,notnull"`
|
|
CreatedBy string `bun:"created_by,type:text,notnull"`
|
|
UpdatedAt time.Time `bun:"updated_at,type:datetime,notnull"`
|
|
UpdatedBy string `bun:"updated_by,type:text,notnull"`
|
|
Deleted int `bun:"deleted,notnull,default:0"`
|
|
Data string `bun:"data,type:text,notnull"`
|
|
}
|
|
|
|
type newRule struct {
|
|
bun.BaseModel `bun:"table:rule"`
|
|
types.Identifiable
|
|
types.TimeAuditable
|
|
types.UserAuditable
|
|
Deleted int `bun:"deleted,notnull,default:0"`
|
|
Data string `bun:"data,type:text,notnull"`
|
|
OrgID string `bun:"org_id,type:text"`
|
|
}
|
|
|
|
type existingMaintenance struct {
|
|
bun.BaseModel `bun:"table:planned_maintenance"`
|
|
ID int `bun:"id,pk,autoincrement"`
|
|
Name string `bun:"name,type:text,notnull"`
|
|
Description string `bun:"description,type:text"`
|
|
AlertIDs *AlertIds `bun:"alert_ids,type:text"`
|
|
Schedule *ruletypes.Schedule `bun:"schedule,type:text,notnull"`
|
|
CreatedAt time.Time `bun:"created_at,type:datetime,notnull"`
|
|
CreatedBy string `bun:"created_by,type:text,notnull"`
|
|
UpdatedAt time.Time `bun:"updated_at,type:datetime,notnull"`
|
|
UpdatedBy string `bun:"updated_by,type:text,notnull"`
|
|
}
|
|
|
|
type newMaintenance struct {
|
|
bun.BaseModel `bun:"table:planned_maintenance_new"`
|
|
types.Identifiable
|
|
types.TimeAuditable
|
|
types.UserAuditable
|
|
Name string `bun:"name,type:text,notnull"`
|
|
Description string `bun:"description,type:text"`
|
|
Schedule *ruletypes.Schedule `bun:"schedule,type:text,notnull"`
|
|
OrgID string `bun:"org_id,type:text"`
|
|
}
|
|
|
|
type storablePlannedMaintenanceRule struct {
|
|
bun.BaseModel `bun:"table:planned_maintenance_rule"`
|
|
types.Identifiable
|
|
PlannedMaintenanceID valuer.UUID `bun:"planned_maintenance_id,type:text"`
|
|
RuleID valuer.UUID `bun:"rule_id,type:text"`
|
|
}
|
|
|
|
type ruleHistory struct {
|
|
bun.BaseModel `bun:"table:rule_history"`
|
|
RuleID int `bun:"rule_id"`
|
|
RuleUUID valuer.UUID `bun:"rule_uuid"`
|
|
}
|
|
|
|
func NewUpdateRulesFactory(sqlstore sqlstore.SQLStore) factory.ProviderFactory[SQLMigration, Config] {
|
|
return factory.NewProviderFactory(factory.MustNewName("update_rules"), func(ctx context.Context, ps factory.ProviderSettings, c Config) (SQLMigration, error) {
|
|
return newUpdateRules(ctx, ps, c, sqlstore)
|
|
})
|
|
}
|
|
|
|
func newUpdateRules(_ context.Context, _ factory.ProviderSettings, _ Config, store sqlstore.SQLStore) (SQLMigration, error) {
|
|
return &updateRules{store: store}, nil
|
|
}
|
|
|
|
func (migration *updateRules) Register(migrations *migrate.Migrations) error {
|
|
if err := migrations.Register(migration.Up, migration.Down); err != nil {
|
|
return err
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
func (migration *updateRules) Up(ctx context.Context, db *bun.DB) error {
|
|
tx, err := db.BeginTx(ctx, nil)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
defer func() {
|
|
_ = tx.Rollback()
|
|
}()
|
|
|
|
ruleIDToRuleUUIDMap := map[int]valuer.UUID{}
|
|
err = migration.
|
|
store.
|
|
Dialect().
|
|
RenameTableAndModifyModel(ctx, tx, new(existingRule), new(newRule), []string{OrgReference}, func(ctx context.Context) error {
|
|
existingRules := make([]*existingRule, 0)
|
|
err := tx.
|
|
NewSelect().
|
|
Model(&existingRules).
|
|
Scan(ctx)
|
|
if err != nil {
|
|
if err != sql.ErrNoRows {
|
|
return err
|
|
}
|
|
}
|
|
if err == nil && len(existingRules) > 0 {
|
|
var orgID string
|
|
err := migration.
|
|
store.
|
|
BunDB().
|
|
NewSelect().
|
|
Model((*types.Organization)(nil)).
|
|
Column("id").
|
|
Scan(ctx, &orgID)
|
|
if err != nil {
|
|
if err != sql.ErrNoRows {
|
|
return err
|
|
}
|
|
}
|
|
if err == nil {
|
|
newRules, idUUIDMap := migration.CopyExistingRulesToNewRules(existingRules, orgID)
|
|
ruleIDToRuleUUIDMap = idUUIDMap
|
|
_, err = tx.
|
|
NewInsert().
|
|
Model(&newRules).
|
|
Exec(ctx)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
}
|
|
}
|
|
err = migration.store.Dialect().UpdatePrimaryKey(ctx, tx, new(existingMaintenance), new(newMaintenance), OrgReference, func(ctx context.Context) error {
|
|
_, err := tx.
|
|
NewCreateTable().
|
|
IfNotExists().
|
|
Model(new(storablePlannedMaintenanceRule)).
|
|
ForeignKey(`("planned_maintenance_id") REFERENCES "planned_maintenance_new" ("id") ON DELETE CASCADE ON UPDATE CASCADE`).
|
|
ForeignKey(`("rule_id") REFERENCES "rule" ("id")`).
|
|
Exec(ctx)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
existingMaintenances := make([]*existingMaintenance, 0)
|
|
err = tx.
|
|
NewSelect().
|
|
Model(&existingMaintenances).
|
|
Scan(ctx)
|
|
if err != nil {
|
|
if err != sql.ErrNoRows {
|
|
return err
|
|
}
|
|
}
|
|
if err == nil && len(existingMaintenances) > 0 {
|
|
var orgID string
|
|
err := migration.
|
|
store.
|
|
BunDB().
|
|
NewSelect().
|
|
Model((*types.Organization)(nil)).
|
|
Column("id").
|
|
Scan(ctx, &orgID)
|
|
if err != nil {
|
|
if err != sql.ErrNoRows {
|
|
return err
|
|
}
|
|
}
|
|
if err == nil {
|
|
newMaintenances, newMaintenancesRules, err := migration.CopyExistingMaintenancesToNewMaintenancesAndRules(existingMaintenances, orgID, ruleIDToRuleUUIDMap)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
_, err = tx.
|
|
NewInsert().
|
|
Model(&newMaintenances).
|
|
Exec(ctx)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
if len(newMaintenancesRules) > 0 {
|
|
_, err = tx.
|
|
NewInsert().
|
|
Model(&newMaintenancesRules).
|
|
Exec(ctx)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
}
|
|
|
|
}
|
|
|
|
}
|
|
return nil
|
|
})
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
ruleHistories := make([]*ruleHistory, 0)
|
|
for ruleID, ruleUUID := range ruleIDToRuleUUIDMap {
|
|
ruleHistories = append(ruleHistories, &ruleHistory{
|
|
RuleID: ruleID,
|
|
RuleUUID: ruleUUID,
|
|
})
|
|
}
|
|
|
|
_, err = tx.
|
|
NewCreateTable().
|
|
IfNotExists().
|
|
Model(&ruleHistories).
|
|
Exec(ctx)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if len(ruleHistories) > 0 {
|
|
_, err = tx.
|
|
NewInsert().
|
|
Model(&ruleHistories).
|
|
Exec(ctx)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
}
|
|
|
|
return nil
|
|
})
|
|
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
err = tx.Commit()
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
func (migration *updateRules) Down(context.Context, *bun.DB) error {
|
|
return nil
|
|
}
|
|
|
|
func (migration *updateRules) CopyExistingRulesToNewRules(existingRules []*existingRule, orgID string) ([]*newRule, map[int]valuer.UUID) {
|
|
newRules := make([]*newRule, 0)
|
|
idUUIDMap := map[int]valuer.UUID{}
|
|
for _, rule := range existingRules {
|
|
uuid := valuer.GenerateUUID()
|
|
idUUIDMap[rule.ID] = uuid
|
|
newRules = append(newRules, &newRule{
|
|
Identifiable: types.Identifiable{
|
|
ID: uuid,
|
|
},
|
|
TimeAuditable: types.TimeAuditable{
|
|
CreatedAt: rule.CreatedAt,
|
|
UpdatedAt: rule.UpdatedAt,
|
|
},
|
|
UserAuditable: types.UserAuditable{
|
|
CreatedBy: rule.CreatedBy,
|
|
UpdatedBy: rule.UpdatedBy,
|
|
},
|
|
Deleted: rule.Deleted,
|
|
Data: rule.Data,
|
|
OrgID: orgID,
|
|
})
|
|
}
|
|
return newRules, idUUIDMap
|
|
}
|
|
|
|
func (migration *updateRules) CopyExistingMaintenancesToNewMaintenancesAndRules(existingMaintenances []*existingMaintenance, orgID string, ruleIDToRuleUUIDMap map[int]valuer.UUID) ([]*newMaintenance, []*storablePlannedMaintenanceRule, error) {
|
|
newMaintenances := make([]*newMaintenance, 0)
|
|
newMaintenanceRules := make([]*storablePlannedMaintenanceRule, 0)
|
|
|
|
for _, maintenance := range existingMaintenances {
|
|
ruleIDs := maintenance.AlertIDs
|
|
maintenanceUUID := valuer.GenerateUUID()
|
|
newMaintenance := newMaintenance{
|
|
Identifiable: types.Identifiable{
|
|
ID: maintenanceUUID,
|
|
},
|
|
TimeAuditable: types.TimeAuditable{
|
|
CreatedAt: maintenance.CreatedAt,
|
|
UpdatedAt: maintenance.UpdatedAt,
|
|
},
|
|
UserAuditable: types.UserAuditable{
|
|
CreatedBy: maintenance.CreatedBy,
|
|
UpdatedBy: maintenance.UpdatedBy,
|
|
},
|
|
Name: maintenance.Name,
|
|
Description: maintenance.Description,
|
|
Schedule: maintenance.Schedule,
|
|
OrgID: orgID,
|
|
}
|
|
newMaintenances = append(newMaintenances, &newMaintenance)
|
|
for _, ruleIDStr := range *ruleIDs {
|
|
ruleID, err := strconv.Atoi(ruleIDStr)
|
|
if err != nil {
|
|
return nil, nil, err
|
|
}
|
|
|
|
newMaintenanceRules = append(newMaintenanceRules, &storablePlannedMaintenanceRule{
|
|
Identifiable: types.Identifiable{
|
|
ID: valuer.GenerateUUID(),
|
|
},
|
|
PlannedMaintenanceID: maintenanceUUID,
|
|
RuleID: ruleIDToRuleUUIDMap[ruleID],
|
|
})
|
|
}
|
|
}
|
|
return newMaintenances, newMaintenanceRules, nil
|
|
}
|