diff --git a/CHANGELOG b/CHANGELOG index 898cf584..566e813b 100644 --- a/CHANGELOG +++ b/CHANGELOG @@ -1,3 +1,6 @@ +2026-08-13 v10.1.0 +- Support custom HTTP request headers for a source via `request.headers`, including on the file-extension pre-flight request (needed for downloads gated on Referer/etc.) + 2026-03-27 v10.0.0 - Upgrade base Docker image from GDAL 3.7.1 to 3.11.0 https://github.com/openaddresses/batch-machine/pull/99 - Update SSL CA certificates to fix download failures https://github.com/openaddresses/batch-machine/pull/98 diff --git a/openaddr/VERSION b/openaddr/VERSION index a13e7b9c..4149c39e 100644 --- a/openaddr/VERSION +++ b/openaddr/VERSION @@ -1 +1 @@ -10.0.0 +10.1.0 diff --git a/openaddr/__init__.py b/openaddr/__init__.py index b3fc1098..f16cbe9a 100644 --- a/openaddr/__init__.py +++ b/openaddr/__init__.py @@ -75,8 +75,10 @@ def cache(source_config, destdir, extras): source_urls = [source_urls] protocol_string = source_config.data_source.get('protocol') + request_settings = source_config.data_source.get('request') or {} + source_headers = request_settings.get('headers') or {} - task = DownloadTask.from_protocol_string(protocol_string, source_config) + task = DownloadTask.from_protocol_string(protocol_string, source_config, headers=source_headers) downloaded_files = task.download(source_urls, workdir, source_config) # FIXME: I wrote the download stuff to assume multiple files because @@ -128,6 +130,10 @@ def conform(source_config, destdir, extras): if not isinstance(source_urls, list): source_urls = [source_urls] + # source_config.data_source['request'] is intentionally not passed here: + # this re-downloads from the OA-owned cache artifact (S3), not the + # contributor's original host, so contributor-supplied headers don't + # apply. task1 = URLDownloadTask(source_config.data_source_name) downloaded_path = task1.download(source_urls, workdir, source_config) _L.info("Downloaded to %s", downloaded_path) diff --git a/openaddr/cache.py b/openaddr/cache.py index d1ce4f70..939ba5f6 100644 --- a/openaddr/cache.py +++ b/openaddr/cache.py @@ -130,22 +130,23 @@ def __init__(self, source_prefix, params={}, headers={}): @classmethod - def from_protocol_string(clz, protocol_string, source_prefix=None): + def from_protocol_string(clz, protocol_string, source_prefix=None, headers=None): + headers = headers or {} if protocol_string.lower() == 'http': - return URLDownloadTask(source_prefix) + return URLDownloadTask(source_prefix, headers=headers) elif protocol_string.lower() == 'file': - return URLDownloadTask(source_prefix) + return URLDownloadTask(source_prefix, headers=headers) elif protocol_string.lower() == 'ftp': - return URLDownloadTask(source_prefix) + return URLDownloadTask(source_prefix, headers=headers) elif protocol_string.lower() == 'esri': - return EsriRestDownloadTask(source_prefix) + return EsriRestDownloadTask(source_prefix, headers=headers) else: raise KeyError("I don't know how to extract for protocol {}".format(protocol_string)) def download(self, source_urls, workdir, source_config): raise NotImplementedError() -def guess_url_file_extension(url): +def guess_url_file_extension(url, headers=None): ''' Get a filename extension for a URL using various hints. ''' scheme, _, path, _, query, _ = urlparse(url) @@ -172,7 +173,7 @@ def guess_url_file_extension(url): # Get a dictionary of headers and a few bytes of content from the URL. # if scheme in ('http', 'https'): - response = request('GET', url, stream=True) + response = request('GET', url, headers=headers or {}, stream=True) handle, file = mkstemp() for chunk in response.iter_content(chunk_size=8192): @@ -256,7 +257,7 @@ def get_file_path(self, url, dir_path): hash = sha1((host + path_base).encode('utf-8')) name_base = u'{}-{}'.format(self.source_prefix, hash.hexdigest()[:8]) - path_ext = guess_url_file_extension(url) + path_ext = guess_url_file_extension(url, self.headers) _L.debug(u'Guessed {}{} for {}'.format(name_base, path_ext, url)) return os.path.join(dir_path, name_base + path_ext) @@ -391,7 +392,7 @@ def download(self, source_urls, workdir, source_config): _L.debug("File exists %s", file_path) continue - downloader = EsriDumper(source_url, parent_logger=_L, timeout=300) + downloader = EsriDumper(source_url, parent_logger=_L, timeout=300, extra_headers=self.headers) metadata = downloader.get_metadata() diff --git a/openaddr/tests/cache.py b/openaddr/tests/cache.py index 37ccafef..ad2d947e 100644 --- a/openaddr/tests/cache.py +++ b/openaddr/tests/cache.py @@ -3,6 +3,7 @@ import csv from .. import SourceConfig +from .. import cache as cache_fn from urllib.parse import urlparse, parse_qs from os.path import join, dirname @@ -16,7 +17,13 @@ import httmock import tempfile -from ..cache import guess_url_file_extension, EsriRestDownloadTask +import sys +from ..cache import guess_url_file_extension, EsriRestDownloadTask, URLDownloadTask, DownloadTask + +# openaddr/__init__.py defines a `cache` function that shadows the `cache` +# submodule on the `openaddr` package object, so `from .. import cache` +# would grab the function, not the module. Go through sys.modules instead. +cache_module = sys.modules['openaddr.cache'] class TestCacheExtensionGuessing (unittest.TestCase): @@ -392,3 +399,146 @@ def test_handle_feature_server_with_lat_lon_in_conform(self): self.assertEqual(len(all_data), 5) self.assertTrue('oa:geom' in all_data[0]) self.assertEqual(all_data[0]['oa:geom'], 'POINT (-86.82960553 34.18671398)') + +class TestFromProtocolStringHeaders (unittest.TestCase): + + def test_headers_reach_url_download_task(self): + task = DownloadTask.from_protocol_string('http', 'us-il-champaign', headers={'Referer': 'https://example.gov/'}) + self.assertIsInstance(task, URLDownloadTask) + self.assertEqual(task.headers['Referer'], 'https://example.gov/') + # The default User-Agent is still present alongside custom headers. + self.assertIn('User-Agent', task.headers) + + def test_headers_reach_esri_download_task(self): + task = DownloadTask.from_protocol_string('ESRI', 'us-il-champaign', headers={'Referer': 'https://example.gov/'}) + self.assertIsInstance(task, EsriRestDownloadTask) + self.assertEqual(task.headers['Referer'], 'https://example.gov/') + + def test_headers_reach_esri_dumper(self): + ''' EsriRestDownloadTask.download() must pass its headers to + EsriDumper as extra_headers, since pyesridump makes its own + requests independent of openaddr.cache.request(). + ''' + workdir = tempfile.mkdtemp(prefix='testCacheHeaders-') + try: + task = EsriRestDownloadTask('us-fl-palmbeach', headers={'Referer': 'https://example.gov/'}) + c = SourceConfig(dict({ + "schema": 2, + "layers": { + "addresses": [{ + "name": "default", + "conform": None + }] + } + }), "addresses", "default") + + with patch.object(cache_module, 'EsriDumper') as dumper_patch: + dumper_patch.return_value.get_metadata.return_value = {'fields': []} + dumper_patch.return_value.get_feature_count.return_value = 0 + dumper_patch.return_value.__iter__.return_value = iter([]) + + task.download(['http://example.com/'], workdir, c) + + _, kwargs = dumper_patch.call_args + self.assertEqual(kwargs.get('extra_headers'), task.headers) + self.assertEqual(kwargs['extra_headers']['Referer'], 'https://example.gov/') + finally: + shutil.rmtree(workdir) + + def test_no_headers_still_gets_default_user_agent(self): + task = DownloadTask.from_protocol_string('http', 'us-il-champaign') + self.assertEqual(list(task.headers.keys()), ['User-Agent']) + + def test_custom_user_agent_overrides_default(self): + task = DownloadTask.from_protocol_string('http', 'us-il-champaign', headers={'User-Agent': 'custom-agent/1.0'}) + self.assertEqual(task.headers['User-Agent'], 'custom-agent/1.0') + +class TestURLDownloadTaskHeaders (unittest.TestCase): + ''' Confirm that a source's custom headers are sent on both the + file-extension pre-flight request and the real download request. + ''' + + def setUp(self): + self.workdir = tempfile.mkdtemp(prefix='testCacheHeaders-') + self.seen_referers = [] + + def tearDown(self): + shutil.rmtree(self.workdir) + + def response_content(self, url, request): + scheme, host, path, _, query, _ = urlparse(url.geturl()) + + # A query string forces guess_url_file_extension() to make a + # sniffing request instead of trusting the URL's extension, + # so this URL exercises both the pre-flight and real download. + if (host, path, query) == ('headers-test.local', '/addresses.csv', 'download=true'): + self.seen_referers.append(request.headers.get('Referer')) + return httmock.response(200, b'FAKE,FAKE\n', headers={'Content-Type': 'text/csv'}) + + raise NotImplementedError(url.geturl()) + + def test_headers_sent_on_preflight_and_download_requests(self): + task = URLDownloadTask('us-il-champaign', headers={'Referer': 'https://example.gov/gis/'}) + with httmock.HTTMock(self.response_content): + output_files = task.download(['http://headers-test.local/addresses.csv?download=true'], self.workdir, None) + + self.assertEqual(len(output_files), 1) + # One request for the extension-guessing pre-flight, one for the real download. + self.assertEqual(self.seen_referers, ['https://example.gov/gis/', 'https://example.gov/gis/']) + +class TestCacheRequestSettings (unittest.TestCase): + ''' Confirm that openaddr.cache() reads headers from the nested + request.headers key in the source config, and that a source with + no request settings at all still works. + ''' + + def setUp(self): + self.destdir = tempfile.mkdtemp(prefix='testCacheRequestSettings-') + self.seen_referers = [] + + def tearDown(self): + shutil.rmtree(self.destdir) + + def response_content(self, url, request): + scheme, host, path, _, query, _ = urlparse(url.geturl()) + + # A query string forces guess_url_file_extension() to make a + # sniffing request instead of trusting the URL's extension, + # so this URL exercises both the pre-flight and real download. + if (host, path, query) == ('request-settings-test.local', '/addresses.csv', 'download=true'): + self.seen_referers.append(request.headers.get('Referer')) + return httmock.response(200, b'FAKE,FAKE\n', headers={'Content-Type': 'text/csv'}) + + raise NotImplementedError(url.geturl()) + + def make_source_config(self, layersource_extra): + return SourceConfig(dict({ + "schema": 2, + "layers": { + "addresses": [dict({ + "name": "default", + "protocol": "http", + "data": "http://request-settings-test.local/addresses.csv?download=true", + }, **layersource_extra)] + } + }), "addresses", "default") + + def test_cache_reads_headers_from_request_settings(self): + source_config = self.make_source_config({ + "request": { + "headers": {"Referer": "https://example.gov/gis/"} + } + }) + + with httmock.HTTMock(self.response_content): + cache_fn(source_config, self.destdir, {}) + + self.assertEqual(self.seen_referers, ['https://example.gov/gis/', 'https://example.gov/gis/']) + + def test_cache_without_request_settings_sends_no_referer(self): + source_config = self.make_source_config({}) + + with httmock.HTTMock(self.response_content): + cache_fn(source_config, self.destdir, {}) + + self.assertEqual(self.seen_referers, [None, None])