From 39f7ddc2a36935d1bc9a82286a429966e11c773e Mon Sep 17 00:00:00 2001
From: Shubhra Deshpande <shubhrajayant+github.com>
Date: Tue, 15 Oct 2024 22:54:51 -0400
Subject: [PATCH] statefulset: updating controller to read from a statefulset

---
 controllers/quay/features.go                | 70 ++++++++++++++++-----
 controllers/quay/quayregistry_controller.go |  2 +-
 2 files changed, 56 insertions(+), 16 deletions(-)

diff --git a/controllers/quay/features.go b/controllers/quay/features.go
index b8511b6f9..3bcf777a2 100644
--- a/controllers/quay/features.go
+++ b/controllers/quay/features.go
@@ -410,8 +410,8 @@ func (r *QuayRegistryReconciler) checkNeedsPostgresUpgradeForComponent(
 	ctx context.Context, qctx *quaycontext.QuayRegistryContext, quay *v1.QuayRegistry, component v1.ComponentKind,
 ) error {
 	componentInfo := map[v1.ComponentKind]struct {
-		deploymentSuffix string
-		upgradeField     *bool
+		resourceSuffix string
+		upgradeField   *bool
 	}{
 		v1.ComponentClairPostgres: {"clair-postgres", &qctx.NeedsClairPgUpgrade},
 		v1.ComponentPostgres:      {"quay-database", &qctx.NeedsPgUpgrade},
@@ -422,23 +422,63 @@ func (r *QuayRegistryReconciler) checkNeedsPostgresUpgradeForComponent(
 		return fmt.Errorf("invalid component kind: %s", component)
 	}
 
-	deploymentName := fmt.Sprintf("%s-%s", quay.GetName(), info.deploymentSuffix)
+	resourceName := fmt.Sprintf("%s-%s", quay.GetName(), info.resourceSuffix)
 	r.Log.Info(fmt.Sprintf("getting %s version", component))
 
-	postgresDeployment := &appsv1.Deployment{}
-	if err := r.Client.Get(
-		ctx,
-		types.NamespacedName{
-			Name:      deploymentName,
-			Namespace: quay.GetNamespace(),
-		},
-		postgresDeployment,
-	); err != nil {
-		r.Log.Info(fmt.Sprintf("%s deployment not found, skipping", component))
-		return nil
+	var deployedImageName string
+
+	if component == v1.ComponentClairPostgres {
+		statefulSet := &appsv1.StatefulSet{}
+		err := r.Client.Get(
+			ctx,
+			types.NamespacedName{
+				Name:      resourceName,
+				Namespace: quay.GetNamespace(),
+			},
+			statefulSet,
+		)
+		if err != nil {
+			if !errors.IsNotFound(err) {
+				return err
+			}
+			// NOTE: Check for Deployment to support migration from Deployment to StatefulSet.
+			// This ensures compatibility with both old and new setups during the upgrade process.
+			deployment := &appsv1.Deployment{}
+			err = r.Client.Get(
+				ctx,
+				types.NamespacedName{
+					Name:      resourceName,
+					Namespace: quay.GetNamespace(),
+				},
+				deployment,
+			)
+			if err != nil {
+				if errors.IsNotFound(err) {
+					r.Log.Info(fmt.Sprintf("%s statefulset and deployment not found, skipping", component))
+					return nil
+				}
+				return err
+			}
+			deployedImageName = deployment.Spec.Template.Spec.Containers[0].Image
+		} else {
+			deployedImageName = statefulSet.Spec.Template.Spec.Containers[0].Image
+		}
+	} else {
+		deployment := &appsv1.Deployment{}
+		if err := r.Client.Get(
+			ctx,
+			types.NamespacedName{
+				Name:      resourceName,
+				Namespace: quay.GetNamespace(),
+			},
+			deployment,
+		); err != nil {
+			r.Log.Info(fmt.Sprintf("%s deployment not found, skipping", component))
+			return nil
+		}
+		deployedImageName = deployment.Spec.Template.Spec.Containers[0].Image
 	}
 
-	deployedImageName := postgresDeployment.Spec.Template.Spec.Containers[0].Image
 	r.Log.Info(fmt.Sprintf("%s deployment found", component), "image", deployedImageName)
 
 	expectedImage, err := kustomize.ComponentImageFor(component)
diff --git a/controllers/quay/quayregistry_controller.go b/controllers/quay/quayregistry_controller.go
index 75fadd2a8..03a09a4d6 100644
--- a/controllers/quay/quayregistry_controller.go
+++ b/controllers/quay/quayregistry_controller.go
@@ -548,7 +548,7 @@ func (r *QuayRegistryReconciler) Reconcile(ctx context.Context, req ctrl.Request
 				v1.ConditionTypeRolloutBlocked,
 				metav1.ConditionTrue,
 				v1.ConditionReasonPostgresUpgradeFailed,
-				fmt.Sprintf("error checking for pg upgrade: %s", err),
+				fmt.Sprintf("error checking for clair pg upgrade: %s", err),
 			)
 		}
 	}