From 24d9ea61ce44952829b64d40f13961108e4eb85d Mon Sep 17 00:00:00 2001 From: Bert Blommers Date: Tue, 3 Oct 2023 21:06:07 +0000 Subject: [PATCH] LakeFormation: add_lf_tags_to_resource() (#6879) --- .github/workflows/tests_real_aws.yml | 2 +- IMPLEMENTATION_COVERAGE.md | 10 +- docs/docs/services/lakeformation.rst | 8 +- moto/lakeformation/models.py | 135 ++++++- moto/lakeformation/responses.py | 42 ++ moto/lakeformation/urls.py | 4 + tests/test_lakeformation/__init__.py | 89 +++++ .../test_lakeformation/test_lakeformation.py | 72 ++-- .../test_lakeformation/test_resource_tags.py | 361 ++++++++++++++++++ 9 files changed, 677 insertions(+), 46 deletions(-) create mode 100644 tests/test_lakeformation/test_resource_tags.py diff --git a/.github/workflows/tests_real_aws.yml b/.github/workflows/tests_real_aws.yml index 342769732..1c726cfdc 100644 --- a/.github/workflows/tests_real_aws.yml +++ b/.github/workflows/tests_real_aws.yml @@ -42,4 +42,4 @@ jobs: env: MOTO_TEST_ALLOW_AWS_REQUEST: ${{ true }} run: | - pytest -sv tests/test_ec2/ tests/test_ses/ tests/test_s3 -m aws_verified + pytest -sv tests/test_ec2/ tests/test_lakeformation/ tests/test_ses/ tests/test_s3 -m aws_verified diff --git a/IMPLEMENTATION_COVERAGE.md b/IMPLEMENTATION_COVERAGE.md index d80f6b0b2..053a5dc86 100644 --- a/IMPLEMENTATION_COVERAGE.md +++ b/IMPLEMENTATION_COVERAGE.md @@ -4246,9 +4246,9 @@ ## lakeformation
-32% implemented +40% implemented -- [ ] add_lf_tags_to_resource +- [X] add_lf_tags_to_resource - [ ] assume_decorated_role_with_saml - [X] batch_grant_permissions - [X] batch_revoke_permissions @@ -4271,7 +4271,7 @@ - [X] get_lf_tag - [ ] get_query_state - [ ] get_query_statistics -- [ ] get_resource_lf_tags +- [X] get_resource_lf_tags - [ ] get_table_objects - [ ] get_temporary_glue_partition_credentials - [ ] get_temporary_glue_table_credentials @@ -4287,14 +4287,14 @@ - [ ] list_transactions - [X] put_data_lake_settings - [X] register_resource -- [ ] remove_lf_tags_from_resource +- [X] remove_lf_tags_from_resource - [X] revoke_permissions - [ ] search_databases_by_lf_tags - [ ] search_tables_by_lf_tags - [ ] start_query_planning - [ ] start_transaction - [ ] update_data_cells_filter -- [ ] update_lf_tag +- [X] update_lf_tag - [ ] update_resource - [ ] update_table_objects - [ ] update_table_storage_optimizer diff --git a/docs/docs/services/lakeformation.rst b/docs/docs/services/lakeformation.rst index 642f639a8..e328c98c2 100644 --- a/docs/docs/services/lakeformation.rst +++ b/docs/docs/services/lakeformation.rst @@ -25,7 +25,7 @@ lakeformation |start-h3| Implemented features for this service |end-h3| -- [ ] add_lf_tags_to_resource +- [X] add_lf_tags_to_resource - [ ] assume_decorated_role_with_saml - [X] batch_grant_permissions - [X] batch_revoke_permissions @@ -48,7 +48,7 @@ lakeformation - [X] get_lf_tag - [ ] get_query_state - [ ] get_query_statistics -- [ ] get_resource_lf_tags +- [X] get_resource_lf_tags - [ ] get_table_objects - [ ] get_temporary_glue_partition_credentials - [ ] get_temporary_glue_table_credentials @@ -72,14 +72,14 @@ lakeformation - [ ] list_transactions - [X] put_data_lake_settings - [X] register_resource -- [ ] remove_lf_tags_from_resource +- [X] remove_lf_tags_from_resource - [X] revoke_permissions - [ ] search_databases_by_lf_tags - [ ] search_tables_by_lf_tags - [ ] start_query_planning - [ ] start_transaction - [ ] update_data_cells_filter -- [ ] update_lf_tag +- [X] update_lf_tag - [ ] update_resource - [ ] update_table_objects - [ ] update_table_storage_optimizer diff --git a/moto/lakeformation/models.py b/moto/lakeformation/models.py index 61005f6cf..18cd77abb 100644 --- a/moto/lakeformation/models.py +++ b/moto/lakeformation/models.py @@ -1,5 +1,5 @@ from collections import defaultdict -from typing import Any, Dict, List +from typing import Any, Dict, List, Tuple from moto.core import BaseBackend, BackendDict, BaseModel from moto.utilities.tagging_service import TaggingService @@ -46,6 +46,9 @@ class LakeFormationBackend(BaseBackend): self.settings: Dict[str, Dict[str, Any]] = defaultdict(default_settings) self.grants: Dict[str, List[Dict[str, Any]]] = defaultdict(list) self.tagger = TaggingService() + self.lf_database_tags: Dict[Tuple[str, str], List[Dict[str, str]]] = {} + self.lf_table_tags: Dict[Tuple[str, str, str], List[Dict[str, str]]] = {} + self.lf_columns_tags: Dict[Tuple[str, ...], List[Dict[str, str]]] = {} def describe_resource(self, resource_arn: str) -> Resource: if resource_arn not in self.resources: @@ -132,11 +135,37 @@ class LakeFormationBackend(BaseBackend): arn = f"arn:lakeformation:{catalog_id}" self.tagger.untag_resource_using_names(arn, tag_names=[key]) + # Also remove any LF resource tags that used this tag-key + for db_name in self.lf_database_tags: + self.lf_database_tags[db_name] = [ + tag for tag in self.lf_database_tags[db_name] if tag["TagKey"] != key + ] + for table in self.lf_table_tags: + self.lf_table_tags[table] = [ + tag for tag in self.lf_table_tags[table] if tag["TagKey"] != key + ] + for column in self.lf_columns_tags: + self.lf_columns_tags[column] = [ + tag for tag in self.lf_columns_tags[column] if tag["TagKey"] != key + ] + def list_lf_tags(self, catalog_id: str) -> Dict[str, str]: # There is no ARN that we can use, so just create another unique identifier that's easy to recognize and reproduce arn = f"arn:lakeformation:{catalog_id}" return self.tagger.get_tag_dict_for_resource(arn=arn) + def update_lf_tag( + self, catalog_id: str, tag_key: str, to_delete: List[str], to_add: List[str] + ) -> None: + arn = f"arn:lakeformation:{catalog_id}" + existing_tags = self.list_lf_tags(catalog_id) + existing_tags[tag_key].extend(to_add or []) # type: ignore + for tag in to_delete or []: + existing_tags[tag_key].remove(tag) # type: ignore + self.tagger.tag_resource( + arn, TaggingService.convert_dict_to_tags_input(existing_tags) + ) + def list_data_cells_filter(self) -> List[Dict[str, Any]]: """ This currently just returns an empty list, as the corresponding Create is not yet implemented @@ -169,5 +198,109 @@ class LakeFormationBackend(BaseBackend): ), ) + def add_lf_tags_to_resource( + self, catalog_id: str, resource: Dict[str, Any], tags: List[Dict[str, str]] + ) -> List[Dict[str, Any]]: + existing_lf_tags = self.list_lf_tags(catalog_id) + failures = [] + + for tag in tags: + if "CatalogId" not in tag: + tag["CatalogId"] = catalog_id + if tag["TagKey"] not in existing_lf_tags: + failures.append( + { + "LFTag": tag, + "Error": { + "ErrorCode": "EntityNotFoundException", + "ErrorMessage": "Tag or tag value does not exist.", + }, + } + ) + + if failures: + return failures + + if "Database" in resource: + db_catalog_id = resource["Database"].get("CatalogId", self.account_id) + db_name = resource["Database"]["Name"] + self.lf_database_tags[(db_catalog_id, db_name)] = tags + if "Table" in resource: + db_catalog_id = resource["Table"].get("CatalogId", self.account_id) + db_name = resource["Table"]["DatabaseName"] + name = resource["Table"]["Name"] + self.lf_table_tags[(db_catalog_id, db_name, name)] = tags + if "TableWithColumns" in resource: + db_catalog_id = resource["TableWithColumns"].get( + "CatalogId", self.account_id + ) + db_name = resource["TableWithColumns"]["DatabaseName"] + name = resource["TableWithColumns"]["Name"] + for column in resource["TableWithColumns"]["ColumnNames"]: + self.lf_columns_tags[(db_catalog_id, db_name, name, column)] = tags + return failures + + def get_resource_lf_tags( + self, + catalog_id: str, # pylint: disable=unused-argument + resource: Dict[str, Any], + ) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]], List[Dict[str, Any]]]: + database_tags = [] + table_tags = [] + column_tags = [] + if "Database" in resource: + database_catalog_id = resource["Database"].get("CatalogId", self.account_id) + database_name = resource["Database"]["Name"] + database_tags = self.lf_database_tags[(database_catalog_id, database_name)] + if "Table" in resource: + db_catalog_id = resource["Table"].get("CatalogId", self.account_id) + db_name = resource["Table"]["DatabaseName"] + name = resource["Table"]["Name"] + table_tags = self.lf_table_tags[(db_catalog_id, db_name, name)] + if "TableWithColumns" in resource: + for column in resource["TableWithColumns"]["ColumnNames"]: + db_catalog_id = resource["TableWithColumns"].get( + "CatalogId", self.account_id + ) + db_name = resource["TableWithColumns"]["DatabaseName"] + name = resource["TableWithColumns"]["Name"] + dct_key = (db_catalog_id, db_name, name, column) + if self.lf_columns_tags.get(dct_key): + column_tags.append( + {"Name": column, "LFTags": self.lf_columns_tags[dct_key]} + ) + return database_tags, table_tags, column_tags + + def remove_lf_tags_from_resource( + self, catalog_id: str, resource: Dict[str, Any], tags: List[Dict[str, str]] + ) -> None: + for tag in tags: + if "CatalogId" not in tag: + tag["CatalogId"] = catalog_id + if "Database" in resource: + database_catalog_id = resource["Database"].get("CatalogId", self.account_id) + database_name = resource["Database"]["Name"] + existing_tags = self.lf_database_tags[(database_catalog_id, database_name)] + for tag in tags: + existing_tags.remove(tag) + if "Table" in resource: + db_catalog_id = resource["Table"].get("CatalogId", self.account_id) + db_name = resource["Table"]["DatabaseName"] + name = resource["Table"]["Name"] + existing_tags = self.lf_table_tags[(db_catalog_id, db_name, name)] + for tag in tags: + existing_tags.remove(tag) + if "TableWithColumns" in resource: + for column in resource["TableWithColumns"]["ColumnNames"]: + db_catalog_id = resource["TableWithColumns"].get( + "CatalogId", self.account_id + ) + db_name = resource["TableWithColumns"]["DatabaseName"] + name = resource["TableWithColumns"]["Name"] + dct_key = (db_catalog_id, db_name, name, column) + existing_tags = self.lf_columns_tags[dct_key] + for tag in tags: + existing_tags.remove(tag) + lakeformation_backends = BackendDict(LakeFormationBackend, "lakeformation") diff --git a/moto/lakeformation/responses.py b/moto/lakeformation/responses.py index 4109f0265..628bbb5ba 100644 --- a/moto/lakeformation/responses.py +++ b/moto/lakeformation/responses.py @@ -1,5 +1,6 @@ """Handles incoming lakeformation requests, invokes methods, returns responses.""" import json +from typing import Any, Dict from moto.core.responses import BaseResponse from .models import lakeformation_backends, LakeFormationBackend @@ -122,6 +123,14 @@ class LakeFormationResponse(BaseResponse): } ) + def update_lf_tag(self) -> str: + catalog_id = self._get_param("CatalogId") or self.current_account + tag_key = self._get_param("TagKey") + to_delete = self._get_param("TagValuesToDelete") + to_add = self._get_param("TagValuesToAdd") + self.lakeformation_backend.update_lf_tag(catalog_id, tag_key, to_delete, to_add) + return "{}" + def list_data_cells_filter(self) -> str: data_cells = self.lakeformation_backend.list_data_cells_filter() return json.dumps({"DataCellsFilters": data_cells}) @@ -137,3 +146,36 @@ class LakeFormationResponse(BaseResponse): entries = self._get_param("Entries") self.lakeformation_backend.batch_revoke_permissions(catalog_id, entries) return json.dumps({"Failures": []}) + + def add_lf_tags_to_resource(self) -> str: + catalog_id = self._get_param("CatalogId") or self.current_account + resource = self._get_param("Resource") + tags = self._get_param("LFTags") + failures = self.lakeformation_backend.add_lf_tags_to_resource( + catalog_id, resource, tags + ) + return json.dumps({"Failures": failures}) + + def get_resource_lf_tags(self) -> str: + catalog_id = self._get_param("CatalogId") or self.current_account + resource = self._get_param("Resource") + db, table, columns = self.lakeformation_backend.get_resource_lf_tags( + catalog_id, resource + ) + resp: Dict[str, Any] = {} + if db: + resp["LFTagOnDatabase"] = db + if table: + resp["LFTagsOnTable"] = table + if columns: + resp["LFTagsOnColumns"] = columns + return json.dumps(resp) + + def remove_lf_tags_from_resource(self) -> str: + catalog_id = self._get_param("CatalogId") or self.current_account + resource = self._get_param("Resource") + tags = self._get_param("LFTags") + self.lakeformation_backend.remove_lf_tags_from_resource( + catalog_id, resource, tags + ) + return "{}" diff --git a/moto/lakeformation/urls.py b/moto/lakeformation/urls.py index ed8c4e093..fa6458e10 100644 --- a/moto/lakeformation/urls.py +++ b/moto/lakeformation/urls.py @@ -22,7 +22,11 @@ url_paths = { "{0}/CreateLFTag$": response.dispatch, "{0}/GetLFTag$": response.dispatch, "{0}/DeleteLFTag$": response.dispatch, + "{0}/UpdateLFTag": response.dispatch, "{0}/ListLFTags$": response.dispatch, + "{0}/AddLFTagsToResource": response.dispatch, + "{0}/RemoveLFTagsFromResource": response.dispatch, + "{0}/GetResourceLFTags": response.dispatch, "{0}/ListDataCellsFilter$": response.dispatch, "{0}/BatchGrantPermissions$": response.dispatch, "{0}/BatchRevokePermissions$": response.dispatch, diff --git a/tests/test_lakeformation/__init__.py b/tests/test_lakeformation/__init__.py index e69de29bb..6d2c63bbf 100644 --- a/tests/test_lakeformation/__init__.py +++ b/tests/test_lakeformation/__init__.py @@ -0,0 +1,89 @@ +import boto3 +import os +from functools import wraps +from moto import mock_glue, mock_lakeformation, mock_s3, mock_sts +from uuid import uuid4 + + +def lakeformation_aws_verified(func): + """ + Function that is verified to work against AWS. + Can be run against AWS at any time by setting: + MOTO_TEST_ALLOW_AWS_REQUEST=true + + If this environment variable is not set, the function runs in a `mock_lakeformation`/`mock_sts`/`mock_s3` context. + + Note that LakeFormation is not enabled by default - visit the AWS Console to permit access to the user who executes these tests. + """ + + @wraps(func) + def pagination_wrapper(): + glue = boto3.client("glue", region_name="eu-west-2") + lf = boto3.client("lakeformation", region_name="eu-west-2") + s3 = boto3.client("s3", region_name="us-east-1") + bucket_name = str(uuid4()) + + allow_aws_request = ( + os.environ.get("MOTO_TEST_ALLOW_AWS_REQUEST", "false").lower() == "true" + ) + + if allow_aws_request: + resp = create_glue_infra_and_test(bucket_name, s3, glue, lf) + else: + with mock_glue(), mock_lakeformation(), mock_s3(), mock_sts(): + resp = create_glue_infra_and_test(bucket_name, s3, glue, lf) + return resp + + def create_glue_infra_and_test(bucket_name, s3, glue, lf): + s3.create_bucket(Bucket=bucket_name) + s3.put_bucket_tagging( + Bucket=bucket_name, + Tagging={"TagSet": [{"Key": "environment", "Value": "moto_tests"}]}, + ) + lf.register_resource( + ResourceArn=f"arn:aws:s3:::{bucket_name}", UseServiceLinkedRole=True + ) + + db_name = str(uuid4())[0:6] + table_name = str(uuid4())[0:6] + column_name = str(uuid4())[0:6] + glue.create_database( + DatabaseInput={"Name": db_name}, Tags={"environment": "moto_tests"} + ) + glue.create_table( + DatabaseName=db_name, + TableInput={ + "Name": table_name, + "StorageDescriptor": { + "Columns": [{"Name": column_name, "Type": "string"}] + }, + }, + ) + + try: + resp = func(bucket_name, db_name, table_name, column_name) + finally: + ### CLEANUP ### + + glue.delete_table(DatabaseName=db_name, Name=table_name) + glue.delete_database(Name=db_name) + + lf.deregister_resource(ResourceArn=f"arn:aws:s3:::{bucket_name}") + + versions = s3.list_object_versions(Bucket=bucket_name).get("Versions", []) + for key in versions: + s3.delete_object( + Bucket=bucket_name, Key=key["Key"], VersionId=key.get("VersionId") + ) + delete_markers = s3.list_object_versions(Bucket=bucket_name).get( + "DeleteMarkers", [] + ) + for key in delete_markers: + s3.delete_object( + Bucket=bucket_name, Key=key["Key"], VersionId=key.get("VersionId") + ) + s3.delete_bucket(Bucket=bucket_name) + + return resp + + return pagination_wrapper diff --git a/tests/test_lakeformation/test_lakeformation.py b/tests/test_lakeformation/test_lakeformation.py index 5083a9f5e..aa60456fa 100644 --- a/tests/test_lakeformation/test_lakeformation.py +++ b/tests/test_lakeformation/test_lakeformation.py @@ -4,7 +4,8 @@ import pytest from botocore.exceptions import ClientError from moto import mock_lakeformation -from moto.core import DEFAULT_ACCOUNT_ID + +from . import lakeformation_aws_verified # See our Development Tips on writing tests for hints on how to write good tests: # http://docs.getmoto.org/en/latest/docs/contributing/development_tips/tests.html @@ -145,51 +146,52 @@ def test_revoke_permissions(): ] -@mock_lakeformation -def test_lf_tags(): +@lakeformation_aws_verified +def test_lf_tags( + bucket_name=None, db_name=None, table_name=None, column_name=None +): # pylint: disable=unused-argument client = boto3.client("lakeformation", region_name="eu-west-2") + sts = boto3.client("sts", "eu-west-2") + account_id = sts.get_caller_identity()["Account"] - client.create_lf_tag(TagKey="tag1", TagValues=["1a", "1b"]) + client.create_lf_tag(TagKey="tag1", TagValues=["1a", "1b", "1c"]) client.create_lf_tag(TagKey="tag2", TagValues=["2a", "2b"]) client.create_lf_tag(TagKey="tag3", TagValues=["3a", "3b"]) resp = client.get_lf_tag(TagKey="tag1") - assert resp["CatalogId"] == DEFAULT_ACCOUNT_ID + assert resp["CatalogId"] == account_id assert resp["TagKey"] == "tag1" - assert resp["TagValues"] == ["1a", "1b"] + assert resp["TagValues"] == ["1a", "1b", "1c"] - resp = client.list_lf_tags() - assert len(resp["LFTags"]) == 3 - assert { - "CatalogId": DEFAULT_ACCOUNT_ID, - "TagKey": "tag1", - "TagValues": ["1a", "1b"], - } in resp["LFTags"] - assert { - "CatalogId": DEFAULT_ACCOUNT_ID, - "TagKey": "tag2", - "TagValues": ["2a", "2b"], - } in resp["LFTags"] - assert { - "CatalogId": DEFAULT_ACCOUNT_ID, - "TagKey": "tag3", - "TagValues": ["3a", "3b"], - } in resp["LFTags"] + client.update_lf_tag(TagKey="tag1", TagValuesToDelete=["1a", "1c"]) + + tags = client.list_lf_tags()["LFTags"] + assert set([x["CatalogId"] for x in tags]) == {account_id} + tag_keys = [x["TagKey"] for x in tags] + assert "tag1" in tag_keys + assert "tag2" in tag_keys + assert "tag3" in tag_keys + + assert [x for x in tags if x["TagKey"] == "tag1"][0]["TagValues"] == ["1b"] + assert set([x for x in tags if x["TagKey"] == "tag2"][0]["TagValues"]) == { + "2a", + "2b", + } + assert set([x for x in tags if x["TagKey"] == "tag3"][0]["TagValues"]) == { + "3a", + "3b", + } client.delete_lf_tag(TagKey="tag2") - resp = client.list_lf_tags() - assert len(resp["LFTags"]) == 2 - assert { - "CatalogId": DEFAULT_ACCOUNT_ID, - "TagKey": "tag1", - "TagValues": ["1a", "1b"], - } in resp["LFTags"] - assert { - "CatalogId": DEFAULT_ACCOUNT_ID, - "TagKey": "tag3", - "TagValues": ["3a", "3b"], - } in resp["LFTags"] + tags = client.list_lf_tags()["LFTags"] + tag_keys = [x["TagKey"] for x in tags] + assert "tag1" in tag_keys + assert "tag3" in tag_keys + assert "tag2" not in tag_keys + + client.delete_lf_tag(TagKey="tag1") + client.delete_lf_tag(TagKey="tag3") @mock_lakeformation diff --git a/tests/test_lakeformation/test_resource_tags.py b/tests/test_lakeformation/test_resource_tags.py new file mode 100644 index 000000000..01841eb50 --- /dev/null +++ b/tests/test_lakeformation/test_resource_tags.py @@ -0,0 +1,361 @@ +import boto3 + +from uuid import uuid4 + +from . import lakeformation_aws_verified + + +@lakeformation_aws_verified +def test_add_unknown_lf_tags( + bucket_name=None, # pylint: disable=unused-argument + db_name=None, + table_name=None, # pylint: disable=unused-argument + column_name=None, # pylint: disable=unused-argument +): + client = boto3.client("lakeformation", region_name="eu-west-2") + sts = boto3.client("sts", "eu-west-2") + account_id = sts.get_caller_identity()["Account"] + + failures = client.add_lf_tags_to_resource( + Resource={"Database": {"Name": db_name}}, + LFTags=[{"TagKey": "unknown-tag", "TagValues": ["value"]}], + )["Failures"] + assert len(failures) == 1 + assert failures[0]["LFTag"] == { + "CatalogId": account_id, + "TagKey": "unknown-tag", + "TagValues": ["value"], + } + assert failures[0]["Error"] == { + "ErrorCode": "EntityNotFoundException", + "ErrorMessage": "Tag or tag value does not exist.", + } + + +@lakeformation_aws_verified +def test_tag_lakeformation_database( + bucket_name=None, # pylint: disable=unused-argument + db_name=None, + table_name=None, # pylint: disable=unused-argument + column_name=None, # pylint: disable=unused-argument +): + client = boto3.client("lakeformation", region_name="eu-west-2") + sts = boto3.client("sts", "eu-west-2") + account_id = sts.get_caller_identity()["Account"] + + tag_name = str(uuid4())[0:6] + client.create_lf_tag(TagKey=tag_name, TagValues=["value1"]) + resp = client.add_lf_tags_to_resource( + Resource={"Database": {"Name": db_name}}, + LFTags=[{"TagKey": tag_name, "TagValues": ["value1"]}], + ) + assert resp["Failures"] == [] + + tags = client.get_resource_lf_tags(Resource={"Database": {"Name": db_name}})[ + "LFTagOnDatabase" + ] + assert tags == [ + {"CatalogId": account_id, "TagKey": tag_name, "TagValues": ["value1"]} + ] + + client.update_lf_tag(TagKey=tag_name, TagValuesToAdd=["value2"]) + + all_tags = client.list_lf_tags()["LFTags"] + our_tag = next(tag for tag in all_tags if tag["TagKey"] == tag_name) + assert set(our_tag["TagValues"]) == {"value1", "value2"} + + # The value for this particular resource has not been updated + db_tags = client.get_resource_lf_tags(Resource={"Database": {"Name": db_name}})[ + "LFTagOnDatabase" + ] + assert db_tags == [ + {"CatalogId": account_id, "TagKey": tag_name, "TagValues": ["value1"]} + ] + + # Update the existing tags for this resource + client.add_lf_tags_to_resource( + Resource={"Database": {"Name": db_name}}, + LFTags=[{"TagKey": tag_name, "TagValues": ["value2"]}], + ) + + db_tags = client.get_resource_lf_tags(Resource={"Database": {"Name": db_name}})[ + "LFTagOnDatabase" + ] + assert db_tags == [ + {"CatalogId": account_id, "TagKey": tag_name, "TagValues": ["value2"]} + ] + + # Try remove and re-add + client.remove_lf_tags_from_resource( + Resource={"Database": {"Name": db_name}}, + LFTags=[{"TagKey": tag_name, "TagValues": ["value2"]}], + ) + + assert "LFTagOnDatabase" not in client.get_resource_lf_tags( + Resource={"Database": {"Name": db_name}} + ) + + client.add_lf_tags_to_resource( + Resource={"Database": {"Name": db_name}}, + LFTags=[{"TagKey": tag_name, "TagValues": ["value1"]}], + ) + + db_tags = client.get_resource_lf_tags(Resource={"Database": {"Name": db_name}})[ + "LFTagOnDatabase" + ] + assert db_tags == [ + {"CatalogId": account_id, "TagKey": tag_name, "TagValues": ["value1"]} + ] + + # Deleting the tag automatically deletes it from any resource + client.delete_lf_tag(TagKey=tag_name) + + assert "LFTagOnDatabase" not in client.get_resource_lf_tags( + Resource={"Database": {"Name": db_name}} + ) + + +@lakeformation_aws_verified +def test_tag_lakeformation_table( + bucket_name=None, # pylint: disable=unused-argument + db_name=None, + table_name=None, + column_name=None, # pylint: disable=unused-argument +): + client = boto3.client("lakeformation", region_name="eu-west-2") + sts = boto3.client("sts", "eu-west-2") + account_id = sts.get_caller_identity()["Account"] + + tag_name = str(uuid4())[0:6] + client.create_lf_tag(TagKey=tag_name, TagValues=["value1"]) + resp = client.add_lf_tags_to_resource( + Resource={"Table": {"DatabaseName": db_name, "Name": table_name}}, + LFTags=[{"TagKey": tag_name, "TagValues": ["value1"]}], + ) + assert resp["Failures"] == [] + + tags = client.get_resource_lf_tags( + Resource={"Table": {"DatabaseName": db_name, "Name": table_name}}, + )["LFTagsOnTable"] + assert tags == [ + {"CatalogId": account_id, "TagKey": tag_name, "TagValues": ["value1"]} + ] + + client.update_lf_tag(TagKey=tag_name, TagValuesToAdd=["value2"]) + + all_tags = client.list_lf_tags()["LFTags"] + our_tag = next(tag for tag in all_tags if tag["TagKey"] == tag_name) + assert set(our_tag["TagValues"]) == {"value1", "value2"} + + # The value for this particular resource has not been updated + db_tags = client.get_resource_lf_tags( + Resource={"Table": {"DatabaseName": db_name, "Name": table_name}}, + )["LFTagsOnTable"] + assert db_tags == [ + {"CatalogId": account_id, "TagKey": tag_name, "TagValues": ["value1"]} + ] + + # Update the existing tags for this resource + client.add_lf_tags_to_resource( + Resource={"Table": {"DatabaseName": db_name, "Name": table_name}}, + LFTags=[{"TagKey": tag_name, "TagValues": ["value2"]}], + ) + + db_tags = client.get_resource_lf_tags( + Resource={"Table": {"DatabaseName": db_name, "Name": table_name}} + )["LFTagsOnTable"] + assert db_tags == [ + {"CatalogId": account_id, "TagKey": tag_name, "TagValues": ["value2"]} + ] + + # Try remove and re-add + client.remove_lf_tags_from_resource( + Resource={"Table": {"DatabaseName": db_name, "Name": table_name}}, + LFTags=[{"TagKey": tag_name, "TagValues": ["value2"]}], + ) + + assert "LFTagsOnTable" not in client.get_resource_lf_tags( + Resource={"Table": {"DatabaseName": db_name, "Name": table_name}} + ) + + client.add_lf_tags_to_resource( + Resource={"Table": {"DatabaseName": db_name, "Name": table_name}}, + LFTags=[{"TagKey": tag_name, "TagValues": ["value1"]}], + ) + + db_tags = client.get_resource_lf_tags( + Resource={"Table": {"DatabaseName": db_name, "Name": table_name}} + )["LFTagsOnTable"] + assert db_tags == [ + {"CatalogId": account_id, "TagKey": tag_name, "TagValues": ["value1"]} + ] + + # Deleting the tag automatically deletes it from any resource + client.delete_lf_tag(TagKey=tag_name) + + assert "LFTagsOnTable" not in client.get_resource_lf_tags( + Resource={"Table": {"DatabaseName": db_name, "Name": table_name}} + ) + + +@lakeformation_aws_verified +def test_tag_lakeformation_columns( + bucket_name=None, # pylint: disable=unused-argument + db_name=None, + table_name=None, + column_name=None, +): + client = boto3.client("lakeformation", region_name="eu-west-2") + sts = boto3.client("sts", "eu-west-2") + account_id = sts.get_caller_identity()["Account"] + + tag_name = str(uuid4())[0:6] + client.create_lf_tag(TagKey=tag_name, TagValues=["value1"]) + resp = client.add_lf_tags_to_resource( + Resource={ + "TableWithColumns": { + "DatabaseName": db_name, + "Name": table_name, + "ColumnNames": [column_name], + } + }, + LFTags=[{"TagKey": tag_name, "TagValues": ["value1"]}], + ) + assert resp["Failures"] == [] + + tags = client.get_resource_lf_tags( + Resource={ + "TableWithColumns": { + "DatabaseName": db_name, + "Name": table_name, + "ColumnNames": [column_name], + } + }, + )["LFTagsOnColumns"] + assert tags == [ + { + "Name": column_name, + "LFTags": [ + {"CatalogId": account_id, "TagKey": tag_name, "TagValues": ["value1"]} + ], + } + ] + + client.update_lf_tag(TagKey=tag_name, TagValuesToAdd=["value2"]) + + all_tags = client.list_lf_tags()["LFTags"] + our_tag = next(tag for tag in all_tags if tag["TagKey"] == tag_name) + assert set(our_tag["TagValues"]) == {"value1", "value2"} + + # The value for this particular resource has not been updated + tags = client.get_resource_lf_tags( + Resource={ + "TableWithColumns": { + "DatabaseName": db_name, + "Name": table_name, + "ColumnNames": [column_name], + } + }, + )["LFTagsOnColumns"] + assert tags == [ + { + "Name": column_name, + "LFTags": [ + {"CatalogId": account_id, "TagKey": tag_name, "TagValues": ["value1"]} + ], + } + ] + + # Update the existing tags for this resource + client.add_lf_tags_to_resource( + Resource={ + "TableWithColumns": { + "DatabaseName": db_name, + "Name": table_name, + "ColumnNames": [column_name], + } + }, + LFTags=[{"TagKey": tag_name, "TagValues": ["value2"]}], + ) + + tags = client.get_resource_lf_tags( + Resource={ + "TableWithColumns": { + "DatabaseName": db_name, + "Name": table_name, + "ColumnNames": [column_name], + } + } + )["LFTagsOnColumns"] + assert tags == [ + { + "Name": column_name, + "LFTags": [ + {"CatalogId": account_id, "TagKey": tag_name, "TagValues": ["value2"]} + ], + } + ] + + # Try remove and re-add + client.remove_lf_tags_from_resource( + Resource={ + "TableWithColumns": { + "DatabaseName": db_name, + "Name": table_name, + "ColumnNames": [column_name], + } + }, + LFTags=[{"TagKey": tag_name, "TagValues": ["value2"]}], + ) + + assert "LFTagsOnColumns" not in client.get_resource_lf_tags( + Resource={ + "TableWithColumns": { + "DatabaseName": db_name, + "Name": table_name, + "ColumnNames": [column_name], + } + } + ) + + client.add_lf_tags_to_resource( + Resource={ + "TableWithColumns": { + "DatabaseName": db_name, + "Name": table_name, + "ColumnNames": [column_name], + } + }, + LFTags=[{"TagKey": tag_name, "TagValues": ["value1"]}], + ) + + tags = client.get_resource_lf_tags( + Resource={ + "TableWithColumns": { + "DatabaseName": db_name, + "Name": table_name, + "ColumnNames": [column_name], + } + } + )["LFTagsOnColumns"] + assert tags == [ + { + "Name": column_name, + "LFTags": [ + {"CatalogId": account_id, "TagKey": tag_name, "TagValues": ["value1"]} + ], + } + ] + + # Deleting the tag automatically deletes it from any resource + client.delete_lf_tag(TagKey=tag_name) + + assert "LFTagsOnColumns" not in client.get_resource_lf_tags( + Resource={ + "TableWithColumns": { + "DatabaseName": db_name, + "Name": table_name, + "ColumnNames": [column_name], + } + } + )