Safely iterate over a DB QuerySet yielding ES documents
976898
def iter_qs(qs, adapter):
    ''''''
    for obj in qs.no_cache().no_dereference().timeout(False):
        if adapter.is_indexable(obj):
            try:
                doc = adapter.from_model(obj).to_dict(include_meta=True)
                yield doc
            except Exception as e:
                model = adapter.model.__name__
                log.error('Unable to index %s "%s": %s', model, str(obj.id),
                          str(e), exc_info=True)
189455
def generic_in(self, **kwargs):
        ''''''
        query = {}
        for key, value in kwargs.items():
            if not value:
                continue
            # Optimize query for when there is only one value
            if isinstance(value, (list, tuple)) and len(value) == 1:
                value = value[0]
            if isinstance(value, (list, tuple)):
                if all(isinstance(v, basestring) for v in value):
                    ids = [ObjectId(v) for v in value]
                    query['{0}._ref.$id'.format(key)] = {'$in': ids}
                elif all(isinstance(v, DBRef) for v in value):
                    query['{0}._ref'.format(key)] = {'$in': value}
                elif all(isinstance(v, ObjectId) for v in value):
                    query['{0}._ref.$id'.format(key)] = {'$in': value}
            elif isinstance(value, ObjectId):
                query['{0}._ref.$id'.format(key)] = value
            elif isinstance(value, basestring):
                query['{0}._ref.$id'.format(key)] = ObjectId(value)
            else:
                self.error('expect a list of string, ObjectId or DBRef')
        return self(__raw__=query)
33730
def index(models=None, name=None, force=False, keep=False):
    '''
    
    '''
    index_name = name or default_index_name()

    doc_types_names = [m.__name__.lower() for m in adapter_catalog.keys()]
    models = [model.lower().rstrip('s') for model in (models or [])]
    for model in models:
        if model not in doc_types_names:
            log.error('Unknown model %s', model)
            sys.exit(-1)

    log.info('Initiliazing index "{0}"'.format(index_name))
    if es.indices.exists(index_name):
        if IS_TTY and not force:
            msg = 'Index {0} will be deleted, are you sure?'
            click.confirm(msg.format(index_name), abort=True)
        es.indices.delete(index_name)

    es.initialize(index_name)

    with handle_error(index_name, keep):

        disable_refresh(index_name)
        for adapter in iter_adapters():
            if not models or adapter.doc_type().lower() in models:
                index_model(index_name, adapter)
            else:
                log.info('Copying {0} objects to the new index'.format(
                         adapter.model.__name__))
                # Need upgrade to Elasticsearch-py 5.0.0 to write:
                # es.reindex({
                #     'source': {'index': es.index_name, 'type': adapter.doc_type()},
                #     'dest': {'index': index_name}
                # })
                #
                # http://elasticsearch-py.readthedocs.io/en/master/api.html#elasticsearch.Elasticsearch.reindex
                # This method (introduced in Elasticsearch 2.3 but only in Elasticsearch-py 5.0.0)
                # triggers a server-side documents copy.
                # Instead we use this helper for meant for backward compatibility
                # but with poor performance as copy is client-side (scan+bulk)
                es_reindex(es.client, es.index_name, index_name, scan_kwargs={
                    'doc_type': adapter.doc_type()
                })

        enable_refresh(index_name)
    # At this step, we don't want error handler to delete the index
    # in case of error
    set_alias(index_name, delete=not keep)
65256
def get_docs_sources_from_ES(self):
        """"""
        docs = [doc for doc, _, _, get_from_ES in self.doc_to_update if get_from_ES]
        if docs:
            documents = self.docman.elastic.mget(body={"docs": docs}, realtime=True)
            return iter(documents["docs"])
        else:
            return iter([])
222565
def queryset_iterator(queryset, chunksize=1000):
    
    if queryset.exists():
        primary_key = 0
        last_pk = queryset.order_by('-pk')[0].pk
        queryset = queryset.order_by('pk')
        while primary_key < last_pk:
            for row in queryset.filter(pk__gt=primary_key)[:chunksize]:
                primary_key = row.pk
                yield row
            gc.collect()
1237530
def _execute(self, query, model, adapter, raw=False):
        
        values = self.load(model, adapter)
        return IterableStore(values=values)._execute(query, model=model, adapter=None, raw=raw)
1162536
def build_for_each(self, db, safe_mode=False, extra=None):
        """"""
        result = dict()
        for var, query in iteritems(self.for_each):
            result[var] = db.query(
                query,
                additional_locals=extra,
                safe_mode=safe_mode
            )
        return result
836833
def all_documents(index=INDEX_NAME):
    
    query = {
        'query': {
            'match_all': {}
        }
    }
    for result in raw_query(query, index=index):
        yield result
663371
def get_queryset(self):
        """"""
        self.strain = get_object_or_404(Strain, Strain_slug__iexact=self.kwargs['slug'])
        return PlugEvents.objects.filter(Breeding__Strain=self.strain)
765311
def _query(options, collection_name, num_to_skip,
           num_to_return, query, field_selector, opts, check_keys):
    """"""
    encoded = _dict_to_bson(query, check_keys, opts)
    if field_selector:
        efs = _dict_to_bson(field_selector, False, opts)
    else:
        efs = b""
    max_bson_size = max(len(encoded), len(efs))
    return b"".join([
        _pack_int(options),
        _make_c_string(collection_name),
        _pack_int(num_to_skip),
        _pack_int(num_to_return),
        encoded,
        efs]), max_bson_size
449471
def generate_query_report(self, db_uri, parsed_query, db_name, collection_name):
        """"""
        index_analysis = None
        recommendation = None
        namespace = parsed_query['ns']
        indexStatus = "unknown"

        index_cache_entry = self._ensure_index_cache(db_uri,
                                                     db_name,
                                                     collection_name)


        query_analysis = self._generate_query_analysis(parsed_query,
                                                       db_name,
                                                       collection_name)
        if ((query_analysis['analyzedFields'] != []) and
             query_analysis['supported']):
            index_analysis = self._generate_index_analysis(query_analysis,
                                                           index_cache_entry['indexes'])
            indexStatus = index_analysis['indexStatus']
            if index_analysis['indexStatus'] != 'full':
                recommendation = self._generate_recommendation(query_analysis,
                                                               db_name,
                                                               collection_name)
                # a temporary fix to suppress faulty parsing of $regexes.
                # if the recommendation cannot be re-parsed into yaml, we assume
                # it is invalid.
                if not validate_yaml(recommendation['index']):
                    recommendation = None
                    query_analysis['supported'] = False


        # QUERY REPORT
        return OrderedDict({
            'queryMask': parsed_query['queryMask'],
            'indexStatus': indexStatus,
            'parsed': parsed_query,
            'namespace': namespace,
            'queryAnalysis': query_analysis,
            'indexAnalysis': index_analysis,
            'recommendation': recommendation
        })
578458
def _generate_recommendation(self,
                                 query_analysis,
                                 db_name,
                                 collection_name):
        """"""
        index_rec = '{'
        for query_field in query_analysis['analyzedFields']:
            if query_field['fieldType'] is EQUIV_TYPE:
                if len(index_rec) is not 1:
                    index_rec += ', '
                index_rec += '"' + query_field['fieldName'] + '": 1'
        for query_field in query_analysis['analyzedFields']:
            if query_field['fieldType'] is SORT_TYPE:
                if len(index_rec) is not 1:
                    index_rec += ', '
                index_rec += '"' + query_field['fieldName'] + '": 1'
        for query_field in query_analysis['analyzedFields']:
            if query_field['fieldType'] is RANGE_TYPE:
                if len(index_rec) is not 1:
                    index_rec += ', '
                index_rec += '"' + query_field['fieldName'] + '": 1'
        index_rec += '}'

        # RECOMMENDATION
        return OrderedDict([('index',index_rec),
                            ('shellCommand', self.generate_shell_command(collection_name, index_rec))])
1127516
def params_to_mongo(query_params):
    
    if not query_params:
        return {}
    for key, value in query_params.items():
        if isinstance(value, list):
            query_params[key] = {'$in': value}
    return query_params
1017729
def _find_matching_collections_internally(collections, record):
    
    for name, data in iteritems(collections):
        if _build_query(data['query']).match(record):
            yield data['ancestors']
    raise StopIteration
1310989
def mget(self, ids, index=None, doc_type=None, **query_params):
        
        if not ids:
            return []

        body = []
        for value in ids:
            if isinstance(value, tuple):
                if len(value) == 3:
                    a, b, c = value
                    body.append({"_index": a,
                                 "_type": b,
                                 "_id": c})
                elif len(value) == 4:
                    a, b, c, d = value
                    body.append({"_index": a,
                                 "_type": b,
                                 "_id": c,
                                 "fields": d})

            else:
                if index is None:
                    raise InvalidQuery("index value is required for id")
                if doc_type is None:
                    raise InvalidQuery("doc_type value is required for id")
                body.append({"_index": index,
                             "_type": doc_type,
                             "_id": value})

        results = self._send_request('GET', "/_mget", body={'docs': body},
                                     params=query_params)
        if 'docs' in results:
            model = self.model
            return [model(self, item) for item in results['docs']]
        return []
690860
function(db) {

		function get_documents(data, config, plv8, ERROR, WARNING) {

			function get_document(id, fields) {
				fields = (fields || [{'query':'*'}]);

				var map = {};

				var obj = plv8.execute("SELECT "+fields.map(function(f, i) {
					var k;
					if(f.key) {
						k = 'p__' + i;
						map[k] = f;
						return f.query + ' AS ' + k;
					} else {
						return f.query;
					}
				}).join(', ')+" FROM documents WHERE id = $1 LIMIT 1", [id])[0];

				if(!obj) {
					return obj;
				}

				Object.keys(obj).forEach(function(key) {
					if(!map.hasOwnProperty(key)) {
						return;
					}

					var spec = map[key];
					if(!is_object(spec)) {
						return error('mapped resource missing');
					}
					if(!is_string(spec.key)) {
						return error('key missing');
					}
					var value = obj[key];
					delete obj[key];

					if(is_string(spec.datakey)) {
						if(!obj[spec.datakey]) {
							obj[spec.datakey] = {};
						}
						obj[spec.datakey][spec.key] = value;
					} else {
						obj[spec.key] = value;
					}
				});

				return obj;
			}

			//function get_type(id) {
			//	return plv8.execute("SELECT * FROM types WHERE id = $1 LIMIT 1", [id])[0];
			//}

			function error(msg) {
				plv8.elog(ERROR, msg);
			}

			function warn(msg) {
				plv8.elog(WARNING, msg);
			}

			function is_object(a) {
				return a && (typeof a === 'object');
			}

			function is_array(a) {
				return a && (a instanceof Array);
			}

			function is_string(a) {
				return a && (typeof a === 'string');
			}

			function is_uuid(obj) {
				return /^[a-fA-F0-9]{8}-[a-fA-F0-9]{4}-[a-fA-F0-9]{4}-[a-fA-F0-9]{4}-[a-fA-F0-9]{12}$/.test(obj);
			}

			function get_property_step(last, current) {
				return is_object(last) ? last[current] : undefined;
			}

			/** Get value of property from provided object using a path.
			 * @param obj {object} The object from where to find the value.
			 * @param path {string} The path to the value in object.
			 * @returns {mixed} The value from the object at the end of path or `undefined` if any part is missing.
			 */
			function get_property(obj, path) {
				if(!is_object(obj)) { return error('get_property(obj, ...) not object: '+ obj); }
				if(!is_string(path)) { return error('get_property(..., path) not string: '+ path); }
				return path.split('.').reduce(get_property_step, obj);
			}

			/** Get document by UUID */
			function fetch_object_by_uuid(data, prop, fields, uuid) {
				if(!is_object(data)) { return error('fetch_object_by_uuid(data, ..., ...) not object: '+ data); }
				if(!is_string(prop)) { return error('fetch_object_by_uuid(..., prop, ...) not string: '+ prop); }
				if(!is_uuid(uuid)) { return warn('Property ' + prop + ' was not uuid: ' + uuid); }
				if(data.documents[uuid] === undefined) {
					data.documents[uuid] = get_document(uuid, fields);
				} else {
					return warn('Document already fetched: ' + uuid);
				}
			}

			/* */
			function fetch_object(data, config) {
				var prop, fields;

				if(is_object(config)) {
					prop = config.prop;
					fields = config.fields;
				} else if(is_string(config)) {
					prop = config;
				}

				if(!is_string(prop)) {
					return error('property is not valid: ' + prop);
				}

				fields = fields ? fields : [{'query':'*'}];

				if(!is_array(fields)) {
					return error('fields are not valid: ' + fields);
				}

				var uuid = get_property(data, prop);

				if(is_array(uuid)) {
					// FIXME: This could be implemented with one query
					return uuid.forEach(fetch_object_by_uuid.bind(undefined, data, prop, fields));
				}
				return fetch_object_by_uuid(data, prop, fields, uuid);
			}

			// Check for bad input

			if(!is_object(data)) { return error("get_documents(data, ...) not object"); }
			if(!is_array(config)) { return error("get_documents(..., config) not array"); }

			//var types_id = data.types_id;
			//if(!is_uuid(types_id)) { return error("data.types_id not valid UUID: " + types_id); }

			// Get document type
			//var type = get_type(types_id);
			//if(!is_object(type)) { return error("type not valid object: " + type); }

			//var relations = type.relations || {};
			//if(!is_object(relations)) { return error("type.relations not valid object: " + relations); }

			// Populate documents from `config`
			data.documents = {};
			config.forEach( fetch_object.bind(undefined, data) );
			return data.documents;
		}

		return db.query('CREATE OR REPLACE FUNCTION get_documents(data json, config json) RETURNS json LANGUAGE plv8 VOLATILE AS ' + NoPg._escapeFunction(get_documents, ["data", "config", "plv8", "ERROR", "WARNING"]));
	}
1210921
def get_querysets_for_context_hints(context_hints_per_source):
    
    model_select_relateds = defaultdict(set)
    model_prefetch_relateds = defaultdict(set)
    model_querysets = {}
    for context_hints in context_hints_per_source.values():
        for hints in context_hints.values():
            model = get_model(hints['app_name'], hints['model_name'])
            model_querysets[model] = model.objects
            model_select_relateds[model].update(hints.get('select_related', []))
            model_prefetch_relateds[model].update(hints.get('prefetch_related', []))

    # Attach select and prefetch related parameters to the querysets if needed
    for model, queryset in model_querysets.items():
        if model_select_relateds[model]:
            queryset = queryset.select_related(*model_select_relateds[model])
        if model_prefetch_relateds[model]:
            queryset = queryset.prefetch_related(*model_prefetch_relateds[model])
        model_querysets[model] = queryset

    return model_querysets