From 36a3f90aedd282564f45209047504b6480eb8e95 Mon Sep 17 00:00:00 2001 From: Claudiu Popa Date: Thu, 29 Oct 2026 11:23:18 +0100 Subject: [PATCH 1/4] Add Athena query command Generated by AI tools, and reviewed by Claudiu Popa Co-authored-by: Baris Kurt --- awscli/customizations/athena/__init__.py | 21 + awscli/customizations/athena/run_query.py | 1147 +++++++++++++++++ .../compositereadonly/__init__.py | 33 + .../compositereadonly/helpers.py | 546 ++++++++ .../compositereadonly/output.py | 56 + .../customizations/compositereadonly/poll.py | 180 +++ awscli/handlers_registry.py | 3 + tests/functional/athena/__init__.py | 12 + tests/functional/athena/test_run_query.py | 133 ++ tests/unit/customizations/athena/__init__.py | 12 + .../athena/test_resolve_workgroup.py | 226 ++++ .../customizations/athena/test_run_query.py | 1106 ++++++++++++++++ .../compositereadonly/__init__.py | 12 + .../compositereadonly/test_helpers.py | 382 ++++++ .../compositereadonly/test_output.py | 92 ++ .../compositereadonly/test_poll.py | 116 ++ 16 files changed, 4077 insertions(+) create mode 100644 awscli/customizations/athena/__init__.py create mode 100644 awscli/customizations/athena/run_query.py create mode 100644 awscli/customizations/compositereadonly/__init__.py create mode 100644 awscli/customizations/compositereadonly/helpers.py create mode 100644 awscli/customizations/compositereadonly/output.py create mode 100644 awscli/customizations/compositereadonly/poll.py create mode 100644 tests/functional/athena/__init__.py create mode 100644 tests/functional/athena/test_run_query.py create mode 100644 tests/unit/customizations/athena/__init__.py create mode 100644 tests/unit/customizations/athena/test_resolve_workgroup.py create mode 100644 tests/unit/customizations/athena/test_run_query.py create mode 100644 tests/unit/customizations/compositereadonly/__init__.py create mode 100644 tests/unit/customizations/compositereadonly/test_helpers.py create mode 100644 tests/unit/customizations/compositereadonly/test_output.py create mode 100644 tests/unit/customizations/compositereadonly/test_poll.py diff --git a/awscli/customizations/athena/__init__.py b/awscli/customizations/athena/__init__.py new file mode 100644 index 000000000000..18744a21e176 --- /dev/null +++ b/awscli/customizations/athena/__init__.py @@ -0,0 +1,21 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"). You +# may not use this file except in compliance with the License. A copy of +# the License is located at +# +# http://aws.amazon.com/apache2.0/ +# +# or in the "license" file accompanying this file. This file is +# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF +# ANY KIND, either express or implied. See the License for the specific +# language governing permissions and limitations under the License. +from awscli.customizations.athena.run_query import RunQueryCommand + + +def register(event_handlers): + event_handlers.register('building-command-table.athena', _inject_commands) + + +def _inject_commands(command_table, session, **kwargs): + command_table['run-query'] = RunQueryCommand(session) diff --git a/awscli/customizations/athena/run_query.py b/awscli/customizations/athena/run_query.py new file mode 100644 index 000000000000..63e22ce735c4 --- /dev/null +++ b/awscli/customizations/athena/run_query.py @@ -0,0 +1,1147 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"). You +# may not use this file except in compliance with the License. A copy of +# the License is located at +# +# http://aws.amazon.com/apache2.0/ +# +# or in the "license" file accompanying this file. This file is +# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF +# ANY KIND, either express or implied. See the License for the specific +# language governing permissions and limitations under the License. +"""``aws athena run-query`` -- run one Athena statement and return every row. + +Callers otherwise hand-roll the Athena lifecycle: +``StartQueryExecution`` -> ``GetQueryExecution`` (poll) -> ``GetQueryResults`` +(paginate). This command runs it once and returns the full, typed result set in +one call. + +This command is mutative (read-write): it submits whatever SQL it is given -- +SELECT/WITH/SHOW/DESCRIBE as well as INSERT, CREATE, DROP, MERGE, UNLOAD and +other DDL/write statements -- without inspecting or restricting the SQL. The +shared call bridge still refuses mutating *API operations* by name, but +``StartQueryExecution`` is allowed through, so a statement's effect on your data +is not constrained here; choose the workgroup and database with that in mind. + +Procedure: + +0. Validate the parameters (workgroup/catalog/database shape, S3 URI, row cap and + polling parameters), all before any AWS call. The SQL itself is passed through + verbatim and is neither inspected nor restricted. +1. ``athena:GetWorkGroup`` -- honour the workgroup's enforced output location and + engine version. When no workgroup is given, auto-resolve one whose effective + configuration actually has an output location. +2. ``athena:StartQueryExecution``. +3. ``athena:GetQueryExecution`` -- poll with backoff to SUCCEEDED, FAILED or + CANCELLED, using the shared bounded poller. +4. ``athena:GetQueryResults`` -- paginate on NextToken, drop the header row, and + return column types, ``DataScannedInBytes`` and an exact resume token if + ``--max-rows`` truncates the result. +""" + +from __future__ import annotations + +import time +from typing import Any, Callable, NamedTuple + +from awscli.customizations.commands import BasicCommand +from awscli.customizations.compositereadonly.helpers import ( + CONFLICT, + CREDENTIALS_EXPIRED, + INVALID_INPUT, + NETWORK_ERROR, + NOT_FOUND, + NOT_READY, + PERMISSION_DENIED, + SERVICE_ERROR, + THROTTLED, + TIMEOUT, + call, + check_non_empty_str, + check_positive_int, + check_positive_number, + create_read_client, + emit_result, + err_result, + invalid_input, + ok_result, +) +from awscli.customizations.compositereadonly.poll import drive_poll + +_TERMINAL_STATES = frozenset({"SUCCEEDED", "FAILED", "CANCELLED"}) +_INTEGER_TYPES = frozenset({"tinyint", "smallint", "integer", "int", "bigint"}) +_MAX_PAGE_SIZE = 1000 +_MAX_POLL_INTERVAL = 5.0 +_MAX_POLLS = 10_000 +_MAX_TRANSIENT_POLL_ERRORS = 5 +_NOT_SUPPORTED_MSG = "queries of this type are not supported" + +_VERBATIM_GUIDANCE = ( + "Cell values are verbatim Athena strings (timestamps such as " + "'2024-01-15 00:00:00.000' and ids are untouched); print them exactly as " + "returned, do not round, abbreviate or relabel them." +) +_NULL_GUIDANCE = "SQL NULL is returned as JSON null; render it explicitly as NULL in any report." + + +CallFn = Callable[[str, dict[str, Any]], dict[str, Any]] + +# Bounded paging for workgroup enumeration so auto-resolution can never run away +# on an account with thousands of workgroups. +_MAX_WORKGROUP_PAGES = 10 + + +def _workgroup_output(cfg: dict[str, Any]) -> tuple[str | None, bool]: + """Return (effective_output_location, has_managed_results) for a WG config.""" + loc = (cfg.get("ResultConfiguration") or {}).get("OutputLocation") + managed = bool( + (cfg.get("ManagedQueryResultsConfiguration") or {}).get("Enabled") + ) + return loc, managed + + +def resolve_workgroup_with_output( + call_op: CallFn, requested: str | None +) -> dict[str, Any]: + """Resolve a workgroup whose effective config can place query results. + + ``call_op(operation, params)`` is the read-only call bridge bound to an + Athena client. When ``requested`` is given it is returned untouched (the + caller's explicit choice always wins). When it is ``None`` we read + ``primary`` first and, only if it has neither an output location nor managed + results, enumerate other enabled workgroups (bounded) and pick the first that + does -- reporting the choice in ``guidance``. + + The value is ``{"workgroup": name, "auto_resolved": bool, + "output_location": loc, "managed_results": bool}``. + """ + if requested: + # Honour an explicit workgroup verbatim; its output location (or an + # enforced one) is resolved later when the query runs. + return ok_result( + { + "workgroup": requested, + "auto_resolved": False, + "output_location": None, + "managed_results": False, + } + ) + + guidance: list[str] = [] + primary = call_op("GetWorkGroup", {"WorkGroup": "primary"}) + if primary["ok"]: + cfg = ((primary.get("value") or {}).get("WorkGroup") or {}).get( + "Configuration" + ) or {} + loc, managed = _workgroup_output(cfg) + if loc or managed: + return ok_result( + { + "workgroup": "primary", + "auto_resolved": False, + "output_location": loc, + "managed_results": managed, + } + ) + + # 'primary' cannot place results (or could not be read); look for a workgroup + # that can. Enumerate enabled workgroups and read each effective config. + token: str | None = None + for _ in range(_MAX_WORKGROUP_PAGES): + params: dict[str, Any] = {} + if token: + params["NextToken"] = token + listing = call_op("ListWorkGroups", params) + if not listing["ok"]: + # Fall back to 'primary'; the query step will surface a precise + # "no output location" error if it truly has none. + guidance.append( + "Could not list workgroups to auto-resolve an output location; " + "falling back to 'primary'. Pass --workgroup or --output-location." + ) + return ok_result( + { + "workgroup": "primary", + "auto_resolved": False, + "output_location": None, + "managed_results": False, + }, + guidance=guidance, + ) + summaries = (listing.get("value") or {}).get("WorkGroups") or [] + for summ in summaries: + name = summ.get("Name") + if not name or summ.get("State") == "DISABLED": + continue + wg = call_op("GetWorkGroup", {"WorkGroup": name}) + if not wg["ok"]: + continue + cfg = ((wg.get("value") or {}).get("WorkGroup") or {}).get( + "Configuration" + ) or {} + loc, managed = _workgroup_output(cfg) + if loc or managed: + guidance.append( + f"No --workgroup given and 'primary' has no result output " + f"location; auto-selected workgroup {name!r} " + f"({'managed results' if managed and not loc else loc})." + ) + return ok_result( + { + "workgroup": name, + "auto_resolved": True, + "output_location": loc, + "managed_results": managed, + }, + guidance=guidance, + ) + token = (listing.get("value") or {}).get("NextToken") + if not token: + break + + # Nothing better than 'primary' was found; return it and let the query step + # emit the precise "no output location" error if applicable. + guidance.append( + "No workgroup with a result output location was found; using 'primary'. " + "Pass --output-location='s3:////' to run the query." + ) + return ok_result( + { + "workgroup": "primary", + "auto_resolved": False, + "output_location": None, + "managed_results": False, + }, + guidance=guidance, + ) + + +def _validate_params( + sql: Any, + workgroup: Any, + database: Any, + catalog: Any, + output_location: Any, + max_rows: Any, + poll_interval: Any, + max_wait: Any, +) -> dict[str, Any] | None: + # sql is required and must never be empty; catch it client-side so an empty + # statement is a ParamValidation error (exit 252) rather than slipping + # through to a SERVICE_ERROR from Athena (exit 255). + err = check_non_empty_str("sql", sql, required=True) + if err is not None: + return invalid_input( + err, + ["Pass --sql with a non-empty SQL statement."], + ) + # workgroup may be None here (auto-resolved later); only reject an explicit + # empty string. + err = check_non_empty_str("workgroup", workgroup, required=False) + if err is not None: + return invalid_input( + err, + ["Pass --workgroup or omit it to auto-resolve one."], + ) + err = check_non_empty_str("catalog", catalog, required=True) + if err is not None: + return invalid_input( + err, + ["Pass --catalog AwsDataCatalog unless the task names another."], + ) + err = check_non_empty_str("database", database, required=False) + if err is not None: + return invalid_input( + err, + ["Omit --database or pass the Glue database name."], + ) + if output_location is not None and ( + not isinstance(output_location, str) + or not output_location.startswith("s3://") + or len(output_location) <= 5 + ): + return invalid_input( + f"output_location {output_location!r} is not an s3:// URI.", + [ + "Pass --output-location s3://// or omit it to use " + "the workgroup's location." + ], + ) + err = check_positive_int("max_rows", max_rows) + if err is not None: + return invalid_input( + err, + ["Pass e.g. --max-rows 10000."], + ) + for name, val in ( + ("poll_interval", poll_interval), + ("max_wait", max_wait), + ): + err = check_positive_number(name, val) + if err is not None: + return invalid_input( + err, + [f"Pass a positive {name}."], + ) + return None + + +def _read_workgroup( + call_op, workgroup: str +) -> tuple[dict[str, Any] | None, dict[str, Any]]: + info: dict[str, Any] = { + "read": False, + "output_location": None, + "enforced": False, + "managed_results": False, + "engine_version": None, + "state": None, + "note": None, + } + res = call_op("GetWorkGroup", {"WorkGroup": workgroup}) + if res["ok"]: + wg = (res.get("value") or {}).get("WorkGroup") or {} + cfg = wg.get("Configuration") or {} + info["read"] = True + info["state"] = wg.get("State") + info["output_location"] = (cfg.get("ResultConfiguration") or {}).get( + "OutputLocation" + ) + info["enforced"] = bool(cfg.get("EnforceWorkGroupConfiguration")) + info["managed_results"] = bool( + (cfg.get("ManagedQueryResultsConfiguration") or {}).get("Enabled") + ) + ev = cfg.get("EngineVersion") or {} + info["engine_version"] = ev.get("EffectiveEngineVersion") or ev.get( + "SelectedEngineVersion" + ) + if info["state"] == "DISABLED": + return ( + err_result( + NOT_READY, + f"Workgroup '{workgroup}' is DISABLED; queries cannot run in it.", + key=workgroup, + guidance=[ + "Use another workgroup (athena:ListWorkGroups) or ask an admin to enable it." + ], + ), + info, + ) + return None, info + err = res.get("error") or "" + low = err.lower() + if res.get("code") == NOT_FOUND or ( + "not found" in low and "workgroup" in low + ): + return ( + err_result( + NOT_FOUND, + f"Workgroup '{workgroup}' not found: {err}", + aws_code=res.get("aws_code"), + key=workgroup, + guidance=[ + "List workgroups with athena:ListWorkGroups, or use the workgroup named by the task." + ], + ), + info, + ) + if res.get("code") == CREDENTIALS_EXPIRED: + return ( + err_result( + CREDENTIALS_EXPIRED, + err, + aws_code=res.get("aws_code"), + retriable=True, + guidance=[ + "Credentials may have expired; refresh them and retry." + ], + ), + info, + ) + info["note"] = ( + f"Could not read workgroup '{workgroup}' ({res.get('code')}: {err}); proceeding without its " + "enforced output location / engine version." + ) + return None, info + + +def _coerce(value: str | None, col_type: str | None) -> Any: + if value is None: + return None + t = (col_type or "").lower() + if t in _INTEGER_TYPES: + try: + return int(value) + except ValueError: + return value + if t == "boolean" and value.lower() in ("true", "false"): + return value.lower() == "true" + return value + + +def _row_keys(columns: list[dict[str, Any]]) -> tuple[list[str], bool]: + keys: list[str] = [] + seen: dict[str, int] = {} + renamed = False + for idx, col in enumerate(columns): + base = col.get("label") or col.get("name") or f"_col{idx}" + count = seen.get(base, 0) + 1 + seen[base] = count + if count > 1: + renamed = True + keys.append(f"{base}_{count}") + else: + keys.append(base) + return keys, renamed + + +def _stats(qe: dict[str, Any]) -> dict[str, Any]: + st = qe.get("Statistics") or {} + return { + "data_scanned_bytes": st.get("DataScannedInBytes"), + "engine_execution_time_ms": st.get("EngineExecutionTimeInMillis"), + "total_execution_time_ms": st.get("TotalExecutionTimeInMillis"), + "queue_time_ms": st.get("QueryQueueTimeInMillis"), + } + + +def _failure_code(qe: dict[str, Any], reason: str) -> tuple[str, bool]: + """Best-effort mapping of an Athena query failure to an error code. + + Athena reports a query-level failure as a status of ``FAILED`` plus a + free-text ``StateChangeReason``; the API exposes no structured, + machine-readable error codes for query-level failures. The checks below + therefore string-match that reason (and the optional ``AthenaError`` + category, when present) to hand the caller an actionable error code. This + matching is deliberately best-effort: if Athena changes its error-message + wording, an unmatched reason falls back to ``SERVICE_ERROR``. + + Returns a ``(code, retriable)`` tuple. + """ + low = reason.lower() + if _NOT_SUPPORTED_MSG in low: + return INVALID_INPUT, False + if ( + "access denied" in low + or "not authorized" in low + or "accessdenied" in low + ): + return PERMISSION_DENIED, False + athena_error = (qe.get("Status") or {}).get("AthenaError") or {} + category = athena_error.get("ErrorCategory") + if category == 2: # USER error: bad SQL, missing column/table + return INVALID_INPUT, False + if category == 1: # SYSTEM error: retriable + return SERVICE_ERROR, True + return SERVICE_ERROR, False + + +class _ResolvedLocation(NamedTuple): + """Where a query's results are written, plus how that was decided. + + ``error`` is a pre-dispatch failure result (or ``None`` to proceed); + ``location`` is the effective S3 output location (``None`` for managed + results, or when the workgroup places results itself); ``source`` says where + the location came from ('caller', 'workgroup_enforced', 'workgroup_default', + 'workgroup_managed', 'unknown'); ``send_location`` is True only when the + caller's ``ResultConfiguration`` override should be sent on + StartQueryExecution. The field order matches the historical 4-tuple, so + index-unpacking call sites keep working unchanged. + """ + + error: dict[str, Any] | None + location: str | None + source: str + send_location: bool + + +def _resolve_location( + wg: dict[str, Any], + workgroup: str, + output_location: str | None, + notes: list[str], +) -> _ResolvedLocation: + """Resolve the effective result location (see :class:`_ResolvedLocation`).""" + if wg["enforced"]: + loc = wg["output_location"] + source = ( + "workgroup_managed" + if (wg["managed_results"] and not loc) + else "workgroup_enforced" + ) + if output_location and output_location != loc: + notes.append( + f"Workgroup '{workgroup}' enforces its configuration; output_location " + f"{output_location!r} was ignored and results go to {loc or 'managed query results'}." + ) + return _ResolvedLocation(None, loc, source, False) + if output_location: + return _ResolvedLocation(None, output_location, "caller", True) + if wg["output_location"]: + return _ResolvedLocation( + None, wg["output_location"], "workgroup_default", False + ) + if wg["managed_results"]: + return _ResolvedLocation(None, None, "workgroup_managed", False) + if wg["read"]: + return _ResolvedLocation( + invalid_input( + f"Workgroup '{workgroup}' has no result output location and none was passed.", + [ + "Pass --output-location s3://// (a bucket the role can write to)." + ], + value={"workgroup": workgroup}, + ), + None, + "", + False, + ) + return _ResolvedLocation(None, None, "unknown", False) + + +def _start_failure(start: dict[str, Any]) -> dict[str, Any]: + err = start.get("error") or "" + if _NOT_SUPPORTED_MSG in err.lower(): + return err_result( + INVALID_INPUT, + f"Athena rejected the statement: {err}. Usual causes: more than one statement, an " + "unsupported statement type (USE/SET/transactions/GRANT), or DDL for a table " + "format/catalog that does not support it.", + aws_code=start.get("aws_code"), + guidance=[ + "Send exactly one SELECT/WITH/SHOW/DESCRIBE statement per call.", + "Pass --database/--catalog instead of USE; drop SET/BEGIN/COMMIT.", + "For DDL, check the target catalog/table format supports the operation in this engine version.", + ], + ) + by_code = { + CREDENTIALS_EXPIRED: [ + "Credentials may have expired; refresh them and retry." + ], + PERMISSION_DENIED: [ + "The role lacks athena:StartQueryExecution on this workgroup, or the glue/s3 access it " + "needs; use the role named by the task or another workgroup." + ], + THROTTLED: [ + "Athena throttled StartQueryExecution; wait a few seconds and run again." + ], + } + code = start.get("code") or SERVICE_ERROR + guidance = by_code.get( + code, [f"StartQueryExecution failed ({code}); fix and retry."] + ) + return err_result( + code, + err, + aws_code=start.get("aws_code"), + retriable=bool(start.get("retriable")) + or code in (THROTTLED, CREDENTIALS_EXPIRED), + guidance=guidance, + ) + + +def run_athena_query( + call_op, + sql: str, + *, + workgroup: str | None = None, + database: str | None = None, + catalog: str = "AwsDataCatalog", + output_location: str | None = None, + max_rows: int = 10_000, + coerce_types: bool = True, + poll_interval: float = 1.0, + max_wait: float = 300.0, + cancel_on_timeout: bool = False, + sleep=None, + clock=None, +) -> dict[str, Any]: + """Run one Athena SQL statement to a terminal state and return all rows. + + ``call_op(operation, params)`` is the read-only call bridge bound to an + Athena client. Returns the plain-dict result convention; never raises. When + ``workgroup`` is ``None`` a workgroup with an output location is auto-resolved. + """ + sleep = time.sleep if sleep is None else sleep + clock = time.monotonic if clock is None else clock + try: + bad = _validate_params( + sql, + workgroup, + database, + catalog, + output_location, + max_rows, + poll_interval, + max_wait, + ) + if bad is not None: + return bad + # The SQL is submitted verbatim; run-query no longer inspects or + # restricts it (see the module docstring -- this command is mutative). + statement = sql + resolve_guidance: list[str] = [] + if not workgroup: + resolved = resolve_workgroup_with_output(call_op, None) + # resolve never fails hard; it always yields a workgroup to try. + workgroup = (resolved.get("value") or {}).get( + "workgroup", "primary" + ) + resolve_guidance = resolved.get("guidance") or [] + return _run( + call_op, + statement, + workgroup=workgroup, + database=database, + catalog=catalog, + output_location=output_location, + max_rows=max_rows, + coerce_types=coerce_types, + poll_interval=float(poll_interval), + max_wait=float(max_wait), + cancel_on_timeout=cancel_on_timeout, + sleep=sleep, + clock=clock, + extra_guidance=resolve_guidance, + ) + except Exception as exc: # noqa: BLE001 - the command contract is "never raises" + return err_result( + SERVICE_ERROR, + f"run-query internal error: {exc!r}", + guidance=[ + "Unexpected response shape; retry once, then fall back to athena:GetQueryExecution by hand." + ], + ) + + +def _start_query( + call_op, params: dict[str, Any] +) -> tuple[dict[str, Any] | None, str | None]: + """StartQueryExecution; return ``(error_result_or_None, query_execution_id)``. + + ``params`` is the fully-built StartQueryExecution request (QueryString, + WorkGroup, QueryExecutionContext and an optional ResultConfiguration). Never + raises: a failed start, or a start that returns no QueryExecutionId, comes + back as an error result in the first slot (and ``None`` id); a successful + start returns ``(None, qid)``. + """ + start = call_op("StartQueryExecution", params) + if not start["ok"]: + return _start_failure(start), None + qid = (start.get("value") or {}).get("QueryExecutionId") + if not qid: + return ( + err_result( + SERVICE_ERROR, + "StartQueryExecution returned no QueryExecutionId.", + guidance=[ + "Retry the call; if it persists, check the workgroup in the Athena console." + ], + ), + None, + ) + return None, qid + + +def _fetch_results( + call_op, + qid: str, + *, + max_rows: int, + coerce_types: bool, + expect_header: bool, + base: dict[str, Any], + effective_location: str | None, + notes: list[str], +) -> tuple[dict[str, Any] | None, dict[str, Any]]: + """Paginate GetQueryResults for ``qid``; return ``(error_or_None, data)``. + + Follows NextToken, drops the header row, coerces integer/boolean cells (when + ``coerce_types``) and stops at ``max_rows``. On a page failure it returns an + ``err_result`` (carrying the rows already gathered and a resume token) as the + first element and an empty ``data``; on success the first element is ``None`` + and ``data`` holds the decoded ``columns``/``rows``/``header_row_dropped``/ + ``pages_fetched``/``is_truncated``/``next_token``/``update_count`` fields the + orchestrator folds into the final result. ``notes`` is appended to in place + (duplicate-column and repeated-token notes). + """ + columns: list[dict[str, Any]] | None = None + keys: list[str] = [] + rows: list[dict[str, Any]] = [] + header_dropped = False + pages = 0 + token: str | None = None + update_count: int | None = None + seen_tokens: set[str] = set() + while True: + remaining = max_rows - len(rows) + page_size = min( + _MAX_PAGE_SIZE, + remaining + (1 if pages == 0 and expect_header else 0), + ) + rparams: dict[str, Any] = { + "QueryExecutionId": qid, + "MaxResults": page_size, + } + if token: + rparams["NextToken"] = token + res = call_op("GetQueryResults", rparams) + if not res["ok"]: + err = err_result( + res.get("code") or SERVICE_ERROR, + f"GetQueryResults failed after {pages} page(s) for query {qid}: {res.get('error')}", + aws_code=res.get("aws_code"), + retriable=bool(res.get("retriable")), + key=qid, + is_exhaustive=False, + value={ + **base, + "columns": columns or [], + "rows": rows, + "row_count": len(rows), + "pages_fetched": pages, + "next_token": token, + "is_truncated": True, + }, + guidance=[ + f"Resume with athena:GetQueryResults QueryExecutionId={qid} NextToken={token!r}, or " + f"read the CSV at {effective_location}." + ], + ) + return err, {} + pages += 1 + body = res.get("value") or {} + if update_count is None and body.get("UpdateCount") is not None: + update_count = body.get("UpdateCount") + rs = body.get("ResultSet") or {} + if columns is None: + columns = [ + { + "name": c.get("Name"), + "label": c.get("Label") or c.get("Name"), + "type": c.get("Type"), + "nullable": c.get("Nullable"), + "precision": c.get("Precision"), + "scale": c.get("Scale"), + } + for c in (rs.get("ResultSetMetadata") or {}).get("ColumnInfo") + or [] + ] + keys, renamed = _row_keys(columns) + if renamed: + notes.append( + "Duplicate column labels were suffixed (_2, _3, ...) in row keys; alias them in SQL." + ) + raw_rows = rs.get("Rows") or [] + if pages == 1 and raw_rows: + first = [ + d.get("VarCharValue") for d in raw_rows[0].get("Data") or [] + ] + labels = [c["label"] for c in columns] + names = [c["name"] for c in columns] + if expect_header or ( + columns and (first == labels or first == names) + ): + raw_rows = raw_rows[1:] + header_dropped = True + for raw in raw_rows: + if len(rows) >= max_rows: + break + cells = [d.get("VarCharValue") for d in raw.get("Data") or []] + cells += [None] * (len(keys) - len(cells)) + rows.append( + { + k: (_coerce(v, columns[i]["type"]) if coerce_types else v) + for i, (k, v) in enumerate(zip(keys, cells)) + } + ) + token = body.get("NextToken") + if not token: + break + if token in seen_tokens: + notes.append( + "GetQueryResults repeated a NextToken; pagination stopped to avoid a loop." + ) + token = None + break + seen_tokens.add(token) + if len(rows) >= max_rows: + break + + truncated = bool(token) + return None, { + "columns": columns or [], + "rows": rows, + "header_row_dropped": header_dropped, + "pages_fetched": pages, + "is_truncated": truncated, + "next_token": token, + "update_count": update_count, + } + + +def _run( + call_op, + statement: str, + *, + workgroup: str, + database: str | None, + catalog: str, + output_location: str | None, + max_rows: int, + coerce_types: bool, + poll_interval: float, + max_wait: float, + cancel_on_timeout: bool, + sleep, + clock, + extra_guidance: list[str], +) -> dict[str, Any]: + notes: list[str] = list(extra_guidance) + + # 1. Workgroup. + wg_err, wg = _read_workgroup(call_op, workgroup) + if wg_err is not None: + return wg_err + if wg["note"]: + notes.append(wg["note"]) + resolved = _resolve_location(wg, workgroup, output_location, notes) + if resolved.error is not None: + return resolved.error + effective_location = resolved.location + source = resolved.source + send_location = resolved.send_location + + # 2. Start. + ctx: dict[str, str] = {"Catalog": catalog} + if database: + ctx["Database"] = database + params: dict[str, Any] = { + "QueryString": statement, + "WorkGroup": workgroup, + "QueryExecutionContext": ctx, + } + if send_location and effective_location: + params["ResultConfiguration"] = {"OutputLocation": effective_location} + start_err, qid = _start_query(call_op, params) + if start_err is not None: + return start_err + + base: dict[str, Any] = { + "query_execution_id": qid, + "workgroup": workgroup, + "database": database, + "catalog": catalog, + "engine_version": wg["engine_version"], + "output_location": effective_location, + "output_location_source": source, + "submitted_sql": statement, + } + + # 3. Poll to a terminal state using the shared bounded poller. + qe: dict[str, Any] = {} + state = "QUEUED" + + def _terminal(res: dict[str, Any]) -> str | None: + nonlocal qe, state + qe = (res.get("value") or {}).get("QueryExecution") or {} + # State persists across polls (defaults to the last observed value). + state = (qe.get("Status") or {}).get("State") or state + return state if state in _TERMINAL_STATES else None + + def _is_retriable(res: dict[str, Any]) -> bool: + return bool(res.get("retriable")) or res.get("code") in ( + THROTTLED, + NETWORK_ERROR, + ) + + cancelled = False + + def _on_timeout() -> None: + nonlocal cancelled + if cancel_on_timeout: + # StopQueryExecution cancels a running query (read-only w.r.t. + # your data; it does not create or modify resources). + stop = call_op("StopQueryExecution", {"QueryExecutionId": qid}) + cancelled = stop["ok"] + + poll_result = drive_poll( + call_op, + "GetQueryExecution", + {"QueryExecutionId": qid}, + terminal=_terminal, + is_retriable=_is_retriable, + on_timeout=_on_timeout, + interval=poll_interval, + max_wait=max_wait, + max_polls=_MAX_POLLS, + max_transient=_MAX_TRANSIENT_POLL_ERRORS, + clock=clock, + sleep=sleep, + key=qid, + ) + polls = poll_result.get("polls_count", 0) + + if not poll_result["ok"]: + if poll_result.get("code") == TIMEOUT: + if cancelled: + guidance = [ + f"Query {qid} was cancelled after {max_wait:g}s (--cancel-on-timeout); narrow the " + "query (partition/time predicates) and run again." + ] + else: + guidance = [ + f"Query {qid} is still {state}; call athena:GetQueryExecution later or re-run with a " + "larger --max-wait. Do not re-submit the same SQL (it would be billed twice)." + ] + return err_result( + TIMEOUT, + f"Query {qid} did not finish within {max_wait:g}s (last state {state}).", + retriable=not cancelled, + key=qid, + value={ + **base, + "state": "CANCELLED" if cancelled else state, + "polls": polls, + "cancelled": cancelled, + **_stats(qe), + }, + guidance=guidance, + ) + return err_result( + poll_result.get("code") or SERVICE_ERROR, + poll_result.get("error") or "GetQueryExecution failed.", + aws_code=poll_result.get("aws_code"), + retriable=bool(poll_result.get("retriable")), + key=qid, + value={**base, "state": state, "polls": polls}, + guidance=[ + f"Query {qid} was started; poll athena:GetQueryExecution for it later instead of " + "starting a new one." + ], + ) + + stmt_type = qe.get("StatementType") + effective_location = (qe.get("ResultConfiguration") or {}).get( + "OutputLocation" + ) or effective_location + base["output_location"] = effective_location + base["engine_version"] = (qe.get("EngineVersion") or {}).get( + "EffectiveEngineVersion" + ) or base["engine_version"] + base.update( + { + "state": state, + "statement_type": stmt_type, + "polls": polls, + **_stats(qe), + } + ) + + if state == "FAILED": + reason = (qe.get("Status") or {}).get( + "StateChangeReason" + ) or "no StateChangeReason returned" + code, retriable = _failure_code(qe, reason) + return err_result( + code, + f"Athena query {qid} FAILED: {reason}", + retriable=retriable, + key=qid, + aws_code="QueryFailed", + value={**base, "state_change_reason": reason}, + guidance=[ + f"Fix the cause named in StateChangeReason: {reason}", + "Use DESCRIBE / SHOW COLUMNS to check column and table names before re-running.", + ], + ) + if state == "CANCELLED": + reason = (qe.get("Status") or {}).get( + "StateChangeReason" + ) or "cancelled" + return err_result( + CONFLICT, + f"Athena query {qid} was CANCELLED: {reason}", + key=qid, + value={**base, "state_change_reason": reason}, + guidance=[ + "The query was cancelled (by a user or the workgroup's BytesScannedCutoffPerQuery); " + "narrow it with partition predicates and run again." + ], + ) + + # 4. Paginate results. + fetch_err, data = _fetch_results( + call_op, + qid, + max_rows=max_rows, + coerce_types=coerce_types, + expect_header=stmt_type == "DML", + base=base, + effective_location=effective_location, + notes=notes, + ) + if fetch_err is not None: + return fetch_err + + truncated = data["is_truncated"] + token = data["next_token"] + guidance = [_VERBATIM_GUIDANCE, _NULL_GUIDANCE, *notes] + if truncated: + guidance.insert( + 0, + ( + f"Result truncated at max_rows={max_rows}; more rows exist. Re-run with a larger " + f"--max-rows, add a tighter WHERE, or resume with athena:GetQueryResults " + f"QueryExecutionId={qid} NextToken={token!r}." + ), + ) + value = { + **base, + "columns": data["columns"], + "rows": data["rows"], + "row_count": len(data["rows"]), + "header_row_dropped": data["header_row_dropped"], + "pages_fetched": data["pages_fetched"], + "is_truncated": truncated, + "next_token": token if truncated else None, + "max_rows": max_rows, + "update_count": data["update_count"], + "null_representation": "SQL NULL -> JSON null (render as NULL)", + } + return ok_result( + value, key=qid, is_exhaustive=not truncated, guidance=guidance + ) + + +class RunQueryCommand(BasicCommand): + NAME = "run-query" + DESCRIPTION = ( + "Run a single Athena statement to completion and return every row in one " + "call.\n\n" + "Use this command when you need to run an Athena SQL query and get the " + "results in one step, without managing the poll-and-paginate lifecycle " + "yourself. It handles query submission, polling for completion, result " + "pagination, and output formatting.\n\n" + "Submits the query, polls it to a terminal state, and paginates all " + "result pages. The SQL is passed through verbatim and is not restricted, " + "so this command is mutative (read-write): writes and DDL (INSERT, " + "CREATE, DROP, MERGE, UNLOAD, ...) run against real tables. When no " + "--workgroup is given, a workgroup with a result output location is " + "auto-resolved." + ) + EXAMPLES = ( + "To run a query in the default workgroup::\n\n" + " aws athena run-query --sql \"SELECT * FROM mydb.mytable LIMIT 100\"\n\n" + "To target a workgroup and database with a row cap::\n\n" + " aws athena run-query --sql \"SELECT count(*) FROM logs\" " + "--workgroup analytics --database mydb --max-rows 1\n" + ) + ARG_TABLE = [ + { + "name": "sql", + "required": True, + "help_text": "The SQL statement to run (submitted verbatim; may be a write/DDL statement).", + "cli_type_name": "string", + }, + { + "name": "workgroup", + "help_text": "Athena workgroup. Omit to auto-resolve one with an output location.", + "cli_type_name": "string", + }, + { + "name": "database", + "help_text": "Glue database for unqualified names.", + "cli_type_name": "string", + }, + { + "name": "catalog", + "default": "AwsDataCatalog", + "help_text": "Data catalog (default AwsDataCatalog).", + "cli_type_name": "string", + }, + { + "name": "output-location", + "help_text": "s3:// URI for results; overrides the workgroup default unless enforced.", + "cli_type_name": "string", + }, + { + "name": "max-rows", + "default": 10000, + "cli_type_name": "integer", + "help_text": "Maximum rows to return (default 10000).", + }, + { + "name": "poll-interval", + "default": 1.0, + "cli_type_name": "double", + "help_text": "Initial seconds between GetQueryExecution polls (default 1.0).", + }, + { + "name": "max-wait", + "default": 300.0, + "cli_type_name": "double", + "help_text": "Maximum seconds to wait for the query (default 300).", + }, + { + "name": "cancel-on-timeout", + "action": "store_true", + "default": False, + "help_text": "Cancel the query (StopQueryExecution) if --max-wait is exceeded.", + }, + { + "name": "no-coerce-types", + "action": "store_true", + "default": False, + "help_text": "Return all cells as strings instead of coercing integer/boolean columns.", + }, + ] + + def __init__( + self, + session, + client=None, + stdout=None, + stderr=None, + sleep=None, + clock=None, + ): + super().__init__(session) + self._client = client + self._stdout = stdout + self._stderr = stderr + self._sleep = sleep + self._clock = clock + + def _run_main(self, parsed_args, parsed_globals): + client = self._client or create_read_client( + self._session, "athena", parsed_globals + ) + + def call_op(operation, params): + return call(client, operation, params) + + result = run_athena_query( + call_op, + parsed_args.sql, + workgroup=parsed_args.workgroup, + database=parsed_args.database, + catalog=parsed_args.catalog, + output_location=parsed_args.output_location, + max_rows=parsed_args.max_rows, + coerce_types=not parsed_args.no_coerce_types, + poll_interval=parsed_args.poll_interval, + max_wait=parsed_args.max_wait, + cancel_on_timeout=parsed_args.cancel_on_timeout, + sleep=self._sleep, + clock=self._clock, + ) + return emit_result( + self.NAME, + result, + parsed_globals, + self._session, + self._stdout, + self._stderr, + ) diff --git a/awscli/customizations/compositereadonly/__init__.py b/awscli/customizations/compositereadonly/__init__.py new file mode 100644 index 000000000000..542b63cb025a --- /dev/null +++ b/awscli/customizations/compositereadonly/__init__.py @@ -0,0 +1,33 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"). You +# may not use this file except in compliance with the License. A copy of +# the License is located at +# +# http://aws.amazon.com/apache2.0/ +# +# or in the "license" file accompanying this file. This file is +# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF +# ANY KIND, either express or implied. See the License for the specific +# language governing permissions and limitations under the License. +"""Shared mechanism for read-only, multi-call aggregating CLI commands. + +Some services expose an operational question that no single API answers: "run +this query to completion and give me every row", "investigate this alarm", "tell +me why this service is unhealthy". Answering it means several read-only API +calls, bounded polling, and one summarized result. This package provides that +common machinery so each such command stays a thin, service-scoped verb: + +* :mod:`helpers` -- a read-only botocore ``call`` wrapper (errors as data, refuses + any mutating operation), plain-dict result factories, and a JSON emitter that + maps a result to a CLI exit code. +* :mod:`poll` -- :class:`~.poll.BoundedPoller`, a monotonic, bounded poll clock + with capped backoff for start-query/poll lifecycles that have no botocore waiter. +* :mod:`output` -- output caps (field truncation, series capping) that record + every clamp so trimming is visible, never silent. + +Nothing here is service-specific; the commands live under their owning service's +customization namespace and build on these primitives. The service-specific name +resolvers that used to live here now live with their owning commands (Athena's +workgroup resolver and CloudWatch's alarm resolver). +""" diff --git a/awscli/customizations/compositereadonly/helpers.py b/awscli/customizations/compositereadonly/helpers.py new file mode 100644 index 000000000000..f8087d611688 --- /dev/null +++ b/awscli/customizations/compositereadonly/helpers.py @@ -0,0 +1,546 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"). You +# may not use this file except in compliance with the License. A copy of +# the License is located at +# +# http://aws.amazon.com/apache2.0/ +# +# or in the "license" file accompanying this file. This file is +# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF +# ANY KIND, either express or implied. See the License for the specific +# language governing permissions and limitations under the License. +"""Read-only call bridge and output emitter for aggregating read commands. + +These commands make several read-only API calls and summarize them into one +result. The call bridge returns a result object from every call and never +raises; the CLI, by contrast, returns an exit code and prints model-shaped +JSON. This module bridges the two: + +* :func:`ok_result` / :func:`err_result` build a plain ``dict`` with the + control-flow keys the command logic branches on (``ok``, ``value``, ``code``, + ...). Only the ``value`` payload (plus ``guidance`` / ``is_exhaustive`` as plain + keys) is ever emitted, so the output stays normal CLI JSON that ``--query`` and + ``--output table/text`` understand -- it is not wrapped in a result envelope. +* :func:`call` runs one botocore read, normalizes a ``ClientError`` into the same + dict convention, and REFUSES any mutating operation so the commands cannot + issue a write even by mistake. +* :func:`emit_result` serializes the result and maps it to a CLI exit code. + +Error codes are plain strings (``"NOT_FOUND"`` etc.) that the command logic +branches on directly. +""" + +from __future__ import annotations + +import sys +from typing import Any, Callable, TypedDict + +from botocore import xform_name +from botocore.exceptions import ( + BotoCoreError, + ClientError, + EndpointConnectionError, + NoCredentialsError, +) +from botocore.exceptions import ( + ConnectionError as BotoConnectionError, +) + +from awscli.customizations.exceptions import ParamValidationError +from awscli.customizations.utils import create_client_from_parsed_globals +from awscli.formatter import get_formatter + +# --- CLI exit codes (see awscli/constants.py) -------------------------------- +RC_SUCCESS = 0 +RC_OPERATIONAL_ERROR = 255 +# ParamValidationError is raised for bad input; the clidriver maps it to rc 252. + +# --- Normalized error codes -------------------------------------------------- +NOT_FOUND = "NOT_FOUND" +NOT_READY = "NOT_READY" +THROTTLED = "THROTTLED" +PERMISSION_DENIED = "PERMISSION_DENIED" +INVALID_INPUT = "INVALID_INPUT" +TIMEOUT = "TIMEOUT" +CONFLICT = "CONFLICT" +QUOTA_EXCEEDED = "QUOTA_EXCEEDED" +SERVICE_ERROR = "SERVICE_ERROR" +NETWORK_ERROR = "NETWORK_ERROR" +VALIDATION_ERROR = "VALIDATION_ERROR" +CREDENTIALS_EXPIRED = "CREDENTIALS_EXPIRED" + +# Botocore error-code -> normalized code for the services these commands touch; +# anything unmapped falls back by HTTP status. +_CLIENT_ERROR_MAP = { + "AccessDenied": PERMISSION_DENIED, + "AccessDeniedException": PERMISSION_DENIED, + "UnauthorizedException": PERMISSION_DENIED, + "UnauthorizedOperation": PERMISSION_DENIED, + "NotAuthorized": PERMISSION_DENIED, + "ResourceNotFoundException": NOT_FOUND, + "ResourceNotFound": NOT_FOUND, + "NoSuchEntity": NOT_FOUND, + "WorkGroupNotFound": NOT_FOUND, + "DashboardNotFoundError": NOT_FOUND, + "ThrottlingException": THROTTLED, + "Throttling": THROTTLED, + "TooManyRequestsException": THROTTLED, + "RequestLimitExceeded": THROTTLED, + "LimitExceededException": QUOTA_EXCEEDED, + "ServiceQuotaExceededException": QUOTA_EXCEEDED, + "ServiceUnavailableException": SERVICE_ERROR, + "InternalServerError": SERVICE_ERROR, + "InternalFailure": SERVICE_ERROR, + "ExpiredTokenException": CREDENTIALS_EXPIRED, + "ExpiredToken": CREDENTIALS_EXPIRED, + "RequestExpired": CREDENTIALS_EXPIRED, + "InvalidParameterException": INVALID_INPUT, + "ValidationException": VALIDATION_ERROR, + "ValidationError": VALIDATION_ERROR, + "MalformedQueryException": INVALID_INPUT, + "InvalidRequestException": INVALID_INPUT, +} +_RETRIABLE = frozenset( + {THROTTLED, QUOTA_EXCEEDED, NETWORK_ERROR, CREDENTIALS_EXPIRED} +) + +# Verb prefixes that denote a mutating API. The read-only lifecycle verbs these +# commands DO use (``Start``/``Stop`` for Athena/Logs Insights queries, +# ``Describe``/``Get``/``List``) are deliberately absent, so ``call`` lets them +# through while refusing anything that would create, change or delete state. +_MUTATING_PREFIXES = ( + "Create", + "Delete", + "Put", + "Update", + "Modify", + "Terminate", + "Remove", + "Attach", + "Detach", + "Associate", + "Disassociate", + "Register", + "Deregister", + "Enable", + "Disable", + "Reboot", + "Run", + "Set", + "Cancel", + "Add", + "Revoke", + "Authorize", + "Replace", + "Reset", + "Import", + "Tag", + "Untag", + "Copy", + "Send", + "Publish", + "Invoke", + "Execute", + "Deploy", + "Restore", + "Purchase", + "Accept", + "Reject", + "Allocate", + "Release", + "Assign", + "Unassign", + "Grant", + "Move", + "Apply", + "Batch", + "Change", + "Confirm", + "Initiate", + "Provision", +) +# ``Batch`` and ``Change`` are intentionally conservative: they block mutating +# verbs such as ``BatchDeleteTable`` but would also block a read like +# ``BatchGetItem``. Callers needing a Batch*/Change* read must allowlist it. + + +class ReadOnlyViolation(AssertionError): + """Raised when a command attempts a mutating AWS operation. + + This is a programming-error guard, not user-facing: every command built on + this package is read-only by construction, and the test suite asserts it is + never raised. + """ + + +def assert_read_only(operation: str) -> None: + """Refuse any operation whose name begins with a mutating verb. + + This is a defense-in-depth guard applied at the call-bridge level. Every + high-order command in this package is designed to call read-only APIs + (``Describe*``/``Get*``/``List*``) only; this check ensures that even if a + command has a bug and accidentally issues a write operation, the call is + blocked here before it ever reaches AWS. + + The mutating-prefix list intentionally errs on the side of caution: it + blocks verbs such as ``Batch*`` and ``Change*`` even though some operations + under them are actually reads (for example ``BatchGetItem``). A caller that + legitimately needs a ``Batch*``/``Change*`` read must explicitly allowlist + it rather than have this guard relaxed for everyone. + """ + for prefix in _MUTATING_PREFIXES: + # Require the verb to be followed by an upper-case word boundary (or the + # end of the name) so a read like "Settings"/"Getter" is not mistaken + # for a "Set"/"Get" mutation. + if operation.startswith(prefix) and ( + len(operation) == len(prefix) or operation[len(prefix)].isupper() + ): + raise ReadOnlyViolation( + f"Refusing to call mutating operation {operation!r}; this " + "command is strictly read-only." + ) + + +class _ResultEnvelopeBase(TypedDict): + """Required keys present on every result, success or failure.""" + + ok: bool + key: str | None + + +class ResultEnvelope(_ResultEnvelopeBase, total=False): + """The plain-``dict`` control-flow envelope built by the result helpers. + + Produced by :func:`ok_result` and :func:`err_result` (and extended by + :func:`invalid_input`). ``ok`` and ``key`` form the minimum contract and are + always present; the remaining keys are populated depending on outcome -- the + ``value``/``is_exhaustive``/``guidance`` payload on success, and the + ``error``/``code``/``aws_code``/``retriable`` diagnostics on failure. This is + a typing aid only: the helpers still return ordinary ``dict`` instances, so + the type is structurally dict-compatible and existing call sites are + unaffected. + """ + + # Success payload (also carried through on failures as data). + value: Any + is_exhaustive: bool | None + guidance: list[str] + # Failure diagnostics. + error: str + code: str + aws_code: str | None + retriable: bool + # Set by :func:`invalid_input` to request ParamValidationError (rc 252). + param_error: bool + + +def ok_result( + value: Any, + *, + is_exhaustive: bool | None = None, + guidance: list[str] | None = None, + key: str | None = None, +) -> ResultEnvelope: + """Build a success result dict.""" + return { + "ok": True, + "value": value, + "is_exhaustive": is_exhaustive, + "guidance": guidance or [], + "key": key, + } + + +def err_result( + code: str, + error: str, + *, + value: Any = None, + guidance: list[str] | None = None, + retriable: bool = False, + aws_code: str | None = None, + key: str | None = None, + is_exhaustive: bool | None = None, +) -> ResultEnvelope: + """Build a failure result dict (errors are data, never raised).""" + return { + "ok": False, + "value": value, + "error": error, + "code": code, + "aws_code": aws_code, + "retriable": retriable or code in _RETRIABLE, + "guidance": guidance or [], + "key": key, + "is_exhaustive": is_exhaustive, + } + + +# --- Input validation predicates --------------------------------------------- +# Each predicate returns a human-readable error description when the check fails, +# or ``None`` when it passes. Callers that fail-fast wrap the description in +# :func:`invalid_input`; callers that accumulate append the non-``None`` results +# to a problem list. Sharing one definition of "non-empty string", "positive +# int", etc. keeps the high-order commands from re-deriving the same predicates +# (and the bool-is-a-subclass-of-int trap) by hand in every validator. + + +def check_non_empty_str( + name: str, val: Any, *, required: bool = True +) -> str | None: + """Reject a blank or non-string value. + + ``None`` passes when ``required`` is False (the argument is optional); + otherwise it is reported as required. A non-string, an empty string, or a + whitespace-only string is always rejected. + """ + if val is None: + return None if not required else f"{name} is required." + if not isinstance(val, str): + return f"{name} must be a string, got {val!r}." + if val == "": + return f"{name} is an empty string." + if not val.strip(): + return f"{name} is empty." + return None + + +def check_positive_int( + name: str, val: Any, *, lo: int = 1, hi: int | None = None +) -> str | None: + """Reject a non-integer or an integer outside ``[lo, hi]``. + + ``bool`` is rejected explicitly: it is a subclass of ``int``, so ``True`` + would otherwise sneak through as ``1``. ``lo=0`` accepts non-negative + values; ``hi`` bounds the value from above when provided. + """ + if isinstance(val, bool) or not isinstance(val, int): + return f"{name} must be an integer, got {val!r}." + if hi is not None and not (lo <= val <= hi): + return f"{name} must be an integer between {lo} and {hi}, got {val!r}." + if val < lo: + if lo == 0: + descriptor = "a non-negative integer" + elif lo == 1: + descriptor = "a positive integer" + else: + descriptor = f"an integer >= {lo}" + return f"{name} must be {descriptor}, got {val!r}." + return None + + +def check_positive_number(name: str, val: Any) -> str | None: + """Reject anything that is not an ``int``/``float`` strictly greater than 0. + + ``bool`` is rejected explicitly, for the same reason as + :func:`check_positive_int`. + """ + if isinstance(val, bool) or not isinstance(val, (int, float)) or val <= 0: + return f"{name} must be a positive number, got {val!r}." + return None + + +def check_max_size( + name: str, + val: str, + limit: int, + *, + unit: str = "bytes", + measure: str = "utf-8", +) -> str | None: + """Reject a string whose size exceeds ``limit``. + + ``measure='chars'`` measures with :func:`len`; any other value is treated as + a text encoding and the size is the encoded byte length (UTF-8 by default). + ``unit`` only labels the error message. + """ + size = len(val) if measure == "chars" else len(val.encode(measure)) + if size > limit: + return f"{name} is {size} {unit}; the limit is {limit} {unit}." + return None + + +def invalid_input( + error: str, guidance: list[str] | None = None, *, value: Any = None +) -> dict[str, Any]: + """Shorthand for a pre-dispatch INVALID_INPUT failure result. + + Flagged ``param_error`` so :func:`emit_result` raises ``ParamValidationError`` + (exit code 252), matching how the CLI reports bad input. Server-discovered + invalids use :func:`err_result` directly and stay operational (exit code 255). + """ + result = err_result(INVALID_INPUT, error, guidance=guidance, value=value) + result["param_error"] = True + return result + + +def create_read_client(session, service_name, parsed_globals): + """Create a botocore client honouring --region/--endpoint-url/--verify-ssl. + + Reuses the standard CLI helper so credentials and endpoint resolution match + every other command; the read-only guarantee is enforced at the :func:`call` + boundary, not here. + """ + return create_client_from_parsed_globals( + session, service_name, parsed_globals + ) + + +def _map_client_error(exc: ClientError) -> tuple[str, str | None]: + aws_code = exc.response.get("Error", {}).get("Code") + if aws_code in _CLIENT_ERROR_MAP: + return _CLIENT_ERROR_MAP[aws_code], aws_code + status = exc.response.get("ResponseMetadata", {}).get("HTTPStatusCode") + if status == 404: + return NOT_FOUND, aws_code + if status in (401, 403): + return PERMISSION_DENIED, aws_code + if status == 429: + return THROTTLED, aws_code + if isinstance(status, int) and status >= 500: + return SERVICE_ERROR, aws_code + return SERVICE_ERROR, aws_code + + +def call(client, operation: str, params: dict[str, Any]) -> dict[str, Any]: + """Run ONE read-only botocore operation, returning the plain-dict convention. + + ``operation`` is the modeled PascalCase name (e.g. ``"GetQueryExecution"``); + it is translated to the botocore method name. Any mutating operation is + refused via :func:`assert_read_only`. Errors are returned as data, never + raised, so the caller can branch on ``result["code"]``. + """ + assert_read_only(operation) + method_name = xform_name(operation) + method = getattr(client, method_name) + try: + response = method(**params) + except ClientError as exc: + code, aws_code = _map_client_error(exc) + return err_result( + code, str(exc), aws_code=aws_code, retriable=code in _RETRIABLE + ) + except NoCredentialsError as exc: + return err_result(CREDENTIALS_EXPIRED, str(exc), retriable=True) + except (EndpointConnectionError, BotoConnectionError) as exc: + return err_result(NETWORK_ERROR, str(exc), retriable=True) + except BotoCoreError as exc: + return err_result(SERVICE_ERROR, str(exc)) + return ok_result(response) + + +def paginate( + call_op: Callable[[str, dict[str, Any]], dict[str, Any]], + operation: str, + params: dict[str, Any], + key: str, + max_pages: int, + *, + token_field: str = "NextToken", + item_filter: Callable[[Any], bool] | None = None, +) -> tuple[list[Any], bool, dict[str, Any] | None]: + """Follow a NextToken-style pager to the end or ``max_pages``. + + ``call_op(operation, params)`` is the read-only call bridge; ``key`` names the + response field holding each page's items and ``token_field`` the continuation + token field (``NextToken`` for most services, ``nextToken`` for the + lower-camel services such as ECS). ``item_filter``, when given, keeps only the + items it accepts (e.g. dict-shaped entries, or truthy ARNs); the continuation + token is never sent on the first call. + + Returns ``(items, is_exhaustive, last_error_result)``: + + * ``is_exhaustive`` is True only when a page without a continuation token was + reached, so every page was read. It is False when the bounded loop stopped + on ``max_pages`` with a token still outstanding, or when a call failed. + * ``last_error_result`` is the failing call's result dict when a call failed, + else ``None``; callers branch on it before trusting ``is_exhaustive``. + """ + items: list[Any] = [] + token: str | None = None + for _ in range(max_pages): + res = call_op( + operation, + {**params, token_field: token} if token else dict(params), + ) + if not res["ok"]: + return items, False, res + value = res.get("value") or {} + page = value.get(key) or [] + if item_filter is not None: + items.extend(i for i in page if item_filter(i)) + else: + items.extend(page) + token = value.get(token_field) + if not token: + return items, True, None + return items, False, None + + +def _render_payload(result: dict[str, Any]) -> dict[str, Any]: + """Project the user-facing JSON from a result dict. + + Emits the primary ``value`` payload with ``guidance``/``is_exhaustive`` + folded in as plain keys -- not a result envelope. On a failure the + error/code are included so the caller can branch on them. + """ + value = result.get("value") + if isinstance(value, dict): + payload = dict(value) + elif value is None: + payload = {} + else: + payload = {"result": value} + if result.get("is_exhaustive") is not None: + payload.setdefault("is_exhaustive", result["is_exhaustive"]) + if result.get("guidance"): + payload.setdefault("guidance", result["guidance"]) + if not result.get("ok"): + payload.setdefault("error", result.get("error")) + payload.setdefault("code", result.get("code")) + if result.get("retriable"): + payload.setdefault("retriable", True) + return payload + + +def emit_result( + name, + result: dict[str, Any], + parsed_globals, + session, + stdout=None, + stderr=None, +) -> int: + """Serialize ``result`` as CLI JSON and return the exit code. + + Success -> payload to stdout, rc 0. Operational failure -> payload (if any) + to stdout so ``--query`` still works, a concise ``code: error`` line to + stderr, rc 255. + """ + if stdout is None: + stdout = sys.stdout + if stderr is None: + stderr = sys.stderr + # Pre-dispatch bad input is reported the CLI-native way: ParamValidationError + # (exit code 252), no JSON payload. + if not result.get("ok") and result.get("param_error"): + raise ParamValidationError(result.get("error") or "Invalid input") + output = getattr(parsed_globals, "output", None) + if output is None: + output = session.get_config_variable("output") + formatter = get_formatter(output, parsed_globals) + payload = _render_payload(result) + formatter(name, payload, stream=stdout) + stdout.write("\n") + if result.get("ok"): + return RC_SUCCESS + stderr.write( + f"{result.get('code', SERVICE_ERROR)}: {result.get('error', '')}\n" + ) + return RC_OPERATIONAL_ERROR + + +def raise_invalid(message: str) -> None: + """Raise a ParamValidationError (rc 252) for pre-dispatch bad input.""" + raise ParamValidationError(message) diff --git a/awscli/customizations/compositereadonly/output.py b/awscli/customizations/compositereadonly/output.py new file mode 100644 index 000000000000..d241e0c58245 --- /dev/null +++ b/awscli/customizations/compositereadonly/output.py @@ -0,0 +1,56 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"). You +# may not use this file except in compliance with the License. A copy of +# the License is located at +# +# http://aws.amazon.com/apache2.0/ +# +# or in the "license" file accompanying this file. This file is +# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF +# ANY KIND, either express or implied. See the License for the specific +# language governing permissions and limitations under the License. +"""Central output caps for aggregating read commands. + +Output is bounded to keep responses manageable: row and series counts are capped +and oversized free-text fields (such as a Logs Insights ``@message``) are +truncated. Every clamp is recorded so the caller can see that the result was +trimmed rather than silently losing data. +""" + +from __future__ import annotations + +from typing import Any + +# Default width an individual free-text field is truncated to. Full log lines +# can be many kilobytes; 2000 chars keeps the head (where the error usually is) +# while preventing one row from dominating the response. +DEFAULT_MESSAGE_MAX_CHARS = 2000 +_TRUNCATION_SUFFIX = "...[truncated]" + + +def truncate_field(value: Any, max_chars: int) -> tuple[Any, bool]: + """Truncate a string to ``max_chars`` (keeping the head). + + Returns ``(value, truncated)``. Non-strings and ``max_chars <= 0`` pass + through untouched so callers can disable truncation. + + Note: when *max_chars* is smaller than the suffix length (14 chars), the + returned string may exceed *max_chars*. This is intentional; practical + defaults are 2000+. + """ + if max_chars <= 0 or not isinstance(value, str) or len(value) <= max_chars: + return value, False + head = value[: max(0, max_chars - len(_TRUNCATION_SUFFIX))] + return head + _TRUNCATION_SUFFIX, True + + +def cap_series(items: list[Any], max_count: int) -> tuple[list[Any], int]: + """Cap a list to ``max_count`` entries. + + Returns ``(kept, dropped)`` where ``dropped`` is how many were removed, so + the caller can report the clamp in ``clamps``/``guidance``. + """ + if max_count < 0 or len(items) <= max_count: + return items, 0 + return items[:max_count], len(items) - max_count diff --git a/awscli/customizations/compositereadonly/poll.py b/awscli/customizations/compositereadonly/poll.py new file mode 100644 index 000000000000..2a58c9000bd0 --- /dev/null +++ b/awscli/customizations/compositereadonly/poll.py @@ -0,0 +1,180 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"). You +# may not use this file except in compliance with the License. A copy of +# the License is located at +# +# http://aws.amazon.com/apache2.0/ +# +# or in the "license" file accompanying this file. This file is +# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF +# ANY KIND, either express or implied. See the License for the specific +# language governing permissions and limitations under the License. +"""Bounded polling shared by start-query/poll command lifecycles. + +Some query lifecycles (e.g. ``athena:GetQueryExecution`` and +``logs:GetQueryResults``) have no botocore waiter, so a command must poll a +started query to a terminal state by hand. Commands share the *timing* mechanism +here -- a monotonic deadline, a poll-count cap, and a geometric (1.5x, capped) +backoff -- while each keeps its own terminal-state and error handling so no +behaviour drifts between them. + +``clock`` and ``sleep`` are injectable so tests advance time deterministically +without real sleeping. +""" + +from __future__ import annotations + +import time +from collections.abc import Callable + +# Backoff ceiling applied to the sleep interval between polls. +MAX_POLL_INTERVAL = 5.0 +_GROWTH = 1.5 + + +class BoundedPoller: + """A monotonic, bounded poll loop clock with capped geometric backoff. + + Usage:: + + poller = BoundedPoller(interval=1.0, max_wait=300.0, max_polls=10_000, + clock=clock, sleep=sleep) + while True: + result = do_one_poll() + poller.record_poll() + if terminal(result): + break + if poller.exhausted(): + ... handle timeout (e.g. StopQuery) ... + break + poller.sleep_and_backoff() + """ + + def __init__( + self, + *, + interval: float, + max_wait: float, + max_polls: int, + clock: Callable[[], float], + sleep: Callable[[float], None], + max_interval: float = MAX_POLL_INTERVAL, + growth: float = _GROWTH, + ) -> None: + self._interval = float(interval) + self._max_wait = float(max_wait) + self._max_polls = int(max_polls) + self._clock = clock + self._sleep = sleep + self._max_interval = float(max_interval) + self._growth = float(growth) + # Monotonic deadline anchored at construction. + self._deadline = clock() + self._max_wait + self.polls = 0 + + def record_poll(self) -> int: + """Count one completed poll and return the running total.""" + self.polls += 1 + return self.polls + + def remaining(self) -> float: + """Seconds left before the deadline (clamp to >= 0 for sleep math).""" + return self._deadline - self._clock() + + def exhausted(self) -> bool: + """True once the time budget or the poll-count cap is reached.""" + return self.polls >= self._max_polls or self._clock() >= self._deadline + + def sleep_and_backoff(self) -> None: + """Sleep the current interval (clamped to remaining time), then grow it.""" + self._sleep(min(self._interval, max(0.0, self.remaining()))) + self._interval = min(self._interval * self._growth, self._max_interval) + + +def drive_poll( + call_op: Callable[[str, dict], dict], + poll_op: str, + params: dict, + *, + terminal: Callable[[dict], str | None], + is_retriable: Callable[[dict], bool], + on_timeout: Callable[[], None] | None = None, + interval: float = 1.0, + max_wait: float = 300.0, + max_polls: int = 10_000, + max_transient: int = 5, + clock: Callable[[], float] | None = None, + sleep: Callable[[float], None] | None = None, + key: str = "", +) -> dict: + """Drive a poll-until-terminal loop with bounded retries. + + Returns a result envelope (from helpers.ok_result/err_result): + - On terminal status: the last successful result from call_op + - On timeout: err_result with code 'TIMEOUT' and retriable=True + - On non-retriable error: the error result from call_op + - On transient errors exhausted: err_result with the last transient error + + The *terminal* callback receives a successful result dict and returns + the terminal status string if polling should stop, or None to continue. + + The *is_retriable* callback receives a failed result dict and returns + True if the error is transient (will be counted against max_transient). + + The returned envelope carries ``polls_count`` (the number of polls made) + so callers can surface it in their own diagnostics. + """ + # Late import keeps the dependency one-way (helpers never imports poll) and + # avoids any import-time cycle. + from awscli.customizations.compositereadonly.helpers import ( + TIMEOUT, + err_result, + ) + + sleep = time.sleep if sleep is None else sleep + clock = time.monotonic if clock is None else clock + + poller = BoundedPoller( + interval=interval, + max_wait=max_wait, + max_polls=max_polls, + clock=clock, + sleep=sleep, + max_interval=MAX_POLL_INTERVAL, + ) + transient = 0 + while True: + res = call_op(poll_op, params) + poller.record_poll() + if res["ok"]: + # Terminal callback decides whether this successful poll ends it. + if terminal(res) is not None: + res["polls_count"] = poller.polls + return res + elif is_retriable(res): + # Count transient errors; abort with the last one once exhausted. + if transient >= max_transient: + res["polls_count"] = poller.polls + return res + transient += 1 + else: + # Non-retriable error: hand the caller the raw failure result. + res["polls_count"] = poller.polls + return res + if poller.exhausted(): + # Let the caller cancel the lifecycle (e.g. StopQuery) before we + # report the timeout. + if on_timeout is not None: + on_timeout() + timed_out = err_result( + TIMEOUT, + f"Polling {poll_op} did not reach a terminal state within " + f"{max_wait:g}s after {poller.polls} polls.", + retriable=True, + key=key or None, + value={"polls": poller.polls}, + ) + timed_out["polls_count"] = poller.polls + return timed_out + poller.sleep_and_backoff() diff --git a/awscli/handlers_registry.py b/awscli/handlers_registry.py index d30348649e45..24e1d926a418 100644 --- a/awscli/handlers_registry.py +++ b/awscli/handlers_registry.py @@ -535,6 +535,9 @@ class CommandTableOp(enum.Enum): 'register_agent_toolkit_commands', ) ], + 'building-command-table.athena': [ + ('awscli.customizations.athena', 'register') + ], 'building-command-table.bedrock-agent-runtime': [ ('awscli.customizations.removals', 'register_removals') ], diff --git a/tests/functional/athena/__init__.py b/tests/functional/athena/__init__.py new file mode 100644 index 000000000000..654905217728 --- /dev/null +++ b/tests/functional/athena/__init__.py @@ -0,0 +1,12 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"). You +# may not use this file except in compliance with the License. A copy of +# the License is located at +# +# http://aws.amazon.com/apache2.0/ +# +# or in the "license" file accompanying this file. This file is +# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF +# ANY KIND, either express or implied. See the License for the specific +# language governing permissions and limitations under the License. diff --git a/tests/functional/athena/test_run_query.py b/tests/functional/athena/test_run_query.py new file mode 100644 index 000000000000..f8f372113068 --- /dev/null +++ b/tests/functional/athena/test_run_query.py @@ -0,0 +1,133 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"). You +# may not use this file except in compliance with the License. A copy of +# the License is located at +# +# http://aws.amazon.com/apache2.0/ +# +# or in the "license" file accompanying this file. This file is +# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF +# ANY KIND, either express or implied. See the License for the specific +# language governing permissions and limitations under the License. +import json + +from awscli.testutils import BaseAWSCommandParamsTest + +_MUTATING = { + 'CreateWorkGroup', + 'CreateNamedQuery', + 'UpdateWorkGroup', + 'DeleteWorkGroup', +} + + +class TestRunQueryCommand(BaseAWSCommandParamsTest): + def _lifecycle_responses(self): + return [ + { + 'WorkGroup': { + 'State': 'ENABLED', + 'Configuration': { + 'ResultConfiguration': { + 'OutputLocation': 's3://results/' + }, + 'EngineVersion': { + 'EffectiveEngineVersion': 'Athena engine version 3' + }, + }, + } + }, + {'QueryExecutionId': 'q-123'}, + { + 'QueryExecution': { + 'Status': {'State': 'SUCCEEDED'}, + 'StatementType': 'DML', + 'Statistics': {'DataScannedInBytes': 10}, + 'ResultConfiguration': {'OutputLocation': 's3://r/q.csv'}, + } + }, + { + 'ResultSet': { + 'ResultSetMetadata': { + 'ColumnInfo': [{'Name': 'c1', 'Type': 'varchar'}] + }, + 'Rows': [ + {'Data': [{'VarCharValue': 'c1'}]}, + {'Data': [{'VarCharValue': 'hello'}]}, + ], + } + }, + ] + + def test_runs_full_read_only_lifecycle(self): + self.parsed_responses = self._lifecycle_responses() + stdout, _, rc = self.run_cmd( + [ + 'athena', + 'run-query', + '--sql', + 'SELECT c1 FROM t', + '--workgroup', + 'primary', + '--region', + 'us-east-1', + ], + expected_rc=0, + ) + names = [c[0].name for c in self.operations_called] + self.assertEqual( + names, + [ + 'GetWorkGroup', + 'StartQueryExecution', + 'GetQueryExecution', + 'GetQueryResults', + ], + ) + # No mutating operation is ever issued. + self.assertFalse(_MUTATING.intersection(names)) + parsed = json.loads(stdout) + self.assertEqual(parsed['query_execution_id'], 'q-123') + self.assertEqual(parsed['rows'], [{'c1': 'hello'}]) + + def test_failed_query_returns_operational_error(self): + self.parsed_responses = [ + { + 'WorkGroup': { + 'State': 'ENABLED', + 'Configuration': { + 'ResultConfiguration': { + 'OutputLocation': 's3://results/' + } + }, + } + }, + {'QueryExecutionId': 'q-9'}, + { + 'QueryExecution': { + 'Status': { + 'State': 'FAILED', + 'StateChangeReason': 'bad column', + 'AthenaError': {'ErrorCategory': 2}, + }, + 'StatementType': 'DML', + } + }, + ] + stdout, _, rc = self.run_cmd( + [ + 'athena', + 'run-query', + '--sql', + 'SELECT x FROM t', + '--workgroup', + 'primary', + '--region', + 'us-east-1', + ], + expected_rc=255, + ) + self.assertNotIn( + 'StopQueryExecution', [c[0].name for c in self.operations_called] + ) diff --git a/tests/unit/customizations/athena/__init__.py b/tests/unit/customizations/athena/__init__.py new file mode 100644 index 000000000000..654905217728 --- /dev/null +++ b/tests/unit/customizations/athena/__init__.py @@ -0,0 +1,12 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"). You +# may not use this file except in compliance with the License. A copy of +# the License is located at +# +# http://aws.amazon.com/apache2.0/ +# +# or in the "license" file accompanying this file. This file is +# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF +# ANY KIND, either express or implied. See the License for the specific +# language governing permissions and limitations under the License. diff --git a/tests/unit/customizations/athena/test_resolve_workgroup.py b/tests/unit/customizations/athena/test_resolve_workgroup.py new file mode 100644 index 000000000000..0ce91a6df256 --- /dev/null +++ b/tests/unit/customizations/athena/test_resolve_workgroup.py @@ -0,0 +1,226 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"). You +# may not use this file except in compliance with the License. A copy of +# the License is located at +# +# http://aws.amazon.com/apache2.0/ +# +# or in the "license" file accompanying this file. This file is +# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF +# ANY KIND, either express or implied. See the License for the specific +# language governing permissions and limitations under the License. +"""Unit tests for Athena's workgroup-with-output resolver. + +``resolve_workgroup_with_output`` used to live in the shared compositereadonly +``resolve`` module; it now lives with its owning command in +:mod:`awscli.customizations.athena.run_query`, so its tests live here too. +""" + +from awscli.customizations.athena import run_query as rq +from awscli.customizations.compositereadonly.helpers import ( + NOT_FOUND, + PERMISSION_DENIED, + assert_read_only, + err_result, + ok_result, +) + + +class Scripted: + """A scripted read-only call bridge; records calls, returns queued dicts.""" + + def __init__(self): + self.calls = [] + self._responses = {} + + def queue(self, operation, *results): + self._responses.setdefault(operation, []).extend(results) + return self + + def __call__(self, operation, params): + assert_read_only(operation) + self.calls.append((operation, params)) + queued = self._responses.get(operation) + if queued: + return queued.pop(0) + return ok_result({}) + + def ops(self): + return [op for op, _ in self.calls] + + +def _wg(name="primary", output=None, managed=False, state="ENABLED"): + cfg = {} + if output is not None: + cfg["ResultConfiguration"] = {"OutputLocation": output} + if managed: + cfg["ManagedQueryResultsConfiguration"] = {"Enabled": True} + return ok_result( + {"WorkGroup": {"Name": name, "State": state, "Configuration": cfg}} + ) + + +# ----------------------------------------------- resolve_workgroup_with_output + + +def test_explicit_workgroup_is_honoured_without_any_call(): + call = Scripted() + r = rq.resolve_workgroup_with_output(call, "analytics") + assert r["ok"] is True + assert r["value"]["workgroup"] == "analytics" + assert r["value"]["auto_resolved"] is False + # An explicit choice never triggers a read. + assert call.calls == [] + + +def test_primary_with_output_location_is_used(): + call = Scripted() + call.queue("GetWorkGroup", _wg(output="s3://primary/")) + r = rq.resolve_workgroup_with_output(call, None) + assert r["value"]["workgroup"] == "primary" + assert r["value"]["output_location"] == "s3://primary/" + assert r["value"]["auto_resolved"] is False + assert call.ops() == ["GetWorkGroup"] + + +def test_primary_with_managed_results_is_used(): + call = Scripted() + call.queue("GetWorkGroup", _wg(managed=True)) + r = rq.resolve_workgroup_with_output(call, None) + assert r["value"]["managed_results"] is True + assert r["value"]["workgroup"] == "primary" + + +def test_auto_selects_enabled_workgroup_with_output(): + call = Scripted() + # primary has neither output nor managed results. + call.queue("GetWorkGroup", _wg(output=None)) + call.queue( + "ListWorkGroups", + ok_result( + { + "WorkGroups": [ + {"Name": "disabled-wg", "State": "DISABLED"}, + {"Name": "noname-skip"}, # has a name, no output + {"Name": "analytics", "State": "ENABLED"}, + ] + } + ), + ) + # First enabled candidate with no name key filtered; 'noname-skip' read but + # has no output, 'analytics' has one and wins. + call.queue("GetWorkGroup", _wg(name="noname-skip", output=None)) + call.queue("GetWorkGroup", _wg(name="analytics", output="s3://a/")) + r = rq.resolve_workgroup_with_output(call, None) + assert r["value"]["workgroup"] == "analytics" + assert r["value"]["auto_resolved"] is True + assert r["value"]["output_location"] == "s3://a/" + assert any("auto-selected" in g for g in r["guidance"]) + + +def test_auto_selects_managed_workgroup_reports_managed(): + call = Scripted() + call.queue("GetWorkGroup", _wg(output=None)) + call.queue( + "ListWorkGroups", + ok_result({"WorkGroups": [{"Name": "mgd", "State": "ENABLED"}]}), + ) + call.queue("GetWorkGroup", _wg(name="mgd", managed=True)) + r = rq.resolve_workgroup_with_output(call, None) + assert r["value"]["workgroup"] == "mgd" + assert r["value"]["managed_results"] is True + # Guidance names managed results when there is no output location. + assert any("managed results" in g for g in r["guidance"]) + + +def test_skips_entries_without_name_and_failed_getworkgroup(): + call = Scripted() + call.queue("GetWorkGroup", _wg(output=None)) + call.queue( + "ListWorkGroups", + ok_result( + { + "WorkGroups": [ + {"State": "ENABLED"}, # no Name -> skipped, no read + {"Name": "broken", "State": "ENABLED"}, + {"Name": "good", "State": "ENABLED"}, + ] + } + ), + ) + call.queue( + "GetWorkGroup", err_result(PERMISSION_DENIED, "denied") + ) # 'broken' read fails -> continue + call.queue("GetWorkGroup", _wg(name="good", output="s3://g/")) + r = rq.resolve_workgroup_with_output(call, None) + assert r["value"]["workgroup"] == "good" + + +def test_listworkgroups_failure_falls_back_to_primary(): + call = Scripted() + call.queue("GetWorkGroup", _wg(output=None)) + call.queue("ListWorkGroups", err_result(PERMISSION_DENIED, "denied")) + r = rq.resolve_workgroup_with_output(call, None) + assert r["value"]["workgroup"] == "primary" + assert r["value"]["auto_resolved"] is False + assert any("falling back to 'primary'" in g for g in r["guidance"]) + + +def test_pagination_followed_then_nothing_found(): + call = Scripted() + call.queue("GetWorkGroup", _wg(output=None)) # primary + # Page 1 has one candidate with no output; a NextToken leads to an empty + # page 2; resolver then exhausts and falls back to 'primary'. + call.queue( + "ListWorkGroups", + ok_result( + { + "WorkGroups": [{"Name": "x", "State": "ENABLED"}], + "NextToken": "tok2", + } + ), + ) + call.queue("GetWorkGroup", _wg(name="x", output=None)) + call.queue("ListWorkGroups", ok_result({"WorkGroups": []})) + r = rq.resolve_workgroup_with_output(call, None) + assert r["value"]["workgroup"] == "primary" + assert r["value"]["auto_resolved"] is False + assert any("No workgroup with a result output" in g for g in r["guidance"]) + # Pagination token on page 1 was followed into a second ListWorkGroups. + assert call.ops().count("ListWorkGroups") == 2 + + +def test_workgroup_listing_stops_at_page_cap(): + # Every page carries a NextToken and only unusable workgroups, so the loop + # runs the full _MAX_WORKGROUP_PAGES and then falls back to 'primary' + # instead of paging forever. + call = Scripted() + call.queue("GetWorkGroup", _wg(output=None)) # primary + for _ in range(rq._MAX_WORKGROUP_PAGES + 2): + # No "Name" key -> each candidate is skipped (no GetWorkGroup read). + call.queue( + "ListWorkGroups", + ok_result( + {"WorkGroups": [{"State": "ENABLED"}], "NextToken": "more"} + ), + ) + r = rq.resolve_workgroup_with_output(call, None) + assert r["value"]["workgroup"] == "primary" + assert r["value"]["auto_resolved"] is False + # The cap bounded the enumeration to exactly _MAX_WORKGROUP_PAGES reads. + assert call.ops().count("ListWorkGroups") == rq._MAX_WORKGROUP_PAGES + + +def test_primary_unreadable_still_enumerates(): + call = Scripted() + # primary read fails -> skip the has-output short-circuit, go to listing. + call.queue("GetWorkGroup", err_result(NOT_FOUND, "gone")) + call.queue( + "ListWorkGroups", + ok_result({"WorkGroups": [{"Name": "a", "State": "ENABLED"}]}), + ) + call.queue("GetWorkGroup", _wg(name="a", output="s3://a/")) + r = rq.resolve_workgroup_with_output(call, None) + assert r["value"]["workgroup"] == "a" + assert r["value"]["auto_resolved"] is True diff --git a/tests/unit/customizations/athena/test_run_query.py b/tests/unit/customizations/athena/test_run_query.py new file mode 100644 index 000000000000..54122ba45419 --- /dev/null +++ b/tests/unit/customizations/athena/test_run_query.py @@ -0,0 +1,1106 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"). You +# may not use this file except in compliance with the License. A copy of +# the License is located at +# +# http://aws.amazon.com/apache2.0/ +# +# or in the "license" file accompanying this file. This file is +# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF +# ANY KIND, either express or implied. See the License for the specific +# language governing permissions and limitations under the License. +"""Unit tests for ``aws athena run-query`` (run_athena_query logic).""" + +from awscli.customizations.athena.run_query import run_athena_query +from awscli.customizations.compositereadonly.helpers import ( + CREDENTIALS_EXPIRED, + PERMISSION_DENIED, + THROTTLED, + assert_read_only, + err_result, + ok_result, +) + + +class ScriptedAthena: + """A scripted read-only call bridge: records calls, returns queued results.""" + + def __init__(self): + self.calls = [] + self._responses = {} + + def queue(self, operation, *results): + self._responses.setdefault(operation, []).extend(results) + return self + + def __call__(self, operation, params): + # Every scripted operation must be read-only, mirroring the real bridge. + assert_read_only(operation) + self.calls.append((operation, params)) + queued = self._responses.get(operation) + if queued: + return queued.pop(0) + return ok_result({}) + + def ops(self): + return [op for op, _ in self.calls] + + +def _wg( + output_location="s3://results/", + enforced=False, + state="ENABLED", + engine="Athena engine version 3", + managed=False, +): + cfg = { + "EnforceWorkGroupConfiguration": enforced, + "EngineVersion": {"EffectiveEngineVersion": engine}, + } + if output_location is not None: + cfg["ResultConfiguration"] = {"OutputLocation": output_location} + if managed: + cfg["ManagedQueryResultsConfiguration"] = {"Enabled": True} + return ok_result({"WorkGroup": {"State": state, "Configuration": cfg}}) + + +def _exec(state="SUCCEEDED", stmt_type="DML", reason=None, category=None): + status = {"State": state} + if reason: + status["StateChangeReason"] = reason + if category is not None: + status["AthenaError"] = {"ErrorCategory": category} + return ok_result( + { + "QueryExecution": { + "Status": status, + "StatementType": stmt_type, + "Statistics": {"DataScannedInBytes": 123}, + "ResultConfiguration": { + "OutputLocation": "s3://results/q.csv" + }, + } + } + ) + + +def _results(rows, columns=None, next_token=None, with_header=True): + columns = columns or [{"Name": "c1", "Type": "varchar"}] + data_rows = [] + if with_header: + data_rows.append( + {"Data": [{"VarCharValue": c["Name"]} for c in columns]} + ) + for r in rows: + data_rows.append({"Data": [{"VarCharValue": v} for v in r]}) + body = { + "ResultSet": { + "ResultSetMetadata": {"ColumnInfo": columns}, + "Rows": data_rows, + } + } + if next_token: + body["NextToken"] = next_token + return ok_result(body) + + +def _run(call_op, sql="SELECT 1", **kwargs): + kwargs.setdefault("workgroup", "primary") + return run_athena_query( + call_op, sql, sleep=lambda _s: None, clock=_fake_clock(), **kwargs + ) + + +def _fake_clock(): + state = {"t": 0.0} + + def clock(): + state["t"] += 0.001 + return state["t"] + + return clock + + +def _expired_clock(): + # First call (poller construction) returns 0.0, so the deadline is max_wait; + # every later call returns a time well past it, forcing a timeout. + calls = {"n": 0} + + def clock(): + calls["n"] += 1 + return 0.0 if calls["n"] == 1 else 1000.0 + + return clock + + +# --------------------------------------------------------------- input validation + + +def test_rejects_non_s3_output_location(): + call = ScriptedAthena() + result = _run(call, sql="SELECT 1", output_location="/tmp/x") + assert result["ok"] is False + assert result["code"] == "INVALID_INPUT" + + +def test_rejects_bad_max_rows(): + call = ScriptedAthena() + result = _run(call, sql="SELECT 1", max_rows=0) + assert result["ok"] is False + assert result["code"] == "INVALID_INPUT" + + +def test_rejects_empty_sql(): + # An empty --sql is caught client-side as a ParamValidation error (exit 252, + # param_error=True) instead of slipping through to Athena as a SERVICE_ERROR. + call = ScriptedAthena() + result = _run(call, sql="") + assert result["ok"] is False + assert result["code"] == "INVALID_INPUT" + assert result["param_error"] is True + assert "sql is an empty string" in result["error"] + # No AWS call is made for an empty statement. + assert call.ops() == [] + + +# --------------------------------------------------------------- happy path + + +def test_success_returns_rows_and_only_read_ops(): + call = ScriptedAthena() + call.queue("GetWorkGroup", _wg()) + call.queue("StartQueryExecution", ok_result({"QueryExecutionId": "q1"})) + call.queue("GetQueryExecution", _exec()) + call.queue("GetQueryResults", _results([["a"], ["b"]])) + result = _run(call, sql="SELECT c1 FROM t") + assert result["ok"] is True + value = result["value"] + assert value["query_execution_id"] == "q1" + assert value["row_count"] == 2 + assert value["rows"] == [{"c1": "a"}, {"c1": "b"}] + assert value["header_row_dropped"] is True + assert result["is_exhaustive"] is True + # Read-only: only the query lifecycle read ops were issued. + assert call.ops() == [ + "GetWorkGroup", + "StartQueryExecution", + "GetQueryExecution", + "GetQueryResults", + ] + + +def test_integer_columns_are_coerced(): + call = ScriptedAthena() + call.queue("GetWorkGroup", _wg()) + call.queue("StartQueryExecution", ok_result({"QueryExecutionId": "q1"})) + call.queue("GetQueryExecution", _exec()) + call.queue( + "GetQueryResults", + _results([["42"]], columns=[{"Name": "n", "Type": "bigint"}]), + ) + result = _run(call, sql="SELECT n FROM t") + assert result["value"]["rows"] == [{"n": 42}] + + +def test_no_coerce_keeps_strings(): + call = ScriptedAthena() + call.queue("GetWorkGroup", _wg()) + call.queue("StartQueryExecution", ok_result({"QueryExecutionId": "q1"})) + call.queue("GetQueryExecution", _exec()) + call.queue( + "GetQueryResults", + _results([["42"]], columns=[{"Name": "n", "Type": "bigint"}]), + ) + result = _run(call, sql="SELECT n FROM t", coerce_types=False) + assert result["value"]["rows"] == [{"n": "42"}] + + +def test_empty_result_set(): + call = ScriptedAthena() + call.queue("GetWorkGroup", _wg()) + call.queue("StartQueryExecution", ok_result({"QueryExecutionId": "q1"})) + call.queue("GetQueryExecution", _exec()) + call.queue("GetQueryResults", _results([])) + result = _run(call, sql="SELECT c1 FROM t") + assert result["ok"] is True + assert result["value"]["row_count"] == 0 + + +# --------------------------------------------------------------- pagination / cap + + +def test_pagination_follows_next_token(): + call = ScriptedAthena() + call.queue("GetWorkGroup", _wg()) + call.queue("StartQueryExecution", ok_result({"QueryExecutionId": "q1"})) + call.queue("GetQueryExecution", _exec()) + call.queue("GetQueryResults", _results([["a"]], next_token="t2")) + call.queue("GetQueryResults", _results([["b"]], with_header=False)) + result = _run(call, sql="SELECT c1 FROM t") + assert result["value"]["row_count"] == 2 + assert result["value"]["pages_fetched"] == 2 + + +def test_row_cap_truncates_and_exposes_resume_token(): + call = ScriptedAthena() + call.queue("GetWorkGroup", _wg()) + call.queue("StartQueryExecution", ok_result({"QueryExecutionId": "q1"})) + call.queue("GetQueryExecution", _exec()) + call.queue( + "GetQueryResults", _results([["a"], ["b"], ["c"]], next_token="more") + ) + result = _run(call, sql="SELECT c1 FROM t", max_rows=2) + assert result["value"]["is_truncated"] is True + assert result["value"]["next_token"] == "more" + assert result["is_exhaustive"] is False + assert result["value"]["row_count"] == 2 + + +# --------------------------------------------------------------- failures + + +def test_workgroup_not_found(): + call = ScriptedAthena() + call.queue( + "GetWorkGroup", + err_result( + "NOT_FOUND", + "WorkGroup not found", + aws_code="InvalidRequestException", + ), + ) + result = _run(call, sql="SELECT 1", workgroup="ghost") + assert result["ok"] is False + assert result["code"] == "NOT_FOUND" + # No query was started against a missing workgroup. + assert "StartQueryExecution" not in call.ops() + + +def test_failed_query_surfaces_state_change_reason(): + call = ScriptedAthena() + call.queue("GetWorkGroup", _wg()) + call.queue("StartQueryExecution", ok_result({"QueryExecutionId": "q1"})) + call.queue( + "GetQueryExecution", + _exec( + state="FAILED", reason="Column 'x' cannot be resolved", category=2 + ), + ) + result = _run(call, sql="SELECT x FROM t") + assert result["ok"] is False + assert result["code"] == "INVALID_INPUT" + assert "cannot be resolved" in result["value"]["state_change_reason"] + + +def test_cancelled_query_is_conflict(): + call = ScriptedAthena() + call.queue("GetWorkGroup", _wg()) + call.queue("StartQueryExecution", ok_result({"QueryExecutionId": "q1"})) + call.queue("GetQueryExecution", _exec(state="CANCELLED", reason="by user")) + result = _run(call, sql="SELECT 1 FROM t") + assert result["ok"] is False + assert result["code"] == "CONFLICT" + + +def test_timeout_returns_query_id_and_does_not_resubmit(): + call = ScriptedAthena() + call.queue("GetWorkGroup", _wg()) + call.queue("StartQueryExecution", ok_result({"QueryExecutionId": "q1"})) + # Always RUNNING -> never terminal; deadline forces a timeout. + for _ in range(50): + call.queue("GetQueryExecution", _exec(state="RUNNING")) + result = run_athena_query( + call, + "SELECT 1 FROM t", + workgroup="primary", + max_wait=1.0, + sleep=lambda _s: None, + clock=_expired_clock(), + ) + assert result["ok"] is False + assert result["code"] == "TIMEOUT" + assert result["key"] == "q1" + assert call.ops().count("StartQueryExecution") == 1 + + +def test_cancel_on_timeout_stops_the_query(): + call = ScriptedAthena() + call.queue("GetWorkGroup", _wg()) + call.queue("StartQueryExecution", ok_result({"QueryExecutionId": "q1"})) + for _ in range(50): + call.queue("GetQueryExecution", _exec(state="RUNNING")) + call.queue("StopQueryExecution", ok_result({})) + result = run_athena_query( + call, + "SELECT 1 FROM t", + workgroup="primary", + max_wait=1.0, + cancel_on_timeout=True, + sleep=lambda _s: None, + clock=_expired_clock(), + ) + assert result["code"] == "TIMEOUT" + assert result["value"]["cancelled"] is True + assert "StopQueryExecution" in call.ops() + + +# --- workgroup output-location auto-resolution + + +def test_workgroup_auto_resolved_when_primary_lacks_output_location(): + call = ScriptedAthena() + # resolve step: primary has no output location, then list + pick 'analytics'. + call.queue("GetWorkGroup", _wg(output_location=None)) # resolve: primary + call.queue( + "ListWorkGroups", + ok_result({"WorkGroups": [{"Name": "analytics", "State": "ENABLED"}]}), + ) + call.queue( + "GetWorkGroup", _wg(output_location="s3://picked/") + ) # resolve: analytics + # run step re-reads the chosen workgroup. + call.queue("GetWorkGroup", _wg(output_location="s3://picked/")) + call.queue("StartQueryExecution", ok_result({"QueryExecutionId": "q1"})) + call.queue("GetQueryExecution", _exec()) + call.queue("GetQueryResults", _results([["a"]])) + result = run_athena_query( + call, + "SELECT 1 FROM t", + workgroup=None, + sleep=lambda _s: None, + clock=_fake_clock(), + ) + assert result["ok"] is True + assert result["value"]["workgroup"] == "analytics" + assert any("auto-selected workgroup" in g for g in result["guidance"]) + + +def test_enforced_workgroup_ignores_caller_output_location(): + call = ScriptedAthena() + call.queue( + "GetWorkGroup", + _wg(output_location="s3://enforced/", enforced=True), + ) + call.queue("StartQueryExecution", ok_result({"QueryExecutionId": "q1"})) + call.queue("GetQueryExecution", _exec()) + call.queue("GetQueryResults", _results([["a"]])) + result = _run( + call, + sql="SELECT 1 FROM t", + output_location="s3://caller/", + workgroup="wg", + ) + assert result["ok"] is True + start_params = dict(call.calls)["StartQueryExecution"] + # Enforced config: we do NOT send a ResultConfiguration override. + assert "ResultConfiguration" not in start_params + + +# --------------------------------------------------------------- param validation (direct) + + +def test_empty_workgroup_string_rejected(): + result = run_athena_query( + lambda *a: ok_result({}), + "SELECT 1", + workgroup=" ", + sleep=lambda _s: None, + clock=_fake_clock(), + ) + assert result["ok"] is False + assert "workgroup is empty" in result["error"] + + +def test_empty_catalog_rejected(): + result = run_athena_query( + lambda *a: ok_result({}), + "SELECT 1", + workgroup="primary", + catalog=" ", + sleep=lambda _s: None, + clock=_fake_clock(), + ) + assert result["ok"] is False + assert "catalog is empty" in result["error"] + + +def test_empty_database_string_rejected(): + result = run_athena_query( + lambda *a: ok_result({}), + "SELECT 1", + workgroup="primary", + database="", + sleep=lambda _s: None, + clock=_fake_clock(), + ) + assert result["ok"] is False + assert "database is an empty string" in result["error"] + + +def test_output_location_too_short_rejected(): + call = ScriptedAthena() + result = _run(call, sql="SELECT 1", output_location="s3://") + assert result["ok"] is False + assert result["code"] == "INVALID_INPUT" + + +def test_boolean_max_rows_rejected(): + # bool is a subclass of int; it must be rejected, not accepted as 1. + call = ScriptedAthena() + result = _run(call, sql="SELECT 1", max_rows=True) + assert result["ok"] is False + assert "max_rows" in result["error"] + + +def test_nonpositive_poll_interval_rejected(): + result = run_athena_query( + lambda *a: ok_result({}), + "SELECT 1", + workgroup="primary", + poll_interval=0, + sleep=lambda _s: None, + clock=_fake_clock(), + ) + assert result["ok"] is False + assert "poll_interval" in result["error"] + + +def test_nonpositive_max_wait_rejected(): + result = run_athena_query( + lambda *a: ok_result({}), + "SELECT 1", + workgroup="primary", + max_wait=-1.0, + sleep=lambda _s: None, + clock=_fake_clock(), + ) + assert result["ok"] is False + assert "max_wait" in result["error"] + + +# --------------------------------------------------------------- workgroup read + + +def test_disabled_workgroup_is_not_ready(): + call = ScriptedAthena() + call.queue("GetWorkGroup", _wg(state="DISABLED")) + result = _run(call, sql="SELECT 1", workgroup="wg") + assert result["ok"] is False + assert result["code"] == "NOT_READY" + assert "StartQueryExecution" not in call.ops() + + +def test_workgroup_not_found_via_error_message(): + call = ScriptedAthena() + call.queue( + "GetWorkGroup", + err_result( + "SERVICE_ERROR", + "WorkGroup ghost not found", + ), + ) + result = _run(call, sql="SELECT 1", workgroup="ghost") + assert result["ok"] is False + assert result["code"] == "NOT_FOUND" + + +def test_workgroup_credentials_expired(): + call = ScriptedAthena() + call.queue( + "GetWorkGroup", + err_result(CREDENTIALS_EXPIRED, "token expired"), + ) + result = _run(call, sql="SELECT 1", workgroup="wg") + assert result["ok"] is False + assert result["code"] == CREDENTIALS_EXPIRED + assert result["retriable"] is True + + +def test_unreadable_workgroup_proceeds_with_note(): + # A non-fatal GetWorkGroup error: proceed, but with no enforced location + # the caller must supply one, which we do. + call = ScriptedAthena() + call.queue("GetWorkGroup", err_result("SERVICE_ERROR", "flaky")) + call.queue("StartQueryExecution", ok_result({"QueryExecutionId": "q1"})) + call.queue("GetQueryExecution", _exec()) + call.queue("GetQueryResults", _results([["a"]])) + result = _run( + call, + sql="SELECT 1", + workgroup="wg", + output_location="s3://caller/", + ) + assert result["ok"] is True + assert any("proceeding without" in g for g in result["guidance"]) + + +def test_selected_engine_version_used_when_no_effective(): + call = ScriptedAthena() + wg = ok_result( + { + "WorkGroup": { + "State": "ENABLED", + "Configuration": { + "ResultConfiguration": {"OutputLocation": "s3://r/"}, + "EngineVersion": { + "SelectedEngineVersion": "Athena engine version 2" + }, + }, + } + } + ) + call.queue("GetWorkGroup", wg) + call.queue("StartQueryExecution", ok_result({"QueryExecutionId": "q1"})) + # The terminal execution carries no EngineVersion, so the workgroup's + # selected version survives into the result. + call.queue( + "GetQueryExecution", + ok_result( + { + "QueryExecution": { + "Status": {"State": "SUCCEEDED"}, + "StatementType": "DML", + "Statistics": {"DataScannedInBytes": 1}, + } + } + ), + ) + call.queue("GetQueryResults", _results([["a"]])) + result = _run(call, sql="SELECT 1", workgroup="wg") + assert result["value"]["engine_version"] == "Athena engine version 2" + + +def test_managed_results_workgroup_sends_no_result_config(): + call = ScriptedAthena() + call.queue("GetWorkGroup", _wg(output_location=None, managed=True)) + call.queue("StartQueryExecution", ok_result({"QueryExecutionId": "q1"})) + call.queue("GetQueryExecution", _exec()) + call.queue("GetQueryResults", _results([["a"]])) + result = _run(call, sql="SELECT 1", workgroup="wg") + assert result["ok"] is True + assert result["value"]["output_location_source"] == "workgroup_managed" + start = dict(call.calls)["StartQueryExecution"] + assert "ResultConfiguration" not in start + + +def test_enforced_managed_results_source(): + call = ScriptedAthena() + call.queue( + "GetWorkGroup", + _wg(output_location=None, managed=True, enforced=True), + ) + call.queue("StartQueryExecution", ok_result({"QueryExecutionId": "q1"})) + call.queue("GetQueryExecution", _exec()) + call.queue("GetQueryResults", _results([["a"]])) + result = _run( + call, + sql="SELECT 1", + workgroup="wg", + output_location="s3://caller/", + ) + assert result["ok"] is True + assert result["value"]["output_location_source"] == "workgroup_managed" + assert any("enforces its configuration" in g for g in result["guidance"]) + + +def test_no_output_location_anywhere_is_invalid_input(): + # Workgroup read OK, not enforced, no default/managed location, and no + # caller output -> a precise INVALID_INPUT before starting the query. + call = ScriptedAthena() + call.queue("GetWorkGroup", _wg(output_location=None)) + result = _run(call, sql="SELECT 1", workgroup="wg") + assert result["ok"] is False + assert result["code"] == "INVALID_INPUT" + assert "no result output location" in result["error"] + assert "StartQueryExecution" not in call.ops() + + +# --------------------------------------------------------------- start failures + + +def test_start_not_supported_maps_to_invalid_input(): + call = ScriptedAthena() + call.queue("GetWorkGroup", _wg()) + call.queue( + "StartQueryExecution", + err_result( + "SERVICE_ERROR", + "Queries of this type are not supported", + aws_code="InvalidRequestException", + ), + ) + result = _run(call, sql="SELECT 1", workgroup="wg") + assert result["ok"] is False + assert result["code"] == "INVALID_INPUT" + + +def test_start_permission_denied_guidance(): + call = ScriptedAthena() + call.queue("GetWorkGroup", _wg()) + call.queue( + "StartQueryExecution", + err_result(PERMISSION_DENIED, "not authorized"), + ) + result = _run(call, sql="SELECT 1", workgroup="wg") + assert result["code"] == PERMISSION_DENIED + assert any("StartQueryExecution" in g for g in result["guidance"]) + + +def test_start_throttled_is_retriable(): + call = ScriptedAthena() + call.queue("GetWorkGroup", _wg()) + call.queue("StartQueryExecution", err_result(THROTTLED, "slow down")) + result = _run(call, sql="SELECT 1", workgroup="wg") + assert result["code"] == THROTTLED + assert result["retriable"] is True + + +def test_start_credentials_expired_is_retriable(): + call = ScriptedAthena() + call.queue("GetWorkGroup", _wg()) + call.queue( + "StartQueryExecution", err_result(CREDENTIALS_EXPIRED, "expired") + ) + result = _run(call, sql="SELECT 1", workgroup="wg") + assert result["code"] == CREDENTIALS_EXPIRED + assert result["retriable"] is True + + +def test_start_generic_failure(): + call = ScriptedAthena() + call.queue("GetWorkGroup", _wg()) + call.queue("StartQueryExecution", err_result("SERVICE_ERROR", "boom")) + result = _run(call, sql="SELECT 1", workgroup="wg") + assert result["code"] == "SERVICE_ERROR" + + +def test_start_returns_no_query_id(): + call = ScriptedAthena() + call.queue("GetWorkGroup", _wg()) + call.queue("StartQueryExecution", ok_result({})) # no QueryExecutionId + result = _run(call, sql="SELECT 1", workgroup="wg") + assert result["ok"] is False + assert result["code"] == "SERVICE_ERROR" + assert "no QueryExecutionId" in result["error"] + + +# --------------------------------------------------------------- polling edge cases + + +def test_transient_poll_errors_then_success(): + call = ScriptedAthena() + call.queue("GetWorkGroup", _wg()) + call.queue("StartQueryExecution", ok_result({"QueryExecutionId": "q1"})) + # Two retriable poll errors are tolerated before a terminal success. + call.queue("GetQueryExecution", err_result(THROTTLED, "slow")) + call.queue( + "GetQueryExecution", + err_result("NETWORK_ERROR", "reset", retriable=True), + ) + call.queue("GetQueryExecution", _exec()) + call.queue("GetQueryResults", _results([["a"]])) + result = _run(call, sql="SELECT 1", workgroup="wg") + assert result["ok"] is True + assert call.ops().count("GetQueryExecution") == 3 + + +def test_non_retriable_poll_error_returns_started_guidance(): + call = ScriptedAthena() + call.queue("GetWorkGroup", _wg()) + call.queue("StartQueryExecution", ok_result({"QueryExecutionId": "q1"})) + call.queue("GetQueryExecution", err_result(PERMISSION_DENIED, "denied")) + result = _run(call, sql="SELECT 1", workgroup="wg") + assert result["ok"] is False + assert result["code"] == PERMISSION_DENIED + assert result["key"] == "q1" + assert any("was started" in g for g in result["guidance"]) + + +def test_cancel_on_timeout_stop_fails_stays_retriable(): + call = ScriptedAthena() + call.queue("GetWorkGroup", _wg()) + call.queue("StartQueryExecution", ok_result({"QueryExecutionId": "q1"})) + for _ in range(50): + call.queue("GetQueryExecution", _exec(state="RUNNING")) + # StopQueryExecution itself fails, so the query is NOT confirmed cancelled. + call.queue("StopQueryExecution", err_result("SERVICE_ERROR", "nope")) + result = run_athena_query( + call, + "SELECT 1 FROM t", + workgroup="primary", + max_wait=1.0, + cancel_on_timeout=True, + sleep=lambda _s: None, + clock=_expired_clock(), + ) + assert result["code"] == "TIMEOUT" + assert result["value"]["cancelled"] is False + assert result["retriable"] is True + + +# --------------------------------------------------------------- result decoding + + +def test_boolean_column_is_coerced(): + call = ScriptedAthena() + call.queue("GetWorkGroup", _wg()) + call.queue("StartQueryExecution", ok_result({"QueryExecutionId": "q1"})) + call.queue("GetQueryExecution", _exec()) + call.queue( + "GetQueryResults", + _results( + [["true"], ["false"]], + columns=[{"Name": "flag", "Type": "boolean"}], + ), + ) + result = _run(call, sql="SELECT flag FROM t") + assert result["value"]["rows"] == [{"flag": True}, {"flag": False}] + + +def test_non_integer_value_in_integer_column_stays_string(): + call = ScriptedAthena() + call.queue("GetWorkGroup", _wg()) + call.queue("StartQueryExecution", ok_result({"QueryExecutionId": "q1"})) + call.queue("GetQueryExecution", _exec()) + call.queue( + "GetQueryResults", + _results([["NaN"]], columns=[{"Name": "n", "Type": "bigint"}]), + ) + result = _run(call, sql="SELECT n FROM t") + # int("NaN") raises -> value is left as the verbatim string. + assert result["value"]["rows"] == [{"n": "NaN"}] + + +def test_null_cell_is_preserved_as_none(): + call = ScriptedAthena() + call.queue("GetWorkGroup", _wg()) + call.queue("StartQueryExecution", ok_result({"QueryExecutionId": "q1"})) + call.queue("GetQueryExecution", _exec()) + # A cell with no VarCharValue decodes to SQL NULL -> None. + body = ok_result( + { + "ResultSet": { + "ResultSetMetadata": { + "ColumnInfo": [{"Name": "c1", "Type": "varchar"}] + }, + "Rows": [ + {"Data": [{"VarCharValue": "c1"}]}, + {"Data": [{}]}, + ], + } + } + ) + call.queue("GetQueryResults", body) + result = _run(call, sql="SELECT c1 FROM t") + assert result["value"]["rows"] == [{"c1": None}] + + +def test_duplicate_column_labels_are_suffixed_with_note(): + call = ScriptedAthena() + call.queue("GetWorkGroup", _wg()) + call.queue("StartQueryExecution", ok_result({"QueryExecutionId": "q1"})) + call.queue("GetQueryExecution", _exec()) + call.queue( + "GetQueryResults", + _results( + [["1", "2"]], + columns=[ + {"Name": "c", "Type": "varchar"}, + {"Name": "c", "Type": "varchar"}, + ], + ), + ) + result = _run(call, sql="SELECT a AS c, b AS c FROM t") + assert result["value"]["rows"] == [{"c": "1", "c_2": "2"}] + assert any("Duplicate column" in g for g in result["guidance"]) + + +def test_non_dml_header_dropped_when_first_row_matches_names(): + # SHOW/DESCRIBE are non-DML (expect_header False); the header is dropped + # only because the first row equals the column names. + call = ScriptedAthena() + call.queue("GetWorkGroup", _wg()) + call.queue("StartQueryExecution", ok_result({"QueryExecutionId": "q1"})) + call.queue("GetQueryExecution", _exec(stmt_type="UTILITY")) + call.queue( + "GetQueryResults", + _results( + [["partition_a"]], + columns=[{"Name": "partition", "Type": "varchar"}], + ), + ) + result = _run(call, sql="SHOW PARTITIONS t") + assert result["value"]["header_row_dropped"] is True + assert result["value"]["rows"] == [{"partition": "partition_a"}] + + +def test_get_query_results_failure_midway_returns_partial(): + call = ScriptedAthena() + call.queue("GetWorkGroup", _wg()) + call.queue("StartQueryExecution", ok_result({"QueryExecutionId": "q1"})) + call.queue("GetQueryExecution", _exec()) + call.queue("GetQueryResults", _results([["a"]], next_token="t2")) + call.queue( + "GetQueryResults", + err_result(THROTTLED, "throttled on page 2", retriable=True), + ) + result = _run(call, sql="SELECT c1 FROM t") + assert result["ok"] is False + assert result["code"] == THROTTLED + # The already-fetched page and a resume token are surfaced, not discarded. + assert result["value"]["row_count"] == 1 + assert result["value"]["is_truncated"] is True + assert result["is_exhaustive"] is False + + +def test_repeated_next_token_stops_pagination(): + call = ScriptedAthena() + call.queue("GetWorkGroup", _wg()) + call.queue("StartQueryExecution", ok_result({"QueryExecutionId": "q1"})) + call.queue("GetQueryExecution", _exec()) + # Server repeats the same NextToken; the loop must break, not spin. + call.queue("GetQueryResults", _results([["a"]], next_token="dup")) + call.queue( + "GetQueryResults", + _results([["b"]], next_token="dup", with_header=False), + ) + result = _run(call, sql="SELECT c1 FROM t") + assert result["ok"] is True + assert any("repeated a NextToken" in g for g in result["guidance"]) + + +def test_truncation_guidance_names_resume_token(): + call = ScriptedAthena() + call.queue("GetWorkGroup", _wg()) + call.queue("StartQueryExecution", ok_result({"QueryExecutionId": "q1"})) + call.queue("GetQueryExecution", _exec()) + call.queue( + "GetQueryResults", + _results([["a"], ["b"]], next_token="RESUME"), + ) + result = _run(call, sql="SELECT c1 FROM t", max_rows=1) + assert result["value"]["is_truncated"] is True + assert result["value"]["next_token"] == "RESUME" + assert any("RESUME" in g for g in result["guidance"]) + + +def test_statistics_are_surfaced(): + call = ScriptedAthena() + call.queue("GetWorkGroup", _wg()) + call.queue("StartQueryExecution", ok_result({"QueryExecutionId": "q1"})) + exec_body = ok_result( + { + "QueryExecution": { + "Status": {"State": "SUCCEEDED"}, + "StatementType": "DML", + "Statistics": { + "DataScannedInBytes": 1000, + "EngineExecutionTimeInMillis": 50, + "TotalExecutionTimeInMillis": 70, + "QueryQueueTimeInMillis": 5, + }, + } + } + ) + call.queue("GetQueryExecution", exec_body) + call.queue("GetQueryResults", _results([["a"]])) + result = _run(call, sql="SELECT c1 FROM t") + assert result["value"]["data_scanned_bytes"] == 1000 + assert result["value"]["engine_execution_time_ms"] == 50 + assert result["value"]["queue_time_ms"] == 5 + + +# --------------------------------------------------------------- never-raises guard + + +def test_internal_error_is_normalized_not_raised(): + # call_op raising an unexpected exception is caught and returned as data. + def exploding(operation, params): + raise RuntimeError("kaboom") + + result = run_athena_query( + exploding, + "SELECT 1", + workgroup="primary", + sleep=lambda _s: None, + clock=_fake_clock(), + ) + assert result["ok"] is False + assert result["code"] == "SERVICE_ERROR" + assert "internal error" in result["error"] + + +# --------------------------------------------------------------- failure classification + + +def test_failed_query_access_denied_is_permission_denied(): + call = ScriptedAthena() + call.queue("GetWorkGroup", _wg()) + call.queue("StartQueryExecution", ok_result({"QueryExecutionId": "q1"})) + call.queue( + "GetQueryExecution", + _exec(state="FAILED", reason="Access Denied writing results to S3"), + ) + result = _run(call, sql="SELECT 1 FROM t") + assert result["ok"] is False + assert result["code"] == PERMISSION_DENIED + + +def test_failed_query_not_supported_reason_is_invalid_input(): + # A terminal FAILED whose reason is "not supported" classifies as bad input. + call = ScriptedAthena() + call.queue("GetWorkGroup", _wg()) + call.queue("StartQueryExecution", ok_result({"QueryExecutionId": "q1"})) + call.queue( + "GetQueryExecution", + _exec( + state="FAILED", + reason="Queries of this type are not supported here", + ), + ) + result = _run(call, sql="SELECT 1 FROM t") + assert result["code"] == "INVALID_INPUT" + + +def test_failed_query_not_authorized_is_permission_denied(): + # Matches the 'not authorized' disjunct (no 'access denied' substring). + call = ScriptedAthena() + call.queue("GetWorkGroup", _wg()) + call.queue("StartQueryExecution", ok_result({"QueryExecutionId": "q1"})) + call.queue( + "GetQueryExecution", + _exec( + state="FAILED", + reason="User arn:...:role/x is not authorized to perform glue", + ), + ) + result = _run(call, sql="SELECT 1 FROM t") + assert result["code"] == PERMISSION_DENIED + + +def test_failed_query_system_error_is_retriable(): + # AthenaError ErrorCategory 1 is a SYSTEM error -> retriable SERVICE_ERROR. + call = ScriptedAthena() + call.queue("GetWorkGroup", _wg()) + call.queue("StartQueryExecution", ok_result({"QueryExecutionId": "q1"})) + call.queue( + "GetQueryExecution", + _exec(state="FAILED", reason="internal error", category=1), + ) + result = _run(call, sql="SELECT 1 FROM t") + assert result["code"] == "SERVICE_ERROR" + assert result["retriable"] is True + + +def test_failed_query_no_category_is_non_retriable_service_error(): + call = ScriptedAthena() + call.queue("GetWorkGroup", _wg()) + call.queue("StartQueryExecution", ok_result({"QueryExecutionId": "q1"})) + call.queue( + "GetQueryExecution", + _exec(state="FAILED", reason="something odd happened"), + ) + result = _run(call, sql="SELECT 1 FROM t") + assert result["code"] == "SERVICE_ERROR" + assert result["retriable"] is False + + +# --------------------------------------------------------------- location edge cases + + +def test_enforced_workgroup_without_caller_output_adds_no_note(): + # Enforced config but no caller output_location -> no "ignored" note. + call = ScriptedAthena() + call.queue( + "GetWorkGroup", + _wg(output_location="s3://enforced/", enforced=True), + ) + call.queue("StartQueryExecution", ok_result({"QueryExecutionId": "q1"})) + call.queue("GetQueryExecution", _exec()) + call.queue("GetQueryResults", _results([["a"]])) + result = _run(call, sql="SELECT 1 FROM t", workgroup="wg") + assert result["ok"] is True + assert not any("was ignored" in g for g in result["guidance"]) + start = dict(call.calls)["StartQueryExecution"] + assert "ResultConfiguration" not in start + + +def test_unreadable_workgroup_without_output_uses_unknown_source(): + # Workgroup unreadable AND no caller output: proceed with an "unknown" + # location source and let Athena surface any real "no location" error. + call = ScriptedAthena() + call.queue("GetWorkGroup", err_result("SERVICE_ERROR", "flaky")) + call.queue("StartQueryExecution", ok_result({"QueryExecutionId": "q1"})) + call.queue("GetQueryExecution", _exec()) + call.queue("GetQueryResults", _results([["a"]])) + result = _run(call, sql="SELECT 1 FROM t", workgroup="wg") + assert result["ok"] is True + assert result["value"]["output_location_source"] == "unknown" + start = dict(call.calls)["StartQueryExecution"] + assert "ResultConfiguration" not in start + + +def test_database_is_sent_in_query_execution_context(): + call = ScriptedAthena() + call.queue("GetWorkGroup", _wg()) + call.queue("StartQueryExecution", ok_result({"QueryExecutionId": "q1"})) + call.queue("GetQueryExecution", _exec()) + call.queue("GetQueryResults", _results([["a"]])) + result = _run(call, sql="SELECT 1 FROM t", workgroup="wg", database="mydb") + assert result["ok"] is True + start = dict(call.calls)["StartQueryExecution"] + assert start["QueryExecutionContext"]["Database"] == "mydb" + assert result["value"]["database"] == "mydb" + + +# --------------------------------------------------------------- misc result fields + + +def test_update_count_is_surfaced(): + call = ScriptedAthena() + call.queue("GetWorkGroup", _wg()) + call.queue("StartQueryExecution", ok_result({"QueryExecutionId": "q1"})) + call.queue("GetQueryExecution", _exec()) + body = ok_result( + { + "UpdateCount": 7, + "ResultSet": { + "ResultSetMetadata": { + "ColumnInfo": [{"Name": "c1", "Type": "varchar"}] + }, + "Rows": [ + {"Data": [{"VarCharValue": "c1"}]}, + {"Data": [{"VarCharValue": "a"}]}, + ], + }, + } + ) + call.queue("GetQueryResults", body) + result = _run(call, sql="SELECT c1 FROM t") + assert result["value"]["update_count"] == 7 + + +def test_non_dml_header_not_dropped_when_first_row_differs(): + # A non-DML result whose first row is NOT the column names keeps every row + # (no header is stripped). + call = ScriptedAthena() + call.queue("GetWorkGroup", _wg()) + call.queue("StartQueryExecution", ok_result({"QueryExecutionId": "q1"})) + call.queue("GetQueryExecution", _exec(stmt_type="UTILITY")) + body = ok_result( + { + "ResultSet": { + "ResultSetMetadata": { + "ColumnInfo": [{"Name": "col", "Type": "varchar"}] + }, + "Rows": [ + {"Data": [{"VarCharValue": "actual-data-row"}]}, + ], + } + } + ) + call.queue("GetQueryResults", body) + result = _run(call, sql="SHOW TABLES") + assert result["value"]["header_row_dropped"] is False + assert result["value"]["rows"] == [{"col": "actual-data-row"}] diff --git a/tests/unit/customizations/compositereadonly/__init__.py b/tests/unit/customizations/compositereadonly/__init__.py new file mode 100644 index 000000000000..654905217728 --- /dev/null +++ b/tests/unit/customizations/compositereadonly/__init__.py @@ -0,0 +1,12 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"). You +# may not use this file except in compliance with the License. A copy of +# the License is located at +# +# http://aws.amazon.com/apache2.0/ +# +# or in the "license" file accompanying this file. This file is +# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF +# ANY KIND, either express or implied. See the License for the specific +# language governing permissions and limitations under the License. diff --git a/tests/unit/customizations/compositereadonly/test_helpers.py b/tests/unit/customizations/compositereadonly/test_helpers.py new file mode 100644 index 000000000000..d5e2c596174a --- /dev/null +++ b/tests/unit/customizations/compositereadonly/test_helpers.py @@ -0,0 +1,382 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"). You +# may not use this file except in compliance with the License. A copy of +# the License is located at +# +# http://aws.amazon.com/apache2.0/ +# +# or in the "license" file accompanying this file. This file is +# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF +# ANY KIND, either express or implied. See the License for the specific +# language governing permissions and limitations under the License. +"""Unit tests for the shared read-only call bridge / emitter (helpers).""" + +import io +from types import SimpleNamespace + +import pytest +from botocore.exceptions import ( + BotoCoreError, + ClientError, + EndpointConnectionError, + NoCredentialsError, +) +from botocore.exceptions import ( + ConnectionError as BotoConnectionError, +) + +from awscli.customizations.exceptions import ParamValidationError +from awscli.customizations.compositereadonly import helpers + + +class FakeClient: + """A botocore-shaped client whose methods return or raise on demand.""" + + def __init__(self): + self._returns = {} + self._raises = {} + self.calls = [] + + def set_return(self, method, value): + self._returns[method] = value + + def set_raise(self, method, exc): + self._raises[method] = exc + + def __getattr__(self, method): + # __getattr__ only fires for names not found normally, so the private + # attributes above are untouched. + def _call(**params): + self.calls.append((method, params)) + if method in self._raises: + raise self._raises[method] + return self._returns.get(method, {"ok_marker": method}) + + return _call + + +def _client_error(code, status=400, message="boom"): + return ClientError( + { + "Error": {"Code": code, "Message": message}, + "ResponseMetadata": {"HTTPStatusCode": status}, + }, + "GetThing", + ) + + +def _parsed_globals(output="json", query=None): + return SimpleNamespace(output=output, query=query) + + +# ------------------------------------------------------------- assert_read_only + + +def test_assert_read_only_allows_read_verbs(): + # None of these should raise; read/lifecycle verbs are permitted. + for op in ( + "GetThing", + "DescribeThing", + "ListThings", + "StartQuery", + "StopQuery", + "GetQueryResults", + ): + helpers.assert_read_only(op) + + +def test_assert_read_only_refuses_mutations(): + for op in ( + "CreateThing", + "DeleteThing", + "PutObject", + "UpdateStack", + "TerminateInstances", + "RunInstances", + ): + with pytest.raises(helpers.ReadOnlyViolation): + helpers.assert_read_only(op) + + +def test_assert_read_only_exact_verb_name_is_refused(): + # The verb alone (no suffix) is still a mutation. + with pytest.raises(helpers.ReadOnlyViolation): + helpers.assert_read_only("Create") + + +def test_assert_read_only_respects_word_boundary(): + # "Created"/"Setting" start with a mutating prefix but the next char is + # lower-case, so they are NOT mistaken for Create/Set mutations. + helpers.assert_read_only("Createcountry") + helpers.assert_read_only("Setting") + helpers.assert_read_only("Removalist") + + +# ------------------------------------------------------------- result factories + + +def test_ok_result_shape(): + r = helpers.ok_result( + {"a": 1}, is_exhaustive=True, guidance=["g"], key="k" + ) + assert r["ok"] is True + assert r["value"] == {"a": 1} + assert r["is_exhaustive"] is True + assert r["guidance"] == ["g"] + assert r["key"] == "k" + + +def test_err_result_marks_retriable_codes(): + # THROTTLED is in the retriable set, so retriable is True even unset. + r = helpers.err_result(helpers.THROTTLED, "slow down") + assert r["ok"] is False + assert r["retriable"] is True + # A non-retriable code stays non-retriable unless forced. + r2 = helpers.err_result(helpers.NOT_FOUND, "missing") + assert r2["retriable"] is False + r3 = helpers.err_result(helpers.NOT_FOUND, "missing", retriable=True) + assert r3["retriable"] is True + + +def test_invalid_input_flags_param_error(): + r = helpers.invalid_input("bad", ["fix it"], value={"x": 1}) + assert r["code"] == helpers.INVALID_INPUT + assert r["param_error"] is True + assert r["value"] == {"x": 1} + + +# --------------------------------------------------------------- _map_client_error + + +def test_map_client_error_known_code(): + code, aws_code = helpers._map_client_error( + _client_error("AccessDenied", status=403) + ) + assert code == helpers.PERMISSION_DENIED + assert aws_code == "AccessDenied" + + +def test_map_client_error_status_fallbacks(): + cases = { + 404: helpers.NOT_FOUND, + 401: helpers.PERMISSION_DENIED, + 403: helpers.PERMISSION_DENIED, + 429: helpers.THROTTLED, + 500: helpers.SERVICE_ERROR, + 503: helpers.SERVICE_ERROR, + } + for status, expected in cases.items(): + code, aws_code = helpers._map_client_error( + _client_error("UnmappedCode", status=status) + ) + assert code == expected, status + assert aws_code == "UnmappedCode" + + +def test_map_client_error_unknown_status_is_service_error(): + code, _ = helpers._map_client_error( + _client_error("UnmappedCode", status=418) + ) + assert code == helpers.SERVICE_ERROR + + +# --------------------------------------------------------------------- call() + + +def test_call_refuses_mutation_before_dispatch(): + client = FakeClient() + with pytest.raises(helpers.ReadOnlyViolation): + helpers.call(client, "CreateThing", {}) + # The mutation guard runs before any client method is touched. + assert client.calls == [] + + +def test_call_success_returns_ok_result(): + client = FakeClient() + client.set_return("get_thing", {"Thing": "ok"}) + r = helpers.call(client, "GetThing", {"Id": "x"}) + assert r["ok"] is True + assert r["value"] == {"Thing": "ok"} + # Operation name was xform'd and params forwarded verbatim. + assert client.calls == [("get_thing", {"Id": "x"})] + + +def test_call_client_error_is_mapped(): + client = FakeClient() + client.set_raise("get_thing", _client_error("ThrottlingException", 429)) + r = helpers.call(client, "GetThing", {}) + assert r["ok"] is False + assert r["code"] == helpers.THROTTLED + assert r["retriable"] is True + assert r["aws_code"] == "ThrottlingException" + + +def test_call_no_credentials(): + client = FakeClient() + client.set_raise("get_thing", NoCredentialsError()) + r = helpers.call(client, "GetThing", {}) + assert r["code"] == helpers.CREDENTIALS_EXPIRED + assert r["retriable"] is True + + +def test_call_endpoint_connection_error_is_network(): + client = FakeClient() + client.set_raise( + "get_thing", EndpointConnectionError(endpoint_url="https://x") + ) + r = helpers.call(client, "GetThing", {}) + assert r["code"] == helpers.NETWORK_ERROR + assert r["retriable"] is True + + +def test_call_connection_error_is_network(): + client = FakeClient() + client.set_raise("get_thing", BotoConnectionError(error="reset")) + r = helpers.call(client, "GetThing", {}) + assert r["code"] == helpers.NETWORK_ERROR + + +def test_call_botocore_error_is_service_error(): + client = FakeClient() + client.set_raise("get_thing", BotoCoreError()) + r = helpers.call(client, "GetThing", {}) + assert r["code"] == helpers.SERVICE_ERROR + assert r["retriable"] is False + + +# --------------------------------------------------------------- _render_payload + + +def test_render_payload_dict_value_folds_in_flags(): + result = helpers.ok_result( + {"rows": []}, is_exhaustive=True, guidance=["g1"] + ) + payload = helpers._render_payload(result) + assert payload["rows"] == [] + assert payload["is_exhaustive"] is True + assert payload["guidance"] == ["g1"] + + +def test_render_payload_scalar_value_wrapped(): + payload = helpers._render_payload(helpers.ok_result(42)) + assert payload == {"result": 42} + + +def test_render_payload_none_value_is_empty_dict(): + payload = helpers._render_payload(helpers.ok_result(None)) + assert payload == {} + + +def test_render_payload_failure_includes_error_and_retriable(): + result = helpers.err_result( + helpers.THROTTLED, "slow", value={"partial": True} + ) + payload = helpers._render_payload(result) + assert payload["partial"] is True + assert payload["error"] == "slow" + assert payload["code"] == helpers.THROTTLED + assert payload["retriable"] is True + + +def test_render_payload_does_not_override_existing_value_keys(): + # When value already carries error/code, _render_payload must not clobber. + result = helpers.err_result( + helpers.NOT_FOUND, "outer", value={"error": "inner", "code": "INNER"} + ) + payload = helpers._render_payload(result) + assert payload["error"] == "inner" + assert payload["code"] == "INNER" + + +# ------------------------------------------------------------------ emit_result + + +def test_emit_result_success_writes_payload_rc0(): + out, err = io.StringIO(), io.StringIO() + result = helpers.ok_result({"rows": [1, 2]}) + rc = helpers.emit_result( + "x", result, _parsed_globals(), session=None, stdout=out, stderr=err + ) + assert rc == helpers.RC_SUCCESS + assert '"rows"' in out.getvalue() + assert err.getvalue() == "" + + +def test_emit_result_operational_failure_rc255_and_stderr(): + out, err = io.StringIO(), io.StringIO() + result = helpers.err_result(helpers.NOT_FOUND, "missing thing") + rc = helpers.emit_result( + "x", result, _parsed_globals(), session=None, stdout=out, stderr=err + ) + assert rc == helpers.RC_OPERATIONAL_ERROR + # Payload still printed to stdout so --query keeps working. + assert "NOT_FOUND" in err.getvalue() + assert "missing thing" in err.getvalue() + + +def test_emit_result_param_error_raises_param_validation(): + result = helpers.invalid_input("bad flag") + with pytest.raises(ParamValidationError): + helpers.emit_result( + "x", + result, + _parsed_globals(), + session=None, + stdout=io.StringIO(), + stderr=io.StringIO(), + ) + + +def test_emit_result_output_falls_back_to_session_config(): + # parsed_globals.output is None -> read from session config variable. + out = io.StringIO() + session = SimpleNamespace( + get_config_variable=lambda name: "json" if name == "output" else None + ) + result = helpers.ok_result({"k": "v"}) + rc = helpers.emit_result( + "x", + result, + _parsed_globals(output=None), + session=session, + stdout=out, + stderr=io.StringIO(), + ) + assert rc == helpers.RC_SUCCESS + assert '"k"' in out.getvalue() + + +def test_emit_result_defaults_to_std_streams(monkeypatch): + captured_out = io.StringIO() + captured_err = io.StringIO() + monkeypatch.setattr(helpers.sys, "stdout", captured_out) + monkeypatch.setattr(helpers.sys, "stderr", captured_err) + result = helpers.err_result(helpers.SERVICE_ERROR, "boom") + rc = helpers.emit_result("x", result, _parsed_globals(), session=None) + assert rc == helpers.RC_OPERATIONAL_ERROR + assert "boom" in captured_err.getvalue() + + +# -------------------------------------------------------------- create_read_client + + +def test_create_read_client_delegates(monkeypatch): + seen = {} + + def fake_create(session, service_name, parsed_globals): + seen["args"] = (session, service_name, parsed_globals) + return "a-client" + + monkeypatch.setattr( + helpers, "create_client_from_parsed_globals", fake_create + ) + pg = _parsed_globals() + client = helpers.create_read_client("sess", "athena", pg) + assert client == "a-client" + assert seen["args"] == ("sess", "athena", pg) + + +def test_raise_invalid_raises_param_validation(): + with pytest.raises(ParamValidationError): + helpers.raise_invalid("nope") diff --git a/tests/unit/customizations/compositereadonly/test_output.py b/tests/unit/customizations/compositereadonly/test_output.py new file mode 100644 index 000000000000..db38fefbea2c --- /dev/null +++ b/tests/unit/customizations/compositereadonly/test_output.py @@ -0,0 +1,92 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"). You +# may not use this file except in compliance with the License. A copy of +# the License is located at +# +# http://aws.amazon.com/apache2.0/ +# +# or in the "license" file accompanying this file. This file is +# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF +# ANY KIND, either express or implied. See the License for the specific +# language governing permissions and limitations under the License. +"""Unit tests for the compositereadonly output caps (truncation + series capping).""" + +from awscli.customizations.compositereadonly import output + +# --------------------------------------------------------------- truncate_field + + +def test_truncate_field_keeps_head_and_flags(): + value, truncated = output.truncate_field("x" * 100, 20) + assert truncated is True + assert len(value) <= 20 + assert value.endswith("...[truncated]") + # The kept portion is the head of the original string. + assert value.startswith("x") + + +def test_truncate_field_under_limit_untouched(): + value, truncated = output.truncate_field("short", 20) + assert truncated is False + assert value == "short" + + +def test_truncate_field_exact_length_untouched(): + text = "x" * 20 + value, truncated = output.truncate_field(text, 20) + assert truncated is False + assert value == text + + +def test_truncate_field_disabled_with_nonpositive_max(): + # max_chars <= 0 disables truncation entirely. + value, truncated = output.truncate_field("x" * 100, 0) + assert truncated is False + assert len(value) == 100 + value2, truncated2 = output.truncate_field("x" * 100, -5) + assert truncated2 is False + assert len(value2) == 100 + + +def test_truncate_field_passes_through_non_strings(): + for val in (123, None, ["a"], {"k": "v"}): + result, truncated = output.truncate_field(val, 5) + assert result == val + assert truncated is False + + +def test_truncate_field_tiny_max_smaller_than_suffix(): + # When max_chars is smaller than the suffix, the head is empty but the + # function still returns the suffix (max(0, ...) guards the negative slice). + value, truncated = output.truncate_field("abcdef", 3) + assert truncated is True + assert value == "...[truncated]" + + +# -------------------------------------------------------------------- cap_series + + +def test_cap_series_under_limit_untouched(): + kept, dropped = output.cap_series([1, 2, 3], 5) + assert kept == [1, 2, 3] + assert dropped == 0 + + +def test_cap_series_truncates_and_reports_dropped(): + kept, dropped = output.cap_series([1, 2, 3, 4, 5], 2) + assert kept == [1, 2] + assert dropped == 3 + + +def test_cap_series_exact_length_untouched(): + kept, dropped = output.cap_series([1, 2], 2) + assert kept == [1, 2] + assert dropped == 0 + + +def test_cap_series_negative_max_disables_cap(): + items = [1, 2, 3] + kept, dropped = output.cap_series(items, -1) + assert kept == items + assert dropped == 0 diff --git a/tests/unit/customizations/compositereadonly/test_poll.py b/tests/unit/customizations/compositereadonly/test_poll.py new file mode 100644 index 000000000000..2c57b57e2559 --- /dev/null +++ b/tests/unit/customizations/compositereadonly/test_poll.py @@ -0,0 +1,116 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"). You +# may not use this file except in compliance with the License. A copy of +# the License is located at +# +# http://aws.amazon.com/apache2.0/ +# +# or in the "license" file accompanying this file. This file is +# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF +# ANY KIND, either express or implied. See the License for the specific +# language governing permissions and limitations under the License. +"""Unit tests for the shared BoundedPoller (deterministic, injected clock).""" + +from awscli.customizations.compositereadonly.poll import ( + MAX_POLL_INTERVAL, + BoundedPoller, +) + + +class FakeClock: + """A manually advanced monotonic clock.""" + + def __init__(self): + self.t = 0.0 + + def __call__(self): + return self.t + + def advance(self, dt): + self.t += dt + + +def _poller(clock, sleeps, **kwargs): + kwargs.setdefault("interval", 1.0) + kwargs.setdefault("max_wait", 10.0) + kwargs.setdefault("max_polls", 100) + return BoundedPoller( + clock=clock, sleep=lambda s: sleeps.append(s), **kwargs + ) + + +def test_record_poll_counts_up(): + clock = FakeClock() + poller = _poller(clock, []) + assert poller.record_poll() == 1 + assert poller.record_poll() == 2 + assert poller.polls == 2 + + +def test_exhausted_on_poll_count_cap(): + clock = FakeClock() + poller = _poller(clock, [], max_polls=2) + poller.record_poll() + assert poller.exhausted() is False + poller.record_poll() + # The poll-count cap is reached even though no wall time elapsed. + assert poller.exhausted() is True + + +def test_exhausted_on_deadline(): + clock = FakeClock() + poller = _poller(clock, [], max_wait=5.0) + assert poller.exhausted() is False + clock.advance(5.0) + # Deadline reached -> exhausted regardless of poll count. + assert poller.exhausted() is True + + +def test_remaining_decreases_with_clock(): + clock = FakeClock() + poller = _poller(clock, [], max_wait=10.0) + assert poller.remaining() == 10.0 + clock.advance(3.0) + assert poller.remaining() == 7.0 + + +def test_backoff_grows_geometrically_and_caps(): + clock = FakeClock() + sleeps = [] + poller = _poller( + clock, + sleeps, + interval=1.0, + max_wait=1000.0, + growth=2.0, + max_interval=MAX_POLL_INTERVAL, + ) + for _ in range(6): + poller.sleep_and_backoff() + # 1, 2, 4, then capped at MAX_POLL_INTERVAL (5.0). + assert sleeps[0] == 1.0 + assert sleeps[1] == 2.0 + assert sleeps[2] == 4.0 + assert all(s <= MAX_POLL_INTERVAL for s in sleeps) + assert sleeps[-1] == MAX_POLL_INTERVAL + + +def test_sleep_is_clamped_to_remaining_time(): + clock = FakeClock() + sleeps = [] + # interval (10s) exceeds the remaining budget (3s); sleep is clamped so the + # poller never sleeps past its own deadline. + poller = _poller(clock, sleeps, interval=10.0, max_wait=3.0) + poller.sleep_and_backoff() + assert sleeps[0] == 3.0 + + +def test_sleep_never_negative_past_deadline(): + clock = FakeClock() + sleeps = [] + poller = _poller(clock, sleeps, interval=10.0, max_wait=5.0) + clock.advance(8.0) # already past the deadline + poller.sleep_and_backoff() + # remaining() would be negative; the clamp floors the sleep at 0. + assert sleeps[0] == 0.0 From 034c80a051410caf00ef275a8ceae5a52bf783d3 Mon Sep 17 00:00:00 2001 From: Claudiu Popa Date: Thu, 29 Oct 2026 11:23:26 +0100 Subject: [PATCH 2/4] Add Logs Insights query command Generated by AI tools, and reviewed by Claudiu Popa Co-authored-by: Baris Kurt --- .../customizations/logs/run_insights_query.py | 1240 ++++++++++++++++ awscli/handlers_registry.py | 1 + .../logs/test_run_insights_query.py | 111 ++ .../logs/test_run_insights_query.py | 1303 +++++++++++++++++ 4 files changed, 2655 insertions(+) create mode 100644 awscli/customizations/logs/run_insights_query.py create mode 100644 tests/functional/logs/test_run_insights_query.py create mode 100644 tests/unit/customizations/logs/test_run_insights_query.py diff --git a/awscli/customizations/logs/run_insights_query.py b/awscli/customizations/logs/run_insights_query.py new file mode 100644 index 000000000000..8c6bba17cce9 --- /dev/null +++ b/awscli/customizations/logs/run_insights_query.py @@ -0,0 +1,1240 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"). You +# may not use this file except in compliance with the License. A copy of +# the License is located at +# +# http://aws.amazon.com/apache2.0/ +# +# or in the "license" file accompanying this file. This file is +# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF +# ANY KIND, either express or implied. See the License for the specific +# language governing permissions and limitations under the License. +"""``aws logs run-insights-query`` -- run one Logs Insights query to completion. + +Callers otherwise write their own ``start-query`` -> ``get-query-results`` polling +loops and get recurring things wrong: a blank ``queryId`` passed to +``get-query-results`` after a failed start, a window that falls outside a group's +retention, loops that stop on a partial window or the 10,000-row cap (so counts +come out low), and missed log groups from a guessed name or prefix. + +Procedure (one call, read-only): + +1. ``logs:DescribeLogGroups`` (all pages) resolves each input as an exact name, a + prefix (trailing ``*`` or when no exact name exists), or a case-insensitive + ``logGroupNamePattern`` substring. The window is bounded to the queryable + range: end to now, start to the earliest **retention** edge. Groups with no + data in the window are skipped and reported. +2. The query string is validated locally before dispatch. On a server + ``MalformedQueryException`` the parser line/position and offending token are + returned. +3. ``logs:StartQuery`` (batches of at most 50 groups, retrying on throttling); + the ``queryId`` is checked non-blank; ``logs:GetQueryResults`` is polled with + the shared bounded poller. On the poll cap ``logs:StopQuery`` is called. +4. For a non-aggregate query that hits the row cap, the window is split in half, + each half re-run, and results merged, de-duplicating on ``@ptr``. Any cap hit + is reported through ``is_exhaustive=False`` and ``partial_reasons``. + +Notes on retention and output size: the queryable-start bound uses retention +only -- a log group's *creation time* does not bound a Logs Insights query, so it +is intentionally NOT used to bound the window. The default row output is small +(``--max-rows`` ~50) and each ``@message`` is truncated to a bounded width +(``--max-message-length``); both bounds are reported so nothing is lost silently. +""" + +from __future__ import annotations + +import re +import time +from datetime import datetime, timezone +from typing import Any + +from awscli.customizations.commands import BasicCommand +from awscli.customizations.compositereadonly.helpers import ( + CONFLICT, + CREDENTIALS_EXPIRED, + INVALID_INPUT, + NOT_FOUND, + PERMISSION_DENIED, + QUOTA_EXCEEDED, + SERVICE_ERROR, + THROTTLED, + TIMEOUT, + VALIDATION_ERROR, + call, + check_non_empty_str, + check_positive_int, + check_positive_number, + create_read_client, + emit_result, + err_result, + invalid_input, + ok_result, +) +from awscli.customizations.compositereadonly.output import ( + DEFAULT_MESSAGE_MAX_CHARS, + truncate_field, +) +from awscli.customizations.compositereadonly.poll import drive_poll + +_ROW_CAP = 10_000 +_MAX_GROUPS_PER_QUERY = 50 +_MAX_QUERY_CHARS = 10_000 +_DESCRIBE_PAGE_CAP = 100 +_START_RETRIES = 5 +_THROTTLE_RETRIES = 5 +_DAY_MS = 86_400_000 + +_TERMINAL = frozenset( + {"Complete", "Failed", "Cancelled", "Timeout", "Unknown"} +) +_KNOWN_COMMANDS = frozenset( + { + "fields", + "display", + "filter", + "filterindex", + "stats", + "sort", + "limit", + "parse", + "dedup", + "unmask", + "pattern", + "diff", + "anomaly", + "unnest", + "lookup", + "source", + } +) +_AGGREGATE_COMMANDS = frozenset({"stats", "pattern", "diff", "anomaly"}) +_GROUP_NAME_RE = re.compile(r"^[\.\-_/#A-Za-z0-9]+$") +_ARN_RE = re.compile( + r"^arn:aws[\w-]*:logs:[^:]*:\d*:log-group:([^:]+)(?::\*)?$" +) +_LINE_POS_RE = re.compile( + r"line\s+(\d+)\s*(?:and|,)?\s*position\s+(\d+)", re.IGNORECASE +) +_POS_RE = re.compile(r"position\s+(\d+)", re.IGNORECASE) +_THROTTLE_CODES = frozenset( + { + "LimitExceededException", + "ThrottlingException", + "ServiceUnavailableException", + "TooManyRequestsException", + } +) +_EXPIRED_CODES = frozenset( + {"ExpiredTokenException", "ExpiredToken", "RequestExpired"} +) + + +# --------------------------------------------------------------------------- small helpers + + +def _iso(epoch_seconds: int) -> str: + return ( + datetime.fromtimestamp(epoch_seconds, tz=timezone.utc) + .isoformat() + .replace("+00:00", "Z") + ) + + +def _iso_ms(epoch_ms: int) -> str: + return ( + datetime.fromtimestamp(epoch_ms / 1000, tz=timezone.utc) + .isoformat() + .replace("+00:00", "Z") + ) + + +def _ceil_s(ms: int) -> int: + return -(-ms // 1000) + + +def _parse_time(value: Any, name: str) -> tuple[datetime | None, str | None]: + if isinstance(value, datetime): + dt = value + elif isinstance(value, str) and value.strip(): + try: + dt = datetime.fromisoformat( + value.strip().replace("Z", "+00:00").replace("z", "+00:00") + ) + except ValueError: + return ( + None, + f"{name}={value!r} is not ISO-8601; use e.g. '2024-01-15T12:00:00Z'.", + ) + else: + return ( + None, + f"{name} is required as an ISO-8601 string, e.g. '2024-01-15T12:00:00Z'.", + ) + if dt.tzinfo is None: + dt = dt.replace(tzinfo=timezone.utc) + return dt, None + + +def _is_throttle(res: dict[str, Any]) -> bool: + return ( + res.get("code") in (THROTTLED, QUOTA_EXCEEDED) + or (res.get("aws_code") or "") in _THROTTLE_CODES + ) + + +def _propagate( + res: dict[str, Any], + context: str, + permission: str, + extra: list[str] | None = None, +) -> dict[str, Any]: + """Return an AWS failure as data, adding the fix the caller needs.""" + code = res.get("code") or SERVICE_ERROR + guidance = list(res.get("guidance") or []) + if ( + code == CREDENTIALS_EXPIRED + or (res.get("aws_code") or "") in _EXPIRED_CODES + ): + code = CREDENTIALS_EXPIRED + guidance.append( + "AWS credentials expired: refresh the session, then re-run this command (it is idempotent)." + ) + elif code == PERMISSION_DENIED: + guidance.append( + f"The caller lacks {permission}; use a role that grants it." + ) + guidance.extend(extra or []) + return err_result( + code, + f"{context}: {res.get('error')}", + aws_code=res.get("aws_code"), + retriable=bool(res.get("retriable")) + or code in (THROTTLED, CREDENTIALS_EXPIRED), + guidance=guidance, + ) + + +# --------------------------------------------------------------------------- query validation + + +def _blank_comments(query: str) -> str: + """Replace `#` comment lines with spaces so offsets into the original string stay valid.""" + return "\n".join( + " " * len(ln) if ln.lstrip().startswith("#") else ln + for ln in query.split("\n") + ) + + +def _opens_regex(text: str, seg_start: int, i: int) -> bool: + prev = text[seg_start:i].rstrip() + if not prev: + return True + words = prev.split() + return ( + prev[-1] in "(,=~!" + or words[-1].lower() in {"like", "match", "rlike"} + or words[0].lower() == "parse" + ) + + +def _split_pipeline(query: str) -> list[tuple[int, str]]: + """Split on top-level `|`, ignoring pipes inside quotes, backticks, /regex/ literals and parentheses.""" + text = _blank_comments(query) + segments: list[tuple[int, str]] = [] + seg_start, i, n = 0, 0, len(text) + quote: str | None = None + in_regex = False + depth = 0 + while i < n: + ch = text[i] + if quote is not None or in_regex: + if ch == "\\": + i += 2 + continue + if quote is not None and ch == quote: + quote = None + elif in_regex and ch == "/": + in_regex = False + elif ch in "\"'`": + quote = ch + elif ch == "/" and _opens_regex(text, seg_start, i): + in_regex = True + elif ch == "(": + depth += 1 + elif ch == ")": + depth = max(0, depth - 1) + elif ch == "|" and depth == 0: + segments.append((seg_start, text[seg_start:i])) + seg_start = i + 1 + i += 1 + segments.append((seg_start, text[seg_start:])) + return segments + + +def _analyze_query(query: str) -> dict[str, Any]: + """Analyse the pipeline's commands for guidance -- advisory only, never fatal. + + Returns an analysis dict: the recognized ``commands`` list, ``is_aggregate``, + ``explicit_limit``, an ``unknown_commands`` flag and the ``unknown_tokens`` + that triggered it. A locally-unrecognized command (or an empty/stray stage) + is recorded rather than rejected: the server validates the query, so we never + refuse locally what the server might accept -- a MalformedQueryException from + StartQuery handles a query that is truly invalid. + """ + commands: list[str] = [] + unknown_tokens: list[str] = [] + for offset, seg in _split_pipeline(query): + stripped = seg.strip() + if not stripped: + # A stray '|' / empty stage: record it, do not reject locally. + unknown_tokens.append("|") + continue + m = re.match(r"[A-Za-z_]+", stripped) + token = m.group(0) if m else stripped.split()[0] + if not m or token.lower() not in _KNOWN_COMMANDS: + # An unrecognized leading token: the server may still accept it, so + # flag it for guidance instead of blocking the query. + unknown_tokens.append(token) + continue + commands.append(token.lower()) + return { + "commands": commands, + "is_aggregate": any(c in _AGGREGATE_COMMANDS for c in commands), + "explicit_limit": "limit" in commands, + "unknown_commands": bool(unknown_tokens), + "unknown_tokens": unknown_tokens, + } + + +def _compile_error_detail(res: dict[str, Any], query: str) -> dict[str, Any]: + """Extract parser line/position and the offending token from a MalformedQueryException.""" + detail: dict[str, Any] = {"source": "server", "message": res.get("error")} + offset: int | None = None + body = res.get("value") if isinstance(res.get("value"), dict) else {} + loc = (body.get("queryCompileError") or {}).get("location") or {} + msg = res.get("error") or "" + m_line = _LINE_POS_RE.search(msg) + m_pos = _POS_RE.search(msg) + if isinstance(loc.get("start"), int): + offset = loc["start"] + detail["position"] = offset + elif m_line: + line, col = int(m_line.group(1)), int(m_line.group(2)) + lines = query.split("\n") + offset = sum(len(ln) + 1 for ln in lines[: max(line - 1, 0)]) + col + detail.update({"line": line, "position": col}) + elif m_pos: + offset = int(m_pos.group(1)) + detail["position"] = offset + token = "" + if offset is not None and 0 <= offset < len(query): + mt = re.compile(r"\S+").search(query, offset) + token = mt.group(0) if mt else "" + detail["offset"] = offset + detail["token"] = token + return detail + + +# --------------------------------------------------------------------------- log group resolution + + +def _normalize_entry(raw: Any) -> tuple[str, bool] | None: + if not isinstance(raw, str) or not raw.strip(): + return None + entry = raw.strip() + m = _ARN_RE.match(entry) + if m: + entry = m.group(1) + is_prefix = entry.endswith("*") + name = entry.rstrip("*") + if not name or not _GROUP_NAME_RE.match(name): + return None + return name, is_prefix + + +def _describe( + call_op, selector: dict[str, Any] +) -> tuple[list[dict[str, Any]], bool, dict[str, Any] | None]: + """DescribeLogGroups, following nextToken to the end (bounded). Returns (groups, complete, error).""" + groups: list[dict[str, Any]] = [] + token: str | None = None + for _ in range(_DESCRIBE_PAGE_CAP): + params = dict(selector) + if token: + params["nextToken"] = token + res = call_op("DescribeLogGroups", params) + if not res["ok"]: + return groups, False, res + body = res.get("value") or {} + groups.extend( + g + for g in (body.get("logGroups") or []) + if isinstance(g, dict) and g.get("logGroupName") + ) + token = body.get("nextToken") + if not token: + return groups, True, None + return groups, False, None + + +def _resolve_groups( + call_op, entries: list[tuple[str, bool]] +) -> tuple[ + list[dict[str, Any]], + dict[str, dict[str, Any]], + list[str], + bool, + dict[str, Any] | None, +]: + resolution: list[dict[str, Any]] = [] + groups: dict[str, dict[str, Any]] = {} + notes: list[str] = [] + complete = True + for name, is_prefix in entries: + found, done, err = _describe(call_op, {"logGroupNamePrefix": name}) + if err is not None: + return ( + [], + {}, + [], + False, + _propagate( + err, + f"DescribeLogGroups(prefix={name!r})", + "logs:DescribeLogGroups", + ), + ) + complete = complete and done + exact = [g for g in found if g["logGroupName"] == name] + if is_prefix: + chosen, how = found, "prefix" + elif exact: + chosen, how = exact, "exact" + siblings = len(found) - len(exact) + if siblings: + notes.append( + f"{siblings} other log group(s) share the prefix {name!r}; " + f"pass {name + '*'!r} to include them." + ) + elif found: + chosen, how = found, "prefix" + notes.append( + f"No log group is named exactly {name!r}; expanded it as a prefix to {len(found)} group(s)." + ) + else: + found, done, err = _describe( + call_op, {"logGroupNamePattern": name} + ) + if err is not None and err.get("code") not in ( + INVALID_INPUT, + VALIDATION_ERROR, + ): + return ( + [], + {}, + [], + False, + _propagate( + err, + f"DescribeLogGroups(pattern={name!r})", + "logs:DescribeLogGroups", + ), + ) + chosen = found if err is None else [] + if err is None: + complete = complete and done + how = "pattern" if chosen else "none" + if chosen: + notes.append( + f"{name!r} is not a log group name or prefix; matched {len(chosen)} group(s) " + "containing it: " + + ", ".join(g["logGroupName"] for g in chosen[:5]) + + (" ..." if len(chosen) > 5 else "") + ) + for g in chosen: + groups.setdefault(g["logGroupName"], g) + resolution.append( + { + "input": name + ("*" if is_prefix else ""), + "match": how, + "log_groups": [g["logGroupName"] for g in chosen], + } + ) + return resolution, groups, notes, complete, None + + +# --------------------------------------------------------------------------- one query segment + + +def _run_segment( + call_op, + names: list[str], + start_s: int, + end_s: int, + query: str, + limit: int | None, + *, + sleep, + clock, + poll_interval: float, + max_polls: int, + timeout_seconds: float, +) -> dict[str, Any]: + """Run one StartQuery/poll cycle. Returns segment data, or an err_result (ok=False).""" + params: dict[str, Any] = { + "logGroupNames": names, + "startTime": start_s, + "endTime": end_s, + "queryString": query, + } + if limit is not None: + params["limit"] = limit + attempt = 0 + while True: + res = call_op("StartQuery", params) + if res["ok"]: + break + if _is_throttle(res): + if attempt < _START_RETRIES: + sleep(min(max(poll_interval, 1.0) * (2**attempt), 30.0)) + attempt += 1 + continue + return err_result( + QUOTA_EXCEEDED, + f"StartQuery still throttled after {_START_RETRIES} retries: {res.get('error')}", + aws_code=res.get("aws_code"), + retriable=True, + guidance=[ + "Logs Insights limits concurrent queries per account/region; wait for running queries " + "to finish (logs:DescribeQueries) and retry." + ], + ) + aws_code = res.get("aws_code") or "" + if aws_code == "MalformedQueryException": + detail = _compile_error_detail(res, query) + where = f"line {detail.get('line', 1)}, position {detail.get('position')}" + # Server-discovered malformed query: operational error (not a pre-dispatch param error). + return err_result( + INVALID_INPUT, + f"MalformedQueryException at {where}: token {detail['token']!r}", + value=detail, + guidance=[ + f"Fix the query near {detail['token']!r} ({where}) and re-run; nothing was executed.", + "Check quoting of regexes (/.../), string literals and field names with special " + "characters (wrap them in backticks).", + ], + ) + if aws_code == "InvalidParameterException": + return err_result( + INVALID_INPUT, + f"StartQuery rejected the parameters: {res.get('error')}", + aws_code=aws_code, + guidance=[ + f"Window used was {_iso(start_s)}..{_iso(end_s)} on {len(names)} group(s); " + "narrow the window or check the group names in log_group_resolution." + ], + ) + if aws_code == "ResourceNotFoundException": + return err_result( + NOT_FOUND, + f"StartQuery: {res.get('error')}", + aws_code=aws_code, + guidance=[ + "A log group disappeared between DescribeLogGroups and StartQuery; re-run." + ], + ) + return _propagate(res, "StartQuery", "logs:StartQuery") + + qid = (res.get("value") or {}).get("queryId") + if not isinstance(qid, str) or not qid.strip(): + return err_result( + SERVICE_ERROR, + "StartQuery returned no queryId; refusing to poll GetQueryResults with a blank id.", + retriable=True, + guidance=[ + "Re-run the command; if it persists, check logs:StartQuery permissions." + ], + ) + + body: dict[str, Any] = {} + status = "Unknown" + + def _terminal(res: dict[str, Any]) -> str | None: + nonlocal body, status + body = res.get("value") or {} + status = str(body.get("status") or "Unknown") + return status if status in _TERMINAL else None + + def _on_timeout() -> None: + # StopQuery cancels the query lifecycle; it does not mutate your data. + call_op("StopQuery", {"queryId": qid}) + + poll_result = drive_poll( + call_op, + "GetQueryResults", + {"queryId": qid}, + terminal=_terminal, + is_retriable=_is_throttle, + on_timeout=_on_timeout, + interval=poll_interval, + max_wait=timeout_seconds, + max_polls=max_polls, + max_transient=_THROTTLE_RETRIES, + clock=clock, + sleep=sleep, + key=qid, + ) + polls = poll_result.get("polls_count", 0) + + if not poll_result["ok"]: + if poll_result.get("code") == TIMEOUT: + return err_result( + TIMEOUT, + f"Query {qid} still {status} after {polls} polls; stopped it.", + retriable=True, + is_exhaustive=False, + value={ + "query_id": qid, + "status": status, + "polls": polls, + "window": {"start": _iso(start_s), "end": _iso(end_s)}, + "log_groups": names, + }, + guidance=[ + "Raise --timeout-seconds, narrow the time window, or add a selective " + "'filter' stage before 'stats'." + ], + ) + return _propagate( + poll_result, + f"GetQueryResults(queryId={qid})", + "logs:GetQueryResults", + [ + f"Query {qid} may still be running; call logs:GetQueryResults with it later." + ], + ) + + if status != "Complete": + code = {"Timeout": TIMEOUT, "Cancelled": CONFLICT}.get( + status, SERVICE_ERROR + ) + return err_result( + code, + f"Query {qid} ended with status {status}.", + retriable=status in ("Timeout", "Unknown"), + value={ + "query_id": qid, + "status": status, + "window": {"start": _iso(start_s), "end": _iso(end_s)}, + }, + guidance=[ + "Logs Insights queries time out after 60 minutes; narrow the window or filter earlier." + if status == "Timeout" + else "Re-run; if it fails again, simplify the query." + ], + ) + + rows: list[dict[str, Any]] = [] + for raw in body.get("results") or []: + if isinstance(raw, list): + rows.append( + { + c.get("field"): c.get("value") + for c in raw + if isinstance(c, dict) and c.get("field") + } + ) + stats = body.get("statistics") or {} + # Segment data (no "ok" key) distinguishes success from an err_result. + return { + "query_id": qid, + "rows": rows, + "statistics": { + k: float(v) + for k, v in stats.items() + if isinstance(v, (int, float)) + }, + "polls": polls, + } + + +# --------------------------------------------------------------------------- the command + + +def run_insights_query( + call_op, + log_group_names_or_prefix: list[str] | str, + query_string: str, + start_time: str | datetime, + end_time: str | datetime, + *, + max_splits: int = 6, + max_queries: int = 64, + max_rows: int = 50, + max_message_length: int = DEFAULT_MESSAGE_MAX_CHARS, + poll_interval: float = 1.0, + max_polls: int = 600, + timeout_seconds: float = 900.0, + now: str | datetime | None = None, + sleep=None, + clock=None, +) -> dict[str, Any]: + """Run one Logs Insights query over one or more log groups to completion. + + ``call_op(operation, params)`` is the read-only call bridge bound to a Logs + client. Returns the plain-dict result convention; never raises. The window is + bounded to retention (NOT creation time). Row output is capped at ``max_rows`` + and each ``@message`` is truncated to ``max_message_length``; both bounds are + reported. + """ + sleep = time.sleep if sleep is None else sleep + clock = time.monotonic if clock is None else clock + try: + return _run( + call_op, + log_group_names_or_prefix, + query_string, + start_time, + end_time, + max_splits, + max_queries, + max_rows, + max_message_length, + poll_interval, + max_polls, + timeout_seconds, + now, + sleep, + clock, + ) + except Exception as exc: # noqa: BLE001 - contract: never raise + return err_result( + SERVICE_ERROR, + f"run-insights-query internal error: {exc!r}", + guidance=[ + "Unexpected internal error; report it with the inputs used." + ], + ) + + +def _validate( + groups_in: list[str] | str, + query_string: str, + start_time: Any, + end_time: Any, + now: Any, + ints: dict[str, Any], + max_splits: Any, + floats: dict[str, Any], +) -> tuple[ + list[tuple[str, bool]], + datetime | None, + datetime | None, + datetime | None, + list[str], +]: + problems: list[str] = [] + raw_entries = ( + [groups_in] if isinstance(groups_in, str) else list(groups_in or []) + ) + entries: list[tuple[str, bool]] = [] + if not raw_entries: + problems.append( + "log_group_names is empty; pass at least one log group name, prefix* or ARN." + ) + for raw in raw_entries: + norm = _normalize_entry(raw) + if norm is None: + problems.append( + f"{raw!r} is not a valid log group name/prefix (allowed: letters, digits, . - _ / #, " + "optional trailing '*', or a log-group ARN)." + ) + elif norm not in entries: + entries.append(norm) + q_err = check_non_empty_str("query_string", query_string) + if q_err is not None: + problems.append(q_err) + elif len(query_string) > _MAX_QUERY_CHARS: + problems.append( + f"query_string is {len(query_string)} chars; Logs Insights accepts at most {_MAX_QUERY_CHARS}." + ) + start_dt, err_s = _parse_time(start_time, "start_time") + end_dt, err_e = _parse_time(end_time, "end_time") + problems.extend(e for e in (err_s, err_e) if e) + if start_dt and end_dt and start_dt >= end_dt: + problems.append( + f"start_time ({start_dt.isoformat()}) must be before end_time ({end_dt.isoformat()})." + ) + now_dt: datetime | None = datetime.now(timezone.utc) + if now is not None: + now_dt, err_n = _parse_time(now, "now") + if err_n: + problems.append(err_n) + for pname, pval in ints.items(): + err = check_positive_int(pname, pval) + if err is not None: + problems.append(err) + err = check_positive_int("max_splits", max_splits, lo=0) + if err is not None: + problems.append(err) + for pname, pval in floats.items(): + err = check_positive_number(pname, pval) + if err is not None: + problems.append(err) + return entries, start_dt, end_dt, now_dt, problems + + +def _run( + call_op, + groups_in: list[str] | str, + query_string: str, + start_time: str | datetime, + end_time: str | datetime, + max_splits: int, + max_queries: int, + max_rows: int, + max_message_length: int, + poll_interval: float, + max_polls: int, + timeout_seconds: float, + now: str | datetime | None, + sleep, + clock, +) -> dict[str, Any]: + # ---- 0. validate everything before the first AWS call + entries, start_dt, end_dt, now_dt, problems = _validate( + groups_in, + query_string, + start_time, + end_time, + now, + { + "max_queries": max_queries, + "max_rows": max_rows, + "max_polls": max_polls, + }, + max_splits, + {"poll_interval": poll_interval, "timeout_seconds": timeout_seconds}, + ) + if problems or start_dt is None or end_dt is None or now_dt is None: + return invalid_input( + "Invalid input: " + + (problems[0] if problems else "bad time values"), + problems, + ) + + analysis = _analyze_query(query_string) + commands = analysis["commands"] + is_aggregate = analysis["is_aggregate"] + explicit_limit = analysis["explicit_limit"] + + # ---- 1. resolve groups and bound the window + resolution, groups, notes, resolution_complete, err = _resolve_groups( + call_op, entries + ) + if err is not None: + return err + if not groups: + wanted = ", ".join(repr(n + ("*" if p else "")) for n, p in entries) + return err_result( + NOT_FOUND, + f"No log groups matched {wanted}.", + guidance=[ + "Check the region (log groups are regional) and try a shorter prefix ending in '*' or a " + "distinctive substring such as the function/service name." + ], + value={"log_group_resolution": resolution}, + ) + + guidance: list[str] = list(notes) + if analysis["unknown_commands"]: + guidance.append( + "Query contains commands not recognized by this tool; results may vary." + ) + partial_reasons: list[str] = [] + if not resolution_complete: + partial_reasons.append( + f"DescribeLogGroups listing stopped at the {_DESCRIBE_PAGE_CAP}-page cap; " + "some matching groups may be missing. Use a narrower prefix." + ) + + now_ms = int(now_dt.timestamp() * 1000) + req_start_ms = int(start_dt.timestamp() * 1000) + req_end_ms = int(end_dt.timestamp() * 1000) + clamp_reasons: list[str] = [] + end_ms = req_end_ms + if end_ms > now_ms: + end_ms = now_ms + clamp_reasons.append( + f"end_time is in the future; clamped to now ({_iso_ms(now_ms)})." + ) + in_window: list[tuple[str, int]] = [] + skipped: list[dict[str, Any]] = [] + for name, meta in groups.items(): + # Retention-only bound. A log group's creationTime does NOT bound a Logs + # Insights query (events can exist from before the group's record was + # created), so it is intentionally excluded from the queryable edge. + earliest = 0 + retention = meta.get("retentionInDays") + if retention: + earliest = now_ms - int(retention) * _DAY_MS + if earliest > end_ms: + skipped.append( + { + "log_group": name, + "queryable_from": _iso_ms(earliest), + "retention_in_days": retention, + "reason": "window ends before this group's retention edge", + } + ) + else: + in_window.append((name, earliest)) + if not in_window: + return invalid_input( + "None of the matched log groups hold data in the requested window.", + [ + f"{s['log_group']} is queryable from {s['queryable_from']}" + for s in skipped + ] + + ["Move start_time/end_time into the queryable range."], + value={ + "skipped_log_groups": skipped, + "log_group_resolution": resolution, + }, + ) + start_ms = req_start_ms + earliest_any = min(e for _, e in in_window) + if earliest_any > start_ms: + start_ms = earliest_any + clamp_reasons.append( + f"start_time precedes all retained data; clamped to {_iso_ms(start_ms)} " + "(retention edge). Events before it no longer exist." + ) + start_s, end_s = _ceil_s(start_ms), _ceil_s(end_ms) + if start_s >= end_s: + return invalid_input( + f"Clamped window {_iso(start_s)}..{_iso(end_s)} is empty.", + [ + "Choose a window that ends after start_time and is not entirely in the future." + ], + ) + if skipped: + guidance.append( + f"Skipped {len(skipped)} group(s) with no data in the window: " + + ", ".join(s["log_group"] for s in skipped[:5]) + ) + guidance.extend(clamp_reasons) + + names = [n for n, _ in in_window] + batches = [ + names[i : i + _MAX_GROUPS_PER_QUERY] + for i in range(0, len(names), _MAX_GROUPS_PER_QUERY) + ] + if len(batches) > 1: + guidance.append( + f"{len(names)} groups exceed the 50-per-query limit; ran {len(batches)} batches and merged." + ) + limit = None if explicit_limit else _ROW_CAP + + # ---- 2-4. run segments, splitting row-capped windows + work: list[tuple[int, int, int, int]] = [ + (b, start_s, end_s, 0) for b in range(len(batches)) + ] + segments: list[dict[str, Any]] = [] + rows: list[dict[str, Any]] = [] + seen_ptrs: set[str] = set() + totals: dict[str, float] = {} + queries_run = 0 + split_count = 0 + rows_truncated = False + while work: + if queries_run >= max_queries: + partial_reasons.append( + f"Stopped at max_queries={max_queries}; {len(work)} sub-window(s) not queried: " + + ", ".join(f"{_iso(s)}..{_iso(e)}" for _, s, e, _ in work[:3]) + ) + break + b, s, e, depth = work.pop(0) + seg = _run_segment( + call_op, + batches[b], + s, + e, + query_string, + limit, + sleep=sleep, + clock=clock, + poll_interval=poll_interval, + max_polls=max_polls, + timeout_seconds=timeout_seconds, + ) + queries_run += 1 + if seg.get("ok") is False: + done = [x["query_id"] for x in segments if not x["discarded"]] + if done: + seg["guidance"] = list(seg.get("guidance") or []) + [ + f"{len(done)} sub-query(ies) had completed before the failure: " + + ", ".join(done[:5]) + ] + return seg + seg_rows: list[dict[str, Any]] = seg["rows"] + matched = seg["statistics"].get("recordsMatched", 0.0) + capped = ( + limit is not None + and len(seg_rows) >= limit + and matched > len(seg_rows) + ) + record = { + "query_id": seg["query_id"], + "start": _iso(s), + "end": _iso(e), + "log_group_batch": b, + "rows": len(seg_rows), + "records_matched": matched, + "discarded": False, + } + if capped and not is_aggregate: + if depth < max_splits and e - s >= 1: + mid = (s + e) // 2 + # mid+1 is correct: Insights time range is inclusive on both + # ends, so [s, mid] and [mid+1, e] partition without overlap or + # gap. + work.extend( + [(b, s, mid, depth + 1), (b, mid + 1, e, depth + 1)] + ) + record["discarded"] = True + segments.append(record) + split_count += 1 + continue + partial_reasons.append( + f"Sub-window {_iso(s)}..{_iso(e)} matched {int(matched)} records but returned " + f"{len(seg_rows)} (row cap) and max_splits={max_splits} was reached." + ) + elif capped: + partial_reasons.append( + f"Aggregate query returned {len(seg_rows)} groups (row cap); some groups are " + "missing. Coarsen the 'by' clause or bin() size." + ) + segments.append(record) + for k, v in seg["statistics"].items(): + totals[k] = totals.get(k, 0.0) + v + for row in seg_rows: + ptr = row.get("@ptr") + if isinstance(ptr, str): + if ptr in seen_ptrs: + continue + seen_ptrs.add(ptr) + if len(rows) >= max_rows: + rows_truncated = True + break + rows.append(row) + if rows_truncated: + partial_reasons.append( + f"Row output truncated at max_rows={max_rows}; statistics still reflect all matches." + ) + if explicit_limit: + guidance.append( + "Query has an explicit 'limit'; results honour it and were not window-split." + ) + kept = [x for x in segments if not x["discarded"]] + if not is_aggregate and len(kept) > 1 and "sort" in commands: + guidance.append( + "Rows were merged from several sub-windows; 'sort' order holds within each sub-window only." + ) + if is_aggregate and len(kept) > 1: + guidance.append( + "Aggregate rows come from several batches/sub-windows; sum per-key counts across rows." + ) + + # ---- Truncate oversized @message fields (bounded output), recording the adjustment. + truncated_messages = 0 + for row in rows: + if "@message" in row: + new_value, was_truncated = truncate_field( + row["@message"], max_message_length + ) + if was_truncated: + row["@message"] = new_value + truncated_messages += 1 + if truncated_messages: + clamp_reasons.append( + f"Truncated {truncated_messages} @message field(s) to {max_message_length} chars " + "(--max-message-length); raise it to see full messages." + ) + guidance.append(clamp_reasons[-1]) + + fields: list[str] = [] + for row in rows: + for k in row: + if k not in fields: + fields.append(k) + value = { + "rows": rows, + "row_count": len(rows), + "fields": fields, + "statistics": totals, + "window": { + "start": _iso(start_s), + "end": _iso(end_s), + "start_epoch": start_s, + "end_epoch": end_s, + }, + "requested_window": { + "start": _iso_ms(req_start_ms), + "end": _iso_ms(req_end_ms), + }, + "clamped": bool(clamp_reasons), + "clamp_reasons": clamp_reasons, + "log_groups": names, + "log_group_resolution": resolution, + "skipped_log_groups": skipped, + "segments": segments, + "split_count": split_count, + "queries_run": queries_run, + "is_aggregate": is_aggregate, + "message_truncated_count": truncated_messages, + "partial_reasons": partial_reasons, + "query_string": query_string, + } + return ok_result( + value, + is_exhaustive=not partial_reasons, + guidance=guidance + partial_reasons, + ) + + +class RunInsightsQueryCommand(BasicCommand): + NAME = "run-insights-query" + DESCRIPTION = ( + "Run a single CloudWatch Logs Insights query to completion over one or " + "more log groups and return the full result set in one call.\n\n" + "Use this command when you need to query CloudWatch Logs Insights across " + "one or more log groups and get structured results in one step. It " + "handles log group resolution, query submission, polling, time-window " + "splitting for large result sets, and output capping.\n\n" + "Resolves " + "groups by exact name, prefix (trailing '*'), ARN or substring; bounds " + "the window to retention; polls to a terminal state; and splits " + "row-capped windows to keep counts accurate. Read-only. Uses " + "logGroupNames; cross-account queries (logGroupIdentifiers) are not " + "supported in this version." + ) + EXAMPLES = ( + "To run a query over a log group for the last hour's errors::\n\n" + " aws logs run-insights-query --log-group-names /aws/lambda/my-fn " + "--query-string 'fields @timestamp, @message | filter @message like /ERROR/' " + "--start-time 2024-01-15T00:00:00Z --end-time 2024-01-15T01:00:00Z\n" + ) + ARG_TABLE = [ + { + "name": "log-group-names", + "required": True, + "nargs": "+", + "help_text": "One or more log group names, prefixes (trailing '*'), ARNs or substrings.", + "cli_type_name": "string", + }, + { + "name": "query-string", + "required": True, + "help_text": "The Logs Insights query to run.", + "cli_type_name": "string", + }, + { + "name": "start-time", + "required": True, + "help_text": "Window start (ISO-8601, e.g. 2024-01-15T00:00:00Z).", + "cli_type_name": "string", + }, + { + "name": "end-time", + "required": True, + "help_text": "Window end (ISO-8601).", + "cli_type_name": "string", + }, + { + "name": "max-rows", + "default": 50, + "cli_type_name": "integer", + "help_text": "Maximum rows to return (default 50); statistics still reflect all matches.", + }, + { + "name": "max-message-length", + "default": DEFAULT_MESSAGE_MAX_CHARS, + "cli_type_name": "integer", + "help_text": f"Truncate each @message to this many chars (default {DEFAULT_MESSAGE_MAX_CHARS}).", + }, + { + "name": "max-splits", + "default": 6, + "cli_type_name": "integer", + "help_text": "Maximum times a row-capped window is split to recover counts (default 6).", + }, + { + "name": "poll-interval", + "default": 1.0, + "cli_type_name": "double", + "help_text": "Initial seconds between GetQueryResults polls (default 1.0).", + }, + { + "name": "timeout-seconds", + "default": 900.0, + "cli_type_name": "double", + "help_text": "Maximum seconds to wait per sub-query (default 900).", + }, + ] + + def __init__( + self, + session, + client=None, + stdout=None, + stderr=None, + sleep=None, + clock=None, + ): + super().__init__(session) + self._client = client + self._stdout = stdout + self._stderr = stderr + self._sleep = sleep + self._clock = clock + + def _run_main(self, parsed_args, parsed_globals): + client = self._client or create_read_client( + self._session, "logs", parsed_globals + ) + + def call_op(operation, params): + return call(client, operation, params) + + result = run_insights_query( + call_op, + parsed_args.log_group_names, + parsed_args.query_string, + parsed_args.start_time, + parsed_args.end_time, + max_splits=parsed_args.max_splits, + max_rows=parsed_args.max_rows, + max_message_length=parsed_args.max_message_length, + poll_interval=parsed_args.poll_interval, + timeout_seconds=parsed_args.timeout_seconds, + sleep=self._sleep, + clock=self._clock, + ) + return emit_result( + self.NAME, + result, + parsed_globals, + self._session, + self._stdout, + self._stderr, + ) + + +def register(event_handlers): + event_handlers.register('building-command-table.logs', _inject_commands) + + +def _inject_commands(command_table, session, **kwargs): + command_table['run-insights-query'] = RunInsightsQueryCommand(session) diff --git a/awscli/handlers_registry.py b/awscli/handlers_registry.py index 24e1d926a418..5a1710888774 100644 --- a/awscli/handlers_registry.py +++ b/awscli/handlers_registry.py @@ -667,6 +667,7 @@ class CommandTableOp(enum.Enum): 'building-command-table.logs': [ ('awscli.customizations.removals', 'register_removals'), ('awscli.customizations.logs', 'register_logs_commands'), + ('awscli.customizations.logs.run_insights_query', 'register'), ], 'building-command-table.main': [ ('awscli.customizations.s3.s3', 'register_s3_main'), diff --git a/tests/functional/logs/test_run_insights_query.py b/tests/functional/logs/test_run_insights_query.py new file mode 100644 index 000000000000..c797da5be4d8 --- /dev/null +++ b/tests/functional/logs/test_run_insights_query.py @@ -0,0 +1,111 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"). You +# may not use this file except in compliance with the License. A copy of +# the License is located at +# +# http://aws.amazon.com/apache2.0/ +# +# or in the "license" file accompanying this file. This file is +# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF +# ANY KIND, either express or implied. See the License for the specific +# language governing permissions and limitations under the License. +import json + +from awscli.testutils import BaseAWSCommandParamsTest + +# Operations that would create, modify or delete state. The command must never +# issue any of these; the functional test asserts it. +_MUTATING = { + 'CreateLogGroup', + 'DeleteLogGroup', + 'PutLogEvents', + 'PutRetentionPolicy', + 'DeleteQueryDefinition', +} + +# A window well in the past so it is never clamped to "now" and the single log +# group (no retention policy) always holds data in it, keeping the test +# deterministic regardless of the wall clock. +_START = '2024-01-01T00:00:00Z' +_END = '2024-01-01T01:00:00Z' + + +class TestRunInsightsQueryCommand(BaseAWSCommandParamsTest): + def _lifecycle_responses(self): + return [ + # 1. DescribeLogGroups resolves the group (no retention -> queryable). + {'logGroups': [{'logGroupName': '/aws/lambda/fn'}]}, + # 2. StartQuery returns a non-blank queryId. + {'queryId': 'q-1'}, + # 3. GetQueryResults returns a terminal, Complete result on the first + # poll so the bounded poller never sleeps. + { + 'status': 'Complete', + 'results': [ + [ + {'field': '@message', 'value': 'hello'}, + {'field': '@ptr', 'value': 'p1'}, + ] + ], + 'statistics': {}, + }, + ] + + def test_runs_full_read_only_lifecycle(self): + self.parsed_responses = self._lifecycle_responses() + stdout, _, rc = self.run_cmd( + [ + 'logs', + 'run-insights-query', + '--log-group-names', + '/aws/lambda/fn', + '--query-string', + 'fields @timestamp, @message', + '--start-time', + _START, + '--end-time', + _END, + '--region', + 'us-east-1', + ], + expected_rc=0, + ) + names = [c[0].name for c in self.operations_called] + self.assertEqual( + names, + ['DescribeLogGroups', 'StartQuery', 'GetQueryResults'], + ) + # No mutating operation is ever issued. + self.assertFalse(_MUTATING.intersection(names)) + parsed = json.loads(stdout) + self.assertEqual(parsed['row_count'], 1) + self.assertEqual(parsed['rows'], [{'@message': 'hello', '@ptr': 'p1'}]) + + def test_unknown_query_command_runs_instead_of_rejecting(self): + # An unrecognized leading token is advisory now: the command dispatches + # the full read-only lifecycle instead of rejecting locally (the server + # would reject a query that is truly malformed). + self.parsed_responses = self._lifecycle_responses() + self.run_cmd( + [ + 'logs', + 'run-insights-query', + '--log-group-names', + '/aws/lambda/fn', + '--query-string', + 'boguscmd @message', + '--start-time', + _START, + '--end-time', + _END, + '--region', + 'us-east-1', + ], + expected_rc=0, + ) + names = [c[0].name for c in self.operations_called] + self.assertEqual( + names, + ['DescribeLogGroups', 'StartQuery', 'GetQueryResults'], + ) diff --git a/tests/unit/customizations/logs/test_run_insights_query.py b/tests/unit/customizations/logs/test_run_insights_query.py new file mode 100644 index 000000000000..d53ca15cf209 --- /dev/null +++ b/tests/unit/customizations/logs/test_run_insights_query.py @@ -0,0 +1,1303 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"). You +# may not use this file except in compliance with the License. A copy of +# the License is located at +# +# http://aws.amazon.com/apache2.0/ +# +# or in the "license" file accompanying this file. This file is +# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF +# ANY KIND, either express or implied. See the License for the specific +# language governing permissions and limitations under the License. +"""Unit tests for ``aws logs run-insights-query`` (run_insights_query logic).""" + +from awscli.customizations.compositereadonly.helpers import ( + CONFLICT, + CREDENTIALS_EXPIRED, + NOT_FOUND, + PERMISSION_DENIED, + QUOTA_EXCEEDED, + SERVICE_ERROR, + THROTTLED, + TIMEOUT, + VALIDATION_ERROR, + assert_read_only, + err_result, + ok_result, +) +from awscli.customizations.logs import run_insights_query as mod + +NOW = "2026-10-28T12:00:00Z" + + +class ScriptedLogs: + def __init__(self): + self.calls = [] + self._responses = {} + + def queue(self, operation, *results): + self._responses.setdefault(operation, []).extend(results) + return self + + def __call__(self, operation, params): + assert_read_only(operation) + self.calls.append((operation, params)) + queued = self._responses.get(operation) + if queued: + return queued.pop(0) + return ok_result({}) + + def ops(self): + return [op for op, _ in self.calls] + + +def _group(name, creation_ms=0, retention=None): + g = {"logGroupName": name, "creationTime": creation_ms} + if retention: + g["retentionInDays"] = retention + return g + + +def _dlg(*groups): + return ok_result({"logGroups": list(groups)}) + + +def _qr(status="Complete", rows=None, stats=None): + results = [ + [{"field": k, "value": v} for k, v in row.items()] + for row in (rows or []) + ] + return ok_result( + {"status": status, "results": results, "statistics": stats or {}} + ) + + +def _run( + call, + groups="/aws/lambda/fn", + query="fields @timestamp, @message", + start="2026-10-28T10:00:00Z", + end="2026-10-28T11:00:00Z", + **kwargs, +): + return mod.run_insights_query( + call, + groups, + query, + start, + end, + now=NOW, + sleep=lambda _s: None, + clock=_clock(), + **kwargs, + ) + + +def _clock(): + state = {"t": 0.0} + + def clock(): + state["t"] += 0.001 + return state["t"] + + return clock + + +# --------------------------------------------------------------- happy path / read-only + + +def test_success_only_read_ops(): + call = ScriptedLogs() + call.queue("DescribeLogGroups", _dlg(_group("/aws/lambda/fn"))) + call.queue("StartQuery", ok_result({"queryId": "q1"})) + call.queue( + "GetQueryResults", _qr(rows=[{"@message": "hello", "@ptr": "p1"}]) + ) + result = _run(call) + assert result["ok"] is True + assert result["value"]["row_count"] == 1 + assert result["value"]["rows"][0]["@message"] == "hello" + # Read-only lifecycle only. + assert call.ops() == ["DescribeLogGroups", "StartQuery", "GetQueryResults"] + + +def test_empty_results(): + call = ScriptedLogs() + call.queue("DescribeLogGroups", _dlg(_group("/aws/lambda/fn"))) + call.queue("StartQuery", ok_result({"queryId": "q1"})) + call.queue("GetQueryResults", _qr(rows=[])) + result = _run(call) + assert result["ok"] is True + assert result["value"]["row_count"] == 0 + + +def test_polls_until_terminal(): + call = ScriptedLogs() + call.queue("DescribeLogGroups", _dlg(_group("/aws/lambda/fn"))) + call.queue("StartQuery", ok_result({"queryId": "q1"})) + call.queue("GetQueryResults", _qr(status="Running")) + call.queue("GetQueryResults", _qr(status="Scheduled")) + call.queue("GetQueryResults", _qr(rows=[{"@ptr": "p1"}])) + result = _run(call) + assert result["ok"] is True + assert call.ops().count("GetQueryResults") == 3 + + +# --- creation time must NOT clamp the query window + + +def test_creation_time_does_not_clamp_window(): + # Group was recorded (creationTime) AFTER start_time, but with no retention + # its data is still queryable from start_time. The window must NOT be clamped + # up to creationTime. + call = ScriptedLogs() + creation = _ms("2026-10-28T10:30:00Z") # after the 10:00 start + call.queue( + "DescribeLogGroups", + _dlg(_group("/aws/lambda/fn", creation_ms=creation)), + ) + call.queue("StartQuery", ok_result({"queryId": "q1"})) + call.queue("GetQueryResults", _qr(rows=[{"@ptr": "p1"}])) + result = _run(call) + assert result["ok"] is True + # Window start stays at the requested 10:00, not clamped to 10:30 creation. + assert result["value"]["window"]["start"] == "2026-10-28T10:00:00Z" + assert not any( + "creation" in r.lower() for r in result["value"]["clamp_reasons"] + ) + + +# --- retention period clamps the query window + + +def test_retention_clamps_start(): + # start far in the past; a 1-day retention edge is after it, so the window + # start is clamped UP to the retention edge (this clamp is kept). + call = ScriptedLogs() + call.queue( + "DescribeLogGroups", _dlg(_group("/aws/lambda/fn", retention=1)) + ) + call.queue("StartQuery", ok_result({"queryId": "q1"})) + call.queue("GetQueryResults", _qr(rows=[{"@ptr": "p1"}])) + result = mod.run_insights_query( + call, + "/aws/lambda/fn", + "fields @message", + "2026-10-01T00:00:00Z", + "2026-10-28T11:00:00Z", + now=NOW, + sleep=lambda _s: None, + clock=_clock(), + ) + assert result["ok"] is True + assert result["value"]["clamped"] is True + assert any("retention edge" in r for r in result["value"]["clamp_reasons"]) + # Clamped to now - 1 day = 2026-10-27T12:00:00Z. + assert result["value"]["window"]["start"] == "2026-10-27T12:00:00Z" + + +def test_group_skipped_when_window_before_retention_edge(): + call = ScriptedLogs() + call.queue( + "DescribeLogGroups", _dlg(_group("/aws/lambda/fn", retention=1)) + ) + # Window entirely before the 1-day retention edge (now-1day = 10-27T12:00). + result = mod.run_insights_query( + call, + "/aws/lambda/fn", + "fields @message", + "2026-10-20T00:00:00Z", + "2026-10-21T00:00:00Z", + now=NOW, + sleep=lambda _s: None, + clock=_clock(), + ) + assert result["ok"] is False + assert result["code"] == "INVALID_INPUT" + assert "StartQuery" not in call.ops() + + +# --- default max_rows caps at ~50 rows + + +def test_default_max_rows_is_50(): + call = ScriptedLogs() + call.queue("DescribeLogGroups", _dlg(_group("/aws/lambda/fn"))) + call.queue("StartQuery", ok_result({"queryId": "q1"})) + many = [{"@ptr": f"p{i}", "@message": f"m{i}"} for i in range(60)] + call.queue("GetQueryResults", _qr(rows=many, stats={"recordsMatched": 60})) + result = _run(call) + assert result["value"]["row_count"] == 50 + assert result["is_exhaustive"] is False + assert any("max_rows=50" in r for r in result["value"]["partial_reasons"]) + + +# --- @message field is truncated at limit + + +def test_message_is_truncated_and_recorded(): + call = ScriptedLogs() + call.queue("DescribeLogGroups", _dlg(_group("/aws/lambda/fn"))) + call.queue("StartQuery", ok_result({"queryId": "q1"})) + long_msg = "x" * 500 + call.queue( + "GetQueryResults", _qr(rows=[{"@message": long_msg, "@ptr": "p1"}]) + ) + result = _run(call, max_message_length=20) + msg = result["value"]["rows"][0]["@message"] + assert len(msg) <= 20 + assert msg.endswith("...[truncated]") + assert result["value"]["message_truncated_count"] == 1 + assert any("@message" in r for r in result["value"]["clamp_reasons"]) + + +# --------------------------------------------------------------- resolution + + +def test_prefix_resolution_when_no_exact_match(): + call = ScriptedLogs() + call.queue( + "DescribeLogGroups", + _dlg(_group("/aws/lambda/fn-a"), _group("/aws/lambda/fn-b")), + ) + call.queue("StartQuery", ok_result({"queryId": "q1"})) + call.queue("GetQueryResults", _qr(rows=[{"@ptr": "p1"}])) + result = _run(call, groups="/aws/lambda/fn") + assert result["ok"] is True + res = result["value"]["log_group_resolution"][0] + assert res["match"] == "prefix" + assert len(res["log_groups"]) == 2 + + +def test_no_groups_matched_is_not_found(): + call = ScriptedLogs() + call.queue("DescribeLogGroups", _dlg()) # prefix + call.queue("DescribeLogGroups", _dlg()) # pattern fallback + result = _run(call, groups="nonexistent") + assert result["ok"] is False + assert result["code"] == "NOT_FOUND" + assert "StartQuery" not in call.ops() + + +# --------------------------------------------------------------- validation / malformed + + +def test_unknown_command_is_advisory_not_fatal(): + call = ScriptedLogs() + call.queue("DescribeLogGroups", _dlg(_group("/aws/lambda/fn"))) + call.queue("StartQuery", ok_result({"queryId": "q1"})) + call.queue("GetQueryResults", _qr(rows=[{"@ptr": "p1"}])) + result = _run(call, query="boguscmd @message") + # An unrecognized leading token no longer blocks the query locally; it runs + # and the server would reject it if it were truly invalid. + assert result["ok"] is True + assert "StartQuery" in call.ops() + assert any("not recognized by this tool" in g for g in result["guidance"]) + + +def test_empty_query_rejected(): + call = ScriptedLogs() + result = _run(call, query=" ") + assert result["ok"] is False + assert result["param_error"] is True + assert call.calls == [] + + +def test_malformed_query_from_server_is_operational_error(): + call = ScriptedLogs() + call.queue("DescribeLogGroups", _dlg(_group("/aws/lambda/fn"))) + call.queue( + "StartQuery", + err_result( + "INVALID_INPUT", + "line 1, position 5: bad", + aws_code="MalformedQueryException", + ), + ) + result = _run(call) + assert result["ok"] is False + assert result["code"] == "INVALID_INPUT" + # Server-discovered: NOT flagged as a pre-dispatch param error. + assert result.get("param_error") is not True + + +def test_blank_query_id_refuses_to_poll(): + call = ScriptedLogs() + call.queue("DescribeLogGroups", _dlg(_group("/aws/lambda/fn"))) + call.queue("StartQuery", ok_result({"queryId": " "})) + result = _run(call) + assert result["ok"] is False + assert result["code"] == "SERVICE_ERROR" + assert "GetQueryResults" not in call.ops() + + +# --------------------------------------------------------------- splitting (row cap) + + +def test_row_capped_window_splits(monkeypatch): + # Shrink the per-query row cap so a small fixture triggers a split. + monkeypatch.setattr(mod, "_ROW_CAP", 2) + call = ScriptedLogs() + call.queue("DescribeLogGroups", _dlg(_group("/aws/lambda/fn"))) + # First segment hits the cap (2 rows, more matched) -> split into two halves. + call.queue("StartQuery", ok_result({"queryId": "q1"})) + call.queue( + "GetQueryResults", + _qr(rows=[{"@ptr": "a"}, {"@ptr": "b"}], stats={"recordsMatched": 10}), + ) + # Two sub-windows, each returns a small, non-capped result. + call.queue("StartQuery", ok_result({"queryId": "q2"})) + call.queue( + "GetQueryResults", + _qr(rows=[{"@ptr": "c"}], stats={"recordsMatched": 1}), + ) + call.queue("StartQuery", ok_result({"queryId": "q3"})) + call.queue( + "GetQueryResults", + _qr(rows=[{"@ptr": "d"}], stats={"recordsMatched": 1}), + ) + result = _run(call, query="fields @message") + assert result["ok"] is True + assert result["value"]["split_count"] >= 1 + # @ptr de-duplication across merged sub-windows. + ptrs = [r["@ptr"] for r in result["value"]["rows"]] + assert len(ptrs) == len(set(ptrs)) + + +def _ms(iso): + from datetime import datetime, timezone + + dt = datetime.fromisoformat(iso.replace("Z", "+00:00")) + if dt.tzinfo is None: + dt = dt.replace(tzinfo=timezone.utc) + return int(dt.timestamp() * 1000) + + +def _run_window(call, start, end, **kwargs): + kwargs.setdefault("groups", "/aws/lambda/fn") + kwargs.setdefault("query", "fields @timestamp, @message") + return mod.run_insights_query( + call, + kwargs.pop("groups"), + kwargs.pop("query"), + start, + end, + now=NOW, + sleep=lambda _s: None, + clock=_clock(), + **kwargs, + ) + + +# --------------------------------------------------------------- input validation + + +def test_accepts_datetime_objects_for_times(): + from datetime import datetime, timezone + + call = ScriptedLogs() + call.queue("DescribeLogGroups", _dlg(_group("/aws/lambda/fn"))) + call.queue("StartQuery", ok_result({"queryId": "q1"})) + call.queue("GetQueryResults", _qr(rows=[{"@ptr": "p1"}])) + start = datetime(2026, 10, 28, 10, tzinfo=timezone.utc) + end = datetime(2026, 10, 28, 11, tzinfo=timezone.utc) + result = _run_window(call, start, end) + assert result["ok"] is True + + +def test_naive_datetime_is_treated_as_utc(): + from datetime import datetime + + call = ScriptedLogs() + call.queue("DescribeLogGroups", _dlg(_group("/aws/lambda/fn"))) + call.queue("StartQuery", ok_result({"queryId": "q1"})) + call.queue("GetQueryResults", _qr(rows=[{"@ptr": "p1"}])) + # No tzinfo -> assumed UTC; the query still runs. + result = _run_window( + call, + datetime(2026, 10, 28, 10), + datetime(2026, 10, 28, 11), + ) + assert result["ok"] is True + + +def test_non_iso_start_time_rejected(): + call = ScriptedLogs() + result = _run(call, start="not-a-time") + assert result["ok"] is False + assert result["param_error"] is True + assert call.calls == [] + + +def test_missing_start_time_rejected(): + call = ScriptedLogs() + result = _run(call, start="") + assert result["ok"] is False + assert result["param_error"] is True + + +def test_start_not_before_end_rejected(): + call = ScriptedLogs() + result = _run( + call, + start="2026-10-28T11:00:00Z", + end="2026-10-28T10:00:00Z", + ) + assert result["ok"] is False + assert "before end_time" in result["error"] + + +def test_invalid_now_rejected(): + call = ScriptedLogs() + result = mod.run_insights_query( + call, + "/aws/lambda/fn", + "fields @message", + "2026-10-28T10:00:00Z", + "2026-10-28T11:00:00Z", + now="garbage", + sleep=lambda _s: None, + clock=_clock(), + ) + assert result["ok"] is False + + +def test_empty_group_list_rejected(): + call = ScriptedLogs() + result = _run(call, groups=[]) + assert result["ok"] is False + assert any("log_group_names is empty" in p for p in result["guidance"]) + + +def test_invalid_group_name_rejected(): + call = ScriptedLogs() + result = _run(call, groups="has a space") + assert result["ok"] is False + assert result["param_error"] is True + assert call.calls == [] + + +def test_nonpositive_max_rows_rejected(): + call = ScriptedLogs() + result = _run(call, max_rows=0) + assert result["ok"] is False + + +def test_boolean_max_queries_rejected(): + call = ScriptedLogs() + result = _run(call, max_queries=True) + assert result["ok"] is False + + +def test_negative_max_splits_rejected(): + call = ScriptedLogs() + result = _run(call, max_splits=-1) + assert result["ok"] is False + + +def test_nonpositive_timeout_rejected(): + call = ScriptedLogs() + result = _run(call, timeout_seconds=0) + assert result["ok"] is False + + +def test_query_too_long_rejected(): + call = ScriptedLogs() + result = _run(call, query="fields @message " + "x" * 11000) + assert result["ok"] is False + + +# --------------------------------------------------------------- query analysis + + +def test_pipe_inside_regex_is_not_a_stage_separator(): + call = ScriptedLogs() + call.queue("DescribeLogGroups", _dlg(_group("/aws/lambda/fn"))) + call.queue("StartQuery", ok_result({"queryId": "q1"})) + call.queue("GetQueryResults", _qr(rows=[{"@ptr": "p1"}])) + # The '|' lives inside the /ERROR|WARN/ regex, so there are two stages + # (fields, filter), both valid -- not three. + result = _run( + call, + query="fields @message | filter @message like /ERROR|WARN/", + ) + assert result["ok"] is True + + +def test_pipe_inside_quotes_is_not_a_stage_separator(): + call = ScriptedLogs() + call.queue("DescribeLogGroups", _dlg(_group("/aws/lambda/fn"))) + call.queue("StartQuery", ok_result({"queryId": "q1"})) + call.queue("GetQueryResults", _qr(rows=[{"@ptr": "p1"}])) + result = _run( + call, + query="fields @message | filter @message = 'a|b'", + ) + assert result["ok"] is True + + +def test_pipe_inside_parens_is_not_a_stage_separator(): + call = ScriptedLogs() + call.queue("DescribeLogGroups", _dlg(_group("/aws/lambda/fn"))) + call.queue("StartQuery", ok_result({"queryId": "q1"})) + call.queue("GetQueryResults", _qr(rows=[{"@ptr": "p1"}])) + result = _run( + call, + query="fields @message | stats count() by (a)", + ) + assert result["ok"] is True + + +def test_parse_command_opens_a_regex_literal(): + call = ScriptedLogs() + call.queue("DescribeLogGroups", _dlg(_group("/aws/lambda/fn"))) + call.queue("StartQuery", ok_result({"queryId": "q1"})) + call.queue("GetQueryResults", _qr(rows=[{"@ptr": "p1"}])) + # 'parse' opens a /.../ literal that itself contains a '|'. + result = _run( + call, + query="fields @message | parse @message /a|b/ as x", + ) + assert result["ok"] is True + + +def test_empty_pipeline_stage_is_advisory_not_fatal(): + call = ScriptedLogs() + call.queue("DescribeLogGroups", _dlg(_group("/aws/lambda/fn"))) + call.queue("StartQuery", ok_result({"queryId": "q1"})) + call.queue("GetQueryResults", _qr(rows=[{"@ptr": "p1"}])) + result = _run(call, query="fields @message | | stats count()") + # A stray '|' is no longer rejected locally; it is flagged as advisory. + assert result["ok"] is True + assert any("not recognized by this tool" in g for g in result["guidance"]) + + +# --------------------------------------------------------------- group resolution + + +def test_arn_input_is_normalized_to_group_name(): + call = ScriptedLogs() + arn = "arn:aws:logs:us-east-1:123456789012:log-group:/aws/lambda/fn:*" + call.queue("DescribeLogGroups", _dlg(_group("/aws/lambda/fn"))) + call.queue("StartQuery", ok_result({"queryId": "q1"})) + call.queue("GetQueryResults", _qr(rows=[{"@ptr": "p1"}])) + result = _run(call, groups=arn) + assert result["ok"] is True + # The ARN was reduced to the bare group name for DescribeLogGroups. + describe = call.calls[0][1] + assert describe["logGroupNamePrefix"] == "/aws/lambda/fn" + + +def test_exact_match_with_siblings_adds_note(): + call = ScriptedLogs() + call.queue( + "DescribeLogGroups", + _dlg(_group("/aws/lambda/fn"), _group("/aws/lambda/fn-2")), + ) + call.queue("StartQuery", ok_result({"queryId": "q1"})) + call.queue("GetQueryResults", _qr(rows=[{"@ptr": "p1"}])) + result = _run(call, groups="/aws/lambda/fn") + res = result["value"]["log_group_resolution"][0] + assert res["match"] == "exact" + # Only the exact group is queried, but a note flags the siblings. + assert res["log_groups"] == ["/aws/lambda/fn"] + assert any("share the prefix" in g for g in result["guidance"]) + + +def test_explicit_prefix_star_matches_all(): + call = ScriptedLogs() + call.queue( + "DescribeLogGroups", + _dlg(_group("/aws/lambda/fn-a"), _group("/aws/lambda/fn-b")), + ) + call.queue("StartQuery", ok_result({"queryId": "q1"})) + call.queue("GetQueryResults", _qr(rows=[{"@ptr": "p1"}])) + result = _run(call, groups="/aws/lambda/fn*") + res = result["value"]["log_group_resolution"][0] + assert res["match"] == "prefix" + assert len(res["log_groups"]) == 2 + + +def test_substring_pattern_fallback(): + call = ScriptedLogs() + call.queue("DescribeLogGroups", _dlg()) # prefix miss + call.queue( + "DescribeLogGroups", + _dlg(_group("/aws/lambda/my-fn"), _group("/aws/ecs/my-fn")), + ) + call.queue("StartQuery", ok_result({"queryId": "q1"})) + call.queue("GetQueryResults", _qr(rows=[{"@ptr": "p1"}])) + result = _run(call, groups="my-fn") + res = result["value"]["log_group_resolution"][0] + assert res["match"] == "pattern" + assert any("containing it" in g for g in result["guidance"]) + + +def test_substring_pattern_many_matches_elided(): + call = ScriptedLogs() + call.queue("DescribeLogGroups", _dlg()) # prefix miss + call.queue( + "DescribeLogGroups", + _dlg(*[_group(f"/svc/fn-{i}") for i in range(7)]), + ) + call.queue("StartQuery", ok_result({"queryId": "q1"})) + call.queue("GetQueryResults", _qr(rows=[{"@ptr": "p1"}])) + result = _run(call, groups="fn") + # More than 5 matches -> the note lists the first 5 and elides the rest. + assert any("..." in g for g in result["guidance"]) + + +def test_pattern_unsupported_is_treated_as_no_match(): + call = ScriptedLogs() + call.queue("DescribeLogGroups", _dlg()) # prefix miss + # logGroupNamePattern unsupported in the region -> VALIDATION error is + # swallowed and the input resolves to no groups (NOT_FOUND), not an error. + call.queue( + "DescribeLogGroups", + err_result(VALIDATION_ERROR, "pattern not supported"), + ) + result = _run(call, groups="weird") + assert result["ok"] is False + assert result["code"] == NOT_FOUND + + +def test_describe_prefix_error_propagates(): + call = ScriptedLogs() + call.queue("DescribeLogGroups", err_result(PERMISSION_DENIED, "denied")) + result = _run(call, groups="/aws/lambda/fn") + assert result["ok"] is False + assert result["code"] == PERMISSION_DENIED + assert any("lacks logs:DescribeLogGroups" in g for g in result["guidance"]) + + +def test_describe_pattern_hard_error_propagates(): + call = ScriptedLogs() + call.queue("DescribeLogGroups", _dlg()) # prefix miss + call.queue("DescribeLogGroups", err_result(PERMISSION_DENIED, "denied")) + result = _run(call, groups="substr") + assert result["ok"] is False + assert result["code"] == PERMISSION_DENIED + + +def test_describe_paging_is_followed(): + call = ScriptedLogs() + call.queue( + "DescribeLogGroups", + ok_result( + { + "logGroups": [{"logGroupName": "/aws/lambda/fn-a"}], + "nextToken": "t2", + } + ), + ) + call.queue( + "DescribeLogGroups", + ok_result({"logGroups": [{"logGroupName": "/aws/lambda/fn-b"}]}), + ) + call.queue("StartQuery", ok_result({"queryId": "q1"})) + call.queue("GetQueryResults", _qr(rows=[{"@ptr": "p1"}])) + result = _run(call, groups="/aws/lambda/fn*") + assert sorted(result["value"]["log_groups"]) == [ + "/aws/lambda/fn-a", + "/aws/lambda/fn-b", + ] + + +def test_describe_page_cap_reports_incomplete(monkeypatch): + monkeypatch.setattr(mod, "_DESCRIBE_PAGE_CAP", 1) + call = ScriptedLogs() + # One page with a nextToken, but the cap stops listing -> incomplete. + call.queue( + "DescribeLogGroups", + ok_result( + { + "logGroups": [{"logGroupName": "/aws/lambda/fn"}], + "nextToken": "more", + } + ), + ) + call.queue("StartQuery", ok_result({"queryId": "q1"})) + call.queue("GetQueryResults", _qr(rows=[{"@ptr": "p1"}])) + result = _run(call, groups="/aws/lambda/fn*") + assert result["is_exhaustive"] is False + assert any("page cap" in p for p in result["value"]["partial_reasons"]) + + +# --------------------------------------------------------------- window clamping + + +def test_future_end_time_clamped_to_now(): + call = ScriptedLogs() + call.queue("DescribeLogGroups", _dlg(_group("/aws/lambda/fn"))) + call.queue("StartQuery", ok_result({"queryId": "q1"})) + call.queue("GetQueryResults", _qr(rows=[{"@ptr": "p1"}])) + result = _run_window(call, "2026-10-28T11:00:00Z", "2026-10-28T13:00:00Z") + assert result["value"]["clamped"] is True + assert result["value"]["window"]["end"] == "2026-10-28T12:00:00Z" + assert any("future" in r for r in result["value"]["clamp_reasons"]) + + +def test_some_groups_skipped_others_queried(): + call = ScriptedLogs() + # fn-old has 1-day retention (queryable from 10-27T12:00); the window + # 10:00-11:00 on 10-28 is inside retention for both, so both are kept -- + # instead make one group's retention exclude the window. + call.queue( + "DescribeLogGroups", + _dlg( + _group("/aws/lambda/keep"), + _group("/aws/lambda/drop", retention=1), + ), + ) + call.queue("StartQuery", ok_result({"queryId": "q1"})) + call.queue("GetQueryResults", _qr(rows=[{"@ptr": "p1"}])) + # Window entirely before the 1-day retention edge for 'drop' but '0' + # retention 'keep' is always queryable. + result = mod.run_insights_query( + call, + "/aws/lambda/*", + "fields @message", + "2026-10-20T00:00:00Z", + "2026-10-21T00:00:00Z", + now=NOW, + sleep=lambda _s: None, + clock=_clock(), + ) + assert result["ok"] is True + skipped = [s["log_group"] for s in result["value"]["skipped_log_groups"]] + assert "/aws/lambda/drop" in skipped + assert result["value"]["log_groups"] == ["/aws/lambda/keep"] + assert any("Skipped" in g for g in result["guidance"]) + + +# --------------------------------------------------------------- segment lifecycle + + +def test_start_query_throttled_then_succeeds(): + call = ScriptedLogs() + call.queue("DescribeLogGroups", _dlg(_group("/aws/lambda/fn"))) + call.queue("StartQuery", err_result(THROTTLED, "slow")) + call.queue("StartQuery", ok_result({"queryId": "q1"})) + call.queue("GetQueryResults", _qr(rows=[{"@ptr": "p1"}])) + result = _run(call) + assert result["ok"] is True + assert call.ops().count("StartQuery") == 2 + + +def test_start_query_throttled_beyond_retries(): + call = ScriptedLogs() + call.queue("DescribeLogGroups", _dlg(_group("/aws/lambda/fn"))) + for _ in range(mod._START_RETRIES + 1): + call.queue("StartQuery", err_result(THROTTLED, "slow")) + result = _run(call) + assert result["ok"] is False + assert result["code"] == QUOTA_EXCEEDED + assert result["retriable"] is True + + +def test_malformed_query_with_line_position_detail(): + call = ScriptedLogs() + call.queue("DescribeLogGroups", _dlg(_group("/aws/lambda/fn"))) + call.queue( + "StartQuery", + err_result( + SERVICE_ERROR, + "line 1, position 8: unexpected token", + aws_code="MalformedQueryException", + ), + ) + result = _run(call, query="fields @badtoken") + assert result["ok"] is False + assert result["code"] == "INVALID_INPUT" + assert result["value"]["source"] == "server" + assert result["value"]["position"] == 8 + + +def test_malformed_query_with_compile_error_location(): + call = ScriptedLogs() + call.queue("DescribeLogGroups", _dlg(_group("/aws/lambda/fn"))) + call.queue( + "StartQuery", + err_result( + SERVICE_ERROR, + "bad query", + aws_code="MalformedQueryException", + value={"queryCompileError": {"location": {"start": 7}}}, + ), + ) + result = _run(call, query="fields @timestamp") + assert result["value"]["position"] == 7 + # Offset 7 lands on '@timestamp', so that token is extracted. + assert result["value"]["token"] + + +def test_malformed_query_position_only_detail(): + call = ScriptedLogs() + call.queue("DescribeLogGroups", _dlg(_group("/aws/lambda/fn"))) + call.queue( + "StartQuery", + err_result( + SERVICE_ERROR, + "syntax error at position 3", + aws_code="MalformedQueryException", + ), + ) + result = _run(call) + assert result["value"]["position"] == 3 + + +def test_invalid_parameter_from_start_query(): + call = ScriptedLogs() + call.queue("DescribeLogGroups", _dlg(_group("/aws/lambda/fn"))) + call.queue( + "StartQuery", + err_result( + SERVICE_ERROR, + "bad window", + aws_code="InvalidParameterException", + ), + ) + result = _run(call) + assert result["ok"] is False + assert result["code"] == "INVALID_INPUT" + + +def test_resource_not_found_from_start_query(): + call = ScriptedLogs() + call.queue("DescribeLogGroups", _dlg(_group("/aws/lambda/fn"))) + call.queue( + "StartQuery", + err_result( + NOT_FOUND, + "group vanished", + aws_code="ResourceNotFoundException", + ), + ) + result = _run(call) + assert result["ok"] is False + assert result["code"] == NOT_FOUND + + +def test_generic_start_query_error_propagates(): + call = ScriptedLogs() + call.queue("DescribeLogGroups", _dlg(_group("/aws/lambda/fn"))) + call.queue("StartQuery", err_result(PERMISSION_DENIED, "denied")) + result = _run(call) + assert result["ok"] is False + assert result["code"] == PERMISSION_DENIED + assert any("logs:StartQuery" in g for g in result["guidance"]) + + +def test_get_query_results_throttled_then_completes(): + call = ScriptedLogs() + call.queue("DescribeLogGroups", _dlg(_group("/aws/lambda/fn"))) + call.queue("StartQuery", ok_result({"queryId": "q1"})) + call.queue("GetQueryResults", err_result(THROTTLED, "slow")) + call.queue("GetQueryResults", _qr(rows=[{"@ptr": "p1"}])) + result = _run(call) + assert result["ok"] is True + assert call.ops().count("GetQueryResults") == 2 + + +def test_get_query_results_error_propagates(): + call = ScriptedLogs() + call.queue("DescribeLogGroups", _dlg(_group("/aws/lambda/fn"))) + call.queue("StartQuery", ok_result({"queryId": "q1"})) + call.queue("GetQueryResults", err_result(PERMISSION_DENIED, "denied")) + result = _run(call) + assert result["ok"] is False + assert result["code"] == PERMISSION_DENIED + assert any("may still be running" in g for g in result["guidance"]) + + +def test_credentials_expired_is_flagged_retriable(): + call = ScriptedLogs() + call.queue("DescribeLogGroups", _dlg(_group("/aws/lambda/fn"))) + call.queue("StartQuery", ok_result({"queryId": "q1"})) + call.queue( + "GetQueryResults", + err_result(SERVICE_ERROR, "expired", aws_code="ExpiredTokenException"), + ) + result = _run(call) + assert result["code"] == CREDENTIALS_EXPIRED + assert result["retriable"] is True + + +def test_poll_timeout_stops_query(): + call = ScriptedLogs() + call.queue("DescribeLogGroups", _dlg(_group("/aws/lambda/fn"))) + call.queue("StartQuery", ok_result({"queryId": "q1"})) + call.queue("GetQueryResults", _qr(status="Running")) + call.queue("StopQuery", ok_result({})) + # max_polls=1 forces the poller to exhaust after a single non-terminal poll. + result = _run(call, max_polls=1) + assert result["ok"] is False + assert result["code"] == TIMEOUT + assert "StopQuery" in call.ops() + + +def test_terminal_failed_status_is_service_error(): + call = ScriptedLogs() + call.queue("DescribeLogGroups", _dlg(_group("/aws/lambda/fn"))) + call.queue("StartQuery", ok_result({"queryId": "q1"})) + call.queue("GetQueryResults", _qr(status="Failed")) + result = _run(call) + assert result["ok"] is False + assert result["code"] == SERVICE_ERROR + + +def test_terminal_timeout_status(): + call = ScriptedLogs() + call.queue("DescribeLogGroups", _dlg(_group("/aws/lambda/fn"))) + call.queue("StartQuery", ok_result({"queryId": "q1"})) + call.queue("GetQueryResults", _qr(status="Timeout")) + result = _run(call) + assert result["code"] == TIMEOUT + assert result["retriable"] is True + assert any("time out" in g for g in result["guidance"]) + + +def test_terminal_cancelled_status_is_conflict(): + call = ScriptedLogs() + call.queue("DescribeLogGroups", _dlg(_group("/aws/lambda/fn"))) + call.queue("StartQuery", ok_result({"queryId": "q1"})) + call.queue("GetQueryResults", _qr(status="Cancelled")) + result = _run(call) + assert result["code"] == CONFLICT + + +def test_statistics_are_summed_and_floated(): + call = ScriptedLogs() + call.queue("DescribeLogGroups", _dlg(_group("/aws/lambda/fn"))) + call.queue("StartQuery", ok_result({"queryId": "q1"})) + call.queue( + "GetQueryResults", + _qr( + rows=[{"@ptr": "p1"}], + stats={"recordsMatched": 5, "bytesScanned": 100}, + ), + ) + result = _run(call) + assert result["value"]["statistics"]["recordsMatched"] == 5.0 + assert result["value"]["statistics"]["bytesScanned"] == 100.0 + + +# --------------------------------------------------------------- splitting / caps + + +def test_aggregate_query_capped_reports_without_splitting(monkeypatch): + monkeypatch.setattr(mod, "_ROW_CAP", 2) + call = ScriptedLogs() + call.queue("DescribeLogGroups", _dlg(_group("/aws/lambda/fn"))) + call.queue("StartQuery", ok_result({"queryId": "q1"})) + call.queue( + "GetQueryResults", + _qr( + rows=[{"g": "a"}, {"g": "b"}], + stats={"recordsMatched": 99}, + ), + ) + result = _run(call, query="stats count() by g") + assert result["value"]["is_aggregate"] is True + assert result["value"]["split_count"] == 0 + assert any( + "Aggregate query returned" in p + for p in result["value"]["partial_reasons"] + ) + + +def test_max_splits_reached_reports_partial(monkeypatch): + monkeypatch.setattr(mod, "_ROW_CAP", 2) + call = ScriptedLogs() + call.queue("DescribeLogGroups", _dlg(_group("/aws/lambda/fn"))) + call.queue("StartQuery", ok_result({"queryId": "q1"})) + call.queue( + "GetQueryResults", + _qr(rows=[{"@ptr": "a"}, {"@ptr": "b"}], stats={"recordsMatched": 9}), + ) + # max_splits=0 forbids splitting, so the cap is reported as partial. + result = _run(call, query="fields @message", max_splits=0) + assert result["ok"] is True + assert result["is_exhaustive"] is False + assert any( + "max_splits=0 was reached" in p + for p in result["value"]["partial_reasons"] + ) + + +def test_max_queries_reached_breaks_with_partial(monkeypatch): + monkeypatch.setattr(mod, "_ROW_CAP", 2) + call = ScriptedLogs() + call.queue("DescribeLogGroups", _dlg(_group("/aws/lambda/fn"))) + call.queue("StartQuery", ok_result({"queryId": "q1"})) + call.queue( + "GetQueryResults", + _qr(rows=[{"@ptr": "a"}, {"@ptr": "b"}], stats={"recordsMatched": 9}), + ) + # The first segment caps and schedules a split, but max_queries=1 stops us + # before the sub-windows run. + result = _run(call, query="fields @message", max_queries=1) + assert result["ok"] is True + assert any( + "max_queries=1" in p for p in result["value"]["partial_reasons"] + ) + + +def test_explicit_limit_disables_splitting(monkeypatch): + monkeypatch.setattr(mod, "_ROW_CAP", 2) + call = ScriptedLogs() + call.queue("DescribeLogGroups", _dlg(_group("/aws/lambda/fn"))) + call.queue("StartQuery", ok_result({"queryId": "q1"})) + call.queue( + "GetQueryResults", + _qr(rows=[{"@ptr": "a"}, {"@ptr": "b"}], stats={"recordsMatched": 9}), + ) + result = _run(call, query="fields @message | limit 5") + assert result["value"]["split_count"] == 0 + assert any("explicit 'limit'" in g for g in result["guidance"]) + # No limit parameter is forced onto StartQuery when the query sets one. + start = call.calls[1][1] + assert "limit" not in start + + +def test_row_output_truncated_at_max_rows(): + call = ScriptedLogs() + call.queue("DescribeLogGroups", _dlg(_group("/aws/lambda/fn"))) + call.queue("StartQuery", ok_result({"queryId": "q1"})) + rows = [{"@ptr": f"p{i}"} for i in range(5)] + call.queue("GetQueryResults", _qr(rows=rows, stats={"recordsMatched": 5})) + result = _run(call, max_rows=2) + assert result["value"]["row_count"] == 2 + assert any( + "Row output truncated" in p for p in result["value"]["partial_reasons"] + ) + + +def test_sort_across_subwindows_warns(monkeypatch): + monkeypatch.setattr(mod, "_ROW_CAP", 2) + call = ScriptedLogs() + call.queue("DescribeLogGroups", _dlg(_group("/aws/lambda/fn"))) + # First segment caps and splits into two sub-windows, each small. + call.queue("StartQuery", ok_result({"queryId": "q1"})) + call.queue( + "GetQueryResults", + _qr(rows=[{"@ptr": "a"}, {"@ptr": "b"}], stats={"recordsMatched": 9}), + ) + call.queue("StartQuery", ok_result({"queryId": "q2"})) + call.queue( + "GetQueryResults", + _qr(rows=[{"@ptr": "c"}], stats={"recordsMatched": 1}), + ) + call.queue("StartQuery", ok_result({"queryId": "q3"})) + call.queue( + "GetQueryResults", + _qr(rows=[{"@ptr": "d"}], stats={"recordsMatched": 1}), + ) + result = _run(call, query="fields @message | sort @timestamp desc") + assert any("sort' order holds" in g for g in result["guidance"]) + + +def test_more_than_50_groups_are_batched(monkeypatch): + call = ScriptedLogs() + many = [_group(f"/svc/fn-{i}") for i in range(51)] + call.queue("DescribeLogGroups", _dlg(*many)) + # One StartQuery + GetQueryResults per batch (51 groups -> 2 batches). + call.queue("StartQuery", ok_result({"queryId": "q1"})) + call.queue("GetQueryResults", _qr(rows=[{"c": "1"}])) + call.queue("StartQuery", ok_result({"queryId": "q2"})) + call.queue("GetQueryResults", _qr(rows=[{"c": "2"}])) + result = _run(call, groups="/svc/*", query="stats count()") + assert result["ok"] is True + assert call.ops().count("StartQuery") == 2 + assert any("exceed the 50-per-query" in g for g in result["guidance"]) + assert any("several batches" in g for g in result["guidance"]) + + +# --------------------------------------------------------------- never-raises guard + + +def test_internal_error_is_normalized_not_raised(): + def exploding(operation, params): + raise RuntimeError("kaboom") + + result = mod.run_insights_query( + exploding, + "/aws/lambda/fn", + "fields @message", + "2026-10-28T10:00:00Z", + "2026-10-28T11:00:00Z", + now=NOW, + sleep=lambda _s: None, + clock=_clock(), + ) + assert result["ok"] is False + assert result["code"] == SERVICE_ERROR + assert "internal error" in result["error"] + + +# --------------------------------------------------------------- remaining branches + + +def test_propagate_generic_code_has_no_extra_guidance(): + # A plain SERVICE_ERROR from StartQuery goes through _propagate without the + # credentials/permission-specific guidance branches. + call = ScriptedLogs() + call.queue("DescribeLogGroups", _dlg(_group("/aws/lambda/fn"))) + call.queue("StartQuery", err_result(SERVICE_ERROR, "transient blip")) + result = _run(call) + assert result["ok"] is False + assert result["code"] == SERVICE_ERROR + assert not any("lacks" in g for g in result["guidance"]) + + +def test_stage_beginning_with_regex_is_advisory_not_fatal(): + # A '/' at the very start of a stage opens a regex literal (prev is empty); + # the resulting stage is not a known command, but that is advisory now, not + # a local rejection -- the query still runs. + call = ScriptedLogs() + call.queue("DescribeLogGroups", _dlg(_group("/aws/lambda/fn"))) + call.queue("StartQuery", ok_result({"queryId": "q1"})) + call.queue("GetQueryResults", _qr(rows=[{"@ptr": "p1"}])) + result = _run(call, query="fields @message | /oops/") + assert result["ok"] is True + assert "StartQuery" in call.ops() + assert any("not recognized by this tool" in g for g in result["guidance"]) + + +def test_escaped_slash_inside_regex_literal(): + call = ScriptedLogs() + call.queue("DescribeLogGroups", _dlg(_group("/aws/lambda/fn"))) + call.queue("StartQuery", ok_result({"queryId": "q1"})) + call.queue("GetQueryResults", _qr(rows=[{"@ptr": "p1"}])) + # The backslash escapes the '/', so the regex does not close early and the + # trailing '|' is still inside it (two stages, both valid). + result = _run( + call, + query="fields @message | filter @message like /a\\/b|c/", + ) + assert result["ok"] is True + + +def test_malformed_query_without_position_info(): + call = ScriptedLogs() + call.queue("DescribeLogGroups", _dlg(_group("/aws/lambda/fn"))) + call.queue( + "StartQuery", + err_result( + SERVICE_ERROR, + "the query could not be compiled", + aws_code="MalformedQueryException", + ), + ) + result = _run(call) + assert result["ok"] is False + assert result["code"] == "INVALID_INPUT" + # No position anywhere -> offset is None and no token is extracted. + assert result["value"]["offset"] is None + assert result["value"]["token"] == "" + + +def test_whitespace_group_entry_rejected(): + call = ScriptedLogs() + result = _run(call, groups=[" "]) + assert result["ok"] is False + assert result["param_error"] is True + assert call.calls == [] + + +def test_non_list_result_rows_are_skipped(): + call = ScriptedLogs() + call.queue("DescribeLogGroups", _dlg(_group("/aws/lambda/fn"))) + call.queue("StartQuery", ok_result({"queryId": "q1"})) + # The second entry is not a list (malformed) and must be skipped, not crash. + body = ok_result( + { + "status": "Complete", + "results": [ + [{"field": "@ptr", "value": "p1"}], + "garbage-not-a-list", + ], + "statistics": {}, + } + ) + call.queue("GetQueryResults", body) + result = _run(call) + assert result["ok"] is True + assert result["value"]["row_count"] == 1 + + +def test_duplicate_group_entries_are_deduped(): + call = ScriptedLogs() + call.queue("DescribeLogGroups", _dlg(_group("/aws/lambda/fn"))) + call.queue("StartQuery", ok_result({"queryId": "q1"})) + call.queue("GetQueryResults", _qr(rows=[{"@ptr": "p1"}])) + # The same group listed twice is normalized once -> one DescribeLogGroups. + result = _run(call, groups=["/aws/lambda/fn", "/aws/lambda/fn"]) + assert result["ok"] is True + assert call.ops().count("DescribeLogGroups") == 1 + + +def test_subsecond_window_clamps_to_empty(): + call = ScriptedLogs() + call.queue("DescribeLogGroups", _dlg(_group("/aws/lambda/fn"))) + # start < end by 800ms, but both ceil to the same whole second -> the + # clamped integer-second window is empty. + result = _run_window( + call, + "2026-10-28T10:00:00.100000Z", + "2026-10-28T10:00:00.900000Z", + ) + assert result["ok"] is False + assert "empty" in result["error"] + + +def test_later_subwindow_failure_notes_completed_subqueries(monkeypatch): + monkeypatch.setattr(mod, "_ROW_CAP", 2) + call = ScriptedLogs() + call.queue("DescribeLogGroups", _dlg(_group("/aws/lambda/fn"))) + # Parent segment caps and splits into two sub-windows. + call.queue("StartQuery", ok_result({"queryId": "q1"})) + call.queue( + "GetQueryResults", + _qr(rows=[{"@ptr": "a"}, {"@ptr": "b"}], stats={"recordsMatched": 9}), + ) + # First sub-window succeeds... + call.queue("StartQuery", ok_result({"queryId": "q2"})) + call.queue( + "GetQueryResults", + _qr(rows=[{"@ptr": "c"}], stats={"recordsMatched": 1}), + ) + # ...the second sub-window fails after a completed one exists. + call.queue("StartQuery", ok_result({"queryId": "q3"})) + call.queue("GetQueryResults", err_result(PERMISSION_DENIED, "denied")) + result = _run(call, query="fields @message") + assert result["ok"] is False + assert any( + "had completed before the failure" in g for g in result["guidance"] + ) + + +def test_duplicate_ptr_across_subwindows_deduped(monkeypatch): + monkeypatch.setattr(mod, "_ROW_CAP", 2) + call = ScriptedLogs() + call.queue("DescribeLogGroups", _dlg(_group("/aws/lambda/fn"))) + call.queue("StartQuery", ok_result({"queryId": "q1"})) + call.queue( + "GetQueryResults", + _qr(rows=[{"@ptr": "a"}, {"@ptr": "b"}], stats={"recordsMatched": 9}), + ) + # Both sub-windows return the SAME @ptr; the second occurrence is skipped. + call.queue("StartQuery", ok_result({"queryId": "q2"})) + call.queue( + "GetQueryResults", + _qr(rows=[{"@ptr": "dup"}], stats={"recordsMatched": 1}), + ) + call.queue("StartQuery", ok_result({"queryId": "q3"})) + call.queue( + "GetQueryResults", + _qr(rows=[{"@ptr": "dup"}], stats={"recordsMatched": 1}), + ) + result = _run(call, query="fields @message") + ptrs = [r["@ptr"] for r in result["value"]["rows"]] + assert ptrs == ["dup"] From 3bbf9c097d2f2571cf0ed68d0915c56b6670d869 Mon Sep 17 00:00:00 2001 From: Claudiu Popa Date: Thu, 29 Oct 2026 11:23:33 +0100 Subject: [PATCH 3/4] Add CloudWatch alarm metric sweep Generated by AI tools, and reviewed by Claudiu Popa Co-authored-by: Baris Kurt --- .../cloudwatch_alarm_metric_sweep.py | 1552 +++++++++++++++++ awscli/handlers_registry.py | 6 +- .../cloudwatch/test_alarm_metric_sweep.py | 148 ++ .../customizations/cloudwatch/__init__.py | 12 + .../cloudwatch/test_alarm_metric_sweep.py | 808 +++++++++ .../cloudwatch/test_resolve_alarm.py | 212 +++ 6 files changed, 2737 insertions(+), 1 deletion(-) create mode 100644 awscli/customizations/cloudwatch_alarm_metric_sweep.py create mode 100644 tests/functional/cloudwatch/test_alarm_metric_sweep.py create mode 100644 tests/unit/customizations/cloudwatch/__init__.py create mode 100644 tests/unit/customizations/cloudwatch/test_alarm_metric_sweep.py create mode 100644 tests/unit/customizations/cloudwatch/test_resolve_alarm.py diff --git a/awscli/customizations/cloudwatch_alarm_metric_sweep.py b/awscli/customizations/cloudwatch_alarm_metric_sweep.py new file mode 100644 index 000000000000..92852e76999d --- /dev/null +++ b/awscli/customizations/cloudwatch_alarm_metric_sweep.py @@ -0,0 +1,1552 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"). You +# may not use this file except in compliance with the License. A copy of +# the License is located at +# +# http://aws.amazon.com/apache2.0/ +# +# or in the "license" file accompanying this file. This file is +# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF +# ANY KIND, either express or implied. See the License for the specific +# language governing permissions and limitations under the License. +"""``aws cloudwatch alarm-metric-sweep``: investigate a CloudWatch alarm in ONE call. + +Replaces the hand-rolled ``describe-alarms -> describe-alarm-history -> list-metrics +-> get-metric-data`` chain that callers otherwise assemble by hand. A single +read-only verb resolves the alarm, reads its state history, enumerates every +series that shares the alarm metric's dimension-name set (plus error-rate +siblings), and pulls per-series timelines on a period-aligned grid. + +This command returns a result object from every call and never raises; it keeps +an "errors as data" control flow via the shared +:mod:`awscli.customizations.compositereadonly.helpers` bridge and maps the final result +to a CLI exit code. + +Behaviour contract: + +* **Resolution** -- an exact alarm name is tried first and never silently + rerouted; a miss falls back to a bounded, deterministic prefix lookup (shared + :func:`resolve_alarm`), returning a clear not-found / ambiguous result rather + than a broad fuzzy scan. +* **Pagination / bounds** -- ``DescribeAlarmHistory``, ``ListMetrics`` and + ``GetMetricData`` all follow ``NextToken`` under a ``--max-pages`` cap; a + partial read degrades ``is_exhaustive`` instead of failing. +* **Period / time-window** -- the window defaults to the last 3h, is aligned to + the period, and the period is raised for retention and the 1,440-points-per- + series budget; every adjustment is reported in ``guidance``. +* **Output capping** -- series are capped by ``--max-series`` and per-series + datapoints by ``--max-datapoints``; a global ``--max-total-datapoints`` budget + drops later datapoint arrays (keeping their summaries) so a wide sweep returns + a bounded response. Every adjustment is recorded, never silent. + +The command is strictly read-only: it only issues ``Describe*`` / ``Get*`` / +``List*`` calls, and the shared :func:`call` bridge refuses any mutating +operation by construction. +""" + +from __future__ import annotations + +import difflib +import json +import math +import re +import time +from datetime import datetime, timedelta, timezone +from typing import Any, Callable + +from awscli.customizations.commands import BasicCommand +from awscli.customizations.compositereadonly.helpers import ( + NOT_FOUND, + SERVICE_ERROR, + call, + check_non_empty_str, + check_positive_int, + create_read_client, + emit_result, + err_result, + invalid_input, + ok_result, + paginate, +) + +UTC = timezone.utc + +# --- GetMetricData bounds ---------------------------------------------------- +# GetMetricData accepts at most 500 MetricDataQueries and ~100,800 datapoints +# per call; a single series is capped at 1,440 points (the window is re-perioded +# to fit). These keep one call from running away. +_MAX_QUERIES = 500 +_MAX_POINTS_PER_CALL = 100_800 +_MAX_POINTS = 1_440 +# Default global budget for emitted datapoints across the enumerated series. A +# wide alarm (many dimension values x sibling metrics) can list hundreds of +# series; without a ceiling the JSON grows without bound. 0 disables. +_DEFAULT_MAX_TOTAL_DATAPOINTS = 50_000 +# Period ladder used when re-perioding to honour the per-series point budget. +_LADDER = (60, 120, 300, 600, 900, 1800, 3600, 10800, 21600, 43200, 86400) +# Accepted CloudWatch statistics (standard + percentile/trimmed forms). +_STAT_RE = re.compile( + r"^(SampleCount|Average|Sum|Minimum|Maximum|IQM|p\d{1,2}(\.\d+)?|p100|" + r"(tm|wm|tc|ts)\d{1,2}(\.\d+)?|(TM|WM|TC|TS|PR)\([^)]*\))$" +) +# An alarm/metric whose name looks like an error-rate target; only then do we +# also sweep its sibling status-code/success metrics (keeps the sweep focused). +_ERR_RE = re.compile( + r"(error|fault|fail|5xx|4xx|5\d\d|success|availability)", re.IGNORECASE +) +_SIB_RE = re.compile( + r"(error|fault|fail|5xx|4xx|5\d\d|4\d\d|2xx|success|request|count|" + r"invocation|throttl)", + re.IGNORECASE, +) +# Pull child alarm names out of a composite alarm's rule expression. +_CHILD_RE = re.compile( + r'(?:ALARM|OK|INSUFFICIENT_DATA)\(\s*"?([^")]+?)"?\s*\)' +) +# GetMetricData per-page status precedence (a later page's status supersedes an +# earlier PartialData, but never downgrades a hard Forbidden/InternalError). +_RANK = { + "Complete": 0, + "PartialData": 1, + "Forbidden": 2, + "InternalError": 3, + "NoResult": 3, +} +# Alarm fields surfaced verbatim in the target definition. +_ALARM_FIELDS = ( + "AlarmName", + "AlarmArn", + "AlarmDescription", + "StateValue", + "StateReason", + "StateUpdatedTimestamp", + "Namespace", + "MetricName", + "Dimensions", + "Statistic", + "ExtendedStatistic", + "Period", + "Unit", + "EvaluationPeriods", + "DatapointsToAlarm", + "Threshold", + "ComparisonOperator", + "TreatMissingData", + "Metrics", + "AlarmRule", +) + +CallFn = Callable[[str, dict[str, Any]], dict[str, Any]] + + +# ----------------------------------------------------------------- time helpers +def _iso(epoch: float) -> str: + return datetime.fromtimestamp(int(epoch), UTC).strftime( + "%Y-%m-%dT%H:%M:%SZ" + ) + + +def _dt(v: Any) -> datetime | None: + if isinstance(v, datetime): + return v if v.tzinfo else v.replace(tzinfo=UTC) + if isinstance(v, str) and v.strip(): + text = v.strip() + text = text[:-1] + "+00:00" if text[-1] in "Zz" else text + try: + d = datetime.fromisoformat(text) + except ValueError: + return None + return d if d.tzinfo else d.replace(tzinfo=UTC) + return None + + +def _js(o: Any) -> Any: + """Recursively JSON-normalize datetimes so the payload is plain CLI JSON.""" + if isinstance(o, datetime): + return _iso(_dt(o).timestamp()) # type: ignore[union-attr] + if isinstance(o, dict): + return {str(k): _js(v) for k, v in o.items()} + if isinstance(o, (list, tuple)): + return [_js(v) for v in o] + return o + + +def _pages( + call_op: CallFn, op: str, params: dict[str, Any], key: str, max_pages: int +) -> tuple[list[dict[str, Any]], bool, dict[str, Any] | None]: + """Follow NextToken to the end (or max_pages). Returns (items, complete, err). + + ``complete`` is False when the bounded loop stopped on ``max_pages`` with a + token still outstanding, so the caller can mark the sweep partial. + """ + # Thin service-scoped wrapper over the shared helper: ListMetrics / + # DescribeAlarmHistory / GetMetricData all return dict-shaped items, so a + # non-dict entry is filtered out as it was before. + return paginate( + call_op, + op, + params, + key, + max_pages, + item_filter=lambda i: isinstance(i, dict), + ) + + +# Bounded paging / fuzzy cutoff for alarm resolution so a lookup can never run +# away on an account with thousands of alarms. +_MAX_ALARM_PAGES = 5 +_FUZZY_CUTOFF = 0.6 +_ALARM_TYPES = ["MetricAlarm", "CompositeAlarm"] + + +def _alarms_from(value: dict[str, Any]) -> list[dict[str, Any]]: + """Flatten DescribeAlarms metric + composite alarms into one list.""" + return list(value.get("MetricAlarms") or []) + list( + value.get("CompositeAlarms") or [] + ) + + +def _describe_prefix( + call_op: CallFn, prefix: str +) -> tuple[list[dict[str, Any]], dict[str, Any] | None]: + """DescribeAlarms by name prefix, following NextToken (bounded).""" + alarms: list[dict[str, Any]] = [] + token: str | None = None + for _ in range(_MAX_ALARM_PAGES): + params: dict[str, Any] = { + "AlarmNamePrefix": prefix, + "AlarmTypes": _ALARM_TYPES, + } + if token: + params["NextToken"] = token + res = call_op("DescribeAlarms", params) + if not res["ok"]: + return alarms, res + value = res.get("value") or {} + alarms.extend(_alarms_from(value)) + token = value.get("NextToken") + if not token: + break + return alarms, None + + +def resolve_alarm(call_op: CallFn, name: str) -> dict[str, Any]: + """Resolve a CloudWatch alarm by exact name, then prefix, then fuzzy match. + + ``call_op(operation, params)`` is the read-only call bridge bound to a + CloudWatch client. Resolution order: + + 1. **Exact** -- ``DescribeAlarms AlarmNames=[name]``. A hit short-circuits + with no fallback, so an exact name is never silently rerouted. + 2. **Prefix** -- ``DescribeAlarms AlarmNamePrefix=name`` (bounded paging). A + single candidate is used; a case-insensitive exact match among several is + preferred before considering it ambiguous. + 3. **Fuzzy** -- a bounded, deterministic ``difflib`` ranking over the prefix + candidates only (no extra scan). The closest fuzzy match above the cutoff + is used; if no match, NOT_FOUND with the candidate list. + + The value is ``{"alarm": , "kind": "metric_alarm"| + "composite_alarm", "match": "exact"|"prefix"|"fuzzy", "guidance": [...]}``. + """ + # 1. Exact. + exact = call_op( + "DescribeAlarms", {"AlarmNames": [name], "AlarmTypes": _ALARM_TYPES} + ) + if not exact["ok"]: + return exact + metric = (exact.get("value") or {}).get("MetricAlarms") or [] + composite = (exact.get("value") or {}).get("CompositeAlarms") or [] + if metric or composite: + alarm = (metric or composite)[0] + return ok_result( + { + "alarm": alarm, + "kind": "metric_alarm" if metric else "composite_alarm", + "match": "exact", + "guidance": [], + } + ) + + # 2. Prefix. + candidates, err = _describe_prefix(call_op, name) + if err is not None: + return err + if not candidates: + return err_result( + NOT_FOUND, + f"No alarm named or prefixed by {name!r}.", + key=name, + guidance=[ + "Alarm names are exact and case-sensitive; list candidates with " + "'aws cloudwatch describe-alarms --alarm-name-prefix '.", + "Alarms are regional: pass --region.", + ], + ) + by_name = {a.get("AlarmName"): a for a in candidates if a.get("AlarmName")} + if len(candidates) == 1: + return _prefix_hit(candidates[0], name) + # Several share the prefix: prefer a unique case-insensitive exact match. + ci_matches = [ + a + for a in candidates + if (a.get("AlarmName") or "").casefold() == name.casefold() + ] + if len(ci_matches) == 1: + return _prefix_hit(ci_matches[0], name, note="case-insensitive exact") + + # 3. Fuzzy over the already-fetched candidate names (bounded, deterministic). + names_sorted = sorted(by_name) + close = difflib.get_close_matches( + name, names_sorted, n=2, cutoff=_FUZZY_CUTOFF + ) + if close: + chosen = by_name[close[0]] + return _prefix_hit(chosen, name, match="fuzzy") + # Ambiguous: report the candidates rather than guessing. + listed = ", ".join(names_sorted[:10]) + return err_result( + NOT_FOUND, + f"{name!r} is ambiguous: {len(candidates)} alarms share this prefix " + "and none is an unambiguous match.", + key=name, + guidance=[ + f"Matching alarms: {listed}" + + (" ..." if len(names_sorted) > 10 else ""), + "Pass --alarm-name with one of these exact names.", + ], + ) + + +def _prefix_hit( + alarm: dict[str, Any], + requested: str, + *, + match: str = "prefix", + note: str = "", +) -> dict[str, Any]: + kind = "composite_alarm" if alarm.get("AlarmRule") else "metric_alarm" + resolved = alarm.get("AlarmName") + detail = f" ({note})" if note else "" + guidance = ( + [] + if match == "exact" + else [ + f"No alarm named exactly {requested!r}; resolved to {resolved!r} " + f"by {match} match{detail}." + ] + ) + return ok_result( + {"alarm": alarm, "kind": kind, "match": match, "guidance": guidance} + ) + + +# ------------------------------------------------------------------- validation +def _validate( + alarm_name: Any, + stat: Any, + pins: Any, + max_series: Any, + max_pages: Any, + max_datapoints: Any, + max_total_datapoints: Any, +) -> list[str]: + """Validate inputs that are checkable before any API call.""" + p: list[str] = [] + name_err = check_non_empty_str("alarm_name", alarm_name, required=True) + if name_err is not None: + p.append(f"{name_err} Pass the exact alarm name, or a unique prefix.") + if stat is not None and ( + not isinstance(stat, str) or not _STAT_RE.match(stat) + ): + p.append( + f"stat={stat!r} is not a CloudWatch statistic; use Sum, Average, " + "Minimum, Maximum, SampleCount or a percentile like p99 (lowercase p)." + ) + # pin_dimensions is a list of dimension NAMES; a non-string entry is the + # "malformed pin" case and is rejected rather than silently dropped. + if pins is not None and ( + not isinstance(pins, (list, tuple)) + or not all(isinstance(x, str) for x in pins) + ): + p.append( + "pin_dimensions must be a list of dimension names, e.g. " + "['AvailabilityZone']." + ) + for name, v, hi in ( + ("max_series", max_series, 2000), + ("max_pages", max_pages, 500), + ): + err = check_positive_int(name, v, lo=1, hi=hi) + if err is not None: + p.append(err) + for name, v in ( + ("max_datapoints", max_datapoints), + ("max_total_datapoints", max_total_datapoints), + ): + err = check_positive_int(name, v, lo=0) + if err is not None: + p.append(f"{err} (0 disables the cap).") + return p + + +def _window( + start_time: Any, end_time: Any, period_s: Any, now: float +) -> dict[str, Any] | list[str]: + """Resolve/validate the time window and period, returning errors as a list. + + The returned dict carries the aligned epochs, the effective period (raised + for retention and for the per-series point budget) and every adjustment/alignment + note so the caller can surface them in ``guidance``. + """ + if ( + isinstance(period_s, bool) + or not isinstance(period_s, int) + or period_s <= 0 + ): + return [ + f"period_s must be a positive int (got {period_s!r}); use 60, 300 or 3600." + ] + errs: list[str] = [] + if period_s not in (1, 5, 10, 30) and period_s % 60: + errs.append( + f"period_s={period_s} is invalid: use 1, 5, 10, 30 or a multiple of " + f"60 (e.g. {max(60, round(period_s / 60) * 60)})." + ) + s, e = _dt(start_time), _dt(end_time) + for name, raw, parsed in ( + ("start_time", start_time, s), + ("end_time", end_time, e), + ): + if raw not in (None, "") and parsed is None: + errs.append( + f"{name}={raw!r} is not ISO-8601; use e.g. '2024-01-15T21:00:00Z'." + ) + if errs: + return errs + now_dt = datetime.fromtimestamp(now, UTC) + notes: list[str] = [] + if e is None: + e = now_dt + notes.append(f"end_time defaulted to now ({_iso(now)}).") + if s is None: + # Default to a 3h look-back. + s = e - timedelta(hours=3) + notes.append("start_time defaulted to end_time - 3h.") + if s >= e: + return [ + f"start_time ({_iso(s.timestamp())}) must be before end_time " + f"({_iso(e.timestamp())})." + ] + if e > now_dt: + e = now_dt + notes.append( + f"end_time is in the future: clamped to now ({_iso(now)})." + ) + if s >= e: + return [ + "start_time is in the future; choose a window ending at or before now." + ] + age = now - s.timestamp() + if age > 455 * 86400: + return [ + "start_time is older than CloudWatch's 455-day retention; move start_time later." + ] + # CloudWatch only retains sub-minute/minute resolution for a limited age; + # raise the period to the finest resolution still available for this window. + min_p = ( + 3600 + if age > 63 * 86400 + else 300 + if age > 15 * 86400 + else 60 + if age > 3 * 3600 + else 1 + ) + period, clamped = period_s, [] + if period < min_p: + clamped.append( + f"period_s raised from {period} to {min_p}: data this old is only " + f"retained at {min_p}s." + ) + period = min_p + span = e.timestamp() - s.timestamp() + # Keep each series within the per-series point budget by widening the period. + if math.ceil(span / period) > _MAX_POINTS: + new = next( + ( + c + for c in _LADDER + if c >= period and math.ceil(span / c) <= _MAX_POINTS + ), + 86400, + ) + clamped.append( + f"period_s raised from {period} to {new} to keep <= {_MAX_POINTS} " + "points per series; narrow the window for finer resolution." + ) + period = new + # Align the window to period boundaries so timestamps line up across series. + s_ep = int(math.floor(s.timestamp() / period) * period) + e_ep = max(int(math.ceil(e.timestamp() / period) * period), s_ep + period) + if (s_ep, e_ep) != (s.timestamp(), e.timestamp()): + notes.append( + f"Window aligned to {period}s boundaries: {_iso(s_ep)} .. {_iso(e_ep)}." + ) + return { + "start": _iso(s_ep), + "end": _iso(e_ep), + "start_epoch": s_ep, + "end_epoch": e_ep, + "period_s_requested": period_s, + "period_s": period, + "min_period_s": min_p, + "expected_points_per_series": (e_ep - s_ep) // period, + "clamped": clamped, + "notes": notes, + } + + +# ------------------------------------------------------------ step 1: target +def _target_metric(ms: Any, qid: str) -> dict[str, Any] | None: + m = (ms or {}).get("Metric") or {} if isinstance(ms, dict) else {} + if not m.get("Namespace") or not m.get("MetricName"): + return None + return { + "namespace": m["Namespace"], + "metric_name": m["MetricName"], + "dimensions": list(m.get("Dimensions") or []), + "stat": ms.get("Stat") or "Average", + "unit": ms.get("Unit"), + "source_id": qid, + } + + +def _load_from_alarm( + alarm: dict[str, Any], + kind: str, + match: str, + resolve_guidance: list[str], + requested: str, +) -> dict[str, Any]: + """Build the sweep target + replay queries from a resolved alarm. + + ``resolve_alarm`` has already issued ``DescribeAlarms``; this takes the + returned alarm dict and shapes it into the command's target descriptor. + """ + name = alarm.get("AlarmName", requested) + out: dict[str, Any] = { + "target": { + "kind": kind, + "alarm_name": name, + "alarm_arn": alarm.get("AlarmArn"), + "match": match, + "definition": _js( + { + k: alarm[k] + for k in _ALARM_FIELDS + if alarm.get(k) is not None + } + ), + }, + "originals": [], + "metrics": [], + "history": [name], + "notes": list(resolve_guidance), + } + if kind == "composite_alarm": + # A composite alarm owns no metrics of its own; surface its children so + # the caller can re-run the sweep against one of them. + children = list( + dict.fromkeys(_CHILD_RE.findall(alarm.get("AlarmRule") or "")) + ) + out["target"]["child_alarms"] = children + out["notes"].append( + "Composite alarm has no metrics of its own; re-run with " + f"--alarm-name set to one of its children: {children}." + ) + return out + if alarm.get("MetricName"): + # Single-metric alarm: one replay query. + ms = { + "Metric": { + "Namespace": alarm.get("Namespace"), + "MetricName": alarm["MetricName"], + "Dimensions": list(alarm.get("Dimensions") or []), + }, + "Period": int(alarm.get("Period") or 60), + "Stat": alarm.get("ExtendedStatistic") + or alarm.get("Statistic") + or "Average", + **({"Unit": alarm["Unit"]} if alarm.get("Unit") else {}), + } + out["originals"].append( + { + "Id": "alarm_metric", + "MetricStat": ms, + "Label": name, + "ReturnData": True, + } + ) + else: + # Metric-math alarm: replay each query (expose the inputs behind the math). + for m in alarm.get("Metrics") or []: + q = { + k: json.loads(json.dumps(m[k], default=str)) + for k in ("Id", "MetricStat", "Expression", "Label", "Period") + if k in m + } + out["originals"].append({**q, "ReturnData": True}) + out["metrics"] = [ + t + for q in out["originals"] + if (t := _target_metric(q.get("MetricStat"), q["Id"])) + ] + return out + + +# ----------------------------------------------------------- step 2: history +def _history( + call_op: CallFn, names: list[str], w: dict[str, Any], max_pages: int +) -> tuple[dict[str, Any], list[str]]: + """Read alarm state transitions in the window (bounded NextToken paging).""" + transitions: list[dict[str, Any]] = [] + warns: list[str] = [] + for name in names: + items, done, err = _pages( + call_op, + "DescribeAlarmHistory", + { + "AlarmName": name, + "HistoryItemType": "StateUpdate", + "StartDate": w["start"], + "EndDate": w["end"], + "ScanBy": "TimestampAscending", + }, + "AlarmHistoryItems", + max_pages, + ) + if err is not None: + # A history read failure degrades the sweep to partial; it does not + # abort it (the metric data is still worth returning). + warns.append( + f"DescribeAlarmHistory({name}) failed ({err['code']}: " + f"{err['error']}); state history missing." + ) + elif not done: + warns.append( + f"DescribeAlarmHistory({name}) stopped at max_pages={max_pages}; " + "history is partial." + ) + for it in items: + try: + data = json.loads(it.get("HistoryData") or "{}") + except ValueError: + data = {} + new, old = data.get("newState") or {}, data.get("oldState") or {} + ts = _dt(it.get("Timestamp")) + transitions.append( + { + "alarm_name": it.get("AlarmName"), + "timestamp": _iso(ts.timestamp()) if ts else None, + "old_state": old.get("stateValue"), + "new_state": new.get("stateValue"), + "reason": new.get("stateReason"), + "summary": it.get("HistorySummary"), + "recent_datapoints": _js( + (new.get("stateReasonData") or {}).get( + "recentDatapoints" + ) + ), + } + ) + transitions.sort(key=lambda t: t["timestamp"] or "") + return ( + { + "alarm_names": names, + "transitions": transitions, + "count": len(transitions), + "is_exhaustive": not warns, + }, + warns, + ) + + +# --------------------------------------------------------- step 3: enumerate +def _spec( + ns: str, + name: str, + dims: list[dict[str, Any]], + stat: str, + unit: str | None, + role: str, +) -> dict[str, Any]: + d = sorted( + ( + {"Name": str(x.get("Name")), "Value": str(x.get("Value"))} + for x in dims + ), + key=lambda x: x["Name"], + ) + label = name + ( + " " + ",".join(f"{x['Name']}={x['Value']}" for x in d) if d else "" + ) + return { + "namespace": ns, + "metric_name": name, + "dimensions": d, + "stat": stat, + "unit": unit, + "role": role, + "is_target": role == "target", + "label": label, + } + + +def _enumerate( + call_op: CallFn, + targets: list[dict[str, Any]], + stat: str | None, + pins: set[str], + siblings: bool, + max_series: int, + max_pages: int, +) -> dict[str, Any]: + """Enumerate every series sharing each target's dimension-name set. + + ``pins`` are dimension names kept at their target value while the rest are + wildcarded, so the ListMetrics filter finds sibling dimension values. For an + error-rate target the sibling status-code/success metrics are added too. + """ + specs: list[dict[str, Any]] = [] + seen: set[tuple[Any, ...]] = set() + warns: list[str] = [] + guidance: list[str] = [] + report: list[dict[str, Any]] = [] + cache: dict[ + str, tuple[list[dict[str, Any]], bool, dict[str, Any] | None] + ] = {} + complete = True + + def add(s: dict[str, Any]) -> None: + k = ( + s["namespace"], + s["metric_name"], + tuple((d["Name"], d["Value"]) for d in s["dimensions"]), + s["stat"], + ) + if k not in seen: + seen.add(k) + specs.append(s) + + def listed(params): + ck = json.dumps(params, sort_keys=True) + if ck not in cache: + cache[ck] = _pages( + call_op, "ListMetrics", params, "Metrics", max_pages + ) + return cache[ck] + + # The target series themselves always come first. + for t in targets: + add( + _spec( + t["namespace"], + t["metric_name"], + t["dimensions"], + stat or t["stat"], + t["unit"], + "target", + ) + ) + for t in targets: + names = {str(d.get("Name")) for d in t["dimensions"]} + filt = [ + {"Name": d["Name"], "Value": d["Value"]} + if d.get("Name") in pins + else {"Name": d["Name"]} + for d in t["dimensions"] + ] + found, done, err = listed( + { + "Namespace": t["namespace"], + "MetricName": t["metric_name"], + "Dimensions": filt, + } + ) + if err is not None or not done: + complete = False + warns.append( + f"ListMetrics({t['namespace']}/{t['metric_name']}) incomplete " + f"({err['code'] if err else f'max_pages={max_pages}'}); " + "dimension enumeration is partial." + ) + other: dict[str, int] = {} + swept = 0 + for m in found: + dn = {str(d.get("Name")) for d in m.get("Dimensions") or []} + if dn == names: + swept += 1 + add( + _spec( + t["namespace"], + t["metric_name"], + m.get("Dimensions") or [], + stat or t["stat"], + t["unit"], + "dimension_sibling", + ) + ) + else: + other[",".join(sorted(dn))] = ( + other.get(",".join(sorted(dn)), 0) + 1 + ) + if other: + guidance.append( + f"{sum(other.values())} more {t['metric_name']} series use other " + f"dimension sets {sorted(other)}; they were not swept." + ) + report.append( + { + "namespace": t["namespace"], + "metric_name": t["metric_name"], + "dimension_names": sorted(names), + "pinned": sorted(names & pins), + "series_listed": swept, + "list_complete": err is None and done, + "unswept_dimension_sets": [ + { + "dimension_names": k.split(",") if k else [], + "series_count": v, + } + for k, v in sorted(other.items()) + ], + } + ) + sib_names: set[str] = set() + if siblings: + target_names = {(t["namespace"], t["metric_name"]) for t in targets} + groups = { + json.dumps( + [t["namespace"], t["dimensions"]], sort_keys=True, default=str + ): t + for t in targets + } + for t in groups.values(): + dims = [ + {"Name": d["Name"], "Value": d["Value"]} + for d in t["dimensions"] + ] + found, done, err = listed( + {"Namespace": t["namespace"], "Dimensions": dims} + ) + if err is not None or not done: + complete = False + warns.append( + f"Sibling ListMetrics in {t['namespace']} incomplete; " + "status-code metrics may be missing." + ) + want = {(d["Name"], d["Value"]) for d in dims} + for m in found: + n = str(m.get("MetricName") or "") + got = { + (str(d.get("Name")), str(d.get("Value"))) + for d in m.get("Dimensions") or [] + } + if ( + n + and (t["namespace"], n) not in target_names + and _SIB_RE.search(n) + and got == want + ): + sib_names.add(n) + add( + _spec( + t["namespace"], + n, + m.get("Dimensions") or [], + "Sum", + None, + "metric_sibling", + ) + ) + # Cap the series count last (siblings dropped first via list order). + dropped = max(0, len(specs) - max_series) + if dropped: + complete = False + guidance.append( + f"{dropped} series dropped by max_series={max_series} (siblings " + "last); raise --max-series or pin a dimension with --pin-dimensions." + ) + return { + "series": specs[:max_series], + "complete": complete, + "warnings": warns, + "guidance": guidance, + "report": { + "metrics": report, + "sibling_metric_names": sorted(sib_names), + "series_enumerated": len(specs), + "series_dropped": dropped, + }, + } + + +# -------------------------------------------------------------- step 4: data +def _fetch( + call_op: CallFn, + queries: list[dict[str, Any]], + w: dict[str, Any], + max_pages: int, +) -> tuple[dict[tuple[str, str], dict[str, Any]], bool, dict[str, Any] | None]: + """Run GetMetricData for a batch of queries, following NextToken (bounded).""" + items, done, err = _pages( + call_op, + "GetMetricData", + { + "MetricDataQueries": queries, + "StartTime": w["start"], + "EndTime": w["end"], + "ScanBy": "TimestampAscending", + }, + "MetricDataResults", + max_pages, + ) + out: dict[tuple[str, str], dict[str, Any]] = {} + for r in items: + rid = str(r.get("Id") or "") + e = out.setdefault( + (rid, str(r.get("Label") or rid)), + { + "label": str(r.get("Label") or rid), + "timestamps": [], + "values": [], + "status": "Complete", + "messages": [], + }, + ) + e["timestamps"] += r.get("Timestamps") or [] + e["values"] += r.get("Values") or [] + # PartialData on an earlier page is superseded by the final page's status. + if _RANK.get(e["status"], 0) < 2: + e["status"] = str(r.get("StatusCode") or "Complete") + e["messages"] += [ + str(m.get("Value") or m.get("Code")) + for m in r.get("Messages") or [] + ] + return out, done, err + + +def _timeline( + ts: list[Any], + vals: list[Any], + period: int, + s_ep: int, + e_ep: int, + include: bool, + max_datapoints: int, +) -> dict[str, Any]: + """Build a per-series timeline on the aligned grid. + + Missing timestamps are reported separately from zeros; trailing gaps, peak + and longest non-zero run are computed. When ``include`` is set the raw + datapoints are emitted, truncated to ``max_datapoints`` (0 = no per-series + cap) with the adjustment flagged so trimming is never silent. + """ + g0 = s_ep // period * period + grid = list(range(g0, -(-e_ep // period) * period, period)) + pts: dict[int, float] = {} + off_grid = 0 + for t, v in zip(ts, vals): + d = _dt(t) + if d is not None and isinstance(v, (int, float)) and not math.isnan(v): + ep = int(d.timestamp()) + off_grid += bool((ep - g0) % period) + pts[ep] = float(v) + missing = [g for g in grid if g not in pts] + ranges: list[list[Any]] = [] + for g in missing: + if ranges and g - ranges[-1][1] == period: + ranges[-1][1], ranges[-1][2] = g, ranges[-1][2] + 1 + else: + ranges.append([g, g, 1]) + trailing = next( + (i for i, g in enumerate(reversed(grid)) if g in pts), len(grid) + ) + run = best = 0 + for g in grid: + run = run + 1 if pts.get(g, 0.0) != 0.0 else 0 + best = max(best, run) + seq = sorted(pts.items()) + peak = max(seq, key=lambda kv: kv[1]) if seq else None + out = { + "period_s": period, + "expected_count": len(grid), + "datapoint_count": len(pts), + "missing_count": len(missing), + "missing_ranges": [[_iso(a), _iso(b), n] for a, b, n in ranges], + "zero_count": sum(v == 0.0 for v in pts.values()), + "off_grid_count": off_grid, + "trailing_missing_count": trailing if seq else len(grid), + "first_datapoint": _iso(seq[0][0]) if seq else None, + "last_datapoint": _iso(seq[-1][0]) if seq else None, + "sum": sum(v for _, v in seq) if seq else None, + "min": min(v for _, v in seq) if seq else None, + "max": peak[1] if peak else None, + "peak_timestamp": _iso(peak[0]) if peak else None, + "longest_nonzero_run": best, + "longest_nonzero_run_s": best * period, + } + if include: + dps = [[_iso(t), v] for t, v in seq] + truncated = False + # Per-series datapoint cap: keep the head (where a spike usually is). + if max_datapoints and len(dps) > max_datapoints: + dps = dps[:max_datapoints] + truncated = True + out["datapoints"] = dps + out["datapoints_truncated"] = truncated + return out + + +def _qperiod(q: dict[str, Any], fallback: int) -> int: + p = ( + (q.get("MetricStat") or {}).get("Period") + if isinstance(q.get("MetricStat"), dict) + else q.get("Period") + ) + return p if isinstance(p, int) and p > 0 else fallback + + +# ------------------------------------------------------------------- the sweep +def alarm_metric_sweep( + call_op: CallFn, + alarm_name: str | None, + *, + start_time: str | None = None, + end_time: str | None = None, + period_s: int = 300, + stat: str | None = None, + pin_dimensions: list[str] | None = None, + include_sibling_metrics: bool = True, + include_datapoints: bool = True, + max_series: int = 200, + max_pages: int = 20, + max_datapoints: int = _MAX_POINTS, + max_total_datapoints: int = _DEFAULT_MAX_TOTAL_DATAPOINTS, + clock: Callable[[], float] | None = None, + region: str | None = None, +) -> dict[str, Any]: + """Sweep a CloudWatch alarm: definition, state history, every matching series. + + ``call_op(operation, params)`` is the shared read-only call bridge. Returns a + result dict (see :mod:`...compositereadonly.helpers`); never raises. ``is_exhaustive`` + is False whenever a page cap, series/datapoint cap, or failed sub-call left + the sweep partial. + """ + clock = time.time if clock is None else clock + # Outer never-raise boundary (mirrors run_athena_query / run_insights_query / + # diagnose_ecs_service): the sweep makes many bounded sub-calls over botocore + # responses whose exact shape it cannot guarantee. Any unexpected exception + # is converted into an errors-as-data SERVICE_ERROR result so the documented + # "never raises" contract holds instead of a traceback escaping to the caller. + try: + return _run_sweep( + call_op, + alarm_name, + start_time=start_time, + end_time=end_time, + period_s=period_s, + stat=stat, + pin_dimensions=pin_dimensions, + include_sibling_metrics=include_sibling_metrics, + include_datapoints=include_datapoints, + max_series=max_series, + max_pages=max_pages, + max_datapoints=max_datapoints, + max_total_datapoints=max_total_datapoints, + clock=clock, + region=region, + ) + except Exception as exc: # noqa: BLE001 - contract: never raise + return err_result( + SERVICE_ERROR, + f"alarm-metric-sweep internal error: {exc!r}", + guidance=[ + "Unexpected response shape; retry once, then fall back to " + "cloudwatch:DescribeAlarms / get-metric-data by hand." + ], + ) + + +def _run_sweep( + call_op: CallFn, + alarm_name: str | None, + *, + start_time: str | None = None, + end_time: str | None = None, + period_s: int = 300, + stat: str | None = None, + pin_dimensions: list[str] | None = None, + include_sibling_metrics: bool = True, + include_datapoints: bool = True, + max_series: int = 200, + max_pages: int = 20, + max_datapoints: int = _MAX_POINTS, + max_total_datapoints: int = _DEFAULT_MAX_TOTAL_DATAPOINTS, + clock: Callable[[], float] = time.time, + region: str | None = None, +) -> dict[str, Any]: + """Run the sweep body; wrapped by :func:`alarm_metric_sweep`'s never-raise boundary.""" + problems = _validate( + alarm_name, + stat, + pin_dimensions, + max_series, + max_pages, + max_datapoints, + max_total_datapoints, + ) + # Window/period validation runs only once input problems are cleared. + w = ( + _window(start_time, end_time, period_s, float(clock())) + if not problems + else problems + ) + if isinstance(w, list): + # Pre-dispatch bad input -> ParamValidationError (exit code 252). + return invalid_input("; ".join(w), guidance=list(w)) + + # 1. Resolve the alarm: exact name first, then a bounded prefix lookup. + resolved = resolve_alarm(call_op, alarm_name.strip()) + if not resolved["ok"]: + # NOT_FOUND / ambiguous / permission errors are returned verbatim. + return resolved + rv = resolved["value"] + loaded = _load_from_alarm( + rv["alarm"], + rv["kind"], + rv["match"], + rv.get("guidance") or [], + alarm_name.strip(), + ) + guidance = [*w["clamped"], *w["notes"], *loaded["notes"]] + originals: list[dict[str, Any]] = loaded["originals"] + for q in originals: + holder = ( + q["MetricStat"] if isinstance(q.get("MetricStat"), dict) else q + ) + if ( + isinstance(holder.get("Period"), int) + and holder["Period"] < w["min_period_s"] + ): + guidance.append( + f"Query {q['Id']} period raised from {holder['Period']}s to " + f"{w['min_period_s']}s (retention)." + ) + holder["Period"] = w["min_period_s"] + + # 2. State history (bounded paging). + history, warnings = _history(call_op, loaded["history"], w, max_pages) + + # 3. Enumerate sibling series; sweep metric siblings only for error-rate targets. + texts = [ + loaded["target"]["alarm_name"], + *(t["metric_name"] for t in loaded["metrics"]), + *(str(q.get("Expression") or "") for q in originals), + ] + error_rate = any(_ERR_RE.search(t) for t in texts if t) + enum = _enumerate( + call_op, + loaded["metrics"], + stat, + set(pin_dimensions or []), + include_sibling_metrics and error_rate, + max_series, + max_pages, + ) + warnings += enum["warnings"] + guidance += enum["guidance"] + + # 4. Fetch metric data for the replay queries and the enumerated series. + period, s_ep, e_ep = w["period_s"], w["start_epoch"], w["end_epoch"] + errors: list[dict[str, Any]] = [] + complete, attempted = True, 0 + orig_res: dict[tuple[str, str], dict[str, Any]] = {} + if originals: + attempted += 1 + orig_res, done, err = _fetch(call_op, originals, w, max_pages) + complete &= done + if err is not None: + errors.append(err) + warnings.append( + f"GetMetricData for the original queries failed " + f"({err['code']}: {err['error']})." + ) + specs = enum["series"] + queries = [ + { + "Id": f"zs{i}", + "Label": s["label"], + "ReturnData": True, + "MetricStat": { + "Metric": { + "Namespace": s["namespace"], + "MetricName": s["metric_name"], + "Dimensions": s["dimensions"], + }, + "Period": period, + "Stat": s["stat"], + **({"Unit": s["unit"]} if s["unit"] else {}), + }, + } + for i, s in enumerate(specs) + ] + # Batch so no single GetMetricData call exceeds the per-call point budget. + batch = max( + 1, + min( + _MAX_QUERIES, + _MAX_POINTS_PER_CALL // max(1, w["expected_points_per_series"]), + ), + ) + series_res: dict[str, dict[str, Any]] = {} + for b in range(0, len(queries), batch): + attempted += 1 + res, done, err = _fetch(call_op, queries[b : b + batch], w, max_pages) + complete &= done + for (rid, _), entry in res.items(): + series_res.setdefault(rid, entry) + if err is not None: + errors.append(err) + warnings.append( + f"GetMetricData batch {b // batch + 1} failed " + f"({err['code']}: {err['error']})." + ) + if not complete: + guidance.append( + f"GetMetricData pagination stopped at max_pages={max_pages}; later " + "datapoints are missing. Raise --max-pages or narrow the window." + ) + + empty = { + "timestamps": [], + "values": [], + "status": "NoResult", + "messages": [], + } + fallback = next( + (_qperiod(q, period) for q in originals if "MetricStat" in q), period + ) + original_out = [] + for q in originals: + matches = [ + r for (rid, _), r in orig_res.items() if rid == q["Id"] + ] or [{**empty, "label": q["Id"]}] + for r in matches: + original_out.append( + { + "id": q["Id"], + "label": r["label"], + "query": _js(q), + "status_code": r["status"], + "role": "expression" + if "Expression" in q + else "original_metric", + "messages": r["messages"], + **_timeline( + r["timestamps"], + r["values"], + _qperiod(q, fallback), + s_ep, + e_ep, + include_datapoints, + max_datapoints, + ), + } + ) + + # Global emitted-datapoint budget across the enumerated series (the fan-out + # that can overflow). 0 disables. Originals are O(1) and exempt. + remaining = max_total_datapoints if max_total_datapoints > 0 else None + output_capped = False + series_out = [] + for i, s in enumerate(specs): + r = series_res.get(f"zs{i}", empty) + tl = _timeline( + r["timestamps"], + r["values"], + period, + s_ep, + e_ep, + include_datapoints, + max_datapoints, + ) + if include_datapoints and "datapoints" in tl and remaining is not None: + n = len(tl["datapoints"]) + if n > remaining: + # Budget spent: drop this (and later) datapoint arrays, keep stats. + del tl["datapoints"] + tl["datapoints_truncated"] = True + output_capped = True + remaining = 0 + else: + remaining -= n + series_out.append( + { + "id": f"zs{i}", + **s, + "status_code": r["status"], + "messages": r["messages"], + **tl, + } + ) + + any_truncated = any( + x.get("datapoints_truncated") for x in original_out + series_out + ) + if any_truncated and not output_capped: + guidance.append( + f"Per-series datapoints truncated to max_datapoints={max_datapoints}; " + "raise --max-datapoints or narrow the window for the full series." + ) + if output_capped: + guidance.append( + f"Total emitted datapoints reached max_total_datapoints=" + f"{max_total_datapoints}; later series return summary stats only. " + "Raise --max-total-datapoints or narrow the sweep." + ) + + with_data = [s for s in series_out if s["datapoint_count"]] + not_complete = [ + s["id"] + for s in original_out + series_out + if s["status_code"] != "Complete" + ] + if not_complete: + warnings.append(f"Series without a Complete status: {not_complete}.") + summary = { + "series_count": len(series_out), + "series_with_data": len(with_data), + "series_without_data": [ + s["label"] for s in series_out if not s["datapoint_count"] + ], + "series_with_gaps": sum(bool(s["missing_count"]) for s in series_out), + "series_stopped_reporting": [ + { + "id": s["id"], + "label": s["label"], + "last_datapoint": s["last_datapoint"], + "trailing_missing_count": s["trailing_missing_count"], + } + for s in with_data + if s["trailing_missing_count"] + ], + "total_datapoints": sum(s["datapoint_count"] for s in series_out), + "top_by_peak": [ + { + "id": s["id"], + "label": s["label"], + "peak_value": s["max"], + "peak_timestamp": s["peak_timestamp"], + } + for s in sorted(with_data, key=lambda s: s["max"], reverse=True)[ + :10 + ] + ], + "state_transitions": history["count"], + "error_rate_target": error_rate, + "output_capped": output_capped, + } + value = { + "target": loaded["target"], + "region": region, + "window": {k: v for k, v in w.items() if k != "notes"}, + "state_history": history, + "original_queries": original_out, + "series": series_out, + "summary": summary, + "enumeration": enum["report"], + "warnings": warnings, + } + # If every GetMetricData batch failed, this is an operational error (the + # partial value is still attached so the caller can inspect what was gathered). + if attempted and len(errors) == attempted: + first = errors[0] + return err_result( + first.get("code") or SERVICE_ERROR, + f"GetMetricData failed for every batch: {first.get('error')}", + value=value, + aws_code=first.get("aws_code"), + retriable=bool(first.get("retriable")), + guidance=[*guidance, *warnings], + is_exhaustive=False, + key=loaded["target"].get("alarm_arn"), + ) + exhaustive = ( + enum["complete"] + and history["is_exhaustive"] + and complete + and not errors + and not not_complete + and not any_truncated + and not output_capped + ) + return ok_result( + value, + is_exhaustive=exhaustive, + guidance=guidance + warnings, + key=loaded["target"].get("alarm_arn"), + ) + + +# ------------------------------------------------------------- CLI plumbing +def register(event_handlers): + """Inject the ``alarm-metric-sweep`` verb into the cloudwatch command table.""" + event_handlers.register( + "building-command-table.cloudwatch", _inject_commands + ) + + +def _inject_commands(command_table, session, **kwargs): + command_table["alarm-metric-sweep"] = AlarmMetricSweepCommand(session) + + +class AlarmMetricSweepCommand(BasicCommand): + NAME = "alarm-metric-sweep" + DESCRIPTION = ( + "Investigate a CloudWatch alarm in one read-only call: resolve it (exact " + "name, else a bounded prefix lookup), read its state history, enumerate " + "every series sharing the alarm metric's dimension-name set (plus error-" + "rate siblings), and return per-series timelines with gaps, peaks and " + "trailing-gap detection.\n\n" + "Use this command when you need to understand what metric data is driving " + "an alarm's state. It resolves the alarm (including composite alarm " + "children), extracts the metric configuration, and pulls bounded metric " + "data in one call.\n\n" + "Output is capped (series, per-series datapoints, " + "and a global datapoint budget) so a wide sweep cannot overflow. Strictly " + "read-only: only Describe/Get/List calls are ever issued." + ) + EXAMPLES = ( + "To sweep an alarm over the last 3 hours::\n\n" + " aws cloudwatch alarm-metric-sweep --alarm-name my-5xx-alarm\n\n" + "To pin a dimension and bound output over an explicit window::\n\n" + " aws cloudwatch alarm-metric-sweep --alarm-name my-5xx-alarm " + "--start-time 2024-01-15T18:00:00Z --end-time 2024-01-15T21:00:00Z " + "--pin-dimensions LoadBalancer --max-series 50 --max-datapoints 500\n" + ) + ARG_TABLE = [ + { + "name": "alarm-name", + "required": True, + "cli_type_name": "string", + "help_text": "Exact alarm name (or a unique prefix) to investigate.", + }, + { + "name": "start-time", + "cli_type_name": "string", + "help_text": "Window start (ISO-8601). Default: end-time minus 3h.", + }, + { + "name": "end-time", + "cli_type_name": "string", + "help_text": "Window end (ISO-8601). Default: now.", + }, + { + "name": "period-s", + "default": 300, + "cli_type_name": "integer", + "help_text": "Requested period in seconds (raised for retention / point budget).", + }, + { + "name": "stat", + "cli_type_name": "string", + "help_text": "Override the statistic (Sum, Average, Maximum, p99, ...).", + }, + { + "name": "pin-dimensions", + "nargs": "+", + "cli_type_name": "string", + "help_text": "Dimension names to pin to the alarm's value when enumerating siblings.", + }, + { + "name": "no-sibling-metrics", + "action": "store_true", + "default": False, + "help_text": "Do not sweep error-rate sibling metrics (status-code/success).", + }, + { + "name": "no-datapoints", + "action": "store_true", + "default": False, + "help_text": "Return summary stats only; omit raw per-series datapoint arrays.", + }, + { + "name": "max-series", + "default": 200, + "cli_type_name": "integer", + "help_text": "Maximum number of series to return (default 200).", + }, + { + "name": "max-pages", + "default": 20, + "cli_type_name": "integer", + "help_text": "Maximum NextToken pages per paginated read (default 20).", + }, + { + "name": "max-datapoints", + "default": _MAX_POINTS, + "cli_type_name": "integer", + "help_text": "Maximum datapoints per series (default 1440; 0 disables).", + }, + { + "name": "max-total-datapoints", + "default": _DEFAULT_MAX_TOTAL_DATAPOINTS, + "cli_type_name": "integer", + "help_text": "Global emitted-datapoint budget across series (default 50000; 0 disables).", + }, + ] + + def __init__( + self, session, client=None, stdout=None, stderr=None, clock=None + ): + super().__init__(session) + self._client = client + self._stdout = stdout + self._stderr = stderr + self._clock = clock + + def _run_main(self, parsed_args, parsed_globals): + client = self._client or create_read_client( + self._session, "cloudwatch", parsed_globals + ) + + def call_op(operation, params): + return call(client, operation, params) + + result = alarm_metric_sweep( + call_op, + parsed_args.alarm_name, + start_time=parsed_args.start_time, + end_time=parsed_args.end_time, + period_s=parsed_args.period_s, + stat=parsed_args.stat, + pin_dimensions=parsed_args.pin_dimensions, + include_sibling_metrics=not parsed_args.no_sibling_metrics, + include_datapoints=not parsed_args.no_datapoints, + max_series=parsed_args.max_series, + max_pages=parsed_args.max_pages, + max_datapoints=parsed_args.max_datapoints, + max_total_datapoints=parsed_args.max_total_datapoints, + region=getattr(parsed_globals, "region", None), + clock=self._clock, + ) + return emit_result( + self.NAME, + result, + parsed_globals, + self._session, + self._stdout, + self._stderr, + ) diff --git a/awscli/handlers_registry.py b/awscli/handlers_registry.py index 5a1710888774..d2730b16dc82 100644 --- a/awscli/handlers_registry.py +++ b/awscli/handlers_registry.py @@ -560,7 +560,11 @@ class CommandTableOp(enum.Enum): ('awscli.customizations.cloudtrail', 'initialize') ], 'building-command-table.cloudwatch': [ - ('awscli.customizations.cloudwatch', 'register_rename_otel_commands') + ('awscli.customizations.cloudwatch', 'register_rename_otel_commands'), + ( + 'awscli.customizations.cloudwatch_alarm_metric_sweep', + 'register', + ), ], 'building-command-table.codeartifact': [ ( diff --git a/tests/functional/cloudwatch/test_alarm_metric_sweep.py b/tests/functional/cloudwatch/test_alarm_metric_sweep.py new file mode 100644 index 000000000000..f2475e60765e --- /dev/null +++ b/tests/functional/cloudwatch/test_alarm_metric_sweep.py @@ -0,0 +1,148 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"). You +# may not use this file except in compliance with the License. A copy of +# the License is located at +# +# http://aws.amazon.com/apache2.0/ +# +# or in the "license" file accompanying this file. This file is +# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF +# ANY KIND, either express or implied. See the License for the specific +# language governing permissions and limitations under the License. +import json + +from awscli.testutils import BaseAWSCommandParamsTest + +# Any of these appearing in operations_called would mean the "read-only" command +# issued a write; the suite asserts none ever do. +_MUTATING = { + 'PutMetricAlarm', + 'DeleteAlarms', + 'SetAlarmState', + 'PutDashboard', + 'PutMetricData', + 'TagResource', +} + + +class TestAlarmMetricSweepCommand(BaseAWSCommandParamsTest): + def _metric_alarm_responses(self): + return [ + # resolve_alarm exact hit + { + 'MetricAlarms': [ + { + 'AlarmName': 'cpu-high', + 'AlarmArn': ( + 'arn:aws:cloudwatch:us-east-1:1:alarm:cpu-high' + ), + 'StateValue': 'ALARM', + 'Namespace': 'AWS/EC2', + 'MetricName': 'CPUUtilization', + 'Dimensions': [{'Name': 'InstanceId', 'Value': 'i-1'}], + 'Statistic': 'Average', + 'Period': 300, + 'Threshold': 80.0, + 'ComparisonOperator': 'GreaterThanThreshold', + 'EvaluationPeriods': 1, + } + ], + 'CompositeAlarms': [], + }, + {'AlarmHistoryItems': []}, + { + 'Metrics': [ + { + 'Namespace': 'AWS/EC2', + 'MetricName': 'CPUUtilization', + 'Dimensions': [{'Name': 'InstanceId', 'Value': 'i-1'}], + } + ] + }, + { + 'MetricDataResults': [ + { + 'Id': 'alarm_metric', + 'Label': 'cpu-high', + 'Timestamps': ['2023-11-14T19:00:00Z'], + 'Values': [42.0], + 'StatusCode': 'Complete', + } + ] + }, + { + 'MetricDataResults': [ + { + 'Id': 'zs0', + 'Label': 'CPUUtilization', + 'Timestamps': ['2023-11-14T19:00:00Z'], + 'Values': [42.0], + 'StatusCode': 'Complete', + } + ] + }, + ] + + def test_command_is_registered_and_runs_read_only(self): + self.parsed_responses = self._metric_alarm_responses() + # Default window (last 3h from "now") keeps the request within + # CloudWatch's 455-day retention regardless of when the test runs. + stdout, _, rc = self.run_cmd( + [ + 'cloudwatch', + 'alarm-metric-sweep', + '--alarm-name', + 'cpu-high', + '--region', + 'us-east-1', + ], + expected_rc=0, + ) + names = [c[0].name for c in self.operations_called] + # Resolution begins with an exact DescribeAlarms and only read ops run. + self.assertEqual(names[0], 'DescribeAlarms') + self.assertTrue( + all(n.startswith(('Describe', 'Get', 'List')) for n in names) + ) + self.assertFalse(_MUTATING.intersection(names)) + parsed = json.loads(stdout) + self.assertEqual(parsed['target']['alarm_name'], 'cpu-high') + self.assertEqual(parsed['target']['match'], 'exact') + + def test_missing_alarm_is_operational_error(self): + # Exact miss then prefix miss -> NOT_FOUND -> rc 255. + self.parsed_responses = [ + {'MetricAlarms': [], 'CompositeAlarms': []}, + {'MetricAlarms': [], 'CompositeAlarms': []}, + ] + _, stderr, rc = self.run_cmd( + [ + 'cloudwatch', + 'alarm-metric-sweep', + '--alarm-name', + 'ghost', + '--region', + 'us-east-1', + ], + expected_rc=255, + ) + self.assertIn('NOT_FOUND', stderr) + + def test_invalid_parameter_rejected_before_any_call(self): + self.parsed_responses = [] + self.run_cmd( + [ + 'cloudwatch', + 'alarm-metric-sweep', + '--alarm-name', + 'cpu-high', + '--period-s', + '0', + '--region', + 'us-east-1', + ], + expected_rc=252, + ) + # Nothing was dispatched to AWS for bad input. + self.assertEqual(self.operations_called, []) diff --git a/tests/unit/customizations/cloudwatch/__init__.py b/tests/unit/customizations/cloudwatch/__init__.py new file mode 100644 index 000000000000..654905217728 --- /dev/null +++ b/tests/unit/customizations/cloudwatch/__init__.py @@ -0,0 +1,12 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"). You +# may not use this file except in compliance with the License. A copy of +# the License is located at +# +# http://aws.amazon.com/apache2.0/ +# +# or in the "license" file accompanying this file. This file is +# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF +# ANY KIND, either express or implied. See the License for the specific +# language governing permissions and limitations under the License. diff --git a/tests/unit/customizations/cloudwatch/test_alarm_metric_sweep.py b/tests/unit/customizations/cloudwatch/test_alarm_metric_sweep.py new file mode 100644 index 000000000000..cb17e8d7357c --- /dev/null +++ b/tests/unit/customizations/cloudwatch/test_alarm_metric_sweep.py @@ -0,0 +1,808 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"). You +# may not use this file except in compliance with the License. A copy of +# the License is located at +# +# http://aws.amazon.com/apache2.0/ +# +# or in the "license" file accompanying this file. This file is +# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF +# ANY KIND, either express or implied. See the License for the specific +# language governing permissions and limitations under the License. +"""Unit tests for ``aws cloudwatch alarm-metric-sweep`` (core sweep logic).""" + +from datetime import datetime, timezone + +from awscli.customizations.cloudwatch_alarm_metric_sweep import ( + alarm_metric_sweep, +) +from awscli.customizations.compositereadonly.helpers import ( + assert_read_only, + err_result, + ok_result, +) + +UTC = timezone.utc + + +class ScriptedCW: + """A scripted read-only call bridge: records calls, returns queued results. + + Mirrors the real :func:`call` bridge by refusing any mutating operation, so + a test fails loudly if the sweep ever attempts a write. + """ + + def __init__(self): + self.calls = [] + self._responses = {} + + def queue(self, operation, *results): + self._responses.setdefault(operation, []).extend(results) + return self + + def __call__(self, operation, params): + assert_read_only(operation) # read-only guarantee under test + self.calls.append((operation, params)) + queued = self._responses.get(operation) + if queued: + return queued.pop(0) + return ok_result({}) + + def ops(self): + return [op for op, _ in self.calls] + + +# ------------------------------------------------------------------- fixtures + + +def _metric_alarm(name="cpu-high", dims=None): + dims = ( + dims if dims is not None else [{"Name": "InstanceId", "Value": "i-1"}] + ) + return { + "AlarmName": name, + "AlarmArn": f"arn:aws:cloudwatch:us-east-1:1:alarm:{name}", + "StateValue": "ALARM", + "Namespace": "AWS/EC2", + "MetricName": "CPUUtilization", + "Dimensions": dims, + "Statistic": "Average", + "Period": 300, + "Threshold": 80.0, + "ComparisonOperator": "GreaterThanThreshold", + "EvaluationPeriods": 1, + } + + +def _describe(metric=None, composite=None, next_token=None): + body = { + "MetricAlarms": list(metric or []), + "CompositeAlarms": list(composite or []), + } + if next_token: + body["NextToken"] = next_token + return ok_result(body) + + +def _history(items=None, next_token=None): + body = {"AlarmHistoryItems": list(items or [])} + if next_token: + body["NextToken"] = next_token + return ok_result(body) + + +def _list_metrics(metrics=None, next_token=None): + body = {"Metrics": list(metrics or [])} + if next_token: + body["NextToken"] = next_token + return ok_result(body) + + +def _series(id_, values, start=1_700_000_000, period=300, status="Complete"): + # GetMetricData returns ISO-8601 / datetime timestamps (not raw epochs), so + # the fixture mirrors that; _dt() only parses datetimes and ISO strings. + ts = [ + datetime.fromtimestamp(start + i * period, UTC).strftime( + "%Y-%m-%dT%H:%M:%SZ" + ) + for i in range(len(values)) + ] + return { + "Id": id_, + "Label": id_, + "Timestamps": ts, + "Values": list(values), + "StatusCode": status, + } + + +def _metric_data(results=None, next_token=None): + body = {"MetricDataResults": list(results or [])} + if next_token: + body["NextToken"] = next_token + return ok_result(body) + + +def _fixed_clock(t=1_700_030_000.0): + return lambda: t + + +def _run(call, alarm_name="cpu-high", **kwargs): + kwargs.setdefault("clock", _fixed_clock()) + # Deterministic window so period/time-window math is stable in assertions. + kwargs.setdefault("start_time", "2023-11-14T18:00:00Z") + kwargs.setdefault("end_time", "2023-11-14T21:00:00Z") + return alarm_metric_sweep(call, alarm_name, **kwargs) + + +# ------------------------------------------------------- resolution: exact + + +def test_exact_match_resolves_and_issues_only_read_ops(): + call = ScriptedCW() + call.queue("DescribeAlarms", _describe(metric=[_metric_alarm()])) + call.queue("DescribeAlarmHistory", _history()) + call.queue( + "ListMetrics", + _list_metrics( + [ + { + "Namespace": "AWS/EC2", + "MetricName": "CPUUtilization", + "Dimensions": [{"Name": "InstanceId", "Value": "i-1"}], + } + ] + ), + ) + call.queue("GetMetricData", _metric_data([_series("alarm_metric", [1.0])])) + call.queue("GetMetricData", _metric_data([_series("zs0", [1.0])])) + result = _run(call) + assert result["ok"] is True + assert result["value"]["target"]["match"] == "exact" + # First resolution call is an exact DescribeAlarms by name (no prefix). + first_op, first_params = call.calls[0] + assert first_op == "DescribeAlarms" + assert first_params["AlarmNames"] == ["cpu-high"] + # Strictly read-only: no mutating verb ever reached the bridge. + assert all(op.startswith(("Describe", "Get", "List")) for op in call.ops()) + + +# ------------------------------------------------------- resolution: prefix + + +def test_unique_prefix_resolves_with_guidance(): + call = ScriptedCW() + # Exact miss, then a single prefix candidate. + call.queue("DescribeAlarms", _describe()) + call.queue( + "DescribeAlarms", _describe(metric=[_metric_alarm("cpu-high-1")]) + ) + call.queue("DescribeAlarmHistory", _history()) + call.queue("ListMetrics", _list_metrics()) + result = _run(call, alarm_name="cpu-high") + assert result["ok"] is True + assert result["value"]["target"]["match"] == "prefix" + assert any("prefix" in g for g in result["guidance"]) + + +def test_ambiguous_prefix_returns_not_found_with_candidates(): + call = ScriptedCW() + call.queue("DescribeAlarms", _describe()) # exact miss + call.queue( + "DescribeAlarms", + _describe( + metric=[ + _metric_alarm("db-writer-latency"), + _metric_alarm("db-reader-errors"), + ] + ), + ) + result = _run(call, alarm_name="db") + assert result["ok"] is False + assert result["code"] == "NOT_FOUND" + assert "ambiguous" in result["error"] + # No data calls after an unresolved name. + assert "GetMetricData" not in call.ops() + + +def test_missing_alarm_returns_not_found(): + call = ScriptedCW() + call.queue("DescribeAlarms", _describe()) # exact miss + call.queue("DescribeAlarms", _describe()) # prefix miss + result = _run(call, alarm_name="ghost") + assert result["ok"] is False + assert result["code"] == "NOT_FOUND" + assert "GetMetricData" not in call.ops() + + +# ------------------------------------------------------- param validation + + +def test_malformed_pin_dimensions_rejected_before_dispatch(): + call = ScriptedCW() + result = _run(call, pin_dimensions=[123]) + assert result["ok"] is False + assert result["code"] == "INVALID_INPUT" + assert result.get("param_error") is True + assert call.calls == [] # nothing dispatched + + +def test_bad_period_rejected_before_dispatch(): + call = ScriptedCW() + result = _run(call, period_s=0) + assert result["ok"] is False + assert result["code"] == "INVALID_INPUT" + assert call.calls == [] + + +def test_start_after_end_rejected(): + call = ScriptedCW() + result = alarm_metric_sweep( + call, + "cpu-high", + start_time="2023-11-14T21:00:00Z", + end_time="2023-11-14T18:00:00Z", + clock=_fixed_clock(), + ) + assert result["ok"] is False + assert result["code"] == "INVALID_INPUT" + assert call.calls == [] + + +# ------------------------------------------------------- pagination / bounds + + +def test_history_pagination_follows_next_token(): + call = ScriptedCW() + call.queue("DescribeAlarms", _describe(metric=[_metric_alarm()])) + item = { + "AlarmName": "cpu-high", + "Timestamp": "2023-11-14T19:00:00Z", + "HistoryItemType": "StateUpdate", + "HistoryData": "{}", + } + call.queue("DescribeAlarmHistory", _history([item], next_token="h2")) + call.queue("DescribeAlarmHistory", _history([item])) + call.queue("ListMetrics", _list_metrics()) + result = _run(call) + assert result["ok"] is True + assert result["value"]["state_history"]["count"] == 2 + assert call.ops().count("DescribeAlarmHistory") == 2 + + +def test_listmetrics_pagination_bounded_marks_partial(): + call = ScriptedCW() + call.queue("DescribeAlarms", _describe(metric=[_metric_alarm()])) + call.queue("DescribeAlarmHistory", _history()) + # ListMetrics never stops returning a token; the bounded loop must stop. + for _ in range(10): + call.queue( + "ListMetrics", + _list_metrics( + [ + { + "Namespace": "AWS/EC2", + "MetricName": "CPUUtilization", + "Dimensions": [{"Name": "InstanceId", "Value": "i-1"}], + } + ], + next_token="more", + ), + ) + call.queue("GetMetricData", _metric_data([_series("alarm_metric", [1.0])])) + call.queue("GetMetricData", _metric_data([_series("zs0", [1.0])])) + result = _run(call, max_pages=3) + assert result["ok"] is True + assert result["is_exhaustive"] is False + assert call.ops().count("ListMetrics") == 3 # capped at max_pages + + +# ------------------------------------------------------- empty / composite + + +def test_composite_alarm_has_no_metrics_lists_children(): + call = ScriptedCW() + composite = { + "AlarmName": "parent", + "AlarmArn": "arn:aws:cloudwatch:us-east-1:1:alarm:parent", + "StateValue": "ALARM", + "AlarmRule": 'ALARM("child-a") OR ALARM("child-b")', + } + call.queue("DescribeAlarms", _describe(composite=[composite])) + # History is still fetched for the composite and its children. + call.queue("DescribeAlarmHistory", _history()) + result = _run(call) + assert result["ok"] is True + assert result["value"]["target"]["kind"] == "composite_alarm" + assert result["value"]["target"]["child_alarms"] == ["child-a", "child-b"] + assert result["value"]["series"] == [] + assert "GetMetricData" not in call.ops() + + +def test_metric_alarm_with_no_datapoints_reports_series_without_data(): + call = ScriptedCW() + call.queue("DescribeAlarms", _describe(metric=[_metric_alarm()])) + call.queue("DescribeAlarmHistory", _history()) + call.queue( + "ListMetrics", + _list_metrics( + [ + { + "Namespace": "AWS/EC2", + "MetricName": "CPUUtilization", + "Dimensions": [{"Name": "InstanceId", "Value": "i-1"}], + } + ] + ), + ) + call.queue("GetMetricData", _metric_data([])) # originals: no data + call.queue("GetMetricData", _metric_data([])) # series: no data + result = _run(call) + assert result["ok"] is True + assert result["value"]["summary"]["series_with_data"] == 0 + assert result["value"]["summary"]["series_without_data"] + + +# ------------------------------------------------------- output capping + + +def test_per_series_datapoint_cap_truncates_and_reports(): + call = ScriptedCW() + call.queue("DescribeAlarms", _describe(metric=[_metric_alarm()])) + call.queue("DescribeAlarmHistory", _history()) + call.queue( + "ListMetrics", + _list_metrics( + [ + { + "Namespace": "AWS/EC2", + "MetricName": "CPUUtilization", + "Dimensions": [{"Name": "InstanceId", "Value": "i-1"}], + } + ] + ), + ) + many = [float(i) for i in range(12)] + call.queue("GetMetricData", _metric_data([_series("alarm_metric", many)])) + call.queue("GetMetricData", _metric_data([_series("zs0", many)])) + result = _run(call, period_s=300, max_datapoints=5) + series = result["value"]["series"][0] + assert len(series["datapoints"]) == 5 + assert series["datapoints_truncated"] is True + assert result["is_exhaustive"] is False + assert any( + "max-datapoints" in g or "max_datapoints" in g + for g in result["guidance"] + ) + + +def test_total_output_cap_drops_later_datapoint_arrays(): + call = ScriptedCW() + call.queue("DescribeAlarms", _describe(metric=[_metric_alarm()])) + call.queue("DescribeAlarmHistory", _history()) + # Two sibling series under the same dimension-name set. + call.queue( + "ListMetrics", + _list_metrics( + [ + { + "Namespace": "AWS/EC2", + "MetricName": "CPUUtilization", + "Dimensions": [{"Name": "InstanceId", "Value": "i-1"}], + }, + { + "Namespace": "AWS/EC2", + "MetricName": "CPUUtilization", + "Dimensions": [{"Name": "InstanceId", "Value": "i-2"}], + }, + ] + ), + ) + vals = [1.0, 2.0, 3.0, 4.0] + call.queue("GetMetricData", _metric_data([_series("alarm_metric", vals)])) + call.queue( + "GetMetricData", + _metric_data([_series("zs0", vals), _series("zs1", vals)]), + ) + # Budget only allows the first series' datapoints to be emitted. + result = _run(call, max_total_datapoints=5) + emitted = [("datapoints" in s) for s in result["value"]["series"]] + assert emitted[0] is True + assert emitted[-1] is False + assert result["value"]["summary"]["output_capped"] is True + + +def test_no_datapoints_flag_omits_arrays_but_keeps_summary(): + call = ScriptedCW() + call.queue("DescribeAlarms", _describe(metric=[_metric_alarm()])) + call.queue("DescribeAlarmHistory", _history()) + call.queue( + "ListMetrics", + _list_metrics( + [ + { + "Namespace": "AWS/EC2", + "MetricName": "CPUUtilization", + "Dimensions": [{"Name": "InstanceId", "Value": "i-1"}], + } + ] + ), + ) + call.queue( + "GetMetricData", _metric_data([_series("alarm_metric", [1.0, 2.0])]) + ) + call.queue("GetMetricData", _metric_data([_series("zs0", [1.0, 2.0])])) + result = _run(call, include_datapoints=False) + series = result["value"]["series"][0] + assert "datapoints" not in series + assert series["datapoint_count"] == 2 # summary stats preserved + + +# ------------------------------------------------------- partial failure + + +def test_history_failure_degrades_but_sweep_continues(): + call = ScriptedCW() + call.queue("DescribeAlarms", _describe(metric=[_metric_alarm()])) + call.queue( + "DescribeAlarmHistory", + err_result("PERMISSION_DENIED", "no cloudwatch:DescribeAlarmHistory"), + ) + call.queue( + "ListMetrics", + _list_metrics( + [ + { + "Namespace": "AWS/EC2", + "MetricName": "CPUUtilization", + "Dimensions": [{"Name": "InstanceId", "Value": "i-1"}], + } + ] + ), + ) + call.queue("GetMetricData", _metric_data([_series("alarm_metric", [1.0])])) + call.queue("GetMetricData", _metric_data([_series("zs0", [1.0])])) + result = _run(call) + assert result["ok"] is True # sweep still returns a result + assert result["is_exhaustive"] is False + assert any( + "DescribeAlarmHistory" in w for w in result["value"]["warnings"] + ) + + +def test_all_metric_data_batches_failing_is_operational_error(): + call = ScriptedCW() + call.queue("DescribeAlarms", _describe(metric=[_metric_alarm()])) + call.queue("DescribeAlarmHistory", _history()) + call.queue( + "ListMetrics", + _list_metrics( + [ + { + "Namespace": "AWS/EC2", + "MetricName": "CPUUtilization", + "Dimensions": [{"Name": "InstanceId", "Value": "i-1"}], + } + ] + ), + ) + call.queue( + "GetMetricData", + err_result("THROTTLED", "rate exceeded", retriable=True), + ) + call.queue( + "GetMetricData", + err_result("THROTTLED", "rate exceeded", retriable=True), + ) + result = _run(call) + assert result["ok"] is False + assert result["code"] == "THROTTLED" + assert result["retriable"] is True + # The partial value is still attached so an agent can inspect it. + assert result["value"]["target"]["alarm_name"] == "cpu-high" + + +# ------------------------------------------------------- error-rate siblings + + +def _lm_entry(metric_name, dims): + return { + "Namespace": "AWS/ApplicationELB", + "MetricName": metric_name, + "Dimensions": dims, + } + + +def test_error_rate_alarm_sweeps_sibling_metrics(): + call = ScriptedCW() + dims = [{"Name": "LoadBalancer", "Value": "app/x"}] + alarm = _metric_alarm("target-5xx-errors", dims=dims) + alarm["Namespace"] = "AWS/ApplicationELB" + alarm["MetricName"] = "HTTPCode_Target_5XX_Count" + call.queue("DescribeAlarms", _describe(metric=[alarm])) + call.queue("DescribeAlarmHistory", _history()) + # 1st ListMetrics: dimension enumeration for the target metric. + call.queue( + "ListMetrics", + _list_metrics([_lm_entry("HTTPCode_Target_5XX_Count", dims)]), + ) + # 2nd ListMetrics: namespace-wide sibling scan at the same dimensions. + call.queue( + "ListMetrics", + _list_metrics( + [ + _lm_entry("HTTPCode_Target_5XX_Count", dims), + _lm_entry("RequestCount", dims), + ] + ), + ) + call.queue("GetMetricData", _metric_data([_series("alarm_metric", [1.0])])) + call.queue( + "GetMetricData", + _metric_data([_series("zs0", [1.0]), _series("zs1", [2.0])]), + ) + result = _run(call, alarm_name="target-5xx-errors") + assert result["ok"] is True + assert result["value"]["summary"]["error_rate_target"] is True + assert result["value"]["enumeration"]["sibling_metric_names"] == [ + "RequestCount" + ] + roles = {s["role"] for s in result["value"]["series"]} + assert "metric_sibling" in roles + assert call.ops().count("ListMetrics") == 2 + + +def test_no_sibling_metrics_flag_skips_sibling_scan(): + call = ScriptedCW() + dims = [{"Name": "LoadBalancer", "Value": "app/x"}] + alarm = _metric_alarm("target-5xx-errors", dims=dims) + alarm["Namespace"] = "AWS/ApplicationELB" + alarm["MetricName"] = "HTTPCode_Target_5XX_Count" + call.queue("DescribeAlarms", _describe(metric=[alarm])) + call.queue("DescribeAlarmHistory", _history()) + call.queue( + "ListMetrics", + _list_metrics([_lm_entry("HTTPCode_Target_5XX_Count", dims)]), + ) + call.queue("GetMetricData", _metric_data([_series("alarm_metric", [1.0])])) + call.queue("GetMetricData", _metric_data([_series("zs0", [1.0])])) + result = _run( + call, alarm_name="target-5xx-errors", include_sibling_metrics=False + ) + assert result["ok"] is True + # Only the dimension-enumeration ListMetrics ran; no sibling scan. + assert call.ops().count("ListMetrics") == 1 + assert result["value"]["enumeration"]["sibling_metric_names"] == [] + + +# --------------------------------------------------- metric-math alarm + + +def test_metric_math_alarm_replays_expression_and_inputs(): + call = ScriptedCW() + math_alarm = { + "AlarmName": "math-alarm", + "AlarmArn": "arn:aws:cloudwatch:us-east-1:1:alarm:math-alarm", + "StateValue": "OK", + "Metrics": [ + { + "Id": "m1", + "MetricStat": { + "Metric": { + "Namespace": "AWS/EC2", + "MetricName": "CPUUtilization", + "Dimensions": [{"Name": "InstanceId", "Value": "i-1"}], + }, + "Period": 300, + "Stat": "Average", + }, + }, + {"Id": "e1", "Expression": "m1/100", "Label": "ratio"}, + ], + } + call.queue("DescribeAlarms", _describe(metric=[math_alarm])) + call.queue("DescribeAlarmHistory", _history()) + call.queue( + "ListMetrics", + _list_metrics( + [ + { + "Namespace": "AWS/EC2", + "MetricName": "CPUUtilization", + "Dimensions": [{"Name": "InstanceId", "Value": "i-1"}], + } + ] + ), + ) + call.queue( + "GetMetricData", + _metric_data([_series("m1", [50.0]), _series("e1", [0.5])]), + ) + call.queue("GetMetricData", _metric_data([_series("zs0", [50.0])])) + result = _run(call, alarm_name="math-alarm") + assert result["ok"] is True + roles = {q["role"] for q in result["value"]["original_queries"]} + assert roles == {"original_metric", "expression"} + + +# --------------------------------------------- period / window clamps + + +def test_old_window_raises_period_for_retention_and_point_budget(): + # ~40 days old: retention forces >=300s, and a 40d span forces a wider + # period to stay within the 1,440-points-per-series budget. + now = 1_700_030_000.0 + start = datetime.fromtimestamp(now - 40 * 86400, UTC).strftime( + "%Y-%m-%dT%H:%M:%SZ" + ) + call = ScriptedCW() + call.queue("DescribeAlarms", _describe(metric=[_metric_alarm()])) + call.queue("DescribeAlarmHistory", _history()) + call.queue("ListMetrics", _list_metrics()) + result = alarm_metric_sweep( + call, + "cpu-high", + start_time=start, + end_time=None, + period_s=60, + clock=lambda: now, + ) + assert result["ok"] is True + assert result["value"]["window"]["period_s"] > 60 + joined = " ".join(result["guidance"]) + assert "retained at" in joined + assert "points per series" in joined + + +# --------------------------------------------- timeline analytics + + +def _series_iso(id_, points, status="Complete"): + return { + "Id": id_, + "Label": id_, + "Timestamps": [p[0] for p in points], + "Values": [p[1] for p in points], + "StatusCode": status, + } + + +def test_timeline_reports_gaps_zero_peak_and_trailing(): + call = ScriptedCW() + call.queue("DescribeAlarms", _describe(metric=[_metric_alarm()])) + call.queue("DescribeAlarmHistory", _history()) + call.queue( + "ListMetrics", + _list_metrics( + [ + { + "Namespace": "AWS/EC2", + "MetricName": "CPUUtilization", + "Dimensions": [{"Name": "InstanceId", "Value": "i-1"}], + } + ] + ), + ) + # Window 18:00..19:00 @300s = 13 grid points; data only at the first three, + # one of which is a zero and one a peak, so the rest are trailing gaps. + pts = [ + ("2023-11-14T18:00:00Z", 5.0), + ("2023-11-14T18:05:00Z", 99.0), + ("2023-11-14T18:10:00Z", 0.0), + ] + call.queue( + "GetMetricData", _metric_data([_series_iso("alarm_metric", pts)]) + ) + call.queue("GetMetricData", _metric_data([_series_iso("zs0", pts)])) + result = alarm_metric_sweep( + call, + "cpu-high", + start_time="2023-11-14T18:00:00Z", + end_time="2023-11-14T19:00:00Z", + period_s=300, + clock=lambda: 1_700_030_000.0, + ) + series = result["value"]["series"][0] + assert series["datapoint_count"] == 3 + assert series["zero_count"] == 1 + assert series["max"] == 99.0 + assert series["peak_timestamp"] == "2023-11-14T18:05:00Z" + assert series["missing_count"] > 0 + assert series["trailing_missing_count"] > 0 + assert series["longest_nonzero_run"] == 2 + summary = result["value"]["summary"] + assert summary["series_with_gaps"] >= 1 + assert summary["series_stopped_reporting"] + + +# --------------------------------------------- enumeration edges + + +def test_other_dimension_sets_are_reported_not_swept(): + call = ScriptedCW() + call.queue("DescribeAlarms", _describe(metric=[_metric_alarm()])) + call.queue("DescribeAlarmHistory", _history()) + call.queue( + "ListMetrics", + _list_metrics( + [ + { + "Namespace": "AWS/EC2", + "MetricName": "CPUUtilization", + "Dimensions": [{"Name": "InstanceId", "Value": "i-1"}], + }, + # Different dimension-name set: counted, not swept. + { + "Namespace": "AWS/EC2", + "MetricName": "CPUUtilization", + "Dimensions": [ + {"Name": "AutoScalingGroupName", "Value": "g"} + ], + }, + ] + ), + ) + call.queue("GetMetricData", _metric_data([_series("alarm_metric", [1.0])])) + call.queue("GetMetricData", _metric_data([_series("zs0", [1.0])])) + result = _run(call) + report = result["value"]["enumeration"]["metrics"][0] + assert report["unswept_dimension_sets"] + assert any("other dimension sets" in g for g in result["guidance"]) + + +def test_max_series_cap_drops_and_marks_partial(): + call = ScriptedCW() + call.queue("DescribeAlarms", _describe(metric=[_metric_alarm()])) + call.queue("DescribeAlarmHistory", _history()) + call.queue( + "ListMetrics", + _list_metrics( + [ + { + "Namespace": "AWS/EC2", + "MetricName": "CPUUtilization", + "Dimensions": [{"Name": "InstanceId", "Value": f"i-{n}"}], + } + for n in range(3) + ] + ), + ) + call.queue("GetMetricData", _metric_data([_series("alarm_metric", [1.0])])) + call.queue("GetMetricData", _metric_data([_series("zs0", [1.0])])) + result = _run(call, max_series=1) + assert len(result["value"]["series"]) == 1 + assert result["value"]["enumeration"]["series_dropped"] == 2 + assert result["is_exhaustive"] is False + assert any("max-series" in g for g in result["guidance"]) + + +# --------------------------------------------- internal error boundary + + +def test_internal_exception_is_returned_as_service_error(): + # Contract: alarm_metric_sweep "never raises". If the injected call bridge + # blows up with something that is NOT an errors-as-data result (an + # unexpected botocore shape, a bug), the outer boundary must convert it into + # a SERVICE_ERROR result -- mirroring the Athena/Logs/ECS siblings -- rather + # than let the traceback escape to the caller. + def exploding(operation, params): + raise RuntimeError("kaboom") + + result = _run(exploding, alarm_name="cpu-high") + assert result["ok"] is False + assert result["code"] == "SERVICE_ERROR" + assert "internal error" in result["error"] + assert "kaboom" in result["error"] # original cause preserved for triage + + +def test_normal_structured_errors_are_not_swallowed_by_the_boundary(): + # The never-raise boundary must NOT turn a legitimate structured error (an + # errors-as-data result returned by the bridge) into a generic SERVICE_ERROR: + # a NOT_FOUND from resolution must still surface verbatim. + call = ScriptedCW() + call.queue("DescribeAlarms", _describe()) # exact miss + call.queue("DescribeAlarms", _describe()) # prefix miss + result = _run(call, alarm_name="ghost") + assert result["ok"] is False + assert result["code"] == "NOT_FOUND" diff --git a/tests/unit/customizations/cloudwatch/test_resolve_alarm.py b/tests/unit/customizations/cloudwatch/test_resolve_alarm.py new file mode 100644 index 000000000000..6052f029a8ee --- /dev/null +++ b/tests/unit/customizations/cloudwatch/test_resolve_alarm.py @@ -0,0 +1,212 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"). You +# may not use this file except in compliance with the License. A copy of +# the License is located at +# +# http://aws.amazon.com/apache2.0/ +# +# or in the "license" file accompanying this file. This file is +# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF +# ANY KIND, either express or implied. See the License for the specific +# language governing permissions and limitations under the License. +"""Unit tests for CloudWatch's alarm resolver. + +``resolve_alarm`` used to live in the shared compositereadonly ``resolve`` +module; it now lives with its owning command in +:mod:`awscli.customizations.cloudwatch_alarm_metric_sweep`, so its tests live +here too. +""" + +from awscli.customizations import cloudwatch_alarm_metric_sweep as cw +from awscli.customizations.compositereadonly.helpers import ( + NOT_FOUND, + PERMISSION_DENIED, + assert_read_only, + err_result, + ok_result, +) + + +class Scripted: + """A scripted read-only call bridge; records calls, returns queued dicts.""" + + def __init__(self): + self.calls = [] + self._responses = {} + + def queue(self, operation, *results): + self._responses.setdefault(operation, []).extend(results) + return self + + def __call__(self, operation, params): + assert_read_only(operation) + self.calls.append((operation, params)) + queued = self._responses.get(operation) + if queued: + return queued.pop(0) + return ok_result({}) + + def ops(self): + return [op for op, _ in self.calls] + + +# --------------------------------------------------------------- resolve_alarm + + +def _metric_alarm(name, rule=False): + a = {"AlarmName": name} + if rule: + a["AlarmRule"] = "ALARM(x)" + return a + + +def _describe(metric=None, composite=None, token=None): + body = { + "MetricAlarms": list(metric or []), + "CompositeAlarms": list(composite or []), + } + if token: + body["NextToken"] = token + return ok_result(body) + + +def test_alarm_prefix_paging_stops_at_page_cap(): + call = Scripted() + call.queue("DescribeAlarms", _describe()) # exact miss + # Every prefix page returns a candidate and a NextToken; _describe_prefix + # must stop after _MAX_ALARM_PAGES reads rather than chase the last token, + # whatever the eventual match decision is. + for i in range(cw._MAX_ALARM_PAGES + 2): + call.queue( + "DescribeAlarms", + _describe(metric=[_metric_alarm(f"svc-{i}")], token="more"), + ) + cw.resolve_alarm(call, "svc") + # 1 exact + exactly _MAX_ALARM_PAGES prefix reads (the cap held); the + # surplus queued pages were never requested. + assert call.ops().count("DescribeAlarms") == (1 + cw._MAX_ALARM_PAGES) + + +def test_alarm_exact_metric_hit_short_circuits(): + call = Scripted() + call.queue("DescribeAlarms", _describe(metric=[_metric_alarm("cpu-high")])) + r = cw.resolve_alarm(call, "cpu-high") + assert r["ok"] is True + assert r["value"]["match"] == "exact" + assert r["value"]["kind"] == "metric_alarm" + # Exact hit means no prefix fallback was attempted. + assert call.ops() == ["DescribeAlarms"] + + +def test_alarm_exact_composite_hit(): + call = Scripted() + call.queue( + "DescribeAlarms", + _describe(composite=[_metric_alarm("agg", rule=True)]), + ) + r = cw.resolve_alarm(call, "agg") + assert r["value"]["kind"] == "composite_alarm" + assert r["value"]["match"] == "exact" + + +def test_alarm_exact_describe_error_is_returned(): + call = Scripted() + call.queue("DescribeAlarms", err_result(PERMISSION_DENIED, "denied")) + r = cw.resolve_alarm(call, "cpu") + assert r["ok"] is False + assert r["code"] == PERMISSION_DENIED + + +def test_alarm_prefix_single_candidate(): + call = Scripted() + call.queue("DescribeAlarms", _describe()) # exact miss + call.queue("DescribeAlarms", _describe(metric=[_metric_alarm("cpu-high")])) + r = cw.resolve_alarm(call, "cpu") + assert r["ok"] is True + assert r["value"]["match"] == "prefix" + assert any("prefix match" in g for g in r["value"]["guidance"]) + + +def test_alarm_prefix_describe_error_propagates(): + call = Scripted() + call.queue("DescribeAlarms", _describe()) # exact miss + call.queue("DescribeAlarms", err_result(PERMISSION_DENIED, "denied")) + r = cw.resolve_alarm(call, "cpu") + assert r["ok"] is False + assert r["code"] == PERMISSION_DENIED + + +def test_alarm_no_candidates_is_not_found(): + call = Scripted() + call.queue("DescribeAlarms", _describe()) # exact miss + call.queue("DescribeAlarms", _describe()) # prefix miss + r = cw.resolve_alarm(call, "ghost") + assert r["ok"] is False + assert r["code"] == NOT_FOUND + assert r["key"] == "ghost" + + +def test_alarm_prefix_paging_is_followed(): + call = Scripted() + call.queue("DescribeAlarms", _describe()) # exact miss + call.queue( + "DescribeAlarms", + _describe(metric=[_metric_alarm("cpu-a")], token="more"), + ) + call.queue("DescribeAlarms", _describe(metric=[_metric_alarm("cpu-b")])) + r = cw.resolve_alarm(call, "cpu-a") + # Two prefix candidates fetched across pages; exact-ci picks 'cpu-a'. + assert r["ok"] is True + assert r["value"]["alarm"]["AlarmName"] == "cpu-a" + + +def test_alarm_case_insensitive_exact_among_several(): + call = Scripted() + call.queue("DescribeAlarms", _describe()) # exact miss (case differs) + call.queue( + "DescribeAlarms", + _describe( + metric=[_metric_alarm("CPU-High"), _metric_alarm("cpu-low")] + ), + ) + r = cw.resolve_alarm(call, "cpu-high") + assert r["ok"] is True + assert r["value"]["alarm"]["AlarmName"] == "CPU-High" + assert any("case-insensitive" in g for g in r["value"]["guidance"]) + + +def test_alarm_fuzzy_single_close_match(): + call = Scripted() + call.queue("DescribeAlarms", _describe()) # exact miss + call.queue( + "DescribeAlarms", + _describe( + metric=[ + _metric_alarm("cpu-utilization-high"), + _metric_alarm("totally-different-name"), + ] + ), + ) + # 'cpu-utilisation-high' (British spelling) fuzzily matches the first only. + r = cw.resolve_alarm(call, "cpu-utilisation-high") + assert r["ok"] is True + assert r["value"]["match"] == "fuzzy" + assert r["value"]["alarm"]["AlarmName"] == "cpu-utilization-high" + + +def test_alarm_ambiguous_prefix_is_not_found_with_candidates(): + call = Scripted() + call.queue("DescribeAlarms", _describe()) # exact miss + # Many similar names, none an unambiguous match. + names = [f"svc-alarm-{i}" for i in range(12)] + call.queue( + "DescribeAlarms", + _describe(metric=[_metric_alarm(n) for n in names]), + ) + r = cw.resolve_alarm(call, "svc") + assert r["ok"] is False + assert r["code"] == NOT_FOUND + assert "ambiguous" in r["error"] + # More than 10 candidates -> the guidance list is elided with "...". + assert any("..." in g for g in r["guidance"]) From db84a5b991f511cb13f8b50af49baf8d82544543 Mon Sep 17 00:00:00 2001 From: Claudiu Popa Date: Thu, 29 Oct 2026 11:23:42 +0100 Subject: [PATCH 4/4] Add ECS service diagnosis command Generated by AI tools, and reviewed by Claudiu Popa Co-authored-by: Baris Kurt --- awscli/customizations/ecs/diagnose_service.py | 1243 +++++++++++++++++ awscli/handlers_registry.py | 3 +- tests/functional/ecs/test_diagnose_service.py | 145 ++ .../ecs/test_diagnose_service.py | 970 +++++++++++++ 4 files changed, 2360 insertions(+), 1 deletion(-) create mode 100644 awscli/customizations/ecs/diagnose_service.py create mode 100644 tests/functional/ecs/test_diagnose_service.py create mode 100644 tests/unit/customizations/ecs/test_diagnose_service.py diff --git a/awscli/customizations/ecs/diagnose_service.py b/awscli/customizations/ecs/diagnose_service.py new file mode 100644 index 000000000000..8cddee50eb6b --- /dev/null +++ b/awscli/customizations/ecs/diagnose_service.py @@ -0,0 +1,1243 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"). You +# may not use this file except in compliance with the License. A copy of +# the License is located at +# +# http://aws.amazon.com/apache2.0/ +# +# or in the "license" file accompanying this file. This file is +# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF +# ANY KIND, either express or implied. See the License for the specific +# language governing permissions and limitations under the License. +"""``aws ecs diagnose-service`` -- explain why an ECS service is unhealthy. + +Callers otherwise hand-roll a long traversal to answer "why is my service not +running?": DescribeServices, DescribeTaskDefinition, ListTasks/DescribeTasks for +both running and stopped tasks, ELBv2 target health, the execution/task role +policies, and the container log stream -- each read paginated and bounded by +hand. This command runs that traversal once, read-only, and returns one +structured diagnosis: the blockers it found, non-fatal warnings, the checks that +passed, ordered remediations, and a per-section availability map so one flaky +optional read never sinks the whole answer. + +This command returns a result object from every call and never raises; it keeps +an "errors as data" control flow via the shared +:mod:`awscli.customizations.compositereadonly.helpers` bridge and maps the final result +to a CLI exit code. + +Scope note: this command's read surface is deliberately ECS + ELBv2 + +CloudWatch Logs + IAM (the APIs these commands declare). ECR image-existence, +EC2/VPC egress and Auto Scaling capacity checks need clients outside that +surface and are therefore not included here; the accuracy decisions that fit +this surface ARE kept -- +the desiredCount=0 "false green", the conservative IAM check (only an outright +missing execution role or a zero-policy role is flagged), target health that is +only a blocker when the service is meant to run tasks, and stopped-task reasons +as ground-truth corroboration. A log-evidence section is added (Logs scope). + +Behaviour contract: + +* **Gating vs optional reads** -- ``DescribeServices`` is the single gating read; + a transport/permission failure there is returned verbatim and a missing or + INACTIVE service is a structured NOT_FOUND. Every other section is optional: a + failure degrades that section (recorded in ``sections`` and as a warning) and + the diagnosis continues rather than discarding the evidence already gathered. +* **Bounded fan-out / pagination** -- ``ListTasks`` follows ``nextToken`` under a + ``--max-pages`` cap; stopped tasks, target groups, log groups and log lines are + each capped; ``DescribeTasks`` is issued in batches of 100. Every clamp is + reflected in ``is_exhaustive``. +* **Redaction** -- container environment values are never emitted (names only), + secret references are surfaced by name only, and every surfaced log line is run + through :func:`redact_sensitive` before truncation so a password or access key + in a log message is scrubbed. + +Strictly read-only: it only issues ``Describe*`` / ``Get*`` / ``List*`` / +``Filter*`` calls, and the shared :func:`call` bridge refuses any mutating +operation by construction. +""" + +from __future__ import annotations + +import re +import time +from typing import Any, Callable + +from awscli.customizations.commands import BasicCommand +from awscli.customizations.compositereadonly.helpers import ( + NOT_FOUND, + SERVICE_ERROR, + call, + check_non_empty_str, + check_positive_int, + create_read_client, + emit_result, + err_result, + invalid_input, + ok_result, + paginate, +) +from awscli.customizations.compositereadonly.output import ( + DEFAULT_MESSAGE_MAX_CHARS, + cap_series, + truncate_field, +) + +# A call bridge is (service, operation, params) -> result dict. Unlike the +# single-service sibling commands, diagnose-service spans four services, so the +# bridge is keyed by service and each client is created lazily. +CallFn = Callable[[str, str, dict[str, Any]], dict[str, Any]] + +# --- Bounds (defaults; all overridable on the CLI) --------------------------- +_MAX_DESCRIBE_TASKS = 100 # ECS DescribeTasks accepts at most 100 task ARNs. +# Stopped-task reasons that mean the task could not start/stay up. A match +# corroborates an image/network/IAM problem with ground-truth evidence. +_FATAL_TASK_REASONS = ( + "CannotPull", + "ResourceInitialization", + "AccessDenied", + "CannotStart", + "CannotCreate", + "OutOfMemory", + "EssentialContainerExited", +) +# A CloudWatch Logs filter pattern that matches common error markers. The fake +# bridge ignores it in tests; in production it narrows the fetch to error lines. +_ERROR_FILTER = ( + "?ERROR ?Error ?error ?Exception ?Fatal ?FATAL ?fatal ?Traceback" +) + +# --- Redaction patterns ------------------------------------------------------ +# AWS access-key-id style tokens (AKIA/ASIA/... + 16 upper-alnum chars). +_AWS_KEY_RE = re.compile( + r"\b(?:AKIA|ASIA|AGPA|AIDA|AROA|AIPA|ANPA|ANVA|ABIA|ACCA)[0-9A-Z]{16}\b" +) +# A key=value / key: value pair whose key looks sensitive -> scrub the value. +# The key and separator are preserved so the shape of the log stays legible. The +# optional identifier-prefix/suffix lets a sensitive word match even when it is +# embedded in a longer name (e.g. DB_PASSWORD, APP_SECRET_KEY). +_SECRET_KV_RE = re.compile( + r"(?i)([a-z0-9_]*" + r"(?:password|passwd|pwd|secret|token|apikey|api[_-]key|" + r"accesskey|access[_-]key|authorization|bearer)" + r"[a-z0-9_]*)" + r"(\s*[=:]\s*|\s+)" + r"([^\s,;]+)" +) +# The ``\s+`` branch of the separator group matches a secret-like key followed +# only by whitespace and then any token, so free text such as "password is +# required" redacts the following word. This over-redaction of legitimate text +# is intentional: a false positive is preferred over leaking a real secret. + + +def redact_sensitive(text: Any) -> Any: + """Scrub secret-looking substrings from free text. + + Redacts AWS access-key-id tokens and the value of any ``key=value`` / + ``key: value`` pair whose key names a credential (password, secret, token, + api-key, authorization, ...). Non-strings pass through untouched so callers + can apply it blindly. This is intentionally conservative -- it never tries to + guess that an arbitrary opaque string is a secret -- to avoid mangling + legitimate log content. + + Redaction is best-effort: unlabeled secrets (e.g. bare secret access keys or + session tokens without a key= prefix) may not be caught. + """ + if not isinstance(text, str): + return text + text = _AWS_KEY_RE.sub("[REDACTED]", text) + # Replace only the value group, keeping the key and separator verbatim. + text = _SECRET_KV_RE.sub( + lambda m: f"{m.group(1)}{m.group(2)}[REDACTED]", text + ) + return text + + +def _blocker( + kind: str, severity: str, evidence: str, fix: str +) -> dict[str, Any]: + """Build a blocker/warning dict with a stable shape.""" + return { + "type": kind, + "severity": severity, + "evidence": evidence, + "fix": fix, + } + + +def _chunks(items: list[Any], size: int) -> list[list[Any]]: + return [items[i : i + size] for i in range(0, len(items), size)] + + +def _task_id(arn: str) -> str: + return arn.rsplit("/", 1)[-1] if "/" in arn else arn + + +class _Diagnosis: + """Accumulator for findings, per-section availability and the degraded flag. + + ``degraded`` records that SOME evidence was lost -- an optional read failed or + output was capped/truncated -- which is exactly what flips ``is_exhaustive`` + to False. A section deliberately skipped (e.g. ``--no-logs``) is recorded as + unavailable but does NOT set ``degraded``, because nothing was lost. + """ + + def __init__(self) -> None: + self.blockers: list[dict[str, Any]] = [] + self.warnings: list[dict[str, Any]] = [] + self.verified_ok: list[str] = [] + self.sections: dict[str, dict[str, Any]] = {} + self.degraded = False + + def merge(self, findings: tuple[list, list, list]) -> None: + self.blockers.extend(findings[0]) + self.warnings.extend(findings[1]) + self.verified_ok.extend(findings[2]) + + def section_ok(self, name: str) -> None: + self.sections[name] = {"available": True, "error": None, "code": None} + + def section_skipped(self, name: str) -> None: + # Intentional skip: unavailable but not a loss of evidence. + self.sections[name] = { + "available": False, + "error": None, + "code": "SKIPPED", + } + + def section_failed(self, name: str, res: dict[str, Any]) -> None: + self.sections[name] = { + "available": False, + "error": res.get("error"), + "code": res.get("code"), + } + self.degraded = True + + +# --------------------------------------------------------------- findings + + +def _check_desired_count(svc: dict[str, Any]) -> tuple[list, list, list]: + """desiredCount=0 is the single most common blocker and a 'false green': + a deployment can report rolloutState=COMPLETED while running zero tasks + (including the false-positive healthy note when data is partial).""" + blockers: list[dict[str, Any]] = [] + warnings: list[dict[str, Any]] = [] + ok: list[str] = [] + desired = svc.get("desiredCount", 0) + running = svc.get("runningCount", 0) + if desired == 0: + rollout = "" + deployments = svc.get("deployments", []) + if deployments: + rollout = deployments[0].get("rolloutState", "") + false_green = ( + " (deployment reports COMPLETED, but that is vacuously true at " + "desiredCount=0)" + if rollout == "COMPLETED" + else "" + ) + blockers.append( + _blocker( + "DESIRED_COUNT_ZERO", + "critical", + f"desiredCount=0, runningCount={running}{false_green}", + "Set desiredCount > 0 (UpdateService). For CODE_DEPLOY " + "services, scale via the deployment, not UpdateService.", + ) + ) + else: + ok.append("desired_count") + if svc.get("deploymentController", {}).get("type") == "CODE_DEPLOY": + warnings.append( + _blocker( + "CODE_DEPLOY_SERVICE", + "info", + "Service uses the CODE_DEPLOY deployment controller.", + "Scale/roll out via CodeDeploy, not a plain UpdateService " + "desiredCount change.", + ) + ) + return blockers, warnings, ok + + +def _role_has_no_policies( + role_arn: str, call_op: CallFn +) -> tuple[bool | None, bool]: + """Return (is_empty, read_ok). + + ``is_empty`` is True when the role has neither attached nor inline policies, + None when it could not be determined. ``read_ok`` is False when either IAM + read failed, so the caller can mark the IAM section degraded. + """ + role_name = role_arn.rsplit("/", 1)[-1] + attached = call_op( + "iam", "ListAttachedRolePolicies", {"RoleName": role_name} + ) + inline = call_op("iam", "ListRolePolicies", {"RoleName": role_name}) + if not attached["ok"] or not inline["ok"]: + return None, False + has_attached = bool((attached.get("value") or {}).get("AttachedPolicies")) + has_inline = bool((inline.get("value") or {}).get("PolicyNames")) + return (not (has_attached or has_inline)), True + + +def _check_iam( + task_def: dict[str, Any], call_op: CallFn +) -> tuple[list, list, list, bool, bool]: + """Conservative IAM check (ported). Only the unambiguous cases are flagged: + a missing execution role (blocker) or a role with zero policies (warning). + Deep policy simulation is false-positive prone, so it is deliberately + avoided. Returns findings plus (attempted, failed) for section bookkeeping. + """ + blockers: list[dict[str, Any]] = [] + warnings: list[dict[str, Any]] = [] + ok: list[str] = [] + attempted = False + failed = False + exec_arn = task_def.get("executionRoleArn", "") + if not exec_arn: + blockers.append( + _blocker( + "EXECUTION_ROLE_MISSING", + "critical", + "Task definition has no executionRoleArn; ECS cannot " + "pull images or write logs.", + "Attach an execution role with ecr:GetAuthorizationToken, ECR " + "pull, and logs:CreateLogStream/PutLogEvents.", + ) + ) + else: + attempted = True + empty, read_ok = _role_has_no_policies(exec_arn, call_op) + if not read_ok: + failed = True + warnings.append( + _blocker( + "IAM_CHECK_DEGRADED", + "info", + f"Could not read policies for execution role " + f"{exec_arn.rsplit('/', 1)[-1]}.", + "Re-run with iam:ListAttachedRolePolicies and " + "iam:ListRolePolicies permission.", + ) + ) + elif empty is True: + warnings.append( + _blocker( + "EXECUTION_ROLE_MISSING_PERMS", + "warning", + f"Execution role {exec_arn.rsplit('/', 1)[-1]} has no " + "attached or inline policies.", + "Grant ecr:GetAuthorizationToken, ECR pull, and logs write " + "permissions.", + ) + ) + else: + ok.append("execution_role") + task_arn = task_def.get("taskRoleArn", "") + if task_arn: + attempted = True + empty, read_ok = _role_has_no_policies(task_arn, call_op) + if not read_ok: + failed = True + elif empty is True: + warnings.append( + _blocker( + "TASK_ROLE_MISSING_PERMS", + "warning", + f"Task role {task_arn.rsplit('/', 1)[-1]} has zero " + "policies; the app will have no AWS permissions.", + "Attach the permissions the container needs, or remove the " + "task role if unused.", + ) + ) + elif empty is False: + ok.append("task_role") + return blockers, warnings, ok, attempted, failed + + +# --------------------------------------------------------------- tasks + + +def _list_task_arns( + call_op: CallFn, + cluster: str, + service: str, + desired_status: str, + max_pages: int, +) -> tuple[list[str], bool, dict[str, Any] | None]: + """ListTasks for a service, following nextToken under ``max_pages``. + + Returns (task_arns, pages_truncated, error). ``pages_truncated`` is True when + the bounded loop stopped with a token still outstanding, so the caller can + mark the diagnosis non-exhaustive rather than spin forever. + """ + # The ECS call bridge is ``call_op(service, operation, params)``; the shared + # helper expects ``call_op(operation, params)``, so bind the "ecs" service. + # taskArns are strings (not dicts), so keep only the truthy ones, and ECS + # uses the lower-camel ``nextToken`` field. + arns, is_exhaustive, err = paginate( + lambda operation, params: call_op("ecs", operation, params), + "ListTasks", + { + "cluster": cluster, + "serviceName": service, + "desiredStatus": desired_status, + }, + "taskArns", + max_pages, + token_field="nextToken", + item_filter=lambda a: bool(a), + ) + # A clean finish or an error is not a truncation; only a page-cap stop with a + # token still pending is. (Callers check ``err`` first, so this value is only + # consulted on the no-error path.) + pages_truncated = err is None and not is_exhaustive + return arns, pages_truncated, err + + +def _describe_tasks( + call_op: CallFn, cluster: str, arns: list[str] +) -> tuple[list[dict[str, Any]], dict[str, Any] | None]: + """DescribeTasks in batches of 100 (the API maximum).""" + tasks: list[dict[str, Any]] = [] + for batch in _chunks(arns, _MAX_DESCRIBE_TASKS): + res = call_op( + "ecs", "DescribeTasks", {"cluster": cluster, "tasks": batch} + ) + if not res["ok"]: + return tasks, res + tasks.extend((res.get("value") or {}).get("tasks", [])) + return tasks, None + + +def _running_tasks( + call_op: CallFn, + cluster: str, + service: str, + diag: _Diagnosis, + *, + max_tasks: int, + max_pages: int, +) -> dict[str, Any]: + """Running-task status/health (optional section).""" + arns, pages_truncated, err = _list_task_arns( + call_op, cluster, service, "RUNNING", max_pages + ) + if err is not None: + diag.section_failed("running_tasks", err) + return {"count": 0, "tasks": [], "truncated": False} + capped = len(arns) > max_tasks + arns, _ = cap_series(arns, max_tasks) + truncated = capped or pages_truncated + tasks_out: list[dict[str, Any]] = [] + if arns: + described, derr = _describe_tasks(call_op, cluster, arns) + if derr is not None: + diag.section_failed("running_tasks", derr) + return {"count": 0, "tasks": [], "truncated": truncated} + # Cap defensively so the bound stays authoritative. + described = described[:max_tasks] + for t in described: + tasks_out.append( + { + "task_id": _task_id(t.get("taskArn", "")), + "status": t.get("lastStatus", "UNKNOWN"), + "health": t.get("healthStatus", "UNKNOWN"), + "started_at": str(t.get("startedAt", "")), + "containers": [ + { + "name": c.get("name"), + "last_status": c.get("lastStatus"), + "exit_code": c.get("exitCode"), + "reason": c.get("reason"), + } + for c in t.get("containers", []) + ], + } + ) + diag.section_ok("running_tasks") + if truncated: + diag.degraded = True + return { + "count": len(tasks_out), + "tasks": tasks_out, + "truncated": truncated, + } + + +def _stopped_tasks( + call_op: CallFn, + cluster: str, + service: str, + diag: _Diagnosis, + *, + max_stopped_tasks: int, + max_pages: int, +) -> dict[str, Any]: + """Stopped tasks + their stop reasons/exit codes (optional section). + + Also derives the STOPPED_TASK_ERRORS blocker from the same data (no extra + calls): a stopped reason naming an image-pull, init or permission failure is + ground-truth corroboration of a blocker. + """ + arns, pages_truncated, err = _list_task_arns( + call_op, cluster, service, "STOPPED", max_pages + ) + if err is not None: + diag.section_failed("stopped_tasks", err) + return {"count": 0, "tasks": [], "truncated": False} + capped = len(arns) > max_stopped_tasks + arns, _ = cap_series(arns, max_stopped_tasks) + truncated = capped or pages_truncated + tasks_out: list[dict[str, Any]] = [] + fatal_reasons: list[str] = [] + if arns: + described, derr = _describe_tasks(call_op, cluster, arns) + if derr is not None: + diag.section_failed("stopped_tasks", derr) + return {"count": 0, "tasks": [], "truncated": truncated} + # Cap the described tasks defensively (DescribeTasks returns only the + # ARNs we asked for, but the cap stays authoritative regardless). + described = described[:max_stopped_tasks] + for t in described: + reason = t.get("stoppedReason", "") or "" + if any(k in reason for k in _FATAL_TASK_REASONS): + fatal_reasons.append(reason) + tasks_out.append( + { + "task_id": _task_id(t.get("taskArn", "")), + "last_status": t.get("lastStatus", "UNKNOWN"), + "stopped_reason": t.get("stoppedReason"), + "stop_code": t.get("stopCode"), + "task_definition": t.get("taskDefinitionArn"), + "containers": [ + { + "name": c.get("name"), + "exit_code": c.get("exitCode"), + "reason": c.get("reason"), + "last_status": c.get("lastStatus"), + } + for c in t.get("containers", []) + ], + } + ) + diag.section_ok("stopped_tasks") + if truncated: + diag.degraded = True + if fatal_reasons: + diag.blockers.append( + _blocker( + "STOPPED_TASK_ERRORS", + "critical", + f"Recent stopped tasks failed to start: {fatal_reasons[:3]}", + "Resolve the stopped-task reason (image pull, init, or " + "permissions) shown in the evidence.", + ) + ) + return { + "count": len(tasks_out), + "tasks": tasks_out, + "truncated": truncated, + } + + +# --------------------------------------------------------------- target health + + +def _target_health( + svc: dict[str, Any], + call_op: CallFn, + diag: _Diagnosis, + *, + max_target_groups: int, +) -> dict[str, Any]: + """ELBv2 target-group health for the service's load balancers (optional). + + Also derives the TARGET_GROUP_UNHEALTHY blocker, but only when the service is + meant to run tasks (desiredCount>0) -- an empty target group at + desiredCount=0 is expected, not a blocker (ported accuracy guard). + """ + tg_arns = [ + lb.get("targetGroupArn") + for lb in svc.get("loadBalancers", []) + if lb.get("targetGroupArn") + ] + if not tg_arns: + # No load balancer wired: not an error, nothing to query. + diag.section_ok("target_health") + return {"wired": False, "target_groups": []} + capped = len(tg_arns) > max_target_groups + tg_arns, _ = cap_series(tg_arns, max_target_groups) + if capped: + diag.degraded = True + desired = svc.get("desiredCount", 0) + groups: list[dict[str, Any]] = [] + any_failed = False + for tg in tg_arns: + th = call_op("elbv2", "DescribeTargetHealth", {"TargetGroupArn": tg}) + if not th["ok"]: + any_failed = True + groups.append({"target_group_arn": tg, "error": th.get("error")}) + continue + descriptions = (th.get("value") or {}).get( + "TargetHealthDescriptions", [] + ) + healthy = 0 + unhealthy: list[dict[str, Any]] = [] + for d in descriptions: + health = d.get("TargetHealth", {}) or {} + state = health.get("State", "") + if state == "healthy": + healthy += 1 + else: + unhealthy.append( + { + "target": (d.get("Target", {}) or {}).get("Id"), + "state": state, + "reason": health.get("Reason"), + "description": health.get("Description"), + } + ) + groups.append( + { + "target_group_arn": tg, + "healthy": healthy, + "total": len(descriptions), + "unhealthy": unhealthy, + } + ) + # Only a blocker when the service should have running, healthy targets. + if desired > 0 and healthy == 0: + diag.blockers.append( + _blocker( + "TARGET_GROUP_UNHEALTHY", + "critical", + f"Target group {tg.rsplit('/', 1)[-1]} has no healthy " + f"targets ({len(descriptions)} registered).", + "Fix the task health so targets pass health checks; verify " + "the health-check path/port and SG rules.", + ) + ) + if any_failed: + # Record a degraded section but keep the groups that did resolve. + diag.sections["target_health"] = { + "available": False, + "error": "one or more DescribeTargetHealth calls failed", + "code": SERVICE_ERROR, + } + diag.degraded = True + else: + diag.section_ok("target_health") + return {"wired": True, "target_groups": groups} + + +# --------------------------------------------------------------- log evidence + + +def _awslogs_groups(containers: list[dict[str, Any]], cap: int) -> list[str]: + """Distinct awslogs log-group names from the container definitions (capped). + + Order-preserving de-dup so the output is stable; only the ``awslogs`` driver + is understood (other drivers write elsewhere). + """ + groups: list[str] = [] + for c in containers: + log_cfg = c.get("logConfiguration") or {} + if log_cfg.get("logDriver") != "awslogs": + continue + group = (log_cfg.get("options") or {}).get("awslogs-group") + if group and group not in groups: + groups.append(group) + return groups[:cap] + + +def _log_evidence( + containers: list[dict[str, Any]], + call_op: CallFn, + diag: _Diagnosis, + *, + include_logs: bool, + clock: Callable[[], float], + max_log_groups: int, + max_log_lines: int, + log_window_minutes: int, + max_message_length: int, +) -> dict[str, Any]: + """Recent error-ish log lines from the task's awslogs groups (optional). + + Every surfaced message is redacted (:func:`redact_sensitive`) and then + truncated to ``max_message_length`` so a secret cannot leak and one huge line + cannot dominate the output. The total is capped at ``max_log_lines``. + """ + if not include_logs: + diag.section_skipped("log_evidence") + return {"groups": [], "lines": [], "truncated": False} + groups = _awslogs_groups(containers, max_log_groups) + if not groups: + # Nothing to fetch (no awslogs driver); report as unavailable, no loss. + diag.section_skipped("log_evidence") + return {"groups": [], "lines": [], "truncated": False} + now_ms = int(clock() * 1000) + start_ms = now_ms - log_window_minutes * 60 * 1000 + lines: list[dict[str, Any]] = [] + truncated = False + any_failed = False + for group in groups: + if len(lines) >= max_log_lines: + truncated = True + break + res = call_op( + "logs", + "FilterLogEvents", + { + "logGroupName": group, + "startTime": start_ms, + "endTime": now_ms, + "filterPattern": _ERROR_FILTER, + "limit": max_log_lines, + }, + ) + if not res["ok"]: + any_failed = True + continue + for ev in (res.get("value") or {}).get("events", []): + if len(lines) >= max_log_lines: + truncated = True + break + # Redact THEN truncate: redaction must see the whole token, and the + # truncation keeps the (already-scrubbed) head where the error is. + message = redact_sensitive(ev.get("message", "")) + message, _ = truncate_field(message, max_message_length) + lines.append( + { + "log_group": group, + "log_stream": ev.get("logStreamName"), + "timestamp": ev.get("timestamp"), + "message": message, + } + ) + if any_failed: + diag.sections["log_evidence"] = { + "available": False, + "error": "one or more FilterLogEvents calls failed", + "code": SERVICE_ERROR, + } + diag.degraded = True + else: + diag.section_ok("log_evidence") + if truncated: + diag.degraded = True + return {"groups": groups, "lines": lines, "truncated": truncated} + + +# --------------------------------------------------------------- task definition + + +def _task_definition_section( + task_def: dict[str, Any], +) -> dict[str, Any]: + """Shape the task definition for output: images, roles and a redacted + container config (environment NAMES only -- values are never emitted -- and + secret references by name only).""" + containers = [] + for c in task_def.get("containerDefinitions", []): + log_cfg = c.get("logConfiguration") or {} + containers.append( + { + "name": c.get("name"), + "image": c.get("image"), + # Names only: an environment value may be a secret, so it is + # never surfaced. + "environment_names": [ + e.get("name") + for e in c.get("environment", []) + if e.get("name") + ], + "secret_names": [ + s.get("name") + for s in c.get("secrets", []) + if s.get("name") + ], + "log_configuration": { + "log_driver": log_cfg.get("logDriver"), + "awslogs_group": (log_cfg.get("options") or {}).get( + "awslogs-group" + ), + }, + } + ) + return { + "family": task_def.get("family"), + "revision": task_def.get("revision"), + "arn": task_def.get("taskDefinitionArn"), + "execution_role": task_def.get("executionRoleArn"), + "task_role": task_def.get("taskRoleArn"), + "containers": containers, + } + + +# --------------------------------------------------------------- the diagnosis + + +def _validate( + service: Any, + cluster: Any, + caps: dict[str, Any], +) -> dict[str, Any] | None: + """Pre-dispatch input validation (makes no AWS calls).""" + svc_err = check_non_empty_str("service", service, required=True) + if svc_err is not None: + return invalid_input( + f"{svc_err} (the ECS service name or ARN).", + ["Pass --service ."], + ) + cluster_err = check_non_empty_str("cluster", cluster, required=False) + if cluster_err is not None: + return invalid_input( + cluster_err, + ["Pass --cluster (default 'default')."], + ) + for name, val in caps.items(): + err = check_positive_int(name, val) + if err is not None: + return invalid_input( + err, + [f"Pass a positive {name}."], + ) + return None + + +def diagnose_ecs_service( + call_op: CallFn, + service: str, + *, + cluster: str = "default", + region: str | None = None, + include_logs: bool = True, + max_events: int = 10, + max_stopped_tasks: int = 20, + max_tasks: int = 100, + max_target_groups: int = 10, + max_log_lines: int = 20, + max_log_groups: int = 5, + log_window_minutes: int = 60, + max_message_length: int = DEFAULT_MESSAGE_MAX_CHARS, + max_pages: int = 10, + clock: Callable[[], float] | None = None, +) -> dict[str, Any]: + """Diagnose why an ECS service is unhealthy -- in ONE read-only call. + + ``call_op(service, operation, params)`` is the shared read-only call bridge + keyed by service (ecs/elbv2/logs/iam). Returns the plain-dict result + convention (see :mod:`...compositereadonly.helpers`); never raises. ``DescribeServices`` + is the only gating read: a failure there is returned verbatim and a missing or + INACTIVE service is a structured NOT_FOUND. Every other section is optional + and degrades independently. ``is_exhaustive`` is False whenever any optional + read failed or output was capped/truncated. + """ + clock = time.time if clock is None else clock + try: + bad = _validate( + service, + cluster, + { + "max_events": max_events, + "max_stopped_tasks": max_stopped_tasks, + "max_tasks": max_tasks, + "max_target_groups": max_target_groups, + "max_log_lines": max_log_lines, + "max_log_groups": max_log_groups, + "log_window_minutes": log_window_minutes, + "max_message_length": max_message_length, + "max_pages": max_pages, + }, + ) + if bad is not None: + return bad + return _run( + call_op, + service.strip(), + cluster=cluster.strip(), + region=region, + include_logs=include_logs, + max_events=max_events, + max_stopped_tasks=max_stopped_tasks, + max_tasks=max_tasks, + max_target_groups=max_target_groups, + max_log_lines=max_log_lines, + max_log_groups=max_log_groups, + log_window_minutes=log_window_minutes, + max_message_length=max_message_length, + max_pages=max_pages, + clock=clock, + ) + except Exception as exc: # noqa: BLE001 - contract: never raise + return err_result( + SERVICE_ERROR, + f"diagnose-service internal error: {exc!r}", + guidance=[ + "Unexpected response shape; retry once, then fall back to " + "ecs:DescribeServices by hand." + ], + ) + + +def _run( + call_op: CallFn, + service: str, + *, + cluster: str, + region: str | None, + include_logs: bool, + max_events: int, + max_stopped_tasks: int, + max_tasks: int, + max_target_groups: int, + max_log_lines: int, + max_log_groups: int, + log_window_minutes: int, + max_message_length: int, + max_pages: int, + clock: Callable[[], float], +) -> dict[str, Any]: + # --- gating read: DescribeServices ------------------------------------- + svc_res = call_op( + "ecs", + "DescribeServices", + {"cluster": cluster, "services": [service]}, + ) + if not svc_res["ok"]: + return svc_res # propagate transport/permission errors verbatim + services = (svc_res.get("value") or {}).get("services", []) + if not services or services[0].get("status") == "INACTIVE": + return err_result( + NOT_FOUND, + f"Service '{service}' not found in cluster '{cluster}'", + key=service, + guidance=[ + "Check the cluster (--cluster) and region; list services with " + "ecs:ListServices." + ], + ) + svc = services[0] + diag = _Diagnosis() + + # 1. Cheapest + most common blocker first (pure logic, no API call). + diag.merge(_check_desired_count(svc)) + + # 2. Task definition drives the IAM check, the config summary and the log + # groups. Its failure is a warning, not fatal. + task_def: dict[str, Any] = {} + container_defs: list[dict[str, Any]] = [] + td_ref = svc.get("taskDefinition", "") + task_definition_out: dict[str, Any] = {} + if td_ref: + td = call_op( + "ecs", "DescribeTaskDefinition", {"taskDefinition": td_ref} + ) + if td["ok"]: + task_def = (td.get("value") or {}).get("taskDefinition", {}) + container_defs = task_def.get("containerDefinitions", []) + task_definition_out = _task_definition_section(task_def) + diag.section_ok("task_definition") + else: + diag.section_failed("task_definition", td) + diag.warnings.append( + _blocker( + "TASK_DEF_UNREADABLE", + "info", + f"Could not read task definition ({td.get('code')}: " + f"{td.get('error')}).", + "Re-run with ecs:DescribeTaskDefinition permission.", + ) + ) + + # 3. IAM check (only if the task definition was readable). + if task_def: + blockers, warnings, ok, attempted, failed = _check_iam( + task_def, call_op + ) + diag.merge((blockers, warnings, ok)) + if attempted and failed: + diag.sections["iam"] = { + "available": False, + "error": "an IAM policy read failed", + "code": SERVICE_ERROR, + } + diag.degraded = True + elif attempted: + diag.section_ok("iam") + + # 4. Stopped tasks (ground-truth failure reasons + blocker) then running + # tasks. Each is optional and degrades independently. + stopped = _stopped_tasks( + call_op, + cluster, + service, + diag, + max_stopped_tasks=max_stopped_tasks, + max_pages=max_pages, + ) + running = _running_tasks( + call_op, + cluster, + service, + diag, + max_tasks=max_tasks, + max_pages=max_pages, + ) + + # 5. Target health (+ blocker when desiredCount>0 and no healthy targets). + target_health = _target_health( + svc, call_op, diag, max_target_groups=max_target_groups + ) + + # 6. Log evidence from the task's awslogs groups (redacted + capped). + log_evidence = _log_evidence( + container_defs, + call_op, + diag, + include_logs=include_logs, + clock=clock, + max_log_groups=max_log_groups, + max_log_lines=max_log_lines, + log_window_minutes=log_window_minutes, + max_message_length=max_message_length, + ) + + # --- assemble the diagnosis -------------------------------------------- + deployment_config = svc.get("deploymentConfiguration", {}) or {} + cb = deployment_config.get("deploymentCircuitBreaker", {}) or {} + deployment = { + "service_status": svc.get("status", "UNKNOWN"), + "circuit_breaker": { + "enabled": bool(cb.get("enable", False)), + "rollback": bool(cb.get("rollback", False)), + }, + "deployments": [ + { + "id": d.get("id"), + "status": d.get("status"), + "task_definition": d.get("taskDefinition"), + "desired_count": d.get("desiredCount"), + "running_count": d.get("runningCount"), + "pending_count": d.get("pendingCount"), + "failed_tasks": d.get("failedTasks"), + "rollout_state": d.get("rolloutState"), + "rollout_state_reason": d.get("rolloutStateReason"), + } + for d in svc.get("deployments", []) + ], + } + recent_events, events_dropped = cap_series( + [ + { + "created_at": str(e.get("createdAt", "")), + "message": e.get("message", ""), + } + for e in svc.get("events", []) + ], + max_events, + ) + + recommendations = [f"[{b['type']}] {b['fix']}" for b in diag.blockers] + value = { + "service": service, + "cluster": cluster, + "region": region, + "healthy": not diag.blockers, + "service_summary": { + "desired_count": svc.get("desiredCount", 0), + "running_count": svc.get("runningCount", 0), + "pending_count": svc.get("pendingCount", 0), + "launch_type": svc.get("launchType", ""), + "status": svc.get("status", ""), + "task_definition": svc.get("taskDefinition"), + }, + "blockers": diag.blockers, + "warnings": diag.warnings, + "verified_ok": diag.verified_ok, + "recommendations": recommendations, + "deployment": deployment, + "recent_events": recent_events, + "stopped_tasks": stopped, + "running_tasks": running, + "target_health": target_health, + "task_definition": task_definition_out, + "log_evidence": log_evidence, + "sections": diag.sections, + } + return ok_result( + value, + key=service, + is_exhaustive=not diag.degraded, + guidance=recommendations, + ) + + +# ------------------------------------------------------------- CLI plumbing + + +def register(event_handlers): + """Inject the ``diagnose-service`` verb into the ecs command table.""" + event_handlers.register("building-command-table.ecs", _inject_commands) + + +def _inject_commands(command_table, session, **kwargs): + command_table["diagnose-service"] = DiagnoseServiceCommand(session) + + +class DiagnoseServiceCommand(BasicCommand): + NAME = "diagnose-service" + DESCRIPTION = ( + "Explain why an ECS service is unhealthy in one read-only call.\n\n" + "Use this command when you need a full diagnostic view of an ECS " + "service. It gathers the service configuration, task definitions, " + "running and stopped tasks, recent events, " + "log evidence, and load balancer health in one call with automatic " + "secret redaction.\n\n" + "Reads " + "the service (the gating read), its task definition, running and " + "stopped tasks (with stop reasons and container exit codes), ELBv2 " + "target health, the execution/task role policies and recent error log " + "lines, and returns one structured diagnosis: blockers, warnings, the " + "checks that passed, ordered remediations and a per-section " + "availability map. Environment values and secret-looking log tokens are " + "redacted; output is bounded. A missing/INACTIVE service is NOT_FOUND; " + "any optional read failing degrades only its section. Strictly " + "read-only: only Describe/Get/List/Filter calls are ever issued." + ) + EXAMPLES = ( + "To diagnose a service in the default cluster::\n\n" + " aws ecs diagnose-service --service my-svc\n\n" + "To diagnose a service in a named cluster without fetching logs::\n\n" + " aws ecs diagnose-service --service my-svc --cluster prod " + "--no-logs\n" + ) + ARG_TABLE = [ + { + "name": "service", + "required": True, + "cli_type_name": "string", + "help_text": "ECS service name (or ARN) to diagnose.", + }, + { + "name": "cluster", + "default": "default", + "cli_type_name": "string", + "help_text": "ECS cluster name (default 'default').", + }, + { + "name": "max-events", + "default": 10, + "cli_type_name": "integer", + "help_text": "Maximum recent service events to return (default 10).", + }, + { + "name": "max-stopped-tasks", + "default": 20, + "cli_type_name": "integer", + "help_text": "Maximum stopped tasks to inspect (default 20).", + }, + { + "name": "max-tasks", + "default": 100, + "cli_type_name": "integer", + "help_text": "Maximum running tasks to inspect (default 100).", + }, + { + "name": "max-target-groups", + "default": 10, + "cli_type_name": "integer", + "help_text": "Maximum target groups to check (default 10).", + }, + { + "name": "max-log-lines", + "default": 20, + "cli_type_name": "integer", + "help_text": "Maximum log lines to return (default 20).", + }, + { + "name": "max-log-groups", + "default": 5, + "cli_type_name": "integer", + "help_text": "Maximum awslogs groups to read (default 5).", + }, + { + "name": "log-window-minutes", + "default": 60, + "cli_type_name": "integer", + "help_text": "Look-back window for log evidence in minutes (default 60).", + }, + { + "name": "max-message-length", + "default": DEFAULT_MESSAGE_MAX_CHARS, + "cli_type_name": "integer", + "help_text": ( + "Truncate each log message to this many chars " + f"(default {DEFAULT_MESSAGE_MAX_CHARS})." + ), + }, + { + "name": "max-pages", + "default": 10, + "cli_type_name": "integer", + "help_text": "Maximum nextToken pages per paginated read (default 10).", + }, + { + "name": "no-logs", + "action": "store_true", + "default": False, + "help_text": "Skip the CloudWatch Logs evidence fetch.", + }, + ] + + def __init__( + self, session, clients=None, stdout=None, stderr=None, clock=None + ): + super().__init__(session) + self._clients = clients or {} + self._stdout = stdout + self._stderr = stderr + self._clock = clock + + def _run_main(self, parsed_args, parsed_globals): + # Lazily create one read-only client per service, honouring + # --region/--endpoint-url/--verify-ssl via the shared helper. + cache = dict(self._clients) + + def call_op(service, operation, params): + client = cache.get(service) + if client is None: + client = cache[service] = create_read_client( + self._session, service, parsed_globals + ) + return call(client, operation, params) + + result = diagnose_ecs_service( + call_op, + parsed_args.service, + cluster=parsed_args.cluster, + region=getattr(parsed_globals, "region", None), + include_logs=not parsed_args.no_logs, + max_events=parsed_args.max_events, + max_stopped_tasks=parsed_args.max_stopped_tasks, + max_tasks=parsed_args.max_tasks, + max_target_groups=parsed_args.max_target_groups, + max_log_lines=parsed_args.max_log_lines, + max_log_groups=parsed_args.max_log_groups, + log_window_minutes=parsed_args.log_window_minutes, + max_message_length=parsed_args.max_message_length, + max_pages=parsed_args.max_pages, + clock=self._clock, + ) + return emit_result( + self.NAME, + result, + parsed_globals, + self._session, + self._stdout, + self._stderr, + ) diff --git a/awscli/handlers_registry.py b/awscli/handlers_registry.py index d2730b16dc82..ef21f7331bf2 100644 --- a/awscli/handlers_registry.py +++ b/awscli/handlers_registry.py @@ -625,7 +625,8 @@ class CommandTableOp(enum.Enum): ('awscli.customizations.ecr_public', 'register_ecr_public_commands') ], 'building-command-table.ecs': [ - ('awscli.customizations.ecs', 'initialize') + ('awscli.customizations.ecs', 'initialize'), + ('awscli.customizations.ecs.diagnose_service', 'register'), ], 'building-command-table.eks': [ ('awscli.customizations.eks', 'initialize') diff --git a/tests/functional/ecs/test_diagnose_service.py b/tests/functional/ecs/test_diagnose_service.py new file mode 100644 index 000000000000..e01be04bc35c --- /dev/null +++ b/tests/functional/ecs/test_diagnose_service.py @@ -0,0 +1,145 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"). You +# may not use this file except in compliance with the License. A copy of +# the License is located at +# +# http://aws.amazon.com/apache2.0/ +# +# or in the "license" file accompanying this file. This file is +# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF +# ANY KIND, either express or implied. See the License for the specific +# language governing permissions and limitations under the License. +"""Functional tests for ``aws ecs diagnose-service``. + +These drive the command through the real CLI command table (so they also assert +the verb is registered alongside the existing ``ecs deploy`` et al.) with +scripted botocore responses, and verify that only read operations are ever +issued and that exit codes follow the shared convention. +""" + +import json + +from awscli.testutils import BaseAWSCommandParamsTest + +# Operations that would create, modify or delete state across the four services +# this command touches. It must never issue any of these. +_MUTATING = { + 'UpdateService', + 'CreateService', + 'DeleteService', + 'RegisterTaskDefinition', + 'RunTask', + 'StopTask', + 'AttachRolePolicy', + 'PutRetentionPolicy', + 'RegisterTargets', +} + +_TASK_DEF_ARN = 'arn:aws:ecs:us-east-1:111122223333:task-definition/web:5' + + +class TestDiagnoseServiceCommand(BaseAWSCommandParamsTest): + def _healthy_responses(self): + return [ + # 1. DescribeServices (gating) -> one ACTIVE service, no LB. + { + 'services': [ + { + 'serviceName': 'web', + 'status': 'ACTIVE', + 'desiredCount': 1, + 'runningCount': 1, + 'pendingCount': 0, + 'launchType': 'FARGATE', + 'taskDefinition': _TASK_DEF_ARN, + 'deployments': [ + { + 'id': 'ecs-svc/1', + 'status': 'PRIMARY', + 'rolloutState': 'COMPLETED', + } + ], + 'events': [{'message': 'steady state'}], + 'loadBalancers': [], + } + ], + 'failures': [], + }, + # 2. DescribeTaskDefinition (execution role present). + { + 'taskDefinition': { + 'family': 'web', + 'revision': 5, + 'taskDefinitionArn': _TASK_DEF_ARN, + 'executionRoleArn': ( + 'arn:aws:iam::111122223333:role/ecsTaskExecutionRole' + ), + 'containerDefinitions': [ + {'name': 'web', 'image': 'web:1'} + ], + } + }, + # 3. IAM: execution role has an attached policy. + {'AttachedPolicies': [{'PolicyName': 'exec'}]}, + {'PolicyNames': []}, + # 4. ListTasks STOPPED (none). + {'taskArns': []}, + # 5. ListTasks RUNNING (none). + {'taskArns': []}, + ] + + def test_runs_read_only_and_reports_healthy(self): + self.parsed_responses = self._healthy_responses() + stdout, _, rc = self.run_cmd( + [ + 'ecs', + 'diagnose-service', + '--service', + 'web', + '--cluster', + 'prod', + '--no-logs', + '--region', + 'us-east-1', + ], + expected_rc=0, + ) + names = [c[0].name for c in self.operations_called] + # The gating read runs first; no mutating operation is ever issued. + self.assertEqual(names[0], 'DescribeServices') + self.assertFalse(_MUTATING.intersection(names)) + parsed = json.loads(stdout) + self.assertTrue(parsed['healthy']) + self.assertEqual(parsed['blockers'], []) + self.assertEqual(parsed['service_summary']['running_count'], 1) + + def test_missing_service_is_operational_error(self): + # DescribeServices returns no services -> NOT_FOUND -> rc 255. + self.parsed_responses = [ + {'services': [], 'failures': [{'reason': 'MISSING'}]} + ] + _, stderr, rc = self.run_cmd( + [ + 'ecs', + 'diagnose-service', + '--service', + 'ghost', + '--region', + 'us-east-1', + ], + expected_rc=255, + ) + self.assertIn('NOT_FOUND', stderr) + names = [c[0].name for c in self.operations_called] + # Nothing beyond the gating read was attempted. + self.assertEqual(names, ['DescribeServices']) + + def test_missing_required_service_arg_is_param_error(self): + self.parsed_responses = [] + # argparse rejects the missing required --service before dispatch. + self.run_cmd( + ['ecs', 'diagnose-service', '--region', 'us-east-1'], + expected_rc=252, + ) + self.assertEqual(self.operations_called, []) diff --git a/tests/unit/customizations/ecs/test_diagnose_service.py b/tests/unit/customizations/ecs/test_diagnose_service.py new file mode 100644 index 000000000000..6afce406ad56 --- /dev/null +++ b/tests/unit/customizations/ecs/test_diagnose_service.py @@ -0,0 +1,970 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"). You +# may not use this file except in compliance with the License. A copy of +# the License is located at +# +# http://aws.amazon.com/apache2.0/ +# +# or in the "license" file accompanying this file. This file is +# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF +# ANY KIND, either express or implied. See the License for the specific +# language governing permissions and limitations under the License. +"""Unit tests for ``aws ecs diagnose-service`` (diagnose_ecs_service logic). + +These exercise the pure, read-only diagnosis function with a scripted call +bridge so no network, clients or wall clock are involved. The scripted bridge +mirrors the real one: it refuses any mutating operation, so a test fails if the +command ever reaches for a write API. +""" + +from awscli.customizations.ecs.diagnose_service import ( + diagnose_ecs_service, + redact_sensitive, +) +from awscli.customizations.compositereadonly.helpers import ( + assert_read_only, + err_result, + ok_result, +) + +# Operations that would create, modify or delete state. No test ever queues one; +# the scripted bridge's assert_read_only would raise, and _no_writes() asserts +# the recorded calls never include one. +_MUTATING = { + ("ecs", "UpdateService"), + ("ecs", "CreateService"), + ("ecs", "RegisterTaskDefinition"), + ("ecs", "RunTask"), + ("ecs", "DeleteService"), + ("iam", "AttachRolePolicy"), + ("elbv2", "RegisterTargets"), + ("logs", "PutLogEvents"), +} + +_TASK_DEF_ARN = "arn:aws:ecs:us-east-1:111122223333:task-definition/web:5" + + +class ScriptedECS: + """A scripted (service, operation, params) read-only call bridge. + + Records every call and returns queued results keyed by (service, operation); + an unqueued read returns an empty ok_result so optional sections degrade to + "empty" rather than erroring. + """ + + def __init__(self): + self.calls = [] + self._responses = {} + + def queue(self, service, operation, *results): + self._responses.setdefault((service, operation), []).extend(results) + return self + + def __call__(self, service, operation, params): + # Mirror the real bridge: every operation must be read-only. + assert_read_only(operation) + self.calls.append((service, operation, params)) + queued = self._responses.get((service, operation)) + if queued: + return queued.pop(0) + return ok_result({}) + + def ops(self): + return [(s, o) for s, o, _ in self.calls] + + +def _fake_clock(): + # Fixed epoch so the log-evidence time window is deterministic. + return lambda: 1_700_000_000.0 + + +def _svc( + *, + status="ACTIVE", + desired=2, + running=2, + pending=0, + launch_type="FARGATE", + task_def=_TASK_DEF_ARN, + deployments=None, + events=None, + load_balancers=None, + circuit=None, +): + if deployments is None: + deployments = [ + { + "id": "ecs-svc/1", + "status": "PRIMARY", + "taskDefinition": task_def, + "desiredCount": desired, + "runningCount": running, + "pendingCount": pending, + "failedTasks": 0, + "rolloutState": "COMPLETED", + "rolloutStateReason": "ECS deployment completed.", + } + ] + if events is None: + events = [ + { + "createdAt": "2026-01-01T00:00:00Z", + "message": "has reached a steady state.", + } + ] + svc = { + "serviceName": "web", + "serviceArn": "arn:aws:ecs:us-east-1:111122223333:service/prod/web", + "status": status, + "desiredCount": desired, + "runningCount": running, + "pendingCount": pending, + "launchType": launch_type, + "taskDefinition": task_def, + "deployments": deployments, + "events": events, + "loadBalancers": load_balancers or [], + } + if circuit is not None: + svc["deploymentConfiguration"] = {"deploymentCircuitBreaker": circuit} + return svc + + +def _describe_services(*svcs, failures=None): + return ok_result({"services": list(svcs), "failures": failures or []}) + + +def _task_def( + *, + exec_role="arn:aws:iam::111122223333:role/ecsTaskExecutionRole", + task_role=None, + containers=None, +): + if containers is None: + containers = [ + { + "name": "web", + "image": "111122223333.dkr.ecr.us-east-1.amazonaws.com/web:1", + "environment": [ + {"name": "LOG_LEVEL", "value": "INFO"}, + {"name": "DB_PASSWORD", "value": "hunter2"}, + ], + "secrets": [ + { + "name": "API_TOKEN", + "valueFrom": ( + "arn:aws:ssm:us-east-1:111122223333:" + "parameter/api-token" + ), + } + ], + "logConfiguration": { + "logDriver": "awslogs", + "options": { + "awslogs-group": "/ecs/web", + "awslogs-stream-prefix": "ecs", + "awslogs-region": "us-east-1", + }, + }, + } + ] + td = {"family": "web", "revision": 5, "taskDefinitionArn": _TASK_DEF_ARN} + if exec_role: + td["executionRoleArn"] = exec_role + if task_role: + td["taskRoleArn"] = task_role + td["containerDefinitions"] = containers + return ok_result({"taskDefinition": td}) + + +def _iam_with_policies(): + """IAM reads for a role that HAS policies (attached).""" + return ( + ok_result({"AttachedPolicies": [{"PolicyName": "exec"}]}), + ok_result({"PolicyNames": []}), + ) + + +def _diag(call, **kwargs): + kwargs.setdefault("cluster", "prod") + kwargs.setdefault("clock", _fake_clock()) + return diagnose_ecs_service(call, "web", **kwargs) + + +def _no_writes(call): + assert not _MUTATING.intersection(call.ops()) + + +def _healthy_call(): + """A fully healthy service with every optional section populated.""" + call = ScriptedECS() + call.queue("ecs", "DescribeServices", _describe_services(_svc())) + call.queue("ecs", "DescribeTaskDefinition", _task_def()) + attached, inline = _iam_with_policies() + call.queue("iam", "ListAttachedRolePolicies", attached) + call.queue("iam", "ListRolePolicies", inline) + # STOPPED tasks: none. + call.queue("ecs", "ListTasks", ok_result({"taskArns": []})) + # RUNNING tasks: one healthy. + call.queue( + "ecs", + "ListTasks", + ok_result({"taskArns": ["arn:aws:ecs:us-east-1:1:task/prod/t1"]}), + ) + call.queue( + "ecs", + "DescribeTasks", + ok_result( + { + "tasks": [ + { + "taskArn": "arn:aws:ecs:us-east-1:1:task/prod/t1", + "lastStatus": "RUNNING", + "healthStatus": "HEALTHY", + "startedAt": "2026-01-01T00:00:00Z", + "containers": [ + {"name": "web", "lastStatus": "RUNNING"} + ], + } + ] + } + ), + ) + return call + + +# --------------------------------------------------------------- redaction + + +def test_redact_sensitive_scrubs_keyed_values_and_aws_keys(): + assert redact_sensitive("DB_PASSWORD=hunter2") == "DB_PASSWORD=[REDACTED]" + assert "[REDACTED]" in redact_sensitive("token: abc.def.ghi") + assert redact_sensitive("key AKIAIOSFODNN7EXAMPLE here") == ( + "key [REDACTED] here" + ) + # Non-sensitive text is returned unchanged. + assert redact_sensitive("plain log line") == "plain log line" + # Non-strings pass through. + assert redact_sensitive(None) is None + + +# --------------------------------------------------------------- happy path + + +def test_healthy_service_reports_no_blockers_and_only_reads(): + call = _healthy_call() + result = _diag(call) + assert result["ok"] is True + value = result["value"] + assert value["healthy"] is True + assert value["blockers"] == [] + assert value["service_summary"]["desired_count"] == 2 + assert value["service_summary"]["running_count"] == 2 + # Gating read comes first. + assert call.ops()[0] == ("ecs", "DescribeServices") + _no_writes(call) + + +def test_healthy_service_populates_optional_sections(): + call = _healthy_call() + value = _diag(call)["value"] + sections = value["sections"] + # Every attempted optional section is available. + assert sections["task_definition"]["available"] is True + assert sections["stopped_tasks"]["available"] is True + assert sections["running_tasks"]["available"] is True + assert value["running_tasks"]["count"] == 1 + assert value["stopped_tasks"]["count"] == 0 + + +# ------------------------------------------------------ redaction in config + + +def test_task_definition_config_emits_env_names_only_no_values(): + call = _healthy_call() + value = _diag(call)["value"] + container = value["task_definition"]["containers"][0] + # Names surfaced; values never emitted. + assert "DB_PASSWORD" in container["environment_names"] + assert "LOG_LEVEL" in container["environment_names"] + serialized = repr(value) + assert "hunter2" not in serialized + # Secret references surfaced by name only. + assert container["secret_names"] == ["API_TOKEN"] + + +# ---------------------------------------------------- service missing / inactive + + +def test_missing_service_returns_not_found(): + call = ScriptedECS() + call.queue( + "ecs", + "DescribeServices", + _describe_services(failures=[{"reason": "MISSING"}]), + ) + result = _diag(call) + assert result["ok"] is False + assert result["code"] == "NOT_FOUND" + # No optional section was queried once the service is absent. + assert call.ops() == [("ecs", "DescribeServices")] + + +def test_inactive_service_returns_not_found(): + call = ScriptedECS() + call.queue( + "ecs", "DescribeServices", _describe_services(_svc(status="INACTIVE")) + ) + result = _diag(call) + assert result["ok"] is False + assert result["code"] == "NOT_FOUND" + + +# --------------------------------------------------------------- gating failure + + +def test_gating_describe_services_failure_propagates_and_stops(): + call = ScriptedECS() + call.queue( + "ecs", + "DescribeServices", + err_result( + "PERMISSION_DENIED", "not authorized", aws_code="AccessDenied" + ), + ) + result = _diag(call) + assert result["ok"] is False + assert result["code"] == "PERMISSION_DENIED" + # Gating failure means nothing else is attempted. + assert call.ops() == [("ecs", "DescribeServices")] + + +# ---------------------------------------------------- partial optional failure + + +def test_optional_section_failure_degrades_but_continues(): + call = ScriptedECS() + call.queue("ecs", "DescribeServices", _describe_services(_svc())) + # Task definition read fails -> warning + section unavailable, run continues. + call.queue( + "ecs", + "DescribeTaskDefinition", + err_result("PERMISSION_DENIED", "denied", aws_code="AccessDenied"), + ) + call.queue("ecs", "ListTasks", ok_result({"taskArns": []})) + call.queue("ecs", "ListTasks", ok_result({"taskArns": []})) + result = _diag(call) + assert result["ok"] is True + value = result["value"] + assert value["sections"]["task_definition"]["available"] is False + # Stopped/running task sections were still attempted after the failure. + assert value["sections"]["stopped_tasks"]["available"] is True + # A degraded section makes the result non-exhaustive. + assert result["is_exhaustive"] is False + assert any(w["type"] == "TASK_DEF_UNREADABLE" for w in value["warnings"]) + + +# --------------------------------------------------------- stopped / failed tasks + + +def test_stopped_tasks_surface_reasons_and_blocker(): + call = ScriptedECS() + call.queue("ecs", "DescribeServices", _describe_services(_svc(running=0))) + call.queue("ecs", "DescribeTaskDefinition", _task_def(exec_role=None)) + # STOPPED tasks with a CannotPull reason. + call.queue( + "ecs", + "ListTasks", + ok_result({"taskArns": ["arn:aws:ecs:us-east-1:1:task/prod/s1"]}), + ) + call.queue( + "ecs", + "DescribeTasks", + ok_result( + { + "tasks": [ + { + "taskArn": "arn:aws:ecs:us-east-1:1:task/prod/s1", + "lastStatus": "STOPPED", + "stoppedReason": "CannotPullContainerError: not found", + "stopCode": "TaskFailedToStart", + "taskDefinitionArn": _TASK_DEF_ARN, + "containers": [ + { + "name": "web", + "exitCode": 1, + "reason": "CannotPullContainerError", + "lastStatus": "STOPPED", + } + ], + } + ] + } + ), + ) + call.queue("ecs", "ListTasks", ok_result({"taskArns": []})) + value = _diag(call)["value"] + assert value["stopped_tasks"]["count"] == 1 + task = value["stopped_tasks"]["tasks"][0] + assert task["stop_code"] == "TaskFailedToStart" + assert task["containers"][0]["exit_code"] == 1 + # Ground-truth stopped reason becomes a blocker. + assert any(b["type"] == "STOPPED_TASK_ERRORS" for b in value["blockers"]) + + +# --------------------------------------------------------------- deployment / events + + +def test_desired_count_zero_is_false_green_blocker(): + call = ScriptedECS() + call.queue( + "ecs", + "DescribeServices", + _describe_services(_svc(desired=0, running=0)), + ) + call.queue("ecs", "DescribeTaskDefinition", _task_def()) + attached, inline = _iam_with_policies() + call.queue("iam", "ListAttachedRolePolicies", attached) + call.queue("iam", "ListRolePolicies", inline) + call.queue("ecs", "ListTasks", ok_result({"taskArns": []})) + call.queue("ecs", "ListTasks", ok_result({"taskArns": []})) + value = _diag(call)["value"] + blockers = {b["type"]: b for b in value["blockers"]} + assert "DESIRED_COUNT_ZERO" in blockers + # The false-green note is present because rolloutState is COMPLETED. + assert "vacuously true" in blockers["DESIRED_COUNT_ZERO"]["evidence"] + + +def test_circuit_breaker_and_events_surfaced(): + call = ScriptedECS() + svc = _svc( + circuit={"enable": True, "rollback": True}, + events=[ + { + "createdAt": "2026-01-02T00:00:00Z", + "message": "unable to place a task", + } + ], + ) + call.queue("ecs", "DescribeServices", _describe_services(svc)) + call.queue("ecs", "DescribeTaskDefinition", _task_def()) + attached, inline = _iam_with_policies() + call.queue("iam", "ListAttachedRolePolicies", attached) + call.queue("iam", "ListRolePolicies", inline) + call.queue("ecs", "ListTasks", ok_result({"taskArns": []})) + call.queue("ecs", "ListTasks", ok_result({"taskArns": []})) + value = _diag(call)["value"] + assert value["deployment"]["circuit_breaker"] == { + "enabled": True, + "rollback": True, + } + assert value["recent_events"][0]["message"] == "unable to place a task" + + +# --------------------------------------------------------------- target health + + +def test_target_group_without_healthy_targets_is_blocker(): + lb = [ + { + "targetGroupArn": ( + "arn:aws:elasticloadbalancing:us-east-1:1:targetgroup/tg/abc" + ), + "containerName": "web", + "containerPort": 8080, + } + ] + call = ScriptedECS() + call.queue( + "ecs", "DescribeServices", _describe_services(_svc(load_balancers=lb)) + ) + call.queue("ecs", "DescribeTaskDefinition", _task_def()) + attached, inline = _iam_with_policies() + call.queue("iam", "ListAttachedRolePolicies", attached) + call.queue("iam", "ListRolePolicies", inline) + call.queue("ecs", "ListTasks", ok_result({"taskArns": []})) + call.queue("ecs", "ListTasks", ok_result({"taskArns": []})) + call.queue( + "elbv2", + "DescribeTargetHealth", + ok_result( + { + "TargetHealthDescriptions": [ + { + "Target": {"Id": "10.0.0.1"}, + "TargetHealth": { + "State": "unhealthy", + "Reason": "Target.FailedHealthChecks", + }, + } + ] + } + ), + ) + value = _diag(call)["value"] + assert value["target_health"]["wired"] is True + tg = value["target_health"]["target_groups"][0] + assert tg["healthy"] == 0 + assert tg["total"] == 1 + assert any( + b["type"] == "TARGET_GROUP_UNHEALTHY" for b in value["blockers"] + ) + + +def test_no_load_balancer_reports_not_wired(): + call = _healthy_call() + value = _diag(call)["value"] + assert value["target_health"]["wired"] is False + # No target-health call was issued when no LB is attached. + assert ("elbv2", "DescribeTargetHealth") not in call.ops() + + +# --------------------------------------------------------------- log evidence + + +def test_log_evidence_is_fetched_redacted_and_truncated(): + call = _healthy_call() + call.queue( + "logs", + "FilterLogEvents", + ok_result( + { + "events": [ + { + "timestamp": 1_700_000_000_000, + "logStreamName": "ecs/web/t1", + "message": "ERROR DB_PASSWORD=hunter2 " + "x" * 5000, + } + ] + } + ), + ) + value = _diag(call, max_message_length=100)["value"] + assert value["sections"]["log_evidence"]["available"] is True + line = value["log_evidence"]["lines"][0] + # Redacted and truncated. + assert "hunter2" not in line["message"] + assert "[REDACTED]" in line["message"] + assert len(line["message"]) <= 100 + + +def test_no_logs_flag_skips_log_fetch(): + call = _healthy_call() + value = _diag(call, include_logs=False)["value"] + assert ("logs", "FilterLogEvents") not in call.ops() + assert value["sections"]["log_evidence"]["available"] is False + + +# --------------------------------------------------------------- pagination / bounds + + +def test_list_tasks_pagination_is_bounded(): + call = ScriptedECS() + call.queue("ecs", "DescribeServices", _describe_services(_svc())) + call.queue("ecs", "DescribeTaskDefinition", _task_def()) + attached, inline = _iam_with_policies() + call.queue("iam", "ListAttachedRolePolicies", attached) + call.queue("iam", "ListRolePolicies", inline) + # STOPPED ListTasks keeps returning a nextToken; the loop must stop at + # max_pages rather than spinning forever. + for _ in range(50): + call.queue( + "ecs", + "ListTasks", + ok_result({"taskArns": [], "nextToken": "more"}), + ) + call.queue("ecs", "ListTasks", ok_result({"taskArns": []})) + result = _diag(call, max_pages=3) + assert result["ok"] is True + stopped_list_calls = sum( + 1 + for s, o, p in call.calls + if (s, o) == ("ecs", "ListTasks") + and p.get("desiredStatus") == "STOPPED" + ) + assert stopped_list_calls == 3 + + +def test_stopped_tasks_capped(): + call = ScriptedECS() + call.queue("ecs", "DescribeServices", _describe_services(_svc(running=0))) + call.queue("ecs", "DescribeTaskDefinition", _task_def()) + attached, inline = _iam_with_policies() + call.queue("iam", "ListAttachedRolePolicies", attached) + call.queue("iam", "ListRolePolicies", inline) + arns = [f"arn:aws:ecs:us-east-1:1:task/prod/s{i}" for i in range(10)] + call.queue("ecs", "ListTasks", ok_result({"taskArns": arns})) + call.queue( + "ecs", + "DescribeTasks", + ok_result( + { + "tasks": [ + { + "taskArn": a, + "lastStatus": "STOPPED", + "stoppedReason": "Essential container exited", + "stopCode": "EssentialContainerExited", + "containers": [], + } + for a in arns + ] + } + ), + ) + call.queue("ecs", "ListTasks", ok_result({"taskArns": []})) + value = _diag(call, max_stopped_tasks=3)["value"] + assert value["stopped_tasks"]["count"] == 3 + assert value["stopped_tasks"]["truncated"] is True + + +# --------------------------------------------------------------- empty optional data + + +def test_empty_optional_data_still_healthy(): + call = ScriptedECS() + call.queue("ecs", "DescribeServices", _describe_services(_svc())) + call.queue("ecs", "DescribeTaskDefinition", _task_def()) + attached, inline = _iam_with_policies() + call.queue("iam", "ListAttachedRolePolicies", attached) + call.queue("iam", "ListRolePolicies", inline) + call.queue("ecs", "ListTasks", ok_result({"taskArns": []})) + call.queue("ecs", "ListTasks", ok_result({"taskArns": []})) + result = _diag(call) + value = result["value"] + assert value["healthy"] is True + assert value["stopped_tasks"]["tasks"] == [] + assert value["target_health"]["wired"] is False + # Nothing failed or truncated -> exhaustive. + assert result["is_exhaustive"] is True + + +# --------------------------------------------------------------- IAM findings + + +def test_missing_execution_role_is_blocker(): + call = ScriptedECS() + call.queue("ecs", "DescribeServices", _describe_services(_svc())) + call.queue("ecs", "DescribeTaskDefinition", _task_def(exec_role=None)) + call.queue("ecs", "ListTasks", ok_result({"taskArns": []})) + call.queue("ecs", "ListTasks", ok_result({"taskArns": []})) + value = _diag(call)["value"] + assert any( + b["type"] == "EXECUTION_ROLE_MISSING" for b in value["blockers"] + ) + # No IAM call is made when there is no role to inspect. + assert ("iam", "ListAttachedRolePolicies") not in call.ops() + + +def test_execution_role_without_policies_is_warning(): + call = ScriptedECS() + call.queue("ecs", "DescribeServices", _describe_services(_svc())) + call.queue("ecs", "DescribeTaskDefinition", _task_def()) + call.queue( + "iam", "ListAttachedRolePolicies", ok_result({"AttachedPolicies": []}) + ) + call.queue("iam", "ListRolePolicies", ok_result({"PolicyNames": []})) + call.queue("ecs", "ListTasks", ok_result({"taskArns": []})) + call.queue("ecs", "ListTasks", ok_result({"taskArns": []})) + value = _diag(call)["value"] + assert any( + w["type"] == "EXECUTION_ROLE_MISSING_PERMS" for w in value["warnings"] + ) + + +# --------------------------------------------------------------- no-write guarantee + + +def test_full_run_issues_no_mutating_operation(): + call = _healthy_call() + call.queue("logs", "FilterLogEvents", ok_result({"events": []})) + _diag(call) + # Every recorded op is a read; not one mutating verb slipped through. + _no_writes(call) + for _, operation, _ in call.calls: + # assert_read_only would have raised already; re-assert for clarity. + assert_read_only(operation) + + +# ------------------------------------------------ extra findings / branches + + +def test_code_deploy_controller_adds_info_warning(): + call = ScriptedECS() + svc = _svc() + svc["deploymentController"] = {"type": "CODE_DEPLOY"} + call.queue("ecs", "DescribeServices", _describe_services(svc)) + call.queue("ecs", "DescribeTaskDefinition", _task_def()) + attached, inline = _iam_with_policies() + call.queue("iam", "ListAttachedRolePolicies", attached) + call.queue("iam", "ListRolePolicies", inline) + call.queue("ecs", "ListTasks", ok_result({"taskArns": []})) + call.queue("ecs", "ListTasks", ok_result({"taskArns": []})) + value = _diag(call)["value"] + assert any(w["type"] == "CODE_DEPLOY_SERVICE" for w in value["warnings"]) + + +def test_task_role_without_policies_is_warning(): + call = ScriptedECS() + call.queue("ecs", "DescribeServices", _describe_services(_svc())) + call.queue( + "ecs", + "DescribeTaskDefinition", + _task_def(task_role="arn:aws:iam::111122223333:role/app-task"), + ) + # Execution role: has policies. Task role: no policies -> warning. + call.queue( + "iam", + "ListAttachedRolePolicies", + ok_result({"AttachedPolicies": [{"PolicyName": "exec"}]}), + ) + call.queue("iam", "ListRolePolicies", ok_result({"PolicyNames": []})) + call.queue( + "iam", "ListAttachedRolePolicies", ok_result({"AttachedPolicies": []}) + ) + call.queue("iam", "ListRolePolicies", ok_result({"PolicyNames": []})) + call.queue("ecs", "ListTasks", ok_result({"taskArns": []})) + call.queue("ecs", "ListTasks", ok_result({"taskArns": []})) + value = _diag(call)["value"] + assert any( + w["type"] == "TASK_ROLE_MISSING_PERMS" for w in value["warnings"] + ) + assert "task_role" not in value["verified_ok"] + + +def test_task_role_with_policies_is_verified_ok(): + call = ScriptedECS() + call.queue("ecs", "DescribeServices", _describe_services(_svc())) + call.queue( + "ecs", + "DescribeTaskDefinition", + _task_def(task_role="arn:aws:iam::111122223333:role/app-task"), + ) + for _ in range(2): + call.queue( + "iam", + "ListAttachedRolePolicies", + ok_result({"AttachedPolicies": [{"PolicyName": "p"}]}), + ) + call.queue("iam", "ListRolePolicies", ok_result({"PolicyNames": []})) + call.queue("ecs", "ListTasks", ok_result({"taskArns": []})) + call.queue("ecs", "ListTasks", ok_result({"taskArns": []})) + value = _diag(call)["value"] + assert "task_role" in value["verified_ok"] + assert value["sections"]["iam"]["available"] is True + + +def test_iam_read_failure_degrades_iam_section(): + call = ScriptedECS() + call.queue("ecs", "DescribeServices", _describe_services(_svc())) + call.queue("ecs", "DescribeTaskDefinition", _task_def()) + call.queue( + "iam", + "ListAttachedRolePolicies", + err_result("PERMISSION_DENIED", "denied", aws_code="AccessDenied"), + ) + call.queue("iam", "ListRolePolicies", ok_result({"PolicyNames": []})) + call.queue("ecs", "ListTasks", ok_result({"taskArns": []})) + call.queue("ecs", "ListTasks", ok_result({"taskArns": []})) + result = _diag(call) + value = result["value"] + assert value["sections"]["iam"]["available"] is False + assert any(w["type"] == "IAM_CHECK_DEGRADED" for w in value["warnings"]) + assert result["is_exhaustive"] is False + + +def test_no_task_definition_ref_skips_task_sections(): + call = ScriptedECS() + svc = _svc(task_def="") + call.queue("ecs", "DescribeServices", _describe_services(svc)) + call.queue("ecs", "ListTasks", ok_result({"taskArns": []})) + call.queue("ecs", "ListTasks", ok_result({"taskArns": []})) + value = _diag(call)["value"] + # No task definition was referenced, so that section is absent and no IAM + # read was issued. + assert "task_definition" not in value["sections"] + assert ("ecs", "DescribeTaskDefinition") not in call.ops() + + +# --------------------------------------------- optional-read failures + + +def test_stopped_tasks_list_failure_degrades_section(): + call = ScriptedECS() + call.queue("ecs", "DescribeServices", _describe_services(_svc())) + call.queue("ecs", "DescribeTaskDefinition", _task_def()) + attached, inline = _iam_with_policies() + call.queue("iam", "ListAttachedRolePolicies", attached) + call.queue("iam", "ListRolePolicies", inline) + call.queue( + "ecs", + "ListTasks", + err_result("THROTTLED", "slow down", aws_code="ThrottlingException"), + ) + call.queue("ecs", "ListTasks", ok_result({"taskArns": []})) + result = _diag(call) + assert result["value"]["sections"]["stopped_tasks"]["available"] is False + assert result["is_exhaustive"] is False + + +def test_running_tasks_describe_failure_degrades_section(): + call = ScriptedECS() + call.queue("ecs", "DescribeServices", _describe_services(_svc())) + call.queue("ecs", "DescribeTaskDefinition", _task_def()) + attached, inline = _iam_with_policies() + call.queue("iam", "ListAttachedRolePolicies", attached) + call.queue("iam", "ListRolePolicies", inline) + call.queue("ecs", "ListTasks", ok_result({"taskArns": []})) + call.queue( + "ecs", + "ListTasks", + ok_result({"taskArns": ["arn:aws:ecs:us-east-1:1:task/prod/r1"]}), + ) + call.queue( + "ecs", + "DescribeTasks", + err_result("SERVICE_ERROR", "boom"), + ) + result = _diag(call) + assert result["value"]["sections"]["running_tasks"]["available"] is False + + +def test_target_health_read_failure_degrades_and_keeps_others(): + lb = [ + {"targetGroupArn": "arn:aws:elasticloadbalancing:us-east-1:1:tg/a/1"}, + {"targetGroupArn": "arn:aws:elasticloadbalancing:us-east-1:1:tg/b/2"}, + ] + call = ScriptedECS() + call.queue( + "ecs", "DescribeServices", _describe_services(_svc(load_balancers=lb)) + ) + call.queue("ecs", "DescribeTaskDefinition", _task_def()) + attached, inline = _iam_with_policies() + call.queue("iam", "ListAttachedRolePolicies", attached) + call.queue("iam", "ListRolePolicies", inline) + call.queue("ecs", "ListTasks", ok_result({"taskArns": []})) + call.queue("ecs", "ListTasks", ok_result({"taskArns": []})) + call.queue( + "elbv2", + "DescribeTargetHealth", + ok_result( + { + "TargetHealthDescriptions": [ + {"TargetHealth": {"State": "healthy"}} + ] + } + ), + ) + call.queue( + "elbv2", + "DescribeTargetHealth", + err_result("SERVICE_ERROR", "elb down"), + ) + result = _diag(call) + value = result["value"] + assert value["sections"]["target_health"]["available"] is False + # The group that resolved is still present alongside the errored one. + groups = value["target_health"]["target_groups"] + assert any("error" in g for g in groups) + assert any(g.get("healthy") == 1 for g in groups) + + +def test_target_groups_capped_marks_degraded(): + lb = [ + { + "targetGroupArn": f"arn:aws:elasticloadbalancing:us-east-1:1:tg/{i}/x" + } + for i in range(3) + ] + call = ScriptedECS() + call.queue( + "ecs", "DescribeServices", _describe_services(_svc(load_balancers=lb)) + ) + call.queue("ecs", "DescribeTaskDefinition", _task_def()) + attached, inline = _iam_with_policies() + call.queue("iam", "ListAttachedRolePolicies", attached) + call.queue("iam", "ListRolePolicies", inline) + call.queue("ecs", "ListTasks", ok_result({"taskArns": []})) + call.queue("ecs", "ListTasks", ok_result({"taskArns": []})) + call.queue( + "elbv2", + "DescribeTargetHealth", + ok_result( + { + "TargetHealthDescriptions": [ + {"TargetHealth": {"State": "healthy"}} + ] + } + ), + ) + result = _diag(call, max_target_groups=1) + assert result["is_exhaustive"] is False + # Only one target group was queried under the cap. + assert sum(1 for s, o, _ in call.calls if o == "DescribeTargetHealth") == 1 + + +def test_log_evidence_fetch_failure_degrades_section(): + call = _healthy_call() + call.queue( + "logs", + "FilterLogEvents", + err_result("PERMISSION_DENIED", "no logs:FilterLogEvents"), + ) + result = _diag(call) + assert result["value"]["sections"]["log_evidence"]["available"] is False + assert result["is_exhaustive"] is False + + +def test_log_evidence_skipped_when_no_awslogs_driver(): + containers = [ + { + "name": "web", + "image": "web:1", + "logConfiguration": {"logDriver": "json-file"}, + } + ] + call = ScriptedECS() + call.queue("ecs", "DescribeServices", _describe_services(_svc())) + call.queue( + "ecs", "DescribeTaskDefinition", _task_def(containers=containers) + ) + attached, inline = _iam_with_policies() + call.queue("iam", "ListAttachedRolePolicies", attached) + call.queue("iam", "ListRolePolicies", inline) + call.queue("ecs", "ListTasks", ok_result({"taskArns": []})) + call.queue("ecs", "ListTasks", ok_result({"taskArns": []})) + result = _diag(call) + # No awslogs group -> section unavailable but not a loss -> still exhaustive. + assert result["value"]["sections"]["log_evidence"]["available"] is False + assert ("logs", "FilterLogEvents") not in call.ops() + assert result["is_exhaustive"] is True + + +# --------------------------------------------- validation / internal error + + +def test_empty_service_rejected_before_dispatch(): + call = ScriptedECS() + result = diagnose_ecs_service(call, " ", cluster="prod") + assert result["ok"] is False + assert result["code"] == "INVALID_INPUT" + assert call.calls == [] + + +def test_non_positive_cap_rejected_before_dispatch(): + call = ScriptedECS() + result = _diag(call, max_events=0) + assert result["ok"] is False + assert result["code"] == "INVALID_INPUT" + assert call.calls == [] + + +def test_internal_exception_is_returned_as_service_error(): + def exploding(service, operation, params): + raise RuntimeError("kaboom") + + result = diagnose_ecs_service(exploding, "web", cluster="prod") + assert result["ok"] is False + assert result["code"] == "SERVICE_ERROR" + assert "internal error" in result["error"]