store recipe similarity data on create/update

Author: elfkuzco

backend/maint-scripts/update_scraper_version.py

@@ -1,4 +1,5 @@
 #!/usr/bin/env python3
+
 """
 Script to update schedules and requested tasks with new image tags or offliner
 definitions.
@@ -54,6 +55,7 @@
 from sqlalchemy.orm.attributes import flag_modified
 
 from zimfarm_backend import logger
+from zimfarm_backend.common.schemas.offliners.builder import generate_similarity_data
 from zimfarm_backend.common.schemas.orms import OfflinerDefinitionSchema, OfflinerSchema
 from zimfarm_backend.db import Session
 from zimfarm_backend.db.models import RequestedTask, Schedule
@@ -103,23 +105,32 @@ def update_entries(
             model.config["offliner"]["offliner_id"].astext == offliner.id
         )
     ).scalars():
+        obj_name = getattr(obj, "name", obj.id)
         if image_tag:
-            logger.info(f"setting {offliner.id} image tag to {args.image_tag}...")
+            logger.info(
+                f"setting {offliner.id} image tag for {model.__tablename__} {obj_name} "
+                f"to {args.image_tag}..."
+            )
             obj.config["image"]["tag"] = image_tag
             flag_modified(obj, "config")
 
         if offliner_definition:
-            logger.info(f"setting offliner defintion for {model.__tablename__}...")
+            logger.info(
+                f"setting offliner defintion for {model.__tablename__} {obj_name} ..."
+            )
             if name_mappings:
                 obj.config["offliner"] = update_offliner_flags(
                     offliner=offliner,
                     offliner_definition=offliner_definition,
                     data=obj.config["offliner"],
                     name_mappings=name_mappings,
                 )
-                flag_modified(obj, "config")
             obj.offliner_definition_id = offliner_definition.id
-
+            if isinstance(obj, Schedule):
+                obj.similarity_data = generate_similarity_data(
+                    obj.config["offliner"], offliner, offliner_definition.schema_
+                )
+            flag_modified(obj, "config")
         # Just needed to ensure our model still works
         if isinstance(obj, Schedule):
             create_schedule_full_schema(obj, offliner)

backend/maint-scripts/update_similarity_data.py

@@ -0,0 +1,135 @@
+#!/usr/bin/env python3
+
+"""
+Script to update similarity data of offliner definitions and schedules
+
+EXAMPLES:
+
+1. Update the similarity data for all offliner definitions. This updates the recipes
+   similarity data too. Input can be a JSON file or sent via stdin.
+   ./update_scraper_version.py -o mwoffliner -s simialrity_data.in
+
+2 To update via stdin:
+   echo '[
+      {
+        "flag": "mwUrl",
+        "transformers": [
+          {
+            "name": "hostname",
+            "operand": null
+          }
+        ]
+      }
+    ]' \
+    | ./update_scraper_version.py \
+      -o mwoffliner \
+      -s
+"""
+
+import argparse
+import json
+import sys
+
+import sqlalchemy as sa
+from sqlalchemy import select
+from sqlalchemy.orm import Session as OrmSession
+from sqlalchemy.orm.attributes import flag_modified
+
+from zimfarm_backend import logger
+from zimfarm_backend.common.schemas.offliners.builder import generate_similarity_data
+from zimfarm_backend.common.schemas.offliners.models import (
+    OfflinerSpecSchema,
+    SimilarityDataSchema,
+)
+from zimfarm_backend.common.schemas.orms import OfflinerDefinitionSchema, OfflinerSchema
+from zimfarm_backend.db import Session
+from zimfarm_backend.db.models import OfflinerDefinition, Schedule
+from zimfarm_backend.db.offliner import get_offliner
+from zimfarm_backend.db.offliner_definition import (
+    create_offliner_definition_schema,
+)
+from zimfarm_backend.db.schedule import create_schedule_full_schema
+
+
+def update_schedules(
+    session: OrmSession,
+    *,
+    offliner: OfflinerSchema,
+    offliner_definition: OfflinerDefinitionSchema,
+) -> int:
+    nb_modified: int = 0
+
+    for schedule in session.execute(
+        sa.select(Schedule).where(
+            Schedule.config["offliner"]["offliner_id"].astext == offliner.id,
+            Schedule.offliner_definition_id == offliner_definition.id,
+        )
+    ).scalars():
+        logger.info(f"seting similarity data for schedule {schedule.name}")
+        schedule.similarity_data = generate_similarity_data(
+            schedule.config["offliner"], offliner, offliner_definition.schema_
+        )
+        flag_modified(schedule, "config")
+        create_schedule_full_schema(schedule, offliner)
+        nb_modified += 1
+    return nb_modified
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(
+        description="Update schedules and requested tasks image tag "
+        "or offliner definition",
+    )
+
+    # Required offliner specification
+    parser.add_argument(
+        "-o", "--offliner", required=True, help="Specify which offliner to update"
+    )
+
+    parser.add_argument(
+        "-s",
+        "--similarity-data",
+        metavar="SIMILARITY DATA",
+        type=argparse.FileType("r", encoding="utf-8"),
+        const=sys.stdin,
+        nargs="?",
+        help=(
+            "List of similarity data transformers to apply to all definitions of the "
+            "offliner. This updates the similarity data array of the recipes "
+            "associated with the definitions too."
+        ),
+    )
+
+    args = parser.parse_args()
+
+    with Session.begin() as session:
+        offliner = get_offliner(session, args.offliner)
+
+        if args.similarity_data:
+            similarity_data = [
+                SimilarityDataSchema.model_validate(data)
+                for data in json.loads(args.similarity_data.read())
+            ]
+            for offliner_definition in session.scalars(
+                select(OfflinerDefinition).where(
+                    OfflinerDefinition.offliner == args.offliner
+                )
+            ):
+                # update the spec for the offliner definition to use the new
+                # similarity data transformers list
+                offliner_spec = OfflinerSpecSchema.model_validate(
+                    offliner_definition.schema
+                )
+                offliner_spec.similarity_data = similarity_data
+                # update the db offliner definition and all of its associated recipes
+                offliner_definition.schema = offliner_spec.model_dump(mode="json")
+                nb_schedules_modified = update_schedules(
+                    session,
+                    offliner=offliner,
+                    offliner_definition=create_offliner_definition_schema(
+                        offliner_definition
+                    ),
+                )
+                logger.info(f"updated {nb_schedules_modified} schedule(s) ")
+
+    logger.info("FINISH!")

backend/src/zimfarm_backend/common/schemas/offliners/builder.py

@@ -1,6 +1,7 @@
 # ruff: noqa: UP007
 from collections.abc import Callable
 from enum import Enum, StrEnum
+from itertools import chain
 from typing import Annotated, Any, Literal, Optional, cast
 
 from pydantic import (
@@ -12,6 +13,7 @@
     model_validator,
 )
 
+from zimfarm_backend import logger
 from zimfarm_backend.common.constants import getenv, parse_bool
 from zimfarm_backend.common.schemas import CamelModel, DashModel
 from zimfarm_backend.common.schemas.fields import (
@@ -30,6 +32,7 @@
     FlagSchema,
     OfflinerSpecSchema,
 )
+from zimfarm_backend.common.schemas.offliners.transformers import transform_data
 from zimfarm_backend.common.schemas.offliners.validators import (
     check_exclusive_fields,
     validate_ted_links,
@@ -169,6 +172,34 @@ def generate_field_type(offliner: str, flag: FlagSchema, label: str):
         return Annotated[py_type, pydantic_field, WrapValidator(skip_validation)]
 
 
+def generate_similarity_data(
+    flags: dict[str, Any], offliner: OfflinerSchema, spec: OfflinerSpecSchema
+) -> list[str]:
+    """Generate the similarity list of flags data."""
+    schema_cls = build_offliner_model(offliner, spec)
+    result: list[list[str]] = []
+    for similarity_data in spec.similarity_data:
+        # find the pydantic field info from the schema class
+        field_info = schema_cls.model_fields[similarity_data.flag]
+        # find the value of the flag in the data. This is typically at the alias of the
+        # field given we always dump with aliases. Sometimes, it turns out to be at the
+        # name of the python identifier because we originally dumped without aliases.
+        value: Any | None = None
+        if field_info.alias:
+            value = flags.get(field_info.alias, flags.get(similarity_data.flag))
+        else:
+            value = flags.get(similarity_data.flag)
+
+        if value is None:
+            logger.warning(
+                f"Could not find value in data that matched the keys: "
+                f"'{field_info.alias}', '{similarity_data.flag}'"
+            )
+            continue
+        result.append(transform_data([value], similarity_data.transformers))
+    return list(set(chain.from_iterable(result)))
+
+
 def build_offliner_model(
     offliner: OfflinerSchema,
     schema: OfflinerSpecSchema,

backend/src/zimfarm_backend/common/schemas/offliners/models.py

@@ -78,10 +78,29 @@ class ModelValidatorSchema(CamelModel):
     model_config = ConfigDict(extra="forbid")
 
 
+class TransformerSchema(CamelModel):
+    # the name of the transformer function to use for the field. If None, the field
+    # will be used as it is
+    name: Literal["split", "hostname"] | None = None
+    # the operand to use for the transformer function (if the function
+    # takes an operand)
+    operand: str | None = None
+
+
+class SimilarityDataSchema(CamelModel):
+    # the name of the flag to use for the similarity data
+    flag: str
+    # transformers are applied in sequential order
+    transformers: list[TransformerSchema]
+
+
 class OfflinerSpecSchema(CamelModel):
     flags: dict[str, FlagSchema]
     model_validators: list[ModelValidatorSchema] = Field(  # pyright: ignore
         default_factory=list,
     )
     std_output: str | bool = Field(default=False)
     std_stats: str | bool = Field(default=False)
+    similarity_data: list[SimilarityDataSchema] = Field(  # pyright: ignore
+        default_factory=list
+    )

backend/src/zimfarm_backend/common/schemas/offliners/transformers.py

@@ -0,0 +1,49 @@
+from collections.abc import Callable
+from functools import partial
+from itertools import chain
+from urllib.parse import urlparse
+
+from zimfarm_backend.common.schemas.offliners.models import TransformerSchema
+
+
+def get_transformer_function(
+    transformer: TransformerSchema,
+) -> Callable[[str], list[str]]:
+    def _get_hostname(url: str) -> list[str]:
+        if not url.startswith(("https://", "http://")):
+            url = "https://" + url
+        return [urlparse(url).hostname or ""]
+
+    match transformer.name:
+        case "split":
+            return partial(str.split, sep=transformer.operand)
+        case "hostname":
+            return _get_hostname
+        case None:
+            # return the value as it is if there is no transformer associated
+            return lambda x: [x]
+        case _:
+            raise ValueError(
+                f"No transformer function registered for '{transformer.name}'"
+            )
+
+
+def transform_data(data: list[str], transformers: list[TransformerSchema]) -> list[str]:
+    """Generate the output by applying each tranformer to data"""
+    if not transformers:
+        return data
+    head, *tail = transformers
+    return transform_data(
+        list(
+            chain.from_iterable(
+                # apply the transformer function to each entry in the list
+                # and feed the new list as input to the next transformer function
+                [
+                    get_transformer_function(head)(entry)
+                    for entry in data
+                    if entry.strip()
+                ]
+            )
+        ),
+        tail,
+    )

backend/src/zimfarm_backend/common/schemas/orms.py

@@ -239,6 +239,7 @@ class ScheduleFullSchema(BaseModel):
     archived: bool
     context: str
     offliner_definition_id: UUID = Field(exclude=True)
+    similarity_data: list[str] = Field(exclude=True)
     offliner: str
     version: str
 

backend/src/zimfarm_backend/db/models.py

@@ -208,6 +208,9 @@ class Schedule(Base):
     # context that a worker must have to run this schedule
     context: Mapped[str] = mapped_column(default="", server_default="", index=True)
     archived: Mapped[bool] = mapped_column(default=False, server_default=false())
+    similarity_data: Mapped[list[str]] = mapped_column(
+        default_factory=list, server_default="{}", index=True
+    )
 
     # use_alter is mandatory for alembic to break the dependency cycle
     # but it is still not totally handled automatically, the migration