From 4b806c3510c5ab6e257150cd3108c37e336aced2 Mon Sep 17 00:00:00 2001 From: UnniKohonen Date: Fri, 14 Aug 2026 11:30:49 +0000 Subject: [PATCH 1/7] Add loading PDF files from URL with REST API and cli --- bibra/api/v0/routes.py | 101 ++++++++++++++++- bibra/cli.py | 50 +++++++++ tests/test_api_routes.py | 57 +++++++++- tests/test_cli.py | 214 ++++++++++++++++++++++++++++++++++++- tests/test_schemathesis.py | 35 ++++++ uv.lock | 4 + 6 files changed, 457 insertions(+), 4 deletions(-) diff --git a/bibra/api/v0/routes.py b/bibra/api/v0/routes.py index 864b5ef..40c6d37 100644 --- a/bibra/api/v0/routes.py +++ b/bibra/api/v0/routes.py @@ -3,9 +3,12 @@ import logging import os import tempfile +import urllib.request from typing import Annotated +from urllib.error import HTTPError, URLError -from fastapi import APIRouter, Depends, File, HTTPException, Request, UploadFile +from fastapi import APIRouter, Depends, File, Form, HTTPException, Request, UploadFile +from pydantic import HttpUrl from bibra import __version__ from bibra.config import ConfigError, ProjectNotFoundError, ProjectRegistry @@ -94,3 +97,99 @@ async def extract( logger.debug( "Failed to remove temporary file: %s", tmp_path, exc_info=True ) + + +@router.post( + "/projects/{project_id}/extract-url", + responses={400: {"description": "Bad Request - malformed data"}}, +) +async def extract_url( + project_id: str, + registry: Annotated[ProjectRegistry, Depends(get_registry)], + urls: list[HttpUrl] = Form(...), +) -> PublicationMetadata: + """ + Extract publication metadata from PDF or image files at given URLs for a specific project. + + Args: + project_id: The ID of the project to extract metadata for + urls: List of URLs, each pointing to a file to process + + Returns: + PublicationMetadata: Extracted metadata as JSON + """ + temp_files: list[str] = [] + try: + for url in urls: + # Create a temporary file to save the downloaded PDF for the current URL + with tempfile.NamedTemporaryFile(suffix=".pdf", delete=False) as tmp: + temp_files.append(tmp.name) + + try: + # Download content from URL + with urllib.request.urlopen(str(url)) as response: + content_type = response.info().get_content_type() + if content_type != "application/pdf": + if os.path.exists(tmp.name): + os.unlink(tmp.name) + temp_files.remove(tmp.name) + raise HTTPException( + status_code=400, + detail=f"'{url}' does not point to a PDF file. Expected 'application/pdf', got '{content_type}'.", + ) + + while chunk := response.read(1024 * 1024): + tmp.write(chunk) + + except HTTPException: + raise + except HTTPError as e: + logger.exception( + f"HTTP Error downloading {url}: {e.code} - {e.reason}" + ) + # Clean up the temporary file associated with the failed download + if os.path.exists(tmp.name): + os.unlink(tmp.name) + temp_files.remove(tmp.name) + raise HTTPException(status_code=e.code, detail=str(e)) + except URLError as e: + logger.exception(f"URL Error downloading {url}: {e.reason}") + # Clean up the temporary file associated with the failed download + if os.path.exists(tmp.name): + os.unlink(tmp.name) + temp_files.remove(tmp.name) + raise HTTPException(status_code=400, detail=str(e)) + except Exception as e: + logger.exception( + f"An unexpected error occurred during download from {url}: {e}", + exc_info=True, + ) + # Clean up the temporary file associated with the failed download + if os.path.exists(tmp.name): + os.unlink(tmp.name) + temp_files.remove(tmp.name) + raise HTTPException(status_code=500, detail=str(e)) + + # Get backend for the project + try: + backend = registry.get_backend(project_id) + except ProjectNotFoundError as e: + raise HTTPException(status_code=404, detail=str(e)) + except ConfigError as e: + logger.exception("Configuration error") + raise HTTPException(status_code=500, detail=str(e)) + # Extract metadata using the backend + result = await backend.extract(temp_files) + return result + finally: + # Clean up all temporary files that were successfully added to temp_files + for tmp_path in temp_files: + try: + if os.path.exists( + tmp_path + ): # Check if it still exists (might have been removed by an earlier error) + os.unlink(tmp_path) + except OSError: + logger.debug( + "Failed to remove temporary file: %s", tmp_path, exc_info=True + ) diff --git a/bibra/cli.py b/bibra/cli.py index 8faf46a..436686b 100644 --- a/bibra/cli.py +++ b/bibra/cli.py @@ -1,6 +1,8 @@ """CLI interface for BIBRA.""" import asyncio +import tempfile +import urllib.request import click from dotenv import load_dotenv @@ -103,3 +105,51 @@ def extract(project_id: str, file_path: str, config: str | None, output: str | N click.echo(f"Output written to {output}") else: click.echo(json_output) + + +@cli.command("extract-url") +@click.argument("project_id") +@click.argument("url") +@click.option( + "--config", + "-c", + default=None, + help="Path to the project configuration file (overrides BIBRA_CONFIG).", +) +@click.option( + "--output", + "-o", + type=click.Path(dir_okay=False, writable=True, resolve_path=True), + default=None, + help="Write JSON output to file instead of stdout", +) +def extract_url(project_id: str, url: str, config: str | None, output: str | None): + """Extract publication metadata from a PDF or image file at a URL.""" + registry = ProjectRegistry(config) + + try: + backend = registry.get_backend(project_id) + except ProjectNotFoundError as e: + raise click.UsageError(str(e)) from None + except ConfigError as e: + raise click.ClickException(str(e)) from None + + try: + with tempfile.NamedTemporaryFile(suffix=".pdf") as tmp: + with urllib.request.urlopen(url) as response: + while chunk := response.read(1024 * 1024): + tmp.write(chunk) + + tmp.flush() + result = asyncio.run(backend.extract([tmp.name])) + except Exception as e: + raise click.ClickException(f"Extraction failed: {e}") from e + + json_output = result.model_dump_json(indent=2) + + if output: + with open(output, "w", encoding="utf-8") as f: + f.write(json_output + "\n") + click.echo(f"Output written to {output}") + else: + click.echo(json_output) diff --git a/tests/test_api_routes.py b/tests/test_api_routes.py index 5c9fa0b..29ee185 100644 --- a/tests/test_api_routes.py +++ b/tests/test_api_routes.py @@ -1,6 +1,6 @@ """Tests for API routes.""" -from unittest.mock import MagicMock +from unittest.mock import MagicMock, patch import pytest from fastapi import HTTPException, Request @@ -8,6 +8,7 @@ from bibra.api.v0.routes import ( extract, + extract_url, get_registry, list_projects, router, @@ -80,6 +81,60 @@ async def test_extract_returns_example_metadata(self): assert result.p_isbn == [] assert result.e_issn is None + def test_extract_url_route_exists(self): + """The router should have a project-specific extract-url route.""" + routes = [str(r.path) for r in router.routes] + assert "/projects/{project_id}/extract-url" in routes + + def test_extract_url_route_is_post_method(self): + """The extract-url route should use POST method.""" + extract_url_routes = [ + r + for r in router.routes + if str(r.path) == "/projects/{project_id}/extract-url" + ] + assert len(extract_url_routes) >= 1 + # Check that the route uses POST method + route = extract_url_routes[0] + assert isinstance(route, APIRoute) + + async def test_extract_returns_example_metadata(self): + """The /projects/{project_id}/extract-url endpoint should return example + publication metadata.""" + registry = ProjectRegistry() + + # Mock pdf file download + mock_response = MagicMock() + mock_response.info.return_value.get_content_type.return_value = ( + "application/pdf" + ) + mock_response.read.side_effect = [b"%PDF-1.4 mock content", b""] + mock_response.__enter__.return_value = mock_response + mock_response.__exit__.return_value = False + + with patch("urllib.request.urlopen", return_value=mock_response): + result = await extract_url( + project_id="dummy", + registry=registry, + urls=["https://example.com/paper.pdf"], + ) + + assert isinstance(result, PublicationMetadata) + assert result.language == "en" + assert ( + result.title == "Machine Learning Approaches for Software Defect Prediction" + ) + assert result.creator == ["Smith, John", "Johnson, Emily"] + assert result.year == "2023" + assert result.publisher == ["Springer", "ACM"] + assert result.doi == "10.1234/example.doi.12345" + assert result.e_isbn == ["978-0-123456-78-9"] + assert result.type_coar == "article" + # Verify fields that don't have values are None or empty lists + assert result.alt_title is None + assert result.p_isbn == [] + assert result.e_issn is None + def test_get_registry_returns_existing_registry(self): """get_registry should return the registry already on app.state.""" mock_state = MagicMock() diff --git a/tests/test_cli.py b/tests/test_cli.py index 5575e60..65ede7d 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -7,8 +7,8 @@ import pytest from click.testing import CliRunner -from bibra.cli import _make_list_template, cli, extract, list_projects -from bibra.config import ConfigError +from bibra.cli import _make_list_template, cli, extract, extract_url, list_projects +from bibra.config import ConfigError, ProjectNotFoundError class TestCli: @@ -247,6 +247,216 @@ def test_extract_config_error_converted_to_click_exception(self, tmp_path): assert "Invalid config syntax" in result.output +"""Tests for the extract-url command.""" + + +def _make_urlopen_mock(chunks=(b"%PDF-1.4 dummy content", b"")): + """Build a mock for urllib.request.urlopen that yields the given chunks + from .read() and supports use as a context manager.""" + mock_response = MagicMock() + mock_response.read.side_effect = list(chunks) + + mock_cm = MagicMock() + mock_cm.__enter__.return_value = mock_response + mock_cm.__exit__.return_value = False + return mock_cm + + +def _make_backend(json_payload=None): + """Build a mock backend whose .extract() returns an object with a + model_dump_json method, matching what asyncio.run(backend.extract(...)) + is expected to produce.""" + if json_payload is None: + json_payload = {"title": "Some Paper", "authors": ["A. Author"]} + + mock_result = MagicMock() + mock_result.model_dump_json.return_value = json.dumps(json_payload, indent=2) + + mock_backend = MagicMock() + + async def _extract(*args, **kwargs): + return mock_result + + mock_backend.extract.side_effect = _extract + return mock_backend + + +class TestExtractUrl: + """Tests for the extract-url command.""" + + def setup_method(self): + """Set up test fixtures.""" + self.runner = CliRunner() + + def test_extract_url_help(self): + """Test extract-url help output.""" + result = self.runner.invoke(extract_url, ["--help"]) + assert result.exit_code == 0 + assert not result.exception + assert "Extract publication metadata" in result.output + assert "PROJECT_ID" in result.output + assert "URL" in result.output + assert "--output" in result.output + assert "-o" in result.output + + def test_extract_url_missing_url(self): + """Test extract-url command with only a project id (missing URL).""" + result = self.runner.invoke(extract_url, ["test-project"]) + assert result.exit_code != 0 + assert result.exception + + def test_extract_url_with_valid_url(self): + """Test extract-url command with a valid URL and successful extraction.""" + with ( + patch("bibra.cli.ProjectRegistry") as mock_registry_cls, + patch("bibra.cli.urllib.request.urlopen") as mock_urlopen, + ): + mock_registry = MagicMock() + mock_registry_cls.return_value = mock_registry + mock_registry.get_backend.return_value = _make_backend() + mock_urlopen.return_value = _make_urlopen_mock() + + result = self.runner.invoke( + extract_url, ["dummy", "https://example.com/paper.pdf"] + ) + + assert result.exit_code == 0 + assert not result.exception + + output = result.output.strip() + data = json.loads(output) + assert "title" in data or "authors" in data + + def test_extract_url_with_output_option(self, tmp_path): + """Test extract-url command with --output option to write JSON to file.""" + output_file = tmp_path / "output.json" + + with ( + patch("bibra.cli.ProjectRegistry") as mock_registry_cls, + patch("bibra.cli.urllib.request.urlopen") as mock_urlopen, + ): + mock_registry = MagicMock() + mock_registry_cls.return_value = mock_registry + mock_registry.get_backend.return_value = _make_backend() + mock_urlopen.return_value = _make_urlopen_mock() + + result = self.runner.invoke( + extract_url, + [ + "dummy", + "https://example.com/paper.pdf", + "--output", + str(output_file), + ], + ) + + assert result.exit_code == 0 + assert not result.exception + assert "Output written to" in result.output + + assert output_file.exists() + content = output_file.read_text(encoding="utf-8") + data = json.loads(content.strip()) + assert data is not None + + def test_extract_url_with_short_output_option(self, tmp_path): + """Test extract-url command with -o short option.""" + output_file = tmp_path / "output.json" + + with ( + patch("bibra.cli.ProjectRegistry") as mock_registry_cls, + patch("bibra.cli.urllib.request.urlopen") as mock_urlopen, + ): + mock_registry = MagicMock() + mock_registry_cls.return_value = mock_registry + mock_registry.get_backend.return_value = _make_backend() + mock_urlopen.return_value = _make_urlopen_mock() + + result = self.runner.invoke( + extract_url, + ["dummy", "https://example.com/paper.pdf", "-o", str(output_file)], + ) + + assert result.exit_code == 0 + assert not result.exception + assert "Output written to" in result.output + + def test_extract_url_with_nonexistent_project(self): + """Test extract-url command with a project that isn't found.""" + with patch("bibra.cli.ProjectRegistry") as mock_registry_cls: + mock_registry = MagicMock() + mock_registry_cls.return_value = mock_registry + mock_registry.get_backend.side_effect = ProjectNotFoundError( + "Project 'nonexistent-project' not found" + ) + + result = self.runner.invoke( + extract_url, + ["nonexistent-project", "https://example.com/paper.pdf"], + ) + + assert result.exit_code != 0 + assert result.exception + assert "not found" in result.output + + def test_extract_url_config_error_converted_to_click_exception(self): + """Test that ConfigError while resolving the backend becomes a ClickException.""" + with patch("bibra.cli.ProjectRegistry") as mock_registry_cls: + mock_registry = MagicMock() + mock_registry_cls.return_value = mock_registry + mock_registry.get_backend.side_effect = ConfigError("Invalid config syntax") + + result = self.runner.invoke( + extract_url, ["dummy", "https://example.com/paper.pdf"] + ) + + assert result.exit_code != 0 + assert "Invalid config syntax" in result.output + + def test_extract_url_download_failure_converted_to_click_exception(self): + """Test that a failure while downloading the URL is wrapped as 'Extraction failed:'.""" + with ( + patch("bibra.cli.ProjectRegistry") as mock_registry_cls, + patch("bibra.cli.urllib.request.urlopen") as mock_urlopen, + ): + mock_registry = MagicMock() + mock_registry_cls.return_value = mock_registry + mock_registry.get_backend.return_value = _make_backend() + mock_urlopen.side_effect = OSError("Name or service not known") + + result = self.runner.invoke( + extract_url, ["dummy", "https://bad.example.invalid/paper.pdf"] + ) + + assert result.exit_code != 0 + assert "Extraction failed:" in result.output + + def test_extract_url_generic_exception_converted_to_click_exception(self): + """Test that a generic Exception during extraction is wrapped in + ClickException with the 'Extraction failed:' prefix.""" + with ( + patch("bibra.cli.ProjectRegistry") as mock_registry_cls, + patch("bibra.cli.urllib.request.urlopen") as mock_urlopen, + ): + mock_registry = MagicMock() + mock_registry_cls.return_value = mock_registry + mock_backend = MagicMock() + + async def _extract(*args, **kwargs): + raise RuntimeError("PDF corrupted") + + mock_backend.extract.side_effect = _extract + mock_registry.get_backend.return_value = mock_backend + mock_urlopen.return_value = _make_urlopen_mock() + + result = self.runner.invoke( + extract_url, ["dummy", "https://example.com/paper.pdf"] + ) + + assert result.exit_code != 0 + assert "Extraction failed: PDF corrupted" in result.output + + class TestMakeListTemplate: """Tests for the _make_list_template helper function.""" diff --git a/tests/test_schemathesis.py b/tests/test_schemathesis.py index 2083c35..aa6112c 100644 --- a/tests/test_schemathesis.py +++ b/tests/test_schemathesis.py @@ -1,3 +1,5 @@ +from unittest.mock import MagicMock, patch + import schemathesis from bibra.main import app @@ -6,6 +8,15 @@ schema = schemathesis.openapi.from_asgi("/openapi.json", app) +def _mock_pdf_response(*args, **kwargs): + mock_response = MagicMock() + mock_response.info.return_value.get_content_type.return_value = "application/pdf" + mock_response.read.side_effect = [b"%PDF-1.4 mock content", b""] + mock_response.__enter__.return_value = mock_response + mock_response.__exit__.return_value = False + return mock_response + + @schema.parametrize() def test_api(case): # Skip extract cases missing the required `files` field. @@ -23,4 +34,28 @@ def test_api(case): if hasattr(case, "path") and case.path == "/v0/projects/{project_id}/extract": # Modify the path to use dummy project case.path = "/v0/projects/dummy/extract" + # Skip extract-url cases missing the required `urls` field. + is_extract_url = ( + case.path == "/v0/projects/{project_id}/extract-url" + and case.method.upper() == "POST" + ) + if is_extract_url: + body = case.body + # Skip if `urls` is absent or empty + has_urls = isinstance(body, dict) and bool(body.get("urls")) + if not has_urls: + return + # Use dummy backend for testing to avoid real network downloads/API calls + if ( + hasattr(case, "path") + and case.path == "/v0/projects/{project_id}/extract-url" + ): + # Modify the path to use dummy project + case.path = "/v0/projects/dummy/extract-url" + + # Mock pdf file download + with patch("urllib.request.urlopen", side_effect=_mock_pdf_response): + case.call_and_validate() + return + case.call_and_validate() diff --git a/uv.lock b/uv.lock index 576d841..445ca84 100644 --- a/uv.lock +++ b/uv.lock @@ -2,6 +2,10 @@ version = 1 revision = 3 requires-python = ">=3.11" +[options] +exclude-newer = "2026-08-06T09:44:44.7887759Z" +exclude-newer-span = "P1W" + [[package]] name = "annotated-doc" version = "0.0.4" From 535a7c853d4014124fa65b41d906102bc12077ed Mon Sep 17 00:00:00 2001 From: Osma Suominen Date: Fri, 14 Aug 2026 16:30:50 +0300 Subject: [PATCH 2/7] feat(api): Switch URL downloading to httpx for async support Replaces synchronous `urllib.request` with `httpx.AsyncClient` in `extract_url`. This allows for non-blocking HTTP requests, improving performance when downloading multiple files concurrently. The refactoring handles file downloading, content type checking, and error management asynchronously, ensuring robust handling of network and file system operations. --- bibra/api/v0/routes.py | 138 +++++++++++++++++++------------------ bibra/cli.py | 11 +-- tests/test_api_routes.py | 28 +++++--- tests/test_cli.py | 43 ++++++------ tests/test_schemathesis.py | 28 +++++--- 5 files changed, 136 insertions(+), 112 deletions(-) diff --git a/bibra/api/v0/routes.py b/bibra/api/v0/routes.py index 40c6d37..a052269 100644 --- a/bibra/api/v0/routes.py +++ b/bibra/api/v0/routes.py @@ -3,10 +3,9 @@ import logging import os import tempfile -import urllib.request from typing import Annotated -from urllib.error import HTTPError, URLError +import httpx from fastapi import APIRouter, Depends, File, Form, HTTPException, Request, UploadFile from pydantic import HttpUrl @@ -106,10 +105,11 @@ async def extract( async def extract_url( project_id: str, registry: Annotated[ProjectRegistry, Depends(get_registry)], - urls: list[HttpUrl] = Form(...), + urls: list[HttpUrl] = Form(...), # noqa: B008 ) -> PublicationMetadata: """ - Extract publication metadata from PDF or image files at given URLs for a specific project. + Extract publication metadata from PDF or image files at given URLs for a + specific project. Args: project_id: The ID of the project to extract metadata for @@ -120,74 +120,76 @@ async def extract_url( """ temp_files: list[str] = [] try: - for url in urls: - # Create a temporary file to save the downloaded PDF for the current URL - with tempfile.NamedTemporaryFile(suffix=".pdf", delete=False) as tmp: - temp_files.append(tmp.name) - - try: - # Download content from URL - with urllib.request.urlopen(str(url)) as response: - content_type = response.info().get_content_type() - if content_type != "application/pdf": - if os.path.exists(tmp.name): - os.unlink(tmp.name) - temp_files.remove(tmp.name) - raise HTTPException( - status_code=400, - detail=f"'{url}' does not point to a PDF file. Expected 'application/pdf', got '{content_type}'.", - ) - - while chunk := response.read(1024 * 1024): - tmp.write(chunk) - - except HTTPException: - raise - except HTTPError as e: - logger.exception( - f"HTTP Error downloading {url}: {e.code} - {e.reason}" - ) - # Clean up the temporary file associated with the failed download - if os.path.exists(tmp.name): - os.unlink(tmp.name) - temp_files.remove(tmp.name) - raise HTTPException(status_code=e.code, detail=str(e)) - except URLError as e: - logger.exception(f"URL Error downloading {url}: {e.reason}") - # Clean up the temporary file associated with the failed download - if os.path.exists(tmp.name): - os.unlink(tmp.name) - temp_files.remove(tmp.name) - raise HTTPException(status_code=400, detail=str(e)) - except Exception as e: - logger.exception( - f"An unexpected error occurred during download from {url}: {e}", - exc_info=True, - ) - # Clean up the temporary file associated with the failed download - if os.path.exists(tmp.name): - os.unlink(tmp.name) - temp_files.remove(tmp.name) - raise HTTPException(status_code=500, detail=str(e)) - - # Get backend for the project - try: - backend = registry.get_backend(project_id) - except ProjectNotFoundError as e: - raise HTTPException(status_code=404, detail=str(e)) - except ConfigError as e: - logger.exception("Configuration error") - raise HTTPException(status_code=500, detail=str(e)) - # Extract metadata using the backend - result = await backend.extract(temp_files) - return result + async with httpx.AsyncClient() as client: + for url in urls: + # Create a temporary file to save the downloaded PDF for the current URL + with tempfile.NamedTemporaryFile(suffix=".pdf", delete=False) as tmp: + temp_files.append(tmp.name) + + try: + url_str = str(url) + async with client.stream("GET", url_str) as response: + content_type = response.headers.get("content-type", "") + if content_type != "application/pdf": + if os.path.exists(tmp.name): + os.unlink(tmp.name) + temp_files.remove(tmp.name) + expected = "application/pdf" + detail = ( + f"'{url}' does not point to a PDF file. " + f"Expected '{expected}', got '{content_type}'." + ) + raise HTTPException(status_code=400, detail=detail) + + status_code = response.status_code + if status_code >= 400: + if os.path.exists(tmp.name): + os.unlink(tmp.name) + temp_files.remove(tmp.name) + raise HTTPException( + status_code=status_code, + detail=str(response.reason_phrase), + ) + + async for chunk in response.aiter_bytes( + chunk_size=1024 * 1024 + ): + tmp.write(chunk) + + except HTTPException: + raise + except httpx.HTTPError as e: + logger.exception("HTTP Error downloading %s", url_str) + if os.path.exists(tmp.name): + os.unlink(tmp.name) + temp_files.remove(tmp.name) + raise HTTPException(status_code=500, detail=str(e)) + except Exception as e: + logger.exception( + "Unexpected error during download from %s", url + ) + if os.path.exists(tmp.name): + os.unlink(tmp.name) + temp_files.remove(tmp.name) + raise HTTPException(status_code=500, detail=str(e)) + + # Get backend for the project + try: + backend = registry.get_backend(project_id) + except ProjectNotFoundError as e: + raise HTTPException(status_code=404, detail=str(e)) + except ConfigError as e: + logger.exception("Configuration error") + raise HTTPException(status_code=500, detail=str(e)) + # Extract metadata using the backend + result = await backend.extract(temp_files) + return result finally: # Clean up all temporary files that were successfully added to temp_files for tmp_path in temp_files: try: - if os.path.exists( - tmp_path - ): # Check if it still exists (might have been removed by an earlier error) + # Check if it still exists (might have been removed earlier) + if os.path.exists(tmp_path): os.unlink(tmp_path) except OSError: logger.debug( diff --git a/bibra/cli.py b/bibra/cli.py index 436686b..6901a8e 100644 --- a/bibra/cli.py +++ b/bibra/cli.py @@ -2,9 +2,9 @@ import asyncio import tempfile -import urllib.request import click +import httpx from dotenv import load_dotenv from bibra.config import ( @@ -136,11 +136,12 @@ def extract_url(project_id: str, url: str, config: str | None, output: str | Non try: with tempfile.NamedTemporaryFile(suffix=".pdf") as tmp: - with urllib.request.urlopen(url) as response: - while chunk := response.read(1024 * 1024): + with httpx.stream("GET", url) as response: + if response.status_code >= 400: + raise httpx.HTTPError(f"HTTP {response.status_code} for {url}") + for chunk in response.iter_bytes(chunk_size=1024 * 1024): tmp.write(chunk) - - tmp.flush() + tmp.flush() result = asyncio.run(backend.extract([tmp.name])) except Exception as e: raise click.ClickException(f"Extraction failed: {e}") from e diff --git a/tests/test_api_routes.py b/tests/test_api_routes.py index 29ee185..5fa9de0 100644 --- a/tests/test_api_routes.py +++ b/tests/test_api_routes.py @@ -98,21 +98,29 @@ def test_extract_url_route_is_post_method(self): route = extract_url_routes[0] assert isinstance(route, APIRoute) - async def test_extract_returns_example_metadata(self): + async def test_extract_url_returns_example_metadata(self): """The /projects/{project_id}/extract-url endpoint should return example publication metadata.""" registry = ProjectRegistry() - # Mock pdf file download + # Mock httpx.AsyncClient stream response mock_response = MagicMock() - mock_response.info.return_value.get_content_type.return_value = ( - "application/pdf" - ) - mock_response.read.side_effect = [b"%PDF-1.4 mock content", b""] - mock_response.__enter__.return_value = mock_response - mock_response.__exit__.return_value = False - - with patch("urllib.request.urlopen", return_value=mock_response): + mock_response.headers.get.return_value = "application/pdf" + mock_response.status_code = 200 + + async def mock_aiter_bytes(*args, **kwargs): + yield b"%PDF-1.4 mock content" + + mock_response.aiter_bytes.return_value = mock_aiter_bytes() + mock_response.__aenter__.return_value = mock_response + mock_response.__aexit__.return_value = False + + with patch("httpx.AsyncClient") as mock_client_cls: + mock_client = MagicMock() + mock_client.__aenter__.return_value = mock_client + mock_client.__aexit__.return_value = False + mock_client.stream = MagicMock(return_value=mock_response) + mock_client_cls.return_value = mock_client result = await extract_url( project_id="dummy", registry=registry, diff --git a/tests/test_cli.py b/tests/test_cli.py index 65ede7d..53ffcce 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -250,16 +250,15 @@ def test_extract_config_error_converted_to_click_exception(self, tmp_path): """Tests for the extract-url command.""" -def _make_urlopen_mock(chunks=(b"%PDF-1.4 dummy content", b"")): - """Build a mock for urllib.request.urlopen that yields the given chunks - from .read() and supports use as a context manager.""" +def _make_httpx_stream_mock(chunks=(b"%PDF-1.4 dummy content",)): + """Build a mock for httpx stream response that yields the given chunks.""" mock_response = MagicMock() - mock_response.read.side_effect = list(chunks) - - mock_cm = MagicMock() - mock_cm.__enter__.return_value = mock_response - mock_cm.__exit__.return_value = False - return mock_cm + mock_response.headers.get.return_value = "application/pdf" + mock_response.status_code = 200 + mock_response.iter_bytes.return_value = chunks + mock_response.__enter__.return_value = mock_response + mock_response.__exit__.return_value = False + return mock_response def _make_backend(json_payload=None): @@ -309,12 +308,12 @@ def test_extract_url_with_valid_url(self): """Test extract-url command with a valid URL and successful extraction.""" with ( patch("bibra.cli.ProjectRegistry") as mock_registry_cls, - patch("bibra.cli.urllib.request.urlopen") as mock_urlopen, + patch("bibra.cli.httpx.stream") as mock_stream, ): mock_registry = MagicMock() mock_registry_cls.return_value = mock_registry mock_registry.get_backend.return_value = _make_backend() - mock_urlopen.return_value = _make_urlopen_mock() + mock_stream.return_value = _make_httpx_stream_mock() result = self.runner.invoke( extract_url, ["dummy", "https://example.com/paper.pdf"] @@ -333,12 +332,12 @@ def test_extract_url_with_output_option(self, tmp_path): with ( patch("bibra.cli.ProjectRegistry") as mock_registry_cls, - patch("bibra.cli.urllib.request.urlopen") as mock_urlopen, + patch("bibra.cli.httpx.stream") as mock_stream, ): mock_registry = MagicMock() mock_registry_cls.return_value = mock_registry mock_registry.get_backend.return_value = _make_backend() - mock_urlopen.return_value = _make_urlopen_mock() + mock_stream.return_value = _make_httpx_stream_mock() result = self.runner.invoke( extract_url, @@ -365,12 +364,12 @@ def test_extract_url_with_short_output_option(self, tmp_path): with ( patch("bibra.cli.ProjectRegistry") as mock_registry_cls, - patch("bibra.cli.urllib.request.urlopen") as mock_urlopen, + patch("bibra.cli.httpx.stream") as mock_stream, ): mock_registry = MagicMock() mock_registry_cls.return_value = mock_registry mock_registry.get_backend.return_value = _make_backend() - mock_urlopen.return_value = _make_urlopen_mock() + mock_stream.return_value = _make_httpx_stream_mock() result = self.runner.invoke( extract_url, @@ -400,7 +399,7 @@ def test_extract_url_with_nonexistent_project(self): assert "not found" in result.output def test_extract_url_config_error_converted_to_click_exception(self): - """Test that ConfigError while resolving the backend becomes a ClickException.""" + """Test ConfigError during backend resolution becomes ClickException.""" with patch("bibra.cli.ProjectRegistry") as mock_registry_cls: mock_registry = MagicMock() mock_registry_cls.return_value = mock_registry @@ -414,15 +413,17 @@ def test_extract_url_config_error_converted_to_click_exception(self): assert "Invalid config syntax" in result.output def test_extract_url_download_failure_converted_to_click_exception(self): - """Test that a failure while downloading the URL is wrapped as 'Extraction failed:'.""" + """Test download failure is wrapped as 'Extraction failed:'.""" with ( patch("bibra.cli.ProjectRegistry") as mock_registry_cls, - patch("bibra.cli.urllib.request.urlopen") as mock_urlopen, + patch("bibra.cli.httpx.stream") as mock_stream, ): + import httpx as _httpx + mock_registry = MagicMock() mock_registry_cls.return_value = mock_registry mock_registry.get_backend.return_value = _make_backend() - mock_urlopen.side_effect = OSError("Name or service not known") + mock_stream.side_effect = _httpx.HTTPError("Name or service not known") result = self.runner.invoke( extract_url, ["dummy", "https://bad.example.invalid/paper.pdf"] @@ -436,7 +437,7 @@ def test_extract_url_generic_exception_converted_to_click_exception(self): ClickException with the 'Extraction failed:' prefix.""" with ( patch("bibra.cli.ProjectRegistry") as mock_registry_cls, - patch("bibra.cli.urllib.request.urlopen") as mock_urlopen, + patch("bibra.cli.httpx.stream") as mock_stream, ): mock_registry = MagicMock() mock_registry_cls.return_value = mock_registry @@ -447,7 +448,7 @@ async def _extract(*args, **kwargs): mock_backend.extract.side_effect = _extract mock_registry.get_backend.return_value = mock_backend - mock_urlopen.return_value = _make_urlopen_mock() + mock_stream.return_value = _make_httpx_stream_mock() result = self.runner.invoke( extract_url, ["dummy", "https://example.com/paper.pdf"] diff --git a/tests/test_schemathesis.py b/tests/test_schemathesis.py index aa6112c..3b0a741 100644 --- a/tests/test_schemathesis.py +++ b/tests/test_schemathesis.py @@ -8,12 +8,19 @@ schema = schemathesis.openapi.from_asgi("/openapi.json", app) -def _mock_pdf_response(*args, **kwargs): +async def _mock_aiter_bytes(*args, **kwargs): + """Async generator that yields mock PDF content.""" + yield b"%PDF-1.4 mock content" + + +def _mock_httpx_response(*args, **kwargs): + """Build a mock httpx stream response for PDF content.""" mock_response = MagicMock() - mock_response.info.return_value.get_content_type.return_value = "application/pdf" - mock_response.read.side_effect = [b"%PDF-1.4 mock content", b""] - mock_response.__enter__.return_value = mock_response - mock_response.__exit__.return_value = False + mock_response.headers.get.return_value = "application/pdf" + mock_response.status_code = 200 + mock_response.aiter_bytes.return_value = _mock_aiter_bytes() + mock_response.__aenter__.return_value = mock_response + mock_response.__aexit__.return_value = False return mock_response @@ -53,9 +60,14 @@ def test_api(case): # Modify the path to use dummy project case.path = "/v0/projects/dummy/extract-url" - # Mock pdf file download - with patch("urllib.request.urlopen", side_effect=_mock_pdf_response): - case.call_and_validate() + # Mock httpx.AsyncClient stream response + with patch("httpx.AsyncClient") as mock_client: + mock_client.__aenter__.return_value = mock_client + mock_client.__aexit__.return_value = False + mock_response = _mock_httpx_response() + mock_client.stream = MagicMock(return_value=mock_response) + with patch("httpx.AsyncClient", return_value=mock_client): + case.call_and_validate() return case.call_and_validate() From 8b5db1aa6309f6f5df0d15479a4cbb18f5288f71 Mon Sep 17 00:00:00 2001 From: Osma Suominen Date: Fri, 14 Aug 2026 16:57:37 +0300 Subject: [PATCH 3/7] feat(api): Refactor URL extraction endpoint to handle single file input Refactors the `extract_url` endpoint to accept a single `HttpUrl` instead of a list of URLs. This simplifies the API contract and logic, focusing on processing one file at a time. The implementation is updated to streamline the file download and validation process, ensuring that the input URL points to a PDF file and handling HTTP errors gracefully. It also introduces proper error handling for missing projects during backend retrieval. --- bibra/api/v0/routes.py | 117 +++++++++++++------------------------ tests/test_api_routes.py | 4 +- tests/test_schemathesis.py | 8 +-- 3 files changed, 49 insertions(+), 80 deletions(-) diff --git a/bibra/api/v0/routes.py b/bibra/api/v0/routes.py index a052269..a288003 100644 --- a/bibra/api/v0/routes.py +++ b/bibra/api/v0/routes.py @@ -105,93 +105,60 @@ async def extract( async def extract_url( project_id: str, registry: Annotated[ProjectRegistry, Depends(get_registry)], - urls: list[HttpUrl] = Form(...), # noqa: B008 + url: HttpUrl = Form(...), # noqa: B008 ) -> PublicationMetadata: """ - Extract publication metadata from PDF or image files at given URLs for a + Extract publication metadata from a PDF or image file at a given URL for a specific project. Args: project_id: The ID of the project to extract metadata for - urls: List of URLs, each pointing to a file to process + url: URL pointing to a file to process Returns: PublicationMetadata: Extracted metadata as JSON """ - temp_files: list[str] = [] try: async with httpx.AsyncClient() as client: - for url in urls: - # Create a temporary file to save the downloaded PDF for the current URL - with tempfile.NamedTemporaryFile(suffix=".pdf", delete=False) as tmp: - temp_files.append(tmp.name) - + url_str = str(url) + async with client.stream("GET", url_str) as response: + content_type = response.headers.get("content-type", "") + if content_type != "application/pdf": + expected = "application/pdf" + detail = ( + f"'{url}' does not point to a PDF file. " + f"Expected '{expected}', got '{content_type}'." + ) + raise HTTPException(status_code=400, detail=detail) + + status_code = response.status_code + if status_code >= 400: + raise HTTPException( + status_code=status_code, + detail=str(response.reason_phrase), + ) + + with tempfile.NamedTemporaryFile(suffix=".pdf", delete=True) as tmp: + async for chunk in response.aiter_bytes(chunk_size=1024 * 1024): + tmp.write(chunk) + tmp.flush() + + # Get backend for the project try: - url_str = str(url) - async with client.stream("GET", url_str) as response: - content_type = response.headers.get("content-type", "") - if content_type != "application/pdf": - if os.path.exists(tmp.name): - os.unlink(tmp.name) - temp_files.remove(tmp.name) - expected = "application/pdf" - detail = ( - f"'{url}' does not point to a PDF file. " - f"Expected '{expected}', got '{content_type}'." - ) - raise HTTPException(status_code=400, detail=detail) - - status_code = response.status_code - if status_code >= 400: - if os.path.exists(tmp.name): - os.unlink(tmp.name) - temp_files.remove(tmp.name) - raise HTTPException( - status_code=status_code, - detail=str(response.reason_phrase), - ) - - async for chunk in response.aiter_bytes( - chunk_size=1024 * 1024 - ): - tmp.write(chunk) - - except HTTPException: - raise - except httpx.HTTPError as e: - logger.exception("HTTP Error downloading %s", url_str) - if os.path.exists(tmp.name): - os.unlink(tmp.name) - temp_files.remove(tmp.name) - raise HTTPException(status_code=500, detail=str(e)) - except Exception as e: - logger.exception( - "Unexpected error during download from %s", url - ) - if os.path.exists(tmp.name): - os.unlink(tmp.name) - temp_files.remove(tmp.name) + backend = registry.get_backend(project_id) + except ProjectNotFoundError as e: + raise HTTPException(status_code=404, detail=str(e)) + except ConfigError as e: + logger.exception("Configuration error") raise HTTPException(status_code=500, detail=str(e)) - # Get backend for the project - try: - backend = registry.get_backend(project_id) - except ProjectNotFoundError as e: - raise HTTPException(status_code=404, detail=str(e)) - except ConfigError as e: - logger.exception("Configuration error") - raise HTTPException(status_code=500, detail=str(e)) - # Extract metadata using the backend - result = await backend.extract(temp_files) - return result - finally: - # Clean up all temporary files that were successfully added to temp_files - for tmp_path in temp_files: - try: - # Check if it still exists (might have been removed earlier) - if os.path.exists(tmp_path): - os.unlink(tmp_path) - except OSError: - logger.debug( - "Failed to remove temporary file: %s", tmp_path, exc_info=True - ) + # Extract metadata using the backend + result = await backend.extract([tmp.name]) + return result + + except httpx.HTTPError as e: + logger.exception("HTTP Error downloading %s", url_str) + raise HTTPException(status_code=500, detail=str(e)) + except Exception as e: + logger.exception("Unexpected error during download from %s", url) + raise HTTPException(status_code=500, detail=str(e)) diff --git a/tests/test_api_routes.py b/tests/test_api_routes.py index 5fa9de0..0dbb799 100644 --- a/tests/test_api_routes.py +++ b/tests/test_api_routes.py @@ -101,6 +101,8 @@ def test_extract_url_route_is_post_method(self): async def test_extract_url_returns_example_metadata(self): """The /projects/{project_id}/extract-url endpoint should return example publication metadata.""" + from pydantic import HttpUrl + registry = ProjectRegistry() # Mock httpx.AsyncClient stream response @@ -124,7 +126,7 @@ async def mock_aiter_bytes(*args, **kwargs): result = await extract_url( project_id="dummy", registry=registry, - urls=["https://example.com/paper.pdf"], + url=HttpUrl("https://example.com/paper.pdf"), ) assert isinstance(result, PublicationMetadata) diff --git a/tests/test_schemathesis.py b/tests/test_schemathesis.py index 3b0a741..f220d3f 100644 --- a/tests/test_schemathesis.py +++ b/tests/test_schemathesis.py @@ -41,16 +41,16 @@ def test_api(case): if hasattr(case, "path") and case.path == "/v0/projects/{project_id}/extract": # Modify the path to use dummy project case.path = "/v0/projects/dummy/extract" - # Skip extract-url cases missing the required `urls` field. + # Skip extract-url cases missing the required `url` field. is_extract_url = ( case.path == "/v0/projects/{project_id}/extract-url" and case.method.upper() == "POST" ) if is_extract_url: body = case.body - # Skip if `urls` is absent or empty - has_urls = isinstance(body, dict) and bool(body.get("urls")) - if not has_urls: + # Skip if `url` is absent or empty + has_url = isinstance(body, dict) and bool(body.get("url")) + if not has_url: return # Use dummy backend for testing to avoid real network downloads/API calls if ( From 00ee1577588778d9c1872af0e1cc8176f42b5783 Mon Sep 17 00:00:00 2001 From: Osma Suominen Date: Fri, 14 Aug 2026 17:12:00 +0300 Subject: [PATCH 4/7] feat(api): Refactor URL extraction to use project backend for metadata extraction Refactors the `extract_url` function to first retrieve the project's backend using `registry.get_backend`. This ensures that the metadata extraction process is handled by the configured backend, improving modularity and error handling. The logic is updated to: 1. Fail fast if the project or configuration is missing. 2. Stream the PDF content into a temporary file. 3. Pass the temporary file path to the backend for metadata extraction. This change centralizes the extraction logic and improves robustness against configuration and project errors. --- bibra/api/v0/routes.py | 30 +++++++++++++----------------- 1 file changed, 13 insertions(+), 17 deletions(-) diff --git a/bibra/api/v0/routes.py b/bibra/api/v0/routes.py index a288003..4ff1f58 100644 --- a/bibra/api/v0/routes.py +++ b/bibra/api/v0/routes.py @@ -119,8 +119,17 @@ async def extract_url( PublicationMetadata: Extracted metadata as JSON """ try: - async with httpx.AsyncClient() as client: - url_str = str(url) + backend = registry.get_backend(project_id) + except ProjectNotFoundError as e: + raise HTTPException(status_code=404, detail=str(e)) + except ConfigError as e: + logger.exception("Configuration error") + raise HTTPException(status_code=500, detail=str(e)) + + url_str = str(url) + + try: + async with httpx.AsyncClient() as client: # noqa: SIM117 async with client.stream("GET", url_str) as response: content_type = response.headers.get("content-type", "") if content_type != "application/pdf": @@ -138,24 +147,11 @@ async def extract_url( detail=str(response.reason_phrase), ) - with tempfile.NamedTemporaryFile(suffix=".pdf", delete=True) as tmp: + with tempfile.NamedTemporaryFile(suffix=".pdf") as tmp: async for chunk in response.aiter_bytes(chunk_size=1024 * 1024): tmp.write(chunk) tmp.flush() - - # Get backend for the project - try: - backend = registry.get_backend(project_id) - except ProjectNotFoundError as e: - raise HTTPException(status_code=404, detail=str(e)) - except ConfigError as e: - logger.exception("Configuration error") - raise HTTPException(status_code=500, detail=str(e)) - - # Extract metadata using the backend - result = await backend.extract([tmp.name]) - return result - + return await backend.extract([tmp.name]) except httpx.HTTPError as e: logger.exception("HTTP Error downloading %s", url_str) raise HTTPException(status_code=500, detail=str(e)) From 67c991740eb5794981b42a8fff586622a4207f0d Mon Sep 17 00:00:00 2001 From: Osma Suominen Date: Fri, 14 Aug 2026 17:18:52 +0300 Subject: [PATCH 5/7] refactor: Simplify URL extraction by using client.get() instead of streaming Replaces the use of `httpx.AsyncClient.stream` with a direct `client.get()` call in `extract_url`. This simplifies the asynchronous HTTP request handling logic. The corresponding tests in `test_api_routes.py` are updated to mock the new `client.get()` behavior instead of the previous streaming interface, ensuring compatibility with the refactored function. --- bibra/api/v0/routes.py | 47 ++++++++++++++++++++-------------------- tests/test_api_routes.py | 19 +++++++--------- 2 files changed, 32 insertions(+), 34 deletions(-) diff --git a/bibra/api/v0/routes.py b/bibra/api/v0/routes.py index 4ff1f58..4fee494 100644 --- a/bibra/api/v0/routes.py +++ b/bibra/api/v0/routes.py @@ -129,29 +129,30 @@ async def extract_url( url_str = str(url) try: - async with httpx.AsyncClient() as client: # noqa: SIM117 - async with client.stream("GET", url_str) as response: - content_type = response.headers.get("content-type", "") - if content_type != "application/pdf": - expected = "application/pdf" - detail = ( - f"'{url}' does not point to a PDF file. " - f"Expected '{expected}', got '{content_type}'." - ) - raise HTTPException(status_code=400, detail=detail) - - status_code = response.status_code - if status_code >= 400: - raise HTTPException( - status_code=status_code, - detail=str(response.reason_phrase), - ) - - with tempfile.NamedTemporaryFile(suffix=".pdf") as tmp: - async for chunk in response.aiter_bytes(chunk_size=1024 * 1024): - tmp.write(chunk) - tmp.flush() - return await backend.extract([tmp.name]) + async with httpx.AsyncClient() as client: + response = await client.get(url_str) + + content_type = response.headers.get("content-type", "") + if content_type != "application/pdf": + expected = "application/pdf" + detail = ( + f"'{url}' does not point to a PDF file. " + f"Expected '{expected}', got '{content_type}'." + ) + raise HTTPException(status_code=400, detail=detail) + + status_code = response.status_code + if status_code >= 400: + raise HTTPException( + status_code=status_code, + detail=str(response.reason_phrase), + ) + + with tempfile.NamedTemporaryFile(suffix=".pdf") as tmp: + async for chunk in response.aiter_bytes(chunk_size=1024 * 1024): + tmp.write(chunk) + tmp.flush() + return await backend.extract([tmp.name]) except httpx.HTTPError as e: logger.exception("HTTP Error downloading %s", url_str) raise HTTPException(status_code=500, detail=str(e)) diff --git a/tests/test_api_routes.py b/tests/test_api_routes.py index 0dbb799..daa8a82 100644 --- a/tests/test_api_routes.py +++ b/tests/test_api_routes.py @@ -1,10 +1,11 @@ """Tests for API routes.""" -from unittest.mock import MagicMock, patch +from unittest.mock import AsyncMock, MagicMock, patch import pytest from fastapi import HTTPException, Request from fastapi.routing import APIRoute +from httpx import Headers from bibra.api.v0.routes import ( extract, @@ -105,23 +106,19 @@ async def test_extract_url_returns_example_metadata(self): registry = ProjectRegistry() - # Mock httpx.AsyncClient stream response - mock_response = MagicMock() - mock_response.headers.get.return_value = "application/pdf" - mock_response.status_code = 200 - async def mock_aiter_bytes(*args, **kwargs): yield b"%PDF-1.4 mock content" + mock_response = MagicMock() + mock_response.headers = Headers({"content-type": "application/pdf"}) + mock_response.status_code = 200 mock_response.aiter_bytes.return_value = mock_aiter_bytes() - mock_response.__aenter__.return_value = mock_response - mock_response.__aexit__.return_value = False with patch("httpx.AsyncClient") as mock_client_cls: mock_client = MagicMock() - mock_client.__aenter__.return_value = mock_client - mock_client.__aexit__.return_value = False - mock_client.stream = MagicMock(return_value=mock_response) + mock_client.__aenter__ = AsyncMock(return_value=mock_client) + mock_client.__aexit__ = AsyncMock(return_value=False) + mock_client.get = AsyncMock(return_value=mock_response) mock_client_cls.return_value = mock_client result = await extract_url( project_id="dummy", From 97e2fff908e7a92bc3bf873c0690f4ce4aff77f4 Mon Sep 17 00:00:00 2001 From: Osma Suominen Date: Fri, 14 Aug 2026 17:25:12 +0300 Subject: [PATCH 6/7] fix(testing): Improve mock setup for Schemathesis API tests Refactors the mocking logic in `test_schemathesis.py` to correctly simulate asynchronous HTTP responses when testing API endpoints. This ensures that the test environment accurately reflects how `httpx.AsyncClient` handles streaming responses, leading to more reliable and accurate test coverage. --- tests/test_schemathesis.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/tests/test_schemathesis.py b/tests/test_schemathesis.py index f220d3f..ccb544f 100644 --- a/tests/test_schemathesis.py +++ b/tests/test_schemathesis.py @@ -61,11 +61,16 @@ def test_api(case): case.path = "/v0/projects/dummy/extract-url" # Mock httpx.AsyncClient stream response + mock_response = _mock_httpx_response() + + async def async_get(*args, **kwargs): + return mock_response + with patch("httpx.AsyncClient") as mock_client: mock_client.__aenter__.return_value = mock_client mock_client.__aexit__.return_value = False - mock_response = _mock_httpx_response() mock_client.stream = MagicMock(return_value=mock_response) + mock_client.get = async_get with patch("httpx.AsyncClient", return_value=mock_client): case.call_and_validate() return From c216db9ef9fbc8ece1bedde517984df181ff5dcd Mon Sep 17 00:00:00 2001 From: Osma Suominen Date: Fri, 14 Aug 2026 17:48:34 +0300 Subject: [PATCH 7/7] feat: Implement URL download proxy for SSRF mitigation Introduces support for routing all external URL downloads through a configurable proxy. This change adds `BIBRA_URL_PROXY` environment variable support, allowing users to specify a proxy URL. The proxy is integrated into: - `bibra/config.py` via `get_url_proxy()` - `bibra/api/v0/routes.py` (async client) - `bibra/cli.py` (synchronous stream) This feature helps mitigate Server-Side Request Forgery (SSRF) risks by controlling outbound network traffic. Corresponding tests are added to verify proxy usage. --- .env.example | 4 +++ bibra/api/v0/routes.py | 10 +++++-- bibra/cli.py | 4 ++- bibra/config.py | 9 ++++++ tests/test_api_routes.py | 62 ++++++++++++++++++++++++++++++++++++++++ tests/test_cli.py | 47 ++++++++++++++++++++++++++++++ tests/test_config.py | 22 ++++++++++++++ 7 files changed, 155 insertions(+), 3 deletions(-) diff --git a/.env.example b/.env.example index f0ac87e..8f5375b 100644 --- a/.env.example +++ b/.env.example @@ -13,5 +13,9 @@ NUEXTRACT_MODEL=nuextract3 #NUEXTRACT_INSTRUCTIONS= #NUEXTRACT_DPI= +# URL Download Proxy (for SSRF mitigation) +# All URL downloads will be routed through this proxy. +#BIBRA_URL_PROXY=http://proxy.example.com:8080 + # System Settings (optional) # LOG_LEVEL=INFO diff --git a/bibra/api/v0/routes.py b/bibra/api/v0/routes.py index 4fee494..e5618a9 100644 --- a/bibra/api/v0/routes.py +++ b/bibra/api/v0/routes.py @@ -10,7 +10,12 @@ from pydantic import HttpUrl from bibra import __version__ -from bibra.config import ConfigError, ProjectNotFoundError, ProjectRegistry +from bibra.config import ( + ConfigError, + ProjectNotFoundError, + ProjectRegistry, + get_url_proxy, +) from bibra.types import PublicationMetadata logger = logging.getLogger(__name__) @@ -129,7 +134,8 @@ async def extract_url( url_str = str(url) try: - async with httpx.AsyncClient() as client: + proxy = get_url_proxy() + async with httpx.AsyncClient(proxy=proxy) as client: response = await client.get(url_str) content_type = response.headers.get("content-type", "") diff --git a/bibra/cli.py b/bibra/cli.py index 6901a8e..8e681da 100644 --- a/bibra/cli.py +++ b/bibra/cli.py @@ -11,6 +11,7 @@ ConfigError, ProjectNotFoundError, ProjectRegistry, + get_url_proxy, ) @@ -136,7 +137,8 @@ def extract_url(project_id: str, url: str, config: str | None, output: str | Non try: with tempfile.NamedTemporaryFile(suffix=".pdf") as tmp: - with httpx.stream("GET", url) as response: + proxy = get_url_proxy() + with httpx.stream("GET", url, proxy=proxy) as response: if response.status_code >= 400: raise httpx.HTTPError(f"HTTP {response.status_code} for {url}") for chunk in response.iter_bytes(chunk_size=1024 * 1024): diff --git a/bibra/config.py b/bibra/config.py index 308080b..8417edb 100644 --- a/bibra/config.py +++ b/bibra/config.py @@ -277,3 +277,12 @@ def list_projects(self) -> list[dict[str, Any]]: } for project in self._projects.values() ] + + +def get_url_proxy() -> str | None: + """Return the BIBRA_URL_PROXY environment variable value. + + Returns: + The proxy URL string, or None if not set. + """ + return os.environ.get("BIBRA_URL_PROXY") diff --git a/tests/test_api_routes.py b/tests/test_api_routes.py index daa8a82..7ae85ed 100644 --- a/tests/test_api_routes.py +++ b/tests/test_api_routes.py @@ -217,3 +217,65 @@ async def test_extract_handles_project_not_found_error(self): assert exc_info.value.status_code == 404 assert exc_info.value.detail == "Project 'unknown' not found" + + async def test_extract_url_passes_proxy_when_set(self, monkeypatch): + """Test that extract-url passes the proxy to httpx.AsyncClient when set.""" + from pydantic import HttpUrl + + monkeypatch.setenv("BIBRA_URL_PROXY", "http://proxy.example.com:8080") + + registry = ProjectRegistry() + + async def mock_aiter_bytes(*args, **kwargs): + yield b"%PDF-1.4 mock content" + + mock_response = MagicMock() + mock_response.headers = Headers({"content-type": "application/pdf"}) + mock_response.status_code = 200 + mock_response.aiter_bytes.return_value = mock_aiter_bytes() + + with patch("httpx.AsyncClient") as mock_client_cls: + mock_client = MagicMock() + mock_client.__aenter__ = AsyncMock(return_value=mock_client) + mock_client.__aexit__ = AsyncMock(return_value=False) + mock_client.get = AsyncMock(return_value=mock_response) + mock_client_cls.return_value = mock_client + result = await extract_url( + project_id="dummy", + registry=registry, + url=HttpUrl("https://example.com/paper.pdf"), + ) + + mock_client_cls.assert_called_once_with(proxy="http://proxy.example.com:8080") + assert isinstance(result, PublicationMetadata) + + async def test_extract_url_no_proxy_when_not_set(self, monkeypatch): + """Test that extract-url passes proxy=None when env var is not set.""" + from pydantic import HttpUrl + + monkeypatch.delenv("BIBRA_URL_PROXY", raising=False) + + registry = ProjectRegistry() + + async def mock_aiter_bytes(*args, **kwargs): + yield b"%PDF-1.4 mock content" + + mock_response = MagicMock() + mock_response.headers = Headers({"content-type": "application/pdf"}) + mock_response.status_code = 200 + mock_response.aiter_bytes.return_value = mock_aiter_bytes() + + with patch("httpx.AsyncClient") as mock_client_cls: + mock_client = MagicMock() + mock_client.__aenter__ = AsyncMock(return_value=mock_client) + mock_client.__aexit__ = AsyncMock(return_value=False) + mock_client.get = AsyncMock(return_value=mock_response) + mock_client_cls.return_value = mock_client + result = await extract_url( + project_id="dummy", + registry=registry, + url=HttpUrl("https://example.com/paper.pdf"), + ) + + mock_client_cls.assert_called_once_with(proxy=None) + assert isinstance(result, PublicationMetadata) diff --git a/tests/test_cli.py b/tests/test_cli.py index 53ffcce..fd2786b 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -457,6 +457,53 @@ async def _extract(*args, **kwargs): assert result.exit_code != 0 assert "Extraction failed: PDF corrupted" in result.output + def test_extract_url_passes_proxy_when_set(self): + """Test that extract-url passes the proxy to httpx.stream when set.""" + with ( + patch("bibra.cli.ProjectRegistry") as mock_registry_cls, + patch("bibra.cli.httpx.stream") as mock_stream, + ): + mock_registry = MagicMock() + mock_registry_cls.return_value = mock_registry + mock_registry.get_backend.return_value = _make_backend() + mock_stream.return_value = _make_httpx_stream_mock() + + result = self.runner.invoke( + extract_url, + ["dummy", "https://example.com/paper.pdf"], + env={"BIBRA_URL_PROXY": "http://proxy.example.com:8080"}, + ) + + assert result.exit_code == 0 + mock_stream.assert_called_once_with( + "GET", + "https://example.com/paper.pdf", + proxy="http://proxy.example.com:8080", + ) + + def test_extract_url_no_proxy_when_not_set(self): + """Test that extract-url passes proxy=None when env var is not set.""" + with ( + patch("bibra.cli.ProjectRegistry") as mock_registry_cls, + patch("bibra.cli.httpx.stream") as mock_stream, + ): + mock_registry = MagicMock() + mock_registry_cls.return_value = mock_registry + mock_registry.get_backend.return_value = _make_backend() + mock_stream.return_value = _make_httpx_stream_mock() + + result = self.runner.invoke( + extract_url, + ["dummy", "https://example.com/paper.pdf"], + ) + + assert result.exit_code == 0 + mock_stream.assert_called_once_with( + "GET", + "https://example.com/paper.pdf", + proxy=None, + ) + class TestMakeListTemplate: """Tests for the _make_list_template helper function.""" diff --git a/tests/test_config.py b/tests/test_config.py index d241e28..7534730 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -607,3 +607,25 @@ def test_unclosed_placeholder_preserved(self): def test_unclosed_placeholder_with_trailing_text(self): """Test that an unclosed ${ followed by text is left as-is.""" assert _interpolate_env_vars("prefix${BAR suffix") == "prefix${BAR suffix" + + +class TestGetUrlProxy: + """Tests for get_url_proxy helper function.""" + + def test_returns_proxy_when_set(self, monkeypatch): + """Test that get_url_proxy returns the proxy URL when set.""" + monkeypatch.setenv("BIBRA_URL_PROXY", "http://proxy.example.com:8080") + from bibra.config import get_url_proxy + + assert get_url_proxy() == "http://proxy.example.com:8080" + + def test_returns_none_when_not_set(self, monkeypatch): + """Test that get_url_proxy returns None when the env var is not set.""" + monkeypatch.delenv("BIBRA_URL_PROXY", raising=False) + # Re-import to get fresh import + import importlib + + import bibra.config + + importlib.reload(bibra.config) + assert bibra.config.get_url_proxy() is None