strukto-ai--mirage/python/mirage/core/s3/stat.py

# ========= Copyright 2026 @ Strukto.AI All Rights Reserved. =========
# 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.
# ========= Copyright 2026 @ Strukto.AI All Rights Reserved. =========

from mirage.accessor.s3 import S3Accessor
from mirage.cache.index import IndexCacheStore, ResourceType
from mirage.core.s3._client import _client_kwargs, _key, async_session
from mirage.core.timeutil import to_iso_z
from mirage.types import FileStat, FileType, PathSpec
from mirage.utils.errors import enoent
from mirage.utils.filetype import guess_type
from mirage.utils.key_prefix import mount_prefix_of


def _is_not_found(exc: Exception) -> bool:
    if hasattr(exc, "response"):
        code = exc.response.get("Error", {}).get("Code")
        return code in ("404", "NoSuchKey")
    return False


async def stat(accessor: S3Accessor,
               path: PathSpec,
               index: IndexCacheStore = None) -> FileStat:
    if isinstance(path, str):
        path = PathSpec(virtual=path,
                        directory=path,
                        resource_path=path.strip("/"))
    virtual = path.virtual if isinstance(path, PathSpec) else path
    original_prefix = ""
    if isinstance(path, PathSpec):
        original_prefix = mount_prefix_of(path.virtual, path.resource_path)
        path = path.virtual
    if original_prefix and path.startswith(original_prefix):
        path = path[len(original_prefix):] or "/"

    # A trailing slash ("/s3/csv/") signals the caller treats it as a
    # directory. S3 allows both an object at key "csv" AND a prefix "csv/"
    # to coexist; without this hint head_object would return the file and
    # `ls /s3/csv/` would list the file itself instead of the prefix.
    hints_directory = path.endswith("/")

    stripped = path.strip("/")

    if not stripped:
        return FileStat(name="/", type=FileType.DIRECTORY)

    # Fast path: check the index cache populated by readdir().
    # readdir() stores entries with resource_type="folder" or "file"
    # and file sizes, so stat can return instantly for known paths.
    if index is not None:
        virtual_key = (original_prefix + "/" +
                       stripped if original_prefix else "/" + stripped)
        lookup = await index.get(virtual_key)
        if lookup.entry is not None:
            entry = lookup.entry
            # S3 "folders" are synthetic common-prefixes with no object,
            # so readdir() records no time or size for them.
            if entry.resource_type == ResourceType.FOLDER:
                return FileStat(name=entry.name, type=FileType.DIRECTORY)
            # TODO: propagate ETag into IndexCacheEntry so this fast
            # path can also carry fingerprint.
            return FileStat(
                name=entry.name,
                size=entry.size,
                modified=entry.remote_time or None,
                type=guess_type(entry.name),
            )
        # If the parent directory was already listed by readdir() but
        # this path is not among its children, it does not exist.
        # This avoids expensive network calls for paths that shells
        # probe speculatively (e.g. .git, HEAD, .hg during cd).
        parent = virtual_key.rsplit("/", 1)[0] or "/"
        parent_listing = await index.list_dir(parent)
        if parent_listing.entries is not None:
            raise enoent(virtual)

    # Slow path: no index cache available, or parent directory not yet
    # listed. Hit the network.
    config = accessor.config
    key = _key(path, config)
    session = async_session(config)
    async with session.client(**_client_kwargs(config)) as client:
        # Try head_object first — works for files. Skipped when the path
        # hints a directory (trailing slash), so a coexisting object of the
        # same name does not shadow the prefix.
        if not hints_directory:
            try:
                resp = await client.head_object(Bucket=config.bucket, Key=key)
                modified = to_iso_z(resp["LastModified"])
                etag_raw = resp.get("ETag", "").strip('"')
                vid_raw = resp.get("VersionId")
                if vid_raw == "null":
                    vid_raw = None
                return FileStat(
                    name=path.rstrip("/").rsplit("/", 1)[-1],
                    size=resp["ContentLength"],
                    modified=modified,
                    type=guess_type(path),
                    fingerprint=etag_raw or None,
                    revision=vid_raw or None,
                    extra={"etag": etag_raw},
                )
            except Exception as exc:
                if not _is_not_found(exc):
                    raise

        # head_object returned 404 (or was skipped) — check if the path is a
        # valid
        # prefix (directory). S3/GCS don't have real directory objects,
        # so we probe with list_objects_v2 using MaxKeys=1.
        pfx = key.rstrip("/") + "/" if key else ""
        resp = await client.list_objects_v2(
            Bucket=config.bucket,
            Prefix=pfx,
            Delimiter="/",
            MaxKeys=1,
        )
        if resp.get("CommonPrefixes") or resp.get("Contents"):
            return FileStat(
                name=path.rstrip("/").rsplit("/", 1)[-1] or "/",
                type=FileType.DIRECTORY,
            )

        raise enoent(virtual)