Able to select samples and sort them before generating genlabIDs.

src/genlab_bestilling/libs/genlabid.py

@@ -1,196 +1,83 @@
-from typing import Any
-
-import sqlglot
-import sqlglot.expressions
-from django.db import connection, transaction
-from sqlglot.expressions import (
-    EQ,
-    Alias,
-    Cast,
-    DataType,
-    Extract,
-    From,
-    Literal,
-    Subquery,
-    and_,
-    column,
-)
-
-from ..models import ExtractionOrder, Order, Sample, Species
-
-
-def get_replica_for_sample() -> None:
-    """
-    TODO: implement
-    """
-    pass
-
-
-def get_current_sequences(order_id: int | str) -> Any:
-    """
-    Invoke a Postgres function to get the current sequence number
-    for a specific combination of year and species.
-    """
+from django.db import transaction
+from django.db.models import Case, QuerySet, When
+
+from ..models import Sample
+
+
+# Format the genlab ID as GYYCODEXXXX
+def generate_genlab_id(code: str, year: int, count: int) -> str:
+    return f"G{str(year)[-2:]}{code.upper()}{count:04d}"
+
+
+# Order samples by their names, assuming they are numeric strings
+def order(samples: QuerySet) -> QuerySet:
+    name_to_sample = {sample.name.strip(): sample for sample in samples}
+    names = list(name_to_sample.keys())
+
+    # Check all are digits
+    if not all(name.isdigit() for name in names):
+        return samples
+
+    # Sort names as integers
+    sorted_names = sorted(names, key=lambda x: int(x))
+
+    # Build a CASE statement to preserve order in SQL
+    when_statements = [
+        When(name=name, then=pos) for pos, name in enumerate(sorted_names)
+    ]
+
+    return samples.order_by(Case(*when_statements))
+
+
+# New entry point for generating genlab IDs
+@transaction.atomic
+def generate(
+    order_id: int,
+    sorting_order: list[str] | None = None,
+    selected_samples: list[int] | None = None,
+) -> None:
+    if selected_samples is None:
+        return
+
+    # Sort order, default is by name
+    order_by = sorting_order or ["name"]
+
+    # Fetch samples
     samples = (
-        Sample.objects.select_related("order", "species")
-        .filter(order=order_id)
-        .values("order__created_at", "species__code")
-        .distinct()
+        Sample.objects.select_related("species", "order")
+        .filter(order_id=order_id, id__in=selected_samples, genlab_id__isnull=True)
+        .order_by(*order_by)
     )
 
-    with connection.cursor() as cursor:
-        sequences = {}
-        for sample in samples:
-            query = sqlglot.select(
-                sqlglot.expressions.func(
-                    "get_genlab_sequence_name",
-                    sqlglot.expressions.Literal.string(sample["species__code"]),
-                    sqlglot.expressions.Literal.number(
-                        sample["order__created_at"].year
-                    ),
-                    dialect="postgres",
-                )
-            ).sql(dialect="postgres")
-            seq = cursor.execute(
-                query,
-            ).fetchone()[0]
-            sequences[seq] = 0
-
-        for k in sequences.keys():
-            query = sqlglot.select("last_value").from_(k).sql(dialect="postgres")
-            sequences[k] = cursor.execute(query).fetchone()[0]
-
-        return sequences
-
-
-def generate(order_id: int | str) -> None:
-    """
-    wrapper to handle errors and reset the sequence to the current sequence value
-    """
-    sequences = get_current_sequences(order_id)
-    print(sequences)
-
-    with connection.cursor() as cursor:
-        try:
-            with transaction.atomic():
-                cursor.execute(update_genlab_id_query(order_id))
-        except Exception:
-            # if there is an error, reset the sequence
-            # NOTE: this is unsafe unless this function is executed in a queue
-            # by just one worker to prevent concurrency
-            with connection.cursor() as cursor:  # noqa: PLW2901 # Was this intentional?
-                for k, v in sequences.items():
-                    cursor.execute("SELECT setval(%s, %s)", [k, v])
-
-    sequences = get_current_sequences(order_id)
-    print(sequences)
-
-
-def update_genlab_id_query(order_id: int | str) -> Any:
-    """
-    Safe generation of a SQL raw query using sqlglot
-    The query runs an update on all the rows with a specific order_id
-    and set genlab_id = generate_genlab_id(code, year)
-    """
-    samples_table = Sample._meta.db_table
-    extraction_order_table = ExtractionOrder._meta.db_table
-    order_table = Order._meta.db_table
-    species_table = Species._meta.db_table
-
-    return sqlglot.expressions.update(
-        samples_table,
-        properties={
-            "genlab_id": column(
-                "genlab_id",
-                table="order_samples",
-            )
-        },
-        where=sqlglot.expressions.EQ(
-            this=column("id", table=samples_table),
-            expression="order_samples.id",
-        ),
-        from_=From(
-            this=Alias(
-                this=Subquery(
-                    this=(
-                        sqlglot.select(
-                            column(
-                                "id",
-                                table=samples_table,
-                            ),
-                            Alias(
-                                this=sqlglot.expressions.func(
-                                    "generate_genlab_id",
-                                    column("code", table=species_table),
-                                    Cast(
-                                        this=Extract(
-                                            this="YEAR",
-                                            expression=column(
-                                                "confirmed_at",
-                                                table=order_table,
-                                            ),
-                                        ),
-                                        to=DataType(
-                                            this=DataType.Type.INT, nested=False
-                                        ),
-                                    ),
-                                ),
-                                alias="genlab_id",
-                            ),
-                        )
-                        .join(
-                            extraction_order_table,
-                            on=EQ(
-                                this=column(
-                                    "order_id",
-                                    table=samples_table,
-                                ),
-                                expression=column(
-                                    "order_ptr_id",
-                                    table=extraction_order_table,
-                                ),
-                            ),
-                        )
-                        .join(
-                            order_table,
-                            on=EQ(
-                                this=column(
-                                    "order_ptr_id",
-                                    table=extraction_order_table,
-                                ),
-                                expression=column(
-                                    "id",
-                                    table=order_table,
-                                ),
-                            ),
-                        )
-                        .from_(samples_table)
-                        .join(
-                            species_table,
-                            on=EQ(
-                                this=column(
-                                    "species_id",
-                                    table=samples_table,
-                                ),
-                                expression=column(
-                                    "id",
-                                    table=species_table,
-                                ),
-                            ),
-                        )
-                        .where(
-                            and_(
-                                EQ(
-                                    this=column(col="order_id"),
-                                    expression=Literal.number(order_id),
-                                ),
-                                column(col="genlab_id").is_(None),
-                            )
-                        )
-                        .order_by(column("name", table=samples_table))
-                    ),
-                ),
-                alias="order_samples",
-            )
-        ),
-    ).sql(dialect="postgres")
+    # If sorting order is by name, we might need to order them numerically
+    if order_by == ["name"]:
+        samples = order(samples)
+
+    # Group counts by species + year
+    counts = {}
+    updates = []
+
+    # Iterate through samples to generate genlab IDs in correct order
+    for sample in list(samples):
+        # Extract species code and year from the sample
+        species_code = sample.species.code
+        year = sample.order.confirmed_at.year
+        key = (species_code, year)
+
+        # Initialize count
+        if key not in counts:
+            existing_count = Sample.objects.filter(
+                species__code=species_code,
+                order__confirmed_at__year=year,
+                genlab_id__isnull=False,
+            ).count()
+            counts[key] = existing_count + 1
+        else:
+            counts[key] += 1
+
+        genlab_id = generate_genlab_id(species_code, year, counts[key])
+        sample.genlab_id = genlab_id
+        updates.append(sample)
+
+    # Bulk update samples with new genlab IDs
+    Sample.objects.bulk_update(updates, ["genlab_id"])

src/genlab_bestilling/models.py

@@ -295,6 +295,10 @@ def to_draft(self) -> None:
     def get_type(self) -> str:
         return "order"
 
+    @property
+    def filled_genlab_count(self) -> int:
+        return self.samples.filter(genlab_id__isnull=False).count()
+
     @property
     def next_status(self) -> OrderStatus | None:
         current_index = self.STATUS_ORDER.index(self.status)
@@ -452,6 +456,26 @@ def order_manually_checked(self) -> None:
         self.save()
         app.configure_task(name="generate-genlab-ids").defer(order_id=self.id)
 
+    def order_selected_checked(
+        self,
+        sorting_order: list[str] | None = None,
+        selected_samples: list[str] | None = None,
+    ) -> None:
+        """
+        Partially set the order as checked by the lab staff, generate a genlab id
+        """
+        self.internal_status = self.Status.CHECKED
+        self.status = self.OrderStatus.PROCESSING
+        self.save()
+
+        selected_sample_names = list(selected_samples.values_list("id", flat=True))
+
+        app.configure_task(name="generate-genlab-ids").defer(
+            order_id=self.id,
+            sorting_order=sorting_order,
+            selected_samples=selected_sample_names,
+        )
+
 
 class AnalysisOrder(Order):
     samples = models.ManyToManyField(

src/genlab_bestilling/tasks.py

@@ -1,4 +1,6 @@
 # from .libs.isolation import isolate
+from typing import Any
+
 from django.db.utils import OperationalError
 from procrastinate import RetryStrategy
 from procrastinate.contrib.django import app
@@ -12,6 +14,14 @@
         max_attempts=5, linear_wait=5, retry_exceptions={OperationalError}
     ),
 )
-def generate_ids(order_id: str | int) -> None:
-    generate_genlab_id(order_id=order_id)
+def generate_ids(
+    order_id: int | str,
+    sorting_order: list[str] | None = None,
+    selected_samples: list[Any] | None = None,
+) -> None:
+    generate_genlab_id(
+        order_id=order_id,
+        sorting_order=sorting_order,
+        selected_samples=selected_samples,
+    )
     # isolate(order_id=order_id)

src/staff/tables.py

@@ -1,6 +1,8 @@
 from typing import Any
 
 import django_tables2 as tables
+from django.db.models import IntegerField
+from django.db.models.functions import Cast
 from django.utils.safestring import mark_safe
 
 from genlab_bestilling.models import (
@@ -123,6 +125,16 @@ class SampleBaseTable(tables.Table):
     plate_positions = tables.Column(
         empty_values=(), orderable=False, verbose_name="Extraction position"
     )
+    checked = tables.CheckBoxColumn(
+        attrs={
+            "th__input": {"type": "checkbox", "id": "select-all-checkbox"},
+        },
+        accessor="pk",
+        orderable=False,
+        empty_values=(),
+    )
+
+    name = tables.Column(order_by=("name_as_int",))
 
     class Meta:
         model = Sample
@@ -139,15 +151,28 @@ class Meta:
             "plate_positions",
         ]
         attrs = {"class": "w-full table-auto tailwind-table table-sm"}
+        sequence = ("checked",)
+        order_by = ("genlab_id", "name_as_int", "species")
 
         empty_text = "No Samples"
 
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+        if hasattr(self.data, "data"):
+            self.data.data = self.data.data.annotate(
+                name_as_int=Cast("name", output_field=IntegerField())
+            )
+
     def render_plate_positions(self, value: Any) -> str:
         if value:
             return ", ".join([str(v) for v in value.all()])
 
         return ""
 
+    def render_checked(self, record: Any) -> str:
+        return mark_safe(f'<input type="checkbox" name="checked" value="{record.id}">')
+
 
 class OrderExtractionSampleTable(SampleBaseTable):
     class Meta(SampleBaseTable.Meta):