diff --git a/sqlx/__init__.py b/sqlx/__init__.py new file mode 100644 index 00000000..a1fb5a48 --- /dev/null +++ b/sqlx/__init__.py @@ -0,0 +1,4 @@ +from sqlx.magic import load_ipython_extension + + +__all__ = ['load_ipython_extension'] diff --git a/sqlx/magic.py b/sqlx/magic.py new file mode 100644 index 00000000..f49f2b6c --- /dev/null +++ b/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)) diff --git a/test-requirements.txt b/test-requirements.txt index d63e274d..e0ce25a4 100644 --- a/test-requirements.txt +++ b/test-requirements.txt @@ -1,6 +1,8 @@ coverage dash fastapi +ipython +jupysql pandas parameterized polars