From 9e0eccf3726ba6603084089af140364f2f0c73e9 Mon Sep 17 00:00:00 2001 From: "yuichi.nasukawa" Date: Fri, 31 Jul 2026 06:41:06 +0900 Subject: [PATCH 1/2] feat: add DynamoDB query support to DynamoDBRead Extend DynamoDBRead to use the DynamoDB query operation when partition_key/partition_value (and optionally sort_key/sort_value) are specified, instead of always scanning the whole table. Falls back to scan when no key is specified, so existing scenarios keep working unchanged. Closes #529 --- cliboa/scenario/extract/aws.py | 151 +++++++++++++++------- docs/modules/dynamodb_read.md | 42 +++++- tests/scenario/extract/test_aws.py | 200 ++++++++++++++++++++++++++++- 3 files changed, 344 insertions(+), 49 deletions(-) diff --git a/cliboa/scenario/extract/aws.py b/cliboa/scenario/extract/aws.py index 90ace413..4b1a670c 100644 --- a/cliboa/scenario/extract/aws.py +++ b/cliboa/scenario/extract/aws.py @@ -16,8 +16,11 @@ import os import re from decimal import Decimal +from typing import Any, Literal import boto3 +from boto3.dynamodb.conditions import Key +from pydantic import model_validator from cliboa.adapter.aws import S3Adapter from cliboa.scenario.aws import BaseAws, BaseS3 @@ -208,70 +211,94 @@ class DynamoDBRead(BaseAws): Download data from DynamoDB and save as a CSV or JSONL file """ - def __init__(self): - super().__init__() - self._table_name = None - self._dest_dir = "." - self._file_name = None - self._file_format = "csv" - - def table_name(self, table_name): - self._table_name = table_name + class Arguments(BaseAws.Arguments): + table_name: str + dest_dir: str = "." + file_name: str + file_format: Literal["csv", "jsonl"] = "csv" + partition_key: str | None = None + partition_value: Any = None + sort_key: str | None = None + sort_value: Any = None + + @model_validator(mode="before") + def check_key_conditions(cls, data: dict) -> dict: + if not isinstance(data, dict): + raise ValueError(f"arguments is not dict: {data}") + + partition_key_present = "partition_key" in data + partition_value_present = "partition_value" in data + if partition_key_present != partition_value_present: + raise InvalidParameter( + "Both 'partition_key' and 'partition_value' must be specified together." + ) - def dest_dir(self, dest_dir): - self._dest_dir = dest_dir + sort_key_present = "sort_key" in data + sort_value_present = "sort_value" in data + if sort_key_present != sort_value_present: + raise InvalidParameter( + "Both 'sort_key' and 'sort_value' must be specified together." + ) - def file_name(self, file_name): - self._file_name = file_name + if sort_key_present and not partition_key_present: + raise InvalidParameter( + "'sort_key'/'sort_value' require 'partition_key'/'partition_value' " + "to also be specified." + ) - def file_format(self, file_format): - if file_format not in ["csv", "jsonl"]: - raise InvalidParameter("file_format must be either 'csv' or 'jsonl'") - self._file_format = file_format + return data def execute(self, *args): """ - DynamoDBからデータをダウンロードし、指定されたフォーマットでファイルに保存します。 - """ - super().execute() + Download items from a DynamoDB table and save them to a CSV or JSONL file. - valid = EssentialParameters(self.__class__.__name__, [self._table_name, self._file_name]) - valid() - - os.makedirs(self._dest_dir, exist_ok=True) + If 'partition_key'/'partition_value' are specified, a query operation is used + (optionally narrowed further by 'sort_key'/'sort_value'). Otherwise, a scan + operation is used to retrieve the whole table, matching the prior behavior. + """ + os.makedirs(self.args.dest_dir, exist_ok=True) dynamodb = boto3.resource( "dynamodb", - aws_access_key_id=self._access_key, - aws_secret_access_key=self._secret_key, - region_name=self._region, + aws_access_key_id=self.args.access_key, + aws_secret_access_key=self.args.secret_key, + region_name=self.args.region, ) - table = dynamodb.Table(self._table_name) + table = dynamodb.Table(self.args.table_name) - file_path = os.path.join(self._dest_dir, self._file_name) - if self._file_format == "jsonl": - self._write_jsonl(self._scan_table(table), file_path) + if self.args.partition_key: + items = self._query_table(table) else: - self._write_csv(self._scan_table(table), file_path) + items = self._scan_table(table) - self._logger.info(f"Downloaded items from DynamoDB table {self._table_name} to {file_path}") + file_path = os.path.join(self.args.dest_dir, self.args.file_name) + if self.args.file_format == "jsonl": + self._write_jsonl(items, file_path) + else: + self._write_csv(items, file_path) - def _scan_table(self, table): + self.logger.info( + f"Downloaded items from DynamoDB table {self.args.table_name} to {file_path}" + ) + + def _paginate(self, operation, **kwargs): """ - DynamoDBテーブルをスキャンし、全アイテムを取得するジェネレータ関数。 + Generator that repeatedly calls a boto3 Table operation (scan or query), + following DynamoDB's ExclusiveStartKey/LastEvaluatedKey pagination. Args: - table (boto3.resources.factory.dynamodb.Table): スキャン対象のDynamoDBテーブル + operation: bound Table method to call, e.g. table.scan or table.query + **kwargs: extra arguments passed to the operation, e.g. KeyConditionExpression Yields: - dict: テーブルの各アイテム + dict: each item returned by the operation """ last_evaluated_key = None while True: if last_evaluated_key: - response = table.scan(ExclusiveStartKey=last_evaluated_key) + response = operation(ExclusiveStartKey=last_evaluated_key, **kwargs) else: - response = table.scan() + response = operation(**kwargs) for item in response["Items"]: yield item @@ -280,12 +307,42 @@ def _scan_table(self, table): if not last_evaluated_key: break + def _scan_table(self, table): + """ + Generator function that scans a DynamoDB table and retrieves all items. + + Args: + table (boto3.resources.factory.dynamodb.Table): DynamoDB table to scan + + Yields: + dict: each item from the table + """ + yield from self._paginate(table.scan) + + def _query_table(self, table): + """ + Generator function that queries a DynamoDB table by partition key + (and optionally sort key), retrieving all matching items. + + Args: + table (boto3.resources.factory.dynamodb.Table): DynamoDB table to query + + Yields: + dict: each matching item from the table + """ + key_condition = Key(self.args.partition_key).eq(self.args.partition_value) + if self.args.sort_key: + key_condition &= Key(self.args.sort_key).eq(self.args.sort_value) + + yield from self._paginate(table.query, KeyConditionExpression=key_condition) + def _write_jsonl(self, items, file_path): """ - アイテムをJSONL形式でファイルに書き込みます。 + Write items to a file in JSONL format. + Args: - items (iterator): 書き込むアイテムのイテレータ - file_path (str): 書き込み先のファイルパス + items (iterator): iterator of items to write + file_path (str): destination file path """ with open(file_path, "w") as f: for item in items: @@ -296,7 +353,7 @@ def _write_jsonl(self, items, file_path): def _json_serial(self, obj): """ - JSONシリアライズ関数 + JSON serialization helper for types not natively supported by json.dumps. """ if isinstance(obj, Decimal): return int(obj) if obj % 1 == 0 else float(obj) @@ -304,11 +361,11 @@ def _json_serial(self, obj): def _write_csv(self, items, file_path): """ - アイテムをCSV形式でファイルに書き込みます。 + Write items to a file in CSV format. Args: - items (iterator): 書き込むアイテムのイテレータ - file_path (str): 書き込み先のファイルパス + items (iterator): iterator of items to write + file_path (str): destination file path """ with open(file_path, "w", newline="") as f: writer = None @@ -319,7 +376,7 @@ def _write_csv(self, items, file_path): for key, value in item.items(): if isinstance(value, (dict, list)): - # ネストされた属性値はJSON形式に変換 + # Nested attribute values are converted to JSON item[key] = json.dumps( value, default=self._json_serial, sort_keys=False, ensure_ascii=False ) diff --git a/docs/modules/dynamodb_read.md b/docs/modules/dynamodb_read.md index edc7b32a..4ccd70d2 100644 --- a/docs/modules/dynamodb_read.md +++ b/docs/modules/dynamodb_read.md @@ -8,13 +8,18 @@ Reads data from a DynamoDB table and saves it as a CSV or JSONL file. |dest_dir|Output directory|No|"." (current directory)|If a non-existent directory path is specified, it will be automatically created.| |file_name|Output file name|Yes|None|| |file_format|Output file format|No|"csv"|Can be either "csv" or "jsonl".| +|partition_key|Partition key attribute name|No|None|If specified together with `partition_value`, a `query` operation is used instead of `scan`.| +|partition_value|Partition key value to match (equality)|No|None|Required together with `partition_key`.| +|sort_key|Sort key attribute name|No|None|Can only be specified together with `partition_key`/`partition_value`.| +|sort_value|Sort key value to match (equality)|No|None|Required together with `sort_key`.| |region|AWS region|No|None|If not specified, the default region will be used.| |access_key|AWS access key|No|None|If not specified, environment variables or IAM role will be used.| |secret_key|AWS secret key|No|None|If not specified, environment variables or IAM role will be used.| |profile|AWS profile|No|None|Section name of ~/.aws/config| -# Example +# Examples ```yaml +# Read the whole table (scan operation) scenario: step: class: DynamoDBRead @@ -26,8 +31,43 @@ scenario: region: us-west-2 ``` +```yaml +# Read items matching a partition key (query operation) +scenario: + step: + class: DynamoDBRead + arguments: + table_name: your_dynamodb_table + dest_dir: /path/to/destination + file_name: dynamodb_data.csv + file_format: csv + region: us-west-2 + partition_key: user_id + partition_value: "12345" +``` + +```yaml +# Read items matching a partition key and sort key (query operation) +scenario: + step: + class: DynamoDBRead + arguments: + table_name: your_dynamodb_table + dest_dir: /path/to/destination + file_name: dynamodb_data.csv + file_format: csv + region: us-west-2 + partition_key: user_id + partition_value: "12345" + sort_key: created_at + sort_value: "2026-01-01" +``` + # Notes +- When `partition_key`/`partition_value` are specified, a `query` operation is used for more efficient and lower-cost retrieval than `scan`. Otherwise, the whole table is read via `scan`, matching the previous behavior. +- `sort_key`/`sort_value` narrow a query further, but require `partition_key`/`partition_value` to also be specified. +- Only equality conditions are supported for `partition_value`/`sort_value`. Range conditions (`<`, `>=`, `BETWEEN`, etc.) are not supported yet. - Conversion to CSV might be complex for certain DynamoDB attribute types (sets, lists, maps, etc.). - If the output file already exists, it will be overwritten. - Partition and sort keys are not guaranteed to line up before other attributes. \ No newline at end of file diff --git a/tests/scenario/extract/test_aws.py b/tests/scenario/extract/test_aws.py index 5d21ada1..dccc309d 100644 --- a/tests/scenario/extract/test_aws.py +++ b/tests/scenario/extract/test_aws.py @@ -18,6 +18,8 @@ from decimal import Decimal from unittest.mock import Mock, patch +import pytest + from cliboa.adapter.aws import S3Adapter from cliboa.scenario.extract.aws import ( DynamoDBRead, @@ -26,6 +28,7 @@ S3DownloadFileDelete, S3FileExistsCheck, ) +from cliboa.util.exception import InvalidParameter from tests import BaseCliboaTest @@ -309,6 +312,201 @@ def test_execute_jsonl_without_nested_data(self, mock_boto_resource): self._run_test(mock_boto_resource, test_data, expected_jsonl, "jsonl") + @patch("boto3.resource") + def test_execute_query_with_partition_key_only(self, mock_boto_resource): + test_data = { + "Items": [ + {"id": "1", "name": "Item 1", "value": Decimal("100")}, + ], + "Count": 1, + "ScannedCount": 1, + "LastEvaluatedKey": None, + } + expected_csv = [ + ["id", "name", "value"], + ["1", "Item 1", "100"], + ] + + mock_table = mock_boto_resource.return_value.Table.return_value + mock_table.query.return_value = test_data + + with tempfile.TemporaryDirectory() as temp_dir: + instance = DynamoDBRead() + instance._set_arguments( + { + "table_name": "test_table", + "file_name": "output.csv", + "dest_dir": temp_dir, + "region": "us-east-1", + "partition_key": "id", + "partition_value": "1", + } + ) + instance.execute() + + output_file_path = os.path.join(temp_dir, instance.args.file_name) + self._verify_csv(output_file_path, expected_csv) + + mock_table.query.assert_called_once() + mock_table.scan.assert_not_called() + + @patch("boto3.resource") + def test_execute_query_with_partition_and_sort_key(self, mock_boto_resource): + test_data = { + "Items": [ + {"id": "1", "sort": "a", "value": Decimal("100")}, + ], + "Count": 1, + "ScannedCount": 1, + "LastEvaluatedKey": None, + } + expected_csv = [ + ["id", "sort", "value"], + ["1", "a", "100"], + ] + + mock_table = mock_boto_resource.return_value.Table.return_value + mock_table.query.return_value = test_data + + with tempfile.TemporaryDirectory() as temp_dir: + instance = DynamoDBRead() + instance._set_arguments( + { + "table_name": "test_table", + "file_name": "output.csv", + "dest_dir": temp_dir, + "region": "us-east-1", + "partition_key": "id", + "partition_value": "1", + "sort_key": "sort", + "sort_value": "a", + } + ) + instance.execute() + + output_file_path = os.path.join(temp_dir, instance.args.file_name) + self._verify_csv(output_file_path, expected_csv) + mock_table.query.assert_called_once() + + @patch("boto3.resource") + def test_execute_query_pagination(self, mock_boto_resource): + first_page = { + "Items": [{"id": "1", "value": Decimal("100")}], + "LastEvaluatedKey": {"id": "1"}, + } + second_page = { + "Items": [{"id": "2", "value": Decimal("200")}], + "LastEvaluatedKey": None, + } + expected_csv = [ + ["id", "value"], + ["1", "100"], + ["2", "200"], + ] + + mock_table = mock_boto_resource.return_value.Table.return_value + mock_table.query.side_effect = [first_page, second_page] + + with tempfile.TemporaryDirectory() as temp_dir: + instance = DynamoDBRead() + instance._set_arguments( + { + "table_name": "test_table", + "file_name": "output.csv", + "dest_dir": temp_dir, + "region": "us-east-1", + "partition_key": "id", + "partition_value": "1", + } + ) + instance.execute() + + output_file_path = os.path.join(temp_dir, instance.args.file_name) + self._verify_csv(output_file_path, expected_csv) + assert mock_table.query.call_count == 2 + + @patch("boto3.resource") + def test_execute_scan_pagination(self, mock_boto_resource): + first_page = { + "Items": [{"id": "1", "value": Decimal("100")}], + "LastEvaluatedKey": {"id": "1"}, + } + second_page = { + "Items": [{"id": "2", "value": Decimal("200")}], + "LastEvaluatedKey": None, + } + expected_csv = [ + ["id", "value"], + ["1", "100"], + ["2", "200"], + ] + + mock_table = mock_boto_resource.return_value.Table.return_value + mock_table.scan.side_effect = [first_page, second_page] + + with tempfile.TemporaryDirectory() as temp_dir: + instance = DynamoDBRead() + instance._set_arguments( + { + "table_name": "test_table", + "file_name": "output.csv", + "dest_dir": temp_dir, + "region": "us-east-1", + } + ) + instance.execute() + + output_file_path = os.path.join(temp_dir, instance.args.file_name) + self._verify_csv(output_file_path, expected_csv) + assert mock_table.scan.call_count == 2 + + def test_execute_ng_partition_value_without_partition_key(self): + with pytest.raises(InvalidParameter) as execinfo: + instance = DynamoDBRead() + instance._set_arguments( + { + "table_name": "test_table", + "file_name": "output.csv", + "region": "us-east-1", + "partition_value": "1", + } + ) + assert "Both 'partition_key' and 'partition_value' must be specified together." in str( + execinfo.value + ) + + def test_execute_ng_sort_key_without_partition_key(self): + with pytest.raises(InvalidParameter) as execinfo: + instance = DynamoDBRead() + instance._set_arguments( + { + "table_name": "test_table", + "file_name": "output.csv", + "region": "us-east-1", + "sort_key": "sort", + "sort_value": "a", + } + ) + assert ( + "'sort_key'/'sort_value' require 'partition_key'/'partition_value' " + "to also be specified." in str(execinfo.value) + ) + + def test_execute_ng_sort_value_without_sort_key(self): + with pytest.raises(InvalidParameter) as execinfo: + instance = DynamoDBRead() + instance._set_arguments( + { + "table_name": "test_table", + "file_name": "output.csv", + "region": "us-east-1", + "partition_key": "id", + "partition_value": "1", + "sort_key": "sort", + } + ) + assert "Both 'sort_key' and 'sort_value' must be specified together." in str(execinfo.value) + def _run_test(self, mock_boto_resource, test_data, expected_data, file_format): mock_table = mock_boto_resource.return_value.Table.return_value mock_table.scan.return_value = test_data @@ -326,7 +524,7 @@ def _run_test(self, mock_boto_resource, test_data, expected_data, file_format): ) instance.execute() - output_file_path = os.path.join(temp_dir, instance._file_name) + output_file_path = os.path.join(temp_dir, instance.args.file_name) assert os.path.exists(output_file_path) if file_format == "csv": From 9383f8b1855e0c0e103ddd602df7165ca5589b96 Mon Sep 17 00:00:00 2001 From: "yuichi.nasukawa" Date: Fri, 31 Jul 2026 06:58:28 +0900 Subject: [PATCH 2/2] fix: make region optional for DynamoDBRead, matching documented behavior DynamoDBRead.Arguments inherited the required `region: str` from BaseAws.Arguments, so omitting region raised a pydantic ValidationError even though the docs describe it as optional (falls back to the default AWS region). Override it as optional, same as BaseS3.Arguments does. --- cliboa/scenario/extract/aws.py | 1 + 1 file changed, 1 insertion(+) diff --git a/cliboa/scenario/extract/aws.py b/cliboa/scenario/extract/aws.py index 4b1a670c..c99d3543 100644 --- a/cliboa/scenario/extract/aws.py +++ b/cliboa/scenario/extract/aws.py @@ -212,6 +212,7 @@ class DynamoDBRead(BaseAws): """ class Arguments(BaseAws.Arguments): + region: str | None = None table_name: str dest_dir: str = "." file_name: str