diff --git a/rest-api/api/pkg/api/handler/allocation.go b/rest-api/api/pkg/api/handler/allocation.go index 87309e4717..9ff9ff93e1 100644 --- a/rest-api/api/pkg/api/handler/allocation.go +++ b/rest-api/api/pkg/api/handler/allocation.go @@ -1470,30 +1470,39 @@ func (dah DeleteAllocationHandler) Handle(c echo.Context) error { return cutil.NewAPIError(http.StatusInternalServerError, "Error deleting Allocation, DB error retrieving remaining Allocations for Tenant", nil) } if acount == 0 { - tsDAO := cdbm.NewTenantSiteDAO(dah.dbSession) - tss, tscount, serr := tsDAO.GetAll( - ctx, - tx, - cdbm.TenantSiteFilterInput{ - TenantIDs: []uuid.UUID{a.TenantID}, - SiteIDs: []uuid.UUID{a.SiteID}, - }, - cdbp.PageInput{}, - nil, - ) - - if serr != nil { - logger.Error().Err(serr).Msg("error getting count of Tenant/Site associations") - return cutil.NewAPIError(http.StatusInternalServerError, "Error deleting Allocation, DB error retrieving Tenant/Site associations", nil) - } - if tscount > 0 { - derr := tsDAO.Delete(ctx, tx, tss[0].ID) - if derr != nil { - logger.Error().Err(derr).Msg("error deleting Tenant/Site association") - return cutil.NewAPIError(http.StatusInternalServerError, "Error deleting Allocation, DB error deleting Tenant/Site association", nil) + // Privileged tenants (and service accounts, which are marked + // the same way) keep their TenantSite association even after + // their last allocation on a site is removed, since they + // don't rely on allocations to access the site in the first + // place. a.Tenant is already loaded via includeRelations on + // the earlier GetByID for this allocation. + isPrivileged := a.Tenant != nil && a.Tenant.Config != nil && a.Tenant.Config.TargetedInstanceCreation + + if !isPrivileged { + tsDAO := cdbm.NewTenantSiteDAO(dah.dbSession) + tss, tscount, serr := tsDAO.GetAll( + ctx, + tx, + cdbm.TenantSiteFilterInput{ + TenantIDs: []uuid.UUID{a.TenantID}, + SiteIDs: []uuid.UUID{a.SiteID}, + }, + cdbp.PageInput{}, + nil, + ) + if serr != nil { + logger.Error().Err(serr).Msg("error getting count of Tenant/Site associations") + return cutil.NewAPIError(http.StatusInternalServerError, "Error deleting Allocation, DB error retrieving Tenant/Site associations", nil) + } + if tscount > 0 { + derr := tsDAO.Delete(ctx, tx, tss[0].ID) + if derr != nil { + logger.Error().Err(derr).Msg("error deleting Tenant/Site association") + return cutil.NewAPIError(http.StatusInternalServerError, "Error deleting Allocation, DB error deleting Tenant/Site association", nil) + } + } } } - } return nil }) if err != nil { diff --git a/rest-api/api/pkg/api/handler/allocation_test.go b/rest-api/api/pkg/api/handler/allocation_test.go index b815cf1586..65c3925050 100644 --- a/rest-api/api/pkg/api/handler/allocation_test.go +++ b/rest-api/api/pkg/api/handler/allocation_test.go @@ -2796,3 +2796,102 @@ func TestInstanceTypeAllocationForMultipleTenants(t *testing.T) { }) } } + +func TestAllocationHandler_Delete_PrivilegedTenantKeepsTenantSite(t *testing.T) { + ctx := context.Background() + dbSession := common.TestInitDB(t) + defer dbSession.Close() + + common.TestSetupSchema(t, dbSession) + + ipOrg := "test-privileged-ip-org" + tnOrg := "test-privileged-tn-org" + + ipu := common.TestBuildUser(t, dbSession, uuid.NewString(), ipOrg, []string{authz.ProviderAdminRole}) + tnu := common.TestBuildUser(t, dbSession, uuid.NewString(), tnOrg, []string{authz.TenantAdminRole}) + + ip := common.TestBuildInfrastructureProvider(t, dbSession, "test-privileged-ip", ipOrg, ipu) + site := common.TestBuildSite(t, dbSession, ip, "test-privileged-site", ipu) + + tenant := common.TestBuildTenant(t, dbSession, "test-privileged-tenant", tnOrg, tnu) + + // Mark this tenant as privileged. + tnDAO := cdbm.NewTenantDAO(dbSession) + updatedTenant, err := tnDAO.Update(ctx, nil, cdbm.TenantUpdateInput{ + TenantID: tenant.ID, + Config: &cdbm.TenantConfig{TargetedInstanceCreation: true}, + }) + require.NoError(t, err) + require.True(t, updatedTenant.Config.TargetedInstanceCreation) + + it := common.TestBuildInstanceType(t, dbSession, "test-privileged-it", cutil.GetPtr(uuid.New()), site, map[string]string{ + "name": "test-privileged-instance-type", + }, ipu) + + for i := 1; i <= 1; i++ { + mc := testInstanceBuildMachine(t, dbSession, ip.ID, site.ID, cutil.GetPtr(false), nil) + require.NotNil(t, mc) + mit := testInstanceBuildMachineInstanceType(t, dbSession, mc, it) + require.NotNil(t, mit) + } + + ac := model.APIAllocationConstraintCreateRequest{ + ResourceType: cdbm.AllocationResourceTypeInstanceType, + ResourceTypeID: it.ID.String(), + ConstraintType: cdbm.AllocationConstraintTypeReserved, + ConstraintValue: 1, + } + body, err := json.Marshal(model.APIAllocationCreateRequest{ + Name: "privileged-tenant-allocation", + Description: cutil.GetPtr(""), + TenantID: tenant.ID.String(), + SiteID: site.ID.String(), + AllocationConstraints: []model.APIAllocationConstraintCreateRequest{ac}, + }) + require.NoError(t, err) + + ipamStorage := ipam.NewIpamStorage(dbSession.DB, nil) + allocation := testCreateAllocation(t, dbSession, ipamStorage, ipu, ipOrg, string(body)) + require.NotNil(t, allocation) + + tsDAO := cdbm.NewTenantSiteDAO(dbSession) + + // Sanity check: creating the allocation should have created a TenantSite. + _, tscount, err := tsDAO.GetAll(ctx, nil, cdbm.TenantSiteFilterInput{ + TenantIDs: []uuid.UUID{tenant.ID}, + SiteIDs: []uuid.UUID{site.ID}, + }, cdbp.PageInput{}, nil) + require.NoError(t, err) + require.Equal(t, 1, tscount) + + // Delete the allocation -- it's the tenant's only/last one on this site. + e := echo.New() + req := httptest.NewRequest(http.MethodPost, "/", nil) + req.Header.Set(echo.HeaderContentType, echo.MIMEApplicationJSON) + rec := httptest.NewRecorder() + + ec := e.NewContext(req, rec) + ec.SetParamNames("orgName", "id") + ec.SetParamValues(ipOrg, allocation.ID) + ec.Set("user", ipu) + ec.SetRequest(ec.Request().WithContext(ctx)) + + cfg := common.GetTestConfig() + dah := DeleteAllocationHandler{ + dbSession: dbSession, + tc: &tmocks.Client{}, + cfg: cfg, + } + err = dah.Handle(ec) + require.NoError(t, err) + require.Equal(t, http.StatusAccepted, rec.Code) + + // Even though this was the tenant's last allocation on the site, the + // TenantSite association should persist because the tenant is privileged. + _, tscount, err = tsDAO.GetAll(ctx, nil, cdbm.TenantSiteFilterInput{ + TenantIDs: []uuid.UUID{tenant.ID}, + SiteIDs: []uuid.UUID{site.ID}, + }, cdbp.PageInput{}, nil) + require.NoError(t, err) + require.Equal(t, 1, tscount, "privileged tenant should keep its TenantSite association after its last allocation is deleted") +} \ No newline at end of file diff --git a/rest-api/api/pkg/api/handler/serviceaccount.go b/rest-api/api/pkg/api/handler/serviceaccount.go index 7004790958..9702756c53 100644 --- a/rest-api/api/pkg/api/handler/serviceaccount.go +++ b/rest-api/api/pkg/api/handler/serviceaccount.go @@ -177,6 +177,18 @@ func (gcsah GetCurrentServiceAccountHandler) Handle(c echo.Context) error { } } + // Ensure this service account's tenant has TenantSite records for every + // site under this org's infrastructure provider, since service accounts + // should be able to operate on any site without needing an allocation. + serr = cdb.WithTx(ctx, gcsah.dbSession, func(tx *cdb.Tx) error { + return EnsureTenantSitesForInfrastructureProvider(ctx, tx, gcsah.dbSession, tn, ip.ID, dbUser.ID) + }) + if serr != nil { + return common.HandleTxError(c, logger, serr, "Failed to ensure TenantSite records for service account") + } + + // Create response + // Create response apiServiceAccount := model.NewAPIServiceAccount(serviceAccountEnabled, ip, tn) diff --git a/rest-api/api/pkg/api/handler/tenant.go b/rest-api/api/pkg/api/handler/tenant.go index 4b5ab28bf8..ac2c45279b 100644 --- a/rest-api/api/pkg/api/handler/tenant.go +++ b/rest-api/api/pkg/api/handler/tenant.go @@ -206,6 +206,28 @@ func (gcth GetCurrentTenantHandler) Handle(c echo.Context) error { return common.HandleTxError(c, logger, err, "Failed to retrieve current Tenant, DB transaction error") } + // If this tenant is privileged (targeted instance creation enabled), + // ensure it has TenantSite records for every site under its org's + // infrastructure provider, since privileged tenants should be able to + // operate on any site without needing an allocation first. + if tn.Config != nil && tn.Config.TargetedInstanceCreation { + ipDAO := cdbm.NewInfrastructureProviderDAO(gcth.dbSession) + ips, ipErr := ipDAO.GetAllByOrg(ctx, nil, org, nil) + if ipErr != nil { + logger.Error().Err(ipErr).Msg("error retrieving Infrastructure Provider for this org") + return cutil.NewAPIErrorResponse(c, http.StatusInternalServerError, "Failed to retrieve Infrastructure Provider for org, DB error", nil) + } + if len(ips) > 0 { + ip := ips[0] + err = cdb.WithTx(ctx, gcth.dbSession, func(tx *cdb.Tx) error { + return EnsureTenantSitesForInfrastructureProvider(ctx, tx, gcth.dbSession, tn, ip.ID, dbUser.ID) + }) + if err != nil { + return common.HandleTxError(c, logger, err, "Failed to ensure TenantSite records for privileged tenant") + } + } + } + // Create response apiInstance := model.NewAPITenant(tn) diff --git a/rest-api/api/pkg/api/handler/tenantsite_helpers.go b/rest-api/api/pkg/api/handler/tenantsite_helpers.go new file mode 100644 index 0000000000..00b0f178c2 --- /dev/null +++ b/rest-api/api/pkg/api/handler/tenantsite_helpers.go @@ -0,0 +1,90 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +package handler + +import ( + "context" + "fmt" + + cdb "github.com/NVIDIA/infra-controller/rest-api/db/pkg/db" + cdbm "github.com/NVIDIA/infra-controller/rest-api/db/pkg/db/model" + cdbp "github.com/NVIDIA/infra-controller/rest-api/db/pkg/db/paginator" + cutil "github.com/NVIDIA/infra-controller/rest-api/common/pkg/util" + "github.com/google/uuid" +) + +// EnsureTenantSitesForInfrastructureProvider auto-creates TenantSite records +// linking the given tenant to every registered Site under the given +// InfrastructureProvider. This is used for tenants that should be able to +// operate on any site belonging to their infrastructure provider without +// first needing a compute or network allocation there -- privileged tenants +// (TargetedInstanceCreation enabled) and service accounts. +func EnsureTenantSitesForInfrastructureProvider( + ctx context.Context, + tx *cdb.Tx, + dbSession *cdb.Session, + tenant *cdbm.Tenant, + infrastructureProviderID uuid.UUID, + createdBy uuid.UUID, +) error { + siteDAO := cdbm.NewSiteDAO(dbSession) + sites, _, err := siteDAO.GetAll( + ctx, + tx, + cdbm.SiteFilterInput{InfrastructureProviderIDs: []uuid.UUID{infrastructureProviderID}}, + cdbp.PageInput{Limit: cutil.GetPtr(cdbp.TotalLimit)}, + nil, + ) + if err != nil { + return fmt.Errorf("error retrieving sites for infrastructure provider: %w", err) + } + if len(sites) == 0 { + return nil + } + + siteIDs := make([]uuid.UUID, len(sites)) + for i, site := range sites { + siteIDs[i] = site.ID + } + + tsDAO := cdbm.NewTenantSiteDAO(dbSession) + existing, _, err := tsDAO.GetAll( + ctx, + tx, + cdbm.TenantSiteFilterInput{ + TenantIDs: []uuid.UUID{tenant.ID}, + SiteIDs: siteIDs, + }, + cdbp.PageInput{Limit: cutil.GetPtr(cdbp.TotalLimit)}, + nil, + ) + if err != nil { + return fmt.Errorf("error retrieving existing TenantSite entries: %w", err) + } + + existingSiteIDs := make(map[uuid.UUID]bool, len(existing)) + for _, ts := range existing { + existingSiteIDs[ts.SiteID] = true + } + + for _, site := range sites { + if existingSiteIDs[site.ID] { + continue + } + _, err := tsDAO.Create( + ctx, + tx, + cdbm.TenantSiteCreateInput{ + TenantID: tenant.ID, + TenantOrg: tenant.Org, + SiteID: site.ID, + CreatedBy: createdBy, + }, + ) + if err != nil { + return fmt.Errorf("error creating TenantSite entry for site %s: %w", site.ID, err) + } + } + return nil +} \ No newline at end of file diff --git a/rest-api/api/pkg/api/handler/tenantsite_helpers_test.go b/rest-api/api/pkg/api/handler/tenantsite_helpers_test.go new file mode 100644 index 0000000000..86064a439b --- /dev/null +++ b/rest-api/api/pkg/api/handler/tenantsite_helpers_test.go @@ -0,0 +1,68 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +package handler + +import ( + "context" + "testing" + + "github.com/NVIDIA/infra-controller/rest-api/api/pkg/api/handler/util/common" + authz "github.com/NVIDIA/infra-controller/rest-api/auth/pkg/authorization" + cdb "github.com/NVIDIA/infra-controller/rest-api/db/pkg/db" + cdbm "github.com/NVIDIA/infra-controller/rest-api/db/pkg/db/model" + cdbp "github.com/NVIDIA/infra-controller/rest-api/db/pkg/db/paginator" + "github.com/google/uuid" + "github.com/stretchr/testify/require" +) + +func TestEnsureTenantSitesForInfrastructureProvider(t *testing.T) { + ctx := context.Background() + + dbSession := common.TestInitDB(t) + defer dbSession.Close() + + common.TestSetupSchema(t, dbSession) + + org := "test-org-ensure-sites" + user := common.TestBuildUser(t, dbSession, uuid.NewString(), org, []string{authz.ProviderAdminRole, authz.TenantAdminRole}) + ip := common.TestBuildInfrastructureProvider(t, dbSession, "test-provider-ensure-sites", org, user) + tn := common.TestBuildTenant(t, dbSession, "test-tenant-ensure-sites", org, user) + + site1 := common.TestBuildSite(t, dbSession, ip, "site-1", user) + site2 := common.TestBuildSite(t, dbSession, ip, "site-2", user) + + tsDAO := cdbm.NewTenantSiteDAO(dbSession) + + // Sanity check: no TenantSite records exist yet. + _, count, err := tsDAO.GetAll(ctx, nil, cdbm.TenantSiteFilterInput{TenantIDs: []uuid.UUID{tn.ID}}, cdbp.PageInput{}, nil) + require.NoError(t, err) + require.Equal(t, 0, count) + + // First call should create one TenantSite record per site. + err = cdb.WithTx(ctx, dbSession, func(tx *cdb.Tx) error { + return EnsureTenantSitesForInfrastructureProvider(ctx, tx, dbSession, tn, ip.ID, user.ID) + }) + require.NoError(t, err) + + sites, count, err := tsDAO.GetAll(ctx, nil, cdbm.TenantSiteFilterInput{TenantIDs: []uuid.UUID{tn.ID}}, cdbp.PageInput{}, nil) + require.NoError(t, err) + require.Equal(t, 2, count) + + siteIDs := map[uuid.UUID]bool{} + for _, s := range sites { + siteIDs[s.SiteID] = true + } + require.True(t, siteIDs[site1.ID]) + require.True(t, siteIDs[site2.ID]) + + // Calling it again should be idempotent -- no duplicate TenantSite rows. + err = cdb.WithTx(ctx, dbSession, func(tx *cdb.Tx) error { + return EnsureTenantSitesForInfrastructureProvider(ctx, tx, dbSession, tn, ip.ID, user.ID) + }) + require.NoError(t, err) + + _, count, err = tsDAO.GetAll(ctx, nil, cdbm.TenantSiteFilterInput{TenantIDs: []uuid.UUID{tn.ID}}, cdbp.PageInput{}, nil) + require.NoError(t, err) + require.Equal(t, 2, count, "second call should not create duplicate TenantSite records") +} \ No newline at end of file