#使用带认证和 RBAC 的 MLflow Python SDK

import base64, hashlib, json, os, secrets
from urllib.parse import urlparse, parse_qs
import requests
from cryptography.hazmat.primitives.asymmetric import padding
from cryptography.hazmat.primitives.serialization import load_pem_public_key

PLATFORM      = os.environ["PLATFORM_ADDRESS"].rstrip("/")    # https://<platform>
CLIENT_ID     = os.environ["DEX_CLIENT_ID"]                   # the MLflow proxy's client, e.g. alauda-auth
CLIENT_SECRET = os.environ["DEX_CLIENT_SECRET"]
USERNAME      = os.environ["MLFLOW_USERNAME"]
PASSWORD      = os.environ["MLFLOW_PASSWORD"]
REDIRECT_URI  = f"{PLATFORM}/oauth2/callback"                 # any URI the client has registered
VERIFY_TLS    = os.environ.get("PLATFORM_CA", False)         # CA bundle path, or False to skip (lab only)

s = requests.Session(); s.verify = VERIFY_TLS
_b64url = lambda b: base64.urlsafe_b64encode(b).rstrip(b"=").decode()

def get_tokens() -> dict:
    """Run the authorization-code + PKCE flow headlessly. Returns the Dex token response."""
    verifier  = _b64url(secrets.token_bytes(48))
    challenge = _b64url(hashlib.sha256(verifier.encode()).digest())
    # 1) start the flow -> auth-request id
    req = s.get(f"{PLATFORM}/dex/api/v1/authorize", params={
        "client_id": CLIENT_ID, "redirect_uri": REDIRECT_URI, "response_type": "code",
        "scope": "openid email groups offline_access", "state": "cli",
        "code_challenge": challenge, "code_challenge_method": "S256"}).json()["req"]
    # 2) RSA-encrypt the password, then log in via the local connector -> auth code
    pk  = s.get(f"{PLATFORM}/dex/pubkey").json()              # {"ts": ..., "pubkey": "<PEM>"}
    payload = json.dumps({"ts": pk["ts"], "password": PASSWORD}, separators=(",", ":")).encode()
    enc = base64.b64encode(load_pem_public_key(pk["pubkey"].encode()).encrypt(payload, padding.PKCS1v15())).decode()
    redirect = s.post(f"{PLATFORM}/dex/api/v1/authorize/local", params={"req": req},
        json={"account": USERNAME, "password": enc}).json()["redirect_url"]
    code = parse_qs(urlparse(redirect).query)["code"][0]
    # 3) exchange the code (with the PKCE verifier) -> id_token + refresh_token
    return s.post(f"{PLATFORM}/dex/token", data={
        "grant_type": "authorization_code", "code": code, "redirect_uri": REDIRECT_URI,
        "code_verifier": verifier, "client_id": CLIENT_ID, "client_secret": CLIENT_SECRET}).json()

def refresh(refresh_token: str) -> str:
    """Mint a fresh id token from a refresh token — no login, no browser."""
    return s.post(f"{PLATFORM}/dex/token", data={
        "grant_type": "refresh_token", "refresh_token": refresh_token,
        "client_id": CLIENT_ID, "client_secret": CLIENT_SECRET,
        "scope": "openid email groups"}).json()["id_token"]

import os, mlflow

tok = get_tokens()
os.environ["MLFLOW_TRACKING_TOKEN"] = tok["id_token"].strip()           # → Authorization: Bearer
mlflow.set_tracking_uri("http://mlflow-tracking-server.kubeflow:5000")  # in-cluster Service (fronted by the OAuth proxy)
mlflow.set_workspace("team-a")                                          # workspace namespace → X-MLFLOW-WORKSPACE
mlflow.set_experiment("my-experiment")

with mlflow.start_run(run_name="sdk-quickstart") as run:
    mlflow.log_param("learning_rate", 2e-4)
    mlflow.log_metric("loss", 0.123)
    print("run:", run.info.run_id)

mlflow.set_workspace("team-a")
with mlflow.start_run():
    mlflow.sklearn.log_model(sk_model, name="model", registered_model_name="fraud-detector")

import os, mlflow
from mlflow.tracking.request_header.abstract_request_header_provider import RequestHeaderProvider
from mlflow.tracking.request_header.registry import _request_header_provider_registry

class ProxySessionHeader(RequestHeaderProvider):
    def in_context(self):
        return bool(os.environ.get("MLFLOW_PROXY_COOKIE"))     # export MLFLOW_PROXY_COOKIE='_oauth2_proxy=<value>'
    def request_headers(self):
        return {"Cookie": os.environ["MLFLOW_PROXY_COOKIE"]}

_request_header_provider_registry.register(ProxySessionHeader)
mlflow.set_tracking_uri("https://<platform>/clusters/<cluster>/mlflow")
mlflow.set_workspace("team-a")