Yasuke/moto
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

LakeFormation: extend functionality of list_permissions (#7051)

Browse Source
This commit is contained in:
JohannesKoenigTMH 2023-11-20 23:42:01 +01:00 committed by GitHub
parent 3eed5c2d68
commit 7c5f0ef7a3
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
3 changed files with 130 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

31
moto/lakeformation/models.py
View File

@ -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:

20
moto/lakeformation/responses.py
View File

@ -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,

85
tests/test_lakeformation/test_lakeformation.py
View File

@ -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")