Fix type error when timeout is provided in connection string

pinotdb/db.py

@@ -181,7 +181,7 @@ def cursor(self):
         if not self.session or self.session.is_closed:
             self.session = httpx.Client(
                 verify=self._kwargs.get('verify_ssl'),
-                timeout=self._kwargs.get('timeout'))
+                timeout=int(self._kwargs.get('timeout')) if self._kwargs.get('timeout') else None)
 
         self._kwargs['session'] = self.session
         cursor = Cursor(*self._args, **self._kwargs)
@@ -213,7 +213,7 @@ def cursor(self):
         if not self.session or self.session.is_closed:
             self.session = httpx.AsyncClient(
                 verify=self._kwargs.get('verify_ssl'),
-                timeout=self._kwargs.get('timeout'))
+                timeout=int(self._kwargs.get('timeout')) if self._kwargs.get('timeout') else None)
 
         self._kwargs['session'] = self.session
         cursor = AsyncCursor(*self._args, **self._kwargs)

pinotdb/sqlalchemy.py

@@ -182,7 +182,7 @@ def update_from_kwargs(self, givenkw):
             kwargs["database"] = self._database = kwargs.pop("database")
         kwargs["debug"] = self._debug = bool(kwargs.get("debug", False))
         kwargs["verify_ssl"] = self._verify_ssl = (str(kwargs.get("verify_ssl", "true")).lower() in ['true'])
-        kwargs["timeout"] = self._timeout = kwargs.get("timeout", None)
+        kwargs["timeout"] = self._timeout = int(kwargs.get('timeout')) if kwargs.get('timeout') else None
         logger.info(
             "Updated pinot dialect args from %s: %s and %s",
             dict(map(lambda kv: (kv[0], mask_value(kv[0], kv[1], ['password'])), kwargs.items())),
@@ -209,7 +209,7 @@ def create_connect_args(self, url):
             "username": url.username,
             "password": url.password,
             "verify_ssl": self._verify_ssl or True,
-            "timeout": self._timeout or 10.0,
+            "timeout": int(self._timeout) if self._timeout else 10.0,
         }
         if self.engine_type == "multi_stage":
             kwargs.update({"use_multistage_engine": True})

tests/unit/test_db.py

@@ -41,6 +41,17 @@ def test_verifies_httpx_session(self):
         connection.verify_session()
         # All good, no errors.
 
+    def test_verifies_timeout_kwargs(self):
+        connection = db.Connection(host='localhost', timeout=100)
+        cursor = connection.cursor()
+
+        self.assertDictEqual(
+            connection.session.timeout.as_dict(), 
+            {'connect': 100.0, 'read': 100.0, 'write': 100.0, 'pool': 100.0}
+            )
+
+
+
     def test_fails_to_verify_session_if_unexpected_type(self):
         connection = db.Connection()
         connection.session = object()

tests/unit/test_sqlalchemy.py

@@ -59,6 +59,26 @@ def test_creates_connection_args(self):
             'timeout': 10.0,
         })
 
+    def test_checks_timeout_in_query_params(self):
+        url = make_url(
+            'pinot://localhost:8000/query/sql?controller='
+            'http://localhost:9000/&&verify_ssl=True&&timeout=100')
+
+        cargs, cparams = self.dialect.create_connect_args(url)
+
+        self.assertEqual(cargs, [])
+        self.assertEqual(cparams, {
+            'debug': False,
+            'scheme': 'http',
+            'host': 'localhost',
+            'port': 8000,
+            'path': 'query/sql',
+            'username': None,
+            'password': None,
+            'verify_ssl': True,
+            'timeout': 100,
+        })
+
     def test_creates_connection_args_without_query(self):
         url = make_url('pinot://localhost:8000/query/sql')