new script and launch file

Author: weblorin

.vscode/launch.json

@@ -516,6 +516,20 @@
         // "--delete"
       ],
     },
+    {
+      "name": "🆕 Scrape rs_context to Athena",
+      "type": "debugpy",
+      "request": "launch",
+      "program": "${workspaceFolder}/pipelines/rscontext_to_athena/rscontext_to_athena.py",
+      "console": "integratedTerminal",
+      "env": {
+        "PYTHONPATH": "${workspaceFolder}"
+      },
+      "args": [
+        "${input:environment}",
+        "{env:DATA_ROOT}/huc10_athena",
+      ],
+    },
     {
       "name": "Add ChaMP Aux Measurements to Topo - SQLite",
       "type": "debugpy",

pipelines/rscontext_to_athena/rscontext_to_athena.py

@@ -1,7 +1,11 @@
 """Scrape rs_context projects (HUC10) from Data Exchange and load the data S3 for Athena
 This version queries the Athena index of Data Exchange projects instead of using graphql API
 
-Lorin 2026-03-23
+Downloads specific files and uses geo to bin rasters.
+Requires geo extras
+`uv sync --extra geo`
+
+Lorin 2026-March-23
 
 """
 
@@ -21,7 +25,7 @@
 from pydex.lib.athena import query_to_dataframe
 from pydex.lib.raster import Raster
 
-# RegEx for finding DEM files
+# RegEx for finding DEM, Vegetation and Metrics files
 REGEXES = {"DEM_REGEX": r'.*\/dem\.tif$', "METRICS_REGEX": r'.*rscontext_metrics\.json$', "VEG_REGEX": r'.*\/existing_veg\.tif$'}
 S3_BUCKET = 'riverscapes-athena'
 S3_BASE_PATH = 'data_exchange/rs-context'
@@ -53,29 +57,11 @@ def join_s3_key(*parts: str) -> str:
     return str(PurePosixPath(*parts))
 
 
-def get_matching_file(parent_dir: str, regex: str) -> str | None:
-    """
-    Get the path to the first file in the parent directory that matches the regex.
-    Returns None if no file is found.
-    This is used to check if the output GeoPackage has already been downloaded and
-    to avoid downloading it again.
-    """
-
-    regex = re.compile(regex)
-    for root, __dirs, files in os.walk(parent_dir):
-        for file_name in files:
-            # Check if the file name matches the regex
-            if regex.match(file_name):
-                return os.path.join(root, file_name)
-
-    return None
-
-
 def scrape_rscontext_project(s3, rs_api: RiverscapesAPI, project: RiverscapesProject, download_dir: Path, skip_overwrite: bool):
     """Scrape (download, transform, upload) a single project"""
     DOWNLOAD_RETRIES = 3
     log = Logger("Scrape RSContext project")
-    # upload to s3 key
+    # S3 key for upload
     s3_key = join_s3_key(S3_BASE_PATH, f'{project.huc}.json')
     if project.huc is None or project.huc == '':
         log.warning(f'Project {project.id} does not have a HUC. Skipping.')
@@ -106,20 +92,20 @@ def scrape_rscontext_project(s3, rs_api: RiverscapesAPI, project: RiverscapesPro
         complete = False
         while retry < DOWNLOAD_RETRIES and complete is False:
             try:
-                rs_api.download_files(project_id=project.id, download_dir=huc_dir, re_filter=list(REGEXES.values()))
+                rs_api.download_files(project_id=project.id, download_dir=str(huc_dir), re_filter=list(REGEXES.values()))
                 complete = True
                 break
             except Exception as e:
                 log.error(f'Error downloading files for project {project.id}: {e}')
                 traceback.print_exc(file=sys.stdout)
                 retry += 1
-            return
+            continue
 
-        dem_tif_path = huc_dir / 'topography' / 'dem_tif'
-        if not dem_tif_path.exists:
+        dem_tif_path = huc_dir / 'topography' / 'dem.tif'
+        if not dem_tif_path.exists():
             raise FileNotFoundError(f'Could not find DEM file for project {project.id}')
         veg_tif_path = huc_dir / 'vegetation' / 'existing_veg.tif'
-        if not veg_tif_path.exists:
+        if not veg_tif_path.exists():
             raise FileNotFoundError(f'Could not find vegetation file for project {project.id}')
         metrics_json_path = huc_dir / 'rscontext_metrics.json'
         try:
@@ -137,27 +123,27 @@ def scrape_rscontext_project(s3, rs_api: RiverscapesAPI, project: RiverscapesPro
             metrics['rs_context'] = {}
         metrics['rs_context']['dem_bins'] = dem_bins
         metrics['rs_context']['existing_veg_bins'] = veg_bins
-        log.info(f'Writing HUC10 metrics to {huc10_json_path}')
 
         # Add the project ID to the metrics so we can trace this back to its source
         metrics['rs_context']['project_id'] = project.id
         metrics['rs_context']['model_version'] = str(project.model_version)
 
-        # Write the JSON back to `huc10code.json` (just for debugging purposes really)
+        log.info(f'Writing HUC10 metrics to {huc10_json_path}')
+        # Write the JSON back to `huc10_{huc}.json` (just for debugging purposes really)
         with open(huc10_json_path, 'w', encoding='utf-8') as f:
             json.dump(metrics, f, indent=2)
 
         # Now use boto3 to upload the file to S3
-        log.info(f'Uploading {huc10_json_path} to s3://{S3_BUCKET}/{s3_key}')
+        # log.info(f'Uploading metrics to s3://{S3_BUCKET}/{s3_key}')
 
-        s3.put_object(Bucket=S3_BUCKET, Key=s3_key, Body=json.dumps(metrics['rs_context']))
+        # s3.put_object(Bucket=S3_BUCKET, Key=s3_key, Body=json.dumps(metrics['rs_context']))
 
     except Exception as e:
         log.error(f'Error scraping HUC {project.huc}: {e}')
         traceback.print_exc(file=sys.stdout)
 
 
-def scrape_rsprojects(rs_api: RiverscapesAPI, download_dir: Path, delete_downloads: bool):
+def scrape_rsprojects(rs_api: RiverscapesAPI, download_dir: Path, delete_downloads: bool, skip_overwrite: bool):
     """Scrape all projects matching criteria"""
     log = Logger('Scrape RSContext')
     projects_to_add_df = query_to_dataframe(missing_projects_query, 'identify new projects')
@@ -176,12 +162,12 @@ def scrape_rsprojects(rs_api: RiverscapesAPI, download_dir: Path, delete_downloa
         count += 1
         prg.update(count)
 
-    if delete_downloads is True and os.path.isdir(huc_dir):
+    if delete_downloads is True and download_dir.is_dir():
         try:
-            log.info(f'Deleting download directory {huc_dir}')
-            shutil.rmtree(huc_dir)
+            log.info(f'Deleting download directory {download_dir}')
+            shutil.rmtree(download_dir)
         except Exception as e:
-            log.error(f'Error deleting download directory {huc_dir}: {e}')
+            log.error(f'Error deleting download directory {download_dir}: {e}')
 
 
 def main():
@@ -191,25 +177,29 @@ def main():
     parser = argparse.ArgumentParser()
     parser.add_argument('stage', help='Environment: staging or production', type=str)
     parser.add_argument('working_folder', help='top level folder for downloads and output', type=str)
-    parser.add_argument('--delete', help='Whether or not to delete downloaded GeoPackages', action='store_true', default=False)
+    parser.add_argument('--delete', help='Delete downloaded files after processing', action='store_true', default=False)
+    parser.add_argument('--skip-overwrite', help='Whether or not to skip overwriting existing S3 files', action='store_true', default=False)
 
     args = dotenv.parse_args_env(parser)
 
     # Set up some reasonable folders to store things
     working_folder = Path(args.working_folder)
     download_folder = working_folder / 'downloads'
+    safe_makedirs(str(working_folder))
 
     log = Logger('Setup')
     log.setup(log_path=working_folder / 'rscontext_to_athena.log', log_level=logging.DEBUG)
     try:
         with RiverscapesAPI(stage=args.stage) as rs_api:
-            scrape_rsprojects(rs_api, download_folder, args.delete)
+            scrape_rsprojects(rs_api, download_folder, args.delete, args.skip_overwrite)
 
     except Exception as e:
         log.error(e)
         traceback.print_exc(file=sys.stdout)
         sys.exit(1)
 
+    log.info('Process complete')
+
 
 if __name__ == '__main__':
     main()