netbird/management/server/store/sql_store_idp_migration.go

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
}