Finish url verificiation, download, and geospatial extent for RFCs

blobfish/aorc/mirror_utils/aio.py

@@ -0,0 +1,85 @@
+""" Asynchronous operations for creating AORC data mirror on s3"""
+# Make sure script can access common classes
+import sys
+
+sys.path.append("../classes")
+
+import asyncio
+import io
+import datetime
+from collections.abc import Iterator
+
+from aiohttp import ClientSession, ClientTimeout
+from classes.common import AORCDataURL
+
+
+async def async_verify_url(session: ClientSession, url: str, **kwargs) -> dict | None:
+    print(f"Verifying {url}")
+    async with session.head(url) as resp:
+        if resp.ok:
+            last_modification_str = resp.headers.get("Last-Modified")
+            last_modification_dt = datetime.datetime.strptime(last_modification_str, '%a, %d %b %Y %H:%M:%S %Z')
+            result = {
+                "url": url,
+                "last_modified": last_modification_dt
+            }
+            if kwargs:
+                result.update(kwargs)
+            return result
+
+
+async def async_stream_zip_to_s3(zip_url: AORCDataURL, client_session: ClientSession, s3_resource, bucket: str) -> AORCDataURL:
+    bucket = s3_resource.Bucket(bucket)
+    print(f"Starting download for {zip_url.url}")
+    async with client_session.get(zip_url.url) as resp:
+        data = io.BytesIO()
+        async for chunk in  resp.content.iter_chunked(1024):
+            data.write(chunk)
+        data.seek(0)
+        print(f"All data loaded for {zip_url.url}")
+        obj = bucket.Object(zip_url.s3_key())
+        print(f"Starting upload for {zip_url.s3_key()}")
+        obj.upload_fileobj(data)
+        s3_url = zip_url
+        s3_url.url = f"s3://{bucket}/{zip_url.s3_key()}"
+        obj.load()
+        s3_url.last_modified = obj.last_modified
+    return s3_url
+
+
+async def async_verify_url_iter(potential_urls: Iterator[AORCDataURL], limit: int, timeout: float):
+    timeout = ClientTimeout(timeout)
+    semaphore = asyncio.BoundedSemaphore(limit)
+    async with ClientSession(timeout=timeout) as session:
+        tasks = []
+        for aorc_data_url in potential_urls:
+            async with semaphore:
+                task = async_verify_url(session, aorc_data_url.url, data_start_dt=aorc_data_url.data_start_dt, rfc_alias=aorc_data_url.rfc_alias)
+                tasks.append(task)
+        results = await asyncio.gather(*tasks)
+    return results
+
+async def async_stream_zip_to_s3_iter(zip_urls: Iterator[AORCDataURL], s3_resource, bucket: str, limit: int, timeout: float):
+    timeout = ClientTimeout(timeout)
+    semaphore = asyncio.BoundedSemaphore(limit)
+    async with ClientSession(timeout=timeout) as session:
+        tasks = []
+        for aorc_data_url in zip_urls:
+            async with semaphore:
+                task = async_stream_zip_to_s3(aorc_data_url, session, s3_resource, bucket)
+                tasks.append(task)
+        results = await asyncio.gather(*tasks)
+    return results
+
+
+def verify_urls(potential_urls: Iterator[AORCDataURL], limit: int = 5, timeout: float = 2000) -> Iterator[AORCDataURL]:
+    results_list = asyncio.run(async_verify_url_iter(potential_urls, limit, timeout))
+    for result in results_list:
+        if result != None:
+            yield AORCDataURL(**result)
+
+def stream_zips_to_s3(zip_urls: Iterator[AORCDataURL], s3_resource, bucket: str, limit: int = 5, timeout: float = 10000) -> Iterator[AORCDataURL]:
+    results_list = asyncio.run(async_stream_zip_to_s3_iter(zip_urls, s3_resource, bucket, limit, timeout))
+    for result in results_list:
+        if result != None:
+            yield result

blobfish/aorc/mirror_utils/general.py

@@ -1,16 +1,57 @@
 """ Synchronous operations for creating AORC data mirror on s3"""
-import requests
-import shapely
+# Make sure script can access common classes
+import sys
+
+sys.path.append("../classes")
+
 import datetime
+import io
+import os
+import tarfile
+from collections.abc import Iterator
+from tempfile import TemporaryDirectory
+
+import fiona
+import requests
+from classes.common import AORCDataURL, RFCInfo
+from shapely.geometry import (GeometryCollection, LinearRing, LineString,
+                              MultiLineString, MultiPoint, MultiPolygon, Point,
+                              Polygon, shape)
 
 
-def create_rfc_list(rfc_tar_shp_url: str) -> list:
-    return []
+def create_rfc_list(rfc_tar_shp_url: str) -> list[tuple[str, Point | MultiPoint | LineString | MultiLineString | Polygon | MultiPolygon | LinearRing | GeometryCollection]]:
+    with requests.get(rfc_tar_shp_url, stream=True) as resp:
+        if resp.ok:
+            with TemporaryDirectory() as tmpdir:
+                with io.BytesIO(resp.content) as stream:
+                    with tarfile.open(fileobj=stream, mode="r:gz") as tar:
+                        tar.extractall(tmpdir)
+                        for member in tar.getmembers():
+                            if member.name.endswith(".shp"):
+                                shp_path = os.path.join(tmpdir, member.name)
+                                with fiona.open(shp_path, driver="ESRI Shapefile") as f:
+                                    rfc_features = [(feature["properties"]["NAME"], shape(feature["geometry"])) for feature in f]
+                                    return rfc_features
+                        else:
+                            raise FileNotFoundError("No shapefile found in tarfile")
+        else:
+            raise requests.exceptions.ConnectionError(f"Request response status indicates faillure: {resp.status_code}")
 
 
-def create_potential_urls(rfcs: list, start_dt: datetime.datetime) -> list:
-    return []
+def create_potential_urls(rfcs: list[RFCInfo], start_dt: datetime.datetime, base_url: str) -> Iterator[AORCDataURL]:
+    end_dt = datetime.datetime.now()
+    end_dt = datetime.datetime(end_dt.year, end_dt.month, 1)
+    start_dt = datetime.datetime(start_dt.year, start_dt.month, 1)
+    current_dt = start_dt
+    while current_dt < end_dt:
+        for rfc in rfcs:
+            potential_url = f"{base_url}/AORC_{rfc.alias}RFC_4km/{rfc.alias}RFC_precip_partition/AORC_APCP_4KM_{rfc.alias}RFC_{current_dt.strftime('%Y%m')}.zip"
+            yield AORCDataURL(potential_url, current_dt, rfc.alias)
+        if current_dt.month < 12:
+            current_dt = datetime.datetime(current_dt.year, current_dt.month + 1, 1)
+        else:
+            current_dt = datetime.datetime(current_dt.year + 1, 1, 1)
 
 
-def check_zip_metadata(zip_fn: str) -> dict:
+def check_zipped_nc_meta(zip_fn: str) -> dict:
     return {}