refactor(BA-5650): add owner_id helpers and rename session data types [1/4]#11046
Closed
jopemachine wants to merge 22 commits into
Closed
refactor(BA-5650): add owner_id helpers and rename session data types [1/4]#11046jopemachine wants to merge 22 commits into
jopemachine wants to merge 22 commits into
Conversation
This was referenced Apr 14, 2026
This was referenced Apr 14, 2026
Two independent additions that other parts of the BA-5650 stack build on: - ``UserRepository.get_main_access_key_by_id`` (+ the matching ``UserDBSource`` method) returns ``main_access_key`` for a given user UUID with a single readonly query. Used by downstream refactors that need to resolve the keypair access_key on demand from the owner. - ``SessionConditions.by_access_key_*`` filters in ``models/session/conditions.py`` now resolve through a subquery on ``UserRow.main_access_key`` keyed by ``SessionRow.user_uuid``, instead of the ``sessions.access_key`` column. A new ``_owners_where_main_access_key`` staticmethod centralises the subquery shape. Behavior is unchanged on ``main`` where ``sessions.access_key`` still mirrors the owner's main_access_key; this change just lets us drop the redundant column later.
Slice A scope is only the new get_main_access_key_by_id helper and the SessionConditions rewrite. Accidentally dropping the target_main_access_key argument from delegate_endpoint_ownership belongs to a later slice that also updates EndpointRow and the user service; restoring the signature here so this PR builds standalone.
0b27565 to
caa8001
Compare
719ff4b to
f811d03
Compare
Contributor
There was a problem hiding this comment.
Pull request overview
Updates the session repository/data-source layer to use an owner UUID (owner_id) instead of an access key for session lookups and dependency resolution, aligning with the ongoing BA-5650 ownership refactor.
Changes:
- Refactor
SessionRepositoryandSessionDBSourcemethod signatures to acceptowner_id: UUID(e.g., session lookup/matching, rename, dependency traversal). - Update dependency-graph helpers to traverse and validate sessions using
owner_id. - Adjust related documentation/changelog entry for the signature shift.
Reviewed changes
Copilot reviewed 5 out of 5 changed files in this pull request and generated 3 comments.
Show a summary per file
| File | Description |
|---|---|
src/ai/backend/manager/repositories/session/repository.py |
Switch repository APIs from owner_access_key to owner_id and forward to DB source. |
src/ai/backend/manager/repositories/session/db_source/db_source.py |
Switch DB source session lookup/update/dependency APIs to owner_id; update internal call sites. |
src/ai/backend/manager/repositories/session/dependency_graph.py |
Replace access-key scoping with owner UUID scoping for dependency traversal. |
src/ai/backend/manager/repositories/session/creators.py |
Update documentation wording related to ownership scope (needs correction per comment). |
changes/11046.misc.md |
Add changelog note describing the signature collapse to owner_id. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
jopemachine
commented
Apr 14, 2026
caa8001 to
7bac45f
Compare
f811d03 to
c2d0d94
Compare
Drop the redundant UserRow.main_access_key IS NOT NULL guard inside SessionConditions._owners_where_main_access_key. The subquery selects users.uuid (non-null PK), so NOT IN is well-defined; NULL main_access_key fails the caller-supplied condition on its own. The behaviour is unchanged — this just makes the helper contract clearer. Also drop the stale misc news fragment (skip:changelog applies to this purely additive refactor).
Rename UserPermission.user_uuid to owner_id and add main_access_key: str | None. KernelRow.to_kernel_info populates the new field from the eagerly-loaded user_row relationship; search_kernels now selectinload's user_row so the scalar is available without N+1. KernelRow.delegate_ownership drops its access_key argument (resolved from the owner). recalc_concurrency_used switches from KernelRow.access_key to a KernelRow.user_uuid subquery keyed by the owner's main_access_key as a shim for the eventual access_key column drop. Test fixtures that construct UserPermission are updated.
Earlier the branch pulled test fixtures that also referenced later slices' renames (SessionMetadata.owner_id, ScheduledSessionData .main_access_key, SessionDataForPull/Start.main_access_key, SessionData.owner_id). Revert those fixtures to main state and re-apply only the UserPermission field renames (user_uuid -> owner_id, access_key -> main_access_key) which are in this slice's scope. Also update the two remaining live readers of UserPermission: - sokovan/scheduler/fair_share/aggregator.py - api/adapters/session.py (kernel_info_to_node) And fix the SessionRow.delegate_ownership caller of KernelRow.delegate_ownership to match the new single-argument signature.
2caf306 to
c1a5f32
Compare
jopemachine
commented
Apr 14, 2026
jopemachine
commented
Apr 14, 2026
After renaming SessionData/SessionMetadata fields and the SessionRow.get_session / match_sessions helpers, in-repo callers that used the old keyword argument (``access_key=``) or positional signature need the matching updates so slice C builds standalone: - manager/registry.py, repositories/stream/db_source/db_source.py, repositories/session/dependency_graph.py: switch to the new ``owner_id`` keyword argument on SessionRow.get_session / match_sessions. - models/endpoint/row.py: drop the ``target_access_key`` argument on SessionRow.delegate_ownership (its signature was trimmed by slice B). - tests/unit/manager/sokovan/scheduler/handlers/conftest.py: move SessionMetadata fixture fields from ``user_uuid``/``access_key`` to ``owner_id``; downstream access via ``metadata.owner_id`` and the rows now read ``main_access_key`` via the user record.
e2ff93b to
765d164
Compare
Slice C renames SessionData fields but left several downstream call sites broken in isolation (intentional per the original commit message). Fixing those so each PR in the stack passes its own CI: - ``SessionRow.match_sessions`` / ``get_session`` accept an optional ``owner_access_key`` keyword that resolves to the owner via the ``users.main_access_key`` sub-query — lets existing callers keep their AccessKey-based filter while the column rename completes. - ``find_dependency_sessions`` now takes ``owner: UUID | AccessKey`` and dispatches to the matching filter. - session/db_source call sites updated: positional AccessKey → keyword ``owner_access_key=…``. - Scheduler handlers test conftest temporarily uses ``access_key`` / ``user_uuid`` field names until the rename lands in slice F. - ``SessionAdapter._session_data_to_node`` reads ``data.owner_id`` and drops the obsolete ``data.access_key`` lookup. - ``ModelServingRepository.get_session_by_id`` drops the obsolete positional arg. - ``UserDBSource.delegate_endpoint_ownership`` discards the ``target_main_access_key`` arg (the ORM helper no longer uses it). - ``AgentRegistry`` stops threading ``user_repository`` into the ``SessionLifecycleManager`` constructor (the parameter is added in slice G). Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Data-layer rename of SessionData / SessionMetadata user_uuid to owner_id, plus the matching SessionRow adapters (to_dataclass/from_dataclass/to_session_info/from_session_info). ComputeSessionNode.from_dataclass becomes async and resolves main_access_key from the owner via UserRepository when owner is not eagerly loaded. models/resource_usage.py sources user_id from SessionRow.user_uuid through the session relationship. SessionRow._build_session_fetch_query / _match_sessions_by_* also rename the filter parameter to owner_id; repository and db_source callers are updated in the next slice so intermediate builds may fail in isolation.
``SessionRepository`` and the underlying ``SessionDBSource`` now take ``owner_id: UUID`` on every method that previously accepted ``owner_access_key: AccessKey``. Affects: - ``get_session_validated`` - ``match_sessions`` - ``update_session_name`` - ``find_dependency_sessions`` / ``_find_dependent_sessions`` - ``get_target_session_ids`` - ``get_session_with_group`` The matching ``dependency_graph`` helpers and ``creators`` are updated in lockstep. Service-layer callers still pass ``owner_access_key`` temporarily; they will be migrated in the next slice.
- repositories/session/creators.py: docstring pointed to a non-existent ``owner_id`` attribute on SessionRow; clarify that the scope_id comes from ``user_uuid`` (the owner's UUID). - repositories/session/db_source/db_source.py: drop the redundant ``session_row.user_uuid is not None`` guard (the column is non-nullable); update the ``_find_dependent_sessions`` docstring to reference ``owner_id`` instead of the removed ``access_key`` param.
- Drop the leftover ``changes/BA-5650-C.misc.md`` file that was superseded when slice C's fragment was renamed to ``changes/11045.enhance.md``. - Remove the explanatory comment about ``user_uuid`` being non-nullable from the ``update_session_name`` path — it restates what the surrounding code already makes clear.
After rebasing on slice C, three residual references in ``session/db_source/db_source.py`` still used the old ``access_key`` variable name. This commit finishes the rename in: - ``get_session_validated`` / ``match_sessions`` / ``update_session_name`` (now forward ``owner_id=...`` instead of the leftover ``owner_access_key`` keyword). - ``_find_dependent_sessions`` and ``get_target_session_ids`` (now pass the local ``owner_id`` UUID instead of the non-existent ``access_key``). Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
4ecaf01 to
82ea2d0
Compare
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Member
Author
|
Squashed into PR #11051 — stacked PRs cannot pass CI independently with cross-cutting type changes. |
Member
Author
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Resolves #10914 (BA-5650)
Summary
main_access_keyresolver helpers for user/keypair lookupmain_access_keythroughUserPermissionreplacing rawAccessKeySessionData.user_uuid→owner_id, removeaccess_keyfieldowner_idconsistentlyTest plan
pants checkpasses on this sliceStack
🤖 Generated with Claude Code