diff --git a/scripts/download_all.py b/scripts/download_all.py index 4188b20..f22c674 100644 --- a/scripts/download_all.py +++ b/scripts/download_all.py @@ -7,20 +7,16 @@ import os import pkgutil import random -import sys - -# Add rivretrieve to path -sys.path.append(os.path.join(os.path.dirname(__file__), "rivretrieve")) import rivretrieve from rivretrieve import constants # Configuration ROOT_DIR = "downloaded_data" -VARIABLE = constants.DISCHARGE +VARIABLE = constants.DISCHARGE_DAILY_MEAN # Default variable to download START_DATE = "1950-01-01" END_DATE = "2025-10-03" -N_WORKERS = 8 +N_WORKERS = 1 # Japan specific dates JAPAN_START_DATE = "1980-01-01" @@ -51,40 +47,43 @@ def get_fetcher_classes(): return fetchers -def download_gauge_data(country, fetcher_class, gauge_id, start_date, end_date): +def download_gauge_data(country, fetcher_instance, gauge_id, variable, start_date, end_date): """Downloads and saves data for a single gauge.""" output_dir = os.path.join(ROOT_DIR, country) os.makedirs(output_dir, exist_ok=True) # Sanitize gauge_id to be used as a filename sanitized_gauge_id = "".join(c if c.isalnum() or c in ["-", "_"] else "_" for c in gauge_id) - output_file = os.path.join(output_dir, f"{sanitized_gauge_id}.csv") + output_file = os.path.join(output_dir, f"{sanitized_gauge_id}_{variable}.csv") if os.path.exists(output_file): - logging.info(f"Skipping {country} - {gauge_id} (already downloaded)") - return f"SKIPPED: {country} - {gauge_id}" + logging.info(f"Skipping {country} - {gauge_id} - {variable} (already downloaded)") + return f"SKIPPED: {country} - {gauge_id} - {variable}" - logging.info(f"Processing {country} - {gauge_id} with dates {start_date} to {end_date}") + logging.info(f"Processing {country} - {gauge_id} - {variable} with dates {start_date} to {end_date}") try: - fetcher = fetcher_class() - data = fetcher.get_data( + if variable not in fetcher_instance.get_available_variables(): + logging.warning(f"Variable {variable} not supported by {country} fetcher, skipping gauge {gauge_id}") + return f"UNSUPPORTED VARIABLE: {country} - {gauge_id} - {variable}" + + data = fetcher_instance.get_data( gauge_id=gauge_id, - variable=VARIABLE, + variable=variable, start_date=start_date, end_date=end_date, ) if data is not None and not data.empty: - data.to_csv(output_file, index=False) - logging.info(f"Successfully downloaded and saved {country} - {gauge_id}") - return f"SUCCESS: {country} - {gauge_id}" + data.to_csv(output_file) + logging.info(f"Successfully downloaded and saved {country} - {gauge_id} - {variable}") + return f"SUCCESS: {country} - {gauge_id} - {variable}" else: - logging.info(f"No data returned for {country} - {gauge_id}") - return f"NO DATA: {country} - {gauge_id}" + logging.info(f"No data returned for {country} - {gauge_id} - {variable}") + return f"NO DATA: {country} - {gauge_id} - {variable}" except Exception as e: - logging.error(f"Error downloading {country} - {gauge_id}: {e}", exc_info=False) - return f"FAILED: {country} - {gauge_id} - {e}" + logging.error(f"Error downloading {country} - {gauge_id} - {variable}: {e}", exc_info=False) + return f"FAILED: {country} - {gauge_id} - {variable} - {e}" def main(): @@ -101,41 +100,60 @@ def main(): default=["all"], help=f"Specify which fetchers to use. Choices are {fetcher_names + ['all']}", ) + parser.add_argument( + "--variable", + type=str, + default=VARIABLE, + help=f"Variable to download (e.g., {constants.DISCHARGE_DAILY_MEAN}).", + ) + parser.add_argument("--start_date", type=str, default=START_DATE, help="Start date in YYYY-MM-DD.") + parser.add_argument("--end_date", type=str, default=END_DATE, help="End date in YYYY-MM-DD.") + parser.add_argument("--n_workers", type=int, default=N_WORKERS, help="Number of worker threads.") args = parser.parse_args() selected_fetchers = args.fetchers + variable_to_download = args.variable logging.info(f"Selected fetchers: {selected_fetchers}") + logging.info(f"Variable to download: {variable_to_download}") - fetcher_classes_to_run = {} + fetcher_instances = {} if "all" in selected_fetchers: - fetcher_classes_to_run = all_fetcher_classes + for country, fetcher_class in all_fetcher_classes.items(): + try: + fetcher_instances[country] = fetcher_class() + except Exception as e: + logging.error(f"Failed to instantiate fetcher for {country}: {e}") else: for fetcher_name in selected_fetchers: if fetcher_name in all_fetcher_classes: - fetcher_classes_to_run[fetcher_name] = all_fetcher_classes[fetcher_name] + try: + fetcher_instances[fetcher_name] = all_fetcher_classes[fetcher_name]() + except Exception as e: + logging.error(f"Failed to instantiate fetcher for {fetcher_name}: {e}") else: logging.warning(f"Fetcher '{fetcher_name}' not found, skipping.") tasks = [] - for country, fetcher_class in fetcher_classes_to_run.items(): + for country, fetcher_instance in fetcher_instances.items(): try: - sites = fetcher_class.get_gauge_ids() + sites = fetcher_instance.get_cached_metadata() if sites is None or sites.empty: logging.warning(f"No sites found for {country}") continue - current_start_date = START_DATE - current_end_date = END_DATE + current_start_date = args.start_date + current_end_date = args.end_date if country == "japan": current_start_date = JAPAN_START_DATE current_end_date = JAPAN_END_DATE - for gauge_id in sites[constants.GAUGE_ID]: + for gauge_id in sites.index: tasks.append( ( country, - fetcher_class, + fetcher_instance, gauge_id, + variable_to_download, current_start_date, current_end_date, ) @@ -144,32 +162,13 @@ def main(): logging.error(f"Error getting sites for {country}: {e}") random.shuffle(tasks) - logging.info(f"Found {len(tasks)} total sites to process for fetchers: {list(fetcher_classes_to_run.keys())}.") + logging.info(f"Found {len(tasks)} total sites to process for fetchers: {list(fetcher_instances.keys())}.") if not tasks: logging.info("No tasks to process. Exiting.") return - if "poland" in fetcher_classes_to_run: - logging.info("Poland fetcher selected, ensuring cache exists...") - try: - poland_fetcher = fetcher_classes_to_run["poland"]() - # This call will block if the cache needs to be built - poland_fetcher.get_data( - gauge_id="dummy", - variable=constants.DISCHARGE, - start_date="2000-01-01", - end_date="2000-01-01", - ) - logging.info("Poland cache check complete.") - except Exception as e: - logging.error(f"Error during Poland cache pre-check: {e}") - # Optionally, remove poland from fetcher_classes_to_run if pre-check fails - if "poland" in fetcher_classes_to_run: - del fetcher_classes_to_run["poland"] - logging.info("Removed Poland from fetchers to run due to pre-check error.") - - with concurrent.futures.ThreadPoolExecutor(max_workers=N_WORKERS) as executor: + with concurrent.futures.ThreadPoolExecutor(max_workers=args.n_workers) as executor: futures = [executor.submit(download_gauge_data, *task) for task in tasks] for future in concurrent.futures.as_completed(futures): try: