Extract metadata tree controller

services/api-server/src/api_server/main.py

@@ -94,6 +94,12 @@ from api_server.html5_setup_controller import (
     html5_setup_source as _html5_setup_source,
     html5_setup_summary as _html5_setup_summary,
 )
+from api_server.metadata_tree_controller import (
+    metadata_tree as _metadata_tree,
+    metadata_tree_children as _metadata_tree_children,
+    metadata_tree_path as _metadata_tree_path,
+    metadata_tree_search as _metadata_tree_search,
+)
 from impact_engine import object_impact, routine_impact
 from incremental_indexer import rebuild_changed_file
 from integration_topology import IntegrationKind, build_integration_topology
@@ -3204,16 +3210,14 @@ async def metadata_catalog() -> MetadataCatalogResponse:
 
 @app.get("/projects/{project_id}/metadata/tree", response_model=ProjectMetadataTreeResponse)
 async def project_metadata_tree(project_id: str, object_limit_per_branch: int = 200) -> ProjectMetadataTreeResponse:
-    snapshot = _project_snapshot_or_404(project_id)
+    return _metadata_tree(
-    normalized = _load_normalized_project(project_id)
-    root = (
-        _project_metadata_tree_response_from_normalized(normalized, object_limit_per_branch=max(0, object_limit_per_branch))
-        if normalized is not None
-        else _project_metadata_tree_response(snapshot, object_limit_per_branch=max(0, object_limit_per_branch))
-    )
-    return ProjectMetadataTreeResponse(
         project_id=project_id,
-        root=root,
+        object_limit_per_branch=object_limit_per_branch,
+        project_snapshot=_project_snapshot_or_404,
+        normalized_project=_load_normalized_project,
+        normalized_tree=_project_metadata_tree_response_from_normalized,
+        snapshot_tree=_project_metadata_tree_response,
+        response_model=ProjectMetadataTreeResponse,
     )
 
 
@@ -3224,78 +3228,43 @@ async def project_metadata_tree_children(
     offset: int = 0,
     limit: int = 50,
 ) -> MetadataTreeChildrenResponse:
-    snapshot = _project_snapshot_or_404(project_id)
+    return _metadata_tree_children(
-    normalized = _load_normalized_project(project_id)
-    normalized_offset = max(0, offset)
-    normalized_limit = min(max(1, limit), 250)
-    normalized_children = (
-        _normalized_metadata_tree_children_for_node(
-            normalized,
-            node_id=node_id,
-            offset=normalized_offset,
-            limit=normalized_limit,
-        )
-        if normalized is not None
-        else None
-    )
-    if normalized_children is None:
-        children, total = _metadata_tree_children_for_node(
-            snapshot,
-            node_id=node_id,
-            offset=normalized_offset,
-            limit=normalized_limit,
-        )
-    else:
-        children, total = normalized_children
-    return MetadataTreeChildrenResponse(
         project_id=project_id,
-        parent_id=node_id,
+        node_id=node_id,
-        offset=normalized_offset,
+        offset=offset,
-        limit=normalized_limit,
+        limit=limit,
-        total=total,
+        project_snapshot=_project_snapshot_or_404,
-        has_more=normalized_offset + len(children) < total,
+        normalized_project=_load_normalized_project,
-        children=children,
+        normalized_children_for_node=_normalized_metadata_tree_children_for_node,
+        snapshot_children_for_node=_metadata_tree_children_for_node,
+        response_model=MetadataTreeChildrenResponse,
     )
 
 
 @app.get("/projects/{project_id}/metadata/tree/search", response_model=MetadataTreeSearchResponse)
 async def project_metadata_tree_search(project_id: str, q: str, limit: int = 80) -> MetadataTreeSearchResponse:
-    snapshot = _project_snapshot_or_404(project_id)
+    return _metadata_tree_search(
-    normalized_query = q.strip().casefold()
-    normalized_limit = min(max(1, limit), 250)
-    if len(normalized_query) < 2:
-        return MetadataTreeSearchResponse(project_id=project_id, q=q, total=0, results=[])
-    matches = [
-        node
-        for node in snapshot.nodes
-        if _is_metadata_tree_search_node(node)
-        and (
-            normalized_query in node.name.casefold()
-            or normalized_query in node.qualified_name.casefold()
-        )
-    ]
-    matches.sort(key=lambda item: _metadata_search_rank(item, normalized_query))
-    page = matches[:normalized_limit]
-    child_count_index = _metadata_child_count_index(snapshot, [node.lineage_id for node in page])
-    return MetadataTreeSearchResponse(
         project_id=project_id,
         q=q,
-        total=len(matches),
+        limit=limit,
-        results=[_metadata_node_for_search_result(snapshot, node, child_count_index) for node in page],
+        project_snapshot=_project_snapshot_or_404,
+        is_search_node=_is_metadata_tree_search_node,
+        search_rank=_metadata_search_rank,
+        child_count_index=_metadata_child_count_index,
+        node_for_search_result=_metadata_node_for_search_result,
+        response_model=MetadataTreeSearchResponse,
     )
 
 
 @app.get("/projects/{project_id}/metadata/tree/path", response_model=MetadataTreePathResponse)
 async def project_metadata_tree_path(project_id: str, node_id: str) -> MetadataTreePathResponse:
-    snapshot = _project_snapshot_or_404(project_id)
+    return _metadata_tree_path(
-    path = _metadata_tree_path_for_node(snapshot, node_id)
-    if not path:
-        raise HTTPException(status_code=404, detail=f"Metadata tree path not found: {node_id}")
-    return MetadataTreePathResponse(
         project_id=project_id,
         node_id=node_id,
-        path=path,
+        project_snapshot=_project_snapshot_or_404,
-        steps=_metadata_tree_path_steps(snapshot, path),
+        tree_path_for_node=_metadata_tree_path_for_node,
+        tree_path_steps=_metadata_tree_path_steps,
+        response_model=MetadataTreePathResponse,
     )
 
 

services/api-server/src/api_server/metadata_tree_controller.py

@@ -0,0 +1,130 @@
+from __future__ import annotations
+
+from collections.abc import Callable
+from typing import Any
+
+from fastapi import HTTPException
+
+
+def metadata_tree(
+    *,
+    project_id: str,
+    object_limit_per_branch: int,
+    project_snapshot: Callable[[str], object],
+    normalized_project: Callable[[str], object | None],
+    normalized_tree: Callable[..., object],
+    snapshot_tree: Callable[..., object],
+    response_model: Callable[..., object],
+) -> object:
+    snapshot = project_snapshot(project_id)
+    normalized = normalized_project(project_id)
+    root = (
+        normalized_tree(normalized, object_limit_per_branch=max(0, object_limit_per_branch))
+        if normalized is not None
+        else snapshot_tree(snapshot, object_limit_per_branch=max(0, object_limit_per_branch))
+    )
+    return response_model(project_id=project_id, root=root)
+
+
+def metadata_tree_children(
+    *,
+    project_id: str,
+    node_id: str,
+    offset: int,
+    limit: int,
+    project_snapshot: Callable[[str], object],
+    normalized_project: Callable[[str], object | None],
+    normalized_children_for_node: Callable[..., tuple[list[object], int] | None],
+    snapshot_children_for_node: Callable[..., tuple[list[object], int]],
+    response_model: Callable[..., object],
+) -> object:
+    snapshot = project_snapshot(project_id)
+    normalized = normalized_project(project_id)
+    normalized_offset = max(0, offset)
+    normalized_limit = min(max(1, limit), 250)
+    normalized_children = (
+        normalized_children_for_node(
+            normalized,
+            node_id=node_id,
+            offset=normalized_offset,
+            limit=normalized_limit,
+        )
+        if normalized is not None
+        else None
+    )
+    if normalized_children is None:
+        children, total = snapshot_children_for_node(
+            snapshot,
+            node_id=node_id,
+            offset=normalized_offset,
+            limit=normalized_limit,
+        )
+    else:
+        children, total = normalized_children
+    return response_model(
+        project_id=project_id,
+        parent_id=node_id,
+        offset=normalized_offset,
+        limit=normalized_limit,
+        total=total,
+        has_more=normalized_offset + len(children) < total,
+        children=children,
+    )
+
+
+def metadata_tree_search(
+    *,
+    project_id: str,
+    q: str,
+    limit: int,
+    project_snapshot: Callable[[str], object],
+    is_search_node: Callable[[Any], bool],
+    search_rank: Callable[[Any, str], object],
+    child_count_index: Callable[[object, list[str]], object],
+    node_for_search_result: Callable[[object, Any, object], object],
+    response_model: Callable[..., object],
+) -> object:
+    snapshot = project_snapshot(project_id)
+    normalized_query = q.strip().casefold()
+    normalized_limit = min(max(1, limit), 250)
+    if len(normalized_query) < 2:
+        return response_model(project_id=project_id, q=q, total=0, results=[])
+    matches = [
+        node
+        for node in snapshot.nodes
+        if is_search_node(node)
+        and (
+            normalized_query in node.name.casefold()
+            or normalized_query in node.qualified_name.casefold()
+        )
+    ]
+    matches.sort(key=lambda item: search_rank(item, normalized_query))
+    page = matches[:normalized_limit]
+    counts = child_count_index(snapshot, [node.lineage_id for node in page])
+    return response_model(
+        project_id=project_id,
+        q=q,
+        total=len(matches),
+        results=[node_for_search_result(snapshot, node, counts) for node in page],
+    )
+
+
+def metadata_tree_path(
+    *,
+    project_id: str,
+    node_id: str,
+    project_snapshot: Callable[[str], object],
+    tree_path_for_node: Callable[[object, str], list[str]],
+    tree_path_steps: Callable[[object, list[str]], list[object]],
+    response_model: Callable[..., object],
+) -> object:
+    snapshot = project_snapshot(project_id)
+    path = tree_path_for_node(snapshot, node_id)
+    if not path:
+        raise HTTPException(status_code=404, detail=f"Metadata tree path not found: {node_id}")
+    return response_model(
+        project_id=project_id,
+        node_id=node_id,
+        path=path,
+        steps=tree_path_steps(snapshot, path),
+    )