Do not count API keys as users (#3022)

* don't count api keys as users

* typing
This commit is contained in:
pablodanswer 2024-11-01 12:34:30 -07:00 committed by GitHub
parent ccdc09e2d4
commit 6d543f3d4f
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 15 additions and 6 deletions

View File

@ -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

View File

@ -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(

View File

@ -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]: