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

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

management/internals/network_map_db/pgsql/account_settings.go

@@ -28,17 +28,8 @@ const (
 	`
 )
 
-func (pg *PgStore) GetAccountSettings(ctx context.Context, accountId string) (nmdata.AccountSettingsInfo, error) {
-	c, err := pg.Pool.Acquire(ctx)
-	if err != nil {
-		return nmdata.AccountSettingsInfo{}, err
-	}
-	return GetAccountSettingsViaPgxConnection(ctx, c.Conn(), accountId)
-
-}
-
-func GetAccountSettingsViaPgxConnection(ctx context.Context, con *pgx.Conn, accountId string) (nmdata.AccountSettingsInfo, error) {
-	rows, err := con.Query(ctx, GetAccountSettingsQuery, accountId)
+func (pgc *PgStoreConn) GetAccountSettings(ctx context.Context, accountId string) (nmdata.AccountSettingsInfo, error) {
+	rows, err := pgc.Conn.Query(ctx, GetAccountSettingsQuery, accountId)
 	if err != nil {
 		return nmdata.AccountSettingsInfo{}, err
 	}

management/internals/network_map_db/pgsql/dns.go

@@ -27,16 +27,8 @@ const (
 	`
 )
 
-func (pg *PgStore) GetAppliedZoneCandidates(ctx context.Context, accountId string) ([]networkmap.AppliedZoneCandidate, error) {
-	c, err := pg.Pool.Acquire(ctx)
-	if err != nil {
-		return nil, err
-	}
-	return GetAppliedZoneCandidatesViaPgxConnection(ctx, c.Conn(), accountId)
-}
-
-func GetAppliedZoneCandidatesViaPgxConnection(ctx context.Context, conn *pgx.Conn, accountId string) ([]networkmap.AppliedZoneCandidate, error) {
-	rows, err := conn.Query(ctx, GetAccountZonesQuery, accountId)
+func (pgc *PgStoreConn) GetAppliedZoneCandidates(ctx context.Context, accountId string) ([]networkmap.AppliedZoneCandidate, error) {
+	rows, err := pgc.Conn.Query(ctx, GetAccountZonesQuery, accountId)
 	if err != nil {
 		return nil, err
 	}

management/internals/network_map_db/pgsql/dns_settings.go

@@ -16,16 +16,8 @@ const (
 	`
 )
 
-func (pg *PgStore) GetDnsSettings(ctx context.Context, accountId string) (nmdata.DNSSettings, error) {
-	c, err := pg.Pool.Acquire(ctx)
-	if err != nil {
-		return nmdata.DNSSettings{}, err
-	}
-	return GetDnsSettingsViaPgxConnection(ctx, c.Conn(), accountId)
-}
-
-func GetDnsSettingsViaPgxConnection(ctx context.Context, con *pgx.Conn, accountId string) (nmdata.DNSSettings, error) {
-	rows, err := con.Query(ctx, GetDnsSettingsQuery, accountId)
+func (pgc *PgStoreConn) GetDnsSettings(ctx context.Context, accountId string) (nmdata.DNSSettings, error) {
+	rows, err := pgc.Conn.Query(ctx, GetDnsSettingsQuery, accountId)
 	if err != nil {
 		return nmdata.DNSSettings{}, err
 	}

management/internals/network_map_db/pgsql/domain.go

@@ -2,9 +2,9 @@ package networkmap_pgsql
 
 import (
 	"context"
-	"database/sql"
 
 	"github.com/jackc/pgx/v5"
+	networkmapdb "github.com/netbirdio/netbird/management/internals/network_map_db"
 )
 
 const (
@@ -15,24 +15,11 @@ const (
 	`
 )
 
-func (pg *PgStore) GetDomains(ctx context.Context, accountId string) ([]Domain, error) {
-	c, err := pg.Pool.Acquire(ctx)
-	if err != nil {
-		return nil, err
-	}
-	return GetDomainsViaPgxConnection(ctx, c.Conn(), accountId)
-}
-
-func GetDomainsViaPgxConnection(ctx context.Context, conn *pgx.Conn, accountId string) ([]Domain, error) {
-	rows, err := conn.Query(ctx, GetDomainsQuery, accountId)
+func (pgc *PgStoreConn) GetDomains(ctx context.Context, accountId string) ([]networkmapdb.Domain, error) {
+	rows, err := pgc.Conn.Query(ctx, GetDomainsQuery, accountId)
 	if err != nil {
 		return nil, err
 	}
 
-	return pgx.CollectRows(rows, pgx.RowToStructByName[Domain])
-}
-
-type Domain struct {
-	Domain        sql.NullString
-	TargetCluster sql.NullString
+	return pgx.CollectRows(rows, pgx.RowToStructByName[networkmapdb.Domain])
 }

management/internals/network_map_db/pgsql/group.go

@@ -26,16 +26,8 @@ const (
 // we also return a resource-to-group index.
 // an alternative is to add json indexes, query this directly. Not sure how expensive
 // json indexes are. TODO (dmitri) verify and maybe change the implementation here.
-func (pg *PgStore) GetGroups(ctx context.Context, accountId string) ([]nmdata.Group, map[string]map[string]any, error) {
-	c, err := pg.Pool.Acquire(ctx)
-	if err != nil {
-		return nil, nil, err
-	}
-	return GetGroupsViaPgxConnection(ctx, c.Conn(), accountId)
-}
-
-func GetGroupsViaPgxConnection(ctx context.Context, con *pgx.Conn, accountId string) ([]nmdata.Group, map[string]map[string]any, error) {
-	rows, err := con.Query(ctx, GetGroupsQuery, accountId)
+func (pgc *PgStoreConn) GetGroups(ctx context.Context, accountId string) ([]nmdata.Group, map[string]map[string]any, error) {
+	rows, err := pgc.Conn.Query(ctx, GetGroupsQuery, accountId)
 	if err != nil {
 		return nil, nil, err
 	}

management/internals/network_map_db/pgsql/nameserver.go

@@ -19,16 +19,8 @@ const (
 	`
 )
 
-func (pg *PgStore) GetNameServerGroups(ctx context.Context, accountId string) ([]nmdata.NameServerGroup, error) {
-	c, err := pg.Pool.Acquire(ctx)
-	if err != nil {
-		return nil, err
-	}
-	return GetNameServerGroupsViaPgxConnection(ctx, c.Conn(), accountId)
-}
-
-func GetNameServerGroupsViaPgxConnection(ctx context.Context, con *pgx.Conn, accountId string) ([]nmdata.NameServerGroup, error) {
-	rows, err := con.Query(ctx, GetNameserversQuery, accountId)
+func (pgc *PgStoreConn) GetNameServerGroups(ctx context.Context, accountId string) ([]nmdata.NameServerGroup, error) {
+	rows, err := pgc.Conn.Query(ctx, GetNameserversQuery, accountId)
 	if err != nil {
 		return nil, err
 	}

management/internals/network_map_db/pgsql/network.go

@@ -19,16 +19,8 @@ const (
 	`
 )
 
-func (pg *PgStore) GetNetwork(ctx context.Context, accountId string) (nmdata.Network, error) {
-	c, err := pg.Pool.Acquire(ctx)
-	if err != nil {
-		return nmdata.Network{}, err
-	}
-	return GetNetworkViaPgxConnection(ctx, c.Conn(), accountId)
-}
-
-func GetNetworkViaPgxConnection(ctx context.Context, con *pgx.Conn, accountId string) (nmdata.Network, error) {
-	rows, err := con.Query(ctx, GetNetworkQuery, accountId)
+func (pgc *PgStoreConn) GetNetwork(ctx context.Context, accountId string) (nmdata.Network, error) {
+	rows, err := pgc.Conn.Query(ctx, GetNetworkQuery, accountId)
 	if err != nil {
 		return nmdata.Network{}, err
 	}

management/internals/network_map_db/pgsql/network_map_data.go

@@ -9,6 +9,7 @@ import (
 	"github.com/miekg/dns"
 	log "github.com/sirupsen/logrus"
 
+	networkmapdb "github.com/netbirdio/netbird/management/internals/network_map_db"
 	"github.com/netbirdio/netbird/shared/management/networkmap"
 	"github.com/netbirdio/netbird/shared/management/networkmap/nmdata"
 )
@@ -19,71 +20,73 @@ func (pg *PgStore) GetNetworkMapData(ctx context.Context, accountId string) (*ne
 		return nil, err
 	}
 
-	acctSettings, err := GetAccountSettingsViaPgxConnection(ctx, tx.Conn(), accountId)
+	conn := pg.UsingConnection(tx.Conn())
+
+	acctSettings, err := conn.GetAccountSettings(ctx, accountId)
 	if err != nil {
 		return rollbackAndReturnError(ctx, tx, fmt.Errorf("failed to get account settings: %w", err))
 	}
-	dnsZones, err := GetAppliedZoneCandidatesViaPgxConnection(ctx, tx.Conn(), accountId)
+	dnsZones, err := conn.GetAppliedZoneCandidates(ctx, accountId)
 	if err != nil {
 		return rollbackAndReturnError(ctx, tx, fmt.Errorf("failed to get applied zone candidates: %w", err))
 	}
-	groups, resourceToGroupIdx, err := GetGroupsViaPgxConnection(ctx, tx.Conn(), accountId)
+	groups, resourceToGroupIdx, err := conn.GetGroups(ctx, accountId)
 	if err != nil {
 		return rollbackAndReturnError(ctx, tx, fmt.Errorf("failed to get groups: %w", err))
 	}
-	nsGroups, err := GetNameServerGroupsViaPgxConnection(ctx, tx.Conn(), accountId)
+	nsGroups, err := conn.GetNameServerGroups(ctx, accountId)
 	if err != nil {
 		return rollbackAndReturnError(ctx, tx, fmt.Errorf("failed to get nameserver groups: %w", err))
 	}
-	networkResources, err := GetNetworkResourcesViaPgxConnection(ctx, tx.Conn(), accountId)
+	networkResources, err := conn.GetNetworkResources(ctx, accountId)
 	if err != nil {
 		return rollbackAndReturnError(ctx, tx, fmt.Errorf("failed to get network resources: %w", err))
 	}
-	routers, err := GetNetworkRoutersViaPgxConnection(ctx, tx.Conn(), accountId)
+	routers, err := conn.GetNetworkRouters(ctx, accountId)
 	if err != nil {
 		return rollbackAndReturnError(ctx, tx, fmt.Errorf("failed to get network routers: %w", err))
 	}
-	network, err := GetNetworkViaPgxConnection(ctx, tx.Conn(), accountId)
+	network, err := conn.GetNetwork(ctx, accountId)
 	if err != nil {
 		return rollbackAndReturnError(ctx, tx, fmt.Errorf("failed to get network: %w", err))
 	}
-	peers, proxyPeers, err := GetPeersViaPgxConnection(ctx, tx.Conn(), accountId)
+	peers, proxyPeers, err := conn.GetPeers(ctx, accountId)
 	if err != nil {
 		return rollbackAndReturnError(ctx, tx, fmt.Errorf("failed to get peers: %w", err))
 	}
-	policies, policyToDestinationResourceIdx, policyToDestinationGroupIdx, err := GetPoliciesViaPgxConnection(ctx, tx.Conn(), accountId)
+	policies, policyToDestinationResourceIdx, policyToDestinationGroupIdx, err := conn.GetPolicies(ctx, accountId)
 	if err != nil {
 		return rollbackAndReturnError(ctx, tx, fmt.Errorf("failed to get policies: %w", err))
 	}
-	postureChecks, postureCheckXIDToPublicID, err := GetPostureChecksViaPgxConnection(ctx, tx.Conn(), accountId)
+	postureChecks, postureCheckXIDToPublicID, err := conn.GetPostureChecks(ctx, accountId)
 	if err != nil {
 		return rollbackAndReturnError(ctx, tx, fmt.Errorf("failed to get posture checks: %w", err))
 	}
-	routes, err := GetRoutesViaPgxConnection(ctx, tx.Conn(), accountId)
+	routes, err := conn.GetRoutes(ctx, accountId)
 	if err != nil {
 		return rollbackAndReturnError(ctx, tx, fmt.Errorf("failed to get routes: %w", err))
 	}
-	networkXIDToPublicID, err := GetNetworkXIDToPublicIdMapViaPgxConnection(ctx, tx.Conn(), accountId)
+	networkXIDToPublicID, err := conn.GetNetworkXIDToPublicIdMap(ctx, accountId)
 	if err != nil {
 		return rollbackAndReturnError(ctx, tx, fmt.Errorf("failed to get network xid to public id map: %w", err))
 	}
-	allowedUserIds, groupsToUserIds, err := GetAllowedUsersViaPgxConnection(ctx, tx.Conn(), accountId)
+	allowedUserIds, groupsToUserIds, err := conn.GetAllowedUsers(ctx, accountId)
 	if err != nil {
 		return rollbackAndReturnError(ctx, tx, fmt.Errorf("failed to get allowed users: %w", err))
 	}
-	dnsSettings, err := GetDnsSettingsViaPgxConnection(ctx, tx.Conn(), accountId)
+	dnsSettings, err := conn.GetDnsSettings(ctx, accountId)
 	if err != nil {
 		return rollbackAndReturnError(ctx, tx, fmt.Errorf("failed to get dns settings: %w", err))
 	}
-	domains, err := GetDomainsViaPgxConnection(ctx, tx.Conn(), accountId)
+	domains, err := conn.GetDomains(ctx, accountId)
 	if err != nil {
 		return rollbackAndReturnError(ctx, tx, err)
 	}
-	services, err := GetPrivateServicesViaPgxConnection(ctx, tx.Conn(), accountId)
+	services, err := conn.GetPrivateServices(ctx, accountId)
 	if err != nil {
 		return rollbackAndReturnError(ctx, tx, err)
 	}
-	proxyTargetedDomainResourceIDs, err := GetProxyTargetedDomainResourceIDsViaPgxConnection(ctx, tx.Conn(), accountId)
+	proxyTargetedDomainResourceIDs, err := conn.GetProxyTargetedDomainResourceIDs(ctx, accountId)
 	if err != nil {
 		return rollbackAndReturnError(ctx, tx, fmt.Errorf("failed to get proxy targeted domain resources: %w", err))
 	}
@@ -165,7 +168,7 @@ func toSliceOfPtrs[T any](all []T) []*T {
 	return toret
 }
 
-func serviceDomainZone(svc Service, ds []Domain) string {
+func serviceDomainZone(svc networkmapdb.Service, ds []networkmapdb.Domain) string {
 	if domainFromSuffix(svc.Domain.String, svc.ProxyCluster.String) {
 		return svc.ProxyCluster.String
 	}
@@ -190,7 +193,7 @@ func domainFromSuffix(domain, suffix string) bool {
 	return domain == suffix || strings.HasSuffix(domain, "."+suffix)
 }
 
-func buildPrivateServiceCandidates(svcs []Service, domains []Domain, proxyPeersByCluster map[string][]*nmdata.Peer) []networkmap.PrivateServiceCandidate {
+func buildPrivateServiceCandidates(svcs []networkmapdb.Service, domains []networkmapdb.Domain, proxyPeersByCluster map[string][]*nmdata.Peer) []networkmap.PrivateServiceCandidate {
 	var out []networkmap.PrivateServiceCandidate
 
 	if len(proxyPeersByCluster) == 0 {

management/internals/network_map_db/pgsql/network_resource.go

@@ -19,16 +19,8 @@ const (
 	`
 )
 
-func (pg *PgStore) GetNetworkResources(ctx context.Context, accountId string) ([]nmdata.NetworkResource, error) {
-	c, err := pg.Pool.Acquire(ctx)
-	if err != nil {
-		return nil, err
-	}
-	return GetNetworkResourcesViaPgxConnection(ctx, c.Conn(), accountId)
-}
-
-func GetNetworkResourcesViaPgxConnection(ctx context.Context, con *pgx.Conn, accountId string) ([]nmdata.NetworkResource, error) {
-	rows, err := con.Query(ctx, GetNetworkResourcesQuery, accountId)
+func (pgc *PgStoreConn) GetNetworkResources(ctx context.Context, accountId string) ([]nmdata.NetworkResource, error) {
+	rows, err := pgc.Conn.Query(ctx, GetNetworkResourcesQuery, accountId)
 	if err != nil {
 		return nil, err
 	}

management/internals/network_map_db/pgsql/network_router.go

@@ -25,16 +25,8 @@ const (
 	`
 )
 
-func (pg *PgStore) GetNetworkRouters(ctx context.Context, accountId string) (map[string]map[string]*nmdata.NetworkRouter, error) {
-	c, err := pg.Pool.Acquire(ctx)
-	if err != nil {
-		return nil, err
-	}
-	return GetNetworkRoutersViaPgxConnection(ctx, c.Conn(), accountId)
-}
-
-func GetNetworkRoutersViaPgxConnection(ctx context.Context, con *pgx.Conn, accountId string) (map[string]map[string]*nmdata.NetworkRouter, error) {
-	rows, err := con.Query(ctx, GetNetworkRouterQuery, accountId)
+func (pgc *PgStoreConn) GetNetworkRouters(ctx context.Context, accountId string) (map[string]map[string]*nmdata.NetworkRouter, error) {
+	rows, err := pgc.Conn.Query(ctx, GetNetworkRouterQuery, accountId)
 	if err != nil {
 		return nil, err
 	}

management/internals/network_map_db/pgsql/networks.go

@@ -14,24 +14,16 @@ const (
 	`
 )
 
-func (pg *PgStore) GetNetworks(ctx context.Context, accountId string) ([]network, error) {
-	c, err := pg.Pool.Acquire(ctx)
-	if err != nil {
-		return nil, err
-	}
-	return GetNetworksViaPgxConnection(ctx, c.Conn(), accountId)
-}
-
-func GetNetworksViaPgxConnection(ctx context.Context, con *pgx.Conn, accountId string) ([]network, error) {
-	rows, err := con.Query(ctx, GetNetworksQuery, accountId)
+func (pgc *PgStoreConn) GetNetworks(ctx context.Context, accountId string) ([]network, error) {
+	rows, err := pgc.Conn.Query(ctx, GetNetworksQuery, accountId)
 	if err != nil {
 		return nil, err
 	}
 	return pgx.CollectRows(rows, pgx.RowToStructByName[network])
 }
 
-func GetNetworkXIDToPublicIdMapViaPgxConnection(ctx context.Context, con *pgx.Conn, accountId string) (map[string]string, error) {
-	networks, err := GetNetworksViaPgxConnection(ctx, con, accountId)
+func (pgc *PgStoreConn) GetNetworkXIDToPublicIdMap(ctx context.Context, accountId string) (map[string]string, error) {
+	networks, err := pgc.GetNetworks(ctx, accountId)
 	if err != nil {
 		return nil, err
 	}

management/internals/network_map_db/pgsql/peer.go

@@ -22,16 +22,8 @@ const (
 	`
 )
 
-func (pg *PgStore) GetPeers(ctx context.Context, accountId string) ([]nmdata.Peer, map[string][]*nmdata.Peer, error) {
-	c, err := pg.Pool.Acquire(ctx)
-	if err != nil {
-		return nil, nil, err
-	}
-	return GetPeersViaPgxConnection(ctx, c.Conn(), accountId)
-}
-
-func GetPeersViaPgxConnection(ctx context.Context, con *pgx.Conn, accountId string) ([]nmdata.Peer, map[string][]*nmdata.Peer, error) {
-	rows, err := con.Query(ctx, GetPeersQuery, accountId)
+func (pgc *PgStoreConn) GetPeers(ctx context.Context, accountId string) ([]nmdata.Peer, map[string][]*nmdata.Peer, error) {
+	rows, err := pgc.Conn.Query(ctx, GetPeersQuery, accountId)
 	if err != nil {
 		return nil, nil, err
 	}

management/internals/network_map_db/pgsql/pg_store.go

@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"time"
 
+	"github.com/jackc/pgx/v5"
 	"github.com/jackc/pgx/v5/pgxpool"
 	networkmapdb "github.com/netbirdio/netbird/management/internals/network_map_db"
 )
@@ -22,6 +23,12 @@ type PgStore struct {
 	Pool *pgxpool.Pool
 }
 
+type PgStoreConn struct {
+	Conn *pgx.Conn
+}
+
+var _ networkmapdb.NetworkMapDBStoreConn = &PgStoreConn{}
+
 func NewPostgresqlStore(ctx context.Context, dsn string) (*PgStore, error) {
 	pool, err := connectToPgDb(ctx, dsn)
 	if err != nil {
@@ -31,6 +38,10 @@ func NewPostgresqlStore(ctx context.Context, dsn string) (*PgStore, error) {
 	return &PgStore{Pool: pool}, nil
 }
 
+func (p *PgStore) UsingConnection(c *pgx.Conn) networkmapdb.NetworkMapDBStoreConn {
+	return &PgStoreConn{Conn: c}
+}
+
 func connectToPgDb(ctx context.Context, dsn string) (*pgxpool.Pool, error) {
 	config, err := pgxpool.ParseConfig(dsn)
 	if err != nil {

management/internals/network_map_db/pgsql/policy.go

@@ -22,16 +22,8 @@ const (
 	`
 )
 
-func (pg *PgStore) GetPolicies(ctx context.Context, accountId string) ([]nmdata.Policy, map[string]map[string]any, map[string]map[string]any, error) {
-	c, err := pg.Pool.Acquire(ctx)
-	if err != nil {
-		return nil, nil, nil, err
-	}
-	return GetPoliciesViaPgxConnection(ctx, c.Conn(), accountId)
-}
-
-func GetPoliciesViaPgxConnection(ctx context.Context, con *pgx.Conn, accountId string) ([]nmdata.Policy, map[string]map[string]any, map[string]map[string]any, error) {
-	rows, err := con.Query(ctx, GetPoliciesQuery, accountId)
+func (pgc *PgStoreConn) GetPolicies(ctx context.Context, accountId string) ([]nmdata.Policy, map[string]map[string]any, map[string]map[string]any, error) {
+	rows, err := pgc.Conn.Query(ctx, GetPoliciesQuery, accountId)
 	if err != nil {
 		return nil, nil, nil, err
 	}

management/internals/network_map_db/pgsql/posture.go

@@ -19,16 +19,8 @@ const (
 	`
 )
 
-func (pg *PgStore) GetPostureChecks(ctx context.Context, accountId string) ([]nmdata.PostureChecks, map[string]string, error) {
-	c, err := pg.Pool.Acquire(ctx)
-	if err != nil {
-		return nil, nil, err
-	}
-	return GetPostureChecksViaPgxConnection(ctx, c.Conn(), accountId)
-}
-
-func GetPostureChecksViaPgxConnection(ctx context.Context, con *pgx.Conn, accountId string) ([]nmdata.PostureChecks, map[string]string, error) {
-	rows, err := con.Query(ctx, GetPostureChecksQuery, accountId)
+func (pgc *PgStoreConn) GetPostureChecks(ctx context.Context, accountId string) ([]nmdata.PostureChecks, map[string]string, error) {
+	rows, err := pgc.Conn.Query(ctx, GetPostureChecksQuery, accountId)
 	if err != nil {
 		return nil, nil, err
 	}

management/internals/network_map_db/pgsql/route.go

@@ -21,16 +21,8 @@ const (
 	`
 )
 
-func (pg *PgStore) GetRoutes(ctx context.Context, accountId string) ([]nmdata.Route, error) {
-	c, err := pg.Pool.Acquire(ctx)
-	if err != nil {
-		return nil, err
-	}
-	return GetRoutesViaPgxConnection(ctx, c.Conn(), accountId)
-}
-
-func GetRoutesViaPgxConnection(ctx context.Context, con *pgx.Conn, accountId string) ([]nmdata.Route, error) {
-	rows, err := con.Query(ctx, GetRoutesQuery, accountId)
+func (pgc *PgStoreConn) GetRoutes(ctx context.Context, accountId string) ([]nmdata.Route, error) {
+	rows, err := pgc.Conn.Query(ctx, GetRoutesQuery, accountId)
 	if err != nil {
 		return nil, err
 	}

management/internals/network_map_db/pgsql/service.go

@@ -2,9 +2,9 @@ package networkmap_pgsql
 
 import (
 	"context"
-	"database/sql"
 
 	"github.com/jackc/pgx/v5"
+	networkmapdb "github.com/netbirdio/netbird/management/internals/network_map_db"
 )
 
 const (
@@ -23,25 +23,17 @@ const (
 	`
 )
 
-func (pg *PgStore) GetPrivateServices(ctx context.Context, accountId string) ([]Service, error) {
-	c, err := pg.Pool.Acquire(ctx)
-	if err != nil {
-		return nil, err
-	}
-	return GetPrivateServicesViaPgxConnection(ctx, c.Conn(), accountId)
-}
-
-func GetPrivateServicesViaPgxConnection(ctx context.Context, conn *pgx.Conn, accountId string) ([]Service, error) {
-	rows, err := conn.Query(ctx, GetServicesQuery, accountId)
+func (pgc *PgStoreConn) GetPrivateServices(ctx context.Context, accountId string) ([]networkmapdb.Service, error) {
+	rows, err := pgc.Conn.Query(ctx, GetServicesQuery, accountId)
 	if err != nil {
 		return nil, err
 	}
 
-	return pgx.CollectRows(rows, pgx.RowToStructByName[Service])
+	return pgx.CollectRows(rows, pgx.RowToStructByName[networkmapdb.Service])
 }
 
-func GetProxyTargetedDomainResourceIDsViaPgxConnection(ctx context.Context, conn *pgx.Conn, accountId string) (map[string]struct{}, error) {
-	rows, err := conn.Query(ctx, GetProxyTargetedDomainResourcesQuery, accountId)
+func (pgc *PgStoreConn) GetProxyTargetedDomainResourceIDs(ctx context.Context, accountId string) (map[string]struct{}, error) {
+	rows, err := pgc.Conn.Query(ctx, GetProxyTargetedDomainResourcesQuery, accountId)
 	if err != nil {
 		return nil, err
 	}
@@ -57,11 +49,3 @@ func GetProxyTargetedDomainResourceIDsViaPgxConnection(ctx context.Context, conn
 	}
 	return toret, nil
 }
-
-type Service struct {
-	Enabled      sql.NullBool
-	Private      sql.NullBool
-	AccessGroups []string
-	ProxyCluster sql.NullString
-	Domain       sql.NullString
-}

management/internals/network_map_db/pgsql/user.go

@@ -19,16 +19,8 @@ const (
 	`
 )
 
-func (pg *PgStore) GetAllowedUsers(ctx context.Context, accountId string) (map[string]struct{}, map[string][]string, error) {
-	c, err := pg.Pool.Acquire(ctx)
-	if err != nil {
-		return nil, nil, err
-	}
-	return GetAllowedUsersViaPgxConnection(ctx, c.Conn(), accountId)
-}
-
-func GetAllowedUsersViaPgxConnection(ctx context.Context, con *pgx.Conn, accountId string) (map[string]struct{}, map[string][]string, error) {
-	rows, err := con.Query(ctx, GetAllowedUserIdsQuery, accountId)
+func (pgc *PgStoreConn) GetAllowedUsers(ctx context.Context, accountId string) (map[string]struct{}, map[string][]string, error) {
+	rows, err := pgc.Conn.Query(ctx, GetAllowedUserIdsQuery, accountId)
 	if err != nil {
 		return nil, nil, err
 	}
@@ -38,7 +30,7 @@ func GetAllowedUsersViaPgxConnection(ctx context.Context, con *pgx.Conn, account
 		return nil, nil, err
 	}
 
-	rows, err = con.Query(ctx, GetAllGroupIdQuery, accountId)
+	rows, err = pgc.Conn.Query(ctx, GetAllGroupIdQuery, accountId)
 	if err != nil {
 		return nil, nil, err
 	}