diff --git a/.gitignore b/.gitignore index 2419db0..da3b690 100644 --- a/.gitignore +++ b/.gitignore @@ -1,2 +1,8 @@ *.pyc .venv + +# IDE files +.idea/ +.vscode/ +*.swp +*.swo diff --git a/README.md b/README.md index fb066f2..01216be 100644 --- a/README.md +++ b/README.md @@ -103,16 +103,29 @@ Note: the uri is api endpoint associated with your environment: - For Dremio cloud based in the EMEA region (https://app.eu.dremio.cloud) use `https://api.eu.dremio.cloud` or use the short hand `prodemea` - For SW/K8S deployments use https://:<9047 or custom port> -Note: For security purposes, if you don't want the PAT to leak into your shell history file, create a file with your PAT in it and give it as an argument to the dremio config. +**Authentication Options:** -Example: +The Dremio MCP supports two authentication methods: +**Option 1: Personal Access Token (PAT)** ```shell $ uv run dremio-mcp-server config create dremioai \ --uri \ - --pat @/path/to/tokenfile \ + --pat ``` +**Option 2: Username/Password** +```shell +$ uv run dremio-mcp-server config create dremioai \ + --uri \ + --username \ + --password +``` + +Note: For security purposes, if you don't want credentials to leak into your shell history file, you can: +- For PAT: Create a file with your PAT and reference it with `--pat @/path/to/tokenfile` +- For passwords: Use environment variables or interactive prompts (future enhancement) + 2. Download and install Claude Desktop ([Claude](https://claude.ai/download)) Note: Claude has system requirements, such as node.js, please validate your system requirements with Claude official documentation. @@ -171,16 +184,23 @@ This file is located by default at `$HOME/.config/dremioai/config.yaml` but can #### Format ```yaml -# The dremio section contains 3 main things - the URI to connect, PAT to use -# and optionally the project_id if using with Dremio Cloud +# The dremio section contains the URI to connect and authentication credentials +# Choose either PAT or username/password authentication dremio: uri: https://.... # the Dremio URI + + # Option 1: Personal Access Token (PAT) pat: "@~/ws/tokens/idl.token" # PAT can be put in a file and used here with @ prefix + + # Option 2: Username/Password (alternative to PAT) + # username: your_username + # password: your_password + # optional - # project_id: .... + # project_id: .... # Required for Dremio Cloud tools: - server_mode: FOR_DATA_PATTERNS # the serverm + server_mode: FOR_DATA_PATTERNS # the server mode # Optionally the MCP server can also connect and use a prometheus configuration if it # has been enabled for your Dremio cluster (typically useful for SW installations) diff --git a/src/dremioai/api/transport.py b/src/dremioai/api/transport.py index b3597d3..b3c1ed2 100644 --- a/src/dremioai/api/transport.py +++ b/src/dremioai/api/transport.py @@ -137,9 +137,50 @@ def __init__(self, uri: Optional[str] = None, pat: Optional[str] = None): if uri is None: uri = dremio.uri + + # Determine authentication method if pat is None: - pat = dremio.pat + if dremio.pat is not None: + # Use PAT from settings + pat = dremio.pat + elif dremio.has_username_password: + # Use username/password authentication + pat = self._authenticate_with_username_password(uri, dremio.username, dremio.password) + else: + raise RuntimeError("No authentication method configured. Either 'pat' or 'username/password' must be provided") if uri is None or pat is None: raise RuntimeError(f"uri={uri} pat={pat} are required") super().__init__(uri, pat) + + def _authenticate_with_username_password(self, uri: str, username: str, password: str) -> str: + """Authenticate with Dremio using username/password and return the token""" + import asyncio + return asyncio.run(self._get_auth_token(uri, username, password)) + + async def _get_auth_token(self, uri: str, username: str, password: str) -> str: + """Get authentication token from Dremio using username/password""" + login_url = f"{uri}/apiv2/login" + login_data = { + "userName": username, + "password": password + } + + async with ClientSession() as session: + logger().info(f"Authenticating with Dremio at {login_url}") + async with session.post( + login_url, + json=login_data, + headers={"content-type": "application/json"}, + ssl=False + ) as response: + if response.status != 200: + error_text = await response.text() + raise RuntimeError(f"Authentication failed: {response.status} - {error_text}") + + auth_response = await response.json() + if "token" not in auth_response: + raise RuntimeError("Authentication response missing token") + + logger().info("Successfully authenticated with username/password") + return auth_response["token"] diff --git a/src/dremioai/config/settings.py b/src/dremioai/config/settings.py index 5089384..a52c95e 100644 --- a/src/dremioai/config/settings.py +++ b/src/dremioai/config/settings.py @@ -21,6 +21,7 @@ BaseModel, ConfigDict, field_serializer, + model_validator, ) from pydantic_settings import BaseSettings, SettingsConfigDict from typing import Optional, Union, Annotated, Self, List, Dict, Any, Callable @@ -119,6 +120,8 @@ class Dremio(BaseModel): Union[str, HttpUrl, DremioCloudUri], AfterValidator(_resolve_dremio_uri) ] raw_pat: Optional[str] = Field(default=None, alias="pat") + username: Optional[str] = None + password: Optional[str] = None project_id: Optional[str] = None enable_experimental: Optional[bool] = False # enable experimental tools oauth2: Optional[OAuth2] = None @@ -129,6 +132,22 @@ class Dremio(BaseModel): def serialize_pat(self, pat: str): return self.raw_pat if pat != self.raw_pat else pat + @model_validator(mode='after') + def validate_auth_method(self) -> 'Dremio': + """Validate authentication method configuration""" + has_pat = self.raw_pat is not None + has_user_pass = self.username is not None and self.password is not None + + # Allow configurations with no authentication for backward compatibility + # The actual authentication requirement will be enforced at runtime + if has_pat and has_user_pass: + raise ValueError("Cannot specify both 'pat' and 'username/password' authentication methods") + + if (self.username is None) != (self.password is None): + raise ValueError("Both 'username' and 'password' must be provided together") + + return self + @property def oauth_configured(self) -> bool: return self.oauth2 is not None @@ -137,6 +156,11 @@ def oauth_configured(self) -> bool: def oauth_supported(self) -> bool: return self.project_id is not None + @property + def has_username_password(self) -> bool: + """Check if username/password authentication is configured""" + return self.username is not None and self.password is not None + # @field_validator("_pat", mode="wrap") # @classmethod # def validate_pat(cls, v: str, handler: ValidatorFunctionWrapHandler) -> str: diff --git a/src/dremioai/servers/mcp.py b/src/dremioai/servers/mcp.py index d21bed9..e62a24e 100644 --- a/src/dremioai/servers/mcp.py +++ b/src/dremioai/servers/mcp.py @@ -91,6 +91,8 @@ def _mode() -> List[str]: def main( dremio_uri: Annotated[Optional[str], Option(help="Dremio URI")] = None, dremio_pat: Annotated[Optional[str], Option(help="Dremio PAT")] = None, + dremio_username: Annotated[Optional[str], Option(help="Dremio username")] = None, + dremio_password: Annotated[Optional[str], Option(help="Dremio password")] = None, dremio_project_id: Annotated[ Optional[str], Option(help="Dremio Project Id") ] = None, @@ -123,6 +125,8 @@ def main( { "dremio.uri": dremio_uri, "dremio.pat": dremio_pat, + "dremio.username": dremio_username, + "dremio.password": dremio_password, "dremio.project_id": dremio_project_id, "tools.server_mode": mode, } @@ -267,11 +271,23 @@ def create_default_config( ), ], pat: Annotated[ - str, + Optional[str], Option( - help="The Dremio PAT. If it starts with @ then treat the rest is treated as a filename" + help="The Dremio PAT. If it starts with @ then treat the rest is treated as a filename. Cannot be used with --username/--password." ), - ], + ] = None, + username: Annotated[ + Optional[str], + Option( + help="Dremio username for authentication. Must be used with --password. Cannot be used with --pat." + ), + ] = None, + password: Annotated[ + Optional[str], + Option( + help="Dremio password for authentication. Must be used with --username. Cannot be used with --pat." + ), + ] = None, project_id: Annotated[ Optional[str], Option(help="The Dremio project id, only if connecting to Dremio Cloud"), @@ -291,20 +307,39 @@ def create_default_config( bool, Option(help="Dry run, do not overwrite the config file. Just print it") ] = False, ): + # Validate authentication method + has_pat = pat is not None + has_user_pass = username is not None and password is not None + + if not has_pat and not has_user_pass: + raise BadParameter("Either --pat or both --username and --password must be provided") + + if has_pat and has_user_pass: + raise BadParameter("Cannot specify both --pat and --username/--password authentication methods") + + if (username is None) != (password is None): + raise BadParameter("Both --username and --password must be provided together") + mode = "|".join([tools.ToolType[m.upper()].name for m in mode]) - dremio = settings.Dremio.model_validate( - { - "uri": uri, - "pat": pat, - "project_id": project_id, - "enable_experimental": enable_experimental, - "oauth": ( - settings.OAuth2.model_validate({"client_id": oauth_client_id}) - if oauth_client_id - else None - ), - } - ) + dremio_config = { + "uri": uri, + "project_id": project_id, + "enable_experimental": enable_experimental, + "oauth": ( + settings.OAuth2.model_validate({"client_id": oauth_client_id}) + if oauth_client_id + else None + ), + } + + # Add authentication method + if has_pat: + dremio_config["pat"] = pat + else: + dremio_config["username"] = username + dremio_config["password"] = password + + dremio = settings.Dremio.model_validate(dremio_config) ts = settings.Tools.model_validate({"server_mode": mode}) settings.configure(settings.default_config(), force=True) settings.instance().dremio = dremio