def get_iterator(self, missing_items=None): def data_provider(batch_number): try: return self.batches[batch_number] except IndexError: return [] itr = ResumableFunctionIterator('test', data_provider, TestArgsProvider()) itr.couch_db = self.couch_db return itr
def get_endkey_docid(domain, doc_type, migration_id): resume_key = "%s.%s.%s" % (domain, doc_type, migration_id) state = ResumableFunctionIterator(resume_key, None, None, None).state assert getattr(state, '_rev', None), "rebuild not necessary (no resume state)" assert not state.complete, "iteration is complete" state_json = state.to_json() assert not state_json['args'] kwargs = state_json['kwargs'] return kwargs['startkey'], kwargs['startkey_docid']
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
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 get_iterator(self, missing_items=None): def data_provider(batch_number): try: return self.batches[batch_number] except IndexError: return [] def item_getter(item_id): if missing_items and item_id in missing_items: return None return int(item_id) itr = ResumableFunctionIterator('test', data_provider, TestArgsProvider(), item_getter) itr.couch_db = self.couch_db return itr
def _iter_missing_ids(db, min_tries, resume_key, view_name, view_params, repair): def data_function(**view_kwargs): @retry_on_couch_error def get_doc_ids(): results = list(db.view(view_name, **view_kwargs)) if "limit" in view_kwargs and results: nonlocal last_result last_result = results[-1] replace_limit_with_endkey(view_kwargs, last_result) return {r["id"] for r in results} def replace_limit_with_endkey(view_kwargs, last_result): assert "endkey_docid" not in view_kwargs, view_kwargs view_kwargs.pop("limit") view_kwargs["endkey"] = last_result["key"] view_kwargs["endkey_docid"] = last_result["id"] last_result = None missing, tries = find_missing_ids(get_doc_ids, min_tries=min_tries) if last_result is None: log.debug("no results %s - %s", view_kwargs['startkey'], view_kwargs['endkey']) assert not missing return [] if missing and repair: missing, tries2, repaired = repair_couch_docs(db, missing, get_doc_ids, min_tries) tries += tries2 else: repaired = 0 log.debug(f"{len(missing)}/{tries} start={view_kwargs['startkey']} {missing or ''}") last_result["missing_info"] = missing, tries, repaired return [last_result] args_provider = NoSkipArgsProvider(view_params) return ResumableFunctionIterator(resume_key, data_function, args_provider)
def get_iterator(self, missing_items=None): def data_provider(batch_number): try: return self.batches[batch_number] except IndexError: return [] def item_getter(item_id): if missing_items and item_id in missing_items: return None return int(item_id) itr = ResumableFunctionIterator('test', data_provider, TestArgsProvider(), item_getter) itr.couch_db = self.couch_db return itr
def _iter_docs(domain, doc_type, resume_key, stopper): def data_function(**view_kwargs): return couch_db.view('by_domain_doc_type_date/view', **view_kwargs) if "." in doc_type: doc_type, row_key = doc_type.split(".") else: row_key = "doc" if stopper.clean_break: return [] couch_db = XFormInstance.get_db() args_provider = NoSkipArgsProvider({ 'startkey': [domain, doc_type], 'endkey': [domain, doc_type, {}], 'limit': _iter_docs.chunk_size, 'include_docs': row_key == "doc", 'reduce': False, }) rows = ResumableFunctionIterator( resume_key, data_function, args_provider, item_getter=None, event_handler=MigrationPaginationEventHandler(domain, stopper)) return (row[row_key] for row in rows)
def get_main_forms_iteration_stop_date(domain_name, migration_id): resume_key = "%s.%s.%s" % (domain_name, "XFormInstance", migration_id) itr = ResumableFunctionIterator(resume_key, None, None, None) kwargs = itr.state.kwargs assert kwargs, f"migration state not found: {resume_key}" # this is tightly coupled to by_domain_doc_type_date/view in couch: # the last key element is expected to be a datetime return kwargs["startkey"][-1]
def get_main_forms_iteration_stop_date(statedb): resume_key = f"{statedb.domain}.XFormInstance.{statedb.unique_id}" itr = ResumableFunctionIterator(resume_key, None, None, None) kwargs = itr.state.kwargs assert kwargs, f"migration state not found: {resume_key}" # this is tightly coupled to by_domain_doc_type_date/view in couch: # the last key element is expected to be a datetime return kwargs["startkey"][-1]
def _iter_skipped_form_ids(domain, migration_id, stopper, with_progress): resume_key = "%s.%s.%s" % (domain, "XFormInstance.id", migration_id) couch_ids = _iter_docs(domain, "XFormInstance.id", resume_key, stopper) couch_ids = with_progress(["XFormInstance"], couch_ids, "Scanning", offset_key="XFormInstance.id") for batch in chunked(couch_ids, _iter_skipped_form_ids.chunk_size, list): yield from _drop_sql_form_ids(batch, domain) if not stopper.clean_break: # discard iteration state on successful completion so it is possible # to run another skipped forms iteration later ResumableFunctionIterator(resume_key, None, None, None).discard_state()
def _get_paginated_iterable(data_function, args_provider, event_handler=None, resumable_key=None): if resumable_key: return ResumableFunctionIterator(resumable_key, data_function, args_provider, lambda x: x.id, event_handler=event_handler) else: return paginate_function(data_function, args_provider, event_handler=event_handler)
def _iter_docs(domain, doc_type, resume_key, stopper): @retry_on_couch_error def data_function(**view_kwargs): view_name = 'by_domain_doc_type_date/view' results = list(couch_db.view(view_name, **view_kwargs)) assert all(r['key'][0] == domain for r in results), \ _repr_bad_results(view_name, view_kwargs, results, domain) return results if "." in doc_type: doc_type, row_key = doc_type.split(".") else: row_key = "doc" if stopper.clean_break: return [] couch_db = XFormInstance.get_db() args_provider = NoSkipArgsProvider({ 'startkey': [domain, doc_type], 'endkey': [domain, doc_type, {}], 'limit': _iter_docs.chunk_size, 'include_docs': row_key == "doc", 'reduce': False, }) rows = ResumableFunctionIterator( resume_key, data_function, args_provider, item_getter=None, event_handler=MigrationPaginationEventHandler(domain, stopper)) if rows.state.is_resume(): log.info("iteration state: %r", rows.state.to_json()) row = None try: for row in rows: yield row[row_key] finally: if row is not None: row_copy = dict(row) row_copy.pop("doc", None) log.info("last item: %r", row_copy) log.info("final iteration state: %r", rows.state.to_json())
def discard_iteration_state(resume_key): ResumableFunctionIterator(resume_key, None, None, None).discard_state()
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)
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)