metadata-crawler 2510.1.0__py3-none-any.whl

metadata_crawler/data_collector.py

@@ -0,0 +1,278 @@
+"""Gather metadata and for adding them to a temporary metadata store."""
+
+from __future__ import annotations
+
+import asyncio
+import os
+from multiprocessing import Event, Value
+from pathlib import Path
+from types import TracebackType
+from typing import (
+    Any,
+    AsyncIterator,
+    Callable,
+    Coroutine,
+    Dict,
+    Iterator,
+    Optional,
+    Tuple,
+    Type,
+    Union,
+    cast,
+)
+
+import tomlkit
+
+from .api.config import CrawlerSettings, DRSConfig
+from .api.metadata_stores import CatalogueWriter, IndexName
+from .api.storage_backend import PathTemplate
+from .logger import logger
+from .utils import (
+    Counter,
+    MetadataCrawlerException,
+    create_async_iterator,
+    print_performance,
+)
+
+ScanItem = Tuple[str, str, bool, bool]
+
+
+class DataCollector:
+    """Collect file objects from a given directory object and search for files.
+
+    Parameters
+    ----------
+    config_file:
+        Path to the drs-config file / loaded configuration.
+    *search_objects:
+        Paths of the search directories. e.g. `root_path` attr in drs_config
+    uri: str
+        the uir of the metadata store.
+    password: str
+        Password for the ingestion
+    batch_size: int
+        Batch size for the ingestion
+    """
+
+    def __init__(
+        self,
+        config_file: Union[Path, str, Dict[str, Any], tomlkit.TOMLDocument],
+        metadata_store: Optional[
+            Union[Path, str, Dict[str, Any], tomlkit.TOMLDocument]
+        ],
+        index_name: IndexName,
+        *search_objects: CrawlerSettings,
+        **kwargs: Any,
+    ):
+        self._search_objects = search_objects
+        if not search_objects:
+            raise MetadataCrawlerException("You have to give search directories")
+        self._num_files: Counter = Value("i", 0)
+        self.index_name = index_name
+        self.config = DRSConfig.load(config_file)
+        kwargs.setdefault("scan_concurrency", os.getenv("SCAN_CONCURRENCY", "64"))
+        self._scan_concurrency: int = int(kwargs.pop("scan_concurrency", 64))
+        self._scan_queue: asyncio.Queue[Optional[ScanItem]] = asyncio.Queue(
+            maxsize=int(kwargs.pop("scan_queue_size", 10_000))
+        )
+        self._print_status = Event()
+        self.ingest_queue = CatalogueWriter(
+            str(metadata_store or "metadata.yaml"),
+            index_name=index_name,
+            config=config_file,
+            **kwargs,
+        )
+        self.ingest_queue.run_consumer()
+        self._max_files = int(cast(str, os.getenv("MDC_MAX_FILES", "-1")))
+
+    @property
+    def crawled_files(
+        self,
+    ) -> int:
+        """Get the total number of crawled files."""
+        return self._num_files.value
+
+    @property
+    def ingested_objects(self) -> int:
+        """Get the number of ingested objects."""
+        return self.ingest_queue.ingested_objects
+
+    @property
+    def search_objects(self) -> Iterator[tuple[str, str]]:
+        """Async iterator for the search directories."""
+        for cfg in self._search_objects:
+            yield cfg.name, str(cfg.search_path)
+
+    async def __aenter__(self) -> "DataCollector":
+        return self
+
+    async def __aexit__(
+        self,
+        exc_type: Optional[Type[BaseException]],
+        exc: Optional[BaseException],
+        tb: TracebackType,
+    ) -> None:
+        self._print_status.clear()
+        self.ingest_queue.join_all_tasks()
+        await self.ingest_queue.close()
+
+        async def _safe_close(b: PathTemplate) -> None:
+            try:
+                await asyncio.wait_for(b.close(), timeout=3)
+            except Exception:
+                pass
+
+        await asyncio.gather(
+            *[_safe_close(ds.backend) for ds in self.config.datasets.values()],
+            return_exceptions=True,
+        )
+
+    def _test_env(self) -> bool:
+        return (
+            True
+            if self._max_files > 0 and self._max_files < self.crawled_files
+            else False
+        )
+
+    async def _ingest_dir(
+        self,
+        drs_type: str,
+        search_dir: str,
+        iterable: bool = True,
+        is_versioned: bool = True,
+    ) -> None:
+        if iterable:
+            try:
+                sub_dirs = self.config.datasets[drs_type].backend.iterdir(
+                    search_dir
+                )
+            except Exception as error:
+                logger.error(error)
+                return
+        else:
+            sub_dirs = cast(
+                AsyncIterator[str], create_async_iterator([search_dir])
+            )
+        rank = 0
+        async for _dir in sub_dirs:
+            async for _inp in self.config.datasets[drs_type].backend.rglob(
+                _dir, self.config.datasets[drs_type].glob_pattern
+            ):
+                if self._test_env():
+                    return
+                await self.ingest_queue.put(
+                    _inp, drs_type, name=self.index_name.all
+                )
+                if rank == 0 or is_versioned is False:
+                    await self.ingest_queue.put(
+                        _inp, drs_type, name=self.index_name.latest
+                    )
+                self._num_files.value += 1
+            rank += 1
+        return None
+
+    async def _scan_worker(self) -> None:
+        """Drain _scan_queue and run _ingest_dir concurrently (bounded pool)."""
+        while True:
+            item = await self._scan_queue.get()  # blocks
+            if item is None:  # sentinel -> exit
+                # do not task_done() for sentinel
+                break
+            drs_type, path, iterable, is_versioned = item
+            try:
+                await self._ingest_dir(
+                    drs_type, path, iterable=iterable, is_versioned=is_versioned
+                )
+            except Exception as error:
+                logger.error(error)
+            finally:
+                self._scan_queue.task_done()
+
+    async def _iter_content(
+        self,
+        drs_type: str,
+        inp_dir: str,
+        pos: int = 0,
+        is_versioned: bool = True,
+    ) -> None:
+        """Walk recursively until files or the version level is reached."""
+        store = self.config.datasets[drs_type].backend
+        if self._test_env():
+            return
+        try:
+            is_file, iterable, suffix = await asyncio.gather(
+                store.is_file(inp_dir),
+                store.is_dir(inp_dir),
+                store.suffix(inp_dir),
+            )
+        except Exception as error:
+            logger.error("Error checking file %s", error)
+            return
+
+        iterable = False if suffix == ".zarr" else iterable
+        op: Optional[Callable[..., Coroutine[Any, Any, None]]] = None
+        if is_file and suffix in self.config.suffixes:
+            op = self._ingest_dir
+        elif pos <= 0 or suffix == ".zarr":
+            op = self._ingest_dir
+
+        if op is not None:
+            # enqueue the heavy scan; workers will run _ingest_dir concurrently
+            await self._scan_queue.put(
+                (drs_type, inp_dir, iterable, is_versioned)
+            )
+            return
+
+        # otherwise, recurse sequentially (cheap) — no task per directory
+        try:
+            async for sub in store.iterdir(inp_dir):
+                await self._iter_content(
+                    drs_type, sub, pos - 1, is_versioned=is_versioned
+                )
+        except Exception as error:
+            logger.error(error)
+
+    async def ingest_data(self) -> None:
+        """Produce scan tasks and process them with a bounded worker pool."""
+        self._print_status.set()
+        self._num_files.value = 0
+        print_performance(
+            self._print_status,
+            self._num_files,
+            self.ingest_queue.queue,
+            self.ingest_queue.num_objects,
+        )
+
+        async with asyncio.TaskGroup() as tg:
+            # start scan workers
+            for _ in range(self._scan_concurrency):
+                tg.create_task(self._scan_worker())
+
+            # produce scan items by walking roots sequentially
+            for drs_type, path in self.search_objects:  # <- property is sync
+                pos, is_versioned = self.config.max_directory_tree_level(
+                    path, drs_type=drs_type
+                )
+                if pos < 0:
+                    logger.warning(
+                        "Can't define latest version of versioned dataset."
+                        " This might lead to unexpected results. Try adjusting"
+                        " your search path."
+                    )
+
+                await self._iter_content(
+                    drs_type, path, pos, is_versioned=is_versioned
+                )
+
+            # wait until all queued scan items are processed
+            await self._scan_queue.join()
+
+            # stop workers (one sentinel per worker)
+            for _ in range(self._scan_concurrency):
+                await self._scan_queue.put(None)
+
+        logger.info(
+            "%i ingestion tasks have been completed", len(self._search_objects)
+        )
+        self.ingest_queue.join_all_tasks()
+        self._print_status.clear()

metadata_crawler/ingester/__init__.py

@@ -0,0 +1 @@
+"""Module for ingesting data to the metadata index."""

metadata_crawler/ingester/mongo.py

@@ -0,0 +1,206 @@
+"""Collection of aync data ingest classes."""
+
+from __future__ import annotations
+
+import asyncio
+import re
+from functools import cached_property
+from typing import Annotated, Any, Dict, List, Optional, Tuple
+from urllib.parse import ParseResult, parse_qs, urlencode, urlparse, urlunparse
+
+from motor.motor_asyncio import (
+    AsyncIOMotorClient,
+    AsyncIOMotorCollection,
+    AsyncIOMotorDatabase,
+)
+from pymongo import DeleteMany, UpdateOne
+
+from ..api.cli import cli_function, cli_parameter
+from ..api.index import BaseIndex
+from ..logger import logger
+
+
+class MongoIndex(BaseIndex):
+    """Ingest metadata into a mongoDB server."""
+
+    def __post_init__(self) -> None:
+        self._raw_uri = ""
+        self._url = ""
+        self._client: Optional[AsyncIOMotorClient[Any]] = None
+
+    @property
+    def uri(self) -> str:
+        """Create the connection uri for the mongoDB."""
+        if self._url:
+            return self._url
+        parsed_url = urlparse(self._raw_uri)
+        query = parse_qs(parsed_url.query)
+        if "timeout" not in parsed_url.query.lower():
+            query["timeoutMS"] = ["5000"]
+        new_query = urlencode(query, doseq=True)
+        self._url = urlunparse(
+            ParseResult(
+                parsed_url.scheme or "mongodb",
+                parsed_url.netloc,
+                parsed_url.path.rstrip("/"),
+                parsed_url.params,
+                new_query,
+                parsed_url.fragment,
+            )
+        )
+        return self._url
+
+    @cached_property
+    def unique_index(self) -> str:
+        """Get the index."""
+        for name, schema in self.index_schema.items():
+            if schema.unique:
+                return name
+        raise ValueError("The schema doesn't define a unique value.")
+
+    @property
+    def client(self) -> AsyncIOMotorClient[Any]:
+        """Get the mongoDB client."""
+        if self._client is None:
+            logger.debug("Creating async mongoDB client: %s", self.uri)
+            self._client = AsyncIOMotorClient(self.uri)
+        return self._client
+
+    async def _bulk_upsert(
+        self, chunk: List[Dict[str, Any]], collection: AsyncIOMotorCollection[Any]
+    ) -> None:
+        ops = [
+            UpdateOne(
+                {self.unique_index: m[self.unique_index]},
+                {"$set": m},
+                upsert=True,
+            )
+            for m in chunk
+        ]
+        await collection.bulk_write(ops, ordered=False)
+
+    async def _index_collection(
+        self, db: AsyncIOMotorDatabase[Any], collection: str, suffix: str = ""
+    ) -> None:
+        """Index a collection."""
+        col = collection + suffix
+        await db[col].create_index(self.unique_index, unique=True)
+        async for chunk in self.get_metadata(collection):
+            await self._bulk_upsert(chunk, db[col])
+
+    async def _prep_db_connection(
+        self, database: str, url: str
+    ) -> AsyncIOMotorDatabase[Any]:
+
+        await self.close()
+        self._raw_uri = url or ""
+        return self.client[database]
+
+    @cli_function(
+        help="Add metadata to the mongoDB metadata server.",
+    )
+    async def index(
+        self,
+        *,
+        url: Annotated[
+            Optional[str],
+            cli_parameter(
+                "--url",
+                help="The <host>:<port> to the mngoDB server",
+                type=str,
+            ),
+        ] = None,
+        database: Annotated[
+            str,
+            cli_parameter(
+                "--database",
+                "--db",
+                help="The DB name holding the metadata.",
+                type=str,
+                default="metadata",
+            ),
+        ] = "metadata",
+        index_suffix: Annotated[
+            Optional[str],
+            cli_parameter(
+                "--index-suffix",
+                help="Suffix for the latest and all version collections.",
+                type=str,
+            ),
+        ] = None,
+    ) -> None:
+        """Add metadata to the mongoDB metadata server."""
+        db = await self._prep_db_connection(database, url or "")
+        async with asyncio.TaskGroup() as tg:
+            for collection in self.index_names:
+                tg.create_task(
+                    self._index_collection(
+                        db, collection, suffix=index_suffix or ""
+                    )
+                )
+
+    async def close(self) -> None:
+        """Close the mongoDB connection."""
+        self._client.close() if self._client is not None else None
+        self._url = ""
+        self._raw_uri = ""
+
+    @cli_function(
+        help="Remove metadata from the mongoDB metadata server.",
+    )
+    async def delete(
+        self,
+        *,
+        url: Annotated[
+            Optional[str],
+            cli_parameter(
+                "--url",
+                help="The <host>:<port> to the mngoDB server",
+                type=str,
+            ),
+        ] = None,
+        database: Annotated[
+            str,
+            cli_parameter(
+                "--database",
+                "--db",
+                help="The DB name holding the metadata.",
+                type=str,
+                default="metadata",
+            ),
+        ] = "metadata",
+        facets: Annotated[
+            Optional[List[Tuple[str, str]]],
+            cli_parameter(
+                "-f",
+                "--facets",
+                type=str,
+                nargs=2,
+                action="append",
+                help="Search facets matching the delete query.",
+            ),
+        ] = None,
+    ) -> None:
+        """Remove metadata from the mongoDB metadata server."""
+        db = await self._prep_db_connection(database, url or "")
+        if not facets:
+            logger.info("Nothing to delete")
+            return
+
+        def glob_to_regex(glob: str) -> str:
+            """Turn a shell‐style glob into a anchored mongo regex."""
+            # escape everything, then un-escape our wildcards
+            esc = re.escape(glob)
+            esc = esc.replace(r"\*", ".*").replace(r"\?", ".")
+            return f"^{esc}$"
+
+        ops: List[DeleteMany] = []
+        for field, val in facets:
+            if "*" in val or "?" in val:
+                pattern = glob_to_regex(val)
+                ops.append(DeleteMany({field: {"$regex": pattern}}))
+            else:
+                ops.append(DeleteMany({field: val}))
+        logger.debug("Deleting entries matching %s", ops)
+        for collection in await db.list_collection_names():
+            await db[collection].bulk_write(ops, ordered=False)