Merge pull request #7 from ydb-platform/auth_methods

Add several auth methods

Author: vgvoleg

README.md

@@ -31,58 +31,67 @@
 }
 ```
 
-#### Example: Using Login/Password Authentication
+### Via pipx
 
-To use login/password authentication, specify the `--ydb-auth-mode`, `--ydb-login`, and `--ydb-password` arguments:
+[pipx](https://pipx.pypa.io/stable/) allows you to run various applications from PyPI without explicitly installing each one. However, it must be [installed](https://pipx.pypa.io/stable/#install-pipx) first. Below are examples of how to configure YDB MCP using `pipx`.
+
+#### Example: Using Anonymous Authentication
 
 ```json
 {
   "mcpServers": {
     "ydb": {
-      "command": "uvx",
+      "command": "pipx",
       "args": [
-        "ydb-mcp",
-        "--ydb-endpoint", "grpc://localhost:2136/local",
-        "--ydb-auth-mode", "login-password",
-        "--ydb-login", "<your-username>",
-        "--ydb-password", "<your-password>"
+        "run", "ydb-mcp",
+        "--ydb-endpoint", "grpc://localhost:2136/local"
       ]
     }
   }
 }
 ```
 
-### Via pipx
+### Via pip
 
-[pipx](https://pipx.pypa.io/stable/) allows you to run various applications from PyPI without explicitly installing each one. However, it must be [installed](https://pipx.pypa.io/stable/#install-pipx) first. Below are examples of how to configure YDB MCP using `pipx`.
+YDB MCP can be installed using `pip`, [Python's package installer](https://pypi.org/project/pip/). The package is [available on PyPI](https://pypi.org/project/ydb-mcp/) and includes all necessary dependencies.
+
+```bash
+pip install ydb-mcp
+```
+
+To get started with YDB MCP, you'll need to configure your MCP client to communicate with the YDB instance. Below are example configuration files that you can customize according to your setup and then put into MCP client's settings. Path to the Python interpreter might also need to be adjusted to the correct virtual environment that has the `ydb-mcp` package installed.
 
 #### Example: Using Anonymous Authentication
 
 ```json
 {
   "mcpServers": {
     "ydb": {
-      "command": "pipx",
+      "command": "python3",
       "args": [
-        "run", "ydb-mcp",
+        "-m", "ydb_mcp",
         "--ydb-endpoint", "grpc://localhost:2136/local"
       ]
     }
   }
 }
 ```
 
-#### Example: Using Login/Password Authentication
+### Authentication
+
+Regardless of the usage method (`uvx`, `pipx` or `pip`), you can configure authentication for your YDB installation. To do this, pass special command line arguments.
+
+#### Using Login/Password Authentication
 
 To use login/password authentication, specify the `--ydb-auth-mode`, `--ydb-login`, and `--ydb-password` arguments:
 
 ```json
 {
   "mcpServers": {
     "ydb": {
-      "command": "pipx",
+      "command": "uvx",
       "args": [
-        "run", "ydb-mcp",
+        "ydb-mcp",
         "--ydb-endpoint", "grpc://localhost:2136/local",
         "--ydb-auth-mode", "login-password",
         "--ydb-login", "<your-username>",
@@ -93,47 +102,40 @@ To use login/password authentication, specify the `--ydb-auth-mode`, `--ydb-logi
 }
 ```
 
-### Via pip
-
-YDB MCP can be installed using `pip`, [Python's package installer](https://pypi.org/project/pip/). The package is [available on PyPI](https://pypi.org/project/ydb-mcp/) and includes all necessary dependencies.
+#### Using Access Token Authentication
 
-```bash
-pip install ydb-mcp
-```
-
-To get started with YDB MCP, you'll need to configure your MCP client to communicate with the YDB instance. Below are example configuration files that you can customize according to your setup and then put into MCP client's settings. Path to the Python interpreter might also need to be adjusted to the correct virtual environment that has the `ydb-mcp` package installed.
-
-#### Example: Using Anonymous Authentication
+To use access token authentication, specify the `--ydb-auth-mode` and `--ydb-access-token` arguments:
 
 ```json
 {
   "mcpServers": {
     "ydb": {
-      "command": "python3",
+      "command": "uvx",
       "args": [
-        "-m", "ydb_mcp",
-        "--ydb-endpoint", "grpc://localhost:2136/local"
+        "ydb-mcp",
+        "--ydb-endpoint", "grpc://localhost:2136/local",
+        "--ydb-auth-mode", "access-token",
+        "--ydb-access-token", "qwerty123"
       ]
     }
   }
 }
 ```
 
-#### Example: Using Login/Password Authentication
+#### Using Service Account Authentication
 
-To use login/password authentication, specify the `--ydb-auth-mode`, `--ydb-login`, and `--ydb-password` arguments:
+To use service account authentication, specify the `--ydb-auth-mode` and `--ydb-sa-key-file` arguments:
 
 ```json
 {
   "mcpServers": {
     "ydb": {
-      "command": "python3",
+      "command": "uvx",
       "args": [
-        "-m", "ydb_mcp",
+        "ydb-mcp",
         "--ydb-endpoint", "grpc://localhost:2136/local",
-        "--ydb-auth-mode", "login-password",
-        "--ydb-login", "<your-username>",
-        "--ydb-password", "<your-password>"
+        "--ydb-auth-mode", "service-account",
+        "--ydb-sa-key-file", "~/sa_key.json"
       ]
     }
   }

ydb_mcp/__main__.py

@@ -5,7 +5,13 @@
 import os
 import sys
 
-from ydb_mcp.server import AUTH_MODE_ANONYMOUS, AUTH_MODE_LOGIN_PASSWORD, YDBMCPServer
+from ydb_mcp.server import (
+    AUTH_MODE_ACCESS_TOKEN,
+    AUTH_MODE_ANONYMOUS,
+    AUTH_MODE_LOGIN_PASSWORD,
+    AUTH_MODE_SERVICE_ACCOUNT,
+    YDBMCPServer,
+)
 
 
 def parse_args():
@@ -66,7 +72,12 @@ def main():
     )
 
     # Validate auth mode and required credentials
-    supported_auth_modes = {AUTH_MODE_ANONYMOUS, AUTH_MODE_LOGIN_PASSWORD}
+    supported_auth_modes = {
+        AUTH_MODE_ANONYMOUS,
+        AUTH_MODE_LOGIN_PASSWORD,
+        AUTH_MODE_ACCESS_TOKEN,
+        AUTH_MODE_SERVICE_ACCOUNT,
+    }
     auth_mode = args.ydb_auth_mode or AUTH_MODE_ANONYMOUS
     if auth_mode not in supported_auth_modes:
         print(
@@ -81,26 +92,46 @@ def main():
                 file=sys.stderr,
             )
             exit(1)
+    if auth_mode == AUTH_MODE_ACCESS_TOKEN:
+        if not args.ydb_access_token:
+            print(
+                "Error: --ydb-access-token is required for access-token authentication mode.",
+                file=sys.stderr,
+            )
+            exit(1)
+    if auth_mode == AUTH_MODE_SERVICE_ACCOUNT:
+        if not args.ydb_sa_key_file:
+            print(
+                "Error: --ydb-sa-key-file is required for service-account authentication mode.",
+                file=sys.stderr,
+            )
+            exit(1)
 
     # Set environment variables for YDB if provided via arguments
     if args.ydb_endpoint:
         os.environ["YDB_ENDPOINT"] = args.ydb_endpoint
     if args.ydb_database:
         os.environ["YDB_DATABASE"] = args.ydb_database
+    if args.ydb_auth_mode:
+        os.environ["YDB_AUTH_MODE"] = args.ydb_auth_mode
     if args.ydb_login:
         os.environ["YDB_LOGIN"] = args.ydb_login
     if args.ydb_password:
         os.environ["YDB_PASSWORD"] = args.ydb_password
-    if args.ydb_auth_mode:
-        os.environ["YDB_AUTH_MODE"] = args.ydb_auth_mode
+    if args.ydb_sa_key_file:
+        os.environ["YDB_SA_KEY_FILE"] = args.ydb_sa_key_file
+    if args.ydb_access_token:
+        os.environ["YDB_ACCESS_TOKEN"] = args.ydb_access_token
 
     # Create and run the server
     server = YDBMCPServer(
         endpoint=args.ydb_endpoint,
         database=args.ydb_database,
+        auth_mode=auth_mode,
         login=args.ydb_login,
         password=args.ydb_password,
-        auth_mode=auth_mode,
+        access_token=args.ydb_access_token,
+        sa_key_file=args.ydb_sa_key_file,
     )
 
     print("Starting YDB MCP server with stdio transport")

ydb_mcp/server.py

@@ -22,6 +22,8 @@
 # Authentication mode constants
 AUTH_MODE_ANONYMOUS = "anonymous"
 AUTH_MODE_LOGIN_PASSWORD = "login-password"
+AUTH_MODE_ACCESS_TOKEN = "access-token"
+AUTH_MODE_SERVICE_ACCOUNT = "service-account"
 
 
 class CustomJSONEncoder(json.JSONEncoder):
@@ -99,6 +101,8 @@ def __init__(
         auth_mode: str | None = None,
         login: str | None = None,
         password: str | None = None,
+        access_token: str | None = None,
+        sa_key_file: str | None = None,
         root_certificates: str | None = None,
         *args,
         **kwargs,
@@ -134,7 +138,12 @@ def __init__(
         self._original_methods: Dict = {}
 
         # Authentication settings
-        supported_auth_modes = {AUTH_MODE_ANONYMOUS, AUTH_MODE_LOGIN_PASSWORD}
+        supported_auth_modes = {
+            AUTH_MODE_ANONYMOUS,
+            AUTH_MODE_LOGIN_PASSWORD,
+            AUTH_MODE_ACCESS_TOKEN,
+            AUTH_MODE_SERVICE_ACCOUNT,
+        }
         self.auth_mode = auth_mode or AUTH_MODE_ANONYMOUS
         if self.auth_mode not in supported_auth_modes:
             raise ValueError(
@@ -143,6 +152,9 @@ def __init__(
         self.login = login
         self.password = password
 
+        self.sa_key_file = sa_key_file
+        self.access_token = access_token
+
         # Initialize logging
         logging.basicConfig(level=logging.INFO)
 
@@ -171,6 +183,16 @@ def _login_password_credentials(self) -> ydb.Credentials:
         logger.info(f"Using login-password authentication with login: {self.login}")
         return ydb.credentials.StaticCredentials.from_user_password(self.login, self.password)
 
+    def _access_token_credentials(self) -> ydb.Credentials:
+        """Create access token credentials."""
+        logger.info("Using access token authentication")
+        return ydb.credentials.AccessTokenCredentials(self.access_token)
+
+    def _service_account_credentials(self) -> ydb.Credentials:
+        """Create service account credentials."""
+        logger.info(f"Using service account authentication with key file: {self.sa_key_file}")
+        return ydb.iam.ServiceAccountCredentials.from_file(self.sa_key_file)
+
     async def create_driver(self):
         """Create a YDB driver with the current settings.
 
@@ -995,6 +1017,18 @@ def get_credentials_factory(self) -> Optional[Callable[[], ydb.Credentials]]:
                 return None
             logger.info(f"Using login/password authentication with user '{self.login}'")
             return self._login_password_credentials
+        if self.auth_mode == AUTH_MODE_SERVICE_ACCOUNT:
+            if not self.sa_key_file:
+                self.auth_error = "Service account key file must be provided for service-account authentication mode."
+                return None
+            logger.info(f"Using service-account authentication with key file '{self.sa_key_file}'")
+            return self._service_account_credentials
+        if self.auth_mode == AUTH_MODE_ACCESS_TOKEN:
+            if not self.access_token:
+                self.auth_error = "Access token must be provided for access-token authentication mode."
+                return None
+            logger.info("Using access-token authentication with token")
+            return self._access_token_credentials
         else:
             # Default to anonymous auth
             logger.info("Using anonymous authentication")