Batch Comparison: Consider extra batches shared across contests

server/api/batches.py

@@ -5,7 +5,7 @@
 from flask import jsonify, request, session
 from werkzeug.exceptions import BadRequest, Conflict
 from sqlalchemy.orm import Query, joinedload
-from sqlalchemy import func
+from sqlalchemy import func, or_
 
 from . import api
 from ..auth import get_loggedin_user, get_support_user, restrict_access, UserType
@@ -46,14 +46,21 @@ def replace_combined_batches_with_representative_batches(
 
 
 def already_audited_batches(jurisdiction: Jurisdiction, round: Round) -> Query:
-    query: Query = (
+    sampled_in_prior_rounds = (
         Batch.query.filter_by(jurisdiction_id=jurisdiction.id)
         .join(SampledBatchDraw)
         .join(Round)
         .filter(Round.round_num < round.round_num)
         .with_entities(Batch.id)
-        .subquery()
     )
+    extra_in_prior_rounds = (
+        Batch.query.filter_by(jurisdiction_id=jurisdiction.id)
+        .join(ExtraBatchDraw)
+        .join(Round)
+        .filter(Round.round_num < round.round_num)
+        .with_entities(Batch.id)
+    )
+    query: Query = sampled_in_prior_rounds.union(extra_in_prior_rounds).subquery()
     return query
 
 
@@ -67,8 +74,12 @@ def get_batch_retrieval_list(
 ):
     batches = (
         Batch.query.filter_by(jurisdiction_id=jurisdiction.id)
-        .join(SampledBatchDraw)
-        .filter_by(round_id=round.id)
+        .filter(
+            or_(
+                Batch.draws.any(SampledBatchDraw.round_id == round.id),
+                Batch.extra_draws.any(ExtraBatchDraw.round_id == round.id),
+            )
+        )
         .filter(Batch.id.notin_(already_audited_batches(jurisdiction, round)))
         .group_by(Batch.id)
         .order_by(func.human_sort(Batch.name))
@@ -139,8 +150,12 @@ def list_batches_for_jurisdiction(
 ):
     batches = (
         Batch.query.filter_by(jurisdiction_id=jurisdiction.id)
-        .join(SampledBatchDraw)
-        .filter_by(round_id=round.id)
+        .filter(
+            or_(
+                Batch.draws.any(SampledBatchDraw.round_id == round.id),
+                Batch.extra_draws.any(ExtraBatchDraw.round_id == round.id),
+            )
+        )
         .filter(Batch.id.notin_(already_audited_batches(jurisdiction, round)))
         .order_by(func.human_sort(Batch.name))
         .options(
@@ -366,8 +381,12 @@ def finalize_batch_results(
 
     num_batches_without_results = (
         Batch.query.filter_by(jurisdiction_id=jurisdiction.id)
-        .join(SampledBatchDraw)
-        .filter_by(round_id=round.id)
+        .filter(
+            or_(
+                Batch.draws.any(SampledBatchDraw.round_id == round.id),
+                Batch.extra_draws.any(ExtraBatchDraw.round_id == round.id),
+            )
+        )
         .outerjoin(BatchResultTallySheet)
         .group_by(Batch.id)
         .having(func.count(BatchResultTallySheet.batch_id) == 0)

server/api/discrepancies.py

@@ -74,6 +74,12 @@ def get_batch_comparison_discrepancies_by_jurisdiction(
         .join(Jurisdiction)
         .with_entities(Jurisdiction.name, Batch.name)
         .all()
+    ) | set(
+        ExtraBatchDraw.query.filter_by(round_id=round_id)
+        .join(Batch)
+        .join(Jurisdiction)
+        .with_entities(Jurisdiction.name, Batch.name)
+        .all()
     )
 
     jurisdiction_name_to_id = dict(

server/api/jurisdictions.py

@@ -501,6 +501,22 @@ def batch_round_status(election: Election, round: Round) -> dict[str, JSONDict]:
         .values(Batch.jurisdiction_id, func.count(Batch.id.distinct()))
     )
 
+    extra_batch_count_by_jurisdiction = dict(
+        Batch.query.join(ExtraBatchDraw)
+        .filter(ExtraBatchDraw.round_id == round.id)
+        .group_by(Batch.jurisdiction_id)
+        .values(Batch.jurisdiction_id, func.count(Batch.id))
+    )
+
+    audited_extra_batch_count_by_jurisdiction = dict(
+        Batch.query.join(ExtraBatchDraw)
+        .filter(ExtraBatchDraw.round_id == round.id)
+        .join(BatchResultTallySheet)
+        .group_by(Batch.jurisdiction_id)
+        .having(func.count(BatchResultTallySheet.id) > 0)
+        .values(Batch.jurisdiction_id, func.count(Batch.id.distinct()))
+    )
+
     finalized_jurisdiction_ids = {
         jurisdiction_id
         for (jurisdiction_id,) in BatchResultsFinalized.query.filter_by(
@@ -509,16 +525,24 @@ def batch_round_status(election: Election, round: Round) -> dict[str, JSONDict]:
     }
 
     def num_samples(jurisdiction_id: str) -> int:
-        return sample_count_by_jurisdiction.get(jurisdiction_id, 0)
+        return sample_count_by_jurisdiction.get(
+            jurisdiction_id, 0
+        ) + extra_batch_count_by_jurisdiction.get(jurisdiction_id, 0)
 
     def num_samples_audited(jurisdiction_id: str) -> int:
-        return audited_sample_count_by_jurisdiction.get(jurisdiction_id, 0)
+        return audited_sample_count_by_jurisdiction.get(
+            jurisdiction_id, 0
+        ) + audited_extra_batch_count_by_jurisdiction.get(jurisdiction_id, 0)
 
     def num_batches(jurisdiction_id: str) -> int:
-        return batch_count_by_jurisdiction.get(jurisdiction_id, 0)
+        return batch_count_by_jurisdiction.get(
+            jurisdiction_id, 0
+        ) + extra_batch_count_by_jurisdiction.get(jurisdiction_id, 0)
 
     def num_batches_audited(jurisdiction_id: str) -> int:
-        return audited_batch_count_by_jurisdiction.get(jurisdiction_id, 0)
+        return audited_batch_count_by_jurisdiction.get(
+            jurisdiction_id, 0
+        ) + audited_extra_batch_count_by_jurisdiction.get(jurisdiction_id, 0)
 
     # NOT_STARTED = the jurisdiction hasn’t audited any batches yet
     # IN_PROGRESS = the jurisdiction is auditing batches

server/api/reports.py

@@ -519,7 +519,6 @@ def round_rows(election: Election):
                 .join(Jurisdiction)
                 .filter_by(election_id=election.id)
                 .filter(SampledBatchDraw.round_id == round.id)
-                .filter(SampledBatchDraw.ticket_number != EXTRA_TICKET_NUMBER)
                 .all()
             )
             sampled_batch_ids = {batch.id for batch in sampled_batches}

server/api/rounds.py

@@ -3,7 +3,7 @@
 from flask import jsonify, request
 from jsonschema import validate
 from werkzeug.exceptions import BadRequest, Conflict
-from sqlalchemy import and_, func, not_
+from sqlalchemy import and_, func, not_, or_
 
 
 from . import api
@@ -64,11 +64,11 @@ def is_round_ready_to_finish(election: Election, round: Round) -> bool:
         num_jurisdictions_not_finalized: int = (
             Jurisdiction.query.filter_by(election_id=election.id)
             .filter(
-                Jurisdiction.id.in_(
-                    SampledBatchDraw.query.filter_by(round_id=round.id)
-                    .join(Batch)
-                    .with_entities(Batch.jurisdiction_id)
-                    .subquery()
+                Jurisdiction.batches.any(
+                    or_(
+                        Batch.draws.any(SampledBatchDraw.round_id == round.id),
+                        Batch.extra_draws.any(ExtraBatchDraw.round_id == round.id),
+                    )
                 )
             )
             .filter(
@@ -195,8 +195,6 @@ def count_audited_votes(election: Election, round: Round):
                     BatchResult.tally_sheet_id.in_(
                         BatchResultTallySheet.query.join(Batch)
                         .join(SampledBatchDraw)
-                        # Special case: don't include extra sampled batches
-                        .filter(SampledBatchDraw.ticket_number != EXTRA_TICKET_NUMBER)
                         .filter_by(round_id=round.id)
                         .with_entities(BatchResultTallySheet.id)
                         .subquery()
@@ -356,8 +354,10 @@ def draw_sample_batches(
     round: Round,
     contest_sample_sizes: list[tuple[Contest, SampleSize]],
 ):
-    sample = compute_sample_batches(election, round.round_num, contest_sample_sizes)
-    for batch_draw in sample:
+    rla_sample_batches, extra_batches = compute_sample_batches(
+        election, round.round_num, contest_sample_sizes
+    )
+    for batch_draw in rla_sample_batches:
         sampled_batch_draw = SampledBatchDraw(
             batch_id=batch_draw["batch_id"],
             round_id=round.id,
@@ -366,6 +366,13 @@ def draw_sample_batches(
         )
         db_session.add(sampled_batch_draw)
 
+    for extra_batch in extra_batches:
+        round_extra_batch = ExtraBatchDraw(
+            round_id=round.id,
+            batch_id=extra_batch["batch_id"],
+        )
+        db_session.add(round_extra_batch)
+
 
 def draw_sample_ballots(
     election: Election,