LakeFormation: extend functionality of list_permissions (#7051)

moto/lakeformation/models.py

@@ -123,7 +123,14 @@ class ListPermissionsResource:
         lf_tag: Optional[ListPermissionsResourceLFTag],
         lf_tag_policy: Optional[ListPermissionsResourceLFTagPolicy],
     ):
-        if catalog is None and database is None and table is None:
+        if (
+            catalog is None
+            and database is None
+            and table is None
+            and data_location is None
+        ):
+            # Error message is the exact string returned by the AWS-CLI eventhough it is valid
+            # to not populate the respective fields as long as data_location is given.
             raise InvalidInput(
                 "Resource must have either the catalog, table or database field populated."
             )
@@ -276,10 +283,11 @@ class LakeFormationBackend(BaseBackend):
                 only matching permissions with resource-type "Database" are returned;
             if catalog and database are not provided and table is provided:
                 only matching permissions with resource-type "Table" are returned;
+            if catalog and database and table are not provided and data location is provided:
+                only matching permissions with resource-type "DataLocation" are returned;
             """
             if resource is None:  # Check for linter
                 return False
-
             permission_resource = permission["Resource"]
             catalog = resource.catalog
             if catalog is not None and "Catalog" in permission_resource:
@@ -291,7 +299,7 @@ class LakeFormationBackend(BaseBackend):
                 if database.catalog_id is not None:
                     equals = equals and (
                         database.catalog_id
-                        == permission_resource["Database"]["CatalogId"]
+                        == permission_resource["Database"].get("CatalogId")
                     )
                 return equals
 
@@ -302,7 +310,8 @@ class LakeFormationBackend(BaseBackend):
                 )
                 if table.catalog_id is not None:
                     equals = equals and (
-                        table.catalog_id == permission_resource["Table"]["CatalogId"]
+                        table.catalog_id
+                        == permission_resource["Table"].get("CatalogId")
                     )
                 if table.name is not None and table.table_wildcard is None:
                     equals = equals and (
@@ -314,6 +323,20 @@ class LakeFormationBackend(BaseBackend):
                         == permission_resource["Table"]["TableWildcard"]
                     )
                 return equals
+
+            data_location = resource.data_location
+            if data_location is not None and "DataLocation" in permission_resource:
+                equals = (
+                    data_location.resource_arn
+                    == permission_resource["DataLocation"]["ResourceArn"]
+                )
+                if data_location.catalog_id is not None:
+                    equals = equals and (
+                        data_location.catalog_id
+                        == permission_resource["DataLocation"].get("CatalogId")
+                    )
+                return equals
+
             return False
 
         if resource is not None:

moto/lakeformation/responses.py

@@ -8,10 +8,13 @@ from .models import (
     LakeFormationBackend,
     ListPermissionsResource,
     ListPermissionsResourceDatabase,
+    ListPermissionsResourceDataLocation,
     ListPermissionsResourceTable,
     RessourceType,
 )
 
+from .exceptions import InvalidInput
+
 
 class LakeFormationResponse(BaseResponse):
     """Handler for LakeFormation requests and responses."""
@@ -97,6 +100,12 @@ class LakeFormationResponse(BaseResponse):
         principal = self._get_param("Principal")
         resource = self._get_param("Resource")
         resource_type_param = self._get_param("ResourceType")
+        if principal is not None and resource is None:
+            # Error message is the exact string returned by the AWS-CLI
+            raise InvalidInput(
+                "An error occurred (InvalidInputException) when calling the ListPermissions operation: Resource is mandatory if Principal is set in the input."
+            )
+
         if resource_type_param is None:
             resource_type = None
         else:
@@ -108,6 +117,7 @@ class LakeFormationResponse(BaseResponse):
             database_sub_dictionary = resource.get("Database")
             table_sub_dictionary = resource.get("Table")
             catalog_sub_dictionary = resource.get("Catalog")
+            data_location_sub_dictionary = resource.get("DataLocation")
 
             if database_sub_dictionary is None:
                 database = None
@@ -127,12 +137,20 @@ class LakeFormationResponse(BaseResponse):
                     table_wildcard=table_sub_dictionary.get("TableWildcard"),
                 )
 
+            if data_location_sub_dictionary is None:
+                data_location = None
+            else:
+                data_location = ListPermissionsResourceDataLocation(
+                    resource_arn=data_location_sub_dictionary.get("ResourceArn"),
+                    catalog_id=data_location_sub_dictionary.get("CatalogId"),
+                )
+
             list_permission_resource = ListPermissionsResource(
                 catalog=catalog_sub_dictionary,
                 database=database,
                 table=table,
                 table_with_columns=None,
-                data_location=None,
+                data_location=data_location,
                 data_cells_filter=None,
                 lf_tag=None,
                 lf_tag_policy=None,

tests/test_lakeformation/test_lakeformation.py

@@ -1,5 +1,5 @@
 """Unit tests for lakeformation-supported APIs."""
-from typing import Dict
+from typing import Dict, Optional
 import boto3
 import pytest
 
@@ -118,6 +118,24 @@ def test_list_permissions():
     ]
 
 
+@mock_lakeformation
+def test_list_permissions_invalid_input():
+    client = boto3.client("lakeformation", region_name="eu-west-2")
+
+    client.grant_permissions(
+        Principal={"DataLakePrincipalIdentifier": "asdf"},
+        Resource={"Database": {"Name": "db"}},
+        Permissions=["ALL"],
+        PermissionsWithGrantOption=["SELECT"],
+    )
+
+    with pytest.raises(client.exceptions.InvalidInputException):
+        client.list_permissions(Principal={"DataLakePrincipalIdentifier": "asdf"})
+
+    with pytest.raises(client.exceptions.InvalidInputException):
+        client.list_permissions(Resource={})
+
+
 def grant_table_permissions(
     client: boto3.client, catalog_id: str, principal: str, db: str, table: str
 ):
@@ -159,6 +177,18 @@ def grant_catalog_permissions(client: boto3.client, catalog_id: str, principal:
     )
 
 
+def grant_data_location_permissions(
+    client: boto3.client, resource_arn: str, catalog_id: str, principal: str
+):
+    client.grant_permissions(
+        CatalogId=catalog_id,
+        Principal={"DataLakePrincipalIdentifier": principal},
+        Resource={"DataLocation": {"ResourceArn": resource_arn}},
+        Permissions=["DATA_LOCATION_ACCESS"],
+        PermissionsWithGrantOption=[],
+    )
+
+
 def db_response(principal: str, catalog_id: str, db: str) -> Dict:
     return {
         "Principal": {"DataLakePrincipalIdentifier": principal},
@@ -198,6 +228,20 @@ def catalog_response(principal: str) -> Dict:
     }
 
 
+def data_location_response(
+    principal: str, resource_arn: str, catalog_id: Optional[str] = None
+) -> Dict:
+    response = {
+        "Principal": {"DataLakePrincipalIdentifier": principal},
+        "Resource": {"DataLocation": {"ResourceArn": resource_arn}},
+        "Permissions": ["DATA_LOCATION_ACCESS"],
+        "PermissionsWithGrantOption": [],
+    }
+    if catalog_id is not None:
+        response["Resource"]["CatalogId"] = catalog_id
+    return response
+
+
 @mock_lakeformation
 def test_list_permissions_filtered_for_catalog_id():
     client = boto3.client("lakeformation", region_name="eu-west-2")
@@ -342,6 +386,45 @@ def test_list_permissions_filtered_for_resource_table():
     ]
 
 
+@mock_lakeformation
+def test_list_permissions_filtered_for_resource_data_location():
+    client = boto3.client("lakeformation", region_name="eu-west-2")
+    catalog_id = "000000000000"
+    principal = "principal"
+    data_location_resource_arn_1 = "resource_arn_1"
+    data_location_resource_arn_2 = "resource_arn_2"
+
+    grant_data_location_permissions(
+        catalog_id=catalog_id,
+        client=client,
+        resource_arn=data_location_resource_arn_1,
+        principal=principal,
+    )
+    grant_data_location_permissions(
+        catalog_id=catalog_id,
+        client=client,
+        resource_arn=data_location_resource_arn_2,
+        principal=principal,
+    )
+
+    res = client.list_permissions(
+        CatalogId=catalog_id,
+        Resource={"DataLocation": {"ResourceArn": data_location_resource_arn_1}},
+    )
+    assert res["PrincipalResourcePermissions"] == [
+        data_location_response(principal, resource_arn=data_location_resource_arn_1)
+    ]
+
+    res = client.list_permissions(
+        CatalogId=catalog_id,
+        Resource={"DataLocation": {"ResourceArn": data_location_resource_arn_2}},
+    )
+
+    assert res["PrincipalResourcePermissions"] == [
+        data_location_response(principal, resource_arn=data_location_resource_arn_2)
+    ]
+
+
 @mock_lakeformation
 def test_revoke_permissions():
     client = boto3.client("lakeformation", region_name="eu-west-2")