Skip to content
Merged
Show file tree
Hide file tree
Changes from 6 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions packages/grz-common/src/grz_common/workers/worker.py
Original file line number Diff line number Diff line change
Expand Up @@ -370,13 +370,13 @@ def populate(self, s3_options: S3Options, db: SubmissionDb, submission_id: str,
submission_diff, donors_diff = db.diff(submission_id, metadata, submission_date)

# check for destructive changes
if not force and (len(submission_diff.deleted) + len(submission_diff.updated)) > 0:
if not force and submission_diff.has_pending_destructive:
raise RuntimeError(
f"Would update/delete existing submission data in the database, but `force` not set. "
f"submission_id={submission_id!r}, submission_diff={submission_diff!r}"
)

if not force and (len(donors_diff.deleted) + len(donors_diff.updated)) > 0:
if not force and donors_diff.has_pending_destructive:
raise RuntimeError(
f"Would update/delete existing donors in the database, but `force` not set. "
f"submission_id={submission_id!r}, donors_diff={donors_diff!r}"
Expand Down
18 changes: 18 additions & 0 deletions packages/grz-db/src/grz_db/models/submission/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -919,6 +919,24 @@ def get_submission(self, submission_id: str) -> Submission | None:
submission = session.exec(statement).first()
return submission

def get_submissions(self, submission_ids: Sequence[str]) -> list[Submission | None]:
"""Fetch a specific set of submissions by their IDs in a single query.

:param submission_ids: IDs to fetch.
:returns: A list of the same length as *submission_ids*, where each element is the
matching :class:`Submission` or ``None`` when the ID does not exist.
"""
if not submission_ids:
return []
with self._get_session() as session:
statement = (
select(Submission)
.where(Submission.id.in_(submission_ids)) # type: ignore[attr-defined]
.options(selectinload(Submission.states)) # type: ignore[arg-type]
)
found = {s.id: s for s in session.exec(statement).all()}
return [found.get(sid) for sid in submission_ids]

def list_submissions(
self,
limit: int | None,
Expand Down
25 changes: 21 additions & 4 deletions packages/grz-db/src/grz_db/models/submission/diff.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
from __future__ import annotations

from collections.abc import Generator
from dataclasses import dataclass, field
from enum import StrEnum
from typing import TYPE_CHECKING, Self, cast
Expand Down Expand Up @@ -98,14 +99,22 @@ class SubmissionDiffCollection:
unchanged: list[FieldDiff] = field(default_factory=list)

@property
def pending(self) -> list[FieldDiff]:
def pending(self) -> Generator[FieldDiff, None, None]:
"""All diffs that need to be written to the database (added + updated + deleted)."""
return self.added + self.updated + self.deleted
yield from self.added
yield from self.updated
yield from self.deleted

@property
def has_pending(self) -> bool:
"""True if any field needs to be written to the database (added, updated, or deleted)."""
return len(self.added) > 0 or len(self.updated) > 0 or len(self.deleted) > 0

@property
def has_pending_destructive(self) -> bool:
"""True if any field will overwrite or remove an existing database value (updated or deleted)."""
return len(self.updated) > 0 or len(self.deleted) > 0

def append(self, field_diff: FieldDiff):
match field_diff.diff.state:
case DiffState.UPDATED:
Expand Down Expand Up @@ -179,14 +188,22 @@ class DonorsDiffCollection:
unchanged: list[DonorDiff] = field(default_factory=list)

@property
def pending(self) -> list[DonorDiff]:
def pending(self) -> Generator[DonorDiff, None, None]:
"""All diffs that need to be written to the database (added + updated + deleted)."""
return self.added + self.updated + self.deleted
yield from self.added
yield from self.updated
yield from self.deleted

@property
def has_pending(self) -> bool:
"""True if any donor needs to be written to the database (added, updated, or deleted)."""
return len(self.added) > 0 or len(self.updated) > 0 or len(self.deleted) > 0

@property
def has_pending_destructive(self) -> bool:
"""True if any donor will overwrite or remove an existing database record (updated or deleted)."""
return len(self.updated) > 0 or len(self.deleted) > 0

def append(self, donor_diff: DonorDiff):
match donor_diff.state:
case DiffState.UPDATED:
Expand Down
53 changes: 53 additions & 0 deletions packages/grz-db/tests/test_submission.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,8 @@

TWO_TB = 2 * 1024**4 # 2,199,023,255,552 bytes
SUBMISSION_ID = "123456789_2024-01-01_abcdef01"
SUBMISSION_ID_2 = "123456789_2024-01-02_abcdef02"
SUBMISSION_ID_3 = "123456789_2024-01-03_abcdef03"


@pytest.fixture(scope="function")
Expand Down Expand Up @@ -94,3 +96,54 @@ def test_from_metadata_sets_fields_from_metadata(metadata: GrzSubmissionMetadata
assert system_fields.isdisjoint(submission.model_fields_set), (
f"System fields unexpectedly set: {system_fields & submission.model_fields_set}"
)


def test_get_submissions_returns_matching(db: SubmissionDb) -> None:
"""get_submissions returns one entry per requested ID in the same order."""
db.add_submission(SUBMISSION_ID)
db.add_submission(SUBMISSION_ID_2)
db.add_submission(SUBMISSION_ID_3)

result = db.get_submissions([SUBMISSION_ID, SUBMISSION_ID_3])

assert len(result) == 2
assert result[0] is not None and result[0].id == SUBMISSION_ID
assert result[1] is not None and result[1].id == SUBMISSION_ID_3


def test_get_submissions_empty_input(db: SubmissionDb) -> None:
"""get_submissions with an empty list returns an empty list."""
assert db.get_submissions([]) == []


def test_get_submissions_unknown_ids_map_to_none(db: SubmissionDb) -> None:
"""get_submissions maps unknown IDs to None, preserving position."""
db.add_submission(SUBMISSION_ID)
missing = "000000000_2000-01-01_deadbeef"

result = db.get_submissions([SUBMISSION_ID, missing])

assert len(result) == 2
assert result[0] is not None and result[0].id == SUBMISSION_ID
assert result[1] is None


def test_get_submissions_all_unknown(db: SubmissionDb) -> None:
"""get_submissions returns all-None when none of the IDs exist."""
ids = ["000000000_2000-01-01_deadbeef", "000000000_2000-01-01_cafebabe"]
result = db.get_submissions(ids)

assert result == [None, None]


def test_get_submissions_includes_states(db: SubmissionDb) -> None:
"""States relationship is eagerly loaded so it can be accessed outside the session."""
from grz_db.models.submission import SubmissionStateEnum

db.add_submission(SUBMISSION_ID)
db.update_submission_state(SUBMISSION_ID, SubmissionStateEnum.UPLOADED)

result = db.get_submissions([SUBMISSION_ID])

assert result[0] is not None
assert len(result[0].states) >= 1
Loading
Loading