From c08b80c9c76b512a01a5dbf61a6ae13f8a197f9a Mon Sep 17 00:00:00 2001 From: Jeffrey Aven Date: Wed, 18 Oct 2023 13:31:22 +1100 Subject: [PATCH] v3.2.0 updates to magic ext --- CHANGELOG.md | 8 +++ docs/source/conf.py | 2 +- docs/source/magic_ext.rst | 19 ++++++- pystackql/__init__.py | 3 +- pystackql/base_stackql_magic.py | 73 +++++++++++++++++++++++++ pystackql/stackql_magic.py | 91 +++---------------------------- pystackql/stackql_server_magic.py | 13 +++++ setup.py | 2 +- tests/pystackql_tests.py | 80 +++++++++++++++------------ tests/test_params.py | 2 +- 10 files changed, 168 insertions(+), 125 deletions(-) create mode 100644 pystackql/base_stackql_magic.py create mode 100644 pystackql/stackql_server_magic.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 9c82013..2574a64 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,13 @@ # Changelog +## v3.2.0 (2023-10-18) + +* implemented non `server_mode` magic extension + +### Updates + + * `pandas` type fixes + ## v3.1.2 (2023-10-16) ### Updates diff --git a/docs/source/conf.py b/docs/source/conf.py index c67206b..084d47c 100644 --- a/docs/source/conf.py +++ b/docs/source/conf.py @@ -26,7 +26,7 @@ # The short X.Y version version = '' # The full version, including alpha/beta/rc tags -release = '3.1.2' +release = '3.2.0' # -- General configuration --------------------------------------------------- diff --git a/docs/source/magic_ext.rst b/docs/source/magic_ext.rst index fed427c..695ba91 100644 --- a/docs/source/magic_ext.rst +++ b/docs/source/magic_ext.rst @@ -1,7 +1,7 @@ StackqlMagic Extension for Jupyter ================================== -The ``StackqlMagic`` extension for Jupyter notebooks provides a convenient interface to run SQL queries against the StackQL database directly from within the notebook environment. Results can be visualized in a tabular format using Pandas DataFrames. +The ``StackqlMagic`` extension for Jupyter notebooks provides a convenient interface to run StackQL queries against cloud or SaaS providers directly from within the notebook environment. Results can be visualized in a tabular format using Pandas DataFrames. Setup ----- @@ -10,7 +10,13 @@ To enable the `StackqlMagic` extension in your Jupyter notebook, use the followi .. code-block:: python - %load_ext pystackql + %load_ext stackql_magic + +To use the `StackqlMagic` extension in your Jupyter notebook to run queries against a StackQL server, use the following command: + +.. code-block:: python + + %load_ext stackql_server_magic Usage ----- @@ -54,9 +60,10 @@ Example: .. code-block:: python %%stackql --no-display - SELECT name, status + SELECT SPLIT_PART(machineType, '/', -1) as machine_type, count(*) as num_instances FROM google.compute.instances WHERE project = '$project' AND zone = '$zone' + GROUP BY machine_type This will run the query but won't display the results in the notebook. Instead, you can later access the results via the `stackql_df` variable. @@ -64,6 +71,12 @@ This will run the query but won't display the results in the notebook. Instead, The results of the queries are always saved in a Pandas DataFrame named `stackql_df` in the notebook's current namespace. This allows you to further process or visualize the data as needed. +An example of visualizing the results using Pandas is shown below: + +.. code-block:: python + + stackql_df.plot(kind='pie', y='num_instances', labels=_['machine_type'], title='Instances by Type', autopct='%1.1f%%') + -------- This documentation provides a basic overview and usage guide for the `StackqlMagic` extension. For advanced usage or any additional features provided by the extension, refer to the source code or any other accompanying documentation. diff --git a/pystackql/__init__.py b/pystackql/__init__.py index 9a56aa5..b34d155 100644 --- a/pystackql/__init__.py +++ b/pystackql/__init__.py @@ -1,2 +1,3 @@ from .stackql import StackQL -from .stackql_magic import load_ipython_extension \ No newline at end of file +from .stackql_magic import StackqlMagic, load_non_server_magic +from .stackql_server_magic import StackqlServerMagic, load_server_magic \ No newline at end of file diff --git a/pystackql/base_stackql_magic.py b/pystackql/base_stackql_magic.py new file mode 100644 index 0000000..697aa6f --- /dev/null +++ b/pystackql/base_stackql_magic.py @@ -0,0 +1,73 @@ +from __future__ import print_function +import pandas as pd +import json, argparse +from IPython.core.magic import (Magics, line_cell_magic) +from string import Template + +class BaseStackqlMagic(Magics): + """Base Jupyter magic extension enabling running StackQL queries. + + This extension allows users to conveniently run StackQL queries against cloud + or SaaS reources directly from Jupyter notebooks, and visualize the results in a tabular + format using Pandas DataFrames. + """ + def __init__(self, shell, server_mode): + """Initialize the StackqlMagic class. + + :param shell: The IPython shell instance. + """ + from . import StackQL + super(BaseStackqlMagic, self).__init__(shell) + self.stackql_instance = StackQL(server_mode=server_mode, output='pandas') + + def get_rendered_query(self, data): + """Substitute placeholders in a query template with variables from the current namespace. + + :param data: SQL query template containing placeholders. + :type data: str + :return: A SQL query with placeholders substituted. + :rtype: str + """ + t = Template(data) + return t.substitute(self.shell.user_ns) + + def run_query(self, query): + """Execute a StackQL query + + :param query: StackQL query to be executed. + :type query: str + :return: Query results, returned as a Pandas DataFrame. + :rtype: pandas.DataFrame + """ + return self.stackql_instance.execute(query) + + @line_cell_magic + def stackql(self, line, cell=None): + """A Jupyter magic command to run StackQL queries. + + Can be used as both line and cell magic: + - As a line magic: `%stackql QUERY` + - As a cell magic: `%%stackql [OPTIONS]` followed by the QUERY in the next line. + + :param line: The arguments and/or StackQL query when used as line magic. + :param cell: The StackQL query when used as cell magic. + :return: StackQL query results as a named Pandas DataFrame (`stackql_df`). + """ + is_cell_magic = cell is not None + + if is_cell_magic: + parser = argparse.ArgumentParser() + parser.add_argument("--no-display", action="store_true", help="Suppress result display.") + args = parser.parse_args(line.split()) + query_to_run = self.get_rendered_query(cell) + else: + args = None + query_to_run = self.get_rendered_query(line) + + results = self.run_query(query_to_run) + self.shell.user_ns['stackql_df'] = results + + if is_cell_magic and args and not args.no_display: + return results + elif not is_cell_magic: + return results \ No newline at end of file diff --git a/pystackql/stackql_magic.py b/pystackql/stackql_magic.py index 9f8f9f8..245a17f 100644 --- a/pystackql/stackql_magic.py +++ b/pystackql/stackql_magic.py @@ -1,89 +1,12 @@ -from __future__ import print_function -import pandas as pd -import json, argparse -from IPython.core.magic import (Magics, magics_class, line_cell_magic) -from string import Template +# stackql_magic.py +from IPython.core.magic import magics_class +from .base_stackql_magic import BaseStackqlMagic @magics_class -class StackqlMagic(Magics): - """ - A Jupyter magic extension enabling SQL querying against a StackQL database. - - This extension allows users to conveniently run SQL queries against the StackQL - database directly from Jupyter notebooks, and visualize the results in a tabular - format using Pandas DataFrames. - """ - +class StackqlMagic(BaseStackqlMagic): def __init__(self, shell): - """ - Initialize the StackqlMagic class. - - :param shell: The IPython shell instance. - """ - from . import StackQL - super(StackqlMagic, self).__init__(shell) - self.stackql_instance = StackQL(server_mode=True, output='pandas') - - def get_rendered_query(self, data): - """ - Substitute placeholders in a query template with variables from the current namespace. - - :param data: SQL query template containing placeholders. - :type data: str - :return: A SQL query with placeholders substituted. - :rtype: str - """ - t = Template(data) - return t.substitute(self.shell.user_ns) - - def run_query(self, query): - """ - Execute a StackQL query - - :param query: StackQL query to be executed. - :type query: str - :return: Query results, returned as a Pandas DataFrame. - :rtype: pandas.DataFrame - """ - return self.stackql_instance.execute(query) - - @line_cell_magic - def stackql(self, line, cell=None): - """ - A Jupyter magic command to run SQL queries against the StackQL database. - - Can be used as both line and cell magic: - - As a line magic: `%stackql QUERY` - - As a cell magic: `%%stackql [OPTIONS]` followed by the QUERY in the next line. - - :param line: The arguments and/or StackQL query when used as line magic. - :param cell: The StackQL query when used as cell magic. - :return: StackQL query results as a named Pandas DataFrame (`stackql_df`). - """ - is_cell_magic = cell is not None - - if is_cell_magic: - parser = argparse.ArgumentParser() - parser.add_argument("--no-display", action="store_true", help="Suppress result display.") - args = parser.parse_args(line.split()) - query_to_run = self.get_rendered_query(cell) - else: - args = None - query_to_run = self.get_rendered_query(line) - - results = self.run_query(query_to_run) - self.shell.user_ns['stackql_df'] = results - - if is_cell_magic and args and not args.no_display: - return results - elif not is_cell_magic: - return results + super().__init__(shell, server_mode=False) -def load_ipython_extension(ipython): - """ - Enable the StackqlMagic extension in IPython. - - This function allows the extension to be loaded via the `%load_ext` command or - be automatically loaded by IPython at startup. - """ +def load_non_server_magic(ipython): + """Load the non-server magic in IPython.""" ipython.register_magics(StackqlMagic) diff --git a/pystackql/stackql_server_magic.py b/pystackql/stackql_server_magic.py new file mode 100644 index 0000000..c6bb7f5 --- /dev/null +++ b/pystackql/stackql_server_magic.py @@ -0,0 +1,13 @@ +# stackql_server_magic.py +from IPython.core.magic import magics_class +from .base_stackql_magic import BaseStackqlMagic + +@magics_class +class StackqlServerMagic(BaseStackqlMagic): + + def __init__(self, shell): + super().__init__(shell, server_mode=True) + +def load_server_magic(ipython): + """Load the extension in IPython.""" + ipython.register_magics(StackqlServerMagic) \ No newline at end of file diff --git a/setup.py b/setup.py index eb163e2..6254b52 100644 --- a/setup.py +++ b/setup.py @@ -10,7 +10,7 @@ setup( name='pystackql', - version='3.1.2', + version='3.2.0', description='A Python interface for StackQL', long_description=readme, author='Jeffrey Aven', diff --git a/tests/pystackql_tests.py b/tests/pystackql_tests.py index d4d887a..b1c4d6c 100644 --- a/tests/pystackql_tests.py +++ b/tests/pystackql_tests.py @@ -1,8 +1,7 @@ import sys, os, unittest, asyncio from unittest.mock import MagicMock sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))) -from pystackql import StackQL -from pystackql.stackql_magic import StackqlMagic, load_ipython_extension +from pystackql import StackQL, load_non_server_magic, load_server_magic, StackqlMagic, StackqlServerMagic from .test_params import * def pystackql_test_setup(**kwargs): @@ -275,48 +274,61 @@ def instance(): """Return a mock instance of the shell.""" return MockInteractiveShell() -class StackQLMagicTests(PyStackQLTestsBase): - +class BaseStackQLMagicTests: + MAGIC_CLASS = None # To be overridden by child classes + server_mode = None # To be overridden by child classes def setUp(self): """Set up for the magic tests.""" + assert self.MAGIC_CLASS, "MAGIC_CLASS should be set by child classes" self.shell = MockInteractiveShell.instance() - load_ipython_extension(self.shell) - self.stackql_magic = StackqlMagic(shell=self.shell) + if self.server_mode: + load_server_magic(self.shell) + else: + load_non_server_magic(self.shell) + self.stackql_magic = self.MAGIC_CLASS(shell=self.shell) self.query = "SELECT 1 as fred" self.expected_result = pd.DataFrame({"fred": [1]}) - def test_23_line_magic_query(self): - # Mock the run_query method to return a known DataFrame. - self.stackql_magic.run_query = MagicMock(return_value=self.expected_result) - # Execute the line magic with our query. - result = self.stackql_magic.stackql(line=self.query, cell=None) - # Check if the result is as expected and if 'stackql_df' is set in the namespace. - self.assertTrue(result.equals(self.expected_result)) - self.assertTrue('stackql_df' in self.shell.user_ns) - self.assertTrue(self.shell.user_ns['stackql_df'].equals(self.expected_result)) - print_test_result(f"""Line magic test""", True, True, True) - - def test_24_cell_magic_query(self): - # Mock the run_query method to return a known DataFrame. - self.stackql_magic.run_query = MagicMock(return_value=self.expected_result) - # Execute the cell magic with our query. - result = self.stackql_magic.stackql(line="", cell=self.query) - # Validate the outcome. - self.assertTrue(result.equals(self.expected_result)) - self.assertTrue('stackql_df' in self.shell.user_ns) - self.assertTrue(self.shell.user_ns['stackql_df'].equals(self.expected_result)) - print_test_result(f"""Cell magic test""", True, True, True) + def print_test_result(self, test_name, *checks): + all_passed = all(checks) + print_test_result(f"{test_name}, server_mode: {self.server_mode}", all_passed, True, True) - def test_25_cell_magic_query_no_output(self): + def run_magic_test(self, line, cell, expect_none=False): # Mock the run_query method to return a known DataFrame. self.stackql_magic.run_query = MagicMock(return_value=self.expected_result) - # Execute the cell magic with our query and the no-display argument. - result = self.stackql_magic.stackql(line="--no-display", cell=self.query) + # Execute the magic with our query. + result = self.stackql_magic.stackql(line=line, cell=cell) # Validate the outcome. - self.assertIsNone(result) - self.assertTrue('stackql_df' in self.shell.user_ns) - self.assertTrue(self.shell.user_ns['stackql_df'].equals(self.expected_result)) - print_test_result(f"""Cell magic test (with --no-display)""", True, True, True) + checks = [] + if expect_none: + checks.append(result is None) + else: + checks.append(result.equals(self.expected_result)) + checks.append('stackql_df' in self.shell.user_ns) + checks.append(self.shell.user_ns['stackql_df'].equals(self.expected_result)) + return checks + + def test_line_magic_query(self): + checks = self.run_magic_test(line=self.query, cell=None) + self.print_test_result("Line magic test", *checks) + + def test_cell_magic_query(self): + checks = self.run_magic_test(line="", cell=self.query) + self.print_test_result("Cell magic test", *checks) + + def test_cell_magic_query_no_output(self): + checks = self.run_magic_test(line="--no-display", cell=self.query, expect_none=True) + self.print_test_result("Cell magic test (with --no-display)", *checks) + + +class StackQLMagicTests(BaseStackQLMagicTests, unittest.TestCase): + + MAGIC_CLASS = StackqlMagic + server_mode = False + +class StackQLServerMagicTests(BaseStackQLMagicTests, unittest.TestCase): + MAGIC_CLASS = StackqlServerMagic + server_mode = True def main(): unittest.main(verbosity=0) diff --git a/tests/test_params.py b/tests/test_params.py index b786cb3..cbb46a5 100644 --- a/tests/test_params.py +++ b/tests/test_params.py @@ -62,7 +62,7 @@ def registry_pull_resp_pattern(provider): for region in regions ] -def print_test_result(test_name, condition, server_mode=False, is_ipython=False): +def print_test_result(test_name, condition=True, server_mode=False, is_ipython=False): status_header = colored("[PASSED] ", 'green') if condition else colored("[FAILED] ", 'red') headers = [status_header]