feat: add OAuth URL generation tool

src/keboola_mcp_server/server.py

@@ -20,6 +20,7 @@
 from keboola_mcp_server.tools.doc import add_doc_tools
 from keboola_mcp_server.tools.flow.tools import add_flow_tools
 from keboola_mcp_server.tools.jobs import add_job_tools
+from keboola_mcp_server.tools.oauth import add_oauth_tools
 from keboola_mcp_server.tools.project import add_project_tools
 from keboola_mcp_server.tools.search import add_search_tools
 from keboola_mcp_server.tools.sql import add_sql_tools
@@ -163,6 +164,7 @@ async def oauth_callback_handler(request: Request) -> Response:
     add_doc_tools(mcp)
     add_flow_tools(mcp)
     add_job_tools(mcp)
+    add_oauth_tools(mcp)
     add_project_tools(mcp)
     add_search_tools(mcp)
     add_sql_tools(mcp)

src/keboola_mcp_server/tools/oauth.py

@@ -0,0 +1,67 @@
+"""OAuth URL generation tools for the MCP server."""
+
+import logging
+from typing import Annotated
+
+from fastmcp import Context
+from fastmcp.tools import FunctionTool
+from pydantic import Field
+
+from keboola_mcp_server.client import KeboolaClient
+from keboola_mcp_server.errors import tool_errors
+from keboola_mcp_server.mcp import KeboolaMcpServer, with_session_state
+
+LOG = logging.getLogger(__name__)
+
+TOOL_GROUP_NAME = 'OAUTH'
+
+
+def add_oauth_tools(mcp: KeboolaMcpServer) -> None:
+    """Adds OAuth tools to the MCP server."""
+    mcp.add_tool(FunctionTool.from_function(create_oauth_url))
+    LOG.info('OAuth tools added to the MCP server.')
+
+
+@tool_errors()
+@with_session_state()
+async def create_oauth_url(
+    component_id: Annotated[
+        str, Field(description='The component ID to grant access to (e.g., "keboola.ex-google-analytics-v4").')
+    ],
+    config_id: Annotated[str, Field(description='The configuration ID for the component.')],
+    ctx: Context,
+) -> str:
+    """
+    Generates an OAuth authorization URL for a Keboola component configuration.
+
+    When using this tool, be very concise in your response. Just guide the user to click the
+    authorization link.
+
+    Note that this tool should be called specifically for these OAuth-requiring components after their
+    configuration is created: keboola.ex-google-analytics-v4 and keboola.ex-gmail.
+    """
+    client = KeboolaClient.from_state(ctx.session.state)
+
+    # Create short-lived SAPI token
+    token_data = {
+        'description': f'Short-lived token for OAuth URL - {component_id}/{config_id}',
+        'componentAccess': [component_id],
+        'expiresIn': 3600,  # 1 hour expiration
+    }
+
+    # Create the token using the storage client
+    token_response = await client.storage_client.post(endpoint='tokens', data=token_data)
+
+    # Extract the token from response
+    sapi_token = token_response['token']
+
+    # Get the storage API URL from client
+    storage_api_url = client.storage_client.base_api_url
+
+    # Generate OAuth URL
+    oauth_url = (
+        f'https://external.keboola.com/oauth/index.html?token={sapi_token}'
+        f'&sapiUrl={storage_api_url}#/{component_id}/{config_id}'
+    )
+
+    return oauth_url

tests/test_server.py

@@ -23,6 +23,7 @@ async def test_list_tools(self):
             'add_config_row',
             'create_config',
             'create_flow',
+            'create_oauth_url',
             'create_sql_transformation',
             'docs_query',
             'find_component_id',

tests/tools/test_oauth.py

@@ -0,0 +1,117 @@
+"""Tests for OAuth URL generation tools."""
+
+from typing import Any, Mapping
+
+import pytest
+from mcp.server.fastmcp import Context
+from pytest_mock import MockerFixture
+
+from keboola_mcp_server.tools.oauth import create_oauth_url
+
+
+@pytest.fixture
+def mock_token_response() -> Mapping[str, Any]:
+    """Mock valid response from the token creation endpoint."""
+    return {
+        'token': 'KBC_TOKEN_12345',
+        'description': 'Short-lived token for OAuth URL - keboola.ex-google-analytics-v4/config-123',
+        'expiresIn': 3600,
+    }
+
+
+@pytest.mark.asyncio
+async def test_create_oauth_url_success(
+    mocker: MockerFixture, mcp_context_client: Context, mock_token_response: Mapping[str, Any]
+) -> None:
+    """Test successful OAuth URL creation."""
+    # Mock the storage client's post method to return the token response
+    mcp_context_client.session.state['sapi_client'].storage_client.post.return_value = mock_token_response
+    mcp_context_client.session.state['sapi_client'].storage_client.base_api_url = 'https://connection.test.keboola.com'
+
+    component_id = 'keboola.ex-google-analytics-v4'
+    config_id = 'config-123'
+
+    result = await create_oauth_url(component_id=component_id, config_id=config_id, ctx=mcp_context_client)
+
+    # Verify the storage client was called with correct parameters
+    mcp_context_client.session.state['sapi_client'].storage_client.post.assert_called_once_with(
+        endpoint='tokens',
+        data={
+            'description': f'Short-lived token for OAuth URL - {component_id}/{config_id}',
+            'componentAccess': [component_id],
+            'expiresIn': 3600,
+        },
+    )
+
+    # Verify the response is the URL string
+    assert isinstance(result, str)
+
+    expected_url = (
+        f'https://external.keboola.com/oauth/index.html'
+        f'?token=KBC_TOKEN_12345'
+        f'&sapiUrl=https://connection.test.keboola.com'
+        f'#/{component_id}/{config_id}'
+    )
+    assert result == expected_url
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    ('component_id', 'config_id'),
+    [
+        ('keboola.ex-google-analytics-v4', 'my-config-123'),
+        ('keboola.ex-gmail', 'gmail-config-456'),
+        ('other.component', 'test-config'),
+    ],
+)
+async def test_create_oauth_url_different_components(
+    mocker: MockerFixture,
+    mcp_context_client: Context,
+    mock_token_response: Mapping[str, Any],
+    component_id: str,
+    config_id: str,
+) -> None:
+    """Test OAuth URL creation for different components."""
+    # Mock the storage client
+    mcp_context_client.session.state['sapi_client'].storage_client.post.return_value = mock_token_response
+    mcp_context_client.session.state['sapi_client'].storage_client.base_api_url = 'https://connection.test.keboola.com'
+
+    result = await create_oauth_url(component_id=component_id, config_id=config_id, ctx=mcp_context_client)
+
+    # Verify component-specific parameters were used
+    assert isinstance(result, str)
+    assert f'#/{component_id}/{config_id}' in result
+
+    # Verify the API call included the correct component access
+    call_args = mcp_context_client.session.state['sapi_client'].storage_client.post.call_args
+    assert call_args[1]['data']['componentAccess'] == [component_id]
+    assert component_id in call_args[1]['data']['description']
+    assert config_id in call_args[1]['data']['description']
+
+
+@pytest.mark.asyncio
+async def test_create_oauth_url_token_creation_failure(mocker: MockerFixture, mcp_context_client: Context) -> None:
+    """Test OAuth URL creation when token creation fails."""
+    # Mock the storage client to raise an exception
+    mcp_context_client.session.state['sapi_client'].storage_client.post.side_effect = Exception('Token creation failed')
+
+    with pytest.raises(Exception, match='Token creation failed'):
+        await create_oauth_url(
+            component_id='keboola.ex-google-analytics-v4', config_id='config-123', ctx=mcp_context_client
+        )
+
+
+@pytest.mark.asyncio
+async def test_create_oauth_url_missing_token_in_response(mocker: MockerFixture, mcp_context_client: Context) -> None:
+    """Test OAuth URL creation when token is missing from response."""
+    # Mock response without token field
+    invalid_response = {
+        'description': 'Short-lived token for OAuth URL',
+        'expiresIn': 3600,
+    }
+    mcp_context_client.session.state['sapi_client'].storage_client.post.return_value = invalid_response
+
+    with pytest.raises(KeyError):
+        await create_oauth_url(
+            component_id='keboola.ex-google-analytics-v4', config_id='config-123', ctx=mcp_context_client
+        )