noaa-goes-visualization/puller_fits.py

import os
import sys
import urllib.request
import urllib.parse
import re
import time
import random
import datetime
from threading import Thread
import queue
import math
from functools import partial
from queue import Empty

import tqdm

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from suvi import db, index, paths, vfs

directory_url = r"https://data.ngdc.noaa.gov/platforms/solar-space-observing-satellites/goes/"
stored_images_dir = paths.data_root()
ignore_folder_names = ["Parent Directory", "l1b", "goes17", "2017", "2018", "2019", "2020", "2021", "2022", "2023"]
fetch_interval = 0 # 60*60
nfetchworkers = 2 # Be nice to the servers, this value is how many threads will be asking for links and file info at the same time
ndownloadworkers = 3 # Be nice to the servers, this value is how many threads will be downloading files at the same time
randomize_order = False

links_regex_pattern = r'(?<=<a href=")([^ ?:]*)(?=">.*\d{4}-\d{2}-\d{2} \d{2}:\d{2})' # Find href links that do not contain question marks or whitespace
times_regex_pattern = r'(?<=<td align="right">)(\d{4}-\d{2}-\d{2} \d{2}:\d{2})(?= )' # Find timestamps in UTC in the format YYYY-MM-DD HH:mm
sizes_regex_pattern = r'(?<=\d{4}-\d{2}-\d{2} \d{2}:\d{2}  <\/td><td align="right">)([ .\d]{1,3})([KMG]|- |0 | )(?=<\/td>)' # Find file size markers in one of many possible formats
links_matcher = re.compile(links_regex_pattern)
times_matcher = re.compile(times_regex_pattern)
sizes_matcher = re.compile(sizes_regex_pattern)

def find_links_worker(query_work_queue, query_result_queue, attempt_count = 3, randomize_order = True):
    while True:
        job = query_work_queue.get()
        if job == None:
            return
        url= job
    
        attempts = 0
        html_content = None
        while attempts < attempt_count:
            try:
                with urllib.request.urlopen(url) as response:
                    html_content = response.read().decode('utf-8')
                    break
            except Exception as e:
                tqdm.tqdm.write(f"Exception while fetching links: {e}")
                time.sleep(1 + random.random())
                attempts += 1
            if (attempt_count > 1) and (attempts == attempt_count):
                tqdm.tqdm.write(f'\nAfter {attempt_count} retries, could not fetch: {url}')

        if html_content is None: # We failed to fetch the link for some reason, continue to the next job
            continue
        else:
            links = links_matcher.findall(html_content)
            times = times_matcher.findall(html_content)
            sizes = sizes_matcher.findall(html_content)
            for i in range(len(times)):
                dt = datetime.datetime.strptime(times[i].strip(), "%Y-%m-%d %H:%M")
                dt.replace(tzinfo=datetime.timezone.utc)
                times[i] = time.mktime(dt.timetuple())
            for i in range(len(sizes)):
                match sizes[i][1].strip():
                    case '-':
                        sizes[i] = 0
                    case 'K':
                        sizes[i] = int(float(sizes[i][0].strip()) * 1024)
                    case 'M':
                        sizes[i] = int(float(sizes[i][0].strip()) * 1024 * 1024)
                    case 'G':
                        sizes[i] = int(float(sizes[i][0].strip()) * 1024 * 1024 * 1024)
                    case "0":
                        sizes[i] = 0
                    case "":
                        sizes[i] = int(float(sizes[i][0].strip()))
                    case _:
                        raise(ValueError(f"Unexpected symbol while parsing links page: {_}"))

            
            if (len(links) != len(times)) or (len(times) != len(sizes)):
                raise(ValueError("Links parsing error, mismatched numbers of links, times, or sizes!"))

            results = list(zip(links, times, sizes))
            if randomize_order:
                random.shuffle(results)
            for _link, _time, _size in results:
                if _link.endswith("/"):
                    if _link.split(r"/")[-2] in ignore_folder_names:
                        continue
                    else:
                        # print(f"Queueing folder: {urllib.parse.urljoin(url,_link)}")
                        query_work_queue.put(urllib.parse.urljoin(url,_link))
                else:
                    # print(f"Queueing file: {urllib.parse.urljoin(url,_link)}")
                    query_result_queue.put((urllib.parse.urljoin(url,_link), _time, _size))
    

# Fetch file from url and store it to path, retrying on failure
def file_download_worker(download_work_queue, download_result_queue, attempt_count = 2):
    while True:
        job = download_work_queue.get()
        if job == None:
            return
        url, path, t, s = job

        # print(f"Next job: {url}, {path}, {t}, {s}")
        attempts = 0
        while attempts < attempt_count:
            try:
                starttime = time.time()
                req = urllib.request.Request(url, data=None)
                image_data = urllib.request.urlopen(req, timeout=10.0).read()
                endtime = time.time()
                tqdm.tqdm.write(f"Downloaded {url} in {endtime-starttime:0.2f} s | {len(image_data)/1024/1024/(endtime-starttime):0.2f} MB/s")
                if math.isclose(len(image_data), s, rel_tol=0.05):
                    os.makedirs(os.path.split(path)[0], exist_ok=True)
                    open(path, 'wb').write(image_data)
                    download_result_queue.put((True, url, t))
                    break
                else:
                    raise ValueError("Downloaded file is the wrong size!")
            except Exception as e:
                if hasattr(e, "code") and e.code == 404: # This is expected if the file has been removed from the site (at least for swpc.noaa.gov)
                    attempts += attempt_count
                elif (0 < attempts) and (attempts < attempt_count):
                    # tqdm.tqdm.write(f"\nA problem occurred on file: {url} | {e}")
                    time.sleep(1 + random.random())
                    attempts += 1
                if (attempt_count > 1) and (attempts == attempt_count):
                    tqdm.tqdm.write(f'\nAfter {attempt_count} retries, could not fetch: {url}')
                    tqdm.tqdm.write(f"Exception: {e}")
                    download_result_queue.put((False, url, t))
                    break


if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser(
        description="Mirror NOAA GOES SUVI L2 FITS into the local archive.",
    )
    parser.add_argument("--db", default=None, help="SQLite index (default: $SUVI_DB)")
    parser.add_argument("--root", default=None, help="archive root (default: $SUVI_DATA_ROOT)")
    parser.add_argument("--interval", type=int, default=fetch_interval,
                        help="seconds to sleep between passes; 0 runs once and exits")
    args = parser.parse_args()

    if args.root:
        stored_images_dir = os.path.abspath(args.root)
    fetch_interval = args.interval

    # Download bookkeeping lives in the index. It used to be an 800 MB JSON object
    # parsed into memory on every run; see migrate_urlcache.py for the conversion.
    # Losing this state means re-fetching the whole archive, so writes are committed
    # as they happen rather than only at the end of a pass.
    conn = db.connect(args.db)
    known = conn.execute("SELECT count(*) c FROM remote_file").fetchone()["c"]
    print(f"Download records in index: {known}")

    query_work_queue = queue.Queue()
    query_result_queue = queue.Queue()
    download_work_queue = queue.Queue(maxsize=ndownloadworkers)
    download_result_queue = queue.Queue(maxsize=ndownloadworkers)
    workers = []
    for _ in range(nfetchworkers):
        t = Thread(target=find_links_worker, args=(query_work_queue, query_result_queue, 3, randomize_order), daemon=True)
        t.start()
        workers.append(t)
    for _ in range(ndownloadworkers):
        t = Thread(target=file_download_worker, args=(download_work_queue, download_result_queue), daemon=True)
        t.start()
        workers.append(t)

    fetched_image_count = already_had_image_count = failed_image_count = 0
    pending_records = []

    def flush_records(force=False):
        """Commit download records in batches, and always before a long pause."""
        if pending_records and (force or len(pending_records) >= 200):
            db.record_remote_files(conn, pending_records)
            conn.commit()
            pending_records.clear()

    reliever = vfs.Reliever(label="puller")

    def index_frame(local_path):
        reliever.tick()
        """Record a freshly written frame in the index.

        The downloader already knows this file exists, so indexing it here means the
        index never has to be rebuilt by traversing the archive -- which on this
        virtiofs mount is an operation to be avoided rather than merely optimised
        (see suvi/index.py).  Failure to index is not worth losing a download over;
        `filter_FITS.py index` will pick the frame up from its directory mtime.
        """
        try:
            index.record_downloaded(conn, local_path, stored_images_dir)
        except Exception as exc:
            tqdm.tqdm.write(f"Could not index {local_path}: {exc}")

    try:
        while True:
            fetched_image_count = 0
            already_had_image_count = 0
            failed_image_count = 0
            urllen = len(directory_url)

            query_work_queue.put(directory_url)
            for l, t, s in tqdm.tqdm(iter(partial(query_result_queue.get, timeout=30.0), None), desc="Downloading files"):
                # Collect completed jobs and record completion status
                while True:
                    try:
                        r_success, r_url, r_t = download_result_queue.get_nowait()
                        if r_success:
                            fetched_image_count += 1
                            pending_records.append((r_url, r_t, None, None, time.time()))
                            index_frame(os.path.join(stored_images_dir,
                                                     r_url[urllen:].replace("/", os.sep)))
                        else:
                            failed_image_count += 1
                    except queue.Empty:
                        break
                flush_records()

                file_portion_of_link = l[urllen:]
                filepath = os.path.join(stored_images_dir, file_portion_of_link.replace("/", os.sep))
                recorded_mtime = db.get_remote_mtime(conn, l)

                # The archive no longer carries _f/_e suffixes, so a frame is either
                # at its published filename or absent -- no variant probing needed.
                if recorded_mtime is None or t > recorded_mtime:
                    if os.path.exists(filepath):
                        if recorded_mtime is not None:
                            # We have a record, so the remote copy is newer: replace it.
                            try:
                                os.remove(filepath)
                            except OSError:
                                pass
                        else:
                            # No record, but the file is here. Adopt it if the size
                            # matches rather than re-downloading the whole archive.
                            if math.isclose(os.path.getsize(filepath), s, rel_tol=0.05):
                                pending_records.append((l, t, s, filepath, time.time()))
                                index_frame(filepath)
                                already_had_image_count += 1
                                continue
                            tqdm.tqdm.write(f"Size mismatch on {filepath}; redownloading")
                    download_work_queue.put((l, filepath, t, s))
                elif os.path.exists(filepath):
                    already_had_image_count += 1
                else:
                    # Recorded as downloaded but gone from disk; fetch it again.
                    download_work_queue.put((l, filepath, t, s))

            flush_records(force=True)
            print(f"Downloaded {fetched_image_count} | Already had {already_had_image_count} | Failed {failed_image_count}")
            if fetch_interval > 0:
                print(f"Run complete!, sleeping for {fetch_interval} seconds.")
                time.sleep(fetch_interval)
            else:
                print("Run complete!, exiting...")
                break
    except KeyboardInterrupt:
        print("Saving download records and shutting down.")
    except Empty:
        print("Work Complete, Shutting down...")

    for _ in range(nfetchworkers):
        try:
            query_work_queue.put(None)
        except Exception:
            break

    time.sleep(1)
    print(f"Downloaded {fetched_image_count} | Already had {already_had_image_count} | Failed {failed_image_count}")

    while len(download_work_queue.queue) > 0:
        print(f"Waiting for {len(download_work_queue.queue)} downloads in queue...")
        time.sleep(1)

    for _ in range(ndownloadworkers):
        try:
            download_work_queue.put(None, timeout=5.0)
        except Exception:
            break

    # Drain any results the workers finished after the main loop ended.
    while True:
        try:
            r_success, r_url, r_t = download_result_queue.get_nowait()
            if r_success:
                pending_records.append((r_url, r_t, None, None, time.time()))
                index_frame(os.path.join(stored_images_dir,
                                         r_url[len(directory_url):].replace("/", os.sep)))
        except queue.Empty:
            break
    flush_records(force=True)
    conn.close()

    print("Waiting for workers to shutdown...")
    for w in workers:
        w.join(5.0)