diff --git a/changes/12706.feature.md b/changes/12706.feature.md new file mode 100644 index 00000000000..3280f7fa543 --- /dev/null +++ b/changes/12706.feature.md @@ -0,0 +1 @@ +Add the AppConfigFragment visible-fragments repository query: fetch the public / domain / user fragments visible to a (domain, user) principal for one or many config names, rank-ordered via the allow-list, in a single query. diff --git a/src/ai/backend/manager/models/app_config_fragment/conditions.py b/src/ai/backend/manager/models/app_config_fragment/conditions.py index 53d3d792131..a8f4d836c13 100644 --- a/src/ai/backend/manager/models/app_config_fragment/conditions.py +++ b/src/ai/backend/manager/models/app_config_fragment/conditions.py @@ -72,6 +72,19 @@ def inner() -> sa.sql.expression.ColumnElement[bool]: by_config_name_in = staticmethod(make_string_in_factory(AppConfigFragmentRow.config_name)) + @staticmethod + def by_config_names(config_names: Sequence[str]) -> QueryCondition: + """Fragments whose ``config_name`` is one of ``config_names``. + + The plain-name membership filter the merged read AND-combines with a visibility + filter — kept separate from the ``by_*_visibility`` scope builders. + """ + + def inner() -> sa.sql.expression.ColumnElement[bool]: + return AppConfigFragmentRow.config_name.in_(list(config_names)) + + return inner + # --- scope_type enum filters --- @staticmethod @@ -111,6 +124,41 @@ def inner() -> sa.sql.expression.ColumnElement[bool]: return inner + # --- per-scope visibility filters (one scope_type each), independent of config_name --- + + @staticmethod + def by_public_visibility() -> QueryCondition: + """The ``public`` scope (public has no per-entity scope_id).""" + + def inner() -> sa.sql.expression.ColumnElement[bool]: + return AppConfigFragmentRow.scope_type == AppConfigScopeType.PUBLIC + + return inner + + @staticmethod + def by_domain_visibility(domain: str) -> QueryCondition: + """The ``domain`` scope for ``domain``.""" + + def inner() -> sa.sql.expression.ColumnElement[bool]: + return sa.and_( + AppConfigFragmentRow.scope_type == AppConfigScopeType.DOMAIN, + AppConfigFragmentRow.scope_id == domain, + ) + + return inner + + @staticmethod + def by_user_visibility(user_id: str) -> QueryCondition: + """The ``user`` scope for ``user_id``.""" + + def inner() -> sa.sql.expression.ColumnElement[bool]: + return sa.and_( + AppConfigFragmentRow.scope_type == AppConfigScopeType.USER, + AppConfigFragmentRow.scope_id == user_id, + ) + + return inner + # --- created_at / updated_at datetime filters --- @staticmethod diff --git a/src/ai/backend/manager/repositories/app_config_fragment/db_source/db_source.py b/src/ai/backend/manager/repositories/app_config_fragment/db_source/db_source.py index ecdfedbb71c..45513ab2872 100644 --- a/src/ai/backend/manager/repositories/app_config_fragment/db_source/db_source.py +++ b/src/ai/backend/manager/repositories/app_config_fragment/db_source/db_source.py @@ -21,12 +21,18 @@ from ai.backend.manager.errors.app_config import ( AppConfigFragmentNotFound, ) +from ai.backend.manager.models.app_config_allow_list.row import AppConfigAllowListRow +from ai.backend.manager.models.app_config_fragment.conditions import AppConfigFragmentConditions from ai.backend.manager.models.app_config_fragment.row import AppConfigFragmentRow from ai.backend.manager.models.scopes import SearchScope +from ai.backend.manager.repositories.app_config_fragment.types import ( + AppConfigScopeArguments, +) from ai.backend.manager.repositories.base import ( BatchQuerier, BulkCreator, Creator, + NoPagination, Purger, Querier, Updater, @@ -193,3 +199,44 @@ async def scoped_search( has_next_page=result.has_next_page, has_previous_page=result.has_previous_page, ) + + @app_config_fragment_db_source_resilience.apply() + async def list_visible_fragments_bulk( + self, config_names: list[str], scope: AppConfigScopeArguments + ) -> list[AppConfigFragmentData]: + """Visible fragments for several ``config_names`` at once, in a single query. + + Selects the requested names AND any one of the principal's visible scopes (public, + its domain, or its own user). The scope filter is name-independent, so it is a single + OR group AND-combined with the name membership. Merge priority (``rank``) lives on the + joined allow-list entry; the result is always ordered by ascending ``rank`` so the + caller can group by name (each name's subset stays rank-ordered) and deep-merge each + name's fragments in order. + """ + if not config_names: + return [] + scope_visibility = [ + AppConfigFragmentConditions.by_public_visibility(), + AppConfigFragmentConditions.by_domain_visibility(str(scope.domain_id)), + AppConfigFragmentConditions.by_user_visibility(str(scope.user_id)), + ] + querier = BatchQuerier( + pagination=NoPagination(), + conditions=[ + AppConfigFragmentConditions.by_config_names(config_names), + lambda: sa.or_(*(visibility() for visibility in scope_visibility)), + ], + orders=[AppConfigAllowListRow.rank.asc()], + ) + # Join each fragment to its allow-list entry (indexed ``(config_name, scope_type)`` FK + # pair), which carries the merge ``rank`` the result is ordered by. + selector = sa.select(AppConfigFragmentRow).join( + AppConfigAllowListRow, + sa.and_( + AppConfigAllowListRow.config_name == AppConfigFragmentRow.config_name, + AppConfigAllowListRow.scope_type == AppConfigFragmentRow.scope_type, + ), + ) + async with self._ops.read_ops() as r: + result = await r.batch_query_in_global(selector, querier) + return [row.AppConfigFragmentRow.to_data() for row in result.rows] diff --git a/src/ai/backend/manager/repositories/app_config_fragment/repository.py b/src/ai/backend/manager/repositories/app_config_fragment/repository.py index 0ade21b67fd..43334e915c5 100644 --- a/src/ai/backend/manager/repositories/app_config_fragment/repository.py +++ b/src/ai/backend/manager/repositories/app_config_fragment/repository.py @@ -18,6 +18,9 @@ from ai.backend.manager.repositories.app_config_fragment.db_source import ( AppConfigFragmentDBSource, ) +from ai.backend.manager.repositories.app_config_fragment.types import ( + AppConfigScopeArguments, +) from ai.backend.manager.repositories.base import ( BatchQuerier, BulkCreator, @@ -103,3 +106,9 @@ async def bulk_purge( purgers: Sequence[Purger[AppConfigFragmentRow]], ) -> AppConfigFragmentBulkResult: return await self._db_source.bulk_purge(purgers) + + @app_config_fragment_repository_resilience.apply() + async def list_visible_fragments_bulk( + self, config_names: list[str], scope: AppConfigScopeArguments + ) -> list[AppConfigFragmentData]: + return await self._db_source.list_visible_fragments_bulk(config_names, scope) diff --git a/src/ai/backend/manager/repositories/app_config_fragment/types.py b/src/ai/backend/manager/repositories/app_config_fragment/types.py index 190dd9def78..39e7e3c7750 100644 --- a/src/ai/backend/manager/repositories/app_config_fragment/types.py +++ b/src/ai/backend/manager/repositories/app_config_fragment/types.py @@ -1,4 +1,4 @@ -"""Types for app config fragment repository operations (search scopes).""" +"""Types for app config fragment repository operations (search scopes, resolve arguments).""" from __future__ import annotations @@ -16,11 +16,25 @@ from ai.backend.manager.models.scopes import ExistenceCheck, SearchScope __all__ = ( + "AppConfigScopeArguments", "DomainAppConfigFragmentSearchScope", "UserAppConfigFragmentSearchScope", ) +@dataclass(frozen=True) +class AppConfigScopeArguments: + """The principal an ``AppConfig`` is resolved for: the resolving user and its domain. + + Bundles the scope-identifying arguments so they travel together (add new principal + dimensions here rather than growing method signatures). Plain value object — not a + :class:`SearchScope`. + """ + + domain_id: DomainID + user_id: UserID + + @dataclass(frozen=True) class DomainAppConfigFragmentSearchScope(SearchScope): """Fragments written at one domain scope (``scope_type == domain``, ``scope_id == domain_id``). diff --git a/tests/unit/manager/repositories/app_config_fragment/test_repository.py b/tests/unit/manager/repositories/app_config_fragment/test_repository.py index 78300007086..cf1a624040f 100644 --- a/tests/unit/manager/repositories/app_config_fragment/test_repository.py +++ b/tests/unit/manager/repositories/app_config_fragment/test_repository.py @@ -31,6 +31,7 @@ AppConfigFragmentRepository, ) from ai.backend.manager.repositories.app_config_fragment.types import ( + AppConfigScopeArguments, DomainAppConfigFragmentSearchScope, UserAppConfigFragmentSearchScope, ) @@ -569,3 +570,137 @@ async def test_partial_when_one_missing( assert [f.index for f in result.failed] == [1] with pytest.raises(AppConfigFragmentNotFound): await repository.get_by_id(two_fragments[0].id) + + +class TestVisibilityConditions: + async def test_public_visibility_selects_only_public( + self, + repository: AppConfigFragmentRepository, + fragments_across_scopes: list[AppConfigFragmentData], + ) -> None: + result = await repository.admin_search( + BatchQuerier( + pagination=OffsetPagination(limit=10, offset=0), + conditions=[AppConfigFragmentConditions.by_public_visibility()], + ) + ) + # Visibility is scope-only (name-independent): every public fragment, any name. + expected = { + f.id for f in fragments_across_scopes if f.scope_type is AppConfigScopeType.PUBLIC + } + assert {item.id for item in result.items} == expected + + async def test_domain_visibility_selects_only_that_domain( + self, + repository: AppConfigFragmentRepository, + fragments_across_scopes: list[AppConfigFragmentData], + ) -> None: + result = await repository.admin_search( + BatchQuerier( + pagination=OffsetPagination(limit=10, offset=0), + conditions=[AppConfigFragmentConditions.by_domain_visibility(_DOMAIN_ID)], + ) + ) + expected = { + f.id + for f in fragments_across_scopes + if f.scope_type is AppConfigScopeType.DOMAIN and f.scope_id == _DOMAIN_ID + } + assert {item.id for item in result.items} == expected + + async def test_user_visibility_selects_only_that_user( + self, + repository: AppConfigFragmentRepository, + fragments_across_scopes: list[AppConfigFragmentData], + ) -> None: + result = await repository.admin_search( + BatchQuerier( + pagination=OffsetPagination(limit=10, offset=0), + conditions=[AppConfigFragmentConditions.by_user_visibility(_USER_ID)], + ) + ) + expected = { + f.id + for f in fragments_across_scopes + if f.scope_type is AppConfigScopeType.USER and f.scope_id == _USER_ID + } + assert {item.id for item in result.items} == expected + + +class TestApplicableFragments: + async def test_one_query_returns_public_domain_user_rank_ordered( + self, + repository: AppConfigFragmentRepository, + fragments_across_scopes: list[AppConfigFragmentData], + ) -> None: + applicable = await repository.list_visible_fragments_bulk( + ["theme"], + AppConfigScopeArguments(domain_id=DomainID(_DOMAIN_UUID), user_id=UserID(_USER_UUID)), + ) + # public + the caller's domain + the caller's own user fragment, ordered by the + # allow-list entries' ranks (scope-type defaults: public < domain < user). + expected = [ + f + for f in fragments_across_scopes + if f.config_name == "theme" + and ( + f.scope_type is AppConfigScopeType.PUBLIC + or (f.scope_type is AppConfigScopeType.DOMAIN and f.scope_id == _DOMAIN_ID) + or (f.scope_type is AppConfigScopeType.USER and f.scope_id == _USER_ID) + ) + ] + assert [f.id for f in applicable] == [f.id for f in expected] + assert [f.scope_type for f in applicable] == [ + AppConfigScopeType.PUBLIC, + AppConfigScopeType.DOMAIN, + AppConfigScopeType.USER, + ] + + async def test_unknown_config_name_returns_empty( + self, + repository: AppConfigFragmentRepository, + fragments_across_scopes: list[AppConfigFragmentData], + ) -> None: + applicable = await repository.list_visible_fragments_bulk( + ["unregistered"], + AppConfigScopeArguments(domain_id=DomainID(_DOMAIN_UUID), user_id=UserID(_USER_UUID)), + ) + assert applicable == [] + + async def test_bulk_returns_visible_fragments_for_all_names_ordered( + self, + repository: AppConfigFragmentRepository, + fragments_across_scopes: list[AppConfigFragmentData], + ) -> None: + applicable = await repository.list_visible_fragments_bulk( + ["theme", "menu"], + AppConfigScopeArguments(domain_id=DomainID(_DOMAIN_UUID), user_id=UserID(_USER_UUID)), + ) + # public + the caller's domain + the caller's own user fragment, for both names. + expected = { + f.id + for f in fragments_across_scopes + if f.config_name in ("theme", "menu") + and ( + f.scope_type is AppConfigScopeType.PUBLIC + or (f.scope_type is AppConfigScopeType.DOMAIN and f.scope_id == _DOMAIN_ID) + or (f.scope_type is AppConfigScopeType.USER and f.scope_id == _USER_ID) + ) + } + assert {f.id for f in applicable} == expected + # Rank-ordered globally, so each config_name's subset stays rank-ordered + # (public < domain < user) for the caller to group by name and deep-merge in order. + for name in ("theme", "menu"): + ranks = [f.scope_type.default_rank() for f in applicable if f.config_name == name] + assert ranks == sorted(ranks) + + async def test_bulk_empty_names_returns_empty( + self, + repository: AppConfigFragmentRepository, + fragments_across_scopes: list[AppConfigFragmentData], + ) -> None: + applicable = await repository.list_visible_fragments_bulk( + [], + AppConfigScopeArguments(domain_id=DomainID(_DOMAIN_UUID), user_id=UserID(_USER_UUID)), + ) + assert applicable == []