Proper user deletion / organization leaving (#3460)

* Proper user deletion / organization leaving

* minor nit

* update

* udpate provisioning

* minor cleanup

* typing

* post rebase

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

@@ -3,6 +3,7 @@ from fastapi import APIRouter
 from fastapi import Depends
 from fastapi import HTTPException
 from fastapi import Response
+from sqlalchemy.orm import Session
 
 from ee.onyx.auth.users import current_cloud_superuser
 from ee.onyx.configs.app_configs import STRIPE_SECRET_KEY
@@ -12,15 +13,23 @@ 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 ImpersonateRequest
 from ee.onyx.server.tenants.models import ProductGatingRequest
+from ee.onyx.server.tenants.provisioning import delete_user_from_control_plane
 from ee.onyx.server.tenants.user_mapping import get_tenant_id_for_email
+from ee.onyx.server.tenants.user_mapping import remove_all_users_from_tenant
+from ee.onyx.server.tenants.user_mapping import remove_users_from_tenant
 from onyx.auth.users import auth_backend
 from onyx.auth.users import current_admin_user
 from onyx.auth.users import get_jwt_strategy
 from onyx.auth.users import User
 from onyx.configs.app_configs import WEB_DOMAIN
+from onyx.db.auth import get_user_count
+from onyx.db.engine import get_current_tenant_id
+from onyx.db.engine import get_session
 from onyx.db.engine import get_session_with_tenant
 from onyx.db.notification import create_notification
+from onyx.db.users import delete_user_from_db
 from onyx.db.users import get_user_by_email
+from onyx.server.manage.models import UserByEmail
 from onyx.server.settings.store import load_settings
 from onyx.server.settings.store import store_settings
 from onyx.utils.logger import setup_logger
@@ -114,3 +123,48 @@ async def impersonate_user(
         samesite="lax",
     )
     return response
+
+
+@router.post("/leave-organization")
+async def leave_organization(
+    user_email: UserByEmail,
+    current_user: User | None = Depends(current_admin_user),
+    db_session: Session = Depends(get_session),
+    tenant_id: str = Depends(get_current_tenant_id),
+) -> None:
+    if current_user is None or current_user.email != user_email.user_email:
+        raise HTTPException(
+            status_code=403, detail="You can only leave the organization as yourself"
+        )
+
+    user_to_delete = get_user_by_email(user_email.user_email, db_session)
+    if user_to_delete is None:
+        raise HTTPException(status_code=404, detail="User not found")
+
+    num_admin_users = await get_user_count(only_admin_users=True)
+
+    should_delete_tenant = num_admin_users == 1
+
+    if should_delete_tenant:
+        logger.info(
+            "Last admin user is leaving the organization. Deleting tenant from control plane."
+        )
+        try:
+            await delete_user_from_control_plane(tenant_id, user_to_delete.email)
+            logger.debug("User deleted from control plane")
+        except Exception as e:
+            logger.exception(
+                f"Failed to delete user from control plane for tenant {tenant_id}: {e}"
+            )
+            raise HTTPException(
+                status_code=500,
+                detail=f"Failed to remove user from control plane: {str(e)}",
+            )
+
+    db_session.expunge(user_to_delete)
+    delete_user_from_db(user_to_delete, db_session)
+
+    if should_delete_tenant:
+        remove_all_users_from_tenant(tenant_id)
+    else:
+        remove_users_from_tenant([user_to_delete.email], tenant_id)

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

@@ -39,3 +39,8 @@ class TenantCreationPayload(BaseModel):
     tenant_id: str
     email: str
     referral_source: str | None = None
+
+
+class TenantDeletionPayload(BaseModel):
+    tenant_id: str
+    email: str

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

@@ -15,6 +15,7 @@ from ee.onyx.configs.app_configs import HUBSPOT_TRACKING_URL
 from ee.onyx.configs.app_configs import OPENAI_DEFAULT_API_KEY
 from ee.onyx.server.tenants.access import generate_data_plane_token
 from ee.onyx.server.tenants.models import TenantCreationPayload
+from ee.onyx.server.tenants.models import TenantDeletionPayload
 from ee.onyx.server.tenants.schema_management import create_schema_if_not_exists
 from ee.onyx.server.tenants.schema_management import drop_schema
 from ee.onyx.server.tenants.schema_management import run_alembic_migrations
@@ -185,6 +186,7 @@ async def rollback_tenant_provisioning(tenant_id: str) -> None:
     try:
         # Drop the tenant's schema to rollback provisioning
         drop_schema(tenant_id)
+
         # Remove tenant mapping
         with Session(get_sqlalchemy_engine()) as db_session:
             db_session.query(UserTenantMapping).filter(
@@ -320,3 +322,26 @@ async def submit_to_hubspot(
 
     if response.status_code != 200:
         logger.error(f"Failed to submit to HubSpot: {response.text}")
+
+
+async def delete_user_from_control_plane(tenant_id: str, email: str) -> None:
+    token = generate_data_plane_token()
+    headers = {
+        "Authorization": f"Bearer {token}",
+        "Content-Type": "application/json",
+    }
+    payload = TenantDeletionPayload(tenant_id=tenant_id, email=email)
+
+    async with aiohttp.ClientSession() as session:
+        async with session.delete(
+            f"{CONTROL_PLANE_API_BASE_URL}/tenants/delete",
+            headers=headers,
+            json=payload.model_dump(),
+        ) as response:
+            print(response)
+            if response.status != 200:
+                error_text = await response.text()
+                logger.error(f"Control plane tenant creation failed: {error_text}")
+                raise Exception(
+                    f"Failed to delete tenant on control plane: {error_text}"
+                )

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

@@ -68,3 +68,11 @@ def remove_users_from_tenant(emails: list[str], tenant_id: str) -> None:
                 f"Failed to remove users from tenant {tenant_id}: {str(e)}"
             )
             db_session.rollback()
+
+
+def remove_all_users_from_tenant(tenant_id: str) -> None:
+    with get_session_with_tenant(POSTGRES_DEFAULT_SCHEMA) as db_session:
+        db_session.query(UserTenantMapping).filter(
+            UserTenantMapping.tenant_id == tenant_id
+        ).delete()
+        db_session.commit()

backend/onyx/db/auth.py

@@ -54,9 +54,11 @@ def get_total_users_count(db_session: Session) -> int:
     return user_count + invited_users
 
 
-async def get_user_count() -> int:
+async def get_user_count(only_admin_users: bool = False) -> int:
     async with get_async_session_with_tenant() as session:
         stmt = select(func.count(User.id))
+        if only_admin_users:
+            stmt = stmt.where(User.role == UserRole.ADMIN)
         result = await session.execute(stmt)
         user_count = result.scalar()
         if user_count is None:

backend/onyx/db/users.py

@@ -7,8 +7,15 @@ from sqlalchemy import func
 from sqlalchemy import select
 from sqlalchemy.orm import Session
 
+from onyx.auth.invited_users import get_invited_users
+from onyx.auth.invited_users import write_invited_users
 from onyx.auth.schemas import UserRole
+from onyx.db.models import DocumentSet__User
+from onyx.db.models import Persona__User
+from onyx.db.models import SamlAccount
 from onyx.db.models import User
+from onyx.db.models import User__UserGroup
+from onyx.utils.variable_functionality import fetch_ee_implementation_or_noop
 
 
 def validate_user_role_update(requested_role: UserRole, current_role: UserRole) -> None:
@@ -185,3 +192,43 @@ def batch_add_ext_perm_user_if_not_exists(
     db_session.commit()
 
     return found_users + new_users
+
+
+def delete_user_from_db(
+    user_to_delete: User,
+    db_session: Session,
+) -> None:
+    for oauth_account in user_to_delete.oauth_accounts:
+        db_session.delete(oauth_account)
+
+    fetch_ee_implementation_or_noop(
+        "onyx.db.external_perm",
+        "delete_user__ext_group_for_user__no_commit",
+    )(
+        db_session=db_session,
+        user_id=user_to_delete.id,
+    )
+    db_session.query(SamlAccount).filter(
+        SamlAccount.user_id == user_to_delete.id
+    ).delete()
+    db_session.query(DocumentSet__User).filter(
+        DocumentSet__User.user_id == user_to_delete.id
+    ).delete()
+    db_session.query(Persona__User).filter(
+        Persona__User.user_id == user_to_delete.id
+    ).delete()
+    db_session.query(User__UserGroup).filter(
+        User__UserGroup.user_id == user_to_delete.id
+    ).delete()
+    db_session.delete(user_to_delete)
+    db_session.commit()
+
+    # NOTE: edge case may exist with race conditions
+    # with this `invited user` scheme generally.
+    user_emails = get_invited_users()
+    remaining_users = [
+        remaining_user_email
+        for remaining_user_email in user_emails
+        if remaining_user_email != user_to_delete.email
+    ]
+    write_invited_users(remaining_users)

backend/onyx/server/manage/users.py

@@ -42,11 +42,8 @@ from onyx.db.auth import get_total_users_count
 from onyx.db.engine import CURRENT_TENANT_ID_CONTEXTVAR
 from onyx.db.engine import get_session
 from onyx.db.models import AccessToken
-from onyx.db.models import DocumentSet__User
-from onyx.db.models import Persona__User
-from onyx.db.models import SamlAccount
 from onyx.db.models import User
-from onyx.db.models import User__UserGroup
+from onyx.db.users import delete_user_from_db
 from onyx.db.users import get_user_by_email
 from onyx.db.users import list_users
 from onyx.db.users import validate_user_role_update
@@ -370,45 +367,10 @@ async def delete_user(
     db_session.expunge(user_to_delete)
 
     try:
-        for oauth_account in user_to_delete.oauth_accounts:
-            db_session.delete(oauth_account)
-
-        fetch_ee_implementation_or_noop(
-            "onyx.db.external_perm",
-            "delete_user__ext_group_for_user__no_commit",
-        )(
-            db_session=db_session,
-            user_id=user_to_delete.id,
-        )
-        db_session.query(SamlAccount).filter(
-            SamlAccount.user_id == user_to_delete.id
-        ).delete()
-        db_session.query(DocumentSet__User).filter(
-            DocumentSet__User.user_id == user_to_delete.id
-        ).delete()
-        db_session.query(Persona__User).filter(
-            Persona__User.user_id == user_to_delete.id
-        ).delete()
-        db_session.query(User__UserGroup).filter(
-            User__UserGroup.user_id == user_to_delete.id
-        ).delete()
-        db_session.delete(user_to_delete)
-        db_session.commit()
-
-        # NOTE: edge case may exist with race conditions
-        # with this `invited user` scheme generally.
-        user_emails = get_invited_users()
-        remaining_users = [
-            user for user in user_emails if user != user_email.user_email
-        ]
-        write_invited_users(remaining_users)
-
+        delete_user_from_db(user_to_delete, db_session)
         logger.info(f"Deleted user {user_to_delete.email}")
-    except Exception as e:
-        import traceback
 
-        full_traceback = traceback.format_exc()
-        logger.error(f"Full stack trace:\n{full_traceback}")
+    except Exception as e:
         db_session.rollback()
         logger.error(f"Error deleting user {user_to_delete.email}: {str(e)}")
         raise HTTPException(status_code=500, detail="Error deleting user")

web/src/components/admin/users/SignedUpUserTable.tsx

@@ -13,6 +13,9 @@ import { TableHeader } from "@/components/ui/table";
 import { UserRoleDropdown } from "./buttons/UserRoleDropdown";
 import { DeleteUserButton } from "./buttons/DeleteUserButton";
 import { DeactivaterButton } from "./buttons/DeactivaterButton";
+import { useUser } from "@/components/user/UserProvider";
+import { LeaveOrganizationButton } from "./buttons/LeaveOrganizationButton";
+import { NEXT_PUBLIC_CLOUD_ENABLED } from "@/lib/constants";
 
 interface Props {
   users: Array<User>;
@@ -28,6 +31,8 @@ const SignedUpUserTable = ({
   onPageChange,
   mutate,
 }: Props & PageSelectorProps) => {
+  const { user: currentUser } = useUser();
+
   if (!users.length) return null;
 
   const handlePopup = (message: string, type: "success" | "error") => {
@@ -81,18 +86,30 @@ const SignedUpUserTable = ({
                 </TableCell>
                 <TableCell>
                   <div className="flex justify-end  gap-x-2">
-                    <DeactivaterButton
-                      user={user}
-                      deactivate={user.status === UserStatus.live}
-                      setPopup={setPopup}
-                      mutate={mutate}
-                    />
-                    {user.status == UserStatus.deactivated && (
-                      <DeleteUserButton
+                    {NEXT_PUBLIC_CLOUD_ENABLED &&
+                    user.id === currentUser?.id ? (
+                      <LeaveOrganizationButton
                         user={user}
                         setPopup={setPopup}
                         mutate={mutate}
                       />
+                    ) : (
+                      <>
+                        <DeactivaterButton
+                          user={user}
+                          deactivate={user.status === UserStatus.live}
+                          setPopup={setPopup}
+                          mutate={mutate}
+                        />
+
+                        {user.status == UserStatus.deactivated && (
+                          <DeleteUserButton
+                            user={user}
+                            setPopup={setPopup}
+                            mutate={mutate}
+                          />
+                        )}
+                      </>
                     )}
                   </div>
                 </TableCell>

web/src/components/admin/users/buttons/DeactivaterButton.tsx

@@ -1,36 +1,8 @@
-import {
-  type User,
-  UserStatus,
-  UserRole,
-  USER_ROLE_LABELS,
-  INVALID_ROLE_HOVER_TEXT,
-} from "@/lib/types";
-import { type PageSelectorProps } from "@/components/PageSelector";
-import { HidableSection } from "@/app/admin/assistants/HidableSection";
+import { type User } from "@/lib/types";
 import { PopupSpec } from "@/components/admin/connectors/Popup";
 import userMutationFetcher from "@/lib/admin/users/userMutationFetcher";
 import useSWRMutation from "swr/mutation";
-import {
-  Table,
-  TableHead,
-  TableRow,
-  TableBody,
-  TableCell,
-} from "@/components/ui/table";
-
-import {
-  Select,
-  SelectContent,
-  SelectItem,
-  SelectTrigger,
-  SelectValue,
-} from "@/components/ui/select";
 import { Button } from "@/components/ui/button";
-import { GenericConfirmModal } from "@/components/modals/GenericConfirmModal";
-import { useState } from "react";
-import { usePaidEnterpriseFeaturesEnabled } from "@/components/settings/usePaidEnterpriseFeaturesEnabled";
-import { DeleteEntityModal } from "@/components/modals/DeleteEntityModal";
-import { TableHeader } from "@/components/ui/table";
 
 export const DeactivaterButton = ({
   user,

web/src/components/admin/users/buttons/LeaveOrganizationButton.tsx

@@ -0,0 +1,70 @@
+import { type User } from "@/lib/types";
+import { PopupSpec } from "@/components/admin/connectors/Popup";
+import userMutationFetcher from "@/lib/admin/users/userMutationFetcher";
+import useSWRMutation from "swr/mutation";
+import { Button } from "@/components/ui/button";
+import { useState } from "react";
+import { DeleteEntityModal } from "@/components/modals/DeleteEntityModal";
+import { useRouter } from "next/navigation";
+
+export const LeaveOrganizationButton = ({
+  user,
+  setPopup,
+  mutate,
+}: {
+  user: User;
+  setPopup: (spec: PopupSpec) => void;
+  mutate: () => void;
+}) => {
+  const router = useRouter();
+  const { trigger, isMutating } = useSWRMutation(
+    "/api/tenants/leave-organization",
+    userMutationFetcher,
+    {
+      onSuccess: () => {
+        mutate();
+        setPopup({
+          message: "Successfully left the organization!",
+          type: "success",
+        });
+      },
+      onError: (errorMsg) =>
+        setPopup({
+          message: `Unable to leave organization - ${errorMsg}`,
+          type: "error",
+        }),
+    }
+  );
+
+  const [showLeaveModal, setShowLeaveModal] = useState(false);
+
+  const handleLeaveOrganization = async () => {
+    await trigger({ user_email: user.email, method: "POST" });
+    router.push("/");
+  };
+
+  return (
+    <>
+      {showLeaveModal && (
+        <DeleteEntityModal
+          deleteButtonText="Leave"
+          entityType="organization"
+          entityName="your organization"
+          onClose={() => setShowLeaveModal(false)}
+          onSubmit={handleLeaveOrganization}
+          additionalDetails="You will lose access to all organization data and resources."
+        />
+      )}
+
+      <Button
+        className="w-min"
+        onClick={() => setShowLeaveModal(true)}
+        disabled={isMutating}
+        size="sm"
+        variant="destructive"
+      >
+        Leave Organization
+      </Button>
+    </>
+  );
+};

web/src/components/modals/DeleteEntityModal.tsx

@@ -8,21 +8,26 @@ export const DeleteEntityModal = ({
   entityType,
   entityName,
   additionalDetails,
+  deleteButtonText,
 }: {
   entityType: string;
   entityName: string;
   onClose: () => void;
   onSubmit: () => void;
   additionalDetails?: string;
+  deleteButtonText?: string;
 }) => {
   return (
     <Modal onOutsideClick={onClose}>
       <>
         <div className="flex mb-4">
-          <h2 className="my-auto text-2xl font-bold">Delete {entityType}?</h2>
+          <h2 className="my-auto text-2xl font-bold">
+            {deleteButtonText || `Delete`} {entityType}
+          </h2>
         </div>
         <p className="mb-4">
-          Click below to confirm that you want to delete <b>{entityName}</b>
+          Click below to confirm that you want to {deleteButtonText || "delete"}{" "}
+          <b>{entityName}</b>
         </p>
         {additionalDetails && <p className="mb-4">{additionalDetails}</p>}
         <div className="flex">
@@ -30,7 +35,7 @@ export const DeleteEntityModal = ({
             <BasicClickable onClick={onSubmit}>
               <div className="flex mx-2">
                 <FiTrash className="my-auto mr-2" />
-                Delete
+                {deleteButtonText || "Delete"}
               </div>
             </BasicClickable>
           </div>

web/src/lib/userSS.ts

@@ -115,6 +115,7 @@ export const getAuthUrlSS = async (
       return await getGoogleOAuthUrlSS(nextUrl);
     }
     case "cloud": {
+      console.log("LLpp");
       return await getGoogleOAuthUrlSS(nextUrl);
     }
     case "saml": {