Update subscriber to accept combination of inputs

.github/workflows/ci.yaml

@@ -20,7 +20,7 @@ jobs:
       - run: wget https://github.com/mikefarah/yq/releases/download/v$YQ_VERSION/yq_linux_amd64.tar.gz -O - | tar xz && mv yq_linux_amd64 /usr/local/bin/yq
       - run: |-
           mkdir build
-          yq ".Resources.LogsLambda.Properties.Code.ZipFile = \"$(sed  's/\"/\\\"/g' ./handler.py)\""  cloudformation-stacks/forwarder.template.yaml > build/axiom-cloudwatch-forwarder-cloudformation-stack.yaml
+          yq ".Resources.AxiomCloudWatchForwarder.Properties.Code.ZipFile = \"$(sed  's/\"/\\\"/g' ./handler.py)\""  cloudformation-stacks/forwarder.template.yaml > build/axiom-cloudwatch-forwarder-cloudformation-stack.yaml
           yq ".Resources.SubscriberLambda.Properties.Code.ZipFile = \"$(sed  's/\"/\\\"/g' ./subscriber.py)\""  cloudformation-stacks/subscriber.template.yaml > build/axiom-cloudwatch-subscriber-cloudformation-stack.yaml
           yq ".Resources.AxiomCloudWatchLogGroupsListener.Properties.Code.ZipFile = \"$(sed  's/\"/\\\"/g' ./logs_subscriber.py)\""  cloudformation-stacks/log-groups-listener.template.yaml > build/axiom-cloudwatch-log-groups-listener-cloudformation-stack.yaml
       - run: cat build/*

.github/workflows/release.yaml

@@ -21,7 +21,7 @@ jobs:
       - run: wget https://github.com/mikefarah/yq/releases/download/v$YQ_VERSION/yq_linux_amd64.tar.gz -O - | tar xz && mv yq_linux_amd64 /usr/local/bin/yq
       - run: |-
           mkdir build
-          yq ".Resources.LogsLambda.Properties.Code.ZipFile = \"$(sed  's/\"/\\\"/g' ./handler.py)\""  cloudformation-stacks/forwarder.template.yaml > build/axiom-cloudwatch-forwarder-cloudformation-stack.yaml
+          yq ".Resources.AxiomCloudWatchForwarder.Properties.Code.ZipFile = \"$(sed  's/\"/\\\"/g' ./handler.py)\""  cloudformation-stacks/forwarder.template.yaml > build/axiom-cloudwatch-forwarder-cloudformation-stack.yaml
           yq ".Resources.SubscriberLambda.Properties.Code.ZipFile = \"$(sed  's/\"/\\\"/g' ./subscriber.py)\"" cloudformation-stacks/subscriber.template.yaml > build/axiom-cloudwatch-subscriber-cloudformation-stack.yaml
           yq ".Resources.AxiomCloudWatchLogGroupsListener.Properties.Code.ZipFile = \"$(sed  's/\"/\\\"/g' ./logs_subscriber.py)\""  cloudformation-stacks/subscriber.template.yaml > build/axiom-cloudwatch-log-groups-listener-cloudformation-stack.yaml
       - name: Configure AWS Credentials

cloudformation-stacks/forwarder.template.yaml

@@ -15,10 +15,6 @@ Parameters:
     Type: String
     Description: The Name of the dataset in Axiom to push events to.
     AllowedPattern: ".+" # required
-  CloudWatchLogGroupNames:
-    Type: CommaDelimitedList
-    Description: The names of the AWS CloudWatch log groups to subscribe to. Comma Separated string of CloudWatch log group names.
-    AllowedPattern: ".*" # optional
   LambdaFunctionName:
     Type: String
     Description: Name of the AWS Lambda function.
@@ -27,25 +23,8 @@ Parameters:
   DataTags:
     Type: String
     Description: Tags to be included with the data ingested into axiom. e.g. <key1>=<value1>,<key2>=<value2>.
-Conditions:
-  HasCloudWatchLogGroupNames: !Not
-    - !Equals
-      - !Join ["", !Ref CloudWatchLogGroupNames]
-      - ""
 Resources:
-  "Fn::ForEach::SubscriptionFilters":
-    - GroupName
-    - !Ref CloudWatchLogGroupNames
-    - "LGSF&{GroupName}":
-        Type: AWS::Logs::SubscriptionFilter
-        Condition: HasCloudWatchLogGroupNames
-        Properties:
-          DestinationArn: !GetAtt
-            - LogsLambda
-            - Arn
-          FilterPattern: ""
-          LogGroupName: !Ref GroupName
-  LogsRole:
+  AxiomCloudWatchForwarderRole:
     Type: AWS::IAM::Role
     Properties:
       AssumeRolePolicyDocument:
@@ -58,7 +37,7 @@ Resources:
                 - lambda.amazonaws.com
       ManagedPolicyArns:
         - "arn:aws:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole"
-  LogsLambda:
+  AxiomCloudWatchForwarder:
     Type: AWS::Lambda::Function
     Properties:
       FunctionName: !Ref LambdaFunctionName
@@ -69,7 +48,7 @@ Resources:
           # DO NOT EDIT
           # CI will replace these comments with the code from ./handler.py
       Role: !GetAtt
-        - LogsRole
+        - AxiomCloudWatchForwarderRole
         - Arn
       Environment:
         Variables:
@@ -78,27 +57,7 @@ Resources:
           AXIOM_URL: !Ref "AxiomURL"
           DISABLE_JSON: !Ref DisableJSON
           DATA_TAGS: !Ref DataTags
-  "Fn::ForEach::LambdaPermissions":
-    - GroupName
-    - !Ref CloudWatchLogGroupNames
-    - "LogsLambdaPermission&{GroupName}":
-        Type: AWS::Lambda::Permission
-        Condition: HasCloudWatchLogGroupNames
-        DependsOn:
-          - LogsLambda
-        Properties:
-          Action: lambda:InvokeFunction
-          FunctionName: !Ref "LogsLambda"
-          Principal: !Sub
-            - "logs.${Region}.amazonaws.com"
-            - Region: !Ref "AWS::Region"
-          SourceAccount: !Ref "AWS::AccountId"
-          SourceArn: !Sub
-            - "arn:aws:logs:${Region}:${AccountID}:log-group:${LogGroupName}:*"
-            - AccountID: !Ref "AWS::AccountId"
-              Region: !Ref "AWS::Region"
-              LogGroupName: !Ref GroupName
 Outputs:
-  LogsLambdaARN:
+  AxiomCloudWatchForwarderARN:
     Description: The ARN of the created Forwarder Lambda
-    Value: !GetAtt LogsLambda.Arn
+    Value: !GetAtt AxiomCloudWatchForwarder.Arn

cloudformation-stacks/subscriber.template.yaml

@@ -8,10 +8,21 @@ Parameters:
     Type: String
     Description: The ARN of the Axiom CloudWatch Forwarder Lambda function used to ship logs to Axiom.
     AllowedPattern: ".+" # required
+  CloudWatchLogGroupsNames:
+    Type: String
+    Description: A comma separated list of CloudWatch log groups to subscribe to.
+    Default: "" # all
+    required: false
   CloudWatchLogGroupsPrefix:
     Type: String
-    Description: The Prefix of CloudWatch log groups to trigger the Axiom CloudWatch Forwarder lambda.
+    Description: The Prefix of CloudWatch log groups to subscribe to.
+    Default: "" # all
+    required: false
+  CloudWatchLogGroupsPattern:
+    Type: String
+    Description: A regular expression pattern of CloudWatch log groups to subscribe to.
     Default: "" # all
+    required: false
 Resources:
   SubscriberPolicy:
     Type: AWS::IAM::Policy
@@ -59,7 +70,9 @@ Resources:
       Environment:
         Variables:
           AXIOM_CLOUDWATCH_FORWARDER_LAMBDA_ARN: !Ref 'AxiomCloudWatchForwarderLambdaARN'
+          LOG_GROUP_NAMES: !Ref 'CloudWatchLogGroupsNames'
           LOG_GROUP_PREFIX: !Ref 'CloudWatchLogGroupsPrefix'
+          LOG_GROUP_PATTERN: !Ref 'CloudWatchLogGroupsPattern'
   SubscriberInvoker:
     Type: AWS::CloudFormation::CustomResource
     DependsOn: SubscriberLambda

log_groups_listener.py

@@ -4,7 +4,9 @@
 import logging
 
 # Set environment variables.
-axiom_cloudwatch_forwarder_lambda_arn = os.getenv("AXIOM_CLOUDWATCH_FORWARDER_LAMBDA_ARN")
+axiom_cloudwatch_forwarder_lambda_arn = os.getenv(
+    "AXIOM_CLOUDWATCH_FORWARDER_LAMBDA_ARN"
+)
 log_group_prefix = os.getenv("LOG_GROUP_PREFIX", "")
 
 # set logger

subscriber.py

@@ -1,3 +1,4 @@
+import re
 import boto3
 import os
 import logging
@@ -12,33 +13,49 @@
 cloudwatch_logs_client = boto3.client("logs")
 lambda_client = boto3.client("lambda")
 
-axiom_cloudwatch_forwarder_lambda_arn = os.getenv("AXIOM_CLOUDWATCH_FORWARDER_LAMBDA_ARN")
+axiom_cloudwatch_forwarder_lambda_arn = os.getenv(
+    "AXIOM_CLOUDWATCH_FORWARDER_LAMBDA_ARN"
+)
+log_group_names = os.getenv("LOG_GROUP_NAMES", "")
 log_group_prefix = os.getenv("LOG_GROUP_PREFIX", "")
-log_groups_return_limit = int(os.getenv("LOG_GROUPS_LIMIT", 10))
-
+log_group_pattern = os.getenv("LOG_GROUP_PATTERN", "")
+log_groups_return_limit = int(os.getenv("LOG_GROUPS_LIMIT", 50))
+
+
+def build_groups_list(all_groups, names, pattern, prefix):
+    # filter out the log groups based on the names, pattern, and prefix provided in the environment variables
+    groups = []
+    for g in all_groups:
+        group = {"name": g["logGroupName"], "arn": g["arn"]}
+        if names is not None and group["name"] in names:
+            groups.append(group)
+            continue
+        elif prefix is not None and group["name"].startswith(prefix):
+            groups.append(group)
+            continue
+        elif pattern is not None and re.match(pattern, group["name"]):
+            groups.append(group)
+
+    return groups
+
+
+def get_log_groups(nextToken=None):
+    # check docs:
+    # 1. boto3 https://boto3.amazonaws.com/v1/documentation/api/1.9.42/reference/services/logs.html#CloudWatchLogs.Client.describe_log_groups
+    # 2. AWS API https://docs.aws.amazon.com/AmazonCloudWatchLogs/latest/APIReference/API_DescribeLogGroups.html#API_DescribeLogGroups_RequestSyntax
+    resp = cloudwatch_logs_client.describe_log_groups(limit=log_groups_return_limit)
+    all_groups = resp["logGroups"]
+    nextToken = resp["nextToken"]
+    # continue fetching log groups until nextToken is None
+    while nextToken is not None:
+        resp = cloudwatch_logs_client.describe_log_groups(
+            limit=log_groups_return_limit, nextToken=nextToken
+        )
+        all_groups.extend(resp["logGroups"])
+        # print(f'got ${len(all_groups)} groups')
+        nextToken = resp["nextToken"] if "nextToken" in resp else None
 
-def get_log_groups(token=None):
-    if token is None:
-        if log_group_prefix != "":
-            return cloudwatch_logs_client.describe_log_groups(
-                logGroupNamePrefix=log_group_prefix, limit=log_groups_return_limit
-            )
-        else:
-            return cloudwatch_logs_client.describe_log_groups(
-                limit=log_groups_return_limit
-            )
-    else:
-        if log_group_prefix != "":
-            return cloudwatch_logs_client.describe_log_groups(
-                logGroupNamePrefix=log_group_prefix,
-                nextToken=token,
-                limit=log_groups_return_limit,
-            )
-        else:
-            return cloudwatch_logs_client.describe_log_groups(
-                nextToken=token,
-                limit=log_groups_return_limit,
-            )
+    return all_groups
 
 
 def remove_permission(lambda_arn):
@@ -117,21 +134,19 @@ def lambda_handler(event: dict, context=None):
 
     create_statement(region, aws_account_id, axiom_cloudwatch_forwarder_lambda_arn)
 
-    ingester_lambda_group_name = (
+    forwarder_lambda_group_name = (
         "/aws/lambda/" + axiom_cloudwatch_forwarder_lambda_arn.split(":")[-1]
     )
 
-    def log_groups(token=None):
-        groups_response = get_log_groups(token)
-        groups = groups_response["logGroups"]
-        token = groups_response["nextToken"] if "nextToken" in groups_response else None
-
-        if len(groups) == 0:
-            return
+    log_groups = build_groups_list(
+        get_log_groups(), log_group_names, log_group_pattern, log_group_prefix
+    )
 
-        for group in groups:
-            # skip the ingester lambda log group to avoid circular logging
-            if group["logGroupName"] == ingester_lambda_group_name:
+    responseData = {}
+    try:
+        for group in log_groups:
+            # skip the Forwarder lambda log group to avoid circular logging
+            if group["logGroupName"] == forwarder_lambda_group_name:
                 continue
 
             try:
@@ -152,18 +167,6 @@ def log_groups(token=None):
                 )
                 logger.error(error)
                 continue
-
-        if token is None:
-            return
-
-        try:
-            log_groups(token)
-        except Exception as e:
-            raise e
-
-    responseData = {}
-    try:
-        log_groups()
     except Exception as e:
         responseData["success"] = "False"
         if "ResponseURL" in event: