Improve custom policy usage reconciliation

platform-api/internal/repository/custom_policy.go

@@ -20,6 +20,7 @@ package repository
 import (
 	"database/sql"
 	"errors"
+	"strings"
 	"time"
 
 	"github.com/wso2/api-platform/platform-api/internal/apperror"
@@ -32,6 +33,11 @@ type CustomPolicyRepo struct {
 	db *database.DB
 }
 
+// Each custom-policy usage row uses two bind parameters. Keeping batches at
+// 499 rows stays below SQLite's conservative 999-parameter limit and therefore
+// also below the limits of the other supported databases.
+const maxCustomPolicyUsageRowsPerInsert = 499
+
 // NewCustomPolicyRepo creates a new CustomPolicyRepo
 func NewCustomPolicyRepo(db *database.DB) CustomPolicyRepository {
 	return &CustomPolicyRepo{db: db}
@@ -239,13 +245,30 @@ func replaceCustomPolicyUsagesTx(tx *sql.Tx, db *database.DB, artifactUUID strin
 	if _, err := tx.Exec(db.Rebind(`DELETE FROM gateway_custom_policy_usages WHERE artifact_uuid = ?`), artifactUUID); err != nil {
 		return err
 	}
+
 	seen := make(map[string]struct{}, len(policyUUIDs))
+	uniquePolicyUUIDs := make([]string, 0, len(policyUUIDs))
 	for _, policyUUID := range policyUUIDs {
 		if _, exists := seen[policyUUID]; exists {
 			continue
 		}
 		seen[policyUUID] = struct{}{}
-		if _, err := tx.Exec(db.Rebind(`INSERT INTO gateway_custom_policy_usages (policy_uuid, artifact_uuid) VALUES (?, ?)`), policyUUID, artifactUUID); err != nil {
+		uniquePolicyUUIDs = append(uniquePolicyUUIDs, policyUUID)
+	}
+
+	for start := 0; start < len(uniquePolicyUUIDs); start += maxCustomPolicyUsageRowsPerInsert {
+		end := min(start+maxCustomPolicyUsageRowsPerInsert, len(uniquePolicyUUIDs))
+		batch := uniquePolicyUUIDs[start:end]
+		valuePlaceholders := make([]string, len(batch))
+		args := make([]any, 0, len(batch)*2)
+		for i, policyUUID := range batch {
+			valuePlaceholders[i] = "(?, ?)"
+			args = append(args, policyUUID, artifactUUID)
+		}
+
+		query := `INSERT INTO gateway_custom_policy_usages (policy_uuid, artifact_uuid) VALUES ` +
+			strings.Join(valuePlaceholders, ", ")
+		if _, err := tx.Exec(db.Rebind(query), args...); err != nil {
 			return err
 		}
 	}

platform-api/internal/repository/llm.go

@@ -880,16 +880,16 @@ func NewLLMProviderRepo(db *database.DB) LLMProviderRepository {
 }
 
 func (r *LLMProviderRepo) Create(p *model.LLMProvider) error {
-	return r.create(p, nil, false)
+	return r.create(p, nil)
 }
 
 // CreateWithCustomPolicyUsages creates an LLM provider and its custom-policy
 // deletion guards in one transaction.
 func (r *LLMProviderRepo) CreateWithCustomPolicyUsages(p *model.LLMProvider, policyUUIDs []string) error {
-	return r.create(p, policyUUIDs, true)
+	return r.create(p, policyUUIDs)
 }
 
-func (r *LLMProviderRepo) create(p *model.LLMProvider, policyUUIDs []string, reconcilePolicyUsages bool) error {
+func (r *LLMProviderRepo) create(p *model.LLMProvider, policyUUIDs []string) error {
 	uuidStr, err := utils.GenerateUUID()
 	if err != nil {
 		return fmt.Errorf("failed to generate LLM provider ID: %w", err)
@@ -957,10 +957,8 @@ func (r *LLMProviderRepo) create(p *model.LLMProvider, policyUUIDs []string, rec
 	if err := insertArtifactGatewayAssociations(tx, r.db, p.UUID, p.OrganizationUUID, p.CreatedBy, p.AssociatedGateways, now); err != nil {
 		return err
 	}
-	if reconcilePolicyUsages {
-		if err := replaceCustomPolicyUsagesTx(tx, r.db, p.UUID, policyUUIDs); err != nil {
-			return fmt.Errorf("failed to persist custom policy usages: %w", err)
-		}
+	if err := replaceCustomPolicyUsagesTx(tx, r.db, p.UUID, policyUUIDs); err != nil {
+		return fmt.Errorf("failed to persist custom policy usages: %w", err)
 	}
 
 	if err := tx.Commit(); err != nil {
@@ -1084,16 +1082,16 @@ func (r *LLMProviderRepo) Count(orgUUID string) (int, error) {
 }
 
 func (r *LLMProviderRepo) Update(p *model.LLMProvider) error {
-	return r.update(p, nil, false)
+	return r.update(p, nil)
 }
 
 // UpdateWithCustomPolicyUsages updates an LLM provider and replaces its
 // custom-policy deletion guards in one transaction.
 func (r *LLMProviderRepo) UpdateWithCustomPolicyUsages(p *model.LLMProvider, policyUUIDs []string) error {
-	return r.update(p, policyUUIDs, true)
+	return r.update(p, policyUUIDs)
 }
 
-func (r *LLMProviderRepo) update(p *model.LLMProvider, policyUUIDs []string, reconcilePolicyUsages bool) error {
+func (r *LLMProviderRepo) update(p *model.LLMProvider, policyUUIDs []string) error {
 	now := time.Now().UTC()
 	p.UpdatedAt = now
 
@@ -1166,10 +1164,8 @@ func (r *LLMProviderRepo) update(p *model.LLMProvider, policyUUIDs []string, rec
 			return err
 		}
 	}
-	if reconcilePolicyUsages {
-		if err := replaceCustomPolicyUsagesTx(tx, r.db, providerUUID, policyUUIDs); err != nil {
-			return fmt.Errorf("failed to persist custom policy usages: %w", err)
-		}
+	if err := replaceCustomPolicyUsagesTx(tx, r.db, providerUUID, policyUUIDs); err != nil {
+		return fmt.Errorf("failed to persist custom policy usages: %w", err)
 	}
 
 	if err := tx.Commit(); err != nil {

platform-api/internal/repository/llm_test.go

@@ -27,7 +27,7 @@ import (
 	_ "github.com/mattn/go-sqlite3"
 )
 
-func TestLLMProviderRepoUpdateWithCustomPolicyUsagesRollsBackOnInsertFailure(t *testing.T) {
+func TestLLMProviderRepoCustomPolicyUsageReconciliation(t *testing.T) {
 	db, cleanup := setupTestDB(t)
 	t.Cleanup(cleanup)
 
@@ -60,6 +60,15 @@ func TestLLMProviderRepoUpdateWithCustomPolicyUsagesRollsBackOnInsertFailure(t *
 	if err := customPolicyRepo.InsertCustomPolicy(policy); err != nil {
 		t.Fatalf("create custom policy: %v", err)
 	}
+	secondPolicy := &model.CustomPolicy{
+		UUID:             "policy-second",
+		OrganizationUUID: orgUUID,
+		Name:             "second-custom-policy",
+		Version:          "v1.0.0",
+	}
+	if err := customPolicyRepo.InsertCustomPolicy(secondPolicy); err != nil {
+		t.Fatalf("create second custom policy: %v", err)
+	}
 
 	providerRepo := NewLLMProviderRepo(db)
 	provider := &model.LLMProvider{
@@ -69,7 +78,7 @@ func TestLLMProviderRepoUpdateWithCustomPolicyUsagesRollsBackOnInsertFailure(t *
 		Version:          "v1.0",
 		TemplateUUID:     template.UUID,
 	}
-	if err := providerRepo.CreateWithCustomPolicyUsages(provider, []string{policy.UUID}); err != nil {
+	if err := providerRepo.CreateWithCustomPolicyUsages(provider, []string{policy.UUID, policy.UUID, secondPolicy.UUID}); err != nil {
 		t.Fatalf("create provider: %v", err)
 	}
 
@@ -90,8 +99,31 @@ func TestLLMProviderRepoUpdateWithCustomPolicyUsagesRollsBackOnInsertFailure(t *
 	if err != nil {
 		t.Fatalf("get usages after failed update: %v", err)
 	}
-	if len(usages) != 1 || usages[0] != policy.UUID {
-		t.Fatalf("policy usages = %v, want [%s]", usages, policy.UUID)
+	usageSet := make(map[string]struct{}, len(usages))
+	for _, policyUUID := range usages {
+		usageSet[policyUUID] = struct{}{}
+	}
+	if len(usages) != 2 {
+		t.Fatalf("policy usages = %v, want [%s %s]", usages, policy.UUID, secondPolicy.UUID)
+	}
+	if _, exists := usageSet[policy.UUID]; !exists {
+		t.Fatalf("policy usages = %v, missing %s", usages, policy.UUID)
+	}
+	if _, exists := usageSet[secondPolicy.UUID]; !exists {
+		t.Fatalf("policy usages = %v, missing %s", usages, secondPolicy.UUID)
+	}
+
+	// The plain update path must also reconcile usages. With no policy UUIDs,
+	// it removes any usages previously stored for the provider.
+	if err := providerRepo.Update(provider); err != nil {
+		t.Fatalf("plain provider update: %v", err)
+	}
+	usages, err = customPolicyRepo.GetCustomPolicyUsagesByAPIUUID(provider.UUID)
+	if err != nil {
+		t.Fatalf("get usages after plain update: %v", err)
+	}
+	if len(usages) != 0 {
+		t.Fatalf("policy usages after plain update = %v, want none", usages)
 	}
 }