Merge pull request #388 from skalish/stream-records

TC: Stream records from a dataset

CHANGELOG.md

@@ -6,6 +6,7 @@
   - [#367](https://github.com/Datatamer/tamr-client/issues/367) Support for projects:
     - generic projects via `tc.project`
     - Mastering projects via `tc.mastering.project`
+  - Support for streaming records from a dataset via `tc.record.stream`
 
   **BUG FIXES**
   - `from_geo_features` now returns information on the operation.  

docs/beta/dataset/record.rst

@@ -7,6 +7,7 @@ Record
 .. autofunction:: tamr_client.record.upsert
 .. autofunction:: tamr_client.record.delete
 .. autofunction:: tamr_client.record._update
+.. autofunction:: tamr_client.record.stream
 
 Exceptions
 ----------

tamr_client/dataset/record.py

@@ -1,11 +1,11 @@
 """
 See https://docs.tamr.com/reference/record
-"The recommended approach for interacting with records is to use the :func:`~tamr_client.record.upsert` and
+"The recommended approach for modifying records is to use the :func:`~tamr_client.record.upsert` and
 :func:`~tamr_client.record.delete` functions for all use cases they can handle. For more advanced use cases, the
-underlying _update function can be used directly."
+underlying :func:`~tamr_client.record._update` function can be used directly."
 """
 import json
-from typing import cast, Dict, IO, Iterable, Optional
+from typing import cast, Dict, IO, Iterable, Iterator, Optional
 
 from tamr_client import response
 from tamr_client.dataset.dataset import Dataset
@@ -143,3 +143,16 @@ def _delete_command(record: Dict, *, primary_key_name: str) -> Dict:
         The DELETE command in the proper format
     """
     return {"action": "DELETE", "recordId": record[primary_key_name]}
+
+
+def stream(session: Session, dataset: Dataset) -> Iterator[JsonDict]:
+    """Stream the records in this dataset as Python dictionaries.
+
+    Args:
+        dataset: Dataset from which to stream records
+
+    Returns:
+        Python generator yielding records
+    """
+    with session.get(str(dataset.url) + "/records", stream=True) as r:
+        return response.ndjson(r)

tests/tamr_client/datasets/test_record.py → tests/tamr_client/dataset/test_record.py

@@ -1,4 +1,5 @@
 from functools import partial
+import json
 from typing import Dict
 
 import pytest
@@ -154,6 +155,20 @@ def test_delete_infer_primary_key():
     assert snoop["payload"] == utils.stringify(deletes)
 
 
+@responses.activate
+def test_stream():
+    s = utils.session()
+    dataset = utils.dataset()
+
+    url = tc.URL(path="datasets/1/records")
+    responses.add(
+        responses.GET, str(url), body="\n".join(json.dumps(x) for x in _records_json)
+    )
+
+    records = tc.record.stream(s, dataset)
+    assert list(records) == _records_json
+
+
 _records_json = [{"primary_key": 1}, {"primary_key": 2}]
 
 _response_json = {