<commit_before># -*- coding: utf-8 -*-
"""
    Copyright 2013-2014 Olivier Cortès <oc@1flow.io>

    This file is part of the 1flow project.

    1flow is free software: you can redistribute it and/or modify
    it under the terms of the GNU Affero General Public License as
    published by the Free Software Foundation, either version 3 of
    the License, or (at your option) any later version.

    1flow is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU Affero General Public License for more details.

    You should have received a copy of the GNU Affero General Public
    License along with 1flow.  If not, see http://www.gnu.org/licenses/

"""
import logging

import requests

from statsd import statsd
from constance import config

# from json_field import JSONField
# from celery import chain as tasks_chain

# from django.conf import settings
from django.db import models, IntegrityError
from django.utils.translation import ugettext_lazy as _
# from django.utils.text import slugify

from oneflow.base.utils import register_task_method
from oneflow.base.utils.http import clean_url
# from oneflow.base.utils.dateutils import now

from ...common import REQUEST_BASE_HEADERS

from ..base import BaseItem, BaseItemQuerySet

LOGGER = logging.getLogger(__name__)

__all__ = [
    'UrlItem',

    # tasks will be added by register_task_method()
]

# ————————————————————————————————————————————————————————————————————— Classes

class DuplicateUrl(models.Model):

    """ A simple model to store duplicates URLs pointing to real ones. """

    # We can have extremely long URLs, notably via
    # google which wraps the entire world in their URLs.
    url = models.URLField(max_length=784,
                          verbose_name=_(u'URL'),
                          primary_key=True)

    real_url = models.URLField(max_length=512,
                               verbose_name=_(u'Real URL'))

    @classmethod
    def resolve(cls, url, clean=False):
        """ return the real URL of :param:`url` if it is a dupe.

        Return ``None`` if not registered as duplicate.
        """

        if clean:
            url = clean_url(url)

        try:
            return cls.objects.get(url=url).values('real_url')

        except:
            return None

    @classmethod
    def record_duplicate(cls, dupe, real):
        """ Register :param:`dupe` as pointing to :param:`real`. """

        url = cls(url=dupe, real_url=real)
        url.save()

# ——————————————————————————————————————————————————————————— QuerySet patching

def BaseItemQuerySet_absolute_method(self):
    """ Patch BaseItemQuerySet to know how to return absolute items. """

    return self.filter(url_absolute=True)

def BaseItemQuerySet_non_absolute_method(self):
    """ Patch BaseItemQuerySet to know how to return non_absolute content. """

    return self.filter(url_absolute=True)

def BaseItemQuerySet_orphaned_method(self):
    """ Patch BaseItemQuerySet to know how to return non_absolute content. """

    return self.filter(is_orphaned=True)

def BaseItemQuerySet_non_orphaned_method(self):
    """ Patch BaseItemQuerySet to know how to return non_absolute content. """

    return self.filter(is_orphaned=False)

def BaseItemQuerySet_url_error_method(self):
    """ Patch BaseItemQuerySet to know how to return non_absolute content. """

    return self.exclude(url_error=None)

BaseItemQuerySet.absolute     = BaseItemQuerySet_absolute_method
BaseItemQuerySet.non_absolute = BaseItemQuerySet_non_absolute_method
BaseItemQuerySet.orphaned     = BaseItemQuerySet_orphaned_method
BaseItemQuerySet.non_orphaned = BaseItemQuerySet_non_orphaned_method
BaseItemQuerySet.url_error    = BaseItemQuerySet_url_error_method

# ——————————————————————————————————————————————————————————————————————— Model

class UrlItem(models.Model):

    """ Abstract Item that knows how to handle an URL. """

    class Meta:
        abstract = True
        app_label = 'core'
        verbose_name = _(u'URL item')
        verbose_name_plural = _(u'URL items')

    url = models.URLField(unique=True,
                          max_length=512,
                          verbose_name=_(u'Public URL'))

    comments_feed_url = models.URLField(max_length=512,
                                        null=True, blank=True,
                                        verbose_name=_(u'Comments feed URL'))

    url_absolute = models.BooleanField(
        verbose_name=_(u'Absolute URL'),
        default=False, blank=True,
        help_text=_(u'The article URL has been successfully absolutized '
                    u'to its unique and final location.'))

    url_error = models.TextField(
        verbose_name=_(u'URL fetch error'),
        null=True, blank=True,
        help_text=_(u'Error when absolutizing the URL'))

    is_orphaned = models.BooleanField(
        verbose_name=_(u'Orphaned'),
        default=False, blank=True,
        help_text=_(u'This article has no public URL anymore, or is '
                    u'unfetchable for some reason.'))

    # —————————————————————————————————————————————————————————————— Properties

    @property
    def is_good(self):
        """ Return True if the current item is considered good. """

        #
        # NOTE: sync the conditions with @Feed.good_articles
        #       and core.views.read_with_endless_pagination::search
        #

        # HEADS UP: 20141111, we consider orphaned items to be good.
        #           although they don't have an URL, their title should
        #           be good, and even if we don't know yet how to display
        #           them correctly, they are not duplicates, and should
        #           be pushed to the user in the interface.
        # if self.is_orphaned:
        #    return False

        if not self.url_absolute:
            return False

        return True

    def absolutize_url_must_abort(self, force=False, commit=True):

        if config.ARTICLE_ABSOLUTIZING_DISABLED:
            LOGGER.info(u'Absolutizing disabled by configuration, aborting.')
            return True

        if self.url_absolute and not force:
            LOGGER.info(u'URL of article %s is already absolute!', self)
            return True

        if self.is_orphaned and not force:
            LOGGER.warning(u'Item %s is orphaned, absolutization aborted.',
                           self)
            return True

        if self.url_error:
            if force:
                self.url_error = None

                if commit:
                    self.save()
            else:
                LOGGER.warning(u'Item %s already has an URL error, '
                               u'aborting absolutization (currently: %s).',
                               self, self.url_error)
                return True

        return False

    def absolutize_url(self, requests_response=None, force=False, commit=True):
        """ Make the current article URL absolute.

        Eg. transform:

        http://feedproxy.google.com/~r/francaistechcrunch/~3/hEIhLwVyEEI/

        into:

        http://techcrunch.com/2013/05/18/hell-no-tumblr-users-wont-go-to-yahoo/ # NOQA
            ?utm_source=feeurner&utm_medium=feed&utm_campaign=Feed%3A+francaistechcrunch+%28TechCrunch+en+Francais%29 # NOQA

        and then remove all these F*G utm_* parameters to get a clean
        final URL for the current article.

        Returns ``True`` if the operation succeeded, ``False`` if the
        absolutization pointed out that the current article is a
        duplicate of another. In this case the caller should stop its
        processing because the current article will be marked for deletion.

        Can also return ``None`` if absolutizing is disabled globally
        in ``constance`` configuration.
        """

        # Another example: http://rss.lefigaro.fr/~r/lefigaro/laune/~3/7jgyrQ-PmBA/story01.htm # NOQA

        if self.absolutize_url_must_abort(force=force, commit=commit):
            return

        if requests_response is None:
            try:
                requests_response = requests.get(self.url,
                                                 headers=REQUEST_BASE_HEADERS)

            except requests.ConnectionError as e:
                statsd.gauge('articles.counts.url_errors', 1, delta=True)
                message = u'Connection error while absolutizing “%s”: %s'
                args = (self.url, str(e), )

                self.url_error = message % args
                self.save()

                LOGGER.error(message, *args)
                return

        if not requests_response.ok or requests_response.status_code != 200:

            message = u'HTTP Error %s while absolutizing “%s”: %s'
            args = (
                requests_response.status_code,
                requests_response.url,
                requests_response.reason
            )

            with statsd.pipeline() as spipe:
                # spipe.gauge('articles.counts.orphaned', 1, delta=True)
                spipe.gauge('articles.counts.url_errors', 1, delta=True)

                if requests_response.status_code in (404, ):
                    self.is_orphaned = True

                    # This is not handled by the post_save()
                    # which acts only at article creation.
                    spipe.gauge('articles.counts.orphaned', 1, delta=True)

            self.url_error = message % args
            self.save()

            LOGGER.error(message, *args)
            return

        #
        # NOTE: we could also get it eventually from r.headers['link'],
        #       which contains '<another_url>'. We need to strip out
        #       the '<>', and re-absolutize this link, because in the
        #       example it's another redirector. Also r.links is a good
        #       candidate but in the example I used, it contains the
        #       shortlink, which must be re-resolved too.
        #
        #       So: as we already are at the final address *now*, no need
        #       bothering re-following another which would lead us to the
        #       the same final place.
        #

        final_url = clean_url(requests_response.url)

        if final_url != self.url:

            # Just for displaying purposes, see below.
            old_url = self.url

            if self.url_error:
                statsd.gauge('articles.counts.url_errors', -1, delta=True)

            # Even if we are a duplicate, we came until here and everything
            # went fine. We won't need to lookup again the absolute URL.
            statsd.gauge('articles.counts.absolutes', 1, delta=True)
            self.url_absolute = True
            self.url_error = None

            self.url = final_url

            try:
                self.save()

            except IntegrityError as e:
                # LOGGER.exception(u'IntegrityError occured while saving %s',
                #                  final_url)

                # The following will legitimely crash on articles  missing
                # their BaseItem. We have a repair script dedicated to this.
                original = BaseItem.objects.get(Article___url=final_url)

                # Just to display the right "old" one in sentry errors and logs.
                self.url = old_url

                LOGGER.info(u'Item %s is a duplicate of %s, '
                            u'registering as such.', self, original)

                original.register_duplicate(self)
                return False

            # Any other exception will raise. This is intentional.
            else:
                LOGGER.info(u'Item %s (#%s) successfully absolutized URL '
                            u'from %s to %s.', self.name, self.id,
                            old_url, final_url)

        else:
            # Don't do the job twice.
            if self.url_error:
                statsd.gauge('articles.counts.url_errors', -1, delta=True)

            statsd.gauge('articles.counts.absolutes', 1, delta=True)
            self.url_absolute = True
            self.url_error = None
            self.save()

        return True

# ———————————————————————————————————————————————————————————————— Celery Tasks

# HEADS UP: we need to register against BaseItem, because UrlItem is abstract
#           and cannot run .objects.get() in register_task_method().
register_task_method(BaseItem, UrlItem.absolutize_url,
                     globals(), queue=u'swarm', default_retry_delay=3600)
<commit_msg>When absolutizing an URL and it changes, update the item name if it contains the URL, else some imported articles don't get a correct title if they have an initial url that's resolved to another one during the process.<commit_after># -*- coding: utf-8 -*-
"""
    Copyright 2013-2014 Olivier Cortès <oc@1flow.io>

    This file is part of the 1flow project.

    1flow is free software: you can redistribute it and/or modify
    it under the terms of the GNU Affero General Public License as
    published by the Free Software Foundation, either version 3 of
    the License, or (at your option) any later version.

    1flow is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU Affero General Public License for more details.

    You should have received a copy of the GNU Affero General Public
    License along with 1flow.  If not, see http://www.gnu.org/licenses/

"""
import logging

import requests

from statsd import statsd
from constance import config

# from json_field import JSONField
# from celery import chain as tasks_chain

# from django.conf import settings
from django.db import models, IntegrityError
from django.utils.translation import ugettext_lazy as _
# from django.utils.text import slugify

from oneflow.base.utils import register_task_method
from oneflow.base.utils.http import clean_url
# from oneflow.base.utils.dateutils import now

from ...common import REQUEST_BASE_HEADERS

from ..base import BaseItem, BaseItemQuerySet

LOGGER = logging.getLogger(__name__)

__all__ = [
    'UrlItem',

    # tasks will be added by register_task_method()
]

# ————————————————————————————————————————————————————————————————————— Classes

class DuplicateUrl(models.Model):

    """ A simple model to store duplicates URLs pointing to real ones. """

    # We can have extremely long URLs, notably via
    # google which wraps the entire world in their URLs.
    url = models.URLField(max_length=784,
                          verbose_name=_(u'URL'),
                          primary_key=True)

    real_url = models.URLField(max_length=512,
                               verbose_name=_(u'Real URL'))

    @classmethod
    def resolve(cls, url, clean=False):
        """ return the real URL of :param:`url` if it is a dupe.

        Return ``None`` if not registered as duplicate.
        """

        if clean:
            url = clean_url(url)

        try:
            return cls.objects.get(url=url).values('real_url')

        except:
            return None

    @classmethod
    def record_duplicate(cls, dupe, real):
        """ Register :param:`dupe` as pointing to :param:`real`. """

        url = cls(url=dupe, real_url=real)
        url.save()

# ——————————————————————————————————————————————————————————— QuerySet patching

def BaseItemQuerySet_absolute_method(self):
    """ Patch BaseItemQuerySet to know how to return absolute items. """

    return self.filter(url_absolute=True)

def BaseItemQuerySet_non_absolute_method(self):
    """ Patch BaseItemQuerySet to know how to return non_absolute content. """

    return self.filter(url_absolute=True)

def BaseItemQuerySet_orphaned_method(self):
    """ Patch BaseItemQuerySet to know how to return non_absolute content. """

    return self.filter(is_orphaned=True)

def BaseItemQuerySet_non_orphaned_method(self):
    """ Patch BaseItemQuerySet to know how to return non_absolute content. """

    return self.filter(is_orphaned=False)

def BaseItemQuerySet_url_error_method(self):
    """ Patch BaseItemQuerySet to know how to return non_absolute content. """

    return self.exclude(url_error=None)

BaseItemQuerySet.absolute     = BaseItemQuerySet_absolute_method
BaseItemQuerySet.non_absolute = BaseItemQuerySet_non_absolute_method
BaseItemQuerySet.orphaned     = BaseItemQuerySet_orphaned_method
BaseItemQuerySet.non_orphaned = BaseItemQuerySet_non_orphaned_method
BaseItemQuerySet.url_error    = BaseItemQuerySet_url_error_method

# ——————————————————————————————————————————————————————————————————————— Model

class UrlItem(models.Model):

    """ Abstract Item that knows how to handle an URL. """

    class Meta:
        abstract = True
        app_label = 'core'
        verbose_name = _(u'URL item')
        verbose_name_plural = _(u'URL items')

    url = models.URLField(unique=True,
                          max_length=512,
                          verbose_name=_(u'Public URL'))

    comments_feed_url = models.URLField(max_length=512,
                                        null=True, blank=True,
                                        verbose_name=_(u'Comments feed URL'))

    url_absolute = models.BooleanField(
        verbose_name=_(u'Absolute URL'),
        default=False, blank=True,
        help_text=_(u'The article URL has been successfully absolutized '
                    u'to its unique and final location.'))

    url_error = models.TextField(
        verbose_name=_(u'URL fetch error'),
        null=True, blank=True,
        help_text=_(u'Error when absolutizing the URL'))

    is_orphaned = models.BooleanField(
        verbose_name=_(u'Orphaned'),
        default=False, blank=True,
        help_text=_(u'This article has no public URL anymore, or is '
                    u'unfetchable for some reason.'))

    # —————————————————————————————————————————————————————————————— Properties

    @property
    def is_good(self):
        """ Return True if the current item is considered good. """

        #
        # NOTE: sync the conditions with @Feed.good_articles
        #       and core.views.read_with_endless_pagination::search
        #

        # HEADS UP: 20141111, we consider orphaned items to be good.
        #           although they don't have an URL, their title should
        #           be good, and even if we don't know yet how to display
        #           them correctly, they are not duplicates, and should
        #           be pushed to the user in the interface.
        # if self.is_orphaned:
        #    return False

        if not self.url_absolute:
            return False

        return True

    def absolutize_url_must_abort(self, force=False, commit=True):

        if config.ARTICLE_ABSOLUTIZING_DISABLED:
            LOGGER.info(u'Absolutizing disabled by configuration, aborting.')
            return True

        if self.url_absolute and not force:
            LOGGER.info(u'URL of article %s is already absolute!', self)
            return True

        if self.is_orphaned and not force:
            LOGGER.warning(u'Item %s is orphaned, absolutization aborted.',
                           self)
            return True

        if self.url_error:
            if force:
                self.url_error = None

                if commit:
                    self.save()
            else:
                LOGGER.warning(u'Item %s already has an URL error, '
                               u'aborting absolutization (currently: %s).',
                               self, self.url_error)
                return True

        return False

    def absolutize_url(self, requests_response=None, force=False, commit=True):
        """ Make the current article URL absolute.

        Eg. transform:

        http://feedproxy.google.com/~r/francaistechcrunch/~3/hEIhLwVyEEI/

        into:

        http://techcrunch.com/2013/05/18/hell-no-tumblr-users-wont-go-to-yahoo/ # NOQA
            ?utm_source=feeurner&utm_medium=feed&utm_campaign=Feed%3A+francaistechcrunch+%28TechCrunch+en+Francais%29 # NOQA

        and then remove all these F*G utm_* parameters to get a clean
        final URL for the current article.

        Returns ``True`` if the operation succeeded, ``False`` if the
        absolutization pointed out that the current article is a
        duplicate of another. In this case the caller should stop its
        processing because the current article will be marked for deletion.

        Can also return ``None`` if absolutizing is disabled globally
        in ``constance`` configuration.
        """

        # Another example: http://rss.lefigaro.fr/~r/lefigaro/laune/~3/7jgyrQ-PmBA/story01.htm # NOQA

        if self.absolutize_url_must_abort(force=force, commit=commit):
            return

        if requests_response is None:
            try:
                requests_response = requests.get(self.url,
                                                 headers=REQUEST_BASE_HEADERS)

            except requests.ConnectionError as e:
                statsd.gauge('articles.counts.url_errors', 1, delta=True)
                message = u'Connection error while absolutizing “%s”: %s'
                args = (self.url, str(e), )

                self.url_error = message % args
                self.save()

                LOGGER.error(message, *args)
                return

        if not requests_response.ok or requests_response.status_code != 200:

            message = u'HTTP Error %s while absolutizing “%s”: %s'
            args = (
                requests_response.status_code,
                requests_response.url,
                requests_response.reason
            )

            with statsd.pipeline() as spipe:
                # spipe.gauge('articles.counts.orphaned', 1, delta=True)
                spipe.gauge('articles.counts.url_errors', 1, delta=True)

                if requests_response.status_code in (404, ):
                    self.is_orphaned = True

                    # This is not handled by the post_save()
                    # which acts only at article creation.
                    spipe.gauge('articles.counts.orphaned', 1, delta=True)

            self.url_error = message % args
            self.save()

            LOGGER.error(message, *args)
            return

        #
        # NOTE: we could also get it eventually from r.headers['link'],
        #       which contains '<another_url>'. We need to strip out
        #       the '<>', and re-absolutize this link, because in the
        #       example it's another redirector. Also r.links is a good
        #       candidate but in the example I used, it contains the
        #       shortlink, which must be re-resolved too.
        #
        #       So: as we already are at the final address *now*, no need
        #       bothering re-following another which would lead us to the
        #       the same final place.
        #

        final_url = clean_url(requests_response.url)

        if final_url != self.url:

            # Just for displaying purposes, see below.
            old_url = self.url

            if self.url_error:
                statsd.gauge('articles.counts.url_errors', -1, delta=True)

            # Even if we are a duplicate, we came until here and everything
            # went fine. We won't need to lookup again the absolute URL.
            statsd.gauge('articles.counts.absolutes', 1, delta=True)
            self.url_absolute = True
            self.url_error = None

            self.url = final_url

            try:
                if self.name.endswith(old_url):
                    self.name = self.name.replace(old_url, final_url)
            except:
                LOGGER.exception(u'Could not replace URL in name of %s', self)

            try:
                self.save()

            except IntegrityError as e:
                # LOGGER.exception(u'IntegrityError occured while saving %s',
                #                  final_url)

                # The following will legitimely crash on articles  missing
                # their BaseItem. We have a repair script dedicated to this.
                original = BaseItem.objects.get(Article___url=final_url)

                # Just to display the right "old" one in sentry errors and logs.
                self.url = old_url

                LOGGER.info(u'Item %s is a duplicate of %s, '
                            u'registering as such.', self, original)

                original.register_duplicate(self)
                return False

            # Any other exception will raise. This is intentional.
            else:
                LOGGER.info(u'Item %s (#%s) successfully absolutized URL '
                            u'from %s to %s.', self.name, self.id,
                            old_url, final_url)

        else:
            # Don't do the job twice.
            if self.url_error:
                statsd.gauge('articles.counts.url_errors', -1, delta=True)

            statsd.gauge('articles.counts.absolutes', 1, delta=True)
            self.url_absolute = True
            self.url_error = None
            self.save()

        return True

# ———————————————————————————————————————————————————————————————— Celery Tasks

# HEADS UP: we need to register against BaseItem, because UrlItem is abstract
#           and cannot run .objects.get() in register_task_method().
register_task_method(BaseItem, UrlItem.absolutize_url,
                     globals(), queue=u'swarm', default_retry_delay=3600)
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-

# OpenFisca -- A versatile microsimulation software
# By: OpenFisca Team <contact@openfisca.fr>
#
# Copyright (C) 2011, 2012, 2013, 2014, 2015 OpenFisca Team
# https://github.com/openfisca
#
# This file is part of OpenFisca.
#
# OpenFisca is free software; you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as
# published by the Free Software Foundation, either version 3 of the
# License, or (at your option) any later version.
#
# OpenFisca is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

from ..base import *  # noqa

### VARIABLES MANQUANTES dans OpenFisca : ###
## 8TC 2011-    8WE -2006   7UP, 7UQ -2009  7LS 2010    7JY 2009-2011   7XN 2007-2010
## 7XJ, 7XK -2007   7XD -2006   7WQ 2006-2008   7UL -2009   7IY 2009-2011   7GS 2009-2011
## 6CH -2006    3VT ?

# Socio-economic data
# Données d'entrée de la simulation à fournir à partir d'une enquête ou générées par le générateur de cas type
build_column('idmen', IntCol(is_permanent = True, label = u"Identifiant du ménage"))
build_column('idfoy', IntCol(is_permanent = True, label = u"Identifiant du foyer"))
build_column('idfam', IntCol(is_permanent = True, label = u"Identifiant de la famille"))

build_column('quimen', EnumCol(QUIMEN, is_permanent = True))
build_column('quifoy', EnumCol(QUIFOY, is_permanent = True))
build_column('quifam', EnumCol(QUIFAM, is_permanent = True))

build_column('birth', DateCol(is_permanent = True, label = u"Date de naissance"))

build_column(Familles.name_key, StrCol(entity = 'fam', is_permanent = True, label = u"Nom"))
build_column(FoyersFiscaux.name_key, StrCol(entity = 'foy', is_permanent = True, label = u"Nom"))
build_column(Individus.name_key, StrCol(is_permanent = True, label = u"Prénom"))
build_column(Menages.name_key, StrCol(entity = 'men', is_permanent = True, label = u"Nom"))

build_column('enceinte', BoolCol(entity = 'ind', label = u"Est enceinte"))

# Prestations familiales
build_column('inactif', BoolCol(entity = 'fam',
                    label = u"Parent inactif (PAJE-CLCA)"))

build_column('partiel1', BoolCol(entity = 'fam',
                     label = u"Parent actif à moins de 50% (PAJE-CLCA)"))

build_column('partiel2', BoolCol(entity = 'fam',
                     label = u"Parent actif entre 50% et 80% (PAJE-CLCA)"))

build_column('categ_inv', PeriodSizeIndependentIntCol(label = u"Catégorie de handicap (AEEH)"))

build_column('opt_colca', BoolCol(entity = 'fam',
                      label = u"Opte pour le COLCA"))

build_column('empl_dir', BoolCol(entity = 'fam',
                     label = u"Emploi direct (CLCMG)"))

build_column('ass_mat', BoolCol(entity = 'fam',
                    label = u"Assistante maternelle (CLCMG)"))

build_column('gar_dom', BoolCol(entity = 'fam',
                    label = u"Garde à domicile (CLCMG)"))

# Autres
build_column('coloc', BoolCol(label = u"Vie en colocation"))
build_column(
    'csg_rempl',
    EnumCol(
        label = u"Taux retenu sur la CSG des revenus de remplacment",
        entity = 'ind',
        enum = Enum([
            u"Non renseigné/non pertinent",
            u"Exonéré",
            u"Taux réduit",
            u"Taux plein",
            ]),
        default = 3,
        ),
    )

reference_input_variable(
    column = IntCol(),
    entity_class = Individus,
    label = u"Chômage brut",
    name = 'chobrut',
    )

reference_input_variable(
    column = IntCol(),
    entity_class = Individus,
    label = u"Retraite brute",
    name = 'rstbrut',
    )

build_column('aer', IntCol(label = u"Allocation équivalent retraite (AER)"))  # L'AER est remplacée depuis le 1er juillet 2011 par l'allocation transitoire de solidarité (ATS).
build_column('f5sq', IntCol())

build_column('zthabm', IntCol(entity = 'men'))  # TODO: Devrait être renommée tax_hab

build_column(
    'proprietaire_proche_famille',
    BoolCol(
        entity = "fam",
        label = u"Le propriétaire du logement a un lien de parenté avec la personne de référence ou son conjoint",
        ),
    )

build_column('adoption', BoolCol(entity = "ind", label = u"Enfant adopté"))

build_column('ass_precondition_remplie', BoolCol(entity = "ind", label = u"Éligible à l'ASS"))

build_column('elig_creimp_jeunes', BoolCol(entity = "ind", label = u"Éligible au crédit d'impôt jeunes",
                    start = date(2005, 1, 1),
                    end = date(2008, 1, 1))) #Sert à savoir si son secteur d'activité permet au jeune de bénéficier du crédit impôts jeunes

# ('tax_hab', IntCol())
<commit_msg>passage chobrut et rstbrut en floatcol<commit_after># -*- coding: utf-8 -*-

# OpenFisca -- A versatile microsimulation software
# By: OpenFisca Team <contact@openfisca.fr>
#
# Copyright (C) 2011, 2012, 2013, 2014, 2015 OpenFisca Team
# https://github.com/openfisca
#
# This file is part of OpenFisca.
#
# OpenFisca is free software; you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as
# published by the Free Software Foundation, either version 3 of the
# License, or (at your option) any later version.
#
# OpenFisca is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

from ..base import *  # noqa

### VARIABLES MANQUANTES dans OpenFisca : ###
## 8TC 2011-    8WE -2006   7UP, 7UQ -2009  7LS 2010    7JY 2009-2011   7XN 2007-2010
## 7XJ, 7XK -2007   7XD -2006   7WQ 2006-2008   7UL -2009   7IY 2009-2011   7GS 2009-2011
## 6CH -2006    3VT ?

# Socio-economic data
# Données d'entrée de la simulation à fournir à partir d'une enquête ou générées par le générateur de cas type
build_column('idmen', IntCol(is_permanent = True, label = u"Identifiant du ménage"))
build_column('idfoy', IntCol(is_permanent = True, label = u"Identifiant du foyer"))
build_column('idfam', IntCol(is_permanent = True, label = u"Identifiant de la famille"))

build_column('quimen', EnumCol(QUIMEN, is_permanent = True))
build_column('quifoy', EnumCol(QUIFOY, is_permanent = True))
build_column('quifam', EnumCol(QUIFAM, is_permanent = True))

build_column('birth', DateCol(is_permanent = True, label = u"Date de naissance"))

build_column(Familles.name_key, StrCol(entity = 'fam', is_permanent = True, label = u"Nom"))
build_column(FoyersFiscaux.name_key, StrCol(entity = 'foy', is_permanent = True, label = u"Nom"))
build_column(Individus.name_key, StrCol(is_permanent = True, label = u"Prénom"))
build_column(Menages.name_key, StrCol(entity = 'men', is_permanent = True, label = u"Nom"))

build_column('enceinte', BoolCol(entity = 'ind', label = u"Est enceinte"))

# Prestations familiales
build_column('inactif', BoolCol(entity = 'fam',
                    label = u"Parent inactif (PAJE-CLCA)"))

build_column('partiel1', BoolCol(entity = 'fam',
                     label = u"Parent actif à moins de 50% (PAJE-CLCA)"))

build_column('partiel2', BoolCol(entity = 'fam',
                     label = u"Parent actif entre 50% et 80% (PAJE-CLCA)"))

build_column('categ_inv', PeriodSizeIndependentIntCol(label = u"Catégorie de handicap (AEEH)"))

build_column('opt_colca', BoolCol(entity = 'fam',
                      label = u"Opte pour le COLCA"))

build_column('empl_dir', BoolCol(entity = 'fam',
                     label = u"Emploi direct (CLCMG)"))

build_column('ass_mat', BoolCol(entity = 'fam',
                    label = u"Assistante maternelle (CLCMG)"))

build_column('gar_dom', BoolCol(entity = 'fam',
                    label = u"Garde à domicile (CLCMG)"))

# Autres
build_column('coloc', BoolCol(label = u"Vie en colocation"))
build_column(
    'csg_rempl',
    EnumCol(
        label = u"Taux retenu sur la CSG des revenus de remplacment",
        entity = 'ind',
        enum = Enum([
            u"Non renseigné/non pertinent",
            u"Exonéré",
            u"Taux réduit",
            u"Taux plein",
            ]),
        default = 3,
        ),
    )

reference_input_variable(
    column = FloatCol(),
    entity_class = Individus,
    label = u"Chômage brut",
    name = 'chobrut',
    )

reference_input_variable(
    column = FloatCol(),
    entity_class = Individus,
    label = u"Retraite brute",
    name = 'rstbrut',
    )

build_column('aer', IntCol(label = u"Allocation équivalent retraite (AER)"))  # L'AER est remplacée depuis le 1er juillet 2011 par l'allocation transitoire de solidarité (ATS).
build_column('f5sq', IntCol())

build_column('zthabm', IntCol(entity = 'men'))  # TODO: Devrait être renommée tax_hab

build_column(
    'proprietaire_proche_famille',
    BoolCol(
        entity = "fam",
        label = u"Le propriétaire du logement a un lien de parenté avec la personne de référence ou son conjoint",
        ),
    )

build_column('adoption', BoolCol(entity = "ind", label = u"Enfant adopté"))

build_column('ass_precondition_remplie', BoolCol(entity = "ind", label = u"Éligible à l'ASS"))

build_column('elig_creimp_jeunes', BoolCol(entity = "ind", label = u"Éligible au crédit d'impôt jeunes",
                    start = date(2005, 1, 1),
                    end = date(2008, 1, 1))) #Sert à savoir si son secteur d'activité permet au jeune de bénéficier du crédit impôts jeunes

# ('tax_hab', IntCol())
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-

# Copyright (c) 2010-2011, GEM Foundation.
#
# OpenQuake is free software: you can redistribute it and/or modify
# it under the terms of the GNU Lesser General Public License version 3
# only, as published by the Free Software Foundation.
#
# OpenQuake is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Lesser General Public License version 3 for more details
# (a copy is included in the LICENSE file that accompanied this code).
#
# You should have received a copy of the GNU Lesser General Public License
# version 3 along with OpenQuake.  If not, see
# <http://www.gnu.org/licenses/lgpl-3.0.txt> for a copy of the LGPLv3 License.

# vim: tabstop=4 shiftwidth=4 softtabstop=4

"""Core functionality for Event-Based Risk calculations."""

import os

from numpy import zeros

from celery.exceptions import TimeoutError

from openquake import kvs
from openquake import logs
from openquake import shapes
from openquake.db import models
from openquake.parser import vulnerability
from openquake.output import curve
from openquake.calculators.risk import general

LOGGER = logs.LOG

def _filename(job_id):
    """Return the name of the generated file."""
    return "%s-aggregate-loss-curve.svg" % job_id

def _for_plotting(loss_curve, time_span):
    """Translate a loss curve into a dictionary compatible to
    the interface defined in CurvePlot.write."""
    data = {}

    data["AggregateLossCurve"] = {}
    data["AggregateLossCurve"]["abscissa"] = tuple(loss_curve.abscissae)
    data["AggregateLossCurve"]["ordinate"] = tuple(loss_curve.ordinates)
    data["AggregateLossCurve"]["abscissa_property"] = "Economic Losses"
    data["AggregateLossCurve"]["ordinate_property"] = \
            "PoE in %s years" % (str(time_span))

    data["AggregateLossCurve"]["curve_title"] = "Aggregate Loss Curve"

    return data

def plot_aggregate_curve(calculator, aggregate_curve):
    """Plot an aggreate loss curve.

    This function is triggered only if the AGGREGATE_LOSS_CURVE
    parameter is specified in the configuration file.

    :param calculator:
        :py:class:`EventBasedRiskCalculator` instance for an in-progress
        calculation.
    :param aggregate_curve: the aggregate curve to plot.
    :type aggregate_curve: :py:class:`openquake.shapes.Curve`
    """

    if not calculator.calc_proxy.has("AGGREGATE_LOSS_CURVE"):
        LOGGER.debug("AGGREGATE_LOSS_CURVE parameter not specified, " \
                "skipping aggregate loss curve computation...")

        return

    path = os.path.join(
            calculator.calc_proxy.params["BASE_PATH"],
            calculator.calc_proxy.params["OUTPUT_DIR"],
            _filename(calculator.calc_proxy.job_id))

    plotter = curve.CurvePlot(path)
    plotter.write(_for_plotting(aggregate_curve,
            calculator.calc_proxy.params["INVESTIGATION_TIME"]),
            autoscale_y=False)

    plotter.close()
    LOGGER.debug("Aggregate loss curve stored at %s" % path)

class EventBasedRiskCalculator(general.ProbabilisticRiskCalculator):
    """Calculator for Event-Based Risk computations."""

    def __init__(self, job_profile):
        super(EventBasedRiskCalculator, self).__init__(job_profile)
        self.vuln_curves = None

    def execute(self):
        """Execute the job."""
        general.preload(self)

        aggregate_curve = general.AggregateLossCurve()

        tasks = []
        for block_id in self.calc_proxy.blocks_keys:
            LOGGER.debug("Starting task block, block_id = %s of %s"
                    % (block_id, len(self.calc_proxy.blocks_keys)))

            tasks.append(general.compute_risk.delay(self.calc_proxy.job_id,
                                                    block_id))

        for task in tasks:
            try:
                task.wait()

                aggregate_curve.append(task.result)
            except TimeoutError:
                # TODO(jmc): Cancel and respawn this task
                return

        if self.is_benefit_cost_ratio_mode():
            general.write_output_bcr(self)
            return

        agg_curve = aggregate_curve.compute(self._tses(), self._time_span())
        plot_aggregate_curve(self, agg_curve)

        general.write_output(self)

    def _tses(self):
        """Return the time representative of the Stochastic Event Set
        specified for this job."""

        # TODO (ac): Confirm this works regardless of the method of hazard calc
        histories = int(
            self.calc_proxy.params["NUMBER_OF_SEISMICITY_HISTORIES"])
        realizations = int(
            self.calc_proxy.params["NUMBER_OF_LOGIC_TREE_SAMPLES"])
        num_ses = histories * realizations

        return num_ses * self._time_span()

    def _time_span(self):
        """Return the time span specified for this job."""
        return float(self.calc_proxy.params["INVESTIGATION_TIME"])

    def _gmf_db_list(self, job_id):  # pylint: disable=R0201
        """Returns a list of the output IDs of all computed GMFs"""

        ids = models.Output.objects.filter(
            oq_calculation=job_id, output_type='gmf').values_list('id',
                                                                  flat=True)

        return list(ids)

    def _get_db_gmf(self, gmf_id):
        """Returns a field for the given GMF"""
        grid = self.calc_proxy.region.grid
        field = zeros((grid.rows, grid.columns))

        gmf_sites = models.GmfData.objects.filter(output=gmf_id)

        for gmf_site in gmf_sites:
            loc = gmf_site.location
            site = shapes.Site(loc.x, loc.y)
            grid_point = grid.point_at(site)

            field[grid_point.row][grid_point.column] = gmf_site.ground_motion

        return shapes.Field(field)

    def _sites_to_gmf_keys(self, sites):
        """Returns the GMF keys "row!col" for the given site list"""
        keys = []

        for site in sites:
            risk_point = self.calc_proxy.region.grid.point_at(site)
            key = "%s!%s" % (risk_point.row, risk_point.column)
            keys.append(key)

        return keys

    def _get_db_gmfs(self, sites, job_id):
        """Aggregates GMF data from the DB by site"""
        all_gmfs = self._gmf_db_list(job_id)
        gmf_keys = self._sites_to_gmf_keys(sites)
        gmfs = dict((k, []) for k in gmf_keys)

        for gmf_id in all_gmfs:
            field = self._get_db_gmf(gmf_id)

            for key in gmfs.keys():
                (row, col) = key.split("!")
                gmfs[key].append(field.get(int(row), int(col)))

        return gmfs

    def _get_kvs_gmfs(self, sites, histories, realizations):
        """Aggregates GMF data from the KVS by site"""
        gmf_keys = self._sites_to_gmf_keys(sites)
        gmfs = dict((k, []) for k in gmf_keys)

        for i in range(0, histories):
            for j in range(0, realizations):
                key = kvs.tokens.stochastic_set_key(
                    self.calc_proxy.job_id, i, j)
                fieldset = shapes.FieldSet.from_json(kvs.get(key),
                    self.calc_proxy.region.grid)

                for field in fieldset:
                    for key in gmfs.keys():
                        (row, col) = key.split("!")
                        gmfs[key].append(field.get(int(row), int(col)))

        return gmfs

    def slice_gmfs(self, block_id):
        """Load and collate GMF values for all sites in this block. """
        block = general.Block.from_kvs(block_id)
        gmfs = self._get_db_gmfs(block.sites, self.calc_proxy.job_id)

        for key, gmf_slice in gmfs.items():
            (row, col) = key.split("!")
            key_gmf = kvs.tokens.gmf_set_key(self.calc_proxy.job_id, col, row)
            LOGGER.debug("GMF_SLICE for %s X %s : \n\t%s" % (
                    col, row, gmf_slice))
            gmf = {"IMLs": gmf_slice, "TSES": self._tses(),
                    "TimeSpan": self._time_span()}
            kvs.set_value_json_encoded(key_gmf, gmf)

    def _compute_loss(self, block_id):
        """Compute risk for a block of sites, that means:

        * loss ratio curves
        * loss curves
        * conditional losses
        * (partial) aggregate loss curve
        """

        self.slice_gmfs(block_id)

        self.vuln_curves = vulnerability.load_vuln_model_from_kvs(
            self.calc_proxy.job_id)

        block = general.Block.from_kvs(block_id)

        # aggregate the losses for this block
        aggregate_curve = general.AggregateLossCurve()

        for point in block.grid(self.calc_proxy.region):
            key = kvs.tokens.gmf_set_key(self.calc_proxy.job_id, point.column,
                                         point.row)
            gmf_slice = kvs.get_value_json_decoded(key)

            asset_key = kvs.tokens.asset_key(
                self.calc_proxy.job_id, point.row, point.column)

            for asset in kvs.get_list_json_decoded(asset_key):
                LOGGER.debug("Processing asset %s" % (asset))

                # loss ratios, used both to produce the curve
                # and to aggregate the losses
                loss_ratios = self.compute_loss_ratios(asset, gmf_slice)

                loss_ratio_curve = self.compute_loss_ratio_curve(
                    point.column, point.row, asset, gmf_slice, loss_ratios)

                aggregate_curve.append(loss_ratios * asset["assetValue"])

                if loss_ratio_curve:
                    loss_curve = self.compute_loss_curve(
                        point.column, point.row, loss_ratio_curve, asset)

                    for loss_poe in general.conditional_loss_poes(
                        self.calc_proxy.params):

                        general.compute_conditional_loss(
                                self.calc_proxy.job_id, point.column,
                                point.row, loss_curve, asset, loss_poe)

        return aggregate_curve.losses

    def _compute_bcr(self, block_id):
        """
        Calculate and store in the kvs the benefit-cost ratio data for block.

        A value is stored with key :func:`openquake.kvs.tokens.bcr_block_key`.
        See :func:`openquake.risk.job.general.compute_bcr_for_block` for result
        data structure spec.
        """
        self.slice_gmfs(block_id)

        points = list(general.Block.from_kvs(block_id).grid(
            self.calc_proxy.region))
        gmf_slices = dict(
            (point.site, kvs.get_value_json_decoded(
                 kvs.tokens.gmf_set_key(self.calc_proxy.job_id, point.column,
                                        point.row)
            ))
            for point in points
        )
        epsilon_provider = general.EpsilonProvider(self.calc_proxy.params)

        def get_loss_curve(point, vuln_function, asset):
            "Compute loss curve basing on GMF data"
            gmf_slice = gmf_slices[point.site]
            loss_ratios = general.compute_loss_ratios(
                vuln_function, gmf_slice, epsilon_provider, asset)
            loss_ratio_curve = general.compute_loss_ratio_curve(
                vuln_function, gmf_slice, epsilon_provider, asset,
                None, loss_ratios=loss_ratios)
            return loss_ratio_curve.rescale_abscissae(asset["assetValue"])

        result = general.compute_bcr_for_block(self.calc_proxy.job_id, points,
            get_loss_curve, float(self.calc_proxy.params['INTEREST_RATE']),
            float(self.calc_proxy.params['ASSET_LIFE_EXPECTANCY'])
        )

        bcr_block_key = kvs.tokens.bcr_block_key(self.calc_proxy.job_id,
                                                 block_id)
        kvs.set_value_json_encoded(bcr_block_key, result)
        LOGGER.debug('bcr result for block %s: %r', block_id, result)
        return True

    def compute_loss_ratios(self, asset, gmf_slice):
        """For a given asset and ground motion field, computes
        the loss ratios used to obtain the related loss ratio curve
        and aggregate loss curve."""

        epsilon_provider = general.EpsilonProvider(self.calc_proxy.params)

        vuln_function = self.vuln_curves.get(
            asset["taxonomy"], None)

        if not vuln_function:
            LOGGER.error(
                "Unknown vulnerability function %s for asset %s"
                % (asset["taxonomy"], asset["assetID"]))

            return None

        return general.compute_loss_ratios(
            vuln_function, gmf_slice, epsilon_provider, asset)

    def compute_loss_ratio_curve(
            self, col, row, asset, gmf_slice, loss_ratios):
        """Compute the loss ratio curve for a single asset."""

        vuln_function = self.vuln_curves.get(
            asset["taxonomy"], None)

        if not vuln_function:
            LOGGER.error(
                "Unknown vulnerability function %s for asset %s"
                % (asset["taxonomy"], asset["assetID"]))

            return None

        epsilon_provider = general.EpsilonProvider(self.calc_proxy.params)

        loss_ratio_curve = general.compute_loss_ratio_curve(
                vuln_function, gmf_slice, epsilon_provider, asset,
                None, loss_ratios=loss_ratios)

        # NOTE (jmc): Early exit if the loss ratio is all zeros
        if not False in (loss_ratio_curve.ordinates == 0.0):
            return None

        key = kvs.tokens.loss_ratio_key(
            self.calc_proxy.job_id, row, col, asset["assetID"])

        kvs.get_client().set(key, loss_ratio_curve.to_json())

        LOGGER.debug("Loss ratio curve is %s, write to key %s" %
                (loss_ratio_curve, key))

        return loss_ratio_curve

    def compute_loss_curve(self, column, row, loss_ratio_curve, asset):
        """Compute the loss curve for a single asset."""

        if asset is None:
            return None

        loss_curve = loss_ratio_curve.rescale_abscissae(asset["assetValue"])

        key = kvs.tokens.loss_curve_key(
            self.calc_proxy.job_id, row, column, asset["assetID"])

        LOGGER.debug("Loss curve is %s, write to key %s" % (loss_curve, key))
        kvs.get_client().set(key, loss_curve.to_json())

        return loss_curve
<commit_msg>pep8<commit_after># -*- coding: utf-8 -*-

# Copyright (c) 2010-2011, GEM Foundation.
#
# OpenQuake is free software: you can redistribute it and/or modify
# it under the terms of the GNU Lesser General Public License version 3
# only, as published by the Free Software Foundation.
#
# OpenQuake is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Lesser General Public License version 3 for more details
# (a copy is included in the LICENSE file that accompanied this code).
#
# You should have received a copy of the GNU Lesser General Public License
# version 3 along with OpenQuake.  If not, see
# <http://www.gnu.org/licenses/lgpl-3.0.txt> for a copy of the LGPLv3 License.

# vim: tabstop=4 shiftwidth=4 softtabstop=4

"""Core functionality for Event-Based Risk calculations."""

import os

from numpy import zeros

from celery.exceptions import TimeoutError

from openquake import kvs
from openquake import logs
from openquake import shapes
from openquake.db import models
from openquake.parser import vulnerability
from openquake.output import curve
from openquake.calculators.risk import general

LOGGER = logs.LOG

def _filename(job_id):
    """Return the name of the generated file."""
    return "%s-aggregate-loss-curve.svg" % job_id

def _for_plotting(loss_curve, time_span):
    """Translate a loss curve into a dictionary compatible to
    the interface defined in CurvePlot.write."""
    data = {}

    data["AggregateLossCurve"] = {}
    data["AggregateLossCurve"]["abscissa"] = tuple(loss_curve.abscissae)
    data["AggregateLossCurve"]["ordinate"] = tuple(loss_curve.ordinates)
    data["AggregateLossCurve"]["abscissa_property"] = "Economic Losses"
    data["AggregateLossCurve"]["ordinate_property"] = \
            "PoE in %s years" % (str(time_span))

    data["AggregateLossCurve"]["curve_title"] = "Aggregate Loss Curve"

    return data

def plot_aggregate_curve(calculator, aggregate_curve):
    """Plot an aggreate loss curve.

    This function is triggered only if the AGGREGATE_LOSS_CURVE
    parameter is specified in the configuration file.

    :param calculator:
        :py:class:`EventBasedRiskCalculator` instance for an in-progress
        calculation.
    :param aggregate_curve: the aggregate curve to plot.
    :type aggregate_curve: :py:class:`openquake.shapes.Curve`
    """

    if not calculator.calc_proxy.has("AGGREGATE_LOSS_CURVE"):
        LOGGER.debug("AGGREGATE_LOSS_CURVE parameter not specified, " \
                "skipping aggregate loss curve computation...")

        return

    path = os.path.join(
            calculator.calc_proxy.params["BASE_PATH"],
            calculator.calc_proxy.params["OUTPUT_DIR"],
            _filename(calculator.calc_proxy.job_id))

    plotter = curve.CurvePlot(path)
    plotter.write(_for_plotting(aggregate_curve,
            calculator.calc_proxy.params["INVESTIGATION_TIME"]),
            autoscale_y=False)

    plotter.close()
    LOGGER.debug("Aggregate loss curve stored at %s" % path)

class EventBasedRiskCalculator(general.ProbabilisticRiskCalculator):
    """Calculator for Event-Based Risk computations."""

    def __init__(self, job_profile):
        super(EventBasedRiskCalculator, self).__init__(job_profile)
        self.vuln_curves = None

    def execute(self):
        """Execute the job."""
        general.preload(self)

        aggregate_curve = general.AggregateLossCurve()

        tasks = []
        for block_id in self.calc_proxy.blocks_keys:
            LOGGER.debug("Starting task block, block_id = %s of %s"
                    % (block_id, len(self.calc_proxy.blocks_keys)))

            tasks.append(general.compute_risk.delay(self.calc_proxy.job_id,
                                                    block_id))

        for task in tasks:
            try:
                task.wait()

                aggregate_curve.append(task.result)
            except TimeoutError:
                # TODO(jmc): Cancel and respawn this task
                return

        if self.is_benefit_cost_ratio_mode():
            general.write_output_bcr(self)
            return

        agg_curve = aggregate_curve.compute(self._tses(), self._time_span())
        plot_aggregate_curve(self, agg_curve)

        general.write_output(self)

    def _tses(self):
        """Return the time representative of the Stochastic Event Set
        specified for this job."""

        # TODO (ac): Confirm this works regardless of the method of hazard calc
        histories = int(
            self.calc_proxy.params["NUMBER_OF_SEISMICITY_HISTORIES"])
        realizations = int(
            self.calc_proxy.params["NUMBER_OF_LOGIC_TREE_SAMPLES"])
        num_ses = histories * realizations

        return num_ses * self._time_span()

    def _time_span(self):
        """Return the time span specified for this job."""
        return float(self.calc_proxy.params["INVESTIGATION_TIME"])

    def _gmf_db_list(self, job_id):  # pylint: disable=R0201
        """Returns a list of the output IDs of all computed GMFs"""

        ids = models.Output.objects.filter(
            oq_calculation=job_id, output_type='gmf').values_list('id',
                                                                  flat=True)

        return list(ids)

    def _get_db_gmf(self, gmf_id):
        """Returns a field for the given GMF"""
        grid = self.calc_proxy.region.grid
        field = zeros((grid.rows, grid.columns))

        gmf_sites = models.GmfData.objects.filter(output=gmf_id)

        for gmf_site in gmf_sites:
            loc = gmf_site.location
            site = shapes.Site(loc.x, loc.y)
            grid_point = grid.point_at(site)

            field[grid_point.row][grid_point.column] = gmf_site.ground_motion

        return shapes.Field(field)

    def _sites_to_gmf_keys(self, sites):
        """Returns the GMF keys "row!col" for the given site list"""
        keys = []

        for site in sites:
            risk_point = self.calc_proxy.region.grid.point_at(site)
            key = "%s!%s" % (risk_point.row, risk_point.column)
            keys.append(key)

        return keys

    def _get_db_gmfs(self, sites, job_id):
        """Aggregates GMF data from the DB by site"""
        all_gmfs = self._gmf_db_list(job_id)
        gmf_keys = self._sites_to_gmf_keys(sites)
        gmfs = dict((k, []) for k in gmf_keys)

        for gmf_id in all_gmfs:
            field = self._get_db_gmf(gmf_id)

            for key in gmfs.keys():
                (row, col) = key.split("!")
                gmfs[key].append(field.get(int(row), int(col)))

        return gmfs

    def _get_kvs_gmfs(self, sites, histories, realizations):
        """Aggregates GMF data from the KVS by site"""
        gmf_keys = self._sites_to_gmf_keys(sites)
        gmfs = dict((k, []) for k in gmf_keys)

        for i in range(0, histories):
            for j in range(0, realizations):
                key = kvs.tokens.stochastic_set_key(
                    self.calc_proxy.job_id, i, j)
                fieldset = shapes.FieldSet.from_json(kvs.get(key),
                    self.calc_proxy.region.grid)

                for field in fieldset:
                    for key in gmfs.keys():
                        (row, col) = key.split("!")
                        gmfs[key].append(field.get(int(row), int(col)))

        return gmfs

    def slice_gmfs(self, block_id):
        """Load and collate GMF values for all sites in this block. """
        block = general.Block.from_kvs(block_id)
        gmfs = self._get_db_gmfs(block.sites, self.calc_proxy.job_id)

        for key, gmf_slice in gmfs.items():
            (row, col) = key.split("!")
            key_gmf = kvs.tokens.gmf_set_key(self.calc_proxy.job_id, col, row)
            LOGGER.debug("GMF_SLICE for %s X %s : \n\t%s" % (
                    col, row, gmf_slice))
            gmf = {"IMLs": gmf_slice, "TSES": self._tses(),
                    "TimeSpan": self._time_span()}
            kvs.set_value_json_encoded(key_gmf, gmf)

    def _compute_loss(self, block_id):
        """Compute risk for a block of sites, that means:

        * loss ratio curves
        * loss curves
        * conditional losses
        * (partial) aggregate loss curve
        """

        self.slice_gmfs(block_id)

        self.vuln_curves = vulnerability.load_vuln_model_from_kvs(
            self.calc_proxy.job_id)

        block = general.Block.from_kvs(block_id)

        # aggregate the losses for this block
        aggregate_curve = general.AggregateLossCurve()

        for point in block.grid(self.calc_proxy.region):
            key = kvs.tokens.gmf_set_key(self.calc_proxy.job_id, point.column,
                                         point.row)
            gmf_slice = kvs.get_value_json_decoded(key)

            asset_key = kvs.tokens.asset_key(
                self.calc_proxy.job_id, point.row, point.column)

            for asset in kvs.get_list_json_decoded(asset_key):
                LOGGER.debug("Processing asset %s" % (asset))

                # loss ratios, used both to produce the curve
                # and to aggregate the losses
                loss_ratios = self.compute_loss_ratios(asset, gmf_slice)

                loss_ratio_curve = self.compute_loss_ratio_curve(
                    point.column, point.row, asset, gmf_slice, loss_ratios)

                aggregate_curve.append(loss_ratios * asset["assetValue"])

                if loss_ratio_curve:
                    loss_curve = self.compute_loss_curve(
                        point.column, point.row, loss_ratio_curve, asset)

                    for loss_poe in general.conditional_loss_poes(
                        self.calc_proxy.params):

                        general.compute_conditional_loss(
                                self.calc_proxy.job_id, point.column,
                                point.row, loss_curve, asset, loss_poe)

        return aggregate_curve.losses

    def _compute_bcr(self, block_id):
        """
        Calculate and store in the kvs the benefit-cost ratio data for block.

        A value is stored with key :func:`openquake.kvs.tokens.bcr_block_key`.
        See :func:`openquake.risk.job.general.compute_bcr_for_block` for result
        data structure spec.
        """
        self.slice_gmfs(block_id)

        points = list(general.Block.from_kvs(block_id).grid(
            self.calc_proxy.region))
        gmf_slices = dict(
            (point.site, kvs.get_value_json_decoded(
                 kvs.tokens.gmf_set_key(self.calc_proxy.job_id, point.column,
                                        point.row)
            ))
            for point in points
        )
        epsilon_provider = general.EpsilonProvider(self.calc_proxy.params)

        def get_loss_curve(point, vuln_function, asset):
            "Compute loss curve basing on GMF data"
            gmf_slice = gmf_slices[point.site]
            loss_ratios = general.compute_loss_ratios(
                vuln_function, gmf_slice, epsilon_provider, asset)
            loss_ratio_curve = general.compute_loss_ratio_curve(
                vuln_function, gmf_slice, epsilon_provider, asset,
                None, loss_ratios=loss_ratios)
            return loss_ratio_curve.rescale_abscissae(asset["assetValue"])

        result = general.compute_bcr_for_block(self.calc_proxy.job_id, points,
            get_loss_curve, float(self.calc_proxy.params['INTEREST_RATE']),
            float(self.calc_proxy.params['ASSET_LIFE_EXPECTANCY'])
        )

        bcr_block_key = kvs.tokens.bcr_block_key(self.calc_proxy.job_id,
                                                 block_id)
        kvs.set_value_json_encoded(bcr_block_key, result)
        LOGGER.debug('bcr result for block %s: %r', block_id, result)
        return True

    def compute_loss_ratios(self, asset, gmf_slice):
        """For a given asset and ground motion field, computes
        the loss ratios used to obtain the related loss ratio curve
        and aggregate loss curve."""

        epsilon_provider = general.EpsilonProvider(self.calc_proxy.params)

        vuln_function = self.vuln_curves.get(
            asset["taxonomy"], None)

        if not vuln_function:
            LOGGER.error(
                "Unknown vulnerability function %s for asset %s"
                % (asset["taxonomy"], asset["assetID"]))

            return None

        return general.compute_loss_ratios(
            vuln_function, gmf_slice, epsilon_provider, asset)

    def compute_loss_ratio_curve(
            self, col, row, asset, gmf_slice, loss_ratios):
        """Compute the loss ratio curve for a single asset."""

        vuln_function = self.vuln_curves.get(
            asset["taxonomy"], None)

        if not vuln_function:
            LOGGER.error(
                "Unknown vulnerability function %s for asset %s"
                % (asset["taxonomy"], asset["assetID"]))

            return None

        epsilon_provider = general.EpsilonProvider(self.calc_proxy.params)

        loss_ratio_curve = general.compute_loss_ratio_curve(
                vuln_function, gmf_slice, epsilon_provider, asset,
                None, loss_ratios=loss_ratios)

        # NOTE (jmc): Early exit if the loss ratio is all zeros
        if not False in (loss_ratio_curve.ordinates == 0.0):
            return None

        key = kvs.tokens.loss_ratio_key(
            self.calc_proxy.job_id, row, col, asset["assetID"])

        kvs.get_client().set(key, loss_ratio_curve.to_json())

        LOGGER.debug("Loss ratio curve is %s, write to key %s" %
                (loss_ratio_curve, key))

        return loss_ratio_curve

    def compute_loss_curve(self, column, row, loss_ratio_curve, asset):
        """Compute the loss curve for a single asset."""

        if asset is None:
            return None

        loss_curve = loss_ratio_curve.rescale_abscissae(asset["assetValue"])

        key = kvs.tokens.loss_curve_key(
            self.calc_proxy.job_id, row, column, asset["assetID"])

        LOGGER.debug("Loss curve is %s, write to key %s" % (loss_curve, key))
        kvs.get_client().set(key, loss_curve.to_json())

        return loss_curve
<|endoftext|>


<commit_before># Licensed to Cloudera, Inc. under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  Cloudera, Inc. 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 cm_api.endpoints.types import *
from cm_api.endpoints import services, parcels, host_templates

__docformat__ = "epytext"

CLUSTERS_PATH = "/clusters"

def create_cluster(resource_root, name, version):
  """
  Create a cluster
  @param resource_root: The root Resource object.
  @param name: Cluster name
  @param version: Cluster CDH version
  @return: An ApiCluster object
  """
  apicluster = ApiCluster(resource_root, name, version)
  return call(resource_root.post, CLUSTERS_PATH, ApiCluster, True,
      data=[apicluster])[0]

def get_cluster(resource_root, name):
  """
  Lookup a cluster by name
  @param resource_root: The root Resource object.
  @param name: Cluster name
  @return: An ApiCluster object
  """
  return call(resource_root.get, "%s/%s" % (CLUSTERS_PATH, name), ApiCluster)

def get_all_clusters(resource_root, view=None):
  """
  Get all clusters
  @param resource_root: The root Resource object.
  @return: A list of ApiCluster objects.
  """
  return call(resource_root.get, CLUSTERS_PATH, ApiCluster, True,
      params=view and dict(view=view) or None)

def delete_cluster(resource_root, name):
  """
  Delete a cluster by name
  @param resource_root: The root Resource object.
  @param name: Cluster name
  @return: The deleted ApiCluster object
  """
  return call(resource_root.delete, "%s/%s" % (CLUSTERS_PATH, name), ApiCluster)

class ApiCluster(BaseApiResource):
  _ATTRIBUTES = {
    'name'              : None,
    'version'           : None,
    'maintenanceMode'   : ROAttr(),
    'maintenanceOwners' : ROAttr(),
  }

  def __init__(self, resource_root, name=None, version=None):
    BaseApiObject.init(self, resource_root, locals())

  def __str__(self):
    return "<ApiCluster>: %s; version: %s" % (self.name, self.version)

  def _path(self):
    return "%s/%s" % (CLUSTERS_PATH, self.name)

  def _put(self, dic, params=None):
    """Change cluster attributes"""
    cluster = self._put('', ApiCluster, data=dic, params=params)
    self._update(cluster)
    return self

  def get_commands(self, view=None):
    """
    Retrieve a list of running commands for this cluster.

    @param view: View to materialize ('full' or 'summary')
    @return: A list of running commands.
    """
    return self._get("commands", ApiCommand, True,
        params = view and dict(view=view) or None)

  def rename(self, newname):
    """
    Rename a cluster.

    @param newname: New cluster name
    @return: An ApiCluster object
    @since: API v2
    """
    dic = self.to_json_dict()
    dic['name'] = newname
    return self._put(dic)

  def create_service(self, name, service_type):
    """
    Create a service.

    @param name: Service name
    @param service_type: Service type
    @return: An ApiService object
    """
    return services.create_service(self._get_resource_root(), name,
        service_type, self.name)

  def delete_service(self, name):
    """
    Delete a service by name.

    @param name Service name
    @return The deleted ApiService object
    """
    return services.delete_service(self._get_resource_root(), name, self.name)

  def get_service(self, name):
    """
    Lookup a service by name.

    @param name: Service name
    @return: An ApiService object
    """
    return services.get_service(self._get_resource_root(), name, self.name)

  def get_all_services(self, view = None):
    """
    Get all services in this cluster.

    @return: A list of ApiService objects.
    """
    return services.get_all_services(self._get_resource_root(), self.name, view)

  def get_parcel(self, product, version):
    """
    Lookup a parcel by product and version.

    @param product: the product name
    @param version: the product version
    @return: An ApiParcel object
    """
    return parcels.get_parcel(self._get_resource_root(), product, version, self.name)

  def get_all_parcels(self, view = None):
    """
    Get all parcels in this cluster.

    @return: A list of ApiParcel objects.
    """
    return parcels.get_all_parcels(self._get_resource_root(), self.name, view)

  def list_hosts(self):
    """
    Lists all the hosts that are associated with this cluster.

    @return: A list of ApiHostRef objects of the hosts in the cluster.
    @since: API v3
    """
    return self._get("hosts", ApiHostRef, True, api_version=3)

  def remove_host(self, hostId):
    """
    Removes the association of the host with the cluster.

    @return: A ApiHostRef of the host that was removed.
    @since: API v3
    """
    return self._delete("hosts/" + hostId, ApiHostRef, api_version=3)

  def remove_all_hosts(self):
    """
    Removes the association of all the hosts with the cluster.

    @return: A list of ApiHostRef objects of the hosts that were removed.
    @since: API v3
    """
    return self._delete("hosts", ApiHostRef, True, api_version=3)

  def add_hosts(self, hostIds):
    """
    Adds a host to the cluster.

    @param hostIds: List of IDs of hosts to add to cluster.
    @return: A list of ApiHostRef objects of the new
             hosts that were added to the cluster
    @since: API v3
    """
    hostRefList = [ApiHostRef(self._get_resource_root(), x) for x in hostIds]
    return self._post("hosts", ApiHostRef, True, data=hostRefList,
        api_version=3)

  def start(self):
    """
    Start all services in a cluster, respecting dependencies.

    @return: Reference to the submitted command.
    """
    return self._cmd('start')

  def stop(self):
    """
    Stop all services in a cluster, respecting dependencies.

    @return: Reference to the submitted command.
    """
    return self._cmd('stop')

  def deploy_client_config(self):
    """
    Deploys client configuration to the hosts on the cluster.

    @return: Reference to the submitted command.
    @since: API v2
    """
    return self._cmd('deployClientConfig')

  def upgrade_services(self):
    """
    Upgrades the services in the cluster to CDH5 version.
    This command requires that the CDH packages in the hosts used by the
    cluster be upgraded to CDH5 before this command is issued. Once issued,
    this command will stop all running services before proceeding.

    If parcels are used instead of CDH system packages then the following
    steps need to happen in order:
    1. Stop all services manually
    2. Activate parcel
    3. Run this upgrade command

    The command will upgrade the services and their configuration to the
    version available in the CDH5 distribution.

    @return: Reference to the submitted command.
    """
    return self._cmd('upgradeServices')

  def enter_maintenance_mode(self):
    """
    Put the cluster in maintenance mode.

    @return: Reference to the completed command.
    @since: API v2
    """
    cmd = self._cmd('enterMaintenanceMode')
    if cmd.success:
      self._update(get_cluster(self._get_resource_root(), self.name))
    return cmd

  def exit_maintenance_mode(self):
    """
    Take the cluster out of maintenance mode.

    @return: Reference to the completed command.
    @since: API v2
    """
    cmd = self._cmd('exitMaintenanceMode')
    if cmd.success:
      self._update(get_cluster(self._get_resource_root(), self.name))
    return cmd

  def get_all_host_templates(self):
    """
    Retrieves all host templates in the cluster.
    @return ApiList of ApiHostTemplate objects.
    """
    return host_templates.get_all_host_templates(self._get_resource_root(), self.name)

  def get_host_template(self, name):
    """
    Retrieves a host templates by name.
    @param name: Host template name.
    @return An ApiHostTemplate object.
    """
    return host_templates.get_host_template(self._get_resource_root(), name, self.name)

  def create_host_template(self, name):
    """
    Creates a host template.
    @param name: Name of the host template to create.
    @return An ApiHostTemplate object.
    """
    return host_templates.create_host_template(self._get_resource_root(), name, self.name)

  def delete_host_template(self, name):
    """
    Deletes a host template.
    @param name: Name of the host template to delete.
    @return An ApiHostTemplate object.
    """
    return host_templates.delete_host_template(self._get_resource_root(), name, self.name)

  def rolling_restart(self, slave_batch_size=None,
                      slave_fail_count_threshold=None,
                      sleep_seconds=None,
                      stale_configs_only=None,
                      unupgraded_only=None,
                      roles_to_include=None,
                      restart_service_names=None):
    """
    Command to do a "best-effort" rolling restart of the given cluster,
    i.e. it does plain restart of services that cannot be rolling restarted,
    followed by first rolling restarting non-slaves and then rolling restarting
    the slave roles of services that can be rolling restarted. The slave restarts
    are done host-by-host.
    @param: slave_batch_size Number of hosts with slave roles to restart at a time
            Must be greater than 0. Default is 1.
    @param: slave_fail_count_threshold The threshold for number of slave host batches that
            are allowed to fail to restart before the entire command is considered failed.
            Must be >= 0. Default is 0.
    @param: sleep_seconds Number of seconds to sleep between restarts of slave host batches.
            Must be >=0. Default is 0.
    @param: stale_configs_only Restart roles with stale configs only. Default is false.
    @param: unupgraded_only Restart roles that haven't been upgraded yet. Default is false.
    @param: roles_to_include Role types to restart. Default is slave roles only.
    @param: restart_service_names List of specific services to restart.
    @return: Reference to the submitted command.
    @since: API v4
    """
    args = dict()
    if slave_batch_size:
      args['slaveBatchSize'] = slave_batch_size
    if slave_fail_count_threshold:
      args['slaveFailCountThreshold'] = slave_fail_count_threshold
    if sleep_seconds:
      args['sleepSeconds'] = sleep_seconds
    if stale_configs_only:
      args['staleConfigsOnly'] = stale_configs_only
    if unupgraded_only:
      args['unUpgradedOnly'] = unupgraded_only
    if roles_to_include:
      args['rolesToInclude'] = roles_to_include
    if restart_service_names:
      args['restartServiceNames'] = restart_service_names

    return self._cmd('rollingRestart', data=args, api_version=4)
<commit_msg>[cm_api] Expose cluster restart<commit_after># Licensed to Cloudera, Inc. under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  Cloudera, Inc. 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 cm_api.endpoints.types import *
from cm_api.endpoints import services, parcels, host_templates

__docformat__ = "epytext"

CLUSTERS_PATH = "/clusters"

def create_cluster(resource_root, name, version):
  """
  Create a cluster
  @param resource_root: The root Resource object.
  @param name: Cluster name
  @param version: Cluster CDH version
  @return: An ApiCluster object
  """
  apicluster = ApiCluster(resource_root, name, version)
  return call(resource_root.post, CLUSTERS_PATH, ApiCluster, True,
      data=[apicluster])[0]

def get_cluster(resource_root, name):
  """
  Lookup a cluster by name
  @param resource_root: The root Resource object.
  @param name: Cluster name
  @return: An ApiCluster object
  """
  return call(resource_root.get, "%s/%s" % (CLUSTERS_PATH, name), ApiCluster)

def get_all_clusters(resource_root, view=None):
  """
  Get all clusters
  @param resource_root: The root Resource object.
  @return: A list of ApiCluster objects.
  """
  return call(resource_root.get, CLUSTERS_PATH, ApiCluster, True,
      params=view and dict(view=view) or None)

def delete_cluster(resource_root, name):
  """
  Delete a cluster by name
  @param resource_root: The root Resource object.
  @param name: Cluster name
  @return: The deleted ApiCluster object
  """
  return call(resource_root.delete, "%s/%s" % (CLUSTERS_PATH, name), ApiCluster)

class ApiCluster(BaseApiResource):
  _ATTRIBUTES = {
    'name'              : None,
    'version'           : None,
    'maintenanceMode'   : ROAttr(),
    'maintenanceOwners' : ROAttr(),
  }

  def __init__(self, resource_root, name=None, version=None):
    BaseApiObject.init(self, resource_root, locals())

  def __str__(self):
    return "<ApiCluster>: %s; version: %s" % (self.name, self.version)

  def _path(self):
    return "%s/%s" % (CLUSTERS_PATH, self.name)

  def _put(self, dic, params=None):
    """Change cluster attributes"""
    cluster = self._put('', ApiCluster, data=dic, params=params)
    self._update(cluster)
    return self

  def get_commands(self, view=None):
    """
    Retrieve a list of running commands for this cluster.

    @param view: View to materialize ('full' or 'summary')
    @return: A list of running commands.
    """
    return self._get("commands", ApiCommand, True,
        params = view and dict(view=view) or None)

  def rename(self, newname):
    """
    Rename a cluster.

    @param newname: New cluster name
    @return: An ApiCluster object
    @since: API v2
    """
    dic = self.to_json_dict()
    dic['name'] = newname
    return self._put(dic)

  def create_service(self, name, service_type):
    """
    Create a service.

    @param name: Service name
    @param service_type: Service type
    @return: An ApiService object
    """
    return services.create_service(self._get_resource_root(), name,
        service_type, self.name)

  def delete_service(self, name):
    """
    Delete a service by name.

    @param name Service name
    @return The deleted ApiService object
    """
    return services.delete_service(self._get_resource_root(), name, self.name)

  def get_service(self, name):
    """
    Lookup a service by name.

    @param name: Service name
    @return: An ApiService object
    """
    return services.get_service(self._get_resource_root(), name, self.name)

  def get_all_services(self, view = None):
    """
    Get all services in this cluster.

    @return: A list of ApiService objects.
    """
    return services.get_all_services(self._get_resource_root(), self.name, view)

  def get_parcel(self, product, version):
    """
    Lookup a parcel by product and version.

    @param product: the product name
    @param version: the product version
    @return: An ApiParcel object
    """
    return parcels.get_parcel(self._get_resource_root(), product, version, self.name)

  def get_all_parcels(self, view = None):
    """
    Get all parcels in this cluster.

    @return: A list of ApiParcel objects.
    """
    return parcels.get_all_parcels(self._get_resource_root(), self.name, view)

  def list_hosts(self):
    """
    Lists all the hosts that are associated with this cluster.

    @return: A list of ApiHostRef objects of the hosts in the cluster.
    @since: API v3
    """
    return self._get("hosts", ApiHostRef, True, api_version=3)

  def remove_host(self, hostId):
    """
    Removes the association of the host with the cluster.

    @return: A ApiHostRef of the host that was removed.
    @since: API v3
    """
    return self._delete("hosts/" + hostId, ApiHostRef, api_version=3)

  def remove_all_hosts(self):
    """
    Removes the association of all the hosts with the cluster.

    @return: A list of ApiHostRef objects of the hosts that were removed.
    @since: API v3
    """
    return self._delete("hosts", ApiHostRef, True, api_version=3)

  def add_hosts(self, hostIds):
    """
    Adds a host to the cluster.

    @param hostIds: List of IDs of hosts to add to cluster.
    @return: A list of ApiHostRef objects of the new
             hosts that were added to the cluster
    @since: API v3
    """
    hostRefList = [ApiHostRef(self._get_resource_root(), x) for x in hostIds]
    return self._post("hosts", ApiHostRef, True, data=hostRefList,
        api_version=3)

  def start(self):
    """
    Start all services in a cluster, respecting dependencies.

    @return: Reference to the submitted command.
    """
    return self._cmd('start')

  def stop(self):
    """
    Stop all services in a cluster, respecting dependencies.

    @return: Reference to the submitted command.
    """
    return self._cmd('stop')

  def restart(self):
    """
    Restart all services in the cluster.
    Services are restarted in the appropriate order given their dependencies.

    @return: Reference to the submitted command.
    """
    return self._cmd('restart')

def deploy_client_config(self):
    """
    Deploys client configuration to the hosts on the cluster.

    @return: Reference to the submitted command.
    @since: API v2
    """
    return self._cmd('deployClientConfig')

  def upgrade_services(self):
    """
    Upgrades the services in the cluster to CDH5 version.
    This command requires that the CDH packages in the hosts used by the
    cluster be upgraded to CDH5 before this command is issued. Once issued,
    this command will stop all running services before proceeding.

    If parcels are used instead of CDH system packages then the following
    steps need to happen in order:
    1. Stop all services manually
    2. Activate parcel
    3. Run this upgrade command

    The command will upgrade the services and their configuration to the
    version available in the CDH5 distribution.

    @return: Reference to the submitted command.
    """
    return self._cmd('upgradeServices')

  def enter_maintenance_mode(self):
    """
    Put the cluster in maintenance mode.

    @return: Reference to the completed command.
    @since: API v2
    """
    cmd = self._cmd('enterMaintenanceMode')
    if cmd.success:
      self._update(get_cluster(self._get_resource_root(), self.name))
    return cmd

  def exit_maintenance_mode(self):
    """
    Take the cluster out of maintenance mode.

    @return: Reference to the completed command.
    @since: API v2
    """
    cmd = self._cmd('exitMaintenanceMode')
    if cmd.success:
      self._update(get_cluster(self._get_resource_root(), self.name))
    return cmd

  def get_all_host_templates(self):
    """
    Retrieves all host templates in the cluster.
    @return ApiList of ApiHostTemplate objects.
    """
    return host_templates.get_all_host_templates(self._get_resource_root(), self.name)

  def get_host_template(self, name):
    """
    Retrieves a host templates by name.
    @param name: Host template name.
    @return An ApiHostTemplate object.
    """
    return host_templates.get_host_template(self._get_resource_root(), name, self.name)

  def create_host_template(self, name):
    """
    Creates a host template.
    @param name: Name of the host template to create.
    @return An ApiHostTemplate object.
    """
    return host_templates.create_host_template(self._get_resource_root(), name, self.name)

  def delete_host_template(self, name):
    """
    Deletes a host template.
    @param name: Name of the host template to delete.
    @return An ApiHostTemplate object.
    """
    return host_templates.delete_host_template(self._get_resource_root(), name, self.name)

  def rolling_restart(self, slave_batch_size=None,
                      slave_fail_count_threshold=None,
                      sleep_seconds=None,
                      stale_configs_only=None,
                      unupgraded_only=None,
                      roles_to_include=None,
                      restart_service_names=None):
    """
    Command to do a "best-effort" rolling restart of the given cluster,
    i.e. it does plain restart of services that cannot be rolling restarted,
    followed by first rolling restarting non-slaves and then rolling restarting
    the slave roles of services that can be rolling restarted. The slave restarts
    are done host-by-host.
    @param: slave_batch_size Number of hosts with slave roles to restart at a time
            Must be greater than 0. Default is 1.
    @param: slave_fail_count_threshold The threshold for number of slave host batches that
            are allowed to fail to restart before the entire command is considered failed.
            Must be >= 0. Default is 0.
    @param: sleep_seconds Number of seconds to sleep between restarts of slave host batches.
            Must be >=0. Default is 0.
    @param: stale_configs_only Restart roles with stale configs only. Default is false.
    @param: unupgraded_only Restart roles that haven't been upgraded yet. Default is false.
    @param: roles_to_include Role types to restart. Default is slave roles only.
    @param: restart_service_names List of specific services to restart.
    @return: Reference to the submitted command.
    @since: API v4
    """
    args = dict()
    if slave_batch_size:
      args['slaveBatchSize'] = slave_batch_size
    if slave_fail_count_threshold:
      args['slaveFailCountThreshold'] = slave_fail_count_threshold
    if sleep_seconds:
      args['sleepSeconds'] = sleep_seconds
    if stale_configs_only:
      args['staleConfigsOnly'] = stale_configs_only
    if unupgraded_only:
      args['unUpgradedOnly'] = unupgraded_only
    if roles_to_include:
      args['rolesToInclude'] = roles_to_include
    if restart_service_names:
      args['restartServiceNames'] = restart_service_names

    return self._cmd('rollingRestart', data=args, api_version=4)
<|endoftext|>


<commit_before>"""
Data resource for upload data
"""
from flask_restful import Resource
from flask import request, Response
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker
import json
import os
import uuid
import datetime
import logging
from pprint import pprint

from meerkat_nest import model
from meerkat_nest import config
from meerkat_nest.util import scramble, format_form_field_key, validate_request
from meerkat_nest import message_service

db_url = os.environ['MEERKAT_NEST_DB_URL']
engine = create_engine(db_url)

class UploadData(Resource):
    """
    Receives JSON data and stores it in Meerkat Nest database

    Returns:\n
        HTTP return code\n
    """

    def get(self):
        return "upload data GET"

    def post(self):

        logging.debug("received upload request")
        logging.debug(str(request.headers))
        data_entry = request.get_json()
        logging.debug(str(data_entry))

        # Validate the request
        try:
            validate_request(data_entry)
        except AssertionError as e:
            msg = "Input was not a valid Meerkat Nest JSON object: " + e.args[0]
            logging.error(msg)
            return Response(json.dumps({"message": msg}),
                            status=400,
                            mimetype='application/json')
        # Upload the data entry in to raw data storage
        try:
            uuid_pk = upload_to_raw_data(data_entry)
            # data_entry['uuid'] = uuid_pk
        except AssertionError as e:
            msg = "Raw input type '" + data_entry['content'] + "' is not supported"
            logging.error(msg)
            return Response(json.dumps({"message": msg}),
                            status=400,
                            mimetype='application/json')
        except Exception as e:
            msg = "Error in uploading data: " + e.args[0]
            logging.error(msg)
            return Response(json.dumps({"message": msg}),
                            status=502,
                            mimetype='application/json')

        # Process the data entry
        try:
            processed_data_entry = process(data_entry)
        except AssertionError as e:
            msg = "Data type '" + data_entry['formId'] + "' is not supported for input type '"\
                + data_entry['content'] + "'"
            logging.error(msg)
            return Response(json.dumps({"message": msg}),
                            status=400,
                            mimetype='application/json')

        # Send processed data forward to the cloud
        try:
            sent = message_service.send_data(processed_data_entry)
        except AssertionError as e:
            msg = "Error in forwarding data to message queue: " + str(e)
            logging.error(msg)
            return Response(json.dumps({"message": msg}),
                            status=502,
                            mimetype='application/json')

        logging.debug("processed upload request")
        return Response(json.dumps(processed_data_entry),
                        status=200,
                        mimetype='application/json')

def upload_to_raw_data(data_entry):
    """
    Stores raw data in Meerkat Nest database

    Returns:\n
        uuid for the PK of the raw data row\n
    """

    uuid_pk = str(uuid.uuid4())

    insert_row = None

    if data_entry['content'] == 'record':
        insert_row = model.RawDataOdkCollect.__table__.insert().values(
                uuid=uuid_pk,
                received_on=datetime.datetime.now(),
                active_from=datetime.datetime.now(),
                authentication_token=data_entry['token'],
                content=data_entry['content'],
                formId=data_entry['formId'],
                formVersion=data_entry['formVersion'],
                data=data_entry['data']
        )

    assert insert_row is not None, "Content handling not implemented"

    connection = engine.connect()
    connection.execute(insert_row)
    connection.close()

    return uuid_pk

def process(data_entry):
    """
    Processes raw data and stores the processed data entry in in Meerkat Nest database

    Returns:\n
        processed data_entry if processing was successful, False otherwise
    """
    processed_data_entry = restructure_aggregate_data(data_entry)
    processed_data_entry = scramble_fields(processed_data_entry)
    processed_data_entry = format_field_keys(processed_data_entry)

    assert data_entry['content'] in ['form', 'record'], "Content not supported"
    assert data_entry['formId'] in config.country_config['tables'], "Form not supported"

    insert_row = model.data_type_tables[processed_data_entry['formId']].__table__.insert().values(
           uuid=processed_data_entry['uuid'],
           data=processed_data_entry['data']
        )

    try:
        connection = engine.connect()
        connection.execute(insert_row)
        connection.close()
        return processed_data_entry
    except Exception as e:
        raise

def restructure_aggregate_data(data_entry):
    """
    Restructures data from aggregate JSON feed

    Returns:\n
        restructured data entry
    """

    restructured_data = data_entry['data'][0]
    data_entry['data'] = restructured_data

    data_entry['uuid'] = data_entry['data']['*meta-instance-id*']

    return data_entry

def scramble_fields(data_entry):
    """
    Scrambles fields in data entry based on configurations

    Returns:\n
        data entry structure with scrambled fields
    """

    data_entry_scrambled = data_entry

    try:
        fields = config.country_config['scramble_fields'][data_entry['formId']]
    except KeyError as e:
        return data_entry_scrambled

    for field in fields:
        if field in data_entry_scrambled['data'].keys():
            data_entry_scrambled['data'][field] = scramble(data_entry_scrambled['data'][field])

    return data_entry_scrambled

def format_field_keys(data_entry):
    """
    Formats the field names in the data entry

    Returns:\n
        data entry structure with formatted field namess
    """

    for key in data_entry['data'].keys():
        new_key = format_form_field_key(key)
        if new_key != key:
            data_entry['data'].update({new_key: data_entry['data'][key]})
            data_entry['data'].pop(key)

    return data_entry
<commit_msg>refactor key handling<commit_after>"""
Data resource for upload data
"""
from flask_restful import Resource
from flask import request, Response
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker
import json
import os
import uuid
import datetime
import logging
from pprint import pprint

from meerkat_nest import model
from meerkat_nest import config
from meerkat_nest.util import scramble, format_form_field_key, validate_request
from meerkat_nest import message_service

db_url = os.environ['MEERKAT_NEST_DB_URL']
engine = create_engine(db_url)

class UploadData(Resource):
    """
    Receives JSON data and stores it in Meerkat Nest database

    Returns:\n
        HTTP return code\n
    """

    def get(self):
        return "upload data GET"

    def post(self):

        logging.debug("received upload request")
        logging.debug(str(request.headers))
        data_entry = request.get_json()
        logging.debug(str(data_entry))

        # Validate the request
        try:
            validate_request(data_entry)
        except AssertionError as e:
            msg = "Input was not a valid Meerkat Nest JSON object: " + e.args[0]
            logging.error(msg)
            return Response(json.dumps({"message": msg}),
                            status=400,
                            mimetype='application/json')
        # Upload the data entry in to raw data storage
        try:
            uuid_pk = upload_to_raw_data(data_entry)
            # data_entry['uuid'] = uuid_pk
        except AssertionError as e:
            msg = "Raw input type '" + data_entry['content'] + "' is not supported"
            logging.error(msg)
            return Response(json.dumps({"message": msg}),
                            status=400,
                            mimetype='application/json')
        except Exception as e:
            msg = "Error in uploading data: " + e.args[0]
            logging.error(msg)
            return Response(json.dumps({"message": msg}),
                            status=502,
                            mimetype='application/json')

        # Process the data entry
        try:
            processed_data_entry = process(data_entry)
        except AssertionError as e:
            msg = "Data type '" + data_entry['formId'] + "' is not supported for input type '"\
                + data_entry['content'] + "'"
            logging.error(msg)
            return Response(json.dumps({"message": msg}),
                            status=400,
                            mimetype='application/json')

        # Send processed data forward to the cloud
        try:
            sent = message_service.send_data(processed_data_entry)
        except AssertionError as e:
            msg = "Error in forwarding data to message queue: " + str(e)
            logging.error(msg)
            return Response(json.dumps({"message": msg}),
                            status=502,
                            mimetype='application/json')

        logging.debug("processed upload request")
        return Response(json.dumps(processed_data_entry),
                        status=200,
                        mimetype='application/json')

def upload_to_raw_data(data_entry):
    """
    Stores raw data in Meerkat Nest database

    Returns:\n
        uuid for the PK of the raw data row\n
    """

    uuid_pk = str(uuid.uuid4())

    insert_row = None

    if data_entry['content'] == 'record':
        insert_row = model.RawDataOdkCollect.__table__.insert().values(
                uuid=uuid_pk,
                received_on=datetime.datetime.now(),
                active_from=datetime.datetime.now(),
                authentication_token=data_entry['token'],
                content=data_entry['content'],
                formId=data_entry['formId'],
                formVersion=data_entry['formVersion'],
                data=data_entry['data']
        )

    assert insert_row is not None, "Content handling not implemented"

    connection = engine.connect()
    connection.execute(insert_row)
    connection.close()

    return uuid_pk

def process(data_entry):
    """
    Processes raw data and stores the processed data entry in in Meerkat Nest database

    Returns:\n
        processed data_entry if processing was successful, False otherwise
    """
    processed_data_entry = restructure_aggregate_data(data_entry)
    processed_data_entry = scramble_fields(processed_data_entry)
    processed_data_entry = format_field_keys(processed_data_entry)

    assert data_entry['content'] in ['form', 'record'], "Content not supported"
    assert data_entry['formId'] in config.country_config['tables'], "Form not supported"

    insert_row = model.data_type_tables[processed_data_entry['formId']].__table__.insert().values(
           uuid=processed_data_entry['uuid'],
           data=processed_data_entry['data']
        )

    try:
        connection = engine.connect()
        connection.execute(insert_row)
        connection.close()
        return processed_data_entry
    except Exception as e:
        raise

def restructure_aggregate_data(data_entry):
    """
    Restructures data from aggregate JSON feed

    Returns:\n
        restructured data entry
    """

    restructured_data = data_entry['data'][0]
    data_entry['data'] = restructured_data

    data_entry['uuid'] = data_entry['data']['*meta-instance-id*']

    return data_entry

def scramble_fields(data_entry):
    """
    Scrambles fields in data entry based on configurations

    Returns:\n
        data entry structure with scrambled fields
    """

    data_entry_scrambled = data_entry

    try:
        fields = config.country_config['scramble_fields'][data_entry['formId']]
    except KeyError as e:
        return data_entry_scrambled

    for field in fields:
        if field in data_entry_scrambled['data'].keys():
            data_entry_scrambled['data'][field] = scramble(data_entry_scrambled['data'][field])

    return data_entry_scrambled

def format_field_keys(data_entry):
    """
    Formats the field names in the data entry

    Returns:\n
        data entry structure with formatted field namess
    """

    new_data_entry = {}

    for key in data_entry['data'].keys():
        new_key = format_form_field_key(key)
        new_data_entry['data'].update({new_key: data_entry['data'][key]})

    return data_entry
<|endoftext|>