feat: auto-generate TenantSite entries for privileged tenants and service accounts

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 {

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")
+}

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)
 

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)
 

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
+}

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")
+}