mc_425_agency_search_endpoint

alembic/versions/2025_09_26_1357-b9317c6836e7_add_agency_and_jurisdiction_type.py

@@ -0,0 +1,67 @@
+"""Add agency and jurisdiction type
+
+Revision ID: b9317c6836e7
+Revises: 7b955c783e27
+Create Date: 2025-09-26 13:57:42.357788
+
+"""
+from typing import Sequence, Union
+
+from alembic import op
+import sqlalchemy as sa
+
+
+# revision identifiers, used by Alembic.
+revision: str = 'b9317c6836e7'
+down_revision: Union[str, None] = '7b955c783e27'
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def _add_agency_type_column():
+    agency_type_enum = sa.Enum(
+                "unknown",
+                "incarceration",
+                "law enforcement",
+                "court",
+                "aggregated",
+                name="agency_type_enum",
+                create_type=True,
+            )
+    agency_type_enum.create(op.get_bind())
+
+    op.add_column(
+        table_name="agencies",
+        column=sa.Column(
+            "agency_type",
+            agency_type_enum,
+            server_default="unknown",
+            nullable=False,
+        )
+    )
+
+
+def _add_jurisdiction_type_column():
+    jurisdiction_type_enum = sa.Enum(
+                'school', 'county', 'local', 'port', 'tribal', 'transit', 'state', 'federal',
+                name="jurisdiction_type_enum",
+            )
+    jurisdiction_type_enum.create(op.get_bind())
+
+    op.add_column(
+        table_name="agencies",
+        column=sa.Column(
+            "jurisdiction_type",
+            jurisdiction_type_enum,
+            nullable=True,
+        )
+    )
+
+
+def upgrade() -> None:
+    _add_agency_type_column()
+    _add_jurisdiction_type_column()
+
+
+def downgrade() -> None:
+    pass

src/api/endpoints/review/approve/query_/core.py

@@ -148,12 +148,12 @@ async def _add_new_agencies(self, existing_agency_ids, new_agency_ids, session):
             existing_agency = await session.execute(query)
             existing_agency = existing_agency.scalars().first()
             if existing_agency is None:
-                # If not, create it
-                agency = Agency(
-                    agency_id=new_agency_id,
-                    name=PLACEHOLDER_AGENCY_NAME,
+                # If not, raise an error
+                raise HTTPException(
+                    status_code=HTTP_400_BAD_REQUEST,
+                    detail="Agency not found"
                 )
-                session.add(agency)
+
 
             # If the new agency id is not in the existing agency ids, add it
             confirmed_url_agency = LinkURLAgency(

src/db/client/async_.py

@@ -75,6 +75,7 @@
 from src.db.dtos.url.raw_html import RawHTMLInfo
 from src.db.enums import TaskType
 from src.db.helpers.session import session_helper as sh
+from src.db.models.impl.agency.enums import AgencyType, JurisdictionType
 from src.db.models.impl.agency.sqlalchemy import Agency
 from src.db.models.impl.backlog_snapshot import BacklogSnapshot
 from src.db.models.impl.batch.pydantic.info import BatchInfo
@@ -652,6 +653,7 @@
             agency.state = suggestion.state
             agency.county = suggestion.county
             agency.locality = suggestion.locality
+            agency.agency_type = AgencyType.UNKNOWN
             session.add(agency)
 
     @session_manager
@@ -686,7 +688,8 @@
             if len(result.all()) == 0:
                 agency = Agency(
                     agency_id=agency_id,
-                    name=PLACEHOLDER_AGENCY_NAME
+                    name=PLACEHOLDER_AGENCY_NAME,
+                    agency_type=AgencyType.UNKNOWN,
                 )
                 await session.merge(agency)
 

src/db/models/impl/agency/enums.py

@@ -0,0 +1,19 @@
+from enum import Enum
+
+
+class AgencyType(Enum):
+    UNKNOWN = "unknown"
+    INCARCERATION = "incarceration"
+    LAW_ENFORCEMENT = "law enforcement"
+    COURT = "court"
+    AGGREGATED = "aggregated"
+
+class JurisdictionType(Enum):
+    SCHOOL = "school"
+    COUNTY = "county"
+    LOCAL = "local"
+    PORT = "port"
+    TRIBAL = "tribal"
+    TRANSIT = "transit"
+    STATE = "state"
+    FEDERAL = "federal"

src/db/models/impl/agency/sqlalchemy.py

@@ -5,6 +5,8 @@
 from sqlalchemy import Column, Integer, String, DateTime
 from sqlalchemy.orm import relationship
 
+from src.db.models.helpers import enum_column
+from src.db.models.impl.agency.enums import AgencyType, JurisdictionType
 from src.db.models.mixins import UpdatedAtMixin, CreatedAtMixin
 from src.db.models.templates_.with_id import WithIDBase
 
@@ -23,6 +25,8 @@ class Agency(
     state = Column(String, nullable=True)
     county = Column(String, nullable=True)
     locality = Column(String, nullable=True)
+    agency_type = enum_column(AgencyType, name="agency_type_enum")
+    jurisdiction_type = enum_column(JurisdictionType, name="jurisdiction_type_enum")
 
     # Relationships
     automated_suggestions = relationship("AgencyIDSubtaskSuggestion")

tests/automated/integration/api/review/test_approve_and_get_next_source.py

@@ -29,11 +29,8 @@ async def test_approve_and_get_next_source_for_review(api_test_helper):
     # Add confirmed agency
     await db_data_creator.confirmed_suggestions([url_mapping.url_id])
 
-    # Additionally, include an agency not yet included in the database
-    additional_agency = 999999
 
     agency_ids = [await db_data_creator.agency() for _ in range(3)]
-    agency_ids.append(additional_agency)
 
     result: GetNextURLForFinalReviewOuterResponse = await ath.request_validator.approve_and_get_next_source_for_review(
         approval_info=FinalReviewApprovalInfo(
@@ -73,15 +70,10 @@ async def test_approve_and_get_next_source_for_review(api_test_helper):
 
     # Get agencies
     confirmed_agencies = await adb_client.get_all(LinkURLAgency)
-    assert len(confirmed_agencies) == 4
+    assert len(confirmed_agencies) == 3
     for agency in confirmed_agencies:
         assert agency.agency_id in agency_ids
 
-    # Check that created agency has placeholder
-    agencies = await adb_client.get_all(Agency)
-    for agency in agencies:
-        if agency.agency_id == additional_agency:
-            assert agency.name == PLACEHOLDER_AGENCY_NAME
 
     # Confirm presence of FlagURLValidated
     flag_url_validated = await adb_client.get_all(FlagURLValidated)

tests/automated/integration/tasks/url/impl/submit_approved/setup.py

@@ -4,7 +4,7 @@
 from tests.helpers.data_creator.models.creation_info.batch.v1 import BatchURLCreationInfo
 
 
-async def setup_validated_urls(db_data_creator: DBDataCreator) -> list[str]:
+async def setup_validated_urls(db_data_creator: DBDataCreator, agency_id: int) -> list[str]:
     creation_info: BatchURLCreationInfo = await db_data_creator.batch_and_urls(
         url_count=3,
         with_html_content=True
@@ -17,7 +17,7 @@
         approval_info=FinalReviewApprovalInfo(
             url_id=url_1,
             record_type=RecordType.ACCIDENT_REPORTS,
-            agency_ids=[1, 2],
+            agency_ids=[agency_id],
             name="URL 1 Name",
             description=None,
             record_formats=["Record Format 1", "Record Format 2"],
@@ -30,7 +30,7 @@
         approval_info=FinalReviewApprovalInfo(
             url_id=url_2,
             record_type=RecordType.INCARCERATION_RECORDS,
-            agency_ids=[3, 4],
+            agency_ids=[agency_id],
             name="URL 2 Name",
             description="URL 2 Description",
         ),
@@ -40,7 +40,7 @@
         approval_info=FinalReviewApprovalInfo(
             url_id=url_3,
             record_type=RecordType.ACCIDENT_REPORTS,
-            agency_ids=[5, 6],
+            agency_ids=[agency_id],
             name="URL 3 Name",
             description="URL 3 Description",
         ),

tests/automated/integration/tasks/url/impl/submit_approved/test_submit_approved_url_task.py

@@ -37,7 +37,8 @@ async def test_submit_approved_url_task(
 
     # Create URLs with status 'validated' in database and all requisite URL values
     # Ensure they have optional metadata as well
-    urls: list[str] = await setup_validated_urls(db_data_creator)
+    agency_id = await db_data_creator.agency()
+    urls: list[str] = await setup_validated_urls(db_data_creator, agency_id=agency_id)
     mock_make_request(mock_pdap_client, urls)
 
     # Check Task Operator does meet pre-requisites
@@ -107,7 +108,7 @@ async def test_submit_approved_url_task(
                     "data_portal_type": "Data Portal Type 1",
                     "last_approval_editor": 1,
                     "supplying_entity": "Supplying Entity 1",
-                    "agency_ids": [1, 2]
+                    "agency_ids": [agency_id]
                 },
                 {
                     "name": "URL 2 Name",
@@ -118,7 +119,7 @@ async def test_submit_approved_url_task(
                     "supplying_entity": None,
                     "record_formats": None,
                     "data_portal_type": None,
-                    "agency_ids": [3, 4]
+                    "agency_ids": [agency_id]
                 },
                 {
                     "name": "URL 3 Name",
@@ -129,7 +130,7 @@ async def test_submit_approved_url_task(
                     "supplying_entity": None,
                     "record_formats": None,
                     "data_portal_type": None,
-                    "agency_ids": [5, 6]
+                    "agency_ids": [agency_id]
                 }
             ]
         }

tests/helpers/data_creator/core.py

@@ -12,6 +12,7 @@
 from src.db.dtos.url.insert import InsertURLsInfo
 from src.db.dtos.url.mapping import URLMapping
 from src.db.enums import TaskType
+from src.db.models.impl.agency.enums import AgencyType
 from src.db.models.impl.agency.sqlalchemy import Agency
 from src.db.models.impl.duplicate.pydantic.insert import DuplicateInsertInfo
 from src.db.models.impl.flag.root_url.sqlalchemy import FlagRootURL
@@ -514,7 +515,8 @@ async def create_agency(self, agency_id: int = 1) -> None:
             name=generate_test_name(agency_id),
             state=None,
             county=None,
-            locality=None
+            locality=None,
+            agency_type=AgencyType.UNKNOWN
         )
         await self.adb_client.add_all([agency])
 
@@ -528,7 +530,8 @@ async def create_agencies(self, count: int = 3) -> list[int]:
                 name=generate_test_name(agency_id),
                 state=None,
                 county=None,
-                locality=None
+                locality=None,
+                agency_type=AgencyType.UNKNOWN
             )
             agencies.append(agency)
             agency_ids.append(agency_id)