diff --git a/esbuild/gdc_elasticsearch.py b/esbuild/gdc_elasticsearch.py index 90a68ceb..e41de1ec 100644 --- a/esbuild/gdc_elasticsearch.py +++ b/esbuild/gdc_elasticsearch.py @@ -19,6 +19,7 @@ from gdcdatamodel.models import File from progressbar import ProgressBar, Percentage, Bar, ETA from psqlgraph import PsqlGraphDriver +from threading import Thread # TODO this could probably be bumped now that the number of bulk # threads in the config is higher, c.f. @@ -50,6 +51,98 @@ def shouldnt_delete(node): return False +def progress_bar(title, maxval): + """Create and initialize a custom progressbar + + :param str title: The text of the progress bar + "param int maxva': The maximumum value of the progress bar + + """ + + pbar = ProgressBar(widgets=[ + title, + Percentage(), ' ', + Bar(marker='#', left='[', right=']'), + ' ', + ETA(), + ' ', + ], maxval=maxval) + + pbar.update(0) + + return pbar + + +def bulk_upload(es, index, doc_type, docs, batch_size=BATCH_SIZE): + """Chunk and upload docs to Elasticsearch. This function will raise + an exception of there were errors inserting any of the + documents + + :param es: Elasticsearch client + :param str index: The index to upload documents to + :param str doc_type: The type of document to pload as + :param list docs: The documents to upload + :param int batch_size: The number of docs per batch + + """ + + if not docs: + return + + instruction = {"index": {"_index": index, "_type": doc_type}} + pbar = progress_bar('{} upload '.format(doc_type), len(docs)) + + def body(): + start = pbar.currval + for doc in docs[start:start+batch_size]: + yield instruction + yield doc + pbar.update(pbar.currval+1) + while pbar.currval < len(docs): + res = es.bulk(body=body()) + if res['errors']: + raise RuntimeError(json.dumps([ + d for d in res['items'] if d['index']['status'] != 100 + ], indent=2)) + pbar.finish() + + +def upload_to_index(es, index, case_docs, ann_docs, file_docs, + project_docs, batch_size=BATCH_SIZE): + """Upload {case,ann,file,project}_docs to `index` with Elasticsearch + client `es` + + """ + + bulk_upload(es, index, 'project', project_docs, batch_size) + bulk_upload(es, index, 'annotation', ann_docs, batch_size) + bulk_upload(es, index, 'case', case_docs, batch_size) + bulk_upload(es, index, 'file', file_docs, batch_size) + + +def async_upload_to_index(es, index, case_docs, ann_docs, file_docs, + project_docs, batch_size=BATCH_SIZE): + + """Spawn a thread to upload {case,ann,file,project}_docs to `index` + with Elasticsearch client `es` asynchronously + + """ + + thread = Thread(target=upload_to_index, args=( + es, + index, + case_docs, + ann_docs, + file_docs, + project_docs, + batch_size, + )) + + thread.start() + + return thread + + class GDCElasticsearch(object): """ @@ -84,89 +177,134 @@ def __init__(self, converter_class, es=None, self.converter = converter_class(self.graph) - def go(self, roll_alias=True): - self.log.info("Caching database") - # having a transation out here is important, since it ensures - # that the cached database and which nodes get deleted is - # consistent - with self.graph.session_scope() as session: + def denormalize_database(self): + """Cache the database, stash the to_delete files, and build the + index. Returns a generator yielding (case, file, annotation, + project) docs for each project + + """ + + with self.graph.session_scope(must_inherit=True): + self.log.info("Caching database") self.converter.cache_database() + + self.log.info("Denormalizing database into JSON docs") + + project_indexes = self.converter.denormalize_all_by_projects() + + for project_index in project_indexes: + # (case_docs, file_docs, ann_docs, project_docs) + yield project_index + + + def get_to_delete_nodes(self): + """Store list of nodes to be deleted after index is successfully + released. See also :func:`self.delete_to_delete_nodes`. + + """ + + with self.graph.session_scope(must_inherit=True): self.log.info("Querying for old nodes to delete") to_delete = self.graph.nodes().sysan({"to_delete": True}).all() to_delete = [n for n in to_delete if not shouldnt_delete(n)] self.log.info("Found %s to_delete nodes, saving for later", len(to_delete)) - self.log.info("Denormalizing database into JSON docs") - case_docs, file_docs, ann_docs, project_docs = self.converter.denormalize_all() - self.log.info("%s case docs, %s file docs, %s annotation docs, %s project docs", - len(case_docs), - len(file_docs), - len(ann_docs), - len(project_docs)) - self.log.info("Validating docs produced") - self.converter.validate_docs(case_docs, file_docs, ann_docs, project_docs) - self.log.info("Deploying new ES index with new docs and bumping alias") - new_index = self.deploy(case_docs, file_docs, ann_docs, project_docs, - roll_alias=roll_alias) - with self.graph.session_scope() as session: - for expired_node in to_delete: - node = self.graph.nodes(expired_node.__class__)\ - .ids(expired_node.node_id)\ - .scalar() - if node: - self.log.info("Deleting %s", node) - session.delete(node) + + return to_delete + + def go(self, roll_alias=True, batch_size=BATCH_SIZE): + """Create a new index with an incremented name based on + self.index_base, populates it with a new index + + """ + + last_upload_thread = None + + # Create index and upload mapping + index = self.setup_new_index() + + case_count = 0 + file_count = 0 + annotation_count = 0 + project_count = 0 + + # having a transation out here is important, since it ensures + # that the cached database and which nodes get deleted is + # consistent + with self.graph.session_scope(): + + # Store nodes that should be deleted later + to_delete = self.get_to_delete_nodes() + + # Create generator that returns partial indexes + project_indexes = self.denormalize_database() + + # Consume generator and load into ES asynchronously + for project_index in project_indexes: + case_docs, file_docs, ann_docs, project_docs = project_index + + case_count += len(case_docs) + annotation_count += len(ann_docs) + file_count += len(file_docs) + project_count += len(project_docs) + + # Only allow one index to upload asynchronously at + # one time, so join the last one if it exists + if last_upload_thread is not None: + self.log.info("Waiting for previous upload thread...") + last_upload_thread.join() + + # Asynchronously upload this project + last_upload_thread = async_upload_to_index( + self.es, index, case_docs, ann_docs, file_docs, + project_docs, batch_size) + self.log.info("Spawned async upload thread") + + # Wait for last upload thread + if last_upload_thread is not None: + self.log.info("Waiting on final upload thread.") + last_upload_thread.join() + else: + raise RuntimeError( + "No async thread was used, so no docs were uploaded.") + + # Check the total number of docs are in ES + self.check_document_counts( + index, case_count, file_count, annotation_count, project_count) + + # Update the alias to point to the new index + if roll_alias: + self.roll_alias(index) + else: + self.log.info("Skipping alias roll / old index deletion") + + self.delete_to_delete_nodes(to_delete) + statsd.event( "esbuild finished", - "successfully built index {}".format(new_index), + "successfully built index {}".format(index), source_type_name="esbuild", alert_type="success", - tags=["es_index:{}".format(new_index)], + tags=["es_index:{}".format(index)], ) - def pbar(self, title, maxval): - """Create and initialize a custom progressbar - - :param str title: The text of the progress bar - "param int maxva': The maximumum value of the progress bar + def delete_to_delete_nodes(self, to_delete_nodes): + """To make the index release appear atomic, files can be marked + `to_delete == True` in system_annotations. If it is, esbuild + will delete it as the last step in an index release. This + keeps postgres and the Elasticsearch index as consistant as + possible. """ - pbar = ProgressBar(widgets=[ - title, Percentage(), ' ', - Bar(marker='#', left='[', right=']'), ' ', - ETA(), ' '], maxval=maxval) - pbar.update(0) - return pbar - - def bulk_upload(self, index, doc_type, docs, batch_size=BATCH_SIZE): - """Chunk and upload docs to Elasticsearch. This function will raise - an exception of there were errors inserting any of the - documents - - :param str index: The index to upload documents to - :param str doc_type: The type of document to pload as - :param list docs: The documents to upload - :param int batch_size: The number of docs per batch - """ - if not docs: - return - instruction = {"index": {"_index": index, "_type": doc_type}} - pbar = self.pbar('{} upload '.format(doc_type), len(docs)) - - def body(): - start = pbar.currval - for doc in docs[start:start+batch_size]: - yield instruction - yield doc - pbar.update(pbar.currval+1) - while pbar.currval < len(docs): - res = self.es.bulk(body=body()) - if res['errors']: - raise RuntimeError(json.dumps([ - d for d in res['items'] if d['index']['status'] != 100 - ], indent=2)) - pbar.finish() + with self.graph.session_scope() as session: + for expired_node in to_delete_nodes: + node = (self.graph.nodes(expired_node.__class__) + .ids(expired_node.node_id) + .scalar()) + if node: + self.log.info("Deleting %s", node) + session.delete(node) def put_mappings(self, index): """Add mappings to index. @@ -193,43 +331,39 @@ def put_mappings(self, index): body=self.converter.mapper.get_annotation_es_mapping()), ] - def index_populate(self, index, case_docs=[], file_docs=[], - ann_docs=[], project_docs=[], - batch_size=BATCH_SIZE): - self.bulk_upload(index, 'project', project_docs, batch_size) - self.bulk_upload(index, 'annotation', ann_docs, batch_size) - self.bulk_upload(index, 'case', case_docs, batch_size) - self.bulk_upload(index, 'file', file_docs, batch_size) - - def index_create_and_populate(self, index, case_docs=[], - file_docs=[], ann_docs=[], project_docs=[], - batch_size=BATCH_SIZE): - """Create a new index with name `index` and add given documents to it. - `case_docs` or `project_docs` are empty, the will be generated - automatically. - - :param list case_docs: The case docs to upload. - :param list file_docs: - The file docs to upload. If case_docs is empty, - `file_docs` will be overwritten when case_docs are - produced. - :param list ann_docs: - The annotation docs to upload. If case_docs is empty, - `ann_docs` will be overwritten when case_docs are - produced. - :param list project_docs: The project docs to upload. + def get_next_index_name(self): + """Generates and returns the string with the index name of the next + index based on the highest numbered existing index """ + current_numbers = self.get_index_numbers() + self.log.info("Currently deployed indices are %s", current_numbers) + + if not current_numbers: + index_number = 1 + else: + index_number = max(current_numbers)+1 + + new_index = INDEX_PATTERN.format(base=self.index_base, n=index_number) + + return new_index + + def setup_new_index(self): + """Create a new index by incrementing the naming scheme. + + :returns: A string name of the index + + """ + + new_index = self.get_next_index_name() + + self.log.info("Deploying to index %s", new_index) index_settings = self.converter.mapper.index_settings() - self.es.indices.create(index=index, body=index_settings) - self.put_mappings(index) - if not case_docs: - self.log.warning("There were no case docs passed to populate with!") - if not project_docs: - self.log.warning("There were no case docs passed to populate with!") - self.index_populate(index, case_docs, file_docs, ann_docs, - project_docs, batch_size) + self.es.indices.create(index=new_index, body=index_settings) + self.put_mappings(new_index) + + return new_index def swap_index(self, old_index, new_index): """Atomically switch the resolution of `alias` from `old_index` to @@ -246,10 +380,12 @@ def swap_index(self, old_index, new_index): {'add': {'index': new_index, 'alias': self.index_base}}]}) def get_index_numbers(self): - """Return the numbers of the current set of indices. So concretely if we - have gdc_from_graph_23, gdc_from_graph_24, and + """Return the numbers of the current set of indices. So concretely if + we have gdc_from_graph_23, gdc_from_graph_24, and gdc_from_graph_25, this will return [23, 24, 25]. + """ + indices = set(self.es.indices.get_aliases().keys()) p = re.compile(INDEX_PATTERN.format(base=self.index_base, n='(\d+)')+'$') matches = [p.match(index) for index in indices if p.match(index)] @@ -269,15 +405,67 @@ def lookup_index_by_alias(self): except NotFoundError: return None + def check_document_counts(self, index, file_count, case_count, ann_count, + project_count): + """sanity checks that there are the correct number of docs in the new + index and logs warnings of counts don't match + + """ + + msg = 'There appears to be the wrong number of {0} files. {1} != {2}' + + es_file_count = self.es.count(index=index, doc_type="file")["count"] + es_case_count = self.es.count(index=index, doc_type="case")["count"] + es_ann_count = self.es.count(index=index, doc_type="annotation")["count"] + es_project_count = self.es.count(index=index, doc_type="project")["count"] + + if es_file_count != file_count: + self.log.warning(msg.format('file', es_file_count, file_count)) + + if es_case_count != case_count: + self.log.warning(msg.format('case', es_case_count, case_count)) + + if es_ann_count != ann_count: + self.log.warning(msg.format('annotation', es_ann_count, ann_count)) + + if es_project_count != project_count: + self.log.warning(msg.format('project', es_project_count, project_count)) + + def roll_alias(self, new_index): + """atomically switch the alias to point to the new index, and delete + anything older than the last 5 versions of this index. + + """ + + # ensure all writes are visible + self.es.indices.refresh(index=new_index) + + # Roll indices + self.log.info("Rolling alias and deleting old indices") + old_index = self.lookup_index_by_alias() + + if old_index: + assert old_index.startswith(self.index_base) + self.swap_index(old_index, new_index) + else: + self.es.indices.put_alias(index=new_index, name=self.index_base) + + self.cleanup_old_indices([old_index, new_index]) + def cleanup_old_indices(self, kept): self.log.info("Deleting old indices") - numbers = self.get_index_numbers() - indices = [INDEX_PATTERN.format(base=self.index_base, n=n) - for n in numbers] - if len(numbers) <= 5: + + index_numbers = self.get_index_numbers() + indices = [ + INDEX_PATTERN.format(base=self.index_base, n=n) + for n in index_numbers + ] + + if len(index_numbers) <= 5: self.log.info("less than 5 matching indices found, not deleting anything") to_close = indices to_delete = [] + else: to_delete = indices[0:-5] to_close = indices[-5:] @@ -286,80 +474,24 @@ def cleanup_old_indices(self, kept): for index in to_delete: self.log.info("Deleting %s", index) self.es.indices.delete(index=index) + for index in to_close: if index not in kept: self.log.info("Closing %s", index) try: self.es.indices.flush(index=index) - except AuthorizationException as e: - # authorization exception will be raised if it's already closed - if "IndexClosedException" in e.error: - self.log.info("%s is already closed" % index) + except AuthorizationException as exception: + # authorization exception will be raised if it's + # already closed + if "IndexClosedException" in exception.error: + self.log.info("%s is already closed", index) continue else: - self.log.exception("Fail to flush %s" % index) + self.log.exception("Fail to flush %s", index) except: - self.log.exception("Can't flush index %s" % index) + self.log.exception("Can't flush index %s", index) + try: self.es.indices.close(index=index) except: - self.log.error("Can't close index %s" % index) - - - def deploy(self, case_docs, file_docs, ann_docs, - project_docs, roll_alias=True, - batch_size=BATCH_SIZE): - """Create a new index with an incremented name based on - self.index_base, populate it with :func - index_create_and_populate:, atomically switch the alias to - point to the new index, and delete anything older than the - last 5 versions of this index. - - """ - current_numbers = self.get_index_numbers() - self.log.info("Currently deployed indices are %s", current_numbers) - if not current_numbers: - n = 1 - else: - n = max(current_numbers)+1 - new_index = INDEX_PATTERN.format(base=self.index_base, n=n) - self.log.info("Deploying to index %s", new_index) - self.index_create_and_populate(new_index, case_docs, - file_docs, ann_docs, - project_docs, batch_size) - if roll_alias: - # ensure all writes are visible - self.es.indices.refresh(index=new_index) - - # sanity checks that there are the correct number of docs in the new index - msg = ('There appears to be the wrong number of {0} files. {1} != {2}') - - file_count = self.es.count(index=new_index, doc_type="file")["count"] - case_count = self.es.count(index=new_index, doc_type="case")["count"] - ann_count = self.es.count(index=new_index, doc_type="annotation")["count"] - project_count = self.es.count(index=new_index, doc_type="project")["count"] - - if file_count != len(file_docs): - self.log.warning(msg.format('file', file_count, len(file_docs))) - - if case_count != len(case_docs): - self.log.warning(msg.format('case', case_count, len(case_docs))) - - if ann_count != len(ann_docs): - self.log.warning(msg.format('annotation', ann_count, len(ann_docs))) - - if project_count != len(project_docs): - self.log.warning(msg.format('project', project_count, len(project_docs))) - - # Roll indices - self.log.info("Rolling alias and deleting old indices") - old_index = self.lookup_index_by_alias() - if old_index: - assert old_index.startswith(self.index_base) - self.swap_index(old_index, new_index) - else: - self.es.indices.put_alias(index=new_index, name=self.index_base) - self.cleanup_old_indices([old_index, new_index]) - else: - self.log.info("Skipping alias roll / old index deletion") - return new_index + self.log.error("Can't close index %s", index) diff --git a/esbuild/graph/common/builder.py b/esbuild/graph/common/builder.py index 0c907abe..8038e5fe 100644 --- a/esbuild/graph/common/builder.py +++ b/esbuild/graph/common/builder.py @@ -1255,6 +1255,32 @@ def denormalize_cases(self, cases=None): pbar.finish() return case_docs, file_docs.values(), ann_docs.values() + def denormalize_cases_from_project(self, project): + """Produce documents for a single project""" + + cases = list(self.neighbors_labeled(project, "case")) + return self.denormalize_cases(cases) + + def denormalize_cases_by_projects(self): + """Yields the (case_docs, file_docs, annotation_docs) for each project + in the index""" + + for project in self.nodes_labeled('project'): + yield self.denormalize_cases_from_project(project) + + def denormalize_all_by_projects(self): + """Yields the (case_docs, file_docs, annotation_docs, project_docs) + for each project in the index + + """ + + for project in self.nodes_labeled('project'): + index = self.denormalize_cases_from_project(project) + case_docs, file_docs, ann_docs = index + project_docs = [self.denormalize_project(project)] + + yield case_docs, file_docs, ann_docs, project_docs + def denormalize_projects(self, projects=None): """If projects is not specified, denormalize all projects in the graph. If projects is specified, denormalize only those diff --git a/tests/test_suite_data.gv.pdf b/tests/test_suite_data.gv.pdf index d4a61b87..7e1ef716 100644 Binary files a/tests/test_suite_data.gv.pdf and b/tests/test_suite_data.gv.pdf differ