Merge pull request #404 from ynput/403-improve-performance-of-handlin…

…g-sites-in-apiinfo

/api/info performance improvements

api/onboarding/release.py

@@ -12,6 +12,7 @@
 from ayon_server.helpers.cloud import get_cloud_api_headers
 from ayon_server.installer.models import DependencyPackageManifest, InstallerManifest
 from ayon_server.lib.postgres import Postgres
+from ayon_server.lib.redis import Redis
 from ayon_server.types import OPModel
 
 from .router import router
@@ -86,6 +87,7 @@ async def abort_onboarding(request: Request, user: CurrentUser) -> EmptyResponse
         """
     )
 
+    await Redis.set("global", "onboardingFinished", "1")
     return EmptyResponse()
 
 
@@ -96,12 +98,9 @@ async def restart_onboarding(request: Request, user: CurrentUser) -> EmptyRespon
     if not user.is_admin:
         raise ForbiddenException()
 
-    await Postgres().execute(
-        """
-        DELETE FROM config WHERE key = 'onboardingFinished'
-        """
-    )
-
+    q = "DELETE FROM config WHERE key = 'onboardingFinished'"
+    await Postgres().execute(q)
+    await Redis.delete("global", "onboardingFinished")
     return EmptyResponse()
 
 

api/system/info.py

@@ -2,9 +2,10 @@
 from typing import Any
 from urllib.parse import urlparse
 
+import aiocache
 from attributes.attributes import AttributeModel  # type: ignore
 from fastapi import Request
-from nxtools import log_traceback
+from nxtools import log_traceback, logging
 from pydantic import ValidationError
 
 from ayon_server.addons import AddonLibrary, SSOOption
@@ -15,6 +16,7 @@
 from ayon_server.helpers.email import is_mailing_enabled
 from ayon_server.info import ReleaseInfo, get_release_info, get_uptime, get_version
 from ayon_server.lib.postgres import Postgres
+from ayon_server.lib.redis import Redis
 from ayon_server.types import Field, OPModel
 
 from .router import router
@@ -71,14 +73,34 @@ class InfoResponseModel(OPModel):
     sso_options: list[SSOOption] = Field(default_factory=list, title="SSO options")
 
 
+# Ensure that an admin user exists
+# This is used to determine if the 'Create admin user' form should be displayed
+# We raise an exception in ensure_admin_user_exists, so the False value is not cached
+# and when the admin user is created, we use the cache to avoid unnecessary queries
+
+
+@aiocache.cached()
+async def ensure_admin_user_exists() -> None:
+    res = await Postgres.fetch(
+        "SELECT name FROM users WHERE (data->'isAdmin')::boolean"
+    )
+    if not res:
+        raise ValueError("No admin user exists")
+    return None
+
+
 async def admin_exists() -> bool:
-    async for row in Postgres.iterate(
-        "SELECT name FROM users WHERE data->>'isAdmin' = 'true'"
-    ):
+    try:
+        await ensure_admin_user_exists()
         return True
-    return False
+    except ValueError:
+        return False
 
 
+# Get all SSO options from the active addons
+
+
+@aiocache.cached(ttl=10)
 async def get_sso_options(request: Request) -> list[SSOOption]:
     referer = request.headers.get("referer")
     if referer:
@@ -114,47 +136,76 @@ async def get_sso_options(request: Request) -> list[SSOOption]:
     return result
 
 
-async def get_additional_info(user: UserEntity, request: Request):
-    current_site: SiteInfo | None = None
+async def get_user_sites(
+    user_name: str, current_site: SiteInfo | None
+) -> list[SiteInfo]:
+    """Return a list of sites the user is registered to
 
-    with contextlib.suppress(ValidationError):
-        current_site = SiteInfo(
-            id=request.headers.get("x-ayon-site-id"),
-            platform=request.headers.get("x-ayon-platform"),
-            hostname=request.headers.get("x-ayon-hostname"),
-            version=request.headers.get("x-ayon-version"),
-            users=[user.name],
-        )
+    If site information in the request headers, it will be added to the
+    top of the listand updated in the database if necessary.
+    """
+    sites: list[SiteInfo] = []
+    current_needs_update = False
+    current_site_exists = False
 
-    sites = []
-    async for row in Postgres.iterate("SELECT id, data FROM sites"):
-        site = SiteInfo(id=row["id"], **row["data"])
+    query_id = current_site.id if current_site else ""
 
-        if current_site and site.id == current_site.id:
-            current_site.users = list(set(current_site.users + site.users))
-            continue
+    # Get all sites the user is registered to or the current site
+    query = """
+        SELECT id, data FROM sites
+        WHERE id = $1 OR data->'users' ? $2
+    """
 
-        if user.name not in site.users:
+    async for row in Postgres.iterate(query, query_id, user_name):
+        site = SiteInfo(id=row["id"], **row["data"])
+        if current_site and site.id == current_site.id:
+            # record matches the current site
+            current_site_exists = True
+            if user_name not in site.users:
+                current_site.users.update(site.users)
+                current_needs_update = True
+            # we can use elif here, because we only need to check one condition
+            elif site.platform != current_site.platform:
+                current_needs_update = True
+            elif site.hostname != current_site.hostname:
+                current_needs_update = True
+            elif site.version != current_site.version:
+                current_needs_update = True
+            # do not add the current site to the list,
+            # we'll insert it at the beginning at the end of the loop
             continue
-
         sites.append(site)
 
     if current_site:
-        mdata = current_site.dict()
-        mid = mdata.pop("id")
-        await Postgres.execute(
-            """
-            INSERT INTO sites (id, data)
-            VALUES ($1, $2) ON CONFLICT (id)
-            DO UPDATE SET data = EXCLUDED.data
-            """,
-            mid,
-            mdata,
-        )
+        # if the current site is not in the database
+        # or has been changed, upsert it
+        if current_needs_update or not current_site_exists:
+            logging.debug(f"Registering to site {current_site.id}", user=user_name)
+            mdata = current_site.dict()
+            mid = mdata.pop("id")
+            await Postgres.execute(
+                """
+                INSERT INTO sites (id, data)
+                VALUES ($1, $2) ON CONFLICT (id)
+                DO UPDATE SET data = EXCLUDED.data
+                """,
+                mid,
+                mdata,
+            )
 
+        # insert the current site at the beginning of the list
         sites.insert(0, current_site)
+    return sites
+
+
+@aiocache.cached(ttl=5)
+async def get_attributes() -> list[AttributeModel]:
+    """Return a list of available attributes
+
+    populate enum fields with values from the database
+    in the case dynamic enums are used.
+    """
 
-    # load dynamic_enums
     enums: dict[str, Any] = {}
     async for row in Postgres.iterate(
         "SELECT name, data FROM attributes WHERE data->'enum' is not null"
@@ -171,12 +222,55 @@ async def get_additional_info(user: UserEntity, request: Request):
         except ValidationError:
             log_traceback(f"Invalid attribute data: {row}")
             continue
+    return attr_list
+
+
+async def get_additional_info(user: UserEntity, request: Request):
+    """Return additional information for the user
+
+    This is returned only if the user is logged in.
+    """
+
+    # Handle site information
+
+    current_site: SiteInfo | None = None
+    with contextlib.suppress(ValidationError):
+        current_site = SiteInfo(
+            id=request.headers.get("x-ayon-site-id"),
+            platform=request.headers.get("x-ayon-platform"),
+            hostname=request.headers.get("x-ayon-hostname"),
+            version=request.headers.get("x-ayon-version"),
+            users=[user.name],
+        )
+
+    sites = await get_user_sites(user.name, current_site)
+
+    attr_list = await get_attributes()
+
     return {
         "attributes": attr_list,
         "sites": sites,
     }
 
 
+async def is_onboarding_finished() -> bool:
+    r = await Redis.get("global", "onboardingFinished")
+    if r is None:
+        query = "SELECT * FROM config where key = 'onboardingFinished'"
+        rdb = await Postgres.fetch(query)
+        if rdb:
+            await Redis.set("global", "onboardingFinished", "1")
+            return True
+    elif r:
+        return True
+    return False
+
+
+#
+# The actual endpoint
+#
+
+
 @router.get("/info", response_model_exclude_none=True, tags=["System"])
 async def get_site_info(
     request: Request,
@@ -194,13 +288,9 @@ async def get_site_info(
     if current_user:
         additional_info = await get_additional_info(current_user, request)
 
-        if current_user.is_admin:
-            res = await Postgres.fetch(
-                """SELECT * FROM config where key = 'onboardingFinished'"""
-            )
-            if not res:
+        if current_user.is_admin and not current_user.is_service:
+            if not await is_onboarding_finished():
                 additional_info["onboarding"] = True
-
     else:
         sso_options = await get_sso_options(request)
         has_admin_user = await admin_exists()

api/system/sites.py

@@ -12,7 +12,7 @@ class SiteInfo(OPModel):
     platform: Platform = Field(...)
     hostname: str = Field(..., title="Machine hostname")
     version: str = Field(..., title="Ayon version")
-    users: list[str] = Field(..., title="List of users")
+    users: set[str] = Field(..., title="List of users")
 
 
 @router.get("/system/sites", tags=["System"])

ayon_server/utils.py

@@ -49,6 +49,9 @@ def json_default_handler(value: Any) -> Any:
     if isinstance(value, datetime.datetime):
         return value.isoformat()
 
+    if isinstance(value, set):
+        return list(value)
+
     raise TypeError(f"Type {type(value)} is not JSON serializable")