From b02736adc316de8a2f37172039948bc25220b1bb Mon Sep 17 00:00:00 2001 From: Dmitri Dolguikh Date: Sat, 1 Aug 2026 11:54:14 +0200 Subject: [PATCH] cleanup network routers retrieval Signed-off-by: Dmitri Dolguikh --- .../internals/network_map_db/db_store.go | 2 +- .../network_map_db/pgsql/group_test.go | 3 +- .../network_map_db/pgsql/network_map_data.go | 44 +-------- .../network_map_db/pgsql/network_router.go | 92 ++++++++++++------- 4 files changed, 64 insertions(+), 77 deletions(-) diff --git a/management/internals/network_map_db/db_store.go b/management/internals/network_map_db/db_store.go index ff2ee03d0..55a6798c4 100644 --- a/management/internals/network_map_db/db_store.go +++ b/management/internals/network_map_db/db_store.go @@ -25,7 +25,7 @@ type NetworkMapDBStore interface { GetRoutes(ctx context.Context, accountId string) ([]nmdata.Route, error) GetNameServerGroups(ctx context.Context, accountId string) ([]nmdata.NameServerGroup, error) GetNetworkResources(ctx context.Context, accountId string) ([]nmdata.NetworkResource, error) - GetNetworkRouters(ctx context.Context, accountId string) ([]nmdata.NetworkRouter, error) + GetNetworkRouters(ctx context.Context, accountId string) (map[string]map[string]*nmdata.NetworkRouter, error) GetNetwork(ctx context.Context, accountId string) (nmdata.Network, error) GetAccountZones(ctx context.Context, accountId string) ([]nmdata.CustomZone, error) GetAccountSettings(ctx context.Context, accountId string) (nmdata.AccountSettingsInfo, error) diff --git a/management/internals/network_map_db/pgsql/group_test.go b/management/internals/network_map_db/pgsql/group_test.go index 79f0ef00d..3bc75c10c 100644 --- a/management/internals/network_map_db/pgsql/group_test.go +++ b/management/internals/network_map_db/pgsql/group_test.go @@ -155,7 +155,8 @@ func TestGetNetworkRouters(t *testing.T) { // err = loadSQL(ctx, s.pool, initDb) //assert.NoError(t, err) - res, err := s.GetNetworkRouters(ctx, "d29f99jl0ubs73cm8ce0") //"ckd7ee2fic3c73dtendg") + c, _ := s.Pool.Acquire(ctx) + res, err := GetNetworkRoutersViaPgxConnection(ctx, c.Conn(), "d29f99jl0ubs73cm8ce0") //"ckd7ee2fic3c73dtendg") assert.NoError(t, err) fmt.Print(res) diff --git a/management/internals/network_map_db/pgsql/network_map_data.go b/management/internals/network_map_db/pgsql/network_map_data.go index 6a493f143..34a23d339 100644 --- a/management/internals/network_map_db/pgsql/network_map_data.go +++ b/management/internals/network_map_db/pgsql/network_map_data.go @@ -2,11 +2,8 @@ package networkmap_pgsql import ( "context" - "fmt" - "reflect" "github.com/jackc/pgx/v5" - 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" ) @@ -37,7 +34,6 @@ func (pg *PgStore) GetNetworkMapData(ctx context.Context, accountId string) (*ne if err != nil { return rollbackAndReturnError(ctx, tx, err) } - // TODO (dmitri) this needs cleaning up -- returns an internal struct routers, err := GetNetworkRoutersViaPgxConnection(ctx, tx.Conn(), accountId) if err != nil { return rollbackAndReturnError(ctx, tx, err) @@ -78,49 +74,13 @@ func (pg *PgStore) GetNetworkMapData(ctx context.Context, accountId string) (*ne Groups: toMap(groups, func(g nmdata.Group) string { return g.PublicID }), Policies: toSliceOfPtrs(policies), Routes: toSliceOfPtrs(routes), + Routers: routers, NameServerGroups: toSliceOfPtrs(nsGroups), NetworkResources: toSliceOfPtrs(networkResources), PostureChecks: toMap(postureChecks, func(pc nmdata.PostureChecks) string { return pc.ID }), - NetworkXIDToPublicID: networkXIDToPublicID, // TODO (dmitri-d) do we still need it now? + NetworkXIDToPublicID: networkXIDToPublicID, // TODO (dmitri) maybe we can switch to public ids everywhere? } - networktoRouters := make(map[string]map[string]*nmdata.NetworkRouter) - for _, router := range routers { - if !router.Enabled.Bool { - continue - } - - if router.NetworkID.String == "" { - return nil, fmt.Errorf("router with public_id %s doesn't have network_id set", router.PublicID.String) - } - networkPublicId := networkXIDToPublicID[router.NetworkID.String] - if networkPublicId == "" { - return nil, fmt.Errorf("network with id %s has no public_id", router.NetworkID.String) - } - - nmdatarouter := nmdata.NetworkRouter{} - err := networkmapdb.FromSqlTypesToSharedTypes(reflect.ValueOf(&router), reflect.ValueOf(&nmdatarouter)) - if err != nil { - return nil, err - } - - if networktoRouters[networkPublicId] == nil { - networktoRouters[networkPublicId] = make(map[string]*nmdata.NetworkRouter) - } - if router.Peer.String != "" { - networktoRouters[networkPublicId][router.Peer.String] = &nmdatarouter - } - for _, peerGroup := range nmdatarouter.PeerGroups { - g := toret.Groups[peerGroup] - if g != nil { - for _, peerID := range g.Peers { - networktoRouters[networkPublicId][peerID] = &nmdatarouter - } - } - } - } - toret.Routers = networktoRouters - return &toret, nil } diff --git a/management/internals/network_map_db/pgsql/network_router.go b/management/internals/network_map_db/pgsql/network_router.go index 9c0fc1445..8d6b594fc 100644 --- a/management/internals/network_map_db/pgsql/network_router.go +++ b/management/internals/network_map_db/pgsql/network_router.go @@ -3,57 +3,83 @@ package networkmap_pgsql import ( "context" "database/sql" - "encoding/json" + "fmt" + "reflect" "github.com/jackc/pgx/v5" + networkmapdb "github.com/netbirdio/netbird/management/internals/network_map_db" + "github.com/netbirdio/netbird/shared/management/networkmap/nmdata" ) const ( GetNetworkRouterQuery = ` - select public_id, peer, peer_groups, network_id, masquerade, metric, enabled, + select public_id, peer, network_id, masquerade, metric, enabled, + ( + select array_agg(group_peers.peer_id) + from group_peers + where group_peers.group_id in (select json_array_elements_text(peer_groups::json)) + ) as peers_via_groups from network_routers where account_id=$1 ` ) -// func (pg *PgStore) GetNetworkRouters(ctx context.Context, accountId string) ([]nmdata.NetworkRouter, error) { -// c, err := pg.Pool.Acquire(ctx) -// if err != nil { -// return nil, err -// } -// return GetNetworkRoutersViaPgxConnection(ctx, c.Conn(), accountId) -// } +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) ([]networkrouter, error) { +func GetNetworkRoutersViaPgxConnection(ctx context.Context, con *pgx.Conn, accountId string) (map[string]map[string]*nmdata.NetworkRouter, error) { rows, err := con.Query(ctx, GetNetworkRouterQuery, accountId) if err != nil { return nil, err } - return pgx.CollectRows(rows, pgx.RowToStructByName[networkrouter]) - // if err != nil { - // return nil, err - // } - // - // toret := make([]nmdata.NetworkRouter, 0, len(netrouters)) - // for _, nrt := range netrouters { - // router := nmdata.NetworkRouter{} - // err := networkmapdb.FromSqlTypesToSharedTypes( - // reflect.ValueOf(&nrt), reflect.ValueOf(&router)) - // if err != nil { - // return nil, err - // } - // toret = append(toret, router) - // } - // return toret, nil + routers, err := pgx.CollectRows(rows, pgx.RowToStructByName[networkrouter]) + if err != nil { + return nil, err + } + + toret := make(map[string]map[string]*nmdata.NetworkRouter) + for _, router := range routers { + if !router.Enabled.Bool { + continue + } + + networkId := router.NetworkID.String + if networkId == "" { + return nil, fmt.Errorf("router with public_id %s doesn't have network_id set", router.PublicID.String) + } + + nmdatarouter := nmdata.NetworkRouter{} + err := networkmapdb.FromSqlTypesToSharedTypes(reflect.ValueOf(&router), reflect.ValueOf(&nmdatarouter)) + if err != nil { + return nil, err + } + + if toret[networkId] == nil { + toret[networkId] = make(map[string]*nmdata.NetworkRouter) + } + if router.Peer.String != "" { + toret[networkId][router.Peer.String] = &nmdatarouter + } + for _, peerId := range router.PeersViaGroups { + toret[networkId][peerId] = &nmdatarouter + } + } + + return toret, nil } type networkrouter struct { - PublicID sql.NullString - NetworkID sql.NullString `nmap:"skip"` - Peer sql.NullString `nmap:"skip"` - PeerGroups json.RawMessage - Masquerade sql.NullBool - Metric sql.NullInt64 - Enabled sql.NullBool + PublicID sql.NullString + NetworkID sql.NullString `nmap:"skip"` + Peer sql.NullString `nmap:"skip"` + PeersViaGroups []string `nmap:"skip"` + Masquerade sql.NullBool + Metric sql.NullInt64 + Enabled sql.NullBool }