Merge pull request #11 from singlestore-labs/enable_multi_statements

Add connection option to enable multi statements

Author: kesmit13

.gitignore

@@ -76,3 +76,6 @@ docs/src/_build/doctrees
 docs/src/generated
 
 dask-worker-space
+
+## Test configuration
+singlestoredb/mysql/tests/databases.json

singlestoredb/config.py

@@ -90,6 +90,12 @@
     environ='SINGLESTOREDB_LOCAL_INFILE',
 )
 
+register_option(
+    'multi_statements', 'bool', check_bool, False,
+    'Should it be possible use multiple statements in one query?',
+    environ='SINGLESTOREDB_MULTI_STATEMENTS',
+)
+
 register_option(
     'ssl_key', 'str', check_str, None,
     'File containing SSL key',

singlestoredb/connection.py

@@ -1276,6 +1276,7 @@ def connect(
     results_format: Optional[str] = None,
     program_name: Optional[str] = None,
     conn_attrs: Optional[Dict[str, str]] = None,
+    multi_statements: Optional[bool] = None,
 ) -> Connection:
     """
     Return a SingleStoreDB connection.

singlestoredb/mysql/connection.py

@@ -292,6 +292,7 @@ def __init__(  # noqa: C901
         db=None,  # deprecated
         driver=None,  # internal use
         conn_attrs=None,
+        multi_statements=None,
     ):
         BaseConnection.__init__(**dict(locals()))
 
@@ -316,6 +317,8 @@ def __init__(  # noqa: C901
         self._local_infile = bool(local_infile)
         if self._local_infile:
             client_flag |= CLIENT.LOCAL_FILES
+        if multi_statements:
+            client_flag |= CLIENT.MULTI_STATEMENTS
 
         if read_default_group and not read_default_file:
             if sys.platform.startswith('win'):

singlestoredb/tests/test_connection.py

@@ -1357,6 +1357,19 @@ def test_results_format(self):
                     out = cur.fetchall()
                     assert type(out[0]) is dict, type(out)
 
+    def test_multi_statements(self):
+        if self.conn.driver not in ['http', 'https']:
+            with s2.connect(database=type(self).dbname, multi_statements=True) as conn:
+                with conn.cursor() as cur:
+                    cur.execute('SELECT 1; SELECT 2;')
+                    self.assertEqual([(1,)], list(cur))
+
+                    r = cur.nextset()
+                    self.assertTrue(r)
+
+                    self.assertEqual([(2,)], list(cur))
+                    self.assertIsNone(cur.nextset())
+
     def test_show_accessors(self):
         out = self.conn.show.columns('data')
         assert out.columns == [