singlestoredb 1.7.2__cp38-abi3-win_amd64.whl → 1.9.0__cp38-abi3-win_amd64.whl

singlestoredb/mysql/connection.py

@@ -226,6 +226,8 @@ class Connection(BaseConnection):
         Set to true to check the server certificate's validity.
     ssl_verify_identity : bool, optional
         Set to true to check the server's identity.
+    tls_sni_servername: str, optional
+        Set server host name for TLS connection
     read_default_group : str, optional
         Group to read from in the configuration file.
     autocommit : bool, optional
@@ -295,6 +297,7 @@ class Connection(BaseConnection):
     _auth_plugin_name = ''
     _closed = False
     _secure = False
+    _tls_sni_servername = None
 
     def __init__(  # noqa: C901
         self,
@@ -335,6 +338,7 @@ class Connection(BaseConnection):
         ssl_key=None,
         ssl_verify_cert=None,
         ssl_verify_identity=None,
+        tls_sni_servername=None,
         parse_json=True,
         invalid_values=None,
         pure_python=None,
@@ -638,6 +642,7 @@ class Connection(BaseConnection):
 
         self._is_committable = True
         self._in_sync = False
+        self._tls_sni_servername = tls_sni_servername
         self._track_env = bool(track_env) or self.host == 'singlestore.com'
         self._enable_extended_data_types = enable_extended_data_types
         if vector_data_format.lower() in ['json', 'binary']:
@@ -1364,7 +1369,10 @@ class Connection(BaseConnection):
         if self.ssl and self.server_capabilities & CLIENT.SSL:
             self.write_packet(data_init)
 
-            self._sock = self.ctx.wrap_socket(self._sock, server_hostname=self.host)
+            hostname = self.host
+            if self._tls_sni_servername:
+                hostname = self._tls_sni_servername
+            self._sock = self.ctx.wrap_socket(self._sock, server_hostname=hostname)
             self._rfile = self._sock.makefile('rb')
             self._secure = True
 

singlestoredb/tests/test.ipynb

@@ -0,0 +1,18 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Test Notebook"
+   ]
+  }
+ ],
+ "metadata": {
+  "language_info": {
+   "name": "python"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

singlestoredb/tests/test2.ipynb

@@ -0,0 +1,18 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Test Notebook 2"
+   ]
+  }
+ ],
+ "metadata": {
+  "language_info": {
+   "name": "python"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

singlestoredb/tests/test_fusion.py

@@ -465,10 +465,6 @@ class TestWorkspaceFusion(unittest.TestCase):
                 pass
 
 
-@unittest.skipIf(
-    os.environ.get('SINGLESTOREDB_FUSION_ENABLE_HIDDEN', '0') == '0',
-    'Hidden Fusion commands are not enabled.',
-)
 @pytest.mark.management
 class TestJobsFusion(unittest.TestCase):
 

singlestoredb/tests/test_management.py

@@ -740,7 +740,7 @@ class TestStage(unittest.TestCase):
             'rename_test_2/nest_1/nested_rename_test_3.sql', overwrite=True,
         )
 
-    def test_stage_object(self):
+    def test_file_object(self):
         st = self.wg.stage
 
         st.mkdir('obj_test')
@@ -1028,3 +1028,275 @@ class TestJob(unittest.TestCase):
         assert deleted
         job = job_manager.get(job.job_id)
         assert job.terminated_at is not None
+
+
+@pytest.mark.management
+class TestFileSpaces(unittest.TestCase):
+
+    manager = None
+    personal_space = None
+    shared_space = None
+
+    @classmethod
+    def setUpClass(cls):
+        cls.manager = s2.manage_files()
+        cls.personal_space = cls.manager.personal_space
+        cls.shared_space = cls.manager.shared_space
+
+    @classmethod
+    def tearDownClass(cls):
+        cls.manager = None
+        cls.personal_space = None
+        cls.shared_space = None
+
+    def test_upload_file(self):
+        for space in [self.personal_space, self.shared_space]:
+            root = space.info('/')
+            assert str(root.path) == '/'
+            assert root.type == 'directory'
+
+            # Upload files
+            f = space.upload_file(
+                TEST_DIR / 'test.ipynb',
+                'upload_test.ipynb',
+            )
+            assert str(f.path) == 'upload_test.ipynb'
+            assert f.type == 'notebook'
+
+            # Download and compare to original
+            txt = f.download(encoding='utf-8')
+            assert txt == open(TEST_DIR / 'test.ipynb').read()
+
+            # Make sure we can't overwrite
+            with self.assertRaises(OSError):
+                space.upload_file(
+                    TEST_DIR / 'test.ipynb',
+                    'upload_test.ipynb',
+                )
+
+            # Force overwrite with new content
+            f = space.upload_file(
+                TEST_DIR / 'test2.ipynb',
+                'upload_test.ipynb', overwrite=True,
+            )
+            assert str(f.path) == 'upload_test.ipynb'
+            assert f.type == 'notebook'
+
+            # Verify new content
+            txt = f.download(encoding='utf-8')
+            assert txt == open(TEST_DIR / 'test2.ipynb').read()
+
+            # Make sure we can't upload a folder
+            with self.assertRaises(s2.ManagementError):
+                space.upload_folder(TEST_DIR, 'test')
+
+            # Cleanup
+            space.remove('upload_test.ipynb')
+
+    def test_open(self):
+        for space in [self.personal_space, self.shared_space]:
+            # See if error is raised for non-existent file
+            with self.assertRaises(s2.ManagementError):
+                space.open('open_test.ipynb', 'r')
+
+            # Load test file
+            space.upload_file(TEST_DIR / 'test.ipynb', 'open_test.ipynb')
+
+            # Read file using `open`
+            with space.open('open_test.ipynb', 'r') as rfile:
+                assert rfile.read() == open(TEST_DIR / 'test.ipynb').read()
+
+            # Read file using `open` with 'rt' mode
+            with space.open('open_test.ipynb', 'rt') as rfile:
+                assert rfile.read() == open(TEST_DIR / 'test.ipynb').read()
+
+            # Read file using `open` with 'rb' mode
+            with space.open('open_test.ipynb', 'rb') as rfile:
+                assert rfile.read() == open(TEST_DIR / 'test.ipynb', 'rb').read()
+
+            # Read file using `open` with 'rb' mode
+            with self.assertRaises(ValueError):
+                with space.open('open_test.ipynb', 'b') as rfile:
+                    pass
+
+            # Attempt overwrite file using `open` with mode 'x'
+            with self.assertRaises(OSError):
+                with space.open('open_test.ipynb', 'x') as wfile:
+                    pass
+
+            # Attempt overwrite file using `open` with mode 'w'
+            with space.open('open_test.ipynb', 'w') as wfile:
+                wfile.write(open(TEST_DIR / 'test2.ipynb').read())
+
+            txt = space.download_file('open_test.ipynb', encoding='utf-8')
+
+            assert txt == open(TEST_DIR / 'test2.ipynb').read()
+
+            # Test writer without context manager
+            wfile = space.open('open_raw_test.ipynb', 'w')
+            for line in open(TEST_DIR / 'test.ipynb'):
+                wfile.write(line)
+            wfile.close()
+
+            txt = space.download_file(
+                'open_raw_test.ipynb',
+                encoding='utf-8',
+            )
+
+            assert txt == open(TEST_DIR / 'test.ipynb').read()
+
+            # Test reader without context manager
+            rfile = space.open('open_raw_test.ipynb', 'r')
+            txt = ''
+            for line in rfile:
+                txt += line
+            rfile.close()
+
+            assert txt == open(TEST_DIR / 'test.ipynb').read()
+
+            # Cleanup
+            space.remove('open_test.ipynb')
+            space.remove('open_raw_test.ipynb')
+
+    def test_obj_open(self):
+        for space in [self.personal_space, self.shared_space]:
+            # Load test file
+            f = space.upload_file(
+                TEST_DIR / 'test.ipynb',
+                'obj_open_test.ipynb',
+            )
+
+            # Read file using `open`
+            with f.open() as rfile:
+                assert rfile.read() == open(TEST_DIR / 'test.ipynb').read()
+
+            # Make sure directories error out
+            with self.assertRaises(s2.ManagementError):
+                space.mkdir('obj_open_dir')
+
+            # Write file using `open`
+            with f.open('w', encoding='utf-8') as wfile:
+                wfile.write(open(TEST_DIR / 'test2.ipynb').read())
+
+            assert f.download(encoding='utf-8') == open(TEST_DIR / 'test2.ipynb').read()
+
+            # Test writer without context manager
+            wfile = f.open('w')
+            for line in open(TEST_DIR / 'test.ipynb'):
+                wfile.write(line)
+            wfile.close()
+
+            txt = space.download_file(f.path, encoding='utf-8')
+
+            assert txt == open(TEST_DIR / 'test.ipynb').read()
+
+            # Test reader without context manager
+            rfile = f.open('r')
+            txt = ''
+            for line in rfile:
+                txt += line
+            rfile.close()
+
+            assert txt == open(TEST_DIR / 'test.ipynb').read()
+
+            # Cleanup
+            space.remove('obj_open_test.ipynb')
+
+    def test_os_directories(self):
+        for space in [self.personal_space, self.shared_space]:
+            # Make sure directories error out
+            with self.assertRaises(s2.ManagementError):
+                space.mkdir('mkdir_test_1')
+
+            with self.assertRaises(s2.ManagementError):
+                space.exists('mkdir_test_1/')
+
+            out = space.listdir('/')
+            assert 'mkdir_test_1/' not in out
+
+            with self.assertRaises(s2.ManagementError):
+                space.rmdir('mkdir_test_1/')
+
+    def test_os_rename(self):
+        for space in [self.personal_space, self.shared_space]:
+            space.upload_file(
+                TEST_DIR / 'test.ipynb',
+                'rename_test.ipynb',
+            )
+            assert 'rename_test.ipynb' in space.listdir('/')
+            assert 'rename_test_2.ipynb' not in space.listdir('/')
+
+            space.rename(
+                'rename_test.ipynb',
+                'rename_test_2.ipynb',
+            )
+            assert 'rename_test.ipynb' not in space.listdir('/')
+            assert 'rename_test_2.ipynb' in space.listdir('/')
+
+            # non-existent file
+            with self.assertRaises(OSError):
+                space.rename('rename_foo.ipynb', 'rename_foo_2.ipynb')
+
+            space.upload_file(
+                TEST_DIR / 'test.ipynb',
+                'rename_test_3.ipynb',
+            )
+
+            # overwrite
+            with self.assertRaises(OSError):
+                space.rename(
+                    'rename_test_2.ipynb',
+                    'rename_test_3.ipynb',
+                )
+
+            space.rename(
+                'rename_test_2.ipynb',
+                'rename_test_3.ipynb', overwrite=True,
+            )
+
+            # Cleanup
+            space.remove('rename_test_3.ipynb')
+
+    def test_file_object(self):
+        for space in [self.personal_space, self.shared_space]:
+            f = space.upload_file(
+                TEST_DIR / 'test.ipynb',
+                'obj_test.ipynb',
+            )
+
+            assert not f.is_dir()
+            assert f.is_file()
+
+            # abspath / basename / dirname / exists
+            assert f.abspath() == 'obj_test.ipynb'
+            assert f.basename() == 'obj_test.ipynb'
+            assert f.dirname() == '/'
+            assert f.exists()
+
+            # download
+            assert f.download(encoding='utf-8') == \
+                open(TEST_DIR / 'test.ipynb', 'r').read()
+            assert f.download() == open(TEST_DIR / 'test.ipynb', 'rb').read()
+
+            assert space.is_file('obj_test.ipynb')
+            f.remove()
+            assert not space.is_file('obj_test.ipynb')
+
+            # mtime / ctime
+            assert f.getmtime() > 0
+            assert f.getctime() > 0
+
+            # rename
+            f = space.upload_file(
+                TEST_DIR / 'test.ipynb',
+                'obj_test.ipynb',
+            )
+            assert space.exists('obj_test.ipynb')
+            assert not space.exists('obj_test_2.ipynb')
+            f.rename('obj_test_2.ipynb')
+            assert not space.exists('obj_test.ipynb')
+            assert space.exists('obj_test_2.ipynb')
+            assert f.abspath() == 'obj_test_2.ipynb'
+
+            # Cleanup
+            space.remove('obj_test_2.ipynb')

{singlestoredb-1.7.2.dist-info → singlestoredb-1.9.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: singlestoredb
-Version: 1.7.2
+Version: 1.9.0
 Summary: Interface to the SingleStoreDB database and workspace management APIs
 Home-page: https://github.com/singlestore-labs/singlestoredb-python
 Author: SingleStore

{singlestoredb-1.7.2.dist-info → singlestoredb-1.9.0.dist-info}/RECORD

@@ -1,10 +1,11 @@
-_singlestoredb_accel.pyd,sha256=74y_1TpqSVFQ-rTN7XE73b1dJ16kA-s1K57dV-NmkEY,59392
-singlestoredb/__init__.py,sha256=jSf2uwMGpvrTEf60dTZMxvO63yzzmKGE12bZUuCMOQA,1697
+_singlestoredb_accel.pyd,sha256=garYUCMiztjRATI49oXgFrKMrTJfMkYHoT9CsdRVVTw,59392
+singlestoredb/__init__.py,sha256=uPTzIZk0aLMH0g8uIDrCqusEFDP5oLCegTc5h6XwM3U,1711
 singlestoredb/auth.py,sha256=RmYiH0Wlc2RXc4pTlRMysxtBI445ggCIwojWKC_eDLE,7844
-singlestoredb/config.py,sha256=LlrwKor_23uA9u7jWBYb-IaOKs30CjWIM7o9xCXEPcc,12721
-singlestoredb/connection.py,sha256=QC5YQemwJOhdW_-ZBFpNLE15xFxwI1fB2LuvE1hNi9k,46812
+singlestoredb/config.py,sha256=n6ludREIoiZDEzSGmv0xouv_zHFnznKNKxSvjzgQ3Lk,12876
+singlestoredb/connection.py,sha256=I3A0VkA_E6pfUNCzIvj0KTb5zn6-htBAPsu6OVXO9-I,47150
 singlestoredb/converters.py,sha256=7_Of1f3Ow-JUoY1pHFlMPYxvt8llzdk-7X8qk5z2xJM,21631
 singlestoredb/exceptions.py,sha256=WCCJrNSsU-hD-621Jpd6bwmvGftQ7byXkk-XKXlaxpg,3354
+singlestoredb/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 singlestoredb/pytest.py,sha256=TH364xRCN7_QaN0oRQDHixrEcDx_ZBgu3bmY0tvKrYU,9357
 singlestoredb/types.py,sha256=Lv0BEQl6aSZBiAe0OSI07FEJhcHZ9HX45iT9NU_mxHQ,10334
 singlestoredb/ai/__init__.py,sha256=nT048t90xqjaNhz7KJ10KfSVW4RcZRoujyC6po6Nmb8,61
@@ -31,29 +32,33 @@ singlestoredb/functions/ext/rowdat_1.py,sha256=KYj_y5JWm3_B2-QC47HK-CNOrzujBqGUw
 singlestoredb/functions/ext/utils.py,sha256=OPMFD-tTCx2Kk9jguQkrTr7e4AgNkt15YsvaT1YSmN8,5480
 singlestoredb/fusion/__init__.py,sha256=FHWtrg6OJFTf6Ye197V5sU6ssryr2h6FBcDIgXP7-H4,367
 singlestoredb/fusion/graphql.py,sha256=SHqsPe4xgawdsTPHEtJGQlybYGWqPrGMmyK-v20RLac,5420
-singlestoredb/fusion/handler.py,sha256=iILH6aeVWOPbqbkJ8m8CVYV61uJnw8ohYIv8U7bXklM,26231
+singlestoredb/fusion/handler.py,sha256=dZ5goR9aseeGbSRTVZpxIYeLsraWE5VTNtA7nskt3rk,28507
 singlestoredb/fusion/registry.py,sha256=_eT1gd38VPlFKs5f9Pu6lqQyoDQ_ixW5O56QwYLQ89Y,6361
 singlestoredb/fusion/result.py,sha256=EcFY5Qv43ySlQsfl_JB-I3ko7PzVdjuhhoKN96uHSAM,12171
 singlestoredb/fusion/handlers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-singlestoredb/fusion/handlers/job.py,sha256=TE57c3IDsKb6C5QH4AKAZBMMeZCIaQ4NrZPzXIq0KaI,21959
-singlestoredb/fusion/handlers/stage.py,sha256=uPqawMvchnpyrPYLhB0joTremCURNYKOvYntFc3zTRU,14133
-singlestoredb/fusion/handlers/utils.py,sha256=7xWb_1mJzxW0po9iHVY2ZVnRvHIQgOlKZQZ1zllJBjk,5271
+singlestoredb/fusion/handlers/export.py,sha256=mY-lEbi3WZLEHJ91hAFPYcasP4g1VeQkP0xBj3casts,7045
+singlestoredb/fusion/handlers/files.py,sha256=sXO5OFS5QsmwNOnMY3G5xpMPqh4_WkssW3dOB6uI4VQ,19669
+singlestoredb/fusion/handlers/job.py,sha256=3enfxHwERH7T4u0FEwOPN0IL0GtepaCYgEsisiy3Df4,21753
+singlestoredb/fusion/handlers/stage.py,sha256=PP-SSP204lwpmnycSXXSmFPzoN535JVuwglDCbaQ8Lw,14789
+singlestoredb/fusion/handlers/utils.py,sha256=OffWWj8FkZgqr5ijY3WQ6cwaAB7a7Z5cYxLeOAfU_JE,10356
 singlestoredb/fusion/handlers/workspace.py,sha256=NxoEY5xd5lCQmXiim4nhAYCL0agHo1H_rGPpqa31hiw,28397
 singlestoredb/http/__init__.py,sha256=4cEDvLloGc3LSpU-PnIwacyu0n5oIIIE6xk2SPyWD_w,939
 singlestoredb/http/connection.py,sha256=LFUeWx7maS7xhQLqEX3pgvIGoosqyJTovtWwJ1DyDSA,40721
-singlestoredb/management/__init__.py,sha256=pQbsOffl-i6zIzR63MCxSjxPY6TNmy7Kg0BCTExt3mk,244
+singlestoredb/management/__init__.py,sha256=A66ZnFyX--PsAZ2tvtYUfIUBvVGDBFQsnVc6nGTlX60,277
 singlestoredb/management/billing_usage.py,sha256=0UHFSPCrN0nyeGFFM-HXS3NP8pYmYo2BCCahDEPXvzg,3883
 singlestoredb/management/cluster.py,sha256=XfdBuTlrAG-mnW1BFKeoAr4YSE5IVgxLjbuBSpqIySo,14823
+singlestoredb/management/export.py,sha256=Ksrb8_sxeqva4NElaGxdPQUWQga2yOEkf0wnw5M-iGc,4302
+singlestoredb/management/files.py,sha256=PmXEeJuIYthGJpjgTaGmb5eMD1lIfJSevOVnf8QDXYk,29204
 singlestoredb/management/job.py,sha256=Npfe1JLYJlggGBrXLniPKwKUKF1i3alvSY1SFtvauSs,25498
 singlestoredb/management/manager.py,sha256=uGNrUe3zhuP-HUqdfwvy4MdEXTCmq-FZKjIwZSc3hOM,9096
 singlestoredb/management/organization.py,sha256=JBsNC4R3boUKdYvyCZyfGoVMC1mD6SPuMI1UssBVoOM,5611
 singlestoredb/management/region.py,sha256=oGoLLS88dE1GmY7GCc0BV7X3f7bWwKQyeXOVBFmK9Pk,1678
 singlestoredb/management/utils.py,sha256=BP-Wb8Sg16GbdLI_DeBz-3ttMklz6ZjYyMOz-sxElz8,13594
-singlestoredb/management/workspace.py,sha256=BBNa3Af5IaUAwtzHWrNTJXoc8anXJ7FiWH9yIvtg79Q,63610
+singlestoredb/management/workspace.py,sha256=EKSmu7sFBgotcgk0zlQP1Ke6ZnfQxwl_4tY7hev1yO4,58018
 singlestoredb/mysql/__init__.py,sha256=CbpwzNUJPAmKPpIobC0-ugBta_RgHCMq7X7N75QLReY,4669
 singlestoredb/mysql/_auth.py,sha256=YaqqyvAHmeraBv3BM207rNveUVPM-mPnW20ts_ynVWg,8341
 singlestoredb/mysql/charset.py,sha256=mnCdMpvdub1S2mm2PSk2j5JddgsWRjsVLtGx-y9TskE,10724
-singlestoredb/mysql/connection.py,sha256=M7SH2QJbD6IG-E6UOzLuneGtXKmKbvjhvYn18EMnA_A,74379
+singlestoredb/mysql/connection.py,sha256=mmTF9X7cPbmTS7yPHHbkjwQGpOPSbpDgZ3MNd2MyXkA,74716
 singlestoredb/mysql/converters.py,sha256=vebFFm6IrC0WgY-5Eh-esaPizY5cq3vDOUlEKGaYM-U,7771
 singlestoredb/mysql/cursors.py,sha256=pkrP-1t8IhBJRnYpdM7Rdm-332nOq1RYTDJ_yg_q5HI,27682
 singlestoredb/mysql/err.py,sha256=aDbmfq08gWVmfgIea735wSeiFdvYbB5wusgd3qTVq1s,2480
@@ -97,7 +102,9 @@ singlestoredb/notebook/_portal.py,sha256=ip3MkJ51syxppUVGRjgqLfR0He4KCZVklFwWT_X
 singlestoredb/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 singlestoredb/tests/empty.sql,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 singlestoredb/tests/local_infile.csv,sha256=0fYxcZcTvcwS2McF4sktFsipRY1G-ClGmCRR1eCqdEQ,45
+singlestoredb/tests/test.ipynb,sha256=IEgXbByXyWDplZvod1J2SqNHZiPOGdD4oNVMd0ArP7s,234
 singlestoredb/tests/test.sql,sha256=winJzhZ2W52PcQ1j8X_NNIDRBmEa-xMAYrS_hoUtAP8,18337
+singlestoredb/tests/test2.ipynb,sha256=_kBQVvEoinQ1zInNcWFKpbdw-djkLsEO8l3g2MEU_Zs,236
 singlestoredb/tests/test2.sql,sha256=CEM8_lX189iQU65G3Pod7lig6osfrveQyoDz6HC35YQ,38
 singlestoredb/tests/test_basics.py,sha256=tLiR46qUy8-OHHRSHsQTBp5q9NAwNPR9HvfutI6YgnM,47629
 singlestoredb/tests/test_config.py,sha256=Ad0PDmCnJMOyy9f7WTKiRasSR_3mYRByUlSb7k5ZySg,11502
@@ -106,9 +113,9 @@ singlestoredb/tests/test_dbapi.py,sha256=cNJoTEZvYG7ckcwT7xqlkJX-2TDEYGTDDU1Iguc
 singlestoredb/tests/test_exceptions.py,sha256=vscMYmdOJr0JmkTAJrNI2w0Q96Nfugjkrt5_lYnw8i0,1176
 singlestoredb/tests/test_ext_func.py,sha256=gQErR-wAN8BqLNG5U4pNbg4qkQEo6Re8Hd9_Ztqo1RM,38550
 singlestoredb/tests/test_ext_func_data.py,sha256=9kn8BWmCjkbnP6hSbFhmhcdW4OmVT-GSvBTIzFBLEys,48796
-singlestoredb/tests/test_fusion.py,sha256=Qab2TTVwUPo7NPso6nIQgDkR6hHtDLp5j1gO7X4bDsk,24685
+singlestoredb/tests/test_fusion.py,sha256=ckATjXKcDBvej68PZgTRJirdoywtUJ7WXkZmfR0la_4,24544
 singlestoredb/tests/test_http.py,sha256=7hwXe61hlUes3nji0MTTZweo94tJAlJ-vA5ct9geXFQ,8868
-singlestoredb/tests/test_management.py,sha256=B4NkK7J0luuS7T-7OR5qzu-v8gkViIiXie-58bHQIDQ,35334
+singlestoredb/tests/test_management.py,sha256=y6K9dm-s9wxi1yk7fGt-h6hYlB9W4VZy4FgxfVbgZcE,44782
 singlestoredb/tests/test_plugin.py,sha256=P1nXLnTafaHkHN-6bVbGryxTu7OWJPU9SYFZ_WQUwq8,845
 singlestoredb/tests/test_results.py,sha256=Zg1ynZFRZqalAMfNLOU5C6BDXaox6JxrKm_XZwVNFcg,6753
 singlestoredb/tests/test_types.py,sha256=YeVE6KPqlqzJke-4hbRmc8ko1E7RLHu5S8qLg04Bl5Y,4632
@@ -125,9 +132,11 @@ singlestoredb/utils/events.py,sha256=rC9cHAetua_E1f-EiFkFM-gJzQSQIH5Uk-4sspC3KjI
 singlestoredb/utils/mogrify.py,sha256=gCcn99-vgsGVjTUV7RHJ6hH4vCNrsGB_Xo4z8kiSPDQ,4201
 singlestoredb/utils/results.py,sha256=wR70LhCqlobniZf52r67zYLBOKjWHQm68NAskdRQND8,15862
 singlestoredb/utils/xdict.py,sha256=-wi1lSPTnY99fhVMBhPKJ8cCsQhNG4GMUfkEBDKYgCw,13321
-singlestoredb-1.7.2.dist-info/LICENSE,sha256=Bojenzui8aPNjlF3w4ojguDP7sTf8vFV_9Gc2UAG1sg,11542
-singlestoredb-1.7.2.dist-info/METADATA,sha256=EcxKSeGY473TFRuGP6W42122Mz_hry8hPLXAVseir0A,5710
-singlestoredb-1.7.2.dist-info/WHEEL,sha256=UyMHzmWA0xVqVPKfTiLs2eN3OWWZUl-kQemNbpIqlKo,100
-singlestoredb-1.7.2.dist-info/entry_points.txt,sha256=bSLaTWB5zGjpVYPAaI46MkkDup0su-eb3uAhCNYuRV0,48
-singlestoredb-1.7.2.dist-info/top_level.txt,sha256=SDtemIXf-Kp-_F2f_S6x0db33cHGOILdAEsIQZe2LZc,35
-singlestoredb-1.7.2.dist-info/RECORD,,
+sqlx/__init__.py,sha256=4Sdn8HN-Hf8v0_wCt60DCckCg8BvgM3-9r4YVfZycRE,89
+sqlx/magic.py,sha256=6VBlotgjautjev599tHaTYOfcfOA9m6gV_-P1_Qc4lI,3622
+singlestoredb-1.9.0.dist-info/LICENSE,sha256=Bojenzui8aPNjlF3w4ojguDP7sTf8vFV_9Gc2UAG1sg,11542
+singlestoredb-1.9.0.dist-info/METADATA,sha256=-R-w0V56XC6oKl6i-EgytHVnL1Z1_vhuJNG4aLAnj5o,5710
+singlestoredb-1.9.0.dist-info/WHEEL,sha256=UyMHzmWA0xVqVPKfTiLs2eN3OWWZUl-kQemNbpIqlKo,100
+singlestoredb-1.9.0.dist-info/entry_points.txt,sha256=bSLaTWB5zGjpVYPAaI46MkkDup0su-eb3uAhCNYuRV0,48
+singlestoredb-1.9.0.dist-info/top_level.txt,sha256=lA65Vf4qAMfg_s1oG3LEO90h4t1Z-SPDbRqkevI3bSY,40
+singlestoredb-1.9.0.dist-info/RECORD,,

{singlestoredb-1.7.2.dist-info → singlestoredb-1.9.0.dist-info}/top_level.txt

@@ -1,2 +1,3 @@
 _singlestoredb_accel
 singlestoredb
+sqlx

sqlx/__init__.py

@@ -0,0 +1,4 @@
+from sqlx.magic import load_ipython_extension
+
+
+__all__ = ['load_ipython_extension']

sqlx/magic.py

@@ -0,0 +1,113 @@
+import os
+from typing import Any
+from typing import Optional
+
+from IPython.core.interactiveshell import InteractiveShell
+from IPython.core.magic import cell_magic
+from IPython.core.magic import line_magic
+from IPython.core.magic import Magics
+from IPython.core.magic import magics_class
+from IPython.core.magic import needs_local_scope
+from IPython.core.magic import no_var_expand
+from sql.magic import SqlMagic
+from sqlalchemy import create_engine
+from sqlalchemy import Engine
+from sqlalchemy import PoolProxiedConnection
+
+DEFAULT_POOL_SIZE = 10  # Maximum number of connections in the pool
+DEFAULT_MAX_OVERFLOW = 5  # additional connections (temporary overflow)
+DEFAULT_POOL_TIMEOUT = 30  # Wait time for a connection from the pool
+
+
+@magics_class
+class SqlxMagic(Magics):
+    def __init__(self, shell: InteractiveShell):
+        Magics.__init__(self, shell=shell)
+        self.magic = SqlMagic(shell)
+        self.engine: Optional['Engine'] = None
+
+    @no_var_expand
+    @needs_local_scope
+    @line_magic('sqlx')
+    @cell_magic('sqlx')
+    def sqlx(self, line: str, cell: str = '', local_ns: Any = None) -> Any:
+        """
+        Runs SQL statement against a database, specified by
+        SQLAlchemy connect string present in DATABASE_URL environment variable.
+
+        The magic can be used both as a cell magic `%%sqlx` and
+        line magic `%sqlx` (see examples below).
+
+        This is a thin wrapper around the [jupysql](https://jupysql.ploomber.io/) magic,
+        allowing multi-threaded execution.
+        A connection pool will be maintained internally.
+
+        Examples::
+
+          # Line usage
+
+          %sqlx SELECT * FROM mytable
+
+          result = %sqlx SELECT 1
+
+
+          # Cell usage
+
+          %%sqlx
+          DELETE FROM mytable
+
+          %%sqlx
+          DROP TABLE mytable
+
+        """
+
+        connection = self.get_connection()
+        try:
+            result = self.magic.execute(line, cell, local_ns, connection)
+        finally:
+            connection.close()
+
+        return result
+
+    def get_connection(self) -> PoolProxiedConnection:
+        if self.engine is None:
+            if 'DATABASE_URL' not in os.environ:
+                raise RuntimeError(
+                    'Cannot create connection pool, environment variable'
+                    " 'DATABASE_URL' is missing.",
+                )
+
+            # TODO: allow configuring engine
+            # idea: %sqlx engine
+            # idea: %%sqlx engine
+            self.engine = create_engine(
+                os.environ['DATABASE_URL'],
+                pool_size=DEFAULT_POOL_SIZE,
+                max_overflow=DEFAULT_MAX_OVERFLOW,
+                pool_timeout=DEFAULT_POOL_TIMEOUT,
+            )
+
+        return self.engine.raw_connection()
+
+
+# In order to actually use these magics, you must register them with a
+# running IPython.
+
+
+def load_ipython_extension(ip: InteractiveShell) -> None:
+    """
+    Any module file that define a function named `load_ipython_extension`
+    can be loaded via `%load_ext module.path` or be configured to be
+    autoloaded by IPython at startup time.
+    """
+
+    # Load jupysql extension
+    # This is necessary for jupysql to initialize internal state
+    # required to render messages
+    assert ip.extension_manager is not None
+    result = ip.extension_manager.load_extension('sql')
+    if result == 'no load function':
+        raise RuntimeError('Could not load sql extension. Is jupysql installed?')
+
+    # Register sqlx
+    ip.register_magics(SqlxMagic(ip))