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

1import datetime 

2import logging 

3import os 

4from typing import Optional 

5 

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 

30 

31from bbconf.const import LICENSES_CSV_URL, PKG_NAME 

32 

33_LOGGER = logging.getLogger(PKG_NAME) 

34 

35 

36POSTGRES_DIALECT = "postgresql+psycopg" 

37 

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 

43 

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} 

53 

54# tables, that were created in this execution 

55tables_initialized: list = [] 

56 

57 

58class SchemaError(Exception): 

59 def __init__(self): 

60 super().__init__( 

61 """ 

62 The database schema is incorrect, can't connect to the database!""" 

63 ) 

64 

65 

66class BIGSERIAL(BigInteger): 

67 pass 

68 

69 

70@compiles(BIGSERIAL, POSTGRES_DIALECT) 

71def compile_bigserial_pg(type_, compiler, **kw): 

72 return "BIGSERIAL" 

73 

74 

75@compiles(JSON, POSTGRES_DIALECT) 

76def compile_jsonb_pg(type_, compiler, **kw): 

77 return "JSONB" 

78 

79 

80@compiles(ARRAY, POSTGRES_DIALECT) 

81def compile_array_pg(type_, compiler, **kw): 

82 return "ARRAY" 

83 

84 

85class Base(DeclarativeBase): 

86 type_annotation_map = {datetime.datetime: TIMESTAMP(timezone=True)} 

87 

88 

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") 

100 

101 

102def deliver_update_date(context): 

103 return datetime.datetime.now(datetime.timezone.utc) 

104 

105 

106class Bed(Base): 

107 __tablename__ = "bed" 

108 

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) 

118 

119 header: Mapped[Optional[str]] = mapped_column( 

120 nullable=True, comment="Header of the bed file, it if was provided." 

121 ) 

122 

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 ) 

130 

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) 

139 

140 files: Mapped[list["Files"]] = relationship( 

141 "Files", back_populates="bedfile", cascade="all, delete-orphan" 

142 ) 

143 

144 bedsets: Mapped[list["BedFileBedSetRelation"]] = relationship( 

145 "BedFileBedSetRelation", back_populates="bedfile", cascade="all, delete-orphan" 

146 ) 

147 

148 annotations: Mapped["BedMetadata"] = relationship( 

149 back_populates="bed", cascade="all, delete-orphan", lazy="joined" 

150 ) 

151 

152 stats: Mapped["BedStats"] = relationship( 

153 back_populates="bed", cascade="all, delete-orphan" 

154 ) 

155 

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") 

166 

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 ) 

173 

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 ) 

197 

198 

199class BedMetadata(Base): 

200 __tablename__ = "bed_metadata" 

201 

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 ) 

206 

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 ) 

213 

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 ) 

246 

247 original_file_name: Mapped[Optional[str]] = mapped_column( 

248 nullable=True, comment="Original file name" 

249 ) 

250 

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 ) 

263 

264 id: Mapped[str] = mapped_column( 

265 ForeignKey("bed.id", ondelete="CASCADE"), 

266 primary_key=True, 

267 index=True, 

268 ) 

269 

270 bed: Mapped["Bed"] = relationship("Bed", back_populates="annotations") 

271 

272 

273class BedStats(Base): 

274 __tablename__ = "bed_stats" 

275 

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]] 

300 

301 distributions: Mapped[Optional[dict]] = mapped_column( 

302 JSONB, 

303 nullable=True, 

304 comment="Full distribution arrays from gtars genomicdist (JSONB)", 

305 ) 

306 

307 bed: Mapped["Bed"] = relationship("Bed", back_populates="stats") 

308 

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 ) 

317 

318 

319class Files(Base): 

320 __tablename__ = "files" 

321 

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") 

339 

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 ) 

346 

347 bedfile: Mapped["Bed"] = relationship("Bed", back_populates="files") 

348 bedset: Mapped["BedSets"] = relationship("BedSets", back_populates="files") 

349 

350 __table_args__ = ( 

351 UniqueConstraint( 

352 "name", "bedfile_id" 

353 ), # TODO: add in the future file_digest here 

354 UniqueConstraint("name", "bedset_id"), 

355 ) 

356 

357 

358class BedFileBedSetRelation(Base): 

359 __tablename__ = "bedfile_bedset_relation" 

360 

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 ) 

367 

368 bedset: Mapped["BedSets"] = relationship("BedSets", back_populates="bedfiles") 

369 bedfile: Mapped["Bed"] = relationship("Bed", back_populates="bedsets") 

370 

371 

372class BedSets(Base): 

373 __tablename__ = "bedsets" 

374 

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") 

391 

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 ) 

403 

404 bedfile_count: Mapped[int] = mapped_column( 

405 default=0, comment="Number of bedfiles in the bedset (denormalized count)" 

406 ) 

407 

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") 

413 

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") 

416 

417 processed: Mapped[bool] = mapped_column( 

418 default=False, comment="Whether the bedset was processed" 

419 ) 

420 

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 ) 

444 

445 

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) 

454 

455 

456class Universes(Base): 

457 __tablename__ = "universes" 

458 

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 ) 

472 

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 ) 

479 

480 

481class TokenizedBed(Base): 

482 __tablename__ = "tokenized_bed" 

483 

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 ) 

499 

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 ) 

506 

507 

508class License(Base): 

509 __tablename__ = "licenses" 

510 

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 ) 

517 

518 bed: Mapped[list["Bed"]] = relationship("Bed", back_populates="license_mapping") 

519 

520 

521class ReferenceGenome(Base): 

522 __tablename__ = "reference_genomes" 

523 

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 ) 

528 

529 bed_reference: Mapped[list["GenomeRefStats"]] = relationship( 

530 "GenomeRefStats", 

531 back_populates="genome_object", 

532 cascade="all, delete-orphan", 

533 ) 

534 

535 

536class GenomeRefStats(Base): 

537 __tablename__ = "genome_ref_stats" 

538 

539 id: Mapped[int] = mapped_column(primary_key=True, index=True) 

540 

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 ) 

553 

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) 

559 

560 bed: Mapped["Bed"] = relationship("Bed", back_populates="ref_classifier") 

561 

562 genome_object: Mapped["ReferenceGenome"] = relationship( 

563 "ReferenceGenome", 

564 back_populates="bed_reference", 

565 lazy="joined", 

566 ) 

567 

568 __table_args__ = (UniqueConstraint("bed_id", "compared_genome"),) 

569 

570 

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() 

578 

579 

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() 

586 

587 

588class GeoGseStatus(Base): 

589 __tablename__ = "geo_gse_status" 

590 

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") 

605 

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") 

610 

611 

612class GeoGsmStatus(Base): 

613 __tablename__ = "geo_gsm_status" 

614 

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") 

625 

626 source_submission_date: Mapped[datetime.datetime] = mapped_column( 

627 nullable=True, comment="Submission date of the source" 

628 ) 

629 

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 ) 

636 

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") 

641 

642 gse_status_mapper: Mapped["GeoGseStatus"] = relationship( 

643 "GeoGseStatus", back_populates="gsm_status_mapper" 

644 ) 

645 

646 

647class UsageBedMeta(Base): 

648 __tablename__ = "usage_bed_meta" 

649 

650 id: Mapped[int] = mapped_column(primary_key=True, index=True, autoincrement=True) 

651 

652 bed_id: Mapped[str] = mapped_column( 

653 ForeignKey("bed.id", ondelete="CASCADE"), nullable=True, index=True 

654 ) 

655 

656 count: Mapped[int] = mapped_column(default=0, comment="Number of visits") 

657 

658 date_from: Mapped[datetime.datetime] = mapped_column(comment="Date from") 

659 date_to: Mapped[datetime.datetime] = mapped_column(comment="Date to") 

660 

661 

662class UsageBedSetMeta(Base): 

663 __tablename__ = "usage_bedset_meta" 

664 

665 id: Mapped[int] = mapped_column(primary_key=True, index=True, autoincrement=True) 

666 

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") 

671 

672 date_from: Mapped[datetime.datetime] = mapped_column(comment="Date from") 

673 date_to: Mapped[datetime.datetime] = mapped_column(comment="Date to") 

674 

675 

676class UsageFiles(Base): 

677 __tablename__ = "usage_files" 

678 

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") 

682 

683 date_from: Mapped[datetime.datetime] = mapped_column(comment="Date from") 

684 date_to: Mapped[datetime.datetime] = mapped_column(comment="Date to") 

685 

686 

687class UsageSearch(Base): 

688 __tablename__ = "usage_search" 

689 

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") 

696 

697 date_from: Mapped[datetime.datetime] = mapped_column(comment="Date from") 

698 date_to: Mapped[datetime.datetime] = mapped_column(comment="Date to") 

699 

700 

701class BedSnapshot(Base): 

702 """ 

703 Index of bulk metadata exports published to S3. 

704 

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 """ 

710 

711 __tablename__ = "bed_snapshots" 

712 

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 ) 

735 

736 

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. 

741 

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 """ 

746 

747 __tablename__ = "analysis_files" 

748 

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 ) 

777 

778 

779class BaseEngine: 

780 """ 

781 A class with base methods, that are used in several classes. 

782 """ 

783 

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. 

800 

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 ) 

823 

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) 

830 

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() 

840 

841 def create_schema(self, engine=None): 

842 """ 

843 Create sql schema in the database. 

844 

845 Args: 

846 engine: Sqlalchemy engine [Default: None]. 

847 

848 Returns: 

849 None. 

850 """ 

851 if not engine: 

852 engine = self._engine 

853 Base.metadata.create_all(engine) 

854 

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 

862 

863 def delete_schema(self, engine=None) -> None: 

864 """ 

865 Delete sql schema in the database. 

866 

867 Args: 

868 engine: Sqlalchemy engine [Default: None]. 

869 

870 Returns: 

871 None. 

872 """ 

873 if not engine: 

874 engine = self._engine 

875 Base.metadata.drop_all(engine) 

876 return None 

877 

878 def run_db_migration(self, database_url: str) -> None: 

879 """ 

880 Upgrade the database to the latest Alembic revision (``head``). 

881 

882 The Alembic config is built programmatically so the package does not 

883 depend on the repo-root ``alembic.ini`` at runtime. 

884 

885 Args: 

886 database_url: SQLAlchemy connection URL (with password) to migrate. 

887 """ 

888 script_location = os.path.join(os.path.dirname(__file__), "alembic") 

889 

890 alembic_cfg = Config() 

891 alembic_cfg.set_main_option("script_location", script_location) 

892 alembic_cfg.set_main_option("sqlalchemy.url", database_url) 

893 

894 _LOGGER.info("Running database migrations to the latest revision...") 

895 command.upgrade(alembic_cfg, "head") 

896 

897 def session_execute(self, statement: Select) -> Result: 

898 """ 

899 Execute statement using sqlalchemy statement. 

900 

901 Args: 

902 statement: SQL query or a SQL expression that is constructed using 

903 SQLAlchemy's SQL expression language. 

904 

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) 

911 

912 return query_result 

913 

914 @property 

915 def session(self): 

916 """ 

917 Get a started sqlalchemy session. 

918 

919 Returns: 

920 Started sqlalchemy session. 

921 """ 

922 return self._start_session() 

923 

924 @property 

925 def engine(self) -> Engine: 

926 """ 

927 Get sqlalchemy engine. 

928 

929 Returns: 

930 Sqlalchemy engine. 

931 """ 

932 return self._engine 

933 

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() 

940 

941 return session 

942 

943 def check_db_connection(self): 

944 try: 

945 self.session_execute(select(Bed).limit(1)) 

946 except ProgrammingError: 

947 raise SchemaError() 

948 

949 def create_schema_graph(self, output_file: str = "schema.svg"): 

950 """ 

951 Create schema graph of the database. 

952 

953 Args: 

954 output_file: Path to the output file. 

955 

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 

962 

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 

973 

974 _LOGGER.info("Uploading licenses to the database...") 

975 df = pd.read_csv(LICENSES_CSV_URL) 

976 

977 df.to_sql(License.__tablename__, self.engine, if_exists="append", index=False) 

978 

979 _LOGGER.info("Licenses uploaded successfully!")