signoz/pkg/sqlmigration/023_update_apdex_ttl.go

package sqlmigration

import (
	"context"
	"database/sql"
	"fmt"
	"reflect"
	"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 updateApdexTtl struct {
	store sqlstore.SQLStore
}

type existingApdexSettings struct {
	bun.BaseModel      `bun:"table:apdex_settings"`
	OrgID              string  `bun:"org_id,pk,type:text" json:"orgId"`
	ServiceName        string  `bun:"service_name,pk,type:text" json:"serviceName"`
	Threshold          float64 `bun:"threshold,type:float,notnull" json:"threshold"`
	ExcludeStatusCodes string  `bun:"exclude_status_codes,type:text,notnull" json:"excludeStatusCodes"`
}

type newApdexSettings struct {
	bun.BaseModel `bun:"table:apdex_setting"`
	types.Identifiable
	OrgID              string  `bun:"org_id,type:text" json:"orgId"`
	ServiceName        string  `bun:"service_name,type:text" json:"serviceName"`
	Threshold          float64 `bun:"threshold,type:float,notnull" json:"threshold"`
	ExcludeStatusCodes string  `bun:"exclude_status_codes,type:text,notnull" json:"excludeStatusCodes"`
}

type existingTTLStatus struct {
	bun.BaseModel  `bun:"table:ttl_status"`
	ID             int       `bun:"id,pk,autoincrement"`
	TransactionID  string    `bun:"transaction_id,type:text,notnull"`
	CreatedAt      time.Time `bun:"created_at,type:datetime,notnull"`
	UpdatedAt      time.Time `bun:"updated_at,type:datetime,notnull"`
	TableName      string    `bun:"table_name,type:text,notnull"`
	TTL            int       `bun:"ttl,notnull,default:0"`
	ColdStorageTTL int       `bun:"cold_storage_ttl,notnull,default:0"`
	Status         string    `bun:"status,type:text,notnull"`
}

type newTTLStatus struct {
	bun.BaseModel `bun:"table:ttl_setting"`
	types.Identifiable
	types.TimeAuditable
	TransactionID  string `bun:"transaction_id,type:text,notnull"`
	TableName      string `bun:"table_name,type:text,notnull"`
	TTL            int    `bun:"ttl,notnull,default:0"`
	ColdStorageTTL int    `bun:"cold_storage_ttl,notnull,default:0"`
	Status         string `bun:"status,type:text,notnull"`
	OrgID          string `json:"-" bun:"org_id,notnull"`
}

func NewUpdateApdexTtlFactory(sqlstore sqlstore.SQLStore) factory.ProviderFactory[SQLMigration, Config] {
	return factory.NewProviderFactory(factory.MustNewName("update_apdex_ttl"), func(ctx context.Context, ps factory.ProviderSettings, c Config) (SQLMigration, error) {
		return newUpdateApdexTtl(ctx, ps, c, sqlstore)
	})
}

func newUpdateApdexTtl(_ context.Context, _ factory.ProviderSettings, _ Config, store sqlstore.SQLStore) (SQLMigration, error) {
	return &updateApdexTtl{store: store}, nil
}

func (migration *updateApdexTtl) Register(migrations *migrate.Migrations) error {
	if err := migrations.Register(migration.Up, migration.Down); err != nil {
		return err
	}

	return nil
}

func (migration *updateApdexTtl) Up(ctx context.Context, db *bun.DB) error {
	tx, err := db.BeginTx(ctx, nil)
	if err != nil {
		return err
	}

	defer tx.Rollback()

	err = migration.
		store.
		Dialect().
		RenameTableAndModifyModel(ctx, tx, new(existingApdexSettings), new(newApdexSettings), []string{OrgReference}, func(ctx context.Context) error {
			existingApdexSettings := make([]*existingApdexSettings, 0)
			err = tx.
				NewSelect().
				Model(&existingApdexSettings).
				Scan(ctx)
			if err != nil {
				if err != sql.ErrNoRows {
					return err
				}
			}

			if err == nil && len(existingApdexSettings) > 0 {
				newSettings := migration.
					CopyExistingApdexSettingsToNewApdexSettings(existingApdexSettings)
				_, err = tx.
					NewInsert().
					Model(&newSettings).
					Exec(ctx)
				if err != nil {
					return err
				}
			}

			tableName := tx.Dialect().Tables().Get(reflect.TypeOf(new(newApdexSettings))).Name
			_, err = tx.
				ExecContext(ctx, fmt.Sprintf("CREATE UNIQUE INDEX IF NOT EXISTS %s_unique_idx ON %s (service_name, org_id)", tableName, tableName))
			if err != nil {
				return err
			}
			return nil
		})
	if err != nil {
		return err
	}

	err = migration.
		store.
		Dialect().
		RenameTableAndModifyModel(ctx, tx, new(existingTTLStatus), new(newTTLStatus), []string{OrgReference}, func(ctx context.Context) error {
			existingTTLStatus := make([]*existingTTLStatus, 0)
			err = tx.
				NewSelect().
				Model(&existingTTLStatus).
				Scan(ctx)
			if err != nil {
				if err != sql.ErrNoRows {
					return err
				}
			}

			if err == nil && len(existingTTLStatus) > 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 {
					newTTLStatus := migration.CopyExistingTTLStatusToNewTTLStatus(existingTTLStatus, orgID)
					_, err = tx.
						NewInsert().
						Model(&newTTLStatus).
						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 *updateApdexTtl) Down(context.Context, *bun.DB) error {
	return nil
}

func (migration *updateApdexTtl) CopyExistingApdexSettingsToNewApdexSettings(existingApdexSettings []*existingApdexSettings) []*newApdexSettings {
	newSettings := make([]*newApdexSettings, 0)
	for _, apdexSetting := range existingApdexSettings {
		newSettings = append(newSettings, &newApdexSettings{
			Identifiable: types.Identifiable{
				ID: valuer.GenerateUUID(),
			},
			ServiceName:        apdexSetting.ServiceName,
			Threshold:          apdexSetting.Threshold,
			ExcludeStatusCodes: apdexSetting.ExcludeStatusCodes,
			OrgID:              apdexSetting.OrgID,
		})
	}

	return newSettings
}

func (migration *updateApdexTtl) CopyExistingTTLStatusToNewTTLStatus(existingTTLStatus []*existingTTLStatus, orgID string) []*newTTLStatus {
	newTTLStatuses := make([]*newTTLStatus, 0)

	for _, ttl := range existingTTLStatus {
		newTTLStatuses = append(newTTLStatuses, &newTTLStatus{
			Identifiable: types.Identifiable{
				ID: valuer.GenerateUUID(),
			},
			TimeAuditable: types.TimeAuditable{
				CreatedAt: ttl.CreatedAt,
				UpdatedAt: ttl.UpdatedAt,
			},
			TransactionID:  ttl.TransactionID,
			TTL:            ttl.TTL,
			TableName:      ttl.TableName,
			ColdStorageTTL: ttl.ColdStorageTTL,
			Status:         ttl.Status,
			OrgID:          orgID,
		})
	}

	return newTTLStatuses
}