Coverage for python / lsst / daf / butler / registry / interfaces / _database.py: 21%
462 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-01-20 01:25 -0800
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 software is dual licensed under the GNU General Public License and also
10# under a 3-clause BSD license. Recipients may choose which of these licenses
11# to use; please see the files gpl-3.0.txt and/or bsd_license.txt,
12# respectively. If you choose the GPL option then the following text applies
13# (but note that there is still no warranty even if you opt for BSD instead):
14#
15# This program is free software: you can redistribute it and/or modify
16# it under the terms of the GNU General Public License as published by
17# the Free Software Foundation, either version 3 of the License, or
18# (at your option) any later version.
19#
20# This program is distributed in the hope that it will be useful,
21# but WITHOUT ANY WARRANTY; without even the implied warranty of
22# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
23# GNU General Public License for more details.
24#
25# You should have received a copy of the GNU General Public License
26# along with this program. If not, see <http://www.gnu.org/licenses/>.
27from __future__ import annotations
29from ... import ddl, time_utils
31__all__ = [
32 "Database",
33 "DatabaseConflictError",
34 "DatabaseInsertMode",
35 "DatabaseMetadata",
36 "ReadOnlyDatabaseError",
37 "SchemaAlreadyDefinedError",
38 "StaticTablesContext",
39]
41import enum
42import os
43import sys
44import uuid
45import warnings
46from abc import ABC, abstractmethod
47from collections import defaultdict
48from collections.abc import Callable, Iterable, Iterator, Sequence
49from contextlib import contextmanager
50from threading import Lock
51from typing import Any, cast, final
53import astropy.time
54import sqlalchemy
56from ..._named import NamedValueAbstractSet
57from ...name_shrinker import NameShrinker
58from ...timespan_database_representation import TimespanDatabaseRepresentation
59from .._exceptions import ConflictingDefinitionError
60from ._database_explain import get_query_plan
63class DatabaseInsertMode(enum.Enum):
64 """Mode options available for inserting database records."""
66 INSERT = enum.auto()
67 """Insert records, failing if they already exist."""
69 REPLACE = enum.auto()
70 """Replace records, overwriting existing."""
72 ENSURE = enum.auto()
73 """Insert records, skipping any that already exist."""
76# TODO: method is called with list[ReflectedColumn] in SA 2, and
77# ReflectedColumn does not exist in 1.4.
78def _checkExistingTableDefinition(name: str, spec: ddl.TableSpec, inspection: list) -> None:
79 """Test that the definition of a table in a `ddl.TableSpec` and from
80 database introspection are consistent.
82 Parameters
83 ----------
84 name : `str`
85 Name of the table (only used in error messages).
86 spec : `ddl.TableSpec`
87 Specification of the table.
88 inspection : `dict`
89 Dictionary returned by
90 `sqlalchemy.engine.reflection.Inspector.get_columns`.
92 Raises
93 ------
94 DatabaseConflictError
95 Raised if the definitions are inconsistent.
96 """
97 columnNames = [c["name"] for c in inspection]
98 if spec.fields.names != set(columnNames):
99 raise DatabaseConflictError(
100 f"Table '{name}' exists but is defined differently in the database; "
101 f"specification has columns {list(spec.fields.names)}, while the "
102 f"table in the database has {columnNames}."
103 )
106class ReadOnlyDatabaseError(RuntimeError):
107 """Exception raised when a write operation is called on a read-only
108 `Database`.
109 """
112class DatabaseConflictError(ConflictingDefinitionError):
113 """Exception raised when database content (row values or schema entities)
114 are inconsistent with what this client expects.
115 """
118class SchemaAlreadyDefinedError(RuntimeError):
119 """Exception raised when trying to initialize database schema when some
120 tables already exist.
121 """
124class StaticTablesContext:
125 """Helper class used to declare the static schema for a registry layer
126 in a database.
128 An instance of this class is returned by `Database.declareStaticTables`,
129 which should be the only way it should be constructed.
131 Parameters
132 ----------
133 db : `Database`
134 The database.
135 connection : `sqlalchemy.engine.Connection`
136 The connection object.
137 """
139 def __init__(self, db: Database, connection: sqlalchemy.engine.Connection):
140 self._db = db
141 self._inspector = sqlalchemy.inspect(connection)
142 self._tableNames = frozenset(self._inspector.get_table_names(schema=self._db.namespace))
143 self._initializers: list[Callable[[Database], None]] = []
145 def addTable(self, name: str, spec: ddl.TableSpec) -> sqlalchemy.schema.Table:
146 """Add a new table to the schema, returning its sqlalchemy
147 representation.
149 Parameters
150 ----------
151 name : `str`
152 The name of the table.
153 spec : `ddl.TableSpec`
154 The specification of the table.
156 Returns
157 -------
158 table : `sqlalchemy.schema.Table`
159 The created table.
161 Notes
162 -----
163 The new table may not actually be created until the end of the
164 context created by `Database.declareStaticTables`, allowing tables
165 to be declared in any order even in the presence of foreign key
166 relationships.
167 """
168 metadata = self._db._metadata
169 assert metadata is not None, "Guaranteed by context manager that returns this object."
170 return metadata.add_table(self._db, name, spec)
172 def addTableTuple(self, specs: tuple[ddl.TableSpec, ...]) -> tuple[sqlalchemy.schema.Table, ...]:
173 """Add a named tuple of tables to the schema, returning their
174 SQLAlchemy representations in a named tuple of the same type.
176 The new tables may not actually be created until the end of the
177 context created by `Database.declareStaticTables`, allowing tables
178 to be declared in any order even in the presence of foreign key
179 relationships.
181 Parameters
182 ----------
183 specs : `tuple` of `ddl.TableSpec`
184 Specifications of multiple tables.
186 Returns
187 -------
188 tables : `tuple` of `sqlalchemy.schema.Table`
189 All the tables created.
191 Notes
192 -----
193 ``specs`` *must* be an instance of a type created by
194 `collections.namedtuple`, not just regular tuple, and the returned
195 object is guaranteed to be the same. Because `~collections.namedtuple`
196 is just a factory for `type` objects, not an actual type itself,
197 we cannot represent this with type annotations.
198 """
199 return specs._make( # type: ignore
200 self.addTable(name, spec)
201 for name, spec in zip(specs._fields, specs, strict=True) # type: ignore
202 )
204 def addInitializer(self, initializer: Callable[[Database], None]) -> None:
205 """Add a method that does one-time initialization of a database.
207 Initialization can mean anything that changes state of a database
208 and needs to be done exactly once after database schema was created.
209 An example for that could be population of schema attributes.
211 Parameters
212 ----------
213 initializer : `~collections.abc.Callable`
214 Method of a single argument which is a `Database` instance.
215 """
216 self._initializers.append(initializer)
219class Database(ABC):
220 """An abstract interface that represents a particular database engine's
221 representation of a single schema/namespace/database.
223 Parameters
224 ----------
225 origin : `int`
226 An integer ID that should be used as the default for any datasets,
227 quanta, or other entities that use a (autoincrement, origin) compound
228 primary key.
229 engine : `sqlalchemy.engine.Engine`
230 The SQLAlchemy engine for this `Database`.
231 namespace : `str`, optional
232 Name of the schema or namespace this instance is associated with.
233 This is passed as the ``schema`` argument when constructing a
234 `sqlalchemy.schema.MetaData` instance. We use ``namespace`` instead to
235 avoid confusion between "schema means namespace" and "schema means
236 table definitions".
237 metadata : `sqlalchemy.schema.MetaData`, optional
238 Object representing the tables and other schema entities. If not
239 provided, will be generated during the next call to
240 ``declareStaticTables``.
241 allow_temporary_tables : `bool`, optional
242 If `True`, database operations will be allowed to use temporary tables.
243 If `False`, other SQL constructs will be used instead of temporary
244 tables when possible.
246 Notes
247 -----
248 `Database` requires all write operations to go through its special named
249 methods. Our write patterns are sufficiently simple that we don't really
250 need the full flexibility of SQL insert/update/delete syntax, and we need
251 non-standard (but common) functionality in these operations sufficiently
252 often that it seems worthwhile to provide our own generic API.
254 In contrast, `Database.query` allows arbitrary ``SELECT`` queries (via
255 their SQLAlchemy representation) to be run, as we expect these to require
256 significantly more sophistication while still being limited to standard
257 SQL.
259 `Database` itself has several underscore-prefixed attributes:
261 - ``_engine``: SQLAlchemy object representing its engine.
262 - ``_connection``: method returning a context manager for
263 `sqlalchemy.engine.Connection` object.
264 - ``_metadata``: the `sqlalchemy.schema.MetaData` object representing
265 the tables and other schema entities.
267 These are considered protected (derived classes may access them, but other
268 code should not), and read-only, aside from executing SQL via
269 ``_connection``.
270 """
272 def __init__(
273 self,
274 *,
275 origin: int,
276 engine: sqlalchemy.engine.Engine,
277 namespace: str | None = None,
278 metadata: DatabaseMetadata | None = None,
279 allow_temporary_tables: bool = True,
280 ):
281 self.origin = origin
282 self.name_shrinker = NameShrinker(engine.dialect.max_identifier_length)
283 self.namespace = namespace
284 self._engine = engine
285 self._session_connection: sqlalchemy.engine.Connection | None = None
286 self._temp_tables: set[str] = set()
287 self._metadata = metadata
288 self._allow_temporary_tables = allow_temporary_tables
290 def dispose(self) -> None:
291 """Close all open database connections held by this `Database`
292 instance.
293 """
294 self._engine.dispose()
296 def __repr__(self) -> str:
297 # Rather than try to reproduce all the parameters used to create
298 # the object, instead report the more useful information of the
299 # connection URL.
300 if self._engine.url.password is not None:
301 uri = str(self._engine.url.set(password="***"))
302 else:
303 uri = str(self._engine.url)
304 if self.namespace:
305 uri += f"#{self.namespace}"
306 return f'{type(self).__name__}("{uri}")'
308 @classmethod
309 def makeDefaultUri(cls, root: str) -> str | None:
310 """Create a default connection URI appropriate for the given root
311 directory, or `None` if there can be no such default.
313 Parameters
314 ----------
315 root : `str`
316 Root string to use to build connection URI.
318 Returns
319 -------
320 uri : `str` or `None`
321 The URI string or `None`.
322 """
323 return None
325 @classmethod
326 def fromUri(
327 cls,
328 uri: str | sqlalchemy.engine.URL,
329 *,
330 origin: int,
331 namespace: str | None = None,
332 writeable: bool = True,
333 allow_temporary_tables: bool = True,
334 ) -> Database:
335 """Construct a database from a SQLAlchemy URI.
337 Parameters
338 ----------
339 uri : `str` or `sqlalchemy.engine.URL`
340 A SQLAlchemy URI connection string.
341 origin : `int`
342 An integer ID that should be used as the default for any datasets,
343 quanta, or other entities that use a (autoincrement, origin)
344 compound primary key.
345 namespace : `str`, optional
346 A database namespace (i.e. schema) the new instance should be
347 associated with. If `None` (default), the namespace (if any) is
348 inferred from the URI.
349 writeable : `bool`, optional
350 If `True`, allow write operations on the database, including
351 ``CREATE TABLE``.
352 allow_temporary_tables : `bool`, optional
353 If `True`, database operations will be allowed to use temporary
354 tables.
355 If `False`, other SQL constructs will be used instead of temporary
356 tables when possible.
358 Returns
359 -------
360 db : `Database`
361 A new `Database` instance.
362 """
363 db = cls.fromEngine(
364 cls.makeEngine(uri, writeable=writeable), origin=origin, namespace=namespace, writeable=writeable
365 )
366 db._allow_temporary_tables = allow_temporary_tables
367 return db
369 @abstractmethod
370 def clone(self) -> Database:
371 """Make an independent copy of this `Database` object.
373 Returns
374 -------
375 db : `Database`
376 A new `Database` instance with the same configuration as this
377 instance.
378 """
379 raise NotImplementedError()
381 @classmethod
382 @abstractmethod
383 def makeEngine(
384 cls, uri: str | sqlalchemy.engine.URL, *, writeable: bool = True
385 ) -> sqlalchemy.engine.Engine:
386 """Create a `sqlalchemy.engine.Engine` from a SQLAlchemy URI.
388 Parameters
389 ----------
390 uri : `str` or `sqlalchemy.engine.URL`
391 A SQLAlchemy URI connection string.
392 writeable : `bool`, optional
393 If `True`, allow write operations on the database, including
394 ``CREATE TABLE``.
396 Returns
397 -------
398 engine : `sqlalchemy.engine.Engine`
399 A database engine.
401 Notes
402 -----
403 Subclasses that support other ways to connect to a database are
404 encouraged to add optional arguments to their implementation of this
405 method, as long as they maintain compatibility with the base class
406 call signature.
407 """
408 raise NotImplementedError()
410 @classmethod
411 @abstractmethod
412 def fromEngine(
413 cls,
414 engine: sqlalchemy.engine.Engine,
415 *,
416 origin: int,
417 namespace: str | None = None,
418 writeable: bool = True,
419 ) -> Database:
420 """Create a new `Database` from an existing `sqlalchemy.engine.Engine`.
422 Parameters
423 ----------
424 engine : `sqlalchemy.engine.Engine`
425 The engine for the database. May be shared between `Database`
426 instances.
427 origin : `int`
428 An integer ID that should be used as the default for any datasets,
429 quanta, or other entities that use a (autoincrement, origin)
430 compound primary key.
431 namespace : `str`, optional
432 A different database namespace (i.e. schema) the new instance
433 should be associated with. If `None` (default), the namespace
434 (if any) is inferred from the connection.
435 writeable : `bool`, optional
436 If `True`, allow write operations on the database, including
437 ``CREATE TABLE``.
439 Returns
440 -------
441 db : `Database`
442 A new `Database` instance.
444 Notes
445 -----
446 This method allows different `Database` instances to share the same
447 engine, which is desirable when they represent different namespaces
448 can be queried together.
449 """
450 raise NotImplementedError()
452 @final
453 @contextmanager
454 def session(self) -> Iterator[None]:
455 """Return a context manager that represents a session (persistent
456 connection to a database).
458 Returns
459 -------
460 context : `AbstractContextManager` [ `None` ]
461 A context manager that does not return a value when entered.
463 Notes
464 -----
465 This method should be used when a sequence of read-only SQL operations
466 will be performed in rapid succession *without* a requirement that they
467 yield consistent results in the presence of concurrent writes (or, more
468 rarely, when conflicting concurrent writes are rare/impossible and the
469 session will be open long enough that a transaction is inadvisable).
470 """
471 with self._session():
472 yield
474 @final
475 @contextmanager
476 def transaction(
477 self,
478 *,
479 interrupting: bool = False,
480 savepoint: bool = False,
481 lock: Iterable[sqlalchemy.schema.Table] = (),
482 for_temp_tables: bool = False,
483 ) -> Iterator[None]:
484 """Return a context manager that represents a transaction.
486 Parameters
487 ----------
488 interrupting : `bool`, optional
489 If `True` (`False` is default), this transaction block may not be
490 nested within an outer one, and attempting to do so is a logic
491 (i.e. assertion) error.
492 savepoint : `bool`, optional
493 If `True` (`False` is default), create a `SAVEPOINT`, allowing
494 exceptions raised by the database (e.g. due to constraint
495 violations) during this transaction's context to be caught outside
496 it without also rolling back all operations in an outer transaction
497 block. If `False`, transactions may still be nested, but a
498 rollback may be generated at any level and affects all levels, and
499 commits are deferred until the outermost block completes. If any
500 outer transaction block was created with ``savepoint=True``, all
501 inner blocks will be as well (regardless of the actual value
502 passed). This has no effect if this is the outermost transaction.
503 lock : `~collections.abc.Iterable` [ `sqlalchemy.schema.Table` ], \
504 optional
505 A list of tables to lock for the duration of this transaction.
506 These locks are guaranteed to prevent concurrent writes and allow
507 this transaction (only) to acquire the same locks (others should
508 block), but only prevent concurrent reads if the database engine
509 requires that in order to block concurrent writes.
510 for_temp_tables : `bool`, optional
511 If `True`, this transaction may involve creating temporary tables.
513 Returns
514 -------
515 context : `AbstractContextManager` [ `None` ]
516 A context manager that commits the transaction when it is exited
517 without error and rolls back the transactoin when it is exited via
518 an exception.
520 Notes
521 -----
522 All transactions on a connection managed by one or more `Database`
523 instances _must_ go through this method, or transaction state will not
524 be correctly managed.
525 """
526 with self._transaction(
527 interrupting=interrupting, savepoint=savepoint, lock=lock, for_temp_tables=for_temp_tables
528 ):
529 yield
531 @contextmanager
532 def temporary_table(
533 self, spec: ddl.TableSpec, name: str | None = None
534 ) -> Iterator[sqlalchemy.schema.Table]:
535 """Return a context manager that creates and then drops a temporary
536 table.
538 Parameters
539 ----------
540 spec : `ddl.TableSpec`
541 Specification for the columns. Unique and foreign key constraints
542 may be ignored.
543 name : `str`, optional
544 If provided, the name of the SQL construct. If not provided, an
545 opaque but unique identifier is generated.
547 Returns
548 -------
549 context : `AbstractContextManager` [ `sqlalchemy.schema.Table` ]
550 A context manager that returns a SQLAlchemy representation of the
551 temporary table when entered.
553 Notes
554 -----
555 Temporary tables may be created, dropped, and written to even in
556 read-only databases - at least according to the Python-level
557 protections in the `Database` classes. Server permissions may say
558 otherwise, but in that case they probably need to be modified to
559 support the full range of expected read-only butler behavior.
560 """
561 assert self._metadata is not None, "Static tables must be created before temporary tables"
562 if not self.supports_temporary_tables:
563 raise ReadOnlyDatabaseError("Creation of temporary tables is not supported by this database.")
564 with self._session() as connection:
565 table = self._make_temporary_table(connection, spec=spec, name=name)
566 self._temp_tables.add(table.key)
567 try:
568 yield table
569 finally:
570 self._temp_tables.remove(table.key)
571 self._metadata.remove_table(table.name)
573 @contextmanager
574 def _session(self) -> Iterator[sqlalchemy.engine.Connection]:
575 """Protected implementation for `session` that actually returns the
576 connection.
578 This method is for internal `Database` calls that need the actual
579 SQLAlchemy connection object. It should be overridden by subclasses
580 instead of `session` itself.
582 Returns
583 -------
584 context : `AbstractContextManager` [ `sqlalchemy.engine.Connection` ]
585 A context manager that returns a SQLALchemy connection when
586 entered.
588 """
589 if self._session_connection is not None:
590 # session already started, just reuse that
591 yield self._session_connection
592 else:
593 try:
594 # open new connection and close it when done
595 self._session_connection = self._engine.connect()
596 yield self._session_connection
597 finally:
598 if self._session_connection is not None:
599 self._session_connection.close()
600 self._session_connection = None
601 # Temporary tables only live within session
602 self._temp_tables = set()
604 @contextmanager
605 def _transaction(
606 self,
607 *,
608 interrupting: bool = False,
609 savepoint: bool = False,
610 lock: Iterable[sqlalchemy.schema.Table] = (),
611 for_temp_tables: bool = False,
612 ) -> Iterator[tuple[bool, sqlalchemy.engine.Connection]]:
613 """Protected implementation for `transaction` that actually returns the
614 connection and whether this is a new outermost transaction.
616 This method is for internal `Database` calls that need the actual
617 SQLAlchemy connection object. It should be overridden by subclasses
618 instead of `transaction` itself.
620 Parameters
621 ----------
622 interrupting : `bool`, optional
623 If `True` (`False` is default), this transaction block may not be
624 nested without an outer one, and attempting to do so is a logic
625 (i.e. assertion) error.
626 savepoint : `bool`, optional
627 If `True` (`False` is default), create a `SAVEPOINT`, allowing
628 exceptions raised by the database (e.g. due to constraint
629 violations) during this transaction's context to be caught outside
630 it without also rolling back all operations in an outer transaction
631 block. If `False`, transactions may still be nested, but a
632 rollback may be generated at any level and affects all levels, and
633 commits are deferred until the outermost block completes. If any
634 outer transaction block was created with ``savepoint=True``, all
635 inner blocks will be as well (regardless of the actual value
636 passed). This has no effect if this is the outermost transaction.
637 lock : `~collections.abc.Iterable` [ `sqlalchemy.schema.Table` ], \
638 optional
639 A list of tables to lock for the duration of this transaction.
640 These locks are guaranteed to prevent concurrent writes and allow
641 this transaction (only) to acquire the same locks (others should
642 block), but only prevent concurrent reads if the database engine
643 requires that in order to block concurrent writes.
644 for_temp_tables : `bool`, optional
645 If `True`, this transaction may involve creating temporary tables.
647 Returns
648 -------
649 context : `AbstractContextManager` [ `tuple` [ `bool`,
650 `sqlalchemy.engine.Connection` ] ]
651 A context manager that commits the transaction when it is exited
652 without error and rolls back the transactoin when it is exited via
653 an exception. When entered, it returns a tuple of:
655 - ``is_new`` (`bool`): whether this is a new (outermost)
656 transaction;
657 - ``connection`` (`sqlalchemy.engine.Connection`): the connection.
658 """
659 with self._session() as connection:
660 already_in_transaction = self.isInTransaction()
661 assert not (interrupting and already_in_transaction), (
662 "Logic error in transaction nesting: an operation that would "
663 "interrupt the active transaction context has been requested."
664 )
665 savepoint = savepoint or connection.in_nested_transaction()
666 trans: sqlalchemy.engine.Transaction | None
667 if already_in_transaction:
668 if savepoint:
669 trans = connection.begin_nested()
670 else:
671 # Nested non-savepoint transactions don't do anything.
672 trans = None
673 else:
674 # Use a regular (non-savepoint) transaction always for the
675 # outermost context.
676 trans = connection.begin()
677 self._lockTables(connection, lock)
678 try:
679 yield not already_in_transaction, connection
680 if trans is not None:
681 trans.commit()
682 except BaseException:
683 if trans is not None:
684 trans.rollback()
685 raise
687 @abstractmethod
688 def _lockTables(
689 self, connection: sqlalchemy.engine.Connection, tables: Iterable[sqlalchemy.schema.Table] = ()
690 ) -> None:
691 """Acquire locks on the given tables.
693 This is an implementation hook for subclasses, called by `transaction`.
694 It should not be called directly by other code.
696 Parameters
697 ----------
698 connection : `sqlalchemy.engine.Connection`
699 Database connection object. It is guaranteed that transaction is
700 already in a progress for this connection.
701 tables : `~collections.abc.Iterable` [ `sqlalchemy.schema.Table` ], \
702 optional
703 A list of tables to lock for the duration of this transaction.
704 These locks are guaranteed to prevent concurrent writes and allow
705 this transaction (only) to acquire the same locks (others should
706 block), but only prevent concurrent reads if the database engine
707 requires that in order to block concurrent writes.
708 """
709 raise NotImplementedError()
711 def isTableWriteable(self, table: sqlalchemy.schema.Table) -> bool:
712 """Check whether a table is writeable, either because the database
713 connection is read-write or the table is a temporary table.
715 Parameters
716 ----------
717 table : `sqlalchemy.schema.Table`
718 SQLAlchemy table object to check.
720 Returns
721 -------
722 writeable : `bool`
723 Whether this table is writeable.
724 """
725 return self.isWriteable() or table.key in self._temp_tables
727 def assertTableWriteable(self, table: sqlalchemy.schema.Table, msg: str) -> None:
728 """Raise if the given table is not writeable, either because the
729 database connection is read-write or the table is a temporary table.
731 Parameters
732 ----------
733 table : `sqlalchemy.schema.Table`
734 SQLAlchemy table object to check.
735 msg : `str`, optional
736 If provided, raise `ReadOnlyDatabaseError` instead of returning
737 `False`, with this message.
738 """
739 if not self.isTableWriteable(table):
740 raise ReadOnlyDatabaseError(msg)
742 @contextmanager
743 def declareStaticTables(self, *, create: bool) -> Iterator[StaticTablesContext]:
744 """Return a context manager in which the database's static DDL schema
745 can be declared.
747 Parameters
748 ----------
749 create : `bool`
750 If `True`, attempt to create all tables at the end of the context.
751 If `False`, they will be assumed to already exist.
753 Returns
754 -------
755 schema : `StaticTablesContext`
756 A helper object that is used to add new tables.
758 Raises
759 ------
760 ReadOnlyDatabaseError
761 Raised if ``create`` is `True`, `Database.isWriteable` is `False`,
762 and one or more declared tables do not already exist.
764 Examples
765 --------
766 Given a `Database` instance ``db``::
768 with db.declareStaticTables(create=True) as schema:
769 schema.addTable("table1", TableSpec(...))
770 schema.addTable("table2", TableSpec(...))
772 Notes
773 -----
774 A database's static DDL schema must be declared before any dynamic
775 tables are managed via calls to `ensureTableExists` or
776 `getExistingTable`. The order in which static schema tables are added
777 inside the context block is unimportant; they will automatically be
778 sorted and added in an order consistent with their foreign key
779 relationships.
780 """
781 if create and not self.isWriteable():
782 raise ReadOnlyDatabaseError(f"Cannot create tables in read-only database {self}.")
784 self._metadata = DatabaseMetadata(self.namespace)
785 try:
786 with self._transaction() as (_, connection):
787 context = StaticTablesContext(self, connection)
788 if create and context._tableNames:
789 # Looks like database is already initalized, to avoid
790 # danger of modifying/destroying valid schema we refuse to
791 # do anything in this case
792 raise SchemaAlreadyDefinedError(f"Cannot create tables in non-empty database {self}.")
793 yield context
794 if create:
795 if (
796 self.namespace is not None
797 and self.namespace not in context._inspector.get_schema_names()
798 ):
799 connection.execute(sqlalchemy.schema.CreateSchema(self.namespace))
800 # In our tables we have columns that make use of sqlalchemy
801 # Sequence objects. There is currently a bug in sqlalchemy
802 # that causes a deprecation warning to be thrown on a
803 # property of the Sequence object when the repr for the
804 # sequence is created. Here a filter is used to catch these
805 # deprecation warnings when tables are created.
806 with warnings.catch_warnings():
807 warnings.simplefilter("ignore", category=sqlalchemy.exc.SADeprecationWarning)
808 self._metadata.create_all(connection)
809 # call all initializer methods sequentially
810 for init in context._initializers:
811 init(self)
812 except BaseException:
813 self._metadata = None
814 raise
816 @abstractmethod
817 def isWriteable(self) -> bool:
818 """Return `True` if this database can be modified by this client."""
819 raise NotImplementedError()
821 def isInTransaction(self) -> bool:
822 """Return `True` if there is currently a database connection open with
823 an active transaction; `False` otherwise.
824 """
825 session = self._session_connection
826 return session is not None and session.in_transaction()
828 @abstractmethod
829 def __str__(self) -> str:
830 """Return a human-readable identifier for this `Database`, including
831 any namespace or schema that identifies its names within a `Registry`.
832 """
833 raise NotImplementedError()
835 @property
836 def dialect(self) -> sqlalchemy.engine.Dialect:
837 """The SQLAlchemy dialect for this database engine
838 (`sqlalchemy.engine.Dialect`).
839 """
840 return self._engine.dialect
842 def shrinkDatabaseEntityName(self, original: str) -> str:
843 """Return a version of the given name that fits within this database
844 engine's length limits for table, constraint, indexes, and sequence
845 names.
847 Implementations should not assume that simple truncation is safe,
848 because multiple long names often begin with the same prefix.
850 The default implementation simply returns the given name.
852 Parameters
853 ----------
854 original : `str`
855 The original name.
857 Returns
858 -------
859 shrunk : `str`
860 The new, possibly shortened name.
861 """
862 return original
864 def expandDatabaseEntityName(self, shrunk: str) -> str:
865 """Retrieve the original name for a database entity that was too long
866 to fit within the database engine's limits.
868 Parameters
869 ----------
870 shrunk : `str`
871 The original name.
873 Returns
874 -------
875 shrunk : `str`
876 The new, possibly shortened name.
877 """
878 return shrunk
880 def _makeColumnConstraints(self, table: str, spec: ddl.FieldSpec) -> list[sqlalchemy.CheckConstraint]:
881 """Create constraints based on this spec.
883 Parameters
884 ----------
885 table : `str`
886 Name of the table this column is being added to.
887 spec : `FieldSpec`
888 Specification for the field to be added.
890 Returns
891 -------
892 constraint : `list` of `sqlalchemy.CheckConstraint`
893 Constraint added for this column.
894 """
895 # By default we return no additional constraints
896 return []
898 def _convertFieldSpec(
899 self, table: str, spec: ddl.FieldSpec, metadata: sqlalchemy.MetaData, **kwargs: Any
900 ) -> sqlalchemy.schema.Column:
901 """Convert a `FieldSpec` to a `sqlalchemy.schema.Column`.
903 Parameters
904 ----------
905 table : `str`
906 Name of the table this column is being added to.
907 spec : `FieldSpec`
908 Specification for the field to be added.
909 metadata : `sqlalchemy.MetaData`
910 SQLAlchemy representation of the DDL schema this field's table is
911 being added to.
912 **kwargs
913 Additional keyword arguments to forward to the
914 `sqlalchemy.schema.Column` constructor. This is provided to make
915 it easier for derived classes to delegate to ``super()`` while
916 making only minor changes.
918 Returns
919 -------
920 column : `sqlalchemy.schema.Column`
921 SQLAlchemy representation of the field.
922 """
923 args = []
924 if spec.autoincrement:
925 # Generate a sequence to use for auto incrementing for databases
926 # that do not support it natively. This will be ignored by
927 # sqlalchemy for databases that do support it.
928 args.append(
929 sqlalchemy.Sequence(
930 self.shrinkDatabaseEntityName(f"{table}_seq_{spec.name}"), metadata=metadata
931 )
932 )
933 assert spec.doc is None or isinstance(spec.doc, str), f"Bad doc for {table}.{spec.name}."
934 return sqlalchemy.schema.Column(
935 spec.name,
936 spec.getSizedColumnType(),
937 *args,
938 nullable=spec.nullable,
939 primary_key=spec.primaryKey,
940 comment=spec.doc,
941 server_default=spec.default,
942 **kwargs,
943 )
945 def _convertForeignKeySpec(
946 self, table: str, spec: ddl.ForeignKeySpec, metadata: sqlalchemy.MetaData, **kwargs: Any
947 ) -> sqlalchemy.schema.ForeignKeyConstraint:
948 """Convert a `ForeignKeySpec` to a
949 `sqlalchemy.schema.ForeignKeyConstraint`.
951 Parameters
952 ----------
953 table : `str`
954 Name of the table this foreign key is being added to.
955 spec : `ForeignKeySpec`
956 Specification for the foreign key to be added.
957 metadata : `sqlalchemy.MetaData`
958 SQLAlchemy representation of the DDL schema this constraint is
959 being added to.
960 **kwargs
961 Additional keyword arguments to forward to the
962 `sqlalchemy.schema.ForeignKeyConstraint` constructor. This is
963 provided to make it easier for derived classes to delegate to
964 ``super()`` while making only minor changes.
966 Returns
967 -------
968 constraint : `sqlalchemy.schema.ForeignKeyConstraint`
969 SQLAlchemy representation of the constraint.
970 """
971 name = self.shrinkDatabaseEntityName(
972 "_".join(["fkey", table, spec.table] + list(spec.target) + list(spec.source))
973 )
974 return sqlalchemy.schema.ForeignKeyConstraint(
975 spec.source,
976 [f"{spec.table}.{col}" for col in spec.target],
977 name=name,
978 ondelete=spec.onDelete,
979 )
981 def _convertExclusionConstraintSpec(
982 self,
983 table: str,
984 spec: tuple[str | type[TimespanDatabaseRepresentation], ...],
985 metadata: sqlalchemy.MetaData,
986 ) -> sqlalchemy.schema.Constraint:
987 """Convert a `tuple` from `ddl.TableSpec.exclusion` into a SQLAlchemy
988 constraint representation.
990 Parameters
991 ----------
992 table : `str`
993 Name of the table this constraint is being added to.
994 spec : `tuple` [ `str` or `type` ]
995 A tuple of `str` column names and the `type` object returned by
996 `getTimespanRepresentation` (which must appear exactly once),
997 indicating the order of the columns in the index used to back the
998 constraint.
999 metadata : `sqlalchemy.MetaData`
1000 SQLAlchemy representation of the DDL schema this constraint is
1001 being added to.
1003 Returns
1004 -------
1005 constraint : `sqlalchemy.schema.Constraint`
1006 SQLAlchemy representation of the constraint.
1008 Raises
1009 ------
1010 NotImplementedError
1011 Raised if this database does not support exclusion constraints.
1012 """
1013 raise NotImplementedError(f"Database {self} does not support exclusion constraints.")
1015 def _convertTableSpec(
1016 self, name: str, spec: ddl.TableSpec, metadata: sqlalchemy.MetaData, **kwargs: Any
1017 ) -> sqlalchemy.schema.Table:
1018 """Convert a `TableSpec` to a `sqlalchemy.schema.Table`.
1020 Parameters
1021 ----------
1022 name : `str`
1023 The name of the table.
1024 spec : `TableSpec`
1025 Specification for the foreign key to be added.
1026 metadata : `sqlalchemy.MetaData`
1027 SQLAlchemy representation of the DDL schema this table is being
1028 added to.
1029 **kwargs
1030 Additional keyword arguments to forward to the
1031 `sqlalchemy.schema.Table` constructor. This is provided to make it
1032 easier for derived classes to delegate to ``super()`` while making
1033 only minor changes.
1035 Returns
1036 -------
1037 table : `sqlalchemy.schema.Table`
1038 SQLAlchemy representation of the table.
1040 Notes
1041 -----
1042 This method does not handle ``spec.foreignKeys`` at all, in order to
1043 avoid circular dependencies. These are added by higher-level logic in
1044 `ensureTableExists`, `getExistingTable`, and `declareStaticTables`.
1045 """
1046 args: list[sqlalchemy.schema.SchemaItem] = [
1047 self._convertFieldSpec(name, fieldSpec, metadata) for fieldSpec in spec.fields
1048 ]
1050 # Add any column constraints
1051 for fieldSpec in spec.fields:
1052 args.extend(self._makeColumnConstraints(name, fieldSpec))
1054 # Track indexes added for primary key and unique constraints, to make
1055 # sure we don't add duplicate explicit or foreign key indexes for
1056 # those.
1057 allIndexes = {tuple(fieldSpec.name for fieldSpec in spec.fields if fieldSpec.primaryKey)}
1058 args.extend(
1059 sqlalchemy.schema.UniqueConstraint(
1060 *columns, name=self.shrinkDatabaseEntityName("_".join([name, "unq"] + list(columns)))
1061 )
1062 for columns in spec.unique
1063 )
1064 allIndexes.update(spec.unique)
1065 args.extend(
1066 sqlalchemy.schema.Index(
1067 self.shrinkDatabaseEntityName("_".join([name, "idx"] + list(index.columns))),
1068 *index.columns,
1069 unique=(index.columns in spec.unique),
1070 **index.kwargs,
1071 )
1072 for index in spec.indexes
1073 if index.columns not in allIndexes
1074 )
1075 allIndexes.update(index.columns for index in spec.indexes)
1076 args.extend(
1077 sqlalchemy.schema.Index(
1078 self.shrinkDatabaseEntityName("_".join((name, "fkidx") + fk.source)),
1079 *fk.source,
1080 )
1081 for fk in spec.foreignKeys
1082 if fk.addIndex and fk.source not in allIndexes
1083 )
1085 args.extend(self._convertExclusionConstraintSpec(name, excl, metadata) for excl in spec.exclusion)
1087 assert spec.doc is None or isinstance(spec.doc, str), f"Bad doc for {name}."
1088 return sqlalchemy.schema.Table(name, metadata, *args, comment=spec.doc, info={"spec": spec}, **kwargs)
1090 def ensureTableExists(self, name: str, spec: ddl.TableSpec) -> sqlalchemy.schema.Table:
1091 """Ensure that a table with the given name and specification exists,
1092 creating it if necessary.
1094 Parameters
1095 ----------
1096 name : `str`
1097 Name of the table (not including namespace qualifiers).
1098 spec : `TableSpec`
1099 Specification for the table. This will be used when creating the
1100 table, and *may* be used when obtaining an existing table to check
1101 for consistency, but no such check is guaranteed.
1103 Returns
1104 -------
1105 table : `sqlalchemy.schema.Table`
1106 SQLAlchemy representation of the table.
1108 Raises
1109 ------
1110 ReadOnlyDatabaseError
1111 Raised if `isWriteable` returns `False`, and the table does not
1112 already exist.
1113 DatabaseConflictError
1114 Raised if the table exists but ``spec`` is inconsistent with its
1115 definition.
1117 Notes
1118 -----
1119 This method may not be called within transactions. It may be called on
1120 read-only databases if and only if the table does in fact already
1121 exist.
1123 Subclasses may override this method, but usually should not need to.
1124 """
1125 # TODO: if _engine is used to make a table then it uses separate
1126 # connection and should not interfere with current transaction
1127 assert not self.isInTransaction(), "Table creation interrupts transactions."
1128 assert self._metadata is not None, "Static tables must be declared before dynamic tables."
1129 table = self.getExistingTable(name, spec)
1130 if table is not None:
1131 return table
1132 if not self.isWriteable():
1133 raise ReadOnlyDatabaseError(
1134 f"Table {name} does not exist, and cannot be created because database {self} is read-only."
1135 )
1137 table = self._metadata.add_table(self, name, spec)
1138 try:
1139 with self._transaction() as (_, connection):
1140 table.create(connection)
1141 except sqlalchemy.exc.DatabaseError:
1142 # Some other process could have created the table meanwhile, which
1143 # usually causes OperationalError or ProgrammingError. We cannot
1144 # use IF NOT EXISTS clause in this case due to PostgreSQL race
1145 # condition on server side which causes IntegrityError. Instead we
1146 # catch these exceptions (they all inherit DatabaseError) and
1147 # re-check whether table is now there.
1148 table = self.getExistingTable(name, spec)
1149 if table is None:
1150 raise
1151 return table
1153 def getExistingTable(self, name: str, spec: ddl.TableSpec) -> sqlalchemy.schema.Table | None:
1154 """Obtain an existing table with the given name and specification.
1156 Parameters
1157 ----------
1158 name : `str`
1159 Name of the table (not including namespace qualifiers).
1160 spec : `TableSpec`
1161 Specification for the table. This will be used when creating the
1162 SQLAlchemy representation of the table, and it is used to
1163 check that the actual table in the database is consistent.
1165 Returns
1166 -------
1167 table : `sqlalchemy.schema.Table` or `None`
1168 SQLAlchemy representation of the table, or `None` if it does not
1169 exist.
1171 Raises
1172 ------
1173 DatabaseConflictError
1174 Raised if the table exists but ``spec`` is inconsistent with its
1175 definition.
1177 Notes
1178 -----
1179 This method can be called within transactions and never modifies the
1180 database.
1182 Subclasses may override this method, but usually should not need to.
1183 """
1184 assert self._metadata is not None, "Static tables must be declared before dynamic tables."
1185 table = self._metadata.get_table(name)
1186 if table is not None:
1187 if spec.fields.names != set(table.columns.keys()):
1188 raise DatabaseConflictError(
1189 f"Table '{name}' has already been defined differently; the new "
1190 f"specification has columns {list(spec.fields.names)}, while "
1191 f"the previous definition has {list(table.columns.keys())}."
1192 )
1193 else:
1194 inspector = sqlalchemy.inspect(
1195 self._engine if self._session_connection is None else self._session_connection, raiseerr=True
1196 )
1197 if name in inspector.get_table_names(schema=self.namespace):
1198 _checkExistingTableDefinition(name, spec, inspector.get_columns(name, schema=self.namespace))
1199 return self._metadata.add_table(self, name, spec)
1200 return table
1202 def _make_temporary_table(
1203 self,
1204 connection: sqlalchemy.engine.Connection,
1205 spec: ddl.TableSpec,
1206 name: str | None = None,
1207 **kwargs: Any,
1208 ) -> sqlalchemy.schema.Table:
1209 """Create a temporary table.
1211 Parameters
1212 ----------
1213 connection : `sqlalchemy.engine.Connection`
1214 Connection to use when creating the table.
1215 spec : `TableSpec`
1216 Specification for the table.
1217 name : `str`, optional
1218 A unique (within this session/connetion) name for the table.
1219 Subclasses may override to modify the actual name used. If not
1220 provided, a unique name will be generated.
1221 **kwargs
1222 Additional keyword arguments to forward to the
1223 `sqlalchemy.schema.Table` constructor. This is provided to make it
1224 easier for derived classes to delegate to ``super()`` while making
1225 only minor changes.
1227 Returns
1228 -------
1229 table : `sqlalchemy.schema.Table`
1230 SQLAlchemy representation of the table.
1231 """
1232 if name is None:
1233 name = f"tmp_{uuid.uuid4().hex}"
1234 if self._metadata is None:
1235 raise RuntimeError("Cannot create temporary table before static schema is defined.")
1236 table = self._metadata.add_table(
1237 self, name, spec, prefixes=["TEMPORARY"], schema=sqlalchemy.schema.BLANK_SCHEMA, **kwargs
1238 )
1239 if table.key in self._temp_tables and table.key != name:
1240 raise ValueError(
1241 f"A temporary table with name {name} (transformed to {table.key} by Database) already exists."
1242 )
1243 with self._transaction():
1244 table.create(connection)
1245 return table
1247 @classmethod
1248 def getTimespanRepresentation(cls) -> type[TimespanDatabaseRepresentation]:
1249 """Return a `type` that encapsulates the way `Timespan` objects are
1250 stored in this database.
1252 `Database` does not automatically use the return type of this method
1253 anywhere else; calling code is responsible for making sure that DDL
1254 and queries are consistent with it.
1256 Returns
1257 -------
1258 TimespanReprClass : `type` (`TimespanDatabaseRepresention` subclass)
1259 A type that encapsulates the way `Timespan` objects should be
1260 stored in this database.
1262 Notes
1263 -----
1264 There are two big reasons we've decided to keep timespan-mangling logic
1265 outside the `Database` implementations, even though the choice of
1266 representation is ultimately up to a `Database` implementation:
1268 - Timespans appear in relatively few tables and queries in our
1269 typical usage, and the code that operates on them is already aware
1270 that it is working with timespans. In contrast, a
1271 timespan-representation-aware implementation of, say, `insert`,
1272 would need to have extra logic to identify when timespan-mangling
1273 needed to occur, which would usually be useless overhead.
1275 - SQLAlchemy's rich SELECT query expression system has no way to wrap
1276 multiple columns in a single expression object (the ORM does, but
1277 we are not using the ORM). So we would have to wrap _much_ more of
1278 that code in our own interfaces to encapsulate timespan
1279 representations there.
1280 """
1281 return TimespanDatabaseRepresentation.Compound
1283 def sync(
1284 self,
1285 table: sqlalchemy.schema.Table,
1286 *,
1287 keys: dict[str, Any],
1288 compared: dict[str, Any] | None = None,
1289 extra: dict[str, Any] | None = None,
1290 returning: Sequence[str] | None = None,
1291 update: bool = False,
1292 ) -> tuple[dict[str, Any] | None, bool | dict[str, Any]]:
1293 """Insert into a table as necessary to ensure database contains
1294 values equivalent to the given ones.
1296 Parameters
1297 ----------
1298 table : `sqlalchemy.schema.Table`
1299 Table to be queried and possibly inserted into.
1300 keys : `dict`
1301 Column name-value pairs used to search for an existing row; must
1302 be a combination that can be used to select a single row if one
1303 exists. If such a row does not exist, these values are used in
1304 the insert.
1305 compared : `dict`, optional
1306 Column name-value pairs that are compared to those in any existing
1307 row. If such a row does not exist, these rows are used in the
1308 insert.
1309 extra : `dict`, optional
1310 Column name-value pairs that are ignored if a matching row exists,
1311 but used in an insert if one is necessary.
1312 returning : `~collections.abc.Sequence` of `str`, optional
1313 The names of columns whose values should be returned.
1314 update : `bool`, optional
1315 If `True` (`False` is default), update the existing row with the
1316 values in ``compared`` instead of raising `DatabaseConflictError`.
1318 Returns
1319 -------
1320 row : `dict`, optional
1321 The value of the fields indicated by ``returning``, or `None` if
1322 ``returning`` is `None`.
1323 inserted_or_updated : `bool` or `dict`
1324 If `True`, a new row was inserted; if `False`, a matching row
1325 already existed. If a `dict` (only possible if ``update=True``),
1326 then an existing row was updated, and the dict maps the names of
1327 the updated columns to their *old* values (new values can be
1328 obtained from ``compared``).
1330 Raises
1331 ------
1332 DatabaseConflictError
1333 Raised if the values in ``compared`` do not match the values in the
1334 database.
1335 ReadOnlyDatabaseError
1336 Raised if `isWriteable` returns `False`, and no matching record
1337 already exists.
1339 Notes
1340 -----
1341 May be used inside transaction contexts, so implementations may not
1342 perform operations that interrupt transactions.
1344 It may be called on read-only databases if and only if the matching row
1345 does in fact already exist.
1346 """
1348 def check() -> tuple[int, dict[str, Any] | None, list | None]:
1349 """Query for a row that matches the ``key`` argument, and compare
1350 to what was given by the caller.
1352 Returns
1353 -------
1354 n : `int`
1355 Number of matching rows. ``n != 1`` is always an error, but
1356 it's a different kind of error depending on where `check` is
1357 being called.
1358 bad : `dict` or `None`
1359 The subset of the keys of ``compared`` for which the existing
1360 values did not match the given one, mapped to the existing
1361 values in the database. Once again, ``not bad`` is always an
1362 error, but a different kind on context. `None` if ``n != 1``.
1363 result : `list` or `None`
1364 Results in the database that correspond to the columns given
1365 in ``returning``, or `None` if ``returning is None``.
1366 """
1367 toSelect: set[str] = set()
1368 if compared is not None:
1369 toSelect.update(compared.keys())
1370 if returning is not None:
1371 toSelect.update(returning)
1372 if not toSelect:
1373 # Need to select some column, even if we just want to see
1374 # how many rows we get back.
1375 toSelect.add(next(iter(keys.keys())))
1376 selectSql = (
1377 sqlalchemy.sql.select(*[table.columns[k].label(k) for k in toSelect])
1378 .select_from(table)
1379 .where(sqlalchemy.sql.and_(*[table.columns[k] == v for k, v in keys.items()]))
1380 )
1381 with self._transaction() as (_, connection):
1382 fetched = list(connection.execute(selectSql).mappings())
1383 if len(fetched) != 1:
1384 return len(fetched), None, None
1385 existing = fetched[0]
1386 if compared is not None:
1388 def safeNotEqual(a: Any, b: Any) -> bool:
1389 if isinstance(a, astropy.time.Time):
1390 return not time_utils.TimeConverter().times_equal(a, b)
1391 return a != b
1393 inconsistencies = {
1394 k: existing[k] for k, v in compared.items() if safeNotEqual(existing[k], v)
1395 }
1396 else:
1397 inconsistencies = {}
1398 if returning is not None:
1399 toReturn: list | None = [existing[k] for k in returning]
1400 else:
1401 toReturn = None
1402 return 1, inconsistencies, toReturn
1404 def _format_bad(inconsistencies: dict[str, Any]) -> str:
1405 """Format the 'bad' dictionary of existing values returned by
1406 ``check`` into a string suitable for an error message.
1407 """
1408 assert compared is not None, "Should not be able to get inconsistencies without comparing."
1409 return ", ".join(f"{k}: {v!r} != {compared[k]!r}" for k, v in inconsistencies.items())
1411 if self.isTableWriteable(table):
1412 # Try an insert first, but allow it to fail (in only specific
1413 # ways).
1414 row = keys.copy()
1415 if compared is not None:
1416 row.update(compared)
1417 if extra is not None:
1418 row.update(extra)
1419 with self.transaction():
1420 inserted = bool(self.ensure(table, row))
1421 inserted_or_updated: bool | dict[str, Any]
1422 # Need to perform check() for this branch inside the
1423 # transaction, so we roll back an insert that didn't do
1424 # what we expected. That limits the extent to which we
1425 # can reduce duplication between this block and the other
1426 # ones that perform similar logic.
1427 n, bad, result = check()
1428 if n < 1:
1429 raise ConflictingDefinitionError(
1430 f"Attempted to ensure {row} exists by inserting it with ON CONFLICT IGNORE, "
1431 f"but a post-insert query on {keys} returned no results. "
1432 f"Insert was {'' if inserted else 'not '}reported as successful. "
1433 "This can occur if the insert violated a database constraint other than the "
1434 "unique constraint or primary key used to identify the row in this call."
1435 )
1436 elif n > 1:
1437 raise RuntimeError(
1438 f"Keys passed to sync {keys.keys()} do not comprise a "
1439 f"unique constraint for table {table.name}."
1440 )
1441 elif bad:
1442 assert compared is not None, (
1443 "Should not be able to get inconsistencies without comparing."
1444 )
1445 if inserted:
1446 raise RuntimeError(
1447 f"Conflict ({bad}) in sync after successful insert; this is "
1448 "possible if the same table is being updated by a concurrent "
1449 "process that isn't using sync, but it may also be a bug in "
1450 "daf_butler."
1451 )
1452 elif update:
1453 with self._transaction() as (_, connection):
1454 connection.execute(
1455 table.update()
1456 .where(sqlalchemy.sql.and_(*[table.columns[k] == v for k, v in keys.items()]))
1457 .values(**{k: compared[k] for k in bad})
1458 )
1459 inserted_or_updated = bad
1460 else:
1461 raise DatabaseConflictError(
1462 f"Conflict in sync for table {table.name} on column(s) {_format_bad(bad)}."
1463 )
1464 else:
1465 inserted_or_updated = inserted
1466 else:
1467 # Database is not writeable; just see if the row exists.
1468 n, bad, result = check()
1469 if n < 1:
1470 raise ReadOnlyDatabaseError("sync needs to insert, but database is read-only.")
1471 elif n > 1:
1472 raise RuntimeError("Keys passed to sync do not comprise a unique constraint.")
1473 elif bad:
1474 if update:
1475 raise ReadOnlyDatabaseError("sync needs to update, but database is read-only.")
1476 else:
1477 raise DatabaseConflictError(
1478 f"Conflict in sync for table {table.name} on column(s) {_format_bad(bad)}."
1479 )
1480 inserted_or_updated = False
1481 if returning is None:
1482 return None, inserted_or_updated
1483 else:
1484 assert result is not None
1485 return dict(zip(returning, result, strict=True)), inserted_or_updated
1487 def insert(
1488 self,
1489 table: sqlalchemy.schema.Table,
1490 *rows: dict,
1491 returnIds: bool = False,
1492 select: sqlalchemy.sql.expression.SelectBase | None = None,
1493 names: Iterable[str] | None = None,
1494 ) -> list[int] | None:
1495 """Insert one or more rows into a table, optionally returning
1496 autoincrement primary key values.
1498 Parameters
1499 ----------
1500 table : `sqlalchemy.schema.Table`
1501 Table rows should be inserted into.
1502 *rows : `dict`
1503 Positional arguments are the rows to be inserted, as dictionaries
1504 mapping column name to value. The keys in all dictionaries must
1505 be the same.
1506 returnIds : `bool`, optional
1507 If `True` (`False` is default), return the values of the table's
1508 autoincrement primary key field (which much exist).
1509 select : `sqlalchemy.sql.SelectBase`, optional
1510 A SELECT query expression to insert rows from. Cannot be provided
1511 with either ``rows`` or ``returnIds=True``.
1512 names : `~collections.abc.Iterable` [ `str` ], optional
1513 Names of columns in ``table`` to be populated, ordered to match the
1514 columns returned by ``select``. Ignored if ``select`` is `None`.
1515 If not provided, the columns returned by ``select`` must be named
1516 to match the desired columns of ``table``.
1518 Returns
1519 -------
1520 ids : `None`, or `list` of `int`
1521 If ``returnIds`` is `True`, a `list` containing the inserted
1522 values for the table's autoincrement primary key.
1524 Raises
1525 ------
1526 ReadOnlyDatabaseError
1527 Raised if `isWriteable` returns `False` when this method is called.
1529 Notes
1530 -----
1531 The default implementation uses bulk insert syntax when ``returnIds``
1532 is `False`, and a loop over single-row insert operations when it is
1533 `True`.
1535 Derived classes should reimplement when they can provide a more
1536 efficient implementation (especially for the latter case).
1538 May be used inside transaction contexts, so implementations may not
1539 perform operations that interrupt transactions.
1540 """
1541 self.assertTableWriteable(table, f"Cannot insert into read-only table {table}.")
1542 if select is not None and (rows or returnIds):
1543 raise TypeError("'select' is incompatible with passing value rows or returnIds=True.")
1544 if not rows and select is None:
1545 if returnIds:
1546 return []
1547 else:
1548 return None
1549 with self._transaction() as (_, connection):
1550 if not returnIds:
1551 if select is not None:
1552 if names is None:
1553 # columns() is deprecated since 1.4, but
1554 # selected_columns() method did not exist in 1.3.
1555 if hasattr(select, "selected_columns"):
1556 names = select.selected_columns.keys()
1557 else:
1558 names = select.columns.keys()
1559 connection.execute(table.insert().from_select(list(names), select))
1560 else:
1561 connection.execute(table.insert(), rows)
1562 return None
1563 else:
1564 sql = table.insert()
1565 ids = []
1566 for row in rows:
1567 key = connection.execute(sql, row).inserted_primary_key
1568 assert key is not None
1569 ids.append(key[0])
1570 return ids
1572 @abstractmethod
1573 def replace(self, table: sqlalchemy.schema.Table, *rows: dict) -> None:
1574 """Insert one or more rows into a table, replacing any existing rows
1575 for which insertion of a new row would violate the primary key
1576 constraint.
1578 Parameters
1579 ----------
1580 table : `sqlalchemy.schema.Table`
1581 Table rows should be inserted into.
1582 *rows
1583 Positional arguments are the rows to be inserted, as dictionaries
1584 mapping column name to value. The keys in all dictionaries must
1585 be the same.
1587 Raises
1588 ------
1589 ReadOnlyDatabaseError
1590 Raised if `isWriteable` returns `False` when this method is called.
1592 Notes
1593 -----
1594 May be used inside transaction contexts, so implementations may not
1595 perform operations that interrupt transactions.
1597 Implementations should raise a `sqlalchemy.exc.IntegrityError`
1598 exception when a constraint other than the primary key would be
1599 violated.
1601 Implementations are not required to support `replace` on tables
1602 with autoincrement keys.
1603 """
1604 raise NotImplementedError()
1606 @abstractmethod
1607 def ensure(self, table: sqlalchemy.schema.Table, *rows: dict, primary_key_only: bool = False) -> int:
1608 """Insert one or more rows into a table, skipping any rows for which
1609 insertion would violate a unique constraint.
1611 Parameters
1612 ----------
1613 table : `sqlalchemy.schema.Table`
1614 Table rows should be inserted into.
1615 *rows
1616 Positional arguments are the rows to be inserted, as dictionaries
1617 mapping column name to value. The keys in all dictionaries must
1618 be the same.
1619 primary_key_only : `bool`, optional
1620 If `True` (`False` is default), only skip rows that violate the
1621 primary key constraint, and raise an exception (and rollback
1622 transactions) for other constraint violations.
1624 Returns
1625 -------
1626 count : `int`
1627 The number of rows actually inserted.
1629 Raises
1630 ------
1631 ReadOnlyDatabaseError
1632 Raised if `isWriteable` returns `False` when this method is called.
1633 This is raised even if the operation would do nothing even on a
1634 writeable database.
1636 Notes
1637 -----
1638 May be used inside transaction contexts, so implementations may not
1639 perform operations that interrupt transactions.
1641 Implementations are not required to support `ensure` on tables
1642 with autoincrement keys.
1643 """
1644 raise NotImplementedError()
1646 def delete(self, table: sqlalchemy.schema.Table, columns: Iterable[str], *rows: dict) -> int:
1647 """Delete one or more rows from a table.
1649 Parameters
1650 ----------
1651 table : `sqlalchemy.schema.Table`
1652 Table that rows should be deleted from.
1653 columns : `~collections.abc.Iterable` of `str`
1654 The names of columns that will be used to constrain the rows to
1655 be deleted; these will be combined via ``AND`` to form the
1656 ``WHERE`` clause of the delete query.
1657 *rows
1658 Positional arguments are the keys of rows to be deleted, as
1659 dictionaries mapping column name to value. The keys in all
1660 dictionaries must be exactly the names in ``columns``.
1662 Returns
1663 -------
1664 count : `int`
1665 Number of rows deleted.
1667 Raises
1668 ------
1669 ReadOnlyDatabaseError
1670 Raised if `isWriteable` returns `False` when this method is called.
1672 Notes
1673 -----
1674 May be used inside transaction contexts, so implementations may not
1675 perform operations that interrupt transactions.
1677 The default implementation should be sufficient for most derived
1678 classes.
1679 """
1680 self.assertTableWriteable(table, f"Cannot delete from read-only table {table}.")
1681 if columns and not rows:
1682 # If there are no columns, this operation is supposed to delete
1683 # everything (so we proceed as usual). But if there are columns,
1684 # but no rows, it was a constrained bulk operation where the
1685 # constraint is that no rows match, and we should short-circuit
1686 # while reporting that no rows were affected.
1687 return 0
1688 sql = table.delete()
1689 columns = list(columns) # Force iterators to list
1691 # More efficient to use IN operator if there is only one
1692 # variable changing across all rows.
1693 content: dict[str, set] = defaultdict(set)
1694 if len(columns) == 1:
1695 # Nothing to calculate since we can always use IN
1696 column = columns[0]
1697 changing_columns = [column]
1698 content[column] = {row[column] for row in rows}
1699 else:
1700 for row in rows:
1701 for k, v in row.items():
1702 content[k].add(v)
1703 changing_columns = [col for col, values in content.items() if len(values) > 1]
1705 if len(changing_columns) != 1:
1706 # More than one column changes each time so do explicit bind
1707 # parameters and have each row processed separately.
1708 whereTerms = [table.columns[name] == sqlalchemy.sql.bindparam(name) for name in columns]
1709 if whereTerms:
1710 sql = sql.where(sqlalchemy.sql.and_(*whereTerms))
1711 with self._transaction() as (_, connection):
1712 return connection.execute(sql, rows).rowcount
1713 else:
1714 # One of the columns has changing values but any others are
1715 # fixed. In this case we can use an IN operator and be more
1716 # efficient.
1717 name = changing_columns.pop()
1719 # Simple where clause for the unchanging columns
1720 clauses = []
1721 for k, v in content.items():
1722 if k == name:
1723 continue
1724 # The set only has one element
1725 clauses.append(table.columns[k] == v.pop())
1727 # The IN operator will not work for "infinite" numbers of
1728 # rows so must batch it up into distinct calls.
1729 in_content = list(content[name])
1730 n_elements = len(in_content)
1732 rowcount = 0
1733 iposn = 0
1734 n_per_loop = 1_000 # Controls how many items to put in IN clause
1735 with self._transaction() as (_, connection):
1736 for iposn in range(0, n_elements, n_per_loop):
1737 endpos = iposn + n_per_loop
1738 in_clause = table.columns[name].in_(in_content[iposn:endpos])
1740 newsql = sql.where(sqlalchemy.sql.and_(*clauses, in_clause))
1741 rowcount += connection.execute(newsql).rowcount
1742 return rowcount
1744 def deleteWhere(self, table: sqlalchemy.schema.Table, where: sqlalchemy.sql.ColumnElement) -> int:
1745 """Delete rows from a table with pre-constructed WHERE clause.
1747 Parameters
1748 ----------
1749 table : `sqlalchemy.schema.Table`
1750 Table that rows should be deleted from.
1751 where : `sqlalchemy.sql.ClauseElement`
1752 The names of columns that will be used to constrain the rows to
1753 be deleted; these will be combined via ``AND`` to form the
1754 ``WHERE`` clause of the delete query.
1756 Returns
1757 -------
1758 count : `int`
1759 Number of rows deleted.
1761 Raises
1762 ------
1763 ReadOnlyDatabaseError
1764 Raised if `isWriteable` returns `False` when this method is called.
1766 Notes
1767 -----
1768 May be used inside transaction contexts, so implementations may not
1769 perform operations that interrupt transactions.
1771 The default implementation should be sufficient for most derived
1772 classes.
1773 """
1774 self.assertTableWriteable(table, f"Cannot delete from read-only table {table}.")
1776 sql = table.delete().where(where)
1777 with self._transaction() as (_, connection):
1778 return connection.execute(sql).rowcount
1780 def update(self, table: sqlalchemy.schema.Table, where: dict[str, str], *rows: dict) -> int:
1781 """Update one or more rows in a table.
1783 Parameters
1784 ----------
1785 table : `sqlalchemy.schema.Table`
1786 Table containing the rows to be updated.
1787 where : `dict` [`str`, `str`]
1788 A mapping from the names of columns that will be used to search for
1789 existing rows to the keys that will hold these values in the
1790 ``rows`` dictionaries. Note that these may not be the same due to
1791 SQLAlchemy limitations.
1792 *rows
1793 Positional arguments are the rows to be updated. The keys in all
1794 dictionaries must be the same, and may correspond to either a
1795 value in the ``where`` dictionary or the name of a column to be
1796 updated.
1798 Returns
1799 -------
1800 count : `int`
1801 Number of rows matched (regardless of whether the update actually
1802 modified them).
1804 Raises
1805 ------
1806 ReadOnlyDatabaseError
1807 Raised if `isWriteable` returns `False` when this method is called.
1809 Notes
1810 -----
1811 May be used inside transaction contexts, so implementations may not
1812 perform operations that interrupt transactions.
1814 The default implementation should be sufficient for most derived
1815 classes.
1816 """
1817 self.assertTableWriteable(table, f"Cannot update read-only table {table}.")
1818 if not rows:
1819 return 0
1820 sql = table.update().where(
1821 sqlalchemy.sql.and_(*[table.columns[k] == sqlalchemy.sql.bindparam(v) for k, v in where.items()])
1822 )
1823 with self._transaction() as (_, connection):
1824 return connection.execute(sql, rows).rowcount
1826 @contextmanager
1827 def query(
1828 self,
1829 sql: sqlalchemy.sql.expression.Executable | sqlalchemy.sql.expression.SelectBase,
1830 *args: Any,
1831 **kwargs: Any,
1832 ) -> Iterator[sqlalchemy.engine.CursorResult]:
1833 """Run a SELECT query against the database.
1835 Parameters
1836 ----------
1837 sql : `sqlalchemy.sql.expression.SelectBase`
1838 A SQLAlchemy representation of a ``SELECT`` query.
1839 *args
1840 Additional positional arguments are forwarded to
1841 `sqlalchemy.engine.Connection.execute`.
1842 **kwargs
1843 Additional keyword arguments are forwarded to
1844 `sqlalchemy.engine.Connection.execute`.
1846 Returns
1847 -------
1848 result_context : `sqlalchemy.engine.CursorResults`
1849 Context manager that returns the query result object when entered.
1850 These results are invalidated when the context is exited.
1851 """
1852 if self._session_connection is None:
1853 connection = self._engine.connect()
1854 else:
1855 connection = self._session_connection
1856 try:
1857 self._log_query_and_plan(connection, sql)
1858 # TODO: SelectBase is not good for execute(), but it used
1859 # everywhere, e.g. in daf_relation. We should switch to Executable
1860 # at some point.
1861 with connection.execute(
1862 cast(sqlalchemy.sql.expression.Executable, sql), *args, **kwargs
1863 ) as result:
1864 yield result
1865 finally:
1866 if connection is not self._session_connection:
1867 connection.close()
1869 def _log_query_and_plan(
1870 self,
1871 connection: sqlalchemy.Connection,
1872 sql: sqlalchemy.sql.expression.Executable | sqlalchemy.sql.expression.SelectBase,
1873 ) -> None:
1874 """Log the given SQL statement and the DB's plan for executing it if
1875 the environment variable DAF_BUTLER_DEBUG_QUERIES is set to a truthy
1876 value.
1877 """
1878 if os.environ.get("DAF_BUTLER_DEBUG_QUERIES", False):
1879 assert isinstance(sql, sqlalchemy.SelectBase)
1880 compiled = sql.compile(connection)
1881 print(
1882 f"Executing SQL statement:\n{compiled}\nBind parameters: {compiled.params}", file=sys.stderr
1883 )
1885 query_plan = get_query_plan(connection, sql)
1886 print(f"Query plan:\n{query_plan}", file=sys.stderr)
1888 @abstractmethod
1889 def constant_rows(
1890 self,
1891 fields: NamedValueAbstractSet[ddl.FieldSpec],
1892 *rows: dict,
1893 name: str | None = None,
1894 ) -> sqlalchemy.sql.FromClause:
1895 """Return a SQLAlchemy object that represents a small number of
1896 constant-valued rows.
1898 Parameters
1899 ----------
1900 fields : `NamedValueAbstractSet` [ `ddl.FieldSpec` ]
1901 The columns of the rows. Unique and foreign key constraints are
1902 ignored.
1903 *rows : `dict`
1904 Values for the rows.
1905 name : `str`, optional
1906 If provided, the name of the SQL construct. If not provided, an
1907 opaque but unique identifier is generated.
1909 Returns
1910 -------
1911 from_clause : `sqlalchemy.sql.FromClause`
1912 SQLAlchemy object representing the given rows. This is guaranteed
1913 to be something that can be directly joined into a ``SELECT``
1914 query's ``FROM`` clause, and will not involve a temporary table
1915 that needs to be cleaned up later.
1917 Notes
1918 -----
1919 The default implementation uses the SQL-standard ``VALUES`` construct,
1920 but support for that construct is varied enough across popular RDBMSs
1921 that the method is still marked abstract to force explicit opt-in via
1922 delegation to `super`.
1923 """
1924 if name is None:
1925 name = f"tmp_{uuid.uuid4().hex}"
1926 return sqlalchemy.sql.values(
1927 *[sqlalchemy.Column(field.name, field.getSizedColumnType()) for field in fields],
1928 name=name,
1929 ).data([tuple(row[name] for name in fields.names) for row in rows])
1931 def get_constant_rows_max(self) -> int:
1932 """Return the maximum number of rows that should be passed to
1933 `constant_rows` for this backend.
1935 Returns
1936 -------
1937 max : `int`
1938 Maximum number of rows.
1940 Notes
1941 -----
1942 This should reflect typical performance profiles (or a guess at these),
1943 not just hard database engine limits.
1944 """
1945 return 1000
1947 @abstractmethod
1948 def glob_expression(
1949 self, expression: sqlalchemy.ColumnElement[Any], pattern: str
1950 ) -> sqlalchemy.ColumnElement[bool]:
1951 """Produce boolean expression for expression match against glob-like
1952 pattern.
1954 Parameters
1955 ----------
1956 expression : `sqlalchemy.ColumnElement`
1957 Column expression that evaluates to string.
1958 pattern : `str`
1959 GLob pattern string.
1961 Returns
1962 -------
1963 glob : `sqlalchemy.ColumnElement`
1964 Boolean expression that matches ``expression`` against ``pattern``.
1965 """
1966 raise NotImplementedError()
1968 @property
1969 @abstractmethod
1970 def has_distinct_on(self) -> bool:
1971 """Whether this database supports the ``DISTINCT ON`` SQL construct."""
1972 raise NotImplementedError()
1974 @property
1975 @abstractmethod
1976 def has_any_aggregate(self) -> bool:
1977 """Whether this database supports the ``ANY_VALUE`` aggregate function
1978 or something equivalent.
1979 """
1980 raise NotImplementedError()
1982 @property
1983 def supports_temporary_tables(self) -> bool:
1984 return self._allow_temporary_tables
1986 @abstractmethod
1987 def apply_any_aggregate(self, column: sqlalchemy.ColumnElement[Any]) -> sqlalchemy.ColumnElement[Any]:
1988 """Wrap the given SQLAlchemy column in the ``ANY_VALUE`` aggregate
1989 function or its equivalent.
1991 Parameters
1992 ----------
1993 column : `sqlalchemy.ColumnElement`
1994 Original column to wrap.
1996 Returns
1997 -------
1998 wrapped : `sqlalchemy.ColumnElement`
1999 A column element of the same SQL type that can appear in the
2000 ``SELECT`` clause even when this column does not appear in the
2001 ``GROUP BY`` clause.
2003 Notes
2004 -----
2005 This method's behavior is unspecified when `has_any_aggregate` is
2006 `False`; the caller is responsible for checking that property first.
2007 """
2008 raise NotImplementedError()
2010 origin: int
2011 """An integer ID that should be used as the default for any datasets,
2012 quanta, or other entities that use a (autoincrement, origin) compound
2013 primary key (`int`).
2014 """
2016 namespace: str | None
2017 """The schema or namespace this database instance is associated with
2018 (`str` or `None`).
2019 """
2021 name_shrinker: NameShrinker
2022 """An object that can be used to shrink field names to fit within the
2023 identifier limit of the database engine (`NameShrinker`).
2024 """
2027class DatabaseMetadata:
2028 """Wrapper around SqlAlchemy MetaData object to ensure threadsafety.
2030 Parameters
2031 ----------
2032 namespace : `str` or `None`
2033 Name of the schema or namespace this instance is associated with.
2035 Notes
2036 -----
2037 `sqlalchemy.MetaData` is documented to be threadsafe for reads, but not
2038 with concurrent modifications. We add tables dynamically at runtime,
2039 and the MetaData object is shared by all Database instances sharing
2040 the same connection pool.
2041 """
2043 def __init__(self, namespace: str | None) -> None:
2044 self._lock = Lock()
2045 self._metadata = sqlalchemy.MetaData(schema=namespace)
2046 self._tables: dict[str, sqlalchemy.Table] = {}
2048 def add_table(
2049 self, db: Database, name: str, spec: ddl.TableSpec, **kwargs: Any
2050 ) -> sqlalchemy.schema.Table:
2051 """Add a new table to the MetaData object, returning its sqlalchemy
2052 representation. This does not physically create the table in the
2053 database -- it only sets up its definition.
2055 Parameters
2056 ----------
2057 db : `Database`
2058 Database connection associated with the table definition.
2059 name : `str`
2060 The name of the table.
2061 spec : `ddl.TableSpec`
2062 The specification of the table.
2063 **kwargs
2064 Additional keyword arguments to forward to the
2065 `sqlalchemy.schema.Table` constructor.
2067 Returns
2068 -------
2069 table : `sqlalchemy.schema.Table`
2070 The created table.
2071 """
2072 with self._lock:
2073 if (table := self._tables.get(name)) is not None:
2074 return table
2076 table = db._convertTableSpec(name, spec, self._metadata, **kwargs)
2077 for foreignKeySpec in spec.foreignKeys:
2078 table.append_constraint(db._convertForeignKeySpec(name, foreignKeySpec, self._metadata))
2080 self._tables[name] = table
2081 return table
2083 def get_table(self, name: str) -> sqlalchemy.schema.Table | None:
2084 """Return the definition of a table that was previously added to this
2085 MetaData object.
2087 Parameters
2088 ----------
2089 name : `str`
2090 Name of the table.
2092 Returns
2093 -------
2094 table : `sqlalchemy.schema.Table` or `None`
2095 The table definition, or `None` if the table is not known to this
2096 MetaData instance.
2097 """
2098 with self._lock:
2099 return self._tables.get(name)
2101 def remove_table(self, name: str) -> None:
2102 """Remove a table that was previously added to this MetaData object.
2104 Parameters
2105 ----------
2106 name : `str`
2107 Name of the table.
2108 """
2109 with self._lock:
2110 table = self._tables.pop(name, None)
2111 if table is not None:
2112 self._metadata.remove(table)
2114 def create_all(self, connection: sqlalchemy.engine.Connection) -> None:
2115 """Create all tables known to this MetaData object in the database.
2116 Same as `sqlalchemy.MetaData.create_all`.
2118 Parameters
2119 ----------
2120 connection : `sqlalchemy.engine.connection`
2121 Database connection that will be used to create tables.
2122 """
2123 with self._lock:
2124 self._metadata.create_all(connection)