Coverage for bbconf/modules/snapshots.py: 28%
87 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-16 22:05 +0000
1import logging
2import os
3from datetime import datetime, timezone
5from sqlalchemy import func, select
6from sqlalchemy.orm import Session
8from bbconf.config_parser import BedBaseConfig
9from bbconf.const import PKG_NAME
10from bbconf.db_utils import BedSnapshot
11from bbconf.exceptions import SnapshotNotFoundError
12from bbconf.models.base_models import (
13 BedSnapshotArtifact,
14 BedSnapshotListResult,
15 BedSnapshotResult,
16)
18_LOGGER = logging.getLogger(PKG_NAME)
20# All snapshots live under this single S3 prefix. Not configurable.
21SNAPSHOT_S3_PREFIX = "snapshot"
24class BedAgentSnapshot:
25 """
26 Class that manages bulk-export snapshots (the ``bed_snapshots`` index).
28 One row per published artifact (metadata / bedsets / bedset_membership /
29 manifest). Adding a snapshot always uploads the file to S3 *and* records it
30 in the database; both writes live here in bbconf. This class also exposes
31 read (``list`` / ``get``) and ``delete`` helpers.
32 """
34 def __init__(self, config: BedBaseConfig):
35 """
36 Initialize BedAgentSnapshot.
38 Args:
39 config: Config object.
40 """
41 self.config = config
42 self._db_engine = self.config.db_engine
44 def add(
45 self,
46 artifacts: BedSnapshotArtifact | list[BedSnapshotArtifact],
47 creation_date: datetime | None = None,
48 ) -> BedSnapshotListResult:
49 """
50 Add snapshot artifacts: upload each to S3 and record it in the database.
52 Every artifact is uploaded under the fixed ``snapshot/`` prefix and then
53 recorded in ``bed_snapshots``. The index rows are written only after all
54 uploads succeed, so a partial upload never leaves dangling rows.
56 Args:
57 artifacts: One artifact or a list of them. Each carries the local
58 ``path`` to upload plus its ``file_type`` and file metadata.
59 creation_date: Build date recorded on every row
60 (defaults to now, UTC).
62 Returns:
63 The created snapshot rows.
64 """
65 if isinstance(artifacts, BedSnapshotArtifact):
66 artifacts = [artifacts]
67 if creation_date is None:
68 creation_date = datetime.now(timezone.utc)
70 # Upload everything first; only record rows once all uploads succeed.
71 results: list[BedSnapshotResult] = []
72 for artifact in artifacts:
73 key = f"{SNAPSHOT_S3_PREFIX}/{os.path.basename(artifact.path)}"
74 self.config.upload_s3(artifact.path, s3_path=key)
75 results.append(
76 BedSnapshotResult(
77 file_path=key,
78 file_type=artifact.file_type,
79 creation_date=creation_date,
80 record_count=artifact.record_count,
81 file_size=artifact.file_size,
82 checksum=artifact.checksum,
83 schema_version=artifact.schema_version,
84 )
85 )
87 with Session(self._db_engine.engine) as session:
88 for result in results:
89 session.add(
90 BedSnapshot(
91 file_path=result.file_path,
92 file_type=result.file_type,
93 creation_date=result.creation_date,
94 record_count=result.record_count,
95 file_size=result.file_size,
96 checksum=result.checksum,
97 schema_version=result.schema_version,
98 )
99 )
100 session.commit()
102 _LOGGER.info(f"Recorded {len(results)} rows in bed_snapshots")
103 return BedSnapshotListResult(count=len(results), results=results)
105 def delete(self, id: int, remove_s3: bool = True) -> None:
106 """
107 Delete a snapshot index row.
109 Args:
110 id: Primary key of the snapshot row.
111 remove_s3: Also delete the underlying S3 object.
113 Returns:
114 None.
116 Raises:
117 SnapshotNotFoundError: If no row with this id exists.
118 """
119 with Session(self._db_engine.engine) as session:
120 row = session.scalar(select(BedSnapshot).where(BedSnapshot.id == id))
121 if row is None:
122 raise SnapshotNotFoundError(f"Snapshot with id '{id}' not found.")
123 file_path = row.file_path
124 session.delete(row)
125 session.commit()
127 if remove_s3:
128 self.config.delete_s3(file_path)
130 def list(
131 self,
132 file_type: str | None = None,
133 limit: int | None = 100,
134 offset: int = 0,
135 ) -> BedSnapshotListResult:
136 """
137 List all snapshot index rows in the database, newest first.
139 Args:
140 file_type: Optional filter on file type.
141 limit: Maximum number of rows to return. ``None`` returns all rows.
142 offset: Number of rows to skip.
144 Returns:
145 List of snapshots and the total matching count.
146 """
147 statement = select(BedSnapshot)
148 count_statement = select(func.count()).select_from(BedSnapshot)
149 if file_type is not None:
150 statement = statement.where(BedSnapshot.file_type == file_type)
151 count_statement = count_statement.where(BedSnapshot.file_type == file_type)
152 statement = statement.order_by(
153 BedSnapshot.creation_date.desc(), BedSnapshot.id.desc()
154 )
155 if limit is not None:
156 statement = statement.limit(limit).offset(offset)
157 elif offset:
158 statement = statement.offset(offset)
160 with Session(self._db_engine.engine) as session:
161 total = session.execute(count_statement).scalar_one()
162 rows = session.scalars(statement).all()
163 results = [self._to_result(row) for row in rows]
165 return BedSnapshotListResult(count=total, results=results)
167 def get_by_filename(self, filename: str) -> BedSnapshotResult:
168 """
169 Resolve a snapshot by its file name (the basename of its S3 key).
171 Returns the newest row whose ``file_path`` basename equals ``filename``.
172 Used to round-trip an export's DRS object-id back to its row.
174 Args:
175 filename: The bare file name, e.g.
176 ``bedbase_metadata_2026_08_03.parquet``.
178 Returns:
179 The matching snapshot row.
181 Raises:
182 SnapshotNotFoundError: If no row matches.
183 """
184 filename = os.path.basename(filename)
185 with Session(self._db_engine.engine) as session:
186 rows = session.scalars(
187 select(BedSnapshot)
188 .where(BedSnapshot.file_path.like(f"%{filename}"))
189 .order_by(BedSnapshot.creation_date.desc(), BedSnapshot.id.desc())
190 ).all()
191 for row in rows:
192 if os.path.basename(row.file_path) == filename:
193 return self._to_result(row)
194 raise SnapshotNotFoundError(f"Snapshot '{filename}' not found.")
196 def delete_by_checksum(self, checksum: str, remove_s3: bool = True) -> None:
197 """
198 Delete snapshot index rows by their checksum.
200 Deletes every ``bed_snapshots`` row whose ``checksum`` matches (a checksum
201 identifies one file's content) and optionally removes the underlying S3
202 objects.
204 Args:
205 checksum: SHA256 checksum of the snapshot file.
206 remove_s3: Also delete the underlying S3 object(s).
208 Returns:
209 None.
211 Raises:
212 SnapshotNotFoundError: If no row matches the checksum.
213 """
214 with Session(self._db_engine.engine) as session:
215 rows = session.scalars(
216 select(BedSnapshot).where(BedSnapshot.checksum == checksum)
217 ).all()
218 if not rows:
219 raise SnapshotNotFoundError(
220 f"Snapshot with checksum '{checksum}' not found."
221 )
222 file_paths = {row.file_path for row in rows}
223 for row in rows:
224 session.delete(row)
225 session.commit()
227 if remove_s3:
228 for file_path in file_paths:
229 self.config.delete_s3(file_path)
231 def get(self, id: int) -> BedSnapshotResult:
232 """
233 Get a single snapshot index row by id.
235 Args:
236 id: Primary key of the snapshot row.
238 Returns:
239 The snapshot row.
241 Raises:
242 SnapshotNotFoundError: If no row with this id exists.
243 """
244 with Session(self._db_engine.engine) as session:
245 row = session.scalar(select(BedSnapshot).where(BedSnapshot.id == id))
246 if row is None:
247 raise SnapshotNotFoundError(f"Snapshot with id '{id}' not found.")
248 return self._to_result(row)
250 @staticmethod
251 def _to_result(row: BedSnapshot) -> BedSnapshotResult:
252 return BedSnapshotResult(
253 file_path=row.file_path,
254 file_type=row.file_type,
255 creation_date=row.creation_date,
256 record_count=row.record_count,
257 file_size=row.file_size,
258 checksum=row.checksum,
259 schema_version=row.schema_version,
260 )