Skip to content
Merged
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Prev Previous commit
Next Next commit
Implement kwargs override feature for execute and executeStmt methods
Co-authored-by: jeffreyaven <[email protected]>
  • Loading branch information
Copilot and jeffreyaven committed Oct 2, 2025
commit cc087c033009c9fd8a6355bed93391976123f2c3
86 changes: 86 additions & 0 deletions docs/source/intro.rst
Original file line number Diff line number Diff line change
Expand Up @@ -168,6 +168,92 @@ Here is an example of using the ``json_extract`` function to extract a field fro
res = stackql.execute(query)
print(res)

Overriding Parameters per Query
================================

The :meth:`pystackql.StackQL.execute` and :meth:`pystackql.StackQL.executeStmt` methods support keyword arguments that can override parameters set in the constructor for individual query executions. This is useful when you need to:

- Change the output format for specific queries
- Adjust CSV formatting (separator, headers) for specific exports
- Override authentication for specific providers
- Change other execution parameters on a per-query basis

**Example: Overriding Output Format**

You can create a StackQL instance with a default output format, then override it for specific queries:

.. code-block:: python

from pystackql import StackQL

# Create instance with CSV output by default
provider_auth = {
"github": {
"credentialsenvvar": "GITHUBCREDS",
"type": "basic"
}
}
stackql = StackQL(auth=provider_auth, output="csv")

# This returns CSV format (default)
csv_result = stackql.execute("select id, name from github.repos.repos where org = 'stackql'")
print(csv_result)
# Output:
# id,name
# 443987542,stackql
# 441087132,stackql-provider-registry
# ...

# This overrides to JSON/dict format for this query only
json_result = stackql.execute("select id, name from github.repos.repos where org = 'stackql'", output="json")
print(json_result)
# Output:
# [{"id":"443987542","name":"stackql"},{"id":"441087132","name":"stackql-provider-registry"},...]

# Subsequent calls without override use the original CSV format
csv_result2 = stackql.execute("select id, name from github.repos.repos where org = 'stackql' limit 1")

**Example: Overriding CSV Formatting**

You can also override CSV-specific parameters like separator and headers:

.. code-block:: python

from pystackql import StackQL

# Create instance with default CSV settings
stackql = StackQL(output="csv", sep=",", header=False)

# Override to use pipe separator and include headers for this query
result = stackql.execute(
"select id, name from github.repos.repos where org = 'stackql' limit 3",
sep="|",
header=True
)

**Supported Override Parameters**

The following parameters can be overridden in :meth:`pystackql.StackQL.execute` and :meth:`pystackql.StackQL.executeStmt`:

- ``output``: Output format ('dict', 'pandas', or 'csv')
- ``sep``: CSV delimiter/separator (when output='csv')
- ``header``: Include headers in CSV output (when output='csv')
- ``auth``: Custom authentication for providers
- ``custom_registry``: Custom StackQL provider registry URL
- ``max_results``: Maximum results per HTTP request
- ``page_limit``: Maximum pages per resource
- ``max_depth``: Maximum depth for indirect queries
- ``api_timeout``: API request timeout
- ``http_debug``: Enable HTTP debug logging
- Proxy settings: ``proxy_host``, ``proxy_port``, ``proxy_user``, ``proxy_password``, ``proxy_scheme``
- Backend settings: ``backend_storage_mode``, ``backend_file_storage_location``, ``app_root``
- Execution settings: ``execution_concurrency_limit``, ``dataflow_dependency_max``, ``dataflow_components_max``

.. note::

Parameter overrides only affect the specific query execution and do not modify the StackQL instance's configuration. Subsequent queries will use the original constructor parameters unless overridden again.


Using the Jupyter Magic Extension
=================================

Expand Down
5 changes: 3 additions & 2 deletions pystackql/core/query.py
Original file line number Diff line number Diff line change
Expand Up @@ -50,18 +50,19 @@ def _debug_log(self, message):
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):
def execute(self, query, custom_auth=None, env_vars=None, override_params=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.
override_params (list, optional): Override parameters for this execution. Defaults to None.

Returns:
dict: The query results
"""
local_params = self.params.copy()
local_params = (override_params if override_params is not None else self.params).copy()
script_path = None

# Format query for platform
Expand Down
83 changes: 70 additions & 13 deletions pystackql/core/stackql.py
Original file line number Diff line number Diff line change
Expand Up @@ -263,7 +263,7 @@ def upgrade(self, showprogress=True):

return message

def executeStmt(self, query, custom_auth=None, env_vars=None):
def executeStmt(self, query, custom_auth=None, env_vars=None, **kwargs):
"""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
Expand All @@ -279,6 +279,12 @@ def executeStmt(self, query, custom_auth=None, env_vars=None):
:type custom_auth: dict, optional
:param env_vars: Command-specific environment variables for this execution.
:type env_vars: dict, optional
:param kwargs: Additional keyword arguments that override constructor parameters for this execution.
Supported overrides: output, sep, header, auth, custom_registry, max_results, page_limit,
max_depth, api_timeout, http_debug, proxy_host, proxy_port, proxy_user, proxy_password,
proxy_scheme, backend_storage_mode, backend_file_storage_location, app_root,
execution_concurrency_limit, dataflow_dependency_max, dataflow_components_max
:type kwargs: optional

:return: The output result of the query in string format. If in `server_mode`, it
returns a JSON string representation of the result.
Expand All @@ -292,25 +298,47 @@ def executeStmt(self, query, custom_auth=None, env_vars=None):
>>> result
"""
if self.server_mode:
# Server mode: handle output override
output_format = kwargs.get('output', self.output)

result = self.server_connection.execute_query(query, is_statement=True)

# Format result based on output type
if self.output == 'pandas':
if output_format == 'pandas':
import pandas as pd
return pd.DataFrame(result)
elif self.output == 'csv':
elif output_format == 'csv':
# Return the string representation of the result
return result[0]['message']
else:
return result
else:
# Local mode: handle parameter overrides
override_params = None
output_format = kwargs.get('output', self.output)

# If custom_auth is provided as kwarg, use it
if 'auth' in kwargs:
custom_auth = kwargs['auth']

# Generate override params if kwargs provided
if kwargs:
from ..utils import generate_params_for_execution
override_params = generate_params_for_execution(self._base_kwargs, kwargs)

# Execute the query
result = self.local_query_executor.execute(query, custom_auth=custom_auth, env_vars=env_vars)
result = self.local_query_executor.execute(query, custom_auth=custom_auth, env_vars=env_vars, override_params=override_params)

# Format the result
return self.local_output_formatter.format_statement_result(result)
# Format the result with appropriate output formatter
if output_format != self.output:
# Create a temporary formatter for this execution
from .output import OutputFormatter
temp_formatter = OutputFormatter(output_format)
return temp_formatter.format_statement_result(result)
else:
return self.local_output_formatter.format_statement_result(result)

def execute(self, query, suppress_errors=True, custom_auth=None, env_vars=None):
def execute(self, query, suppress_errors=True, custom_auth=None, env_vars=None, **kwargs):
"""
Executes a StackQL query and returns the output based on the specified output format.

Expand All @@ -325,6 +353,12 @@ def execute(self, query, suppress_errors=True, custom_auth=None, env_vars=None):
:type custom_auth: dict, optional
:param env_vars: Command-specific environment variables for this execution.
:type env_vars: dict, optional
:param kwargs: Additional keyword arguments that override constructor parameters for this execution.
Supported overrides: output, sep, header, auth, custom_registry, max_results, page_limit,
max_depth, api_timeout, http_debug, proxy_host, proxy_port, proxy_user, proxy_password,
proxy_scheme, backend_storage_mode, backend_file_storage_location, app_root,
execution_concurrency_limit, dataflow_dependency_max, dataflow_components_max
:type kwargs: 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.
Expand All @@ -344,29 +378,52 @@ def execute(self, query, suppress_errors=True, custom_auth=None, env_vars=None):
>>> result = stackql.execute(query)
"""
if self.server_mode:
# Server mode: handle output override
output_format = kwargs.get('output', self.output)

result = self.server_connection.execute_query(query)

# Format result based on output type
if self.output == 'pandas':
if output_format == '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':
elif output_format == 'csv':
raise ValueError("CSV output is not supported in server_mode.")
else: # Assume 'dict' output
return result
else:
# Local mode: handle parameter overrides
override_params = None
output_format = kwargs.get('output', self.output)
http_debug = kwargs.get('http_debug', self.http_debug)

# If custom_auth is provided as kwarg, use it
if 'auth' in kwargs:
custom_auth = kwargs['auth']

# Generate override params if kwargs provided
if kwargs:
from ..utils import generate_params_for_execution
override_params = generate_params_for_execution(self._base_kwargs, kwargs)

# Apply HTTP debug setting
if self.http_debug:
if http_debug:
suppress_errors = False

# Execute the query
output = self.local_query_executor.execute(query, custom_auth=custom_auth, env_vars=env_vars)
output = self.local_query_executor.execute(query, custom_auth=custom_auth, env_vars=env_vars, override_params=override_params)

# Format the result
return self.local_output_formatter.format_query_result(output, suppress_errors)
# Format the result with appropriate output formatter
if output_format != self.output:
# Create a temporary formatter for this execution
from .output import OutputFormatter
temp_formatter = OutputFormatter(output_format)
return temp_formatter.format_query_result(output, suppress_errors)
else:
return self.local_output_formatter.format_query_result(output, suppress_errors)

async def executeQueriesAsync(self, queries):
"""Executes multiple StackQL queries asynchronously using the current StackQL instance.
Expand Down
5 changes: 3 additions & 2 deletions pystackql/utils/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,7 @@
)

from .auth import format_auth
from .params import setup_local_mode
from .params import setup_local_mode, generate_params_for_execution

__all__ = [
# Platform utilities
Expand All @@ -45,5 +45,6 @@
'format_auth',

# Parameter utilities
'setup_local_mode'
'setup_local_mode',
'generate_params_for_execution'
]
Loading