From 6a03cbda9143ac888f0a876a5fd3f8557106eecd Mon Sep 17 00:00:00 2001 From: mariankrotil Date: Tue, 12 Aug 2025 17:36:09 +0200 Subject: [PATCH 01/35] AI-1343 feat: add data app component id --- src/keboola_mcp_server/client.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/keboola_mcp_server/client.py b/src/keboola_mcp_server/client.py index d5dcfa3a..e724d56b 100644 --- a/src/keboola_mcp_server/client.py +++ b/src/keboola_mcp_server/client.py @@ -40,6 +40,7 @@ ORCHESTRATOR_COMPONENT_ID = 'keboola.orchestrator' CONDITIONAL_FLOW_COMPONENT_ID = 'keboola.flow' +DATA_APP_COMPONENT_ID = 'keboola.data-apps' FlowType = Literal['keboola.flow', 'keboola.orchestrator'] FLOW_TYPES: Sequence[FlowType] = (CONDITIONAL_FLOW_COMPONENT_ID, ORCHESTRATOR_COMPONENT_ID) From a7b25c739adce137aebd882152eda129a02ea538 Mon Sep 17 00:00:00 2001 From: mariankrotil Date: Thu, 14 Aug 2025 09:22:54 +0200 Subject: [PATCH 02/35] AI-1343 feat: create client for data science service, create client subpackage --- src/keboola_mcp_server/clients/__init__.py | 0 .../clients/data_science.py | 151 ++++++++++++++++++ 2 files changed, 151 insertions(+) create mode 100644 src/keboola_mcp_server/clients/__init__.py create mode 100644 src/keboola_mcp_server/clients/data_science.py diff --git a/src/keboola_mcp_server/clients/__init__.py b/src/keboola_mcp_server/clients/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/src/keboola_mcp_server/clients/data_science.py b/src/keboola_mcp_server/clients/data_science.py new file mode 100644 index 00000000..371cd338 --- /dev/null +++ b/src/keboola_mcp_server/clients/data_science.py @@ -0,0 +1,151 @@ +from typing import Any, Optional + +from pydantic import AliasChoices, BaseModel, Field + +from keboola_mcp_server.clients.base import KeboolaServiceClient, RawKeboolaClient + + +class DataAppResponse(BaseModel): + id: str = Field(validation_alias=AliasChoices('id', 'data_app_id'), description='The data app ID') + project_id: str = Field(validation_alias=AliasChoices('projectId', 'project_id'), description='The project ID') + component_id: str = Field( + validation_alias=AliasChoices('componentId', 'component_id'), description='The component ID' + ) + branch_id: Optional[str] = Field( + validation_alias=AliasChoices('branchId', 'branch_id'), description='The branch ID' + ) + config_id: str = Field( + validation_alias=AliasChoices('configId', 'config_id'), description='The component config ID' + ) + config_version: str = Field( + validation_alias=AliasChoices('configVersion', 'config_version'), description='The config version' + ) + type: str = Field(description='The type of the data app') + state: str = Field(description='The state of the data app') + desired_state: str = Field( + validation_alias=AliasChoices('desiredState', 'desired_state'), description='The desired state' + ) + last_request_timestamp: str = Field( + validation_alias=AliasChoices('lastRequestTimestamp', 'last_request_timestamp'), + description='The last request timestamp', + ) + last_start_timestamp: str = Field( + validation_alias=AliasChoices('lastStartTimestamp', 'last_start_timestamp'), + description='The last start timestamp', + ) + url: str = Field(validation_alias=AliasChoices('url', 'url'), description='The URL of the running data app') + auto_suspend_after_seconds: int = Field( + validation_alias=AliasChoices('autoSuspendAfterSeconds', 'auto_suspend_after_seconds'), + description='The auto suspend after seconds', + ) + size: Optional[str] = Field(validation_alias=AliasChoices('size', 'size'), description='The size of the data app') + + +class DataAppConfig(BaseModel): + class Parameters(BaseModel): + class DataApp(BaseModel): + slug: str = Field(description='The slug of the data app') + streamlit: dict[str, str] = Field(description='The streamlit config.toml file') + + size: str = Field(description='The size of the data app') + auto_suspend_after_seconds: int = Field( + validation_alias=AliasChoices('autoSuspendAfterSeconds', 'auto_suspend_after_seconds'), + serialization_alias='autoSuspendAfterSeconds', + description='The auto suspend after seconds', + ) + data_app: DataApp = Field(description='The data app sub config', serialization_alias='dataApp') + id: Optional[str] = Field(description='The id of the data app', default=None) + script: Optional[list[str]] = Field(description='The script of the data app', default=None) + packages: Optional[list[str]] = Field( + description='The python packages needed to be installed in the data app', default=None + ) + + class Authorization(BaseModel): + class AppProxy(BaseModel): + auth_providers: list[dict[str, Any]] = Field(description='The auth providers') + auth_rules: list[dict[str, Any]] = Field(description='The auth rules') + + app_proxy: AppProxy = Field(description='The app proxy') + + parameters: Parameters = Field(description='The parameters of the data app') + authorization: Authorization = Field(description='The authorization of the data app') + + +class AsyncDataScienceClient(KeboolaServiceClient): + + def __init__(self, raw_client: RawKeboolaClient) -> None: + """ + Creates an AsyncStorageClient from a RawKeboolaClient and a branch id. + + :param raw_client: The raw client to use + :param branch_id: The id of the branch + """ + super().__init__(raw_client=raw_client) + + @property + def base_api_url(self) -> str: + return self.raw_client.base_api_url.split('/apps')[0] + + @classmethod + def create( + cls, + root_url: str, + token: str, + headers: dict[str, Any] | None = None, + ) -> 'AsyncDataScienceClient': + """ + Creates an AsyncStorageClient from a Keboola Storage API token. + + :param root_url: The root URL of the service API + :param token: The Keboola Storage API token + :param version: The version of the API to use (default: 'v2') + :param branch_id: The id of the branch + :param headers: Additional headers for the requests + :return: A new instance of AsyncStorageClient + """ + return cls( + raw_client=RawKeboolaClient( + base_api_url=f'{root_url}/apps', + api_token=token, + headers=headers, + ) + ) + + async def get_data_app(self, data_app_id: str) -> DataAppResponse: + """ + Get a data app by its ID. + + :param data_app_id: The ID of the data app + :return: The data app + """ + response = await self.raw_client.get(endpoint=data_app_id) + return DataAppResponse.model_validate(response) + + async def create_data_app( + self, + name: str, + description: str, + parameters: dict[str, Any], + authorization: dict[str, Any], + ) -> dict[str, Any]: + """ + Create a data app. + + :param data_app_id: The ID of the data app + :return: The data app + """ + _params = DataAppConfig.Parameters.model_validate(parameters).model_dump(exclude_none=True, by_alias=True) + _authorization = DataAppConfig.Authorization.model_validate(authorization).model_dump( + exclude_none=True, by_alias=True + ) + params = { + 'name': name, + 'type': 'streamlit', + 'description': description, + 'config': { + 'parameters': _params, + 'authorization': _authorization, + }, + } + response = await self.raw_client.post('', params=params) + return response From 1c0352617a53f1d8d3a79ae867f5a77940712e47 Mon Sep 17 00:00:00 2001 From: mariankrotil Date: Thu, 14 Aug 2025 09:23:54 +0200 Subject: [PATCH 03/35] AI-1343 refactor: move client file to client subpackage and create base client preventing from circ imports --- src/keboola_mcp_server/clients/base.py | 257 +++++++++++++++++ .../{ => clients}/client.py | 271 ++---------------- 2 files changed, 273 insertions(+), 255 deletions(-) create mode 100644 src/keboola_mcp_server/clients/base.py rename src/keboola_mcp_server/{ => clients}/client.py (85%) diff --git a/src/keboola_mcp_server/clients/base.py b/src/keboola_mcp_server/clients/base.py new file mode 100644 index 00000000..70d0007a --- /dev/null +++ b/src/keboola_mcp_server/clients/base.py @@ -0,0 +1,257 @@ +from typing import Any, Optional, Union, cast + +import httpx + +JsonPrimitive = Union[int, float, str, bool, None] +JsonDict = dict[str, Union[JsonPrimitive, 'JsonStruct']] +JsonList = list[Union[JsonPrimitive, 'JsonStruct']] +JsonStruct = Union[JsonDict, JsonList] + + +class RawKeboolaClient: + """ + Raw async client for Keboola services. + + Implements the basic HTTP methods (GET, POST, PUT, DELETE) + and can be used to implement high-level functions in clients for individual services. + """ + + def __init__( + self, + base_api_url: str, + api_token: str, + headers: dict[str, Any] | None = None, + timeout: httpx.Timeout | None = None, + ) -> None: + self.base_api_url = base_api_url + self.headers = { + 'Content-Type': 'application/json', + 'Accept-encoding': 'gzip', + } + if api_token.startswith('Bearer '): + self.headers['Authorization'] = api_token + else: + self.headers['X-StorageAPI-Token'] = api_token + self.timeout = timeout or httpx.Timeout(connect=5.0, read=60.0, write=10.0, pool=5.0) + if headers: + self.headers.update(headers) + + @staticmethod + def _raise_for_status(response: httpx.Response) -> None: + """ + Checks the HTTP response status code and raises an exception with a detailed message. The message will + include "error" and "exceptionId" fields if they are present in the response. + """ + try: + response.raise_for_status() + except httpx.HTTPStatusError as e: + message_parts = [str(e)] + + try: + error_data = response.json() + if error_msg := error_data.get('error'): + message_parts.append(f'API error: {error_msg}') + if exception_id := error_data.get('exceptionId'): + message_parts.append(f'Exception ID: {exception_id}') + message_parts.append('When contacting Keboola support please provide the exception ID.') + + except ValueError: + try: + if response.text: + message_parts.append(f'API error: {response.text}') + except Exception: + pass # should never get here + + raise httpx.HTTPStatusError('\n'.join(message_parts), request=response.request, response=response) from e + + async def get( + self, + endpoint: str, + params: dict[str, Any] | None = None, + headers: dict[str, Any] | None = None, + ) -> JsonStruct: + """ + Makes a GET request to the service API. + + :param endpoint: API endpoint to call + :param params: Query parameters for the request + :param headers: Additional headers for the request + :return: API response as dictionary + """ + headers = self.headers | (headers or {}) + async with httpx.AsyncClient(timeout=self.timeout) as client: + response = await client.get( + f'{self.base_api_url}/{endpoint}', + params=params, + headers=headers, + ) + self._raise_for_status(response) + return cast(JsonStruct, response.json()) + + async def post( + self, + endpoint: str, + data: dict[str, Any] | None = None, + params: dict[str, Any] | None = None, + headers: dict[str, Any] | None = None, + ) -> JsonStruct: + """ + Makes a POST request to the service API. + + :param endpoint: API endpoint to call + :param data: Request payload + :param params: Query parameters for the request + :param headers: Additional headers for the request + :return: API response as dictionary + """ + headers = self.headers | (headers or {}) + async with httpx.AsyncClient(timeout=self.timeout) as client: + response = await client.post( + f'{self.base_api_url}/{endpoint}', + params=params, + headers=headers, + json=data or {}, + ) + self._raise_for_status(response) + return cast(JsonStruct, response.json()) + + async def put( + self, + endpoint: str, + data: dict[str, Any] | None = None, + params: dict[str, Any] | None = None, + headers: dict[str, Any] | None = None, + ) -> JsonStruct: + """ + Makes a PUT request to the service API. + + :param endpoint: API endpoint to call + :param data: Request payload + :param params: Query parameters for the request + :param headers: Additional headers for the request + :return: API response as dictionary + """ + headers = self.headers | (headers or {}) + async with httpx.AsyncClient(timeout=self.timeout) as client: + response = await client.put( + f'{self.base_api_url}/{endpoint}', + params=params, + headers=headers, + json=data or {}, + ) + self._raise_for_status(response) + return cast(JsonStruct, response.json()) + + async def delete( + self, + endpoint: str, + headers: dict[str, Any] | None = None, + ) -> JsonStruct | None: + """ + Makes a DELETE request to the service API. + + :param endpoint: API endpoint to call + :param headers: Additional headers for the request + :return: API response as dictionary + """ + headers = self.headers | (headers or {}) + async with httpx.AsyncClient(timeout=self.timeout) as client: + response = await client.delete( + f'{self.base_api_url}/{endpoint}', + headers=headers, + ) + self._raise_for_status(response) + + if response.content: + return cast(JsonStruct, response.json()) + + return None + + +class KeboolaServiceClient: + """ + Base class for Keboola service clients. + + Implements the basic HTTP methods (GET, POST, PUT, DELETE) + and is used as a base class for clients for individual services. + """ + + def __init__(self, raw_client: RawKeboolaClient) -> None: + """ + Creates a client instance. + + The inherited classes should implement the `create` method + rather than overriding this constructor. + + :param raw_client: The raw client to use + """ + self.raw_client = raw_client + + @classmethod + def create(cls, root_url: str, token: str) -> 'KeboolaServiceClient': + """ + Creates a KeboolaServiceClient from a Keboola Storage API token. + + :param root_url: The root URL of the service API + :param token: The Keboola Storage API token + :return: A new instance of KeboolaServiceClient + """ + return cls(raw_client=RawKeboolaClient(base_api_url=root_url, api_token=token)) + + async def get( + self, + endpoint: str, + params: Optional[dict[str, Any]] = None, + ) -> JsonStruct: + """ + Makes a GET request to the service API. + + :param endpoint: API endpoint to call + :param params: Query parameters for the request + :return: API response as dictionary + """ + return await self.raw_client.get(endpoint=endpoint, params=params) + + async def post( + self, + endpoint: str, + data: Optional[dict[str, Any]] = None, + params: Optional[dict[str, Any]] = None, + ) -> JsonStruct: + """ + Makes a POST request to the service API. + + :param endpoint: API endpoint to call + :param data: Request payload + :param params: Query parameters for the request + :return: API response as dictionary + """ + return await self.raw_client.post(endpoint=endpoint, data=data, params=params) + + async def put( + self, + endpoint: str, + data: Optional[dict[str, Any]] = None, + params: Optional[dict[str, Any]] = None, + ) -> JsonStruct: + """ + Makes a PUT request to the service API. + + :param endpoint: API endpoint to call + :param data: Request payload + :param params: Query parameters for the request + :return: API response as dictionary + """ + return await self.raw_client.put(endpoint=endpoint, data=data, params=params) + + async def delete( + self, + endpoint: str, + ) -> JsonStruct | None: + """ + Makes a DELETE request to the service API. + + :param endpoint: API endpoint to call + :return: API response as dictionary + """ + return await self.raw_client.delete(endpoint=endpoint) diff --git a/src/keboola_mcp_server/client.py b/src/keboola_mcp_server/clients/client.py similarity index 85% rename from src/keboola_mcp_server/client.py rename to src/keboola_mcp_server/clients/client.py index e724d56b..a70290b5 100644 --- a/src/keboola_mcp_server/client.py +++ b/src/keboola_mcp_server/clients/client.py @@ -5,18 +5,21 @@ import math import os from datetime import datetime -from typing import Any, Iterable, Literal, Mapping, Optional, Sequence, TypeVar, Union, cast +from typing import Any, Iterable, Literal, Mapping, Optional, Sequence, TypeVar, cast -import httpx from pydantic import AliasChoices, BaseModel, Field, field_validator +from keboola_mcp_server.clients.base import ( + JsonDict, + JsonList, + KeboolaServiceClient, + RawKeboolaClient, +) +from keboola_mcp_server.clients.data_science import AsyncDataScienceClient + LOG = logging.getLogger(__name__) T = TypeVar('T') -JsonPrimitive = Union[int, float, str, bool, None] -JsonDict = dict[str, Union[JsonPrimitive, 'JsonStruct']] -JsonList = list[Union[JsonPrimitive, 'JsonStruct']] -JsonStruct = Union[JsonDict, JsonList] ComponentResource = Literal['configuration', 'rows', 'state'] StorageEventType = Literal['info', 'success', 'warn', 'error'] @@ -77,6 +80,7 @@ class KeboolaClient: _PREFIX_STORAGE_API_URL = 'connection.' _PREFIX_QUEUE_API_URL = 'https://queue.' _PREFIX_AISERVICE_API_URL = 'https://ai.' + _PREFIX_DATA_SCIENCE_API_URL = 'https://data-science.' @classmethod def from_state(cls, state: Mapping[str, Any]) -> 'KeboolaClient': @@ -103,6 +107,9 @@ def __init__(self, storage_api_token: str, storage_api_url: str, bearer_token: s # and add the prefix for the queue API https://queue.REGION.keboola.com queue_api_url = f'{self._PREFIX_QUEUE_API_URL}{storage_api_url.split(self._PREFIX_STORAGE_API_URL)[1]}' ai_service_api_url = f'{self._PREFIX_AISERVICE_API_URL}{storage_api_url.split(self._PREFIX_STORAGE_API_URL)[1]}' + data_science_api_url = ( + f'{self._PREFIX_DATA_SCIENCE_API_URL}{storage_api_url.split(self._PREFIX_STORAGE_API_URL)[1]}' + ) # Initialize clients for individual services bearer_or_sapi_token = f'Bearer {bearer_token}' if bearer_token else storage_api_token @@ -115,6 +122,9 @@ def __init__(self, storage_api_token: str, storage_api_url: str, bearer_token: s self.ai_service_client = AIServiceClient.create( root_url=ai_service_api_url, token=self.token, headers=self._get_headers() ) + self.data_science_client = AsyncDataScienceClient.create( + root_url=data_science_api_url, token=self.token, headers=self._get_headers() + ) @classmethod def _get_user_agent(cls) -> str: @@ -137,255 +147,6 @@ def _get_headers(cls) -> dict[str, Any]: return {'User-Agent': cls._get_user_agent()} -class RawKeboolaClient: - """ - Raw async client for Keboola services. - - Implements the basic HTTP methods (GET, POST, PUT, DELETE) - and can be used to implement high-level functions in clients for individual services. - """ - - def __init__( - self, - base_api_url: str, - api_token: str, - headers: dict[str, Any] | None = None, - timeout: httpx.Timeout | None = None, - ) -> None: - self.base_api_url = base_api_url - self.headers = { - 'Content-Type': 'application/json', - 'Accept-encoding': 'gzip', - } - if api_token.startswith('Bearer '): - self.headers['Authorization'] = api_token - else: - self.headers['X-StorageAPI-Token'] = api_token - self.timeout = timeout or httpx.Timeout(connect=5.0, read=60.0, write=10.0, pool=5.0) - if headers: - self.headers.update(headers) - - @staticmethod - def _raise_for_status(response: httpx.Response) -> None: - """ - Checks the HTTP response status code and raises an exception with a detailed message. The message will - include "error" and "exceptionId" fields if they are present in the response. - """ - try: - response.raise_for_status() - except httpx.HTTPStatusError as e: - message_parts = [str(e)] - - try: - error_data = response.json() - if error_msg := error_data.get('error'): - message_parts.append(f'API error: {error_msg}') - if exception_id := error_data.get('exceptionId'): - message_parts.append(f'Exception ID: {exception_id}') - message_parts.append('When contacting Keboola support please provide the exception ID.') - - except ValueError: - try: - if response.text: - message_parts.append(f'API error: {response.text}') - except Exception: - pass # should never get here - - raise httpx.HTTPStatusError('\n'.join(message_parts), request=response.request, response=response) from e - - async def get( - self, - endpoint: str, - params: dict[str, Any] | None = None, - headers: dict[str, Any] | None = None, - ) -> JsonStruct: - """ - Makes a GET request to the service API. - - :param endpoint: API endpoint to call - :param params: Query parameters for the request - :param headers: Additional headers for the request - :return: API response as dictionary - """ - headers = self.headers | (headers or {}) - async with httpx.AsyncClient(timeout=self.timeout) as client: - response = await client.get( - f'{self.base_api_url}/{endpoint}', - params=params, - headers=headers, - ) - self._raise_for_status(response) - return cast(JsonStruct, response.json()) - - async def post( - self, - endpoint: str, - data: dict[str, Any] | None = None, - params: dict[str, Any] | None = None, - headers: dict[str, Any] | None = None, - ) -> JsonStruct: - """ - Makes a POST request to the service API. - - :param endpoint: API endpoint to call - :param data: Request payload - :param params: Query parameters for the request - :param headers: Additional headers for the request - :return: API response as dictionary - """ - headers = self.headers | (headers or {}) - async with httpx.AsyncClient(timeout=self.timeout) as client: - response = await client.post( - f'{self.base_api_url}/{endpoint}', - params=params, - headers=headers, - json=data or {}, - ) - self._raise_for_status(response) - return cast(JsonStruct, response.json()) - - async def put( - self, - endpoint: str, - data: dict[str, Any] | None = None, - params: dict[str, Any] | None = None, - headers: dict[str, Any] | None = None, - ) -> JsonStruct: - """ - Makes a PUT request to the service API. - - :param endpoint: API endpoint to call - :param data: Request payload - :param params: Query parameters for the request - :param headers: Additional headers for the request - :return: API response as dictionary - """ - headers = self.headers | (headers or {}) - async with httpx.AsyncClient(timeout=self.timeout) as client: - response = await client.put( - f'{self.base_api_url}/{endpoint}', - params=params, - headers=headers, - json=data or {}, - ) - self._raise_for_status(response) - return cast(JsonStruct, response.json()) - - async def delete( - self, - endpoint: str, - headers: dict[str, Any] | None = None, - ) -> JsonStruct | None: - """ - Makes a DELETE request to the service API. - - :param endpoint: API endpoint to call - :param headers: Additional headers for the request - :return: API response as dictionary - """ - headers = self.headers | (headers or {}) - async with httpx.AsyncClient(timeout=self.timeout) as client: - response = await client.delete( - f'{self.base_api_url}/{endpoint}', - headers=headers, - ) - self._raise_for_status(response) - - if response.content: - return cast(JsonStruct, response.json()) - - return None - - -class KeboolaServiceClient: - """ - Base class for Keboola service clients. - - Implements the basic HTTP methods (GET, POST, PUT, DELETE) - and is used as a base class for clients for individual services. - """ - - def __init__(self, raw_client: RawKeboolaClient) -> None: - """ - Creates a client instance. - - The inherited classes should implement the `create` method - rather than overriding this constructor. - - :param raw_client: The raw client to use - """ - self.raw_client = raw_client - - @classmethod - def create(cls, root_url: str, token: str) -> 'KeboolaServiceClient': - """ - Creates a KeboolaServiceClient from a Keboola Storage API token. - - :param root_url: The root URL of the service API - :param token: The Keboola Storage API token - :return: A new instance of KeboolaServiceClient - """ - return cls(raw_client=RawKeboolaClient(base_api_url=root_url, api_token=token)) - - async def get( - self, - endpoint: str, - params: Optional[dict[str, Any]] = None, - ) -> JsonStruct: - """ - Makes a GET request to the service API. - - :param endpoint: API endpoint to call - :param params: Query parameters for the request - :return: API response as dictionary - """ - return await self.raw_client.get(endpoint=endpoint, params=params) - - async def post( - self, - endpoint: str, - data: Optional[dict[str, Any]] = None, - params: Optional[dict[str, Any]] = None, - ) -> JsonStruct: - """ - Makes a POST request to the service API. - - :param endpoint: API endpoint to call - :param data: Request payload - :param params: Query parameters for the request - :return: API response as dictionary - """ - return await self.raw_client.post(endpoint=endpoint, data=data, params=params) - - async def put( - self, - endpoint: str, - data: Optional[dict[str, Any]] = None, - params: Optional[dict[str, Any]] = None, - ) -> JsonStruct: - """ - Makes a PUT request to the service API. - - :param endpoint: API endpoint to call - :param data: Request payload - :param params: Query parameters for the request - :return: API response as dictionary - """ - return await self.raw_client.put(endpoint=endpoint, data=data, params=params) - - async def delete( - self, - endpoint: str, - ) -> JsonStruct | None: - """ - Makes a DELETE request to the service API. - - :param endpoint: API endpoint to call - :return: API response as dictionary - """ - return await self.raw_client.delete(endpoint=endpoint) - - class GlobalSearchResponse(BaseModel): """The SAPI global search response.""" From d18d013a0a56c34b5e4756af529133826726ad2a Mon Sep 17 00:00:00 2001 From: mariankrotil Date: Thu, 14 Aug 2025 09:24:18 +0200 Subject: [PATCH 04/35] AI-1343 refactor: update imports wrt client submodule changes --- integtests/conftest.py | 2 +- integtests/test_client.py | 2 +- integtests/test_errors.py | 2 +- integtests/test_mcp_server.py | 2 +- integtests/test_validate.py | 2 +- integtests/test_workspace.py | 2 +- integtests/tools/components/test_tools.py | 2 +- integtests/tools/flow/test_tools.py | 2 +- integtests/tools/test_jobs.py | 2 +- integtests/tools/test_search.py | 2 +- integtests/tools/test_storage.py | 2 +- src/keboola_mcp_server/errors.py | 2 +- src/keboola_mcp_server/links.py | 2 +- src/keboola_mcp_server/mcp.py | 2 +- src/keboola_mcp_server/tools/components/model.py | 2 +- src/keboola_mcp_server/tools/components/tools.py | 2 +- src/keboola_mcp_server/tools/components/utils.py | 2 +- src/keboola_mcp_server/tools/doc.py | 2 +- src/keboola_mcp_server/tools/flow/model.py | 2 +- src/keboola_mcp_server/tools/flow/tools.py | 2 +- src/keboola_mcp_server/tools/flow/utils.py | 2 +- src/keboola_mcp_server/tools/jobs.py | 2 +- src/keboola_mcp_server/tools/oauth.py | 2 +- src/keboola_mcp_server/tools/project.py | 2 +- src/keboola_mcp_server/tools/search.py | 2 +- src/keboola_mcp_server/tools/storage.py | 2 +- src/keboola_mcp_server/tools/validation.py | 5 ++--- src/keboola_mcp_server/workspace.py | 2 +- tests/conftest.py | 2 +- tests/test_client.py | 2 +- tests/test_errors.py | 2 +- tests/test_server.py | 6 +++--- tests/tools/components/conftest.py | 2 +- tests/tools/components/test_tools.py | 2 +- tests/tools/components/test_validation.py | 2 +- tests/tools/flow/test_model.py | 2 +- tests/tools/flow/test_tools.py | 2 +- tests/tools/test_doc.py | 2 +- tests/tools/test_jobs.py | 2 +- tests/tools/test_oauth.py | 2 +- tests/tools/test_project.py | 2 +- tests/tools/test_search.py | 2 +- tests/tools/test_sql.py | 2 +- tests/tools/test_storage.py | 2 +- 44 files changed, 47 insertions(+), 48 deletions(-) diff --git a/integtests/conftest.py b/integtests/conftest.py index cf2269b1..7c859ce9 100644 --- a/integtests/conftest.py +++ b/integtests/conftest.py @@ -20,7 +20,7 @@ from mcp.server.session import ServerSession from mcp.shared.context import RequestContext -from keboola_mcp_server.client import KeboolaClient +from keboola_mcp_server.clients.client import KeboolaClient from keboola_mcp_server.config import Config from keboola_mcp_server.mcp import ServerState from keboola_mcp_server.workspace import WorkspaceManager diff --git a/integtests/test_client.py b/integtests/test_client.py index 64a95f51..1c8521c9 100644 --- a/integtests/test_client.py +++ b/integtests/test_client.py @@ -3,7 +3,7 @@ import pytest from integtests.conftest import ProjectDef, TableDef -from keboola_mcp_server.client import AsyncStorageClient, GlobalSearchResponse, KeboolaClient +from keboola_mcp_server.clients.client import AsyncStorageClient, GlobalSearchResponse, KeboolaClient LOG = logging.getLogger(__name__) diff --git a/integtests/test_errors.py b/integtests/test_errors.py index 04507947..3830b2ed 100644 --- a/integtests/test_errors.py +++ b/integtests/test_errors.py @@ -11,7 +11,7 @@ from fastmcp import Context from mcp.types import ClientCapabilities, Implementation, InitializeRequestParams -from keboola_mcp_server.client import KeboolaClient +from keboola_mcp_server.clients.client import KeboolaClient from keboola_mcp_server.errors import tool_errors from keboola_mcp_server.tools.doc import docs_query from keboola_mcp_server.tools.jobs import get_job diff --git a/integtests/test_mcp_server.py b/integtests/test_mcp_server.py index 243a3ce9..9f37a83a 100644 --- a/integtests/test_mcp_server.py +++ b/integtests/test_mcp_server.py @@ -6,7 +6,7 @@ from mcp.types import TextContent from integtests.conftest import AsyncContextClientRunner, AsyncContextServerRemoteRunner, ConfigDef -from keboola_mcp_server.client import KeboolaClient +from keboola_mcp_server.clients.client import KeboolaClient from keboola_mcp_server.config import Config from keboola_mcp_server.server import create_server from keboola_mcp_server.tools.components.model import Configuration diff --git a/integtests/test_validate.py b/integtests/test_validate.py index 1eaaaf57..49574004 100644 --- a/integtests/test_validate.py +++ b/integtests/test_validate.py @@ -16,7 +16,7 @@ import jsonschema import pytest -from keboola_mcp_server.client import ComponentAPIResponse, JsonDict, KeboolaClient +from keboola_mcp_server.clients.client import ComponentAPIResponse, JsonDict, KeboolaClient from keboola_mcp_server.tools.components.model import Component from keboola_mcp_server.tools.validation import KeboolaParametersValidator diff --git a/integtests/test_workspace.py b/integtests/test_workspace.py index 4583c983..2448afc1 100644 --- a/integtests/test_workspace.py +++ b/integtests/test_workspace.py @@ -5,7 +5,7 @@ import requests from kbcstorage.client import Client as SyncStorageClient -from keboola_mcp_server.client import KeboolaClient +from keboola_mcp_server.clients.client import KeboolaClient from keboola_mcp_server.workspace import WorkspaceManager LOG = logging.getLogger(__name__) diff --git a/integtests/tools/components/test_tools.py b/integtests/tools/components/test_tools.py index b907a590..e32686e9 100644 --- a/integtests/tools/components/test_tools.py +++ b/integtests/tools/components/test_tools.py @@ -7,7 +7,7 @@ from mcp.server.fastmcp import Context from integtests.conftest import ConfigDef, ProjectDef -from keboola_mcp_server.client import KeboolaClient, get_metadata_property +from keboola_mcp_server.clients.client import KeboolaClient, get_metadata_property from keboola_mcp_server.config import Config, MetadataField from keboola_mcp_server.links import Link from keboola_mcp_server.server import create_server diff --git a/integtests/tools/flow/test_tools.py b/integtests/tools/flow/test_tools.py index aed7a3bd..b7a6d2f6 100644 --- a/integtests/tools/flow/test_tools.py +++ b/integtests/tools/flow/test_tools.py @@ -8,7 +8,7 @@ from pydantic import ValidationError from integtests.conftest import ConfigDef, ProjectDef -from keboola_mcp_server.client import ( +from keboola_mcp_server.clients.client import ( CONDITIONAL_FLOW_COMPONENT_ID, ORCHESTRATOR_COMPONENT_ID, FlowType, diff --git a/integtests/tools/test_jobs.py b/integtests/tools/test_jobs.py index c0ff61b6..f278d8df 100644 --- a/integtests/tools/test_jobs.py +++ b/integtests/tools/test_jobs.py @@ -5,7 +5,7 @@ from mcp.server.fastmcp import Context from integtests.conftest import ConfigDef, ProjectDef -from keboola_mcp_server.client import KeboolaClient +from keboola_mcp_server.clients.client import KeboolaClient from keboola_mcp_server.links import Link from keboola_mcp_server.tools.components import create_config from keboola_mcp_server.tools.jobs import JobDetail, ListJobsOutput, get_job, list_jobs, run_job diff --git a/integtests/tools/test_search.py b/integtests/tools/test_search.py index a904a68b..aa1d755b 100644 --- a/integtests/tools/test_search.py +++ b/integtests/tools/test_search.py @@ -4,7 +4,7 @@ from fastmcp import Context from integtests.conftest import BucketDef, ConfigDef, TableDef -from keboola_mcp_server.client import KeboolaClient, SuggestedComponent +from keboola_mcp_server.clients.client import KeboolaClient, SuggestedComponent from keboola_mcp_server.tools.search import GlobalSearchOutput, find_component_id, search LOG = logging.getLogger(__name__) diff --git a/integtests/tools/test_storage.py b/integtests/tools/test_storage.py index 38860475..6eca2734 100644 --- a/integtests/tools/test_storage.py +++ b/integtests/tools/test_storage.py @@ -5,7 +5,7 @@ from fastmcp import Context from integtests.conftest import BucketDef, TableDef -from keboola_mcp_server.client import KeboolaClient +from keboola_mcp_server.clients.client import KeboolaClient from keboola_mcp_server.config import MetadataField from keboola_mcp_server.tools.storage import ( BucketDetail, diff --git a/src/keboola_mcp_server/errors.py b/src/keboola_mcp_server/errors.py index 5dae02fe..993093cd 100644 --- a/src/keboola_mcp_server/errors.py +++ b/src/keboola_mcp_server/errors.py @@ -9,7 +9,7 @@ from fastmcp.utilities.types import find_kwarg_by_type from pydantic import BaseModel -from keboola_mcp_server.client import KeboolaClient +from keboola_mcp_server.clients.client import KeboolaClient from keboola_mcp_server.mcp import ServerState, get_http_request_or_none LOG = logging.getLogger(__name__) diff --git a/src/keboola_mcp_server/links.py b/src/keboola_mcp_server/links.py index 3bd42405..c0c8d444 100644 --- a/src/keboola_mcp_server/links.py +++ b/src/keboola_mcp_server/links.py @@ -2,7 +2,7 @@ from pydantic import BaseModel, ConfigDict, Field -from keboola_mcp_server.client import CONDITIONAL_FLOW_COMPONENT_ID, FlowType, KeboolaClient +from keboola_mcp_server.clients.client import CONDITIONAL_FLOW_COMPONENT_ID, FlowType, KeboolaClient URLType = Literal['ui-detail', 'ui-dashboard', 'docs'] diff --git a/src/keboola_mcp_server/mcp.py b/src/keboola_mcp_server/mcp.py index f504b3d3..5300955c 100644 --- a/src/keboola_mcp_server/mcp.py +++ b/src/keboola_mcp_server/mcp.py @@ -26,7 +26,7 @@ from starlette.requests import Request from starlette.types import ASGIApp, Receive, Scope, Send -from keboola_mcp_server.client import KeboolaClient +from keboola_mcp_server.clients.client import KeboolaClient from keboola_mcp_server.config import Config from keboola_mcp_server.oauth import ProxyAccessToken from keboola_mcp_server.workspace import WorkspaceManager diff --git a/src/keboola_mcp_server/tools/components/model.py b/src/keboola_mcp_server/tools/components/model.py index 94f01261..b50bd40a 100644 --- a/src/keboola_mcp_server/tools/components/model.py +++ b/src/keboola_mcp_server/tools/components/model.py @@ -37,7 +37,7 @@ from pydantic import AliasChoices, BaseModel, Field -from keboola_mcp_server.client import ComponentAPIResponse, ConfigurationAPIResponse +from keboola_mcp_server.clients.client import ComponentAPIResponse, ConfigurationAPIResponse from keboola_mcp_server.links import Link # ============================================================================ diff --git a/src/keboola_mcp_server/tools/components/tools.py b/src/keboola_mcp_server/tools/components/tools.py index 2433f89a..e5a10001 100644 --- a/src/keboola_mcp_server/tools/components/tools.py +++ b/src/keboola_mcp_server/tools/components/tools.py @@ -36,7 +36,7 @@ from httpx import HTTPStatusError from pydantic import Field -from keboola_mcp_server.client import ConfigurationAPIResponse, JsonDict, KeboolaClient +from keboola_mcp_server.clients.client import ConfigurationAPIResponse, JsonDict, KeboolaClient from keboola_mcp_server.errors import tool_errors from keboola_mcp_server.links import ProjectLinksManager from keboola_mcp_server.mcp import KeboolaMcpServer, exclude_none_serializer diff --git a/src/keboola_mcp_server/tools/components/utils.py b/src/keboola_mcp_server/tools/components/utils.py index a04a8cc0..b128e0fc 100644 --- a/src/keboola_mcp_server/tools/components/utils.py +++ b/src/keboola_mcp_server/tools/components/utils.py @@ -28,7 +28,7 @@ from httpx import HTTPStatusError from pydantic import AliasChoices, BaseModel, Field -from keboola_mcp_server.client import ComponentAPIResponse, ConfigurationAPIResponse, JsonDict, KeboolaClient +from keboola_mcp_server.clients.client import ComponentAPIResponse, ConfigurationAPIResponse, JsonDict, KeboolaClient from keboola_mcp_server.config import MetadataField from keboola_mcp_server.tools.components.model import ( AllComponentTypes, diff --git a/src/keboola_mcp_server/tools/doc.py b/src/keboola_mcp_server/tools/doc.py index 43a47b01..ffea2863 100644 --- a/src/keboola_mcp_server/tools/doc.py +++ b/src/keboola_mcp_server/tools/doc.py @@ -5,7 +5,7 @@ from fastmcp.tools import FunctionTool from pydantic import BaseModel, Field -from keboola_mcp_server.client import KeboolaClient +from keboola_mcp_server.clients.client import KeboolaClient from keboola_mcp_server.errors import tool_errors LOG = logging.getLogger(__name__) diff --git a/src/keboola_mcp_server/tools/flow/model.py b/src/keboola_mcp_server/tools/flow/model.py index 00e65d07..c5a85313 100644 --- a/src/keboola_mcp_server/tools/flow/model.py +++ b/src/keboola_mcp_server/tools/flow/model.py @@ -7,7 +7,7 @@ from pydantic import AliasChoices, BaseModel, Field -from keboola_mcp_server.client import ORCHESTRATOR_COMPONENT_ID, APIFlowResponse, FlowType +from keboola_mcp_server.clients.client import ORCHESTRATOR_COMPONENT_ID, APIFlowResponse, FlowType from keboola_mcp_server.links import Link # ============================================================================= diff --git a/src/keboola_mcp_server/tools/flow/tools.py b/src/keboola_mcp_server/tools/flow/tools.py index a50b659b..30956e3a 100644 --- a/src/keboola_mcp_server/tools/flow/tools.py +++ b/src/keboola_mcp_server/tools/flow/tools.py @@ -11,7 +11,7 @@ from pydantic import Field from keboola_mcp_server import resources -from keboola_mcp_server.client import ( +from keboola_mcp_server.clients.client import ( CONDITIONAL_FLOW_COMPONENT_ID, ORCHESTRATOR_COMPONENT_ID, CreateConfigurationAPIResponse, diff --git a/src/keboola_mcp_server/tools/flow/utils.py b/src/keboola_mcp_server/tools/flow/utils.py index d48c5edf..a4d89da4 100644 --- a/src/keboola_mcp_server/tools/flow/utils.py +++ b/src/keboola_mcp_server/tools/flow/utils.py @@ -5,7 +5,7 @@ from importlib import resources from typing import Any, Mapping, Sequence -from keboola_mcp_server.client import ( +from keboola_mcp_server.clients.client import ( CONDITIONAL_FLOW_COMPONENT_ID, FLOW_TYPES, ORCHESTRATOR_COMPONENT_ID, diff --git a/src/keboola_mcp_server/tools/jobs.py b/src/keboola_mcp_server/tools/jobs.py index 50b14a4d..229e5181 100644 --- a/src/keboola_mcp_server/tools/jobs.py +++ b/src/keboola_mcp_server/tools/jobs.py @@ -6,7 +6,7 @@ from fastmcp.tools import FunctionTool from pydantic import AliasChoices, BaseModel, Field, field_validator -from keboola_mcp_server.client import KeboolaClient +from keboola_mcp_server.clients.client import KeboolaClient from keboola_mcp_server.errors import tool_errors from keboola_mcp_server.links import Link, ProjectLinksManager from keboola_mcp_server.mcp import KeboolaMcpServer, exclude_none_serializer diff --git a/src/keboola_mcp_server/tools/oauth.py b/src/keboola_mcp_server/tools/oauth.py index 683b039a..a50efa3b 100644 --- a/src/keboola_mcp_server/tools/oauth.py +++ b/src/keboola_mcp_server/tools/oauth.py @@ -8,7 +8,7 @@ from fastmcp.tools import FunctionTool from pydantic import Field -from keboola_mcp_server.client import KeboolaClient +from keboola_mcp_server.clients.client import KeboolaClient from keboola_mcp_server.errors import tool_errors from keboola_mcp_server.mcp import KeboolaMcpServer diff --git a/src/keboola_mcp_server/tools/project.py b/src/keboola_mcp_server/tools/project.py index 0af228e1..907a18f0 100644 --- a/src/keboola_mcp_server/tools/project.py +++ b/src/keboola_mcp_server/tools/project.py @@ -5,7 +5,7 @@ from fastmcp.tools import FunctionTool from pydantic import BaseModel, Field -from keboola_mcp_server.client import JsonDict, KeboolaClient +from keboola_mcp_server.clients.client import JsonDict, KeboolaClient from keboola_mcp_server.config import MetadataField from keboola_mcp_server.errors import tool_errors from keboola_mcp_server.links import Link, ProjectLinksManager diff --git a/src/keboola_mcp_server/tools/search.py b/src/keboola_mcp_server/tools/search.py index ca8c9885..9179a606 100644 --- a/src/keboola_mcp_server/tools/search.py +++ b/src/keboola_mcp_server/tools/search.py @@ -7,7 +7,7 @@ from fastmcp.tools import FunctionTool from pydantic import BaseModel, Field -from keboola_mcp_server.client import GlobalSearchResponse, ItemType, KeboolaClient, SuggestedComponent +from keboola_mcp_server.clients.client import GlobalSearchResponse, ItemType, KeboolaClient, SuggestedComponent from keboola_mcp_server.errors import tool_errors LOG = logging.getLogger(__name__) diff --git a/src/keboola_mcp_server/tools/storage.py b/src/keboola_mcp_server/tools/storage.py index 23e1fe06..acd11dba 100644 --- a/src/keboola_mcp_server/tools/storage.py +++ b/src/keboola_mcp_server/tools/storage.py @@ -8,7 +8,7 @@ from fastmcp.tools import FunctionTool from pydantic import AliasChoices, BaseModel, Field, model_validator -from keboola_mcp_server.client import JsonDict, KeboolaClient, get_metadata_property +from keboola_mcp_server.clients.client import JsonDict, KeboolaClient, get_metadata_property from keboola_mcp_server.config import MetadataField from keboola_mcp_server.errors import tool_errors from keboola_mcp_server.links import Link, ProjectLinksManager diff --git a/src/keboola_mcp_server/tools/validation.py b/src/keboola_mcp_server/tools/validation.py index 255752a6..324f4fab 100644 --- a/src/keboola_mcp_server/tools/validation.py +++ b/src/keboola_mcp_server/tools/validation.py @@ -11,13 +11,12 @@ import jsonschema import jsonschema.validators -from keboola_mcp_server.client import ( - ORCHESTRATOR_COMPONENT_ID, - FlowType, +from keboola_mcp_server.clients.base import ( JsonDict, JsonPrimitive, JsonStruct, ) +from keboola_mcp_server.clients.client import ORCHESTRATOR_COMPONENT_ID, FlowType from keboola_mcp_server.tools.components.model import Component from keboola_mcp_server.tools.components.utils import BIGQUERY_TRANSFORMATION_ID, SNOWFLAKE_TRANSFORMATION_ID diff --git a/src/keboola_mcp_server/workspace.py b/src/keboola_mcp_server/workspace.py index eb3ca908..4227ea8f 100644 --- a/src/keboola_mcp_server/workspace.py +++ b/src/keboola_mcp_server/workspace.py @@ -9,7 +9,7 @@ from pydantic import Field, TypeAdapter from pydantic.dataclasses import dataclass -from keboola_mcp_server.client import KeboolaClient +from keboola_mcp_server.clients.client import KeboolaClient LOG = logging.getLogger(__name__) diff --git a/tests/conftest.py b/tests/conftest.py index eac469ed..01312a84 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -3,7 +3,7 @@ from mcp.server.session import ServerSession from mcp.shared.context import RequestContext -from keboola_mcp_server.client import ( +from keboola_mcp_server.clients.client import ( AIServiceClient, AsyncStorageClient, JobsQueueClient, diff --git a/tests/test_client.py b/tests/test_client.py index 6d7723ca..94592d22 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -5,7 +5,7 @@ import httpx import pytest -from keboola_mcp_server.client import KeboolaClient, RawKeboolaClient +from keboola_mcp_server.clients.client import KeboolaClient, RawKeboolaClient @pytest.fixture diff --git a/tests/test_errors.py b/tests/test_errors.py index 25319b13..e8c617c9 100644 --- a/tests/test_errors.py +++ b/tests/test_errors.py @@ -7,7 +7,7 @@ from fastmcp import Context from mcp.shared.context import RequestContext -from keboola_mcp_server.client import KeboolaClient +from keboola_mcp_server.clients.client import KeboolaClient from keboola_mcp_server.config import Config from keboola_mcp_server.errors import ToolException, tool_errors from keboola_mcp_server.mcp import ServerState diff --git a/tests/test_server.py b/tests/test_server.py index e9659b90..35bfcaee 100644 --- a/tests/test_server.py +++ b/tests/test_server.py @@ -7,7 +7,7 @@ from mcp.types import TextContent from pydantic import Field -from keboola_mcp_server.client import KeboolaClient +from keboola_mcp_server.clients.client import KeboolaClient from keboola_mcp_server.config import Config from keboola_mcp_server.mcp import ServerState from keboola_mcp_server.server import create_server @@ -150,7 +150,7 @@ async def assessed_function( os_mock.environ = envs mocker.patch( - 'keboola_mcp_server.client.AsyncStorageClient.verify_token', + 'keboola_mcp_server.clients.client.AsyncStorageClient.verify_token', return_value={'owner': {'features': ['global-search', 'waii-integration', 'hide-conditional-flows']}}, ) @@ -204,7 +204,7 @@ async def test_keboola_injection_and_lifespan( mocker.patch('keboola_mcp_server.server.os.environ', os_environ_params) mocker.patch( - 'keboola_mcp_server.client.AsyncStorageClient.verify_token', + 'keboola_mcp_server.clients.client.AsyncStorageClient.verify_token', return_value={'owner': {'features': ['global-search', 'waii-integration', 'conditional-flows']}}, ) diff --git a/tests/tools/components/conftest.py b/tests/tools/components/conftest.py index 43423d23..4eda08f7 100644 --- a/tests/tools/components/conftest.py +++ b/tests/tools/components/conftest.py @@ -4,7 +4,7 @@ from mcp.server.fastmcp import Context from pytest_mock import MockerFixture -from keboola_mcp_server.client import KeboolaClient +from keboola_mcp_server.clients.client import KeboolaClient @pytest.fixture diff --git a/tests/tools/components/test_tools.py b/tests/tools/components/test_tools.py index 294e2191..47e575de 100644 --- a/tests/tools/components/test_tools.py +++ b/tests/tools/components/test_tools.py @@ -4,7 +4,7 @@ from mcp.server.fastmcp import Context from pytest_mock import MockerFixture -from keboola_mcp_server.client import KeboolaClient +from keboola_mcp_server.clients.client import KeboolaClient from keboola_mcp_server.links import Link from keboola_mcp_server.tools.components import ( add_config_row, diff --git a/tests/tools/components/test_validation.py b/tests/tools/components/test_validation.py index b6df3f9b..cbd7c5b8 100644 --- a/tests/tools/components/test_validation.py +++ b/tests/tools/components/test_validation.py @@ -5,7 +5,7 @@ import jsonschema import pytest -from keboola_mcp_server.client import ORCHESTRATOR_COMPONENT_ID, ComponentAPIResponse, JsonDict +from keboola_mcp_server.clients.client import ORCHESTRATOR_COMPONENT_ID, ComponentAPIResponse, JsonDict from keboola_mcp_server.tools import validation from keboola_mcp_server.tools.components.model import Component diff --git a/tests/tools/flow/test_model.py b/tests/tools/flow/test_model.py index 2a608882..3b7132b4 100644 --- a/tests/tools/flow/test_model.py +++ b/tests/tools/flow/test_model.py @@ -1,6 +1,6 @@ from typing import Any -from keboola_mcp_server.client import ORCHESTRATOR_COMPONENT_ID, APIFlowResponse +from keboola_mcp_server.clients.client import ORCHESTRATOR_COMPONENT_ID, APIFlowResponse from keboola_mcp_server.tools.flow.model import ( Flow, FlowConfiguration, diff --git a/tests/tools/flow/test_tools.py b/tests/tools/flow/test_tools.py index 64be9c4d..e275ca6e 100644 --- a/tests/tools/flow/test_tools.py +++ b/tests/tools/flow/test_tools.py @@ -7,7 +7,7 @@ from mcp.server.fastmcp import Context from pytest_mock import MockerFixture -from keboola_mcp_server.client import CONDITIONAL_FLOW_COMPONENT_ID, ORCHESTRATOR_COMPONENT_ID, KeboolaClient +from keboola_mcp_server.clients.client import CONDITIONAL_FLOW_COMPONENT_ID, ORCHESTRATOR_COMPONENT_ID, KeboolaClient from keboola_mcp_server.tools.flow.model import ( ConditionalFlowPhase, ConditionalFlowTask, diff --git a/tests/tools/test_doc.py b/tests/tools/test_doc.py index 531ed292..1cdaa2f1 100644 --- a/tests/tools/test_doc.py +++ b/tests/tools/test_doc.py @@ -2,7 +2,7 @@ from mcp.server.fastmcp import Context from pytest_mock import MockerFixture -from keboola_mcp_server.client import DocsQuestionResponse, KeboolaClient +from keboola_mcp_server.clients.client import DocsQuestionResponse, KeboolaClient from keboola_mcp_server.tools.doc import DocsAnswer, docs_query diff --git a/tests/tools/test_jobs.py b/tests/tools/test_jobs.py index a4acb45a..afad52ce 100644 --- a/tests/tools/test_jobs.py +++ b/tests/tools/test_jobs.py @@ -6,7 +6,7 @@ from mcp.server.fastmcp import Context from pytest_mock import MockerFixture -from keboola_mcp_server.client import KeboolaClient +from keboola_mcp_server.clients.client import KeboolaClient from keboola_mcp_server.links import Link from keboola_mcp_server.tools.jobs import ( JobDetail, diff --git a/tests/tools/test_oauth.py b/tests/tools/test_oauth.py index a05dcc11..b02c85cb 100644 --- a/tests/tools/test_oauth.py +++ b/tests/tools/test_oauth.py @@ -5,7 +5,7 @@ import pytest from mcp.server.fastmcp import Context -from keboola_mcp_server.client import KeboolaClient +from keboola_mcp_server.clients.client import KeboolaClient from keboola_mcp_server.tools.oauth import create_oauth_url diff --git a/tests/tools/test_project.py b/tests/tools/test_project.py index c1b2c28c..15906414 100644 --- a/tests/tools/test_project.py +++ b/tests/tools/test_project.py @@ -2,7 +2,7 @@ from mcp.server.fastmcp import Context from pytest_mock import MockerFixture -from keboola_mcp_server.client import KeboolaClient +from keboola_mcp_server.clients.client import KeboolaClient from keboola_mcp_server.config import MetadataField from keboola_mcp_server.links import Link from keboola_mcp_server.tools.project import ProjectInfo, get_project_info diff --git a/tests/tools/test_search.py b/tests/tools/test_search.py index a0d9d896..3d879040 100644 --- a/tests/tools/test_search.py +++ b/tests/tools/test_search.py @@ -5,7 +5,7 @@ from fastmcp import Context from pytest_mock import MockerFixture -from keboola_mcp_server.client import GlobalSearchResponse, KeboolaClient +from keboola_mcp_server.clients.client import GlobalSearchResponse, KeboolaClient from keboola_mcp_server.tools.search import ( DEFAULT_GLOBAL_SEARCH_LIMIT, GlobalSearchOutput, diff --git a/tests/tools/test_sql.py b/tests/tools/test_sql.py index 2bdbf498..e0fd8d4b 100644 --- a/tests/tools/test_sql.py +++ b/tests/tools/test_sql.py @@ -5,7 +5,7 @@ from mcp.server.fastmcp import Context from pydantic import TypeAdapter -from keboola_mcp_server.client import KeboolaClient +from keboola_mcp_server.clients.client import KeboolaClient from keboola_mcp_server.tools.sql import QueryDataOutput, get_sql_dialect, query_data from keboola_mcp_server.workspace import ( QueryResult, diff --git a/tests/tools/test_storage.py b/tests/tools/test_storage.py index 78e590d5..6541c3ae 100644 --- a/tests/tools/test_storage.py +++ b/tests/tools/test_storage.py @@ -6,7 +6,7 @@ from mcp.server.fastmcp import Context from pytest_mock import MockerFixture -from keboola_mcp_server.client import KeboolaClient +from keboola_mcp_server.clients.client import KeboolaClient from keboola_mcp_server.config import MetadataField from keboola_mcp_server.links import Link from keboola_mcp_server.tools.storage import ( From 971cd97b28546927f1f267cb88d4f617ae6e257f Mon Sep 17 00:00:00 2001 From: mariankrotil Date: Thu, 14 Aug 2025 12:27:49 +0200 Subject: [PATCH 05/35] AI-1343 feat: add encryption client to encrpyt values --- src/keboola_mcp_server/clients/encryption.py | 63 ++++++++++++++++++++ 1 file changed, 63 insertions(+) create mode 100644 src/keboola_mcp_server/clients/encryption.py diff --git a/src/keboola_mcp_server/clients/encryption.py b/src/keboola_mcp_server/clients/encryption.py new file mode 100644 index 00000000..c7beb914 --- /dev/null +++ b/src/keboola_mcp_server/clients/encryption.py @@ -0,0 +1,63 @@ +from typing import Any, Optional + +from pydantic import AliasChoices, BaseModel, Field + +from keboola_mcp_server.clients.base import JsonStruct, KeboolaServiceClient, RawKeboolaClient + + +class AsyncEncryptionClient(KeboolaServiceClient): + + def __init__(self, raw_client: RawKeboolaClient) -> None: + """ + Creates an AsyncEncryptionClient from a RawKeboolaClient. + + :param raw_client: The raw client to use + """ + super().__init__(raw_client=raw_client) + + @property + def base_api_url(self) -> str: + return self.raw_client.base_api_url.split('/apps')[0] + + @classmethod + def create( + cls, + root_url: str, + token: str, + headers: dict[str, Any] | None = None, + ) -> 'AsyncEncryptionClient': + """ + Creates an AsyncStorageClient from a Keboola Storage API token. + + :param root_url: The root URL of the service API + :param token: The Keboola Storage API token + :param headers: Additional headers for the requests + :return: A new instance of AsyncEncryptionClient + """ + return cls( + raw_client=RawKeboolaClient( + base_api_url=root_url, + api_token=token, + headers=headers, + ) + ) + + async def encrypt( + self, value: Any, component_id: str, project_id: str, config_id: Optional[str] = None + ) -> JsonStruct: + """ + Get a data app by its ID. + + :param data_app_id: The ID of the data app + :return: The data app + """ + response = await self.raw_client.post( + endpoint='encrypt', + params={ + 'componentId': component_id, + 'projectId': project_id, + 'configId': config_id, + }, + data=value, + ) + return response From cea8f3fa8695462031964b902b13ad1d7bee6536 Mon Sep 17 00:00:00 2001 From: mariankrotil Date: Thu, 14 Aug 2025 12:29:57 +0200 Subject: [PATCH 06/35] AI-1343 style: update code wrt flake8 using black and isort --- src/keboola_mcp_server/clients/client.py | 8 ++++++++ src/keboola_mcp_server/clients/data_science.py | 10 ++++------ src/keboola_mcp_server/clients/encryption.py | 2 -- 3 files changed, 12 insertions(+), 8 deletions(-) diff --git a/src/keboola_mcp_server/clients/client.py b/src/keboola_mcp_server/clients/client.py index a70290b5..bb96d4a8 100644 --- a/src/keboola_mcp_server/clients/client.py +++ b/src/keboola_mcp_server/clients/client.py @@ -16,6 +16,7 @@ RawKeboolaClient, ) from keboola_mcp_server.clients.data_science import AsyncDataScienceClient +from keboola_mcp_server.clients.encryption import AsyncEncryptionClient LOG = logging.getLogger(__name__) @@ -81,6 +82,7 @@ class KeboolaClient: _PREFIX_QUEUE_API_URL = 'https://queue.' _PREFIX_AISERVICE_API_URL = 'https://ai.' _PREFIX_DATA_SCIENCE_API_URL = 'https://data-science.' + _PREFIX_ENCRYPTION_API_URL = 'https://encryption.' @classmethod def from_state(cls, state: Mapping[str, Any]) -> 'KeboolaClient': @@ -110,6 +112,9 @@ def __init__(self, storage_api_token: str, storage_api_url: str, bearer_token: s data_science_api_url = ( f'{self._PREFIX_DATA_SCIENCE_API_URL}{storage_api_url.split(self._PREFIX_STORAGE_API_URL)[1]}' ) + encryption_api_url = ( + f'{self._PREFIX_ENCRYPTION_API_URL}{storage_api_url.split(self._PREFIX_STORAGE_API_URL)[1]}' + ) # Initialize clients for individual services bearer_or_sapi_token = f'Bearer {bearer_token}' if bearer_token else storage_api_token @@ -125,6 +130,9 @@ def __init__(self, storage_api_token: str, storage_api_url: str, bearer_token: s self.data_science_client = AsyncDataScienceClient.create( root_url=data_science_api_url, token=self.token, headers=self._get_headers() ) + self.encryption_client = AsyncEncryptionClient.create( + root_url=encryption_api_url, token=self.token, headers=self._get_headers() + ) @classmethod def _get_user_agent(cls) -> str: diff --git a/src/keboola_mcp_server/clients/data_science.py b/src/keboola_mcp_server/clients/data_science.py index 371cd338..2cdb2485 100644 --- a/src/keboola_mcp_server/clients/data_science.py +++ b/src/keboola_mcp_server/clients/data_science.py @@ -46,6 +46,7 @@ class Parameters(BaseModel): class DataApp(BaseModel): slug: str = Field(description='The slug of the data app') streamlit: dict[str, str] = Field(description='The streamlit config.toml file') + secrets: Optional[dict[str, str]] = Field(description='The secrets of the data app', default=None) size: str = Field(description='The size of the data app') auto_suspend_after_seconds: int = Field( @@ -75,10 +76,9 @@ class AsyncDataScienceClient(KeboolaServiceClient): def __init__(self, raw_client: RawKeboolaClient) -> None: """ - Creates an AsyncStorageClient from a RawKeboolaClient and a branch id. + Creates an AsyncDataScienceClient from a RawKeboolaClient and a branch id. :param raw_client: The raw client to use - :param branch_id: The id of the branch """ super().__init__(raw_client=raw_client) @@ -94,14 +94,12 @@ def create( headers: dict[str, Any] | None = None, ) -> 'AsyncDataScienceClient': """ - Creates an AsyncStorageClient from a Keboola Storage API token. + Creates an AsyncDataScienceClient from a Keboola Storage API token. :param root_url: The root URL of the service API :param token: The Keboola Storage API token - :param version: The version of the API to use (default: 'v2') - :param branch_id: The id of the branch :param headers: Additional headers for the requests - :return: A new instance of AsyncStorageClient + :return: A new instance of AsyncDataScienceClient """ return cls( raw_client=RawKeboolaClient( diff --git a/src/keboola_mcp_server/clients/encryption.py b/src/keboola_mcp_server/clients/encryption.py index c7beb914..2a7dae53 100644 --- a/src/keboola_mcp_server/clients/encryption.py +++ b/src/keboola_mcp_server/clients/encryption.py @@ -1,7 +1,5 @@ from typing import Any, Optional -from pydantic import AliasChoices, BaseModel, Field - from keboola_mcp_server.clients.base import JsonStruct, KeboolaServiceClient, RawKeboolaClient From a3e98def55c4f1e3c5734f93af2a54fe799eea7d Mon Sep 17 00:00:00 2001 From: mariankrotil Date: Thu, 14 Aug 2025 15:52:52 +0200 Subject: [PATCH 07/35] AI-1343 fix: update endpoints, docs, refactor --- .../clients/data_science.py | 48 +++++++++++++------ src/keboola_mcp_server/clients/encryption.py | 22 +++++---- 2 files changed, 47 insertions(+), 23 deletions(-) diff --git a/src/keboola_mcp_server/clients/data_science.py b/src/keboola_mcp_server/clients/data_science.py index 2cdb2485..3c444830 100644 --- a/src/keboola_mcp_server/clients/data_science.py +++ b/src/keboola_mcp_server/clients/data_science.py @@ -25,23 +25,33 @@ class DataAppResponse(BaseModel): desired_state: str = Field( validation_alias=AliasChoices('desiredState', 'desired_state'), description='The desired state' ) - last_request_timestamp: str = Field( + last_request_timestamp: Optional[str] = Field( validation_alias=AliasChoices('lastRequestTimestamp', 'last_request_timestamp'), + default=None, description='The last request timestamp', ) - last_start_timestamp: str = Field( + last_start_timestamp: Optional[str] = Field( validation_alias=AliasChoices('lastStartTimestamp', 'last_start_timestamp'), + default=None, description='The last start timestamp', ) - url: str = Field(validation_alias=AliasChoices('url', 'url'), description='The URL of the running data app') + url: Optional[str] = Field( + validation_alias=AliasChoices('url', 'url'), description='The URL of the running data app', default=None + ) auto_suspend_after_seconds: int = Field( validation_alias=AliasChoices('autoSuspendAfterSeconds', 'auto_suspend_after_seconds'), description='The auto suspend after seconds', ) - size: Optional[str] = Field(validation_alias=AliasChoices('size', 'size'), description='The size of the data app') + size: Optional[str] = Field( + validation_alias=AliasChoices('size', 'size'), description='The size of the data app', default=None + ) class DataAppConfig(BaseModel): + """ + The simplified data app config model, which is used for creating a data app within the mcp server. + """ + class Parameters(BaseModel): class DataApp(BaseModel): slug: str = Field(description='The slug of the data app') @@ -54,9 +64,13 @@ class DataApp(BaseModel): serialization_alias='autoSuspendAfterSeconds', description='The auto suspend after seconds', ) - data_app: DataApp = Field(description='The data app sub config', serialization_alias='dataApp') + data_app: DataApp = Field( + description='The data app sub config', + serialization_alias='dataApp', + validation_alias=AliasChoices('dataApp', 'data_app'), + ) id: Optional[str] = Field(description='The id of the data app', default=None) - script: Optional[list[str]] = Field(description='The script of the data app', default=None) + script: Optional[str] = Field(description='The script of the data app', default=None) packages: Optional[list[str]] = Field( description='The python packages needed to be installed in the data app', default=None ) @@ -84,7 +98,7 @@ def __init__(self, raw_client: RawKeboolaClient) -> None: @property def base_api_url(self) -> str: - return self.raw_client.base_api_url.split('/apps')[0] + return self.raw_client.base_api_url @classmethod def create( @@ -103,7 +117,7 @@ def create( """ return cls( raw_client=RawKeboolaClient( - base_api_url=f'{root_url}/apps', + base_api_url=f'{root_url}', api_token=token, headers=headers, ) @@ -116,7 +130,7 @@ async def get_data_app(self, data_app_id: str) -> DataAppResponse: :param data_app_id: The ID of the data app :return: The data app """ - response = await self.raw_client.get(endpoint=data_app_id) + response = await self.raw_client.get(endpoint=f'apps/{data_app_id}') return DataAppResponse.model_validate(response) async def create_data_app( @@ -125,18 +139,22 @@ async def create_data_app( description: str, parameters: dict[str, Any], authorization: dict[str, Any], - ) -> dict[str, Any]: + ) -> DataAppResponse: """ Create a data app. - - :param data_app_id: The ID of the data app + :param name: The name of the data app + :param description: The description of the data app + :param parameters: The parameters of the data app + :param authorization: The authorization of the data app :return: The data app """ + # Validate the parameters and authorization _params = DataAppConfig.Parameters.model_validate(parameters).model_dump(exclude_none=True, by_alias=True) _authorization = DataAppConfig.Authorization.model_validate(authorization).model_dump( exclude_none=True, by_alias=True ) - params = { + data = { + 'branchId': None, 'name': name, 'type': 'streamlit', 'description': description, @@ -145,5 +163,5 @@ async def create_data_app( 'authorization': _authorization, }, } - response = await self.raw_client.post('', params=params) - return response + response = await self.raw_client.post('apps', data=data) + return DataAppResponse.model_validate(response) diff --git a/src/keboola_mcp_server/clients/encryption.py b/src/keboola_mcp_server/clients/encryption.py index 2a7dae53..46785870 100644 --- a/src/keboola_mcp_server/clients/encryption.py +++ b/src/keboola_mcp_server/clients/encryption.py @@ -1,6 +1,6 @@ from typing import Any, Optional -from keboola_mcp_server.clients.base import JsonStruct, KeboolaServiceClient, RawKeboolaClient +from keboola_mcp_server.clients.base import KeboolaServiceClient, RawKeboolaClient class AsyncEncryptionClient(KeboolaServiceClient): @@ -15,7 +15,7 @@ def __init__(self, raw_client: RawKeboolaClient) -> None: @property def base_api_url(self) -> str: - return self.raw_client.base_api_url.split('/apps')[0] + return self.raw_client.base_api_url @classmethod def create( @@ -25,7 +25,7 @@ def create( headers: dict[str, Any] | None = None, ) -> 'AsyncEncryptionClient': """ - Creates an AsyncStorageClient from a Keboola Storage API token. + Creates an AsyncEncryptionClient from a Keboola Storage API token. :param root_url: The root URL of the service API :param token: The Keboola Storage API token @@ -41,13 +41,19 @@ def create( ) async def encrypt( - self, value: Any, component_id: str, project_id: str, config_id: Optional[str] = None - ) -> JsonStruct: + self, value: Any, project_id: str, component_id: Optional[str] = None, config_id: Optional[str] = None + ) -> Any: """ - Get a data app by its ID. + Encrypt a value using the encryption service, returns encrypted value. + if value is a dict, values whose keys start with '#' are encrypted. + if value is a str, it is encrypted. + if value contains already encrypted values, they are returned as is. - :param data_app_id: The ID of the data app - :return: The data app + :param value: The value to encrypt + :param project_id: The project ID + :param component_id: The component ID (optional) + :param config_id: The config ID (optional) + :return: The encrypted value, same type as input """ response = await self.raw_client.post( endpoint='encrypt', From 6dfc84d7cbc9028d8531634cd761f84f80909f77 Mon Sep 17 00:00:00 2001 From: mariankrotil Date: Fri, 15 Aug 2025 14:34:29 +0200 Subject: [PATCH 08/35] AI-1343 feat: update data science client, add password, deploy, and suspend data app methods --- .../clients/data_science.py | 41 +++++++++++++++++-- 1 file changed, 37 insertions(+), 4 deletions(-) diff --git a/src/keboola_mcp_server/clients/data_science.py b/src/keboola_mcp_server/clients/data_science.py index 3c444830..df3c47cd 100644 --- a/src/keboola_mcp_server/clients/data_science.py +++ b/src/keboola_mcp_server/clients/data_science.py @@ -1,4 +1,4 @@ -from typing import Any, Optional +from typing import Any, Optional, cast from pydantic import AliasChoices, BaseModel, Field @@ -70,7 +70,7 @@ class DataApp(BaseModel): validation_alias=AliasChoices('dataApp', 'data_app'), ) id: Optional[str] = Field(description='The id of the data app', default=None) - script: Optional[str] = Field(description='The script of the data app', default=None) + script: Optional[list[str]] = Field(description='The script of the data app', default=None) packages: Optional[list[str]] = Field( description='The python packages needed to be installed in the data app', default=None ) @@ -130,9 +130,42 @@ async def get_data_app(self, data_app_id: str) -> DataAppResponse: :param data_app_id: The ID of the data app :return: The data app """ - response = await self.raw_client.get(endpoint=f'apps/{data_app_id}') + response = await self.get(endpoint=f'apps/{data_app_id}') return DataAppResponse.model_validate(response) + async def deploy_data_app(self, data_app_id: str, config_version: str) -> DataAppResponse: + """ + Deploy a data app by its ID. + + :param data_app_id: The ID of the data app + :param config_version: The version of the config to deploy + :return: The data app + """ + data = { + 'desiredState': 'running', + 'configVersion': config_version, + 'restartIfRunning': True, + 'updateDependencies': True, + } + response = await self.patch(endpoint=f'apps/{data_app_id}', data=data) + return DataAppResponse.model_validate(response) + + async def suspend_data_app(self, data_app_id: str) -> DataAppResponse: + """ + Suspend a data app by its ID. + """ + data = {'desiredState': 'stopped'} + response = await self.patch(endpoint=f'apps/{data_app_id}', data=data) + return DataAppResponse.model_validate(response) + + async def get_data_app_password(self, data_app_id: str) -> str: + """ + Get the password for a data app by its ID. + """ + response = await self.get(endpoint=f'apps/{data_app_id}/password') + assert isinstance(response, dict) + return cast(str, response['password']) + async def create_data_app( self, name: str, @@ -163,5 +196,5 @@ async def create_data_app( 'authorization': _authorization, }, } - response = await self.raw_client.post('apps', data=data) + response = await self.post(endpoint='apps', data=data) return DataAppResponse.model_validate(response) From 4146e066292ee38e8020d06619b1e8bc3105ba5f Mon Sep 17 00:00:00 2001 From: mariankrotil Date: Fri, 15 Aug 2025 14:35:08 +0200 Subject: [PATCH 09/35] AI-1343 feat: add patch method to the base client --- src/keboola_mcp_server/clients/base.py | 43 ++++++++++++++++++++++++++ 1 file changed, 43 insertions(+) diff --git a/src/keboola_mcp_server/clients/base.py b/src/keboola_mcp_server/clients/base.py index 70d0007a..3fe86626 100644 --- a/src/keboola_mcp_server/clients/base.py +++ b/src/keboola_mcp_server/clients/base.py @@ -167,6 +167,33 @@ async def delete( return None + async def patch( + self, + endpoint: str, + data: Optional[dict[str, Any]] = None, + params: Optional[dict[str, Any]] = None, + headers: Optional[dict[str, Any]] = None, + ) -> JsonStruct: + """ + Makes a PATCH request to the service API. + + :param endpoint: API endpoint to call + :param data: Request payload + :param params: Query parameters for the request + :param headers: Additional headers for the request + :return: API response as dictionary + """ + headers = self.headers | (headers or {}) + async with httpx.AsyncClient(timeout=self.timeout) as client: + response = await client.patch( + f'{self.base_api_url}/{endpoint}', + params=params, + headers=headers, + json=data or {}, + ) + self._raise_for_status(response) + return cast(JsonStruct, response.json()) + class KeboolaServiceClient: """ @@ -255,3 +282,19 @@ async def delete( :return: API response as dictionary """ return await self.raw_client.delete(endpoint=endpoint) + + async def patch( + self, + endpoint: str, + data: Optional[dict[str, Any]] = None, + params: Optional[dict[str, Any]] = None, + ) -> JsonStruct: + """ + Makes a PATCH request to the service API. + + :param endpoint: API endpoint to call + :param data: Request payload + :param params: Query parameters for the request + :return: API response as dictionary + """ + return await self.raw_client.patch(endpoint=endpoint, data=data, params=params) From f9a9e229ea55de9f5dc8f416259b49e5496cf0e3 Mon Sep 17 00:00:00 2001 From: mariankrotil Date: Fri, 15 Aug 2025 14:36:05 +0200 Subject: [PATCH 10/35] AI-1343 feat: add method getting versions of component from versions endpoint to the storage client --- src/keboola_mcp_server/clients/client.py | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/src/keboola_mcp_server/clients/client.py b/src/keboola_mcp_server/clients/client.py index bb96d4a8..df2d7bb3 100644 --- a/src/keboola_mcp_server/clients/client.py +++ b/src/keboola_mcp_server/clients/client.py @@ -594,6 +594,25 @@ async def configuration_row_detail(self, component_id: str, config_id: str, conf endpoint = f'branch/{self.branch_id}/components/{component_id}/configs/{config_id}/rows/{configuration_row_id}' return cast(JsonDict, await self.get(endpoint=endpoint)) + async def configuration_versions(self, component_id: str, config_id: str) -> list[JsonDict]: + """ + Retrieves details of a specific configuration version. + """ + endpoint = f'branch/{self.branch_id}/components/{component_id}/configs/{config_id}/versions' + return cast(list[JsonDict], await self.get(endpoint=endpoint)) + + async def configuration_version_latest(self, component_id: str, config_id: str) -> int: + """ + Retrieves details of the last configuration version. + """ + versions = await self.configuration_versions(component_id, config_id) + latest_version = 0 + for data in versions: + assert isinstance(data, dict) and isinstance(data['version'], int) + if latest_version is None or data['version'] > latest_version: + latest_version = data['version'] + return latest_version + async def job_detail(self, job_id: str | int) -> JsonDict: """ NOTE: To get info for regular jobs, use the Job Queue API. From 4ecb6d1c6f329b8de8475a55ad5d998ecdcd14a0 Mon Sep 17 00:00:00 2001 From: mariankrotil Date: Fri, 15 Aug 2025 14:36:59 +0200 Subject: [PATCH 11/35] AI-1343 style: apply flake8 changes --- src/keboola_mcp_server/clients/client.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/keboola_mcp_server/clients/client.py b/src/keboola_mcp_server/clients/client.py index df2d7bb3..17233da6 100644 --- a/src/keboola_mcp_server/clients/client.py +++ b/src/keboola_mcp_server/clients/client.py @@ -608,7 +608,8 @@ async def configuration_version_latest(self, component_id: str, config_id: str) versions = await self.configuration_versions(component_id, config_id) latest_version = 0 for data in versions: - assert isinstance(data, dict) and isinstance(data['version'], int) + assert isinstance(data, dict) + assert isinstance(data['version'], int) if latest_version is None or data['version'] > latest_version: latest_version = data['version'] return latest_version From 92042e378eb7605eb3e2fa72c9941329d49e6c57 Mon Sep 17 00:00:00 2001 From: mariankrotil Date: Mon, 18 Aug 2025 11:21:01 +0200 Subject: [PATCH 12/35] AI-1343 feat(client): add handlers for listing, getting logs and deleting given the corresponding endpoints --- .../clients/data_science.py | 53 +++++++++++++++++++ 1 file changed, 53 insertions(+) diff --git a/src/keboola_mcp_server/clients/data_science.py b/src/keboola_mcp_server/clients/data_science.py index df3c47cd..460bcaf4 100644 --- a/src/keboola_mcp_server/clients/data_science.py +++ b/src/keboola_mcp_server/clients/data_science.py @@ -1,9 +1,13 @@ +import logging +from datetime import datetime, timedelta, timezone from typing import Any, Optional, cast from pydantic import AliasChoices, BaseModel, Field from keboola_mcp_server.clients.base import KeboolaServiceClient, RawKeboolaClient +LOG = logging.getLogger(__name__) + class DataAppResponse(BaseModel): id: str = Field(validation_alias=AliasChoices('id', 'data_app_id'), description='The data app ID') @@ -84,6 +88,7 @@ class AppProxy(BaseModel): parameters: Parameters = Field(description='The parameters of the data app') authorization: Authorization = Field(description='The authorization of the data app') + storage: dict[str, Any] = Field(description='The storage of the data app', default_factory=dict) class AsyncDataScienceClient(KeboolaServiceClient): @@ -198,3 +203,51 @@ async def create_data_app( } response = await self.post(endpoint='apps', data=data) return DataAppResponse.model_validate(response) + + async def delete_data_app(self, data_app_id: str) -> None: + """ + Delete a data app by its ID. + """ + await self.delete(endpoint=f'apps/{data_app_id}') + + async def list_data_apps(self, limit: int = 100, offset: int = 0) -> list[DataAppResponse]: + """ + List all data apps. + """ + response = await self.get(endpoint='apps', params={'limit': limit, 'offset': offset}) + return [DataAppResponse.model_validate(app) for app in response] + + async def tail_app_logs(self, app_id: str, since: Optional[str] = None, *, lines: Optional[int] = None) -> str: + """ + Tail application logs. Either `since` or `lines` must be provided but not both at the same time, otherwise it + uses the `lines` parameter. In case when none of the parameters are provided, it uses the `lines` parameter with + the last 100 lines. + :param app_id: ID of the app. + :param since: ISO-8601 timestamp with nanoseconds. + E.g: since = (datetime.now(timezone.utc) - timedelta(days=1)).isoformat() + :param lines: Number of log lines from the end. Defaults to 100. + :return: Logs as plain text. + :raise requests.HTTPError: For non-200 status codes. + """ + if since and lines: + LOG.warning( + 'You cannot use both "since" and "lines" query parameters together. Using the "lines" parameter.' + ) + since = None + + elif not since and not lines: + LOG.info( + 'No "since" or "lines" query parameters provided. Using "lines" with the last 100 lines as default.' + ) + lines = 100 + + if lines: + lines = max(lines, 1) # Ensure lines is at least 1 + params = {'lines': lines} + elif since: + params = {'since': since} + else: + raise ValueError('Either "since" or "lines" must be provided.') + + response = await self.get_text(endpoint=f'apps/{app_id}/logs/tail', params=params) + return cast(str, response) From 2eab15aff0c794d8ee319fa27366b48fa074cea3 Mon Sep 17 00:00:00 2001 From: mariankrotil Date: Mon, 18 Aug 2025 11:23:39 +0200 Subject: [PATCH 13/35] AI-1343 feat: add get method returning text value instead of json --- src/keboola_mcp_server/clients/base.py | 38 ++++++++++++++++++++++++++ 1 file changed, 38 insertions(+) diff --git a/src/keboola_mcp_server/clients/base.py b/src/keboola_mcp_server/clients/base.py index 3fe86626..bdce266c 100644 --- a/src/keboola_mcp_server/clients/base.py +++ b/src/keboola_mcp_server/clients/base.py @@ -88,6 +88,30 @@ async def get( self._raise_for_status(response) return cast(JsonStruct, response.json()) + async def get_text( + self, + endpoint: str, + params: dict[str, Any] | None = None, + headers: dict[str, Any] | None = None, + ) -> str: + """ + Makes a GET request to the service API. + + :param endpoint: API endpoint to call + :param params: Query parameters for the request + :param headers: Additional headers for the request + :return: API response as dictionary + """ + headers = self.headers | (headers or {}) + async with httpx.AsyncClient(timeout=self.timeout) as client: + response = await client.get( + f'{self.base_api_url}/{endpoint}', + params=params, + headers=headers, + ) + self._raise_for_status(response) + return cast(str, response.text) + async def post( self, endpoint: str, @@ -239,6 +263,20 @@ async def get( """ return await self.raw_client.get(endpoint=endpoint, params=params) + async def get_text( + self, + endpoint: str, + params: Optional[dict[str, Any]] = None, + ) -> str: + """ + Makes a GET request to the service API. + + :param endpoint: API endpoint to call + :param params: Query parameters for the request + :return: API response as text + """ + return await self.raw_client.get_text(endpoint=endpoint, params=params) + async def post( self, endpoint: str, From 9b08f892e780d1908a04761b417bbaebbc7dfe25 Mon Sep 17 00:00:00 2001 From: mariankrotil Date: Mon, 18 Aug 2025 11:24:25 +0200 Subject: [PATCH 14/35] AI-1343 test: Move tests for client to the corresponding dir, add tests for data science client methods --- integtests/{ => clients}/test_client.py | 0 integtests/clients/test_data_science.py | 106 ++++++++++++++++++++++++ 2 files changed, 106 insertions(+) rename integtests/{ => clients}/test_client.py (100%) create mode 100644 integtests/clients/test_data_science.py diff --git a/integtests/test_client.py b/integtests/clients/test_client.py similarity index 100% rename from integtests/test_client.py rename to integtests/clients/test_client.py diff --git a/integtests/clients/test_data_science.py b/integtests/clients/test_data_science.py new file mode 100644 index 00000000..542ee534 --- /dev/null +++ b/integtests/clients/test_data_science.py @@ -0,0 +1,106 @@ +import logging + +import httpx +import pytest + +from keboola_mcp_server.clients.client import DATA_APP_COMPONENT_ID, KeboolaClient +from keboola_mcp_server.clients.data_science import AsyncDataScienceClient, DataAppResponse + +LOG = logging.getLogger(__name__) + + +def _minimal_parameters(slug: str) -> dict[str, object]: + """Build minimal valid parameters for a code-based Streamlit data app.""" + return { + 'size': 'tiny', + 'autoSuspendAfterSeconds': 600, + 'dataApp': { + 'slug': slug, + 'streamlit': { + 'config.toml': '[theme]\nbase = "light"', + }, + }, + 'script': [ + 'import streamlit as st', + "st.write('Hello from integration test')", + ], + } + + +def _public_access_authorization() -> dict[str, object]: + """Allow public access to all paths; no providers required.""" + return { + 'app_proxy': { + 'auth_providers': [], + 'auth_rules': [ + {'type': 'pathPrefix', 'value': '/', 'auth_required': False}, + ], + } + } + + +@pytest.fixture +def ds_client(keboola_client: KeboolaClient) -> AsyncDataScienceClient: + return keboola_client.data_science_client + + +@pytest.mark.asyncio +async def test_create_and_fetch_data_app( + ds_client: AsyncDataScienceClient, unique_id: str, keboola_client: KeboolaClient +) -> None: + """Test creating a data app and fetching it from detail and list endpoints""" + slug = f'test-app-{unique_id}' + created: DataAppResponse = await ds_client.create_data_app( + name=f'IntegTest {slug}', + description='Created by integration tests', + parameters=_minimal_parameters(slug), + authorization=_public_access_authorization(), + ) + + try: + # Check if the created data app is valid + assert isinstance(created, DataAppResponse) + assert created.id + assert created.type == 'streamlit' + assert created.component_id == DATA_APP_COMPONENT_ID + + # Deploy the data app + response = await ds_client.deploy_data_app(created.id, created.config_version) + assert response.id == created.id + + # Fetch the data app from data science + fethced_ds = await ds_client.get_data_app(created.id) + assert fethced_ds.id == created.id + assert fethced_ds.type == created.type + assert fethced_ds.component_id == created.component_id + assert fethced_ds.project_id == created.project_id + assert fethced_ds.config_id == created.config_id + assert fethced_ds.config_version == created.config_version + + # Fetch the data app config from storage + fetched_s = await keboola_client.storage_client.configuration_detail( + component_id=DATA_APP_COMPONENT_ID, + configuration_id=created.config_id, + ) + + # check if the data app ids are the same (data app from data science and config from storage) + assert 'configuration' in fetched_s + assert isinstance(fetched_s['configuration'], dict) + assert 'parameters' in fetched_s['configuration'] + assert isinstance(fetched_s['configuration']['parameters'], dict) + assert 'id' in fetched_s['configuration']['parameters'] + assert fethced_ds.id == fetched_s['configuration']['parameters']['id'] + + # Fetch the all data apps and check if the created data app is in the list + data_apps = await ds_client.list_data_apps() + assert isinstance(data_apps, list) + assert len(data_apps) > 0 + assert any(app.id == created.id for app in data_apps) + + finally: + for _ in range(2): # Delete configuration 2 times (from storage and then from temporal bin) + try: + await ds_client.delete_data_app(created.id) + except Exception as e: + LOG.exception(f'Failed to delete data app: {e}') + pass From ad226b5314dd398d3fc3775bcafd5acd1005125a Mon Sep 17 00:00:00 2001 From: mariankrotil Date: Mon, 18 Aug 2025 11:25:45 +0200 Subject: [PATCH 15/35] AI-1343 refactor: move client tests to the cor. dir --- tests/{ => clients}/test_client.py | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename tests/{ => clients}/test_client.py (100%) diff --git a/tests/test_client.py b/tests/clients/test_client.py similarity index 100% rename from tests/test_client.py rename to tests/clients/test_client.py From 4c9a436853fee4e2bc3da5a18a9b81ef19187074 Mon Sep 17 00:00:00 2001 From: mariankrotil Date: Mon, 18 Aug 2025 11:29:57 +0200 Subject: [PATCH 16/35] AI-1343 style: apply black, isort and flake8 changes --- integtests/clients/test_data_science.py | 1 - src/keboola_mcp_server/clients/data_science.py | 1 - 2 files changed, 2 deletions(-) diff --git a/integtests/clients/test_data_science.py b/integtests/clients/test_data_science.py index 542ee534..8e01fe68 100644 --- a/integtests/clients/test_data_science.py +++ b/integtests/clients/test_data_science.py @@ -1,6 +1,5 @@ import logging -import httpx import pytest from keboola_mcp_server.clients.client import DATA_APP_COMPONENT_ID, KeboolaClient diff --git a/src/keboola_mcp_server/clients/data_science.py b/src/keboola_mcp_server/clients/data_science.py index 460bcaf4..6db24e1f 100644 --- a/src/keboola_mcp_server/clients/data_science.py +++ b/src/keboola_mcp_server/clients/data_science.py @@ -1,5 +1,4 @@ import logging -from datetime import datetime, timedelta, timezone from typing import Any, Optional, cast from pydantic import AliasChoices, BaseModel, Field From 9a31c75129577e1aedac7a74240dadcf4cfa6aa2 Mon Sep 17 00:00:00 2001 From: mariankrotil Date: Mon, 18 Aug 2025 17:22:51 +0200 Subject: [PATCH 17/35] AI-1343 test: add client ecryption integ tests --- integtests/clients/test_encryption.py | 37 +++++++++++++++++++++++++++ 1 file changed, 37 insertions(+) create mode 100644 integtests/clients/test_encryption.py diff --git a/integtests/clients/test_encryption.py b/integtests/clients/test_encryption.py new file mode 100644 index 00000000..ceb0b433 --- /dev/null +++ b/integtests/clients/test_encryption.py @@ -0,0 +1,37 @@ +from typing import Any + +import pytest + +from keboola_mcp_server.clients.client import DATA_APP_COMPONENT_ID, KeboolaClient + + +@pytest.mark.asyncio +async def test_encrypt_string_not_equal(keboola_client: KeboolaClient) -> None: + project_id = await keboola_client.storage_client.project_id() + plaintext = 'my-plain-text' + encrypted = await keboola_client.encryption_client.encrypt( + value=plaintext, + project_id=str(project_id), + component_id=DATA_APP_COMPONENT_ID, + ) + assert isinstance(encrypted, str) + assert encrypted != plaintext + + +@pytest.mark.asyncio +async def test_encrypt_dict_hash_keys_only(keboola_client: KeboolaClient) -> None: + project_id = await keboola_client.storage_client.project_id() + payload: dict[str, Any] = { + '#secret': 'sensitive-value', + 'public': 'visible-value', + } + result = await keboola_client.encryption_client.encrypt( + value=payload, + project_id=str(project_id), + component_id=DATA_APP_COMPONENT_ID, + ) + assert isinstance(result, dict) + # Values under keys beginning with '#' should be encrypted (changed) + assert result['#secret'] != payload['#secret'] + # Non-secret values should remain the same + assert result['public'] == payload['public'] From 86489a9ea5bcfa0377d7025a0c81d8ec07c6c6a3 Mon Sep 17 00:00:00 2001 From: mariankrotil Date: Mon, 18 Aug 2025 17:23:50 +0200 Subject: [PATCH 18/35] AI-1343 refactor: move clients to separate submodules --- integtests/clients/test_client.py | 3 +- integtests/test_validate.py | 4 +- integtests/tools/test_search.py | 3 +- src/keboola_mcp_server/clients/ai_service.py | 85 ++ src/keboola_mcp_server/clients/client.py | 1156 +---------------- src/keboola_mcp_server/clients/jobs_queue.py | 125 ++ src/keboola_mcp_server/clients/storage.py | 952 ++++++++++++++ .../tools/components/model.py | 2 +- .../tools/components/tools.py | 3 +- .../tools/components/utils.py | 4 +- src/keboola_mcp_server/tools/flow/model.py | 3 +- src/keboola_mcp_server/tools/flow/tools.py | 8 +- src/keboola_mcp_server/tools/flow/utils.py | 3 +- src/keboola_mcp_server/tools/project.py | 3 +- src/keboola_mcp_server/tools/search.py | 4 +- src/keboola_mcp_server/tools/storage.py | 3 +- tests/clients/test_client.py | 3 +- tests/conftest.py | 12 +- tests/tools/components/test_validation.py | 3 +- tests/tools/flow/test_model.py | 3 +- tests/tools/test_doc.py | 3 +- tests/tools/test_search.py | 3 +- 22 files changed, 1209 insertions(+), 1179 deletions(-) create mode 100644 src/keboola_mcp_server/clients/ai_service.py create mode 100644 src/keboola_mcp_server/clients/jobs_queue.py create mode 100644 src/keboola_mcp_server/clients/storage.py diff --git a/integtests/clients/test_client.py b/integtests/clients/test_client.py index 1c8521c9..6b91d73d 100644 --- a/integtests/clients/test_client.py +++ b/integtests/clients/test_client.py @@ -3,7 +3,8 @@ import pytest from integtests.conftest import ProjectDef, TableDef -from keboola_mcp_server.clients.client import AsyncStorageClient, GlobalSearchResponse, KeboolaClient +from keboola_mcp_server.clients.client import KeboolaClient +from keboola_mcp_server.clients.storage import AsyncStorageClient, GlobalSearchResponse LOG = logging.getLogger(__name__) diff --git a/integtests/test_validate.py b/integtests/test_validate.py index 49574004..53b7c97f 100644 --- a/integtests/test_validate.py +++ b/integtests/test_validate.py @@ -16,7 +16,9 @@ import jsonschema import pytest -from keboola_mcp_server.clients.client import ComponentAPIResponse, JsonDict, KeboolaClient +from keboola_mcp_server.clients.client import KeboolaClient +from keboola_mcp_server.clients.storage import ComponentAPIResponse +from keboola_mcp_server.clients.base import JsonDict from keboola_mcp_server.tools.components.model import Component from keboola_mcp_server.tools.validation import KeboolaParametersValidator diff --git a/integtests/tools/test_search.py b/integtests/tools/test_search.py index aa1d755b..19140d54 100644 --- a/integtests/tools/test_search.py +++ b/integtests/tools/test_search.py @@ -4,7 +4,8 @@ from fastmcp import Context from integtests.conftest import BucketDef, ConfigDef, TableDef -from keboola_mcp_server.clients.client import KeboolaClient, SuggestedComponent +from keboola_mcp_server.clients.client import KeboolaClient +from keboola_mcp_server.clients.ai_service import SuggestedComponent from keboola_mcp_server.tools.search import GlobalSearchOutput, find_component_id, search LOG = logging.getLogger(__name__) diff --git a/src/keboola_mcp_server/clients/ai_service.py b/src/keboola_mcp_server/clients/ai_service.py new file mode 100644 index 00000000..f16fe9ff --- /dev/null +++ b/src/keboola_mcp_server/clients/ai_service.py @@ -0,0 +1,85 @@ +from typing import Any, cast + +from pydantic import BaseModel, Field + +from keboola_mcp_server.clients.base import JsonDict, KeboolaServiceClient, RawKeboolaClient + + +class DocsQuestionResponse(BaseModel): + """ + The AI service response to a request to `/docs/question` endpoint. + """ + + text: str = Field(description='Text of the answer to a documentation query.') + source_urls: list[str] = Field( + description='List of URLs to the sources of the answer.', + default_factory=list, + alias='sourceUrls', + ) + + +class SuggestedComponent(BaseModel): + """The AI service response to a /docs/suggest-component request.""" + + component_id: str = Field(description='The component ID.', alias='componentId') + score: float = Field(description='Score of the component suggestion.') + source: str = Field(description='Source of the component suggestion.') + + +class ComponentSuggestionResponse(BaseModel): + """The AI service response to a /suggest/component request.""" + + components: list[SuggestedComponent] = Field(description='List of suggested components.', default_factory=list) + + +class AIServiceClient(KeboolaServiceClient): + """Async client for Keboola AI Service.""" + + @classmethod + def create(cls, root_url: str, token: str, headers: dict[str, Any] | None = None) -> 'AIServiceClient': + """ + Creates an AIServiceClient from a Keboola Storage API token. + + :param root_url: The root URL of the AI service API. + :param token: The Keboola Storage API token. + :param headers: Additional headers for the requests. + :return: A new instance of AIServiceClient. + """ + return cls(raw_client=RawKeboolaClient(base_api_url=root_url, api_token=token, headers=headers)) + + async def get_component_detail(self, component_id: str) -> JsonDict: + """ + Retrieves information about a given component. + + :param component_id: The id of the component. + :return: Component details as dictionary. + """ + return cast(JsonDict, await self.get(endpoint=f'docs/components/{component_id}')) + + async def docs_question(self, query: str) -> DocsQuestionResponse: + """ + Answers a question using the Keboola documentation as a source. + :param query: The query to answer. + :return: Response containing the answer and source URLs. + """ + response = await self.raw_client.post( + endpoint='docs/question', + data={'query': query}, + headers={'Accept': 'application/json'}, + ) + + return DocsQuestionResponse.model_validate(response) + + async def suggest_component(self, query: str) -> ComponentSuggestionResponse: + """ + Provides list of component suggestions based on natural language query. + :param query: The query to answer. + :return: Response containing the list of suggested component IDs, their score and source. + """ + response = await self.raw_client.post( + endpoint='suggest/component', + data={'prompt': query}, + headers={'Accept': 'application/json'}, + ) + + return ComponentSuggestionResponse.model_validate(response) diff --git a/src/keboola_mcp_server/clients/client.py b/src/keboola_mcp_server/clients/client.py index 17233da6..ac1b1dbd 100644 --- a/src/keboola_mcp_server/clients/client.py +++ b/src/keboola_mcp_server/clients/client.py @@ -2,45 +2,22 @@ import importlib.metadata import logging -import math import os -from datetime import datetime -from typing import Any, Iterable, Literal, Mapping, Optional, Sequence, TypeVar, cast +from typing import Any, Literal, Mapping, Optional, Sequence, TypeVar -from pydantic import AliasChoices, BaseModel, Field, field_validator - -from keboola_mcp_server.clients.base import ( - JsonDict, - JsonList, - KeboolaServiceClient, - RawKeboolaClient, -) +from keboola_mcp_server.clients.ai_service import AIServiceClient from keboola_mcp_server.clients.data_science import AsyncDataScienceClient from keboola_mcp_server.clients.encryption import AsyncEncryptionClient +from keboola_mcp_server.clients.jobs_queue import JobsQueueClient +from keboola_mcp_server.clients.storage import AsyncStorageClient LOG = logging.getLogger(__name__) T = TypeVar('T') -ComponentResource = Literal['configuration', 'rows', 'state'] -StorageEventType = Literal['info', 'success', 'warn', 'error'] - -# Project features that can be checked with the is_enabled method -ProjectFeature = Literal['global-search'] # Input types for the global search endpoint parameters BranchType = Literal['production', 'development'] -ItemType = Literal[ - 'flow', - 'bucket', - 'table', - 'transformation', - 'configuration', - 'configuration-row', - 'workspace', - 'shared-code', - 'rows', - 'state', -] + ORCHESTRATOR_COMPONENT_ID = 'keboola.orchestrator' CONDITIONAL_FLOW_COMPONENT_ID = 'keboola.flow' @@ -153,1126 +130,3 @@ def _get_headers(cls) -> dict[str, Any]: :return: Additional headers for the requests, namely the user agent. """ return {'User-Agent': cls._get_user_agent()} - - -class GlobalSearchResponse(BaseModel): - """The SAPI global search response.""" - - class Item(BaseModel): - id: str = Field(description='The id of the item.') - name: str = Field(description='The name of the item.') - type: ItemType = 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[Item] = Field(description='List of search results 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: - """ - Creates an AsyncStorageClient from a RawKeboolaClient and a branch id. - - :param raw_client: The raw client to use - :param branch_id: The id of the branch - """ - super().__init__(raw_client=raw_client) - self._branch_id: str = branch_id - - @property - def branch_id(self) -> str: - return self._branch_id - - @property - def base_api_url(self) -> str: - return self.raw_client.base_api_url.split('/v2')[0] - - @classmethod - def create( - cls, - root_url: str, - token: str, - version: str = 'v2', - branch_id: str = 'default', - headers: dict[str, Any] | None = None, - ) -> 'AsyncStorageClient': - """ - Creates an AsyncStorageClient from a Keboola Storage API token. - - :param root_url: The root URL of the service API - :param token: The Keboola Storage API token - :param version: The version of the API to use (default: 'v2') - :param branch_id: The id of the branch - :param headers: Additional headers for the requests - :return: A new instance of AsyncStorageClient - """ - return cls( - raw_client=RawKeboolaClient( - base_api_url=f'{root_url}/{version}/storage', - api_token=token, - headers=headers, - ), - branch_id=branch_id, - ) - - async def branch_metadata_get(self) -> list[JsonDict]: - """ - Retrieves metadata for the current branch. - - :return: Branch metadata as a list of dictionaries. Each dictionary contains the 'key' and 'value' keys. - """ - return cast(list[JsonDict], await self.get(endpoint=f'branch/{self.branch_id}/metadata')) - - async def branch_metadata_update(self, metadata: dict[str, Any]) -> list[JsonDict]: - """ - Updates metadata for the current branch. - - :param metadata: The metadata to update. - :return: The SAPI call response - updated metadata or raise an error. - """ - payload = { - 'metadata': [{'key': key, 'value': value} for key, value in metadata.items()], - } - return cast(list[JsonDict], await self.post(endpoint=f'branch/{self.branch_id}/metadata', data=payload)) - - async def bucket_detail(self, bucket_id: str) -> JsonDict: - """ - Retrieves information about a given bucket. - - :param bucket_id: The id of the bucket - :return: Bucket details as dictionary - """ - return cast(JsonDict, await self.get(endpoint=f'buckets/{bucket_id}')) - - async def bucket_list(self, include: list[str] | None = None) -> list[JsonDict]: - """ - Lists all buckets. - - :param include: List of fields to include in the response ('metadata' or 'linkedBuckets') - :return: List of buckets as dictionary - """ - params = {} - if include is not None and isinstance(include, list): - params['include'] = ','.join(include) - return cast(list[JsonDict], await self.get(endpoint='buckets', params=params)) - - async def bucket_metadata_delete(self, bucket_id: str, metadata_id: str) -> None: - """ - Deletes metadata for a given bucket. - - :param bucket_id: The id of the bucket - :param metadata_id: The id of the metadata - """ - await self.delete(endpoint=f'buckets/{bucket_id}/metadata/{metadata_id}') - - async def bucket_metadata_get(self, bucket_id: str) -> list[JsonDict]: - """ - Retrieves metadata for a given bucket. - - :param bucket_id: The id of the bucket - :return: Bucket metadata as a list of dictionaries. Each dictionary contains the 'key' and 'value' keys. - """ - return cast(list[JsonDict], await self.get(endpoint=f'buckets/{bucket_id}/metadata')) - - async def bucket_metadata_update( - self, - bucket_id: str, - metadata: dict[str, Any], - provider: str = 'user', - ) -> list[JsonDict]: - """ - Updates metadata for a given bucket. - - :param bucket_id: The id of the bucket - :param metadata: The metadata to update. - :param provider: The provider of the metadata ('user' by default). - :return: Bucket metadata as a list of dictionaries. Each dictionary contains the 'key' and 'value' keys. - """ - payload = { - 'provider': provider, - 'metadata': [{'key': key, 'value': value} for key, value in metadata.items()], - } - return cast(list[JsonDict], await self.post(endpoint=f'buckets/{bucket_id}/metadata', data=payload)) - - async def bucket_table_list(self, bucket_id: str, include: list[str] | None = None) -> list[JsonDict]: - """ - Lists all tables in a given bucket. - - :param bucket_id: The id of the bucket - :param include: List of fields to include in the response - :return: List of tables as dictionary - """ - params = {} - if include is not None and isinstance(include, list): - params['include'] = ','.join(include) - return cast(list[JsonDict], await self.get(endpoint=f'buckets/{bucket_id}/tables', params=params)) - - async def component_detail(self, component_id: str) -> JsonDict: - """ - Retrieves information about a given component. - - :param component_id: The id of the component - :return: Component details as a dictionary - """ - return cast(JsonDict, await self.get(endpoint=f'branch/{self.branch_id}/components/{component_id}')) - - async def component_list( - self, component_type: str, include: list[ComponentResource] | None = None - ) -> list[JsonDict]: - """ - Lists all components of a given type. - - :param component_type: The type of the component (extractor, writer, application, etc.) - :param include: Comma separated list of resources to include. - Available resources: configuration, rows and state. - :return: List of components as dictionary - """ - endpoint = f'branch/{self.branch_id}/components' - params = {'componentType': component_type} - if include is not None and isinstance(include, list): - params['include'] = ','.join(include) - - return cast(list[JsonDict], await self.get(endpoint=endpoint, params=params)) - - async def configuration_create( - self, - component_id: str, - name: str, - description: str, - configuration: dict[str, Any], - ) -> JsonDict: - """ - Creates a new configuration for a component. - - :param component_id: The id of the component for which to create the configuration. - :param name: The name of the configuration. - :param description: The description of the configuration. - :param configuration: The configuration definition as a dictionary. - - :return: The SAPI call response - created configuration or raise an error. - """ - endpoint = f'branch/{self.branch_id}/components/{component_id}/configs' - - payload = { - 'name': name, - 'description': description, - 'configuration': configuration, - } - return cast(JsonDict, await self.post(endpoint=endpoint, data=payload)) - - async def configuration_delete(self, component_id: str, configuration_id: str, skip_trash: bool = False) -> None: - """ - Deletes a configuration. - - :param component_id: The id of the component. - :param configuration_id: The id of the configuration. - :param skip_trash: If True, the configuration is deleted without moving to the trash. - (Technically it means the API endpoint is called twice.) - :raises httpx.HTTPStatusError: If the (component_id, configuration_id) is not found. - """ - endpoint = f'branch/{self.branch_id}/components/{component_id}/configs/{configuration_id}' - await self.delete(endpoint=endpoint) - if skip_trash: - await self.delete(endpoint=endpoint) - - async def configuration_detail(self, component_id: str, configuration_id: str) -> JsonDict: - """ - Retrieves information about a given configuration. - - :param component_id: The id of the component. - :param configuration_id: The id of the configuration. - :return: The parsed json from the HTTP response. - :raises ValueError: If the component_id or configuration_id is invalid. - """ - if not isinstance(component_id, str) or component_id == '': - raise ValueError(f"Invalid component_id '{component_id}'.") - if not isinstance(configuration_id, str) or configuration_id == '': - raise ValueError(f"Invalid configuration_id '{configuration_id}'.") - endpoint = f'branch/{self.branch_id}/components/{component_id}/configs/{configuration_id}' - - return cast(JsonDict, await self.get(endpoint=endpoint)) - - async def configuration_list(self, component_id: str) -> list[JsonDict]: - """ - Lists configurations of the given component. - - :param component_id: The id of the component. - :return: List of configurations. - :raises ValueError: If the component_id is invalid. - """ - if not isinstance(component_id, str) or component_id == '': - raise ValueError(f"Invalid component_id '{component_id}'.") - endpoint = f'branch/{self.branch_id}/components/{component_id}/configs' - - return cast(list[JsonDict], await self.get(endpoint=endpoint)) - - async def configuration_metadata_get(self, component_id: str, configuration_id: str) -> list[JsonDict]: - """ - Retrieves metadata for a given configuration. - - :param component_id: The id of the component. - :param configuration_id: The id of the configuration. - :return: Configuration metadata as a list of dictionaries. Each dictionary contains the 'key' and 'value' keys. - """ - endpoint = f'branch/{self.branch_id}/components/{component_id}/configs/{configuration_id}/metadata' - return cast(list[JsonDict], await self.get(endpoint=endpoint)) - - async def configuration_metadata_update( - self, - component_id: str, - configuration_id: str, - metadata: dict[str, Any], - ) -> list[JsonDict]: - """ - Updates metadata for the given configuration. - - :param component_id: The id of the component. - :param configuration_id: The id of the configuration. - :param metadata: The metadata to update. - :return: Configuration metadata as a list of dictionaries. Each dictionary contains the 'key' and 'value' keys. - """ - endpoint = f'branch/{self.branch_id}/components/{component_id}/configs/{configuration_id}/metadata' - payload = { - 'metadata': [{'key': key, 'value': value} for key, value in metadata.items()], - } - return cast(list[JsonDict], await self.post(endpoint=endpoint, data=payload)) - - async def configuration_update( - self, - component_id: str, - configuration_id: str, - configuration: dict[str, Any], - change_description: str, - updated_name: Optional[str] = None, - updated_description: Optional[str] = None, - is_disabled: bool = False, - ) -> JsonDict: - """ - Updates a component configuration. - - :param component_id: The id of the component. - :param configuration_id: The id of the configuration. - :param configuration: The updated configuration dictionary. - :param change_description: The description of the modification to the configuration. - :param updated_name: The updated name of the configuration, if None, the original - name is preserved. - :param updated_description: The entire description of the updated configuration, if None, the original - description is preserved. - :param is_disabled: Whether the configuration should be disabled. - :return: The SAPI call response - updated configuration or raise an error. - """ - endpoint = f'branch/{self.branch_id}/components/{component_id}/configs/{configuration_id}' - - payload = { - 'configuration': configuration, - 'changeDescription': change_description, - } - if updated_name: - payload['name'] = updated_name - - if updated_description: - payload['description'] = updated_description - - if is_disabled: - payload['isDisabled'] = is_disabled - - return cast(JsonDict, await self.put(endpoint=endpoint, data=payload)) - - async def configuration_row_create( - self, - component_id: str, - config_id: str, - name: str, - description: str, - configuration: dict[str, Any], - ) -> JsonDict: - """ - Creates a new row configuration for a component configuration. - - :param component_id: The ID of the component. - :param config_id: The ID of the configuration. - :param name: The name of the row configuration. - :param description: The description of the row configuration. - :param configuration: The configuration data to create row configuration. - :return: The SAPI call response - created row configuration or raise an error. - """ - payload = { - 'name': name, - 'description': description, - 'configuration': configuration, - } - - return cast( - JsonDict, - await self.post( - endpoint=f'branch/{self.branch_id}/components/{component_id}/configs/{config_id}/rows', - data=payload, - ), - ) - - async def configuration_row_update( - self, - component_id: str, - config_id: str, - configuration_row_id: str, - configuration: dict[str, Any], - change_description: str, - updated_name: Optional[str] = None, - updated_description: Optional[str] = None, - ) -> JsonDict: - """ - Updates a row configuration for a component configuration. - - :param configuration: The configuration data to update row configuration. - :param component_id: The ID of the component. - :param config_id: The ID of the configuration. - :param configuration_row_id: The ID of the row. - :param change_description: The description of the changes made. - :param updated_name: The updated name of the configuration, if None, the original - name is preserved. - :param updated_description: The updated description of the configuration, if None, the original - description is preserved. - :return: The SAPI call response - updated row configuration or raise an error. - """ - - payload = { - 'configuration': configuration, - 'changeDescription': change_description, - } - if updated_name: - payload['name'] = updated_name - - if updated_description: - payload['description'] = updated_description - - return cast( - JsonDict, - await self.put( - endpoint=f'branch/{self.branch_id}/components/{component_id}/configs/{config_id}' - f'/rows/{configuration_row_id}', - data=payload, - ), - ) - - async def configuration_row_detail(self, component_id: str, config_id: str, configuration_row_id: str) -> JsonDict: - """ - Retrieves details of a specific configuration row. - - :param component_id: The id of the component. - :param config_id: The id of the configuration. - :param configuration_row_id: The id of the configuration row. - :return: Configuration row details. - """ - endpoint = f'branch/{self.branch_id}/components/{component_id}/configs/{config_id}/rows/{configuration_row_id}' - return cast(JsonDict, await self.get(endpoint=endpoint)) - - async def configuration_versions(self, component_id: str, config_id: str) -> list[JsonDict]: - """ - Retrieves details of a specific configuration version. - """ - endpoint = f'branch/{self.branch_id}/components/{component_id}/configs/{config_id}/versions' - return cast(list[JsonDict], await self.get(endpoint=endpoint)) - - async def configuration_version_latest(self, component_id: str, config_id: str) -> int: - """ - Retrieves details of the last configuration version. - """ - versions = await self.configuration_versions(component_id, config_id) - latest_version = 0 - for data in versions: - assert isinstance(data, dict) - assert isinstance(data['version'], int) - if latest_version is None or data['version'] > latest_version: - latest_version = data['version'] - return latest_version - - async def job_detail(self, job_id: str | int) -> JsonDict: - """ - NOTE: To get info for regular jobs, use the Job Queue API. - Retrieves information about a given job. - - :param job_id: The id of the job - :return: Job details as dictionary - """ - return cast(JsonDict, await self.get(endpoint=f'jobs/{job_id}')) - - async def global_search( - self, - query: str, - limit: int = 100, - offset: int = 0, - types: Sequence[ItemType] = tuple(), - ) -> 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 - ensure confidentiality. We restrict the search to the project and branch production type of the user. - - :param query: The query to search for. - :param limit: The maximum number of items to return. - :param offset: The offset to start from, pagination parameter. - :param types: The types of items to search for. - """ - params: dict[str, Any] = { - 'query': query, - 'projectIds[]': [await self.project_id()], - 'branchTypes[]': 'production', - 'types[]': types, - 'limit': limit, - 'offset': offset, - } - params = {k: v for k, v in params.items() if v} - 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: - """ - Retrieves information about a given table. - - :param table_id: The id of the table - :return: Table details as dictionary - """ - return cast(JsonDict, await self.get(endpoint=f'tables/{table_id}')) - - async def table_metadata_delete(self, table_id: str, metadata_id: str) -> None: - """ - Deletes metadata for a given table. - - :param table_id: The id of the table - :param metadata_id: The id of the metadata - """ - await self.delete(endpoint=f'tables/{table_id}/metadata/{metadata_id}') - - async def table_metadata_get(self, table_id: str) -> list[JsonDict]: - """ - Retrieves metadata for a given table. - - :param table_id: The id of the table - :return: Table metadata as a list of dictionaries. Each dictionary contains the 'key' and 'value' keys. - """ - return cast(list[JsonDict], await self.get(endpoint=f'tables/{table_id}/metadata')) - - async def table_metadata_update( - self, - table_id: str, - metadata: dict[str, Any] | None = None, - columns_metadata: dict[str, list[dict[str, Any]]] | None = None, - provider: str = 'user', - ) -> JsonDict: - """ - Updates metadata for a given table. At least one of the `metadata` or `columns_metadata` arguments - must be provided. - - :param table_id: The id of the table - :param metadata: The metadata to update. - :param columns_metadata: The column metadata to update. Mapping of column names to a list of dictionaries. - Each dictionary contains the 'key' and 'value' keys. - :param provider: The provider of the metadata ('user' by default). - :return: Dictionary with 'metadata' key under which the table metadata is stored as a list of dictionaries. - Each dictionary contains the 'key' and 'value' keys. Under 'columnsMetadata' key, the column metadata - is stored as a mapping of column names to a list of dictionaries. - """ - if not metadata and not columns_metadata: - raise ValueError('At least one of the `metadata` or `columns_metadata` arguments must be provided.') - - payload: dict[str, Any] = {'provider': provider} - if metadata: - payload['metadata'] = [{'key': key, 'value': value} for key, value in metadata.items()] - if columns_metadata: - payload['columnsMetadata'] = columns_metadata - - return cast(JsonDict, await self.post(endpoint=f'tables/{table_id}/metadata', data=payload)) - - async def trigger_event( - self, - message: str, - component_id: str, - configuration_id: str | None = None, - event_type: StorageEventType | None = None, - params: Mapping[str, Any] | None = None, - results: Mapping[str, Any] | None = None, - duration: float | None = None, - run_id: str | None = None, - ) -> JsonDict: - """ - Sends a Storage API event. - - :param message: The event message. - :param component_id: The ID of the component triggering the event. - :param configuration_id: The ID of the component configuration triggering the event. - :param event_type: The type of event. - :param params: The component parameters. The structure of the params object must follow the JSON schema - registered for the component_id. - :param results: The component results. The structure of the results object must follow the JSON schema - registered for the component_id. - :param duration: The component processing duration in seconds. - :param run_id: The ID of the associated component job. - - :return: Dictionary with the new event ID. - """ - payload: dict[str, Any] = { - 'message': message, - 'component': component_id, - } - if configuration_id: - payload['configurationId'] = configuration_id - if event_type: - payload['type'] = event_type - if params: - payload['params'] = params - if results: - payload['results'] = results - if duration is not None: - # The events API ignores floats, so we round up to the nearest integer. - payload['duration'] = int(math.ceil(duration)) - if run_id: - payload['runId'] = run_id - - LOG.info(f'[trigger_event] payload={payload}') - - return cast(JsonDict, await self.post(endpoint='events', data=payload)) - - async def workspace_create( - self, - login_type: str, - backend: str, - async_run: bool = True, - read_only_storage_access: bool = False, - ) -> JsonDict: - """ - Creates a new workspace. - - :param async_run: If True, the workspace creation is run asynchronously. - :param read_only_storage_access: If True, the workspace has read-only access to the storage. - :return: The SAPI call response - created workspace or raise an error. - """ - return cast( - JsonDict, - await self.post( - endpoint=f'branch/{self.branch_id}/workspaces', - params={'async': async_run}, - data={ - 'readOnlyStorageAccess': read_only_storage_access, - 'loginType': login_type, - 'backend': backend, - }, - ), - ) - - async def workspace_detail(self, workspace_id: int) -> JsonDict: - """ - Retrieves information about a given workspace. - - :param workspace_id: The id of the workspace - :return: Workspace details as dictionary - """ - return cast(JsonDict, await self.get(endpoint=f'branch/{self.branch_id}/workspaces/{workspace_id}')) - - async def workspace_query(self, workspace_id: int, query: str) -> JsonDict: - """ - Executes a query in a given workspace. - - :param workspace_id: The id of the workspace - :param query: The query to execute - :return: The SAPI call response - query result or raise an error. - """ - return cast( - JsonDict, - await self.post( - endpoint=f'branch/{self.branch_id}/workspaces/{workspace_id}/query', - data={'query': query}, - ), - ) - - async def workspace_list(self) -> list[JsonDict]: - """ - Lists all workspaces in the project. - - :return: List of workspaces - """ - return cast(list[JsonDict], await self.get(endpoint=f'branch/{self.branch_id}/workspaces')) - - async def verify_token(self) -> JsonDict: - """ - Checks the token privileges and returns information about the project to which the token belongs. - - :return: Token and project information - """ - return cast(JsonDict, await self.get(endpoint='tokens/verify')) - - async def project_id(self) -> str: - """ - Retrieves the project id. - :return: Project id. - """ - raw_data = cast(JsonDict, await self.get(endpoint='tokens/verify')) - return str(raw_data['owner']['id']) - - async def is_enabled(self, features: ProjectFeature | Iterable[ProjectFeature]) -> bool: - """ - Checks if the features are enabled in the project - conjunction of features. - :param features: The features to check. - :return: True if the features are enabled, False otherwise. - """ - features = [features] if isinstance(features, str) else features - verified_info = await self.verify_token() - project_data = cast(JsonDict, verified_info['owner']) - project_features = cast(list[str], project_data.get('features', [])) - return all(feature in project_features for feature in features) - - async def token_create( - self, - description: str, - component_access: list[str] | None = None, - expires_in: int | None = None, - ) -> JsonDict: - """ - Creates a new Storage API token. - - :param description: Description of the token - :param component_access: List of component IDs the token should have access to - :param expires_in: Token expiration time in seconds - :return: Token creation response containing the token and its details - """ - token_data: dict[str, Any] = {'description': description} - - if component_access: - token_data['componentAccess'] = component_access - - if expires_in: - token_data['expiresIn'] = expires_in - - return cast(JsonDict, await self.post(endpoint='tokens', data=token_data)) - - -class JobsQueueClient(KeboolaServiceClient): - """ - Async client for Keboola Job Queue API. - """ - - @classmethod - def create(cls, root_url: str, token: str, headers: dict[str, Any] | None = None) -> 'JobsQueueClient': - """ - Creates a JobsQueue client. - - :param root_url: Root url of API. e.g. "https://queue.keboola.com/". - :param token: A key for the Storage API. Can be found in the storage console. - :param headers: Additional headers for the requests. - :return: A new instance of JobsQueueClient. - """ - return cls(raw_client=RawKeboolaClient(base_api_url=root_url, api_token=token, headers=headers)) - - async def get_job_detail(self, job_id: str) -> JsonDict: - """ - Retrieves information about a given job. - - :param job_id: The id of the job. - :return: Job details as dictionary. - """ - - return cast(JsonDict, await self.get(endpoint=f'jobs/{job_id}')) - - async def search_jobs_by( - self, - component_id: Optional[str] = None, - config_id: Optional[str] = None, - status: Optional[list[str]] = None, - limit: int = 100, - offset: int = 0, - sort_by: Optional[str] = 'startTime', - sort_order: Optional[str] = 'desc', - ) -> JsonList: - """ - Searches for jobs based on the provided parameters. - - :param component_id: The id of the component. - :param config_id: The id of the configuration. - :param status: The status of the jobs to filter by. - :param limit: The number of jobs to return. - :param offset: The offset of the jobs to return. - :param sort_by: The field to sort the jobs by. - :param sort_order: The order to sort the jobs by. - :return: Dictionary containing matching jobs. - """ - params = { - 'componentId': component_id, - 'configId': config_id, - 'status': status, - 'limit': limit, - 'offset': offset, - 'sortBy': sort_by, - 'sortOrder': sort_order, - } - params = {k: v for k, v in params.items() if v is not None} - return await self._search(params=params) - - async def create_job( - self, - component_id: str, - configuration_id: str, - ) -> JsonDict: - """ - Creates a new job. - - :param component_id: The id of the component. - :param configuration_id: The id of the configuration. - :return: The response from the API call - created job or raise an error. - """ - payload = { - 'component': component_id, - 'config': configuration_id, - 'mode': 'run', - } - return cast(JsonDict, await self.post(endpoint='jobs', data=payload)) - - async def _search(self, params: dict[str, Any]) -> JsonList: - """ - Searches for jobs based on the provided parameters. - - :param params: The parameters to search for. - :return: Dictionary containing matching jobs. - - Parameters (copied from the API docs): - - id str/list[str]: Search jobs by id - - runId str/list[str]: Search jobs by runId - - branchId str/list[str]: Search jobs by branchId - - tokenId str/list[str]: Search jobs by tokenId - - tokenDescription str/list[str]: Search jobs by tokenDescription - - componentId str/list[str]: Search jobs by componentId - - component str/list[str]: Search jobs by componentId, alias for componentId - - configId str/list[str]: Search jobs by configId - - config str/list[str]: Search jobs by configId, alias for configId - - configRowIds str/list[str]: Search jobs by configRowIds - - status str/list[str]: Search jobs by status - - createdTimeFrom str: The jobs that were created after the given date - e.g. "2021-01-01, -8 hours, -1 week,..." - - createdTimeTo str: The jobs that were created before the given date - e.g. "2021-01-01, today, last monday,..." - - startTimeFrom str: The jobs that were started after the given date - e.g. "2021-01-01, -8 hours, -1 week,..." - - startTimeTo str: The jobs that were started before the given date - e.g. "2021-01-01, today, last monday,..." - - endTimeTo str: The jobs that were finished before the given date - e.g. "2021-01-01, today, last monday,..." - - endTimeFrom str: The jobs that were finished after the given date - e.g. "2021-01-01, -8 hours, -1 week,..." - - limit int: The number of jobs returned, default 100 - - offset int: The jobs page offset, default 0 - - sortBy str: The jobs sorting field, default "id" - values: id, runId, projectId, branchId, componentId, configId, tokenDescription, status, createdTime, - updatedTime, startTime, endTime, durationSeconds - - sortOrder str: The jobs sorting order, default "desc" - values: asc, desc - """ - return cast(JsonList, await self.get(endpoint='search/jobs', params=params)) - - -class DocsQuestionResponse(BaseModel): - """ - The AI service response to a request to `/docs/question` endpoint. - """ - - text: str = Field(description='Text of the answer to a documentation query.') - source_urls: list[str] = Field( - description='List of URLs to the sources of the answer.', - default_factory=list, - alias='sourceUrls', - ) - - -class SuggestedComponent(BaseModel): - """The AI service response to a /docs/suggest-component request.""" - - component_id: str = Field(description='The component ID.', alias='componentId') - score: float = Field(description='Score of the component suggestion.') - source: str = Field(description='Source of the component suggestion.') - - -class ComponentSuggestionResponse(BaseModel): - """The AI service response to a /suggest/component request.""" - - components: list[SuggestedComponent] = Field(description='List of suggested components.', default_factory=list) - - -class APIFlowResponse(BaseModel): - """ - Raw API response for configuration endpoints. - - Note: will be removed soon due to removal of flow specific client methods. - """ - - # Core identification fields - configuration_id: str = Field( - description='The ID of the flow configuration', - validation_alias=AliasChoices('id', 'configuration_id', 'configurationId', 'configuration-id'), - serialization_alias='id', - ) - name: str = Field(description='The name of the flow configuration') - description: Optional[str] = Field(default=None, description='The description of the flow configuration') - - # Versioning and state - version: int = Field(description='The version of the flow configuration') - is_disabled: bool = Field( - default=False, - description='Whether the flow configuration is disabled', - validation_alias=AliasChoices('isDisabled', 'is_disabled', 'is-disabled'), - serialization_alias='isDisabled', - ) - is_deleted: bool = Field( - default=False, - description='Whether the flow configuration is deleted', - validation_alias=AliasChoices('isDeleted', 'is_deleted', 'is-deleted'), - serialization_alias='isDeleted', - ) - - # Flow-specific configuration data (as returned by API) - configuration: dict[str, Any] = Field( - description='The nested flow configuration object containing phases and tasks' - ) - - # Change tracking - change_description: Optional[str] = Field( - default=None, - description='The description of the latest changes', - validation_alias=AliasChoices('changeDescription', 'change_description', 'change-description'), - serialization_alias='changeDescription', - ) - - # Metadata - metadata: list[dict[str, Any]] = Field( - default_factory=list, - description='Flow configuration metadata', - validation_alias=AliasChoices('metadata', 'configuration_metadata', 'configurationMetadata'), - ) - - # Timestamps - created: Optional[str] = Field(None, description='Creation timestamp') - updated: Optional[str] = Field(None, description='Last update timestamp') - - -class ComponentAPIResponse(BaseModel): - """ - Raw component response that can handle both Storage API and AI Service API responses. - - Storage API (/v2/storage/components/{id}) returns just the core fields. - AI Service API (/docs/components/{id}) returns core fields + optional documentation metadata. - - The optional fields will be None when parsing Storage API responses. - """ - - # Core fields present in both APIs (SAPI and AI service) - component_id: str = Field( - description='The ID of the component', - validation_alias=AliasChoices('component_id', 'id', 'componentId', 'component-id'), - ) - component_name: str = Field( - description='The name of the component', - validation_alias=AliasChoices( - 'name', - 'component_name', - 'componentName', - 'component-name', - ), - ) - type: str = Field( - description='Component type (extractor, writer, application)', - validation_alias=AliasChoices('type', 'component_type', 'componentType', 'component-type'), - ) - flags: list[str] = Field( - default_factory=list, - description='Developer portal flags', - validation_alias=AliasChoices('flags', 'component_flags', 'componentFlags', 'component-flags'), - ) - categories: list[str] = Field( - default_factory=list, - description='Component categories', - validation_alias=AliasChoices( - 'categories', - 'component_categories', - 'componentCategories', - 'component-categories', - ), - ) - - # Optional metadata fields only present in AI Service API responses - documentation_url: str | None = Field( - default=None, - description='Documentation URL', - validation_alias=AliasChoices('documentationUrl', 'documentation_url', 'documentation-url'), - ) - documentation: str | None = Field( - default=None, - description='Component documentation', - validation_alias=AliasChoices('documentation'), - ) - configuration_schema: dict[str, Any] | None = Field( - default=None, - description='Configuration schema', - validation_alias=AliasChoices('configurationSchema', 'configuration_schema', 'configuration-schema'), - ) - configuration_row_schema: dict[str, Any] | None = Field( - default=None, - description='Configuration row schema', - validation_alias=AliasChoices('configurationRowSchema', 'configuration_row_schema', 'configuration-row-schema'), - ) - - -class ConfigurationAPIResponse(BaseModel): - """ - Raw API response for configuration endpoints. - - Mirrors the actual JSON structure returned by Keboola Storage API for: - - configuration_detail() - - configuration_list() - - configuration_create() - - configuration_update() - """ - - component_id: str = Field( - description='The ID of the component', - validation_alias=AliasChoices('component_id', 'componentId', 'component-id'), - ) - configuration_id: str = Field( - description='The ID of the configuration', - validation_alias=AliasChoices('configuration_id', 'id', 'configurationId', 'configuration-id'), - ) - name: str = Field(description='The name of the configuration') - description: Optional[str] = Field(default=None, description='The description of the configuration') - version: int = Field(description='The version of the configuration') - is_disabled: bool = Field( - default=False, - description='Whether the configuration is disabled', - validation_alias=AliasChoices('isDisabled', 'is_disabled', 'is-disabled'), - ) - is_deleted: bool = Field( - default=False, - description='Whether the configuration is deleted', - validation_alias=AliasChoices('isDeleted', 'is_deleted', 'is-deleted'), - ) - configuration: dict[str, Any] = Field( - description='The nested configuration object containing parameters and storage' - ) - rows: Optional[list[dict[str, Any]]] = Field( - default=None, description='The row configurations within this configuration' - ) - change_description: Optional[str] = Field( - default=None, - description='The description of the latest changes', - validation_alias=AliasChoices('changeDescription', 'change_description', 'change-description'), - ) - metadata: list[dict[str, Any]] = Field( - default_factory=list, - description='Configuration metadata', - validation_alias=AliasChoices('metadata', 'configuration_metadata', 'configurationMetadata'), - ) - - -class CreateConfigurationAPIResponse(BaseModel): - id: str = Field(description='Unique identifier of the newly created configuration.') - name: str = Field(description='Human-readable name of the configuration.') - description: Optional[str] = Field(default='', description='Optional description of the configuration.') - created: datetime = Field(description='Timestamp when the configuration was created (ISO 8601).') - creator_token: dict[str, Any] = Field( - description='Metadata about the token that created the configuration.', alias='creatorToken' - ) - version: int = Field(description='Version number of the configuration.') - change_description: Optional[str] = Field( - description='Optional description of the change that introduced this configuration version.', - alias='changeDescription', - ) - is_disabled: bool = Field( - description='Indicates whether the configuration is currently disabled.', alias='isDisabled' - ) - is_deleted: bool = Field( - description='Indicates whether the configuration has been marked as deleted.', alias='isDeleted' - ) - configuration: Optional[dict[str, Any]] = Field( - description='User-defined configuration payload (key-value structure).' - ) - state: Optional[dict[str, Any]] = Field( - description='Internal runtime state data associated with the configuration.' - ) - current_version: Optional[dict[str, Any]] = Field( - description='Metadata about the currently deployed version of the configuration.', alias='currentVersion' - ) - - -class AIServiceClient(KeboolaServiceClient): - """Async client for Keboola AI Service.""" - - @classmethod - def create(cls, root_url: str, token: str, headers: dict[str, Any] | None = None) -> 'AIServiceClient': - """ - Creates an AIServiceClient from a Keboola Storage API token. - - :param root_url: The root URL of the AI service API. - :param token: The Keboola Storage API token. - :param headers: Additional headers for the requests. - :return: A new instance of AIServiceClient. - """ - return cls(raw_client=RawKeboolaClient(base_api_url=root_url, api_token=token, headers=headers)) - - async def get_component_detail(self, component_id: str) -> JsonDict: - """ - Retrieves information about a given component. - - :param component_id: The id of the component. - :return: Component details as dictionary. - """ - return cast(JsonDict, await self.get(endpoint=f'docs/components/{component_id}')) - - async def docs_question(self, query: str) -> DocsQuestionResponse: - """ - Answers a question using the Keboola documentation as a source. - :param query: The query to answer. - :return: Response containing the answer and source URLs. - """ - response = await self.raw_client.post( - endpoint='docs/question', - data={'query': query}, - headers={'Accept': 'application/json'}, - ) - - return DocsQuestionResponse.model_validate(response) - - async def suggest_component(self, query: str) -> ComponentSuggestionResponse: - """ - Provides list of component suggestions based on natural language query. - :param query: The query to answer. - :return: Response containing the list of suggested component IDs, their score and source. - """ - response = await self.raw_client.post( - endpoint='suggest/component', - data={'prompt': query}, - headers={'Accept': 'application/json'}, - ) - - return ComponentSuggestionResponse.model_validate(response) diff --git a/src/keboola_mcp_server/clients/jobs_queue.py b/src/keboola_mcp_server/clients/jobs_queue.py new file mode 100644 index 00000000..83c80eff --- /dev/null +++ b/src/keboola_mcp_server/clients/jobs_queue.py @@ -0,0 +1,125 @@ +from typing import Any, Optional, cast + +from keboola_mcp_server.clients.base import JsonDict, JsonList, KeboolaServiceClient, RawKeboolaClient + + +class JobsQueueClient(KeboolaServiceClient): + """ + Async client for Keboola Job Queue API. + """ + + @classmethod + def create(cls, root_url: str, token: str, headers: dict[str, Any] | None = None) -> 'JobsQueueClient': + """ + Creates a JobsQueue client. + + :param root_url: Root url of API. e.g. "https://queue.keboola.com/". + :param token: A key for the Storage API. Can be found in the storage console. + :param headers: Additional headers for the requests. + :return: A new instance of JobsQueueClient. + """ + return cls(raw_client=RawKeboolaClient(base_api_url=root_url, api_token=token, headers=headers)) + + async def get_job_detail(self, job_id: str) -> JsonDict: + """ + Retrieves information about a given job. + + :param job_id: The id of the job. + :return: Job details as dictionary. + """ + + return cast(JsonDict, await self.get(endpoint=f'jobs/{job_id}')) + + async def search_jobs_by( + self, + component_id: Optional[str] = None, + config_id: Optional[str] = None, + status: Optional[list[str]] = None, + limit: int = 100, + offset: int = 0, + sort_by: Optional[str] = 'startTime', + sort_order: Optional[str] = 'desc', + ) -> JsonList: + """ + Searches for jobs based on the provided parameters. + + :param component_id: The id of the component. + :param config_id: The id of the configuration. + :param status: The status of the jobs to filter by. + :param limit: The number of jobs to return. + :param offset: The offset of the jobs to return. + :param sort_by: The field to sort the jobs by. + :param sort_order: The order to sort the jobs by. + :return: Dictionary containing matching jobs. + """ + params = { + 'componentId': component_id, + 'configId': config_id, + 'status': status, + 'limit': limit, + 'offset': offset, + 'sortBy': sort_by, + 'sortOrder': sort_order, + } + params = {k: v for k, v in params.items() if v is not None} + return await self._search(params=params) + + async def create_job( + self, + component_id: str, + configuration_id: str, + ) -> JsonDict: + """ + Creates a new job. + + :param component_id: The id of the component. + :param configuration_id: The id of the configuration. + :return: The response from the API call - created job or raise an error. + """ + payload = { + 'component': component_id, + 'config': configuration_id, + 'mode': 'run', + } + return cast(JsonDict, await self.post(endpoint='jobs', data=payload)) + + async def _search(self, params: dict[str, Any]) -> JsonList: + """ + Searches for jobs based on the provided parameters. + + :param params: The parameters to search for. + :return: Dictionary containing matching jobs. + + Parameters (copied from the API docs): + - id str/list[str]: Search jobs by id + - runId str/list[str]: Search jobs by runId + - branchId str/list[str]: Search jobs by branchId + - tokenId str/list[str]: Search jobs by tokenId + - tokenDescription str/list[str]: Search jobs by tokenDescription + - componentId str/list[str]: Search jobs by componentId + - component str/list[str]: Search jobs by componentId, alias for componentId + - configId str/list[str]: Search jobs by configId + - config str/list[str]: Search jobs by configId, alias for configId + - configRowIds str/list[str]: Search jobs by configRowIds + - status str/list[str]: Search jobs by status + - createdTimeFrom str: The jobs that were created after the given date + e.g. "2021-01-01, -8 hours, -1 week,..." + - createdTimeTo str: The jobs that were created before the given date + e.g. "2021-01-01, today, last monday,..." + - startTimeFrom str: The jobs that were started after the given date + e.g. "2021-01-01, -8 hours, -1 week,..." + - startTimeTo str: The jobs that were started before the given date + e.g. "2021-01-01, today, last monday,..." + - endTimeTo str: The jobs that were finished before the given date + e.g. "2021-01-01, today, last monday,..." + - endTimeFrom str: The jobs that were finished after the given date + e.g. "2021-01-01, -8 hours, -1 week,..." + - limit int: The number of jobs returned, default 100 + - offset int: The jobs page offset, default 0 + - sortBy str: The jobs sorting field, default "id" + values: id, runId, projectId, branchId, componentId, configId, tokenDescription, status, createdTime, + updatedTime, startTime, endTime, durationSeconds + - sortOrder str: The jobs sorting order, default "desc" + values: asc, desc + """ + return cast(JsonList, await self.get(endpoint='search/jobs', params=params)) diff --git a/src/keboola_mcp_server/clients/storage.py b/src/keboola_mcp_server/clients/storage.py new file mode 100644 index 00000000..df1796d8 --- /dev/null +++ b/src/keboola_mcp_server/clients/storage.py @@ -0,0 +1,952 @@ +import logging +import math +from datetime import datetime +from typing import Any, Iterable, Literal, Mapping, Optional, Sequence, cast + +from pydantic import AliasChoices, BaseModel, Field, field_validator + +from keboola_mcp_server.clients.base import JsonDict, KeboolaServiceClient, RawKeboolaClient + +LOG = logging.getLogger(__name__) + + +ComponentResource = Literal['configuration', 'rows', 'state'] +StorageEventType = Literal['info', 'success', 'warn', 'error'] + +# Project features that can be checked with the is_enabled method +ProjectFeature = Literal['global-search'] + +ItemType = Literal[ + 'flow', + 'bucket', + 'table', + 'transformation', + 'configuration', + 'configuration-row', + 'workspace', + 'shared-code', + 'rows', + 'state', +] + + +class GlobalSearchResponse(BaseModel): + """The SAPI global search response.""" + + class Item(BaseModel): + id: str = Field(description='The id of the item.') + name: str = Field(description='The name of the item.') + type: ItemType = 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[Item] = Field(description='List of search results 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 APIFlowResponse(BaseModel): + """ + Raw API response for configuration endpoints. + + Note: will be removed soon due to removal of flow specific client methods. + """ + + # Core identification fields + configuration_id: str = Field( + description='The ID of the flow configuration', + validation_alias=AliasChoices('id', 'configuration_id', 'configurationId', 'configuration-id'), + serialization_alias='id', + ) + name: str = Field(description='The name of the flow configuration') + description: Optional[str] = Field(default=None, description='The description of the flow configuration') + + # Versioning and state + version: int = Field(description='The version of the flow configuration') + is_disabled: bool = Field( + default=False, + description='Whether the flow configuration is disabled', + validation_alias=AliasChoices('isDisabled', 'is_disabled', 'is-disabled'), + serialization_alias='isDisabled', + ) + is_deleted: bool = Field( + default=False, + description='Whether the flow configuration is deleted', + validation_alias=AliasChoices('isDeleted', 'is_deleted', 'is-deleted'), + serialization_alias='isDeleted', + ) + + # Flow-specific configuration data (as returned by API) + configuration: dict[str, Any] = Field( + description='The nested flow configuration object containing phases and tasks' + ) + + # Change tracking + change_description: Optional[str] = Field( + default=None, + description='The description of the latest changes', + validation_alias=AliasChoices('changeDescription', 'change_description', 'change-description'), + serialization_alias='changeDescription', + ) + + # Metadata + metadata: list[dict[str, Any]] = Field( + default_factory=list, + description='Flow configuration metadata', + validation_alias=AliasChoices('metadata', 'configuration_metadata', 'configurationMetadata'), + ) + + # Timestamps + created: Optional[str] = Field(None, description='Creation timestamp') + updated: Optional[str] = Field(None, description='Last update timestamp') + + +class ComponentAPIResponse(BaseModel): + """ + Raw component response that can handle both Storage API and AI Service API responses. + + Storage API (/v2/storage/components/{id}) returns just the core fields. + AI Service API (/docs/components/{id}) returns core fields + optional documentation metadata. + + The optional fields will be None when parsing Storage API responses. + """ + + # Core fields present in both APIs (SAPI and AI service) + component_id: str = Field( + description='The ID of the component', + validation_alias=AliasChoices('component_id', 'id', 'componentId', 'component-id'), + ) + component_name: str = Field( + description='The name of the component', + validation_alias=AliasChoices( + 'name', + 'component_name', + 'componentName', + 'component-name', + ), + ) + type: str = Field( + description='Component type (extractor, writer, application)', + validation_alias=AliasChoices('type', 'component_type', 'componentType', 'component-type'), + ) + flags: list[str] = Field( + default_factory=list, + description='Developer portal flags', + validation_alias=AliasChoices('flags', 'component_flags', 'componentFlags', 'component-flags'), + ) + categories: list[str] = Field( + default_factory=list, + description='Component categories', + validation_alias=AliasChoices( + 'categories', + 'component_categories', + 'componentCategories', + 'component-categories', + ), + ) + + # Optional metadata fields only present in AI Service API responses + documentation_url: str | None = Field( + default=None, + description='Documentation URL', + validation_alias=AliasChoices('documentationUrl', 'documentation_url', 'documentation-url'), + ) + documentation: str | None = Field( + default=None, + description='Component documentation', + validation_alias=AliasChoices('documentation'), + ) + configuration_schema: dict[str, Any] | None = Field( + default=None, + description='Configuration schema', + validation_alias=AliasChoices('configurationSchema', 'configuration_schema', 'configuration-schema'), + ) + configuration_row_schema: dict[str, Any] | None = Field( + default=None, + description='Configuration row schema', + validation_alias=AliasChoices('configurationRowSchema', 'configuration_row_schema', 'configuration-row-schema'), + ) + + +class ConfigurationAPIResponse(BaseModel): + """ + Raw API response for configuration endpoints. + + Mirrors the actual JSON structure returned by Keboola Storage API for: + - configuration_detail() + - configuration_list() + - configuration_create() + - configuration_update() + """ + + component_id: str = Field( + description='The ID of the component', + validation_alias=AliasChoices('component_id', 'componentId', 'component-id'), + ) + configuration_id: str = Field( + description='The ID of the configuration', + validation_alias=AliasChoices('configuration_id', 'id', 'configurationId', 'configuration-id'), + ) + name: str = Field(description='The name of the configuration') + description: Optional[str] = Field(default=None, description='The description of the configuration') + version: int = Field(description='The version of the configuration') + is_disabled: bool = Field( + default=False, + description='Whether the configuration is disabled', + validation_alias=AliasChoices('isDisabled', 'is_disabled', 'is-disabled'), + ) + is_deleted: bool = Field( + default=False, + description='Whether the configuration is deleted', + validation_alias=AliasChoices('isDeleted', 'is_deleted', 'is-deleted'), + ) + configuration: dict[str, Any] = Field( + description='The nested configuration object containing parameters and storage' + ) + rows: Optional[list[dict[str, Any]]] = Field( + default=None, description='The row configurations within this configuration' + ) + change_description: Optional[str] = Field( + default=None, + description='The description of the latest changes', + validation_alias=AliasChoices('changeDescription', 'change_description', 'change-description'), + ) + metadata: list[dict[str, Any]] = Field( + default_factory=list, + description='Configuration metadata', + validation_alias=AliasChoices('metadata', 'configuration_metadata', 'configurationMetadata'), + ) + + +class CreateConfigurationAPIResponse(BaseModel): + id: str = Field(description='Unique identifier of the newly created configuration.') + name: str = Field(description='Human-readable name of the configuration.') + description: Optional[str] = Field(default='', description='Optional description of the configuration.') + created: datetime = Field(description='Timestamp when the configuration was created (ISO 8601).') + creator_token: dict[str, Any] = Field( + description='Metadata about the token that created the configuration.', alias='creatorToken' + ) + version: int = Field(description='Version number of the configuration.') + change_description: Optional[str] = Field( + description='Optional description of the change that introduced this configuration version.', + alias='changeDescription', + ) + is_disabled: bool = Field( + description='Indicates whether the configuration is currently disabled.', alias='isDisabled' + ) + is_deleted: bool = Field( + description='Indicates whether the configuration has been marked as deleted.', alias='isDeleted' + ) + configuration: Optional[dict[str, Any]] = Field( + description='User-defined configuration payload (key-value structure).' + ) + state: Optional[dict[str, Any]] = Field( + description='Internal runtime state data associated with the configuration.' + ) + current_version: Optional[dict[str, Any]] = Field( + description='Metadata about the currently deployed version of the configuration.', alias='currentVersion' + ) + + +class AsyncStorageClient(KeboolaServiceClient): + + def __init__(self, raw_client: RawKeboolaClient, branch_id: str = 'default') -> None: + """ + Creates an AsyncStorageClient from a RawKeboolaClient and a branch id. + + :param raw_client: The raw client to use + :param branch_id: The id of the branch + """ + super().__init__(raw_client=raw_client) + self._branch_id: str = branch_id + + @property + def branch_id(self) -> str: + return self._branch_id + + @property + def base_api_url(self) -> str: + return self.raw_client.base_api_url.split('/v2')[0] + + @classmethod + def create( + cls, + root_url: str, + token: str, + version: str = 'v2', + branch_id: str = 'default', + headers: dict[str, Any] | None = None, + ) -> 'AsyncStorageClient': + """ + Creates an AsyncStorageClient from a Keboola Storage API token. + + :param root_url: The root URL of the service API + :param token: The Keboola Storage API token + :param version: The version of the API to use (default: 'v2') + :param branch_id: The id of the branch + :param headers: Additional headers for the requests + :return: A new instance of AsyncStorageClient + """ + return cls( + raw_client=RawKeboolaClient( + base_api_url=f'{root_url}/{version}/storage', + api_token=token, + headers=headers, + ), + branch_id=branch_id, + ) + + async def branch_metadata_get(self) -> list[JsonDict]: + """ + Retrieves metadata for the current branch. + + :return: Branch metadata as a list of dictionaries. Each dictionary contains the 'key' and 'value' keys. + """ + return cast(list[JsonDict], await self.get(endpoint=f'branch/{self.branch_id}/metadata')) + + async def branch_metadata_update(self, metadata: dict[str, Any]) -> list[JsonDict]: + """ + Updates metadata for the current branch. + + :param metadata: The metadata to update. + :return: The SAPI call response - updated metadata or raise an error. + """ + payload = { + 'metadata': [{'key': key, 'value': value} for key, value in metadata.items()], + } + return cast(list[JsonDict], await self.post(endpoint=f'branch/{self.branch_id}/metadata', data=payload)) + + async def bucket_detail(self, bucket_id: str) -> JsonDict: + """ + Retrieves information about a given bucket. + + :param bucket_id: The id of the bucket + :return: Bucket details as dictionary + """ + return cast(JsonDict, await self.get(endpoint=f'buckets/{bucket_id}')) + + async def bucket_list(self, include: list[str] | None = None) -> list[JsonDict]: + """ + Lists all buckets. + + :param include: List of fields to include in the response ('metadata' or 'linkedBuckets') + :return: List of buckets as dictionary + """ + params = {} + if include is not None and isinstance(include, list): + params['include'] = ','.join(include) + return cast(list[JsonDict], await self.get(endpoint='buckets', params=params)) + + async def bucket_metadata_delete(self, bucket_id: str, metadata_id: str) -> None: + """ + Deletes metadata for a given bucket. + + :param bucket_id: The id of the bucket + :param metadata_id: The id of the metadata + """ + await self.delete(endpoint=f'buckets/{bucket_id}/metadata/{metadata_id}') + + async def bucket_metadata_get(self, bucket_id: str) -> list[JsonDict]: + """ + Retrieves metadata for a given bucket. + + :param bucket_id: The id of the bucket + :return: Bucket metadata as a list of dictionaries. Each dictionary contains the 'key' and 'value' keys. + """ + return cast(list[JsonDict], await self.get(endpoint=f'buckets/{bucket_id}/metadata')) + + async def bucket_metadata_update( + self, + bucket_id: str, + metadata: dict[str, Any], + provider: str = 'user', + ) -> list[JsonDict]: + """ + Updates metadata for a given bucket. + + :param bucket_id: The id of the bucket + :param metadata: The metadata to update. + :param provider: The provider of the metadata ('user' by default). + :return: Bucket metadata as a list of dictionaries. Each dictionary contains the 'key' and 'value' keys. + """ + payload = { + 'provider': provider, + 'metadata': [{'key': key, 'value': value} for key, value in metadata.items()], + } + return cast(list[JsonDict], await self.post(endpoint=f'buckets/{bucket_id}/metadata', data=payload)) + + async def bucket_table_list(self, bucket_id: str, include: list[str] | None = None) -> list[JsonDict]: + """ + Lists all tables in a given bucket. + + :param bucket_id: The id of the bucket + :param include: List of fields to include in the response + :return: List of tables as dictionary + """ + params = {} + if include is not None and isinstance(include, list): + params['include'] = ','.join(include) + return cast(list[JsonDict], await self.get(endpoint=f'buckets/{bucket_id}/tables', params=params)) + + async def component_detail(self, component_id: str) -> JsonDict: + """ + Retrieves information about a given component. + + :param component_id: The id of the component + :return: Component details as a dictionary + """ + return cast(JsonDict, await self.get(endpoint=f'branch/{self.branch_id}/components/{component_id}')) + + async def component_list( + self, component_type: str, include: list[ComponentResource] | None = None + ) -> list[JsonDict]: + """ + Lists all components of a given type. + + :param component_type: The type of the component (extractor, writer, application, etc.) + :param include: Comma separated list of resources to include. + Available resources: configuration, rows and state. + :return: List of components as dictionary + """ + endpoint = f'branch/{self.branch_id}/components' + params = {'componentType': component_type} + if include is not None and isinstance(include, list): + params['include'] = ','.join(include) + + return cast(list[JsonDict], await self.get(endpoint=endpoint, params=params)) + + async def configuration_create( + self, + component_id: str, + name: str, + description: str, + configuration: dict[str, Any], + ) -> JsonDict: + """ + Creates a new configuration for a component. + + :param component_id: The id of the component for which to create the configuration. + :param name: The name of the configuration. + :param description: The description of the configuration. + :param configuration: The configuration definition as a dictionary. + + :return: The SAPI call response - created configuration or raise an error. + """ + endpoint = f'branch/{self.branch_id}/components/{component_id}/configs' + + payload = { + 'name': name, + 'description': description, + 'configuration': configuration, + } + return cast(JsonDict, await self.post(endpoint=endpoint, data=payload)) + + async def configuration_delete(self, component_id: str, configuration_id: str, skip_trash: bool = False) -> None: + """ + Deletes a configuration. + + :param component_id: The id of the component. + :param configuration_id: The id of the configuration. + :param skip_trash: If True, the configuration is deleted without moving to the trash. + (Technically it means the API endpoint is called twice.) + :raises httpx.HTTPStatusError: If the (component_id, configuration_id) is not found. + """ + endpoint = f'branch/{self.branch_id}/components/{component_id}/configs/{configuration_id}' + await self.delete(endpoint=endpoint) + if skip_trash: + await self.delete(endpoint=endpoint) + + async def configuration_detail(self, component_id: str, configuration_id: str) -> JsonDict: + """ + Retrieves information about a given configuration. + + :param component_id: The id of the component. + :param configuration_id: The id of the configuration. + :return: The parsed json from the HTTP response. + :raises ValueError: If the component_id or configuration_id is invalid. + """ + if not isinstance(component_id, str) or component_id == '': + raise ValueError(f"Invalid component_id '{component_id}'.") + if not isinstance(configuration_id, str) or configuration_id == '': + raise ValueError(f"Invalid configuration_id '{configuration_id}'.") + endpoint = f'branch/{self.branch_id}/components/{component_id}/configs/{configuration_id}' + + return cast(JsonDict, await self.get(endpoint=endpoint)) + + async def configuration_list(self, component_id: str) -> list[JsonDict]: + """ + Lists configurations of the given component. + + :param component_id: The id of the component. + :return: List of configurations. + :raises ValueError: If the component_id is invalid. + """ + if not isinstance(component_id, str) or component_id == '': + raise ValueError(f"Invalid component_id '{component_id}'.") + endpoint = f'branch/{self.branch_id}/components/{component_id}/configs' + + return cast(list[JsonDict], await self.get(endpoint=endpoint)) + + async def configuration_metadata_get(self, component_id: str, configuration_id: str) -> list[JsonDict]: + """ + Retrieves metadata for a given configuration. + + :param component_id: The id of the component. + :param configuration_id: The id of the configuration. + :return: Configuration metadata as a list of dictionaries. Each dictionary contains the 'key' and 'value' keys. + """ + endpoint = f'branch/{self.branch_id}/components/{component_id}/configs/{configuration_id}/metadata' + return cast(list[JsonDict], await self.get(endpoint=endpoint)) + + async def configuration_metadata_update( + self, + component_id: str, + configuration_id: str, + metadata: dict[str, Any], + ) -> list[JsonDict]: + """ + Updates metadata for the given configuration. + + :param component_id: The id of the component. + :param configuration_id: The id of the configuration. + :param metadata: The metadata to update. + :return: Configuration metadata as a list of dictionaries. Each dictionary contains the 'key' and 'value' keys. + """ + endpoint = f'branch/{self.branch_id}/components/{component_id}/configs/{configuration_id}/metadata' + payload = { + 'metadata': [{'key': key, 'value': value} for key, value in metadata.items()], + } + return cast(list[JsonDict], await self.post(endpoint=endpoint, data=payload)) + + async def configuration_update( + self, + component_id: str, + configuration_id: str, + configuration: dict[str, Any], + change_description: str, + updated_name: Optional[str] = None, + updated_description: Optional[str] = None, + is_disabled: bool = False, + ) -> JsonDict: + """ + Updates a component configuration. + + :param component_id: The id of the component. + :param configuration_id: The id of the configuration. + :param configuration: The updated configuration dictionary. + :param change_description: The description of the modification to the configuration. + :param updated_name: The updated name of the configuration, if None, the original + name is preserved. + :param updated_description: The entire description of the updated configuration, if None, the original + description is preserved. + :param is_disabled: Whether the configuration should be disabled. + :return: The SAPI call response - updated configuration or raise an error. + """ + endpoint = f'branch/{self.branch_id}/components/{component_id}/configs/{configuration_id}' + + payload = { + 'configuration': configuration, + 'changeDescription': change_description, + } + if updated_name: + payload['name'] = updated_name + + if updated_description: + payload['description'] = updated_description + + if is_disabled: + payload['isDisabled'] = is_disabled + + return cast(JsonDict, await self.put(endpoint=endpoint, data=payload)) + + async def configuration_row_create( + self, + component_id: str, + config_id: str, + name: str, + description: str, + configuration: dict[str, Any], + ) -> JsonDict: + """ + Creates a new row configuration for a component configuration. + + :param component_id: The ID of the component. + :param config_id: The ID of the configuration. + :param name: The name of the row configuration. + :param description: The description of the row configuration. + :param configuration: The configuration data to create row configuration. + :return: The SAPI call response - created row configuration or raise an error. + """ + payload = { + 'name': name, + 'description': description, + 'configuration': configuration, + } + + return cast( + JsonDict, + await self.post( + endpoint=f'branch/{self.branch_id}/components/{component_id}/configs/{config_id}/rows', + data=payload, + ), + ) + + async def configuration_row_update( + self, + component_id: str, + config_id: str, + configuration_row_id: str, + configuration: dict[str, Any], + change_description: str, + updated_name: Optional[str] = None, + updated_description: Optional[str] = None, + ) -> JsonDict: + """ + Updates a row configuration for a component configuration. + + :param configuration: The configuration data to update row configuration. + :param component_id: The ID of the component. + :param config_id: The ID of the configuration. + :param configuration_row_id: The ID of the row. + :param change_description: The description of the changes made. + :param updated_name: The updated name of the configuration, if None, the original + name is preserved. + :param updated_description: The updated description of the configuration, if None, the original + description is preserved. + :return: The SAPI call response - updated row configuration or raise an error. + """ + + payload = { + 'configuration': configuration, + 'changeDescription': change_description, + } + if updated_name: + payload['name'] = updated_name + + if updated_description: + payload['description'] = updated_description + + return cast( + JsonDict, + await self.put( + endpoint=f'branch/{self.branch_id}/components/{component_id}/configs/{config_id}' + f'/rows/{configuration_row_id}', + data=payload, + ), + ) + + async def configuration_row_detail(self, component_id: str, config_id: str, configuration_row_id: str) -> JsonDict: + """ + Retrieves details of a specific configuration row. + + :param component_id: The id of the component. + :param config_id: The id of the configuration. + :param configuration_row_id: The id of the configuration row. + :return: Configuration row details. + """ + endpoint = f'branch/{self.branch_id}/components/{component_id}/configs/{config_id}/rows/{configuration_row_id}' + return cast(JsonDict, await self.get(endpoint=endpoint)) + + async def configuration_versions(self, component_id: str, config_id: str) -> list[JsonDict]: + """ + Retrieves details of a specific configuration version. + """ + endpoint = f'branch/{self.branch_id}/components/{component_id}/configs/{config_id}/versions' + return cast(list[JsonDict], await self.get(endpoint=endpoint)) + + async def configuration_version_latest(self, component_id: str, config_id: str) -> int: + """ + Retrieves details of the last configuration version. + """ + versions = await self.configuration_versions(component_id, config_id) + latest_version = 0 + for data in versions: + assert isinstance(data, dict) + assert isinstance(data['version'], int) + if latest_version is None or data['version'] > latest_version: + latest_version = data['version'] + return latest_version + + async def job_detail(self, job_id: str | int) -> JsonDict: + """ + NOTE: To get info for regular jobs, use the Job Queue API. + Retrieves information about a given job. + + :param job_id: The id of the job + :return: Job details as dictionary + """ + return cast(JsonDict, await self.get(endpoint=f'jobs/{job_id}')) + + async def global_search( + self, + query: str, + limit: int = 100, + offset: int = 0, + types: Sequence[ItemType] = tuple(), + ) -> 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 + ensure confidentiality. We restrict the search to the project and branch production type of the user. + + :param query: The query to search for. + :param limit: The maximum number of items to return. + :param offset: The offset to start from, pagination parameter. + :param types: The types of items to search for. + """ + params: dict[str, Any] = { + 'query': query, + 'projectIds[]': [await self.project_id()], + 'branchTypes[]': 'production', + 'types[]': types, + 'limit': limit, + 'offset': offset, + } + params = {k: v for k, v in params.items() if v} + 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: + """ + Retrieves information about a given table. + + :param table_id: The id of the table + :return: Table details as dictionary + """ + return cast(JsonDict, await self.get(endpoint=f'tables/{table_id}')) + + async def table_metadata_delete(self, table_id: str, metadata_id: str) -> None: + """ + Deletes metadata for a given table. + + :param table_id: The id of the table + :param metadata_id: The id of the metadata + """ + await self.delete(endpoint=f'tables/{table_id}/metadata/{metadata_id}') + + async def table_metadata_get(self, table_id: str) -> list[JsonDict]: + """ + Retrieves metadata for a given table. + + :param table_id: The id of the table + :return: Table metadata as a list of dictionaries. Each dictionary contains the 'key' and 'value' keys. + """ + return cast(list[JsonDict], await self.get(endpoint=f'tables/{table_id}/metadata')) + + async def table_metadata_update( + self, + table_id: str, + metadata: dict[str, Any] | None = None, + columns_metadata: dict[str, list[dict[str, Any]]] | None = None, + provider: str = 'user', + ) -> JsonDict: + """ + Updates metadata for a given table. At least one of the `metadata` or `columns_metadata` arguments + must be provided. + + :param table_id: The id of the table + :param metadata: The metadata to update. + :param columns_metadata: The column metadata to update. Mapping of column names to a list of dictionaries. + Each dictionary contains the 'key' and 'value' keys. + :param provider: The provider of the metadata ('user' by default). + :return: Dictionary with 'metadata' key under which the table metadata is stored as a list of dictionaries. + Each dictionary contains the 'key' and 'value' keys. Under 'columnsMetadata' key, the column metadata + is stored as a mapping of column names to a list of dictionaries. + """ + if not metadata and not columns_metadata: + raise ValueError('At least one of the `metadata` or `columns_metadata` arguments must be provided.') + + payload: dict[str, Any] = {'provider': provider} + if metadata: + payload['metadata'] = [{'key': key, 'value': value} for key, value in metadata.items()] + if columns_metadata: + payload['columnsMetadata'] = columns_metadata + + return cast(JsonDict, await self.post(endpoint=f'tables/{table_id}/metadata', data=payload)) + + async def trigger_event( + self, + message: str, + component_id: str, + configuration_id: str | None = None, + event_type: StorageEventType | None = None, + params: Mapping[str, Any] | None = None, + results: Mapping[str, Any] | None = None, + duration: float | None = None, + run_id: str | None = None, + ) -> JsonDict: + """ + Sends a Storage API event. + + :param message: The event message. + :param component_id: The ID of the component triggering the event. + :param configuration_id: The ID of the component configuration triggering the event. + :param event_type: The type of event. + :param params: The component parameters. The structure of the params object must follow the JSON schema + registered for the component_id. + :param results: The component results. The structure of the results object must follow the JSON schema + registered for the component_id. + :param duration: The component processing duration in seconds. + :param run_id: The ID of the associated component job. + + :return: Dictionary with the new event ID. + """ + payload: dict[str, Any] = { + 'message': message, + 'component': component_id, + } + if configuration_id: + payload['configurationId'] = configuration_id + if event_type: + payload['type'] = event_type + if params: + payload['params'] = params + if results: + payload['results'] = results + if duration is not None: + # The events API ignores floats, so we round up to the nearest integer. + payload['duration'] = int(math.ceil(duration)) + if run_id: + payload['runId'] = run_id + + LOG.info(f'[trigger_event] payload={payload}') + + return cast(JsonDict, await self.post(endpoint='events', data=payload)) + + async def workspace_create( + self, + login_type: str, + backend: str, + async_run: bool = True, + read_only_storage_access: bool = False, + ) -> JsonDict: + """ + Creates a new workspace. + + :param async_run: If True, the workspace creation is run asynchronously. + :param read_only_storage_access: If True, the workspace has read-only access to the storage. + :return: The SAPI call response - created workspace or raise an error. + """ + return cast( + JsonDict, + await self.post( + endpoint=f'branch/{self.branch_id}/workspaces', + params={'async': async_run}, + data={ + 'readOnlyStorageAccess': read_only_storage_access, + 'loginType': login_type, + 'backend': backend, + }, + ), + ) + + async def workspace_detail(self, workspace_id: int) -> JsonDict: + """ + Retrieves information about a given workspace. + + :param workspace_id: The id of the workspace + :return: Workspace details as dictionary + """ + return cast(JsonDict, await self.get(endpoint=f'branch/{self.branch_id}/workspaces/{workspace_id}')) + + async def workspace_query(self, workspace_id: int, query: str) -> JsonDict: + """ + Executes a query in a given workspace. + + :param workspace_id: The id of the workspace + :param query: The query to execute + :return: The SAPI call response - query result or raise an error. + """ + return cast( + JsonDict, + await self.post( + endpoint=f'branch/{self.branch_id}/workspaces/{workspace_id}/query', + data={'query': query}, + ), + ) + + async def workspace_list(self) -> list[JsonDict]: + """ + Lists all workspaces in the project. + + :return: List of workspaces + """ + return cast(list[JsonDict], await self.get(endpoint=f'branch/{self.branch_id}/workspaces')) + + async def verify_token(self) -> JsonDict: + """ + Checks the token privileges and returns information about the project to which the token belongs. + + :return: Token and project information + """ + return cast(JsonDict, await self.get(endpoint='tokens/verify')) + + async def project_id(self) -> str: + """ + Retrieves the project id. + :return: Project id. + """ + raw_data = cast(JsonDict, await self.get(endpoint='tokens/verify')) + assert isinstance(raw_data['owner'], dict) + return str(raw_data['owner']['id']) + + async def is_enabled(self, features: ProjectFeature | Iterable[ProjectFeature]) -> bool: + """ + Checks if the features are enabled in the project - conjunction of features. + :param features: The features to check. + :return: True if the features are enabled, False otherwise. + """ + features = [features] if isinstance(features, str) else features + verified_info = await self.verify_token() + project_data = cast(JsonDict, verified_info['owner']) + project_features = cast(list[str], project_data.get('features', [])) + return all(feature in project_features for feature in features) + + async def token_create( + self, + description: str, + component_access: list[str] | None = None, + expires_in: int | None = None, + ) -> JsonDict: + """ + Creates a new Storage API token. + + :param description: Description of the token + :param component_access: List of component IDs the token should have access to + :param expires_in: Token expiration time in seconds + :return: Token creation response containing the token and its details + """ + token_data: dict[str, Any] = {'description': description} + + if component_access: + token_data['componentAccess'] = component_access + + if expires_in: + token_data['expiresIn'] = expires_in + + return cast(JsonDict, await self.post(endpoint='tokens', data=token_data)) diff --git a/src/keboola_mcp_server/tools/components/model.py b/src/keboola_mcp_server/tools/components/model.py index b50bd40a..86de9f13 100644 --- a/src/keboola_mcp_server/tools/components/model.py +++ b/src/keboola_mcp_server/tools/components/model.py @@ -37,7 +37,7 @@ from pydantic import AliasChoices, BaseModel, Field -from keboola_mcp_server.clients.client import ComponentAPIResponse, ConfigurationAPIResponse +from keboola_mcp_server.clients.storage import ComponentAPIResponse, ConfigurationAPIResponse from keboola_mcp_server.links import Link # ============================================================================ diff --git a/src/keboola_mcp_server/tools/components/tools.py b/src/keboola_mcp_server/tools/components/tools.py index 4892c2ed..a36e21d0 100644 --- a/src/keboola_mcp_server/tools/components/tools.py +++ b/src/keboola_mcp_server/tools/components/tools.py @@ -37,7 +37,8 @@ from mcp.types import ToolAnnotations from pydantic import Field -from keboola_mcp_server.clients.client import ConfigurationAPIResponse, JsonDict, KeboolaClient +from keboola_mcp_server.clients.client import KeboolaClient +from keboola_mcp_server.clients.storage import ConfigurationAPIResponse, JsonDict from keboola_mcp_server.errors import tool_errors from keboola_mcp_server.links import ProjectLinksManager from keboola_mcp_server.mcp import KeboolaMcpServer, exclude_none_serializer diff --git a/src/keboola_mcp_server/tools/components/utils.py b/src/keboola_mcp_server/tools/components/utils.py index b128e0fc..27889dda 100644 --- a/src/keboola_mcp_server/tools/components/utils.py +++ b/src/keboola_mcp_server/tools/components/utils.py @@ -28,7 +28,9 @@ from httpx import HTTPStatusError from pydantic import AliasChoices, BaseModel, Field -from keboola_mcp_server.clients.client import ComponentAPIResponse, ConfigurationAPIResponse, JsonDict, KeboolaClient +from keboola_mcp_server.clients.base import JsonDict +from keboola_mcp_server.clients.client import KeboolaClient +from keboola_mcp_server.clients.storage import ComponentAPIResponse, ConfigurationAPIResponse from keboola_mcp_server.config import MetadataField from keboola_mcp_server.tools.components.model import ( AllComponentTypes, diff --git a/src/keboola_mcp_server/tools/flow/model.py b/src/keboola_mcp_server/tools/flow/model.py index c5a85313..7afae2b6 100644 --- a/src/keboola_mcp_server/tools/flow/model.py +++ b/src/keboola_mcp_server/tools/flow/model.py @@ -7,7 +7,8 @@ from pydantic import AliasChoices, BaseModel, Field -from keboola_mcp_server.clients.client import ORCHESTRATOR_COMPONENT_ID, APIFlowResponse, FlowType +from keboola_mcp_server.clients.client import ORCHESTRATOR_COMPONENT_ID, FlowType +from keboola_mcp_server.clients.storage import APIFlowResponse from keboola_mcp_server.links import Link # ============================================================================= diff --git a/src/keboola_mcp_server/tools/flow/tools.py b/src/keboola_mcp_server/tools/flow/tools.py index bbc0c70d..f7dcbc7e 100644 --- a/src/keboola_mcp_server/tools/flow/tools.py +++ b/src/keboola_mcp_server/tools/flow/tools.py @@ -12,14 +12,14 @@ from pydantic import Field from keboola_mcp_server import resources +from keboola_mcp_server.clients.base import JsonDict from keboola_mcp_server.clients.client import ( CONDITIONAL_FLOW_COMPONENT_ID, ORCHESTRATOR_COMPONENT_ID, - CreateConfigurationAPIResponse, FlowType, - JsonDict, KeboolaClient, ) +from keboola_mcp_server.clients.storage import CreateConfigurationAPIResponse from keboola_mcp_server.errors import tool_errors from keboola_mcp_server.links import ProjectLinksManager from keboola_mcp_server.mcp import exclude_none_serializer @@ -204,7 +204,7 @@ async def create_flow( flow_links = links_manager.get_flow_links(flow_id=api_config.id, flow_name=api_config.name, flow_type=flow_type) tool_response = FlowToolResponse( id=api_config.id, - description=api_config.description, + description=api_config.description or '', timestamp=datetime.now(timezone.utc), success=True, links=flow_links, @@ -268,7 +268,7 @@ async def create_conditional_flow( flow_links = links_manager.get_flow_links(flow_id=api_config.id, flow_name=api_config.name, flow_type=flow_type) tool_response = FlowToolResponse( id=api_config.id, - description=api_config.description, + description=api_config.description or '', timestamp=datetime.now(timezone.utc), success=True, links=flow_links, diff --git a/src/keboola_mcp_server/tools/flow/utils.py b/src/keboola_mcp_server/tools/flow/utils.py index a4d89da4..1342b8b5 100644 --- a/src/keboola_mcp_server/tools/flow/utils.py +++ b/src/keboola_mcp_server/tools/flow/utils.py @@ -9,11 +9,10 @@ CONDITIONAL_FLOW_COMPONENT_ID, FLOW_TYPES, ORCHESTRATOR_COMPONENT_ID, - APIFlowResponse, FlowType, - JsonDict, KeboolaClient, ) +from keboola_mcp_server.clients.storage import APIFlowResponse, JsonDict from keboola_mcp_server.tools.flow.model import ( FlowPhase, FlowSummary, diff --git a/src/keboola_mcp_server/tools/project.py b/src/keboola_mcp_server/tools/project.py index ec93db77..6b75df89 100644 --- a/src/keboola_mcp_server/tools/project.py +++ b/src/keboola_mcp_server/tools/project.py @@ -6,7 +6,8 @@ from mcp.types import ToolAnnotations from pydantic import BaseModel, Field -from keboola_mcp_server.clients.client import JsonDict, KeboolaClient +from keboola_mcp_server.clients.base import JsonDict +from keboola_mcp_server.clients.client import KeboolaClient from keboola_mcp_server.config import MetadataField from keboola_mcp_server.errors import tool_errors from keboola_mcp_server.links import Link, ProjectLinksManager diff --git a/src/keboola_mcp_server/tools/search.py b/src/keboola_mcp_server/tools/search.py index 3b865568..837e6a81 100644 --- a/src/keboola_mcp_server/tools/search.py +++ b/src/keboola_mcp_server/tools/search.py @@ -8,7 +8,9 @@ from mcp.types import ToolAnnotations from pydantic import BaseModel, Field -from keboola_mcp_server.clients.client import GlobalSearchResponse, ItemType, KeboolaClient, SuggestedComponent +from keboola_mcp_server.clients.ai_service import SuggestedComponent +from keboola_mcp_server.clients.client import KeboolaClient +from keboola_mcp_server.clients.storage import GlobalSearchResponse, ItemType from keboola_mcp_server.errors import tool_errors LOG = logging.getLogger(__name__) diff --git a/src/keboola_mcp_server/tools/storage.py b/src/keboola_mcp_server/tools/storage.py index a85d069f..13ec0c4f 100644 --- a/src/keboola_mcp_server/tools/storage.py +++ b/src/keboola_mcp_server/tools/storage.py @@ -9,7 +9,8 @@ from mcp.types import ToolAnnotations from pydantic import AliasChoices, BaseModel, Field, model_validator -from keboola_mcp_server.clients.client import JsonDict, KeboolaClient, get_metadata_property +from keboola_mcp_server.clients.base import JsonDict +from keboola_mcp_server.clients.client import KeboolaClient, get_metadata_property from keboola_mcp_server.config import MetadataField from keboola_mcp_server.errors import tool_errors from keboola_mcp_server.links import Link, ProjectLinksManager diff --git a/tests/clients/test_client.py b/tests/clients/test_client.py index 94592d22..e86cded1 100644 --- a/tests/clients/test_client.py +++ b/tests/clients/test_client.py @@ -5,7 +5,8 @@ import httpx import pytest -from keboola_mcp_server.clients.client import KeboolaClient, RawKeboolaClient +from keboola_mcp_server.clients.base import RawKeboolaClient +from keboola_mcp_server.clients.client import KeboolaClient @pytest.fixture diff --git a/tests/conftest.py b/tests/conftest.py index 01312a84..acb182aa 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -3,13 +3,11 @@ from mcp.server.session import ServerSession from mcp.shared.context import RequestContext -from keboola_mcp_server.clients.client import ( - AIServiceClient, - AsyncStorageClient, - JobsQueueClient, - KeboolaClient, - RawKeboolaClient, -) +from keboola_mcp_server.clients.ai_service import AIServiceClient +from keboola_mcp_server.clients.base import RawKeboolaClient +from keboola_mcp_server.clients.client import KeboolaClient +from keboola_mcp_server.clients.jobs_queue import JobsQueueClient +from keboola_mcp_server.clients.storage import AsyncStorageClient from keboola_mcp_server.config import Config from keboola_mcp_server.mcp import ServerState from keboola_mcp_server.workspace import WorkspaceManager diff --git a/tests/tools/components/test_validation.py b/tests/tools/components/test_validation.py index cbd7c5b8..44a07a78 100644 --- a/tests/tools/components/test_validation.py +++ b/tests/tools/components/test_validation.py @@ -5,7 +5,8 @@ import jsonschema import pytest -from keboola_mcp_server.clients.client import ORCHESTRATOR_COMPONENT_ID, ComponentAPIResponse, JsonDict +from keboola_mcp_server.clients.client import ORCHESTRATOR_COMPONENT_ID +from keboola_mcp_server.clients.storage import ComponentAPIResponse, JsonDict from keboola_mcp_server.tools import validation from keboola_mcp_server.tools.components.model import Component diff --git a/tests/tools/flow/test_model.py b/tests/tools/flow/test_model.py index 3b7132b4..22cf3ce8 100644 --- a/tests/tools/flow/test_model.py +++ b/tests/tools/flow/test_model.py @@ -1,6 +1,7 @@ from typing import Any -from keboola_mcp_server.clients.client import ORCHESTRATOR_COMPONENT_ID, APIFlowResponse +from keboola_mcp_server.clients.client import ORCHESTRATOR_COMPONENT_ID +from keboola_mcp_server.clients.storage import APIFlowResponse from keboola_mcp_server.tools.flow.model import ( Flow, FlowConfiguration, diff --git a/tests/tools/test_doc.py b/tests/tools/test_doc.py index 1cdaa2f1..5b9a69db 100644 --- a/tests/tools/test_doc.py +++ b/tests/tools/test_doc.py @@ -2,7 +2,8 @@ from mcp.server.fastmcp import Context from pytest_mock import MockerFixture -from keboola_mcp_server.clients.client import DocsQuestionResponse, KeboolaClient +from keboola_mcp_server.clients.ai_service import DocsQuestionResponse +from keboola_mcp_server.clients.client import KeboolaClient from keboola_mcp_server.tools.doc import DocsAnswer, docs_query diff --git a/tests/tools/test_search.py b/tests/tools/test_search.py index 3d879040..48a7c01f 100644 --- a/tests/tools/test_search.py +++ b/tests/tools/test_search.py @@ -5,7 +5,8 @@ from fastmcp import Context from pytest_mock import MockerFixture -from keboola_mcp_server.clients.client import GlobalSearchResponse, KeboolaClient +from keboola_mcp_server.clients.client import KeboolaClient +from keboola_mcp_server.clients.storage import GlobalSearchResponse from keboola_mcp_server.tools.search import ( DEFAULT_GLOBAL_SEARCH_LIMIT, GlobalSearchOutput, From 892b024235c215e5e037925e1998dd81109cda75 Mon Sep 17 00:00:00 2001 From: mariankrotil Date: Mon, 18 Aug 2025 17:26:10 +0200 Subject: [PATCH 19/35] AI-1343 test: skip global serarch test for client method as it is unstable due to long db sync --- integtests/clients/test_client.py | 1 + 1 file changed, 1 insertion(+) diff --git a/integtests/clients/test_client.py b/integtests/clients/test_client.py index 6b91d73d..8afe28e1 100644 --- a/integtests/clients/test_client.py +++ b/integtests/clients/test_client.py @@ -26,6 +26,7 @@ async def test_global_search(self, storage_client: AsyncStorageClient): assert ret.by_project == {} @pytest.mark.asyncio + @pytest.mark.skip(reason='Writing to the global search index takes too long, making the test unstable.') 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') From 51f2bfb03bc3bf317e2353bfe8d17a83b5953ce1 Mon Sep 17 00:00:00 2001 From: mariankrotil Date: Mon, 18 Aug 2025 17:32:12 +0200 Subject: [PATCH 20/35] AI-1343 style: apply isort --- integtests/test_validate.py | 2 +- integtests/tools/test_search.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/integtests/test_validate.py b/integtests/test_validate.py index 53b7c97f..8757e17d 100644 --- a/integtests/test_validate.py +++ b/integtests/test_validate.py @@ -16,9 +16,9 @@ import jsonschema import pytest +from keboola_mcp_server.clients.base import JsonDict from keboola_mcp_server.clients.client import KeboolaClient from keboola_mcp_server.clients.storage import ComponentAPIResponse -from keboola_mcp_server.clients.base import JsonDict from keboola_mcp_server.tools.components.model import Component from keboola_mcp_server.tools.validation import KeboolaParametersValidator diff --git a/integtests/tools/test_search.py b/integtests/tools/test_search.py index 19140d54..bb815e84 100644 --- a/integtests/tools/test_search.py +++ b/integtests/tools/test_search.py @@ -4,8 +4,8 @@ from fastmcp import Context from integtests.conftest import BucketDef, ConfigDef, TableDef -from keboola_mcp_server.clients.client import KeboolaClient from keboola_mcp_server.clients.ai_service import SuggestedComponent +from keboola_mcp_server.clients.client import KeboolaClient from keboola_mcp_server.tools.search import GlobalSearchOutput, find_component_id, search LOG = logging.getLogger(__name__) From ad951c71eedc0c5a3e686f3f79090cbffab44a9e Mon Sep 17 00:00:00 2001 From: mariankrotil Date: Tue, 19 Aug 2025 14:35:38 +0200 Subject: [PATCH 21/35] AI-1343 refactor: validate parameters in encryption client for encrypt, improve docs --- src/keboola_mcp_server/clients/encryption.py | 34 +++++++++++++------- 1 file changed, 22 insertions(+), 12 deletions(-) diff --git a/src/keboola_mcp_server/clients/encryption.py b/src/keboola_mcp_server/clients/encryption.py index 46785870..9b1d51f0 100644 --- a/src/keboola_mcp_server/clients/encryption.py +++ b/src/keboola_mcp_server/clients/encryption.py @@ -21,14 +21,14 @@ def base_api_url(self) -> str: def create( cls, root_url: str, - token: str, + token: Optional[str] = None, headers: dict[str, Any] | None = None, ) -> 'AsyncEncryptionClient': """ Creates an AsyncEncryptionClient from a Keboola Storage API token. :param root_url: The root URL of the service API - :param token: The Keboola Storage API token + :param token: The Keboola Storage API token. If None, the client will not send any authorization header. :param headers: Additional headers for the requests :return: A new instance of AsyncEncryptionClient """ @@ -41,13 +41,16 @@ def create( ) async def encrypt( - self, value: Any, project_id: str, component_id: Optional[str] = None, config_id: Optional[str] = None + self, value: Any, project_id: Optional[str], component_id: Optional[str] = None, config_id: Optional[str] = None ) -> Any: """ - Encrypt a value using the encryption service, returns encrypted value. - if value is a dict, values whose keys start with '#' are encrypted. - if value is a str, it is encrypted. - if value contains already encrypted values, they are returned as is. + Encrypt a value using the encryption service, returns encrypted value. Parameters are optional and the ciphers + created by the service are dependent on those parameters when decrypting. Decryption is done automatically + when using encrypted values in a request to Storage API (for components) + See: https://developers.keboola.com/overview/encryption/ + If value is a dict, values whose keys start with '#' are encrypted. + If value is a str, it is encrypted. + If value contains already encrypted values, they are returned as is. :param value: The value to encrypt :param project_id: The project ID @@ -55,13 +58,20 @@ async def encrypt( :param config_id: The config ID (optional) :return: The encrypted value, same type as input """ + if component_id and project_id is None: + raise ValueError('project_id is required if component_id is provided') + if config_id and not component_id: + raise ValueError('component_id is required if config_id is provided') + + params = { + 'componentId': component_id, + 'projectId': project_id, + 'configId': config_id, + } + params = {k: v for k, v in params.items() if v is not None} response = await self.raw_client.post( endpoint='encrypt', - params={ - 'componentId': component_id, - 'projectId': project_id, - 'configId': config_id, - }, + params=params, data=value, ) return response From 7ddeea2928b6182864b3f3c5e9c4b7480be92aa9 Mon Sep 17 00:00:00 2001 From: mariankrotil Date: Tue, 19 Aug 2025 14:38:16 +0200 Subject: [PATCH 22/35] AI-1343 fix: make clients token parameter optional and handle headers when no token provided as encryption client does not use storage token in the headers --- integtests/clients/test_encryption.py | 6 ++++++ src/keboola_mcp_server/clients/ai_service.py | 6 +++--- src/keboola_mcp_server/clients/base.py | 17 +++++++++-------- src/keboola_mcp_server/clients/client.py | 3 ++- src/keboola_mcp_server/clients/data_science.py | 4 ++-- src/keboola_mcp_server/clients/jobs_queue.py | 5 +++-- src/keboola_mcp_server/clients/storage.py | 4 ++-- tests/clients/test_client.py | 4 ++-- 8 files changed, 29 insertions(+), 20 deletions(-) diff --git a/integtests/clients/test_encryption.py b/integtests/clients/test_encryption.py index ceb0b433..a0acd5df 100644 --- a/integtests/clients/test_encryption.py +++ b/integtests/clients/test_encryption.py @@ -5,6 +5,12 @@ from keboola_mcp_server.clients.client import DATA_APP_COMPONENT_ID, KeboolaClient +@pytest.mark.asyncio +async def test_encrypt_without_project_id(keboola_client: KeboolaClient) -> None: + """Check that the encryption client does not send any authorization headers.""" + assert 'Authorization' not in keboola_client.encryption_client.raw_client.headers + assert 'X-StorageAPI-Token' not in keboola_client.encryption_client.raw_client.headers + @pytest.mark.asyncio async def test_encrypt_string_not_equal(keboola_client: KeboolaClient) -> None: project_id = await keboola_client.storage_client.project_id() diff --git a/src/keboola_mcp_server/clients/ai_service.py b/src/keboola_mcp_server/clients/ai_service.py index f16fe9ff..534f7049 100644 --- a/src/keboola_mcp_server/clients/ai_service.py +++ b/src/keboola_mcp_server/clients/ai_service.py @@ -1,4 +1,4 @@ -from typing import Any, cast +from typing import Any, Optional, cast from pydantic import BaseModel, Field @@ -36,12 +36,12 @@ class AIServiceClient(KeboolaServiceClient): """Async client for Keboola AI Service.""" @classmethod - def create(cls, root_url: str, token: str, headers: dict[str, Any] | None = None) -> 'AIServiceClient': + def create(cls, root_url: str, token: Optional[str], headers: dict[str, Any] | None = None) -> 'AIServiceClient': """ Creates an AIServiceClient from a Keboola Storage API token. :param root_url: The root URL of the AI service API. - :param token: The Keboola Storage API token. + :param token: The Keboola Storage API token. If None, the client will not send any authorization header. :param headers: Additional headers for the requests. :return: A new instance of AIServiceClient. """ diff --git a/src/keboola_mcp_server/clients/base.py b/src/keboola_mcp_server/clients/base.py index bdce266c..fdec90b6 100644 --- a/src/keboola_mcp_server/clients/base.py +++ b/src/keboola_mcp_server/clients/base.py @@ -19,19 +19,20 @@ class RawKeboolaClient: def __init__( self, base_api_url: str, - api_token: str, + api_token: Optional[str], headers: dict[str, Any] | None = None, timeout: httpx.Timeout | None = None, ) -> None: self.base_api_url = base_api_url self.headers = { 'Content-Type': 'application/json', - 'Accept-encoding': 'gzip', + 'Accept-Encoding': 'gzip', } - if api_token.startswith('Bearer '): - self.headers['Authorization'] = api_token - else: - self.headers['X-StorageAPI-Token'] = api_token + if api_token: + if api_token.startswith('Bearer '): + self.headers['Authorization'] = api_token + else: + self.headers['X-StorageAPI-Token'] = api_token self.timeout = timeout or httpx.Timeout(connect=5.0, read=60.0, write=10.0, pool=5.0) if headers: self.headers.update(headers) @@ -239,12 +240,12 @@ def __init__(self, raw_client: RawKeboolaClient) -> None: self.raw_client = raw_client @classmethod - def create(cls, root_url: str, token: str) -> 'KeboolaServiceClient': + def create(cls, root_url: str, token: Optional[str]) -> 'KeboolaServiceClient': """ Creates a KeboolaServiceClient from a Keboola Storage API token. :param root_url: The root URL of the service API - :param token: The Keboola Storage API token + :param token: The Keboola Storage API token. If None, the client will not send any authorization header. :return: A new instance of KeboolaServiceClient """ return cls(raw_client=RawKeboolaClient(base_api_url=root_url, api_token=token)) diff --git a/src/keboola_mcp_server/clients/client.py b/src/keboola_mcp_server/clients/client.py index ac1b1dbd..feecaf45 100644 --- a/src/keboola_mcp_server/clients/client.py +++ b/src/keboola_mcp_server/clients/client.py @@ -107,8 +107,9 @@ def __init__(self, storage_api_token: str, storage_api_url: str, bearer_token: s self.data_science_client = AsyncDataScienceClient.create( root_url=data_science_api_url, token=self.token, headers=self._get_headers() ) + # The encryption service does not require an authorization header, so we pass None as the token self.encryption_client = AsyncEncryptionClient.create( - root_url=encryption_api_url, token=self.token, headers=self._get_headers() + root_url=encryption_api_url, token=None, headers=self._get_headers() ) @classmethod diff --git a/src/keboola_mcp_server/clients/data_science.py b/src/keboola_mcp_server/clients/data_science.py index 6db24e1f..51ca3beb 100644 --- a/src/keboola_mcp_server/clients/data_science.py +++ b/src/keboola_mcp_server/clients/data_science.py @@ -108,14 +108,14 @@ def base_api_url(self) -> str: def create( cls, root_url: str, - token: str, + token: Optional[str], headers: dict[str, Any] | None = None, ) -> 'AsyncDataScienceClient': """ Creates an AsyncDataScienceClient from a Keboola Storage API token. :param root_url: The root URL of the service API - :param token: The Keboola Storage API token + :param token: The Keboola Storage API token. If None, the client will not send any authorization header. :param headers: Additional headers for the requests :return: A new instance of AsyncDataScienceClient """ diff --git a/src/keboola_mcp_server/clients/jobs_queue.py b/src/keboola_mcp_server/clients/jobs_queue.py index 83c80eff..9af9fd3e 100644 --- a/src/keboola_mcp_server/clients/jobs_queue.py +++ b/src/keboola_mcp_server/clients/jobs_queue.py @@ -9,12 +9,13 @@ class JobsQueueClient(KeboolaServiceClient): """ @classmethod - def create(cls, root_url: str, token: str, headers: dict[str, Any] | None = None) -> 'JobsQueueClient': + def create(cls, root_url: str, token: Optional[str], headers: dict[str, Any] | None = None) -> 'JobsQueueClient': """ Creates a JobsQueue client. :param root_url: Root url of API. e.g. "https://queue.keboola.com/". - :param token: A key for the Storage API. Can be found in the storage console. + :param token: A key for the Storage API. Can be found in the storage console. If None, the client will not send + any authorization header. :param headers: Additional headers for the requests. :return: A new instance of JobsQueueClient. """ diff --git a/src/keboola_mcp_server/clients/storage.py b/src/keboola_mcp_server/clients/storage.py index df1796d8..49c3a822 100644 --- a/src/keboola_mcp_server/clients/storage.py +++ b/src/keboola_mcp_server/clients/storage.py @@ -297,7 +297,7 @@ def base_api_url(self) -> str: def create( cls, root_url: str, - token: str, + token: Optional[str], version: str = 'v2', branch_id: str = 'default', headers: dict[str, Any] | None = None, @@ -306,7 +306,7 @@ def create( Creates an AsyncStorageClient from a Keboola Storage API token. :param root_url: The root URL of the service API - :param token: The Keboola Storage API token + :param token: The Keboola Storage API token, If None, the client will not send any authorization header. :param version: The version of the API to use (default: 'v2') :param branch_id: The id of the branch :param headers: Additional headers for the requests diff --git a/tests/clients/test_client.py b/tests/clients/test_client.py index e86cded1..6876937b 100644 --- a/tests/clients/test_client.py +++ b/tests/clients/test_client.py @@ -192,7 +192,7 @@ async def test_trigger_event( params=None, headers={ 'Content-Type': 'application/json', - 'Accept-encoding': 'gzip', + 'Accept-Encoding': 'gzip', 'X-StorageAPI-Token': 'test-token', 'User-Agent': f'Keboola MCP Server/{version} app_env=local', }, @@ -277,7 +277,7 @@ async def test_token_create( params=None, headers={ 'Content-Type': 'application/json', - 'Accept-encoding': 'gzip', + 'Accept-Encoding': 'gzip', 'X-StorageAPI-Token': 'test-token', 'User-Agent': f'Keboola MCP Server/{version} app_env=local', }, From 0fb290f8029b9d602c46dbd5805bed35bbb66cbc Mon Sep 17 00:00:00 2001 From: mariankrotil Date: Tue, 19 Aug 2025 14:39:39 +0200 Subject: [PATCH 23/35] AI-1343 style: apply black --- integtests/clients/test_encryption.py | 1 + 1 file changed, 1 insertion(+) diff --git a/integtests/clients/test_encryption.py b/integtests/clients/test_encryption.py index a0acd5df..e04aa6ae 100644 --- a/integtests/clients/test_encryption.py +++ b/integtests/clients/test_encryption.py @@ -11,6 +11,7 @@ async def test_encrypt_without_project_id(keboola_client: KeboolaClient) -> None assert 'Authorization' not in keboola_client.encryption_client.raw_client.headers assert 'X-StorageAPI-Token' not in keboola_client.encryption_client.raw_client.headers + @pytest.mark.asyncio async def test_encrypt_string_not_equal(keboola_client: KeboolaClient) -> None: project_id = await keboola_client.storage_client.project_id() From a7b124e98ccdc90d86dd89cc08e6237202e170e3 Mon Sep 17 00:00:00 2001 From: mariankrotil Date: Tue, 19 Aug 2025 14:55:49 +0200 Subject: [PATCH 24/35] AI-1343 refactor: update description, rename classes wrt reviews --- integtests/clients/test_data_science.py | 6 +- src/keboola_mcp_server/clients/client.py | 8 +-- .../clients/data_science.py | 59 +++++++++++-------- src/keboola_mcp_server/clients/encryption.py | 10 ++-- 4 files changed, 47 insertions(+), 36 deletions(-) diff --git a/integtests/clients/test_data_science.py b/integtests/clients/test_data_science.py index 8e01fe68..05edf02a 100644 --- a/integtests/clients/test_data_science.py +++ b/integtests/clients/test_data_science.py @@ -3,7 +3,7 @@ import pytest from keboola_mcp_server.clients.client import DATA_APP_COMPONENT_ID, KeboolaClient -from keboola_mcp_server.clients.data_science import AsyncDataScienceClient, DataAppResponse +from keboola_mcp_server.clients.data_science import DataAppResponse, DataScienceClient LOG = logging.getLogger(__name__) @@ -39,13 +39,13 @@ def _public_access_authorization() -> dict[str, object]: @pytest.fixture -def ds_client(keboola_client: KeboolaClient) -> AsyncDataScienceClient: +def ds_client(keboola_client: KeboolaClient) -> DataScienceClient: return keboola_client.data_science_client @pytest.mark.asyncio async def test_create_and_fetch_data_app( - ds_client: AsyncDataScienceClient, unique_id: str, keboola_client: KeboolaClient + ds_client: DataScienceClient, unique_id: str, keboola_client: KeboolaClient ) -> None: """Test creating a data app and fetching it from detail and list endpoints""" slug = f'test-app-{unique_id}' diff --git a/src/keboola_mcp_server/clients/client.py b/src/keboola_mcp_server/clients/client.py index feecaf45..863cb732 100644 --- a/src/keboola_mcp_server/clients/client.py +++ b/src/keboola_mcp_server/clients/client.py @@ -6,8 +6,8 @@ from typing import Any, Literal, Mapping, Optional, Sequence, TypeVar from keboola_mcp_server.clients.ai_service import AIServiceClient -from keboola_mcp_server.clients.data_science import AsyncDataScienceClient -from keboola_mcp_server.clients.encryption import AsyncEncryptionClient +from keboola_mcp_server.clients.data_science import DataScienceClient +from keboola_mcp_server.clients.encryption import EncryptionClient from keboola_mcp_server.clients.jobs_queue import JobsQueueClient from keboola_mcp_server.clients.storage import AsyncStorageClient @@ -104,11 +104,11 @@ def __init__(self, storage_api_token: str, storage_api_url: str, bearer_token: s self.ai_service_client = AIServiceClient.create( root_url=ai_service_api_url, token=self.token, headers=self._get_headers() ) - self.data_science_client = AsyncDataScienceClient.create( + self.data_science_client = DataScienceClient.create( root_url=data_science_api_url, token=self.token, headers=self._get_headers() ) # The encryption service does not require an authorization header, so we pass None as the token - self.encryption_client = AsyncEncryptionClient.create( + self.encryption_client = EncryptionClient.create( root_url=encryption_api_url, token=None, headers=self._get_headers() ) diff --git a/src/keboola_mcp_server/clients/data_science.py b/src/keboola_mcp_server/clients/data_science.py index 51ca3beb..133aabe6 100644 --- a/src/keboola_mcp_server/clients/data_science.py +++ b/src/keboola_mcp_server/clients/data_science.py @@ -1,4 +1,5 @@ import logging +from datetime import datetime from typing import Any, Optional, cast from pydantic import AliasChoices, BaseModel, Field @@ -58,7 +59,12 @@ class DataAppConfig(BaseModel): class Parameters(BaseModel): class DataApp(BaseModel): slug: str = Field(description='The slug of the data app') - streamlit: dict[str, str] = Field(description='The streamlit config.toml file') + streamlit: dict[str, str] = Field( + description=( + 'The streamlit configuration, expected to have a key with TOML file name and the value with the ' + 'file content' + ) + ) secrets: Optional[dict[str, str]] = Field(description='The secrets of the data app', default=None) size: str = Field(description='The size of the data app') @@ -90,11 +96,11 @@ class AppProxy(BaseModel): storage: dict[str, Any] = Field(description='The storage of the data app', default_factory=dict) -class AsyncDataScienceClient(KeboolaServiceClient): +class DataScienceClient(KeboolaServiceClient): def __init__(self, raw_client: RawKeboolaClient) -> None: """ - Creates an AsyncDataScienceClient from a RawKeboolaClient and a branch id. + Creates an DataScienceClient from a RawKeboolaClient. :param raw_client: The raw client to use """ @@ -110,18 +116,18 @@ def create( root_url: str, token: Optional[str], headers: dict[str, Any] | None = None, - ) -> 'AsyncDataScienceClient': + ) -> 'DataScienceClient': """ - Creates an AsyncDataScienceClient from a Keboola Storage API token. + Creates an DataScienceClient from a Keboola Storage API token. :param root_url: The root URL of the service API :param token: The Keboola Storage API token. If None, the client will not send any authorization header. :param headers: Additional headers for the requests - :return: A new instance of AsyncDataScienceClient + :return: A new instance of DataScienceClient """ return cls( raw_client=RawKeboolaClient( - base_api_url=f'{root_url}', + base_api_url=root_url, api_token=token, headers=headers, ) @@ -156,7 +162,9 @@ async def deploy_data_app(self, data_app_id: str, config_version: str) -> DataAp async def suspend_data_app(self, data_app_id: str) -> DataAppResponse: """ - Suspend a data app by its ID. + Suspend a data app by setting its desired state to 'stopped'. + :param data_app_id: data app ID to suspend + :return: Updated data app response with new state """ data = {'desiredState': 'stopped'} response = await self.patch(endpoint=f'apps/{data_app_id}', data=data) @@ -206,6 +214,7 @@ async def create_data_app( async def delete_data_app(self, data_app_id: str) -> None: """ Delete a data app by its ID. + :param data_app_id: ID of the data app to delete """ await self.delete(endpoint=f'apps/{data_app_id}') @@ -216,35 +225,37 @@ async def list_data_apps(self, limit: int = 100, offset: int = 0) -> list[DataAp response = await self.get(endpoint='apps', params={'limit': limit, 'offset': offset}) return [DataAppResponse.model_validate(app) for app in response] - async def tail_app_logs(self, app_id: str, since: Optional[str] = None, *, lines: Optional[int] = None) -> str: + async def tail_app_logs( + self, + app_id: str, + *, + since: Optional[datetime], + lines: Optional[int], + ) -> str: """ - Tail application logs. Either `since` or `lines` must be provided but not both at the same time, otherwise it - uses the `lines` parameter. In case when none of the parameters are provided, it uses the `lines` parameter with + Tail application logs. Either `since` or `lines` must be provided but not both at the same time. + In case when none of the parameters are provided, it uses the `lines` parameter with the last 100 lines. :param app_id: ID of the app. - :param since: ISO-8601 timestamp with nanoseconds. - E.g: since = (datetime.now(timezone.utc) - timedelta(days=1)).isoformat() + :param since: ISO-8601 timestamp with nanoseconds as a datetime object. + E.g: since = datetime.now(timezone.utc) - timedelta(days=1) :param lines: Number of log lines from the end. Defaults to 100. :return: Logs as plain text. - :raise requests.HTTPError: For non-200 status codes. + :raise ValueError: If both "since" and "lines" are provided. + :raise ValueError: If neither "since" nor "lines" are provided. + :raise httpx.HTTPStatusError: For non-200 status codes. """ if since and lines: - LOG.warning( - 'You cannot use both "since" and "lines" query parameters together. Using the "lines" parameter.' - ) - since = None - + raise ValueError('You cannot use both "since" and "lines" query parameters together.') elif not since and not lines: - LOG.info( - 'No "since" or "lines" query parameters provided. Using "lines" with the last 100 lines as default.' - ) - lines = 100 + raise ValueError('Either "since" or "lines" must be provided.') if lines: lines = max(lines, 1) # Ensure lines is at least 1 params = {'lines': lines} elif since: - params = {'since': since} + iso_since = since.isoformat(timespec='nanoseconds') + params = {'since': iso_since} else: raise ValueError('Either "since" or "lines" must be provided.') diff --git a/src/keboola_mcp_server/clients/encryption.py b/src/keboola_mcp_server/clients/encryption.py index 9b1d51f0..a9e478d7 100644 --- a/src/keboola_mcp_server/clients/encryption.py +++ b/src/keboola_mcp_server/clients/encryption.py @@ -3,11 +3,11 @@ from keboola_mcp_server.clients.base import KeboolaServiceClient, RawKeboolaClient -class AsyncEncryptionClient(KeboolaServiceClient): +class EncryptionClient(KeboolaServiceClient): def __init__(self, raw_client: RawKeboolaClient) -> None: """ - Creates an AsyncEncryptionClient from a RawKeboolaClient. + Creates an EncryptionClient from a RawKeboolaClient. :param raw_client: The raw client to use """ @@ -23,14 +23,14 @@ def create( root_url: str, token: Optional[str] = None, headers: dict[str, Any] | None = None, - ) -> 'AsyncEncryptionClient': + ) -> 'EncryptionClient': """ - Creates an AsyncEncryptionClient from a Keboola Storage API token. + Creates an EncryptionClient from a Keboola Storage API token. :param root_url: The root URL of the service API :param token: The Keboola Storage API token. If None, the client will not send any authorization header. :param headers: Additional headers for the requests - :return: A new instance of AsyncEncryptionClient + :return: A new instance of EncryptionClient """ return cls( raw_client=RawKeboolaClient( From 7cda4d8a3b4433857c08d5270aaafe56e7796c76 Mon Sep 17 00:00:00 2001 From: mariankrotil Date: Tue, 19 Aug 2025 14:57:10 +0200 Subject: [PATCH 25/35] AI-1343 style: apply tox --- src/keboola_mcp_server/cli.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/src/keboola_mcp_server/cli.py b/src/keboola_mcp_server/cli.py index 24847bdd..071ef04c 100644 --- a/src/keboola_mcp_server/cli.py +++ b/src/keboola_mcp_server/cli.py @@ -125,10 +125,7 @@ async def lifespan(app: Starlette): async with sse_app.lifespan(app): yield - app = Starlette( - middleware=[Middleware(ForwardSlashMiddleware)], - lifespan=lifespan - ) + app = Starlette(middleware=[Middleware(ForwardSlashMiddleware)], lifespan=lifespan) app.mount('/mcp', http_app) app.mount('/sse', sse_app) # serves /sse/ and /messages custom_routes.add_to_starlette(app) From 008f229fe643f956c25c95eeaf7bdce91feec15a Mon Sep 17 00:00:00 2001 From: mariankrotil Date: Tue, 19 Aug 2025 15:00:57 +0200 Subject: [PATCH 26/35] AI-1343 docs: improve client method get_text description to follow return type --- src/keboola_mcp_server/clients/base.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/keboola_mcp_server/clients/base.py b/src/keboola_mcp_server/clients/base.py index fdec90b6..4fd7f78d 100644 --- a/src/keboola_mcp_server/clients/base.py +++ b/src/keboola_mcp_server/clients/base.py @@ -96,12 +96,12 @@ async def get_text( headers: dict[str, Any] | None = None, ) -> str: """ - Makes a GET request to the service API. + Makes a GET request to the service API and returns the response as text. :param endpoint: API endpoint to call :param params: Query parameters for the request :param headers: Additional headers for the request - :return: API response as dictionary + :return: API response as text """ headers = self.headers | (headers or {}) async with httpx.AsyncClient(timeout=self.timeout) as client: From 5e890170a63fb2b1413bf5c807e3ec5ad6314540 Mon Sep 17 00:00:00 2001 From: mariankrotil Date: Tue, 19 Aug 2025 15:41:54 +0200 Subject: [PATCH 27/35] AI-1343 refactor: update parameters for client methods - structured, handle edge cases for None values --- .../clients/data_science.py | 60 ++++++++----------- 1 file changed, 25 insertions(+), 35 deletions(-) diff --git a/src/keboola_mcp_server/clients/data_science.py b/src/keboola_mcp_server/clients/data_science.py index 133aabe6..12998bfa 100644 --- a/src/keboola_mcp_server/clients/data_science.py +++ b/src/keboola_mcp_server/clients/data_science.py @@ -98,18 +98,6 @@ class AppProxy(BaseModel): class DataScienceClient(KeboolaServiceClient): - def __init__(self, raw_client: RawKeboolaClient) -> None: - """ - Creates an DataScienceClient from a RawKeboolaClient. - - :param raw_client: The raw client to use - """ - super().__init__(raw_client=raw_client) - - @property - def base_api_url(self) -> str: - return self.raw_client.base_api_url - @classmethod def create( cls, @@ -143,19 +131,28 @@ async def get_data_app(self, data_app_id: str) -> DataAppResponse: response = await self.get(endpoint=f'apps/{data_app_id}') return DataAppResponse.model_validate(response) - async def deploy_data_app(self, data_app_id: str, config_version: str) -> DataAppResponse: + async def deploy_data_app( + self, + data_app_id: str, + config_version: str, + *, + restart_if_running: bool = True, + update_dependencies: bool = True, + ) -> DataAppResponse: """ Deploy a data app by its ID. :param data_app_id: The ID of the data app :param config_version: The version of the config to deploy + :param restart_if_running: Whether to restart the data app if it is already running + :param update_dependencies: Whether to update the dependencies of the data app :return: The data app """ data = { 'desiredState': 'running', 'configVersion': config_version, - 'restartIfRunning': True, - 'updateDependencies': True, + 'restartIfRunning': restart_if_running, + 'updateDependencies': update_dependencies, } response = await self.patch(endpoint=f'apps/{data_app_id}', data=data) return DataAppResponse.model_validate(response) @@ -182,31 +179,23 @@ async def create_data_app( self, name: str, description: str, - parameters: dict[str, Any], - authorization: dict[str, Any], + configuration: DataAppConfig, + branch_id: Optional[str] = None, ) -> DataAppResponse: """ - Create a data app. + Create a data app from a simplified config used in the MCP server. :param name: The name of the data app :param description: The description of the data app - :param parameters: The parameters of the data app - :param authorization: The authorization of the data app + :param configuration: The simplified configuration of the data app + :param branch_id: The branch ID of the data app :return: The data app """ - # Validate the parameters and authorization - _params = DataAppConfig.Parameters.model_validate(parameters).model_dump(exclude_none=True, by_alias=True) - _authorization = DataAppConfig.Authorization.model_validate(authorization).model_dump( - exclude_none=True, by_alias=True - ) data = { - 'branchId': None, + 'branchId': branch_id, 'name': name, 'type': 'streamlit', 'description': description, - 'config': { - 'parameters': _params, - 'authorization': _authorization, - }, + 'config': configuration.model_dump(exclude_none=True, by_alias=True), } response = await self.post(endpoint='apps', data=data) return DataAppResponse.model_validate(response) @@ -237,7 +226,8 @@ async def tail_app_logs( In case when none of the parameters are provided, it uses the `lines` parameter with the last 100 lines. :param app_id: ID of the app. - :param since: ISO-8601 timestamp with nanoseconds as a datetime object. + :param since: ISO-8601 timestamp with nanoseconds as a datetime object + Providing microseconds is enough, nanoseconds are not supported via datetime E.g: since = datetime.now(timezone.utc) - timedelta(days=1) :param lines: Number of log lines from the end. Defaults to 100. :return: Logs as plain text. @@ -247,14 +237,14 @@ async def tail_app_logs( """ if since and lines: raise ValueError('You cannot use both "since" and "lines" query parameters together.') - elif not since and not lines: + elif since is None and lines is None: raise ValueError('Either "since" or "lines" must be provided.') - if lines: + if lines is not None: lines = max(lines, 1) # Ensure lines is at least 1 params = {'lines': lines} - elif since: - iso_since = since.isoformat(timespec='nanoseconds') + elif since is not None: + iso_since = since.isoformat(timespec='microseconds') params = {'since': iso_since} else: raise ValueError('Either "since" or "lines" must be provided.') From 8247f128b9ee1612b6f3334d19b6778e7a151299 Mon Sep 17 00:00:00 2001 From: mariankrotil Date: Tue, 19 Aug 2025 15:42:36 +0200 Subject: [PATCH 28/35] AI-1343 test: add unit tests for data science client tail_app_log method --- tests/clients/test_data_science.py | 59 ++++++++++++++++++++++++++++++ 1 file changed, 59 insertions(+) create mode 100644 tests/clients/test_data_science.py diff --git a/tests/clients/test_data_science.py b/tests/clients/test_data_science.py new file mode 100644 index 00000000..875656e8 --- /dev/null +++ b/tests/clients/test_data_science.py @@ -0,0 +1,59 @@ +from __future__ import annotations + +from datetime import datetime, timedelta, timezone +from unittest.mock import AsyncMock + +import pytest + +from keboola_mcp_server.clients.data_science import DataScienceClient + + +@pytest.mark.asyncio +async def test_tail_app_logs_with_lines_calls_get_text_with_lines() -> None: + client = DataScienceClient.create('https://api.example.com', token=None) + client.get_text = AsyncMock(return_value='LOGS') # type: ignore[assignment] + + result = await client.tail_app_logs('app-123', since=None, lines=5) + + assert result == 'LOGS' + client.get_text.assert_awaited_once_with(endpoint='apps/app-123/logs/tail', params={'lines': 5}) + + +@pytest.mark.asyncio +async def test_tail_app_logs_with_lines_minimum_enforced() -> None: + client = DataScienceClient.create('https://api.example.com', token=None) + client.get_text = AsyncMock(return_value='LOGS') # type: ignore[assignment] + + _ = await client.tail_app_logs('app-123', since=None, lines=0) + + client.get_text.assert_awaited_once_with(endpoint='apps/app-123/logs/tail', params={'lines': 1}) + + +@pytest.mark.asyncio +async def test_tail_app_logs_with_since_calls_get_text_with_since_param() -> None: + client = DataScienceClient.create('https://api.example.com', token=None) + client.get_text = AsyncMock(return_value='LOGS') # type: ignore[assignment] + + since = datetime.now(timezone.utc) - timedelta(days=1) + result = await client.tail_app_logs('app-xyz', since=since, lines=None) + + assert result == 'LOGS' + client.get_text.assert_awaited_once_with( + endpoint='apps/app-xyz/logs/tail', params={'since': since.isoformat(timespec='microseconds')} + ) + + +@pytest.mark.asyncio +async def test_tail_app_logs_raises_when_both_since_and_lines_provided() -> None: + client = DataScienceClient.create('https://api.example.com', token=None) + + with pytest.raises(ValueError, match='You cannot use both "since" and "lines"'): + await client.tail_app_logs('app-123', since=datetime.now(timezone.utc), lines=10) + + +@pytest.mark.asyncio +async def test_tail_app_logs_raises_when_neither_param_provided() -> None: + client = DataScienceClient.create('https://api.example.com', token=None) + + with pytest.raises(ValueError, match='Either "since" or "lines" must be provided.'): + await client.tail_app_logs('app-123', since=None, lines=None) From b503441eca6e2e3aade3c8152612bafa8850a62c Mon Sep 17 00:00:00 2001 From: mariankrotil Date: Tue, 19 Aug 2025 15:43:29 +0200 Subject: [PATCH 29/35] AI-1343 test: improve tests with asynccontext for initial data app tests --- integtests/clients/test_data_science.py | 121 +++++++++++++----------- 1 file changed, 66 insertions(+), 55 deletions(-) diff --git a/integtests/clients/test_data_science.py b/integtests/clients/test_data_science.py index 05edf02a..ef0687e2 100644 --- a/integtests/clients/test_data_science.py +++ b/integtests/clients/test_data_science.py @@ -1,9 +1,11 @@ import logging +from typing import AsyncGenerator import pytest +import pytest_asyncio from keboola_mcp_server.clients.client import DATA_APP_COMPONENT_ID, KeboolaClient -from keboola_mcp_server.clients.data_science import DataAppResponse, DataScienceClient +from keboola_mcp_server.clients.data_science import DataAppConfig, DataAppResponse, DataScienceClient LOG = logging.getLogger(__name__) @@ -43,63 +45,72 @@ def ds_client(keboola_client: KeboolaClient) -> DataScienceClient: return keboola_client.data_science_client +@pytest_asyncio.fixture +async def initial_data_app(ds_client: DataScienceClient, unique_id: str) -> AsyncGenerator[DataAppResponse, None]: + data_app: DataAppResponse | None = None + try: + slug = f'test-app-{unique_id}' + config = DataAppConfig.model_validate( + {'parameters': _minimal_parameters(slug), 'authorization': _public_access_authorization()} + ) + data_app = await ds_client.create_data_app( + name=f'IntegTest {slug}', + description='Created by integration tests', + configuration=config, + ) + assert isinstance(data_app, DataAppResponse) + yield data_app + finally: + if data_app: + for _ in range(2): # Delete configuration 2 times (from storage and then from temporal bin) + try: + await ds_client.delete_data_app(data_app.id) + except Exception as e: + LOG.exception(f'Failed to delete data app: {e}') + pass + + @pytest.mark.asyncio async def test_create_and_fetch_data_app( - ds_client: DataScienceClient, unique_id: str, keboola_client: KeboolaClient + ds_client: DataScienceClient, initial_data_app: DataAppResponse, keboola_client: KeboolaClient ) -> None: """Test creating a data app and fetching it from detail and list endpoints""" - slug = f'test-app-{unique_id}' - created: DataAppResponse = await ds_client.create_data_app( - name=f'IntegTest {slug}', - description='Created by integration tests', - parameters=_minimal_parameters(slug), - authorization=_public_access_authorization(), + # Check if the created data app is valid + created = initial_data_app + assert isinstance(created, DataAppResponse) + assert created.id + assert created.type == 'streamlit' + assert created.component_id == DATA_APP_COMPONENT_ID + + # Deploy the data app + response = await ds_client.deploy_data_app(created.id, created.config_version) + assert response.id == created.id + + # Fetch the data app from data science + fethced_ds = await ds_client.get_data_app(created.id) + assert fethced_ds.id == created.id + assert fethced_ds.type == created.type + assert fethced_ds.component_id == created.component_id + assert fethced_ds.project_id == created.project_id + assert fethced_ds.config_id == created.config_id + assert fethced_ds.config_version == created.config_version + + # Fetch the data app config from storage + fetched_s = await keboola_client.storage_client.configuration_detail( + component_id=DATA_APP_COMPONENT_ID, + configuration_id=created.config_id, ) - try: - # Check if the created data app is valid - assert isinstance(created, DataAppResponse) - assert created.id - assert created.type == 'streamlit' - assert created.component_id == DATA_APP_COMPONENT_ID - - # Deploy the data app - response = await ds_client.deploy_data_app(created.id, created.config_version) - assert response.id == created.id - - # Fetch the data app from data science - fethced_ds = await ds_client.get_data_app(created.id) - assert fethced_ds.id == created.id - assert fethced_ds.type == created.type - assert fethced_ds.component_id == created.component_id - assert fethced_ds.project_id == created.project_id - assert fethced_ds.config_id == created.config_id - assert fethced_ds.config_version == created.config_version - - # Fetch the data app config from storage - fetched_s = await keboola_client.storage_client.configuration_detail( - component_id=DATA_APP_COMPONENT_ID, - configuration_id=created.config_id, - ) - - # check if the data app ids are the same (data app from data science and config from storage) - assert 'configuration' in fetched_s - assert isinstance(fetched_s['configuration'], dict) - assert 'parameters' in fetched_s['configuration'] - assert isinstance(fetched_s['configuration']['parameters'], dict) - assert 'id' in fetched_s['configuration']['parameters'] - assert fethced_ds.id == fetched_s['configuration']['parameters']['id'] - - # Fetch the all data apps and check if the created data app is in the list - data_apps = await ds_client.list_data_apps() - assert isinstance(data_apps, list) - assert len(data_apps) > 0 - assert any(app.id == created.id for app in data_apps) - - finally: - for _ in range(2): # Delete configuration 2 times (from storage and then from temporal bin) - try: - await ds_client.delete_data_app(created.id) - except Exception as e: - LOG.exception(f'Failed to delete data app: {e}') - pass + # check if the data app ids are the same (data app from data science and config from storage) + assert 'configuration' in fetched_s + assert isinstance(fetched_s['configuration'], dict) + assert 'parameters' in fetched_s['configuration'] + assert isinstance(fetched_s['configuration']['parameters'], dict) + assert 'id' in fetched_s['configuration']['parameters'] + assert fethced_ds.id == fetched_s['configuration']['parameters']['id'] + + # Fetch the all data apps and check if the created data app is in the list + data_apps = await ds_client.list_data_apps() + assert isinstance(data_apps, list) + assert len(data_apps) > 0 + assert any(app.id == created.id for app in data_apps) From b48ff3734c63882fc0bcf254b63793b392addf35 Mon Sep 17 00:00:00 2001 From: mariankrotil Date: Tue, 19 Aug 2025 15:43:53 +0200 Subject: [PATCH 30/35] AI-1343 fix: increase waiting time for global search --- integtests/conftest.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/integtests/conftest.py b/integtests/conftest.py index 7c859ce9..aa8bc4a3 100644 --- a/integtests/conftest.py +++ b/integtests/conftest.py @@ -248,7 +248,7 @@ def keboola_project(env_init: bool, storage_api_token: str, storage_api_url: str if 'global-search' in token_info['owner'].get('fetaures', []): # Give the global search time to catch up on the changes done in the testing project. # See https://help.keboola.com/management/global-search/#limitations for moe info. - time.sleep(5) + time.sleep(10) LOG.info(f'Test setup for project {project_id} complete') yield ProjectDef(project_id=project_id, buckets=buckets, tables=tables, configs=configs) From 4457d0c4458cd26050a9e57929493b85bc6d6948 Mon Sep 17 00:00:00 2001 From: mariankrotil Date: Tue, 19 Aug 2025 15:44:29 +0200 Subject: [PATCH 31/35] AI-1343 refactor: update parameter types for encrypt method wrt reviews --- src/keboola_mcp_server/clients/encryption.py | 29 ++++++++++---------- 1 file changed, 14 insertions(+), 15 deletions(-) diff --git a/src/keboola_mcp_server/clients/encryption.py b/src/keboola_mcp_server/clients/encryption.py index a9e478d7..d54b94b7 100644 --- a/src/keboola_mcp_server/clients/encryption.py +++ b/src/keboola_mcp_server/clients/encryption.py @@ -1,17 +1,11 @@ -from typing import Any, Optional +from typing import Any, Optional, cast -from keboola_mcp_server.clients.base import KeboolaServiceClient, RawKeboolaClient +from keboola_mcp_server.clients.base import JsonDict, KeboolaServiceClient, RawKeboolaClient +EncValue = str | JsonDict -class EncryptionClient(KeboolaServiceClient): - - def __init__(self, raw_client: RawKeboolaClient) -> None: - """ - Creates an EncryptionClient from a RawKeboolaClient. - :param raw_client: The raw client to use - """ - super().__init__(raw_client=raw_client) +class EncryptionClient(KeboolaServiceClient): @property def base_api_url(self) -> str: @@ -41,8 +35,13 @@ def create( ) async def encrypt( - self, value: Any, project_id: Optional[str], component_id: Optional[str] = None, config_id: Optional[str] = None - ) -> Any: + self, + value: EncValue, + *, + project_id: Optional[str] = None, + component_id: Optional[str] = None, + config_id: Optional[str] = None, + ) -> EncValue: """ Encrypt a value using the encryption service, returns encrypted value. Parameters are optional and the ciphers created by the service are dependent on those parameters when decrypting. Decryption is done automatically @@ -60,7 +59,7 @@ async def encrypt( """ if component_id and project_id is None: raise ValueError('project_id is required if component_id is provided') - if config_id and not component_id: + if config_id and component_id is None: raise ValueError('component_id is required if config_id is provided') params = { @@ -72,6 +71,6 @@ async def encrypt( response = await self.raw_client.post( endpoint='encrypt', params=params, - data=value, + data=cast(dict[str, Any], value), ) - return response + return cast(EncValue, response) From 812a341f654e42db7914f06ac1ad6e16e7c40d2e Mon Sep 17 00:00:00 2001 From: mariankrotil Date: Tue, 19 Aug 2025 15:44:54 +0200 Subject: [PATCH 32/35] AI-1343 test: unskip the global search test --- integtests/clients/test_client.py | 1 - 1 file changed, 1 deletion(-) diff --git a/integtests/clients/test_client.py b/integtests/clients/test_client.py index 8afe28e1..6b91d73d 100644 --- a/integtests/clients/test_client.py +++ b/integtests/clients/test_client.py @@ -26,7 +26,6 @@ async def test_global_search(self, storage_client: AsyncStorageClient): assert ret.by_project == {} @pytest.mark.asyncio - @pytest.mark.skip(reason='Writing to the global search index takes too long, making the test unstable.') 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') From a9359bf4e75045e506a39c3d76f8185c24b0af33 Mon Sep 17 00:00:00 2001 From: mariankrotil Date: Tue, 19 Aug 2025 15:53:12 +0200 Subject: [PATCH 33/35] AI-1343 refactor: remove async annot for non-async test, improve test naming --- integtests/clients/test_encryption.py | 3 +-- src/keboola_mcp_server/clients/encryption.py | 4 ---- 2 files changed, 1 insertion(+), 6 deletions(-) diff --git a/integtests/clients/test_encryption.py b/integtests/clients/test_encryption.py index e04aa6ae..00a822a2 100644 --- a/integtests/clients/test_encryption.py +++ b/integtests/clients/test_encryption.py @@ -5,8 +5,7 @@ from keboola_mcp_server.clients.client import DATA_APP_COMPONENT_ID, KeboolaClient -@pytest.mark.asyncio -async def test_encrypt_without_project_id(keboola_client: KeboolaClient) -> None: +def test_client_does_not_send_authorization_headers(keboola_client: KeboolaClient) -> None: """Check that the encryption client does not send any authorization headers.""" assert 'Authorization' not in keboola_client.encryption_client.raw_client.headers assert 'X-StorageAPI-Token' not in keboola_client.encryption_client.raw_client.headers diff --git a/src/keboola_mcp_server/clients/encryption.py b/src/keboola_mcp_server/clients/encryption.py index d54b94b7..740d076f 100644 --- a/src/keboola_mcp_server/clients/encryption.py +++ b/src/keboola_mcp_server/clients/encryption.py @@ -7,10 +7,6 @@ class EncryptionClient(KeboolaServiceClient): - @property - def base_api_url(self) -> str: - return self.raw_client.base_api_url - @classmethod def create( cls, From a16b0def7c52a37a4d5e2e71fe117cdc83388129 Mon Sep 17 00:00:00 2001 From: mariankrotil Date: Tue, 19 Aug 2025 15:59:58 +0200 Subject: [PATCH 34/35] AI-1343 refactor: add shortcut imports --- src/keboola_mcp_server/clients/__init__.py | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/src/keboola_mcp_server/clients/__init__.py b/src/keboola_mcp_server/clients/__init__.py index e69de29b..aaad313d 100644 --- a/src/keboola_mcp_server/clients/__init__.py +++ b/src/keboola_mcp_server/clients/__init__.py @@ -0,0 +1,16 @@ +from keboola_mcp_server.clients.ai_service import AIServiceClient +from keboola_mcp_server.clients.base import KeboolaServiceClient, RawKeboolaClient +from keboola_mcp_server.clients.client import KeboolaClient +from keboola_mcp_server.clients.encryption import EncryptionClient +from keboola_mcp_server.clients.jobs_queue import JobsQueueClient +from keboola_mcp_server.clients.storage import AsyncStorageClient + +__all__ = [ + 'KeboolaClient', + 'EncryptionClient', + 'AsyncStorageClient', + 'AIServiceClient', + 'JobsQueueClient', + 'RawKeboolaClient', + 'KeboolaServiceClient', +] From 558a387ab1546e94d0977bf21833bdbeb1464f73 Mon Sep 17 00:00:00 2001 From: mariankrotil Date: Wed, 20 Aug 2025 08:57:57 +0200 Subject: [PATCH 35/35] AI-1343 refactor: change Optional[s] to s | None in clients, ref code --- integtests/clients/test_data_science.py | 1 - .../clients/data_science.py | 30 +++++++++---------- src/keboola_mcp_server/clients/encryption.py | 10 +++---- 3 files changed, 19 insertions(+), 22 deletions(-) diff --git a/integtests/clients/test_data_science.py b/integtests/clients/test_data_science.py index ef0687e2..171bfdbf 100644 --- a/integtests/clients/test_data_science.py +++ b/integtests/clients/test_data_science.py @@ -67,7 +67,6 @@ async def initial_data_app(ds_client: DataScienceClient, unique_id: str) -> Asyn await ds_client.delete_data_app(data_app.id) except Exception as e: LOG.exception(f'Failed to delete data app: {e}') - pass @pytest.mark.asyncio diff --git a/src/keboola_mcp_server/clients/data_science.py b/src/keboola_mcp_server/clients/data_science.py index 12998bfa..b9273601 100644 --- a/src/keboola_mcp_server/clients/data_science.py +++ b/src/keboola_mcp_server/clients/data_science.py @@ -1,6 +1,6 @@ import logging from datetime import datetime -from typing import Any, Optional, cast +from typing import Any, cast from pydantic import AliasChoices, BaseModel, Field @@ -15,9 +15,7 @@ class DataAppResponse(BaseModel): component_id: str = Field( validation_alias=AliasChoices('componentId', 'component_id'), description='The component ID' ) - branch_id: Optional[str] = Field( - validation_alias=AliasChoices('branchId', 'branch_id'), description='The branch ID' - ) + branch_id: str | None = Field(validation_alias=AliasChoices('branchId', 'branch_id'), description='The branch ID') config_id: str = Field( validation_alias=AliasChoices('configId', 'config_id'), description='The component config ID' ) @@ -29,24 +27,24 @@ class DataAppResponse(BaseModel): desired_state: str = Field( validation_alias=AliasChoices('desiredState', 'desired_state'), description='The desired state' ) - last_request_timestamp: Optional[str] = Field( + last_request_timestamp: str | None = Field( validation_alias=AliasChoices('lastRequestTimestamp', 'last_request_timestamp'), default=None, description='The last request timestamp', ) - last_start_timestamp: Optional[str] = Field( + last_start_timestamp: str | None = Field( validation_alias=AliasChoices('lastStartTimestamp', 'last_start_timestamp'), default=None, description='The last start timestamp', ) - url: Optional[str] = Field( + url: str | None = Field( validation_alias=AliasChoices('url', 'url'), description='The URL of the running data app', default=None ) auto_suspend_after_seconds: int = Field( validation_alias=AliasChoices('autoSuspendAfterSeconds', 'auto_suspend_after_seconds'), description='The auto suspend after seconds', ) - size: Optional[str] = Field( + size: str | None = Field( validation_alias=AliasChoices('size', 'size'), description='The size of the data app', default=None ) @@ -65,7 +63,7 @@ class DataApp(BaseModel): 'file content' ) ) - secrets: Optional[dict[str, str]] = Field(description='The secrets of the data app', default=None) + secrets: dict[str, str] | None = Field(description='The secrets of the data app', default=None) size: str = Field(description='The size of the data app') auto_suspend_after_seconds: int = Field( @@ -78,9 +76,9 @@ class DataApp(BaseModel): serialization_alias='dataApp', validation_alias=AliasChoices('dataApp', 'data_app'), ) - id: Optional[str] = Field(description='The id of the data app', default=None) - script: Optional[list[str]] = Field(description='The script of the data app', default=None) - packages: Optional[list[str]] = Field( + id: str | None = Field(description='The id of the data app', default=None) + script: list[str] | None = Field(description='The script of the data app', default=None) + packages: list[str] | None = Field( description='The python packages needed to be installed in the data app', default=None ) @@ -102,7 +100,7 @@ class DataScienceClient(KeboolaServiceClient): def create( cls, root_url: str, - token: Optional[str], + token: str | None, headers: dict[str, Any] | None = None, ) -> 'DataScienceClient': """ @@ -180,7 +178,7 @@ async def create_data_app( name: str, description: str, configuration: DataAppConfig, - branch_id: Optional[str] = None, + branch_id: str | None = None, ) -> DataAppResponse: """ Create a data app from a simplified config used in the MCP server. @@ -218,8 +216,8 @@ async def tail_app_logs( self, app_id: str, *, - since: Optional[datetime], - lines: Optional[int], + since: datetime | None, + lines: int | None, ) -> str: """ Tail application logs. Either `since` or `lines` must be provided but not both at the same time. diff --git a/src/keboola_mcp_server/clients/encryption.py b/src/keboola_mcp_server/clients/encryption.py index 740d076f..98c9c404 100644 --- a/src/keboola_mcp_server/clients/encryption.py +++ b/src/keboola_mcp_server/clients/encryption.py @@ -1,4 +1,4 @@ -from typing import Any, Optional, cast +from typing import Any, cast from keboola_mcp_server.clients.base import JsonDict, KeboolaServiceClient, RawKeboolaClient @@ -11,7 +11,7 @@ class EncryptionClient(KeboolaServiceClient): def create( cls, root_url: str, - token: Optional[str] = None, + token: str | None = None, headers: dict[str, Any] | None = None, ) -> 'EncryptionClient': """ @@ -34,9 +34,9 @@ async def encrypt( self, value: EncValue, *, - project_id: Optional[str] = None, - component_id: Optional[str] = None, - config_id: Optional[str] = None, + project_id: str | None = None, + component_id: str | None = None, + config_id: str | None = None, ) -> EncValue: """ Encrypt a value using the encryption service, returns encrypted value. Parameters are optional and the ciphers