freva-client 2508.0.0__py3-none-any.whl → 2509.0.0__py3-none-any.whl

freva_client/query.py

@@ -21,6 +21,7 @@ from typing import (
 
 import intake
 import intake_esm
+import pandas as pd
 import requests
 import xarray as xr
 import yaml
@@ -28,7 +29,12 @@ from rich import print as pprint
 
 from .auth import Auth
 from .utils import logger
-from .utils.auth_utils import Token, choose_token_strategy, load_token
+from .utils.auth_utils import (
+    Token,
+    choose_token_strategy,
+    load_token,
+    requires_authentication,
+)
 from .utils.databrowser_utils import Config, UserDataHandler
 
 __all__ = ["databrowser"]
@@ -218,6 +224,16 @@ class databrowser:
            }
         )
 
+    You can also filter the metadata to only include specific facets.
+
+    .. code-block:: python
+
+        from freva_client import databrowser
+        db = databrowser(
+            "era5*",
+            realm="atmos",
+        )[['project', 'model', 'experiment']]
+        print(db.metadata)
 
     """
 
@@ -225,9 +241,7 @@ class databrowser:
         self,
         *facets: str,
         uniq_key: Literal["file", "uri"] = "file",
-        flavour: Literal[
-            "freva", "cmip6", "cmip5", "cordex", "nextgems", "user"
-        ] = "freva",
+        flavour: Optional[str] = None,
         time: Optional[str] = None,
         host: Optional[str] = None,
         time_select: Literal["flexible", "strict", "file"] = "flexible",
@@ -241,8 +255,9 @@ class databrowser:
         self._auth = Auth()
         self._fail_on_error = fail_on_error
         self._cfg = Config(host, uniq_key=uniq_key, flavour=flavour)
-        self._flavour = flavour
+        self._flavour = self._cfg.flavour
         self._stream_zarr = stream_zarr
+        self.builtin_flavours = {"freva", "cmip6", "cmip5", "cordex", "user"}
         facet_search: Dict[str, List[str]] = defaultdict(list)
         for key, value in search_keys.items():
             if isinstance(value, str):
@@ -289,12 +304,10 @@ class databrowser:
 
     def __iter__(self) -> Iterator[str]:
         query_url = self._cfg.search_url
-        headers = {}
         if self._stream_zarr:
             query_url = self._cfg.zarr_loader_url
-            token = self._auth.authenticate(config=self._cfg)
-            headers = {"Authorization": f"Bearer {token['access_token']}"}
-        result = self._request("GET", query_url, headers=headers, stream=True)
+
+        result = self._request("GET", query_url, stream=True)
         if result is not None:
             try:
                 for res in result.iter_lines():
@@ -368,11 +381,7 @@ class databrowser:
         kwargs: Dict[str, Any] = {"stream": True}
         url = self._cfg.intake_url
         if self._stream_zarr:
-            token = self._auth.authenticate(config=self._cfg)
             url = self._cfg.zarr_loader_url
-            kwargs["headers"] = {
-                "Authorization": f"Bearer {token['access_token']}"
-            }
             kwargs["params"] = {"catalogue-type": "intake"}
         result = self._request("GET", url, **kwargs)
         if result is None:
@@ -539,9 +548,7 @@ class databrowser:
     def count_values(
         cls,
         *facets: str,
-        flavour: Literal[
-            "freva", "cmip6", "cmip5", "cordex", "nextgems", "user"
-        ] = "freva",
+        flavour: Optional[str] = None,
         time: Optional[str] = None,
         host: Optional[str] = None,
         time_select: Literal["flexible", "strict", "file"] = "flexible",
@@ -643,6 +650,15 @@ class databrowser:
             from freva_client import databrowser
             print(databrowser.count_values("reana*", realm="ocean", flavour="cmip6"))
 
+        Count only specific facets:
+
+        .. code-block:: python
+
+            from freva_client import databrowser
+            era5_counts = databrowser.count_values(
+                "era5*",
+            )[['project', 'model']]
+            print(era5_counts)
         """
         this = cls(
             *facets,
@@ -667,7 +683,7 @@ class databrowser:
         return counts
 
     @cached_property
-    def metadata(self) -> Dict[str, List[str]]:
+    def metadata(self) -> pd.DataFrame:
         """Get the metadata (facets) for the current databrowser query.
 
         You can retrieve all information that is associated with your current
@@ -685,20 +701,34 @@ class databrowser:
             db = databrowser(uri="slk:///arch/*/CPC/*")
             print(db.metadata)
 
+        To retrieve only a limited set of metadata you can
+        specify the facets you are interested in:
+
+        .. code-block:: python
+
+            from freva_client import databrowser
+            db = databrowser(
+                "era5*",
+                realm="atmos",
+            )
+            print(db.metadata[['project', 'model', 'experiment']])
+
 
         """
-        return {
-            k: v[::2]
-            for (k, v) in self._facet_search(extended_search=True).items()
-        }
+        return (
+            pd.DataFrame([
+                (k, v[::2])
+                for k, v in self._facet_search(extended_search=True).items()
+            ], columns=["facet", "values"])
+            .explode("values")
+            .groupby("facet")["values"].apply(list)
+        )
 
     @classmethod
     def metadata_search(
         cls,
         *facets: str,
-        flavour: Literal[
-            "freva", "cmip6", "cmip5", "cordex", "nextgems", "user"
-        ] = "freva",
+        flavour: Optional[str] = None,
         time: Optional[str] = None,
         host: Optional[str] = None,
         time_select: Literal["flexible", "strict", "file"] = "flexible",
@@ -708,7 +738,7 @@ class databrowser:
         fail_on_error: bool = False,
         extended_search: bool = False,
         **search_keys: Union[str, List[str]],
-    ) -> Dict[str, List[str]]:
+    ) -> pd.DataFrame:
         """Search for data attributes (facets) in the databrowser.
 
         The method queries the databrowser for available search facets (keys)
@@ -816,6 +846,16 @@ class databrowser:
             res = databrowser.metadata_search(file="/arch/*CPC/*")
             print(res)
 
+        Return only specific facets: for example project and realm:
+
+        .. code-block:: python
+
+            from freva_client import databrowser
+            selected = databrowser.metadata_search(
+                "era5*",
+            )[['project', 'realm']]
+            print(selected)
+
         Sometimes you don't exactly know the exact names of the search keys and
         want retrieve all file objects that match a certain category. For
         example for getting all ocean reanalysis datasets you can apply the
@@ -857,12 +897,18 @@ class databrowser:
             stream_zarr=False,
             **search_keys,
         )
-        return {
-            k: v[::2]
-            for (k, v) in this._facet_search(
-                extended_search=extended_search
-            ).items()
-        }
+        return (
+            pd.DataFrame(
+                [
+                    (k, v[::2])
+                    for k, v in this._facet_search(
+                        extended_search=extended_search
+                    ).items()
+                ], columns=["facet", "values"]
+            )
+            .explode("values")
+            .groupby("facet")["values"].apply(list)
+        )
 
     @classmethod
     def overview(cls, host: Optional[str] = None) -> str:
@@ -894,9 +940,12 @@ class databrowser:
             print(databrowser.overview())
         """
         overview = Config(host).overview.copy()
+        note = overview.pop("Note", None)
+        if note:
+            overview["Note"] = note
         overview["Available search flavours"] = overview.pop("flavours")
         overview["Search attributes by flavour"] = overview.pop("attributes")
-        return yaml.safe_dump(overview)
+        return yaml.safe_dump(overview, sort_keys=False)
 
     @property
     def url(self) -> str:
@@ -1038,8 +1087,6 @@ class databrowser:
                 if result is not None:
                     response_data = result.json()
                     status_message = response_data.get("status")
-                else:
-                    raise ValueError("Failed to add user data")
                 pprint(f"[b][green]{status_message}[green][b]")
             else:
                 raise ValueError("No metadata generated from the input data.")
@@ -1052,11 +1099,153 @@ class databrowser:
                 )
 
             result = this._request("DELETE", url, data=metadata, headers=headers)
-
-            if result is None:
-                raise ValueError("Failed to delete user data")
             pprint("[b][green]User data deleted successfully[green][b]")
 
+    @classmethod
+    def flavour(
+        cls,
+        action: Literal["add", "delete", "list"],
+        name: Optional[str] = None,
+        mapping: Optional[Dict[str, str]] = None,
+        is_global: bool = False,
+        host: Optional[str] = None,
+        fail_on_error: bool = False,
+    ) -> Union[None, Dict[str, Any]]:
+        """Manage custom flavours in the databrowser system.
+
+        This method allows user to add, delete, or list custom flavours that
+        define how search facets are mapped to different Data Reference Syntax
+        (DRS) standards.
+
+        Parameters
+        ~~~~~~~~~~
+        action : Literal["add", "delete", "list"]
+            The action to perform: "add" to create a new flavour, "delete"
+            to remove an existing flavour, or "list" to retrieve all available
+            custom flavours.
+        name : str, optional
+            The name of the flavour to add or delete (required for "add" and
+            "delete" actions).
+        mapping : Dict[str, str], optional
+            A dictionary mapping facet names to their corresponding values in
+            the new flavour (required for "add" action).
+        is_global : bool, default: False
+            Whether to make the flavour available to all users (requires admin
+            privileges) or just the current user.
+        host : str, optional
+            Override the host name of the databrowser server. This is usually
+            the url where the freva web site can be found. Such as
+            www.freva.dkrz.de. By default no host name is given and the host
+            name will be taken from the freva config file.
+        fail_on_error : bool, optional
+            Make the call fail if the connection to the databrowser could not
+            be established.
+
+        Returns
+        ~~~~~~~
+        Union[None, List[Dict[str, Any]]]
+            For "list" action, returns a list of dictionaries containing flavour
+            information. For "add" and "delete" actions, returns None.
+
+        Raises
+        ~~~~~~
+        ValueError
+            If the operation fails, required parameters are missing, or the
+            flavour name conflicts with built-in flavours.
+
+        Example
+        ~~~~~~~
+
+        Adding a custom flavour:
+
+        .. code-block:: python
+
+            from freva_client import databrowser
+            databrowser.flavour(
+                action="add",
+                name="klimakataster",
+                mapping={"project": "Projekt", "model": "Modell"},
+                is_global=False
+            )
+
+        Listing all custom flavours:
+
+        .. code-block:: python
+
+            from freva_client import databrowser
+            flavours = databrowser.flavour(action="list")
+            print(flavours)
+
+        Deleting a custom flavour:
+
+        .. code-block:: python
+
+            from freva_client import databrowser
+            databrowser.flavour(action="delete", name="klimakataster")
+        """
+        this = cls(
+            host=host,
+            fail_on_error=fail_on_error,
+        )
+        cfg = Config(host)
+        if action == "add":
+            token = this._auth.authenticate(config=this._cfg)
+            headers = {"Authorization": f"Bearer {token['access_token']}"}
+            if not name or not mapping:
+                raise ValueError(
+                    "Both 'name' and 'mapping' are required for add action"
+                )
+            payload = {
+                "flavour_name": name,
+                "mapping": mapping,
+                "is_global": is_global,
+            }
+            result = this._request(
+                "POST",
+                f"{this._cfg.databrowser_url}/flavours",
+                data=payload,
+                headers=headers,
+            )
+            if result is not None:
+                msg = result.json().get("status", "Flavour added successfully")
+                pprint(f"[b][green] {msg} [/green][/b]")
+
+        elif action == "delete":
+            token = this._auth.authenticate(config=this._cfg)
+            headers = {"Authorization": f"Bearer {token['access_token']}"}
+            if not name:
+                raise ValueError("'name' is required for delete action")
+            params = {"is_global": "true" if is_global else "false"}
+
+            result = this._request(
+                "DELETE",
+                f"{this._cfg.databrowser_url}/flavours/{name}",
+                headers=headers,
+                params=params,
+            )
+            if result is not None:
+                msg = result.json().get("status", "Flavour deleted successfully")
+                pprint(f"[b][green] {msg} [/green][/b]")
+        elif action == "list":
+            headers = cast(Dict[str, str], this._cfg._get_headers)
+            flavours: List[Dict[str, Any]] = []
+            result = this._request(
+                "GET", f"{cfg.databrowser_url}/flavours", headers=headers
+            )
+            if result is not None:
+                flavours = result.json().get("flavours", [])
+            result_data: Dict[str, Any] = {
+                "flavours": flavours,
+            }
+            if not headers:
+                result_data["Note"] = (
+                    "Displaying only global flavours. "
+                    "Authenticate to see custom user flavours as well."
+                )
+
+            return result_data
+        return None
+
     def _request(
         self,
         method: Literal["GET", "POST", "PUT", "PATCH", "DELETE"],
@@ -1065,10 +1254,23 @@ class databrowser:
         **kwargs: Any,
     ) -> Optional[requests.models.Response]:
         """Request method to handle CRUD operations (GET, POST, PUT, PATCH, DELETE)."""
+        self._cfg.validate_server
         method_upper = method.upper()
         timeout = kwargs.pop("timeout", 30)
         params = kwargs.pop("params", {})
         stream = kwargs.pop("stream", False)
+        kwargs.setdefault("headers", {})
+
+        if (
+            requires_authentication(
+                self._flavour,
+                self._stream_zarr,
+                self._cfg.databrowser_url
+            )
+            and "Authorization" not in kwargs["headers"]
+        ):
+            token = self._auth.authenticate(config=self._cfg)
+            kwargs["headers"]["Authorization"] = f"Bearer {token['access_token']}"
 
         logger.debug(
             "%s request to %s with data: %s and parameters: %s",
@@ -1097,8 +1299,23 @@ class databrowser:
         except (
             requests.exceptions.ConnectionError,
             requests.exceptions.HTTPError,
+            requests.exceptions.InvalidURL,
         ) as error:
-            msg = f"{method_upper} request failed with {error}"
+            server_msg = ""
+            if hasattr(error, 'response') and error.response is not None:
+                try:
+                    error_data = error.response.json()
+                    error_var = {error_data.get(
+                        'detail', error_data.get(
+                            'message', error_data.get('error', '')
+                        )
+                    )}
+                    server_msg = (
+                        f" - {error_var}"
+                    )
+                except Exception:
+                    pass
+            msg = f"{method_upper} request failed with: {error}{server_msg}"
             if self._fail_on_error:
                 raise ValueError(msg) from None
             logger.warning(msg)

freva_client/utils/auth_utils.py

@@ -8,6 +8,7 @@ import time
 from pathlib import Path
 from typing import Literal, Optional, TypedDict, Union, cast
 
+import requests
 from appdirs import user_cache_dir
 
 TOKEN_EXPIRY_BUFFER = 60  # seconds
@@ -238,3 +239,40 @@ def wait_for_port(host: str, port: int, timeout: float = 5.0) -> None:
     raise TimeoutError(
         f"Port {port} on {host} did not open within {timeout} seconds."
     )
+
+
+def requires_authentication(
+        flavour: Optional[str],
+        zarr: bool = False,
+        databrowser_url: Optional[str] = None
+) -> bool:
+    """Check if authentication is required.
+
+    Parameters
+    ----------
+    flavour : str or None
+        The data flavour to check.
+    zarr : bool, default: False
+        Whether the request is for zarr data.
+    databrowser_url : str or None
+        The URL of the databrowser to query for available flavours.
+        If None, the function will skip querying and assume authentication
+        is required for non-default flavours.
+    """
+    if zarr:
+        return True
+    if flavour in {"freva", "cmip6", "cmip5", "cordex", "user", None}:
+        return False
+    try:
+        response = requests.get(f"{databrowser_url}/flavours", timeout=30)
+        response.raise_for_status()
+        result = {"flavours": response.json().get("flavours", [])}
+        if "flavours" in result:
+            global_flavour_names = {
+                f["flavour_name"] for f in result["flavours"]
+            }
+            return flavour not in global_flavour_names
+    except Exception:
+        pass
+
+    return True