lamindb_setup 1.14.0__tar.gz → 1.15.0__tar.gz

{lamindb_setup-1.14.0 → lamindb_setup-1.15.0}/.gitignore

@@ -113,3 +113,5 @@ _docs_tmp*
 *.db
 storage_uid.txt
 test.ipynb
+test2.ipynb
+lamindb_export

{lamindb_setup-1.14.0 → lamindb_setup-1.15.0}/.pre-commit-config.yaml

@@ -42,5 +42,7 @@ repos:
       - id: mypy
         exclude: |
           (?x)(
-              tests/hub-local/conftest.py
+              tests/hub-local/conftest.py|
+              tests/hub-prod/conftest.py|
+              tests/storage/conftest.py
           )

{lamindb_setup-1.14.0 → lamindb_setup-1.15.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: lamindb_setup
-Version: 1.14.0
+Version: 1.15.0
 Summary: Setup & configure LaminDB.
 Author-email: Lamin Labs <open-source@lamin.ai>
 Requires-Python: >=3.10

{lamindb_setup-1.14.0 → lamindb_setup-1.15.0}/lamindb_setup/__init__.py

@@ -35,7 +35,7 @@ Migration management
 
 """
 
-__version__ = "1.14.0"  # denote a release candidate for 0.1.0 with 0.1rc1
+__version__ = "1.15.0"  # denote a release candidate for 0.1.0 with 0.1rc1
 
 import os
 import warnings
@@ -47,7 +47,7 @@ warnings.filterwarnings("ignore", category=DeprecationWarning, module="postgrest
 
 from packaging import version as packaging_version
 
-from . import core, errors, types
+from . import core, errors, io, types
 from ._check_setup import _check_instance_setup
 from ._connect_instance import connect
 from ._delete import delete

{lamindb_setup-1.14.0 → lamindb_setup-1.15.0}/lamindb_setup/_connect_instance.py

@@ -204,18 +204,12 @@ def reset_django_module_variables():
     app_names = {app.name for app in apps.get_app_configs()}
     # always copy before iterations over sys.modules
     # see https://docs.python.org/3/library/sys.html#sys.modules
+    # this whole thing runs about 50ms in a big env
     for name, module in sys.modules.copy().items():
         if (
             module is not None
             and (not name.startswith("__") or name == "__main__")
             and name not in sys.builtin_module_names
-            and not (
-                hasattr(module, "__file__")
-                and module.__file__
-                and any(
-                    path in module.__file__ for path in ["/lib/python", "\\lib\\python"]
-                )
-            )
         ):
             try:
                 for k, v in vars(module).items():
@@ -260,6 +254,10 @@ def _connect_cli(
         connect(_write_settings=False, _reload_lamindb=False)
     else:
         logger.important(f"connected lamindb: {isettings.slug}")
+    if settings_.dev_dir is None:
+        logger.important_hint(
+            "to map a local dev directory, call: lamin settings set dev-dir ."
+        )
     return None
 
 
@@ -427,6 +425,10 @@ def connect(instance: str | None = None, **kwargs: Any) -> str | tuple | None:
                 isettings._get_settings_file().unlink(missing_ok=True)  # type: ignore
             settings._instance_settings = None
         raise e
+    if settings.dev_dir is None:
+        logger.important_hint(
+            "to map a local dev directory, set: ln.setup.settings.dev_dir = '.'"
+        )
     return None
 
 

{lamindb_setup-1.14.0 → lamindb_setup-1.15.0}/lamindb_setup/_init_instance.py

@@ -7,8 +7,6 @@ from typing import TYPE_CHECKING, Literal
 from uuid import UUID
 
 import click
-from django.core.exceptions import FieldError
-from django.db.utils import IntegrityError, OperationalError, ProgrammingError
 from lamin_utils import logger
 
 from ._disconnect import disconnect
@@ -71,6 +69,10 @@ def register_storage_in_instance(ssettings: StorageSettings) -> Storage:
 
 
 def register_user(usettings: UserSettings, update_user: bool = True) -> None:
+    # we have to import this here dynamically because otherwise
+    # the except below will fail on re-connect due to reset
+    from django.core.exceptions import FieldError
+    from django.db.utils import IntegrityError, OperationalError, ProgrammingError
     from lamindb.models import User
 
     if not update_user and User.objects.filter(uid=usettings.uid).exists():

{lamindb_setup-1.14.0 → lamindb_setup-1.15.0}/lamindb_setup/_setup_user.py

@@ -7,6 +7,7 @@ from lamin_utils import logger
 
 from ._check_setup import _check_instance_setup
 from ._init_instance import register_user
+from .core._aws_options import reset_aws_options_cache
 from .core._settings import settings
 from .core._settings_load import load_user_settings
 from .core._settings_save import save_user_settings
@@ -147,6 +148,8 @@ def login(
         register_user(user_settings)
 
     settings._user_settings = None
+    # aws s3 credentials are scoped to the user
+    reset_aws_options_cache()
     return user_settings
 
 
@@ -163,6 +166,8 @@ def logout():
     if current_user_settings_file().exists():
         current_user_settings_file().unlink()
         settings._user_settings = None
+        # aws s3 credentials are scoped to the user
+        reset_aws_options_cache()
         logger.success("logged out")
     else:
         logger.important("already logged out")

{lamindb_setup-1.14.0 → lamindb_setup-1.15.0}/lamindb_setup/core/__init__.py

@@ -1,8 +1,23 @@
 """Core setup library.
 
+General
+-------
+
 .. autoclass:: SetupSettings
+
+User
+----
+
 .. autoclass:: UserSettings
+
+Instance
+--------
+
 .. autoclass:: InstanceSettings
+
+Storage
+-------
+
 .. autoclass:: StorageSettings
 
 """

{lamindb_setup-1.14.0 → lamindb_setup-1.15.0}/lamindb_setup/core/_aws_options.py

@@ -245,3 +245,11 @@ def get_aws_options_manager() -> AWSOptionsManager:
         _aws_options_manager = AWSOptionsManager()
 
     return _aws_options_manager
+
+
+def reset_aws_options_cache():
+    global _aws_options_manager
+
+    if _aws_options_manager is not None:
+        _aws_options_manager._credentials_cache = {}
+        _aws_options_manager._parameters_cache = {}

{lamindb_setup-1.14.0 → lamindb_setup-1.15.0}/lamindb_setup/core/_settings.py

@@ -59,7 +59,6 @@ class SetupSettings:
 
     _auto_connect_path: Path = settings_dir / "auto_connect"
     _private_django_api_path: Path = settings_dir / "private_django_api"
-    _work_dir: Path = settings_dir / "work_dir.txt"
 
     _cache_dir: Path | None = None
 
@@ -70,25 +69,6 @@ class SetupSettings:
     def _instance_settings_path(self) -> Path:
         return current_instance_settings_file()
 
-    @property
-    def work_dir(self) -> Path | None:
-        """Get or set the current working directory.
-
-        If setting it to `None`, the working directory is unset
-        """
-        if not self._work_dir.exists():
-            return None
-        return Path(self._work_dir.read_text())
-
-    @work_dir.setter
-    def work_dir(self, value: str | Path | None) -> None:
-        if value is None:
-            if self._work_dir.exists():
-                self._work_dir.unlink()
-        else:
-            value_str = Path(value).expanduser().resolve().as_posix()
-            self._work_dir.write_text(value_str)
-
     @property
     def settings_dir(self) -> Path:
         """The directory that holds locally persisted settings."""
@@ -112,6 +92,31 @@ class SetupSettings:
         else:
             self._auto_connect_path.unlink(missing_ok=True)
 
+    @property
+    def _dev_dir_path(self) -> Path:
+        return (
+            settings_dir / f"dev-dir--{self.instance.owner}--{self.instance.name}.txt"
+        )
+
+    @property
+    def dev_dir(self) -> Path | None:
+        """Get or set the local development directory for the current instance.
+
+        If setting it to `None`, the working development directory is unset.
+        """
+        if not self._dev_dir_path.exists():
+            return None
+        return Path(self._dev_dir_path.read_text())
+
+    @dev_dir.setter
+    def dev_dir(self, value: str | Path | None) -> None:
+        if value is None:
+            if self._dev_dir_path.exists():
+                self._dev_dir_path.unlink()
+        else:
+            value_str = Path(value).expanduser().resolve().as_posix()
+            self._dev_dir_path.write_text(value_str)
+
     @property
     def _branch_path(self) -> Path:
         return (
@@ -361,9 +366,9 @@ class SetupSettings:
         if self._instance_exists:
             instance_rep = self.instance.__repr__().split("\n")
             repr += f"{colors.cyan('Instance:')} {instance_rep[0].replace('Instance: ', '')}\n"
-            repr += f" - work-dir: {self.work_dir}\n"
             repr += f" - branch: {self._read_branch_idlike_name()[1]}\n"
-            repr += f" - space: {self._read_space_idlike_name()[1]}"
+            repr += f" - space: {self._read_space_idlike_name()[1]}\n"
+            repr += f" - dev-dir: {self.dev_dir}"
             repr += f"\n{colors.yellow('Details:')}\n"
             repr += "\n".join(instance_rep[1:])
         else:

lamindb_setup-1.15.0/lamindb_setup/io.py

@@ -0,0 +1,194 @@
+from __future__ import annotations
+
+import json
+import warnings
+from importlib import import_module
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+import pandas as pd
+from django.db import models, transaction
+from rich.progress import Progress
+
+if TYPE_CHECKING:
+    from collections.abc import Sequence
+    from typing import Literal
+
+
+def _get_registries(module_name: str) -> list[str]:
+    """Get registry class names from a module."""
+    schema_module = import_module(module_name)
+    exclude = {"SQLRecord", "BaseSQLRecord"}
+
+    if module_name == "lamindb":
+        module_filter = lambda cls, name: cls.__module__.startswith(
+            f"{module_name}.models."
+        ) and name in dir(schema_module)
+    else:
+        module_filter = (
+            lambda cls, name: cls.__module__ == f"{module_name}.models"
+            and name in dir(schema_module)
+        )
+
+    return [
+        name
+        for name in dir(schema_module.models)
+        if (
+            name[0].isupper()
+            and isinstance(cls := getattr(schema_module.models, name, None), type)
+            and issubclass(cls, models.Model)
+            and module_filter(cls, name)
+            and name not in exclude
+        )
+    ]
+
+
+def _export_registry_to_parquet(registry: type[models.Model], directory: Path) -> None:
+    """Export a single registry table to parquet."""
+    import lamindb_setup as ln_setup
+
+    table_name = registry._meta.db_table
+    with warnings.catch_warnings():
+        warnings.filterwarnings("ignore", message="Skipped unsupported reflection")
+        df = pd.read_sql_table(table_name, ln_setup.settings.instance.db)
+    df.to_parquet(directory / f"{table_name}.parquet", compression=None)
+
+
+def export_db(
+    module_names: Sequence[str] | None = None,
+    *,
+    output_dir: str | Path = "./lamindb_export/",
+) -> None:
+    """Export registry tables and many-to-many link tables to parquet files.
+
+    Ensure that you connect to postgres instances using `use_root_db_user=True`.
+
+    Args:
+        module_names: Module names to export (e.g., ["lamindb", "bionty", "wetlab"]).
+            Defaults to "lamindb" if not provided.
+        output_dir: Directory path for exported parquet files.
+    """
+    directory = Path(output_dir)
+    directory.mkdir(parents=True, exist_ok=True)
+
+    module_names = module_names or ["lamindb"]
+    modules = {name: _get_registries(name) for name in module_names}
+    total_models = sum(len(models) for models in modules.values())
+
+    with Progress() as progress:
+        task = progress.add_task("Exporting", total=total_models)
+        for module_name, model_names in modules.items():
+            schema_module = import_module(module_name)
+            for model_name in model_names:
+                progress.update(task, description=f"[cyan]{module_name}.{model_name}")
+                registry = getattr(schema_module, model_name)
+                _export_registry_to_parquet(registry, directory)
+                for field in registry._meta.many_to_many:
+                    link_orm = getattr(registry, field.name).through
+                    _export_registry_to_parquet(link_orm, directory)
+                progress.advance(task)
+
+
+def _import_registry(
+    registry: type[models.Model],
+    directory: Path,
+    if_exists: Literal["fail", "replace", "append"] = "replace",
+) -> None:
+    """Import a single registry table from parquet.
+
+    Uses raw SQL export instead of django to later circumvent FK constraints.
+    """
+    table_name = registry._meta.db_table
+    parquet_file = directory / f"{table_name}.parquet"
+
+    if not parquet_file.exists():
+        print(f"Skipped {table_name} (file not found)")
+        return
+
+    df = pd.read_parquet(parquet_file)
+
+    old_foreign_key_columns = [col for col in df.columns if col.endswith("_old")]
+    if old_foreign_key_columns:
+        df = df.drop(columns=old_foreign_key_columns)
+
+    for col in df.columns:
+        if df[col].dtype == "object":
+            df[col] = df[col].apply(
+                lambda x: json.dumps(x) if isinstance(x, (dict, list)) else x
+            )
+
+    from django.db import connection
+
+    df.to_sql(table_name, connection.connection, if_exists=if_exists, index=False)
+
+
+def import_db(
+    module_names: Sequence[str] | None = None,
+    *,
+    input_dir: str | Path = "./lamindb_export/",
+    if_exists: Literal["fail", "replace", "append"] = "replace",
+) -> None:
+    """Import registry and link tables from parquet files.
+
+    Temporarily disables FK constraints to allow insertion in arbitrary order.
+    Requires superuser/RDS admin privileges for postgres databases.
+
+    Args:
+        input_dir: Directory containing parquet files to import.
+        module_names: Module names to import (e.g., ["lamindb", "bionty", "wetlab"]).
+        if_exists: How to behave if table exists: 'fail', 'replace', or 'append'.
+    """
+    from django.db import connection
+
+    import lamindb_setup as ln_setup
+
+    directory = Path(input_dir)
+
+    if not directory.exists():
+        raise ValueError(f"Directory does not exist: {directory}")
+
+    if module_names is None:
+        parquet_files = list(directory.glob("*.parquet"))
+        detected_modules = {
+            f.name.split("_")[0] for f in parquet_files if "_" in f.name
+        }
+        module_names = sorted(detected_modules)
+
+    modules = {name: _get_registries(name) for name in module_names}
+    total_models = sum(len(models) for models in modules.values())
+
+    # Disable FK constraints to allow insertion in arbitrary order
+    if ln_setup.settings.instance.dialect == "sqlite":
+        with connection.cursor() as cursor:
+            if ln_setup.settings.instance.dialect == "postgresql":
+                cursor.execute("SET session_replication_role = 'replica'")
+            elif ln_setup.settings.instance.dialect == "sqlite":
+                cursor.execute("PRAGMA foreign_keys = OFF")
+
+    with transaction.atomic():
+        if ln_setup.settings.instance.dialect == "postgresql":
+            with connection.cursor() as cursor:
+                cursor.execute("SET CONSTRAINTS ALL DEFERRED")
+
+        with Progress() as progress:
+            task = progress.add_task("Importing", total=total_models)
+            for module_name, model_names in modules.items():
+                schema_module = import_module(module_name)
+                for model_name in model_names:
+                    progress.update(
+                        task, description=f"[cyan]{module_name}.{model_name}"
+                    )
+                    registry = getattr(schema_module, model_name)
+                    _import_registry(registry, directory, if_exists=if_exists)
+                    for field in registry._meta.many_to_many:
+                        link_orm = getattr(registry, field.name).through
+                        _import_registry(link_orm, directory, if_exists=if_exists)
+                    progress.advance(task)
+
+    # Re-enable FK constraints again
+    if ln_setup.settings.instance.dialect == "sqlite":
+        with connection.cursor() as cursor:
+            if ln_setup.settings.instance.dialect == "postgresql":
+                cursor.execute("SET session_replication_role = 'origin'")
+            elif ln_setup.settings.instance.dialect == "sqlite":
+                cursor.execute("PRAGMA foreign_keys = ON")

{lamindb_setup-1.14.0 → lamindb_setup-1.15.0}/noxfile.py

@@ -40,7 +40,7 @@ uv pip install --system git+https://github.com/laminlabs/bionty
     elif group == "docs":
         cmds = modules_deps.strip()
     elif group == "storage":
-        cmds = modules_deps + "uv pip install --system gcsfs huggingface_hub"
+        cmds = modules_deps + "uv pip install --system gcsfs huggingface_hub sqlalchemy"
     elif group == "hub-prod":
         # cmds = "git clone --depth 1 https://github.com/django/django\n"
         # cmds += "uv pip install --system -e ./django\n"

{lamindb_setup-1.14.0 → lamindb_setup-1.15.0}/pyproject.toml

@@ -158,4 +158,6 @@ filterwarnings = [
     "ignore:Jupyter is migrating its paths to use standard platformdirs:DeprecationWarning",
     "ignore:The 'timeout' parameter is deprecated. Please configure it in the http client instead.:DeprecationWarning",
     "ignore:The 'verify' parameter is deprecated. Please configure it in the http client instead.:DeprecationWarning",
+    "ignore:There is no current event loop:DeprecationWarning",
+    "ignore:DateTimeField.*received a naive datetime.*while time zone support is active:RuntimeWarning",
 ]

{lamindb_setup-1.14.0 → lamindb_setup-1.15.0}/tests/hub-local/test_update_schema_in_hub.py

@@ -44,7 +44,7 @@ def test_update_schema_in_hub(setup_instance):
     assert not schema["schema_json"]["core"]["artifact"]["is_link_table"]
     assert schema["schema_json"]["core"]["artifactulabel"]["is_link_table"]
 
-    assert schema["schema_json"]["core"]["artifact"]["name_field"] is None
+    assert schema["schema_json"]["core"]["artifact"]["name_field"] == "key"
     assert schema["schema_json"]["core"]["artifact"]["ontology_id_field"] is None
     assert schema["schema_json"]["bionty"]["gene"]["name_field"] == "symbol"
     assert (

lamindb_setup-1.15.0/tests/storage/conftest.py

@@ -0,0 +1,10 @@
+import pytest
+
+
+@pytest.fixture(scope="session")
+def simple_instance():
+    import lamindb_setup as ln_setup
+
+    ln_setup.init(storage="./testdb", modules="bionty,wetlab")
+    yield
+    ln_setup.delete("testdb", force=True)

lamindb_setup-1.15.0/tests/storage/test_db_import_export.py

@@ -0,0 +1,215 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import pandas as pd
+import pytest
+from lamindb_setup.io import _get_registries, export_db, import_db
+
+if TYPE_CHECKING:
+    from collections.abc import Callable, Generator
+    from pathlib import Path
+
+
+@pytest.fixture
+def cleanup_export_dir(tmp_path) -> Generator[Path, None, None]:
+    output_dir = tmp_path / "test_export"
+    yield output_dir
+
+
+def test_get_registries_lamindb(simple_instance: Callable):
+    registries = _get_registries("lamindb")
+
+    assert "Artifact" in registries
+    assert "Collection" in registries
+    assert "Run" in registries
+    assert "Transform" in registries
+    assert "Record" in registries
+
+    assert "SQLRecord" not in registries
+
+
+def test_get_registries_bionty(simple_instance: Callable):
+    registries = _get_registries("bionty")
+
+    assert "Gene" in registries
+    assert "Protein" in registries
+    assert "CellType" in registries
+    # ... and a few more
+    assert len(registries) > 0
+
+
+def test_exportdb_creates_directory(
+    simple_instance: Callable, cleanup_export_dir: Path
+):
+    export_db(
+        module_names=["lamindb"],
+        output_dir=cleanup_export_dir,
+    )
+
+    assert cleanup_export_dir.exists()
+    assert cleanup_export_dir.is_dir()
+
+
+def test_exportdb_exports_parquet_files(
+    simple_instance: Callable, cleanup_export_dir: Path
+):
+    export_db(
+        module_names=["lamindb"],
+        output_dir=cleanup_export_dir,
+    )
+
+    parquet_files = list(cleanup_export_dir.glob("*.parquet"))
+    assert len(parquet_files) > 0
+
+    for file in parquet_files:
+        df = pd.read_parquet(file)
+        assert isinstance(df, pd.DataFrame)
+
+
+def test_exportdb_multiple_modules(simple_instance: Callable, cleanup_export_dir: Path):
+    import bionty as bt
+    import lamindb as ln
+
+    artifact = ln.Artifact.from_dataframe(
+        pd.DataFrame({"col": [1, 2, 3]}), key="test_artifact.parquet"
+    ).save()
+    gene = bt.Gene.from_source(symbol="TCF7").save()
+    artifact.genes.add(gene)
+
+    export_db(
+        module_names=["lamindb", "bionty"],
+        output_dir=cleanup_export_dir,
+    )
+
+    lamindb_files = list(cleanup_export_dir.glob("lamindb_*.parquet"))
+    bionty_files = list(cleanup_export_dir.glob("bionty_*.parquet"))
+
+    assert len(lamindb_files) > 0
+    assert len(bionty_files) > 0
+
+    gene_df = pd.read_parquet(cleanup_export_dir / "bionty_gene.parquet")
+    assert "TCF7" in gene_df["symbol"].values
+
+    artifact_df = pd.read_parquet(cleanup_export_dir / "lamindb_artifact.parquet")
+    assert artifact.uid in artifact_df["uid"].values
+
+    link_df = pd.read_parquet(cleanup_export_dir / "bionty_artifactgene.parquet")
+    assert (
+        len(
+            link_df[
+                (link_df["artifact_id"] == artifact.id)
+                & (link_df["gene_id"] == gene.id)
+            ]
+        )
+        == 1
+    )
+
+    artifact.delete(permanent=True)
+    gene.delete(permanent=True)
+
+
+def test_exportdb_default_module(simple_instance: Callable, cleanup_export_dir: Path):
+    export_db(output_dir=cleanup_export_dir)
+
+    lamindb_files = list(cleanup_export_dir.glob("lamindb_*.parquet"))
+    assert len(lamindb_files) > 0
+
+
+def test_exportdb_exports_link_tables(
+    simple_instance: Callable, cleanup_export_dir: Path
+):
+    export_db(module_names=["lamindb"], output_dir=cleanup_export_dir)
+
+    parquet_files = [f.name for f in cleanup_export_dir.glob("*.parquet")]
+    link_tables = [f for f in parquet_files if "_" in f and "artifact" in f.lower()]
+
+    assert len(link_tables) > 0
+
+
+def test_import_db_from_parquet(simple_instance: Callable, tmp_path):
+    """Tests imports of a parquet file.
+
+    Implicitly also tests whether `import_db` can deal with FK constraints.
+    """
+    import bionty as bt
+    import lamindb as ln
+    import lamindb_setup as ln_setup
+
+    export_dir = tmp_path / "export"
+    export_dir.mkdir()
+
+    artifact_data = pd.DataFrame(
+        {
+            "id": [888],
+            "uid": ["test_artifact_uid"],
+            "key": ["test_key"],
+            "_key_is_virtual": [False],
+            "_overwrite_versions": [False],
+            "description": ["Test artifact"],
+            "suffix": [".txt"],
+            "kind": ["dataset"],
+            "size": [1024],
+            "hash": ["testhash123"],
+            "is_latest": [True],
+            "is_locked": [False],
+            "storage_id": [1],
+            "created_by_id": [ln_setup.settings.user.id],
+            "created_at": [pd.Timestamp.now()],
+            "updated_at": [pd.Timestamp.now()],
+        }
+    )
+    artifact_data.to_parquet(export_dir / "lamindb_artifact.parquet", index=False)
+
+    gene_data = pd.DataFrame(
+        {
+            "id": [999],
+            "uid": ["test_uid_999"],
+            "symbol": ["TESTGENE"],
+            "ensembl_gene_id": ["ENSG00000999"],
+            "ncbi_gene_ids": [None],
+            "biotype": ["protein_coding"],
+            "description": ["Test gene for import"],
+            "synonyms": [None],
+            "organism_id": [1],
+            "source_id": [1],
+            "created_by_id": [ln_setup.settings.user.id],
+            "created_at": [pd.Timestamp.now()],
+            "updated_at": [pd.Timestamp.now()],
+        }
+    )
+    gene_data.to_parquet(export_dir / "bionty_gene.parquet", index=False)
+
+    link_data = pd.DataFrame(
+        {
+            "id": [1],
+            "artifact_id": [888],
+            "gene_id": [999],
+            "feature_id": [None],
+            "feature_ref_is_name": [None],
+            "label_ref_is_name": [None],
+            "created_at": [pd.Timestamp.now()],
+            "created_by_id": [ln_setup.settings.user.id],
+            "run_id": [None],
+        }
+    )
+    link_data.to_parquet(export_dir / "bionty_artifactgene.parquet", index=False)
+
+    import_db(
+        input_dir=export_dir,
+        module_names=["lamindb", "bionty"],
+        if_exists="append",
+    )
+
+    # gene and artifact should exist after the import
+    imported_gene = bt.Gene.get(id=999)
+    assert imported_gene.symbol == "TESTGENE"
+    assert imported_gene.ensembl_gene_id == "ENSG00000999"
+    imported_artifact = ln.Artifact.get(id=888)
+    assert imported_artifact.key == "test_key"
+    assert imported_artifact.genes.count() == 1
+
+    # they should also be linked
+    linked_gene = imported_artifact.genes.first()
+    assert linked_gene.id == 999
+    assert linked_gene.symbol == "TESTGENE"

lamindb_setup-1.14.0/lamindb_setup/_exportdb.py

@@ -1,68 +0,0 @@
-from __future__ import annotations
-
-from importlib import import_module
-from pathlib import Path
-
-MODELS = {
-    "lamindb": {
-        "Collection": False,
-        "Artifact": False,
-        "Transform": False,
-        "Run": True,
-        "User": False,
-        "Storage": False,
-        "Feature": False,
-        "Schema": False,
-        "ULabel": False,
-    },
-    # "bionty": {
-    #     "Organism": False,
-    #     "Gene": False,
-    #     "Protein": False,
-    #     "CellMarker": False,
-    #     "Tissue": False,
-    #     "CellType": False,
-    #     "Disease": False,
-    #     "CellLine": False,
-    #     "Phenotype": False,
-    #     "Pathway": False,
-    #     "ExperimentalFactor": False,
-    #     "DevelopmentalStage": False,
-    #     "Ethnicity": False,
-    #     "Source": False,
-    # },
-    # "wetlab": {
-    #     "ExperimentType": False,
-    #     "Experiment": False,
-    #     "Well": False,
-    #     "TreatmentTarget": False,
-    #     "Treatment": False,
-    #     "Biosample": False,
-    #     "Techsample": False,
-    # },
-}
-
-
-def exportdb() -> None:
-    directory = Path("./lamindb_export/")
-    directory.mkdir(parents=True, exist_ok=True)
-    import pandas as pd
-
-    import lamindb_setup as ln_setup
-
-    def export_registry(registry, directory):
-        table_name = registry._meta.db_table
-        df = pd.read_sql_table(table_name, ln_setup.settings.instance.db)
-        df.to_parquet(directory / f"{table_name}.parquet", compression=None)
-
-    # export data to parquet files
-    print(f"\nexporting data to parquet files in: {directory}\n")
-    for module_name, models in MODELS.items():
-        for model_name in models.keys():
-            schema_module = import_module(f"lnschema_{module_name}")
-            registry = getattr(schema_module, model_name)
-            export_registry(registry, directory)
-            many_to_many_names = [field.name for field in registry._meta.many_to_many]
-            for many_to_many_name in many_to_many_names:
-                link_orm = getattr(registry, many_to_many_name).through
-                export_registry(link_orm, directory)

lamindb_setup-1.14.0/lamindb_setup/_importdb.py

@@ -1,50 +0,0 @@
-from __future__ import annotations
-
-from importlib import import_module
-from pathlib import Path
-
-from ._exportdb import MODELS
-
-
-def import_registry(registry, directory, connection):
-    import pandas as pd
-
-    table_name = registry._meta.db_table
-    df = pd.read_parquet(directory / f"{table_name}.parquet")
-    old_foreign_key_columns = [
-        column for column in df.columns if column.endswith("_old")
-    ]
-    for column in old_foreign_key_columns:
-        df.drop(column, axis=1, inplace=True)
-    df.to_sql(table_name, connection, if_exists="append", index=False)
-
-
-def importdb() -> None:
-    # import data from parquet files
-    directory = Path("./lamindb_export/")
-    if directory.exists():
-        response = input(
-            f"\n\nDo you want to import registries from here: {directory}? (y/n)\n"
-        )
-        if response != "y":
-            return None
-    from sqlalchemy import create_engine, text
-
-    import lamindb_setup as ln_setup
-
-    engine = create_engine(ln_setup.settings.instance.db, echo=False)
-    with engine.begin() as connection:
-        if ln_setup.settings.instance.dialect == "postgresql":
-            connection.execute(text("SET CONSTRAINTS ALL DEFERRED;"))
-        for module_name, models in MODELS.items():
-            for model_name in models.keys():
-                print(model_name)
-                schema_module = import_module(f"lnschema_{module_name}")
-                registry = getattr(schema_module, model_name)
-                import_registry(registry, directory, connection)
-                many_to_many_names = [
-                    field.name for field in registry._meta.many_to_many
-                ]
-                for many_to_many_name in many_to_many_names:
-                    link_orm = getattr(registry, many_to_many_name).through
-                    import_registry(link_orm, directory, connection)