metadata-crawler 2509.0.0__py3-none-any.whl

metadata_crawler/logger.py

@@ -0,0 +1,142 @@
+"""Logging utilities."""
+
+import logging
+import logging.config
+import os
+from logging.handlers import RotatingFileHandler
+from pathlib import Path
+from typing import Any, Optional, cast
+
+import appdirs
+from rich.console import Console
+from rich.logging import RichHandler
+
+THIS_NAME = "data-crawler"
+
+logging.basicConfig(
+    level=logging.WARNING,
+    format="%(asctime)s %(levelname)s: %(name)s - %(message)s",
+)
+
+logging.config.dictConfig(
+    {
+        "version": 1,
+        # keep existing handlers
+        "disable_existing_loggers": False,
+        "root": {
+            "level": "WARNING",
+            "handlers": ["default"],
+        },
+        "formatters": {
+            "standard": {
+                "format": "%(asctime)s %(levelname)s: %(name)s - %(message)s",
+            },
+        },
+        "handlers": {
+            "default": {
+                "class": "logging.StreamHandler",
+                "formatter": "standard",
+                "level": "WARNING",
+            },
+        },
+    }
+)
+
+logging.getLogger("sqlalchemy").setLevel(logging.WARNING)
+logging.getLogger("sqlalchemy.engine").setLevel(logging.WARNING)
+logging.getLogger("sqlalchemy.pool").setLevel(logging.WARNING)
+
+
+class Logger(logging.Logger):
+    """Custom Logger defining the logging behaviour."""
+
+    logfmt: str = "%(name)s: %(message)s"
+    filelogfmt: str = "%(asctime)s %(levelname)s: %(name)s - %(message)s"
+    datefmt: str = "%Y-%m-%dT%H:%M:%S"
+    no_debug: list[str] = ["watchfiles", "httpcore", "pymongo", "pika"]
+
+    def __init__(
+        self, name: Optional[str] = None, level: Optional[int] = None
+    ) -> None:
+        """Instantiate this logger only once and for all."""
+        level = level or int(
+            cast(str, os.getenv("MDC_LOG_LEVEL", str(logging.WARNING)))
+        )
+        name = name or THIS_NAME
+        logger_format = logging.Formatter(self.logfmt, self.datefmt)
+        self.file_format = logging.Formatter(self.filelogfmt, self.datefmt)
+        self._logger_file_handle: Optional[RotatingFileHandler] = None
+        self._logger_stream_handle = RichHandler(
+            rich_tracebacks=True,
+            tracebacks_max_frames=10,
+            tracebacks_extra_lines=5,
+            show_path=True,
+            console=Console(
+                soft_wrap=False,
+                force_jupyter=False,
+                stderr=True,
+            ),
+        )
+        self._logger_stream_handle.setFormatter(logger_format)
+        self._logger_stream_handle.setLevel(level)
+        super().__init__(name, level)
+
+        self.propagate = False
+        self.handlers = [self._logger_stream_handle]
+
+    def set_level(self, level: int) -> None:
+        """Set the logger level to level."""
+        for handler in self.handlers:
+            log_level = level
+            if isinstance(handler, RotatingFileHandler):
+                log_level = min(level, logging.CRITICAL)
+            handler.setLevel(log_level)
+        self.setLevel(level)
+        logger.level = level
+
+    def error(
+        self,
+        msg: object,
+        *args: Any,
+        **kwargs: Any,
+    ) -> None:
+        """Log an error. When log level is smaller than INFO, log exceptions."""
+        if self.level < logging.INFO:
+            kwargs.setdefault("exc_info", True)
+        self._log(logging.ERROR, msg, args, **kwargs)
+
+
+logger = Logger()
+
+
+def add_file_handle(
+    suffix: Optional[str], log_level: int = logging.CRITICAL
+) -> None:
+    """Add a file log handle to the logger."""
+    base_name = f"{THIS_NAME}-{suffix}" if suffix else THIS_NAME
+    log_dir = Path(appdirs.user_log_dir(THIS_NAME))
+    log_dir.mkdir(exist_ok=True, parents=True)
+    logger_file_handle = RotatingFileHandler(
+        log_dir / f"{base_name}.log",
+        mode="a",
+        maxBytes=5 * 1024**2,
+        backupCount=5,
+        encoding="utf-8",
+        delay=False,
+    )
+    logger_file_handle.setFormatter(logger.file_format)
+    logger_file_handle.setLevel(min(log_level, logging.CRITICAL))
+    logger.addHandler(logger_file_handle)
+
+
+def get_level_from_verbosity(verbosity: int) -> int:
+    """Calculate the log level from a verbosity."""
+    return max(logging.CRITICAL - 10 * verbosity, -1)
+
+
+def apply_verbosity(level: int) -> int:
+    """Set the logging level of the handlers to a certain level."""
+    old_level = logger.level
+    level = get_level_from_verbosity(level)
+    logger.set_level(level)
+    return old_level

metadata_crawler/run.py

@@ -0,0 +1,373 @@
+"""Apply the metadata collector."""
+
+import asyncio
+import os
+import time
+from fnmatch import fnmatch
+from pathlib import Path
+from types import NoneType
+from typing import Any, Collection, Dict, List, Optional, Sequence, Union, cast
+
+import tomlkit
+from rich.prompt import Prompt
+
+from .api.config import CrawlerSettings, DRSConfig, strip_protocol
+from .api.metadata_stores import CatalogueBackendType, IndexName
+from .data_collector import DataCollector
+from .logger import apply_verbosity, get_level_from_verbosity, logger
+from .utils import (
+    Console,
+    EmptyCrawl,
+    MetadataCrawlerException,
+    find_closest,
+    load_plugins,
+    timedelta_to_str,
+)
+
+FilesArg = Optional[Union[str, Path, Sequence[Union[str, Path]]]]
+
+
+def _norm_files(catalogue_files: FilesArg) -> List[str]:
+    if catalogue_files is None:
+        return [""]
+    return (
+        [str(catalogue_files)]
+        if isinstance(catalogue_files, (str, Path))
+        else [str(p) for p in catalogue_files]
+    )
+
+
+def _match(match: str, items: Collection[str]) -> List[str]:
+    out: List[str] = []
+    for item in items:
+        if fnmatch(item, match):
+            out.append(item)
+
+    if not out:
+        msg = find_closest(f"No such dataset: {match}", match, items)
+        raise MetadataCrawlerException(msg) from None
+    return out
+
+
+def _get_search(
+    config_file: Union[str, Path, Dict[str, Any], tomlkit.TOMLDocument],
+    search_dirs: Optional[List[str]] = None,
+    datasets: Optional[List[str]] = None,
+) -> list[CrawlerSettings]:
+    _search_items = []
+    search_dirs = search_dirs or []
+    datasets = datasets or []
+    config = DRSConfig.load(config_file).datasets
+    if not datasets and not search_dirs:
+        return [
+            CrawlerSettings(name=k, search_path=cfg.root_path)
+            for (k, cfg) in config.items()
+        ]
+    for item in datasets or []:
+        for ds in _match(item, config.keys()):
+            logger.debug("Adding dataset %s", ds)
+            _search_items.append(
+                CrawlerSettings(name=ds, search_path=config[ds].root_path)
+            )
+    for num, _dir in enumerate(map(strip_protocol, search_dirs or [])):
+        for name, cfg in config.items():
+            if _dir.is_relative_to(strip_protocol(cfg.root_path)):
+                logger.debug("Adding dataset %s", name)
+                _search_items.append(
+                    CrawlerSettings(name=name, search_path=str(search_dirs[num]))
+                )
+
+    return _search_items
+
+
+async def async_call(
+    index_system: str,
+    method: str,
+    batch_size: int = 2500,
+    catalogue_files: Optional[Sequence[Union[Path, str]]] = None,
+    verbosity: int = 0,
+    *args: Any,
+    **kwargs: Any,
+) -> None:
+    """Index metadata."""
+    env = cast(os._Environ[str], os.environ.copy())
+    old_level = apply_verbosity(verbosity)
+    try:
+        os.environ["MDC_LOG_LEVEL"] = str(get_level_from_verbosity(verbosity))
+        backends = load_plugins("metadata_crawler.ingester")
+        try:
+            cls = backends[index_system]
+        except KeyError:
+            msg = find_closest(
+                f"No such backend: {index_system}", index_system, backends.keys()
+            )
+            raise ValueError(msg) from None
+        flat_files = _norm_files(catalogue_files)
+        _event_loop = asyncio.get_event_loop()
+        flat_files = flat_files or [""]
+        futures = []
+        storage_options = kwargs.pop("storage_options", {})
+        for cf in flat_files:
+            obj = cls(
+                batch_size=batch_size,
+                catalogue_file=cf or None,
+                storage_options=storage_options,
+            )
+            func = getattr(obj, method)
+            future = _event_loop.create_task(func(**kwargs))
+            futures.append(future)
+        await asyncio.gather(*futures)
+    finally:
+        os.environ = env
+        logger.set_level(old_level)
+
+
+async def async_index(
+    index_system: str,
+    *catalogue_files: Union[Path, str, List[str], List[Path]],
+    batch_size: int = 2500,
+    verbosity: int = 0,
+    **kwargs: Any,
+) -> None:
+    """Index metadata in the indexing system.
+
+    Parameters
+    ^^^^^^^^^^
+
+    index_system:
+        The index server where the metadata is indexed.
+    catalogue_file:
+        Path to the file where the metadata was stored.
+    batch_size:
+        If the index system supports batch-sizes, the size of the batches.
+    verbosity:
+        Set the verbosity of the system.
+
+    Other Parameters
+    ^^^^^^^^^^^^^^^^
+
+    **kwargs:
+        Keyword arguments used to delete data from the index.
+
+
+    Example
+    ^^^^^^^
+
+    .. code-block:: python
+
+        await async_index(
+           "solr"
+            "/tmp/catalog.yaml",
+            server="localhost:8983",
+            batch_size=1000,
+        )
+    """
+    kwargs.setdefault("catalogue_files", catalogue_files)
+    await async_call(
+        index_system,
+        "index",
+        batch_size=batch_size,
+        verbosity=verbosity,
+        **kwargs,
+    )
+
+
+async def async_delete(
+    index_system: str,
+    batch_size: int = 2500,
+    verbosity: int = 0,
+    **kwargs: Any,
+) -> None:
+    """Delete metadata from the indexing system.
+
+    Parameters
+    ^^^^^^^^^^^
+    index_system:
+        The index server where the metadata is indexed.
+    batch_size:
+        If the index system supports batch-sizes, the size of the batches.
+    verbosity:
+        Set the verbosity of the system.
+
+    Other Parameters
+    ^^^^^^^^^^^^^^^^^
+
+    **kwargs:
+        Keyword arguments used to delete data from the index.
+
+    Examples
+    ^^^^^^^^
+
+    .. code-block:: python
+
+        await async_delete(
+            "solr"
+            server="localhost:8983",
+            latest_version="latest",
+            facets=[("file", "*.nc"), ("project", "OBS")],
+        )
+    """
+    await async_call(
+        index_system,
+        "delete",
+        batch_size=batch_size,
+        verbosity=verbosity,
+        **kwargs,
+    )
+
+
+async def async_add(
+    store: Optional[
+        Union[str, Path, Dict[str, Any], tomlkit.TOMLDocument]
+    ] = None,
+    config_file: Optional[
+        Union[Path, str, Dict[str, Any], tomlkit.TOMLDocument]
+    ] = None,
+    data_object: Optional[Union[str, List[str]]] = None,
+    data_set: Optional[Union[List[str], str]] = None,
+    data_store_prefix: str = "metadata",
+    batch_size: int = 25_000,
+    comp_level: int = 4,
+    storage_options: Optional[Dict[str, Any]] = None,
+    shadow: Optional[Union[str, List[str]]] = None,
+    catalogue_backend: CatalogueBackendType = "jsonlines",
+    latest_version: str = IndexName().latest,
+    all_versions: str = IndexName().all,
+    password: bool = False,
+    n_procs: Optional[int] = None,
+    verbosity: int = 0,
+    fail_under: int = -1,
+    **kwargs: Any,
+) -> None:
+    """Harvest metadata from storage systems and add them to an intake catalogue.
+
+    Parameters
+    ^^^^^^^^^^
+
+    store:
+        Path to the intake catalogue.
+    config_file:
+        Path to the drs-config file / loaded configuration.
+    data_objects:
+        Instead of defining datasets that are to be crawled you can crawl
+        data based on their directories. The directories must be a root dirs
+        given in the drs-config file. By default all root dirs are crawled.
+    data_object:
+        Objects (directories or catalogue files) that are processed.
+    data_set:
+        Dataset(s) that should be crawled. The datasets need to be defined
+        in the drs-config file. By default all datasets are crawled.
+        Names can contain wildcards such as ``xces-*``.
+    data_store_prefix: str
+        Absolute path or relative path to intake catalogue source
+    batch_size:
+        Batch size that is used to collect the meta data. This can affect
+        performance.
+    comp_level:
+        Compression level used to write the meta data to csv.gz
+    storage_options:
+        Set additional storage options for adding metadata to the metadata store
+    shadow:
+        'Shadow' this storage options. This is useful to hide secrets in public
+        data catalogues.
+    catalogue_backend:
+        Intake catalogue backend
+    latest_version:
+        Name of the core holding 'latest' metadata.
+    all_versions:
+        Name of the core holding 'all' metadata versions.
+    password:
+        Display a password prompt before beginning
+    n_procs:
+        Set the number of parallel processes for collecting.
+    verbosity:
+        Set the verbosity of the system.
+    fail_under:
+        Fail if less than X of the discovered files could be indexed.
+
+    Other Parameters
+    ^^^^^^^^^^^^^^^^
+
+    **kwargs:
+        Additional keyword arguments.
+
+
+    Examples
+    ^^^^^^^^
+
+     .. code-block:: python
+
+        await async_add(
+            store="my-data.yaml",
+             config_file="~/data/drs-config.toml",
+             data_set=["cmip6", "cordex"],
+        )
+
+    """
+    env = cast(os._Environ[str], os.environ.copy())
+    old_level = apply_verbosity(verbosity)
+    try:
+        os.environ["MDC_LOG_LEVEL"] = str(get_level_from_verbosity(verbosity))
+        config_file = config_file or os.environ.get(
+            "EVALUATION_SYSTEM_CONFIG_DIR"
+        )
+        if not config_file:
+            raise MetadataCrawlerException(
+                "You must give a config file/directory"
+            )
+        st = time.time()
+        passwd = ""
+        if password:  # pragma: no cover
+            passwd = Prompt.ask(
+                "[b]Enter the password", password=True
+            )  # pragma: no cover
+
+        if passwd:
+            os.environ["DRS_STORAGE_PASSWD"] = passwd
+        data_object = (
+            data_object
+            if isinstance(data_object, (NoneType, list))
+            else [str(data_object)]
+        )
+        data_set = (
+            data_set
+            if isinstance(data_set, (NoneType, list))
+            else [str(data_set)]
+        )
+        async with DataCollector(
+            config_file,
+            store,
+            IndexName(latest=latest_version, all=all_versions),
+            *_get_search(config_file, data_object, data_set),
+            batch_size=batch_size,
+            comp_level=comp_level,
+            backend=catalogue_backend,
+            data_store_prefix=data_store_prefix,
+            n_procs=n_procs,
+            storage_options=storage_options or {},
+            shadow=shadow,
+            **kwargs,
+        ) as data_col:
+            await data_col.ingest_data()
+            num_files = data_col.ingested_objects
+            files_discovered = data_col.crawled_files
+            dt = timedelta_to_str(time.time() - st)
+        logger.info("Discovered: %s files", f"{files_discovered:10,.0f}")
+        logger.info("Ingested: %s files", f"{num_files:10,.0f}")
+        logger.info("Spend: %s", dt)
+        Console.print(" " * Console.width, end="\r")
+        Console.print(
+            (
+                f"[bold]Ingested [green]{num_files:10,.0f}[/green] "
+                f"within [green]{dt}[/green][/bold]"
+            )
+        )
+
+        if (
+            files_discovered >= fail_under and num_files < fail_under
+        ) or files_discovered == 0:
+            await data_col.ingest_queue.delete()
+            raise EmptyCrawl("Could not fulfill discovery threshold!") from None
+    finally:
+        os.environ = env
+        logger.set_level(old_level)