handle streaming results in some streams

Signed-off-by: John Seekins <[email protected]>

Author: johnseekins

enricher.py

@@ -1,9 +1,9 @@
 from concurrent.futures import ProcessPoolExecutor
 import copy
 from enrichers import (
+    openstreetmap,
     wikidata,
     wikipedia,
-    openstreetmap,
 )
 from schemas import (
     default_coords,
@@ -43,14 +43,11 @@ def enrich_facility(facility_data: tuple) -> tuple:
     logger.info("Enriching facility %s...", facility_name)
     enriched_facility = copy.deepcopy(facility)
 
-    wiki = wikipedia.Wikipedia(facility_name=facility_name, wait_time=WIKIPEDIA_DELAY)
-    wiki_res = wiki.search()
-    wd = wikidata.Wikidata(facility_name=facility_name, wait_time=WIKIDATA_DELAY)
-    wd_res = wd.search()
-    osm = openstreetmap.OpenStreetMap(
+    wiki_res = wikipedia.Wikipedia(facility_name=facility_name, wait_time=WIKIPEDIA_DELAY).search()
+    wd_res = wikidata.Wikidata(facility_name=facility_name, wait_time=WIKIDATA_DELAY).search()
+    osm_res = openstreetmap.OpenStreetMap(
         facility_name=facility_name, address=facility.get("address", {}), wait_time=OSM_DELAY
-    )
-    osm_res = osm.search()
+    ).search()
     enriched_facility["wikipedia"]["page_url"] = wiki_res.get("url", "")
     enriched_facility["wikipedia"]["search_query"] = wiki_res.get("search_query_steps", "")
     enriched_facility["wikidata"]["page_url"] = wd_res.get("url", "")

enrichers/__init__.py

@@ -3,6 +3,7 @@
 from schemas import enrich_resp_schema
 import time
 from utils import (
+    default_headers,
     session,
 )
 
@@ -24,10 +25,21 @@ def search(self) -> dict:
         """Child objects should implement this"""
         return {}
 
-    def _req(self, url: str, params: dict = {}, timeout: int = 10) -> requests.Response:
+    def _req(
+        self, url: str, params: dict = {}, timeout: int = 10, stream: bool = False, headers: dict = default_headers
+    ) -> requests.Response:
         """requests response wrapper to ensure we honor waits"""
 
-        response = session.get(url, allow_redirects=True, timeout=timeout, params=params)
+        # ensure we get all headers configured correctly
+        # but manually applied headers win the argument
+        for k, v in default_headers.items():
+            if k in headers.keys():
+                continue
+            headers[k] = v
+
+        response = session.get(
+            url, allow_redirects=True, timeout=timeout, params=params, stream=stream, headers=headers
+        )
         response.raise_for_status()
         time.sleep(self.wait_time)
         return response

scraper.py

@@ -58,7 +58,7 @@ def _download_sheet(self) -> None:
         logger.debug("Found sheet at: %s", actual_link)
         self.sheet_url = actual_link
         logger.info("Downloading detention stats sheet from %s", self.sheet_url)
-        resp = session.get(self.sheet_url, timeout=120)
+        resp = session.get(self.sheet_url, timeout=120, stream=True)
         with open(self.filename, "wb") as f:
             for chunk in resp.iter_content(chunk_size=1024):
                 if chunk:

utils.py

@@ -13,11 +13,12 @@
     total=4,
     backoff_factor=1,
 )
+default_headers = {"User-Agent": "ICE-Facilities-Research/1.0 (Educational Research Purpose)"}
 _adapter = HTTPAdapter(max_retries=_retry_strategy)
 session = requests.Session()
 session.mount("https://", _adapter)
 session.mount("http://", _adapter)
-session.headers.update({"User-Agent": "ICE-Facilities-Research/1.0 (Educational Research Purpose)"})
+session.headers.update(default_headers)
 
 default_timestamp = "1970-01-01T00:00:00-+0000"
 timestamp_format = "%Y-%m-%dT%H:%M:%S-%z"