flyrain commented on code in PR #83:
URL: https://github.com/apache/polaris-tools/pull/83#discussion_r2586367067
##########
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:
> I was thinking maybe users will use same credential for diff realm and not
wanting to set credential for each of them (which is currently a support
deployed model with polaris server).
That's a valid point. Thanks for pointing it out. Do we want the config
without realm to handle both: 1. global setting to all realms, 2. REST calls
without realm header? My minor concern is that it may confuse users. For use
case 1, we mighty introduce a new config like, `POLARIS_GLOBAL_CLIENT_ID`. I
don't have a strong option on this though. Happy to discuss further.
Thanks for the change. LGTM!
--
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]
