Repository: boto/boto3
File: boto3/dynamodb/table.py (Python)

Code before:


        """
        self._table_name = table_name
        self._client = client
        self._items_buffer = []
        self._flush_amount = flush_amount

    def put_item(self, Item):
        self._items_buffer.append({'PutRequest': {'Item': Item}})
        self._flush_if_needed()

    def delete_item(self, Key):
        self._items_buffer.append({'DeleteRequest': {'Key': Key}})
        self._flush_if_needed()

    def _flush_if_needed(self):
        if len(self._items_buffer) >= self._flush_amount:
            self._flush()

    def _flush(self):
        logger.debug("Batch writer flushing pending items.")
        response = self._client.batch_write_item(
            RequestItems={self._table_name: self._items_buffer})
        unprocessed_items = response['UnprocessedItems']

        if unprocessed_items and unprocessed_items[self._table_name]:
            # Any unprocessed_items are immediately added to the
            # next batch we send.
            self._items_buffer = unprocessed_items[self._table_name]
            logger.debug("Batch writer received %s unprocessed items.",
                         len(self._items_buffer))
        else:
            self._items_buffer = []
        logger.debug("Batch write sent %s, unprocessed: %s",
                     self._flush_amount, len(self._items_buffer))

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, tb):
        # When we exit, we need to keep flushing whatever's left
        # until there's nothing left in our items buffer.
        while self._items_buffer:
            self._flush()


Review comment (suggestion):
Is it guaranteed that the table_name is in the unprocessed_items? I would be wary of accessing it directly. Probably would use a `get()`.


Code after:


        """
        self._table_name = table_name
        self._client = client
        self._items_buffer = []
        self._flush_amount = flush_amount

    def put_item(self, Item):
        self._items_buffer.append({'PutRequest': {'Item': Item}})
        self._flush_if_needed()

    def delete_item(self, Key):
        self._items_buffer.append({'DeleteRequest': {'Key': Key}})
        self._flush_if_needed()

    def _flush_if_needed(self):
        if len(self._items_buffer) >= self._flush_amount:
            self._flush()

    def _flush(self):
        response = self._client.batch_write_item(
            RequestItems={self._table_name: self._items_buffer})
        unprocessed_items = response['UnprocessedItems']

        if unprocessed_items and unprocessed_items[self._table_name]:
            # Any unprocessed_items are immediately added to the
            # next batch we send.
            self._items_buffer = unprocessed_items[self._table_name]
        else:
            self._items_buffer = []
        logger.debug("Batch write sent %s, unprocessed: %s",
                     self._flush_amount, len(self._items_buffer))

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, tb):
        # When we exit, we need to keep flushing whatever's left
        # until there's nothing left in our items buffer.
        while self._items_buffer:
            self._flush()