import logging
import pickle
import warnings
from datetime import datetime, timedelta, timezone
from functools import wraps
import numpy as np
from dateutil import parser
from email.utils import format_datetime
from packaging.version import Version
from speasy.config import inventories as inventories_cfg
from speasy.config import proxy as proxy_cfg
from .. import http, make_utc_datetime
from ..index import index
from ..inventory.indexes import from_dict as inventory_from_dict
from ... import SpeasyIndex
from ...products.variable import from_dictionary as var_from_dict
from ..cache import CacheCall
log = logging.getLogger(__name__)
PROXY_ALLOWED_KWARGS = ['disable_proxy']
MINIMUM_REQUIRED_PROXY_VERSION = Version("0.13.0")
_CURRENT_PROXY_SERVER_VERSION = None
if proxy_cfg.url() == "" or proxy_cfg.enabled() == False:
warnings.warn("""Proxy server is disabled you might want to use it both to improve Speasy performances and to reduce pressure on remote servers
use the following python snippet to configure proxy server:
===========================================================================
import speasy as spz
spz.config.proxy.url.set("https://sciqlop.lpp.polytechnique.fr/cache")
spz.config.proxy.enabled.set(True)
===========================================================================
""", stacklevel=0)
[docs]
class ProxyError(Exception):
pass
[docs]
def query_proxy_version():
global _CURRENT_PROXY_SERVER_VERSION
if _CURRENT_PROXY_SERVER_VERSION is None:
url = proxy_cfg.url()
if url.endswith("/"):
url = url[:-1]
if url != "":
resp = http.get(f"{url}/get_version")
if resp.status_code == 200:
_CURRENT_PROXY_SERVER_VERSION = Version(resp.text.strip())
return _CURRENT_PROXY_SERVER_VERSION
return _CURRENT_PROXY_SERVER_VERSION
try:
import pyzstd
except ImportError:
pyzstd = None
try:
from numcodecs import Blosc
except ImportError:
Blosc = None
zstd_compression = 'true' if pyzstd is not None else 'false'
def _raw_view(array: np.ndarray) -> np.ndarray:
# datetime64/timedelta64 don't expose the buffer protocol, the proxy sends their bytes as int64
return array.view("i8") if array.dtype.kind in "mM" else array
def _unblosc_array(marker: dict) -> np.ndarray:
array = np.empty(marker["shape"], dtype=marker["dtype"])
Blosc().decode(marker["__blosc__"], out=_raw_view(array))
return array
[docs]
def unblosc_arrays(obj):
"""Replace every {"__blosc__": chunk, "dtype", "shape"} marker sent by the proxy by its numpy array."""
if isinstance(obj, dict):
return _unblosc_array(obj) if "__blosc__" in obj else {k: unblosc_arrays(v) for k, v in obj.items()}
if isinstance(obj, list):
return [unblosc_arrays(v) for v in obj]
return obj
_DECODERS = {
"application/x-zstd-compressed": lambda body: pickle.loads(pyzstd.decompress(body)),
"application/x-speasy-blosc-pickle": lambda body: unblosc_arrays(pickle.loads(body)),
}
[docs]
def load_response(resp):
"""Unpickle a proxy response according to its Content-Type: the server may ignore the compression we asked for."""
# pickle has always been the proxy wire format; the proxy is a server the user configured and trusts
mime = resp.headers.get("Content-Type", "").split(";")[0].strip()
return _DECODERS.get(mime, pickle.loads)(resp.bytes)
[docs]
@CacheCall(cache_retention=timedelta(minutes=10), is_pure=True)
def is_proxy_up() -> bool:
if http.is_server_up(proxy_cfg.url()):
try:
url = proxy_cfg.url()
if url.endswith("/"):
url = url[:-1]
r = http.get(f"{url}/get_inventory", params={"provider": "ssc"}, timeout=1)
return r.status_code == 200
except: # lgtm [py/catch-base-exception]
pass
return False
[docs]
class GetProduct:
def __init__(self):
pass
[docs]
@staticmethod
def get(path: str, start_time: str, stop_time: str, **kwargs):
url = proxy_cfg.url()
if url.endswith("/"):
url = url[:-1]
kwargs['path'] = path
kwargs['start_time'] = start_time
kwargs['stop_time'] = stop_time
kwargs['format'] = 'python_dict'
kwargs['zstd_compression'] = zstd_compression
if Blosc is not None:
kwargs['compression'] = 'blosc'
resp = http.get(f"{url}/get_data", params=kwargs, timeout=60 * 5)
log.debug(f"Asking data from proxy {resp.url}, {resp.headers}")
if resp.status_code == 200:
var = var_from_dict(load_response(resp))
return var
return None
[docs]
class GetInventory:
[docs]
@staticmethod
def get(provider: str, **kwargs):
saved_inventory: SpeasyIndex = index.get("proxy_inventories", provider, None)
saved_inventory_dt: datetime = make_utc_datetime(
index.get("proxy_inventories_save_date", provider, datetime.fromtimestamp(0, tz=timezone.utc)))
if saved_inventory is not None and \
saved_inventory_dt + timedelta(days=inventories_cfg.cache_retention_days.get()) > datetime.now(
tz=timezone.utc):
return saved_inventory
url = proxy_cfg.url()
if url.endswith("/"):
url = url[:-1]
kwargs['provider'] = provider
kwargs['format'] = 'python_dict'
kwargs['zstd_compression'] = zstd_compression
kwargs['version'] = 2
headers = {}
if saved_inventory is not None:
# If-Modified-Since must be a valid HTTP-date (RFC 7231), datetime.ctime() is not one
headers["If-Modified-Since"] = format_datetime(
make_utc_datetime(parser.parse(saved_inventory.build_date)), usegmt=True)
resp = http.get(f"{url}/get_inventory", params=kwargs, headers=headers)
log.debug(f"Asking {provider} inventory from proxy {resp.url}, {resp.headers}")
if resp.status_code == 200:
inventory = inventory_from_dict(load_response(resp), version=2)
index.set("proxy_inventories", provider, inventory)
index.set("proxy_inventories_save_date", provider, datetime.now(tz=timezone.utc))
return inventory
if resp.status_code == 304:
return saved_inventory
raise ProxyError(f"Can't get inventory for provider {provider} from proxy server {proxy_cfg.url()}, status code {resp.status_code}")
[docs]
class Proxyfiable(object):
def __init__(self, request, arg_builder, min_version=MINIMUM_REQUIRED_PROXY_VERSION):
self.request = request
self.arg_builder = arg_builder
if not isinstance(min_version, Version):
min_version = Version(min_version)
self.min_version = max(min_version, MINIMUM_REQUIRED_PROXY_VERSION)
def __call__(self, func):
@wraps(func)
def wrapped(*args, **kwargs):
disable_proxy = kwargs.pop("disable_proxy", False)
min_version = kwargs.pop("min_proxy_version", self.min_version)
if proxy_cfg.enabled() and not disable_proxy and is_proxy_up():
try:
proxy_version = query_proxy_version()
if proxy_version is not None and proxy_version >= min_version:
return self.request.get(**self.arg_builder(**kwargs))
else:
log.warning(
f"You are using an incompatible proxy server {proxy_cfg.url()} which is {proxy_version} while minimun required version is {min_version}")
except Exception as e: # lgtm [py/catch-base-exception]
log.error(f"Can't get data from proxy server {proxy_cfg.url()}")
print(e)
return func(*args, **kwargs)
return wrapped