diff --git a/pinotdb/db.py b/pinotdb/db.py index 31cbb3f..4a18ea7 100644 --- a/pinotdb/db.py +++ b/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=float(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=float(self._kwargs.get('timeout')) if self._kwargs.get('timeout') else None) self._kwargs['session'] = self.session cursor = AsyncCursor(*self._args, **self._kwargs) @@ -286,7 +286,7 @@ def __init__( # TODO: Remove this unused parameter when we can afford to break the # interface (e.g. new minor version). verify_ssl=True, - timeout=10, + timeout=10.0, extra_request_headers="", debug=False, preserve_types=False, diff --git a/pinotdb/sqlalchemy.py b/pinotdb/sqlalchemy.py index 0487beb..a5786d1 100644 --- a/pinotdb/sqlalchemy.py +++ b/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 = float(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": float(self._timeout) if self._timeout else 10.0, } if self.engine_type == "multi_stage": kwargs.update({"use_multistage_engine": True}) diff --git a/tests/unit/test_db.py b/tests/unit/test_db.py index a05b610..28c8386 100644 --- a/tests/unit/test_db.py +++ b/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() diff --git a/tests/unit/test_sqlalchemy.py b/tests/unit/test_sqlalchemy.py index 1809393..de346fa 100644 --- a/tests/unit/test_sqlalchemy.py +++ b/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')