singlestoredb 1.8.0__py3-none-any.whl → 1.9.0__py3-none-any.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_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.8.0.dist-info → singlestoredb-1.9.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: singlestoredb
-Version: 1.8.0
+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.8.0.dist-info → singlestoredb-1.9.0.dist-info}/RECORD

@@ -1,7 +1,7 @@
-singlestoredb/__init__.py,sha256=Kvu0WMBtP9TgYwY6eu7vuBfMfx-sbawDTrZDGewMCuo,1634
+singlestoredb/__init__.py,sha256=2NT8djBgcovcYu3ttVecdDhhDZit6vnMyAtm_HsFuHQ,1648
 singlestoredb/auth.py,sha256=u8D9tpKzrqa4ssaHjyZnGDX1q8XBpGtuoOkTkSv7B28,7599
-singlestoredb/config.py,sha256=NtONv4Etpraoy1nenHqRAS08xHJZmho00J95uDjLxQM,12290
-singlestoredb/connection.py,sha256=x5lINBa9kB_GoEEeL2uUZi9G8pNwKuFtA1uqJirR6HI,45352
+singlestoredb/config.py,sha256=ZGmnG37Ug1hT5jdTSdMJBCGDvvz2GwxrZZ_ROVV2wUE,12439
+singlestoredb/connection.py,sha256=PlOD-4Qx3Q-lVPPdjru-kTC8kOWDos3eiWVo7pPKeKM,45679
 singlestoredb/converters.py,sha256=t1hRMZfccWJs_WyOw-W-Kh87fxsOkpOnKXAeh_Nr-zU,20681
 singlestoredb/exceptions.py,sha256=HuoA6sMRL5qiCiee-_5ddTGmFbYC9Euk8TYUsh5GvTw,3234
 singlestoredb/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -31,31 +31,33 @@ singlestoredb/functions/ext/rowdat_1.py,sha256=JgKRsVSQYczFD6cmo2xLilbNPYpyLL2tP
 singlestoredb/functions/ext/utils.py,sha256=2-B8YU_Iekv8JcpI-ochs9TIeuyatLaLAH-AyYyUUIg,5311
 singlestoredb/fusion/__init__.py,sha256=Qo7SuqGw-l-vE8-EI2jhm6hXJkYfOLUKIws9c7LFNX0,356
 singlestoredb/fusion/graphql.py,sha256=ZA3HcDq5rER-dCEavwTqnF7KM0D2LCYIY7nLQk7lSso,5207
-singlestoredb/fusion/handler.py,sha256=ohoGzWLJ-t8co08FUmpe-2ELh5eFY5bLJw7rgxXslso,27540
+singlestoredb/fusion/handler.py,sha256=5DjY8bUf8x-d0kVr9mba_o0W5vLRFj5ImKk-on_UtYk,27607
 singlestoredb/fusion/registry.py,sha256=jjdRTYZ3ylhy6gAoW5xBj0tkxGFBT-2yLQ0tztTgDIY,6112
 singlestoredb/fusion/result.py,sha256=Bd3KbRpqWqQcWp_Chd4bzBy8Kfc8nXLS_Pn_GGbPO6o,11772
 singlestoredb/fusion/handlers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-singlestoredb/fusion/handlers/export.py,sha256=jzPD6b7oTTWYMt25bFNxZVaw6PUfrKpAk-VjvrOQomQ,7097
+singlestoredb/fusion/handlers/export.py,sha256=GJ-tykXHK0iMNAvJrJvtrGCgRpRtlFwFt-e7hxKhp2U,6808
+singlestoredb/fusion/handlers/files.py,sha256=c-QaWzgLGRMheFvG2uUeYfVgLA5WfSE3e9XeBrsFNh0,18979
 singlestoredb/fusion/handlers/job.py,sha256=r0KdOD55VUDw-SymC__5Mn-fzJTZE_xcBgH-O8DYVHc,21095
-singlestoredb/fusion/handlers/stage.py,sha256=4BSwqcmA4PI3oksTdYaTmQ6bY-AzmL0mZfGnP3U0ldM,13643
-singlestoredb/fusion/handlers/utils.py,sha256=oYbf13Y3orEkJfHMNnO7B_W1anEdK-0S9vVVkF2pPFk,5109
+singlestoredb/fusion/handlers/stage.py,sha256=kYVjbPys83kf3jX6jWwN8Ju0oEocKVZ3TIOt2HiC5Ew,14287
+singlestoredb/fusion/handlers/utils.py,sha256=yXggiwnKq7IrCFXQwY-ZNc8gHxvftUUq4NXc8LCP7eM,10046
 singlestoredb/fusion/handlers/workspace.py,sha256=4xN2TFO4yF7KZB2Fcht7IuvoDdAT6fDfDLjixiHZN8w,27506
 singlestoredb/http/__init__.py,sha256=A_2ZUCCpvRYIA6YDpPy57wL5R1eZ5SfP6I1To5nfJ2s,912
 singlestoredb/http/connection.py,sha256=dU0a72pMpyq9l9ADKs5jpB-GAJScBxgd83NOlGreIdc,39473
-singlestoredb/management/__init__.py,sha256=mhWXjLhp5-t8dhl0vl7SjayKrvJlDb5_hl1YWvDgiMA,237
+singlestoredb/management/__init__.py,sha256=ofNTPCdkZ1dS_aX2aUujd8aMHQi8Lle5Ced0aaO3RH4,269
 singlestoredb/management/billing_usage.py,sha256=9ighjIpcopgIyJOktBYQ6pahBZmWGHOPyyCW4gu9FGs,3735
 singlestoredb/management/cluster.py,sha256=i23Smr1PBrDZ8NO_VPd_-bEYkyHvVe9CCRGUjHn_1yQ,14362
-singlestoredb/management/export.py,sha256=jAqpoyKHyT9rPZhxn1OQdq4l7U3V-RG0074XiCyjMZc,8658
+singlestoredb/management/export.py,sha256=_kr5grRVFe0RGuk_XoDGYq_4VXMvGCWvkan9dZO-jGA,4156
+singlestoredb/management/files.py,sha256=AwqcYn1x0uWKzV_RYQgcpSRcMblPP5jKvgcmesyCP00,28166
 singlestoredb/management/job.py,sha256=4-xLWzbE8odQogVVaFer80UEoTAZY1T28VZ9Ug4rbmM,24611
 singlestoredb/management/manager.py,sha256=sFP1vZGS8WpN8E0XLu1N7Mxtq6Sixalln44HlTQEyXI,8800
 singlestoredb/management/organization.py,sha256=hqMaM7H-naMjNbxDl_f7G_2o5TkiGKyzPhxuzDveJAw,5402
 singlestoredb/management/region.py,sha256=HnLcWUh7r_aLECliplCDHak4a_F3B7LOSXEYMW66qD0,1611
 singlestoredb/management/utils.py,sha256=P4fp8a7EwaYiag_hvpILcgwXtdFNYKKO70dsKjmxn1A,13171
-singlestoredb/management/workspace.py,sha256=9oamNIaE5vLZNAlpl5SK-xu27qPqx2Ff3ZIdKckNfmc,61673
+singlestoredb/management/workspace.py,sha256=MdVuCxNv6C1QgiT-1V8kqZeWLeOYWnnDSYMUvj7EtRU,56268
 singlestoredb/mysql/__init__.py,sha256=olUTAvkiERhDW41JXQMawkg-i0tvBEkoTkII1tt6lxU,4492
 singlestoredb/mysql/_auth.py,sha256=AugRitoUwgRIDFuJxuAH4MWIAmckY7Ji2pP6r_Ng9dY,8043
 singlestoredb/mysql/charset.py,sha256=-FlONDS_oAUF5B3mIgeHBPb_SCt4zHD33arUeBNctU0,10510
-singlestoredb/mysql/connection.py,sha256=6N7ZG3rNMrpQhv4eixA686NdJ5GBe4HknSpy46qlZ9g,72367
+singlestoredb/mysql/connection.py,sha256=MfPjQfAadI8xtn4xbaock34bCCJ2P1uaTqAw8o1PqS8,72696
 singlestoredb/mysql/converters.py,sha256=CVe8SDmjbIAhy1xpQ2N5OKWw6t5eWpw-EU3QTlA0Hh0,7500
 singlestoredb/mysql/cursors.py,sha256=Eqe7jITRvOo4P_TxIarTumg_2PG1DcCfZ4Uo9IFdDa8,26794
 singlestoredb/mysql/err.py,sha256=-m5rqXi8yhq6b8SCEJ2h0E5Rudh_15dlAU_WbJ1YrM8,2388
@@ -99,7 +101,9 @@ singlestoredb/notebook/_portal.py,sha256=DLerIEQmAUymtYcx8RBeuYJ4pJSy_xl1K6t1Oc-
 singlestoredb/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 singlestoredb/tests/empty.sql,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 singlestoredb/tests/local_infile.csv,sha256=sBtqjvfkS9aoOVx8nMXYgYv4rDuT4OuYhqUhNRu0O68,42
+singlestoredb/tests/test.ipynb,sha256=jrkI2WoSsUA9xQpKTBCHnsDptryQhPdM5QaxfvYRGpg,216
 singlestoredb/tests/test.sql,sha256=dfMehVCQ9wObSVTQKyQi-fRFDZeqRxV4Cj8doBCPEFM,17679
+singlestoredb/tests/test2.ipynb,sha256=yd1PE1VK-DwiRd6mYS4_0cPBtuVkvcDtycvTwD-YnDo,218
 singlestoredb/tests/test2.sql,sha256=D4U2GSlOVeo39U8-RMM4YziJzYFfi4Ztm2YXJVJVAS8,37
 singlestoredb/tests/test_basics.py,sha256=1__lEF7FmQF4_pFi5R53TtJidtQznmQ592Ci6aDVgrc,46368
 singlestoredb/tests/test_config.py,sha256=63lyIQ2KrvGE6C9403B_4Mc90mX4tp42ys5Bih2sXrE,11184
@@ -110,7 +114,7 @@ singlestoredb/tests/test_ext_func.py,sha256=OWd-CJ1Owhx72nikSWWEF2EQFCJk7vEXZM2O
 singlestoredb/tests/test_ext_func_data.py,sha256=yTADD93nPxX6_rZXXLZaOWEI_yPvYyir9psn5PK9ctU,47695
 singlestoredb/tests/test_fusion.py,sha256=W3aRfBeu8HBGm1CIQWFIeWUPBUlfHBCbJy8vejPHdRs,23828
 singlestoredb/tests/test_http.py,sha256=RXasTqBWRn__omj0eLFTJYIbZjd0PPdIV2d4Cqz0MC8,8580
-singlestoredb/tests/test_management.py,sha256=89hKu82qiH1YTbLzKl5FOPEapNB5qo8k06d3fkoYajc,34304
+singlestoredb/tests/test_management.py,sha256=biqwuHfzaH5Jay1yj-nHnX7fhvLB9DPcjujmJR_T1vM,43480
 singlestoredb/tests/test_plugin.py,sha256=qpO9wmWc62VaijN1sJ97YSYIX7I7Y5C6sY-WzwrutDQ,812
 singlestoredb/tests/test_results.py,sha256=wg93sujwt-R9_eJCgSCElgAZhLDkIiAo3qPkPydOv78,6582
 singlestoredb/tests/test_types.py,sha256=jqoAaSjhbgwB3vt0KsTcl7XBWoMMIa0mPFKhEi5bBjo,4500
@@ -127,9 +131,11 @@ singlestoredb/utils/events.py,sha256=9IB84T3pKQjs7aaoSSJCw7soNngnhoTDWIC52M51R9Y
 singlestoredb/utils/mogrify.py,sha256=-a56IF70U6CkfadeaZgfjRSVsAD3PuqRrzPpjZlgbwY,4050
 singlestoredb/utils/results.py,sha256=bJtaUaDiFq26IsPAKZ2FHGB7csMn94EAxLKrP4HaEEA,15277
 singlestoredb/utils/xdict.py,sha256=S9HKgrPrnu_6b7iOwa2KrW8CmU1Uqx0BWdEyogFzWbE,12896
-singlestoredb-1.8.0.dist-info/LICENSE,sha256=Mlq78idURT-9G026aMYswwwnnrLcgzTLuXeAs5hjDLM,11341
-singlestoredb-1.8.0.dist-info/METADATA,sha256=TOVRXyXoPew6AtTrB7b9oM3PBLYtgsb19SEBBKvOeic,5557
-singlestoredb-1.8.0.dist-info/WHEEL,sha256=bFJAMchF8aTQGUgMZzHJyDDMPTO3ToJ7x23SLJa1SVo,92
-singlestoredb-1.8.0.dist-info/entry_points.txt,sha256=bSLaTWB5zGjpVYPAaI46MkkDup0su-eb3uAhCNYuRV0,48
-singlestoredb-1.8.0.dist-info/top_level.txt,sha256=eet8bVPNRqiGeY0PrO5ERH2UpamwlrKHEQCffz4dOh8,14
-singlestoredb-1.8.0.dist-info/RECORD,,
+sqlx/__init__.py,sha256=aBYiU8DZXCogvWu3yWafOz7bZS5WWwLZXj7oL0dXGyU,85
+sqlx/magic.py,sha256=JsS9_9aBFaOt91Torm1JPN0c8qB2QmYJmNSKtbSQIY0,3509
+singlestoredb-1.9.0.dist-info/LICENSE,sha256=Mlq78idURT-9G026aMYswwwnnrLcgzTLuXeAs5hjDLM,11341
+singlestoredb-1.9.0.dist-info/METADATA,sha256=gCoMBtMdsViRJC1mdvHhMjhKd0IleF7zYsg6LqfNSl0,5557
+singlestoredb-1.9.0.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
+singlestoredb-1.9.0.dist-info/entry_points.txt,sha256=bSLaTWB5zGjpVYPAaI46MkkDup0su-eb3uAhCNYuRV0,48
+singlestoredb-1.9.0.dist-info/top_level.txt,sha256=DfFGz7bM4XrshloiCeTABgylT3BUnS8T5pJam3ewT6Q,19
+singlestoredb-1.9.0.dist-info/RECORD,,

{singlestoredb-1.8.0.dist-info → singlestoredb-1.9.0.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: bdist_wheel (0.45.0)
+Generator: bdist_wheel (0.45.1)
 Root-Is-Purelib: true
 Tag: py3-none-any
 

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

@@ -1 +1,2 @@
 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))