diff --git a/management/internals/network_map_db/pgsql/account_settings.go b/management/internals/network_map_db/pgsql/account_settings.go index 620a8b600..5aebbac0d 100644 --- a/management/internals/network_map_db/pgsql/account_settings.go +++ b/management/internals/network_map_db/pgsql/account_settings.go @@ -6,7 +6,6 @@ import ( "time" "github.com/jackc/pgx/v5" - "github.com/jackc/pgx/v5/pgxpool" "github.com/netbirdio/netbird/shared/management/networkmap/nmdata" ) @@ -26,11 +25,11 @@ func (pg *PgStore) GetAccountSettings(ctx context.Context, accountId string) (nm if err != nil { return nmdata.AccountSettingsInfo{}, err } - return GetAccountSettingsViaConnection(ctx, c, accountId) + return GetAccountSettingsViaPgxConnection(ctx, c.Conn(), accountId) } -func GetAccountSettingsViaConnection(ctx context.Context, con *pgxpool.Conn, accountId string) (nmdata.AccountSettingsInfo, error) { +func GetAccountSettingsViaPgxConnection(ctx context.Context, con *pgx.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 b0593fdda..f22cc2569 100644 --- a/management/internals/network_map_db/pgsql/dns.go +++ b/management/internals/network_map_db/pgsql/dns.go @@ -8,7 +8,6 @@ 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" @@ -31,10 +30,10 @@ func (pg *PgStore) GetAccountZones(ctx context.Context, accountId string) ([]nmd if err != nil { return nil, err } - return GetAccountZonesViaConnection(ctx, c, accountId) + return GetAccountZonesViaPgxConnection(ctx, c.Conn(), accountId) } -func GetAccountZonesViaConnection(ctx context.Context, conn *pgxpool.Conn, accountId string) ([]nmdata.CustomZone, error) { +func GetAccountZonesViaPgxConnection(ctx context.Context, conn *pgx.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 9d849af6c..5fa42b22f 100644 --- a/management/internals/network_map_db/pgsql/group.go +++ b/management/internals/network_map_db/pgsql/group.go @@ -7,7 +7,6 @@ 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" ) @@ -29,10 +28,10 @@ func (pg *PgStore) GetGroups(ctx context.Context, accountId string) ([]nmdata.Gr if err != nil { return nil, err } - return GetGroupsViaConnection(ctx, c, accountId) + return GetGroupsViaPgxConnection(ctx, c.Conn(), accountId) } -func GetGroupsViaConnection(ctx context.Context, con *pgxpool.Conn, accountId string) ([]nmdata.Group, error) { +func GetGroupsViaPgxConnection(ctx context.Context, con *pgx.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 236430436..718285153 100644 --- a/management/internals/network_map_db/pgsql/nameserver.go +++ b/management/internals/network_map_db/pgsql/nameserver.go @@ -7,7 +7,6 @@ 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" ) @@ -25,10 +24,10 @@ func (pg *PgStore) GetNameServerGroups(ctx context.Context, accountId string) ([ if err != nil { return nil, err } - return GetNameServerGroupsViaConnection(ctx, c, accountId) + return GetNameServerGroupsViaPgxConnection(ctx, c.Conn(), accountId) } -func GetNameServerGroupsViaConnection(ctx context.Context, con *pgxpool.Conn, accountId string) ([]nmdata.NameServerGroup, error) { +func GetNameServerGroupsViaPgxConnection(ctx context.Context, con *pgx.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 ad46d2a80..a5d61e3ee 100644 --- a/management/internals/network_map_db/pgsql/network.go +++ b/management/internals/network_map_db/pgsql/network.go @@ -7,7 +7,6 @@ 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" ) @@ -25,10 +24,10 @@ func (pg *PgStore) GetNetwork(ctx context.Context, accountId string) (nmdata.Net if err != nil { return nmdata.Network{}, err } - return GetNetworkViaConnection(ctx, c, accountId) + return GetNetworkViaPgxConnection(ctx, c.Conn(), accountId) } -func GetNetworkViaConnection(ctx context.Context, con *pgxpool.Conn, accountId string) (nmdata.Network, error) { +func GetNetworkViaPgxConnection(ctx context.Context, con *pgx.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_map_data.go b/management/internals/network_map_db/pgsql/network_map_data.go new file mode 100644 index 000000000..76efeea1e --- /dev/null +++ b/management/internals/network_map_db/pgsql/network_map_data.go @@ -0,0 +1,102 @@ +package networkmap_pgsql + +import ( + "context" + + "github.com/jackc/pgx/v5" + "github.com/netbirdio/netbird/shared/management/networkmap" + "github.com/netbirdio/netbird/shared/management/networkmap/nmdata" +) + +func (pg *PgStore) GetNetworkMapData(ctx context.Context, accountId string) (*networkmap.NetworkMapData, error) { + tx, err := pg.Pool.BeginTx(ctx, pgx.TxOptions{IsoLevel: pgx.RepeatableRead, AccessMode: pgx.ReadOnly}) + if err != nil { + return nil, err + } + + // acctSettings, err := GetAccountSettingsViaPgxConnection(ctx, tx.Conn(), accountId) + // if err != nil { + // return rollbackAndReturnError(ctx, tx, err) + // } + // dnsZones, err := GetAccountZonesViaPgxConnection(ctx, tx.Conn(), accountId) + // if err != nil { + // return rollbackAndReturnError(ctx, tx, err) + // } + groups, err := GetGroupsViaPgxConnection(ctx, tx.Conn(), accountId) + if err != nil { + return rollbackAndReturnError(ctx, tx, err) + } + nsGroups, err := GetNameServerGroupsViaPgxConnection(ctx, tx.Conn(), accountId) + if err != nil { + return rollbackAndReturnError(ctx, tx, err) + } + networkResources, err := GetNetworkResourcesViaPgxConnection(ctx, tx.Conn(), accountId) + if err != nil { + return rollbackAndReturnError(ctx, tx, err) + } + // routers, err := GetNetworkRoutersViaPgxConnection(ctx, tx.Conn(), accountId) + // if err != nil { + // return rollbackAndReturnError(ctx, tx, err) + // } + network, err := GetNetworkViaPgxConnection(ctx, tx.Conn(), accountId) + if err != nil { + return rollbackAndReturnError(ctx, tx, err) + } + peers, err := GetPeersViaPgxConnection(ctx, tx.Conn(), accountId) + if err != nil { + return rollbackAndReturnError(ctx, tx, err) + } + policies, err := GetPoliciesViaPgxConnection(ctx, tx.Conn(), accountId) + if err != nil { + return rollbackAndReturnError(ctx, tx, err) + } + postureChecks, err := GetPostureChecksViaPgxConnection(ctx, tx.Conn(), accountId) + if err != nil { + return rollbackAndReturnError(ctx, tx, err) + } + routes, err := GetRoutesViaPgxConnection(ctx, tx.Conn(), accountId) + if err != nil { + return rollbackAndReturnError(ctx, tx, err) + } + + err = tx.Commit(ctx) + if err != nil { + // TODO log and ignore? + } + + toret := networkmap.NetworkMapData{ + Network: &network, + Peers: toMap(peers, func(p nmdata.Peer) string { return p.ID }), + Groups: toMap(groups, func(g nmdata.Group) string { return g.PublicID }), + Policies: toSliceOfPtrs(policies), + Routes: toSliceOfPtrs(routes), + NameServerGroups: toSliceOfPtrs(nsGroups), + NetworkResources: toSliceOfPtrs(networkResources), + PostureChecks: toMap(postureChecks, func(pc nmdata.PostureChecks) string { return pc.ID }), + } + + return &toret, nil +} + +func rollbackAndReturnError(ctx context.Context, tx pgx.Tx, err error) (*networkmap.NetworkMapData, error) { + if errr := tx.Rollback(ctx); errr != nil { + // TODO log and ignore? + } + return nil, err +} + +func toMap[T any](all []T, id func(t T) string) map[string]*T { + toret := make(map[string]*T, len(all)) + for _, t := range all { + toret[id(t)] = &t + } + return toret +} + +func toSliceOfPtrs[T any](all []T) []*T { + toret := make([]*T, len(all)) + for _, t := range all { + toret = append(toret, &t) + } + return toret +} diff --git a/management/internals/network_map_db/pgsql/network_resource.go b/management/internals/network_map_db/pgsql/network_resource.go index 9badbcd27..cf08d35a7 100644 --- a/management/internals/network_map_db/pgsql/network_resource.go +++ b/management/internals/network_map_db/pgsql/network_resource.go @@ -7,7 +7,6 @@ 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" ) @@ -25,10 +24,10 @@ func (pg *PgStore) GetNetworkResources(ctx context.Context, accountId string) ([ if err != nil { return nil, err } - return GetNetworkResourcesViaConnection(ctx, c, accountId) + return GetNetworkResourcesViaPgxConnection(ctx, c.Conn(), accountId) } -func GetNetworkResourcesViaConnection(ctx context.Context, con *pgxpool.Conn, accountId string) ([]nmdata.NetworkResource, error) { +func GetNetworkResourcesViaPgxConnection(ctx context.Context, con *pgx.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 234b96de8..517e55856 100644 --- a/management/internals/network_map_db/pgsql/network_router.go +++ b/management/internals/network_map_db/pgsql/network_router.go @@ -7,7 +7,6 @@ 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" ) @@ -25,10 +24,10 @@ func (pg *PgStore) GetNetworkRouters(ctx context.Context, accountId string) ([]n if err != nil { return nil, err } - return GetNetworkRoutersViaConnection(ctx, c, accountId) + return GetNetworkRoutersViaPgxConnection(ctx, c.Conn(), accountId) } -func GetNetworkRoutersViaConnection(ctx context.Context, con *pgxpool.Conn, accountId string) ([]nmdata.NetworkRouter, error) { +func GetNetworkRoutersViaPgxConnection(ctx context.Context, con *pgx.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 4736337d9..40bafd3d5 100644 --- a/management/internals/network_map_db/pgsql/peer.go +++ b/management/internals/network_map_db/pgsql/peer.go @@ -7,7 +7,6 @@ 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" ) @@ -27,10 +26,10 @@ func (pg *PgStore) GetPeers(ctx context.Context, accountId string) ([]nmdata.Pee if err != nil { return nil, err } - return GetPeersViaConnection(ctx, c, accountId) + return GetPeersViaPgxConnection(ctx, c.Conn(), accountId) } -func GetPeersViaConnection(ctx context.Context, con *pgxpool.Conn, accountId string) ([]nmdata.Peer, error) { +func GetPeersViaPgxConnection(ctx context.Context, con *pgx.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 3a0b8735e..81e889d15 100644 --- a/management/internals/network_map_db/pgsql/policy.go +++ b/management/internals/network_map_db/pgsql/policy.go @@ -7,7 +7,6 @@ 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" ) @@ -28,10 +27,10 @@ func (pg *PgStore) GetPolicies(ctx context.Context, accountId string) ([]nmdata. if err != nil { return nil, err } - return GetPoliciesViaConnection(ctx, c, accountId) + return GetPoliciesViaPgxConnection(ctx, c.Conn(), accountId) } -func GetPoliciesViaConnection(ctx context.Context, con *pgxpool.Conn, accountId string) ([]nmdata.Policy, error) { +func GetPoliciesViaPgxConnection(ctx context.Context, con *pgx.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 a9c6d8e2b..7b3c97ae7 100644 --- a/management/internals/network_map_db/pgsql/posture.go +++ b/management/internals/network_map_db/pgsql/posture.go @@ -6,14 +6,13 @@ 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" ) const ( GetPostureChecksQuery = ` - select id, checks + select public_id as id, checks from posture_checks where account_id=$1 ` @@ -24,10 +23,10 @@ func (pg *PgStore) GetPostureChecks(ctx context.Context, accountId string) ([]nm if err != nil { return nil, err } - return GetPostureChecksViaConnection(ctx, c, accountId) + return GetPostureChecksViaPgxConnection(ctx, c.Conn(), accountId) } -func GetPostureChecksViaConnection(ctx context.Context, con *pgxpool.Conn, accountId string) ([]nmdata.PostureChecks, error) { +func GetPostureChecksViaPgxConnection(ctx context.Context, con *pgx.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 bb07e7dc7..60752b50f 100644 --- a/management/internals/network_map_db/pgsql/route.go +++ b/management/internals/network_map_db/pgsql/route.go @@ -7,7 +7,6 @@ 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" ) @@ -27,10 +26,10 @@ func (pg *PgStore) GetRoutes(ctx context.Context, accountId string) ([]nmdata.Ro if err != nil { return nil, err } - return GetRoutesViaConnection(ctx, c, accountId) + return GetRoutesViaPgxConnection(ctx, c.Conn(), accountId) } -func GetRoutesViaConnection(ctx context.Context, con *pgxpool.Conn, accountId string) ([]nmdata.Route, error) { +func GetRoutesViaPgxConnection(ctx context.Context, con *pgx.Conn, accountId string) ([]nmdata.Route, error) { rows, err := con.Query(ctx, GetRoutesQuery, accountId) if err != nil { return nil, err