From 578de3d47f1ff514f42b19d934b1d5480f31dda6 Mon Sep 17 00:00:00 2001 From: Bert Blommers Date: Mon, 2 May 2022 20:15:19 +0000 Subject: [PATCH] TimestreamWrite - support MagneticStoreWriteProps-argument (#5088) --- moto/timestreamwrite/models.py | 56 +++++++++++++--- moto/timestreamwrite/responses.py | 17 ++++- .../terraform-tests.failures.txt | 1 - .../terraform-tests.success.txt | 1 + .../test_timestreamwrite_table.py | 66 +++++++++++++++++++ 5 files changed, 128 insertions(+), 13 deletions(-) diff --git a/moto/timestreamwrite/models.py b/moto/timestreamwrite/models.py index f3fb877a1..a8ed0b5d2 100644 --- a/moto/timestreamwrite/models.py +++ b/moto/timestreamwrite/models.py @@ -5,7 +5,14 @@ from .exceptions import ResourceNotFound class TimestreamTable(BaseModel): - def __init__(self, region_name, table_name, db_name, retention_properties): + def __init__( + self, + region_name, + table_name, + db_name, + retention_properties, + magnetic_store_write_properties, + ): self.region_name = region_name self.name = table_name self.db_name = db_name @@ -13,10 +20,13 @@ class TimestreamTable(BaseModel): "MemoryStoreRetentionPeriodInHours": 123, "MagneticStoreRetentionPeriodInDays": 123, } + self.magnetic_store_write_properties = magnetic_store_write_properties or {} self.records = [] - def update(self, retention_properties): + def update(self, retention_properties, magnetic_store_write_properties): self.retention_properties = retention_properties + if magnetic_store_write_properties is not None: + self.magnetic_store_write_properties = magnetic_store_write_properties def write_records(self, records): self.records.extend(records) @@ -32,6 +42,7 @@ class TimestreamTable(BaseModel): "DatabaseName": self.db_name, "TableStatus": "ACTIVE", "RetentionProperties": self.retention_properties, + "MagneticStoreWriteProperties": self.magnetic_store_write_properties, } @@ -47,23 +58,31 @@ class TimestreamDatabase(BaseModel): def update(self, kms_key_id): self.kms_key_id = kms_key_id - def create_table(self, table_name, retention_properties): + def create_table( + self, table_name, retention_properties, magnetic_store_write_properties + ): table = TimestreamTable( region_name=self.region_name, table_name=table_name, db_name=self.name, retention_properties=retention_properties, + magnetic_store_write_properties=magnetic_store_write_properties, ) self.tables[table_name] = table return table - def update_table(self, table_name, retention_properties): + def update_table( + self, table_name, retention_properties, magnetic_store_write_properties + ): table = self.tables[table_name] - table.update(retention_properties=retention_properties) + table.update( + retention_properties=retention_properties, + magnetic_store_write_properties=magnetic_store_write_properties, + ) return table def delete_table(self, table_name): - del self.tables[table_name] + self.tables.pop(table_name, None) def describe_table(self, table_name): if table_name not in self.tables: @@ -116,9 +135,18 @@ class TimestreamWriteBackend(BaseBackend): database.update(kms_key_id=kms_key_id) return database - def create_table(self, database_name, table_name, retention_properties, tags): + def create_table( + self, + database_name, + table_name, + retention_properties, + tags, + magnetic_store_write_properties, + ): database = self.describe_database(database_name) - table = database.create_table(table_name, retention_properties) + table = database.create_table( + table_name, retention_properties, magnetic_store_write_properties + ) self.tagging_service.tag_resource(table.arn, tags) return table @@ -136,9 +164,17 @@ class TimestreamWriteBackend(BaseBackend): tables = database.list_tables() return tables - def update_table(self, database_name, table_name, retention_properties): + def update_table( + self, + database_name, + table_name, + retention_properties, + magnetic_store_write_properties, + ): database = self.describe_database(database_name) - table = database.update_table(table_name, retention_properties) + table = database.update_table( + table_name, retention_properties, magnetic_store_write_properties + ) return table def write_records(self, database_name, table_name, records): diff --git a/moto/timestreamwrite/responses.py b/moto/timestreamwrite/responses.py index 70fbcc692..232202029 100644 --- a/moto/timestreamwrite/responses.py +++ b/moto/timestreamwrite/responses.py @@ -51,8 +51,15 @@ class TimestreamWriteResponse(BaseResponse): table_name = self._get_param("TableName") retention_properties = self._get_param("RetentionProperties") tags = self._get_param("Tags") + magnetic_store_write_properties = self._get_param( + "MagneticStoreWriteProperties" + ) table = self.timestreamwrite_backend.create_table( - database_name, table_name, retention_properties, tags + database_name, + table_name, + retention_properties, + tags, + magnetic_store_write_properties, ) return json.dumps(dict(Table=table.description())) @@ -77,8 +84,14 @@ class TimestreamWriteResponse(BaseResponse): database_name = self._get_param("DatabaseName") table_name = self._get_param("TableName") retention_properties = self._get_param("RetentionProperties") + magnetic_store_write_properties = self._get_param( + "MagneticStoreWriteProperties" + ) table = self.timestreamwrite_backend.update_table( - database_name, table_name, retention_properties + database_name, + table_name, + retention_properties, + magnetic_store_write_properties, ) return json.dumps(dict(Table=table.description())) diff --git a/tests/terraformtests/terraform-tests.failures.txt b/tests/terraformtests/terraform-tests.failures.txt index 427668434..4993b1804 100644 --- a/tests/terraformtests/terraform-tests.failures.txt +++ b/tests/terraformtests/terraform-tests.failures.txt @@ -59,7 +59,6 @@ TestAccSQSQueuePolicy TestAccSSMDocument TestAccSsmDocumentDataSource TestAccSsmParameterDataSource -TestAccTimestreamWriteTable TestAccDataSourceLambdaLayerVersion TestAccDataSourceLambdaInvocation TestAccDataSourceNetworkInterface_ diff --git a/tests/terraformtests/terraform-tests.success.txt b/tests/terraformtests/terraform-tests.success.txt index 2afba913a..4ba891f07 100644 --- a/tests/terraformtests/terraform-tests.success.txt +++ b/tests/terraformtests/terraform-tests.success.txt @@ -128,3 +128,4 @@ sqs: - TestAccSQSQueue_FIFOQueue_ timestreamwrite: - TestAccTimestreamWriteDatabase + - TestAccTimestreamWriteTable diff --git a/tests/test_timestreamwrite/test_timestreamwrite_table.py b/tests/test_timestreamwrite/test_timestreamwrite_table.py index 411f1d210..33ad966da 100644 --- a/tests/test_timestreamwrite/test_timestreamwrite_table.py +++ b/tests/test_timestreamwrite/test_timestreamwrite_table.py @@ -36,6 +36,38 @@ def test_create_table(): ) +@mock_timestreamwrite +def test_create_table__with_magnetic_store_write_properties(): + ts = boto3.client("timestream-write", region_name="us-east-1") + ts.create_database(DatabaseName="mydatabase") + + resp = ts.create_table( + DatabaseName="mydatabase", + TableName="mytable", + MagneticStoreWriteProperties={ + "EnableMagneticStoreWrites": True, + "MagneticStoreRejectedDataLocation": { + "S3Configuration": {"BucketName": "hithere"} + }, + }, + ) + table = resp["Table"] + table.should.have.key("Arn").equal( + f"arn:aws:timestream:us-east-1:{ACCOUNT_ID}:database/mydatabase/table/mytable" + ) + table.should.have.key("TableName").equal("mytable") + table.should.have.key("DatabaseName").equal("mydatabase") + table.should.have.key("TableStatus").equal("ACTIVE") + table.should.have.key("MagneticStoreWriteProperties").should.equal( + { + "EnableMagneticStoreWrites": True, + "MagneticStoreRejectedDataLocation": { + "S3Configuration": {"BucketName": "hithere"} + }, + } + ) + + @mock_timestreamwrite def test_create_table_without_retention_properties(): ts = boto3.client("timestream-write", region_name="us-east-1") @@ -188,6 +220,40 @@ def test_update_table(): ) +@mock_timestreamwrite +def test_update_table__with_magnetic_store_write_properties(): + ts = boto3.client("timestream-write", region_name="us-east-1") + ts.create_database(DatabaseName="mydatabase") + + ts.create_table(DatabaseName="mydatabase", TableName="mytable") + + resp = ts.update_table( + DatabaseName="mydatabase", + TableName="mytable", + MagneticStoreWriteProperties={ + "EnableMagneticStoreWrites": True, + "MagneticStoreRejectedDataLocation": { + "S3Configuration": {"BucketName": "hithere"} + }, + }, + ) + table = resp["Table"] + table.should.have.key("Arn").equal( + f"arn:aws:timestream:us-east-1:{ACCOUNT_ID}:database/mydatabase/table/mytable" + ) + table.should.have.key("TableName").equal("mytable") + table.should.have.key("DatabaseName").equal("mydatabase") + table.should.have.key("TableStatus").equal("ACTIVE") + table.should.have.key("MagneticStoreWriteProperties").should.equal( + { + "EnableMagneticStoreWrites": True, + "MagneticStoreRejectedDataLocation": { + "S3Configuration": {"BucketName": "hithere"} + }, + } + ) + + @mock_timestreamwrite def test_write_records(): # The query-feature is not available at the moment,