Add filter alias support to layer/subscription commands (#70)

* Add alias support to layer/subscription commands

Closes #69

* Add test for get_aliased_functions

Author: kolanos

README.md

@@ -119,7 +119,8 @@ newrelic-lambda layers install \
 
 | Option | Required? | Description |
 |--------|-----------|-------------|
-| `--function` or `-f` | Yes | The AWS Lambda function name or ARN in which to add a layer. Can provide multiple `--function` arguments. |
+| `--function` or `-f` | Yes | The AWS Lambda function name or ARN in which to add a layer. Can provide multiple `--function` arguments. Will also accept `all`, `installed` and `not-installed` similar to `newrelic-lambda functions list`. |
+| `--exclude` or `-e` | No | A function name to exclude while installing layers. Can provide multiple `--exclude` arguments. Only checked when `all`, `installed` and `not-installed` are used. See `newrelic-lambda functions list` for function names. |
 | `--nr-account-id` or `-a` | Yes | The [New Relic Account ID](https://docs.newrelic.com/docs/accounts/install-new-relic/account-setup/account-id) this function should use. Can also use the `NEW_RELIC_ACCOUNT_ID` environment variable. |
 | `--layer-arn` or `-l` | No | Specify a specific layer version ARN to use. This is auto detected by default. |
 | `--upgrade` or `-u` | No | Permit upgrade to the latest layer version for this region and runtime. |
@@ -134,7 +135,8 @@ newrelic-lambda layers uninstall --function <name or arn>
 
 | Option | Required? | Description |
 |--------|-----------|-------------|
-| `--function` or `-f` | Yes | The AWS Lambda function name or ARN in which to remove a layer. Can provide multiple `--function` arguments. |
+| `--function` or `-f` | Yes | The AWS Lambda function name or ARN in which to remove a layer. Can provide multiple `--function` arguments. Will also accept `all`, `installed` and `not-installed` similar to `newrelic-lambda functions list`. |
+| `--exclude` or `-e` | No | A function name to exclude while uninstalling layers. Can provide multiple `--exclude` arguments. Only checked when `all`, `installed` and `not-installed` are used. See `newrelic-lambda functions list` for function names. |
 | `--layer-arn` or `-l` | No | Specify a specific layer version ARN to remove. This is auto detected by default. |
 | `--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. |
@@ -170,7 +172,8 @@ 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. Can provide multiple `--function` arguments. |
+| `--function` or `-f` | Yes | The AWS Lambda function name or ARN in which to add a log subscription. Can provide multiple `--function` arguments. Will also accept `all`, `installed` and `not-installed` similar to `newrelic-lambda functions list`. |
+| `--exclude` or `-e` | No | A function name to exclude while installing subscriptions. Can provide multiple `--exclude` arguments. Only checked when `all`, `installed` and `not-installed` are used. See `newrelic-lambda functions list` for function names. |
 | `--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 +185,8 @@ 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. Can provide multiple `--function` arguments. |
+| `--function` or `-f` | Yes | The AWS Lambda function name or ARN in which to remove a log subscription. Can provide multiple `--function` arguments. Will also accept `all`, `installed` and `not-installed` similar to `newrelic-lambda functions list`. |
+| `--exclude` or `-e` | No | A function name to exclude while uninstalling subscriptions. Can provide multiple `--exclude` arguments. Only checked when `all`, `installed` and `not-installed` are used. See `newrelic-lambda functions list` for function names. |
 | `--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/functions.py

@@ -32,8 +32,10 @@ def list(aws_profile, aws_region, aws_permissions_check, filter):
     """List AWS Lambda Functions"""
     _, rows = shutil.get_terminal_size((80, 50))
     session = boto3.Session(profile_name=aws_profile, region_name=aws_region)
+
     if aws_permissions_check:
         permissions.ensure_lambda_list_permissions(session)
+
     funcs = functions.list_functions(session, filter)
 
     def _format(funcs, header=False):
@@ -59,4 +61,5 @@ def _format(funcs, header=False):
             )
             buffer = []
             return
+
     click.echo(_format(buffer, header=True))

newrelic_lambda_cli/cli/layers.py

@@ -6,6 +6,7 @@
 from newrelic_lambda_cli import layers, permissions
 from newrelic_lambda_cli.cliutils import done, failure, success
 from newrelic_lambda_cli.cli.decorators import add_options, AWS_OPTIONS
+from newrelic_lambda_cli.functions import get_aliased_functions
 
 
 @click.group(name="layers")
@@ -40,6 +41,14 @@ def register(group):
     multiple=True,
     required=True,
 )
+@click.option(
+    "excludes",
+    "--exclude",
+    "-e",
+    help="Functions to exclude (if using 'all, 'installed', 'not-installed aliases)",
+    metavar="<name>",
+    multiple=True,
+)
 @click.option(
     "--layer-arn",
     "-l",
@@ -60,6 +69,7 @@ def install(
     aws_region,
     aws_permissions_check,
     functions,
+    excludes,
     layer_arn,
     upgrade,
 ):
@@ -69,6 +79,8 @@ def install(
     if aws_permissions_check:
         permissions.ensure_lambda_install_permissions(session)
 
+    functions = get_aliased_functions(session, functions, excludes)
+
     install_success = True
 
     for function in functions:
@@ -96,14 +108,24 @@ def install(
     multiple=True,
     required=True,
 )
+@click.option(
+    "excludes",
+    "--exclude",
+    "-e",
+    help="Functions to exclude (if using 'all, 'installed', 'not-installed aliases)",
+    metavar="<name>",
+    multiple=True,
+)
 @click.pass_context
-def uninstall(ctx, aws_profile, aws_region, aws_permissions_check, functions):
+def uninstall(ctx, aws_profile, aws_region, aws_permissions_check, functions, excludes):
     """Uninstall New Relic AWS Lambda Layers"""
     session = boto3.Session(profile_name=aws_profile, region_name=aws_region)
 
     if aws_permissions_check:
         permissions.ensure_lambda_uninstall_permissions(session)
 
+    functions = get_aliased_functions(session, functions, excludes)
+
     uninstall_success = True
 
     for function in functions:

newrelic_lambda_cli/cli/subscriptions.py

@@ -4,6 +4,7 @@
 from newrelic_lambda_cli import permissions, subscriptions
 from newrelic_lambda_cli.cliutils import done, failure, success
 from newrelic_lambda_cli.cli.decorators import add_options, AWS_OPTIONS
+from newrelic_lambda_cli.functions import get_aliased_functions
 
 
 @click.group(name="subscriptions")
@@ -29,13 +30,23 @@ def register(group):
     multiple=True,
     required=True,
 )
-def install(aws_profile, aws_region, aws_permissions_check, functions):
+@click.option(
+    "excludes",
+    "--exclude",
+    "-e",
+    help="Functions to exclude (if using 'all, 'installed', 'not-installed aliases)",
+    metavar="<name>",
+    multiple=True,
+)
+def install(aws_profile, aws_region, aws_permissions_check, functions, excludes):
     """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)
 
+    functions = get_aliased_functions(session, functions, excludes)
+
     install_success = True
 
     for function in functions:
@@ -61,13 +72,23 @@ def install(aws_profile, aws_region, aws_permissions_check, functions):
     multiple=True,
     required=True,
 )
-def uninstall(aws_profile, aws_region, aws_permissions_check, functions):
+@click.option(
+    "excludes",
+    "--exclude",
+    "-e",
+    help="Functions to exclude (if using 'all, 'installed', 'not-installed aliases)",
+    metavar="<name>",
+    multiple=True,
+)
+def uninstall(aws_profile, aws_region, aws_permissions_check, functions, excludes):
     """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)
 
+    functions = get_aliased_functions(session, functions, excludes)
+
     uninstall_success = True
 
     for function in functions:

newrelic_lambda_cli/functions.py

@@ -4,29 +4,27 @@
 from newrelic_lambda_cli import utils
 
 
-def list_functions(session, filter_choice):
+def list_functions(session, filter=None):
     client = session.client("lambda")
 
-    # set all if the filter_choice is "all" or there is no filter_choice active.
-    all = filter_choice == "all" or not filter_choice
+    all = filter == "all" or not filter
 
     pager = client.get_paginator("list_functions")
-    for func_resp in pager.paginate():
-        funcs = func_resp.get("Functions", [])
-
-        for f in funcs:
-            f.setdefault("x-new-relic-enabled", False)
-            for layer in f.get("Layers", []):
+    for res in pager.paginate():
+        funcs = res.get("Functions", [])
+        for func in funcs:
+            func.setdefault("x-new-relic-enabled", False)
+            for layer in func.get("Layers", []):
                 if layer.get("Arn", "").startswith(
                     utils.get_arn_prefix(session.region_name)
                 ):
-                    f["x-new-relic-enabled"] = True
+                    func["x-new-relic-enabled"] = True
             if all:
-                yield f
-            elif filter_choice == "installed" and f["x-new-relic-enabled"]:
-                yield f
-            elif filter_choice == "not_installed" and not f["x-new-relic-enabled"]:
-                yield f
+                yield func
+            elif filter == "installed" and func["x-new-relic-enabled"]:
+                yield func
+            elif filter == "not-installed" and not func["x-new-relic-enabled"]:
+                yield func
 
 
 def get_function(session, function_name):
@@ -42,3 +40,37 @@ def get_function(session, function_name):
         ):
             return None
         raise click.UsageError(str(e))
+
+
+def get_aliased_functions(session, functions, excludes):
+    """
+    Retrieves functions for 'all, 'installed' and 'not-installed' aliases and appends
+    them to existing list of functions.
+    """
+    aliases = [
+        function.lower()
+        for function in functions
+        if function.lower() in ("all", "installed", "not-installed")
+    ]
+
+    functions = [
+        function
+        for function in functions
+        if function.lower()
+        not in ("all", "installed", "not-installed", "newrelic-log-ingestion")
+        and function not in excludes
+    ]
+
+    if not aliases:
+        return utils.unique(functions)
+
+    for alias in set(aliases):
+        for function in list_functions(session, alias):
+            if (
+                "FunctionName" in function
+                and "newrelic-log-ingestion" not in function["FunctionName"]
+                and function["FunctionName"] not in excludes
+            ):
+                functions.append(function["FunctionName"])
+
+    return utils.unique(functions)

newrelic_lambda_cli/layers.py

@@ -3,7 +3,7 @@
 import requests
 
 from newrelic_lambda_cli import utils
-from newrelic_lambda_cli.cliutils import failure
+from newrelic_lambda_cli.cliutils import failure, success
 from newrelic_lambda_cli.functions import get_function
 
 
@@ -22,11 +22,12 @@ def _add_new_relic(config, region, layer_arn, account_id, allow_upgrade):
             "Unsupported Lambda runtime for '%s': %s"
             % (config["Configuration"]["FunctionArn"], runtime)
         )
+        return
 
     handler = config["Configuration"]["Handler"]
     runtime_handler = utils.RUNTIME_CONFIG.get(runtime, {}).get("Handler")
     if not allow_upgrade and handler == runtime_handler:
-        failure(
+        success(
             "Already installed on function '%s'. Pass --upgrade (or -u) to allow "
             "upgrade or reinstall to latest layer version."
             % config["Configuration"]["FunctionArn"]
@@ -46,6 +47,13 @@ def _add_new_relic(config, region, layer_arn, account_id, allow_upgrade):
         # discover compatible layers...
         available_layers = index(region, runtime)
 
+        if not available_layers:
+            failure(
+                "No Lambda layers published for '%s' runtime: %s"
+                % (config["Configuration"]["FunctionArn"], runtime)
+            )
+            return
+
         # TODO: MAke this a layer selection screen
         if len(available_layers) > 1:
             message = ["Discovered layers for runtime (%s)" % runtime]

newrelic_lambda_cli/utils.py

@@ -82,3 +82,13 @@ def validate_aws_profile(ctx, param, value):
         raise click.BadParameter(e.fmt)
     else:
         return value
+
+
+def unique(seq):
+    """Returns unique values in a sequence while preserving order"""
+    seen = set()
+    # Why assign seen.add to seen_add instead of just calling seen.add?
+    # Python is a dynamic language, and resolving seen.add each iteration is more costly
+    # than resolving a local variable.
+    seen_add = seen.add
+    return [x for x in seq if not (x in seen or seen_add(x))]

tests/test_functions.py

@@ -0,0 +1,39 @@
+import boto3
+import mock
+from moto import mock_lambda
+from newrelic_lambda_cli.functions import get_aliased_functions
+
+
+@mock_lambda
+@mock.patch("newrelic_lambda_cli.functions.list_functions", autospec=True)
+def test_get_aliased_functions(mock_list_functions, aws_credentials):
+    """
+    Asserts that get_aliased_functions adds functions matching one of the alias filters
+    """
+    session = boto3.Session(region_name="us-east-1")
+
+    assert get_aliased_functions(session, [], []) == []
+    assert get_aliased_functions(session, ["foo"], []) == ["foo"]
+    assert get_aliased_functions(session, ["foo", "bar"], ["bar"]) == ["foo"]
+    assert get_aliased_functions(session, ["foo", "bar", "baz"], ["bar"]) == [
+        "foo",
+        "baz",
+    ]
+
+    mock_list_functions.return_value = [{"FunctionName": "aliased-func"}]
+    assert get_aliased_functions(session, ["foo", "bar", "all"], []) == [
+        "foo",
+        "bar",
+        "aliased-func",
+    ]
+
+    mock_list_functions.return_value = [
+        {"FunctionName": "aliased-func"},
+        {"FunctionName": "ignored-func"},
+        {"FunctionName": "newrelic-log-ingestion"},
+    ]
+    assert get_aliased_functions(session, ["foo", "bar", "all"], ["ignored-func"]) == [
+        "foo",
+        "bar",
+        "aliased-func",
+    ]