From 4bc11897e1b18ad2ad3533a8f49b81d2ae390e17 Mon Sep 17 00:00:00 2001 From: chantelle-cohere Date: Tue, 13 Aug 2024 09:47:33 -0700 Subject: [PATCH 1/9] Initial commit - scaffolded new tool auth deletion route --- src/backend/routers/auth.py | 70 +++++++++++++++++++++++++- src/backend/tools/google_drive/auth.py | 8 +++ 2 files changed, 76 insertions(+), 2 deletions(-) diff --git a/src/backend/routers/auth.py b/src/backend/routers/auth.py index 5e38b6529b..357e53870a 100644 --- a/src/backend/routers/auth.py +++ b/src/backend/routers/auth.py @@ -2,14 +2,14 @@ from typing import Union from urllib.parse import quote -from fastapi import APIRouter, Depends, HTTPException +from fastapi import APIRouter, Depends, HTTPException, status from fastapi.responses import RedirectResponse from starlette.requests import Request from backend.config.auth import ENABLED_AUTH_STRATEGY_MAPPING from backend.config.routers import RouterName from backend.config.settings import Settings -from backend.config.tools import AVAILABLE_TOOLS +from backend.config.tools import AVAILABLE_TOOLS, ToolName from backend.crud import blacklist as blacklist_crud from backend.database_models import Blacklist from backend.database_models.database import DBSessionDep @@ -317,3 +317,69 @@ def log_and_redirect_err(error_message: str): response = RedirectResponse(redirect_uri) return response + + +@router.delete("/tool/auth") +async def delete_tool_auth( + request: Request, session: DBSessionDep, ctx: Context = Depends(get_context) +): + """ + Endpoint to delete Tool Authentication. + + If completed, the corresponding ToolAuth for the requesting user is removed from the DB. + + Args: + request (Request): current Request object. Note that the request body should contain the tool_id parameter as the tool to delete. + session (DBSessionDep): Database session. + ctx (Context): Context object. + + Returns: + {"status": "success"} on successful deletion. + + Raises: + HTTPException: If there was an error deleting the tool auth. + """ + + logger = ctx.get_logger() + logger.error(event=f"CHANTELLE TEST request body {request.body}") + + def log_and_return_error(error_message: str): + logger.error(event=error_message) + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail=f"{error_message}", + ) + + user_id = ctx.get_user_id() + request_body = await request.json() + tool_id = request_body["tool_id"] + + logger.error(event=f"CHANTELLE TEST tool_id {tool_id}, user_id {user_id}") + + if user_id is None or user_id == "" or user_id == "default": + log_and_return_error("User ID not found.") + + if not isinstance(tool_id, str): + log_and_return_error("tool_id must be present in the request body and must be a string.") + + if tool_id != ToolName.Google_Drive: + log_and_return_error(f"Deletion for {tool_id} not implemented.") + + tool = AVAILABLE_TOOLS.get(tool_id) + + if tool.auth_implementation is None: + log_and_return_error(f"Tool {tool.name} does not have an auth_implementation required for Tool Auth Deletion.") + + try: + tool_auth_service = tool.auth_implementation() + tool_auth_service.delete_tool_auth(session, user_id) + is_delete_tool_auth_successful = tool_auth_service.delete_tool_auth(session, user_id) + + if not is_delete_tool_auth_successful: + log_and_return_error("Error deleting Tool Auth.") + + except Exception as e: + log_and_return_error(str(e)) + + + return {"status": "success"} \ No newline at end of file diff --git a/src/backend/tools/google_drive/auth.py b/src/backend/tools/google_drive/auth.py index f3d46de967..956405156e 100644 --- a/src/backend/tools/google_drive/auth.py +++ b/src/backend/tools/google_drive/auth.py @@ -140,3 +140,11 @@ def get_tool_auth(self, session: DBSessionDep, user_id: str) -> ToolAuthModel: def get_token(self, session: DBSessionDep, user_id: str) -> str: tool_auth = tool_auth_crud.get_tool_auth(session, self.TOOL_ID, user_id) return tool_auth.access_token if tool_auth else None + + def delete_tool_auth(self, session: DBSessionDep, user_id: str) -> bool: + try: + tool_auth_crud.delete_tool_auth(session, self.TOOL_ID, user_id) + return true + except Exception as e: + logger.error(event=f"Error while deleting Tool Auth: {str(e)}") + return False From a20e07fdb5f050985017588bf8aba6e9694a16e1 Mon Sep 17 00:00:00 2001 From: chantelle-cohere Date: Tue, 13 Aug 2024 11:26:57 -0700 Subject: [PATCH 2/9] formatting --- src/backend/routers/auth.py | 19 ++++++++++++------- src/backend/tools/google_drive/auth.py | 1 + 2 files changed, 13 insertions(+), 7 deletions(-) diff --git a/src/backend/routers/auth.py b/src/backend/routers/auth.py index 357e53870a..d39fe50618 100644 --- a/src/backend/routers/auth.py +++ b/src/backend/routers/auth.py @@ -358,9 +358,11 @@ def log_and_return_error(error_message: str): if user_id is None or user_id == "" or user_id == "default": log_and_return_error("User ID not found.") - + if not isinstance(tool_id, str): - log_and_return_error("tool_id must be present in the request body and must be a string.") + log_and_return_error( + "tool_id must be present in the request body and must be a string." + ) if tool_id != ToolName.Google_Drive: log_and_return_error(f"Deletion for {tool_id} not implemented.") @@ -368,18 +370,21 @@ def log_and_return_error(error_message: str): tool = AVAILABLE_TOOLS.get(tool_id) if tool.auth_implementation is None: - log_and_return_error(f"Tool {tool.name} does not have an auth_implementation required for Tool Auth Deletion.") + log_and_return_error( + f"Tool {tool.name} does not have an auth_implementation required for Tool Auth Deletion." + ) try: tool_auth_service = tool.auth_implementation() tool_auth_service.delete_tool_auth(session, user_id) - is_delete_tool_auth_successful = tool_auth_service.delete_tool_auth(session, user_id) + is_delete_tool_auth_successful = tool_auth_service.delete_tool_auth( + session, user_id + ) if not is_delete_tool_auth_successful: log_and_return_error("Error deleting Tool Auth.") - + except Exception as e: log_and_return_error(str(e)) - - return {"status": "success"} \ No newline at end of file + return {"status": "success"} diff --git a/src/backend/tools/google_drive/auth.py b/src/backend/tools/google_drive/auth.py index 956405156e..f6e23b53fc 100644 --- a/src/backend/tools/google_drive/auth.py +++ b/src/backend/tools/google_drive/auth.py @@ -142,6 +142,7 @@ def get_token(self, session: DBSessionDep, user_id: str) -> str: return tool_auth.access_token if tool_auth else None def delete_tool_auth(self, session: DBSessionDep, user_id: str) -> bool: + logger.error(event=f"CHANTELLE TEST - GOOGLEDRIVEAUTH class - Deleting Tool Auth: {self.TOOL_ID} for user: {user_id}") try: tool_auth_crud.delete_tool_auth(session, self.TOOL_ID, user_id) return true From ab147da1a01b7ccc9542a266a3f9e3db7ff849cc Mon Sep 17 00:00:00 2001 From: chantelle-cohere Date: Tue, 13 Aug 2024 12:20:13 -0700 Subject: [PATCH 3/9] corrected error in function param order, cleaned up logs --- src/backend/routers/auth.py | 12 ++++++------ src/backend/tools/google_drive/auth.py | 5 ++--- 2 files changed, 8 insertions(+), 9 deletions(-) diff --git a/src/backend/routers/auth.py b/src/backend/routers/auth.py index d39fe50618..5b7f3a2792 100644 --- a/src/backend/routers/auth.py +++ b/src/backend/routers/auth.py @@ -341,7 +341,6 @@ async def delete_tool_auth( """ logger = ctx.get_logger() - logger.error(event=f"CHANTELLE TEST request body {request.body}") def log_and_return_error(error_message: str): logger.error(event=error_message) @@ -351,10 +350,12 @@ def log_and_return_error(error_message: str): ) user_id = ctx.get_user_id() - request_body = await request.json() - tool_id = request_body["tool_id"] - logger.error(event=f"CHANTELLE TEST tool_id {tool_id}, user_id {user_id}") + try: + request_body = await request.json() + tool_id = request_body["tool_id"] + except Exception as e: + log_and_return_error("Error parsing tool_id from request body.") if user_id is None or user_id == "" or user_id == "default": log_and_return_error("User ID not found.") @@ -365,7 +366,7 @@ def log_and_return_error(error_message: str): ) if tool_id != ToolName.Google_Drive: - log_and_return_error(f"Deletion for {tool_id} not implemented.") + log_and_return_error(f"Tool auth deletion for {tool_id} not implemented.") tool = AVAILABLE_TOOLS.get(tool_id) @@ -376,7 +377,6 @@ def log_and_return_error(error_message: str): try: tool_auth_service = tool.auth_implementation() - tool_auth_service.delete_tool_auth(session, user_id) is_delete_tool_auth_successful = tool_auth_service.delete_tool_auth( session, user_id ) diff --git a/src/backend/tools/google_drive/auth.py b/src/backend/tools/google_drive/auth.py index f6e23b53fc..5649393148 100644 --- a/src/backend/tools/google_drive/auth.py +++ b/src/backend/tools/google_drive/auth.py @@ -142,10 +142,9 @@ def get_token(self, session: DBSessionDep, user_id: str) -> str: return tool_auth.access_token if tool_auth else None def delete_tool_auth(self, session: DBSessionDep, user_id: str) -> bool: - logger.error(event=f"CHANTELLE TEST - GOOGLEDRIVEAUTH class - Deleting Tool Auth: {self.TOOL_ID} for user: {user_id}") try: - tool_auth_crud.delete_tool_auth(session, self.TOOL_ID, user_id) - return true + tool_auth_crud.delete_tool_auth(session, user_id, self.TOOL_ID) + return True except Exception as e: logger.error(event=f"Error while deleting Tool Auth: {str(e)}") return False From 6cea7e6f7459fe1195637eb6a32625dab53818c3 Mon Sep 17 00:00:00 2001 From: chantelle-cohere Date: Wed, 14 Aug 2024 10:57:23 -0700 Subject: [PATCH 4/9] PR Feedback: moved tool auth deletion to base auth implementation class, moved tool_id to path parameter --- src/backend/routers/auth.py | 26 +++++++++++++------------- src/backend/tools/base.py | 8 ++++++++ src/backend/tools/google_drive/auth.py | 8 -------- 3 files changed, 21 insertions(+), 21 deletions(-) diff --git a/src/backend/routers/auth.py b/src/backend/routers/auth.py index 5b7f3a2792..848ba86b95 100644 --- a/src/backend/routers/auth.py +++ b/src/backend/routers/auth.py @@ -319,9 +319,12 @@ def log_and_redirect_err(error_message: str): return response -@router.delete("/tool/auth") +@router.delete("/tool/auth/{tool_id}") async def delete_tool_auth( - request: Request, session: DBSessionDep, ctx: Context = Depends(get_context) + tool_id: str, + request: Request, + session: DBSessionDep, + ctx: Context = Depends(get_context), ): """ Endpoint to delete Tool Authentication. @@ -329,7 +332,8 @@ async def delete_tool_auth( If completed, the corresponding ToolAuth for the requesting user is removed from the DB. Args: - request (Request): current Request object. Note that the request body should contain the tool_id parameter as the tool to delete. + tool_id (str): ToolID to be deleted for the user. (eg. google_drive) Should be one of the values listed in the ToolName string enum class. + request (Request): current Request object. session (DBSessionDep): Database session. ctx (Context): Context object. @@ -351,25 +355,21 @@ def log_and_return_error(error_message: str): user_id = ctx.get_user_id() - try: - request_body = await request.json() - tool_id = request_body["tool_id"] - except Exception as e: - log_and_return_error("Error parsing tool_id from request body.") + tool_id = tool_id.lower() if user_id is None or user_id == "" or user_id == "default": log_and_return_error("User ID not found.") - if not isinstance(tool_id, str): + if tool_id not in [tool_name.value for tool_name in ToolName]: log_and_return_error( - "tool_id must be present in the request body and must be a string." + "tool_id must be present in the path of the request and must be a member of the ToolName string enum class." ) - if tool_id != ToolName.Google_Drive: - log_and_return_error(f"Tool auth deletion for {tool_id} not implemented.") - tool = AVAILABLE_TOOLS.get(tool_id) + if tool is None: + log_and_return_error(f"Tool {tool_id} is not available in AVAILABLE_TOOLS.") + if tool.auth_implementation is None: log_and_return_error( f"Tool {tool.name} does not have an auth_implementation required for Tool Auth Deletion." diff --git a/src/backend/tools/base.py b/src/backend/tools/base.py index 3051403b44..e84804c2dd 100644 --- a/src/backend/tools/base.py +++ b/src/backend/tools/base.py @@ -119,6 +119,14 @@ def retrieve_auth_token( def get_token(self, user_id: str, session: DBSessionDep) -> Optional[str]: return None + def delete_tool_auth(self, session: DBSessionDep, user_id: str) -> bool: + try: + tool_auth_crud.delete_tool_auth(session, user_id, self.TOOL_ID) + return True + except Exception as e: + logger.error(event=f"Error while deleting Tool Auth: {str(e)}") + return False + class ToolAuthenticationCacheMixin: def insert_tool_auth_cache(self, user_id: str, tool_id: str) -> str: diff --git a/src/backend/tools/google_drive/auth.py b/src/backend/tools/google_drive/auth.py index 5649393148..f3d46de967 100644 --- a/src/backend/tools/google_drive/auth.py +++ b/src/backend/tools/google_drive/auth.py @@ -140,11 +140,3 @@ def get_tool_auth(self, session: DBSessionDep, user_id: str) -> ToolAuthModel: def get_token(self, session: DBSessionDep, user_id: str) -> str: tool_auth = tool_auth_crud.get_tool_auth(session, self.TOOL_ID, user_id) return tool_auth.access_token if tool_auth else None - - def delete_tool_auth(self, session: DBSessionDep, user_id: str) -> bool: - try: - tool_auth_crud.delete_tool_auth(session, user_id, self.TOOL_ID) - return True - except Exception as e: - logger.error(event=f"Error while deleting Tool Auth: {str(e)}") - return False From 2e6b6045d87452e9ce2552bd1b53f028ccb47003 Mon Sep 17 00:00:00 2001 From: chantelle-cohere Date: Wed, 14 Aug 2024 11:08:44 -0700 Subject: [PATCH 5/9] Add new endpoint to postman collection --- docs/postman/Toolkit.postman_collection.json | 20 ++++++++++++++++++++ 1 file changed, 20 insertions(+) diff --git a/docs/postman/Toolkit.postman_collection.json b/docs/postman/Toolkit.postman_collection.json index 0eb15b8d83..1d2bdac2ef 100644 --- a/docs/postman/Toolkit.postman_collection.json +++ b/docs/postman/Toolkit.postman_collection.json @@ -575,6 +575,26 @@ } ] }, + { + "name": "Tool Auth", + "item": [ + { + "name": "Delete Tool Auth", + "request": { + "auth": { + "type": "bearer", + "bearer": { + "token": "{{auth_token}}" + } + }, + "method": "DELETE", + "header": [], + "url": "http://localhost:8000/v1/tool/auth/{{tool_id}}" + }, + "response": [] + } + ] + }, { "name": "Health", "protocolProfileBehavior": { From b86cec2b51b54c8a276ff181d41acf238603d316 Mon Sep 17 00:00:00 2001 From: chantelle-cohere Date: Wed, 14 Aug 2024 12:27:10 -0700 Subject: [PATCH 6/9] PR feedback: log_and_raise_exception to logger util, return empty DeleteToolAuth response --- src/backend/routers/auth.py | 46 +++++++++++++--------------- src/backend/schemas/tool_auth.py | 4 +++ src/backend/services/logger/utils.py | 10 ++++++ 3 files changed, 35 insertions(+), 25 deletions(-) diff --git a/src/backend/routers/auth.py b/src/backend/routers/auth.py index 848ba86b95..2050cebee2 100644 --- a/src/backend/routers/auth.py +++ b/src/backend/routers/auth.py @@ -2,7 +2,7 @@ from typing import Union from urllib.parse import quote -from fastapi import APIRouter, Depends, HTTPException, status +from fastapi import APIRouter, Depends, HTTPException from fastapi.responses import RedirectResponse from starlette.requests import Request @@ -15,6 +15,7 @@ from backend.database_models.database import DBSessionDep from backend.schemas.auth import JWTResponse, ListAuthStrategy, Login, Logout from backend.schemas.context import Context +from backend.schemas.tool_auth import DeleteToolAuth from backend.services.auth.jwt import JWTService from backend.services.auth.request_validators import validate_authorization from backend.services.auth.utils import ( @@ -23,6 +24,7 @@ ) from backend.services.cache import cache_get_dict from backend.services.context import get_context +from backend.services.logger.utils import log_and_raise_http_exception router = APIRouter(prefix="/v1") router.name = RouterName.AUTH @@ -325,20 +327,20 @@ async def delete_tool_auth( request: Request, session: DBSessionDep, ctx: Context = Depends(get_context), -): +) -> DeleteToolAuth: """ Endpoint to delete Tool Authentication. If completed, the corresponding ToolAuth for the requesting user is removed from the DB. Args: - tool_id (str): ToolID to be deleted for the user. (eg. google_drive) Should be one of the values listed in the ToolName string enum class. + tool_id (str): Tool ID to be deleted for the user. (eg. google_drive) Should be one of the values listed in the ToolName string enum class. request (Request): current Request object. session (DBSessionDep): Database session. ctx (Context): Context object. Returns: - {"status": "success"} on successful deletion. + DeleteToolAuth: Empty response. Raises: HTTPException: If there was an error deleting the tool auth. @@ -346,45 +348,39 @@ async def delete_tool_auth( logger = ctx.get_logger() - def log_and_return_error(error_message: str): - logger.error(event=error_message) - raise HTTPException( - status_code=status.HTTP_400_BAD_REQUEST, - detail=f"{error_message}", - ) - user_id = ctx.get_user_id() - tool_id = tool_id.lower() if user_id is None or user_id == "" or user_id == "default": - log_and_return_error("User ID not found.") + log_and_raise_http_exception("User ID not found.") if tool_id not in [tool_name.value for tool_name in ToolName]: - log_and_return_error( - "tool_id must be present in the path of the request and must be a member of the ToolName string enum class." + log_and_raise_http_exception( + logger, + "tool_id must be present in the path of the request and must be a member of the ToolName string enum class.", ) tool = AVAILABLE_TOOLS.get(tool_id) if tool is None: - log_and_return_error(f"Tool {tool_id} is not available in AVAILABLE_TOOLS.") + log_and_raise_http_exception( + logger, f"Tool {tool_id} is not available in AVAILABLE_TOOLS." + ) if tool.auth_implementation is None: - log_and_return_error( - f"Tool {tool.name} does not have an auth_implementation required for Tool Auth Deletion." + log_and_raise_http_exception( + logger, + f"Tool {tool.name} does not have an auth_implementation required for Tool Auth Deletion.", ) try: tool_auth_service = tool.auth_implementation() - is_delete_tool_auth_successful = tool_auth_service.delete_tool_auth( - session, user_id - ) + is_deleted = tool_auth_service.delete_tool_auth(session, user_id) - if not is_delete_tool_auth_successful: - log_and_return_error("Error deleting Tool Auth.") + if not is_deleted: + log_and_raise_http_exception(logger, "Error deleting Tool Auth.") except Exception as e: - log_and_return_error(str(e)) + log_and_raise_http_exception(logger, str(e)) - return {"status": "success"} + return DeleteToolAuth() diff --git a/src/backend/schemas/tool_auth.py b/src/backend/schemas/tool_auth.py index 19dfd31a85..3dc7026bfa 100644 --- a/src/backend/schemas/tool_auth.py +++ b/src/backend/schemas/tool_auth.py @@ -15,3 +15,7 @@ class UpdateToolAuth(BaseModel): class Config: from_attributes = True use_enum_values = True + + +class DeleteToolAuth(BaseModel): + pass diff --git a/src/backend/services/logger/utils.py b/src/backend/services/logger/utils.py index 49580caa4f..3668718779 100644 --- a/src/backend/services/logger/utils.py +++ b/src/backend/services/logger/utils.py @@ -1,3 +1,5 @@ +from fastapi import HTTPException, status + from backend.config.settings import Settings from backend.services.logger.strategies.base import BaseLogger from backend.services.logger.strategies.structured_log import StructuredLogging @@ -20,3 +22,11 @@ def get_logger(self) -> BaseLogger: else: # Default to StructuredLogging return StructuredLogging(level, renderer) + + +def log_and_raise_http_exception(logger: BaseLogger, error_message: str): + logger.error(event=error_message) + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail=f"{error_message}", + ) From b655af6f416804941faf29da073f28ecc050b783 Mon Sep 17 00:00:00 2001 From: chantelle-cohere Date: Thu, 15 Aug 2024 08:43:02 -0700 Subject: [PATCH 7/9] PR Feedback: error_and_raise_http_exception to logger class --- src/backend/routers/auth.py | 21 ++++++++----------- .../services/logger/strategies/base.py | 3 +++ .../logger/strategies/structured_log.py | 9 ++++++++ src/backend/services/logger/utils.py | 10 --------- src/backend/tools/base.py | 6 ++++-- 5 files changed, 25 insertions(+), 24 deletions(-) diff --git a/src/backend/routers/auth.py b/src/backend/routers/auth.py index 2050cebee2..209ed064e6 100644 --- a/src/backend/routers/auth.py +++ b/src/backend/routers/auth.py @@ -24,7 +24,6 @@ ) from backend.services.cache import cache_get_dict from backend.services.context import get_context -from backend.services.logger.utils import log_and_raise_http_exception router = APIRouter(prefix="/v1") router.name = RouterName.AUTH @@ -352,25 +351,23 @@ async def delete_tool_auth( tool_id = tool_id.lower() if user_id is None or user_id == "" or user_id == "default": - log_and_raise_http_exception("User ID not found.") + logger.error_and_raise_http_exception(event="User ID not found.") if tool_id not in [tool_name.value for tool_name in ToolName]: - log_and_raise_http_exception( - logger, - "tool_id must be present in the path of the request and must be a member of the ToolName string enum class.", + logger.error_and_raise_http_exception( + event="tool_id must be present in the path of the request and must be a member of the ToolName string enum class.", ) tool = AVAILABLE_TOOLS.get(tool_id) if tool is None: - log_and_raise_http_exception( - logger, f"Tool {tool_id} is not available in AVAILABLE_TOOLS." + logger.error_and_raise_http_exception( + event=f"Tool {tool_id} is not available in AVAILABLE_TOOLS." ) if tool.auth_implementation is None: - log_and_raise_http_exception( - logger, - f"Tool {tool.name} does not have an auth_implementation required for Tool Auth Deletion.", + logger.error_and_raise_http_exception( + event=f"Tool {tool.name} does not have an auth_implementation required for Tool Auth Deletion.", ) try: @@ -378,9 +375,9 @@ async def delete_tool_auth( is_deleted = tool_auth_service.delete_tool_auth(session, user_id) if not is_deleted: - log_and_raise_http_exception(logger, "Error deleting Tool Auth.") + logger.error_and_raise_http_exception(event="Error deleting Tool Auth.") except Exception as e: - log_and_raise_http_exception(logger, str(e)) + logger.error_and_raise_http_exception(event=str(e)) return DeleteToolAuth() diff --git a/src/backend/services/logger/strategies/base.py b/src/backend/services/logger/strategies/base.py index d02c83141d..a799fb4ec5 100644 --- a/src/backend/services/logger/strategies/base.py +++ b/src/backend/services/logger/strategies/base.py @@ -14,6 +14,9 @@ def info(self, **kwargs: Any) -> Any: ... @abstractmethod def error(self, **kwargs: Any) -> Any: ... + @abstractmethod + def error_and_raise_http_exception(self, **kwargs: Any) -> Any: ... + @abstractmethod def warning(self, **kwargs: Any) -> Any: ... diff --git a/src/backend/services/logger/strategies/structured_log.py b/src/backend/services/logger/strategies/structured_log.py index c9d9f8a2f9..1e9f962e6e 100644 --- a/src/backend/services/logger/strategies/structured_log.py +++ b/src/backend/services/logger/strategies/structured_log.py @@ -2,6 +2,7 @@ from typing import Any, Dict import structlog +from fastapi import HTTPException, status from backend.services.logger.strategies.base import BaseLogger @@ -84,6 +85,14 @@ def info(self, **kwargs): def error(self, **kwargs): self.logger.error(**kwargs) + @log_context + def error_and_raise_http_exception(self, **kwargs): + self.logger.error(**kwargs) + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail=f"{kwargs}", + ) + @log_context def warning(self, **kwargs): self.logger.warning(**kwargs) diff --git a/src/backend/services/logger/utils.py b/src/backend/services/logger/utils.py index 3668718779..49580caa4f 100644 --- a/src/backend/services/logger/utils.py +++ b/src/backend/services/logger/utils.py @@ -1,5 +1,3 @@ -from fastapi import HTTPException, status - from backend.config.settings import Settings from backend.services.logger.strategies.base import BaseLogger from backend.services.logger.strategies.structured_log import StructuredLogging @@ -22,11 +20,3 @@ def get_logger(self) -> BaseLogger: else: # Default to StructuredLogging return StructuredLogging(level, renderer) - - -def log_and_raise_http_exception(logger: BaseLogger, error_message: str): - logger.error(event=error_message) - raise HTTPException( - status_code=status.HTTP_400_BAD_REQUEST, - detail=f"{error_message}", - ) diff --git a/src/backend/tools/base.py b/src/backend/tools/base.py index e84804c2dd..2f274b9804 100644 --- a/src/backend/tools/base.py +++ b/src/backend/tools/base.py @@ -124,8 +124,10 @@ def delete_tool_auth(self, session: DBSessionDep, user_id: str) -> bool: tool_auth_crud.delete_tool_auth(session, user_id, self.TOOL_ID) return True except Exception as e: - logger.error(event=f"Error while deleting Tool Auth: {str(e)}") - return False + logger.error( + event=f"BaseToolAuthentication: Error while deleting Tool Auth: {str(e)}" + ) + raise Exception(e) class ToolAuthenticationCacheMixin: From ab7cc8faa1b2dba17c1b18a581413b4b7f5745a8 Mon Sep 17 00:00:00 2001 From: chantelle-cohere Date: Thu, 15 Aug 2024 09:17:28 -0700 Subject: [PATCH 8/9] raise bare exception to preserve stacktrace --- src/backend/tools/base.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/backend/tools/base.py b/src/backend/tools/base.py index 2f274b9804..d9d82a7772 100644 --- a/src/backend/tools/base.py +++ b/src/backend/tools/base.py @@ -127,7 +127,7 @@ def delete_tool_auth(self, session: DBSessionDep, user_id: str) -> bool: logger.error( event=f"BaseToolAuthentication: Error while deleting Tool Auth: {str(e)}" ) - raise Exception(e) + raise class ToolAuthenticationCacheMixin: From 2f9d0842cbdd2234b75f25f6cb53f91fe116e531 Mon Sep 17 00:00:00 2001 From: chantelle-cohere Date: Thu, 15 Aug 2024 15:09:58 -0700 Subject: [PATCH 9/9] added crud tests for tool_auth --- src/backend/tests/unit/crud/test_tool_auth.py | 52 +++++++++++++++++++ src/backend/tests/unit/factories/__init__.py | 2 + src/backend/tests/unit/factories/tool_auth.py | 22 ++++++++ 3 files changed, 76 insertions(+) create mode 100644 src/backend/tests/unit/crud/test_tool_auth.py create mode 100644 src/backend/tests/unit/factories/tool_auth.py diff --git a/src/backend/tests/unit/crud/test_tool_auth.py b/src/backend/tests/unit/crud/test_tool_auth.py new file mode 100644 index 0000000000..ea99164070 --- /dev/null +++ b/src/backend/tests/unit/crud/test_tool_auth.py @@ -0,0 +1,52 @@ +from datetime import datetime + +import pytest + +from backend.config.tools import ToolName +from backend.crud import tool_auth as tool_auth_crud +from backend.database_models.tool_auth import ToolAuth +from backend.tests.unit.factories import get_factory + + +def test_create_tool_auth(session, user): + + tool_auth_data = ToolAuth( + user_id=user.id, + tool_id=ToolName.Google_Drive, + token_type="Bearer", + encrypted_access_token=bytes(b"foobar"), + encrypted_refresh_token=bytes(b"foobar"), + expires_at=datetime.strptime("12/31/2124 00:00:00", "%m/%d/%Y %H:%M:%S"), + created_at=datetime.strptime("01/01/2000 00:00:00", "%m/%d/%Y %H:%M:%S"), + updated_at=datetime.strptime("01/01/2010 00:00:00", "%m/%d/%Y %H:%M:%S"), + ) + + tool_auth = tool_auth_crud.create_tool_auth(session, tool_auth_data) + + assert tool_auth.user_id == tool_auth_data.user_id + assert tool_auth.tool_id == tool_auth_data.tool_id + assert tool_auth.token_type == tool_auth_data.token_type + assert tool_auth.encrypted_access_token == tool_auth_data.encrypted_access_token + assert tool_auth.encrypted_refresh_token == tool_auth_data.encrypted_refresh_token + assert tool_auth.expires_at == tool_auth_data.expires_at + assert tool_auth.id == tool_auth_data.id + assert tool_auth.created_at == tool_auth_data.created_at + assert tool_auth.updated_at == tool_auth_data.updated_at + + +def test_delete_tool_auth_by_tool_id(session, user): + tool_auth = get_factory("ToolAuth", session).create( + user_id=user.id, + tool_id=ToolName.Google_Drive, + token_type="Bearer", + encrypted_access_token=bytes(b"foobar"), + encrypted_refresh_token=bytes(b"foobar"), + expires_at=datetime.strptime("12/31/2124 00:00:00", "%m/%d/%Y %H:%M:%S"), + created_at=datetime.strptime("01/01/2000 00:00:00", "%m/%d/%Y %H:%M:%S"), + updated_at=datetime.strptime("01/01/2010 00:00:00", "%m/%d/%Y %H:%M:%S"), + ) + + tool_auth_crud.delete_tool_auth(session, user.id, tool_auth.tool_id) + + tool_auth = tool_auth_crud.get_tool_auth(session, tool_auth.tool_id, user.id) + assert tool_auth is None diff --git a/src/backend/tests/unit/factories/__init__.py b/src/backend/tests/unit/factories/__init__.py index f0916222c6..e5b3cda343 100644 --- a/src/backend/tests/unit/factories/__init__.py +++ b/src/backend/tests/unit/factories/__init__.py @@ -25,6 +25,7 @@ SnapshotFactory, SnapshotLinkFactory, ) +from backend.tests.unit.factories.tool_auth import ToolAuthFactory from backend.tests.unit.factories.tool_call import ToolCallFactory from backend.tests.unit.factories.user import UserFactory @@ -40,6 +41,7 @@ "Agent": AgentFactory, "Organization": OrganizationFactory, "ToolCall": ToolCallFactory, + "ToolAuth": ToolAuthFactory, "Snapshot": SnapshotFactory, "SnapshotLink": SnapshotLinkFactory, "SnapshotAccess": SnapshotAccessFactory, diff --git a/src/backend/tests/unit/factories/tool_auth.py b/src/backend/tests/unit/factories/tool_auth.py new file mode 100644 index 0000000000..8393eae259 --- /dev/null +++ b/src/backend/tests/unit/factories/tool_auth.py @@ -0,0 +1,22 @@ +from datetime import datetime + +import factory + +from backend.config.tools import ToolName +from backend.database_models.tool_auth import ToolAuth + +from .base import BaseFactory + + +class ToolAuthFactory(BaseFactory): + class Meta: + model = ToolAuth + + user_id = factory.Faker("uuid4") + tool_id = ToolName.Google_Drive + token_type = "Bearer" + encrypted_access_token = bytes(b"foobar") + encrypted_refresh_token = bytes(b"foobar") + expires_at = datetime.strptime("12/31/2124 00:00:00", "%m/%d/%Y %H:%M:%S") + created_at = datetime.strptime("01/01/2000 00:00:00", "%m/%d/%Y %H:%M:%S") + updated_at = datetime.strptime("01/01/2010 00:00:00", "%m/%d/%Y %H:%M:%S")