Building a Go Kubernetes Operator for Automated Postgres Failover

package v1

import (
	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
)

// PostgresClusterSpec defines the desired state of PostgresCluster
type PostgresClusterSpec struct {
	// Number of desired pods. This is a pointer to distinguish between explicit
	// zero and not specified. Defaults to 3.
	// +kubebuilder:validation:Minimum=1
	// +kubebuilder:default=3
	Replicas *int32 `json:"replicas"`

	// Image is the container image to use for the PostgreSQL cluster.
	Image string `json:"image"`

	// Resources required by each PostgreSQL instance.
	Resources corev1.ResourceRequirements `json:"resources,omitempty"`

	// StorageSpec defines the storage configuration for the PostgreSQL data.
	StorageSpec corev1.PersistentVolumeClaimSpec `json:"storageSpec"`
}

// PostgresClusterStatus defines the observed state of PostgresCluster
type PostgresClusterStatus struct {
	// ReadyReplicas is the number of pods in the cluster that are in a Ready state.
	ReadyReplicas int32 `json:"readyReplicas"`

	// CurrentPrimary is the name of the pod that is currently acting as the primary.
	CurrentPrimary string `json:"currentPrimary,omitempty"`

	// Conditions represent the latest available observations of the PostgresCluster's state.
	// +optional
	Conditions []metav1.Condition `json:"conditions,omitempty"`
}

//+kubebuilder:object:root=true
//+kubebuilder:subresource:status
//+kubebuilder:printcolumn:name="Replicas",type="integer",JSONPath=".spec.replicas"
//+kubebuilder:printcolumn:name="Ready",type="string",JSONPath=".status.readyReplicas"
//+kubebuilder:printcolumn:name="Primary",type="string",JSONPath=".status.currentPrimary"
//+kubebuilder:printcolumn:name="Age",type="date",JSONPath=".metadata.creationTimestamp"

// PostgresCluster is the Schema for the postgresclusters API
type PostgresCluster struct {
	metav1.TypeMeta   `json:",inline"`
	metav1.ObjectMeta `json:"metadata,omitempty"`

	Spec   PostgresClusterSpec   `json:"spec,omitempty"`
	Status PostgresClusterStatus `json:"status,omitempty"`
}

//+kubebuilder:object:root=true

// PostgresClusterList contains a list of PostgresCluster
type PostgresClusterList struct {
	metav1.TypeMeta `json:",inline"`
	metav1.ListMeta `json:"metadata,omitempty"`
	Items           []PostgresCluster `json:"items"`
}

func init() {
	SchemeBuilder.Register(&PostgresCluster{}, &PostgresClusterList{})
}

// controllers/postgrescluster_controller.go

func (r *PostgresClusterReconciler) Reconcile(ctx context.Context, req ctrl.Request) (ctrl.Result, error) {
	log := log.FromContext(ctx)

	// 1. Fetch the PostgresCluster instance
	var pgCluster dbv1.PostgresCluster
	if err := r.Get(ctx, req.NamespacedName, &pgCluster); err != nil {
		// If the resource is not found, it might have been deleted. We can ignore this.
		return ctrl.Result{}, client.IgnoreNotFound(err)
	}

	// 2. Handle Deletion: Check for a deletion timestamp and run finalizers
	// ... (We will implement this later)

	// 3. Reconcile the StatefulSet for PostgreSQL pods
	sts, err := r.reconcileStatefulSet(ctx, &pgCluster)
	if err != nil {
		log.Error(err, "Failed to reconcile StatefulSet")
		// Update status with a degraded condition
		return ctrl.Result{}, err // Requeue with error
	}

	// 4. Reconcile the Services (Primary and Replica)
	err = r.reconcileServices(ctx, &pgCluster)
	if err != nil {
		log.Error(err, "Failed to reconcile Services")
		return ctrl.Result{}, err
	}

	// 5. Perform Primary Election and Failover Logic
	// This is the most complex part
	primaryPodName, err := r.managePrimaryAndReplicas(ctx, &pgCluster, sts)
	if err != nil {
		log.Error(err, "Failed to manage primary/replica roles")
		return ctrl.Result{}, err
	}

	// 6. Update the Status Subresource
	err = r.updateClusterStatus(ctx, &pgCluster, sts, primaryPodName)
	if err != nil {
		log.Error(err, "Failed to update PostgresCluster status")
		return ctrl.Result{}, err
	}

	// Requeue after a short duration to re-check the state
	return ctrl.Result{RequeueAfter: 30 * time.Second}, nil
}

// controllers/postgrescluster_controller.go (helper function)

func (r *PostgresClusterReconciler) reconcileStatefulSet(ctx context.Context, pgCluster *dbv1.PostgresCluster) (*appsv1.StatefulSet, error) {
	log := log.FromContext(ctx)

	desiredSts := r.defineStatefulSet(pgCluster)

	// Set PostgresCluster instance as the owner and controller
	if err := controllerutil.SetControllerReference(pgCluster, desiredSts, r.Scheme); err != nil {
		return nil, err
	}

	// Check if the StatefulSet already exists
	foundSts := &appsv1.StatefulSet{}
	err := r.Get(ctx, types.NamespacedName{Name: desiredSts.Name, Namespace: desiredSts.Namespace}, foundSts)
	if err != nil && errors.IsNotFound(err) {
		log.Info("Creating a new StatefulSet", "StatefulSet.Namespace", desiredSts.Namespace, "StatefulSet.Name", desiredSts.Name)
		err = r.Create(ctx, desiredSts)
		if err != nil {
			return nil, err
		}
		// StatefulSet created successfully - return the new STS
		return desiredSts, nil
	} else if err != nil {
		return nil, err
	}

	// StatefulSet already exists - check for updates
	// A simple check here: if replica count or image has changed, update it.
	// Production operators use more sophisticated deep-diffing.
	if *foundSts.Spec.Replicas != *desiredSts.Spec.Replicas || foundSts.Spec.Template.Spec.Containers[0].Image != desiredSts.Spec.Template.Spec.Containers[0].Image {
		log.Info("Updating StatefulSet", "StatefulSet.Namespace", foundSts.Namespace, "StatefulSet.Name", foundSts.Name)
		foundSts.Spec.Replicas = desiredSts.Spec.Replicas
		foundSts.Spec.Template.Spec.Containers[0].Image = desiredSts.Spec.Template.Spec.Containers[0].Image
		err = r.Update(ctx, foundSts)
		if err != nil {
			return nil, err
		}
	}

	return foundSts, nil
}

// defineStatefulSet creates the StatefulSet definition based on the CR spec.
func (r *PostgresClusterReconciler) defineStatefulSet(pgCluster *dbv1.PostgresCluster) *appsv1.StatefulSet {
    // ... logic to build the StatefulSet object ...
    // This includes setting labels, container spec, volume mounts, and volume claim templates from pgCluster.Spec.StorageSpec
    // Crucially, we add a label selector that we'll use to find our pods.
    labels := map[string]string{"app": "postgres", "cluster": pgCluster.Name}

    return &appsv1.StatefulSet{
        // ... metadata ...
        Spec: appsv1.StatefulSetSpec{
            Replicas: pgCluster.Spec.Replicas,
            Selector: &metav1.LabelSelector{MatchLabels: labels},
            Template: corev1.PodTemplateSpec{
                ObjectMeta: metav1.ObjectMeta{Labels: labels},
                Spec: corev1.PodSpec{
                    Containers: []corev1.Container{{
                        Name:  "postgres",
                        Image: pgCluster.Spec.Image,
                        // ... ports, env vars, probes, volumeMounts ...
                    }},
                },
            },
            VolumeClaimTemplates: []corev1.PersistentVolumeClaim{{
                ObjectMeta: metav1.ObjectMeta{Name: "pgdata"},
                Spec:       pgCluster.Spec.StorageSpec,
            }},
        },
    }
}

// controllers/postgrescluster_controller.go

const (
	primaryRoleLabel = "postgres-role"
	primaryValue     = "primary"
	replicaValue     = "replica"
)

func (r *PostgresClusterReconciler) managePrimaryAndReplicas(ctx context.Context, pgCluster *dbv1.PostgresCluster, sts *appsv1.StatefulSet) (string, error) {
	log := log.FromContext(ctx)

	// 1. List all pods for this cluster
	podList := &corev1.PodList{}
	labels := map[string]string{"app": "postgres", "cluster": pgCluster.Name}
	if err := r.List(ctx, podList, client.InNamespace(pgCluster.Namespace), client.MatchingLabels(labels)); err != nil {
		return "", err
	}

	var primaryPod *corev1.Pod
	var healthyReplicas []*corev1.Pod

	// 2. Categorize pods
	for i := range podList.Items {
		pod := &podList.Items[i]
		role, exists := pod.Labels[primaryRoleLabel]

		if exists && role == primaryValue {
			primaryPod = pod
		} else if isPodReady(pod) {
			healthyReplicas = append(healthyReplicas, pod)
		}
	}

	// 3. The State Machine
	if primaryPod == nil {
		// Case 1: No primary. Elect a new one.
		log.Info("No primary found. Electing a new one.")
		if len(podList.Items) > 0 {
			// For simplicity, elect the pod with the lowest ordinal index.
			// A production system would check WAL positions.
			sort.Slice(podList.Items, func(i, j int) bool {
				return podList.Items[i].Name < podList.Items[j].Name
			})
			newPrimary := &podList.Items[0]
			if err := r.promotePod(ctx, newPrimary); err != nil {
				return "", err
			}
			return newPrimary.Name, nil
		} else {
			log.Info("No pods available to elect a primary.")
			return "", nil // No pods yet, will reconcile again
		}
	} else if !isPodReady(primaryPod) {
		// Case 3: Primary exists but is unhealthy. Initiate failover.
		log.Info("Primary is not ready. Initiating failover.", "primaryPod", primaryPod.Name)
		if len(healthyReplicas) > 0 {
			// Sort replicas to have a deterministic promotion candidate.
			sort.Slice(healthyReplicas, func(i, j int) bool {
				return healthyReplicas[i].Name < healthyReplicas[j].Name
			})
			newPrimary := healthyReplicas[0]
			log.Info("Promoting new primary", "newPrimaryPod", newPrimary.Name)

			// First, demote the old primary to prevent split-brain
			if err := r.demotePod(ctx, primaryPod); err != nil {
				log.Error(err, "Failed to demote old primary, retrying...")
				return "", err
			}

			if err := r.promotePod(ctx, newPrimary); err != nil {
				return "", err
			}
			return newPrimary.Name, nil
		} else {
			log.Error(nil, "Primary is unhealthy, but no healthy replicas are available for failover.")
			return primaryPod.Name, fmt.Errorf("failover impossible: no healthy replicas")
		}
	} else {
		// Case 2: Primary is healthy.
		log.V(1).Info("Primary is healthy.", "primaryPod", primaryPod.Name)
		// In a real operator, we would now ensure all other pods are configured
		// to replicate from this primary. This would involve exec-ing into pods
		// or updating a ConfigMap they read from.
		return primaryPod.Name, nil
	}
}

func (r *PostgresClusterReconciler) promotePod(ctx context.Context, pod *corev1.Pod) error {
	// Update pod label to mark as primary
	pod.Labels[primaryRoleLabel] = primaryValue
	if err := r.Update(ctx, pod); err != nil {
		return err
	}

	// In a real operator, you would `exec` into the pod and run `pg_ctl promote`
	// or a similar command for your replication manager.
	log.FromContext(ctx).Info("Simulating promotion command for pod", "podName", pod.Name)

	return nil
}

func (r *PostgresClusterReconciler) demotePod(ctx context.Context, pod *corev1.Pod) error {
	pod.Labels[primaryRoleLabel] = replicaValue
	return r.Update(ctx, pod)
}

func isPodReady(pod *corev1.Pod) bool {
	for _, cond := range pod.Status.Conditions {
		if cond.Type == corev1.PodReady && cond.Status == corev1.ConditionTrue {
			return true
		}
	}
	return false
}

    // controllers/postgrescluster_controller.go - inside Reconcile

    const postgresClusterFinalizer = "db.example.com/finalizer"

    // ... after fetching the pgCluster ...

    // Handle deletion
    if pgCluster.ObjectMeta.DeletionTimestamp.IsZero() {
        // The object is not being deleted, so if it does not have our finalizer, 
        // then lets add it and update the object.
        if !controllerutil.ContainsFinalizer(&pgCluster, postgresClusterFinalizer) {
            controllerutil.AddFinalizer(&pgCluster, postgresClusterFinalizer)
            if err := r.Update(ctx, &pgCluster); err != nil {
                return ctrl.Result{}, err
            }
        }
    } else {
        // The object is being deleted
        if controllerutil.ContainsFinalizer(&pgCluster, postgresClusterFinalizer) {
            // Our finalizer is present, so lets handle any external dependency
            if err := r.cleanupPostgresCluster(ctx, &pgCluster); err != nil {
                // if fail to delete the external dependency here, return with error
                // so that it can be retried
                return ctrl.Result{}, err
            }

            // remove our finalizer from the list and update it.
            controllerutil.RemoveFinalizer(&pgCluster, postgresClusterFinalizer)
            if err := r.Update(ctx, &pgCluster); err != nil {
                return ctrl.Result{}, err
            }
        }

        // Stop reconciliation as the item is being deleted
        return ctrl.Result{}, nil
    }

    // ... continue with normal reconciliation ...

    func (r *PostgresClusterReconciler) cleanupPostgresCluster(ctx context.Context, pgCluster *dbv1.PostgresCluster) error {
        log := log.FromContext(ctx)
        log.Info("Performing cleanup for PostgresCluster")

        // This is where you would implement your cleanup logic.
        // For example, triggering a final backup to S3.
        // Or, if you wanted to delete the PVCs, you would list them and delete them here.

        // Example: Triggering a backup job (pseudo-code)
        // backupJob := createBackupJob(pgCluster)
        // if err := r.Create(ctx, backupJob); err != nil {
        //     return err
        // }
        // log.Info("Started final backup job.")
        // You would then need logic to wait for the job to complete.

        log.Info("Cleanup finished. Finalizer can be removed.")
        return nil
    }