diff --git a/README.md b/README.md index e172fd7..96770ed 100644 --- a/README.md +++ b/README.md @@ -192,6 +192,7 @@ newrelic-lambda layers install \ | `--nr-env-delimite` | No | Set `NR_ENV_DELIMITER` environment variable for your Lambda Function | | `--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`. | +| `--java-agent` | No | For Java runtimes only (`java17`, `java21`). Attaches the New Relic Java Agent layer (`NewRelicAgentJava`) instead of the default OpenTracing layer. Sets `AWS_LAMBDA_EXEC_WRAPPER=/opt/newrelic-java-handler` and leaves the function handler unchanged. Use `--java-agent true` to enable. | | `--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. | diff --git a/newrelic_lambda_cli/cli/layers.py b/newrelic_lambda_cli/cli/layers.py index 2ceb782..447b77b 100644 --- a/newrelic_lambda_cli/cli/layers.py +++ b/newrelic_lambda_cli/cli/layers.py @@ -183,6 +183,13 @@ def register(group): "log output in CloudWatch, reducing CloudWatch log volume without affecting " "telemetry delivery to New Relic", ) +@click.option( + "--java-agent", + "java_agent", + default=False, + type=bool, + help="Java runtimes only - Use New Relic Java Agent layer (sets AWS_LAMBDA_EXEC_WRAPPER, keeps original handler)", +) @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 f6c029b..e3e153d 100644 --- a/newrelic_lambda_cli/layers.py +++ b/newrelic_lambda_cli/layers.py @@ -17,6 +17,7 @@ NEW_RELIC_ENV_VARS = ( + "AWS_LAMBDA_EXEC_WRAPPER", "NEW_RELIC_ACCOUNT_ID", "NEW_RELIC_EXTENSION_LOGS_ENABLED", "NEW_RELIC_EXTENSION_SEND_EXTENSION_LOGS", @@ -57,10 +58,20 @@ def layer_selection( existing_layer_arn=None, slim=False, ): + layer_options = [ + layer["LatestMatchingVersion"]["LayerVersionArn"] for layer in available_layers + ] + + if slim: + for arn in layer_options: + if "-slim:" in arn: + success("Layer %s selected (slim)" % arn) + return arn + if upgrade and existing_layer_arn: base_arn = existing_layer_arn.rsplit(":", 1)[0] - for i, layer in enumerate(available_layers): + for layer in available_layers: candidate_arn = layer["LatestMatchingVersion"]["LayerVersionArn"] candidate_base_arn = candidate_arn.rsplit(":", 1)[0] if candidate_base_arn == base_arn: @@ -69,14 +80,6 @@ def layer_selection( if len(available_layers) == 1: return available_layers[0]["LatestMatchingVersion"]["LayerVersionArn"] - layer_options = [ - layer["LatestMatchingVersion"]["LayerVersionArn"] for layer in available_layers - ] - if slim: - for arn in layer_options: - if "-slim:" in arn: - success("Layer %s selected (slim)" % arn) - return arn if sys.stdout.isatty(): output = "\n".join( [ @@ -126,9 +129,14 @@ def _add_new_relic(input, config, nr_license_key): handler = config["Configuration"]["Handler"] runtime_handler = utils.RUNTIME_CONFIG.get(runtime, {}).get("Handler") + use_java_agent = input.java_agent and "java" in runtime + if "java" in runtime: - postfix = input.java_handler_method or "handleRequest" - runtime_handler = runtime_handler + postfix + if use_java_agent: + runtime_handler = None + else: + postfix = input.java_handler_method or "handleRequest" + runtime_handler = runtime_handler + postfix if "nodejs" in runtime: prefix = ( "/opt/nodejs/node_modules/newrelic-esm-lambda-wrapper/index" @@ -179,6 +187,20 @@ def _add_new_relic(input, config, nr_license_key): # discover compatible layers... available_layers = index(aws_region, runtime, architecture) + if "java" in runtime: + if use_java_agent: + available_layers = [ + l + for l in available_layers + if l.get("LayerName", "").lower().startswith("newrelicagent") + ] + else: + available_layers = [ + l + for l in available_layers + if not l.get("LayerName", "").lower().startswith("newrelicagent") + ] + if not available_layers: failure( "No Lambda layers published for %s (%s) runtime: %s" @@ -217,6 +239,18 @@ def _add_new_relic(input, config, nr_license_key): if runtime_handler: update_kwargs["Handler"] = runtime_handler + if use_java_agent: + update_kwargs["Environment"]["Variables"][ + "AWS_LAMBDA_EXEC_WRAPPER" + ] = "/opt/newrelic-java-handler" + original_handler = update_kwargs["Environment"]["Variables"].pop( + "NEW_RELIC_LAMBDA_HANDLER", None + ) + if original_handler: + update_kwargs["Handler"] = original_handler + elif "java" in runtime: + update_kwargs["Environment"]["Variables"].pop("AWS_LAMBDA_EXEC_WRAPPER", None) + # Update the account id update_kwargs["Environment"]["Variables"]["NEW_RELIC_ACCOUNT_ID"] = str( input.nr_account_id @@ -508,19 +542,29 @@ def _remove_new_relic(input, config): handler = config["Configuration"]["Handler"] - # For java runtimes we need to remove the method name before - # validating because method names are variable - if "java" in runtime: - handler = handler.split("::", 1)[0] + "::" + is_java_agent = ( + "java" in runtime + and config["Configuration"] + .get("Environment", {}) + .get("Variables", {}) + .get("AWS_LAMBDA_EXEC_WRAPPER") + == "/opt/newrelic-java-handler" + ) - # Detect non-New Relic handler and error if necessary. - if not utils.is_valid_handler(runtime, handler): - failure( - "New Relic installation (via layers) not auto-detected for the specified " - "function '%s'. Unrecognized handler in deployed function." - % config["Configuration"]["FunctionArn"] - ) - return False + if not is_java_agent: + # For java runtimes we need to remove the method name before + # validating because method names are variable + if "java" in runtime: + handler = handler.split("::", 1)[0] + "::" + + # Detect non-New Relic handler and error if necessary. + if not utils.is_valid_handler(runtime, handler): + failure( + "New Relic installation (via layers) not auto-detected for the specified " + "function '%s'. Unrecognized handler in deployed function." + % config["Configuration"]["FunctionArn"] + ) + return False env_handler = ( config["Configuration"] diff --git a/newrelic_lambda_cli/types.py b/newrelic_lambda_cli/types.py index f83445a..9c0de60 100644 --- a/newrelic_lambda_cli/types.py +++ b/newrelic_lambda_cli/types.py @@ -118,6 +118,7 @@ "send_platform_logs", "disable_platform_logs", "java_handler_method", + "java_agent", "esm", "slim", "extension_logs_enabled", diff --git a/tests/test_layers.py b/tests/test_layers.py index 8ec99b5..2ab267e 100644 --- a/tests/test_layers.py +++ b/tests/test_layers.py @@ -1787,6 +1787,273 @@ def test_extension_logs_not_set_by_default(aws_credentials, mock_function_config ) +@mock_aws +def test_add_new_relic_java_agent(aws_credentials, mock_function_config): + session = boto3.Session(region_name="us-east-1") + + with patch("newrelic_lambda_cli.layers.index") as mock_index, patch( + "newrelic_lambda_cli.layers.layer_selection" + ) as mock_layer_selection: + mock_index.return_value = [ + { + "LayerName": "NewRelicJava17", + "LatestMatchingVersion": { + "LayerVersionArn": "arn:aws:lambda:us-east-1:123456789:layer/NewRelicJava17:1", + "CompatibleArchitectures": ["x86_64"], + }, + }, + { + "LayerName": "NewRelicAgentJava", + "LatestMatchingVersion": { + "LayerVersionArn": "arn:aws:lambda:us-east-1:123456789:layer/NewRelicAgentJava:1", + "CompatibleArchitectures": ["x86_64"], + }, + }, + { + "LayerName": "NewRelicAgentJava-slim", + "LatestMatchingVersion": { + "LayerVersionArn": "arn:aws:lambda:us-east-1:123456789:layer/NewRelicAgentJava-slim:1", + "CompatibleArchitectures": ["x86_64"], + }, + }, + ] + mock_layer_selection.return_value = ( + "arn:aws:lambda:us-east-1:123456789:layer/NewRelicAgentJava:1" + ) + + config = mock_function_config("java17") + + update_kwargs = _add_new_relic( + layer_install( + session=session, + aws_region="us-east-1", + nr_account_id=12345, + enable_extension=True, + java_agent=True, + ), + config, + nr_license_key=None, + ) + + # Only agent layers passed to layer_selection + called_layers = mock_layer_selection.call_args[0][0] + assert all( + l.get("LayerName", "").lower().startswith("newrelicagent") + for l in called_layers + ) + # Handler unchanged — not in update_kwargs + assert "Handler" not in update_kwargs + # AWS_LAMBDA_EXEC_WRAPPER set + assert ( + update_kwargs["Environment"]["Variables"]["AWS_LAMBDA_EXEC_WRAPPER"] + == "/opt/newrelic-java-handler" + ) + # NEW_RELIC_LAMBDA_HANDLER not set + assert ( + "NEW_RELIC_LAMBDA_HANDLER" not in update_kwargs["Environment"]["Variables"] + ) + + +@mock_aws +def test_add_new_relic_java_ot_excludes_agent_layers( + aws_credentials, mock_function_config +): + session = boto3.Session(region_name="us-east-1") + + with patch("newrelic_lambda_cli.layers.index") as mock_index, patch( + "newrelic_lambda_cli.layers.layer_selection" + ) as mock_layer_selection: + mock_index.return_value = [ + { + "LayerName": "NewRelicJava17", + "LatestMatchingVersion": { + "LayerVersionArn": "arn:aws:lambda:us-east-1:123456789:layer/NewRelicJava17:1", + "CompatibleArchitectures": ["x86_64"], + }, + }, + { + "LayerName": "NewRelicAgentJava", + "LatestMatchingVersion": { + "LayerVersionArn": "arn:aws:lambda:us-east-1:123456789:layer/NewRelicAgentJava:1", + "CompatibleArchitectures": ["x86_64"], + }, + }, + ] + mock_layer_selection.return_value = ( + "arn:aws:lambda:us-east-1:123456789:layer/NewRelicJava17:1" + ) + + config = mock_function_config("java17") + + update_kwargs = _add_new_relic( + layer_install( + session=session, + aws_region="us-east-1", + nr_account_id=12345, + enable_extension=True, + ), + config, + nr_license_key=None, + ) + + # Only non-agent layers passed to layer_selection + called_layers = mock_layer_selection.call_args[0][0] + assert all( + not l.get("LayerName", "").lower().startswith("newrelicagent") + for l in called_layers + ) + # Handler set to OT wrapper + assert "HandlerWrapper" in update_kwargs["Handler"] + # NEW_RELIC_LAMBDA_HANDLER set + assert "NEW_RELIC_LAMBDA_HANDLER" in update_kwargs["Environment"]["Variables"] + # AWS_LAMBDA_EXEC_WRAPPER not set + assert ( + "AWS_LAMBDA_EXEC_WRAPPER" not in update_kwargs["Environment"]["Variables"] + ) + + +@mock_aws +def test_add_new_relic_java_ot_to_agent_switch(aws_credentials, mock_function_config): + session = boto3.Session(region_name="us-east-1") + + with patch("newrelic_lambda_cli.layers.index") as mock_index, patch( + "newrelic_lambda_cli.layers.layer_selection" + ) as mock_layer_selection: + mock_index.return_value = [ + { + "LayerName": "NewRelicAgentJava", + "LatestMatchingVersion": { + "LayerVersionArn": "arn:aws:lambda:us-east-1:123456789:layer/NewRelicAgentJava:1", + "CompatibleArchitectures": ["x86_64"], + }, + }, + ] + mock_layer_selection.return_value = ( + "arn:aws:lambda:us-east-1:123456789:layer/NewRelicAgentJava:1" + ) + + config = mock_function_config("java17") + config["Configuration"][ + "Handler" + ] = "com.newrelic.java.HandlerWrapper::handleRequest" + config["Configuration"]["Environment"]["Variables"][ + "NEW_RELIC_LAMBDA_HANDLER" + ] = "original_handler" + + update_kwargs = _add_new_relic( + layer_install( + session=session, + aws_region="us-east-1", + nr_account_id=12345, + enable_extension=True, + upgrade=True, + java_agent=True, + ), + config, + nr_license_key=None, + ) + + # Handler restored to original + assert update_kwargs["Handler"] == "original_handler" + # NEW_RELIC_LAMBDA_HANDLER removed + assert ( + "NEW_RELIC_LAMBDA_HANDLER" not in update_kwargs["Environment"]["Variables"] + ) + # AWS_LAMBDA_EXEC_WRAPPER set + assert ( + update_kwargs["Environment"]["Variables"]["AWS_LAMBDA_EXEC_WRAPPER"] + == "/opt/newrelic-java-handler" + ) + + +@mock_aws +def test_add_new_relic_java_agent_to_ot_switch(aws_credentials, mock_function_config): + session = boto3.Session(region_name="us-east-1") + + with patch("newrelic_lambda_cli.layers.index") as mock_index, patch( + "newrelic_lambda_cli.layers.layer_selection" + ) as mock_layer_selection: + mock_index.return_value = [ + { + "LayerName": "NewRelicJava17", + "LatestMatchingVersion": { + "LayerVersionArn": "arn:aws:lambda:us-east-1:123456789:layer/NewRelicJava17:1", + "CompatibleArchitectures": ["x86_64"], + }, + }, + ] + mock_layer_selection.return_value = ( + "arn:aws:lambda:us-east-1:123456789:layer/NewRelicJava17:1" + ) + + config = mock_function_config("java17") + config["Configuration"]["Environment"]["Variables"][ + "AWS_LAMBDA_EXEC_WRAPPER" + ] = "/opt/newrelic-java-handler" + + update_kwargs = _add_new_relic( + layer_install( + session=session, + aws_region="us-east-1", + nr_account_id=12345, + enable_extension=True, + upgrade=True, + ), + config, + nr_license_key=None, + ) + + # AWS_LAMBDA_EXEC_WRAPPER removed + assert ( + "AWS_LAMBDA_EXEC_WRAPPER" not in update_kwargs["Environment"]["Variables"] + ) + # Handler set to OT wrapper + assert "Handler" in update_kwargs + + +@mock_aws +def test_remove_new_relic_java_agent(aws_credentials, mock_function_config): + session = boto3.Session(region_name="us-east-1") + + config = mock_function_config("java17") + config["Configuration"]["Handler"] = "original_handler" + config["Configuration"]["Environment"]["Variables"][ + "AWS_LAMBDA_EXEC_WRAPPER" + ] = "/opt/newrelic-java-handler" + + update_kwargs = _remove_new_relic( + layer_uninstall(session=session, aws_region="us-east-1"), config + ) + + assert update_kwargs is not False + # AWS_LAMBDA_EXEC_WRAPPER removed + assert "AWS_LAMBDA_EXEC_WRAPPER" not in update_kwargs["Environment"]["Variables"] + # Handler unchanged + assert update_kwargs["Handler"] == "original_handler" + + +def test_layer_selection_slim(): + mock_layers = [ + { + "LayerName": "NewRelicAgentJava", + "LatestMatchingVersion": { + "LayerVersionArn": "arn:aws:lambda:us-east-1:123456789:layer/NewRelicAgentJava:1" + }, + }, + { + "LayerName": "NewRelicAgentJava-slim", + "LatestMatchingVersion": { + "LayerVersionArn": "arn:aws:lambda:us-east-1:123456789:layer/NewRelicAgentJava-slim:1" + }, + }, + ] + + selected = layer_selection(mock_layers, "java17", "x86_64", slim=True) + assert ( + selected == "arn:aws:lambda:us-east-1:123456789:layer/NewRelicAgentJava-slim:1" + ) + + @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"""