add timeout options to cli/module

README.md

@@ -36,6 +36,8 @@ entire dataset for the specified type will be downloaded
 * `--type`/`-t` (required): The Overture map data type to be downloaded. Examples of types are `building`
 for building footprints, `place` for POI places data, etc. Run `overturemaps download --help` for the
 complete list of allowed types
+* `--connect_timeout` (optional): Socket connection timeout, in seconds. If omitted, the AWS SDK default value is used (typically 1 second).
+* `--request_timeout` (optional): Socket read timeouts on Windows and macOS, in seconds. If omitted, the AWS SDK default value is used (typically 3 seconds). This option is ignored on non-Windows, non-macOS systems.
 
 This downloads data directly from Overture's S3 bucket without interacting with any other servers. 
 By including bounding box extents on each row in the Overture distribution, the underlying Parquet

overturemaps/cli.py

@@ -99,11 +99,13 @@ def cli():
     type=click.Choice(get_all_overture_types()),
     required=True,
 )
-def download(bbox, output_format, output, type_):
+@click.option("--connect_timeout", required=False, type=int)
+@click.option("--request_timeout", required=False, type=int)
+def download(bbox, output_format, output, type_, connect_timeout, request_timeout):
     if output is None:
         output = sys.stdout
 
-    reader = record_batch_reader(type_, bbox)
+    reader = record_batch_reader(type_, bbox, connect_timeout, request_timeout)
     if reader is None:
         return
 

overturemaps/core.py

@@ -6,15 +6,19 @@
 import pyarrow.fs as fs
 
 # Allows for optional import of additional dependencies
-try: 
+try:
     import geopandas as gpd
     from geopandas import GeoDataFrame
+
     HAS_GEOPANDAS = True
 except ImportError:
     HAS_GEOPANDAS = False
     GeoDataFrame = None
 
-def record_batch_reader(overture_type, bbox=None) -> Optional[pa.RecordBatchReader]:
+
+def record_batch_reader(
+    overture_type, bbox=None, connect_timeout=None, request_timeout=None
+) -> Optional[pa.RecordBatchReader]:
     """
     Return a pyarrow RecordBatchReader for the desired bounding box and s3 path
     """
@@ -32,7 +36,13 @@ def record_batch_reader(overture_type, bbox=None) -> Optional[pa.RecordBatchRead
         filter = None
 
     dataset = ds.dataset(
-        path, filesystem=fs.S3FileSystem(anonymous=True, region="us-west-2")
+        path,
+        filesystem=fs.S3FileSystem(
+            anonymous=True,
+            region="us-west-2",
+            connect_timeout=connect_timeout,
+            request_timeout=request_timeout,
+        ),
     )
     batches = dataset.to_batches(filter=filter)
 
@@ -48,14 +58,22 @@ def record_batch_reader(overture_type, bbox=None) -> Optional[pa.RecordBatchRead
     reader = pa.RecordBatchReader.from_batches(geoarrow_schema, non_empty_batches)
     return reader
 
-def geodataframe(overture_type: str, bbox: (float, float, float, float) = None) -> GeoDataFrame:
+
+def geodataframe(
+    overture_type: str,
+    bbox: (float, float, float, float) = None,
+    connect_timeout: int = None,
+    request_timeout: int = None,
+) -> GeoDataFrame:
     """
     Loads geoparquet for specified type into a geopandas dataframe
 
     Parameters
     ----------
     overture_type: type to load
     bbox: optional bounding box for data fetch (xmin, ymin, xmax, ymax)
+    connect_timeout: optional connection timeout in seconds
+    request_timeout: optional request timeout in seconds
 
     Returns
     -------
@@ -65,9 +83,10 @@ def geodataframe(overture_type: str, bbox: (float, float, float, float) = None)
     if not HAS_GEOPANDAS:
         raise ImportError("geopandas is required to use this function")
 
-    reader = record_batch_reader(overture_type, bbox)
+    reader = record_batch_reader(overture_type, bbox, connect_timeout, request_timeout)
     return gpd.GeoDataFrame.from_arrow(reader)
 
+
 def geoarrow_schema_adapter(schema: pa.Schema) -> pa.Schema:
     """
     Convert a geoarrow-compatible schema to a proper geoarrow schema