Merge pull request #24 from stackql/feature/refactor

v3.2.0 updates to magic ext

Author: jeffreyaven

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

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

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,16 +60,23 @@ 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.
 
 .. note::
 
     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.

pystackql/__init__.py

@@ -1,2 +1,3 @@
 from .stackql import StackQL
-from .stackql_magic import load_ipython_extension
+from .stackql_magic import StackqlMagic, load_non_server_magic
+from .stackql_server_magic import StackqlServerMagic, load_server_magic

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

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)

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)

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',

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)

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]