diff --git a/alembic/versions/2025_09_26_1357-b9317c6836e7_add_agency_and_jurisdiction_type.py b/alembic/versions/2025_09_26_1357-b9317c6836e7_add_agency_and_jurisdiction_type.py new file mode 100644 index 00000000..7d917fbf --- /dev/null +++ b/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 diff --git a/src/api/endpoints/review/approve/query_/core.py b/src/api/endpoints/review/approve/query_/core.py index a624f53d..15641764 100644 --- a/src/api/endpoints/review/approve/query_/core.py +++ b/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( diff --git a/src/db/client/async_.py b/src/db/client/async_.py index 18ac2a29..6066a2e5 100644 --- a/src/db/client/async_.py +++ b/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 @@ async def upsert_new_agencies( 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 @@ async def add_agency_manual_suggestion( 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) diff --git a/src/db/models/impl/agency/enums.py b/src/db/models/impl/agency/enums.py new file mode 100644 index 00000000..80ed9780 --- /dev/null +++ b/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" \ No newline at end of file diff --git a/src/db/models/impl/agency/sqlalchemy.py b/src/db/models/impl/agency/sqlalchemy.py index 032dc397..20cd5f12 100644 --- a/src/db/models/impl/agency/sqlalchemy.py +++ b/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") diff --git a/tests/automated/integration/api/review/test_approve_and_get_next_source.py b/tests/automated/integration/api/review/test_approve_and_get_next_source.py index 2483921f..858df360 100644 --- a/tests/automated/integration/api/review/test_approve_and_get_next_source.py +++ b/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) diff --git a/tests/automated/integration/tasks/url/impl/submit_approved/setup.py b/tests/automated/integration/tasks/url/impl/submit_approved/setup.py index c1a1d4f4..1f9d8915 100644 --- a/tests/automated/integration/tasks/url/impl/submit_approved/setup.py +++ b/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 @@ async def setup_validated_urls(db_data_creator: DBDataCreator) -> list[str]: 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 @@ async def setup_validated_urls(db_data_creator: DBDataCreator) -> list[str]: 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 @@ async def setup_validated_urls(db_data_creator: DBDataCreator) -> list[str]: 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", ), diff --git a/tests/automated/integration/tasks/url/impl/submit_approved/test_submit_approved_url_task.py b/tests/automated/integration/tasks/url/impl/submit_approved/test_submit_approved_url_task.py index acb0005e..44b70d53 100644 --- a/tests/automated/integration/tasks/url/impl/submit_approved/test_submit_approved_url_task.py +++ b/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] } ] } diff --git a/tests/helpers/data_creator/core.py b/tests/helpers/data_creator/core.py index 0efe279d..5fb700b7 100644 --- a/tests/helpers/data_creator/core.py +++ b/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)