added support for queries via pgx connection

Signed-off-by: Dmitri Dolguikh <dmitri.external@netbird.io>
This commit is contained in:
Dmitri Dolguikh
2026-07-31 15:03:15 +02:00
parent acbe22b831
commit ea1b4d56e8
12 changed files with 113 additions and 11 deletions

View File

@@ -8,6 +8,7 @@ import (
"reflect"
"strings"
"github.com/netbirdio/netbird/shared/management/networkmap"
"github.com/netbirdio/netbird/shared/management/networkmap/nmdata"
)
@@ -29,6 +30,7 @@ type NetworkMapDBStore interface {
GetAccountZones(ctx context.Context, accountId string) ([]nmdata.CustomZone, error)
GetAccountSettings(ctx context.Context, accountId string) (nmdata.AccountSettingsInfo, error)
GetPostureChecks(ctx context.Context, accountId string) ([]nmdata.PostureChecks, error)
GetNetworkMapData(ctx context.Context, accountId string) (*networkmap.NetworkMapData, error)
}
type NetworkMapDBStoreImpl struct {

View File

@@ -6,6 +6,7 @@ import (
"time"
"github.com/jackc/pgx/v5"
"github.com/jackc/pgx/v5/pgxpool"
"github.com/netbirdio/netbird/shared/management/networkmap/nmdata"
)
@@ -21,7 +22,16 @@ const (
)
func (pg *PgStore) GetAccountSettings(ctx context.Context, accountId string) (nmdata.AccountSettingsInfo, error) {
rows, err := pg.Pool.Query(ctx, GetAccountSettingsQuery, accountId)
c, err := pg.Pool.Acquire(ctx)
if err != nil {
return nmdata.AccountSettingsInfo{}, err
}
return GetAccountSettingsViaConnection(ctx, c, accountId)
}
func GetAccountSettingsViaConnection(ctx context.Context, con *pgxpool.Conn, accountId string) (nmdata.AccountSettingsInfo, error) {
rows, err := con.Query(ctx, GetAccountSettingsQuery, accountId)
if err != nil {
return nmdata.AccountSettingsInfo{}, err
}

View File

@@ -8,6 +8,7 @@ import (
"reflect"
"github.com/jackc/pgx/v5"
"github.com/jackc/pgx/v5/pgxpool"
"github.com/miekg/dns"
networkmapdb "github.com/netbirdio/netbird/management/internals/network_map_db"
"github.com/netbirdio/netbird/shared/management/networkmap/nmdata"
@@ -26,7 +27,15 @@ const (
)
func (pg *PgStore) GetAccountZones(ctx context.Context, accountId string) ([]nmdata.CustomZone, error) {
rows, err := pg.Pool.Query(ctx, GetAccountZonesQuery, accountId)
c, err := pg.Pool.Acquire(ctx)
if err != nil {
return nil, err
}
return GetAccountZonesViaConnection(ctx, c, accountId)
}
func GetAccountZonesViaConnection(ctx context.Context, conn *pgxpool.Conn, accountId string) ([]nmdata.CustomZone, error) {
rows, err := conn.Query(ctx, GetAccountZonesQuery, accountId)
if err != nil {
return nil, err
}

View File

@@ -7,6 +7,7 @@ import (
"reflect"
"github.com/jackc/pgx/v5"
"github.com/jackc/pgx/v5/pgxpool"
networkmapdb "github.com/netbirdio/netbird/management/internals/network_map_db"
"github.com/netbirdio/netbird/shared/management/networkmap/nmdata"
)
@@ -24,7 +25,15 @@ const (
)
func (pg *PgStore) GetGroups(ctx context.Context, accountId string) ([]nmdata.Group, error) {
rows, err := pg.Pool.Query(ctx, GetGroupsQuery, accountId)
c, err := pg.Pool.Acquire(ctx)
if err != nil {
return nil, err
}
return GetGroupsViaConnection(ctx, c, accountId)
}
func GetGroupsViaConnection(ctx context.Context, con *pgxpool.Conn, accountId string) ([]nmdata.Group, error) {
rows, err := con.Query(ctx, GetGroupsQuery, accountId)
if err != nil {
return nil, err
}

View File

@@ -7,6 +7,7 @@ import (
"reflect"
"github.com/jackc/pgx/v5"
"github.com/jackc/pgx/v5/pgxpool"
networkmapdb "github.com/netbirdio/netbird/management/internals/network_map_db"
"github.com/netbirdio/netbird/shared/management/networkmap/nmdata"
)
@@ -20,7 +21,15 @@ const (
)
func (pg *PgStore) GetNameServerGroups(ctx context.Context, accountId string) ([]nmdata.NameServerGroup, error) {
rows, err := pg.Pool.Query(ctx, GetNameserversQuery, accountId)
c, err := pg.Pool.Acquire(ctx)
if err != nil {
return nil, err
}
return GetNameServerGroupsViaConnection(ctx, c, accountId)
}
func GetNameServerGroupsViaConnection(ctx context.Context, con *pgxpool.Conn, accountId string) ([]nmdata.NameServerGroup, error) {
rows, err := con.Query(ctx, GetNameserversQuery, accountId)
if err != nil {
return nil, err
}

View File

@@ -7,6 +7,7 @@ import (
"reflect"
"github.com/jackc/pgx/v5"
"github.com/jackc/pgx/v5/pgxpool"
networkmapdb "github.com/netbirdio/netbird/management/internals/network_map_db"
"github.com/netbirdio/netbird/shared/management/networkmap/nmdata"
)
@@ -20,7 +21,15 @@ const (
)
func (pg *PgStore) GetNetwork(ctx context.Context, accountId string) (nmdata.Network, error) {
rows, err := pg.Pool.Query(ctx, GetNetworkQuery, accountId)
c, err := pg.Pool.Acquire(ctx)
if err != nil {
return nmdata.Network{}, err
}
return GetNetworkViaConnection(ctx, c, accountId)
}
func GetNetworkViaConnection(ctx context.Context, con *pgxpool.Conn, accountId string) (nmdata.Network, error) {
rows, err := con.Query(ctx, GetNetworkQuery, accountId)
if err != nil {
return nmdata.Network{}, err
}

View File

@@ -7,6 +7,7 @@ import (
"reflect"
"github.com/jackc/pgx/v5"
"github.com/jackc/pgx/v5/pgxpool"
networkmapdb "github.com/netbirdio/netbird/management/internals/network_map_db"
"github.com/netbirdio/netbird/shared/management/networkmap/nmdata"
)
@@ -20,7 +21,15 @@ const (
)
func (pg *PgStore) GetNetworkResources(ctx context.Context, accountId string) ([]nmdata.NetworkResource, error) {
rows, err := pg.Pool.Query(ctx, GetNetworkResourcesQuery, accountId)
c, err := pg.Pool.Acquire(ctx)
if err != nil {
return nil, err
}
return GetNetworkResourcesViaConnection(ctx, c, accountId)
}
func GetNetworkResourcesViaConnection(ctx context.Context, con *pgxpool.Conn, accountId string) ([]nmdata.NetworkResource, error) {
rows, err := con.Query(ctx, GetNetworkResourcesQuery, accountId)
if err != nil {
return nil, err
}

View File

@@ -7,6 +7,7 @@ import (
"reflect"
"github.com/jackc/pgx/v5"
"github.com/jackc/pgx/v5/pgxpool"
networkmapdb "github.com/netbirdio/netbird/management/internals/network_map_db"
"github.com/netbirdio/netbird/shared/management/networkmap/nmdata"
)
@@ -20,7 +21,15 @@ const (
)
func (pg *PgStore) GetNetworkRouters(ctx context.Context, accountId string) ([]nmdata.NetworkRouter, error) {
rows, err := pg.Pool.Query(ctx, GetNetworkRouterQuery, accountId)
c, err := pg.Pool.Acquire(ctx)
if err != nil {
return nil, err
}
return GetNetworkRoutersViaConnection(ctx, c, accountId)
}
func GetNetworkRoutersViaConnection(ctx context.Context, con *pgxpool.Conn, accountId string) ([]nmdata.NetworkRouter, error) {
rows, err := con.Query(ctx, GetNetworkRouterQuery, accountId)
if err != nil {
return nil, err
}

View File

@@ -7,6 +7,7 @@ import (
"reflect"
"github.com/jackc/pgx/v5"
"github.com/jackc/pgx/v5/pgxpool"
networkmapdb "github.com/netbirdio/netbird/management/internals/network_map_db"
"github.com/netbirdio/netbird/shared/management/networkmap/nmdata"
)
@@ -22,7 +23,15 @@ const (
)
func (pg *PgStore) GetPeers(ctx context.Context, accountId string) ([]nmdata.Peer, error) {
rows, err := pg.Pool.Query(ctx, GetPeersQuery, accountId)
c, err := pg.Pool.Acquire(ctx)
if err != nil {
return nil, err
}
return GetPeersViaConnection(ctx, c, accountId)
}
func GetPeersViaConnection(ctx context.Context, con *pgxpool.Conn, accountId string) ([]nmdata.Peer, error) {
rows, err := con.Query(ctx, GetPeersQuery, accountId)
if err != nil {
return nil, err
}

View File

@@ -7,6 +7,7 @@ import (
"reflect"
"github.com/jackc/pgx/v5"
"github.com/jackc/pgx/v5/pgxpool"
networkmapdb "github.com/netbirdio/netbird/management/internals/network_map_db"
"github.com/netbirdio/netbird/shared/management/networkmap/nmdata"
)
@@ -23,7 +24,15 @@ const (
)
func (pg *PgStore) GetPolicies(ctx context.Context, accountId string) ([]nmdata.Policy, error) {
rows, err := pg.Pool.Query(ctx, GetPoliciesQuery, accountId)
c, err := pg.Pool.Acquire(ctx)
if err != nil {
return nil, err
}
return GetPoliciesViaConnection(ctx, c, accountId)
}
func GetPoliciesViaConnection(ctx context.Context, con *pgxpool.Conn, accountId string) ([]nmdata.Policy, error) {
rows, err := con.Query(ctx, GetPoliciesQuery, accountId)
if err != nil {
return nil, err
}

View File

@@ -6,6 +6,7 @@ import (
"reflect"
"github.com/jackc/pgx/v5"
"github.com/jackc/pgx/v5/pgxpool"
networkmapdb "github.com/netbirdio/netbird/management/internals/network_map_db"
"github.com/netbirdio/netbird/shared/management/networkmap/nmdata"
)
@@ -19,7 +20,15 @@ const (
)
func (pg *PgStore) GetPostureChecks(ctx context.Context, accountId string) ([]nmdata.PostureChecks, error) {
rows, err := pg.Pool.Query(ctx, GetPostureChecksQuery, accountId)
c, err := pg.Pool.Acquire(ctx)
if err != nil {
return nil, err
}
return GetPostureChecksViaConnection(ctx, c, accountId)
}
func GetPostureChecksViaConnection(ctx context.Context, con *pgxpool.Conn, accountId string) ([]nmdata.PostureChecks, error) {
rows, err := con.Query(ctx, GetPostureChecksQuery, accountId)
if err != nil {
return nil, err
}

View File

@@ -7,6 +7,7 @@ import (
"reflect"
"github.com/jackc/pgx/v5"
"github.com/jackc/pgx/v5/pgxpool"
networkmapdb "github.com/netbirdio/netbird/management/internals/network_map_db"
"github.com/netbirdio/netbird/shared/management/networkmap/nmdata"
)
@@ -22,7 +23,15 @@ const (
)
func (pg *PgStore) GetRoutes(ctx context.Context, accountId string) ([]nmdata.Route, error) {
rows, err := pg.Pool.Query(ctx, GetRoutesQuery, accountId)
c, err := pg.Pool.Acquire(ctx)
if err != nil {
return nil, err
}
return GetRoutesViaConnection(ctx, c, accountId)
}
func GetRoutesViaConnection(ctx context.Context, con *pgxpool.Conn, accountId string) ([]nmdata.Route, error) {
rows, err := con.Query(ctx, GetRoutesQuery, accountId)
if err != nil {
return nil, err
}