Make changes to collections to support publicly listed collections (#2164)

Fixes #2158 

- Adds `Organization.listPublicCollections` field and API endpoint to
update it
- Replaces `Collection.isPublic` boolean with `Collection.access`
(values: `private`, `unlisted`, `public`) and add database migration
- Update frontend to use `Collection.access` instead of `isPublic`,
otherwise not changing current behavior

---------

Co-authored-by: sua yoo <[email protected]>

Author: tw4l

backend/btrixcloud/colls.py

@@ -89,7 +89,7 @@ async def add_collection(self, oid: UUID, coll_in: CollIn):
             name=coll_in.name,
             description=coll_in.description,
             modified=modified,
-            isPublic=coll_in.isPublic,
+            access=coll_in.access,
         )
         try:
             await self.collections.insert_one(coll.to_dict())
@@ -189,7 +189,7 @@ async def get_collection(
         """Get collection by id"""
         query: dict[str, object] = {"_id": coll_id}
         if public_only:
-            query["isPublic"] = True
+            query["access"] = {"$in": ["public", "unlisted"]}
 
         result = await self.collections.find_one(query)
         if not result:
@@ -210,6 +210,7 @@ async def list_collections(
         sort_direction: int = 1,
         name: Optional[str] = None,
         name_prefix: Optional[str] = None,
+        access: Optional[str] = None,
     ):
         """List all collections for org"""
         # pylint: disable=too-many-locals, duplicate-code
@@ -226,6 +227,9 @@ async def list_collections(
             regex_pattern = f"^{name_prefix}"
             match_query["name"] = {"$regex": regex_pattern, "$options": "i"}
 
+        if access:
+            match_query["access"] = access
+
         aggregate = [{"$match": match_query}]
 
         if sort_by:
@@ -427,6 +431,7 @@ async def list_collection_all(
         sortDirection: int = 1,
         name: Optional[str] = None,
         namePrefix: Optional[str] = None,
+        access: Optional[str] = None,
     ):
         collections, total = await colls.list_collections(
             org.id,
@@ -436,6 +441,7 @@ async def list_collection_all(
             sort_direction=sortDirection,
             name=name,
             name_prefix=namePrefix,
+            access=access,
         )
         return paginated_format(collections, total, page, pageSize)
 

backend/btrixcloud/db.py

@@ -17,7 +17,7 @@
 from .migrations import BaseMigration
 
 
-CURR_DB_VERSION = "0035"
+CURR_DB_VERSION = "0036"
 
 
 # ============================================================================

backend/btrixcloud/migrations/migration_0036_coll_visibility.py

@@ -0,0 +1,49 @@
+"""
+Migration 0036 -- collection access
+"""
+
+from btrixcloud.migrations import BaseMigration
+
+
+MIGRATION_VERSION = "0036"
+
+
+class Migration(BaseMigration):
+    """Migration class."""
+
+    # pylint: disable=unused-argument
+    def __init__(self, mdb, **kwargs):
+        super().__init__(mdb, migration_version=MIGRATION_VERSION)
+
+    async def migrate_up(self):
+        """Perform migration up.
+
+        Move from Collection.isPublic cool to Collection.access enum
+        """
+        colls_mdb = self.mdb["collections"]
+
+        # Set non-public collections to private
+        try:
+            await colls_mdb.update_many(
+                {"isPublic": False},
+                {"$set": {"access": "private"}, "$unset": {"isPublic": 1}},
+            )
+        # pylint: disable=broad-exception-caught
+        except Exception as err:
+            print(
+                f"Error migrating private collections: {err}",
+                flush=True,
+            )
+
+        # Set public collections to unlisted
+        try:
+            await colls_mdb.update_many(
+                {"isPublic": True},
+                {"$set": {"access": "unlisted"}, "$unset": {"isPublic": 1}},
+            )
+        # pylint: disable=broad-exception-caught
+        except Exception as err:
+            print(
+                f"Error migrating public unlisted collections: {err}",
+                flush=True,
+            )

backend/btrixcloud/models.py

@@ -1055,6 +1055,15 @@ class UpdateUpload(UpdateCrawl):
 ### COLLECTIONS ###
 
 
+# ============================================================================
+class CollAccessType(str, Enum):
+    """Collection access types"""
+
+    PRIVATE = "private"
+    UNLISTED = "unlisted"
+    PUBLIC = "public"
+
+
 # ============================================================================
 class Collection(BaseMongoModel):
     """Org collection structure"""
@@ -1071,7 +1080,7 @@ class Collection(BaseMongoModel):
     # Sorted by count, descending
     tags: Optional[List[str]] = []
 
-    isPublic: Optional[bool] = False
+    access: CollAccessType = CollAccessType.PRIVATE
 
 
 # ============================================================================
@@ -1082,7 +1091,7 @@ class CollIn(BaseModel):
     description: Optional[str] = None
     crawlIds: Optional[List[str]] = []
 
-    isPublic: bool = False
+    access: CollAccessType = CollAccessType.PRIVATE
 
 
 # ============================================================================
@@ -1098,7 +1107,7 @@ class UpdateColl(BaseModel):
 
     name: Optional[str] = None
     description: Optional[str] = None
-    isPublic: Optional[bool] = None
+    access: Optional[CollAccessType] = None
 
 
 # ============================================================================
@@ -1371,6 +1380,13 @@ class OrgReadOnlyUpdate(BaseModel):
     readOnlyReason: Optional[str] = None
 
 
+# ============================================================================
+class OrgListPublicCollectionsUpdate(BaseModel):
+    """Organization listPublicCollections update"""
+
+    listPublicCollections: bool
+
+
 # ============================================================================
 class OrgWebhookUrls(BaseModel):
     """Organization webhook URLs"""
@@ -1439,6 +1455,8 @@ class OrgOut(BaseMongoModel):
     allowedProxies: list[str] = []
     crawlingDefaults: Optional[CrawlConfigDefaults] = None
 
+    listPublicCollections: bool = False
+
 
 # ============================================================================
 class Organization(BaseMongoModel):
@@ -1494,6 +1512,8 @@ class Organization(BaseMongoModel):
     allowedProxies: list[str] = []
     crawlingDefaults: Optional[CrawlConfigDefaults] = None
 
+    listPublicCollections: bool = False
+
     def is_owner(self, user):
         """Check if user is owner"""
         return self._is_auth(user, UserRole.OWNER)

backend/btrixcloud/orgs.py

@@ -78,6 +78,7 @@
     RemovedResponse,
     OrgSlugsResponse,
     OrgImportResponse,
+    OrgListPublicCollectionsUpdate,
 )
 from .pagination import DEFAULT_PAGE_SIZE, paginated_format
 from .utils import (
@@ -940,7 +941,7 @@ async def get_org_metrics(self, org: Organization) -> dict[str, int]:
         )
         collections_count = await self.colls_db.count_documents({"oid": org.id})
         public_collections_count = await self.colls_db.count_documents(
-            {"oid": org.id, "isPublic": True}
+            {"oid": org.id, "access": {"$in": ["public", "unlisted"]}}
         )
 
         return {
@@ -997,6 +998,16 @@ async def update_read_only_on_cancel(
         )
         return res is not None
 
+    async def update_list_public_collections(
+        self, org: Organization, list_public_collections: bool
+    ):
+        """Update listPublicCollections field on organization"""
+        res = await self.orgs.find_one_and_update(
+            {"_id": org.id},
+            {"$set": {"listPublicCollections": list_public_collections}},
+        )
+        return res is not None
+
     async def export_org(
         self, org: Organization, user_manager: UserManager
     ) -> StreamingResponse:
@@ -1553,6 +1564,19 @@ async def update_read_only_on_cancel(
 
         return {"updated": True}
 
+    @router.post(
+        "/list-public-collections",
+        tags=["organizations", "collections"],
+        response_model=UpdatedResponse,
+    )
+    async def update_list_public_collections(
+        update: OrgListPublicCollectionsUpdate,
+        org: Organization = Depends(org_owner_dep),
+    ):
+        await ops.update_list_public_collections(org, update.listPublicCollections)
+
+        return {"updated": True}
+
     @router.post(
         "/event-webhook-urls", tags=["organizations"], response_model=UpdatedResponse
     )

backend/test/test_collections.py

@@ -58,7 +58,7 @@ def test_create_public_collection(
         json={
             "crawlIds": [crawler_crawl_id],
             "name": PUBLIC_COLLECTION_NAME,
-            "isPublic": True,
+            "access": "public",
         },
     )
     assert r.status_code == 200
@@ -73,7 +73,7 @@ def test_create_public_collection(
         f"{API_PREFIX}/orgs/{default_org_id}/collections/{_public_coll_id}",
         headers=crawler_auth_headers,
     )
-    assert r.json()["isPublic"]
+    assert r.json()["access"] == "public"
 
 
 def test_create_collection_taken_name(
@@ -311,12 +311,31 @@ def test_collection_public(crawler_auth_headers, default_org_id):
     )
     assert r.status_code == 404
 
-    # make public
+    # make public and test replay headers
     r = requests.patch(
         f"{API_PREFIX}/orgs/{default_org_id}/collections/{_coll_id}",
         headers=crawler_auth_headers,
         json={
-            "isPublic": True,
+            "access": "public",
+        },
+    )
+    assert r.status_code == 200
+    assert r.json()["updated"]
+
+    r = requests.get(
+        f"{API_PREFIX}/orgs/{default_org_id}/collections/{_coll_id}/public/replay.json",
+        headers=crawler_auth_headers,
+    )
+    assert r.status_code == 200
+    assert r.headers["Access-Control-Allow-Origin"] == "*"
+    assert r.headers["Access-Control-Allow-Headers"] == "*"
+
+    # make unlisted and test replay headers
+    r = requests.patch(
+        f"{API_PREFIX}/orgs/{default_org_id}/collections/{_coll_id}",
+        headers=crawler_auth_headers,
+        json={
+            "access": "unlisted",
         },
     )
     assert r.status_code == 200
@@ -335,7 +354,7 @@ def test_collection_public(crawler_auth_headers, default_org_id):
         f"{API_PREFIX}/orgs/{default_org_id}/collections/{_coll_id}",
         headers=crawler_auth_headers,
         json={
-            "isPublic": False,
+            "access": "private",
         },
     )
 
@@ -346,6 +365,24 @@ def test_collection_public(crawler_auth_headers, default_org_id):
     assert r.status_code == 404
 
 
+def test_collection_access_invalid_value(crawler_auth_headers, default_org_id):
+    r = requests.patch(
+        f"{API_PREFIX}/orgs/{default_org_id}/collections/{_coll_id}",
+        headers=crawler_auth_headers,
+        json={
+            "access": "invalid",
+        },
+    )
+    assert r.status_code == 422
+
+    r = requests.get(
+        f"{API_PREFIX}/orgs/{default_org_id}/collections/{_coll_id}",
+        headers=crawler_auth_headers,
+    )
+    assert r.status_code == 200
+    assert r.json()["access"] == "private"
+
+
 def test_add_upload_to_collection(crawler_auth_headers, default_org_id):
     with open(os.path.join(curr_dir, "data", "example.wacz"), "rb") as fh:
         r = requests.put(
@@ -429,6 +466,7 @@ def test_list_collections(
     assert first_coll["totalSize"] > 0
     assert first_coll["modified"]
     assert first_coll["tags"] == ["wr-test-2", "wr-test-1"]
+    assert first_coll["access"] == "private"
 
     second_coll = [coll for coll in items if coll["name"] == SECOND_COLLECTION_NAME][0]
     assert second_coll["id"]
@@ -440,6 +478,7 @@ def test_list_collections(
     assert second_coll["totalSize"] > 0
     assert second_coll["modified"]
     assert second_coll["tags"] == ["wr-test-2"]
+    assert second_coll["access"] == "private"
 
 
 def test_remove_upload_from_collection(crawler_auth_headers, default_org_id):
@@ -525,6 +564,26 @@ def test_filter_sort_collections(
     assert coll["oid"] == default_org_id
     assert coll.get("description") is None
 
+    # Test filtering by access
+    name_prefix = name_prefix.upper()
+    r = requests.get(
+        f"{API_PREFIX}/orgs/{default_org_id}/collections?access=public",
+        headers=crawler_auth_headers,
+    )
+    assert r.status_code == 200
+    data = r.json()
+    assert data["total"] == 1
+
+    items = data["items"]
+    assert len(items) == 1
+
+    coll = items[0]
+    assert coll["id"]
+    assert coll["name"] == PUBLIC_COLLECTION_NAME
+    assert coll["oid"] == default_org_id
+    assert coll.get("description") is None
+    assert coll["access"] == "public"
+
     # Test sorting by name, ascending (default)
     r = requests.get(
         f"{API_PREFIX}/orgs/{default_org_id}/collections?sortBy=name",

backend/test/test_org.py

@@ -697,6 +697,24 @@ def test_update_read_only(admin_auth_headers, default_org_id):
     assert data["readOnlyReason"] == ""
 
 
+def test_update_list_public_collections(admin_auth_headers, default_org_id):
+    # Test that default is false
+    r = requests.get(f"{API_PREFIX}/orgs/{default_org_id}", headers=admin_auth_headers)
+    assert r.json()["listPublicCollections"] is False
+
+    # Update
+    r = requests.post(
+        f"{API_PREFIX}/orgs/{default_org_id}/list-public-collections",
+        headers=admin_auth_headers,
+        json={"listPublicCollections": True},
+    )
+    assert r.json()["updated"]
+
+    # Test update is reflected in GET response
+    r = requests.get(f"{API_PREFIX}/orgs/{default_org_id}", headers=admin_auth_headers)
+    assert r.json()["listPublicCollections"]
+
+
 def test_sort_orgs(admin_auth_headers):
     # Create a few new orgs for testing
     r = requests.post(