[management] Legacy to embedded IdP migration tool (#5586)

management/server/store/sql_store_idp_migration.go

@@ -0,0 +1,177 @@
+package store
+
+// This file contains migration-only methods on SqlStore.
+// They satisfy the migration.Store interface via duck typing.
+// Delete this file when migration tooling is no longer needed.
+
+import (
+	"context"
+	"fmt"
+
+	log "github.com/sirupsen/logrus"
+	"gorm.io/gorm"
+
+	"github.com/netbirdio/netbird/management/server/idp/migration"
+	nbpeer "github.com/netbirdio/netbird/management/server/peer"
+	"github.com/netbirdio/netbird/management/server/types"
+	"github.com/netbirdio/netbird/shared/management/status"
+)
+
+func (s *SqlStore) CheckSchema(checks []migration.SchemaCheck) []migration.SchemaError {
+	migrator := s.db.Migrator()
+	var errs []migration.SchemaError
+
+	for _, check := range checks {
+		if !migrator.HasTable(check.Table) {
+			errs = append(errs, migration.SchemaError{Table: check.Table})
+			continue
+		}
+		for _, col := range check.Columns {
+			if !migrator.HasColumn(check.Table, col) {
+				errs = append(errs, migration.SchemaError{Table: check.Table, Column: col})
+			}
+		}
+	}
+
+	return errs
+}
+
+func (s *SqlStore) ListUsers(ctx context.Context) ([]*types.User, error) {
+	tx := s.db
+	var users []*types.User
+	result := tx.Find(&users)
+	if result.Error != nil {
+		log.WithContext(ctx).Errorf("error when listing users from the store: %s", result.Error)
+		return nil, status.Errorf(status.Internal, "issue listing users from store")
+	}
+
+	for _, user := range users {
+		if err := user.DecryptSensitiveData(s.fieldEncrypt); err != nil {
+			return nil, fmt.Errorf("decrypt user: %w", err)
+		}
+	}
+
+	return users, nil
+}
+
+// txDeferFKConstraints defers foreign key constraint checks for the duration of the transaction.
+// MySQL is already handled by s.transaction (SET FOREIGN_KEY_CHECKS = 0).
+func (s *SqlStore) txDeferFKConstraints(tx *gorm.DB) error {
+	if s.storeEngine == types.SqliteStoreEngine {
+		return tx.Exec("PRAGMA defer_foreign_keys = ON").Error
+	}
+
+	if s.storeEngine != types.PostgresStoreEngine {
+		return nil
+	}
+
+	// GORM creates FK constraints as NOT DEFERRABLE by default, so
+	// SET CONSTRAINTS ALL DEFERRED is a no-op unless we ALTER them first.
+	err := tx.Exec(`
+			DO $$ DECLARE r RECORD;
+			BEGIN
+				FOR r IN SELECT conname, conrelid::regclass AS tbl
+				         FROM pg_constraint WHERE contype = 'f' AND NOT condeferrable
+				LOOP
+					EXECUTE format('ALTER TABLE %s ALTER CONSTRAINT %I DEFERRABLE INITIALLY IMMEDIATE', r.tbl, r.conname);
+				END LOOP;
+			END $$
+		`).Error
+	if err != nil {
+		return fmt.Errorf("make FK constraints deferrable: %w", err)
+	}
+	return tx.Exec("SET CONSTRAINTS ALL DEFERRED").Error
+}
+
+// txRestoreFKConstraints reverts FK constraints back to NOT DEFERRABLE after the
+// deferred updates are done but before the transaction commits.
+func (s *SqlStore) txRestoreFKConstraints(tx *gorm.DB) error {
+	if s.storeEngine != types.PostgresStoreEngine {
+		return nil
+	}
+
+	return tx.Exec(`
+		DO $$ DECLARE r RECORD;
+		BEGIN
+			FOR r IN SELECT conname, conrelid::regclass AS tbl
+			         FROM pg_constraint WHERE contype = 'f' AND condeferrable
+			LOOP
+				EXECUTE format('ALTER TABLE %s ALTER CONSTRAINT %I NOT DEFERRABLE', r.tbl, r.conname);
+			END LOOP;
+		END $$
+	`).Error
+}
+
+func (s *SqlStore) UpdateUserInfo(ctx context.Context, userID, email, name string) error {
+	user := &types.User{Email: email, Name: name}
+	if err := user.EncryptSensitiveData(s.fieldEncrypt); err != nil {
+		return fmt.Errorf("encrypt user info: %w", err)
+	}
+
+	result := s.db.Model(&types.User{}).Where("id = ?", userID).Updates(map[string]any{
+		"email": user.Email,
+		"name":  user.Name,
+	})
+	if result.Error != nil {
+		log.WithContext(ctx).Errorf("error updating user info for %s: %s", userID, result.Error)
+		return status.Errorf(status.Internal, "failed to update user info")
+	}
+
+	return nil
+}
+
+func (s *SqlStore) UpdateUserID(ctx context.Context, accountID, oldUserID, newUserID string) error {
+	type fkUpdate struct {
+		model  any
+		column string
+		where  string
+	}
+
+	updates := []fkUpdate{
+		{&types.PersonalAccessToken{}, "user_id", "user_id = ?"},
+		{&types.PersonalAccessToken{}, "created_by", "created_by = ?"},
+		{&nbpeer.Peer{}, "user_id", "user_id = ?"},
+		{&types.UserInviteRecord{}, "created_by", "created_by = ?"},
+		{&types.Account{}, "created_by", "created_by = ?"},
+		{&types.ProxyAccessToken{}, "created_by", "created_by = ?"},
+		{&types.Job{}, "triggered_by", "triggered_by = ?"},
+	}
+
+	log.Info("Updating user ID in the store")
+	err := s.transaction(func(tx *gorm.DB) error {
+		if err := s.txDeferFKConstraints(tx); err != nil {
+			return err
+		}
+
+		for _, u := range updates {
+			if err := tx.Model(u.model).Where(u.where, oldUserID).Update(u.column, newUserID).Error; err != nil {
+				return fmt.Errorf("update %s: %w", u.column, err)
+			}
+		}
+
+		if err := tx.Model(&types.User{}).Where(accountAndIDQueryCondition, accountID, oldUserID).Update("id", newUserID).Error; err != nil {
+			return fmt.Errorf("update users: %w", err)
+		}
+
+		return nil
+	})
+	if err != nil {
+		log.WithContext(ctx).Errorf("failed to update user ID in the store: %s", err)
+		return status.Errorf(status.Internal, "failed to update user ID in store")
+	}
+
+	log.Info("Restoring FK constraints")
+	err = s.transaction(func(tx *gorm.DB) error {
+		if err := s.txRestoreFKConstraints(tx); err != nil {
+			return fmt.Errorf("restore FK constraints: %w", err)
+		}
+
+		return nil
+	})
+	if err != nil {
+		log.WithContext(ctx).Errorf("failed to restore FK constraints after user ID update: %s", err)
+		return status.Errorf(status.Internal, "failed to restore FK constraints after user ID update")
+	}
+
+	return nil
+}