Merge pull request #28 from stackql/feature/refactor

Feature/refactor

Author: jeffreyaven

CHANGELOG.md

@@ -1,10 +1,11 @@
 # Changelog
 
-## v3.2.3 (2023-10-20)
+## v3.2.4 (2023-10-24)
 
 ### Updates
 
  * implemented non `server_mode` magic extension
+ * updated dataframe output for statements
  * `pandas` type updates
  * updated class parameters
  * added additional tests

README.rst

@@ -193,4 +193,4 @@ To publish the package to PyPI, run the following command:
 
 ::
 
-    twine upload dist/pystackql-3.2.3.tar.gz
+    twine upload dist/pystackql-3.2.4.tar.gz

docs/source/conf.py

@@ -26,7 +26,7 @@
 # The short X.Y version
 version = ''
 # The full version, including alpha/beta/rc tags
-release = '3.2.3'
+release = '3.2.4'
 
 
 # -- General configuration ---------------------------------------------------

pystackql/base_stackql_magic.py

@@ -1,6 +1,7 @@
 from __future__ import print_function
 from IPython.core.magic import (Magics)
 from string import Template
+import pandas as pd
 
 class BaseStackqlMagic(Magics):
     """Base Jupyter magic extension enabling running StackQL queries.
@@ -37,4 +38,8 @@ def run_query(self, query):
         :return: Query results, returned as a Pandas DataFrame.
         :rtype: pandas.DataFrame
         """
+        # Check if the query starts with "registry pull" (case insensitive)
+        if query.strip().lower().startswith("registry pull"):
+            return self.stackql_instance.executeStmt(query)
+        
         return self.stackql_instance.execute(query)

pystackql/stackql.py

@@ -124,19 +124,24 @@ def _connect_to_server(self):
 			print(f"Unexpected error while connecting to the server: {e}")
 		return None
 
-	def _run_server_query(self, query):
+	def _run_server_query(self, query, is_statement=False):
 		"""Runs a query against the server using psycopg2.
 		
 		:param query: SQL query to be executed on the server.
 		:type query: str
 		:return: List of result rows if the query fetches results; empty list if there are no results.
-		:rtype: list
+		:rtype: list of dict objects
 		:raises: psycopg2.ProgrammingError for issues related to the SQL query, 
 				unless the error is "no results to fetch", in which case an empty list is returned.
 		"""
 		try:
 			cur = self._conn.cursor(cursor_factory=RealDictCursor)
 			cur.execute(query)
+			if is_statement:
+				# If the query is a statement, there are no results to fetch.
+				result_msg = cur.statusmessage
+				cur.close()
+				return [{'message': result_msg}]
 			rows = cur.fetchall()
 			cur.close()
 			return rows
@@ -146,7 +151,7 @@ def _run_server_query(self, query):
 			else:
 				raise
 
-	def _run_query(self, query, is_statement=False):
+	def _run_query(self, query):
 		"""Internal method to execute a StackQL query using a subprocess.
 
 		The method spawns a subprocess to run the StackQL binary with the specified query and parameters.
@@ -388,32 +393,45 @@ def upgrade(self, showprogress=True):
 	def executeStmt(self, query):
 		"""Executes a query using the StackQL instance and returns the output as a string.  
 			This is intended for operations which do not return a result set, for example a mutation 
-			operation such as an `INSERT` or a `DELETE` or life cycle method such as an `EXEC` operation.
+			operation such as an `INSERT` or a `DELETE` or life cycle method such as an `EXEC` operation
+			or a `REGISTRY PULL` operation.
 
 		This method determines the mode of operation (server_mode or local execution) based 
 		on the `server_mode` attribute of the instance. If `server_mode` is True, it runs the query 
 		against the server. Otherwise, it executes the query using a subprocess.
 
 		:param query: The StackQL query string to be executed.
-		:type query: str
+		:type query: str, list of dict objects, or Pandas DataFrame
 
 		:return: The output result of the query in string format. If in `server_mode`, it 
 				returns a JSON string representation of the result. 
-		:rtype: str
+		:rtype: dict, Pandas DataFrame or str (for `csv` output)
 
 		Example:
 			>>> from pystackql import StackQL
 			>>> stackql = StackQL()
 			>>> stackql_query = "REGISTRY PULL okta"
 			>>> result = stackql.executeStmt(stackql_query)
-			>>> print(result)
+			>>> result
 		"""
 		if self.server_mode:
 			# Use server mode
-			result = self._run_server_query(query)
-			return json.dumps(result)
+			result = self._run_server_query(query, True)
+			if self.output == 'pandas':
+				return pd.DataFrame(result)
+			elif self.output == 'csv':
+				# return the string representation of the result
+				return result[0]['message']
+			else:
+				return result
 		else:
-			return self._run_query(query, is_statement=True)
+			result_msg = self._run_query(query)
+			if self.output == 'pandas':
+				return pd.DataFrame({'message': [result_msg]})
+			elif self.output == 'csv':
+				return result_msg
+			else:
+				return [{'message': result_msg}]			
 
 	def execute(self, query):
 		"""Executes a query using the StackQL instance and returns the output 

setup.py

@@ -10,7 +10,7 @@
 
 setup(
     name='pystackql',
-    version='3.2.3',
+    version='3.2.4',
     description='A Python interface for StackQL',
     long_description=readme,
     author='Jeffrey Aven',

tests/pystackql_tests.py

@@ -1,4 +1,4 @@
-import sys, os, unittest, asyncio
+import sys, os, unittest, asyncio, re
 from unittest.mock import MagicMock
 sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
 from pystackql import StackQL, magic, magics, StackqlMagic, StackqlServerMagic
@@ -155,13 +155,37 @@ def test_09_csv_output_with_header(self):
 
     @pystackql_test_setup()
     def test_10_executeStmt(self):
+        okta_result_dict = self.stackql.executeStmt(registry_pull_okta_query)
+        okta_result = okta_result_dict[0]["message"]
+        expected_pattern = registry_pull_resp_pattern("okta")
+        self.assertTrue(re.search(expected_pattern, okta_result), f"Expected pattern not found in result: {okta_result}")
+        github_result_dict = self.stackql.executeStmt(registry_pull_github_query)
+        github_result = github_result_dict[0]["message"]
+        expected_pattern = registry_pull_resp_pattern("github")
+        self.assertTrue(re.search(expected_pattern, github_result), f"Expected pattern not found in result: {github_result}")
+        print_test_result(f"""Test executeStmt method\nRESULTS:\n{okta_result_dict}\n{github_result_dict}""", True)
+
+    @pystackql_test_setup(output="csv")
+    def test_10a_executeStmt_with_csv_output(self):
         okta_result = self.stackql.executeStmt(registry_pull_okta_query)
         expected_pattern = registry_pull_resp_pattern("okta")
         self.assertTrue(re.search(expected_pattern, okta_result), f"Expected pattern not found in result: {okta_result}")
         github_result = self.stackql.executeStmt(registry_pull_github_query)
         expected_pattern = registry_pull_resp_pattern("github")
         self.assertTrue(re.search(expected_pattern, github_result), f"Expected pattern not found in result: {github_result}")
-        print_test_result(f"""Test executeStmt method\nRESULTS:\n{okta_result}{github_result}""", True)
+        print_test_result(f"""Test executeStmt method with csv output\nRESULTS:\n{okta_result}\n{github_result}""", True)
+
+    @pystackql_test_setup(output="pandas")
+    def test_10b_executeStmt_with_pandas_output(self):
+        okta_result_df = self.stackql.executeStmt(registry_pull_okta_query)
+        okta_result = okta_result_df['message'].iloc[0]
+        expected_pattern = registry_pull_resp_pattern("okta")
+        self.assertTrue(re.search(expected_pattern, okta_result), f"Expected pattern not found in result: {okta_result}")
+        github_result_df = self.stackql.executeStmt(registry_pull_github_query)
+        github_result = github_result_df['message'].iloc[0]
+        expected_pattern = registry_pull_resp_pattern("github")
+        self.assertTrue(re.search(expected_pattern, github_result), f"Expected pattern not found in result: {github_result}")
+        print_test_result(f"""Test executeStmt method with pandas output\nRESULTS:\n{okta_result_df}\n{github_result_df}""", True)
 
     @pystackql_test_setup()
     def test_11_execute_with_defaults(self):
@@ -232,13 +256,16 @@ def test_19_server_mode_connectivity(self):
     @pystackql_test_setup(server_mode=True)
     def test_20_executeStmt_server_mode(self):
         result = self.stackql.executeStmt(registry_pull_google_query)
-        is_valid_json_string_of_empty_list = False
-        try:
-            parsed_result = json.loads(result)
-            is_valid_json_string_of_empty_list = isinstance(parsed_result, list) and len(parsed_result) == 0
-        except json.JSONDecodeError:
-            pass
-        print_test_result("Test executeStmt in server mode", is_valid_json_string_of_empty_list, True)
+        # Checking if the result is a list containing a single dictionary with a key 'message' and value 'OK'
+        is_valid_response = isinstance(result, list) and len(result) == 1 and result[0].get('message') == 'OK'
+        print_test_result(f"Test executeStmt in server mode\n{result}", is_valid_response, True)
+
+    @pystackql_test_setup(server_mode=True, output='pandas')
+    def test_20a_executeStmt_server_mode_with_pandas_output(self):
+        result_df = self.stackql.executeStmt(registry_pull_google_query)
+        # Verifying if the result is a dataframe with a column 'message' containing the value 'OK' in its first row
+        is_valid_response = isinstance(result_df, pd.DataFrame) and 'message' in result_df.columns and result_df['message'].iloc[0] == 'OK'
+        print_test_result(f"Test executeStmt in server mode with pandas output\n{result_df}", is_valid_response, True)
 
     @pystackql_test_setup(server_mode=True)
     def test_21_execute_server_mode_default_output(self):
@@ -288,10 +315,11 @@ def setUp(self):
         self.stackql_magic = self.MAGIC_CLASS(shell=self.shell)
         self.query = "SELECT 1 as fred"
         self.expected_result = pd.DataFrame({"fred": [1]})
+        self.statement = "REGISTRY PULL github"
 
     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)
+        print_test_result(f"{test_name}", all_passed, self.server_mode, True)
 
     def run_magic_test(self, line, cell, expect_none=False):
         # Mock the run_query method to return a known DataFrame.
@@ -310,16 +338,52 @@ def run_magic_test(self, line, cell, expect_none=False):
 
     def test_line_magic_query(self):
         checks = self.run_magic_test(line=self.query, cell=None)
-        self.print_test_result("Line magic test", *checks)
+        self.print_test_result("Line magic query 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)
+        self.print_test_result("Cell magic query 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)
+        self.print_test_result("Cell magic query test (with --no-display)", *checks)
 
+    def run_magic_statement_test(self, line, cell, expect_none=False):
+        # Execute the magic with our statement.
+        result = self.stackql_magic.stackql(line=line, cell=cell)
+        # Validate the outcome.
+        checks = []
+        # Check that the output contains expected content
+        if expect_none:
+            checks.append(result is None)
+        else:
+            if self.server_mode:
+                checks.append("OK" in result["message"].iloc[0])
+            else:
+                pattern = registry_pull_resp_pattern('github')
+                message = result["message"].iloc[0] if "message" in result.columns else ""
+                checks.append(bool(re.search(pattern, message)))        
+        # Check dataframe exists and is populated as expected
+        checks.append('stackql_df' in self.shell.user_ns)
+        if self.server_mode:
+            checks.append("OK" in self.shell.user_ns['stackql_df']["message"].iloc[0])
+        else:
+            pattern = registry_pull_resp_pattern('github')
+            message = self.shell.user_ns['stackql_df']["message"].iloc[0] if 'stackql_df' in self.shell.user_ns else ""
+            checks.append(bool(re.search(pattern, message)))
+        return checks, result
+
+    def test_line_magic_statement(self):
+        checks, result = self.run_magic_statement_test(line=self.statement, cell=None)
+        self.print_test_result(f"Line magic statement test\n{result}", *checks)
+
+    def test_cell_magic_statement(self):
+        checks, result = self.run_magic_statement_test(line="", cell=self.statement)
+        self.print_test_result(f"Cell magic statement test\n{result}", *checks)
+
+    def test_cell_magic_statement_no_output(self):
+        checks, result = self.run_magic_statement_test(line="--no-display", cell=self.statement, expect_none=True)
+        self.print_test_result(f"Cell magic statement test (with --no-display)\n{result}", *checks)
 
 class StackQLMagicTests(BaseStackQLMagicTests, unittest.TestCase):