Hide keyboard shortcuts

Hot-keys on this page

r m x p   toggle line displays

j k   next/prev highlighted chunk

0   (zero) top of page

1   (one) first highlighted chunk

1# This file is part of daf_butler. 

2# 

3# Developed for the LSST Data Management System. 

4# This product includes software developed by the LSST Project 

5# (http://www.lsst.org). 

6# See the COPYRIGHT file at the top-level directory of this distribution 

7# for details of code ownership. 

8# 

9# This program is free software: you can redistribute it and/or modify 

10# it under the terms of the GNU General Public License as published by 

11# the Free Software Foundation, either version 3 of the License, or 

12# (at your option) any later version. 

13# 

14# This program is distributed in the hope that it will be useful, 

15# but WITHOUT ANY WARRANTY; without even the implied warranty of 

16# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 

17# GNU General Public License for more details. 

18# 

19# You should have received a copy of the GNU General Public License 

20# along with this program. If not, see <http://www.gnu.org/licenses/>. 

21from __future__ import annotations 

22 

23__all__ = [ 

24 "Database", 

25 "ReadOnlyDatabaseError", 

26 "DatabaseConflictError", 

27 "StaticTablesContext", 

28] 

29 

30from abc import ABC, abstractmethod 

31from contextlib import contextmanager 

32from typing import ( 

33 Any, 

34 Dict, 

35 Iterable, 

36 List, 

37 Optional, 

38 Sequence, 

39 Tuple, 

40) 

41import warnings 

42 

43import astropy.time 

44import sqlalchemy 

45 

46from ...core import ddl, time_utils 

47from .._exceptions import ConflictingDefinitionError 

48 

49 

50def _checkExistingTableDefinition(name: str, spec: ddl.TableSpec, inspection: Dict[str, Any]): 

51 """Test that the definition of a table in a `ddl.TableSpec` and from 

52 database introspection are consistent. 

53 

54 Parameters 

55 ---------- 

56 name : `str` 

57 Name of the table (only used in error messages). 

58 spec : `ddl.TableSpec` 

59 Specification of the table. 

60 inspection : `dict` 

61 Dictionary returned by 

62 `sqlalchemy.engine.reflection.Inspector.get_columns`. 

63 

64 Raises 

65 ------ 

66 DatabaseConflictError 

67 Raised if the definitions are inconsistent. 

68 """ 

69 columnNames = [c["name"] for c in inspection] 

70 if spec.fields.names != set(columnNames): 

71 raise DatabaseConflictError(f"Table '{name}' exists but is defined differently in the database; " 

72 f"specification has columns {list(spec.fields.names)}, while the " 

73 f"table in the database has {columnNames}.") 

74 

75 

76class ReadOnlyDatabaseError(RuntimeError): 

77 """Exception raised when a write operation is called on a read-only 

78 `Database`. 

79 """ 

80 

81 

82class DatabaseConflictError(ConflictingDefinitionError): 

83 """Exception raised when database content (row values or schema entities) 

84 are inconsistent with what this client expects. 

85 """ 

86 

87 

88class StaticTablesContext: 

89 """Helper class used to declare the static schema for a registry layer 

90 in a database. 

91 

92 An instance of this class is returned by `Database.declareStaticTables`, 

93 which should be the only way it should be constructed. 

94 """ 

95 

96 def __init__(self, db: Database): 

97 self._db = db 

98 self._foreignKeys = [] 

99 self._inspector = sqlalchemy.engine.reflection.Inspector(self._db._connection) 

100 self._tableNames = frozenset(self._inspector.get_table_names(schema=self._db.namespace)) 

101 

102 def addTable(self, name: str, spec: ddl.TableSpec) -> sqlalchemy.schema.Table: 

103 """Add a new table to the schema, returning its sqlalchemy 

104 representation. 

105 

106 The new table may not actually be created until the end of the 

107 context created by `Database.declareStaticTables`, allowing tables 

108 to be declared in any order even in the presence of foreign key 

109 relationships. 

110 """ 

111 name = self._db._mangleTableName(name) 

112 if name in self._tableNames: 

113 _checkExistingTableDefinition(name, spec, self._inspector.get_columns(name, 

114 schema=self._db.namespace)) 

115 table = self._db._convertTableSpec(name, spec, self._db._metadata) 

116 for foreignKeySpec in spec.foreignKeys: 

117 self._foreignKeys.append( 

118 (table, self._db._convertForeignKeySpec(name, foreignKeySpec, self._db._metadata)) 

119 ) 

120 return table 

121 

122 def addTableTuple(self, specs: Tuple[ddl.TableSpec, ...]) -> Tuple[sqlalchemy.schema.Table, ...]: 

123 """Add a named tuple of tables to the schema, returning their 

124 SQLAlchemy representations in a named tuple of the same type. 

125 

126 The new tables may not actually be created until the end of the 

127 context created by `Database.declareStaticTables`, allowing tables 

128 to be declared in any order even in the presence of foreign key 

129 relationships. 

130 

131 Notes 

132 ----- 

133 ``specs`` *must* be an instance of a type created by 

134 `collections.namedtuple`, not just regular tuple, and the returned 

135 object is guaranteed to be the same. Because `~collections.namedtuple` 

136 is just a factory for `type` objects, not an actual type itself, 

137 we cannot represent this with type annotations. 

138 """ 

139 return specs._make(self.addTable(name, spec) for name, spec in zip(specs._fields, specs)) 

140 

141 

142class Database(ABC): 

143 """An abstract interface that represents a particular database engine's 

144 representation of a single schema/namespace/database. 

145 

146 Parameters 

147 ---------- 

148 origin : `int` 

149 An integer ID that should be used as the default for any datasets, 

150 quanta, or other entities that use a (autoincrement, origin) compound 

151 primary key. 

152 connection : `sqlalchemy.engine.Connection` 

153 The SQLAlchemy connection this `Database` wraps. 

154 namespace : `str`, optional 

155 Name of the schema or namespace this instance is associated with. 

156 This is passed as the ``schema`` argument when constructing a 

157 `sqlalchemy.schema.MetaData` instance. We use ``namespace`` instead to 

158 avoid confusion between "schema means namespace" and "schema means 

159 table definitions". 

160 

161 Notes 

162 ----- 

163 `Database` requires all write operations to go through its special named 

164 methods. Our write patterns are sufficiently simple that we don't really 

165 need the full flexibility of SQL insert/update/delete syntax, and we need 

166 non-standard (but common) functionality in these operations sufficiently 

167 often that it seems worthwhile to provide our own generic API. 

168 

169 In contrast, `Database.query` allows arbitrary ``SELECT`` queries (via 

170 their SQLAlchemy representation) to be run, as we expect these to require 

171 significantly more sophistication while still being limited to standard 

172 SQL. 

173 

174 `Database` itself has several underscore-prefixed attributes: 

175 

176 - ``_cs``: SQLAlchemy objects representing the connection and transaction 

177 state. 

178 - ``_metadata``: the `sqlalchemy.schema.MetaData` object representing 

179 the tables and other schema entities. 

180 

181 These are considered protected (derived classes may access them, but other 

182 code should not), and read-only, aside from executing SQL via 

183 ``_connection``. 

184 """ 

185 

186 def __init__(self, *, origin: int, connection: sqlalchemy.engine.Connection, 

187 namespace: Optional[str] = None): 

188 self.origin = origin 

189 self.namespace = namespace 

190 self._connection = connection 

191 self._metadata = None 

192 

193 @classmethod 

194 def makeDefaultUri(cls, root: str) -> Optional[str]: 

195 """Create a default connection URI appropriate for the given root 

196 directory, or `None` if there can be no such default. 

197 """ 

198 return None 

199 

200 @classmethod 

201 def fromUri(cls, uri: str, *, origin: int, namespace: Optional[str] = None, 

202 writeable: bool = True) -> Database: 

203 """Construct a database from a SQLAlchemy URI. 

204 

205 Parameters 

206 ---------- 

207 uri : `str` 

208 A SQLAlchemy URI connection string. 

209 origin : `int` 

210 An integer ID that should be used as the default for any datasets, 

211 quanta, or other entities that use a (autoincrement, origin) 

212 compound primary key. 

213 namespace : `str`, optional 

214 A database namespace (i.e. schema) the new instance should be 

215 associated with. If `None` (default), the namespace (if any) is 

216 inferred from the URI. 

217 writeable : `bool`, optional 

218 If `True`, allow write operations on the database, including 

219 ``CREATE TABLE``. 

220 

221 Returns 

222 ------- 

223 db : `Database` 

224 A new `Database` instance. 

225 """ 

226 return cls.fromConnection(cls.connect(uri, writeable=writeable), 

227 origin=origin, 

228 namespace=namespace, 

229 writeable=writeable) 

230 

231 @classmethod 

232 @abstractmethod 

233 def connect(cls, uri: str, *, writeable: bool = True) -> sqlalchemy.engine.Connection: 

234 """Create a `sqlalchemy.engine.Connection` from a SQLAlchemy URI. 

235 

236 Parameters 

237 ---------- 

238 uri : `str` 

239 A SQLAlchemy URI connection string. 

240 origin : `int` 

241 An integer ID that should be used as the default for any datasets, 

242 quanta, or other entities that use a (autoincrement, origin) 

243 compound primary key. 

244 writeable : `bool`, optional 

245 If `True`, allow write operations on the database, including 

246 ``CREATE TABLE``. 

247 

248 Returns 

249 ------- 

250 connection : `sqlalchemy.engine.Connection` 

251 A database connection. 

252 

253 Notes 

254 ----- 

255 Subclasses that support other ways to connect to a database are 

256 encouraged to add optional arguments to their implementation of this 

257 method, as long as they maintain compatibility with the base class 

258 call signature. 

259 """ 

260 raise NotImplementedError() 

261 

262 @classmethod 

263 @abstractmethod 

264 def fromConnection(cls, connection: sqlalchemy.engine.Connection, *, origin: int, 

265 namespace: Optional[str] = None, writeable: bool = True) -> Database: 

266 """Create a new `Database` from an existing 

267 `sqlalchemy.engine.Connection`. 

268 

269 Parameters 

270 ---------- 

271 connection : `sqllachemy.engine.Connection` 

272 The connection for the the database. May be shared between 

273 `Database` instances. 

274 origin : `int` 

275 An integer ID that should be used as the default for any datasets, 

276 quanta, or other entities that use a (autoincrement, origin) 

277 compound primary key. 

278 namespace : `str`, optional 

279 A different database namespace (i.e. schema) the new instance 

280 should be associated with. If `None` (default), the namespace 

281 (if any) is inferred from the connection. 

282 writeable : `bool`, optional 

283 If `True`, allow write operations on the database, including 

284 ``CREATE TABLE``. 

285 

286 Returns 

287 ------- 

288 db : `Database` 

289 A new `Database` instance. 

290 

291 Notes 

292 ----- 

293 This method allows different `Database` instances to share the same 

294 connection, which is desirable when they represent different namespaces 

295 can be queried together. This also ties their transaction state, 

296 however; starting a transaction in any database automatically starts 

297 on in all other databases. 

298 """ 

299 raise NotImplementedError() 

300 

301 @contextmanager 

302 def transaction(self, *, interrupting: bool = False) -> None: 

303 """Return a context manager that represents a transaction. 

304 

305 Parameters 

306 ---------- 

307 interrupting : `bool` 

308 If `True`, this transaction block needs to be able to interrupt 

309 any existing one in order to yield correct behavior. 

310 """ 

311 assert not (interrupting and self._connection.in_transaction()), ( 

312 "Logic error in transaction nesting: an operation that would " 

313 "interrupt the active transaction context has been requested." 

314 ) 

315 if self._connection.in_transaction(): 

316 trans = self._connection.begin_nested() 

317 else: 

318 # Use a regular (non-savepoint) transaction only for the outermost 

319 # context. 

320 trans = self._connection.begin() 

321 try: 

322 yield 

323 trans.commit() 

324 except BaseException: 

325 trans.rollback() 

326 raise 

327 

328 @contextmanager 

329 def declareStaticTables(self, *, create: bool) -> StaticTablesContext: 

330 """Return a context manager in which the database's static DDL schema 

331 can be declared. 

332 

333 Parameters 

334 ---------- 

335 create : `bool` 

336 If `True`, attempt to create all tables at the end of the context. 

337 If `False`, they will be assumed to already exist. 

338 

339 Returns 

340 ------- 

341 schema : `StaticTablesContext` 

342 A helper object that is used to add new tables. 

343 

344 Raises 

345 ------ 

346 ReadOnlyDatabaseError 

347 Raised if ``create`` is `True`, `Database.isWriteable` is `False`, 

348 and one or more declared tables do not already exist. 

349 

350 Examples 

351 -------- 

352 Given a `Database` instance ``db``:: 

353 

354 with db.declareStaticTables(create=True) as schema: 

355 schema.addTable("table1", TableSpec(...)) 

356 schema.addTable("table2", TableSpec(...)) 

357 

358 Notes 

359 ----- 

360 A database's static DDL schema must be declared before any dynamic 

361 tables are managed via calls to `ensureTableExists` or 

362 `getExistingTable`. The order in which static schema tables are added 

363 inside the context block is unimportant; they will automatically be 

364 sorted and added in an order consistent with their foreign key 

365 relationships. 

366 """ 

367 if create and not self.isWriteable(): 

368 raise ReadOnlyDatabaseError(f"Cannot create tables in read-only database {self}.") 

369 self._metadata = sqlalchemy.MetaData(schema=self.namespace) 

370 try: 

371 context = StaticTablesContext(self) 

372 yield context 

373 for table, foreignKey in context._foreignKeys: 

374 table.append_constraint(foreignKey) 

375 if create: 

376 if self.namespace is not None: 

377 if self.namespace not in context._inspector.get_schema_names(): 

378 self._connection.execute(sqlalchemy.schema.CreateSchema(self.namespace)) 

379 # In our tables we have columns that make use of sqlalchemy 

380 # Sequence objects. There is currently a bug in sqlalchmey that 

381 # causes a deprecation warning to be thrown on a property of 

382 # the Sequence object when the repr for the sequence is 

383 # created. Here a filter is used to catch these deprecation 

384 # warnings when tables are created. 

385 with warnings.catch_warnings(): 

386 warnings.simplefilter("ignore", category=sqlalchemy.exc.SADeprecationWarning) 

387 self._metadata.create_all(self._connection) 

388 except BaseException: 

389 self._metadata.drop_all(self._connection) 

390 self._metadata = None 

391 raise 

392 

393 @abstractmethod 

394 def isWriteable(self) -> bool: 

395 """Return `True` if this database can be modified by this client. 

396 """ 

397 raise NotImplementedError() 

398 

399 @abstractmethod 

400 def __str__(self) -> str: 

401 """Return a human-readable identifier for this `Database`, including 

402 any namespace or schema that identifies its names within a `Registry`. 

403 """ 

404 raise NotImplementedError() 

405 

406 def shrinkDatabaseEntityName(self, original: str) -> str: 

407 """Return a version of the given name that fits within this database 

408 engine's length limits for table, constraint, indexes, and sequence 

409 names. 

410 

411 Implementations should not assume that simple truncation is safe, 

412 because multiple long names often begin with the same prefix. 

413 

414 The default implementation simply returns the given name. 

415 

416 Parameters 

417 ---------- 

418 original : `str` 

419 The original name. 

420 

421 Returns 

422 ------- 

423 shrunk : `str` 

424 The new, possibly shortened name. 

425 """ 

426 return original 

427 

428 def expandDatabaseEntityName(self, shrunk: str) -> str: 

429 """Retrieve the original name for a database entity that was too long 

430 to fit within the database engine's limits. 

431 

432 Parameters 

433 ---------- 

434 original : `str` 

435 The original name. 

436 

437 Returns 

438 ------- 

439 shrunk : `str` 

440 The new, possibly shortened name. 

441 """ 

442 return shrunk 

443 

444 def _mangleTableName(self, name: str) -> str: 

445 """Map a logical, user-visible table name to the true table name used 

446 in the database. 

447 

448 The default implementation returns the given name unchanged. 

449 

450 Parameters 

451 ---------- 

452 name : `str` 

453 Input table name. Should not include a namespace (i.e. schema) 

454 prefix. 

455 

456 Returns 

457 ------- 

458 mangled : `str` 

459 Mangled version of the table name (still with no namespace prefix). 

460 

461 Notes 

462 ----- 

463 Reimplementations of this method must be idempotent - mangling an 

464 already-mangled name must have no effect. 

465 """ 

466 return name 

467 

468 def _convertFieldSpec(self, table: str, spec: ddl.FieldSpec, metadata: sqlalchemy.MetaData, 

469 **kwds) -> sqlalchemy.schema.Column: 

470 """Convert a `FieldSpec` to a `sqlalchemy.schema.Column`. 

471 

472 Parameters 

473 ---------- 

474 table : `str` 

475 Name of the table this column is being added to. 

476 spec : `FieldSpec` 

477 Specification for the field to be added. 

478 metadata : `sqlalchemy.MetaData` 

479 SQLAlchemy representation of the DDL schema this field's table is 

480 being added to. 

481 **kwds 

482 Additional keyword arguments to forward to the 

483 `sqlalchemy.schema.Column` constructor. This is provided to make 

484 it easier for derived classes to delegate to ``super()`` while 

485 making only minor changes. 

486 

487 Returns 

488 ------- 

489 column : `sqlalchemy.schema.Column` 

490 SQLAlchemy representation of the field. 

491 """ 

492 args = [spec.name, spec.getSizedColumnType()] 

493 if spec.autoincrement: 

494 # Generate a sequence to use for auto incrementing for databases 

495 # that do not support it natively. This will be ignored by 

496 # sqlalchemy for databases that do support it. 

497 args.append(sqlalchemy.Sequence(self.shrinkDatabaseEntityName(f"{table}_seq_{spec.name}"), 

498 metadata=metadata)) 

499 assert spec.doc is None or isinstance(spec.doc, str), f"Bad doc for {table}.{spec.name}." 

500 return sqlalchemy.schema.Column(*args, nullable=spec.nullable, primary_key=spec.primaryKey, 

501 comment=spec.doc, **kwds) 

502 

503 def _convertForeignKeySpec(self, table: str, spec: ddl.ForeignKeySpec, metadata: sqlalchemy.MetaData, 

504 **kwds) -> sqlalchemy.schema.ForeignKeyConstraint: 

505 """Convert a `ForeignKeySpec` to a 

506 `sqlalchemy.schema.ForeignKeyConstraint`. 

507 

508 Parameters 

509 ---------- 

510 table : `str` 

511 Name of the table this foreign key is being added to. 

512 spec : `ForeignKeySpec` 

513 Specification for the foreign key to be added. 

514 metadata : `sqlalchemy.MetaData` 

515 SQLAlchemy representation of the DDL schema this constraint is 

516 being added to. 

517 **kwds 

518 Additional keyword arguments to forward to the 

519 `sqlalchemy.schema.ForeignKeyConstraint` constructor. This is 

520 provided to make it easier for derived classes to delegate to 

521 ``super()`` while making only minor changes. 

522 

523 Returns 

524 ------- 

525 constraint : `sqlalchemy.schema.ForeignKeyConstraint` 

526 SQLAlchemy representation of the constraint. 

527 """ 

528 name = self.shrinkDatabaseEntityName( 

529 "_".join(["fkey", table, self._mangleTableName(spec.table)] 

530 + list(spec.target) + list(spec.source)) 

531 ) 

532 return sqlalchemy.schema.ForeignKeyConstraint( 

533 spec.source, 

534 [f"{self._mangleTableName(spec.table)}.{col}" for col in spec.target], 

535 name=name, 

536 ondelete=spec.onDelete 

537 ) 

538 

539 def _convertTableSpec(self, name: str, spec: ddl.TableSpec, metadata: sqlalchemy.MetaData, 

540 **kwds) -> sqlalchemy.schema.Table: 

541 """Convert a `TableSpec` to a `sqlalchemy.schema.Table`. 

542 

543 Parameters 

544 ---------- 

545 spec : `TableSpec` 

546 Specification for the foreign key to be added. 

547 metadata : `sqlalchemy.MetaData` 

548 SQLAlchemy representation of the DDL schema this table is being 

549 added to. 

550 **kwds 

551 Additional keyword arguments to forward to the 

552 `sqlalchemy.schema.Table` constructor. This is provided to make it 

553 easier for derived classes to delegate to ``super()`` while making 

554 only minor changes. 

555 

556 Returns 

557 ------- 

558 table : `sqlalchemy.schema.Table` 

559 SQLAlchemy representation of the table. 

560 

561 Notes 

562 ----- 

563 This method does not handle ``spec.foreignKeys`` at all, in order to 

564 avoid circular dependencies. These are added by higher-level logic in 

565 `ensureTableExists`, `getExistingTable`, and `declareStaticTables`. 

566 """ 

567 name = self._mangleTableName(name) 

568 args = [self._convertFieldSpec(name, fieldSpec, metadata) for fieldSpec in spec.fields] 

569 args.extend( 

570 sqlalchemy.schema.UniqueConstraint( 

571 *columns, 

572 name=self.shrinkDatabaseEntityName("_".join([name, "unq"] + list(columns))) 

573 ) 

574 for columns in spec.unique if columns not in spec.indexes 

575 ) 

576 args.extend( 

577 sqlalchemy.schema.Index( 

578 self.shrinkDatabaseEntityName("_".join([name, "idx"] + list(columns))), 

579 *columns, 

580 unique=(columns in spec.unique) 

581 ) 

582 for columns in spec.indexes 

583 ) 

584 args.extend( 

585 sqlalchemy.schema.Index( 

586 self.shrinkDatabaseEntityName("_".join((name, "fkidx") + fk.source)), 

587 *fk.source, 

588 ) 

589 for fk in spec.foreignKeys if fk.addIndex 

590 ) 

591 assert spec.doc is None or isinstance(spec.doc, str), f"Bad doc for {name}." 

592 return sqlalchemy.schema.Table(name, metadata, *args, comment=spec.doc, info=spec, **kwds) 

593 

594 def ensureTableExists(self, name: str, spec: ddl.TableSpec) -> sqlalchemy.sql.FromClause: 

595 """Ensure that a table with the given name and specification exists, 

596 creating it if necessary. 

597 

598 Parameters 

599 ---------- 

600 name : `str` 

601 Name of the table (not including namespace qualifiers). 

602 spec : `TableSpec` 

603 Specification for the table. This will be used when creating the 

604 table, and *may* be used when obtaining an existing table to check 

605 for consistency, but no such check is guaranteed. 

606 

607 Returns 

608 ------- 

609 table : `sqlalchemy.schema.Table` 

610 SQLAlchemy representation of the table. 

611 

612 Raises 

613 ------ 

614 ReadOnlyDatabaseError 

615 Raised if `isWriteable` returns `False`, and the table does not 

616 already exist. 

617 DatabaseConflictError 

618 Raised if the table exists but ``spec`` is inconsistent with its 

619 definition. 

620 

621 Notes 

622 ----- 

623 This method may not be called within transactions. It may be called on 

624 read-only databases if and only if the table does in fact already 

625 exist. 

626 

627 Subclasses may override this method, but usually should not need to. 

628 """ 

629 assert not self._connection.in_transaction(), "Table creation interrupts transactions." 

630 assert self._metadata is not None, "Static tables must be declared before dynamic tables." 

631 table = self.getExistingTable(name, spec) 

632 if table is not None: 

633 return table 

634 if not self.isWriteable(): 

635 raise ReadOnlyDatabaseError( 

636 f"Table {name} does not exist, and cannot be created " 

637 f"because database {self} is read-only." 

638 ) 

639 table = self._convertTableSpec(name, spec, self._metadata) 

640 for foreignKeySpec in spec.foreignKeys: 

641 table.append_constraint(self._convertForeignKeySpec(name, foreignKeySpec, self._metadata)) 

642 table.create(self._connection) 

643 return table 

644 

645 def getExistingTable(self, name: str, spec: ddl.TableSpec) -> Optional[sqlalchemy.schema.Table]: 

646 """Obtain an existing table with the given name and specification. 

647 

648 Parameters 

649 ---------- 

650 name : `str` 

651 Name of the table (not including namespace qualifiers). 

652 spec : `TableSpec` 

653 Specification for the table. This will be used when creating the 

654 SQLAlchemy representation of the table, and it is used to 

655 check that the actual table in the database is consistent. 

656 

657 Returns 

658 ------- 

659 table : `sqlalchemy.schema.Table` or `None` 

660 SQLAlchemy representation of the table, or `None` if it does not 

661 exist. 

662 

663 Raises 

664 ------ 

665 DatabaseConflictError 

666 Raised if the table exists but ``spec`` is inconsistent with its 

667 definition. 

668 

669 Notes 

670 ----- 

671 This method can be called within transactions and never modifies the 

672 database. 

673 

674 Subclasses may override this method, but usually should not need to. 

675 """ 

676 assert self._metadata is not None, "Static tables must be declared before dynamic tables." 

677 name = self._mangleTableName(name) 

678 table = self._metadata.tables.get(name if self.namespace is None else f"{self.namespace}.{name}") 

679 if table is not None: 

680 if spec.fields.names != set(table.columns.keys()): 

681 raise DatabaseConflictError(f"Table '{name}' has already been defined differently; the new " 

682 f"specification has columns {list(spec.fields.names)}, while " 

683 f"the previous definition has {list(table.columns.keys())}.") 

684 else: 

685 inspector = sqlalchemy.engine.reflection.Inspector(self._connection) 

686 if name in inspector.get_table_names(schema=self.namespace): 

687 _checkExistingTableDefinition(name, spec, inspector.get_columns(name, schema=self.namespace)) 

688 table = self._convertTableSpec(name, spec, self._metadata) 

689 for foreignKeySpec in spec.foreignKeys: 

690 table.append_constraint(self._convertForeignKeySpec(name, foreignKeySpec, self._metadata)) 

691 return table 

692 return table 

693 

694 def sync(self, table: sqlalchemy.schema.Table, *, 

695 keys: Dict[str, Any], 

696 compared: Optional[Dict[str, Any]] = None, 

697 extra: Optional[Dict[str, Any]] = None, 

698 returning: Optional[Sequence[str]] = None, 

699 ) -> Tuple[Optional[Dict[str, Any], bool]]: 

700 """Insert into a table as necessary to ensure database contains 

701 values equivalent to the given ones. 

702 

703 Parameters 

704 ---------- 

705 table : `sqlalchemy.schema.Table` 

706 Table to be queried and possibly inserted into. 

707 keys : `dict` 

708 Column name-value pairs used to search for an existing row; must 

709 be a combination that can be used to select a single row if one 

710 exists. If such a row does not exist, these values are used in 

711 the insert. 

712 compared : `dict`, optional 

713 Column name-value pairs that are compared to those in any existing 

714 row. If such a row does not exist, these rows are used in the 

715 insert. 

716 extra : `dict`, optional 

717 Column name-value pairs that are ignored if a matching row exists, 

718 but used in an insert if one is necessary. 

719 returning : `~collections.abc.Sequence` of `str`, optional 

720 The names of columns whose values should be returned. 

721 

722 Returns 

723 ------- 

724 row : `dict`, optional 

725 The value of the fields indicated by ``returning``, or `None` if 

726 ``returning`` is `None`. 

727 inserted : `bool` 

728 If `True`, a new row was inserted. 

729 

730 Raises 

731 ------ 

732 DatabaseConflictError 

733 Raised if the values in ``compared`` do not match the values in the 

734 database. 

735 ReadOnlyDatabaseError 

736 Raised if `isWriteable` returns `False`, and no matching record 

737 already exists. 

738 

739 Notes 

740 ----- 

741 This method may not be called within transactions. It may be called on 

742 read-only databases if and only if the matching row does in fact 

743 already exist. 

744 """ 

745 

746 def check(): 

747 """Query for a row that matches the ``key`` argument, and compare 

748 to what was given by the caller. 

749 

750 Returns 

751 ------- 

752 n : `int` 

753 Number of matching rows. ``n != 1`` is always an error, but 

754 it's a different kind of error depending on where `check` is 

755 being called. 

756 bad : `list` of `str`, or `None` 

757 The subset of the keys of ``compared`` for which the existing 

758 values did not match the given one. Once again, ``not bad`` 

759 is always an error, but a different kind on context. `None` 

760 if ``n != 1`` 

761 result : `list` or `None` 

762 Results in the database that correspond to the columns given 

763 in ``returning``, or `None` if ``returning is None``. 

764 """ 

765 toSelect = set() 

766 if compared is not None: 

767 toSelect.update(compared.keys()) 

768 if returning is not None: 

769 toSelect.update(returning) 

770 if not toSelect: 

771 # Need to select some column, even if we just want to see 

772 # how many rows we get back. 

773 toSelect.add(next(iter(keys.keys()))) 

774 selectSql = sqlalchemy.sql.select( 

775 [table.columns[k].label(k) for k in toSelect] 

776 ).select_from(table).where( 

777 sqlalchemy.sql.and_(*[table.columns[k] == v for k, v in keys.items()]) 

778 ) 

779 fetched = list(self._connection.execute(selectSql).fetchall()) 

780 if len(fetched) != 1: 

781 return len(fetched), None, None 

782 existing = fetched[0] 

783 if compared is not None: 

784 

785 def safeNotEqual(a, b): 

786 if isinstance(a, astropy.time.Time): 

787 return not time_utils.times_equal(a, b) 

788 return a != b 

789 

790 inconsistencies = [f"{k}: {existing[k]!r} != {v!r}" 

791 for k, v in compared.items() 

792 if safeNotEqual(existing[k], v)] 

793 else: 

794 inconsistencies = [] 

795 if returning is not None: 

796 toReturn = [existing[k] for k in returning] 

797 else: 

798 toReturn = None 

799 return 1, inconsistencies, toReturn 

800 

801 if self.isWriteable(): 

802 # Database is writeable. Try an insert first, but allow it to fail 

803 # (in only specific ways). 

804 row = keys.copy() 

805 if compared is not None: 

806 row.update(compared) 

807 if extra is not None: 

808 row.update(extra) 

809 insertSql = table.insert().values(row) 

810 try: 

811 with self.transaction(interrupting=True): 

812 self._connection.execute(insertSql) 

813 # Need to perform check() for this branch inside the 

814 # transaction, so we roll back an insert that didn't do 

815 # what we expected. That limits the extent to which we 

816 # can reduce duplication between this block and the other 

817 # ones that perform similar logic. 

818 n, bad, result = check() 

819 if n < 1: 

820 raise RuntimeError("Insertion in sync did not seem to affect table. This is a bug.") 

821 elif n > 1: 

822 raise RuntimeError(f"Keys passed to sync {keys.keys()} do not comprise a " 

823 f"unique constraint for table {table.name}.") 

824 elif bad: 

825 raise RuntimeError( 

826 f"Conflict ({bad}) in sync after successful insert; this is " 

827 f"possible if the same table is being updated by a concurrent " 

828 f"process that isn't using sync, but it may also be a bug in " 

829 f"daf_butler." 

830 ) 

831 # No exceptions, so it looks like we inserted the requested row 

832 # successfully. 

833 inserted = True 

834 except sqlalchemy.exc.IntegrityError as err: 

835 # Most likely cause is that an equivalent row already exists, 

836 # but it could also be some other constraint. Query for the 

837 # row we think we matched to resolve that question. 

838 n, bad, result = check() 

839 if n < 1: 

840 # There was no matched row; insertion failed for some 

841 # completely different reason. Just re-raise the original 

842 # IntegrityError. 

843 raise 

844 elif n > 2: 

845 # There were multiple matched rows, which means we 

846 # conflicted *and* the arguments were bad to begin with. 

847 raise RuntimeError(f"Keys passed to sync {keys.keys()} do not comprise a " 

848 f"unique constraint for table {table.name}.") from err 

849 elif bad: 

850 # No logic bug, but data conflicted on the keys given. 

851 raise DatabaseConflictError(f"Conflict in sync for table " 

852 f"{table.name} on column(s) {bad}.") from err 

853 # The desired row is already present and consistent with what 

854 # we tried to insert. 

855 inserted = False 

856 else: 

857 assert not self._connection.in_transaction(), ( 

858 "Calling sync within a transaction block is an error even " 

859 "on a read-only database." 

860 ) 

861 # Database is not writeable; just see if the row exists. 

862 n, bad, result = check() 

863 if n < 1: 

864 raise ReadOnlyDatabaseError("sync needs to insert, but database is read-only.") 

865 elif n > 1: 

866 raise RuntimeError("Keys passed to sync do not comprise a unique constraint.") 

867 elif bad: 

868 raise DatabaseConflictError(f"Conflict in sync on column(s) {bad}.") 

869 inserted = False 

870 return {k: v for k, v in zip(returning, result)} if returning is not None else None, inserted 

871 

872 def insert(self, table: sqlalchemy.schema.Table, *rows: dict, returnIds: bool = False, 

873 ) -> Optional[List[int]]: 

874 """Insert one or more rows into a table, optionally returning 

875 autoincrement primary key values. 

876 

877 Parameters 

878 ---------- 

879 table : `sqlalchemy.schema.Table` 

880 Table rows should be inserted into. 

881 returnIds: `bool` 

882 If `True` (`False` is default), return the values of the table's 

883 autoincrement primary key field (which much exist). 

884 *rows 

885 Positional arguments are the rows to be inserted, as dictionaries 

886 mapping column name to value. The keys in all dictionaries must 

887 be the same. 

888 

889 Returns 

890 ------- 

891 ids : `None`, or `list` of `int` 

892 If ``returnIds`` is `True`, a `list` containing the inserted 

893 values for the table's autoincrement primary key. 

894 

895 Raises 

896 ------ 

897 ReadOnlyDatabaseError 

898 Raised if `isWriteable` returns `False` when this method is called. 

899 

900 Notes 

901 ----- 

902 The default implementation uses bulk insert syntax when ``returnIds`` 

903 is `False`, and a loop over single-row insert operations when it is 

904 `True`. 

905 

906 Derived classes should reimplement when they can provide a more 

907 efficient implementation (especially for the latter case). 

908 

909 May be used inside transaction contexts, so implementations may not 

910 perform operations that interrupt transactions. 

911 """ 

912 if not self.isWriteable(): 

913 raise ReadOnlyDatabaseError(f"Attempt to insert into read-only database '{self}'.") 

914 if not returnIds: 

915 self._connection.execute(table.insert(), *rows) 

916 else: 

917 sql = table.insert() 

918 return [self._connection.execute(sql, row).inserted_primary_key[0] for row in rows] 

919 

920 @abstractmethod 

921 def replace(self, table: sqlalchemy.schema.Table, *rows: dict): 

922 """Insert one or more rows into a table, replacing any existing rows 

923 for which insertion of a new row would violate the primary key 

924 constraint. 

925 

926 Parameters 

927 ---------- 

928 table : `sqlalchemy.schema.Table` 

929 Table rows should be inserted into. 

930 *rows 

931 Positional arguments are the rows to be inserted, as dictionaries 

932 mapping column name to value. The keys in all dictionaries must 

933 be the same. 

934 

935 Raises 

936 ------ 

937 ReadOnlyDatabaseError 

938 Raised if `isWriteable` returns `False` when this method is called. 

939 

940 Notes 

941 ----- 

942 May be used inside transaction contexts, so implementations may not 

943 perform operations that interrupt transactions. 

944 

945 Implementations should raise a `sqlalchemy.exc.IntegrityError` 

946 exception when a constraint other than the primary key would be 

947 violated. 

948 

949 Implementations are not required to support `replace` on tables 

950 with autoincrement keys. 

951 """ 

952 raise NotImplementedError() 

953 

954 def delete(self, table: sqlalchemy.schema.Table, columns: Iterable[str], *rows: dict) -> int: 

955 """Delete one or more rows from a table. 

956 

957 Parameters 

958 ---------- 

959 table : `sqlalchemy.schema.Table` 

960 Table that rows should be deleted from. 

961 columns: `~collections.abc.Iterable` of `str` 

962 The names of columns that will be used to constrain the rows to 

963 be deleted; these will be combined via ``AND`` to form the 

964 ``WHERE`` clause of the delete query. 

965 *rows 

966 Positional arguments are the keys of rows to be deleted, as 

967 dictionaries mapping column name to value. The keys in all 

968 dictionaries must exactly the names in ``columns``. 

969 

970 Returns 

971 ------- 

972 count : `int` 

973 Number of rows deleted. 

974 

975 Raises 

976 ------ 

977 ReadOnlyDatabaseError 

978 Raised if `isWriteable` returns `False` when this method is called. 

979 

980 Notes 

981 ----- 

982 May be used inside transaction contexts, so implementations may not 

983 perform operations that interrupt transactions. 

984 

985 The default implementation should be sufficient for most derived 

986 classes. 

987 """ 

988 if not self.isWriteable(): 

989 raise ReadOnlyDatabaseError(f"Attempt to delete from read-only database '{self}'.") 

990 sql = table.delete() 

991 whereTerms = [table.columns[name] == sqlalchemy.sql.bindparam(name) for name in columns] 

992 if whereTerms: 

993 sql = sql.where(sqlalchemy.sql.and_(*whereTerms)) 

994 return self._connection.execute(sql, *rows).rowcount 

995 

996 def update(self, table: sqlalchemy.schema.Table, where: Dict[str, str], *rows: dict) -> int: 

997 """Update one or more rows in a table. 

998 

999 Parameters 

1000 ---------- 

1001 table : `sqlalchemy.schema.Table` 

1002 Table containing the rows to be updated. 

1003 where : `dict` [`str`, `str`] 

1004 A mapping from the names of columns that will be used to search for 

1005 existing rows to the keys that will hold these values in the 

1006 ``rows`` dictionaries. Note that these may not be the same due to 

1007 SQLAlchemy limitations. 

1008 *rows 

1009 Positional arguments are the rows to be updated. The keys in all 

1010 dictionaries must be the same, and may correspond to either a 

1011 value in the ``where`` dictionary or the name of a column to be 

1012 updated. 

1013 

1014 Returns 

1015 ------- 

1016 count : `int` 

1017 Number of rows matched (regardless of whether the update actually 

1018 modified them). 

1019 

1020 Raises 

1021 ------ 

1022 ReadOnlyDatabaseError 

1023 Raised if `isWriteable` returns `False` when this method is called. 

1024 

1025 Notes 

1026 ----- 

1027 May be used inside transaction contexts, so implementations may not 

1028 perform operations that interrupt transactions. 

1029 

1030 The default implementation should be sufficient for most derived 

1031 classes. 

1032 """ 

1033 if not self.isWriteable(): 

1034 raise ReadOnlyDatabaseError(f"Attempt to update read-only database '{self}'.") 

1035 sql = table.update().where( 

1036 sqlalchemy.sql.and_(*[table.columns[k] == sqlalchemy.sql.bindparam(v) for k, v in where.items()]) 

1037 ) 

1038 return self._connection.execute(sql, *rows).rowcount 

1039 

1040 def query(self, sql: sqlalchemy.sql.FromClause, *args, **kwds) -> sqlalchemy.engine.ResultProxy: 

1041 """Run a SELECT query against the database. 

1042 

1043 Parameters 

1044 ---------- 

1045 sql : `sqlalchemy.sql.FromClause` 

1046 A SQLAlchemy representation of a ``SELECT`` query. 

1047 *args 

1048 Additional positional arguments are forwarded to 

1049 `sqlalchemy.engine.Connection.execute`. 

1050 **kwds 

1051 Additional keyword arguments are forwarded to 

1052 `sqlalchemy.engine.Connection.execute`. 

1053 

1054 Returns 

1055 ------- 

1056 result : `sqlalchemy.engine.ResultProxy` 

1057 Query results. 

1058 

1059 Notes 

1060 ----- 

1061 The default implementation should be sufficient for most derived 

1062 classes. 

1063 """ 

1064 # TODO: should we guard against non-SELECT queries here? 

1065 return self._connection.execute(sql, *args, **kwds) 

1066 

1067 origin: int 

1068 """An integer ID that should be used as the default for any datasets, 

1069 quanta, or other entities that use a (autoincrement, origin) compound 

1070 primary key (`int`). 

1071 """ 

1072 

1073 namespace: Optional[str] 

1074 """The schema or namespace this database instance is associated with 

1075 (`str` or `None`). 

1076 """