Merge branch 'qa' into production

Author: aaronfriedman6

.gitignore

@@ -1,6 +1,7 @@
 .DS_Store
 dist/
 __pycache__/
+.vscode/
 *env/
 *.py[cod]
 *$py.class

CHANGELOG.md

@@ -1,5 +1,8 @@
 # Changelog
 
+## v1.0.4 - 6/28/23
+- Enforce Kinesis stream 1000 records/second write limit
+
 ## v1.0.3 - 5/19/23
 - Add research_catalog_identifier_helper function
 

README.md

@@ -35,7 +35,7 @@ kinesis_client = KinesisClient(...)
 # Do not use any version below 1.0.0
 # All available optional dependencies can be found in pyproject.toml.
 # See the "Managing dependencies" section below for more details.
-nypl-py-utils[kinesis-client,config-helper]==1.0.1
+nypl-py-utils[kinesis-client,config-helper]==1.0.4
 ```
 
 ## Developing locally

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "nypl_py_utils"
-version = "1.0.3"
+version = "1.0.4"
 authors = [
   { name="Aaron Friedman", email="aaronfriedman@nypl.org" },
 ]

src/nypl_py_utils/classes/kinesis_client.py

@@ -39,14 +39,22 @@ def close(self):
     def send_records(self, records):
         """
         Sends list of records (usually represented as Avro-encoded byte
-        strings) to Kinesis in batches of size self.batch_size.
+        strings) to Kinesis in batches of size self.batch_size. Kinesis can
+        only handle 1000 records per second, so this method waits a second
+        between each 1000 records.
         """
+        records_sent_since_pause = 0
         for i in range(0, len(records), self.batch_size):
             encoded_batch = records[i:i + self.batch_size]
             kinesis_records = [{'Data': record, 'PartitionKey':
                                 str(int(time.time() * 1000000000))}
                                for record in encoded_batch]
+
+            if records_sent_since_pause + len(encoded_batch) > 1000:
+                records_sent_since_pause = 0
+                time.sleep(1)
             self._send_kinesis_format_records(kinesis_records, 1)
+            records_sent_since_pause += len(encoded_batch)
 
     def _send_kinesis_format_records(self, kinesis_records, call_count):
         """

tests/test_kinesis_client.py

@@ -26,6 +26,7 @@ def test_send_records(self, test_instance, mocker):
         MOCK_RECORDS = [b'a', b'b', b'c', b'd', b'e']
         mocked_send_method = mocker.patch(
             'nypl_py_utils.classes.kinesis_client.KinesisClient._send_kinesis_format_records')  # noqa: E501
+        mock_sleep = mocker.patch('time.sleep', return_value=None)
 
         test_instance.send_records(MOCK_RECORDS)
         mocked_send_method.assert_has_calls([
@@ -34,6 +35,25 @@ def test_send_records(self, test_instance, mocker):
             mocker.call([_TEST_KINESIS_RECORDS[2],
                         _TEST_KINESIS_RECORDS[3]], 1),
             mocker.call([_TEST_KINESIS_RECORDS[4]], 1)])
+        mock_sleep.assert_not_called()
+
+    def test_send_records_with_pause(self, mocker):
+        mocker.patch('boto3.client')
+        test_instance = KinesisClient('test_stream_arn', 500)
+
+        MOCK_RECORDS = [b'a'] * 2200
+        mocked_send_method = mocker.patch(
+            'nypl_py_utils.classes.kinesis_client.KinesisClient._send_kinesis_format_records')  # noqa: E501
+        mock_sleep = mocker.patch('time.sleep', return_value=None)
+
+        test_instance.send_records(MOCK_RECORDS)
+        mocked_send_method.assert_has_calls([
+            mocker.call([_TEST_KINESIS_RECORDS[0]]*500, 1),
+            mocker.call([_TEST_KINESIS_RECORDS[0]]*500, 1),
+            mocker.call([_TEST_KINESIS_RECORDS[0]]*500, 1),
+            mocker.call([_TEST_KINESIS_RECORDS[0]]*500, 1),
+            mocker.call([_TEST_KINESIS_RECORDS[0]]*200, 1)])
+        assert mock_sleep.call_count == 2
 
     def test_send_kinesis_format_records(self, test_instance):
         test_instance.kinesis_client.put_records.return_value = {