diff --git a/backend/danswer/db/auth.py b/backend/danswer/db/auth.py index 53fba1e52..0a794f3c9 100644 --- a/backend/danswer/db/auth.py +++ b/backend/danswer/db/auth.py @@ -22,6 +22,7 @@ from danswer.db.models import User from danswer.utils.variable_functionality import ( fetch_versioned_implementation_with_fallback, ) +from ee.danswer.db.api_key import get_api_key_email_pattern def get_default_admin_user_emails() -> list[str]: @@ -35,12 +36,16 @@ def get_default_admin_user_emails() -> list[str]: return get_default_admin_user_emails_fn() -def get_total_users(db_session: Session) -> int: +def get_total_users_count(db_session: Session) -> int: """ Returns the total number of users in the system. This is the sum of users and invited users. """ - user_count = db_session.query(User).count() + user_count = ( + db_session.query(User) + .filter(~User.email.endswith(get_api_key_email_pattern())) # type: ignore + .count() + ) invited_users = len(get_invited_users()) return user_count + invited_users diff --git a/backend/danswer/server/manage/users.py b/backend/danswer/server/manage/users.py index 9cb1cec1c..f0675c392 100644 --- a/backend/danswer/server/manage/users.py +++ b/backend/danswer/server/manage/users.py @@ -37,7 +37,7 @@ from danswer.configs.app_configs import ENABLE_EMAIL_INVITES from danswer.configs.app_configs import SESSION_EXPIRE_TIME_SECONDS from danswer.configs.app_configs import VALID_EMAIL_DOMAINS from danswer.configs.constants import AuthType -from danswer.db.auth import get_total_users +from danswer.db.auth import get_total_users_count from danswer.db.engine import CURRENT_TENANT_ID_CONTEXTVAR from danswer.db.engine import get_session from danswer.db.models import AccessToken @@ -226,7 +226,7 @@ def bulk_invite_users( try: logger.info("Registering tenant users") register_tenant_users( - CURRENT_TENANT_ID_CONTEXTVAR.get(), get_total_users(db_session) + CURRENT_TENANT_ID_CONTEXTVAR.get(), get_total_users_count(db_session) ) if ENABLE_EMAIL_INVITES: try: @@ -262,7 +262,7 @@ def remove_invited_user( try: if MULTI_TENANT: register_tenant_users( - CURRENT_TENANT_ID_CONTEXTVAR.get(), get_total_users(db_session) + CURRENT_TENANT_ID_CONTEXTVAR.get(), get_total_users_count(db_session) ) except Exception: logger.error( diff --git a/backend/ee/danswer/db/api_key.py b/backend/ee/danswer/db/api_key.py index c38f32a0f..8bbdf7eaa 100644 --- a/backend/ee/danswer/db/api_key.py +++ b/backend/ee/danswer/db/api_key.py @@ -17,8 +17,12 @@ from ee.danswer.auth.api_key import hash_api_key from ee.danswer.server.api_key.models import APIKeyArgs +def get_api_key_email_pattern() -> str: + return DANSWER_API_KEY_DUMMY_EMAIL_DOMAIN + + def is_api_key_email_address(email: str) -> bool: - return email.endswith(f"{DANSWER_API_KEY_DUMMY_EMAIL_DOMAIN}") + return email.endswith(get_api_key_email_pattern()) def fetch_api_keys(db_session: Session) -> list[ApiKeyDescriptor]: