<commit_before>from collections import namedtuple
import logging

from couchdbkit import ResourceNotFound
from django.db.models import Q
from casexml.apps.case.const import UNOWNED_EXTENSION_OWNER_ID, CASE_INDEX_EXTENSION
from casexml.apps.case.dbaccessors import get_extension_chain
from casexml.apps.case.signals import cases_received
from casexml.apps.case.util import validate_phone_datetime
from casexml.apps.phone.cleanliness import should_create_flags_on_submission
from casexml.apps.phone.models import OwnershipCleanlinessFlag
from corehq.toggles import LOOSE_SYNC_TOKEN_VALIDATION, EXTENSION_CASES_SYNC_ENABLED
from corehq.apps.users.util import SYSTEM_USER_ID
from corehq.form_processor.interfaces.processor import FormProcessorInterface
from couchforms.models import XFormInstance
from casexml.apps.case.exceptions import (
    NoDomainProvided,
)
from django.conf import settings
from dimagi.utils.couch.database import iter_docs

from casexml.apps.case import const
from casexml.apps.case.models import CommCareCase
from casexml.apps.case.xml.parser import case_update_from_block
from dimagi.utils.logging import notify_exception

# Lightweight class used to store the dirtyness of a case/owner pair.
DirtinessFlag = namedtuple('DirtinessFlag', ['case_id', 'owner_id'])

class CaseProcessingResult(object):
    """
    Lightweight class used to collect results of case processing
    """

    def __init__(self, domain, cases, dirtiness_flags, extensions_to_close=None):
        self.domain = domain
        self.cases = cases
        self.dirtiness_flags = dirtiness_flags
        if extensions_to_close is None:
            extensions_to_close = set()
        self.extensions_to_close = extensions_to_close

    def get_clean_owner_ids(self):
        dirty_flags = self.get_flags_to_save()
        return {c.owner_id for c in self.cases if c.owner_id and c.owner_id not in dirty_flags}

    def set_cases(self, cases):
        self.cases = cases

    def get_flags_to_save(self):
        return {f.owner_id: f.case_id for f in self.dirtiness_flags}

    def close_extensions(self):
        from casexml.apps.case.cleanup import close_cases
        extensions_to_close = list(self.extensions_to_close)
        if extensions_to_close:
            return close_cases(list(self.extensions_to_close), self.domain, SYSTEM_USER_ID)

    def commit_dirtiness_flags(self):
        """
        Updates any dirtiness flags in the database.
        """
        if self.domain:
            flags_to_save = self.get_flags_to_save()
            if should_create_flags_on_submission(self.domain):
                assert settings.UNIT_TESTING  # this is currently only true when unit testing
                all_touched_ids = set(flags_to_save.keys()) | self.get_clean_owner_ids()
                to_update = {f.owner_id: f for f in OwnershipCleanlinessFlag.objects.filter(
                    domain=self.domain,
                    owner_id__in=list(all_touched_ids),
                )}
                for owner_id in all_touched_ids:
                    if owner_id not in to_update:
                        # making from scratch - default to clean, but set to dirty if needed
                        flag = OwnershipCleanlinessFlag(domain=self.domain, owner_id=owner_id, is_clean=True)
                        if owner_id in flags_to_save:
                            flag.is_clean = False
                            flag.hint = flags_to_save[owner_id]
                        flag.save()
                    else:
                        # updating - only save if we are marking dirty or setting a hint
                        flag = to_update[owner_id]
                        if owner_id in flags_to_save and (flag.is_clean or not flag.hint):
                            flag.is_clean = False
                            flag.hint = flags_to_save[owner_id]
                            flag.save()
            else:
                # only update the flags that are already in the database
                flags_to_update = OwnershipCleanlinessFlag.objects.filter(
                    Q(domain=self.domain),
                    Q(owner_id__in=flags_to_save.keys()),
                    Q(is_clean=True) | Q(hint__isnull=True)
                )
                for flag in flags_to_update:
                    flag.is_clean = False
                    flag.hint = flags_to_save[flag.owner_id]
                    flag.save()

def process_cases_with_casedb(xforms, case_db, config=None):
    config = config or CaseProcessingConfig()
    case_processing_result = _get_or_update_cases(xforms, case_db)
    cases = case_processing_result.cases
    xform = xforms[0]

    # handle updating the sync records for apps that use sync mode
    try:
        relevant_log = xform.get_sync_token()
    except ResourceNotFound:
        if LOOSE_SYNC_TOKEN_VALIDATION.enabled(xform.domain):
            relevant_log = None
        else:
            raise

    if relevant_log:
        # in reconciliation mode, things can be unexpected
        relevant_log.strict = config.strict_asserts
        from casexml.apps.case.util import update_sync_log_with_checks
        update_sync_log_with_checks(relevant_log, xform, cases, case_db,
                                    case_id_blacklist=config.case_id_blacklist)

    try:
        cases_received.send(sender=None, xform=xform, cases=cases)
    except Exception as e:
        # don't let the exceptions in signals prevent standard case processing
        notify_exception(
            None,
            'something went wrong sending the cases_received signal '
            'for form %s: %s' % (xform.form_id, e)
        )

    for case in cases:
        case_db.post_process_case(case, xform)
        case_db.mark_changed(case)

    case_processing_result.set_cases(cases)
    return case_processing_result

class CaseProcessingConfig(object):

    def __init__(self, strict_asserts=True, case_id_blacklist=None):
        self.strict_asserts = strict_asserts
        self.case_id_blacklist = case_id_blacklist if case_id_blacklist is not None else []

    def __repr__(self):
        return 'strict: {strict}, ids: {ids}'.format(
            strict=self.strict_asserts,
            ids=", ".join(self.case_id_blacklist)
        )

def get_and_check_xform_domain(xform):
    try:
        domain = xform.domain
    except AttributeError:
        domain = None

    if not domain and settings.CASEXML_FORCE_DOMAIN_CHECK:
        raise NoDomainProvided()

    return domain

def _get_or_update_cases(xforms, case_db):
    """
    Given an xform document, update any case blocks found within it,
    returning a dictionary mapping the case ids affected to the
    couch case document objects
    """
    domain = getattr(case_db, 'domain', None)
    touched_cases = FormProcessorInterface(domain).get_cases_from_forms(case_db, xforms)

    # once we've gotten through everything, validate all indices
    # and check for new dirtiness flags
    def _validate_indices(case):
        if case.indices:
            for index in case.indices:
                # call get and not doc_exists to force domain checking
                # see CaseDbCache._validate_case
                referenced_case = case_db.get(index.referenced_id)
                if not referenced_case:
                    # just log, don't raise an error or modify the index
                    logging.error(
                        "Case '%s' references non-existent case '%s'",
                        case.case_id,
                        index.referenced_id,
                    )

    def _get_dirtiness_flags_for_outgoing_indices(case, tree_owners=None):
        """ if the outgoing indices touch cases owned by another user this cases owner is dirty """
        if tree_owners is None:
            tree_owners = set()

        extension_indices = [index for index in case.indices if index.relationship == CASE_INDEX_EXTENSION]

        unowned_host_cases = []
        for index in extension_indices:
            host_case = case_db.get(index.referenced_id)
            if (
                host_case
                and host_case.owner_id == UNOWNED_EXTENSION_OWNER_ID
                and host_case not in unowned_host_cases
            ):
                unowned_host_cases.append(host_case)

        owner_ids = {case_db.get(index.referenced_id).owner_id
                     for index in case.indices if case_db.get(index.referenced_id)} | tree_owners
        potential_clean_owner_ids = owner_ids | set([UNOWNED_EXTENSION_OWNER_ID])
        more_than_one_owner_touched = len(owner_ids) > 1
        touches_different_owner = len(owner_ids) == 1 and case.owner_id not in potential_clean_owner_ids

        if (more_than_one_owner_touched or touches_different_owner):
            yield DirtinessFlag(case.case_id, case.owner_id)
            if extension_indices:
                # If this case is an extension, each of the touched cases is also dirty
                for index in case.indices:
                    referenced_case = case_db.get(index.referenced_id)
                    yield DirtinessFlag(referenced_case.case_id, referenced_case.owner_id)

        if case.owner_id != UNOWNED_EXTENSION_OWNER_ID:
            tree_owners.add(case.owner_id)
        for unowned_host_case in unowned_host_cases:
            # A host case of this extension is unowned, which means it could potentially touch an owned case
            # Check these unowned cases' outgoing indices and mark dirty if appropriate
            for dirtiness_flag in _get_dirtiness_flags_for_outgoing_indices(unowned_host_case,
                                                                            tree_owners=tree_owners):
                yield dirtiness_flag

    def _get_dirtiness_flags_for_child_cases(cases):
        child_cases = case_db.get_reverse_indexed_cases([c.case_id for c in cases])
        case_owner_map = dict((case.case_id, case.owner_id) for case in cases)
        for child_case in child_cases:
            for index in child_case.indices:
                if (index.referenced_id in case_owner_map
                        and child_case.owner_id != case_owner_map[index.referenced_id]):
                    yield DirtinessFlag(child_case.case_id, child_case.owner_id)

    def _get_dirtiness_flags_for_reassigned_case(case_metas):
        # for reassigned cases, we mark them temporarily dirty to allow phones to sync
        # the latest changes. these will get cleaned up when the weekly rebuild triggers
        for case_update_meta in case_metas:
            if _is_change_of_ownership(case_update_meta.previous_owner_id, case_update_meta.case.owner_id):
                yield DirtinessFlag(case_update_meta.case.case_id, case_update_meta.previous_owner_id)

    dirtiness_flags = []
    extensions_to_close = set()

    # process the temporary dirtiness flags first so that any hints for real dirtiness get overridden
    dirtiness_flags += list(_get_dirtiness_flags_for_reassigned_case(touched_cases.values()))
    for case_update_meta in touched_cases.values():
        _validate_indices(case_update_meta.case)
        extensions_to_close = extensions_to_close | get_extensions_to_close(case_update_meta.case, domain)
        dirtiness_flags += list(_get_dirtiness_flags_for_outgoing_indices(case_update_meta.case))
    dirtiness_flags += list(_get_dirtiness_flags_for_child_cases([meta.case for meta in touched_cases.values()]))

    return CaseProcessingResult(
        domain,
        [update.case for update in touched_cases.values()],
        dirtiness_flags,
        extensions_to_close
    )

def _is_change_of_ownership(previous_owner_id, next_owner_id):
    return (
        previous_owner_id
        and previous_owner_id != UNOWNED_EXTENSION_OWNER_ID
        and previous_owner_id != next_owner_id
    )

def get_extensions_to_close(case, domain):
    outgoing_extension_indices = [index.relationship for index in case.indices
                                  if index.relationship == CASE_INDEX_EXTENSION]
    if not outgoing_extension_indices and case.closed and EXTENSION_CASES_SYNC_ENABLED.enabled(domain):
        return get_extension_chain([case.case_id], domain)
    else:
        return set()

def is_device_report(doc):
    """exclude device reports"""
    device_report_xmlns = "http://code.javarosa.org/devicereport"

    def _from_form_dict(doc):
        return isinstance(doc, dict) and "@xmlns" in doc and doc["@xmlns"] == device_report_xmlns

    def _from_xform_instance(doc):
        return getattr(doc, 'xmlns', None) == device_report_xmlns

    return _from_form_dict(doc) or _from_xform_instance(doc)

def has_case_id(case_block):
    return const.CASE_TAG_ID in case_block or const.CASE_ATTR_ID in case_block

CaseBlockWithPath = namedtuple('CaseBlockWithPath', ['caseblock', 'path'])

def extract_case_blocks(doc, include_path=False):
    """
    Extract all case blocks from a document, returning an array of dictionaries
    with the data in each case.

    The json returned is not normalized for casexml version;
    for that get_case_updates is better.

    if `include_path` is True then instead of returning just the case block it will
    return a dict with the following structure:

    {
       "caseblock": caseblock
       "path": ["form", "path", "to", "block"]
    }

    Repeat nodes will all share the same path.
    """
    if isinstance(doc, XFormInstance):
        form = doc.to_json()['form']
    elif isinstance(doc, dict):
        form = doc
    else:
        form = doc.form_data

    return [struct if include_path else struct.caseblock for struct in _extract_case_blocks(form)]

def _extract_case_blocks(data, path=None):
    """
    helper for extract_case_blocks

    data must be json representing a node in an xform submission
    """
    path = path or []
    if isinstance(data, list):
        for item in data:
            for case_block in _extract_case_blocks(item, path=path):
                yield case_block
    elif isinstance(data, dict) and not is_device_report(data):
        for key, value in data.items():
            new_path = path + [key]
            if const.CASE_TAG == key:
                # it's a case block! Stop recursion and add to this value
                if isinstance(value, list):
                    case_blocks = value
                else:
                    case_blocks = [value]

                for case_block in case_blocks:
                    if has_case_id(case_block):
                        validate_phone_datetime(
                            case_block.get('@date_modified'), none_ok=True)
                        yield CaseBlockWithPath(caseblock=case_block, path=path)
            else:
                for case_block in _extract_case_blocks(value, path=new_path):
                    yield case_block

def get_case_updates(xform):
    return [case_update_from_block(cb) for cb in extract_case_blocks(xform)]

def get_case_ids_from_form(xform):
    return set(cu.id for cu in get_case_updates(xform))

def cases_referenced_by_xform(xform):
    """
    Returns a list of CommCareCase or CommCareCaseSQL given a JSON
    representation of an XFormInstance
    """
    from corehq.form_processor.backends.couch.dbaccessors import CaseAccessorCouch
    from corehq.form_processor.interfaces.dbaccessors import CaseAccessors
    case_ids = get_case_ids_from_form(xform)
    domain = get_and_check_xform_domain(xform)
    case_accessor = CaseAccessors(domain)
    if domain is None:
        assert case_accessor.db_accessor == CaseAccessorCouch
    return list(case_accessor.get_cases(list(case_ids)))
<commit_msg>Remove unused imports<commit_after>from collections import namedtuple
import logging

from couchdbkit import ResourceNotFound
from django.db.models import Q
from casexml.apps.case.const import UNOWNED_EXTENSION_OWNER_ID, CASE_INDEX_EXTENSION
from casexml.apps.case.signals import cases_received
from casexml.apps.case.util import validate_phone_datetime
from casexml.apps.phone.cleanliness import should_create_flags_on_submission
from casexml.apps.phone.models import OwnershipCleanlinessFlag
from corehq.toggles import LOOSE_SYNC_TOKEN_VALIDATION, EXTENSION_CASES_SYNC_ENABLED
from corehq.apps.users.util import SYSTEM_USER_ID
from corehq.form_processor.interfaces.processor import FormProcessorInterface
from couchforms.models import XFormInstance
from casexml.apps.case.exceptions import (
    NoDomainProvided,
)
from django.conf import settings

from casexml.apps.case import const
from casexml.apps.case.xml.parser import case_update_from_block
from dimagi.utils.logging import notify_exception

# Lightweight class used to store the dirtyness of a case/owner pair.
DirtinessFlag = namedtuple('DirtinessFlag', ['case_id', 'owner_id'])

class CaseProcessingResult(object):
    """
    Lightweight class used to collect results of case processing
    """

    def __init__(self, domain, cases, dirtiness_flags, extensions_to_close=None):
        self.domain = domain
        self.cases = cases
        self.dirtiness_flags = dirtiness_flags
        if extensions_to_close is None:
            extensions_to_close = set()
        self.extensions_to_close = extensions_to_close

    def get_clean_owner_ids(self):
        dirty_flags = self.get_flags_to_save()
        return {c.owner_id for c in self.cases if c.owner_id and c.owner_id not in dirty_flags}

    def set_cases(self, cases):
        self.cases = cases

    def get_flags_to_save(self):
        return {f.owner_id: f.case_id for f in self.dirtiness_flags}

    def close_extensions(self):
        from casexml.apps.case.cleanup import close_cases
        extensions_to_close = list(self.extensions_to_close)
        if extensions_to_close:
            return close_cases(list(self.extensions_to_close), self.domain, SYSTEM_USER_ID)

    def commit_dirtiness_flags(self):
        """
        Updates any dirtiness flags in the database.
        """
        if self.domain:
            flags_to_save = self.get_flags_to_save()
            if should_create_flags_on_submission(self.domain):
                assert settings.UNIT_TESTING  # this is currently only true when unit testing
                all_touched_ids = set(flags_to_save.keys()) | self.get_clean_owner_ids()
                to_update = {f.owner_id: f for f in OwnershipCleanlinessFlag.objects.filter(
                    domain=self.domain,
                    owner_id__in=list(all_touched_ids),
                )}
                for owner_id in all_touched_ids:
                    if owner_id not in to_update:
                        # making from scratch - default to clean, but set to dirty if needed
                        flag = OwnershipCleanlinessFlag(domain=self.domain, owner_id=owner_id, is_clean=True)
                        if owner_id in flags_to_save:
                            flag.is_clean = False
                            flag.hint = flags_to_save[owner_id]
                        flag.save()
                    else:
                        # updating - only save if we are marking dirty or setting a hint
                        flag = to_update[owner_id]
                        if owner_id in flags_to_save and (flag.is_clean or not flag.hint):
                            flag.is_clean = False
                            flag.hint = flags_to_save[owner_id]
                            flag.save()
            else:
                # only update the flags that are already in the database
                flags_to_update = OwnershipCleanlinessFlag.objects.filter(
                    Q(domain=self.domain),
                    Q(owner_id__in=flags_to_save.keys()),
                    Q(is_clean=True) | Q(hint__isnull=True)
                )
                for flag in flags_to_update:
                    flag.is_clean = False
                    flag.hint = flags_to_save[flag.owner_id]
                    flag.save()

def process_cases_with_casedb(xforms, case_db, config=None):
    config = config or CaseProcessingConfig()
    case_processing_result = _get_or_update_cases(xforms, case_db)
    cases = case_processing_result.cases
    xform = xforms[0]

    # handle updating the sync records for apps that use sync mode
    try:
        relevant_log = xform.get_sync_token()
    except ResourceNotFound:
        if LOOSE_SYNC_TOKEN_VALIDATION.enabled(xform.domain):
            relevant_log = None
        else:
            raise

    if relevant_log:
        # in reconciliation mode, things can be unexpected
        relevant_log.strict = config.strict_asserts
        from casexml.apps.case.util import update_sync_log_with_checks
        update_sync_log_with_checks(relevant_log, xform, cases, case_db,
                                    case_id_blacklist=config.case_id_blacklist)

    try:
        cases_received.send(sender=None, xform=xform, cases=cases)
    except Exception as e:
        # don't let the exceptions in signals prevent standard case processing
        notify_exception(
            None,
            'something went wrong sending the cases_received signal '
            'for form %s: %s' % (xform.form_id, e)
        )

    for case in cases:
        case_db.post_process_case(case, xform)
        case_db.mark_changed(case)

    case_processing_result.set_cases(cases)
    return case_processing_result

class CaseProcessingConfig(object):

    def __init__(self, strict_asserts=True, case_id_blacklist=None):
        self.strict_asserts = strict_asserts
        self.case_id_blacklist = case_id_blacklist if case_id_blacklist is not None else []

    def __repr__(self):
        return 'strict: {strict}, ids: {ids}'.format(
            strict=self.strict_asserts,
            ids=", ".join(self.case_id_blacklist)
        )

def get_and_check_xform_domain(xform):
    try:
        domain = xform.domain
    except AttributeError:
        domain = None

    if not domain and settings.CASEXML_FORCE_DOMAIN_CHECK:
        raise NoDomainProvided()

    return domain

def _get_or_update_cases(xforms, case_db):
    """
    Given an xform document, update any case blocks found within it,
    returning a dictionary mapping the case ids affected to the
    couch case document objects
    """
    domain = getattr(case_db, 'domain', None)
    touched_cases = FormProcessorInterface(domain).get_cases_from_forms(case_db, xforms)

    # once we've gotten through everything, validate all indices
    # and check for new dirtiness flags
    def _validate_indices(case):
        if case.indices:
            for index in case.indices:
                # call get and not doc_exists to force domain checking
                # see CaseDbCache._validate_case
                referenced_case = case_db.get(index.referenced_id)
                if not referenced_case:
                    # just log, don't raise an error or modify the index
                    logging.error(
                        "Case '%s' references non-existent case '%s'",
                        case.case_id,
                        index.referenced_id,
                    )

    def _get_dirtiness_flags_for_outgoing_indices(case, tree_owners=None):
        """ if the outgoing indices touch cases owned by another user this cases owner is dirty """
        if tree_owners is None:
            tree_owners = set()

        extension_indices = [index for index in case.indices if index.relationship == CASE_INDEX_EXTENSION]

        unowned_host_cases = []
        for index in extension_indices:
            host_case = case_db.get(index.referenced_id)
            if (
                host_case
                and host_case.owner_id == UNOWNED_EXTENSION_OWNER_ID
                and host_case not in unowned_host_cases
            ):
                unowned_host_cases.append(host_case)

        owner_ids = {case_db.get(index.referenced_id).owner_id
                     for index in case.indices if case_db.get(index.referenced_id)} | tree_owners
        potential_clean_owner_ids = owner_ids | set([UNOWNED_EXTENSION_OWNER_ID])
        more_than_one_owner_touched = len(owner_ids) > 1
        touches_different_owner = len(owner_ids) == 1 and case.owner_id not in potential_clean_owner_ids

        if (more_than_one_owner_touched or touches_different_owner):
            yield DirtinessFlag(case.case_id, case.owner_id)
            if extension_indices:
                # If this case is an extension, each of the touched cases is also dirty
                for index in case.indices:
                    referenced_case = case_db.get(index.referenced_id)
                    yield DirtinessFlag(referenced_case.case_id, referenced_case.owner_id)

        if case.owner_id != UNOWNED_EXTENSION_OWNER_ID:
            tree_owners.add(case.owner_id)
        for unowned_host_case in unowned_host_cases:
            # A host case of this extension is unowned, which means it could potentially touch an owned case
            # Check these unowned cases' outgoing indices and mark dirty if appropriate
            for dirtiness_flag in _get_dirtiness_flags_for_outgoing_indices(unowned_host_case,
                                                                            tree_owners=tree_owners):
                yield dirtiness_flag

    def _get_dirtiness_flags_for_child_cases(cases):
        child_cases = case_db.get_reverse_indexed_cases([c.case_id for c in cases])
        case_owner_map = dict((case.case_id, case.owner_id) for case in cases)
        for child_case in child_cases:
            for index in child_case.indices:
                if (index.referenced_id in case_owner_map
                        and child_case.owner_id != case_owner_map[index.referenced_id]):
                    yield DirtinessFlag(child_case.case_id, child_case.owner_id)

    def _get_dirtiness_flags_for_reassigned_case(case_metas):
        # for reassigned cases, we mark them temporarily dirty to allow phones to sync
        # the latest changes. these will get cleaned up when the weekly rebuild triggers
        for case_update_meta in case_metas:
            if _is_change_of_ownership(case_update_meta.previous_owner_id, case_update_meta.case.owner_id):
                yield DirtinessFlag(case_update_meta.case.case_id, case_update_meta.previous_owner_id)

    dirtiness_flags = []
    extensions_to_close = set()

    # process the temporary dirtiness flags first so that any hints for real dirtiness get overridden
    dirtiness_flags += list(_get_dirtiness_flags_for_reassigned_case(touched_cases.values()))
    for case_update_meta in touched_cases.values():
        _validate_indices(case_update_meta.case)
        extensions_to_close = extensions_to_close | get_extensions_to_close(case_update_meta.case, domain)
        dirtiness_flags += list(_get_dirtiness_flags_for_outgoing_indices(case_update_meta.case))
    dirtiness_flags += list(_get_dirtiness_flags_for_child_cases([meta.case for meta in touched_cases.values()]))

    return CaseProcessingResult(
        domain,
        [update.case for update in touched_cases.values()],
        dirtiness_flags,
        extensions_to_close
    )

def _is_change_of_ownership(previous_owner_id, next_owner_id):
    return (
        previous_owner_id
        and previous_owner_id != UNOWNED_EXTENSION_OWNER_ID
        and previous_owner_id != next_owner_id
    )

def get_extensions_to_close(case, domain):
    outgoing_extension_indices = [index.relationship for index in case.indices
                                  if index.relationship == CASE_INDEX_EXTENSION]
    if not outgoing_extension_indices and case.closed and EXTENSION_CASES_SYNC_ENABLED.enabled(domain):
        return get_extension_chain([case.case_id], domain)
    else:
        return set()

def is_device_report(doc):
    """exclude device reports"""
    device_report_xmlns = "http://code.javarosa.org/devicereport"

    def _from_form_dict(doc):
        return isinstance(doc, dict) and "@xmlns" in doc and doc["@xmlns"] == device_report_xmlns

    def _from_xform_instance(doc):
        return getattr(doc, 'xmlns', None) == device_report_xmlns

    return _from_form_dict(doc) or _from_xform_instance(doc)

def has_case_id(case_block):
    return const.CASE_TAG_ID in case_block or const.CASE_ATTR_ID in case_block

CaseBlockWithPath = namedtuple('CaseBlockWithPath', ['caseblock', 'path'])

def extract_case_blocks(doc, include_path=False):
    """
    Extract all case blocks from a document, returning an array of dictionaries
    with the data in each case.

    The json returned is not normalized for casexml version;
    for that get_case_updates is better.

    if `include_path` is True then instead of returning just the case block it will
    return a dict with the following structure:

    {
       "caseblock": caseblock
       "path": ["form", "path", "to", "block"]
    }

    Repeat nodes will all share the same path.
    """
    if isinstance(doc, XFormInstance):
        form = doc.to_json()['form']
    elif isinstance(doc, dict):
        form = doc
    else:
        form = doc.form_data

    return [struct if include_path else struct.caseblock for struct in _extract_case_blocks(form)]

def _extract_case_blocks(data, path=None):
    """
    helper for extract_case_blocks

    data must be json representing a node in an xform submission
    """
    path = path or []
    if isinstance(data, list):
        for item in data:
            for case_block in _extract_case_blocks(item, path=path):
                yield case_block
    elif isinstance(data, dict) and not is_device_report(data):
        for key, value in data.items():
            new_path = path + [key]
            if const.CASE_TAG == key:
                # it's a case block! Stop recursion and add to this value
                if isinstance(value, list):
                    case_blocks = value
                else:
                    case_blocks = [value]

                for case_block in case_blocks:
                    if has_case_id(case_block):
                        validate_phone_datetime(
                            case_block.get('@date_modified'), none_ok=True)
                        yield CaseBlockWithPath(caseblock=case_block, path=path)
            else:
                for case_block in _extract_case_blocks(value, path=new_path):
                    yield case_block

def get_case_updates(xform):
    return [case_update_from_block(cb) for cb in extract_case_blocks(xform)]

def get_case_ids_from_form(xform):
    return set(cu.id for cu in get_case_updates(xform))

def cases_referenced_by_xform(xform):
    """
    Returns a list of CommCareCase or CommCareCaseSQL given a JSON
    representation of an XFormInstance
    """
    from corehq.form_processor.backends.couch.dbaccessors import CaseAccessorCouch
    from corehq.form_processor.interfaces.dbaccessors import CaseAccessors
    case_ids = get_case_ids_from_form(xform)
    domain = get_and_check_xform_domain(xform)
    case_accessor = CaseAccessors(domain)
    if domain is None:
        assert case_accessor.db_accessor == CaseAccessorCouch
    return list(case_accessor.get_cases(list(case_ids)))
<|endoftext|>


<commit_before># (C) Datadog, Inc. 2018
# (C) Justin Slattery <Justin.Slattery@fzysqr.com> 2013
# All rights reserved
# Licensed under a 3-clause BSD style license (see LICENSE)

from __future__ import division

import re
import time
from collections import defaultdict

import requests
from six import string_types

from datadog_checks.checks import AgentCheck
from datadog_checks.couchbase.couchbase_consts import (
    BUCKET_STATS,
    COUCHBASE_STATS_PATH,
    COUCHBASE_VITALS_PATH,
    NODE_CLUSTER_SERVICE_CHECK_NAME,
    NODE_HEALTH_SERVICE_CHECK_NAME,
    NODE_HEALTH_TRANSLATION,
    NODE_MEMBERSHIP_TRANSLATION,
    QUERY_STATS,
    SECONDS_VALUE_PATTERN,
    SERVICE_CHECK_NAME,
    SOURCE_TYPE_NAME,
    TO_SECONDS,
)
from datadog_checks.utils.containers import hash_mutable

class Couchbase(AgentCheck):
    """
    Extracts stats from Couchbase via its REST API
    http://docs.couchbase.com/couchbase-manual-2.0/#using-the-rest-api
    """

    HTTP_CONFIG_REMAPPER = {'user': {'name': 'username'}, 'ssl_verify': {'name': 'tls_verify'}}

    class CouchbaseInstanceState(object):
        def __init__(self):
            self.previous_status = None

    def __init__(self, name, init_config, instances):
        super(Couchbase, self).__init__(name, init_config, instances)

        # Keep track of all instances
        self._instance_states = defaultdict(lambda: self.CouchbaseInstanceState())

    def _create_metrics(self, data, instance_state, server, tags=None):
        # Get storage metrics
        storage_totals = data['stats']['storageTotals']
        for key, storage_type in storage_totals.items():
            for metric_name, val in storage_type.items():
                if val is not None:
                    metric_name = 'couchbase.{}.{}'.format(key, self.camel_case_to_joined_lower(metric_name))
                    self.gauge(metric_name, val, tags=tags)

        # Get bucket metrics
        for bucket_name, bucket_stats in data['buckets'].items():
            metric_tags = ['bucket:{}'.format(bucket_name), 'device:{}'.format(bucket_name)]
            if tags:
                metric_tags.extend(tags)
            for metric_name, val in bucket_stats.items():
                if val is not None:
                    norm_metric_name = self.camel_case_to_joined_lower(metric_name)
                    if norm_metric_name in BUCKET_STATS:
                        full_metric_name = 'couchbase.by_bucket.{}'.format(norm_metric_name)
                        self.gauge(full_metric_name, val[0], tags=metric_tags)

        # Get node metrics
        for node_name, node_stats in data['nodes'].items():
            metric_tags = ['node:{}'.format(node_name), 'device:{}'.format(node_name)]
            if tags:
                metric_tags.extend(tags)
            for metric_name, val in node_stats['interestingStats'].items():
                if val is not None:
                    metric_name = 'couchbase.by_node.{}'.format(self.camel_case_to_joined_lower(metric_name))
                    self.gauge(metric_name, val, tags=metric_tags)

            # Get cluster health data
            self._process_cluster_health_data(node_name, node_stats, tags)

        # Get query metrics
        for metric_name, val in data['query'].items():
            if val is not None:
                norm_metric_name = self.camel_case_to_joined_lower(metric_name)
                if norm_metric_name in QUERY_STATS:
                    # for query times, the unit is part of the value, we need to extract it
                    if isinstance(val, string_types):
                        val = self.extract_seconds_value(val)

                    full_metric_name = 'couchbase.query.{}'.format(self.camel_case_to_joined_lower(norm_metric_name))
                    self.gauge(full_metric_name, val, tags=tags)

        # Get tasks, we currently only care about 'rebalance' tasks
        rebalance_status, rebalance_msg = data['tasks'].get('rebalance', (None, None))

        # Only fire an event when the state has changed
        if rebalance_status is not None and instance_state.previous_status != rebalance_status:
            rebalance_event = None

            # If we get an error, we create an error event with the msg we receive
            if rebalance_status == 'error':
                msg_title = 'Encountered an error while rebalancing'
                msg = rebalance_msg

                rebalance_event = self._create_event('error', msg_title, msg, server, tags=tags)

            # We only want to fire a 'completion' of a rebalance so make sure we're not firing an event on first run
            elif rebalance_status == 'notRunning' and instance_state.previous_status is not None:
                msg_title = 'Stopped rebalancing'
                msg = 'stopped rebalancing.'
                rebalance_event = self._create_event('info', msg_title, msg, server, tags=tags)

            # If a rebalance task is running, fire an event. This will also fire an event if a rebalance task was
            #   already running when the check first runs.
            elif rebalance_status == 'gracefulFailover':
                msg_title = 'Failing over gracefully'
                msg = 'is failing over gracefully.'
                rebalance_event = self._create_event('info', msg_title, msg, server, tags=tags)
            elif rebalance_status == 'rebalance':
                msg_title = 'Rebalancing'
                msg = 'is rebalancing.'
                rebalance_event = self._create_event('info', msg_title, msg, server, tags=tags)

            # Send the event
            if rebalance_event is not None:
                self.event(rebalance_event)

            # Update the status of this instance
            instance_state.previous_status = rebalance_status

    def _process_cluster_health_data(self, node_name, node_stats, tags):
        """
        Process and send cluster health data (i.e. cluster membership status and node health
        """

        # Tags for service check
        cluster_health_tags = list(tags) + ['node:{}'.format(node_name)]

        # Get the membership status of the node
        cluster_membership = node_stats.get('clusterMembership', None)
        membership_status = NODE_MEMBERSHIP_TRANSLATION.get(cluster_membership, AgentCheck.UNKNOWN)
        self.service_check(NODE_CLUSTER_SERVICE_CHECK_NAME, membership_status, tags=cluster_health_tags)

        # Get the health status of the node
        health = node_stats.get('status', None)
        health_status = NODE_HEALTH_TRANSLATION.get(health, AgentCheck.UNKNOWN)
        self.service_check(NODE_HEALTH_SERVICE_CHECK_NAME, health_status, tags=cluster_health_tags)

    def _create_event(self, alert_type, msg_title, msg, server, tags=None):
        """
        Create an event object
        """
        msg_title = 'Couchbase {}: {}'.format(server, msg_title)
        msg = 'Couchbase instance {} {}'.format(server, msg)

        return {
            'timestamp': int(time.time()),
            'event_type': 'couchbase_rebalance',
            'msg_text': msg,
            'msg_title': msg_title,
            'alert_type': alert_type,
            'source_type_name': SOURCE_TYPE_NAME,
            'aggregation_key': server,
            'tags': tags,
        }

    def _get_stats(self, url):
        """
        Hit a given URL and return the parsed json.
        """
        r = self.http.get(url)
        r.raise_for_status()
        return r.json()

    def check(self, instance):
        instance_state = self._instance_states[hash_mutable(instance)]

        server = instance.get('server', None)
        if server is None:
            raise Exception("The server must be specified")
        tags = instance.get('tags', [])
        # Clean up tags in case there was a None entry in the instance
        # e.g. if the yaml contains tags: but no actual tags
        if tags is None:
            tags = []
        else:
            tags = list(set(tags))
        tags.append('instance:{}'.format(server))
        data = self.get_data(server, instance)
        self._create_metrics(data, instance_state, server, tags=list(set(tags)))

    def get_data(self, server, instance):
        # The dictionary to be returned.
        couchbase = {'stats': None, 'buckets': {}, 'nodes': {}, 'query': {}, 'tasks': {}}

        # build couchbase stats entry point
        url = '{}{}'.format(server, COUCHBASE_STATS_PATH)

        # Fetch initial stats and capture a service check based on response.
        service_check_tags = instance.get('tags', [])
        if service_check_tags is None:
            service_check_tags = []
        else:
            service_check_tags = list(set(service_check_tags))
        service_check_tags.append('instance:{}'.format(server))
        try:
            overall_stats = self._get_stats(url)
            # No overall stats? bail out now
            if overall_stats is None:
                raise Exception("No data returned from couchbase endpoint: {}".format(url))
        except requests.exceptions.HTTPError as e:
            self.service_check(SERVICE_CHECK_NAME, AgentCheck.CRITICAL, tags=service_check_tags, message=str(e))
            raise
        except Exception as e:
            self.service_check(SERVICE_CHECK_NAME, AgentCheck.CRITICAL, tags=service_check_tags, message=str(e))
            raise
        else:
            self.service_check(SERVICE_CHECK_NAME, AgentCheck.OK, tags=service_check_tags)

        couchbase['stats'] = overall_stats

        nodes = overall_stats['nodes']

        # Next, get all the nodes
        if nodes is not None:
            for node in nodes:
                couchbase['nodes'][node['hostname']] = node

        # Next, get all buckets .
        endpoint = overall_stats['buckets']['uri']

        url = '{}{}'.format(server, endpoint)
        buckets = self._get_stats(url)

        if buckets is not None:
            for bucket in buckets:
                bucket_name = bucket['name']

                # Fetch URI for the stats bucket
                endpoint = bucket['stats']['uri']
                url = '{}{}'.format(server, endpoint)

                try:
                    bucket_stats = self._get_stats(url)
                except requests.exceptions.HTTPError:
                    url_backup = '{}/pools/nodes/buckets/{}/stats'.format(server, bucket_name)
                    bucket_stats = self._get_stats(url_backup)

                bucket_samples = bucket_stats['op']['samples']
                if bucket_samples is not None:
                    couchbase['buckets'][bucket['name']] = bucket_samples

        # Next, get the query monitoring data
        query_data = self._get_query_monitoring_data(instance)
        if query_data is not None:
            couchbase['query'] = query_data

        # Next, get all the tasks
        tasks_url = '{}{}/tasks'.format(server, COUCHBASE_STATS_PATH)
        try:
            tasks = self._get_stats(tasks_url)
            for task in tasks:
                task_type = task['type']

                # We only care about rebalance statuses
                if task_type != 'rebalance':
                    continue

                # Format the status so it's easier to understand
                if 'errorMessage' in task:
                    couchbase['tasks'][task_type] = ('error', task['errorMessage'])
                elif task['status'] == 'notRunning':
                    couchbase['tasks'][task_type] = (task['status'], None)

                # If the status is 'running', we want to differentiate between a regular rebalance and a graceful
                #   failover rebalance, so we use the subtype
                elif task['status'] == 'running':
                    couchbase['tasks'][task_type] = (task['subtype'], None)

                # Should only be 1 rebalance
                break

        except requests.exceptions.HTTPError:
            self.log.error("Error accessing the endpoint {}".format(url))

        return couchbase

    def _get_query_monitoring_data(self, instance):
        query_data = None
        query_monitoring_url = instance.get('query_monitoring_url')
        if query_monitoring_url:
            url = '{}{}'.format(query_monitoring_url, COUCHBASE_VITALS_PATH)
            try:
                query_data = self._get_stats(url)
            except requests.exceptions.RequestException:
                self.log.error(
                    "Error accessing the endpoint {}, make sure you're running at least "
                    "couchbase 4.5 to collect the query monitoring metrics".format(url)
                )

        return query_data

    # Takes a camelCased variable and returns a joined_lower equivalent.
    # Returns input if non-camelCase variable is detected.
    def camel_case_to_joined_lower(self, variable):
        # replace non-word with _
        converted_variable = re.sub(r'\W+', '_', variable)

        # insert _ in front of capital letters and lowercase the string
        converted_variable = re.sub('([A-Z])', r'_\g<1>', converted_variable).lower()

        # remove duplicate _
        converted_variable = re.sub('_+', '_', converted_variable)

        # handle special case of starting/ending underscores
        converted_variable = re.sub('^_|_$', '', converted_variable)

        return converted_variable

    # Takes a string with a time and a unit (e.g '3.45ms') and returns the value in seconds
    def extract_seconds_value(self, value):

        # When couchbase is set up, most of values are equal to 0 and are exposed as "0" and not "0s"
        # This statement is preventing values to be searched by the pattern (and break things)
        if value == '0':
            return 0

        match = SECONDS_VALUE_PATTERN.search(value)

        val, unit = match.group(1, 3)
        # They use the 'micro' symbol for microseconds so there is an encoding problem
        # so let's assume it's microseconds if we don't find the key in unit
        if unit not in TO_SECONDS:
            unit = 'us'

        return float(val) / TO_SECONDS[unit]
<commit_msg>Standardize logging format (#4905)<commit_after># (C) Datadog, Inc. 2018
# (C) Justin Slattery <Justin.Slattery@fzysqr.com> 2013
# All rights reserved
# Licensed under a 3-clause BSD style license (see LICENSE)

from __future__ import division

import re
import time
from collections import defaultdict

import requests
from six import string_types

from datadog_checks.checks import AgentCheck
from datadog_checks.couchbase.couchbase_consts import (
    BUCKET_STATS,
    COUCHBASE_STATS_PATH,
    COUCHBASE_VITALS_PATH,
    NODE_CLUSTER_SERVICE_CHECK_NAME,
    NODE_HEALTH_SERVICE_CHECK_NAME,
    NODE_HEALTH_TRANSLATION,
    NODE_MEMBERSHIP_TRANSLATION,
    QUERY_STATS,
    SECONDS_VALUE_PATTERN,
    SERVICE_CHECK_NAME,
    SOURCE_TYPE_NAME,
    TO_SECONDS,
)
from datadog_checks.utils.containers import hash_mutable

class Couchbase(AgentCheck):
    """
    Extracts stats from Couchbase via its REST API
    http://docs.couchbase.com/couchbase-manual-2.0/#using-the-rest-api
    """

    HTTP_CONFIG_REMAPPER = {'user': {'name': 'username'}, 'ssl_verify': {'name': 'tls_verify'}}

    class CouchbaseInstanceState(object):
        def __init__(self):
            self.previous_status = None

    def __init__(self, name, init_config, instances):
        super(Couchbase, self).__init__(name, init_config, instances)

        # Keep track of all instances
        self._instance_states = defaultdict(lambda: self.CouchbaseInstanceState())

    def _create_metrics(self, data, instance_state, server, tags=None):
        # Get storage metrics
        storage_totals = data['stats']['storageTotals']
        for key, storage_type in storage_totals.items():
            for metric_name, val in storage_type.items():
                if val is not None:
                    metric_name = 'couchbase.{}.{}'.format(key, self.camel_case_to_joined_lower(metric_name))
                    self.gauge(metric_name, val, tags=tags)

        # Get bucket metrics
        for bucket_name, bucket_stats in data['buckets'].items():
            metric_tags = ['bucket:{}'.format(bucket_name), 'device:{}'.format(bucket_name)]
            if tags:
                metric_tags.extend(tags)
            for metric_name, val in bucket_stats.items():
                if val is not None:
                    norm_metric_name = self.camel_case_to_joined_lower(metric_name)
                    if norm_metric_name in BUCKET_STATS:
                        full_metric_name = 'couchbase.by_bucket.{}'.format(norm_metric_name)
                        self.gauge(full_metric_name, val[0], tags=metric_tags)

        # Get node metrics
        for node_name, node_stats in data['nodes'].items():
            metric_tags = ['node:{}'.format(node_name), 'device:{}'.format(node_name)]
            if tags:
                metric_tags.extend(tags)
            for metric_name, val in node_stats['interestingStats'].items():
                if val is not None:
                    metric_name = 'couchbase.by_node.{}'.format(self.camel_case_to_joined_lower(metric_name))
                    self.gauge(metric_name, val, tags=metric_tags)

            # Get cluster health data
            self._process_cluster_health_data(node_name, node_stats, tags)

        # Get query metrics
        for metric_name, val in data['query'].items():
            if val is not None:
                norm_metric_name = self.camel_case_to_joined_lower(metric_name)
                if norm_metric_name in QUERY_STATS:
                    # for query times, the unit is part of the value, we need to extract it
                    if isinstance(val, string_types):
                        val = self.extract_seconds_value(val)

                    full_metric_name = 'couchbase.query.{}'.format(self.camel_case_to_joined_lower(norm_metric_name))
                    self.gauge(full_metric_name, val, tags=tags)

        # Get tasks, we currently only care about 'rebalance' tasks
        rebalance_status, rebalance_msg = data['tasks'].get('rebalance', (None, None))

        # Only fire an event when the state has changed
        if rebalance_status is not None and instance_state.previous_status != rebalance_status:
            rebalance_event = None

            # If we get an error, we create an error event with the msg we receive
            if rebalance_status == 'error':
                msg_title = 'Encountered an error while rebalancing'
                msg = rebalance_msg

                rebalance_event = self._create_event('error', msg_title, msg, server, tags=tags)

            # We only want to fire a 'completion' of a rebalance so make sure we're not firing an event on first run
            elif rebalance_status == 'notRunning' and instance_state.previous_status is not None:
                msg_title = 'Stopped rebalancing'
                msg = 'stopped rebalancing.'
                rebalance_event = self._create_event('info', msg_title, msg, server, tags=tags)

            # If a rebalance task is running, fire an event. This will also fire an event if a rebalance task was
            #   already running when the check first runs.
            elif rebalance_status == 'gracefulFailover':
                msg_title = 'Failing over gracefully'
                msg = 'is failing over gracefully.'
                rebalance_event = self._create_event('info', msg_title, msg, server, tags=tags)
            elif rebalance_status == 'rebalance':
                msg_title = 'Rebalancing'
                msg = 'is rebalancing.'
                rebalance_event = self._create_event('info', msg_title, msg, server, tags=tags)

            # Send the event
            if rebalance_event is not None:
                self.event(rebalance_event)

            # Update the status of this instance
            instance_state.previous_status = rebalance_status

    def _process_cluster_health_data(self, node_name, node_stats, tags):
        """
        Process and send cluster health data (i.e. cluster membership status and node health
        """

        # Tags for service check
        cluster_health_tags = list(tags) + ['node:{}'.format(node_name)]

        # Get the membership status of the node
        cluster_membership = node_stats.get('clusterMembership', None)
        membership_status = NODE_MEMBERSHIP_TRANSLATION.get(cluster_membership, AgentCheck.UNKNOWN)
        self.service_check(NODE_CLUSTER_SERVICE_CHECK_NAME, membership_status, tags=cluster_health_tags)

        # Get the health status of the node
        health = node_stats.get('status', None)
        health_status = NODE_HEALTH_TRANSLATION.get(health, AgentCheck.UNKNOWN)
        self.service_check(NODE_HEALTH_SERVICE_CHECK_NAME, health_status, tags=cluster_health_tags)

    def _create_event(self, alert_type, msg_title, msg, server, tags=None):
        """
        Create an event object
        """
        msg_title = 'Couchbase {}: {}'.format(server, msg_title)
        msg = 'Couchbase instance {} {}'.format(server, msg)

        return {
            'timestamp': int(time.time()),
            'event_type': 'couchbase_rebalance',
            'msg_text': msg,
            'msg_title': msg_title,
            'alert_type': alert_type,
            'source_type_name': SOURCE_TYPE_NAME,
            'aggregation_key': server,
            'tags': tags,
        }

    def _get_stats(self, url):
        """
        Hit a given URL and return the parsed json.
        """
        r = self.http.get(url)
        r.raise_for_status()
        return r.json()

    def check(self, instance):
        instance_state = self._instance_states[hash_mutable(instance)]

        server = instance.get('server', None)
        if server is None:
            raise Exception("The server must be specified")
        tags = instance.get('tags', [])
        # Clean up tags in case there was a None entry in the instance
        # e.g. if the yaml contains tags: but no actual tags
        if tags is None:
            tags = []
        else:
            tags = list(set(tags))
        tags.append('instance:{}'.format(server))
        data = self.get_data(server, instance)
        self._create_metrics(data, instance_state, server, tags=list(set(tags)))

    def get_data(self, server, instance):
        # The dictionary to be returned.
        couchbase = {'stats': None, 'buckets': {}, 'nodes': {}, 'query': {}, 'tasks': {}}

        # build couchbase stats entry point
        url = '{}{}'.format(server, COUCHBASE_STATS_PATH)

        # Fetch initial stats and capture a service check based on response.
        service_check_tags = instance.get('tags', [])
        if service_check_tags is None:
            service_check_tags = []
        else:
            service_check_tags = list(set(service_check_tags))
        service_check_tags.append('instance:{}'.format(server))
        try:
            overall_stats = self._get_stats(url)
            # No overall stats? bail out now
            if overall_stats is None:
                raise Exception("No data returned from couchbase endpoint: {}".format(url))
        except requests.exceptions.HTTPError as e:
            self.service_check(SERVICE_CHECK_NAME, AgentCheck.CRITICAL, tags=service_check_tags, message=str(e))
            raise
        except Exception as e:
            self.service_check(SERVICE_CHECK_NAME, AgentCheck.CRITICAL, tags=service_check_tags, message=str(e))
            raise
        else:
            self.service_check(SERVICE_CHECK_NAME, AgentCheck.OK, tags=service_check_tags)

        couchbase['stats'] = overall_stats

        nodes = overall_stats['nodes']

        # Next, get all the nodes
        if nodes is not None:
            for node in nodes:
                couchbase['nodes'][node['hostname']] = node

        # Next, get all buckets .
        endpoint = overall_stats['buckets']['uri']

        url = '{}{}'.format(server, endpoint)
        buckets = self._get_stats(url)

        if buckets is not None:
            for bucket in buckets:
                bucket_name = bucket['name']

                # Fetch URI for the stats bucket
                endpoint = bucket['stats']['uri']
                url = '{}{}'.format(server, endpoint)

                try:
                    bucket_stats = self._get_stats(url)
                except requests.exceptions.HTTPError:
                    url_backup = '{}/pools/nodes/buckets/{}/stats'.format(server, bucket_name)
                    bucket_stats = self._get_stats(url_backup)

                bucket_samples = bucket_stats['op']['samples']
                if bucket_samples is not None:
                    couchbase['buckets'][bucket['name']] = bucket_samples

        # Next, get the query monitoring data
        query_data = self._get_query_monitoring_data(instance)
        if query_data is not None:
            couchbase['query'] = query_data

        # Next, get all the tasks
        tasks_url = '{}{}/tasks'.format(server, COUCHBASE_STATS_PATH)
        try:
            tasks = self._get_stats(tasks_url)
            for task in tasks:
                task_type = task['type']

                # We only care about rebalance statuses
                if task_type != 'rebalance':
                    continue

                # Format the status so it's easier to understand
                if 'errorMessage' in task:
                    couchbase['tasks'][task_type] = ('error', task['errorMessage'])
                elif task['status'] == 'notRunning':
                    couchbase['tasks'][task_type] = (task['status'], None)

                # If the status is 'running', we want to differentiate between a regular rebalance and a graceful
                #   failover rebalance, so we use the subtype
                elif task['status'] == 'running':
                    couchbase['tasks'][task_type] = (task['subtype'], None)

                # Should only be 1 rebalance
                break

        except requests.exceptions.HTTPError:
            self.log.error("Error accessing the endpoint %s", url)

        return couchbase

    def _get_query_monitoring_data(self, instance):
        query_data = None
        query_monitoring_url = instance.get('query_monitoring_url')
        if query_monitoring_url:
            url = '{}{}'.format(query_monitoring_url, COUCHBASE_VITALS_PATH)
            try:
                query_data = self._get_stats(url)
            except requests.exceptions.RequestException:
                self.log.error(
                    "Error accessing the endpoint %s, make sure you're running at least "
                    "couchbase 4.5 to collect the query monitoring metrics",
                    url,
                )

        return query_data

    # Takes a camelCased variable and returns a joined_lower equivalent.
    # Returns input if non-camelCase variable is detected.
    def camel_case_to_joined_lower(self, variable):
        # replace non-word with _
        converted_variable = re.sub(r'\W+', '_', variable)

        # insert _ in front of capital letters and lowercase the string
        converted_variable = re.sub('([A-Z])', r'_\g<1>', converted_variable).lower()

        # remove duplicate _
        converted_variable = re.sub('_+', '_', converted_variable)

        # handle special case of starting/ending underscores
        converted_variable = re.sub('^_|_$', '', converted_variable)

        return converted_variable

    # Takes a string with a time and a unit (e.g '3.45ms') and returns the value in seconds
    def extract_seconds_value(self, value):

        # When couchbase is set up, most of values are equal to 0 and are exposed as "0" and not "0s"
        # This statement is preventing values to be searched by the pattern (and break things)
        if value == '0':
            return 0

        match = SECONDS_VALUE_PATTERN.search(value)

        val, unit = match.group(1, 3)
        # They use the 'micro' symbol for microseconds so there is an encoding problem
        # so let's assume it's microseconds if we don't find the key in unit
        if unit not in TO_SECONDS:
            unit = 'us'

        return float(val) / TO_SECONDS[unit]
<|endoftext|>


<commit_before># coding=utf-8
# --------------------------------------------------------------------------
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License. See License.txt in the project root for
# license information.
#
# Code generated by Microsoft (R) AutoRest Code Generator.
# Changes may cause incorrect behavior and will be lost if the code is
# regenerated.
# --------------------------------------------------------------------------

from msrest.service_client import ServiceClient
from msrest import Serializer, Deserializer
from msrestazure import AzureConfiguration
from .version import VERSION
from msrest.pipeline import ClientRawResponse
from msrestazure.azure_exceptions import CloudError
from msrestazure.azure_operation import AzureOperationPoller
import uuid
from .operations.application_gateways_operations import ApplicationGatewaysOperations
from .operations.express_route_circuit_authorizations_operations import ExpressRouteCircuitAuthorizationsOperations
from .operations.express_route_circuit_peerings_operations import ExpressRouteCircuitPeeringsOperations
from .operations.express_route_circuits_operations import ExpressRouteCircuitsOperations
from .operations.express_route_service_providers_operations import ExpressRouteServiceProvidersOperations
from .operations.load_balancers_operations import LoadBalancersOperations
from .operations.network_interfaces_operations import NetworkInterfacesOperations
from .operations.network_security_groups_operations import NetworkSecurityGroupsOperations
from .operations.security_rules_operations import SecurityRulesOperations
from .operations.network_watchers_operations import NetworkWatchersOperations
from .operations.packet_captures_operations import PacketCapturesOperations
from .operations.public_ip_addresses_operations import PublicIPAddressesOperations
from .operations.route_filters_operations import RouteFiltersOperations
from .operations.route_filter_rules_operations import RouteFilterRulesOperations
from .operations.route_tables_operations import RouteTablesOperations
from .operations.routes_operations import RoutesOperations
from .operations.bgp_service_communities_operations import BgpServiceCommunitiesOperations
from .operations.usages_operations import UsagesOperations
from .operations.virtual_networks_operations import VirtualNetworksOperations
from .operations.subnets_operations import SubnetsOperations
from .operations.virtual_network_peerings_operations import VirtualNetworkPeeringsOperations
from .operations.virtual_network_gateways_operations import VirtualNetworkGatewaysOperations
from .operations.virtual_network_gateway_connections_operations import VirtualNetworkGatewayConnectionsOperations
from .operations.local_network_gateways_operations import LocalNetworkGatewaysOperations
from . import models

class NetworkManagementClientConfiguration(AzureConfiguration):
    """Configuration for NetworkManagementClient
    Note that all parameters used to create this instance are saved as instance
    attributes.

    :param credentials: Credentials needed for the client to connect to Azure.
    :type credentials: :mod:`A msrestazure Credentials
     object<msrestazure.azure_active_directory>`
    :param subscription_id: The subscription credentials which uniquely
     identify the Microsoft Azure subscription. The subscription ID forms part
     of the URI for every service call.
    :type subscription_id: str
    :param str base_url: Service URL
    """

    def __init__(
            self, credentials, subscription_id, base_url=None):

        if credentials is None:
            raise ValueError("Parameter 'credentials' must not be None.")
        if subscription_id is None:
            raise ValueError("Parameter 'subscription_id' must not be None.")
        if not isinstance(subscription_id, str):
            raise TypeError("Parameter 'subscription_id' must be str.")
        if not base_url:
            base_url = 'https://management.azure.com'

        super(NetworkManagementClientConfiguration, self).__init__(base_url)

        self.add_user_agent('networkmanagementclient/{}'.format(VERSION))
        self.add_user_agent('Azure-SDK-For-Python')

        self.credentials = credentials
        self.subscription_id = subscription_id

class NetworkManagementClient(object):
    """Composite Swagger for Network Client

    :ivar config: Configuration for client.
    :vartype config: NetworkManagementClientConfiguration

    :ivar application_gateways: ApplicationGateways operations
    :vartype application_gateways: .operations.ApplicationGatewaysOperations
    :ivar express_route_circuit_authorizations: ExpressRouteCircuitAuthorizations operations
    :vartype express_route_circuit_authorizations: .operations.ExpressRouteCircuitAuthorizationsOperations
    :ivar express_route_circuit_peerings: ExpressRouteCircuitPeerings operations
    :vartype express_route_circuit_peerings: .operations.ExpressRouteCircuitPeeringsOperations
    :ivar express_route_circuits: ExpressRouteCircuits operations
    :vartype express_route_circuits: .operations.ExpressRouteCircuitsOperations
    :ivar express_route_service_providers: ExpressRouteServiceProviders operations
    :vartype express_route_service_providers: .operations.ExpressRouteServiceProvidersOperations
    :ivar load_balancers: LoadBalancers operations
    :vartype load_balancers: .operations.LoadBalancersOperations
    :ivar network_interfaces: NetworkInterfaces operations
    :vartype network_interfaces: .operations.NetworkInterfacesOperations
    :ivar network_security_groups: NetworkSecurityGroups operations
    :vartype network_security_groups: .operations.NetworkSecurityGroupsOperations
    :ivar security_rules: SecurityRules operations
    :vartype security_rules: .operations.SecurityRulesOperations
    :ivar network_watchers: NetworkWatchers operations
    :vartype network_watchers: .operations.NetworkWatchersOperations
    :ivar packet_captures: PacketCaptures operations
    :vartype packet_captures: .operations.PacketCapturesOperations
    :ivar public_ip_addresses: PublicIPAddresses operations
    :vartype public_ip_addresses: .operations.PublicIPAddressesOperations
    :ivar route_filters: RouteFilters operations
    :vartype route_filters: .operations.RouteFiltersOperations
    :ivar route_filter_rules: RouteFilterRules operations
    :vartype route_filter_rules: .operations.RouteFilterRulesOperations
    :ivar route_tables: RouteTables operations
    :vartype route_tables: .operations.RouteTablesOperations
    :ivar routes: Routes operations
    :vartype routes: .operations.RoutesOperations
    :ivar bgp_service_communities: BgpServiceCommunities operations
    :vartype bgp_service_communities: .operations.BgpServiceCommunitiesOperations
    :ivar usages: Usages operations
    :vartype usages: .operations.UsagesOperations
    :ivar virtual_networks: VirtualNetworks operations
    :vartype virtual_networks: .operations.VirtualNetworksOperations
    :ivar subnets: Subnets operations
    :vartype subnets: .operations.SubnetsOperations
    :ivar virtual_network_peerings: VirtualNetworkPeerings operations
    :vartype virtual_network_peerings: .operations.VirtualNetworkPeeringsOperations
    :ivar virtual_network_gateways: VirtualNetworkGateways operations
    :vartype virtual_network_gateways: .operations.VirtualNetworkGatewaysOperations
    :ivar virtual_network_gateway_connections: VirtualNetworkGatewayConnections operations
    :vartype virtual_network_gateway_connections: .operations.VirtualNetworkGatewayConnectionsOperations
    :ivar local_network_gateways: LocalNetworkGateways operations
    :vartype local_network_gateways: .operations.LocalNetworkGatewaysOperations

    :param credentials: Credentials needed for the client to connect to Azure.
    :type credentials: :mod:`A msrestazure Credentials
     object<msrestazure.azure_active_directory>`
    :param subscription_id: The subscription credentials which uniquely
     identify the Microsoft Azure subscription. The subscription ID forms part
     of the URI for every service call.
    :type subscription_id: str
    :param str base_url: Service URL
    """

    def __init__(
            self, credentials, subscription_id, base_url=None):

        self.config = NetworkManagementClientConfiguration(credentials, subscription_id, base_url)
        self._client = ServiceClient(self.config.credentials, self.config)

        client_models = {k: v for k, v in models.__dict__.items() if isinstance(v, type)}
        self._serialize = Serializer(client_models)
        self._deserialize = Deserializer(client_models)

        self.application_gateways = ApplicationGatewaysOperations(
            self._client, self.config, self._serialize, self._deserialize)
        self.express_route_circuit_authorizations = ExpressRouteCircuitAuthorizationsOperations(
            self._client, self.config, self._serialize, self._deserialize)
        self.express_route_circuit_peerings = ExpressRouteCircuitPeeringsOperations(
            self._client, self.config, self._serialize, self._deserialize)
        self.express_route_circuits = ExpressRouteCircuitsOperations(
            self._client, self.config, self._serialize, self._deserialize)
        self.express_route_service_providers = ExpressRouteServiceProvidersOperations(
            self._client, self.config, self._serialize, self._deserialize)
        self.load_balancers = LoadBalancersOperations(
            self._client, self.config, self._serialize, self._deserialize)
        self.network_interfaces = NetworkInterfacesOperations(
            self._client, self.config, self._serialize, self._deserialize)
        self.network_security_groups = NetworkSecurityGroupsOperations(
            self._client, self.config, self._serialize, self._deserialize)
        self.security_rules = SecurityRulesOperations(
            self._client, self.config, self._serialize, self._deserialize)
        self.network_watchers = NetworkWatchersOperations(
            self._client, self.config, self._serialize, self._deserialize)
        self.packet_captures = PacketCapturesOperations(
            self._client, self.config, self._serialize, self._deserialize)
        self.public_ip_addresses = PublicIPAddressesOperations(
            self._client, self.config, self._serialize, self._deserialize)
        self.route_filters = RouteFiltersOperations(
            self._client, self.config, self._serialize, self._deserialize)
        self.route_filter_rules = RouteFilterRulesOperations(
            self._client, self.config, self._serialize, self._deserialize)
        self.route_tables = RouteTablesOperations(
            self._client, self.config, self._serialize, self._deserialize)
        self.routes = RoutesOperations(
            self._client, self.config, self._serialize, self._deserialize)
        self.bgp_service_communities = BgpServiceCommunitiesOperations(
            self._client, self.config, self._serialize, self._deserialize)
        self.usages = UsagesOperations(
            self._client, self.config, self._serialize, self._deserialize)
        self.virtual_networks = VirtualNetworksOperations(
            self._client, self.config, self._serialize, self._deserialize)
        self.subnets = SubnetsOperations(
            self._client, self.config, self._serialize, self._deserialize)
        self.virtual_network_peerings = VirtualNetworkPeeringsOperations(
            self._client, self.config, self._serialize, self._deserialize)
        self.virtual_network_gateways = VirtualNetworkGatewaysOperations(
            self._client, self.config, self._serialize, self._deserialize)
        self.virtual_network_gateway_connections = VirtualNetworkGatewayConnectionsOperations(
            self._client, self.config, self._serialize, self._deserialize)
        self.local_network_gateways = LocalNetworkGatewaysOperations(
            self._client, self.config, self._serialize, self._deserialize)

    def check_dns_name_availability(
            self, location, domain_name_label=None, custom_headers=None, raw=False, **operation_config):
        """Checks whether a domain name in the cloudapp.net zone is available for
        use.

        :param location: The location of the domain name.
        :type location: str
        :param domain_name_label: The domain name to be verified. It must
         conform to the following regular expression:
         ^[a-z][a-z0-9-]{1,61}[a-z0-9]$.
        :type domain_name_label: str
        :param dict custom_headers: headers that will be added to the request
        :param bool raw: returns the direct response alongside the
         deserialized response
        :param operation_config: :ref:`Operation configuration
         overrides<msrest:optionsforoperations>`.
        :rtype: :class:`DnsNameAvailabilityResult
         <azure.mgmt.network.v2017_03_01.models.DnsNameAvailabilityResult>`
        :rtype: :class:`ClientRawResponse<msrest.pipeline.ClientRawResponse>`
         if raw=true
        :raises: :class:`CloudError<msrestazure.azure_exceptions.CloudError>`
        """
        api_version = "2017-03-01"

        # Construct URL
        url = '/subscriptions/{subscriptionId}/providers/Microsoft.Network/locations/{location}/CheckDnsNameAvailability'
        path_format_arguments = {
            'location': self._serialize.url("location", location, 'str'),
            'subscriptionId': self._serialize.url("self.config.subscription_id", self.config.subscription_id, 'str')
        }
        url = self._client.format_url(url, **path_format_arguments)

        # Construct parameters
        query_parameters = {}
        if domain_name_label is not None:
            query_parameters['domainNameLabel'] = self._serialize.query("domain_name_label", domain_name_label, 'str')
        query_parameters['api-version'] = self._serialize.query("api_version", api_version, 'str')

        # Construct headers
        header_parameters = {}
        header_parameters['Content-Type'] = 'application/json; charset=utf-8'
        if self.config.generate_client_request_id:
            header_parameters['x-ms-client-request-id'] = str(uuid.uuid1())
        if custom_headers:
            header_parameters.update(custom_headers)
        if self.config.accept_language is not None:
            header_parameters['accept-language'] = self._serialize.header("self.config.accept_language", self.config.accept_language, 'str')

        # Construct and send request
        request = self._client.get(url, query_parameters)
        response = self._client.send(request, header_parameters, **operation_config)

        if response.status_code not in [200]:
            exp = CloudError(response)
            exp.request_id = response.headers.get('x-ms-request-id')
            raise exp

        deserialized = None

        if response.status_code == 200:
            deserialized = self._deserialize('DnsNameAvailabilityResult', response)

        if raw:
            client_raw_response = ClientRawResponse(deserialized, response)
            return client_raw_response

        return deserialized
<commit_msg>Generated from 2dc43333a1ba18ebf0213f9d3948172be17d3bc1<commit_after># coding=utf-8
# --------------------------------------------------------------------------
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License. See License.txt in the project root for
# license information.
#
# Code generated by Microsoft (R) AutoRest Code Generator.
# Changes may cause incorrect behavior and will be lost if the code is
# regenerated.
# --------------------------------------------------------------------------

from msrest.service_client import ServiceClient
from msrest import Serializer, Deserializer
from msrestazure import AzureConfiguration
from .version import VERSION
from msrest.pipeline import ClientRawResponse
from msrestazure.azure_exceptions import CloudError
from msrestazure.azure_operation import AzureOperationPoller
import uuid
from .operations.application_gateways_operations import ApplicationGatewaysOperations
from .operations.express_route_circuit_authorizations_operations import ExpressRouteCircuitAuthorizationsOperations
from .operations.express_route_circuit_peerings_operations import ExpressRouteCircuitPeeringsOperations
from .operations.express_route_circuits_operations import ExpressRouteCircuitsOperations
from .operations.express_route_service_providers_operations import ExpressRouteServiceProvidersOperations
from .operations.load_balancers_operations import LoadBalancersOperations
from .operations.network_interfaces_operations import NetworkInterfacesOperations
from .operations.network_security_groups_operations import NetworkSecurityGroupsOperations
from .operations.security_rules_operations import SecurityRulesOperations
from .operations.network_watchers_operations import NetworkWatchersOperations
from .operations.packet_captures_operations import PacketCapturesOperations
from .operations.public_ip_addresses_operations import PublicIPAddressesOperations
from .operations.route_filters_operations import RouteFiltersOperations
from .operations.route_filter_rules_operations import RouteFilterRulesOperations
from .operations.route_tables_operations import RouteTablesOperations
from .operations.routes_operations import RoutesOperations
from .operations.bgp_service_communities_operations import BgpServiceCommunitiesOperations
from .operations.usages_operations import UsagesOperations
from .operations.virtual_networks_operations import VirtualNetworksOperations
from .operations.subnets_operations import SubnetsOperations
from .operations.virtual_network_peerings_operations import VirtualNetworkPeeringsOperations
from .operations.virtual_network_gateways_operations import VirtualNetworkGatewaysOperations
from .operations.virtual_network_gateway_connections_operations import VirtualNetworkGatewayConnectionsOperations
from .operations.local_network_gateways_operations import LocalNetworkGatewaysOperations
from . import models

class NetworkManagementClientConfiguration(AzureConfiguration):
    """Configuration for NetworkManagementClient
    Note that all parameters used to create this instance are saved as instance
    attributes.

    :param credentials: Credentials needed for the client to connect to Azure.
    :type credentials: :mod:`A msrestazure Credentials
     object<msrestazure.azure_active_directory>`
    :param subscription_id: The subscription credentials which uniquely
     identify the Microsoft Azure subscription. The subscription ID forms part
     of the URI for every service call.
    :type subscription_id: str
    :param str base_url: Service URL
    """

    def __init__(
            self, credentials, subscription_id, base_url=None):

        if credentials is None:
            raise ValueError("Parameter 'credentials' must not be None.")
        if subscription_id is None:
            raise ValueError("Parameter 'subscription_id' must not be None.")
        if not isinstance(subscription_id, str):
            raise TypeError("Parameter 'subscription_id' must be str.")
        if not base_url:
            base_url = 'https://management.azure.com'

        super(NetworkManagementClientConfiguration, self).__init__(base_url)

        self.add_user_agent('networkmanagementclient/{}'.format(VERSION))
        self.add_user_agent('Azure-SDK-For-Python')

        self.credentials = credentials
        self.subscription_id = subscription_id

class NetworkManagementClient(object):
    """Composite Swagger for Network Client

    :ivar config: Configuration for client.
    :vartype config: NetworkManagementClientConfiguration

    :ivar application_gateways: ApplicationGateways operations
    :vartype application_gateways: azure.mgmt.network.v2017_03_01.operations.ApplicationGatewaysOperations
    :ivar express_route_circuit_authorizations: ExpressRouteCircuitAuthorizations operations
    :vartype express_route_circuit_authorizations: azure.mgmt.network.v2017_03_01.operations.ExpressRouteCircuitAuthorizationsOperations
    :ivar express_route_circuit_peerings: ExpressRouteCircuitPeerings operations
    :vartype express_route_circuit_peerings: azure.mgmt.network.v2017_03_01.operations.ExpressRouteCircuitPeeringsOperations
    :ivar express_route_circuits: ExpressRouteCircuits operations
    :vartype express_route_circuits: azure.mgmt.network.v2017_03_01.operations.ExpressRouteCircuitsOperations
    :ivar express_route_service_providers: ExpressRouteServiceProviders operations
    :vartype express_route_service_providers: azure.mgmt.network.v2017_03_01.operations.ExpressRouteServiceProvidersOperations
    :ivar load_balancers: LoadBalancers operations
    :vartype load_balancers: azure.mgmt.network.v2017_03_01.operations.LoadBalancersOperations
    :ivar network_interfaces: NetworkInterfaces operations
    :vartype network_interfaces: azure.mgmt.network.v2017_03_01.operations.NetworkInterfacesOperations
    :ivar network_security_groups: NetworkSecurityGroups operations
    :vartype network_security_groups: azure.mgmt.network.v2017_03_01.operations.NetworkSecurityGroupsOperations
    :ivar security_rules: SecurityRules operations
    :vartype security_rules: azure.mgmt.network.v2017_03_01.operations.SecurityRulesOperations
    :ivar network_watchers: NetworkWatchers operations
    :vartype network_watchers: azure.mgmt.network.v2017_03_01.operations.NetworkWatchersOperations
    :ivar packet_captures: PacketCaptures operations
    :vartype packet_captures: azure.mgmt.network.v2017_03_01.operations.PacketCapturesOperations
    :ivar public_ip_addresses: PublicIPAddresses operations
    :vartype public_ip_addresses: azure.mgmt.network.v2017_03_01.operations.PublicIPAddressesOperations
    :ivar route_filters: RouteFilters operations
    :vartype route_filters: azure.mgmt.network.v2017_03_01.operations.RouteFiltersOperations
    :ivar route_filter_rules: RouteFilterRules operations
    :vartype route_filter_rules: azure.mgmt.network.v2017_03_01.operations.RouteFilterRulesOperations
    :ivar route_tables: RouteTables operations
    :vartype route_tables: azure.mgmt.network.v2017_03_01.operations.RouteTablesOperations
    :ivar routes: Routes operations
    :vartype routes: azure.mgmt.network.v2017_03_01.operations.RoutesOperations
    :ivar bgp_service_communities: BgpServiceCommunities operations
    :vartype bgp_service_communities: azure.mgmt.network.v2017_03_01.operations.BgpServiceCommunitiesOperations
    :ivar usages: Usages operations
    :vartype usages: azure.mgmt.network.v2017_03_01.operations.UsagesOperations
    :ivar virtual_networks: VirtualNetworks operations
    :vartype virtual_networks: azure.mgmt.network.v2017_03_01.operations.VirtualNetworksOperations
    :ivar subnets: Subnets operations
    :vartype subnets: azure.mgmt.network.v2017_03_01.operations.SubnetsOperations
    :ivar virtual_network_peerings: VirtualNetworkPeerings operations
    :vartype virtual_network_peerings: azure.mgmt.network.v2017_03_01.operations.VirtualNetworkPeeringsOperations
    :ivar virtual_network_gateways: VirtualNetworkGateways operations
    :vartype virtual_network_gateways: azure.mgmt.network.v2017_03_01.operations.VirtualNetworkGatewaysOperations
    :ivar virtual_network_gateway_connections: VirtualNetworkGatewayConnections operations
    :vartype virtual_network_gateway_connections: azure.mgmt.network.v2017_03_01.operations.VirtualNetworkGatewayConnectionsOperations
    :ivar local_network_gateways: LocalNetworkGateways operations
    :vartype local_network_gateways: azure.mgmt.network.v2017_03_01.operations.LocalNetworkGatewaysOperations

    :param credentials: Credentials needed for the client to connect to Azure.
    :type credentials: :mod:`A msrestazure Credentials
     object<msrestazure.azure_active_directory>`
    :param subscription_id: The subscription credentials which uniquely
     identify the Microsoft Azure subscription. The subscription ID forms part
     of the URI for every service call.
    :type subscription_id: str
    :param str base_url: Service URL
    """

    def __init__(
            self, credentials, subscription_id, base_url=None):

        self.config = NetworkManagementClientConfiguration(credentials, subscription_id, base_url)
        self._client = ServiceClient(self.config.credentials, self.config)

        client_models = {k: v for k, v in models.__dict__.items() if isinstance(v, type)}
        self._serialize = Serializer(client_models)
        self._deserialize = Deserializer(client_models)

        self.application_gateways = ApplicationGatewaysOperations(
            self._client, self.config, self._serialize, self._deserialize)
        self.express_route_circuit_authorizations = ExpressRouteCircuitAuthorizationsOperations(
            self._client, self.config, self._serialize, self._deserialize)
        self.express_route_circuit_peerings = ExpressRouteCircuitPeeringsOperations(
            self._client, self.config, self._serialize, self._deserialize)
        self.express_route_circuits = ExpressRouteCircuitsOperations(
            self._client, self.config, self._serialize, self._deserialize)
        self.express_route_service_providers = ExpressRouteServiceProvidersOperations(
            self._client, self.config, self._serialize, self._deserialize)
        self.load_balancers = LoadBalancersOperations(
            self._client, self.config, self._serialize, self._deserialize)
        self.network_interfaces = NetworkInterfacesOperations(
            self._client, self.config, self._serialize, self._deserialize)
        self.network_security_groups = NetworkSecurityGroupsOperations(
            self._client, self.config, self._serialize, self._deserialize)
        self.security_rules = SecurityRulesOperations(
            self._client, self.config, self._serialize, self._deserialize)
        self.network_watchers = NetworkWatchersOperations(
            self._client, self.config, self._serialize, self._deserialize)
        self.packet_captures = PacketCapturesOperations(
            self._client, self.config, self._serialize, self._deserialize)
        self.public_ip_addresses = PublicIPAddressesOperations(
            self._client, self.config, self._serialize, self._deserialize)
        self.route_filters = RouteFiltersOperations(
            self._client, self.config, self._serialize, self._deserialize)
        self.route_filter_rules = RouteFilterRulesOperations(
            self._client, self.config, self._serialize, self._deserialize)
        self.route_tables = RouteTablesOperations(
            self._client, self.config, self._serialize, self._deserialize)
        self.routes = RoutesOperations(
            self._client, self.config, self._serialize, self._deserialize)
        self.bgp_service_communities = BgpServiceCommunitiesOperations(
            self._client, self.config, self._serialize, self._deserialize)
        self.usages = UsagesOperations(
            self._client, self.config, self._serialize, self._deserialize)
        self.virtual_networks = VirtualNetworksOperations(
            self._client, self.config, self._serialize, self._deserialize)
        self.subnets = SubnetsOperations(
            self._client, self.config, self._serialize, self._deserialize)
        self.virtual_network_peerings = VirtualNetworkPeeringsOperations(
            self._client, self.config, self._serialize, self._deserialize)
        self.virtual_network_gateways = VirtualNetworkGatewaysOperations(
            self._client, self.config, self._serialize, self._deserialize)
        self.virtual_network_gateway_connections = VirtualNetworkGatewayConnectionsOperations(
            self._client, self.config, self._serialize, self._deserialize)
        self.local_network_gateways = LocalNetworkGatewaysOperations(
            self._client, self.config, self._serialize, self._deserialize)

    def check_dns_name_availability(
            self, location, domain_name_label=None, custom_headers=None, raw=False, **operation_config):
        """Checks whether a domain name in the cloudapp.net zone is available for
        use.

        :param location: The location of the domain name.
        :type location: str
        :param domain_name_label: The domain name to be verified. It must
         conform to the following regular expression:
         ^[a-z][a-z0-9-]{1,61}[a-z0-9]$.
        :type domain_name_label: str
        :param dict custom_headers: headers that will be added to the request
        :param bool raw: returns the direct response alongside the
         deserialized response
        :param operation_config: :ref:`Operation configuration
         overrides<msrest:optionsforoperations>`.
        :rtype: :class:`DnsNameAvailabilityResult
         <azure.mgmt.network.v2017_03_01.models.DnsNameAvailabilityResult>`
        :rtype: :class:`ClientRawResponse<msrest.pipeline.ClientRawResponse>`
         if raw=true
        :raises: :class:`CloudError<msrestazure.azure_exceptions.CloudError>`
        """
        api_version = "2017-03-01"

        # Construct URL
        url = '/subscriptions/{subscriptionId}/providers/Microsoft.Network/locations/{location}/CheckDnsNameAvailability'
        path_format_arguments = {
            'location': self._serialize.url("location", location, 'str'),
            'subscriptionId': self._serialize.url("self.config.subscription_id", self.config.subscription_id, 'str')
        }
        url = self._client.format_url(url, **path_format_arguments)

        # Construct parameters
        query_parameters = {}
        if domain_name_label is not None:
            query_parameters['domainNameLabel'] = self._serialize.query("domain_name_label", domain_name_label, 'str')
        query_parameters['api-version'] = self._serialize.query("api_version", api_version, 'str')

        # Construct headers
        header_parameters = {}
        header_parameters['Content-Type'] = 'application/json; charset=utf-8'
        if self.config.generate_client_request_id:
            header_parameters['x-ms-client-request-id'] = str(uuid.uuid1())
        if custom_headers:
            header_parameters.update(custom_headers)
        if self.config.accept_language is not None:
            header_parameters['accept-language'] = self._serialize.header("self.config.accept_language", self.config.accept_language, 'str')

        # Construct and send request
        request = self._client.get(url, query_parameters)
        response = self._client.send(request, header_parameters, **operation_config)

        if response.status_code not in [200]:
            exp = CloudError(response)
            exp.request_id = response.headers.get('x-ms-request-id')
            raise exp

        deserialized = None

        if response.status_code == 200:
            deserialized = self._deserialize('DnsNameAvailabilityResult', response)

        if raw:
            client_raw_response = ClientRawResponse(deserialized, response)
            return client_raw_response

        return deserialized
<|endoftext|>