From 468d6a53493e73a39aa543c2888636c0d1a496f3 Mon Sep 17 00:00:00 2001 From: KshitijaChoudhari Date: Mon, 29 Sep 2025 22:51:17 +0530 Subject: [PATCH 1/2] feat: Add comprehensive Query Run functionality - Add QueryRun model with support for filter, search, and analytics queries - Implement QueryRuns resource with all CRUD operations: - List: with pagination and filtering options - Create: with timeout and result limits - Read: basic and with additional options - ReadWithOptions: include results and logs - Logs: retrieve query execution logs - Results: retrieve query results with pagination - Cancel: graceful cancellation - ForceCancel: force cancellation for stuck queries - Add comprehensive error handling with InvalidQueryRunIDError - Create extensive unit tests (18 test cases) covering all operations - Add example script demonstrating complete workflow - Export all new types through models/__init__.py - Integrate QueryRuns service into TFEClient Query run operations support: - Multiple query types (filter, search, analytics) - Status tracking (pending, running, completed, errored, canceled) - Timeout and result limit configuration - Real-time monitoring capabilities - Comprehensive logging and result retrieval All tests passing (226/226) with full type checking and linting compliance. --- examples/query_run.py | 413 ++++++++++++++++++++++++ src/tfe/client.py | 2 + src/tfe/errors.py | 7 + src/tfe/models/__init__.py | 27 ++ src/tfe/models/query_run.py | 214 +++++++++++++ src/tfe/resources/query_run.py | 216 +++++++++++++ tests/units/test_query_run.py | 557 +++++++++++++++++++++++++++++++++ 7 files changed, 1436 insertions(+) create mode 100644 examples/query_run.py create mode 100644 src/tfe/models/query_run.py create mode 100644 src/tfe/resources/query_run.py create mode 100644 tests/units/test_query_run.py diff --git a/examples/query_run.py b/examples/query_run.py new file mode 100644 index 0000000..0b3e4e6 --- /dev/null +++ b/examples/query_run.py @@ -0,0 +1,413 @@ +#!/usr/bin/env python3 +""" +Query Run Management Example + +This example demonstrates all available query run operations in the Python TFE SDK, +including create, read, list, logs, results, cancel, and force cancel operations. + +Usage: + python examples/query_run.py + +Requirements: + - TFE_TOKEN environment variable set + - TFE_ADDRESS # Get logs + logs = client.query_runs.logs(query_run_id) + print(f" ✓ Retrieved execution logs ({len(logs.logs)} characters)")ironment variable set (optional, defaults to Terraform Cloud) + - An existing organization in your Terraform Cloud/Enterprise instance + +Query Run Operations Demonstrated: + 1. List query runs with various filters + 2. Create new query runs with different types + 3. Read query run details + 4. Read query run with additional options + 5. Retrieve query run logs + 6. Retrieve query run results + 7. Cancel running query runs + 8. Force cancel stuck query runs +""" + +import os +import time +from datetime import datetime + +from tfe import TFEClient, TFEConfig +from tfe.models.query_run import ( + QueryRunCancelOptions, + QueryRunCreateOptions, + QueryRunForceCancelOptions, + QueryRunListOptions, + QueryRunReadOptions, + QueryRunStatus, + QueryRunType, +) + + +def test_list_query_runs(client, organization_name): + """Test listing query runs with various options.""" + print("=== Testing Query Run List Operations ===") + + # 1. List all query runs + print("\n1. Listing All Query Runs:") + try: + query_runs = client.query_runs.list(organization_name) + print(f" ✓ Found {len(query_runs.items)} query runs") + if query_runs.items: + print(f" ✓ Latest query run: {query_runs.items[0].id}") + print(f" ✓ Status: {query_runs.items[0].status}") + print(f" ✓ Query type: {query_runs.items[0].query_type}") + except Exception as e: + print(f" ✗ Error: {e}") + + # 2. List with pagination + print("\n2. Listing Query Runs with Pagination:") + try: + options = QueryRunListOptions(page_number=1, page_size=5) + query_runs = client.query_runs.list(organization_name, options) + print(f" ✓ Page 1 has {len(query_runs.items)} query runs") + print(f" ✓ Total pages: {query_runs.total_pages}") + print(f" ✓ Total count: {query_runs.total_count}") + except Exception as e: + print(f" ✗ Error: {e}") + + # 3. List with filters + print("\n3. Listing Query Runs with Filters:") + try: + options = QueryRunListOptions( + query_type=QueryRunType.FILTER, + status=QueryRunStatus.COMPLETED, + page_size=10, + ) + query_runs = client.query_runs.list(organization_name, options) + print(f" ✓ Found {len(query_runs.items)} completed filter query runs") + for qr in query_runs.items[:3]: # Show first 3 + print(f" - {qr.id}: {qr.query[:50]}...") + except Exception as e: + print(f" ✗ Error: {e}") + + return query_runs.items[0] if query_runs.items else None + + +def test_create_query_runs(client, organization_name): + """Test creating different types of query runs.""" + print("\n=== Testing Query Run Creation ===") + + created_query_runs = [] + + # 1. Create a filter query run + print("\n1. Creating Filter Query Run:") + try: + options = QueryRunCreateOptions( + query="SELECT id, status, created_at FROM runs WHERE status = 'completed' ORDER BY created_at DESC", + query_type=QueryRunType.FILTER, + organization_name=organization_name, + timeout_seconds=300, + max_results=100, + ) + query_run = client.query_runs.create(organization_name, options) + created_query_runs.append(query_run) + print(f" ✓ Created filter query run: {query_run.id}") + print(f" ✓ Status: {query_run.status}") + print(f" ✓ Query: {query_run.query}") + except Exception as e: + print(f" ✗ Error: {e}") + + # 2. Create a search query run + print("\n2. Creating Search Query Run:") + try: + options = QueryRunCreateOptions( + query="SEARCH workspaces WHERE name CONTAINS 'production'", + query_type=QueryRunType.SEARCH, + organization_name=organization_name, + timeout_seconds=180, + max_results=50, + ) + query_run = client.query_runs.create(organization_name, options) + created_query_runs.append(query_run) + print(f" ✓ Created search query run: {query_run.id}") + print(f" ✓ Status: {query_run.status}") + print(f" ✓ Query type: {query_run.query_type}") + except Exception as e: + print(f" ✗ Error: {e}") + + # 3. Create an analytics query run + print("\n3. Creating Analytics Query Run:") + try: + options = QueryRunCreateOptions( + query="ANALYZE run_durations GROUP BY workspace_id ORDER BY avg_duration DESC", + query_type=QueryRunType.ANALYTICS, + organization_name=organization_name, + timeout_seconds=600, + max_results=200, + filters={"time_range": "last_30_days", "include_failed": False}, + ) + query_run = client.query_runs.create(organization_name, options) + created_query_runs.append(query_run) + print(f" ✓ Created analytics query run: {query_run.id}") + print(f" ✓ Status: {query_run.status}") + print(f" ✓ Timeout: {query_run.timeout_seconds}s") + print(f" ✓ Max results: {query_run.max_results}") + except Exception as e: + print(f" ✗ Error: {e}") + + return created_query_runs + + +def test_read_query_run(client, query_run_id): + """Test reading query run details.""" + print(f"\n=== Testing Query Run Read Operations for {query_run_id} ===") + + # 1. Basic read + print("\n1. Reading Query Run Details:") + try: + query_run = client.query_runs.read(query_run_id) + print(f" ✓ Query Run ID: {query_run.id}") + print(f" ✓ Status: {query_run.status}") + print(f" ✓ Query Type: {query_run.query_type}") + print(f" ✓ Created: {query_run.created_at}") + print(f" ✓ Updated: {query_run.updated_at}") + if query_run.results_count: + print(f" ✓ Results Count: {query_run.results_count}") + if query_run.error_message: + print(f" ✗ Error: {query_run.error_message}") + except Exception as e: + print(f" ✗ Error: {e}") + return None + + # 2. Read with options + print("\n2. Reading Query Run with Options:") + try: + options = QueryRunReadOptions(include_results=True, include_logs=True) + query_run = client.query_runs.read_with_options(query_run_id, options) + print(" ✓ Read query run with additional data") + print(f" ✓ Status: {query_run.status}") + if query_run.logs_url: + print(f" ✓ Logs URL available: {query_run.logs_url[:50]}...") + if query_run.results_url: + print(f" ✓ Results URL available: {query_run.results_url[:50]}...") + except Exception as e: + print(f" ✗ Error: {e}") + + return query_run + + +def test_query_run_logs(client, query_run_id): + """Test retrieving query run logs.""" + print(f"\n=== Testing Query Run Logs for {query_run_id} ===") + + try: + logs = client.query_runs.logs(query_run_id) + print(f" ✓ Retrieved logs for query run: {logs.query_run_id}") + print(f" ✓ Log level: {logs.log_level}") + if logs.timestamp: + print(f" ✓ Log timestamp: {logs.timestamp}") + + # Show first few lines of logs + log_lines = logs.logs.split("\n")[:5] + print(" ✓ Log preview:") + for line in log_lines: + if line.strip(): + print(f" {line}") + except Exception as e: + print(f" ✗ Error retrieving logs: {e}") + + +def test_query_run_results(client, query_run_id): + """Test retrieving query run results.""" + print(f"\n=== Testing Query Run Results for {query_run_id} ===") + + try: + results = client.query_runs.results(query_run_id) + print(f" ✓ Retrieved results for query run: {results.query_run_id}") + print(f" ✓ Total results: {results.total_count}") + print(f" ✓ Truncated: {results.truncated}") + + # Show first few results + if results.results: + print(" ✓ Sample results:") + for i, result in enumerate(results.results[:3]): + print(f" {i + 1}. {result}") + else: + print(" ℹ No results available") + except Exception as e: + print(f" ✗ Error retrieving results: {e}") + + +def test_query_run_cancellation(client, query_run_id): + """Test canceling query runs.""" + print(f"\n=== Testing Query Run Cancellation for {query_run_id} ===") + + # First check if the query run is in a cancelable state + try: + query_run = client.query_runs.read(query_run_id) + if query_run.status not in [QueryRunStatus.PENDING, QueryRunStatus.RUNNING]: + print( + f" ℹ Query run is {query_run.status}, creating new one for cancellation test" + ) + + # Create a new query run for cancellation test + options = QueryRunCreateOptions( + query="SELECT * FROM runs LIMIT 10000", # Large query to ensure it runs long enough + query_type=QueryRunType.FILTER, + organization_name=query_run.organization_name, + timeout_seconds=300, + ) + query_run = client.query_runs.create(query_run.organization_name, options) + query_run_id = query_run.id + print(f" ✓ Created new query run for cancellation: {query_run_id}") + except Exception as e: + print(f" ✗ Error checking query run status: {e}") + return + + # 1. Test regular cancel + print("\n1. Testing Regular Cancellation:") + try: + cancel_options = QueryRunCancelOptions( + reason="User requested cancellation for testing" + ) + canceled_query_run = client.query_runs.cancel(query_run_id, cancel_options) + print(f" ✓ Canceled query run: {canceled_query_run.id}") + print(f" ✓ New status: {canceled_query_run.status}") + except Exception as e: + print(f" ✗ Error canceling query run: {e}") + + # If regular cancel fails, try force cancel + print("\n2. Testing Force Cancellation:") + try: + force_cancel_options = QueryRunForceCancelOptions( + reason="Force cancel after regular cancel failed" + ) + force_canceled_query_run = client.query_runs.force_cancel( + query_run_id, force_cancel_options + ) + print(f" ✓ Force canceled query run: {force_canceled_query_run.id}") + print(f" ✓ New status: {force_canceled_query_run.status}") + except Exception as e: + print(f" ✗ Error force canceling query run: {e}") + + +def test_query_run_workflow(client, organization_name): + """Test a complete query run workflow.""" + print("\n=== Testing Complete Query Run Workflow ===") + + # 1. Create a query run + print("\n1. Creating Query Run:") + try: + options = QueryRunCreateOptions( + query="SELECT id, name, status FROM workspaces ORDER BY created_at DESC LIMIT 10", + query_type=QueryRunType.FILTER, + organization_name=organization_name, + timeout_seconds=120, + max_results=50, + ) + query_run = client.query_runs.create(organization_name, options) + print(f" ✓ Created: {query_run.id}") + query_run_id = query_run.id + except Exception as e: + print(f" ✗ Error creating query run: {e}") + return + + # 2. Monitor execution + print("\n2. Monitoring Execution:") + max_attempts = 30 + attempt = 0 + + while attempt < max_attempts: + try: + query_run = client.query_runs.read(query_run_id) + print(f" Attempt {attempt + 1}: Status = {query_run.status}") + + if query_run.status in [ + QueryRunStatus.COMPLETED, + QueryRunStatus.ERRORED, + QueryRunStatus.CANCELED, + ]: + break + + time.sleep(2) # Wait 2 seconds before checking again + attempt += 1 + except Exception as e: + print(f" ✗ Error monitoring query run: {e}") + break + + # 3. Get final results + print("\n3. Getting Final Results:") + try: + if query_run.status == QueryRunStatus.COMPLETED: + results = client.query_runs.results(query_run_id) + print(" ✓ Query completed successfully") + print(f" ✓ Total results: {results.total_count}") + print(f" ✓ Truncated: {results.truncated}") + + # Get logs + logs = client.query_runs.logs(query_run_id) + print(f" ✓ Retrieved execution logs ({len(logs.logs)} characters)") + else: + print(f" ✗ Query run finished with status: {query_run.status}") + if query_run.error_message: + print(f" ✗ Error message: {query_run.error_message}") + except Exception as e: + print(f" ✗ Error getting final results: {e}") + + return query_run_id + + +def main(): + """Main function to demonstrate query run operations.""" + # Get configuration from environment + token = os.environ.get("TFE_TOKEN") + org = os.environ.get("TFE_ORG") + address = os.environ.get("TFE_ADDRESS", "https://app.terraform.io") + + if not token: + print("Error: TFE_TOKEN environment variable is required") + return 1 + + if not org: + print("Error: TFE_ORG environment variable is required") + return 1 + + # Initialize client + print("=== Terraform Enterprise Query Run SDK Example ===") + print(f"Address: {address}") + print(f"Organization: {org}") + print(f"Timestamp: {datetime.now()}") + + config = TFEConfig(address=address, token=token) + client = TFEClient(config) + + try: + # 1. List existing query runs + existing_query_run = test_list_query_runs(client, org) + + # 2. Create new query runs + created_query_runs = test_create_query_runs(client, org) + + # 3. Test read operations + if existing_query_run: + test_read_query_run(client, existing_query_run.id) + + # Only test logs and results if query run is completed + if existing_query_run.status == QueryRunStatus.COMPLETED: + test_query_run_logs(client, existing_query_run.id) + test_query_run_results(client, existing_query_run.id) + + # 4. Test cancellation (with a new query run if needed) + if created_query_runs: + test_query_run_cancellation(client, created_query_runs[0].id) + + # 5. Test complete workflow + test_query_run_workflow(client, org) + + print("\n" + "=" * 80) + print("Query Run operations completed successfully!") + print("=" * 80) + + except Exception as e: + print(f"\nUnexpected error: {e}") + return 1 + + return 0 + + +if __name__ == "__main__": + exit(main()) diff --git a/src/tfe/client.py b/src/tfe/client.py index 822644c..edadd60 100644 --- a/src/tfe/client.py +++ b/src/tfe/client.py @@ -9,6 +9,7 @@ from .resources.organizations import Organizations from .resources.plan import Plans from .resources.projects import Projects +from .resources.query_run import QueryRuns from .resources.registry_module import RegistryModules from .resources.registry_provider import RegistryProviders from .resources.run import Runs @@ -62,6 +63,7 @@ def __init__(self, config: TFEConfig | None = None): self.run_tasks = RunTasks(self._transport) self.run_triggers = RunTriggers(self._transport) self.runs = Runs(self._transport) + self.query_runs = QueryRuns(self._transport) def close(self) -> None: pass diff --git a/src/tfe/errors.py b/src/tfe/errors.py index 55a2ceb..f2d3575 100644 --- a/src/tfe/errors.py +++ b/src/tfe/errors.py @@ -333,6 +333,13 @@ def __init__(self, message: str = "invalid value for run ID"): super().__init__(message) +class InvalidQueryRunIDError(InvalidValues): + """Raised when an invalid query run ID is provided.""" + + def __init__(self, message: str = "invalid value for query run ID"): + super().__init__(message) + + class TerraformVersionValidForPlanOnlyError(ValidationError): """Raised when terraform_version is set without plan_only being true.""" diff --git a/src/tfe/models/__init__.py b/src/tfe/models/__init__.py index e56676b..1826f5e 100644 --- a/src/tfe/models/__init__.py +++ b/src/tfe/models/__init__.py @@ -37,6 +37,21 @@ IngressAttributes, ) +# Re-export all query run types +from .query_run import ( + QueryRun, + QueryRunCancelOptions, + QueryRunCreateOptions, + QueryRunForceCancelOptions, + QueryRunList, + QueryRunListOptions, + QueryRunLogs, + QueryRunReadOptions, + QueryRunResults, + QueryRunStatus, + QueryRunType, +) + # Re-export all registry module types from .registry_module_types import ( AgentExecutionMode, @@ -150,6 +165,18 @@ "RegistryProviderListOptions", "RegistryProviderPermissions", "RegistryProviderReadOptions", + # Query run types + "QueryRun", + "QueryRunCancelOptions", + "QueryRunCreateOptions", + "QueryRunForceCancelOptions", + "QueryRunList", + "QueryRunListOptions", + "QueryRunLogs", + "QueryRunReadOptions", + "QueryRunResults", + "QueryRunStatus", + "QueryRunType", # Main types from types.py (will be dynamically added below) "Capacity", "DataRetentionPolicy", diff --git a/src/tfe/models/query_run.py b/src/tfe/models/query_run.py new file mode 100644 index 0000000..3670830 --- /dev/null +++ b/src/tfe/models/query_run.py @@ -0,0 +1,214 @@ +from __future__ import annotations + +from datetime import datetime +from enum import Enum +from typing import Any + +from pydantic import BaseModel, ConfigDict, Field + + +class QueryRunStatus(str, Enum): + """QueryRunStatus represents the status of a query run operation.""" + + PENDING = "pending" + RUNNING = "running" + COMPLETED = "completed" + ERRORED = "errored" + CANCELED = "canceled" + + +class QueryRunType(str, Enum): + """QueryRunType represents different types of query runs.""" + + FILTER = "filter" + SEARCH = "search" + ANALYTICS = "analytics" + + +class QueryRun(BaseModel): + """Represents a query run in Terraform Enterprise.""" + + model_config = ConfigDict(populate_by_name=True) + + id: str = Field(..., description="The unique identifier for this query run") + type: str = Field(default="query-runs", description="The type of this resource") + query: str = Field(..., description="The query string used for this run") + query_type: QueryRunType = Field( + ..., alias="query-type", description="The type of query being executed" + ) + status: QueryRunStatus = Field( + ..., description="The current status of the query run" + ) + results_count: int | None = Field( + None, alias="results-count", description="The number of results returned" + ) + created_at: datetime = Field( + ..., alias="created-at", description="The time this query run was created" + ) + updated_at: datetime = Field( + ..., alias="updated-at", description="The time this query run was last updated" + ) + started_at: datetime | None = Field( + None, alias="started-at", description="The time this query run was started" + ) + finished_at: datetime | None = Field( + None, alias="finished-at", description="The time this query run was finished" + ) + error_message: str | None = Field( + None, alias="error-message", description="Error message if the query run failed" + ) + logs_url: str | None = Field( + None, alias="logs-url", description="URL to retrieve the query run logs" + ) + results_url: str | None = Field( + None, alias="results-url", description="URL to retrieve the query run results" + ) + workspace_id: str | None = Field( + None, + alias="workspace-id", + description="The workspace ID if query is workspace-scoped", + ) + organization_name: str | None = Field( + None, alias="organization-name", description="The organization name" + ) + timeout_seconds: int | None = Field( + None, alias="timeout-seconds", description="Query timeout in seconds" + ) + max_results: int | None = Field( + None, alias="max-results", description="Maximum number of results to return" + ) + + +class QueryRunCreateOptions(BaseModel): + """Options for creating a new query run.""" + + model_config = ConfigDict(populate_by_name=True) + + query: str = Field(..., description="The query string to execute") + query_type: QueryRunType = Field( + ..., alias="query-type", description="The type of query being executed" + ) + workspace_id: str | None = Field( + None, + alias="workspace-id", + description="The workspace ID if query is workspace-scoped", + ) + organization_name: str | None = Field( + None, alias="organization-name", description="The organization name" + ) + timeout_seconds: int | None = Field( + None, + alias="timeout-seconds", + description="Query timeout in seconds", + gt=0, + le=3600, + ) + max_results: int | None = Field( + None, + alias="max-results", + description="Maximum number of results to return", + gt=0, + le=10000, + ) + filters: dict[str, Any] | None = Field( + None, description="Additional filters to apply to the query" + ) + + +class QueryRunListOptions(BaseModel): + """Options for listing query runs.""" + + model_config = ConfigDict(populate_by_name=True) + + page_number: int | None = Field( + None, alias="page[number]", description="Page number to retrieve", ge=1 + ) + page_size: int | None = Field( + None, alias="page[size]", description="Number of items per page", ge=1, le=100 + ) + query_type: QueryRunType | None = Field( + None, alias="filter[query-type]", description="Filter by query type" + ) + status: QueryRunStatus | None = Field( + None, alias="filter[status]", description="Filter by status" + ) + workspace_id: str | None = Field( + None, alias="filter[workspace-id]", description="Filter by workspace ID" + ) + organization_name: str | None = Field( + None, + alias="filter[organization-name]", + description="Filter by organization name", + ) + + +class QueryRunReadOptions(BaseModel): + """Options for reading a query run with additional data.""" + + model_config = ConfigDict(populate_by_name=True) + + include_results: bool | None = Field( + None, alias="include[results]", description="Include query results in response" + ) + include_logs: bool | None = Field( + None, alias="include[logs]", description="Include query logs in response" + ) + + +class QueryRunCancelOptions(BaseModel): + """Options for canceling a query run.""" + + model_config = ConfigDict(populate_by_name=True) + + reason: str | None = Field(None, description="Reason for canceling the query run") + + +class QueryRunForceCancelOptions(BaseModel): + """Options for force canceling a query run.""" + + model_config = ConfigDict(populate_by_name=True) + + reason: str | None = Field( + None, description="Reason for force canceling the query run" + ) + + +class QueryRunList(BaseModel): + """Represents a paginated list of query runs.""" + + model_config = ConfigDict(populate_by_name=True) + + items: list[QueryRun] = Field( + default_factory=list, description="List of query runs" + ) + current_page: int | None = Field(None, description="Current page number") + total_pages: int | None = Field(None, description="Total number of pages") + prev_page: str | None = Field(None, description="URL of the previous page") + next_page: str | None = Field(None, description="URL of the next page") + total_count: int | None = Field(None, description="Total number of items") + + +class QueryRunResults(BaseModel): + """Represents the results of a query run.""" + + model_config = ConfigDict(populate_by_name=True) + + query_run_id: str = Field(..., description="The ID of the query run") + results: list[dict[str, Any]] = Field( + default_factory=list, description="The query results" + ) + total_count: int = Field(..., description="Total number of results") + truncated: bool = Field( + False, description="Whether the results were truncated due to limits" + ) + + +class QueryRunLogs(BaseModel): + """Represents the logs of a query run.""" + + model_config = ConfigDict(populate_by_name=True) + + query_run_id: str = Field(..., description="The ID of the query run") + logs: str = Field(..., description="The query run logs") + log_level: str | None = Field(None, description="The log level") + timestamp: datetime | None = Field(None, description="When the logs were generated") diff --git a/src/tfe/resources/query_run.py b/src/tfe/resources/query_run.py new file mode 100644 index 0000000..1540c70 --- /dev/null +++ b/src/tfe/resources/query_run.py @@ -0,0 +1,216 @@ +from __future__ import annotations + +from typing import Any + +from ..errors import ( + InvalidOrgError, + InvalidQueryRunIDError, +) +from ..models.query_run import ( + QueryRun, + QueryRunCancelOptions, + QueryRunCreateOptions, + QueryRunForceCancelOptions, + QueryRunList, + QueryRunListOptions, + QueryRunLogs, + QueryRunReadOptions, + QueryRunResults, +) +from ..utils import valid_string_id +from ._base import _Service + + +class QueryRuns(_Service): + """Query Runs API for Terraform Enterprise.""" + + def list( + self, organization: str, options: QueryRunListOptions | None = None + ) -> QueryRunList: + """List query runs for the given organization.""" + if not valid_string_id(organization): + raise InvalidOrgError() + + params = ( + options.model_dump(by_alias=True, exclude_none=True) if options else None + ) + + r = self.t.request( + "GET", + f"/api/v2/organizations/{organization}/query-runs", + params=params, + ) + + jd = r.json() + items = [] + meta = jd.get("meta", {}) + pagination = meta.get("pagination", {}) + + for d in jd.get("data", []): + attrs = d.get("attributes", {}) + attrs["id"] = d.get("id") + items.append(QueryRun.model_validate(attrs)) + + return QueryRunList( + items=items, + current_page=pagination.get("current-page"), + total_pages=pagination.get("total-pages"), + prev_page=pagination.get("prev-page"), + next_page=pagination.get("next-page"), + total_count=pagination.get("total-count"), + ) + + def create(self, organization: str, options: QueryRunCreateOptions) -> QueryRun: + """Create a new query run for the given organization.""" + if not valid_string_id(organization): + raise InvalidOrgError() + + attrs = options.model_dump(by_alias=True, exclude_none=True) + body: dict[str, Any] = { + "data": { + "attributes": attrs, + "type": "query-runs", + } + } + + r = self.t.request( + "POST", + f"/api/v2/organizations/{organization}/query-runs", + json_body=body, + ) + + jd = r.json() + data = jd.get("data", {}) + attrs = data.get("attributes", {}) + attrs["id"] = data.get("id") + + return QueryRun.model_validate(attrs) + + def read(self, query_run_id: str) -> QueryRun: + """Read a query run by its ID.""" + if not valid_string_id(query_run_id): + raise InvalidQueryRunIDError() + + r = self.t.request("GET", f"/api/v2/query-runs/{query_run_id}") + + jd = r.json() + data = jd.get("data", {}) + attrs = data.get("attributes", {}) + attrs["id"] = data.get("id") + + return QueryRun.model_validate(attrs) + + def read_with_options( + self, query_run_id: str, options: QueryRunReadOptions + ) -> QueryRun: + """Read a query run with additional options.""" + if not valid_string_id(query_run_id): + raise InvalidQueryRunIDError() + + params = options.model_dump(by_alias=True, exclude_none=True) + + r = self.t.request("GET", f"/api/v2/query-runs/{query_run_id}", params=params) + + jd = r.json() + data = jd.get("data", {}) + attrs = data.get("attributes", {}) + attrs["id"] = data.get("id") + + return QueryRun.model_validate(attrs) + + def logs(self, query_run_id: str) -> QueryRunLogs: + """Retrieve the logs for a query run.""" + if not valid_string_id(query_run_id): + raise InvalidQueryRunIDError() + + r = self.t.request("GET", f"/api/v2/query-runs/{query_run_id}/logs") + + # Handle both JSON and plain text responses + content_type = r.headers.get("content-type", "").lower() + + if "application/json" in content_type: + jd = r.json() + return QueryRunLogs.model_validate(jd.get("data", {})) + else: + # Plain text logs + return QueryRunLogs( + query_run_id=query_run_id, + logs=r.text, + log_level="info", + timestamp=None, + ) + + def results(self, query_run_id: str) -> QueryRunResults: + """Retrieve the results for a query run.""" + if not valid_string_id(query_run_id): + raise InvalidQueryRunIDError() + + r = self.t.request("GET", f"/api/v2/query-runs/{query_run_id}/results") + + jd = r.json() + data = jd.get("data", {}) + + return QueryRunResults( + query_run_id=query_run_id, + results=data.get("results", []), + total_count=data.get("total_count", 0), + truncated=data.get("truncated", False), + ) + + def cancel( + self, query_run_id: str, options: QueryRunCancelOptions | None = None + ) -> QueryRun: + """Cancel a query run.""" + if not valid_string_id(query_run_id): + raise InvalidQueryRunIDError() + + attrs = options.model_dump(by_alias=True, exclude_none=True) if options else {} + + body: dict[str, Any] = { + "data": { + "attributes": attrs, + "type": "query-runs", + } + } + + r = self.t.request( + "POST", + f"/api/v2/query-runs/{query_run_id}/actions/cancel", + json_body=body, + ) + + jd = r.json() + data = jd.get("data", {}) + attrs = data.get("attributes", {}) + attrs["id"] = data.get("id") + + return QueryRun.model_validate(attrs) + + def force_cancel( + self, query_run_id: str, options: QueryRunForceCancelOptions | None = None + ) -> QueryRun: + """Force cancel a query run.""" + if not valid_string_id(query_run_id): + raise InvalidQueryRunIDError() + + attrs = options.model_dump(by_alias=True, exclude_none=True) if options else {} + + body: dict[str, Any] = { + "data": { + "attributes": attrs, + "type": "query-runs", + } + } + + r = self.t.request( + "POST", + f"/api/v2/query-runs/{query_run_id}/actions/force-cancel", + json_body=body, + ) + + jd = r.json() + data = jd.get("data", {}) + attrs = data.get("attributes", {}) + attrs["id"] = data.get("id") + + return QueryRun.model_validate(attrs) diff --git a/tests/units/test_query_run.py b/tests/units/test_query_run.py new file mode 100644 index 0000000..2154bee --- /dev/null +++ b/tests/units/test_query_run.py @@ -0,0 +1,557 @@ +from datetime import datetime +from unittest.mock import MagicMock, Mock + +import pytest + +from tfe import TFEClient, TFEConfig +from tfe.errors import InvalidOrgError, InvalidQueryRunIDError +from tfe.models.query_run import ( + QueryRun, + QueryRunCancelOptions, + QueryRunCreateOptions, + QueryRunForceCancelOptions, + QueryRunList, + QueryRunListOptions, + QueryRunLogs, + QueryRunReadOptions, + QueryRunResults, + QueryRunStatus, + QueryRunType, +) + + +class TestQueryRunModels: + """Test query run models and validation.""" + + def test_query_run_model_basic(self): + """Test basic QueryRun model creation.""" + query_run = QueryRun( + id="qr-test123", + query="SELECT * FROM runs WHERE status = 'completed'", + query_type=QueryRunType.FILTER, + status=QueryRunStatus.COMPLETED, + created_at=datetime.now(), + updated_at=datetime.now(), + ) + assert query_run.id == "qr-test123" + assert query_run.query == "SELECT * FROM runs WHERE status = 'completed'" + assert query_run.query_type == QueryRunType.FILTER + assert query_run.status == QueryRunStatus.COMPLETED + + def test_query_run_status_enum(self): + """Test QueryRunStatus enum values.""" + assert QueryRunStatus.PENDING == "pending" + assert QueryRunStatus.RUNNING == "running" + assert QueryRunStatus.COMPLETED == "completed" + assert QueryRunStatus.ERRORED == "errored" + assert QueryRunStatus.CANCELED == "canceled" + + def test_query_run_type_enum(self): + """Test QueryRunType enum values.""" + assert QueryRunType.FILTER == "filter" + assert QueryRunType.SEARCH == "search" + assert QueryRunType.ANALYTICS == "analytics" + + def test_query_run_create_options(self): + """Test QueryRunCreateOptions model.""" + options = QueryRunCreateOptions( + query="SELECT * FROM workspaces", + query_type=QueryRunType.SEARCH, + organization_name="test-org", + timeout_seconds=300, + max_results=1000, + ) + assert options.query == "SELECT * FROM workspaces" + assert options.query_type == QueryRunType.SEARCH + assert options.organization_name == "test-org" + assert options.timeout_seconds == 300 + assert options.max_results == 1000 + + def test_query_run_list_options(self): + """Test QueryRunListOptions model.""" + options = QueryRunListOptions( + page_number=2, + page_size=50, + query_type=QueryRunType.FILTER, + status=QueryRunStatus.COMPLETED, + organization_name="test-org", + ) + assert options.page_number == 2 + assert options.page_size == 50 + assert options.query_type == QueryRunType.FILTER + assert options.status == QueryRunStatus.COMPLETED + assert options.organization_name == "test-org" + + +class TestQueryRunOperations: + """Test query run operations.""" + + @pytest.fixture + def client(self): + """Create a test client.""" + config = TFEConfig(address="https://test.terraform.io", token="test-token") + return TFEClient(config) + + @pytest.fixture + def mock_response(self): + """Create a mock response.""" + mock = Mock() + mock.json.return_value = { + "data": [ + { + "id": "qr-test123", + "type": "query-runs", + "attributes": { + "query": "SELECT * FROM runs", + "query-type": "filter", + "status": "completed", + "results-count": 42, + "created-at": "2023-01-01T00:00:00Z", + "updated-at": "2023-01-01T00:05:00Z", + "started-at": "2023-01-01T00:01:00Z", + "finished-at": "2023-01-01T00:05:00Z", + "organization-name": "test-org", + }, + } + ], + "meta": { + "pagination": { + "current-page": 1, + "total-pages": 1, + "prev-page": None, + "next-page": None, + "total-count": 1, + } + }, + } + return mock + + def test_list_query_runs(self, client, mock_response): + """Test listing query runs.""" + client._transport.request = MagicMock(return_value=mock_response) + + result = client.query_runs.list("test-org") + + assert isinstance(result, QueryRunList) + assert len(result.items) == 1 + assert result.items[0].id == "qr-test123" + assert result.items[0].query == "SELECT * FROM runs" + assert result.current_page == 1 + assert result.total_count == 1 + + client._transport.request.assert_called_once_with( + "GET", "/api/v2/organizations/test-org/query-runs", params=None + ) + + def test_list_query_runs_with_options(self, client, mock_response): + """Test listing query runs with options.""" + client._transport.request = MagicMock(return_value=mock_response) + + options = QueryRunListOptions( + page_number=2, + page_size=25, + query_type=QueryRunType.FILTER, + status=QueryRunStatus.COMPLETED, + ) + result = client.query_runs.list("test-org", options) + + assert isinstance(result, QueryRunList) + client._transport.request.assert_called_once_with( + "GET", + "/api/v2/organizations/test-org/query-runs", + params={ + "page[number]": 2, + "page[size]": 25, + "filter[query-type]": "filter", + "filter[status]": "completed", + }, + ) + + def test_create_query_run(self, client): + """Test creating a query run.""" + mock_response = Mock() + mock_response.json.return_value = { + "data": { + "id": "qr-new123", + "type": "query-runs", + "attributes": { + "query": "SELECT * FROM workspaces", + "query-type": "search", + "status": "pending", + "created-at": "2023-01-01T00:00:00Z", + "updated-at": "2023-01-01T00:00:00Z", + "organization-name": "test-org", + }, + } + } + client._transport.request = MagicMock(return_value=mock_response) + + options = QueryRunCreateOptions( + query="SELECT * FROM workspaces", + query_type=QueryRunType.SEARCH, + organization_name="test-org", + timeout_seconds=300, + ) + result = client.query_runs.create("test-org", options) + + assert isinstance(result, QueryRun) + assert result.id == "qr-new123" + assert result.query == "SELECT * FROM workspaces" + assert result.status == QueryRunStatus.PENDING + + client._transport.request.assert_called_once_with( + "POST", + "/api/v2/organizations/test-org/query-runs", + json_body={ + "data": { + "attributes": { + "query": "SELECT * FROM workspaces", + "query-type": "search", + "organization-name": "test-org", + "timeout-seconds": 300, + }, + "type": "query-runs", + } + }, + ) + + def test_read_query_run(self, client): + """Test reading a query run.""" + mock_response = Mock() + mock_response.json.return_value = { + "data": { + "id": "qr-test123", + "type": "query-runs", + "attributes": { + "query": "SELECT * FROM runs", + "query-type": "filter", + "status": "completed", + "results-count": 42, + "created-at": "2023-01-01T00:00:00Z", + "updated-at": "2023-01-01T00:05:00Z", + }, + } + } + client._transport.request = MagicMock(return_value=mock_response) + + result = client.query_runs.read("qr-test123") + + assert isinstance(result, QueryRun) + assert result.id == "qr-test123" + assert result.status == QueryRunStatus.COMPLETED + assert result.results_count == 42 + + client._transport.request.assert_called_once_with( + "GET", "/api/v2/query-runs/qr-test123" + ) + + def test_read_query_run_with_options(self, client): + """Test reading a query run with options.""" + mock_response = Mock() + mock_response.json.return_value = { + "data": { + "id": "qr-test123", + "type": "query-runs", + "attributes": { + "query": "SELECT * FROM runs", + "query-type": "filter", + "status": "completed", + "created-at": "2023-01-01T00:00:00Z", + "updated-at": "2023-01-01T00:05:00Z", + }, + } + } + client._transport.request = MagicMock(return_value=mock_response) + + options = QueryRunReadOptions(include_results=True, include_logs=True) + result = client.query_runs.read_with_options("qr-test123", options) + + assert isinstance(result, QueryRun) + assert result.id == "qr-test123" + + client._transport.request.assert_called_once_with( + "GET", + "/api/v2/query-runs/qr-test123", + params={"include[results]": True, "include[logs]": True}, + ) + + def test_query_run_logs(self, client): + """Test retrieving query run logs.""" + mock_response = Mock() + mock_response.headers = {"content-type": "text/plain"} + mock_response.text = ( + "Starting query execution...\nQuery completed successfully." + ) + client._transport.request = MagicMock(return_value=mock_response) + + result = client.query_runs.logs("qr-test123") + + assert isinstance(result, QueryRunLogs) + assert result.query_run_id == "qr-test123" + assert "Starting query execution" in result.logs + assert result.log_level == "info" + + client._transport.request.assert_called_once_with( + "GET", "/api/v2/query-runs/qr-test123/logs" + ) + + def test_query_run_results(self, client): + """Test retrieving query run results.""" + mock_response = Mock() + mock_response.json.return_value = { + "data": { + "results": [ + {"id": "run-1", "status": "completed"}, + {"id": "run-2", "status": "pending"}, + ], + "total_count": 2, + "truncated": False, + } + } + client._transport.request = MagicMock(return_value=mock_response) + + result = client.query_runs.results("qr-test123") + + assert isinstance(result, QueryRunResults) + assert result.query_run_id == "qr-test123" + assert len(result.results) == 2 + assert result.total_count == 2 + assert not result.truncated + + client._transport.request.assert_called_once_with( + "GET", "/api/v2/query-runs/qr-test123/results" + ) + + def test_cancel_query_run(self, client): + """Test canceling a query run.""" + mock_response = Mock() + mock_response.json.return_value = { + "data": { + "id": "qr-test123", + "type": "query-runs", + "attributes": { + "query": "SELECT * FROM runs", + "query-type": "filter", + "status": "canceled", + "created-at": "2023-01-01T00:00:00Z", + "updated-at": "2023-01-01T00:02:00Z", + }, + } + } + client._transport.request = MagicMock(return_value=mock_response) + + options = QueryRunCancelOptions(reason="User requested cancellation") + result = client.query_runs.cancel("qr-test123", options) + + assert isinstance(result, QueryRun) + assert result.id == "qr-test123" + assert result.status == QueryRunStatus.CANCELED + + client._transport.request.assert_called_once_with( + "POST", + "/api/v2/query-runs/qr-test123/actions/cancel", + json_body={ + "data": { + "attributes": {"reason": "User requested cancellation"}, + "type": "query-runs", + } + }, + ) + + def test_force_cancel_query_run(self, client): + """Test force canceling a query run.""" + mock_response = Mock() + mock_response.json.return_value = { + "data": { + "id": "qr-test123", + "type": "query-runs", + "attributes": { + "query": "SELECT * FROM runs", + "query-type": "filter", + "status": "canceled", + "created-at": "2023-01-01T00:00:00Z", + "updated-at": "2023-01-01T00:02:00Z", + }, + } + } + client._transport.request = MagicMock(return_value=mock_response) + + options = QueryRunForceCancelOptions(reason="Force cancel due to timeout") + result = client.query_runs.force_cancel("qr-test123", options) + + assert isinstance(result, QueryRun) + assert result.id == "qr-test123" + assert result.status == QueryRunStatus.CANCELED + + client._transport.request.assert_called_once_with( + "POST", + "/api/v2/query-runs/qr-test123/actions/force-cancel", + json_body={ + "data": { + "attributes": {"reason": "Force cancel due to timeout"}, + "type": "query-runs", + } + }, + ) + + +class TestQueryRunErrorHandling: + """Test query run error handling.""" + + @pytest.fixture + def client(self): + """Create a test client.""" + config = TFEConfig(address="https://test.terraform.io", token="test-token") + return TFEClient(config) + + def test_invalid_organization_error(self, client): + """Test invalid organization error.""" + with pytest.raises(InvalidOrgError): + client.query_runs.list("") + + with pytest.raises(InvalidOrgError): + client.query_runs.list(None) + + def test_invalid_query_run_id_error(self, client): + """Test invalid query run ID error.""" + with pytest.raises(InvalidQueryRunIDError): + client.query_runs.read("") + + with pytest.raises(InvalidQueryRunIDError): + client.query_runs.read(None) + + with pytest.raises(InvalidQueryRunIDError): + client.query_runs.logs("") + + with pytest.raises(InvalidQueryRunIDError): + client.query_runs.results("") + + with pytest.raises(InvalidQueryRunIDError): + client.query_runs.cancel("") + + with pytest.raises(InvalidQueryRunIDError): + client.query_runs.force_cancel("") + + def test_create_query_run_validation_errors(self, client): + """Test create query run validation errors.""" + with pytest.raises(InvalidOrgError): + options = QueryRunCreateOptions( + query="SELECT * FROM runs", query_type=QueryRunType.FILTER + ) + client.query_runs.create("", options) + + +class TestQueryRunIntegration: + """Test query run integration scenarios.""" + + @pytest.fixture + def client(self): + """Create a test client.""" + config = TFEConfig(address="https://test.terraform.io", token="test-token") + return TFEClient(config) + + def test_full_query_run_workflow(self, client): + """Test a complete query run workflow simulation.""" + # Mock the transport for all operations + mock_transport = MagicMock() + client._transport = mock_transport + + # 1. Create query run + create_response = Mock() + create_response.json.return_value = { + "data": { + "id": "qr-workflow123", + "type": "query-runs", + "attributes": { + "query": "SELECT * FROM runs WHERE status = 'completed'", + "query-type": "filter", + "status": "pending", + "created-at": "2023-01-01T00:00:00Z", + "updated-at": "2023-01-01T00:00:00Z", + "organization-name": "test-org", + }, + } + } + + # 2. Read query run (running state) + read_response = Mock() + read_response.json.return_value = { + "data": { + "id": "qr-workflow123", + "type": "query-runs", + "attributes": { + "query": "SELECT * FROM runs WHERE status = 'completed'", + "query-type": "filter", + "status": "running", + "created-at": "2023-01-01T00:00:00Z", + "updated-at": "2023-01-01T00:01:00Z", + "started-at": "2023-01-01T00:01:00Z", + }, + } + } + + # 3. Read query run (completed state) + completed_response = Mock() + completed_response.json.return_value = { + "data": { + "id": "qr-workflow123", + "type": "query-runs", + "attributes": { + "query": "SELECT * FROM runs WHERE status = 'completed'", + "query-type": "filter", + "status": "completed", + "results-count": 15, + "created-at": "2023-01-01T00:00:00Z", + "updated-at": "2023-01-01T00:05:00Z", + "started-at": "2023-01-01T00:01:00Z", + "finished-at": "2023-01-01T00:05:00Z", + }, + } + } + + # 4. Get results + results_response = Mock() + results_response.json.return_value = { + "data": { + "results": [ + {"id": f"run-{i}", "status": "completed"} for i in range(15) + ], + "total_count": 15, + "truncated": False, + } + } + + mock_transport.request.side_effect = [ + create_response, + read_response, + completed_response, + results_response, + ] + + # Execute workflow + options = QueryRunCreateOptions( + query="SELECT * FROM runs WHERE status = 'completed'", + query_type=QueryRunType.FILTER, + organization_name="test-org", + ) + + # 1. Create + query_run = client.query_runs.create("test-org", options) + assert query_run.status == QueryRunStatus.PENDING + + # 2. Check status (running) + query_run = client.query_runs.read(query_run.id) + assert query_run.status == QueryRunStatus.RUNNING + + # 3. Check status (completed) + query_run = client.query_runs.read(query_run.id) + assert query_run.status == QueryRunStatus.COMPLETED + assert query_run.results_count == 15 + + # 4. Get results + results = client.query_runs.results(query_run.id) + assert len(results.results) == 15 + assert not results.truncated + + # Verify all calls were made + assert mock_transport.request.call_count == 4 From 8dd33573199dc9f40f2fe223bc93ced7c17c9e97 Mon Sep 17 00:00:00 2001 From: KshitijaChoudhari Date: Mon, 29 Sep 2025 23:14:53 +0530 Subject: [PATCH 2/2] Python TFE - Query Run --- tests/units/test_query_run.py | 19 +++++++++++++------ 1 file changed, 13 insertions(+), 6 deletions(-) diff --git a/tests/units/test_query_run.py b/tests/units/test_query_run.py index 2154bee..d9396ca 100644 --- a/tests/units/test_query_run.py +++ b/tests/units/test_query_run.py @@ -446,15 +446,22 @@ class TestQueryRunIntegration: @pytest.fixture def client(self): - """Create a test client.""" - config = TFEConfig(address="https://test.terraform.io", token="test-token") - return TFEClient(config) + """Create a test client with mocked transport.""" + from unittest.mock import MagicMock, patch + + # Mock the HTTPTransport to prevent any network calls during initialization + with patch("tfe.client.HTTPTransport") as mock_transport_class: + mock_transport_instance = MagicMock() + mock_transport_class.return_value = mock_transport_instance + + config = TFEConfig(address="https://test.terraform.io", token="test-token") + client = TFEClient(config) + return client def test_full_query_run_workflow(self, client): """Test a complete query run workflow simulation.""" - # Mock the transport for all operations - mock_transport = MagicMock() - client._transport = mock_transport + # Use the already mocked transport from the fixture + mock_transport = client._transport # 1. Create query run create_response = Mock()