diff --git a/README.md b/README.md index 9e76a58..cfebf12 100644 --- a/README.md +++ b/README.md @@ -191,6 +191,8 @@ newrelic-lambda layers install \ | `--nr-tags` | No | Set `NR_TAGS` environment variable for your Lambda Function | | `--java_handler_method` or `-j` | No | For java runtimes only to specify an aws implementation method. Defaults to RequestHandler. Optional inputs are: handleRequest, handleStreamsRequest `--java_handler_method handleStreamsRequest`. | | `--esm` | No | For Node.js functions using ES Modules (ESM), enable the specific ESM wrapper during installation (e.g., using the --esm flag). This sets the Lambda handler to `/opt/nodejs/node_modules/newrelic-esm-lambda-wrapper/index.handler`. | +| `--extension-logs-enabled` | No | Set `NEW_RELIC_EXTENSION_LOGS_ENABLED=true` to enable `[NR_EXT]` extension log output in CloudWatch. This is the default extension behaviour.| +| `--extension-logs-disabled` | No | Set `NEW_RELIC_EXTENSION_LOGS_ENABLED=false` to suppress `[NR_EXT]` extension log output in CloudWatch. | #### Uninstall Layer diff --git a/newrelic_lambda_cli/cli/layers.py b/newrelic_lambda_cli/cli/layers.py index dbb37cc..1f56c6e 100644 --- a/newrelic_lambda_cli/cli/layers.py +++ b/newrelic_lambda_cli/cli/layers.py @@ -158,6 +158,21 @@ def register(group): is_flag=True, help="Nodejs runtimes only - nodejs implementation runtime handler to /opt/nodejs/node_modules/newrelic-esm-lambda-wrapper/index.handler", ) +@click.option( + "--extension-logs-enabled", + "extension_logs_enabled", + flag_value="true", + help="Set NEW_RELIC_EXTENSION_LOGS_ENABLED=true to enable [NR_EXT] extension " + "log output in CloudWatch (default extension behaviour)", +) +@click.option( + "--extension-logs-disabled", + "extension_logs_enabled", + flag_value="false", + help="Set NEW_RELIC_EXTENSION_LOGS_ENABLED=false to suppress [NR_EXT] extension " + "log output in CloudWatch, reducing CloudWatch log volume without affecting " + "telemetry delivery to New Relic", +) @click.pass_context def install(ctx, **kwargs): """Install New Relic AWS Lambda Layers""" diff --git a/newrelic_lambda_cli/layers.py b/newrelic_lambda_cli/layers.py index c6e21f5..e0ae81a 100644 --- a/newrelic_lambda_cli/layers.py +++ b/newrelic_lambda_cli/layers.py @@ -18,6 +18,7 @@ NEW_RELIC_ENV_VARS = ( "NEW_RELIC_ACCOUNT_ID", + "NEW_RELIC_EXTENSION_LOGS_ENABLED", "NEW_RELIC_EXTENSION_SEND_EXTENSION_LOGS", "NEW_RELIC_EXTENSION_SEND_FUNCTION_LOGS", "NEW_RELIC_LAMBDA_EXTENSION_ENABLED", @@ -274,6 +275,19 @@ def _add_new_relic(input, config, nr_license_key): "Successfully disabled NEW_RELIC_EXTENSION_SEND_EXTENSION_LOGS tag to the function" ) + if input.extension_logs_enabled is not None: + update_kwargs["Environment"]["Variables"][ + "NEW_RELIC_EXTENSION_LOGS_ENABLED" + ] = input.extension_logs_enabled + if input.extension_logs_enabled == "true": + success( + "Successfully enabled NEW_RELIC_EXTENSION_LOGS_ENABLED for the function" + ) + else: + success( + "Successfully disabled NEW_RELIC_EXTENSION_LOGS_ENABLED for the function" + ) + if input.nr_tags: update_kwargs["Environment"]["Variables"]["NR_TAGS"] = input.nr_tags success("Successfully added NR_TAGS tag to the function") diff --git a/newrelic_lambda_cli/types.py b/newrelic_lambda_cli/types.py index 934c43b..50aee6d 100644 --- a/newrelic_lambda_cli/types.py +++ b/newrelic_lambda_cli/types.py @@ -118,6 +118,7 @@ "java_handler_method", "esm", "slim", + "extension_logs_enabled", ] LAYER_UNINSTALL_KEYS = [ diff --git a/tests/test_layers.py b/tests/test_layers.py index d2a8ea3..c23d024 100644 --- a/tests/test_layers.py +++ b/tests/test_layers.py @@ -1525,3 +1525,163 @@ def test_install_account_mismatch_error_with_missing_keys( # This should raise UsageError with pytest.raises(UsageError): mock_install(input_obj, "foobarbaz") + + +@mock_aws +def test_extension_logs_enabled_flag(aws_credentials, mock_function_config): + """Test that --extension-logs-enabled sets NEW_RELIC_EXTENSION_LOGS_ENABLED=true""" + session = boto3.Session(region_name="us-east-1") + nr_account_id = 12345 + + # Test 1: Fresh install with --extension-logs-enabled + config = mock_function_config("python3.12") + update_kwargs = _add_new_relic( + layer_install( + session=session, + aws_region="us-east-1", + nr_account_id=nr_account_id, + enable_extension=True, + extension_logs_enabled="true", + ), + config, + nr_license_key=None, + ) + + assert ( + update_kwargs["Environment"]["Variables"]["NEW_RELIC_EXTENSION_LOGS_ENABLED"] + == "true" + ) + + # Test 2: Upgrade with --extension-logs-enabled + config = mock_function_config("python3.12") + update_kwargs = _add_new_relic( + layer_install( + session=session, + aws_region="us-east-1", + nr_account_id=nr_account_id, + enable_extension=True, + extension_logs_enabled="true", + upgrade=True, + ), + config, + nr_license_key=None, + ) + + assert ( + update_kwargs["Environment"]["Variables"]["NEW_RELIC_EXTENSION_LOGS_ENABLED"] + == "true" + ) + + +@mock_aws +def test_extension_logs_disabled_flag(aws_credentials, mock_function_config): + """Test that --extension-logs-disabled sets NEW_RELIC_EXTENSION_LOGS_ENABLED=false""" + session = boto3.Session(region_name="us-east-1") + nr_account_id = 12345 + + # Test 1: Fresh install with --extension-logs-disabled + config = mock_function_config("python3.12") + update_kwargs = _add_new_relic( + layer_install( + session=session, + aws_region="us-east-1", + nr_account_id=nr_account_id, + enable_extension=True, + extension_logs_enabled="false", + ), + config, + nr_license_key=None, + ) + + assert ( + update_kwargs["Environment"]["Variables"]["NEW_RELIC_EXTENSION_LOGS_ENABLED"] + == "false" + ) + + # Test 2: Upgrade with --extension-logs-disabled + config = mock_function_config("python3.12") + update_kwargs = _add_new_relic( + layer_install( + session=session, + aws_region="us-east-1", + nr_account_id=nr_account_id, + enable_extension=True, + extension_logs_enabled="false", + upgrade=True, + ), + config, + nr_license_key=None, + ) + + assert ( + update_kwargs["Environment"]["Variables"]["NEW_RELIC_EXTENSION_LOGS_ENABLED"] + == "false" + ) + + +@mock_aws +def test_extension_logs_not_set_by_default(aws_credentials, mock_function_config): + """Test that NEW_RELIC_EXTENSION_LOGS_ENABLED is NOT set when neither flag is passed""" + session = boto3.Session(region_name="us-east-1") + nr_account_id = 12345 + + # Fresh install without either --extension-logs-enabled or --extension-logs-disabled + config = mock_function_config("python3.12") + update_kwargs = _add_new_relic( + layer_install( + session=session, + aws_region="us-east-1", + nr_account_id=nr_account_id, + enable_extension=True, + ), + config, + nr_license_key=None, + ) + + assert ( + "NEW_RELIC_EXTENSION_LOGS_ENABLED" + not in update_kwargs["Environment"]["Variables"] + ) + + # Upgrade without either flag - should also not set the var + config = mock_function_config("python3.12") + update_kwargs = _add_new_relic( + layer_install( + session=session, + aws_region="us-east-1", + nr_account_id=nr_account_id, + enable_extension=True, + upgrade=True, + ), + config, + nr_license_key=None, + ) + + assert ( + "NEW_RELIC_EXTENSION_LOGS_ENABLED" + not in update_kwargs["Environment"]["Variables"] + ) + + +@mock_aws +def test_extension_logs_removed_on_uninstall(aws_credentials, mock_function_config): + """Test that NEW_RELIC_EXTENSION_LOGS_ENABLED is removed during uninstall""" + session = boto3.Session(region_name="us-east-1") + + config = mock_function_config("python3.12") + config["Configuration"]["Handler"] = "newrelic_lambda_wrapper.handler" + config["Configuration"]["Environment"]["Variables"][ + "NEW_RELIC_LAMBDA_HANDLER" + ] = "original_handler" + config["Configuration"]["Environment"]["Variables"][ + "NEW_RELIC_EXTENSION_LOGS_ENABLED" + ] = "true" + + update_kwargs = _remove_new_relic( + layer_uninstall(session=session, aws_region="us-east-1"), config + ) + + assert ( + "NEW_RELIC_EXTENSION_LOGS_ENABLED" + not in update_kwargs["Environment"]["Variables"] + )