statement fixes

Author: jeffreyaven

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.
@@ -395,7 +400,7 @@ def executeStmt(self, 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. 
@@ -406,14 +411,26 @@ def executeStmt(self, query):
 			>>> 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 

tests/pystackql_tests.py

@@ -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,6 +315,7 @@ 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)
@@ -320,6 +348,31 @@ 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)
 
+    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 = []
+        if expect_none:
+            checks.append(result is None)
+        else:
+            # Check that the output contains expected content
+            checks.append("OK" in result["message"].iloc[0])
+        checks.append('stackql_df' in self.shell.user_ns)
+        checks.append("OK" in self.shell.user_ns['stackql_df']["message"].iloc[0])
+        return checks
+
+    def test_line_magic_statement(self):
+        checks = self.run_magic_statement_test(line=self.statement, cell=None)
+        self.print_test_result("Line magic statement", *checks)
+
+    def test_cell_magic_statement(self):
+        checks = self.run_magic_statement_test(line="", cell=self.statement)
+        self.print_test_result("Cell magic statement", *checks)
+
+    def test_cell_magic_statement_no_output(self):
+        checks = self.run_magic_statement_test(line="--no-display", cell=self.statement, expect_none=True)
+        self.print_test_result("Cell magic statement (with --no-display)", *checks)
 
 class StackQLMagicTests(BaseStackQLMagicTests, unittest.TestCase):