Coverage for pepdbagent/db_utils.py: 91%

1import datetime 

2import enum 

3import logging 

4import os 

5 

6from alembic import command 

7from alembic.config import Config 

8from sqlalchemy import ( 

9 TIMESTAMP, 

10 BigInteger, 

11 Enum, 

12 FetchedValue, 

13 ForeignKey, 

14 Result, 

15 Select, 

16 String, 

17 UniqueConstraint, 

18 event, 

19 select, 

20) 

21from sqlalchemy.dialects.postgresql import JSON 

22from sqlalchemy.engine import URL, Engine, create_engine 

23from sqlalchemy.exc import ProgrammingError 

24from sqlalchemy.ext.compiler import compiles 

25from sqlalchemy.orm import DeclarativeBase, Mapped, Session, mapped_column, relationship 

26 

27from pepdbagent.const import PKG_NAME, POSTGRES_DIALECT 

28from pepdbagent.exceptions import SchemaError 

29 

30_LOGGER = logging.getLogger(PKG_NAME) 

31 

32 

33class BIGSERIAL(BigInteger): 

34 pass 

35 

36 

37@compiles(BIGSERIAL, POSTGRES_DIALECT) 

38def compile_bigserial_pg(type_, compiler, **kw) -> str: 

39 return "BIGSERIAL" 

40 

41 

42@compiles(JSON, POSTGRES_DIALECT) 

43def compile_jsonb_pg(type_, compiler, **kw) -> str: 

44 return "JSON" 

45 

46 

47class Base(DeclarativeBase): 

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

49 

50 

51@event.listens_for(Base.metadata, "after_create") 

52def receive_after_create(target, connection, tables, **kw) -> None: 

53 """Listen for the 'after_create' event.""" 

54 if tables: 

55 _LOGGER.info("A table was created") 

56 else: 

57 _LOGGER.info("A table was not created") 

58 

59 

60# def deliver_description(context): 

61# return context.get_current_parameters()["config"]["description"] 

62 

63 

64def deliver_update_date(context) -> datetime.datetime: 

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

66 

67 

68class Projects(Base): 

69 """ 

70 Projects table representation in the database 

71 """ 

72 

73 __tablename__ = "projects" 

74 

75 id: Mapped[int] = mapped_column(primary_key=True) 

76 namespace: Mapped[str] = mapped_column( 

77 ForeignKey("users.namespace", ondelete="CASCADE") 

78 ) 

79 name: Mapped[str] = mapped_column() 

80 tag: Mapped[str] = mapped_column() 

81 digest: Mapped[str] = mapped_column(String(32)) 

82 description: Mapped[str | None] 

83 config: Mapped[dict] = mapped_column(JSON, server_default=FetchedValue()) 

84 private: Mapped[bool] 

85 number_of_samples: Mapped[int] 

86 number_of_stars: Mapped[int] = mapped_column(default=0) 

87 submission_date: Mapped[datetime.datetime] 

88 last_update_date: Mapped[datetime.datetime | None] = mapped_column( 

89 default=deliver_update_date, # onupdate=deliver_update_date, # This field should not be updated, while we are adding project to favorites 

90 ) 

91 

92 schema_id: Mapped[int | None] = mapped_column( 

93 ForeignKey("schema_versions.id", ondelete="SET NULL"), nullable=True 

94 ) 

95 schema_mapping: Mapped["SchemaVersions"] = relationship( 

96 "SchemaVersions", lazy="joined" 

97 ) 

98 

99 pop: Mapped[bool | None] = mapped_column(default=False) 

100 samples_mapping: Mapped[list["Samples"]] = relationship( 

101 back_populates="project_mapping", cascade="all, delete-orphan" 

102 ) 

103 subsamples_mapping: Mapped[list["Subsamples"]] = relationship( 

104 back_populates="subsample_mapping", cascade="all, delete-orphan" 

105 ) 

106 stars_mapping: Mapped[list["Stars"]] = relationship( 

107 back_populates="project_mapping", cascade="all, delete-orphan" 

108 ) 

109 views_mapping: Mapped[list["Views"]] = relationship( 

110 back_populates="project_mapping", cascade="all, delete-orphan" 

111 ) 

112 

113 # Self-referential relationship. The parent project is the one that was forked to create this one. 

114 forked_from_id: Mapped[int | None] = mapped_column( 

115 ForeignKey("projects.id", ondelete="SET NULL"), nullable=True 

116 ) 

117 forked_from_mapping = relationship( 

118 "Projects", 

119 back_populates="forked_to_mapping", 

120 remote_side=[id], 

121 single_parent=True, 

122 cascade="save-update, merge, refresh-expire", 

123 ) 

124 

125 forked_to_mapping = relationship( 

126 "Projects", 

127 back_populates="forked_from_mapping", 

128 cascade="save-update, merge, refresh-expire", 

129 ) 

130 

131 namespace_mapping: Mapped["User"] = relationship( 

132 "User", back_populates="projects_mapping" 

133 ) 

134 

135 history_mapping: Mapped[list["HistoryProjects"]] = relationship( 

136 back_populates="project_mapping", cascade="all, delete-orphan" 

137 ) 

138 

139 __table_args__ = (UniqueConstraint("namespace", "name", "tag"),) 

140 

141 

142class Samples(Base): 

143 """ 

144 Samples table representation in the database 

145 """ 

146 

147 __tablename__ = "samples" 

148 

149 id: Mapped[int] = mapped_column(primary_key=True) 

150 sample: Mapped[dict] = mapped_column(JSON, server_default=FetchedValue()) 

151 project_id = mapped_column(ForeignKey("projects.id", ondelete="CASCADE")) 

152 project_mapping: Mapped["Projects"] = relationship(back_populates="samples_mapping") 

153 sample_name: Mapped[str | None] = mapped_column() 

154 guid: Mapped[str | None] = mapped_column(nullable=False, unique=True) 

155 

156 submission_date: Mapped[datetime.datetime] = mapped_column( 

157 default=deliver_update_date 

158 ) 

159 last_update_date: Mapped[datetime.datetime | None] = mapped_column( 

160 default=deliver_update_date, 

161 onupdate=deliver_update_date, 

162 ) 

163 

164 parent_guid: Mapped[str | None] = mapped_column( 

165 ForeignKey("samples.guid", ondelete="CASCADE"), 

166 nullable=True, 

167 doc="Parent sample id. Used to create a hierarchy of samples.", 

168 ) 

169 

170 parent_mapping: Mapped["Samples"] = relationship( 

171 "Samples", remote_side=guid, back_populates="child_mapping" 

172 ) 

173 child_mapping: Mapped["Samples"] = relationship( 

174 "Samples", back_populates="parent_mapping" 

175 ) 

176 

177 views: Mapped[list["ViewSampleAssociation"] | None] = relationship( 

178 back_populates="sample", cascade="all, delete-orphan" 

179 ) 

180 

181 

182class Subsamples(Base): 

183 """ 

184 Subsamples table representation in the database 

185 """ 

186 

187 __tablename__ = "subsamples" 

188 

189 id: Mapped[int] = mapped_column(primary_key=True) 

190 subsample: Mapped[dict] = mapped_column(JSON, server_default=FetchedValue()) 

191 subsample_number: Mapped[int] 

192 row_number: Mapped[int] 

193 project_id = mapped_column(ForeignKey("projects.id", ondelete="CASCADE")) 

194 subsample_mapping: Mapped["Projects"] = relationship( 

195 back_populates="subsamples_mapping" 

196 ) 

197 

198 

199class User(Base): 

200 """ 

201 User table representation in the database 

202 """ 

203 

204 __tablename__ = "users" 

205 

206 id: Mapped[int] = mapped_column(primary_key=True) 

207 namespace: Mapped[str] = mapped_column(nullable=False, unique=True) 

208 stars_mapping: Mapped[list["Stars"]] = relationship( 

209 back_populates="user_mapping", 

210 cascade="all, delete-orphan", 

211 order_by="Stars.star_date.desc()", 

212 ) 

213 number_of_projects: Mapped[int] = mapped_column(default=0) 

214 number_of_schemas: Mapped[int] = mapped_column(default=0) 

215 

216 projects_mapping: Mapped[list["Projects"]] = relationship( 

217 "Projects", back_populates="namespace_mapping" 

218 ) 

219 schemas_mapping: Mapped[list["SchemaRecords"]] = relationship( 

220 "SchemaRecords", back_populates="user_mapping" 

221 ) 

222 

223 

224class Stars(Base): 

225 """ 

226 FavoriteProjects table representation in the database 

227 """ 

228 

229 __tablename__ = "stars" 

230 

231 user_id = mapped_column( 

232 ForeignKey("users.id", ondelete="CASCADE"), primary_key=True 

233 ) 

234 project_id = mapped_column( 

235 ForeignKey("projects.id", ondelete="CASCADE"), primary_key=True 

236 ) 

237 user_mapping: Mapped[list["User"]] = relationship(back_populates="stars_mapping") 

238 project_mapping: Mapped["Projects"] = relationship(back_populates="stars_mapping") 

239 star_date: Mapped[datetime.datetime] = mapped_column( 

240 onupdate=deliver_update_date, default=deliver_update_date 

241 ) 

242 

243 

244class Views(Base): 

245 """ 

246 Views table representation in the database 

247 """ 

248 

249 __tablename__ = "views" 

250 

251 id: Mapped[int] = mapped_column(primary_key=True) 

252 name: Mapped[str] = mapped_column() 

253 description: Mapped[str | None] 

254 

255 project_id = mapped_column(ForeignKey("projects.id", ondelete="CASCADE")) 

256 project_mapping = relationship("Projects", back_populates="views_mapping") 

257 

258 samples: Mapped[list["ViewSampleAssociation"]] = relationship( 

259 back_populates="view", cascade="all, delete-orphan" 

260 ) 

261 

262 _table_args__ = (UniqueConstraint("name", "project_id"),) 

263 

264 

265class ViewSampleAssociation(Base): 

266 """ 

267 Association table between views and samples 

268 """ 

269 

270 __tablename__ = "views_samples" 

271 

272 sample_id = mapped_column( 

273 ForeignKey("samples.id", ondelete="CASCADE"), primary_key=True 

274 ) 

275 view_id = mapped_column( 

276 ForeignKey("views.id", ondelete="CASCADE"), primary_key=True 

277 ) 

278 sample: Mapped["Samples"] = relationship(back_populates="views") 

279 view: Mapped["Views"] = relationship(back_populates="samples") 

280 

281 

282class HistoryProjects(Base): 

283 __tablename__ = "project_history" 

284 

285 id: Mapped[int] = mapped_column(primary_key=True) 

286 project_id: Mapped[int] = mapped_column( 

287 ForeignKey("projects.id", ondelete="CASCADE") 

288 ) 

289 user: Mapped[str] = mapped_column( 

290 ForeignKey("users.namespace", ondelete="SET NULL") 

291 ) 

292 update_time: Mapped[datetime.datetime] = mapped_column( 

293 TIMESTAMP(timezone=True), default=deliver_update_date 

294 ) 

295 project_yaml: Mapped[dict] = mapped_column(JSON, server_default=FetchedValue()) 

296 

297 project_mapping: Mapped["Projects"] = relationship( 

298 "Projects", back_populates="history_mapping" 

299 ) 

300 sample_changes_mapping: Mapped[list["HistorySamples"]] = relationship( 

301 back_populates="history_project_mapping", cascade="all, delete-orphan" 

302 ) 

303 

304 

305class UpdateTypes(enum.Enum): 

306 """ 

307 Enum for the type of update 

308 """ 

309 

310 UPDATE = "update" 

311 INSERT = "insert" 

312 DELETE = "delete" 

313 

314 

315class HistorySamples(Base): 

316 __tablename__ = "sample_history" 

317 

318 id: Mapped[int] = mapped_column(primary_key=True) 

319 history_id: Mapped[int] = mapped_column( 

320 ForeignKey("project_history.id", ondelete="CASCADE") 

321 ) 

322 guid: Mapped[str] = mapped_column(nullable=False) 

323 parent_guid: Mapped[str | None] = mapped_column(nullable=True) 

324 sample_json: Mapped[dict] = mapped_column(JSON, server_default=FetchedValue()) 

325 change_type: Mapped[UpdateTypes] = mapped_column(Enum(UpdateTypes), nullable=False) 

326 

327 history_project_mapping: Mapped["HistoryProjects"] = relationship( 

328 "HistoryProjects", back_populates="sample_changes_mapping" 

329 ) 

330 

331 

332class SchemaRecords(Base): 

333 __tablename__ = "schema_records" 

334 

335 id: Mapped[int] = mapped_column(primary_key=True) 

336 namespace: Mapped[str] = mapped_column( 

337 ForeignKey("users.namespace", ondelete="CASCADE") 

338 ) 

339 name: Mapped[str] = mapped_column(nullable=False) 

340 maintainers: Mapped[str] = mapped_column(nullable=True) 

341 lifecycle_stage: Mapped[str] = mapped_column(nullable=True) 

342 description: Mapped[str | None] = mapped_column(nullable=True) 

343 private: Mapped[bool] = mapped_column(default=False) 

344 last_update_date: Mapped[datetime.datetime | None] = mapped_column( 

345 default=deliver_update_date, onupdate=deliver_update_date 

346 ) 

347 

348 __table_args__ = (UniqueConstraint("namespace", "name"),) 

349 

350 versions_mapping: Mapped[list["SchemaVersions"]] = relationship( 

351 "SchemaVersions", 

352 back_populates="schema_mapping", 

353 cascade="all, delete-orphan", 

354 order_by="SchemaVersions.version.desc()", 

355 ) 

356 user_mapping: Mapped["User"] = relationship( 

357 "User", back_populates="schemas_mapping" 

358 ) 

359 

360 

361class SchemaVersions(Base): 

362 __tablename__ = "schema_versions" 

363 

364 id: Mapped[int] = mapped_column(primary_key=True) 

365 schema_id: Mapped[int] = mapped_column( 

366 ForeignKey("schema_records.id", ondelete="CASCADE") 

367 ) 

368 version: Mapped[str] = mapped_column(nullable=False) 

369 schema_value: Mapped[dict] = mapped_column(JSON, server_default=FetchedValue()) 

370 release_date: Mapped[datetime.datetime] = mapped_column(default=deliver_update_date) 

371 last_update_date: Mapped[datetime.datetime | None] = mapped_column( 

372 default=deliver_update_date, onupdate=deliver_update_date 

373 ) 

374 contributors: Mapped[str | None] = mapped_column(nullable=True) 

375 release_notes: Mapped[str | None] = mapped_column(nullable=True) 

376 

377 __table_args__ = (UniqueConstraint("schema_id", "version"),) 

378 

379 schema_mapping: Mapped["SchemaRecords"] = relationship( 

380 "SchemaRecords", back_populates="versions_mapping" 

381 ) 

382 

383 tags_mapping: Mapped[list["SchemaTags"]] = relationship( 

384 "SchemaTags", 

385 back_populates="schema_mapping", 

386 lazy="joined", 

387 cascade="all, delete-orphan", 

388 ) 

389 

390 

391class SchemaTags(Base): 

392 __tablename__ = "schema_tags" 

393 

394 id: Mapped[int] = mapped_column(primary_key=True) 

395 tag_name: Mapped[str] = mapped_column(nullable=False) 

396 tag_value: Mapped[str] = mapped_column(nullable=True) 

397 schema_version_id: Mapped[int] = mapped_column( 

398 ForeignKey("schema_versions.id", ondelete="CASCADE") 

399 ) 

400 

401 schema_mapping: Mapped["SchemaVersions"] = relationship( 

402 "SchemaVersions", back_populates="tags_mapping" 

403 ) 

404 

405 

406class TarNamespace(Base): 

407 __tablename__ = "namespace_archives" 

408 

409 id: Mapped[int] = mapped_column(primary_key=True) 

410 namespace: Mapped[str] = mapped_column( 

411 ForeignKey("users.namespace", ondelete="CASCADE") 

412 ) 

413 file_path: Mapped[str] = mapped_column(nullable=False) 

414 creation_date: Mapped[datetime.datetime] = mapped_column( 

415 default=deliver_update_date 

416 ) 

417 number_of_projects: Mapped[int] = mapped_column(default=0) 

418 file_size: Mapped[int] = mapped_column(nullable=False) 

419 

420 

421class BedBaseStats(Base): 

422 __tablename__ = "bedbase_stats" 

423 

424 id: Mapped[int] = mapped_column(primary_key=True) 

425 gse: Mapped[str] = mapped_column() 

426 gsm: Mapped[str] = mapped_column() 

427 sample_name: Mapped[str] = mapped_column(nullable=True) 

428 genome: Mapped[str | None] = mapped_column(nullable=True, default="") 

429 last_update_date: Mapped[str | None] = mapped_column() 

430 submission_date: Mapped[str | None] = mapped_column() 

431 

432 

433class BaseEngine: 

434 """ 

435 A class with base methods, that are used in several classes. e.g. fetch_one or fetch_all 

436 """ 

437 

438 def __init__( 

439 self, 

440 *, 

441 host: str = "localhost", 

442 port: int = 5432, 

443 database: str = "pep-db", 

444 user: str | None = None, 

445 password: str | None = None, 

446 drivername: str = POSTGRES_DIALECT, 

447 dsn: str | None = None, 

448 echo: bool = False, 

449 run_migrations: bool = False, 

450 ) -> None: 

451 """Initialize connection to the pep_db database. 

452 

453 Args: 

454 host: Database server address. 

455 port: Port number (default: 5432). 

456 database: Database name. 

457 user: Username for authentication. 

458 password: Password for authentication. 

459 drivername: Database driver. 

460 dsn: libpq connection string, e.g., "postgresql://user:pass@host:port/db". 

461 echo: Log all SQL statements if True. 

462 run_migrations: Run database migrations if True. 

463 """ 

464 if not dsn: 

465 dsn = URL.create( 

466 host=host, 

467 port=port, 

468 database=database, 

469 username=user, 

470 password=password, 

471 drivername=drivername, 

472 ) 

473 if run_migrations: 

474 dsn_with_password = ( 

475 f"{drivername}://{user}:{password}@{host}:{port}/{database}" 

476 ) 

477 self.run_db_migration(dsn_with_password) 

478 

479 self._engine = create_engine(dsn, echo=echo) 

480 self.create_schema(self._engine) 

481 self.check_db_connection() 

482 

483 def create_schema(self, engine=None) -> None: 

484 """Create SQL schema in the database. 

485 

486 Args: 

487 engine: SQLAlchemy engine (default: uses internal engine). 

488 """ 

489 if not engine: 

490 engine = self._engine 

491 Base.metadata.create_all(engine) 

492 return None 

493 

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

495 """Execute a SQLAlchemy statement. 

496 

497 Args: 

498 statement: SQL query or expression to execute. 

499 

500 Returns: 

501 Query result. 

502 """ 

503 _LOGGER.debug(f"Executing statement: {statement}") 

504 with Session(self._engine) as session: 

505 query_result = session.execute(statement) 

506 

507 return query_result 

508 

509 @property 

510 def session(self) -> Session: 

511 """Return a started SQLAlchemy session.""" 

512 return self._start_session() 

513 

514 @property 

515 def engine(self) -> Engine: 

516 return self._engine 

517 

518 def _start_session(self) -> Session: 

519 session = Session(self.engine) 

520 try: 

521 session.execute(select(Projects).limit(1)) 

522 except ProgrammingError: 

523 raise SchemaError() 

524 

525 return session 

526 

527 def check_db_connection(self) -> None: 

528 try: 

529 self.session_execute(select(Projects).limit(1)) 

530 except ProgrammingError: 

531 raise SchemaError() 

532 

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

534 """Delete SQL schema from the database. 

535 

536 Args: 

537 engine: SQLAlchemy engine (default: uses internal engine). 

538 """ 

539 if not engine: 

540 engine = self._engine 

541 Base.metadata.drop_all(engine) 

542 return None 

543 

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

545 """Migrate the database to the latest version.""" 

546 script_directory = os.path.dirname(os.path.abspath(__file__)) 

547 script_location = os.path.join(script_directory, "alembic") 

548 

549 alembic_cfg = Config() 

550 

551 alembic_cfg.set_main_option("script_location", script_location) 

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

553 

554 command.upgrade(alembic_cfg, "head")