signoz/pkg/sqlmigration/036_add_key_organization.go

package sqlmigration

import (
	"context"
	"hash/fnv"

	"github.com/SigNoz/signoz/pkg/factory"
	"github.com/SigNoz/signoz/pkg/sqlstore"
	"github.com/uptrace/bun"
	"github.com/uptrace/bun/migrate"
)

type addKeyOrganization struct {
	sqlstore sqlstore.SQLStore
}

func NewAddKeyOrganizationFactory(sqlstore sqlstore.SQLStore) factory.ProviderFactory[SQLMigration, Config] {
	return factory.NewProviderFactory(factory.MustNewName("add_key_organization"), func(ctx context.Context, providerSettings factory.ProviderSettings, config Config) (SQLMigration, error) {
		return newAddKeyOrganization(ctx, providerSettings, config, sqlstore)
	})
}

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

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

	return nil
}

func (migration *addKeyOrganization) Up(ctx context.Context, db *bun.DB) error {
	ok, err := migration.sqlstore.Dialect().ColumnExists(ctx, db, "organizations", "key")
	if err != nil {
		return err
	}

	if ok {
		return nil
	}

	tx, err := db.BeginTx(ctx, nil)
	if err != nil {
		return err
	}

	defer func() {
		_ = tx.Rollback()
	}()

	if _, err := tx.
		NewAddColumn().
		Table("organizations").
		ColumnExpr("key BIGINT").
		Exec(ctx); err != nil {
		return err
	}

	var existingOrgIDs []string
	if err := tx.NewSelect().
		Table("organizations").
		Column("id").
		Scan(ctx, &existingOrgIDs); err != nil {
		return err
	}

	for _, orgID := range existingOrgIDs {
		key := migration.getHash(ctx, orgID)
		if _, err := tx.
			NewUpdate().
			Table("organizations").
			Set("key = ?", key).
			Where("id = ?", orgID).
			Exec(ctx); err != nil {
			return err
		}
	}

	if _, err := tx.
		NewCreateIndex().
		Unique().
		IfNotExists().
		Index("idx_unique_key").
		Table("organizations").
		Column("key").
		Exec(ctx); err != nil {
		return err
	}

	if err := tx.Commit(); err != nil {
		return err
	}

	return nil
}

func (migration *addKeyOrganization) Down(ctx context.Context, db *bun.DB) error {
	return nil
}

func (migration *addKeyOrganization) getHash(_ context.Context, orgID string) uint32 {
	hasher := fnv.New32a()

	// Hasher never returns err.
	_, _ = hasher.Write([]byte(orgID))
	return hasher.Sum32()
}