Merge pull request #199 from keboola/AI-1172-client-typed-response

Ai 1172 Typed response for global search client method
Make typed response in the global search client method
Add integ tests for global search.

Author: mariankrotil

integtests/test_client.py

@@ -0,0 +1,40 @@
+import logging
+
+import pytest
+
+from integtests.conftest import ProjectDef, TableDef
+from keboola_mcp_server.client import AsyncStorageClient, GlobalSearchResponse, KeboolaClient
+
+LOG = logging.getLogger(__name__)
+
+
+class TestAsyncStorageClient:
+
+    @pytest.fixture
+    def storage_client(self, keboola_client: KeboolaClient, keboola_project: ProjectDef) -> AsyncStorageClient:
+        return keboola_client.storage_client
+
+    @pytest.mark.asyncio
+    async def test_global_search(self, storage_client: AsyncStorageClient):
+        not_existing_id = 'not-existing-id'
+        ret = await storage_client.global_search(query=not_existing_id)
+        assert isinstance(ret, GlobalSearchResponse)
+        assert ret.all == 0
+        assert ret.items == []
+        assert ret.by_type == {'total': 0}
+        assert ret.by_project == {}
+
+    @pytest.mark.asyncio
+    async def test_global_search_with_results(self, storage_client: AsyncStorageClient, tables: list[TableDef]):
+        search_for_name = 'test'
+        is_global_search_enabled = await storage_client.is_enabled('global-search')
+        if not is_global_search_enabled:
+            LOG.warning('Global search is not enabled in the project. Skipping test. Please enable it in the project.')
+            pytest.skip('Global search is not enabled in the project. Skipping test.')
+
+        ret = await storage_client.global_search(query=search_for_name, types=['table'])
+        assert isinstance(ret, GlobalSearchResponse)
+        assert ret.all == len(tables)
+        assert len(ret.items) == len(tables)
+        assert all(isinstance(item, GlobalSearchResponse.GlobalSearchResponseItem) for item in ret.items)
+        assert all(item.type == 'table' for item in ret.items)

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "keboola-mcp-server"
-version = "1.9.0"
+version = "1.9.1"
 description = "MCP server for interacting with Keboola Connection"
 readme = "README.md"
 requires-python = ">=3.10"

src/keboola_mcp_server/client.py

@@ -3,10 +3,11 @@
 import importlib.metadata
 import logging
 import os
+from datetime import datetime
 from typing import Any, Iterable, Literal, Mapping, Optional, Union, cast
 
 import httpx
-from pydantic import BaseModel, Field
+from pydantic import BaseModel, Field, field_validator
 
 LOG = logging.getLogger(__name__)
 
@@ -355,6 +356,48 @@ async def delete(
         return await self.raw_client.delete(endpoint=endpoint)
 
 
+class GlobalSearchResponse(BaseModel):
+    """The SAPI global search response."""
+
+    class GlobalSearchResponseItem(BaseModel):
+        id: str = Field(description='The id of the item.')
+        name: str = Field(description='The name of the item.')
+        type: GlobalSearchTypes = Field(description='The type of the item.')
+        full_path: dict[str, Any] = Field(
+            description=(
+                'The full path of the item containing project, branch and other information depending on the '
+                'type of the item.'
+            ),
+            alias='fullPath',
+        )
+        component_id: Optional[str] = Field(
+            default=None, description='The id of the component the item belongs to.', alias='componentId'
+        )
+        organization_id: int = Field(
+            description='The id of the organization the item belongs to.', alias='organizationId'
+        )
+        project_id: int = Field(description='The id of the project the item belongs to.', alias='projectId')
+        project_name: str = Field(description='The name of the project the item belongs to.', alias='projectName')
+        created: datetime = Field(description='The date and time the item was created in ISO format.')
+
+    all: int = Field(description='Total number of found results.')
+    items: list[GlobalSearchResponseItem] = Field(
+        description='List of search results containing the items of the GlobalSearchType.'
+    )
+    by_type: dict[str, int] = Field(
+        description='Mapping of found types to the number of corresponding results.', alias='byType'
+    )
+    by_project: dict[str, str] = Field(description='Mapping of project id to project name.', alias='byProject')
+
+    @field_validator('by_type', 'by_project', mode='before')
+    @classmethod
+    def validate_dict_fields(cls, current_value: Any) -> Any:
+        # If the value is empty-list/None, return an empty dictionary, otherwise return the value
+        if not current_value:
+            return dict()
+        return current_value
+
+
 class AsyncStorageClient(KeboolaServiceClient):
 
     def __init__(self, raw_client: RawKeboolaClient, branch_id: str = 'default') -> None:
@@ -838,7 +881,7 @@ async def global_search(
         limit: int = 100,
         offset: int = 0,
         types: list[GlobalSearchTypes] | None = None,
-    ) -> JsonDict:
+    ) -> GlobalSearchResponse:
         """
         Searches for items in the storage. It allows you to search for entities by name across all projects within an
         organization, even those you do not have direct access to. The search is conducted only through entity names to
@@ -849,7 +892,7 @@ async def global_search(
         :param offset: The offset to start from, pagination parameter.
         :param types: The types of items to search for.
         """
-        params : dict[str, Any] = {
+        params: dict[str, Any] = {
             'query': query,
             'projectIds[]': [await self.project_id()],
             'branchTypes[]': 'production',
@@ -858,7 +901,8 @@ async def global_search(
             'offset': offset,
         }
         params = {k: v for k, v in params.items() if v}
-        return cast(JsonDict, await self.get(endpoint='global-search', params=params))
+        raw_resp = await self.get(endpoint='global-search', params=params)
+        return GlobalSearchResponse.model_validate(raw_resp)
 
     async def table_detail(self, table_id: str) -> JsonDict:
         """