From ea1b4d56e864610e49061cc311e0b8b04ff2ac7e Mon Sep 17 00:00:00 2001 From: Dmitri Dolguikh Date: Fri, 31 Jul 2026 15:03:15 +0200 Subject: [PATCH] added support for queries via pgx connection Signed-off-by: Dmitri Dolguikh --- management/internals/network_map_db/db_store.go | 2 ++ .../network_map_db/pgsql/account_settings.go | 12 +++++++++++- management/internals/network_map_db/pgsql/dns.go | 11 ++++++++++- management/internals/network_map_db/pgsql/group.go | 11 ++++++++++- .../internals/network_map_db/pgsql/nameserver.go | 11 ++++++++++- management/internals/network_map_db/pgsql/network.go | 11 ++++++++++- .../network_map_db/pgsql/network_resource.go | 11 ++++++++++- .../internals/network_map_db/pgsql/network_router.go | 11 ++++++++++- management/internals/network_map_db/pgsql/peer.go | 11 ++++++++++- management/internals/network_map_db/pgsql/policy.go | 11 ++++++++++- management/internals/network_map_db/pgsql/posture.go | 11 ++++++++++- management/internals/network_map_db/pgsql/route.go | 11 ++++++++++- 12 files changed, 113 insertions(+), 11 deletions(-) diff --git a/management/internals/network_map_db/db_store.go b/management/internals/network_map_db/db_store.go index 5efa469f7..ff2ee03d0 100644 --- a/management/internals/network_map_db/db_store.go +++ b/management/internals/network_map_db/db_store.go @@ -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 { diff --git a/management/internals/network_map_db/pgsql/account_settings.go b/management/internals/network_map_db/pgsql/account_settings.go index 6f4713b96..620a8b600 100644 --- a/management/internals/network_map_db/pgsql/account_settings.go +++ b/management/internals/network_map_db/pgsql/account_settings.go @@ -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 } diff --git a/management/internals/network_map_db/pgsql/dns.go b/management/internals/network_map_db/pgsql/dns.go index cdbcaf298..b0593fdda 100644 --- a/management/internals/network_map_db/pgsql/dns.go +++ b/management/internals/network_map_db/pgsql/dns.go @@ -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 } diff --git a/management/internals/network_map_db/pgsql/group.go b/management/internals/network_map_db/pgsql/group.go index 12386aec3..9d849af6c 100644 --- a/management/internals/network_map_db/pgsql/group.go +++ b/management/internals/network_map_db/pgsql/group.go @@ -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 } diff --git a/management/internals/network_map_db/pgsql/nameserver.go b/management/internals/network_map_db/pgsql/nameserver.go index 57761fceb..236430436 100644 --- a/management/internals/network_map_db/pgsql/nameserver.go +++ b/management/internals/network_map_db/pgsql/nameserver.go @@ -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 } diff --git a/management/internals/network_map_db/pgsql/network.go b/management/internals/network_map_db/pgsql/network.go index 33786b49e..ad46d2a80 100644 --- a/management/internals/network_map_db/pgsql/network.go +++ b/management/internals/network_map_db/pgsql/network.go @@ -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 } diff --git a/management/internals/network_map_db/pgsql/network_resource.go b/management/internals/network_map_db/pgsql/network_resource.go index 1a4073015..9badbcd27 100644 --- a/management/internals/network_map_db/pgsql/network_resource.go +++ b/management/internals/network_map_db/pgsql/network_resource.go @@ -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 } diff --git a/management/internals/network_map_db/pgsql/network_router.go b/management/internals/network_map_db/pgsql/network_router.go index 1415ac6e3..234b96de8 100644 --- a/management/internals/network_map_db/pgsql/network_router.go +++ b/management/internals/network_map_db/pgsql/network_router.go @@ -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 } diff --git a/management/internals/network_map_db/pgsql/peer.go b/management/internals/network_map_db/pgsql/peer.go index 38b2a800c..4736337d9 100644 --- a/management/internals/network_map_db/pgsql/peer.go +++ b/management/internals/network_map_db/pgsql/peer.go @@ -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 } diff --git a/management/internals/network_map_db/pgsql/policy.go b/management/internals/network_map_db/pgsql/policy.go index 35cca1ee8..3a0b8735e 100644 --- a/management/internals/network_map_db/pgsql/policy.go +++ b/management/internals/network_map_db/pgsql/policy.go @@ -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 } diff --git a/management/internals/network_map_db/pgsql/posture.go b/management/internals/network_map_db/pgsql/posture.go index 23010e527..a9c6d8e2b 100644 --- a/management/internals/network_map_db/pgsql/posture.go +++ b/management/internals/network_map_db/pgsql/posture.go @@ -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 } diff --git a/management/internals/network_map_db/pgsql/route.go b/management/internals/network_map_db/pgsql/route.go index ecc176aaf..bb07e7dc7 100644 --- a/management/internals/network_map_db/pgsql/route.go +++ b/management/internals/network_map_db/pgsql/route.go @@ -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 }