earthcatalog 0.2.0__py3-none-any.whl

earthcatalog/stac_hooks.py

@@ -0,0 +1,754 @@
+# stac_hooks.py
+"""STAC fetch hooks for custom item generation.
+
+This module provides a flexible hook system for generating STAC items from URLs
+when the URL doesn't point directly to a STAC JSON. This is useful when:
+
+1. URLs point to raw data files (COG, NetCDF, HDF5) that need metadata extraction
+2. URLs are identifiers that need to be transformed via an external API
+3. Custom STAC generation logic is required for specific data providers
+
+Architecture:
+    Hooks are designed to be serializable for Dask distributed processing.
+    Instead of passing callable objects (which can't be pickled), hooks are
+    configured via module paths that can be imported on remote workers.
+
+Hook Types:
+    - default: Standard STAC JSON fetch from URL (current behavior)
+    - callable: Python callable (local processing only, not serializable)
+    - module: Import path to Python function (serializable for Dask)
+    - script: External executable (serializable, runs as subprocess)
+
+Usage:
+    >>> # Module-based hook (serializable, recommended for Dask)
+    >>> config = ProcessingConfig(
+    ...     input_file='urls.parquet',
+    ...     output_catalog='./catalog',
+    ...     scratch_location='./scratch',
+    ...     stac_hook='module:mypackage.hooks:url_to_stac_item',
+    ... )
+
+    >>> # Script-based hook (serializable)
+    >>> config = ProcessingConfig(
+    ...     stac_hook='script:/path/to/generate_stac.py',
+    ... )
+
+    >>> # Callable hook (local processing only)
+    >>> def my_hook(url: str) -> dict:
+    ...     return {...}
+    >>> config = ProcessingConfig(
+    ...     stac_hook=my_hook,  # Only works with LocalProcessor
+    ... )
+
+Hook Function Signature:
+    All hook functions must implement the following signature:
+
+    def hook_function(url: str, **kwargs) -> dict[str, Any] | pystac.Item | None:
+        '''Generate a STAC item from a URL.
+
+        Args:
+            url: The input URL/identifier to process.
+            **kwargs: Additional context (timeout, retry_attempts, etc.)
+
+        Returns:
+            A valid STAC item dictionary, pystac.Item object, or None if generation failed.
+            pystac.Item objects are automatically converted to dictionaries.
+        '''
+        pass
+
+Batch Hook Signature (Optional):
+    For efficiency, hooks can optionally support batch processing:
+
+    def batch_hook_function(urls: list[str], **kwargs) -> list[dict[str, Any] | pystac.Item | None]:
+        '''Generate STAC items from multiple URLs.
+
+        Args:
+            urls: List of URLs/identifiers to process.
+            **kwargs: Additional context.
+
+        Returns:
+            List of STAC items (dict or pystac.Item) or None for failures, same order as input.
+            pystac.Item objects are automatically converted to dictionaries.
+        '''
+        pass
+"""
+
+from __future__ import annotations
+
+import importlib
+import json
+import logging
+import subprocess  # nosec B404 - subprocess used for external tool integration
+import tempfile
+from abc import ABC, abstractmethod
+from collections.abc import Callable
+from typing import Any, Protocol, runtime_checkable
+
+import requests
+
+logger = logging.getLogger(__name__)
+
+
+# =============================================================================
+# Helper Functions
+# =============================================================================
+
+
+def _normalize_stac_result(result: Any) -> dict[str, Any] | None:
+    """Convert hook result to a dictionary, handling pystac.Item objects.
+
+    Args:
+        result: The result from a hook function. Can be:
+            - dict: Returned as-is
+            - pystac.Item: Converted to dict via to_dict()
+            - None: Returned as-is
+            - Other: Returns None
+
+    Returns:
+        STAC item dictionary or None.
+    """
+    if result is None:
+        return None
+    if isinstance(result, dict):
+        return result
+    # Check for pystac.Item or any object with to_dict method
+    if hasattr(result, "to_dict") and callable(result.to_dict):
+        try:
+            converted = result.to_dict()
+            if isinstance(converted, dict):
+                return converted
+            return None
+        except (ValueError, TypeError, AttributeError, RuntimeError) as e:
+            logger.warning(f"Failed to convert pystac.Item to dict: {e}")
+            return None
+    return None
+
+
+def _normalize_stac_results(results: list[Any]) -> list[dict[str, Any] | None]:
+    """Convert a list of hook results to dictionaries.
+
+    Args:
+        results: List of results from a batch hook function.
+
+    Returns:
+        List of STAC item dictionaries (or None for failures).
+    """
+    return [_normalize_stac_result(r) for r in results]
+
+
+# =============================================================================
+# Protocol Definitions
+# =============================================================================
+
+
+@runtime_checkable
+class STACHookProtocol(Protocol):
+    """Protocol for STAC fetch hooks."""
+
+    def fetch(self, url: str, **kwargs) -> dict[str, Any] | None:
+        """Fetch/generate a STAC item from a URL.
+
+        Args:
+            url: The URL or identifier to process.
+            **kwargs: Additional context like timeout, retry_attempts.
+
+        Returns:
+            STAC item dictionary or None if fetch failed.
+        """
+        ...
+
+
+@runtime_checkable
+class BatchSTACHookProtocol(Protocol):
+    """Protocol for batch STAC fetch hooks."""
+
+    def fetch_batch(self, urls: list[str], **kwargs) -> list[dict[str, Any] | None]:
+        """Fetch/generate STAC items from multiple URLs.
+
+        Args:
+            urls: List of URLs or identifiers to process.
+            **kwargs: Additional context.
+
+        Returns:
+            List of STAC items (or None for failures), in same order as input.
+        """
+        ...
+
+
+# =============================================================================
+# Hook Implementations
+# =============================================================================
+
+
+class BaseSTACHook(ABC):
+    """Base class for STAC fetch hooks."""
+
+    @abstractmethod
+    def fetch(self, url: str, **kwargs) -> dict[str, Any] | None:
+        """Fetch a single STAC item."""
+        pass
+
+    def fetch_batch(self, urls: list[str], **kwargs) -> list[dict[str, Any] | None]:
+        """Fetch multiple STAC items. Default: sequential fetch."""
+        return [self.fetch(url, **kwargs) for url in urls]
+
+    def to_config(self) -> str:
+        """Serialize hook to config string for storage/transmission."""
+        raise NotImplementedError("Subclass must implement to_config()")
+
+
+class DefaultSTACHook(BaseSTACHook):
+    """Default hook: fetch STAC JSON directly from URL.
+
+    This is the standard behavior - assumes the URL points to a valid STAC item JSON.
+    """
+
+    def fetch(
+        self,
+        url: str,
+        timeout: int = 30,
+        retry_attempts: int = 3,
+        **kwargs,
+    ) -> dict[str, Any] | None:
+        """Download and parse STAC item from URL.
+
+        Args:
+            url: URL to STAC item JSON.
+            timeout: Request timeout in seconds.
+            retry_attempts: Number of retry attempts.
+
+        Returns:
+            Parsed STAC item or None if failed.
+        """
+        import time
+
+        for attempt in range(retry_attempts):
+            try:
+                response = requests.get(url, timeout=timeout)
+                response.raise_for_status()
+                return response.json()
+            except requests.RequestException as e:
+                if attempt == retry_attempts - 1:
+                    logger.warning(f"Failed to download {url} after {retry_attempts} attempts: {e}")
+                    return None
+                time.sleep(2**attempt)  # Exponential backoff
+        return None
+
+    def to_config(self) -> str:
+        """Return config string for default hook."""
+        return "default"
+
+
+class ModuleSTACHook(BaseSTACHook):
+    """Hook that calls a Python function specified by module path.
+
+    This is the recommended approach for Dask distributed processing because
+    the module path is a string that can be serialized and the function is
+    imported on each worker.
+
+    The module path format is: 'package.module:function_name'
+
+    Example:
+        >>> hook = ModuleSTACHook('mypackage.stac_generator:url_to_item')
+        >>> item = hook.fetch('https://example.com/data.tif')
+    """
+
+    def __init__(self, module_path: str):
+        """Initialize with module path.
+
+        Args:
+            module_path: Import path in format 'package.module:function_name'
+        """
+        self.module_path = module_path
+        self._func: Callable | None = None
+        self._batch_func: Callable | None = None
+
+    def _load_function(self) -> Callable[..., Any]:
+        """Import and return the hook function."""
+        if self._func is not None:
+            return self._func
+
+        try:
+            module_name, func_name = self.module_path.rsplit(":", 1)
+            module = importlib.import_module(module_name)
+            func = getattr(module, func_name)
+            self._func = func
+
+            # Check for batch function (convention: func_name + '_batch')
+            batch_func_name = f"{func_name}_batch"
+            if hasattr(module, batch_func_name):
+                self._batch_func = getattr(module, batch_func_name)
+
+            return func
+        except (ValueError, ImportError, AttributeError) as e:
+            raise ImportError(f"Failed to import hook function from '{self.module_path}': {e}") from e
+
+    def fetch(self, url: str, **kwargs) -> dict[str, Any] | None:
+        """Fetch STAC item using the imported function.
+
+        Args:
+            url: URL or identifier to process.
+            **kwargs: Additional context passed to the function.
+
+        Returns:
+            STAC item or None. Supports both dict and pystac.Item returns.
+        """
+        func = self._load_function()
+        try:
+            result = func(url, **kwargs)
+            return _normalize_stac_result(result)
+        except (ValueError, TypeError, AttributeError, RuntimeError, OSError) as e:
+            logger.warning(f"Hook function failed for {url}: {e}")
+            return None
+
+    def fetch_batch(self, urls: list[str], **kwargs) -> list[dict[str, Any] | None]:
+        """Fetch multiple STAC items, using batch function if available.
+
+        Args:
+            urls: List of URLs to process.
+            **kwargs: Additional context.
+
+        Returns:
+            List of STAC items (or None for failures). Supports pystac.Item returns.
+        """
+        self._load_function()  # Ensure functions are loaded
+
+        if self._batch_func is not None:
+            try:
+                results = self._batch_func(urls, **kwargs)
+                if isinstance(results, list) and len(results) == len(urls):
+                    return _normalize_stac_results(results)
+                logger.warning("Batch hook returned invalid result, falling back to sequential")
+            except (ValueError, TypeError, AttributeError, RuntimeError, OSError) as e:
+                logger.warning(f"Batch hook failed: {e}, falling back to sequential")
+
+        # Fall back to sequential
+        return super().fetch_batch(urls, **kwargs)
+
+    def to_config(self) -> str:
+        """Return config string for module hook."""
+        return f"module:{self.module_path}"
+
+
+class PassthroughSTACHook(BaseSTACHook):
+    """Hook that treats input URLs as pre-fetched STAC item JSON.
+
+    This hook is designed for processing NDJSON files that already contain
+    complete STAC items, eliminating the need for HTTP fetching. The "URL"
+    in the input file is actually the full STAC item JSON as a string.
+
+    Use Cases:
+        - ITS_LIVE bulk data: NDJSON files with pre-aggregated STAC items
+        - Local STAC collections: STAC items already downloaded
+        - Performance optimization: Skip HTTP fetch for cached data
+
+    Example:
+        >>> hook = PassthroughSTACHook()
+        >>> # url is actually a JSON string
+        >>> item = hook.fetch('{"type": "Feature", "id": "abc", ...}')
+        >>> print(item['id'])
+        'abc'
+
+    Performance:
+        - Eliminates HTTP overhead (no network calls)
+        - Faster processing for local/cached STAC items
+        - Reduces load on STAC catalog servers
+    """
+
+    def fetch(self, url: str, **kwargs) -> dict[str, Any] | None:
+        """Parse input URL as STAC item JSON directly.
+
+        Args:
+            url: JSON string containing a complete STAC item.
+            **kwargs: Additional context (timeout, retry_attempts) - ignored for passthrough.
+
+        Returns:
+            Parsed STAC item dictionary or None if JSON is invalid.
+        """
+        try:
+            # Input "url" is actually a JSON string
+            item = json.loads(url)
+
+            # Validate it's a dictionary with STAC item structure
+            if not isinstance(item, dict):
+                logger.warning(f"Passthrough hook input is not a dict: {type(url)}")
+                return None
+
+            # Basic STAC item validation - should have at minimum 'type': 'Feature'
+            if item.get("type") != "Feature":
+                logger.warning(f"Passthrough hook input missing 'type': 'Feature': {url[:100]}...")
+                return None
+
+            # Should have geometry and properties
+            if "geometry" not in item:
+                logger.warning(f"Passthrough hook input missing 'geometry': {url[:100]}...")
+                return None
+
+            if "properties" not in item:
+                logger.warning(f"Passthrough hook input missing 'properties': {url[:100]}...")
+                return None
+
+            return item
+
+        except json.JSONDecodeError as e:
+            logger.warning(f"Passthrough hook failed to parse JSON: {e}")
+            return None
+        except (ValueError, TypeError, OSError) as e:
+            logger.warning(f"Passthrough hook error: {e}")
+            return None
+
+    def fetch_batch(self, urls: list[str], **kwargs) -> list[dict[str, Any] | None]:
+        """Parse multiple JSON strings as STAC items.
+
+        Args:
+            urls: List of JSON strings containing STAC items.
+
+        Returns:
+            List of STAC item dictionaries (or None for failures), in same order as input.
+        """
+        return [self.fetch(url, **kwargs) for url in urls]
+
+    def to_config(self) -> str:
+        """Return config string for passthrough hook."""
+        return "passthrough"
+
+
+class ScriptSTACHook(BaseSTACHook):
+    """Hook that calls an external script/executable.
+
+    The script receives the URL as an argument and should output valid STAC JSON
+    to stdout. For batch processing, the script can accept multiple URLs.
+
+    Script invocation:
+        Single: script_path URL
+        Batch:  script_path --batch < urls.txt  (one URL per line)
+
+    The script should exit with code 0 on success, non-zero on failure.
+    For batch mode, output should be NDJSON (one JSON object per line).
+
+    Example:
+        >>> hook = ScriptSTACHook('/path/to/generate_stac.py')
+        >>> item = hook.fetch('https://example.com/data.tif')
+    """
+
+    def __init__(self, script_path: str, interpreter: str | None = None):
+        """Initialize with script path.
+
+        Args:
+            script_path: Path to the script/executable.
+            interpreter: Optional interpreter (e.g., 'python', 'python3').
+                        If None, script is executed directly.
+        """
+        self.script_path = script_path
+        self.interpreter = interpreter
+
+    def fetch(
+        self,
+        url: str,
+        timeout: int = 60,
+        **kwargs,
+    ) -> dict[str, Any] | None:
+        """Run script to generate STAC item.
+
+        Args:
+            url: URL to pass to the script.
+            timeout: Script execution timeout in seconds.
+
+        Returns:
+            Parsed STAC item from script stdout, or None if failed.
+        """
+        cmd = self._build_command([url])
+
+        try:
+            result = subprocess.run(  # nosec B603 - command built from validated config
+                cmd,
+                capture_output=True,
+                text=True,
+                timeout=timeout,
+                check=True,
+            )
+            return json.loads(result.stdout.strip())
+        except subprocess.TimeoutExpired:
+            logger.warning(f"Script timed out for {url}")
+            return None
+        except subprocess.CalledProcessError as e:
+            logger.warning(f"Script failed for {url}: {e.stderr}")
+            return None
+        except json.JSONDecodeError as e:
+            logger.warning(f"Script output is not valid JSON for {url}: {e}")
+            return None
+
+    def fetch_batch(
+        self,
+        urls: list[str],
+        timeout: int = 300,
+        **kwargs,
+    ) -> list[dict[str, Any] | None]:
+        """Run script in batch mode for multiple URLs.
+
+        Args:
+            urls: List of URLs to process.
+            timeout: Script execution timeout in seconds.
+
+        Returns:
+            List of STAC items (or None for failures).
+        """
+        if len(urls) == 0:
+            return []
+
+        if len(urls) == 1:
+            return [self.fetch(urls[0], timeout=timeout, **kwargs)]
+
+        # Try batch mode first
+        cmd = self._build_command(["--batch"])
+
+        try:
+            # Write URLs to temp file for stdin
+            with tempfile.NamedTemporaryFile(mode="w", suffix=".txt", delete=False) as f:
+                f.write("\n".join(urls))
+                temp_path = f.name
+
+            with open(temp_path) as stdin_file:
+                result = subprocess.run(  # nosec B603 - command built from validated config
+                    cmd,
+                    stdin=stdin_file,
+                    capture_output=True,
+                    text=True,
+                    timeout=timeout,
+                )
+
+            # Clean up temp file
+            import os
+
+            os.unlink(temp_path)
+
+            if result.returncode == 0:
+                # Parse NDJSON output
+                items: list[dict[str, Any] | None] = []
+                for line in result.stdout.strip().split("\n"):
+                    if line.strip():
+                        try:
+                            items.append(json.loads(line))
+                        except json.JSONDecodeError:
+                            items.append(None)
+                    else:
+                        items.append(None)
+
+                if len(items) == len(urls):
+                    return items
+
+                logger.warning("Batch script returned wrong number of items, falling back to sequential")
+
+        except subprocess.TimeoutExpired:
+            logger.warning("Batch script timed out, falling back to sequential")
+        except (subprocess.CalledProcessError, OSError, ValueError, FileNotFoundError) as e:
+            logger.warning(f"Batch script failed: {e}, falling back to sequential")
+
+        # Fall back to sequential processing
+        return super().fetch_batch(urls, **kwargs)
+
+    def _build_command(self, args: list[str]) -> list[str]:
+        """Build command list for subprocess."""
+        if self.interpreter:
+            return [self.interpreter, self.script_path, *args]
+        return [self.script_path, *args]
+
+    def to_config(self) -> str:
+        """Return config string for script hook."""
+        if self.interpreter:
+            return f"script:{self.interpreter}:{self.script_path}"
+        return f"script:{self.script_path}"
+
+
+class CallableSTACHook(BaseSTACHook):
+    """Hook that wraps a Python callable.
+
+    WARNING: This hook is NOT serializable and will not work with Dask
+    distributed processing. Use ModuleSTACHook for distributed workloads.
+
+    This hook is useful for local processing or testing where you want to
+    pass a function directly without creating a module.
+
+    Example:
+        >>> def my_hook(url, **kwargs):
+        ...     return {'type': 'Feature', 'id': url, ...}
+        >>> hook = CallableSTACHook(my_hook)
+        >>> item = hook.fetch('https://example.com/data.tif')
+    """
+
+    def __init__(
+        self,
+        func: Callable[[str], dict[str, Any] | None],
+        batch_func: Callable[[list[str]], list[dict[str, Any] | None]] | None = None,
+    ):
+        """Initialize with callable.
+
+        Args:
+            func: Function that takes URL and returns STAC item or None.
+            batch_func: Optional batch function for efficient batch processing.
+        """
+        self.func = func
+        self.batch_func = batch_func
+
+    def fetch(self, url: str, **kwargs) -> dict[str, Any] | None:
+        """Invoke the callable with the URL.
+
+        Args:
+            url: URL to process.
+            **kwargs: Additional context.
+
+        Returns:
+            STAC item or None. Supports both dict and pystac.Item returns.
+        """
+        try:
+            result = self.func(url, **kwargs)
+            return _normalize_stac_result(result)
+        except (ValueError, TypeError, AttributeError, RuntimeError, OSError) as e:
+            logger.warning(f"Callable hook failed for {url}: {e}")
+            return None
+
+    def fetch_batch(self, urls: list[str], **kwargs) -> list[dict[str, Any] | None]:
+        """Invoke batch callable if available, otherwise sequential.
+
+        Args:
+            urls: List of URLs to process.
+            **kwargs: Additional context.
+
+        Returns:
+            List of STAC items (or None). Supports pystac.Item returns.
+        """
+        if self.batch_func is not None:
+            try:
+                results = self.batch_func(urls, **kwargs)
+                if isinstance(results, list) and len(results) == len(urls):
+                    return _normalize_stac_results(results)
+            except (ValueError, TypeError, AttributeError, RuntimeError, OSError) as e:
+                logger.warning(f"Batch callable failed: {e}, falling back to sequential")
+
+        return super().fetch_batch(urls, **kwargs)
+
+    def to_config(self) -> str:
+        """Return config string - not meaningful for callables."""
+        func_name = getattr(self.func, "__name__", "anonymous")
+        module = getattr(self.func, "__module__", "unknown")
+        return f"callable:{module}.{func_name}"
+
+
+# =============================================================================
+# Hook Factory
+# =============================================================================
+
+
+def parse_hook_config(config: str | Callable | BaseSTACHook | None) -> BaseSTACHook:
+    """Parse hook configuration and return appropriate hook instance.
+
+    Args:
+        config: Hook configuration, can be:
+            - None or "default": Use DefaultSTACHook
+            - "passthrough": Use PassthroughSTACHook (URLs are pre-fetched STAC JSON)
+            - "module:path.to.module:function": Use ModuleSTACHook
+            - "script:/path/to/script": Use ScriptSTACHook
+            - "script:python:/path/to/script.py": Use ScriptSTACHook with interpreter
+            - Callable: Use CallableSTACHook (local only)
+            - BaseSTACHook instance: Use as-is
+
+    Returns:
+        Configured hook instance.
+
+    Raises:
+        ValueError: If config format is invalid.
+    """
+    if config is None or config == "default":
+        return DefaultSTACHook()
+
+    if isinstance(config, BaseSTACHook):
+        return config
+
+    if callable(config):
+        return CallableSTACHook(config)
+
+    if not isinstance(config, str):
+        raise ValueError(f"Invalid hook config type: {type(config)}")
+
+    # Parse string config
+    if config == "passthrough":
+        return PassthroughSTACHook()
+
+    if config.startswith("module:"):
+        module_path = config[7:]  # Remove 'module:' prefix
+        return ModuleSTACHook(module_path)
+
+    if config.startswith("script:"):
+        script_config = config[7:]  # Remove 'script:' prefix
+        # Check for interpreter:path format
+        if ":" in script_config and not script_config.startswith("/"):
+            interpreter, script_path = script_config.split(":", 1)
+            return ScriptSTACHook(script_path, interpreter=interpreter)
+        return ScriptSTACHook(script_config)
+
+    raise ValueError(
+        f"Invalid hook config: {config}. Expected 'default', 'passthrough', 'module:path:func', "
+        "'script:/path', or a callable."
+    )
+
+
+def get_hook(config: str | Callable | BaseSTACHook | None = None) -> BaseSTACHook:
+    """Get a STAC hook instance from configuration.
+
+    This is the main entry point for getting hooks. It's an alias for
+    parse_hook_config with a friendlier name.
+
+    Args:
+        config: Hook configuration (see parse_hook_config for details).
+
+    Returns:
+        Configured hook instance.
+    """
+    return parse_hook_config(config)
+
+
+def serialize_hook(hook: BaseSTACHook | Callable | None) -> str:
+    """Serialize a hook to a config string for transmission.
+
+    Args:
+        hook: Hook instance or callable.
+
+    Returns:
+        Config string that can be used with parse_hook_config.
+    """
+    if hook is None:
+        return "default"
+
+    if isinstance(hook, BaseSTACHook):
+        return hook.to_config()
+
+    if callable(hook):
+        # Try to create a module path for the callable
+        func_name = getattr(hook, "__name__", None)
+        module = getattr(hook, "__module__", None)
+        if func_name and module and module != "__main__":
+            return f"module:{module}:{func_name}"
+        # Can't serialize anonymous or __main__ functions
+        raise ValueError(
+            f"Cannot serialize callable {hook}. Use ModuleSTACHook with an importable "
+            "function path for distributed processing."
+        )
+
+    raise ValueError(f"Cannot serialize hook: {hook}")
+
+
+__all__ = [
+    "STACHookProtocol",
+    "BatchSTACHookProtocol",
+    "BaseSTACHook",
+    "DefaultSTACHook",
+    "PassthroughSTACHook",
+    "ModuleSTACHook",
+    "ScriptSTACHook",
+    "CallableSTACHook",
+    "parse_hook_config",
+    "get_hook",
+    "serialize_hook",
+    "_normalize_stac_result",
+    "_normalize_stac_results",
+]