class IterationState: statedb = attr.ib() domain = attr.ib() doc_type = attr.ib() def __attrs_post_init__(self): migration_id = self.statedb.unique_id resume_key = "%s.%s.%s" % (self.domain, self.doc_type, migration_id) self.itr = ResumableFunctionIterator(resume_key, None, None, None) @property def value(self): return self.itr.state.to_json() def backup_resume_state(self, value): """Attempt to save iteration state in state db :param value: iteration state dict. See `self.value` :returns: resume key if saved else `None` """ assert isinstance(value, dict), value key = f"resume-{value['timestamp']}" pretty_value = json.dumps(value, indent=2) log.info("saving resume state with key=%s : %s", key, pretty_value) old = self.statedb.get(key) if old is None: self.statedb.set(key, value) log.info("saved.") elif old != value: log.warn("NOT SAVED! refusing to overwrite:\n%s", old) return None return key def restore_resume_state(self, key): """Attempt to restore resume state represented by key :returns: true if restored else false """ new_state = self.statedb.get(key) if new_state is None: return False key = self.backup_resume_state(self.value) if key is None: return False log.info("restoring iteration state: %s", new_state) self.itr._save_state_json(new_state) return True def drop_from_couch(self): """Delete resume state from Couch""" try: self.itr.couch_db.delete_doc(self.itr.iteration_id) except ResourceNotFound: pass
class Rewinder: statedb = attr.ib() domain = attr.ib() doc_type = attr.ib() move_to = attr.ib() def __attrs_post_init__(self): migration_id = self.statedb.unique_id resume_key = "%s.%s.%s" % (self.domain, self.doc_type, migration_id) self.itr = ResumableFunctionIterator(resume_key, None, None, None) for method, regex in [ ("case_rewind", r"^case-(\d+)$"), ("resume_rewind", r"^resume-"), ]: match = re.search(regex, self.move_to) if match: getattr(self, method)(match) break else: raise NotImplementedError(self.move_to) def resume_rewind(self, match): self.offset = None new_state = self.statedb.get(self.move_to) if new_state is None: sys.exit(1, "resume state not found") old_state = self.itr.state self._save_resume_state(old_state.to_json()) log.info("restoring iteration state: %s", new_state) self.itr._save_state_json(new_state) def case_rewind(self, match): self.offset = int(match.group(1)) self.stats = FormStats() def __iter__(self): def data_function(**view_kwargs): return couch_db.view('by_domain_doc_type_date/view', **view_kwargs) log.info("preparing to rewind: %s", self.move_to) state_json = self.itr.state.to_json() self._save_resume_state(state_json) couch_db = XFormInstance.get_db() args_provider = NoSkipArgsProvider({ 'startkey': state_json["kwargs"]["startkey"], 'startkey_docid': state_json["kwargs"]["startkey_docid"], 'endkey': [self.domain, self.doc_type], 'descending': True, 'limit': 1000, 'include_docs': True, 'reduce': False, }) args, kwargs = args_provider.get_initial_args() while True: results = list(data_function(*args, **kwargs)) results = args_provider.adjust_results(results, args, kwargs) if not results: break for result in results: yield get_received_on(result["doc"], self.stats) try: args, kwargs = args_provider.get_next_args(results[-1], *args, **kwargs) except StopIteration: break def save_state(self, received_on): state = self.itr.state startkey = state.kwargs["startkey"] assert len(startkey) == 3, startkey assert isinstance(startkey[-1], type(received_on)), (startkey, received_on) startkey[-1] = received_on assert state.kwargs["startkey"] is startkey, (state.kwargs, startkey) state.kwargs.pop("startkey_docid") state.timestamp = datetime.utcnow() self._save_resume_state(state.to_json()) def _save_resume_state(self, state_json): assert isinstance(state_json, dict), state_json key = f"resume-{state_json['timestamp']}" log.info("saving resume state. restore with: rewind --to=%s\n%s", key, state_json) old = self.statedb.get(key) if old is None: log.info("saved.") self.statedb.set(key, state_json) elif old != state_json: log.warn("NOT SAVED! refusing to overwrite:\n%s", old)