xopr.stac_cache

  1"""
  2STAC catalog caching utilities for xopr.
  3
  4This module provides functions to cache STAC GeoParquet catalogs locally,
  5reducing network latency for repeated queries.
  6"""
  7
  8import json
  9import os
 10import shutil
 11import xml.etree.ElementTree as ET
 12from pathlib import Path
 13from typing import Optional
 14
 15import requests
 16from platformdirs import user_cache_dir
 17
 18# OPR catalog constants
 19OPR_CATALOG_S3_PREFIX = "englacial/xopr/catalog/"
 20OPR_CATALOG_S3_GLOB = (
 21    "s3://us-west-2.opendata.source.coop/englacial/xopr/catalog/**/*.parquet"
 22)
 23S3_LIST_URL = "https://s3.us-west-2.amazonaws.com/us-west-2.opendata.source.coop"
 24OPR_CATALOG_HTTPS_BASE = "https://data.source.coop/"
 25
 26# Cloud URLs for bedmap catalogs
 27BEDMAP_CATALOG_BASE_URL = "https://data.source.coop/englacial/bedmap"
 28BEDMAP_CATALOG_FILES = ["bedmap1.parquet", "bedmap2.parquet", "bedmap3.parquet"]
 29
 30
 31def get_cache_dir() -> Path:
 32    """
 33    Get the xopr cache directory.
 34
 35    Checks $XOPR_CACHE_DIR environment variable first, otherwise uses
 36    platform-specific user cache directory.
 37
 38    Returns
 39    -------
 40    Path
 41        Path to xopr cache directory
 42    """
 43    env_cache = os.environ.get("XOPR_CACHE_DIR")
 44    if env_cache:
 45        cache_path = Path(env_cache)
 46    else:
 47        cache_path = Path(user_cache_dir("xopr", "englacial"))
 48
 49    return cache_path
 50
 51
 52def get_bedmap_catalog_dir() -> Path:
 53    """
 54    Get the bedmap catalog cache directory.
 55
 56    Returns
 57    -------
 58    Path
 59        Path to bedmap catalog directory within cache
 60    """
 61    return get_cache_dir() / "catalogs" / "bedmap"
 62
 63
 64def _download_file(url: str, dest: Path) -> bool:
 65    """
 66    Download a file from URL to destination path.
 67
 68    Parameters
 69    ----------
 70    url : str
 71        URL to download from
 72    dest : Path
 73        Destination file path
 74
 75    Returns
 76    -------
 77    bool
 78        True if download succeeded, False otherwise
 79    """
 80    try:
 81        response = requests.get(url, stream=True, timeout=30)
 82        response.raise_for_status()
 83
 84        dest.parent.mkdir(parents=True, exist_ok=True)
 85
 86        with open(dest, "wb") as f:
 87            for chunk in response.iter_content(chunk_size=8192):
 88                f.write(chunk)
 89
 90        return True
 91    except Exception as e:
 92        print(f"Warning: Failed to download {url}: {e}")
 93        return False
 94
 95
 96def ensure_bedmap_catalogs(force_download: bool = False) -> Optional[Path]:
 97    """
 98    Ensure bedmap catalogs are cached locally, downloading if needed.
 99
100    Parameters
101    ----------
102    force_download : bool, default False
103        If True, re-download catalogs even if they exist
104
105    Returns
106    -------
107    Path or None
108        Path to catalog directory if successful, None if download failed
109    """
110    catalog_dir = get_bedmap_catalog_dir()
111
112    # Check if all catalogs exist
113    all_exist = all((catalog_dir / f).exists() for f in BEDMAP_CATALOG_FILES)
114
115    if all_exist and not force_download:
116        return catalog_dir
117
118    # Download missing catalogs
119    print(f"Downloading bedmap catalogs to {catalog_dir}...")
120    catalog_dir.mkdir(parents=True, exist_ok=True)
121
122    success = True
123    for filename in BEDMAP_CATALOG_FILES:
124        dest = catalog_dir / filename
125        if dest.exists() and not force_download:
126            continue
127
128        url = f"{BEDMAP_CATALOG_BASE_URL}/{filename}"
129        if not _download_file(url, dest):
130            success = False
131
132    if success:
133        print(f"Bedmap catalogs cached successfully")
134        return catalog_dir
135    else:
136        print("Warning: Some catalogs failed to download")
137        # Return catalog_dir anyway - partial cache may still be useful
138        return catalog_dir if any((catalog_dir / f).exists() for f in BEDMAP_CATALOG_FILES) else None
139
140
141def get_bedmap_catalog_path() -> str:
142    """
143    Get the path pattern for bedmap catalogs, downloading if needed.
144
145    This is the main entry point for query functions. It ensures catalogs
146    are cached locally and returns the glob pattern for querying.
147
148    Returns
149    -------
150    str
151        Glob pattern to local bedmap catalog files, or cloud URL as fallback
152    """
153    catalog_dir = ensure_bedmap_catalogs()
154
155    if catalog_dir and any((catalog_dir / f).exists() for f in BEDMAP_CATALOG_FILES):
156        return str(catalog_dir / "bedmap*.parquet")
157    else:
158        # Fallback to cloud URL if local cache failed
159        print("Warning: Using cloud catalogs (local cache unavailable)")
160        return f"{BEDMAP_CATALOG_BASE_URL}/bedmap*.parquet"
161
162
163def clear_bedmap_cache() -> None:
164    """
165    Clear cached bedmap catalogs.
166
167    Useful for forcing a fresh download of catalogs.
168    """
169    catalog_dir = get_bedmap_catalog_dir()
170    if catalog_dir.exists():
171        for f in BEDMAP_CATALOG_FILES:
172            path = catalog_dir / f
173            if path.exists():
174                path.unlink()
175        print(f"Cleared bedmap catalog cache at {catalog_dir}")
176
177
178# ---------------------------------------------------------------------------
179# OPR catalog caching
180# ---------------------------------------------------------------------------
181
182
183def get_opr_catalog_dir() -> Path:
184    """
185    Get the OPR catalog cache directory.
186
187    Returns
188    -------
189    Path
190        Path to OPR catalog directory within cache
191    """
192    return get_cache_dir() / "catalogs" / "opr"
193
194
195def _list_remote_opr_catalogs() -> list[dict]:
196    """
197    List remote OPR STAC catalog files via the S3 ListBucketV2 API.
198
199    Returns
200    -------
201    list[dict]
202        Each dict has keys ``key``, ``etag``, and ``size``.
203    """
204    ns = "{http://s3.amazonaws.com/doc/2006-03-01/}"
205    results: list[dict] = []
206    continuation_token = None
207
208    while True:
209        params = {"list-type": "2", "prefix": OPR_CATALOG_S3_PREFIX}
210        if continuation_token:
211            params["continuation-token"] = continuation_token
212
213        resp = requests.get(S3_LIST_URL, params=params, timeout=30)
214        resp.raise_for_status()
215
216        root = ET.fromstring(resp.text)
217        for content in root.findall(f"{ns}Contents"):
218            key = content.findtext(f"{ns}Key", "")
219            if key.endswith(".parquet"):
220                results.append({
221                    "key": key,
222                    "etag": content.findtext(f"{ns}ETag", "").strip('"'),
223                    "size": int(content.findtext(f"{ns}Size", "0")),
224                })
225
226        if root.findtext(f"{ns}IsTruncated", "false") == "true":
227            continuation_token = root.findtext(f"{ns}NextContinuationToken")
228        else:
229            break
230
231    return results
232
233
234def _load_opr_manifest() -> dict:
235    """
236    Load the local OPR manifest (maps relative path -> etag/size).
237
238    Returns
239    -------
240    dict
241        Manifest dictionary, or empty dict if missing.
242    """
243    manifest_path = get_opr_catalog_dir() / "_manifest.json"
244    if manifest_path.exists():
245        with open(manifest_path) as f:
246            return json.load(f)
247    return {}
248
249
250def _save_opr_manifest(manifest: dict) -> None:
251    """
252    Persist the OPR manifest to disk.
253
254    Parameters
255    ----------
256    manifest : dict
257        Manifest mapping relative paths to ``{"etag", "size"}`` dicts.
258    """
259    catalog_dir = get_opr_catalog_dir()
260    catalog_dir.mkdir(parents=True, exist_ok=True)
261    manifest_path = catalog_dir / "_manifest.json"
262    with open(manifest_path, "w") as f:
263        json.dump(manifest, f)
264
265
266def sync_opr_catalogs() -> None:
267    """
268    Sync OPR STAC catalog parquet files to local cache.
269
270    Compares remote ETags against a local manifest and only downloads
271    new or changed files.  Designed to run in a background thread;
272    silently returns on network errors.
273    """
274    try:
275        remote_files = _list_remote_opr_catalogs()
276    except Exception:
277        return  # silent failure on network errors
278
279    catalog_dir = get_opr_catalog_dir()
280    manifest = _load_opr_manifest()
281    changed = False
282
283    for entry in remote_files:
284        key = entry["key"]
285        # Relative path under the catalog dir (strip the S3 prefix)
286        rel = key[len(OPR_CATALOG_S3_PREFIX):]
287        cached = manifest.get(rel)
288        if cached and cached.get("etag") == entry["etag"]:
289            continue  # unchanged
290
291        # Download via HTTPS with atomic write
292        url = f"{OPR_CATALOG_HTTPS_BASE}{key}"
293        dest = catalog_dir / rel
294        tmp = dest.with_suffix(".tmp")
295        try:
296            resp = requests.get(url, stream=True, timeout=60)
297            resp.raise_for_status()
298            dest.parent.mkdir(parents=True, exist_ok=True)
299            with open(tmp, "wb") as f:
300                for chunk in resp.iter_content(chunk_size=8192):
301                    f.write(chunk)
302            os.replace(tmp, dest)
303            manifest[rel] = {"etag": entry["etag"], "size": entry["size"]}
304            changed = True
305        except Exception:
306            if tmp.exists():
307                tmp.unlink()
308            continue  # skip this file, keep going
309
310    if changed:
311        _save_opr_manifest(manifest)
312
313
314def get_opr_catalog_path() -> str:
315    """
316    Get a path/glob for OPR STAC catalogs, preferring local cache.
317
318    Returns
319    -------
320    str
321        Local glob pattern if cached files exist, otherwise the S3 glob.
322    """
323    catalog_dir = get_opr_catalog_dir()
324    if catalog_dir.exists() and any(catalog_dir.rglob("*.parquet")):
325        return str(catalog_dir / "**" / "*.parquet")
326    return OPR_CATALOG_S3_GLOB
327
328
329def clear_opr_cache() -> None:
330    """
331    Remove the entire OPR catalog cache directory.
332    """
333    catalog_dir = get_opr_catalog_dir()
334    if catalog_dir.exists():
335        shutil.rmtree(catalog_dir)
336        print(f"Cleared OPR catalog cache at {catalog_dir}")