diff --git a/CHANGELOG.md b/CHANGELOG.md index 8331846..614cfee 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,16 @@ and this project adheres to [Semantic Versioning](http://semver.org/spec/v2.0.0. ## [Unreleased] +### Added + +- `request_headers` on a converter: headers sent with every download, e.g. a browser + User-Agent for a server that refuses other clients. + +### Changed + +- The DuckDB converter downloads sources DuckDB cannot read in place, such as archives, also + without a cache: to a temporary folder, as the GeoDataFrame-based converter does. + ### Fixed - The DuckDB converter no longer adds a column for each `key=value` folder in the path of a diff --git a/tests/test_convert_duckdb.py b/tests/test_convert_duckdb.py index e7a468f..ac6fa44 100644 --- a/tests/test_convert_duckdb.py +++ b/tests/test_convert_duckdb.py @@ -1,6 +1,7 @@ import json import re import sys +import zipfile import geopandas as gpd import numpy as np @@ -109,6 +110,18 @@ def test_duckdb_converter(tmp_folder): assert validation.errors == [] +def test_duckdb_converter_extracts_an_archive_without_a_cache(tmp_folder): + src = _source_file(tmp_folder) + archive = tmp_folder / "source.zip" + with zipfile.ZipFile(archive, "w") as zf: + zf.write(src, "source.parquet") + dest = tmp_folder / "converted.parquet" + + Converter().convert(dest, input_files={str(archive): ["source.parquet"]}) + + assert sorted(gpd.read_parquet(dest)["id"]) == ["bowtie", "multi", "square", "with_z"] + + def test_duckdb_converter_ignores_hive_style_folders(tmp_folder): source_folder = tmp_folder / "region=north" source_folder.mkdir() @@ -125,6 +138,20 @@ def test_duckdb_converter_ignores_hive_style_folders(tmp_folder): assert "region" not in names +@pytest.mark.parametrize( + "urls, in_place", + [ + ("https://host/data.parquet?token=abc", True), + ("https://host/data.geoparquet", True), + ({"https://host/x?id=1": "fields.parquet"}, True), + ({"https://host/data.zip": ["fields.parquet"]}, False), + ("https://host/data.fgb", False), + ], +) +def test_only_parquet_is_read_in_place(urls, in_place): + assert duckdb_module._readable_in_place(urls) is in_place + + def test_duckdb_converter_numbers_rows_without_an_id_mapping(tmp_folder): src = _source_file(tmp_folder) dest = tmp_folder / "converted.parquet" diff --git a/tests/test_download.py b/tests/test_download.py index be36761..c211314 100644 --- a/tests/test_download.py +++ b/tests/test_download.py @@ -7,7 +7,7 @@ import vecorel_cli.conversion.base as base from vecorel_cli.conversion.base import BaseConverter -from vecorel_cli.vecorel.util import stream_file +from vecorel_cli.vecorel.util import get_fs, merge_headers, stream_file URI = "https://example.invalid/data/fields.gml" @@ -192,3 +192,56 @@ def stream(fs, src_uri, dst_file, chunk_size=None): # every volume was cached under its own name for part in parts: assert (tmp_folder / part).exists() + + +def test_request_headers_go_with_every_download(tmp_folder, monkeypatch): + seen = [] + source_fs = _FakeSourceFS(b"body", {}) + + def fake_get_fs(uri, **kwargs): + if str(uri).startswith(("http://", "https://")): + seen.append(kwargs.get("client_kwargs")) + return source_fs + return LocalFileSystem() + + monkeypatch.setattr(base, "get_fs", fake_get_fs) + HeaderConverter = type( + "HeaderConverter", (BaseConverter,), {"request_headers": {"X-Test": "1"}} + ) + + HeaderConverter().download_files(URI, cache_folder=str(tmp_folder)) + + assert seen == [{"headers": {"X-Test": "1"}}] + + +def test_request_headers_merge_with_explicit_client_options(tmp_folder, monkeypatch): + seen = [] + source_fs = _FakeSourceFS(b"body", {}) + + def fake_get_fs(uri, **kwargs): + if str(uri).startswith(("http://", "https://")): + seen.append(kwargs.get("client_kwargs")) + return source_fs + return LocalFileSystem() + + monkeypatch.setattr(base, "get_fs", fake_get_fs) + HeaderConverter = type( + "HeaderConverter", (BaseConverter,), {"request_headers": {"X-Test": "1", "X-Kept": "k"}} + ) + client_kwargs = {"trust_env": True, "headers": {"x-test": "2"}} + + HeaderConverter().download_files(URI, cache_folder=str(tmp_folder), client_kwargs=client_kwargs) + + assert seen == [{"trust_env": True, "headers": {"X-Kept": "k", "x-test": "2"}}] + + +def test_headers_merge_case_insensitively(): + merged = merge_headers( + {"User-Agent": "default", "X-Test": "1"}, {"user-agent": "mine", "x-test": "2"} + ) + assert merged == {"user-agent": "mine", "x-test": "2"} + + +def test_a_converter_user_agent_replaces_the_default(): + fs = get_fs("https://example.invalid/x", client_kwargs={"headers": {"user-agent": "browser"}}) + assert fs.client_kwargs["headers"] == {"user-agent": "browser"} diff --git a/vecorel_cli/conversion/base.py b/vecorel_cli/conversion/base.py index 0e942f6..5ef94ee 100644 --- a/vecorel_cli/conversion/base.py +++ b/vecorel_cli/conversion/base.py @@ -33,7 +33,7 @@ from ..vecorel.hilbert import hilbert_sort_geodataframe from ..vecorel.schemas import Schemas from ..vecorel.typing import Sources -from ..vecorel.util import get_fs, name_from_uri, stream_file, suffix_duplicate_ids +from ..vecorel.util import get_fs, merge_headers, name_from_uri, stream_file, suffix_duplicate_ids from .deflate64 import deflate64_zip from .flatdict import FlatDict @@ -58,6 +58,8 @@ class BaseConverter(LoggerMixin): data_access: str = "" open_options: dict = {} avoid_range_request: bool = False + # Sent with every download, e.g. a browser User-Agent for a server that refuses others + request_headers: dict[str, str] = {} variants: dict[str, Sources] = {} variant: Optional[str] = None @@ -365,6 +367,11 @@ def download_files(self, uris, cache_folder=None, **kwargs): if self.avoid_range_request and "block_size" not in kwargs: kwargs["block_size"] = 0 + if self.request_headers: + # headers passed in explicitly win over the converter's + client_kwargs = kwargs.get("client_kwargs") or {} + headers = merge_headers(self.request_headers, client_kwargs.get("headers", {})) + kwargs["client_kwargs"] = {**client_kwargs, "headers": headers} # Multi-volume 7z archives (.7z.001, .7z.002, ...) are one 7z stream split # across several URIs; they must be downloaded and extracted together. diff --git a/vecorel_cli/conversion/duckdb.py b/vecorel_cli/conversion/duckdb.py index bb849e4..806a04d 100644 --- a/vecorel_cli/conversion/duckdb.py +++ b/vecorel_cli/conversion/duckdb.py @@ -21,7 +21,7 @@ report, warn_missing_required, ) -from ..vecorel.util import find_differing_crs +from ..vecorel.util import find_differing_crs, name_from_uri from .base import BaseConverter @@ -32,6 +32,15 @@ def _sql_path(path) -> str: return f"'{escaped}'" +def _readable_in_place(urls) -> bool: + if isinstance(urls, str): + urls = {urls: name_from_uri(urls)} + return all( + isinstance(target, str) and target.lower().endswith((".parquet", ".geoparquet")) + for target in urls.values() + ) + + def _sql_name(name) -> str: escaped = str(name).replace('"', '""') return f'"{escaped}"' @@ -116,7 +125,9 @@ def convert( raise ValueError("No input files provided") self.info("Getting file(s) if not cached yet") - if cache: + # Without a cache DuckDB reads remote parquet files in place; anything else, such as + # an archive, is downloaded to a temporary folder as the GeoDataFrame codepath does + if cache or not _readable_in_place(urls): request_args = {} if self.avoid_range_request: request_args["block_size"] = 0 diff --git a/vecorel_cli/vecorel/util.py b/vecorel_cli/vecorel/util.py index 4e0812d..8ab50c8 100644 --- a/vecorel_cli/vecorel/util.py +++ b/vecorel_cli/vecorel/util.py @@ -120,6 +120,17 @@ def _expected_content_length(f) -> Optional[int]: return None +def merge_headers(*headers: dict) -> dict: + """Merge HTTP headers; names are case-insensitive, and a later one replaces an earlier.""" + merged = {} + for group in headers: + for name, value in group.items(): + for existing in [k for k in merged if k.lower() == name.lower()]: + del merged[existing] + merged[name] = value + return merged + + def get_fs(url_or_path: Union[str, Path, URL], **kwargs) -> AbstractFileSystem: """Choose fsspec filesystem by sniffing input url""" if isinstance(url_or_path, Path): @@ -130,7 +141,7 @@ def get_fs(url_or_path: Union[str, Path, URL], **kwargs) -> AbstractFileSystem: if parsed.scheme in ("http", "https"): client_kwargs = kwargs.pop("client_kwargs", {}) - headers = {"User-Agent": USER_AGENT, **client_kwargs.get("headers", {})} + headers = merge_headers({"User-Agent": USER_AGENT}, client_kwargs.get("headers", {})) return HTTPFileSystem(client_kwargs={**client_kwargs, "headers": headers}, **kwargs) if parsed.scheme == "s3":