Python: Add log group argument to CloudWatchQuery. Addresses #7022 (#7066)

Author: sslogar

python/example_code/cloudwatch-logs/README.md

@@ -38,8 +38,8 @@ python -m pip install -r requirements.txt
 
 Code excerpts that show you how to call individual service functions.
 
-- [GetQueryResults](scenarios/large-query/cloudwatch_query.py#L200)
-- [StartQuery](scenarios/large-query/cloudwatch_query.py#L126)
+- [GetQueryResults](scenarios/large-query/cloudwatch_query.py#L204)
+- [StartQuery](scenarios/large-query/cloudwatch_query.py#L130)
 
 ### Scenarios
 
@@ -106,4 +106,4 @@ in the `python` folder.
 
 Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 
-SPDX-License-Identifier: Apache-2.0
+SPDX-License-Identifier: Apache-2.0

python/example_code/cloudwatch-logs/scenarios/large-query/README.md

@@ -33,6 +33,7 @@ Use the following steps to create the necessary resources in AWS CloudFormation
 1. Run `aws cloudformation deploy --template-file stack.yaml --stack-name CloudWatchLargeQuery`
 1. Run `./make-log-files.sh`. This will output two timestamps for use in the following step.
 1. Run `export QUERY_START_DATE=<QUERY_START_DATE>`. Replace `<QUERY_START_DATE>` with the output from the previous step. Repeat this for `QUERY_END_DATE`.
+1. Optional: Run `export QUERY_LOG_GROUP=<QUERY_LOG_GROUP>`. Replace `<QUERY_LOG_GROUP>` with your preferred log group.
 1. Run `./put-log-events.sh`.
 1. Wait five minutes for logs to settle and to make sure you're not querying for logs that exist in the future.
 

python/example_code/cloudwatch-logs/scenarios/large-query/cloudwatch_query.py

@@ -8,6 +8,8 @@
 
 from date_utilities import DateUtilities
 
+DEFAULT_QUERY = "fields @timestamp, @message | sort @timestamp asc"
+DEFAULT_LOG_GROUP = "/workflows/cloudwatch-logs/large-query"
 
 class DateOutOfBoundsError(Exception):
  """Exception raised when the date range for a query is out of bounds."""
@@ -19,17 +21,18 @@ class CloudWatchQuery:
  """
  A class to query AWS CloudWatch logs within a specified date range.
 
- :ivar date_range: Start and end datetime for the query.
  :vartype date_range: tuple
  :ivar limit: Maximum number of log entries to return.
  :vartype limit: int
+ :log_group str: Name of the log group to query
+ :query_string str: query
  """
 
- def __init__(self, date_range):
+ def __init__(self, log_group: str = DEFAULT_LOG_GROUP, query_string: str=DEFAULT_QUERY) -> None:
  self.lock = threading.Lock()
- self.log_groups = "/workflows/cloudwatch-logs/large-query"
+ self.log_group = log_group
+ self.query_string = query_string
  self.query_results = []
- self.date_range = date_range
  self.query_duration = None
  self.datetime_format = "%Y-%m-%d %H:%M:%S.%f"
  self.date_utilities = DateUtilities()
@@ -50,8 +53,9 @@ def query_logs(self, date_range):
 
  logging.info(
  f"Original query:"
- f"\n START: {start_date}"
- f"\n END: {end_date}"
+ f"\n START: {start_date}"
+ f"\n END: {end_date}"
+ f"\n LOG GROUP: {self.log_group}"
  )
  self.recursive_query((start_date, end_date))
  end_time = datetime.now()
@@ -143,10 +147,10 @@ def perform_query(self, date_range):
  self.date_utilities.convert_iso8601_to_unix_timestamp(date_range[1])
  )
  response = client.start_query(
- logGroupName=self.log_groups,
+ logGroupName=self.log_group,
  startTime=start_time,
  endTime=end_time,
- queryString="fields @timestamp, @message | sort @timestamp asc",
+ queryString=self.query_string,
  limit=self.limit,
  )
  query_id = response["queryId"]
@@ -185,10 +189,10 @@ def _initiate_query(self, client, date_range, max_logs):
  self.date_utilities.convert_iso8601_to_unix_timestamp(date_range[1])
  )
  response = client.start_query(
- logGroupName=self.log_groups,
+ logGroupName=self.log_group,
  startTime=start_time,
  endTime=end_time,
- queryString="fields @timestamp, @message | sort @timestamp asc",
+ queryString=self.query_string,
  limit=max_logs,
  )
  return response["queryId"]

python/example_code/cloudwatch-logs/scenarios/large-query/exec.py

@@ -16,6 +16,8 @@
  format="%(asctime)s - %(levelname)s - %(filename)s:%(lineno)d - %(message)s",
 )
 
+DEFAULT_QUERY_LOG_GROUP = "/workflows/cloudwatch-logs/large-query"
+
 
 class CloudWatchLogsQueryRunner:
  def __init__(self):
@@ -42,8 +44,10 @@ def create_cloudwatch_logs_client(self):
  def fetch_environment_variables(self):
  """
  Fetches and validates required environment variables for query start and end dates.
+ Fetches the environment variable for log group, returning the default value if it
+ does not exist.
 
- :return: Tuple of query start date and end date as integers.
+ :return: Tuple of query start date and end date as integers and the log group.
  :rtype: tuple
  :raises SystemExit: If required environment variables are missing or invalid.
  """
@@ -58,8 +62,14 @@ def fetch_environment_variables(self):
  except ValueError as e:
  logging.error(f"Error parsing date environment variables: {e}")
  sys.exit(1)
+ 
+ try:
+ log_group = os.environ["QUERY_LOG_GROUP"]
+ except KeyError:
+ logging.warning("No QUERY_LOG_GROUP environment variable, using default value")
+ log_group = DEFAULT_QUERY_LOG_GROUP
 
- return query_start_date, query_end_date
+ return query_start_date, query_end_date, log_group
 
  def convert_dates_to_iso8601(self, start_date, end_date):
  """
@@ -85,6 +95,7 @@ def execute_query(
  start_date_iso8601,
  end_date_iso8601,
  log_group="/workflows/cloudwatch-logs/large-query",
+ query="fields @timestamp, @message | sort @timestamp asc"
  ):
  """
  Creates a CloudWatchQuery instance and executes the query with provided date range.
@@ -95,9 +106,12 @@ def execute_query(
  :type end_date_iso8601: str
  :param log_group: Log group to search: "/workflows/cloudwatch-logs/large-query"
  :type log_group: str
+ :param query: Query string to pass to the CloudWatchQuery instance
+ :type query: str
  """
  cloudwatch_query = CloudWatchQuery(
- [start_date_iso8601, end_date_iso8601],
+ log_group=log_group,
+ query_string=query
  )
  cloudwatch_query.query_logs((start_date_iso8601, end_date_iso8601))
  logging.info("Query executed successfully.")
@@ -113,12 +127,12 @@ def main():
  """
  logging.info("Starting a recursive CloudWatch logs query...")
  runner = CloudWatchLogsQueryRunner()
- query_start_date, query_end_date = runner.fetch_environment_variables()
+ query_start_date, query_end_date, log_group = runner.fetch_environment_variables()
  start_date_iso8601 = DateUtilities.convert_unix_timestamp_to_iso8601(
  query_start_date
  )
  end_date_iso8601 = DateUtilities.convert_unix_timestamp_to_iso8601(query_end_date)
- runner.execute_query(start_date_iso8601, end_date_iso8601)
+ runner.execute_query(start_date_iso8601, end_date_iso8601, log_group=log_group)
 
 
 if __name__ == "__main__":