metadata-crawler 2510.1.0__py3-none-any.whl

metadata_crawler/ingester/solr.py

@@ -0,0 +1,282 @@
+"""Collection of aync data ingest classes."""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+import os
+import time
+from concurrent.futures import ThreadPoolExecutor
+from types import TracebackType
+from typing import Annotated, Any, Dict, List, Optional, Tuple, Type, cast
+
+import aiohttp
+import orjson
+
+from ..api.cli import cli_function, cli_parameter
+from ..api.index import BaseIndex
+from ..api.metadata_stores import IndexName
+from ..logger import logger
+
+
+class SolrIndex(BaseIndex):
+    """Ingest metadata into an apache solr server."""
+
+    senteniel: Optional[bytes] = None
+
+    def __post_init__(self) -> None:
+        self.timeout = aiohttp.ClientTimeout(
+            connect=10, sock_connect=10, sock_read=180, total=None
+        )
+        self.semaphore = asyncio.Event()
+        self.max_http_workers: int = 0
+        queue_max: int = 128
+        encode_workers: int = 4
+        self._uri: str = ""
+        self.cpu_pool = ThreadPoolExecutor(max_workers=encode_workers)
+        self.producer_queue: asyncio.Queue[Tuple[str, Optional[bytes]]] = (
+            asyncio.Queue(maxsize=queue_max)
+        )
+        self.connector = aiohttp.TCPConnector(
+            ttl_dns_cache=300,
+            use_dns_cache=True,
+            enable_cleanup_closed=True,
+        )
+
+    async def solr_url(self, server: str, core: str) -> str:
+        """Construct the solr url from a given solr core."""
+        if not self._uri:
+            scheme, _, server = server.rpartition("://")
+            scheme = scheme or "http"
+            solr_server, _, solr_port = server.partition(":")
+            solr_server = solr_server or "localhost"
+            uri = f"{scheme}://{solr_server}"
+            uri = f"{uri}:{solr_port}" if solr_port else uri
+            self._uri = f"{uri}/solr"
+        return f"{self._uri}/{core}/update/json?commit=true"
+
+    @cli_function(
+        help="Remove metadata from the apache solr server.",
+    )
+    async def delete(
+        self,
+        *,
+        server: Annotated[
+            Optional[str],
+            cli_parameter(
+                "-sv",
+                "--server",
+                help="The <host>:<port> to the solr server",
+                type=str,
+            ),
+        ] = None,
+        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,
+        latest_version: Annotated[
+            str,
+            cli_parameter(
+                "--latest-version",
+                type=str,
+                help="Name of the core holding 'latest' metadata.",
+            ),
+        ] = IndexName().latest,
+        all_versions: Annotated[
+            str,
+            cli_parameter(
+                "--all-versions",
+                type=str,
+                help="Name of the core holding 'all' metadata versions.",
+            ),
+        ] = IndexName().all,
+    ) -> None:
+        """Remove metadata from the apache solr server."""
+        query = []
+        for key, value in facets or []:
+            if key.lower() == "file":
+                if value[0] in (os.sep, "/"):
+                    value = f"\\{value}"
+                value = value.replace(":", "\\:")
+            else:
+                value = value.lower()
+            query.append(f"{key.lower()}:{value}")
+        query_str = " AND ".join(query)
+        server = server or ""
+        async with aiohttp.ClientSession(timeout=self.timeout) as session:
+            logger.debug("Deleting entries matching %s", query_str)
+            for core in (all_versions, latest_version):
+                url = await self.solr_url(server, core)
+                async with session.post(
+                    url, json={"delete": {"query": query_str}}
+                ) as resp:
+                    level = (
+                        logging.WARNING
+                        if resp.status not in (200, 201)
+                        else logging.DEBUG
+                    )
+                    logger.log(level, await resp.text())
+
+    def _convert(self, metadata: Dict[str, Any]) -> Dict[str, Any]:
+        for k, v in metadata.items():
+            match self.index_schema[k].type:
+                case "bbox":
+                    metadata[k] = f"ENVELOPE({v[0]}, {v[1]}, {v[3]}, {v[2]})"
+                case "daterange":
+                    metadata[k] = f"[{v[0].isoformat()} TO {v[-1].isoformat()}]"
+
+        return metadata
+
+    def _encode_payload(self, chunk: List[Dict[str, Any]]) -> bytes:
+        """CPU-bound: convert docs and JSON-encode off the event loop."""
+        return orjson.dumps([self._convert(x) for x in chunk])
+
+    async def _post_chunk(
+        self,
+        session: aiohttp.ClientSession,
+        url: str,
+        body: bytes,
+    ) -> None:
+        """POST one batch with minimal overhead and simple retries."""
+        status = 500
+        t0 = time.perf_counter()
+        try:
+            async with session.post(
+                url, data=body, headers={"Content-Type": "application/json"}
+            ) as resp:
+                status = resp.status
+                await resp.read()
+
+        except Exception as error:
+            logger.log(
+                logging.WARNING,
+                error,
+                exc_info=logger.level < logging.INFO,
+            )
+            return
+        logger.debug(
+            "POST %s -> %i (index time: %.3f)",
+            url,
+            status,
+            time.perf_counter() - t0,
+        )
+
+    async def consumer(self, session: aiohttp.ClientSession) -> None:
+        """Consume the metadata read by the porducers."""
+        while True:
+            update_url, body = await self.producer_queue.get()
+            if body is self.senteniel:
+                self.producer_queue.task_done()
+                break
+            try:
+                await self._post_chunk(session, update_url, cast(bytes, body))
+            finally:
+                self.producer_queue.task_done()
+
+    async def _index_core(
+        self,
+        session: aiohttp.ClientSession,
+        server: str,
+        core: str,
+        suffix: str,
+        http_workers: int = 8,
+    ) -> None:
+        """Zero-copy-ish, backpressured, bounded-concurrency indexer.
+
+        - No per-batch commit.
+        - Bounded queue so tasks don't pile up.
+        - Constant number of worker tasks (not O(batches)).
+        """
+        base_url = await self.solr_url(server, core + suffix)
+        update_url = base_url.split("?", 1)[0]  # guard
+        loop = asyncio.get_running_loop()
+        async for batch in self.get_metadata(core):
+            body = await loop.run_in_executor(
+                self.cpu_pool, self._encode_payload, batch
+            )
+            await self.producer_queue.put((update_url, body))
+        commit_url = f"{update_url}?commit=true"
+        async with session.post(
+            commit_url,
+            data=b"[]",
+            headers={"Content-Type": "application/json"},
+        ) as resp:
+            if resp.status >= 400:
+                text = await resp.text()
+                logger.warning(
+                    "COMMIT %s -> %i: %s", commit_url, resp.status, text
+                )
+
+    async def __aexit__(
+        self,
+        exc_type: Optional[Type[BaseException]],
+        exc_val: Optional[BaseException],
+        exc_tb: Optional[TracebackType],
+    ) -> None:
+
+        try:
+            self.producer_queue.shutdown()
+        except AttributeError:  # pragma: no cover
+            pass  # prgama: no cover
+        self.cpu_pool.shutdown()
+
+    @cli_function(
+        help="Add metadata to the apache solr metadata server.",
+    )
+    async def index(
+        self,
+        *,
+        server: Annotated[
+            Optional[str],
+            cli_parameter(
+                "-sv",
+                "--server",
+                help="The <host>:<port> to the solr server",
+                type=str,
+            ),
+        ] = None,
+        index_suffix: Annotated[
+            Optional[str],
+            cli_parameter(
+                "--index-suffix",
+                help="Suffix for the latest and all version collections.",
+                type=str,
+            ),
+        ] = None,
+        http_workers: Annotated[
+            int,
+            cli_parameter(
+                "--http-workers", help="Number of ingestion threads.", type=int
+            ),
+        ] = 8,
+    ) -> None:
+        """Add metadata to the apache solr metadata server."""
+        async with aiohttp.ClientSession(
+            timeout=self.timeout, connector=self.connector, raise_for_status=True
+        ) as session:
+            consumers = [
+                asyncio.create_task(self.consumer(session))
+                for _ in range(http_workers)
+            ]
+            async with asyncio.TaskGroup() as tg:
+                for core in self.index_names:
+                    tg.create_task(
+                        self._index_core(
+                            session,
+                            server or "",
+                            core,
+                            suffix=index_suffix or "",
+                            http_workers=http_workers,
+                        )
+                    )
+            for _ in range(http_workers):
+                await self.producer_queue.put(("", self.senteniel))
+            await self.producer_queue.join()
+            await asyncio.gather(*consumers)

metadata_crawler/logger.py

@@ -0,0 +1,153 @@
+"""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 = "metadata-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": "CRITICAL",
+            "handlers": ["default"],
+        },
+        "formatters": {
+            "standard": {
+                "format": "%(asctime)s %(levelname)s: %(name)s - %(message)s",
+            },
+        },
+        "handlers": {
+            "default": {
+                "class": "logging.StreamHandler",
+                "formatter": "standard",
+                "level": "CRITICAL",
+            },
+        },
+    }
+)
+
+
+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,
+        suffix: Optional[str] = None,
+    ) -> None:
+        """Instantiate this logger only once and for all."""
+        self.level = level or int(
+            cast(str, os.getenv("MDC_LOG_LEVEL", str(logging.CRITICAL)))
+        )
+        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(self.level)
+        super().__init__(name, self.level)
+
+        self.propagate = False
+        self.handlers = [self._logger_stream_handle]
+        (
+            self.add_file_handle(suffix=suffix)
+            if os.getenv("MDC_LOG_INIT", "0") == "1"
+            else None
+        )
+
+    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)
+        self.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)
+
+    def add_file_handle(
+        self,
+        suffix: Optional[str] = None,
+        level: int = logging.CRITICAL,
+    ) -> None:
+        """Add a file log handle to the logger."""
+        suffix = suffix or os.getenv("MDC_LOG_SUFFIX", "")
+        base_name = f"{THIS_NAME}-{suffix}" if suffix else THIS_NAME
+        log_dir = Path(os.getenv("MDC_LOG_DIR", 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(self.file_format)
+        logger_file_handle.setLevel(self.level)
+        self.addHandler(logger_file_handle)
+
+
+logger = Logger()
+
+
+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: Optional[int] = None, suffix: Optional[str] = None
+) -> int:
+    """Set the logging level of the handlers to a certain level."""
+    level = logger.level if level is None else level
+    old_level = logger.level
+    level = get_level_from_verbosity(level)
+    logger.set_level(level)
+    logger.add_file_handle(suffix, level)
+
+    return old_level