<commit_before>from flask_admin import AdminIndexView, BaseView, expose, form
from flask_admin.contrib.sqla import ModelView
from flask_admin.actions import action
from flask_admin.form import rules
from flask import current_app
from flask import session, flash, redirect, url_for, request, send_from_directory
from hashids import Hashids
from models import Houses, Markers, Users
from sqlalchemy.event import listens_for
from wtforms import TextAreaField
from wtforms.widgets import TextArea
import os
import os.path as op

file_path = op.join(op.dirname(__file__), 'static/upload')
try:
    os.mkdir(file_path)
except OSError:
    pass

class MyAdminIndexView(AdminIndexView):
    def is_accessible(self):
        return session.get('ad_login') # This does the trick rendering the view only if the user is authenticated

    def _handle_view(self, name, **kwargs):
        if not self.is_accessible():
            flash("You must be logged in to access this feature")
            return redirect(url_for('admin_login', next=request.url))

class QRGenView(BaseView):
    @expose('/')
    def index(self):
        return self.render('admin/index.html')

    def gen_zip(self, markers):
        import pyqrcode
        from zipfile import ZipFile
        import io, os

        ## This ensures that it is in the package directory rather than just cwd
        __location__ = self.get_location()
        print(__location__)
        zipper = ZipFile(os.path.join(__location__, 'markers.zip'), 'w')
        hashid = Hashids(min_length=6, salt=current_app.config['HASHID_KEY'])
        for i in range(0, len(markers)):
            marker = markers[i]
            marker.url = url_for('scan_marker', scan_id=hashid.encode(marker.id), _external=True)
            buffer = io.BytesIO()
            qrcode = pyqrcode.create(marker.url)
            qrcode.svg(buffer, scale=5, background="white")
            zipper.writestr(str(marker.id)+"-"+marker.name+"-"+str(marker.houses.name)+".svg", buffer.getvalue())
        zipper.close()

    @expose('/gen', methods=['POST'])
    def generate_zip(self):
        markers = Markers.query.all()
        self.gen_zip(markers)
        flash("Currently generating .zip in the background, please wait a few minutes")
        return redirect(url_for('admin.index'))

    @expose('/download', methods=['POST'])
    def download_zip(self):
        try:

            __location__ = self.get_location()
            print(__location__)
            return send_from_directory(__location__, 'markers.zip', as_attachment=True)
        except:
            flash("Please generate the .zip first")
            return redirect(url_for('admin.index'))

    def get_location(self):
        import os
        return os.path.realpath(os.path.join(os.getcwd(), os.path.dirname(__file__)))

class MarkerView(ModelView):
    page_size = 50
    column_searchable_list = ['name', 'batch', 'location']
    list_template = 'admin/marker_list.html'

    @action('set_on', 'Activate', 'Are you sure you want to activate all of these markers?')
    def action_activate(self, ids):
        try:
            query = Markers.query.filter(Markers.id.in_(ids))

            count = 0
            for marker in query.all():
                marker.in_current_use = 1

            flash("Markers set to active")
            db.session.commit()
        except Exception as ex:
            if not self.handle_view_exception(ex):
                raise
            print(ex)
            flash("Failed to update database")

    @action('set_off', 'Deactivate', 'Are you sure you want to set all of these makers to inactive?')
    def action_inactivate(self, ids):
        try:
            query = Markers.query.filter(Markers.id.in_(ids))

            count = 0
            for marker in query.all():
                marker.in_current_use = 0

            flash("Markers set to inactive")
            db.session.commit()
        except Exception as ex:
            if not self.handle_view_exception(ex):
                raise
            print(ex)
            flash("Failed to update database")

    @action('hide', 'Set Hidden', 'Are you sure you want to hide the selected markers from the torch registry page?')
    def action_hide(self, ids):
        try:
            query = Markers.query.filter(Markers.id.in_(ids))

            count = 0
            for marker in query.all():
                marker.is_hidden = 1

            flash("Markers set to hidden")
            db.session.commit()
        except Exception as ex:
            if not self.handle_view_exception(ex):
                raise
            print(ex)
            flash("Failed to update database")

    @action('show', 'Set Visible', 'Are you sure you want to include the selected markers in the torch registry page?')
    def action_make_visble(self, ids):
        try:
            query = Markers.query.filter(Markers.id.in_(ids))

            count = 0
            for marker in query.all():
                marker.is_hidden = 0

            flash("Markers set to hidden")
            db.session.commit()
        except Exception as ex:
            if not self.handle_view_exception(ex):
                raise
            print(ex)
            flash("Failed to update database")

    @expose('/mass')
    def mass_view(self):
        """
            Custom view for mass changing.
        """
        return self.render('admin/marker_mass.html')

    @expose('/mass/can_scan', methods=['POST'])
    def set_all_on(self):

        try:
            for marker in Markers.query.all():
                marker.in_current_use = 1

            flash("Markers set to active")
            db.session.commit()
        except Exception as ex:
            if not self.handle_view_exception(ex):
                raise
            print(ex)
            flash("Failed to update database")

        return redirect(url_for('.mass_view'))

    @expose('/mass/no_scan',  methods=['POST'])
    def set_all_off(self):

        try:
            for marker in Markers.query.all():
                marker.in_current_use = 0

            flash("Markers set to inactive")
            db.session.commit()
        except Exception as ex:
            if not self.handle_view_exception(ex):
                raise
            print(ex)
            flash("Failed to update database")

        return redirect(url_for('.mass_view'))

    @expose('/mass/hide', methods=['POST'])
    def set_all_hidden(self):
        try:
            for marker in Markers.query.all():
                marker.is_hidden = 1

            flash("Markers set to hidden")
            db.session.commit()
        except Exception as ex:
            if not self.handle_view_exception(ex):
                raise
            print(ex)
            flash("Failed to update database")

        return redirect(url_for('.mass_view'))

    @expose('/mass/visble', methods=['POST'])
    def set_all_visible(self):
        try:
            for marker in Markers.query.all():
                marker.is_hidden = 0

            flash("Markers set to hidden")
            db.session.commit()
        except Exception as ex:
            if not self.handle_view_exception(ex):
                raise
            print(ex)
            flash("Failed to update database")

        return redirect(url_for('.mass_view'))

class UserView(ModelView):
    column_exclude_list = ['pwhash']

class APageView(ModelView):
    form_overrides = {
        'content': TextAreaField
    }

class HouseView(ModelView):
    form_overrides = {
        'imagepath': form.FileUploadField
    }

    # Pass additional parameters to 'path' to FileUploadField constructor
    form_args = {
        'imagepath': {
            'label': 'File',
            'base_path': file_path,
            'allow_overwrite': False
        }
    }

@listens_for(Houses, 'after_delete')
def del_image(mapper, connection, target):
    if target.imagepath:
        # Delete image
        try:
            os.remove(op.join(file_path, target.path))
        except OSError:
            pass

        # Delete thumbnail
        try:
            os.remove(op.join(file_path,
                              form.thumbgen_filename(target.path)))
        except OSError:
            pass
<commit_msg>Modified code generation for markers with no house<commit_after>from flask_admin import AdminIndexView, BaseView, expose, form
from flask_admin.contrib.sqla import ModelView
from flask_admin.actions import action
from flask_admin.form import rules
from flask import current_app
from flask import session, flash, redirect, url_for, request, send_from_directory
from hashids import Hashids
from models import Houses, Markers, Users
from sqlalchemy.event import listens_for
from wtforms import TextAreaField
from wtforms.widgets import TextArea
import os
import os.path as op

file_path = op.join(op.dirname(__file__), 'static/upload')
try:
    os.mkdir(file_path)
except OSError:
    pass

class MyAdminIndexView(AdminIndexView):
    def is_accessible(self):
        return session.get('ad_login') # This does the trick rendering the view only if the user is authenticated

    def _handle_view(self, name, **kwargs):
        if not self.is_accessible():
            flash("You must be logged in to access this feature")
            return redirect(url_for('admin_login', next=request.url))

class QRGenView(BaseView):
    @expose('/')
    def index(self):
        return self.render('admin/index.html')

    def gen_zip(self, markers):
        import pyqrcode
        from zipfile import ZipFile
        import io, os

        ## This ensures that it is in the package directory rather than just cwd
        __location__ = self.get_location()
        print(__location__)
        zipper = ZipFile(os.path.join(__location__, 'markers.zip'), 'w')
        hashid = Hashids(min_length=6, salt=current_app.config['HASHID_KEY'])
        for i in range(0, len(markers)):
            marker = markers[i]
            marker.url = url_for('scan_marker', scan_id=hashid.encode(marker.id), _external=True)
            buffer = io.BytesIO()
            qrcode = pyqrcode.create(marker.url)
            qrcode.svg(buffer, scale=5, background="white")
            housename = "nohouse"
            try:
                housename = marker.houses.name
            except:
                pass
            zipper.writestr(str(marker.id)+"-"+marker.name+"-"+housename+".svg", buffer.getvalue())
        zipper.close()

    @expose('/gen', methods=['POST'])
    def generate_zip(self):
        markers = Markers.query.all()
        self.gen_zip(markers)
        flash("Currently generating .zip in the background, please wait a few minutes")
        return redirect(url_for('admin.index'))

    @expose('/download', methods=['POST'])
    def download_zip(self):
        try:

            __location__ = self.get_location()
            print(__location__)
            return send_from_directory(__location__, 'markers.zip', as_attachment=True)
        except:
            flash("Please generate the .zip first")
            return redirect(url_for('admin.index'))

    def get_location(self):
        import os
        return os.path.realpath(os.path.join(os.getcwd(), os.path.dirname(__file__)))

class MarkerView(ModelView):
    page_size = 50
    column_searchable_list = ['name', 'batch', 'location']
    list_template = 'admin/marker_list.html'

    @action('set_on', 'Activate', 'Are you sure you want to activate all of these markers?')
    def action_activate(self, ids):
        try:
            query = Markers.query.filter(Markers.id.in_(ids))

            count = 0
            for marker in query.all():
                marker.in_current_use = 1

            flash("Markers set to active")
            db.session.commit()
        except Exception as ex:
            if not self.handle_view_exception(ex):
                raise
            print(ex)
            flash("Failed to update database")

    @action('set_off', 'Deactivate', 'Are you sure you want to set all of these makers to inactive?')
    def action_inactivate(self, ids):
        try:
            query = Markers.query.filter(Markers.id.in_(ids))

            count = 0
            for marker in query.all():
                marker.in_current_use = 0

            flash("Markers set to inactive")
            db.session.commit()
        except Exception as ex:
            if not self.handle_view_exception(ex):
                raise
            print(ex)
            flash("Failed to update database")

    @action('hide', 'Set Hidden', 'Are you sure you want to hide the selected markers from the torch registry page?')
    def action_hide(self, ids):
        try:
            query = Markers.query.filter(Markers.id.in_(ids))

            count = 0
            for marker in query.all():
                marker.is_hidden = 1

            flash("Markers set to hidden")
            db.session.commit()
        except Exception as ex:
            if not self.handle_view_exception(ex):
                raise
            print(ex)
            flash("Failed to update database")

    @action('show', 'Set Visible', 'Are you sure you want to include the selected markers in the torch registry page?')
    def action_make_visble(self, ids):
        try:
            query = Markers.query.filter(Markers.id.in_(ids))

            count = 0
            for marker in query.all():
                marker.is_hidden = 0

            flash("Markers set to hidden")
            db.session.commit()
        except Exception as ex:
            if not self.handle_view_exception(ex):
                raise
            print(ex)
            flash("Failed to update database")

    @expose('/mass')
    def mass_view(self):
        """
            Custom view for mass changing.
        """
        return self.render('admin/marker_mass.html')

    @expose('/mass/can_scan', methods=['POST'])
    def set_all_on(self):

        try:
            for marker in Markers.query.all():
                marker.in_current_use = 1

            flash("Markers set to active")
            db.session.commit()
        except Exception as ex:
            if not self.handle_view_exception(ex):
                raise
            print(ex)
            flash("Failed to update database")

        return redirect(url_for('.mass_view'))

    @expose('/mass/no_scan',  methods=['POST'])
    def set_all_off(self):

        try:
            for marker in Markers.query.all():
                marker.in_current_use = 0

            flash("Markers set to inactive")
            db.session.commit()
        except Exception as ex:
            if not self.handle_view_exception(ex):
                raise
            print(ex)
            flash("Failed to update database")

        return redirect(url_for('.mass_view'))

    @expose('/mass/hide', methods=['POST'])
    def set_all_hidden(self):
        try:
            for marker in Markers.query.all():
                marker.is_hidden = 1

            flash("Markers set to hidden")
            db.session.commit()
        except Exception as ex:
            if not self.handle_view_exception(ex):
                raise
            print(ex)
            flash("Failed to update database")

        return redirect(url_for('.mass_view'))

    @expose('/mass/visble', methods=['POST'])
    def set_all_visible(self):
        try:
            for marker in Markers.query.all():
                marker.is_hidden = 0

            flash("Markers set to hidden")
            db.session.commit()
        except Exception as ex:
            if not self.handle_view_exception(ex):
                raise
            print(ex)
            flash("Failed to update database")

        return redirect(url_for('.mass_view'))

class UserView(ModelView):
    column_exclude_list = ['pwhash']

class APageView(ModelView):
    form_overrides = {
        'content': TextAreaField
    }

class HouseView(ModelView):
    form_overrides = {
        'imagepath': form.FileUploadField
    }

    # Pass additional parameters to 'path' to FileUploadField constructor
    form_args = {
        'imagepath': {
            'label': 'File',
            'base_path': file_path,
            'allow_overwrite': False
        }
    }

@listens_for(Houses, 'after_delete')
def del_image(mapper, connection, target):
    if target.imagepath:
        # Delete image
        try:
            os.remove(op.join(file_path, target.path))
        except OSError:
            pass

        # Delete thumbnail
        try:
            os.remove(op.join(file_path,
                              form.thumbgen_filename(target.path)))
        except OSError:
            pass
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
'''
Module for running arbitrary tests
'''

# Import Python libs
import os
import sys
import time
import traceback
import random

# Import Salt libs
import salt
import salt.version
import salt.loader

__proxyenabled__ = ['*']

try:
    # ...a hack used to ignore this class when this is used as a salt
    # execution module
    #
    from mapper import *

    class _mapper(mapper):

        sum = 0

        class partializer():
            part_size = 1000

            def __init__(self, upper):
                self.upper = upper
                self.x = 0

            def next(self):
                ret = self.x
                if ret > self.upper:
                    ret = self.upper - (ret - self.part_size)
                if ret <= 0 and self.x > 0:
                    raise StopIteration

                self.x += self.part_size
                return [ret, self.part_size]

            def __iter__(self):
                return self

    def reducer(self, n):
        self.sum += n
        return self.sum

    def statit(self):
        print "sum = " + self.sum

except:
    pass

def echo(*args):
    return args

def fib(num):
    '''
    Return a Fibonacci sequence up to the passed number, and the
    timeit took to compute in seconds. Used for performance tests

    CLI Example:

    .. code-block:: bash

        salt '*' test.fib 3
    '''
    num = int(num)
    start = time.time()
    fib_a, fib_b = 0, 1
    ret = [0]
    while fib_b < num:
        ret.append(fib_b)
        fib_a, fib_b = fib_b, fib_a + fib_b
    return ret, time.time() - start

def sum_nums(upper):
    '''
    Return the sum of the sequence of numbers from 1 to [upper], and the
    time it took to compute in seconds. Useful for validating the mapreduce runner

    CLI Example:

    .. code-block:: bash

        salt '*' mapit.sum_nums 10

    '''
    upper = int(upper)
    start = time.time()
    num = 0
    sum = 0

    while num < upper:
        num += 1
        sum += num
    print "sum = " + str(sum)
    return sum

def partial_result(lower, count):
    '''
    Return the sum of the sequence of numbers in the range [lower, upper], and the
    time it took to compute in seconds. Useful for validating the mapreduce runner

    CLI Example:

    .. code-block:: bash

        salt '*' mapit.sum_nums_partial 10 20

    '''
    # lower = int(lower)
    # num = lower
    # sum = num
    #
    # for a in xrange(0, count):
    #     num += 1
    #     sum += num
    #
    # return sum
    return lower, count

<commit_msg>fix to partializer.next() StopIteration condition<commit_after># -*- coding: utf-8 -*-
'''
Module for running arbitrary tests
'''

# Import Python libs
import os
import sys
import time
import traceback
import random

# Import Salt libs
import salt
import salt.version
import salt.loader

__proxyenabled__ = ['*']

try:
    # ...a hack used to ignore this class when this is used as a salt
    # execution module
    #
    from mapper import *

    class _mapper(mapper):

        sum = 0

        class partializer():
            part_size = 1000

            def __init__(self, upper):
                self.upper = upper[0]
                self.x = 0

            def next(self):
                ret = self.x
                if ret >= self.upper:
                    raise StopIteration
                remain = self.upper - ret
                if remain >= self.part_size:
                    remain = self.part_size
                # elif remain > 0:
                #     pass
                self.x += remain
                return [ret, remain]

            def __iter__(self):
                return self

    def reducer(self, n):
        self.sum += n
        return self.sum

    def statit(self):
        print "sum = " + self.sum

except:
    pass

def echo(*args):
    return args

def fib(num):
    '''
    Return a Fibonacci sequence up to the passed number, and the
    timeit took to compute in seconds. Used for performance tests

    CLI Example:

    .. code-block:: bash

        salt '*' test.fib 3
    '''
    num = int(num)
    start = time.time()
    fib_a, fib_b = 0, 1
    ret = [0]
    while fib_b < num:
        ret.append(fib_b)
        fib_a, fib_b = fib_b, fib_a + fib_b
    return ret, time.time() - start

def sum_nums(upper):
    '''
    Return the sum of the sequence of numbers from 1 to [upper], and the
    time it took to compute in seconds. Useful for validating the mapreduce runner

    CLI Example:

    .. code-block:: bash

        salt '*' mapit.sum_nums 10

    '''
    upper = int(upper)
    start = time.time()
    num = 0
    sum = 0

    while num < upper:
        num += 1
        sum += num
    print "sum = " + str(sum)
    return sum

def partial_result(lower, count):
    '''
    Return the sum of the sequence of numbers in the range [lower, upper], and the
    time it took to compute in seconds. Useful for validating the mapreduce runner

    CLI Example:

    .. code-block:: bash

        salt '*' mapit.sum_nums_partial 10 20

    '''
    # lower = int(lower)
    # num = lower
    # sum = num
    #
    # for a in xrange(0, count):
    #     num += 1
    #     sum += num
    #
    # return sum
    return lower, count

<|endoftext|>


<commit_before>from gooey.gui.components.widgets.bases import TextContainer
import wx

from gooey.gui import formatters
from gooey.gui.lang.i18n import _

class Dropdown(TextContainer):

    def getWidget(self, parent, *args, **options):
        default = _('select_option')
        return wx.ComboBox(
            parent=parent,
            id=-1,
            value=default,
            choices=[default] + self._meta['choices'],
            style=wx.CB_DROPDOWN)

    def setOptions(self, options):
        prevSelection = self.widget.GetSelection()
        self.widget.Clear()
        for option in [_('select_option')] + options:
            self.widget.Append(option)
        self.widget.SetSelection(0)

    def setValue(self, value):
        ## +1 to offset the default placeholder value
        index = self._meta['choices'].index(value) + 1
        self.widget.SetSelection(index)

    def getWidgetValue(self):
        value = self.widget.GetValue()
        # filter out the extra default option that's
        # appended during creation
        if value == _('select_option'):
            return None
        return value

    def formatOutput(self, metadata, value):
        return formatters.dropdown(metadata, value)
<commit_msg>Added a cast to string for the Dropdown default and choices values.<commit_after>from gooey.gui.components.widgets.bases import TextContainer
import wx

from gooey.gui import formatters
from gooey.gui.lang.i18n import _

class Dropdown(TextContainer):

    def getWidget(self, parent, *args, **options):
        default = _('select_option')
        return wx.ComboBox(
            parent=parent,
            id=-1,
            # str conversion allows using stringyfiable values in addition to pure strings
            value=str(default),
            choices=[str(default)] + [str(choice) for choice in self._meta['choices']],
            style=wx.CB_DROPDOWN)

    def setOptions(self, options):
        prevSelection = self.widget.GetSelection()
        self.widget.Clear()
        for option in [_('select_option')] + options:
            self.widget.Append(option)
        self.widget.SetSelection(0)

    def setValue(self, value):
        ## +1 to offset the default placeholder value
        index = self._meta['choices'].index(value) + 1
        self.widget.SetSelection(index)

    def getWidgetValue(self):
        value = self.widget.GetValue()
        # filter out the extra default option that's
        # appended during creation
        if value == _('select_option'):
            return None
        return value

    def formatOutput(self, metadata, value):
        return formatters.dropdown(metadata, value)
<|endoftext|>


<commit_before>from datetime import timedelta
import logging
from sqlalchemy import select, and_
import sys

from hades.common.db import (get_connection, radacct, radpostauth, utcnow)
from hades.config.loader import get_config

logger = logging.getLogger(__name__)

def refresh():
    logger.info("Refreshing materialized views")
    connection = get_connection()
    with connection.begin() as trans:
        # TODO: After updating the nas table, we have to restart (reload?)
        # the freeradius server. Currently, this must be done manually.
        connection.execute("REFRESH MATERIALIZED VIEW nas")
        connection.execute("REFRESH MATERIALIZED VIEW radcheck")
        connection.execute("REFRESH MATERIALIZED VIEW radgroupcheck")
        connection.execute("REFRESH MATERIALIZED VIEW radgroupreply")
        connection.execute("REFRESH MATERIALIZED VIEW radusergroup")

def delete_old():
    logger.info("Deleting old records")
    conf = get_config()
    connection = get_connection()
    with connection.begin() as trans:
        result = connection.execute(radacct.delete().where(and_(
            radacct.c.acctstoptime < utcnow() - conf["HADES_RETENTION_INTERVAL"]
        )))
        result = connection.execute(radpostauth.delete().where(and_(
            radpostauth.c.authdate < utcnow() - timedelta(days=1)
        )))

def get_sessions(self, mac):
    connection = get_connection()
    results = connection.execute(
        select([radacct.c.nasipaddress, radacct.c.nasportid,
                radacct.c.acctstarttime, radacct.c.acctstoptime,
                radacct.c.acctstartdelay, radacct.c.acctstopdelay])
        .where(and_(radacct.c.username == mac,
                    radacct.c.acctstarttime >= utcnow() - timedelta(days=1))))
    return results.fetchall()

if __name__ == '__main__':
    logger.info('Hades Agent was invoked')
    failure = False

    try:
        refresh()
    except:
        failure = True
        logger.exception('Error while refreshing materialized views')

    try:
        delete_old()
    except:
        failure = True
        logger.exception('Error while deleting old records')

    if failure:
        sys.exit(1)
<commit_msg>Refresh dhcphost materialized view<commit_after>from datetime import timedelta
import logging
from sqlalchemy import select, and_
import sys

from hades.common.db import (get_connection, radacct, radpostauth, utcnow)
from hades.config.loader import get_config

logger = logging.getLogger(__name__)

def refresh():
    logger.info("Refreshing materialized views")
    connection = get_connection()
    with connection.begin() as trans:
        connection.execute("REFRESH MATERIALIZED VIEW dhcphost")
        # TODO: After updating the nas table, we have to restart (reload?)
        # the freeradius server. Currently, this must be done manually.
        connection.execute("REFRESH MATERIALIZED VIEW nas")
        connection.execute("REFRESH MATERIALIZED VIEW radcheck")
        connection.execute("REFRESH MATERIALIZED VIEW radgroupcheck")
        connection.execute("REFRESH MATERIALIZED VIEW radgroupreply")
        connection.execute("REFRESH MATERIALIZED VIEW radusergroup")

def delete_old():
    logger.info("Deleting old records")
    conf = get_config()
    connection = get_connection()
    with connection.begin() as trans:
        result = connection.execute(radacct.delete().where(and_(
            radacct.c.acctstoptime < utcnow() - conf["HADES_RETENTION_INTERVAL"]
        )))
        result = connection.execute(radpostauth.delete().where(and_(
            radpostauth.c.authdate < utcnow() - timedelta(days=1)
        )))

def get_sessions(self, mac):
    connection = get_connection()
    results = connection.execute(
        select([radacct.c.nasipaddress, radacct.c.nasportid,
                radacct.c.acctstarttime, radacct.c.acctstoptime,
                radacct.c.acctstartdelay, radacct.c.acctstopdelay])
        .where(and_(radacct.c.username == mac,
                    radacct.c.acctstarttime >= utcnow() - timedelta(days=1))))
    return results.fetchall()

if __name__ == '__main__':
    logger.info('Hades Agent was invoked')
    failure = False

    try:
        refresh()
    except:
        failure = True
        logger.exception('Error while refreshing materialized views')

    try:
        delete_old()
    except:
        failure = True
        logger.exception('Error while deleting old records')

    if failure:
        sys.exit(1)
<|endoftext|>


<commit_before>from __future__ import absolute_import
from collections import namedtuple
from datetime import datetime, timedelta
import pytz
from casexml.apps.case.const import CASE_ACTION_CLOSE
from casexml.apps.case.dbaccessors import get_open_case_docs_in_domain
from casexml.apps.case.mock import CaseBlock
import uuid
from xml.etree import ElementTree
from corehq.apps.app_manager.const import USERCASE_TYPE
from corehq.apps.domain.models import Domain
from corehq.apps.es.domains import DomainES
from corehq.apps.es import filters
from corehq.apps.hqcase.utils import submit_case_blocks, get_case_by_domain_hq_user_id
from corehq.apps.locations.models import SQLLocation
from corehq.feature_previews import CALLCENTER
from corehq.util.quickcache import quickcache
from corehq.util.timezones.conversions import UserTime, ServerTime
from dimagi.utils.couch import CriticalSection

class DomainLite(namedtuple('DomainLite', 'name default_timezone cc_case_type use_fixtures')):
    def midnights(self, utcnow=None):
        """Returns a list containing two datetimes in UTC that corresponds to midnight
        in the domains timezone on either side of the current UTC datetime.
        i.e. [<previous midnight in TZ>, <next midnight in TZ>]

        >>> d = DomainLite('', 'Asia/Kolkata', '', True)
        >>> d.midnights(datetime(2015, 8, 27, 18, 30, 0  ))
        [datetime.datetime(2015, 8, 26, 18, 30), datetime.datetime(2015, 8, 27, 18, 30)]
        >>> d.midnights(datetime(2015, 8, 27, 18, 31, 0  ))
        [datetime.datetime(2015, 8, 27, 18, 30), datetime.datetime(2015, 8, 28, 18, 30)]
        """
        utcnow = utcnow or datetime.utcnow()
        tz = pytz.timezone(self.default_timezone)
        current_time_tz = ServerTime(utcnow).user_time(tz).done()
        midnight_tz1 = current_time_tz.replace(hour=0, minute=0, second=0, microsecond=0)
        midnight_tz_utc1 = UserTime(midnight_tz1).server_time().done()
        midnight_tz_utc2 = midnight_tz_utc1 + timedelta(days=(1 if midnight_tz_utc1 < utcnow else -1))
        return sorted([midnight_tz_utc1, midnight_tz_utc2])

CallCenterCase = namedtuple('CallCenterCase', 'case_id hq_user_id')

def sync_user_case(commcare_user, case_type, owner_id, case=None):
    """
    Each time a CommCareUser is saved this method gets called and creates or updates
    a case associated with the user with the user's details.

    This is also called to create user cases when the usercase is used for the
    first time.
    """
    with CriticalSection(['user_case_%s_for_%s' % (case_type, commcare_user._id)]):
        domain = commcare_user.project
        fields = _get_user_case_fields(commcare_user)
        case = case or get_case_by_domain_hq_user_id(domain.name, commcare_user._id, case_type)
        close = commcare_user.to_be_deleted() or not commcare_user.is_active

        def case_should_be_reopened(case, user_case_should_be_closed):
            return case and case.closed and not user_case_should_be_closed

        if case_should_be_reopened(case, close):
            _re_open_case(case)
            return

        if not case:
            _create_user_case(case_type, commcare_user, domain, fields, owner_id)
        else:
            changed = _user_case_changed(case, case_type, close, fields, owner_id)
            if changed:
                _update_user_case(case, case_type, domain, fields, owner_id)
            if close and not case.closed:
                _close_user_case(case, case_type, domain, owner_id)

def _create_user_case(case_type, commcare_user, domain, fields, owner_id):
    fields['hq_user_id'] = commcare_user._id
    caseblock = CaseBlock(
        create=True,
        case_id=uuid.uuid4().hex,
        owner_id=owner_id,
        user_id=owner_id,
        case_type=case_type,
        update=fields
    )
    _submit_case_block(caseblock, domain)

def _update_user_case(case, case_type, domain, fields, owner_id):
    caseblock = CaseBlock(
        create=False,
        case_id=case._id,
        owner_id=owner_id,
        case_type=case_type,
        close=False,
        update=fields
    )
    _submit_case_block(caseblock, domain)

def _close_user_case(case, case_type, domain, owner_id):
    caseblock = CaseBlock(
        create=False,
        case_id=case._id,
        owner_id=owner_id,
        case_type=case_type,
        close=True,
    )
    _submit_case_block(caseblock, domain)

def _submit_case_block(caseblock, domain_obj):
    casexml = ElementTree.tostring(caseblock.as_xml())
    submit_case_blocks(casexml, domain_obj.name)

def _get_user_case_fields(commcare_user):

    def valid_element_name(name):
        try:
            ElementTree.fromstring('<{}/>'.format(name))
            return True
        except ElementTree.ParseError:
            return False

    # remove any keys that aren't valid XML element names
    fields = {k: v for k, v in commcare_user.user_data.items() if
              valid_element_name(k)}
    # language or phone_number can be null and will break
    # case submission
    fields.update({
        'name': commcare_user.name or commcare_user.raw_username,
        'username': commcare_user.raw_username,
        'email': commcare_user.email,
        'language': commcare_user.language or '',
        'phone_number': commcare_user.phone_number or ''
    })
    return fields

def _re_open_case(case):
    closing_action = None
    for action in reversed(case.actions):
        if action.action_type == CASE_ACTION_CLOSE:
            closing_action = action
            break
    if closing_action:
        closing_action.xform.archive()

def _user_case_changed(case, case_type, close, fields, owner_id):
    props = case.dynamic_case_properties()
    changed = close != case.closed
    changed = changed or case.type != case_type
    changed = changed or case.name != fields['name']
    changed = changed or case.owner_id != owner_id
    if not changed:
        for field, value in fields.items():
            if field != 'name' and props.get(field) != value:
                changed = True
                break
    return changed

def sync_call_center_user_case(user):
    domain = user.project
    if domain and domain.call_center_config.enabled:
        case, owner_id = _get_call_center_case_and_owner(user, domain)
        sync_user_case(user, domain.call_center_config.case_type, owner_id, case)

CallCenterCaseAndOwner = namedtuple('CallCenterCaseAndOwner', 'case owner_id')

def _get_call_center_case_and_owner(user, domain):
    """
    Return the appropriate owner id for the given users call center case.
    """
    case = get_case_by_domain_hq_user_id(user.project.name, user._id, domain.call_center_config.case_type)
    if domain.call_center_config.use_user_location_as_owner:
        owner_id = call_center_location_owner(user, domain.call_center_config.user_location_ancestor_level)
    elif case and case.owner_id:
        owner_id = case.owner_id
    else:
        owner_id = domain.call_center_config.case_owner_id
    return CallCenterCaseAndOwner(case, owner_id)

def call_center_location_owner(user, ancestor_level):
    if ancestor_level == 0:
        owner_id = user.location_id
    else:
        location = SQLLocation.objects.get(location_id=user.location_id)
        ancestors = location.get_ancestors(ascending=True, include_self=True)
        try:
            owner_id = ancestors[ancestor_level].location_id
        except IndexError:
            owner_id = ancestors[-1].location_id
    return owner_id

def sync_usercase(user):
    domain = user.project
    if domain and domain.usercase_enabled:
        sync_user_case(
            user,
            USERCASE_TYPE,
            user.get_id
        )

def is_midnight_for_domain(midnight_form_domain, error_margin=15, current_time=None):
    current_time = current_time or datetime.utcnow()
    diff = current_time - midnight_form_domain
    return diff.days >= 0 and diff < timedelta(minutes=error_margin)

def get_call_center_domains():
    result = (
        DomainES()
            .is_active()
            .is_snapshot(False)
            .filter(filters.term('call_center_config.enabled', True))
            .fields(['name', 'default_timezone', 'call_center_config.case_type', 'call_center_config.use_fixtures'])
            .run()
    )

    def to_domain_lite(hit):
        return DomainLite(
            name=hit['name'],
            default_timezone=hit['default_timezone'],
            cc_case_type=hit.get('call_center_config.case_type', ''),
            use_fixtures=hit.get('call_center_config.use_fixtures', True)
        )
    return [to_domain_lite(hit) for hit in result.hits]

def get_call_center_cases(domain_name, case_type, user=None):
    all_cases = []

    if user:
        docs = (doc for owner_id in user.get_owner_ids()
                for doc in get_open_case_docs_in_domain(domain_name, case_type,
                                                        owner_id=owner_id))
    else:
        docs = get_open_case_docs_in_domain(domain_name, case_type)

    for case_doc in docs:
        hq_user_id = case_doc.get('hq_user_id', None)
        if hq_user_id:
            all_cases.append(CallCenterCase(
                case_id=case_doc['_id'],
                hq_user_id=hq_user_id
            ))
    return all_cases

@quickcache(['domain'])
def get_call_center_case_type_if_enabled(domain):
    if CALLCENTER.enabled(domain):
        return Domain.get_by_name(domain).call_center_config.case_type
<commit_msg>[refactor] Group methods in _UserCaseHelper class<commit_after>from __future__ import absolute_import
from collections import namedtuple
from datetime import datetime, timedelta
import pytz
from casexml.apps.case.const import CASE_ACTION_CLOSE
from casexml.apps.case.dbaccessors import get_open_case_docs_in_domain
from casexml.apps.case.mock import CaseBlock
import uuid
from xml.etree import ElementTree
from corehq.apps.app_manager.const import USERCASE_TYPE
from corehq.apps.domain.models import Domain
from corehq.apps.es.domains import DomainES
from corehq.apps.es import filters
from corehq.apps.hqcase.utils import submit_case_blocks, get_case_by_domain_hq_user_id
from corehq.apps.locations.models import SQLLocation
from corehq.feature_previews import CALLCENTER
from corehq.util.quickcache import quickcache
from corehq.util.timezones.conversions import UserTime, ServerTime
from dimagi.utils.couch import CriticalSection

class DomainLite(namedtuple('DomainLite', 'name default_timezone cc_case_type use_fixtures')):
    def midnights(self, utcnow=None):
        """Returns a list containing two datetimes in UTC that corresponds to midnight
        in the domains timezone on either side of the current UTC datetime.
        i.e. [<previous midnight in TZ>, <next midnight in TZ>]

        >>> d = DomainLite('', 'Asia/Kolkata', '', True)
        >>> d.midnights(datetime(2015, 8, 27, 18, 30, 0  ))
        [datetime.datetime(2015, 8, 26, 18, 30), datetime.datetime(2015, 8, 27, 18, 30)]
        >>> d.midnights(datetime(2015, 8, 27, 18, 31, 0  ))
        [datetime.datetime(2015, 8, 27, 18, 30), datetime.datetime(2015, 8, 28, 18, 30)]
        """
        utcnow = utcnow or datetime.utcnow()
        tz = pytz.timezone(self.default_timezone)
        current_time_tz = ServerTime(utcnow).user_time(tz).done()
        midnight_tz1 = current_time_tz.replace(hour=0, minute=0, second=0, microsecond=0)
        midnight_tz_utc1 = UserTime(midnight_tz1).server_time().done()
        midnight_tz_utc2 = midnight_tz_utc1 + timedelta(days=(1 if midnight_tz_utc1 < utcnow else -1))
        return sorted([midnight_tz_utc1, midnight_tz_utc2])

class _UserCaseHelper(object):

    def __init__(self, domain, owner_id):
        self.domain = domain
        self.owner_id = owner_id

    def _submit_case_block(self, caseblock):
        casexml = ElementTree.tostring(caseblock.as_xml())
        submit_case_blocks(casexml, self.domain.name)

    @staticmethod
    def re_open_case(case):
        closing_action = None
        for action in reversed(case.actions):
            if action.action_type == CASE_ACTION_CLOSE:
                closing_action = action
                break
        if closing_action:
            closing_action.xform.archive()

    def create_user_case(self, case_type, commcare_user, fields):
        fields['hq_user_id'] = commcare_user._id
        caseblock = CaseBlock(
            create=True,
            case_id=uuid.uuid4().hex,
            owner_id=self.owner_id,
            user_id=self.owner_id,
            case_type=case_type,
            update=fields
        )
        self._submit_case_block(caseblock)

    def update_user_case(self, case, case_type, fields):
        caseblock = CaseBlock(
            create=False,
            case_id=case._id,
            owner_id=self.owner_id,
            case_type=case_type,
            close=False,
            update=fields
        )
        self._submit_case_block(caseblock)

    def close_user_case(self, case, case_type):
        caseblock = CaseBlock(
            create=False,
            case_id=case._id,
            owner_id=self.owner_id,
            case_type=case_type,
            close=True,
        )
        self._submit_case_block(caseblock)

CallCenterCase = namedtuple('CallCenterCase', 'case_id hq_user_id')

def sync_user_case(commcare_user, case_type, owner_id, case=None):
    """
    Each time a CommCareUser is saved this method gets called and creates or updates
    a case associated with the user with the user's details.

    This is also called to create user cases when the usercase is used for the
    first time.
    """
    with CriticalSection(['user_case_%s_for_%s' % (case_type, commcare_user._id)]):
        domain = commcare_user.project
        fields = _get_user_case_fields(commcare_user)
        case = case or get_case_by_domain_hq_user_id(domain.name, commcare_user._id, case_type)
        close = commcare_user.to_be_deleted() or not commcare_user.is_active
        user_case_helper = _UserCaseHelper(domain, owner_id)

        def case_should_be_reopened(case, user_case_should_be_closed):
            return case and case.closed and not user_case_should_be_closed

        if case_should_be_reopened(case, close):
            user_case_helper.re_open_case(case)
            return

        if not case:
            user_case_helper.create_user_case(case_type, commcare_user, fields)
        else:
            changed = _user_case_changed(case, case_type, close, fields, owner_id)
            if changed:
                user_case_helper.update_user_case(case, case_type, fields)
            if close and not case.closed:
                user_case_helper.close_user_case(case, case_type)

def _get_user_case_fields(commcare_user):

    def valid_element_name(name):
        try:
            ElementTree.fromstring('<{}/>'.format(name))
            return True
        except ElementTree.ParseError:
            return False

    # remove any keys that aren't valid XML element names
    fields = {k: v for k, v in commcare_user.user_data.items() if
              valid_element_name(k)}
    # language or phone_number can be null and will break
    # case submission
    fields.update({
        'name': commcare_user.name or commcare_user.raw_username,
        'username': commcare_user.raw_username,
        'email': commcare_user.email,
        'language': commcare_user.language or '',
        'phone_number': commcare_user.phone_number or ''
    })
    return fields

def _user_case_changed(case, case_type, close, fields, owner_id):
    props = case.dynamic_case_properties()
    changed = close != case.closed
    changed = changed or case.type != case_type
    changed = changed or case.name != fields['name']
    changed = changed or case.owner_id != owner_id
    if not changed:
        for field, value in fields.items():
            if field != 'name' and props.get(field) != value:
                changed = True
                break
    return changed

def sync_call_center_user_case(user):
    domain = user.project
    if domain and domain.call_center_config.enabled:
        case, owner_id = _get_call_center_case_and_owner(user, domain)
        sync_user_case(user, domain.call_center_config.case_type, owner_id, case)

CallCenterCaseAndOwner = namedtuple('CallCenterCaseAndOwner', 'case owner_id')

def _get_call_center_case_and_owner(user, domain):
    """
    Return the appropriate owner id for the given users call center case.
    """
    case = get_case_by_domain_hq_user_id(user.project.name, user._id, domain.call_center_config.case_type)
    if domain.call_center_config.use_user_location_as_owner:
        owner_id = call_center_location_owner(user, domain.call_center_config.user_location_ancestor_level)
    elif case and case.owner_id:
        owner_id = case.owner_id
    else:
        owner_id = domain.call_center_config.case_owner_id
    return CallCenterCaseAndOwner(case, owner_id)

def call_center_location_owner(user, ancestor_level):
    if ancestor_level == 0:
        owner_id = user.location_id
    else:
        location = SQLLocation.objects.get(location_id=user.location_id)
        ancestors = location.get_ancestors(ascending=True, include_self=True)
        try:
            owner_id = ancestors[ancestor_level].location_id
        except IndexError:
            owner_id = ancestors[-1].location_id
    return owner_id

def sync_usercase(user):
    domain = user.project
    if domain and domain.usercase_enabled:
        sync_user_case(
            user,
            USERCASE_TYPE,
            user.get_id
        )

def is_midnight_for_domain(midnight_form_domain, error_margin=15, current_time=None):
    current_time = current_time or datetime.utcnow()
    diff = current_time - midnight_form_domain
    return diff.days >= 0 and diff < timedelta(minutes=error_margin)

def get_call_center_domains():
    result = (
        DomainES()
            .is_active()
            .is_snapshot(False)
            .filter(filters.term('call_center_config.enabled', True))
            .fields(['name', 'default_timezone', 'call_center_config.case_type', 'call_center_config.use_fixtures'])
            .run()
    )

    def to_domain_lite(hit):
        return DomainLite(
            name=hit['name'],
            default_timezone=hit['default_timezone'],
            cc_case_type=hit.get('call_center_config.case_type', ''),
            use_fixtures=hit.get('call_center_config.use_fixtures', True)
        )
    return [to_domain_lite(hit) for hit in result.hits]

def get_call_center_cases(domain_name, case_type, user=None):
    all_cases = []

    if user:
        docs = (doc for owner_id in user.get_owner_ids()
                for doc in get_open_case_docs_in_domain(domain_name, case_type,
                                                        owner_id=owner_id))
    else:
        docs = get_open_case_docs_in_domain(domain_name, case_type)

    for case_doc in docs:
        hq_user_id = case_doc.get('hq_user_id', None)
        if hq_user_id:
            all_cases.append(CallCenterCase(
                case_id=case_doc['_id'],
                hq_user_id=hq_user_id
            ))
    return all_cases

@quickcache(['domain'])
def get_call_center_case_type_if_enabled(domain):
    if CALLCENTER.enabled(domain):
        return Domain.get_by_name(domain).call_center_config.case_type
<|endoftext|>


<commit_before>#
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.
#
import asyncio
import threading
import uuid

from gremlin_python.driver.client import Client
from gremlin_python.driver.protocol import GremlinServerError
from gremlin_python.driver.request import RequestMessage
from gremlin_python.process.graph_traversal import __
from gremlin_python.process.strategies import OptionsStrategy
from gremlin_python.structure.graph import Graph
from gremlin_python.driver.aiohttp.transport import AiohttpTransport
from gremlin_python.statics import *
from asyncio import TimeoutError

__author__ = 'David M. Brown (davebshow@gmail.com)'

test_no_auth_url = 'ws://localhost:45940/gremlin'

def test_connection(connection):
    g = Graph().traversal()
    t = g.V()
    message = RequestMessage('traversal', 'bytecode', {'gremlin': t.bytecode, 'aliases': {'g': 'gmodern'}})
    results_set = connection.write(message).result()
    future = results_set.all()
    results = future.result()
    assert len(results) == 6
    assert isinstance(results, list)
    assert results_set.done.done()
    assert 'host' in results_set.status_attributes

def test_client_message_too_big(client):
    try:
        client = Client(test_no_auth_url, 'g', max_content_length=1024)
        client.submit("\" \".repeat(2000)").all().result()
        assert False
    except Exception as ex:
        assert ex.args[0].startswith("Received error on read: 'Message size") \
               and ex.args[0].endswith("exceeds limit 1024'")
    finally:
        client.close()

def test_client_simple_eval(client):
    assert client.submit('1 + 1').all().result()[0] == 2

def test_client_simple_eval_bindings(client):
    assert client.submit('x + x', {'x': 2}).all().result()[0] == 4

def test_client_eval_traversal(client):
    assert len(client.submit('g.V()').all().result()) == 6

def test_client_error(client):
    try:
        # should fire an exception
        client.submit('1/0').all().result()
        assert False
    except GremlinServerError as ex:
        assert 'exceptions' in ex.status_attributes
        assert 'stackTrace' in ex.status_attributes

def test_client_connection_pool_after_error(client):
    # Overwrite fixture with pool_size=1 client
    client = Client(test_no_auth_url, 'gmodern', pool_size=1)

    try:
        # should fire an exception
        client.submit('1/0').all().result()
        assert False
    except GremlinServerError as gse:
        # expecting the pool size to be 1 again after query returned
        assert gse.status_code == 597
        assert client.available_pool_size == 1

def test_client_side_timeout_set_for_aiohttp(client):
    client = Client(test_no_auth_url, 'gmodern',
                    transport_factory=lambda: AiohttpTransport(read_timeout=1, write_timeout=1))

    try:
        # should fire an exception
        client.submit('Thread.sleep(1000);1').all().result()
        assert False
    except TimeoutError as err:
        # asyncio TimeoutError has no message.
        assert str(err) == ""

async def async_connect(enable):
    try:
        transport = AiohttpTransport(call_from_event_loop=enable)
        transport.connect(test_no_auth_url)
        transport.close()
        return True
    except RuntimeError:
        return False

def test_from_event_loop():
    assert not asyncio.get_event_loop().run_until_complete(async_connect(False))
    assert asyncio.get_event_loop().run_until_complete(async_connect(True))

def test_client_bytecode(client):
    g = Graph().traversal()
    t = g.V()
    message = RequestMessage('traversal', 'bytecode', {'gremlin': t.bytecode, 'aliases': {'g': 'gmodern'}})
    result_set = client.submit(message)
    assert len(result_set.all().result()) == 6

def test_client_bytecode_options(client):
    # smoke test to validate serialization of OptionsStrategy. no way to really validate this from an integration
    # test perspective because there's no way to access the internals of the strategy via bytecode
    g = Graph().traversal()
    t = g.withStrategies(OptionsStrategy(options={"x": "test", "y": True})).V()
    message = RequestMessage('traversal', 'bytecode', {'gremlin': t.bytecode, 'aliases': {'g': 'gmodern'}})
    result_set = client.submit(message)
    assert len(result_set.all().result()) == 6
    ##
    t = g.with_("x", "test").with_("y", True).V()
    message = RequestMessage('traversal', 'bytecode', {'gremlin': t.bytecode, 'aliases': {'g': 'gmodern'}})
    result_set = client.submit(message)
    assert len(result_set.all().result()) == 6

def test_iterate_result_set(client):
    g = Graph().traversal()
    t = g.V()
    message = RequestMessage('traversal', 'bytecode', {'gremlin': t.bytecode, 'aliases': {'g': 'gmodern'}})
    result_set = client.submit(message)
    results = []
    for result in result_set:
        results += result
    assert len(results) == 6

def test_client_async(client):
    g = Graph().traversal()
    t = g.V()
    message = RequestMessage('traversal', 'bytecode', {'gremlin': t.bytecode, 'aliases': {'g': 'gmodern'}})
    future = client.submit_async(message)
    result_set = future.result()
    assert len(result_set.all().result()) == 6

def test_connection_share(client):
    # Overwrite fixture with pool_size=1 client
    client = Client(test_no_auth_url, 'gmodern', pool_size=1)
    g = Graph().traversal()
    t = g.V()
    message = RequestMessage('traversal', 'bytecode', {'gremlin': t.bytecode, 'aliases': {'g': 'gmodern'}})
    message2 = RequestMessage('traversal', 'bytecode', {'gremlin': t.bytecode, 'aliases': {'g': 'gmodern'}})
    future = client.submit_async(message)
    future2 = client.submit_async(message2)

    result_set2 = future2.result()
    assert len(result_set2.all().result()) == 6

    # This future has to finish for the second to yield result - pool_size=1
    assert future.done()
    result_set = future.result()
    assert len(result_set.all().result()) == 6

def test_multi_conn_pool(client):
    g = Graph().traversal()
    t = g.V()
    message = RequestMessage('traversal', 'bytecode', {'gremlin': t.bytecode, 'aliases': {'g': 'gmodern'}})
    message2 = RequestMessage('traversal', 'bytecode', {'gremlin': t.bytecode, 'aliases': {'g': 'gmodern'}})
    client = Client(test_no_auth_url, 'g', pool_size=1)
    future = client.submit_async(message)
    future2 = client.submit_async(message2)

    result_set2 = future2.result()
    assert len(result_set2.all().result()) == 6

    # with connection pool `future` may or may not be done here
    result_set = future.result()
    assert len(result_set.all().result()) == 6

def test_multi_thread_pool(client):
    g = Graph().traversal()
    traversals = [g.V(),
                  g.V().count(),
                  g.E(),
                  g.E().count()
                  ]
    results = [[] for _ in traversals]

    # Use a condition variable to synchronise a group of threads, which should also inject some
    # non-determinism into the run-time execution order
    condition = threading.Condition()

    def thread_run(tr, result_list):
        message = RequestMessage('traversal', 'bytecode', {'gremlin': tr.bytecode, 'aliases': {'g': 'gmodern'}})
        with condition:
            condition.wait(5)
        result_set = client.submit(message)
        for result in result_set:
            result_list.append(result)

    threads = []
    for i in range(len(results)):
        thread = threading.Thread(target=thread_run,
                                  args=(traversals[i], results[i]),
                                  name="test_multi_thread_pool_%d" % i)
        thread.daemon = True
        threads.append(thread)
        thread.start()
    with condition:
        condition.notify_all()

    for t in threads:
        t.join(5)

    assert len(results[0][0]) == 6
    assert results[1][0][0].object == 6
    assert len(results[2][0]) == 6
    assert results[3][0][0].object == 6

def test_client_bytecode_with_long(client):
    g = Graph().traversal()
    t = g.V().has('age', long(851401972585122)).count()
    message = RequestMessage('traversal', 'bytecode', {'gremlin': t.bytecode, 'aliases': {'g': 'gmodern'}})
    result_set = client.submit(message)
    results = []
    for result in result_set:
        results += result
    assert len(results) == 1

def test_client_bytecode_with_bigint(client):
    g = Graph().traversal()
    t = g.V().has('age', bigint(0x1000_0000_0000_0000_0000)).count()
    message = RequestMessage('traversal', 'bytecode', {'gremlin': t.bytecode, 'aliases': {'g': 'gmodern'}})
    result_set = client.submit(message)
    results = []
    for result in result_set:
        results += result
    assert len(results) == 1

def test_multi_request_in_session(client):
    # Overwrite fixture with session client
    session_id = str(uuid.uuid4())
    client = Client(test_no_auth_url, 'g', session=session_id)

    assert client.submit('x = 1').all().result()[0] == 1
    assert client.submit('x + 2').all().result()[0] == 3

    client.close()

    # attempt reconnect to session and make sure "x" is no longer a thing
    client = Client(test_no_auth_url, 'g', session=session_id)
    try:
        # should fire an exception
        client.submit('x').all().result()
        assert False
    except Exception:
        assert True

def test_client_pool_in_session(client):
    # Overwrite fixture with pool_size=2 client
    try:
        # should fire an exception
        client = Client(test_no_auth_url, 'g', session=str(uuid.uuid4()), pool_size=2)
        assert False
    except Exception:
        assert True

def test_big_result_set(client):
    g = Graph().traversal()
    t = g.inject(1).repeat(__.addV('person').property('name', __.loops())).times(20000).count()
    message = RequestMessage('traversal', 'bytecode', {'gremlin': t.bytecode, 'aliases': {'g': 'g'}})
    result_set = client.submit(message)
    results = []
    for result in result_set:
        results += result
    assert len(results) == 1

    t = g.V().limit(10)
    message = RequestMessage('traversal', 'bytecode', {'gremlin': t.bytecode, 'aliases': {'g': 'g'}})
    result_set = client.submit(message)
    results = []
    for result in result_set:
        results += result
    assert len(results) == 10

    t = g.V().limit(100)
    message = RequestMessage('traversal', 'bytecode', {'gremlin': t.bytecode, 'aliases': {'g': 'g'}})
    result_set = client.submit(message)
    results = []
    for result in result_set:
        results += result
    assert len(results) == 100

    t = g.V().limit(1000)
    message = RequestMessage('traversal', 'bytecode', {'gremlin': t.bytecode, 'aliases': {'g': 'g'}})
    result_set = client.submit(message)
    results = []
    for result in result_set:
        results += result
    assert len(results) == 1000

    t = g.V().limit(10000)
    message = RequestMessage('traversal', 'bytecode', {'gremlin': t.bytecode, 'aliases': {'g': 'g'}})
    result_set = client.submit(message)
    results = []
    for result in result_set:
        results += result
    assert len(results) == 10000

def test_big_result_set_secure(authenticated_client):
    g = Graph().traversal()
    t = g.inject(1).repeat(__.addV('person').property('name', __.loops())).times(20000).count()
    message = RequestMessage('traversal', 'bytecode', {'gremlin': t.bytecode, 'aliases': {'g': 'g'}})
    result_set = authenticated_client.submit(message)
    results = []
    for result in result_set:
        results += result
    assert len(results) == 1

    t = g.V().limit(10)
    message = RequestMessage('traversal', 'bytecode', {'gremlin': t.bytecode, 'aliases': {'g': 'g'}})
    result_set = authenticated_client.submit(message)
    results = []
    for result in result_set:
        results += result
    assert len(results) == 10

    t = g.V().limit(100)
    message = RequestMessage('traversal', 'bytecode', {'gremlin': t.bytecode, 'aliases': {'g': 'g'}})
    result_set = authenticated_client.submit(message)
    results = []
    for result in result_set:
        results += result
    assert len(results) == 100

    t = g.V().limit(1000)
    message = RequestMessage('traversal', 'bytecode', {'gremlin': t.bytecode, 'aliases': {'g': 'g'}})
    result_set = authenticated_client.submit(message)
    results = []
    for result in result_set:
        results += result
    assert len(results) == 1000

    t = g.V().limit(10000)
    message = RequestMessage('traversal', 'bytecode', {'gremlin': t.bytecode, 'aliases': {'g': 'g'}})
    result_set = authenticated_client.submit(message)
    results = []
    for result in result_set:
        results += result
    assert len(results) == 10000

async def asyncio_func():
    return 1

def test_asyncio(client):
    try:
        asyncio.get_event_loop().run_until_complete(asyncio_func())
    except RuntimeError:
        assert False
<commit_msg>Fixing floating gremlin-python test failure<commit_after>#
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.
#
import asyncio
import threading
import uuid

from gremlin_python.driver.client import Client
from gremlin_python.driver.protocol import GremlinServerError
from gremlin_python.driver.request import RequestMessage
from gremlin_python.process.graph_traversal import __
from gremlin_python.process.strategies import OptionsStrategy
from gremlin_python.structure.graph import Graph
from gremlin_python.driver.aiohttp.transport import AiohttpTransport
from gremlin_python.statics import *
from asyncio import TimeoutError

__author__ = 'David M. Brown (davebshow@gmail.com)'

test_no_auth_url = 'ws://localhost:45940/gremlin'

def test_connection(connection):
    g = Graph().traversal()
    t = g.V()
    message = RequestMessage('traversal', 'bytecode', {'gremlin': t.bytecode, 'aliases': {'g': 'gmodern'}})
    results_set = connection.write(message).result()
    future = results_set.all()
    results = future.result()
    assert len(results) == 6
    assert isinstance(results, list)
    assert results_set.done.done()
    assert 'host' in results_set.status_attributes

def test_client_message_too_big(client):
    try:
        client = Client(test_no_auth_url, 'g', max_content_length=1024)
        client.submit("\" \".repeat(2000)").all().result()
        assert False
    except Exception as ex:
        assert ex.args[0].startswith("Received error on read: 'Message size") \
               and ex.args[0].endswith("exceeds limit 1024'")
    finally:
        client.close()

def test_client_simple_eval(client):
    assert client.submit('1 + 1').all().result()[0] == 2

def test_client_simple_eval_bindings(client):
    assert client.submit('x + x', {'x': 2}).all().result()[0] == 4

def test_client_eval_traversal(client):
    assert len(client.submit('g.V()').all().result()) == 6

def test_client_error(client):
    try:
        # should fire an exception
        client.submit('1/0').all().result()
        assert False
    except GremlinServerError as ex:
        assert 'exceptions' in ex.status_attributes
        assert 'stackTrace' in ex.status_attributes

def test_client_connection_pool_after_error(client):
    # Overwrite fixture with pool_size=1 client
    client = Client(test_no_auth_url, 'gmodern', pool_size=1)

    try:
        # should fire an exception
        client.submit('1/0').all().result()
        assert False
    except GremlinServerError as gse:
        # expecting the pool size to be 1 again after query returned
        assert gse.status_code == 597
        assert client.available_pool_size == 1

def test_client_side_timeout_set_for_aiohttp(client):
    client = Client(test_no_auth_url, 'gmodern',
                    transport_factory=lambda: AiohttpTransport(read_timeout=1, write_timeout=1))

    try:
        # should fire an exception
        client.submit('Thread.sleep(2000);1').all().result()
        assert False
    except TimeoutError as err:
        # asyncio TimeoutError has no message.
        assert str(err) == ""

async def async_connect(enable):
    try:
        transport = AiohttpTransport(call_from_event_loop=enable)
        transport.connect(test_no_auth_url)
        transport.close()
        return True
    except RuntimeError:
        return False

def test_from_event_loop():
    assert not asyncio.get_event_loop().run_until_complete(async_connect(False))
    assert asyncio.get_event_loop().run_until_complete(async_connect(True))

def test_client_bytecode(client):
    g = Graph().traversal()
    t = g.V()
    message = RequestMessage('traversal', 'bytecode', {'gremlin': t.bytecode, 'aliases': {'g': 'gmodern'}})
    result_set = client.submit(message)
    assert len(result_set.all().result()) == 6

def test_client_bytecode_options(client):
    # smoke test to validate serialization of OptionsStrategy. no way to really validate this from an integration
    # test perspective because there's no way to access the internals of the strategy via bytecode
    g = Graph().traversal()
    t = g.withStrategies(OptionsStrategy(options={"x": "test", "y": True})).V()
    message = RequestMessage('traversal', 'bytecode', {'gremlin': t.bytecode, 'aliases': {'g': 'gmodern'}})
    result_set = client.submit(message)
    assert len(result_set.all().result()) == 6
    ##
    t = g.with_("x", "test").with_("y", True).V()
    message = RequestMessage('traversal', 'bytecode', {'gremlin': t.bytecode, 'aliases': {'g': 'gmodern'}})
    result_set = client.submit(message)
    assert len(result_set.all().result()) == 6

def test_iterate_result_set(client):
    g = Graph().traversal()
    t = g.V()
    message = RequestMessage('traversal', 'bytecode', {'gremlin': t.bytecode, 'aliases': {'g': 'gmodern'}})
    result_set = client.submit(message)
    results = []
    for result in result_set:
        results += result
    assert len(results) == 6

def test_client_async(client):
    g = Graph().traversal()
    t = g.V()
    message = RequestMessage('traversal', 'bytecode', {'gremlin': t.bytecode, 'aliases': {'g': 'gmodern'}})
    future = client.submit_async(message)
    result_set = future.result()
    assert len(result_set.all().result()) == 6

def test_connection_share(client):
    # Overwrite fixture with pool_size=1 client
    client = Client(test_no_auth_url, 'gmodern', pool_size=1)
    g = Graph().traversal()
    t = g.V()
    message = RequestMessage('traversal', 'bytecode', {'gremlin': t.bytecode, 'aliases': {'g': 'gmodern'}})
    message2 = RequestMessage('traversal', 'bytecode', {'gremlin': t.bytecode, 'aliases': {'g': 'gmodern'}})
    future = client.submit_async(message)
    future2 = client.submit_async(message2)

    result_set2 = future2.result()
    assert len(result_set2.all().result()) == 6

    # This future has to finish for the second to yield result - pool_size=1
    assert future.done()
    result_set = future.result()
    assert len(result_set.all().result()) == 6

def test_multi_conn_pool(client):
    g = Graph().traversal()
    t = g.V()
    message = RequestMessage('traversal', 'bytecode', {'gremlin': t.bytecode, 'aliases': {'g': 'gmodern'}})
    message2 = RequestMessage('traversal', 'bytecode', {'gremlin': t.bytecode, 'aliases': {'g': 'gmodern'}})
    client = Client(test_no_auth_url, 'g', pool_size=1)
    future = client.submit_async(message)
    future2 = client.submit_async(message2)

    result_set2 = future2.result()
    assert len(result_set2.all().result()) == 6

    # with connection pool `future` may or may not be done here
    result_set = future.result()
    assert len(result_set.all().result()) == 6

def test_multi_thread_pool(client):
    g = Graph().traversal()
    traversals = [g.V(),
                  g.V().count(),
                  g.E(),
                  g.E().count()
                  ]
    results = [[] for _ in traversals]

    # Use a condition variable to synchronise a group of threads, which should also inject some
    # non-determinism into the run-time execution order
    condition = threading.Condition()

    def thread_run(tr, result_list):
        message = RequestMessage('traversal', 'bytecode', {'gremlin': tr.bytecode, 'aliases': {'g': 'gmodern'}})
        with condition:
            condition.wait(5)
        result_set = client.submit(message)
        for result in result_set:
            result_list.append(result)

    threads = []
    for i in range(len(results)):
        thread = threading.Thread(target=thread_run,
                                  args=(traversals[i], results[i]),
                                  name="test_multi_thread_pool_%d" % i)
        thread.daemon = True
        threads.append(thread)
        thread.start()
    with condition:
        condition.notify_all()

    for t in threads:
        t.join(5)

    assert len(results[0][0]) == 6
    assert results[1][0][0].object == 6
    assert len(results[2][0]) == 6
    assert results[3][0][0].object == 6

def test_client_bytecode_with_long(client):
    g = Graph().traversal()
    t = g.V().has('age', long(851401972585122)).count()
    message = RequestMessage('traversal', 'bytecode', {'gremlin': t.bytecode, 'aliases': {'g': 'gmodern'}})
    result_set = client.submit(message)
    results = []
    for result in result_set:
        results += result
    assert len(results) == 1

def test_client_bytecode_with_bigint(client):
    g = Graph().traversal()
    t = g.V().has('age', bigint(0x1000_0000_0000_0000_0000)).count()
    message = RequestMessage('traversal', 'bytecode', {'gremlin': t.bytecode, 'aliases': {'g': 'gmodern'}})
    result_set = client.submit(message)
    results = []
    for result in result_set:
        results += result
    assert len(results) == 1

def test_multi_request_in_session(client):
    # Overwrite fixture with session client
    session_id = str(uuid.uuid4())
    client = Client(test_no_auth_url, 'g', session=session_id)

    assert client.submit('x = 1').all().result()[0] == 1
    assert client.submit('x + 2').all().result()[0] == 3

    client.close()

    # attempt reconnect to session and make sure "x" is no longer a thing
    client = Client(test_no_auth_url, 'g', session=session_id)
    try:
        # should fire an exception
        client.submit('x').all().result()
        assert False
    except Exception:
        assert True

def test_client_pool_in_session(client):
    # Overwrite fixture with pool_size=2 client
    try:
        # should fire an exception
        client = Client(test_no_auth_url, 'g', session=str(uuid.uuid4()), pool_size=2)
        assert False
    except Exception:
        assert True

def test_big_result_set(client):
    g = Graph().traversal()
    t = g.inject(1).repeat(__.addV('person').property('name', __.loops())).times(20000).count()
    message = RequestMessage('traversal', 'bytecode', {'gremlin': t.bytecode, 'aliases': {'g': 'g'}})
    result_set = client.submit(message)
    results = []
    for result in result_set:
        results += result
    assert len(results) == 1

    t = g.V().limit(10)
    message = RequestMessage('traversal', 'bytecode', {'gremlin': t.bytecode, 'aliases': {'g': 'g'}})
    result_set = client.submit(message)
    results = []
    for result in result_set:
        results += result
    assert len(results) == 10

    t = g.V().limit(100)
    message = RequestMessage('traversal', 'bytecode', {'gremlin': t.bytecode, 'aliases': {'g': 'g'}})
    result_set = client.submit(message)
    results = []
    for result in result_set:
        results += result
    assert len(results) == 100

    t = g.V().limit(1000)
    message = RequestMessage('traversal', 'bytecode', {'gremlin': t.bytecode, 'aliases': {'g': 'g'}})
    result_set = client.submit(message)
    results = []
    for result in result_set:
        results += result
    assert len(results) == 1000

    t = g.V().limit(10000)
    message = RequestMessage('traversal', 'bytecode', {'gremlin': t.bytecode, 'aliases': {'g': 'g'}})
    result_set = client.submit(message)
    results = []
    for result in result_set:
        results += result
    assert len(results) == 10000

def test_big_result_set_secure(authenticated_client):
    g = Graph().traversal()
    t = g.inject(1).repeat(__.addV('person').property('name', __.loops())).times(20000).count()
    message = RequestMessage('traversal', 'bytecode', {'gremlin': t.bytecode, 'aliases': {'g': 'g'}})
    result_set = authenticated_client.submit(message)
    results = []
    for result in result_set:
        results += result
    assert len(results) == 1

    t = g.V().limit(10)
    message = RequestMessage('traversal', 'bytecode', {'gremlin': t.bytecode, 'aliases': {'g': 'g'}})
    result_set = authenticated_client.submit(message)
    results = []
    for result in result_set:
        results += result
    assert len(results) == 10

    t = g.V().limit(100)
    message = RequestMessage('traversal', 'bytecode', {'gremlin': t.bytecode, 'aliases': {'g': 'g'}})
    result_set = authenticated_client.submit(message)
    results = []
    for result in result_set:
        results += result
    assert len(results) == 100

    t = g.V().limit(1000)
    message = RequestMessage('traversal', 'bytecode', {'gremlin': t.bytecode, 'aliases': {'g': 'g'}})
    result_set = authenticated_client.submit(message)
    results = []
    for result in result_set:
        results += result
    assert len(results) == 1000

    t = g.V().limit(10000)
    message = RequestMessage('traversal', 'bytecode', {'gremlin': t.bytecode, 'aliases': {'g': 'g'}})
    result_set = authenticated_client.submit(message)
    results = []
    for result in result_set:
        results += result
    assert len(results) == 10000

async def asyncio_func():
    return 1

def test_asyncio(client):
    try:
        asyncio.get_event_loop().run_until_complete(asyncio_func())
    except RuntimeError:
        assert False
<|endoftext|>


<commit_before>import calendar
import datetime
import functools
import json
import os
import time

import cherrypy
import pypyodbc

def connect(thread_index):
    cherrypy.thread_data.db = pypyodbc.connect(os.environ['DATABASE'])

def dont_cache():
    cherrypy.response.headers['Expires'] = datetime.datetime.utcnow().strftime(
        '%a, %d %b %Y %H:%M:%S GMT')
    cherrypy.response.headers['Cache-Control'] = ('no-store, '
                                                  'no-cache, '
                                                  'must-revalidate, '
                                                  'post-check=0, '
                                                  'pre-check=0')

def protect(func):
    @functools.wraps(func)
    def decorator(*args, **kwargs):
        if kwargs.pop('app_secret', False) != os.environ['APP_SECRET']:
            raise cherrypy.HTTPError(403)
        return func(*args, **kwargs)
    return decorator

def auth(func):
    @functools.wraps(func)
    def decorator(*args, **kwargs):
        if not kwargs.get('facebook_token', False):
            raise cherrypy.HTTPError(403)
        cursor = kwargs['cursor']
        qry = {'select':   'id',
               'table':    'users',
               'where':    'facebook_token = ?',
               'order_by': 'id',
               'limit':    1}
        cursor.execute(query(**qry), (kwargs['facebook_token'],))
        res = cursor.fetchone()['id']
        if not res:
            raise cherrypy.HTTPError(403)
        kwargs.update(user_id=res)
        return func(*args, **kwargs)
    return decorator

def jsonp(func):
    @functools.wraps(func)
    def decorator(*args, **kwargs):
        dont_cache()
        res = json.dumps(func(*args, **kwargs), separators=(',', ':'))
        callback = kwargs.pop('callback', False)
        if callback:
            cherrypy.response.headers['Content-Type'] = ('text/javascript; '
                                                         'charset=utf-8')
            res = callback + '(' + res + ');'
        else:
            cherrypy.response.headers['Content-Type'] = 'application/json'
        return res
    return decorator

def db(func):
    @functools.wraps(func)
    def decorator(*args, **kwargs):
        cursor = cherrypy.thread_data.db.cursor()
        kwargs.update(cursor=cursor)
        try:
            res = func(*args, **kwargs)
        finally:
            cnx.commit()
            cursor.close()
        return res
    return decorator

def implode(glue, list):
    return list if isinstance(list, str) else glue.join(list)

def query(select=None, table=None, left_join=None, on=None, where=None,
          group_by=None, order_by=None, limit=None,
          insert_into=None, columns=None,
          update=None, set_values=None):
    if select:
        qry = 'SELECT ' + implode(', ', select)
        if table:
            qry += ' FROM ' + implode(', ', table)
        if left_join and on:
            if isinstance(left_join, str):
                left_join = [left_join]
            if isinstance(on, str):
                on = [on]
            for j, o in zip(left_join, on):
                qry += ' LEFT JOIN ' + j + ' ON ' + o
        if where:
            qry += ' WHERE ' + implode(' AND ', where)
        if group_by:
            qry += ' GROUP BY ' + implode(', ', group_by)
        if order_by:
            qry += ' ORDER BY ' + implode(', ', order_by)
        if limit:
            if isinstance(limit, str) or isinstance(limit, int):
                qry += ' OFFSET 0 ROWS FETCH NEXT ' + str(limit) + ' ROWS ONLY'
            else:
                qry += (' OFFSET ' + str(int(limit[0])) + ' ROWS FETCH NEXT ' +
                        str(int(limit[1])) + ' ROWS ONLY')
    elif insert_into:
        qry = 'INSERT INTO ' + insert_into
        if columns:
            qry += (' (' + implode(', ', columns) + ')' + ' VALUES (' +
                    ('?' + ', ?' * (len(columns) - 1)) + ')')
    elif update:
        qry = 'UPDATE ' + update
        if set_values:
            qry += ' SET ' + implode('=?, ', set_values) + '=?'
        if where:
            qry += ' WHERE ' + implode(' AND ', where)
    return qry

expose = cherrypy.expose

def to_int(value):
    return int(value) if value else None

def to_bool(value):
    if value in (0, False, None, '0'): return False
    if isinstance(value, str) and value in ('false', 'no', 'off'): return False
    return True
<commit_msg>Call commit on cursor<commit_after>import calendar
import datetime
import functools
import json
import os
import time

import cherrypy
import pypyodbc

def connect(thread_index):
    cherrypy.thread_data.db = pypyodbc.connect(os.environ['DATABASE'])

def dont_cache():
    cherrypy.response.headers['Expires'] = datetime.datetime.utcnow().strftime(
        '%a, %d %b %Y %H:%M:%S GMT')
    cherrypy.response.headers['Cache-Control'] = ('no-store, '
                                                  'no-cache, '
                                                  'must-revalidate, '
                                                  'post-check=0, '
                                                  'pre-check=0')

def protect(func):
    @functools.wraps(func)
    def decorator(*args, **kwargs):
        if kwargs.pop('app_secret', False) != os.environ['APP_SECRET']:
            raise cherrypy.HTTPError(403)
        return func(*args, **kwargs)
    return decorator

def auth(func):
    @functools.wraps(func)
    def decorator(*args, **kwargs):
        if not kwargs.get('facebook_token', False):
            raise cherrypy.HTTPError(403)
        cursor = kwargs['cursor']
        qry = {'select':   'id',
               'table':    'users',
               'where':    'facebook_token = ?',
               'order_by': 'id',
               'limit':    1}
        cursor.execute(query(**qry), (kwargs['facebook_token'],))
        res = cursor.fetchone()['id']
        if not res:
            raise cherrypy.HTTPError(403)
        kwargs.update(user_id=res)
        return func(*args, **kwargs)
    return decorator

def jsonp(func):
    @functools.wraps(func)
    def decorator(*args, **kwargs):
        dont_cache()
        res = json.dumps(func(*args, **kwargs), separators=(',', ':'))
        callback = kwargs.pop('callback', False)
        if callback:
            cherrypy.response.headers['Content-Type'] = ('text/javascript; '
                                                         'charset=utf-8')
            res = callback + '(' + res + ');'
        else:
            cherrypy.response.headers['Content-Type'] = 'application/json'
        return res
    return decorator

def db(func):
    @functools.wraps(func)
    def decorator(*args, **kwargs):
        cursor = cherrypy.thread_data.db.cursor()
        kwargs.update(cursor=cursor)
        try:
            res = func(*args, **kwargs)
        finally:
            cursor.commit()
            cursor.close()
        return res
    return decorator

def implode(glue, list):
    return list if isinstance(list, str) else glue.join(list)

def query(select=None, table=None, left_join=None, on=None, where=None,
          group_by=None, order_by=None, limit=None,
          insert_into=None, columns=None,
          update=None, set_values=None):
    if select:
        qry = 'SELECT ' + implode(', ', select)
        if table:
            qry += ' FROM ' + implode(', ', table)
        if left_join and on:
            if isinstance(left_join, str):
                left_join = [left_join]
            if isinstance(on, str):
                on = [on]
            for j, o in zip(left_join, on):
                qry += ' LEFT JOIN ' + j + ' ON ' + o
        if where:
            qry += ' WHERE ' + implode(' AND ', where)
        if group_by:
            qry += ' GROUP BY ' + implode(', ', group_by)
        if order_by:
            qry += ' ORDER BY ' + implode(', ', order_by)
        if limit:
            if isinstance(limit, str) or isinstance(limit, int):
                qry += ' OFFSET 0 ROWS FETCH NEXT ' + str(limit) + ' ROWS ONLY'
            else:
                qry += (' OFFSET ' + str(int(limit[0])) + ' ROWS FETCH NEXT ' +
                        str(int(limit[1])) + ' ROWS ONLY')
    elif insert_into:
        qry = 'INSERT INTO ' + insert_into
        if columns:
            qry += (' (' + implode(', ', columns) + ')' + ' VALUES (' +
                    ('?' + ', ?' * (len(columns) - 1)) + ')')
    elif update:
        qry = 'UPDATE ' + update
        if set_values:
            qry += ' SET ' + implode('=?, ', set_values) + '=?'
        if where:
            qry += ' WHERE ' + implode(' AND ', where)
    return qry

expose = cherrypy.expose

def to_int(value):
    return int(value) if value else None

def to_bool(value):
    if value in (0, False, None, '0'): return False
    if isinstance(value, str) and value in ('false', 'no', 'off'): return False
    return True
<|endoftext|>