def load_database(self, database):
"""
Load the given database, whatever it might be.
A connection string: ``sqlite:///database.sqlite``
A dictionary: ``{'engine': 'SqliteDatabase', 'name': 'database.sqlite'}``
A peewee.Database instance: ``peewee.SqliteDatabase('database.sqlite')``
:param database: Connection string, dict, or peewee.Database instance to use.
:raises: peewee.DatabaseError if database connection cannot be established.
:return: Database connection.
:rtype: peewee.Database instance.
"""
# It could be an actual instance...
if isinstance(database, (peewee.Proxy, peewee.Database)):
return database
# It could be a dictionary...
if isinstance(database, dict):
try:
name = database.pop('name')
engine = database.pop('engine')
except KeyError:
error_msg = 'Configuration dict must specify "name" and "engine" keys.'
raise peewee.DatabaseError(error_msg)
db_class = pydoc.locate(engine)
if not db_class:
raise peewee.DatabaseError(
'Unable to import engine class: {}'.format(engine))
return db_class(name, **database)
# Or it could be a database URL.
return url_connect(database)
def reflect(self):
"""Adds fields and indexes to the model using reflection."""
meta = self._meta
database = meta.database
table_name = meta.table_name
schema = meta.schema
if not database or not database.connected:
return
# Lists tables in the schema. This is a bit of a hack but
# faster than using database.table_exists because it's cached.
database.get_fields(table_name, schema) # Force caching of _metadata.
schema_tables = database._metadata[schema].keys()
if table_name not in schema_tables:
return
for index in meta.indexes:
if hasattr(index, 'reflected') and index.reflected:
meta.indexes.remove(index)
if not database.is_connection_usable():
raise peewee.DatabaseError('database not connected.')
if hasattr(meta, 'reflection_options'):
opts = meta.reflection_options
skip_fks = opts.get('skip_foreign_keys', False)
use_peewee_reflection = opts.get('use_peewee_reflection', True)
else:
skip_fks = False
use_peewee_reflection = True
try:
if use_peewee_reflection:
# Check for locks. We only need to do this if using the Peewee
# reflection because one of the queries it does can be blocked
# by a AccessExclusiveLock lock.
locks = is_table_locked(database, table_name)
if locks and 'AccessExclusiveLock' in locks:
warnings.warn(
f'table {schema}.{table_name} is locked and '
'will not be reflected.', SdssdbUserWarning)
return
introspector = database.get_introspector(schema)
reflected_model = introspector.generate_models(
table_names=[table_name])[table_name]
fields = reflected_model._meta.fields
else:
fields = database.get_fields(table_name, schema)
fields = {field.column_name: field for field in fields}
except KeyError as ee:
warnings.warn(
f'reflection failed for {table_name}: '
f'table or column {ee} not found.', SdssdbUserWarning)
return
except Exception as ee:
warnings.warn(f'reflection failed for {table_name}: {ee}',
SdssdbUserWarning)
return
for field_name, field in fields.items():
if field_name in keyword.kwlist:
field_name += '_'
if field_name in meta.fields:
meta_field = meta.fields[field_name]
if not getattr(meta_field, 'reflected', False):
continue
if isinstance(field, peewee.ForeignKeyField) and skip_fks:
continue
if field.primary_key:
meta.set_primary_key(field_name, field)
else:
meta.add_field(field_name, field)
meta.fields[field_name].reflected = True
# Composite keys are not a normal column so if the pk has not been
# set already, check if it exists in the reflected model. We avoid
# adding pks that are foreign keys.
if not meta.primary_key:
if use_peewee_reflection and reflected_model._meta.primary_key:
pk = reflected_model._meta.primary_key
if not isinstance(pk, peewee.ForeignKeyField) or not skip_fks:
meta.set_primary_key(pk.name, pk)
elif not use_peewee_reflection:
pk = database.get_primary_keys(table_name, schema)
if len(pk) > 1:
pk = peewee.CompositeKey(*pk)
meta.set_primary_key('__composite_key__', pk)