lamindb 1.1.0__py3-none-any.whl → 1.2.0__py3-none-any.whl

lamindb/{_query_manager.py → models/query_manager.py}

@@ -8,8 +8,6 @@ from lamindb_setup.core._docs import doc_args
 
 from lamindb.models import Record
 
-from .core._settings import settings
-
 if TYPE_CHECKING:
     from lamindb.base.types import StrField
 
@@ -19,7 +17,7 @@ class QueryManager(models.Manager):
 
     See Also:
 
-        :class:`lamindb.core.QuerySet`
+        :class:`lamindb.models.QuerySet`
         `django Manager <https://docs.djangoproject.com/en/4.2/topics/db/managers/>`__
 
     Examples:
@@ -39,8 +37,12 @@ class QueryManager(models.Manager):
                 self.source_field_name == "collection"
                 and self.target_field_name == "artifact"
             ):
+                from lamindb import settings
                 from lamindb.core._context import context
-                from lamindb.core._data import WARNING_RUN_TRANSFORM, _track_run_input
+                from lamindb.models.artifact import (
+                    WARNING_RUN_TRANSFORM,
+                    _track_run_input,
+                )
 
                 if (
                     context.run is None
@@ -71,14 +73,14 @@ class QueryManager(models.Manager):
     def df(self, **kwargs):
         """Convert to DataFrame.
 
-        For `**kwargs`, see :meth:`lamindb.core.QuerySet.df`.
+        For `**kwargs`, see :meth:`lamindb.models.QuerySet.df`.
         """
         return self.all().df(**kwargs)
 
     def all(self):
         """Return QuerySet of all.
 
-        For `**kwargs`, see :meth:`lamindb.core.QuerySet.df`.
+        For `**kwargs`, see :meth:`lamindb.models.QuerySet.df`.
         """
         self._track_run_input_manager()
         return self._all_base_class()
@@ -86,14 +88,14 @@ class QueryManager(models.Manager):
     @doc_args(Record.search.__doc__)
     def search(self, string: str, **kwargs):
         """{}"""  # noqa: D415
-        from ._record import _search
+        from .record import _search
 
         return _search(cls=self.all(), string=string, **kwargs)
 
     @doc_args(Record.lookup.__doc__)
     def lookup(self, field: StrField | None = None, **kwargs) -> NamedTuple:
         """{}"""  # noqa: D415
-        from ._record import _lookup
+        from .record import _lookup
 
         return _lookup(cls=self.all(), field=field, **kwargs)
 

lamindb/{_query_set.py → models/query_set.py}

@@ -5,37 +5,27 @@ import warnings
 from collections import UserList
 from collections.abc import Iterable
 from collections.abc import Iterable as IterableType
-from typing import TYPE_CHECKING, Any, Generic, NamedTuple, TypeVar
+from typing import TYPE_CHECKING, Any, Generic, NamedTuple, TypeVar, Union
 
 import pandas as pd
 from django.core.exceptions import FieldError
 from django.db import models
-from django.db.models import F, ForeignKey, ManyToManyField
+from django.db.models import F, ForeignKey, ManyToManyField, Subquery
 from django.db.models.fields.related import ForeignObjectRel
 from lamin_utils import logger
 from lamindb_setup.core._docs import doc_args
 
-from lamindb.models import (
-    Artifact,
-    CanCurate,
-    Collection,
-    Feature,
-    IsVersioned,
-    Record,
-    Run,
-    Schema,
-    Transform,
-)
+from lamindb.models._is_versioned import IsVersioned
+from lamindb.models.record import Record
 
-from .errors import DoesNotExist
-
-T = TypeVar("T")
+from ..errors import DoesNotExist
+from .can_curate import CanCurate
 
 if TYPE_CHECKING:
-    from collections.abc import Iterable
-
     from lamindb.base.types import ListLike, StrField
 
+T = TypeVar("T")
+
 
 class MultipleResultsFound(Exception):
     pass
@@ -82,6 +72,13 @@ def one_helper(self):
 
 
 def get_backward_compat_filter_kwargs(queryset, expressions):
+    from lamindb.models import (
+        Artifact,
+        Collection,
+        Schema,
+        Transform,
+    )
+
     if queryset.model in {Collection, Transform}:
         name_mappings = {
             "name": "key",
@@ -190,7 +187,7 @@ def process_expressions(queryset: QuerySet, expressions: dict) -> dict:
 
 
 def get(
-    registry_or_queryset: type[Record] | QuerySet,
+    registry_or_queryset: Union[type[Record], QuerySet],
     idlike: int | str | None = None,
     **expressions,
 ) -> Record:
@@ -214,10 +211,27 @@ def get(
     else:
         assert idlike is None  # noqa: S101
         expressions = process_expressions(qs, expressions)
+        # don't want _branch_code here in .get(), only in .filter()
+        expressions.pop("_branch_code", None)
         # inject is_latest for consistency with idlike
-        if issubclass(registry, IsVersioned) and "is_latest" not in expressions:
+        is_latest_was_not_in_expressions = "is_latest" not in expressions
+        if issubclass(registry, IsVersioned) and is_latest_was_not_in_expressions:
             expressions["is_latest"] = True
-        return registry.objects.using(qs.db).get(**expressions)
+        try:
+            return registry.objects.using(qs.db).get(**expressions)
+        except registry.DoesNotExist:
+            # handle the case in which the is_latest injection led to a missed query
+            if "is_latest" in expressions and is_latest_was_not_in_expressions:
+                expressions.pop("is_latest")
+                result = (
+                    registry.objects.using(qs.db)
+                    .filter(**expressions)
+                    .order_by("-created_at")
+                    .first()
+                )
+                if result is not None:
+                    return result
+            raise registry.DoesNotExist from registry.DoesNotExist
 
 
 class RecordList(UserList, Generic[T]):
@@ -240,7 +254,7 @@ class RecordList(UserList, Generic[T]):
 
     def save(self) -> RecordList[T]:
         """Save all records to the database."""
-        from lamindb._save import save
+        from lamindb.models.save import save
 
         save(self)
         return self
@@ -288,6 +302,11 @@ def get_basic_field_names(
 
 
 def get_feature_annotate_kwargs(show_features: bool | list[str]) -> dict[str, Any]:
+    from lamindb.models import (
+        Artifact,
+        Feature,
+    )
+
     features = Feature.filter()
     if isinstance(show_features, list):
         features.filter(name__in=show_features)
@@ -548,7 +567,17 @@ class QuerySet(models.QuerySet):
             include_kwargs = {s: F(s) for s in include if s not in field_names}
             annotate_kwargs.update(include_kwargs)
         if annotate_kwargs:
-            queryset = self.annotate(**annotate_kwargs)
+            id_subquery = self.values("id")
+            # for annotate, we want the queryset without filters so that joins don't affect the annotations
+            query_set_without_filters = self.model.objects.filter(
+                id__in=Subquery(id_subquery)
+            )
+            if self.query.order_by:
+                # Apply the same ordering to the new queryset
+                query_set_without_filters = query_set_without_filters.order_by(
+                    *self.query.order_by
+                )
+            queryset = query_set_without_filters.annotate(**annotate_kwargs)
         else:
             queryset = self
 
@@ -566,6 +595,8 @@ class QuerySet(models.QuerySet):
 
     def delete(self, *args, **kwargs):
         """Delete all records in the query set."""
+        from lamindb.models import Artifact, Collection, Run, Transform
+
         # both Transform & Run might reference artifacts
         if self.model in {Artifact, Collection, Transform, Run}:
             for record in self:
@@ -641,11 +672,12 @@ class QuerySet(models.QuerySet):
                 and value.strip("-").isalpha()
                 and "__" not in field
                 and hasattr(self.model, field)
-                and getattr(self.model, field).field.related_model
             ):
-                raise FieldError(
-                    f"Invalid lookup '{value}' for {field}. Did you mean {field}__name?"
-                )
+                field_attr = getattr(self.model, field)
+                if hasattr(field_attr, "field") and field_attr.field.related_model:
+                    raise FieldError(
+                        f"Invalid lookup '{value}' for {field}. Did you mean {field}__name?"
+                    )
 
         expressions = process_expressions(self, expressions)
         if len(expressions) > 0:
@@ -689,7 +721,7 @@ class QuerySet(models.QuerySet):
 @doc_args(Record.search.__doc__)
 def search(self, string: str, **kwargs):
     """{}"""  # noqa: D415
-    from ._record import _search
+    from .record import _search
 
     return _search(cls=self, string=string, **kwargs)
 
@@ -697,7 +729,7 @@ def search(self, string: str, **kwargs):
 @doc_args(Record.lookup.__doc__)
 def lookup(self, field: StrField | None = None, **kwargs) -> NamedTuple:
     """{}"""  # noqa: D415
-    from ._record import _lookup
+    from .record import _lookup
 
     return _lookup(cls=self, field=field, **kwargs)
 
@@ -705,7 +737,7 @@ def lookup(self, field: StrField | None = None, **kwargs) -> NamedTuple:
 @doc_args(CanCurate.validate.__doc__)
 def validate(self, values: ListLike, field: str | StrField | None = None, **kwargs):
     """{}"""  # noqa: D415
-    from ._can_curate import _validate
+    from .can_curate import _validate
 
     return _validate(cls=self, values=values, field=field, **kwargs)
 
@@ -713,7 +745,7 @@ def validate(self, values: ListLike, field: str | StrField | None = None, **kwar
 @doc_args(CanCurate.inspect.__doc__)
 def inspect(self, values: ListLike, field: str | StrField | None = None, **kwargs):
     """{}"""  # noqa: D415
-    from ._can_curate import _inspect
+    from .can_curate import _inspect
 
     return _inspect(cls=self, values=values, field=field, **kwargs)
 
@@ -721,7 +753,7 @@ def inspect(self, values: ListLike, field: str | StrField | None = None, **kwarg
 @doc_args(CanCurate.standardize.__doc__)
 def standardize(self, values: Iterable, field: str | StrField | None = None, **kwargs):
     """{}"""  # noqa: D415
-    from ._can_curate import _standardize
+    from .can_curate import _standardize
 
     return _standardize(cls=self, values=values, field=field, **kwargs)