Lambda: allow <fn_name>:<alias_name> (#6977)

moto/awslambda/models.py

@@ -1354,8 +1354,10 @@ class LambdaStorage(object):
         return None
 
     def _get_function_aliases(self, function_name: str) -> Dict[str, LambdaAlias]:
-        fn = self.get_function_by_name_or_arn(function_name)
-        return self._aliases[fn.function_arn]
+        fn = self.get_function_by_name_or_arn_with_qualifier(function_name)
+        # Split ARN to retrieve an ARN without a qualifier present
+        [arn, _, _] = self.split_function_arn(fn.function_arn)
+        return self._aliases[arn]
 
     def delete_alias(self, name: str, function_name: str) -> None:
         aliases = self._get_function_aliases(function_name)
@@ -1377,7 +1379,9 @@ class LambdaStorage(object):
         description: str,
         routing_config: str,
     ) -> LambdaAlias:
-        fn = self.get_function_by_name_or_arn(function_name, function_version)
+        fn = self.get_function_by_name_or_arn_with_qualifier(
+            function_name, function_version
+        )
         aliases = self._get_function_aliases(function_name)
         if name in aliases:
             arn = f"arn:aws:lambda:{self.region_name}:{self.account_id}:function:{function_name}:{name}"
@@ -1406,34 +1410,71 @@ class LambdaStorage(object):
         alias = self.get_alias(name, function_name)
 
         # errors if new function version doesn't exist
-        self.get_function_by_name_or_arn(function_name, function_version)
+        self.get_function_by_name_or_arn_with_qualifier(function_name, function_version)
 
         alias.update(description, function_version, routing_config)
         return alias
 
-    def get_function_by_name(
-        self, name: str, qualifier: Optional[str] = None
-    ) -> Optional[LambdaFunction]:
+    def get_function_by_name_forbid_qualifier(self, name: str) -> LambdaFunction:
+        """
+        Get function by name forbidding a qualifier
+        :raises: UnknownFunctionException if function not found
+        :raises: InvalidParameterValue if qualifier is provided
+        """
+
+        if name.count(":") == 1:
+            raise InvalidParameterValueException("Cannot provide qualifier")
+
         if name not in self._functions:
-            return None
+            raise self.construct_unknown_function_exception(name)
 
+        return self._get_latest(name)
+
+    def get_function_by_name_with_qualifier(
+        self, name: str, qualifier: Optional[str] = None
+    ) -> LambdaFunction:
+        """
+        Get function by name with an optional qualifier
+        :raises: UnknownFunctionException if function not found
+        """
+
+        # Function name may contain an alias
+        # <fn_name>:<alias_name>
+        if ":" in name:
+            # Prefer qualifier in name over qualifier arg
+            [name, qualifier] = name.split(":")
+
+        # Find without qualifier
         if qualifier is None:
-            return self._get_latest(name)
+            return self.get_function_by_name_forbid_qualifier(name)
 
+        if name not in self._functions:
+            raise self.construct_unknown_function_exception(name, qualifier)
+
+        # Find by latest
         if qualifier.lower() == "$latest":
             return self._functions[name]["latest"]
 
+        # Find by version
         found_version = self._get_version(name, qualifier)
         if found_version:
             return found_version
 
+        # Find by alias
         aliases = self._get_function_aliases(name)
-
         if qualifier in aliases:
-            alias = aliases[qualifier]
-            return self._get_version(name, alias.function_version)
+            alias_version = aliases[qualifier].function_version
 
-        return None
+            # Find by alias pointing to latest
+            if alias_version.lower() == "$latest":
+                return self._functions[name]["latest"]
+
+            # Find by alias pointing to version
+            found_alias = self._get_version(name, alias_version)
+            if found_alias:
+                return found_alias
+
+        raise self.construct_unknown_function_exception(name, qualifier)
 
     def list_versions_by_function(self, name: str) -> Iterable[LambdaFunction]:
         if name not in self._functions:
@@ -1448,28 +1489,74 @@ class LambdaStorage(object):
         return sorted(aliases.values(), key=lambda alias: alias.name)
 
     def get_arn(self, arn: str) -> Optional[LambdaFunction]:
+        [arn_without_qualifier, _, _] = self.split_function_arn(arn)
+        return self._arns.get(arn_without_qualifier, None)
+
+    def split_function_arn(self, arn: str) -> Tuple[str, str, Optional[str]]:
+        """
+        Handy utility to parse an ARN into:
+        - ARN without qualifier
+        - Function name
+        - Optional qualifier
+        """
+        qualifier = None
         # Function ARN may contain an alias
         # arn:aws:lambda:region:account_id:function:<fn_name>:<alias_name>
         if ":" in arn.split(":function:")[-1]:
+            qualifier = arn.split(":")[-1]
             # arn = arn:aws:lambda:region:account_id:function:<fn_name>
             arn = ":".join(arn.split(":")[0:-1])
-        return self._arns.get(arn, None)
+        name = arn.split(":")[-1]
+        return arn, name, qualifier
 
-    def get_function_by_name_or_arn(
+    def get_function_by_name_or_arn_forbid_qualifier(
+        self, name_or_arn: str
+    ) -> LambdaFunction:
+        """
+        Get function by name or arn forbidding a qualifier
+        :raises: UnknownFunctionException if function not found
+        :raises: InvalidParameterValue if qualifier is provided
+        """
+
+        if name_or_arn.startswith("arn:aws"):
+            [_, name, qualifier] = self.split_function_arn(name_or_arn)
+
+            if qualifier is not None:
+                raise InvalidParameterValueException("Cannot provide qualifier")
+
+            return self.get_function_by_name_forbid_qualifier(name)
+        else:
+            # name_or_arn is not an arn
+            return self.get_function_by_name_forbid_qualifier(name_or_arn)
+
+    def get_function_by_name_or_arn_with_qualifier(
         self, name_or_arn: str, qualifier: Optional[str] = None
     ) -> LambdaFunction:
-        fn = self.get_function_by_name(name_or_arn, qualifier) or self.get_arn(
-            name_or_arn
-        )
-        if fn is None:
-            if name_or_arn.startswith("arn:aws"):
-                arn = name_or_arn
-            else:
-                arn = make_function_arn(self.region_name, self.account_id, name_or_arn)
-            if qualifier:
+        """
+        Get function by name or arn with an optional qualifier
+        :raises: UnknownFunctionException if function not found
+        """
+
+        if name_or_arn.startswith("arn:aws"):
+            [_, name, qualifier_in_arn] = self.split_function_arn(name_or_arn)
+            return self.get_function_by_name_with_qualifier(
+                name, qualifier_in_arn or qualifier
+            )
+        else:
+            return self.get_function_by_name_with_qualifier(name_or_arn, qualifier)
+
+    def construct_unknown_function_exception(
+        self, name_or_arn: str, qualifier: Optional[str] = None
+    ) -> UnknownFunctionException:
+        if name_or_arn.startswith("arn:aws"):
+            arn = name_or_arn
+        else:
+            # name_or_arn is a function name with optional qualifier <func_name>[:<qualifier>]
+            arn = make_function_arn(self.region_name, self.account_id, name_or_arn)
+            # Append explicit qualifier to arn only if the name doesn't already have it
+            if qualifier and ":" not in name_or_arn:
                 arn = f"{arn}:{qualifier}"
-            raise UnknownFunctionException(arn)
-        return fn
+        return UnknownFunctionException(arn)
 
     def put_function(self, fn: LambdaFunction) -> None:
         valid_role = re.match(InvalidRoleFormat.pattern, fn.role)
@@ -1497,7 +1584,7 @@ class LambdaStorage(object):
     def publish_function(
         self, name_or_arn: str, description: str = ""
     ) -> Optional[LambdaFunction]:
-        function = self.get_function_by_name_or_arn(name_or_arn)
+        function = self.get_function_by_name_or_arn_forbid_qualifier(name_or_arn)
         name = function.function_name
         if name not in self._functions:
             return None
@@ -1522,7 +1609,19 @@ class LambdaStorage(object):
         return fn
 
     def del_function(self, name_or_arn: str, qualifier: Optional[str] = None) -> None:
-        function = self.get_function_by_name_or_arn(name_or_arn, qualifier)
+        # Qualifier may be explicitly passed or part of function name or ARN, extract it here
+        if name_or_arn.startswith("arn:aws"):
+            # Extract from ARN
+            if ":" in name_or_arn.split(":function:")[-1]:
+                qualifier = name_or_arn.split(":")[-1]
+        else:
+            # Extract from function name
+            if ":" in name_or_arn:
+                qualifier = name_or_arn.split(":")[1]
+
+        function = self.get_function_by_name_or_arn_with_qualifier(
+            name_or_arn, qualifier
+        )
         name = function.function_name
         if not qualifier:
             # Something is still reffing this so delete all arns
@@ -1534,8 +1633,11 @@ class LambdaStorage(object):
 
             del self._functions[name]
 
-        elif qualifier == "$LATEST":
-            self._functions[name]["latest"] = None
+        else:
+            if qualifier == "$LATEST":
+                self._functions[name]["latest"] = None
+            else:
+                self._functions[name]["versions"].remove(function)
 
             # If theres no functions left
             if (
@@ -1544,18 +1646,6 @@ class LambdaStorage(object):
             ):
                 del self._functions[name]
 
-        else:
-            fn = self.get_function_by_name(name, qualifier)
-            if fn:
-                self._functions[name]["versions"].remove(fn)
-
-                # If theres no functions left
-                if (
-                    not self._functions[name]["versions"]
-                    and not self._functions[name]["latest"]
-                ):
-                    del self._functions[name]
-
         self._aliases[function.function_arn] = {}
 
     def all(self) -> Iterable[LambdaFunction]:
@@ -1783,21 +1873,27 @@ class LambdaBackend(BaseBackend):
         The Qualifier-parameter is not yet implemented.
         Function URLs are not yet mocked, so invoking them will fail
         """
-        function = self._lambdas.get_function_by_name_or_arn(name_or_arn)
+        function = self._lambdas.get_function_by_name_or_arn_forbid_qualifier(
+            name_or_arn
+        )
         return function.create_url_config(config)
 
     def delete_function_url_config(self, name_or_arn: str) -> None:
         """
         The Qualifier-parameter is not yet implemented
         """
-        function = self._lambdas.get_function_by_name_or_arn(name_or_arn)
+        function = self._lambdas.get_function_by_name_or_arn_forbid_qualifier(
+            name_or_arn
+        )
         function.delete_url_config()
 
     def get_function_url_config(self, name_or_arn: str) -> FunctionUrlConfig:
         """
         The Qualifier-parameter is not yet implemented
         """
-        function = self._lambdas.get_function_by_name_or_arn(name_or_arn)
+        function = self._lambdas.get_function_by_name_or_arn_forbid_qualifier(
+            name_or_arn
+        )
         if not function:
             raise UnknownFunctionException(arn=name_or_arn)
         return function.get_url_config()
@@ -1808,7 +1904,9 @@ class LambdaBackend(BaseBackend):
         """
         The Qualifier-parameter is not yet implemented
         """
-        function = self._lambdas.get_function_by_name_or_arn(name_or_arn)
+        function = self._lambdas.get_function_by_name_or_arn_forbid_qualifier(
+            name_or_arn
+        )
         return function.update_url_config(config)
 
     def create_event_source_mapping(self, spec: Dict[str, Any]) -> EventSourceMapping:
@@ -1818,9 +1916,9 @@ class LambdaBackend(BaseBackend):
                 raise RESTError("InvalidParameterValueException", f"Missing {param}")
 
         # Validate function name
-        func = self._lambdas.get_function_by_name_or_arn(spec.get("FunctionName", ""))
-        if not func:
-            raise RESTError("ResourceNotFoundException", "Invalid FunctionName")
+        func = self._lambdas.get_function_by_name_or_arn_with_qualifier(
+            spec.get("FunctionName", "")
+        )
 
         # Validate queue
         sqs_backend = sqs_backends[self.account_id][self.region_name]
@@ -1887,7 +1985,7 @@ class LambdaBackend(BaseBackend):
     def get_function(
         self, function_name_or_arn: str, qualifier: Optional[str] = None
     ) -> LambdaFunction:
-        return self._lambdas.get_function_by_name_or_arn(
+        return self._lambdas.get_function_by_name_or_arn_with_qualifier(
             function_name_or_arn, qualifier
         )
 
@@ -1912,7 +2010,9 @@ class LambdaBackend(BaseBackend):
 
         for key in spec.keys():
             if key == "FunctionName":
-                func = self._lambdas.get_function_by_name_or_arn(spec[key])
+                func = self._lambdas.get_function_by_name_or_arn_with_qualifier(
+                    spec[key]
+                )
                 esm.function_arn = func.function_arn
             elif key == "BatchSize":
                 esm.batch_size = spec[key]
@@ -2025,7 +2125,9 @@ class LambdaBackend(BaseBackend):
                 }
             ]
         }
-        func = self._lambdas.get_function_by_name_or_arn(function_name, qualifier)
+        func = self._lambdas.get_function_by_name_or_arn_with_qualifier(
+            function_name, qualifier
+        )
         func.invoke(json.dumps(event), {}, {})
 
     def send_dynamodb_items(
@@ -2076,14 +2178,14 @@ class LambdaBackend(BaseBackend):
         func.invoke(json.dumps(event), {}, {})  # type: ignore[union-attr]
 
     def list_tags(self, resource: str) -> Dict[str, str]:
-        return self._lambdas.get_function_by_name_or_arn(resource).tags
+        return self._lambdas.get_function_by_name_or_arn_with_qualifier(resource).tags
 
     def tag_resource(self, resource: str, tags: Dict[str, str]) -> None:
-        fn = self._lambdas.get_function_by_name_or_arn(resource)
+        fn = self._lambdas.get_function_by_name_or_arn_with_qualifier(resource)
         fn.tags.update(tags)
 
     def untag_resource(self, resource: str, tagKeys: List[str]) -> None:
-        fn = self._lambdas.get_function_by_name_or_arn(resource)
+        fn = self._lambdas.get_function_by_name_or_arn_with_qualifier(resource)
         for key in tagKeys:
             fn.tags.pop(key, None)
 
@@ -2104,9 +2206,9 @@ class LambdaBackend(BaseBackend):
         return fn.get_code_signing_config()
 
     def get_policy(self, function_name: str, qualifier: Optional[str] = None) -> str:
-        fn = self._lambdas.get_function_by_name_or_arn(function_name, qualifier)
-        if not fn:
-            raise UnknownFunctionException(function_name)
+        fn = self._lambdas.get_function_by_name_or_arn_with_qualifier(
+            function_name, qualifier
+        )
         return fn.policy.wire_format()  # type: ignore[union-attr]
 
     def update_function_code(
@@ -2125,7 +2227,7 @@ class LambdaBackend(BaseBackend):
     ) -> Optional[Dict[str, Any]]:
         fn = self.get_function(function_name, qualifier)
 
-        return fn.update_configuration(body) if fn else None
+        return fn.update_configuration(body)
 
     def invoke(
         self,
@@ -2139,12 +2241,9 @@ class LambdaBackend(BaseBackend):
         Invoking a Function with PackageType=Image is not yet supported.
         """
         fn = self.get_function(function_name, qualifier)
-        if fn:
-            payload = fn.invoke(body, headers, response_headers)
-            response_headers["Content-Length"] = str(len(payload))
-            return payload
-        else:
-            return None
+        payload = fn.invoke(body, headers, response_headers)
+        response_headers["Content-Length"] = str(len(payload))
+        return payload
 
     def put_function_concurrency(
         self, function_name: str, reserved_concurrency: str

moto/awslambda/responses.py

@@ -398,7 +398,17 @@ class LambdaResponse(BaseResponse):
         return 204, {}, ""
 
     @staticmethod
-    def _set_configuration_qualifier(configuration: Dict[str, Any], qualifier: str) -> Dict[str, Any]:  # type: ignore[misc]
+    def _set_configuration_qualifier(configuration: Dict[str, Any], function_name: str, qualifier: str) -> Dict[str, Any]:  # type: ignore[misc]
+        # Qualifier may be explicitly passed or part of function name or ARN, extract it here
+        if function_name.startswith("arn:aws"):
+            # Extract from ARN
+            if ":" in function_name.split(":function:")[-1]:
+                qualifier = function_name.split(":")[-1]
+        else:
+            # Extract from function name
+            if ":" in function_name:
+                qualifier = function_name.split(":")[1]
+
         if qualifier is None or qualifier == "$LATEST":
             configuration["Version"] = "$LATEST"
         if qualifier == "$LATEST":
@@ -413,7 +423,7 @@ class LambdaResponse(BaseResponse):
 
         code = fn.get_code()
         code["Configuration"] = self._set_configuration_qualifier(
-            code["Configuration"], qualifier
+            code["Configuration"], function_name, qualifier
         )
         return 200, {}, json.dumps(code)
 
@@ -424,7 +434,7 @@ class LambdaResponse(BaseResponse):
         fn = self.backend.get_function(function_name, qualifier)
 
         configuration = self._set_configuration_qualifier(
-            fn.get_configuration(), qualifier
+            fn.get_configuration(), function_name, qualifier
         )
         return 200, {}, json.dumps(configuration)
 

tests/test_awslambda/test_lambda.py

@@ -591,6 +591,22 @@ def test_get_function():
         == f"arn:aws:lambda:us-west-2:{ACCOUNT_ID}:function:{function_name}:$LATEST"
     )
 
+    # Test get function with version
+    result = conn.get_function(FunctionName=function_name, Qualifier="1")
+    assert result["Configuration"]["Version"] == "1"
+    assert (
+        result["Configuration"]["FunctionArn"]
+        == f"arn:aws:lambda:us-west-2:{ACCOUNT_ID}:function:{function_name}:1"
+    )
+
+    # Test get function with version inside of name
+    result = conn.get_function(FunctionName=f"{function_name}:1")
+    assert result["Configuration"]["Version"] == "1"
+    assert (
+        result["Configuration"]["FunctionArn"]
+        == f"arn:aws:lambda:us-west-2:{ACCOUNT_ID}:function:{function_name}:1"
+    )
+
     # Test get function when can't find function name
     with pytest.raises(conn.exceptions.ResourceNotFoundException):
         conn.get_function(FunctionName="junk", Qualifier="$LATEST")
@@ -724,6 +740,22 @@ def test_get_function_by_arn():
     result = conn.get_function(FunctionName=fnc["FunctionArn"])
     assert result["Configuration"]["FunctionName"] == function_name
 
+    # Test with version
+    result = conn.get_function(FunctionName=fnc["FunctionArn"], Qualifier="1")
+    assert (
+        result["Configuration"]["FunctionArn"]
+        == f"arn:aws:lambda:us-east-1:{ACCOUNT_ID}:function:{function_name}:1"
+    )
+    assert result["Configuration"]["Version"] == "1"
+
+    # Test with version inside of ARN
+    result = conn.get_function(FunctionName=f"{fnc['FunctionArn']}:1")
+    assert result["Configuration"]["Version"] == "1"
+    assert (
+        result["Configuration"]["FunctionArn"]
+        == f"arn:aws:lambda:us-east-1:{ACCOUNT_ID}:function:{function_name}:1"
+    )
+
 
 @mock_lambda
 @mock_s3
@@ -760,7 +792,7 @@ def test_delete_function():
 
     assert success_result == {"ResponseMetadata": {"HTTPStatusCode": 204}}
 
-    func_list = conn.list_functions()["Functions"]
+    func_list = conn.list_functions(FunctionVersion="ALL")["Functions"]
     our_functions = [f for f in func_list if f["FunctionName"] == function_name]
     assert len(our_functions) == 0
 
@@ -1592,12 +1624,23 @@ def test_multiple_qualifiers():
     qualis = [fn["FunctionArn"].split(":")[-1] for fn in resp]
     assert qualis == ["$LATEST", "1", "2", "3", "4", "5", "6", "7", "8", "9", "10"]
 
+    # Test delete with function name and qualifier
     client.delete_function(FunctionName=fn_name, Qualifier="4")
-    client.delete_function(FunctionName=fn_name, Qualifier="5")
+    # Test delete with ARN and qualifier
+    client.delete_function(
+        FunctionName=f"arn:aws:lambda:us-east-1:{ACCOUNT_ID}:function:{fn_name}",
+        Qualifier="5",
+    )
+    # Test delete with qualifier part of function name
+    client.delete_function(FunctionName=fn_name + ":8")
+    # Test delete with qualifier inside ARN
+    client.delete_function(
+        FunctionName=f"arn:aws:lambda:us-east-1:{ACCOUNT_ID}:function:{fn_name}:9"
+    )
 
     resp = client.list_versions_by_function(FunctionName=fn_name)["Versions"]
     qualis = [fn["FunctionArn"].split(":")[-1] for fn in resp]
-    assert qualis == ["$LATEST", "1", "2", "3", "6", "7", "8", "9", "10"]
+    assert qualis == ["$LATEST", "1", "2", "3", "6", "7", "10"]
 
     fn = client.get_function(FunctionName=fn_name, Qualifier="6")["Configuration"]
     assert (
@@ -1606,6 +1649,18 @@ def test_multiple_qualifiers():
     )
 
 
+@mock_lambda
+def test_delete_non_existent():
+    client = boto3.client("lambda", "us-east-1")
+
+    with pytest.raises(ClientError) as exc:
+        client.delete_function(
+            FunctionName=f"arn:aws:lambda:us-east-1:{ACCOUNT_ID}:function:nonexistent:9"
+        )
+
+    assert exc.value.response["Error"]["Code"] == "ResourceNotFoundException"
+
+
 def test_get_role_name_utility_race_condition():
     # Play with these variables as needed to reproduce the error.
     max_workers, num_threads = 3, 15

tests/test_awslambda/test_lambda_alias.py

@@ -394,9 +394,11 @@ def test_update_alias_routingconfig():
 
 
 @mock_lambda
-def test_get_function_using_alias():
+@pytest.mark.parametrize("qualifierIn", ["NAME", "SEPARATE", "BOTH"])
+def test_get_function_using_alias(qualifierIn):
     client = boto3.client("lambda", region_name="us-east-2")
     fn_name = str(uuid4())[0:6]
+    fn_qualifier = "live"
 
     client.create_function(
         FunctionName=fn_name,
@@ -408,9 +410,20 @@ def test_get_function_using_alias():
     client.publish_version(FunctionName=fn_name)
     client.publish_version(FunctionName=fn_name)
 
-    client.create_alias(FunctionName=fn_name, Name="live", FunctionVersion="1")
+    client.create_alias(FunctionName=fn_name, Name=fn_qualifier, FunctionVersion="1")
 
-    fn = client.get_function(FunctionName=fn_name, Qualifier="live")["Configuration"]
+    if qualifierIn == "NAME":
+        fn = client.get_function(FunctionName=f"{fn_name}:{fn_qualifier}")[
+            "Configuration"
+        ]
+    elif qualifierIn == "SEPARATE":
+        fn = client.get_function(FunctionName=fn_name, Qualifier=fn_qualifier)[
+            "Configuration"
+        ]
+    elif qualifierIn == "BOTH":
+        fn = client.get_function(
+            FunctionName=f"{fn_name}:{fn_qualifier}", Qualifier=fn_qualifier
+        )["Configuration"]
     assert (
         fn["FunctionArn"]
         == f"arn:aws:lambda:us-east-2:{ACCOUNT_ID}:function:{fn_name}:1"