22
33import logging
44import os
5+ from contextlib import asynccontextmanager
56from uuid import UUID
67
78from fastapi import Depends , FastAPI , Request , status
89from fastapi .responses import JSONResponse
9- from contextlib import asynccontextmanager
10+ from prometheus_fastapi_instrumentator import Instrumentator
1011
12+ import ccx_upgrades_data_eng .metrics as metrics
1113from ccx_upgrades_data_eng .auth import (
14+ SessionManagerError ,
15+ TokenError ,
1216 get_session_manager ,
13- SessionManagerException ,
14- TokenException ,
1517)
16- from ccx_upgrades_data_eng .config import get_settings , Settings
18+ from ccx_upgrades_data_eng .config import Settings , get_settings
1719from ccx_upgrades_data_eng .inference import get_filled_inference_for_predictors
1820from ccx_upgrades_data_eng .models import (
19- ClustersList ,
2021 ClusterPrediction ,
22+ ClustersList ,
2123 MultiClusterUpgradeApiResponse ,
2224 UpgradeApiResponse ,
2325)
2628 perform_rhobs_request_multi_cluster ,
2729)
2830from ccx_upgrades_data_eng .sentry import init_sentry
29- import ccx_upgrades_data_eng .metrics as metrics
30-
31- from prometheus_fastapi_instrumentator import Instrumentator
3231from ccx_upgrades_data_eng .utils import get_retry_decorator
3332
3433logger = logging .getLogger (__name__ )
3534
36- init_sentry (os .environ .get ("SENTRY_DSN" , None ), None , os .environ .get ("SENTRY_ENVIRONMENT" , None ))
35+ init_sentry (
36+ os .environ .get ("SENTRY_DSN" , None ), None , os .environ .get ("SENTRY_ENVIRONMENT" , None )
37+ )
3738
3839
3940def create_lifespan_handler (instrumentator : Instrumentator ):
@@ -77,13 +78,13 @@ async def refresh_sso_token(request: Request, call_next) -> JSONResponse:
7778 """Middleware to ensure SSO token is refreshed before processing the request."""
7879 try :
7980 await get_session_and_refresh_token ()
80- except SessionManagerException as ex :
81+ except SessionManagerError as ex :
8182 logger .error ("Unable to initialize SSO session: %s" , ex )
8283 return JSONResponse (
8384 "Unable to initialize SSO session" ,
8485 status_code = status .HTTP_503_SERVICE_UNAVAILABLE ,
8586 )
86- except TokenException as ex :
87+ except TokenError as ex :
8788 logger .error ("Unable to update SSO token: %s" , ex )
8889 return JSONResponse (
8990 "Unable to update SSO token" ,
@@ -92,15 +93,22 @@ async def refresh_sso_token(request: Request, call_next) -> JSONResponse:
9293 return await call_next (request )
9394
9495
95- @app .get ("/cluster/{cluster_id}/upgrade-risks-prediction" , response_model = UpgradeApiResponse )
96- async def upgrade_risks_prediction (cluster_id : UUID , settings : Settings = Depends (get_settings )):
96+ @app .get (
97+ "/cluster/{cluster_id}/upgrade-risks-prediction" , response_model = UpgradeApiResponse
98+ )
99+ async def upgrade_risks_prediction (
100+ cluster_id : UUID ,
101+ settings : Settings = Depends (get_settings ), # noqa: B008
102+ ):
97103 """Return the predition of an upgrade failure given a set of alerts and focs."""
98104 logger .info (f"Received cluster: { cluster_id } " )
99105 logger .debug ("Getting predictors from RHOBS" )
100106 predictors , console_url = perform_rhobs_request (cluster_id )
101107
102108 if console_url is None or console_url == "" :
103- return JSONResponse ("No data for this cluster" , status_code = status .HTTP_404_NOT_FOUND )
109+ return JSONResponse (
110+ "No data for this cluster" , status_code = status .HTTP_404_NOT_FOUND
111+ )
104112
105113 logger .debug ("Getting inference result" )
106114 inference_result = get_filled_inference_for_predictors (predictors , console_url )
@@ -113,16 +121,19 @@ async def upgrade_risks_prediction(cluster_id: UUID, settings: Settings = Depend
113121
114122@app .post ("/upgrade-risks-prediction" , response_model = MultiClusterUpgradeApiResponse )
115123async def upgrade_risks_multi_cluster_predictions (
116- clusters_list : ClustersList , settings : Settings = Depends (get_settings )
124+ clusters_list : ClustersList ,
125+ settings : Settings = Depends (get_settings ), # noqa: B008
117126):
118127 """Return the upgrade risks predictions for the provided clusters."""
119128 logger .info ("Received clusters list: %s" , clusters_list )
120129 logger .debug ("Getting predictors from RHOBS or cache" )
121130 predictors_per_cluster = perform_rhobs_request_multi_cluster (clusters_list .clusters )
122131
123- results = list ()
132+ results = []
124133 for cluster , prediction in predictors_per_cluster .items ():
125- inference_result = get_filled_inference_for_predictors (prediction [0 ], prediction [1 ])
134+ inference_result = get_filled_inference_for_predictors (
135+ prediction [0 ], prediction [1 ]
136+ )
126137 results .append (
127138 ClusterPrediction (
128139 cluster_id = str (cluster ),
0 commit comments