Coverage for python/lsst/daf/butler/registry/interfaces/_database.py : 14%

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
23__all__ = [
24 "Database",
25 "ReadOnlyDatabaseError",
26 "DatabaseConflictError",
27 "StaticTablesContext",
28]
30from abc import ABC, abstractmethod
31from contextlib import contextmanager
32from typing import (
33 Any,
34 Dict,
35 Iterable,
36 Iterator,
37 List,
38 Optional,
39 Sequence,
40 Set,
41 Tuple,
42)
43import warnings
45import astropy.time
46import sqlalchemy
48from ...core import ddl, time_utils
49from .._exceptions import ConflictingDefinitionError
52def _checkExistingTableDefinition(name: str, spec: ddl.TableSpec, inspection: List[Dict[str, Any]]) -> None:
53 """Test that the definition of a table in a `ddl.TableSpec` and from
54 database introspection are consistent.
56 Parameters
57 ----------
58 name : `str`
59 Name of the table (only used in error messages).
60 spec : `ddl.TableSpec`
61 Specification of the table.
62 inspection : `dict`
63 Dictionary returned by
64 `sqlalchemy.engine.reflection.Inspector.get_columns`.
66 Raises
67 ------
68 DatabaseConflictError
69 Raised if the definitions are inconsistent.
70 """
71 columnNames = [c["name"] for c in inspection]
72 if spec.fields.names != set(columnNames):
73 raise DatabaseConflictError(f"Table '{name}' exists but is defined differently in the database; "
74 f"specification has columns {list(spec.fields.names)}, while the "
75 f"table in the database has {columnNames}.")
78class ReadOnlyDatabaseError(RuntimeError):
79 """Exception raised when a write operation is called on a read-only
80 `Database`.
81 """
84class DatabaseConflictError(ConflictingDefinitionError):
85 """Exception raised when database content (row values or schema entities)
86 are inconsistent with what this client expects.
87 """
90class StaticTablesContext:
91 """Helper class used to declare the static schema for a registry layer
92 in a database.
94 An instance of this class is returned by `Database.declareStaticTables`,
95 which should be the only way it should be constructed.
96 """
98 def __init__(self, db: Database):
99 self._db = db
100 self._foreignKeys: List[Tuple[sqlalchemy.schema.Table, sqlalchemy.schema.ForeignKeyConstraint]] = []
101 self._inspector = sqlalchemy.engine.reflection.Inspector(self._db._connection)
102 self._tableNames = frozenset(self._inspector.get_table_names(schema=self._db.namespace))
104 def addTable(self, name: str, spec: ddl.TableSpec) -> sqlalchemy.schema.Table:
105 """Add a new table to the schema, returning its sqlalchemy
106 representation.
108 The new table may not actually be created until the end of the
109 context created by `Database.declareStaticTables`, allowing tables
110 to be declared in any order even in the presence of foreign key
111 relationships.
112 """
113 name = self._db._mangleTableName(name)
114 if name in self._tableNames:
115 _checkExistingTableDefinition(name, spec, self._inspector.get_columns(name,
116 schema=self._db.namespace))
117 table = self._db._convertTableSpec(name, spec, self._db._metadata)
118 for foreignKeySpec in spec.foreignKeys:
119 self._foreignKeys.append(
120 (table, self._db._convertForeignKeySpec(name, foreignKeySpec, self._db._metadata))
121 )
122 return table
124 def addTableTuple(self, specs: Tuple[ddl.TableSpec, ...]) -> Tuple[sqlalchemy.schema.Table, ...]:
125 """Add a named tuple of tables to the schema, returning their
126 SQLAlchemy representations in a named tuple of the same type.
128 The new tables may not actually be created until the end of the
129 context created by `Database.declareStaticTables`, allowing tables
130 to be declared in any order even in the presence of foreign key
131 relationships.
133 Notes
134 -----
135 ``specs`` *must* be an instance of a type created by
136 `collections.namedtuple`, not just regular tuple, and the returned
137 object is guaranteed to be the same. Because `~collections.namedtuple`
138 is just a factory for `type` objects, not an actual type itself,
139 we cannot represent this with type annotations.
140 """
141 return specs._make(self.addTable(name, spec) # type: ignore
142 for name, spec in zip(specs._fields, specs)) # type: ignore
145class Database(ABC):
146 """An abstract interface that represents a particular database engine's
147 representation of a single schema/namespace/database.
149 Parameters
150 ----------
151 origin : `int`
152 An integer ID that should be used as the default for any datasets,
153 quanta, or other entities that use a (autoincrement, origin) compound
154 primary key.
155 connection : `sqlalchemy.engine.Connection`
156 The SQLAlchemy connection this `Database` wraps.
157 namespace : `str`, optional
158 Name of the schema or namespace this instance is associated with.
159 This is passed as the ``schema`` argument when constructing a
160 `sqlalchemy.schema.MetaData` instance. We use ``namespace`` instead to
161 avoid confusion between "schema means namespace" and "schema means
162 table definitions".
164 Notes
165 -----
166 `Database` requires all write operations to go through its special named
167 methods. Our write patterns are sufficiently simple that we don't really
168 need the full flexibility of SQL insert/update/delete syntax, and we need
169 non-standard (but common) functionality in these operations sufficiently
170 often that it seems worthwhile to provide our own generic API.
172 In contrast, `Database.query` allows arbitrary ``SELECT`` queries (via
173 their SQLAlchemy representation) to be run, as we expect these to require
174 significantly more sophistication while still being limited to standard
175 SQL.
177 `Database` itself has several underscore-prefixed attributes:
179 - ``_connection``: SQLAlchemy object representing the connection.
180 - ``_metadata``: the `sqlalchemy.schema.MetaData` object representing
181 the tables and other schema entities.
183 These are considered protected (derived classes may access them, but other
184 code should not), and read-only, aside from executing SQL via
185 ``_connection``.
186 """
188 def __init__(self, *, origin: int, connection: sqlalchemy.engine.Connection,
189 namespace: Optional[str] = None):
190 self.origin = origin
191 self.namespace = namespace
192 self._connection = connection
193 self._metadata: Optional[sqlalchemy.schema.MetaData] = None
195 @classmethod
196 def makeDefaultUri(cls, root: str) -> Optional[str]:
197 """Create a default connection URI appropriate for the given root
198 directory, or `None` if there can be no such default.
199 """
200 return None
202 @classmethod
203 def fromUri(cls, uri: str, *, origin: int, namespace: Optional[str] = None,
204 writeable: bool = True) -> Database:
205 """Construct a database from a SQLAlchemy URI.
207 Parameters
208 ----------
209 uri : `str`
210 A SQLAlchemy URI connection string.
211 origin : `int`
212 An integer ID that should be used as the default for any datasets,
213 quanta, or other entities that use a (autoincrement, origin)
214 compound primary key.
215 namespace : `str`, optional
216 A database namespace (i.e. schema) the new instance should be
217 associated with. If `None` (default), the namespace (if any) is
218 inferred from the URI.
219 writeable : `bool`, optional
220 If `True`, allow write operations on the database, including
221 ``CREATE TABLE``.
223 Returns
224 -------
225 db : `Database`
226 A new `Database` instance.
227 """
228 return cls.fromConnection(cls.connect(uri, writeable=writeable),
229 origin=origin,
230 namespace=namespace,
231 writeable=writeable)
233 @classmethod
234 @abstractmethod
235 def connect(cls, uri: str, *, writeable: bool = True) -> sqlalchemy.engine.Connection:
236 """Create a `sqlalchemy.engine.Connection` from a SQLAlchemy URI.
238 Parameters
239 ----------
240 uri : `str`
241 A SQLAlchemy URI connection string.
242 origin : `int`
243 An integer ID that should be used as the default for any datasets,
244 quanta, or other entities that use a (autoincrement, origin)
245 compound primary key.
246 writeable : `bool`, optional
247 If `True`, allow write operations on the database, including
248 ``CREATE TABLE``.
250 Returns
251 -------
252 connection : `sqlalchemy.engine.Connection`
253 A database connection.
255 Notes
256 -----
257 Subclasses that support other ways to connect to a database are
258 encouraged to add optional arguments to their implementation of this
259 method, as long as they maintain compatibility with the base class
260 call signature.
261 """
262 raise NotImplementedError()
264 @classmethod
265 @abstractmethod
266 def fromConnection(cls, connection: sqlalchemy.engine.Connection, *, origin: int,
267 namespace: Optional[str] = None, writeable: bool = True) -> Database:
268 """Create a new `Database` from an existing
269 `sqlalchemy.engine.Connection`.
271 Parameters
272 ----------
273 connection : `sqllachemy.engine.Connection`
274 The connection for the the database. May be shared between
275 `Database` instances.
276 origin : `int`
277 An integer ID that should be used as the default for any datasets,
278 quanta, or other entities that use a (autoincrement, origin)
279 compound primary key.
280 namespace : `str`, optional
281 A different database namespace (i.e. schema) the new instance
282 should be associated with. If `None` (default), the namespace
283 (if any) is inferred from the connection.
284 writeable : `bool`, optional
285 If `True`, allow write operations on the database, including
286 ``CREATE TABLE``.
288 Returns
289 -------
290 db : `Database`
291 A new `Database` instance.
293 Notes
294 -----
295 This method allows different `Database` instances to share the same
296 connection, which is desirable when they represent different namespaces
297 can be queried together. This also ties their transaction state,
298 however; starting a transaction in any database automatically starts
299 on in all other databases.
300 """
301 raise NotImplementedError()
303 @contextmanager
304 def transaction(self, *, interrupting: bool = False) -> Iterator:
305 """Return a context manager that represents a transaction.
307 Parameters
308 ----------
309 interrupting : `bool`
310 If `True`, this transaction block needs to be able to interrupt
311 any existing one in order to yield correct behavior.
312 """
313 assert not (interrupting and self._connection.in_transaction()), (
314 "Logic error in transaction nesting: an operation that would "
315 "interrupt the active transaction context has been requested."
316 )
317 if self._connection.in_transaction():
318 trans = self._connection.begin_nested()
319 else:
320 # Use a regular (non-savepoint) transaction only for the outermost
321 # context.
322 trans = self._connection.begin()
323 try:
324 yield
325 trans.commit()
326 except BaseException:
327 trans.rollback()
328 raise
330 @contextmanager
331 def declareStaticTables(self, *, create: bool) -> Iterator[StaticTablesContext]:
332 """Return a context manager in which the database's static DDL schema
333 can be declared.
335 Parameters
336 ----------
337 create : `bool`
338 If `True`, attempt to create all tables at the end of the context.
339 If `False`, they will be assumed to already exist.
341 Returns
342 -------
343 schema : `StaticTablesContext`
344 A helper object that is used to add new tables.
346 Raises
347 ------
348 ReadOnlyDatabaseError
349 Raised if ``create`` is `True`, `Database.isWriteable` is `False`,
350 and one or more declared tables do not already exist.
352 Examples
353 --------
354 Given a `Database` instance ``db``::
356 with db.declareStaticTables(create=True) as schema:
357 schema.addTable("table1", TableSpec(...))
358 schema.addTable("table2", TableSpec(...))
360 Notes
361 -----
362 A database's static DDL schema must be declared before any dynamic
363 tables are managed via calls to `ensureTableExists` or
364 `getExistingTable`. The order in which static schema tables are added
365 inside the context block is unimportant; they will automatically be
366 sorted and added in an order consistent with their foreign key
367 relationships.
368 """
369 if create and not self.isWriteable():
370 raise ReadOnlyDatabaseError(f"Cannot create tables in read-only database {self}.")
371 self._metadata = sqlalchemy.MetaData(schema=self.namespace)
372 try:
373 context = StaticTablesContext(self)
374 yield context
375 for table, foreignKey in context._foreignKeys:
376 table.append_constraint(foreignKey)
377 if create:
378 if self.namespace is not None:
379 if self.namespace not in context._inspector.get_schema_names():
380 self._connection.execute(sqlalchemy.schema.CreateSchema(self.namespace))
381 # In our tables we have columns that make use of sqlalchemy
382 # Sequence objects. There is currently a bug in sqlalchmey that
383 # causes a deprecation warning to be thrown on a property of
384 # the Sequence object when the repr for the sequence is
385 # created. Here a filter is used to catch these deprecation
386 # warnings when tables are created.
387 with warnings.catch_warnings():
388 warnings.simplefilter("ignore", category=sqlalchemy.exc.SADeprecationWarning)
389 self._metadata.create_all(self._connection)
390 except BaseException:
391 self._metadata.drop_all(self._connection)
392 self._metadata = None
393 raise
395 @abstractmethod
396 def isWriteable(self) -> bool:
397 """Return `True` if this database can be modified by this client.
398 """
399 raise NotImplementedError()
401 @abstractmethod
402 def __str__(self) -> str:
403 """Return a human-readable identifier for this `Database`, including
404 any namespace or schema that identifies its names within a `Registry`.
405 """
406 raise NotImplementedError()
408 def shrinkDatabaseEntityName(self, original: str) -> str:
409 """Return a version of the given name that fits within this database
410 engine's length limits for table, constraint, indexes, and sequence
411 names.
413 Implementations should not assume that simple truncation is safe,
414 because multiple long names often begin with the same prefix.
416 The default implementation simply returns the given name.
418 Parameters
419 ----------
420 original : `str`
421 The original name.
423 Returns
424 -------
425 shrunk : `str`
426 The new, possibly shortened name.
427 """
428 return original
430 def expandDatabaseEntityName(self, shrunk: str) -> str:
431 """Retrieve the original name for a database entity that was too long
432 to fit within the database engine's limits.
434 Parameters
435 ----------
436 original : `str`
437 The original name.
439 Returns
440 -------
441 shrunk : `str`
442 The new, possibly shortened name.
443 """
444 return shrunk
446 def _mangleTableName(self, name: str) -> str:
447 """Map a logical, user-visible table name to the true table name used
448 in the database.
450 The default implementation returns the given name unchanged.
452 Parameters
453 ----------
454 name : `str`
455 Input table name. Should not include a namespace (i.e. schema)
456 prefix.
458 Returns
459 -------
460 mangled : `str`
461 Mangled version of the table name (still with no namespace prefix).
463 Notes
464 -----
465 Reimplementations of this method must be idempotent - mangling an
466 already-mangled name must have no effect.
467 """
468 return name
470 def _makeColumnConstraints(self, table: str, spec: ddl.FieldSpec) -> List[sqlalchemy.CheckConstraint]:
471 """Create constraints based on this spec.
473 Parameters
474 ----------
475 table : `str`
476 Name of the table this column is being added to.
477 spec : `FieldSpec`
478 Specification for the field to be added.
480 Returns
481 -------
482 constraint : `list` of `sqlalchemy.CheckConstraint`
483 Constraint added for this column.
484 """
485 # By default we return no additional constraints
486 return []
488 def _convertFieldSpec(self, table: str, spec: ddl.FieldSpec, metadata: sqlalchemy.MetaData,
489 **kwds: Any) -> sqlalchemy.schema.Column:
490 """Convert a `FieldSpec` to a `sqlalchemy.schema.Column`.
492 Parameters
493 ----------
494 table : `str`
495 Name of the table this column is being added to.
496 spec : `FieldSpec`
497 Specification for the field to be added.
498 metadata : `sqlalchemy.MetaData`
499 SQLAlchemy representation of the DDL schema this field's table is
500 being added to.
501 **kwds
502 Additional keyword arguments to forward to the
503 `sqlalchemy.schema.Column` constructor. This is provided to make
504 it easier for derived classes to delegate to ``super()`` while
505 making only minor changes.
507 Returns
508 -------
509 column : `sqlalchemy.schema.Column`
510 SQLAlchemy representation of the field.
511 """
512 args = [spec.name, spec.getSizedColumnType()]
513 if spec.autoincrement:
514 # Generate a sequence to use for auto incrementing for databases
515 # that do not support it natively. This will be ignored by
516 # sqlalchemy for databases that do support it.
517 args.append(sqlalchemy.Sequence(self.shrinkDatabaseEntityName(f"{table}_seq_{spec.name}"),
518 metadata=metadata))
519 assert spec.doc is None or isinstance(spec.doc, str), f"Bad doc for {table}.{spec.name}."
520 return sqlalchemy.schema.Column(*args, nullable=spec.nullable, primary_key=spec.primaryKey,
521 comment=spec.doc, **kwds)
523 def _convertForeignKeySpec(self, table: str, spec: ddl.ForeignKeySpec, metadata: sqlalchemy.MetaData,
524 **kwds: Any) -> sqlalchemy.schema.ForeignKeyConstraint:
525 """Convert a `ForeignKeySpec` to a
526 `sqlalchemy.schema.ForeignKeyConstraint`.
528 Parameters
529 ----------
530 table : `str`
531 Name of the table this foreign key is being added to.
532 spec : `ForeignKeySpec`
533 Specification for the foreign key to be added.
534 metadata : `sqlalchemy.MetaData`
535 SQLAlchemy representation of the DDL schema this constraint is
536 being added to.
537 **kwds
538 Additional keyword arguments to forward to the
539 `sqlalchemy.schema.ForeignKeyConstraint` constructor. This is
540 provided to make it easier for derived classes to delegate to
541 ``super()`` while making only minor changes.
543 Returns
544 -------
545 constraint : `sqlalchemy.schema.ForeignKeyConstraint`
546 SQLAlchemy representation of the constraint.
547 """
548 name = self.shrinkDatabaseEntityName(
549 "_".join(["fkey", table, self._mangleTableName(spec.table)]
550 + list(spec.target) + list(spec.source))
551 )
552 return sqlalchemy.schema.ForeignKeyConstraint(
553 spec.source,
554 [f"{self._mangleTableName(spec.table)}.{col}" for col in spec.target],
555 name=name,
556 ondelete=spec.onDelete
557 )
559 def _convertTableSpec(self, name: str, spec: ddl.TableSpec, metadata: sqlalchemy.MetaData,
560 **kwds: Any) -> sqlalchemy.schema.Table:
561 """Convert a `TableSpec` to a `sqlalchemy.schema.Table`.
563 Parameters
564 ----------
565 spec : `TableSpec`
566 Specification for the foreign key to be added.
567 metadata : `sqlalchemy.MetaData`
568 SQLAlchemy representation of the DDL schema this table is being
569 added to.
570 **kwds
571 Additional keyword arguments to forward to the
572 `sqlalchemy.schema.Table` constructor. This is provided to make it
573 easier for derived classes to delegate to ``super()`` while making
574 only minor changes.
576 Returns
577 -------
578 table : `sqlalchemy.schema.Table`
579 SQLAlchemy representation of the table.
581 Notes
582 -----
583 This method does not handle ``spec.foreignKeys`` at all, in order to
584 avoid circular dependencies. These are added by higher-level logic in
585 `ensureTableExists`, `getExistingTable`, and `declareStaticTables`.
586 """
587 name = self._mangleTableName(name)
588 args = [self._convertFieldSpec(name, fieldSpec, metadata) for fieldSpec in spec.fields]
590 # Add any column constraints
591 for fieldSpec in spec.fields:
592 args.extend(self._makeColumnConstraints(name, fieldSpec))
594 # Track indexes added for primary key and unique constraints, to make
595 # sure we don't add duplicate explicit or foreign key indexes for
596 # those.
597 allIndexes = {tuple(fieldSpec.name for fieldSpec in spec.fields if fieldSpec.primaryKey)}
598 args.extend(
599 sqlalchemy.schema.UniqueConstraint(
600 *columns,
601 name=self.shrinkDatabaseEntityName("_".join([name, "unq"] + list(columns)))
602 )
603 for columns in spec.unique
604 )
605 allIndexes.update(spec.unique)
606 args.extend(
607 sqlalchemy.schema.Index(
608 self.shrinkDatabaseEntityName("_".join([name, "idx"] + list(columns))),
609 *columns,
610 unique=(columns in spec.unique)
611 )
612 for columns in spec.indexes if columns not in allIndexes
613 )
614 allIndexes.update(spec.indexes)
615 args.extend(
616 sqlalchemy.schema.Index(
617 self.shrinkDatabaseEntityName("_".join((name, "fkidx") + fk.source)),
618 *fk.source,
619 )
620 for fk in spec.foreignKeys if fk.addIndex and fk.source not in allIndexes
621 )
622 assert spec.doc is None or isinstance(spec.doc, str), f"Bad doc for {name}."
623 return sqlalchemy.schema.Table(name, metadata, *args, comment=spec.doc, info=spec, **kwds)
625 def ensureTableExists(self, name: str, spec: ddl.TableSpec) -> sqlalchemy.sql.FromClause:
626 """Ensure that a table with the given name and specification exists,
627 creating it if necessary.
629 Parameters
630 ----------
631 name : `str`
632 Name of the table (not including namespace qualifiers).
633 spec : `TableSpec`
634 Specification for the table. This will be used when creating the
635 table, and *may* be used when obtaining an existing table to check
636 for consistency, but no such check is guaranteed.
638 Returns
639 -------
640 table : `sqlalchemy.schema.Table`
641 SQLAlchemy representation of the table.
643 Raises
644 ------
645 ReadOnlyDatabaseError
646 Raised if `isWriteable` returns `False`, and the table does not
647 already exist.
648 DatabaseConflictError
649 Raised if the table exists but ``spec`` is inconsistent with its
650 definition.
652 Notes
653 -----
654 This method may not be called within transactions. It may be called on
655 read-only databases if and only if the table does in fact already
656 exist.
658 Subclasses may override this method, but usually should not need to.
659 """
660 assert not self._connection.in_transaction(), "Table creation interrupts transactions."
661 assert self._metadata is not None, "Static tables must be declared before dynamic tables."
662 table = self.getExistingTable(name, spec)
663 if table is not None:
664 return table
665 if not self.isWriteable():
666 raise ReadOnlyDatabaseError(
667 f"Table {name} does not exist, and cannot be created "
668 f"because database {self} is read-only."
669 )
670 table = self._convertTableSpec(name, spec, self._metadata)
671 for foreignKeySpec in spec.foreignKeys:
672 table.append_constraint(self._convertForeignKeySpec(name, foreignKeySpec, self._metadata))
673 table.create(self._connection)
674 return table
676 def getExistingTable(self, name: str, spec: ddl.TableSpec) -> Optional[sqlalchemy.schema.Table]:
677 """Obtain an existing table with the given name and specification.
679 Parameters
680 ----------
681 name : `str`
682 Name of the table (not including namespace qualifiers).
683 spec : `TableSpec`
684 Specification for the table. This will be used when creating the
685 SQLAlchemy representation of the table, and it is used to
686 check that the actual table in the database is consistent.
688 Returns
689 -------
690 table : `sqlalchemy.schema.Table` or `None`
691 SQLAlchemy representation of the table, or `None` if it does not
692 exist.
694 Raises
695 ------
696 DatabaseConflictError
697 Raised if the table exists but ``spec`` is inconsistent with its
698 definition.
700 Notes
701 -----
702 This method can be called within transactions and never modifies the
703 database.
705 Subclasses may override this method, but usually should not need to.
706 """
707 assert self._metadata is not None, "Static tables must be declared before dynamic tables."
708 name = self._mangleTableName(name)
709 table = self._metadata.tables.get(name if self.namespace is None else f"{self.namespace}.{name}")
710 if table is not None:
711 if spec.fields.names != set(table.columns.keys()):
712 raise DatabaseConflictError(f"Table '{name}' has already been defined differently; the new "
713 f"specification has columns {list(spec.fields.names)}, while "
714 f"the previous definition has {list(table.columns.keys())}.")
715 else:
716 inspector = sqlalchemy.engine.reflection.Inspector(self._connection)
717 if name in inspector.get_table_names(schema=self.namespace):
718 _checkExistingTableDefinition(name, spec, inspector.get_columns(name, schema=self.namespace))
719 table = self._convertTableSpec(name, spec, self._metadata)
720 for foreignKeySpec in spec.foreignKeys:
721 table.append_constraint(self._convertForeignKeySpec(name, foreignKeySpec, self._metadata))
722 return table
723 return table
725 def sync(self, table: sqlalchemy.schema.Table, *,
726 keys: Dict[str, Any],
727 compared: Optional[Dict[str, Any]] = None,
728 extra: Optional[Dict[str, Any]] = None,
729 returning: Optional[Sequence[str]] = None,
730 ) -> Tuple[Optional[Dict[str, Any]], bool]:
731 """Insert into a table as necessary to ensure database contains
732 values equivalent to the given ones.
734 Parameters
735 ----------
736 table : `sqlalchemy.schema.Table`
737 Table to be queried and possibly inserted into.
738 keys : `dict`
739 Column name-value pairs used to search for an existing row; must
740 be a combination that can be used to select a single row if one
741 exists. If such a row does not exist, these values are used in
742 the insert.
743 compared : `dict`, optional
744 Column name-value pairs that are compared to those in any existing
745 row. If such a row does not exist, these rows are used in the
746 insert.
747 extra : `dict`, optional
748 Column name-value pairs that are ignored if a matching row exists,
749 but used in an insert if one is necessary.
750 returning : `~collections.abc.Sequence` of `str`, optional
751 The names of columns whose values should be returned.
753 Returns
754 -------
755 row : `dict`, optional
756 The value of the fields indicated by ``returning``, or `None` if
757 ``returning`` is `None`.
758 inserted : `bool`
759 If `True`, a new row was inserted.
761 Raises
762 ------
763 DatabaseConflictError
764 Raised if the values in ``compared`` do not match the values in the
765 database.
766 ReadOnlyDatabaseError
767 Raised if `isWriteable` returns `False`, and no matching record
768 already exists.
770 Notes
771 -----
772 This method may not be called within transactions. It may be called on
773 read-only databases if and only if the matching row does in fact
774 already exist.
775 """
777 def check() -> Tuple[int, Optional[List[str]], Optional[List]]:
778 """Query for a row that matches the ``key`` argument, and compare
779 to what was given by the caller.
781 Returns
782 -------
783 n : `int`
784 Number of matching rows. ``n != 1`` is always an error, but
785 it's a different kind of error depending on where `check` is
786 being called.
787 bad : `list` of `str`, or `None`
788 The subset of the keys of ``compared`` for which the existing
789 values did not match the given one. Once again, ``not bad``
790 is always an error, but a different kind on context. `None`
791 if ``n != 1``
792 result : `list` or `None`
793 Results in the database that correspond to the columns given
794 in ``returning``, or `None` if ``returning is None``.
795 """
796 toSelect: Set[str] = set()
797 if compared is not None:
798 toSelect.update(compared.keys())
799 if returning is not None:
800 toSelect.update(returning)
801 if not toSelect:
802 # Need to select some column, even if we just want to see
803 # how many rows we get back.
804 toSelect.add(next(iter(keys.keys())))
805 selectSql = sqlalchemy.sql.select(
806 [table.columns[k].label(k) for k in toSelect]
807 ).select_from(table).where(
808 sqlalchemy.sql.and_(*[table.columns[k] == v for k, v in keys.items()])
809 )
810 fetched = list(self._connection.execute(selectSql).fetchall())
811 if len(fetched) != 1:
812 return len(fetched), None, None
813 existing = fetched[0]
814 if compared is not None:
816 def safeNotEqual(a: Any, b: Any) -> bool:
817 if isinstance(a, astropy.time.Time):
818 return not time_utils.times_equal(a, b)
819 return a != b
821 inconsistencies = [f"{k}: {existing[k]!r} != {v!r}"
822 for k, v in compared.items()
823 if safeNotEqual(existing[k], v)]
824 else:
825 inconsistencies = []
826 if returning is not None:
827 toReturn: Optional[list] = [existing[k] for k in returning]
828 else:
829 toReturn = None
830 return 1, inconsistencies, toReturn
832 if self.isWriteable():
833 # Database is writeable. Try an insert first, but allow it to fail
834 # (in only specific ways).
835 row = keys.copy()
836 if compared is not None:
837 row.update(compared)
838 if extra is not None:
839 row.update(extra)
840 insertSql = table.insert().values(row)
841 try:
842 with self.transaction(interrupting=True):
843 self._connection.execute(insertSql)
844 # Need to perform check() for this branch inside the
845 # transaction, so we roll back an insert that didn't do
846 # what we expected. That limits the extent to which we
847 # can reduce duplication between this block and the other
848 # ones that perform similar logic.
849 n, bad, result = check()
850 if n < 1:
851 raise RuntimeError("Insertion in sync did not seem to affect table. This is a bug.")
852 elif n > 1:
853 raise RuntimeError(f"Keys passed to sync {keys.keys()} do not comprise a "
854 f"unique constraint for table {table.name}.")
855 elif bad:
856 raise RuntimeError(
857 f"Conflict ({bad}) in sync after successful insert; this is "
858 f"possible if the same table is being updated by a concurrent "
859 f"process that isn't using sync, but it may also be a bug in "
860 f"daf_butler."
861 )
862 # No exceptions, so it looks like we inserted the requested row
863 # successfully.
864 inserted = True
865 except sqlalchemy.exc.IntegrityError as err:
866 # Most likely cause is that an equivalent row already exists,
867 # but it could also be some other constraint. Query for the
868 # row we think we matched to resolve that question.
869 n, bad, result = check()
870 if n < 1:
871 # There was no matched row; insertion failed for some
872 # completely different reason. Just re-raise the original
873 # IntegrityError.
874 raise
875 elif n > 2:
876 # There were multiple matched rows, which means we
877 # conflicted *and* the arguments were bad to begin with.
878 raise RuntimeError(f"Keys passed to sync {keys.keys()} do not comprise a "
879 f"unique constraint for table {table.name}.") from err
880 elif bad:
881 # No logic bug, but data conflicted on the keys given.
882 raise DatabaseConflictError(f"Conflict in sync for table "
883 f"{table.name} on column(s) {bad}.") from err
884 # The desired row is already present and consistent with what
885 # we tried to insert.
886 inserted = False
887 else:
888 assert not self._connection.in_transaction(), (
889 "Calling sync within a transaction block is an error even "
890 "on a read-only database."
891 )
892 # Database is not writeable; just see if the row exists.
893 n, bad, result = check()
894 if n < 1:
895 raise ReadOnlyDatabaseError("sync needs to insert, but database is read-only.")
896 elif n > 1:
897 raise RuntimeError("Keys passed to sync do not comprise a unique constraint.")
898 elif bad:
899 raise DatabaseConflictError(f"Conflict in sync on column(s) {bad}.")
900 inserted = False
901 if returning is None:
902 return None, inserted
903 else:
904 assert result is not None
905 return {k: v for k, v in zip(returning, result)}, inserted
907 def insert(self, table: sqlalchemy.schema.Table, *rows: dict, returnIds: bool = False,
908 ) -> Optional[List[int]]:
909 """Insert one or more rows into a table, optionally returning
910 autoincrement primary key values.
912 Parameters
913 ----------
914 table : `sqlalchemy.schema.Table`
915 Table rows should be inserted into.
916 returnIds: `bool`
917 If `True` (`False` is default), return the values of the table's
918 autoincrement primary key field (which much exist).
919 *rows
920 Positional arguments are the rows to be inserted, as dictionaries
921 mapping column name to value. The keys in all dictionaries must
922 be the same.
924 Returns
925 -------
926 ids : `None`, or `list` of `int`
927 If ``returnIds`` is `True`, a `list` containing the inserted
928 values for the table's autoincrement primary key.
930 Raises
931 ------
932 ReadOnlyDatabaseError
933 Raised if `isWriteable` returns `False` when this method is called.
935 Notes
936 -----
937 The default implementation uses bulk insert syntax when ``returnIds``
938 is `False`, and a loop over single-row insert operations when it is
939 `True`.
941 Derived classes should reimplement when they can provide a more
942 efficient implementation (especially for the latter case).
944 May be used inside transaction contexts, so implementations may not
945 perform operations that interrupt transactions.
946 """
947 if not self.isWriteable():
948 raise ReadOnlyDatabaseError(f"Attempt to insert into read-only database '{self}'.")
949 if not returnIds:
950 self._connection.execute(table.insert(), *rows)
951 return None
952 else:
953 sql = table.insert()
954 return [self._connection.execute(sql, row).inserted_primary_key[0] for row in rows]
956 @abstractmethod
957 def replace(self, table: sqlalchemy.schema.Table, *rows: dict) -> None:
958 """Insert one or more rows into a table, replacing any existing rows
959 for which insertion of a new row would violate the primary key
960 constraint.
962 Parameters
963 ----------
964 table : `sqlalchemy.schema.Table`
965 Table rows should be inserted into.
966 *rows
967 Positional arguments are the rows to be inserted, as dictionaries
968 mapping column name to value. The keys in all dictionaries must
969 be the same.
971 Raises
972 ------
973 ReadOnlyDatabaseError
974 Raised if `isWriteable` returns `False` when this method is called.
976 Notes
977 -----
978 May be used inside transaction contexts, so implementations may not
979 perform operations that interrupt transactions.
981 Implementations should raise a `sqlalchemy.exc.IntegrityError`
982 exception when a constraint other than the primary key would be
983 violated.
985 Implementations are not required to support `replace` on tables
986 with autoincrement keys.
987 """
988 raise NotImplementedError()
990 def delete(self, table: sqlalchemy.schema.Table, columns: Iterable[str], *rows: dict) -> int:
991 """Delete one or more rows from a table.
993 Parameters
994 ----------
995 table : `sqlalchemy.schema.Table`
996 Table that rows should be deleted from.
997 columns: `~collections.abc.Iterable` of `str`
998 The names of columns that will be used to constrain the rows to
999 be deleted; these will be combined via ``AND`` to form the
1000 ``WHERE`` clause of the delete query.
1001 *rows
1002 Positional arguments are the keys of rows to be deleted, as
1003 dictionaries mapping column name to value. The keys in all
1004 dictionaries must exactly the names in ``columns``.
1006 Returns
1007 -------
1008 count : `int`
1009 Number of rows deleted.
1011 Raises
1012 ------
1013 ReadOnlyDatabaseError
1014 Raised if `isWriteable` returns `False` when this method is called.
1016 Notes
1017 -----
1018 May be used inside transaction contexts, so implementations may not
1019 perform operations that interrupt transactions.
1021 The default implementation should be sufficient for most derived
1022 classes.
1023 """
1024 if not self.isWriteable():
1025 raise ReadOnlyDatabaseError(f"Attempt to delete from read-only database '{self}'.")
1026 sql = table.delete()
1027 whereTerms = [table.columns[name] == sqlalchemy.sql.bindparam(name) for name in columns]
1028 if whereTerms:
1029 sql = sql.where(sqlalchemy.sql.and_(*whereTerms))
1030 return self._connection.execute(sql, *rows).rowcount
1032 def update(self, table: sqlalchemy.schema.Table, where: Dict[str, str], *rows: dict) -> int:
1033 """Update one or more rows in a table.
1035 Parameters
1036 ----------
1037 table : `sqlalchemy.schema.Table`
1038 Table containing the rows to be updated.
1039 where : `dict` [`str`, `str`]
1040 A mapping from the names of columns that will be used to search for
1041 existing rows to the keys that will hold these values in the
1042 ``rows`` dictionaries. Note that these may not be the same due to
1043 SQLAlchemy limitations.
1044 *rows
1045 Positional arguments are the rows to be updated. The keys in all
1046 dictionaries must be the same, and may correspond to either a
1047 value in the ``where`` dictionary or the name of a column to be
1048 updated.
1050 Returns
1051 -------
1052 count : `int`
1053 Number of rows matched (regardless of whether the update actually
1054 modified them).
1056 Raises
1057 ------
1058 ReadOnlyDatabaseError
1059 Raised if `isWriteable` returns `False` when this method is called.
1061 Notes
1062 -----
1063 May be used inside transaction contexts, so implementations may not
1064 perform operations that interrupt transactions.
1066 The default implementation should be sufficient for most derived
1067 classes.
1068 """
1069 if not self.isWriteable():
1070 raise ReadOnlyDatabaseError(f"Attempt to update read-only database '{self}'.")
1071 sql = table.update().where(
1072 sqlalchemy.sql.and_(*[table.columns[k] == sqlalchemy.sql.bindparam(v) for k, v in where.items()])
1073 )
1074 return self._connection.execute(sql, *rows).rowcount
1076 def query(self, sql: sqlalchemy.sql.FromClause,
1077 *args: Any, **kwds: Any) -> sqlalchemy.engine.ResultProxy:
1078 """Run a SELECT query against the database.
1080 Parameters
1081 ----------
1082 sql : `sqlalchemy.sql.FromClause`
1083 A SQLAlchemy representation of a ``SELECT`` query.
1084 *args
1085 Additional positional arguments are forwarded to
1086 `sqlalchemy.engine.Connection.execute`.
1087 **kwds
1088 Additional keyword arguments are forwarded to
1089 `sqlalchemy.engine.Connection.execute`.
1091 Returns
1092 -------
1093 result : `sqlalchemy.engine.ResultProxy`
1094 Query results.
1096 Notes
1097 -----
1098 The default implementation should be sufficient for most derived
1099 classes.
1100 """
1101 # TODO: should we guard against non-SELECT queries here?
1102 return self._connection.execute(sql, *args, **kwds)
1104 origin: int
1105 """An integer ID that should be used as the default for any datasets,
1106 quanta, or other entities that use a (autoincrement, origin) compound
1107 primary key (`int`).
1108 """
1110 namespace: Optional[str]
1111 """The schema or namespace this database instance is associated with
1112 (`str` or `None`).
1113 """