def build_actions_chain(left_schema: Schema,
right_schema: Schema) -> Iterable[BaseAction]:
"""
Build full Action objects chain which suitable for such schema
change.
:param left_schema: current schema
:param right_schema: schema collected from mongoengine models
:return: iterable of Action objects
"""
action_chain = []
# Actions registry sorted by priority
registry = list(sorted(actions_registry.values(),
key=lambda x: x.priority))
left_schema = copy(left_schema)
document_types = get_all_document_types(left_schema, right_schema)
for action_cls in registry:
if issubclass(action_cls, BaseDocumentAction):
new_actions = list(
build_document_action_chain(action_cls, left_schema,
right_schema, document_types))
elif issubclass(action_cls, BaseFieldAction):
new_actions = list(
build_field_action_chain(action_cls, left_schema, right_schema,
document_types))
elif issubclass(action_cls, BaseIndexAction):
new_actions = list(
build_index_action_chain(action_cls, left_schema, right_schema,
document_types))
else:
continue
for action in new_actions:
log.debug('> %s', action)
try:
left_schema = patch(action.to_schema_patch(left_schema),
left_schema)
except (TypeError, ValueError, KeyError) as e:
raise ActionError(
f"Unable to apply schema patch of {action!r}. More likely that the "
f"schema is corrupted. You can use schema repair tools to fix this issue"
) from e
action_chain.extend(new_actions)
document_types = get_all_document_types(left_schema, right_schema)
if right_schema != left_schema:
log.error(
'Schema is still not reached the target state after applying all actions. '
'Changes left to make (diff): %s',
list(diff(left_schema, right_schema)))
raise ActionError(
'Could not reach target schema state after applying whole Action chain. '
'This could be a problem in some Action which does not process schema '
'properly or produces wrong schema diff. This is a programming error'
)
return action_chain
def build_document_action_chain(action_cls: Type[BaseDocumentAction],
left_schema: Schema,
right_schema: Schema,
document_types: Iterable[str]) -> Iterable[BaseAction]:
"""
Walk through schema changes, and produce chain of Action objects
of given type which could handle schema changes from left to right
:param action_cls: Action type to consider
:param left_schema:
:param right_schema:
:param document_types: list of document types to inspect
:return: iterable of suitable Action objects
"""
for document_type in document_types:
action_obj = action_cls.build_object(document_type, left_schema, right_schema)
if action_obj is not None:
try:
left_schema = patch(action_obj.to_schema_patch(left_schema), left_schema)
except (TypeError, ValueError, KeyError) as e:
raise ActionError(
f"Unable to apply schema patch of {action_obj!r}. More likely that the "
f"schema is corrupted. You can use schema repair tools to fix this issue"
) from e
yield action_obj
def makemigrations(self):
"""
Compare current mongoengine documents state and the last db
state and make a migration file if needed
"""
log.debug('Loading migration files...')
graph = self.build_graph()
log.debug('Loading schema from database...')
db_schema = self.load_db_schema()
# Obtain schema changes which migrations would make (including
# unapplied ones)
# If mongoengine models schema was changed regarding db schema
# then try to guess which actions would reflect such changes
for migration in graph.walk_down(graph.initial, unapplied_only=False):
for action_object in migration.get_actions():
try:
db_schema = patch(action_object.to_schema_patch(db_schema),
db_schema)
except (TypeError, ValueError, KeyError) as e:
raise ActionError(
f"Unable to apply schema patch of {action_object!r}. More likely that the "
f"schema is corrupted. You can use schema repair tools to fix this issue"
) from e
log.debug('Collecting schema from mongoengine documents...')
models_schema = collect_models_schema()
if db_schema == models_schema:
log.info('No changes detected')
return
log.debug('Building actions chain...')
actions_chain = build_actions_chain(db_schema, models_schema)
import_expressions = {'from mongoengine_migrate.actions import *'}
for action in actions_chain:
# If `regex` is set in action, then we probably need 're'
if isinstance(action.parameters.get('regex'), re.Pattern):
import_expressions.add('import re')
break
log.debug('Writing migrations file...')
env = Environment()
env.filters['symbol_wrap'] = symbol_wrap
tpl_ctx = {
'graph': graph,
'actions_chain': actions_chain,
'policy_enum': MigrationPolicy,
'import_expressions': import_expressions
}
tpl_path = Path(__file__).parent / 'migration_template.tpl'
tpl = env.from_string(tpl_path.read_text())
migration_source = tpl.render(tpl_ctx)
seq_number = str(len(graph.migrations)).zfill(4)
name = f'{seq_number}_auto_{datetime.now().strftime("%Y%m%d_%H%M")}.py'
migration_file = Path(self.migration_dir) / name
migration_file.write_text(migration_source)
log.info('Migration file "%s" was created', migration_file)
def __init__(self, document_type, *, forward_func=None, backward_func=None, **kwargs):
super().__init__(document_type, **kwargs)
if forward_func is None and backward_func is None:
raise ActionError("forward_func and backward_func are not set")
self.forward_func = forward_func
self.backward_func = backward_func
def __init__(self, document_type: str, field_name: str, **kwargs):
"""
:param document_type: collection name to be affected
:param field_name: changing mongoengine document field name
"""
super().__init__(document_type, **kwargs)
self.field_name = field_name
db_field = kwargs.get('db_field')
if db_field and '.' in db_field:
raise ActionError(f"db_field must not contain dots "
f"{self.document_type}.{self.field_name}")
def upgrade(self,
migration_name: str,
graph: Optional[MigrationsGraph] = None):
"""
Upgrade db to the given migration
:param migration_name: target migration name
:param graph: Optional. Migrations graph. If omitted, then it
will be loaded
:return:
"""
if graph is None:
log.debug('Loading migration files...')
graph = self.build_graph()
log.debug('Loading schema from database...')
left_schema = self.load_db_schema()
if migration_name not in graph.migrations:
raise MigrationGraphError(f'Migration {migration_name} not found')
db = self.db
for migration in graph.walk_down(graph.initial, unapplied_only=True):
log.info('Upgrading %s...', migration.name)
for idx, action_object in enumerate(migration.get_actions(),
start=1):
log.debug('> [%d] %s', idx, str(action_object))
if not action_object.dummy_action and not runtime_flags.schema_only:
action_object.prepare(db, left_schema, migration.policy)
action_object.run_forward()
action_object.cleanup()
try:
left_schema = patch(
action_object.to_schema_patch(left_schema),
left_schema)
except (TypeError, ValueError, KeyError) as e:
raise ActionError(
f"Unable to apply schema patch of {action_object!r}. More likely that the "
f"schema is corrupted. You can use schema repair tools to fix this issue"
) from e
graph.migrations[migration.name].applied = True
if not runtime_flags.dry_run:
log.debug('Writing db schema and migrations graph...')
self.write_db_schema(left_schema)
self.write_db_migrations_graph(graph)
if migration.name == migration_name:
break # We've reached the target migration
self._verify_schema(left_schema)
def build_schema(cls, field_obj: mongoengine.fields.BaseField) -> dict:
"""
Return schema of a given mongoengine field object
'type_key' schema item will get filled with a mongoengine field
class name or one of its parents which have its own type key
in registry
:param field_obj: mongoengine field object
:return: field schema dict
"""
schema_skel = cls.schema_skel()
schema = {
f: getattr(field_obj, f, val)
for f, val in schema_skel.items()
}
if 'default' in schema:
schema['default'] = cls._normalize_default(schema['default'])
if 'choices' in schema:
schema['choices'] = cls._normalize_choices(schema['choices'])
field_class = field_obj.__class__
if field_class.__name__ in type_key_registry:
schema['type_key'] = field_class.__name__
else:
registry_field_cls = get_closest_parent(
field_class, (x.field_cls for x in type_key_registry.values()))
if registry_field_cls is None:
raise ActionError(
f'Could not find {field_class!r} or one of its base classes '
f'in type_key registry')
schema['type_key'] = registry_field_cls.__name__
return schema
def collect_models_schema() -> Schema:
"""
Transform all available mongoengine document objects to db schema
:return:
"""
schema = Schema()
collections: Dict[str,
set] = {} # {collection_name: set(top_level_documents)}
# Retrieve models from mongoengine global document registry
for model_cls in _document_registry.values():
log.debug('> Reading document %s', repr(model_cls))
# NOTE: EmbeddedDocuments are not append 'abstract' in meta if
# `meta` is defined
if model_cls._meta.get('abstract'):
log.debug('> Skip %s since it is an abstract document')
continue
document_type = get_document_type(model_cls)
if document_type is None:
raise ActionError(f'Could not get document type for {model_cls!r}')
if document_type in schema:
raise ActionError(
f'Models with the same document types {document_type!r} found')
schema[document_type] = Schema.Document()
if not document_type.startswith(
runtime_flags.EMBEDDED_DOCUMENT_NAME_PREFIX):
col = schema[document_type].parameters[
'collection'] = model_cls._get_collection_name()
# Determine if unrelated documents have the same collection
# E.g. DropDocument could drop all of these documents
top_lvl_doc = document_type.split(
runtime_flags.DOCUMENT_NAME_SEPARATOR)[0]
collections.setdefault(col, set())
if collections[col] and top_lvl_doc not in collections[col]:
collections[col].add(top_lvl_doc)
log.warning(
f'These mongoengine documents use the same collection '
f'which can be a cause of data corruption: {collections[col]}'
)
else:
collections[col].add(top_lvl_doc)
if model_cls._meta.get('allow_inheritance'):
schema[document_type].parameters['inherit'] = True
if model_cls._dynamic:
schema[document_type].parameters['dynamic'] = True
# {field_cls: TypeKeyRegistryItem}
field_mapping_registry = {
x.field_cls: x
for x in type_key_registry.values()
}
# Collect schema for every field
for field_name, field_obj in model_cls._fields.items():
# Exclude '_id' special MongoDB field since it is immutable
# and should not be a part of migrations
if field_obj.db_field == '_id':
continue
field_cls = field_obj.__class__
if field_cls in field_mapping_registry:
registry_field_cls = field_cls
else:
registry_field_cls = get_closest_parent(
field_cls, field_mapping_registry.keys())
if registry_field_cls is None:
raise ActionError(
f'Could not find {field_cls!r} or one of its base classes '
f'in type_key registry')
handler_cls = field_mapping_registry[
registry_field_cls].field_handler_cls
schema[document_type][field_name] = handler_cls.build_schema(
field_obj)
# TODO: validate default against all field restrictions such as min_length, regex, etc.
log.debug("> Schema '%s' => %s", document_type,
str(schema[document_type]))
return schema
def downgrade(self,
migration_name: str,
graph: Optional[MigrationsGraph] = None):
"""
Downgrade db to the given migration
:param migration_name: target migration name
:param graph: Optional. Migrations graph. If omitted, then it
will be loaded
:return:
"""
if graph is None:
log.debug('Loading migration files...')
graph = self.build_graph()
log.debug('Loading schema from database...')
left_schema = self.load_db_schema()
if migration_name not in graph.migrations:
raise MigrationGraphError(f'Migration {migration_name} not found')
log.debug('Precalculating schema diffs...')
# Collect schema diffs across all migrations
migration_diffs = {} # {migration_name: [action1_diff, ...]}
temp_left_schema = Schema()
for migration in graph.walk_down(graph.initial, unapplied_only=False):
migration_diffs[migration.name] = []
for action in migration.get_actions():
forward_patch = action.to_schema_patch(temp_left_schema)
migration_diffs[migration.name].append(forward_patch)
try:
temp_left_schema = patch(forward_patch, temp_left_schema)
except (TypeError, ValueError, KeyError) as e:
raise ActionError(
f"Unable to apply schema patch of {action!r}. More likely that the "
f"schema is corrupted. You can use schema repair tools to fix this issue"
) from e
db = self.db
for migration in graph.walk_up(graph.last, applied_only=True):
if migration.name == migration_name:
break # We've reached the target migration
log.info('Downgrading %s...', migration.name)
action_diffs = zip(migration.get_actions(),
migration_diffs[migration.name],
range(1,
len(migration.get_actions()) + 1))
for action_object, action_diff, idx in reversed(
list(action_diffs)):
log.debug('> [%d] %s', idx, str(action_object))
try:
left_schema = patch(list(swap(action_diff)), left_schema)
except (TypeError, ValueError, KeyError) as e:
raise ActionError(
f"Unable to apply schema patch of {action_object!r}. More likely that the "
f"schema is corrupted. You can use schema repair tools to fix this issue"
) from e
if not action_object.dummy_action and not runtime_flags.schema_only:
action_object.prepare(db, left_schema, migration.policy)
action_object.run_backward()
action_object.cleanup()
graph.migrations[migration.name].applied = False
if not runtime_flags.dry_run:
log.debug('Writing db schema and migrations graph...')
self.write_db_schema(left_schema)
self.write_db_migrations_graph(graph)
self._verify_schema(left_schema)