From afa34ffd8d727c3e67d0fa8d515f498ba165fbe1 Mon Sep 17 00:00:00 2001 From: Bert Blommers Date: Sat, 26 Mar 2022 20:25:56 -0100 Subject: [PATCH] TimestreamWrite - improvements (#4971) --- .github/workflows/build.yml | 4 +- IMPLEMENTATION_COVERAGE.md | 8 +- docs/docs/services/timestream-write.rst | 6 +- moto/timestreamwrite/exceptions.py | 8 ++ moto/timestreamwrite/models.py | 31 ++++++- moto/timestreamwrite/responses.py | 23 ++++- tests/terraform-tests.success.txt | 2 + .../test_timestreamwrite_database.py | 18 +++- .../test_timestreamwrite_table.py | 22 ++++- .../test_timestreamwrite_tagging.py | 91 +++++++++++++++++++ 10 files changed, 196 insertions(+), 17 deletions(-) create mode 100644 tests/test_timestreamwrite/test_timestreamwrite_tagging.py diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index db97ae15a..4f2ad650b 100644 --- a/.github/workflows/build.yml +++ b/.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: | diff --git a/IMPLEMENTATION_COVERAGE.md b/IMPLEMENTATION_COVERAGE.md index 0fb569ed7..5f788d3ce 100644 --- a/IMPLEMENTATION_COVERAGE.md +++ b/IMPLEMENTATION_COVERAGE.md @@ -5472,7 +5472,7 @@ ## timestream-write
-80% implemented +100% implemented - [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 diff --git a/docs/docs/services/timestream-write.rst b/docs/docs/services/timestream-write.rst index bec942e8d..242230a2d 100644 --- a/docs/docs/services/timestream-write.rst +++ b/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 diff --git a/moto/timestreamwrite/exceptions.py b/moto/timestreamwrite/exceptions.py index 3017dd382..f1598bd65 100644 --- a/moto/timestreamwrite/exceptions.py +++ b/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) diff --git a/moto/timestreamwrite/models.py b/moto/timestreamwrite/models.py index bb87c3cba..f3fb877a1 100644 --- a/moto/timestreamwrite/models.py +++ b/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__ = {} diff --git a/moto/timestreamwrite/responses.py b/moto/timestreamwrite/responses.py index 9da11edfa..70fbcc692 100644 --- a/moto/timestreamwrite/responses.py +++ b/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 "{}" diff --git a/tests/terraform-tests.success.txt b/tests/terraform-tests.success.txt index e2aea1f4f..d51728bba 100644 --- a/tests/terraform-tests.success.txt +++ b/tests/terraform-tests.success.txt @@ -131,6 +131,8 @@ TestAccAWSSQSQueuePolicy TestAccAWSSSMDocument TestAccAWSSsmDocumentDataSource TestAccAWSSsmParameterDataSource +TestAccAWSTimestreamWriteDatabase +TestAccAWSTimestreamWriteTable TestAccAWSUserGroupMembership TestAccAWSUserPolicyAttachment TestAccAWSUserSSHKey diff --git a/tests/test_timestreamwrite/test_timestreamwrite_database.py b/tests/test_timestreamwrite/test_timestreamwrite_database.py index 0b07565ce..cf003b52c 100644 --- a/tests/test_timestreamwrite/test_timestreamwrite_database.py +++ b/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", } ) diff --git a/tests/test_timestreamwrite/test_timestreamwrite_table.py b/tests/test_timestreamwrite/test_timestreamwrite_table.py index 8dad6ddae..411f1d210 100644 --- a/tests/test_timestreamwrite/test_timestreamwrite_table.py +++ b/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") diff --git a/tests/test_timestreamwrite/test_timestreamwrite_tagging.py b/tests/test_timestreamwrite/test_timestreamwrite_tagging.py new file mode 100644 index 000000000..a18ce8f22 --- /dev/null +++ b/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"}] + )