SQLAlchemy: Fix handling URL parameters timeout and pool_size

amotl · amotl · commit d01167393390 · 2023-09-05T06:21:54.000-07:00
diff --git a/CHANGES.txt b/CHANGES.txt
@@ -6,6 +6,8 @@ Unreleased
 ==========
 
 - Properly handle Python-native UUID types in SQL parameters
+- SQLAlchemy: Fix handling URL parameters ``timeout`` and ``pool_size``
+
 
 2023/07/17 0.33.0
 =================
diff --git a/docs/by-example/sqlalchemy/getting-started.rst b/docs/by-example/sqlalchemy/getting-started.rst
@@ -46,8 +46,8 @@ Create an SQLAlchemy :doc:`Session <sa:orm/session_basics>`:
     >>> Base = declarative_base()
 
 
-Connection string
-=================
+Connect
+=======
 
 In SQLAlchemy, a connection is established using the ``create_engine`` function.
 This function takes a connection string, actually an `URL`_, that varies from
@@ -65,7 +65,9 @@ to a different server the following syntax can be used:
     >>> sa.create_engine('crate://otherserver:4200')
     Engine(crate://otherserver:4200)
 
-Since CrateDB is a clustered database running on multiple servers, it is
+Multiple Hosts
+--------------
+Because CrateDB is a clustered database running on multiple servers, it is
 recommended to connect to all of them. This enables the DB-API layer to
 use round-robin to distribute the load and skip a server if it becomes
 unavailable. In order to make the driver aware of multiple servers, use
@@ -76,6 +78,8 @@ the ``connect_args`` parameter like so:
     ... })
     Engine(crate://)
 
+TLS Options
+-----------
 As defined in :ref:`https_connection`, the client validates SSL server
 certificates by default. To configure this further, use e.g. the ``ca_cert``
 attribute within the ``connect_args``, like:
@@ -96,6 +100,37 @@ In order to disable SSL verification, use ``verify_ssl_cert = False``, like:
     ...         'verify_ssl_cert': False,
     ...     })
 
+Timeout Options
+---------------
+In order to configure TCP timeout options, use the ``timeout`` parameter within
+``connect_args``,
+
+    >>> timeout_engine = sa.create_engine('crate://localhost/', connect_args={'timeout': 42.42})
+    >>> timeout_engine.raw_connection().driver_connection.client._pool_kw["timeout"]
+    42.42
+
+or use the ``timeout`` URL parameter within the database connection URL.
+
+    >>> timeout_engine = sa.create_engine('crate://localhost/?timeout=42.42')
+    >>> timeout_engine.raw_connection().driver_connection.client._pool_kw["timeout"]
+    42.42
+
+Pool Size
+---------
+
+In order to configure the database connection pool size, use the ``pool_size``
+parameter within ``connect_args``,
+
+    >>> timeout_engine = sa.create_engine('crate://localhost/', connect_args={'pool_size': 20})
+    >>> timeout_engine.raw_connection().driver_connection.client._pool_kw["maxsize"]
+    20
+
+or use the ``pool_size`` URL parameter within the database connection URL.
+
+    >>> timeout_engine = sa.create_engine('crate://localhost/?pool_size=20')
+    >>> timeout_engine.raw_connection().driver_connection.client._pool_kw["maxsize"]
+    20
+
 
 Basic DDL operations
 ====================
diff --git a/src/crate/client/http.py b/src/crate/client/http.py
@@ -255,10 +255,11 @@ def _pool_kw_args(verify_ssl_cert, ca_cert, client_cert, client_key,
         'cert_reqs': ssl.CERT_REQUIRED if verify_ssl_cert else ssl.CERT_NONE,
         'cert_file': client_cert,
         'key_file': client_key,
-        'timeout': timeout,
     }
+    if timeout is not None:
+        kw['timeout'] = float(timeout)
     if pool_size is not None:
-        kw['maxsize'] = pool_size
+        kw['maxsize'] = int(pool_size)
     return kw
 
 
diff --git a/src/crate/client/sqlalchemy/tests/connection_test.py b/src/crate/client/sqlalchemy/tests/connection_test.py
@@ -83,6 +83,22 @@ def test_connection_server_uri_https_with_credentials(self):
         conn.close()
         engine.dispose()
 
+    def test_connection_server_uri_parameter_timeout(self):
+        engine = sa.create_engine(
+            "crate://otherhost:19201/?timeout=42.42")
+        conn = engine.raw_connection()
+        self.assertEqual(conn.driver_connection.client._pool_kw["timeout"], 42.42)
+        conn.close()
+        engine.dispose()
+
+    def test_connection_server_uri_parameter_pool_size(self):
+        engine = sa.create_engine(
+            "crate://otherhost:19201/?pool_size=20")
+        conn = engine.raw_connection()
+        self.assertEqual(conn.driver_connection.client._pool_kw["maxsize"], 20)
+        conn.close()
+        engine.dispose()
+
     def test_connection_multiple_server_http(self):
         engine = sa.create_engine(
             "crate://", connect_args={
