fix

Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>

backend/ee/onyx/server/tenants/billing_api.py

@@ -10,10 +10,12 @@ from ee.onyx.server.tenants.billing import fetch_billing_information
 from ee.onyx.server.tenants.billing import fetch_stripe_checkout_session
 from ee.onyx.server.tenants.billing import fetch_tenant_stripe_information
 from ee.onyx.server.tenants.models import BillingInformation
+from ee.onyx.server.tenants.models import ProductGatingFullSyncRequest
 from ee.onyx.server.tenants.models import ProductGatingRequest
 from ee.onyx.server.tenants.models import ProductGatingResponse
 from ee.onyx.server.tenants.models import SubscriptionSessionResponse
 from ee.onyx.server.tenants.models import SubscriptionStatusResponse
+from ee.onyx.server.tenants.product_gating import overwrite_full_gated_set
 from ee.onyx.server.tenants.product_gating import store_product_gating
 from onyx.auth.users import User
 from onyx.configs.app_configs import WEB_DOMAIN
@@ -47,6 +49,26 @@ def gate_product(
         return ProductGatingResponse(updated=False, error=str(e))
 
 
+@router.post("/product-gating/full-sync")
+def gate_product_full_sync(
+    product_gating_request: ProductGatingFullSyncRequest,
+    _: None = Depends(control_plane_dep),
+) -> ProductGatingResponse:
+    """
+    Bulk operation to overwrite the entire gated tenant set.
+    This replaces all currently gated tenants with the provided list.
+    Gated tenants are not available to access the product and will be
+    directed to the billing page when their subscription has ended.
+    """
+    try:
+        overwrite_full_gated_set(product_gating_request.gated_tenant_ids)
+        return ProductGatingResponse(updated=True, error=None)
+
+    except Exception as e:
+        logger.exception("Failed to gate products during full sync")
+        return ProductGatingResponse(updated=False, error=str(e))
+
+
 @router.get("/billing-information")
 async def billing_information(
     _: User = Depends(current_admin_user),

backend/ee/onyx/server/tenants/models.py

@@ -19,6 +19,10 @@ class ProductGatingRequest(BaseModel):
     application_status: ApplicationStatus
 
 
+class ProductGatingFullSyncRequest(BaseModel):
+    gated_tenant_ids: list[str]
+
+
 class SubscriptionStatusResponse(BaseModel):
     subscribed: bool
 

backend/ee/onyx/server/tenants/product_gating.py

@@ -16,10 +16,6 @@ logger = setup_logger()
 def update_tenant_gating(tenant_id: str, status: ApplicationStatus) -> None:
     redis_client = get_redis_client(tenant_id=ONYX_CLOUD_TENANT_ID)
 
-    # Store the full status
-    status_key = f"tenant:{tenant_id}:status"
-    redis_client.set(status_key, status.value)
-
     # Maintain the GATED_ACCESS set
     if status == ApplicationStatus.GATED_ACCESS:
         redis_client.sadd(GATED_TENANTS_KEY, tenant_id)
@@ -46,6 +42,22 @@ def store_product_gating(tenant_id: str, application_status: ApplicationStatus)
         raise
 
 
+def overwrite_full_gated_set(tenant_ids: list[str]) -> None:
+    redis_client = get_redis_client(tenant_id=ONYX_CLOUD_TENANT_ID)
+
+    pipeline = redis_client.pipeline()
+
+    # Clear the existing set
+    pipeline.delete(GATED_TENANTS_KEY)
+
+    # Add all tenant IDs to the set and set their status
+    for tenant_id in tenant_ids:
+        pipeline.sadd(GATED_TENANTS_KEY, tenant_id)
+
+    # Execute all commands at once
+    pipeline.execute()
+
+
 def get_gated_tenants() -> set[str]:
     redis_client = get_redis_replica_client(tenant_id=ONYX_CLOUD_TENANT_ID)
     gated_tenants_bytes = cast(set[bytes], redis_client.smembers(GATED_TENANTS_KEY))