add --apm flag to enable APM Lambda mode (#327)

* add --apm flag to enable APM Lambda mode

Author: chaudharysaket

newrelic_lambda_cli/cli/layers.py

@@ -81,6 +81,11 @@ def register(group):
     help="Permit upgrade of function layers to new version.",
     is_flag=True,
 )
+@click.option(
+    "--apm",
+    help="Enable APM Lambda mode",
+    is_flag=True,
+)
 @click.option(
     "--enable-extension/--disable-extension",
     "-x",

newrelic_lambda_cli/layers.py

@@ -118,6 +118,7 @@ def _add_new_relic(input, config, nr_license_key):
         success(
             "Already installed on function '%s'. Pass --upgrade (or -u) to allow "
             "upgrade or reinstall to latest layer version."
+            "Additionally pass --apm to enable APM Lambda mode."
             % config["Configuration"]["FunctionArn"]
         )
         return True
@@ -232,6 +233,13 @@ def _add_new_relic(input, config, nr_license_key):
             "CORECLR_PROFILER_PATH"
         ] = "/opt/lib/newrelic-dotnet-agent/libNewRelicProfiler.so"
 
+    if input.apm:
+        success(
+            "Enabling APM Lambda mode for function '%s' "
+            % config["Configuration"]["FunctionArn"]
+        )
+        update_kwargs["Environment"]["Variables"]["NEW_RELIC_APM_LAMBDA_MODE"] = "True"
+
     return update_kwargs
 
 

newrelic_lambda_cli/types.py

@@ -104,6 +104,7 @@
     "excludes",
     "layer_arn",
     "upgrade",
+    "apm",
     "enable_extension",
     "enable_extension_function_logs",
     "java_handler_method",

tests/test_layers.py

@@ -235,6 +235,34 @@ def test_add_new_relic(aws_credentials, mock_function_config):
     assert update_kwargs["Layers"][0] != get_arn_prefix("us-east-1")
 
 
+@mock_aws
+def test_add_new_relic_apm_lambda_mode(aws_credentials, mock_function_config):
+    session = boto3.Session(region_name="us-east-1")
+
+    config = mock_function_config("python3.12")
+
+    assert config["Configuration"]["Handler"] == "original_handler"
+
+    update_kwargs = _add_new_relic(
+        layer_install(
+            session=session,
+            aws_region="us-east-1",
+            nr_account_id=12345,
+            enable_extension=True,
+            enable_extension_function_logs=True,
+            apm=True,
+        ),
+        config,
+        nr_license_key=None,
+    )
+
+    assert update_kwargs["FunctionName"] == config["Configuration"]["FunctionArn"]
+    assert update_kwargs["Handler"] == "newrelic_lambda_wrapper.handler"
+    assert (
+        update_kwargs["Environment"]["Variables"]["NEW_RELIC_APM_LAMBDA_MODE"] == "True"
+    )
+
+
 @mock_aws
 def test_add_new_relic_dotnet(aws_credentials, mock_function_config):
     session = boto3.Session(region_name="us-east-1")