Merge pull request #215 from lonvia/fix-socket-timeouts-again

Simplify handling of custom request parameters

Author: lonvia

src/osmium/replication/server.py

@@ -32,13 +32,19 @@ class ReplicationServer:
         Replication change files allow to keep local OSM data up-to-date without
         downloading the full dataset again.
 
+        `url` contains the base URL of the replication service. This is the
+        directory that contains the state file with the current state. If the
+        replication service serves something other than osc.gz files, set
+        the `diff_type` to the given file suffix.
+
         ReplicationServer may be used as a context manager. In this case, it
         internally keeps a connection to the server making downloads faster.
     """
 
     def __init__(self, url: str, diff_type: str = 'osc.gz') -> None:
         self.baseurl = url
         self.diff_type = diff_type
+        self.extra_request_params: Mapping[str, Any] = dict(timeout=60, stream=True)
         self.session: Optional[requests.Session] = None
 
     def close(self) -> None:
@@ -55,39 +61,37 @@ def __enter__(self) -> 'ReplicationServer':
     def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None:
         self.close()
 
+    def set_request_parameter(self, key: str, value: Any) -> None:
+        """ Set a parameter which will be handed to the requests library
+            when calling `requests.get()`. This
+            may be used to set custom headers, timeouts and similar parameters.
+            See the `requests documentation <https://requests.readthedocs.io/en/latest/api/?highlight=get#requests.request>`_
+            for possible parameters. Per default, a timeout of 60 sec is set
+            and streaming download enabled.
+        """
+        self.extra_request_params[key] = value
+
     def make_request(self, url: str) -> urlrequest.Request:
         headers = {"User-Agent" : f"pyosmium/{version.pyosmium_release}"}
         return urlrequest.Request(url, headers=headers)
 
     def open_url(self, url: urlrequest.Request) -> Any:
         """ Download a resource from the given URL and return a byte sequence
             of the content.
-
-            This method has no support for cookies or any special authentication
-            methods. If you need these, you have to provide your own custom URL
-            opener. Overwrite open_url() with a method that receives an
-            urlrequest.Request object and returns a ByteIO-like object or a
-            requests.Response.
-
-            Example::
-
-                opener = urlrequest.build_opener()
-                opener.addheaders = [('X-Fancy-Header', 'important_content')]
-
-                svr = ReplicationServer()
-                svr.open_url = opener.open
         """
-        headers = dict()
-        for h in url.header_items():
-            headers[h[0]] = h[1]
+        if 'headers' in self.extra_request_params:
+            get_params = self.extra_request_params
+        else:
+            get_params = dict(self.extra_request_params)
+            get_params['headers'] = {k: v for k,v in url.header_items()}
 
         if self.session is not None:
-            return self.session.get(url.get_full_url(), headers=headers, stream=True)
+            return self.session.get(url.get_full_url(), **get_params)
 
         @contextmanager
         def _get_url_with_session() -> Iterator[requests.Response]:
             with requests.Session() as session:
-                request = session.get(url.get_full_url(), headers=headers, stream=True)
+                request = session.get(url.get_full_url(), **get_params)
                 yield request
 
         return _get_url_with_session()

test/test_pyosmium_get_changes.py

@@ -48,13 +48,6 @@ def mock_get(session, url, **kwargs):
         monkeypatch.setattr(osmium.replication.server.requests.Session, "get", mock_get)
 
 
-    @pytest.fixture
-    def mock_urllib(self, monkeypatch):
-        def mock_get(_, url, **kwargs):
-            return BytesIO(self.urls[url.get_full_url()])
-        monkeypatch.setattr(self.script['urlrequest'].OpenerDirector, "open", mock_get)
-
-
     def url(self, url, result):
         self.urls[url] = dedent(result).encode()
 
@@ -103,7 +96,7 @@ def test_init_from_seq_file(self, tmp_path):
         assert fname.read_text() == '453'
 
 
-    def test_init_date_with_cookie(self, capsys, tmp_path, mock_urllib):
+    def test_init_date_with_cookie(self, capsys, tmp_path, mock_requests):
         self.url('https://planet.osm.org/replication/minute//state.txt',
                  """\
                     sequenceNumber=100

tools/pyosmium-get-changes

@@ -27,18 +27,14 @@ cookies to the server and will save received cookies to the jar file.
 
 from argparse import ArgumentParser, RawDescriptionHelpFormatter, ArgumentTypeError
 import datetime as dt
-import socket
+import http.cookiejar
+
 from osmium.replication import server as rserv
 from osmium.replication import newest_change_from_file
 from osmium.replication.utils import get_replication_header
 from osmium.version import pyosmium_release
 from osmium import SimpleHandler, WriteHandler
 
-try:
-    import http.cookiejar as cookiejarlib
-except ImportError:
-    import cookielib as cookiejarlib
-import urllib.request as urlrequest
 
 import re
 import sys
@@ -211,15 +207,13 @@ def main(args):
             or 'https://planet.osm.org/replication/minute/'
     logging.info("Using replication server at %s" % url)
 
-    socket.setdefaulttimeout(options.socket_timeout)
-
     with rserv.ReplicationServer(url) as svr:
+        svr.set_request_parameter('timeout', options.socket_timeout or None)
+
         if options.cookie is not None:
-            # According to the documentation, the cookie jar loads the file only if FileCookieJar.load is called.
-            cookie_jar = cookiejarlib.MozillaCookieJar(options.cookie)
+            cookie_jar = http.cookiejar.MozillaCookieJar(options.cookie)
             cookie_jar.load(options.cookie)
-            opener = urlrequest.build_opener(urlrequest.HTTPCookieProcessor(cookie_jar))
-            svr.open_url = opener.open
+            svr.set_request_parameter('cookies', cookie_jar)
 
         startseq = options.start.get_sequence(svr)
         if startseq is None:
@@ -245,6 +239,7 @@ def main(args):
         cookie_jar.save(options.cookie)
 
     if endseq is None:
+        log.error("Error while downloading diffs.")
         return 3
 
     if options.outfile != '-' or options.seq_file is not None:

tools/pyosmium-up-to-date

@@ -35,7 +35,7 @@ import re
 import sys
 import traceback
 import logging
-import socket
+import http.cookiejar
 
 from argparse import ArgumentParser, RawDescriptionHelpFormatter
 import datetime as dt
@@ -47,15 +47,6 @@ from textwrap import dedent as msgfmt
 from tempfile import mktemp
 import os.path
 
-try:
-    import http.cookiejar as cookiejarlib
-except ImportError:
-    import cookielib as cookiejarlib
-try:
-    import urllib.request as urlrequest
-except ImportError:
-    import urllib2 as urlrequest
-
 log = logging.getLogger()
 
 def update_from_osm_server(ts, options):
@@ -71,14 +62,14 @@ def update_from_osm_server(ts, options):
 def update_from_custom_server(url, seq, ts, options):
     """Update from a custom URL, simply using the diff sequence as is."""
     with rserv.ReplicationServer(url, "osc.gz") as svr:
+        log.info("Using replication service at %s", url)
+
+        svr.set_request_parameter('timeout', options.socket_timeout or None)
+
         if options.cookie is not None:
-            # According to the documentation, the cookie jar loads the file only if FileCookieJar.load is called.
-            cookie_jar = cookiejarlib.MozillaCookieJar(options.cookie)
+            cookie_jar = http.cookiejar.MozillaCookieJar(options.cookie)
             cookie_jar.load(options.cookie)
-            opener = urlrequest.build_opener(urlrequest.HTTPCookieProcessor(cookie_jar))
-            svr.open_url = opener.open
-
-        log.info("Using replication service at %s", url)
+            svr.set_request_parameter('cookies', cookie_jar)
 
         current = svr.get_state_info()
         if current is None:
@@ -244,8 +235,6 @@ if __name__ == '__main__':
     options = get_arg_parser(from_main=True).parse_args()
     log.setLevel(max(3 - options.loglevel, 0) * 10)
 
-    socket.setdefaulttimeout(options.socket_timeout)
-
     try:
         url, seq, ts = compute_start_point(options)
     except RuntimeError as e: