diff --git a/CHANGELOG.md b/CHANGELOG.md index 72707af94f..6ba26ba760 100644 --- a/CHANGELOG.md +++ b/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) diff --git a/src/fides/api/models/detection_discovery.py b/src/fides/api/models/detection_discovery.py index c70cd7789d..f6bdcd9eb7 100644 --- a/src/fides/api/models/detection_discovery.py +++ b/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] = ( diff --git a/tests/conftest.py b/tests/conftest.py index 78432d453e..4ef33da56d 100644 --- a/tests/conftest.py +++ b/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) diff --git a/tests/ops/models/test_detection_discovery.py b/tests/ops/models/test_detection_discovery.py index 7052b3924a..6e41349d6b 100644 --- a/tests/ops/models/test_detection_discovery.py +++ b/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