From 0f9604d6b3f655412a7742ec5660f9dc784e25f0 Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Fri, 31 Jul 2026 18:49:11 -0400 Subject: [PATCH 1/2] perf(registry): deduplicate manifests in batch queries --- ...est_registry_actions_custom_entitlement.py | 20 +++ tracecat/registry/actions/service.py | 119 ++++++++++++++---- 2 files changed, 117 insertions(+), 22 deletions(-) diff --git a/tests/unit/test_registry_actions_custom_entitlement.py b/tests/unit/test_registry_actions_custom_entitlement.py index e8bf06de25..6c871024ad 100644 --- a/tests/unit/test_registry_actions_custom_entitlement.py +++ b/tests/unit/test_registry_actions_custom_entitlement.py @@ -234,6 +234,26 @@ async def test_get_actions_from_index_filters_custom_and_keeps_platform_fallback assert results[shared_action].origin == DEFAULT_REGISTRY_ORIGIN +@pytest.mark.anyio +async def test_get_actions_from_index_reuses_manifest_for_same_version( + svc_role: Role, + session: AsyncSession, +) -> None: + action_names = ["acme.batch.first", "acme.batch.second"] + await _seed_platform_registry( + session, + origin=DEFAULT_REGISTRY_ORIGIN, + version="platform-shared-manifest", + action_names=action_names, + ) + + service = RegistryActionsService(session, role=svc_role) + results = await service.get_actions_from_index(action_names) + + assert set(results) == set(action_names) + assert results[action_names[0]].manifest is results[action_names[1]].manifest + + @pytest.mark.anyio async def test_list_actions_from_index_by_repository_returns_empty_for_custom_repo_without_entitlement( svc_role: Role, diff --git a/tracecat/registry/actions/service.py b/tracecat/registry/actions/service.py index cdcd80d97e..0343426135 100644 --- a/tracecat/registry/actions/service.py +++ b/tracecat/registry/actions/service.py @@ -103,7 +103,7 @@ class _IndexSelectRow(NamedTuple): class _ActionIndexRow(NamedTuple): - """Action index row with manifest: get_action_from_index, get_actions_from_index.""" + """Action index row with manifest: get_action_from_index.""" id: uuid.UUID namespace: str @@ -122,6 +122,34 @@ class _ActionIndexRow(NamedTuple): source: str +class _ActionMetadataRow(NamedTuple): + """Action metadata row for batched index lookups.""" + + id: uuid.UUID + namespace: str + name: str + action_type: str + description: str + default_title: str | None + display_group: str | None + options: dict[str, object] + doc_url: str | None + author: str | None + deprecated: str | None + registry_version_id: uuid.UUID + origin: str + repo_id: uuid.UUID + source: str + + +class _VersionManifestRow(NamedTuple): + """One manifest payload for a selected registry version.""" + + registry_version_id: uuid.UUID + manifest: dict[str, object] + source: str + + class _RepoIndexRow(NamedTuple): """Repository index row: list_actions_from_index_by_repository (no source).""" @@ -711,7 +739,7 @@ async def get_actions_from_index( RegistryIndex.doc_url, RegistryIndex.author, RegistryIndex.deprecated, - RegistryVersion.manifest, + RegistryIndex.registry_version_id, RegistryRepository.origin, RegistryRepository.id.label("repo_id"), literal("org", type_=String).label("source"), @@ -745,7 +773,7 @@ async def get_actions_from_index( PlatformRegistryIndex.doc_url, PlatformRegistryIndex.author, PlatformRegistryIndex.deprecated, - PlatformRegistryVersion.manifest, + PlatformRegistryIndex.registry_version_id, PlatformRegistryRepository.origin, PlatformRegistryRepository.id.label("repo_id"), literal("platform", type_=String).label("source"), @@ -770,21 +798,43 @@ async def get_actions_from_index( text("source") # "org" < "platform" alphabetically ) result = await self.session.execute(combined) - rows = typing_cast(list[_ActionIndexRow], result.all()) + rows = typing_cast(list[_ActionMetadataRow], result.all()) allow_custom_origins = await self._allow_custom_origins_for_rows( row.origin for row in rows ) - actions: dict[str, IndexedActionResult] = {} + selected_rows: dict[str, _ActionMetadataRow] = {} for row in rows: if not allow_custom_origins and self._is_custom_origin(row.origin): continue action_name = f"{row.namespace}.{row.name}" # Skip if already found (org-scoped takes precedence) - if action_name in actions: + if action_name in selected_rows: continue + selected_rows[action_name] = row - manifest = RegistryVersionManifest.model_validate(row.manifest) + required_any: set[str] = set() + for row in selected_rows.values(): + required_any |= self._normalize_required_entitlements(row.options or {}) + if required_any: + enabled = await self._get_enabled_entitlements() + selected_rows = { + name: row + for name, row in selected_rows.items() + if self._normalize_required_entitlements(row.options or {}).issubset( + enabled + ) + } + + manifests = await self._load_action_manifests(list(selected_rows.values())) + + actions: dict[str, IndexedActionResult] = {} + for action_name, row in selected_rows.items(): + manifest = manifests.get((row.source, row.registry_version_id)) + if manifest is None: + raise RegistryError( + "Manifest missing for a selected registry action version" + ) actions[action_name] = IndexedActionResult( index_entry=IndexEntry( id=row.id, @@ -804,23 +854,48 @@ async def get_actions_from_index( repository_id=row.repo_id, ) - required_any: set[str] = set() - for result in actions.values(): - required_any |= self._normalize_required_entitlements( - result.index_entry.options - ) - if required_any: - enabled = await self._get_enabled_entitlements() - actions = { - name: result - for name, result in actions.items() - if self._normalize_required_entitlements( - result.index_entry.options - ).issubset(enabled) - } - return actions + async def _load_action_manifests( + self, + rows: Sequence[_ActionMetadataRow], + ) -> dict[tuple[str, uuid.UUID], RegistryVersionManifest]: + """Load and validate each registry version manifest once.""" + if not rows: + return {} + + org_version_ids = { + row.registry_version_id for row in rows if row.source == "org" + } + platform_version_ids = { + row.registry_version_id for row in rows if row.source == "platform" + } + + org_statement = select( + RegistryVersion.id.label("registry_version_id"), + RegistryVersion.manifest, + literal("org", type_=String).label("source"), + ).where( + RegistryVersion.organization_id == self.organization_id, + RegistryVersion.id.in_(org_version_ids), + ) + platform_statement = select( + PlatformRegistryVersion.id.label("registry_version_id"), + PlatformRegistryVersion.manifest, + literal("platform", type_=String).label("source"), + ).where(PlatformRegistryVersion.id.in_(platform_version_ids)) + + result = await self.session.execute( + union_all(org_statement, platform_statement) + ) + manifest_rows = typing_cast(list[_VersionManifestRow], result.all()) + return { + (row.source, row.registry_version_id): ( + RegistryVersionManifest.model_validate(row.manifest) + ) + for row in manifest_rows + } + async def search_actions_from_index( self, query: str, From 87b0e9833401f87229ab4f58191eea945f9d1ee8 Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Mon, 3 Aug 2026 13:01:32 -0700 Subject: [PATCH 2/2] fix(registry): retry manifest lookup after version changes --- ...est_registry_actions_custom_entitlement.py | 138 +++++++++++++++++- tracecat/registry/actions/service.py | 23 ++- 2 files changed, 155 insertions(+), 6 deletions(-) diff --git a/tests/unit/test_registry_actions_custom_entitlement.py b/tests/unit/test_registry_actions_custom_entitlement.py index 6c871024ad..92a5145a17 100644 --- a/tests/unit/test_registry_actions_custom_entitlement.py +++ b/tests/unit/test_registry_actions_custom_entitlement.py @@ -1,11 +1,19 @@ from __future__ import annotations +import uuid +from collections.abc import Sequence from unittest.mock import AsyncMock, patch import pytest -from sqlalchemy import select -from sqlalchemy.ext.asyncio import AsyncSession +from sqlalchemy import delete, select +from sqlalchemy.ext.asyncio import ( + AsyncSession, + async_sessionmaker, + create_async_engine, +) +from sqlalchemy.pool import NullPool +from tests.database import TEST_DB_CONFIG from tracecat.auth.types import Role from tracecat.db.models import ( PlatformRegistryIndex, @@ -15,8 +23,12 @@ RegistryRepository, RegistryVersion, ) -from tracecat.registry.actions.service import RegistryActionsService +from tracecat.registry.actions.service import ( + RegistryActionsService, + _ActionMetadataRow, +) from tracecat.registry.constants import DEFAULT_REGISTRY_ORIGIN +from tracecat.registry.versions.schemas import RegistryVersionManifest pytestmark = pytest.mark.usefixtures("db") @@ -254,6 +266,126 @@ async def test_get_actions_from_index_reuses_manifest_for_same_version( assert results[action_names[0]].manifest is results[action_names[1]].manifest +@pytest.mark.anyio +async def test_get_actions_from_index_retries_version_replaced_between_queries( + svc_role: Role, +) -> None: + action_name = "acme.batch.replaced" + origin = f"git+ssh://git@github.com/acme/registry-{uuid.uuid4()}.git" + engine = create_async_engine( + TEST_DB_CONFIG.test_url, + isolation_level="READ COMMITTED", + poolclass=NullPool, + ) + session_factory = async_sessionmaker(engine, expire_on_commit=False) + repository_id: uuid.UUID | None = None + + try: + async with session_factory() as seed_session: + repository = await _seed_org_registry( + seed_session, + role=svc_role, + origin=origin, + version="before-replacement", + action_names=[action_name], + ) + repository_id = repository.id + old_version_id = repository.current_version_id + assert old_version_id is not None + + async with session_factory() as read_session: + service = RegistryActionsService(read_session, role=svc_role) + original_load = service._load_action_manifests + load_calls = 0 + + async def replace_version_before_manifest_load( + rows: Sequence[_ActionMetadataRow], + ) -> dict[tuple[str, uuid.UUID], RegistryVersionManifest]: + nonlocal load_calls + load_calls += 1 + if load_calls == 1: + async with session_factory() as write_session: + repository = await write_session.scalar( + select(RegistryRepository).where( + RegistryRepository.id == repository_id + ) + ) + assert repository is not None + repository.current_version_id = None + await write_session.flush() + await write_session.execute( + delete(RegistryVersion).where( + RegistryVersion.id == old_version_id + ) + ) + + replacement_manifest = _make_manifest( + [action_name], origin=origin + ) + replacement_manifest["actions"][action_name]["description"] = ( + "Replacement action" + ) + replacement = RegistryVersion( + organization_id=svc_role.organization_id, + repository_id=repository.id, + version="after-replacement", + manifest=replacement_manifest, + tarball_uri="s3://org/after-replacement.tar.gz", + ) + write_session.add(replacement) + await write_session.flush() + repository.current_version_id = replacement.id + write_session.add( + RegistryIndex( + organization_id=svc_role.organization_id, + registry_version_id=replacement.id, + namespace="acme.batch", + name="replaced", + action_type="udf", + description="Replacement action", + options={"include_in_schema": True}, + ) + ) + await write_session.commit() + + return await original_load(rows) + + with ( + patch.object( + service, + "has_entitlement", + new=AsyncMock(return_value=True), + ), + patch.object( + service, + "_load_action_manifests", + new=replace_version_before_manifest_load, + ), + ): + results = await service.get_actions_from_index([action_name]) + + assert load_calls == 2 + assert results[action_name].index_entry.description == "Replacement action" + assert ( + results[action_name].manifest.actions[action_name].description + == "Replacement action" + ) + finally: + if repository_id is not None: + async with session_factory() as cleanup_session: + repository = await cleanup_session.scalar( + select(RegistryRepository).where( + RegistryRepository.id == repository_id + ) + ) + if repository is not None: + repository.current_version_id = None + await cleanup_session.flush() + await cleanup_session.delete(repository) + await cleanup_session.commit() + await engine.dispose() + + @pytest.mark.anyio async def test_list_actions_from_index_by_repository_returns_empty_for_custom_repo_without_entitlement( svc_role: Role, diff --git a/tracecat/registry/actions/service.py b/tracecat/registry/actions/service.py index 0343426135..d39c8a2b05 100644 --- a/tracecat/registry/actions/service.py +++ b/tracecat/registry/actions/service.py @@ -698,6 +698,25 @@ async def get_actions_from_index( Dict mapping action_name -> IndexedActionResult. Actions not found are omitted from the result. """ + actions = await self._get_actions_from_index_once(action_names) + if actions is not None: + return actions + + actions = await self._get_actions_from_index_once(action_names) + if actions is None: + raise RegistryError( + "Manifest missing for a selected registry action version" + ) + return actions + + async def _get_actions_from_index_once( + self, + action_names: list[str], + ) -> dict[str, IndexedActionResult] | None: + """Run one metadata and manifest lookup attempt. + + Returns ``None`` when a selected version disappears between statements. + """ if not action_names: return {} @@ -832,9 +851,7 @@ async def get_actions_from_index( for action_name, row in selected_rows.items(): manifest = manifests.get((row.source, row.registry_version_id)) if manifest is None: - raise RegistryError( - "Manifest missing for a selected registry action version" - ) + return None actions[action_name] = IndexedActionResult( index_entry=IndexEntry( id=row.id,