Add kibana upgrade assistant

Author: ericywl

functionaltests/internal/kbclient/client.go

@@ -27,6 +27,7 @@ import (
 	"fmt"
 	"io"
 	"net/http"
+	"slices"
 	"time"
 )
 
@@ -235,3 +236,108 @@ func (c *Client) EnableIntegrationsServer(ctx context.Context) error {
 
 	return nil
 }
+
+func (c *Client) ResolveMigrationDeprecations(ctx context.Context) error {
+	ctx, cancel := context.WithTimeout(ctx, 30*time.Second)
+	defer cancel()
+
+	deprecations, err := c.QueryCriticalESDeprecations(ctx)
+	if err != nil {
+		return fmt.Errorf("failed to query critical deprecations: %w", err)
+	}
+
+	for _, deprecation := range deprecations {
+		switch deprecation.Type {
+		case "index_settings":
+			if err = c.markIndexAsReadOnly(ctx, deprecation.Name); err != nil {
+				return fmt.Errorf("failed to mark index as read-only: %w", err)
+			}
+		case "data_streams":
+			if err = c.markDataStreamAsReadOnly(
+				ctx,
+				deprecation.Name,
+				deprecation.CorrectiveAction.Metadata.IndicesRequiringUpgrade,
+			); err != nil {
+				return fmt.Errorf("failed to mark data stream as read-only: %w", err)
+			}
+		default:
+			return fmt.Errorf("unknown deprecation type: %s", deprecation.Type)
+		}
+	}
+
+	return nil
+}
+
+type MigrationDeprecation struct {
+	Name             string `json:"index"`
+	Type             string `json:"type"`
+	IsCritical       bool   `json:"isCritical"`
+	CorrectiveAction struct {
+		Type     string `json:"type"`
+		Metadata struct {
+			IndicesRequiringUpgrade []string `json:"indicesRequiringUpgrade,omitempty"`
+		}
+	} `json:"correctiveAction"`
+}
+
+type esDeprecationsResponse struct {
+	MigrationDeprecations []MigrationDeprecation `json:"migrationsDeprecations"`
+}
+
+// QueryCriticalESDeprecations retrieves the critical deprecation warnings for Elasticsearch.
+// It is essentially equivalent to `GET _migration/deprecations`, but through Kibana Upgrade
+// Assistant API.
+func (c *Client) QueryCriticalESDeprecations(ctx context.Context) ([]MigrationDeprecation, error) {
+	path := "/api/upgrade_assistant/es_deprecations"
+	b, err := c.sendRequest(ctx, http.MethodGet, path, nil, nil)
+	if err != nil {
+		return nil, err
+	}
+
+	var esDeprecationsResp esDeprecationsResponse
+	if err = json.Unmarshal(b, &esDeprecationsResp); err != nil {
+		return nil, fmt.Errorf("cannot unmarshal response body: %w", err)
+	}
+
+	// Remove all non-critical deprecation info.
+	return slices.DeleteFunc(
+		esDeprecationsResp.MigrationDeprecations,
+		func(dep MigrationDeprecation) bool {
+			return !dep.IsCritical
+		},
+	), nil
+}
+
+type upgradeAssistUpdateIndexRequest struct {
+	Operations []string `json:"operations"`
+}
+
+// markIndexAsReadOnly updates the index to read-only through the Upgrade Assistant API:
+// https://www.elastic.co/guide/en/kibana/current/upgrade-assistant.html.
+func (c *Client) markIndexAsReadOnly(ctx context.Context, index string) error {
+	path := fmt.Sprintf("/api/upgrade_assistant/update_index/%s", index)
+	req := upgradeAssistUpdateIndexRequest{
+		Operations: []string{"blockWrite", "unfreeze"},
+	}
+
+	_, err := c.sendRequest(ctx, http.MethodPost, path, req, nil)
+	return err
+}
+
+type upgradeAssistMigrateDSRequest struct {
+	Indices []string `json:"indices"`
+}
+
+// markDataStreamAsReadOnly marks the backing indices of the data stream as read-only
+// through the Upgrade Assistant API:
+// https://www.elastic.co/guide/en/kibana/current/upgrade-assistant.html.
+func (c *Client) markDataStreamAsReadOnly(ctx context.Context, dataStream string, indices []string) error {
+	// Data stream
+	path := fmt.Sprintf("/api/upgrade_assistant/migrate_data_stream/%s/readonly", dataStream)
+	req := upgradeAssistMigrateDSRequest{
+		Indices: indices,
+	}
+
+	_, err := c.sendRequest(ctx, http.MethodPost, path, req, nil)
+	return err
+}

functionaltests/internal/kbclient/client_test.go

@@ -108,7 +108,7 @@ func getHttpClient(t *testing.T) (*recorder.Recorder, *http.Client) {
 	return rec, hc
 }
 
-func TestClient_GetPackagePolicyByID(t *testing.T) {
+func newRecordedTestClient(t *testing.T) *kbclient.Client {
 	kibanaURL := os.Getenv("KIBANA_URL")
 	username := os.Getenv("KIBANA_USERNAME")
 	password := os.Getenv("KIBANA_PASSWORD")
@@ -117,23 +117,23 @@ func TestClient_GetPackagePolicyByID(t *testing.T) {
 	_, httpc := getHttpClient(t)
 	kbc.Client = *httpc
 
+	return kbc
+}
+
+func TestClient_GetPackagePolicyByID(t *testing.T) {
+	kbc := newRecordedTestClient(t)
+
 	policy, err := kbc.GetPackagePolicyByID(context.Background(), "elastic-cloud-apm")
 	require.NoError(t, err)
 	assert.Equal(t, "Elastic APM", policy.Name)
 }
 
 func TestClient_UpdatePackagePolicyByID(t *testing.T) {
-	kibanaURL := os.Getenv("KIBANA_URL")
-	username := os.Getenv("KIBANA_USERNAME")
-	password := os.Getenv("KIBANA_PASSWORD")
-	kbc, err := kbclient.New(kibanaURL, username, password)
-	require.NoError(t, err)
-	_, httpc := getHttpClient(t)
-	kbc.Client = *httpc
+	kbc := newRecordedTestClient(t)
 
 	ctx := context.Background()
 	policyID := "elastic-cloud-apm"
-	err = kbc.UpdatePackagePolicyByID(ctx, policyID,
+	err := kbc.UpdatePackagePolicyByID(ctx, policyID,
 		kbclient.PackagePolicy{
 			Name:        "Elastic APM",
 			Description: "Hello World",
@@ -149,3 +149,25 @@ func TestClient_UpdatePackagePolicyByID(t *testing.T) {
 	require.NoError(t, err)
 	assert.Equal(t, "Hello World", policy.Description)
 }
+
+func TestClient_ResolveMigrationDeprecations(t *testing.T) {
+	kbc := newRecordedTestClient(t)
+
+	ctx := context.Background()
+	// Check that there are some critical deprecation warnings.
+	deprecations, err := kbc.QueryCriticalESDeprecations(ctx)
+	require.NoError(t, err)
+	require.Greater(t, len(deprecations), 0)
+	for _, deprecation := range deprecations {
+		require.True(t, deprecation.IsCritical)
+	}
+
+	// Resolve them.
+	err = kbc.ResolveMigrationDeprecations(ctx)
+	require.NoError(t, err)
+
+	// Check that there are no more.
+	deprecations, err = kbc.QueryCriticalESDeprecations(ctx)
+	require.NoError(t, err)
+	assert.Len(t, deprecations, 0)
+}