Send custom HTTP request headers when downloading a source

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

openaddr/VERSION

@@ -1 +1 @@
-10.0.0
+10.1.0

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)

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()
 

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])