From 0f894da75e658f8e32050c5fe7dd757f12b80816 Mon Sep 17 00:00:00 2001 From: Gabriel Luiz Freitas Almeida Date: Mon, 13 Nov 2023 11:31:07 -0300 Subject: [PATCH] =?UTF-8?q?=F0=9F=90=9B=20fix(store.py):=20handle=20HTTPSt?= =?UTF-8?q?atusError=20when=20querying=20components=20and=20raise=20approp?= =?UTF-8?q?riate=20exceptions=20=E2=9C=A8=20feat(store.py):=20add=20suppor?= =?UTF-8?q?t=20for=20user=5Fdata=5Fcontext=20to=20fetch=20and=20set=20user?= =?UTF-8?q?=20data=20to=20the=20context=20variable=20=F0=9F=90=9B=20fix(st?= =?UTF-8?q?ore.py):=20handle=20HTTPStatusError=20when=20fetching=20user=20?= =?UTF-8?q?data=20and=20raise=20ValueError=20for=20invalid=20API=20key=20?= =?UTF-8?q?=F0=9F=90=9B=20fix(service.py):=20handle=20HTTPStatusError=20wh?= =?UTF-8?q?en=20fetching=20user=20data=20and=20raise=20ValueError=20for=20?= =?UTF-8?q?invalid=20API=20key=20=F0=9F=90=9B=20fix(service.py):=20return?= =?UTF-8?q?=20count=20as=20int=20in=20count=5Fcomponents=20method?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/backend/langflow/api/v1/store.py | 53 +++++++++++-------- .../langflow/services/store/service.py | 28 ++++++---- 2 files changed, 48 insertions(+), 33 deletions(-) diff --git a/src/backend/langflow/api/v1/store.py b/src/backend/langflow/api/v1/store.py index 707542ee3..3852afdad 100644 --- a/src/backend/langflow/api/v1/store.py +++ b/src/backend/langflow/api/v1/store.py @@ -1,11 +1,13 @@ +from datetime import datetime from typing import List, Optional from uuid import UUID + +from fastapi import APIRouter, Depends, HTTPException, Query +from httpx import HTTPStatusError + from langflow.services.auth import utils as auth_utils from langflow.services.database.models.user.user import User -from langflow.services.deps import ( - get_store_service, - get_settings_service, -) +from langflow.services.deps import get_settings_service, get_store_service from langflow.services.store.schema import ( ComponentResponse, DownloadComponentResponse, @@ -14,14 +16,9 @@ from langflow.services.store.schema import ( TagResponse, UsersLikesResponse, ) - -from fastapi import APIRouter, Depends, HTTPException, Query -from datetime import datetime - from langflow.services.store.service import StoreService, user_data_context from langflow.services.store.utils import update_components_with_user_data - router = APIRouter(prefix="/store", tags=["Components Store"]) @@ -91,19 +88,25 @@ def get_components( store_api_Key: Optional[str] = Depends(get_optional_user_store_api_key), ): try: - with user_data_context(store_api_Key, store_service): - authorized = False - result = store_service.query_components( - api_key=store_api_Key, - page=page, - limit=limit, - filter_by_user=filter_by_user, - is_component=is_component, - search=search, - status=status, - tags=tags, - sort=sort, - ) + with user_data_context(api_key=store_api_Key, store_service=store_service): + try: + authorized = False + result = store_service.query_components( + api_key=store_api_Key, + page=page, + limit=limit, + filter_by_user=filter_by_user, + is_component=is_component, + search=search, + status=status, + tags=tags, + sort=sort, + ) + except HTTPStatusError as exc: + if exc.response.status_code == 403: + raise ValueError( + "You are not authorized to access this public resource" + ) try: comp_count = store_service.count_components( api_key=store_api_Key, @@ -130,6 +133,12 @@ def get_components( results=result, authorized=authorized, count=comp_count ) except Exception as exc: + if isinstance(exc, HTTPStatusError): + if exc.response.status_code == 403: + raise HTTPException(status_code=403, detail="Forbidden") + elif isinstance(exc, ValueError): + raise HTTPException(status_code=403, detail=str(exc)) + raise HTTPException(status_code=500, detail=str(exc)) diff --git a/src/backend/langflow/services/store/service.py b/src/backend/langflow/services/store/service.py index 79e124a2b..6c0e06013 100644 --- a/src/backend/langflow/services/store/service.py +++ b/src/backend/langflow/services/store/service.py @@ -1,11 +1,12 @@ -from datetime import datetime import json +from datetime import datetime +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union from uuid import UUID -from langflow.services.base import Service -from typing import TYPE_CHECKING, List, Dict, Any, Optional, Union -import httpx -from httpx import HTTPError +import httpx +from httpx import HTTPError, HTTPStatusError + +from langflow.services.base import Service from langflow.services.store.schema import ( ComponentResponse, DownloadComponentResponse, @@ -16,6 +17,7 @@ from langflow.services.store.utils import process_tags_for_post if TYPE_CHECKING: from langflow.services.settings.service import SettingsService + from contextlib import contextmanager from contextvars import ContextVar @@ -25,12 +27,16 @@ user_data_var: ContextVar[Optional[Dict[str, Any]]] = ContextVar( @contextmanager -def user_data_context(api_key: str, store_service: "StoreService"): +def user_data_context(store_service: "StoreService", api_key: Optional[str] = None): # Fetch and set user data to the context variable if api_key: - user_data = store_service._get( - f"{store_service.base_url}/users/me", api_key, params={"fields": "id"} - ) + try: + user_data = store_service._get( + f"{store_service.base_url}/users/me", api_key, params={"fields": "id"} + ) + except HTTPStatusError as exc: + if exc.response.status_code == 403: + raise ValueError("Invalid API key") user_data_var.set(user_data) try: yield @@ -199,7 +205,7 @@ class StoreService(Service): params["filter"] = json.dumps({"_and": filter_conditions}) results = self._get(self.components_url, api_key, params) - return results[0].get("count", 0) + return int(results[0].get("count", 0)) @staticmethod def build_search_filter_conditions(query: str): @@ -213,7 +219,7 @@ class StoreService(Service): def query_components( self, - api_key: str, + api_key: Optional[str] = None, search: Optional[str] = None, status: Optional[str] = None, tags: Optional[List[str]] = None,