From 6980022f85a177f11cc34a7221be53f8190556e9 Mon Sep 17 00:00:00 2001
From: Julia Bardi <90178898+juliaElastic@users.noreply.github.com>
Date: Mon, 12 Feb 2024 10:46:33 +0100
Subject: [PATCH] fix process upgrade details when null (#3264)

* fix process upgrade details when null

* fix test

* added changelog

* updated schema to let null be converted to nil automatically

* refactor to reduce complexity
---
 .../fragments/1707481003-fix-upgraded-at.yaml | 32 +++++++++++++++
 internal/pkg/api/handleCheckin.go             | 40 +++++++++++--------
 internal/pkg/api/handleCheckin_test.go        |  2 +-
 internal/pkg/model/schema.go                  |  6 ++-
 model/schema.json                             |  2 +-
 5 files changed, 63 insertions(+), 19 deletions(-)
 create mode 100644 changelog/fragments/1707481003-fix-upgraded-at.yaml

diff --git a/changelog/fragments/1707481003-fix-upgraded-at.yaml b/changelog/fragments/1707481003-fix-upgraded-at.yaml
new file mode 100644
index 000000000..5ea3963b9
--- /dev/null
+++ b/changelog/fragments/1707481003-fix-upgraded-at.yaml
@@ -0,0 +1,32 @@
+# Kind can be one of:
+# - breaking-change: a change to previously-documented behavior
+# - deprecation: functionality that is being removed in a later release
+# - bug-fix: fixes a problem in a previous version
+# - enhancement: extends functionality but does not break or fix existing behavior
+# - feature: new functionality
+# - known-issue: problems that we are aware of in a given version
+# - security: impacts on the security of a product or a user’s deployment.
+# - upgrade: important information for someone upgrading from a prior version
+# - other: does not fit into any of the other categories
+kind: bug-fix
+
+# Change summary; a 80ish characters long description of the change.
+summary: Fixed a bug where agents were stuck in non-upgradeable state after upgrade.
+
+# Long description; in case the summary is not enough to describe the change
+# this field accommodate a description without length limits.
+# NOTE: This field will be rendered only for breaking-change and known-issue kinds at the moment.
+# description: 
+
+# Affected component; usually one of "elastic-agent", "fleet-server", "filebeat", "metricbeat", "auditbeat", "all", etc.
+component:
+
+# PR URL; optional; the PR number that added the changeset.
+# If not present is automatically filled by the tooling finding the PR where this changelog fragment has been added.
+# NOTE: the tooling supports backports, so it's able to fill the original PR number instead of the backport PR number.
+# Please provide it if you are adding a fragment for a different PR.
+pr: 3264
+
+# Issue URL; optional; the GitHub issue related to this changeset (either closes or is part of).
+# If not present is automatically filled by the tooling with the issue linked to the PR number.
+issue: 3263
diff --git a/internal/pkg/api/handleCheckin.go b/internal/pkg/api/handleCheckin.go
index 328f5924a..f9b591e77 100644
--- a/internal/pkg/api/handleCheckin.go
+++ b/internal/pkg/api/handleCheckin.go
@@ -390,25 +390,11 @@ func (ct *CheckinT) ProcessRequest(zlog zerolog.Logger, w http.ResponseWriter, r
 // otherwise the details are validated; action_id is checked and upgrade_details.metadata is validated based on upgrade_details.state and the agent doc is updated.
 func (ct *CheckinT) processUpgradeDetails(ctx context.Context, agent *model.Agent, details *UpgradeDetails) error {
 	if details == nil {
-		// nop if there are no checkin details, and the agent has no details
-		if len(agent.UpgradeDetails) == 0 {
-			return nil
-		}
-		span, ctx := apm.StartSpan(ctx, "Mark update complete", "update")
-		span.Context.SetLabel("agent_id", agent.Agent.ID)
-		defer span.End()
-		// if the checkin had no details, but agent has details treat like a successful upgrade
-		doc := bulk.UpdateFields{
-			dl.FieldUpgradeDetails:   nil,
-			dl.FieldUpgradeStartedAt: nil,
-			dl.FieldUpgradeStatus:    nil,
-			dl.FieldUpgradedAt:       time.Now().UTC().Format(time.RFC3339),
-		}
-		body, err := doc.Marshal()
+		err := ct.markUpgradeComplete(ctx, agent)
 		if err != nil {
 			return err
 		}
-		return ct.bulker.Update(ctx, dl.FleetAgents, agent.Id, body, bulk.WithRefresh(), bulk.WithRetryOnConflict(3))
+		return nil
 	}
 	// update docs with in progress details
 
@@ -506,6 +492,28 @@ func (ct *CheckinT) processUpgradeDetails(ctx context.Context, agent *model.Agen
 	return ct.bulker.Update(ctx, dl.FleetAgents, agent.Id, body, bulk.WithRefresh(), bulk.WithRetryOnConflict(3))
 }
 
+func (ct *CheckinT) markUpgradeComplete(ctx context.Context, agent *model.Agent) error {
+	// nop if there are no checkin details, and the agent has no details
+	if agent.UpgradeDetails == nil {
+		return nil
+	}
+	span, ctx := apm.StartSpan(ctx, "Mark update complete", "update")
+	span.Context.SetLabel("agent_id", agent.Agent.ID)
+	defer span.End()
+	// if the checkin had no details, but agent has details treat like a successful upgrade
+	doc := bulk.UpdateFields{
+		dl.FieldUpgradeDetails:   nil,
+		dl.FieldUpgradeStartedAt: nil,
+		dl.FieldUpgradeStatus:    nil,
+		dl.FieldUpgradedAt:       time.Now().UTC().Format(time.RFC3339),
+	}
+	body, err := doc.Marshal()
+	if err != nil {
+		return err
+	}
+	return ct.bulker.Update(ctx, dl.FleetAgents, agent.Id, body, bulk.WithRefresh(), bulk.WithRetryOnConflict(3))
+}
+
 func (ct *CheckinT) writeResponse(zlog zerolog.Logger, w http.ResponseWriter, r *http.Request, agent *model.Agent, resp CheckinResponse) error {
 	ctx := r.Context()
 	var links []apm.SpanLink
diff --git a/internal/pkg/api/handleCheckin_test.go b/internal/pkg/api/handleCheckin_test.go
index 42f9969a2..7ca1714ee 100644
--- a/internal/pkg/api/handleCheckin_test.go
+++ b/internal/pkg/api/handleCheckin_test.go
@@ -304,7 +304,7 @@ func TestProcessUpgradeDetails(t *testing.T) {
 		err: nil,
 	}, {
 		name:    "agent has details checkin details are nil",
-		agent:   &model.Agent{ESDocument: esd, Agent: &model.AgentMetadata{ID: "test-agent"}, UpgradeDetails: json.RawMessage(`{"action_id":"test"}`)},
+		agent:   &model.Agent{ESDocument: esd, Agent: &model.AgentMetadata{ID: "test-agent"}, UpgradeDetails: &model.UpgradeDetails{}},
 		details: nil,
 		bulk: func() *ftesting.MockBulk {
 			mBulk := ftesting.NewMockBulk()
diff --git a/internal/pkg/model/schema.go b/internal/pkg/model/schema.go
index b8be8577c..9309ee86a 100644
--- a/internal/pkg/model/schema.go
+++ b/internal/pkg/model/schema.go
@@ -196,7 +196,7 @@ type Agent struct {
 	UpdatedAt string `json:"updated_at,omitempty"`
 
 	// Additional upgrade status details.
-	UpgradeDetails json.RawMessage `json:"upgrade_details,omitempty"`
+	UpgradeDetails *UpgradeDetails `json:"upgrade_details,omitempty"`
 
 	// Date/time the Elastic Agent started the current upgrade
 	UpgradeStartedAt string `json:"upgrade_started_at,omitempty"`
@@ -497,3 +497,7 @@ type ToRetireAPIKeyIdsItems struct {
 	// Date/time the API key was retired
 	RetiredAt string `json:"retired_at,omitempty"`
 }
+
+// UpgradeDetails Additional upgrade status details.
+type UpgradeDetails struct {
+}
diff --git a/model/schema.json b/model/schema.json
index 75699d1aa..90614fe6e 100644
--- a/model/schema.json
+++ b/model/schema.json
@@ -628,7 +628,7 @@
         },
         "upgrade_details": {
             "description": "Additional upgrade status details.",
-            "format": "raw"
+            "type": "object"
         }
       },
       "required": [