notebookutils-py 0.1.0__py3-none-any.whl

notebookutils/fs/_mount.py

@@ -0,0 +1,268 @@
+"""Process-safe POSIX mounting layer.
+
+Concurrency safety across processes (e.g. co-scheduled Spark executors on a
+shared host) is enforced with a POSIX ``fcntl`` file-lock protocol:
+
+1. Acquire an exclusive ``flock`` on a lock file derived from a SHA-256 hash
+   of the resolved local mount path (hashing avoids ``/tmp`` name collisions
+   between distinct paths).
+2. Inside the critical section, check ``os.path.ismount``: if already
+   mounted, step over the FUSE launch (idempotency).
+3. Otherwise translate the Global Registry payload + ``extraConfigs`` into an
+   explicit CLI command and execute the FUSE daemon.
+
+An on-disk JSON mount registry (guarded by its own lock) backs ``mounts()``,
+``getMountPath()`` and ``unmount()`` across processes on the same host.
+"""
+
+from __future__ import annotations
+
+import fcntl
+import hashlib
+import json
+import os
+import subprocess
+import tempfile
+from contextlib import contextmanager
+from typing import Any, Dict, List, Optional
+
+from notebookutils.fs import _registry
+from notebookutils.fs._client import get_canonical_protocol
+from notebookutils.fs._fileinfo import MountPointInfo
+
+DEFAULT_MOUNT_TIMEOUT = 30  # seconds, per Fabric documentation
+
+
+def _state_dir() -> str:
+    base = os.environ.get(
+        "NOTEBOOKUTILS_STATE_DIR",
+        os.path.join(tempfile.gettempdir(), f"notebookutils-{os.getuid()}"),
+    )
+    os.makedirs(base, mode=0o700, exist_ok=True)
+    return base
+
+
+def _mount_base_dir() -> str:
+    base = os.environ.get(
+        "NOTEBOOKUTILS_MOUNT_BASE", os.path.join(_state_dir(), "mounts")
+    )
+    os.makedirs(base, exist_ok=True)
+    return base
+
+
+def _registry_path() -> str:
+    return os.path.join(_state_dir(), "mounts.json")
+
+
+def _lock_path_for(key: str) -> str:
+    digest = hashlib.sha256(key.encode("utf-8")).hexdigest()[:32]
+    return os.path.join(_state_dir(), f"lock_{digest}.lock")
+
+
+@contextmanager
+def _flock(lock_path: str):
+    with open(lock_path, "w") as lock_file:
+        try:
+            fcntl.flock(lock_file, fcntl.LOCK_EX)
+            yield
+        finally:
+            fcntl.flock(lock_file, fcntl.LOCK_UN)
+
+
+# ---------------------------------------------------------------------------
+# Mount registry persistence
+# ---------------------------------------------------------------------------
+
+def _read_registry() -> Dict[str, Dict[str, Any]]:
+    path = _registry_path()
+    if not os.path.exists(path):
+        return {}
+    try:
+        with open(path) as f:
+            return json.load(f)
+    except (json.JSONDecodeError, OSError):
+        return {}
+
+
+def _write_registry(data: Dict[str, Dict[str, Any]]) -> None:
+    path = _registry_path()
+    tmp = path + ".tmp"
+    with open(tmp, "w") as f:
+        json.dump(data, f, indent=2)
+    os.replace(tmp, path)
+
+
+@contextmanager
+def _registry_transaction():
+    with _flock(_lock_path_for("__mount_registry__")):
+        data = _read_registry()
+        yield data
+        _write_registry(data)
+
+
+# ---------------------------------------------------------------------------
+# Mount point resolution
+# ---------------------------------------------------------------------------
+
+def _normalize_mount_point(mount_point: str) -> str:
+    mount_point = mount_point.rstrip("/") or "/"
+    if not mount_point.startswith("/"):
+        mount_point = "/" + mount_point
+    return mount_point
+
+
+def _resolve_local_path(mount_point: str) -> str:
+    """Resolve the actual local directory backing a logical mount point.
+
+    Prefer mounting directly at ``mount_point``; if it cannot be created
+    (e.g. unprivileged ``/test``), fall back to a per-user base directory.
+    """
+    try:
+        os.makedirs(mount_point, exist_ok=True)
+        if os.access(mount_point, os.W_OK):
+            return mount_point
+    except OSError:
+        pass
+    fallback = os.path.join(_mount_base_dir(), mount_point.lstrip("/"))
+    os.makedirs(fallback, exist_ok=True)
+    return fallback
+
+
+# ---------------------------------------------------------------------------
+# Public API
+# ---------------------------------------------------------------------------
+
+def mount(
+    source: str,
+    mount_point: str,
+    extra_configs: Optional[Dict[str, Any]] = None,
+) -> bool:
+    """Mount a cloud storage location at a local POSIX path.
+
+    Supported ``extra_configs`` keys (Fabric-compatible):
+
+    - ``accountKey``: Azure storage account key (abfs/abfss only).
+    - ``sasToken``: Azure SAS token (abfs/abfss only).
+    - ``fileCacheTimeout``: local file cache timeout in seconds.
+    - ``timeout``: mount operation timeout in seconds (default 30).
+    """
+    extra_configs = dict(extra_configs or {})
+    protocol = get_canonical_protocol(source)
+    mount_point = _normalize_mount_point(mount_point)
+    local_path = _resolve_local_path(mount_point)
+
+    with _flock(_lock_path_for(local_path)):
+        # Critical section: evaluate filesystem status inside the safe window
+        if os.path.ismount(local_path):
+            _record_mount(mount_point, source, local_path, extra_configs)
+            return True  # already handled by a neighboring process
+
+        config = _registry.get_fsspec_options(protocol)
+        mount_opts = _registry.get_mount_options(protocol)
+
+        if protocol == "abfs":
+            from notebookutils.fs._mount_azure import exec_blobfuse2
+
+            exec_blobfuse2(source, local_path, config, mount_opts, extra_configs)
+        elif protocol == "s3":
+            from notebookutils.fs._mount_s3 import exec_mount_s3
+
+            exec_mount_s3(source, local_path, config, mount_opts, extra_configs)
+        elif protocol == "gcs":
+            from notebookutils.fs._mount_gcs import exec_gcsfuse
+
+            exec_gcsfuse(source, local_path, config, mount_opts, extra_configs)
+        else:
+            raise ValueError(
+                f"Unsupported mount target protocol scheme: {protocol}"
+            )
+
+        _record_mount(mount_point, source, local_path, extra_configs)
+        return True
+
+
+def unmount(mount_point: str, extra_configs: Optional[Dict[str, Any]] = None) -> bool:
+    """Unmount and remove a mount point previously created with :func:`mount`."""
+    mount_point = _normalize_mount_point(mount_point)
+    with _registry_transaction() as data:
+        entry = data.get(mount_point)
+    local_path = entry["localPath"] if entry else _resolve_local_path(mount_point)
+
+    with _flock(_lock_path_for(local_path)):
+        if os.path.ismount(local_path):
+            _exec_unmount(local_path)
+
+    with _registry_transaction() as data:
+        data.pop(mount_point, None)
+    return True
+
+
+def mounts() -> List[MountPointInfo]:
+    """List all existing mount points registered on this host."""
+    result: List[MountPointInfo] = []
+    with _registry_transaction() as data:
+        stale = [
+            mp for mp, e in data.items() if not os.path.ismount(e.get("localPath", ""))
+        ]
+        for mp in stale:
+            data.pop(mp, None)
+        for mp, entry in sorted(data.items()):
+            result.append(
+                MountPointInfo(
+                    mountPoint=mp,
+                    source=entry.get("source", ""),
+                    localPath=entry.get("localPath", ""),
+                    extraConfigs=entry.get("extraConfigs", {}),
+                )
+            )
+    return result
+
+
+def getMountPath(mount_point: str, scope: str = "") -> str:
+    """Get the local file system path backing a mount point."""
+    mount_point = _normalize_mount_point(mount_point)
+    with _registry_transaction() as data:
+        entry = data.get(mount_point)
+    if entry is None:
+        raise ValueError(f"Mount point not found: {mount_point}")
+    return entry["localPath"]
+
+
+# ---------------------------------------------------------------------------
+# Internals
+# ---------------------------------------------------------------------------
+
+def _record_mount(
+    mount_point: str,
+    source: str,
+    local_path: str,
+    extra_configs: Dict[str, Any],
+) -> None:
+    # Never persist raw credentials in the on-disk registry
+    safe_configs = {
+        k: v
+        for k, v in extra_configs.items()
+        if k not in ("accountKey", "sasToken")
+    }
+    with _registry_transaction() as data:
+        data[mount_point] = {
+            "source": source,
+            "localPath": local_path,
+            "extraConfigs": safe_configs,
+        }
+
+
+def _exec_unmount(local_path: str) -> None:
+    timeout = DEFAULT_MOUNT_TIMEOUT
+    for cmd in (["fusermount", "-u", local_path], ["umount", local_path]):
+        try:
+            subprocess.run(
+                cmd,
+                check=True,
+                timeout=timeout,
+                capture_output=True,
+            )
+            return
+        except (FileNotFoundError, subprocess.CalledProcessError):
+            continue
+    raise RuntimeError(f"Failed to unmount {local_path}")

notebookutils/fs/_mount_azure.py

@@ -0,0 +1,188 @@
+"""FUSE execution layer for Azure Blob Storage / ADLS Gen2 via blobfuse2.
+
+Because blobfuse2 cannot consume abstract Python credential classes, this
+module unrolls the ``DefaultAzureCredential`` chain to discover which concrete
+mechanism succeeds locally, then maps it onto an explicit blobfuse2 auth-type
+(``spn``, ``azcli``, ``msi``, ``key``, ``sas``).
+"""
+
+from __future__ import annotations
+
+import logging
+import os
+import subprocess
+import tempfile
+from typing import Any, Dict, Optional, Tuple
+
+from notebookutils.fs._mount import DEFAULT_MOUNT_TIMEOUT
+
+logging.getLogger("azure.identity").setLevel(logging.ERROR)
+
+
+def _parse_azure_source(source: str, config: Dict[str, Any]) -> Tuple[str, str, str]:
+    """Return ``(account, container, prefix)`` from an abfs(s):// URL.
+
+    Supports both ``abfss://container@account.dfs.core.windows.net/path`` and
+    ``abfs://container/path`` (account from registry ``account_name``).
+    """
+    clean = source.split("://", 1)[1]
+    if "@" in clean:
+        container, _, host_and_path = clean.partition("@")
+        host, _, prefix = host_and_path.partition("/")
+        account = host.split(".", 1)[0]
+    else:
+        container, _, prefix = clean.partition("/")
+        account = config.get("account_name", "")
+    account = config.get("account_name") or account
+    if not account:
+        raise ValueError(
+            "Cannot determine the Azure storage account name. Use the "
+            "abfss://container@account.dfs.core.windows.net form or call "
+            "configure('abfs', account_name=...)."
+        )
+    return account, container, prefix
+
+
+def _resolve_auth(
+    account: str, config: Dict[str, Any], extra_configs: Dict[str, Any]
+) -> Tuple[str, Dict[str, str]]:
+    """Resolve blobfuse2 auth mode + parameters.
+
+    Priority: explicit extraConfigs (accountKey/sasToken) > explicit registry
+    keys > DefaultAzureCredential chain unrolling.
+    """
+    if extra_configs.get("accountKey"):
+        return "key", {"account-key": str(extra_configs["accountKey"])}
+    if extra_configs.get("sasToken"):
+        return "sas", {"sas": str(extra_configs["sasToken"])}
+    if config.get("account_key"):
+        return "key", {"account-key": str(config["account_key"])}
+    if config.get("sas_token"):
+        return "sas", {"sas": str(config["sas_token"])}
+
+    auth_type = _unroll_credential_chain(account)
+    params: Dict[str, str] = {}
+    if auth_type == "spn":
+        params = {
+            "clientid": os.environ.get("AZURE_CLIENT_ID", ""),
+            "clientsecret": os.environ.get("AZURE_CLIENT_SECRET", ""),
+            "tenantid": os.environ.get("AZURE_TENANT_ID", ""),
+        }
+    return auth_type, params
+
+
+def _unroll_credential_chain(account: str) -> str:
+    """Probe the DefaultAzureCredential chain and map the first working
+    sub-credential onto a blobfuse2 auth-type."""
+    try:
+        from azure.identity import DefaultAzureCredential
+    except ImportError:
+        return "msi"  # fallback guess without the SDK
+
+    scope = f"https://{account}.dfs.core.windows.net/.default"
+    try:
+        cred = DefaultAzureCredential()
+        for sub_cred in getattr(cred, "credentials", []):
+            try:
+                sub_cred.get_token(scope)
+            except Exception:
+                continue
+            class_name = type(sub_cred).__name__
+            if class_name == "EnvironmentCredential":
+                return "spn"
+            if class_name in ("AzureCliCredential", "AzureDeveloperCliCredential"):
+                return "azcli"
+            if class_name in ("ManagedIdentityCredential",):
+                return "msi"
+            return "msi"
+    except Exception:
+        pass
+    return "msi"
+
+
+def _build_yaml_config(
+    account: str,
+    container: str,
+    prefix: str,
+    mount_point: str,
+    cache_dir: str,
+    auth_type: str,
+    auth_params: Dict[str, str],
+    mount_opts: Dict[str, Any],
+    extra_configs: Dict[str, Any],
+) -> str:
+    file_cache_timeout = int(
+        extra_configs.get(
+            "fileCacheTimeout", mount_opts.get("fileCacheTimeout", 120)
+        )
+    )
+    allow_other = str(bool(mount_opts.get("allow_other", True))).lower()
+
+    lines = [
+        f"allow-other: {allow_other}",
+        "components:",
+        "  - libfuse",
+        "  - file_cache",
+        "  - attr_cache",
+        "  - azstorage",
+        "libfuse:",
+        f"  mount-path: {mount_point}",
+        "file_cache:",
+        f"  path: {cache_dir}",
+        f"  timeout-sec: {file_cache_timeout}",
+        "azstorage:",
+        "  type: adls",
+        f"  account-name: {account}",
+        f"  container: {container}",
+        f"  mode: {auth_type}",
+    ]
+    if prefix:
+        lines.append(f"  subdirectory: {prefix.rstrip('/')}")
+    for key, value in auth_params.items():
+        if value:
+            lines.append(f"  {key}: {value}")
+    return "\n".join(lines) + "\n"
+
+
+def exec_blobfuse2(
+    source: str,
+    mount_point: str,
+    config: Dict[str, Any],
+    mount_opts: Dict[str, Any],
+    extra_configs: Dict[str, Any],
+    cache_dir: Optional[str] = None,
+) -> None:
+    account, container, prefix = _parse_azure_source(source, config)
+    auth_type, auth_params = _resolve_auth(account, config, extra_configs)
+
+    cache_dir = cache_dir or mount_opts.get("cache_path") or tempfile.mkdtemp(
+        prefix=f"blobfuse_cache_{container}_"
+    )
+    os.makedirs(cache_dir, exist_ok=True)
+
+    yaml_config = _build_yaml_config(
+        account,
+        container,
+        prefix,
+        mount_point,
+        cache_dir,
+        auth_type,
+        auth_params,
+        mount_opts,
+        extra_configs,
+    )
+
+    fd, cfg_name = tempfile.mkstemp(suffix=".yaml", prefix="blobfuse2_cfg_")
+    try:
+        os.fchmod(fd, 0o600)
+        with os.fdopen(fd, "w") as tmp:
+            tmp.write(yaml_config)
+        timeout = int(extra_configs.get("timeout", DEFAULT_MOUNT_TIMEOUT))
+        subprocess.run(
+            ["blobfuse2", "mount", mount_point, f"--config-file={cfg_name}"],
+            check=True,
+            timeout=timeout,
+        )
+    finally:
+        if os.path.exists(cfg_name):
+            os.remove(cfg_name)

notebookutils/fs/_mount_gcs.py

@@ -0,0 +1,51 @@
+"""FUSE execution layer for Google Cloud Storage via gcsfuse."""
+
+from __future__ import annotations
+
+import os
+import subprocess
+from typing import Any, Dict
+
+from notebookutils.fs._mount import DEFAULT_MOUNT_TIMEOUT
+
+
+def exec_gcsfuse(
+    source: str,
+    mount_point: str,
+    config: Dict[str, Any],
+    mount_opts: Dict[str, Any],
+    extra_configs: Dict[str, Any],
+) -> None:
+    """Translate registry configuration into a ``gcsfuse`` CLI invocation.
+
+    Credential resolution follows GCP Application Default Credentials:
+    explicit key file (registry ``token`` or ``GOOGLE_APPLICATION_CREDENTIALS``)
+    else the GCE metadata server, handled natively by gcsfuse.
+    """
+    remainder = source.split("://", 1)[1]
+    bucket, _, prefix = remainder.partition("/")
+
+    cmd = ["gcsfuse"]
+
+    key_file = os.environ.get("GOOGLE_APPLICATION_CREDENTIALS") or config.get("token")
+    if key_file and isinstance(key_file, str) and os.path.exists(key_file):
+        cmd.extend(["--key-file", key_file])
+
+    if prefix:
+        cmd.extend(["--only-dir", prefix.rstrip("/")])
+
+    if mount_opts.get("allow_other", True):
+        cmd.extend(["-o", "allow_other"])
+    if "uid" in mount_opts:
+        cmd.extend(["--uid", str(mount_opts["uid"])])
+    if "gid" in mount_opts:
+        cmd.extend(["--gid", str(mount_opts["gid"])])
+
+    if "fileCacheTimeout" in extra_configs:
+        ttl = int(extra_configs["fileCacheTimeout"])
+        cmd.extend(["--stat-cache-ttl", f"{ttl}s", "--type-cache-ttl", f"{ttl}s"])
+
+    cmd.extend([bucket, mount_point])
+
+    timeout = int(extra_configs.get("timeout", DEFAULT_MOUNT_TIMEOUT))
+    subprocess.run(cmd, check=True, timeout=timeout)

notebookutils/fs/_mount_s3.py

@@ -0,0 +1,70 @@
+"""FUSE execution layer for Amazon S3 (and S3-compatible) via mount-s3."""
+
+from __future__ import annotations
+
+import os
+import subprocess
+from typing import Any, Dict
+
+from notebookutils.fs._mount import DEFAULT_MOUNT_TIMEOUT
+
+
+def exec_mount_s3(
+    source: str,
+    mount_point: str,
+    config: Dict[str, Any],
+    mount_opts: Dict[str, Any],
+    extra_configs: Dict[str, Any],
+) -> None:
+    """Translate registry configuration into a ``mount-s3`` CLI invocation.
+
+    Credential resolution mirrors the AWS default chain: explicit registry
+    keys > profile > environment variables > IMDS instance profile (handled
+    natively by mount-s3).
+    """
+    remainder = source.split("://", 1)[1]
+    bucket, _, prefix = remainder.partition("/")
+
+    cmd = ["mount-s3", bucket, mount_point]
+    env = os.environ.copy()
+
+    if prefix:
+        prefix = prefix.rstrip("/") + "/"
+        cmd.extend(["--prefix", prefix])
+
+    if "profile" in config:
+        cmd.extend(["--profile", str(config["profile"])])
+    if config.get("anon", False):
+        cmd.append("--no-sign-request")
+    if "key" in config and "secret" in config:
+        env["AWS_ACCESS_KEY_ID"] = str(config["key"])
+        env["AWS_SECRET_ACCESS_KEY"] = str(config["secret"])
+        if config.get("token"):
+            env["AWS_SESSION_TOKEN"] = str(config["token"])
+
+    # Custom / S3-compatible endpoints (OVH, Hetzner, DigitalOcean, MinIO...)
+    endpoint = config.get("endpoint_url") or (
+        config.get("client_kwargs", {}) or {}
+    ).get("endpoint_url")
+    if endpoint:
+        cmd.extend(["--endpoint-url", str(endpoint)])
+        cmd.append("--force-path-style")
+
+    region = config.get("region") or (config.get("client_kwargs", {}) or {}).get(
+        "region_name"
+    )
+    if region:
+        cmd.extend(["--region", str(region)])
+
+    # OS-level mount tuning
+    if mount_opts.get("allow_other", True):
+        cmd.append("--allow-other")
+    if "uid" in mount_opts:
+        cmd.extend(["--uid", str(mount_opts["uid"])])
+    if "gid" in mount_opts:
+        cmd.extend(["--gid", str(mount_opts["gid"])])
+    if not mount_opts.get("read_only", False):
+        cmd.append("--allow-delete")
+
+    timeout = int(extra_configs.get("timeout", DEFAULT_MOUNT_TIMEOUT))
+    subprocess.run(cmd, env=env, check=True, timeout=timeout)

notebookutils/fs/_registry.py

@@ -0,0 +1,86 @@
+"""Thread-safe Global Registry for per-protocol configuration.
+
+The registry tracks two groups of parameters per protocol:
+
+- ``fsspec_options``: kwargs passed directly to ``fsspec.filesystem()``
+  (memory-space client wrappers).
+- ``mount_options``: OS-specific tuning parameters consumed during FUSE
+  mount command generation (uid, gid, cache settings, allow_other, ...).
+"""
+
+from __future__ import annotations
+
+import threading
+from typing import Any, Dict
+
+_lock = threading.RLock()
+
+_fs_registry: Dict[str, Dict[str, Any]] = {}
+_fs_instances: Dict[str, Any] = {}
+
+
+def configure(protocol: str, **kwargs: Any) -> None:
+    """Explicit setup interface to override environment discovery.
+
+    Strips out system-level mount configurations (``mount_options``) from
+    pure fsspec parameters. Invalidates any cached fsspec client instance
+    for the protocol.
+
+    Example::
+
+        notebookutils.fs.configure("s3", key="AKIA...", secret="...",
+                                   mount_options={"allow_other": False})
+    """
+    protocol = normalize_protocol(protocol)
+    mount_options = kwargs.pop("mount_options", {})
+    with _lock:
+        _fs_registry[protocol] = {
+            "fsspec_options": dict(kwargs),
+            "mount_options": dict(mount_options),
+        }
+        _fs_instances.pop(protocol, None)
+
+
+def normalize_protocol(protocol: str) -> str:
+    """Collapse protocol aliases to their canonical fsspec name."""
+    aliases = {
+        "s3a": "s3",
+        "s3n": "s3",
+        "abfss": "abfs",
+        "adl": "abfs",
+        "gs": "gcs",
+        "": "file",
+        "local": "file",
+    }
+    return aliases.get(protocol, protocol)
+
+
+def get_fsspec_options(protocol: str) -> Dict[str, Any]:
+    with _lock:
+        return dict(
+            _fs_registry.get(normalize_protocol(protocol), {}).get("fsspec_options", {})
+        )
+
+
+def get_mount_options(protocol: str) -> Dict[str, Any]:
+    with _lock:
+        return dict(
+            _fs_registry.get(normalize_protocol(protocol), {}).get("mount_options", {})
+        )
+
+
+def get_cached_instance(protocol: str) -> Any:
+    with _lock:
+        return _fs_instances.get(normalize_protocol(protocol))
+
+
+def set_cached_instance(protocol: str, instance: Any) -> None:
+    with _lock:
+        _fs_instances[normalize_protocol(protocol)] = instance
+
+
+def reset() -> None:
+    """Clear all configuration and cached clients (mainly for tests)."""
+    with _lock:
+        _fs_registry.clear()
+        _fs_instances.clear()