🐛 fix(store.py): handle HTTPStatusError when querying components and raise appropriate exceptions

✨ feat(store.py): add support for user_data_context to fetch and set user data to the context variable
🐛 fix(store.py): handle HTTPStatusError when fetching user data and raise ValueError for invalid API key
🐛 fix(service.py): handle HTTPStatusError when fetching user data and raise ValueError for invalid API key
🐛 fix(service.py): return count as int in count_components method
This commit is contained in:
Gabriel Luiz Freitas Almeida 2023-11-13 11:31:07 -03:00
commit 0f894da75e
2 changed files with 48 additions and 33 deletions

View file

@ -1,11 +1,13 @@
from datetime import datetime
from typing import List, Optional from typing import List, Optional
from uuid import UUID 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.auth import utils as auth_utils
from langflow.services.database.models.user.user import User from langflow.services.database.models.user.user import User
from langflow.services.deps import ( from langflow.services.deps import get_settings_service, get_store_service
get_store_service,
get_settings_service,
)
from langflow.services.store.schema import ( from langflow.services.store.schema import (
ComponentResponse, ComponentResponse,
DownloadComponentResponse, DownloadComponentResponse,
@ -14,14 +16,9 @@ from langflow.services.store.schema import (
TagResponse, TagResponse,
UsersLikesResponse, 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.service import StoreService, user_data_context
from langflow.services.store.utils import update_components_with_user_data from langflow.services.store.utils import update_components_with_user_data
router = APIRouter(prefix="/store", tags=["Components Store"]) 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), store_api_Key: Optional[str] = Depends(get_optional_user_store_api_key),
): ):
try: try:
with user_data_context(store_api_Key, store_service): with user_data_context(api_key=store_api_Key, store_service=store_service):
authorized = False try:
result = store_service.query_components( authorized = False
api_key=store_api_Key, result = store_service.query_components(
page=page, api_key=store_api_Key,
limit=limit, page=page,
filter_by_user=filter_by_user, limit=limit,
is_component=is_component, filter_by_user=filter_by_user,
search=search, is_component=is_component,
status=status, search=search,
tags=tags, status=status,
sort=sort, 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: try:
comp_count = store_service.count_components( comp_count = store_service.count_components(
api_key=store_api_Key, api_key=store_api_Key,
@ -130,6 +133,12 @@ def get_components(
results=result, authorized=authorized, count=comp_count results=result, authorized=authorized, count=comp_count
) )
except Exception as exc: 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)) raise HTTPException(status_code=500, detail=str(exc))

View file

@ -1,11 +1,12 @@
from datetime import datetime
import json import json
from datetime import datetime
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
from uuid import UUID 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 ( from langflow.services.store.schema import (
ComponentResponse, ComponentResponse,
DownloadComponentResponse, DownloadComponentResponse,
@ -16,6 +17,7 @@ from langflow.services.store.utils import process_tags_for_post
if TYPE_CHECKING: if TYPE_CHECKING:
from langflow.services.settings.service import SettingsService from langflow.services.settings.service import SettingsService
from contextlib import contextmanager from contextlib import contextmanager
from contextvars import ContextVar from contextvars import ContextVar
@ -25,12 +27,16 @@ user_data_var: ContextVar[Optional[Dict[str, Any]]] = ContextVar(
@contextmanager @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 # Fetch and set user data to the context variable
if api_key: if api_key:
user_data = store_service._get( try:
f"{store_service.base_url}/users/me", api_key, params={"fields": "id"} 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) user_data_var.set(user_data)
try: try:
yield yield
@ -199,7 +205,7 @@ class StoreService(Service):
params["filter"] = json.dumps({"_and": filter_conditions}) params["filter"] = json.dumps({"_and": filter_conditions})
results = self._get(self.components_url, api_key, params) results = self._get(self.components_url, api_key, params)
return results[0].get("count", 0) return int(results[0].get("count", 0))
@staticmethod @staticmethod
def build_search_filter_conditions(query: str): def build_search_filter_conditions(query: str):
@ -213,7 +219,7 @@ class StoreService(Service):
def query_components( def query_components(
self, self,
api_key: str, api_key: Optional[str] = None,
search: Optional[str] = None, search: Optional[str] = None,
status: Optional[str] = None, status: Optional[str] = None,
tags: Optional[List[str]] = None, tags: Optional[List[str]] = None,