Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions changes/12706.feature.md
Original file line number Diff line number Diff line change
@@ -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.
48 changes: 48 additions & 0 deletions src/ai/backend/manager/models/app_config_fragment/conditions.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Comment on lines +138 to +160

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Did you decide that managing it this way is better than generalizing it into a single condition? (I’m asking just to confirm, since this seems like a design choice.)


# --- created_at / updated_at datetime filters ---

@staticmethod
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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)),
Comment on lines +218 to +221

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It doesn't seem necessary to manage this as a separate variable.

]
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,
),
)
Comment on lines +231 to +239

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Isn't the "join" unnecessary?

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]
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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)
Original file line number Diff line number Diff line change
@@ -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

Expand All @@ -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``).
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,7 @@
AppConfigFragmentRepository,
)
from ai.backend.manager.repositories.app_config_fragment.types import (
AppConfigScopeArguments,
DomainAppConfigFragmentSearchScope,
UserAppConfigFragmentSearchScope,
)
Expand Down Expand Up @@ -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 == []
Loading