From 1ee6ba98a874e229fe5965f4929917ad236ea4d4 Mon Sep 17 00:00:00 2001 From: Ankit Saini Date: Wed, 25 Sep 2024 22:51:35 +0530 Subject: [PATCH 1/3] Fix output formatting for run_function_app function --- singlestoredb/apps/_cloud_functions.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/singlestoredb/apps/_cloud_functions.py b/singlestoredb/apps/_cloud_functions.py index d8bf43948..63f12ff57 100644 --- a/singlestoredb/apps/_cloud_functions.py +++ b/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: From 271b1650f7cd783b470236288e0d61c4f3b56609 Mon Sep 17 00:00:00 2001 From: Ankit Saini Date: Tue, 19 Nov 2024 12:53:22 +0530 Subject: [PATCH 2/3] adding sqlx magic --- sqlx/__init__.py | 4 ++ sqlx/magic.py | 113 ++++++++++++++++++++++++++++++++++++++++++ test-requirements.txt | 2 + 3 files changed, 119 insertions(+) create mode 100644 sqlx/__init__.py create mode 100644 sqlx/magic.py diff --git a/sqlx/__init__.py b/sqlx/__init__.py new file mode 100644 index 000000000..a1fb5a482 --- /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 000000000..04bf252da --- /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 `%%sql` and + line magic `%sql` (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 + + %%sql + DELETE FROM mytable + + %%sql + 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 d63e274de..e0ce25a4f 100644 --- a/test-requirements.txt +++ b/test-requirements.txt @@ -1,6 +1,8 @@ coverage dash fastapi +ipython +jupysql pandas parameterized polars From 2dbbcbf9cafbdacf3c3ad0affb4d9f92588785d5 Mon Sep 17 00:00:00 2001 From: Ankit Saini Date: Tue, 3 Dec 2024 21:38:41 +0530 Subject: [PATCH 3/3] fix sql -> sqlx typo in docstrings --- sqlx/magic.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/sqlx/magic.py b/sqlx/magic.py index 04bf252da..f49f2b6cd 100644 --- a/sqlx/magic.py +++ b/sqlx/magic.py @@ -35,8 +35,8 @@ 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 `%%sql` and - line magic `%sql` (see examples below). + 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. @@ -53,10 +53,10 @@ def sqlx(self, line: str, cell: str = '', local_ns: Any = None) -> Any: # Cell usage - %%sql + %%sqlx DELETE FROM mytable - %%sql + %%sqlx DROP TABLE mytable """