Coverage for bbconf/modules/bedsets.py: 72%
210 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-16 22:05 +0000
1import logging
3from geniml.io.utils import compute_md5sum_bedset
4from sqlalchemy import Float, Numeric, func, or_, select
5from sqlalchemy.exc import IntegrityError
6from sqlalchemy.orm import Session
8from bbconf.config_parser import BedBaseConfig
9from bbconf.const import PKG_NAME
10from bbconf.db_utils import Bed, BedFileBedSetRelation, BedSets, BedStats, Files
11from bbconf.exceptions import (
12 BedBaseConfError,
13 BEDFileNotFoundError,
14 BedSetExistsError,
15 BedSetNotFoundError,
16 BedSetTrackHubLimitError,
17)
18from bbconf.models.bed_models import BedStatsModel, StandardMeta
19from bbconf.models.bedset_models import (
20 BedMetadataBasic,
21 BedSetBedFiles,
22 BedSetListResult,
23 BedSetMetadata,
24 BedSetPEP,
25 BedSetPlots,
26 BedSetStats,
27 FileModel,
28)
30_LOGGER = logging.getLogger(PKG_NAME)
33class BedAgentBedSet:
34 """
35 Class that represents Bedset in Database.
37 This class has method to add, delete, get files and metadata from the database.
38 """
40 def __init__(self, config: BedBaseConfig):
41 """
42 Initialize BedAgentBedSet.
44 Args:
45 config: Config object.
46 """
47 self.config = config
48 self._db_engine = self.config.db_engine
50 def get(self, identifier: str, full: bool = False) -> BedSetMetadata:
51 """
52 Get file metadata by identifier.
54 Args:
55 identifier: Bed file identifier.
56 full: Return full record with stats, plots, files and metadata.
58 Returns:
59 BED file metadata.
60 """
62 statement = select(BedSets).where(BedSets.id == identifier)
64 with Session(self._db_engine.engine) as session:
65 bedset_obj = session.scalar(statement)
66 if not bedset_obj:
67 raise BedSetNotFoundError(identifier)
68 list_of_bedfiles = [
69 bedset_obj.bedfile_id for bedset_obj in bedset_obj.bedfiles
70 ]
71 if full:
72 plots = BedSetPlots()
73 for plot in bedset_obj.files:
74 setattr(plots, plot.name, FileModel(**plot.__dict__))
76 stats = BedSetStats(
77 mean=BedStatsModel(**bedset_obj.bedset_means),
78 sd=BedStatsModel(**bedset_obj.bedset_standard_deviation),
79 ).model_dump()
80 else:
81 plots = None
82 stats = None
84 bedset_metadata = BedSetMetadata(
85 id=bedset_obj.id,
86 name=bedset_obj.name,
87 description=bedset_obj.description,
88 md5sum=bedset_obj.md5sum,
89 statistics=stats,
90 plots=plots,
91 bed_ids=list_of_bedfiles,
92 bedfile_count=bedset_obj.bedfile_count,
93 submission_date=bedset_obj.submission_date,
94 last_update_date=bedset_obj.last_update_date,
95 author=bedset_obj.author,
96 source=bedset_obj.source,
97 )
99 return bedset_metadata
101 def get_plots(self, identifier: str) -> BedSetPlots:
102 """
103 Get plots for bedset by identifier.
105 Args:
106 identifier: Bedset identifier.
108 Returns:
109 Bedset plots.
110 """
111 statement = select(BedSets).where(BedSets.id == identifier)
113 with Session(self._db_engine.engine) as session:
114 bedset_object = session.scalar(statement)
115 if not bedset_object:
116 raise BedSetNotFoundError(f"Bed file with id: {identifier} not found.")
117 bedset_files = BedSetPlots()
118 for result in bedset_object.files:
119 if result.name in BedSetPlots.model_fields:
120 setattr(
121 bedset_files,
122 result.name,
123 FileModel(
124 **result.__dict__,
125 object_id=f"bed.{identifier}.{result.name}",
126 access_methods=self.config.construct_access_method_list(
127 result.path
128 ),
129 ),
130 )
131 return bedset_files
133 def get_objects(self, identifier: str) -> dict[str, FileModel]:
134 """
135 Get objects for bedset by identifier.
137 Args:
138 identifier: Bedset identifier.
140 Returns:
141 Bedset objects.
142 """
143 statement = select(BedSets).where(BedSets.id == identifier)
144 return_dict = {}
146 with Session(self._db_engine.engine) as session:
147 bedset_object = session.scalar(statement)
148 if not bedset_object:
149 raise BedSetNotFoundError(f"Bedset with id: {identifier} not found.")
150 for result in bedset_object.files:
151 return_dict[result.name] = FileModel(
152 **result.__dict__,
153 object_id=f"bed.{identifier}.{result.name}",
154 access_methods=self.config.construct_access_method_list(
155 result.path
156 ),
157 )
159 return return_dict
161 def get_statistics(self, identifier: str) -> BedSetStats:
162 """
163 Get statistics for bedset by identifier.
165 Args:
166 identifier: Bedset identifier.
168 Returns:
169 Bedset statistics.
170 """
171 statement = select(BedSets).where(BedSets.id == identifier)
172 with Session(self._db_engine.engine) as session:
173 bedset_object = session.scalar(statement)
174 if not bedset_object:
175 raise BedSetNotFoundError(f"Bedset with id: {identifier} not found.")
176 return BedSetStats(
177 mean=BedStatsModel(**bedset_object.bedset_means),
178 sd=BedStatsModel(**bedset_object.bedset_standard_deviation),
179 )
181 def get_bedset_pep(self, identifier: str) -> dict:
182 """
183 Create pep file for a bedset.
185 Args:
186 identifier: Bedset identifier.
188 Returns:
189 Pep dict.
190 """
192 statement = select(BedFileBedSetRelation).where(
193 BedFileBedSetRelation.bedset_id == identifier
194 )
196 with Session(self._db_engine.engine) as session:
197 bedfile_bedset = session.scalars(statement)
198 bedfiles = [res.bedfile for res in bedfile_bedset]
200 if len(bedfiles) == 0:
201 raise BedSetNotFoundError(identifier)
202 else:
203 bedfile_meta_list = []
205 for bedfile in bedfiles:
206 try:
207 annotation = bedfile.annotations.__dict__
208 except AttributeError:
209 annotation = {}
211 bedfile_metadata = BedSetPEP(
212 sample_name=bedfile.id,
213 original_name=bedfile.name,
214 genome_alias=bedfile.genome_alias,
215 genome_digest=bedfile.genome_digest,
216 bed_compliance=bedfile.bed_compliance,
217 data_format=bedfile.data_format,
218 description=bedfile.description,
219 url=f"https://data2.bedbase.org/files/{bedfile.id[0]}/{bedfile.id[1]}/{bedfile.id}.bed.gz",
220 **annotation,
221 )
222 bedfile_meta_list.append(bedfile_metadata.model_dump())
224 bedset = session.scalar(select(BedSets).where(BedSets.id == identifier))
226 pep_config = {
227 "pep_version": "2.1.0",
228 "name": bedset.id,
229 "description": bedset.description,
230 "md5sum": bedset.md5sum,
231 "author": bedset.author,
232 "source": bedset.source,
233 }
235 return {
236 "_config": pep_config,
237 "_sample_dict": bedfile_meta_list,
238 "_subsample_list": [],
239 }
241 def get_track_hub_file(self, identifier: str) -> str:
242 """
243 Get track hub file for bedset.
245 Args:
246 identifier: Bedset identifier.
248 Returns:
249 Track hub file.
250 """
251 statement = select(BedFileBedSetRelation).where(
252 BedFileBedSetRelation.bedset_id == identifier
253 )
255 trackDb_txt = ""
257 with Session(self._db_engine.engine) as session:
258 bs2bf_objects = session.scalars(statement)
259 if not bs2bf_objects:
260 raise BedSetNotFoundError(f"Bedset with id: {identifier} not found.")
262 relationship_objects = [relationship for relationship in bs2bf_objects]
264 if len(relationship_objects) > 20:
265 raise BedSetTrackHubLimitError(
266 "Number of bedfiles exceeds 20. Unable to process request for track hub."
267 )
269 for bs2bf_obj in relationship_objects:
270 bed_obj = bs2bf_obj.bedfile
272 try:
273 bigbed_url = None
274 for bedfile in bed_obj.files:
275 if bedfile.name == "bigbed_file":
276 bigbed_url = self.config.get_prefixed_uri(
277 postfix=bedfile.path, access_id="http"
278 )
279 break
280 if not bigbed_url:
281 _LOGGER.debug(
282 f"BigBed file for bedfile {bs2bf_obj.bedfile_id} not found."
283 )
284 continue
285 except AttributeError:
286 _LOGGER.debug(
287 f"BigBed file for bedfile {bs2bf_obj.bedfile_id} not found."
288 )
289 continue
290 trackDb_txt = (
291 trackDb_txt + f"track\t {bed_obj.name}\n"
292 "type\t bigBed\n"
293 f"bigDataUrl\t {bigbed_url} \n"
294 f"shortLabel\t {bed_obj.name}\n"
295 f"longLabel\t {bed_obj.description}\n"
296 "visibility\t full\n\n"
297 )
298 return trackDb_txt
300 def create(
301 self,
302 identifier: str,
303 name: str,
304 bedid_list: list[str],
305 description: str | None = None,
306 statistics: bool = False,
307 annotation: dict | None = None,
308 plots: dict | None = None,
309 upload_s3: bool = False,
310 local_path: str = "",
311 no_fail: bool = False,
312 overwrite: bool = False,
313 processed: bool = True,
314 ) -> None:
315 """
316 Create bedset in the database.
318 Args:
319 identifier: Bedset identifier.
320 name: Bedset name.
321 description: Bedset description.
322 bedid_list: List of bed file identifiers.
323 statistics: Calculate statistics for bedset.
324 annotation: Bedset annotation (author, source).
325 plots: Dictionary with plots.
326 upload_s3: Upload bedset to s3.
327 local_path: Local path to the output files.
328 no_fail: Do not raise an error if bedset already exists.
329 overwrite: Overwrite the record in the database.
330 processed: Flag to indicate that bedset is processed. [Default: True].
332 Returns:
333 None.
334 """
335 _LOGGER.info(f"Creating bedset '{identifier}'")
337 if statistics:
338 stats = self._calculate_statistics(bedid_list)
339 else:
340 stats = None
341 if self.exists(identifier):
342 if no_fail and not overwrite:
343 _LOGGER.warning(
344 f"Bedset '{identifier}' already exists. no_fail=True. Skipping updating bedset."
345 )
346 return None
348 if not overwrite:
349 raise BedSetExistsError(
350 f"BEDset already exist in the database: {identifier}"
351 )
353 self.delete(identifier)
355 if not isinstance(annotation, dict):
356 annotation = {}
358 if no_fail:
359 bedid_list = list(set(bedid_list))
361 new_bedset = BedSets(
362 id=identifier,
363 name=name,
364 description=description,
365 summary=annotation.get("summary"),
366 bedset_means=stats.mean.model_dump() if stats else None,
367 bedset_standard_deviation=stats.sd.model_dump() if stats else None,
368 md5sum=compute_md5sum_bedset(bedid_list),
369 author=annotation.get("author"),
370 source=annotation.get("source"),
371 processed=processed,
372 bedfile_count=len(bedid_list),
373 )
375 if upload_s3:
376 plots = BedSetPlots(**plots) if plots else BedSetPlots()
377 plots = self.config.upload_files_s3(
378 identifier, files=plots, base_path=local_path, type="bedsets"
379 )
381 try:
382 with Session(self._db_engine.engine) as session:
383 session.add(new_bedset)
385 for bedfile in bedid_list:
386 session.add(
387 BedFileBedSetRelation(bedset_id=identifier, bedfile_id=bedfile)
388 )
389 if upload_s3:
390 for k, v in plots:
391 if v:
392 new_file = Files(
393 **v.model_dump(exclude_none=True, exclude_unset=True),
394 bedset_id=identifier,
395 type="plot",
396 )
397 session.add(new_file)
399 session.commit()
400 except IntegrityError as _:
401 raise BEDFileNotFoundError(
402 "Failed to create bedset. One of the bedfiles does not exist."
403 )
404 except Exception as _:
405 if not no_fail:
406 raise BedBaseConfError("Failed to create bedset. SQL error.")
408 _LOGGER.info(f"Bedset '{identifier}' was created successfully")
409 return None
411 def _calculate_statistics(self, bed_ids: list[str]) -> BedSetStats:
412 """
413 Calculate statistics for bedset.
415 Args:
416 bed_ids: List of bed file identifiers.
418 Returns:
419 Statistics.
420 """
422 _LOGGER.info("Calculating bedset statistics")
423 numeric_columns = [
424 name
425 for name, field in BedStatsModel.model_fields.items()
426 if field.annotation in (float, float | None)
427 ]
429 bedset_sd = {}
430 bedset_mean = {}
431 with Session(self._db_engine.engine) as session:
432 for column_name in numeric_columns:
433 mean_bedset_statement = select(
434 func.round(
435 func.avg(getattr(BedStats, column_name)).cast(Numeric), 4
436 ).cast(Float)
437 ).where(BedStats.id.in_(bed_ids))
439 sd_bedset_statement = select(
440 func.round(
441 func.stddev(getattr(BedStats, column_name)).cast(Numeric),
442 4,
443 ).cast(Float)
444 ).where(BedStats.id.in_(bed_ids))
446 bedset_sd[column_name] = session.execute(sd_bedset_statement).one()[0]
447 bedset_mean[column_name] = session.execute(mean_bedset_statement).one()[
448 0
449 ]
451 bedset_stats = BedSetStats(
452 mean=bedset_mean,
453 sd=bedset_sd,
454 )
456 _LOGGER.info("Bedset statistics were calculated successfully")
457 return bedset_stats
459 def get_ids_list(
460 self, query: str | None = None, limit: int = 10, offset: int = 0
461 ) -> BedSetListResult:
462 """
463 Find (search) bedsets from the database.
465 Use `get(identifier)` to
466 fetch a single bedset's member ids. `bedfile_count` is populated
467 directly from the denormalized column, so it's free.
469 Args:
470 query: Search query.
471 limit: Limit of results.
472 offset: Offset of results.
474 Returns:
475 List of bedsets.
476 """
477 statement = select(BedSets)
478 count_statement = select(func.count(BedSets.id))
479 if query:
480 query = query.strip()
481 sql_search_str = f"%{query}%"
482 statement = statement.where(
483 or_(
484 BedSets.name.ilike(sql_search_str),
485 BedSets.description.ilike(sql_search_str),
486 )
487 )
488 count_statement = count_statement.where(
489 or_(
490 BedSets.name.ilike(sql_search_str),
491 BedSets.description.ilike(sql_search_str),
492 )
493 )
495 with Session(self._db_engine.engine) as session:
496 bedset_list = session.scalars(
497 statement.order_by(BedSets.id).limit(limit).offset(offset)
498 )
499 bedset_count = session.execute(count_statement).one()
501 result_list = [
502 BedSetMetadata(
503 id=bedset_obj.id,
504 name=bedset_obj.name,
505 description=bedset_obj.description,
506 md5sum=bedset_obj.md5sum,
507 bedfile_count=bedset_obj.bedfile_count,
508 submission_date=bedset_obj.submission_date,
509 last_update_date=bedset_obj.last_update_date,
510 author=bedset_obj.author,
511 source=bedset_obj.source,
512 )
513 for bedset_obj in bedset_list
514 ]
516 return BedSetListResult(
517 count=bedset_count[0],
518 limit=limit,
519 offset=offset,
520 results=result_list,
521 )
523 def get_bedset_bedfiles(self, identifier: str) -> BedSetBedFiles:
524 """
525 Get list of bedfiles in bedset.
527 Args:
528 identifier: Bedset identifier.
530 Returns:
531 List of bedfiles.
532 """
533 sub_statement = select(BedFileBedSetRelation.bedfile_id).where(
534 BedFileBedSetRelation.bedset_id == identifier
535 )
536 statement = select(Bed).where(Bed.id.in_(sub_statement))
538 with Session(self._db_engine.engine) as session:
539 bedfiles_list = session.scalars(statement)
540 results = [
541 BedMetadataBasic(
542 **bedfile_obj.__dict__,
543 annotation=StandardMeta(
544 **(
545 bedfile_obj.annotations.__dict__
546 if bedfile_obj.annotations
547 else {}
548 )
549 ),
550 )
551 for bedfile_obj in bedfiles_list
552 ]
554 return BedSetBedFiles(
555 count=len(results),
556 results=results,
557 )
559 def delete(self, identifier: str) -> None:
560 """
561 Delete bed file from the database.
563 Args:
564 identifier: Bedset identifier.
566 Returns:
567 None.
568 """
569 if not self.exists(identifier):
570 raise BedSetNotFoundError(identifier)
572 _LOGGER.info(f"Deleting bedset '{identifier}'")
574 with Session(self._db_engine.engine) as session:
575 statement = select(BedSets).where(BedSets.id == identifier)
577 bedset_obj = session.scalar(statement)
578 files = [FileModel(**k.__dict__) for k in bedset_obj.files]
580 session.delete(bedset_obj)
581 session.commit()
583 if files:
584 self.config.delete_files_s3(files)
586 def exists(self, identifier: str) -> bool:
587 """
588 Check if bedset exists in the database.
590 Args:
591 identifier: Bedset identifier.
593 Returns:
594 True if bedset exists, False otherwise.
595 """
596 statement = select(BedSets).where(BedSets.id == identifier)
597 with Session(self._db_engine.engine) as session:
598 result = session.execute(statement).one_or_none()
599 if result:
600 return True
601 return False
603 def get_unprocessed(self, limit: int = 100, offset: int = 0) -> BedSetListResult:
604 """
605 Get unprocessed bedset from the database.
607 Args:
608 limit: Limit of results.
609 offset: Offset of results.
611 Returns:
612 Bedset metadata.
613 """
615 with Session(self._db_engine.engine) as session:
616 statement = (
617 select(BedSets)
618 .where(BedSets.processed.is_(False))
619 .limit(limit)
620 .offset(offset)
621 )
622 count_statement = select(func.count()).where(BedSets.processed.is_(False))
624 count = session.execute(count_statement).one()[0]
626 bedset_object_list = session.scalars(statement)
628 results = []
630 for bedset_obj in bedset_object_list:
631 list_of_bedfiles = [
632 bedset_obj.bedfile_id for bedset_obj in bedset_obj.bedfiles
633 ]
635 results.append(
636 BedSetMetadata(
637 id=bedset_obj.id,
638 name=bedset_obj.name,
639 description=bedset_obj.description,
640 md5sum=bedset_obj.md5sum,
641 statistics=None,
642 plots=None,
643 bed_ids=list_of_bedfiles,
644 bedfile_count=bedset_obj.bedfile_count,
645 submission_date=bedset_obj.submission_date,
646 last_update_date=bedset_obj.last_update_date,
647 author=bedset_obj.author,
648 source=bedset_obj.source,
649 )
650 )
652 return BedSetListResult(
653 count=count,
654 limit=limit,
655 offset=offset,
656 results=results,
657 )
659 def add_bedfile(self, identifier: str, bedfile: str) -> None:
660 raise NotImplementedError
662 def delete_bedfile(self, identifier: str, bedfile: str) -> None:
663 raise NotImplementedError