From f24f6c4d74a153b21b87ec915b3788ae1de3c21d Mon Sep 17 00:00:00 2001 From: Max Chis Date: Tue, 21 Oct 2025 11:25:58 -0400 Subject: [PATCH] Add linking to batch logic, remove required user id for batches --- ...4e9f_set_batches_user_id_to_be_nullable.py | 30 +++++++++++++++++++ src/api/endpoints/submit/data_source/query.py | 20 +++++++++++++ src/db/models/impl/batch/sqlalchemy.py | 7 +++-- .../integration/api/batch/test_batch.py | 5 +--- .../api/submit/data_source/test_core.py | 18 ++++++++++- 5 files changed, 72 insertions(+), 8 deletions(-) create mode 100644 alembic/versions/2025_10_21_1123-f32ba7664e9f_set_batches_user_id_to_be_nullable.py diff --git a/alembic/versions/2025_10_21_1123-f32ba7664e9f_set_batches_user_id_to_be_nullable.py b/alembic/versions/2025_10_21_1123-f32ba7664e9f_set_batches_user_id_to_be_nullable.py new file mode 100644 index 00000000..d6076e7a --- /dev/null +++ b/alembic/versions/2025_10_21_1123-f32ba7664e9f_set_batches_user_id_to_be_nullable.py @@ -0,0 +1,30 @@ +"""Set batches.user_id to be nullable + +Revision ID: f32ba7664e9f +Revises: 6adf9d894180 +Create Date: 2025-10-21 11:23:35.611484 + +""" +from typing import Sequence, Union + +from alembic import op +import sqlalchemy as sa + + +# revision identifiers, used by Alembic. +revision: str = 'f32ba7664e9f' +down_revision: Union[str, None] = '6adf9d894180' +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + op.alter_column( + table_name='batches', + column_name='user_id', + nullable=True + ) + + +def downgrade() -> None: + pass diff --git a/src/api/endpoints/submit/data_source/query.py b/src/api/endpoints/submit/data_source/query.py index 405390d3..6d7360f5 100644 --- a/src/api/endpoints/submit/data_source/query.py +++ b/src/api/endpoints/submit/data_source/query.py @@ -5,6 +5,9 @@ from src.api.endpoints.submit.data_source.request import DataSourceSubmissionRequest from src.api.endpoints.submit.data_source.response import SubmitDataSourceURLProposalResponse from src.collectors.enums import URLStatus +from src.core.enums import BatchStatus +from src.db.models.impl.batch.sqlalchemy import Batch +from src.db.models.impl.link.batch_url.sqlalchemy import LinkBatchURL from src.db.models.impl.url.core.enums import URLSource from src.db.models.impl.url.core.sqlalchemy import URL from src.db.models.impl.url.optional_ds_metadata.sqlalchemy import URLOptionalDataSourceMetadata @@ -40,6 +43,23 @@ async def run(self, session: AsyncSession) -> Any: url_id: int = url.id + # Add Batch + batch = Batch( + strategy='manual', + status=BatchStatus.READY_TO_LABEL, + parameters={} + ) + session.add(batch) + await session.flush() + batch_id: int = batch.id + + # Add Batch URL link + batch_url_link = LinkBatchURL( + batch_id=batch_id, + url_id=url_id + ) + session.add(batch_url_link) + # Optionally add Record Type as suggestion if self.request.record_type is not None: record_type_suggestion = AnonymousAnnotationRecordType( diff --git a/src/db/models/impl/batch/sqlalchemy.py b/src/db/models/impl/batch/sqlalchemy.py index 564ce163..fb44396b 100644 --- a/src/db/models/impl/batch/sqlalchemy.py +++ b/src/db/models/impl/batch/sqlalchemy.py @@ -1,7 +1,8 @@ from sqlalchemy import Column, Integer, TIMESTAMP, Float, JSON from sqlalchemy.dialects import postgresql -from sqlalchemy.orm import relationship +from sqlalchemy.orm import relationship, Mapped +from src.core.enums import BatchStatus from src.db.models.helpers import CURRENT_TIME_SERVER_DEFAULT from src.db.models.impl.log.sqlalchemy import Log from src.db.models.templates_.with_id import WithIDBase @@ -23,9 +24,9 @@ class Batch(WithIDBase): 'manual', name='batch_strategy'), nullable=False) - user_id = Column(Integer, nullable=False) + user_id = Column(Integer, nullable=True) # Gives the status of the batch - status = Column( + status: Mapped[BatchStatus] = Column( batch_status_enum, nullable=False ) diff --git a/tests/automated/integration/api/batch/test_batch.py b/tests/automated/integration/api/batch/test_batch.py index f1e3d4f2..f34928d6 100644 --- a/tests/automated/integration/api/batch/test_batch.py +++ b/tests/automated/integration/api/batch/test_batch.py @@ -1,8 +1,5 @@ -from src.api.endpoints.batch.dtos.get.summaries.summary import BatchSummary -from src.db.models.impl.batch.pydantic.info import BatchInfo from src.db.dtos.url.insert import InsertURLsInfo -from src.collectors.impl.example.dtos.input import ExampleInputDTO -from src.core.enums import BatchStatus + def test_get_batch_urls(api_test_helper): diff --git a/tests/automated/integration/api/submit/data_source/test_core.py b/tests/automated/integration/api/submit/data_source/test_core.py index 566ff60a..49df1dd4 100644 --- a/tests/automated/integration/api/submit/data_source/test_core.py +++ b/tests/automated/integration/api/submit/data_source/test_core.py @@ -4,8 +4,10 @@ from src.api.endpoints.submit.data_source.request import DataSourceSubmissionRequest from src.collectors.enums import URLStatus -from src.core.enums import RecordType +from src.core.enums import RecordType, BatchStatus from src.db.client.async_ import AsyncDatabaseClient +from src.db.models.impl.batch.sqlalchemy import Batch +from src.db.models.impl.link.batch_url.sqlalchemy import LinkBatchURL from src.db.models.impl.url.core.enums import URLSource from src.db.models.impl.url.core.sqlalchemy import URL from src.db.models.impl.url.optional_ds_metadata.enums import AgencyAggregationEnum, UpdateMethodEnum, \ @@ -73,6 +75,20 @@ async def test_submit_data_source( assert url.source == URLSource.MANUAL assert url.status == URLStatus.OK + # Check for Batch + batch: Batch = await adb_client.one_or_none_model(Batch) + assert batch is not None + assert batch.user_id is None + assert batch.strategy == 'manual' + assert batch.status == BatchStatus.READY_TO_LABEL.value + assert batch.parameters == {} + + # Check for Batch URL Link + batch_url_link: LinkBatchURL = await adb_client.one_or_none_model(LinkBatchURL) + assert batch_url_link is not None + assert batch_url_link.batch_id == batch.id + assert batch_url_link.url_id == url.id + # Check for Location Suggestion location_suggestion: AnonymousAnnotationLocation = await adb_client.one_or_none_model(AnonymousAnnotationLocation) assert location_suggestion is not None