PROD-2352 helper: get_urn_list method (#5192)

CHANGELOG.md

@@ -23,6 +23,7 @@ The types of changes are:
 - Adds user_read scope to approver role so that they can update their own password [#5178](https://github.com/ethyca/fides/pull/5178)
 - Added PATCH endpoint for partially updating connection secrets [#5172](https://github.com/ethyca/fides/pull/5172)
 - Add success toast on confirming classification in data discovery tables [#5182](https://github.com/ethyca/fides/pull/5182)
+- Add function to return list of StagedResource objs according to list of URNs [#5192](https://github.com/ethyca/fides/pull/5192)
 - Add DSR Support for ScyllaDB [#5140](https://github.com/ethyca/fides/pull/5140)
 - Added support for nested fields in BigQuery in D&D result views [#5175](https://github.com/ethyca/fides/pull/5175)
 

src/fides/api/models/detection_discovery.py

@@ -2,10 +2,11 @@
 
 from datetime import datetime
 from enum import Enum
-from typing import Any, Dict, Iterable, Optional, Type
+from typing import Any, Dict, Iterable, List, Optional, Type
 
-from sqlalchemy import ARRAY, Boolean, Column, DateTime, ForeignKey, String
+from sqlalchemy import ARRAY, Boolean, Column, DateTime, ForeignKey, String, select
 from sqlalchemy.dialects.postgresql import JSONB
+from sqlalchemy.ext.asyncio import AsyncSession
 from sqlalchemy.ext.mutable import MutableDict
 from sqlalchemy.orm import Session, relationship
 
@@ -281,6 +282,38 @@ def get_urn(cls, db: Session, urn: str) -> Optional[StagedResource]:
         """Utility to retrieve the staged resource with the given URN"""
         return cls.get_by(db=db, field="urn", value=urn)
 
+    @classmethod
+    def get_urn_list(cls, db: Session, urns: Iterable[str]) -> Iterable[StagedResource]:
+        """
+        Utility to retrieve all staged resources with the given URNs
+        """
+        results = db.execute(select(StagedResource).where(StagedResource.urn.in_(urns)))  # type: ignore
+        return results.scalars().all()
+
+    @classmethod
+    async def get_urn_async(
+        cls, db: AsyncSession, urn: str
+    ) -> Optional[StagedResource]:
+        """
+        Utility to retrieve the staged resource with the given URN using an async session
+        """
+        results = await db.execute(
+            select(StagedResource).where(StagedResource.urn == urn)  # type: ignore
+        )
+        return results.scalars().first()
+
+    @classmethod
+    async def get_urn_list_async(
+        cls, db: AsyncSession, urns: List[str]
+    ) -> Optional[List[StagedResource]]:
+        """
+        Utility to retrieve the staged resource with the given URN using an async session
+        """
+        results = await db.execute(
+            select(StagedResource).where(StagedResource.urn.in_(urns))
+        )
+        return results.scalars().all()
+
     def add_child_diff_status(self, diff_status: DiffStatus) -> None:
         """Increments the specified child diff status"""
         self.child_diff_statuses[diff_status.value] = (

tests/conftest.py

@@ -121,7 +121,6 @@ async def async_session(test_client):
 @pytest.mark.asyncio
 async def async_session_temp(test_client):
     assert CONFIG.test_mode
-    assert requests.post == test_client.post
 
     create_citext_extension(sync_engine)
 

tests/ops/models/test_detection_discovery.py

@@ -1,6 +1,8 @@
 from datetime import datetime, timezone
+from typing import List
 
 import pytest
+from sqlalchemy.ext.asyncio import AsyncSession
 from sqlalchemy.orm import Session
 from sqlalchemy.orm.attributes import flag_modified
 
@@ -61,6 +63,30 @@ def create_staged_resource(self, db: Session):
         )
         return resource
 
+    def test_get_urn(self, db: Session, create_staged_resource) -> None:
+        urn_list = [create_staged_resource.urn]
+        from_db = StagedResource.get_urn_list(db, urn_list)
+        assert len(from_db) == len(urn_list)
+        assert from_db[0].urn == urn_list[0]
+
+        # single urn
+        from_db_single = StagedResource.get_urn(db, urn_list[0])
+        assert from_db_single.urn == urn_list[0]
+
+    async def test_get_urn_async(
+        self, async_session_temp: AsyncSession, create_staged_resource
+    ) -> None:
+        urn_list: List[str] = [str(create_staged_resource.urn)]
+
+        from_db_single = await StagedResource.get_urn_async(
+            async_session_temp, urn_list[0]
+        )
+        assert from_db_single.urn == urn_list[0]
+
+        from_db = await StagedResource.get_urn_list_async(async_session_temp, urn_list)
+        assert len(from_db) == len(urn_list)
+        assert from_db[0].urn == urn_list[0]
+
     def test_create_staged_resource(self, db: Session, create_staged_resource) -> None:
         """
         Creation fixture creates the resource, this tests that it was created successfully