From 7c17ba61165d1b7abb3ff488a6c67fe6666b2bbe Mon Sep 17 00:00:00 2001 From: Martastain Date: Thu, 24 Oct 2024 09:19:47 +0200 Subject: [PATCH 1/5] fix: cache admin user exists --- api/system/info.py | 32 ++++++++++++++++++++++++++++---- 1 file changed, 28 insertions(+), 4 deletions(-) diff --git a/api/system/info.py b/api/system/info.py index 1a3c0413..487f4418 100644 --- a/api/system/info.py +++ b/api/system/info.py @@ -2,6 +2,7 @@ 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 @@ -71,14 +72,32 @@ 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'") + 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=30) async def get_sso_options(request: Request) -> list[SSOOption]: referer = request.headers.get("referer") if referer: @@ -177,6 +196,11 @@ async def get_additional_info(user: UserEntity, request: Request): } +# +# The actual endpoint +# + + @router.get("/info", response_model_exclude_none=True, tags=["System"]) async def get_site_info( request: Request, From 582824c18e8f521cd8a054f650b857cda7cf0002 Mon Sep 17 00:00:00 2001 From: Martastain Date: Thu, 24 Oct 2024 11:42:08 +0200 Subject: [PATCH 2/5] fix: optimize site update in info --- api/system/info.py | 106 ++++++++++++++++++++++++++++++++------------- 1 file changed, 75 insertions(+), 31 deletions(-) diff --git a/api/system/info.py b/api/system/info.py index 487f4418..bfa780c9 100644 --- a/api/system/info.py +++ b/api/system/info.py @@ -5,7 +5,7 @@ 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 @@ -80,7 +80,9 @@ class InfoResponseModel(OPModel): @aiocache.cached() async def ensure_admin_user_exists() -> None: - res = await Postgres.fetch("SELECT name FROM users WHERE data->'isAdmin'") + res = await Postgres.fetch( + "SELECT name FROM users WHERE (data->'isAdmin')::boolean" + ) if not res: raise ValueError("No admin user exists") return None @@ -133,9 +135,76 @@ async def get_sso_options(request: Request) -> list[SSOOption]: return result +async def get_user_sites( + user_name: str, current_site: SiteInfo | None +) -> list[SiteInfo]: + """Return a list of sites the user is registered to + + 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 + + query_id = current_site.id if current_site else "" + + # 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 + """ + + 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.extend(site.users) + current_needs_update = True + if site.platform != current_site.platform: + current_needs_update = True + if site.hostname != current_site.hostname: + current_needs_update = True + if 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: + # 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 + + async def get_additional_info(user: UserEntity, request: Request): - current_site: SiteInfo | None = None + """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"), @@ -145,35 +214,10 @@ async def get_additional_info(user: UserEntity, request: Request): users=[user.name], ) - sites = [] - async for row in Postgres.iterate("SELECT id, data FROM sites"): - site = SiteInfo(id=row["id"], **row["data"]) - - if current_site and site.id == current_site.id: - current_site.users = list(set(current_site.users + site.users)) - continue - - if user.name not in site.users: - 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, - ) - - sites.insert(0, current_site) + sites = await get_user_sites(user.name, current_site) # load dynamic_enums + enums: dict[str, Any] = {} async for row in Postgres.iterate( "SELECT name, data FROM attributes WHERE data->'enum' is not null" @@ -218,7 +262,7 @@ async def get_site_info( if current_user: additional_info = await get_additional_info(current_user, request) - if current_user.is_admin: + if current_user.is_admin and not current_user.is_service: res = await Postgres.fetch( """SELECT * FROM config where key = 'onboardingFinished'""" ) From b30bed0e50d06453aea84db669fb333e6d250f4a Mon Sep 17 00:00:00 2001 From: Martastain Date: Thu, 24 Oct 2024 12:23:14 +0200 Subject: [PATCH 3/5] fix: cache attribute list and onboardingfinished --- api/onboarding/release.py | 11 +++---- api/system/info.py | 65 ++++++++++++++++++++++++++------------- 2 files changed, 48 insertions(+), 28 deletions(-) diff --git a/api/onboarding/release.py b/api/onboarding/release.py index 5d9dcbf6..2c1e30bc 100644 --- a/api/onboarding/release.py +++ b/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() diff --git a/api/system/info.py b/api/system/info.py index bfa780c9..549efda2 100644 --- a/api/system/info.py +++ b/api/system/info.py @@ -16,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 @@ -196,6 +197,33 @@ async def get_user_sites( 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. + """ + + enums: dict[str, Any] = {} + async for row in Postgres.iterate( + "SELECT name, data FROM attributes WHERE data->'enum' is not null" + ): + enums[row["name"]] = row["data"]["enum"] + + attr_list: list[AttributeModel] = [] + for row in attribute_library.info_data: + row = {**row} + if row["name"] in enums: + row["data"]["enum"] = enums[row["name"]] + try: + attr_list.append(AttributeModel(**row)) + 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 @@ -216,30 +244,27 @@ async def get_additional_info(user: UserEntity, request: Request): sites = await get_user_sites(user.name, current_site) - # load dynamic_enums + attr_list = await get_attributes() - enums: dict[str, Any] = {} - async for row in Postgres.iterate( - "SELECT name, data FROM attributes WHERE data->'enum' is not null" - ): - enums[row["name"]] = row["data"]["enum"] - - attr_list: list[AttributeModel] = [] - for row in attribute_library.info_data: - row = {**row} - if row["name"] in enums: - row["data"]["enum"] = enums[row["name"]] - try: - attr_list.append(AttributeModel(**row)) - except ValidationError: - log_traceback(f"Invalid attribute data: {row}") - continue 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 # @@ -263,12 +288,8 @@ async def get_site_info( additional_info = await get_additional_info(current_user, request) if current_user.is_admin and not current_user.is_service: - res = await Postgres.fetch( - """SELECT * FROM config where key = 'onboardingFinished'""" - ) - if not res: + if not await is_onboarding_finished(): additional_info["onboarding"] = True - else: sso_options = await get_sso_options(request) has_admin_user = await admin_exists() From f2501c2108af3f831a879681da027bcec149798f Mon Sep 17 00:00:00 2001 From: Martastain Date: Thu, 24 Oct 2024 12:30:05 +0200 Subject: [PATCH 4/5] chore: reduced sso list caching ttl --- api/system/info.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/api/system/info.py b/api/system/info.py index 549efda2..89eec38f 100644 --- a/api/system/info.py +++ b/api/system/info.py @@ -100,7 +100,7 @@ async def admin_exists() -> bool: # Get all SSO options from the active addons -@aiocache.cached(ttl=30) +@aiocache.cached(ttl=10) async def get_sso_options(request: Request) -> list[SSOOption]: referer = request.headers.get("referer") if referer: @@ -164,11 +164,12 @@ async def get_user_sites( if user_name not in site.users: current_site.users.extend(site.users) current_needs_update = True - if site.platform != current_site.platform: + # we can use elif here, because we only need to check one condition + elif site.platform != current_site.platform: current_needs_update = True - if site.hostname != current_site.hostname: + elif site.hostname != current_site.hostname: current_needs_update = True - if site.version != current_site.version: + 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 From f210dbf2af10adc5d1bcdaf162bf7b12c60ef535 Mon Sep 17 00:00:00 2001 From: Martastain Date: Thu, 24 Oct 2024 12:54:58 +0200 Subject: [PATCH 5/5] feat: use set serializer for user list in sites --- api/system/info.py | 2 +- api/system/sites.py | 2 +- ayon_server/utils.py | 3 +++ 3 files changed, 5 insertions(+), 2 deletions(-) diff --git a/api/system/info.py b/api/system/info.py index 89eec38f..32729298 100644 --- a/api/system/info.py +++ b/api/system/info.py @@ -162,7 +162,7 @@ async def get_user_sites( # record matches the current site current_site_exists = True if user_name not in site.users: - current_site.users.extend(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: diff --git a/api/system/sites.py b/api/system/sites.py index e4d37c29..7b23a4e7 100644 --- a/api/system/sites.py +++ b/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"]) diff --git a/ayon_server/utils.py b/ayon_server/utils.py index ad312d3d..104b47a1 100644 --- a/ayon_server/utils.py +++ b/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")