Don't recompute on private->public for pro/enterprise (#3111)

* don't recompute on piblic from pro/enterprise * minor * fix tests * again
huggingface · Dec 7, 2024 · f00ed10 · f00ed10
1 parent 05cf9f4
commit f00ed10
Show file tree

Hide file tree

Showing 3 changed files with 35 additions and 12 deletions.
diff --git a/services/webhook/src/webhook/routes/webhook.py b/services/webhook/src/webhook/routes/webhook.py
@@ -40,6 +40,7 @@
 class _MoonWebhookV2PayloadRepo(TypedDict):
     type: Literal["model", "dataset", "space"]
     name: str
+    private: bool
 
 
 class MoonWebhookV2PayloadRepo(_MoonWebhookV2PayloadRepo, total=False):
@@ -82,18 +83,16 @@ def process_payload(
         # ^ it filters out the webhook calls for non-dataset repos and discussions in dataset repos
         return None
     dataset = payload["repo"]["name"]
+    private = payload["repo"]["private"]
     if dataset is None:
         return None
     event = payload["event"]
     if event == "remove":
         delete_dataset(dataset=dataset, storage_clients=storage_clients)
     elif event in ["add", "update", "move"]:
-        if (
-            event == "update"
-            and get_current_revision(dataset) == payload["repo"]["headSha"]
-            and not payload["scope"] == "repo.config"
-        ):
+        if event == "update" and get_current_revision(dataset) == payload["repo"]["headSha"] and not private:
             # ^ it filters out the webhook calls when the refs/convert/parquet branch is updated
+            # ^ it also filters switching from private to public if the headSha is in the cache (i.e. if the user is PRO/Enterprise)
             logging.warning(
                 f"Webhook revision for {dataset} is the same as the current revision in the db - skipping update."
             )

diff --git a/services/webhook/tests/routes/test_webhook.py b/services/webhook/tests/routes/test_webhook.py
@@ -82,23 +82,30 @@ def test_parse_payload(
     "payload,does_update",
     [
         (
-            {"event": "add", "repo": {"type": "dataset", "name": "webhook-test", "gitalyUid": "123"}, "scope": "repo"},
+            {
+                "event": "add",
+                "repo": {"type": "dataset", "name": "webhook-test", "gitalyUid": "123", "private": False},
+                "scope": "repo",
+            },
             True,
         ),
         (
             {
                 "event": "move",
                 "movedTo": "webhook-test",
-                "repo": {"type": "dataset", "name": "previous-name", "gitalyUid": "123"},
+                "repo": {"type": "dataset", "name": "previous-name", "gitalyUid": "123", "private": False},
                 "scope": "repo",
             },
             True,
         ),
-        ({"event": "add", "repo": {"type": "dataset", "name": "webhook-test"}, "scope": "repo"}, True),
+        (
+            {"event": "add", "repo": {"type": "dataset", "name": "webhook-test", "private": False}, "scope": "repo"},
+            True,
+        ),
         (
             {
                 "event": "doesnotexist",
-                "repo": {"type": "dataset", "name": "webhook-test", "gitalyUid": "123"},
+                "repo": {"type": "dataset", "name": "webhook-test", "gitalyUid": "123", "private": False},
                 "scope": "repo",
             },
             False,
@@ -128,7 +135,12 @@ def test_parse_payload(
         (
             {
                 "event": "update",
-                "repo": {"type": "dataset", "name": "AresEkb/prof_standards_sbert_large_mt_nlu_ru", "headSha": "abc"},
+                "repo": {
+                    "type": "dataset",
+                    "name": "AresEkb/prof_standards_sbert_large_mt_nlu_ru",
+                    "headSha": "abc",
+                    "private": False,
+                },
                 "scope": "repo.content",
             },
             True,

diff --git a/services/webhook/tests/test_app_real.py b/services/webhook/tests/test_app_real.py
@@ -46,7 +46,13 @@ def test_webhook_untrusted(
 ) -> None:
     payload = {
         "event": "add",
-        "repo": {"type": "dataset", "name": "nyu-mll/glue", "gitalyUid": "123", "headSha": "revision"},
+        "repo": {
+            "type": "dataset",
+            "name": "nyu-mll/glue",
+            "gitalyUid": "123",
+            "headSha": "revision",
+            "private": False,
+        },
         "scope": "repo",
     }
     response = real_client.post("/webhook", json=payload)
@@ -57,7 +63,13 @@ def test_webhook_untrusted(
 def test_webhook_trusted(real_client: TestClient) -> None:
     payload = {
         "event": "add",
-        "repo": {"type": "dataset", "name": "nyu-mll/glue", "gitalyUid": "123", "headSha": "revision"},
+        "repo": {
+            "type": "dataset",
+            "name": "nyu-mll/glue",
+            "gitalyUid": "123",
+            "headSha": "revision",
+            "private": False,
+        },
         "scope": "repo",
     }
     response = real_client.post("/webhook", json=payload, headers={"x-webhook-secret": API_HF_WEBHOOK_SECRET})