updating the code for existing log groups. Existing is called when user select existing as true in parameters.

sourabh · sourabh · commit c9ec731bc546 · 2021-03-26T16:08:01.000+05:30
diff --git a/loggroup-lambda-connector/sam/packaged.yaml b/loggroup-lambda-connector/sam/packaged.yaml
@@ -6,7 +6,7 @@ Description: '"Lambda Function for automatic subscription of any Sumo Logic lamb
   '
 Globals:
   Function:
-    Timeout: 300
+    Timeout: 900
     MemorySize: 128
 Metadata:
   AWS::ServerlessRepo::Application:
@@ -73,6 +73,10 @@ Conditions:
     Fn::Equals:
     - Ref: DestinationArnType
     - Kinesis
+  invoke_existing:
+    Fn::Equals:
+    - Ref: UseExistingLogs
+    - 'true'
 Rules:
   testRoleArnWithLambda:
     RuleCondition:
@@ -103,17 +107,15 @@ Resources:
   SumoLogGroupLambdaConnector:
     Type: AWS::Serverless::Function
     Properties:
-      CodeUri: s3://cf-templates-1qpf3unpuo1hw-us-east-1/test-log-group-lambda-connector/44a32060109e4c7efa4feed59007b04f
+      CodeUri: s3://cf-templates-1qpf3unpuo1hw-us-east-1/test-log-group-lambda-connector/0d12ca7b2219c4fc3fe3cd29d6ba56d0
       Handler: loggroup-lambda-connector.handler
-      Runtime: nodejs14.x
+      Runtime: nodejs12.x
       Environment:
         Variables:
           DESTINATION_ARN:
             Ref: DestinationArnValue
           LOG_GROUP_PATTERN:
             Ref: LogGroupPattern
-          USE_EXISTING_LOG_GROUPS:
-            Ref: UseExistingLogs
           LOG_GROUP_TAGS:
             Fn::Join:
             - ','
@@ -130,6 +132,12 @@ Resources:
           - logs:PutSubscriptionFilter
           Resource:
           - Fn::Sub: arn:aws:logs:${AWS::Region}:${AWS::AccountId}:log-group:*
+        - Sid: InvokePolicy
+          Effect: Allow
+          Action:
+          - lambda:InvokeFunction
+          Resource:
+          - Fn::Sub: arn:aws:lambda:${AWS::Region}:${AWS::AccountId}:function:*
       Events:
         LambdaTrigger:
           Type: CloudWatchEvent
@@ -170,6 +178,50 @@ Resources:
         Ref: AWS::AccountId
       SourceArn:
         Fn::Sub: arn:aws:logs:${AWS::Region}:${AWS::AccountId}:log-group:*:*
+  SumoLogGroupExistingLambdaConnector:
+    Type: AWS::Serverless::Function
+    Condition: invoke_existing
+    Properties:
+      InlineCode: "var aws = require('aws-sdk');\nvar response = require('cfn-response');\n\
+        \nexports.handler = function(event, context) {\n    var lambda = new aws.Lambda();\n\
+        \    var payload = {\"existingLogs\": \"true\", \"token\": \"\"};\n    var\
+        \ responseStatus = \"FAILED\";\n    var responseData = {};\n    lambda.invoke(\n\
+        \        {\n          InvocationType: 'Event',\n          FunctionName: process.env.FUNCTION_NAME,\n\
+        \          Payload: JSON.stringify(payload),\n        }, function(err, invokeResult)\
+        \ {\n          if (err) {\n              responseData = {Error: \"Invoke call\
+        \ failed\"};\n              console.log(responseData.Error + \":\\n\", err);\n\
+        \          }\n          else {\n            responseStatus = \"SUCCESS\";\n\
+        \          }\n          response.send(event, context, responseStatus, responseData);\n\
+        \    });\n};\n"
+      Handler: index.handler
+      Runtime: nodejs12.x
+      Environment:
+        Variables:
+          FUNCTION_NAME:
+            Ref: SumoLogGroupLambdaConnector
+      Policies:
+      - Statement:
+        - Sid: InvokePolicy
+          Effect: Allow
+          Action:
+          - lambda:InvokeFunction
+          Resource:
+          - Fn::Sub: arn:aws:lambda:${AWS::Region}:${AWS::AccountId}:function:${SumoLogGroupLambdaConnector}
+  primerinvoke:
+    Type: AWS::CloudFormation::CustomResource
+    Version: '1.0'
+    Condition: invoke_existing
+    Properties:
+      ServiceToken:
+        Fn::GetAtt:
+        - SumoLogGroupExistingLambdaConnector
+        - Arn
+      DESTINATION_ARN:
+        Ref: DestinationArnValue
+      LOG_GROUP_PATTERN:
+        Ref: LogGroupPattern
+      ROLE_ARN:
+        Ref: RoleArn
 Outputs:
   SumoLogGroupLambdaConnector:
     Description: SumoLogGroupLambdaConnector Function ARN
diff --git a/loggroup-lambda-connector/sam/template.yaml b/loggroup-lambda-connector/sam/template.yaml
@@ -6,7 +6,7 @@ Description: >
 # More info about Globals: https://github.com/awslabs/serverless-application-model/blob/master/docs/globals.rst
 Globals:
   Function:
-    Timeout: 300
+    Timeout: 900
     MemorySize: 128
 
 Metadata:
@@ -69,6 +69,8 @@ Conditions:
   create_invoke_permission: !Equals [ !Ref DestinationArnType, 'Lambda' ]
   create_pass_role: !Equals [ !Ref DestinationArnType, 'Kinesis' ]
 
+  invoke_existing: !Equals [ !Ref UseExistingLogs, 'true' ]
+
 Rules:
   testRoleArnWithLambda:
     RuleCondition: !Equals
@@ -97,12 +99,11 @@ Resources:
     Properties:
       CodeUri: ../src/
       Handler: "loggroup-lambda-connector.handler"
-      Runtime: nodejs14.x
+      Runtime: nodejs12.x
       Environment:
         Variables:
           DESTINATION_ARN: !Ref "DestinationArnValue"
           LOG_GROUP_PATTERN: !Ref "LogGroupPattern"
-          USE_EXISTING_LOG_GROUPS: !Ref "UseExistingLogs"
           LOG_GROUP_TAGS: !Join [ ",", { "Ref": "LogGroupTags" } ]
           ROLE_ARN: !Ref "RoleArn"
       Policies:
@@ -115,6 +116,12 @@ Resources:
                 - logs:PutSubscriptionFilter
               Resource:
                 - !Sub 'arn:aws:logs:${AWS::Region}:${AWS::AccountId}:log-group:*'
+            - Sid: InvokePolicy
+              Effect: Allow
+              Action:
+                - lambda:InvokeFunction
+              Resource:
+                - !Sub 'arn:aws:lambda:${AWS::Region}:${AWS::AccountId}:function:*'
       Events:
         LambdaTrigger:
           Type: CloudWatchEvent
@@ -153,6 +160,59 @@ Resources:
       SourceAccount: !Ref AWS::AccountId
       SourceArn: !Sub 'arn:aws:logs:${AWS::Region}:${AWS::AccountId}:log-group:*:*'
 
+  SumoLogGroupExistingLambdaConnector:
+    Type: AWS::Serverless::Function
+    Condition: invoke_existing
+    Properties:
+      InlineCode: |
+        var aws = require('aws-sdk');
+        var response = require('cfn-response');
+
+        exports.handler = function(event, context) {
+            var lambda = new aws.Lambda();
+            var payload = {"existingLogs": "true", "token": ""};
+            var responseStatus = "FAILED";
+            var responseData = {};
+            lambda.invoke(
+                {
+                  InvocationType: 'Event',
+                  FunctionName: process.env.FUNCTION_NAME,
+                  Payload: JSON.stringify(payload),
+                }, function(err, invokeResult) {
+                  if (err) {
+                      responseData = {Error: "Invoke call failed"};
+                      console.log(responseData.Error + ":\n", err);
+                  }
+                  else {
+                    responseStatus = "SUCCESS";
+                  }
+                  response.send(event, context, responseStatus, responseData);
+            });
+        };
+      Handler: "index.handler"
+      Runtime: nodejs12.x
+      Environment:
+        Variables:
+          FUNCTION_NAME: !Ref SumoLogGroupLambdaConnector
+      Policies:
+        - Statement:
+            - Sid: InvokePolicy
+              Effect: Allow
+              Action:
+                - lambda:InvokeFunction
+              Resource:
+                - !Sub 'arn:aws:lambda:${AWS::Region}:${AWS::AccountId}:function:${SumoLogGroupLambdaConnector}'
+
+  primerinvoke:
+    Type: AWS::CloudFormation::CustomResource
+    Version: "1.0"
+    Condition: invoke_existing
+    Properties:
+      ServiceToken: !GetAtt SumoLogGroupExistingLambdaConnector.Arn
+      DESTINATION_ARN: !Ref "DestinationArnValue"
+      LOG_GROUP_PATTERN: !Ref "LogGroupPattern"
+      ROLE_ARN: !Ref "RoleArn"
+
 Outputs:
 
   SumoLogGroupLambdaConnector:
diff --git a/loggroup-lambda-connector/src/loggroup-lambda-connector.js b/loggroup-lambda-connector/src/loggroup-lambda-connector.js
@@ -1,11 +1,7 @@
 var AWS = require("aws-sdk");
 var cwl = new AWS.CloudWatchLogs({apiVersion: '2014-03-28'});
-async function sleep(waitTimeInMs) {
-    console.log("sleeping for " + waitTimeInMs + " ms...");
-    return new Promise((resolve) => setTimeout(resolve, waitTimeInMs));
-}
 
-function createSubscriptionFilter(lambdaLogGroupName, destinationArn, roleArn, errorHandler) {
+async function createSubscriptionFilter(lambdaLogGroupName, destinationArn, roleArn, errorHandler) {
     if (destinationArn.startsWith("arn:aws:lambda")){
         var params = {
             destinationArn: destinationArn,
@@ -16,15 +12,15 @@ function createSubscriptionFilter(lambdaLogGroupName, destinationArn, roleArn, e
     } else {
         var params = {
             destinationArn: destinationArn,
-            filterName: 'SumoLGKinesisFilter',
+            filterName: 'SumoLGLBDFilter',
             filterPattern: '',
             logGroupName: lambdaLogGroupName,
             roleArn: roleArn
         };
     }
 
     // handle case where subscription filter exists/case where loggroup generated by target lambda
-    cwl.putSubscriptionFilter(params, errorHandler);
+   await cwl.putSubscriptionFilter(params, errorHandler);
 }
 
 function filterLogGroups(event, logGroupRegex) {
@@ -58,11 +54,7 @@ async function subscribeExistingLogGroups(logGroups) {
     for (var i = logGroups.length - 1; i >= 0; i--) {
         logGroupName = logGroups[i].logGroupName;
         if (logGroupName.match(logGroupRegex)) {
-
-            // sleep time between calls
-            await sleep(1000*(process.env.SUBSCRIBE_DELAY_SECONDS || 2)); // 5 seconds
-
-            createSubscriptionFilter(logGroupName, destinationArn, roleArn, (function(inner_logGroupName) { return function (err, data) {
+            await createSubscriptionFilter(logGroupName, destinationArn, roleArn, (function(inner_logGroupName) { return function (err, data) {
                 if (err) {
                     console.log("Error in subscribing", inner_logGroupName, err);
                 } else {
@@ -75,29 +67,34 @@ async function subscribeExistingLogGroups(logGroups) {
     }
 }
 
-function processExistingLogGroups(token, errorHandler) {
-
-    var params = {
-      limit: 50,
-      // logGroupNamePrefix: '',
-      nextToken: token
-    };
+function processExistingLogGroups(token, context, errorHandler) {
+    var params = {limit: 50};
+    if (token) {
+        params = {
+          limit: 50,
+          // logGroupNamePrefix: '',
+          nextToken: token
+        };
+    }
     var p = new Promise(function(resolve, reject) {
         cwl.describeLogGroups(params, function(err, data) {
             if (err) {
                 console.log("error in fetching logGroups", err, err.stack);
                 reject(err);
             } else {
                 console.log("fetched logGroups: " + data.logGroups.length + " nextToken: " + data.nextToken);
+                subscribeExistingLogGroups(data.logGroups);
                 resolve(data);
             }
         });
     });
     var cb = function (data) {
-        subscribeExistingLogGroups(data.logGroups);
-        if (data.nextToken) {// if next set of log groups exists
-            processExistingLogGroups(data.nextToken, errorHandler)
+        if (data.nextToken) {// if next set of log groups exists, invoke next instance of lambda
+            console.log("Log Groups remaining...Calling the lambda again with token " + data.nextToken);
+            invoke_lambda(context, data.nextToken, errorHandler);
+            console.log("Lambda invoke complete with token " + data.nextToken);
         } else {
+            console.log("ALl Log Groups are subscribed to Destination Type " + process.env.DESTINATION_ARN);
             errorHandler(null, "Success");
         }
     };
@@ -106,6 +103,16 @@ function processExistingLogGroups(token, errorHandler) {
     });
 }
 
+function invoke_lambda(context, token, errorHandler) {
+    var lambda = new AWS.Lambda();
+    var payload = {"existingLogs": "true", "token": token};
+    lambda.invoke({
+        InvocationType: 'Event',
+        FunctionName: context.functionName,
+        Payload: JSON.stringify(payload),
+    }, errorHandler);
+}
+
 function processEvents(env, event, errorHandler) {
 
     var logGroupName = event.detail.requestParameters.logGroupName;
@@ -119,7 +126,7 @@ function processEvents(env, event, errorHandler) {
 }
 
 exports.handler = function (event, context, callback) {
-    console.log("Invoking Log Group connector function")
+    console.log("Invoking Log Group connector function");
     function errorHandler(err, msg) {
         if (err) {
             console.log(err, msg);
@@ -128,10 +135,9 @@ exports.handler = function (event, context, callback) {
             callback(null, "Success");
         }
     }
-    if (process.env.USE_EXISTING_LOG_GROUPS == "true") {
-        processExistingLogGroups(null, errorHandler);
+    if (event.existingLogs  == "true") {
+        processExistingLogGroups(event.token, context, errorHandler);
     } else {
         processEvents(process.env, event, errorHandler);
     }
-
 };
diff --git a/loggroup-lambda-connector/test/test_loggroup_lambda_connector.py b/loggroup-lambda-connector/test/test_loggroup_lambda_connector.py
@@ -54,23 +54,23 @@ def test_2_existing_logs(self):
         self.create_stack(self.stack_name, self.template_data, self.create_stack_parameters("Lambda", "true"))
         print("Testing Stack Creation")
         self.assertTrue(self.stack_exists(self.stack_name))
-        self.invoke_lambda()
+        #self.invoke_lambda()
         self.assert_subscription_filter("SumoLGLBDFilter")
 
     def test_3_kinesis(self):
         self.create_stack(self.stack_name, self.template_data, self.create_stack_parameters("Kinesis", "false"))
         print("Testing Stack Creation")
         self.assertTrue(self.stack_exists(self.stack_name))
         self.create_log_group()
-        self.assert_subscription_filter("SumoLGKinesisFilter")
+        self.assert_subscription_filter("SumoLGLBDFilter")
 
     def test_4_existing_kinesis(self):
         self.create_log_group()
         self.create_stack(self.stack_name, self.template_data, self.create_stack_parameters("Kinesis", "true"))
         print("Testing Stack Creation")
         self.assertTrue(self.stack_exists(self.stack_name))
-        self.invoke_lambda()
-        self.assert_subscription_filter("SumoLGKinesisFilter")
+        #self.invoke_lambda()
+        self.assert_subscription_filter("SumoLGLBDFilter")
 
     def create_stack_parameters(self, destination, existing, pattern='test'):
         return [
