move read-only queries to an interface to reuse in tests

Signed-off-by: Dmitri Dolguikh <dmitri.external@netbird.io>

integration_tests/management/network_map_db/pgsql/account_settings_test.go

@@ -7,7 +7,6 @@ import (
 	"testing"
 	"time"
 
-	networkmap_pgsql "github.com/netbirdio/netbird/management/internals/network_map_db/pgsql"
 	"github.com/netbirdio/netbird/shared/management/networkmap/nmdata"
 	"github.com/stretchr/testify/assert"
 )
@@ -21,7 +20,7 @@ func TestGetAccountSettings(t *testing.T) {
                       settings_lazy_connection_enabled, settings_auto_update_version, settings_auto_update_always, settings_metrics_push_enabled)
 		 values('account-3',null,null,null,null,null,null,null,null,null,null,null)`)
 
-	accountSettings, err := networkmap_pgsql.GetAccountSettingsViaPgxConnection(ctx, conn(t, ctx), "account-1")
+	accountSettings, err := conn(t, ctx).GetAccountSettings(ctx, "account-1")
 	assert.NoError(t, err)
 	assert.Equal(t, accountSettings, nmdata.AccountSettingsInfo{
 		PeerLoginExpirationEnabled:      true,
@@ -37,7 +36,7 @@ func TestGetAccountSettings(t *testing.T) {
 		MetricsPushEnabled:              false,
 	})
 
-	accountSettings, err = networkmap_pgsql.GetAccountSettingsViaPgxConnection(ctx, conn(t, ctx), "account-2")
+	accountSettings, err = conn(t, ctx).GetAccountSettings(ctx, "account-2")
 	assert.NoError(t, err)
 	assert.Equal(t, accountSettings, nmdata.AccountSettingsInfo{
 		PeerLoginExpirationEnabled:      true,
@@ -53,7 +52,7 @@ func TestGetAccountSettings(t *testing.T) {
 		MetricsPushEnabled:              false,
 	})
 
-	accountSettings, err = networkmap_pgsql.GetAccountSettingsViaPgxConnection(ctx, conn(t, ctx), "account-3")
+	accountSettings, err = conn(t, ctx).GetAccountSettings(ctx, "account-3")
 	assert.NoError(t, err)
 	assert.Equal(t, accountSettings, nmdata.AccountSettingsInfo{})
 }

integration_tests/management/network_map_db/pgsql/dns_settings_test.go

@@ -13,13 +13,13 @@ import (
 func TestGetDnsSettings(t *testing.T) {
 	ctx := context.TODO()
 
-	settings, err := pgstore.GetDnsSettings(ctx, "account-1")
+	settings, err := conn(t, ctx).GetDnsSettings(ctx, "account-1")
 	assert.NoError(t, err)
 	assert.Equal(t, settings, nmdata.DNSSettings{
 		DisabledManagementGroups: []string{"disabled-group-1", "disabled-group-2"},
 	})
 
-	settings, err = pgstore.GetDnsSettings(ctx, "account-2")
+	settings, err = conn(t, ctx).GetDnsSettings(ctx, "account-2")
 	assert.NoError(t, err)
 	assert.Equal(t, settings, nmdata.DNSSettings{})
 }

integration_tests/management/network_map_db/pgsql/dns_test.go

@@ -7,7 +7,6 @@ import (
 	"testing"
 
 	"github.com/miekg/dns"
-	networkmap_pgsql "github.com/netbirdio/netbird/management/internals/network_map_db/pgsql"
 	"github.com/netbirdio/netbird/shared/management/networkmap"
 	"github.com/netbirdio/netbird/shared/management/networkmap/nmdata"
 	"github.com/stretchr/testify/assert"
@@ -35,7 +34,7 @@ func TestGetAppliedZoneCandidatesViaPgxConnection(t *testing.T) {
 		`insert into records (id, account_id, zone_id, name, type, ttl, content)
 		VALUES('record-4','account-1','zone-2','test2.test-2.com','CNAME',1800,'test3.test-2.com')`)
 
-	zoneCandidates, err := networkmap_pgsql.GetAppliedZoneCandidatesViaPgxConnection(ctx, conn(t, ctx), "account-1")
+	zoneCandidates, err := conn(t, ctx).GetAppliedZoneCandidates(ctx, "account-1")
 	assert.NoError(t, err)
 
 	assert.Contains(t, zoneCandidates, networkmap.AppliedZoneCandidate{

integration_tests/management/network_map_db/pgsql/domain_test.go

@@ -7,7 +7,7 @@ import (
 	"database/sql"
 	"testing"
 
-	networkmap_pgsql "github.com/netbirdio/netbird/management/internals/network_map_db/pgsql"
+	networkmapdb "github.com/netbirdio/netbird/management/internals/network_map_db"
 	"github.com/stretchr/testify/assert"
 )
 
@@ -24,15 +24,15 @@ func TestGetDomains(t *testing.T) {
 		`insert into domains (id, account_id, domain, target_cluster)
 		VALUES('domain-3','account-1',null,null)`)
 
-	domains, err := pgstore.GetDomains(ctx, "account-1")
+	domains, err := conn(t, ctx).GetDomains(ctx, "account-1")
 	assert.NoError(t, err)
 	assert.Len(t, domains, 2)
 
-	assert.Contains(t, domains, networkmap_pgsql.Domain{
+	assert.Contains(t, domains, networkmapdb.Domain{
 		Domain:        sql.NullString{String: "test-1.com", Valid: true},
 		TargetCluster: sql.NullString{String: "target-1.cluster.local", Valid: true},
 	})
-	assert.Contains(t, domains, networkmap_pgsql.Domain{
+	assert.Contains(t, domains, networkmapdb.Domain{
 		Domain:        sql.NullString{String: "test-2.com", Valid: true},
 		TargetCluster: sql.NullString{String: "target-2.cluster.local", Valid: true},
 	})

integration_tests/management/network_map_db/pgsql/group_test.go

@@ -6,7 +6,6 @@ import (
 	"context"
 	"testing"
 
-	networkmap_pgsql "github.com/netbirdio/netbird/management/internals/network_map_db/pgsql"
 	"github.com/netbirdio/netbird/shared/management/networkmap/nmdata"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
@@ -15,10 +14,7 @@ import (
 func TestGetGroups(t *testing.T) {
 	ctx := context.TODO()
 
-	s, err := networkmap_pgsql.NewPostgresqlStore(ctx, dsn)
-	assert.NoError(t, err)
-
-	groups, resourceToGroupIdx, err := s.GetGroups(ctx, "account-1")
+	groups, resourceToGroupIdx, err := conn(t, ctx).GetGroups(ctx, "account-1")
 	assert.NoError(t, err)
 	assert.Contains(t,
 		groups,
@@ -45,17 +41,13 @@ func TestGetGroups(t *testing.T) {
 func TestGetGroupsWithoutExpectedFields(t *testing.T) {
 	ctx := context.TODO()
 
-	s, err := networkmap_pgsql.NewPostgresqlStore(ctx, dsn)
-	assert.NoError(t, err)
-
 	execQuery(t, ctx,
 		"insert into accounts (id) VALUES('random-id')")
 
 	execQuery(t, ctx,
 		"insert into groups (id, account_id) VALUES('g2-test-group-id-1','random-id')")
-	assert.NoError(t, err)
 
-	groups, _, err := s.GetGroups(ctx, "random-id")
+	groups, _, err := conn(t, ctx).GetGroups(ctx, "random-id")
 	assert.NoError(t, err)
 	require.Len(t, groups, 1)
 	assert.NotEmpty(t, groups[0].PublicID)

integration_tests/management/network_map_db/pgsql/main_test.go

@@ -5,149 +5,51 @@ package networkmap_pgsql
 import (
 	"context"
 	_ "embed"
-	"fmt"
 	"os"
-	"regexp"
-	"strings"
 	"testing"
 	"time"
 
-	"github.com/google/uuid"
-	"github.com/jackc/pgx/v5"
+	log "github.com/sirupsen/logrus"
 	"github.com/stretchr/testify/assert"
 
-	log "github.com/sirupsen/logrus"
-	"gorm.io/driver/postgres"
-	"gorm.io/gorm"
-
+	networkmapdb "github.com/netbirdio/netbird/management/internals/network_map_db"
 	networkmap_pgsql "github.com/netbirdio/netbird/management/internals/network_map_db/pgsql"
-	gormstore "github.com/netbirdio/netbird/management/server/store"
-	"github.com/netbirdio/netbird/management/server/testutil"
+	"github.com/netbirdio/netbird/management/server/types"
 )
 
 //go:embed base_data.sql
 var baseData string
 
 var (
-	dsn     string
 	pgstore *networkmap_pgsql.PgStore
+	engine  string
 )
 
 func TestMain(m *testing.M) {
-	_, tmpdsn, err := testutil.CreatePostgresTestContainer()
-	if err != nil {
-		log.Fatalf("error starting postres container %v", err)
+	kind, _ := os.LookupEnv("NETBIRD_STORE_ENGINE")
+	switch kind {
+	case "":
+		engine = string(types.PostgresStoreEngine)
+	case string(types.PostgresStoreEngine), string(types.SqliteStoreEngine):
+		engine = kind
+	default:
+		log.Fatalf("unsupported db '%s' in NETBIRD_STORE_ENGINE env var", kind)
 	}
 
-	var db *gorm.DB
-	for i := range 5 {
-		db, err = gorm.Open(postgres.Open(tmpdsn), &gorm.Config{})
-
-		if err == nil {
-			break
-		}
-
-		if i < 5 {
-			waitTime := time.Duration(100*(i+1)) * time.Millisecond
-			time.Sleep(waitTime)
-			continue
-		}
-
-		log.Fatalf("error connecting to postres db %v", err)
-	}
-
-	var cleanup func()
-	dsn, cleanup, err = createRandomDB(tmpdsn, db)
-	sqlDB, _ := db.DB()
-	if sqlDB != nil {
-		sqlDB.Close()
-	}
-	if err != nil {
-		log.Fatalf("error creating postres db %v", err)
-	}
-
-	_, err = gormstore.NewPostgresqlStoreForTests(context.TODO(), dsn, nil, false)
-	if err != nil {
-		log.Fatalf("error running migrations %v", err)
-	}
-
-	ctx := context.TODO()
-	s, err := networkmap_pgsql.NewPostgresqlStore(ctx, dsn)
-	if err != nil {
-		log.Fatal("error creating postgres store %w", err)
-	}
-
-	for _, query := range strings.Split(baseData, ";") {
-		if _, err := s.Pool.Exec(ctx, query); err != nil {
-			log.Fatalf("error initializing db: %s", err.Error())
-		}
-	}
-
-	pgstore, err = networkmap_pgsql.NewPostgresqlStore(ctx, dsn)
-	if err != nil {
-		log.Fatalf("error creating pg store %v", err.Error())
-	}
+	store, cleanup := createPGTestStore(baseData)
+	pgstore = store
 
 	code := m.Run()
 
 	cleanup()
-
 	os.Exit(code)
 }
 
-func createRandomDB(dsn string, db *gorm.DB) (string, func(), error) {
-	dbName := fmt.Sprintf("test_db_%s", strings.ReplaceAll(uuid.New().String(), "-", "_"))
-
-	if err := db.Exec(fmt.Sprintf("CREATE DATABASE %s", dbName)).Error; err != nil {
-		return "", nil, fmt.Errorf("failed to create database: %v", err)
-	}
-
-	originalDSN := dsn
-
-	cleanup := func() {
-		var dropDB *gorm.DB
-		var err error
-
-		dropDB, err = gorm.Open(postgres.Open(originalDSN), &gorm.Config{
-			SkipDefaultTransaction: true,
-			PrepareStmt:            false,
-		})
-		if err != nil {
-			log.Errorf("failed to connect for dropping database %s: %v", dbName, err)
-			return
-		}
-		defer func() {
-			if sqlDB, _ := dropDB.DB(); sqlDB != nil {
-				sqlDB.Close()
-			}
-		}()
-
-		if sqlDB, _ := dropDB.DB(); sqlDB != nil {
-			sqlDB.SetMaxOpenConns(1)
-			sqlDB.SetMaxIdleConns(0)
-			sqlDB.SetConnMaxLifetime(time.Second)
-		}
-
-		err = dropDB.Exec(fmt.Sprintf("DROP DATABASE IF EXISTS %s WITH (FORCE)", dbName)).Error
-
-		if err != nil {
-			log.Errorf("failed to drop database %s: %v", dbName, err)
-		}
-	}
-
-	return replaceDBName(dsn, dbName), cleanup, nil
-}
-
-func replaceDBName(dsn, newDBName string) string {
-	re := regexp.MustCompile(`(?P<pre>[:/@])(?P<dbname>[^/?]+)(?P<post>\?|$)`)
-	return re.ReplaceAllString(dsn, `${pre}`+newDBName+`${post}`)
-}
-
-func conn(t *testing.T, ctx context.Context) *pgx.Conn {
+func conn(t *testing.T, ctx context.Context) networkmapdb.NetworkMapDBStoreConn {
 	t.Helper()
 	c, err := pgstore.Pool.Acquire(ctx)
 	assert.NoError(t, err)
-	return c.Conn()
+	return pgstore.UsingConnection(c.Conn())
 }
 
 func execQuery(t *testing.T, ctx context.Context, q string) {

integration_tests/management/network_map_db/pgsql/nameserver_test.go

@@ -7,7 +7,6 @@ import (
 	"net/netip"
 	"testing"
 
-	networkmap_pgsql "github.com/netbirdio/netbird/management/internals/network_map_db/pgsql"
 	"github.com/netbirdio/netbird/shared/management/networkmap/nmdata"
 	"github.com/stretchr/testify/assert"
 )
@@ -25,7 +24,7 @@ func TestGetNameServerGroups(t *testing.T) {
 		`insert into name_server_groups (id, public_id, name, description, name_servers, groups, domains, enabled, search_domains_enabled,"primary",account_id)
 		VALUES('nsgroup-3','nsgroup-3-public',null,null,null,null,null,TRUE,FALSE,FALSE,'account-1')`)
 
-	nsgroups, err := networkmap_pgsql.GetNameServerGroupsViaPgxConnection(ctx, conn(t, ctx), "account-1")
+	nsgroups, err := conn(t, ctx).GetNameServerGroups(ctx, "account-1")
 	assert.NoError(t, err)
 
 	assert.Contains(t, nsgroups, nmdata.NameServerGroup{

integration_tests/management/network_map_db/pgsql/network_resource_test.go

@@ -7,7 +7,6 @@ import (
 	"net/netip"
 	"testing"
 
-	networkmap_pgsql "github.com/netbirdio/netbird/management/internals/network_map_db/pgsql"
 	"github.com/netbirdio/netbird/shared/management/networkmap/nmdata"
 	"github.com/stretchr/testify/assert"
 )
@@ -15,9 +14,6 @@ import (
 func TestGetNetworkResources(t *testing.T) {
 	ctx := context.TODO()
 
-	s, err := networkmap_pgsql.NewPostgresqlStore(ctx, dsn)
-	assert.NoError(t, err)
-
 	execQuery(t, ctx,
 		`insert into network_resources (id, account_id, network_id, public_id, name, description, type, domain, prefix, enabled)
 		VALUES('net-resource-1','account-1','network-1','net-resource-public-1','network-resource-1','network-resource-1','subnet','','"10.0.0.0/16"',TRUE)`)
@@ -28,7 +24,7 @@ func TestGetNetworkResources(t *testing.T) {
 		`insert into network_resources (id, account_id, network_id, public_id, name, description, type, domain, prefix, enabled)
 		VALUES('net-resource-3','account-1','network-3','net-resource-public-3','network-resource-3','network-resource-3','host','','"10.0.0.1/32"',TRUE)`)
 
-	resources, err := s.GetNetworkResources(ctx, "account-1")
+	resources, err := conn(t, ctx).GetNetworkResources(ctx, "account-1")
 	assert.NoError(t, err)
 
 	assert.Contains(t, resources, nmdata.NetworkResource{

integration_tests/management/network_map_db/pgsql/network_router_test.go

@@ -6,7 +6,6 @@ import (
 	"context"
 	"testing"
 
-	networkmap_pgsql "github.com/netbirdio/netbird/management/internals/network_map_db/pgsql"
 	"github.com/netbirdio/netbird/shared/management/networkmap/nmdata"
 	"github.com/stretchr/testify/assert"
 )
@@ -14,9 +13,6 @@ import (
 func TestGetNetworkRouters(t *testing.T) {
 	ctx := context.TODO()
 
-	s, err := networkmap_pgsql.NewPostgresqlStore(ctx, dsn)
-	assert.NoError(t, err)
-
 	execQuery(t, ctx,
 		`insert into network_routers (id, account_id, public_id, peer, network_id, masquerade, metric, enabled, peer_groups)
 		VALUES('test-nr-id-1','account-1','public-id-1','peer-id-1','network-id-1',TRUE,999,TRUE,'["group-one-resource-id"]')`)
@@ -24,7 +20,7 @@ func TestGetNetworkRouters(t *testing.T) {
 		`insert into network_routers (id, account_id, public_id, peer, network_id, masquerade, metric, enabled, peer_groups)
 		VALUES('test-nr-id-2','account-1','public-id-2','','network-id-2',TRUE,333,TRUE,'["group-two-resources-id","group-no-resources-id"]')`)
 
-	routers, err := s.GetNetworkRouters(ctx, "account-1")
+	routers, err := conn(t, ctx).GetNetworkRouters(ctx, "account-1")
 	assert.NoError(t, err)
 	assert.NotEmpty(t, routers)
 

integration_tests/management/network_map_db/pgsql/network_test.go

@@ -8,7 +8,6 @@ import (
 	"net"
 	"testing"
 
-	networkmap_pgsql "github.com/netbirdio/netbird/management/internals/network_map_db/pgsql"
 	"github.com/netbirdio/netbird/shared/management/networkmap/nmdata"
 	"github.com/stretchr/testify/assert"
 )
@@ -16,10 +15,7 @@ import (
 func TestGetNetwork(t *testing.T) {
 	ctx := context.TODO()
 
-	s, err := networkmap_pgsql.NewPostgresqlStore(ctx, dsn)
-	assert.NoError(t, err)
-
-	network, err := s.GetNetwork(ctx, "account-1")
+	network, err := conn(t, ctx).GetNetwork(ctx, "account-1")
 	assert.NoError(t, err)
 	assert.Equal(t, network, nmdata.Network{
 		Identifier: "network-1",
@@ -28,7 +24,7 @@ func TestGetNetwork(t *testing.T) {
 		Serial:     1,
 	})
 
-	network, err = s.GetNetwork(ctx, "account-2")
+	network, err = conn(t, ctx).GetNetwork(ctx, "account-2")
 	assert.NoError(t, err)
 	assert.Equal(t, network, nmdata.Network{
 		Identifier: "network-2",

integration_tests/management/network_map_db/pgsql/networks_test.go

@@ -6,7 +6,6 @@ import (
 	"context"
 	"testing"
 
-	networkmap_pgsql "github.com/netbirdio/netbird/management/internals/network_map_db/pgsql"
 	"github.com/stretchr/testify/assert"
 )
 
@@ -18,7 +17,7 @@ func TestGetNetworks(t *testing.T) {
 	execQuery(t, ctx,
 		`insert into networks (id, account_id, public_id) VALUES('network-2','account-1','network-2-public')`)
 
-	networksIdx, err := networkmap_pgsql.GetNetworkXIDToPublicIdMapViaPgxConnection(ctx, conn(t, ctx), "account-1")
+	networksIdx, err := conn(t, ctx).GetNetworkXIDToPublicIdMap(ctx, "account-1")
 	assert.NoError(t, err)
 	assert.Equal(t, networksIdx, map[string]string{
 		"network-1": "network-1-public",

integration_tests/management/network_map_db/pgsql/peer_test.go

@@ -8,7 +8,6 @@ import (
 	"net/netip"
 	"testing"
 
-	networkmap_pgsql "github.com/netbirdio/netbird/management/internals/network_map_db/pgsql"
 	"github.com/netbirdio/netbird/shared/management/networkmap/nmdata"
 	"github.com/stretchr/testify/assert"
 )
@@ -16,7 +15,7 @@ import (
 func TestGetPeers(t *testing.T) {
 	ctx := context.TODO()
 
-	peers, clusterToPeersIdx, err := networkmap_pgsql.GetPeersViaPgxConnection(ctx, conn(t, ctx), "account-1")
+	peers, clusterToPeersIdx, err := conn(t, ctx).GetPeers(ctx, "account-1")
 	assert.NoError(t, err)
 
 	// shouldn't be returned in the index, as it's not connected

integration_tests/management/network_map_db/pgsql/pg_test_store.go

@@ -0,0 +1,126 @@
+//go:build integration
+
+package networkmap_pgsql
+
+import (
+	"context"
+	"fmt"
+	"regexp"
+	"strings"
+	"time"
+
+	log "github.com/sirupsen/logrus"
+
+	"github.com/google/uuid"
+	networkmap_pgsql "github.com/netbirdio/netbird/management/internals/network_map_db/pgsql"
+	gormstore "github.com/netbirdio/netbird/management/server/store"
+	"github.com/netbirdio/netbird/management/server/testutil"
+	"gorm.io/driver/postgres"
+	"gorm.io/gorm"
+)
+
+func createPGTestStore(baseData string) (*networkmap_pgsql.PgStore, func()) {
+	_, tmpdsn, err := testutil.CreatePostgresTestContainer()
+	if err != nil {
+		log.Fatalf("error starting postres container %v", err)
+	}
+
+	var db *gorm.DB
+	for i := range 5 {
+		db, err = gorm.Open(postgres.Open(tmpdsn), &gorm.Config{})
+
+		if err == nil {
+			break
+		}
+
+		if i < 5 {
+			waitTime := time.Duration(100*(i+1)) * time.Millisecond
+			time.Sleep(waitTime)
+			continue
+		}
+
+		log.Fatalf("error connecting to postres db %v", err)
+	}
+
+	var cleanup func()
+	dsn, cleanup, err := createRandomDB(tmpdsn, db)
+	sqlDB, _ := db.DB()
+	if sqlDB != nil {
+		sqlDB.Close()
+	}
+	if err != nil {
+		log.Fatalf("error creating postres db %v", err)
+	}
+
+	_, err = gormstore.NewPostgresqlStoreForTests(context.TODO(), dsn, nil, false)
+	if err != nil {
+		log.Fatalf("error running migrations %v", err)
+	}
+
+	ctx := context.TODO()
+	s, err := networkmap_pgsql.NewPostgresqlStore(ctx, dsn)
+	if err != nil {
+		log.Fatal("error creating postgres store %w", err)
+	}
+
+	for _, query := range strings.Split(baseData, ";") {
+		if _, err := s.Pool.Exec(ctx, query); err != nil {
+			log.Fatalf("error initializing db: %s", err.Error())
+		}
+	}
+
+	pgstore, err := networkmap_pgsql.NewPostgresqlStore(ctx, dsn)
+	if err != nil {
+		log.Fatalf("error creating pg store %v", err.Error())
+	}
+
+	return pgstore, cleanup
+}
+
+func createRandomDB(dsn string, db *gorm.DB) (string, func(), error) {
+	dbName := fmt.Sprintf("test_db_%s", strings.ReplaceAll(uuid.New().String(), "-", "_"))
+
+	if err := db.Exec(fmt.Sprintf("CREATE DATABASE %s", dbName)).Error; err != nil {
+		return "", nil, fmt.Errorf("failed to create database: %v", err)
+	}
+
+	originalDSN := dsn
+
+	cleanup := func() {
+		var dropDB *gorm.DB
+		var err error
+
+		dropDB, err = gorm.Open(postgres.Open(originalDSN), &gorm.Config{
+			SkipDefaultTransaction: true,
+			PrepareStmt:            false,
+		})
+		if err != nil {
+			log.Errorf("failed to connect for dropping database %s: %v", dbName, err)
+			return
+		}
+		defer func() {
+			if sqlDB, _ := dropDB.DB(); sqlDB != nil {
+				sqlDB.Close()
+			}
+		}()
+
+		if sqlDB, _ := dropDB.DB(); sqlDB != nil {
+			sqlDB.SetMaxOpenConns(1)
+			sqlDB.SetMaxIdleConns(0)
+			sqlDB.SetConnMaxLifetime(time.Second)
+		}
+
+		err = dropDB.Exec(fmt.Sprintf("DROP DATABASE IF EXISTS %s WITH (FORCE)", dbName)).Error
+
+		if err != nil {
+			log.Errorf("failed to drop database %s: %v", dbName, err)
+		}
+	}
+
+	return replaceDBName(dsn, dbName), cleanup, nil
+}
+
+func replaceDBName(dsn, newDBName string) string {
+	re := regexp.MustCompile(`(?P<pre>[:/@])(?P<dbname>[^/?]+)(?P<post>\?|$)`)
+	return re.ReplaceAllString(dsn, `${pre}`+newDBName+`${post}`)
+}

integration_tests/management/network_map_db/pgsql/policy_test.go

@@ -6,7 +6,6 @@ import (
 	"context"
 	"testing"
 
-	networkmap_pgsql "github.com/netbirdio/netbird/management/internals/network_map_db/pgsql"
 	"github.com/netbirdio/netbird/shared/management/networkmap/nmdata"
 	"github.com/stretchr/testify/assert"
 )
@@ -54,7 +53,7 @@ func TestGetPolicies(t *testing.T) {
 		 values('policy-4-rule-1','policy-4',false,null,null,null,null,'["group-two-resources-id"]',
 		        null,'{"ID":"domain-3","Type":"domain"}',null,null,null,null)`)
 
-	policies, policyToDestinationResourceIdx, policyToDestinationGroupIdx, err := networkmap_pgsql.GetPoliciesViaPgxConnection(ctx, conn(t, ctx), "account-1")
+	policies, policyToDestinationResourceIdx, policyToDestinationGroupIdx, err := conn(t, ctx).GetPolicies(ctx, "account-1")
 	assert.NoError(t, err)
 
 	assert.Contains(t, policies, nmdata.Policy{

integration_tests/management/network_map_db/pgsql/posture_test.go

@@ -7,7 +7,6 @@ import (
 	"net/netip"
 	"testing"
 
-	networkmap_pgsql "github.com/netbirdio/netbird/management/internals/network_map_db/pgsql"
 	"github.com/netbirdio/netbird/shared/management/networkmap/nmdata"
 	"github.com/stretchr/testify/assert"
 )
@@ -34,7 +33,7 @@ func TestGetPostureChecks(t *testing.T) {
 		`insert into posture_checks (id, account_id, public_id, checks)
 		VALUES('posturecheck-3','account-1','posturecheck-3-public', null)`)
 
-	postureChecks, idToPublicIDIdx, err := networkmap_pgsql.GetPostureChecksViaPgxConnection(ctx, conn(t, ctx), "account-1")
+	postureChecks, idToPublicIDIdx, err := conn(t, ctx).GetPostureChecks(ctx, "account-1")
 	assert.NoError(t, err)
 	assert.Equal(t, idToPublicIDIdx, map[string]string{
 		"posturecheck-1": "posturecheck-1-public",

integration_tests/management/network_map_db/pgsql/route_test.go

@@ -7,7 +7,6 @@ import (
 	"net/netip"
 	"testing"
 
-	networkmap_pgsql "github.com/netbirdio/netbird/management/internals/network_map_db/pgsql"
 	"github.com/netbirdio/netbird/shared/management/domain"
 	"github.com/netbirdio/netbird/shared/management/networkmap/nmdata"
 	"github.com/stretchr/testify/assert"
@@ -37,7 +36,7 @@ func TestGetRoutes(t *testing.T) {
 		VALUES('route-3','account-1','route-3-public',null,null,null,null,'route-3',
 		        null,null,null,null,null,null,null,null,null)`)
 
-	routes, err := networkmap_pgsql.GetRoutesViaPgxConnection(ctx, conn(t, ctx), "account-1")
+	routes, err := conn(t, ctx).GetRoutes(ctx, "account-1")
 	assert.NoError(t, err)
 	assert.Contains(t, routes, nmdata.Route{
 		ID:                  "route-1",

integration_tests/management/network_map_db/pgsql/service_test.go

@@ -9,7 +9,7 @@ import (
 
 	"github.com/stretchr/testify/assert"
 
-	networkmap_pgsql "github.com/netbirdio/netbird/management/internals/network_map_db/pgsql"
+	networkmapdb "github.com/netbirdio/netbird/management/internals/network_map_db"
 )
 
 func TestGetPrivateServicesViaPgxConnection(t *testing.T) {
@@ -28,23 +28,23 @@ func TestGetPrivateServicesViaPgxConnection(t *testing.T) {
 		 values('service-3','account-1',null,null,null,null,null)`)
 	assert.NoError(t, err)
 
-	services, err := networkmap_pgsql.GetPrivateServicesViaPgxConnection(ctx, conn(t, ctx), "account-1")
+	services, err := conn(t, ctx).GetPrivateServices(ctx, "account-1")
 	assert.NoError(t, err)
-	assert.Contains(t, services, networkmap_pgsql.Service{
+	assert.Contains(t, services, networkmapdb.Service{
 		Enabled:      sql.NullBool{Bool: true, Valid: true},
 		Private:      sql.NullBool{Bool: true, Valid: true},
 		AccessGroups: []string{"group-one-resource-id"},
 		ProxyCluster: sql.NullString{String: "test-1.com", Valid: true},
 		Domain:       sql.NullString{String: "test-2.com", Valid: true},
 	})
-	assert.Contains(t, services, networkmap_pgsql.Service{
+	assert.Contains(t, services, networkmapdb.Service{
 		Enabled:      sql.NullBool{Bool: true, Valid: true},
 		Private:      sql.NullBool{Bool: true, Valid: true},
 		AccessGroups: []string{"group-one-resource-id", "group-two-resources-id"},
 		ProxyCluster: sql.NullString{String: "test-3.com", Valid: true},
 		Domain:       sql.NullString{String: "test-4.com", Valid: true},
 	})
-	assert.Contains(t, services, networkmap_pgsql.Service{
+	assert.Contains(t, services, networkmapdb.Service{
 		Enabled:      sql.NullBool{Bool: false, Valid: false},
 		Private:      sql.NullBool{Bool: false, Valid: false},
 		AccessGroups: []string{},
@@ -102,7 +102,7 @@ func TestGetProxyTargetedDomainResourceIDsViaPgxConnection(t *testing.T) {
 		`insert into targets (target_id, account_id, service_id, enabled, target_type)
 		 values(null,'account-1','service-4',true,'cluster')`)
 
-	servtargetedDomains, err := networkmap_pgsql.GetProxyTargetedDomainResourceIDsViaPgxConnection(ctx, conn(t, ctx), "account-1")
+	servtargetedDomains, err := conn(t, ctx).GetProxyTargetedDomainResourceIDs(ctx, "account-1")
 	assert.NoError(t, err)
 	assert.Equal(t, servtargetedDomains, map[string]struct{}{
 		"target-1": {},

integration_tests/management/network_map_db/pgsql/user_test.go

@@ -6,7 +6,6 @@ import (
 	"context"
 	"testing"
 
-	networkmap_pgsql "github.com/netbirdio/netbird/management/internals/network_map_db/pgsql"
 	"github.com/stretchr/testify/assert"
 )
 
@@ -40,7 +39,7 @@ func TestGetAllowedUsers(t *testing.T) {
 		`insert into groups (id, name, account_id)
 		VALUES('all-group-3','All','account-1')`)
 
-	userIdx, groupIdToUserIds, err := networkmap_pgsql.GetAllowedUsersViaPgxConnection(ctx, conn(t, ctx), "account-1")
+	userIdx, groupIdToUserIds, err := conn(t, ctx).GetAllowedUsers(ctx, "account-1")
 	assert.NoError(t, err)
 
 	assert.Equal(t, userIdx, map[string]struct{}{