added support for cmd specific env vars and auth

Author: jeffreyaven

.github/workflows/test.yaml

@@ -9,20 +9,20 @@ jobs:
     strategy:
       matrix:
         os:
-          # - ubuntu-latest
+          - ubuntu-latest
           - windows-latest
-          # - macos-latest
+          - macos-latest
         python-version:
-          # - "3.7"
-          # - "3.8"
-          # - "3.9"
-          # - "3.10"
-          # - "3.11"
+          - "3.7"
+          - "3.8"
+          - "3.9"
+          - "3.10"
+          - "3.11"
           - "3.12"
           # - "3.13"
-        # exclude:
-        #   - os: macos-latest
-        #     python-version: "3.7"
+        exclude:
+          - os: macos-latest
+            python-version: "3.7"
     runs-on: ${{matrix.os}}
     name: 'Run Tests on ${{matrix.os}} with Python ${{matrix.python-version}}'
 
@@ -44,11 +44,6 @@ jobs:
           python3 -m pip install --upgrade pip
           pip install -r requirements.txt
 
-      # - name: Install psycopg2 for non-Windows OS
-      #   if: matrix.os != 'windows-latest'
-      #   run: |
-      #     pip install psycopg2
-
       # Windows
       - name: Install psycopg2-binary for Windows
         if: matrix.os == 'windows-latest'
@@ -99,26 +94,12 @@ jobs:
       - name: Debug Shell Information
         if: matrix.os == 'windows-latest'
         shell: pwsh
-        env:
-          AWS_ACCESS_KEY_ID: ${{ secrets.AWS_ACCESS_KEY_ID }}
-          AWS_SECRET_ACCESS_KEY: ${{ secrets.AWS_SECRET_ACCESS_KEY }}
-          GOOGLE_CREDENTIALS: ${{ secrets.GOOGLE_CREDENTIALS }}
-          STACKQL_GITHUB_USERNAME: ${{ secrets.STACKQL_GITHUB_USERNAME }}
-          STACKQL_GITHUB_PASSWORD: ${{ secrets.STACKQL_GITHUB_PASSWORD }}
-          CUSTOM_STACKQL_GITHUB_USERNAME: ${{ secrets.CUSTOM_STACKQL_GITHUB_USERNAME }}
-          CUSTOM_STACKQL_GITHUB_PASSWORD: ${{ secrets.CUSTOM_STACKQL_GITHUB_PASSWORD }}
-          AWS_REGION: ${{ vars.AWS_REGION }}
-          AWS_REGIONS: ${{ vars.AWS_REGIONS }}
-          GCP_PROJECT: ${{ vars.GCP_PROJECT }}
-          GCP_ZONE: ${{ vars.GCP_ZONE }}
         run: |
           Write-Host "Checking PowerShell version and environment variables:"
           $PSVersionTable
           Write-Host "Environment variables:"
           Get-ChildItem Env:
           Write-Host "Current Shell: $SHELL"
-          Write-Host "Command-Specific Username: $Env:AWS_REGION"
-          Write-Host "Command-Specific Password: $Env:AWS_ACCESS_KEY_ID"
 
       - name: Run tests on Windows
         env:

pystackql/_util.py

@@ -38,9 +38,9 @@ def _get_download_dir():
 	return site.getuserbase()
 
 def _get_binary_name(platform):
-	if platform == 'Windows':
+	if platform.startswith('Windows'):
 		return r'stackql.exe'
-	elif platform == 'Darwin':
+	elif platform.startswith('Darwin'):
 		return r'stackql/Payload/stackql'
 	else:
 		return r'stackql'
@@ -79,28 +79,71 @@ def _download_file(url, path, showprogress=True):
 		print("ERROR: [_download_file] %s" % (str(e)))
 		exit(1)
 
+# def _setup(download_dir, platform, showprogress=False):
+#     try:
+#         print('installing stackql...')
+#         binary_name = _get_binary_name(platform)
+#         url = _get_url()
+#         print("downloading latest version of stackql from %s to %s" % (url, download_dir))
+#         archive_file_name = os.path.join(download_dir, os.path.basename(url))
+#         _download_file(url, archive_file_name, showprogress)
+#         # if Platform is starting with Darwin, then it is a MacOS
+#         if platform.startswith('Darwin'):
+#             unpacked_file_name = os.path.join(download_dir, 'stackql')
+#             command = 'pkgutil --expand-full {} {}'.format(archive_file_name, unpacked_file_name)
+#             # if there are files in unpacked_file_name, then remove them
+#             if os.path.exists(unpacked_file_name):
+#                 os.system('rm -rf {}'.format(unpacked_file_name))
+#             os.system(command)
+#         else:
+#             with zipfile.ZipFile(archive_file_name, 'r') as zip_ref:
+#                 zip_ref.extractall(download_dir) 
+#         os.chmod(os.path.join(download_dir, binary_name), 0o755)
+#     except Exception as e:
+#         print("ERROR: [_setup] %s" % (str(e)))
+#         exit(1)
+
 def _setup(download_dir, platform, showprogress=False):
     try:
         print('installing stackql...')
-        binary_name = _get_binary_name(platform)
+        binary_name = _get_binary_name(platform)  # Should return 'stackql.exe' for Windows
         url = _get_url()
-        print("downloading latest version of stackql from %s to %s" % (url, download_dir))
+        print(f"Downloading latest version of stackql from {url} to {download_dir}")
+
+        # Paths
         archive_file_name = os.path.join(download_dir, os.path.basename(url))
+        binary_path = os.path.join(download_dir, binary_name)
+
+        # Download and extract
         _download_file(url, archive_file_name, showprogress)
-        # if Platform is starting with Darwin, then it is a MacOS
+
+        # Handle extraction
         if platform.startswith('Darwin'):
             unpacked_file_name = os.path.join(download_dir, 'stackql')
-            command = 'pkgutil --expand-full {} {}'.format(archive_file_name, unpacked_file_name)
-            # if there are files in unpacked_file_name, then remove them
+            command = f'pkgutil --expand-full {archive_file_name} {unpacked_file_name}'
             if os.path.exists(unpacked_file_name):
-                os.system('rm -rf {}'.format(unpacked_file_name))
+                os.system(f'rm -rf {unpacked_file_name}')
             os.system(command)
-        else:
+
+        else:  # Handle Windows and Linux
             with zipfile.ZipFile(archive_file_name, 'r') as zip_ref:
-                zip_ref.extractall(download_dir) 
-        os.chmod(os.path.join(download_dir, binary_name), 0o755)
+                zip_ref.extractall(download_dir)
+            
+            # Specific check for Windows to ensure `stackql.exe` is extracted
+            if platform.startswith("Windows"):
+                if not os.path.exists(binary_path) and os.path.exists(os.path.join(download_dir, "stackql")):
+                    os.rename(os.path.join(download_dir, "stackql"), binary_path)
+
+        # Confirm binary presence and set permissions
+        if os.path.exists(binary_path):
+            print(f"StackQL executable successfully located at: {binary_path}")
+            os.chmod(binary_path, 0o755)
+        else:
+            print(f"ERROR: Expected binary '{binary_path}' not found after extraction.")
+            exit(1)
+
     except Exception as e:
-        print("ERROR: [_setup] %s" % (str(e)))
+        print(f"ERROR: [_setup] {str(e)}")
         exit(1)
 
 def _get_version(bin_path):

pystackql/stackql.py

@@ -207,10 +207,12 @@ def _run_query(self, query, custom_auth=None, env_vars=None):
 		"""
 		local_params = self.params.copy()
 		local_params.insert(1, f'"{query}"')
+		script_path = None
 
 		# Handle custom authentication if provided
 		if custom_auth:
 			if '--auth' in local_params:
+				# override auth set in the constructor with the command-specific auth
 				auth_index = local_params.index('--auth')
 				local_params.pop(auth_index)  # remove --auth
 				local_params.pop(auth_index)  # remove the auth string
@@ -220,50 +222,23 @@ def _run_query(self, query, custom_auth=None, env_vars=None):
 		output = {}
 		env_command_prefix = ""
 
-		# # Determine platform and set environment command prefix accordingly
-		# if env_vars:
-		# 	if self.platform.startswith("Windows"):
-		# 		# For Windows, use PowerShell syntax
-		# 		env_command_prefix = "& { " + " ".join([f'$env:{key} = "{value}";' for key, value in env_vars.items()]) + " "
-		# 		full_command = f"{env_command_prefix}{self.bin_path} " + " ".join(local_params) + " }"
-		# 		is_github_actions = os.environ.get('GITHUB_ACTIONS') == 'true'
-		# 		# ok new approach, if we are in github actions, dump the command to a file and run it
-		# 		# GO!
-		# 	else:
-		# 		# For Linux/Mac, use standard env variable syntax
-		# 		env_command_prefix = "env " + " ".join([f'{key}="{value}"' for key, value in env_vars.items()]) + " "
-		# 		full_command = env_command_prefix + " ".join([self.bin_path] + local_params)
-		# else:
-		# 	full_command = " ".join([self.bin_path] + local_params)
-
-		# print(full_command)  # For debugging
-
 		# Determine platform and set environment command prefix accordingly
 		if env_vars:
 			if self.platform.startswith("Windows"):
-				# For Windows, use PowerShell syntax with GitHub Actions check
-				env_command_prefix = "& { " + " ".join([f'$env:{key} = "{value}";' for key, value in env_vars.items()]) + " "
-				full_command = f"{env_command_prefix}{self.bin_path} " + " ".join(local_params) + " }"
-				is_github_actions = os.environ.get('GITHUB_ACTIONS') == 'true'
-				
-				# If in GitHub Actions, write command to a PowerShell script file and execute it
-				if is_github_actions:
-					with tempfile.NamedTemporaryFile(delete=False, suffix=".ps1", mode="w") as script_file:
-						# Write environment variable setup and command to script file
-						for key, value in env_vars.items():
-							script_file.write(f'$env:{key} = "{value}";\n')
-						script_file.write(f"{self.bin_path} " + " ".join(local_params) + "\n")
-						script_path = script_file.name
-					full_command = f"powershell -File {script_path}"
+				with tempfile.NamedTemporaryFile(delete=False, suffix=".ps1", mode="w") as script_file:
+					# Write environment variable setup and command to script file
+					for key, value in env_vars.items():
+						script_file.write(f'$env:{key} = "{value}";\n')
+					script_file.write(f"{self.bin_path} " + " ".join(local_params) + "\n")
+					script_path = script_file.name
+				full_command = f"powershell -File {script_path}"
 			else:
 				# For Linux/Mac, use standard env variable syntax
 				env_command_prefix = "env " + " ".join([f'{key}="{value}"' for key, value in env_vars.items()]) + " "
 				full_command = env_command_prefix + " ".join([self.bin_path] + local_params)
 		else:
 			full_command = " ".join([self.bin_path] + local_params)
 
-		print(full_command)  # For debugging
-
 		try:
 			with subprocess.Popen(full_command, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE) as iqlPopen:
 				stdout, stderr = iqlPopen.communicate()
@@ -290,8 +265,11 @@ def _run_query(self, query, custom_auth=None, env_vars=None):
 				"stderr": stderr.decode('utf-8') if 'stderr' in locals() and isinstance(stderr, bytes) else ""
 			}
 			output["exception"] = f"ERROR: {json.dumps(error_details)}"
-
-		return output
+		finally:
+			# Clean up the temporary script file
+			if script_path is not None:
+				os.remove(script_path) 
+			return output
 
 	def __init__(self, 
 				 server_mode=False, 

run_tests.ps1

@@ -1,6 +1,6 @@
 # install packages
-pip.exe install -r requirements.txt --user
-pip install psycopg2-binary --user
+# pip.exe install -r requirements.txt --user
+# pip install psycopg2-binary --user
 
 # Load environment variables
 . .\tests\creds\env_vars\test.env.ps1

tests/pystackql_tests.py

@@ -40,10 +40,8 @@ def setUpModule():
 
     print("downloading aws provider for tests...")
     res = PyStackQLTestsBase.stackql.executeStmt(registry_pull_aws_query)
-    print(res)
     print("downloading google provider for tests...")
     res = PyStackQLTestsBase.stackql.executeStmt(registry_pull_google_query)
-    print(res)
     print("starting stackql server...")
     PyStackQLTestsBase.server_process = subprocess.Popen([PyStackQLTestsBase.stackql.bin_path, "srv", "--pgsrv.address", server_address, "--pgsrv.port", str(server_port)])
     time.sleep(10)
@@ -159,40 +157,37 @@ def test_10_executeStmt(self):
         okta_result = okta_result_dict["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["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 10 executeStmt method\nRESULTS:\n{okta_result_dict}\n{github_result_dict}""", True)
+        print_test_result(f"""Test 10 executeStmt method\nRESULTS:\n{okta_result_dict}""", True)
 
     @pystackql_test_setup(output="csv")
     def test_11_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 11 executeStmt method with csv output\nRESULTS:\n{okta_result}\n{github_result}""", True)
+        print_test_result(f"""Test 11 executeStmt method with csv output\nRESULTS:\n{github_result}""", True)
 
     @pystackql_test_setup(output="pandas")
     def test_12_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 12 executeStmt method with pandas output\nRESULTS:\n{okta_result_df}\n{github_result_df}""", True)
+        homebrew_result_df = self.stackql.executeStmt(registry_pull_homebrew_query)
+        homebrew_result = homebrew_result_df['message'].iloc[0]
+        expected_pattern = registry_pull_resp_pattern("homebrew")
+        self.assertTrue(re.search(expected_pattern, homebrew_result), f"Expected pattern not found in result: {homebrew_result}")
+        print_test_result(f"""Test 12 executeStmt method with pandas output\nRESULTS:\n{homebrew_result_df}""", True)
 
     @pystackql_test_setup()
     def test_13_execute_with_defaults(self):
-        result = self.stackql.execute(google_query)
-        is_valid_data_resp = isinstance(result, list) and all(isinstance(item, dict) for item in result)
-        self.assertTrue(is_valid_data_resp, f"Result is not valid: {result}")
-        print_test_result(f"Test 13 execute with defaults\nRESULT: {result}", is_valid_data_resp)
+        result = self.stackql.execute(google_show_services_query)
+        is_valid_data_resp = (
+            isinstance(result, list) 
+            and all(isinstance(item, dict) and 'error' not in item for item in result)
+        )
+        # Truncate the result message if it's too long
+        truncated_result = (
+            str(result)[:500] + '...' if len(str(result)) > 500 else str(result)
+        )
+        self.assertTrue(is_valid_data_resp, f"Result is not valid: {truncated_result}")
+        print_test_result(f"Test 13 execute with defaults\nRESULT: {truncated_result}", is_valid_data_resp)
+
 
     def test_14_execute_with_defaults_null_response(self):
         result = self.stackql.execute("SELECT 1 WHERE 1=0")
@@ -234,29 +229,23 @@ def test_16_execute_with_csv_output(self, mock_execute):
 
     @pystackql_test_setup()
     def test_17_execute_default_auth_dict_output(self):
-        query = "select login from github.users.users"
-        result = self.stackql.execute(query)
-        
+        result = self.stackql.execute(github_query)
         # Expected result based on default auth
         expected_result = [
             {"login": "stackql-devops-1"}
         ]
-        
         self.assertTrue(isinstance(result, list), "Result should be a list")
         self.assertEqual(result, expected_result, f"Expected result: {expected_result}, got: {result}")
         print_test_result(f"Test 17 execute with default auth and dict output\nRESULT: {result}", result == expected_result)
 
 
     @pystackql_test_setup()
     def test_18_execute_custom_auth_env_vars(self):
-        query = "select login from github.users.users"
-        
         # Set up custom environment variables for authentication
         env_vars = {
             'command_specific_username': os.getenv('CUSTOM_STACKQL_GITHUB_USERNAME'),
             'command_specific_password': os.getenv('CUSTOM_STACKQL_GITHUB_PASSWORD')
         }
-        
         # Define custom authentication configuration
         custom_auth = {
             "github": {
@@ -265,14 +254,11 @@ def test_18_execute_custom_auth_env_vars(self):
                 "password_var": "command_specific_password"
             }
         }
-        
-        result = self.stackql.execute(query, custom_auth=custom_auth, env_vars=env_vars)
-        
+        result = self.stackql.execute(github_query, custom_auth=custom_auth, env_vars=env_vars)
         # Expected result based on custom auth
         expected_result = [
             {"login": "stackql-devops-2"}
         ]
-        
         self.assertTrue(isinstance(result, list), "Result should be a list")
         self.assertEqual(result, expected_result, f"Expected result: {expected_result}, got: {result}")
         print_test_result(f"Test 18 execute with custom auth and command-specific environment variables\nRESULT: {result}", result == expected_result)