From 5f65ea831cf726f53372ae8d5c277f4ccf077514 Mon Sep 17 00:00:00 2001
From: Lucas Manuel Rodriguez <lucas@fleetdm.com>
Date: Mon, 10 Jun 2024 16:49:27 -0300
Subject: [PATCH] Disable AI features on non-new installations upgrading to
 4.51.X (#19482)

#19365

Assuming we release this fix in 4.51.0:
- Migration from a version without the feature (< 4.50.0) to 4.51.0:
Should disable (set `ai_features_disabled=true`).
- Migration from a version with the feature (>= 4.50.X < 4.51.0) to
4.51.0: Should keep `ai_features_disabled` as-is.
- New installation of Fleet: Should come with AI features enabled
(`ai_features_disabled=false`).

From
https://github.com/fleetdm/fleet/issues/19365#issuecomment-2145825363.

- [X] Changes file added for user-visible changes in `changes/`,
`orbit/changes/` or `ee/fleetd-chrome/changes`.
See [Changes
files](https://fleetdm.com/docs/contributing/committing-changes#changes-files)
for more information.
- [X] Input data is properly validated, `SELECT *` is avoided, SQL
injection is prevented (using placeholders for values in statements)
- [X] Manual QA for all new/changed functionality
---
 changes/19365-disable-ai-migration            |  1 +
 .../20240430111727_CleanupQueryResults.go     | 64 +++++++++++++++++++
 2 files changed, 65 insertions(+)
 create mode 100644 changes/19365-disable-ai-migration

diff --git a/changes/19365-disable-ai-migration b/changes/19365-disable-ai-migration
new file mode 100644
index 0000000000..41e096d242
--- /dev/null
+++ b/changes/19365-disable-ai-migration
@@ -0,0 +1 @@
+* Disabled AI features on non-new installations upgrading from < 4.50.X to >= 4.51.X.
diff --git a/server/datastore/mysql/migrations/tables/20240430111727_CleanupQueryResults.go b/server/datastore/mysql/migrations/tables/20240430111727_CleanupQueryResults.go
index 063a150617..c0be097472 100644
--- a/server/datastore/mysql/migrations/tables/20240430111727_CleanupQueryResults.go
+++ b/server/datastore/mysql/migrations/tables/20240430111727_CleanupQueryResults.go
@@ -2,7 +2,10 @@ package tables
 
 import (
 	"database/sql"
+	"encoding/json"
 	"fmt"
+
+	"github.com/pkg/errors"
 )
 
 func init() {
@@ -22,6 +25,67 @@ func Up_20240430111727(tx *sql.Tx) error {
 	if err != nil {
 		return fmt.Errorf("failed to delete query_results %w", err)
 	}
+
+	//
+	// The following "fix" was introduced after this migration was released in 4.50.0.
+	// We are adding it here to disable the AI features (set ai_features_disabled=true)
+	// for non-new installations that are upgrading from < 4.50.0 using the new version
+	// of this migration to be released in 4.51.X.
+	//
+	if err := fixDisableAIForNonNewInstallation(tx); err != nil {
+		return fmt.Errorf("failed to update ai_features_disabled: %w", err)
+	}
+
+	return nil
+}
+
+func fixDisableAIForNonNewInstallation(tx *sql.Tx) error {
+	var usersCount int
+	row := tx.QueryRow(`SELECT COUNT(*) FROM users;`)
+	if err := row.Scan(&usersCount); err != nil {
+		return fmt.Errorf("select count users: %w", err)
+	}
+	if usersCount == 0 {
+		return nil
+	}
+
+	//
+	// At least a "setup" user was configured,
+	// thus we assume this is not a new installation.
+	//
+
+	var raw json.RawMessage
+	row = tx.QueryRow(`SELECT json_value FROM app_config_json LIMIT 1;`)
+	if err := row.Scan(&raw); err != nil {
+		if errors.Is(err, sql.ErrNoRows) {
+			return nil
+		}
+		return fmt.Errorf("select app_config_json: %w", err)
+	}
+
+	var config map[string]interface{}
+	if err := json.Unmarshal(raw, &config); err != nil {
+		return fmt.Errorf("unmarshal appconfig: %w", err)
+	}
+
+	ss, ok := config["server_settings"]
+	if !ok {
+		return errors.New("missing server_settings")
+	}
+	serverSettings, ok := ss.(map[string]interface{})
+	if !ok {
+		return fmt.Errorf("invalid type for server_settings: %T", ss)
+	}
+	serverSettings["ai_features_disabled"] = true
+
+	b, err := json.Marshal(config)
+	if err != nil {
+		return fmt.Errorf("marshal updated appconfig: %w", err)
+	}
+	if _, err := tx.Exec(`UPDATE app_config_json SET json_value = ? WHERE id = 1;`, b); err != nil {
+		return fmt.Errorf("update app_config_json: %w", err)
+	}
+
 	return nil
 }