fix postgres migration issue with 0.24 (#2367)

* fix postgres migration issue with 0.24

Fixes #2351

Signed-off-by: Kristoffer Dalby <[email protected]>

* add postgres migration test for 2351

Signed-off-by: Kristoffer Dalby <[email protected]>

* update changelog

Signed-off-by: Kristoffer Dalby <[email protected]>

---------

Signed-off-by: Kristoffer Dalby <[email protected]>

Author: kradalby

CHANGELOG.md

@@ -4,6 +4,8 @@
 
 ### Changes
 
+- Fix migration issue with user table for PostgreSQL
+  [#2367](https://github.com/juanfont/headscale/pull/2367)
 - Relax username validation to allow emails
   [#2364](https://github.com/juanfont/headscale/pull/2364)
 - Remove invalid routes and add stronger constraints for routes to avoid API panic

hscontrol/db/db.go

@@ -478,6 +478,38 @@ func NewHeadscaleDatabase(
 				// populate the user with more interesting information.
 				ID: "202407191627",
 				Migrate: func(tx *gorm.DB) error {
+					// Fix an issue where the automigration in GORM expected a constraint to
+					// exists that didnt, and add the one it wanted.
+					// Fixes https://github.com/juanfont/headscale/issues/2351
+					if cfg.Type == types.DatabasePostgres {
+						err := tx.Exec(`
+BEGIN;
+DO $$
+BEGIN
+    IF NOT EXISTS (
+        SELECT 1 FROM pg_constraint
+        WHERE conname = 'uni_users_name'
+    ) THEN
+        ALTER TABLE users ADD CONSTRAINT uni_users_name UNIQUE (name);
+    END IF;
+END $$;
+
+DO $$
+BEGIN
+    IF EXISTS (
+        SELECT 1 FROM pg_constraint
+        WHERE conname = 'users_name_key'
+    ) THEN
+        ALTER TABLE users DROP CONSTRAINT users_name_key;
+    END IF;
+END $$;
+COMMIT;
+`).Error
+						if err != nil {
+							return fmt.Errorf("failed to rename constraint: %w", err)
+						}
+					}
+
 					err := tx.AutoMigrate(&types.User{})
 					if err != nil {
 						return err

hscontrol/db/db_test.go

@@ -6,6 +6,7 @@ import (
 	"io"
 	"net/netip"
 	"os"
+	"os/exec"
 	"path/filepath"
 	"slices"
 	"sort"
@@ -23,7 +24,10 @@ import (
 	"zgo.at/zcache/v2"
 )
 
-func TestMigrations(t *testing.T) {
+// TestMigrationsSQLite is the main function for testing migrations,
+// we focus on SQLite correctness as it is the main database used in headscale.
+// All migrations that are worth testing should be added here.
+func TestMigrationsSQLite(t *testing.T) {
 	ipp := func(p string) netip.Prefix {
 		return netip.MustParsePrefix(p)
 	}
@@ -375,3 +379,58 @@ func TestConstraints(t *testing.T) {
 		})
 	}
 }
+
+func TestMigrationsPostgres(t *testing.T) {
+	tests := []struct {
+		name     string
+		dbPath   string
+		wantFunc func(*testing.T, *HSDatabase)
+	}{
+		{
+			name:   "user-idx-breaking",
+			dbPath: "testdata/pre-24-postgresdb.pssql.dump",
+			wantFunc: func(t *testing.T, h *HSDatabase) {
+				users, err := Read(h.DB, func(rx *gorm.DB) ([]types.User, error) {
+					return ListUsers(rx)
+				})
+				require.NoError(t, err)
+
+				for _, user := range users {
+					assert.NotEmpty(t, user.Name)
+					assert.Empty(t, user.ProfilePicURL)
+					assert.Empty(t, user.Email)
+				}
+			},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			u := newPostgresDBForTest(t)
+
+			pgRestorePath, err := exec.LookPath("pg_restore")
+			if err != nil {
+				t.Fatal("pg_restore not found in PATH. Please install it and ensure it is accessible.")
+			}
+
+			// Construct the pg_restore command
+			cmd := exec.Command(pgRestorePath, "--verbose", "--if-exists", "--clean", "--no-owner", "--dbname", u.String(), tt.dbPath)
+
+			// Set the output streams
+			cmd.Stdout = os.Stdout
+			cmd.Stderr = os.Stderr
+
+			// Execute the command
+			err = cmd.Run()
+			if err != nil {
+				t.Fatalf("failed to restore postgres database: %s", err)
+			}
+
+			db = newHeadscaleDBFromPostgresURL(t, u)
+
+			if tt.wantFunc != nil {
+				tt.wantFunc(t, db)
+			}
+		})
+	}
+}

hscontrol/db/suite_test.go

@@ -78,13 +78,11 @@ func newSQLiteTestDB() (*HSDatabase, error) {
 func newPostgresTestDB(t *testing.T) *HSDatabase {
 	t.Helper()
 
-	var err error
-	tmpDir, err = os.MkdirTemp("", "headscale-db-test-*")
-	if err != nil {
-		t.Fatal(err)
-	}
+	return newHeadscaleDBFromPostgresURL(t, newPostgresDBForTest(t))
+}
 
-	log.Printf("database path: %s", tmpDir+"/headscale_test.db")
+func newPostgresDBForTest(t *testing.T) *url.URL {
+	t.Helper()
 
 	ctx := context.Background()
 	srv, err := postgrestest.Start(ctx)
@@ -100,10 +98,16 @@ func newPostgresTestDB(t *testing.T) *HSDatabase {
 	t.Logf("created local postgres: %s", u)
 	pu, _ := url.Parse(u)
 
+	return pu
+}
+
+func newHeadscaleDBFromPostgresURL(t *testing.T, pu *url.URL) *HSDatabase {
+	t.Helper()
+
 	pass, _ := pu.User.Password()
 	port, _ := strconv.Atoi(pu.Port())
 
-	db, err = NewHeadscaleDatabase(
+	db, err := NewHeadscaleDatabase(
 		types.DatabaseConfig{
 			Type: types.DatabasePostgres,
 			Postgres: types.PostgresConfig{