itti_adb/lab/mcp/test_mcp_token.py

import asyncio
import base64
import json
import os
import threading
import time
import warnings
from dataclasses import dataclass
from typing import Optional
from dotenv import load_dotenv

import httpx

# Oculta solo el aviso conocido del urllib3 generado por dependencias del OCI SDK.
warnings.filterwarnings(
    "ignore",
    message=r"The 'strict' parameter is no longer needed on Python 3\+.*",
    category=FutureWarning,
    module=r"urllib3\.poolmanager",
)

import oci
from mcp import ClientSession
from mcp.client.streamable_http import streamable_http_client

# Load variables from the .env file.
load_dotenv()

# === Environment-based configuration ===
REGION         = os.environ["ADB_REGION"]        # e.g.: "sa-saopaulo-1"
DB_OCID        = os.environ["ADB_OCID"]          # ADB OCID
DB_USER        = os.environ["ADB_USER"]          # e.g.: "MCP_USER"
SECRET_OCID    = os.environ["ADB_SECRET_OCID"]   # Vault secret OCID

BASE_URL  = f"https://dataaccess.adb.{REGION}.oraclecloudapps.com"
TOKEN_URL = f"{BASE_URL}/adb/auth/v1/databases/{DB_OCID}/token"
MCP_URL   = f"{BASE_URL}/adb/mcp/v1/databases/{DB_OCID}"


# ---------------------------------------------------------------------------
# OCI Vault: password retrieval
# ---------------------------------------------------------------------------

def _build_secrets_client() -> oci.secrets.SecretsClient:
    """
    Tries Instance Principal first (production on an OCI VM);
    falls back to API Key (~/.oci/config) when not running on an OCI instance.
    
    -- uncomment here for using instance principal instead of oci/config file
    
    try:
        signer = oci.auth.signers.InstancePrincipalsSecurityTokenSigner()
        return oci.secrets.SecretsClient(config={}, signer=signer)
    except Exception:
    """
    config = oci.config.from_file()  # reads ~/.oci/config, DEFAULT profile
    config["region"] = os.environ["ADB_REGION"]
    return oci.secrets.SecretsClient(config)


def get_db_password() -> str:
    """Reads the MCP_USER password from OCI Vault."""
    client = _build_secrets_client()
    bundle = client.get_secret_bundle(secret_id=SECRET_OCID).data # type: ignore
    encoded = bundle.secret_bundle_content.content   # base64
    return base64.b64decode(encoded).decode("utf-8")


# ---------------------------------------------------------------------------
# Bearer token cache with automatic refresh
# ---------------------------------------------------------------------------

@dataclass
class _CachedToken:
    value: str
    expires_at: float   # epoch seconds


class TokenManager:
    """
    Caches the bearer token in memory and refreshes it before expiration.
    Thread-safe: multiple threads requesting a token at the same time trigger
    only one refresh.
    """

    # Refresh when the token has less than this remaining lifetime.
    REFRESH_MARGIN_SECONDS = 5 * 60   # 5 minutes
    DEFAULT_TTL_SECONDS    = 60 * 60  # fallback when the API does not return expires_in

    def __init__(self, db_user: str):
        self._db_user = db_user
        self._cached: Optional[_CachedToken] = None
        self._lock = threading.Lock()

    def get_token(self) -> str:
        # Fast path: token is still valid, no heavy lock needed.
        cached = self._cached
        if cached and cached.expires_at - time.time() > self.REFRESH_MARGIN_SECONDS:
            return cached.value

        # Slow path: refresh required.
        with self._lock:
            # Re-check inside the lock because another thread may have refreshed.
            cached = self._cached
            if cached and cached.expires_at - time.time() > self.REFRESH_MARGIN_SECONDS:
                return cached.value

            self._cached = self._fetch_new_token()
            return self._cached.value

    def _fetch_new_token(self) -> _CachedToken:
        password = get_db_password()  # reads from Vault only when refreshing
        payload = {
            "grant_type": "password",
            "username":   self._db_user,
            "password":   password,
        }
        headers = {"Content-Type": "application/json", "Accept": "application/json"}

        resp = httpx.post(TOKEN_URL, json=payload, headers=headers, timeout=30.0)
        resp.raise_for_status()
        data = resp.json()

        token = data["access_token"]
        ttl   = int(data.get("expires_in", self.DEFAULT_TTL_SECONDS))
        expires_at = time.time() + ttl

        print(f"[token] new token obtained, valid for ~{ttl}s")
        return _CachedToken(value=token, expires_at=expires_at)


# ---------------------------------------------------------------------------
# Test case MCP
# ---------------------------------------------------------------------------
async def run_test(token_manager: TokenManager):
    t0 = time.perf_counter()

    auth_headers = {"Authorization": f"Bearer {token_manager.get_token()}"}
    t1 = time.perf_counter()
    print(f"[timing] token obtained: {t1 - t0:.2f}s")

    async with httpx.AsyncClient(
        headers=auth_headers,
        timeout=httpx.Timeout(300.0, connect=30.0),
    ) as http_client:

        async with streamable_http_client(MCP_URL, http_client=http_client) as (read, write, _):
            t2 = time.perf_counter()
            print(f"[timing] transport opened: {t2 - t1:.2f}s")

            async with ClientSession(read, write) as session:
                await session.initialize()
                t3 = time.perf_counter()
                print(f"[timing] session.initialize(): {t3 - t2:.2f}s")

                tools_result = await session.list_tools()
                t4 = time.perf_counter()
                print(f"[timing] list_tools(): {t4 - t3:.2f}s")

                query = "SELECT username FROM all_users WHERE oracle_maintained = 'N'"
                result = await session.call_tool(
                    "RUNSQL",
                    arguments={"SQL_TEXT": query, "OFFSET_ROWS": 0, "FETCH_ROWS": 10},
                )
                t5 = time.perf_counter()
                print(f"[timing] call_tool(): {t5 - t4:.2f}s")
                print(f"[timing] TOTAL: {t5 - t0:.2f}s")

                for block in result.content:
                    if block.type == "text":
                        rows = json.loads(block.text)
                        print(f"\n[result: {len(rows)} row(s)]")
                        for row in rows:
                            print(f"  {row}")


if __name__ == "__main__":
    token_manager = TokenManager(db_user=DB_USER)
    asyncio.run(run_test(token_manager))