mirror of
https://github.com/netbirdio/netbird.git
synced 2026-10-08 06:29:08 +02:00
Merge branch 'main' into poc/certificate-posture
This commit is contained in:
@@ -3516,7 +3516,13 @@ func buildTestManager(t testing.TB, store store.Store, nmdataStore *networkmapdb
|
||||
|
||||
eventStore := &activity.InMemoryEventStore{}
|
||||
|
||||
metrics, err := telemetry.NewDefaultAppMetrics(context.Background())
|
||||
// Everything built here watches this context; cancelling it on cleanup stops
|
||||
// the metrics flushers, caches and controllers instead of leaking them for
|
||||
// the rest of the package run.
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
t.Cleanup(cancel)
|
||||
|
||||
metrics, err := telemetry.NewDefaultAppMetrics(ctx)
|
||||
if err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
@@ -3535,8 +3541,6 @@ func buildTestManager(t testing.TB, store store.Store, nmdataStore *networkmapdb
|
||||
Return(nil).
|
||||
AnyTimes()
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
cacheStore, err := cache.NewStore(ctx, 100*time.Millisecond, 300*time.Millisecond, 100)
|
||||
if err != nil {
|
||||
return nil, nil, err
|
||||
|
||||
@@ -89,6 +89,7 @@ func TestAgentNetwork_UpdateSettings_PreservesImmutableAndTogglesCollection(t *t
|
||||
|
||||
// Bootstrap is an explicit settings create; providers have no settings
|
||||
// side effects anymore.
|
||||
seedPrivateProxyCluster(t, am.Store, clusterAddr)
|
||||
before, err := mgr.CreateSettings(ctx, adminUserID, agenttypes.DefaultSettings(accountID), clusterAddr, "")
|
||||
require.NoError(t, err, "CreateSettings must bootstrap the row")
|
||||
require.Equal(t, clusterAddr, before.ProxyAddress, "proxy address pinned at bootstrap")
|
||||
|
||||
@@ -13,6 +13,7 @@ import (
|
||||
networkmap "github.com/netbirdio/netbird/management/internals/controllers/network_map"
|
||||
"github.com/netbirdio/netbird/management/internals/modules/agentnetwork"
|
||||
agenttypes "github.com/netbirdio/netbird/management/internals/modules/agentnetwork/types"
|
||||
rpproxy "github.com/netbirdio/netbird/management/internals/modules/reverseproxy/proxy"
|
||||
nbpeer "github.com/netbirdio/netbird/management/server/peer"
|
||||
"github.com/netbirdio/netbird/management/server/permissions"
|
||||
"github.com/netbirdio/netbird/management/server/store"
|
||||
@@ -92,6 +93,7 @@ func TestAgentNetwork_ProviderCRUD_FansOutToProxyAndClientPeers(t *testing.T) {
|
||||
// UpdateAccountPeers, which is the path under test.
|
||||
agentMgr := agentnetwork.NewManager(am.Store, permissions.NewManager(am.Store), am, nil)
|
||||
|
||||
seedPrivateProxyCluster(t, am.Store, clusterAddr)
|
||||
_, err = agentMgr.CreateSettings(ctx, adminUserID, agenttypes.DefaultSettings(accountID), clusterAddr, "")
|
||||
require.NoError(t, err, "CreateSettings must bootstrap the endpoint")
|
||||
// The bootstrap itself reconciles and queues updates on both channels;
|
||||
@@ -222,3 +224,22 @@ func synthZoneRData(sync *nbproto.SyncResponse, clusterAddr, fqdn string) string
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
// seedPrivateProxyCluster registers a connected proxy with private capabilities in a
|
||||
// netbird client for clusterAddr, matching what a real deployment looks like
|
||||
// when the account bootstraps: the agent-network gateway service is always
|
||||
// private, so its cluster has to be one that can serve private services.
|
||||
func seedPrivateProxyCluster(t *testing.T, st store.Store, clusterAddr string) {
|
||||
t.Helper()
|
||||
private := true
|
||||
now := time.Now().UTC()
|
||||
require.NoError(t, st.SaveProxy(context.Background(), &rpproxy.Proxy{
|
||||
ID: "agent-net-proxy-" + clusterAddr,
|
||||
SessionID: "agent-net-session",
|
||||
ClusterAddress: clusterAddr,
|
||||
LastSeen: now,
|
||||
ConnectedAt: &now,
|
||||
Status: rpproxy.StatusConnected,
|
||||
Capabilities: rpproxy.Capabilities{Private: &private},
|
||||
}), "seeding the proxy cluster must succeed")
|
||||
}
|
||||
|
||||
@@ -3157,7 +3157,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 {
|
||||
@@ -3216,23 +3221,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
|
||||
}
|
||||
|
||||
@@ -3244,11 +3246,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
|
||||
@@ -3282,22 +3287,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
|
||||
}
|
||||
@@ -6449,6 +6474,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).
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -4,6 +4,7 @@ package store
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"errors"
|
||||
"fmt"
|
||||
"net"
|
||||
@@ -15,6 +16,7 @@ import (
|
||||
"runtime"
|
||||
"slices"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/google/uuid"
|
||||
@@ -340,6 +342,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)
|
||||
@@ -732,6 +737,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)
|
||||
}
|
||||
|
||||
@@ -758,14 +764,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
|
||||
@@ -781,6 +788,11 @@ 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
|
||||
@@ -805,19 +817,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
|
||||
}
|
||||
|
||||
@@ -850,7 +866,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()
|
||||
|
||||
@@ -858,14 +880,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
|
||||
@@ -891,10 +1099,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)
|
||||
}
|
||||
|
||||
@@ -960,6 +1174,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}`)
|
||||
|
||||
@@ -3065,6 +3065,51 @@ func (mr *MockStoreMockRecorder) HasActiveProxyAtClusterAddress(ctx, clusterAddr
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "HasActiveProxyAtClusterAddress", reflect.TypeOf((*MockStore)(nil).HasActiveProxyAtClusterAddress), ctx, clusterAddress)
|
||||
}
|
||||
|
||||
// HasForeignAccountProxyAtHost mocks base method.
|
||||
func (m *MockStore) HasForeignAccountProxyAtHost(ctx context.Context, host, accountID string) (bool, error) {
|
||||
m.ctrl.T.Helper()
|
||||
ret := m.ctrl.Call(m, "HasForeignAccountProxyAtHost", ctx, host, accountID)
|
||||
ret0, _ := ret[0].(bool)
|
||||
ret1, _ := ret[1].(error)
|
||||
return ret0, ret1
|
||||
}
|
||||
|
||||
// HasForeignAccountProxyAtHost indicates an expected call of HasForeignAccountProxyAtHost.
|
||||
func (mr *MockStoreMockRecorder) HasForeignAccountProxyAtHost(ctx, host, accountID any) *gomock.Call {
|
||||
mr.mock.ctrl.T.Helper()
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "HasForeignAccountProxyAtHost", reflect.TypeOf((*MockStore)(nil).HasForeignAccountProxyAtHost), ctx, host, accountID)
|
||||
}
|
||||
|
||||
// HasGatewayClusterPinnedByOtherAccount mocks base method.
|
||||
func (m *MockStore) HasGatewayClusterPinnedByOtherAccount(ctx context.Context, host, accountID string) (bool, error) {
|
||||
m.ctrl.T.Helper()
|
||||
ret := m.ctrl.Call(m, "HasGatewayClusterPinnedByOtherAccount", ctx, host, accountID)
|
||||
ret0, _ := ret[0].(bool)
|
||||
ret1, _ := ret[1].(error)
|
||||
return ret0, ret1
|
||||
}
|
||||
|
||||
// HasGatewayClusterPinnedByOtherAccount indicates an expected call of HasGatewayClusterPinnedByOtherAccount.
|
||||
func (mr *MockStoreMockRecorder) HasGatewayClusterPinnedByOtherAccount(ctx, host, accountID any) *gomock.Call {
|
||||
mr.mock.ctrl.T.Helper()
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "HasGatewayClusterPinnedByOtherAccount", reflect.TypeOf((*MockStore)(nil).HasGatewayClusterPinnedByOtherAccount), ctx, host, accountID)
|
||||
}
|
||||
|
||||
// HasGatewayEndpointByOtherAccount mocks base method.
|
||||
func (m *MockStore) HasGatewayEndpointByOtherAccount(ctx context.Context, host, accountID string) (bool, error) {
|
||||
m.ctrl.T.Helper()
|
||||
ret := m.ctrl.Call(m, "HasGatewayEndpointByOtherAccount", ctx, host, accountID)
|
||||
ret0, _ := ret[0].(bool)
|
||||
ret1, _ := ret[1].(error)
|
||||
return ret0, ret1
|
||||
}
|
||||
|
||||
// HasGatewayEndpointByOtherAccount indicates an expected call of HasGatewayEndpointByOtherAccount.
|
||||
func (mr *MockStoreMockRecorder) HasGatewayEndpointByOtherAccount(ctx, host, accountID any) *gomock.Call {
|
||||
mr.mock.ctrl.T.Helper()
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "HasGatewayEndpointByOtherAccount", reflect.TypeOf((*MockStore)(nil).HasGatewayEndpointByOtherAccount), ctx, host, accountID)
|
||||
}
|
||||
|
||||
// IncrementAgentNetworkConsumption mocks base method.
|
||||
func (m *MockStore) IncrementAgentNetworkConsumption(ctx context.Context, accountID string, kind types.ConsumptionDimension, dimID string, windowSeconds int64, windowStart time.Time, tokensIn, tokensOut int64, costUSD float64) error {
|
||||
m.ctrl.T.Helper()
|
||||
|
||||
@@ -37,6 +37,18 @@ func CreateMysqlTestContainer() (func(), string, error) {
|
||||
mysql.WithDatabase("testing"),
|
||||
mysql.WithUsername("root"),
|
||||
mysql.WithPassword("testing"),
|
||||
// Every test creates and drops a database with about 40 tables, so with
|
||||
// the server defaults the run is dominated by durability work: each
|
||||
// CREATE TABLE fsyncs the redo log, the binary log and the doublewrite
|
||||
// buffer. None of it protects anything in a container that is discarded
|
||||
// after the run. Tables stay in per-table files on purpose: in the shared
|
||||
// system tablespace the cost of every CREATE and DROP grew with the number
|
||||
// of databases the run had already created.
|
||||
testcontainers.WithCmd("mysqld",
|
||||
"--innodb-flush-log-at-trx-commit=0",
|
||||
"--innodb-doublewrite=OFF",
|
||||
"--skip-log-bin",
|
||||
),
|
||||
testcontainers.WithWaitStrategy(
|
||||
wait.ForLog("/usr/sbin/mysqld: ready for connections").
|
||||
WithOccurrence(1).WithStartupTimeout(15*time.Second).WithPollInterval(100*time.Millisecond),
|
||||
|
||||
Reference in New Issue
Block a user