Merge branch 'main' into chore/upgrade-052

databendlabs · Dec 19, 2023 · 1f32b6c · 1f32b6c
2 parents 9a21642 + 98c0e54
commit 1f32b6c
Show file tree

Hide file tree

Showing 4 changed files with 8 additions and 3 deletions.
diff --git a/databend_py/client.py b/databend_py/client.py
@@ -187,6 +187,7 @@ def from_url(cls, url):
 
             timeouts = {
                 'connect_timeout',
+                'read_timeout',
                 'send_receive_timeout',
                 'sync_request_timeout'
             }

diff --git a/databend_py/connection.py b/databend_py/connection.py
@@ -73,7 +73,7 @@ class Connection(object):
     #   'database': 'default'
     # }
     def __init__(self, host, port=None, user=defines.DEFAULT_USER, password=defines.DEFAULT_PASSWORD,
-                 connect_timeout=defines.DEFAULT_CONNECT_TIMEOUT,
+                 connect_timeout=defines.DEFAULT_CONNECT_TIMEOUT, read_timeout=defines.DEFAULT_READ_TIMEOUT,
                  database=defines.DEFAULT_DATABASE, secure=False, copy_purge=False, session_settings=None,
                  persist_cookies=False):
         self.host = host
@@ -82,6 +82,7 @@ def __init__(self, host, port=None, user=defines.DEFAULT_USER, password=defines.
         self.password = password
         self.database = database
         self.connect_timeout = connect_timeout
+        self.read_timeout = read_timeout
         self.secure = secure
         self.copy_purge = copy_purge
         self.session_max_idle_time = defines.DEFAULT_SESSION_IDLE_TIME
@@ -126,7 +127,7 @@ def do_query(self, url, query_sql):
                                               data=json.dumps(query_sql),
                                               headers=self.make_headers(),
                                               auth=HTTPBasicAuth(self.user, self.password),
-                                              timeout=self.connect_timeout,
+                                              timeout=(self.connect_timeout, self.read_timeout),
                                               verify=True)
         try:
             resp_dict = json.loads(response.content)

diff --git a/databend_py/defines.py b/databend_py/defines.py
@@ -3,6 +3,7 @@
 DEFAULT_PASSWORD = ''
 DEFAULT_SESSION_IDLE_TIME = 30
 DEFAULT_CONNECT_TIMEOUT = 20
+DEFAULT_READ_TIMEOUT = 20
 
 DBMS_NAME = 'Databend'
 CLIENT_NAME = 'databend-py'

diff --git a/tests/test_client.py b/tests/test_client.py
@@ -47,9 +47,11 @@ def test_simple(self):
         c = Client.from_url("databend://root:root@localhost:8000/default?compress=True")
         self.assertEqual(c._uploader._compress, True)
         self.assertEqual(c.connection.connect_timeout, 20)
+        self.assertEqual(c.connection.read_timeout, 20)
 
-        c = Client.from_url("databend://root:root@localhost:8000/default?connect_timeout=30")
+        c = Client.from_url("databend://root:root@localhost:8000/default?connect_timeout=30&read_timeout=30")
         self.assertEqual(c.connection.connect_timeout, 30)
+        self.assertEqual(c.connection.read_timeout, 30)
 
         self.assertEqual(c.connection.persist_cookies, False)
         c = Client.from_url('https://root:root@localhost:8000?persist_cookies=True')