ChangeRevisionMongoDBAccess.insert() - Code Metrics - Inspection of "Add database model to support coordinated write" - StackStorm/st2 - Measure and Improve Code Quality continuously with Scrutinizer

Test Failed

Pull Request — master (#4023)

by W

created 2018-03-06 08:28 UTC

ChangeRevisionMongoDBAccess.insert() A

↳ Parent: ChangeRevisionMongoDBAccess

Complexity

Conditions

Size

Total Lines

Duplication

Lines	0
Ratio	0 %

Importance

Changes

Metric	Value
cc	1
c	0
b	0
f	0
dl	0
loc	5
rs	9.4285

# Licensed to the StackStorm, Inc ('StackStorm') 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 absolute_import

import copy
import importlib
import traceback
import ssl as ssl_lib

import six
import mongoengine
from mongoengine.queryset import visitor
from pymongo import uri_parser
from pymongo.errors import OperationFailure

from st2common import log as logging
from st2common.util import isotime
from st2common.models.db import stormbase
from st2common.models.utils.profiling import log_query_and_profile_data_for_queryset
from st2common.exceptions.db import StackStormDBObjectNotFoundError


LOG = logging.getLogger(__name__)

MODEL_MODULE_NAMES = [
    'st2common.models.db.auth',
    'st2common.models.db.action',
    'st2common.models.db.actionalias',
    'st2common.models.db.keyvalue',
    'st2common.models.db.execution',
    'st2common.models.db.executionstate',
    'st2common.models.db.liveaction',
    'st2common.models.db.notification',
    'st2common.models.db.pack',
    'st2common.models.db.policy',
    'st2common.models.db.rbac',
    'st2common.models.db.rule',
    'st2common.models.db.rule_enforcement',
    'st2common.models.db.runner',
    'st2common.models.db.sensor',
    'st2common.models.db.trace',
    'st2common.models.db.trigger',
    'st2common.models.db.webhook'
]

# A list of model names for which we don't perform extra index cleanup
INDEX_CLEANUP_MODEL_NAMES_BLACKLIST = [
    'PermissionGrantDB'
]


def get_model_classes():
    """
    Retrieve a list of all the defined model classes.

    :rtype: ``list``
    """
    result = []
    for module_name in MODEL_MODULE_NAMES:
        module = importlib.import_module(module_name)
        model_classes = getattr(module, 'MODELS', [])
        result.extend(model_classes)

    return result


def _db_connect(db_name, db_host, db_port, username=None, password=None,
             ssl=False, ssl_keyfile=None, ssl_certfile=None,
             ssl_cert_reqs=None, ssl_ca_certs=None, ssl_match_hostname=True):

    if '://' in db_host:
        # Hostname is provided as a URI string. Make sure we don't log the password in case one is
        # included as part of the URI string.
        uri_dict = uri_parser.parse_uri(db_host)
        username_string = uri_dict.get('username', username) or username

        if uri_dict.get('username', None) and username:
            # Username argument has precedence over connection string username
            username_string = username

        hostnames = get_host_names_for_uri_dict(uri_dict=uri_dict)

        if len(uri_dict['nodelist']) > 1:
            host_string = '%s (replica set)' % (hostnames)
        else:
            host_string = hostnames
    else:
        host_string = '%s:%s' % (db_host, db_port)
        username_string = username

    LOG.info('Connecting to database "%s" @ "%s" as user "%s".' % (db_name, host_string,
                                                                   str(username_string)))

    ssl_kwargs = _get_ssl_kwargs(ssl=ssl, ssl_keyfile=ssl_keyfile, ssl_certfile=ssl_certfile,
                                 ssl_cert_reqs=ssl_cert_reqs, ssl_ca_certs=ssl_ca_certs,
                                 ssl_match_hostname=ssl_match_hostname)

    connection = mongoengine.connection.connect(db_name, host=db_host,
                                                port=db_port, tz_aware=True,
                                                username=username, password=password,
                                                **ssl_kwargs)
    return connection


def db_setup(db_name, db_host, db_port, username=None, password=None, ensure_indexes=True,
             ssl=False, ssl_keyfile=None, ssl_certfile=None,
             ssl_cert_reqs=None, ssl_ca_certs=None, ssl_match_hostname=True):

    connection = _db_connect(db_name, db_host, db_port, username=username,
                             password=password, ssl=ssl, ssl_keyfile=ssl_keyfile,
                             ssl_certfile=ssl_certfile,
                             ssl_cert_reqs=ssl_cert_reqs, ssl_ca_certs=ssl_ca_certs,
                             ssl_match_hostname=ssl_match_hostname)

    # Create all the indexes upfront to prevent race-conditions caused by
    # lazy index creation
    if ensure_indexes:
        db_ensure_indexes()

    return connection


def db_ensure_indexes():
    """
    This function ensures that indexes for all the models have been created and the
    extra indexes cleaned up.

    Note #1: When calling this method database connection already needs to be
    established.

    Note #2: This method blocks until all the index have been created (indexes
    are created in real-time and not in background).
    """
    LOG.debug('Ensuring database indexes...')
    model_classes = get_model_classes()

    for model_class in model_classes:
        class_name = model_class.__name__

        # Note: We need to ensure / create new indexes before removing extra ones
        try:
            model_class.ensure_indexes()
        except OperationFailure as e:
            # Special case for "uid" index. MongoDB 3.4 has dropped "_types" index option so we
            # need to re-create the index to make it work and avoid "index with different options
            # already exists" error.
            # Note: This condition would only be encountered when upgrading existing StackStorm
            # installation from MongoDB 3.2 to 3.4.
            msg = str(e)
            if 'already exists with different options' in msg and 'uid_1' in msg:
                drop_obsolete_types_indexes(model_class=model_class)
            else:
                raise e
        except Exception as e:
            tb_msg = traceback.format_exc()
            msg = 'Failed to ensure indexes for model "%s": %s' % (class_name, str(e))
            msg += '\n\n' + tb_msg
            exc_cls = type(e)
            raise exc_cls(msg)

        if model_class.__name__ in INDEX_CLEANUP_MODEL_NAMES_BLACKLIST:
            LOG.debug('Skipping index cleanup for blacklisted model "%s"...' % (class_name))
            continue

        removed_count = cleanup_extra_indexes(model_class=model_class)
        if removed_count:
            LOG.debug('Removed "%s" extra indexes for model "%s"' % (removed_count, class_name))

    LOG.debug('Indexes are ensured for models: %s' %
              ', '.join(sorted((model_class.__name__ for model_class in model_classes))))


def cleanup_extra_indexes(model_class):
    """
    Finds any extra indexes and removes those from mongodb.
    """
    extra_indexes = model_class.compare_indexes().get('extra', None)
    if not extra_indexes:
        return 0

    # mongoengine does not have the necessary method so we need to drop to
    # pymongo interfaces via some private methods.
    removed_count = 0
    c = model_class._get_collection()
    for extra_index in extra_indexes:
        try:
            c.drop_index(extra_index)
            LOG.debug('Dropped index %s for model %s.', extra_index, model_class.__name__)
            removed_count += 1
        except OperationFailure:
            LOG.warning('Attempt to cleanup index %s failed.', extra_index, exc_info=True)

    return removed_count


def drop_obsolete_types_indexes(model_class):
    """
    Special class for droping offending "types" indexes for which support has
    been removed in mongoengine and MongoDB 3.4.
    For more info, see: http://docs.mongoengine.org/upgrade.html#inheritance
    """
    class_name = model_class.__name__

    LOG.debug('Dropping obsolete types index for model "%s"' % (class_name))
    collection = model_class._get_collection()
    collection.update({}, {'$unset': {'_types': 1}}, multi=True)

    info = collection.index_information()
    indexes_to_drop = [key for key, value in six.iteritems(info)
                       if '_types' in dict(value['key']) or 'types' in value]

    LOG.debug('Will drop obsolete types indexes for model "%s": %s' % (class_name,
                                                                       str(indexes_to_drop)))

    for index in indexes_to_drop:
        collection.drop_index(index)

    LOG.debug('Recreating indexes for model "%s"' % (class_name))
    model_class.ensure_indexes()


def db_teardown():
    mongoengine.connection.disconnect()


def db_cleanup(db_name, db_host, db_port, username=None, password=None,
               ssl=False, ssl_keyfile=None, ssl_certfile=None,
               ssl_cert_reqs=None, ssl_ca_certs=None, ssl_match_hostname=True):

    connection = _db_connect(db_name, db_host, db_port, username=username,
                             password=password, ssl=ssl, ssl_keyfile=ssl_keyfile,
                             ssl_certfile=ssl_certfile,
                             ssl_cert_reqs=ssl_cert_reqs, ssl_ca_certs=ssl_ca_certs,
                             ssl_match_hostname=ssl_match_hostname)

    LOG.info('Dropping database "%s" @ "%s:%s" as user "%s".',
             db_name, db_host, db_port, str(username))

    connection.drop_database(db_name)
    return connection


def _get_ssl_kwargs(ssl=False, ssl_keyfile=None, ssl_certfile=None, ssl_cert_reqs=None,
                    ssl_ca_certs=None, ssl_match_hostname=True):
    ssl_kwargs = {
        'ssl': ssl,
    }
    if ssl_keyfile:
        ssl_kwargs['ssl'] = True
        ssl_kwargs['ssl_keyfile'] = ssl_keyfile
    if ssl_certfile:
        ssl_kwargs['ssl'] = True
        ssl_kwargs['ssl_certfile'] = ssl_certfile
    if ssl_cert_reqs:
        if ssl_cert_reqs is 'none':
            ssl_cert_reqs = ssl_lib.CERT_NONE
        elif ssl_cert_reqs is 'optional':
            ssl_cert_reqs = ssl_lib.CERT_OPTIONAL
        elif ssl_cert_reqs is 'required':
            ssl_cert_reqs = ssl_lib.CERT_REQUIRED
        ssl_kwargs['ssl_cert_reqs'] = ssl_cert_reqs
    if ssl_ca_certs:
        ssl_kwargs['ssl'] = True
        ssl_kwargs['ssl_ca_certs'] = ssl_ca_certs
    if ssl_kwargs.get('ssl', False):
        # pass in ssl_match_hostname only if ssl is True. The right default value
        # for ssl_match_hostname in almost all cases is True.
        ssl_kwargs['ssl_match_hostname'] = ssl_match_hostname
    return ssl_kwargs


class MongoDBAccess(object):
    """Database object access class that provides general functions for a model type."""

    def __init__(self, model):
        self.model = model

    def get_by_name(self, value):
        return self.get(name=value, raise_exception=True)

    def get_by_id(self, value):
        return self.get(id=value, raise_exception=True)

    def get_by_uid(self, value):
        return self.get(uid=value, raise_exception=True)

    def get_by_ref(self, value):
        return self.get(ref=value, raise_exception=True)

    def get_by_pack(self, value):
        return self.get(pack=value, raise_exception=True)

    def get(self, *args, **kwargs):
        exclude_fields = kwargs.pop('exclude_fields', None)
        raise_exception = kwargs.pop('raise_exception', False)
        only_fields = kwargs.pop('only_fields', None)

        args = self._process_arg_filters(args)

        instances = self.model.objects(*args, **kwargs)

        if exclude_fields:
            instances = instances.exclude(*exclude_fields)

        if only_fields:
            try:
                instances = instances.only(*only_fields)
            except mongoengine.errors.LookUpError as e:
                msg = ('Invalid or unsupported include attribute specified: %s' % str(e))
                raise ValueError(msg)

        instance = instances[0] if instances else None
        log_query_and_profile_data_for_queryset(queryset=instances)

        if not instance and raise_exception:
            msg = 'Unable to find the %s instance. %s' % (self.model.__name__, kwargs)
            raise StackStormDBObjectNotFoundError(msg)

        return instance

    def get_all(self, *args, **kwargs):
        return self.query(*args, **kwargs)

    def count(self, *args, **kwargs):
        result = self.model.objects(*args, **kwargs).count()
        log_query_and_profile_data_for_queryset(queryset=result)
        return result

    # TODO: PEP-3102 introduced keyword-only arguments, so once we support Python 3+, we can change
    #       this definition to have explicit keyword-only arguments:
    #
    #           def query(self, *args, offset=0, limit=None, order_by=None, exclude_fields=None,
    #                     **filters):
    def query(self, *args, **filters):
        # Python 2: Pop keyword parameters that aren't actually filters off of the kwargs
        offset = filters.pop('offset', 0)
        limit = filters.pop('limit', None)
        order_by = filters.pop('order_by', None)
        exclude_fields = filters.pop('exclude_fields', None)
        only_fields = filters.pop('only_fields', None)
        no_dereference = filters.pop('no_dereference', None)

        order_by = order_by or []
        exclude_fields = exclude_fields or []
        eop = offset + int(limit) if limit else None

        args = self._process_arg_filters(args)
        # Process the filters
        # Note: Both of those functions manipulate "filters" variable so the order in which they
        # are called matters
        filters, order_by = self._process_datetime_range_filters(filters=filters, order_by=order_by)
        filters = self._process_null_filters(filters=filters)

        result = self.model.objects(*args, **filters)

        if exclude_fields:
            result = result.exclude(*exclude_fields)

        if only_fields:
            try:
                result = result.only(*only_fields)
            except mongoengine.errors.LookUpError as e:
                msg = ('Invalid or unsupported include attribute specified: %s' % str(e))
                raise ValueError(msg)

        if no_dereference:
            result = result.no_dereference()

        result = result.order_by(*order_by)
        result = result[offset:eop]
        log_query_and_profile_data_for_queryset(queryset=result)

        return result

    def distinct(self, *args, **kwargs):
        field = kwargs.pop('field')
        result = self.model.objects(**kwargs).distinct(field)
        log_query_and_profile_data_for_queryset(queryset=result)
        return result

    def aggregate(self, *args, **kwargs):
        return self.model.objects(**kwargs)._collection.aggregate(*args, **kwargs)

    def insert(self, instance):
        instance = self.model.objects.insert(instance)
        return self._undo_dict_field_escape(instance)

    def add_or_update(self, instance):
        instance.save()
        return self._undo_dict_field_escape(instance)

    def update(self, instance, **kwargs):
        return instance.update(**kwargs)

    def delete(self, instance):
        return instance.delete()

    def delete_by_query(self, *args, **query):
        """
        Delete objects by query and return number of deleted objects.
        """
        qs = self.model.objects.filter(*args, **query)
        count = qs.delete()
        log_query_and_profile_data_for_queryset(queryset=qs)

        return count

    def _undo_dict_field_escape(self, instance):
        for attr, field in six.iteritems(instance._fields):
            if isinstance(field, stormbase.EscapedDictField):
                value = getattr(instance, attr)
                setattr(instance, attr, field.to_python(value))
        return instance

    def _process_arg_filters(self, args):
        """
        Fix filter arguments in nested Q objects
        """
        _args = tuple()

        for arg in args:
            # Unforunately mongoengine doesn't expose any visitors other than Q, so we have to
            # extract QCombination from the module itself
            if isinstance(arg, visitor.Q):
                # Note: Both of those functions manipulate "filters" variable so the order in which
                # they are called matters
                filters, _ = self._process_datetime_range_filters(filters=arg.query)
                filters = self._process_null_filters(filters=filters)

                # Create a new Q object with the same filters as the old one
                _args += (visitor.Q(**filters),)
            elif isinstance(arg, visitor.QCombination):
                # Recurse if we need to
                children = self._process_arg_filters(arg.children)

                # Create a new QCombination object with the same operation and fixed filters
                _args += (visitor.QCombination(arg.operation, children),)
            else:
                raise TypeError("Unknown argument type '%s' of argument '%s'"
                    % (type(arg), repr(arg)))

        return _args

    def _process_null_filters(self, filters):
        result = copy.deepcopy(filters)

        null_filters = {k: v for k, v in six.iteritems(filters)
                        if v is None or
                        (type(v) in [str, six.text_type] and str(v.lower()) == 'null')}

        for key in null_filters.keys():
            result['%s__exists' % (key)] = False
            del result[key]

        return result

    def _process_datetime_range_filters(self, filters, order_by=None):
        ranges = {k: v for k, v in six.iteritems(filters)
                  if type(v) in [str, six.text_type] and '..' in v}

        order_by_list = copy.deepcopy(order_by) if order_by else []
        for k, v in six.iteritems(ranges):
            values = v.split('..')
            dt1 = isotime.parse(values[0])
            dt2 = isotime.parse(values[1])

            k__gte = '%s__gte' % k
            k__lte = '%s__lte' % k
            if dt1 < dt2:
                query = {k__gte: dt1, k__lte: dt2}
                sort_key, reverse_sort_key = k, '-' + k
            else:
                query = {k__gte: dt2, k__lte: dt1}
                sort_key, reverse_sort_key = '-' + k, k
            del filters[k]
            filters.update(query)

            if reverse_sort_key in order_by_list:
                idx = order_by_list.index(reverse_sort_key)
                order_by_list.pop(idx)
                order_by_list.insert(idx, sort_key)
            elif sort_key not in order_by_list:
                order_by_list = [sort_key] + order_by_list

        return filters, order_by_list


class ChangeRevisionMongoDBAccess(MongoDBAccess):

    def get(self, *args, **kwargs):
        instance = super(ChangeRevisionMongoDBAccess, self).get(*args, **kwargs)
        instance._original_rev = instance.rev
        return instance

    def insert(self, instance):
        instance = self.model.objects.insert(instance)
        instance._original_rev = instance.rev

        return self._undo_dict_field_escape(instance)

    def add_or_update(self, instance):
        if not hasattr(instance, '_original_rev'):
            return self.insert(instance)
        else:
            save_condition = {
                'id': instance.id,
                'rev': instance._original_rev
            }

            instance.rev = instance._original_rev + 1
            instance.save(save_condition=save_condition)

            return self._undo_dict_field_escape(instance)

    def update(self, instance, **kwargs):
        for k, v in six.iteritems(kwargs):
            setattr(instance, k, v)

        save_condition = {
            'id': instance.id,
            'rev': instance._original_rev
        }

        instance.rev = instance._original_rev + 1
        instance.save(save_condition=save_condition)

        return self._undo_dict_field_escape(instance)


def get_host_names_for_uri_dict(uri_dict):
    hosts = []

    for host, port in uri_dict['nodelist']:
        hosts.append('%s:%s' % (host, port))

    hosts = ','.join(hosts)
    return hosts


1			# Licensed to the StackStorm, Inc ('StackStorm') under one or more
2			# contributor license agreements. See the NOTICE file distributed with
3			# this work for additional information regarding copyright ownership.
4			# The ASF licenses this file to You under the Apache License, Version 2.0
5			# (the "License"); you may not use this file except in compliance with
6			# the License. You may obtain a copy of the License at
7			#
8			# http://www.apache.org/licenses/LICENSE-2.0
9			#
10			# Unless required by applicable law or agreed to in writing, software
11			# distributed under the License is distributed on an "AS IS" BASIS,
12			# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13			# See the License for the specific language governing permissions and
14			# limitations under the License.
15
16			from __future__ import absolute_import
17
18			import copy
19			import importlib
20			import traceback
21			import ssl as ssl_lib
22
23			import six
24			import mongoengine
25			from mongoengine.queryset import visitor
26			from pymongo import uri_parser
27			from pymongo.errors import OperationFailure
28
29			from st2common import log as logging
30			from st2common.util import isotime
31			from st2common.models.db import stormbase
32			from st2common.models.utils.profiling import log_query_and_profile_data_for_queryset
33			from st2common.exceptions.db import StackStormDBObjectNotFoundError
34
35
36			LOG = logging.getLogger(__name__)
37
38			MODEL_MODULE_NAMES = [
39			'st2common.models.db.auth',
40			'st2common.models.db.action',
41			'st2common.models.db.actionalias',
42			'st2common.models.db.keyvalue',
43			'st2common.models.db.execution',
44			'st2common.models.db.executionstate',
45			'st2common.models.db.liveaction',
46			'st2common.models.db.notification',
47			'st2common.models.db.pack',
48			'st2common.models.db.policy',
49			'st2common.models.db.rbac',
50			'st2common.models.db.rule',
51			'st2common.models.db.rule_enforcement',
52			'st2common.models.db.runner',
53			'st2common.models.db.sensor',
54			'st2common.models.db.trace',
55			'st2common.models.db.trigger',
56			'st2common.models.db.webhook'
57			]
58
59			# A list of model names for which we don't perform extra index cleanup
60			INDEX_CLEANUP_MODEL_NAMES_BLACKLIST = [
61			'PermissionGrantDB'
62			]
63
64
65			def get_model_classes():
66			"""
67			Retrieve a list of all the defined model classes.
68
69			:rtype: ``list``
70			"""
71			result = []
72			for module_name in MODEL_MODULE_NAMES:
73			module = importlib.import_module(module_name)
74			model_classes = getattr(module, 'MODELS', [])
75			result.extend(model_classes)
76
77			return result
78
79
80			def _db_connect(db_name, db_host, db_port, username=None, password=None,
81			ssl=False, ssl_keyfile=None, ssl_certfile=None,
82			ssl_cert_reqs=None, ssl_ca_certs=None, ssl_match_hostname=True):
83
84			if '://' in db_host:
85			# Hostname is provided as a URI string. Make sure we don't log the password in case one is
86			# included as part of the URI string.
87			uri_dict = uri_parser.parse_uri(db_host)
88			username_string = uri_dict.get('username', username) or username
89
90			if uri_dict.get('username', None) and username:
91			# Username argument has precedence over connection string username
92			username_string = username
93
94			hostnames = get_host_names_for_uri_dict(uri_dict=uri_dict)
95
96			if len(uri_dict['nodelist']) > 1:
97			host_string = '%s (replica set)' % (hostnames)
98			else:
99			host_string = hostnames
100			else:
101			host_string = '%s:%s' % (db_host, db_port)
102			username_string = username
103
104			LOG.info('Connecting to database "%s" @ "%s" as user "%s".' % (db_name, host_string,
105			str(username_string)))
106
107			ssl_kwargs = _get_ssl_kwargs(ssl=ssl, ssl_keyfile=ssl_keyfile, ssl_certfile=ssl_certfile,
108			ssl_cert_reqs=ssl_cert_reqs, ssl_ca_certs=ssl_ca_certs,
109			ssl_match_hostname=ssl_match_hostname)
110
111			connection = mongoengine.connection.connect(db_name, host=db_host,
112			port=db_port, tz_aware=True,
113			username=username, password=password,
114			**ssl_kwargs)
115			return connection
116
117
118			def db_setup(db_name, db_host, db_port, username=None, password=None, ensure_indexes=True,
119			ssl=False, ssl_keyfile=None, ssl_certfile=None,
120			ssl_cert_reqs=None, ssl_ca_certs=None, ssl_match_hostname=True):
121
122			connection = _db_connect(db_name, db_host, db_port, username=username,
123			password=password, ssl=ssl, ssl_keyfile=ssl_keyfile,
124			ssl_certfile=ssl_certfile,
125			ssl_cert_reqs=ssl_cert_reqs, ssl_ca_certs=ssl_ca_certs,
126			ssl_match_hostname=ssl_match_hostname)
127
128			# Create all the indexes upfront to prevent race-conditions caused by
129			# lazy index creation
130			if ensure_indexes:
131			db_ensure_indexes()
132
133			return connection
134
135
136			def db_ensure_indexes():
137			"""
138			This function ensures that indexes for all the models have been created and the
139			extra indexes cleaned up.
140
141			Note #1: When calling this method database connection already needs to be
142			established.
143
144			Note #2: This method blocks until all the index have been created (indexes
145			are created in real-time and not in background).
146			"""
147			LOG.debug('Ensuring database indexes...')
148			model_classes = get_model_classes()
149
150			for model_class in model_classes:
151			class_name = model_class.__name__
152
153			# Note: We need to ensure / create new indexes before removing extra ones
154			try:
155			model_class.ensure_indexes()
156			except OperationFailure as e:
157			# Special case for "uid" index. MongoDB 3.4 has dropped "_types" index option so we
158			# need to re-create the index to make it work and avoid "index with different options
159			# already exists" error.
160			# Note: This condition would only be encountered when upgrading existing StackStorm
161			# installation from MongoDB 3.2 to 3.4.
162			msg = str(e)
163			if 'already exists with different options' in msg and 'uid_1' in msg:
164			drop_obsolete_types_indexes(model_class=model_class)
165			else:
166			raise e
167			except Exception as e:
168			tb_msg = traceback.format_exc()
169			msg = 'Failed to ensure indexes for model "%s": %s' % (class_name, str(e))
170			msg += '\n\n' + tb_msg
171			exc_cls = type(e)
172			raise exc_cls(msg)
173
174			if model_class.__name__ in INDEX_CLEANUP_MODEL_NAMES_BLACKLIST:
175			LOG.debug('Skipping index cleanup for blacklisted model "%s"...' % (class_name))
176			continue
177
178			removed_count = cleanup_extra_indexes(model_class=model_class)
179			if removed_count:
180			LOG.debug('Removed "%s" extra indexes for model "%s"' % (removed_count, class_name))
181
182			LOG.debug('Indexes are ensured for models: %s' %
183			', '.join(sorted((model_class.__name__ for model_class in model_classes))))
184
185
186			def cleanup_extra_indexes(model_class):
187			"""
188			Finds any extra indexes and removes those from mongodb.
189			"""
190			extra_indexes = model_class.compare_indexes().get('extra', None)
191			if not extra_indexes:
192			return 0
193
194			# mongoengine does not have the necessary method so we need to drop to
195			# pymongo interfaces via some private methods.
196			removed_count = 0
197			c = model_class._get_collection()
198			for extra_index in extra_indexes:
199			try:
200			c.drop_index(extra_index)
201			LOG.debug('Dropped index %s for model %s.', extra_index, model_class.__name__)
202			removed_count += 1
203			except OperationFailure:
204			LOG.warning('Attempt to cleanup index %s failed.', extra_index, exc_info=True)
205
206			return removed_count
207
208
209			def drop_obsolete_types_indexes(model_class):
210			"""
211			Special class for droping offending "types" indexes for which support has
212			been removed in mongoengine and MongoDB 3.4.
213			For more info, see: http://docs.mongoengine.org/upgrade.html#inheritance
214			"""
215			class_name = model_class.__name__
216
217			LOG.debug('Dropping obsolete types index for model "%s"' % (class_name))
218			collection = model_class._get_collection()
219			collection.update({}, {'$unset': {'_types': 1}}, multi=True)
220
221			info = collection.index_information()
222			indexes_to_drop = [key for key, value in six.iteritems(info)
223			if '_types' in dict(value['key']) or 'types' in value]
224
225			LOG.debug('Will drop obsolete types indexes for model "%s": %s' % (class_name,
226			str(indexes_to_drop)))
227
228			for index in indexes_to_drop:
229			collection.drop_index(index)
230
231			LOG.debug('Recreating indexes for model "%s"' % (class_name))
232			model_class.ensure_indexes()
233
234
235			def db_teardown():
236			mongoengine.connection.disconnect()
237
238
239			def db_cleanup(db_name, db_host, db_port, username=None, password=None,
240			ssl=False, ssl_keyfile=None, ssl_certfile=None,
241			ssl_cert_reqs=None, ssl_ca_certs=None, ssl_match_hostname=True):
242
243			connection = _db_connect(db_name, db_host, db_port, username=username,
244			password=password, ssl=ssl, ssl_keyfile=ssl_keyfile,
245			ssl_certfile=ssl_certfile,
246			ssl_cert_reqs=ssl_cert_reqs, ssl_ca_certs=ssl_ca_certs,
247			ssl_match_hostname=ssl_match_hostname)
248
249			LOG.info('Dropping database "%s" @ "%s:%s" as user "%s".',
250			db_name, db_host, db_port, str(username))
251
252			connection.drop_database(db_name)
253			return connection
254
255
256			def _get_ssl_kwargs(ssl=False, ssl_keyfile=None, ssl_certfile=None, ssl_cert_reqs=None,
257			ssl_ca_certs=None, ssl_match_hostname=True):
258			ssl_kwargs = {
259			'ssl': ssl,
260			}
261			if ssl_keyfile:
262			ssl_kwargs['ssl'] = True
263			ssl_kwargs['ssl_keyfile'] = ssl_keyfile
264			if ssl_certfile:
265			ssl_kwargs['ssl'] = True
266			ssl_kwargs['ssl_certfile'] = ssl_certfile
267			if ssl_cert_reqs:
268			if ssl_cert_reqs is 'none':
269			ssl_cert_reqs = ssl_lib.CERT_NONE
270			elif ssl_cert_reqs is 'optional':
271			ssl_cert_reqs = ssl_lib.CERT_OPTIONAL
272			elif ssl_cert_reqs is 'required':
273			ssl_cert_reqs = ssl_lib.CERT_REQUIRED
274			ssl_kwargs['ssl_cert_reqs'] = ssl_cert_reqs
275			if ssl_ca_certs:
276			ssl_kwargs['ssl'] = True
277			ssl_kwargs['ssl_ca_certs'] = ssl_ca_certs
278			if ssl_kwargs.get('ssl', False):
279			# pass in ssl_match_hostname only if ssl is True. The right default value
280			# for ssl_match_hostname in almost all cases is True.
281			ssl_kwargs['ssl_match_hostname'] = ssl_match_hostname
282			return ssl_kwargs
283
284
285			class MongoDBAccess(object):
286			"""Database object access class that provides general functions for a model type."""
287
288			def __init__(self, model):
289			self.model = model
290
291			def get_by_name(self, value):
292			return self.get(name=value, raise_exception=True)
293
294			def get_by_id(self, value):
295			return self.get(id=value, raise_exception=True)
296
297			def get_by_uid(self, value):
298			return self.get(uid=value, raise_exception=True)
299
300			def get_by_ref(self, value):
301			return self.get(ref=value, raise_exception=True)
302
303			def get_by_pack(self, value):
304			return self.get(pack=value, raise_exception=True)
305
306			def get(self, args, *kwargs):
307			exclude_fields = kwargs.pop('exclude_fields', None)
308			raise_exception = kwargs.pop('raise_exception', False)
309			only_fields = kwargs.pop('only_fields', None)
310
311			args = self._process_arg_filters(args)
312
313			instances = self.model.objects(args, *kwargs)
314
315			if exclude_fields:
316			instances = instances.exclude(*exclude_fields)
317
318			if only_fields:
319			try:
320			instances = instances.only(*only_fields)
321			except mongoengine.errors.LookUpError as e:
322			msg = ('Invalid or unsupported include attribute specified: %s' % str(e))
323			raise ValueError(msg)
324
325			instance = instances[0] if instances else None
326			log_query_and_profile_data_for_queryset(queryset=instances)
327
328			if not instance and raise_exception:
329			msg = 'Unable to find the %s instance. %s' % (self.model.__name__, kwargs)
330			raise StackStormDBObjectNotFoundError(msg)
331
332			return instance
333
334			def get_all(self, args, *kwargs):
335			return self.query(args, *kwargs)
336
337			def count(self, args, *kwargs):
338			result = self.model.objects(args, *kwargs).count()
339			log_query_and_profile_data_for_queryset(queryset=result)
340			return result
341
342			# TODO: PEP-3102 introduced keyword-only arguments, so once we support Python 3+, we can change
343			# this definition to have explicit keyword-only arguments:
344			#
345			# def query(self, *args, offset=0, limit=None, order_by=None, exclude_fields=None,
346			# **filters):
347			def query(self, args, *filters):
348			# Python 2: Pop keyword parameters that aren't actually filters off of the kwargs
349			offset = filters.pop('offset', 0)
350			limit = filters.pop('limit', None)
351			order_by = filters.pop('order_by', None)
352			exclude_fields = filters.pop('exclude_fields', None)
353			only_fields = filters.pop('only_fields', None)
354			no_dereference = filters.pop('no_dereference', None)
355
356			order_by = order_by or []
357			exclude_fields = exclude_fields or []
358			eop = offset + int(limit) if limit else None
359
360			args = self._process_arg_filters(args)
361			# Process the filters
362			# Note: Both of those functions manipulate "filters" variable so the order in which they
363			# are called matters
364			filters, order_by = self._process_datetime_range_filters(filters=filters, order_by=order_by)
365			filters = self._process_null_filters(filters=filters)
366
367			result = self.model.objects(args, *filters)
368
369			if exclude_fields:
370			result = result.exclude(*exclude_fields)
371
372			if only_fields:
373			try:
374			result = result.only(*only_fields)
375			except mongoengine.errors.LookUpError as e:
376			msg = ('Invalid or unsupported include attribute specified: %s' % str(e))
377			raise ValueError(msg)
378
379			if no_dereference:
380			result = result.no_dereference()
381
382			result = result.order_by(*order_by)
383			result = result[offset:eop]
384			log_query_and_profile_data_for_queryset(queryset=result)
385
386			return result
387
388			def distinct(self, args, *kwargs):
389			field = kwargs.pop('field')
390			result = self.model.objects(**kwargs).distinct(field)
391			log_query_and_profile_data_for_queryset(queryset=result)
392			return result
393
394			def aggregate(self, args, *kwargs):
395			return self.model.objects(*kwargs)._collection.aggregate(args, **kwargs)
396
397			def insert(self, instance):
398			instance = self.model.objects.insert(instance)
399			return self._undo_dict_field_escape(instance)
400
401			def add_or_update(self, instance):
402			instance.save()
403			return self._undo_dict_field_escape(instance)
404
405			def update(self, instance, **kwargs):
406			return instance.update(**kwargs)
407
408			def delete(self, instance):
409			return instance.delete()
410
411			def delete_by_query(self, args, *query):
412			"""
413			Delete objects by query and return number of deleted objects.
414			"""
415			qs = self.model.objects.filter(args, *query)
416			count = qs.delete()
417			log_query_and_profile_data_for_queryset(queryset=qs)
418
419			return count
420
421			def _undo_dict_field_escape(self, instance):
422			for attr, field in six.iteritems(instance._fields):
423			if isinstance(field, stormbase.EscapedDictField):
424			value = getattr(instance, attr)
425			setattr(instance, attr, field.to_python(value))
426			return instance
427
428			def _process_arg_filters(self, args):
429			"""
430			Fix filter arguments in nested Q objects
431			"""
432			_args = tuple()
433
434			for arg in args:
435			# Unforunately mongoengine doesn't expose any visitors other than Q, so we have to
436			# extract QCombination from the module itself
437			if isinstance(arg, visitor.Q):
438			# Note: Both of those functions manipulate "filters" variable so the order in which
439			# they are called matters
440			filters, _ = self._process_datetime_range_filters(filters=arg.query)
441			filters = self._process_null_filters(filters=filters)
442
443			# Create a new Q object with the same filters as the old one
444			_args += (visitor.Q(**filters),)
445			elif isinstance(arg, visitor.QCombination):
446			# Recurse if we need to
447			children = self._process_arg_filters(arg.children)
448
449			# Create a new QCombination object with the same operation and fixed filters
450			_args += (visitor.QCombination(arg.operation, children),)
451			else:
452			raise TypeError("Unknown argument type '%s' of argument '%s'"
453			% (type(arg), repr(arg)))
454
455			return _args
456
457			def _process_null_filters(self, filters):
458			result = copy.deepcopy(filters)
459
460			null_filters = {k: v for k, v in six.iteritems(filters)
461			if v is None or
462			(type(v) in [str, six.text_type] and str(v.lower()) == 'null')}
463
464			for key in null_filters.keys():
465			result['%s__exists' % (key)] = False
466			del result[key]
467
468			return result
469
470			def _process_datetime_range_filters(self, filters, order_by=None):
471			ranges = {k: v for k, v in six.iteritems(filters)
472			if type(v) in [str, six.text_type] and '..' in v}
473
474			order_by_list = copy.deepcopy(order_by) if order_by else []
475			for k, v in six.iteritems(ranges):
476			values = v.split('..')
477			dt1 = isotime.parse(values[0])
478			dt2 = isotime.parse(values[1])
479
480			k__gte = '%s__gte' % k
481			k__lte = '%s__lte' % k
482			if dt1 < dt2:
483			query = {k__gte: dt1, k__lte: dt2}
484			sort_key, reverse_sort_key = k, '-' + k
485			else:
486			query = {k__gte: dt2, k__lte: dt1}
487			sort_key, reverse_sort_key = '-' + k, k
488			del filters[k]
489			filters.update(query)
490
491			if reverse_sort_key in order_by_list:
492			idx = order_by_list.index(reverse_sort_key)
493			order_by_list.pop(idx)
494			order_by_list.insert(idx, sort_key)
495			elif sort_key not in order_by_list:
496			order_by_list = [sort_key] + order_by_list
497
498			return filters, order_by_list
499
500
501			class ChangeRevisionMongoDBAccess(MongoDBAccess):
502
503			def get(self, args, *kwargs):
504			instance = super(ChangeRevisionMongoDBAccess, self).get(args, *kwargs)
505			instance._original_rev = instance.rev
506			return instance
507
508			def insert(self, instance):
509			instance = self.model.objects.insert(instance)
510			instance._original_rev = instance.rev
511
512			return self._undo_dict_field_escape(instance)
513
514			def add_or_update(self, instance):
515			if not hasattr(instance, '_original_rev'):
516			return self.insert(instance)
517			else:
518			save_condition = {
519			'id': instance.id,
520			'rev': instance._original_rev
521			}
522
523			instance.rev = instance._original_rev + 1
524			instance.save(save_condition=save_condition)
525
526			return self._undo_dict_field_escape(instance)
527
528			def update(self, instance, **kwargs):
529			for k, v in six.iteritems(kwargs):
530			setattr(instance, k, v)
531
532			save_condition = {
533			'id': instance.id,
534			'rev': instance._original_rev
535			}
536
537			instance.rev = instance._original_rev + 1
538			instance.save(save_condition=save_condition)
539
540			return self._undo_dict_field_escape(instance)
541
542
543			def get_host_names_for_uri_dict(uri_dict):
544			hosts = []
545
546			for host, port in uri_dict['nodelist']:
547			hosts.append('%s:%s' % (host, port))
548
549			hosts = ','.join(hosts)
550			return hosts
551

StackStorm / st2

Pull Request — master (#4023)

ChangeRevisionMongoDBAccess.insert() A

Complexity

Size

Duplication

Importance

Duplication Side-by-Side

Filter issues like