def _write_snapshots_for_existing_entities(
self,
session: Session,
context: "_SnapshotContext",
snapshot_time: datetime,
schema: ModuleType,
) -> None:
"""Writes snapshot updates for entities that already have snapshots
present in the database
"""
historical_class = _get_historical_class(type(context.schema_object),
schema)
new_historical_snapshot = historical_class()
self._copy_entity_fields_to_historical_snapshot(
context.schema_object, new_historical_snapshot)
new_historical_snapshot.valid_from = snapshot_time
# Snapshot must be merged separately from record tree, as they are not
# included in the ORM model relationships (to avoid needing to load
# the entire snapshot chain at once)
session.merge(new_historical_snapshot)
# Close last snapshot if one is present
if context.most_recent_snapshot is not None:
if not isinstance(context.most_recent_snapshot,
HistoryTableSharedColumns):
raise ValueError(
f"Snapshot class [{type(context.most_recent_snapshot)}] "
f"must be a subclass of "
f"[{HistoryTableSharedColumns.__name__}]")
context.most_recent_snapshot.valid_to = snapshot_time
session.merge(context.most_recent_snapshot)
def _write_snapshots_for_new_entities(
self,
session: Session,
context: "_SnapshotContext",
snapshot_time: datetime,
schema,
) -> None:
"""Writes snapshots for any new entities, including any required manual
adjustments based on provided start and end times
"""
historical_class = _get_historical_class(type(context.schema_object),
schema)
new_historical_snapshot = historical_class()
self._copy_entity_fields_to_historical_snapshot(
context.schema_object, new_historical_snapshot)
provided_start_time = None
provided_end_time = None
# Validate provided start and end times
if (context.provided_start_time
and context.provided_start_time.date() < snapshot_time.date()):
provided_start_time = context.provided_start_time
if (context.provided_end_time
and context.provided_end_time.date() < snapshot_time.date()):
provided_end_time = context.provided_end_time
if (provided_start_time and provided_end_time
and provided_start_time >= provided_end_time):
provided_start_time = None
if provided_start_time is not None and provided_end_time is None:
new_historical_snapshot.valid_from = provided_start_time
elif provided_end_time is not None:
new_historical_snapshot.valid_from = provided_end_time
else:
new_historical_snapshot.valid_from = snapshot_time
# Snapshot must be merged separately from record tree, as they are not
# included in the ORM model relationships (to avoid needing to load
# the entire snapshot chain at once)
session.merge(new_historical_snapshot)
# If both start and end time were provided, an earlier snapshot needs to
# be created, reflecting the state of the entity before its current
# completed state
if provided_start_time and provided_end_time:
initial_snapshot = historical_class()
self._copy_entity_fields_to_historical_snapshot(
context.schema_object, initial_snapshot)
initial_snapshot.valid_from = provided_start_time
initial_snapshot.valid_to = provided_end_time
self.post_process_initial_snapshot(context, initial_snapshot)
session.merge(initial_snapshot)