signoz/pkg/sqlmigration/027_update_rules.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"
	"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      *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    *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"`
}

type schedule struct {
	Timezone   string      `json:"timezone"`
	StartTime  time.Time   `json:"startTime"`
	EndTime    time.Time   `json:"endTime,omitzero"`
	Recurrence *recurrence `json:"recurrence"`
}

type recurrence struct {
	StartTime  time.Time           `json:"startTime"`
	EndTime    time.Time           `json:"endTime,omitzero"`
	Duration   valuer.TextDuration `json:"duration"`
	RepeatType string              `json:"repeatType"`
	RepeatOn   []string            `json:"repeatOn"`
}

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
}