From 22ac1ecfe992b332fe644bcbe9fa9ab736bcdb6a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Rapha=C3=ABl=20Bournhonesque?= <raphael@bournhonesque.eu>
Date: Mon, 19 Aug 2024 15:28:17 +0200
Subject: [PATCH] fix: add logo_annotation.server_type field

This is the same as ImageModel.server_type, but we store it here to
avoid performing a double join with ImageModel table when we need to
filter logos by server_type
(LogoAnnotation > ImagePrediction > ImageModel)
---
 .../006_add_logo_annotation_server_type.py    | 23 +++++++++++++++++++
 robotoff/cli/logos.py                         |  1 +
 robotoff/models.py                            | 15 ++++++++++--
 robotoff/workers/tasks/import_image.py        |  1 +
 scripts/insert_image_predictions.py           |  1 +
 5 files changed, 39 insertions(+), 2 deletions(-)
 create mode 100644 migrations/006_add_logo_annotation_server_type.py

diff --git a/migrations/006_add_logo_annotation_server_type.py b/migrations/006_add_logo_annotation_server_type.py
new file mode 100644
index 0000000000..a50b5f4649
--- /dev/null
+++ b/migrations/006_add_logo_annotation_server_type.py
@@ -0,0 +1,23 @@
+import peewee as pw
+from peewee_migrate import Migrator
+
+
+def migrate(migrator: Migrator, database: pw.Database, *, fake=False):
+    """Write your migrations here."""
+
+    migrator.add_fields(
+        "logo_annotation",
+        server_type=pw.CharField(
+            null=True,
+            max_length=10,
+            help_text="project associated with the logo annotation, "
+            "one of 'off', 'obf', 'opff', 'opf', 'off-pro'",
+            index=False,
+        ),
+    )
+
+
+def rollback(migrator: Migrator, database: pw.Database, *, fake=False):
+    """Write your rollback migrations here."""
+
+    migrator.remove_fields("logo_annotation", "server_type")
diff --git a/robotoff/cli/logos.py b/robotoff/cli/logos.py
index 9fad302dff..b5e5301dd0 100644
--- a/robotoff/cli/logos.py
+++ b/robotoff/cli/logos.py
@@ -89,6 +89,7 @@ def import_logos(
                         bounding_box=item["bounding_box"],
                         barcode=image_instance.barcode,
                         source_image=image_instance.source_image,
+                        server_type=server_type.name,
                     )
                 seen_set.add(key)
 
diff --git a/robotoff/models.py b/robotoff/models.py
index 6c9f2ec713..64866ccfcb 100644
--- a/robotoff/models.py
+++ b/robotoff/models.py
@@ -224,8 +224,8 @@ class Prediction(BaseModel):
     server_type = peewee.CharField(
         null=False,
         max_length=10,
-        help_text="project associated with the insight, "
-        "one of 'off', 'obf', 'opff', 'opf'",
+        help_text="project associated with the prediction, "
+        "one of 'off', 'obf', 'opff', 'opf', 'off-pro'",
         index=True,
         default="off",
     )
@@ -343,6 +343,17 @@ class LogoAnnotation(BaseModel):
     source_image = peewee.TextField(null=True, index=True)
     # The logo text extracted from the image using OCR
     text = peewee.TextField(null=True)
+    # This is the same as ImageModel.server_type, but we store it here to
+    # avoid performing a double join with ImageModel table when we need to
+    # filter logos by server_type
+    # (LogoAnnotation > ImagePrediction > ImageModel)
+    server_type = peewee.CharField(
+        null=True,
+        max_length=10,
+        help_text="project associated with the logo annotation, "
+        "one of 'off', 'obf', 'opff', 'opf', 'off-pro'",
+        index=False,
+    )
 
     class Meta:
         constraints = [peewee.SQL("UNIQUE(image_prediction_id, index)")]
diff --git a/robotoff/workers/tasks/import_image.py b/robotoff/workers/tasks/import_image.py
index d9fe71f566..fcc43a6adb 100644
--- a/robotoff/workers/tasks/import_image.py
+++ b/robotoff/workers/tasks/import_image.py
@@ -507,6 +507,7 @@ def run_logo_object_detection(
                         barcode=image_model.barcode,
                         source_image=image_model.source_image,
                         text=text,
+                        server_type=product_id.server_type.name,
                     )
                 )
             logger.info("%s logos found for image %s", len(logos), source_image)
diff --git a/scripts/insert_image_predictions.py b/scripts/insert_image_predictions.py
index 6352ba47c7..8c85bb9f44 100644
--- a/scripts/insert_image_predictions.py
+++ b/scripts/insert_image_predictions.py
@@ -85,6 +85,7 @@ def insert_batch(
                     bounding_box=item["bounding_box"],
                     barcode=image_instance.barcode,
                     source_image=image_instance.source_image,
+                    server_type=server_type.name,
                 )
         seen_set.add(key)