diff --git a/README.md b/README.md index bee7d0788..f7cdcc3c8 100644 --- a/README.md +++ b/README.md @@ -159,6 +159,12 @@ This identifies your workspace in Keboola and is used for SQL queries. However, **Note**: KBC_WORKSPACE_SCHEMA is called Dataset Name in BigQuery workspaces, you simply click connect and copy the Dataset Name +### KBC_WORKSPACE_ID + +Pins queries to one specific, already-existing workspace by its ID instead of the schema-based lookup above, and takes precedence over `KBC_WORKSPACE_SCHEMA` when both are set. This is the option a Data App / kai-agent caller supplies, as the `X-Workspace-Id` header, so that Kai embedded in that app queries only through its own workspace. + +Set via the `KBC_WORKSPACE_ID` environment variable, the `--workspace-id` CLI flag, or (per-request, for multi-user deployments) the `X-Workspace-Id` header. + ### KBC_STORAGE_API_URL (Keboola Region) Your Keboola Region API URL depends on your deployment region. You can determine your region by looking at the URL in your browser when logged into your Keboola project: @@ -438,7 +444,7 @@ For a complete list of available tools with detailed descriptions, parameters, a | Issue | Solution | |-------|----------| | **Authentication Errors** | Verify `KBC_STORAGE_TOKEN` is valid | -| **Workspace Issues** | Confirm `KBC_WORKSPACE_SCHEMA` is correct | +| **Workspace Issues** | Confirm `KBC_WORKSPACE_SCHEMA` / `KBC_WORKSPACE_ID` is correct | | **Connection Timeout** | Check network connectivity | ## Development diff --git a/pyproject.toml b/pyproject.toml index 5698f3616..ee9192ff7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "hatchling.build" [project] name = "keboola-mcp-server" -version = "1.75.4" +version = "1.76.0" description = "MCP server for interacting with Keboola Connection" readme = "README.md" requires-python = ">=3.10" diff --git a/src/keboola_mcp_server/cli.py b/src/keboola_mcp_server/cli.py index 21f56be9c..9daf091ec 100644 --- a/src/keboola_mcp_server/cli.py +++ b/src/keboola_mcp_server/cli.py @@ -53,6 +53,7 @@ def parse_args(args: list[str] | None = None) -> argparse.Namespace: ) parser.add_argument('--storage-token', metavar='STR', help='Keboola Storage API token.') parser.add_argument('--workspace-schema', metavar='STR', help='Keboola Storage API workspace schema.') + parser.add_argument('--workspace-id', metavar='STR', help='Keboola Storage API workspace ID.') parser.add_argument('--host', default='localhost', metavar='STR', help='The host to listen on.') parser.add_argument('--port', type=int, default=8000, metavar='INT', help='The port to listen on.') parser.add_argument('--log-config', type=pathlib.Path, metavar='PATH', help='Logging config file.') @@ -129,6 +130,7 @@ async def run_server(args: list[str] | None = None) -> None: storage_api_url=parsed_args.api_url, storage_token=parsed_args.storage_token, workspace_schema=parsed_args.workspace_schema, + workspace_id=parsed_args.workspace_id, ) try: diff --git a/src/keboola_mcp_server/config.py b/src/keboola_mcp_server/config.py index 6168aeb55..46424f9dc 100644 --- a/src/keboola_mcp_server/config.py +++ b/src/keboola_mcp_server/config.py @@ -27,6 +27,18 @@ class Config: """The branch ID to access the storage API using the MCP tools.""" workspace_schema: str | None = None """Workspace schema to access the buckets, tables and execute sql queries.""" + workspace_id: str | None = field(default=None, metadata={'empty_means_absent': True, 'require_prefix': True}) + """Workspace ID to access the buckets, tables and execute sql queries (e.g. a Data App's own + workspace, supplied per-request via the 'X-Workspace-Id' header). Takes precedence over + `workspace_schema` when both are set. + + `require_prefix` is set because the bare `WORKSPACE_ID` env var is what Keboola injects into + Data App containers -- without it, that variable would pin every session on such a server. + `empty_means_absent` is set so an unset header template (`X-Workspace-Id:`) forwarded as an + empty string is not mistaken for an explicit pin to override a server-side default with. This + intentionally does NOT apply to `workspace_schema` (an empty `X-Workspace-Schema` header must + keep clearing it back to the MCP-managed workspace, same as `branch_id` below) nor to most + other fields -- it is opt-in per field precisely to avoid that kind of regression.""" oauth_client_id: str | None = None """OAuth client ID registered in the Keboola OAuth Server.""" oauth_client_secret: str | None = None @@ -67,6 +79,9 @@ def __post_init__(self) -> None: if self.branch_id is not None and self.branch_id.lower() in ['', 'none', 'null', 'default', 'production']: object.__setattr__(self, 'branch_id', None) + if self.workspace_id is not None and not self.workspace_id.isdigit(): + raise ValueError(f'Invalid workspace_id: {self.workspace_id!r}') + @staticmethod def _normalize(name: str) -> str: """Removes dashes and underscores from the input string and turns it into lowercase.""" @@ -79,19 +94,27 @@ def _read_options(cls, d: Mapping[str, str]) -> Mapping[str, Any]: for f in dataclasses.fields(cls): field_names = [f.name] + f.metadata.get('aliases', []) + require_prefix = f.metadata.get('require_prefix', False) + empty_means_absent = f.metadata.get('empty_means_absent', False) + for name in field_names: value: str | None = _NO_VALUE_MARKER - - if (dict_name := cls._normalize(name)) in data: - value = data[dict_name] - - elif (dict_name := cls._normalize(f'KBC_{name}')) in data: - # environment variables start with KBC_ - value = data[dict_name] - - elif (dict_name := cls._normalize(f'X-{name}')) in data: - # HTTP headers start with X- - value = data[dict_name] + # `require_prefix` skips the bare field name -- only the KBC_/X- prefixed forms + # count as "provided". Needed for fields whose bare name collides with a + # variable set by something other than this server's own config (see + # `workspace_id`'s docstring). + candidates = (f'KBC_{name}', f'X-{name}') if require_prefix else (name, f'KBC_{name}', f'X-{name}') + + for candidate in candidates: + if (dict_name := cls._normalize(candidate)) in data: + candidate_value = data[dict_name] + # An empty value means "not provided" for opted-in fields only -- an + # unset header template (e.g. `X-Workspace-Id:`) must not be mistaken + # for an explicit request to override a server-side default with ''. + if empty_means_absent and candidate_value == '': + continue + value = candidate_value + break if value is not _NO_VALUE_MARKER: if f.type == (bool | None): @@ -118,8 +141,18 @@ def replace_by(self, d: Mapping[str, str]) -> 'Config': Creates new `Config` instance from the existing one by replacing the values from the input mapping. The keys in the input mapping can either be the names of the fields in `Config` class or their uppercase variant prefixed with 'KBC_'. + + A malformed `workspace_id` (the only field `__post_init__` validates) degrades to "not + provided" rather than raising: `d` is untrusted per-request input (an HTTP header), so a + junk value from a client should drop the pin, not turn into an unhandled server error. """ - return dataclasses.replace(self, **self._read_options(d)) + options = self._read_options(d) + try: + return dataclasses.replace(self, **options) + except ValueError as e: + LOG.warning(f'Ignoring invalid request header value(s): {e}') + options.pop('workspace_id', None) + return dataclasses.replace(self, **options) def __repr__(self) -> str: params: list[str] = [] diff --git a/src/keboola_mcp_server/mcp.py b/src/keboola_mcp_server/mcp.py index c419e0fab..f77cb8951 100644 --- a/src/keboola_mcp_server/mcp.py +++ b/src/keboola_mcp_server/mcp.py @@ -288,8 +288,29 @@ def apply_request_config(cls, http_rq: Request, config: Config, *, own_stack_sto :return: The configuration to use for this request. """ LOG.debug(f'Injecting headers: http_rq={http_rq}, headers={http_rq.headers}') + server_config = config config = config.replace_by(http_rq.headers) + # A workspace pin (`workspace_id`/`workspace_schema`) configured on the server fences + # off which project data a request can reach -- mirror the `own_stack_storage_api_url` + # reasoning above: a deployment that pinned itself must not be overridable by a header. + # A server with no pin of its own (the shared multi-tenant case, e.g. AJDA-3052's Data + # App flow) keeps taking the pin from the request, which is the only source it has. + if (server_config.workspace_id or server_config.workspace_schema) and ( + config.workspace_id != server_config.workspace_id + or config.workspace_schema != server_config.workspace_schema + ): + LOG.warning( + f'Ignoring the requested workspace pin (workspace_id={config.workspace_id!r}, ' + f'workspace_schema={config.workspace_schema!r}); this server is pinned to ' + f'workspace_id={server_config.workspace_id!r}, workspace_schema={server_config.workspace_schema!r}.' + ) + config = dataclasses.replace( + config, + workspace_id=server_config.workspace_id, + workspace_schema=server_config.workspace_schema, + ) + if own_stack_storage_api_url and not is_same_stack(config.storage_api_url, own_stack_storage_api_url): LOG.warning( f'Ignoring the requested Storage API URL "{config.storage_api_url}"; ' @@ -367,7 +388,10 @@ async def create_session_state( # therefore attaches the JWT only when the target is this server's own stack. kubernetes_token_path = os.environ.get('KBC_KUBERNETES_TOKEN_PATH') workspace_manager = await WorkspaceManager.create( - client, config.workspace_schema, kubernetes_token_path=kubernetes_token_path + client, + config.workspace_schema, + kubernetes_token_path=kubernetes_token_path, + workspace_id=config.workspace_id, ) state[WorkspaceManager.STATE_KEY] = workspace_manager LOG.info('Successfully initialized Storage API Workspace manager.') diff --git a/src/keboola_mcp_server/tools/data_apps.py b/src/keboola_mcp_server/tools/data_apps.py index 114b4c9dc..9e00d6156 100644 --- a/src/keboola_mcp_server/tools/data_apps.py +++ b/src/keboola_mcp_server/tools/data_apps.py @@ -552,9 +552,9 @@ async def modify_streamlit_data_app( links_manager = await ProjectLinksManager.from_client(client) project_id = await client.storage_client.project_id() - workspace_id = await workspace_manager.get_workspace_id() - sql_dialect = await workspace_manager.get_sql_dialect() - branch_id = await workspace_manager.get_branch_id() + workspace_id = await workspace_manager.get_data_app_workspace_id() + sql_dialect = await workspace_manager.get_data_app_sql_dialect() + branch_id = await workspace_manager.get_data_app_branch_id() secrets = _get_secrets( workspace_id=str(workspace_id), @@ -708,8 +708,8 @@ async def modify_streamlit_data_app_internal( folder: str | None = None, ) -> tuple[DataApp, JsonDict, dict | None]: secrets = _get_secrets( - workspace_id=str(await workspace_manager.get_workspace_id()), - branch_id=str(await workspace_manager.get_branch_id()), + workspace_id=str(await workspace_manager.get_data_app_workspace_id()), + branch_id=str(await workspace_manager.get_data_app_branch_id()), ) data_app = await _fetch_data_app(client, configuration_id=configuration_id, data_app_id=None) existing_config = data_app.configuration @@ -720,7 +720,7 @@ async def modify_streamlit_data_app_internal( packages, authentication_type, secrets, - await workspace_manager.get_sql_dialect(), + await workspace_manager.get_data_app_sql_dialect(), ) updated_config = cast( JsonDict, @@ -1096,7 +1096,7 @@ async def modify_python_js_data_app( legacy_secrets: dict[str, Any] | None = None if not has_storage_workspace: workspace_manager = WorkspaceManager.from_state(ctx.session.state) - legacy_secrets = {SECRET_WORKSPACE_ID: str(await workspace_manager.get_workspace_id())} + legacy_secrets = {SECRET_WORKSPACE_ID: str(await workspace_manager.get_data_app_workspace_id())} if configuration_id: # Update existing python-js data app diff --git a/src/keboola_mcp_server/workspace.py b/src/keboola_mcp_server/workspace.py index 0ec4ac7d3..c0c13f49a 100644 --- a/src/keboola_mcp_server/workspace.py +++ b/src/keboola_mcp_server/workspace.py @@ -544,7 +544,7 @@ def _format_error_message(self, message: str | None) -> str | None: return message -@dataclass(frozen=True) +@dataclass(frozen=True, repr=False) class _WspInfo: id: int schema: str @@ -552,6 +552,14 @@ class _WspInfo: credentials: str | None # the backend credentials; it can contain serialized JSON data readonly: bool | None + def __repr__(self) -> str: + # Redact `credentials` (the backend's credential blob, e.g. a service-account JSON for + # BigQuery) so a bare `LOG.info(f'... {info}')` anywhere can never leak it. + return ( + f'_WspInfo(id={self.id!r}, schema={self.schema!r}, backend={self.backend!r}, ' + f'credentials={"****" if self.credentials else None}, readonly={self.readonly!r})' + ) + @staticmethod def from_sapi_info(sapi_wsp_info: Mapping[str, Any]) -> '_WspInfo': _id = sapi_wsp_info.get('id') @@ -578,6 +586,7 @@ async def create( client: KeboolaClient, workspace_schema: str | None = None, kubernetes_token_path: str | None = None, + workspace_id: str | int | None = None, ) -> 'WorkspaceManager': # On projects with the `storage-branches` feature, each dev branch needs its own # workspace so the agent's queries (FQN paths, `query_data`) see that branch's @@ -588,15 +597,18 @@ async def create( # `has_storage_branches` already requires `branch_id is not None`, so the default # branch always takes the prod-client path. if await has_storage_branches(client): - return cls(client, workspace_schema, kubernetes_token_path=kubernetes_token_path) + return cls(client, workspace_schema, kubernetes_token_path=kubernetes_token_path, workspace_id=workspace_id) prod_client = await client.with_branch_id(None) - return cls(prod_client, workspace_schema, kubernetes_token_path=kubernetes_token_path) + return cls( + prod_client, workspace_schema, kubernetes_token_path=kubernetes_token_path, workspace_id=workspace_id + ) def __init__( self, client: KeboolaClient, workspace_schema: str | None = None, kubernetes_token_path: str | None = None, + workspace_id: str | int | None = None, ): """ Initializes the WorkspaceManager. @@ -611,12 +623,17 @@ def __init__( JWT as the X-Kubernetes-Authorization step-up header alongside the user's own token, so Connection can waive permissions the user's token lacks (e.g. read-only users). + :param workspace_id: The ID of the workspace to use (e.g. a Data App's own workspace). + Takes precedence over `workspace_schema` when both are set. """ self._client = client self._workspace_schema = workspace_schema + self._workspace_id = workspace_id self._kubernetes_token_path = kubernetes_token_path self._provisioning_client: AsyncStorageClient | None = None self._workspace: _Workspace | None = None + # Separate cache for the pin-agnostic managed workspace -- see `_get_managed_workspace`. + self._managed_workspace: _Workspace | None = None self._table_info_cache: dict[str, DbTableInfo] = {} async def _provisioning_storage_client(self) -> AsyncStorageClient: @@ -659,24 +676,56 @@ async def _find_ws_by_schema(self, schema: str) -> _WspInfo | None: return None - async def _find_ws_by_id(self, workspace_id: str | int) -> _WspInfo | None: - """Finds the workspace info by its ID.""" - + @staticmethod + async def _fetch_ws(client: KeboolaClient, workspace_id: str | int, *, strict: bool = False) -> _WspInfo | None: + """Fetches the workspace info by its ID from the given client, or None if the id does + not resolve on it. + + :param strict: if False (default), a 400/403 is treated the same as a 404: the id is + caller-supplied and unvalidated, so a malformed or inaccessible id means the same + thing -- "not usable" here. Pass True when the id is already known-good (e.g. a + workspace this server just created), so a permission or validation error on the + follow-up lookup is raised instead of being mistaken for "workspace doesn't exist". + """ try: - sapi_wsp_info = await self._client.storage_client.workspace_detail(workspace_id) - assert isinstance(sapi_wsp_info, dict) - wi = _WspInfo.from_sapi_info(sapi_wsp_info) # type: ignore[attr-defined] - - if wi.id and wi.backend and wi.schema: - return wi - else: - raise ValueError(f'Invalid workspace info: {sapi_wsp_info}') - + sapi_wsp_info = await client.storage_client.workspace_detail(workspace_id) except HTTPStatusError as e: - if e.response.status_code == 404: + not_found_codes = (404,) if strict else (400, 403, 404) + if e.response.status_code in not_found_codes: return None - else: - raise + raise + + assert isinstance(sapi_wsp_info, dict) + wi = _WspInfo.from_sapi_info(sapi_wsp_info) # type: ignore[attr-defined] + if wi.id and wi.backend and wi.schema: + return wi + raise ValueError(f'Invalid workspace info: {sapi_wsp_info}') + + async def _find_ws_by_id( + self, workspace_id: str | int, *, strict: bool = False + ) -> tuple[_WspInfo, KeboolaClient] | None: + """Finds the workspace info by its ID, together with the client it was actually resolved + on. + + Tries this manager's own (possibly branch-bound) client first, then -- since a + workspace is not necessarily tied to the branch this manager happens to be bound to + (e.g. a Data App's workspace is project-wide) -- falls back to the production-branch + client if that differs. `workspace_detail` is branch-scoped, so skipping this fallback + would 404 for a workspace that only lives on the other branch even though the id exists + in the project. The resolving client is returned rather than assigned to `self._client`, + so that resolving one pinned workspace this way does not silently rebind *this manager's* + client -- and with it every other lookup this manager makes (e.g. the MCP-managed + workspace via `_find_ws_in_branch`) -- onto a different branch. + + :param strict: see `_fetch_ws`. + """ + if info := await self._fetch_ws(self._client, workspace_id, strict=strict): + return info, self._client + if self._client.branch_id is not None: + prod_client = await self._client.with_branch_id(None) + if info := await self._fetch_ws(prod_client, workspace_id, strict=strict): + return info, prod_client + return None async def _find_ws_in_branch(self) -> _WspInfo | None: """Finds the shared read-only MCP workspace in the current branch. @@ -700,16 +749,18 @@ async def _find_ws_in_branch(self) -> _WspInfo | None: return None - async def _create_ws(self, *, timeout_sec: float = 300.0) -> _WspInfo | None: + async def _create_ws(self, *, timeout_sec: float = 300.0) -> tuple[_WspInfo, KeboolaClient] | None: """ - Creates a new workspace under a component configuration and returns its info. + Creates a new workspace under a component configuration and returns its info, together + with the client it was resolved on (see `_find_ws_by_id`). The workspace is created under the MCP_WORKSPACE_COMPONENT_ID component so that it is correctly attributed for billing. This method creates the configuration, creates the workspace under it, and cleans up the configuration on failure. :param timeout_sec: The number of seconds to wait for the workspace creation job to finish. - :return: The workspace info if the workspace was created successfully, None otherwise. + :return: The workspace info and resolving client if the workspace was created successfully, + None otherwise. """ # Verify token before creating workspace to ensure it has proper permissions @@ -784,7 +835,7 @@ async def _create_ws(self, *, timeout_sec: float = 300.0) -> _WspInfo | None: workspace_id = job_results['id'] LOG.info(f'Created workspace: {workspace_id}') - return await self._find_ws_by_id(workspace_id) + return await self._find_ws_by_id(workspace_id, strict=True) elif duration > timeout_sec: LOG.info(f'Workspace creation timed out after {duration:.2f} seconds.') @@ -794,11 +845,17 @@ async def _create_ws(self, *, timeout_sec: float = 300.0) -> _WspInfo | None: remaining_time = max(0.0, timeout_sec - duration) await asyncio.sleep(min(5.0, remaining_time)) - def _init_workspace(self, info: _WspInfo) -> _Workspace: - """Creates a new `Workspace` instance based on the workspace info.""" + def _init_workspace(self, info: _WspInfo, *, client: KeboolaClient | None = None) -> _Workspace: + """Creates a new `Workspace` instance based on the workspace info. + + :param client: the client the workspace's queries should run through; defaults to this + manager's own client. Pass the client the info was actually resolved on when it + differs (see `_find_ws_by_id`), without changing this manager's own client. + """ + client = client if client is not None else self._client if info.backend == 'snowflake': - return _SnowflakeWorkspace(workspace_id=info.id, schema=info.schema, client=self._client) + return _SnowflakeWorkspace(workspace_id=info.id, schema=info.schema, client=client) elif info.backend == 'bigquery': credentials = json.loads(info.credentials or '{}') @@ -807,7 +864,7 @@ def _init_workspace(self, info: _WspInfo) -> _Workspace: workspace_id=info.id, dataset_id=info.schema, project_id=project_id, - client=self._client, + client=client, ) else: @@ -817,17 +874,62 @@ def _init_workspace(self, info: _WspInfo) -> _Workspace: raise ValueError(f'Unexpected backend type "{info.backend}" in workspace: {info.schema}') async def _get_workspace(self) -> _Workspace: + """The workspace queries run against -- honors an explicit `workspace_id` pin (e.g. a + Data App's own workspace) ahead of `workspace_schema` and the default MCP-managed + workspace. Pin-aware: do not use this for anything attributed back to the *session's + own* identity rather than the query target (e.g. a data app's own persisted config) -- + use `_get_managed_workspace()`/`get_data_app_workspace_id()`/`get_data_app_branch_id()`/ + `get_data_app_sql_dialect()` for that instead. + """ if self._workspace: return self._workspace + if self._workspace_id is not None: + # use the workspace that was explicitly requested (e.g. a Data App's own workspace) + # this workspace must never be written to the default branch metadata + LOG.info(f'Looking up workspace by id: {self._workspace_id}') + result = await self._find_ws_by_id(self._workspace_id) + if result is None: + raise ValueError( + f'No Keboola workspace found: workspace_id={self._workspace_id}, branch_id={self._client.branch_id}' + ) + info, resolved_client = result + if not info.readonly: + # Every other resolution path enforces read-only storage access; this is a + # caller-supplied id, so it *could* widen `query_data` from read-only to + # read-write. Not hard-enforced here -- whether a Data App's platform-provisioned + # workspace is actually read-only is unconfirmed (needs checking against a real + # stack); enforcing blindly could break the feature outright. Warn for now; if it + # turns out these workspaces are read-only, promote this to a raise. Either way, + # `tools/sql.py` has no SELECT-only guard of its own, which is the real missing + # control and worth a follow-up regardless of this policy. + LOG.warning(f'Pinned workspace {self._workspace_id} has no read-only storage access.') + LOG.info(f'Found workspace: {info}') + self._workspace = self._init_workspace(info, client=resolved_client) + return self._workspace + + self._workspace = await self._get_managed_workspace() + return self._workspace + + async def _get_managed_workspace(self) -> _Workspace: + """The MCP-managed workspace, honoring `workspace_schema` but ignoring any `workspace_id` + pin: `workspace_schema` first, then the default per-branch MCP-managed workspace. + + Data apps persist this id into their own configuration (`SECRET_WORKSPACE_ID`), so it + must never be the caller-supplied `workspace_id` pin of whichever session happened to + create/update the app -- see `get_data_app_workspace_id()`/`get_data_app_branch_id()`. + """ + if self._managed_workspace: + return self._managed_workspace + if self._workspace_schema: # use the workspace that was explicitly requested # this workspace must never be written to the default branch metadata LOG.info(f'Looking up workspace by schema: {self._workspace_schema}') if info := await self._find_ws_by_schema(self._workspace_schema): LOG.info(f'Found workspace: {info}') - self._workspace = self._init_workspace(info) - return self._workspace + self._managed_workspace = self._init_workspace(info) + return self._managed_workspace else: raise ValueError( f'No Keboola workspace found or the workspace has no read-only storage access: ' @@ -838,19 +940,32 @@ async def _get_workspace(self) -> _Workspace: if info := await self._find_ws_in_branch(): # use the workspace that has already been created by the MCP server and noted to the branch LOG.info(f'Found workspace: {info}') - self._workspace = self._init_workspace(info) - return self._workspace + self._managed_workspace = self._init_workspace(info) + return self._managed_workspace + + if self._workspace_id is not None: + # A pinned session (e.g. a Data App session using `X-Workspace-Id`) has no reason to + # provision a brand new MCP-managed workspace of its own -- that workspace would be + # billed and owned by this session's token, not the caller who actually needs it, and + # provisioning can outright fail on a read-only token. Only an already-existing + # managed workspace is usable here; if none exists yet, that is a real "not + # available" case, not something to paper over by creating one. + raise ValueError( + f'No MCP-managed workspace exists for this project/branch, and one will not be ' + f'created for a session pinned to workspace_id={self._workspace_id}.' + ) # create a new workspace under the MCP component LOG.info('Creating workspace in the default branch.') - if info := await self._create_ws(): + if result := await self._create_ws(): # All tokens share the same read-only workspace, rediscovered by its # component id (see _find_ws_in_branch) — no branch-metadata pointer is # written, so no elevated metadata write is needed. Concurrent first-use # may create more than one workspace; that is acceptable, _find_ws_in_branch # returns the first match on the next lookup. - self._workspace = self._init_workspace(info) - return self._workspace + info, resolved_client = result + self._managed_workspace = self._init_workspace(info, client=resolved_client) + return self._managed_workspace else: raise ValueError('Failed to initialize Keboola Workspace.') @@ -898,3 +1013,28 @@ async def get_workspace_id(self) -> int: async def get_branch_id(self) -> str: workspace = await self._get_workspace() return await workspace.get_branch_id() + + async def get_data_app_workspace_id(self) -> int: + """The MCP-managed workspace's id, ignoring any `workspace_id` pin. + + Use this (not `get_workspace_id()`) for anything written into a data app's own + persisted configuration (e.g. `SECRET_WORKSPACE_ID`) -- a session pinned to one Data + App's workspace must not leak that id into a *different* app's config when creating or + updating it. + """ + workspace = await self._get_managed_workspace() + return workspace.id + + async def get_data_app_branch_id(self) -> str: + """The MCP-managed workspace's branch id, ignoring any `workspace_id` pin. See + `get_data_app_workspace_id()`.""" + workspace = await self._get_managed_workspace() + return await workspace.get_branch_id() + + async def get_data_app_sql_dialect(self) -> str: + """The MCP-managed workspace's SQL dialect, ignoring any `workspace_id` pin. See + `get_data_app_workspace_id()` -- the dialect baked into a data app's generated source + code must match the workspace whose id/branch are persisted into that same app, not + whichever workspace the creating/updating session happened to be pinned to.""" + workspace = await self._get_managed_workspace() + return workspace.get_sql_dialect() diff --git a/tests/test_cli.py b/tests/test_cli.py new file mode 100644 index 000000000..cd9a2bab4 --- /dev/null +++ b/tests/test_cli.py @@ -0,0 +1,11 @@ +from keboola_mcp_server.cli import parse_args + + +def test_parse_args_workspace_id() -> None: + parsed = parse_args(['--workspace-id', 'ws-123']) + assert parsed.workspace_id == 'ws-123' + + +def test_parse_args_workspace_id_defaults_to_none() -> None: + parsed = parse_args([]) + assert parsed.workspace_id is None diff --git a/tests/test_config.py b/tests/test_config.py index 22bbf9e9e..693b74922 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -26,6 +26,26 @@ class TestConfig: {'X-StorageApi_Token': 'foo', 'KBC_WORKSPACE_SCHEMA': 'bar'}, Config(storage_token='foo', workspace_schema='bar'), ), + ( + # workspace_id requires the KBC_/X- prefix -- a bare `workspace_id`/`WORKSPACE_ID` + # key must NOT be picked up (it collides with the variable Keboola injects into + # Data App containers). + {'storage_token': 'foo', 'workspace_id': '123'}, + Config(storage_token='foo'), + ), + ( + {'storage_token': 'foo', 'KBC_WORKSPACE_ID': '123'}, + Config(storage_token='foo', workspace_id='123'), + ), + ( + {'X-Workspace-Id': '123'}, + Config(workspace_id='123'), + ), + ( + # An empty value means "not provided" for workspace_id. + {'X-Workspace-Id': ''}, + Config(), + ), ( {'foo': 'bar', 'storage_api_url': 'http://nowhere'}, Config(storage_api_url='http://nowhere'), @@ -67,6 +87,45 @@ def test_from_dict(self, d: Mapping[str, str], expected: Config) -> None: (Config(branch_id='foo'), {'branch-id': 'Null'}, Config()), (Config(branch_id='foo'), {'branch-id': 'Default'}, Config()), (Config(branch_id='foo'), {'branch-id': 'pRoDuCtIoN'}, Config()), + ( + Config(), + {'storage_token': 'foo', 'workspace_id': '123'}, + Config(storage_token='foo'), + ), + ( + Config(), + {'storage_token': 'foo', 'KBC_WORKSPACE_ID': '123'}, + Config(storage_token='foo', workspace_id='123'), + ), + ( + # An empty header must not un-pin a server-configured workspace_id. + Config(workspace_id='999'), + {'X-Workspace-Id': ''}, + Config(workspace_id='999'), + ), + ( + # Unlike workspace_id, an empty workspace_schema header must still clear a + # server default (to the falsy '', same as pre-existing main behavior) -- this is + # the multi-user opt-out the README describes for X-Workspace-Schema, and must + # not regress into keeping the server's pin. + Config(workspace_schema='SERVER'), + {'X-Workspace-Schema': ''}, + Config(workspace_schema=''), + ), + ( + # A malformed workspace_id header must degrade to "not provided" rather than + # raising -- it is untrusted per-request input, so a junk value from a client + # should drop the pin, not turn into an unhandled server error. + Config(), + {'X-Workspace-Id': 'abc'}, + Config(), + ), + ( + # A malformed header must not clear an existing server-configured pin either. + Config(workspace_id='999'), + {'X-Workspace-Id': 'abc'}, + Config(workspace_id='999'), + ), ], ) def test_replace_by(self, orig: Config, d: Mapping[str, str], expected: Config) -> None: @@ -81,11 +140,16 @@ def test_no_token_password_in_repr(self) -> None: config = Config(storage_token='foo') assert str(config) == ( "Config(storage_api_url=None, storage_token='****', branch_id=None, workspace_schema=None, " + 'workspace_id=None, ' 'oauth_client_id=None, oauth_client_secret=None, ' 'oauth_server_url=None, oauth_scope=None, mcp_server_url=None, ' 'jwt_secret=None, bearer_token=None, conversation_id=None)' ) + def test_workspace_id_must_be_numeric(self) -> None: + with pytest.raises(ValueError, match='Invalid workspace_id'): + Config(workspace_id='not-a-valid-id') + @pytest.mark.parametrize( ('url', 'expected'), [ diff --git a/tests/test_mcp.py b/tests/test_mcp.py index 8525e5c0c..42ffc9518 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -797,3 +797,51 @@ def test_apply_request_config_pins_storage_api_url( # Only the Storage API URL is pinned; the other per-request headers keep working. assert applied.storage_token == headers.get('X-Storage-Api-Token', 'server-token') assert applied.branch_id == headers.get('X-Branch-Id') + + @pytest.mark.parametrize( + ('server_kwargs', 'headers', 'expected_workspace_id', 'expected_workspace_schema', 'expect_warning'), + [ + # server pinned by id: header asking for a different id/schema is ignored outright. + ({'workspace_id': '111'}, {'X-Workspace-Id': '222'}, '111', None, True), + ({'workspace_id': '111'}, {'X-Workspace-Schema': 'OTHER'}, '111', None, True), + ({'workspace_id': '111'}, {'X-Workspace-Id': '111'}, '111', None, False), + ({'workspace_id': '111'}, {}, '111', None, False), + # server pinned by schema: same treatment. + ({'workspace_schema': 'SERVER_SCHEMA'}, {'X-Workspace-Schema': 'OTHER'}, None, 'SERVER_SCHEMA', True), + # no server-side pin (the shared multi-tenant / AJDA-3052 Data App flow): the header + # is the only source, so it keeps working. + ({}, {'X-Workspace-Id': '222'}, '222', None, False), + ], + ids=[ + 'id_overridden', + 'id_overridden_by_schema_header', + 'id_unchanged', + 'id_no_header', + 'schema_overridden', + 'no_server_pin', + ], + ) + def test_apply_request_config_pins_workspace( + self, + caplog: pytest.LogCaptureFixture, + server_kwargs: dict[str, str], + headers: dict[str, str], + expected_workspace_id: str | None, + expected_workspace_schema: str | None, + expect_warning: bool, + ): + """A workspace pin configured on the server must be authoritative over a request header + -- mirroring the Storage API URL check above -- but a server with no pin of its own must + keep taking it from the request (AI-3669 review, workspace.py:836 thread).""" + config = Config(storage_token='server-token', **server_kwargs) + http_rq = MagicMock(spec=Request) + http_rq.headers = headers + http_rq.scope = {} + + with caplog.at_level('WARNING'): + applied = SessionStateMiddleware.apply_request_config(http_rq, config, own_stack_storage_api_url=None) + + assert applied.workspace_id == expected_workspace_id + assert applied.workspace_schema == expected_workspace_schema + warned = any('is pinned' in r.message for r in caplog.records) + assert warned is expect_warning diff --git a/tests/test_preview.py b/tests/test_preview.py index 68937cf23..6dfff4e1e 100644 --- a/tests/test_preview.py +++ b/tests/test_preview.py @@ -942,9 +942,9 @@ async def mock_encrypt(*args, **kwargs): # Mock WorkspaceManager mock_workspace_manager = mocker.AsyncMock() - mock_workspace_manager.get_workspace_id = mocker.AsyncMock(return_value=123) - mock_workspace_manager.get_branch_id = mocker.AsyncMock(return_value=456) - mock_workspace_manager.get_sql_dialect = mocker.AsyncMock(return_value='snowflake') + mock_workspace_manager.get_data_app_workspace_id = mocker.AsyncMock(return_value=123) + mock_workspace_manager.get_data_app_branch_id = mocker.AsyncMock(return_value=456) + mock_workspace_manager.get_data_app_sql_dialect = mocker.AsyncMock(return_value='snowflake') mocker.patch( 'keboola_mcp_server.preview.WorkspaceManager.from_state', diff --git a/tests/test_workspace.py b/tests/test_workspace.py index 766fee096..594030587 100644 --- a/tests/test_workspace.py +++ b/tests/test_workspace.py @@ -7,7 +7,7 @@ from keboola_mcp_server.clients.client import KeboolaClient from keboola_mcp_server.clients.query import QueryServiceClient -from keboola_mcp_server.workspace import JobSubmittedInfo, WorkspaceManager, _SnowflakeWorkspace +from keboola_mcp_server.workspace import JobSubmittedInfo, WorkspaceManager, _SnowflakeWorkspace, _WspInfo @pytest.mark.parametrize( @@ -173,22 +173,28 @@ async def test_workspace_creation_cleans_up_config_on_failure(): @pytest.mark.asyncio @pytest.mark.parametrize( - ('input_branch_id', 'has_sb_feature', 'workspace_schema', 'expected_bound_branch_id'), + ('input_branch_id', 'has_sb_feature', 'workspace_schema', 'workspace_id', 'expected_bound_branch_id'), [ # default branch: always production, regardless of feature - (None, True, None, None), - (None, False, None, None), + (None, True, None, None, None), + (None, False, None, None, None), # dev branch + storage-branches feature on: keep dev branch - ('456', True, None, '456'), + ('456', True, None, None, '456'), # dev branch without storage-branches (legacy): fall back to production - ('456', False, None, None), + ('456', False, None, None, None), # dev branch + storage-branches + explicit workspace_schema (KBC_WORKSPACE_SCHEMA): # stay branch-aware. The user is responsible for ensuring the named workspace # exists in the explicitly-bound branch — there is no carve-out for explicit schemas. - ('456', True, 'WORKSPACE_XYZ', '456'), + ('456', True, 'WORKSPACE_XYZ', None, '456'), # dev branch + legacy + explicit workspace_schema: still rebinds to production, # since branched workspaces don't exist on legacy projects. - ('456', False, 'WORKSPACE_XYZ', None), + ('456', False, 'WORKSPACE_XYZ', None, None), + # dev branch + storage-branches + explicit workspace_id: the pin must reach `cls(...)` + # on the storage-branches construction path too (mutation-tested: dropping + # `workspace_id=workspace_id` there leaves the rest of the suite green). + ('456', True, None, '123', '456'), + # default branch + explicit workspace_id: same, on the prod-client construction path. + (None, False, None, '123', None), ], ids=[ 'default_branch_with_sb', @@ -197,20 +203,24 @@ async def test_workspace_creation_cleans_up_config_on_failure(): 'dev_branch_legacy', 'dev_branch_with_sb_explicit_schema', 'dev_branch_legacy_explicit_schema', + 'dev_branch_with_sb_explicit_id', + 'default_branch_without_sb_explicit_id', ], ) async def test_workspace_manager_create_is_branch_aware( input_branch_id: str | None, has_sb_feature: bool, workspace_schema: str | None, + workspace_id: str | None, expected_bound_branch_id: str | None, ): """ WorkspaceManager.create() must keep the client on the dev branch only when the project has the `storage-branches` feature; otherwise it must rebind to the production branch. The rule applies uniformly whether the workspace is auto-managed or pinned via an - explicit `workspace_schema` (KBC_WORKSPACE_SCHEMA) — branch context is governed solely - by KBC_BRANCH_ID and the project's `storage-branches` feature. + explicit `workspace_schema` (KBC_WORKSPACE_SCHEMA) or `workspace_id` (KBC_WORKSPACE_ID) — + branch context is governed solely by KBC_BRANCH_ID and the project's `storage-branches` + feature, and both pins must reach `cls(...)` on either construction path. """ input_client = Mock(spec=KeboolaClient) input_client.branch_id = input_branch_id @@ -227,13 +237,15 @@ def _rebind(target_branch_id: str | None) -> Mock: input_client.with_branch_id = AsyncMock(side_effect=_rebind) - manager = await WorkspaceManager.create(input_client, workspace_schema=workspace_schema) + manager = await WorkspaceManager.create(input_client, workspace_schema=workspace_schema, workspace_id=workspace_id) # noinspection PyProtectedMember bound_client = manager._client assert bound_client.branch_id == expected_bound_branch_id # noinspection PyProtectedMember assert manager._workspace_schema == workspace_schema + # noinspection PyProtectedMember + assert manager._workspace_id == workspace_id # has_feature is only meaningful when the client is on a dev branch — the helper # short-circuits otherwise, so on the default branch we should not even ask. @@ -243,6 +255,259 @@ def _rebind(target_branch_id: str | None) -> Mock: input_client.has_feature.assert_awaited_once() +@pytest.mark.asyncio +async def test_get_workspace_resolves_by_id_when_set(): + """An explicit `workspace_id` (e.g. a Data App's own workspace) must be looked up by ID + and take precedence over `workspace_schema`, instead of falling back to the default + per-branch MCP-managed workspace.""" + mock_client = Mock(spec=KeboolaClient) + manager = WorkspaceManager(mock_client, workspace_schema='SOME_SCHEMA', workspace_id='123') + + ws_info = _WspInfo(id=123, schema='APP_SCHEMA', backend='snowflake', credentials=None, readonly=True) + manager._find_ws_by_id = AsyncMock(return_value=(ws_info, mock_client)) # type: ignore[method-assign] + manager._find_ws_by_schema = AsyncMock() # type: ignore[method-assign] + + workspace = await manager._get_workspace() + + assert workspace.id == 123 + manager._find_ws_by_id.assert_awaited_once_with('123') + manager._find_ws_by_schema.assert_not_called() + + +@pytest.mark.asyncio +async def test_get_workspace_raises_when_id_not_found(): + """A `workspace_id` that resolves to no workspace must fail loudly rather than silently + falling back to the default workspace — the caller asked for a specific workspace.""" + mock_client = Mock(spec=KeboolaClient) + manager = WorkspaceManager(mock_client, workspace_id='999') + manager._find_ws_by_id = AsyncMock(return_value=None) # type: ignore[method-assign] + + with pytest.raises(ValueError, match='workspace_id=999'): + await manager._get_workspace() + + +@pytest.mark.asyncio +async def test_get_workspace_warns_when_pinned_workspace_is_not_readonly(caplog: pytest.LogCaptureFixture) -> None: + """A `workspace_id` pin resolving to a writable workspace must not silently pass through -- + at minimum it needs a warning (whether a Data App's platform-provisioned workspace is + actually read-only is unconfirmed against a real stack; hard-enforcing here without knowing + that could break the feature outright -- see AI-3669 review, workspace.py:834 thread).""" + mock_client = Mock(spec=KeboolaClient) + manager = WorkspaceManager(mock_client, workspace_id='123') + writable_info = _WspInfo(id=123, schema='APP_SCHEMA', backend='snowflake', credentials=None, readonly=False) + manager._find_ws_by_id = AsyncMock(return_value=(writable_info, mock_client)) # type: ignore[method-assign] + + with caplog.at_level('WARNING'): + workspace = await manager._get_workspace() + + assert workspace.id == 123 + assert any('no read-only storage access' in r.message for r in caplog.records) + + +@pytest.mark.asyncio +async def test_get_data_app_workspace_id_ignores_the_pin(): + """`get_data_app_workspace_id()`/`get_data_app_branch_id()`/`get_data_app_sql_dialect()` must + resolve the *managed* workspace even when a session is pinned via `workspace_id` -- they feed + into `tools/data_apps.py`, which writes the result into a *different* data app's own + persisted `WORKSPACE_ID` secret and bakes the dialect into that same app's generated source + code. If they returned the pinned workspace, creating/updating data app B from a session + pinned to app A's workspace would permanently point B at A's workspace/dialect.""" + mock_client = Mock(spec=KeboolaClient) + manager = WorkspaceManager(mock_client, workspace_id='999') + + pinned_info = _WspInfo( + id=999, schema='PINNED_SCHEMA', backend='bigquery', credentials='{"project_id": "proj"}', readonly=True + ) + managed_info = _WspInfo(id=111, schema='MANAGED_SCHEMA', backend='snowflake', credentials=None, readonly=True) + manager._find_ws_by_id = AsyncMock(return_value=(pinned_info, mock_client)) # type: ignore[method-assign] + manager._find_ws_in_branch = AsyncMock(return_value=managed_info) # type: ignore[method-assign] + + assert await manager.get_data_app_workspace_id() == 111 + assert await manager.get_data_app_sql_dialect() == 'Snowflake' + pinned_workspace = await manager._get_workspace() + assert pinned_workspace.id == 999 + + +@pytest.mark.asyncio +async def test_get_managed_workspace_raises_instead_of_provisioning_when_pinned(): + """A session pinned via `workspace_id` (e.g. a Data App session) must not provision a brand + new MCP-managed workspace on demand when none exists yet -- that workspace would be billed + to this session's token, not whoever actually needs it, and provisioning can outright fail on + a read-only token. It should fail loudly instead of silently creating one.""" + mock_client = Mock(spec=KeboolaClient) + manager = WorkspaceManager(mock_client, workspace_id='999') + manager._find_ws_in_branch = AsyncMock(return_value=None) # type: ignore[method-assign] + manager._create_ws = AsyncMock() # type: ignore[method-assign] + + with pytest.raises(ValueError, match='workspace_id=999'): + await manager._get_managed_workspace() + manager._create_ws.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_found_workspace_repr_never_prints_credentials() -> None: + """`_WspInfo.credentials` holds the backend's credential blob (a service-account JSON for + BigQuery) -- its repr, which every `LOG.info(f'... {info}')` call site relies on, must + never include it.""" + secret = 'super-secret-service-account-json' + info = _WspInfo(id=123, schema='APP_SCHEMA', backend='snowflake', credentials=secret, readonly=True) + + assert secret not in repr(info) + assert 'credentials=****' in repr(info) + + +@pytest.mark.asyncio +async def test_find_ws_by_id_falls_back_to_production_branch() -> None: + """A Data App workspace is not tied to any particular branch, but `workspace_detail` is + branch-scoped -- a dev-branch session pinned to a workspace that lives on the default + branch must not 404 just because the first lookup used the wrong branch prefix.""" + dev_client = Mock(spec=KeboolaClient) + dev_client.branch_id = '456' + dev_client.storage_client = AsyncMock() + mock_response = Mock(spec=Response) + mock_response.status_code = 404 + mock_request = Mock(spec=Request) + dev_client.storage_client.workspace_detail = AsyncMock( + side_effect=HTTPStatusError('not found', request=mock_request, response=mock_response) + ) + + prod_client = Mock(spec=KeboolaClient) + prod_client.branch_id = None + prod_client.storage_client = AsyncMock() + prod_client.storage_client.workspace_detail = AsyncMock( + return_value={ + 'id': 123, + 'connection': {'backend': 'snowflake', 'schema': 'APP_SCHEMA', 'user': None}, + 'readOnlyStorageAccess': True, + } + ) + dev_client.with_branch_id = AsyncMock(return_value=prod_client) + + manager = WorkspaceManager(dev_client, workspace_id='123') + + result = await manager._find_ws_by_id('123') + + assert result is not None + info, resolved_client = result + assert info.id == 123 + dev_client.storage_client.workspace_detail.assert_awaited_once_with('123') + dev_client.with_branch_id.assert_awaited_once_with(None) + prod_client.storage_client.workspace_detail.assert_awaited_once_with('123') + # The workspace was only resolvable via the prod-branch client -- that client is returned for + # use on this one workspace, but this manager's own client must NOT change: mutating + # `manager._client` would also redirect every other lookup this manager makes (e.g. the + # MCP-managed workspace) onto the wrong branch. + assert resolved_client is prod_client + assert manager._client is dev_client + + +@pytest.mark.asyncio +async def test_pin_resolution_via_prod_fallback_does_not_leak_into_managed_lookup() -> None: + """Regression test for a real bug caught in review: resolving a `workspace_id` pin through + the prod-branch fallback must not affect the *managed* workspace lookup afterwards -- that + lookup (feeding `get_data_app_workspace_id()`/`get_data_app_branch_id()`/ + `get_data_app_sql_dialect()`, which get persisted into a Data App's own config) must keep + running on the manager's original (dev-branch) client, not the prod client the pin happened + to resolve on.""" + dev_client = Mock(spec=KeboolaClient) + dev_client.branch_id = '456' + dev_client.storage_client = AsyncMock() + mock_response = Mock(spec=Response) + mock_response.status_code = 404 + mock_request = Mock(spec=Request) + dev_client.storage_client.workspace_detail = AsyncMock( + side_effect=HTTPStatusError('not found', request=mock_request, response=mock_response) + ) + + prod_client = Mock(spec=KeboolaClient) + prod_client.branch_id = None + prod_client.storage_client = AsyncMock() + prod_client.storage_client.workspace_detail = AsyncMock( + return_value={ + 'id': 123, + 'connection': {'backend': 'snowflake', 'schema': 'PINNED_SCHEMA', 'user': None}, + 'readOnlyStorageAccess': True, + } + ) + dev_client.with_branch_id = AsyncMock(return_value=prod_client) + + manager = WorkspaceManager(dev_client, workspace_id='123') + managed_info = _WspInfo(id=111, schema='MANAGED_SCHEMA', backend='snowflake', credentials=None, readonly=True) + manager._find_ws_in_branch = AsyncMock(return_value=managed_info) # type: ignore[method-assign] + + pinned_workspace = await manager._get_workspace() + assert pinned_workspace.id == 123 + + managed_workspace_id = await manager.get_data_app_workspace_id() + + assert managed_workspace_id == 111 + assert manager._client is dev_client + + +@pytest.mark.asyncio +@pytest.mark.parametrize('status_code', [400, 403, 404]) +async def test_find_ws_by_id_treats_400_403_404_alike(status_code: int) -> None: + """The header value is unvalidated: a non-numeric id (400), an id the token can't read + (403), and a nonexistent id (404) must all mean the same thing -- "not usable" -- rather + than 400/403 bypassing the intended `ValueError` and surfacing a raw `HTTPStatusError`.""" + mock_client = Mock(spec=KeboolaClient) + mock_client.branch_id = None + mock_client.storage_client = AsyncMock() + mock_response = Mock(spec=Response) + mock_response.status_code = status_code + mock_request = Mock(spec=Request) + mock_client.storage_client.workspace_detail = AsyncMock( + side_effect=HTTPStatusError('error', request=mock_request, response=mock_response) + ) + + manager = WorkspaceManager(mock_client, workspace_id='not-a-valid-id') + + assert await manager._find_ws_by_id('not-a-valid-id') is None + + +@pytest.mark.asyncio +@pytest.mark.parametrize('status_code', [400, 403]) +async def test_find_ws_by_id_strict_reraises_400_403(status_code: int) -> None: + """`strict=True` (used for the id of a workspace this server just created, e.g. in + `_create_ws`) must only treat a 404 as "not found" -- a 400/403 on that known-good id is a + real failure (e.g. the creating token can't read what it just provisioned) that must not be + silently swallowed into a generic "workspace creation failed" error.""" + mock_client = Mock(spec=KeboolaClient) + mock_client.branch_id = None + mock_client.storage_client = AsyncMock() + mock_response = Mock(spec=Response) + mock_response.status_code = status_code + mock_request = Mock(spec=Request) + mock_client.storage_client.workspace_detail = AsyncMock( + side_effect=HTTPStatusError('error', request=mock_request, response=mock_response) + ) + + manager = WorkspaceManager(mock_client, workspace_id='123') + + with pytest.raises(HTTPStatusError): + await manager._find_ws_by_id('123', strict=True) + + +@pytest.mark.asyncio +async def test_find_ws_by_id_reraises_unexpected_status() -> None: + """A 5xx (or any other unexpected status) is a real failure, not an absent/inaccessible + workspace -- it must propagate rather than being swallowed into a misleading "not found".""" + mock_client = Mock(spec=KeboolaClient) + mock_client.branch_id = None + mock_client.storage_client = AsyncMock() + mock_response = Mock(spec=Response) + mock_response.status_code = 500 + mock_request = Mock(spec=Request) + mock_client.storage_client.workspace_detail = AsyncMock( + side_effect=HTTPStatusError('error', request=mock_request, response=mock_response) + ) + + manager = WorkspaceManager(mock_client, workspace_id='123') + + with pytest.raises(HTTPStatusError): + await manager._find_ws_by_id('123') + + def _make_snowflake_workspace_with_mocked_qs(job_id: str = 'job-abc-123') -> tuple[_SnowflakeWorkspace, AsyncMock]: """Builds a _SnowflakeWorkspace whose QueryServiceClient is fully mocked to run a one-row query end to end. diff --git a/tests/tools/test_data_apps.py b/tests/tools/test_data_apps.py index a5e6d9471..81b47bd09 100644 --- a/tests/tools/test_data_apps.py +++ b/tests/tools/test_data_apps.py @@ -814,9 +814,9 @@ async def test_modify_streamlit_data_app_folder( """Test folder metadata and change_summary hint for modify_streamlit_data_app (create and update paths).""" keboola_client = KeboolaClient.from_state(mcp_context_client.session.state) - workspace_manager.get_workspace_id = mocker.AsyncMock(return_value=1) - workspace_manager.get_sql_dialect = mocker.AsyncMock(return_value='snowflake') - workspace_manager.get_branch_id = mocker.AsyncMock(return_value='default') + workspace_manager.get_data_app_workspace_id = mocker.AsyncMock(return_value=1) + workspace_manager.get_data_app_sql_dialect = mocker.AsyncMock(return_value='snowflake') + workspace_manager.get_data_app_branch_id = mocker.AsyncMock(return_value='default') keboola_client.storage_client.project_id = mocker.AsyncMock(return_value='proj-1') @@ -936,9 +936,9 @@ async def test_modify_streamlit_data_app_partial_success_when_response_building_ regardless of failure point, and that the response wording still reflects the app state.""" keboola_client = KeboolaClient.from_state(mcp_context_client.session.state) - workspace_manager.get_workspace_id = mocker.AsyncMock(return_value=1) - workspace_manager.get_sql_dialect = mocker.AsyncMock(return_value='snowflake') - workspace_manager.get_branch_id = mocker.AsyncMock(return_value='default') + workspace_manager.get_data_app_workspace_id = mocker.AsyncMock(return_value=1) + workspace_manager.get_data_app_sql_dialect = mocker.AsyncMock(return_value='snowflake') + workspace_manager.get_data_app_branch_id = mocker.AsyncMock(return_value='default') keboola_client.storage_client.project_id = mocker.AsyncMock(return_value='proj-1') encrypted_config = { @@ -1079,9 +1079,9 @@ async def test_modify_streamlit_data_app_update_skips_metadata_when_version_miss (review hardening on AJDA-2852). The tool still returns a normal success.""" keboola_client = KeboolaClient.from_state(mcp_context_client.session.state) - workspace_manager.get_workspace_id = mocker.AsyncMock(return_value=1) - workspace_manager.get_sql_dialect = mocker.AsyncMock(return_value='snowflake') - workspace_manager.get_branch_id = mocker.AsyncMock(return_value='default') + workspace_manager.get_data_app_workspace_id = mocker.AsyncMock(return_value=1) + workspace_manager.get_data_app_sql_dialect = mocker.AsyncMock(return_value='snowflake') + workspace_manager.get_data_app_branch_id = mocker.AsyncMock(return_value='default') keboola_client.storage_client.project_id = mocker.AsyncMock(return_value='proj-1') encrypted_config = { @@ -1208,7 +1208,7 @@ async def test_modify_python_js_data_app_create_prod_derives_or_honors_slug( keboola_client = KeboolaClient.from_state(mcp_context_client.session.state) keboola_client.data_science_client = mocker.AsyncMock() keboola_client.has_feature = mocker.AsyncMock(return_value=True) - workspace_manager.get_branch_id = mocker.AsyncMock(return_value='branch-1') + workspace_manager.get_data_app_branch_id = mocker.AsyncMock(return_value='branch-1') keboola_client.data_science_client.create_data_app = mocker.AsyncMock( return_value=_make_python_js_data_app_response() @@ -1270,7 +1270,7 @@ async def test_modify_python_js_data_app_create_calls_full_provisioning_chain( keboola_client.data_science_client = mocker.AsyncMock() keboola_client.has_feature = mocker.AsyncMock(return_value=True) - workspace_manager.get_branch_id = mocker.AsyncMock(return_value='branch-1') + workspace_manager.get_data_app_branch_id = mocker.AsyncMock(return_value='branch-1') app_response = _make_python_js_data_app_response() keboola_client.data_science_client.create_data_app = mocker.AsyncMock(return_value=app_response) @@ -1347,7 +1347,7 @@ async def test_modify_python_js_data_app_create_authentication_type( keboola_client.data_science_client = mocker.AsyncMock() keboola_client.has_feature = mocker.AsyncMock(return_value=True) - workspace_manager.get_branch_id = mocker.AsyncMock(return_value='branch-1') + workspace_manager.get_data_app_branch_id = mocker.AsyncMock(return_value='branch-1') app_response = _make_python_js_data_app_response() keboola_client.data_science_client.create_data_app = mocker.AsyncMock(return_value=app_response) @@ -1390,7 +1390,7 @@ async def test_modify_python_js_data_app_update_patches_storage_config( keboola_client = KeboolaClient.from_state(mcp_context_client.session.state) keboola_client.has_feature = mocker.AsyncMock(return_value=True) - workspace_manager.get_branch_id = mocker.AsyncMock(return_value='branch-1') + workspace_manager.get_data_app_branch_id = mocker.AsyncMock(return_value='branch-1') existing_data_app = DataApp( name='Old', @@ -1462,7 +1462,7 @@ async def test_modify_python_js_data_app_update_repoints_external_git_branch( preserved, no re-encryption happens, and the change_summary hints at redeploy (CFTL-714).""" keboola_client = KeboolaClient.from_state(mcp_context_client.session.state) keboola_client.has_feature = mocker.AsyncMock(return_value=True) - workspace_manager.get_branch_id = mocker.AsyncMock(return_value='branch-1') + workspace_manager.get_data_app_branch_id = mocker.AsyncMock(return_value='branch-1') git_block = { 'repository': 'https://github.com/org/repo.git', @@ -1535,7 +1535,7 @@ async def test_modify_python_js_data_app_update_branch_rejected_on_managed_repo_ is written (CFTL-714 review).""" keboola_client = KeboolaClient.from_state(mcp_context_client.session.state) keboola_client.has_feature = mocker.AsyncMock(return_value=True) - workspace_manager.get_branch_id = mocker.AsyncMock(return_value='branch-1') + workspace_manager.get_data_app_branch_id = mocker.AsyncMock(return_value='branch-1') managed_repo_app = DataApp( name='Prod App', @@ -1587,7 +1587,7 @@ async def test_modify_python_js_data_app_create_without_workspace_feature( keboola_client.data_science_client = mocker.AsyncMock() keboola_client.has_feature = mocker.AsyncMock(return_value=False) - workspace_manager.get_workspace_id = mocker.AsyncMock(return_value='wid-legacy') + workspace_manager.get_data_app_workspace_id = mocker.AsyncMock(return_value='wid-legacy') app_response = _make_python_js_data_app_response() keboola_client.data_science_client.create_data_app = mocker.AsyncMock(return_value=app_response) @@ -1628,7 +1628,7 @@ async def test_modify_python_js_data_app_update_injects_workspace_id_without_fea keboola_client = KeboolaClient.from_state(mcp_context_client.session.state) keboola_client.has_feature = mocker.AsyncMock(return_value=False) - workspace_manager.get_workspace_id = mocker.AsyncMock(return_value='wid-legacy') + workspace_manager.get_data_app_workspace_id = mocker.AsyncMock(return_value='wid-legacy') existing_data_app = DataApp( name='Old', @@ -1689,7 +1689,7 @@ async def test_modify_python_js_data_app_create_passes_storage_through( keboola_client = KeboolaClient.from_state(mcp_context_client.session.state) keboola_client.data_science_client = mocker.AsyncMock() - workspace_manager.get_branch_id = mocker.AsyncMock(return_value='branch-1') + workspace_manager.get_data_app_branch_id = mocker.AsyncMock(return_value='branch-1') app_response = _make_python_js_data_app_response() keboola_client.data_science_client.create_data_app = mocker.AsyncMock(return_value=app_response) @@ -1737,7 +1737,7 @@ async def test_modify_python_js_data_app_create_omits_empty_storage( keboola_client = KeboolaClient.from_state(mcp_context_client.session.state) keboola_client.data_science_client = mocker.AsyncMock() - workspace_manager.get_branch_id = mocker.AsyncMock(return_value='branch-1') + workspace_manager.get_data_app_branch_id = mocker.AsyncMock(return_value='branch-1') app_response = _make_python_js_data_app_response() keboola_client.data_science_client.create_data_app = mocker.AsyncMock(return_value=app_response) @@ -1774,7 +1774,7 @@ async def test_modify_python_js_data_app_update_replaces_storage( """Update path: a non-empty `storage` argument replaces the entire stored storage block.""" keboola_client = KeboolaClient.from_state(mcp_context_client.session.state) - workspace_manager.get_branch_id = mocker.AsyncMock(return_value='branch-1') + workspace_manager.get_data_app_branch_id = mocker.AsyncMock(return_value='branch-1') existing_data_app = DataApp( name='Old', @@ -2176,7 +2176,7 @@ async def test_modify_python_js_data_app_create_draft_uses_external_git( keboola_client = KeboolaClient.from_state(mcp_context_client.session.state) keboola_client.data_science_client = mocker.AsyncMock() keboola_client.has_feature = mocker.AsyncMock(return_value=True) - workspace_manager.get_branch_id = mocker.AsyncMock(return_value='branch-1') + workspace_manager.get_data_app_branch_id = mocker.AsyncMock(return_value='branch-1') parent_repo = 'https://managed.repo/org/prod.git' parent_data_app_id = 'app-prod-1' @@ -2273,7 +2273,7 @@ async def test_modify_python_js_data_app_create_draft_defaults_branch_to_init( keboola_client = KeboolaClient.from_state(mcp_context_client.session.state) keboola_client.data_science_client = mocker.AsyncMock() keboola_client.has_feature = mocker.AsyncMock(return_value=True) - workspace_manager.get_branch_id = mocker.AsyncMock(return_value='branch-1') + workspace_manager.get_data_app_branch_id = mocker.AsyncMock(return_value='branch-1') parent = _make_python_js_parent_data_app() mocker.patch('keboola_mcp_server.tools.data_apps._fetch_data_app', mocker.AsyncMock(return_value=parent)) @@ -2337,7 +2337,7 @@ async def test_modify_python_js_data_app_create_draft_auto_derives_slug( keboola_client = KeboolaClient.from_state(mcp_context_client.session.state) keboola_client.data_science_client = mocker.AsyncMock() keboola_client.has_feature = mocker.AsyncMock(return_value=True) - workspace_manager.get_branch_id = mocker.AsyncMock(return_value='branch-1') + workspace_manager.get_data_app_branch_id = mocker.AsyncMock(return_value='branch-1') parent = _make_python_js_parent_data_app() mocker.patch('keboola_mcp_server.tools.data_apps._fetch_data_app', mocker.AsyncMock(return_value=parent)) @@ -2406,7 +2406,7 @@ async def test_modify_python_js_data_app_create_draft_rejects_when_parent_is_str keboola_client = KeboolaClient.from_state(mcp_context_client.session.state) keboola_client.data_science_client = mocker.AsyncMock() keboola_client.has_feature = mocker.AsyncMock(return_value=True) - workspace_manager.get_branch_id = mocker.AsyncMock(return_value='branch-1') + workspace_manager.get_data_app_branch_id = mocker.AsyncMock(return_value='branch-1') streamlit_parent = _make_python_js_parent_data_app(type='streamlit', repo_url=None) mocker.patch('keboola_mcp_server.tools.data_apps._fetch_data_app', mocker.AsyncMock(return_value=streamlit_parent)) @@ -2432,7 +2432,7 @@ async def test_modify_python_js_data_app_create_draft_rejects_when_parent_is_dra keboola_client = KeboolaClient.from_state(mcp_context_client.session.state) keboola_client.data_science_client = mocker.AsyncMock() keboola_client.has_feature = mocker.AsyncMock(return_value=True) - workspace_manager.get_branch_id = mocker.AsyncMock(return_value='branch-1') + workspace_manager.get_data_app_branch_id = mocker.AsyncMock(return_value='branch-1') # A draft has no repo_url of its own; the guard must fire before the repo_url check below. draft_parent = _make_python_js_parent_data_app(is_draft=True, repo_url=None) @@ -2460,7 +2460,7 @@ async def test_modify_python_js_data_app_create_draft_rejects_when_parent_missin keboola_client = KeboolaClient.from_state(mcp_context_client.session.state) keboola_client.data_science_client = mocker.AsyncMock() keboola_client.has_feature = mocker.AsyncMock(return_value=True) - workspace_manager.get_branch_id = mocker.AsyncMock(return_value='branch-1') + workspace_manager.get_data_app_branch_id = mocker.AsyncMock(return_value='branch-1') parent = _make_python_js_parent_data_app(repo_url=None) mocker.patch('keboola_mcp_server.tools.data_apps._fetch_data_app', mocker.AsyncMock(return_value=parent)) @@ -2485,7 +2485,7 @@ async def test_modify_python_js_data_app_create_prod_calls_get_app_git_repo_for_ keboola_client = KeboolaClient.from_state(mcp_context_client.session.state) keboola_client.data_science_client = mocker.AsyncMock() keboola_client.has_feature = mocker.AsyncMock(return_value=True) - workspace_manager.get_branch_id = mocker.AsyncMock(return_value='branch-1') + workspace_manager.get_data_app_branch_id = mocker.AsyncMock(return_value='branch-1') keboola_client.data_science_client.create_data_app = mocker.AsyncMock( return_value=_make_python_js_data_app_response() diff --git a/uv.lock b/uv.lock index f1c06f2f7..f3c3b9664 100644 --- a/uv.lock +++ b/uv.lock @@ -1172,7 +1172,7 @@ wheels = [ [[package]] name = "keboola-mcp-server" -version = "1.75.4" +version = "1.76.0" source = { editable = "." } dependencies = [ { name = "cryptography" },