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

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"]
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 parse_listing(html_content):
    """Parse one Apache index page into [(link, mtime, size), ...].

    Listing timestamps are UTC, so they are read as UTC.  Reading them as local
    time -- which is what a bare ``mktime`` does -- makes every recorded mtime an
    artefact of the machine's timezone, and moving the archive to a box further
    east would then make the whole archive look newer than its records and be
    re-downloaded on top of itself.

    Raises ValueError if the page does not parse, which the caller treats as a
    directory it failed to list rather than as a fatal error.
    """
    links = links_matcher.findall(html_content)
    times = times_matcher.findall(html_content)
    sizes = sizes_matcher.findall(html_content)

    if (len(links) != len(times)) or (len(times) != len(sizes)):
        raise ValueError(
            f"mismatched links/times/sizes: {len(links)}/{len(times)}/{len(sizes)}"
        )

    for i in range(len(times)):
        dt = datetime.datetime.strptime(times[i].strip(), "%Y-%m-%d %H:%M")
        times[i] = dt.replace(tzinfo=datetime.timezone.utc).timestamp()
    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 size symbol while parsing links page: {_}")

    return list(zip(links, times, sizes))


def find_links_worker(query_work_queue, query_result_queue, crawl_errors,
                      attempt_count = 3, randomize_order = True):
    """Crawl directory listings, queueing subdirectories and file records.

    Every directory this fails to fetch or parse is appended to `crawl_errors`
    instead of raising.  An exception here used to escape into a daemon thread and
    kill it silently; with only `nfetchworkers` of them, one malformed page could
    end the crawl while the run still reported success.
    """
    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
            crawl_errors.append((url, "could not be fetched"))
            continue

        try:
            results = parse_listing(html_content)
        except Exception as e:
            tqdm.tqdm.write(f"\nCould not parse listing {url}: {e}")
            crawl_errors.append((url, f"could not be parsed: {e}"))
            continue

        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):
    """Fetch each queued URL, writing it into place only once it is complete.

    Exactly `attempt_count` attempts are made and exactly one result is reported,
    so a URL that keeps failing releases the worker instead of holding it.  The
    counter used to go un-incremented on the first failure, which turned any
    non-404 error into an unbounded, undelayed retry loop: three such URLs pinned
    all three workers and stalled the pass with no output.

    The body is written to a sibling `.part` file and renamed over the target, so
    a frame is only ever replaced by a complete download.  An interrupted fetch
    leaves the existing copy intact.
    """
    while True:
        job = download_work_queue.get()
        if job == None:
            return
        url, path, t, s = job

        # print(f"Next job: {url}, {path}, {t}, {s}")
        for attempt in range(attempt_count):
            partial = path + ".part"
            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 not math.isclose(len(image_data), s, rel_tol=0.05):
                    raise ValueError(f"Downloaded file is the wrong size: {len(image_data)} bytes, listing said {s}")
                os.makedirs(os.path.split(path)[0], exist_ok=True)
                with open(partial, 'wb') as handle:
                    handle.write(image_data)
                os.replace(partial, path)
                download_result_queue.put((True, url, t))
                break
            except Exception as e:
                try:
                    os.remove(partial)
                except OSError:
                    pass
                # A 404 is expected once a file has been withdrawn from the site, and
                # retrying will not bring it back.
                expired = getattr(e, "code", None) == 404
                if not expired and (attempt + 1 < attempt_count):
                    # tqdm.tqdm.write(f"\nA problem occurred on file: {url} | {e}")
                    time.sleep(1 + random.random())
                    continue
                if not expired:
                    tqdm.tqdm.write(f'\nAfter {attempt_count} attempts, 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}")

    # Listing times are now read as UTC; the records written before that fix hold
    # local-time epochs.  Until migrate_mtime_utc.py has normalised them, every
    # recorded mtime reads as later than the listing says, so an updated file is
    # skipped rather than re-fetched.  That is the harmless direction to fail in --
    # nothing is deleted or re-downloaded -- but it is not a state to sit in quietly.
    if known and db.get_meta(conn, "remote_mtime_utc") != "done":
        print("WARNING: remote_file mtimes have not been migrated to UTC.")
        print("         Updated files will not be re-fetched until you run:")
        print("             migrate_mtime_utc.py --apply     (with the puller stopped)")

    crawl_errors = []
    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, crawl_errors, 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 = []
    urllen = len(directory_url)

    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()

    def index_frame(local_path):
        """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}")

    def drain_results():
        """Record every download the workers have finished reporting on."""
        global fetched_image_count, failed_image_count
        while True:
            try:
                r_success, r_url, r_t = download_result_queue.get_nowait()
            except queue.Empty:
                return
            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

    try:
        while True:
            fetched_image_count = 0
            already_had_image_count = 0
            failed_image_count = 0
            crawl_errors.clear()

            query_work_queue.put(directory_url)
            # A quiet result queue is how a pass ends: nobody queues a sentinel, so
            # the crawl is over once no worker has produced a file for 30 seconds.
            # Catching that here rather than around the whole loop is what lets a
            # pass report its own result -- and lets --interval actually repeat.
            try:
                for l, t, s in tqdm.tqdm(iter(partial(query_result_queue.get, timeout=30.0), None), desc="Downloading files"):
                    drain_results()
                    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 None:
                                # 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")
                            # Otherwise the remote copy is newer.  The existing file is
                            # left alone: the worker renames the replacement over it once
                            # the download is complete, so a failed fetch costs us nothing.
                        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))
            except Empty:
                pass

            drain_results()
            flush_records(force=True)
            print(f"Downloaded {fetched_image_count} | Already had {already_had_image_count} | Failed {failed_image_count}")

            # Say so when the pass did not cover the archive.  Unlisted directories
            # are invisible in the counts above -- their files were never considered
            # at all -- so without this an aborted crawl looks like a clean one.
            if crawl_errors:
                print(f"WARNING: pass was incomplete -- {len(crawl_errors)} directories could not be listed:")
                for bad_url, reason in crawl_errors[:10]:
                    print(f"    {bad_url}  ({reason})")
                if len(crawl_errors) > 10:
                    print(f"    ... and {len(crawl_errors) - 10} more")
            if not query_work_queue.empty():
                print(f"WARNING: pass ended with {query_work_queue.qsize()} directories still queued; "
                      "the crawl was quiet for 30 s while work remained.")

            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.")

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

    # Shut the download workers down, draining results throughout: a worker that
    # finishes while the result queue is full would otherwise block there forever,
    # and a frame whose result is never read is a frame we re-discover next pass.
    print("Waiting for in-flight downloads...")
    deadline = time.time() + 60.0
    sentinels = ndownloadworkers
    while sentinels and time.time() < deadline:
        drain_results()
        try:
            download_work_queue.put(None, timeout=1.0)
            sentinels -= 1
        except queue.Full:
            continue

    while any(w.is_alive() for w in workers) and time.time() < deadline:
        drain_results()
        time.sleep(0.5)
    drain_results()

    flush_records(force=True)
    conn.close()
    print(f"Downloaded {fetched_image_count} | Already had {already_had_image_count} | Failed {failed_image_count}")

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