<commit_before>from dateutil.parser import parse
from django.conf import settings
import requests
from requests.exceptions import HTTPError
from tracking.models import DepartmentUser
from tracking.utils import logger_setup

HEADERS_JSON = {'Content-Type': 'application/json'}

def get_freshdesk_object(obj_type, id):
    """Query the Freshdesk v2 API for a single object.
    """
    url = settings.FRESHDESK_ENDPOINT + '/{}/{}'.format(obj_type, id)
    r = requests.get(url, auth=settings.FRESHDESK_AUTH)
    if not r.status_code == 200:
        r.raise_for_status()
    return r.json()

def update_freshdesk_object(obj_type, data, id=None):
    """Use the Freshdesk v2 API to create or update an object.
    Accepts a dict of fields.
    Ref: https://developer.freshdesk.com/api/#create_contact
    """
    if not id:  # Assume creation of new object.
        url = settings.FRESHDESK_ENDPOINT + '/{}'.format(obj_type)
        r = requests.post(url, auth=settings.FRESHDESK_AUTH, data=data)
    else:
        url = settings.FRESHDESK_ENDPOINT + '/{}/{}'.format(obj_type, id)
        r = requests.put(url, auth=settings.FRESHDESK_AUTH, data=data)
    return r  # Return the response, so we can handle non-200 gracefully.

def get_freshdesk_objects(obj_type, progress=True, limit=False, params={}):
    """Query the Freshdesk v2 API for all objects of a defined type.
    ``limit`` should be an integer (maximum number of objects to return).
    May take some time, depending on the number of objects.
    """
    url = settings.FRESHDESK_ENDPOINT + '/{}'.format(obj_type)
    if 'page' not in params:
        params['page'] = 1
    if 'per_page' not in params:
        params['per_page'] = 100
    objects = []
    further_results = True

    while further_results:
        if progress:
            print('Querying page {}'.format(params['page']))

        r = requests.get(url, auth=settings.FRESHDESK_AUTH, params=params)
        if not r.status_code == 200:
            r.raise_for_status()

        if 'link' not in r.headers:  # No further paginated results.
            further_results = False
            if progress:
                print('Done!')

        objects.extend(r.json())
        params['page'] += 1

        if limit and len(objects) >= limit:
            further_results = False
            objects = objects[:limit]

    # Return the full list of objects returned.
    return objects

def freshdesk_sync_contacts(contacts=None, companies=None, agents=None):
    """Iterate through all DepartmentUser objects, and ensure that each user's
    information is synced correctly to a Freshdesk contact.
    May optionally be passed in dicts of contacts & companies.
    """
    logger = logger_setup('freshdesk_sync_contacts')

    try:
        if not contacts:
            logger.info('Querying Freshdesk for current contacts')
            contacts = get_freshdesk_objects(obj_type='contacts', progress=False)
            contacts = {c['email'].lower(): c for c in contacts if c['email']}
        if not companies:
            logger.info('Querying Freshdesk for current companies')
            companies = get_freshdesk_objects(obj_type='companies', progress=False)
            companies = {c['name']: c for c in companies}
        if not agents:
            logger.info('Querying Freshdesk for current agents')
            agents = get_freshdesk_objects(obj_type='agents', progress=False)
            agents = {a['contact']['email'].lower(): a['contact'] for a in agents if a['contact']['email']}
    except Exception as e:
        logger.exception(e)
        return False

    # Filter DepartmentUsers: valid email (contains @), not -admin, DN contains 'OU=Users'
    d_users = DepartmentUser.objects.filter(email__contains='@', ad_dn__contains='OU=Users').exclude(email__contains='-admin')
    logger.info('Syncing details for {} DepartmentUsers to Freshdesk'.format(d_users.count()))
    for user in d_users:
        if user.email.lower() in contacts:
            # The DepartmentUser exists in Freshdesk; verify and update details.
            fd = contacts[user.email.lower()]
            data = {}
            user_sync = False
            # use extra attributes from org_data, if available
            cost_centre = user.org_data.get('cost_centre', {}).get('code', "") if user.org_data else ""
            physical_location = user.org_data.get('location', {}).get('name', "") if user.org_data else ""
            department = user.org_data.get('units', []) if user.org_data else []
            department = department[0].get('name', "") if len(department) > 0 else ""
            changes = []

            if user.name != fd['name']:
                user_sync = True
                data['name'] = user.name
                changes.append('name')
            if user.telephone != fd['phone']:
                user_sync = True
                data['phone'] = user.telephone
                changes.append('phone')
            if user.title != fd['job_title']:
                user_sync = True
                data['job_title'] = user.title
                changes.append('job_title')
            if department and department in companies and fd['company_id'] != companies[department]['id']:
                user_sync = True
                data['company_id'] = companies[department]['id']
                changes.append('company_id')
            # Custom fields in Freshdesk: Cost Centre no.
            if 'custom_field' in fd:  # Field may not exist in the API obj.
                if fd['custom_field']['cf_cost_centre'] != cost_centre:
                    user_sync = True
                    data['custom_field'] = {'cf_cost_centre': cost_centre}
                    changes.append('cost_centre')
                # Custom fields in Freshdesk: Physical location
                if fd['custom_field']['cf_location'] != physical_location:
                    user_sync = True
                    if 'custom_field' in data:
                        data['custom_field']['cf_location'] = physical_location
                    else:
                        data['custom_field'] = {'cf_location': physical_location}
                    changes.append('physical_location')
            if user_sync:  # Sync user details to their Freshdesk contact.
                r = update_freshdesk_object('contacts', data, fd['id'])
                if r.status_code == 403:  # Forbidden
                    # A 403 response probably means that we hit the API throttle limit.
                    # Abort the synchronisation.
                    logger.error('HTTP403 received from Freshdesk API, aborting')
                    return False
                logger.info('{} was updated in Freshdesk (status {}), changed: {}'.format(
                    user.email.lower(), r.status_code, ', '.join(changes)))
            else:
                logger.info('{} already up to date in Freshdesk'.format(user.email.lower()))
        elif user.email.lower() in agents:
            # The DepartmentUser is an agent; skip (can't update Agent details via the API).
            logger.info('{} is an agent, skipping sync'.format(user.email.lower()))
            continue
        else:
            # The DepartmentUser does not exist in Freshdesk; create them as a Contact.
            data = {'name': user.name, 'email': user.email.lower(),
                    'phone': user.telephone, 'job_title': user.title}
            if department and department in companies:
                data['company_id'] = companies[department]['id']
            r = update_freshdesk_object('contacts', data)
            if not r.status_code == 200:  # Error, unable to process request.
                logger.warn('{} not created in Freshdesk (status {})'.format(user.email.lower(), r.status_code))
            else:
                logger.info('{} created in Freshdesk (status {})'.format(user.email.lower(), r.status_code))

    return True

def freshdesk_cache_agents():
    """Cache a list of Freshdesk agents as contacts, as the API treats Agents
    differently to Contacts.
    """
    from registers.models import FreshdeskContact
    logger = logger_setup('freshdesk_cache_agents')
    agents = get_freshdesk_objects(obj_type='agents', progress=False)
    for i in agents:
        data = i['contact']
        data['contact_id'] = i['id']
        data['created_at'] = parse(data['created_at'])
        data['updated_at'] = parse(data['updated_at'])
        data.pop('last_login_at', None)
        fc, create = FreshdeskContact.objects.update_or_create(contact_id=data['contact_id'], defaults=data)
        if create:
            logger.info('{} created'.format(fc))
        else:
            logger.info('{} updated'.format(fc))
        # Attempt to match with a DepartmentUser.
        if DepartmentUser.objects.filter(email__iexact=fc.email).exists():
            fc.du_user = DepartmentUser.objects.get(email__iexact=fc.email)
            fc.save()

def freshdesk_cache_tickets(tickets=None):
    """Cache passed-in list of Freshdesk tickets in the database. If no tickets
    are passed in, query the API for the newest tickets.
    """
    from registers.models import FreshdeskTicket, FreshdeskConversation, FreshdeskContact
    logger = logger_setup('freshdesk_cache_tickets')

    if not tickets:
        try:
            logger.info('Querying Freshdesk for current tickets')
            tickets = get_freshdesk_objects(obj_type='tickets', progress=False, limit=30)
        except Exception as e:
            logger.exception(e)
            return False

    # Tweak the passed-in list of ticket values, prior to caching.
    for t in tickets:
        # Rename key 'id'.
        t['ticket_id'] = t.pop('id')
        # Parse ISO8601-formatted date strings into datetime objs.
        t['created_at'] = parse(t['created_at'])
        t['due_by'] = parse(t['due_by'])
        t['fr_due_by'] = parse(t['fr_due_by'])
        t['updated_at'] = parse(t['updated_at'])
        # Pop unused fields from the dict.
        t.pop('company_id', None)
        t.pop('email_config_id', None)
        t.pop('product_id', None)

    created, updated = 0, 0
    # Iterate through tickets; determine if a cached FreshdeskTicket should be
    # created or updated.
    for t in tickets:
        try:
            ft, create = FreshdeskTicket.objects.update_or_create(ticket_id=t['ticket_id'], defaults=t)
            if create:
                logger.info('{} created'.format(ft))
                created += 1
            else:
                logger.info('{} updated'.format(ft))
                updated += 1
            # Sync contact objects (requester and responder).
            # Check local cache first, to reduce the no. of API calls.
            if ft.requester_id:
                if FreshdeskContact.objects.filter(contact_id=ft.requester_id).exists():
                    ft.freshdesk_requester = FreshdeskContact.objects.get(contact_id=ft.requester_id)
                else:  # Attempt to cache the Freshdesk contact.
                    try:
                        c = get_freshdesk_object(obj_type='contacts', id=ft.requester_id)
                        c['contact_id'] = c.pop('id')
                        c['created_at'] = parse(c['created_at'])
                        c['updated_at'] = parse(c['updated_at'])
                        c.pop('avatar', None)
                        c.pop('company_id', None)
                        c.pop('twitter_id', None)
                        c.pop('deleted', None)
                        con = FreshdeskContact.objects.create(**c)
                        logger.info('Created {}'.format(con))
                        ft.freshdesk_requester = con
                    except HTTPError:  # The GET might fail if the contact is an agent.
                        logger.error('HTTP 404 Freshdesk contact not found: {}'.format(ft.requester_id))
                        pass
            if ft.responder_id:
                if FreshdeskContact.objects.filter(contact_id=ft.responder_id).exists():
                    ft.freshdesk_responder = FreshdeskContact.objects.get(contact_id=ft.responder_id)
                else:  # Attempt to cache the Freshdesk contact.
                    try:
                        c = get_freshdesk_object(obj_type='contacts', id=ft.responder_id)
                        c['contact_id'] = c.pop('id')
                        c['created_at'] = parse(c['created_at'])
                        c['updated_at'] = parse(c['updated_at'])
                        c.pop('avatar', None)
                        c.pop('company_id', None)
                        c.pop('twitter_id', None)
                        c.pop('deleted', None)
                        con = FreshdeskContact.objects.create(**c)
                        logger.info('Created {}'.format(con))
                        ft.freshdesk_responder = con
                    except HTTPError:  # The GET might fail if the contact is an agent.
                        logger.error('HTTP 404 Freshdesk contact not found: {}'.format(ft.responder_id))
                        pass
            ft.save()

            # Sync ticket conversation objects.
            obj = 'tickets/{}/conversations'.format(t['ticket_id'])
            convs = get_freshdesk_objects(obj_type=obj, progress=False)
            for c in convs:
                # Rename key 'id'.
                c['conversation_id'] = c.pop('id')
                # Date ISO8601-formatted date strings into datetimes.
                c['created_at'] = parse(c['created_at'])
                c['updated_at'] = parse(c['updated_at'])
                # Pop unused fields from the dict.
                c.pop('bcc_emails', None)
                c.pop('support_email', None)
                fc, create = FreshdeskConversation.objects.update_or_create(conversation_id=c['conversation_id'], defaults=c)
                if create:
                    logger.info('{} created'.format(fc))
                else:
                    logger.info('{} updated'.format(fc))
                # Link parent ticket, DepartmentUser, etc.
                fc.freshdesk_ticket = ft
                if FreshdeskContact.objects.filter(contact_id=fc.user_id).exists():
                    fc.freshdesk_contact = FreshdeskContact.objects.get(contact_id=fc.user_id)
                else:  # Attempt to cache the Freshdesk contact.
                    try:
                        f_con = get_freshdesk_object(obj_type='contacts', id=fc.user_id)
                        f_con['contact_id'] = f_con.pop('id')
                        f_con['created_at'] = parse(f_con['created_at'])
                        f_con['updated_at'] = parse(f_con['updated_at'])
                        f_con.pop('avatar', None)
                        f_con.pop('company_id', None)
                        f_con.pop('twitter_id', None)
                        f_con.pop('deleted', None)
                        contact = FreshdeskContact.objects.create(**f_con)
                        logger.info('Created {}'.format(contact))
                        fc.freshdesk_contact = contact
                        # Attempt to match contact with a DepartmentUser.
                        if DepartmentUser.objects.filter(email__iexact=contact.email).exists():
                            contact.du_user = DepartmentUser.objects.get(email__iexact=contact.email)
                            contact.save()
                    except HTTPError:  # The GET might fail if the contact is an agent.
                        logger.error('HTTP 404 Freshdesk contact not found: {}'.format(ft.requester_id))
                        pass
                # Attempt to match conversation with a DepartmentUser.
                if fc.freshdesk_contact and DepartmentUser.objects.filter(email__iexact=fc.freshdesk_contact.email).exists():
                    fc.du_user = DepartmentUser.objects.get(email__iexact=fc.freshdesk_contact.email)
                fc.save()
        except Exception as e:
            logger.exception(e)
            return False

    logger.info('Ticket sync: {} created, {} updated'.format(created, updated))
    print('{} created, {} updated'.format(created, updated))

    return True
<commit_msg>Added match_it_system() call to Freshdesk ticket sync command.<commit_after>from dateutil.parser import parse
from django.conf import settings
import requests
from requests.exceptions import HTTPError
from tracking.models import DepartmentUser
from tracking.utils import logger_setup

HEADERS_JSON = {'Content-Type': 'application/json'}

def get_freshdesk_object(obj_type, id):
    """Query the Freshdesk v2 API for a single object.
    """
    url = settings.FRESHDESK_ENDPOINT + '/{}/{}'.format(obj_type, id)
    r = requests.get(url, auth=settings.FRESHDESK_AUTH)
    if not r.status_code == 200:
        r.raise_for_status()
    return r.json()

def update_freshdesk_object(obj_type, data, id=None):
    """Use the Freshdesk v2 API to create or update an object.
    Accepts a dict of fields.
    Ref: https://developer.freshdesk.com/api/#create_contact
    """
    if not id:  # Assume creation of new object.
        url = settings.FRESHDESK_ENDPOINT + '/{}'.format(obj_type)
        r = requests.post(url, auth=settings.FRESHDESK_AUTH, data=data)
    else:
        url = settings.FRESHDESK_ENDPOINT + '/{}/{}'.format(obj_type, id)
        r = requests.put(url, auth=settings.FRESHDESK_AUTH, data=data)
    return r  # Return the response, so we can handle non-200 gracefully.

def get_freshdesk_objects(obj_type, progress=True, limit=False, params={}):
    """Query the Freshdesk v2 API for all objects of a defined type.
    ``limit`` should be an integer (maximum number of objects to return).
    May take some time, depending on the number of objects.
    """
    url = settings.FRESHDESK_ENDPOINT + '/{}'.format(obj_type)
    if 'page' not in params:
        params['page'] = 1
    if 'per_page' not in params:
        params['per_page'] = 100
    objects = []
    further_results = True

    while further_results:
        if progress:
            print('Querying page {}'.format(params['page']))

        r = requests.get(url, auth=settings.FRESHDESK_AUTH, params=params)
        if not r.status_code == 200:
            r.raise_for_status()

        if 'link' not in r.headers:  # No further paginated results.
            further_results = False
            if progress:
                print('Done!')

        objects.extend(r.json())
        params['page'] += 1

        if limit and len(objects) >= limit:
            further_results = False
            objects = objects[:limit]

    # Return the full list of objects returned.
    return objects

def freshdesk_sync_contacts(contacts=None, companies=None, agents=None):
    """Iterate through all DepartmentUser objects, and ensure that each user's
    information is synced correctly to a Freshdesk contact.
    May optionally be passed in dicts of contacts & companies.
    """
    logger = logger_setup('freshdesk_sync_contacts')

    try:
        if not contacts:
            logger.info('Querying Freshdesk for current contacts')
            contacts = get_freshdesk_objects(obj_type='contacts', progress=False)
            contacts = {c['email'].lower(): c for c in contacts if c['email']}
        if not companies:
            logger.info('Querying Freshdesk for current companies')
            companies = get_freshdesk_objects(obj_type='companies', progress=False)
            companies = {c['name']: c for c in companies}
        if not agents:
            logger.info('Querying Freshdesk for current agents')
            agents = get_freshdesk_objects(obj_type='agents', progress=False)
            agents = {a['contact']['email'].lower(): a['contact'] for a in agents if a['contact']['email']}
    except Exception as e:
        logger.exception(e)
        return False

    # Filter DepartmentUsers: valid email (contains @), not -admin, DN contains 'OU=Users'
    d_users = DepartmentUser.objects.filter(email__contains='@', ad_dn__contains='OU=Users').exclude(email__contains='-admin')
    logger.info('Syncing details for {} DepartmentUsers to Freshdesk'.format(d_users.count()))
    for user in d_users:
        if user.email.lower() in contacts:
            # The DepartmentUser exists in Freshdesk; verify and update details.
            fd = contacts[user.email.lower()]
            data = {}
            user_sync = False
            # use extra attributes from org_data, if available
            cost_centre = user.org_data.get('cost_centre', {}).get('code', "") if user.org_data else ""
            physical_location = user.org_data.get('location', {}).get('name', "") if user.org_data else ""
            department = user.org_data.get('units', []) if user.org_data else []
            department = department[0].get('name', "") if len(department) > 0 else ""
            changes = []

            if user.name != fd['name']:
                user_sync = True
                data['name'] = user.name
                changes.append('name')
            if user.telephone != fd['phone']:
                user_sync = True
                data['phone'] = user.telephone
                changes.append('phone')
            if user.title != fd['job_title']:
                user_sync = True
                data['job_title'] = user.title
                changes.append('job_title')
            if department and department in companies and fd['company_id'] != companies[department]['id']:
                user_sync = True
                data['company_id'] = companies[department]['id']
                changes.append('company_id')
            # Custom fields in Freshdesk: Cost Centre no.
            if 'custom_field' in fd:  # Field may not exist in the API obj.
                if fd['custom_field']['cf_cost_centre'] != cost_centre:
                    user_sync = True
                    data['custom_field'] = {'cf_cost_centre': cost_centre}
                    changes.append('cost_centre')
                # Custom fields in Freshdesk: Physical location
                if fd['custom_field']['cf_location'] != physical_location:
                    user_sync = True
                    if 'custom_field' in data:
                        data['custom_field']['cf_location'] = physical_location
                    else:
                        data['custom_field'] = {'cf_location': physical_location}
                    changes.append('physical_location')
            if user_sync:  # Sync user details to their Freshdesk contact.
                r = update_freshdesk_object('contacts', data, fd['id'])
                if r.status_code == 403:  # Forbidden
                    # A 403 response probably means that we hit the API throttle limit.
                    # Abort the synchronisation.
                    logger.error('HTTP403 received from Freshdesk API, aborting')
                    return False
                logger.info('{} was updated in Freshdesk (status {}), changed: {}'.format(
                    user.email.lower(), r.status_code, ', '.join(changes)))
            else:
                logger.info('{} already up to date in Freshdesk'.format(user.email.lower()))
        elif user.email.lower() in agents:
            # The DepartmentUser is an agent; skip (can't update Agent details via the API).
            logger.info('{} is an agent, skipping sync'.format(user.email.lower()))
            continue
        else:
            # The DepartmentUser does not exist in Freshdesk; create them as a Contact.
            data = {'name': user.name, 'email': user.email.lower(),
                    'phone': user.telephone, 'job_title': user.title}
            if department and department in companies:
                data['company_id'] = companies[department]['id']
            r = update_freshdesk_object('contacts', data)
            if not r.status_code == 200:  # Error, unable to process request.
                logger.warn('{} not created in Freshdesk (status {})'.format(user.email.lower(), r.status_code))
            else:
                logger.info('{} created in Freshdesk (status {})'.format(user.email.lower(), r.status_code))

    return True

def freshdesk_cache_agents():
    """Cache a list of Freshdesk agents as contacts, as the API treats Agents
    differently to Contacts.
    """
    from registers.models import FreshdeskContact
    logger = logger_setup('freshdesk_cache_agents')
    agents = get_freshdesk_objects(obj_type='agents', progress=False)
    for i in agents:
        data = i['contact']
        data['contact_id'] = i['id']
        data['created_at'] = parse(data['created_at'])
        data['updated_at'] = parse(data['updated_at'])
        data.pop('last_login_at', None)
        fc, create = FreshdeskContact.objects.update_or_create(contact_id=data['contact_id'], defaults=data)
        if create:
            logger.info('{} created'.format(fc))
        else:
            logger.info('{} updated'.format(fc))
        # Attempt to match with a DepartmentUser.
        if DepartmentUser.objects.filter(email__iexact=fc.email).exists():
            fc.du_user = DepartmentUser.objects.get(email__iexact=fc.email)
            fc.save()

def freshdesk_cache_tickets(tickets=None):
    """Cache passed-in list of Freshdesk tickets in the database. If no tickets
    are passed in, query the API for the newest tickets.
    """
    from registers.models import FreshdeskTicket, FreshdeskConversation, FreshdeskContact
    logger = logger_setup('freshdesk_cache_tickets')

    if not tickets:
        try:
            logger.info('Querying Freshdesk for current tickets')
            tickets = get_freshdesk_objects(obj_type='tickets', progress=False, limit=30)
        except Exception as e:
            logger.exception(e)
            return False

    # Tweak the passed-in list of ticket values, prior to caching.
    for t in tickets:
        # Rename key 'id'.
        t['ticket_id'] = t.pop('id')
        # Parse ISO8601-formatted date strings into datetime objs.
        t['created_at'] = parse(t['created_at'])
        t['due_by'] = parse(t['due_by'])
        t['fr_due_by'] = parse(t['fr_due_by'])
        t['updated_at'] = parse(t['updated_at'])
        # Pop unused fields from the dict.
        t.pop('company_id', None)
        t.pop('email_config_id', None)
        t.pop('product_id', None)

    created, updated = 0, 0
    # Iterate through tickets; determine if a cached FreshdeskTicket should be
    # created or updated.
    for t in tickets:
        try:
            ft, create = FreshdeskTicket.objects.update_or_create(ticket_id=t['ticket_id'], defaults=t)
            if create:
                logger.info('{} created'.format(ft))
                created += 1
            else:
                logger.info('{} updated'.format(ft))
                updated += 1
            # Sync contact objects (requester and responder).
            # Check local cache first, to reduce the no. of API calls.
            if ft.requester_id:
                if FreshdeskContact.objects.filter(contact_id=ft.requester_id).exists():
                    ft.freshdesk_requester = FreshdeskContact.objects.get(contact_id=ft.requester_id)
                else:  # Attempt to cache the Freshdesk contact.
                    try:
                        c = get_freshdesk_object(obj_type='contacts', id=ft.requester_id)
                        c['contact_id'] = c.pop('id')
                        c['created_at'] = parse(c['created_at'])
                        c['updated_at'] = parse(c['updated_at'])
                        c.pop('avatar', None)
                        c.pop('company_id', None)
                        c.pop('twitter_id', None)
                        c.pop('deleted', None)
                        con = FreshdeskContact.objects.create(**c)
                        logger.info('Created {}'.format(con))
                        ft.freshdesk_requester = con
                    except HTTPError:  # The GET might fail if the contact is an agent.
                        logger.error('HTTP 404 Freshdesk contact not found: {}'.format(ft.requester_id))
                        pass
            if ft.responder_id:
                if FreshdeskContact.objects.filter(contact_id=ft.responder_id).exists():
                    ft.freshdesk_responder = FreshdeskContact.objects.get(contact_id=ft.responder_id)
                else:  # Attempt to cache the Freshdesk contact.
                    try:
                        c = get_freshdesk_object(obj_type='contacts', id=ft.responder_id)
                        c['contact_id'] = c.pop('id')
                        c['created_at'] = parse(c['created_at'])
                        c['updated_at'] = parse(c['updated_at'])
                        c.pop('avatar', None)
                        c.pop('company_id', None)
                        c.pop('twitter_id', None)
                        c.pop('deleted', None)
                        con = FreshdeskContact.objects.create(**c)
                        logger.info('Created {}'.format(con))
                        ft.freshdesk_responder = con
                    except HTTPError:  # The GET might fail if the contact is an agent.
                        logger.error('HTTP 404 Freshdesk contact not found: {}'.format(ft.responder_id))
                        pass
            ft.save()

            # Try matching the ticket to an ITSystem object.
            ft.match_it_system()

            # Sync ticket conversation objects.
            obj = 'tickets/{}/conversations'.format(t['ticket_id'])
            convs = get_freshdesk_objects(obj_type=obj, progress=False)
            for c in convs:
                # Rename key 'id'.
                c['conversation_id'] = c.pop('id')
                # Date ISO8601-formatted date strings into datetimes.
                c['created_at'] = parse(c['created_at'])
                c['updated_at'] = parse(c['updated_at'])
                # Pop unused fields from the dict.
                c.pop('bcc_emails', None)
                c.pop('support_email', None)
                fc, create = FreshdeskConversation.objects.update_or_create(conversation_id=c['conversation_id'], defaults=c)
                if create:
                    logger.info('{} created'.format(fc))
                else:
                    logger.info('{} updated'.format(fc))
                # Link parent ticket, DepartmentUser, etc.
                fc.freshdesk_ticket = ft
                if FreshdeskContact.objects.filter(contact_id=fc.user_id).exists():
                    fc.freshdesk_contact = FreshdeskContact.objects.get(contact_id=fc.user_id)
                else:  # Attempt to cache the Freshdesk contact.
                    try:
                        f_con = get_freshdesk_object(obj_type='contacts', id=fc.user_id)
                        f_con['contact_id'] = f_con.pop('id')
                        f_con['created_at'] = parse(f_con['created_at'])
                        f_con['updated_at'] = parse(f_con['updated_at'])
                        f_con.pop('avatar', None)
                        f_con.pop('company_id', None)
                        f_con.pop('twitter_id', None)
                        f_con.pop('deleted', None)
                        contact = FreshdeskContact.objects.create(**f_con)
                        logger.info('Created {}'.format(contact))
                        fc.freshdesk_contact = contact
                        # Attempt to match contact with a DepartmentUser.
                        if DepartmentUser.objects.filter(email__iexact=contact.email).exists():
                            contact.du_user = DepartmentUser.objects.get(email__iexact=contact.email)
                            contact.save()
                    except HTTPError:  # The GET might fail if the contact is an agent.
                        logger.error('HTTP 404 Freshdesk contact not found: {}'.format(ft.requester_id))
                        pass
                # Attempt to match conversation with a DepartmentUser.
                if fc.freshdesk_contact and DepartmentUser.objects.filter(email__iexact=fc.freshdesk_contact.email).exists():
                    fc.du_user = DepartmentUser.objects.get(email__iexact=fc.freshdesk_contact.email)
                fc.save()
        except Exception as e:
            logger.exception(e)
            return False

    logger.info('Ticket sync: {} created, {} updated'.format(created, updated))
    print('{} created, {} updated'.format(created, updated))

    return True
<|endoftext|>


<commit_before>"""
#-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=#
  This file is part of the Smart Developer Hub Project:
    http://www.smartdeveloperhub.org

  Center for Open Middleware
        http://www.centeropenmiddleware.com/
#-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=#
  Copyright (C) 2015 Center for Open Middleware.
#-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=#
  Licensed under the Apache License, Version 2.0 (the "License");
  you may not use this file except in compliance with the License.
  You may obtain a copy of the License at

            http://www.apache.org/licenses/LICENSE-2.0

  Unless required by applicable law or agreed to in writing, software
  distributed under the License is distributed on an "AS IS" BASIS,
  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  See the License for the specific language governing permissions and
  limitations under the License.
#-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=#
"""
from redis.exceptions import RedisError, BusyLoadingError

__author__ = 'Fernando Serena'

import redis
import agora.fountain.vocab.schema as sch
from agora.fountain.exceptions import FountainError
import agora.fountain.vocab.onto as vocs
from datetime import datetime as dt
from concurrent.futures.thread import ThreadPoolExecutor
import logging
from agora.fountain.server import app
import sys

log = logging.getLogger('agora.fountain.index')

redis_conf = app.config['REDIS']
pool = redis.ConnectionPool(host=redis_conf.get('host'), port=6379, db=redis_conf.get('db'))
r = redis.StrictRedis(connection_pool=pool)
tpool = ThreadPoolExecutor(20)

# Ping redis to check if it's ready
requests = 0
while True:
    try:
        r.keys('*')
        break
    except BusyLoadingError as re:
        log.warning(re.message)
    except RedisError:
        print 'Redis is not available'
        sys.exit(-1)

def __get_by_pattern(pattern, func):
    def get_all():
        for k in pkeys:
            yield func(k)

    try:
        pkeys = r.keys(pattern)
        return list(get_all())
    except RedisError as e:
        raise FountainError(e.message)

def __remove_from_sets(values, *args):
    try:
        for pattern in args:
            keys = r.keys(pattern)
            for dk in keys:
                key_parts = dk.split(':')
                ef_values = values
                if len(key_parts) > 1:
                    ef_values = filter(lambda x: x.split(':')[0] != key_parts[1], values)
                if len(ef_values):
                    r.srem(dk, *ef_values)
    except RedisError as e:
        raise FountainError(e.message)

def __get_vocab_set(pattern, vid=None):
    try:
        if vid is not None:
            pattern = pattern.replace(':*:', ':%s:' % vid)
        all_sets = map(lambda x: r.smembers(x), r.keys(pattern))
        return list(reduce(set.union, all_sets, set([])))
    except RedisError as e:
        raise FountainError(e.message)

def __extract_type(t, vid):
    try:
        with r.pipeline() as pipe:
            pipe.multi()
            pipe.sadd('vocabs:{}:types'.format(vid), t)
            for s in sch.get_supertypes(t):
                pipe.sadd('vocabs:{}:types:{}:super'.format(vid, t), s)
            for s in sch.get_subtypes(t):
                pipe.sadd('vocabs:{}:types:{}:sub'.format(vid, t), s)
            for s in sch.get_type_properties(t):
                pipe.sadd('vocabs:{}:types:{}:props'.format(vid, t), s)
            for s in sch.get_type_references(t):
                pipe.sadd('vocabs:{}:types:{}:refs'.format(vid, t), s)
            pipe.execute()
    except RedisError as e:
        raise FountainError(e.message)

def __extract_property(p, vid):
    def p_type():
        if sch.is_object_property(p):
            return 'object'
        else:
            return 'data'

    try:
        with r.pipeline() as pipe:
            pipe.multi()
            pipe.sadd('vocabs:{}:properties'.format(vid), p)
            pipe.hset('vocabs:{}:properties:{}'.format(vid, p), 'uri', p)
            for dc in list(sch.get_property_domain(p)):
                pipe.sadd('vocabs:{}:properties:{}:domain'.format(vid, p), dc)
            for dc in list(sch.get_property_range(p)):
                pipe.sadd('vocabs:{}:properties:{}:range'.format(vid, p), dc)
            for dc in list(sch.get_property_inverses(p)):
                pipe.sadd('vocabs:{}:properties:{}:inverse'.format(vid, p), dc)
            pipe.set('vocabs:{}:properties:{}:type'.format(vid, p), p_type())
            pipe.execute()
    except RedisError as e:
        raise FountainError(e.message)

def __extract_types(vid):
    types = sch.get_types(vid)

    other_vocabs = filter(lambda x: x != vid, vocs.get_vocabularies())
    dependent_types = set([])
    dependent_props = set([])
    for ovid in other_vocabs:
        o_types = [t for t in get_types(ovid) if t not in types]
        for oty in o_types:
            otype = get_type(oty)
            if set.intersection(types, otype.get('super')) or set.intersection(types, otype.get('sub')):
                dependent_types.add((ovid, oty))
        o_props = [t for t in get_properties(ovid)]
        for op in o_props:
            oprop = get_property(op)
            if set.intersection(types, oprop.get('domain')) or set.intersection(types, oprop.get('range')):
                dependent_props.add((ovid, op))

    types = set.union(set([(vid, t) for t in types]), dependent_types)
    futures = []
    # TODO: Catch exceptions from threadpool
    for v, t in types:
        futures.append(tpool.submit(__extract_type, t, v))
    for v, p in dependent_props:
        futures.append(tpool.submit(__extract_property, p, v))
    return types, futures

def __extract_properties(vid):
    properties = sch.get_properties(vid)

    other_vocabs = filter(lambda x: x != vid, vocs.get_vocabularies())
    dependent_types = set([])
    for ovid in other_vocabs:
        o_types = [t for t in get_types(ovid)]
        for oty in o_types:
            otype = get_type(oty)
            if set.intersection(properties, otype.get('refs')) or set.intersection(properties, otype.get('properties')):
                dependent_types.add((ovid, oty))

    futures = []
    for p in properties:
        futures.append(tpool.submit(__extract_property, p, vid))
    for v, ty in dependent_types:
        futures.append(tpool.submit(__extract_type, ty, v))
    return properties, futures

def delete_vocabulary(vid):
    v_types = get_types(vid)
    if len(v_types):
        __remove_from_sets(v_types, '*:domain', '*:range', '*:sub', '*:super')
    v_props = get_properties(vid)
    if len(v_props):
        __remove_from_sets(v_props, '*:refs', '*:props')
    try:
        v_keys = r.keys('vocabs:{}:*'.format(vid))
        if len(v_keys):
            r.delete(*v_keys)
    except RedisError as e:
        raise FountainError(e.message)

def extract_vocabulary(vid):
    log.info('Extracting vocabulary {}...'.format(vid))
    delete_vocabulary(vid)
    start_time = dt.now()
    types, t_futures = __extract_types(vid)
    properties, p_futures = __extract_properties(vid)
    for f in t_futures + p_futures:
        f.result()
    log.info('Done (in {}ms)'.format((dt.now() - start_time).total_seconds() * 1000))
    return types, properties

def get_types(vid=None):
    return __get_vocab_set('vocabs:*:types', vid)

def get_properties(vid=None):
    return __get_vocab_set('vocabs:*:properties', vid)

def get_property(prop):
    def get_inverse_domain(ip):
        return reduce(set.union, __get_by_pattern('*:properties:{}:domain'.format(ip), r.smembers), set([]))

    def get_inverse_range(ip):
        return reduce(set.union, __get_by_pattern('*:properties:{}:range'.format(ip), r.smembers), set([]))

    all_prop_keys = r.keys('*:properties')
    if not filter(lambda k: r.sismember(k, prop), all_prop_keys):
        raise TypeError('Unknown property')

    domain = reduce(set.union, __get_by_pattern('*:properties:{}:domain'.format(prop), r.smembers), set([]))
    rang = reduce(set.union, __get_by_pattern('*:properties:{}:range'.format(prop), r.smembers), set([]))
    inv = reduce(set.union, __get_by_pattern('*:properties:{}:inverse'.format(prop), r.smembers), set([]))

    if len(inv):
        inverse_dr = [(get_inverse_domain(i), get_inverse_range(i)) for i in inv]
        for dom, ra in inverse_dr:
            domain.update(ra)
            rang.update(dom)

    ty = __get_by_pattern('*:properties:{}:type'.format(prop), r.get)
    try:
        ty = ty.pop()
    except IndexError:
        ty = 'object'

    return {'domain': list(domain), 'range': list(rang), 'inverse': list(inv), 'type': ty}

def is_property(prop):
    try:
        return len(r.keys('*:properties:{}:*'.format(prop)))
    except RedisError as e:
        raise FountainError(e.message)

def is_type(ty):
    try:
        return len(r.keys('*:types:{}:*'.format(ty)))
    except RedisError as e:
        raise FountainError(e.message)

def get_type(ty):
    try:
        all_type_keys = r.keys('*:types')
        if not filter(lambda k: r.sismember(k, ty), all_type_keys):
            raise TypeError('Unknown type')

        super_types = reduce(set.union, __get_by_pattern('*:types:{}:super'.format(ty), r.smembers), set([]))
        sub_types = reduce(set.union, __get_by_pattern('*:types:{}:sub'.format(ty), r.smembers), set([]))
        type_props = reduce(set.union, __get_by_pattern('*:types:{}:props'.format(ty), r.smembers), set([]))
        type_refs = reduce(set.union, __get_by_pattern('*:types:{}:refs'.format(ty), r.smembers), set([]))

        return {'super': list(super_types),
                'sub': list(sub_types),
                'properties': list(type_props),
                'refs': list(type_refs)}
    except RedisError as e:
        raise FountainError(e.message)
<commit_msg>Redis port in service config class<commit_after>"""
#-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=#
  This file is part of the Smart Developer Hub Project:
    http://www.smartdeveloperhub.org

  Center for Open Middleware
        http://www.centeropenmiddleware.com/
#-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=#
  Copyright (C) 2015 Center for Open Middleware.
#-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=#
  Licensed under the Apache License, Version 2.0 (the "License");
  you may not use this file except in compliance with the License.
  You may obtain a copy of the License at

            http://www.apache.org/licenses/LICENSE-2.0

  Unless required by applicable law or agreed to in writing, software
  distributed under the License is distributed on an "AS IS" BASIS,
  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  See the License for the specific language governing permissions and
  limitations under the License.
#-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=#
"""
from redis.exceptions import RedisError, BusyLoadingError

__author__ = 'Fernando Serena'

import redis
import agora.fountain.vocab.schema as sch
from agora.fountain.exceptions import FountainError
import agora.fountain.vocab.onto as vocs
from datetime import datetime as dt
from concurrent.futures.thread import ThreadPoolExecutor
import logging
from agora.fountain.server import app
import sys

log = logging.getLogger('agora.fountain.index')

redis_conf = app.config['REDIS']
pool = redis.ConnectionPool(host=redis_conf.get('host'), port=redis_conf.get('port'), db=redis_conf.get('db'))
r = redis.StrictRedis(connection_pool=pool)
tpool = ThreadPoolExecutor(20)

# Ping redis to check if it's ready
requests = 0
while True:
    try:
        r.keys('*')
        break
    except BusyLoadingError as re:
        log.warning(re.message)
    except RedisError:
        print 'Redis is not available'
        sys.exit(-1)

def __get_by_pattern(pattern, func):
    def get_all():
        for k in pkeys:
            yield func(k)

    try:
        pkeys = r.keys(pattern)
        return list(get_all())
    except RedisError as e:
        raise FountainError(e.message)

def __remove_from_sets(values, *args):
    try:
        for pattern in args:
            keys = r.keys(pattern)
            for dk in keys:
                key_parts = dk.split(':')
                ef_values = values
                if len(key_parts) > 1:
                    ef_values = filter(lambda x: x.split(':')[0] != key_parts[1], values)
                if len(ef_values):
                    r.srem(dk, *ef_values)
    except RedisError as e:
        raise FountainError(e.message)

def __get_vocab_set(pattern, vid=None):
    try:
        if vid is not None:
            pattern = pattern.replace(':*:', ':%s:' % vid)
        all_sets = map(lambda x: r.smembers(x), r.keys(pattern))
        return list(reduce(set.union, all_sets, set([])))
    except RedisError as e:
        raise FountainError(e.message)

def __extract_type(t, vid):
    try:
        with r.pipeline() as pipe:
            pipe.multi()
            pipe.sadd('vocabs:{}:types'.format(vid), t)
            for s in sch.get_supertypes(t):
                pipe.sadd('vocabs:{}:types:{}:super'.format(vid, t), s)
            for s in sch.get_subtypes(t):
                pipe.sadd('vocabs:{}:types:{}:sub'.format(vid, t), s)
            for s in sch.get_type_properties(t):
                pipe.sadd('vocabs:{}:types:{}:props'.format(vid, t), s)
            for s in sch.get_type_references(t):
                pipe.sadd('vocabs:{}:types:{}:refs'.format(vid, t), s)
            pipe.execute()
    except RedisError as e:
        raise FountainError(e.message)

def __extract_property(p, vid):
    def p_type():
        if sch.is_object_property(p):
            return 'object'
        else:
            return 'data'

    try:
        with r.pipeline() as pipe:
            pipe.multi()
            pipe.sadd('vocabs:{}:properties'.format(vid), p)
            pipe.hset('vocabs:{}:properties:{}'.format(vid, p), 'uri', p)
            for dc in list(sch.get_property_domain(p)):
                pipe.sadd('vocabs:{}:properties:{}:domain'.format(vid, p), dc)
            for dc in list(sch.get_property_range(p)):
                pipe.sadd('vocabs:{}:properties:{}:range'.format(vid, p), dc)
            for dc in list(sch.get_property_inverses(p)):
                pipe.sadd('vocabs:{}:properties:{}:inverse'.format(vid, p), dc)
            pipe.set('vocabs:{}:properties:{}:type'.format(vid, p), p_type())
            pipe.execute()
    except RedisError as e:
        raise FountainError(e.message)

def __extract_types(vid):
    types = sch.get_types(vid)

    other_vocabs = filter(lambda x: x != vid, vocs.get_vocabularies())
    dependent_types = set([])
    dependent_props = set([])
    for ovid in other_vocabs:
        o_types = [t for t in get_types(ovid) if t not in types]
        for oty in o_types:
            otype = get_type(oty)
            if set.intersection(types, otype.get('super')) or set.intersection(types, otype.get('sub')):
                dependent_types.add((ovid, oty))
        o_props = [t for t in get_properties(ovid)]
        for op in o_props:
            oprop = get_property(op)
            if set.intersection(types, oprop.get('domain')) or set.intersection(types, oprop.get('range')):
                dependent_props.add((ovid, op))

    types = set.union(set([(vid, t) for t in types]), dependent_types)
    futures = []
    # TODO: Catch exceptions from threadpool
    for v, t in types:
        futures.append(tpool.submit(__extract_type, t, v))
    for v, p in dependent_props:
        futures.append(tpool.submit(__extract_property, p, v))
    return types, futures

def __extract_properties(vid):
    properties = sch.get_properties(vid)

    other_vocabs = filter(lambda x: x != vid, vocs.get_vocabularies())
    dependent_types = set([])
    for ovid in other_vocabs:
        o_types = [t for t in get_types(ovid)]
        for oty in o_types:
            otype = get_type(oty)
            if set.intersection(properties, otype.get('refs')) or set.intersection(properties, otype.get('properties')):
                dependent_types.add((ovid, oty))

    futures = []
    for p in properties:
        futures.append(tpool.submit(__extract_property, p, vid))
    for v, ty in dependent_types:
        futures.append(tpool.submit(__extract_type, ty, v))
    return properties, futures

def delete_vocabulary(vid):
    v_types = get_types(vid)
    if len(v_types):
        __remove_from_sets(v_types, '*:domain', '*:range', '*:sub', '*:super')
    v_props = get_properties(vid)
    if len(v_props):
        __remove_from_sets(v_props, '*:refs', '*:props')
    try:
        v_keys = r.keys('vocabs:{}:*'.format(vid))
        if len(v_keys):
            r.delete(*v_keys)
    except RedisError as e:
        raise FountainError(e.message)

def extract_vocabulary(vid):
    log.info('Extracting vocabulary {}...'.format(vid))
    delete_vocabulary(vid)
    start_time = dt.now()
    types, t_futures = __extract_types(vid)
    properties, p_futures = __extract_properties(vid)
    for f in t_futures + p_futures:
        f.result()
    log.info('Done (in {}ms)'.format((dt.now() - start_time).total_seconds() * 1000))
    return types, properties

def get_types(vid=None):
    return __get_vocab_set('vocabs:*:types', vid)

def get_properties(vid=None):
    return __get_vocab_set('vocabs:*:properties', vid)

def get_property(prop):
    def get_inverse_domain(ip):
        return reduce(set.union, __get_by_pattern('*:properties:{}:domain'.format(ip), r.smembers), set([]))

    def get_inverse_range(ip):
        return reduce(set.union, __get_by_pattern('*:properties:{}:range'.format(ip), r.smembers), set([]))

    all_prop_keys = r.keys('*:properties')
    if not filter(lambda k: r.sismember(k, prop), all_prop_keys):
        raise TypeError('Unknown property')

    domain = reduce(set.union, __get_by_pattern('*:properties:{}:domain'.format(prop), r.smembers), set([]))
    rang = reduce(set.union, __get_by_pattern('*:properties:{}:range'.format(prop), r.smembers), set([]))
    inv = reduce(set.union, __get_by_pattern('*:properties:{}:inverse'.format(prop), r.smembers), set([]))

    if len(inv):
        inverse_dr = [(get_inverse_domain(i), get_inverse_range(i)) for i in inv]
        for dom, ra in inverse_dr:
            domain.update(ra)
            rang.update(dom)

    ty = __get_by_pattern('*:properties:{}:type'.format(prop), r.get)
    try:
        ty = ty.pop()
    except IndexError:
        ty = 'object'

    return {'domain': list(domain), 'range': list(rang), 'inverse': list(inv), 'type': ty}

def is_property(prop):
    try:
        return len(r.keys('*:properties:{}:*'.format(prop)))
    except RedisError as e:
        raise FountainError(e.message)

def is_type(ty):
    try:
        return len(r.keys('*:types:{}:*'.format(ty)))
    except RedisError as e:
        raise FountainError(e.message)

def get_type(ty):
    try:
        all_type_keys = r.keys('*:types')
        if not filter(lambda k: r.sismember(k, ty), all_type_keys):
            raise TypeError('Unknown type')

        super_types = reduce(set.union, __get_by_pattern('*:types:{}:super'.format(ty), r.smembers), set([]))
        sub_types = reduce(set.union, __get_by_pattern('*:types:{}:sub'.format(ty), r.smembers), set([]))
        type_props = reduce(set.union, __get_by_pattern('*:types:{}:props'.format(ty), r.smembers), set([]))
        type_refs = reduce(set.union, __get_by_pattern('*:types:{}:refs'.format(ty), r.smembers), set([]))

        return {'super': list(super_types),
                'sub': list(sub_types),
                'properties': list(type_props),
                'refs': list(type_refs)}
    except RedisError as e:
        raise FountainError(e.message)
<|endoftext|>


<commit_before>from . import protocols
from os import urandom

def encryptFileObj(
    reader,
    writer,
    cipher,
    validate=False,
    _prechunk=False
):
    if _prechunk:
        writer.write(cipher.encrypt(urandom(16)))
    if validate:
        writer.write(cipher.encrypt(b'\x00'*16))
    intake = reader.read(4095)
    while len(intake):
        writer.write(_encrypt_chunk(intake, cipher))
        writer.flush()
        intake = reader.read(4095)

def encryptFile(
    input_filename,
    output_filename,
    cipher,
    validate=False,
    _prechunk=False
):
    reader = open(input_filename, mode='rb')
    writer = open(output_filename, mode='wb')
    encryptFileObj(reader, writer, cipher, validate=validate, _prechunk=_prechunk)
    reader.close()
    writer.close()

def decryptFileObj(
    reader,
    writer,
    cipher,
    validate=False,
    _prechunk=False
):
    if _prechunk:
        cipher.decrypt(reader.read(16))
    if validate:
        if cipher.decrypt(reader.read(16)) != b'\x00'*16:
            reader.close()
            writer.close()
            raise KeyError(
                "Decryption Failed! The cipher used may be incorrect"
            )
    intake = reader.read(4096)
    while len(intake):
        writer.write(_decrypt_chunk(intake, cipher))
        writer.flush()
        intake = reader.read(4096)

def decryptFile(
    input_filename,
    output_filename,
    cipher,
    validate=False,
    _prechunk=False
):
    reader = open(input_filename, mode='rb')
    writer = open(output_filename, mode='wb')
    decryptFileObj(reader, writer, cipher, validate=validate, _prechunk=_prechunk)
    reader.close()
    writer.close()

def _encrypt_chunk(chunk, cipher):
    return cipher.encrypt(protocols.padstring(chunk))

def _decrypt_chunk(chunk, cipher):
    return protocols.unpadstring(cipher.decrypt(chunk))
<commit_msg>PEP8 compliance<commit_after>from . import protocols
from os import urandom

def encryptFileObj(
    reader,
    writer,
    cipher,
    validate=False,
    _prechunk=False
):
    if _prechunk:
        writer.write(cipher.encrypt(urandom(16)))
    if validate:
        writer.write(cipher.encrypt(b'\x00'*16))
    intake = reader.read(4095)
    while len(intake):
        writer.write(_encrypt_chunk(intake, cipher))
        writer.flush()
        intake = reader.read(4095)

def encryptFile(
    input_filename,
    output_filename,
    cipher,
    validate=False,
    _prechunk=False
):
    reader = open(input_filename, mode='rb')
    writer = open(output_filename, mode='wb')
    encryptFileObj(
        reader,
        writer,
        cipher,
        validate=validate,
        _prechunk=_prechunk
    )
    reader.close()
    writer.close()

def decryptFileObj(
    reader,
    writer,
    cipher,
    validate=False,
    _prechunk=False
):
    if _prechunk:
        cipher.decrypt(reader.read(16))
    if validate:
        if cipher.decrypt(reader.read(16)) != b'\x00'*16:
            reader.close()
            writer.close()
            raise KeyError(
                "Decryption Failed! The cipher used may be incorrect"
            )
    intake = reader.read(4096)
    while len(intake):
        writer.write(_decrypt_chunk(intake, cipher))
        writer.flush()
        intake = reader.read(4096)

def decryptFile(
    input_filename,
    output_filename,
    cipher,
    validate=False,
    _prechunk=False
):
    reader = open(input_filename, mode='rb')
    writer = open(output_filename, mode='wb')
    decryptFileObj(
        reader,
        writer,
        cipher,
        validate=validate,
        _prechunk=_prechunk
    )
    reader.close()
    writer.close()

def _encrypt_chunk(chunk, cipher):
    return cipher.encrypt(protocols.padstring(chunk))

def _decrypt_chunk(chunk, cipher):
    return protocols.unpadstring(cipher.decrypt(chunk))
<|endoftext|>


<commit_before>import asyncio

import toolz

import aioredux
from aioredux.tests import base

class TestAioredux(base.TestCase):

    def setUp(self):
        self.loop = asyncio.new_event_loop()
        asyncio.set_event_loop(None)
        super().setUp()

    def tearDown(self):
        self.loop.close()
        self.loop = None
        super().tearDown()

    def test_todo(self):

        @asyncio.coroutine
        def go():
            initial_state = {
                'todos': (),
            }
            store = yield from aioredux.create_store(lambda state, action: state, initial_state, loop=self.loop)
            self.assertIsNotNone(store)
            self.assertIsNotNone(store.state)

        self.loop.run_until_complete(go())

    def test_todo_reducer(self):

        @asyncio.coroutine
        def go():
            initial_state = {
                'todos': (),
            }

            def add_todo(text):
                return {'type': 'ADD_TODO', 'text': text}

            def reducer(state, action):
                if action['type'] == 'ADD_TODO':
                    todos = state['todos'] + (action['text'],)
                    return toolz.assoc(state, 'todos', todos)
                return state

            store = yield from aioredux.create_store(reducer, initial_state, loop=self.loop)
            self.assertIsNotNone(store.state)
            self.assertIsNotNone(store.state['todos'])
            self.assertEqual(len(store.state['todos']), 0)
            yield from store.dispatch(add_todo('todo text'))
            self.assertIsNotNone(store.state)
            self.assertIsNotNone(store.state['todos'])
            self.assertEqual(len(store.state['todos']), 1)

        self.loop.run_until_complete(go())
<commit_msg>TST: add test for subscribe<commit_after>import asyncio

import toolz

import aioredux
from aioredux.tests import base

class TestBasic(base.TestCase):

    def setUp(self):
        self.loop = asyncio.new_event_loop()
        asyncio.set_event_loop(None)
        super().setUp()

    def tearDown(self):
        self.loop.close()
        self.loop = None
        super().tearDown()

    def test_todo(self):

        @asyncio.coroutine
        def go():
            initial_state = {
                'todos': (),
            }
            store = yield from aioredux.create_store(lambda state, action: state, initial_state, loop=self.loop)
            self.assertIsNotNone(store)
            self.assertIsNotNone(store.state)

        self.loop.run_until_complete(go())

    def test_todo_reducer(self):

        @asyncio.coroutine
        def go():
            initial_state = {
                'todos': (),
            }

            def add_todo(text):
                return {'type': 'ADD_TODO', 'text': text}

            def reducer(state, action):
                if action['type'] == 'ADD_TODO':
                    todos = state['todos'] + (action['text'],)
                    return toolz.assoc(state, 'todos', todos)
                return state

            store = yield from aioredux.create_store(reducer, initial_state, loop=self.loop)
            self.assertIsNotNone(store.state)
            self.assertIsNotNone(store.state['todos'])
            self.assertEqual(len(store.state['todos']), 0)
            yield from store.dispatch(add_todo('todo text'))
            self.assertIsNotNone(store.state)
            self.assertIsNotNone(store.state['todos'])
            self.assertEqual(len(store.state['todos']), 1)

        self.loop.run_until_complete(go())

    def test_subscribe(self):

        @asyncio.coroutine
        def go():
            initial_state = {
                'todos': (),
            }
            store = yield from aioredux.create_store(lambda state, action: state, initial_state, loop=self.loop)
            self.assertEqual(len(store.listeners), 0)
            unsubscribe = store.subscribe(lambda: None)
            self.assertEqual(len(store.listeners), 1)
            unsubscribe()
            self.assertEqual(len(store.listeners), 0)

        self.loop.run_until_complete(go())
<|endoftext|>


<commit_before>from datetime import datetime
import logging
from urlparse import urlparse
from time import sleep

from airflow import settings
from airflow import hooks
from airflow.models import BaseOperator
from airflow.models import Connection as DB
from airflow.models import State
from airflow.models import TaskInstance
from airflow.utils import apply_defaults

from snakebite.client import HAClient, Namenode

class BaseSensorOperator(BaseOperator):
    '''
    Sensor operators are derived from this class an inherit these attributes.

    Sensor operators keep executing at a time interval and succeed when
        a criteria is met and fail if and when they time out.

    :param poke_interval: Time in seconds that the job should wait in
        between each tries
    :type poke_interval: int
    :param timeout: Time, in seconds before the task times out and fails.
    :type timeout: int
    '''
    ui_color = '#e6f1f2'

    @apply_defaults
    def __init__(
            self,
            poke_interval=60,
            timeout=60*60*24*7,
            *args, **kwargs):
        super(BaseSensorOperator, self).__init__(*args, **kwargs)
        self.poke_interval = poke_interval
        self.timeout = timeout

    def poke(self):
        '''
        Function that the sensors defined while deriving this class should
        override.
        '''
        raise Exception('Override me.')

    def execute(self, context):
        started_at = datetime.now()
        while not self.poke():
            sleep(self.poke_interval)
            if (datetime.now() - started_at).seconds > self.timeout:
                raise Exception('Snap. Time is OUT.')
        logging.info("Success criteria met. Exiting.")

class SqlSensor(BaseSensorOperator):
    """
    Runs a sql statement until a criteria is met. It will keep trying until
    sql returns no row, or if the first cell in (0, '0', '').

    :param conn_id: The connection to run the sensor agains
    :type conn_id: string
    :param sql: The sql to run. To pass, it needs to return at least one cell
        that contains a non-zero / empty string value.
    """
    template_fields = ('sql',)
    template_ext = ('.hql', '.sql',)

    __mapper_args__ = {
        'polymorphic_identity': 'SqlSensor'
    }

    @apply_defaults
    def __init__(self, conn_id, sql, *args, **kwargs):

        super(SqlSensor, self).__init__(*args, **kwargs)

        self.sql = sql
        self.conn_id = conn_id

        session = settings.Session()
        db = session.query(DB).filter(DB.conn_id == conn_id).first()
        if not db:
            raise Exception("conn_id doesn't exist in the repository")
        self.hook = db.get_hook()
        session.commit()
        session.close()

    def poke(self):
        logging.info('Poking: ' + self.sql)
        records = self.hook.get_records(self.sql)
        if not records:
            return False
        else:
            if str(records[0][0]) in ('0', '',):
                return False
            else:
                return True
            print(records[0][0])

class ExternalTaskSensor(BaseSensorOperator):
    """
    Waits for a task to complete in a different DAG

    :param external_dag_id: The dag_id that contains the task you want to
        wait for
    :type external_dag_id: string
    :param external_task_id: The task_id that contains the task you want to
        wait for
    :type external_task_id: string
    """
    template_fields = ('execution_date',)
    __mapper_args__ = {
        'polymorphic_identity': 'ExternalTaskSensor'
    }

    @apply_defaults
    def __init__(self, external_dag_id, external_task_id, *args, **kwargs):
        super(ExternalTaskSensor, self).__init__(*args, **kwargs)
        self.external_dag_id = external_dag_id
        self.external_task_id = external_task_id
        self.execution_date = "{{ execution_date }}"

    def poke(self):
        logging.info(
            'Poking for '
            '{self.external_dag_id}.'
            '{self.external_task_id} on '
            '{self.execution_date} ... '.format(**locals()))
        TI = TaskInstance
        session = settings.Session()
        count = session.query(TI).filter(
            TI.dag_id == self.external_dag_id,
            TI.task_id == self.external_task_id,
            TI.state == State.SUCCESS,
            TI.execution_date == self.execution_date
        ).count()
        session.commit()
        session.close()
        return count

class HivePartitionSensor(BaseSensorOperator):
    """
    Waits for a partition to show up in Hive

    :param table: The name of the table to wait for, supports the dot
        notation (my_database.my_table)
    :type table: string
    :param partition: The partition clause to wait for. This is passed as
        is to the Metastor Thrift client "get_partitions_by_filter" method,
        and apparently supports SQL like notation as in `ds='2015-01-01'
        AND type='value'` and > < sings as in "ds>=2015-01-01"
    :type partition: string
    """
    template_fields = ('table', 'partition',)
    __mapper_args__ = {
        'polymorphic_identity': 'HivePartitionSensor'
    }

    @apply_defaults
    def __init__(
            self,
            table, partition="ds='{{ ds }}'",
            metastore_conn_id='metastore_default',
            schema='default',
            poke_interval=60*3,
            *args, **kwargs):
        super(HivePartitionSensor, self).__init__(
            poke_interval=poke_interval, *args, **kwargs)
        if '.' in table:
            schema, table = table.split('.')
        if not partition:
            partition = "ds='{{ ds }}'"
        self.metastore_conn_id = metastore_conn_id
        self.table = table
        self.partition = partition
        self.schema = schema

    def poke(self):
        logging.info(
            'Poking for table {self.schema}.{self.table}, '
            'partition {self.partition}'.format(**locals()))
        self.hook = hooks.HiveMetastoreHook(
            metastore_conn_id=self.metastore_conn_id)
        return self.hook.check_for_partition(
            self.schema, self.table, self.partition)

class HdfsSensor(BaseSensorOperator):
    """
    Waits for a file or folder to land in HDFS
    """
    template_fields = ('filepath',)
    __mapper_args__ = {
        'polymorphic_identity': 'HdfsSensor'
    }

    @apply_defaults
    def __init__(
            self,
            filepath,
            hdfs_conn_id='hdfs_default',
            *args, **kwargs):
        super(HdfsSensor, self).__init__(*args, **kwargs)
        self.filepath = filepath
        self.sb = HDFSHook(hdfs_conn_id).get_conn()

    def poke(self):
        logging.getLogger("snakebite").setLevel(logging.WARNING)
        logging.info(
            'Poking for file {self.filepath} '.format(**locals()))
        try:
            files = [f for f in self.sb.ls([self.filepath])]
        except:
            return False
        return True

class S3KeySensor(BaseSensorOperator):
    """
    Waits for a key (a file-like instance on S3) to be present in a S3 bucket.
    S3 being a key/value it does not support folders. The path is just a key
    a resource.

    :param bucket_key: The key being waited on. Supports full s3:// style url
        or relative path from root level.
    :type bucket_key: str
    :param bucket_name: Name of the S3 bucket
    :type bucket_name: str
    """
    template_fields = ('bucket_key', 'bucket_name')
    __mapper_args__ = {
        'polymorphic_identity': 'S3KeySensor'
    }

    @apply_defaults
    def __init__(
            self, bucket_key,
            bucket_name=None,
            s3_conn_id='s3_default',
            *args, **kwargs):
        super(S3KeySensor, self).__init__(*args, **kwargs)
        session = settings.Session()
        db = session.query(DB).filter(DB.conn_id == s3_conn_id).first()
        if not db:
            raise Exception("conn_id doesn't exist in the repository")
        # Parse
        if bucket_name is None:
            parsed_url = urlparse(bucket_key)
            if parsed_url.netloc == '':
                raise Exception('Please provide a bucket_name')
            else:
                bucket_name = parsed_url.netloc
                bucket_key = parsed_url.path
        self.bucket_name = bucket_name
        self.bucket_key = bucket_key
        self.s3_conn_id = s3_conn_id
        session.commit()
        session.close()

    def poke(self):
        hook = hooks.S3Hook(s3_conn_id=self.s3_conn_id)
        full_url = "s3://" + self.bucket_name + self.bucket_key
        logging.info('Poking for key : {full_url}'.format(**locals()))
        return hook.check_for_key(self.bucket_key, self.bucket_name)

class S3PrefixSensor(BaseSensorOperator):
    """
    Waits for a prefix to exist. A prefix is the first part of a key,
    thus enabling checking of constructs similar to glob airfl* or
    SQL LIKE 'airfl%'. There is the possibility to precise a delimiter to
    indicate the hierarchy or keys, meaning that the match will stop at that
    delimiter. Current code accepts sane delimiters, i.e. characters that
    are NOT special characters in the Python regex engine.

    :param prefix: The key being waited on. Supports full s3:// style url or
        or relative path from root level.
    :type prefix: str
    :param delimiter: The delimiter intended to show hierarchy.
        Defaults to '/'.
    :type delimiter: str
    :param bucket_name: Name of the S3 bucket
    :type bucket_name: str
    """
    template_fields = ('bucket_key', 'bucket_name')
    __mapper_args__ = {
        'polymorphic_identity': 'S3PrefixSensor'
    }

    @apply_defaults
    def __init__(
            self, bucket_name,
            prefix, delimiter='/',
            s3_conn_id='s3_default',
            *args, **kwargs):
        super(S3PrefixSensor, self).__init__(*args, **kwargs)
        session = settings.Session()
        db = session.query(DB).filter(DB.conn_id == s3_conn_id).first()
        if not db:
            raise Exception("conn_id doesn't exist in the repository")
        # Parse
        self.bucket_name = bucket_name
        self.prefix = prefix
        self.delimiter = delimiter
        self.full_url = "s3://" + bucket_name + '/' + prefix
        self.s3_conn_id = s3_conn_id
        session.commit()
        session.close()

    def poke(self):
        logging.info('Poking for prefix : {self.prefix}\n'
                     'in bucket s3://{self.bucket_name}'.format(**locals()))
        hook = hooks.S3Hook(s3_conn_id=self.s3_conn_id)
        return hook.check_for_prefix(
            prefix=self.prefix,
            delimiter=self.delimiter,
            bucket_name=self.bucket_name)

class TimeSensor(BaseSensorOperator):
    """
    Waits until the specified time of the day.

    :param target_time: time after which the job succeeds
    :type target_time: datetime.time
    """
    template_fields = tuple()
    __mapper_args__ = {
        'polymorphic_identity': 'TimeSensor'
    }

    @apply_defaults
    def __init__(self, target_time, *args, **kwargs):
        super(TimeSensor, self).__init__(*args, **kwargs)
        self.target_time = target_time

    def poke(self):
        logging.info(
            'Checking if the time ({0}) has come'.format(self.target_time))
        return datetime.now().time() > self.target_time
<commit_msg>Fixes around HDFSSensor<commit_after>from datetime import datetime
import logging
from urlparse import urlparse
from time import sleep

from airflow import settings
from airflow import hooks
from airflow.models import BaseOperator
from airflow.models import Connection as DB
from airflow.models import State
from airflow.models import TaskInstance
from airflow.utils import apply_defaults

class BaseSensorOperator(BaseOperator):
    '''
    Sensor operators are derived from this class an inherit these attributes.

    Sensor operators keep executing at a time interval and succeed when
        a criteria is met and fail if and when they time out.

    :param poke_interval: Time in seconds that the job should wait in
        between each tries
    :type poke_interval: int
    :param timeout: Time, in seconds before the task times out and fails.
    :type timeout: int
    '''
    ui_color = '#e6f1f2'

    @apply_defaults
    def __init__(
            self,
            poke_interval=60,
            timeout=60*60*24*7,
            *args, **kwargs):
        super(BaseSensorOperator, self).__init__(*args, **kwargs)
        self.poke_interval = poke_interval
        self.timeout = timeout

    def poke(self):
        '''
        Function that the sensors defined while deriving this class should
        override.
        '''
        raise Exception('Override me.')

    def execute(self, context):
        started_at = datetime.now()
        while not self.poke():
            sleep(self.poke_interval)
            if (datetime.now() - started_at).seconds > self.timeout:
                raise Exception('Snap. Time is OUT.')
        logging.info("Success criteria met. Exiting.")

class SqlSensor(BaseSensorOperator):
    """
    Runs a sql statement until a criteria is met. It will keep trying until
    sql returns no row, or if the first cell in (0, '0', '').

    :param conn_id: The connection to run the sensor agains
    :type conn_id: string
    :param sql: The sql to run. To pass, it needs to return at least one cell
        that contains a non-zero / empty string value.
    """
    template_fields = ('sql',)
    template_ext = ('.hql', '.sql',)

    __mapper_args__ = {
        'polymorphic_identity': 'SqlSensor'
    }

    @apply_defaults
    def __init__(self, conn_id, sql, *args, **kwargs):

        super(SqlSensor, self).__init__(*args, **kwargs)

        self.sql = sql
        self.conn_id = conn_id

        session = settings.Session()
        db = session.query(DB).filter(DB.conn_id == conn_id).first()
        if not db:
            raise Exception("conn_id doesn't exist in the repository")
        self.hook = db.get_hook()
        session.commit()
        session.close()

    def poke(self):
        logging.info('Poking: ' + self.sql)
        records = self.hook.get_records(self.sql)
        if not records:
            return False
        else:
            if str(records[0][0]) in ('0', '',):
                return False
            else:
                return True
            print(records[0][0])

class ExternalTaskSensor(BaseSensorOperator):
    """
    Waits for a task to complete in a different DAG

    :param external_dag_id: The dag_id that contains the task you want to
        wait for
    :type external_dag_id: string
    :param external_task_id: The task_id that contains the task you want to
        wait for
    :type external_task_id: string
    """
    template_fields = ('execution_date',)
    __mapper_args__ = {
        'polymorphic_identity': 'ExternalTaskSensor'
    }

    @apply_defaults
    def __init__(self, external_dag_id, external_task_id, *args, **kwargs):
        super(ExternalTaskSensor, self).__init__(*args, **kwargs)
        self.external_dag_id = external_dag_id
        self.external_task_id = external_task_id
        self.execution_date = "{{ execution_date }}"

    def poke(self):
        logging.info(
            'Poking for '
            '{self.external_dag_id}.'
            '{self.external_task_id} on '
            '{self.execution_date} ... '.format(**locals()))
        TI = TaskInstance
        session = settings.Session()
        count = session.query(TI).filter(
            TI.dag_id == self.external_dag_id,
            TI.task_id == self.external_task_id,
            TI.state == State.SUCCESS,
            TI.execution_date == self.execution_date
        ).count()
        session.commit()
        session.close()
        return count

class HivePartitionSensor(BaseSensorOperator):
    """
    Waits for a partition to show up in Hive

    :param table: The name of the table to wait for, supports the dot
        notation (my_database.my_table)
    :type table: string
    :param partition: The partition clause to wait for. This is passed as
        is to the Metastor Thrift client "get_partitions_by_filter" method,
        and apparently supports SQL like notation as in `ds='2015-01-01'
        AND type='value'` and > < sings as in "ds>=2015-01-01"
    :type partition: string
    """
    template_fields = ('table', 'partition',)
    __mapper_args__ = {
        'polymorphic_identity': 'HivePartitionSensor'
    }

    @apply_defaults
    def __init__(
            self,
            table, partition="ds='{{ ds }}'",
            metastore_conn_id='metastore_default',
            schema='default',
            poke_interval=60*3,
            *args, **kwargs):
        super(HivePartitionSensor, self).__init__(
            poke_interval=poke_interval, *args, **kwargs)
        if '.' in table:
            schema, table = table.split('.')
        if not partition:
            partition = "ds='{{ ds }}'"
        self.metastore_conn_id = metastore_conn_id
        self.table = table
        self.partition = partition
        self.schema = schema

    def poke(self):
        logging.info(
            'Poking for table {self.schema}.{self.table}, '
            'partition {self.partition}'.format(**locals()))
        self.hook = hooks.HiveMetastoreHook(
            metastore_conn_id=self.metastore_conn_id)
        return self.hook.check_for_partition(
            self.schema, self.table, self.partition)

class HdfsSensor(BaseSensorOperator):
    """
    Waits for a file or folder to land in HDFS
    """
    template_fields = ('filepath',)
    __mapper_args__ = {
        'polymorphic_identity': 'HdfsSensor'
    }

    @apply_defaults
    def __init__(
            self,
            filepath,
            hdfs_conn_id='hdfs_default',
            *args, **kwargs):
        super(HdfsSensor, self).__init__(*args, **kwargs)
        self.filepath = filepath
        self.hdfs_conn_id = hdfs_conn_id

    def poke(self):
        sb = hooks.HDFSHook(self.hdfs_conn_id).get_conn()
        logging.getLogger("snakebite").setLevel(logging.WARNING)
        logging.info(
            'Poking for file {self.filepath} '.format(**locals()))
        try:
            files = [f for f in sb.ls([self.filepath])]
        except:
            return False
        return True

class S3KeySensor(BaseSensorOperator):
    """
    Waits for a key (a file-like instance on S3) to be present in a S3 bucket.
    S3 being a key/value it does not support folders. The path is just a key
    a resource.

    :param bucket_key: The key being waited on. Supports full s3:// style url
        or relative path from root level.
    :type bucket_key: str
    :param bucket_name: Name of the S3 bucket
    :type bucket_name: str
    """
    template_fields = ('bucket_key', 'bucket_name')
    __mapper_args__ = {
        'polymorphic_identity': 'S3KeySensor'
    }

    @apply_defaults
    def __init__(
            self, bucket_key,
            bucket_name=None,
            s3_conn_id='s3_default',
            *args, **kwargs):
        super(S3KeySensor, self).__init__(*args, **kwargs)
        session = settings.Session()
        db = session.query(DB).filter(DB.conn_id == s3_conn_id).first()
        if not db:
            raise Exception("conn_id doesn't exist in the repository")
        # Parse
        if bucket_name is None:
            parsed_url = urlparse(bucket_key)
            if parsed_url.netloc == '':
                raise Exception('Please provide a bucket_name')
            else:
                bucket_name = parsed_url.netloc
                bucket_key = parsed_url.path
        self.bucket_name = bucket_name
        self.bucket_key = bucket_key
        self.s3_conn_id = s3_conn_id
        session.commit()
        session.close()

    def poke(self):
        hook = hooks.S3Hook(s3_conn_id=self.s3_conn_id)
        full_url = "s3://" + self.bucket_name + self.bucket_key
        logging.info('Poking for key : {full_url}'.format(**locals()))
        return hook.check_for_key(self.bucket_key, self.bucket_name)

class S3PrefixSensor(BaseSensorOperator):
    """
    Waits for a prefix to exist. A prefix is the first part of a key,
    thus enabling checking of constructs similar to glob airfl* or
    SQL LIKE 'airfl%'. There is the possibility to precise a delimiter to
    indicate the hierarchy or keys, meaning that the match will stop at that
    delimiter. Current code accepts sane delimiters, i.e. characters that
    are NOT special characters in the Python regex engine.

    :param prefix: The key being waited on. Supports full s3:// style url or
        or relative path from root level.
    :type prefix: str
    :param delimiter: The delimiter intended to show hierarchy.
        Defaults to '/'.
    :type delimiter: str
    :param bucket_name: Name of the S3 bucket
    :type bucket_name: str
    """
    template_fields = ('bucket_key', 'bucket_name')
    __mapper_args__ = {
        'polymorphic_identity': 'S3PrefixSensor'
    }

    @apply_defaults
    def __init__(
            self, bucket_name,
            prefix, delimiter='/',
            s3_conn_id='s3_default',
            *args, **kwargs):
        super(S3PrefixSensor, self).__init__(*args, **kwargs)
        session = settings.Session()
        db = session.query(DB).filter(DB.conn_id == s3_conn_id).first()
        if not db:
            raise Exception("conn_id doesn't exist in the repository")
        # Parse
        self.bucket_name = bucket_name
        self.prefix = prefix
        self.delimiter = delimiter
        self.full_url = "s3://" + bucket_name + '/' + prefix
        self.s3_conn_id = s3_conn_id
        session.commit()
        session.close()

    def poke(self):
        logging.info('Poking for prefix : {self.prefix}\n'
                     'in bucket s3://{self.bucket_name}'.format(**locals()))
        hook = hooks.S3Hook(s3_conn_id=self.s3_conn_id)
        return hook.check_for_prefix(
            prefix=self.prefix,
            delimiter=self.delimiter,
            bucket_name=self.bucket_name)

class TimeSensor(BaseSensorOperator):
    """
    Waits until the specified time of the day.

    :param target_time: time after which the job succeeds
    :type target_time: datetime.time
    """
    template_fields = tuple()
    __mapper_args__ = {
        'polymorphic_identity': 'TimeSensor'
    }

    @apply_defaults
    def __init__(self, target_time, *args, **kwargs):
        super(TimeSensor, self).__init__(*args, **kwargs)
        self.target_time = target_time

    def poke(self):
        logging.info(
            'Checking if the time ({0}) has come'.format(self.target_time))
        return datetime.now().time() > self.target_time
<|endoftext|>


<commit_before>from abc import abstractmethod
import cProfile, pstats, io

import datetime
from meerkat_abacus import model
from meerkat_abacus import logger

class ProcessingStep(object):
    """
    Base class for all ProcessingSteps

    """

    def __init__(self):
        self.step_name = "processing_step"
        self.session = None
        self.start = None
        self.end = None

    def __repr__(self):
        return f'<{self.__class__}, step_name="{self.step_name}">'

    @property
    def duration(self):
        if not self.start or not self.end:
            return None
        return self.end - self.start

    @abstractmethod
    def run(self, form, data):
        pass

    def start_step(self):
        self.start = datetime.datetime.now()
        # self.profiler = cProfile.Profile()
        # self.profiler.enable()

    def end_step(self, n):
        self.end = datetime.datetime.now()
        self._write_monitoring_data(n)
        # self.profiler.disable()
        s = io.StringIO()
        sortby = 'cumulative'
        # ps = pstats.Stats(self.profiler, stream=s).sort_stats(sortby)
        # ps.print_stats(30)
        logger.info(s.getvalue())

    def _write_monitoring_data(self, n=None):
        monitoring = model.StepMonitoring(
            step=self.step_name,
            end=self.end,
            start=self.start,
            duration=self.duration.total_seconds(),
            n=n)
        self.session.add(monitoring)
        self.session.commit()

class DoNothing(ProcessingStep):
    def __init__(self, session):
        super().__init__()
        self.step_name = "do_nothing"
        self.session = session

    def run(self, form, data):
        return [{"form": form,
                 "data": data}]
<commit_msg>Disable noisy logging.<commit_after>from abc import abstractmethod
import cProfile, pstats, io

import datetime
from meerkat_abacus import model
from meerkat_abacus import logger

class ProcessingStep(object):
    """
    Base class for all ProcessingSteps

    """

    def __init__(self):
        self.step_name = "processing_step"
        self.session = None
        self.start = None
        self.end = None

    def __repr__(self):
        return f'<{self.__class__}, step_name="{self.step_name}">'

    @property
    def duration(self):
        if not self.start or not self.end:
            return None
        return self.end - self.start

    @abstractmethod
    def run(self, form, data):
        pass

    def start_step(self):
        self.start = datetime.datetime.now()
        # self.profiler = cProfile.Profile()
        # self.profiler.enable()

    def end_step(self, n):
        self.end = datetime.datetime.now()
        self._write_monitoring_data(n)
        # self.profiler.disable()
        s = io.StringIO()
        sortby = 'cumulative'
        # ps = pstats.Stats(self.profiler, stream=s).sort_stats(sortby)
        # ps.print_stats(30)
        # logger.info(s.getvalue())

    def _write_monitoring_data(self, n=None):
        monitoring = model.StepMonitoring(
            step=self.step_name,
            end=self.end,
            start=self.start,
            duration=self.duration.total_seconds(),
            n=n)
        self.session.add(monitoring)
        self.session.commit()

class DoNothing(ProcessingStep):
    def __init__(self, session):
        super().__init__()
        self.step_name = "do_nothing"
        self.session = session

    def run(self, form, data):
        return [{"form": form,
                 "data": data}]
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
# © 2016 ACSONE SA/NV <https://acsone.eu>
# License AGPL-3.0 or later (http://www.gnu.org/licenses/agpl.html).

from openerp import models, api
from openerp import tools
from openerp.tools.translate import _

class MailNotification(models.Model):

    _inherit = 'mail.notification'

    @api.model
    def _get_partner_names(self, partner_ids):
        """
        :type partner_ids: [integer]
        :param partner_ids: ids of the partner followers
        :rparam: list of the partner'name that are also a user or having
        notify_email attribute not none
        """
        partners = self.env['res.partner'].browse(partner_ids)
        partners_name = [
            partner.name for partner in partners if
            partner.user_ids or partner.notify_email != 'none'
        ]
        return partners_name

    @api.model
    def get_signature_footer(
            self, user_id, res_model=None, res_id=None, user_signature=True):
        """
        Override this method to add name of notified partners into the mail
        footer
        """
        footer = super(MailNotification, self).get_signature_footer(
            user_id, res_model=res_model, res_id=res_id,
            user_signature=user_signature)
        partner_ids = self.env.context.get('partners_to_notify')
        if footer and partner_ids:
            partners_name = self._get_partner_names(partner_ids)
            additional_footer = u'
%s%s.
' %\
                (_('Also notified: '),
                 ', '.join(partners_name))
            footer = tools.append_content_to_html(
                additional_footer, footer, plaintext=False,
                container_tag='div')

        return footer

    @api.one
    def _notify(
            self, partners_to_notify=None, force_send=False,
            user_signature=True):
        ctx = self.env.context.copy()
        if not self.env.context.get('mail_notify_noemail'):
            ctx.update({
                'partners_to_notify': partners_to_notify,
            })
        return super(MailNotification, self._model)._notify(
            self.env.cr, self.env.uid, self.id,
            partners_to_notify=partners_to_notify,
            force_send=force_send, user_signature=user_signature, context=ctx)
<commit_msg>[FIX][mail_footer_notified_partners] Use @api.model instead of @api.one.<commit_after># -*- coding: utf-8 -*-
# © 2016 ACSONE SA/NV <https://acsone.eu>
# License AGPL-3.0 or later (http://www.gnu.org/licenses/agpl.html).

from openerp import models, api
from openerp import tools
from openerp.tools.translate import _

class MailNotification(models.Model):

    _inherit = 'mail.notification'

    @api.model
    def _get_partner_names(self, partner_ids):
        """
        :type partner_ids: [integer]
        :param partner_ids: ids of the partner followers
        :rparam: list of the partner'name that are also a user or having
        notify_email attribute not none
        """
        partners = self.env['res.partner'].browse(partner_ids)
        partners_name = [
            partner.name for partner in partners if
            partner.user_ids or partner.notify_email != 'none'
        ]
        return partners_name

    @api.model
    def get_signature_footer(
            self, user_id, res_model=None, res_id=None, user_signature=True):
        """
        Override this method to add name of notified partners into the mail
        footer
        """
        footer = super(MailNotification, self).get_signature_footer(
            user_id, res_model=res_model, res_id=res_id,
            user_signature=user_signature)
        partner_ids = self.env.context.get('partners_to_notify')
        if footer and partner_ids:
            partners_name = self._get_partner_names(partner_ids)
            additional_footer = u'
%s%s.
' %\
                (_('Also notified: '),
                 ', '.join(partners_name))
            footer = tools.append_content_to_html(
                additional_footer, footer, plaintext=False,
                container_tag='div')

        return footer

    @api.model
    def _notify(
            self, message_id, partners_to_notify=None, force_send=False,
            user_signature=True):
        ctx = self.env.context.copy()
        if not self.env.context.get('mail_notify_noemail'):
            ctx.update({
                'partners_to_notify': partners_to_notify,
            })
        return super(MailNotification, self)._notify(
            message_id,
            partners_to_notify=partners_to_notify,
            force_send=force_send, user_signature=user_signature, context=ctx)
<|endoftext|>


<commit_before>#! /usr/bin/env python

import sys
import os
from copy import deepcopy
import rospy

from mdp_plan_exec.top_map_mdp import TopMapMdp
from mdp_plan_exec.prism_java_talker import PrismJavaTalker

from strands_executive_msgs.srv import GetGuaranteesForCoSafeTask, GetGuaranteesForCoSafeTaskResponse

class MdpTaskGuaranteesEstimator(object):

    def __init__(self,top_map):

        self.top_map_mdp=TopMapMdp(top_map, explicit_doors=True)
        self.current_extended_mdp=None
        self.directory = os.path.expanduser("~") + '/tmp/prism/guarantees_estimator/'
        try:
            os.makedirs(self.directory)
        except OSError as ex:
            print 'error creating PRISM directory:',  ex
        self.file_name=top_map+".mdp"
        self.prism_estimator=PrismJavaTalker(8087,self.directory, self.file_name)
        self.last_epoch=-1
        self.get_guarantees_service = rospy.Service('/mdp_plan_exec/get_guarantees_for_co_safe_task',
                                                              GetGuaranteesForCoSafeTask,
                                                              self.get_guarantees_cb)
        rospy.loginfo("MDP task guarantees estimator initialised.")

    def generate_prism_specification(self, ltl_spec):
        return 'partial(R{"time"}min=? [ (' + ltl_spec + ') ])'

    def get_guarantees_cb(self,req):
        self.current_extended_mdp=deepcopy(self.top_map_mdp)
        self.current_extended_mdp.add_extra_domain(req.spec.vars, req.spec.actions)
        if req.epoch != self.last_epoch:
            self.last_epoch=req.epoch
        self.current_extended_mdp.set_mdp_action_durations(self.directory+self.file_name, req.epoch, set_initial_state=False) #TODO we need this here to update the MDP model file, but we shouldn't need to get new action durations to do it when the epoch didnt change.
        specification=self.generate_prism_specification(req.spec.ltl_task)
        rospy.loginfo("The specification is " + specification)
        state_vector=self.prism_estimator.get_state_vector(specification, is_partial=True)
        (waypoint_names, waypoint_probs, waypoint_progs, waypoint_times)=self.current_extended_mdp.map_state_vectors_to_waypoint_expectations(self.directory+'original.sta', state_vector)
        return GetGuaranteesForCoSafeTaskResponse(initial_waypoints=waypoint_names,
                                                  probabilities=waypoint_probs,
                                                  prog_rewards=waypoint_progs,
                                                  expected_times=map(rospy.Duration, waypoint_times))

    def main(self):
       # Wait for control-c
        rospy.spin()
        if rospy.is_shutdown():
            self.prism_estimator.shutdown(True)

if __name__ == '__main__':
    rospy.init_node('mdp_task_guarantees_estimator')

    while not rospy.has_param("/topological_map_name") and not rospy.is_shutdown():
        rospy.sleep(0.1)

    if not rospy.is_shutdown():
        top_map_name=rospy.get_param("/topological_map_name")
        mdp_estimator =  MdpTaskGuaranteesEstimator(top_map_name)
        mdp_estimator.main()

<commit_msg>Added locking to estimator to prevent errors due to concurrent access.<commit_after>#! /usr/bin/env python

from __future__ import with_statement
import sys
import os
from copy import deepcopy
import rospy
import threading
from mdp_plan_exec.top_map_mdp import TopMapMdp
from mdp_plan_exec.prism_java_talker import PrismJavaTalker

from strands_executive_msgs.srv import GetGuaranteesForCoSafeTask, GetGuaranteesForCoSafeTaskResponse

class MdpTaskGuaranteesEstimator(object):

    def __init__(self,top_map):

        self.top_map_mdp=TopMapMdp(top_map, explicit_doors=True)
        self.current_extended_mdp=None
        self.directory = os.path.expanduser("~") + '/tmp/prism/guarantees_estimator/'
        self.service_lock = threading.Lock()
        try:
            os.makedirs(self.directory)
        except OSError as ex:
            print 'error creating PRISM directory:',  ex
        self.file_name=top_map+".mdp"
        self.prism_estimator=PrismJavaTalker(8087,self.directory, self.file_name)
        self.last_epoch=-1
        self.get_guarantees_service = rospy.Service('/mdp_plan_exec/get_guarantees_for_co_safe_task',
                                                              GetGuaranteesForCoSafeTask,
                                                              self.get_guarantees_cb)
        rospy.loginfo("MDP task guarantees estimator initialised.")

    def generate_prism_specification(self, ltl_spec):
        return 'partial(R{"time"}min=? [ (' + ltl_spec + ') ])'

    def get_guarantees_cb(self,req):
        with self.service_lock:
            self.current_extended_mdp=deepcopy(self.top_map_mdp)
            self.current_extended_mdp.add_extra_domain(req.spec.vars, req.spec.actions)
            if req.epoch != self.last_epoch:
                self.last_epoch=req.epoch
            self.current_extended_mdp.set_mdp_action_durations(self.directory+self.file_name, req.epoch, set_initial_state=False) #TODO we need this here to update the MDP model file, but we shouldn't need to get new action durations to do it when the epoch didnt change.
            specification=self.generate_prism_specification(req.spec.ltl_task)
            rospy.loginfo("The specification is " + specification)
            state_vector=self.prism_estimator.get_state_vector(specification, is_partial=True)
            (waypoint_names, waypoint_probs, waypoint_progs, waypoint_times)=self.current_extended_mdp.map_state_vectors_to_waypoint_expectations(self.directory+'original.sta', state_vector)
            return GetGuaranteesForCoSafeTaskResponse(initial_waypoints=waypoint_names,
                                                      probabilities=waypoint_probs,
                                                      prog_rewards=waypoint_progs,
                                                      expected_times=map(rospy.Duration, waypoint_times))

    def main(self):
       # Wait for control-c
        rospy.spin()
        if rospy.is_shutdown():
            self.prism_estimator.shutdown(True)

if __name__ == '__main__':
    rospy.init_node('mdp_task_guarantees_estimator')

    while not rospy.has_param("/topological_map_name") and not rospy.is_shutdown():
        rospy.sleep(0.1)

    if not rospy.is_shutdown():
        top_map_name=rospy.get_param("/topological_map_name")
        mdp_estimator =  MdpTaskGuaranteesEstimator(top_map_name)
        mdp_estimator.main()

<|endoftext|>