feat: Support IP Usage stats in VPC Prefix and Subnet

api/pkg/api/handler/subnet.go

@@ -25,9 +25,12 @@ import (
 	"fmt"
 	"net/http"
 	"slices"
+	"strconv"
 
 	"github.com/labstack/echo/v4"
 
+	cip "github.com/NVIDIA/infra-controller-rest/ipam"
+
 	"go.opentelemetry.io/otel/attribute"
 	temporalClient "go.temporal.io/sdk/client"
 	tp "go.temporal.io/sdk/temporal"
@@ -399,7 +402,7 @@ func (csh CreateSubnetHandler) Handle(c echo.Context) error {
 	txCommitted = true
 
 	// create response
-	apiInstance := model.NewAPISubnet(subnet, []cdbm.StatusDetail{*ssd})
+	apiInstance := model.NewAPISubnet(subnet, []cdbm.StatusDetail{*ssd}, nil)
 	logger.Info().Msg("finishing API handler")
 	return c.JSON(http.StatusCreated, apiInstance)
 }
@@ -437,6 +440,7 @@ func NewGetAllSubnetHandler(dbSession *cdb.Session, tc temporalClient.Client, cf
 // @Param status query string false "Filter by status" e.g. 'Pending', 'Error'"
 // @Param query query string false "Query input for full text search"
 // @Param includeRelation query string false "Related entities to include in response e.g. 'Site', 'Vpc', 'Tenant', 'IPv4Block', 'IPv6Block'"
+// @Param includeUsageStats query boolean false "Subnet IPv4 usage (interface/instance-derived; same shape as IP Block usage)"
 // @Param pageNumber query integer false "Page number of results returned"
 // @Param pageSize query integer false "Number of results per page"
 // @Param orderBy query string false "Order by field"
@@ -501,6 +505,19 @@ func (gash GetAllSubnetHandler) Handle(c echo.Context) error {
 		return cutil.NewAPIErrorResponse(c, http.StatusBadRequest, errMsg, nil)
 	}
 
+	includeUsageStats := false
+	qius := c.QueryParam("includeUsageStats")
+	if qius != "" {
+		includeUsageStats, err = strconv.ParseBool(qius)
+		if err != nil {
+			return cutil.NewAPIErrorResponse(c, http.StatusBadRequest, "Invalid value specified for `includeUsageStats` query param", nil)
+		}
+	}
+	queryIncludeRelations := slices.Clone(qIncludeRelations)
+	if includeUsageStats && !slices.Contains(queryIncludeRelations, cdbm.IPv4BlockRelationName) {
+		queryIncludeRelations = append(queryIncludeRelations, cdbm.IPv4BlockRelationName)
+	}
+
 	// Get site ID from query param
 	tsDAO := cdbm.NewTenantSiteDAO(gash.dbSession)
 	var siteID *uuid.UUID
@@ -570,12 +587,29 @@ func (gash GetAllSubnetHandler) Handle(c echo.Context) error {
 		Limit:   pageRequest.Limit,
 		Offset:  pageRequest.Offset,
 		OrderBy: pageRequest.OrderBy,
-	}, qIncludeRelations)
+	}, queryIncludeRelations)
 	if err != nil {
 		logger.Error().Err(err).Msg("error getting subnets from db")
 		return cutil.NewAPIErrorResponse(c, http.StatusInternalServerError, "Failed to retrieve Subnets", nil)
 	}
 
+	sbusageMap := map[uuid.UUID]*cip.Usage{}
+	if includeUsageStats {
+		for i := range subnets {
+			sn := &subnets[i]
+			if sn.IPv4Block == nil {
+				logger.Error().Str("subnetId", sn.ID.String()).Msg("Subnet missing IPv4 Block relation for usage stats")
+				continue
+			}
+			prefixUsage, serr := sDAO.GetPrefixUsage(ctx, nil, sn)
+			if serr != nil {
+				logger.Error().Err(serr).Str("subnetId", sn.ID.String()).Msg("error retrieving usage stats for Subnet")
+				continue
+			}
+			sbusageMap[sn.ID] = prefixUsage
+		}
+	}
+
 	// Get status details
 	sdDAO := cdbm.NewStatusDetailDAO(gash.dbSession)
 
@@ -600,7 +634,8 @@ func (gash GetAllSubnetHandler) Handle(c echo.Context) error {
 	// get status details
 	for _, sn := range subnets {
 		cursn := sn
-		apiSubnet := model.NewAPISubnet(&cursn, ssdMap[sn.ID.String()])
+		snusage := sbusageMap[sn.ID]
+		apiSubnet := model.NewAPISubnet(&cursn, ssdMap[sn.ID.String()], snusage)
 		apiSubnets = append(apiSubnets, apiSubnet)
 	}
 
@@ -649,6 +684,7 @@ func NewGetSubnetHandler(dbSession *cdb.Session, tc temporalClient.Client, cfg *
 // @Param org path string true "Name of NGC organization"
 // @Param id path string true "ID of Subnet"
 // @Param includeRelation query string false "Related entities to include in response e.g. 'Site', 'Vpc', 'Tenant', 'IPv4Block', 'IPv6Block'"
+// @Param includeUsageStats query boolean false "Subnet IPv4 usage (interface/instance-derived; same shape as IP Block usage)"
 // @Success 200 {object} model.APISubnet
 // @Router /v2/org/{org}/nico/subnet/{id} [get]
 func (gsh GetSubnetHandler) Handle(c echo.Context) error {
@@ -686,6 +722,20 @@ func (gsh GetSubnetHandler) Handle(c echo.Context) error {
 		return cutil.NewAPIErrorResponse(c, http.StatusBadRequest, errMsg, nil)
 	}
 
+	includeUsageStats := false
+	qius := c.QueryParam("includeUsageStats")
+	if qius != "" {
+		includeUsageStats, err = strconv.ParseBool(qius)
+		if err != nil {
+			return cutil.NewAPIErrorResponse(c, http.StatusBadRequest, "Invalid value specified for `includeUsageStats` query param", nil)
+		}
+	}
+
+	queryIncludeRelations := slices.Clone(qIncludeRelations)
+	if includeUsageStats && !slices.Contains(queryIncludeRelations, cdbm.IPv4BlockRelationName) {
+		queryIncludeRelations = append(queryIncludeRelations, cdbm.IPv4BlockRelationName)
+	}
+
 	// Get subnet ID from URL param
 	sStrID := c.Param("id")
 
@@ -707,7 +757,7 @@ func (gsh GetSubnetHandler) Handle(c echo.Context) error {
 	}
 
 	// Check that subnet exists
-	subnet, err := sDAO.GetByID(ctx, nil, sID, qIncludeRelations)
+	subnet, err := sDAO.GetByID(ctx, nil, sID, queryIncludeRelations)
 	if err != nil {
 		logger.Error().Err(err).Msg("error retrieving Subnet DB entity")
 		if err == cdb.ErrDoesNotExist {
@@ -730,8 +780,21 @@ func (gsh GetSubnetHandler) Handle(c echo.Context) error {
 		return cutil.NewAPIErrorResponse(c, http.StatusInternalServerError, "Failed to retrieve Status Details for subnet", nil)
 	}
 
+	var sbusage *cip.Usage
+	if includeUsageStats {
+		if subnet.IPv4Block == nil {
+			logger.Error().Str("subnetId", subnet.ID.String()).Msg("Subnet missing IPv4 Block relation for usage stats")
+			return cutil.NewAPIErrorResponse(c, http.StatusInternalServerError, "Failed to retrieve Usage Stats for Subnet", nil)
+		}
+		sbusage, err = sDAO.GetPrefixUsage(ctx, nil, subnet)
+		if err != nil {
+			logger.Error().Err(err).Msg("error retrieving usage stats for Subnet")
+			return cutil.NewAPIErrorResponse(c, http.StatusInternalServerError, "Failed to retrieve Usage Stats for Subnet", nil)
+		}
+	}
+
 	// Send response
-	apiInstance := model.NewAPISubnet(subnet, ssds)
+	apiInstance := model.NewAPISubnet(subnet, ssds, sbusage)
 	logger.Info().Msg("finishing API handler")
 	return c.JSON(http.StatusOK, apiInstance)
 }
@@ -890,7 +953,7 @@ func (ush UpdateSubnetHandler) Handle(c echo.Context) error {
 	txCommitted = true
 
 	// Send response
-	apiInstance := model.NewAPISubnet(subnet, ssds)
+	apiInstance := model.NewAPISubnet(subnet, ssds, nil)
 	logger.Info().Msg("finishing API handler")
 	return c.JSON(http.StatusOK, apiInstance)
 }

api/pkg/api/handler/subnet_test.go

@@ -77,6 +77,18 @@ func testSubnetBuildInterface(t *testing.T, dbSession *cdb.Session, instanceID,
 	return is
 }
 
+func testSubnetCIDREntity(t *testing.T, sn *cdbm.Subnet) string {
+	t.Helper()
+	require.NotNil(t, sn)
+	require.NotNil(t, sn.IPv4Prefix)
+	p := strings.TrimSpace(*sn.IPv4Prefix)
+	require.NotEmpty(t, p)
+	if strings.Contains(p, "/") {
+		return p
+	}
+	return fmt.Sprintf("%s/%d", p, sn.PrefixLength)
+}
+
 func testSubnetBuildVpc(t *testing.T, dbSession *cdb.Session, ip *cdbm.InfrastructureProvider, tenant *cdbm.Tenant, site *cdbm.Site, org, name string, networkVtType *string, status string, controllerVpcID *uuid.UUID) *cdbm.Vpc {
 	vpc := &cdbm.Vpc{
 		ID:                        uuid.New(),
@@ -1124,6 +1136,8 @@ func TestSubnetHandler_Get(t *testing.T) {
 	assert.NotNil(t, tenant2)
 	vpc1 := testSubnetBuildVpc(t, dbSession, ip, tenant1, site, tnOrg1, "testVPC", cdb.GetStrPtr(cdbm.VpcEthernetVirtualizer), cdbm.VpcStatusReady, cdb.GetUUIDPtr(uuid.New()))
 
+	_ = common.TestBuildTenantSite(t, dbSession, tenant1, site, tnu)
+
 	cfg := common.GetTestConfig()
 	tempClient := &tmocks.Client{}
 	ipamStorage := ipam.NewIpamStorage(dbSession.DB, nil)
@@ -1148,6 +1162,35 @@ func TestSubnetHandler_Get(t *testing.T) {
 
 	subnet := testCreateSubnet(t, dbSession, scp, ipamStorage, tnu, tnOrg1, string(parentIpbBody))
 
+	ifaceSubnetBody, err := json.Marshal(model.APISubnetCreateRequest{
+		Name: "iface-subnet-usage", Description: cdb.GetStrPtr(""), VpcID: vpc1.ID.String(),
+		IPv4BlockID: cdb.GetStrPtr(ipb1.ID.String()), PrefixLength: prefixLen})
+	require.NoError(t, err)
+	subnetWithIfaceWorkload := testCreateSubnet(t, dbSession, scp, ipamStorage, tnu, tnOrg1, string(ifaceSubnetBody))
+	var wantUsageAcquiredPrefixes0 uint64
+	wantUsageAcquiredPrefixes1 := uint64(0)
+
+	alWorkload := common.TestBuildAllocation(t, dbSession, site, tenant1, "get-subnet-usage-iface-alloc", ipu)
+	itWorkload := common.TestBuildInstanceType(t, dbSession, "get-subnet-iface-it", cdb.GetUUIDPtr(uuid.New()), site, nil, ipu)
+	common.TestBuildAllocationConstraint(t, dbSession, alWorkload, itWorkload, nil, 5, ipu)
+	mWorkload := common.TestBuildMachine(t, dbSession, ip, site, &itWorkload.ID, cdb.GetStrPtr("get-subnet-iface-mt"), cdbm.MachineStatusReady)
+	_ = common.TestBuildMachineInstanceType(t, dbSession, mWorkload, itWorkload)
+	osWorkload := common.TestBuildOperatingSystem(t, dbSession, "get-subnet-iface-os", tenant1, cdbm.OperatingSystemStatusReady, tnu)
+
+	snWorkloadUUID := uuid.MustParse(subnetWithIfaceWorkload.ID)
+	snEnt, err := cdbm.NewSubnetDAO(dbSession).GetByID(ctx, nil, snWorkloadUUID, nil)
+	require.NoError(t, err)
+	cidrWorkload := testSubnetCIDREntity(t, snEnt)
+	chosenIfaceIPv4 := testIPv4NthUsableInCIDR(t, cidrWorkload, 20)
+
+	instWorkload := common.TestBuildInstance(t, dbSession, "get-subnet-iface-inst", tenant1.ID, ip.ID, site.ID, itWorkload.ID, vpc1.ID, cdb.GetStrPtr(mWorkload.ID), osWorkload.ID)
+	ifWorkload := common.TestBuildInterface(t, dbSession, instWorkload.ID, &snWorkloadUUID, nil, true, nil, nil, nil, cdb.GetStrPtr(cdbm.InterfaceStatusReady), tnu)
+	_, err = cdbm.NewInterfaceDAO(dbSession).Update(ctx, nil, cdbm.InterfaceUpdateInput{
+		InterfaceID: ifWorkload.ID,
+		IpAddresses: []string{chosenIfaceIPv4},
+	})
+	require.NoError(t, err)
+
 	// OTEL Spanner configuration
 	tracer, _, ctx := common.TestCommonTraceProviderSetup(t, ctx)
 
@@ -1161,9 +1204,12 @@ func TestSubnetHandler_Get(t *testing.T) {
 		queryIncludeRelations1 *string
 		queryIncludeRelations2 *string
 		queryIncludeRelations3 *string
+		queryIncludeUsageStats *string
 		expectedVpcName        *string
 		expectetIPv4Name       *string
 		expectedIPv6Name       *string
+		expectUsageStatsNonNil bool
+		wantAcquiredPrefixes   *uint64
 		verifyChildSpanner     bool
 	}{
 		{
@@ -1243,6 +1289,37 @@ func TestSubnetHandler_Get(t *testing.T) {
 			queryIncludeRelations1: cdb.GetStrPtr(cdbm.IPv4BlockRelationName),
 			expectetIPv4Name:       &ipb1.Name,
 		},
+		{
+			name:                   "error when includeUsageStats query invalid",
+			reqOrgName:             tnOrg1,
+			user:                   tnu,
+			id:                     subnet.ID,
+			expectedErr:            true,
+			expectedStatus:         http.StatusBadRequest,
+			queryIncludeUsageStats: cdb.GetStrPtr("not-a-bool"),
+		},
+		{
+			name:                   "success when includeUsageStats true and no ethernet interfaces",
+			reqOrgName:             tnOrg1,
+			user:                   tnu,
+			id:                     subnet.ID,
+			expectedErr:            false,
+			expectedStatus:         http.StatusOK,
+			queryIncludeUsageStats: cdb.GetStrPtr("true"),
+			expectUsageStatsNonNil: true,
+			wantAcquiredPrefixes:   &wantUsageAcquiredPrefixes0,
+		},
+		{
+			name:                   "success when includeUsageStats true with ethernet interface and IPv4 IPs",
+			reqOrgName:             tnOrg1,
+			user:                   tnu,
+			id:                     subnetWithIfaceWorkload.ID,
+			expectedErr:            false,
+			expectedStatus:         http.StatusOK,
+			queryIncludeUsageStats: cdb.GetStrPtr("true"),
+			expectUsageStatsNonNil: true,
+			wantAcquiredPrefixes:   &wantUsageAcquiredPrefixes1,
+		},
 	}
 	for _, tc := range tests {
 		t.Run(tc.name, func(t *testing.T) {
@@ -1251,6 +1328,9 @@ func TestSubnetHandler_Get(t *testing.T) {
 			e := echo.New()
 
 			q := url.Values{}
+			if tc.queryIncludeUsageStats != nil {
+				q.Set("includeUsageStats", *tc.queryIncludeUsageStats)
+			}
 			if tc.queryIncludeRelations1 != nil {
 				q.Add("includeRelation", *tc.queryIncludeRelations1)
 			}
@@ -1294,7 +1374,9 @@ func TestSubnetHandler_Get(t *testing.T) {
 
 				assert.Equal(t, cdbm.IPBlockRoutingTypeDatacenterOnly, *rsp.RoutingType)
 
-				if tc.queryIncludeRelations1 != nil || tc.queryIncludeRelations2 != nil || tc.queryIncludeRelations3 != nil {
+				hasRelations := tc.queryIncludeRelations1 != nil || tc.queryIncludeRelations2 != nil || tc.queryIncludeRelations3 != nil
+				hasUsageStats := tc.expectUsageStatsNonNil
+				if hasRelations || hasUsageStats {
 					if tc.expectedVpcName != nil {
 						assert.Equal(t, *tc.expectedVpcName, rsp.Vpc.Name)
 					}
@@ -1304,11 +1386,23 @@ func TestSubnetHandler_Get(t *testing.T) {
 					if tc.expectedIPv6Name != nil {
 						assert.Equal(t, *tc.expectedIPv6Name, rsp.IPv6Block.Name)
 					}
+					if hasUsageStats {
+						require.NotNil(t, rsp.IPv4Block)
+					}
 				} else {
 					assert.Nil(t, rsp.Vpc)
 					assert.Nil(t, rsp.IPv4Block)
 					assert.Nil(t, rsp.IPv6Block)
 				}
+				if tc.expectUsageStatsNonNil {
+					require.NotNil(t, rsp.UsageStats)
+				} else {
+					assert.Nil(t, rsp.UsageStats)
+				}
+				if tc.wantAcquiredPrefixes != nil {
+					require.NotNil(t, rsp.UsageStats)
+					assert.Equal(t, *tc.wantAcquiredPrefixes, rsp.UsageStats.AcquiredPrefixes)
+				}
 			}
 
 			if tc.verifyChildSpanner {