From de5d74bbb244baf93761eb375af8e7f48e531013 Mon Sep 17 00:00:00 2001
From: Anderson Queiroz <anderson.queiroz@elastic.co>
Date: Thu, 8 Sep 2022 21:49:19 +0200
Subject: [PATCH] Fix v8.5.0 migration painless script (#1839)

* fix v8.5.0 migration painless script
---
 internal/pkg/dl/migration.go                  |  4 +-
 internal/pkg/dl/migration_integration_test.go | 95 ++++++++++++++++++-
 2 files changed, 96 insertions(+), 3 deletions(-)

diff --git a/internal/pkg/dl/migration.go b/internal/pkg/dl/migration.go
index 5b1722824..4107dfd38 100644
--- a/internal/pkg/dl/migration.go
+++ b/internal/pkg/dl/migration.go
@@ -229,7 +229,9 @@ ctx._source['` + fieldOutputs + `']['default'].to_retire_api_key_ids=new ArrayLi
 
 // copy 'default_api_key_history' to new 'outputs' field
 ctx._source['` + fieldOutputs + `']['default'].type="elasticsearch";
-ctx._source['` + fieldOutputs + `']['default'].to_retire_api_key_ids=ctx._source.default_api_key_history;
+if (ctx._source.default_api_key_history != null && ctx._source.default_api_key_history.length > 0) {
+    ctx._source['` + fieldOutputs + `']['default'].to_retire_api_key_ids=ctx._source.default_api_key_history;
+}
 
 Map map = new HashMap();
 map.put("retired_at", params.` + fieldRetiredAt + `);
diff --git a/internal/pkg/dl/migration_integration_test.go b/internal/pkg/dl/migration_integration_test.go
index 916bf478b..183502b94 100644
--- a/internal/pkg/dl/migration_integration_test.go
+++ b/internal/pkg/dl/migration_integration_test.go
@@ -132,7 +132,7 @@ func TestPolicyCoordinatorIdx(t *testing.T) {
 	}
 }
 
-func TestMigrateOutputs(t *testing.T) {
+func TestMigrateOutputs_withDefaultAPIKeyHistory(t *testing.T) {
 	now, err := time.Parse(time.RFC3339, nowStr)
 	require.NoError(t, err, "could not parse time "+nowStr)
 	timeNow = func() time.Time {
@@ -153,7 +153,7 @@ func TestMigrateOutputs(t *testing.T) {
 	assert.Equal(t, len(agentIDs), migratedAgents)
 
 	for i, id := range agentIDs {
-		wantOutputType := "elasticsearch"
+		wantOutputType := "elasticsearch" //nolint:goconst // test cases have some duplication
 
 		got, err := FindAgent(
 			context.Background(), bulker, QueryAgentByID, FieldID, id, WithIndexName(index))
@@ -202,3 +202,94 @@ func TestMigrateOutputs(t *testing.T) {
 		assert.Nil(t, got.DefaultAPIKeyHistory)
 	}
 }
+
+func TestMigrateOutputs_nil_DefaultAPIKeyHistory(t *testing.T) {
+	wantOutputType := "elasticsearch" //nolint:goconst // test cases have some duplication
+
+	now, err := time.Parse(time.RFC3339, nowStr)
+	require.NoError(t, err, "could not parse time "+nowStr)
+	timeNow = func() time.Time {
+		return now
+	}
+
+	index, bulker := ftesting.SetupCleanIndex(context.Background(), t, FleetAgents)
+	apiKey := bulk.APIKey{
+		ID:  "testAgent_",
+		Key: "testAgent_key_",
+	}
+
+	i := 0
+	outputAPIKey := bulk.APIKey{
+		ID:  fmt.Sprint(apiKey.ID, i),
+		Key: fmt.Sprint(apiKey.Key, i),
+	}
+
+	agentID := uuid.Must(uuid.NewV4()).String()
+	policyID := uuid.Must(uuid.NewV4()).String()
+
+	agentModel := model.Agent{
+		PolicyID:                    policyID,
+		Active:                      true,
+		LastCheckin:                 nowStr,
+		LastCheckinStatus:           "",
+		UpdatedAt:                   nowStr,
+		EnrolledAt:                  nowStr,
+		DefaultAPIKeyID:             outputAPIKey.ID,
+		DefaultAPIKey:               outputAPIKey.Agent(),
+		PolicyOutputPermissionsHash: fmt.Sprint("a_output_permission_SHA_", i),
+	}
+
+	body, err := json.Marshal(agentModel)
+	require.NoError(t, err)
+
+	_, err = bulker.Create(
+		context.Background(), index, agentID, body, bulk.WithRefresh())
+	require.NoError(t, err)
+
+	migratedAgents, err := migrate(context.Background(), bulker, migrateAgentOutputs)
+	require.NoError(t, err)
+
+	got, err := FindAgent(
+		context.Background(), bulker, QueryAgentByID, FieldID, agentID, WithIndexName(index))
+	require.NoError(t, err, "failed to find agent ID %q", agentID) // we want to continue even if a single agent fails
+
+	assert.Equal(t, 1, migratedAgents)
+
+	// Assert new fields
+	require.Len(t, got.Outputs, 1)
+	// Default API key is empty to force fleet-server to regenerate them.
+	assert.Empty(t, got.Outputs["default"].APIKey)
+	assert.Empty(t, got.Outputs["default"].APIKeyID)
+	assert.Equal(t, wantOutputType, got.Outputs["default"].Type)
+	assert.Equal(t,
+		fmt.Sprint("a_output_permission_SHA_", i),
+		got.Outputs["default"].PermissionsHash)
+
+	// Assert ToRetireAPIKeyIds contains the expected values, regardless of the order.
+	if assert.Len(t, got.Outputs["default"].ToRetireAPIKeyIds, 1) {
+		assert.Equal(t,
+			model.ToRetireAPIKeyIdsItems{ID: outputAPIKey.ID, RetiredAt: nowStr},
+			got.Outputs["default"].ToRetireAPIKeyIds[0])
+	}
+
+	// Assert deprecated fields
+	assert.Empty(t, got.DefaultAPIKey)
+	assert.Empty(t, got.DefaultAPIKey)
+	assert.Empty(t, got.PolicyOutputPermissionsHash)
+	assert.Nil(t, got.DefaultAPIKeyHistory)
+}
+
+func TestMigrateOutputs_no_agent_document(t *testing.T) {
+	now, err := time.Parse(time.RFC3339, nowStr)
+	require.NoError(t, err, "could not parse time "+nowStr)
+	timeNow = func() time.Time {
+		return now
+	}
+
+	_, bulker := ftesting.SetupCleanIndex(context.Background(), t, FleetAgents)
+
+	migratedAgents, err := migrate(context.Background(), bulker, migrateAgentOutputs)
+	require.NoError(t, err)
+
+	assert.Equal(t, 0, migratedAgents)
+}