Skip to content

Commit ae0c6d1

Browse files
committed
Completed API request on behalf of users.
Signed-off-by: Eric Evans <[email protected]>
1 parent 2120976 commit ae0c6d1

17 files changed

+651
-202
lines changed

src/aiq/authentication/authentication_manager.py

Lines changed: 6 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -32,7 +32,7 @@
3232
from aiq.authentication.response_manager import ResponseManager
3333

3434

35-
class AuthenticationManager: # TODO EE: Add Tests
35+
class AuthenticationManager:
3636

3737
def __init__(self, request_manager: "RequestManager", response_manager: "ResponseManager") -> None:
3838
self._oauth2_authenticator: OAuth2Authenticator = OAuth2Authenticator(request_manager=request_manager,
@@ -57,11 +57,6 @@ async def _validate_auth_provider_credentials(self, authentication_provider_name
5757
if provider is None:
5858
raise ValueError(f"Authentication provider {authentication_provider_name} not found.")
5959

60-
if not isinstance(provider, (OAuth2Config | APIKeyConfig)):
61-
raise TypeError(
62-
f"Authentication type for {authentication_provider_name} not supported. Supported types can "
63-
f"be found in \"aiq.data_models.authentication\" ")
64-
6560
if isinstance(provider, OAuth2Config):
6661
# Set the provider of interest.
6762
self._oauth2_authenticator.authentication_provider = provider
@@ -107,7 +102,7 @@ async def _validate_api_key_credentials(self, api_key_name: str, api_key_provide
107102

108103
async def _set_auth_provider_credentials(self, authentication_provider: str) -> bool:
109104
"""
110-
Gets the authentication provider credentials for the registered provider.
105+
Gets and persists the authentication provider credentials for the registered provider.
111106
112107
Args:
113108
authentication_provider (str): The name of the registered authentication provider.
@@ -133,7 +128,7 @@ async def _set_auth_provider_credentials(self, authentication_provider: str) ->
133128
return credentials_set
134129

135130
except (Exception, ValueError, TypeError) as e:
136-
logger.error("Failed to validate authentication provider credentials: %s Error: %s",
131+
logger.error("Failed get and persist authentication provider credentials: %s Error: %s",
137132
authentication_provider,
138133
str(e),
139134
exc_info=True)
@@ -168,10 +163,11 @@ async def _construct_authentication_header(self, authentication_provider: str |
168163
return auth_header
169164

170165
else:
171-
raise ValidationError(f"Authentication provider: {authentication_provider} is not authenticated.")
166+
raise ValidationError(
167+
f"Authentication provider: {authentication_provider} credentials can not be validated.")
172168

173169
except (ValidationError) as e:
174-
logger.error("Authentication provider: %s is not authenticated. %s",
170+
logger.error("Authentication provider: %s credentials can not be validated: %s",
175171
authentication_provider,
176172
str(e),
177173
exc_info=True)

src/aiq/authentication/exceptions.py

Lines changed: 11 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -29,16 +29,26 @@ class BaseUrlValidationError(Exception):
2929
pass
3030

3131

32+
class HTTPMethodValidationError(Exception):
33+
"""Raised when HTTP Method validation fails unexpectedly."""
34+
pass
35+
36+
3237
class QueryParameterValidationError(Exception):
3338
"""Raised when HTTP Query Parameter validation fails unexpectedly."""
3439
pass
3540

3641

37-
class HeaderValidationError(Exception):
42+
class HTTPHeaderValidationError(Exception):
3843
"""Raised when HTTP Header validation fails unexpectedly."""
3944
pass
4045

4146

4247
class BodyValidationError(Exception):
4348
"""Raised when HTTP Body validation fails unexpectedly."""
4449
pass
50+
51+
52+
class APIRequestError(Exception):
53+
"""Raised when making an API request fails unexpectedly."""
54+
pass

src/aiq/authentication/jira.py

Lines changed: 0 additions & 29 deletions
This file was deleted.

src/aiq/authentication/oauth2_authenticator.py

Lines changed: 22 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -115,8 +115,16 @@ async def _get_credentials(self) -> bool:
115115

116116
return await self._validate_credentials()
117117

118-
except ValueError as e:
119-
logger.error("Failed to get OAuth2.0 credentials: %s", str(e), exc_info=True)
118+
except OAuthCodeFlowError as e:
119+
logger.error("Failed to get OAuth2.0 credentials using code flow: %s", str(e), exc_info=True)
120+
return False
121+
122+
except OAuthRefreshTokenError as e:
123+
logger.error("Failed to get OAuth2.0 credentials using refresh token: %s", str(e), exc_info=True)
124+
return False
125+
126+
except Exception as e:
127+
logger.error("Failed to get OAuth2.0 credentials due to unexpected error: %s", str(e), exc_info=True)
120128
return False
121129

122130
async def _initiate_authorization_code_flow_console(self) -> None:
@@ -139,6 +147,7 @@ async def _initiate_authorization_code_flow_console(self) -> None:
139147
except OAuthCodeFlowError as e:
140148
logger.error("Failed to complete OAuth2.0 authentication for provider Error: %s", str(e), exc_info=True)
141149
await self._shutdown_code_flow_dispatch[self.execution_mode]()
150+
raise e
142151

143152
async def _initiate_authorization_code_flow_server(self) -> None:
144153
"""
@@ -155,6 +164,7 @@ async def _initiate_authorization_code_flow_server(self) -> None:
155164
except OAuthCodeFlowError as e:
156165
logger.error("Failed to complete OAuth2.0 authentication for provider Error: %s", str(e), exc_info=True)
157166
await self._shutdown_code_flow_dispatch[self.execution_mode]()
167+
raise e
158168

159169
async def _send_oauth_authorization_request(self) -> None:
160170
"""
@@ -184,13 +194,13 @@ async def _get_access_token_with_refresh_token(self) -> None:
184194
"""
185195
try:
186196
if not self.authentication_provider.refresh_token:
187-
raise ValueError("Refresh token is missing during the refresh token request.")
197+
raise OAuthRefreshTokenError("Refresh token is missing during the refresh token request.")
188198

189199
if not self.authentication_provider.client_id:
190-
raise ValueError("Client ID is missing during the refresh token request.")
200+
raise OAuthRefreshTokenError("Client ID is missing during the refresh token request.")
191201

192202
if not self.authentication_provider.client_secret:
193-
raise ValueError("Client secret is missing during the refresh token request.")
203+
raise OAuthRefreshTokenError("Client secret is missing during the refresh token request.")
194204

195205
body_data: RefreshTokenRequest = RefreshTokenRequest(
196206
client_id=self._authentication_provider.client_id,
@@ -206,20 +216,21 @@ async def _get_access_token_with_refresh_token(self) -> None:
206216
body_data=body_data.model_dump())
207217

208218
if response is None:
209-
raise ValueError("Invalid response received while making refresh token request.")
219+
raise OAuthRefreshTokenError("Invalid response received while making refresh token request.")
210220

211221
if not response.status_code == 200:
212222
await self._response_manager._handle_oauth_authorization_response_codes(
213223
response, self._authentication_provider)
214224

215225
if response.json().get("access_token") is None:
216-
raise ValueError("Access token not in successful token request response payload.")
226+
raise OAuthRefreshTokenError("Access token not in successful token request response payload.")
217227

218228
if response.json().get("expires_in") is None:
219-
raise ValueError("Access token expiration time not in successful token request response payload.")
229+
raise OAuthRefreshTokenError(
230+
"Access token expiration time not in successful token request response payload.")
220231

221232
if response.json().get("refresh_token") is None:
222-
raise ValueError("Refresh token not in successful token request response payload.")
233+
raise OAuthRefreshTokenError("Refresh token not in successful token request response payload.")
223234

224235
self.authentication_provider.access_token = response.json().get("access_token")
225236

@@ -228,8 +239,9 @@ async def _get_access_token_with_refresh_token(self) -> None:
228239

229240
self.authentication_provider.refresh_token = response.json().get("refresh_token")
230241

231-
except OAuthRefreshTokenError as e:
242+
except Exception as e:
232243
logger.error("Failed to complete OAuth2.0 authentication for provider Error: %s", str(e), exc_info=True)
244+
raise OAuthRefreshTokenError("Failed to get OAuth2.0 credentials using refresh token.") from e
233245

234246
async def _spawn_oauth_client_server(self) -> None:
235247
"""

0 commit comments

Comments
 (0)