From 1f07a6a7a251fa7cf0616f4cc7161865b1fbdbcb Mon Sep 17 00:00:00 2001
From: Andreas Motl <andreas.motl@crate.io>
Date: Fri, 12 Jan 2024 04:17:00 +0100
Subject: [PATCH] Stopped decoding the `timeout` parameter strictly as scalar

Instead, also permit `urllib3.Timeout` instances.
---
 CHANGES.txt                         |  3 +++
 docs/by-example/client.rst          | 15 ++++++++++++++-
 docs/by-example/http.rst            | 15 +++++++++++++--
 src/crate/client/http.py            |  4 +++-
 src/crate/client/test_connection.py | 24 ++++++++++++++++++++++++
 5 files changed, 57 insertions(+), 4 deletions(-)

diff --git a/CHANGES.txt b/CHANGES.txt
index 3ccfd634..35f87870 100644
--- a/CHANGES.txt
+++ b/CHANGES.txt
@@ -5,6 +5,9 @@ Changes for crate
 Unreleased
 ==========
 
+- Stopped decoding the ``timeout`` parameter strictly as scalar to also
+  permit ``urllib3.Timeout`` instances
+
 
 2023/09/29 0.34.0
 =================
diff --git a/docs/by-example/client.rst b/docs/by-example/client.rst
index e053d73f..6e8f08df 100644
--- a/docs/by-example/client.rst
+++ b/docs/by-example/client.rst
@@ -48,12 +48,25 @@ traceback if a server error occurs:
     >>> connection = client.connect([crate_host], error_trace=True)
     >>> connection.close()
 
+Network Timeouts
+----------------
+
 It's possible to define a default timeout value in seconds for all servers
-using the optional parameter ``timeout``:
+using the optional parameter ``timeout``. In this case, it will serve as a
+total timeout (connect and read):
 
     >>> connection = client.connect([crate_host, invalid_host], timeout=5)
     >>> connection.close()
 
+If you want to adjust the connect- vs. read-timeout values individually,
+please use the ``urllib3.Timeout`` object like:
+
+    >>> import urllib3
+    >>> connection = client.connect(
+    ...     [crate_host, invalid_host],
+    ...     timeout=urllib3.Timeout(connect=5, read=None))
+    >>> connection.close()
+
 Authentication
 --------------
 
diff --git a/docs/by-example/http.rst b/docs/by-example/http.rst
index 494e7b65..07486dee 100644
--- a/docs/by-example/http.rst
+++ b/docs/by-example/http.rst
@@ -199,8 +199,8 @@ timeout exception:
     {...}
     >>> http_client.close()
 
-It's possible to define a HTTP timeout in seconds on client instantiation, so
-an exception is raised when the timeout is reached:
+It is possible to define a HTTP timeout in seconds when creating a client
+object, so an exception is raised when the timeout expired:
 
     >>> http_client = HttpClient(crate_host, timeout=0.01)
     >>> http_client.sql('select fib(32)')
@@ -209,6 +209,17 @@ an exception is raised when the timeout is reached:
     crate.client.exceptions.ConnectionError: No more Servers available, exception from last server: ...
     >>> http_client.close()
 
+In order to adjust the connect- vs. read-timeout values individually,
+please use the ``urllib3.Timeout`` object like:
+
+    >>> import urllib3
+    >>> http_client = HttpClient(crate_host, timeout=urllib3.Timeout(connect=0.01, read=None))
+    >>> http_client.sql('select fib(32)')
+    Traceback (most recent call last):
+    ...
+    crate.client.exceptions.ConnectionError: No more Servers available, exception from last server: ...
+    >>> http_client.close()
+
 When connecting to non-CrateDB servers, the HttpClient will raise a ConnectionError like this:
 
     >>> http_client = HttpClient(["https://example.org/"])
diff --git a/src/crate/client/http.py b/src/crate/client/http.py
index a22a1ff0..78e0e594 100644
--- a/src/crate/client/http.py
+++ b/src/crate/client/http.py
@@ -273,7 +273,9 @@ def _pool_kw_args(verify_ssl_cert, ca_cert, client_cert, client_key,
         'key_file': client_key,
     }
     if timeout is not None:
-        kw['timeout'] = float(timeout)
+        if isinstance(timeout, str):
+            timeout = float(timeout)
+        kw['timeout'] = timeout
     if pool_size is not None:
         kw['maxsize'] = int(pool_size)
     return kw
diff --git a/src/crate/client/test_connection.py b/src/crate/client/test_connection.py
index 3b5c294c..e2b0bd1e 100644
--- a/src/crate/client/test_connection.py
+++ b/src/crate/client/test_connection.py
@@ -1,5 +1,7 @@
 import datetime
 
+from urllib3 import Timeout
+
 from .connection import Connection
 from .http import Client
 from crate.client import connect
@@ -72,3 +74,25 @@ def test_with_timezone(self):
         cursor = connection.cursor()
         self.assertEqual(cursor.time_zone.tzname(None), "UTC")
         self.assertEqual(cursor.time_zone.utcoffset(None), datetime.timedelta(0))
+
+    def test_timeout_float(self):
+        """
+        Verify setting the timeout value as a scalar (float) works.
+        """
+        with connect('localhost:4200', timeout=2.42) as conn:
+            self.assertEqual(conn.client._pool_kw["timeout"], 2.42)
+
+    def test_timeout_string(self):
+        """
+        Verify setting the timeout value as a scalar (string) works.
+        """
+        with connect('localhost:4200', timeout="2.42") as conn:
+            self.assertEqual(conn.client._pool_kw["timeout"], 2.42)
+
+    def test_timeout_object(self):
+        """
+        Verify setting the timeout value as a Timeout object works.
+        """
+        timeout = Timeout(connect=2.42, read=999)
+        with connect('localhost:4200', timeout=timeout) as conn:
+            self.assertEqual(conn.client._pool_kw["timeout"], timeout)