Add linking to batch logic, remove required user id for batches

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

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(

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
     )

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):
 

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