Support multiple --function arguments for subscription commands (#68)

Author: kolanos

README.md

@@ -170,7 +170,7 @@ newrelic-lambda subscriptions install \--function <name or arn>
 
 | Option | Required? | Description |
 |--------|-----------|-------------|
-| `--function` or `-f` | Yes | The AWS Lambda function name or ARN in which to add a log subscription. |
+| `--function` or `-f` | Yes | The AWS Lambda function name or ARN in which to add a log subscription. Can provide multiple `--function` arguments. |
 | `--aws-profile` or `-p` | No | The AWS profile to use for this command. Can also use `AWS_PROFILE`. Will also check `AWS_ACCESS_KEY_ID` and `AWS_SECRET_ACCESS_KEY` environment variables if not using AWS CLI. |
 | `--aws-region` or `-r` | No | The AWS region this function is located. Can use `AWS_DEFAULT_REGION` environment variable. Defaults to AWS session region. |
 
@@ -182,7 +182,7 @@ newrelic-lambda subscriptions uninstall --function <name or arn>
 
 | Option | Required? | Description |
 |--------|-----------|-------------|
-| `--function` or `-f` | Yes | The AWS Lambda function name or ARN in which to remove a log subscription. |
+| `--function` or `-f` | Yes | The AWS Lambda function name or ARN in which to remove a log subscription. Can provide multiple `--function` arguments. |
 | `--aws-profile` or `-p` | No | The AWS profile to use for this command. Can also use `AWS_PROFILE`. Will also check `AWS_ACCESS_KEY_ID` and `AWS_SECRET_ACCESS_KEY` environment variables if not using AWS CLI. |
 | `--aws-region` or `-r` | No | The AWS region this function is located. Can use `AWS_DEFAULT_REGION` environment variable. Defaults to AWS session region. |
 

newrelic_lambda_cli/cli/layers.py

@@ -39,14 +39,12 @@ def register(group):
     metavar="<arn>",
     multiple=True,
     required=True,
-    type=click.STRING,
 )
 @click.option(
     "--layer-arn",
     "-l",
     help="ARN for New Relic layer (default: auto-detect)",
     metavar="<arn>",
-    type=click.STRING,
 )
 @click.option(
     "--upgrade",
@@ -65,7 +63,7 @@ def install(
     layer_arn,
     upgrade,
 ):
-    """Install New Relic AWS Lambda Layer"""
+    """Install New Relic AWS Lambda Layers"""
     session = boto3.Session(profile_name=aws_profile, region_name=aws_region)
 
     if aws_permissions_check:
@@ -77,7 +75,7 @@ def install(
         res = layers.install(session, function, layer_arn, nr_account_id, upgrade)
         install_success = res and install_success
         if res:
-            success("Successfully installed layer on %s" % res["FunctionArn"])
+            success("Successfully installed layer on %s" % function)
             if ctx.obj["VERBOSE"]:
                 click.echo(json.dumps(res, indent=2))
 
@@ -97,11 +95,10 @@ def install(
     metavar="<arn>",
     multiple=True,
     required=True,
-    type=click.STRING,
 )
 @click.pass_context
 def uninstall(ctx, aws_profile, aws_region, aws_permissions_check, functions):
-    """Uninstall New Relic AWS Lambda Layer"""
+    """Uninstall New Relic AWS Lambda Layers"""
     session = boto3.Session(profile_name=aws_profile, region_name=aws_region)
 
     if aws_permissions_check:
@@ -113,7 +110,7 @@ def uninstall(ctx, aws_profile, aws_region, aws_permissions_check, functions):
         res = layers.uninstall(session, function)
         uninstall_success = res and uninstall_success
         if res:
-            success("Successfully uninstalled layer on %s" % res["FunctionArn"])
+            success("Successfully uninstalled layer on %s" % function)
             if ctx.obj["VERBOSE"]:
                 click.echo(json.dumps(res, indent=2))
 

newrelic_lambda_cli/cli/subscriptions.py

@@ -2,7 +2,7 @@
 import click
 
 from newrelic_lambda_cli import permissions, subscriptions
-from newrelic_lambda_cli.cliutils import done
+from newrelic_lambda_cli.cliutils import done, failure, success
 from newrelic_lambda_cli.cli.decorators import add_options, AWS_OPTIONS
 
 
@@ -21,38 +21,62 @@ def register(group):
 @click.command(name="install")
 @add_options(AWS_OPTIONS)
 @click.option(
+    "functions",
     "--function",
     "-f",
     help="AWS Lambda function name or ARN",
     metavar="<arn>",
+    multiple=True,
     required=True,
 )
-def install(aws_profile, aws_region, aws_permissions_check, function):
-    """Install New Relic AWS Lambda Log Subscription"""
+def install(aws_profile, aws_region, aws_permissions_check, functions):
+    """Install New Relic AWS Lambda Log Subscriptions"""
     session = boto3.Session(profile_name=aws_profile, region_name=aws_region)
 
     if aws_permissions_check:
         permissions.ensure_subscription_install_permissions(session)
 
-    subscriptions.create_log_subscription(session, function)
-    done("Install Complete")
+    install_success = True
+
+    for function in functions:
+        result = subscriptions.create_log_subscription(session, function)
+        install_success = result and install_success
+        if result:
+            success("Successfully installed log subscription on %s" % function)
+
+    if install_success:
+        done("Install Complete")
+    else:
+        failure("Install Incomplete. See messages above for details.", exit=True)
 
 
 @click.command(name="uninstall")
 @add_options(AWS_OPTIONS)
 @click.option(
+    "functions",
     "--function",
     "-f",
     help="Lambda function name or ARN",
     metavar="<arn>",
+    multiple=True,
     required=True,
 )
-def uninstall(aws_profile, aws_region, aws_permissions_check, function):
-    """Uninstall New Relic AWS Lambda Log Subscription"""
+def uninstall(aws_profile, aws_region, aws_permissions_check, functions):
+    """Uninstall New Relic AWS Lambda Log Subscriptions"""
     session = boto3.Session(profile_name=aws_profile, region_name=aws_region)
 
     if aws_permissions_check:
         permissions.ensure_subscription_uninstall_permissions(session)
 
-    subscriptions.remove_log_subscription(session, function)
-    done("Uninstall Complete")
+    uninstall_success = True
+
+    for function in functions:
+        result = subscriptions.remove_log_subscription(session, function)
+        uninstall_success = result and uninstall_success
+        if result:
+            success("Successfully uninstalled log subscription on %s" % function)
+
+    if uninstall_success:
+        done("Uninstall Complete")
+    else:
+        failure("Uninstall Incomplete. See messages above for details.", exit=True)

newrelic_lambda_cli/subscriptions.py

@@ -22,31 +22,44 @@ def get_subscription_filters(session, function_name):
             and e.response["ResponseMetadata"]["HTTPStatusCode"] == 404
         ):
             return []
-        raise click.UsageError(str(e))
+        failure(
+            "Error retrieving log subscription filters for '%s': %s"
+            % (function_name, e)
+        )
     else:
         return res.get("subscriptionFilters", [])
 
 
 def create_subscription_filter(session, function_name, destination_arn):
     try:
-        return session.client("logs").put_subscription_filter(
+        session.client("logs").put_subscription_filter(
             logGroupName="/aws/lambda/%s" % function_name,
             filterName="NewRelicLogStreaming",
             filterPattern=DEFAULT_FILTER_PATTERN,
             destinationArn=destination_arn,
         )
     except botocore.exceptions.ClientError as e:
-        raise click.UsageError(str(e))
+        failure(
+            "Error creating log subscription filter for '%s': %s" % (function_name, e)
+        )
+        return False
+    else:
+        return True
 
 
 def remove_subscription_filter(session, function_name):
     try:
-        return session.client("logs").delete_subscription_filter(
+        session.client("logs").delete_subscription_filter(
             logGroupName="/aws/lambda/%s" % function_name,
             filterName="NewRelicLogStreaming",
         )
     except botocore.exceptions.ClientError as e:
-        raise click.UsageError(str(e))
+        failure(
+            "Error removing log subscription filter for '%s': %s" % (function_name, e)
+        )
+        return False
+    else:
+        return True
 
 
 def create_log_subscription(session, function_name):
@@ -56,11 +69,11 @@ def create_log_subscription(session, function_name):
             "Could not find 'newrelic-log-ingestion' function. Is the New Relic AWS "
             "integration installed?"
         )
-        return
+        return False
     destination_arn = destination["Configuration"]["FunctionArn"]
-    subscription_filters = [
-        filter for filter in get_subscription_filters(session, function_name)
-    ]
+    subscription_filters = get_subscription_filters(session, function_name)
+    if subscription_filters is None:
+        return False
     newrelic_filters = [
         filter
         for filter in subscription_filters
@@ -79,27 +92,31 @@ def create_log_subscription(session, function_name):
         )
     if not newrelic_filters:
         click.echo("Adding New Relic log subscription to '%s'" % function_name)
-        create_subscription_filter(session, function_name, destination_arn)
+        return create_subscription_filter(session, function_name, destination_arn)
     else:
         click.echo(
             "Found log subscription for '%s', verifying configuration" % function_name
         )
         newrelic_filter = newrelic_filters[0]
-        if newrelic_filter["filterPattern"] == "":
-            remove_subscription_filter(session, function_name)
-            create_subscription_filter(session, function_name, destination_arn)
+        if newrelic_filter["filterPattern"] != DEFAULT_FILTER_PATTERN:
+            return remove_subscription_filter(
+                session, function_name
+            ) and create_subscription_filter(session, function_name, destination_arn)
 
 
 def remove_log_subscription(session, function_name):
-    subscription_filters = [
+    subscription_filters = get_subscription_filters(session, function_name)
+    if subscription_filters is None:
+        return False
+    newrelic_filters = [
         filter
-        for filter in get_subscription_filters(session, function_name)
+        for filter in subscription_filters
         if filter["filterName"] == "NewRelicLogStreaming"
     ]
-    if not subscription_filters:
+    if not newrelic_filters:
         click.echo(
             "No New Relic subscription filters found for '%s', skipping" % function_name
         )
-        return
+        return False
     click.echo("Removing New Relic log subscription from '%s'" % function_name)
-    remove_subscription_filter(session, function_name)
+    return remove_subscription_filter(session, function_name)

tests/cli/test_subscriptions.py

@@ -0,0 +1,90 @@
+from moto import mock_lambda, mock_logs
+
+from newrelic_lambda_cli.cli import cli, register_groups
+
+
+@mock_lambda
+@mock_logs
+def test_subscriptions_install(aws_credentials, cli_runner):
+    """
+    Assert that 'newrelic-lambda subscriptions install' attempts to install the
+    New Relic log subscription on a function.
+    """
+    register_groups(cli)
+
+    result = cli_runner.invoke(
+        cli,
+        [
+            "subscriptions",
+            "install",
+            "--no-aws-permissions-check",
+            "--function",
+            "foobar",
+        ],
+        env={"AWS_DEFAULT_REGION": "us-east-1"},
+    )
+
+    assert result.exit_code == 1
+    assert result.stdout == ""
+    assert "✖️ Install Incomplete. See messages above for details." in result.stderr
+
+    result2 = cli_runner.invoke(
+        cli,
+        [
+            "subscriptions",
+            "install",
+            "--no-aws-permissions-check",
+            "--function",
+            "foobar",
+            "--function",
+            "barbaz",
+        ],
+        env={"AWS_DEFAULT_REGION": "us-east-1"},
+    )
+
+    assert result2.exit_code == 1
+    assert result2.stdout == ""
+    assert "✖️ Install Incomplete. See messages above for details." in result2.stderr
+
+
+@mock_lambda
+def test_subscriptions_uninstall(aws_credentials, cli_runner):
+    """
+    Assert that 'newrelic-lambda subscriptions uninstall' attempts to uninstall the
+    New Relic log subscription on a function.
+    """
+    register_groups(cli)
+
+    result = cli_runner.invoke(
+        cli,
+        [
+            "subscriptions",
+            "uninstall",
+            "--no-aws-permissions-check",
+            "--function",
+            "foobar",
+        ],
+        env={"AWS_DEFAULT_REGION": "us-east-1"},
+    )
+
+    assert result.exit_code == 1
+    assert result.stdout == ""
+    assert "✖️ Uninstall Incomplete. See messages above for details." in result.stderr
+
+    result2 = cli_runner.invoke(
+        cli,
+        [
+            "subscriptions",
+            "uninstall",
+            "--no-aws-permissions-check",
+            "--function",
+            "foobar",
+            "--function",
+            "barbaz",
+        ],
+        env={"AWS_DEFAULT_REGION": "us-east-1"},
+    )
+
+    assert result2.exit_code == 1
+    assert result2.stdout == ""
+    assert "✖️ Uninstall Incomplete. See messages above for details." in result2.stderr