Coverage for bbconf/modules/analysis_files.py: 96%
107 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 AnalysisFile
11from bbconf.exceptions import AnalysisFileNotFoundError
12from bbconf.models.base_models import (
13 AnalysisFileArtifact,
14 AnalysisFileListResult,
15 AnalysisFileResult,
16)
18_LOGGER = logging.getLogger(PKG_NAME)
20# All analysis files live under this single S3 prefix. Not configurable.
21ANALYSIS_FILES_S3_PREFIX = "analysis_files"
24class BedAgentAnalysisFile:
25 """
26 Class that manages standalone analysis files (the ``analysis_files`` index).
28 One row per uploaded file (openSignalMatrix, models, other analysis inputs).
29 These files are global: they are not tied to any bed file or bedset. Adding
30 a file always uploads it to S3 *and* records it in the database; both writes
31 live here in bbconf. This class also exposes read (``list`` / ``get`` /
32 ``get_by_name`` / ``get_by_filename``) and ``delete`` helpers.
33 """
35 def __init__(self, config: BedBaseConfig):
36 """
37 Initialize BedAgentAnalysisFile.
39 Args:
40 config: Config object.
41 """
42 self.config = config
43 self._db_engine = self.config.db_engine
45 def add(
46 self,
47 artifacts: AnalysisFileArtifact | list[AnalysisFileArtifact],
48 creation_date: datetime | None = None,
49 ) -> AnalysisFileListResult:
50 """
51 Add analysis files: upload each to S3 and record it in the database.
53 Every artifact is uploaded under the fixed ``analysis_files/`` prefix and
54 then recorded in ``analysis_files``. The index rows are written only
55 after all uploads succeed, so a partial upload never leaves dangling
56 rows.
58 Args:
59 artifacts: One artifact or a list of them. Each carries the local
60 ``path`` to upload plus its ``name`` and file metadata.
61 creation_date: Upload date recorded on every row
62 (defaults to now, UTC).
64 Returns:
65 The created analysis-file rows.
66 """
67 if isinstance(artifacts, AnalysisFileArtifact):
68 artifacts = [artifacts]
69 if creation_date is None:
70 creation_date = datetime.now(timezone.utc)
72 # Upload everything first; only record rows once all uploads succeed.
73 uploads: list[tuple[AnalysisFileArtifact, str]] = []
74 for artifact in artifacts:
75 key = f"{ANALYSIS_FILES_S3_PREFIX}/{os.path.basename(artifact.path)}"
76 self.config.upload_s3(artifact.path, s3_path=key)
77 uploads.append((artifact, key))
79 results: list[AnalysisFileResult] = []
80 with Session(self._db_engine.engine) as session:
81 for artifact, key in uploads:
82 row = AnalysisFile(
83 name=artifact.name,
84 file_path=key,
85 file_type=artifact.file_type,
86 genome=artifact.genome,
87 description=artifact.description,
88 tags=artifact.tags,
89 file_size=artifact.file_size,
90 checksum=artifact.checksum,
91 creation_date=creation_date,
92 )
93 session.add(row)
94 session.flush()
95 results.append(self._to_result(row))
96 session.commit()
98 _LOGGER.info(f"Recorded {len(results)} rows in analysis_files")
99 return AnalysisFileListResult(count=len(results), results=results)
101 def delete(self, id: int, remove_s3: bool = True) -> None:
102 """
103 Delete an analysis-file index row.
105 Args:
106 id: Primary key of the analysis-file row.
107 remove_s3: Also delete the underlying S3 object.
109 Returns:
110 None.
112 Raises:
113 AnalysisFileNotFoundError: If no row with this id exists.
114 """
115 with Session(self._db_engine.engine) as session:
116 row = session.scalar(select(AnalysisFile).where(AnalysisFile.id == id))
117 if row is None:
118 raise AnalysisFileNotFoundError(
119 f"Analysis file with id '{id}' not found."
120 )
121 file_path = row.file_path
122 session.delete(row)
123 session.commit()
125 if remove_s3:
126 self.config.delete_s3(file_path)
128 def list(
129 self,
130 file_type: str | None = None,
131 genome: str | None = None,
132 tag: str | None = None,
133 limit: int | None = 100,
134 offset: int = 0,
135 ) -> AnalysisFileListResult:
136 """
137 List analysis-file index rows in the database, newest first.
139 Args:
140 file_type: Optional filter on file type.
141 genome: Optional filter on genome/assembly.
142 tag: Optional filter; keep only rows whose ``tags`` contain this tag.
143 limit: Maximum number of rows to return. ``None`` returns all rows.
144 offset: Number of rows to skip.
146 Returns:
147 List of analysis files and the total matching count.
148 """
149 statement = select(AnalysisFile)
150 count_statement = select(func.count()).select_from(AnalysisFile)
151 if file_type is not None:
152 statement = statement.where(AnalysisFile.file_type == file_type)
153 count_statement = count_statement.where(AnalysisFile.file_type == file_type)
154 if genome is not None:
155 statement = statement.where(AnalysisFile.genome == genome)
156 count_statement = count_statement.where(AnalysisFile.genome == genome)
157 if tag is not None:
158 statement = statement.where(AnalysisFile.tags.any(tag))
159 count_statement = count_statement.where(AnalysisFile.tags.any(tag))
160 statement = statement.order_by(
161 AnalysisFile.creation_date.desc(), AnalysisFile.id.desc()
162 )
163 if limit is not None:
164 statement = statement.limit(limit).offset(offset)
165 elif offset:
166 statement = statement.offset(offset)
168 with Session(self._db_engine.engine) as session:
169 total = session.execute(count_statement).scalar_one()
170 rows = session.scalars(statement).all()
171 results = [self._to_result(row) for row in rows]
173 return AnalysisFileListResult(count=total, results=results)
175 def get(self, id: int) -> AnalysisFileResult:
176 """
177 Get a single analysis-file index row by id.
179 Args:
180 id: Primary key of the analysis-file row.
182 Returns:
183 The analysis-file row.
185 Raises:
186 AnalysisFileNotFoundError: If no row with this id exists.
187 """
188 with Session(self._db_engine.engine) as session:
189 row = session.scalar(select(AnalysisFile).where(AnalysisFile.id == id))
190 if row is None:
191 raise AnalysisFileNotFoundError(
192 f"Analysis file with id '{id}' not found."
193 )
194 return self._to_result(row)
196 def get_by_name(self, name: str, genome: str | None = None) -> AnalysisFileResult:
197 """
198 Resolve an analysis file by its logical name (newest matching row).
200 Args:
201 name: Logical name/key, e.g. ``openSignalMatrix``.
202 genome: Optional genome/assembly to disambiguate, e.g. ``hg38``.
204 Returns:
205 The newest matching analysis-file row.
207 Raises:
208 AnalysisFileNotFoundError: If no row matches.
209 """
210 statement = select(AnalysisFile).where(AnalysisFile.name == name)
211 if genome is not None:
212 statement = statement.where(AnalysisFile.genome == genome)
213 statement = statement.order_by(
214 AnalysisFile.creation_date.desc(), AnalysisFile.id.desc()
215 )
216 with Session(self._db_engine.engine) as session:
217 row = session.scalars(statement).first()
218 if row is None:
219 raise AnalysisFileNotFoundError(f"Analysis file '{name}' not found.")
220 return self._to_result(row)
222 def get_by_filename(self, filename: str) -> AnalysisFileResult:
223 """
224 Resolve an analysis file by its file name (the basename of its S3 key).
226 Returns the newest row whose ``file_path`` basename equals ``filename``.
227 Used to round-trip a DRS object-id back to its row.
229 Args:
230 filename: The bare file name, e.g. ``openSignalMatrix_hg38.txt.gz``.
232 Returns:
233 The matching analysis-file row.
235 Raises:
236 AnalysisFileNotFoundError: If no row matches.
237 """
238 filename = os.path.basename(filename)
239 with Session(self._db_engine.engine) as session:
240 rows = session.scalars(
241 select(AnalysisFile)
242 .where(AnalysisFile.file_path.like(f"%{filename}"))
243 .order_by(AnalysisFile.creation_date.desc(), AnalysisFile.id.desc())
244 ).all()
245 for row in rows:
246 if os.path.basename(row.file_path) == filename:
247 return self._to_result(row)
248 raise AnalysisFileNotFoundError(f"Analysis file '{filename}' not found.")
250 def delete_by_checksum(self, checksum: str, remove_s3: bool = True) -> None:
251 """
252 Delete analysis-file index rows by their checksum.
254 Deletes every ``analysis_files`` row whose ``checksum`` matches (a
255 checksum identifies one file's content) and optionally removes the
256 underlying S3 objects.
258 Args:
259 checksum: SHA256 checksum of the analysis file.
260 remove_s3: Also delete the underlying S3 object(s).
262 Returns:
263 None.
265 Raises:
266 AnalysisFileNotFoundError: If no row matches the checksum.
267 """
268 with Session(self._db_engine.engine) as session:
269 rows = session.scalars(
270 select(AnalysisFile).where(AnalysisFile.checksum == checksum)
271 ).all()
272 if not rows:
273 raise AnalysisFileNotFoundError(
274 f"Analysis file with checksum '{checksum}' not found."
275 )
276 file_paths = {row.file_path for row in rows}
277 for row in rows:
278 session.delete(row)
279 session.commit()
281 if remove_s3:
282 for file_path in file_paths:
283 self.config.delete_s3(file_path)
285 @staticmethod
286 def _to_result(row: AnalysisFile) -> AnalysisFileResult:
287 return AnalysisFileResult(
288 id=row.id,
289 name=row.name,
290 file_path=row.file_path,
291 file_type=row.file_type,
292 genome=row.genome,
293 description=row.description,
294 tags=list(row.tags) if row.tags is not None else None,
295 file_size=row.file_size,
296 checksum=row.checksum,
297 creation_date=row.creation_date,
298 )