LakeFormation: add_lf_tags_to_resource() (#6879)

.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

IMPLEMENTATION_COVERAGE.md

@@ -4246,9 +4246,9 @@
 
 ## lakeformation
 <details>
-<summary>32% implemented</summary>
+<summary>40% implemented</summary>
 
-- [ ] 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

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

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

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 "{}"

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,

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

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
+@lakeformation_aws_verified
-def test_lf_tags():
+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()
+    client.update_lf_tag(TagKey="tag1", TagValuesToDelete=["1a", "1c"])
-    assert len(resp["LFTags"]) == 3
+
-    assert {
+    tags = client.list_lf_tags()["LFTags"]
-        "CatalogId": DEFAULT_ACCOUNT_ID,
+    assert set([x["CatalogId"] for x in tags]) == {account_id}
-        "TagKey": "tag1",
+    tag_keys = [x["TagKey"] for x in tags]
-        "TagValues": ["1a", "1b"],
+    assert "tag1" in tag_keys
-    } in resp["LFTags"]
+    assert "tag2" in tag_keys
-    assert {
+    assert "tag3" in tag_keys
-        "CatalogId": DEFAULT_ACCOUNT_ID,
+
-        "TagKey": "tag2",
+    assert [x for x in tags if x["TagKey"] == "tag1"][0]["TagValues"] == ["1b"]
-        "TagValues": ["2a", "2b"],
+    assert set([x for x in tags if x["TagKey"] == "tag2"][0]["TagValues"]) == {
-    } in resp["LFTags"]
+        "2a",
-    assert {
+        "2b",
-        "CatalogId": DEFAULT_ACCOUNT_ID,
+    }
-        "TagKey": "tag3",
+    assert set([x for x in tags if x["TagKey"] == "tag3"][0]["TagValues"]) == {
-        "TagValues": ["3a", "3b"],
+        "3a",
-    } in resp["LFTags"]
+        "3b",
+    }
 
     client.delete_lf_tag(TagKey="tag2")
 
-    resp = client.list_lf_tags()
+    tags = client.list_lf_tags()["LFTags"]
-    assert len(resp["LFTags"]) == 2
+    tag_keys = [x["TagKey"] for x in tags]
-    assert {
+    assert "tag1" in tag_keys
-        "CatalogId": DEFAULT_ACCOUNT_ID,
+    assert "tag3" in tag_keys
-        "TagKey": "tag1",
+    assert "tag2" not in tag_keys
-        "TagValues": ["1a", "1b"],
+
-    } in resp["LFTags"]
+    client.delete_lf_tag(TagKey="tag1")
-    assert {
+    client.delete_lf_tag(TagKey="tag3")
-        "CatalogId": DEFAULT_ACCOUNT_ID,
-        "TagKey": "tag3",
-        "TagValues": ["3a", "3b"],
-    } in resp["LFTags"]
 
 
 @mock_lakeformation

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],
+            }
+        }
+    )