Provide sqlx magic

singlestoredb/apps/_cloud_functions.py

@@ -72,7 +72,7 @@ def ping() -> str:
     if app_config.running_interactively:
         if app_config.is_gateway_enabled:
             print(
-                'Cloud function available at'
+                'Cloud function available at '
                 f'{app_config.base_url}docs?authToken={app_config.token}',
             )
         else:

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))

test-requirements.txt

@@ -1,6 +1,8 @@
 coverage
 dash
 fastapi
+ipython
+jupysql
 pandas
 parameterized
 polars