Coverage for bbconf/db_utils.py: 91%
362 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-16 22:05 +0000
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-16 22:05 +0000
1import datetime
2import logging
3import os
4from typing import Optional
6import pandas as pd
7from alembic import command
8from alembic.config import Config
9from sqlalchemy import (
10 DDL,
11 TIMESTAMP,
12 BigInteger,
13 ForeignKey,
14 Index,
15 Result,
16 Select,
17 String,
18 UniqueConstraint,
19 event,
20 select,
21 text,
22)
23from sqlalchemy.dialects.postgresql import ARRAY, JSON, JSONB
24from sqlalchemy.engine import URL, Engine, create_engine
25from sqlalchemy.event import listens_for
26from sqlalchemy.exc import IntegrityError, ProgrammingError
27from sqlalchemy.ext.compiler import compiles
28from sqlalchemy.orm import DeclarativeBase, Mapped, Session, mapped_column, relationship
29from sqlalchemy_schemadisplay import create_schema_graph
31from bbconf.const import LICENSES_CSV_URL, PKG_NAME
33_LOGGER = logging.getLogger(PKG_NAME)
36POSTGRES_DIALECT = "postgresql+psycopg"
38# Connection-pool hardening. Without these, a database host that reboots
39# without sending a FIN or RST leaves every pooled socket blocked in recv()
40# until the Linux default TCP keepalive fires (~2 hours), so every request
41# waits out pool_timeout and returns 500. See the 2026-09-18 outage.
42POOL_RECYCLE_SECONDS = 1800 # drop connections older than 30 minutes
44# libpq connection keywords; psycopg 3 passes these straight through.
45CONNECT_ARGS = {
46 "connect_timeout": 5,
47 "keepalives": 1,
48 "keepalives_idle": 30,
49 "keepalives_interval": 10,
50 "keepalives_count": 3, # idle dead peer detected in ~60 s, not ~2 h
51 "tcp_user_timeout": 30000, # unacked data gives up after 30 s, not ~15 min
52}
54# tables, that were created in this execution
55tables_initialized: list = []
58class SchemaError(Exception):
59 def __init__(self):
60 super().__init__(
61 """
62 The database schema is incorrect, can't connect to the database!"""
63 )
66class BIGSERIAL(BigInteger):
67 pass
70@compiles(BIGSERIAL, POSTGRES_DIALECT)
71def compile_bigserial_pg(type_, compiler, **kw):
72 return "BIGSERIAL"
75@compiles(JSON, POSTGRES_DIALECT)
76def compile_jsonb_pg(type_, compiler, **kw):
77 return "JSONB"
80@compiles(ARRAY, POSTGRES_DIALECT)
81def compile_array_pg(type_, compiler, **kw):
82 return "ARRAY"
85class Base(DeclarativeBase):
86 type_annotation_map = {datetime.datetime: TIMESTAMP(timezone=True)}
89@event.listens_for(Base.metadata, "after_create")
90def receive_after_create(target, connection, tables, **kw):
91 """
92 listen for the 'after_create' event
93 """
94 global tables_initialized
95 if tables:
96 _LOGGER.info("A table was created")
97 tables_initialized = [name.fullname for name in tables]
98 else:
99 _LOGGER.info("A table was not created")
102def deliver_update_date(context):
103 return datetime.datetime.now(datetime.timezone.utc)
106class Bed(Base):
107 __tablename__ = "bed"
109 id: Mapped[str] = mapped_column(primary_key=True, index=True)
110 name: Mapped[Optional[str]]
111 genome_alias: Mapped[Optional[str]]
112 genome_digest: Mapped[Optional[str]]
113 description: Mapped[Optional[str]]
114 bed_compliance: Mapped[str] = mapped_column(default="bed3+0")
115 data_format: Mapped[str] = mapped_column(default="bed_like")
116 compliant_columns: Mapped[int] = mapped_column(default=3)
117 non_compliant_columns: Mapped[int] = mapped_column(default=0)
119 header: Mapped[Optional[str]] = mapped_column(
120 nullable=True, comment="Header of the bed file, it if was provided."
121 )
123 indexed: Mapped[bool] = mapped_column(
124 default=False, comment="Whether sample was added to qdrant"
125 )
126 file_indexed: Mapped[bool] = mapped_column(
127 default=False,
128 comment="Whether file was tokenized and added to the vector database",
129 )
131 submission_date: Mapped[datetime.datetime] = mapped_column(
132 default=deliver_update_date
133 )
134 last_update_date: Mapped[Optional[datetime.datetime]] = mapped_column(
135 default=deliver_update_date,
136 onupdate=deliver_update_date,
137 )
138 is_universe: Mapped[Optional[bool]] = mapped_column(default=False)
140 files: Mapped[list["Files"]] = relationship(
141 "Files", back_populates="bedfile", cascade="all, delete-orphan"
142 )
144 bedsets: Mapped[list["BedFileBedSetRelation"]] = relationship(
145 "BedFileBedSetRelation", back_populates="bedfile", cascade="all, delete-orphan"
146 )
148 annotations: Mapped["BedMetadata"] = relationship(
149 back_populates="bed", cascade="all, delete-orphan", lazy="joined"
150 )
152 stats: Mapped["BedStats"] = relationship(
153 back_populates="bed", cascade="all, delete-orphan"
154 )
156 universe: Mapped["Universes"] = relationship(
157 back_populates="bed", cascade="all, delete-orphan"
158 )
159 tokenized: Mapped["TokenizedBed"] = relationship(
160 back_populates="bed", cascade="all, delete-orphan"
161 )
162 license_id: Mapped["str"] = mapped_column(
163 ForeignKey("licenses.id", ondelete="CASCADE"), nullable=True, index=True
164 )
165 license_mapping: Mapped["License"] = relationship("License", back_populates="bed")
167 ref_classifier: Mapped[list["GenomeRefStats"]] = relationship(
168 "GenomeRefStats", back_populates="bed", cascade="all, delete-orphan"
169 )
170 processed: Mapped[bool] = mapped_column(
171 default=False, comment="Whether the bed file was processed"
172 )
174 __table_args__ = (
175 # Backs the genome filter (Bed.genome_alias == genome) and the
176 # GROUP BY genome_alias aggregations. Historically created by hand in the
177 # live DB; declared here so a fresh create_all() reproduces it exactly
178 # (name and btree deduplication included).
179 Index(
180 "genome_alias_index",
181 "genome_alias",
182 postgresql_with={"deduplicate_items": "true"},
183 ),
184 # Backs get_recent_beds / list_beds(order_by="submission_date"):
185 # ORDER BY submission_date DESC, id ASC LIMIT n.
186 Index("ix_bed_submission_date", text("submission_date DESC"), text("id")),
187 # Partial indexes for the background-worker backlog scans. They stay
188 # small and get faster as each queue drains toward empty.
189 Index("ix_bed_unprocessed", "id", postgresql_where=text("processed = false")),
190 Index("ix_bed_not_indexed", "id", postgresql_where=text("indexed = false")),
191 Index(
192 "ix_bed_not_file_indexed",
193 "id",
194 postgresql_where=text("file_indexed = false"),
195 ),
196 )
199class BedMetadata(Base):
200 __tablename__ = "bed_metadata"
202 species_name: Mapped[str] = mapped_column(default=None, comment="Organism name")
203 species_id: Mapped[str] = mapped_column(
204 default=None, nullable=True, comment="Organism taxon id"
205 )
207 genotype: Mapped[str] = mapped_column(
208 default=None, nullable=True, comment="Genotype of the sample"
209 )
210 phenotype: Mapped[str] = mapped_column(
211 default=None, nullable=True, comment="Phenotype of the sample"
212 )
214 cell_type: Mapped[str] = mapped_column(
215 default=None,
216 nullable=True,
217 comment="Specific kind of cell with distinct characteristics found in an organism. e.g. Neurons, Hepatocytes, Adipocytes",
218 )
219 cell_line: Mapped[str] = mapped_column(
220 default=None,
221 nullable=True,
222 comment="Population of cells derived from a single cell and cultured in the lab for extended use, e.g. HeLa, HepG2, k562",
223 )
224 tissue: Mapped[str] = mapped_column(
225 default=None, nullable=True, comment="Tissue type"
226 )
227 library_source: Mapped[str] = mapped_column(
228 default=None,
229 nullable=True,
230 comment="Library source (e.g. genomic, transcriptomic)",
231 )
232 assay: Mapped[str] = mapped_column(
233 default=None, nullable=True, comment="Experimental protocol (e.g. ChIP-seq)"
234 )
235 antibody: Mapped[str] = mapped_column(
236 default=None, nullable=True, comment="Antibody used in the assay"
237 )
238 target: Mapped[str] = mapped_column(
239 default=None, nullable=True, comment="Target of the assay (e.g. H3K4me3)"
240 )
241 treatment: Mapped[str] = mapped_column(
242 default=None,
243 nullable=True,
244 comment="Treatment of the sample (e.g. drug treatment)",
245 )
247 original_file_name: Mapped[Optional[str]] = mapped_column(
248 nullable=True, comment="Original file name"
249 )
251 global_sample_id: Mapped[list] = mapped_column(
252 ARRAY(String),
253 default=None,
254 nullable=True,
255 comment="Global sample identifier. e.g. GSM000",
256 )
257 global_experiment_id: Mapped[list] = mapped_column(
258 ARRAY(String),
259 default=None,
260 nullable=True,
261 comment="Global experiment identifier. e.g. GSE000",
262 )
264 id: Mapped[str] = mapped_column(
265 ForeignKey("bed.id", ondelete="CASCADE"),
266 primary_key=True,
267 index=True,
268 )
270 bed: Mapped["Bed"] = relationship("Bed", back_populates="annotations")
273class BedStats(Base):
274 __tablename__ = "bed_stats"
276 id: Mapped[str] = mapped_column(
277 ForeignKey("bed.id", ondelete="CASCADE"),
278 primary_key=True,
279 index=True,
280 )
281 number_of_regions: Mapped[Optional[float]]
282 gc_content: Mapped[Optional[float]]
283 median_tss_dist: Mapped[Optional[float]]
284 mean_region_width: Mapped[Optional[float]]
285 exon_frequency: Mapped[Optional[float]]
286 intron_frequency: Mapped[Optional[float]]
287 promoterprox_frequency: Mapped[Optional[float]]
288 intergenic_frequency: Mapped[Optional[float]]
289 promotercore_frequency: Mapped[Optional[float]]
290 fiveutr_frequency: Mapped[Optional[float]]
291 threeutr_frequency: Mapped[Optional[float]]
292 fiveutr_percentage: Mapped[Optional[float]]
293 threeutr_percentage: Mapped[Optional[float]]
294 promoterprox_percentage: Mapped[Optional[float]]
295 exon_percentage: Mapped[Optional[float]]
296 intron_percentage: Mapped[Optional[float]]
297 intergenic_percentage: Mapped[Optional[float]]
298 promotercore_percentage: Mapped[Optional[float]]
299 tssdist: Mapped[Optional[float]]
301 distributions: Mapped[Optional[dict]] = mapped_column(
302 JSONB,
303 nullable=True,
304 comment="Full distribution arrays from gtars genomicdist (JSONB)",
305 )
307 bed: Mapped["Bed"] = relationship("Bed", back_populates="stats")
309 __table_args__ = (
310 # Backs the "beds missing computed stats" worker scan.
311 Index(
312 "ix_bed_stats_missing_regions",
313 "id",
314 postgresql_where=text("number_of_regions IS NULL"),
315 ),
316 )
319class Files(Base):
320 __tablename__ = "files"
322 id: Mapped[int] = mapped_column(primary_key=True, index=True)
323 name: Mapped[str] = mapped_column(
324 nullable=False, comment="Name of the file, e.g. bed, bigBed"
325 )
326 file_digest: Mapped[str] = mapped_column(
327 nullable=True, comment="Digest of the file. Mainly used for bed file."
328 )
329 title: Mapped[Optional[str]]
330 type: Mapped[str] = mapped_column(
331 default="file", comment="Type of the object, e.g. file, plot, ..."
332 )
333 path: Mapped[str]
334 path_thumbnail: Mapped[str] = mapped_column(
335 nullable=True, comment="Thumbnail path of the file"
336 )
337 description: Mapped[Optional[str]]
338 size: Mapped[Optional[int]] = mapped_column(default=0, comment="Size of the file")
340 bedfile_id: Mapped[str] = mapped_column(
341 ForeignKey("bed.id", ondelete="CASCADE"), nullable=True, index=True
342 )
343 bedset_id: Mapped[str] = mapped_column(
344 ForeignKey("bedsets.id", ondelete="CASCADE"), nullable=True, index=True
345 )
347 bedfile: Mapped["Bed"] = relationship("Bed", back_populates="files")
348 bedset: Mapped["BedSets"] = relationship("BedSets", back_populates="files")
350 __table_args__ = (
351 UniqueConstraint(
352 "name", "bedfile_id"
353 ), # TODO: add in the future file_digest here
354 UniqueConstraint("name", "bedset_id"),
355 )
358class BedFileBedSetRelation(Base):
359 __tablename__ = "bedfile_bedset_relation"
361 bedset_id: Mapped[str] = mapped_column(
362 ForeignKey("bedsets.id", ondelete="CASCADE"), primary_key=True
363 )
364 bedfile_id: Mapped[str] = mapped_column(
365 ForeignKey("bed.id", ondelete="CASCADE"), primary_key=True, index=True
366 )
368 bedset: Mapped["BedSets"] = relationship("BedSets", back_populates="bedfiles")
369 bedfile: Mapped["Bed"] = relationship("Bed", back_populates="bedsets")
372class BedSets(Base):
373 __tablename__ = "bedsets"
375 id: Mapped[str] = mapped_column(primary_key=True, index=True)
376 name: Mapped[str] = mapped_column(nullable=False, comment="Name of the bedset")
377 description: Mapped[Optional[str]] = mapped_column(
378 comment="Description of the bedset"
379 )
380 summary: Mapped[Optional[str]] = mapped_column(
381 nullable=True, comment="Summary of the bedset"
382 )
383 submission_date: Mapped[datetime.datetime] = mapped_column(
384 default=deliver_update_date
385 )
386 last_update_date: Mapped[Optional[datetime.datetime]] = mapped_column(
387 default=deliver_update_date,
388 onupdate=deliver_update_date,
389 )
390 md5sum: Mapped[Optional[str]] = mapped_column(comment="MD5 sum of the bedset")
392 bedset_means: Mapped[Optional[dict]] = mapped_column(
393 JSON, comment="Mean values of the bedset"
394 )
395 bedset_standard_deviation: Mapped[Optional[dict]] = mapped_column(
396 JSON, comment="Median values of the bedset"
397 )
398 bedset_stats: Mapped[Optional[dict]] = mapped_column(
399 JSONB,
400 nullable=True,
401 comment="Pre-aggregated distribution statistics from gtars (JSONB)",
402 )
404 bedfile_count: Mapped[int] = mapped_column(
405 default=0, comment="Number of bedfiles in the bedset (denormalized count)"
406 )
408 bedfiles: Mapped[list["BedFileBedSetRelation"]] = relationship(
409 "BedFileBedSetRelation", back_populates="bedset", cascade="all, delete-orphan"
410 )
411 files: Mapped[list["Files"]] = relationship("Files", back_populates="bedset")
412 universe: Mapped["Universes"] = relationship("Universes", back_populates="bedset")
414 author: Mapped[str] = mapped_column(nullable=True, comment="Author of the bedset")
415 source: Mapped[str] = mapped_column(nullable=True, comment="Source of the bedset")
417 processed: Mapped[bool] = mapped_column(
418 default=False, comment="Whether the bedset was processed"
419 )
421 __table_args__ = (
422 # Backs the "unprocessed bedsets" worker scan.
423 Index(
424 "ix_bedsets_unprocessed", "id", postgresql_where=text("processed = false")
425 ),
426 # Trigram GIN indexes for the bedset search (get_ids_list), which filters
427 # on name/description with ILIKE '%query%'. A leading-wildcard ILIKE
428 # cannot use a btree index at all, so pg_trgm is the only thing that
429 # avoids a full table scan here. Needs the pg_trgm extension, which the
430 # before_create listener below creates on first creation of this table.
431 Index(
432 "ix_bedsets_name_trgm",
433 "name",
434 postgresql_using="gin",
435 postgresql_ops={"name": "gin_trgm_ops"},
436 ),
437 Index(
438 "ix_bedsets_description_trgm",
439 "description",
440 postgresql_using="gin",
441 postgresql_ops={"description": "gin_trgm_ops"},
442 ),
443 )
446# Make the pg_trgm extension available before the bedsets trigram GIN indexes are
447# built. Scoped to this table's creation so it runs only on first-time schema
448# creation (not on every startup) and only on PostgreSQL.
449event.listen(
450 BedSets.__table__,
451 "before_create",
452 DDL("CREATE EXTENSION IF NOT EXISTS pg_trgm").execute_if(dialect="postgresql"),
453)
456class Universes(Base):
457 __tablename__ = "universes"
459 id: Mapped[str] = mapped_column(
460 ForeignKey("bed.id", ondelete="CASCADE"),
461 primary_key=True,
462 index=True,
463 )
464 method: Mapped[str] = mapped_column(
465 nullable=True, comment="Method used to create the universe"
466 )
467 bedset_id: Mapped[str] = mapped_column(
468 ForeignKey("bedsets.id", ondelete="CASCADE"),
469 index=True,
470 nullable=True,
471 )
473 bed: Mapped["Bed"] = relationship("Bed", back_populates="universe")
474 bedset: Mapped["BedSets"] = relationship("BedSets", back_populates="universe")
475 tokenized: Mapped["TokenizedBed"] = relationship(
476 "TokenizedBed",
477 back_populates="universe",
478 )
481class TokenizedBed(Base):
482 __tablename__ = "tokenized_bed"
484 bed_id: Mapped[str] = mapped_column(
485 ForeignKey("bed.id", ondelete="CASCADE"),
486 primary_key=True,
487 index=True,
488 nullable=False,
489 )
490 universe_id: Mapped[str] = mapped_column(
491 ForeignKey("universes.id", ondelete="CASCADE"),
492 primary_key=True,
493 index=True,
494 nullable=False,
495 )
496 path: Mapped[str] = mapped_column(
497 nullable=False, comment="Path to the tokenized bed file"
498 )
500 bed: Mapped["Bed"] = relationship("Bed", back_populates="tokenized")
501 universe: Mapped["Universes"] = relationship(
502 "Universes",
503 back_populates="tokenized",
504 passive_deletes="all",
505 )
508class License(Base):
509 __tablename__ = "licenses"
511 id: Mapped[str] = mapped_column(primary_key=True, index=True)
512 shorthand: Mapped[str] = mapped_column(nullable=True, comment="License shorthand")
513 label: Mapped[str] = mapped_column(nullable=False, comment="License label")
514 description: Mapped[str] = mapped_column(
515 nullable=False, comment="License description"
516 )
518 bed: Mapped[list["Bed"]] = relationship("Bed", back_populates="license_mapping")
521class ReferenceGenome(Base):
522 __tablename__ = "reference_genomes"
524 digest: Mapped[str] = mapped_column(primary_key=True, index=True)
525 alias: Mapped[str] = mapped_column(
526 nullable=False, comment="Name of the reference genome"
527 )
529 bed_reference: Mapped[list["GenomeRefStats"]] = relationship(
530 "GenomeRefStats",
531 back_populates="genome_object",
532 cascade="all, delete-orphan",
533 )
536class GenomeRefStats(Base):
537 __tablename__ = "genome_ref_stats"
539 id: Mapped[int] = mapped_column(primary_key=True, index=True)
541 bed_id: Mapped[str] = mapped_column(
542 ForeignKey("bed.id", ondelete="CASCADE"),
543 index=True,
544 nullable=False,
545 )
546 provided_genome: Mapped[str]
547 compared_genome: Mapped[str] = mapped_column(
548 nullable=False, comment="Compared Genome"
549 )
550 genome_digest: Mapped[str] = mapped_column(
551 ForeignKey("reference_genomes.digest", ondelete="CASCADE"),
552 )
554 xs: Mapped[float] = mapped_column(nullable=True, default=None)
555 oobr: Mapped[float] = mapped_column(nullable=True, default=None)
556 sequence_fit: Mapped[float] = mapped_column(nullable=True, default=None)
557 assigned_points: Mapped[int] = mapped_column(nullable=False)
558 tier_ranking: Mapped[int] = mapped_column(nullable=False)
560 bed: Mapped["Bed"] = relationship("Bed", back_populates="ref_classifier")
562 genome_object: Mapped["ReferenceGenome"] = relationship(
563 "ReferenceGenome",
564 back_populates="bed_reference",
565 lazy="joined",
566 )
568 __table_args__ = (UniqueConstraint("bed_id", "compared_genome"),)
571@listens_for(Universes, "after_insert")
572@listens_for(Universes, "after_update")
573def add_bed_universe(mapper, connection, target):
574 with Session(connection) as session:
575 bed = session.scalar(select(Bed).where(Bed.id == target.id))
576 bed.is_universe = True
577 session.commit()
580@listens_for(Universes, "after_delete")
581def delete_bed_universe(mapper, connection, target):
582 with Session(connection) as session:
583 bed = session.scalar(select(Bed).where(Bed.id == target.id))
584 bed.is_universe = False
585 session.commit()
588class GeoGseStatus(Base):
589 __tablename__ = "geo_gse_status"
591 id: Mapped[int] = mapped_column(primary_key=True, index=True, autoincrement=True)
592 gse: Mapped[str] = mapped_column(nullable=False, comment="GSE number", unique=True)
593 status: Mapped[str] = mapped_column(
594 nullable=False, comment="Status of the GEO project"
595 )
596 submission_date: Mapped[datetime.datetime] = mapped_column(
597 default=deliver_update_date, onupdate=deliver_update_date
598 )
599 number_of_files: Mapped[int] = mapped_column(default=0, comment="Number of files")
600 number_of_success: Mapped[int] = mapped_column(
601 default=0, comment="Number of success"
602 )
603 number_of_skips: Mapped[int] = mapped_column(default=0, comment="Number of skips")
604 number_of_fails: Mapped[int] = mapped_column(default=0, comment="Number of fails")
606 gsm_status_mapper: Mapped[list["GeoGsmStatus"]] = relationship(
607 "GeoGsmStatus", back_populates="gse_status_mapper"
608 )
609 error: Mapped[str] = mapped_column(nullable=True, comment="Error message")
612class GeoGsmStatus(Base):
613 __tablename__ = "geo_gsm_status"
615 id: Mapped[int] = mapped_column(primary_key=True, index=True, autoincrement=True)
616 gse_status_id: Mapped[str] = mapped_column(
617 ForeignKey("geo_gse_status.id", ondelete="CASCADE"), nullable=False, index=True
618 )
619 gsm: Mapped[str] = mapped_column(nullable=False, comment="GSM number", unique=False)
620 sample_name: Mapped[str] = mapped_column()
621 status: Mapped[str] = mapped_column(
622 nullable=False, comment="Status of the GEO sample"
623 )
624 error: Mapped[str] = mapped_column(nullable=True, comment="Error message")
626 source_submission_date: Mapped[datetime.datetime] = mapped_column(
627 nullable=True, comment="Submission date of the source"
628 )
630 submission_date: Mapped[datetime.datetime] = mapped_column(
631 default=deliver_update_date, onupdate=deliver_update_date
632 )
633 bed_id: Mapped[str] = mapped_column(
634 nullable=True, index=True, comment="Bed identifier"
635 )
637 file_size: Mapped[int] = mapped_column(
638 BigInteger, default=0, comment="Size of the file"
639 )
640 genome: Mapped[str] = mapped_column(nullable=True, comment="Genome")
642 gse_status_mapper: Mapped["GeoGseStatus"] = relationship(
643 "GeoGseStatus", back_populates="gsm_status_mapper"
644 )
647class UsageBedMeta(Base):
648 __tablename__ = "usage_bed_meta"
650 id: Mapped[int] = mapped_column(primary_key=True, index=True, autoincrement=True)
652 bed_id: Mapped[str] = mapped_column(
653 ForeignKey("bed.id", ondelete="CASCADE"), nullable=True, index=True
654 )
656 count: Mapped[int] = mapped_column(default=0, comment="Number of visits")
658 date_from: Mapped[datetime.datetime] = mapped_column(comment="Date from")
659 date_to: Mapped[datetime.datetime] = mapped_column(comment="Date to")
662class UsageBedSetMeta(Base):
663 __tablename__ = "usage_bedset_meta"
665 id: Mapped[int] = mapped_column(primary_key=True, index=True, autoincrement=True)
667 bedset_id: Mapped[str] = mapped_column(
668 ForeignKey("bedsets.id", ondelete="CASCADE"), nullable=True, index=True
669 )
670 count: Mapped[int] = mapped_column(default=0, comment="Number of visits")
672 date_from: Mapped[datetime.datetime] = mapped_column(comment="Date from")
673 date_to: Mapped[datetime.datetime] = mapped_column(comment="Date to")
676class UsageFiles(Base):
677 __tablename__ = "usage_files"
679 id: Mapped[int] = mapped_column(primary_key=True, index=True, autoincrement=True)
680 file_path: Mapped[str] = mapped_column(nullable=False, comment="Path to the file")
681 count: Mapped[int] = mapped_column(default=0, comment="Number of downloads")
683 date_from: Mapped[datetime.datetime] = mapped_column(comment="Date from")
684 date_to: Mapped[datetime.datetime] = mapped_column(comment="Date to")
687class UsageSearch(Base):
688 __tablename__ = "usage_search"
690 id: Mapped[int] = mapped_column(primary_key=True, index=True, autoincrement=True)
691 query: Mapped[str] = mapped_column(nullable=False, comment="Search query")
692 type: Mapped[str] = mapped_column(
693 nullable=False, comment="Type of the search. Bed/Bedset"
694 )
695 count: Mapped[int] = mapped_column(default=0, comment="Number of searches")
697 date_from: Mapped[datetime.datetime] = mapped_column(comment="Date from")
698 date_to: Mapped[datetime.datetime] = mapped_column(comment="Date to")
701class BedSnapshot(Base):
702 """
703 Index of bulk metadata exports published to S3.
705 One row per published artifact (metadata / bedsets / membership / manifest).
706 The exporter writes a row after a successful upload; the /v1/bed/exports
707 endpoint reads them newest-first. This is a new table, so
708 Base.metadata.create_all() creates it on the next connection.
709 """
711 __tablename__ = "bed_snapshots"
713 id: Mapped[int] = mapped_column(primary_key=True, index=True, autoincrement=True)
714 file_path: Mapped[str] = mapped_column(
715 nullable=False, comment="S3 object key, relative to the bucket root"
716 )
717 file_type: Mapped[str] = mapped_column(
718 nullable=False, comment="metadata | bedsets | bedset_membership | manifest"
719 )
720 creation_date: Mapped[datetime.datetime] = mapped_column(
721 default=deliver_update_date, comment="Build date of the export"
722 )
723 record_count: Mapped[Optional[int]] = mapped_column(
724 nullable=True, comment="Rows actually written to the file"
725 )
726 file_size: Mapped[Optional[int]] = mapped_column(
727 nullable=True, comment="Size of the file in bytes"
728 )
729 checksum: Mapped[Optional[str]] = mapped_column(
730 nullable=True, comment="SHA256 of the file"
731 )
732 schema_version: Mapped[Optional[int]] = mapped_column(
733 nullable=True, comment="Export schema version"
734 )
737class AnalysisFile(Base):
738 """
739 Registry of standalone analysis files (openSignalMatrix, models, other
740 analysis inputs) stored in S3. Not tied to any bed file or bedset.
742 Append-only: one row per uploaded file, so name-based lookups resolve the
743 newest matching row (same model as ``bed_snapshots``). This is a new table,
744 so ``Base.metadata.create_all()`` creates it on the next connection.
745 """
747 __tablename__ = "analysis_files"
749 id: Mapped[int] = mapped_column(primary_key=True, index=True, autoincrement=True)
750 name: Mapped[str] = mapped_column(
751 nullable=False, index=True, comment="Logical name/key, e.g. openSignalMatrix"
752 )
753 file_path: Mapped[str] = mapped_column(
754 nullable=False, comment="S3 object key, relative to the bucket root"
755 )
756 file_type: Mapped[Optional[str]] = mapped_column(
757 nullable=True,
758 index=True,
759 comment="Category, e.g. openSignalMatrix | reference | model",
760 )
761 genome: Mapped[Optional[str]] = mapped_column(
762 nullable=True, index=True, comment="Genome/assembly, e.g. hg38 (optional)"
763 )
764 description: Mapped[Optional[str]] = mapped_column(nullable=True)
765 tags: Mapped[Optional[list]] = mapped_column(
766 ARRAY(String), nullable=True, comment="Free-form tags"
767 )
768 file_size: Mapped[Optional[int]] = mapped_column(
769 nullable=True, comment="Size of the file in bytes"
770 )
771 checksum: Mapped[Optional[str]] = mapped_column(
772 nullable=True, comment="SHA256 of the file"
773 )
774 creation_date: Mapped[datetime.datetime] = mapped_column(
775 default=deliver_update_date, comment="Upload date"
776 )
779class BaseEngine:
780 """
781 A class with base methods, that are used in several classes.
782 """
784 def __init__(
785 self,
786 *,
787 host: str = "localhost",
788 port: int = 5432,
789 database: str = "bedbase",
790 user: str | None = None,
791 password: str | None = None,
792 drivername: str = POSTGRES_DIALECT,
793 dsn: str | None = None,
794 echo: bool = False,
795 run_migrations: bool = False,
796 ):
797 """
798 Initialize connection to the bedbase database. You can use the basic connection parameters
799 or libpq connection string.
801 Args:
802 host: Database server address e.g., localhost or an IP address.
803 port: The port number that defaults to 5432 if it is not provided.
804 database: The name of the database that you want to connect.
805 user: The username used to authenticate.
806 password: Password used to authenticate.
807 drivername: Driver used in connection.
808 dsn: Libpq connection string using the dsn parameter
809 (e.g. 'postgresql://user_name:password@host_name:port/db_name').
810 run_migrations: Upgrade the database to the latest Alembic revision
811 (``head``) before connecting. Safe on an already-migrated or
812 pre-existing database (the initial revision is idempotent).
813 """
814 if not dsn:
815 dsn = URL.create(
816 host=host,
817 port=port,
818 database=database,
819 username=user,
820 password=password,
821 drivername=drivername,
822 )
824 if run_migrations:
825 if isinstance(dsn, str):
826 migration_url = dsn
827 else:
828 migration_url = dsn.render_as_string(hide_password=False)
829 self.run_db_migration(migration_url)
831 self._engine = create_engine(
832 dsn,
833 echo=echo,
834 pool_pre_ping=True, # test idle-in-pool connections before reuse
835 pool_recycle=POOL_RECYCLE_SECONDS,
836 connect_args=dict(CONNECT_ARGS),
837 )
838 self.create_schema(self._engine)
839 self.check_db_connection()
841 def create_schema(self, engine=None):
842 """
843 Create sql schema in the database.
845 Args:
846 engine: Sqlalchemy engine [Default: None].
848 Returns:
849 None.
850 """
851 if not engine:
852 engine = self._engine
853 Base.metadata.create_all(engine)
855 global tables_initialized
856 if License.__tablename__ in tables_initialized:
857 try:
858 # It is weired, but tables are sometimes initialized twice. Or it says like that...
859 self._upload_licenses()
860 except IntegrityError:
861 pass
863 def delete_schema(self, engine=None) -> None:
864 """
865 Delete sql schema in the database.
867 Args:
868 engine: Sqlalchemy engine [Default: None].
870 Returns:
871 None.
872 """
873 if not engine:
874 engine = self._engine
875 Base.metadata.drop_all(engine)
876 return None
878 def run_db_migration(self, database_url: str) -> None:
879 """
880 Upgrade the database to the latest Alembic revision (``head``).
882 The Alembic config is built programmatically so the package does not
883 depend on the repo-root ``alembic.ini`` at runtime.
885 Args:
886 database_url: SQLAlchemy connection URL (with password) to migrate.
887 """
888 script_location = os.path.join(os.path.dirname(__file__), "alembic")
890 alembic_cfg = Config()
891 alembic_cfg.set_main_option("script_location", script_location)
892 alembic_cfg.set_main_option("sqlalchemy.url", database_url)
894 _LOGGER.info("Running database migrations to the latest revision...")
895 command.upgrade(alembic_cfg, "head")
897 def session_execute(self, statement: Select) -> Result:
898 """
899 Execute statement using sqlalchemy statement.
901 Args:
902 statement: SQL query or a SQL expression that is constructed using
903 SQLAlchemy's SQL expression language.
905 Returns:
906 Query result represented with declarative base.
907 """
908 _LOGGER.debug(f"Executing statement: {statement}")
909 with Session(self._engine) as session:
910 query_result = session.execute(statement)
912 return query_result
914 @property
915 def session(self):
916 """
917 Get a started sqlalchemy session.
919 Returns:
920 Started sqlalchemy session.
921 """
922 return self._start_session()
924 @property
925 def engine(self) -> Engine:
926 """
927 Get sqlalchemy engine.
929 Returns:
930 Sqlalchemy engine.
931 """
932 return self._engine
934 def _start_session(self):
935 session = Session(self.engine)
936 try:
937 session.execute(select(Bed).limit(1))
938 except ProgrammingError:
939 raise SchemaError()
941 return session
943 def check_db_connection(self):
944 try:
945 self.session_execute(select(Bed).limit(1))
946 except ProgrammingError:
947 raise SchemaError()
949 def create_schema_graph(self, output_file: str = "schema.svg"):
950 """
951 Create schema graph of the database.
953 Args:
954 output_file: Path to the output file.
956 Returns:
957 None.
958 """
959 graph = create_schema_graph(engine=self.engine, metadata=Base.metadata)
960 graph.write(output_file, format="svg", prog="dot")
961 return None
963 def _upload_licenses(self):
964 """
965 Upload licenses to the database.
966 """
967 # Check if licenses already exist to avoid duplicate key errors
968 with Session(self.engine) as session:
969 existing = session.scalar(select(License).limit(1))
970 if existing:
971 _LOGGER.info("Licenses already exist in the database, skipping upload.")
972 return
974 _LOGGER.info("Uploading licenses to the database...")
975 df = pd.read_csv(LICENSES_CSV_URL)
977 df.to_sql(License.__tablename__, self.engine, if_exists="append", index=False)
979 _LOGGER.info("Licenses uploaded successfully!")