Support dynamic project id (#50)

* Supporting e2e server cleanup and middleware for dynamic project id

* Support dynamic_project id as a setting and project id validations

* Correcting env

* Correcting env

Author: aniket-s-kulkarni

.github/workflows/docker.yaml

@@ -20,9 +20,10 @@ jobs:
           tags: dremio-mcp-server:${{ github.sha }}
       - name: Validate docker
         run: |
-          docker run -e TOOLS_MODE=FOR_DATA_PATTERNS \
-              -e DREMIO_URI=https://fake \
-              -e DREMIO_OAUTH_SUPPORTED=false \
+          docker run \
+              -e DREMIOAI_TOOLS__SERVER_MODE=FOR_DATA_PATTERNS \
+              -e DREMIOAI_DREMIO__URI=https://fake \
+              -e DREMIOAI_DREMIO__OAUTH_SUPPORTED=false \
           dremio-mcp-server:${{ github.sha }} \
             dremio-mcp-server tools list
       - name: Install uv
@@ -37,9 +38,10 @@ jobs:
 
       - name: Start container
         run: |
-          docker run -e TOOLS_MODE=FOR_DATA_PATTERNS \
-            -e DREMIO_URI=https://fake \
-            -e DREMIO_OAUTH_SUPPORTED=false \
+          docker run \
+            -e DREMIOAI_TOOLS__SERVER_MODE=FOR_DATA_PATTERNS \
+            -e DREMIOAI_DREMIO__URI=https://fake \
+            -e DREMIOAI_DREMIO__OAUTH_SUPPORTED=false \
             -p 6789:6789 \
             --name mcp --rm \
             --network host \

.github/workflows/pytest.yml

@@ -26,4 +26,4 @@ jobs:
 
       - name: Run tests
         run: |
-          make test
+          uv run pytest tests

Makefile

@@ -8,4 +8,4 @@ asyncio_default_fixture_loop_scope = function
 
 # Display summary info for skipped, xfailed, xpassed tests
 # along with the percentage of passing tests at the end
-addopts = -v --showlocals
+addopts = -v --showlocals -x

pytest.ini

@@ -13,6 +13,8 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 #
+import uuid
+from uuid import UUID
 
 from pydantic import (
     Field,
@@ -24,7 +26,7 @@
     AliasChoices,
 )
 from pydantic_settings import BaseSettings, SettingsConfigDict
-from typing import Optional, Union, Annotated, Self, List, Dict, Any, Callable
+from typing import Optional, Union, Annotated, Self, List, Dict, Any, Callable, Literal
 from dremioai.config.tools import ToolType
 from enum import auto, StrEnum
 from pathlib import Path
@@ -37,6 +39,8 @@
 from importlib.util import find_spec
 from datetime import datetime
 
+ProjectId = Union[UUID, Literal["DREMIO_DYNAMIC"]]
+
 
 def _resolve_tools_settings(server_mode: Union[ToolType, int, str]) -> ToolType:
     if isinstance(server_mode, str):
@@ -120,7 +124,7 @@ class Dremio(BaseModel):
         Union[str, HttpUrl, DremioCloudUri], AfterValidator(_resolve_dremio_uri)
     ]
     raw_pat: Optional[str] = Field(default=None, alias="pat")
-    project_id: Optional[str] = None
+    raw_project_id: Optional[ProjectId] = Field(default=None, alias="project_id")
     enable_search: Optional[bool] = Field(
         default=False,
         alias=AliasChoices("enable_search", "enable_experimental"),
@@ -142,11 +146,13 @@ def oauth_configured(self) -> bool:
     def oauth_supported(self) -> bool:
         return self.project_id is not None
 
-    # @field_validator("_pat", mode="wrap")
-    # @classmethod
-    # def validate_pat(cls, v: str, handler: ValidatorFunctionWrapHandler) -> str:
-    #    v = _resolve_token_file(v)
-    #    return handler(v)
+    @property
+    def project_id(self) -> Optional[str]:
+        return str(self.raw_project_id) if self.raw_project_id else None
+
+    @project_id.setter
+    def project_id(self, v: str):
+        self.raw_project_id = uuid.UUID(v)
 
     @property
     def pat(self) -> str:
@@ -229,7 +235,8 @@ class Settings(BaseSettings):
     beeai: Optional[BeeAI] = Field(default=None)
     model_config = SettingsConfigDict(
         env_file=".env",
-        env_nested_delimiter="_",
+        env_nested_delimiter="__",
+        env_prefix="DREMIOAI_",
         env_extra="allow",
         use_enum_values=True,
     )

src/dremioai/config/settings.py

@@ -46,6 +46,8 @@
 from mcp.server.auth.provider import AccessToken, TokenVerifier
 from starlette.middleware.authentication import AuthenticationMiddleware
 
+from dremioai.tools.tools import ProjectIdMiddleware
+
 
 class Transports(StrEnum):
     stdio = auto()
@@ -73,16 +75,23 @@ def streamable_http_app(self):
         app.add_middleware(
             AuthenticationMiddleware, backend=BearerAuthBackend(token_verifier)
         )
-        log.logger("streamable_http_app").info(
-            f"Adding auth middleware {app.user_middleware}"
-        )
+        if self.support_project_id_endpoints:
+            # this means, dynamically allow endpoints
+            # like ../mcp/{project_id}/..  and extract that project id as
+            # context var
+            app.add_middleware(ProjectIdMiddleware)
         return app
 
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.support_project_id_endpoints = False
+
 
 def init(
     mode: Union[tools.ToolType, List[tools.ToolType]] = None,
     transport: Transports = Transports.stdio,
     port: int = None,
+    support_project_id_endpoints: bool = False,
 ) -> FastMCP:
     mcp_cls = FastMCP if transport == Transports.stdio else FastMCPServerWithAuthToken
     log.logger("init").info(
@@ -92,6 +101,8 @@ def init(
     if port is not None:
         opts["port"] = port
     mcp = mcp_cls("Dremio", **opts)
+    if transport == Transports.streamable_http and support_project_id_endpoints:
+        mcp.support_project_id_endpoints = support_project_id_endpoints
     mode = reduce(ior, mode) if mode is not None else None
     for tool in tools.get_tools(For=mode):
         tool_instance = tool()

src/dremioai/servers/mcp.py

@@ -13,14 +13,14 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 #
-
+import json
+from contextvars import ContextVar
 from typing import (
     List,
     Dict,
     Any,
     Optional,
     Literal,
-    TypeAlias,
     Union,
     Annotated,
     ClassVar,
@@ -32,18 +32,18 @@
     Awaitable,
 )
 
-from pathlib import Path
 from dataclasses import dataclass, asdict, field
-from enum import auto, IntFlag
+
+from starlette.datastructures import URL
+from starlette.middleware.base import BaseHTTPMiddleware, RequestResponseEndpoint
+from starlette.requests import Request
+
 from dremioai import log
 import re
 import functools
 
 import pandas as pd
-
-from pathlib import Path
-
-from dremioai.api.dremio import sql, projects, usage, engines, search
+from dremioai.api.dremio import sql, usage, search
 from dremioai.config import settings
 from dremioai.config.tools import ToolType
 from dremioai.api.prometheus import vm
@@ -53,7 +53,6 @@
 from io import StringIO
 from sqlglot import parse_one
 from sqlglot import expressions
-from mcp.server.fastmcp.server import Context
 from mcp.server.auth.middleware.auth_context import get_access_token
 from mcp.server.auth.provider import AccessToken
 
@@ -106,6 +105,31 @@ async def invoke(self):
         raise NotImplementedError("Subclasses should implement this method")
 
 
+class ProjectIdMiddleware(BaseHTTPMiddleware):
+    pat = re.compile(r"/mcp/([\da-z-]+)(/?.*)")
+    logger = log.logger("ProjectIdMiddleware")
+
+    # Context variable to store the current project ID
+    project_id_context: ContextVar[str | None] = ContextVar("project_id", default=None)
+
+    @classmethod
+    def get_project_id(cls) -> Optional[str]:
+        return cls.project_id_context.get()
+
+    async def dispatch(self, request: Request, call_next: RequestResponseEndpoint):
+        ProjectIdMiddleware.logger.info(
+            f"Request {request.url.path} body = {await request.body()!s}"
+        )
+        if m := ProjectIdMiddleware.pat.search(request.url.path):
+            ProjectIdMiddleware.project_id_context.set(m.group(1))
+        else:
+            ProjectIdMiddleware.logger.debug(
+                f"Path {request.url.path} ({request.url!r}) doesn't match"
+            )
+
+        return await call_next(request)
+
+
 # A decorator to ensure a tool that needs to access Dremio runs with the correct token
 # if invoked through streamable HTTP transport _with_ a valid Dremio bearer token
 # It is a no-op if the tool is invoked through stdio transport, as MCP server ensures
@@ -114,11 +138,22 @@ def secured(fn: Callable[P, Awaitable[T]]) -> Callable[P, Awaitable[T]]:
 
     @functools.wraps(fn)
     async def _impl(self, *args: P.args, **kw: P.kwargs) -> T:
+        overrides = {}
         if isinstance((token := get_access_token()), AccessToken):
-            return await settings.run_with(
-                fn, {"dremio.pat": token.token}, (self,) + args, kw
+            overrides["dremio.pat"] = token.token
+            logger.debug(
+                f"Overriding PAT with token from request: {token.token[:4]}..."
             )
-        return await fn(self, *args, **kw)
+
+        if project_id := ProjectIdMiddleware.get_project_id():
+            overrides["dremio.project_id"] = project_id
+            logger.debug(f"Overriding project_id with {project_id}")
+
+        return (
+            await settings.run_with(fn, overrides, (self,) + args, kw)
+            if overrides
+            else await fn(self, *args, **kw)
+        )
 
     return _impl
 

src/dremioai/tools/tools.py

@@ -15,12 +15,17 @@
 #
 
 import os
+import uuid
+
+import pydantic
 import pytest
 import yaml
 from pathlib import Path
 from tempfile import TemporaryDirectory
 from unittest.mock import patch
 
+from pydantic_core import ValidationError
+
 from dremioai.config import settings
 from dremioai.config.tools import ToolType
 
@@ -48,7 +53,7 @@ def test_configure_creates_default_config(mock_config_dir):
 def test_create_default_config(mock_config_dir):
     uri = settings.DremioCloudUri.PRODEMEA.value
     pat = "test-pat"
-    project_id = "test-project"
+    project_id = uuid.uuid4()
     mode = ToolType.FOR_DATA_PATTERNS
     settings.configure(force=True)
     settings._settings.set(
@@ -70,7 +75,7 @@ def test_create_default_config(mock_config_dir):
     assert (
         dremio.uri == "https://api.eu.dremio.cloud"
         and dremio.pat == pat
-        and dremio.project_id == project_id
+        and dremio.project_id == str(project_id)
     )
     tools = settings.instance().tools
     assert tools.server_mode == mode
@@ -89,3 +94,43 @@ def test_experimental_rename(name: str, value: bool):
         {name: value, "uri": "https://foo", "pat": "bar"}
     )
     assert d.enable_search == value
+
+
+@pytest.mark.parametrize(
+    "name,project_id,error",
+    [
+        ["valid project id", str(uuid.uuid4()), False],
+        ["no project id", None, False],
+        ["invalid project id", "asdfsa safsa", True],
+        ["invalid project id", str(uuid.uuid4())[:-1] + "a", True],
+        ["dynamic project id", "DREMIO_DYNAMIC", False],
+    ],
+)
+def test_projects(name: str, project_id: str | None, error: bool):
+    val = {"uri": "https://foo", "project_id": project_id}
+    if error:
+        try:
+            settings.Dremio.model_validate(val)
+            assert False
+        except:
+            pass
+    else:
+        d = settings.Dremio.model_validate(val)
+        assert d.project_id == project_id or d.project_id is None and project_id is None
+
+
+def test_env_file(mock_config_dir):
+    try:
+        os.environ["DREMIOAI_DREMIO__URI"] = "https://foo"
+        os.environ["DREMIOAI_DREMIO__PAT"] = "bar"
+        os.environ["DREMIOAI_TOOLS__SERVER_MODE"] = "FOR_DATA_PATTERNS"
+        settings.configure(force=True)
+        from rich import print as pp
+
+        pp(settings.instance().model_dump())
+        assert settings.instance().dremio.uri == "https://foo"
+        assert settings.instance().dremio.pat == "bar"
+        assert settings.instance().tools.server_mode == ToolType.FOR_DATA_PATTERNS
+    finally:
+        os.environ.pop("DREMIOAI_DREMIO_URI", None)
+        os.environ.pop("DREMIOAI_DREMIO_PAT", None)