MonkeyCanCode commented on code in PR #83:
URL: https://github.com/apache/polaris-tools/pull/83#discussion_r2587065007
##########
mcp-server/polaris_mcp/authorization.py:
##########
@@ -54,59 +55,95 @@ class
ClientCredentialsAuthorizationProvider(AuthorizationProvider):
def __init__(
self,
- token_endpoint: str,
- client_id: str,
- client_secret: str,
- scope: Optional[str],
+ base_url: str,
http: urllib3.PoolManager,
refresh_buffer_seconds: float,
timeout: urllib3.Timeout,
) -> None:
- self._token_endpoint = token_endpoint
- self._client_id = client_id
- self._client_secret = client_secret
- self._scope = scope
+ self._base_url = base_url
self._http = http
+ self._refresh_buffer_seconds = max(refresh_buffer_seconds, 0.0)
self._timeout = timeout
self._lock = threading.Lock()
- self._cached: Optional[tuple[str, float]] = None # (token,
expires_at_epoch)
- self._refresh_buffer_seconds = max(refresh_buffer_seconds, 0.0)
+ # {realm: (token, expires_at_epoch)}
+ self._cached: dict[str, tuple[str, float]] = {}
- def authorization_header(self) -> Optional[str]:
- token = self._current_token()
+ def authorization_header(self, realm: Optional[str] = None) ->
Optional[str]:
+ token = self._get_token_from_realm(realm)
return f"Bearer {token}" if token else None
- def _current_token(self) -> Optional[str]:
- now = time.time()
- cached = self._cached
- if not cached or cached[1] - self._refresh_buffer_seconds <= now:
- with self._lock:
- cached = self._cached
- if (
- not cached
- or cached[1] - self._refresh_buffer_seconds <= time.time()
- ):
- self._cached = cached = self._fetch_token()
- return cached[0] if cached else None
-
- def _fetch_token(self) -> tuple[str, float]:
+ def _get_token_from_realm(self, realm: Optional[str]) -> Optional[str]:
+ def needs_refresh(cached):
+ return (
+ cached is None
+ or cached[1] - self._refresh_buffer_seconds <= time.time()
+ )
+
+ cache_key = realm or ""
+ token = self._cached.get(cache_key)
+ # Token not expired
+ if not needs_refresh(token):
+ return token[0]
+ # Acquire lock and verify again if token expired
+ with self._lock:
+ token = self._cached.get(cache_key)
+ if needs_refresh(token):
+ credentials = self._get_credentials_from_realm(realm)
+ if not credentials:
+ return None
+ token = self._fetch_token(realm, credentials)
+ self._cached[cache_key] = token
+ return token[0] if token else None
+
+ def _get_credentials_from_realm(
+ self, realm: Optional[str]
+ ) -> Optional[dict[str, str]]:
+ def get_env(key: str) -> Optional[str]:
+ val = os.getenv(key)
+ return val.strip() or None if val else None
+
+ def load_creds(realm: Optional[str] = None) -> dict[str,
Optional[str]]:
+ prefix = f"POLARIS_REALM_{realm}_" if realm else "POLARIS_"
+ return {
+ "client_id": get_env(f"{prefix}CLIENT_ID"),
+ "client_secret": get_env(f"{prefix}CLIENT_SECRET"),
+ "scope": get_env(f"{prefix}TOKEN_SCOPE"),
+ "token_url": get_env(f"{prefix}TOKEN_URL"),
+ }
+
+ # Try realm specific first then global
+ for _ in (realm, None):
+ creds = load_creds(_)
Review Comment:
Yeah, maybe let's proceed with this for now (not fall back and quick fail if
missing). So the behavior is matching to the CLI.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
