Di Snapshot Queries

1"""CloudWatch Logs Insights query helpers for snapshot tools."""
2
3import time
4
5import di_logs_client as aws_clients
6from botocore.exceptions import BotoCoreError, ClientError
7
8
9def _execute_cloudwatch_query(
10    query_string: str,
11    start_epoch: int,
12    end_epoch: int,
13    log_group_name: str,
14    max_timeout: int = 30,
15) -> dict:
16    """Execute a CloudWatch Logs Insights query and poll for results."""
17    logs = aws_clients.logs_client
18
19    try:
20        start_response = logs.start_query(
21            logGroupName=log_group_name,
22            startTime=start_epoch,
23            endTime=end_epoch,
24            queryString=query_string,
25        )
26    except ClientError as exc:
27        return {
28            "status": "Error",
29            "error": f"Failed to start query: {exc}",
30            "results": [],
31        }
32    except BotoCoreError as exc:
33        return {"status": "Error", "error": str(exc), "results": []}
34
35    query_id = start_response.get("queryId")
36    if not query_id:
37        return {
38            "status": "Error",
39            "error": f"start_query did not return a queryId (response: {start_response})",
40            "results": [],
41        }
42
43    poll_start = time.time()
44    while poll_start + max_timeout > time.time():
45        try:
46            response = logs.get_query_results(queryId=query_id)
47        except ClientError as exc:
48            return {
49                "status": "Error",
50                "error": f"Failed to get results: {exc}",
51                "results": [],
52                "queryId": query_id,
53            }
54        except BotoCoreError as exc:
55            return {
56                "status": "Error",
57                "error": str(exc),
58                "results": [],
59                "queryId": query_id,
60            }
61
62        status = response.get("status", "Unknown")
63        if status in {"Complete", "Failed", "Cancelled"}:
64            results = [
65                {field.get("field", ""): field.get("value", "") for field in line}
66                for line in response.get("results", [])
67            ]
68            return {
69                "status": status,
70                "queryId": query_id,
71                "results": results,
72                "messages": response.get("messages", []),
73            }
74
75        time.sleep(1)
76
77    return {
78        "status": "Polling Timeout",
79        "queryId": query_id,
80        "results": [],
81        "error": f"Query did not complete within {max_timeout} seconds.",
82    }