lamindb 0.77.2__py3-none-any.whl → 1.0rc1__py3-none-any.whl

lamindb/migrations/0070_lamindbv1_migrate_data.py

@@ -0,0 +1,78 @@
+# Generated by Django 5.2 on 2025-01-05 11:58
+
+from pathlib import Path
+
+import lamindb_setup as ln_setup
+import psycopg2
+from django.db import migrations
+
+
+def get_artifact_path_psycopg2(artifact_id):
+    """Get artifact path using psycopg2."""
+    query = """
+        SELECT
+            s.root || '/.lamindb/' || a.uid || a.suffix AS full_path
+        FROM
+            lamindb_artifact a
+            JOIN lamindb_storage s ON a.storage_id = s.id
+        WHERE
+            a.id = %s
+    """
+
+    with psycopg2.connect(ln_setup.settings.instance.db) as conn:
+        with conn.cursor() as cur:
+            cur.execute(query, (artifact_id,))
+            return cur.fetchone()[0]
+
+
+def transfer_source_code(apps, schema_editor):
+    from lamindb._finish import notebook_to_script
+
+    Transform = apps.get_model("lamindb", "Transform")
+    transforms = Transform.objects.filter(
+        _source_code_artifact__isnull=False,
+    ).select_related("_source_code_artifact")
+
+    for transform in transforms:
+        print(f"migrating source code of transform {transform}")
+        artifact = transform._source_code_artifact
+        print("artifact", artifact.uid)
+
+        path_str = get_artifact_path_psycopg2(artifact.id)
+        print(ln_setup.settings.storage.root_as_str)
+        print(path_str)
+        if path_str.startswith(ln_setup.settings.storage.root_as_str):
+            path = (
+                ln_setup.settings.storage.root
+                / f".lamindb/{artifact.uid}{artifact.suffix}"
+            )
+        else:
+            path = ln_setup.core.upath.UPath(path_str)
+        if path.exists():
+            if path_str.startswith("s3://"):
+                local_path = Path(f"temp{path.suffix}")
+                path.download_to(local_path)
+            else:
+                local_path = path
+
+            if artifact.suffix == ".ipynb":
+                transform.source_code = notebook_to_script(transform, local_path)
+            else:
+                transform.source_code = local_path.read_text()
+            transform.hash = artifact.hash
+            path.unlink()
+        else:
+            print(f"path did not exist: {path_str}")
+        transform._source_code_artifact = None
+        transform.save()
+        artifact.delete()
+
+
+class Migration(migrations.Migration):
+    dependencies = [
+        ("lamindb", "0069_squashed"),
+    ]
+
+    operations = [
+        migrations.RunPython(transfer_source_code),
+    ]