feat: support cookie

Author: youngsofun

databend_py/connection.py

@@ -4,7 +4,9 @@
 import time
 import uuid
 
+from http.cookiejar import Cookie
 from requests.auth import HTTPBasicAuth
+from requests.cookies import RequestsCookieJar
 
 import environs
 import requests
@@ -75,6 +77,17 @@ def get_error(response):
     return ServerException(response["error"]["message"], response["error"]["code"])
 
 
+class GlobalCookieJar(RequestsCookieJar):
+
+    def __init__(self):
+        super().__init__()
+
+    def set_cookie(self, cookie: Cookie, *args, **kwargs):
+        cookie.domain = ""
+        cookie.path = "/"
+        super().set_cookie(cookie, *args, **kwargs)
+
+
 class Connection(object):
     # Databend http handler doc: https://databend.rs/doc/reference/api/rest
 
@@ -120,6 +133,10 @@ def __init__(
         self.context = Context()
         self.requests_session = requests.Session()
         self.schema = "http"
+        cookie_jar = GlobalCookieJar()
+        cookie_jar.set("cookie_enabled", "true")
+        self.requests_session.cookies = cookie_jar
+        self.schema = 'http'
         if self.secure:
             self.schema = "https"
         e = environs.Env()

tests/test_client.py

@@ -293,6 +293,14 @@ def test_cast_bool(self):
         _, data = client.execute("select 'False'::boolean union select 'True'::boolean")
         self.assertEqual(data, [(True,), (False,)])
 
+    def test_temp_table(self):
+        client = Client.from_url(self.databend_url)
+        client.execute("create temp table t1(a int)")
+        client.execute("insert into t1 values (1)")
+        _, data = client.execute("select * from t1")
+        self.assertEqual(data, [(1,)])
+        client.execute("drop table t1")
+
 
 if __name__ == "__main__":
     unittest.main()

tests/test_simple.py

@@ -16,15 +16,15 @@ def __setattr__(self, key, value):
 
 
 class TestDict(unittest.TestCase):
-    databend_url = None  # 使用类属性来存储 databend_url
+    databend_url = None
 
     @classmethod
     def setUpClass(cls):
-        cls.databend_url = "test_url"  # 在类级别设置 databend_url
+        cls.databend_url = "test_url"
 
     def test_init(self):
         d = Dict(a=1, b="test")
-        self.assertEqual(self.databend_url, "test_url")  # 使用类属性
+        self.assertEqual(self.databend_url, "test_url")
         self.assertEqual(d.a, 1)
         self.assertEqual(d.b, "test")
         self.assertTrue(isinstance(d, dict))