Coverage for python/lsst/daf/butler/registry/databases/sqlite.py : 17%

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
# This file is part of daf_butler. # # Developed for the LSST Data Management System. # This product includes software developed by the LSST Project # (http://www.lsst.org). # See the COPYRIGHT file at the top-level directory of this distribution # for details of code ownership. # # This program is free software: you can redistribute it and/or modify # it under the terms of the GNU General Public License as published by # the Free Software Foundation, either version 3 of the License, or # (at your option) any later version. # # This program is distributed in the hope that it will be useful, # but WITHOUT ANY WARRANTY; without even the implied warranty of # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the # GNU General Public License for more details. # # You should have received a copy of the GNU General Public License # along with this program. If not, see <http://www.gnu.org/licenses/>.
assert isinstance(dbapiConnection, sqlite3.Connection) # Prevent pysqlite from emitting BEGIN and COMMIT statements. dbapiConnection.isolation_level = None # Enable foreign keys with closing(dbapiConnection.cursor()) as cursor: cursor.execute("PRAGMA foreign_keys=ON;") cursor.execute("PRAGMA busy_timeout = 300000;") # in ms, so 5min (way longer than should be needed)
assert connection.dialect.name == "sqlite" # Replace pysqlite's buggy transaction handling that never BEGINs with our # own that does, and tell SQLite to try to acquire a lock as soon as we # start a transaction (this should lead to more blocking and fewer # deadlocks). connection.execute("BEGIN IMMEDIATE") return connection
"""A SQLAlchemy query that compiles to INSERT ... ON CONFLICT REPLACE on the primary key constraint for the table. """
def _replace(insert, compiler, **kw): """Generate an INSERT ... ON CONFLICT REPLACE query. """ # SQLite and PostgreSQL use similar syntax for their ON CONFLICT extension, # but SQLAlchemy only knows about PostgreSQL's, so we have to compile some # custom text SQL ourselves. result = compiler.visit_insert(insert, **kw) preparer = compiler.preparer pk_columns = ", ".join([preparer.format_column(col) for col in insert.table.primary_key]) result += f" ON CONFLICT ({pk_columns})" columns = [preparer.format_column(col) for col in insert.table.columns if col.name not in insert.table.primary_key] updates = ", ".join([f"{col} = excluded.{col}" for col in columns]) result += f" DO UPDATE SET {updates}" return result
fields=[ddl.FieldSpec(name="id", dtype=sqlalchemy.Integer, primaryKey=True)] )
class _AutoincrementCompoundKeyWorkaround: """A workaround for SQLite's lack of support for compound primary keys that include an autoincrement field. """
"""A single-column internal table that can be inserted into to yield autoincrement values (`sqlalchemy.schema.Table`). """
with values from the internal table (`str`). """
"""An implementation of the `Database` interface for SQLite3.
Parameters ---------- connection : `sqlalchemy.engine.Connection` An existing connection created by a previous call to `connect`. origin : `int` An integer ID that should be used as the default for any datasets, quanta, or other entities that use a (autoincrement, origin) compound primary key. namespace : `str`, optional The namespace (schema) this database is associated with. If `None`, the default schema for the connection is used (which may be `None`). writeable : `bool`, optional If `True`, allow write operations on the database, including ``CREATE TABLE``.
Notes ----- The case where ``namespace is not None`` is not yet tested, and may be broken; we need an API for attaching to different databases in order to write those tests, but haven't yet worked out what is common/different across databases well enough to define it. """
namespace: Optional[str] = None, writeable: bool = True): super().__init__(origin=origin, connection=connection, namespace=namespace) # Get the filename from a call to 'PRAGMA database_list'. with closing(connection.connection.cursor()) as cursor: dbList = list(cursor.execute("PRAGMA database_list").fetchall()) if len(dbList) == 0: raise RuntimeError("No database in connection.") if namespace is None: namespace = "main" for _, dbname, filename in dbList: if dbname == namespace: break else: raise RuntimeError(f"No '{namespace}' database in connection.") if not filename: self.filename = None else: self.filename = filename self._writeable = writeable self._autoincr = {}
writeable: bool = True) -> sqlalchemy.engine.Connection: """Create a `sqlalchemy.engine.Connection` from a SQLAlchemy URI or filename.
Parameters ---------- uri : `str` A SQLAlchemy URI connection string. filename : `str` Name of the SQLite database file, or `None` to use an in-memory database. Ignored if ``uri is not None``. origin : `int` An integer ID that should be used as the default for any datasets, quanta, or other entities that use a (autoincrement, origin) compound primary key. writeable : `bool`, optional If `True`, allow write operations on the database, including ``CREATE TABLE``.
Returns ------- cs : `sqlalchemy.engine.Connection` A database connection and transaction state. """ # In order to be able to tell SQLite that we want a read-only or # read-write connection, we need to make the SQLite DBAPI connection # with a "URI"-based connection string. SQLAlchemy claims it can do # this # (https://docs.sqlalchemy.org/en/13/dialects/sqlite.html#uri-connections), # but it doesn't seem to work as advertised. To work around this, we # use the 'creator' argument to sqlalchemy.engine.create_engine, which # lets us pass a callable that creates the DBAPI connection. if uri is None: if filename is None: target = ":memory:" uri = "sqlite://" else: target = f"file:{filename}" uri = f"sqlite:///{filename}" else: parsed = urllib.parse.urlparse(uri) queries = parsed.query.split("&") if "uri=true" in queries: # This is a SQLAlchemy URI that is already trying to make a # SQLite connection via a SQLite URI, and hence there may # be URI components for both SQLite and SQLAlchemy. We # don't need to support that, and it'd be a # reimplementation of all of the (broken) logic in # SQLAlchemy for doing this, so we just don't. raise NotImplementedError("SQLite connection strings with 'uri=true' are not supported.") # This is just a SQLAlchemy URI with a non-URI SQLite # connection string inside it. Pull that out so we can use it # in the creator call. if parsed.path.startswith("/"): filename = parsed.path[1:] target = f"file:{filename}" else: filename = None target = ":memory:" if filename is None: if not writeable: raise NotImplementedError("Read-only :memory: databases are not supported.") else: if writeable: target += '?mode=rwc&uri=true' else: target += '?mode=ro&uri=true'
def creator(): return sqlite3.connect(target, check_same_thread=False, uri=True)
engine = sqlalchemy.engine.create_engine(uri, poolclass=sqlalchemy.pool.NullPool, creator=creator)
sqlalchemy.event.listen(engine, "connect", _onSqlite3Connect) sqlalchemy.event.listen(engine, "begin", _onSqlite3Begin) return engine.connect()
namespace: Optional[str] = None, writeable: bool = True) -> Database: return cls(connection=connection, origin=origin, writeable=writeable, namespace=namespace)
return self._writeable
return f"SQLite3@{self.filename}"
**kwds) -> sqlalchemy.schema.Column: if spec.autoincrement: if not spec.primaryKey: raise RuntimeError(f"Autoincrement field {table}.{spec.name} that is not a " f"primary key is not supported.") if spec.dtype != sqlalchemy.Integer: # SQLite's autoincrement is really limited; it only works if # the column type is exactly "INTEGER". But it also doesn't # care about the distinctions between different integer types, # so it's safe to change it. spec = copy.copy(spec) spec.dtype = sqlalchemy.Integer return super()._convertFieldSpec(table, spec, metadata, **kwds)
**kwds) -> sqlalchemy.schema.Table: primaryKeyFieldNames = set(field.name for field in spec.fields if field.primaryKey) autoincrFieldNames = set(field.name for field in spec.fields if field.autoincrement) if len(autoincrFieldNames) > 1: raise RuntimeError("At most one autoincrement field per table is allowed.") if len(primaryKeyFieldNames) > 1 and len(autoincrFieldNames) > 0: # SQLite's default rowid-based autoincrement doesn't work if the # field is just one field in a compound primary key. As a # workaround, we create an extra table with just one column that # we'll insert into to generate those IDs. That's only safe if # that single-column table's records are already unique with just # the autoincrement field, not the rest of the primary key. In # practice, that means the single-column table's records are those # for which origin == self.origin. autoincrFieldName, = autoincrFieldNames otherPrimaryKeyFieldNames = primaryKeyFieldNames - autoincrFieldNames if otherPrimaryKeyFieldNames != {"origin"}: # We need the only other field in the key to be 'origin'. raise NotImplementedError( "Compound primary keys with an autoincrement are only supported in SQLite " "if the only non-autoincrement primary key field is 'origin'." ) self._autoincr[name] = _AutoincrementCompoundKeyWorkaround( table=self._convertTableSpec(f"_autoinc_{name}", _AUTOINCR_TABLE_SPEC, metadata, **kwds), column=autoincrFieldName ) return super()._convertTableSpec(name, spec, metadata, **kwds)
) -> Optional[List[int]]: autoincr = self._autoincr.get(table.name) if autoincr is not None: # This table has a compound primary key that includes an # autoincrement. That doesn't work natively in SQLite, so we # insert into a single-column table and use those IDs. if not rows: return [] if returnIds else None if autoincr.column in rows[0]: # Caller passed the autoincrement key values explicitly in the # first row. They had better have done the same for all rows, # or SQLAlchemy would have a problem, even if we didn't. assert all(autoincr.column in row for row in rows) # We need to insert only the values that correspond to # ``origin == self.origin`` into the single-column table, to # make sure we don't generate conflicting keys there later. rowsForAutoincrTable = [dict(id=row[autoincr.column]) for row in rows if row["origin"] == self.origin] # Insert into the autoincr table and the target table inside # a transaction. The main-table insertion can take care of # returnIds for us. with self.transaction(): self._connection.execute(autoincr.table.insert(), *rowsForAutoincrTable) return super().insert(table, *rows, returnIds=returnIds) else: # Caller did not pass autoincrement key values on the first # row. Make sure they didn't ever do that, and also make # sure the origin that was passed in is always self.origin, # because we can't safely generate autoincrement values # otherwise. assert all(autoincr.column not in row and row["origin"] == self.origin for row in rows) # Insert into the autoincr table one by one to get the # primary key values back, then insert into the target table # in the same transaction. with self.transaction(): newRows = [] ids = [] for row in rows: newRow = row.copy() id = self._connection.execute(autoincr.table.insert()).inserted_primary_key[0] newRow[autoincr.column] = id newRows.append(newRow) ids.append(id) # Don't ever ask to returnIds here, because we've already # got them. super().insert(table, *newRows) if returnIds: return ids else: return None else: return super().insert(table, *rows, returnIds=returnIds)
if not self.isWriteable(): raise ReadOnlyDatabaseError(f"Attempt to replace into read-only database '{self}'.") if table.name in self._autoincr: raise NotImplementedError( "replace does not support compound primary keys with autoincrement fields." ) self._connection.execute(_Replace(table), *rows)
Set to `None` for in-memory databases. """ |