Add incremental export and cross account export functionality in DynamoDBToS3Operator

airflow/providers/amazon/aws/transfers/dynamodb_to_s3.py

@@ -32,6 +32,7 @@
 from airflow.providers.amazon.aws.hooks.dynamodb import DynamoDBHook
 from airflow.providers.amazon.aws.hooks.s3 import S3Hook
 from airflow.providers.amazon.aws.transfers.base import AwsToAwsBaseOperator
+from airflow.utils.helpers import prune_dict
 
 if TYPE_CHECKING:
     from airflow.utils.context import Context
@@ -89,10 +90,13 @@ class DynamoDBToS3Operator(AwsToAwsBaseOperator):
         <https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/dynamodb.html#DynamoDB.Table.scan>
     :param s3_key_prefix: Prefix of s3 object key
     :param process_func: How we transform a dynamodb item to bytes. By default, we dump the json
+    :param point_in_time_export: Boolean value indicating the operator to use 'scan' or 'point in time export'
     :param export_time: Time in the past from which to export table data, counted in seconds from the start of
      the Unix epoch. The table export will be a snapshot of the table's state at this point in time.
     :param export_format: The format for the exported data. Valid values for ExportFormat are DYNAMODB_JSON
      or ION.
+    :param export_table_to_point_in_time_kwargs: extra parameters for the boto3
+    `export_table_to_point_in_time` function all. e.g. `ExportType`, `IncrementalExportSpecification`
     :param check_interval: The amount of time in seconds to wait between attempts. Only if ``export_time`` is
         provided.
     :param max_attempts: The maximum number of attempts to be made. Only if ``export_time`` is provided.
@@ -107,25 +111,29 @@ class DynamoDBToS3Operator(AwsToAwsBaseOperator):
         "s3_key_prefix",
         "export_time",
         "export_format",
+        "export_table_to_point_in_time_kwargs",
         "check_interval",
         "max_attempts",
     )
 
     template_fields_renderers = {
         "dynamodb_scan_kwargs": "json",
+        "export_table_to_point_in_time_kwargs": "json",
     }
 
     def __init__(
         self,
         *,
         dynamodb_table_name: str,
         s3_bucket_name: str,
-        file_size: int,
+        file_size: int = 1000,
         dynamodb_scan_kwargs: dict[str, Any] | None = None,
         s3_key_prefix: str = "",
         process_func: Callable[[dict[str, Any]], bytes] = _convert_item_to_json_bytes,
+        point_in_time_export: bool = False,
         export_time: datetime | None = None,
         export_format: str = "DYNAMODB_JSON",
+        export_table_to_point_in_time_kwargs: dict | None = None,
         check_interval: int = 30,
         max_attempts: int = 60,
         **kwargs,
@@ -137,8 +145,10 @@ def __init__(
         self.dynamodb_scan_kwargs = dynamodb_scan_kwargs
         self.s3_bucket_name = s3_bucket_name
         self.s3_key_prefix = s3_key_prefix
+        self.point_in_time_export = point_in_time_export
         self.export_time = export_time
         self.export_format = export_format
+        self.export_table_to_point_in_time_kwargs = export_table_to_point_in_time_kwargs
         self.check_interval = check_interval
         self.max_attempts = max_attempts
 
@@ -148,29 +158,50 @@ def hook(self):
         return DynamoDBHook(aws_conn_id=self.source_aws_conn_id)
 
     def execute(self, context: Context) -> None:
-        if self.export_time:
+        # There are 2 separate export to point in time configuration:
+        # 1. Full export, which takes the export_time arg.
+        # 2. Incremental export, which takes the incremental_export_... args
+        # Hence export time could not be used as the proper indicator for the `_export_table_to_point_in_time`
+        # function. This change introduces a new boolean, as the indicator for whether the operator scans
+        # and export entire data or using the point in time functionality.
+        if self.point_in_time_export or self.export_time:
             self._export_table_to_point_in_time()
         else:
             self._export_entire_data()
 
     def _export_table_to_point_in_time(self):
         """
-        Export data from start of epoc till `export_time`.
+        Export data to point in time.
 
+        Full export exports data from start of epoc till `export_time`.
         Table export will be a snapshot of the table's state at this point in time.
+
+        Incremental export exports the data from a specific datetime to a specific datetime
+
+
+        Note: S3BucketOwner is a required parameter when exporting to a S3 bucket in another account.
         """
         if self.export_time and self.export_time > datetime.now(self.export_time.tzinfo):
             raise ValueError("The export_time parameter cannot be a future time.")
 
         client = self.hook.conn.meta.client
         table_description = client.describe_table(TableName=self.dynamodb_table_name)
-        response = client.export_table_to_point_in_time(
-            TableArn=table_description.get("Table", {}).get("TableArn"),
-            ExportTime=self.export_time,
-            S3Bucket=self.s3_bucket_name,
-            S3Prefix=self.s3_key_prefix,
-            ExportFormat=self.export_format,
-        )
+
+        export_table_to_point_in_time_base_args = {
+            "TableArn": table_description.get("Table", {}).get("TableArn"),
+            "ExportTime": self.export_time,
+            "S3Bucket": self.s3_bucket_name,
+            "S3Prefix": self.s3_key_prefix,
+            "ExportFormat": self.export_format,
+        }
+        export_table_to_point_in_time_args = {
+            **export_table_to_point_in_time_base_args,
+            **self.export_table_to_point_in_time_kwargs,
+        }
+
+        args_filtered = prune_dict(export_table_to_point_in_time_args)
+
+        response = client.export_table_to_point_in_time(**args_filtered)
         waiter = self.hook.get_waiter("export_table")
         export_arn = response.get("ExportDescription", {}).get("ExportArn")
         waiter.wait(

docs/apache-airflow-providers-amazon/transfer/dynamodb_to_s3.rst

@@ -63,13 +63,23 @@ To parallelize the replication, users can create multiple ``DynamoDBToS3Operator
     :start-after: [START howto_transfer_dynamodb_to_s3_segmented]
     :end-before: [END howto_transfer_dynamodb_to_s3_segmented]
 
-Users can also pass in ``export_time`` param to ``DynamoDBToS3Operator`` to recover data from a point in time.
+Users can also pass in ``point_in_time_export`` boolean param to ``DynamoDBToS3Operator`` to recover data from a point in time.
+
+Full export example usage:
+
+.. exampleinclude:: /../../tests/system/providers/amazon/aws/example_dynamodb_to_s3.py
+    :language: python
+    :dedent: 4
+    :start-after: [START howto_transfer_dynamodb_to_s3_in_some_point_in_time_full_export]
+    :end-before: [END howto_transfer_dynamodb_to_s3_in_some_point_in_time_full_export]
+
+Incremental export example usage:
 
 .. exampleinclude:: /../../tests/system/providers/amazon/aws/example_dynamodb_to_s3.py
     :language: python
     :dedent: 4
-    :start-after: [START howto_transfer_dynamodb_to_s3_in_some_point_in_time]
-    :end-before: [END howto_transfer_dynamodb_to_s3_in_some_point_in_time]
+    :start-after: [START howto_transfer_dynamodb_to_s3_in_some_point_in_time_incremental_export]
+    :end-before: [END howto_transfer_dynamodb_to_s3_in_some_point_in_time_incremental_export]
 
 Reference
 ---------

tests/providers/amazon/aws/transfers/test_dynamodb_to_s3.py

@@ -348,6 +348,7 @@ def test_dynamodb_execute_calling_export_table_to_point_in_time(self, _export_ta
             dynamodb_table_name="airflow_rocks",
             s3_bucket_name="airflow-bucket",
             file_size=4000,
+            point_in_time_export=True,
             export_time=datetime(year=1983, month=1, day=1),
         )
         dynamodb_to_s3_operator.execute(context={})
@@ -362,5 +363,6 @@ def test_dynamodb_with_future_date(self):
                 dynamodb_table_name="airflow_rocks",
                 s3_bucket_name="airflow-bucket",
                 file_size=4000,
+                point_in_time_export=True,
                 export_time=datetime(year=3000, month=1, day=1),
             ).execute(context={})

tests/system/providers/amazon/aws/example_dynamodb_to_s3.py

@@ -92,6 +92,15 @@ def get_export_time(table_name: str):
     return r["ContinuousBackupsDescription"]["PointInTimeRecoveryDescription"]["EarliestRestorableDateTime"]
 
 
+@task
+def get_latest_export_time(table_name: str):
+    r = boto3.client("dynamodb").describe_continuous_backups(
+        TableName=table_name,
+    )
+
+    return r["ContinuousBackupsDescription"]["PointInTimeRecoveryDescription"]["LatestRestorableDateTime"]
+
+
 @task
 def wait_for_bucket(s3_bucket_name):
     waiter = boto3.client("s3").get_waiter("bucket_exists")
@@ -162,18 +171,39 @@ def delete_dynamodb_table(table_name: str):
     # [END howto_transfer_dynamodb_to_s3_segmented]
 
     export_time = get_export_time(table_name)
-    # [START howto_transfer_dynamodb_to_s3_in_some_point_in_time]
-    backup_db_to_point_in_time = DynamoDBToS3Operator(
-        task_id="backup_db_to_point_in_time",
+    latest_export_time = get_latest_export_time(table_name)
+    # [START howto_transfer_dynamodb_to_s3_in_some_point_in_time_full_export]
+    backup_db_to_point_in_time_full_export = DynamoDBToS3Operator(
+        task_id="backup_db_to_point_in_time_full_export",
         dynamodb_table_name=table_name,
-        file_size=1000,
         s3_bucket_name=bucket_name,
+        point_in_time_export=True,
         export_time=export_time,
         s3_key_prefix=f"{S3_KEY_PREFIX}-3-",
     )
-    # [END howto_transfer_dynamodb_to_s3_in_some_point_in_time]
+    # [END howto_transfer_dynamodb_to_s3_in_some_point_in_time_full_export]
+
+    # [START howto_transfer_dynamodb_to_s3_in_some_point_in_time_incremental_export]
+    backup_db_to_point_in_time_incremental_export = DynamoDBToS3Operator(
+        task_id="backup_db_to_point_in_time_incremental_export",
+        dynamodb_table_name=table_name,
+        s3_bucket_name=bucket_name,
+        point_in_time_export=True,
+        s3_key_prefix=f"{S3_KEY_PREFIX}-4-",
+        export_table_to_point_in_time_kwargs={
+            "ExportType": "INCREMENTAL_EXPORT",
+            "IncrementalExportSpecification": {
+                "ExportFromTime": export_time,
+                "ExportToTime": latest_export_time,
+                "ExportViewType": "NEW_AND_OLD_IMAGES",
+            },
+        },
+    )
+    # [END howto_transfer_dynamodb_to_s3_in_some_point_in_time_incremental_export]
+
     # This operation can take a long time to complete
-    backup_db_to_point_in_time.max_attempts = 90
+    backup_db_to_point_in_time_full_export.max_attempts = 90
+    backup_db_to_point_in_time_incremental_export.max_attempts = 90
 
     delete_table = delete_dynamodb_table(table_name=table_name)
 
@@ -195,7 +225,8 @@ def delete_dynamodb_table(table_name: str):
         backup_db_segment_1,
         backup_db_segment_2,
         export_time,
-        backup_db_to_point_in_time,
+        backup_db_to_point_in_time_full_export,
+        backup_db_to_point_in_time_incremental_export,
         # TEST TEARDOWN
         delete_table,
         delete_bucket,