MonkeyCanCode commented on code in PR #83:
URL: https://github.com/apache/polaris-tools/pull/83#discussion_r2583573526
##########
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:
Good point. Yeah, 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). But I think the
fail-fast may be better route in this case and this also matches to what we are
asking users to do via current polaris CLI:
```
Loading profiles from /Users/yong/.polaris/.polaris.json
Polaris Client ID: xxxxx
Polaris Client Secret: xxxxx
Polaris Host [localhost]:
Polaris Port [8181]:
Polaris Context Realm:
Polaris Context Header Name [Polaris-Realm]:
Polaris profile new created successfully.
```
Regarding the comment mentioned above, do we want to realm an implicit
config as it can be optional and server can still support.
Let me update this to reflect this change.
--
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]
