From 5f4737b7e64434e3b64c6a546a2f3aa00758cba0 Mon Sep 17 00:00:00 2001 From: Jeffrey Aven Date: Wed, 4 Jun 2025 16:16:51 +1000 Subject: [PATCH 01/31] refactor --- .github/workflows/test.yaml | 80 +- .gitignore | 16 +- CHANGELOG.md | 7 + launch_venv.sh | 140 +++ pystackql/__init__.py | 20 +- pystackql/_util.py | 165 ---- pystackql/core/__init__.py | 16 + pystackql/core/binary.py | 79 ++ pystackql/core/output.py | 161 ++++ pystackql/core/query.py | 219 +++++ pystackql/core/server.py | 164 ++++ pystackql/core/stackql.py | 435 ++++++++++ pystackql/magic.py | 50 +- pystackql/magic_ext/__init__.py | 14 + .../base.py} | 20 +- pystackql/magic_ext/local.py | 66 ++ pystackql/magic_ext/server.py | 60 ++ pystackql/magics.py | 50 +- pystackql/stackql.py | 809 ------------------ pystackql/utils/__init__.py | 49 ++ pystackql/utils/auth.py | 39 + pystackql/utils/binary.py | 125 +++ pystackql/utils/download.py | 79 ++ pystackql/utils/helpers.py | 284 ++++++ pystackql/utils/package.py | 31 + pystackql/utils/params.py | 160 ++++ pystackql/utils/platform.py | 40 + requirements.txt | 29 +- run_async_server_tests | 3 - run_server_tests.py | 53 ++ run_tests | 11 - run_tests.ps1 | 11 - run_tests.py | 60 ++ server-status.sh | 20 + setup.py | 8 +- start-stackql-server.sh | 9 + stop-stackql-server.sh | 12 + tests/README.md | 158 ++++ tests/conftest.py | 132 +++ tests/creds/env_vars/.gitignore | 4 - tests/creds/keys/.gitignore | 4 - tests/pystackql_async_server_tests.py | 58 -- tests/pystackql_tests.py | 483 ----------- tests/test_async.py | 141 +++ tests/test_constants.py | 120 +++ tests/test_core.py | 143 ++++ tests/test_magic.py | 122 +++ tests/test_output_formats.py | 213 +++++ tests/test_params.py | 85 -- tests/test_query_execution.py | 270 ++++++ tests/test_server.py | 224 +++++ tests/test_server_magic.py | 120 +++ 52 files changed, 4080 insertions(+), 1791 deletions(-) create mode 100644 launch_venv.sh delete mode 100644 pystackql/_util.py create mode 100644 pystackql/core/__init__.py create mode 100644 pystackql/core/binary.py create mode 100644 pystackql/core/output.py create mode 100644 pystackql/core/query.py create mode 100644 pystackql/core/server.py create mode 100644 pystackql/core/stackql.py create mode 100644 pystackql/magic_ext/__init__.py rename pystackql/{base_stackql_magic.py => magic_ext/base.py} (72%) create mode 100644 pystackql/magic_ext/local.py create mode 100644 pystackql/magic_ext/server.py delete mode 100644 pystackql/stackql.py create mode 100644 pystackql/utils/__init__.py create mode 100644 pystackql/utils/auth.py create mode 100644 pystackql/utils/binary.py create mode 100644 pystackql/utils/download.py create mode 100644 pystackql/utils/helpers.py create mode 100644 pystackql/utils/package.py create mode 100644 pystackql/utils/params.py create mode 100644 pystackql/utils/platform.py delete mode 100644 run_async_server_tests create mode 100644 run_server_tests.py delete mode 100644 run_tests delete mode 100644 run_tests.ps1 create mode 100644 run_tests.py create mode 100644 server-status.sh create mode 100644 start-stackql-server.sh create mode 100644 stop-stackql-server.sh create mode 100644 tests/README.md create mode 100644 tests/conftest.py delete mode 100644 tests/creds/env_vars/.gitignore delete mode 100644 tests/creds/keys/.gitignore delete mode 100644 tests/pystackql_async_server_tests.py delete mode 100644 tests/pystackql_tests.py create mode 100644 tests/test_async.py create mode 100644 tests/test_constants.py create mode 100644 tests/test_core.py create mode 100644 tests/test_magic.py create mode 100644 tests/test_output_formats.py delete mode 100644 tests/test_params.py create mode 100644 tests/test_query_execution.py create mode 100644 tests/test_server.py create mode 100644 tests/test_server_magic.py diff --git a/.github/workflows/test.yaml b/.github/workflows/test.yaml index 09111f5..8bf5a0a 100644 --- a/.github/workflows/test.yaml +++ b/.github/workflows/test.yaml @@ -28,79 +28,51 @@ jobs: name: 'Run Tests on ${{matrix.os}} with Python ${{matrix.python-version}}' steps: - - uses: actions/checkout@v4.1.1 + - uses: actions/checkout@v4 - name: Set up Python ${{ matrix.python-version }} - uses: actions/setup-python@v5.1.0 + uses: actions/setup-python@v5 with: python-version: ${{ matrix.python-version }} - - name: Install dependencies from requirements.txt + - name: Upgrade pip shell: bash run: | python3 -m pip install --upgrade pip - pip install -r requirements.txt - # Windows - - name: Install psycopg - if: matrix.os == 'windows-latest' + - name: Install pystackql with all dependencies run: | - pip install psycopg[binary] - shell: powershell - # End Windows + pip install -e . - # Linux - - name: Install PostgreSQL dev libraries on Ubuntu - if: matrix.os == 'ubuntu-latest' + - name: Install test dependencies run: | - sudo apt-get update - pip install psycopg - # End Linux + pip install pytest>=6.2.5 pytest-cov>=2.12.0 nose>=1.3.7 - # macOS - - name: Install PostgreSQL dev libraries on macOS - if: matrix.os == 'macos-latest' + - name: Run non-server tests + env: + GITHUB_ACTIONS: 'true' run: | - brew install postgresql@14 - pip install psycopg - # End macOS + python3 run_tests.py + shell: bash - - name: Install pystackql + - name: Setup StackQL + uses: stackql/setup-stackql@v2 + with: + platform: ${{ matrix.os == 'windows-latest' && 'windows' || (matrix.os == 'macos-latest' && 'darwin' || 'linux') }} + + - name: Start StackQL server run: | - pip install . + sh start-stackql-server.sh + shell: bash - - name: Run tests + - name: Run server tests 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 }} + GITHUB_ACTIONS: 'true' run: | - python3 -m tests.pystackql_tests + python3 run_server_tests.py shell: bash - if: matrix.os != 'windows-latest' - - name: Run tests on Windows - 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 }} + - name: Stop StackQL server run: | - python3 -m tests.pystackql_tests - shell: pwsh - if: matrix.os == 'windows-latest' + sh stop-stackql-server.sh + shell: bash \ No newline at end of file diff --git a/.gitignore b/.gitignore index 4e6ad3e..173aff7 100644 --- a/.gitignore +++ b/.gitignore @@ -1,11 +1,21 @@ # Miscellaneous -/.stackql -/.stackql/* .pypirc - /.vscode /.vscode/* +# Virtual environments +.venv/ +venv/ +env/ +ENV/ + +# stackql +.stackql/ +stackql +stackql-*.sh +.env +nohup.out + # Byte-compiled / optimized / DLL files __pycache__ *.py[cod] diff --git a/CHANGELOG.md b/CHANGELOG.md index 2c45e53..6e2dd13 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,12 @@ # Changelog +## v3.8.0 (2025-06-04) + +### Updates + +- Refactor +- Enhanced test coverage + ## v3.7.2 (2024-11-19) ### Updates diff --git a/launch_venv.sh b/launch_venv.sh new file mode 100644 index 0000000..741d635 --- /dev/null +++ b/launch_venv.sh @@ -0,0 +1,140 @@ +#!/bin/bash +# launch_venv.sh - Script to create, set up and activate a Python virtual environment for PyStackQL + +# Use simpler code without colors when running with sh +if [ -n "$BASH_VERSION" ]; then + # Color definitions for bash + GREEN='\033[0;32m' + YELLOW='\033[1;33m' + BLUE='\033[0;34m' + RED='\033[0;31m' + NC='\033[0m' # No Color + + # Function to print colored text in bash + cecho() { + printf "%b%s%b\n" "$1" "$2" "$NC" + } +else + # No colors for sh + cecho() { + echo "$2" + } +fi + +# Default virtual environment name +VENV_NAME=".venv" +REQUIREMENTS_FILE="requirements.txt" + +# Function to check if command exists +command_exists() { + command -v "$1" >/dev/null 2>&1 +} + +# Banner +cecho "$BLUE" "=======================================" +cecho "$BLUE" " PyStackQL Development Environment " +cecho "$BLUE" "=======================================" +echo "" + +# Check for Python +if ! command_exists python3; then + cecho "$RED" "Error: Python 3 is not installed." + echo "Please install Python 3 and try again." + exit 1 +fi + +# Print Python version +cecho "$YELLOW" "Using Python:" +python3 --version +echo "" + +# Create virtual environment if it doesn't exist +if [ ! -d "$VENV_NAME" ]; then + cecho "$YELLOW" "Creating virtual environment in ${VENV_NAME}..." + python3 -m venv "$VENV_NAME" + if [ $? -ne 0 ]; then + cecho "$RED" "Error: Failed to create virtual environment." + exit 1 + fi + cecho "$GREEN" "Virtual environment created successfully." +else + cecho "$YELLOW" "Using existing virtual environment in ${VENV_NAME}" +fi + +# Determine the activate script based on OS +case "$OSTYPE" in + msys*|win*|cygwin*) + # Windows + ACTIVATE_SCRIPT="$VENV_NAME/Scripts/activate" + ;; + *) + # Unix-like (Linux, macOS) + ACTIVATE_SCRIPT="$VENV_NAME/bin/activate" + ;; +esac + +# Check if activation script exists +if [ ! -f "$ACTIVATE_SCRIPT" ]; then + cecho "$RED" "Error: Activation script not found at $ACTIVATE_SCRIPT" + echo "The virtual environment may be corrupt. Try removing the $VENV_NAME directory and running this script again." + exit 1 +fi + +# Source the activation script +cecho "$YELLOW" "Activating virtual environment..." +. "$ACTIVATE_SCRIPT" +if [ $? -ne 0 ]; then + cecho "$RED" "Error: Failed to activate virtual environment." + exit 1 +fi + +# Install/upgrade pip, setuptools, and wheel +cecho "$YELLOW" "Upgrading pip, setuptools, and wheel..." +pip install --upgrade pip setuptools wheel +if [ $? -ne 0 ]; then + cecho "$RED" "Warning: Failed to upgrade pip, setuptools, or wheel. Continuing anyway." +fi + +# Check if requirements.txt exists +if [ ! -f "$REQUIREMENTS_FILE" ]; then + cecho "$RED" "Error: $REQUIREMENTS_FILE not found." + echo "Please make sure the file exists in the current directory." + cecho "$YELLOW" "Continuing with an activated environment without installing dependencies." +else + # Install requirements + cecho "$YELLOW" "Installing dependencies from $REQUIREMENTS_FILE..." + pip install -r "$REQUIREMENTS_FILE" + if [ $? -ne 0 ]; then + cecho "$RED" "Warning: Some dependencies may have failed to install." + else + cecho "$GREEN" "Dependencies installed successfully." + fi +fi + +# Install the package in development mode if setup.py exists +if [ -f "setup.py" ]; then + cecho "$YELLOW" "Installing PyStackQL in development mode..." + pip install . + if [ $? -ne 0 ]; then + cecho "$RED" "Warning: Failed to install package in development mode." + else + cecho "$GREEN" "Package installed in development mode." + fi +fi + +# Success message +echo "" +cecho "$GREEN" "Virtual environment is now set up and activated!" +cecho "$YELLOW" "You can use PyStackQL and run tests." +echo "" +cecho "$BLUE" "To run tests:" +echo " python run_tests.py" +echo "" +cecho "$BLUE" "To deactivate the virtual environment when done:" +echo " deactivate" +echo "" +cecho "$BLUE" "=======================================" + +# Keep the terminal open with the activated environment +# The script will be source'd, so the environment stays active +exec "${SHELL:-bash}" \ No newline at end of file diff --git a/pystackql/__init__.py b/pystackql/__init__.py index 87af198..74d3d3f 100644 --- a/pystackql/__init__.py +++ b/pystackql/__init__.py @@ -1,3 +1,17 @@ -from .stackql import StackQL -from .magic import StackqlMagic -from .magics import StackqlServerMagic \ No newline at end of file +# pystackql/__init__.py + +""" +PyStackQL - Python wrapper for StackQL + +This package provides a Python interface to the StackQL query language +for cloud resource querying. +""" + +# Import the core StackQL class +from .core import StackQL + +# Import the magic classes for Jupyter integration +from .magic_ext import StackqlMagic, StackqlServerMagic + +# Define the public API +__all__ = ['StackQL', 'StackqlMagic', 'StackqlServerMagic'] \ No newline at end of file diff --git a/pystackql/_util.py b/pystackql/_util.py deleted file mode 100644 index 086e421..0000000 --- a/pystackql/_util.py +++ /dev/null @@ -1,165 +0,0 @@ -import subprocess, platform, json, site, os, requests, zipfile - -# Conditional import for package metadata retrieval -try: - from importlib.metadata import version, PackageNotFoundError -except ImportError: - # This is for Python versions earlier than 3.8 - from importlib_metadata import version, PackageNotFoundError - -def _is_binary_local(platform): - """Checks if the binary exists at the specified local path.""" - if platform == 'Linux' and os.path.exists('/usr/local/bin/stackql'): - return True - return False - -def _get_package_version(package_name): - try: - pkg_version = version(package_name) - if pkg_version is None: - print(f"Warning: Retrieved version for '{package_name}' is None!") - return pkg_version - except PackageNotFoundError: - print(f"Warning: Package '{package_name}' not found!") - return None - -def _get_platform(): - system_val = platform.system() - machine_val = platform.machine() - platform_val = platform.platform() - python_version_val = platform.python_version() - return "%s %s (%s), Python %s" % (system_val, machine_val, platform_val, python_version_val), system_val - -def _get_download_dir(): - # check if site.getuserbase() dir exists - if not os.path.exists(site.getuserbase()): - # if not, create it - os.makedirs(site.getuserbase()) - return site.getuserbase() - -def _get_binary_name(platform): - if platform.startswith('Windows'): - return r'stackql.exe' - elif platform.startswith('Darwin'): - return r'stackql/Payload/stackql' - else: - return r'stackql' - -def _get_url(): - system_val = platform.system() - machine_val = platform.machine() - - if system_val == 'Linux' and machine_val == 'x86_64': - return 'https://releases.stackql.io/stackql/latest/stackql_linux_amd64.zip' - elif system_val == 'Windows': - return 'https://releases.stackql.io/stackql/latest/stackql_windows_amd64.zip' - elif system_val == 'Darwin': - return 'https://storage.googleapis.com/stackql-public-releases/latest/stackql_darwin_multiarch.pkg' - else: - raise Exception(f"ERROR: [_get_url] unsupported OS type: {system_val} {machine_val}") - -def _download_file(url, path, showprogress=True): - try: - r = requests.get(url, stream=True) - r.raise_for_status() - total_size_in_bytes = int(r.headers.get('content-length', 0)) - block_size = 1024 - with open(path, 'wb') as f: - chunks = 0 - for data in r.iter_content(block_size): - chunks += 1 - f.write(data) - downloaded_size = chunks * block_size - progress_bar = '#' * int(downloaded_size / total_size_in_bytes * 20) - if showprogress: - print(f'\r[{progress_bar.ljust(20)}] {int(downloaded_size / total_size_in_bytes * 100)}%', end='') - - print("\nDownload complete.") - except Exception as e: - 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) # Should return 'stackql.exe' for Windows - url = _get_url() - 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) - - # Handle extraction - if platform.startswith('Darwin'): - unpacked_file_name = os.path.join(download_dir, 'stackql') - command = f'pkgutil --expand-full {archive_file_name} {unpacked_file_name}' - if os.path.exists(unpacked_file_name): - os.system(f'rm -rf {unpacked_file_name}') - os.system(command) - - else: # Handle Windows and Linux - with zipfile.ZipFile(archive_file_name, 'r') as zip_ref: - 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(f"ERROR: [_setup] {str(e)}") - exit(1) - -def _get_version(bin_path): - try: - iqlPopen = subprocess.Popen([bin_path] + ["--version"], - stdout=subprocess.PIPE, stderr=subprocess.STDOUT) - # Use communicate to fetch the outputs and wait for the process to finish - output, _ = iqlPopen.communicate() - # Decode the output - decoded_output = output.decode('utf-8') - # Split to get the version tokens - version_tokens = decoded_output.split('\n')[0].split(' ') - version = version_tokens[1] - sha = version_tokens[3].replace('(', '').replace(')', '') - return version, sha - except FileNotFoundError: - print("ERROR: [_get_version] %s not found" % (bin_path)) - exit(1) - except Exception as e: - error_message = e.args[0] - print("ERROR: [_get_version] %s" % (error_message)) - exit(1) - finally: - # Ensure the subprocess is terminated and streams are closed - iqlPopen.terminate() - iqlPopen.stdout.close() - -def _format_auth(auth): - try: - if auth is not None: - if isinstance(auth, str): - authobj = json.loads(auth) - authstr = auth - elif isinstance(auth, dict): - authobj = auth - authstr = json.dumps(auth) - return authobj, authstr - else: - raise Exception("ERROR: [_format_auth] auth key supplied with no value") - except Exception as e: - error_message = e.args[0] - print("ERROR: [_format_auth] %s" % (error_message)) - exit(1) diff --git a/pystackql/core/__init__.py b/pystackql/core/__init__.py new file mode 100644 index 0000000..b1231b2 --- /dev/null +++ b/pystackql/core/__init__.py @@ -0,0 +1,16 @@ +# pystackql/core/__init__.py + +""" +Core functionality for PyStackQL. + +This module provides the core functionality for the PyStackQL package, +including the main StackQL class. +""" + +from .binary import BinaryManager +from .server import ServerConnection +from .query import QueryExecutor, AsyncQueryExecutor +from .output import OutputFormatter +from .stackql import StackQL + +__all__ = ['StackQL'] \ No newline at end of file diff --git a/pystackql/core/binary.py b/pystackql/core/binary.py new file mode 100644 index 0000000..023f671 --- /dev/null +++ b/pystackql/core/binary.py @@ -0,0 +1,79 @@ +# pystackql/core/binary.py + +""" +Binary management module for PyStackQL. + +This module handles the installation, version checking, and management +of the StackQL binary executable. +""" + +import os +from ..utils import ( + is_binary_local, + get_platform, + get_download_dir, + get_binary_name, + setup_binary, + get_binary_version +) + +class BinaryManager: + """Manages the StackQL binary installation and versions. + + This class is responsible for ensuring the StackQL binary is available + and correctly configured for use. + """ + + def __init__(self, download_dir=None): + """Initialize the BinaryManager. + + Args: + download_dir (str, optional): Directory to store the binary. Defaults to None. + """ + self.platform_info, self.system = get_platform() + + # Determine binary location + if self.system == 'Linux' and is_binary_local(self.system) and download_dir is None: + self.bin_path = '/usr/local/bin/stackql' + self.download_dir = '/usr/local/bin' + else: + # Use provided download_dir or default + self.download_dir = download_dir if download_dir else get_download_dir() + self.bin_path = os.path.join(self.download_dir, get_binary_name(self.system)) + + # Check if binary exists and get version + self._ensure_binary_exists() + + def _ensure_binary_exists(self): + """Ensure the binary exists, download it if not.""" + if os.path.exists(self.bin_path): + # Binary exists, get version + self.version, self.sha = get_binary_version(self.bin_path) + else: + # Binary doesn't exist, download it + setup_binary(self.download_dir, self.system) + self.version, self.sha = get_binary_version(self.bin_path) + + def upgrade(self, showprogress=True): + """Upgrade the StackQL binary to the latest version. + + Args: + showprogress (bool, optional): Whether to show download progress. Defaults to True. + + Returns: + str: A message indicating the new version + """ + setup_binary(self.download_dir, self.system, showprogress) + self.version, self.sha = get_binary_version(self.bin_path) + return f"stackql upgraded to version {self.version}" + + def get_version_info(self): + """Get the version information for the binary. + + Returns: + dict: Version information including version and sha + """ + return { + "version": self.version, + "sha": self.sha + } \ No newline at end of file diff --git a/pystackql/core/output.py b/pystackql/core/output.py new file mode 100644 index 0000000..26387eb --- /dev/null +++ b/pystackql/core/output.py @@ -0,0 +1,161 @@ +# pystackql/core/output.py + +""" +Output formatting module for PyStackQL. + +This module handles the formatting of query results into different output formats. +""" + +import json +from io import StringIO + +class OutputFormatter: + """Formats query results into different output formats. + + This class is responsible for converting raw query results into + the desired output format (dict, pandas, or csv). + """ + + def __init__(self, output_format='dict'): + """Initialize the OutputFormatter. + + Args: + output_format (str, optional): The output format. Defaults to 'dict'. + Allowed values: 'dict', 'pandas', 'csv' + + Raises: + ValueError: If an invalid output format is specified + """ + ALLOWED_OUTPUTS = {'dict', 'pandas', 'csv'} + if output_format.lower() not in ALLOWED_OUTPUTS: + raise ValueError(f"Invalid output format. Expected one of {ALLOWED_OUTPUTS}, got {output_format}.") + self.output_format = output_format.lower() + + def format_query_result(self, result, suppress_errors=True): + """Format a query result. + + Args: + result (dict): The raw query result from the executor + suppress_errors (bool, optional): Whether to suppress errors. Defaults to True. + + Returns: + The formatted result in the specified output format + """ + # Handle exceptions + if "exception" in result: + exception_msg = result["exception"] + return self._format_exception(exception_msg) + + # Handle data + if "data" in result: + data = result["data"] + return self._format_data(data) + + # Handle errors + if "error" in result and not suppress_errors: + err_msg = result["error"] + return self._format_error(err_msg) + + # No data, no error, return empty result + return self._format_empty() + + def _format_exception(self, exception_msg): + """Format an exception message. + + Args: + exception_msg (str): The exception message + + Returns: + The formatted exception in the specified output format + """ + if self.output_format == 'pandas': + import pandas as pd + return pd.DataFrame({'error': [exception_msg]}) if exception_msg else pd.DataFrame({'error': []}) + elif self.output_format == 'csv': + return exception_msg + else: # dict + return [{"error": exception_msg}] + + def _format_error(self, error_msg): + """Format an error message. + + Args: + error_msg (str): The error message + + Returns: + The formatted error in the specified output format + """ + if self.output_format == 'pandas': + import pandas as pd + return pd.DataFrame({'error': [error_msg]}) if error_msg else pd.DataFrame({'error': []}) + elif self.output_format == 'csv': + return error_msg + else: # dict + return [{"error": error_msg}] + + def _format_data(self, data): + """Format data. + + Args: + data (str): The data string + + Returns: + The formatted data in the specified output format + """ + if self.output_format == 'csv': + return data + elif self.output_format == 'pandas': + import pandas as pd + try: + return pd.read_json(StringIO(data)) + except ValueError: + return pd.DataFrame([{"error": "Invalid JSON output"}]) + else: # dict + try: + retval = json.loads(data) + return retval if retval else [] + except ValueError: + return [{"error": f"Invalid JSON output : {data}"}] + + def _format_empty(self): + """Format an empty result. + + Returns: + An empty result in the specified output format + """ + if self.output_format == 'pandas': + import pandas as pd + return pd.DataFrame() + elif self.output_format == 'csv': + return "" + else: # dict + return [] + + def format_statement_result(self, result): + """Format a statement result. + + Args: + result (dict): The raw statement result from the executor + + Returns: + The formatted result in the specified output format + """ + # Handle exceptions + if "exception" in result: + exception_msg = result["exception"] + return self._format_exception(exception_msg) + + # Message on stderr or empty message + message = result.get("error", "") + + if self.output_format == 'pandas': + import pandas as pd + return pd.DataFrame({'message': [message]}) if message else pd.DataFrame({'message': []}) + elif self.output_format == 'csv': + return message + else: # dict + # Count number of rows in the message + try: + return {'message': message, 'rowsaffected': message.count('\n')} + except Exception: + return {'message': message, 'rowsaffected': 0} \ No newline at end of file diff --git a/pystackql/core/query.py b/pystackql/core/query.py new file mode 100644 index 0000000..5cf2b28 --- /dev/null +++ b/pystackql/core/query.py @@ -0,0 +1,219 @@ +# pystackql/core/query.py + +""" +Query execution module for PyStackQL. + +This module handles the execution of StackQL queries via the binary or server. +""" + +import json +import os +import shlex +import subprocess +import tempfile +from io import StringIO +import asyncio +from concurrent.futures import ThreadPoolExecutor + +class QueryExecutor: + """Executes StackQL queries using a subprocess. + + This class is responsible for executing StackQL queries using either + a local binary or a server connection. + """ + + def __init__(self, binary_path, params=None, debug=False, debug_log_file=None): + """Initialize the QueryExecutor. + + Args: + binary_path (str): Path to the StackQL binary + params (list, optional): Additional parameters for the binary. Defaults to None. + debug (bool, optional): Whether to enable debug logging. Defaults to False. + debug_log_file (str, optional): Path to debug log file. Defaults to None. + """ + self.bin_path = binary_path + self.params = params or [] + self.debug = debug + self.debug_log_file = debug_log_file + + # Determine platform for command formatting + import platform + self.platform = platform.system() + + def _debug_log(self, message): + """Log a debug message. + + Args: + message (str): The message to log + """ + if self.debug and self.debug_log_file: + with open(self.debug_log_file, "a") as log_file: + log_file.write(message + "\n") + + def execute(self, query, custom_auth=None, env_vars=None): + """Execute a StackQL query. + + Args: + query (str): The query to execute + custom_auth (dict, optional): Custom authentication dictionary. Defaults to None. + env_vars (dict, optional): Environment variables for the subprocess. Defaults to None. + + Returns: + dict: The query results + """ + local_params = self.params.copy() + script_path = None + + # Format query for platform + if self.platform.startswith("Windows"): + # Escape double quotes and wrap in double quotes for Windows + escaped_query = query.replace('"', '\\"') + safe_query = f'"{escaped_query}"' + else: + # Use shlex.quote for Unix-like systems + safe_query = shlex.quote(query) + + local_params.insert(1, safe_query) + + # 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 + authstr = json.dumps(custom_auth) + local_params.extend(["--auth", f"'{authstr}'"]) + + output = {} + env_command_prefix = "" + + # Determine platform and set environment command prefix accordingly + if env_vars: + if self.platform.startswith("Windows"): + 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) + + try: + # Replace newlines to ensure command works in shell + full_command = full_command.replace("\n", " ") + + # Execute the command + result = subprocess.run( + full_command, + shell=True, + text=True, + capture_output=True + ) + + stdout = result.stdout + stderr = result.stderr + returncode = result.returncode + + # Log debug information if enabled + if self.debug: + self._debug_log(f"fullcommand: {full_command}") + self._debug_log(f"returncode: {returncode}") + self._debug_log(f"stdout: {stdout}") + self._debug_log(f"stderr: {stderr}") + + # Process stdout and stderr + if stderr: + output["error"] = stderr.decode('utf-8') if isinstance(stderr, bytes) else str(stderr) + if stdout: + output["data"] = stdout.decode('utf-8') if isinstance(stdout, bytes) else str(stdout) + + except FileNotFoundError: + output["exception"] = f"ERROR: {self.bin_path} not found" + except Exception as e: + error_details = { + "exception": str(e), + "doc": e.__doc__, + "params": local_params, + "stdout": stdout.decode('utf-8') if 'stdout' in locals() and isinstance(stdout, bytes) else "", + "stderr": stderr.decode('utf-8') if 'stderr' in locals() and isinstance(stderr, bytes) else "" + } + output["exception"] = f"ERROR: {json.dumps(error_details)}" + finally: + # Clean up the temporary script file + if script_path is not None: + os.remove(script_path) + return output + + +class AsyncQueryExecutor: + """Executes StackQL queries asynchronously. + + This class provides methods for executing multiple StackQL queries + concurrently using asyncio. + """ + + def __init__(self, sync_query_func, server_mode=False, output_format='dict'): + """Initialize the AsyncQueryExecutor. + + Args: + sync_query_func (callable): Function to execute a single query synchronously + server_mode (bool, optional): Whether to use server mode. Defaults to False. + output_format (str, optional): Output format (dict or pandas). Defaults to 'dict'. + """ + self.sync_query_func = sync_query_func + self.server_mode = server_mode + self.output_format = output_format + + async def execute_queries(self, queries): + """Execute multiple queries asynchronously. + + Args: + queries (list): List of query strings to execute + + Returns: + list or DataFrame: Results of all queries + + Raises: + ValueError: If output_format is not supported + """ + if self.output_format not in ['dict', 'pandas']: + raise ValueError("executeQueriesAsync supports only 'dict' or 'pandas' output modes.") + + async def main(): + with ThreadPoolExecutor() as executor: + # New connection is created for each query in server_mode, reused otherwise + new_connection = self.server_mode + + # Create tasks for each query + loop = asyncio.get_event_loop() + futures = [ + loop.run_in_executor( + executor, + lambda q=query: self.sync_query_func(q, new_connection), + # Pass query as a default argument to avoid late binding issues + ) + for query in queries + ] + + # Gather results from all the async calls + results = await asyncio.gather(*futures) + + return results + + # Run the async function and process results + results = await main() + + # Process results based on output format + if self.output_format == 'pandas': + import pandas as pd + return pd.concat(results, ignore_index=True) + else: + # Flatten the list of results + return [item for sublist in results for item in sublist] \ No newline at end of file diff --git a/pystackql/core/server.py b/pystackql/core/server.py new file mode 100644 index 0000000..c42d5b5 --- /dev/null +++ b/pystackql/core/server.py @@ -0,0 +1,164 @@ +# pystackql/core/server.py + +""" +Server connection management for PyStackQL. + +This module handles connections to a StackQL server using the Postgres wire protocol. +""" + +class ServerConnection: + """Manages connections to a StackQL server. + + This class handles connecting to and querying a StackQL server + using the Postgres wire protocol. + """ + + def __init__(self, server_address='127.0.0.1', server_port=5466): + """Initialize the ServerConnection. + + Args: + server_address (str, optional): Address of the server. Defaults to '127.0.0.1'. + server_port (int, optional): Port of the server. Defaults to 5466. + """ + self.server_address = server_address + self.server_port = server_port + self._conn = None + + # Import psycopg on demand to avoid dependency issues + try: + global psycopg, dict_row + import psycopg + from psycopg.rows import dict_row + except ImportError: + raise ImportError("psycopg is required in server mode but is not installed. " + "Please install psycopg and try again.") + + # Connect to the server + self._connect() + + def _connect(self): + """Connect to the StackQL server. + + Returns: + bool: True if connection successful, False otherwise + """ + try: + self._conn = psycopg.connect( + dbname='stackql', + user='stackql', + host=self.server_address, + port=self.server_port, + autocommit=True, + row_factory=dict_row + ) + return True + except psycopg.OperationalError as oe: + print(f"OperationalError while connecting to the server: {oe}") + except Exception as e: + print(f"Unexpected error while connecting to the server: {e}") + return False + + def is_connected(self): + """Check if the connection to the server is active. + + Returns: + bool: True if connected, False otherwise + """ + return self._conn is not None and not self._conn.closed + + def ensure_connected(self): + """Ensure the connection to the server is active. + + If the connection is closed, attempt to reconnect. + + Returns: + bool: True if connected, False otherwise + """ + if not self.is_connected(): + return self._connect() + return True + + def execute_query(self, query, is_statement=False): + """Execute a query on the server. + + Args: + query (str): The query to execute + is_statement (bool, optional): Whether this is a statement (non-SELECT). Defaults to False. + + Returns: + list: Results of the query as a list of dictionaries + + Raises: + ConnectionError: If no active connection is available + """ + if not self.ensure_connected(): + raise ConnectionError("No active connection to the server") + + try: + with self._conn.cursor() as cur: + cur.execute(query) + if is_statement: + # Return status message for non-SELECT statements + result_msg = cur.statusmessage + return [{'message': result_msg}] + try: + # Fetch results for SELECT queries + rows = cur.fetchall() + return rows + except psycopg.ProgrammingError as e: + # Handle cases with no results + if "no results to fetch" in str(e): + return [] + else: + raise + except psycopg.OperationalError as oe: + print(f"OperationalError during query execution: {oe}") + # Try to reconnect and retry once + if self._connect(): + return self.execute_query(query, is_statement) + except Exception as e: + print(f"Unexpected error during query execution: {e}") + + return [] + + def execute_query_with_new_connection(self, query): + """Execute a query with a new connection. + + This method creates a new connection to the server, executes the query, + and then closes the connection. + + Args: + query (str): The query to execute + + Returns: + list: Results of the query as a list of dictionaries + """ + try: + with psycopg.connect( + dbname='stackql', + user='stackql', + host=self.server_address, + port=self.server_port, + row_factory=dict_row + ) as conn: + with conn.cursor() as cur: + cur.execute(query) + try: + rows = cur.fetchall() + except psycopg.ProgrammingError as e: + if str(e) == "no results to fetch": + rows = [] + else: + raise + return rows + except psycopg.OperationalError as oe: + print(f"OperationalError while connecting to the server: {oe}") + except Exception as e: + print(f"Unexpected error while connecting to the server: {e}") + + return [] + + def close(self): + """Close the connection to the server.""" + if self._conn and not self._conn.closed: + self._conn.close() \ No newline at end of file diff --git a/pystackql/core/stackql.py b/pystackql/core/stackql.py new file mode 100644 index 0000000..2a681d5 --- /dev/null +++ b/pystackql/core/stackql.py @@ -0,0 +1,435 @@ +# pystackql/core/stackql.py + +""" +Main StackQL class for PyStackQL. + +This module provides the main StackQL class that serves as the primary +interface for executing StackQL queries. +""" + +import os +import json +from .server import ServerConnection +from .query import QueryExecutor, AsyncQueryExecutor +from .output import OutputFormatter +from ..utils import setup_local_mode + +class StackQL: + """A class representing an instance of the StackQL query engine. + + :param server_mode: Connect to a StackQL server + (defaults to `False`) + :type server_mode: bool, optional + :param server_address: The address of the StackQL server + (`server_mode` only, defaults to `'127.0.0.1'`) + :type server_address: str, optional + :param server_port: The port of the StackQL server + (`server_mode` only, defaults to `5444`) + :type server_port: int, optional + :param backend_storage_mode: Specifies backend storage mode, options are 'memory' and 'file' + (defaults to `'memory'`, this option is ignored in `server_mode`) + :type backend_storage_mode: str, optional + :param backend_file_storage_location: Specifies location for database file, only applicable when `backend_storage_mode` is 'file' + (defaults to `'{cwd}/stackql.db'`, this option is ignored in `server_mode`) + :type backend_file_storage_location: str, optional + :param output: Determines the format of the output, options are 'dict', 'pandas', and 'csv' + (defaults to `'dict'`, `'csv'` is not supported in `server_mode`) + :type output: str, optional + :param sep: Seperator for values in CSV output + (defaults to `','`, `output='csv'` only) + :type sep: str, optional + :param header: Show column headers in CSV output + (defaults to `False`, `output='csv'` only) + :type header: bool, optional + :param download_dir: The download directory for the StackQL executable + (defaults to `site.getuserbase()`, not supported in `server_mode`) + :type download_dir: str, optional + :param app_root: Application config and cache root path + (defaults to `{cwd}/.stackql`) + :type app_root: str, optional + :param execution_concurrency_limit: Concurrency limit for query execution + (defaults to `-1` - unlimited) + :type execution_concurrency_limit: int, optional + :param dataflow_dependency_max: Max dataflow weakly connected components for a given query + (defaults to `50`) + :type dataflow_dependency_max: int, optional + :param dataflow_components_max: Max dataflow components for a given query + (defaults to `50`) + :type dataflow_components_max: int, optional + :param api_timeout: API timeout + (defaults to `45`, not supported in `server_mode`) + :type api_timeout: int, optional + :param proxy_host: HTTP proxy host + (not supported in `server_mode`) + :type proxy_host: str, optional + :param proxy_password: HTTP proxy password + (only applicable when `proxy_host` is set) + :type proxy_password: str, optional + :param proxy_port: HTTP proxy port + (defaults to `-1`, only applicable when `proxy_host` is set) + :type proxy_port: int, optional + :param proxy_scheme: HTTP proxy scheme + (defaults to `'http'`, only applicable when `proxy_host` is set) + :type proxy_scheme: str, optional + :param proxy_user: HTTP proxy user + (only applicable when `proxy_host` is set) + :type proxy_user: str, optional + :param max_results: Max results per HTTP request + (defaults to `-1` for no limit, not supported in `server_mode`) + :type max_results: int, optional + :param page_limit: Max pages of results that will be returned per resource + (defaults to `20`, not supported in `server_mode`) + :type page_limit: int, optional + :param max_depth: Max depth for indirect queries: views and subqueries + (defaults to `5`, not supported in `server_mode`) + :type max_depth: int, optional + :param custom_registry: Custom StackQL provider registry URL + (e.g. https://registry-dev.stackql.app/providers) supplied using the class constructor + :type custom_registry: str, optional + :param custom_auth: Custom StackQL provider authentication object supplied using the class constructor + (not supported in `server_mode`) + :type custom_auth: dict, optional + :param debug: Enable debug logging + (defaults to `False`) + :type debug: bool, optional + :param debug_log_file: Path to debug log file + (defaults to `~/.pystackql/debug.log`, only available if debug is `True`) + :type debug_log_file: str, optional + + --- Read-Only Attributes --- + + :param platform: The operating system platform + :type platform: str, readonly + :param package_version: The version number of the `pystackql` Python package + :type package_version: str, readonly + :param version: The version number of the `stackql` executable + (not supported in `server_mode`) + :type version: str, readonly + :param params: A list of command-line parameters passed to the `stackql` executable + (not supported in `server_mode`) + :type params: list, readonly + :param bin_path: The full path of the `stackql` executable + (not supported in `server_mode`). + :type bin_path: str, readonly + :param sha: The commit (short) sha for the installed `stackql` binary build + (not supported in `server_mode`). + :type sha: str, readonly + """ + + def __init__(self, + server_mode=False, + server_address='127.0.0.1', + server_port=5444, + output='dict', + sep=',', + header=False, + debug=False, + debug_log_file=None, + **kwargs): + """Constructor method + """ + + # Get package information from utils + from ..utils import get_platform, get_package_version + self.platform, this_os = get_platform() + self.package_version = get_package_version("pystackql") + + # Setup debug logging + self.debug = debug + if debug: + if debug_log_file is None: + self.debug_log_file = os.path.join(os.path.expanduser("~"), '.pystackql', 'debug.log') + else: + self.debug_log_file = debug_log_file + # Check if the path exists. If not, try to create it. + log_dir = os.path.dirname(self.debug_log_file) + if not os.path.exists(log_dir): + try: + os.makedirs(log_dir, exist_ok=True) + except OSError as e: + raise ValueError(f"Unable to create the log directory {log_dir}: {str(e)}") + else: + self.debug_log_file = None + + # Setup output formatter + self.output_formatter = OutputFormatter(output) + self.output = output.lower() + + # Server mode setup + self.server_mode = server_mode + + if self.server_mode and self.output == 'csv': + raise ValueError("CSV output is not supported in server mode, use 'dict' or 'pandas' instead.") + elif self.output == 'csv': + self.sep = sep + self.header = header + + if self.server_mode: + # Server mode - connect to a server via the postgres wire protocol + self.server_address = server_address + self.server_port = server_port + self.server_connection = ServerConnection(server_address, server_port) + else: + # Local mode - execute the binary locally + # Get all parameters from local variables (excluding 'self') + local_params = locals().copy() + local_params.pop('self') + + # Set up local mode - this sets the instance attributes and returns params + self.params = setup_local_mode(self, **local_params) + + # Initialize query executor + self.query_executor = QueryExecutor( + self.bin_path, + self.params, + self.debug, + self.debug_log_file + ) + + # Initialize async query executor + self.async_executor = AsyncQueryExecutor( + self._sync_query_wrapper, + self.server_mode, + self.output + ) + + def _sync_query_wrapper(self, query, new_connection=False): + """Wrapper for synchronous query execution. + + Args: + query (str): The query to execute + new_connection (bool, optional): Whether to use a new connection. Defaults to False. + + Returns: + The query result + """ + if self.server_mode: + if new_connection: + result = self.server_connection.execute_query_with_new_connection(query) + else: + result = self.server_connection.execute_query(query) + + # Format server result if needed + if self.output == 'pandas': + import pandas as pd + return pd.DataFrame(result) + return result + else: + # Execute query and format result + query_result = self.query_executor.execute(query) + + if "exception" in query_result: + result = [{"error": query_result["exception"]}] + elif "error" in query_result: + result = [{"error": query_result["error"]}] + elif "data" in query_result: + try: + result = json.loads(query_result["data"]) + except Exception: + result = [{"error": f"Invalid JSON output: {query_result['data']}"}] + else: + result = [] + + # Format local result if needed + if self.output == 'pandas': + import pandas as pd + return pd.DataFrame(result) + return result + + def properties(self): + """Retrieves the properties of the StackQL instance. + + This method collects all the attributes of the StackQL instance and + returns them in a dictionary format. + + :return: A dictionary containing the properties of the StackQL instance. + :rtype: dict + + Example: + :: + + { + "platform": "Darwin x86_64 (macOS-12.0.1-x86_64-i386-64bit), Python 3.10.9", + "output": "dict", + ... + } + """ + props = {} + for var in vars(self): + # Skip internal objects + if var.startswith('_') or var in ['output_formatter', 'query_executor', 'async_executor', 'binary_manager', 'server_connection']: + continue + props[var] = getattr(self, var) + return props + + def upgrade(self, showprogress=True): + """Upgrades the StackQL binary to the latest version available. + + This method initiates an upgrade of the StackQL binary. Post-upgrade, + it updates the `version` and `sha` attributes of the StackQL instance + to reflect the newly installed version. + + :param showprogress: Indicates if progress should be displayed during the upgrade. Defaults to True. + :type showprogress: bool, optional + + :return: A message indicating the new version of StackQL post-upgrade. + :rtype: str + """ + if self.server_mode: + raise ValueError("The upgrade method is not supported in server mode.") + + # Use the binary manager to upgrade + message = self.binary_manager.upgrade(showprogress) + + # Update the version and sha attributes + self.version = self.binary_manager.version + self.sha = self.binary_manager.sha + + return message + + def executeStmt(self, query, custom_auth=None, env_vars=None): + """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 + 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, list of dict objects, or Pandas DataFrame + :param custom_auth: Custom authentication dictionary. + :type custom_auth: dict, optional + :param env_vars: Command-specific environment variables for this execution. + :type env_vars: dict, optional + + :return: The output result of the query in string format. If in `server_mode`, it + returns a JSON string representation of the result. + :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) + >>> result + """ + if self.server_mode: + result = self.server_connection.execute_query(query, is_statement=True) + + # Format result based on output type + if self.output == 'pandas': + import pandas as pd + return pd.DataFrame(result) + elif self.output == 'csv': + # Return the string representation of the result + return result[0]['message'] + else: + return result + else: + # Execute the query + result = self.query_executor.execute(query, custom_auth=custom_auth, env_vars=env_vars) + + # Format the result + return self.output_formatter.format_statement_result(result) + + def execute(self, query, suppress_errors=True, custom_auth=None, env_vars=None): + """ + Executes a StackQL query and returns the output based on the specified output format. + + This method supports execution both in server mode and locally using a subprocess. In server mode, + the query is sent to a StackQL server, while in local mode, it runs the query using a local binary. + + :param query: The StackQL query string to be executed. + :type query: str + :param suppress_errors: If set to True, the method will return an empty list if an error occurs. + :type suppress_errors: bool, optional + :param custom_auth: Custom authentication dictionary. + :type custom_auth: dict, optional + :param env_vars: Command-specific environment variables for this execution. + :type env_vars: dict, optional + + :return: The output of the query, which can be a list of dictionary objects, a Pandas DataFrame, + or a raw CSV string, depending on the configured output format. + :rtype: list(dict) | pd.DataFrame | str + + :raises ValueError: If an unsupported output format is specified. + + :example: + + >>> stackql = StackQL() + >>> query = ''' + ... SELECT SPLIT_PART(machineType, '/', -1) as machine_type, status, COUNT(*) as num_instances + ... FROM google.compute.instances + ... WHERE project = 'stackql-demo' AND zone = 'australia-southeast1-a' + ... GROUP BY machine_type, status HAVING COUNT(*) > 2 + ... ''' + >>> result = stackql.execute(query) + """ + if self.server_mode: + result = self.server_connection.execute_query(query) + + # Format result based on output type + if self.output == 'pandas': + import pandas as pd + import json + from io import StringIO + json_str = json.dumps(result) + return pd.read_json(StringIO(json_str)) + elif self.output == 'csv': + raise ValueError("CSV output is not supported in server_mode.") + else: # Assume 'dict' output + return result + else: + # Apply HTTP debug setting + if self.http_debug: + suppress_errors = False + + # Execute the query + output = self.query_executor.execute(query, custom_auth=custom_auth, env_vars=env_vars) + + # Format the result + return self.output_formatter.format_query_result(output, suppress_errors) + + async def executeQueriesAsync(self, queries): + """Executes multiple StackQL queries asynchronously using the current StackQL instance. + + This method utilizes an asyncio event loop to concurrently run a list of provided + StackQL queries. Each query is executed independently, and the combined results of + all the queries are returned as a list of JSON objects if 'dict' output mode is selected, + or as a concatenated DataFrame if 'pandas' output mode is selected. + + The order of the results in the returned list or DataFrame may not necessarily + correspond to the order of the queries in the input list due to the asynchronous nature + of execution. + + :param queries: A list of StackQL query strings to be executed concurrently. + :type queries: list[str], required + :return: A list of results corresponding to each query. Each result is a JSON object or a DataFrame. + :rtype: list[dict] or pd.DataFrame + :raises ValueError: If method is used in `server_mode` on an unsupported OS (anything other than Linux). + :raises ValueError: If an unsupported output mode is selected (anything other than 'dict' or 'pandas'). + + Example: + >>> from pystackql import StackQL + >>> stackql = StackQL() + >>> queries = [ + >>> \"\"\"SELECT '%s' as region, instanceType, COUNT(*) as num_instances + ... FROM aws.ec2.instances + ... WHERE region = '%s' + ... GROUP BY instanceType\"\"\" % (region, region) + >>> for region in regions ] + >>> result = stackql.executeQueriesAsync(queries) + + Note: + - When operating in `server_mode`, this method is not supported. + """ + if self.server_mode: + raise ValueError( + "The executeQueriesAsync method is not supported in server mode. " + "Please use the standard execute method with individual queries instead, " + "or switch to local mode if you need to run multiple queries concurrently." + ) + + return await self.async_executor.execute_queries(queries) \ No newline at end of file diff --git a/pystackql/magic.py b/pystackql/magic.py index 13f7432..8fa63b5 100644 --- a/pystackql/magic.py +++ b/pystackql/magic.py @@ -1,44 +1,10 @@ -# `%load_ext pystackql.magic` - loads the stackql magic with server_mode=False -from IPython.core.magic import (magics_class, line_cell_magic) -from .base_stackql_magic import BaseStackqlMagic -import argparse +# pystackql/magic.py -@magics_class -class StackqlMagic(BaseStackqlMagic): - def __init__(self, shell): - super().__init__(shell, server_mode=False) +""" +StackQL Jupyter magic extension (non-server mode). +""" +# Import and re-export the load_ipython_extension function +from .magic_ext.local import load_ipython_extension - @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 - -def load_ipython_extension(ipython): - """Load the non-server magic in IPython.""" - ipython.register_magics(StackqlMagic) +# For direct imports (though less common) +from .magic_ext.local import StackqlMagic \ No newline at end of file diff --git a/pystackql/magic_ext/__init__.py b/pystackql/magic_ext/__init__.py new file mode 100644 index 0000000..b3d503b --- /dev/null +++ b/pystackql/magic_ext/__init__.py @@ -0,0 +1,14 @@ +# pystackql/magic_ext/__init__.py + +""" +Jupyter magic extensions for PyStackQL. + +This module provides Jupyter magic commands for running StackQL queries +directly in Jupyter notebooks. +""" + +from .base import BaseStackqlMagic +from .local import StackqlMagic +from .server import StackqlServerMagic + +__all__ = ['BaseStackqlMagic', 'StackqlMagic', 'StackqlServerMagic'] \ No newline at end of file diff --git a/pystackql/base_stackql_magic.py b/pystackql/magic_ext/base.py similarity index 72% rename from pystackql/base_stackql_magic.py rename to pystackql/magic_ext/base.py index d109e39..7a647f9 100644 --- a/pystackql/base_stackql_magic.py +++ b/pystackql/magic_ext/base.py @@ -1,21 +1,29 @@ +# pystackql/magic_ext/base.py + +""" +Base Jupyter magic extension for PyStackQL. + +This module provides the base class for PyStackQL Jupyter magic extensions. +""" + from __future__ import print_function -from IPython.core.magic import (Magics) +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. 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 + or SaaS resources 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. + """Initialize the BaseStackqlMagic class. :param shell: The IPython shell instance. + :param server_mode: Whether to use server mode. """ - from . import StackQL + from ..core import StackQL super(BaseStackqlMagic, self).__init__(shell) self.stackql_instance = StackQL(server_mode=server_mode, output='pandas') @@ -42,4 +50,4 @@ def run_query(self, query): if query.strip().lower().startswith("registry pull"): return self.stackql_instance.executeStmt(query) - return self.stackql_instance.execute(query) + return self.stackql_instance.execute(query) \ No newline at end of file diff --git a/pystackql/magic_ext/local.py b/pystackql/magic_ext/local.py new file mode 100644 index 0000000..0c5952c --- /dev/null +++ b/pystackql/magic_ext/local.py @@ -0,0 +1,66 @@ +# pystackql/magic_ext/local.py + +""" +Local Jupyter magic extension for PyStackQL. + +This module provides a Jupyter magic command for running StackQL queries +using a local StackQL binary. +""" + +from IPython.core.magic import (magics_class, line_cell_magic) +from .base import BaseStackqlMagic +import argparse + +@magics_class +class StackqlMagic(BaseStackqlMagic): + """Jupyter magic command for running StackQL queries in local mode.""" + + def __init__(self, shell): + """Initialize the StackqlMagic class. + + :param shell: The IPython shell instance. + """ + super().__init__(shell, server_mode=False) + + @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 + +def load_ipython_extension(ipython): + """Load the non-server magic in IPython. + + This is called when running %load_ext pystackql.magic in a notebook. + It registers the %stackql and %%stackql magic commands. + + :param ipython: The IPython shell instance + """ + # Create an instance of the magic class and register it + magic_instance = StackqlMagic(ipython) + ipython.register_magics(magic_instance) \ No newline at end of file diff --git a/pystackql/magic_ext/server.py b/pystackql/magic_ext/server.py new file mode 100644 index 0000000..d7ccdcf --- /dev/null +++ b/pystackql/magic_ext/server.py @@ -0,0 +1,60 @@ +# pystackql/magic_ext/server.py + +""" +Server Jupyter magic extension for PyStackQL. + +This module provides a Jupyter magic command for running StackQL queries +using a StackQL server connection. +""" + +from IPython.core.magic import (magics_class, line_cell_magic) +from .base import BaseStackqlMagic +import argparse + +@magics_class +class StackqlServerMagic(BaseStackqlMagic): + """Jupyter magic command for running StackQL queries in server mode.""" + + def __init__(self, shell): + """Initialize the StackqlServerMagic class. + + :param shell: The IPython shell instance. + """ + super().__init__(shell, server_mode=True) + + @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 args.no_display: + return None + else: + return results + +def load_ipython_extension(ipython): + """Load the server magic in IPython.""" + # Create an instance of the magic class and register it + magic_instance = StackqlServerMagic(ipython) + ipython.register_magics(magic_instance) \ No newline at end of file diff --git a/pystackql/magics.py b/pystackql/magics.py index 6d69b68..2f52ac3 100644 --- a/pystackql/magics.py +++ b/pystackql/magics.py @@ -1,44 +1,10 @@ -# `%load_ext pystackql.magics` - loads the stackql magic with server_mode=True -from IPython.core.magic import (magics_class, line_cell_magic) -from .base_stackql_magic import BaseStackqlMagic -import argparse +# pystackql/magics.py -@magics_class -class StackqlServerMagic(BaseStackqlMagic): - def __init__(self, shell): - super().__init__(shell, server_mode=True) +""" +StackQL Jupyter magic extension (server mode). +""" +# Import and re-export the load_ipython_extension function +from .magic_ext.server import load_ipython_extension - @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 - -def load_ipython_extension(ipython): - """Load the extension in IPython.""" - ipython.register_magics(StackqlServerMagic) \ No newline at end of file +# For direct imports (though less common) +from .magic_ext.server import StackqlServerMagic \ No newline at end of file diff --git a/pystackql/stackql.py b/pystackql/stackql.py deleted file mode 100644 index 300be86..0000000 --- a/pystackql/stackql.py +++ /dev/null @@ -1,809 +0,0 @@ -from ._util import ( - _get_package_version, - _get_platform, - _get_download_dir, - _get_binary_name, - _is_binary_local, - _setup, - _get_version, - _format_auth -) -import sys, subprocess, json, os, asyncio, functools -from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor -import pandas as pd -import tempfile - -from io import StringIO - -class StackQL: - """A class representing an instance of the StackQL query engine. - - :param server_mode: Connect to a StackQL server - (defaults to `False`) - :type server_mode: bool, optional - :param server_address: The address of the StackQL server - (`server_mode` only, defaults to `'127.0.0.1'`) - :type server_address: str, optional - :param server_port: The port of the StackQL server - (`server_mode` only, defaults to `5466`) - :type server_port: int, optional - :param backend_storage_mode: Specifies backend storage mode, options are 'memory' and 'file' - (defaults to `'memory'`, this option is ignored in `server_mode`) - :type backend_storage_mode: str, optional - :param backend_file_storage_location: Specifies location for database file, only applicable when `backend_storage_mode` is 'file' - (defaults to `'{cwd}/stackql.db'`, this option is ignored in `server_mode`) - :type backend_file_storage_location: str, optional - :param output: Determines the format of the output, options are 'dict', 'pandas', and 'csv' - (defaults to `'dict'`, `'csv'` is not supported in `server_mode`) - :type output: str, optional - :param sep: Seperator for values in CSV output - (defaults to `','`, `output='csv'` only) - :type sep: str, optional - :param header: Show column headers in CSV output - (defaults to `False`, `output='csv'` only) - :type header: bool, optional - :param download_dir: The download directory for the StackQL executable - (defaults to `site.getuserbase()`, not supported in `server_mode`) - :type download_dir: str, optional - :param app_root: Application config and cache root path - (defaults to `{cwd}/.stackql`) - :type app_root: str, optional - :param execution_concurrency_limit: Concurrency limit for query execution - (defaults to `-1` - unlimited) - :type execution_concurrency_limit: int, optional - :param dataflow_dependency_max: Max dataflow weakly connected components for a given query - (defaults to `50`) - :type dataflow_dependency_max: int, optional - :param dataflow_components_max: Max dataflow dependency depth for a given query - (defaults to `50`) - :type dataflow_components_max: int, optional - :param api_timeout: API timeout - (defaults to `45`, not supported in `server_mode`) - :type api_timeout: int, optional - :param proxy_host: HTTP proxy host - (not supported in `server_mode`) - :type proxy_host: str, optional - :param proxy_password: HTTP proxy password - (only applicable when `proxy_host` is set) - :type proxy_password: str, optional - :param proxy_port: HTTP proxy port - (defaults to `-1`, only applicable when `proxy_host` is set) - :type proxy_port: int, optional - :param proxy_scheme: HTTP proxy scheme - (defaults to `'http'`, only applicable when `proxy_host` is set) - :type proxy_scheme: str, optional - :param proxy_user: HTTP proxy user - (only applicable when `proxy_host` is set) - :type proxy_user: str, optional - :param max_results: Max results per HTTP request - (defaults to `-1` for no limit, not supported in `server_mode`) - :type max_results: int, optional - :param page_limit: Max pages of results that will be returned per resource - (defaults to `20`, not supported in `server_mode`) - :type page_limit: int, optional - :param max_depth: Max depth for indirect queries: views and subqueries - (defaults to `5`, not supported in `server_mode`) - :type max_depth: int, optional - :param custom_registry: Custom StackQL provider registry URL - (e.g. https://registry-dev.stackql.app/providers) supplied using the class constructor - :type custom_registry: str, optional - :param custom_auth: Custom StackQL provider authentication object supplied using the class constructor - (not supported in `server_mode`) - :type custom_auth: dict, optional - :param debug: Enable debug logging - (defaults to `False`) - :type debug: bool, optional - :param debug_log_file: Path to debug log file - (defaults to `~/.pystackql/debug.log`, only available if debug is `True`) - :type debug_log_file: str, optional - - --- Read-Only Attributes --- - - :param platform: The operating system platform - :type platform: str, readonly - :param package_version: The version number of the `pystackql` Python package - :type package_version: str, readonly - :param version: The version number of the `stackql` executable - (not supported in `server_mode`) - :type version: str, readonly - :param params: A list of command-line parameters passed to the `stackql` executable - (not supported in `server_mode`) - :type params: list, readonly - :param bin_path: The full path of the `stackql` executable - (not supported in `server_mode`). - :type bin_path: str, readonly - :param sha: The commit (short) sha for the installed `stackql` binary build - (not supported in `server_mode`). - :type sha: str, readonly - """ - - def _debug_log(self, message): - with open(self.debug_log_file, "a") as log_file: - log_file.write(message + "\n") - - def _connect_to_server(self): - """Establishes a connection to the server using psycopg. - - :return: Connection object if successful, or `None` if an error occurred. - :rtype: Connection or None - :raises `psycopg.OperationalError`: Failed to connect to the server. - """ - try: - conn = psycopg.connect( - dbname='stackql', - user='stackql', - host=self.server_address, - port=self.server_port, - autocommit=True, - row_factory=dict_row # Use dict_row to get rows as dictionaries - ) - return conn - except psycopg.OperationalError as oe: - print(f"OperationalError while connecting to the server: {oe}") - except Exception as e: - print(f"Unexpected error while connecting to the server: {e}") - return None - - def _run_server_query(self, query, is_statement=False): - """Run a query against the server using the existing connection in server mode.""" - if not self._conn: - raise ConnectionError("No active connection found. Ensure _connect_to_server is called.") - - try: - with self._conn.cursor() as cur: - cur.execute(query) - if is_statement: - # Return status message for non-SELECT statements - result_msg = cur.statusmessage - return [{'message': result_msg}] - try: - # Fetch results for SELECT queries - rows = cur.fetchall() - return rows - except psycopg.ProgrammingError as e: - # Handle cases with no results - if "no results to fetch" in str(e): - return [] - else: - raise - except psycopg.OperationalError as oe: - print(f"OperationalError during query execution: {oe}") - except Exception as e: - print(f"Unexpected error during query execution: {e}") - - - def _run_query(self, query, custom_auth=None, env_vars=None): - """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. - It waits for the subprocess to complete and captures its stdout as the output. This approach ensures - that resources like pipes are properly cleaned up after the subprocess completes. - - :param query: The StackQL query string to be executed. - :type query: str - :param custom_auth: Custom authentication dictionary. - :type custom_auth: dict, optional - :param env_vars: Command-specific environment variables for the subprocess. - :type env_vars: dict, optional - - :return: The output result of the query, which can either be the actual query result or an error message. - :rtype: dict - - Example: - :: - - { - "data": "[{\"machine_type\": \"n1-standard-1\", \"status\": \"RUNNING\", \"num_instances\": 3}, ...]", - "error": "stderr message if present", - "exception": "ERROR: {\"exception\": \"\", \"doc\": \"\", \"params\": \"\", \"stdout\": \"\", \"stderr\": \"\"} - } - - Possible error messages include: - - Indications that the StackQL binary wasn't found. - - Generic error messages for other exceptions encountered during the query execution. - - :raises FileNotFoundError: If the StackQL binary isn't found. - :raises Exception: For any other exceptions during the execution, providing a generic error message. - """ - - local_params = self.params.copy() - script_path = None - - if self.platform.startswith("Windows"): - # Escape double quotes and wrap in double quotes for Windows - escaped_query = query.replace('"', '\\"') # Escape double quotes properly - safe_query = f'"{escaped_query}"' - else: - # Use shlex.quote for Unix-like systems - import shlex - safe_query = shlex.quote(query) - - local_params.insert(1, safe_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 - authstr = json.dumps(custom_auth) - local_params.extend(["--auth", f"'{authstr}'"]) - - output = {} - env_command_prefix = "" - - # Determine platform and set environment command prefix accordingly - if env_vars: - if self.platform.startswith("Windows"): - 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) - - try: - - full_command = full_command.replace("\n", " ") - - result = subprocess.run( - full_command, - shell=True, - text=True, - capture_output=True - ) - - stdout = result.stdout - stderr = result.stderr - returncode = result.returncode - - if self.debug: - self._debug_log(f"fullcommand: {full_command}") - self._debug_log(f"returncode: {returncode}") - self._debug_log(f"stdout: {stdout}") - self._debug_log(f"stderr: {stderr}") - - # Process stdout and stderr - if stderr: - output["error"] = stderr.decode('utf-8') if isinstance(stderr, bytes) else str(stderr) - if stdout: - output["data"] = stdout.decode('utf-8') if isinstance(stdout, bytes) else str(stdout) - - - except FileNotFoundError: - output["exception"] = f"ERROR: {self.bin_path} not found" - except Exception as e: - error_details = { - "exception": str(e), - "doc": e.__doc__, - "params": local_params, - "stdout": stdout.decode('utf-8') if 'stdout' in locals() and isinstance(stdout, bytes) else "", - "stderr": stderr.decode('utf-8') if 'stderr' in locals() and isinstance(stderr, bytes) else "" - } - output["exception"] = f"ERROR: {json.dumps(error_details)}" - 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, - server_address='127.0.0.1', - server_port=5466, - backend_storage_mode='memory', - backend_file_storage_location='stackql.db', - download_dir=None, - app_root=None, - execution_concurrency_limit=-1, - dataflow_dependency_max=50, - dataflow_components_max=50, - output='dict', - custom_registry=None, - custom_auth=None, - sep=',', - header=False, - api_timeout=45, - proxy_host=None, - proxy_password=None, - proxy_port=-1, - proxy_scheme='http', - proxy_user=None, - max_results=-1, - page_limit=20, - max_depth=5, - debug=False, - http_debug=False, - debug_log_file=None): - """Constructor method - """ - # read only properties - self.platform, this_os = _get_platform() - self.package_version = _get_package_version("pystackql") - - # common constructor args - # Check and assign the output if it is allowed, else raise ValueError - ALLOWED_OUTPUTS = {'dict', 'pandas', 'csv'} - if output.lower() not in ALLOWED_OUTPUTS: - raise ValueError(f"Invalid output. Expected one of {ALLOWED_OUTPUTS}, got {output}.") - self.output = output.lower() - self.server_mode = server_mode - if self.server_mode and self.output == 'csv': - raise ValueError("CSV output is not supported in server mode, use 'dict' or 'pandas' instead.") - - self.debug = debug - - if debug: - if debug_log_file is None: - self.debug_log_file = os.path.join(os.path.expanduser("~"), '.pystackql', 'debug.log') - else: - self.debug_log_file = debug_log_file - # Check if the path exists. If not, try to create it. - log_dir = os.path.dirname(self.debug_log_file) - if not os.path.exists(log_dir): - try: - os.makedirs(log_dir, exist_ok=True) # exist_ok=True will not raise an error if the directory exists. - except OSError as e: - raise ValueError(f"Unable to create the log directory {log_dir}: {str(e)}") - - if self.server_mode: - # server mode, connect to a server via the postgres wire protocol - # Attempt to import psycopg only if server_mode is True - global psycopg, dict_row - try: - import psycopg - from psycopg.rows import dict_row # For returning results as dictionaries - except ImportError: - raise ImportError("psycopg is required in server mode but is not installed. Please install psycopg and try again.") - - self.server_address = server_address - self.server_port = server_port - # establish the connection - self._conn = self._connect_to_server() - else: - # local mode, executes the binary locally - self.params = [] - self.params.append("exec") - - self.params.append("--output") - if self.output == "csv": - self.params.append("csv") - else: - self.params.append("json") - - # backend storage settings - if backend_storage_mode == 'file': - self.params.append("--sqlBackend") - self.params.append(json.dumps({ "dsn": f"file:{backend_file_storage_location}" })) - - # get or download the stackql binary - binary = _get_binary_name(this_os) - - # check if the binary exists locally for Linux - if this_os == 'Linux' and _is_binary_local(this_os) and download_dir is None: - self.bin_path = '/usr/local/bin/stackql' - self.download_dir = '/usr/local/bin' - # get and set version - self.version, self.sha = _get_version(self.bin_path) - else: - # if download_dir not set, use site.getuserbase() or the provided path - if download_dir is None: - self.download_dir = _get_download_dir() - else: - self.download_dir = download_dir - self.bin_path = os.path.join(self.download_dir, binary) - # get and set version - if os.path.exists(self.bin_path): - self.version, self.sha = _get_version(self.bin_path) - else: - # not installed, download - _setup(self.download_dir, this_os) - self.version, self.sha = _get_version(self.bin_path) - - # if app_root is set, use it - if app_root is not None: - self.app_root = app_root - self.params.append("--approot") - self.params.append(app_root) - - # set execution_concurrency_limit - self.execution_concurrency_limit = execution_concurrency_limit - self.params.append("--execution.concurrency.limit") - self.params.append(str(execution_concurrency_limit)) - - # set dataflow_dependency_max and dataflow_components_max - self.dataflow_dependency_max = dataflow_dependency_max - self.params.append("--dataflow.dependency.max") - self.params.append(str(dataflow_dependency_max)) - - self.dataflow_components_max = dataflow_components_max - self.params.append("--dataflow.components.max") - self.params.append(str(dataflow_components_max)) - - # if custom_auth is set, use it - if custom_auth is not None: - authobj, authstr = _format_auth(custom_auth) - self.auth = authobj - self.params.append("--auth") - self.params.append(authstr) - - # if custom_registry is set, use it - if custom_registry is not None: - self.custom_registry = custom_registry - self.params.append("--registry") - self.params.append(json.dumps({ "url": custom_registry })) - - # csv output - if self.output == "csv": - self.sep = sep - self.params.append("--delimiter") - self.params.append(sep) - - self.header = header - if not self.header: - self.params.append("--hideheaders") - - # app behavioural properties - self.max_results = max_results - self.params.append("--http.response.maxResults") - self.params.append(str(max_results)) - - self.page_limit = page_limit - self.params.append("--http.response.pageLimit") - self.params.append(str(page_limit)) - - self.max_depth = max_depth - self.params.append("--indirect.depth.max") - self.params.append(str(max_depth)) - - self.api_timeout = api_timeout - self.params.append("--apirequesttimeout") - self.params.append(str(api_timeout)) - - if http_debug: - self.http_debug = True - self.params.append("--http.log.enabled") - else: - self.http_debug = False - - # proxy settings - if proxy_host is not None: - self.proxy_host = proxy_host - self.params.append("--http.proxy.host") - self.params.append(proxy_host) - - self.proxy_port = proxy_port - self.params.append("--http.proxy.port") - self.params.append(proxy_port) - - self.proxy_user = proxy_user - self.params.append("--http.proxy.user") - self.params.append(proxy_user) - - self.proxy_password = proxy_password - self.params.append("--http.proxy.password") - self.params.append(proxy_password) - - # Check and assign the proxy_scheme if it is allowed, else raise ValueError - ALLOWED_PROXY_SCHEMES = {'http', 'https'} - if proxy_scheme.lower() not in ALLOWED_PROXY_SCHEMES: - raise ValueError(f"Invalid proxy_scheme. Expected one of {ALLOWED_PROXY_SCHEMES}, got {proxy_scheme}.") - self.proxy_scheme = proxy_scheme.lower() - self.params.append("--http.proxy.scheme") - self.params.append(proxy_scheme) - - def properties(self): - """Retrieves the properties of the StackQL instance. - - This method collects all the attributes of the StackQL instance and - returns them in a dictionary format. - - :return: A dictionary containing the properties of the StackQL instance. - :rtype: dict - - Example: - :: - - { - "platform": "Darwin x86_64 (macOS-12.0.1-x86_64-i386-64bit), Python 3.10.9", - "parse_json": True, - ... - } - """ - props = {} - for var in vars(self): - props[var] = getattr(self, var) - return props - - def upgrade(self, showprogress=True): - """Upgrades the StackQL binary to the latest version available. - - This method initiates an upgrade of the StackQL binary. Post-upgrade, - it updates the `version` and `sha` attributes of the StackQL instance - to reflect the newly installed version. - - :param showprogress: Indicates if progress should be displayed during the upgrade. Defaults to True. - :type showprogress: bool, optional - - :return: A message indicating the new version of StackQL post-upgrade. - :rtype: str - """ - _setup(self.download_dir, self.platform, showprogress) - self.version, self.sha = _get_version(self.bin_path) - print("stackql upgraded to version %s" % (self.version)) - - def executeStmt(self, query, custom_auth=None, env_vars=None): - """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 - 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, list of dict objects, or Pandas DataFrame - :param custom_auth: Custom authentication dictionary. - :type custom_auth: dict, optional - :param env_vars: Command-specific environment variables for this execution. - :type env_vars: dict, optional - - :return: The output result of the query in string format. If in `server_mode`, it - returns a JSON string representation of the result. - :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) - >>> result - """ - if self.server_mode: - result = self._run_server_query(query, is_statement=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: - - # returns either... - # {'error': ''} if something went wrong; or - # {'message': ''} if the statement was executed successfully - - result = self._run_query(query, custom_auth=custom_auth, env_vars=env_vars) - - if "exception" in result: - exception_msg = result["exception"] - if self.output == 'pandas': - return pd.DataFrame({'error': [exception_msg]}) if exception_msg else pd.DataFrame({'error': []}) - elif self.output == 'csv': - return exception_msg - else: - return {"error": exception_msg} - - # message on stderr - message = result.get("error", "") - - if self.output == 'pandas': - return pd.DataFrame({'message': [message]}) if message else pd.DataFrame({'message': []}) - elif self.output == 'csv': - return message - else: - # count number of rows in the message - try: - return {'message': message, 'rowsaffected': message.count('\n')} - except Exception as e: - return {'message': message, 'rowsaffected': 0} - - def execute(self, query, suppress_errors=True, custom_auth=None, env_vars=None): - """ - Executes a StackQL query and returns the output based on the specified output format. - - This method supports execution both in server mode and locally using a subprocess. In server mode, - the query is sent to a StackQL server, while in local mode, it runs the query using a local binary. - - :param query: The StackQL query string to be executed. - :type query: str - :param suppress_errors: If set to True, the method will return an empty list if an error occurs. - :type suppress_errors: bool, optional - :param custom_auth: Custom authentication dictionary. - :type custom_auth: dict, optional - :param env_vars: Command-specific environment variables for this execution. - :type env_vars: dict, optional - - :return: The output of the query, which can be a list of dictionary objects, a Pandas DataFrame, - or a raw CSV string, depending on the configured output format. - :rtype: list(dict) | pd.DataFrame | str - - :raises ValueError: If an unsupported output format is specified. - - :example: - - >>> stackql = StackQL() - >>> query = ''' - ... SELECT SPLIT_PART(machineType, '/', -1) as machine_type, status, COUNT(*) as num_instances - ... FROM google.compute.instances - ... WHERE project = 'stackql-demo' AND zone = 'australia-southeast1-a' - ... GROUP BY machine_type, status HAVING COUNT(*) > 2 - ... ''' - >>> result = stackql.execute(query) - """ - if self.server_mode: - result = self._run_server_query(query) - if self.output == 'pandas': - json_str = json.dumps(result) - return pd.read_json(StringIO(json_str)) - elif self.output == 'csv': - raise ValueError("CSV output is not supported in server_mode.") - else: # Assume 'dict' output - return result - else: - - # returns either... - # [{'error': }] if something went wrong; or - # [{},...] if the statement was executed successfully, messages to stderr - - if self.http_debug: - suppress_errors = False - - output = self._run_query(query, custom_auth=custom_auth, env_vars=env_vars) - - if "exception" in output: - exception_msg = output["exception"] - if self.output == 'pandas': - return pd.DataFrame({'error': [exception_msg]}) if exception_msg else pd.DataFrame({'error': []}) - elif self.output == 'csv': - return exception_msg - else: - return [{"error": exception_msg}] - - if "data" in output: - data = output["data"] - # theres data, return it - if self.output == 'csv': - return data - elif self.output == 'pandas': - try: - return pd.read_json(StringIO(data)) - except ValueError: - return pd.DataFrame([{"error": "Invalid JSON output"}]) - else: # Assume 'dict' output - try: - retval = json.loads(data) - return retval if retval else [] - except ValueError: - return [{"error": f"Invalid JSON output : {data}"}] - - if "error" in output and not suppress_errors: - # theres no data but there is stderr from the request, could be an expected error like a 404 - err_msg = output["error"] - if self.output == 'csv': - return err_msg - elif self.output == 'pandas': - return pd.DataFrame([{"error": err_msg}]) - else: - return [{"error": err_msg}] - - return [] - - # asnyc query support - # - - def _run_server_query_with_new_connection(self, query): - """Run a query against a StackQL postgres wire protocol server with a new connection. - """ - try: - # Establish a new connection using credentials and configurations - with psycopg.connect( - dbname='stackql', - user='stackql', - host=self.server_address, - port=self.server_port, - row_factory=dict_row - ) as conn: - # Execute the query with a new cursor - with conn.cursor() as cur: - cur.execute(query) - try: - rows = cur.fetchall() - except psycopg.ProgrammingError as e: - if str(e) == "no results to fetch": - rows = [] - else: - raise - return rows - except psycopg.OperationalError as oe: - print(f"OperationalError while connecting to the server: {oe}") - except Exception as e: - print(f"Unexpected error while connecting to the server: {e}") - - def _sync_query(self, query, new_connection=False): - """Synchronous function to perform the query. - """ - if self.server_mode and new_connection: - # Directly get the list of dicts; no JSON string conversion needed. - result = self._run_server_query_with_new_connection(query) - elif self.server_mode: - # Also directly get the list of dicts here. - result = self._run_server_query(query) # Assuming this is a method that exists - else: - # Convert the JSON string to a Python object (list of dicts). - query_results = self._run_query(query) - if "exception" in query_results: - result = [{"error": query_results["exception"]}] - if "error" in query_results: - result = [{"error": query_results["error"]}] - if "data" in query_results: - result = json.loads(query_results["data"]) - # Convert the result to a DataFrame if necessary. - if self.output == 'pandas': - return pd.DataFrame(result) - else: - return result - - async def executeQueriesAsync(self, queries): - """Executes multiple StackQL queries asynchronously using the current StackQL instance. - - This method utilizes an asyncio event loop to concurrently run a list of provided - StackQL queries. Each query is executed independently, and the combined results of - all the queries are returned as a list of JSON objects if 'dict' output mode is selected, - or as a concatenated DataFrame if 'pandas' output mode is selected. - - The order of the results in the returned list or DataFrame may not necessarily - correspond to the order of the queries in the input list due to the asynchronous nature - of execution. - - :param queries: A list of StackQL query strings to be executed concurrently. - :type queries: list[str], required - :return: A list of results corresponding to each query. Each result is a JSON object or a DataFrame. - :rtype: list[dict] or pd.DataFrame - :raises ValueError: If method is used in `server_mode` on an unsupported OS (anything other than Linux). - :raises ValueError: If an unsupported output mode is selected (anything other than 'dict' or 'pandas'). - - Example: - >>> from pystackql import StackQL - >>> stackql = StackQL() - >>> queries = [ - >>> \"\"\"SELECT '%s' as region, instanceType, COUNT(*) as num_instances - ... FROM aws.ec2.instances - ... WHERE region = '%s' - ... GROUP BY instanceType\"\"\" % (region, region) - >>> for region in regions ] - >>> result = stackql.executeQueriesAsync(queries) - - Note: - - When operating in `server_mode`, this method is not supported. - """ - if self.server_mode: - raise ValueError("executeQueriesAsync are not supported in sever_mode.") - if self.output not in ['dict', 'pandas']: - raise ValueError("executeQueriesAsync supports only 'dict' or 'pandas' output modes.") - async def main(): - with ThreadPoolExecutor() as executor: - # New connection is created for each query in server_mode, reused otherwise. - new_connection = self.server_mode - # Gather results from all the async calls. - loop = asyncio.get_event_loop() - futures = [loop.run_in_executor(executor, self._sync_query, query, new_connection) for query in queries] - results = await asyncio.gather(*futures) - # Concatenate DataFrames if output mode is 'pandas'. - if self.output == 'pandas': - return pd.concat(results, ignore_index=True) - else: - return [item for sublist in results for item in sublist] - # Running the async function - return await main() diff --git a/pystackql/utils/__init__.py b/pystackql/utils/__init__.py new file mode 100644 index 0000000..27f6e7c --- /dev/null +++ b/pystackql/utils/__init__.py @@ -0,0 +1,49 @@ +# pystackql/utils/__init__.py + +""" +Utility functions for PyStackQL package. +""" +from .package import get_package_version + +from .platform import ( + get_platform, + is_binary_local +) + +from .binary import ( + get_binary_name, + get_binary_version, + setup_binary +) + +from .download import ( + get_download_dir, + get_download_url, + download_file +) + +from .auth import format_auth +from .params import setup_local_mode + +__all__ = [ + # Platform utilities + 'get_platform', + 'get_package_version', + 'is_binary_local', + + # Binary utilities + 'get_binary_name', + 'get_binary_version', + 'setup_binary', + + # Download utilities + 'get_download_dir', + 'get_download_url', + 'download_file', + + # Auth utilities + 'format_auth', + + # Parameter utilities + 'setup_local_mode' +] \ No newline at end of file diff --git a/pystackql/utils/auth.py b/pystackql/utils/auth.py new file mode 100644 index 0000000..809557e --- /dev/null +++ b/pystackql/utils/auth.py @@ -0,0 +1,39 @@ +# pystackql/utils/auth.py + +""" +Authentication utility functions for PyStackQL. + +This module contains functions for handling authentication. +""" + +import json + +def format_auth(auth): + """Formats an authentication object for use with stackql. + + Args: + auth: The authentication object, can be a string or a dict + + Returns: + tuple: (auth_obj, auth_str) + - auth_obj: The authentication object as a dict + - auth_str: The authentication object as a JSON string + + Raises: + Exception: If the authentication object is invalid + """ + try: + if auth is not None: + if isinstance(auth, str): + authobj = json.loads(auth) + authstr = auth + elif isinstance(auth, dict): + authobj = auth + authstr = json.dumps(auth) + return authobj, authstr + else: + raise Exception("ERROR: [format_auth] auth key supplied with no value") + except Exception as e: + error_message = e.args[0] + print(f"ERROR: [format_auth] {error_message}") + exit(1) \ No newline at end of file diff --git a/pystackql/utils/binary.py b/pystackql/utils/binary.py new file mode 100644 index 0000000..2411da1 --- /dev/null +++ b/pystackql/utils/binary.py @@ -0,0 +1,125 @@ +# pystackql/utils/binary.py + +""" +Binary management utility functions for PyStackQL. + +This module contains functions for managing the StackQL binary. +""" + +import os +import subprocess +from .download import get_download_url, download_file, get_download_dir +from .platform import get_platform + + +def get_binary_name(system_platform): + """Gets the binary name based on the platform. + + Args: + system_platform (str): The operating system platform + + Returns: + str: The name of the binary + """ + if system_platform.startswith('Windows'): + return r'stackql.exe' + elif system_platform.startswith('Darwin'): + return r'stackql/Payload/stackql' + else: + return r'stackql' + + +def get_binary_version(bin_path): + """Gets the version of the stackql binary. + + Args: + bin_path (str): The path to the binary + + Returns: + tuple: (version, sha) + - version: The version number + - sha: The git commit sha + + Raises: + FileNotFoundError: If the binary is not found + Exception: If the version cannot be determined + """ + try: + iqlPopen = subprocess.Popen([bin_path] + ["--version"], + stdout=subprocess.PIPE, stderr=subprocess.STDOUT) + # Use communicate to fetch the outputs and wait for the process to finish + output, _ = iqlPopen.communicate() + # Decode the output + decoded_output = output.decode('utf-8') + # Split to get the version tokens + version_tokens = decoded_output.split('\n')[0].split(' ') + version = version_tokens[1] + sha = version_tokens[3].replace('(', '').replace(')', '') + return version, sha + except FileNotFoundError: + print(f"ERROR: [get_binary_version] {bin_path} not found") + exit(1) + except Exception as e: + error_message = e.args[0] + print(f"ERROR: [get_binary_version] {error_message}") + exit(1) + finally: + # Ensure the subprocess is terminated and streams are closed + iqlPopen.terminate() + if hasattr(iqlPopen, 'stdout') and iqlPopen.stdout: + iqlPopen.stdout.close() + + +def setup_binary(download_dir, system_platform, showprogress=False): + """Sets up the stackql binary by downloading and extracting it. + + Args: + download_dir (str): The directory to download to + system_platform (str): The operating system platform + showprogress (bool, optional): Whether to show download progress. Defaults to False. + + Raises: + Exception: If the setup fails + """ + try: + print('installing stackql...') + binary_name = get_binary_name(system_platform) + url = get_download_url() + 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) + + # Handle extraction + if system_platform.startswith('Darwin'): + unpacked_file_name = os.path.join(download_dir, 'stackql') + command = f'pkgutil --expand-full {archive_file_name} {unpacked_file_name}' + if os.path.exists(unpacked_file_name): + os.system(f'rm -rf {unpacked_file_name}') + os.system(command) + + else: # Handle Windows and Linux + import zipfile + with zipfile.ZipFile(archive_file_name, 'r') as zip_ref: + zip_ref.extractall(download_dir) + + # Specific check for Windows to ensure `stackql.exe` is extracted + if system_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(f"ERROR: [setup_binary] {str(e)}") + exit(1) \ No newline at end of file diff --git a/pystackql/utils/download.py b/pystackql/utils/download.py new file mode 100644 index 0000000..5fd71ea --- /dev/null +++ b/pystackql/utils/download.py @@ -0,0 +1,79 @@ +# pystackql/utils/download.py + +""" +Download-related utility functions for PyStackQL. + +This module contains functions for downloading and managing the StackQL binary. +""" + +import os +import site +import platform +import requests + + +def get_download_dir(): + """Gets the directory to download the stackql binary. + + Returns: + str: The directory path + """ + # check if site.getuserbase() dir exists + if not os.path.exists(site.getuserbase()): + # if not, create it + os.makedirs(site.getuserbase()) + return site.getuserbase() + + +def get_download_url(): + """Gets the download URL for the stackql binary based on the platform. + + Returns: + str: The download URL + + Raises: + Exception: If the platform is not supported + """ + system_val = platform.system() + machine_val = platform.machine() + + if system_val == 'Linux' and machine_val == 'x86_64': + return 'https://releases.stackql.io/stackql/latest/stackql_linux_amd64.zip' + elif system_val == 'Windows': + return 'https://releases.stackql.io/stackql/latest/stackql_windows_amd64.zip' + elif system_val == 'Darwin': + return 'https://storage.googleapis.com/stackql-public-releases/latest/stackql_darwin_multiarch.pkg' + else: + raise Exception(f"ERROR: [get_download_url] unsupported OS type: {system_val} {machine_val}") + + +def download_file(url, path, showprogress=True): + """Downloads a file from a URL to a local path. + + Args: + url (str): The URL to download from + path (str): The local path to save the file to + showprogress (bool, optional): Whether to show a progress bar. Defaults to True. + + Raises: + Exception: If the download fails + """ + try: + r = requests.get(url, stream=True) + r.raise_for_status() + total_size_in_bytes = int(r.headers.get('content-length', 0)) + block_size = 1024 + with open(path, 'wb') as f: + chunks = 0 + for data in r.iter_content(block_size): + chunks += 1 + f.write(data) + downloaded_size = chunks * block_size + progress_bar = '#' * int(downloaded_size / total_size_in_bytes * 20) + if showprogress: + print(f'\r[{progress_bar.ljust(20)}] {int(downloaded_size / total_size_in_bytes * 100)}%', end='') + + print("\nDownload complete.") + except Exception as e: + print(f"ERROR: [download_file] {str(e)}") + exit(1) \ No newline at end of file diff --git a/pystackql/utils/helpers.py b/pystackql/utils/helpers.py new file mode 100644 index 0000000..de0c329 --- /dev/null +++ b/pystackql/utils/helpers.py @@ -0,0 +1,284 @@ +# pystackql/utils/helpers.py + +""" +Utility functions for PyStackQL package. + +This module contains helper functions for binary management, platform detection, +and other utilities needed by the PyStackQL package. +""" + +import subprocess +import platform +import json +import site +import os +import requests +import zipfile + +# Conditional import for package metadata retrieval +try: + from importlib.metadata import version, PackageNotFoundError +except ImportError: + # This is for Python versions earlier than 3.8 + from importlib_metadata import version, PackageNotFoundError + + +def is_binary_local(system_platform): + """Checks if the binary exists at the specified local path. + + Args: + system_platform (str): The operating system platform + + Returns: + bool: True if the binary exists at the expected local path + """ + if system_platform == 'Linux' and os.path.exists('/usr/local/bin/stackql'): + return True + return False + + +def get_package_version(package_name): + """Gets the version of the specified package. + + Args: + package_name (str): The name of the package + + Returns: + str: The version of the package or None if not found + """ + try: + pkg_version = version(package_name) + if pkg_version is None: + print(f"Warning: Retrieved version for '{package_name}' is None!") + return pkg_version + except PackageNotFoundError: + print(f"Warning: Package '{package_name}' not found!") + return None + + +def get_platform(): + """Gets the current platform information. + + Returns: + tuple: (platform_string, system_value) + - platform_string: A string with platform details + - system_value: The operating system name + """ + system_val = platform.system() + machine_val = platform.machine() + platform_val = platform.platform() + python_version_val = platform.python_version() + return ( + f"{system_val} {machine_val} ({platform_val}), Python {python_version_val}", + system_val + ) + + +def get_download_dir(): + """Gets the directory to download the stackql binary. + + Returns: + str: The directory path + """ + # check if site.getuserbase() dir exists + if not os.path.exists(site.getuserbase()): + # if not, create it + os.makedirs(site.getuserbase()) + return site.getuserbase() + + +def get_binary_name(system_platform): + """Gets the binary name based on the platform. + + Args: + system_platform (str): The operating system platform + + Returns: + str: The name of the binary + """ + if system_platform.startswith('Windows'): + return r'stackql.exe' + elif system_platform.startswith('Darwin'): + return r'stackql/Payload/stackql' + else: + return r'stackql' + + +def get_download_url(): + """Gets the download URL for the stackql binary based on the platform. + + Returns: + str: The download URL + + Raises: + Exception: If the platform is not supported + """ + system_val = platform.system() + machine_val = platform.machine() + + if system_val == 'Linux' and machine_val == 'x86_64': + return 'https://releases.stackql.io/stackql/latest/stackql_linux_amd64.zip' + elif system_val == 'Windows': + return 'https://releases.stackql.io/stackql/latest/stackql_windows_amd64.zip' + elif system_val == 'Darwin': + return 'https://storage.googleapis.com/stackql-public-releases/latest/stackql_darwin_multiarch.pkg' + else: + raise Exception(f"ERROR: [get_download_url] unsupported OS type: {system_val} {machine_val}") + + +def download_file(url, path, showprogress=True): + """Downloads a file from a URL to a local path. + + Args: + url (str): The URL to download from + path (str): The local path to save the file to + showprogress (bool, optional): Whether to show a progress bar. Defaults to True. + + Raises: + Exception: If the download fails + """ + try: + r = requests.get(url, stream=True) + r.raise_for_status() + total_size_in_bytes = int(r.headers.get('content-length', 0)) + block_size = 1024 + with open(path, 'wb') as f: + chunks = 0 + for data in r.iter_content(block_size): + chunks += 1 + f.write(data) + downloaded_size = chunks * block_size + progress_bar = '#' * int(downloaded_size / total_size_in_bytes * 20) + if showprogress: + print(f'\r[{progress_bar.ljust(20)}] {int(downloaded_size / total_size_in_bytes * 100)}%', end='') + + print("\nDownload complete.") + except Exception as e: + print(f"ERROR: [download_file] {str(e)}") + exit(1) + + +def setup_binary(download_dir, system_platform, showprogress=False): + """Sets up the stackql binary by downloading and extracting it. + + Args: + download_dir (str): The directory to download to + system_platform (str): The operating system platform + showprogress (bool, optional): Whether to show download progress. Defaults to False. + + Raises: + Exception: If the setup fails + """ + try: + print('installing stackql...') + binary_name = get_binary_name(system_platform) + url = get_download_url() + 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) + + # Handle extraction + if system_platform.startswith('Darwin'): + unpacked_file_name = os.path.join(download_dir, 'stackql') + command = f'pkgutil --expand-full {archive_file_name} {unpacked_file_name}' + if os.path.exists(unpacked_file_name): + os.system(f'rm -rf {unpacked_file_name}') + os.system(command) + + else: # Handle Windows and Linux + with zipfile.ZipFile(archive_file_name, 'r') as zip_ref: + zip_ref.extractall(download_dir) + + # Specific check for Windows to ensure `stackql.exe` is extracted + if system_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(f"ERROR: [setup_binary] {str(e)}") + exit(1) + + +def get_binary_version(bin_path): + """Gets the version of the stackql binary. + + Args: + bin_path (str): The path to the binary + + Returns: + tuple: (version, sha) + - version: The version number + - sha: The git commit sha + + Raises: + FileNotFoundError: If the binary is not found + Exception: If the version cannot be determined + """ + try: + iqlPopen = subprocess.Popen([bin_path] + ["--version"], + stdout=subprocess.PIPE, stderr=subprocess.STDOUT) + # Use communicate to fetch the outputs and wait for the process to finish + output, _ = iqlPopen.communicate() + # Decode the output + decoded_output = output.decode('utf-8') + # Split to get the version tokens + version_tokens = decoded_output.split('\n')[0].split(' ') + version = version_tokens[1] + sha = version_tokens[3].replace('(', '').replace(')', '') + return version, sha + except FileNotFoundError: + print(f"ERROR: [get_binary_version] {bin_path} not found") + exit(1) + except Exception as e: + error_message = e.args[0] + print(f"ERROR: [get_binary_version] {error_message}") + exit(1) + finally: + # Ensure the subprocess is terminated and streams are closed + iqlPopen.terminate() + if hasattr(iqlPopen, 'stdout') and iqlPopen.stdout: + iqlPopen.stdout.close() + + +def format_auth(auth): + """Formats an authentication object for use with stackql. + + Args: + auth: The authentication object, can be a string or a dict + + Returns: + tuple: (auth_obj, auth_str) + - auth_obj: The authentication object as a dict + - auth_str: The authentication object as a JSON string + + Raises: + Exception: If the authentication object is invalid + """ + try: + if auth is not None: + if isinstance(auth, str): + authobj = json.loads(auth) + authstr = auth + elif isinstance(auth, dict): + authobj = auth + authstr = json.dumps(auth) + return authobj, authstr + else: + raise Exception("ERROR: [format_auth] auth key supplied with no value") + except Exception as e: + error_message = e.args[0] + print(f"ERROR: [format_auth] {error_message}") + exit(1) \ No newline at end of file diff --git a/pystackql/utils/package.py b/pystackql/utils/package.py new file mode 100644 index 0000000..3f1baf1 --- /dev/null +++ b/pystackql/utils/package.py @@ -0,0 +1,31 @@ +# pystackql/utils/package.py + +""" +Package related utility functions for PyStackQL. + +""" + +# Conditional import for package metadata retrieval +try: + from importlib.metadata import version, PackageNotFoundError +except ImportError: + # This is for Python versions earlier than 3.8 + from importlib_metadata import version, PackageNotFoundError + +def get_package_version(package_name): + """Gets the version of the specified package. + + Args: + package_name (str): The name of the package + + Returns: + str: The version of the package or None if not found + """ + try: + pkg_version = version(package_name) + if pkg_version is None: + print(f"Warning: Retrieved version for '{package_name}' is None!") + return pkg_version + except PackageNotFoundError: + print(f"Warning: Package '{package_name}' not found!") + return None diff --git a/pystackql/utils/params.py b/pystackql/utils/params.py new file mode 100644 index 0000000..1b039b5 --- /dev/null +++ b/pystackql/utils/params.py @@ -0,0 +1,160 @@ +# pystackql/utils/params.py + +""" +Parameter generation utility for StackQL local mode. + +This module provides functions to generate command-line parameters for the StackQL binary +and helps set instance attributes. +""" + +import json +from .auth import format_auth + +def _set_param(params, param_name, value): + """Add a parameter and its value to the params list. + + :param params: List of parameters to append to + :param param_name: Parameter name to add + :param value: Value to add + :return: Updated params list + """ + params.append(f"--{param_name}") + params.append(str(value)) + return params + +def setup_local_mode(instance, **kwargs): + """Set up local mode for a StackQL instance. + + This function generates parameters and sets instance attributes + for local mode operation. + + :param instance: The StackQL instance + :param kwargs: Keyword arguments from the constructor + :return: List of parameters for StackQL binary + """ + # Initialize parameter list + params = ["exec"] + + # Extract parameters from kwargs with defaults matching the StackQL.__init__ defaults + output = kwargs.get('output', 'dict') + backend_storage_mode = kwargs.get('backend_storage_mode', 'memory') + backend_file_storage_location = kwargs.get('backend_file_storage_location', 'stackql.db') + app_root = kwargs.get('app_root', None) + execution_concurrency_limit = kwargs.get('execution_concurrency_limit', -1) + dataflow_dependency_max = kwargs.get('dataflow_dependency_max', 50) + dataflow_components_max = kwargs.get('dataflow_components_max', 50) + custom_registry = kwargs.get('custom_registry', None) + custom_auth = kwargs.get('custom_auth', None) + sep = kwargs.get('sep', ',') + header = kwargs.get('header', False) + max_results = kwargs.get('max_results', -1) + page_limit = kwargs.get('page_limit', 20) + max_depth = kwargs.get('max_depth', 5) + api_timeout = kwargs.get('api_timeout', 45) + http_debug = kwargs.get('http_debug', False) + proxy_host = kwargs.get('proxy_host', None) + proxy_port = kwargs.get('proxy_port', -1) + proxy_user = kwargs.get('proxy_user', None) + proxy_password = kwargs.get('proxy_password', None) + proxy_scheme = kwargs.get('proxy_scheme', 'http') + download_dir = kwargs.get('download_dir', None) + debug = kwargs.get('debug', False) + debug_log_file = kwargs.get('debug_log_file', None) + + # Set output format + params.append("--output") + if output.lower() == "csv": + params.append("csv") + else: + params.append("json") + + # Backend storage settings + if backend_storage_mode == 'file': + params.append("--sqlBackend") + params.append(json.dumps({ "dsn": f"file:{backend_file_storage_location}" })) + + # If app_root is set, use it + if app_root is not None: + instance.app_root = app_root + _set_param(params, 'approot', app_root) + + # Set execution parameters + instance.execution_concurrency_limit = execution_concurrency_limit + _set_param(params, 'execution.concurrency.limit', execution_concurrency_limit) + + instance.dataflow_dependency_max = dataflow_dependency_max + _set_param(params, 'dataflow.dependency.max', dataflow_dependency_max) + + instance.dataflow_components_max = dataflow_components_max + _set_param(params, 'dataflow.components.max', dataflow_components_max) + + # If custom_auth is set, use it + if custom_auth is not None: + authobj, authstr = format_auth(custom_auth) + instance.auth = authobj + params.append("--auth") + params.append(authstr) + + # If custom_registry is set, use it + if custom_registry is not None: + instance.custom_registry = custom_registry + params.append("--registry") + params.append(json.dumps({ "url": custom_registry })) + + # CSV output settings + if output.lower() == "csv": + instance.sep = sep + _set_param(params, 'delimiter', sep) + + instance.header = header + if not header: + params.append("--hideheaders") + + # App behavioral properties + instance.max_results = max_results + _set_param(params, 'http.response.maxResults', max_results) + + instance.page_limit = page_limit + _set_param(params, 'http.response.pageLimit', page_limit) + + instance.max_depth = max_depth + _set_param(params, 'indirect.depth.max', max_depth) + + instance.api_timeout = api_timeout + _set_param(params, 'apirequesttimeout', api_timeout) + + instance.http_debug = bool(http_debug) + if http_debug: + params.append("--http.log.enabled") + + # Proxy settings + if proxy_host is not None: + # Set attributes + instance.proxy_host = proxy_host + instance.proxy_port = proxy_port + instance.proxy_user = proxy_user + instance.proxy_password = proxy_password + + # Set basic proxy parameters + _set_param(params, 'http.proxy.host', proxy_host) + _set_param(params, 'http.proxy.port', proxy_port) + _set_param(params, 'http.proxy.user', proxy_user) + _set_param(params, 'http.proxy.password', proxy_password) + + # Validate and set proxy scheme + ALLOWED_PROXY_SCHEMES = {'http', 'https'} + if proxy_scheme.lower() not in ALLOWED_PROXY_SCHEMES: + raise ValueError(f"Invalid proxy_scheme. Expected one of {ALLOWED_PROXY_SCHEMES}, got {proxy_scheme}.") + + instance.proxy_scheme = proxy_scheme.lower() + _set_param(params, 'http.proxy.scheme', proxy_scheme.lower()) + + # Initialize binary manager + from ..core.binary import BinaryManager # Import here to avoid circular imports + instance.binary_manager = BinaryManager(download_dir) + instance.bin_path = instance.binary_manager.bin_path + instance.version = instance.binary_manager.version + instance.sha = instance.binary_manager.sha + + # Return the params list + return params \ No newline at end of file diff --git a/pystackql/utils/platform.py b/pystackql/utils/platform.py new file mode 100644 index 0000000..91889d7 --- /dev/null +++ b/pystackql/utils/platform.py @@ -0,0 +1,40 @@ +# pystackql/utils/platform.py + +""" +Platform-related utility functions for PyStackQL. + +This module contains functions for platform detection and package information. +""" + +import os +import platform + +def is_binary_local(system_platform): + """Checks if the binary exists at the specified local path. + + Args: + system_platform (str): The operating system platform + + Returns: + bool: True if the binary exists at the expected local path + """ + if system_platform == 'Linux' and os.path.exists('/usr/local/bin/stackql'): + return True + return False + +def get_platform(): + """Gets the current platform information. + + Returns: + tuple: (platform_string, system_value) + - platform_string: A string with platform details + - system_value: The operating system name + """ + system_val = platform.system() + machine_val = platform.machine() + platform_val = platform.platform() + python_version_val = platform.python_version() + return ( + f"{system_val} {machine_val} ({platform_val}), Python {python_version_val}", + system_val + ) \ No newline at end of file diff --git a/requirements.txt b/requirements.txt index df3235c..8895f3d 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,6 +1,23 @@ -nose -sphinx -pandas -termcolor -requests -IPython \ No newline at end of file +# Core dependencies +pandas>=1.3.0 +requests>=2.25.0 +IPython>=7.0.0 +termcolor>=1.1.0 + +# Documentation +sphinx>=4.0.0 +sphinx-rtd-theme>=1.0.0 + +# Testing +pytest>=6.2.5 +pytest-cov>=2.12.0 +nose>=1.3.7 # For backward compatibility + +# Platform-independent psycopg installation +psycopg[binary]>=3.1.0 # Uses binary wheels where available + +# Async support +nest-asyncio>=1.5.5 # For running async code in Jupyter + +# Optional utilities +tqdm>=4.61.0 # For progress bars \ No newline at end of file diff --git a/run_async_server_tests b/run_async_server_tests deleted file mode 100644 index 5c66a63..0000000 --- a/run_async_server_tests +++ /dev/null @@ -1,3 +0,0 @@ -#!/bin/bash -. tests/creds/env_vars/test.env -python3 -m tests.pystackql_async_server_tests \ No newline at end of file diff --git a/run_server_tests.py b/run_server_tests.py new file mode 100644 index 0000000..ce9e554 --- /dev/null +++ b/run_server_tests.py @@ -0,0 +1,53 @@ +#!/usr/bin/env python3 +""" +Test runner script for PyStackQL. + +This script runs all the PyStackQL tests in server mode. It can be used to run +individual test files or all tests. + +A running instance of the stackql server is required to run the server tests. + +Examples: + # Run all tests + python run_server_tests.py + + # Run specific test files + python run_server_tests.py tests/test_server.py + + # Run with verbose output + python run_server_tests.py -v +""" + +import sys +import os +import pytest +from termcolor import colored + +# Add the current directory to the Python path +sys.path.insert(0, os.path.abspath(os.path.dirname(__file__))) + +# Add the tests directory to the Python path +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), 'tests'))) + +def main(): + """Run the tests.""" + print(colored("\n===== PyStackQL Server Test Runner =====\n", "cyan")) + + # Default pytest arguments + args = ["-v"] + + # Add any specific test files passed as arguments + if len(sys.argv) > 1: + args.extend(sys.argv[1:]) + else: + # If no specific tests were requested, run all non-server test files + args.extend([ + "tests/test_server.py", + "tests/test_server_magic.py" + ]) + + # Run pytest with the arguments + return pytest.main(args) + +if __name__ == "__main__": + sys.exit(main()) \ No newline at end of file diff --git a/run_tests b/run_tests deleted file mode 100644 index a0cea99..0000000 --- a/run_tests +++ /dev/null @@ -1,11 +0,0 @@ -#!/bin/bash - -# Install packages if they aren't already installed -# pip3 install -r requirements.txt --user -# pip3 install psycopg --user - -# Load environment variables -source ./tests/creds/env_vars/test.env.sh - -# Run tests -python3 -m tests.pystackql_tests diff --git a/run_tests.ps1 b/run_tests.ps1 deleted file mode 100644 index 2b255bd..0000000 --- a/run_tests.ps1 +++ /dev/null @@ -1,11 +0,0 @@ -# install packages -# pip.exe install -r requirements.txt --user -# pip install psycopg[binary] - -# Load environment variables -. .\tests\creds\env_vars\test.env.ps1 - -$env:PYTHONPATH = "C:\LocalGitRepos\stackql\python-packages\pystackql" - -# Run tests -python.exe -m tests.pystackql_tests diff --git a/run_tests.py b/run_tests.py new file mode 100644 index 0000000..d268ea2 --- /dev/null +++ b/run_tests.py @@ -0,0 +1,60 @@ +#!/usr/bin/env python3 +""" +Test runner script for PyStackQL. + +This script runs all the PyStackQL tests. It can be used to run +individual test files or all tests. + +Examples: + # Run all tests + python run_tests.py + + # Run specific test files + python run_tests.py tests/test_core.py tests/test_query_execution.py + + # Run with verbose output + python run_tests.py -v +""" + +import sys +import os +import pytest +from termcolor import colored + +# Add the current directory to the Python path +sys.path.insert(0, os.path.abspath(os.path.dirname(__file__))) + +# Add the tests directory to the Python path +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), 'tests'))) + +def main(): + """Run the tests.""" + print(colored("\n===== PyStackQL Test Runner =====\n", "cyan")) + + # Default pytest arguments + args = ["-v"] + + # Add any specific test files passed as arguments + if len(sys.argv) > 1: + args.extend(sys.argv[1:]) + else: + # If no specific tests were requested, run all non-server test files + args.extend([ + "tests/test_core.py", + "tests/test_query_execution.py", + "tests/test_output_formats.py", + "tests/test_magic.py", + "tests/test_async.py" + ]) + + # Skip server tests by default as they require a running server + # Uncomment to run server tests + # args.append("tests/test_server.py") + # args.append("tests/test_server_magic.py") + + # Run pytest with the arguments + return pytest.main(args) + + +if __name__ == "__main__": + sys.exit(main()) \ No newline at end of file diff --git a/server-status.sh b/server-status.sh new file mode 100644 index 0000000..bdedf52 --- /dev/null +++ b/server-status.sh @@ -0,0 +1,20 @@ +#!/bin/bash + +# Search for the stackql process +stackql_process=$(ps -ef | grep '[s]tackql') + +# Check if the process is running +if [ -z "$stackql_process" ]; then + echo "Server is not running." +else + # Extract the port and PID using awk/sed + port=$(echo "$stackql_process" | sed -n 's/.*--pgsrv.port=\([0-9]*\).*/\1/p') + pid=$(echo "$stackql_process" | awk '{print $2}') + + # Check if port extraction was successful + if [ -z "$port" ]; then + echo "Server is running but could not detect the port (PID $pid)" + else + echo "Server is running on port $port (PID $pid)" + fi +fi diff --git a/setup.py b/setup.py index 003cab8..cadda4b 100644 --- a/setup.py +++ b/setup.py @@ -10,7 +10,7 @@ setup( name='pystackql', - version='v3.7.2', + version='v3.8.0', description='A Python interface for StackQL', long_description=readme, author='Jeffrey Aven', @@ -23,7 +23,11 @@ 'requests', 'pandas', 'IPython', - ], + 'psycopg[binary]>=3.1.0', # Added psycopg with binary wheels for all platforms + 'nest-asyncio>=1.5.5', # For async support in Jupyter + 'termcolor>=1.1.0', # For colored output in test runner + 'tqdm>=4.61.0', # For progress bars in download method + ], # entry_points={ # 'console_scripts': [ # 'stackql = pystackql:setup' diff --git a/start-stackql-server.sh b/start-stackql-server.sh new file mode 100644 index 0000000..bea44a9 --- /dev/null +++ b/start-stackql-server.sh @@ -0,0 +1,9 @@ +# start server if not running +echo "checking if server is running" +if [ -z "$(ps | grep stackql)" ]; then + echo "starting server with registry: $REG" + nohup ./stackql -v --pgsrv.port=5444 srv & + sleep 5 +else + echo "server is already running" +fi \ No newline at end of file diff --git a/stop-stackql-server.sh b/stop-stackql-server.sh new file mode 100644 index 0000000..762f6e8 --- /dev/null +++ b/stop-stackql-server.sh @@ -0,0 +1,12 @@ +#!/bin/bash + +# Find the process ID of the StackQL server +PID=$(pgrep -f "stackql") + +if [ -z "$PID" ]; then + echo "stackql server is not running." +else + echo "stopping stackql server (PID: $PID)..." + kill $PID + echo "stackql server stopped." +fi \ No newline at end of file diff --git a/tests/README.md b/tests/README.md new file mode 100644 index 0000000..6491481 --- /dev/null +++ b/tests/README.md @@ -0,0 +1,158 @@ +# PyStackQL Testing Guide + +## Overview + +This guide explains the PyStackQL testing framework. The tests have been designed to: + +1. Focus on provider-agnostic queries where possible +2. Use the Homebrew provider for provider-specific tests (no authentication required) +3. Be organized into logical modules based on functionality +4. Support both local execution and GitHub Codespaces + +## Test Structure + +The tests are organized into these main files: + +- `test_constants.py`: Common constants and helper functions +- `conftest.py`: Test fixtures and setup +- `test_core.py`: Core functionality tests +- `test_query_execution.py`: Query execution tests +- `test_output_formats.py`: Output format tests +- `test_magic.py`: Magic extension tests +- `test_async.py`: Async functionality tests +- `test_server.py`: Server mode tests + +## Running Tests + +### Running All Tests + +To run all tests: + +```bash +python run_tests.py +``` + +### Running Specific Tests + +To run specific test files: + +```bash +python run_tests.py tests/test_core.py tests/test_query_execution.py +``` + +### Running with Extra Verbosity + +```bash +python run_tests.py -v +``` + +### Running Server Tests + +Server tests are skipped by default because they require a running StackQL server. To run these tests: + +1. Start a StackQL server: + ```bash + stackql srv --pgsrv.address 127.0.0.1 --pgsrv.port 5466 + ``` + +2. Run the server tests: + ```bash + python run_tests.py tests/test_server.py -v + ``` + +## Test Categories + +### Core Tests + +Tests the basic properties and attributes of the `StackQL` class: + +- `properties()` method +- `version`, `package_version`, `platform` attributes +- Binary path and download directory +- Upgrade functionality + +### Query Execution Tests + +Tests the query execution functionality with provider-agnostic queries: + +- Literal values (integers, strings, floats) +- Expressions +- JSON extraction +- Homebrew provider queries +- Registry pull operations + +### Output Format Tests + +Tests the different output formats: + +- Dict output +- Pandas output with type checking +- CSV output with different separators and headers +- Error handling for invalid configurations + +### Magic Tests + +Tests the Jupyter magic extensions: + +- Line and cell magic in non-server mode +- Line and cell magic in server mode +- Result storage in user namespace +- Display options + +### Async Tests + +Tests the async query execution functionality: + +- `executeQueriesAsync` with different output formats +- Concurrent queries with the Homebrew provider +- Error handling + +### Server Tests + +Tests the server mode functionality (requires a running server): + +- Server connectivity +- Query execution in server mode +- Statement execution in server mode +- Different output formats in server mode + +## Test Data + +The tests use: + +1. **Simple literals and expressions**: + ```sql + SELECT 1 as literal_int_value + SELECT 1.001 as literal_float_value + SELECT 'test' as literal_string_value + SELECT 1=1 as expression + ``` + +2. **Homebrew provider queries**: + ```sql + SELECT name, full_name, tap FROM homebrew.formula.formula WHERE formula_name = 'stackql' + SELECT * FROM homebrew.formula.vw_usage_metrics WHERE formula_name = 'stackql' + ``` + +3. **Registry operations**: + ```sql + REGISTRY PULL homebrew + ``` + +## Testing in GitHub Codespaces + +When running in GitHub Codespaces: + +1. The tests automatically detect if they're running in GitHub Actions and skip the binary upgrade +2. The server tests are skipped by default (can be enabled if needed) +3. Async tests might be skipped on Windows due to asyncio issues + +## Adding New Tests + +When adding new tests: + +1. Use provider-agnostic queries where possible +2. For provider-specific tests, prefer the Homebrew provider +3. Add new tests to the appropriate test file based on functionality +4. Update `run_tests.py` if adding a new test file +5. Follow the existing patterns for consistency \ No newline at end of file diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000..b7d285e --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,132 @@ +# tests/conftest.py + +""" +Common test setup and fixtures for PyStackQL tests. +""" + +import os +import sys +import time +import pytest +import subprocess +import signal +from unittest.mock import MagicMock + +# Add the parent directory to the path so we can import from pystackql +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))) + +# Add the current directory to the path so we can import test_constants +sys.path.insert(0, os.path.abspath(os.path.dirname(__file__))) + +from pystackql import StackQL +from tests.test_constants import SERVER_ADDRESS, SERVER_PORT, REGISTRY_PULL_HOMEBREW_QUERY + +# Global variables to store server process +server_process = None + +@pytest.fixture(scope="session", autouse=True) +def setup_stackql(): + """ + Session-wide fixture to download stackql binary and setup server. + This runs once before all tests. + """ + print("\nDownloading and setting up stackql binary...") + stackql = StackQL() + + # Check if we're running in GitHub Actions + is_github_actions = os.environ.get('GITHUB_ACTIONS') == 'true' + if not is_github_actions: + print("Running tests outside of GitHub Actions, upgrading stackql binary...") + stackql.upgrade() + + # Pull the homebrew provider for provider-specific tests + print("Pulling homebrew provider for tests...") + result = stackql.executeStmt(REGISTRY_PULL_HOMEBREW_QUERY) + print(result) + + # Return the StackQL instance for use in tests + return stackql + +@pytest.fixture(scope="session") +def stackql_server(setup_stackql): + """ + Session-level fixture to start and stop a StackQL server. + This runs once for all tests that request it. + + This improved version: + 1. Checks if a server is already running before starting one + 2. Uses process groups for better cleanup + 3. Handles errors more gracefully + """ + global server_process + + # Check if server is already running + print("\nChecking if server is running...") + ps_output = subprocess.run( + ["ps", "aux"], + capture_output=True, + text=True + ).stdout + + if "stackql" in ps_output and f"--pgsrv.port={SERVER_PORT}" in ps_output: + print("Server is already running") + # No need to start a server or set server_process + else: + # Start the server + print(f"Starting stackql server on {SERVER_ADDRESS}:{SERVER_PORT}...") + + # Get the registry setting from environment variable if available + registry = os.environ.get('REG', '') + registry_arg = f"--registry {registry}" if registry else "" + + # Build the command + cmd = f"{setup_stackql.bin_path} srv --pgsrv.address {SERVER_ADDRESS} --pgsrv.port {SERVER_PORT} {registry_arg}" + + # Start the server process with process group for better cleanup + server_process = subprocess.Popen( + cmd, + shell=True, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + preexec_fn=os.setsid # Use process group for cleaner termination + ) + + # Wait for server to start + print("Waiting for server to initialize...") + time.sleep(5) + + # Check if server started successfully + if server_process.poll() is not None: + # Process has terminated + stdout, stderr = server_process.communicate() + pytest.fail(f"Server failed to start: {stderr.decode()}") + + # Yield to run tests + yield + + # Clean up server at the end if we started it + if server_process and server_process.poll() is None: + print("\nShutting down stackql server...") + try: + # Kill the entire process group + os.killpg(os.getpgid(server_process.pid), signal.SIGTERM) + server_process.wait(timeout=5) + print("Server shutdown complete") + except subprocess.TimeoutExpired: + print("Server did not terminate gracefully, forcing shutdown...") + os.killpg(os.getpgid(server_process.pid), signal.SIGKILL) + +@pytest.fixture +def mock_interactive_shell(): + """Create a mock IPython shell for testing.""" + class MockInteractiveShell: + def __init__(self): + self.user_ns = {} + self.register_magics_called = False + + def register_magics(self, magic_instance): + """Mock for registering magics.""" + self.magics = magic_instance + self.register_magics_called = True + + return MockInteractiveShell() \ No newline at end of file diff --git a/tests/creds/env_vars/.gitignore b/tests/creds/env_vars/.gitignore deleted file mode 100644 index 86d0cb2..0000000 --- a/tests/creds/env_vars/.gitignore +++ /dev/null @@ -1,4 +0,0 @@ -# Ignore everything in this directory -* -# Except this file -!.gitignore \ No newline at end of file diff --git a/tests/creds/keys/.gitignore b/tests/creds/keys/.gitignore deleted file mode 100644 index 86d0cb2..0000000 --- a/tests/creds/keys/.gitignore +++ /dev/null @@ -1,4 +0,0 @@ -# Ignore everything in this directory -* -# Except this file -!.gitignore \ No newline at end of file diff --git a/tests/pystackql_async_server_tests.py b/tests/pystackql_async_server_tests.py deleted file mode 100644 index e985b35..0000000 --- a/tests/pystackql_async_server_tests.py +++ /dev/null @@ -1,58 +0,0 @@ -import sys, os, unittest, asyncio -sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))) -from pystackql import StackQL -from .test_params import * - -def async_test_decorator(func): - def wrapper(*args, **kwargs): - if asyncio.iscoroutinefunction(func): - return asyncio.run(func(*args, **kwargs)) - else: - return func(*args, **kwargs) - return wrapper - -class PyStackQLTestsBase(unittest.TestCase): - pass - -def setUpModule(): - print("downloading stackql binary...") - PyStackQLTestsBase.stackql = StackQL() - 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(30) - -def tearDownModule(): - print("stopping stackql server...") - if PyStackQLTestsBase.server_process: - PyStackQLTestsBase.server_process.terminate() - PyStackQLTestsBase.server_process.wait() - -class PyStackQLAsyncTests(PyStackQLTestsBase): - - @async_test_decorator - async def test_executeQueriesAsync_server_mode_default_output(self): - stackql = StackQL(server_mode=True) - result = await stackql.executeQueriesAsync(async_queries) - is_valid_result = isinstance(result, list) and all(isinstance(res, dict) for res in result) - self.assertTrue(is_valid_result, f"Result is not a valid list of dicts: {result}") - print_test_result(f"[ASYNC] Test executeQueriesAsync in server_mode with default output\nRESULT_COUNT: {len(result)}", is_valid_result, True) - - @async_test_decorator - async def test_executeQueriesAsync_server_mode_pandas_output(self): - stackql = StackQL(server_mode=True, output='pandas') - result = await stackql.executeQueriesAsync(async_queries) - is_valid_dataframe = isinstance(result, pd.DataFrame) and not result.empty - self.assertTrue(is_valid_dataframe, f"Result is not a valid DataFrame: {result}") - print_test_result(f"[ASYNC] Test executeQueriesAsync in server_mode with pandas output\nRESULT_COUNT: {len(result)}", is_valid_dataframe, True) - -def main(): - unittest.main(verbosity=0) - -if __name__ == '__main__': - main() diff --git a/tests/pystackql_tests.py b/tests/pystackql_tests.py deleted file mode 100644 index aeaff14..0000000 --- a/tests/pystackql_tests.py +++ /dev/null @@ -1,483 +0,0 @@ -import sys, os, unittest, asyncio, re -from unittest.mock import MagicMock, patch -sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))) -from pystackql import StackQL, magic, magics, StackqlMagic, StackqlServerMagic -from .test_params import * - -def pystackql_test_setup(**kwargs): - def decorator(func): - def wrapper(self, *args): - try: - del self.stackql - except AttributeError: - pass - self.stackql = StackQL(**kwargs) - func(self, *args) - return wrapper - return decorator - -def async_test_decorator(func): - def wrapper(*args, **kwargs): - if asyncio.iscoroutinefunction(func): - return asyncio.run(func(*args, **kwargs)) - else: - return func(*args, **kwargs) - return wrapper - -class PyStackQLTestsBase(unittest.TestCase): - pass - -def setUpModule(): - print("downloading stackql binary...") - PyStackQLTestsBase.stackql = StackQL() - # Check whether code is running in GitHub Actions - is_github_actions = os.environ.get('GITHUB_ACTIONS') == 'true' - - if not is_github_actions: - # Ensure you have the latest version of stackql, only when running locally - print("Running tests outside of GitHub Actions, upgrading stackql binary...") - PyStackQLTestsBase.stackql.upgrade() - - print("downloading aws provider for tests...") - res = PyStackQLTestsBase.stackql.executeStmt(registry_pull_aws_query) - print("downloading google provider for tests...") - res = PyStackQLTestsBase.stackql.executeStmt(registry_pull_google_query) - 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) - -def tearDownModule(): - print("stopping stackql server...") - if PyStackQLTestsBase.server_process: - PyStackQLTestsBase.server_process.terminate() - PyStackQLTestsBase.server_process.wait() - -class PyStackQLNonServerModeTests(PyStackQLTestsBase): - - @pystackql_test_setup() - def test_01_properties_class_method(self): - properties = self.stackql.properties() - # Check that properties is a dictionary - self.assertTrue(isinstance(properties, dict), "properties should be a dictionary") - # List of keys we expect to be in the properties - missing_keys = [key for key in expected_properties if key not in properties] - self.assertTrue(len(missing_keys) == 0, f"Missing keys in properties: {', '.join(missing_keys)}") - # Further type checks (as examples) - self.assertIsInstance(properties["bin_path"], str, "bin_path should be of type str") - self.assertIsInstance(properties["params"], list, "params should be of type list") - self.assertIsInstance(properties["server_mode"], bool, "server_mode should be of type bool") - self.assertIsInstance(properties["output"], str, "output should be of type str") - # If all the assertions pass, then the properties are considered valid. - print_test_result(f"""Test 01 properties method\nPROPERTIES: {properties}""", True) - - @pystackql_test_setup() - def test_02_version_attribute(self): - version = self.stackql.version - self.assertIsNotNone(version) - is_valid_semver = bool(re.match(expected_version_pattern, version)) - self.assertTrue(is_valid_semver) - print_test_result(f"""Test 02 version attribute\nVERSION: {version}""", is_valid_semver) - - @pystackql_test_setup() - def test_03_package_version_attribute(self): - package_version = self.stackql.package_version - self.assertIsNotNone(package_version) - is_valid_semver = bool(re.match(expected_package_version_pattern, package_version)) - self.assertTrue(is_valid_semver) - print_test_result(f"""Test 03 package_version attribute\nPACKAGE VERSION: {package_version}""", is_valid_semver) - - @pystackql_test_setup() - def test_04_platform_attribute(self): - platform_string = self.stackql.platform - self.assertIsNotNone(platform_string) - is_valid_platform = bool(re.match(expected_platform_pattern, platform_string)) - self.assertTrue(is_valid_platform) - print_test_result(f"""Test 04 platform attribute\nPLATFORM: {platform_string}""", is_valid_platform) - - @pystackql_test_setup() - def test_05_bin_path_attribute(self): - self.assertTrue(os.path.exists(self.stackql.bin_path)) - print_test_result(f"""Test 05 bin_path attribute with default download path\nBINARY PATH: {self.stackql.bin_path}""", os.path.exists(self.stackql.bin_path)) - - @pystackql_test_setup(download_dir=get_custom_download_dir(platform.system().lower())) - def test_06_set_custom_download_dir(self): - # Checking that version is not None - version = self.stackql.version - self.assertIsNotNone(version) - # Checking that download_dir is correctly set - expected_download_dir = get_custom_download_dir(platform.system().lower()) - self.assertEqual(self.stackql.download_dir, expected_download_dir, "Download directory is not set correctly.") - # Checking that the binary exists at the expected location - binary_name = 'stackql' if platform.system().lower() != 'windows' else 'stackql.exe' - expected_binary_path = os.path.join(expected_download_dir, binary_name) - self.assertTrue(os.path.exists(expected_binary_path), f"No binary found at {expected_binary_path}") - # Final test result print - print_test_result(f"""Test 06 setting a custom download_dir\nCUSTOM_DOWNLOAD_DIR: {expected_download_dir}""", version is not None and os.path.exists(expected_binary_path)) - - @pystackql_test_setup(output="csv") - def test_07_csv_output_with_defaults(self): - # Check if output is set correctly - self.assertEqual(self.stackql.output, "csv", "Output type is not set to 'csv'") - # Check if sep is set to default (',') - self.assertEqual(self.stackql.sep, ",", "Separator is not set to default ','") - # Check if header is set to default (should be False) - self.assertFalse(self.stackql.header, "Header is not set to default (False)") - # Check if params list has --output and csv - self.assertIn("--output", self.stackql.params) - self.assertIn("csv", self.stackql.params) - # Check if params list has default --delimiter and , - self.assertIn("--delimiter", self.stackql.params) - self.assertIn(",", self.stackql.params) - # Check if params list has --hideheaders (default header value is False) - self.assertIn("--hideheaders", self.stackql.params) - print_test_result(f"""Test 07 csv output with defaults (comma delimited without headers)\nPARAMS: {self.stackql.params}""", True) - - @pystackql_test_setup(output="csv", sep="|") - def test_08_csv_output_with_pipe_separator(self): - # Check if sep is set to '|' - self.assertEqual(self.stackql.sep, "|", "Separator is not set to '|'") - # Check if params list has --delimiter and | - self.assertIn("--delimiter", self.stackql.params) - self.assertIn("|", self.stackql.params) - # Check if --hideheaders is in params list - self.assertIn("--hideheaders", self.stackql.params) - print_test_result(f"""Test 08 csv output with custom sep (pipe delimited without headers)\nPARAMS: {self.stackql.params}""", True) - - @pystackql_test_setup(output="csv", header=True) - def test_09_csv_output_with_header(self): - # Check if header is set to True - self.assertTrue(self.stackql.header, "Header is not set to True") - # Check if params list does not have --hideheaders - self.assertNotIn("--hideheaders", self.stackql.params) - print_test_result(f"""Test 09 csv output with headers (comma delimited with headers)\nPARAMS: {self.stackql.params}""", True) - - @pystackql_test_setup() - def test_10_executeStmt(self): - okta_result_dict = self.stackql.executeStmt(registry_pull_okta_query) - 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}") - 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): - 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{github_result}""", True) - - @pystackql_test_setup(output="pandas") - def test_12_executeStmt_with_pandas_output(self): - 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_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") - is_valid_empty_resp = isinstance(result, list) and len(result) == 0 - self.assertTrue(is_valid_empty_resp, f"Result is not a empty list: {result}") - print_test_result(f"Test 14 execute with defaults (empty response)\nRESULT: {result}", is_valid_empty_resp) - - @pystackql_test_setup(output='pandas') - @patch('pystackql.StackQL.execute') - def test_15_execute_with_pandas_output(self, mock_execute): - # mocking the response for pandas DataFrame - mock_execute.return_value = pd.DataFrame({ - 'status': ['RUNNING', 'TERMINATED'], - 'num_instances': [2, 1] - }) - - result = self.stackql.execute(google_query) - is_valid_dataframe = isinstance(result, pd.DataFrame) - self.assertTrue(is_valid_dataframe, f"Result is not a valid DataFrame: {result}") - # Check datatypes of the columns - expected_dtypes = { - 'status': 'object', - 'num_instances': 'int64' - } - for col, expected_dtype in expected_dtypes.items(): - actual_dtype = result[col].dtype - self.assertEqual(actual_dtype, expected_dtype, f"Column '{col}' has dtype '{actual_dtype}' but expected '{expected_dtype}'") - print_test_result(f"Test 15 execute with pandas output\nRESULT COUNT: {len(result)}", is_valid_dataframe) - - @pystackql_test_setup(output='csv') - @patch('pystackql.StackQL.execute') - def test_16_execute_with_csv_output(self, mock_execute): - # mocking the response for csv output - mock_execute.return_value = "status,num_instances\nRUNNING,2\nTERMINATED,1\n" - result = self.stackql.execute(google_query) - is_valid_csv = isinstance(result, str) and result.count("\n") >= 1 and result.count(",") >= 1 - self.assertTrue(is_valid_csv, f"Result is not a valid CSV: {result}") - print_test_result(f"Test 16 execute with csv output\nRESULT_COUNT: {len(result.splitlines())}", is_valid_csv) - - @pystackql_test_setup() - def test_17_execute_default_auth_dict_output(self): - 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): - # 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": { - "type": "basic", - "username_var": "command_specific_username", - "password_var": "command_specific_password" - } - } - 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) - - @pystackql_test_setup() - def test_19_json_extract_function(self): - query = """ - SELECT - json_extract('{"Key":"StackName","Value":"aws-stack"}', '$.Key') as key, - json_extract('{"Key":"StackName","Value":"aws-stack"}', '$.Value') as value - """ - expected_result = [ - {'key': {'String': 'StackName', 'Valid': True}, 'value': {'String': 'aws-stack', 'Valid': True}} - ] - result = self.stackql.execute(query) - self.assertEqual(result, expected_result, f"Expected result: {expected_result}, got: {result}") - print_test_result(f"Test 19 complex object handling\nRESULT: {result}", result == expected_result) - - -@unittest.skipIf(platform.system() == "Windows", "Skipping async tests on Windows") -class PyStackQLAsyncTests(PyStackQLTestsBase): - - @async_test_decorator - async def test_01_async_executeQueriesAsync(self): - stackql = StackQL() - results = await stackql.executeQueriesAsync(async_queries) - is_valid_results = all(isinstance(res, dict) for res in results) - print_test_result(f"Test 01 executeQueriesAsync with default (dict) output\nRESULT_COUNT: {len(results)}", is_valid_results, is_async=True) - - @async_test_decorator - async def test_02_async_executeQueriesAsync_with_pandas_output(self): - stackql = StackQL(output='pandas') - result = await stackql.executeQueriesAsync(async_queries) - is_valid_dataframe = isinstance(result, pd.DataFrame) and not result.empty - print_test_result(f"Test 02 executeQueriesAsync with pandas output\nRESULT_COUNT: {len(result)}", is_valid_dataframe, is_async=True) - - @async_test_decorator - async def test_03_async_executeQueriesAsync_with_csv_output(self): - stackql = StackQL(output='csv') - exception_caught = False - try: - # This should raise a ValueError since 'csv' output mode is not supported - await stackql.executeQueriesAsync(async_queries) - except ValueError as ve: - exception_caught = str(ve) == "executeQueriesAsync supports only 'dict' or 'pandas' output modes." - except Exception as e: - pass - print_test_result(f"Test 03 executeQueriesAsync with unsupported csv output", exception_caught, is_async=True) - -class PyStackQLServerModeNonAsyncTests(PyStackQLTestsBase): - - @pystackql_test_setup(server_mode=True) - def test_01_server_mode_connectivity(self): - self.assertTrue(self.stackql.server_mode, "StackQL should be in server mode") - self.assertIsNotNone(self.stackql._conn, "Connection object should not be None") - print_test_result("Test 01 server mode connectivity", True, True) - - @pystackql_test_setup(server_mode=True) - def test_02_server_mode_executeStmt(self): - result = self.stackql.executeStmt(registry_pull_google_query) - # 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 02 executeStmt in server mode\n{result}", is_valid_response, True) - - @pystackql_test_setup(server_mode=True, output='pandas') - def test_03_server_mode_executeStmt_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 03 executeStmt in server mode with pandas output\n{result_df}", is_valid_response, True) - - @pystackql_test_setup(server_mode=True) - @patch('pystackql.stackql.StackQL._run_server_query') - def test_04_server_mode_execute_default_output(self, mock_run_server_query): - # Mocking the response as a list of dictionaries - mock_result = [ - {'status': 'RUNNING', 'num_instances': 2}, - {'status': 'TERMINATED', 'num_instances': 1} - ] - mock_run_server_query.return_value = mock_result - - result = self.stackql.execute(google_query) - is_valid_dict_output = isinstance(result, list) and all(isinstance(row, dict) for row in result) - print_test_result(f"""Test 04 execute in server_mode with default output\nRESULT_COUNT: {len(result)}""", is_valid_dict_output, True) - # Check `_run_server_query` method - mock_run_server_query.assert_called_once_with(google_query) - - @pystackql_test_setup(server_mode=True, output='pandas') - @patch('pystackql.stackql.StackQL._run_server_query') - def test_05_server_mode_execute_pandas_output(self, mock_run_server_query): - # Mocking the response for pandas DataFrame - mock_df = pd.DataFrame({ - 'status': ['RUNNING', 'TERMINATED'], - 'num_instances': [2, 1] - }) - mock_run_server_query.return_value = mock_df.to_dict(orient='records') - result = self.stackql.execute(google_query) - is_valid_dataframe = isinstance(result, pd.DataFrame) - self.assertTrue(is_valid_dataframe, f"Result is not a valid DataFrame: {result}") - # Check datatypes of the columns - expected_dtypes = { - 'status': 'object', - 'num_instances': 'int64' - } - for col, expected_dtype in expected_dtypes.items(): - actual_dtype = result[col].dtype - self.assertEqual(actual_dtype, expected_dtype, f"Column '{col}' has dtype '{actual_dtype}' but expected '{expected_dtype}'") - print_test_result(f"Test 05 execute in server_mode with pandas output\nRESULT COUNT: {len(result)}", is_valid_dataframe) - # Check `_run_server_query` method - mock_run_server_query.assert_called_once_with(google_query) - -class MockInteractiveShell: - """A mock class for IPython's InteractiveShell.""" - user_ns = {} # Mock user namespace - - def register_magics(self, magics): - """Mock method to 'register' magics.""" - self.magics = magics - - @staticmethod - def instance(): - """Return a mock instance of the shell.""" - return MockInteractiveShell() - -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() - if self.server_mode: - magics.load_ipython_extension(self.shell) - else: - magic.load_ipython_extension(self.shell) - 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}", 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. - self.stackql_magic.run_query = MagicMock(return_value=self.expected_result) - # Execute the magic with our query. - result = self.stackql_magic.stackql(line=line, cell=cell) - # Validate the outcome. - 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_01_line_magic_query(self): - checks = self.run_magic_test(line=self.query, cell=None) - self.print_test_result("Test 01 Line magic query test", *checks) - - def test_02_cell_magic_query(self): - checks = self.run_magic_test(line="", cell=self.query) - self.print_test_result("Test 02 Cell magic query test", *checks) - - def test_03_cell_magic_query_no_output(self): - checks = self.run_magic_test(line="--no-display", cell=self.query, expect_none=True) - self.print_test_result("Test 03 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_04_line_magic_statement(self): - checks, result = self.run_magic_statement_test(line=self.statement, cell=None) - self.print_test_result(f"Test 04 Line magic statement test\n{result}", *checks) - - def test_05_cell_magic_statement(self): - checks, result = self.run_magic_statement_test(line="", cell=self.statement) - self.print_test_result(f"Test 05 Cell magic statement test\n{result}", *checks) - - def test_06_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"Test 06 Cell magic statement test (with --no-display)\n{result}", *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) - -if __name__ == '__main__': - main() diff --git a/tests/test_async.py b/tests/test_async.py new file mode 100644 index 0000000..f74a32d --- /dev/null +++ b/tests/test_async.py @@ -0,0 +1,141 @@ +# tests/test_async.py + +""" +Async functionality tests for PyStackQL in non-server mode. + +This module tests the async query execution functionality of the StackQL class. +""" + +import os +import sys +import platform +import pytest +import pandas as pd + +# Add the parent directory to the path so we can import from pystackql +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))) + +# Add the current directory to the path so we can import test_constants +sys.path.insert(0, os.path.abspath(os.path.dirname(__file__))) + +from pystackql import StackQL +from tests.test_constants import ( + ASYNC_QUERIES, + print_test_result, + async_test_decorator +) + +# Skip all tests on Windows due to asyncio issues +pytestmark = pytest.mark.skipif( + platform.system() == "Windows", + reason="Skipping async tests on Windows" +) + +class TestAsyncFunctionality: + """Tests for PyStackQL async functionality in non-server mode.""" + + # Helper method to extract value from response objects + def _get_value(self, obj): + """Extract actual value from response objects that might be wrapped in a dict.""" + if isinstance(obj, dict) and 'String' in obj and 'Valid' in obj: + return obj['String'] + return obj + + @async_test_decorator + async def test_execute_queries_async_dict_output(self): + """Test executeQueriesAsync with dict output format.""" + stackql = StackQL() + results = await stackql.executeQueriesAsync(ASYNC_QUERIES) + + # Check result structure + assert isinstance(results, list), "Results should be a list" + assert all(isinstance(item, dict) for item in results), "Each item in results should be a dict" + + # Check result content + assert len(results) > 0, "Results should not be empty" + assert all("formula_name" in item for item in results), "Each item should have 'formula_name' column" + + # Extract formula names, handling possible dictionary format + formula_names = [] + for item in results: + if "formula_name" in item: + formula_name = self._get_value(item["formula_name"]) + formula_names.append(formula_name) + + # Check that we have the expected formula names + assert any("stackql" in str(name) for name in formula_names), "Results should include 'stackql'" + assert any("terraform" in str(name) for name in formula_names), "Results should include 'terraform'" + + print_test_result(f"Async executeQueriesAsync with dict output test\nRESULT COUNT: {len(results)}", + isinstance(results, list) and all(isinstance(item, dict) for item in results), + is_async=True) + + @async_test_decorator + async def test_execute_queries_async_pandas_output(self): + """Test executeQueriesAsync with pandas output format.""" + stackql = StackQL(output='pandas') + result = await stackql.executeQueriesAsync(ASYNC_QUERIES) + + # Check result structure + assert isinstance(result, pd.DataFrame), "Result should be a pandas DataFrame" + assert not result.empty, "DataFrame should not be empty" + assert "formula_name" in result.columns, "DataFrame should have 'formula_name' column" + + # Extract formula names, handling possible dictionary format + formula_values = [] + for i in range(len(result)): + formula_name = result["formula_name"].iloc[i] + if isinstance(formula_name, dict) and 'String' in formula_name: + formula_name = formula_name['String'] + formula_values.append(formula_name) + + # Check that we have the expected formula names + assert any("stackql" in str(name) for name in formula_values), "Results should include 'stackql'" + assert any("terraform" in str(name) for name in formula_values), "Results should include 'terraform'" + + # Check that numeric columns exist + numeric_columns = [ + "installs_30d", "installs_90d", "installs_365d", + "install_on_requests_30d", "install_on_requests_90d", "install_on_requests_365d" + ] + for col in numeric_columns: + assert col in result.columns, f"DataFrame should have '{col}' column" + + # Check that the column can be converted to numeric + try: + if isinstance(result[col].iloc[0], dict) and 'String' in result[col].iloc[0]: + # If it's a dictionary with a String key, try to convert that string to numeric + pd.to_numeric(result[col].iloc[0]['String']) + else: + # Otherwise try to convert the column directly + pd.to_numeric(result[col]) + numeric_conversion_success = True + except (ValueError, TypeError): + numeric_conversion_success = False + + assert numeric_conversion_success, f"Column '{col}' should be convertible to numeric" + + print_test_result(f"Async executeQueriesAsync with pandas output test\nRESULT COUNT: {len(result)}", + isinstance(result, pd.DataFrame) and not result.empty, + is_async=True) + + @async_test_decorator + async def test_execute_queries_async_csv_output(self): + """Test that executeQueriesAsync with csv output raises ValueError.""" + stackql = StackQL(output='csv') + + with pytest.raises(ValueError) as exc_info: + await stackql.executeQueriesAsync(ASYNC_QUERIES) + + # Check exception message + error_msg = str(exc_info.value) + assert "executeQueriesAsync supports only" in error_msg, "Error message should mention supported formats" + assert "dict" in error_msg, "Error message should mention 'dict'" + assert "pandas" in error_msg, "Error message should mention 'pandas'" + + print_test_result(f"Async executeQueriesAsync with csv output test", + "executeQueriesAsync supports only" in error_msg, + is_async=True) + +if __name__ == "__main__": + pytest.main(["-v", __file__]) \ No newline at end of file diff --git a/tests/test_constants.py b/tests/test_constants.py new file mode 100644 index 0000000..86caf1c --- /dev/null +++ b/tests/test_constants.py @@ -0,0 +1,120 @@ +# tests/test_constants.py + +""" +Test constants and helper functions for PyStackQL tests. +""" + +import os +import re +import sys +import time +import platform +import subprocess +from termcolor import colored +import pandas as pd + +# Server connection settings +SERVER_PORT = 5466 +SERVER_ADDRESS = "127.0.0.1" + +# Expected properties and patterns for validation +EXPECTED_PROPERTIES = [ + "bin_path", "params", + "output", "platform", "server_mode", "sha", "version", + "package_version" # Modified: removed "download_dir" as it's no longer exposed +] + +EXPECTED_VERSION_PATTERN = r'^v?(\d+\.\d+\.\d+)$' +EXPECTED_PACKAGE_VERSION_PATTERN = r'^(\d+\.\d+\.\d+)$' +EXPECTED_PLATFORM_PATTERN = r'^(Windows|Linux|Darwin) (\w+) \(([^)]+)\), Python (\d+\.\d+\.\d+)$' + +# Get custom download directory based on platform +def get_custom_download_dir(): + """Return a platform-specific custom download directory.""" + custom_download_dirs = { + 'windows': 'C:\\temp', + 'darwin': '/tmp', + 'linux': '/tmp' + } + return custom_download_dirs.get(platform.system().lower(), '/tmp') + +# Basic test queries that don't require authentication +LITERAL_INT_QUERY = "SELECT 1 as literal_int_value" +LITERAL_FLOAT_QUERY = "SELECT 1.001 as literal_float_value" +LITERAL_STRING_QUERY = "SELECT 'test' as literal_string_value" +EXPRESSION_TRUE_QUERY = "SELECT 1=1 as expression" +EXPRESSION_FALSE_QUERY = "SELECT 1=0 as expression" +EMPTY_RESULT_QUERY = "SELECT 1 WHERE 1=0" +JSON_EXTRACT_QUERY = """ +SELECT + json_extract('{"Key":"StackName","Value":"aws-stack"}', '$.Key') as key, + json_extract('{"Key":"StackName","Value":"aws-stack"}', '$.Value') as value +""" + +# Homebrew provider queries (no authentication required) +HOMEBREW_FORMULA_QUERY = "SELECT name, full_name, tap FROM homebrew.formula.formula WHERE formula_name = 'stackql'" +HOMEBREW_METRICS_QUERY = "SELECT * FROM homebrew.formula.vw_usage_metrics WHERE formula_name = 'stackql'" + +# Registry pull queries +REGISTRY_PULL_HOMEBREW_QUERY = "REGISTRY PULL homebrew" + +# Async test queries +ASYNC_QUERIES = [ + "SELECT * FROM homebrew.formula.vw_usage_metrics WHERE formula_name = 'stackql'", + "SELECT * FROM homebrew.formula.vw_usage_metrics WHERE formula_name = 'terraform'" +] + +# Pattern to match registry pull response +def registry_pull_resp_pattern(provider): + """Returns a regex pattern to match a successful registry pull message.""" + return r"%s provider, version 'v\d+\.\d+\.\d+' successfully installed\s*" % provider + +# Test result printer +def print_test_result(test_name, condition=True, server_mode=False, is_ipython=False, is_async=False): + """Prints a formatted test result. + + Args: + test_name: Name or description of the test + condition: Whether the test passed (True) or failed (False) + server_mode: Whether the test was run in server mode + is_ipython: Whether the test involved IPython magic + is_async: Whether the test involved async functionality + """ + status_header = colored("[PASSED] ", 'green') if condition else colored("[FAILED] ", 'red') + headers = [status_header] + + if server_mode: + headers.append(colored("[SERVER MODE]", 'yellow')) + if is_ipython: + headers.append(colored("[MAGIC EXT]", 'blue')) + if is_async: + headers.append(colored("[ASYNC]", 'magenta')) + + headers.append(test_name) + message = " ".join(headers) + + print("\n" + message) + +# Decorators for test setup +def pystackql_test_setup(**kwargs): + """Decorator to set up a StackQL instance with specified parameters.""" + def decorator(func): + def wrapper(self, *args): + try: + del self.stackql + except AttributeError: + pass + self.stackql = self.StackQL(**kwargs) + func(self, *args) + return wrapper + return decorator + +def async_test_decorator(func): + """Decorator to run async tests with asyncio.""" + def wrapper(*args, **kwargs): + import asyncio + if asyncio.iscoroutinefunction(func): + return asyncio.run(func(*args, **kwargs)) + else: + return func(*args, **kwargs) + return wrapper \ No newline at end of file diff --git a/tests/test_core.py b/tests/test_core.py new file mode 100644 index 0000000..8295dab --- /dev/null +++ b/tests/test_core.py @@ -0,0 +1,143 @@ +# tests/test_core.py + +""" +Core functionality tests for PyStackQL. + +This module tests the basic attributes and properties of the StackQL class. +""" + +import os +import re +import sys +import platform +import pytest + +# Add the parent directory to the path so we can import from pystackql +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))) + +# Add the current directory to the path so we can import test_constants +sys.path.insert(0, os.path.abspath(os.path.dirname(__file__))) + +from pystackql import StackQL +from tests.test_constants import ( + EXPECTED_PROPERTIES, + EXPECTED_VERSION_PATTERN, + EXPECTED_PACKAGE_VERSION_PATTERN, + EXPECTED_PLATFORM_PATTERN, + get_custom_download_dir, + print_test_result, + pystackql_test_setup +) + +class TestStackQLCore: + """Tests for core PyStackQL functionality.""" + + StackQL = StackQL # For use with pystackql_test_setup decorator + + @pystackql_test_setup() + def test_properties_method(self): + """Test that the properties() method returns the expected properties.""" + properties = self.stackql.properties() + + # Check that properties is a dictionary + assert isinstance(properties, dict), "properties should be a dictionary" + + # Check that all expected properties are present + missing_keys = [key for key in EXPECTED_PROPERTIES if key not in properties] + assert len(missing_keys) == 0, f"Missing keys in properties: {', '.join(missing_keys)}" + + # Check property types + assert isinstance(properties["bin_path"], str), "bin_path should be of type str" + assert isinstance(properties["params"], list), "params should be of type list" + assert isinstance(properties["server_mode"], bool), "server_mode should be of type bool" + assert isinstance(properties["output"], str), "output should be of type str" + + print_test_result(f"Properties method test\nPROPERTIES: {properties}", True) + + @pystackql_test_setup() + def test_version_attribute(self): + """Test that the version attribute contains a valid version string.""" + version = self.stackql.version + assert version is not None, "version should not be None" + + is_valid_semver = bool(re.match(EXPECTED_VERSION_PATTERN, version)) + assert is_valid_semver, f"version '{version}' does not match expected pattern" + + print_test_result(f"Version attribute test\nVERSION: {version}", is_valid_semver) + + @pystackql_test_setup() + def test_package_version_attribute(self): + """Test that the package_version attribute contains a valid version string.""" + package_version = self.stackql.package_version + assert package_version is not None, "package_version should not be None" + + is_valid_semver = bool(re.match(EXPECTED_PACKAGE_VERSION_PATTERN, package_version)) + assert is_valid_semver, f"package_version '{package_version}' does not match expected pattern" + + print_test_result(f"Package version attribute test\nPACKAGE VERSION: {package_version}", is_valid_semver) + + @pystackql_test_setup() + def test_platform_attribute(self): + """Test that the platform attribute contains valid platform information.""" + platform_string = self.stackql.platform + assert platform_string is not None, "platform should not be None" + + is_valid_platform = bool(re.match(EXPECTED_PLATFORM_PATTERN, platform_string)) + assert is_valid_platform, f"platform '{platform_string}' does not match expected pattern" + + print_test_result(f"Platform attribute test\nPLATFORM: {platform_string}", is_valid_platform) + + @pystackql_test_setup() + def test_bin_path_attribute(self): + """Test that the bin_path attribute points to an existing binary.""" + assert os.path.exists(self.stackql.bin_path), f"Binary not found at {self.stackql.bin_path}" + + print_test_result(f"Binary path attribute test\nBINARY PATH: {self.stackql.bin_path}", + os.path.exists(self.stackql.bin_path)) + + @pystackql_test_setup(download_dir=get_custom_download_dir()) + def test_custom_download_dir(self): + """Test that a custom download_dir is used correctly.""" + # Check that version is not None (binary was found) + version = self.stackql.version + assert version is not None, "version should not be None" + + # Check that the binary exists at the expected location in the custom directory + expected_download_dir = get_custom_download_dir() + binary_name = 'stackql' if platform.system().lower() != 'windows' else 'stackql.exe' + expected_binary_path = os.path.join(expected_download_dir, binary_name) + + # Check if binary exists + if not os.path.exists(expected_binary_path): + # Give it time to download if needed + import time + time.sleep(5) + + assert os.path.exists(expected_binary_path), f"No binary found at {expected_binary_path}" + + print_test_result(f"Custom download directory test\nCUSTOM_DOWNLOAD_DIR: {expected_download_dir}", + version is not None and os.path.exists(expected_binary_path)) + + @pytest.mark.skip(reason="Skipping upgrade test to avoid unnecessary downloads during regular testing") + @pystackql_test_setup() + def test_upgrade_method(self): + """Test that the upgrade method updates the binary.""" + initial_version = self.stackql.version + initial_sha = self.stackql.sha + + # Perform the upgrade + upgrade_message = self.stackql.upgrade() + + # Check that we got a valid message + assert "stackql upgraded to version" in upgrade_message, "Upgrade message not as expected" + + # Verify that the version attributes were updated + assert self.stackql.version is not None, "version should not be None after upgrade" + assert self.stackql.sha is not None, "sha should not be None after upgrade" + + print_test_result(f"Upgrade method test\nINITIAL VERSION: {initial_version}, SHA: {initial_sha}\n" + f"NEW VERSION: {self.stackql.version}, SHA: {self.stackql.sha}", + "stackql upgraded to version" in upgrade_message) + +if __name__ == "__main__": + pytest.main(["-v", __file__]) \ No newline at end of file diff --git a/tests/test_magic.py b/tests/test_magic.py new file mode 100644 index 0000000..a32ab7c --- /dev/null +++ b/tests/test_magic.py @@ -0,0 +1,122 @@ +# tests/test_magic.py + +""" +Non-server magic extension tests for PyStackQL. + +This module tests the Jupyter magic extensions for StackQL in non-server mode. +""" + +import os +import sys +import re +import pytest +import pandas as pd +from unittest.mock import MagicMock + +# Add the parent directory to the path so we can import from pystackql +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))) + +# Add the current directory to the path so we can import test_constants +sys.path.insert(0, os.path.abspath(os.path.dirname(__file__))) + +# Import directly from the original modules - this is what notebooks would do +from pystackql import magic +from pystackql import StackqlMagic + +from tests.test_constants import ( + LITERAL_INT_QUERY, + REGISTRY_PULL_HOMEBREW_QUERY, + registry_pull_resp_pattern, + print_test_result +) + +class TestStackQLMagic: + """Tests for the non-server mode magic extension.""" + + @pytest.fixture(autouse=True) + def setup_method(self, mock_interactive_shell): + """Set up the test environment.""" + self.shell = mock_interactive_shell + + # Load the magic extension + magic.load_ipython_extension(self.shell) + + # Create the magic instance + self.stackql_magic = StackqlMagic(shell=self.shell) + + # Set up test data + self.query = LITERAL_INT_QUERY + self.expected_result = pd.DataFrame({"literal_int_value": [1]}) + self.statement = REGISTRY_PULL_HOMEBREW_QUERY + + def test_line_magic_query(self): + """Test line magic with a query.""" + # Mock the run_query method to return a known DataFrame + self.stackql_magic.run_query = MagicMock(return_value=self.expected_result) + + # Execute the magic with our query + result = self.stackql_magic.stackql(line=self.query, cell=None) + + # Validate the outcome + assert result.equals(self.expected_result), "Result should match expected DataFrame" + assert 'stackql_df' in self.shell.user_ns, "stackql_df should be in user namespace" + assert self.shell.user_ns['stackql_df'].equals(self.expected_result), "stackql_df should match expected DataFrame" + + print_test_result("Line magic query test", + result.equals(self.expected_result) and + 'stackql_df' in self.shell.user_ns and + self.shell.user_ns['stackql_df'].equals(self.expected_result), + False, True) + + def test_cell_magic_query(self): + """Test cell magic with a query.""" + # Mock the run_query method to return a known DataFrame + self.stackql_magic.run_query = MagicMock(return_value=self.expected_result) + + # Execute the magic with our query + result = self.stackql_magic.stackql(line="", cell=self.query) + + # Validate the outcome + assert result.equals(self.expected_result), "Result should match expected DataFrame" + assert 'stackql_df' in self.shell.user_ns, "stackql_df should be in user namespace" + assert self.shell.user_ns['stackql_df'].equals(self.expected_result), "stackql_df should match expected DataFrame" + + print_test_result("Cell magic query test", + result.equals(self.expected_result) and + 'stackql_df' in self.shell.user_ns and + self.shell.user_ns['stackql_df'].equals(self.expected_result), + False, True) + + def test_cell_magic_query_no_display(self): + """Test cell magic with a query and --no-display option.""" + # Mock the run_query method to return a known DataFrame + self.stackql_magic.run_query = MagicMock(return_value=self.expected_result) + + # Execute the magic with our query and --no-display option + result = self.stackql_magic.stackql(line="--no-display", cell=self.query) + + # Validate the outcome + assert result is None, "Result should be None with --no-display option" + assert 'stackql_df' in self.shell.user_ns, "stackql_df should still be in user namespace" + assert self.shell.user_ns['stackql_df'].equals(self.expected_result), "stackql_df should match expected DataFrame" + + print_test_result("Cell magic query test (with --no-display)", + result is None and + 'stackql_df' in self.shell.user_ns and + self.shell.user_ns['stackql_df'].equals(self.expected_result), + False, True) + +def test_magic_extension_loading(mock_interactive_shell): + """Test that non-server magic extension can be loaded.""" + # Test loading non-server magic + magic.load_ipython_extension(mock_interactive_shell) + assert hasattr(mock_interactive_shell, 'magics'), "Magic should be registered" + assert isinstance(mock_interactive_shell.magics, StackqlMagic), "Registered magic should be StackqlMagic" + + print_test_result("Magic extension loading test", + hasattr(mock_interactive_shell, 'magics') and + isinstance(mock_interactive_shell.magics, StackqlMagic), + False, True) + +if __name__ == "__main__": + pytest.main(["-v", __file__]) \ No newline at end of file diff --git a/tests/test_output_formats.py b/tests/test_output_formats.py new file mode 100644 index 0000000..c5e3912 --- /dev/null +++ b/tests/test_output_formats.py @@ -0,0 +1,213 @@ +# tests/test_output_formats.py + +""" +Output format tests for PyStackQL. + +This module tests the different output formats of the StackQL class. +""" + +import os +import sys +import pytest +import pandas as pd +from unittest.mock import patch + +# Add the parent directory to the path so we can import from pystackql +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))) + +# Add the current directory to the path so we can import test_constants +sys.path.insert(0, os.path.abspath(os.path.dirname(__file__))) + +from pystackql import StackQL +from tests.test_constants import ( + LITERAL_INT_QUERY, + LITERAL_STRING_QUERY, + HOMEBREW_METRICS_QUERY, + print_test_result, + pystackql_test_setup +) + +class TestOutputFormats: + """Tests for PyStackQL output format functionality.""" + + StackQL = StackQL # For use with pystackql_test_setup decorator + + # Helper method to extract value from response objects + def _get_value(self, obj): + """Extract actual value from response objects that might be wrapped in a dict.""" + if isinstance(obj, dict) and 'String' in obj and 'Valid' in obj: + return obj['String'] + return obj + + @pystackql_test_setup() + def test_dict_output_format(self): + """Test that dict output format returns a list of dictionaries.""" + result = self.stackql.execute(LITERAL_INT_QUERY) + + # Check result structure + assert isinstance(result, list), "Result should be a list" + assert all(isinstance(item, dict) for item in result), "Each item in result should be a dict" + + print_test_result(f"Dict output format test\nRESULT: {result}", + isinstance(result, list) and all(isinstance(item, dict) for item in result)) + + @pystackql_test_setup(output='pandas') + def test_pandas_output_format(self): + """Test that pandas output format returns a pandas DataFrame.""" + result = self.stackql.execute(LITERAL_STRING_QUERY) + + # Check result structure + assert isinstance(result, pd.DataFrame), "Result should be a pandas DataFrame" + assert not result.empty, "DataFrame should not be empty" + assert "literal_string_value" in result.columns, "DataFrame should have 'literal_string_value' column" + + # Extract the value, handling possible dictionary format + value = result["literal_string_value"].iloc[0] + if isinstance(value, dict) and 'String' in value and 'Valid' in value: + value = value['String'] + + assert value == "test" or value == '"test"', f"Value should be 'test', got {value}" + + print_test_result(f"Pandas output format test\nRESULT: {result}", + isinstance(result, pd.DataFrame) and "literal_string_value" in result.columns) + + @pystackql_test_setup(output='pandas') + def test_pandas_output_with_numeric_types(self): + """Test that pandas output format handles numeric types correctly.""" + result = self.stackql.execute(HOMEBREW_METRICS_QUERY) + + # Check result structure + assert isinstance(result, pd.DataFrame), "Result should be a pandas DataFrame" + assert not result.empty, "DataFrame should not be empty" + assert "formula_name" in result.columns, "DataFrame should have 'formula_name' column" + + # Check numeric columns - either directly numeric or string representation + numeric_columns = [ + "installs_30d", "installs_90d", "installs_365d", + "install_on_requests_30d", "install_on_requests_90d", "install_on_requests_365d" + ] + + # Validate formula name + formula_name = result["formula_name"].iloc[0] + if isinstance(formula_name, dict) and 'String' in formula_name: + formula_name = formula_name['String'] + + assert "stackql" in str(formula_name), f"Formula name should contain 'stackql', got {formula_name}" + + # Verify numeric columns exist + for col in numeric_columns: + assert col in result.columns, f"DataFrame should have '{col}' column" + + # Try to convert to numeric if possible + try: + if isinstance(result[col].iloc[0], dict) and 'String' in result[col].iloc[0]: + # If it's a dictionary with a String key, try to convert that string to numeric + pd.to_numeric(result[col].iloc[0]['String']) + else: + # Otherwise try to convert the column directly + pd.to_numeric(result[col]) + numeric_conversion_success = True + except (ValueError, TypeError): + numeric_conversion_success = False + + assert numeric_conversion_success, f"Column '{col}' should be convertible to numeric" + + print_test_result(f"Pandas output with numeric types test\nCOLUMNS: {list(result.columns)}", + isinstance(result, pd.DataFrame) and + all(col in result.columns for col in numeric_columns)) + + @pystackql_test_setup(output='csv') + def test_csv_output_format(self): + """Test that csv output format returns a string.""" + result = self.stackql.execute(LITERAL_INT_QUERY) + + # Check result structure + assert isinstance(result, str), "Result should be a string" + # The CSV output might just contain the value (1) or might include the column name + # We'll check for either possibility + assert "1" in result, "Result should contain the value '1'" + + print_test_result(f"CSV output format test\nRESULT: {result}", + isinstance(result, str) and "1" in result) + + @pystackql_test_setup(output='csv') + def test_csv_output_with_pipe_separator(self): + """Test that csv output format with custom separator is configured correctly.""" + # Create a new instance with pipe separator + stackql_with_pipe = StackQL(output='csv', sep='|') + + # Verify that the separator setting is correct + assert stackql_with_pipe.sep == "|", "Separator should be '|'" + assert "--delimiter" in stackql_with_pipe.params, "Params should include '--delimiter'" + assert "|" in stackql_with_pipe.params, "Params should include '|'" + + # Instead of checking the output (which might be affected by other factors), + # we'll focus on verifying that the parameters are set correctly + print_test_result(f"CSV output with pipe separator test\nPARAMS: {stackql_with_pipe.params}", + stackql_with_pipe.sep == "|" and + "--delimiter" in stackql_with_pipe.params and + "|" in stackql_with_pipe.params) + + @pystackql_test_setup(output='csv', header=True) + def test_csv_output_with_header(self): + """Test that csv output format with header works correctly.""" + result = self.stackql.execute(LITERAL_INT_QUERY) + + # Check result structure + assert isinstance(result, str), "Result should be a string" + + # Check that params are set correctly + assert self.stackql.header is True, "Header should be True" + assert "--hideheaders" not in self.stackql.params, "Params should not include '--hideheaders'" + + print_test_result(f"CSV output with header test\nRESULT: {result}", + isinstance(result, str)) + + @pystackql_test_setup(output='csv', header=False) + def test_csv_output_without_header(self): + """Test that csv output format without header works correctly.""" + result = self.stackql.execute(LITERAL_INT_QUERY) + + # Check result structure + assert isinstance(result, str), "Result should be a string" + + # Check that params are set correctly + assert self.stackql.header is False, "Header should be False" + assert "--hideheaders" in self.stackql.params, "Params should include '--hideheaders'" + + print_test_result(f"CSV output without header test\nRESULT: {result}", + isinstance(result, str)) + + def test_invalid_output_format(self): + """Test that an invalid output format raises a ValueError.""" + with pytest.raises(ValueError) as exc_info: + StackQL(output='invalid') + + # Check that the exception message contains the expected elements + # rather than checking for an exact match, which is brittle + error_msg = str(exc_info.value) + assert "Invalid output" in error_msg, "Error message should mention 'Invalid output'" + assert "Expected one of" in error_msg, "Error message should mention 'Expected one of'" + assert "dict" in error_msg, "Error message should mention 'dict'" + assert "pandas" in error_msg, "Error message should mention 'pandas'" + assert "csv" in error_msg, "Error message should mention 'csv'" + assert "invalid" in error_msg, "Error message should mention 'invalid'" + + print_test_result(f"Invalid output format test\nERROR: {error_msg}", + all(text in error_msg for text in ["Invalid output", "Expected one of", "dict", "pandas", "csv", "invalid"])) + + def test_csv_output_in_server_mode(self): + """Test that csv output in server mode raises a ValueError.""" + with pytest.raises(ValueError) as exc_info: + StackQL(server_mode=True, output='csv') + + # Check that the exception message contains the expected elements + error_msg = str(exc_info.value) + assert "CSV output is not supported in server mode" in error_msg, "Error message should mention CSV not supported" + assert "use 'dict' or 'pandas' instead" in error_msg, "Error message should suggest alternatives" + + print_test_result(f"CSV output in server mode test\nERROR: {error_msg}", + "CSV output is not supported in server mode" in error_msg) + +if __name__ == "__main__": + pytest.main(["-v", __file__]) \ No newline at end of file diff --git a/tests/test_params.py b/tests/test_params.py deleted file mode 100644 index a99f19c..0000000 --- a/tests/test_params.py +++ /dev/null @@ -1,85 +0,0 @@ -import json, sys, platform, re, time, unittest, subprocess -import pandas as pd -from termcolor import colored -import unittest.mock as mock - -server_port = 5466 -server_address = "127.0.0.1" - -expected_properties = [ - "bin_path", "download_dir", "package_version", "params", - "output", "platform", "server_mode", "sha", "version" -] - -expected_version_pattern = r'^v?(\d+\.\d+\.\d+)$' -expected_package_version_pattern = r'^(\d+\.\d+\.\d+)$' - -expected_platform_pattern = r'^(Windows|Linux|Darwin) (\w+) \(([^)]+)\), Python (\d+\.\d+\.\d+)$' -# custom_windows_download_dir = 'C:\\temp' -# custom_mac_linux_download_dir = '/tmp' -def get_custom_download_dir(platform_name): - custom_download_dirs = { - 'windows': 'C:\\temp', - 'darwin': '/tmp', - 'linux': '/tmp' - } - return custom_download_dirs.get(platform_name, '/tmp') - -registry_pull_google_query = "REGISTRY PULL google" -registry_pull_aws_query = "REGISTRY PULL aws" -registry_pull_okta_query = "REGISTRY PULL okta" -registry_pull_github_query = "REGISTRY PULL github" -registry_pull_homebrew_query = "REGISTRY PULL homebrew" - -def registry_pull_resp_pattern(provider): - return r"%s provider, version 'v\d+\.\d+\.\d+' successfully installed\s*" % provider - -test_gcp_project_id = "test-gcp-project" -test_gcp_zone = "australia-southeast2-a" - -github_query = "select login from github.users.users" - -google_query = f""" -SELECT status, count(*) as num_instances -FROM google.compute.instances -WHERE project = '{test_gcp_project_id}' -AND zone = '{test_gcp_zone}' -GROUP BY status -""" -google_show_services_query = "SHOW SERVICES IN google" - -aws_query = f""" -SELECT -instance_type, -count(*) as num_instances -FROM aws.ec2.instances -WHERE region = 'ap-southeast-2' -GROUP BY instance_type -""" - -test_aws_regions = ["ap-southeast-2", "ap-southeast-4"] - -async_queries = [ - f""" - SELECT region, COUNT(*) as num_functions - FROM aws.lambda.functions - WHERE region = '{region}' - """ - for region in test_aws_regions -] - -def print_test_result(test_name, condition=True, server_mode=False, is_ipython=False, is_async=False): - status_header = colored("[PASSED] ", 'green') if condition else colored("[FAILED] ", 'red') - headers = [status_header] - - if server_mode: - headers.append(colored("[SERVER MODE]", 'yellow')) - if is_ipython: - headers.append(colored("[MAGIC EXT]", 'blue')) - if is_async: - headers.append(colored("[ASYNC]", 'magenta')) - - headers.append(test_name) - message = " ".join(headers) - - print("\n" + message) diff --git a/tests/test_query_execution.py b/tests/test_query_execution.py new file mode 100644 index 0000000..1bd4402 --- /dev/null +++ b/tests/test_query_execution.py @@ -0,0 +1,270 @@ +# tests/test_query_execution.py + +""" +Query execution tests for PyStackQL. + +This module tests the query execution functionality of the StackQL class. +""" + +import os +import re +import sys +import pytest + +# Add the parent directory to the path so we can import from pystackql +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))) + +# Add the current directory to the path so we can import test_constants +sys.path.insert(0, os.path.abspath(os.path.dirname(__file__))) + +from pystackql import StackQL +from tests.test_constants import ( + LITERAL_INT_QUERY, + LITERAL_FLOAT_QUERY, + LITERAL_STRING_QUERY, + EXPRESSION_TRUE_QUERY, + EXPRESSION_FALSE_QUERY, + EMPTY_RESULT_QUERY, + JSON_EXTRACT_QUERY, + HOMEBREW_FORMULA_QUERY, + HOMEBREW_METRICS_QUERY, + REGISTRY_PULL_HOMEBREW_QUERY, + registry_pull_resp_pattern, + print_test_result, + pystackql_test_setup +) + +class TestQueryExecution: + """Tests for PyStackQL query execution functionality.""" + + StackQL = StackQL # For use with pystackql_test_setup decorator + + # Helper method to extract value from response objects + def _get_value(self, obj): + """Extract actual value from response objects that might be wrapped in a dict.""" + if isinstance(obj, dict) and 'String' in obj and 'Valid' in obj: + return obj['String'] + return obj + + # Helper method to check if a value is numeric + def _is_numeric(self, obj): + """Check if a value is numeric, handling both direct numbers and string representations.""" + if isinstance(obj, (int, float)): + return True + if isinstance(obj, dict) and 'String' in obj and 'Valid' in obj: + try: + float(obj['String']) + return True + except (ValueError, TypeError): + return False + if isinstance(obj, str): + try: + float(obj) + return True + except (ValueError, TypeError): + return False + return False + + @pystackql_test_setup() + def test_execute_literal_int(self): + """Test executing a query with a literal integer value.""" + result = self.stackql.execute(LITERAL_INT_QUERY) + + # Check result structure + assert isinstance(result, list), "Result should be a list" + assert len(result) == 1, "Result should have exactly one row" + assert "literal_int_value" in result[0], "Result should have 'literal_int_value' column" + + # Check the value - allow for either direct int or wrapped dict + value = self._get_value(result[0]["literal_int_value"]) + assert value == "1" or value == 1, f"Result value should be 1, got {value}" + + print_test_result(f"Execute literal int query test\nRESULT: {result}", + value == "1" or value == 1) + + @pystackql_test_setup() + def test_execute_literal_float(self): + """Test executing a query with a literal float value.""" + result = self.stackql.execute(LITERAL_FLOAT_QUERY) + + # Check result structure + assert isinstance(result, list), "Result should be a list" + assert len(result) == 1, "Result should have exactly one row" + assert "literal_float_value" in result[0], "Result should have 'literal_float_value' column" + + # Check the value - allow for either direct float or wrapped dict + value = self._get_value(result[0]["literal_float_value"]) + assert value == "1.001" or value == 1.001, f"Result value should be 1.001, got {value}" + + print_test_result(f"Execute literal float query test\nRESULT: {result}", + value == "1.001" or value == 1.001) + + @pystackql_test_setup() + def test_execute_literal_string(self): + """Test executing a query with a literal string value.""" + result = self.stackql.execute(LITERAL_STRING_QUERY) + + # Check result structure + assert isinstance(result, list), "Result should be a list" + assert len(result) == 1, "Result should have exactly one row" + assert "literal_string_value" in result[0], "Result should have 'literal_string_value' column" + + # Check the value - allow for either direct string or wrapped dict + value = self._get_value(result[0]["literal_string_value"]) + assert value == "test", f"Result value should be 'test', got {value}" + + print_test_result(f"Execute literal string query test\nRESULT: {result}", + value == "test") + + @pystackql_test_setup() + def test_execute_expression_true(self): + """Test executing a query with a true expression.""" + result = self.stackql.execute(EXPRESSION_TRUE_QUERY) + + # Check result structure + assert isinstance(result, list), "Result should be a list" + assert len(result) == 1, "Result should have exactly one row" + assert "expression" in result[0], "Result should have 'expression' column" + + # Check the value - allow for either direct int or wrapped dict + value = self._get_value(result[0]["expression"]) + assert value == "1" or value == 1, f"Result value should be 1 (true), got {value}" + + print_test_result(f"Execute true expression query test\nRESULT: {result}", + value == "1" or value == 1) + + @pystackql_test_setup() + def test_execute_expression_false(self): + """Test executing a query with a false expression.""" + result = self.stackql.execute(EXPRESSION_FALSE_QUERY) + + # Check result structure + assert isinstance(result, list), "Result should be a list" + assert len(result) == 1, "Result should have exactly one row" + assert "expression" in result[0], "Result should have 'expression' column" + + # Check the value - allow for either direct int or wrapped dict + value = self._get_value(result[0]["expression"]) + assert value == "0" or value == 0, f"Result value should be 0 (false), got {value}" + + print_test_result(f"Execute false expression query test\nRESULT: {result}", + value == "0" or value == 0) + + @pystackql_test_setup() + def test_execute_empty_result(self): + """Test executing a query that returns an empty result.""" + result = self.stackql.execute(EMPTY_RESULT_QUERY) + + # Check result structure + assert isinstance(result, list), "Result should be a list" + assert len(result) == 0, "Result should be empty" + + print_test_result(f"Execute empty result query test\nRESULT: {result}", len(result) == 0) + + @pystackql_test_setup() + def test_execute_json_extract(self): + """Test executing a query that uses the json_extract function.""" + result = self.stackql.execute(JSON_EXTRACT_QUERY) + + # Check result structure + assert isinstance(result, list), "Result should be a list" + assert len(result) == 1, "Result should have exactly one row" + assert "key" in result[0], "Result should have 'key' column" + assert "value" in result[0], "Result should have 'value' column" + + # Get the extracted values - complex objects might be returned directly or as wrapped dicts + key_value = self._get_value(result[0]["key"]) + value_value = self._get_value(result[0]["value"]) + + # Check for either JSON objects or string values + if isinstance(key_value, dict): + assert key_value.get("String") == "StackName" or key_value == "StackName", "Key should be 'StackName'" + assert value_value.get("String") == "aws-stack" or value_value == "aws-stack", "Value should be 'aws-stack'" + else: + assert "StackName" in str(key_value), "Key should contain 'StackName'" + assert "aws-stack" in str(value_value), "Value should contain 'aws-stack'" + + print_test_result(f"Execute JSON extract query test\nRESULT: {result}", + "StackName" in str(key_value) and "aws-stack" in str(value_value)) + + @pystackql_test_setup() + def test_execute_homebrew_formula(self): + """Test executing a query against the homebrew.formula.formula table.""" + result = self.stackql.execute(HOMEBREW_FORMULA_QUERY) + + # Check result structure + assert isinstance(result, list), "Result should be a list" + assert len(result) == 1, "Result should have exactly one row" + assert "name" in result[0], "Result should have 'name' column" + assert "full_name" in result[0], "Result should have 'full_name' column" + assert "tap" in result[0], "Result should have 'tap' column" + + # Check formula values - allowing for either direct values or wrapped dicts + name_value = self._get_value(result[0]["name"]) + full_name_value = self._get_value(result[0]["full_name"]) + tap_value = self._get_value(result[0]["tap"]) + + assert name_value == "stackql" or name_value == '"stackql"', f"Name should be 'stackql', got {name_value}" + assert full_name_value == "stackql" or full_name_value == '"stackql"', f"Full name should be 'stackql', got {full_name_value}" + assert tap_value == "homebrew/core" or tap_value == '"homebrew/core"', f"Tap should be 'homebrew/core', got {tap_value}" + + print_test_result(f"Execute homebrew formula query test\nRESULT: {result}", + "stackql" in str(name_value) and + "stackql" in str(full_name_value) and + "homebrew/core" in str(tap_value)) + + @pystackql_test_setup() + def test_execute_homebrew_metrics(self): + """Test executing a query against the homebrew.formula.vw_usage_metrics view.""" + result = self.stackql.execute(HOMEBREW_METRICS_QUERY) + + # Check result structure + assert isinstance(result, list), "Result should be a list" + assert len(result) == 1, "Result should have exactly one row" + + # Check column names (not values as they change over time) + expected_columns = [ + "formula_name", "installs_30d", "installs_90d", "installs_365d", + "install_on_requests_30d", "install_on_requests_90d", "install_on_requests_365d" + ] + for col in expected_columns: + assert col in result[0], f"Result should have '{col}' column" + + # Check formula name + formula_name = self._get_value(result[0]["formula_name"]) + assert formula_name == "stackql" or formula_name == '"stackql"', f"Formula name should be 'stackql', got {formula_name}" + + # Check data types - should be numeric or string representations of numbers + for col in expected_columns[1:]: # Skip formula_name + assert self._is_numeric(result[0][col]), f"Column '{col}' should be numeric or string representation of a number" + + print_test_result(f"Execute homebrew metrics query test\nCOLUMNS: {list(result[0].keys())}", + all(col in result[0] for col in expected_columns) and + "stackql" in str(formula_name)) + + @pystackql_test_setup() + def test_execute_stmt_registry_pull(self): + """Test executing a registry pull statement.""" + result = self.stackql.executeStmt(REGISTRY_PULL_HOMEBREW_QUERY) + + # Check result structure (depends on output format) + if self.stackql.output == 'dict': + assert 'message' in result, "Result should have 'message' key" + message = result['message'] + elif self.stackql.output == 'pandas': + assert 'message' in result.columns, "Result should have 'message' column" + message = result['message'].iloc[0] + elif self.stackql.output == 'csv': + message = result + else: + message = str(result) + + # Check that the message matches the expected pattern + expected_pattern = registry_pull_resp_pattern("homebrew") + assert re.search(expected_pattern, message), f"Message '{message}' does not match expected pattern" + + print_test_result(f"Execute registry pull statement test\nRESULT: {result}", + re.search(expected_pattern, message) is not None) + +if __name__ == "__main__": + pytest.main(["-v", __file__]) \ No newline at end of file diff --git a/tests/test_server.py b/tests/test_server.py new file mode 100644 index 0000000..312f0f1 --- /dev/null +++ b/tests/test_server.py @@ -0,0 +1,224 @@ +# tests/test_server.py + +""" +Server mode tests for PyStackQL. + +This module tests the server mode functionality of the StackQL class. +""" + +import re +import pytest +import pandas as pd +from unittest.mock import patch +from pystackql import StackQL +from test_constants import ( + LITERAL_INT_QUERY, + LITERAL_STRING_QUERY, + HOMEBREW_FORMULA_QUERY, + REGISTRY_PULL_HOMEBREW_QUERY, + print_test_result, + pystackql_test_setup +) + +@pytest.mark.usefixtures("stackql_server") +class TestServerMode: + """Tests for PyStackQL server mode functionality.""" + + StackQL = StackQL # For use with pystackql_test_setup decorator + + @pystackql_test_setup(server_mode=True) + def test_server_mode_connectivity(self): + """Test that server mode connects successfully.""" + assert self.stackql.server_mode, "StackQL should be in server mode" + # Updated assertion to check server_connection attribute instead of _conn + assert hasattr(self.stackql, 'server_connection'), "StackQL should have a server_connection attribute" + assert self.stackql.server_connection is not None, "Server connection object should not be None" + + print_test_result("Server mode connectivity test", + self.stackql.server_mode and + hasattr(self.stackql, 'server_connection') and + self.stackql.server_connection is not None, + True) + + @pystackql_test_setup(server_mode=True) + def test_server_mode_execute_stmt(self): + """Test executeStmt in server mode.""" + result = self.stackql.executeStmt(REGISTRY_PULL_HOMEBREW_QUERY) + + # Check result structure + assert isinstance(result, list), "Result should be a list" + assert len(result) == 1, "Result should have exactly one item" + assert "message" in result[0], "Result should have a 'message' key" + assert result[0]["message"] == "OK", "Message should be 'OK'" + + print_test_result(f"Server mode executeStmt test\nRESULT: {result}", + isinstance(result, list) and + len(result) == 1 and + result[0]["message"] == "OK", + True) + + @pystackql_test_setup(server_mode=True, output='pandas') + def test_server_mode_execute_stmt_pandas(self): + """Test executeStmt in server mode with pandas output.""" + result = self.stackql.executeStmt(REGISTRY_PULL_HOMEBREW_QUERY) + + # Check result structure + assert isinstance(result, pd.DataFrame), "Result should be a pandas DataFrame" + assert not result.empty, "DataFrame should not be empty" + assert "message" in result.columns, "DataFrame should have a 'message' column" + assert result["message"].iloc[0] == "OK", "Message should be 'OK'" + + print_test_result(f"Server mode executeStmt with pandas output test\nRESULT: {result}", + isinstance(result, pd.DataFrame) and + not result.empty and + result["message"].iloc[0] == "OK", + True) + + @pystackql_test_setup(server_mode=True) + def test_server_mode_execute(self): + """Test execute in server mode.""" + result = self.stackql.execute(LITERAL_INT_QUERY) + + # Check result structure + assert isinstance(result, list), "Result should be a list" + assert len(result) == 1, "Result should have exactly one item" + assert "literal_int_value" in result[0], "Result should have a 'literal_int_value' key" + # Update assertion to handle string value from server + literal_value = result[0]["literal_int_value"] + if isinstance(literal_value, str): + literal_value = int(literal_value) + assert literal_value == 1, "Value should be 1" + + print_test_result(f"Server mode execute test\nRESULT: {result}", + isinstance(result, list) and + len(result) == 1 and + int(result[0]["literal_int_value"]) == 1, + True) + + @pystackql_test_setup(server_mode=True, output='pandas') + def test_server_mode_execute_pandas(self): + """Test execute in server mode with pandas output.""" + result = self.stackql.execute(LITERAL_STRING_QUERY) + + # Check result structure + assert isinstance(result, pd.DataFrame), "Result should be a pandas DataFrame" + assert not result.empty, "DataFrame should not be empty" + assert "literal_string_value" in result.columns, "DataFrame should have a 'literal_string_value' column" + assert result["literal_string_value"].iloc[0] == "test", "Value should be 'test'" + + print_test_result(f"Server mode execute with pandas output test\nRESULT: {result}", + isinstance(result, pd.DataFrame) and + not result.empty and + result["literal_string_value"].iloc[0] == "test", + True) + + @pystackql_test_setup(server_mode=True) + def test_server_mode_provider_query(self): + """Test querying a provider in server mode.""" + result = self.stackql.execute(HOMEBREW_FORMULA_QUERY) + + # Check result structure + assert isinstance(result, list), "Result should be a list" + assert len(result) == 1, "Result should have exactly one item" + assert "name" in result[0], "Result should have a 'name' key" + assert "full_name" in result[0], "Result should have a 'full_name' key" + assert "tap" in result[0], "Result should have a 'tap' key" + assert result[0]["name"] == "stackql", "Name should be 'stackql'" + + print_test_result(f"Server mode provider query test\nRESULT: {result}", + isinstance(result, list) and + len(result) == 1 and + result[0]["name"] == "stackql", + True) + + # Update mocked tests to use execute_query instead of _run_server_query + @patch('pystackql.core.server.ServerConnection.execute_query') + def test_server_mode_execute_mocked(self, mock_execute_query): + """Test execute in server mode with mocked server response.""" + # Create a StackQL instance in server mode + stackql = StackQL(server_mode=True) + + # Mock the server response + mock_result = [{"literal_int_value": 1}] + mock_execute_query.return_value = mock_result + + # Execute the query + result = stackql.execute(LITERAL_INT_QUERY) + + # Check that the mock was called with the correct query + mock_execute_query.assert_called_once_with(LITERAL_INT_QUERY) + + # Check result structure + assert result == mock_result, "Result should match the mocked result" + + print_test_result(f"Server mode execute test (mocked)\nRESULT: {result}", + result == mock_result, + True) + + @patch('pystackql.core.server.ServerConnection.execute_query') + def test_server_mode_execute_pandas_mocked(self, mock_execute_query): + """Test execute in server mode with pandas output and mocked server response.""" + # Create a StackQL instance in server mode with pandas output + stackql = StackQL(server_mode=True, output='pandas') + + # Mock the server response + mock_result = [{"literal_string_value": "test"}] + mock_execute_query.return_value = mock_result + + # Execute the query + result = stackql.execute(LITERAL_STRING_QUERY) + + # Check that the mock was called with the correct query + mock_execute_query.assert_called_once_with(LITERAL_STRING_QUERY) + + # Check result structure + assert isinstance(result, pd.DataFrame), "Result should be a pandas DataFrame" + assert not result.empty, "DataFrame should not be empty" + assert "literal_string_value" in result.columns, "DataFrame should have a 'literal_string_value' column" + assert result["literal_string_value"].iloc[0] == "test", "Value should be 'test'" + + print_test_result(f"Server mode execute with pandas output test (mocked)\nRESULT: {result}", + isinstance(result, pd.DataFrame) and + not result.empty and + result["literal_string_value"].iloc[0] == "test", + True) + + @patch('pystackql.core.server.ServerConnection.execute_query') + def test_server_mode_execute_stmt_mocked(self, mock_execute_query): + """Test executeStmt in server mode with mocked server response.""" + # Create a StackQL instance in server mode + stackql = StackQL(server_mode=True) + + # Mock the server response + mock_result = [{"message": "OK"}] + mock_execute_query.return_value = mock_result + + # Execute the statement + result = stackql.executeStmt(REGISTRY_PULL_HOMEBREW_QUERY) + + # Check that the mock was called with the correct query and is_statement=True + mock_execute_query.assert_called_once_with(REGISTRY_PULL_HOMEBREW_QUERY, is_statement=True) + + # Check result structure + assert result == mock_result, "Result should match the mocked result" + + print_test_result(f"Server mode executeStmt test (mocked)\nRESULT: {result}", + result == mock_result, + True) + + def test_server_mode_csv_output_error(self): + """Test that server mode with csv output raises an error.""" + with pytest.raises(ValueError) as exc_info: + StackQL(server_mode=True, output='csv') + + # Check exception message + expected_message = "CSV output is not supported in server mode, use 'dict' or 'pandas' instead." + assert str(exc_info.value) == expected_message, f"Exception message '{str(exc_info.value)}' does not match expected" + + print_test_result(f"Server mode with csv output error test", + str(exc_info.value) == expected_message, + True) + + +if __name__ == "__main__": + pytest.main(["-v", __file__]) \ No newline at end of file diff --git a/tests/test_server_magic.py b/tests/test_server_magic.py new file mode 100644 index 0000000..fa0a9e2 --- /dev/null +++ b/tests/test_server_magic.py @@ -0,0 +1,120 @@ +""" +Server-mode magic extension tests for PyStackQL. + +This module tests the Jupyter magic extensions for StackQL in server mode. +""" + +import os +import sys +import re +import pytest +import pandas as pd +from unittest.mock import MagicMock + +# Add the parent directory to the path so we can import from pystackql +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))) + +# Add the current directory to the path so we can import test_constants +sys.path.insert(0, os.path.abspath(os.path.dirname(__file__))) + +# Import directly from the original modules - this is what notebooks would do +from pystackql import magics +from pystackql import StackqlServerMagic + +from tests.test_constants import ( + LITERAL_INT_QUERY, + REGISTRY_PULL_HOMEBREW_QUERY, + registry_pull_resp_pattern, + print_test_result +) + +class TestStackQLServerMagic: + """Tests for the server mode magic extension.""" + + @pytest.fixture(autouse=True) + def setup_method(self, mock_interactive_shell): + """Set up the test environment.""" + self.shell = mock_interactive_shell + + # Load the magic extension + magics.load_ipython_extension(self.shell) + + # Create the magic instance + self.stackql_magic = StackqlServerMagic(shell=self.shell) + + # Set up test data + self.query = LITERAL_INT_QUERY + self.expected_result = pd.DataFrame({"literal_int_value": [1]}) + self.statement = REGISTRY_PULL_HOMEBREW_QUERY + + def test_line_magic_query(self): + """Test line magic with a query in server mode.""" + # Mock the run_query method to return a known DataFrame + self.stackql_magic.run_query = MagicMock(return_value=self.expected_result) + + # Execute the magic with our query + result = self.stackql_magic.stackql(line=self.query, cell=None) + + # Validate the outcome + assert result.equals(self.expected_result), "Result should match expected DataFrame" + assert 'stackql_df' in self.shell.user_ns, "stackql_df should be in user namespace" + assert self.shell.user_ns['stackql_df'].equals(self.expected_result), "stackql_df should match expected DataFrame" + + print_test_result("Line magic query test (server mode)", + result.equals(self.expected_result) and + 'stackql_df' in self.shell.user_ns and + self.shell.user_ns['stackql_df'].equals(self.expected_result), + True, True) + + def test_cell_magic_query(self): + """Test cell magic with a query in server mode.""" + # Mock the run_query method to return a known DataFrame + self.stackql_magic.run_query = MagicMock(return_value=self.expected_result) + + # Execute the magic with our query + result = self.stackql_magic.stackql(line="", cell=self.query) + + # Validate the outcome + assert result.equals(self.expected_result), "Result should match expected DataFrame" + assert 'stackql_df' in self.shell.user_ns, "stackql_df should be in user namespace" + assert self.shell.user_ns['stackql_df'].equals(self.expected_result), "stackql_df should match expected DataFrame" + + print_test_result("Cell magic query test (server mode)", + result.equals(self.expected_result) and + 'stackql_df' in self.shell.user_ns and + self.shell.user_ns['stackql_df'].equals(self.expected_result), + True, True) + + def test_cell_magic_query_no_display(self): + """Test cell magic with a query and --no-display option in server mode.""" + # Mock the run_query method to return a known DataFrame + self.stackql_magic.run_query = MagicMock(return_value=self.expected_result) + + # Execute the magic with our query and --no-display option + result = self.stackql_magic.stackql(line="--no-display", cell=self.query) + + # Validate the outcome + assert result is None, "Result should be None with --no-display option" + assert 'stackql_df' in self.shell.user_ns, "stackql_df should still be in user namespace" + assert self.shell.user_ns['stackql_df'].equals(self.expected_result), "stackql_df should match expected DataFrame" + + print_test_result("Cell magic query test with --no-display (server mode)", + result is None and + 'stackql_df' in self.shell.user_ns and + self.shell.user_ns['stackql_df'].equals(self.expected_result), + True, True) + +def test_server_magic_extension_loading(mock_interactive_shell): + """Test that server magic extension can be loaded.""" + # Test loading server magic + magics.load_ipython_extension(mock_interactive_shell) + assert hasattr(mock_interactive_shell, 'magics'), "Magic should be registered" + assert isinstance(mock_interactive_shell.magics, StackqlServerMagic), "Registered magic should be StackqlServerMagic" + + print_test_result("Server magic extension loading test", + hasattr(mock_interactive_shell, 'magics') and + isinstance(mock_interactive_shell.magics, StackqlServerMagic), + True, True) + +if __name__ == "__main__": + pytest.main(["-v", __file__]) \ No newline at end of file From 32c46112dfa9356fc44edf859acb012415e023b2 Mon Sep 17 00:00:00 2001 From: Jeffrey Aven Date: Wed, 4 Jun 2025 16:20:31 +1000 Subject: [PATCH 02/31] ci update --- .github/workflows/test.yaml | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/.github/workflows/test.yaml b/.github/workflows/test.yaml index 8bf5a0a..eb48153 100644 --- a/.github/workflows/test.yaml +++ b/.github/workflows/test.yaml @@ -28,10 +28,11 @@ jobs: name: 'Run Tests on ${{matrix.os}} with Python ${{matrix.python-version}}' steps: - - uses: actions/checkout@v4 + - name: Checkout + uses: actions/checkout@v4.2.2 - name: Set up Python ${{ matrix.python-version }} - uses: actions/setup-python@v5 + uses: actions/setup-python@v5.6.0 with: python-version: ${{ matrix.python-version }} @@ -55,8 +56,8 @@ jobs: python3 run_tests.py shell: bash - - name: Setup StackQL - uses: stackql/setup-stackql@v2 + - name: setup-stackql + uses: stackql/setup-stackql@v2.2.3 with: platform: ${{ matrix.os == 'windows-latest' && 'windows' || (matrix.os == 'macos-latest' && 'darwin' || 'linux') }} From 3f760b91f08a263fccd3e27e12109ec5c42cbf0b Mon Sep 17 00:00:00 2001 From: Jeffrey Aven Date: Wed, 4 Jun 2025 16:43:24 +1000 Subject: [PATCH 03/31] ci update --- .github/workflows/test.yaml | 30 +++++++++++++++++++++++++----- run_tests.py | 7 +------ 2 files changed, 26 insertions(+), 11 deletions(-) diff --git a/.github/workflows/test.yaml b/.github/workflows/test.yaml index eb48153..25bd701 100644 --- a/.github/workflows/test.yaml +++ b/.github/workflows/test.yaml @@ -49,17 +49,37 @@ jobs: run: | pip install pytest>=6.2.5 pytest-cov>=2.12.0 nose>=1.3.7 + - name: setup-stackql + uses: stackql/setup-stackql@v2.2.3 + with: + platform: ${{ matrix.os == 'windows-latest' && 'windows' || (matrix.os == 'macos-latest' && 'darwin' || 'linux') }} + + - name: List dir + run: | + ls -la + + - name: Move stackql binary to temp dir (Unix/macOS) + if: matrix.os != 'windows-latest' + shell: bash + run: | + mkdir -p /tmp || true + cp stackql /tmp/stackql + echo "StackQL binary moved to /tmp/stackql" + + - name: Move stackql binary to temp dir (Windows) + if: matrix.os == 'windows-latest' + shell: cmd + run: | + mkdir "C:\Temp" 2>nul || echo "Temp directory already exists" + copy stackql.exe "C:\Temp\stackql.exe" + echo StackQL binary moved to C:\Temp\stackql.exe + - name: Run non-server tests env: GITHUB_ACTIONS: 'true' run: | python3 run_tests.py shell: bash - - - name: setup-stackql - uses: stackql/setup-stackql@v2.2.3 - with: - platform: ${{ matrix.os == 'windows-latest' && 'windows' || (matrix.os == 'macos-latest' && 'darwin' || 'linux') }} - name: Start StackQL server run: | diff --git a/run_tests.py b/run_tests.py index d268ea2..930ba58 100644 --- a/run_tests.py +++ b/run_tests.py @@ -46,12 +46,7 @@ def main(): "tests/test_magic.py", "tests/test_async.py" ]) - - # Skip server tests by default as they require a running server - # Uncomment to run server tests - # args.append("tests/test_server.py") - # args.append("tests/test_server_magic.py") - + # Run pytest with the arguments return pytest.main(args) From fda73256919cf60caf228ff9ccc2e7b16e344958 Mon Sep 17 00:00:00 2001 From: Jeffrey Aven Date: Wed, 4 Jun 2025 16:44:55 +1000 Subject: [PATCH 04/31] ci update --- .github/workflows/test.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/test.yaml b/.github/workflows/test.yaml index 25bd701..6086fcb 100644 --- a/.github/workflows/test.yaml +++ b/.github/workflows/test.yaml @@ -56,7 +56,7 @@ jobs: - name: List dir run: | - ls -la + which stackql || echo "StackQL binary not found" - name: Move stackql binary to temp dir (Unix/macOS) if: matrix.os != 'windows-latest' From b3d374d53413c555bf2ed852d984cbb86ad0a982 Mon Sep 17 00:00:00 2001 From: Jeffrey Aven Date: Wed, 4 Jun 2025 16:48:41 +1000 Subject: [PATCH 05/31] ci update --- .github/workflows/test.yaml | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/.github/workflows/test.yaml b/.github/workflows/test.yaml index 6086fcb..9176850 100644 --- a/.github/workflows/test.yaml +++ b/.github/workflows/test.yaml @@ -52,7 +52,10 @@ jobs: - name: setup-stackql uses: stackql/setup-stackql@v2.2.3 with: - platform: ${{ matrix.os == 'windows-latest' && 'windows' || (matrix.os == 'macos-latest' && 'darwin' || 'linux') }} + use_wrapper: true + + # with: + # platform: ${{ matrix.os == 'windows-latest' && 'windows' || (matrix.os == 'macos-latest' && 'darwin' || 'linux') }} - name: List dir run: | From 29b18461f52f7ccaa7b43c0810d68b59dde6e4e2 Mon Sep 17 00:00:00 2001 From: Jeffrey Aven Date: Wed, 4 Jun 2025 16:51:49 +1000 Subject: [PATCH 06/31] ci update --- .github/workflows/test.yaml | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/.github/workflows/test.yaml b/.github/workflows/test.yaml index 9176850..52f44bf 100644 --- a/.github/workflows/test.yaml +++ b/.github/workflows/test.yaml @@ -61,6 +61,10 @@ jobs: run: | which stackql || echo "StackQL binary not found" + - name: Call stackql + run: | + stackql --version + - name: Move stackql binary to temp dir (Unix/macOS) if: matrix.os != 'windows-latest' shell: bash From f676e857f922fc8945b04da6d1da08768c4c8d54 Mon Sep 17 00:00:00 2001 From: Jeffrey Aven Date: Thu, 5 Jun 2025 11:13:47 +1000 Subject: [PATCH 07/31] ci update --- .github/workflows/test.yaml | 18 +++++++----------- 1 file changed, 7 insertions(+), 11 deletions(-) diff --git a/.github/workflows/test.yaml b/.github/workflows/test.yaml index 52f44bf..4637b6d 100644 --- a/.github/workflows/test.yaml +++ b/.github/workflows/test.yaml @@ -54,14 +54,7 @@ jobs: with: use_wrapper: true - # with: - # platform: ${{ matrix.os == 'windows-latest' && 'windows' || (matrix.os == 'macos-latest' && 'darwin' || 'linux') }} - - - name: List dir - run: | - which stackql || echo "StackQL binary not found" - - - name: Call stackql + - name: Show stackql version run: | stackql --version @@ -69,16 +62,19 @@ jobs: if: matrix.os != 'windows-latest' shell: bash run: | + STACKQL_PATH=$(which stackql) mkdir -p /tmp || true - cp stackql /tmp/stackql - echo "StackQL binary moved to /tmp/stackql" + cp "$STACKQL_PATH" /tmp/stackql + echo "StackQL binary moved from ${STACKQL_PATH} to /tmp/stackql" - name: Move stackql binary to temp dir (Windows) if: matrix.os == 'windows-latest' shell: cmd run: | + where stackql > stackql_path.txt + set /p STACKQL_PATH=< stackql_path.txt mkdir "C:\Temp" 2>nul || echo "Temp directory already exists" - copy stackql.exe "C:\Temp\stackql.exe" + copy "%STACKQL_PATH%" "C:\Temp\stackql.exe" echo StackQL binary moved to C:\Temp\stackql.exe - name: Run non-server tests From 6e3b851884359afb3974ad341b9846c70f9b4773 Mon Sep 17 00:00:00 2001 From: Jeffrey Aven Date: Thu, 5 Jun 2025 11:25:44 +1000 Subject: [PATCH 08/31] ci update --- .github/workflows/test.yaml | 17 +++++++++++++---- start-stackql-server.sh | 1 - 2 files changed, 13 insertions(+), 5 deletions(-) diff --git a/.github/workflows/test.yaml b/.github/workflows/test.yaml index 4637b6d..25214d6 100644 --- a/.github/workflows/test.yaml +++ b/.github/workflows/test.yaml @@ -58,7 +58,7 @@ jobs: run: | stackql --version - - name: Move stackql binary to temp dir (Unix/macOS) + - name: Move stackql binary to temp dir (Linux/macOS) if: matrix.os != 'windows-latest' shell: bash run: | @@ -84,10 +84,19 @@ jobs: python3 run_tests.py shell: bash - - name: Start StackQL server + - name: Start StackQL server (Linux/macOS) + if: matrix.os != 'windows-latest' + shell: bash run: | - sh start-stackql-server.sh - shell: bash + nohup /tmp/stackql -v --pgsrv.port=5444 srv & + sleep 5 + + - name: Start StackQL server (Windows) + if: matrix.os == 'windows-latest' + shell: cmd + run: | + start /b C:\Temp\stackql.exe -v --pgsrv.port=5444 srv > stackql-server.log 2>&1 + timeout /t 5 /nobreak - name: Run server tests env: diff --git a/start-stackql-server.sh b/start-stackql-server.sh index bea44a9..f04e668 100644 --- a/start-stackql-server.sh +++ b/start-stackql-server.sh @@ -1,7 +1,6 @@ # start server if not running echo "checking if server is running" if [ -z "$(ps | grep stackql)" ]; then - echo "starting server with registry: $REG" nohup ./stackql -v --pgsrv.port=5444 srv & sleep 5 else From 1807554f0c527b3510ecd3ed94a67078fc934d56 Mon Sep 17 00:00:00 2001 From: Jeffrey Aven Date: Thu, 5 Jun 2025 11:33:53 +1000 Subject: [PATCH 09/31] ci update --- .github/workflows/test.yaml | 25 +++++++++++++++++++++---- 1 file changed, 21 insertions(+), 4 deletions(-) diff --git a/.github/workflows/test.yaml b/.github/workflows/test.yaml index 25214d6..6a8813c 100644 --- a/.github/workflows/test.yaml +++ b/.github/workflows/test.yaml @@ -25,7 +25,7 @@ jobs: - os: macos-latest python-version: "3.13" runs-on: ${{matrix.os}} - name: 'Run Tests on ${{matrix.os}} with Python ${{matrix.python-version}}' + name: '${{matrix.os}} Python ${{matrix.python-version}}' steps: - name: Checkout @@ -105,7 +105,24 @@ jobs: python3 run_server_tests.py shell: bash - - name: Stop StackQL server + - name: Stop StackQL server (Linux/macOS) + if: matrix.os != 'windows-latest' + shell: bash + run: | + echo "Stopping StackQL server on Unix/macOS..." + PID=$(pgrep -f "/tmp/stackql.*srv" || pgrep -f "stackql.*srv" || echo "") + if [ -z "$PID" ]; then + echo "No stackql server process found." + else + echo "stopping stackql server (PID: $PID)..." + kill -9 $PID + echo "stackql server stopped." + fi + + - name: Stop StackQL server (Windows) + if: matrix.os == 'windows-latest' + shell: cmd run: | - sh stop-stackql-server.sh - shell: bash \ No newline at end of file + echo "Stopping StackQL server on Windows..." + taskkill /F /IM stackql.exe 2>nul || echo "No stackql.exe process found" + echo "StackQL server stopped (Windows)" From 7a2aa3b2de7f7457f648abe80cd7997c88692a92 Mon Sep 17 00:00:00 2001 From: Jeffrey Aven Date: Thu, 5 Jun 2025 11:40:32 +1000 Subject: [PATCH 10/31] ci update --- .github/workflows/test.yaml | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/.github/workflows/test.yaml b/.github/workflows/test.yaml index 6a8813c..50c3fc5 100644 --- a/.github/workflows/test.yaml +++ b/.github/workflows/test.yaml @@ -93,10 +93,12 @@ jobs: - name: Start StackQL server (Windows) if: matrix.os == 'windows-latest' - shell: cmd + shell: cmd run: | - start /b C:\Temp\stackql.exe -v --pgsrv.port=5444 srv > stackql-server.log 2>&1 - timeout /t 5 /nobreak + echo @echo off > start_server.bat + echo C:\Temp\stackql.exe -v --pgsrv.port=5444 srv ^> stackql-server.log 2^>^&1 >> start_server.bat + start /b "" cmd /c start_server.bat + timeout /t 5 /nobreak > NUL - name: Run server tests env: From 1a71b46bd6890ad6604cb7c5710c329c5606d466 Mon Sep 17 00:00:00 2001 From: Jeffrey Aven Date: Thu, 5 Jun 2025 11:43:34 +1000 Subject: [PATCH 11/31] ci update --- .github/workflows/test.yaml | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/.github/workflows/test.yaml b/.github/workflows/test.yaml index 50c3fc5..d32d3a8 100644 --- a/.github/workflows/test.yaml +++ b/.github/workflows/test.yaml @@ -93,12 +93,11 @@ jobs: - name: Start StackQL server (Windows) if: matrix.os == 'windows-latest' - shell: cmd + shell: pwsh run: | - echo @echo off > start_server.bat - echo C:\Temp\stackql.exe -v --pgsrv.port=5444 srv ^> stackql-server.log 2^>^&1 >> start_server.bat - start /b "" cmd /c start_server.bat - timeout /t 5 /nobreak > NUL + $process = Start-Process -FilePath "C:\Temp\stackql.exe" -ArgumentList "-v", "--pgsrv.port=5444", "srv" -RedirectStandardOutput "stackql-server.log" -RedirectStandardError "stackql-server-error.log" -NoNewWindow -PassThru + Write-Host "StackQL server started with PID: $($process.Id)" + Start-Sleep -Seconds 5 - name: Run server tests env: From eb25d2d39496cd4e8e486b86db8cb2bc36764fdb Mon Sep 17 00:00:00 2001 From: Jeffrey Aven Date: Thu, 5 Jun 2025 11:47:52 +1000 Subject: [PATCH 12/31] ci update --- .github/workflows/test.yaml | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/.github/workflows/test.yaml b/.github/workflows/test.yaml index d32d3a8..45b81b6 100644 --- a/.github/workflows/test.yaml +++ b/.github/workflows/test.yaml @@ -9,9 +9,9 @@ jobs: strategy: matrix: os: - - ubuntu-latest + # - ubuntu-latest - windows-latest - - macos-latest + # - macos-latest python-version: - "3.8" - "3.9" @@ -93,7 +93,7 @@ jobs: - name: Start StackQL server (Windows) if: matrix.os == 'windows-latest' - shell: pwsh + shell: powershell run: | $process = Start-Process -FilePath "C:\Temp\stackql.exe" -ArgumentList "-v", "--pgsrv.port=5444", "srv" -RedirectStandardOutput "stackql-server.log" -RedirectStandardError "stackql-server-error.log" -NoNewWindow -PassThru Write-Host "StackQL server started with PID: $($process.Id)" From d0e2c8c1182069643ac36e829b50566f6da70823 Mon Sep 17 00:00:00 2001 From: Jeffrey Aven Date: Thu, 5 Jun 2025 11:52:32 +1000 Subject: [PATCH 13/31] ci update --- .github/workflows/test.yaml | 60 ++++++++++++++++++------------------- 1 file changed, 30 insertions(+), 30 deletions(-) diff --git a/.github/workflows/test.yaml b/.github/workflows/test.yaml index 45b81b6..05703c8 100644 --- a/.github/workflows/test.yaml +++ b/.github/workflows/test.yaml @@ -31,23 +31,23 @@ jobs: - name: Checkout uses: actions/checkout@v4.2.2 - - name: Set up Python ${{ matrix.python-version }} - uses: actions/setup-python@v5.6.0 - with: - python-version: ${{ matrix.python-version }} + # - name: Set up Python ${{ matrix.python-version }} + # uses: actions/setup-python@v5.6.0 + # with: + # python-version: ${{ matrix.python-version }} - - name: Upgrade pip - shell: bash - run: | - python3 -m pip install --upgrade pip + # - name: Upgrade pip + # shell: bash + # run: | + # python3 -m pip install --upgrade pip - - name: Install pystackql with all dependencies - run: | - pip install -e . + # - name: Install pystackql with all dependencies + # run: | + # pip install -e . - - name: Install test dependencies - run: | - pip install pytest>=6.2.5 pytest-cov>=2.12.0 nose>=1.3.7 + # - name: Install test dependencies + # run: | + # pip install pytest>=6.2.5 pytest-cov>=2.12.0 nose>=1.3.7 - name: setup-stackql uses: stackql/setup-stackql@v2.2.3 @@ -77,12 +77,12 @@ jobs: copy "%STACKQL_PATH%" "C:\Temp\stackql.exe" echo StackQL binary moved to C:\Temp\stackql.exe - - name: Run non-server tests - env: - GITHUB_ACTIONS: 'true' - run: | - python3 run_tests.py - shell: bash + # - name: Run non-server tests + # env: + # GITHUB_ACTIONS: 'true' + # run: | + # python3 run_tests.py + # shell: bash - name: Start StackQL server (Linux/macOS) if: matrix.os != 'windows-latest' @@ -93,18 +93,18 @@ jobs: - name: Start StackQL server (Windows) if: matrix.os == 'windows-latest' - shell: powershell + shell: cmd run: | - $process = Start-Process -FilePath "C:\Temp\stackql.exe" -ArgumentList "-v", "--pgsrv.port=5444", "srv" -RedirectStandardOutput "stackql-server.log" -RedirectStandardError "stackql-server-error.log" -NoNewWindow -PassThru - Write-Host "StackQL server started with PID: $($process.Id)" - Start-Sleep -Seconds 5 + start /b cmd /c "C:\Temp\stackql.exe -v --pgsrv.port=5444 srv > stackql-server.log 2>&1" + echo StackQL server started from C:\Temp\stackql.exe + timeout /t 5 /nobreak > NUL - - name: Run server tests - env: - GITHUB_ACTIONS: 'true' - run: | - python3 run_server_tests.py - shell: bash + # - name: Run server tests + # env: + # GITHUB_ACTIONS: 'true' + # run: | + # python3 run_server_tests.py + # shell: bash - name: Stop StackQL server (Linux/macOS) if: matrix.os != 'windows-latest' From ae8eafc6aef092b49f1435316af34e2c1b248f14 Mon Sep 17 00:00:00 2001 From: Jeffrey Aven Date: Thu, 5 Jun 2025 11:56:08 +1000 Subject: [PATCH 14/31] ci update --- .github/workflows/test.yaml | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/.github/workflows/test.yaml b/.github/workflows/test.yaml index 05703c8..5049710 100644 --- a/.github/workflows/test.yaml +++ b/.github/workflows/test.yaml @@ -95,10 +95,11 @@ jobs: if: matrix.os == 'windows-latest' shell: cmd run: | - start /b cmd /c "C:\Temp\stackql.exe -v --pgsrv.port=5444 srv > stackql-server.log 2>&1" + C:\Temp\stackql.exe -v --pgsrv.port=5444 srv > stackql-server.log 2>&1 & echo StackQL server started from C:\Temp\stackql.exe timeout /t 5 /nobreak > NUL + # - name: Run server tests # env: # GITHUB_ACTIONS: 'true' From e103b7e1cba1fc21263aafce2a0f7f59094ae428 Mon Sep 17 00:00:00 2001 From: Jeffrey Aven Date: Thu, 5 Jun 2025 11:58:10 +1000 Subject: [PATCH 15/31] ci update --- .github/workflows/test.yaml | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/.github/workflows/test.yaml b/.github/workflows/test.yaml index 5049710..ba11a88 100644 --- a/.github/workflows/test.yaml +++ b/.github/workflows/test.yaml @@ -93,12 +93,15 @@ jobs: - name: Start StackQL server (Windows) if: matrix.os == 'windows-latest' - shell: cmd + shell: pwsh run: | - C:\Temp\stackql.exe -v --pgsrv.port=5444 srv > stackql-server.log 2>&1 & - echo StackQL server started from C:\Temp\stackql.exe - timeout /t 5 /nobreak > NUL - + Start-Process -FilePath "C:\Temp\stackql.exe" ` + -ArgumentList "-v", "--pgsrv.port=5444", "srv" ` + -RedirectStandardOutput "stackql-server.log" ` + -RedirectStandardError "stackql-server.log" ` + -NoNewWindow + Write-Host "StackQL server started from C:\Temp\stackql.exe" + Start-Sleep -Seconds 5 # - name: Run server tests # env: From 01b056cedc7d3a23e9b52a4037252aba8e3d6539 Mon Sep 17 00:00:00 2001 From: Jeffrey Aven Date: Thu, 5 Jun 2025 12:01:53 +1000 Subject: [PATCH 16/31] ci update --- .github/workflows/test.yaml | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/.github/workflows/test.yaml b/.github/workflows/test.yaml index ba11a88..54b8dc8 100644 --- a/.github/workflows/test.yaml +++ b/.github/workflows/test.yaml @@ -95,10 +95,12 @@ jobs: if: matrix.os == 'windows-latest' shell: pwsh run: | + $outFile = "stackql-server.out.log" + $errFile = "stackql-server.err.log" Start-Process -FilePath "C:\Temp\stackql.exe" ` -ArgumentList "-v", "--pgsrv.port=5444", "srv" ` - -RedirectStandardOutput "stackql-server.log" ` - -RedirectStandardError "stackql-server.log" ` + -RedirectStandardOutput $outFile ` + -RedirectStandardError $errFile ` -NoNewWindow Write-Host "StackQL server started from C:\Temp\stackql.exe" Start-Sleep -Seconds 5 From b17ae98e319f0492274dc792609409da5934d664 Mon Sep 17 00:00:00 2001 From: Jeffrey Aven Date: Thu, 5 Jun 2025 12:04:34 +1000 Subject: [PATCH 17/31] ci update --- .github/workflows/test.yaml | 9 +-------- 1 file changed, 1 insertion(+), 8 deletions(-) diff --git a/.github/workflows/test.yaml b/.github/workflows/test.yaml index 54b8dc8..6494bae 100644 --- a/.github/workflows/test.yaml +++ b/.github/workflows/test.yaml @@ -95,14 +95,7 @@ jobs: if: matrix.os == 'windows-latest' shell: pwsh run: | - $outFile = "stackql-server.out.log" - $errFile = "stackql-server.err.log" - Start-Process -FilePath "C:\Temp\stackql.exe" ` - -ArgumentList "-v", "--pgsrv.port=5444", "srv" ` - -RedirectStandardOutput $outFile ` - -RedirectStandardError $errFile ` - -NoNewWindow - Write-Host "StackQL server started from C:\Temp\stackql.exe" + Start-Process -FilePath "C:\Temp\stackql.exe" -ArgumentList "-v", "--pgsrv.port=5444", "srv" Start-Sleep -Seconds 5 # - name: Run server tests From 5dbd355f92c18f67526145ce5a7c74192210f641 Mon Sep 17 00:00:00 2001 From: Jeffrey Aven Date: Thu, 5 Jun 2025 12:11:12 +1000 Subject: [PATCH 18/31] ci update --- .github/workflows/test.yaml | 25 +++++++++++++++++++------ 1 file changed, 19 insertions(+), 6 deletions(-) diff --git a/.github/workflows/test.yaml b/.github/workflows/test.yaml index 6494bae..ee1e22e 100644 --- a/.github/workflows/test.yaml +++ b/.github/workflows/test.yaml @@ -67,15 +67,27 @@ jobs: cp "$STACKQL_PATH" /tmp/stackql echo "StackQL binary moved from ${STACKQL_PATH} to /tmp/stackql" + # - name: Move stackql binary to temp dir (Windows) + # if: matrix.os == 'windows-latest' + # shell: cmd + # run: | + # where stackql > stackql_path.txt + # set /p STACKQL_PATH=< stackql_path.txt + # mkdir "C:\Temp" 2>nul || echo "Temp directory already exists" + # copy "%STACKQL_PATH%" "C:\Temp\stackql.exe" + # echo StackQL binary moved to C:\Temp\stackql.exe + - name: Move stackql binary to temp dir (Windows) if: matrix.os == 'windows-latest' - shell: cmd + shell: pwsh run: | - where stackql > stackql_path.txt - set /p STACKQL_PATH=< stackql_path.txt - mkdir "C:\Temp" 2>nul || echo "Temp directory already exists" - copy "%STACKQL_PATH%" "C:\Temp\stackql.exe" - echo StackQL binary moved to C:\Temp\stackql.exe + $stackqlBin = "$env:USERPROFILE\AppData\Local\stackql\bin\stackql.exe" + if (-Not (Test-Path $stackqlBin)) { + throw "Expected stackql binary not found at $stackqlBin" + } + Copy-Item $stackqlBin -Destination "C:\Temp\stackql.exe" -Force + Write-Host "StackQL binary moved to C:\Temp\stackql.exe" + # - name: Run non-server tests # env: @@ -98,6 +110,7 @@ jobs: Start-Process -FilePath "C:\Temp\stackql.exe" -ArgumentList "-v", "--pgsrv.port=5444", "srv" Start-Sleep -Seconds 5 + # - name: Run server tests # env: # GITHUB_ACTIONS: 'true' From f70b6cd6c6eedb153e0c7c814fc5a730d3ca5400 Mon Sep 17 00:00:00 2001 From: Jeffrey Aven Date: Thu, 5 Jun 2025 12:19:19 +1000 Subject: [PATCH 19/31] ci update --- .github/workflows/test.yaml | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/.github/workflows/test.yaml b/.github/workflows/test.yaml index ee1e22e..4963ec7 100644 --- a/.github/workflows/test.yaml +++ b/.github/workflows/test.yaml @@ -54,10 +54,19 @@ jobs: with: use_wrapper: true - - name: Show stackql version + - name: Show stackql version (Linux/macOS) + if: matrix.os != 'windows-latest' + shell: bash run: | stackql --version + - name: Show stackql version (Windows) + if: matrix.os == 'windows-latest' + shell: cmd + run: | + where stackql.exe + stackql.exe --version + - name: Move stackql binary to temp dir (Linux/macOS) if: matrix.os != 'windows-latest' shell: bash From 6ec525f240670fa63f9e39a9a1315fafe057ec84 Mon Sep 17 00:00:00 2001 From: Jeffrey Aven Date: Thu, 5 Jun 2025 12:20:27 +1000 Subject: [PATCH 20/31] ci update --- .github/workflows/test.yaml | 1 - 1 file changed, 1 deletion(-) diff --git a/.github/workflows/test.yaml b/.github/workflows/test.yaml index 4963ec7..a8b2f9a 100644 --- a/.github/workflows/test.yaml +++ b/.github/workflows/test.yaml @@ -64,7 +64,6 @@ jobs: if: matrix.os == 'windows-latest' shell: cmd run: | - where stackql.exe stackql.exe --version - name: Move stackql binary to temp dir (Linux/macOS) From 4e2768b99c310968174cec84e525b20925469ddd Mon Sep 17 00:00:00 2001 From: Jeffrey Aven Date: Thu, 5 Jun 2025 12:24:12 +1000 Subject: [PATCH 21/31] ci update --- .github/workflows/test.yaml | 27 ++++++++------------------- 1 file changed, 8 insertions(+), 19 deletions(-) diff --git a/.github/workflows/test.yaml b/.github/workflows/test.yaml index a8b2f9a..81da646 100644 --- a/.github/workflows/test.yaml +++ b/.github/workflows/test.yaml @@ -64,7 +64,7 @@ jobs: if: matrix.os == 'windows-latest' shell: cmd run: | - stackql.exe --version + stackql-bin.exe --version - name: Move stackql binary to temp dir (Linux/macOS) if: matrix.os != 'windows-latest' @@ -75,27 +75,16 @@ jobs: cp "$STACKQL_PATH" /tmp/stackql echo "StackQL binary moved from ${STACKQL_PATH} to /tmp/stackql" - # - name: Move stackql binary to temp dir (Windows) - # if: matrix.os == 'windows-latest' - # shell: cmd - # run: | - # where stackql > stackql_path.txt - # set /p STACKQL_PATH=< stackql_path.txt - # mkdir "C:\Temp" 2>nul || echo "Temp directory already exists" - # copy "%STACKQL_PATH%" "C:\Temp\stackql.exe" - # echo StackQL binary moved to C:\Temp\stackql.exe - - name: Move stackql binary to temp dir (Windows) if: matrix.os == 'windows-latest' shell: pwsh run: | - $stackqlBin = "$env:USERPROFILE\AppData\Local\stackql\bin\stackql.exe" - if (-Not (Test-Path $stackqlBin)) { - throw "Expected stackql binary not found at $stackqlBin" + $bin = Join-Path $Env:STACKQL_CLI_PATH 'stackql-bin.exe' + if (-Not (Test-Path $bin)) { + throw "Binary not found at $bin" } - Copy-Item $stackqlBin -Destination "C:\Temp\stackql.exe" -Force - Write-Host "StackQL binary moved to C:\Temp\stackql.exe" - + Copy-Item $bin -Destination "C:\Temp\stackql.exe" -Force + Write-Host "Moved real StackQL binary to C:\Temp\stackql.exe" # - name: Run non-server tests # env: @@ -115,10 +104,10 @@ jobs: if: matrix.os == 'windows-latest' shell: pwsh run: | - Start-Process -FilePath "C:\Temp\stackql.exe" -ArgumentList "-v", "--pgsrv.port=5444", "srv" + Start-Process -FilePath "C:\Temp\stackql.exe" ` + -ArgumentList "-v", "--pgsrv.port=5444", "srv" Start-Sleep -Seconds 5 - # - name: Run server tests # env: # GITHUB_ACTIONS: 'true' From 66731b593c00fccdc8440a88991b9f8d943573ce Mon Sep 17 00:00:00 2001 From: Jeffrey Aven Date: Thu, 5 Jun 2025 12:26:30 +1000 Subject: [PATCH 22/31] ci update --- .github/workflows/test.yaml | 50 ++++++++++++++++++------------------- 1 file changed, 24 insertions(+), 26 deletions(-) diff --git a/.github/workflows/test.yaml b/.github/workflows/test.yaml index 81da646..74d4b86 100644 --- a/.github/workflows/test.yaml +++ b/.github/workflows/test.yaml @@ -31,23 +31,23 @@ jobs: - name: Checkout uses: actions/checkout@v4.2.2 - # - name: Set up Python ${{ matrix.python-version }} - # uses: actions/setup-python@v5.6.0 - # with: - # python-version: ${{ matrix.python-version }} + - name: Set up Python ${{ matrix.python-version }} + uses: actions/setup-python@v5.6.0 + with: + python-version: ${{ matrix.python-version }} - # - name: Upgrade pip - # shell: bash - # run: | - # python3 -m pip install --upgrade pip + - name: Upgrade pip + shell: bash + run: | + python3 -m pip install --upgrade pip - # - name: Install pystackql with all dependencies - # run: | - # pip install -e . + - name: Install pystackql with all dependencies + run: | + pip install -e . - # - name: Install test dependencies - # run: | - # pip install pytest>=6.2.5 pytest-cov>=2.12.0 nose>=1.3.7 + - name: Install test dependencies + run: | + pip install pytest>=6.2.5 pytest-cov>=2.12.0 nose>=1.3.7 - name: setup-stackql uses: stackql/setup-stackql@v2.2.3 @@ -86,12 +86,11 @@ jobs: Copy-Item $bin -Destination "C:\Temp\stackql.exe" -Force Write-Host "Moved real StackQL binary to C:\Temp\stackql.exe" - # - name: Run non-server tests - # env: - # GITHUB_ACTIONS: 'true' - # run: | - # python3 run_tests.py - # shell: bash + - name: Run non-server tests + env: + GITHUB_ACTIONS: 'true' + run: | + python3 run_tests.py - name: Start StackQL server (Linux/macOS) if: matrix.os != 'windows-latest' @@ -108,12 +107,11 @@ jobs: -ArgumentList "-v", "--pgsrv.port=5444", "srv" Start-Sleep -Seconds 5 - # - name: Run server tests - # env: - # GITHUB_ACTIONS: 'true' - # run: | - # python3 run_server_tests.py - # shell: bash + - name: Run server tests + env: + GITHUB_ACTIONS: 'true' + run: | + python3 run_server_tests.py - name: Stop StackQL server (Linux/macOS) if: matrix.os != 'windows-latest' From 08b74aba1109b75450bcd3cae46c15a01263d60a Mon Sep 17 00:00:00 2001 From: Jeffrey Aven Date: Thu, 5 Jun 2025 14:43:11 +1000 Subject: [PATCH 23/31] ci update --- tests/conftest.py | 87 +++++++++++++---------------------------------- 1 file changed, 23 insertions(+), 64 deletions(-) diff --git a/tests/conftest.py b/tests/conftest.py index b7d285e..ff9280b 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -6,6 +6,7 @@ import os import sys +import platform import time import pytest import subprocess @@ -47,74 +48,32 @@ def setup_stackql(): # Return the StackQL instance for use in tests return stackql +def stackql_process_running(): + try: + if platform.system() == "Windows": + # Use `tasklist` to look for stackql.exe with correct port in args (may not include args, so fallback is loose match) + output = subprocess.check_output(['tasklist', '/FI', 'IMAGENAME eq stackql.exe'], text=True) + return "stackql.exe" in output + else: + # Use `ps aux` to search for 'stackql' process with the correct port + output = subprocess.check_output(['ps', 'aux'], text=True) + return f"--pgsrv.port={SERVER_PORT}" in output and "stackql" in output + except subprocess.CalledProcessError: + return False + @pytest.fixture(scope="session") -def stackql_server(setup_stackql): +def stackql_server(): """ - Session-level fixture to start and stop a StackQL server. - This runs once for all tests that request it. - - This improved version: - 1. Checks if a server is already running before starting one - 2. Uses process groups for better cleanup - 3. Handles errors more gracefully + Verifies that a StackQL server process is running with the expected port. + Does not attempt to start or stop the process. """ - global server_process - - # Check if server is already running - print("\nChecking if server is running...") - ps_output = subprocess.run( - ["ps", "aux"], - capture_output=True, - text=True - ).stdout - - if "stackql" in ps_output and f"--pgsrv.port={SERVER_PORT}" in ps_output: - print("Server is already running") - # No need to start a server or set server_process - else: - # Start the server - print(f"Starting stackql server on {SERVER_ADDRESS}:{SERVER_PORT}...") - - # Get the registry setting from environment variable if available - registry = os.environ.get('REG', '') - registry_arg = f"--registry {registry}" if registry else "" - - # Build the command - cmd = f"{setup_stackql.bin_path} srv --pgsrv.address {SERVER_ADDRESS} --pgsrv.port {SERVER_PORT} {registry_arg}" - - # Start the server process with process group for better cleanup - server_process = subprocess.Popen( - cmd, - shell=True, - stdout=subprocess.PIPE, - stderr=subprocess.PIPE, - preexec_fn=os.setsid # Use process group for cleaner termination - ) - - # Wait for server to start - print("Waiting for server to initialize...") - time.sleep(5) - - # Check if server started successfully - if server_process.poll() is not None: - # Process has terminated - stdout, stderr = server_process.communicate() - pytest.fail(f"Server failed to start: {stderr.decode()}") - - # Yield to run tests + print(f"\nšŸ” Checking for running StackQL server process (port {SERVER_PORT})...") + + if not stackql_process_running(): + pytest.exit(f"āŒ No running StackQL server process found for port {SERVER_PORT}", returncode=1) + + print("āœ… StackQL server process is running.") yield - - # Clean up server at the end if we started it - if server_process and server_process.poll() is None: - print("\nShutting down stackql server...") - try: - # Kill the entire process group - os.killpg(os.getpgid(server_process.pid), signal.SIGTERM) - server_process.wait(timeout=5) - print("Server shutdown complete") - except subprocess.TimeoutExpired: - print("Server did not terminate gracefully, forcing shutdown...") - os.killpg(os.getpgid(server_process.pid), signal.SIGKILL) @pytest.fixture def mock_interactive_shell(): From 6870874335e4275dcca68c2d92f1cc3ac720ecf5 Mon Sep 17 00:00:00 2001 From: Jeffrey Aven Date: Thu, 5 Jun 2025 14:49:27 +1000 Subject: [PATCH 24/31] ci update --- .github/workflows/test.yaml | 13 +++++-------- setup.py | 1 - 2 files changed, 5 insertions(+), 9 deletions(-) diff --git a/.github/workflows/test.yaml b/.github/workflows/test.yaml index 74d4b86..ad5c61f 100644 --- a/.github/workflows/test.yaml +++ b/.github/workflows/test.yaml @@ -10,20 +10,17 @@ jobs: matrix: os: # - ubuntu-latest - - windows-latest - # - macos-latest + # - windows-latest + - macos-latest python-version: - - "3.8" - "3.9" - "3.10" - "3.11" - "3.12" - "3.13" - exclude: - - os: macos-latest - python-version: "3.8" - - os: macos-latest - python-version: "3.13" + # exclude: + # - os: macos-latest + # python-version: "3.13" runs-on: ${{matrix.os}} name: '${{matrix.os}} Python ${{matrix.python-version}}' diff --git a/setup.py b/setup.py index cadda4b..f69d19d 100644 --- a/setup.py +++ b/setup.py @@ -37,7 +37,6 @@ 'Operating System :: Microsoft :: Windows', 'Operating System :: MacOS', 'Operating System :: POSIX :: Linux', - 'Programming Language :: Python :: 3.8', 'Programming Language :: Python :: 3.9', 'Programming Language :: Python :: 3.10', 'Programming Language :: Python :: 3.11', From 8b730f95222345e5d2b5a25e47bc2ead81e5f82d Mon Sep 17 00:00:00 2001 From: Jeffrey Aven Date: Thu, 5 Jun 2025 14:52:47 +1000 Subject: [PATCH 25/31] ci update --- tests/test_constants.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_constants.py b/tests/test_constants.py index 86caf1c..60a734c 100644 --- a/tests/test_constants.py +++ b/tests/test_constants.py @@ -14,7 +14,7 @@ import pandas as pd # Server connection settings -SERVER_PORT = 5466 +SERVER_PORT = 5444 SERVER_ADDRESS = "127.0.0.1" # Expected properties and patterns for validation From 81d8e0ee3e4b64166c31cd36d7356a68e2e26d11 Mon Sep 17 00:00:00 2001 From: Jeffrey Aven Date: Thu, 5 Jun 2025 15:00:38 +1000 Subject: [PATCH 26/31] ci update --- .github/workflows/test.yaml | 11 ++++------- pystackql/core/stackql.py | 4 ++-- start-stackql-server.sh | 2 +- tests/test_constants.py | 2 +- 4 files changed, 8 insertions(+), 11 deletions(-) diff --git a/.github/workflows/test.yaml b/.github/workflows/test.yaml index ad5c61f..42b698a 100644 --- a/.github/workflows/test.yaml +++ b/.github/workflows/test.yaml @@ -9,8 +9,8 @@ jobs: strategy: matrix: os: - # - ubuntu-latest - # - windows-latest + - ubuntu-latest + - windows-latest - macos-latest python-version: - "3.9" @@ -18,9 +18,6 @@ jobs: - "3.11" - "3.12" - "3.13" - # exclude: - # - os: macos-latest - # python-version: "3.13" runs-on: ${{matrix.os}} name: '${{matrix.os}} Python ${{matrix.python-version}}' @@ -93,7 +90,7 @@ jobs: if: matrix.os != 'windows-latest' shell: bash run: | - nohup /tmp/stackql -v --pgsrv.port=5444 srv & + nohup /tmp/stackql -v --pgsrv.port=5466 srv & sleep 5 - name: Start StackQL server (Windows) @@ -101,7 +98,7 @@ jobs: shell: pwsh run: | Start-Process -FilePath "C:\Temp\stackql.exe" ` - -ArgumentList "-v", "--pgsrv.port=5444", "srv" + -ArgumentList "-v", "--pgsrv.port=5466", "srv" Start-Sleep -Seconds 5 - name: Run server tests diff --git a/pystackql/core/stackql.py b/pystackql/core/stackql.py index 2a681d5..25e8149 100644 --- a/pystackql/core/stackql.py +++ b/pystackql/core/stackql.py @@ -24,7 +24,7 @@ class StackQL: (`server_mode` only, defaults to `'127.0.0.1'`) :type server_address: str, optional :param server_port: The port of the StackQL server - (`server_mode` only, defaults to `5444`) + (`server_mode` only, defaults to `5466`) :type server_port: int, optional :param backend_storage_mode: Specifies backend storage mode, options are 'memory' and 'file' (defaults to `'memory'`, this option is ignored in `server_mode`) @@ -119,7 +119,7 @@ class StackQL: def __init__(self, server_mode=False, server_address='127.0.0.1', - server_port=5444, + server_port=5466, output='dict', sep=',', header=False, diff --git a/start-stackql-server.sh b/start-stackql-server.sh index f04e668..82f8d96 100644 --- a/start-stackql-server.sh +++ b/start-stackql-server.sh @@ -1,7 +1,7 @@ # start server if not running echo "checking if server is running" if [ -z "$(ps | grep stackql)" ]; then - nohup ./stackql -v --pgsrv.port=5444 srv & + nohup ./stackql -v --pgsrv.port=5466 srv & sleep 5 else echo "server is already running" diff --git a/tests/test_constants.py b/tests/test_constants.py index 60a734c..86caf1c 100644 --- a/tests/test_constants.py +++ b/tests/test_constants.py @@ -14,7 +14,7 @@ import pandas as pd # Server connection settings -SERVER_PORT = 5444 +SERVER_PORT = 5466 SERVER_ADDRESS = "127.0.0.1" # Expected properties and patterns for validation From e5265085115d652ba129be742399bf903aaab0d7 Mon Sep 17 00:00:00 2001 From: Jeffrey Aven Date: Thu, 5 Jun 2025 15:07:44 +1000 Subject: [PATCH 27/31] ci update --- .github/workflows/test.yaml | 6 +++--- start-stackql-server.sh | 2 +- tests/README.md | 2 +- tests/test_constants.py | 2 +- 4 files changed, 6 insertions(+), 6 deletions(-) diff --git a/.github/workflows/test.yaml b/.github/workflows/test.yaml index 42b698a..d041ae0 100644 --- a/.github/workflows/test.yaml +++ b/.github/workflows/test.yaml @@ -85,12 +85,12 @@ jobs: GITHUB_ACTIONS: 'true' run: | python3 run_tests.py - + - name: Start StackQL server (Linux/macOS) if: matrix.os != 'windows-latest' shell: bash run: | - nohup /tmp/stackql -v --pgsrv.port=5466 srv & + nohup /tmp/stackql -v --pgsrv.port=5444 srv & sleep 5 - name: Start StackQL server (Windows) @@ -98,7 +98,7 @@ jobs: shell: pwsh run: | Start-Process -FilePath "C:\Temp\stackql.exe" ` - -ArgumentList "-v", "--pgsrv.port=5466", "srv" + -ArgumentList "-v", "--pgsrv.port=5444", "srv" Start-Sleep -Seconds 5 - name: Run server tests diff --git a/start-stackql-server.sh b/start-stackql-server.sh index 82f8d96..f04e668 100644 --- a/start-stackql-server.sh +++ b/start-stackql-server.sh @@ -1,7 +1,7 @@ # start server if not running echo "checking if server is running" if [ -z "$(ps | grep stackql)" ]; then - nohup ./stackql -v --pgsrv.port=5466 srv & + nohup ./stackql -v --pgsrv.port=5444 srv & sleep 5 else echo "server is already running" diff --git a/tests/README.md b/tests/README.md index 6491481..e7cf697 100644 --- a/tests/README.md +++ b/tests/README.md @@ -52,7 +52,7 @@ Server tests are skipped by default because they require a running StackQL serve 1. Start a StackQL server: ```bash - stackql srv --pgsrv.address 127.0.0.1 --pgsrv.port 5466 + stackql srv --pgsrv.address 127.0.0.1 --pgsrv.port 5444 ``` 2. Run the server tests: diff --git a/tests/test_constants.py b/tests/test_constants.py index 86caf1c..60a734c 100644 --- a/tests/test_constants.py +++ b/tests/test_constants.py @@ -14,7 +14,7 @@ import pandas as pd # Server connection settings -SERVER_PORT = 5466 +SERVER_PORT = 5444 SERVER_ADDRESS = "127.0.0.1" # Expected properties and patterns for validation From bdd0da38d94258c78514aeab764eb0c783c1335d Mon Sep 17 00:00:00 2001 From: Jeffrey Aven Date: Thu, 5 Jun 2025 15:11:40 +1000 Subject: [PATCH 28/31] ci update --- .github/workflows/test.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/test.yaml b/.github/workflows/test.yaml index d041ae0..3744ce1 100644 --- a/.github/workflows/test.yaml +++ b/.github/workflows/test.yaml @@ -91,7 +91,7 @@ jobs: shell: bash run: | nohup /tmp/stackql -v --pgsrv.port=5444 srv & - sleep 5 + sleep 10 - name: Start StackQL server (Windows) if: matrix.os == 'windows-latest' From d1a4806939d43ed357a5134119c9f79cf8932d9b Mon Sep 17 00:00:00 2001 From: Jeffrey Aven Date: Thu, 5 Jun 2025 15:14:30 +1000 Subject: [PATCH 29/31] ci update --- .github/workflows/test.yaml | 13 +++++-------- 1 file changed, 5 insertions(+), 8 deletions(-) diff --git a/.github/workflows/test.yaml b/.github/workflows/test.yaml index 3744ce1..369f157 100644 --- a/.github/workflows/test.yaml +++ b/.github/workflows/test.yaml @@ -86,12 +86,15 @@ jobs: run: | python3 run_tests.py - - name: Start StackQL server (Linux/macOS) + - name: Start StackQL server and run tests (Linux/macOS) if: matrix.os != 'windows-latest' shell: bash + env: + GITHUB_ACTIONS: 'true' run: | nohup /tmp/stackql -v --pgsrv.port=5444 srv & - sleep 10 + sleep 5 + python3 run_server_tests.py - name: Start StackQL server (Windows) if: matrix.os == 'windows-latest' @@ -101,12 +104,6 @@ jobs: -ArgumentList "-v", "--pgsrv.port=5444", "srv" Start-Sleep -Seconds 5 - - name: Run server tests - env: - GITHUB_ACTIONS: 'true' - run: | - python3 run_server_tests.py - - name: Stop StackQL server (Linux/macOS) if: matrix.os != 'windows-latest' shell: bash From fde69d8bac81e7277eae7044722eab0e681265b6 Mon Sep 17 00:00:00 2001 From: Jeffrey Aven Date: Thu, 5 Jun 2025 15:23:04 +1000 Subject: [PATCH 30/31] ci update --- tests/conftest.py | 28 +++++++++++++++++++++++----- 1 file changed, 23 insertions(+), 5 deletions(-) diff --git a/tests/conftest.py b/tests/conftest.py index ff9280b..adc82bf 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -48,16 +48,34 @@ def setup_stackql(): # Return the StackQL instance for use in tests return stackql +# def stackql_process_running(): +# try: +# if platform.system() == "Windows": +# # Use `tasklist` to look for stackql.exe with correct port in args (may not include args, so fallback is loose match) +# output = subprocess.check_output(['tasklist', '/FI', 'IMAGENAME eq stackql.exe'], text=True) +# return "stackql.exe" in output +# else: +# # Use `ps aux` to search for 'stackql' process with the correct port +# output = subprocess.check_output(['ps', 'aux'], text=True) +# return f"--pgsrv.port={SERVER_PORT}" in output and "stackql" in output +# except subprocess.CalledProcessError: +# return False + def stackql_process_running(): try: if platform.system() == "Windows": - # Use `tasklist` to look for stackql.exe with correct port in args (may not include args, so fallback is loose match) - output = subprocess.check_output(['tasklist', '/FI', 'IMAGENAME eq stackql.exe'], text=True) + output = subprocess.check_output( + ['tasklist', '/FI', 'IMAGENAME eq stackql.exe'], text=True + ) return "stackql.exe" in output else: - # Use `ps aux` to search for 'stackql' process with the correct port - output = subprocess.check_output(['ps', 'aux'], text=True) - return f"--pgsrv.port={SERVER_PORT}" in output and "stackql" in output + # More reliable: use pgrep + full argument check + output = subprocess.check_output( + f"ps aux | grep '[s]tackql' | grep -- '--pgsrv.port={SERVER_PORT}'", + shell=True, + text=True + ) + return bool(output.strip()) except subprocess.CalledProcessError: return False From a069c53dac1dacb99e6440c7e20f2b048d9f1b6c Mon Sep 17 00:00:00 2001 From: Jeffrey Aven Date: Wed, 11 Jun 2025 14:39:52 +1000 Subject: [PATCH 31/31] updated test_connection method --- .github/workflows/test.yaml | 4 +- pystackql/core/stackql.py | 27 ++++++++- start-stackql-server.sh | 2 +- tests/README.md | 2 +- tests/conftest.py | 45 -------------- tests/test_constants.py | 2 +- tests/test_server.py | 118 ++++++++++++++++++++++++++++++++---- 7 files changed, 136 insertions(+), 64 deletions(-) diff --git a/.github/workflows/test.yaml b/.github/workflows/test.yaml index 369f157..d3623cf 100644 --- a/.github/workflows/test.yaml +++ b/.github/workflows/test.yaml @@ -92,7 +92,7 @@ jobs: env: GITHUB_ACTIONS: 'true' run: | - nohup /tmp/stackql -v --pgsrv.port=5444 srv & + nohup /tmp/stackql -v --pgsrv.port=5466 srv & sleep 5 python3 run_server_tests.py @@ -101,7 +101,7 @@ jobs: shell: pwsh run: | Start-Process -FilePath "C:\Temp\stackql.exe" ` - -ArgumentList "-v", "--pgsrv.port=5444", "srv" + -ArgumentList "-v", "--pgsrv.port=5466", "srv" Start-Sleep -Seconds 5 - name: Stop StackQL server (Linux/macOS) diff --git a/pystackql/core/stackql.py b/pystackql/core/stackql.py index 25e8149..457528d 100644 --- a/pystackql/core/stackql.py +++ b/pystackql/core/stackql.py @@ -432,4 +432,29 @@ async def executeQueriesAsync(self, queries): "or switch to local mode if you need to run multiple queries concurrently." ) - return await self.async_executor.execute_queries(queries) \ No newline at end of file + return await self.async_executor.execute_queries(queries) + + def test_connection(self): + """Tests if the server connection is working by executing a simple query. + + This method is only valid when server_mode=True. + + Returns: + bool: True if the connection is working, False otherwise. + + Raises: + ValueError: If called when not in server mode. + """ + if not self.server_mode: + raise ValueError("The test_connectivity method is only available in server mode.") + + try: + result = self.server_connection.execute_query("SELECT 'test' as test_value") + return (isinstance(result, list) and + len(result) == 1 and + 'test_value' in result[0] and + result[0]['test_value'] == 'test') + except Exception as e: + if self.debug: + print(f"Connection test failed: {str(e)}") + return False diff --git a/start-stackql-server.sh b/start-stackql-server.sh index f04e668..82f8d96 100644 --- a/start-stackql-server.sh +++ b/start-stackql-server.sh @@ -1,7 +1,7 @@ # start server if not running echo "checking if server is running" if [ -z "$(ps | grep stackql)" ]; then - nohup ./stackql -v --pgsrv.port=5444 srv & + nohup ./stackql -v --pgsrv.port=5466 srv & sleep 5 else echo "server is already running" diff --git a/tests/README.md b/tests/README.md index e7cf697..6491481 100644 --- a/tests/README.md +++ b/tests/README.md @@ -52,7 +52,7 @@ Server tests are skipped by default because they require a running StackQL serve 1. Start a StackQL server: ```bash - stackql srv --pgsrv.address 127.0.0.1 --pgsrv.port 5444 + stackql srv --pgsrv.address 127.0.0.1 --pgsrv.port 5466 ``` 2. Run the server tests: diff --git a/tests/conftest.py b/tests/conftest.py index adc82bf..ecc13cd 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -48,51 +48,6 @@ def setup_stackql(): # Return the StackQL instance for use in tests return stackql -# def stackql_process_running(): -# try: -# if platform.system() == "Windows": -# # Use `tasklist` to look for stackql.exe with correct port in args (may not include args, so fallback is loose match) -# output = subprocess.check_output(['tasklist', '/FI', 'IMAGENAME eq stackql.exe'], text=True) -# return "stackql.exe" in output -# else: -# # Use `ps aux` to search for 'stackql' process with the correct port -# output = subprocess.check_output(['ps', 'aux'], text=True) -# return f"--pgsrv.port={SERVER_PORT}" in output and "stackql" in output -# except subprocess.CalledProcessError: -# return False - -def stackql_process_running(): - try: - if platform.system() == "Windows": - output = subprocess.check_output( - ['tasklist', '/FI', 'IMAGENAME eq stackql.exe'], text=True - ) - return "stackql.exe" in output - else: - # More reliable: use pgrep + full argument check - output = subprocess.check_output( - f"ps aux | grep '[s]tackql' | grep -- '--pgsrv.port={SERVER_PORT}'", - shell=True, - text=True - ) - return bool(output.strip()) - except subprocess.CalledProcessError: - return False - -@pytest.fixture(scope="session") -def stackql_server(): - """ - Verifies that a StackQL server process is running with the expected port. - Does not attempt to start or stop the process. - """ - print(f"\nšŸ” Checking for running StackQL server process (port {SERVER_PORT})...") - - if not stackql_process_running(): - pytest.exit(f"āŒ No running StackQL server process found for port {SERVER_PORT}", returncode=1) - - print("āœ… StackQL server process is running.") - yield - @pytest.fixture def mock_interactive_shell(): """Create a mock IPython shell for testing.""" diff --git a/tests/test_constants.py b/tests/test_constants.py index 60a734c..86caf1c 100644 --- a/tests/test_constants.py +++ b/tests/test_constants.py @@ -14,7 +14,7 @@ import pandas as pd # Server connection settings -SERVER_PORT = 5444 +SERVER_PORT = 5466 SERVER_ADDRESS = "127.0.0.1" # Expected properties and patterns for validation diff --git a/tests/test_server.py b/tests/test_server.py index 312f0f1..80e91f6 100644 --- a/tests/test_server.py +++ b/tests/test_server.py @@ -7,6 +7,7 @@ """ import re +import os import pytest import pandas as pd from unittest.mock import patch @@ -20,29 +21,80 @@ pystackql_test_setup ) -@pytest.mark.usefixtures("stackql_server") +# @pytest.mark.usefixtures("stackql_server") class TestServerMode: """Tests for PyStackQL server mode functionality.""" StackQL = StackQL # For use with pystackql_test_setup decorator - + server_available = False # Class-level flag to track server availability + + # @pystackql_test_setup(server_mode=True) + # def test_server_mode_connectivity(self): + # """Test that server mode connects successfully.""" + # # Check server_mode flag is set correctly + # assert self.stackql.server_mode, "StackQL should be in server mode" + + # # Check server connection object exists + # assert hasattr(self.stackql, 'server_connection'), "StackQL should have a server_connection attribute" + # assert self.stackql.server_connection is not None, "Server connection object should not be None" + + # # IMPORTANT: Actually test the connection works + # connection_working = self.stackql.test_connection() + + # # Print detailed results for debugging + # if not connection_working: + # print("āš ļø Server connection test failed: unable to execute a simple query") + # print(f"Server address: {self.stackql.server_address}") + # print(f"Server port: {self.stackql.server_port}") + # print("\nāŒ SERVER CONNECTION FAILED - SKIPPING REMAINING SERVER TESTS") + # else: + # # Set flag indicating server is available + # TestServerMode.server_available = True + + # print_test_result("Server mode connectivity test", + # self.stackql.server_mode and + # hasattr(self.stackql, 'server_connection') and + # self.stackql.server_connection is not None and + # connection_working, # Include connection check in the pass criteria + # True) + + # # Add additional output about the actual connection status + # print(f" - Connection status: {'āœ… WORKING' if connection_working else 'āŒ NOT WORKING'}") + # print(f" - Expected status: āœ… WORKING") # Always expected to be working + + # # Always assert that the connection is working + # assert connection_working, "Server connection should be working" + @pystackql_test_setup(server_mode=True) def test_server_mode_connectivity(self): """Test that server mode connects successfully.""" - assert self.stackql.server_mode, "StackQL should be in server mode" - # Updated assertion to check server_connection attribute instead of _conn - assert hasattr(self.stackql, 'server_connection'), "StackQL should have a server_connection attribute" - assert self.stackql.server_connection is not None, "Server connection object should not be None" - - print_test_result("Server mode connectivity test", - self.stackql.server_mode and - hasattr(self.stackql, 'server_connection') and - self.stackql.server_connection is not None, - True) - + # Initialize class variable + TestServerMode.server_available = False + + # Perform basic server connection test + connection_working = self.stackql.test_connection() + + if not connection_working: + # Log minimal diagnostic info + print("\nāš ļø Server connection failed") + print(f"Address: {self.stackql.server_address}:{self.stackql.server_port}") + print("āŒ Skipping remaining server tests") + + # Fail with a concise message - this will be what shows in the error summary + pytest.fail("Server connection failed - please start stackql server") + + # Connection succeeded + TestServerMode.server_available = True + print("āœ… Server connection successful") + @pystackql_test_setup(server_mode=True) def test_server_mode_execute_stmt(self): """Test executeStmt in server mode.""" + + # Skip if server is not available + if not TestServerMode.server_available: + pytest.skip("Server is not available, skipping test") + result = self.stackql.executeStmt(REGISTRY_PULL_HOMEBREW_QUERY) # Check result structure @@ -60,6 +112,11 @@ def test_server_mode_execute_stmt(self): @pystackql_test_setup(server_mode=True, output='pandas') def test_server_mode_execute_stmt_pandas(self): """Test executeStmt in server mode with pandas output.""" + + # Skip if server is not available + if not TestServerMode.server_available: + pytest.skip("Server is not available, skipping test") + result = self.stackql.executeStmt(REGISTRY_PULL_HOMEBREW_QUERY) # Check result structure @@ -77,6 +134,11 @@ def test_server_mode_execute_stmt_pandas(self): @pystackql_test_setup(server_mode=True) def test_server_mode_execute(self): """Test execute in server mode.""" + + # Skip if server is not available + if not TestServerMode.server_available: + pytest.skip("Server is not available, skipping test") + result = self.stackql.execute(LITERAL_INT_QUERY) # Check result structure @@ -98,6 +160,11 @@ def test_server_mode_execute(self): @pystackql_test_setup(server_mode=True, output='pandas') def test_server_mode_execute_pandas(self): """Test execute in server mode with pandas output.""" + + # Skip if server is not available + if not TestServerMode.server_available: + pytest.skip("Server is not available, skipping test") + result = self.stackql.execute(LITERAL_STRING_QUERY) # Check result structure @@ -115,6 +182,11 @@ def test_server_mode_execute_pandas(self): @pystackql_test_setup(server_mode=True) def test_server_mode_provider_query(self): """Test querying a provider in server mode.""" + + # Skip if server is not available + if not TestServerMode.server_available: + pytest.skip("Server is not available, skipping test") + result = self.stackql.execute(HOMEBREW_FORMULA_QUERY) # Check result structure @@ -135,6 +207,11 @@ def test_server_mode_provider_query(self): @patch('pystackql.core.server.ServerConnection.execute_query') def test_server_mode_execute_mocked(self, mock_execute_query): """Test execute in server mode with mocked server response.""" + + # Skip if server is not available + if not TestServerMode.server_available: + pytest.skip("Server is not available, skipping test") + # Create a StackQL instance in server mode stackql = StackQL(server_mode=True) @@ -158,6 +235,11 @@ def test_server_mode_execute_mocked(self, mock_execute_query): @patch('pystackql.core.server.ServerConnection.execute_query') def test_server_mode_execute_pandas_mocked(self, mock_execute_query): """Test execute in server mode with pandas output and mocked server response.""" + + # Skip if server is not available + if not TestServerMode.server_available: + pytest.skip("Server is not available, skipping test") + # Create a StackQL instance in server mode with pandas output stackql = StackQL(server_mode=True, output='pandas') @@ -186,6 +268,11 @@ def test_server_mode_execute_pandas_mocked(self, mock_execute_query): @patch('pystackql.core.server.ServerConnection.execute_query') def test_server_mode_execute_stmt_mocked(self, mock_execute_query): """Test executeStmt in server mode with mocked server response.""" + + # Skip if server is not available + if not TestServerMode.server_available: + pytest.skip("Server is not available, skipping test") + # Create a StackQL instance in server mode stackql = StackQL(server_mode=True) @@ -208,6 +295,11 @@ def test_server_mode_execute_stmt_mocked(self, mock_execute_query): def test_server_mode_csv_output_error(self): """Test that server mode with csv output raises an error.""" + + # Skip if server is not available + if not TestServerMode.server_available: + pytest.skip("Server is not available, skipping test") + with pytest.raises(ValueError) as exc_info: StackQL(server_mode=True, output='csv')