ocientmcp 1.2.2__py3-none-any.whl

ocientmcp/__init__.py

@@ -0,0 +1,25 @@
+"""Ocient MCP Server - Model Context Protocol integration for Ocient databases
+
+This MCP server provides AI agents and tools like Windsurf/Cascade with the ability
+to interact with Ocient databases through the Model Context Protocol.
+
+Features:
+- Execute SQL queries and return results in JSON format
+- Manage database connections with DSN strings
+- List tables and schemas
+- Get table metadata and column information
+- Execute queries with parameters
+- Transaction support
+
+The server exposes tools that can be called by MCP clients to:
+1. Connect to Ocient databases
+2. Execute queries and get structured results
+3. Explore database schema
+4. Manage connections
+"""
+
+from ocientmcp.pkg_version import __commit__, __version__
+
+__version_info__ = tuple(int(i) for i in __version__.split(".") if i.isdigit())
+
+__all__ = ["__commit__", "__version__", "__version_info__"]

ocientmcp/models.py

@@ -0,0 +1,20 @@
+import dataclasses
+from typing import Dict, List
+
+# Slotted classes expose slot descriptors instead of real default values. Because FastMCP (through Pydantic) cannot serialize these,
+# using slots=True breaks schema generation, so dataclasses must avoid slots when used with FastMCP
+
+
+@dataclasses.dataclass(frozen=True, kw_only=True)
+class StatementResult:
+    message: str | None = None
+    """Optional message"""
+
+    columns: List[str] = dataclasses.field(default_factory=list)
+    """List of column names"""
+
+    rows: List[Dict[str, object]] = dataclasses.field(default_factory=list)
+    """List of row data"""
+
+    row_count: int = 0
+    """Number of returned rows"""

ocientmcp/pkg_version.py

@@ -0,0 +1,4 @@
+# ocientmcp version
+#
+__version__ = '1.2.2'
+__commit__: str | None = '0d466e5a70a4e28122676eaee26d431c12f13af6+dirty'

ocientmcp/provider.py

@@ -0,0 +1,198 @@
+"""Provider that dynamically exposes SQL MCP tools from the Ocient database."""
+
+import inspect
+import json
+import logging
+import math
+from collections.abc import Sequence
+from types import FunctionType
+from typing import Optional
+
+from fastmcp import Context
+from fastmcp.server.providers import Provider
+from fastmcp.tools.base import Tool
+from fastmcp.tools.function_tool import FunctionTool
+
+from ocientmcp.models import StatementResult
+from ocientmcp.server import OcientMCPServer, TransportMode
+from ocientmcp.tools import execute_statement
+
+LOGGER = logging.getLogger(__name__)
+
+
+def _escape_identifier(identifier: str) -> str:
+    """Escape a SQL identifier by doubling any double quotes."""
+    escaped = identifier.replace('"', '""')
+    return f'"{escaped}"'
+
+
+def _escape_string(value: str) -> str:
+    """Escape a SQL string by doubling any single quotes."""
+    escaped = value.replace("'", "''")
+    return f"'{escaped}'"
+
+
+def _convert_python_to_sql(value: object) -> str:
+    """Format a Python value as a SQL literal for use in CALL MCP TOOL."""
+    if value is None:
+        return "NULL"
+    if isinstance(value, bool):
+        return "TRUE" if value else "FALSE"
+    if isinstance(value, (int, float)):
+        if isinstance(value, float) and not math.isfinite(value):
+            msg = f"Non-finite float value {value!r} cannot be represented as a SQL literal"
+            raise ValueError(msg)
+        return str(value)
+    if isinstance(value, (dict, list)):
+        return _escape_string(json.dumps(value))
+    return _escape_string(str(value))
+
+
+# Maps the base SQL type name to a Python type.  Complex types (DATE, UUID, etc.)
+# round-trip cleanly as strings.
+_SQL_TYPE_TO_PYTHON: dict[str, type] = {
+    "VARCHAR": str,
+    "CHAR": str,
+    "CHARACTER": str,
+    "TEXT": str,
+    "INT": int,
+    "INTEGER": int,
+    "BIGINT": int,
+    "SMALLINT": int,
+    "TINYINT": int,
+    "BYTE": int,
+    "FLOAT": float,
+    "REAL": float,
+    "DOUBLE": float,
+    "DECIMAL": float,
+    "NUMERIC": float,
+    "BOOLEAN": bool,
+}
+
+
+def _sql_type_to_python(sql_type: str, nullable: bool) -> object:
+    """Return the Python type for an Ocient SQL type string."""
+    base = sql_type.split("(")[0].strip().upper()
+    py_type: type = _SQL_TYPE_TO_PYTHON.get(base, str)
+    return Optional[py_type] if nullable else py_type
+
+
+def _make_db_tool(
+    schema: str,
+    name: str,
+    description: str,
+    arg_names: list[str],
+    arg_types: list[str],
+    arg_nullability: list[bool],
+) -> FunctionTool:
+    """Create a FunctionTool that dispatches to a database SQL MCP tool.
+
+    The returned tool has the MCP name ``{schema}__{name}`` and a proper
+    parameter schema derived from the tool's argument metadata in
+    ``sys.mcp_tools``.
+    """
+    tool_name = f"{schema}__{name}"
+
+    # Build the tool function as a **kwargs closure so the body is generic,
+    # but give it a custom __signature__ so FastMCP sees the named parameters
+    # and generates the correct JSON schema.
+    def _closure(s: str, n: str, names: list[str]) -> FunctionType:
+        async def tool_fn(ctx: Context, **kwargs: object) -> StatementResult:
+            args = [kwargs.get(arg_name) for arg_name in names]
+            formatted_args = ", ".join(_convert_python_to_sql(a) for a in args)
+            sql = f"CALL MCP TOOL {_escape_identifier(s)}.{_escape_identifier(n)}({formatted_args})"
+            return await execute_statement(ctx, sql)
+
+        return tool_fn  # type: ignore[return-value]
+
+    fn = _closure(schema, name, arg_names)
+    fn.__name__ = tool_name
+    fn.__qualname__ = tool_name
+
+    # Construct a signature with proper named parameters so FastMCP derives the
+    # right JSON schema.  inspect.signature() respects __signature__, which is
+    # used by FunctionTool.from_function → ParsedFunction.from_function.
+    params: list[inspect.Parameter] = [
+        inspect.Parameter("ctx", inspect.Parameter.POSITIONAL_OR_KEYWORD, annotation=Context)
+    ]
+    annotations: dict[str, object] = {"ctx": Context, "return": StatementResult}
+
+    for arg_name, sql_type, nullable in zip(arg_names, arg_types, arg_nullability):
+        py_type = _sql_type_to_python(sql_type, nullable)
+        params.append(
+            inspect.Parameter(
+                arg_name,
+                inspect.Parameter.POSITIONAL_OR_KEYWORD,
+                annotation=py_type,
+                default=None if nullable else inspect.Parameter.empty,
+            )
+        )
+        annotations[arg_name] = py_type
+
+    fn.__signature__ = inspect.Signature(params, return_annotation=StatementResult)  # type: ignore[attr-defined]
+    fn.__annotations__ = annotations
+
+    return FunctionTool.from_function(fn, name=tool_name, description=description)
+
+
+class OcientMcpToolProvider(Provider):
+    """Dynamically exposes SQL MCP tools defined in ``sys.mcp_tools`` as MCP tools.
+
+    Each database tool is exposed with the name ``{schema}__{name}`` so clients
+    can call it directly without going through a separate discovery step.
+    """
+
+    def __init__(self, server: OcientMCPServer) -> None:
+        super().__init__()
+        self._server = server
+
+    async def _list_tools(self) -> Sequence[Tool]:
+        is_http = self._server.get_transport() == TransportMode.HTTP
+        if is_http:
+            conn = self._server.connect_for_discovery()
+        else:
+            conn = self._server.fetch_stdio_connection()
+
+        if conn is None:
+            LOGGER.debug("No connection available; skipping dynamic SQL MCP tool listing")
+            return []
+
+        try:
+            with conn.cursor() as cursor:
+                cursor.execute(
+                    "SELECT schema, name, description, "
+                    "argument_names, argument_types, argument_nullability "
+                    "FROM sys.mcp_tools ORDER BY schema, name"
+                )
+                if cursor.description is None:
+                    return []
+                rows = cursor.fetchall()
+        except Exception as e:
+            LOGGER.warning("Failed to query sys.mcp_tools: %s", e)
+            return []
+        finally:
+            if is_http:
+                conn.close()
+
+        tools: list[Tool] = []
+        for row in rows:
+            db_schema, db_name, description, arg_names_raw, arg_types_raw, arg_nullability_raw = row
+            arg_names = list(arg_names_raw or [])
+            arg_types = list(arg_types_raw or [])
+            arg_nullability = [bool(x) for x in (arg_nullability_raw or [])]
+
+            try:
+                tool = _make_db_tool(
+                    db_schema,
+                    db_name,
+                    description or "",
+                    arg_names,
+                    arg_types,
+                    arg_nullability,
+                )
+                tools.append(tool)
+            except Exception as e:
+                LOGGER.warning("Failed to create tool for %s.%s: %s", db_schema, db_name, e)
+
+        LOGGER.info("Discovered %d SQL MCP tools from sys.mcp_tools", len(tools))
+        return tools

ocientmcp/server.py

@@ -0,0 +1,344 @@
+import contextlib
+import logging
+import sys
+from collections.abc import Iterator
+from dataclasses import dataclass
+from datetime import datetime, timedelta, timezone
+from enum import Enum
+from typing import Any, Callable, TypeVar
+
+from fastmcp import Context, FastMCP
+from fastmcp.server.dependencies import get_http_headers
+from fastmcp.server.middleware import CallNext, Middleware, MiddlewareContext
+from mcp.types import AnyFunction
+from pydantic import BaseModel, ConfigDict, model_validator
+
+import pyocient
+from ocientmcp.sso import get_google_id_token
+from ocientmcp.token_cache import clear_token, load_token, save_token
+from pyocient.api import TLSArgType
+
+LOGGER = logging.getLogger(__name__)
+
+F = TypeVar("F", bound=AnyFunction)
+
+
+class TransportMode(Enum):
+    STDIO = "stdio"
+    STDIO_SSO = "stdio_sso"
+    HTTP = "http"
+
+
+@dataclass(frozen=True, kw_only=True, slots=True)
+class SSOConfig:
+    """Configuration for SSO-based STDIO connections."""
+
+    hosts: str
+    database: str
+    port: int = 4050
+    tls: TLSArgType = "unverified"
+    identity_provider: str | None = None
+
+
+class ConnectionCredentials(BaseModel):
+    dsn: str
+    token_source: str | None = None
+
+    model_config = ConfigDict(frozen=True)
+
+    @model_validator(mode="before")
+    @classmethod
+    def coerce_input(cls, data: object) -> object:
+        if isinstance(data, str):
+            return {"dsn": data}
+        return data
+
+
+@dataclass
+class HttpSessionInfo:
+    connection_credentials: ConnectionCredentials
+    connection: pyocient.Connection | None = None
+    last_used: datetime | None = None
+
+
+class OcientMCPServer:
+    def __init__(self) -> None:
+        self._transport: TransportMode = TransportMode.STDIO
+        self._stdio_dsn: str | None = None
+        self._stdio_connection: pyocient.Connection | None = None
+        self._sso_config: SSOConfig | None = None
+        self._http_apiKey_to_connectionCredentials: dict[str, ConnectionCredentials] = {}
+        self._http_sessions: dict[str, HttpSessionInfo] = {}
+        self._http_connection_ttl = timedelta(minutes=30)
+
+        self.MCP: FastMCP[Context] | None = None
+
+    # -----------------------------------------
+    # Initialization order:
+    # 1. Create MCP
+    # 2. Set transport mode (stdio or http)
+    # 3. Register Ocient provider
+    # 4. Run the MCP
+    # -----------------------------------------
+
+    def initialize_stdio(self, dsn: str | None) -> None:
+        self._transport = TransportMode.STDIO
+        self._stdio_dsn = dsn
+
+    def initialize_stdio_sso(self, sso_config: SSOConfig) -> None:
+        self._transport = TransportMode.STDIO_SSO
+        self._sso_config = sso_config
+
+    class HttpMiddleware(Middleware):
+        def __init__(self, server: "OcientMCPServer"):
+            super().__init__()
+            self.server = server
+
+        async def on_request(self, context: MiddlewareContext[Any], call_next: CallNext[Any, Any]) -> Any:  # type: ignore [explicit-any]
+            # Close any idle connections based on TTL
+            now = datetime.now(timezone.utc)
+            for session_id, session in self.server._http_sessions.items():
+                if session.connection is not None:
+                    assert session.last_used is not None
+                    if now - session.last_used > self.server._http_connection_ttl:
+                        try:
+                            idle_time = now - session.last_used
+                            LOGGER.info(f"Closing idle connection for session {session_id} (idle for: {idle_time})")
+                            session.connection.close()
+                        except Exception as e:
+                            LOGGER.warning(f"Failed to close idle connection for session {session_id}: {e}")
+                        session.connection = None
+
+            # Use the request header to see if we have an api key
+            headers = get_http_headers(include_all=True)
+            auth = headers.get("Authorization") or headers.get("authorization")
+
+            # Use the session id to see if we have an existing session
+            fastmcp_context = context.fastmcp_context
+            if fastmcp_context is None:
+                LOGGER.error("HTTP request missing request context")
+                raise ValueError("HTTP request missing request context")
+
+            # https://gofastmcp.com/servers/context#request-context-availability
+            if fastmcp_context.request_context is None:
+                LOGGER.debug("Session has not been established yet, doing nothing")
+                return await call_next(context)
+
+            session_id = fastmcp_context.session_id
+            if not session_id:
+                LOGGER.error("HTTP request missing session ID")
+                raise ValueError("HTTP request missing session ID")
+
+            session_info = self.server._http_sessions.get(session_id)
+
+            if session_info is None:
+                if auth is None:
+                    # If we have no auth header and no existing session, deny access
+                    LOGGER.warning("Unauthorized: No auth header and no existing session")
+                    raise PermissionError("Unauthorized: No auth header and no existing session")
+                else:
+                    # If we have an auth header but no existing session, create a new session
+                    if not auth.startswith("Bearer "):
+                        LOGGER.warning("Unauthorized: Invalid auth header")
+                        raise PermissionError("Unauthorized: Invalid auth header")
+
+                    api_key = auth.removeprefix("Bearer ")
+
+                    connection_credentials = self.server._http_apiKey_to_connectionCredentials.get(api_key)
+                    if connection_credentials is None:
+                        LOGGER.warning("Unauthorized: Invalid API key")
+                        raise PermissionError("Unauthorized: Invalid API key")
+
+                    session_info = HttpSessionInfo(connection_credentials=connection_credentials)
+                    self.server._http_sessions[session_id] = session_info
+
+            await fastmcp_context.set_state("session_id", session_id)
+            return await call_next(context)
+
+    def initialize_http(self, auths: dict[str, ConnectionCredentials]) -> None:
+        self._transport = TransportMode.HTTP
+        self._http_apiKey_to_connectionCredentials = auths
+
+        if self.MCP is None:
+            raise RuntimeError("MCP must be initialized before running HTTP server.")
+
+        self.MCP.add_middleware(OcientMCPServer.HttpMiddleware(server=self))
+
+    def tool(self) -> Callable[[F], F]:
+        m = self.MCP
+        if m is None:
+            raise RuntimeError("MCP must be initialized before registering tools.")
+
+        fastmcp_decorator = m.tool()
+
+        def wrapping_decorator(fn: F) -> F:
+            fastmcp_decorator(fn)
+            return fn
+
+        return wrapping_decorator
+
+    # -----------------------------------------
+    # Connections
+    # -----------------------------------------
+
+    @staticmethod
+    @contextlib.contextmanager
+    def _protect_stdout() -> Iterator[None]:
+        """Redirect stdout to stderr during SSO flow.
+
+        MCP STDIO transport uses stdout as the protocol channel. Any stray
+        print() from pyocient's SSO flow (e.g. device-flow prompts, browser
+        failure warnings) would corrupt the MCP protocol stream.
+        """
+        original = sys.stdout
+        sys.stdout = sys.stderr
+        try:
+            yield
+        finally:
+            sys.stdout = original
+
+    def get_transport(self) -> TransportMode:
+        return self._transport
+
+    def _credentials_from_request_headers(self) -> ConnectionCredentials | None:
+        """Return the ConnectionCredentials for the API key in the current request headers, or None."""
+        headers = get_http_headers(include_all=True)
+        auth = headers.get("Authorization") or headers.get("authorization")
+        if not auth or not auth.startswith("Bearer "):
+            return None
+        api_key = auth.removeprefix("Bearer ")
+        return self._http_apiKey_to_connectionCredentials.get(api_key)
+
+    def connect_for_discovery(self) -> pyocient.Connection | None:
+        """Open a fresh connection for HTTP tool discovery using the request's API key.
+
+        Looks up credentials from the Authorization header so that discovery
+        always uses the calling user's credentials rather than an arbitrary
+        shared credential.
+        """
+        creds = self._credentials_from_request_headers()
+        if creds is None:
+            LOGGER.warning("No valid credentials in request headers; cannot establish discovery connection")
+            return None
+        try:
+            return self._connect_with_credentials(creds)
+        except Exception as e:
+            LOGGER.warning("Failed to establish discovery connection: %s", e)
+            return None
+
+    def _connect_with_credentials(self, creds: ConnectionCredentials) -> pyocient.Connection:
+        """Open a pyocient connection using the given credentials.
+
+        Handles both plain DSN connections and Google ID token auth so that all
+        pyocient.connect() calls go through a single credential management path.
+        """
+        if creds.token_source is not None:
+            return pyocient.connect(dsn=creds.dsn, user="id_token", password=get_google_id_token(creds.token_source))
+        return pyocient.connect(dsn=creds.dsn)
+
+    def fetch_stdio_connection(self) -> pyocient.Connection | None:
+        """One global connection for STDIO mode (DSN-based)"""
+
+        if self._stdio_connection:
+            return self._stdio_connection
+
+        try:
+            self._stdio_connection = pyocient.connect(self._stdio_dsn)
+            return self._stdio_connection
+        except Exception as e:
+            LOGGER.error(f"Failed to connect (stdio): {e}")
+            return None
+
+    @staticmethod
+    def _build_sso_dsn(cfg: SSOConfig) -> str:
+        """Build a DSN string for SSO connections.
+
+        pyocient parses identity_provider and handshake from the DSN query
+        string, so we construct a DSN rather than passing keyword args.
+        """
+        params = f"tls={cfg.tls}&handshake=sso"
+        if cfg.identity_provider:
+            params += f"&identityprovider={cfg.identity_provider}"
+        return f"ocient://@{cfg.hosts}:{cfg.port}/{cfg.database}?{params}"
+
+    def fetch_stdio_sso_connection(self) -> pyocient.Connection | None:
+        """SSO-based STDIO connection with token caching."""
+
+        if self._stdio_connection:
+            return self._stdio_connection
+
+        assert self._sso_config is not None
+        cfg = self._sso_config
+
+        # Try cached security token first
+        cached_token = load_token(cfg.hosts, cfg.database, cfg.port)
+        if cached_token is not None:
+            try:
+                LOGGER.info("Attempting connection with cached security token")
+                conn = pyocient.connect(
+                    host=f"{cfg.hosts}:{cfg.port}",
+                    database=cfg.database,
+                    tls=cfg.tls,
+                    security_token=cached_token,
+                )
+                # Update cache with the (possibly refreshed) token from the server
+                if conn.security_token is not None:
+                    save_token(cfg.hosts, cfg.database, cfg.port, conn.security_token)
+                self._stdio_connection = conn
+                LOGGER.info("Connected using cached security token")
+                return conn
+            except Exception as e:
+                LOGGER.info("Cached token failed (%s), falling back to SSO flow", e)
+                clear_token(cfg.hosts, cfg.database, cfg.port)
+
+        # Full SSO browser flow — protect stdout from stray prints
+        try:
+            dsn = self._build_sso_dsn(cfg)
+            LOGGER.info("Starting SSO authentication for %s:%d/%s", cfg.hosts, cfg.port, cfg.database)
+            with self._protect_stdout():
+                conn = pyocient.connect(dsn)
+
+            # Cache the security token for future connections
+            if conn.security_token is not None:
+                save_token(cfg.hosts, cfg.database, cfg.port, conn.security_token)
+
+            self._stdio_connection = conn
+            LOGGER.info("Connected via SSO")
+            return conn
+        except Exception as e:
+            LOGGER.error("SSO connection failed: %s", e)
+            return None
+
+    async def fetch_http_connection(self, ctx: Context) -> pyocient.Connection | None:
+        """One connection per API key, managed by HttpMiddleware"""
+
+        session_id = await ctx.get_state("session_id")
+        if session_id is None:
+            raise ValueError("Session ID not found in context")
+
+        session_info = self._http_sessions.get(session_id)
+        assert session_info is not None
+
+        if session_info.connection is None:
+            LOGGER.debug("Detected unopened or closed connection, (re)connecting")
+            try:
+                session_info.connection = self._connect_with_credentials(session_info.connection_credentials)
+                session_info.last_used = datetime.now(timezone.utc)
+            except Exception as e:
+                LOGGER.error(f"Failed to reconnect to Ocient database: {e}")
+                raise Exception("Failed to reconnect to Ocient database")
+        else:
+            # If connection is open, just update last used time
+            session_info.last_used = datetime.now(timezone.utc)
+
+        return session_info.connection
+
+    async def get_active_connection(self, ctx: Context) -> pyocient.Connection | None:
+        if self.get_transport() == TransportMode.HTTP:
+            return await self.fetch_http_connection(ctx)
+        if self.get_transport() == TransportMode.STDIO_SSO:
+            return self.fetch_stdio_sso_connection()
+        return self.fetch_stdio_connection()
+
+
+OCIENTMCP = OcientMCPServer()

ocientmcp/server_http.py

@@ -0,0 +1,55 @@
+import argparse
+import logging
+
+from ocientmcp.transports import run_http
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--verbose",
+        "-v",
+        action="store_true",
+        help="Enable verbose logging",
+        default=False,
+    )
+    parser.add_argument(
+        "--auth-file",
+        "-a",
+        default="/var/opt/ocient/mcp/auth.yaml",
+        help="Path to the secure authentication file",
+    )
+    parser.add_argument(
+        "--host",
+        default="127.0.0.1",
+        help="Address to bind the HTTP server to. Use 0.0.0.0 to accept connections on every interface",
+    )
+    parser.add_argument(
+        "--port",
+        "-p",
+        type=int,
+        default=8000,
+        help="Port to run the HTTP server on",
+    )
+
+    parser.add_argument(
+        "--tls-cert",
+        help="Path to a PEM-encoded certificate chain. Serves HTTPS instead of HTTP when supplied",
+    )
+    parser.add_argument(
+        "--tls-key",
+        help="Path to the PEM-encoded private key for --tls-cert, if the key is not in the certificate file",
+    )
+
+    args = parser.parse_args()
+
+    if args.tls_key and not args.tls_cert:
+        parser.error("--tls-key requires --tls-cert")
+
+    logging.basicConfig(level=logging.DEBUG if args.verbose else logging.INFO)
+
+    run_http(args.auth_file, args.host, args.port, args.tls_cert, args.tls_key)
+
+
+if __name__ == "__main__":
+    main()

ocientmcp/server_http_proxy.py

@@ -0,0 +1,70 @@
+import argparse
+import logging
+import os
+
+import httpx
+from fastmcp import FastMCP
+from fastmcp.client.transports import StreamableHttpTransport
+from fastmcp.server.proxy import ProxyClient
+from mcp.shared._httpx_utils import MCP_DEFAULT_SSE_READ_TIMEOUT, MCP_DEFAULT_TIMEOUT
+
+LOGGER = logging.getLogger(__name__)
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--verbose",
+        "-v",
+        action="store_true",
+        help="Enable verbose logging",
+        default=False,
+    )
+    parser.add_argument("--url", required=True)
+    args = parser.parse_args()
+
+    logging.basicConfig(level=logging.DEBUG if args.verbose else logging.INFO)
+
+    api_key = os.environ.get("OCIENT_MCP_API_KEY")
+    if not api_key:
+        raise RuntimeError("OCIENT_MCP_API_KEY must be set")
+
+    # Normalize url to /mcp
+    url = args.url.rstrip("/")
+    if not url.endswith("/mcp"):
+        url = url + "/mcp"
+
+    LOGGER.info(f"Proxy connecting to MCP endpoint: {url}")
+
+    # StreamableHttpTransport creates an httpx.AsyncClient with no explicit timeout,
+    # inheriting httpx's default of 5 seconds. The MCP server's SSE keepalive pings
+    # are sent every 15 seconds, so the client times out before receiving any data
+    # during long-running operations. Use MCP-recommended timeouts instead.
+    def _create_http_client(
+        headers: dict[str, str] | None = None,
+        timeout: httpx.Timeout | None = None,
+        auth: httpx.Auth | None = None,
+        **_extra: object,
+    ) -> httpx.AsyncClient:
+        if timeout is None:
+            timeout = httpx.Timeout(MCP_DEFAULT_TIMEOUT, read=MCP_DEFAULT_SSE_READ_TIMEOUT)
+        return httpx.AsyncClient(
+            headers=headers,
+            timeout=timeout,
+            auth=auth,
+            follow_redirects=True,
+        )
+
+    transport = StreamableHttpTransport(
+        url=url,
+        headers={"Authorization": f"Bearer {api_key}"},
+        httpx_client_factory=_create_http_client,
+    )
+
+    proxy = FastMCP.as_proxy(ProxyClient(transport), name="Proxy Http Server")
+
+    proxy.run(transport="stdio", show_banner=False)
+
+
+if __name__ == "__main__":
+    main()

ocientmcp/server_stdio.py

@@ -0,0 +1,90 @@
+import argparse
+import logging
+import os
+from typing import get_args
+
+from ocientmcp.server import SSOConfig
+from ocientmcp.transports import run_stdio, run_stdio_sso
+from pyocient.api import TLSArgType
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Ocient MCP Server (stdio)")
+    parser.add_argument(
+        "--verbose",
+        "-v",
+        action="store_true",
+        help="Enable verbose logging",
+        default=False,
+    )
+    parser.add_argument(
+        "--hosts",
+        help="Comma-separated list of Ocient hosts (e.g. host1.ocient.com,host2.ocient.com). "
+        "Can also be set via OCIENT_HOSTS env var.",
+    )
+    parser.add_argument(
+        "--database",
+        help="Ocient database name. Can also be set via OCIENT_DATABASE env var.",
+    )
+    parser.add_argument(
+        "--port",
+        type=int,
+        default=None,
+        help="Ocient port (default: 4050). Can also be set via OCIENT_PORT env var.",
+    )
+    parser.add_argument(
+        "--tls",
+        default=None,
+        choices=["unverified", "on"],
+        help="TLS mode (default: unverified). Can also be set via OCIENT_TLS env var.",
+    )
+    parser.add_argument(
+        "--identity-provider",
+        default=None,
+        help="SSO identity provider name (e.g. 'okta', 'google'). "
+        "Can also be set via OCIENT_IDENTITY_PROVIDER env var.",
+    )
+    args = parser.parse_args()
+
+    logging.basicConfig(level=logging.DEBUG if args.verbose else logging.INFO)
+
+    dsn = os.environ.get("OCIENT_DSN")
+
+    if dsn:
+        # Legacy mode: full DSN with embedded credentials
+        run_stdio(dsn)
+        return
+
+    # SSO mode: hosts + database, credentials via browser SSO
+    hosts = args.hosts or os.environ.get("OCIENT_HOSTS")
+    database = args.database or os.environ.get("OCIENT_DATABASE")
+
+    if not hosts or not database:
+        raise RuntimeError(
+            "Either OCIENT_DSN must be set, or both OCIENT_HOSTS and OCIENT_DATABASE "
+            "(via env vars or --hosts/--database flags)"
+        )
+
+    port = args.port if args.port is not None else int(os.environ.get("OCIENT_PORT", "4050"))
+    tls_raw = args.tls if args.tls is not None else os.environ.get("OCIENT_TLS", "unverified")
+    valid_tls_values = get_args(TLSArgType)
+    if tls_raw not in valid_tls_values:
+        raise RuntimeError(f"Invalid TLS mode '{tls_raw}'. Must be one of: {', '.join(valid_tls_values)}")
+    tls: TLSArgType = tls_raw  # type: ignore[assignment]
+    identity_provider = (
+        args.identity_provider if args.identity_provider is not None else os.environ.get("OCIENT_IDENTITY_PROVIDER")
+    )
+
+    sso_config = SSOConfig(
+        hosts=hosts,
+        database=database,
+        port=port,
+        tls=tls,
+        identity_provider=identity_provider,
+    )
+
+    run_stdio_sso(sso_config)
+
+
+if __name__ == "__main__":
+    main()

ocientmcp/sso.py

@@ -0,0 +1,16 @@
+import requests  # required by google.auth.transport.requests.Request
+from google.auth.transport.requests import Request
+from google.oauth2 import service_account
+
+
+def get_google_id_token(service_account_file: str) -> str:
+    creds = service_account.IDTokenCredentials.from_service_account_file(  # type:ignore [no-untyped-call]
+        service_account_file,
+        target_audience="https://accounts.google.com",
+    )
+
+    if not creds.token or creds.expired:
+        creds.refresh(Request())
+
+    assert isinstance(creds.token, str)
+    return creds.token

ocientmcp/token_cache.py

@@ -0,0 +1,106 @@
+"""Local token cache for SSO security tokens.
+
+Caches pyocient SecurityTokens to disk so users don't need to re-authenticate
+through the browser on every MCP server restart. Tokens are stored in
+$XDG_CACHE_HOME/ocientmcp/mcp-token-cache.json (or ~/.cache/ocientmcp/) with 0600
+permissions on POSIX systems.
+"""
+
+import hashlib
+import json
+import logging
+import os
+import stat
+from dataclasses import asdict
+from pathlib import Path
+from typing import cast
+
+from pyocient.api import SecurityToken
+
+LOGGER = logging.getLogger(__name__)
+
+if xdg_cache_home := os.getenv("XDG_CACHE_HOME"):
+    _HOME_CACHE = Path(xdg_cache_home)
+else:
+    _HOME_CACHE = Path.home() / ".cache"
+_CACHE_DIR = _HOME_CACHE / "ocientmcp"
+_CACHE_FILE = _CACHE_DIR / "mcp-token-cache.json"
+
+
+def _cache_key(hosts: str, database: str, port: int) -> str:
+    """Deterministic key for a given cluster + database."""
+    raw = f"{hosts}:{port}/{database}"
+    return hashlib.sha256(raw.encode()).hexdigest()[:16]
+
+
+def _ensure_cache_dir() -> None:
+    _CACHE_DIR.mkdir(mode=0o700, parents=True, exist_ok=True)
+
+
+def _read_cache() -> dict[str, object]:
+    if not _CACHE_FILE.exists():
+        return {}
+
+    # Validate permissions before reading (POSIX only — Windows uses ACLs)
+    if os.name == "posix":
+        st = _CACHE_FILE.stat()
+        mode = stat.S_IMODE(st.st_mode)
+        if mode != 0o600:
+            LOGGER.warning("Token cache file has insecure permissions %s, ignoring", oct(mode))
+            return {}
+
+    try:
+        with open(_CACHE_FILE, "r") as f:
+            return cast(dict[str, object], json.load(f))
+    except (json.JSONDecodeError, OSError) as e:
+        LOGGER.warning("Failed to read token cache: %s", e)
+        return {}
+
+
+def _write_cache(data: dict[str, object]) -> None:
+    _ensure_cache_dir()
+
+    # Write atomically via temp file
+    tmp = _CACHE_FILE.with_suffix(".tmp")
+    try:
+        fd = os.open(str(tmp), os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600)
+        with os.fdopen(fd, "w") as f:
+            json.dump(data, f, indent=2)
+        tmp.rename(_CACHE_FILE)
+    except OSError as e:
+        LOGGER.warning("Failed to write token cache: %s", e)
+        tmp.unlink(missing_ok=True)
+
+
+def load_token(hosts: str, database: str, port: int) -> SecurityToken | None:
+    """Load a cached SecurityToken for the given cluster, or None if not found."""
+    cache = _read_cache()
+    key = _cache_key(hosts, database, port)
+    entry = cache.get(key)
+    if entry is None:
+        return None
+
+    try:
+        return SecurityToken.from_json(entry)  # type: ignore[arg-type]
+    except (KeyError, TypeError) as e:
+        LOGGER.warning("Corrupt cache entry for key %s: %s", key, e)
+        return None
+
+
+def save_token(hosts: str, database: str, port: int, token: SecurityToken) -> None:
+    """Persist a SecurityToken to the local cache."""
+    cache = _read_cache()
+    key = _cache_key(hosts, database, port)
+    cache[key] = asdict(token)
+    _write_cache(cache)
+    LOGGER.info("Cached security token for %s:%d/%s", hosts, port, database)
+
+
+def clear_token(hosts: str, database: str, port: int) -> None:
+    """Remove a cached token (e.g. after it fails to authenticate)."""
+    cache = _read_cache()
+    key = _cache_key(hosts, database, port)
+    if key in cache:
+        del cache[key]
+        _write_cache(cache)
+        LOGGER.info("Cleared cached token for %s:%d/%s", hosts, port, database)

ocientmcp/tools.py

@@ -0,0 +1,59 @@
+import logging
+from typing import Union
+
+from fastmcp import Context
+from fastmcp.exceptions import ToolError
+
+from ocientmcp.models import StatementResult
+from ocientmcp.server import OCIENTMCP
+from ocientmcp.utils import serialize_result
+
+LOGGER = logging.getLogger(__name__)
+
+
+@OCIENTMCP.tool()
+async def execute_statement(
+    ctx: Context,
+    statement: str,
+    parameters: Union[dict[str, object], tuple[object, ...]] | None = None,
+) -> StatementResult:
+    """Executes an arbitrary SQL query or command and returns the result set or
+    row count. Use this only when no specific tool (visible in tools/list)
+    covers the request."""
+    if not statement or not statement.strip():
+        raise ToolError("Statement must not be empty.")
+
+    conn = await OCIENTMCP.get_active_connection(ctx)
+    if not conn:
+        raise ToolError("No active connection.")
+
+    try:
+        with conn.cursor() as cursor:
+            if parameters:
+                cursor.execute(statement, parameters)
+            else:
+                cursor.execute(statement)
+
+            # Handle DDL/INSERT statements with no result set
+            if cursor.description is None:
+                return StatementResult(
+                    message="Statement executed.",
+                    row_count=cursor.rowcount,
+                    columns=[],
+                    rows=[],
+                )
+
+            columns = [desc[0] for desc in cursor.description]
+            rows_raw = cursor.fetchall()
+
+            rows = [{col: serialize_result(val) for col, val in zip(columns, row)} for row in rows_raw]
+
+            return StatementResult(
+                columns=columns,
+                rows=rows,
+                row_count=len(rows),
+            )
+
+    except Exception as e:
+        LOGGER.error("Statement failed: %s", e)
+        raise ToolError(f"Statement failed: {e}") from e

ocientmcp/transports.py

@@ -0,0 +1,127 @@
+import logging
+import os
+import stat
+
+import yaml
+from fastmcp import FastMCP
+from pydantic import TypeAdapter
+
+from ocientmcp.server import OCIENTMCP, ConnectionCredentials, SSOConfig
+
+LOGGER = logging.getLogger(__name__)
+
+
+def run_stdio(dsn: str | None = None) -> None:
+    LOGGER.info("Starting Ocient MCP server (stdio mode)...")
+
+    OCIENTMCP.MCP = FastMCP(name="Ocient")
+    OCIENTMCP.initialize_stdio(dsn)
+
+    # Import decorators/provider after OcientMCP is initialized
+    import ocientmcp.tools
+    from ocientmcp.provider import OcientMcpToolProvider
+
+    OCIENTMCP.MCP.add_provider(OcientMcpToolProvider(OCIENTMCP))
+
+    OCIENTMCP.MCP.run(transport="stdio", show_banner=False)
+
+
+def run_stdio_sso(sso_config: SSOConfig) -> None:
+    LOGGER.info("Starting Ocient MCP server (stdio SSO mode)...")
+
+    OCIENTMCP.MCP = FastMCP(name="Ocient")
+    OCIENTMCP.initialize_stdio_sso(sso_config)
+
+    # Import decorators/provider after OcientMCP is initialized
+    import ocientmcp.tools
+    from ocientmcp.provider import OcientMcpToolProvider
+
+    OCIENTMCP.MCP.add_provider(OcientMcpToolProvider(OCIENTMCP))
+
+    OCIENTMCP.MCP.run(transport="stdio", show_banner=False)
+
+
+def _validate_auth_file_permissions(path: str) -> None:
+    """Sanity checks to ensure the auth file is secure."""
+
+    if os.path.islink(path):
+        raise PermissionError(f"Auth file '{path}' must not be a symlink")
+
+    st = os.stat(path)
+
+    if st.st_uid != os.getuid():
+        raise PermissionError(f"Auth file '{path}' must be owned by UID {os.getuid()}, but is owned by UID {st.st_uid}")
+
+    mode = stat.S_IMODE(st.st_mode)
+    if mode != 0o600:
+        raise PermissionError(f"Auth file '{path}' must have permissions 600 (rw-------), but has {oct(mode)}")
+
+
+def _tls_uvicorn_config(tls_cert: str | None, tls_key: str | None) -> dict[str, str]:
+    """Builds the uvicorn TLS settings, or an empty mapping to serve plain HTTP."""
+
+    if tls_cert is None:
+        if tls_key is not None:
+            msg = "A TLS key was supplied without a TLS certificate"
+            raise ValueError(msg)
+        return {}
+
+    config = {"ssl_certfile": tls_cert}
+    if tls_key is not None:
+        config["ssl_keyfile"] = tls_key
+
+    for path in config.values():
+        if not os.path.isfile(path):
+            msg = f"TLS file '{path}' does not exist"
+            raise FileNotFoundError(msg)
+
+    # The key is as sensitive as the auth file, but unlike the auth file it is
+    # managed by external certificate tooling, so warn rather than refuse.
+    key_holder = config.get("ssl_keyfile", config["ssl_certfile"])
+    mode = stat.S_IMODE(os.stat(key_holder).st_mode)
+    if mode & 0o077:
+        LOGGER.warning(
+            "TLS key file '%s' has permissions %s and is readable beyond its owner; consider 600 (rw-------)",
+            key_holder,
+            oct(mode),
+        )
+
+    return config
+
+
+def run_http(
+    auth_file: str,
+    host: str,
+    port: int,
+    tls_cert: str | None = None,
+    tls_key: str | None = None,
+) -> None:
+    LOGGER.info("Starting Ocient MCP server (streamable http mode)...")
+
+    uvicorn_config = _tls_uvicorn_config(tls_cert, tls_key)
+
+    # Load an auth file that should map API_KEY -> DSN
+    _validate_auth_file_permissions(auth_file)
+
+    with open(auth_file, "r") as f:
+        raw_auths = yaml.safe_load(f)
+
+    adapter = TypeAdapter(dict[str, ConnectionCredentials])
+    auths = adapter.validate_python(raw_auths)
+
+    OCIENTMCP.MCP = FastMCP(name="Ocient")
+    OCIENTMCP.initialize_http(auths)
+
+    # Import decorators/provider after OcientMCP is initialized
+    import ocientmcp.tools
+    from ocientmcp.provider import OcientMcpToolProvider
+
+    OCIENTMCP.MCP.add_provider(OcientMcpToolProvider(OCIENTMCP))
+
+    OCIENTMCP.MCP.run(
+        transport="streamable-http",
+        show_banner=False,
+        host=host,
+        port=port,
+        uvicorn_config=uvicorn_config,
+    )

ocientmcp/utils.py

@@ -0,0 +1,11 @@
+def serialize_result(obj: object) -> object:
+    """Convert Python/DB objects into JSON-serializable structures"""
+    if isinstance(obj, (str, int, float, bool, type(None))):
+        return obj
+    if isinstance(obj, (list, tuple)):
+        return [serialize_result(x) for x in obj]
+    if isinstance(obj, dict):
+        return {k: serialize_result(v) for k, v in obj.items()}
+    if hasattr(obj, "isoformat"):
+        return obj.isoformat()
+    return str(obj)

ocientmcp-1.2.2.dist-info/METADATA

@@ -0,0 +1,32 @@
+Metadata-Version: 2.1
+Name: ocientmcp
+Author: Ocient Inc
+Author-email: info@ocient.com
+Home-page: https://www.ocient.com/
+License: Apache License, Version 2.0
+Description-Content-Type: text/markdown
+Summary: Ocient MCP Server - Model Context Protocol integration for the OcientAIQ™ Unified Data Platform
+Project-URL: documentation, https://docs.ocient.com/
+Classifier: Development Status :: 5 - Production/Stable
+Classifier: Topic :: Database
+Classifier: Topic :: Database :: Front-Ends
+Classifier: Topic :: Software Development
+Classifier: Topic :: Software Development :: Libraries :: Python Modules
+Classifier: Intended Audience :: Developers
+Classifier: License :: OSI Approved :: Apache Software License
+Classifier: Typing :: Typed
+Requires-Python: >=3.10,<4
+Requires-Dist: fastmcp>=3.4.0,<4
+Requires-Dist: google-auth
+Requires-Dist: httpx>=0.28.1
+Requires-Dist: mcp>=1.29.0,<2
+Requires-Dist: pydantic
+Requires-Dist: pyyaml>=6.0
+Requires-Dist: requests
+Requires-Dist: pyocient>=3.7.0
+Version: 1.2.2
+
+# OcientMCP
+
+The OcientMCP server enables AI agents and tools such as Windsurf or Cascade to interact with the OcientAIQ™ Unified Data Platform using the Model Context Protocol (MCP). This server provides a standardized interface for querying, exploring, and analyzing data in the Ocient System.
+

ocientmcp-1.2.2.dist-info/RECORD

@@ -0,0 +1,18 @@
+ocientmcp/__init__.py,sha256=uc26kjoc4Ek4C1yYtdcxZ7Apv5aY4JtbcE965GxMyGI,869
+ocientmcp/models.py,sha256=WJi0e832DDE3fM_K2NyRbEN8ZBnunxq_A-KiVKTS0KE,675
+ocientmcp/pkg_version.py,sha256=oyfmQVJgxP0NR5vs9rKQyTTg8V_VjOBJ8hU4OYivRwA,118
+ocientmcp/provider.py,sha256=-5mogZ_yhXPXtIqCMBG2EIpKotUHfm5fhQf6MQotM9Y,7033
+ocientmcp/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+ocientmcp/server.py,sha256=v6BJL93xuwI1426hwWXkDRoHMB_M-wMuf6yfgFRk1BE,13956
+ocientmcp/server_http.py,sha256=hF4luc6x8lI3z1xqPBd_OByBleypXopDhxucogmTYzQ,1444
+ocientmcp/server_http_proxy.py,sha256=ZFgfKQunWpuiazwN8KI7dxRZd-2i6YBaaOQiqLwsXzc,2174
+ocientmcp/server_stdio.py,sha256=Lxujb9-oP-Y4ODQUrhFfxxFZxsW0DBFapHq_ce2xGOQ,2868
+ocientmcp/sso.py,sha256=r0jfheosQJYS1DwSZgbR1xhhJSl8TedDLl7DvxRPCLA,561
+ocientmcp/token_cache.py,sha256=uLOcZbG7UHIs8YzlwZ6XNGRFWqvF5nEGCrGLkEW7aBw,3414
+ocientmcp/tools.py,sha256=gE5yOyXCgNWC6YSt5y-6ZwJ_j2uW1Ga014a246i52yg,1872
+ocientmcp/transports.py,sha256=sqKRhEMdvkziJitMdDe-2fH5cpdeShJ8J0lxPIzy9Vs,3908
+ocientmcp/utils.py,sha256=8xSoRk9__Vc4tMgl_Gggaom-fusy7DnxndbvdYnDkdU,463
+ocientmcp-1.2.2.dist-info/WHEEL,sha256=sobxWSyDDkdg_rinUth-jxhXHqoNqlmNMJY3aTZn2Us,91
+ocientmcp-1.2.2.dist-info/METADATA,sha256=kyC0LWBE_oO1CTRA4nhkWFpaiX5l5fsr1j7fzYcfQDU,1306
+ocientmcp-1.2.2.dist-info/entry_points.txt,sha256=WYtwwQeGYhGMF6_OjfmJ4GR-VXJL1_9239cqla3J4VA,157
+ocientmcp-1.2.2.dist-info/RECORD,,

ocientmcp-1.2.2.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: bazel-wheelmaker 1.0
+Root-Is-Purelib: true
+Tag: py3-none-any

ocientmcp-1.2.2.dist-info/entry_points.txt

@@ -0,0 +1,4 @@
+[console_scripts]
+ocientmcp-http-proxy=ocientmcp.server_http_proxy:main
+ocientmcp-http=ocientmcp.server_http:main
+ocientmcp-stdio=ocientmcp.server_stdio:main