fix: make all tools async (#34)

Author: hsheth2

.github/workflows/ci.yml

@@ -26,6 +26,6 @@ jobs:
       - name: Checks
         run: make lint-check
 
-      # Tests don't run without credentials yet.
-      # - name: Test
-      #   run: make test
+      # Note that some tests require credentials that aren't available in CI environment.
+      - name: Test
+        run: make test

pyproject.toml

@@ -6,6 +6,7 @@ readme = "README.md"
 requires-python = ">=3.10"
 dependencies = [
     "acryl-datahub==1.2.0.1",
+    "asyncer>=0.0.8",
     "fastmcp==2.10.5",
     "jmespath~=1.0.1",
 ]

src/mcp_server_datahub/mcp_server.py

@@ -1,8 +1,21 @@
 import contextlib
 import contextvars
+import functools
+import inspect
 import pathlib
-from typing import Any, Dict, Iterator, List, Optional
+from typing import (
+    Any,
+    Awaitable,
+    Callable,
+    Dict,
+    Iterator,
+    List,
+    Optional,
+    ParamSpec,
+    TypeVar,
+)
 
+import asyncer
 import jmespath
 from datahub.errors import ItemNotFoundError
 from datahub.ingestion.graph.client import DataHubGraph
@@ -13,6 +26,23 @@
 from fastmcp import FastMCP
 from pydantic import BaseModel
 
+_P = ParamSpec("_P")
+_R = TypeVar("_R")
+
+
+# See https://github.com/jlowin/fastmcp/issues/864#issuecomment-3103678258
+# for why we need to wrap sync functions with asyncify.
+def async_background(fn: Callable[_P, _R]) -> Callable[_P, Awaitable[_R]]:
+    if inspect.iscoroutinefunction(fn):
+        raise RuntimeError("async_background can only be used on non-async functions")
+
+    @functools.wraps(fn)
+    async def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> _R:
+        return await asyncer.asyncify(fn)(*args, **kwargs)
+
+    return wrapper
+
+
 mcp = FastMCP[None](name="datahub")
 
 
@@ -132,6 +162,7 @@ def _clean_get_entity_response(raw_response: dict) -> dict:
 
 
 @mcp.tool(description="Get an entity by its DataHub URN.")
+@async_background
 def get_entity(urn: str) -> dict:
     client = get_datahub_client()
 
@@ -184,6 +215,7 @@ def get_entity(urn: str) -> dict:
 ```
 """
 )
+@async_background
 def search(
     query: str = "*",
     filters: Optional[Filter] = None,
@@ -215,6 +247,7 @@ def search(
 
 
 @mcp.tool(description="Use this tool to get the SQL queries associated with a dataset.")
+@async_background
 def get_dataset_queries(dataset_urn: str, start: int = 0, count: int = 10) -> dict:
     client = get_datahub_client()
 
@@ -330,6 +363,7 @@ def get_lineage(
 Use this tool to get upstream or downstream lineage for any entity, including datasets, schemaFields, dashboards, charts, etc. \
 Set upstream to True for upstream lineage, False for downstream lineage."""
 )
+@async_background
 def get_lineage(urn: str, upstream: bool, max_hops: int = 1) -> dict:
     client = get_datahub_client()
     lineage_api = AssetLineageAPI(client._graph)

tests/test_async_setup.py

@@ -0,0 +1,34 @@
+import inspect
+import time
+
+import anyio
+import fastmcp.tools.tool
+import pytest
+
+from mcp_server_datahub.mcp_server import async_background, mcp
+
+
+@pytest.mark.anyio
+async def test_async_background() -> None:
+    @async_background
+    def my_sleep(sec: float) -> None:
+        time.sleep(sec)
+
+    start_time = time.time()
+
+    async with anyio.create_task_group() as tg:
+        tg.start_soon(my_sleep, 0.5)
+        tg.start_soon(my_sleep, 0.6)
+        tg.start_soon(my_sleep, 0.7)
+
+    end_time = time.time()
+    duration = end_time - start_time
+    # The calls should not be serialized, so the duration should be less than the sum of the durations.
+    assert 0.5 <= duration < 1.8
+
+
+def test_all_tools_are_async() -> None:
+    # If any tools are sync, the tool execution will block the main event loop.
+    for tool in mcp._tool_manager._tools.values():
+        assert isinstance(tool, fastmcp.tools.tool.FunctionTool)
+        assert inspect.iscoroutinefunction(tool.fn)

tests/test_mcp_server.py

@@ -25,7 +25,13 @@
 
 @pytest.fixture(autouse=True, scope="session")
 def setup_client() -> Iterable[None]:
-    with with_datahub_client(DataHubClient.from_env()):
+    try:
+        client = DataHubClient.from_env()
+    except Exception as e:
+        if "`datahub init`" in str(e):
+            pytest.skip("No credentials available, skipping tests")
+        raise
+    with with_datahub_client(client):
         yield
 
 
@@ -95,9 +101,3 @@ async def test_search(mcp_client: Client) -> None:
     )
     assert res.is_error is False
     assert res.data is not None
-
-
-if __name__ == "__main__":
-    import pytest
-
-    pytest.main()