<commit_before>import os
from typing import Optional

from bs4 import BeautifulSoup

import config
from webweeder import configutils
from webweeder import utils
from webweeder.domainconfig import DomainConfig
from webweeder.models import PageMetadata, page_metadata_from_json

class MonsterWeeder:
    # TODO: better logging, log each page as we process it or something
    # TODO: better error handling, one screwy HTML file shouldn't prevent weeding the rest

    def weed(self):
        self._weed_subdirectory(config.OUTPUT_DIRECTORY)

    def _weed_file(self, file: str):
        # Read and parse JSON file
        json_str = utils.read_text_file(file, missing_ok=False)
        meta: PageMetadata = page_metadata_from_json(json_str)

        domain_config = configutils.get_config_for_domain(meta.domain)
        if domain_config is None:
            raise RuntimeError('Unconfigured domain "%s" for page: %r' % (meta.domain, meta))

        raw_html_path = os.path.join(config.OUTPUT_DIRECTORY, meta.directory, meta.file_raw_html)
        raw_html = utils.read_text_file(raw_html_path, missing_ok=False)

        soup = BeautifulSoup(raw_html, config.HTML_PARSER)
        plaintext = self._extract_article_content(soup, domain_config)

        plaintext_path = os.path.join(config.OUTPUT_DIRECTORY, meta.directory, meta.file_article_plaintext)
        utils.write_text_file(plaintext_path, plaintext)

    def _weed_subdirectory(self, dir: str):
        for item_name in os.listdir(dir):
            item_path = os.path.join(dir, item_name)
            if os.path.isfile(item_path):
                if item_name == 'metadata.json':
                    self._weed_file(item_path)
            elif os.path.isdir(item_path):
                self._weed_subdirectory(item_path)
            else:
                raise TypeError('Not a file or a directory: %s' % item_path)

    @staticmethod
    def _extract_article_content(soup: BeautifulSoup, domain: DomainConfig) -> Optional[str]:
        # TODO: Similar to DomainConfig.scrape_* methods, refactor needed?
        """
        IMPORTANT: This modifies the passed soup

        If configured selector is None, None is returned
        If configured selector is not None, the scraped article content is returned
        If configured selector is not None but the article content cannot be scraped, an error is raised
        """
        if domain.article_content_selector is None:
            return None
        for crap in soup.find_all(name='script'):
            crap.extract()
        for crap in soup.find_all(name='style'):
            crap.extract()
        element = soup.select_one(domain.article_content_selector)
        text = element.get_text()
        return text.strip()
<commit_msg>Major refactor for weeding - Split finding metadata out of actual weeding code - Will help with error handling, logging in future<commit_after>import os
from typing import Iterable, List, Optional

from bs4 import BeautifulSoup

import config
from webweeder import configutils
from webweeder import utils
from webweeder.domainconfig import DomainConfig
from webweeder.models import PageMetadata, page_metadata_from_json

class MonsterWeeder:
    # TODO: better logging, log each page as we process it or something
    # TODO: better error handling, one screwy HTML file shouldn't prevent weeding the rest

    def find_page_metadatas(self, directory: str) -> Iterable[str]:
        """
        :param directory: Where to look for metadata files
        :return: List of paths for the metadata files of each found page in the given directory, recursively
        """
        found_pages: List[str] = []
        for item_name in os.listdir(directory):
            item_path = os.path.join(directory, item_name)
            if os.path.isfile(item_path):
                if item_name == 'metadata.json':
                    found_pages.append(item_path)
            elif os.path.isdir(item_path):
                found_pages += self.find_page_metadatas(item_path)
            else:
                raise TypeError('Not a file or a directory: %s' % item_path)
        return found_pages

    def weed_page(self, base_dir: str, file: str):
        """
        :param base_dir: Top-level directory containing the raw HTML and plaintext files referenced by the metadata
        :param file: Full path to the metadata (i.e. already joined to base_dir)
        :return:
        """
        # Read and parse JSON file
        json_str = utils.read_text_file(file, missing_ok=False)
        meta: PageMetadata = page_metadata_from_json(json_str)

        domain_config = configutils.get_config_for_domain(meta.domain)
        if domain_config is None:
            raise RuntimeError('Unconfigured domain "%s" for page: %r' % (meta.domain, meta))

        # Read raw HTML for the page
        raw_html_path = os.path.join(base_dir, meta.directory, meta.file_raw_html)
        raw_html = utils.read_text_file(raw_html_path, missing_ok=False)

        # Parse HTML and extract plaintext from article
        soup = BeautifulSoup(raw_html, config.HTML_PARSER)
        plaintext = self._extract_article_content(soup, domain_config)

        # Write plaintext to file
        plaintext_path = os.path.join(base_dir, meta.directory, meta.file_article_plaintext)
        utils.write_text_file(plaintext_path, plaintext)

    @staticmethod
    def _extract_article_content(soup: BeautifulSoup, domain: DomainConfig) -> Optional[str]:
        # TODO: Similar to DomainConfig.scrape_* methods, refactor needed?
        """
        IMPORTANT: This modifies the passed soup

        If configured selector is None, None is returned
        If configured selector is not None, the scraped article content is returned
        If configured selector is not None but the article content cannot be scraped, an error is raised
        """
        if domain.article_content_selector is None:
            return None
        for crap in soup.find_all(name='script'):
            crap.extract()
        for crap in soup.find_all(name='style'):
            crap.extract()
        element = soup.select_one(domain.article_content_selector)
        text = element.get_text()
        return text.strip()
<|endoftext|>


<commit_before>import datetime

from memoized import memoized

from dimagi.ext.couchdbkit import (
    BooleanProperty,
    DateTimeProperty,
    Document,
    StringProperty,
)
from django.db import DEFAULT_DB_ALIAS, models

from corehq.apps.domain.models import Domain

class RegistrationRequest(models.Model):
    tos_confirmed = models.BooleanField(default=False)
    request_time = models.DateTimeField()
    request_ip = models.CharField(max_length=31)
    activation_guid = models.CharField(max_length=126, unique=True)
    confirm_time = models.DateTimeField(null=True)
    confirm_ip = models.CharField(max_length=31, null=True)
    domain = models.CharField(max_length=255, null=True)
    new_user_username = models.CharField(max_length=255, null=True)
    requesting_user_username = models.CharField(max_length=255, null=True)

    class Meta:
        db_table = "registration_registrationrequest"

    @property
    @memoized
    def project(self):
        return Domain.get_by_name(self.domain)

    @classmethod
    def get_by_guid(cls, guid):
        return RegistrationRequest.objects.filter(activation_guid=guid).first()

    @classmethod
    def get_requests_today(cls):
        today = datetime.datetime.utcnow()
        yesterday = today - datetime.timedelta(1)
        return RegistrationRequest.objects.filter(
            request_time__gte=yesterday.isoformat(),
            request_time__lte=today.isoformat(),
        ).count()

    @classmethod
    def get_requests_24hrs_ago(cls):
        today = datetime.datetime.utcnow()
        yesterday = today - datetime.timedelta(1)
        join_on_start = datetime.datetime(
            yesterday.year, yesterday.month, yesterday.day, yesterday.hour, 0, 0, 0)
        join_on_end = datetime.datetime(
            yesterday.year, yesterday.month, yesterday.day, yesterday.hour, 59, 59, 59)
        requests = RegistrationRequest.objects.filter(
            request_time__gte=join_on_start,
            request_time__lte=join_on_end,
            confirm_time__isnull=True
        )
        return [req for req in requests if req.new_user_username == req.requesting_user_username]

    @classmethod
    def get_request_for_username(cls, username):
        return RegistrationRequest.objects.filter(new_user_username=username).first()
<commit_msg>Removed unused import (PR3)<commit_after>import datetime

from memoized import memoized

from dimagi.ext.couchdbkit import (
    BooleanProperty,
    DateTimeProperty,
    Document,
    StringProperty,
)
from django.db import models

from corehq.apps.domain.models import Domain

class RegistrationRequest(models.Model):
    tos_confirmed = models.BooleanField(default=False)
    request_time = models.DateTimeField()
    request_ip = models.CharField(max_length=31)
    activation_guid = models.CharField(max_length=126, unique=True)
    confirm_time = models.DateTimeField(null=True)
    confirm_ip = models.CharField(max_length=31, null=True)
    domain = models.CharField(max_length=255, null=True)
    new_user_username = models.CharField(max_length=255, null=True)
    requesting_user_username = models.CharField(max_length=255, null=True)

    class Meta:
        db_table = "registration_registrationrequest"

    @property
    @memoized
    def project(self):
        return Domain.get_by_name(self.domain)

    @classmethod
    def get_by_guid(cls, guid):
        return RegistrationRequest.objects.filter(activation_guid=guid).first()

    @classmethod
    def get_requests_today(cls):
        today = datetime.datetime.utcnow()
        yesterday = today - datetime.timedelta(1)
        return RegistrationRequest.objects.filter(
            request_time__gte=yesterday.isoformat(),
            request_time__lte=today.isoformat(),
        ).count()

    @classmethod
    def get_requests_24hrs_ago(cls):
        today = datetime.datetime.utcnow()
        yesterday = today - datetime.timedelta(1)
        join_on_start = datetime.datetime(
            yesterday.year, yesterday.month, yesterday.day, yesterday.hour, 0, 0, 0)
        join_on_end = datetime.datetime(
            yesterday.year, yesterday.month, yesterday.day, yesterday.hour, 59, 59, 59)
        requests = RegistrationRequest.objects.filter(
            request_time__gte=join_on_start,
            request_time__lte=join_on_end,
            confirm_time__isnull=True
        )
        return [req for req in requests if req.new_user_username == req.requesting_user_username]

    @classmethod
    def get_request_for_username(cls, username):
        return RegistrationRequest.objects.filter(new_user_username=username).first()
<|endoftext|>


<commit_before>from decimal import Decimal

from django.core.exceptions import ObjectDoesNotExist
from django.db import models

from corehq.apps.accounting import models as accounting
from corehq.apps.accounting.models import Currency
from corehq.apps.accounting.utils import EXCHANGE_RATE_DECIMAL_PLACES
from corehq.apps.sms.models import DIRECTION_CHOICES, SQLMobileBackend
from corehq.apps.sms.phonenumbers_helper import get_country_code_and_national_number
from corehq.apps.smsbillables.utils import log_smsbillables_error
from corehq.messaging.smsbackends.test.models import SQLTestSMSBackend
from corehq.apps.sms.util import clean_phone_number
from corehq.apps.smsbillables.exceptions import AmbiguousPrefixException
from corehq.util.quickcache import quickcache

class SmsGatewayFeeCriteria(models.Model):
    """
    These are the parameters we'll use to try and calculate the cost of sending a message through
    our gateways. We configure the SMS fee criteria based on parameters given to us by specific
    gateway providers.

    Nullable fields indicate criteria that can be applied globally to all messages with no specific matches
    for that field.
    """
    backend_api_id = models.CharField(max_length=100, db_index=True)
    backend_instance = models.CharField(max_length=255, db_index=True, null=True)
    direction = models.CharField(max_length=10, db_index=True, choices=DIRECTION_CHOICES)
    country_code = models.IntegerField(max_length=5, null=True, blank=True, db_index=True)
    prefix = models.CharField(max_length=10, blank=True, default="", db_index=True)

    class Meta:
        app_label = 'smsbillables'

    @classmethod
    def get_most_specific(cls, backend_api_id, direction,
                          backend_instance=None, country_code=None, national_number=None):
        """
        Gets the most specific criteria available based on (and in order of preference for optional):
        - backend_api_id
        - direction
        - backend_instance (optional)
        - country_code and prefix (optional)
        """
        all_possible_criteria = cls.objects.filter(backend_api_id=backend_api_id, direction=direction)

        if all_possible_criteria.count() == 0:
            return None

        national_number = national_number or ""

        def get_criteria_with_longest_matching_prefix(criteria_list):
            if len(set(criteria.prefix for criteria in criteria_list)) != len(criteria_list):
                raise AmbiguousPrefixException(
                    ", ".join(
                        "%(country_code)d, '%(prefix)s'" % {
                            "country_code": criteria.country_code,
                            "prefix": criteria.prefix,
                        } for criteria in criteria_list
                    )
                )
            criteria_list.sort(key=(lambda criteria: len(criteria.prefix)), reverse=True)
            for criteria in criteria_list:
                if national_number.startswith(criteria.prefix):
                    return criteria
            raise ObjectDoesNotExist

        try:
            return get_criteria_with_longest_matching_prefix(
                list(all_possible_criteria.filter(country_code=country_code, backend_instance=backend_instance))
            )
        except ObjectDoesNotExist:
            pass
        try:
            return all_possible_criteria.get(country_code=None, backend_instance=backend_instance)
        except ObjectDoesNotExist:
            pass
        try:
            return get_criteria_with_longest_matching_prefix(
                list(all_possible_criteria.filter(country_code=country_code, backend_instance=None))
            )
        except ObjectDoesNotExist:
            pass
        try:
            return all_possible_criteria.get(country_code=None, backend_instance=None)
        except ObjectDoesNotExist:
            pass

        return None

class SmsGatewayFee(models.Model):
    """
    The fee for sending or receiving an SMS Message based on gateway.
    When an SmsBillable is calculated, it will use the most recent SmsFee available from the criteria
    to determine the gateway_charge.

    Once an SmsFee is created, it cannot be modified.
    """
    criteria = models.ForeignKey(SmsGatewayFeeCriteria, on_delete=models.PROTECT)
    amount = models.DecimalField(default=0.0, max_digits=10, decimal_places=4)
    currency = models.ForeignKey(accounting.Currency, on_delete=models.PROTECT)
    date_created = models.DateTimeField(auto_now_add=True)

    class Meta:
        app_label = 'smsbillables'

    @classmethod
    def create_new(cls, backend_api_id, direction, amount,
                   currency=None, backend_instance=None, country_code=None, prefix=None,
                   save=True, fee_class=None, criteria_class=None):
        fee_class = fee_class or cls
        criteria_class = criteria_class or SmsGatewayFeeCriteria
        currency = currency or Currency.get_default()

        if 'prefix' in [
            field.name
            for field, _ in criteria_class._meta.get_fields_with_model()
        ]:
            prefix = prefix or ''
            criteria, _ = criteria_class.objects.get_or_create(
                backend_api_id=backend_api_id,
                direction=direction,
                backend_instance=backend_instance,
                country_code=country_code,
                prefix=prefix,
            )
        else:
            criteria, _ = criteria_class.objects.get_or_create(
                backend_api_id=backend_api_id,
                direction=direction,
                backend_instance=backend_instance,
                country_code=country_code,
            )
        new_fee = fee_class(
            currency=currency,
            amount=amount,
            criteria=criteria
        )
        if save:
            new_fee.save()
        return new_fee

    @classmethod
    def get_by_criteria(cls, backend_api_id, direction,
                        backend_instance=None, country_code=None, national_number=None):
        criteria = SmsGatewayFeeCriteria.get_most_specific(
            backend_api_id,
            direction,
            backend_instance=backend_instance,
            country_code=country_code,
            national_number=national_number,
        )
        return cls.get_by_criteria_obj(criteria)

    @classmethod
    def get_by_criteria_obj(cls, criteria):
        if not criteria:
            return None
        return cls.objects.filter(criteria=criteria.id).latest('date_created')

class SmsUsageFeeCriteria(models.Model):
    """
    Criteria for determining a usage fee applied for each SMS message sent or received.

    Nullable fields indicate criteria that can be applied globally to all messages with no specific matches
    for that field.
    """
    direction = models.CharField(max_length=10, db_index=True, choices=DIRECTION_CHOICES)
    domain = models.CharField(max_length=25, db_index=True, null=True)

    class Meta:
        app_label = 'smsbillables'

    @classmethod
    def get_most_specific(cls, direction, domain=None):
        """
        Gets the most specific criteria available based on (and in order of preference for optional):
        - direction
        - domain (optional)
        """
        all_possible_criteria = cls.objects.filter(direction=direction)

        if all_possible_criteria.count() == 0:
            return None

        try:
            return all_possible_criteria.get(domain=domain)
        except ObjectDoesNotExist:
            pass
        try:
            return all_possible_criteria.get(domain=None)
        except ObjectDoesNotExist:
            pass

        return None

class SmsUsageFee(models.Model):
    """
    The usage fee, with version information, based on domain or globally.
    When an SmsBillable is calculated, it will use the most recent SmsUsageFee available from the
    criteria to determine the usage_charge.

    Currency is always in USD since this is something we control.

    Once an SmsUsageFee is created, it cannot be modified.
    """
    criteria = models.ForeignKey(SmsUsageFeeCriteria, on_delete=models.PROTECT)
    amount = models.DecimalField(default=0.0, max_digits=10, decimal_places=4)
    date_created = models.DateTimeField(auto_now_add=True)

    class Meta:
        app_label = 'smsbillables'

    @classmethod
    def create_new(cls, direction, amount, domain=None, save=True):
        criteria, _ = SmsUsageFeeCriteria.objects.get_or_create(
            domain=domain, direction=direction,
        )
        new_fee = SmsUsageFee(
            amount=amount,
            criteria=criteria
        )
        if save:
            new_fee.save()
        return new_fee

    @classmethod
    def get_by_criteria(cls, direction, domain=None):
        criteria = SmsUsageFeeCriteria.get_most_specific(direction, domain=domain)
        if not criteria:
            return None
        return cls.objects.filter(criteria=criteria.id).latest('date_created')

@quickcache(['sms_backend_id'])
def _sms_backend_is_global(sms_backend_id):
    return SQLMobileBackend.load(
        sms_backend_id,
        is_couch_id=True,
        include_deleted=True
    ).is_global

class SmsBillable(models.Model):
    """
    A record of matching a fee to a particular SMS.

    If on closer inspection we determine a particular SmsBillable is invalid (whether something is
    awry with the api_response, or we used the incorrect fee and want to recalculate) we can set
    this billable to is_valid = False and it will not be used toward calculating the SmsLineItem in
    the monthly Invoice.
    """
    gateway_fee = models.ForeignKey(SmsGatewayFee, null=True, on_delete=models.PROTECT)
    gateway_fee_conversion_rate = models.DecimalField(default=Decimal('1.0'), null=True, max_digits=20,
                                                      decimal_places=EXCHANGE_RATE_DECIMAL_PLACES)
    usage_fee = models.ForeignKey(SmsUsageFee, null=True, on_delete=models.PROTECT)
    log_id = models.CharField(max_length=50, db_index=True)
    phone_number = models.CharField(max_length=50)
    api_response = models.TextField(null=True, blank=True)
    is_valid = models.BooleanField(default=True, db_index=True)
    domain = models.CharField(max_length=25, db_index=True)
    direction = models.CharField(max_length=10, db_index=True, choices=DIRECTION_CHOICES)
    date_sent = models.DateField()
    date_created = models.DateField(auto_now_add=True)

    class Meta:
        app_label = 'smsbillables'

    @property
    def gateway_charge(self):
        if self.gateway_fee is not None:
            try:
                charge = SmsGatewayFee.objects.get(id=self.gateway_fee.id)
                if self.gateway_fee_conversion_rate is not None:
                    return charge.amount / self.gateway_fee_conversion_rate
                return charge.amount
            except ObjectDoesNotExist:
                pass
        return Decimal('0.0')

    @property
    def usage_charge(self):
        if self.usage_fee is not None:
            try:
                charge = SmsUsageFee.objects.get(id=self.usage_fee.id)
                return charge.amount
            except ObjectDoesNotExist:
                pass
        return Decimal('0.0')

    @classmethod
    def create(cls, message_log, api_response=None):
        phone_number = clean_phone_number(message_log.phone_number)
        direction = message_log.direction
        domain = message_log.domain
        log_id = message_log.couch_id

        billable = cls(
            log_id=log_id,
            phone_number=phone_number,
            direction=direction,
            date_sent=message_log.date,
            domain=domain,
        )
        billable.gateway_fee, billable.gateway_fee_conversion_rate = cls._get_gateway_fee(
            message_log.backend_api, message_log.backend_id, phone_number, direction, log_id
        )
        billable.usage_fee = cls._get_usage_fee(domain, direction)

        if api_response is not None:
            billable.api_response = api_response

        if message_log.backend_api == SQLTestSMSBackend.get_api_id():
            billable.is_valid = False

        billable.save()
        return billable

    @classmethod
    def _get_gateway_fee(cls, backend_api_id, backend_instance, phone_number, direction, couch_id):
        country_code, national_number = get_country_code_and_national_number(phone_number)
        is_gateway_billable = backend_instance is None or _sms_backend_is_global(backend_instance)

        if is_gateway_billable:
            gateway_fee = SmsGatewayFee.get_by_criteria(
                backend_api_id,
                direction,
                backend_instance=backend_instance,
                country_code=country_code,
                national_number=national_number,
            )
            if gateway_fee:
                conversion_rate = gateway_fee.currency.rate_to_default
                if conversion_rate != 0:
                    return gateway_fee, conversion_rate
                else:
                    log_smsbillables_error(
                        "Gateway fee conversion rate for currency %s is 0"
                        % gateway_fee.currency.code
                    )
                    return gateway_fee, None
            else:
                log_smsbillables_error(
                    "No matching gateway fee criteria for SMS %s" % couch_id
                )
        return None, None

    @classmethod
    def _get_usage_fee(cls, domain, direction):
        usage_fee = SmsUsageFee.get_by_criteria(
            direction, domain=domain
        )
        if not usage_fee:
            log_smsbillables_error(
                "Did not find usage fee for direction %s and domain %s"
                % (direction, domain)
            )
        return usage_fee
<commit_msg>remove references to api_response<commit_after>from decimal import Decimal

from django.core.exceptions import ObjectDoesNotExist
from django.db import models

from corehq.apps.accounting import models as accounting
from corehq.apps.accounting.models import Currency
from corehq.apps.accounting.utils import EXCHANGE_RATE_DECIMAL_PLACES
from corehq.apps.sms.models import DIRECTION_CHOICES, SQLMobileBackend
from corehq.apps.sms.phonenumbers_helper import get_country_code_and_national_number
from corehq.apps.smsbillables.utils import log_smsbillables_error
from corehq.messaging.smsbackends.test.models import SQLTestSMSBackend
from corehq.apps.sms.util import clean_phone_number
from corehq.apps.smsbillables.exceptions import AmbiguousPrefixException
from corehq.util.quickcache import quickcache

class SmsGatewayFeeCriteria(models.Model):
    """
    These are the parameters we'll use to try and calculate the cost of sending a message through
    our gateways. We configure the SMS fee criteria based on parameters given to us by specific
    gateway providers.

    Nullable fields indicate criteria that can be applied globally to all messages with no specific matches
    for that field.
    """
    backend_api_id = models.CharField(max_length=100, db_index=True)
    backend_instance = models.CharField(max_length=255, db_index=True, null=True)
    direction = models.CharField(max_length=10, db_index=True, choices=DIRECTION_CHOICES)
    country_code = models.IntegerField(max_length=5, null=True, blank=True, db_index=True)
    prefix = models.CharField(max_length=10, blank=True, default="", db_index=True)

    class Meta:
        app_label = 'smsbillables'

    @classmethod
    def get_most_specific(cls, backend_api_id, direction,
                          backend_instance=None, country_code=None, national_number=None):
        """
        Gets the most specific criteria available based on (and in order of preference for optional):
        - backend_api_id
        - direction
        - backend_instance (optional)
        - country_code and prefix (optional)
        """
        all_possible_criteria = cls.objects.filter(backend_api_id=backend_api_id, direction=direction)

        if all_possible_criteria.count() == 0:
            return None

        national_number = national_number or ""

        def get_criteria_with_longest_matching_prefix(criteria_list):
            if len(set(criteria.prefix for criteria in criteria_list)) != len(criteria_list):
                raise AmbiguousPrefixException(
                    ", ".join(
                        "%(country_code)d, '%(prefix)s'" % {
                            "country_code": criteria.country_code,
                            "prefix": criteria.prefix,
                        } for criteria in criteria_list
                    )
                )
            criteria_list.sort(key=(lambda criteria: len(criteria.prefix)), reverse=True)
            for criteria in criteria_list:
                if national_number.startswith(criteria.prefix):
                    return criteria
            raise ObjectDoesNotExist

        try:
            return get_criteria_with_longest_matching_prefix(
                list(all_possible_criteria.filter(country_code=country_code, backend_instance=backend_instance))
            )
        except ObjectDoesNotExist:
            pass
        try:
            return all_possible_criteria.get(country_code=None, backend_instance=backend_instance)
        except ObjectDoesNotExist:
            pass
        try:
            return get_criteria_with_longest_matching_prefix(
                list(all_possible_criteria.filter(country_code=country_code, backend_instance=None))
            )
        except ObjectDoesNotExist:
            pass
        try:
            return all_possible_criteria.get(country_code=None, backend_instance=None)
        except ObjectDoesNotExist:
            pass

        return None

class SmsGatewayFee(models.Model):
    """
    The fee for sending or receiving an SMS Message based on gateway.
    When an SmsBillable is calculated, it will use the most recent SmsFee available from the criteria
    to determine the gateway_charge.

    Once an SmsFee is created, it cannot be modified.
    """
    criteria = models.ForeignKey(SmsGatewayFeeCriteria, on_delete=models.PROTECT)
    amount = models.DecimalField(default=0.0, max_digits=10, decimal_places=4)
    currency = models.ForeignKey(accounting.Currency, on_delete=models.PROTECT)
    date_created = models.DateTimeField(auto_now_add=True)

    class Meta:
        app_label = 'smsbillables'

    @classmethod
    def create_new(cls, backend_api_id, direction, amount,
                   currency=None, backend_instance=None, country_code=None, prefix=None,
                   save=True, fee_class=None, criteria_class=None):
        fee_class = fee_class or cls
        criteria_class = criteria_class or SmsGatewayFeeCriteria
        currency = currency or Currency.get_default()

        if 'prefix' in [
            field.name
            for field, _ in criteria_class._meta.get_fields_with_model()
        ]:
            prefix = prefix or ''
            criteria, _ = criteria_class.objects.get_or_create(
                backend_api_id=backend_api_id,
                direction=direction,
                backend_instance=backend_instance,
                country_code=country_code,
                prefix=prefix,
            )
        else:
            criteria, _ = criteria_class.objects.get_or_create(
                backend_api_id=backend_api_id,
                direction=direction,
                backend_instance=backend_instance,
                country_code=country_code,
            )
        new_fee = fee_class(
            currency=currency,
            amount=amount,
            criteria=criteria
        )
        if save:
            new_fee.save()
        return new_fee

    @classmethod
    def get_by_criteria(cls, backend_api_id, direction,
                        backend_instance=None, country_code=None, national_number=None):
        criteria = SmsGatewayFeeCriteria.get_most_specific(
            backend_api_id,
            direction,
            backend_instance=backend_instance,
            country_code=country_code,
            national_number=national_number,
        )
        return cls.get_by_criteria_obj(criteria)

    @classmethod
    def get_by_criteria_obj(cls, criteria):
        if not criteria:
            return None
        return cls.objects.filter(criteria=criteria.id).latest('date_created')

class SmsUsageFeeCriteria(models.Model):
    """
    Criteria for determining a usage fee applied for each SMS message sent or received.

    Nullable fields indicate criteria that can be applied globally to all messages with no specific matches
    for that field.
    """
    direction = models.CharField(max_length=10, db_index=True, choices=DIRECTION_CHOICES)
    domain = models.CharField(max_length=25, db_index=True, null=True)

    class Meta:
        app_label = 'smsbillables'

    @classmethod
    def get_most_specific(cls, direction, domain=None):
        """
        Gets the most specific criteria available based on (and in order of preference for optional):
        - direction
        - domain (optional)
        """
        all_possible_criteria = cls.objects.filter(direction=direction)

        if all_possible_criteria.count() == 0:
            return None

        try:
            return all_possible_criteria.get(domain=domain)
        except ObjectDoesNotExist:
            pass
        try:
            return all_possible_criteria.get(domain=None)
        except ObjectDoesNotExist:
            pass

        return None

class SmsUsageFee(models.Model):
    """
    The usage fee, with version information, based on domain or globally.
    When an SmsBillable is calculated, it will use the most recent SmsUsageFee available from the
    criteria to determine the usage_charge.

    Currency is always in USD since this is something we control.

    Once an SmsUsageFee is created, it cannot be modified.
    """
    criteria = models.ForeignKey(SmsUsageFeeCriteria, on_delete=models.PROTECT)
    amount = models.DecimalField(default=0.0, max_digits=10, decimal_places=4)
    date_created = models.DateTimeField(auto_now_add=True)

    class Meta:
        app_label = 'smsbillables'

    @classmethod
    def create_new(cls, direction, amount, domain=None, save=True):
        criteria, _ = SmsUsageFeeCriteria.objects.get_or_create(
            domain=domain, direction=direction,
        )
        new_fee = SmsUsageFee(
            amount=amount,
            criteria=criteria
        )
        if save:
            new_fee.save()
        return new_fee

    @classmethod
    def get_by_criteria(cls, direction, domain=None):
        criteria = SmsUsageFeeCriteria.get_most_specific(direction, domain=domain)
        if not criteria:
            return None
        return cls.objects.filter(criteria=criteria.id).latest('date_created')

@quickcache(['sms_backend_id'])
def _sms_backend_is_global(sms_backend_id):
    return SQLMobileBackend.load(
        sms_backend_id,
        is_couch_id=True,
        include_deleted=True
    ).is_global

class SmsBillable(models.Model):
    """
    A record of matching a fee to a particular SMS.

    If on closer inspection we determine a particular SmsBillable is invalid
    (we used the incorrect fee and want to recalculate)
    we can set this billable to is_valid = False and it will not be used toward
    calculating the SmsLineItem in the monthly Invoice.
    """
    gateway_fee = models.ForeignKey(SmsGatewayFee, null=True, on_delete=models.PROTECT)
    gateway_fee_conversion_rate = models.DecimalField(default=Decimal('1.0'), null=True, max_digits=20,
                                                      decimal_places=EXCHANGE_RATE_DECIMAL_PLACES)
    usage_fee = models.ForeignKey(SmsUsageFee, null=True, on_delete=models.PROTECT)
    log_id = models.CharField(max_length=50, db_index=True)
    phone_number = models.CharField(max_length=50)
    api_response = models.TextField(null=True, blank=True)
    is_valid = models.BooleanField(default=True, db_index=True)
    domain = models.CharField(max_length=25, db_index=True)
    direction = models.CharField(max_length=10, db_index=True, choices=DIRECTION_CHOICES)
    date_sent = models.DateField()
    date_created = models.DateField(auto_now_add=True)

    class Meta:
        app_label = 'smsbillables'

    @property
    def gateway_charge(self):
        if self.gateway_fee is not None:
            try:
                charge = SmsGatewayFee.objects.get(id=self.gateway_fee.id)
                if self.gateway_fee_conversion_rate is not None:
                    return charge.amount / self.gateway_fee_conversion_rate
                return charge.amount
            except ObjectDoesNotExist:
                pass
        return Decimal('0.0')

    @property
    def usage_charge(self):
        if self.usage_fee is not None:
            try:
                charge = SmsUsageFee.objects.get(id=self.usage_fee.id)
                return charge.amount
            except ObjectDoesNotExist:
                pass
        return Decimal('0.0')

    @classmethod
    def create(cls, message_log):
        phone_number = clean_phone_number(message_log.phone_number)
        direction = message_log.direction
        domain = message_log.domain
        log_id = message_log.couch_id

        billable = cls(
            log_id=log_id,
            phone_number=phone_number,
            direction=direction,
            date_sent=message_log.date,
            domain=domain,
        )
        billable.gateway_fee, billable.gateway_fee_conversion_rate = cls._get_gateway_fee(
            message_log.backend_api, message_log.backend_id, phone_number, direction, log_id
        )
        billable.usage_fee = cls._get_usage_fee(domain, direction)

        if message_log.backend_api == SQLTestSMSBackend.get_api_id():
            billable.is_valid = False

        billable.save()
        return billable

    @classmethod
    def _get_gateway_fee(cls, backend_api_id, backend_instance, phone_number, direction, couch_id):
        country_code, national_number = get_country_code_and_national_number(phone_number)
        is_gateway_billable = backend_instance is None or _sms_backend_is_global(backend_instance)

        if is_gateway_billable:
            gateway_fee = SmsGatewayFee.get_by_criteria(
                backend_api_id,
                direction,
                backend_instance=backend_instance,
                country_code=country_code,
                national_number=national_number,
            )
            if gateway_fee:
                conversion_rate = gateway_fee.currency.rate_to_default
                if conversion_rate != 0:
                    return gateway_fee, conversion_rate
                else:
                    log_smsbillables_error(
                        "Gateway fee conversion rate for currency %s is 0"
                        % gateway_fee.currency.code
                    )
                    return gateway_fee, None
            else:
                log_smsbillables_error(
                    "No matching gateway fee criteria for SMS %s" % couch_id
                )
        return None, None

    @classmethod
    def _get_usage_fee(cls, domain, direction):
        usage_fee = SmsUsageFee.get_by_criteria(
            direction, domain=domain
        )
        if not usage_fee:
            log_smsbillables_error(
                "Did not find usage fee for direction %s and domain %s"
                % (direction, domain)
            )
        return usage_fee
<|endoftext|>


<commit_before>from django.core.exceptions import ObjectDoesNotExist
from django.db import models

from corehq.apps.accounting import models as accounting
from corehq.apps.sms.models import DIRECTION_CHOICES

class SmsGatewayFeeCriteria(models.Model):
    """
    These are the parameters we'll use to try and calculate the cost of sending a message through
    our gateways. We configure the SMS fee criteria based on parameters given to us by specific
    gateway providers.

    Nullable fields indicate criteria that can be applied globally to all messages with no specific matches
    for that field.
    """
    backend_api_id = models.CharField(max_length=100, db_index=True)
    backend_instance = models.CharField(max_length=255, db_index=True, null=True)
    direction = models.CharField(max_length=10, db_index=True, choices=DIRECTION_CHOICES)
    country_code = models.IntegerField(max_length=5, null=True, blank=True, db_index=True)

class SmsGatewayFee(models.Model):
    """
    The fee for sending or receiving an SMS Message based on gateway.
    When an SmsBillable is calculated, it will use the most recent SmsFee available from the criteria
    to determine the gateway_charge.

    Once an SmsFee is created, it cannot be modified.
    """
    criteria = models.ForeignKey(SmsGatewayFeeCriteria, on_delete=models.PROTECT)
    amount = models.DecimalField(default=0.0, max_digits=7, decimal_places=4)
    currency = models.ForeignKey(accounting.Currency, on_delete=models.PROTECT)
    date_created = models.DateField(auto_now_add=True)

class SmsUsageFeeCriteria(models.Model):
    """
    Criteria for determining a usage fee applied for each SMS message sent or received.

    Nullable fields indicate criteria that can be applied globally to all messages with no specific matches
    for that field.
    """
    direction = models.CharField(max_length=10, db_index=True, choices=DIRECTION_CHOICES)
    domain = models.CharField(max_length=25, db_index=True, null=True)

class SmsUsageFee(models.Model):
    """
    The usage fee, with version information, based on domain or globally.
    When an SmsBillable is calculated, it will use the most recent SmsUsageFee available from the
    criteria to determine the usage_charge.

    Currency is always in USD since this is something we control.

    Once an SmsUsageFee is created, it cannot be modified.
    """
    criteria = models.ForeignKey(SmsUsageFeeCriteria, on_delete=models.PROTECT)
    amount = models.DecimalField(default=0.0, max_digits=7, decimal_places=4)
    date_created = models.DateField(auto_now_add=True)

class SmsBillable(models.Model):
    """
    A record of matching a fee to a particular MessageLog (or SMSLog).

    If on closer inspection we determine a particular SmsBillable is invalid (whether something is
    awry with the api_response, or we used the incorrect fee and want to recalculate) we can set
    this billable to is_valid = False and it will not be used toward calculating the SmsLineItem in
    the monthly Invoice.
    """
    gateway_fee = models.ForeignKey(SmsGatewayFee, null=True, on_delete=models.PROTECT)
    usage_fee = models.ForeignKey(SmsUsageFee, null=True, on_delete=models.PROTECT)
    log_id = models.CharField(max_length=50)
    phone_number = models.CharField(max_length=50)
    api_response = models.TextField(blank=True)
    is_valid = models.BooleanField(default=True, db_index=True)
    domain = models.CharField(max_length=25, db_index=True)
    direction = models.CharField(max_length=10, db_index=True, choices=DIRECTION_CHOICES)
    date_sent = models.DateField()
    date_created = models.DateField(auto_now_add=True)

    @property
    def gateway_charge(self):
        try:
            charge = SmsGatewayFee.objects.get(id=self.gateway_fee.id)
            return charge.amount
        except ObjectDoesNotExist:
            return 0.0

    @property
    def usage_charge(self):
        try:
            charge = SmsUsageFee.objects.get(id=self.usage_fee.id)
            return charge.amount
        except ObjectDoesNotExist:
            return 0.0
<commit_msg>api_response can be null as well, “” should be different from None. perhaps the API sends nothing unless there’s an error<commit_after>from django.core.exceptions import ObjectDoesNotExist
from django.db import models

from corehq.apps.accounting import models as accounting
from corehq.apps.sms.models import DIRECTION_CHOICES

class SmsGatewayFeeCriteria(models.Model):
    """
    These are the parameters we'll use to try and calculate the cost of sending a message through
    our gateways. We configure the SMS fee criteria based on parameters given to us by specific
    gateway providers.

    Nullable fields indicate criteria that can be applied globally to all messages with no specific matches
    for that field.
    """
    backend_api_id = models.CharField(max_length=100, db_index=True)
    backend_instance = models.CharField(max_length=255, db_index=True, null=True)
    direction = models.CharField(max_length=10, db_index=True, choices=DIRECTION_CHOICES)
    country_code = models.IntegerField(max_length=5, null=True, blank=True, db_index=True)

class SmsGatewayFee(models.Model):
    """
    The fee for sending or receiving an SMS Message based on gateway.
    When an SmsBillable is calculated, it will use the most recent SmsFee available from the criteria
    to determine the gateway_charge.

    Once an SmsFee is created, it cannot be modified.
    """
    criteria = models.ForeignKey(SmsGatewayFeeCriteria, on_delete=models.PROTECT)
    amount = models.DecimalField(default=0.0, max_digits=7, decimal_places=4)
    currency = models.ForeignKey(accounting.Currency, on_delete=models.PROTECT)
    date_created = models.DateField(auto_now_add=True)

class SmsUsageFeeCriteria(models.Model):
    """
    Criteria for determining a usage fee applied for each SMS message sent or received.

    Nullable fields indicate criteria that can be applied globally to all messages with no specific matches
    for that field.
    """
    direction = models.CharField(max_length=10, db_index=True, choices=DIRECTION_CHOICES)
    domain = models.CharField(max_length=25, db_index=True, null=True)

class SmsUsageFee(models.Model):
    """
    The usage fee, with version information, based on domain or globally.
    When an SmsBillable is calculated, it will use the most recent SmsUsageFee available from the
    criteria to determine the usage_charge.

    Currency is always in USD since this is something we control.

    Once an SmsUsageFee is created, it cannot be modified.
    """
    criteria = models.ForeignKey(SmsUsageFeeCriteria, on_delete=models.PROTECT)
    amount = models.DecimalField(default=0.0, max_digits=7, decimal_places=4)
    date_created = models.DateField(auto_now_add=True)

class SmsBillable(models.Model):
    """
    A record of matching a fee to a particular MessageLog (or SMSLog).

    If on closer inspection we determine a particular SmsBillable is invalid (whether something is
    awry with the api_response, or we used the incorrect fee and want to recalculate) we can set
    this billable to is_valid = False and it will not be used toward calculating the SmsLineItem in
    the monthly Invoice.
    """
    gateway_fee = models.ForeignKey(SmsGatewayFee, null=True, on_delete=models.PROTECT)
    usage_fee = models.ForeignKey(SmsUsageFee, null=True, on_delete=models.PROTECT)
    log_id = models.CharField(max_length=50)
    phone_number = models.CharField(max_length=50)
    api_response = models.TextField(null=True)
    is_valid = models.BooleanField(default=True, db_index=True)
    domain = models.CharField(max_length=25, db_index=True)
    direction = models.CharField(max_length=10, db_index=True, choices=DIRECTION_CHOICES)
    date_sent = models.DateField()
    date_created = models.DateField(auto_now_add=True)

    @property
    def gateway_charge(self):
        try:
            charge = SmsGatewayFee.objects.get(id=self.gateway_fee.id)
            return charge.amount
        except ObjectDoesNotExist:
            return 0.0

    @property
    def usage_charge(self):
        try:
            charge = SmsUsageFee.objects.get(id=self.usage_fee.id)
            return charge.amount
        except ObjectDoesNotExist:
            return 0.0
<|endoftext|>


<commit_before>from collections import defaultdict
from sqlalchemy import Table, MetaData, Column, types, func, sql
from corehq.db import Session
from .models import FormData

# A SQLAlchemy representation of the FormData table.
# This skips the SQLAlchemy ORM and uses their SQL Expression Language
# http://docs.sqlalchemy.org/en/rel_0_8/core/tutorial.html#selecting
SAFormData = Table(
    FormData._meta.db_table,
    MetaData(),
    Column('domain', types.String),
    Column('received_on', types.DateTime),
    Column('time_end', types.DateTime),
    Column('xmlns', types.String),
    Column('user_id', types.String),
    Column('app_id', types.String),
)

def get_form_counts_by_user_xmlns(domain, startdate, enddate, user_ids=None,
                                  xmlnss=None, by_submission_time=True):
    # This kind of COUNT and GROUP BY query isn't possible with the Django ORM
    col = SAFormData.columns
    date_field = col.received_on if by_submission_time else col.time_end
    query = (sql.select([func.count(), col.xmlns, col.user_id, col.app_id])
             .where(col.domain == domain)
             .where(startdate <= date_field)
             .where(date_field < enddate)
             .group_by(col.xmlns, col.user_id, col.app_id))
    if user_ids:
        query = query.where(col.user_id.in_(user_ids))
    if xmlnss:
        query = query.where(col.xmlns.in_(xmlnss))
    return defaultdict(lambda: 0, {
        (user_id, xmlns, app_id): count
        for count, xmlns, user_id, app_id in Session.execute(query)
    })
<commit_msg>*sigh*, use the Django ORM<commit_after>from collections import defaultdict
from django.db.models import Count
from .models import FormData

def get_form_counts_by_user_xmlns(domain, startdate, enddate, user_ids=None,
                                  xmlnss=None, by_submission_time=True):
    date_field = 'received_on' if by_submission_time else 'time_end'
    query = (FormData.objects
             .filter(domain=domain,
                     **{'{}__gte'.format(date_field): startdate,
                        '{}__lt'.format(date_field): enddate})
             .values('xmlns', 'user_id', 'app_id')
             .annotate(count=Count('pk')))

    if user_ids:
        query = query.filter(user_id__in=user_ids)
    if xmlnss:
        query = query.filter(xmlns__in=xmlnss)

    return defaultdict(lambda: 0, {
        (r['user_id'], r['xmlns'], r['app_id']): r['count'] for r in query
    })
<|endoftext|>


<commit_before>import hashlib
from abc import abstractmethod, ABCMeta, abstractproperty
from datetime import datetime

import six
from couchdbkit import ResourceNotFound

from corehq.util.couch_helpers import PaginateViewLogHandler, paginate_view

class ResumableDocsByTypeIterator(object):
    """Perform one-time resumable iteration over documents by type

    Iteration can be efficiently stopped and resumed. The iteration may
    omit documents that are added after the iteration begins or resumes
    and may include deleted documents.

    :param db: Couchdb database.
    :param doc_types: A list of doc type names to iterate on.
    :param iteration_key: A unique key identifying the iteration. This
    key will be used in combination with `doc_types` to maintain state
    about an iteration that is in progress. The state will be maintained
    indefinitely unless it is removed with `discard_state()`.
    :param chunk_size: Number of documents to yield before updating the
    iteration checkpoint. In the worst case about this many documents
    that were previously yielded may be yielded again if the iteration
    is stopped and later resumed.
    """

    def __init__(self, db, doc_types, iteration_key, chunk_size=100, log_handler=None):
        if isinstance(doc_types, str):
            raise TypeError("expected list of strings, got %r" % (doc_types,))
        self.db = db
        self.original_doc_types = doc_types = sorted(doc_types)
        self.iteration_key = iteration_key
        self.chunk_size = chunk_size
        iteration_name = "{}/{}".format(iteration_key, " ".join(doc_types))
        self.iteration_id = hashlib.sha1(iteration_name).hexdigest()
        try:
            self.state = db.get(self.iteration_id)
        except ResourceNotFound:
            # new iteration
            self.state = {
                "_id": self.iteration_id,
                "doc_type": "ResumableDocsByTypeIteratorState",
                "retry": {},

                # for humans
                "name": iteration_name,
                "timestamp": datetime.utcnow().isoformat()
            }
            args = {}
        else:
            # resume iteration
            args = self.state.get("offset", {}).copy()
            if args:
                assert args.get("startkey"), args
                doc_type = args["startkey"][0]
                # skip doc types before offset
                doc_types = doc_types[doc_types.index(doc_type):]
            else:
                # non-retry phase of iteration is complete
                doc_types = []
        args.update(
            view_name='all_docs/by_doc_type',
            chunk_size=chunk_size,
            log_handler=ResumableDocsByTypeLogHandler(self, child=log_handler),
            include_docs=True,
            reduce=False,
        )
        self.view_args = args
        self.doc_types = doc_types

    def __iter__(self):
        args = self.view_args
        for doc_type in self.doc_types:
            if args.get("startkey", [None])[0] != doc_type:
                args.pop("startkey_docid", None)
                args["startkey"] = [doc_type]
            args["endkey"] = [doc_type, {}]
            for result in paginate_view(self.db, **args):
                yield result['doc']

        retried = {}
        while self.state["retry"] != retried:
            for doc_id, retries in list(self.state["retry"].iteritems()):
                if retries == retried.get(doc_id):
                    continue  # skip already retried (successfully)
                retried[doc_id] = retries
                try:
                    yield self.db.get(doc_id)
                except ResourceNotFound:
                    pass

        # save iteration state without offset to signal completion
        self.state.pop("offset", None)
        self.state["retry"] = {}
        self._save_state()

    def retry(self, doc, max_retry=3):
        """Add document to be yielded at end of iteration

        Iteration order of retry documents is undefined. All retry
        documents will be yielded after the initial non-retry phase of
        iteration has completed, and every retry document will be
        yielded each time the iterator is stopped and resumed during the
        retry phase. This method is relatively inefficient because it
        forces the iteration state to be saved to couch. If you find
        yourself calling this for many documents during the iteration
        you may want to consider a different retry strategy.

        :param doc: The doc dict to retry. It will be re-fetched from
        the database before being yielded from the iteration.
        :param max_retry: Maximum number of times a given document may
        be retried.
        :raises: `TooManyRetries` if this method has been called too
        many times with a given document.
        """
        doc_id = doc["_id"]
        retries = self.state["retry"].get(doc_id, 0) + 1
        if retries > max_retry:
            raise TooManyRetries(doc_id)
        self.state["retry"][doc_id] = retries
        self._save_state()

    @property
    def progress_info(self):
        """Extra progress information

        This property can be used to store and retrieve extra progress
        information associated with the iteration. The information is
        persisted with the iteration state in couch.
        """
        return self.state.get("progress_info")

    @progress_info.setter
    def progress_info(self, info):
        self.state["progress_info"] = info
        self._save_state()

    def _save_state(self):
        self.state["timestamp"] = datetime.utcnow().isoformat()
        self.db.save_doc(self.state)

    def discard_state(self):
        try:
            self.db.delete_doc(self.iteration_id)
        except ResourceNotFound:
            pass
        self.__init__(
            self.db,
            self.original_doc_types,
            self.iteration_key,
            self.chunk_size,
        )

class ResumableDocsByTypeLogHandler(PaginateViewLogHandler):

    def __init__(self, iterator, child=None):
        self.iterator = iterator
        self.child = child

    def view_starting(self, db, view_name, kwargs, total_emitted):
        if self.child:
            self.child.view_starting(db, view_name, kwargs, total_emitted)

        offset = {k: v for k, v in kwargs.items() if k.startswith("startkey")}
        self.iterator.state["offset"] = offset
        self.iterator._save_state()

    def view_ending(self, db, view_name, kwargs, total_emitted, time):
        if self.child:
            self.child.view_ending(db, view_name, kwargs, total_emitted, time)

class TooManyRetries(Exception):
    pass

class BulkProcessingFailed(Exception):
    pass

DOCS_SKIPPED_WARNING = """
        WARNING {} documents were not processed due to concurrent modification
        during migration. Run the migration again until you do not see this
        message.
        """

class BaseDocProcessor(six.with_metaclass(ABCMeta)):

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

    def __enter__(self):
        pass

    def __exit__(self, exc_type, exc_val, exc_tb):
        pass

    @abstractproperty
    def unique_key(self):
        return self.slug

    def process_doc(self, doc, couchdb):
        """Process a single document

        :param doc: The document dict to be processed.
        :param couchdb: Couchdb database associated with the docs.
        :returns: True if doc was processed successfully else False. If this returns False
        the document migration will be retried later.
        """
        raise NotImplementedError

    def process_bulk_docs(self, docs, couchdb):
        """Process a batch of documents. The default implementation passes
        each doc in turn to ``process_doc``.

        :param docs: A list of document dicts to be processed.
        :param couchdb: Couchdb database associated with the docs.
        :returns: True if doc was processed successfully else False.
        If this returns False the processing will be halted.
        """
        return all(self.process_doc(doc, couchdb) for doc in docs)

    def handle_skip(self, doc):
        """Called when a document is going to be skipped i.e. it has been
        retried > max_retries.

        :returns: True to indicate that the skip has been handled
                  or False to stop execution
        """
        return False

    def processing_complete(self, skipped):
        pass

    def should_process(self, doc):
        """
        :param doc: the document to filter
        :return: True if this doc should be migrated
        """
        return True

class CouchDocumentProcessor(object):
    """Process Couch Docs

    :param doc_type_map: Dict of `doc_type_name: model_class` pairs.
    :param doc_processor: A `BaseDocProcessor` object used to
    process documents.
    :param reset: Reset existing processor state (if any), causing all
    documents to be reconsidered for processing, if this is true.
    :param max_retry: Number of times to retry processing a document
    before giving up.
    :param chunk_size: Maximum number of records to read from couch at
    one time. It may be necessary to use a smaller chunk size if the
    records being processed are very large and the default chunk size of
    100 would exceed available memory.
    :param log_handler: instance of PaginateViewLogHandler to be notified of
    view events.
    """
    def __init__(self, doc_type_map, doc_processor, reset=False, max_retry=2, chunk_size=100, log_handler=None):
        self.doc_type_map = doc_type_map
        self.doc_processor = doc_processor
        self.reset = reset
        self.max_retry = max_retry
        self.chunk_size = chunk_size

        self.couchdb = next(iter(doc_type_map.values())).get_db()
        assert all(m.get_db() is self.couchdb for m in doc_type_map.values()), \
            "documents must live in same couch db: %s" % repr(doc_type_map)

        self.docs_by_type = ResumableDocsByTypeIterator(
            self.couchdb, doc_type_map, doc_processor.unique_key, chunk_size=chunk_size,
            log_handler=log_handler
        )

        self.visited = 0
        self.previously_visited = 0
        self.total = 0

        self.processed = 0
        self.skipped = 0

        self.start = None

    def has_started(self):
        return bool(self.docs_by_type.progress_info)

    @property
    def session_visited(self):
        return self.visited - self.previously_visited

    @property
    def session_total(self):
        return self.total - self.previously_visited

    @property
    def attempted(self):
        return self.processed + self.skipped

    @property
    def timing(self):
        """Returns a tuple of (elapsed, remaining)"""
        elapsed = datetime.now() - self.start
        if self.session_visited > self.session_total:
            remaining = "?"
        else:
            session_remaining = self.session_total - self.session_visited
            remaining = elapsed / self.session_visited * session_remaining
        return elapsed, remaining

    def _setup(self):
        from corehq.dbaccessors.couchapps.all_docs import get_doc_count_by_type

        self.total = sum(get_doc_count_by_type(self.couchdb, doc_type)
                    for doc_type in self.doc_type_map)

        if self.reset:
            self.docs_by_type.discard_state()
        elif self.docs_by_type.progress_info:
            info = self.docs_by_type.progress_info
            old_total = info["total"]
            # Estimate already visited based on difference of old/new
            # totals. The theory is that new or deleted records will be
            # evenly distributed across the entire set.
            self.visited = int(round(float(self.total) / old_total * info["visited"]))
            self.previously_visited = self.visited
        print("Processing {} documents{}: {}...".format(
            self.total,
            " (~{} already processed)".format(self.visited) if self.visited else "",
            ", ".join(sorted(self.doc_type_map))
        ))

        self.start = datetime.now()

    def run(self):
        """
        :returns: A tuple `(<num processed>, <num skipped>)`
        """
        self._setup()
        with self.doc_processor:
            for doc in self.docs_by_type:
                self._process_doc(doc)
                self._update_progress()

        self._processing_complete()

        return self.processed, self.skipped

    def _process_doc(self, doc):
        if not self.doc_processor.should_process(doc):
            return

        ok = self.doc_processor.process_doc(doc, self.couchdb)
        if ok:
            self.processed += 1
        else:
            try:
                self.docs_by_type.retry(doc, self.max_retry)
            except TooManyRetries:
                if not self.doc_processor.handle_skip(doc):
                    raise
                else:
                    print("Skip: {doc_type} {_id}".format(**doc))
                    self.skipped += 1

    def _update_progress(self):
        self.visited += 1
        if self.visited % self.chunk_size == 0:
            self.docs_by_type.progress_info = {"visited": self.visited, "total": self.total}

        if self.attempted % self.chunk_size == 0:
            elapsed, remaining = self.timing
            print("Processed {}/{} of {} documents in {} ({} remaining)"
                  .format(self.processed, self.visited, self.total, elapsed, remaining))

    def _processing_complete(self):
        if self.session_visited:
            self.docs_by_type.progress_info = {"visited": self.visited, "total": self.total}
        self.doc_processor.processing_complete(self.skipped)
        print("Processed {}/{} of {} documents ({} previously processed, {} filtered out).".format(
            self.processed,
            self.visited,
            self.total,
            self.previously_visited,
            self.session_visited - self.attempted
        ))
        if self.skipped:
            print(DOCS_SKIPPED_WARNING.format(self.skipped))

class BulkDocProcessorLogHandler(PaginateViewLogHandler):

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

    def view_starting(self, db, view_name, kwargs, total_emitted):
        self.processor.process_chunk()

class BulkDocProcessor(CouchDocumentProcessor):
    """Process couch docs for a single doc type in batches

    The bulk doc processor will send a batch of documents to the document
    processor. If the processor does not respond with True then
    the iteration is halted. Restarting the iteration will start by
    re-sending the previous chunk to the processor.

    The size of the batches passed to the document processor may vary
    depending on how they are being filtered by the
    document processor but will never exceed ``chunk_size``.

    :param doc_type_map: Dict containing a single `doc_type_name: model_class` pair.
    :param doc_processor: A `BaseDocProcessor` object used to
    process documents.
    :param reset: Reset existing processor state (if any), causing all
    documents to be reconsidered for processing, if this is true.
    :param chunk_size: Maximum number of records to read from couch at
    one time. It may be necessary to use a smaller chunk size if the
    records being processed are very large and the default chunk size of
    100 would exceed available memory.
    """
    def __init__(self, doc_type_map, doc_processor, reset=False, chunk_size=100):
        assert len(doc_type_map) == 1
        super(BulkDocProcessor, self).__init__(
            doc_type_map, doc_processor, reset=reset, chunk_size=chunk_size,
            log_handler=BulkDocProcessorLogHandler(self)
        )
        self.changes = []

    def _process_doc(self, doc):
        if self.doc_processor.should_process(doc):
            self.changes.append(doc)

    def process_chunk(self):
        """Called by the BulkDocProcessorLogHandler"""
        ok = self.doc_processor.process_bulk_docs(self.changes, self.couchdb)
        if ok:
            self.processed += len(self.changes)
            self.changes = []
        else:
            raise BulkProcessingFailed("Processing batch failed")

    def _update_progress(self):
        self.visited += 1
        if self.visited % self.chunk_size == 0:
            self.docs_by_type.progress_info = {"visited": self.visited, "total": self.total}

            elapsed, remaining = self.timing
            print("Processed {}/{} of {} documents in {} ({} remaining)"
                  .format(self.processed, self.visited, self.total, elapsed, remaining))

    def _processing_complete(self):
        if len(self.changes):
            self._process_chunk()

        super(BulkDocProcessor, self)._processing_complete()
<commit_msg>move process chunk to view_ending<commit_after>import hashlib
from abc import abstractmethod, ABCMeta, abstractproperty
from datetime import datetime

import six
from couchdbkit import ResourceNotFound

from corehq.util.couch_helpers import PaginateViewLogHandler, paginate_view

class ResumableDocsByTypeIterator(object):
    """Perform one-time resumable iteration over documents by type

    Iteration can be efficiently stopped and resumed. The iteration may
    omit documents that are added after the iteration begins or resumes
    and may include deleted documents.

    :param db: Couchdb database.
    :param doc_types: A list of doc type names to iterate on.
    :param iteration_key: A unique key identifying the iteration. This
    key will be used in combination with `doc_types` to maintain state
    about an iteration that is in progress. The state will be maintained
    indefinitely unless it is removed with `discard_state()`.
    :param chunk_size: Number of documents to yield before updating the
    iteration checkpoint. In the worst case about this many documents
    that were previously yielded may be yielded again if the iteration
    is stopped and later resumed.
    """

    def __init__(self, db, doc_types, iteration_key, chunk_size=100, log_handler=None):
        if isinstance(doc_types, str):
            raise TypeError("expected list of strings, got %r" % (doc_types,))
        self.db = db
        self.original_doc_types = doc_types = sorted(doc_types)
        self.iteration_key = iteration_key
        self.chunk_size = chunk_size
        iteration_name = "{}/{}".format(iteration_key, " ".join(doc_types))
        self.iteration_id = hashlib.sha1(iteration_name).hexdigest()
        try:
            self.state = db.get(self.iteration_id)
        except ResourceNotFound:
            # new iteration
            self.state = {
                "_id": self.iteration_id,
                "doc_type": "ResumableDocsByTypeIteratorState",
                "retry": {},

                # for humans
                "name": iteration_name,
                "timestamp": datetime.utcnow().isoformat()
            }
            args = {}
        else:
            # resume iteration
            args = self.state.get("offset", {}).copy()
            if args:
                assert args.get("startkey"), args
                doc_type = args["startkey"][0]
                # skip doc types before offset
                doc_types = doc_types[doc_types.index(doc_type):]
            else:
                # non-retry phase of iteration is complete
                doc_types = []
        args.update(
            view_name='all_docs/by_doc_type',
            chunk_size=chunk_size,
            log_handler=ResumableDocsByTypeLogHandler(self, child=log_handler),
            include_docs=True,
            reduce=False,
        )
        self.view_args = args
        self.doc_types = doc_types

    def __iter__(self):
        args = self.view_args
        for doc_type in self.doc_types:
            if args.get("startkey", [None])[0] != doc_type:
                args.pop("startkey_docid", None)
                args["startkey"] = [doc_type]
            args["endkey"] = [doc_type, {}]
            for result in paginate_view(self.db, **args):
                yield result['doc']

        retried = {}
        while self.state["retry"] != retried:
            for doc_id, retries in list(self.state["retry"].iteritems()):
                if retries == retried.get(doc_id):
                    continue  # skip already retried (successfully)
                retried[doc_id] = retries
                try:
                    yield self.db.get(doc_id)
                except ResourceNotFound:
                    pass

        # save iteration state without offset to signal completion
        self.state.pop("offset", None)
        self.state["retry"] = {}
        self._save_state()

    def retry(self, doc, max_retry=3):
        """Add document to be yielded at end of iteration

        Iteration order of retry documents is undefined. All retry
        documents will be yielded after the initial non-retry phase of
        iteration has completed, and every retry document will be
        yielded each time the iterator is stopped and resumed during the
        retry phase. This method is relatively inefficient because it
        forces the iteration state to be saved to couch. If you find
        yourself calling this for many documents during the iteration
        you may want to consider a different retry strategy.

        :param doc: The doc dict to retry. It will be re-fetched from
        the database before being yielded from the iteration.
        :param max_retry: Maximum number of times a given document may
        be retried.
        :raises: `TooManyRetries` if this method has been called too
        many times with a given document.
        """
        doc_id = doc["_id"]
        retries = self.state["retry"].get(doc_id, 0) + 1
        if retries > max_retry:
            raise TooManyRetries(doc_id)
        self.state["retry"][doc_id] = retries
        self._save_state()

    @property
    def progress_info(self):
        """Extra progress information

        This property can be used to store and retrieve extra progress
        information associated with the iteration. The information is
        persisted with the iteration state in couch.
        """
        return self.state.get("progress_info")

    @progress_info.setter
    def progress_info(self, info):
        self.state["progress_info"] = info
        self._save_state()

    def _save_state(self):
        self.state["timestamp"] = datetime.utcnow().isoformat()
        self.db.save_doc(self.state)

    def discard_state(self):
        try:
            self.db.delete_doc(self.iteration_id)
        except ResourceNotFound:
            pass
        self.__init__(
            self.db,
            self.original_doc_types,
            self.iteration_key,
            self.chunk_size,
        )

class ResumableDocsByTypeLogHandler(PaginateViewLogHandler):

    def __init__(self, iterator, child=None):
        self.iterator = iterator
        self.child = child

    def view_starting(self, db, view_name, kwargs, total_emitted):
        if self.child:
            self.child.view_starting(db, view_name, kwargs, total_emitted)

        offset = {k: v for k, v in kwargs.items() if k.startswith("startkey")}
        self.iterator.state["offset"] = offset
        self.iterator._save_state()

    def view_ending(self, db, view_name, kwargs, total_emitted, time):
        if self.child:
            self.child.view_ending(db, view_name, kwargs, total_emitted, time)

class TooManyRetries(Exception):
    pass

class BulkProcessingFailed(Exception):
    pass

DOCS_SKIPPED_WARNING = """
        WARNING {} documents were not processed due to concurrent modification
        during migration. Run the migration again until you do not see this
        message.
        """

class BaseDocProcessor(six.with_metaclass(ABCMeta)):

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

    def __enter__(self):
        pass

    def __exit__(self, exc_type, exc_val, exc_tb):
        pass

    @abstractproperty
    def unique_key(self):
        return self.slug

    def process_doc(self, doc, couchdb):
        """Process a single document

        :param doc: The document dict to be processed.
        :param couchdb: Couchdb database associated with the docs.
        :returns: True if doc was processed successfully else False. If this returns False
        the document migration will be retried later.
        """
        raise NotImplementedError

    def process_bulk_docs(self, docs, couchdb):
        """Process a batch of documents. The default implementation passes
        each doc in turn to ``process_doc``.

        :param docs: A list of document dicts to be processed.
        :param couchdb: Couchdb database associated with the docs.
        :returns: True if doc was processed successfully else False.
        If this returns False the processing will be halted.
        """
        return all(self.process_doc(doc, couchdb) for doc in docs)

    def handle_skip(self, doc):
        """Called when a document is going to be skipped i.e. it has been
        retried > max_retries.

        :returns: True to indicate that the skip has been handled
                  or False to stop execution
        """
        return False

    def processing_complete(self, skipped):
        pass

    def should_process(self, doc):
        """
        :param doc: the document to filter
        :return: True if this doc should be migrated
        """
        return True

class CouchDocumentProcessor(object):
    """Process Couch Docs

    :param doc_type_map: Dict of `doc_type_name: model_class` pairs.
    :param doc_processor: A `BaseDocProcessor` object used to
    process documents.
    :param reset: Reset existing processor state (if any), causing all
    documents to be reconsidered for processing, if this is true.
    :param max_retry: Number of times to retry processing a document
    before giving up.
    :param chunk_size: Maximum number of records to read from couch at
    one time. It may be necessary to use a smaller chunk size if the
    records being processed are very large and the default chunk size of
    100 would exceed available memory.
    :param log_handler: instance of PaginateViewLogHandler to be notified of
    view events.
    """
    def __init__(self, doc_type_map, doc_processor, reset=False, max_retry=2, chunk_size=100, log_handler=None):
        self.doc_type_map = doc_type_map
        self.doc_processor = doc_processor
        self.reset = reset
        self.max_retry = max_retry
        self.chunk_size = chunk_size

        self.couchdb = next(iter(doc_type_map.values())).get_db()
        assert all(m.get_db() is self.couchdb for m in doc_type_map.values()), \
            "documents must live in same couch db: %s" % repr(doc_type_map)

        self.docs_by_type = ResumableDocsByTypeIterator(
            self.couchdb, doc_type_map, doc_processor.unique_key, chunk_size=chunk_size,
            log_handler=log_handler
        )

        self.visited = 0
        self.previously_visited = 0
        self.total = 0

        self.processed = 0
        self.skipped = 0

        self.start = None

    def has_started(self):
        return bool(self.docs_by_type.progress_info)

    @property
    def session_visited(self):
        return self.visited - self.previously_visited

    @property
    def session_total(self):
        return self.total - self.previously_visited

    @property
    def attempted(self):
        return self.processed + self.skipped

    @property
    def timing(self):
        """Returns a tuple of (elapsed, remaining)"""
        elapsed = datetime.now() - self.start
        if self.session_visited > self.session_total:
            remaining = "?"
        else:
            session_remaining = self.session_total - self.session_visited
            remaining = elapsed / self.session_visited * session_remaining
        return elapsed, remaining

    def _setup(self):
        from corehq.dbaccessors.couchapps.all_docs import get_doc_count_by_type

        self.total = sum(get_doc_count_by_type(self.couchdb, doc_type)
                    for doc_type in self.doc_type_map)

        if self.reset:
            self.docs_by_type.discard_state()
        elif self.docs_by_type.progress_info:
            info = self.docs_by_type.progress_info
            old_total = info["total"]
            # Estimate already visited based on difference of old/new
            # totals. The theory is that new or deleted records will be
            # evenly distributed across the entire set.
            self.visited = int(round(float(self.total) / old_total * info["visited"]))
            self.previously_visited = self.visited
        print("Processing {} documents{}: {}...".format(
            self.total,
            " (~{} already processed)".format(self.visited) if self.visited else "",
            ", ".join(sorted(self.doc_type_map))
        ))

        self.start = datetime.now()

    def run(self):
        """
        :returns: A tuple `(<num processed>, <num skipped>)`
        """
        self._setup()
        with self.doc_processor:
            for doc in self.docs_by_type:
                self._process_doc(doc)
                self._update_progress()

        self._processing_complete()

        return self.processed, self.skipped

    def _process_doc(self, doc):
        if not self.doc_processor.should_process(doc):
            return

        ok = self.doc_processor.process_doc(doc, self.couchdb)
        if ok:
            self.processed += 1
        else:
            try:
                self.docs_by_type.retry(doc, self.max_retry)
            except TooManyRetries:
                if not self.doc_processor.handle_skip(doc):
                    raise
                else:
                    print("Skip: {doc_type} {_id}".format(**doc))
                    self.skipped += 1

    def _update_progress(self):
        self.visited += 1
        if self.visited % self.chunk_size == 0:
            self.docs_by_type.progress_info = {"visited": self.visited, "total": self.total}

        if self.attempted % self.chunk_size == 0:
            elapsed, remaining = self.timing
            print("Processed {}/{} of {} documents in {} ({} remaining)"
                  .format(self.processed, self.visited, self.total, elapsed, remaining))

    def _processing_complete(self):
        if self.session_visited:
            self.docs_by_type.progress_info = {"visited": self.visited, "total": self.total}
        self.doc_processor.processing_complete(self.skipped)
        print("Processed {}/{} of {} documents ({} previously processed, {} filtered out).".format(
            self.processed,
            self.visited,
            self.total,
            self.previously_visited,
            self.session_visited - self.attempted
        ))
        if self.skipped:
            print(DOCS_SKIPPED_WARNING.format(self.skipped))

class BulkDocProcessorLogHandler(PaginateViewLogHandler):

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

    def view_ending(self, db, view_name, kwargs, total_emitted, time):
        self.processor.process_chunk()

class BulkDocProcessor(CouchDocumentProcessor):
    """Process couch docs for a single doc type in batches

    The bulk doc processor will send a batch of documents to the document
    processor. If the processor does not respond with True then
    the iteration is halted. Restarting the iteration will start by
    re-sending the previous chunk to the processor.

    The size of the batches passed to the document processor may vary
    depending on how they are being filtered by the
    document processor but will never exceed ``chunk_size``.

    :param doc_type_map: Dict containing a single `doc_type_name: model_class` pair.
    :param doc_processor: A `BaseDocProcessor` object used to
    process documents.
    :param reset: Reset existing processor state (if any), causing all
    documents to be reconsidered for processing, if this is true.
    :param chunk_size: Maximum number of records to read from couch at
    one time. It may be necessary to use a smaller chunk size if the
    records being processed are very large and the default chunk size of
    100 would exceed available memory.
    """
    def __init__(self, doc_type_map, doc_processor, reset=False, chunk_size=100):
        assert len(doc_type_map) == 1
        super(BulkDocProcessor, self).__init__(
            doc_type_map, doc_processor, reset=reset, chunk_size=chunk_size,
            log_handler=BulkDocProcessorLogHandler(self)
        )
        self.changes = []

    def _process_doc(self, doc):
        if self.doc_processor.should_process(doc):
            self.changes.append(doc)

    def process_chunk(self):
        """Called by the BulkDocProcessorLogHandler"""
        ok = self.doc_processor.process_bulk_docs(self.changes, self.couchdb)
        if ok:
            self.processed += len(self.changes)
            self.changes = []
        else:
            raise BulkProcessingFailed("Processing batch failed")

    def _update_progress(self):
        self.visited += 1
        if self.visited % self.chunk_size == 0:
            self.docs_by_type.progress_info = {"visited": self.visited, "total": self.total}

            elapsed, remaining = self.timing
            print("Processed {}/{} of {} documents in {} ({} remaining)"
                  .format(self.processed, self.visited, self.total, elapsed, remaining))

    def _processing_complete(self):
        if len(self.changes):
            self._process_chunk()

        super(BulkDocProcessor, self)._processing_complete()
<|endoftext|>