Source code for multistorageclient.shortcuts

# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import builtins
import logging
import os
import re
import threading
from collections.abc import Callable, Iterator, Sequence
from typing import IO, Any
from urllib.parse import urlparse

from .client import StorageClient
from .config import RESERVED_POSIX_PROFILE_NAME, SUPPORTED_IMPLICIT_PROFILE_PROTOCOLS, PathMapping, StorageClientConfig
from .file import ObjectFile, PosixFile
from .telemetry import Telemetry
from .types import MSC_PROTOCOL, ExecutionMode, ObjectMetadata, PatternList, SignerType, SymlinkHandling, SyncResult
from .utils import build_full_path

_TELEMETRY_PROVIDER: Callable[[], Telemetry] | None = None
_TELEMETRY_PROVIDER_LOCK = threading.Lock()
_STORAGE_CLIENT_CACHE: dict[str, StorageClient] = {}
_STORAGE_CLIENT_CACHE_LOCK = threading.Lock()
_PATH_MAPPING_CACHE: dict[str | None, PathMapping] = {}
_PATH_MAPPING_CACHE_LOCK = threading.Lock()
_PROCESS_ID = os.getpid()

logger = logging.getLogger(__name__)


def _reinitialize_after_fork() -> None:
    """
    Reinitialize module state after fork to ensure fork-safety.

    This function is called automatically after a fork to:
    1. Clear the storage client cache (cached clients may have invalid state)
    2. Reinitialize locks (parent's lock state must not be inherited)
    3. Update process ID tracking

    Note: The telemetry provider is intentionally inherited by child processes,
    only its lock is reinitialized.
    """
    global _STORAGE_CLIENT_CACHE, _STORAGE_CLIENT_CACHE_LOCK  # noqa: PLW0602
    global _PATH_MAPPING_CACHE, _PATH_MAPPING_CACHE_LOCK  # noqa: PLW0602
    global _TELEMETRY_PROVIDER_LOCK
    global _PROCESS_ID

    _STORAGE_CLIENT_CACHE.clear()
    _STORAGE_CLIENT_CACHE_LOCK = threading.Lock()
    _PATH_MAPPING_CACHE.clear()
    _PATH_MAPPING_CACHE_LOCK = threading.Lock()
    # we don't need to reset telemetry provider as it is supposed to be a top-level Python function
    _TELEMETRY_PROVIDER_LOCK = threading.Lock()
    _PROCESS_ID = os.getpid()


def _check_and_reinitialize_if_forked() -> None:
    """
    Check if the current process is a fork and reinitialize if needed.

    This provides fork-safety for systems where os.register_at_fork is not available
    or as a fallback mechanism.
    """
    global _PROCESS_ID  # noqa: PLW0602

    current_pid = os.getpid()
    if current_pid != _PROCESS_ID:
        _reinitialize_after_fork()


if hasattr(os, "register_at_fork"):
    os.register_at_fork(after_in_child=_reinitialize_after_fork)




def get_telemetry_provider() -> Callable[[], Telemetry] | None:
    """
    Get the function used to create :py:class:`Telemetry` instances for storage clients created by shortcuts.

    :return: A function that provides a telemetry instance.
    """
    global _TELEMETRY_PROVIDER  # noqa: PLW0602

    return _TELEMETRY_PROVIDER





def set_telemetry_provider(telemetry_provider: Callable[[], Telemetry] | None) -> None:
    """
    Set the function used to create :py:class:`Telemetry` instances for storage clients created by shortcuts.

    :param telemetry_provider: A function that provides a telemetry instance. The function must be defined at the top level of a module to work with pickling.
    """
    global _TELEMETRY_PROVIDER
    global _TELEMETRY_PROVIDER_LOCK  # noqa: PLW0602

    with _TELEMETRY_PROVIDER_LOCK:
        _TELEMETRY_PROVIDER = telemetry_provider



def _resolve_msc_url(url: str) -> tuple[str, str]:
    """
    Resolve an MSC URL to a profile name and path.

    :param url: The MSC URL to resolve (msc://profile/path)
    :return: A tuple of (profile_name, path)
    """
    pr = urlparse(url)
    profile = pr.netloc
    # Normalize only the object path so the msc:// scheme separator and profile stay intact.
    pr = pr._replace(path=re.sub(r"/+", "/", pr.path))
    path = build_full_path(url, pr)
    path = path.removeprefix("/")
    return profile, path


def _read_cached_path_mapping() -> PathMapping:
    """
    Read path mapping once per ``MSC_CONFIG`` value for shortcut URL resolution.

    Path mapping checks happen on every non-MSC shortcut call, including POSIX paths. Caching here keeps that hot path
    from repeatedly loading and validating the full MSC config while preserving ``StorageClientConfig.read_path_mapping``
    behavior for direct callers.
    """
    cache_key = os.getenv("MSC_CONFIG", None)
    if cache_key in _PATH_MAPPING_CACHE:
        return _PATH_MAPPING_CACHE[cache_key]

    with _PATH_MAPPING_CACHE_LOCK:
        if cache_key in _PATH_MAPPING_CACHE:
            return _PATH_MAPPING_CACHE[cache_key]

        path_mapping = StorageClientConfig.read_path_mapping()
        _PATH_MAPPING_CACHE[cache_key] = path_mapping
        return path_mapping


def _resolve_non_msc_url(url: str) -> tuple[str, str]:
    """
    Resolve a non-MSC URL to a profile name and path.

    Resolution process:
    1. First check if MSC config exists
    2. If config exists, check for possible path mapping
    3. If no mapping is found, fall back to the reserved POSIX profile (``__filesystem__``) for file paths or create an implicit profile based on URL

    :param url: The non-MSC URL to resolve
    :return: A tuple of (profile_name, path)
    """
    # Check if we have a valid path mapping, if so check if there is a matching mapping
    path_mapping = _read_cached_path_mapping()
    if path_mapping:
        # Look for a matching mapping
        possible_mapping = path_mapping.find_mapping(url)
        if possible_mapping:
            return possible_mapping  # return the profile name and path

    # For file paths, use the default POSIX profile
    if url.startswith("file://"):
        pr = urlparse(url)
        return RESERVED_POSIX_PROFILE_NAME, build_full_path(url, pr)
    elif url.startswith("/"):
        url = os.path.normpath(url)
        return RESERVED_POSIX_PROFILE_NAME, url

    # For other URL protocol, create an implicit profile name
    pr = urlparse(url)
    protocol = pr.scheme.lower()

    # Translate relative paths to absolute paths
    if not protocol:
        return RESERVED_POSIX_PROFILE_NAME, os.path.realpath(url)

    # Validate the protocol is supported
    if protocol not in SUPPORTED_IMPLICIT_PROFILE_PROTOCOLS:
        supported_protocols = ", ".join([f"{p}://" for p in SUPPORTED_IMPLICIT_PROFILE_PROTOCOLS])
        raise ValueError(
            f'Unknown URL "{url}", expecting "{MSC_PROTOCOL}" or a supported protocol ({supported_protocols}) or a POSIX path'
        )

    # Build the implicit profile name using the format _protocol-bucket
    bucket = pr.netloc
    if not bucket:
        raise ValueError(f'Invalid URL "{url}", bucket name is required for {protocol}:// URLs')

    profile_name = f"_{protocol}-{bucket}"

    # Return normalized path with leading slash removed
    path = build_full_path(url, pr)
    path = path.removeprefix("/")

    return profile_name, path




def resolve_storage_client(url: str) -> tuple[StorageClient, str]:
    """
    Build and return a :py:class:`multistorageclient.StorageClient` instance based on the provided URL or path.

    This function parses the given URL or path and determines the appropriate storage profile and path.
    It supports URLs with the protocol ``msc://``, as well as POSIX paths or ``file://`` URLs for local file
    system access. If the profile has already been instantiated, it returns the cached client. Otherwise,
    it creates a new :py:class:`StorageClient` and caches it.

    The function also supports implicit profiles for non-MSC URLs. When a non-MSC URL is provided (like s3://,
    gs://, ais://, file://), MSC will infer the storage provider based on the URL protocol and create an implicit
    profile with the naming convention "_protocol-bucket" (e.g., "_s3-bucket1", "_gs-bucket1").

    Path mapping defined in the MSC configuration are also applied before creating implicit profiles.
    This allows for explicit mappings between source paths and destination MSC profiles.

    This function is fork-safe: after a fork, the cache is automatically cleared and new client instances
    are created in the child process to avoid sharing stale connections or file descriptors.

    :param url: The storage location, which can be:
                - A URL in the format ``msc://profile/path`` for object storage.
                - A local file system path (absolute POSIX path) or a ``file://`` URL.
                - A non-MSC URL with a supported protocol (s3://, gs://, ais://).

    :return: A tuple containing the :py:class:`multistorageclient.StorageClient` instance and the parsed path.

    :raises ValueError: If the URL's protocol is neither ``msc`` nor a valid local file system path
                        or a supported non-MSC protocol.
    """
    global _STORAGE_CLIENT_CACHE  # noqa: PLW0602
    global _STORAGE_CLIENT_CACHE_LOCK  # noqa: PLW0602

    _check_and_reinitialize_if_forked()

    # Normalize the path for msc:/ prefix due to pathlib.Path('msc://')
    if url.startswith("msc:/") and not url.startswith("msc://"):
        url = url.replace("msc:/", "msc://")

    # Resolve the URL to a profile name and path
    profile, path = _resolve_msc_url(url) if url.startswith(MSC_PROTOCOL) else _resolve_non_msc_url(url)

    # Check if the profile has already been instantiated
    if profile in _STORAGE_CLIENT_CACHE:
        return _STORAGE_CLIENT_CACHE[profile], path

    # Create a new StorageClient instance and cache it
    with _STORAGE_CLIENT_CACHE_LOCK:
        if profile in _STORAGE_CLIENT_CACHE:
            return _STORAGE_CLIENT_CACHE[profile], path
        else:
            client = StorageClient(
                config=StorageClientConfig.from_file(profile=profile, telemetry_provider=get_telemetry_provider())
            )
            _STORAGE_CLIENT_CACHE[profile] = client

    return client, path





def open(url: str, mode: str = "rb", **kwargs: Any) -> PosixFile | ObjectFile:
    """
    Open a file at the given URL using the specified mode.

    The function utilizes the :py:class:`multistorageclient.StorageClient` to open a file at the provided path.
    The URL is parsed, and the corresponding :py:class:`multistorageclient.StorageClient` is retrieved or built.

    :param url: The URL of the file to open. (example: ``msc://profile/prefix/dataset.tar``)
    :param mode: The file mode to open the file in.

    :return: A file-like object that allows interaction with the file.

    :raises ValueError: If the URL's protocol does not match the expected protocol ``msc``.
    """
    client, path = resolve_storage_client(url)
    return client.open(path, mode, **kwargs)





def glob(pattern: str, attribute_filter_expression: str | None = None) -> builtins.list[str]:
    """
    Return a list of files matching a pattern.

    This function supports glob-style patterns for matching multiple files within a storage system. The pattern is
    parsed, and the associated :py:class:`multistorageclient.StorageClient` is used to retrieve the
    list of matching files.

    :param pattern: The glob-style pattern to match files. (example: ``msc://profile/prefix/**/*.tar``)
    :param attribute_filter_expression: The attribute filter expression to apply to the result.

    :return: A list of file paths matching the pattern.

    :raises ValueError: If the URL's protocol does not match the expected protocol ``msc``.
    """
    client, path = resolve_storage_client(pattern)
    if not pattern.startswith(MSC_PROTOCOL) and client.profile == RESERVED_POSIX_PROFILE_NAME:
        return client.glob(path, include_url_prefix=False, attribute_filter_expression=attribute_filter_expression)
    else:
        return client.glob(path, include_url_prefix=True, attribute_filter_expression=attribute_filter_expression)





def upload_file(url: str, local_path: str | IO, attributes: dict[str, Any] | None = None) -> None:
    """
    Upload a file to the given URL from a local path.

    The function utilizes the :py:class:`multistorageclient.StorageClient` to upload a file (object) to the
    provided path. The URL is parsed, and the corresponding :py:class:`multistorageclient.StorageClient`
    is retrieved or built.

    :param url: The URL of the file. (example: ``msc://profile/prefix/dataset.tar``)
    :param local_path: The local path or file-like object to read from.

    :raises ValueError: If the URL's protocol does not match the expected protocol ``msc``.
    """
    client, path = resolve_storage_client(url)
    return client.upload_file(remote_path=path, local_path=local_path, attributes=attributes)





def download_file(url: str, local_path: str | IO) -> None:
    """
    Download a file in a given remote_path to a local path

    The function utilizes the :py:class:`multistorageclient.StorageClient` to download a file (object) at the
    provided path. The URL is parsed, and the corresponding :py:class:`multistorageclient.StorageClient`
    is retrieved or built.

    :param url: The URL of the file to download. (example: ``msc://profile/prefix/dataset.tar``)
    :param local_path: The local path or file-like object to write to.

    :raises ValueError: If the URL's protocol does not match the expected protocol ``msc``.
    """
    client, path = resolve_storage_client(url)
    return client.download_file(remote_path=path, local_path=local_path)





def is_empty(url: str) -> bool:
    """
    Checks whether the specified URL contains any objects.

    :param url: The URL to check, typically pointing to a storage location.
    :return: ``True`` if there are no objects/files under this URL, ``False`` otherwise.

    :raises ValueError: If the URL's protocol does not match the expected protocol ``msc``.
    """
    client, path = resolve_storage_client(url)
    return client.is_empty(path)





def is_file(url: str) -> bool:
    """
    Checks whether the specified url points to a file (rather than a directory or folder).

    The function utilizes the :py:class:`multistorageclient.StorageClient` to check if a file (object) exists
    at the provided path. The URL is parsed, and the corresponding :py:class:`multistorageclient.StorageClient`
    is retrieved or built.

    :param url: The URL to check the existence of a file. (example: ``msc://profile/prefix/dataset.tar``)
    """
    client, path = resolve_storage_client(url)
    return client.is_file(path=path)





def sync(
    source_url: str,
    target_url: str,
    delete_unmatched_files: bool = False,
    execution_mode: ExecutionMode = ExecutionMode.LOCAL,
    patterns: PatternList | None = None,
    preserve_source_attributes: bool = False,
    ignore_hidden: bool = True,
    dryrun: bool = False,
    dryrun_output_path: str | None = None,
    symlink_handling: SymlinkHandling = SymlinkHandling.FOLLOW,
) -> SyncResult:
    """
    Syncs files from the source storage to the target storage.

    :param source_url: The URL for the source storage.
    :param target_url: The URL for the target storage.
    :param delete_unmatched_files: Whether to delete files at the target that are not present at the source.
    :param execution_mode: The execution mode to use. Currently supports "local" and "ray".
    :param patterns: PatternList for include/exclude filtering. If None, all files are included.
    :param preserve_source_attributes: Whether to preserve source file metadata attributes during synchronization.
        When False (default), only file content is copied. When True, custom metadata attributes are also preserved.

        .. warning::
            **Performance Impact**: When enabled without a ``metadata_provider`` configured, this will make a HEAD
            request for each object to retrieve attributes, which can significantly impact performance on large-scale
            sync operations. For production use at scale, configure a ``metadata_provider`` in your storage profile.
    :param ignore_hidden: Whether to ignore hidden files and directories (starting with dot). Default is True.
    :param dryrun: If True, only enumerate and compare objects without performing any copy/delete operations.
        The returned :py:class:`SyncResult` will include a :py:class:`DryrunResult` with paths to JSONL files.
    :param dryrun_output_path: Directory to write dryrun JSONL files into. If None (default), a temporary
        directory is created automatically. Ignored when dryrun is False.
    :param symlink_handling: How to handle symbolic links during sync.
        :py:attr:`SymlinkHandling.FOLLOW` (default) dereferences symlinks and copies the target's bytes.
        :py:attr:`SymlinkHandling.SKIP` excludes symlinks from the sync.
        :py:attr:`SymlinkHandling.PRESERVE` recreates symlinks on the target via
        :py:meth:`AbstractStorageClient.make_symlink` instead of copying bytes (required for
        round-trip preservation of symlinks).
    """
    source_client, source_path = resolve_storage_client(source_url)
    target_client, target_path = resolve_storage_client(target_url)
    return target_client.sync_from(
        source_client,
        source_path,
        target_path,
        delete_unmatched_files,
        execution_mode=execution_mode,
        patterns=patterns,
        preserve_source_attributes=preserve_source_attributes,
        ignore_hidden=ignore_hidden,
        dryrun=dryrun,
        dryrun_output_path=dryrun_output_path,
        symlink_handling=symlink_handling,
    )





def sync_replicas(
    source_url: str,
    replica_indices: Sequence[int] | None = None,
    delete_unmatched_files: bool = False,
    execution_mode: ExecutionMode = ExecutionMode.LOCAL,
    patterns: PatternList | None = None,
    ignore_hidden: bool = True,
    symlink_handling: SymlinkHandling = SymlinkHandling.FOLLOW,
) -> None:
    """
    Syncs files from the source storage to all the replicas.

    :param source_url: The URL for the source storage.
    :param replica_indices: Specify the indices of the replicas to sync to. If not provided, all replicas will be synced. Index starts from 0.
    :param delete_unmatched_files: Whether to delete files at the replicas that are not present at the source.
    :param execution_mode: The execution mode to use. Currently supports "local" and "ray".
    :param patterns: PatternList for include/exclude filtering. If None, all files are included.
    :param ignore_hidden: Whether to ignore hidden files and directories (starting with dot). Default is True.
    :param symlink_handling: How to handle symbolic links during sync.
        :py:attr:`SymlinkHandling.FOLLOW` (default) dereferences symlinks and copies the target's bytes.
        :py:attr:`SymlinkHandling.SKIP` excludes symlinks from the sync.
        :py:attr:`SymlinkHandling.PRESERVE` recreates symlinks on each replica via
        :py:meth:`AbstractStorageClient.make_symlink` instead of copying bytes.
    """
    source_client, source_path = resolve_storage_client(source_url)
    source_client.sync_replicas(
        source_path,
        replica_indices=replica_indices,
        delete_unmatched_files=delete_unmatched_files,
        execution_mode=execution_mode,
        patterns=patterns,
        ignore_hidden=ignore_hidden,
        symlink_handling=symlink_handling,
    )





def list(
    url: str,
    start_after: str | None = None,
    end_at: str | None = None,
    include_directories: bool = False,
    attribute_filter_expression: str | None = None,
    show_attributes: bool = False,
    patterns: PatternList | None = None,
    symlink_handling: SymlinkHandling = SymlinkHandling.FOLLOW,
) -> Iterator[ObjectMetadata]:
    """
    Lists the contents of the specified URL prefix.

    This function retrieves the corresponding :py:class:`multistorageclient.StorageClient`
    for the given URL and returns an iterator of objects (files or directories) stored under the provided prefix.

    :param url: The prefix to list objects under.
    :param start_after: The key to start after (i.e. exclusive). An object with this key doesn't have to exist.
    :param end_at: The key to end at (i.e. inclusive). An object with this key doesn't have to exist.
    :param include_directories: Whether to include directories in the result. When True, directories are returned alongside objects.
    :param attribute_filter_expression: The attribute filter expression to apply to the result.
    :param show_attributes: Whether to return attributes in the result.
    :param patterns: PatternList for include/exclude filtering. If None, all files are included.
    :param symlink_handling: How to handle symbolic links. Only applicable for POSIX file storage.
    :return: An iterator of :py:class:`ObjectMetadata` objects representing the files (and optionally directories)
             accessible under the specified URL prefix. The returned keys will always be prefixed with msc://.
    """
    client, path = resolve_storage_client(url)
    return client.list(
        path=path,
        start_after=start_after,
        end_at=end_at,
        include_directories=include_directories,
        include_url_prefix=True,
        attribute_filter_expression=attribute_filter_expression,
        show_attributes=show_attributes,
        patterns=patterns,
        symlink_handling=symlink_handling,
    )





def list_recursive(
    url: str,
    start_after: str | None = None,
    end_at: str | None = None,
    max_workers: int = 32,
    look_ahead: int = 2,
    patterns: PatternList | None = None,
    symlink_handling: SymlinkHandling = SymlinkHandling.FOLLOW,
) -> Iterator[ObjectMetadata]:
    """
    Lists files recursively under the specified URL.

    This function retrieves the corresponding :py:class:`multistorageclient.StorageClient`
    for the given URL and returns an iterator of files under the provided path.

    :param url: The path to list objects under.
    :param start_after: The key to start after (i.e. exclusive). An object with this key doesn't have to exist.
    :param end_at: The key to end at (i.e. inclusive). An object with this key doesn't have to exist.
    :param max_workers: Maximum concurrent workers for provider-level recursive listing.
    :param look_ahead: Prefixes to buffer per worker for provider-level recursive listing.
    :param patterns: PatternList for include/exclude filtering. If None, all files are included.
    :param symlink_handling: How to handle symbolic links during listing.
    :return: An iterator of :py:class:`ObjectMetadata` objects representing files accessible under the specified URL path.
             The returned keys use the same URL-prefix behavior as :py:meth:`multistorageclient.list`.
    """
    client, path = resolve_storage_client(url)
    return client.list_recursive(
        path=path,
        start_after=start_after,
        end_at=end_at,
        max_workers=max_workers,
        look_ahead=look_ahead,
        include_url_prefix=True,
        patterns=patterns,
        symlink_handling=symlink_handling,
    )





def write(url: str, body: bytes, attributes: dict[str, Any] | None = None) -> None:
    """
    Writes an object to the storage provider at the specified path.

    :param url: The path where the object should be written.
    :param body: The content to write to the object.
    """
    client, path = resolve_storage_client(url)
    client.write(path=path, body=body, attributes=attributes)





def make_symlink(url: str, target_url: str) -> None:
    """
    Creates a symbolic link at ``url`` pointing to ``target_url``.

    Both URLs must resolve to the same storage profile.

    :param url: The URL where the symlink will be created.
    :param target_url: The URL of the target that the symlink points to.
    :raises ValueError: If the two URLs resolve to different storage profiles.
    """
    client, path = resolve_storage_client(url)
    target_client, target_path = resolve_storage_client(target_url)
    if client is not target_client:
        raise ValueError("Cannot create cross-profile symlink: url and target_url must belong to the same profile.")
    client.make_symlink(path=path, target=target_path)





def delete(url: str, recursive: bool = False) -> None:
    """
    Deletes the specified object(s) from the storage provider.

    This function retrieves the corresponding :py:class:`multistorageclient.StorageClient`
    for the given URL and deletes the object(s) at the specified path.

    :param url: The URL of the object to delete. (example: ``msc://profile/prefix/file.txt``)
    :param recursive: Whether to delete objects in the path recursively.
    """
    client, path = resolve_storage_client(url)
    client.delete(path, recursive=recursive)





def info(url: str) -> ObjectMetadata:
    """
    Retrieves metadata or information about an object stored at the specified path.

    :param url: The URL of the object to retrieve information about. (example: ``msc://profile/prefix/file.txt``)

    :return: An :py:class:`ObjectMetadata` object representing the object's metadata.
    """
    client, path = resolve_storage_client(url)
    return client.info(path)





def commit_metadata(url: str) -> None:
    """
    Commits the metadata updates for the specified storage client profile.

    :param url: The URL of the path to commit metadata for.
    """
    client, path = resolve_storage_client(url)
    client.commit_metadata(prefix=path)





def generate_presigned_url(
    url: str,
    *,
    method: str = "GET",
    signer_type: SignerType | None = None,
    signer_options: dict[str, Any] | None = None,
) -> str:
    """
    Generate a pre-signed URL granting temporary access to the object at *url*.

    :param url: The storage URL. (example: ``msc://profile/prefix/file.bin``)
    :param method: The HTTP method the URL should authorise (e.g. ``"GET"``, ``"PUT"``).
    :param signer_type: The signing backend to use. ``None`` means the provider's native signer.
    :param signer_options: Backend-specific options forwarded to the signer.
    :return: A pre-signed URL string.
    """
    client, path = resolve_storage_client(url)
    return client.generate_presigned_url(path, method=method, signer_type=signer_type, signer_options=signer_options)