From 3b3bb4f5e4214bbbc10350411b193f79b311833c Mon Sep 17 00:00:00 2001 From: Eugene Yurtsev Date: Thu, 21 Mar 2024 21:53:31 -0400 Subject: [PATCH] Swap to x-key (#91) Swap to api key in headers --- backend/db/models.py | 5 +- backend/server/api/api_key.py | 6 +++ backend/server/api/examples.py | 12 +++-- backend/server/api/extract.py | 5 +- backend/server/api/extractors.py | 19 ++++--- backend/server/main.py | 6 +-- .../api/test_api_defining_extractors.py | 52 +++++++++---------- .../tests/unit_tests/api/test_api_examples.py | 26 +++++----- .../tests/unit_tests/api/test_api_extract.py | 24 +++++---- frontend/app/components/Sidebar.tsx | 9 ++-- frontend/app/utils/api.tsx | 45 ++++++++++++---- 11 files changed, 126 insertions(+), 83 deletions(-) create mode 100644 backend/server/api/api_key.py diff --git a/backend/db/models.py b/backend/db/models.py index 0156fd9..5832d8d 100644 --- a/backend/db/models.py +++ b/backend/db/models.py @@ -182,7 +182,4 @@ def validate_extractor_owner( extractor = ( session.query(Extractor).filter_by(uuid=extractor_id, owner_id=user_id).first() ) - if extractor is None: - return False - else: - return True + return extractor is not None diff --git a/backend/server/api/api_key.py b/backend/server/api/api_key.py new file mode 100644 index 0000000..037f549 --- /dev/null +++ b/backend/server/api/api_key.py @@ -0,0 +1,6 @@ +from fastapi.security import APIKeyHeader + +# For actual auth, you'd need to check the key against a database or some other +# data store. Here, we don't need actual auth, just a key that matches +# a UUID +UserToken = APIKeyHeader(name="x-key") diff --git a/backend/server/api/examples.py b/backend/server/api/examples.py index b0fbdc8..5b9f4a1 100644 --- a/backend/server/api/examples.py +++ b/backend/server/api/examples.py @@ -2,11 +2,12 @@ from typing import Any, List from uuid import UUID -from fastapi import APIRouter, Cookie, Depends, HTTPException +from fastapi import APIRouter, Depends, HTTPException from sqlalchemy.orm import Session from typing_extensions import Annotated, TypedDict from db.models import Example, get_session, validate_extractor_owner +from server.api.api_key import UserToken router = APIRouter( prefix="/examples", @@ -36,7 +37,7 @@ def create( create_request: CreateExample, *, session: Session = Depends(get_session), - user_id: UUID = Cookie(...), + user_id: UUID = Depends(UserToken), ) -> CreateExampleResponse: """Endpoint to create an example.""" if not validate_extractor_owner(session, create_request["extractor_id"], user_id): @@ -59,7 +60,7 @@ def list( limit: int = 10, offset: int = 0, session=Depends(get_session), - user_id: UUID = Cookie(...), + user_id: UUID = Depends(UserToken), ) -> List[Any]: """Endpoint to get all examples.""" if not validate_extractor_owner(session, extractor_id, user_id): @@ -76,7 +77,10 @@ def list( @router.delete("/{uuid}") def delete( - uuid: UUID, *, session: Session = Depends(get_session), user_id: UUID = Cookie(...) + uuid: UUID, + *, + session: Session = Depends(get_session), + user_id: UUID = Depends(UserToken), ) -> None: """Endpoint to delete an example.""" extractor_id = session.query(Example).filter_by(uuid=str(uuid)).first().extractor_id diff --git a/backend/server/api/extract.py b/backend/server/api/extract.py index ced707c..d46c245 100644 --- a/backend/server/api/extract.py +++ b/backend/server/api/extract.py @@ -1,12 +1,13 @@ from typing import Literal, Optional from uuid import UUID -from fastapi import APIRouter, Cookie, Depends, File, Form, HTTPException, UploadFile +from fastapi import APIRouter, Depends, File, Form, HTTPException, UploadFile from sqlalchemy.orm import Session from typing_extensions import Annotated from db.models import Extractor, SharedExtractors, get_session from extraction.parsing import parse_binary_input +from server.api.api_key import UserToken from server.extraction_runnable import ExtractResponse, extract_entire_document from server.models import DEFAULT_MODEL from server.retrieval import extract_from_content @@ -27,7 +28,7 @@ async def extract_using_existing_extractor( file: Optional[UploadFile] = File(None), model_name: Optional[str] = Form(DEFAULT_MODEL), session: Session = Depends(get_session), - user_id: UUID = Cookie(...), + user_id: UUID = Depends(UserToken), ) -> ExtractResponse: """Endpoint that is used with an existing extractor. diff --git a/backend/server/api/extractors.py b/backend/server/api/extractors.py index 64f6f52..9ffc3e9 100644 --- a/backend/server/api/extractors.py +++ b/backend/server/api/extractors.py @@ -2,12 +2,13 @@ from typing import Any, Dict, List from uuid import UUID, uuid4 -from fastapi import APIRouter, Cookie, Depends, HTTPException +from fastapi import APIRouter, Depends, HTTPException from pydantic import BaseModel, Field, validator from sqlalchemy.exc import IntegrityError from sqlalchemy.orm import Session from db.models import Extractor, SharedExtractors, get_session, validate_extractor_owner +from server.api.api_key import UserToken from server.validators import validate_json_schema router = APIRouter( @@ -60,7 +61,7 @@ def share( uuid: UUID, *, session: Session = Depends(get_session), - user_id: UUID = Cookie(...), + user_id: UUID = Depends(UserToken), ) -> ShareExtractorResponse: """Endpoint to share an extractor. @@ -110,7 +111,7 @@ def create( create_request: CreateExtractor, *, session: Session = Depends(get_session), - user_id: UUID = Cookie(...), + user_id: UUID = Depends(UserToken), ) -> CreateExtractorResponse: """Endpoint to create an extractor.""" @@ -128,7 +129,10 @@ def create( @router.get("/{uuid}") def get( - uuid: UUID, *, session: Session = Depends(get_session), user_id: UUID = Cookie(...) + uuid: UUID, + *, + session: Session = Depends(get_session), + user_id: UUID = Depends(UserToken), ) -> Dict[str, Any]: """Endpoint to get an extractor.""" extractor = ( @@ -151,7 +155,7 @@ def list( limit: int = 10, offset: int = 0, session=Depends(get_session), - user_id: UUID = Cookie(...), + user_id: UUID = Depends(UserToken), ) -> List[Any]: """Endpoint to get all extractors.""" return ( @@ -165,7 +169,10 @@ def list( @router.delete("/{uuid}") def delete( - uuid: UUID, *, session: Session = Depends(get_session), user_id: UUID = Cookie(...) + uuid: UUID, + *, + session: Session = Depends(get_session), + user_id: UUID = Depends(UserToken), ) -> None: """Endpoint to delete an extractor.""" session.query(Extractor).filter_by(uuid=str(uuid), owner_id=user_id).delete() diff --git a/backend/server/main.py b/backend/server/main.py index befdd41..7b2839f 100644 --- a/backend/server/main.py +++ b/backend/server/main.py @@ -68,10 +68,10 @@ def ready() -> str: # Serve the frontend -ui_dir = str(ROOT / "ui") +UI_DIR = str(ROOT / "ui") -if os.path.exists(ui_dir): - app.mount("/", StaticFiles(directory=ui_dir, html=True), name="ui") +if os.path.exists(UI_DIR): + app.mount("/", StaticFiles(directory=UI_DIR, html=True), name="ui") else: logger.warning("No UI directory found, serving API only.") diff --git a/backend/tests/unit_tests/api/test_api_defining_extractors.py b/backend/tests/unit_tests/api/test_api_defining_extractors.py index c9fa661..fe05194 100644 --- a/backend/tests/unit_tests/api/test_api_defining_extractors.py +++ b/backend/tests/unit_tests/api/test_api_defining_extractors.py @@ -9,8 +9,8 @@ async def test_extractors_api() -> None: # First verify that the database is empty async with get_async_client() as client: user_id = str(uuid.uuid4()) - cookies = {"user_id": user_id} - response = await client.get("/extractors", cookies=cookies) + headers = {"x-key": user_id} + response = await client.get("/extractors", headers=headers) assert response.status_code == 200 assert response.json() == [] @@ -21,38 +21,36 @@ async def test_extractors_api() -> None: "instruction": "Test Instruction", } response = await client.post( - "/extractors", json=create_request, cookies=cookies + "/extractors", json=create_request, headers=headers ) assert response.status_code == 200 # Verify that the extractor was created - response = await client.get("/extractors", cookies=cookies) + response = await client.get("/extractors", headers=headers) assert response.status_code == 200 get_response = response.json() assert len(get_response) == 1 - # Check cookies - bad_cookies = {"user_id": str(uuid.uuid4())} - bad_response = await client.get("/extractors", cookies=bad_cookies) + # Check headers + bad_headers = {"x-key": str(uuid.uuid4())} + bad_response = await client.get("/extractors", headers=bad_headers) assert bad_response.status_code == 200 assert len(bad_response.json()) == 0 # Check we need cookie to delete uuid_str = get_response[0]["uuid"] _ = uuid.UUID(uuid_str) # assert valid uuid - bad_response = await client.delete( - f"/extractors/{uuid_str}", cookies=bad_cookies - ) + await client.delete(f"/extractors/{uuid_str}", headers=bad_headers) # Check extractor was not deleted - response = await client.get("/extractors", cookies=cookies) + response = await client.get("/extractors", headers=headers) assert len(response.json()) == 1 # Verify that we can delete an extractor _ = uuid.UUID(uuid_str) # assert valid uuid - response = await client.delete(f"/extractors/{uuid_str}", cookies=cookies) + response = await client.delete(f"/extractors/{uuid_str}", headers=headers) assert response.status_code == 200 - get_response = await client.get("/extractors", cookies=cookies) + get_response = await client.get("/extractors", headers=headers) assert get_response.status_code == 200 assert get_response.json() == [] @@ -63,12 +61,12 @@ async def test_extractors_api() -> None: "instruction": "Test Instruction", } response = await client.post( - "/extractors", json=create_request, cookies=cookies + "/extractors", json=create_request, headers=headers ) assert response.status_code == 200 # Verify that the extractor was created - response = await client.get("/extractors", cookies=cookies) + response = await client.get("/extractors", headers=headers) assert response.status_code == 200 assert len(response.json()) == 1 @@ -76,10 +74,10 @@ async def test_extractors_api() -> None: get_response = response.json() uuid_str = get_response[0]["uuid"] _ = uuid.UUID(uuid_str) # assert valid uuid - response = await client.delete(f"/extractors/{uuid_str}", cookies=cookies) + response = await client.delete(f"/extractors/{uuid_str}", headers=headers) assert response.status_code == 200 - get_response = await client.get("/extractors", cookies=cookies) + get_response = await client.get("/extractors", headers=headers) assert get_response.status_code == 200 assert get_response.json() == [] @@ -92,11 +90,11 @@ async def test_extractors_api() -> None: "instruction": "Test Instruction", } response = await client.post( - "/extractors", json=create_request, cookies=cookies + "/extractors", json=create_request, headers=headers ) extractor_uuid = response.json()["uuid"] assert response.status_code == 200 - response = await client.get(f"/extractors/{extractor_uuid}", cookies=cookies) + response = await client.get(f"/extractors/{extractor_uuid}", headers=headers) response_data = response.json() assert extractor_uuid == response_data["uuid"] assert "my extractor" == response_data["name"] @@ -107,8 +105,8 @@ async def test_sharing_extractor() -> None: """Test sharing an extractor.""" async with get_async_client() as client: user_id = str(uuid.uuid4()) - cookies = {"user_id": user_id} - response = await client.get("/extractors", cookies=cookies) + headers = {"x-key": user_id} + response = await client.get("/extractors", headers=headers) assert response.status_code == 200 assert response.json() == [] # Verify that we can create an extractor @@ -119,28 +117,28 @@ async def test_sharing_extractor() -> None: "instruction": "Test Instruction", } response = await client.post( - "/extractors", json=create_request, cookies=cookies + "/extractors", json=create_request, headers=headers ) assert response.status_code == 200 uuid_str = response.json()["uuid"] # Generate a share uuid - response = await client.post(f"/extractors/{uuid_str}/share", cookies=cookies) + response = await client.post(f"/extractors/{uuid_str}/share", headers=headers) assert response.status_code == 200 assert "share_uuid" in response.json() share_uuid = response.json()["share_uuid"] # Test idempotency - response = await client.post(f"/extractors/{uuid_str}/share", cookies=cookies) + response = await client.post(f"/extractors/{uuid_str}/share", headers=headers) assert response.status_code == 200 assert "share_uuid" in response.json() assert response.json()["share_uuid"] == share_uuid - # Check cookies - bad_cookies = {"user_id": str(uuid.uuid4())} + # Check headers + bad_headers = {"x-key": str(uuid.uuid4())} response = await client.post( - f"/extractors/{uuid_str}/share", cookies=bad_cookies + f"/extractors/{uuid_str}/share", headers=bad_headers ) assert response.status_code == 404 diff --git a/backend/tests/unit_tests/api/test_api_examples.py b/backend/tests/unit_tests/api/test_api_examples.py index 23affc7..d3ccab1 100644 --- a/backend/tests/unit_tests/api/test_api_examples.py +++ b/backend/tests/unit_tests/api/test_api_examples.py @@ -16,7 +16,7 @@ async def test_examples_api() -> None: async with get_async_client() as client: # First create an extractor user_id = str(uuid.uuid4()) - cookies = {"user_id": user_id} + headers = {"x-key": user_id} create_request = { "description": "Test Description", "name": "Test Name", @@ -24,7 +24,7 @@ async def test_examples_api() -> None: "instruction": "Test Instruction", } response = await client.post( - "/extractors", json=create_request, cookies=cookies + "/extractors", json=create_request, headers=headers ) assert response.status_code == 200 # Get the extractor id @@ -32,7 +32,7 @@ async def test_examples_api() -> None: # Let's verify that there are no examples response = await client.get( - "/examples?extractor_id=" + extractor_id, cookies=cookies + "/examples?extractor_id=" + extractor_id, headers=headers ) assert response.status_code == 200 assert response.json() == [] @@ -48,20 +48,20 @@ async def test_examples_api() -> None: } ], } - response = await client.post("/examples", json=create_request, cookies=cookies) + response = await client.post("/examples", json=create_request, headers=headers) assert response.status_code == 200 example_id = response.json()["uuid"] - # Check cookies - bad_cookies = {"user_id": str(uuid.uuid4())} + # Check headers + bad_headers = {"x-key": str(uuid.uuid4())} response = await client.post( - "/examples", json=create_request, cookies=bad_cookies + "/examples", json=create_request, headers=bad_headers ) assert response.status_code == 404 # Verify that the example was created response = await client.get( - "/examples?extractor_id=" + extractor_id, cookies=cookies + "/examples?extractor_id=" + extractor_id, headers=headers ) assert response.status_code == 200 assert len(response.json()) == 1 @@ -82,23 +82,23 @@ async def test_examples_api() -> None: "uuid": example_id, } - # Check cookies + # Check headers response = await client.get( - "/examples?extractor_id=" + extractor_id, cookies=bad_cookies + "/examples?extractor_id=" + extractor_id, headers=bad_headers ) assert response.status_code == 404 # Check we need cookie to delete - response = await client.delete(f"/examples/{example_id}", cookies=bad_cookies) + response = await client.delete(f"/examples/{example_id}", headers=bad_headers) assert response.status_code == 404 # Verify that we can delete an example - response = await client.delete(f"/examples/{example_id}", cookies=cookies) + response = await client.delete(f"/examples/{example_id}", headers=headers) assert response.status_code == 200 # Verify that the example was deleted response = await client.get( - "/examples?extractor_id=" + extractor_id, cookies=cookies + "/examples?extractor_id=" + extractor_id, headers=headers ) assert response.status_code == 200 assert response.json() == [] diff --git a/backend/tests/unit_tests/api/test_api_extract.py b/backend/tests/unit_tests/api/test_api_extract.py index f999d64..2860c6d 100644 --- a/backend/tests/unit_tests/api/test_api_extract.py +++ b/backend/tests/unit_tests/api/test_api_extract.py @@ -42,7 +42,7 @@ async def test_extract_from_file() -> None: """Test extract from file API.""" async with get_async_client() as client: user_id = str(uuid4()) - cookies = {"user_id": user_id} + headers = {"x-key": user_id} # Test with invalid extractor extractor_id = UUID(int=1027) # 1027 is a good number. response = await client.post( @@ -51,7 +51,7 @@ async def test_extract_from_file() -> None: "extractor_id": str(extractor_id), "text": "Test Content", }, - cookies=cookies, + headers=headers, ) assert response.status_code == 404, response.text @@ -63,7 +63,9 @@ async def test_extract_from_file() -> None: "instruction": "Test Instruction", } response = await client.post( - "/extractors", json=create_request, cookies=cookies + "/extractors", + json=create_request, + headers=headers, ) assert response.status_code == 200, response.text # Get the extractor id @@ -78,7 +80,7 @@ async def test_extract_from_file() -> None: "text": "Test Content", "mode": "entire_document", }, - cookies=cookies, + headers=headers, ) assert response.status_code == 200 assert response.json() == {"data": ["Test Conte"]} @@ -92,7 +94,7 @@ async def test_extract_from_file() -> None: "mode": "entire_document", "model_name": "gpt-3.5-turbo", }, - cookies=cookies, + headers=headers, ) assert response.status_code == 200 assert response.json() == {"data": ["Test Conte"]} @@ -105,7 +107,7 @@ async def test_extract_from_file() -> None: "text": "Test Content", "mode": "retrieval", }, - cookies=cookies, + headers=headers, ) assert response.status_code == 200 assert response.json() == {"data": ["Test Conte"]} @@ -123,7 +125,7 @@ async def test_extract_from_file() -> None: "mode": "entire_document", }, files={"file": f}, - cookies=cookies, + headers=headers, ) assert response.status_code == 200, response.text @@ -132,7 +134,7 @@ async def test_extract_from_file() -> None: async def test_extract_from_large_file() -> None: user_id = str(uuid4()) - cookies = {"user_id": user_id} + headers = {"x-key": user_id} async with get_async_client() as client: # First create an extractor create_request = { @@ -142,7 +144,7 @@ async def test_extract_from_large_file() -> None: "instruction": "Test Instruction", } response = await client.post( - "/extractors", json=create_request, cookies=cookies + "/extractors", json=create_request, headers=headers ) assert response.status_code == 200, response.text # Get the extractor id @@ -161,7 +163,7 @@ async def test_extract_from_large_file() -> None: "mode": "entire_document", }, files={"file": f}, - cookies=cookies, + headers=headers, ) assert response.status_code == 413 @@ -181,6 +183,6 @@ async def test_extract_from_large_file() -> None: "mode": "entire_document", }, files={"file": f.name}, - cookies=cookies, + headers=headers, ) assert response.status_code == 413 diff --git a/frontend/app/components/Sidebar.tsx b/frontend/app/components/Sidebar.tsx index 2eae4ad..b9c1733 100644 --- a/frontend/app/components/Sidebar.tsx +++ b/frontend/app/components/Sidebar.tsx @@ -17,7 +17,6 @@ import { useDisclosure, } from "@chakra-ui/react"; import React from "react"; -import axios from "axios"; import { ArrowTopRightOnSquareIcon, EllipsisVerticalIcon, @@ -26,7 +25,11 @@ import { } from "@heroicons/react/24/outline"; import { useRouter, useParams } from "next/navigation"; import { useMutation } from "@tanstack/react-query"; -import { useDeleteExtractor, useGetExtractors } from "../utils/api"; +import { + useDeleteExtractor, + useGetExtractors, + axiosClient, +} from "../utils/api"; import { getBaseApiUrl } from "../utils/api_url"; import { ShareModal } from "./ShareModal"; @@ -43,7 +46,7 @@ export function Sidebar() { const baseUrl = getBaseApiUrl(); const mutateShare = useMutation({ mutationFn: async (uuid: string) => - axios.post(`${baseUrl}/extractors/${uuid}/share`), + axiosClient.post(`${baseUrl}/extractors/${uuid}/share`), onSuccess: (onSuccessData) => { setShareUUID(onSuccessData.data.share_uuid); onOpen(); diff --git a/frontend/app/utils/api.tsx b/frontend/app/utils/api.tsx index 93820b8..e9976e9 100644 --- a/frontend/app/utils/api.tsx +++ b/frontend/app/utils/api.tsx @@ -1,3 +1,4 @@ +"use client"; /* Expose API hooks for use in components */ import axios from "axios"; import { @@ -7,6 +8,7 @@ import { MutationFunction, QueryFunctionContext, } from "@tanstack/react-query"; +import { v4 as uuidv4 } from "uuid"; import { getBaseApiUrl } from "./api_url"; type ExtractorData = { @@ -17,36 +19,58 @@ type ExtractorData = { schema: any; }; +const getApiKey = (): string => { + if (typeof window === "undefined") { + return uuidv4(); + } + + const key = localStorage.getItem("lc-extract-key"); + if (!key) { + // Generate key + const newKey = uuidv4(); + localStorage.setItem("lc-extract-key", newKey); + return newKey; + } + return key; +}; + +// Create an instance with custom headers +export const axiosClient = axios.create({ + headers: { + "x-key": getApiKey(), + }, +}); + type GetExtractorQueryKey = [string, string, boolean]; // [queryKey, uuid, isShared] type OnSuccessFn = (data: { uuid: string }) => void; -axios.defaults.withCredentials = true; - const getExtractor = async ({ queryKey, }: QueryFunctionContext): Promise => { const [, uuid, isShared] = queryKey; const baseUrl = getBaseApiUrl(); if (isShared) { - const response = await axios.get(`${baseUrl}/shared/extractors/${uuid}`); + const response = await axiosClient.get( + `${baseUrl}/shared/extractors/${uuid}`, + ); return response.data; } else { - const response = await axios.get(`${baseUrl}/extractors/${uuid}`); + const response = await axiosClient.get(`${baseUrl}/extractors/${uuid}`); return response.data; } }; const listExtractors = async () => { const baseUrl = getBaseApiUrl(); - const response = await axios.get(`${baseUrl}/extractors`); + const response = await axiosClient.get(`${baseUrl}/extractors`); return response.data; }; // eslint-disable-next-line @typescript-eslint/no-explicit-any const createExtractor: MutationFunction = async (extractor) => { const baseUrl = getBaseApiUrl(); - const response = await axios.post(`${baseUrl}/extractors`, extractor); + const response = await axiosClient.post(`${baseUrl}/extractors`, extractor); return response.data; }; @@ -58,7 +82,7 @@ export type ServerConfiguration = { const getConfiguration = async (): Promise => { const baseUrl = getBaseApiUrl(); - const response = await axios.get(`${baseUrl}/configuration`); + const response = await axiosClient.get(`${baseUrl}/configuration`); return response.data; }; @@ -80,7 +104,7 @@ export const suggestExtractor = async ({ return {}; } const baseUrl = getBaseApiUrl(); - const response = await axios.post(`${baseUrl}/suggest`, { + const response = await axiosClient.post(`${baseUrl}/suggest`, { description, jsonSchema, }); @@ -103,7 +127,7 @@ export const runExtraction: MutationFunction< > = async ([extractionRequest, isShared]) => { const endpoint = isShared ? "extract/shared" : "extract"; const baseUrl = getBaseApiUrl(); - const response = await axios.postForm( + const response = await axiosClient.postForm( `${baseUrl}/${endpoint}`, extractionRequest, ); @@ -129,7 +153,8 @@ export const useDeleteExtractor = () => { const baseUrl = getBaseApiUrl(); const queryClient = useQueryClient(); return useMutation({ - mutationFn: (uuid: string) => axios.delete(`${baseUrl}/extractors/${uuid}`), + mutationFn: (uuid: string) => + axiosClient.delete(`${baseUrl}/extractors/${uuid}`), onSuccess: () => { queryClient.invalidateQueries({ queryKey: ["getExtractors"] }); },