TimestreamWrite - improvements (#4971)

.github/workflows/build.yml

@@ -224,7 +224,7 @@ jobs:
       fail-fast: false
       matrix:
         python-version: [ 3.8 ]
-        part: ["aa", "ab", "ac", "ad", "ae", "af"]
+        part: ["aa", "ab", "ac", "ad", "ae", "af", "ag"]
 
     steps:
     - uses: actions/checkout@v2
@@ -259,7 +259,7 @@ jobs:
       run: |
         cd moto-terraform-tests
         bin/list-tests -i ../tests/terraform-tests.success.txt -e ../tests/terraform-tests.failures.txt > tftestlist.txt
-        split -n l/6 tftestlist.txt tf-split-
+        split -n l/7 tftestlist.txt tf-split-
         cd ..
     - name: Run Terraform Tests
       run: |

IMPLEMENTATION_COVERAGE.md

@@ -5472,7 +5472,7 @@
 
 ## timestream-write
 <details>
-<summary>80% implemented</summary>
+<summary>100% implemented</summary>
 
 - [X] create_database
 - [X] create_table
@@ -5483,9 +5483,9 @@
 - [X] describe_table
 - [X] list_databases
 - [X] list_tables
-- [ ] list_tags_for_resource
-- [ ] tag_resource
-- [ ] untag_resource
+- [X] list_tags_for_resource
+- [X] tag_resource
+- [X] untag_resource
 - [X] update_database
 - [X] update_table
 - [X] write_records

docs/docs/services/timestream-write.rst

@@ -34,9 +34,9 @@ timestream-write
 - [X] describe_table
 - [X] list_databases
 - [X] list_tables
-- [ ] list_tags_for_resource
-- [ ] tag_resource
-- [ ] untag_resource
+- [X] list_tags_for_resource
+- [X] tag_resource
+- [X] untag_resource
 - [X] update_database
 - [X] update_table
 - [X] write_records

moto/timestreamwrite/exceptions.py

@@ -1 +1,9 @@
 """Exceptions raised by the timestreamwrite service."""
+from moto.core.exceptions import JsonRESTError
+
+
+class ResourceNotFound(JsonRESTError):
+    error_type = "com.amazonaws.timestream.v20181101#ResourceNotFoundException"
+
+    def __init__(self, msg):
+        super().__init__(ResourceNotFound.error_type, msg)

moto/timestreamwrite/models.py

@@ -1,5 +1,7 @@
 from moto.core import ACCOUNT_ID, BaseBackend, BaseModel
 from moto.core.utils import BackendDict
+from moto.utilities.tagging_service import TaggingService
+from .exceptions import ResourceNotFound
 
 
 class TimestreamTable(BaseModel):
@@ -7,7 +9,10 @@ class TimestreamTable(BaseModel):
         self.region_name = region_name
         self.name = table_name
         self.db_name = db_name
-        self.retention_properties = retention_properties
+        self.retention_properties = retention_properties or {
+            "MemoryStoreRetentionPeriodInHours": 123,
+            "MagneticStoreRetentionPeriodInDays": 123,
+        }
         self.records = []
 
     def update(self, retention_properties):
@@ -34,7 +39,9 @@ class TimestreamDatabase(BaseModel):
     def __init__(self, region_name, database_name, kms_key_id):
         self.region_name = region_name
         self.name = database_name
-        self.kms_key_id = kms_key_id
+        self.kms_key_id = (
+            kms_key_id or f"arn:aws:kms:{region_name}:{ACCOUNT_ID}:key/default_key"
+        )
         self.tables = dict()
 
     def update(self, kms_key_id):
@@ -59,6 +66,8 @@ class TimestreamDatabase(BaseModel):
         del self.tables[table_name]
 
     def describe_table(self, table_name):
+        if table_name not in self.tables:
+            raise ResourceNotFound(f"The table {table_name} does not exist.")
         return self.tables[table_name]
 
     def list_tables(self):
@@ -83,16 +92,20 @@ class TimestreamWriteBackend(BaseBackend):
     def __init__(self, region_name):
         self.region_name = region_name
         self.databases = dict()
+        self.tagging_service = TaggingService()
 
-    def create_database(self, database_name, kms_key_id):
+    def create_database(self, database_name, kms_key_id, tags):
         database = TimestreamDatabase(self.region_name, database_name, kms_key_id)
         self.databases[database_name] = database
+        self.tagging_service.tag_resource(database.arn, tags)
         return database
 
     def delete_database(self, database_name):
         del self.databases[database_name]
 
     def describe_database(self, database_name):
+        if database_name not in self.databases:
+            raise ResourceNotFound(f"The database {database_name} does not exist.")
         return self.databases[database_name]
 
     def list_databases(self):
@@ -103,9 +116,10 @@ class TimestreamWriteBackend(BaseBackend):
         database.update(kms_key_id=kms_key_id)
         return database
 
-    def create_table(self, database_name, table_name, retention_properties):
+    def create_table(self, database_name, table_name, retention_properties, tags):
         database = self.describe_database(database_name)
         table = database.create_table(table_name, retention_properties)
+        self.tagging_service.tag_resource(table.arn, tags)
         return table
 
     def delete_table(self, database_name, table_name):
@@ -147,6 +161,15 @@ class TimestreamWriteBackend(BaseBackend):
             ]
         }
 
+    def list_tags_for_resource(self, resource_arn):
+        return self.tagging_service.list_tags_for_resource(resource_arn)
+
+    def tag_resource(self, resource_arn, tags):
+        self.tagging_service.tag_resource(resource_arn, tags)
+
+    def untag_resource(self, resource_arn, tag_keys):
+        self.tagging_service.untag_resource_using_names(resource_arn, tag_keys)
+
     def reset(self):
         region_name = self.region_name
         self.__dict__ = {}

moto/timestreamwrite/responses.py

@@ -16,8 +16,9 @@ class TimestreamWriteResponse(BaseResponse):
     def create_database(self):
         database_name = self._get_param("DatabaseName")
         kms_key_id = self._get_param("KmsKeyId")
+        tags = self._get_param("Tags")
         database = self.timestreamwrite_backend.create_database(
-            database_name=database_name, kms_key_id=kms_key_id
+            database_name=database_name, kms_key_id=kms_key_id, tags=tags
         )
         return json.dumps(dict(Database=database.description()))
 
@@ -49,8 +50,9 @@ class TimestreamWriteResponse(BaseResponse):
         database_name = self._get_param("DatabaseName")
         table_name = self._get_param("TableName")
         retention_properties = self._get_param("RetentionProperties")
+        tags = self._get_param("Tags")
         table = self.timestreamwrite_backend.create_table(
-            database_name, table_name, retention_properties
+            database_name, table_name, retention_properties, tags
         )
         return json.dumps(dict(Table=table.description()))
 
@@ -97,3 +99,20 @@ class TimestreamWriteResponse(BaseResponse):
     def describe_endpoints(self):
         resp = self.timestreamwrite_backend.describe_endpoints()
         return json.dumps(resp)
+
+    def list_tags_for_resource(self):
+        resource_arn = self._get_param("ResourceARN")
+        tags = self.timestreamwrite_backend.list_tags_for_resource(resource_arn)
+        return json.dumps(tags)
+
+    def tag_resource(self):
+        resource_arn = self._get_param("ResourceARN")
+        tags = self._get_param("Tags")
+        self.timestreamwrite_backend.tag_resource(resource_arn, tags)
+        return "{}"
+
+    def untag_resource(self):
+        resource_arn = self._get_param("ResourceARN")
+        tag_keys = self._get_param("TagKeys")
+        self.timestreamwrite_backend.untag_resource(resource_arn, tag_keys)
+        return "{}"

tests/terraform-tests.success.txt

@@ -131,6 +131,8 @@ TestAccAWSSQSQueuePolicy
 TestAccAWSSSMDocument
 TestAccAWSSsmDocumentDataSource
 TestAccAWSSsmParameterDataSource
+TestAccAWSTimestreamWriteDatabase
+TestAccAWSTimestreamWriteTable
 TestAccAWSUserGroupMembership
 TestAccAWSUserPolicyAttachment
 TestAccAWSUserSSHKey

tests/test_timestreamwrite/test_timestreamwrite_database.py

@@ -1,5 +1,8 @@
 import boto3
+import pytest
 import sure  # noqa # pylint: disable=unused-import
+
+from botocore.exceptions import ClientError
 from moto import mock_timestreamwrite
 from moto.core import ACCOUNT_ID
 
@@ -15,7 +18,9 @@ def test_create_database_simple():
     )
     database.should.have.key("DatabaseName").equals("mydatabase")
     database.should.have.key("TableCount").equals(0)
-    database.shouldnt.have.key("KmsKeyId")
+    database.should.have.key("KmsKeyId").equals(
+        f"arn:aws:kms:us-east-1:{ACCOUNT_ID}:key/default_key"
+    )
 
 
 @mock_timestreamwrite
@@ -51,6 +56,16 @@ def test_describe_database():
     database.should.have.key("KmsKeyId").equal("mykey")
 
 
+@mock_timestreamwrite
+def test_describe_unknown_database():
+    ts = boto3.client("timestream-write", region_name="us-east-1")
+    with pytest.raises(ClientError) as exc:
+        ts.describe_database(DatabaseName="unknown")
+    err = exc.value.response["Error"]
+    err["Code"].should.equal("ResourceNotFoundException")
+    err["Message"].should.equal("The database unknown does not exist.")
+
+
 @mock_timestreamwrite
 def test_list_databases():
     ts = boto3.client("timestream-write", region_name="us-east-1")
@@ -73,6 +88,7 @@ def test_list_databases():
             "Arn": f"arn:aws:timestream:us-east-1:{ACCOUNT_ID}:database/db_without",
             "DatabaseName": "db_without",
             "TableCount": 0,
+            "KmsKeyId": f"arn:aws:kms:us-east-1:{ACCOUNT_ID}:key/default_key",
         }
     )
 

tests/test_timestreamwrite/test_timestreamwrite_table.py

@@ -1,6 +1,9 @@
 import time
 import boto3
+import pytest
 import sure  # noqa # pylint: disable=unused-import
+
+from botocore.exceptions import ClientError
 from moto import mock_timestreamwrite, settings
 from moto.core import ACCOUNT_ID
 
@@ -46,7 +49,12 @@ def test_create_table_without_retention_properties():
     table.should.have.key("TableName").equal("mytable")
     table.should.have.key("DatabaseName").equal("mydatabase")
     table.should.have.key("TableStatus").equal("ACTIVE")
-    table.shouldnt.have.key("RetentionProperties")
+    table.should.have.key("RetentionProperties").equals(
+        {
+            "MemoryStoreRetentionPeriodInHours": 123,
+            "MagneticStoreRetentionPeriodInDays": 123,
+        }
+    )
 
 
 @mock_timestreamwrite
@@ -78,6 +86,18 @@ def test_describe_table():
     )
 
 
+@mock_timestreamwrite
+def test_describe_unknown_database():
+    ts = boto3.client("timestream-write", region_name="us-east-1")
+    ts.create_database(DatabaseName="mydatabase")
+
+    with pytest.raises(ClientError) as exc:
+        ts.describe_table(DatabaseName="mydatabase", TableName="unknown")
+    err = exc.value.response["Error"]
+    err["Code"].should.equal("ResourceNotFoundException")
+    err["Message"].should.equal("The table unknown does not exist.")
+
+
 @mock_timestreamwrite
 def test_create_multiple_tables():
     ts = boto3.client("timestream-write", region_name="us-east-1")

tests/test_timestreamwrite/test_timestreamwrite_tagging.py

@@ -0,0 +1,91 @@
+import boto3
+
+from moto import mock_timestreamwrite
+
+
+@mock_timestreamwrite
+def test_list_tagging_for_table_without_tags():
+    ts = boto3.client("timestream-write", region_name="us-east-1")
+    ts.create_database(DatabaseName="mydatabase")
+
+    resp = ts.create_table(
+        DatabaseName="mydatabase",
+        TableName="mytable",
+        RetentionProperties={
+            "MemoryStoreRetentionPeriodInHours": 7,
+            "MagneticStoreRetentionPeriodInDays": 42,
+        },
+    )
+    table_arn = resp["Table"]["Arn"]
+
+    resp = ts.list_tags_for_resource(ResourceARN=table_arn)
+    resp.should.have.key("Tags").equals([])
+
+
+@mock_timestreamwrite
+def test_list_tagging_for_table_with_tags():
+    ts = boto3.client("timestream-write", region_name="us-east-1")
+    ts.create_database(DatabaseName="mydatabase")
+
+    resp = ts.create_table(
+        DatabaseName="mydatabase",
+        TableName="mytable",
+        RetentionProperties={
+            "MemoryStoreRetentionPeriodInHours": 7,
+            "MagneticStoreRetentionPeriodInDays": 42,
+        },
+        Tags=[{"Key": "k1", "Value": "v1"}],
+    )
+    table_arn = resp["Table"]["Arn"]
+
+    resp = ts.list_tags_for_resource(ResourceARN=table_arn)
+    resp.should.have.key("Tags").equals([{"Key": "k1", "Value": "v1"}])
+
+
+@mock_timestreamwrite
+def test_list_tagging_for_database_without_tags():
+    ts = boto3.client("timestream-write", region_name="us-east-1")
+    db_arn = ts.create_database(DatabaseName="mydatabase")["Database"]["Arn"]
+
+    resp = ts.list_tags_for_resource(ResourceARN=db_arn)
+    resp.should.have.key("Tags").equals([])
+
+
+@mock_timestreamwrite
+def test_list_tagging_for_database_with_tags():
+    ts = boto3.client("timestream-write", region_name="us-east-1")
+    db_arn = ts.create_database(
+        DatabaseName="mydatabase", Tags=[{"Key": "k1", "Value": "v1"}]
+    )["Database"]["Arn"]
+
+    resp = ts.list_tags_for_resource(ResourceARN=db_arn)
+    resp.should.have.key("Tags").equals([{"Key": "k1", "Value": "v1"}])
+
+
+@mock_timestreamwrite
+def test_tag_and_untag_database():
+    ts = boto3.client("timestream-write", region_name="us-east-1")
+    db_arn = ts.create_database(
+        DatabaseName="mydatabase", Tags=[{"Key": "k1", "Value": "v1"}]
+    )["Database"]["Arn"]
+
+    ts.tag_resource(
+        ResourceARN=db_arn,
+        Tags=[{"Key": "k2", "Value": "v2"}, {"Key": "k3", "Value": "v3"}],
+    )
+
+    resp = ts.list_tags_for_resource(ResourceARN=db_arn)
+    resp.should.have.key("Tags").equals(
+        [
+            {"Key": "k1", "Value": "v1"},
+            {"Key": "k2", "Value": "v2"},
+            {"Key": "k3", "Value": "v3"},
+        ]
+    )
+
+    ts.untag_resource(ResourceARN=db_arn, TagKeys=["k2"])
+
+    resp = ts.list_tags_for_resource(ResourceARN=db_arn)
+    resp.should.have.key("Tags").equals(
+        [{"Key": "k1", "Value": "v1"}, {"Key": "k3", "Value": "v3"}]
+    )