Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

v3.2.0 updates to magic ext #24

Merged
merged 1 commit into from
Oct 18, 2023
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 8 additions & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
@@ -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
Expand Down
2 changes: 1 addition & 1 deletion docs/source/conf.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 ---------------------------------------------------
Expand Down
19 changes: 16 additions & 3 deletions docs/source/magic_ext.rst
Original file line number Diff line number Diff line change
@@ -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
-----
Expand All @@ -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
-----
Expand Down Expand Up @@ -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.
3 changes: 2 additions & 1 deletion pystackql/__init__.py
Original file line number Diff line number Diff line change
@@ -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
73 changes: 73 additions & 0 deletions pystackql/base_stackql_magic.py
Original file line number Diff line number Diff line change
@@ -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
91 changes: 7 additions & 84 deletions pystackql/stackql_magic.py
Original file line number Diff line number Diff line change
@@ -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)
13 changes: 13 additions & 0 deletions pystackql/stackql_server_magic.py
Original file line number Diff line number Diff line change
@@ -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)
2 changes: 1 addition & 1 deletion setup.py
Original file line number Diff line number Diff line change
Expand Up @@ -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',
Expand Down
80 changes: 46 additions & 34 deletions tests/pystackql_tests.py
Original file line number Diff line number Diff line change
@@ -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):
Expand Down Expand Up @@ -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)
Expand Down
2 changes: 1 addition & 1 deletion tests/test_params.py
Original file line number Diff line number Diff line change
Expand Up @@ -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]

Expand Down