From 5f009ef3cce67991c03032ebf9881326909be7f0 Mon Sep 17 00:00:00 2001 From: eL0ck Date: Mon, 6 Aug 2018 09:39:31 +1000 Subject: [PATCH] Testing for scan dump. Unfortunately not generators. Issue #9 logged for this. --- README.md | 24 +++++----- tests/dynamodb/test_batch.py | 47 +++++++++++++++++++ tlx/dynamodb/batch.py | 19 ++------ tlx/dynamodb/cli_apps/dynamodb_batch_write.py | 6 +-- 4 files changed, 64 insertions(+), 32 deletions(-) diff --git a/README.md b/README.md index 24e9c92..703cc75 100644 --- a/README.md +++ b/README.md @@ -13,27 +13,25 @@ pip install tlx All CLI applications have their own detailed help menus. Currently available tools are: - `get-aws-creds` -- `dynamo-batch-prepare` - `dynamo-batch-write` ```bash -$ dynamo-batch-prepare --help -Usage: dynamo-batch-prepare [OPTIONS] +$ dynamo-batch-write --help +Usage: dynamo-batch-write [OPTIONS] - DYNAMO BATCH PREPARE (dbp) + DYNAMO BATCH WRITE - `dynamodb-batch-prepare --help` - OR - `dbp --help` + Loads the results of a scan opperation into a table. - Takes the output of a `scan` operation such as: `aws dynamodb scan --table-name ` and formats for use - by: - - `aws dynamodb batch-write-item --request-items file://output.json` + Details: + Takes the output of a `scan` operation such as: `aws dynamodb scan --table-name ` + and writes to an existing table. Similar to the `aws dynamodb batch-write-item` command except: + - No limit to amount of items in the upload (25 with awscli) + - Take the output of a table scan, requiring no reformatting Options: - -d, --dump-file FILENAME File dumped from dynamodb with. [required] - -n, --table-name TEXT Table to send data to. Table must exist and key schema must match. Use `aws dynamodb + -d, --dump-file FILENAME File dumped from dynamodb with a `scan`. [required] + -t, --table TEXT Table to send data to. Table must exist and key schema must match. Use `aws dynamodb describe-table --table-name ` [required] -h, --help Show this message and exit. ``` diff --git a/tests/dynamodb/test_batch.py b/tests/dynamodb/test_batch.py index 3339732..f5b0ce9 100644 --- a/tests/dynamodb/test_batch.py +++ b/tests/dynamodb/test_batch.py @@ -24,6 +24,24 @@ class TestBatchLoad(TestCase): {"ID": 2, "Name": "Hubert", "Last": "McFuddle"}, ] + scan_dump_data = { + "Items": [ + { + "ID": {"N": "2"}, + "Last": {"S": "McFuddle"}, + "Name": {"S": "Hubert"}, + }, + { + "ID": {"N": "1"}, + "Last": {"S": "Jones"}, + "Name": {"S": "Flojo"}, + } + ], + "Count": 2, + "ScannedCount": 2, + "ConsumedCapacity": None, + } + expected_batch_output = [ { 'ID': Decimal('1'), @@ -120,3 +138,32 @@ def test_load_from_json_with_id(self, batch_write, get_ddb_table): batch_write.assert_called_with('table1', self.expected_batch_output) finally: os.remove(path) + + def test_load_scan_dump(self, batch_write, get_ddb_table): + msg = "should format items of dump file correctly for batch upload" + + get_ddb_table.return_value = 'table1' + + # 1. Get tempfile and write csv data + _, path = tempfile.mkstemp() + try: + # Write to scan style dump file + with open(path, 'w') as f: + json.dump(self.scan_dump_data, f) + + # 2. Call function + with open(path, 'r') as f: + tlx.dynamodb.batch.load_scan_dump(f, 'table1') + get_ddb_table.assert_called_once() + + # Check the items sent to batch_write where correct + self.assertEqual(batch_write.call_args[0][0], 'table1') + self.assertListEqual( + sorted( + batch_write.call_args[0][1], + key=lambda x: x['ID']), + self.expected_batch_output, + msg, + ) + finally: + os.remove(path) diff --git a/tlx/dynamodb/batch.py b/tlx/dynamodb/batch.py index a3ae984..1d7f5ed 100644 --- a/tlx/dynamodb/batch.py +++ b/tlx/dynamodb/batch.py @@ -26,19 +26,6 @@ def _set_types(v): return _func_map[v_key](returned_value) -def _batch_write1(table, items): - """This batch writer takes `items` from a scan. Scan results are dictionaries with their keys being - the data types. Changed the form of the scaned item into the form accepted by `put_item`""" - - # TODO: Dont do the type conversion in _pull_values. Just collapse it and use json.dump/load to deal with - # floats and ints - with table.batch_writer() as batch: - for item in items: - batch.put_item( - Item=_pull_values(item), - ) - - def batch_write(table, items): with table.batch_writer() as batch: for item in items: @@ -47,7 +34,7 @@ def batch_write(table, items): ) -def load_data(dump_file, table=None): +def load_scan_dump(dump_file, table=None): """ Loads the results of a scan opperation into a table. @@ -60,8 +47,8 @@ def load_data(dump_file, table=None): table = get_ddb_table(table) items = json.load(dump_file)['Items'] - # TODO: should run items through _pull_values here instead of inside the batch write - _batch_write1(table, items) + # TODO: Make this a generator (problem is testing it) + batch_write(table, [_pull_values(item) for item in items]) def get_ddb_table(table): diff --git a/tlx/dynamodb/cli_apps/dynamodb_batch_write.py b/tlx/dynamodb/cli_apps/dynamodb_batch_write.py index 5e2d114..e3851fb 100644 --- a/tlx/dynamodb/cli_apps/dynamodb_batch_write.py +++ b/tlx/dynamodb/cli_apps/dynamodb_batch_write.py @@ -1,6 +1,6 @@ import sys import click -from tlx.dynamodb.batch import load_data +from tlx.dynamodb.batch import load_scan_dump @click.command(context_settings=dict(max_content_width=120, help_option_names=['-h', '--help'])) @@ -10,7 +10,7 @@ def dbw(dump_file, table=None): """ DYNAMO BATCH WRITE - Loads the results of a scan opperation into a table. + Loads the results of a scan operation into a table. \b Details: @@ -21,7 +21,7 @@ def dbw(dump_file, table=None): """ try: # Surpress all exceptions for CLI app - load_data(dump_file, table) + load_scan_dump(dump_file, table) except Exception as e: print("{}: {}".format(type(e).__name__, e)) sys.exit(1)