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 ( from danswer.utils.variable_functionality import (
fetch_versioned_implementation_with_fallback, 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]: 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() 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. Returns the total number of users in the system.
This is the sum of users and invited users. 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()) invited_users = len(get_invited_users())
return user_count + 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 SESSION_EXPIRE_TIME_SECONDS
from danswer.configs.app_configs import VALID_EMAIL_DOMAINS from danswer.configs.app_configs import VALID_EMAIL_DOMAINS
from danswer.configs.constants import AuthType 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 CURRENT_TENANT_ID_CONTEXTVAR
from danswer.db.engine import get_session from danswer.db.engine import get_session
from danswer.db.models import AccessToken from danswer.db.models import AccessToken
@ -226,7 +226,7 @@ def bulk_invite_users(
try: try:
logger.info("Registering tenant users") logger.info("Registering tenant users")
register_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: if ENABLE_EMAIL_INVITES:
try: try:
@ -262,7 +262,7 @@ def remove_invited_user(
try: try:
if MULTI_TENANT: if MULTI_TENANT:
register_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)
) )
except Exception: except Exception:
logger.error( 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 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: 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]: def fetch_api_keys(db_session: Session) -> list[ApiKeyDescriptor]: