Merge remote-tracking branch 'origin/main' into HEAD

# Conflicts:
#	shared/management/proto/proxy_service.pb.go

management/server/store/sql_store.go

@@ -58,6 +58,7 @@ const (
 	keyQueryCondition              = "key = ?"
 	mysqlKeyQueryCondition         = "`key` = ?"
 	accountAndIDQueryCondition     = "account_id = ? and id = ?"
+	accountAndAnyIDQueryCondition  = "account_id = ? and (id = ? or public_id = ?)"
 	accountAndPeerIDQueryCondition = "account_id = ? and peer_id = ?"
 	accountAndIDsQueryCondition    = "account_id = ? AND id IN ?"
 	accountIDCondition             = "account_id = ?"
@@ -3154,7 +3155,12 @@ func NewMysqlStore(ctx context.Context, dsn string, metrics telemetry.AppMetrics
 		return nil, err
 	}
 
-	return NewSqlStore(ctx, db, types.MysqlStoreEngine, metrics, skipMigration)
+	store, err := NewSqlStore(ctx, db, types.MysqlStoreEngine, metrics, skipMigration)
+	if err != nil {
+		closeGormDB(db)
+		return nil, err
+	}
+	return store, nil
 }
 
 func getGormConfig() *gorm.Config {
@@ -3166,9 +3172,9 @@ func getGormConfig() *gorm.Config {
 
 // newPostgresStore initializes a new Postgres store.
 func newPostgresStore(ctx context.Context, metrics telemetry.AppMetrics, skipMigration bool) (Store, error) {
-	dsn, ok := lookupDSNEnv(postgresDsnEnv, postgresDsnEnvLegacy)
+	dsn, ok := lookupDSNEnv(PostgresDsnEnv, PostgresDsnEnvLegacy)
 	if !ok {
-		return nil, fmt.Errorf("%s is not set", postgresDsnEnv)
+		return nil, fmt.Errorf("%s is not set", PostgresDsnEnv)
 	}
 	return NewPostgresqlStore(ctx, dsn, metrics, skipMigration)
 }
@@ -3213,23 +3219,20 @@ func NewSqliteStoreFromFileStore(ctx context.Context, fileStore *FileStore, data
 
 // NewPostgresqlStoreFromSqlStore restores a store from SqlStore and stores Postgres DB.
 func NewPostgresqlStoreFromSqlStore(ctx context.Context, sqliteStore *SqlStore, dsn string, metrics telemetry.AppMetrics) (*SqlStore, error) {
-	store, err := NewPostgresqlStoreForTests(ctx, dsn, metrics, false)
+	return newPostgresqlStoreFromSqlStore(ctx, sqliteStore, dsn, metrics, false)
+}
+
+func newPostgresqlStoreFromSqlStore(ctx context.Context, sqliteStore *SqlStore, dsn string, metrics telemetry.AppMetrics, skipMigration bool) (*SqlStore, error) {
+	store, err := NewPostgresqlStoreForTests(ctx, dsn, metrics, skipMigration)
 	if err != nil {
 		return nil, err
 	}
 
-	err = store.SaveInstallationID(ctx, sqliteStore.GetInstallationID())
-	if err != nil {
+	if err := seedFromSqliteStore(ctx, store, sqliteStore); err != nil {
+		closeStore(ctx, store)
 		return nil, err
 	}
 
-	for _, account := range sqliteStore.GetAllAccounts(ctx) {
-		err := store.SaveAccount(ctx, account)
-		if err != nil {
-			return nil, err
-		}
-	}
-
 	return store, nil
 }
 
@@ -3241,11 +3244,14 @@ func NewPostgresqlStoreForTests(ctx context.Context, dsn string, metrics telemet
 	}
 	pool, err := connectToPgDbForTests(context.Background(), dsn)
 	if err != nil {
+		closeGormDB(db)
 		return nil, err
 	}
 	store, err := NewSqlStore(ctx, db, types.PostgresStoreEngine, metrics, skipMigration)
 	if err != nil {
+		// Release the sessions, or the caller cannot drop the database.
 		pool.Close()
+		closeGormDB(db)
 		return nil, err
 	}
 	store.pool = pool
@@ -3279,22 +3285,42 @@ func connectToPgDbForTests(ctx context.Context, dsn string) (*pgxpool.Pool, erro
 
 // NewMysqlStoreFromSqlStore restores a store from SqlStore and stores MySQL DB.
 func NewMysqlStoreFromSqlStore(ctx context.Context, sqliteStore *SqlStore, dsn string, metrics telemetry.AppMetrics) (*SqlStore, error) {
-	store, err := NewMysqlStore(ctx, dsn, metrics, false)
-	if err != nil {
-		return nil, err
-	}
+	return newMysqlStoreFromSqlStore(ctx, sqliteStore, dsn, metrics, false)
+}
 
-	err = store.SaveInstallationID(ctx, sqliteStore.GetInstallationID())
-	if err != nil {
-		return nil, err
+// seedFromSqliteStore copies the installation ID and the accounts of the
+// sqlite seed store into a freshly created engine store.
+func seedFromSqliteStore(ctx context.Context, store, sqliteStore *SqlStore) error {
+	if err := store.SaveInstallationID(ctx, sqliteStore.GetInstallationID()); err != nil {
+		return err
 	}
-
 	for _, account := range sqliteStore.GetAllAccounts(ctx) {
-		err := store.SaveAccount(ctx, account)
-		if err != nil {
-			return nil, err
+		if err := store.SaveAccount(ctx, account); err != nil {
+			return err
 		}
 	}
+	return nil
+}
+
+// closeStore releases a store that is not handed to the caller, so a failed
+// seed does not leak its connection and pool.
+func closeStore(ctx context.Context, store *SqlStore) {
+	store.Close(ctx)
+	if store.pool != nil {
+		store.pool.Close()
+	}
+}
+
+func newMysqlStoreFromSqlStore(ctx context.Context, sqliteStore *SqlStore, dsn string, metrics telemetry.AppMetrics, skipMigration bool) (*SqlStore, error) {
+	store, err := NewMysqlStore(ctx, dsn, metrics, skipMigration)
+	if err != nil {
+		return nil, err
+	}
+
+	if err := seedFromSqliteStore(ctx, store, sqliteStore); err != nil {
+		closeStore(ctx, store)
+		return nil, err
+	}
 
 	return store, nil
 }
@@ -3473,7 +3499,7 @@ func (s *SqlStore) GetPeerGroups(ctx context.Context, lockStrength LockingStreng
 	var groups []*types.Group
 	query := tx.
 		Joins("JOIN group_peers ON group_peers.group_id = groups.id").
-		Where("group_peers.peer_id = ?", peerId).
+		Where("groups.account_id = ? AND group_peers.peer_id = ?", accountId, peerId).
 		Preload(clause.Associations).
 		Find(&groups)
 
@@ -4038,6 +4064,30 @@ func (s *SqlStore) GetPolicyByID(ctx context.Context, lockStrength LockingStreng
 	return policy, nil
 }
 
+// GetPolicyByIDOrPublicID retrieves a policy by either its ID or its PublicID. Peers report
+// whichever of the two the network map they were served carries, so callers resolving a
+// peer-reported reference cannot know upfront which namespace it belongs to.
+func (s *SqlStore) GetPolicyByIDOrPublicID(ctx context.Context, lockStrength LockingStrength, accountID, policyID string) (*types.Policy, error) {
+	tx := s.db
+	if lockStrength != LockingStrengthNone {
+		tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
+	}
+
+	var policy *types.Policy
+
+	result := tx.Preload(clause.Associations).
+		Take(&policy, accountAndAnyIDQueryCondition, accountID, policyID, policyID)
+	if err := result.Error; err != nil {
+		if errors.Is(err, gorm.ErrRecordNotFound) {
+			return nil, status.NewPolicyNotFoundError(policyID)
+		}
+		log.WithContext(ctx).Errorf("failed to get policy from store: %s", err)
+		return nil, status.Errorf(status.Internal, "failed to get policy from store")
+	}
+
+	return policy, nil
+}
+
 func (s *SqlStore) CreatePolicy(ctx context.Context, policy *types.Policy) error {
 	result := s.db.Create(policy)
 	if result.Error != nil {
@@ -4223,6 +4273,27 @@ func (s *SqlStore) GetRouteByID(ctx context.Context, lockStrength LockingStrengt
 	return route, nil
 }
 
+// GetRouteByIDOrPublicID retrieves a route by either its ID or its PublicID. See
+// GetPolicyByIDOrPublicID for why peer-reported references need both.
+func (s *SqlStore) GetRouteByIDOrPublicID(ctx context.Context, lockStrength LockingStrength, accountID string, routeID string) (*route.Route, error) {
+	tx := s.db
+	if lockStrength != LockingStrengthNone {
+		tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
+	}
+
+	var route *route.Route
+	result := tx.Take(&route, accountAndAnyIDQueryCondition, accountID, routeID, routeID)
+	if err := result.Error; err != nil {
+		if errors.Is(err, gorm.ErrRecordNotFound) {
+			return nil, status.NewRouteNotFoundError(routeID)
+		}
+		log.WithContext(ctx).Errorf("failed to get route from the store: %s", err)
+		return nil, status.Errorf(status.Internal, "failed to get route from store")
+	}
+
+	return route, nil
+}
+
 // SaveRoute saves a route to the database.
 func (s *SqlStore) SaveRoute(ctx context.Context, route *route.Route) error {
 	result := s.db.Save(route)
@@ -4617,6 +4688,28 @@ func (s *SqlStore) GetNetworkResourceByID(ctx context.Context, lockStrength Lock
 	return netResources, nil
 }
 
+// GetNetworkResourceByIDOrPublicID retrieves a network resource by either its ID or its
+// PublicID. See GetPolicyByIDOrPublicID for why peer-reported references need both.
+func (s *SqlStore) GetNetworkResourceByIDOrPublicID(ctx context.Context, lockStrength LockingStrength, accountID, resourceID string) (*resourceTypes.NetworkResource, error) {
+	tx := s.db
+	if lockStrength != LockingStrengthNone {
+		tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
+	}
+
+	var netResources *resourceTypes.NetworkResource
+	result := tx.
+		Take(&netResources, accountAndAnyIDQueryCondition, accountID, resourceID, resourceID)
+	if result.Error != nil {
+		if errors.Is(result.Error, gorm.ErrRecordNotFound) {
+			return nil, status.NewNetworkResourceNotFoundError(resourceID)
+		}
+		log.WithContext(ctx).Errorf("failed to get network resource from store: %v", result.Error)
+		return nil, status.Errorf(status.Internal, "failed to get network resource from store")
+	}
+
+	return netResources, nil
+}
+
 func (s *SqlStore) GetNetworkResourceByName(ctx context.Context, lockStrength LockingStrength, accountID, resourceName string) (*resourceTypes.NetworkResource, error) {
 	tx := s.db
 	if lockStrength != LockingStrengthNone {
@@ -5053,7 +5146,7 @@ func (s *SqlStore) GetPeersByGroupIDs(ctx context.Context, accountID string, gro
 		Select("DISTINCT peer_id").
 		Where("account_id = ? AND group_id IN ?", accountID, groupIDs)
 
-	result := s.db.Where("id IN (?)", peerIDsSubquery).Find(&peers)
+	result := s.db.Where("account_id = ? AND id IN (?)", accountID, peerIDsSubquery).Find(&peers)
 	if result.Error != nil {
 		log.WithContext(ctx).Errorf("failed to get peers by group IDs: %s", result.Error)
 		return nil, status.Errorf(status.Internal, "failed to get peers by group IDs")
@@ -5686,6 +5779,23 @@ func (s *SqlStore) ListCustomDomains(ctx context.Context, accountID string) ([]*
 	return domains, nil
 }
 
+// GetCustomDomainByName returns the custom domain row holding the given name,
+// regardless of which account owns it.
+func (s *SqlStore) GetCustomDomainByName(ctx context.Context, domainName string) (*domain.Domain, error) {
+	customDomain := &domain.Domain{}
+	result := s.db.Take(customDomain, "domain = ?", domainName)
+	if result.Error != nil {
+		if errors.Is(result.Error, gorm.ErrRecordNotFound) {
+			return nil, status.Errorf(status.NotFound, "custom domain %s not found", domainName)
+		}
+
+		log.WithContext(ctx).Errorf("failed to get custom domain by name from store: %v", result.Error)
+		return nil, status.Errorf(status.Internal, "failed to get custom domain from store")
+	}
+
+	return customDomain, nil
+}
+
 func (s *SqlStore) CreateCustomDomain(ctx context.Context, accountID string, domainName string, targetCluster string, validated bool) (*domain.Domain, error) {
 	newDomain := &domain.Domain{
 		ID:            xid.New().String(), // Generate our own ID because gorm doesn't always configure the database to handle this for us.
@@ -5695,8 +5805,24 @@ func (s *SqlStore) CreateCustomDomain(ctx context.Context, accountID string, dom
 		Type:          domain.TypeCustom,
 		Validated:     validated,
 	}
+	if !validated {
+		expiresAt := time.Now().UTC().Add(domain.ValidationTTL)
+		newDomain.ValidationExpiresAt = &expiresAt
+	}
 	result := s.db.Create(newDomain)
 	if result.Error != nil {
+		// The unique index is the last guard when two requests clear the
+		// manager's availability check at the same time. The one that loses the
+		// insert is a conflict, not an internal failure.
+		var count int64
+		if err := s.db.Model(&domain.Domain{}).Where("domain = ?", domainName).Count(&count).Error; err == nil && count > 0 {
+			// The insert error is logged even on this path: the name being taken
+			// is what the caller has to act on, but if the insert also failed for
+			// an unrelated reason the operator still needs to see it.
+			log.WithContext(ctx).Warnf("create reverse proxy custom domain %s rejected, name already registered: %v", domainName, result.Error)
+			return nil, status.Errorf(status.AlreadyExists, "domain %s is already registered", domainName)
+		}
+
 		log.WithContext(ctx).Errorf("failed to create reverse proxy custom domain to store: %v", result.Error)
 		return nil, status.Errorf(status.Internal, "failed to create reverse proxy custom domain to store")
 	}
@@ -5704,12 +5830,21 @@ func (s *SqlStore) CreateCustomDomain(ctx context.Context, accountID string, dom
 	return newDomain, nil
 }
 
+// UpdateCustomDomain completes validation only while the original registration is pending.
 func (s *SqlStore) UpdateCustomDomain(ctx context.Context, accountID string, d *domain.Domain) (*domain.Domain, error) {
-	d.AccountID = accountID
-	result := s.db.Select("*").Save(d)
+	if !d.Validated {
+		return nil, status.Errorf(status.InvalidArgument, "custom domain update must complete validation")
+	}
+	result := s.db.WithContext(ctx).Model(&domain.Domain{}).
+		Where(accountAndIDQueryCondition, accountID, d.ID).
+		Where("domain = ? AND target_cluster = ?", d.Domain, d.TargetCluster).
+		Where("validated = ? AND validation_expires_at > ?", false, time.Now().UTC()).
+		Update("validated", true)
 	if result.Error != nil {
-		log.WithContext(ctx).Errorf("failed to update reverse proxy custom domain to store: %v", result.Error)
-		return nil, status.Errorf(status.Internal, "failed to update reverse proxy custom domain to store")
+		return nil, fmt.Errorf("validate custom domain in store: %w", result.Error)
+	}
+	if result.RowsAffected == 0 {
+		return nil, status.Errorf(status.PreconditionFailed, "custom domain registration is no longer pending validation")
 	}
 
 	return d, nil
@@ -6404,6 +6539,25 @@ func (s *SqlStore) IsClusterAddressConflicting(ctx context.Context, clusterAddre
 	return count > 0, nil
 }
 
+// HasForeignAccountProxyAtHost reports whether a proxy owned by a different
+// account declares this host. Shared proxies (account_id IS NULL) are not
+// foreign: a shared cluster is what most accounts pin their agent network
+// gateway to. The match folds case because proxies declare their address as
+// the operator spelled it while the caller's host is normalised; that costs a
+// scan of the proxies table, taken once per account when its gateway is
+// bootstrapped, not on the per-connect path IsClusterAddressConflicting serves.
+func (s *SqlStore) HasForeignAccountProxyAtHost(ctx context.Context, host, accountID string) (bool, error) {
+	var count int64
+	result := s.db.
+		Model(&proxy.Proxy{}).
+		Where("LOWER(cluster_address) = LOWER(?) AND account_id IS NOT NULL AND account_id != ?", host, accountID).
+		Count(&count)
+	if result.Error != nil {
+		return false, status.Errorf(status.Internal, "check proxy host ownership: %v", result.Error)
+	}
+	return count > 0, nil
+}
+
 func (s *SqlStore) DeleteAccountCluster(ctx context.Context, clusterAddress, accountID string) error {
 	result := s.db.
 		Where("cluster_address = ? AND account_id = ?", clusterAddress, accountID).

management/server/store/sql_store_agentnetwork.go

@@ -315,6 +315,36 @@ func (s *SqlStore) GetAllAgentNetworkSettings(ctx context.Context, lockStrength
 	return settings, nil
 }
 
+// HasGatewayClusterPinnedByOtherAccount reports whether another account has a
+// labeled agent network gateway pinned beneath host, making host its cluster.
+// A self-addressed endpoint on the very same hostname is not counted: that
+// collision is the domain unique index's to refuse, as a conflict. Case-folded,
+// since a settings row written before hostnames were normalised may carry
+// capitals; one row per account, so the scan is cheap.
+func (s *SqlStore) HasGatewayClusterPinnedByOtherAccount(ctx context.Context, host, accountID string) (bool, error) {
+	return s.countGatewayRowsByOtherAccount(ctx, "LOWER(proxy_address) = LOWER(?) AND LOWER(domain) <> LOWER(proxy_address)", host, accountID)
+}
+
+// HasGatewayEndpointByOtherAccount reports whether host is another account's
+// agent network endpoint hostname (domain). Case-folded for the same reason as
+// HasGatewayClusterPinnedByOtherAccount.
+func (s *SqlStore) HasGatewayEndpointByOtherAccount(ctx context.Context, host, accountID string) (bool, error) {
+	return s.countGatewayRowsByOtherAccount(ctx, "LOWER(domain) = LOWER(?)", host, accountID)
+}
+
+func (s *SqlStore) countGatewayRowsByOtherAccount(ctx context.Context, predicate, host, accountID string) (bool, error) {
+	var count int64
+	result := s.db.
+		Model(&agentNetworkTypes.Settings{}).
+		Where(predicate+" AND account_id != ?", host, accountID).
+		Count(&count)
+	if result.Error != nil {
+		log.WithContext(ctx).Errorf("failed to check agent network gateway claims at host: %v", result.Error)
+		return false, status.Errorf(status.Internal, "check agent network gateway claims at host")
+	}
+	return count > 0, nil
+}
+
 // GetAgentNetworkSettingsByProxyAddress returns every Settings row whose
 // gateway is served by the proxy declaring the given cluster address. Used by
 // cluster-scoped synthesis to find the accounts a shared proxy serves.

management/server/store/sql_store_domain_expiration.go

@@ -0,0 +1,60 @@
+package store
+
+import (
+	"context"
+	"fmt"
+	"strings"
+	"time"
+
+	"gorm.io/gorm"
+
+	"github.com/netbirdio/netbird/management/internals/modules/reverseproxy/domain"
+	rpservice "github.com/netbirdio/netbird/management/internals/modules/reverseproxy/service"
+	"github.com/netbirdio/netbird/shared/management/status"
+)
+
+// GetExpiredCustomDomains lists pending registrations in stable batches across accounts.
+func (s *SqlStore) GetExpiredCustomDomains(ctx context.Context, now time.Time, afterID domain.ID, limit int) ([]*domain.Domain, error) {
+	var domains []*domain.Domain
+	result := s.db.WithContext(ctx).
+		Where("validated = ? AND validation_expires_at <= ? AND id > ?", false, now, string(afterID)).
+		Order("id").Limit(limit).Find(&domains)
+	if result.Error != nil {
+		return nil, fmt.Errorf("list expired custom domains: %w", result.Error)
+	}
+	return domains, nil
+}
+
+// DeleteExpiredCustomDomain deletes an expired registration only if no service uses its namespace.
+func (s *SqlStore) DeleteExpiredCustomDomain(ctx context.Context, d *domain.Domain, now time.Time) (bool, error) {
+	db := s.db.WithContext(ctx)
+	services := customDomainServices(db, d)
+	result := db.Where(accountAndIDQueryCondition, d.AccountID, d.ID).
+		Where("domain = ? AND validated = ? AND validation_expires_at <= ?", d.Domain, false, now).
+		Where("NOT EXISTS (?)", services.Select("1")).Delete(&domain.Domain{})
+	if result.Error != nil {
+		return false, fmt.Errorf("delete expired custom domain: %w", result.Error)
+	}
+	if result.RowsAffected > 0 {
+		return true, nil
+	}
+	var count int64
+	if err := customDomainServices(db, d).Count(&count).Error; err != nil {
+		return false, fmt.Errorf("check expired custom domain services: %w", err)
+	}
+	if count > 0 {
+		return false, status.Errorf(status.PreconditionFailed, "expired custom domain still has dependent services")
+	}
+	return false, nil
+}
+
+func customDomainServices(db *gorm.DB, d *domain.Domain) *gorm.DB {
+	name := strings.ToLower(strings.TrimSuffix(d.Domain, "."))
+	// Shared domain validation permits underscores, and older rows may contain
+	// other LIKE metacharacters.
+	escaped := strings.NewReplacer("!", "!!", "%", "!%", "_", "!_").Replace(name)
+	return db.Model(&rpservice.Service{}).Where(
+		"LOWER(domain) IN ? OR LOWER(domain) LIKE ? ESCAPE '!' OR LOWER(domain) LIKE ? ESCAPE '!'",
+		[]string{name, name + "."}, "%."+escaped, "%."+escaped+".",
+	)
+}

management/server/store/sql_store_domain_expiration_test.go

@@ -0,0 +1,78 @@
+package store
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"github.com/netbirdio/netbird/management/internals/modules/reverseproxy/domain"
+	rpservice "github.com/netbirdio/netbird/management/internals/modules/reverseproxy/service"
+)
+
+func TestDeleteExpiredCustomDomain_ServiceDependencies(t *testing.T) {
+	runTestForAllEngines(t, "", func(t *testing.T, store Store) {
+		ctx := context.Background()
+		now := time.Now().UTC()
+		db := store.(*SqlStore).db
+		require.NoError(t, store.SaveAccount(ctx, newAccountWithId(ctx, "owner", "admin", "")))
+		for _, tt := range []struct {
+			name        string
+			domainName  string
+			serviceHost string
+			protected   bool
+		}{
+			{"exact", "example.com", "example.com", true},
+			{"subdomain", "example.com", "deep.app.example.com", true},
+			{"case", "example.com", "APP.EXAMPLE.COM.", true},
+			{"suffix-boundary", "example.com", "notexample.com", false},
+			{"literal underscore", "a_b.example.com", "app.a_b.example.com", true},
+			{"underscore wildcard", "a_b.example.com", "app.axb.example.com", false},
+			{"legacy percent wildcard", "a%b.example.com", "app.axxb.example.com", false},
+			{"legacy escape character", "a!b.example.com", "app.ab.example.com", false},
+		} {
+			t.Run(tt.name, func(t *testing.T) {
+				d, err := store.CreateCustomDomain(ctx, "owner", tt.domainName, "cluster", false)
+				require.NoError(t, err)
+				require.NoError(t, db.Model(d).Update("validation_expires_at", now.Add(-time.Hour)).Error)
+				svc := &rpservice.Service{ID: "legacy", AccountID: "owner", Domain: tt.serviceHost}
+				require.NoError(t, store.CreateService(ctx, svc))
+				deleted, err := store.DeleteExpiredCustomDomain(ctx, d, now)
+				if tt.protected {
+					require.Error(t, err)
+					assert.False(t, deleted, "service namespaces must remain reserved")
+				} else {
+					require.NoError(t, err)
+					assert.True(t, deleted, "a hostname outside the namespace must not prevent cleanup")
+				}
+				require.NoError(t, db.Delete(svc).Error)
+				require.NoError(t, db.Delete(d).Error)
+			})
+		}
+	})
+}
+
+func TestDeleteExpiredCustomDomain_RechecksValidation(t *testing.T) {
+	runTestForAllEngines(t, "", func(t *testing.T, store Store) {
+		ctx := context.Background()
+		require.NoError(t, store.SaveAccount(ctx, newAccountWithId(ctx, "owner", "admin", "")))
+		d, err := store.CreateCustomDomain(ctx, "owner", "validated.example.com", "cluster", false)
+		require.NoError(t, err)
+		d, err = store.GetCustomDomain(ctx, "owner", d.ID)
+		require.NoError(t, err)
+		stale := d.Copy()
+		d.Validated = true
+		_, err = store.UpdateCustomDomain(ctx, "owner", d)
+		require.NoError(t, err)
+		deleted, err := store.DeleteExpiredCustomDomain(ctx, stale, time.Now().Add(domain.ValidationTTL))
+		require.NoError(t, err)
+		assert.False(t, deleted, "a stale cleanup candidate must not delete a validated registration")
+		stored, err := store.GetCustomDomain(ctx, "owner", d.ID)
+		require.NoError(t, err)
+		assert.True(t, stored.Validated, "the validated registration must remain usable")
+		require.NotNil(t, stored.ValidationExpiresAt)
+		assert.Equal(t, stale.ValidationExpiresAt, stored.ValidationExpiresAt, "validation must preserve the original deadline")
+	})
+}

management/server/store/sql_store_get_account_test.go

@@ -13,7 +13,6 @@ import (
 	"github.com/stretchr/testify/require"
 
 	nbdns "github.com/netbirdio/netbird/dns"
-	"github.com/netbirdio/netbird/management/server/integration_reference"
 	resourceTypes "github.com/netbirdio/netbird/management/server/networks/resources/types"
 	routerTypes "github.com/netbirdio/netbird/management/server/networks/routers/types"
 	networkTypes "github.com/netbirdio/netbird/management/server/networks/types"
@@ -21,6 +20,7 @@ import (
 	"github.com/netbirdio/netbird/management/server/posture"
 	"github.com/netbirdio/netbird/management/server/types"
 	"github.com/netbirdio/netbird/route"
+	"github.com/netbirdio/netbird/shared/management/integration_reference"
 )
 
 // TestGetAccount_LoadsCustomDomains verifies GetAccount populates account.Domains.
@@ -64,7 +64,7 @@ func assertGetAccountLoadsCustomDomains(t *testing.T, store Store) {
 
 	_, err := store.CreateCustomDomain(ctx, accountID, "example.com", "eu.proxy.netbird.io", true)
 	require.NoError(t, err, "creating the first custom domain must succeed")
-	_, err = store.CreateCustomDomain(ctx, accountID, "apps.acme.io", "us.proxy.netbird.io", false)
+	pending, err := store.CreateCustomDomain(ctx, accountID, "apps.acme.io", "us.proxy.netbird.io", false)
 	require.NoError(t, err, "creating the second custom domain must succeed")
 
 	account, err := store.GetAccount(ctx, accountID)
@@ -75,6 +75,10 @@ func assertGetAccountLoadsCustomDomains(t *testing.T, store Store) {
 	for _, d := range account.Domains {
 		require.NotNil(t, d)
 		byDomain[d.Domain] = d.TargetCluster
+		if d.ID == pending.ID {
+			require.NotNil(t, d.ValidationExpiresAt)
+			assert.WithinDuration(t, *pending.ValidationExpiresAt, *d.ValidationExpiresAt, time.Millisecond, "both account loaders must preserve the validation deadline")
+		}
 	}
 	assert.Equal(t, "eu.proxy.netbird.io", byDomain["example.com"], "custom domain must carry its target cluster")
 	assert.Equal(t, "us.proxy.netbird.io", byDomain["apps.acme.io"], "custom domain must carry its target cluster")

management/server/store/sql_store_test.go

@@ -573,7 +573,7 @@ func TestSqlStore_SavePeer(t *testing.T) {
 
 		numOfFields, err := populateFields.PopulateAll(reflectedMetadata)
 		assert.NoError(t, err)
-		assert.Equal(t, 32, numOfFields)
+		assert.Equal(t, 33, numOfFields)
 
 		// save status of non-existing peer
 		peer := &nbpeer.Peer{
@@ -1972,6 +1972,32 @@ func TestSqlStore_GetPolicyByID(t *testing.T) {
 	}
 }
 
+func TestSqlStore_GetPolicyByIDOrPublicID(t *testing.T) {
+	store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/store.sql", t.TempDir())
+	t.Cleanup(cleanup)
+	require.NoError(t, err)
+
+	accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
+	policyID := "cs1tnh0hhcjnqoiuebf0"
+
+	policy, err := store.GetPolicyByID(context.Background(), LockingStrengthNone, accountID, policyID)
+	require.NoError(t, err)
+	require.NotEmpty(t, policy.PublicID)
+
+	for _, id := range []string{policyID, policy.PublicID} {
+		policy, err := store.GetPolicyByIDOrPublicID(context.Background(), LockingStrengthNone, accountID, id)
+		require.NoError(t, err)
+		require.Equal(t, policyID, policy.ID)
+	}
+
+	policy, err = store.GetPolicyByIDOrPublicID(context.Background(), LockingStrengthNone, accountID, "non-existing")
+	require.Error(t, err)
+	sErr, ok := status.FromError(err)
+	require.True(t, ok)
+	require.Equal(t, sErr.Type(), status.NotFound)
+	require.Nil(t, policy)
+}
+
 func TestSqlStore_CreatePolicy(t *testing.T) {
 	store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/store.sql", t.TempDir())
 	t.Cleanup(cleanup)
@@ -2631,6 +2657,32 @@ func TestSqlStore_GetNetworkResourceByID(t *testing.T) {
 	}
 }
 
+func TestSqlStore_GetNetworkResourceByIDOrPublicID(t *testing.T) {
+	store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/store.sql", t.TempDir())
+	t.Cleanup(cleanup)
+	require.NoError(t, err)
+
+	accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
+	netResourceID := "ctc4nci7qv9061u6ilfg"
+
+	netResource, err := store.GetNetworkResourceByID(context.Background(), LockingStrengthNone, accountID, netResourceID)
+	require.NoError(t, err)
+	require.NotEmpty(t, netResource.PublicID)
+
+	for _, id := range []string{netResourceID, netResource.PublicID} {
+		netResource, err := store.GetNetworkResourceByIDOrPublicID(context.Background(), LockingStrengthNone, accountID, id)
+		require.NoError(t, err)
+		require.Equal(t, netResourceID, netResource.ID)
+	}
+
+	netResource, err = store.GetNetworkResourceByIDOrPublicID(context.Background(), LockingStrengthNone, accountID, "non-existing")
+	require.Error(t, err)
+	sErr, ok := status.FromError(err)
+	require.True(t, ok)
+	require.Equal(t, sErr.Type(), status.NotFound)
+	require.Nil(t, netResource)
+}
+
 func TestSqlStore_SaveNetworkResource(t *testing.T) {
 	store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/store.sql", t.TempDir())
 	t.Cleanup(cleanup)
@@ -2844,6 +2896,14 @@ func TestSqlStore_GetPeerGroups(t *testing.T) {
 	groups, err = store.GetPeerGroups(context.Background(), LockingStrengthNone, accountID, peerID)
 	require.NoError(t, err)
 	assert.Len(t, groups, 2)
+
+	foreignPeerID := "foreign-peer"
+	err = store.AddPeerToGroup(context.Background(), accountID, foreignPeerID, "cfefqs706sqkneg59g4h")
+	require.NoError(t, err)
+
+	groups, err = store.GetPeerGroups(context.Background(), LockingStrengthNone, "other-account", foreignPeerID)
+	require.NoError(t, err)
+	assert.Empty(t, groups, "groups of another account must not be returned")
 }
 
 func TestSqlStore_GetAccountPeers(t *testing.T) {
@@ -3748,6 +3808,32 @@ func TestSqlStore_GetRouteByID(t *testing.T) {
 	}
 }
 
+func TestSqlStore_GetRouteByIDOrPublicID(t *testing.T) {
+	store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/extended-store.sql", t.TempDir())
+	t.Cleanup(cleanup)
+	require.NoError(t, err)
+
+	accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
+	routeID := "ct03t427qv97vmtmglog"
+
+	route, err := store.GetRouteByID(context.Background(), LockingStrengthNone, accountID, routeID)
+	require.NoError(t, err)
+	require.NotEmpty(t, route.PublicID)
+
+	for _, id := range []string{routeID, route.PublicID} {
+		route, err := store.GetRouteByIDOrPublicID(context.Background(), LockingStrengthNone, accountID, id)
+		require.NoError(t, err)
+		require.Equal(t, routeID, string(route.ID))
+	}
+
+	route, err = store.GetRouteByIDOrPublicID(context.Background(), LockingStrengthNone, accountID, "non-existing")
+	require.Error(t, err)
+	sErr, ok := status.FromError(err)
+	require.True(t, ok)
+	require.Equal(t, sErr.Type(), status.NotFound)
+	require.Nil(t, route)
+}
+
 func TestSqlStore_SaveRoute(t *testing.T) {
 	store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/extended-store.sql", t.TempDir())
 	t.Cleanup(cleanup)
@@ -4039,9 +4125,15 @@ func TestSqlStore_GetPeersByGroupIDs(t *testing.T) {
 			}
 			require.NoError(t, store.CreateGroups(ctx, accountID, groups))
 
+			otherAccount := newAccountWithId(ctx, "other-account", "other-user", "")
+			require.NoError(t, store.SaveAccount(ctx, otherAccount))
+			foreignPeer := &nbpeer.Peer{ID: "foreign-peer", AccountID: otherAccount.Id}
+			require.NoError(t, store.AddPeerToAccount(ctx, foreignPeer))
+
 			require.NoError(t, store.AddPeerToGroup(ctx, accountID, peer1, group1ID))
 			require.NoError(t, store.AddPeerToGroup(ctx, accountID, peer2, group1ID))
 			require.NoError(t, store.AddPeerToGroup(ctx, accountID, peer1, group2ID))
+			require.NoError(t, store.AddPeerToGroup(ctx, accountID, foreignPeer.ID, group1ID))
 
 			peers, err := store.GetPeersByGroupIDs(ctx, accountID, tt.groupIDs)
 			require.NoError(t, err)

management/server/store/store.go

@@ -1,9 +1,10 @@
 package store
 
-//go:generate go run github.com/golang/mock/mockgen -package store -destination=store_mock.go -source=./store.go -build_flags=-mod=mod
+//go:generate go tool mockgen -package store -destination=store_mock.go -source=./store.go -build_flags=-mod=mod
 
 import (
 	"context"
+	"database/sql"
 	"errors"
 	"fmt"
 	"net"
@@ -15,6 +16,7 @@ import (
 	"runtime"
 	"slices"
 	"strings"
+	"sync"
 	"time"
 
 	"github.com/google/uuid"
@@ -136,6 +138,7 @@ type Store interface {
 
 	GetAccountPolicies(ctx context.Context, lockStrength LockingStrength, accountID string) ([]*types.Policy, error)
 	GetPolicyByID(ctx context.Context, lockStrength LockingStrength, accountID, policyID string) (*types.Policy, error)
+	GetPolicyByIDOrPublicID(ctx context.Context, lockStrength LockingStrength, accountID, policyID string) (*types.Policy, error)
 	CreatePolicy(ctx context.Context, policy *types.Policy) error
 	SavePolicy(ctx context.Context, policy *types.Policy) error
 	DeletePolicy(ctx context.Context, accountID, policyID string) error
@@ -206,6 +209,7 @@ type Store interface {
 
 	GetAccountRoutes(ctx context.Context, lockStrength LockingStrength, accountID string) ([]*route.Route, error)
 	GetRouteByID(ctx context.Context, lockStrength LockingStrength, accountID, routeID string) (*route.Route, error)
+	GetRouteByIDOrPublicID(ctx context.Context, lockStrength LockingStrength, accountID, routeID string) (*route.Route, error)
 	SaveRoute(ctx context.Context, route *route.Route) error
 	DeleteRoute(ctx context.Context, accountID, routeID string) error
 
@@ -246,6 +250,7 @@ type Store interface {
 	GetNetworkResourcesByNetID(ctx context.Context, lockStrength LockingStrength, accountID, netID string) ([]*resourceTypes.NetworkResource, error)
 	GetNetworkResourcesByAccountID(ctx context.Context, lockStrength LockingStrength, accountID string) ([]*resourceTypes.NetworkResource, error)
 	GetNetworkResourceByID(ctx context.Context, lockStrength LockingStrength, accountID, resourceID string) (*resourceTypes.NetworkResource, error)
+	GetNetworkResourceByIDOrPublicID(ctx context.Context, lockStrength LockingStrength, accountID, resourceID string) (*resourceTypes.NetworkResource, error)
 	GetNetworkResourceByName(ctx context.Context, lockStrength LockingStrength, accountID, resourceName string) (*resourceTypes.NetworkResource, error)
 	SaveNetworkResource(ctx context.Context, resource *resourceTypes.NetworkResource) error
 	DeleteNetworkResource(ctx context.Context, accountID, resourceID string) error
@@ -302,8 +307,11 @@ type Store interface {
 	GetCustomDomain(ctx context.Context, accountID string, domainID string) (*domain.Domain, error)
 	ListFreeDomains(ctx context.Context, accountID string) ([]string, error)
 	ListCustomDomains(ctx context.Context, accountID string) ([]*domain.Domain, error)
+	GetCustomDomainByName(ctx context.Context, domainName string) (*domain.Domain, error)
 	CreateCustomDomain(ctx context.Context, accountID string, domainName string, targetCluster string, validated bool) (*domain.Domain, error)
 	UpdateCustomDomain(ctx context.Context, accountID string, d *domain.Domain) (*domain.Domain, error)
+	GetExpiredCustomDomains(ctx context.Context, now time.Time, afterID domain.ID, limit int) ([]*domain.Domain, error)
+	DeleteExpiredCustomDomain(ctx context.Context, d *domain.Domain, now time.Time) (bool, error)
 	DeleteCustomDomain(ctx context.Context, accountID string, domainID string) error
 
 	CreateAccessLog(ctx context.Context, log *accesslogs.AccessLogEntry) error
@@ -337,6 +345,9 @@ type Store interface {
 	CountProxiesByAccountID(ctx context.Context, accountID string) (int64, error)
 	IsClusterAddressConflicting(ctx context.Context, clusterAddress, accountID string) (bool, error)
 	HasActiveProxyAtClusterAddress(ctx context.Context, clusterAddress string) (bool, error)
+	HasForeignAccountProxyAtHost(ctx context.Context, host, accountID string) (bool, error)
+	HasGatewayClusterPinnedByOtherAccount(ctx context.Context, host, accountID string) (bool, error)
+	HasGatewayEndpointByOtherAccount(ctx context.Context, host, accountID string) (bool, error)
 	DeleteAccountCluster(ctx context.Context, clusterAddress, accountID string) error
 
 	GetCustomDomainsCounts(ctx context.Context) (total int64, validated int64, err error)
@@ -436,8 +447,8 @@ type AgentNetworkMetrics struct {
 }
 
 const (
-	postgresDsnEnv       = "NB_STORE_ENGINE_POSTGRES_DSN"
-	postgresDsnEnvLegacy = "NETBIRD_STORE_ENGINE_POSTGRES_DSN"
+	PostgresDsnEnv       = "NB_STORE_ENGINE_POSTGRES_DSN"
+	PostgresDsnEnvLegacy = "NETBIRD_STORE_ENGINE_POSTGRES_DSN"
 	mysqlDsnEnv          = "NB_STORE_ENGINE_MYSQL_DSN"
 	mysqlDsnEnvLegacy    = "NETBIRD_STORE_ENGINE_MYSQL_DSN"
 )
@@ -641,6 +652,9 @@ func migratePostAuto(ctx context.Context, db *gorm.DB) error {
 
 func getMigrationsPostAuto(ctx context.Context) []migrationFunc {
 	return []migrationFunc{
+		func(db *gorm.DB) error {
+			return migration.MigrateCustomDomainValidationExpiry(ctx, db)
+		},
 		func(db *gorm.DB) error {
 			return migration.CreateIndexIfNotExists[nbpeer.Peer](ctx, db, "idx_account_ip", "account_id", "ip")
 		},
@@ -726,6 +740,7 @@ func NewTestStoreFromSQL(ctx context.Context, filename string, dataDir string) (
 			time.Sleep(100 * time.Millisecond)
 		}
 	}
+	store.Close(ctx)
 	return nil, nil, fmt.Errorf("failed to create test store after %d attempts: %v", maxRetries, err)
 }
 
@@ -752,14 +767,15 @@ func addAllGroupToAccount(ctx context.Context, store Store) error {
 	return nil
 }
 
-func getSqlStoreEngine(ctx context.Context, store *SqlStore, kind types.Engine) (Store, func(), error) {
+func getSqlStoreEngine(ctx context.Context, sqliteStore *SqlStore, kind types.Engine) (Store, func(), error) {
+	store := sqliteStore
 	var cleanup func()
 	var err error
 	switch kind {
 	case types.PostgresStoreEngine:
-		store, cleanup, err = newReusedPostgresStore(ctx, store, kind)
+		store, cleanup, err = newReusedPostgresStore(ctx, sqliteStore, kind)
 	case types.MysqlStoreEngine:
-		store, cleanup, err = newReusedMysqlStore(ctx, store, kind)
+		store, cleanup, err = newReusedMysqlStore(ctx, sqliteStore, kind)
 	default:
 		cleanup = func() {
 			// sqlite doesn't need to be cleaned up
@@ -775,13 +791,18 @@ func getSqlStoreEngine(ctx context.Context, store *SqlStore, kind types.Engine)
 		if store.pool != nil {
 			store.pool.Close()
 		}
+		if store != sqliteStore {
+			// The sqlite store only seeded the engine under test; without this
+			// every test leaks its connection and the opener goroutines.
+			sqliteStore.Close(ctx)
+		}
 	}
 
 	return store, closeConnection, nil
 }
 
 func newReusedPostgresStore(ctx context.Context, store *SqlStore, kind types.Engine) (*SqlStore, func(), error) {
-	dsn, ok := lookupDSNEnv(postgresDsnEnv, postgresDsnEnvLegacy)
+	dsn, ok := lookupDSNEnv(PostgresDsnEnv, PostgresDsnEnvLegacy)
 	if !ok || dsn == "" {
 		var err error
 		_, dsn, err = testutil.CreatePostgresTestContainer()
@@ -791,7 +812,7 @@ func newReusedPostgresStore(ctx context.Context, store *SqlStore, kind types.Eng
 	}
 
 	if dsn == "" {
-		return nil, nil, fmt.Errorf("%s is not set", postgresDsnEnv)
+		return nil, nil, fmt.Errorf("%s is not set", PostgresDsnEnv)
 	}
 
 	db, err := openDBWithRetry(dsn, kind, 5)
@@ -799,19 +820,23 @@ func newReusedPostgresStore(ctx context.Context, store *SqlStore, kind types.Eng
 		return nil, nil, fmt.Errorf("failed to open postgres connection: %v", err)
 	}
 
-	dsn, cleanup, err := createRandomDB(dsn, db, kind)
-
-	sqlDB, _ := db.DB()
-	if sqlDB != nil {
-		sqlDB.Close()
+	template, err := postgresSchemaTemplate(ctx, dsn, db)
+	if err != nil {
+		closeGormDB(db)
+		return nil, nil, err
 	}
 
+	dsn, cleanup, err := createRandomDB(dsn, db, kind, template)
+
+	closeGormDB(db)
+
 	if err != nil {
 		return nil, nil, err
 	}
 
-	store, err = NewPostgresqlStoreFromSqlStore(ctx, store, dsn, nil)
+	store, err = newPostgresqlStoreFromSqlStore(ctx, store, dsn, nil, true)
 	if err != nil {
+		cleanup()
 		return nil, nil, err
 	}
 
@@ -844,7 +869,13 @@ func newReusedMysqlStore(ctx context.Context, store *SqlStore, kind types.Engine
 	sqlDB.SetMaxOpenConns(1)
 	sqlDB.SetMaxIdleConns(1)
 
-	dsn, cleanup, err := createRandomDB(dsn, db, kind)
+	tableDDL, err := mysqlSchemaTemplate(ctx, dsn, db)
+	if err != nil {
+		sqlDB.Close()
+		return nil, nil, err
+	}
+
+	dsn, cleanup, err := createRandomDB(dsn, db, kind, "")
 
 	sqlDB.Close()
 
@@ -852,14 +883,200 @@ func newReusedMysqlStore(ctx context.Context, store *SqlStore, kind types.Engine
 		return nil, nil, err
 	}
 
-	store, err = NewMysqlStoreFromSqlStore(ctx, store, dsn, nil)
+	if err := cloneMysqlSchema(ctx, dsn, tableDDL); err != nil {
+		cleanup()
+		return nil, nil, err
+	}
+
+	store, err = newMysqlStoreFromSqlStore(ctx, store, dsn, nil, true)
 	if err != nil {
+		cleanup()
 		return nil, nil, err
 	}
 
 	return store, cleanup, nil
 }
 
+// schemaTemplates remembers, per engine and server, a database that went
+// through the full migration once in this process. Every later test database
+// is cloned from it, so a test pays for CREATE DATABASE and a schema copy
+// instead of the 40-table AutoMigrate plus every pre and post migration, which
+// is what made each MySQL test store cost well over a second in CI.
+var (
+	schemaTemplatesMu sync.Mutex
+	schemaTemplates   = map[string]*schemaTemplate{}
+)
+
+type schemaTemplate struct {
+	dbName string
+	// tableDDL holds the CREATE TABLE statements of the template. MySQL has no
+	// server-side database template, so the schema is replayed statement by
+	// statement into each test database.
+	tableDDL []string
+}
+
+func schemaTemplateKey(engine types.Engine, dsn string) string {
+	return string(engine) + "|" + dsn
+}
+
+func newTestDBName(prefix string) string {
+	return fmt.Sprintf("%s_%s", prefix, strings.ReplaceAll(uuid.New().String(), "-", "_"))
+}
+
+// postgresSchemaTemplate returns the name of a fully migrated database that
+// CREATE DATABASE ... TEMPLATE can copy, creating it on first use.
+func postgresSchemaTemplate(ctx context.Context, baseDSN string, admin *gorm.DB) (string, error) {
+	schemaTemplatesMu.Lock()
+	defer schemaTemplatesMu.Unlock()
+
+	key := schemaTemplateKey(types.PostgresStoreEngine, baseDSN)
+	if tpl, ok := schemaTemplates[key]; ok {
+		return tpl.dbName, nil
+	}
+
+	name := newTestDBName("test_template")
+	if err := admin.Exec(fmt.Sprintf("CREATE DATABASE %s", name)).Error; err != nil {
+		return "", fmt.Errorf("create postgres template database: %w", err)
+	}
+
+	tplStore, err := NewPostgresqlStoreForTests(ctx, replaceDBName(baseDSN, name), nil, false)
+	if err != nil {
+		dropDatabase(admin, name)
+		return "", fmt.Errorf("migrate postgres template database: %w", err)
+	}
+	// TEMPLATE refuses a source that still has sessions, so release both handles
+	// before the first clone.
+	tplStore.Close(ctx)
+	if tplStore.pool != nil {
+		tplStore.pool.Close()
+	}
+
+	schemaTemplates[key] = &schemaTemplate{dbName: name}
+	return name, nil
+}
+
+// mysqlSchemaTemplate returns the CREATE TABLE statements of a fully migrated
+// database, migrating one on first use.
+func mysqlSchemaTemplate(ctx context.Context, baseDSN string, admin *gorm.DB) ([]string, error) {
+	schemaTemplatesMu.Lock()
+	defer schemaTemplatesMu.Unlock()
+
+	key := schemaTemplateKey(types.MysqlStoreEngine, baseDSN)
+	if tpl, ok := schemaTemplates[key]; ok {
+		return tpl.tableDDL, nil
+	}
+
+	name := newTestDBName("test_template")
+	if err := admin.Exec(fmt.Sprintf("CREATE DATABASE %s", name)).Error; err != nil {
+		return nil, fmt.Errorf("create mysql template database: %w", err)
+	}
+
+	tplStore, err := NewMysqlStore(ctx, replaceDBName(baseDSN, name), nil, false)
+	if err != nil {
+		dropDatabase(admin, name)
+		return nil, fmt.Errorf("migrate mysql template database: %w", err)
+	}
+	tableDDL, err := mysqlTableDDL(ctx, tplStore.db, name)
+	tplStore.Close(ctx)
+	if err != nil {
+		dropDatabase(admin, name)
+		return nil, err
+	}
+
+	schemaTemplates[key] = &schemaTemplate{dbName: name, tableDDL: tableDDL}
+	return tableDDL, nil
+}
+
+func mysqlTableDDL(ctx context.Context, db *gorm.DB, dbName string) ([]string, error) {
+	sqlDB, err := db.DB()
+	if err != nil {
+		return nil, err
+	}
+
+	tables, err := mysqlTableNames(ctx, sqlDB, dbName)
+	if err != nil {
+		return nil, err
+	}
+
+	tableDDL := make([]string, 0, len(tables))
+	for _, table := range tables {
+		var name, createStmt string
+		row := sqlDB.QueryRowContext(ctx, fmt.Sprintf("SHOW CREATE TABLE %s.%s", dbName, table))
+		if err := row.Scan(&name, &createStmt); err != nil {
+			return nil, fmt.Errorf("read create statement of %s: %w", table, err)
+		}
+		tableDDL = append(tableDDL, createStmt)
+	}
+	return tableDDL, nil
+}
+
+func mysqlTableNames(ctx context.Context, sqlDB *sql.DB, dbName string) ([]string, error) {
+	rows, err := sqlDB.QueryContext(ctx, fmt.Sprintf("SHOW TABLES FROM %s", dbName))
+	if err != nil {
+		return nil, fmt.Errorf("list template tables: %w", err)
+	}
+	defer rows.Close()
+
+	var tables []string
+	for rows.Next() {
+		var table string
+		if err := rows.Scan(&table); err != nil {
+			return nil, fmt.Errorf("scan template table name: %w", err)
+		}
+		tables = append(tables, table)
+	}
+	if err := rows.Err(); err != nil {
+		return nil, fmt.Errorf("list template tables: %w", err)
+	}
+	return tables, nil
+}
+
+// cloneMysqlSchema replays the template's CREATE TABLE statements into the
+// database the DSN points at.
+func cloneMysqlSchema(ctx context.Context, dsn string, tableDDL []string) error {
+	db, err := gorm.Open(mysql.Open(dsn+"?charset=utf8&parseTime=True&loc=Local"), getGormConfig())
+	if err != nil {
+		return fmt.Errorf("connect to test database: %w", err)
+	}
+	sqlDB, err := db.DB()
+	if err != nil {
+		return err
+	}
+	defer sqlDB.Close()
+
+	// The statements come out of SHOW TABLES in name order, not dependency
+	// order, and their foreign keys reference tables of the session's default
+	// database. Pin a single connection so the session setting below covers
+	// every statement, and connect straight to the new database so unqualified
+	// references land there.
+	sqlDB.SetMaxOpenConns(1)
+	if _, err := sqlDB.ExecContext(ctx, "SET FOREIGN_KEY_CHECKS = 0"); err != nil {
+		return fmt.Errorf("disable foreign key checks: %w", err)
+	}
+	for _, stmt := range tableDDL {
+		if _, err := sqlDB.ExecContext(ctx, stmt); err != nil {
+			return fmt.Errorf("replay table definition: %w", err)
+		}
+	}
+	return nil
+}
+
+// dropDatabase removes a template that never became usable, so a failed setup
+// does not leave it behind on a shared server. The server may still be tearing
+// down the sessions the failed migration held, so the drop retries while
+// Postgres reports the database as in use.
+func dropDatabase(admin *gorm.DB, name string) {
+	if err := execWithTemplateRetry(admin, fmt.Sprintf("DROP DATABASE IF EXISTS %s", name)); err != nil {
+		log.Warnf("failed to drop template database %s: %v", name, err)
+	}
+}
+
+func closeGormDB(db *gorm.DB) {
+	if sqlDB, _ := db.DB(); sqlDB != nil {
+		sqlDB.Close()
+	}
+}
+
 func openDBWithRetry(dsn string, engine types.Engine, maxRetries int) (*gorm.DB, error) {
 	var db *gorm.DB
 	var err error
@@ -885,10 +1102,16 @@ func openDBWithRetry(dsn string, engine types.Engine, maxRetries int) (*gorm.DB,
 	return nil, err
 }
 
-func createRandomDB(dsn string, db *gorm.DB, engine types.Engine) (string, func(), error) {
-	dbName := fmt.Sprintf("test_db_%s", strings.ReplaceAll(uuid.New().String(), "-", "_"))
+// createRandomDB creates a uniquely named database for one test. On postgres a
+// non-empty template is copied server-side with CREATE DATABASE ... TEMPLATE.
+func createRandomDB(dsn string, db *gorm.DB, engine types.Engine, template string) (string, func(), error) {
+	dbName := newTestDBName("test_db")
 
-	if err := db.Exec(fmt.Sprintf("CREATE DATABASE %s", dbName)).Error; err != nil {
+	createStmt := fmt.Sprintf("CREATE DATABASE %s", dbName)
+	if template != "" && engine == types.PostgresStoreEngine {
+		createStmt = fmt.Sprintf("CREATE DATABASE %s TEMPLATE %s", dbName, template)
+	}
+	if err := execWithTemplateRetry(db, createStmt); err != nil {
 		return "", nil, fmt.Errorf("failed to create database: %v", err)
 	}
 
@@ -954,6 +1177,20 @@ func createRandomDB(dsn string, db *gorm.DB, engine types.Engine) (string, func(
 	return replaceDBName(dsn, dbName), cleanup, nil
 }
 
+// execWithTemplateRetry runs a statement, retrying briefly when postgres still
+// sees the template's just-closed sessions and refuses to copy it.
+func execWithTemplateRetry(db *gorm.DB, stmt string) error {
+	var err error
+	for attempt := 0; attempt < 20; attempt++ {
+		err = db.Exec(stmt).Error
+		if err == nil || !strings.Contains(err.Error(), "is being accessed by other users") {
+			return err
+		}
+		time.Sleep(100 * time.Millisecond)
+	}
+	return err
+}
+
 func replaceDBName(dsn, newDBName string) string {
 	re := regexp.MustCompile(`(?P<pre>[:/@])(?P<dbname>[^/?]+)(?P<post>\?|$)`)
 	return re.ReplaceAllString(dsn, `${pre}`+newDBName+`${post}`)