WIP: In process of boto3 upgrades for latest/ put_item.

Author: ABPLMC

ingester/datalake_ingester/storage.py

@@ -12,11 +12,13 @@
 # License for the specific language governing permissions and limitations under
 # the License.
 
+import boto3
+from boto3.dynamodb.conditions import Attr
+
 from memoized_property import memoized_property
 import boto.dynamodb2
 from boto.dynamodb2.table import Table
-from boto.dynamodb2.exceptions import (ConditionalCheckFailedException,
-                                       ItemNotFound)
+from boto.dynamodb2.exceptions import ConditionalCheckFailedException
 import os
 from datalake.common.errors import InsufficientConfiguration
 
@@ -26,6 +28,9 @@ class DynamoDBStorage(object):
 
     def __init__(self, table_name, connection=None):
         self.table_name = table_name
+        self.latest_table_name = os.environ.get("DATALAKE_DNAMODB_LATEST_TABLE",
+                                                f"{self.table_name}-latest")
+        self.use_latest = os.environ.get("DATALAKE_LATEST_FLAG", False)
         self._prepare_connection(connection)
 
     @classmethod
@@ -48,27 +53,29 @@ def _prepare_connection(self, connection):
     @memoized_property
     def _table(self):
         return Table(self.table_name, connection=self._connection)
+    
+    @memoized_property
+    def _latest_table(self):
+        dynamodb = boto3.resource('dynamodb')
+        return dynamodb.Table(self.latest_table_name, connection=self._connection)
 
     def store(self, record):
         try:
-            primary_key = {
-                'time_index_key': record['time_index_key'],
-                'range_key': record['range_key']
-            }
-            existing_record = self._table.get_item(**primary_key)
-
-            if existing_record and \
-                existing_record['metadata']['start'] < record['metadata']['start']:
-                 self._table.put_item(data=record, overwrite=True)
-            else:
-                print(f'Existing record in latest table has later or same start time. No update performed')
-
-        except ItemNotFound:
-            # Item doesn't exist, lets insert
             self._table.put_item(data=record)
         except ConditionalCheckFailedException:
             # Tolerate duplicate stores
             pass
 
     def update(self, record):
         self._table.put_item(data=record, overwrite=True)
+
+    def store_latest(self, record):
+        item_attrs = {'time_index_key': record['time_index_key'],
+                      'range_key': record['range_key']}
+        condition = (Attr('metadata.start').lt(record['metadata']['start']))
+        try:
+            self._latest_table.put_item(item_attrs,
+                                        condition)
+        except ConditionalCheckFailedException:
+            pass
+

ingester/tests/conftest.py

@@ -14,17 +14,31 @@
 import boto.sns
 import boto.sqs
 
+
+import boto3
+from boto3 import client
+from botocore.exceptions import ClientError
+
 from datalake.tests import *  # noqa
 
 from datalake_ingester import SQSQueue
 
+import logging
+logging.basicConfig(level=logging.DEBUG)
 
 @pytest.fixture
 def dynamodb_connection(aws_connector):
     return aws_connector(mock_dynamodb2,
                          lambda: boto.dynamodb2.connect_to_region('us-west-1'))
 
 
+@pytest.fixture
+def dynamodb_latest_connection(aws_connector):
+    return aws_connector(mock_dynamodb2,
+                         lambda: boto3.resource('dynamodb',                          
+                          region_name='us-west-1'))
+
+
 def _delete_table_if_exists(conn, name):
     try:
         table = Table(name, connection=conn)
@@ -33,6 +47,16 @@ def _delete_table_if_exists(conn, name):
         if e.status == 400 and e.error_code == 'ResourceNotFoundException':
             return
         raise e
+    
+def _delete_latest_if_exists(dynamodb, name):
+    try:
+        table = dynamodb.Table(name)
+        table.delete()
+        table.wait_until_not_exists()
+    except ClientError as e:
+        if e.response['Error']['Code'] == 'ResourceNotFoundException':
+            return
+        raise e
 
 
 @pytest.fixture
@@ -55,6 +79,32 @@ def tear_down():
     return table_maker
 
 
+@pytest.fixture
+def dynamodb_latest_table_maker(request, dynamodb_latest_connection):
+
+    def table_maker(name, key_schema, attributes):
+        _delete_latest_if_exists(dynamodb_latest_connection, name)
+        table = dynamodb_latest_connection.create_table(
+            TableName=name,
+            KeySchema=key_schema,
+            AttributeDefinitions=attributes,
+            ProvisionedThroughput={
+                'ReadCapacityUnits': 5,
+                'WriteCapacityUnits': 5
+            }
+        )
+        table.wait_until_exists()
+
+        def tear_down():
+            table.delete()
+            table.wait_until_not_exists()
+
+        request.addfinalizer(tear_down)
+        return table
+
+    return table_maker
+
+
 @pytest.fixture
 def dynamodb_users_table(dynamodb_table_maker):
     schema = [HashKey('name'), RangeKey('last_name')]
@@ -68,9 +118,16 @@ def dynamodb_records_table(dynamodb_table_maker):
 
 
 @pytest.fixture
-def dynamodb_latest_table(dynamodb_table_maker):
-    schema = [HashKey('time_index_key'), RangeKey('range_key')]
-    return dynamodb_table_maker('latest', schema)
+def dynamodb_latest_table(dynamodb_latest_table_maker):
+    schema = [
+        {'AttributeName': 'time_index_key', 'KeyType': 'HASH'},
+        {'AttributeName': 'range_key', 'KeyType': 'RANGE'}
+    ]
+    attributes = [
+        {'AttributeName': 'time_index_key', 'AttributeType': 'S'},
+        {'AttributeName': 'range_key', 'AttributeType': 'S'}
+    ]
+    return dynamodb_latest_table_maker('latest', schema, attributes)
 
 
 @pytest.fixture

ingester/tests/test_storage.py

@@ -1,7 +1,9 @@
 from datalake_ingester import DynamoDBStorage
 from decimal import Decimal
+import boto3
 
 def test_insert_new_record(dynamodb_latest_table, dynamodb_connection):
+    boto3.setup_default_session(fake_credentials=True)
     storage = DynamoDBStorage('latest', connection=dynamodb_connection)
     new_record = {
         'time_index_key': '15225:newlog',
@@ -13,7 +15,7 @@ def test_insert_new_record(dynamodb_latest_table, dynamodb_connection):
         'create_time': 1500000000000
     }
 
-    storage.store(new_record)
+    storage.store_latest(new_record)
 
     stored_record = dynamodb_latest_table.get_item(
         time_index_key='15225:newlog',
@@ -64,9 +66,9 @@ def test_store_conditional_put_latest_multiple_files(dynamodb_latest_table, dyna
         'size': 1048576
     }
 
-    storage.store(file3)
-    storage.store(file1)
-    storage.store(file2) # same what:where, but should replace file1 b/c newer
+    storage.store_latest(file3)
+    storage.store_latest(file1)
+    storage.store_latest(file2) # same what:where, but should replace file1 b/c newer
 
     records = [dict(i) for i in dynamodb_latest_table.scan()]
 
@@ -99,8 +101,8 @@ def test_concurrent_updates(dynamodb_latest_table, dynamodb_connection):
     updated_record2['metadata']['start'] += 5
 
 
-    storage.store(updated_record1)
-    storage.store(updated_record2)
+    storage.store_latest(updated_record1)
+    storage.store_latest(updated_record2)
 
     stored_record = dynamodb_latest_table.get_item(
         time_index_key='15219:zlcdzvawsp',