cloudfoundry-incubator/eirini

package migrations

import (
    "context"
    "fmt"

    "github.com/pkg/errors"
    appsv1 "k8s.io/api/apps/v1"
    "k8s.io/api/policy/v1beta1"
    "k8s.io/apimachinery/pkg/runtime"
)

//counterfeiter:generate . PDBClient

type PDBClient interface {
    Get(ctx context.Context, namespace, name string) (*v1beta1.PodDisruptionBudget, error)
    SetOwner(ctx context.Context, pdb *v1beta1.PodDisruptionBudget, owner *appsv1.StatefulSet) (*v1beta1.PodDisruptionBudget, error)
}

type AdoptPDB struct {
    pdbClient PDBClient
}

func NewAdoptPDB(pdbClient PDBClient) AdoptPDB {
    return AdoptPDB{
        pdbClient: pdbClient,
    }
}

func (m AdoptPDB) Apply(ctx context.Context, obj runtime.Object) error {
    stSet, ok := obj.(*appsv1.StatefulSet)
    if !ok {
        return fmt.Errorf("expected *v1.StatefulSet, got: %T", obj)
    }

    if *stSet.Spec.Replicas <= 1 {
        return nil
    }

    pdb, err := m.pdbClient.Get(ctx, stSet.Namespace, stSet.Name)
    if err != nil {
        return errors.Wrap(err, "adopt-pdb-migration-get-pdb-failed")
    }

    if _, err := m.pdbClient.SetOwner(ctx, pdb, stSet); err != nil {
        return errors.Wrap(err, "adopt-pdb-migration-set-owner-failed")
    }

    return nil
}

func (m AdoptPDB) SequenceID() int {
    return AdoptPDBSequenceID
}

func (m AdoptPDB) AppliesTo() ObjectType {
    return StatefulSetObjectType
}