| # Licensed to the Apache Software Foundation (ASF) under one |
| # or more contributor license agreements. See the NOTICE file |
| # distributed with this work for additional information |
| # regarding copyright ownership. The ASF licenses this file |
| # to you under the Apache License, Version 2.0 (the |
| # "License"); you may not use this file except in compliance |
| # with the License. You may obtain a copy of the License at |
| # |
| # http://www.apache.org/licenses/LICENSE-2.0 |
| # |
| # Unless required by applicable law or agreed to in writing, |
| # software distributed under the License is distributed on an |
| # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY |
| # KIND, either express or implied. See the License for the |
| # specific language governing permissions and limitations |
| # under the License. |
| |
| from __future__ import unicode_literals |
| from __future__ import print_function |
| from __future__ import absolute_import |
| import sys |
| from collections import defaultdict |
| from contextlib import contextmanager |
| from itertools import groupby |
| |
| from paste.deploy.converters import asbool |
| from tg import tmpl_context as c, app_globals as g |
| from pymongo.errors import DuplicateKeyError, InvalidDocument, OperationFailure |
| |
| from ming.orm import mapper, session, Mapper |
| from ming.orm.declarative import MappedClass |
| |
| from allura.tasks.index_tasks import add_artifacts |
| from allura.lib.exceptions import CompoundError |
| from allura.lib import helpers as h |
| from allura.lib import utils |
| from . import base |
| import six |
| |
| |
| class ShowModelsCommand(base.Command): |
| min_args = 1 |
| max_args = 1 |
| usage = '<ini file>' |
| summary = 'Show the inheritance graph of all Ming models' |
| parser = base.Command.standard_parser(verbose=True) |
| |
| def command(self): |
| self.basic_setup() |
| graph = build_model_inheritance_graph() |
| for depth, cls in dfs(MappedClass, graph): |
| for line in dump_cls(depth, cls): |
| print(line) |
| |
| |
| class ReindexCommand(base.Command): |
| min_args = 1 |
| max_args = 1 |
| usage = '<ini file>' |
| summary = 'Reindex into solr and re-shortlink all artifacts' |
| parser = base.Command.standard_parser(verbose=True) |
| parser.add_option('-p', '--project', dest='project', default=None, |
| help='project to reindex') |
| parser.add_option('--project-regex', dest='project_regex', default='', |
| help='Restrict reindex to projects for which the shortname matches ' |
| 'the provided regex.') |
| parser.add_option( |
| '-n', '--neighborhood', dest='neighborhood', default=None, |
| help='neighborhood to reindex (e.g. p)') |
| |
| parser.add_option('--solr', action='store_true', dest='solr', |
| help='Solr needs artifact references to already exist.') |
| parser.add_option( |
| '--skip-solr-delete', action='store_true', dest='skip_solr_delete', |
| help='Skip clearing solr index.') |
| parser.add_option('--refs', action='store_true', dest='refs', |
| help='Update artifact references and shortlinks') |
| parser.add_option('--tasks', action='store_true', dest='tasks', |
| help='Run each individual index operation as a background task. ' |
| 'Note: this is often better, since tasks have "request" objects ' |
| 'which are needed for some markdown macros to run properly') |
| parser.add_option('--solr-hosts', dest='solr_hosts', |
| help='Override the solr host(s) to post to. Comma-separated list of solr server URLs') |
| parser.add_option( |
| '--max-chunk', dest='max_chunk', type=int, default=100 * 1000, |
| help='Max number of artifacts to index in one Solr update command') |
| parser.add_option('--ming-config', dest='ming_config', help='Path (absolute, or relative to ' |
| 'Allura root) to .ini file defining ming configuration.') |
| |
| def command(self): |
| from allura import model as M |
| self.basic_setup() |
| graph = build_model_inheritance_graph() |
| if self.options.project: |
| q_project = dict(shortname=self.options.project) |
| elif self.options.project_regex: |
| q_project = dict(shortname={'$regex': self.options.project_regex}) |
| elif self.options.neighborhood: |
| neighborhood_id = M.Neighborhood.query.get( |
| url_prefix='/%s/' % self.options.neighborhood)._id |
| q_project = dict(neighborhood_id=neighborhood_id) |
| else: |
| q_project = {} |
| |
| # if none specified, do all |
| if not self.options.solr and not self.options.refs: |
| self.options.solr = self.options.refs = True |
| |
| for projects in utils.chunked_find(M.Project, q_project): |
| for p in projects: |
| c.project = p |
| base.log.info('Reindex project %s', p.shortname) |
| # Clear index for this project |
| if self.options.solr and not self.options.skip_solr_delete: |
| g.solr.delete(q='project_id_s:%s' % p._id) |
| if self.options.refs: |
| M.ArtifactReference.query.remove( |
| {'artifact_reference.project_id': p._id}) |
| M.Shortlink.query.remove({'project_id': p._id}) |
| app_config_ids = [ac._id for ac in p.app_configs] |
| # Traverse the inheritance graph, finding all artifacts that |
| # belong to this project |
| for _, a_cls in dfs(M.Artifact, graph): |
| base.log.info(' %s', a_cls) |
| ref_ids = [] |
| # Create artifact references and shortlinks |
| for a in a_cls.query.find(dict(app_config_id={'$in': app_config_ids})): |
| if self.options.verbose: |
| base.log.info(' %s', a.shorthand_id()) |
| if self.options.refs: |
| try: |
| M.ArtifactReference.from_artifact(a) |
| M.Shortlink.from_artifact(a) |
| except Exception: |
| base.log.exception( |
| 'Making ArtifactReference/Shortlink from %s', a) |
| continue |
| ref_ids.append(a.index_id()) |
| M.main_orm_session.flush() |
| M.artifact_orm_session.clear() |
| try: |
| self._chunked_add_artifacts(ref_ids) |
| except CompoundError as err: |
| base.log.exception( |
| 'Error indexing artifacts:\n%r', err) |
| base.log.error('%s', err.format_error()) |
| M.main_orm_session.flush() |
| M.main_orm_session.clear() |
| base.log.info('Reindex %s', 'queued' if self.options.tasks else 'done') |
| |
| @property |
| def add_artifact_kwargs(self): |
| if self.options.solr_hosts: |
| return {'solr_hosts': self.options.solr_hosts.split(',')} |
| return {} |
| |
| def _chunked_add_artifacts(self, ref_ids): |
| # ref_ids contains solr index ids which can easily be over |
| # 100 bytes. Here we allow for 160 bytes avg, plus |
| # room for other document overhead. |
| for chunk in utils.chunked_list(ref_ids, self.options.max_chunk): |
| if self.options.tasks: |
| self._post_add_artifacts(chunk) |
| else: |
| add_artifacts(chunk, |
| update_solr=self.options.solr, |
| update_refs=self.options.refs, |
| **self.add_artifact_kwargs) |
| |
| def _post_add_artifacts(self, chunk): |
| """ |
| Post task, recursively splitting and re-posting if the resulting |
| mongo document is too large. |
| """ |
| try: |
| with self.ming_config(self.options.ming_config): |
| add_artifacts.post(chunk, |
| update_solr=self.options.solr, |
| update_refs=self.options.refs, |
| **self.add_artifact_kwargs) |
| except InvalidDocument as e: |
| # there are many types of InvalidDocument, only recurse if its |
| # expected to help |
| if e.args[0].startswith('BSON document too large'): |
| self._post_add_artifacts(chunk[:len(chunk) // 2]) |
| self._post_add_artifacts(chunk[len(chunk) // 2:]) |
| else: |
| raise |
| |
| @property |
| def ming_config(self): |
| """Return a contextmanager for the provided ming config, if there is |
| one. |
| |
| Otherwise return a no-op contextmanager. |
| |
| """ |
| def noop_cm(*args): |
| yield |
| |
| if self.options.ming_config: |
| return h.ming_config_from_ini |
| return contextmanager(noop_cm) |
| |
| |
| class EnsureIndexCommand(base.Command): |
| min_args = 1 |
| max_args = 1 |
| usage = '[<ini file>]' |
| summary = 'Create all the Mongo indexes specified by Ming models' |
| parser = base.Command.standard_parser(verbose=True) |
| parser.add_option('--clean', action='store_true', dest='clean', |
| help='Drop any unneeded indexes') |
| |
| def command(self): |
| from allura import model as M |
| main_session_classes = [M.main_orm_session, M.repository_orm_session, |
| M.task_orm_session, M.main_explicitflush_orm_session] |
| if asbool(self.config.get('activitystream.recording.enabled', False)): |
| from activitystream.storage.mingstorage import activity_odm_session |
| main_session_classes.append(activity_odm_session) |
| self.basic_setup() |
| # by db, then collection name |
| main_indexes = defaultdict(lambda: defaultdict(list)) |
| project_indexes = defaultdict(list) # by collection name |
| base.log.info('Collecting indexes...') |
| for m in Mapper.all_mappers(): |
| mgr = m.collection.m |
| cname = mgr.collection_name |
| cls = m.mapped_class |
| if cname is None: |
| base.log.info('... skipping abstract class %s', cls) |
| continue |
| base.log.info('... for class %s', cls) |
| if session(cls) in main_session_classes: |
| idx = main_indexes[session(cls)][cname] |
| else: |
| idx = project_indexes[cname] |
| idx.extend(mgr.indexes) |
| base.log.info('Updating indexes for main DB') |
| for odm_session, db_indexes in six.iteritems(main_indexes): |
| db = odm_session.impl.db |
| for name, indexes in six.iteritems(db_indexes): |
| self._update_indexes(db[name], indexes) |
| base.log.info('Updating indexes for project DB') |
| db = M.project_doc_session.db |
| base.log.info('... DB: %s', db) |
| for name, indexes in six.iteritems(project_indexes): |
| self._update_indexes(db[name], indexes) |
| base.log.info('Done updating indexes') |
| |
| def _update_indexes(self, collection, indexes): |
| uindexes = dict( |
| # convert list to tuple so it's hashable for 'set' |
| (tuple(i.index_spec), i) |
| for i in indexes |
| if i.unique) |
| indexes = dict( |
| (tuple(i.index_spec), i) |
| for i in indexes |
| if not i.unique) |
| prev_indexes = {} |
| prev_uindexes = {} |
| unique_flag_drop = {} |
| unique_flag_add = {} |
| try: |
| existing_indexes = six.iteritems(collection.index_information()) |
| except OperationFailure: |
| # exception is raised if db or collection doesn't exist yet |
| existing_indexes = {} |
| for iname, fields in existing_indexes: |
| if iname == '_id_': |
| continue |
| keys = tuple(fields['key']) |
| if fields.get('unique'): |
| if keys in indexes: |
| unique_flag_drop[iname] = keys |
| else: |
| prev_uindexes[iname] = keys |
| else: |
| if keys in uindexes: |
| unique_flag_add[iname] = keys |
| else: |
| prev_indexes[iname] = keys |
| |
| for iname, keys in six.iteritems(unique_flag_drop): |
| self._recreate_index(collection, iname, list(keys), unique=False) |
| for iname, keys in six.iteritems(unique_flag_add): |
| self._recreate_index(collection, iname, list(keys), unique=True) |
| |
| # Ensure all indexes |
| for keys, idx in six.iteritems(uindexes): |
| base.log.info('...... ensure %s:%s', collection.name, idx) |
| while True: |
| try: |
| index_options = idx.index_options.copy() |
| if idx.fields == ('_id',): |
| # as of mongo 3.4 _id fields can't have these options set |
| # _id is always non-sparse and unique anyway |
| del index_options['sparse'] |
| del index_options['unique'] |
| collection.ensure_index(idx.index_spec, **index_options) |
| break |
| except DuplicateKeyError as err: |
| base.log.info('Found dupe key(%s), eliminating dupes', err) |
| self._remove_dupes(collection, idx.index_spec) |
| for keys, idx in six.iteritems(indexes): |
| base.log.info('...... ensure %s:%s', collection.name, idx) |
| collection.ensure_index(idx.index_spec, background=True, **idx.index_options) |
| # Drop obsolete indexes |
| for iname, keys in six.iteritems(prev_indexes): |
| if keys not in indexes: |
| if self.options.clean: |
| base.log.info('...... drop index %s:%s', collection.name, iname) |
| collection.drop_index(iname) |
| else: |
| base.log.info('...... potentially unneeded index, could be removed by running with --clean %s:%s', |
| collection.name, iname) |
| for iname, keys in six.iteritems(prev_uindexes): |
| if keys not in uindexes: |
| if self.options.clean: |
| base.log.info('...... drop index %s:%s', collection.name, iname) |
| collection.drop_index(iname) |
| else: |
| base.log.info('...... potentially unneeded index, could be removed by running with --clean %s:%s', |
| collection.name, iname) |
| |
| def _recreate_index(self, collection, iname, keys, **creation_options): |
| '''Recreate an index with new creation options, using a temporary index |
| so that at no time is an index missing from the specified keys''' |
| superset_keys = keys + [('temporary_extra_field_for_indexing', 1)] |
| base.log.info('...... ensure index %s:%s', |
| collection.name, superset_keys) |
| superset_index = collection.ensure_index(superset_keys) |
| base.log.info('...... drop index %s:%s', collection.name, iname) |
| collection.drop_index(iname) |
| base.log.info('...... ensure index %s:%s %s', |
| collection.name, keys, creation_options) |
| collection.ensure_index(keys, **creation_options) |
| base.log.info('...... drop index %s:%s', |
| collection.name, superset_index) |
| collection.drop_index(superset_index) |
| |
| def _remove_dupes(self, collection, spec): |
| iname = collection.create_index(spec) |
| fields = [f[0] for f in spec] |
| q = collection.find({}, projection=fields).sort(spec) |
| |
| def keyfunc(doc): |
| return tuple(doc.get(f, None) for f in fields) |
| dupes = [] |
| for key, doc_iter in groupby(q, key=keyfunc): |
| docs = list(doc_iter) |
| if len(docs) > 1: |
| base.log.info('Found dupes with %s', key) |
| dupes += [doc['_id'] for doc in docs[1:]] |
| collection.drop_index(iname) |
| collection.remove(dict(_id={'$in': dupes})) |
| |
| |
| def build_model_inheritance_graph(): |
| graph = dict((m.mapped_class, ([], [])) for m in Mapper.all_mappers()) |
| for cls, (parents, children) in six.iteritems(graph): |
| for b in cls.__bases__: |
| if b not in graph: |
| continue |
| parents.append(b) |
| graph[b][1].append(cls) |
| return graph |
| |
| |
| def dump_cls(depth, cls): |
| indent = ' ' * 4 * depth |
| yield indent + '%s.%s' % (cls.__module__, cls.__name__) |
| m = mapper(cls) |
| for p in m.properties: |
| s = indent * 2 + ' - ' + str(p) |
| if hasattr(p, 'field_type'): |
| s += ' (%s)' % p.field_type |
| yield s |
| |
| |
| def dfs(root, graph, depth=0): |
| yield depth, root |
| for node in graph[root][1]: |
| for r in dfs(node, graph, depth + 1): |
| yield r |
