From b2b6217265f55bca50d20d04aedd5182ca92e434 Mon Sep 17 00:00:00 2001
From: Patrick Pacher <patrick.pacher@gmail.com>
Date: Fri, 20 Oct 2023 11:55:18 +0200
Subject: [PATCH] Migrate profile IDs in history database when merging profiles

---
 netquery/database.go   |  9 +++++++++
 netquery/module_api.go | 15 +++++++++++++++
 profile/merge.go       |  2 ++
 profile/module.go      |  2 ++
 4 files changed, 28 insertions(+)

diff --git a/netquery/database.go b/netquery/database.go
index 7adae82b..e3287345 100644
--- a/netquery/database.go
+++ b/netquery/database.go
@@ -432,6 +432,15 @@ func (db *Database) RemoveHistoryForProfile(ctx context.Context, profileID strin
 	}))
 }
 
+// MigrateProfileID migrates the given profile IDs in the history database.
+// This needs to be done when profiles are deleted and replaced by a different profile.
+func (db *Database) MigrateProfileID(ctx context.Context, from string, to string) error {
+	return db.ExecuteWrite(ctx, "UPDATE history.connections SET profile = :to WHERE profile = :from", orm.WithNamedArgs(map[string]any{
+		":from": from,
+		":to":   to,
+	}))
+}
+
 // dumpTo is a simple helper method that dumps all rows stored in the SQLite database
 // as JSON to w.
 // Any error aborts dumping rows and is returned.
diff --git a/netquery/module_api.go b/netquery/module_api.go
index 63a3d07f..b4e56b02 100644
--- a/netquery/module_api.go
+++ b/netquery/module_api.go
@@ -272,6 +272,21 @@ func (m *module) start() error {
 		}
 	}
 
+	// Migrate profile IDs in history database when profiles are migrated/merged.
+	if err := m.RegisterEventHook(
+		"profiles",
+		"profile migrated",
+		"migrate profile IDs in history database",
+		func(ctx context.Context, data interface{}) error {
+			if profileIDs, ok := data.([]string); ok && len(profileIDs) == 2 {
+				return m.Store.MigrateProfileID(ctx, profileIDs[0], profileIDs[1])
+			}
+			return nil
+		},
+	); err != nil {
+		return err
+	}
+
 	return nil
 }
 
diff --git a/profile/merge.go b/profile/merge.go
index 2c0d5cfc..29e274dc 100644
--- a/profile/merge.go
+++ b/profile/merge.go
@@ -77,10 +77,12 @@ func MergeProfiles(name string, primary *Profile, secondaries ...*Profile) (newP
 	if err := primary.delete(); err != nil {
 		return nil, fmt.Errorf("failed to delete primary profile %s: %w", primary.ScopedID(), err)
 	}
+	module.TriggerEvent(MigratedEvent, []string{primary.ScopedID(), newProfile.ScopedID()})
 	for _, sp := range secondaries {
 		if err := sp.delete(); err != nil {
 			return nil, fmt.Errorf("failed to delete secondary profile %s: %w", sp.ScopedID(), err)
 		}
+		module.TriggerEvent(MigratedEvent, []string{sp.ScopedID(), newProfile.ScopedID()})
 	}
 
 	return newProfile, nil
diff --git a/profile/module.go b/profile/module.go
index 280579d0..a462cd7e 100644
--- a/profile/module.go
+++ b/profile/module.go
@@ -22,12 +22,14 @@ var (
 const (
 	ConfigChangeEvent = "profile config change"
 	DeletedEvent      = "profile deleted"
+	MigratedEvent     = "profile migrated"
 )
 
 func init() {
 	module = modules.Register("profiles", prep, start, stop, "base", "updates")
 	module.RegisterEvent(ConfigChangeEvent, true)
 	module.RegisterEvent(DeletedEvent, true)
+	module.RegisterEvent(MigratedEvent, true)
 }
 
 func prep() error {