lamindb_setup 0.78.0__py3-none-any.whl

lamindb_setup/core/cloud_sqlite_locker.py

@@ -0,0 +1,237 @@
+from __future__ import annotations
+
+from datetime import datetime, timezone
+from functools import wraps
+from typing import TYPE_CHECKING
+
+from lamin_utils import logger
+
+from .upath import UPath, create_mapper, infer_filesystem
+
+if TYPE_CHECKING:
+    from pathlib import Path
+    from uuid import UUID
+
+    from ._settings_instance import InstanceSettings
+    from ._settings_user import UserSettings
+
+EXPIRATION_TIME = 24 * 60 * 60 * 7  # 7 days
+
+MAX_MSG_COUNTER = 100  # print the msg after this number of iterations
+
+
+# raise if an instance is already locked
+# ignored by unlock_cloud_sqlite_upon_exception
+class InstanceLockedException(Exception):
+    pass
+
+
+class empty_locker:
+    has_lock = True
+
+    @classmethod
+    def lock(cls):
+        pass
+
+    @classmethod
+    def unlock(cls):
+        pass
+
+
+class Locker:
+    def __init__(self, user_uid: str, storage_root: UPath | Path, instance_id: UUID):
+        logger.debug(
+            f"init cloud sqlite locker: {user_uid}, {storage_root}, {instance_id}."
+        )
+
+        self._counter = 0
+
+        self.user = user_uid
+        self.instance_id = instance_id
+
+        self.root = storage_root
+        self.fs, root_str = infer_filesystem(storage_root)
+
+        exclusion_path = storage_root / f".lamindb/_exclusion/{instance_id.hex}"
+
+        self.mapper = create_mapper(self.fs, str(exclusion_path), create=True)
+
+        priorities_path = str(exclusion_path / "priorities")
+        if self.fs.exists(priorities_path):
+            self.users = self.mapper["priorities"].decode().split("*")
+
+            if self.user not in self.users:
+                self.priority = len(self.users)
+                self.users.append(self.user)
+                # potential problem here if 2 users join at the same time
+                # can be avoided by using separate files for each user
+                # and giving priority by timestamp
+                # here writing the whole list back because gcs
+                # does not support the append mode
+                self.mapper["priorities"] = "*".join(self.users).encode()
+            else:
+                self.priority = self.users.index(self.user)
+        else:
+            self.mapper["priorities"] = self.user.encode()
+            self.users = [self.user]
+            self.priority = 0
+
+        self.mapper[f"numbers/{self.user}"] = b"0"
+        self.mapper[f"entering/{self.user}"] = b"0"
+
+        # clean up failures
+        for user in self.users:
+            for endpoint in ("numbers", "entering"):
+                user_endpoint = f"{endpoint}/{user}"
+                user_path = str(exclusion_path / user_endpoint)
+                if not self.fs.exists(user_path):
+                    continue
+                if self.mapper[user_endpoint] == b"0":
+                    continue
+                period = (datetime.now() - self.modified(user_path)).total_seconds()
+                if period > EXPIRATION_TIME:
+                    logger.info(
+                        f"the lock of the user {user} seems to be stale, clearing"
+                        f" {endpoint}."
+                    )
+                    self.mapper[user_endpoint] = b"0"
+
+        self._has_lock = None
+        self._locked_by = None
+
+    def modified(self, path):
+        mtime = self.fs.modified(path)
+        # always convert to the local timezone before returning
+        # assume in utc if the time zone is not specified
+        if mtime.tzinfo is None:
+            mtime = mtime.replace(tzinfo=timezone.utc)
+        return mtime.astimezone().replace(tzinfo=None)
+
+    def _msg_on_counter(self, user):
+        if self._counter == MAX_MSG_COUNTER:
+            logger.warning(f"competing for the lock with the user {user}.")
+
+        if self._counter <= MAX_MSG_COUNTER:
+            self._counter += 1
+
+    # Lamport's bakery algorithm
+    def _lock_unsafe(self):
+        if self._has_lock:
+            return None
+
+        self._has_lock = True
+        self._locked_by = self.user
+
+        self.users = self.mapper["priorities"].decode().split("*")
+
+        self.mapper[f"entering/{self.user}"] = b"1"
+
+        numbers = [int(self.mapper[f"numbers/{user}"]) for user in self.users]
+        number = 1 + max(numbers)
+        self.mapper[f"numbers/{self.user}"] = str(number).encode()
+
+        self.mapper[f"entering/{self.user}"] = b"0"
+
+        for i, user in enumerate(self.users):
+            if i == self.priority:
+                continue
+
+            while self.mapper[f"entering/{user}"] == b"1":
+                self._msg_on_counter(user)
+
+            c_number = int(self.mapper[f"numbers/{user}"])
+
+            if c_number == 0:
+                continue
+
+            if (number > c_number) or (number == c_number and self.priority > i):
+                self._has_lock = False
+                self._locked_by = user
+                self.mapper[f"numbers/{self.user}"] = b"0"
+                return None
+
+    def lock(self):
+        try:
+            self._lock_unsafe()
+        except BaseException as e:
+            self.unlock()
+            self._clear()
+            raise e
+
+    def unlock(self):
+        self.mapper[f"numbers/{self.user}"] = b"0"
+        self._has_lock = None
+        self._locked_by = None
+        self._counter = 0
+
+    def _clear(self):
+        self.mapper[f"entering/{self.user}"] = b"0"
+
+    @property
+    def has_lock(self):
+        if self._has_lock is None:
+            logger.info("the lock has not been initialized, trying to obtain the lock.")
+            self.lock()
+
+        return self._has_lock
+
+
+_locker: Locker | None = None
+
+
+def get_locker(
+    isettings: InstanceSettings, usettings: UserSettings | None = None
+) -> Locker:
+    from ._settings import settings
+
+    global _locker
+
+    user_uid = settings.user.uid if usettings is None else usettings.uid
+    storage_root = isettings.storage.root
+
+    if (
+        _locker is None
+        or _locker.user != user_uid
+        or _locker.root != storage_root
+        or _locker.instance_id != isettings._id
+    ):
+        _locker = Locker(user_uid, storage_root, isettings._id)
+
+    return _locker
+
+
+def clear_locker():
+    global _locker
+
+    _locker = None
+
+
+# decorator
+def unlock_cloud_sqlite_upon_exception(ignore_prev_locker: bool = False):
+    """Decorator to unlock a cloud sqlite instance upon an exception.
+
+    Ignores `InstanceLockedException`.
+
+    Args:
+        ignore_prev_locker: `bool` - Do not unlock if locker hasn't changed.
+    """
+
+    def wrap_with_args(func):
+        # https://stackoverflow.com/questions/1782843/python-decorator-handling-docstrings
+        @wraps(func)
+        def wrapper(*args, **kwargs):
+            prev_locker = _locker
+            try:
+                return func(*args, **kwargs)
+            except Exception as exc:
+                if isinstance(exc, InstanceLockedException):
+                    raise exc
+                if ignore_prev_locker and _locker is prev_locker:
+                    raise exc
+                if _locker is not None and _locker._has_lock:
+                    _locker.unlock()
+                raise exc
+
+        return wrapper
+
+    return wrap_with_args

lamindb_setup/core/django.py

@@ -0,0 +1,115 @@
+from __future__ import annotations
+
+# flake8: noqa
+import builtins
+import os
+from pathlib import Path
+import time
+from lamin_utils import logger
+from ._settings_store import current_instance_settings_file
+from ._settings_instance import InstanceSettings
+import sys
+
+IS_RUN_FROM_IPYTHON = getattr(builtins, "__IPYTHON__", False)
+IS_SETUP = False
+IS_MIGRATING = False
+CONN_MAX_AGE = 299
+
+
+def close_if_health_check_failed(self) -> None:
+    if self.close_at is not None:
+        if time.monotonic() >= self.close_at:
+            self.close()
+        self.close_at = time.monotonic() + CONN_MAX_AGE
+
+
+# this bundles set up and migration management
+def setup_django(
+    isettings: InstanceSettings,
+    deploy_migrations: bool = False,
+    create_migrations: bool = False,
+    configure_only: bool = False,
+    init: bool = False,
+    view_schema: bool = False,
+):
+    if IS_RUN_FROM_IPYTHON:
+        os.environ["DJANGO_ALLOW_ASYNC_UNSAFE"] = "true"
+
+    import dj_database_url
+    import django
+    from django.conf import settings
+    from django.core.management import call_command
+
+    # configuration
+    if not settings.configured:
+        default_db = dj_database_url.config(
+            env="LAMINDB_DJANGO_DATABASE_URL",
+            default=isettings.db,
+            # see comment next to patching BaseDatabaseWrapper below
+            conn_max_age=CONN_MAX_AGE,
+            conn_health_checks=True,
+        )
+        DATABASES = {
+            "default": default_db,
+        }
+        from .._init_instance import get_schema_module_name
+
+        schema_names = ["core"] + list(isettings.schema)
+        installed_apps = [get_schema_module_name(n) for n in schema_names]
+        if view_schema:
+            installed_apps = installed_apps[::-1]  # to fix how apps appear
+            installed_apps += ["schema_graph", "django.contrib.staticfiles"]
+
+        kwargs = dict(
+            INSTALLED_APPS=installed_apps,
+            DATABASES=DATABASES,
+            DEFAULT_AUTO_FIELD="django.db.models.BigAutoField",
+            TIME_ZONE="UTC",
+            USE_TZ=True,
+        )
+        if view_schema:
+            kwargs.update(
+                DEBUG=True,
+                ROOT_URLCONF="lamindb_setup._schema",
+                SECRET_KEY="dummy",
+                TEMPLATES=[
+                    {
+                        "BACKEND": "django.template.backends.django.DjangoTemplates",
+                        "APP_DIRS": True,
+                    },
+                ],
+                STATIC_ROOT=f"{Path.home().as_posix()}/.lamin/",
+                STATICFILES_FINDERS=[
+                    "django.contrib.staticfiles.finders.AppDirectoriesFinder",
+                ],
+                STATIC_URL="static/",
+            )
+        settings.configure(**kwargs)
+        django.setup(set_prefix=False)
+        # https://laminlabs.slack.com/archives/C04FPE8V01W/p1698239551460289
+        from django.db.backends.base.base import BaseDatabaseWrapper
+
+        BaseDatabaseWrapper.close_if_health_check_failed = close_if_health_check_failed
+
+    if configure_only:
+        return None
+
+    # migrations management
+    if create_migrations:
+        call_command("makemigrations")
+        return None
+
+    if deploy_migrations:
+        call_command("migrate", verbosity=2)
+        isettings._update_cloud_sqlite_file(unlock_cloud_sqlite=False)
+    elif init:
+        global IS_MIGRATING
+        IS_MIGRATING = True
+        call_command("migrate", verbosity=0)
+        IS_MIGRATING = False
+
+    global IS_SETUP
+    IS_SETUP = True
+
+    if isettings.keep_artifacts_local:
+        isettings._search_local_root()

lamindb_setup/core/exceptions.py

@@ -0,0 +1,10 @@
+from __future__ import annotations
+
+
+class DefaultMessageException(Exception):
+    default_message: str | None = None
+
+    def __init__(self, message: str | None = None):
+        if message is None:
+            message = self.default_message
+        super().__init__(message)

lamindb_setup/core/hashing.py

@@ -0,0 +1,116 @@
+from __future__ import annotations
+
+"""Hashing.
+
+.. autosummary::
+   :toctree: .
+
+   hash_set
+   hash_file
+
+"""
+
+import base64
+import hashlib
+from concurrent.futures import ThreadPoolExecutor
+from typing import TYPE_CHECKING
+
+import psutil
+
+HASH_LENGTH = 22
+
+if TYPE_CHECKING:
+    from collections.abc import Iterable
+
+    from .types import Path, UPathStr
+
+
+def hash_and_encode_as_b62(s: str) -> str:
+    from lamin_utils._base62 import encodebytes
+
+    return encodebytes(hashlib.md5(s.encode()).digest())
+
+
+def to_b64_str(bstr: bytes):
+    b64 = base64.urlsafe_b64encode(bstr).decode().strip("=")
+    return b64
+
+
+def b16_to_b64(s: str):
+    return to_b64_str(base64.b16decode(s.strip('"'), casefold=True))
+
+
+# a lot to read about this: lamin-notes/2022/hashing
+def hash_set(s: set[str]) -> str:
+    bstr = ":".join(sorted(s)).encode("utf-8")
+    # as we're truncating at 22 b64, we choose md5 over sha512
+    return to_b64_str(hashlib.md5(bstr).digest())[:HASH_LENGTH]
+
+
+def hash_md5s_from_dir(hashes: Iterable[str]) -> tuple[str, str]:
+    # need to sort below because we don't want the order of parsing the dir to
+    # affect the hash
+    digests = b"".join(
+        hashlib.md5(hash.encode("utf-8")).digest() for hash in sorted(hashes)
+    )
+    digest = hashlib.md5(digests).digest()
+    return to_b64_str(digest)[:HASH_LENGTH], "md5-d"
+
+
+def hash_code(file_path: UPathStr):
+    with open(file_path, "rb") as fp:
+        data = fp.read()
+    data_size = len(data)
+    header = f"blob {data_size}\0".encode()
+    blob = header + data
+    return hashlib.sha1(blob)
+
+
+def hash_file(
+    file_path: Path,
+    file_size: int | None = None,
+    chunk_size: int | None = 50 * 1024 * 1024,
+) -> tuple[str, str]:
+    with open(file_path, "rb") as fp:
+        if file_size is None:
+            fp.seek(0, 2)
+            file_size = fp.tell()
+            fp.seek(0, 0)
+        if chunk_size is None:
+            chunk_size = file_size
+        first_chunk = fp.read(chunk_size)
+        if file_size <= chunk_size:
+            digest = hashlib.md5(first_chunk).digest()
+            hash_type = "md5"
+        else:
+            fp.seek(-chunk_size, 2)
+            last_chunk = fp.read(chunk_size)
+            digest = hashlib.sha1(
+                hashlib.sha1(first_chunk).digest() + hashlib.sha1(last_chunk).digest()
+            ).digest()
+            hash_type = "sha1-fl"
+    return to_b64_str(digest)[:HASH_LENGTH], hash_type
+
+
+def hash_dir(path: Path):
+    files = (subpath for subpath in path.rglob("*") if subpath.is_file())
+
+    def hash_size(file):
+        file_size = file.stat().st_size
+        return hash_file(file, file_size)[0], file_size
+
+    try:
+        n_workers = len(psutil.Process().cpu_affinity())
+    except AttributeError:
+        n_workers = psutil.cpu_count()
+    if n_workers > 1:
+        with ThreadPoolExecutor(n_workers) as pool:
+            hashes_sizes = pool.map(hash_size, files)
+    else:
+        hashes_sizes = map(hash_size, files)
+    hashes, sizes = zip(*hashes_sizes)
+
+    hash, hash_type = hash_md5s_from_dir(hashes)
+    n_objects = len(hashes)
+    size = sum(sizes)
+    return size, hash, hash_type, n_objects

lamindb_setup/core/types.py

@@ -0,0 +1,17 @@
+from __future__ import annotations
+
+"""Types.
+
+.. autosummary::
+   :toctree: .
+
+   UPathStr
+"""
+# we need Union here because __future__ annotations doesn't work with TypeAlias
+from pathlib import Path
+from typing import Union
+
+# UPath is subclass of Path, hence, it's not necessary to list UPath
+# we keep it in the name of the TypeAlias to make it clear to users that
+# cloud paths are allowed / PathStr is often associated with local paths
+UPathStr = Union[str, Path]  # typing.TypeAlias, >3.10 on but already deprecated