diff --git a/lamindb/core/_settings.py b/lamindb/core/_settings.py index e773a76c3..69efc137a 100644 --- a/lamindb/core/_settings.py +++ b/lamindb/core/_settings.py @@ -68,7 +68,7 @@ def __repr__(self) -> str: # pragma: no cover verbosity_color = colors.yellow if self.verbosity == "warning" else colors.green verbosity_str = verbosity_color(self.verbosity) - storage_root = self._storage_settings.root_as_str + storage_root = self.storage.root_as_str storage_str = colors.italic(storage_root) instance_str = colors.italic(self.instance_uid) @@ -123,26 +123,6 @@ def annotation(self) -> AnnotationSettings: FAQ: :doc:`/faq/track-run-inputs` """ - __using_key: str | None = None - _using_storage: str | None = None - - @property - def _using_key(self) -> str | None: - """Key for Django database settings.""" - return self.__using_key - - @_using_key.setter - def _using_key(self, value: str | None): - ln_setup.settings._using_key = value - self.__using_key = value - - @property - def _storage_settings(self) -> ln_setup.core.StorageSettings: - if self._using_storage is None: - storage_settings = ln_setup.settings.storage - else: - storage_settings = ln_setup.core.StorageSettings(root=self._using_storage) - return storage_settings @property def sync_git_repo(self) -> str | None: @@ -205,7 +185,7 @@ def storage(self) -> StorageSettings: ) ln.settings.storage = "s3://some-bucket", kwargs """ - return self._storage_settings + return ln_setup.settings.storage @storage.setter def storage(self, path_kwargs: AnyPathStr | tuple[AnyPathStr, Mapping]): diff --git a/lamindb/core/storage/_backed_access.py b/lamindb/core/storage/_backed_access.py index 740372111..2046ff70e 100644 --- a/lamindb/core/storage/_backed_access.py +++ b/lamindb/core/storage/_backed_access.py @@ -73,7 +73,7 @@ def backed_access( artifact_or_filepath: Artifact | UPath, mode: str = "r", engine: Literal["pyarrow", "polars"] = "pyarrow", - using_key: str | None = None, + using: str | None = None, **kwargs, ) -> ( AnnDataAccessor @@ -91,7 +91,7 @@ def backed_access( if isinstance(artifact_or_filepath, Artifact): artifact = artifact_or_filepath - objectpath, _ = filepath_from_artifact(artifact, using_key=using_key) + objectpath, _ = filepath_from_artifact(artifact, using=using) else: artifact = None objectpath = artifact_or_filepath diff --git a/lamindb/core/storage/paths.py b/lamindb/core/storage/paths.py index 67923f2a2..5ab1b6679 100644 --- a/lamindb/core/storage/paths.py +++ b/lamindb/core/storage/paths.py @@ -69,7 +69,7 @@ def check_path_is_child_of_root(path: AnyPathStr, root: AnyPathStr) -> bool: def attempt_accessing_path( artifact: Artifact, storage_key: str, - using_key: str | None = None, + using: str | None = None, access_token: str | None = None, ) -> tuple[UPath, StorageSettings]: # check whether the file is in the default db and whether storage @@ -78,19 +78,19 @@ def attempt_accessing_path( if ( artifact._state.db in ("default", None) - and artifact.storage_id == settings._storage_settings._id + and artifact.storage_id == settings.storage._id ): if access_token is None: - storage_settings = settings._storage_settings + storage_settings = settings.storage else: storage_settings = StorageSettings( settings.storage.root, access_token=access_token ) else: - if artifact._state.db not in ("default", None) and using_key is None: + if artifact._state.db not in ("default", None) and using is None: storage = Storage.connect(artifact._state.db).get(id=artifact.storage_id) else: - storage = Storage.objects.using(using_key).get(id=artifact.storage_id) + storage = Storage.objects.using(using).get(id=artifact.storage_id) # find a better way than passing None to instance_settings in the future! storage_settings = StorageSettings(storage.root, access_token=access_token) path = storage_settings.key_to_filepath(storage_key) @@ -98,14 +98,12 @@ def attempt_accessing_path( def filepath_from_artifact( - artifact: Artifact, using_key: str | None = None + artifact: Artifact, using: str | None = None ) -> tuple[UPath, StorageSettings | None]: if (local_filepath := getattr(artifact, "_local_filepath", None)) is not None: return local_filepath.resolve(), None storage_key = auto_storage_key_from_artifact(artifact) - path, storage_settings = attempt_accessing_path( - artifact, storage_key, using_key=using_key - ) + path, storage_settings = attempt_accessing_path(artifact, storage_key, using=using) return path, storage_settings @@ -132,9 +130,9 @@ def _cache_key_from_artifact_storage( # return filepath and cache_key if needed def filepath_cache_key_from_artifact( - artifact: Artifact, using_key: str | None = None + artifact: Artifact, using: str | None = None ) -> tuple[UPath, str | None]: - filepath, storage_settings = filepath_from_artifact(artifact, using_key) + filepath, storage_settings = filepath_from_artifact(artifact, using) if isinstance(filepath, LocalPathClasses): return filepath, None cache_key = _cache_key_from_artifact_storage(artifact, storage_settings) @@ -180,13 +178,13 @@ def store_file_or_folder( shutil.copytree(local_path, storage_path) -def delete_storage_using_key( +def delete_storage_using( artifact: Artifact, storage_key: str, raise_file_not_found_error: bool = True, - using_key: str | None = None, + using: str | None = None, ) -> None | str: - filepath, _ = attempt_accessing_path(artifact, storage_key, using_key=using_key) + filepath, _ = attempt_accessing_path(artifact, storage_key, using=using) return delete_storage( filepath, raise_file_not_found_error=raise_file_not_found_error ) diff --git a/lamindb/curators/core.py b/lamindb/curators/core.py index 68dbd9d30..04984cdcd 100644 --- a/lamindb/curators/core.py +++ b/lamindb/curators/core.py @@ -664,17 +664,21 @@ def __init__( schema: Schema, slot: str | None = None, require_saved_schema: bool = True, + using: str | None = None, ) -> None: super().__init__( dataset=dataset, schema=schema, require_saved_schema=require_saved_schema ) + self._using = using categoricals = [] features = [] feature_ids: set[int] = set() if schema.flexible: - features += Feature.filter(name__in=self._dataset.keys()).to_list() + features += ( + Feature.connect(using).filter(name__in=self._dataset.keys()).to_list() + ) feature_ids = {feature.id for feature in features} if schema.n_members and schema.n_members > 0: @@ -847,6 +851,7 @@ def __init__( categoricals=categoricals, index=schema.index, slot=slot, + using=using, maximal_set=schema.maximal_set, schema=schema, ) @@ -978,6 +983,7 @@ def __init__( slot: str | None = None, features: dict[str, Any] | None = None, require_saved_schema: bool = True, + using: str | None = None, ) -> None: # loads or opens dataset, dataset may be an artifact super().__init__( @@ -986,12 +992,14 @@ def __init__( features=features, require_saved_schema=require_saved_schema, ) + self._using = using # uses open dataset at self._dataset self._atomic_curator = ComponentCurator( dataset=self._dataset, schema=schema, slot=slot, require_saved_schema=require_saved_schema, + using=using, ) # Handle (nested) attrs if slot is None and schema.slots: @@ -1013,6 +1021,7 @@ def __init__( slot_schema, slot=slot_name, require_saved_schema=require_saved_schema, + using=using, ) elif slot_name != "__external__": raise ValueError( @@ -1066,6 +1075,7 @@ def __init__( schema: Schema, slot: str | None = None, require_saved_schema: bool = False, + using: str | None = None, ) -> None: if not isinstance(dataset, dict) and not isinstance(dataset, Artifact): raise InvalidArgument("The dataset must be a dict or dict-like artifact.") @@ -1076,7 +1086,11 @@ def __init__( d = dataset df = convert_dict_to_dataframe_for_validation(d, schema) # type: ignore super().__init__( - df, schema, slot=slot, require_saved_schema=require_saved_schema + df, + schema, + slot=slot, + require_saved_schema=require_saved_schema, + using=using, ) @@ -1545,6 +1559,7 @@ def __init__( type_uid: str | None = None, maximal_set: bool = True, # whether unvalidated categoricals cause validation failure. schema: Schema = None, + using: str | None = None, # target instance for cross-instance curation ) -> None: self._values_getter = values_getter self._values_setter = values_setter @@ -1561,6 +1576,7 @@ def __init__( self.records = None self._maximal_set = maximal_set self._type_record = None + self._using = using self._registry = self._field.field.model self._field_name = self._field.field.name self._filter_kwargs = {} @@ -1568,7 +1584,9 @@ def __init__( if filter_str and filter_str != "unsaved": self._filter_kwargs.update( resolve_relation_filters( - parse_filter_string(filter_str), self._registry + parse_filter_string(filter_str), + self._registry, + using=self._using, ) # type: ignore ) if self._registry.__base__.__name__ == "BioRecord": @@ -1590,6 +1608,7 @@ def __init__( self._type_record = get_record_type_from_uid( self._registry, self._type_uid, + using=self._using, ) if hasattr(self._registry, "_name_field"): @@ -1748,7 +1767,9 @@ def _add_validated(self) -> tuple[list, list]: if filter_str: parsed_filters = parse_filter_string(filter_str) filter_kwargs.update( - resolve_relation_filters(parsed_filters, registry) + resolve_relation_filters( + parsed_filters, registry, using=self._using + ) ) if registry.__base__.__name__ == "BioRecord": organism_record = get_organism_record_from_field( @@ -1774,7 +1795,7 @@ def _add_validated(self) -> tuple[list, list]: # When we have a Schema with typed members, # scope the query to the types present in the schema's members (plus untyped features) # to avoid ambiguous matches across different feature types. - qs = registry.filter() + qs = registry.connect(self._using).filter() if self._schema and self._schema.n_members: type_ids = { m.type_id @@ -1782,7 +1803,7 @@ def _add_validated(self) -> tuple[list, list]: if m.type_id is not None } if type_ids: - qs = registry.filter( + qs = registry.connect(self._using).filter( Q(type_id__in=type_ids) | Q(type_id__isnull=True) ) self._subtype_query_set = qs @@ -1837,6 +1858,7 @@ def _add_validated(self) -> tuple[list, list]: remaining_values, field=field, mute=True, + using=self._using, **filter_kwargs, # type: ignore ) existing_and_public_values = [ @@ -1939,12 +1961,20 @@ def _validate( filter_str = result.get("filter_str", "") if filter_str: parsed_filters = parse_filter_string(filter_str) - filter_kwargs.update(resolve_relation_filters(parsed_filters, registry)) + filter_kwargs.update( + resolve_relation_filters( + parsed_filters, registry, using=self._using + ) + ) registry_or_queryset = registry if self._subtype_query_set is not None and registry == self._registry: registry_or_queryset = self._subtype_query_set # first inspect against the registry - inspect_result = registry_or_queryset.filter(**filter_kwargs).inspect( + if registry_or_queryset is registry: + queryset = registry.connect(self._using).filter(**filter_kwargs) + else: + queryset = registry_or_queryset.filter(**filter_kwargs) + inspect_result = queryset.inspect( non_validated, field=field, mute=True, @@ -2074,10 +2104,12 @@ def __init__( slot: str | None = None, maximal_set: bool = False, schema: Schema | None = None, + using: str | None = None, ) -> None: self._non_validated = None self._index = index self._schema = schema + self._using = using self._artifact: Artifact = None # pass the dataset as an artifact self._dataset: Any = df # pass the dataset as an AnyPathStr or data object if isinstance(self._dataset, Artifact): @@ -2115,6 +2147,7 @@ def __init__( if schema.id is None else f"schemas__id={schema.id}", schema=schema, + using=using, ) for feature in self._categoricals: result = parse_dtype(feature._dtype_str)[0] @@ -2135,6 +2168,7 @@ def __init__( cat_manager=self, filter_str=result["filter_str"], type_uid=result.get("type_uid"), + using=using, ) if index is not None and index._dtype_str.startswith("cat"): result = parse_dtype(index._dtype_str)[0] @@ -2150,6 +2184,7 @@ def __init__( cat_manager=self, filter_str=result["filter_str"], type_uid=result.get("type_uid"), + using=using, ) @property diff --git a/lamindb/models/__init__.py b/lamindb/models/__init__.py index 69d26f83a..0736c52be 100644 --- a/lamindb/models/__init__.py +++ b/lamindb/models/__init__.py @@ -239,3 +239,10 @@ ) FeatureValue = JsonValue # backward compatibility + +from .sqlrecord import _ensure_lamindb_router + +# Best-effort early registration of LaminDBRouter. +# This may no-op if Django settings are not configured yet, so we also call +# _ensure_lamindb_router() in add_db_connection() as a late safety net. +_ensure_lamindb_router() diff --git a/lamindb/models/_feature_manager.py b/lamindb/models/_feature_manager.py index 82b7d0018..20ad93655 100644 --- a/lamindb/models/_feature_manager.py +++ b/lamindb/models/_feature_manager.py @@ -1700,7 +1700,10 @@ def add_values( ) schema = Schema(feature_objects) ExperimentalDictCurator( - dictionary, schema, require_saved_schema=False + dictionary, + schema, + require_saved_schema=False, + using=self._host._state.db, ).validate() if host_is_record and schema.index is not None: from .record import strip_index_for_record_persistence @@ -1730,6 +1733,7 @@ def _add_values( from ..base.dtypes import is_iterable_of_sqlrecord from .can_curate import CanCurate + host_db = self._host._state.db host_is_record = self._host.__class__.__name__ == "Record" if host_is_record: feature_json_values: list[SQLRecord] = [] @@ -1746,12 +1750,12 @@ def _add_values( ) self._raise_not_validated_values(record_not_validated_values) if feature_json_values: - save(feature_json_values) + save(feature_json_values, using=host_db) for links in links_by_model.values(): try: - save(links, ignore_conflicts=False) + save(links, ignore_conflicts=False, using=host_db) except Exception: - save(links, ignore_conflicts=True) + save(links, ignore_conflicts=True, using=host_db) from .record import get_type_schema_index, persist_record_name if ( @@ -1918,7 +1922,9 @@ def set_values( if host_is_artifact: schema = self._get_external_schema() if schema is not None: - ExperimentalDictCurator(dictionary, schema).validate() + ExperimentalDictCurator( + dictionary, schema, using=self._host._state.db + ).validate() member_ids = set(schema.members.values_list("id", flat=True)) features_not_in_schema = [ feature.name @@ -2142,7 +2148,6 @@ def _add_from(self, data: Artifact | Collection, transfer_logs: dict = None): transfer_logs = {"mapped": [], "transferred": [], "run": None} from lamindb import settings - using_key = settings._using_key for slot, schema in data.features.slots.items(): # type: ignore try: members = schema.members @@ -2178,7 +2183,7 @@ def _add_from(self, data: Artifact | Collection, transfer_logs: dict = None): if n_new_members > 0: # transfer foreign keys needs to be run before transfer to default db transfer_fk_to_default_db_bulk( - new_members, using_key, transfer_logs=transfer_logs + new_members, None, transfer_logs=transfer_logs ) for feature in new_members: # not calling save=True here as in labels, because want to @@ -2187,7 +2192,7 @@ def _add_from(self, data: Artifact | Collection, transfer_logs: dict = None): # in the previous step transfer_fk_to_default_db_bulk transfer_to_default_db( feature, - using_key, + None, transfer_fk=False, transfer_logs=transfer_logs, ) @@ -2221,7 +2226,9 @@ def _add_from(self, data: Artifact | Collection, transfer_logs: dict = None): self._host.features._add_schema(schema_self, slot) -def bulk_set_features_in_records(records: Iterable[Record]) -> None: +def bulk_set_features_in_records( + records: Iterable[Record], using: str | None = None +) -> None: import numpy as np """Bulk-set lazy feature dictionaries for records. @@ -2241,6 +2248,8 @@ def bulk_set_features_in_records(records: Iterable[Record]) -> None: if len(records_with_features) == 0: return None + using = using if using not in (None, "default") else None + batch_schema: Schema | None = None batch_schema_index: Feature | None = None prepared_records: list[ @@ -2308,6 +2317,7 @@ def bulk_set_features_in_records(records: Iterable[Record]) -> None: ) data: dict[str, pd.Series] = {} + multivalued_columns: list[str] = [] for column in ordered_columns: # None from entirely-null columns is not a valid sentinel for extension # dtypes (StringDtype, BooleanDtype, Int64Dtype) — convert to pd.NA @@ -2316,22 +2326,30 @@ def bulk_set_features_in_records(records: Iterable[Record]) -> None: for v in (row.get(column, pd.NA) for row in prepared_rows) ] target_dtype = feature_dtype_by_name.get(column) - if target_dtype is not None: + is_multivalued = any( + isinstance(v, (list, tuple, set, np.ndarray)) for v in values + ) + if target_dtype is not None and not is_multivalued: data[column] = pd.Series(values, dtype=target_dtype) else: - data[column] = pd.Series(values) + data[column] = pd.Series(values, dtype="object") + if is_multivalued: + multivalued_columns.append(column) if data: dataframe = pd.DataFrame(data) else: dataframe = pd.DataFrame(index=range(len(prepared_rows))) dataframe = move_schema_index_column_to_dataframe_index(dataframe, batch_schema) + # set after the index move, which does not preserve `.attrs` + for column in multivalued_columns: + dataframe.attrs[column] = "list_of_categories" # Single-pass dataframe curation: # validate schema and resolve categoricals once for the entire batch. # # The resolved label records are then reused below when creating per-record # link rows, avoiding repeated registry calls for each row. - curator = DataFrameCurator(dataframe, batch_schema) + curator = DataFrameCurator(dataframe, batch_schema, using=using) curator.validate() members_by_name: dict[str, list[Feature]] = defaultdict(list) @@ -2412,18 +2430,18 @@ def bulk_set_features_in_records(records: Iterable[Record]) -> None: ) FeatureManager._raise_not_validated_values(not_validated_values) if feature_json_values: - save(feature_json_values) + save(feature_json_values, using=using) for links in links_by_model.values(): try: - save(links, ignore_conflicts=False) + save(links, ignore_conflicts=False, using=using) except Exception: - save(links, ignore_conflicts=True) + save(links, ignore_conflicts=True, using=using) from .save import bulk_update if batch_schema_index is not None: # only `name` was modified (via strip_index_for_record_persistence) # updating all fields generates a massive CASE WHEN SQL for large batches - bulk_update(records_with_features, update_fields=["name"]) + bulk_update(records_with_features, update_fields=["name"], using=using) for record in records_with_features: del record._features return None diff --git a/lamindb/models/_from_values.py b/lamindb/models/_from_values.py index d7123f2ff..d56802d46 100644 --- a/lamindb/models/_from_values.py +++ b/lamindb/models/_from_values.py @@ -57,6 +57,7 @@ def _from_values( standardize: bool = True, from_source: bool = True, mute: bool = False, + using: str | None = None, **filter_kwargs, ) -> SQLRecordList: """Get or create records from iterables.""" @@ -84,6 +85,7 @@ def _from_values( field=field, organism=organism_record, mute=mute, + using=using, **filter_kwargs, ) @@ -135,6 +137,7 @@ def get_existing_records( organism: SQLRecord | None = None, standardize: bool = True, mute: bool = False, + using: str | None = None, **filter_kwargs, ) -> tuple[list, Index, str]: """Get existing records from the database.""" @@ -144,7 +147,7 @@ def get_existing_records( # NOTE: existing records matching is agnostic to the source registry = field.field.model # type: ignore - queryset = registry.filter(**filter_kwargs) + queryset = registry.connect(using).filter(**filter_kwargs) if standardize: # log synonyms mapped terms @@ -391,7 +394,7 @@ def get_organism_record_from_field( # type: ignore field: FieldAttr, organism: str | SQLRecord | None = None, values: ListLike = None, - using_key: str | None = None, + using: str | None = None, ) -> SQLRecord | None: """Get organism record based on which field is used in from_values. @@ -399,7 +402,7 @@ def get_organism_record_from_field( # type: ignore field: the field of the registry for from_values organism: the organism to get the organism record for values: the values passed to from_values - using_key: the db to get the organism record from + using: the db to get the organism record from Returns: The organism record if both conditions are met: @@ -447,7 +450,7 @@ def get_organism_record_from_field( # type: ignore (v for v in values if isinstance(v, str) and v.startswith("ENS")), "" ) if first_ensembl: - return infer_organism_from_ensembl_id(first_ensembl, using_key) + return infer_organism_from_ensembl_id(first_ensembl, using) return create_or_get_organism_record( organism=organism, registry=registry, field=field diff --git a/lamindb/models/_label_manager.py b/lamindb/models/_label_manager.py index 2cf88f2bd..df255b5aa 100644 --- a/lamindb/models/_label_manager.py +++ b/lamindb/models/_label_manager.py @@ -240,9 +240,8 @@ def add_from(self, data: Artifact | Collection, transfer_logs: dict = None) -> N """ if transfer_logs is None: transfer_logs = {"mapped": [], "transferred": [], "run": None} - from lamindb import settings + using = data._state.db if data._state.db not in (None, "default") else None - using_key = settings._using_key for related_name, labels in _get_labels(data, instance=data._state.db).items(): labels = labels.all() if not labels.exists(): @@ -254,7 +253,7 @@ def add_from(self, data: Artifact | Collection, transfer_logs: dict = None) -> N new_labels = save_validated_records(labels) if len(new_labels) > 0: transfer_fk_to_default_db_bulk( - new_labels, using_key, transfer_logs=transfer_logs + new_labels, None, transfer_logs=transfer_logs ) for label in labels: keys: list = [] @@ -275,7 +274,7 @@ def add_from(self, data: Artifact | Collection, transfer_logs: dict = None) -> N keys.append(key) label_returned = transfer_to_default_db( label, - using_key, + using, transfer_logs=transfer_logs, transfer_fk=False, save=True, @@ -289,12 +288,12 @@ def add_from(self, data: Artifact | Collection, transfer_logs: dict = None) -> N new_features = save_validated_records(list(features)) if len(new_features) > 0: transfer_fk_to_default_db_bulk( - new_features, using_key, transfer_logs=transfer_logs + new_features, using, transfer_logs=transfer_logs ) for feature in new_features: transfer_to_default_db( feature, # type: ignore - using_key, + using, transfer_logs=transfer_logs, transfer_fk=False, ) diff --git a/lamindb/models/artifact.py b/lamindb/models/artifact.py index c9518d84e..b793281d0 100644 --- a/lamindb/models/artifact.py +++ b/lamindb/models/artifact.py @@ -227,7 +227,7 @@ def _identify_zarr_type(storepath, *, check: bool = True): def process_pathlike( filepath: UPath, storage: Storage, - using_key: str | None, + using: str | None, skip_existence_check: bool = False, ) -> tuple[Storage, bool]: """Determines the appropriate storage for a given path and whether to use an existing storage key.""" @@ -255,7 +255,7 @@ def process_pathlike( # already-registered storage locations result = None # within the hub, we don't want to perform check_path_in_existing_storage - if using_key is None: + if using is None: result = check_path_in_existing_storage( filepath, check_hub_register_storage=setup_settings.instance.is_on_hub ) @@ -346,7 +346,7 @@ def process_data( format: str | None, key: str | None, storage: Storage, - using_key: str | None, + using: str | None, skip_existence_check: bool = False, is_replace: bool = False, to_disk_kwargs: dict[str, Any] | None = None, @@ -387,7 +387,7 @@ def process_data( storage, use_existing_storage_key = process_pathlike( path, storage=storage, - using_key=using_key, + using=using, skip_existence_check=skip_existence_check, ) suffix, raw_suffix = CanonicalSuffix.extract_from_path(path) @@ -560,9 +560,9 @@ def get_stat_or_artifact( def check_path_in_existing_storage( path: Path | UPath, check_hub_register_storage: bool = False, - using_key: str | None = None, + using: str | None = None, ) -> Storage | None: - for storage in Storage.objects.using(using_key).order_by(Length("root").desc()): + for storage in Storage.objects.using(using).order_by(Length("root").desc()): # if path is part of storage, return it if _s().check_path_is_child_of_root(path, root=storage.root): return storage @@ -601,7 +601,7 @@ def get_artifact_kwargs_from_data( provisional_uid: str, version_tag: str | None, storage: Storage, - using_key: str | None = None, + using: str | None = None, is_replace: bool = False, skip_check_exists: bool = False, overwrite_versions: bool | None = None, @@ -617,7 +617,7 @@ def get_artifact_kwargs_from_data( format, key, storage, - using_key, + using, skip_check_exists, is_replace=is_replace, to_disk_kwargs=to_disk_kwargs, @@ -648,7 +648,7 @@ def get_artifact_kwargs_from_data( path=path, storage=storage, key=key, - instance=using_key, + instance=using, is_replace=is_replace, skip_hash_lookup=effective_skip_hash_lookup, skip_key_revises_lookup=skip_key_revises_lookup, @@ -1125,10 +1125,10 @@ def add_labels( ) -def delete_permanently(artifact: Artifact, storage: bool | None, using_key: str): +def delete_permanently(artifact: Artifact, storage: bool | None, using: str): # need to grab file path before deletion try: - path, _ = _s().filepath_from_artifact(artifact, using_key) + path, _ = _s().filepath_from_artifact(artifact, using) except OSError: # we can still delete the record logger.warning("Could not get path") @@ -1817,7 +1817,7 @@ def __init__( kind: str = kwargs.pop("kind", None) key: str | None = kwargs.pop("key", None) - using_key = kwargs.pop("using_key", None) + using = kwargs.pop("using", None) description: str | None = kwargs.pop("description", None) revises: Artifact | None = kwargs.pop("revises", None) refresh_revises_if_stale = revises is None @@ -2028,7 +2028,7 @@ def __init__( provisional_uid=provisional_uid, version_tag=version_tag, storage=storage, - using_key=using_key, + using=using, skip_check_exists=skip_check_exists, overwrite_versions=overwrite_versions, skip_hash_lookup=skip_hash_lookup, @@ -2204,14 +2204,12 @@ def path(self) -> UPath: artifact.path #> PosixPath('/home/runner/work/lamindb/lamindb/docs/guide/mydata/myfile.csv') """ - filepath, _ = _s().filepath_from_artifact(self, using_key=settings._using_key) + filepath, _ = _s().filepath_from_artifact(self) return filepath @property def _cache_path(self) -> UPath: - filepath, cache_key = _s().filepath_cache_key_from_artifact( - self, using_key=settings._using_key - ) + filepath, cache_key = _s().filepath_cache_key_from_artifact(self) if isinstance(filepath, LocalPathClasses): return filepath return setup_settings.paths.cloud_to_local_no_update( @@ -2819,8 +2817,7 @@ def from_dir( """ folderpath: UPath = create_path(path) # returns Path for local storage = settings.storage.record - using_key = settings._using_key - storage, use_existing_storage = process_pathlike(folderpath, storage, using_key) + storage, use_existing_storage = process_pathlike(folderpath, storage, None) folder_key_path: PurePath | Path if key is None: if not use_existing_storage: @@ -3123,10 +3120,7 @@ def open( f" Or no suffix for a folder with {', '.join(df_suffixes)} files" " (no mixing allowed)." ) - using_key = settings._using_key - filepath, cache_key = _s().filepath_cache_key_from_artifact( - self, using_key=using_key - ) + filepath, cache_key = _s().filepath_cache_key_from_artifact(self) is_tiledbsoma_w = ( filepath.name == "soma" or suffix == ".tiledbsoma" @@ -3156,11 +3150,10 @@ def open( open_cache = not isinstance( filepath, LocalPathClasses ) and not filepath.synchronize_to(localpath, just_check=True) + using = self._state.db if self._state.db not in (None, "default") else None if open_cache: try: - access = backed_access( - localpath, mode, engine, using_key=using_key, **kwargs - ) + access = backed_access(localpath, mode, engine, using, **kwargs) except Exception as e: # also ignore ValueError here because # such errors most probably just imply an incorrect argument @@ -3171,9 +3164,7 @@ def open( logger.warning( f"The cache might be corrupted: {e}. Trying to open directly." ) - access = backed_access( - filepath, mode, engine, using_key=using_key, **kwargs - ) + access = backed_access(filepath, mode, engine, using, **kwargs) # happens only if backed_access has been successful # delete the corrupted cache if localpath.is_dir(): @@ -3181,7 +3172,7 @@ def open( else: localpath.unlink(missing_ok=True) else: - access = backed_access(self, mode, engine, using_key=using_key, **kwargs) + access = backed_access(self, mode, engine, using, **kwargs) if is_tiledbsoma_w: def finalize(): @@ -3250,9 +3241,7 @@ def load( if access_memory.__class__.__name__ == "SpatialData": access_memory.path = self._cache_path else: - filepath, cache_key = _s().filepath_cache_key_from_artifact( - self, using_key=settings._using_key - ) + filepath, cache_key = _s().filepath_cache_key_from_artifact(self) cache_path = _synchronize_cleanup_on_error( filepath, cache_key=cache_key, print_progress=not mute ) @@ -3308,9 +3297,7 @@ def cache( if self._overwrite_versions and not self.is_latest: raise ValueError(OUTDATED_ARTIFACT_FILES_OVERWRITTEN_MSG) - filepath, cache_key = _s().filepath_cache_key_from_artifact( - self, using_key=settings._using_key - ) + filepath, cache_key = _s().filepath_cache_key_from_artifact(self) if mute: kwargs["print_progress"] = False cache_path = _synchronize_cleanup_on_error( @@ -3324,7 +3311,7 @@ def delete( self, permanent: bool | None = None, storage: bool | None = None, - using_key: str | None = None, + using: str | None = None, ) -> None: """Trash or permanently delete. @@ -3358,7 +3345,7 @@ def delete( artifact = ln.Artifact.get(key="folder.zarr". is_latest=True) artifact.delete() # delete all versions, the data will be deleted or prompted for deletion. """ - super().delete(permanent=permanent, storage=storage, using_key=using_key) + super().delete(permanent=permanent, storage=storage, using=using) # TODO: consider renaming the transfer argument to sync def save( @@ -3600,12 +3587,12 @@ def save( self._save_skip_storage(**kwargs) - using_key = None + using = None if "using" in kwargs: - using_key = kwargs["using"] + using = kwargs["using"] exception_upload = check_and_attempt_upload( self, - using_key, + using, access_token=access_token, print_progress=print_progress, **store_kwargs, @@ -3623,7 +3610,7 @@ def save( exception_clear = check_and_attempt_clearing( self, raise_file_not_found_error=raise_file_not_found_error, - using_key=using_key, + using=using, ) if exception_upload is not None: raise exception_upload diff --git a/lamindb/models/artifact_set.py b/lamindb/models/artifact_set.py index 0664f2932..af3d90c3c 100644 --- a/lamindb/models/artifact_set.py +++ b/lamindb/models/artifact_set.py @@ -225,7 +225,13 @@ def to_dataframe( ) # `type_id` points to record types (`is_type=True`) by model design. - record_type = Record.get(id=type_ids[0]) + # `type_ids` were read from `qs.db`; resolve the type on the same + # instance so cross-instance reads don't hit the default connection. + record_type = ( + Record.get(id=type_ids[0]) + if qs.db in (None, "default") + else Record.connect(qs.db).get(id=type_ids[0]) + ) qs._record_export_type = record_type logger.important(f"exporting {qs.count()} records of '{record_type.name}'") diff --git a/lamindb/models/can_curate.py b/lamindb/models/can_curate.py index b143668be..3031809fb 100644 --- a/lamindb/models/can_curate.py +++ b/lamindb/models/can_curate.py @@ -26,13 +26,13 @@ from .query_set import SQLRecordList -def _check_if_record_in_db(record: str | SQLRecord | None, using_key: str | None): - """Check if the record is from the using_key DB.""" +def _check_if_record_in_db(record: str | SQLRecord | None, using: str | None): + """Check if the record is from the target DB.""" if isinstance(record, SQLRecord): - if using_key is not None and using_key != "default": - if record._state.db != using_key: + if using is not None and using != "default": + if record._state.db != using: raise ValueError( - f"record must be a {record.__class__.__get_name_with_module__()} record from instance '{using_key}'!" + f"record must be a {record.__class__.__get_name_with_module__()} record from instance '{using}'!" ) diff --git a/lamindb/models/feature.py b/lamindb/models/feature.py index e5c86d1e1..dd3986061 100644 --- a/lamindb/models/feature.py +++ b/lamindb/models/feature.py @@ -44,13 +44,13 @@ TracksUpdates, ) from .sqlrecord import ( + UNSET, BaseSQLRecord, Branch, HasType, Registry, Space, SQLRecord, - UNSET, _get_record_kwargs, pop_space_branch_kwargs, ) @@ -130,7 +130,7 @@ def parse_dtype(dtype_str: str, check_exists: bool = False) -> list[dict[str, An def transfer_feature_dtypes( - feature: Feature, using_key: str | None, transfer_logs: dict + feature: Feature, using: str | None, transfer_logs: dict ) -> None: from .sqlrecord import transfer_to_default_db @@ -152,7 +152,7 @@ def transfer_feature_dtypes( source_type = registry.objects.using(feature._state.db).get(uid=source_type_uid) source_type_id = source_type.id transferred_type = transfer_to_default_db( - source_type, using_key, transfer_logs=transfer_logs, save=True + source_type, using, transfer_logs=transfer_logs, save=True ) if getattr(source_type, "is_type", False): source_typed_children = source_type.__class__.objects.using( @@ -160,7 +160,7 @@ def transfer_feature_dtypes( ).filter(type_id=source_type_id) for source_record in source_typed_children: transfer_to_default_db( - source_record, using_key, transfer_logs=transfer_logs, save=True + source_record, using, transfer_logs=transfer_logs, save=True ) assert transferred_type is None or transferred_type.uid == source_type_uid, ( "transfer_feature_dtypes() expected UID invariance for dtype type " @@ -172,8 +172,13 @@ def transfer_feature_dtypes( def get_record_type_from_uid( registry: Registry, type_uid: str, + using: str | None = None, ) -> SQLRecord: - type_record: SQLRecord = registry.get(type_uid) + type_record: SQLRecord = ( + registry.get(type_uid) + if using is None + else registry.connect(using).get(type_uid) + ) if type_record.branch_id == -1: warning_msg = f"retrieving {registry.__name__} type '{type_record.name}' (uid='{type_uid}') from trash" @@ -636,13 +641,17 @@ def parse_filter_string(filter_str: str) -> dict[str, tuple[str, str | None, str def resolve_relation_filters( - parsed_filters: dict[str, tuple[str, str | None, str]], registry: SQLRecord + parsed_filters: dict[str, tuple[str, str | None, str]], + registry: SQLRecord, + using: str | None = None, ) -> dict[str, str | SQLRecord]: """Resolve relation filters actual model objects. Args: parsed_filters: Django filters like output from :func:`lamindb.models.feature.parse_filter_string` registry: Model class to resolve relationships against + using: Target instance to resolve related objects on (per-instance PKs); + ``None`` resolves against the default connection unchanged. Returns: Dict with resolved objects for successful relations, original values for direct fields and failed resolutions. @@ -657,7 +666,11 @@ def resolve_relation_filters( and relation_field.field.is_relation ): related_model = relation_field.field.related_model - related_obj = related_model.get(**{field_name: value}) + related_obj = ( + related_model.get(**{field_name: value}) + if using is None + else related_model.connect(using).get(**{field_name: value}) + ) resolved[relation_name] = related_obj else: resolved[filter_key] = value diff --git a/lamindb/models/query_manager.py b/lamindb/models/query_manager.py index b31ecb3b3..2474298bd 100644 --- a/lamindb/models/query_manager.py +++ b/lamindb/models/query_manager.py @@ -176,7 +176,7 @@ def _lookup( cls, field: StrField | None = None, return_field: StrField | None = None, - using_key: str | None = None, + using: str | None = None, keep: Literal["first", "last", False] = "first", ) -> NamedTuple: """Return an auto-complete object for a field. diff --git a/lamindb/models/record.py b/lamindb/models/record.py index 13d50722a..baae2b83f 100644 --- a/lamindb/models/record.py +++ b/lamindb/models/record.py @@ -17,7 +17,7 @@ TextField, ) from lamindb.base.utils import class_and_instance_method, strict_classmethod -from lamindb.errors import FieldValidationError +from lamindb.errors import FieldValidationError, InvalidArgument from ..base.uids import base62_16 from .artifact import Artifact @@ -520,14 +520,41 @@ def _build_records(self) -> list[Record]: records.append(self._cls(name=name, **record_kwargs)) return records - def save(self) -> SQLRecordList[Record]: - """Persist all records and their feature values.""" + def save(self, using: str | None = None) -> SQLRecordList[Record]: + """Persist all records and their feature values. + + Args: + using: Optional slug of a target instance to write the whole batch + (records, scalar features, and multi-valued link rows) to, with + the same semantics as ``ln.save(..., using=...)``. The batch's + record type must already exist on that instance. + """ from .query_set import SQLRecordList from .save import save as ln_save + # primary keys are per-instance; a cross-instance write reuses the + # resolved type's in-memory pk as `type_id`, so the type must already + # live on the target instance (`from_dataframe(type="...")` creates it on + # the default instance). Refuse rather than write a dangling FK. + if using is not None and using != "default": + type_db = self._resolved_type._state.db + if type_db != using: + type_location = ( + "the default instance" + if type_db in (None, "default") + else f"instance '{type_db}'" + ) + raise InvalidArgument( + f"Cannot save this batch to instance '{using}' because its " + f"record type '{self._resolved_type.name}' lives on " + f"{type_location}. Pass a type that exists on '{using}' — " + f"e.g. `type=ln.DB('{using}').Record.get(name=...)` — or " + f"create the type on '{using}' first." + ) + if self._records is None: self._records = self._build_records() - ln_save(self._records) + ln_save(self._records, using=using) return SQLRecordList(self._records) diff --git a/lamindb/models/save.py b/lamindb/models/save.py index 5c44ac6de..4521e23e7 100644 --- a/lamindb/models/save.py +++ b/lamindb/models/save.py @@ -8,7 +8,7 @@ from datetime import datetime from typing import TYPE_CHECKING -from django.db import IntegrityError, transaction +from django.db import IntegrityError, connections, transaction from django.utils.functional import partition from lamin_utils import logger from lamindb_setup.core.upath import LocalPathClasses, UPath @@ -26,16 +26,30 @@ from .artifact import Artifact +def _ensure_using_connection(registry: type[SQLRecord], using: str | None) -> None: + if using is None or using == "default" or using in connections: + return + registry.connect(using) + + +def _prepare_cross_instance_create(record: SQLRecord, using: str | None) -> None: + if using is None or using == "default": + return + if (record._state.adding or record.pk is None) and hasattr(record, "run_id"): + record.run = None + record.run_id = None + + def save( records: Iterable[SQLRecord], ignore_conflicts: bool | None = False, batch_size: int = 10000, + using: str | None = None, ) -> None: """Bulk save objects. - Note: - - This is a much faster than saving objects using `object.save()`. + This is a much faster than saving a list of objects + and repeatedly calling `object.save()`. Warning: @@ -49,6 +63,7 @@ def save( If you need records with ids, you need to query them from the database. batch_size: Number of records to process in each batch. Large batch sizes can improve performance but may lead to memory issues. + using: Optional database slug for a target database that differs from the default database. Examples -------- @@ -81,14 +96,19 @@ def save( # for artifacts, we want to bulk-upload rather than upload one-by-one non_artifacts, artifacts = partition(lambda r: isinstance(r, Artifact), records) if non_artifacts: + for record in non_artifacts: + _prepare_cross_instance_create(record, using) non_artifacts_old, non_artifacts_new = partition( lambda r: r._state.adding or r.pk is None, non_artifacts ) bulk_create( - non_artifacts_new, ignore_conflicts=ignore_conflicts, batch_size=batch_size + non_artifacts_new, + ignore_conflicts=ignore_conflicts, + batch_size=batch_size, + using=using, ) if non_artifacts_old: - bulk_update(non_artifacts_old, batch_size=batch_size) + bulk_update(non_artifacts_old, batch_size=batch_size, using=using) non_artifacts_with_parents = [ r for r in non_artifacts_new if hasattr(r, "_parents") ] @@ -106,9 +126,12 @@ def save( if records_with_lazy_features: from ._feature_manager import bulk_set_features_in_records - bulk_set_features_in_records(records_with_lazy_features) + bulk_set_features_in_records(records_with_lazy_features, using=using) if artifacts: + for record in artifacts: + _prepare_cross_instance_create(record, using) + _ensure_using_connection(Artifact, using) with transaction.atomic(): for record in artifacts: # will switch to True after the successful upload / saving @@ -116,9 +139,8 @@ def save( record, "_to_store", False ): record._storage_ongoing = True - record._save_skip_storage() - using_key = settings._using_key - store_artifacts(artifacts, using_key=using_key) + record._save_skip_storage(using=using) + store_artifacts(artifacts, using=using) # this function returns None as potentially 10k records might be saved # refreshing all of them from the DB would mean a severe performance penalty @@ -130,6 +152,7 @@ def bulk_create( records: Iterable[SQLRecord], ignore_conflicts: bool | None = False, batch_size: int = 10000, + using: str | None = None, ): """Create records in batches for safety and performance. @@ -137,14 +160,19 @@ def bulk_create( records: Iterable of SQLRecord objects to create ignore_conflicts: Whether to ignore conflicts during creation batch_size: Number of records to process in each batch. + using: Optional database alias used for bulk operations. """ records_by_orm = defaultdict(list) for record in records: records_by_orm[record.__class__].append(record) for registry, records_list in records_by_orm.items(): + _ensure_using_connection(registry, using) total_records = len(records_list) model_name = registry.__name__ + manager = ( + registry.objects.using(using) if using is not None else registry.objects + ) if total_records > batch_size: logger.important( f"starting creation of {total_records} {model_name} records in batches of {batch_size}" @@ -161,7 +189,7 @@ def bulk_create( f"processing batch {batch_num}/{total_batches} for {model_name}: {len(batch)} records" ) try: - registry.objects.bulk_create(batch, ignore_conflicts=ignore_conflicts) + manager.bulk_create(batch, ignore_conflicts=ignore_conflicts) # handle unique constraint violations due to non-default branches except IntegrityError as e: error_msg = str(e) @@ -189,7 +217,7 @@ def bulk_create( q_objects |= Q(**field_kwargs) # Query against non-default branches - pre_existing_records_not_main_branch = registry.objects.filter( + pre_existing_records_not_main_branch = manager.filter( q_objects ).exclude(branch_id=1) @@ -206,7 +234,7 @@ def bulk_create( if tuple(getattr(r, field) for field in unique_fields) not in pre_existing_value_tuples ] - save(records_main_branch) + save(records_main_branch, using=using) # Now move the pre-existing records to the main branch if pre_existing_value_tuples: @@ -221,7 +249,10 @@ def bulk_create( in pre_existing_value_tuples ] for record in pre_existing_records_to_move: - record.save() + if using is None: + record.save() + else: + record.save(using=using) else: raise e @@ -231,6 +262,7 @@ def bulk_update( ignore_conflicts: bool | None = False, batch_size: int = 10000, update_fields: list[str] | None = None, + using: str | None = None, ): """Update records in batches for safety and performance. @@ -239,14 +271,19 @@ def bulk_update( ignore_conflicts: Whether to ignore conflicts during update (currently unused but kept for consistency) batch_size: Number of records to process in each batch. If None, processes all at once. update_fields: Specific fields to update. If None, updates all fields except created_at and id. + using: Optional database alias used for bulk operations. """ records_by_orm = defaultdict(list) for record in records: records_by_orm[record.__class__].append(record) for registry, records_list in records_by_orm.items(): + _ensure_using_connection(registry, using) total_records = len(records_list) model_name = registry.__name__ + manager = ( + registry.objects.using(using) if using is not None else registry.objects + ) if total_records > batch_size: logger.warning( f"starting update for {total_records} {model_name} records in batches of {batch_size}" @@ -268,13 +305,13 @@ def bulk_update( logger.info( f"processing batch {batch_num}/{total_batches} for {model_name}: {len(batch)} records" ) - registry.objects.bulk_update(batch, field_names) + manager.bulk_update(batch, field_names) # This is also used within Artifact.save() def check_and_attempt_upload( artifact: Artifact, - using_key: str | None = None, + using: str | None = None, access_token: str | None = None, print_progress: bool = True, **kwargs, @@ -286,7 +323,7 @@ def check_and_attempt_upload( try: storage_path, cache_path = upload_artifact( artifact, - using_key, + using, access_token=access_token, print_progress=print_progress, **kwargs, @@ -383,7 +420,7 @@ def copy_or_move_to_cache( def check_and_attempt_clearing( artifact: Artifact, raise_file_not_found_error: bool = True, - using_key: str | None = None, + using: str | None = None, ) -> Exception | None: # this is a clean-up operation after replace() was called # or if there was an exception during upload @@ -393,11 +430,11 @@ def check_and_attempt_clearing( # avoid root-level import of core.storage module from ..core.storage import paths - delete_msg = paths.delete_storage_using_key( + delete_msg = paths.delete_storage_using( artifact, artifact._clear_storagekey, # type: ignore raise_file_not_found_error=raise_file_not_found_error, - using_key=using_key, + using=using, ) if delete_msg != "did-not-delete": logger.success( @@ -410,9 +447,7 @@ def check_and_attempt_clearing( return None -def store_artifacts( - artifacts: Iterable[Artifact], using_key: str | None = None -) -> None: +def store_artifacts(artifacts: Iterable[Artifact], using: str | None = None) -> None: """Upload artifacts in a list of database-committed artifacts to storage. If any upload fails, subsequent artifacts are cleaned up from the DB. @@ -427,7 +462,7 @@ def store_artifacts( for artifact in artifacts: # failure here sets ._clear_storagekey # for cleanup below - exception = check_and_attempt_upload(artifact, using_key) + exception = check_and_attempt_upload(artifact, using) if exception is not None: break @@ -441,11 +476,11 @@ def store_artifacts( if artifact._storage_ongoing: artifact._storage_ongoing = False # each .save() is a separate transaction below - super(Artifact, artifact).save() + super(Artifact, artifact).save(using=using) # if check_and_attempt_upload was successful # then this can have only ._clear_storagekey from .replace exception = check_and_attempt_clearing( - artifact, raise_file_not_found_error=True, using_key=using_key + artifact, raise_file_not_found_error=True, using=using ) if exception is not None: logger.warning(f"clean up of {artifact._clear_storagekey} failed") # type: ignore @@ -456,10 +491,10 @@ def store_artifacts( with transaction.atomic(): for artifact in artifacts: if artifact not in stored_artifacts: - artifact._delete_skip_storage() + artifact._delete_skip_storage(using=using) # clean up storage after failure in check_and_attempt_upload exception_clear = check_and_attempt_clearing( - artifact, raise_file_not_found_error=False, using_key=using_key + artifact, raise_file_not_found_error=False, using=using ) if exception_clear is not None: logger.warning( @@ -494,7 +529,7 @@ def prepare_error_message(records, stored_artifacts, exception) -> str: def upload_artifact( artifact, - using_key: str | None = None, + using: str | None = None, access_token: str | None = None, print_progress: bool = True, **kwargs, @@ -507,7 +542,7 @@ def upload_artifact( storage_key = paths.auto_storage_key_from_artifact(artifact) storage_path, storage_settings = paths.attempt_accessing_path( - artifact, storage_key, using_key=using_key, access_token=access_token + artifact, storage_key, using=using, access_token=access_token ) if getattr(artifact, "_to_store", False): logger.save(f"storing artifact '{artifact.uid}' at '{storage_path}'") diff --git a/lamindb/models/schema.py b/lamindb/models/schema.py index 42225807c..c3b7d0a09 100644 --- a/lamindb/models/schema.py +++ b/lamindb/models/schema.py @@ -41,6 +41,7 @@ from .query_set import QuerySet, SQLRecordList from .run import TracksRun, TracksUpdates from .sqlrecord import ( + UNSET, BaseSQLRecord, Branch, HasType, @@ -48,7 +49,6 @@ Registry, Space, SQLRecord, - UNSET, _get_record_kwargs, init_self_from_db, pop_space_branch_kwargs, @@ -94,6 +94,15 @@ def validate_features(features: list[SQLRecord]) -> SQLRecord: return next(iter(feature_types)) # return value in set of cardinality 1 +def _resolve_pks_on_instance( + registry: type[SQLRecord], uids: list[str], using: str | None +) -> list[int]: + uid_to_pk = dict( + registry.objects.using(using).filter(uid__in=uids).values_list("uid", "id") + ) + return [uid_to_pk[uid] for uid in uids] + + def get_features_config( features: list[SQLRecord] | tuple[SQLRecord, dict], ) -> tuple[list[SQLRecord], list[tuple[SQLRecord, dict]]]: @@ -116,7 +125,7 @@ def transfer_schema_members( schema: Schema, source_db: str, source_pk: int | None, - using_key: str | None, + using: str | None, *, transfer_logs: dict, ) -> None: @@ -144,11 +153,9 @@ def transfer_schema_members( ) if source_index_feature is not None: index_feature = copy(source_index_feature) - transfer_feature_dtypes( - index_feature, using_key, transfer_logs=transfer_logs - ) + transfer_feature_dtypes(index_feature, using, transfer_logs=transfer_logs) transferred_index_feature = transfer_to_default_db( - index_feature, using_key, transfer_logs=transfer_logs, save=True + index_feature, using, transfer_logs=transfer_logs, save=True ) transferred_index_uid = ( transferred_index_feature.uid @@ -169,9 +176,9 @@ def transfer_schema_members( for source_member in members: member = copy(source_member) if isinstance(member, Feature): - transfer_feature_dtypes(member, using_key, transfer_logs=transfer_logs) + transfer_feature_dtypes(member, using, transfer_logs=transfer_logs) transferred_member = transfer_to_default_db( - member, using_key, transfer_logs=transfer_logs, save=True + member, using, transfer_logs=transfer_logs, save=True ) if transferred_member is not None: transferred_members.append(transferred_member) @@ -206,7 +213,7 @@ def transfer_schema_members( def transfer_schema_with_members( - schema: Schema, using_key: str | None, *, transfer_logs: dict + schema: Schema, using: str | None, *, transfer_logs: dict ) -> Schema: from copy import copy @@ -217,7 +224,7 @@ def transfer_schema_with_members( schema_copy = copy(schema) with transaction.atomic(): record_on_default = transfer_to_default_db( - schema_copy, using_key, transfer_logs=transfer_logs, save=True + schema_copy, using, transfer_logs=transfer_logs, save=True ) schema_default = ( record_on_default if record_on_default is not None else schema_copy @@ -226,7 +233,7 @@ def transfer_schema_with_members( schema_default, source_db, source_pk, - using_key, + using, transfer_logs=transfer_logs, ) return schema_default @@ -1094,6 +1101,25 @@ def from_df( ) -> Schema | None: return cls.from_dataframe(df, field, name, mute, organism, source) + def _infer_members_instance(self) -> str | None: + candidates: set[str] = set() + if hasattr(self, "_features"): + for record in self._features[1]: + db = record._state.db + if db is not None and db != "default": + candidates.add(db) + if hasattr(self, "_slots"): + for component in self._slots.values(): + db = component._state.db + if db is not None and db != "default": + candidates.add(db) + if len(candidates) > 1: + raise InvalidArgument( + f"schema members live on multiple instances ({sorted(candidates)}); " + "a schema must be saved to a single instance." + ) + return next(iter(candidates)) if candidates else None + def save(self, *args, **kwargs) -> Schema: """Save schema. @@ -1122,6 +1148,18 @@ def save(self, *args, **kwargs) -> Schema: index_name_conflict = kwargs.pop("index_name_conflict", None) using = kwargs.get("using") or self._state.db + members_instance = self._infer_members_instance() + if members_instance is not None: + if using is None or using == "default": + using = members_instance + kwargs["using"] = members_instance + elif using != members_instance: + raise InvalidArgument( + f"schema is being saved to '{using}' but its members live on " + f"'{members_instance}'; save the schema to the same instance as " + "its members." + ) + if self.pk is not None: existing_features = self.members.to_list() if self.members.exists() else [] if hasattr(self, "_features"): @@ -1192,14 +1230,27 @@ def save(self, *args, **kwargs) -> Schema: index_name_conflict=index_name_conflict, ) super().save(*args, **kwargs) + cross_instance = using is not None and using != "default" + schema_id = ( + _resolve_pks_on_instance(type(self), [self.uid], using)[0] + if cross_instance + else self.id + ) if hasattr(self, "_slots"): # analogous to save_schema_links in core._data.py # which is called to save feature sets in artifact.save() + slots = list(self._slots.items()) + if cross_instance: + component_ids = _resolve_pks_on_instance( + Schema, [component.uid for _, component in slots], using + ) + else: + component_ids = [component.id for _, component in slots] links = [] - for slot, component in self._slots.items(): + for (slot, _), component_id in zip(slots, component_ids): kwargs = { - "composite_id": self.id, - "component_id": component.id, + "composite_id": schema_id, + "component_id": component_id, "slot": slot, } links.append(Schema.components.through(**kwargs)) @@ -1224,18 +1275,25 @@ def save(self, *args, **kwargs) -> Schema: else: related_field = related_model_split[1].lower() related_field_id = f"{related_field}_id" - new_member_ids = [record.id for record in records] + if cross_instance: + new_member_ids = _resolve_pks_on_instance( + records[0].__class__, [record.uid for record in records], using + ) + else: + new_member_ids = [record.id for record in records] existing_member_ids = list( through_model.objects.using(using) - .filter(schema_id=self.id) + .filter(schema_id=schema_id) .order_by("id") .values_list(related_field_id, flat=True) ) if new_member_ids != existing_member_ids: - through_model.objects.using(using).filter(schema_id=self.id).delete() + through_model.objects.using(using).filter(schema_id=schema_id).delete() links = [ - through_model(**{"schema_id": self.id, related_field_id: record.id}) - for record in records + through_model( + **{"schema_id": schema_id, related_field_id: member_id} + ) + for member_id in new_member_ids ] through_model.objects.using(using).bulk_create(links) delattr(self, "_features") diff --git a/lamindb/models/sqlrecord.py b/lamindb/models/sqlrecord.py index 50c6fe429..a82233383 100644 --- a/lamindb/models/sqlrecord.py +++ b/lamindb/models/sqlrecord.py @@ -828,11 +828,11 @@ def filter(cls, *queries, **expressions) -> QuerySet: """ from .query_set import QuerySet - _using_key = None - if "_using_key" in expressions: - _using_key = expressions.pop("_using_key") + using = None + if "using" in expressions: + using = expressions.pop("using") - return QuerySet(model=cls, using=_using_key).filter(*queries, **expressions) + return QuerySet(model=cls, using=using).filter(*queries, **expressions) def get( cls: type[T], @@ -1376,11 +1376,23 @@ def _field_changed(self, field_name: str, check_is_saved: bool = True) -> bool: def save(self: T, *args, **kwargs) -> T: """Save. - Always saves to the default database. + Args: + using: Optional database slug for a target database that differs from the default database. """ - using_key = None + using = None if "using" in kwargs: - using_key = kwargs["using"] + using = kwargs["using"] + if using != "default" and using not in connections: + self.__class__.connect(using) + # cross-instance writes can inherit a run from the active local context; + # that run id does not exist in the target instance. + if ( + using != "default" + and (self._state.adding or self.pk is None) + and hasattr(self, "run_id") + ): + self.run = None + self.run_id = None transfer_config = kwargs.pop("transfer", None) db = self._state.db pk_on_db = self.pk @@ -1389,7 +1401,7 @@ def save(self: T, *args, **kwargs) -> T: and transfer_config is None and db is not None and db != "default" - and using_key is None + and using is None ): transfer_config = "annotations" artifacts: list = [] @@ -1403,14 +1415,14 @@ def save(self: T, *args, **kwargs) -> T: "transferred": [], "run": None, } - if db is not None and db != "default" and using_key is None: + if db is not None and db != "default" and using is None: if isinstance(self, IsVersioned): if not self.is_latest: raise NotImplementedError( "You are attempting to transfer a record that's not the latest in its version history. This is currently not supported." ) pre_existing_record = transfer_to_default_db( - self, using_key, transfer_logs=transfer_logs + self, using, transfer_logs=transfer_logs ) self._revises: IsVersioned if pre_existing_record is not None: @@ -1584,7 +1596,7 @@ def save(self: T, *args, **kwargs) -> T: track_current_name_value(self) # perform transfer of many-to-many fields # only supported for Artifact and Collection records - if db is not None and db != "default" and using_key is None: + if db is not None and db != "default" and using is None: if self.__class__.__name__ == "Collection": if len(artifacts) > 0: logger.info("transfer artifacts") @@ -1595,7 +1607,7 @@ def save(self: T, *args, **kwargs) -> T: from .schema import transfer_schema_members transfer_schema_members( - self, db, pk_on_db, using_key, transfer_logs=transfer_logs + self, db, pk_on_db, using, transfer_logs=transfer_logs ) if hasattr(self, "labels") and transfer_config == "annotations": from copy import copy @@ -2272,7 +2284,7 @@ def delete(self, permanent: bool | None = None, **kwargs): from .artifact import delete_permanently delete_permanently( - self, storage=kwargs["storage"], using_key=kwargs["using_key"] + self, storage=kwargs["storage"], using=kwargs["using"] ) return None return super().delete() @@ -2417,6 +2429,40 @@ def get_name_field( return field +class LaminDBRouter: + def allow_relation(self, obj1, obj2, **hints) -> bool | None: + # Permit assigning a relation when one side is not yet saved, e.g. + # `ln.Record(type=remote_type)` while building a record for a + # cross-instance `using=` save: the new record's `_state.db` may have + # been pinned to "default" by its default branch/space FKs before the + # cross-instance type FK is assigned, which Django's default + # (same-database) check would otherwise reject at construction time. + # + # For two *already-saved* records we defer to Django's default check + # (return None) so that genuinely cross-instance links, such as + # `artifact.records.add(label)` across two different instances, remain + # blocked with the usual friendly error. + if isinstance(obj1, SQLRecord) and isinstance(obj2, SQLRecord): + if obj1._state.adding or obj2._state.adding: + return True + return None + + +def _ensure_lamindb_router() -> None: + from django.conf import settings as django_settings + + if not django_settings.configured: + return + from django.db import router as django_router + + router_path = "lamindb.models.sqlrecord.LaminDBRouter" + existing = list(getattr(django_settings, "DATABASE_ROUTERS", [])) + if router_path in existing or any(isinstance(r, LaminDBRouter) for r in existing): + return + django_settings.DATABASE_ROUTERS = [*existing, router_path] + django_router.__dict__.pop("routers", None) + + def add_db_connection(db: str, using: str): db_config = dj_database_url.config( default=db, conn_max_age=600, conn_health_checks=True @@ -2425,6 +2471,10 @@ def add_db_connection(db: str, using: str): db_config["OPTIONS"] = {} db_config["AUTOCOMMIT"] = True connections.settings[using] = db_config + # Ensure router registration when dynamic DB aliases are added. + # Covers init orders where the import-time call happened before + # Django settings were configured and therefore returned early. + _ensure_lamindb_router() REGISTRY_UNIQUE_FIELD = {"storage": "root", "ulabel": "name"} @@ -2433,7 +2483,7 @@ def add_db_connection(db: str, using: str): def update_fk_to_default_db( records: SQLRecord | list[SQLRecord] | DjangoQuerySet, fk: str, - using_key: str | None, + using: str | None, transfer_logs: dict, ): # here in case it is an iterable, we are checking only a single record @@ -2464,11 +2514,11 @@ def update_fk_to_default_db( from .schema import transfer_schema_with_members fk_record_default = transfer_schema_with_members( - fk_record_default, using_key, transfer_logs=transfer_logs + fk_record_default, using, transfer_logs=transfer_logs ) elif pre_existing_fk_record_default is None: transfer_to_default_db( - fk_record_default, using_key, save=True, transfer_logs=transfer_logs + fk_record_default, using, save=True, transfer_logs=transfer_logs ) else: fk_record_default = pre_existing_fk_record_default @@ -2490,10 +2540,10 @@ def update_fk_to_default_db( def transfer_fk_to_default_db_bulk( - records: list | DjangoQuerySet, using_key: str | None, transfer_logs: dict + records: list | DjangoQuerySet, using: str | None, transfer_logs: dict ): for fk in FKBULK: - update_fk_to_default_db(records, fk, using_key, transfer_logs=transfer_logs) + update_fk_to_default_db(records, fk, using, transfer_logs=transfer_logs) def get_transfer_run(record) -> Run: @@ -2543,7 +2593,7 @@ def get_transfer_run(record) -> Run: def transfer_to_default_db( record: SQLRecord, - using_key: str | None, + using: str | None, *, transfer_logs: dict, save: bool = False, @@ -2586,7 +2636,7 @@ def transfer_to_default_db( # don't transfer fk fields that are already bulk transferred fk_fields = [fk for fk in fk_fields if fk not in FKBULK] for fk in fk_fields: - update_fk_to_default_db(record, fk, using_key, transfer_logs=transfer_logs) + update_fk_to_default_db(record, fk, using, transfer_logs=transfer_logs) # FK ids were remapped to the default DB; drop tracked *_id originals so save # logic does not treat remapping as a user-requested field change. if (original_values := getattr(record, "_original_values", None)) is not None: diff --git a/lamindb/models/transform.py b/lamindb/models/transform.py index 1bdd8e953..4d49e10f1 100644 --- a/lamindb/models/transform.py +++ b/lamindb/models/transform.py @@ -303,7 +303,7 @@ def __init__( reference_type: str | None = kwargs.pop("reference_type", None) space_branch_kwargs = pop_space_branch_kwargs(kwargs) skip_hash_lookup: bool = kwargs.pop("skip_hash_lookup", False) - using_key = kwargs.pop("using_key", None) + using = kwargs.pop("using", None) # below is internal use that we'll hopefully be able to eliminate uid: str | None = kwargs.pop("uid") if "uid" in kwargs else None source_code: str | None = ( @@ -328,14 +328,14 @@ def __init__( # need to check uid before checking key if uid is not None: revises = ( - Transform.objects.using(using_key) + Transform.objects.using(using) .filter(uid__startswith=uid[:-4], is_latest=True) .order_by("-created_at") .first() ) elif key is not None: candidates_for_revises = ( - Transform.objects.using(using_key) + Transform.objects.using(using) .filter(~Q(branch_id=-1), key=key, is_latest=True) .order_by("-created_at") ) diff --git a/sub/lamindb-setup b/sub/lamindb-setup index e4e5d15ca..d71b99b48 160000 --- a/sub/lamindb-setup +++ b/sub/lamindb-setup @@ -1 +1 @@ -Subproject commit e4e5d15ca70869426be78203b2419edd92facac0 +Subproject commit d71b99b48db488dd482944810a8a622ffdc33928 diff --git a/tests/pydata/test_save.py b/tests/pydata/test_save.py index b58ee0e7d..3151f3451 100644 --- a/tests/pydata/test_save.py +++ b/tests/pydata/test_save.py @@ -5,7 +5,10 @@ from _dataset_fixtures import ( # noqa get_mini_csv, ) -from lamindb.models.save import prepare_error_message, store_artifacts +from lamindb.models.save import ( + prepare_error_message, + store_artifacts, +) def test_bulk_save_and_update(): @@ -49,7 +52,7 @@ def test_store_artifacts_acid(get_mini_csv): artifact.save() with pytest.raises(RuntimeError) as error: - store_artifacts([artifact], using_key=None) + store_artifacts([artifact], using=None) assert str(error.exconly()).startswith( "RuntimeError: The following entries have been successfully uploaded" ) diff --git a/tests/storage/test_streaming.py b/tests/storage/test_streaming.py index 159949f9a..1813f3045 100644 --- a/tests/storage/test_streaming.py +++ b/tests/storage/test_streaming.py @@ -82,13 +82,13 @@ def test_backed_access(adata_format): del store with pytest.raises(ValueError): - access = backed_access(fp.with_suffix(".invalid_suffix"), using_key=None) + access = backed_access(fp.with_suffix(".invalid_suffix"), using=None) # can't open anndata in write mode with pytest.raises(ValueError): - access = backed_access(fp, mode="a", using_key=None) + access = backed_access(fp, mode="a", using=None) - access = backed_access(fp, using_key=None) + access = backed_access(fp, using=None) assert not access.closed assert isinstance(access.obs_names, pd.Index) @@ -142,14 +142,14 @@ def test_backed_access(adata_format): assert access.closed del access - with backed_access(fp, using_key=None) as access: + with backed_access(fp, using=None) as access: assert not access.closed sub = access[:10] assert sub[:5].shape == (5, 200) assert sub.layers["test"].shape == sub.shape assert access.closed - with backed_access(fp, using_key=None) as access: + with backed_access(fp, using=None) as access: idx = np.array([3, 1, 2]) assert access[:, idx].to_memory().shape == (30, 3) assert access[idx].to_memory().shape == (3, 200) @@ -231,7 +231,7 @@ def test_write_to_disk(): def test_backed_bad_format(bad_adata_path): - access = backed_access(bad_adata_path, using_key=None) + access = backed_access(bad_adata_path, using=None) assert access.obsp["test"].to_memory().sum() == 30 diff --git a/tests/transfer/conftest.py b/tests/transfer/conftest.py index fa247c8b6..fdfd6bc87 100644 --- a/tests/transfer/conftest.py +++ b/tests/transfer/conftest.py @@ -4,10 +4,17 @@ import pytest +def _close_all_connections(): + from django.db import connections + + connections.close_all() + + @pytest.fixture(scope="session", autouse=True) def setup_testdb1(): ln.setup.init(storage="./testdb1") yield + _close_all_connections() shutil.rmtree("./testdb1") ln.setup.delete("testdb1", force=True) @@ -16,6 +23,7 @@ def setup_testdb1(): def setup_testdb2(): ln.setup.init(storage="./testdb2") yield + _close_all_connections() shutil.rmtree("./testdb2") ln.setup.delete("testdb2", force=True) diff --git a/tests/transfer/test_save_to_another_db.py b/tests/transfer/test_save_to_another_db.py new file mode 100644 index 000000000..53b6fd070 --- /dev/null +++ b/tests/transfer/test_save_to_another_db.py @@ -0,0 +1,542 @@ +import subprocess +import sys + +import lamindb as ln +import pytest + + +def test_cross_instance_m2m_add_still_blocked(): + # the cross-instance relation router must NOT let `.add()` link two + # already-saved records that live on different instances; that guardrail + # (a friendly ValueError) must still fire. Mirrors the network-dependent + # tests/pydata test_unsaved_model_different_instance, but runs locally. + using = f"{ln.setup.settings.user.handle}/testdb1" + db2 = ln.DB(using) + af = db2.Artifact.get(key="README.md") # lives on testdb1 (see conftest) + assert af._state.db == using + + new_label = ln.Record(name="guardrail-testlabel").save() # on default testdb2 + try: + with pytest.raises(ValueError) as excinfo: + af.records.add(new_label) + assert "Cannot label a record from instance" in str(excinfo.value) + assert using in str(excinfo.value) + finally: + new_label.delete(permanent=True) + + +def test_lamindb_router_registered_on_fresh_import(): + code = ( + "import lamindb as ln\n" + "from django.conf import settings\n" + "print('LaminDBRouter' in str(settings.DATABASE_ROUTERS))\n" + ) + result = subprocess.run( + [sys.executable, "-c", code], capture_output=True, text=True + ) + assert "True" in result.stdout, result.stdout + result.stderr + + +def test_construct_relation_to_remote_type_before_any_save(): + using = f"{ln.setup.settings.user.handle}/testdb1" + db2 = ln.DB(using) + type_name = "router-remote-type" + + if db2.Record.filter(name=type_name).exists(): + db2.Record.filter(name=type_name).delete(permanent=True) + try: + remote_type = ln.Record(name=type_name, is_type=True).save(using=using) + # read it back so the FK target is a record whose _state.db is the alias + remote_type = db2.Record.get(name=type_name) + # construction with a cross-instance FK must not raise + child = ln.Record(name="router-remote-child", type=remote_type) + assert child.type.uid == remote_type.uid + finally: + db2.Record.filter(name=type_name).delete(permanent=True) + + +def test_save_ulabel_to_another_db_via_model_save(): + assert ln.setup.settings.instance.name == "testdb2" + + using = f"{ln.setup.settings.user.handle}/testdb1" + db2 = ln.DB(using) + name = "save-using-single" + remote_labels = db2.ULabel.filter(name=name) + if remote_labels.exists(): + remote_labels.delete(permanent=True) + + try: + ulabel = ln.ULabel(name=name).save(using=using) + + assert ulabel._state.db == using + assert db2.ULabel.get(name=name) + assert ln.ULabel.filter(name=name).count() == 0 + finally: + db2.ULabel.filter(name=name).delete(permanent=True) + + +def test_save_ulabels_to_another_db_via_ln_save(): + assert ln.setup.settings.instance.name == "testdb2" + + using = f"{ln.setup.settings.user.handle}/testdb1" + db2 = ln.DB(using) + names = [f"save-using-bulk-{i}" for i in range(3)] + remote_labels = db2.ULabel.filter(name__in=names) + if remote_labels.exists(): + remote_labels.delete(permanent=True) + + try: + ulabels = [ln.ULabel(name=name) for name in names] + ln.save(ulabels, using=using) + + remote_matches = db2.ULabel.filter(name__in=names) + assert remote_matches.count() == len(names) + assert ln.ULabel.filter(name__in=names).count() == 0 + finally: + db2.ULabel.filter(name__in=names).delete(permanent=True) + + +def test_save_record_with_relation_to_another_db_via_model_save(): + assert ln.setup.settings.instance.name == "testdb2" + + using = f"{ln.setup.settings.user.handle}/testdb1" + db2 = ln.DB(using) + type_name = "save-using-relation-type" + rec_name = "save-using-relation-child" + + for qs in (db2.Record.filter(name=rec_name), db2.Record.filter(name=type_name)): + if qs.exists(): + qs.delete(permanent=True) + + try: + rec_type = ln.Record(name=type_name, is_type=True).save(using=using) + child = ln.Record(name=rec_name, type=rec_type).save(using=using) + + assert child._state.db == using + assert db2.Record.get(name=rec_name) + assert ln.Record.filter(name=rec_name).count() == 0 + finally: + for qs in (db2.Record.filter(name=rec_name), db2.Record.filter(name=type_name)): + qs.delete(permanent=True) + +def test_save_record_with_multivalued_relation_to_another_db_via_model_save(): + assert ln.setup.settings.instance.name == "testdb2" + + using = f"{ln.setup.settings.user.handle}/testdb1" + db2 = ln.DB(using) + + type_name = "mv-intervention-type" + gene_type_name = "mv-gene-type" + gene_feat_name = "mv-genes" + gene_names = [f"mv-gene-{i}" for i in range(5)] + rec_name = "mv-child" + + def _clean(): + for qs in ( + db2.Record.filter(name=rec_name), + db2.Record.filter(name__in=gene_names), + db2.Record.filter(name=type_name), + db2.Record.filter(name=gene_type_name), + db2.Feature.filter(name=gene_feat_name), + ): + if qs.exists(): + qs.delete(permanent=True) + + _clean() + try: + intervention_t = ln.Record(name=type_name, is_type=True).save(using=using) + gene_t = ln.Record(name=gene_type_name, is_type=True).save(using=using) + gene_feat = ln.Feature(name=gene_feat_name, dtype=gene_t).save(using=using) + gene_pool = [ln.Record(name=g, type=gene_t).save(using=using) for g in gene_names] + + child = ln.Record( + name=rec_name, + type=intervention_t, + features={gene_feat: gene_pool[:3]}, + ).save(using=using) + + assert child._state.db == using + assert db2.Record.get(name=rec_name) + assert ln.Record.filter(name=rec_name).count() == 0 + finally: + _clean() + + +def test_save_schema_with_features_to_another_db_via_model_save(): + assert ln.setup.settings.instance.name == "testdb2" + + using = f"{ln.setup.settings.user.handle}/testdb1" + db2 = ln.DB(using) + + feat_names = [f"save-using-schema-feat-{i}" for i in range(3)] + schema_name = "save-using-schema" + + def _clean(): + for qs in ( + db2.Schema.filter(name=schema_name), + db2.Feature.filter(name__in=feat_names), + ): + if qs.exists(): + qs.delete(permanent=True) + + _clean() + try: + features = [ + ln.Feature(name=name, dtype=str).save(using=using) for name in feat_names + ] + schema = ln.Schema(name=schema_name, features=features).save(using=using) + + assert schema._state.db == using + remote_schema = db2.Schema.get(name=schema_name) + assert set(remote_schema.members.to_list("name")) == set(feat_names) + assert ln.Schema.filter(name=schema_name).count() == 0 + assert ln.Feature.filter(name__in=feat_names).count() == 0 + finally: + _clean() + + +def test_save_schema_with_remote_features_infers_instance_via_bare_save(): + assert ln.setup.settings.instance.name == "testdb2" + + using = f"{ln.setup.settings.user.handle}/testdb1" + db2 = ln.DB(using) + + feat_names = [f"save-using-schema-infer-feat-{i}" for i in range(3)] + schema_name = "save-using-schema-infer" + + def _clean(): + for qs in ( + db2.Schema.filter(name=schema_name), + db2.Feature.filter(name__in=feat_names), + ): + if qs.exists(): + qs.delete(permanent=True) + + _clean() + try: + features = [ + ln.Feature(name=name, dtype=str).save(using=using) for name in feat_names + ] + # bare save: instance is inferred from the (remote) member features + schema = ln.Schema(name=schema_name, features=features).save() + + assert schema._state.db == using + remote_schema = db2.Schema.get(name=schema_name) + assert set(remote_schema.members.to_list("name")) == set(feat_names) + assert ln.Schema.filter(name=schema_name).count() == 0 + assert ln.Feature.filter(name__in=feat_names).count() == 0 + finally: + _clean() + + +def test_save_record_with_single_valued_feature_to_another_db_via_model_save(): + assert ln.setup.settings.instance.name == "testdb2" + + using = f"{ln.setup.settings.user.handle}/testdb1" + db2 = ln.DB(using) + + gene_type_name = "sv-gene-type" + gene_feat_name = "sv-gene" + gene_name = "sv-gene-value" + rec_name = "sv-child" + + def _clean(): + for qs in ( + db2.Record.filter(name=rec_name), + db2.Record.filter(name=gene_name), + db2.Record.filter(name=gene_type_name), + db2.Feature.filter(name=gene_feat_name), + ): + if qs.exists(): + qs.delete(permanent=True) + + _clean() + try: + gene_t = ln.Record(name=gene_type_name, is_type=True).save(using=using) + gene_feat = ln.Feature(name=gene_feat_name, dtype=gene_t).save(using=using) + gene = ln.Record(name=gene_name, type=gene_t).save(using=using) + + child = ln.Record(name=rec_name, features={gene_feat: gene}).save(using=using) + + assert child._state.db == using + remote_child = db2.Record.get(name=rec_name) + assert remote_child.features.get_values()[gene_feat_name] == gene_name + assert ln.Record.filter(name=rec_name).count() == 0 + finally: + _clean() + + +def test_bulk_save_records_with_multivalued_features_to_another_db(): + assert ln.setup.settings.instance.name == "testdb2" + + using = f"{ln.setup.settings.user.handle}/testdb1" + db2 = ln.DB(using) + + gene_type_name = "bulk-mv-gene-type" + int_type_name = "bulk-mv-int-type" + gene_feat_name = "bulk-mv-genes" + schema_name = "bulk-mv-schema" + gene_names = [f"bulk-mv-gene-{i}" for i in range(5)] + rec_names = [f"bulk-mv-child-{i}" for i in range(3)] + + def _clean(): + for qs in ( + db2.Record.filter(name__in=rec_names), + db2.Record.filter(name__in=gene_names), + db2.Record.filter(name=int_type_name), + db2.Schema.filter(name=schema_name), + db2.Feature.filter(name=gene_feat_name), + db2.Record.filter(name=gene_type_name), + ): + if qs.exists(): + qs.delete(permanent=True) + + _clean() + try: + gene_t = ln.Record(name=gene_type_name, is_type=True).save(using=using) + gene_feat = ln.Feature(name=gene_feat_name, dtype=gene_t).save(using=using) + schema = ln.Schema(name=schema_name, features=[gene_feat]).save(using=using) + int_t = ln.Record( + name=int_type_name, is_type=True, schema=schema + ).save(using=using) + gene_pool = [ + ln.Record(name=g, type=gene_t).save(using=using) for g in gene_names + ] + + records = [ + ln.Record(name=rn, type=int_t, features={gene_feat: gene_pool[:3]}) + for rn in rec_names + ] + ln.save(records, using=using) + + for rn in rec_names: + remote = db2.Record.get(name=rn) + assert set(remote.features.get_values()[gene_feat_name]) == set( + gene_names[:3] + ) + assert ln.Record.filter(name__in=rec_names).count() == 0 + finally: + _clean() + + +def test_read_features_from_another_db_via_to_dataframe(): + # read-path gap: `.to_dataframe(include="features")` on a remote queryset must + # resolve the record type on the queryset's instance, not the default. + assert ln.setup.settings.instance.name == "testdb2" + + using = f"{ln.setup.settings.user.handle}/testdb1" + db2 = ln.DB(using) + + gene_type_name = "rd-gene-type" + int_type_name = "rd-int-type" + gene_feat_name = "rd-genes" + schema_name = "rd-schema" + gene_names = [f"rd-gene-{i}" for i in range(4)] + rec_name = "rd-child" + + def _clean(): + for qs in ( + db2.Record.filter(name=rec_name), + db2.Record.filter(name__in=gene_names), + db2.Record.filter(name=int_type_name), + db2.Schema.filter(name=schema_name), + db2.Feature.filter(name=gene_feat_name), + db2.Record.filter(name=gene_type_name), + ): + if qs.exists(): + qs.delete(permanent=True) + + _clean() + try: + gene_t = ln.Record(name=gene_type_name, is_type=True).save(using=using) + gene_feat = ln.Feature(name=gene_feat_name, dtype=gene_t).save(using=using) + schema = ln.Schema(name=schema_name, features=[gene_feat]).save(using=using) + int_t = ln.Record( + name=int_type_name, is_type=True, schema=schema + ).save(using=using) + gene_pool = [ + ln.Record(name=g, type=gene_t).save(using=using) for g in gene_names + ] + ln.Record( + name=rec_name, type=int_t, features={gene_feat: gene_pool[:3]} + ).save(using=using) + + # cross-instance read with feature reassembly must not raise + df = db2.Record.filter(name=rec_name).to_dataframe(include="features") + assert len(df) == 1 + assert gene_feat_name in df.columns + finally: + _clean() + + +def test_from_dataframe_bulk_save_to_another_db(): + # RecordBatch.save(using=X) must land records, scalar features, AND + # multi-valued (relational) link rows on the target instance. + # + # REQUIRED: validate on Postgres before merging. SQLite aligns per-instance + # id sequences, so a cross-instance FK bug (a link row referencing a + # default-instance pk) is masked here and only reproduces on a Postgres + # testdb pair. A green run here proves routing/placement ONLY, not + # FK-divergence safety — do not treat the SQLite pass as sufficient. + import pandas as pd + + assert ln.setup.settings.instance.name == "testdb2" + + using = f"{ln.setup.settings.user.handle}/testdb1" + db2 = ln.DB(using) + + gene_type_name = "fdf-gene-type" + sheet_name = "fdf-sheet" + gene_feat_name = "fdf-genes" + score_feat_name = "fdf-score" + schema_name = "fdf-schema" + gene_names = [f"fdf-gene-{i}" for i in range(5)] + rec_names = [f"fdf-child-{i}" for i in range(3)] + # keep every list multi-element: a single-element multi-valued cell + # round-trips as a scalar via get_values(), which is orthogonal to `using=` + gene_lists = [gene_names[:3], gene_names[1:3], gene_names[2:4]] + + def _clean(): + for qs in ( + db2.Record.filter(name__in=rec_names), + db2.Record.filter(name__in=gene_names), + db2.Record.filter(name=sheet_name), + db2.Schema.filter(name=schema_name), + db2.Feature.filter(name__in=[gene_feat_name, score_feat_name]), + db2.Record.filter(name=gene_type_name), + ): + if qs.exists(): + qs.delete(permanent=True) + + _clean() + try: + gene_t = ln.Record(name=gene_type_name, is_type=True).save(using=using) + gene_feat = ln.Feature(name=gene_feat_name, dtype=gene_t).save(using=using) + score_feat = ln.Feature(name=score_feat_name, dtype=float).save(using=using) + schema = ln.Schema( + name=schema_name, features=[gene_feat, score_feat] + ).save(using=using) + sheet = ln.Record( + name=sheet_name, is_type=True, schema=schema + ).save(using=using) + [ln.Record(name=g, type=gene_t).save(using=using) for g in gene_names] + + df = pd.DataFrame( + { + "__lamindb_record_name__": rec_names, + score_feat_name: [1.0, 2.0, 3.0], + gene_feat_name: gene_lists, + } + ) + + batch = ln.Record.from_dataframe(df, type=sheet) + result = batch.save(using=using) + + # (a) records land on the target instance, nothing local + assert all(r._state.db == using for r in result) + assert db2.Record.filter(name__in=rec_names).count() == len(rec_names) + assert ln.Record.filter(name__in=rec_names).count() == 0 + + # (b) scalar feature values are correct on the target + for rn, score in zip(rec_names, [1.0, 2.0, 3.0]): + remote = db2.Record.get(name=rn) + assert remote.features.get_values()[score_feat_name] == score + + # (c) multi-valued link rows attach on the target + for rn, genes in zip(rec_names, gene_lists): + remote = db2.Record.get(name=rn) + assert set(remote.features.get_values()[gene_feat_name]) == set(genes) + finally: + _clean() + + +def test_from_dataframe_bulk_save_same_instance_unchanged(): + # (d) using=None must behave exactly like today on the default instance, + # including the multi-valued path. + import pandas as pd + + gene_type_name = "fdf-local-gene-type" + sheet_name = "fdf-local-sheet" + gene_feat_name = "fdf-local-genes" + score_feat_name = "fdf-local-score" + schema_name = "fdf-local-schema" + gene_names = [f"fdf-local-gene-{i}" for i in range(4)] + rec_names = [f"fdf-local-child-{i}" for i in range(2)] + gene_lists = [gene_names[:2], gene_names[1:3]] + + def _clean(): + ln.Record.filter(name__in=rec_names).delete(permanent=True) + ln.Record.filter(name__in=gene_names).delete(permanent=True) + ln.Record.filter(name=sheet_name).delete(permanent=True) + ln.Schema.filter(name=schema_name).delete(permanent=True) + ln.Feature.filter(name__in=[gene_feat_name, score_feat_name]).delete( + permanent=True + ) + ln.Record.filter(name=gene_type_name).delete(permanent=True) + + _clean() + try: + gene_t = ln.Record(name=gene_type_name, is_type=True).save() + gene_feat = ln.Feature(name=gene_feat_name, dtype=gene_t).save() + score_feat = ln.Feature(name=score_feat_name, dtype=float).save() + schema = ln.Schema(name=schema_name, features=[gene_feat, score_feat]).save() + sheet = ln.Record(name=sheet_name, is_type=True, schema=schema).save() + [ln.Record(name=g, type=gene_t).save() for g in gene_names] + + df = pd.DataFrame( + { + "__lamindb_record_name__": rec_names, + score_feat_name: [1.0, 2.0], + gene_feat_name: gene_lists, + } + ) + result = ln.Record.from_dataframe(df, type=sheet).save() + assert len(result) == len(rec_names) + for rn, score, genes in zip(rec_names, [1.0, 2.0], gene_lists): + got = ln.Record.get(name=rn).features.get_values() + assert got[score_feat_name] == score + assert set(got[gene_feat_name]) == set(genes) + finally: + _clean() + + +def test_from_dataframe_bulk_save_using_type_not_on_target_raises(): + # guard: the batch's record type must already exist on the target instance; + # a type resolved on the default instance must raise, not write a dangling FK. + import pandas as pd + + using = f"{ln.setup.settings.user.handle}/testdb1" + + score_feat_name = "fdf-guard-score" + schema_name = "fdf-guard-schema" + sheet_name = "fdf-guard-sheet" + rec_names = ["fdf-guard-a", "fdf-guard-b"] + + def _clean(): + ln.Record.filter(name__in=rec_names).delete(permanent=True) + ln.Record.filter(name=sheet_name).delete(permanent=True) + ln.Schema.filter(name=schema_name).delete(permanent=True) + ln.Feature.filter(name=score_feat_name).delete(permanent=True) + + _clean() + try: + # type + schema created on the DEFAULT instance (testdb2) + score_feat = ln.Feature(name=score_feat_name, dtype=float).save() + schema = ln.Schema(name=schema_name, features=[score_feat]).save() + sheet = ln.Record(name=sheet_name, is_type=True, schema=schema).save() + + df = pd.DataFrame( + { + "__lamindb_record_name__": rec_names, + score_feat_name: [1.0, 2.0], + } + ) + batch = ln.Record.from_dataframe(df, type=sheet) + with pytest.raises(ln.errors.InvalidArgument) as excinfo: + batch.save(using=using) + assert "record type" in str(excinfo.value) + assert using in str(excinfo.value) + # nothing should have been written to the target + assert ln.DB(using).Record.filter(name__in=rec_names).count() == 0 + finally: + _clean()