Merge pull request #14 from databendcloud/feat/support-session

feat: support session settings and copy options

Author: hantmac

.coveragerc

@@ -0,0 +1,3 @@
+[run]
+plugins = Cython.Coverage
+source = databend_py

README.md

@@ -26,6 +26,7 @@ Use the next code to check connection:
 >     user="user",
 >     port="443",
 >     password="password")
+>     settings={"copy_purge":True,"force":True}
 > >>> print(client.execute("SELECT 1"))
 > ```
 
@@ -37,7 +38,7 @@ Pure Client example:
 > ``` python
 > >>> from databend_py import Client
 > >>>
-> >>> client = Client.from_url('http://root@localhost:8000/db')
+> >>> client = Client.from_url('http://root@localhost:8000/db?secure=False&copy_purge=True')
 > >>>
 > >>> client.execute('SHOW TABLES')
 > [('test',)]

databend_py/client.py

@@ -47,7 +47,7 @@ def receive_data(self, next_uri: str):
         return raw_data
 
     def receive_result(self, query, query_id=None, with_column_types=False):
-        raw_data = self.connection.query(query, None)
+        raw_data = self.connection.query(query)
         helper = self.helper()
         helper.response = raw_data
         helper.check_error()
@@ -57,7 +57,7 @@ def receive_result(self, query, query_id=None, with_column_types=False):
         return result.get_result()
 
     def iter_receive_result(self, query, query_id=None, with_column_types=False):
-        raw_data = self.connection.query(query, None)
+        raw_data = self.connection.query(query)
         helper = self.helper()
         helper.response = raw_data
         helper.check_error()
@@ -168,7 +168,7 @@ def from_url(cls, url):
 
         For example::
 
-            http://[user:password]@localhost:8000/default
+            https://[user:password]@localhost:8000/default?secure=True
             http://[user:password]@localhost:8000/default
             databend://[user:password]@localhost:8000/default
 
@@ -195,10 +195,12 @@ def from_url(cls, url):
                 kwargs[name] = value
             elif name == 'secure':
                 kwargs[name] = asbool(value)
+            elif name == 'copy_purge':
+                kwargs[name] = asbool(value)
             elif name in timeouts:
                 kwargs[name] = float(value)
             else:
-                settings[name] = value
+                settings[name] = value  # settings={'copy_purge':False}
         secure = kwargs.get("secure", False)
         kwargs['secure'] = secure
 
@@ -246,6 +248,35 @@ def stage_csv_file(self, filename, data):
     def sync_csv_file_into_table(self, filename, data, table):
         start = time.time()
         stage_path = self.stage_csv_file(filename, data)
-        _, _ = self.execute("COPY INTO %s FROM %s FILE_FORMAT = (type = CSV)" % (table, stage_path))
+        copy_options = self.generate_copy_options()
+        print(copy_options)
+        _, _ = self.execute(
+            f"COPY INTO {table} FROM {stage_path} FILE_FORMAT = (type = CSV)\
+             PURGE = {copy_options['PURGE']} FORCE = {copy_options['FORCE']}\
+              SIZE_LIMIT={copy_options['SIZE_LIMIT']} ON_ERROR = {copy_options['ON_ERROR']}")
         print("sync %s duration:%ss" % (filename, int(time.time() - start)))
         os.remove(filename)
+
+    def generate_copy_options(self):
+        # copy options docs: https://databend.rs/doc/sql-commands/dml/dml-copy-into-table#copyoptions
+        copy_options = {}
+        if "copy_purge" in self.settings:
+            copy_options["PURGE"] = self.settings["copy_purge"]
+        else:
+            copy_options["PURGE"] = False
+
+        if "force" in self.settings:
+            copy_options["FORCE"] = self.settings["force"]
+        else:
+            copy_options["FORCE"] = False
+
+        if "size_limit" in self.settings:
+            copy_options["SIZE_LIMIT"] = self.settings["size_limit"]
+        else:
+            copy_options["SIZE_LIMIT"] = 0
+        if "on_error" in self.settings:
+            copy_options["ON_ERROR"] = self.settings["on_error"]
+
+        else:
+            copy_options["ON_ERROR"] = "abort"
+        return copy_options

databend_py/connection.py

@@ -67,15 +67,16 @@ class Connection(object):
     #   'database': 'default'
     # }
     def __init__(self, host, port=None, user=defines.DEFAULT_USER, password=defines.DEFAULT_PASSWORD,
-                 database=defines.DEFAULT_DATABASE, secure=False, ):
+                 database=defines.DEFAULT_DATABASE, secure=False, copy_purge=False, session_settings=None):
         self.host = host
         self.port = port
         self.user = user
         self.password = password
         self.database = database
         self.secure = secure
+        self.copy_purge = copy_purge
         self.session_max_idle_time = defines.DEFAULT_SESSION_IDLE_TIME
-        self.session = {}
+        self.client_session = session_settings
         self.additional_headers = dict()
         self.query_option = None
         self.context = Context()
@@ -102,26 +103,27 @@ def get_description(self):
         return '{}:{}'.format(self.host, self.port)
 
     def disconnect(self):
-        self.session = {}
+        self.client_session = dict()
 
-    def query(self, statement, session):
+    def query(self, statement):
         url = self.format_url()
         log.logger.debug(f"http sql: {statement}")
         query_sql = {'sql': statement, "string_fields": True}
-        if session is not None:
-            query_sql['session'] = session
+        if self.client_session is not None and len(self.client_session) != 0:
+            query_sql['session'] = self.client_session
         else:
-            session = {"database": self.database}
-            query_sql['session'] = session
+            self.client_session = {"db": self.database}
+            query_sql['session'] = self.client_session
         log.logger.debug(f"http headers {self.make_headers()}")
         response = requests.post(url,
                                  data=json.dumps(query_sql),
                                  headers=self.make_headers(),
                                  auth=HTTPBasicAuth(self.user, self.password),
                                  verify=True)
-
         try:
-            return json.loads(response.content)
+            resp_dict = json.loads(response.content)
+            self.client_session = resp_dict["session"]
+            return resp_dict
         except Exception as err:
             log.logger.error(
                 f"http error on {url}, SQL: {statement} content: {response.content} error msg:{str(err)}"
@@ -136,32 +138,32 @@ def format_url(self):
         return f"{self.schema}://{self.host}:{self.port}/v1/query/"
 
     def reset_session(self):
-        self.session = {}
+        self.client_session = dict()
 
     def next_page(self, next_uri):
         url = "{}://{}:{}{}".format(self.schema, self.host, self.port, next_uri)
         return requests.get(url=url, headers=self.make_headers())
 
     # return a list of response util empty next_uri
     def query_with_session(self, statement):
-        current_session = self.session
+        current_session = self.client_session
         response_list = list()
-        response = self.query(statement, current_session)
+        response = self.query(statement)
         log.logger.debug(f"response content: {response}")
         response_list.append(response)
         start_time = time.time()
         time_limit = 12
         session = response['session']
         if session:
-            self.session = session
+            self.client_session = session
         while response['next_uri'] is not None:
             resp = self.next_page(response['next_uri'])
             response = json.loads(resp.content)
             log.logger.debug(f"Sql in progress, fetch next_uri content: {response}")
             self.check_error(response)
             session = response['session']
             if session:
-                self.session = session
+                self.client_session = session
             response_list.append(response)
             if time.time() - start_time > time_limit:
                 log.logger.warning(

tests/test_client.py

@@ -12,11 +12,12 @@ def assertHostsEqual(self, client, another, msg=None):
         self.assertEqual(client.connection.host, another, msg=msg)
 
     def test_simple(self):
-        c = Client.from_url('https://app.databend.com:443')
+        c = Client.from_url('https://app.databend.com:443?copy_purge=True')
 
         self.assertHostsEqual(c, 'app.databend.com')
         self.assertEqual(c.connection.database, 'default')
         self.assertEqual(c.connection.user, 'root')
+        self.assertEqual(c.connection.copy_purge, True)
 
         c = Client.from_url('https://host:443/db')
 
@@ -31,6 +32,11 @@ def test_simple(self):
         c = Client.from_url("databend://root:root@localhost:8000/default?secure=false")
         self.assertEqual(c.connection.schema, "http")
 
+    def test_session_settings(self):
+        session_settings = {"db": "database"}
+        c = Client(host="localhost", port=8000, user="root", password="root", session_settings={"db": "database"})
+        self.assertEqual(c.connection.client_session, session_settings)
+
     def test_ordinary_query(self):
         select_test = '''
         select
@@ -55,8 +61,8 @@ def test_ordinary_query(self):
         self.assertEqual(r, ([('1', 'UInt8')], [(1,)]))
 
     def test_batch_insert(self):
-        c = Client.from_url(self.databend_url)
-
+        # with copy on purge
+        c = Client(host="localhost", port=8000, user="root", password="root", settings={"copy_purge": True})
         c.execute('DROP TABLE IF EXISTS test')
         c.execute('CREATE TABLE if not exists test (x Int32,y VARCHAR)')
         c.execute('DESC  test')