DynamoDB: Export Table to S3 (#7839)

Author: rita-cheung

AUTHORS.md

@@ -63,3 +63,4 @@ Moto is written by Steve Pulec with contributions from:
 * [Roman Rader](https://github.com/rrader/)
 * [Bryan Chen](https://github.com/bchen1116)
 * [Jonas Bulik](https://github.com/MrGreenTea)
+* [Rita Cheung](https://github.com/rita-cheung)

moto/core/utils.py

@@ -1,7 +1,7 @@
 import datetime
 import inspect
 import re
-from gzip import decompress
+from gzip import compress, decompress
 from typing import Any, Callable, Dict, List, Optional, Tuple
 from urllib.parse import ParseResult, urlparse
 
@@ -403,6 +403,10 @@ def gzip_decompress(body: bytes) -> bytes:
     return decompress(body)
 
 
+def gzip_compress(body: bytes) -> bytes:
+    return compress(body)
+
+
 ISO_REGION_DOMAINS = {
     "iso": "c2s.ic.gov",
     "isob": "sc2s.sgov.gov",

moto/dynamodb/models/__init__.py

@@ -37,6 +37,7 @@
     RestoredTable,
     Table,
 )
+from moto.dynamodb.models.table_export import TableExport
 from moto.dynamodb.models.table_import import TableImport
 from moto.dynamodb.parsing import partiql
 from moto.dynamodb.parsing.executors import UpdateExpressionExecutor
@@ -53,6 +54,7 @@ def __init__(self, region_name: str, account_id: str):
         self.tables: Dict[str, Table] = OrderedDict()
         self.backups: Dict[str, Backup] = OrderedDict()
         self.table_imports: Dict[str, TableImport] = {}
+        self.table_exports: Dict[str, TableExport] = {}
 
     @staticmethod
     def default_vpc_endpoint_service(
@@ -999,5 +1001,41 @@ def import_table(
     def describe_import(self, import_arn: str) -> TableImport:
         return self.table_imports[import_arn]
 
+    def export_table(
+        self,
+        s3_bucket: str,
+        s3_prefix: str,
+        table_arn: str,
+        export_format: str,
+        export_type: str,
+        s3_bucket_owner: str,
+    ) -> TableExport:
+        """Only ExportFormat=DYNAMODB_JSON is supported so far.
+        Only exports one file following DYNAMODB_JSON format to the s3 location. Other files aren't created.
+        Incremental export is also not supported.
+        """
+        table_export = TableExport(
+            s3_bucket=s3_bucket,
+            s3_prefix=s3_prefix,
+            region_name=self.region_name,
+            account_id=s3_bucket_owner if s3_bucket_owner else self.account_id,
+            table_arn=table_arn,
+            export_format=export_format,
+            export_type=export_type,
+        )
+        self.table_exports[table_export.arn] = table_export
+        table_export.start()
+        return table_export
+
+    def describe_export(self, export_arn: str) -> TableExport:
+        return self.table_exports[export_arn]
+
+    def list_exports(self, table_arn: str) -> List[TableExport]:
+        exports = []
+        for export_arn in self.table_exports:
+            if self.table_exports[export_arn].table_arn == table_arn:
+                exports.append(self.table_exports[export_arn])
+        return exports
+
 
 dynamodb_backends = BackendDict(DynamoDBBackend, "dynamodb")

moto/dynamodb/models/table_export.py

@@ -0,0 +1,112 @@
+import json
+from threading import Thread
+from typing import TYPE_CHECKING, Any, Dict, Optional
+from uuid import uuid4
+
+from moto.core.utils import gzip_compress
+from moto.utilities.utils import get_partition
+
+if TYPE_CHECKING:
+    from moto.dynamodb.models import Table
+    from moto.s3.models import S3Backend
+
+
+class TableExport(Thread):
+    def __init__(
+        self,
+        s3_bucket: str,
+        s3_prefix: str,
+        region_name: str,
+        account_id: str,
+        table_arn: str,
+        export_format: str,
+        export_type: str,
+    ):
+        super().__init__()
+        self.partition = get_partition(region_name)
+        self.table_arn = table_arn
+        self.arn = f"arn:{self.partition}:dynamodb:{region_name}:{account_id}:table/{table_arn}/import/{str(uuid4()).replace('-', '')}"
+        self.s3_bucket = s3_bucket
+        self.s3_prefix = s3_prefix
+        self.status = "IN_PROGRESS"
+        self.export_format = export_format
+        self.export_type = export_type
+        self.account_id = account_id
+        self.region_name = region_name
+
+        self.failure_code: Optional[str] = None
+        self.failure_message: Optional[str] = None
+        self.table_name: str = ""
+        self.item_count = 0
+        self.processed_bytes = 0
+        self.error_count = 0
+
+    def run(self) -> None:
+        from moto.dynamodb.models import dynamodb_backends
+
+        dynamodb_backend = dynamodb_backends[self.account_id][self.region_name]
+        try:
+            from moto.s3.models import s3_backends
+
+            s3_backend = s3_backends[self.account_id][self.partition]
+            s3_backend.buckets[self.s3_bucket]
+
+        except KeyError:
+            self.status = "FAILED"
+            self.failure_code = "S3NoSuchBucket"
+            self.failure_message = "The specified bucket does not exist"
+            return
+
+        for key in dynamodb_backend.tables:
+            if dynamodb_backend.tables[key].table_arn == self.table_arn:
+                self.table_name = key
+        if not self.table_name:
+            self.status = "FAILED"
+            self.failure_code = "DynamoDBTableNotFound"
+            self.failure_message = "The specified table does not exist"
+            return
+        table = dynamodb_backend.tables[self.table_name]
+        if (
+            table.continuous_backups["PointInTimeRecoveryDescription"][
+                "PointInTimeRecoveryStatus"
+            ]
+            != "ENABLED"
+        ):
+            self.status = "FAILED"
+            self.failure_code = "PointInTimeRecoveryUnavailable"
+            self.failure_message = "Point in time recovery not enabled for table"
+            return
+        try:
+            self._backup_to_s3_file(s3_backend, table)
+        except Exception as e:
+            self.status = "FAILED"
+            self.failure_code = "UNKNOWN"
+            self.failure_message = str(e)
+
+    def _backup_to_s3_file(self, s3_backend: "S3Backend", table: "Table") -> None:
+        backup = []
+        for item in table.all_items():
+            json_item = item.to_json()
+            backup.append(json_item)
+            self.processed_bytes += len(json_item)
+        self.item_count = len(backup)
+        content = gzip_compress(json.dumps(backup).encode("utf-8"))
+        s3_backend.put_object(
+            bucket_name=self.s3_bucket,
+            key_name=f"{self.s3_prefix}/AWSDynamoDB/{str(uuid4())}/data/{str(uuid4())}.gz",
+            value=content,
+        )
+
+        self.status = "COMPLETED" if self.error_count == 0 else "FAILED"
+
+    def response(self) -> Dict[str, Any]:
+        return {
+            "ExportArn": self.arn,
+            "ExportStatus": self.status,
+            "FailureCode": self.failure_code,
+            "FailureMessage": self.failure_message,
+            "ExportFormat": self.export_format,
+            "ExportType": self.export_type,
+            "ItemCount": self.item_count,
+            "BilledSizeBytes": self.processed_bytes,
+        }

moto/dynamodb/responses.py

@@ -1227,3 +1227,41 @@ def describe_import(self) -> str:
         import_arn = self.body["ImportArn"]
         import_table = self.dynamodb_backend.describe_import(import_arn)
         return json.dumps({"ImportTableDescription": import_table.response()})
+
+    def export_table_to_point_in_time(self) -> str:
+        table_arn = self.body["TableArn"]
+        s3_bucket = self.body["S3Bucket"]
+        s3_prefix = self.body["S3Prefix"]
+        export_type = self.body.get("ExportType", "FULL_EXPORT")
+        export_format = self.body.get("ExportFormat", "DYNAMO_JSON")
+        s3_bucket_owner = self.body.get("S3BucketOwner")
+
+        export_table = self.dynamodb_backend.export_table(
+            s3_bucket=s3_bucket,
+            s3_prefix=s3_prefix,
+            table_arn=table_arn,
+            export_type=export_type,
+            export_format=export_format,
+            s3_bucket_owner=s3_bucket_owner,
+        )
+
+        return dynamo_json_dump({"ExportDescription": export_table.response()})
+
+    def describe_export(self) -> str:
+        export_arn = self.body["ExportArn"]
+        export_table = self.dynamodb_backend.describe_export(export_arn)
+        return json.dumps({"ExportDescription": export_table.response()})
+
+    def list_exports(self) -> str:
+        table_arn = self.body["TableArn"]
+        exports = self.dynamodb_backend.list_exports(table_arn)
+        response = []
+        for export_table in exports:
+            response.append(
+                {
+                    "ExportArn": export_table.arn,
+                    "ExportStatus": export_table.status,
+                    "ExportType": export_table.export_type,
+                }
+            )
+        return json.dumps({"ExportSummaries": response})