<commit_before>from peewee import *

from gouda.bot import DATABASE as db
from gouda.models import Message, Quote

def run_schema():
    Quote.create_table(True)

# main

def main(*args, **kwargs):
    """
    quote stuff

    qsave [tag] quote   save quote
    qsave [tag] ^num    save quote chat backref
    qload [tag]         load random quote from tag
    qfind [pattern]     find a random quote based on a pattern
    """
    # since this module has special syntax we use a main func
    writer = kwargs.pop('writer')
    line = kwargs.pop('line')
    if line[0].startswith("qsave") and len(line) > 2:
        # saving a new quote
        tag = line[1]
        if line[2].startswith("^") and len(line[2]) > 1:
            last_id = Message.select().order_by(Message.id.desc()).get().id
            try:
                num = int(line[2][1:])
                if num >= 0:
                    try:
                        message = Message.select().where(Message.id==last_id-num).get()
                        Quote.create(tag=tag, user=message.name, message=message.message)
                    except Message.DoesNotExist:
                        # just in case
                        pass
            except ValueError as e:
                # naughty user input
                print(e)
        else:
            message = ' '.join(line[2:])
            Quote.create(tag=tag, user=kwargs.pop('nick', '~'), message=message)
    elif line[0].startswith("qload") and len(line) > 1:
        # retrieving a quote
        tag = line[1]
        try:
            quote = Quote.select().where(Quote.tag==tag).order_by(fn.Random()).get()
            writer('<%s> %s' % (quote.user, quote.message))
        except:
            # probably no quotes in this tag
            print("Could not load quote with tag %s" % tag)
    elif line[0].startswith("qfind") and len(line) > 1:
        # search records
        try:
            quote = Quote.select().where(Quote.message.contains(' '.join(line[1:]))).order_by(fn.Random()).get()
            writer('<%s> %s' % (quote.user, quote.message))
        except:
            # probably no quotes in this tag
            print("Could not find quote with pattern %s" % ' '.join(line[1:]))
<commit_msg>qtags added and feedback on saving quote<commit_after>from peewee import *

from gouda.bot import DATABASE as db
from gouda.models import Message, Quote

def run_schema():
    Quote.create_table(True)

# main

def main(*args, **kwargs):
    """
    quote stuff

    qsave [tag] quote   save quote
    qsave [tag] ^num    save quote chat backref
    qload [tag]         load random quote from tag
    qfind [pattern]     find a random quote based on a pattern
    """
    # since this module has special syntax we use a main func
    writer = kwargs.pop('writer')
    line = kwargs.pop('line')
    if line[0] == "qsave" and len(line) > 2:
        # saving a new quote
        tag = line[1]
        saved = False
        if line[2].startswith("^") and len(line[2]) > 1:
            last_id = Message.select().order_by(Message.id.desc()).get().id
            try:
                num = int(line[2][1:])
                if num >= 0:
                    try:
                        message = Message.select().where(Message.id==last_id-num).get()
                        Quote.create(tag=tag, user=message.name, message=message.message)
                        saved = True
                    except Message.DoesNotExist:
                        # just in case
                        pass
            except ValueError as e:
                # naughty user input
                print(e)
        else:
            message = ' '.join(line[2:])
            Quote.create(tag=tag, user=kwargs.pop('nick', '~'), message=message)
            saved = True
        if saved:
            writer("Saved your quote, pal")
    elif line[0] == "qload" and len(line) > 1:
        # retrieving a quote
        tag = line[1]
        try:
            quote = Quote.select().where(Quote.tag==tag).order_by(fn.Random()).get()
            writer('<%s> %s' % (quote.user, quote.message))
        except:
            # probably no quotes in this tag
            print("Could not load quote with tag %s" % tag)
    elif line[0] == "qfind" and len(line) > 1:
        # search records
        try:
            quote = Quote.select().where(Quote.message.contains(' '.join(line[1:]))).order_by(fn.Random()).get()
            writer('<%s> %s' % (quote.user, quote.message))
        except:
            # probably no quotes in this tag
            print("Could not find quote with pattern %s" % ' '.join(line[1:]))
    elif line[0] == "qtags":
        try:
            tags = Quote.select(Quote.tag).distinct()
            writer('Tags: %s.' % ', '.join(t.tag for t in tags))
        except:
            print("Cannot load tags")
<|endoftext|>


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

#  IATI Data Quality, tools for Data QA on IATI-formatted  publications
#  by Mark Brough, Martin Keegan, Ben Webb and Jennifer Smith
#
#  Copyright (C) 2013  Publish What You Fund
#
#  This programme is free software; you may redistribute and/or modify
#  it under the terms of the GNU Affero General Public License v3.0

""" This script is to quickly get started with this tool, by:
        1) creating DB
        2) populating the list of packages from the Registry (will download basic data about all packages)
        3) setting 3 to "active"
"""
#import warnings
#warnings.filterwarnings('error')

import iatidq

from iatidq import models, dqregistry
import iatidq.dqfunctions
import iatidq.dqimporttests
import iatidq.dqdownload
import iatidq.dqcodelists
import iatidq.dqruntests
import iatidq.dqindicators
import iatidq.dqorganisations
import iatidq.dqaggregationtypes
import iatidq.dqtests
import iatidq.dqprocessing
import iatidq.test_level as test_level
import iatidq.inforesult
import optparse
import sys

which_packages = [
    (u'worldbank-tz', True),
    (u'unops-tz', True),
    (u'dfid-tz', True),
    (u'unitedstates-tz', True)
    ]

default_tests_filename="tests/tests.csv"
default_indicator_group_name="pwyf2013"
default_minimal_organisations = [
            {
            'organisation_name': 'World Bank',
            'organisation_code': '44002',
            'packagegroup_name': 'worldbank',
            'condition': None,
            'package_name': 'worldbank-tz'
            },
            {
            'organisation_name': 'USAID',
            'organisation_code': 'US-1',
            'packagegroup_name': 'unitedstates',
            'condition': 'participating-org[@role="Extending"][@ref="US-1"]',
            'package_name': 'unitedstates-tz'
            },
            {
            'organisation_name': 'MCC',
            'organisation_code': 'US-18',
            'packagegroup_name': 'unitedstates',
            'condition': 'participating-org[@role="Extending"][@ref="US-18"]',
            'package_name': 'unitedstates-tz'
            },
            {
            'organisation_name': 'UK, DFID',
            'organisation_code': 'GB-1',
            'packagegroup_name': 'dfid',
            'condition': None,
            'package_name': 'dfid-tz'
            }
        ]

def refresh(options):
    pkg_names = None
    if options.package_name:
        pkg_names = [options.package_name]
    elif options.minimal_pkgs:
        pkg_names = [i[0] for i in which_packages]

    if pkg_names is not None:
        [ dqregistry.refresh_package_by_name(name) for name in pkg_names ]
    else:
        dqregistry.refresh_packages()

def activate_packages(options):
    dqregistry.activate_packages(which_packages, clear_revision_id=True)

def drop_all(options):
    iatidq.db.drop_all()

def init_db(options):
    iatidq.db.create_all()
    iatidq.dqimporttests.hardcodedTests()

def enroll_tests(options):
    assert options.filename
    filename = options.filename.decode()
    result = iatidq.dqimporttests.importTestsFromFile(
        filename=filename,
        level=options.level)
    if not result:
        print "Error importing"

def clear_revisionid(options):
    iatidq.dqfunctions.clear_revisions()

def import_codelists(options):
    iatidq.dqcodelists.importCodelists()

def download(options):
    if options.minimal_pkgs:
        for package_name, _ in which_packages:
            iatidq.dqdownload.run(package_name=package_name)
    else:
        iatidq.dqdownload.run()

def import_indicators(options):
    if options.filename:
        iatidq.dqindicators.importIndicatorsFromFile("pwyf2013",
                                                     options.filename)
    else:
        iatidq.dqindicators.importIndicators()

def import_organisations(options):
    if options.filename:
        iatidq.dqorganisations.importOrganisationPackagesFromFile(options.filename)
    else:
        print "Error: please provide a filename"

def create_aggregation_types(options):
    print "Adding an aggregation type for all data"
    iatidq.dqaggregationtypes.addAggregationType({'name':'All data',
                                                'description': '',
                                                'test_id': None,
                                                'test_result':'1'})
    print "Adding an aggregation type for current data"
    currentdata_test = iatidq.dqtests.test_by_test_name(
        "activity-date[@type='start-planned']/@iso-date or transaction-date/@iso-date (for each transaction) is less than 13 months ago?"
        )
    iatidq.dqaggregationtypes.addAggregationType({'name':'Current data',
                                                'description': '',
                                                'test_id':currentdata_test.id,
                                                'test_result':'1'})

def create_inforesult_types(options):
    print "Adding an aggregation type"
    iatidq.inforesult.add_type("coverage", "Coverage")

def setup_common():
    print "Creating DB"
    iatidq.db.create_all()
    print "Adding hardcoded tests"
    iatidq.dqimporttests.hardcodedTests()
    print "Importing tests"
    iatidq.dqimporttests.importTestsFromFile(
        default_tests_filename,
        test_level.ACTIVITY)
    print "Importing indicators"
    iatidq.dqindicators.importIndicatorsFromFile(
        default_indicator_group_name,
        default_tests_filename)
    print "Importing indicator descriptions"
    iatidq.dqindicators.importIndicatorDescriptionsFromFile("pwyf2013",
                                                            "tests/indicators.csv")
    print "Importing codelists"
    iatidq.dqcodelists.importCodelists()

def setup_minimal(options):
    setup_common()

    print "Creating packages"
    pkg_names = [i[0] for i in which_packages]

    if pkg_names is not None:
        [ dqregistry.refresh_package_by_name(name) for name in pkg_names ]
    else:
        print "No packages are defined in quickstart"

    print "Creating aggregation types."
    create_aggregation_types(options)

    create_inforesult_types(options)

    print "Setup complete."

    for organisation in default_minimal_organisations:
        inserted_organisation = iatidq.dqorganisations.addOrganisation(organisation)
        if inserted_organisation is False:
            inserted_organisation = iatidq.dqorganisations.organisations(organisation['organisation_code'])
        thepackage = models.Package.query.filter_by(package_name=organisation['package_name']
                ).first()

        organisationpackage_data = {"organisation_id": inserted_organisation.id,
                                                       "package_id": thepackage.id,
                                                       "condition": organisation["condition"]
                                                    }
        iatidq.dqorganisations.addOrganisationPackage(organisationpackage_data)
    print "Complete"

def setup(options):
    setup_common()
    print "Refreshing package data from Registry"
    dqregistry.refresh_packages()
    print "Adding organisations"
    iatidq.dqorganisations.importOrganisationPackagesFromFile("tests/organisations_with_identifiers.csv")
    create_aggregation_types(options)
    create_inforesult_types(options)
    print "Setup complete."

def enqueue_test(options):
    assert options.package_name
    assert options.filename

    iatidq.dqruntests.enqueue_package_for_test(options.filename,
                                               options.package_name)
def aggregate_results(options):
    assert options.runtime_id
    assert options.package_id
    iatidq.dqprocessing.aggregate_results(options.runtime_id, options.package_id)

commands = {
    "drop_db": (drop_all, "Delete DB"),
    "init_db": (init_db, "Initialise DB"),
    "enroll_tests": (enroll_tests, "Enroll a CSV file of tests"),
    "clear_revisionid": (clear_revisionid, "Clear CKAN revision ids"),
    "import_codelists": (import_codelists, "Import codelists"),
    "download": (download, "Download packages"),
    "import_indicators": (
        import_indicators,
        "Import indicators. Will try to assign indicators to existing tests."),
    "import_organisations": (
        import_organisations,
        "Import organisations. Will try to create and assign organisations "
        "to existing packages."),
    "setup": (setup, """Quick setup. Will init db, add tests, add codelists,
                      add indicators, refresh package data from Registry."""),
    "enqueue_test": (enqueue_test, "Set a package to be tested (with --package)"),
    "refresh": (refresh, "Refresh"),
    "activate_packages": (activate_packages, "Mark all packages as active"),
    "create_aggregation_types": (create_aggregation_types, "Create basic aggregation types."),
    "aggregate_results": (aggregate_results, "Trigger result aggregation"),
    "setup_minimal": (setup_minimal, "Quick minimal setup"),
    "create_inforesult_types": (create_inforesult_types, "Create basic infroresult types.")
}

def main():
    p = optparse.OptionParser()

    for k, v in commands.iteritems():
        handler, help_text = v
        option_name = "--" + k.replace("_", "-")
        p.add_option(option_name, dest=k, action="store_true", default=False, help=help_text)

    p.add_option("--runtime-id", dest="runtime_id",
                 type=int,
                 help="Runtime id (integer)")
    p.add_option("--package-id", dest="package_id",
                 type=int,
                 help="Package id (integer)")
    p.add_option("--level", dest="level",
                 type="int",
                 default=1,
                 help="Test level (e.g., 1 == Activity)")
    p.add_option("--minimal-pkgs", dest="minimal_pkgs",
                 action="store_true",
                 default=False,
                 help="Operate on a minimal set of packages")
    p.add_option("--package", dest="package_name",
                 help="Set name of package to be tested")
    p.add_option("--filename", dest="filename",
                 help="Set filename of data to test")
    p.add_option("--local-folder", dest="local_folder",
                 help="Set local folder where data to test is stored")

    options, args = p.parse_args()

    for mode, handler_ in commands.iteritems():
        handler, _ = handler_
        if getattr(options, mode, None):
            handler(options)
            return

    usage()

def usage():
    print "You need to specify which mode to run under"
    sys.exit(1)

if __name__ == '__main__':
    main()
<commit_msg>break up long lines<commit_after>#!/usr/bin/env python

#  IATI Data Quality, tools for Data QA on IATI-formatted  publications
#  by Mark Brough, Martin Keegan, Ben Webb and Jennifer Smith
#
#  Copyright (C) 2013  Publish What You Fund
#
#  This programme is free software; you may redistribute and/or modify
#  it under the terms of the GNU Affero General Public License v3.0

""" This script is to quickly get started with this tool, by:
        1) creating DB
        2) populating the list of packages from the Registry (will download basic data about all packages)
        3) setting 3 to "active"
"""
#import warnings
#warnings.filterwarnings('error')

import iatidq

from iatidq import models, dqregistry
import iatidq.dqfunctions
import iatidq.dqimporttests
import iatidq.dqdownload
import iatidq.dqcodelists
import iatidq.dqruntests
import iatidq.dqindicators
import iatidq.dqorganisations
import iatidq.dqaggregationtypes
import iatidq.dqtests
import iatidq.dqprocessing
import iatidq.test_level as test_level
import iatidq.inforesult
import optparse
import sys

which_packages = [
    (u'worldbank-tz', True),
    (u'unops-tz', True),
    (u'dfid-tz', True),
    (u'unitedstates-tz', True)
    ]

default_tests_filename="tests/tests.csv"
default_indicator_group_name="pwyf2013"
default_minimal_organisations = [
            {
            'organisation_name': 'World Bank',
            'organisation_code': '44002',
            'packagegroup_name': 'worldbank',
            'condition': None,
            'package_name': 'worldbank-tz'
            },
            {
            'organisation_name': 'USAID',
            'organisation_code': 'US-1',
            'packagegroup_name': 'unitedstates',
            'condition': 'participating-org[@role="Extending"][@ref="US-1"]',
            'package_name': 'unitedstates-tz'
            },
            {
            'organisation_name': 'MCC',
            'organisation_code': 'US-18',
            'packagegroup_name': 'unitedstates',
            'condition': 'participating-org[@role="Extending"][@ref="US-18"]',
            'package_name': 'unitedstates-tz'
            },
            {
            'organisation_name': 'UK, DFID',
            'organisation_code': 'GB-1',
            'packagegroup_name': 'dfid',
            'condition': None,
            'package_name': 'dfid-tz'
            }
        ]

def refresh(options):
    pkg_names = None
    if options.package_name:
        pkg_names = [options.package_name]
    elif options.minimal_pkgs:
        pkg_names = [i[0] for i in which_packages]

    if pkg_names is not None:
        [ dqregistry.refresh_package_by_name(name) for name in pkg_names ]
    else:
        dqregistry.refresh_packages()

def activate_packages(options):
    dqregistry.activate_packages(which_packages, clear_revision_id=True)

def drop_all(options):
    iatidq.db.drop_all()

def init_db(options):
    iatidq.db.create_all()
    iatidq.dqimporttests.hardcodedTests()

def enroll_tests(options):
    assert options.filename
    filename = options.filename.decode()
    result = iatidq.dqimporttests.importTestsFromFile(
        filename=filename,
        level=options.level)
    if not result:
        print "Error importing"

def clear_revisionid(options):
    iatidq.dqfunctions.clear_revisions()

def import_codelists(options):
    iatidq.dqcodelists.importCodelists()

def download(options):
    if options.minimal_pkgs:
        for package_name, _ in which_packages:
            iatidq.dqdownload.run(package_name=package_name)
    else:
        iatidq.dqdownload.run()

def import_indicators(options):
    if options.filename:
        iatidq.dqindicators.importIndicatorsFromFile("pwyf2013",
                                                     options.filename)
    else:
        iatidq.dqindicators.importIndicators()

def import_organisations(options):
    if options.filename:
        iatidq.dqorganisations.importOrganisationPackagesFromFile(options.filename)
    else:
        print "Error: please provide a filename"

def create_aggregation_types(options):
    print "Adding an aggregation type for all data"
    iatidq.dqaggregationtypes.addAggregationType({'name':'All data',
                                                'description': '',
                                                'test_id': None,
                                                'test_result':'1'})
    print "Adding an aggregation type for current data"
    currentdata_test = iatidq.dqtests.test_by_test_name(
        "activity-date[@type='start-planned']/@iso-date or transaction-date/@iso-date (for each transaction) is less than 13 months ago?"
        )
    iatidq.dqaggregationtypes.addAggregationType({'name':'Current data',
                                                'description': '',
                                                'test_id':currentdata_test.id,
                                                'test_result':'1'})

def create_inforesult_types(options):
    print "Adding an aggregation type"
    iatidq.inforesult.add_type("coverage", "Coverage")

def setup_common():
    print "Creating DB"
    iatidq.db.create_all()
    print "Adding hardcoded tests"
    iatidq.dqimporttests.hardcodedTests()
    print "Importing tests"
    iatidq.dqimporttests.importTestsFromFile(
        default_tests_filename,
        test_level.ACTIVITY)
    print "Importing indicators"
    iatidq.dqindicators.importIndicatorsFromFile(
        default_indicator_group_name,
        default_tests_filename)
    print "Importing indicator descriptions"
    iatidq.dqindicators.importIndicatorDescriptionsFromFile("pwyf2013",
                                                            "tests/indicators.csv")
    print "Importing codelists"
    iatidq.dqcodelists.importCodelists()

def setup_minimal(options):
    setup_common()

    print "Creating packages"
    pkg_names = [i[0] for i in which_packages]

    if pkg_names is not None:
        [ dqregistry.refresh_package_by_name(name) for name in pkg_names ]
    else:
        print "No packages are defined in quickstart"

    print "Creating aggregation types."
    create_aggregation_types(options)

    create_inforesult_types(options)

    print "Setup complete."

    for organisation in default_minimal_organisations:
        inserted_organisation = iatidq.dqorganisations.addOrganisation(
            organisation)
        if inserted_organisation is False:
            inserted_organisation = iatidq.dqorganisations.organisations(
                organisation['organisation_code'])
        thepackage = models.Package.query.filter_by(
            package_name=organisation['package_name']
                ).first()

        organisationpackage_data = {
            "organisation_id": inserted_organisation.id,
            "package_id": thepackage.id,
            "condition": organisation["condition"]
            }
        iatidq.dqorganisations.addOrganisationPackage(organisationpackage_data)
    print "Complete"

def setup(options):
    setup_common()
    print "Refreshing package data from Registry"
    dqregistry.refresh_packages()
    print "Adding organisations"
    iatidq.dqorganisations.importOrganisationPackagesFromFile("tests/organisations_with_identifiers.csv")
    create_aggregation_types(options)
    create_inforesult_types(options)
    print "Setup complete."

def enqueue_test(options):
    assert options.package_name
    assert options.filename

    iatidq.dqruntests.enqueue_package_for_test(options.filename,
                                               options.package_name)
def aggregate_results(options):
    assert options.runtime_id
    assert options.package_id
    iatidq.dqprocessing.aggregate_results(options.runtime_id, options.package_id)

commands = {
    "drop_db": (drop_all, "Delete DB"),
    "init_db": (init_db, "Initialise DB"),
    "enroll_tests": (enroll_tests, "Enroll a CSV file of tests"),
    "clear_revisionid": (clear_revisionid, "Clear CKAN revision ids"),
    "import_codelists": (import_codelists, "Import codelists"),
    "download": (download, "Download packages"),
    "import_indicators": (
        import_indicators,
        "Import indicators. Will try to assign indicators to existing tests."),
    "import_organisations": (
        import_organisations,
        "Import organisations. Will try to create and assign organisations "
        "to existing packages."),
    "setup": (setup, """Quick setup. Will init db, add tests, add codelists,
                      add indicators, refresh package data from Registry."""),
    "enqueue_test": (enqueue_test, "Set a package to be tested (with --package)"),
    "refresh": (refresh, "Refresh"),
    "activate_packages": (activate_packages, "Mark all packages as active"),
    "create_aggregation_types": (create_aggregation_types, "Create basic aggregation types."),
    "aggregate_results": (aggregate_results, "Trigger result aggregation"),
    "setup_minimal": (setup_minimal, "Quick minimal setup"),
    "create_inforesult_types": (create_inforesult_types, "Create basic infroresult types.")
}

def main():
    p = optparse.OptionParser()

    for k, v in commands.iteritems():
        handler, help_text = v
        option_name = "--" + k.replace("_", "-")
        p.add_option(option_name, dest=k, action="store_true", default=False, help=help_text)

    p.add_option("--runtime-id", dest="runtime_id",
                 type=int,
                 help="Runtime id (integer)")
    p.add_option("--package-id", dest="package_id",
                 type=int,
                 help="Package id (integer)")
    p.add_option("--level", dest="level",
                 type="int",
                 default=1,
                 help="Test level (e.g., 1 == Activity)")
    p.add_option("--minimal-pkgs", dest="minimal_pkgs",
                 action="store_true",
                 default=False,
                 help="Operate on a minimal set of packages")
    p.add_option("--package", dest="package_name",
                 help="Set name of package to be tested")
    p.add_option("--filename", dest="filename",
                 help="Set filename of data to test")
    p.add_option("--local-folder", dest="local_folder",
                 help="Set local folder where data to test is stored")

    options, args = p.parse_args()

    for mode, handler_ in commands.iteritems():
        handler, _ = handler_
        if getattr(options, mode, None):
            handler(options)
            return

    usage()

def usage():
    print "You need to specify which mode to run under"
    sys.exit(1)

if __name__ == '__main__':
    main()
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
#
# 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.

import json
import time

from airflow.contrib.hooks.gcs_hook import GoogleCloudStorageHook
from airflow.hooks.mysql_hook import MySqlHook
from airflow.models import BaseOperator
from airflow.utils.decorators import apply_defaults
from datetime import date, datetime
from decimal import Decimal
from MySQLdb.constants import FIELD_TYPE
from tempfile import NamedTemporaryFile

class MySqlToGoogleCloudStorageOperator(BaseOperator):
    """
    Copy data from MySQL to Google cloud storage in JSON format.
    """
    template_fields = ('sql', 'bucket', 'filename', 'schema_filename')
    template_ext = ('.sql',)
    ui_color = '#a0e08c'

    @apply_defaults
    def __init__(self,
                 sql,
                 bucket,
                 filename,
                 schema_filename=None,
                 approx_max_file_size_bytes=1900000000,
                 mysql_conn_id='mysql_default',
                 google_cloud_storage_conn_id='google_cloud_storage_default',
                 delegate_to=None,
                 *args,
                 **kwargs):
        """
        :param sql: The SQL to execute on the MySQL table.
        :type sql: string
        :param bucket: The bucket to upload to.
        :type bucket: string
        :param filename: The filename to use as the object name when uploading
            to Google cloud storage. A {} should be specified in the filename
            to allow the operator to inject file numbers in cases where the
            file is split due to size.
        :type filename: string
        :param schema_filename: If set, the filename to use as the object name
            when uploading a .json file containing the BigQuery schema fields
            for the table that was dumped from MySQL.
        :type schema_filename: string
        :param approx_max_file_size_bytes: This operator supports the ability
            to split large table dumps into multiple files (see notes in the
            filenamed param docs above). Google cloud storage allows for files
            to be a maximum of 4GB. This param allows developers to specify the
            file size of the splits.
        :type approx_max_file_size_bytes: long
        :param mysql_conn_id: Reference to a specific MySQL hook.
        :type mysql_conn_id: string
        :param google_cloud_storage_conn_id: Reference to a specific Google
            cloud storage hook.
        :type google_cloud_storage_conn_id: string
        :param delegate_to: The account to impersonate, if any. For this to
            work, the service account making the request must have domain-wide
            delegation enabled.
        """
        super(MySqlToGoogleCloudStorageOperator, self).__init__(*args, **kwargs)
        self.sql = sql
        self.bucket = bucket
        self.filename = filename
        self.schema_filename = schema_filename
        self.approx_max_file_size_bytes = approx_max_file_size_bytes
        self.mysql_conn_id = mysql_conn_id
        self.google_cloud_storage_conn_id = google_cloud_storage_conn_id
        self.delegate_to = delegate_to

    def execute(self, context):
        cursor = self._query_mysql()
        files_to_upload = self._write_local_data_files(cursor)

        # If a schema is set, create a BQ schema JSON file.
        if self.schema_filename:
            files_to_upload.update(self._write_local_schema_file(cursor))

        # Flush all files before uploading
        for file_handle in files_to_upload.values():
            file_handle.flush()

        self._upload_to_gcs(files_to_upload)

        # Close all temp file handles.
        for file_handle in files_to_upload.values():
            file_handle.close()

    def _query_mysql(self):
        """
        Queries mysql and returns a cursor to the results.
        """
        mysql = MySqlHook(mysql_conn_id=self.mysql_conn_id)
        conn = mysql.get_conn()
        cursor = conn.cursor()
        cursor.execute(self.sql)
        return cursor

    def _write_local_data_files(self, cursor):
        """
        Takes a cursor, and writes results to a local file.

        :return: A dictionary where keys are filenames to be used as object
            names in GCS, and values are file handles to local files that
            contain the data for the GCS objects.
        """
        schema = map(lambda schema_tuple: schema_tuple[0], cursor.description)
        file_no = 0
        tmp_file_handle = NamedTemporaryFile(delete=True)
        tmp_file_handles = {self.filename.format(file_no): tmp_file_handle}

        for row in cursor:
            # Convert datetime objects to utc seconds, and decimals to floats
            row = map(self.convert_types, row)
            row_dict = dict(zip(schema, row))

            # TODO validate that row isn't > 2MB. BQ enforces a hard row size of 2MB.
            json.dump(row_dict, tmp_file_handle)

            # Append newline to make dumps BigQuery compatible.
            tmp_file_handle.write('\n')

            # Stop if the file exceeds the file size limit.
            if tmp_file_handle.tell() >= self.approx_max_file_size_bytes:
                file_no += 1
                tmp_file_handle = NamedTemporaryFile(delete=True)
                tmp_file_handles[self.filename.format(file_no)] = tmp_file_handle

        return tmp_file_handles

    def _write_local_schema_file(self, cursor):
        """
        Takes a cursor, and writes the BigQuery schema for the results to a
        local file system.

        :return: A dictionary where key is a filename to be used as an object
            name in GCS, and values are file handles to local files that
            contains the BigQuery schema fields in .json format.
        """
        schema = []
        for field in cursor.description:
            # See PEP 249 for details about the description tuple.
            field_name = field[0]
            field_type = self.type_map(field[1])
            # Always allow TIMESTAMP to be nullable. MySQLdb returns None types
            # for required fields because some MySQL timestamps can't be
            # represented by Python's datetime (e.g. 0000-00-00 00:00:00).
            field_mode = 'NULLABLE' if field[6] or field_type == 'TIMESTAMP' else 'REQUIRED'
            schema.append({
                'name': field_name,
                'type': field_type,
                'mode': field_mode,
            })

        self.logger.info('Using schema for %s: %s', self.schema_filename, schema)
        tmp_schema_file_handle = NamedTemporaryFile(delete=True)
        json.dump(schema, tmp_schema_file_handle)
        return {self.schema_filename: tmp_schema_file_handle}

    def _upload_to_gcs(self, files_to_upload):
        """
        Upload all of the file splits (and optionally the schema .json file) to
        Google cloud storage.
        """
        hook = GoogleCloudStorageHook(google_cloud_storage_conn_id=self.google_cloud_storage_conn_id,
                                      delegate_to=self.delegate_to)
        for object, tmp_file_handle in files_to_upload.items():
            hook.upload(self.bucket, object, tmp_file_handle.name, 'application/json')

    @classmethod
    def convert_types(cls, value):
        """
        Takes a value from MySQLdb, and converts it to a value that's safe for
        JSON/Google cloud storage/BigQuery. Dates are converted to UTC seconds.
        Decimals are converted to floats.
        """
        if type(value) in (datetime, date):
            return time.mktime(value.timetuple())
        elif isinstance(value, Decimal):
            return float(value)
        else:
            return value

    @classmethod
    def type_map(cls, mysql_type):
        """
        Helper function that maps from MySQL fields to BigQuery fields. Used
        when a schema_filename is set.
        """
        d = {
            FIELD_TYPE.INT24: 'INTEGER',
            FIELD_TYPE.TINY: 'INTEGER',
            FIELD_TYPE.BIT: 'INTEGER',
            FIELD_TYPE.DATETIME: 'TIMESTAMP',
            FIELD_TYPE.DATE: 'TIMESTAMP',
            FIELD_TYPE.DECIMAL: 'FLOAT',
            FIELD_TYPE.NEWDECIMAL: 'FLOAT',
            FIELD_TYPE.DOUBLE: 'FLOAT',
            FIELD_TYPE.FLOAT: 'FLOAT',
            FIELD_TYPE.INT24: 'INTEGER',
            FIELD_TYPE.LONG: 'INTEGER',
            FIELD_TYPE.LONGLONG: 'INTEGER',
            FIELD_TYPE.SHORT: 'INTEGER',
            FIELD_TYPE.TIMESTAMP: 'TIMESTAMP',
            FIELD_TYPE.YEAR: 'INTEGER',
        }
        return d[mysql_type] if mysql_type in d else 'STRING'
<commit_msg>[AIRFLOW-1613] Make MySqlToGoogleCloudStorageOperator compaitible with python3<commit_after># -*- coding: utf-8 -*-
#
# 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.

import json
import time

from airflow.contrib.hooks.gcs_hook import GoogleCloudStorageHook
from airflow.hooks.mysql_hook import MySqlHook
from airflow.models import BaseOperator
from airflow.utils.decorators import apply_defaults
from datetime import date, datetime
from decimal import Decimal
from MySQLdb.constants import FIELD_TYPE
from tempfile import NamedTemporaryFile

class MySqlToGoogleCloudStorageOperator(BaseOperator):
    """
    Copy data from MySQL to Google cloud storage in JSON format.
    """
    template_fields = ('sql', 'bucket', 'filename', 'schema_filename')
    template_ext = ('.sql',)
    ui_color = '#a0e08c'

    @apply_defaults
    def __init__(self,
                 sql,
                 bucket,
                 filename,
                 schema_filename=None,
                 approx_max_file_size_bytes=1900000000,
                 mysql_conn_id='mysql_default',
                 google_cloud_storage_conn_id='google_cloud_storage_default',
                 delegate_to=None,
                 *args,
                 **kwargs):
        """
        :param sql: The SQL to execute on the MySQL table.
        :type sql: string
        :param bucket: The bucket to upload to.
        :type bucket: string
        :param filename: The filename to use as the object name when uploading
            to Google cloud storage. A {} should be specified in the filename
            to allow the operator to inject file numbers in cases where the
            file is split due to size.
        :type filename: string
        :param schema_filename: If set, the filename to use as the object name
            when uploading a .json file containing the BigQuery schema fields
            for the table that was dumped from MySQL.
        :type schema_filename: string
        :param approx_max_file_size_bytes: This operator supports the ability
            to split large table dumps into multiple files (see notes in the
            filenamed param docs above). Google cloud storage allows for files
            to be a maximum of 4GB. This param allows developers to specify the
            file size of the splits.
        :type approx_max_file_size_bytes: long
        :param mysql_conn_id: Reference to a specific MySQL hook.
        :type mysql_conn_id: string
        :param google_cloud_storage_conn_id: Reference to a specific Google
            cloud storage hook.
        :type google_cloud_storage_conn_id: string
        :param delegate_to: The account to impersonate, if any. For this to
            work, the service account making the request must have domain-wide
            delegation enabled.
        """
        super(MySqlToGoogleCloudStorageOperator, self).__init__(*args, **kwargs)
        self.sql = sql
        self.bucket = bucket
        self.filename = filename
        self.schema_filename = schema_filename
        self.approx_max_file_size_bytes = approx_max_file_size_bytes
        self.mysql_conn_id = mysql_conn_id
        self.google_cloud_storage_conn_id = google_cloud_storage_conn_id
        self.delegate_to = delegate_to

    def execute(self, context):
        cursor = self._query_mysql()
        files_to_upload = self._write_local_data_files(cursor)

        # If a schema is set, create a BQ schema JSON file.
        if self.schema_filename:
            files_to_upload.update(self._write_local_schema_file(cursor))

        # Flush all files before uploading
        for file_handle in files_to_upload.values():
            file_handle.flush()

        self._upload_to_gcs(files_to_upload)

        # Close all temp file handles.
        for file_handle in files_to_upload.values():
            file_handle.close()

    def _query_mysql(self):
        """
        Queries mysql and returns a cursor to the results.
        """
        mysql = MySqlHook(mysql_conn_id=self.mysql_conn_id)
        conn = mysql.get_conn()
        cursor = conn.cursor()
        cursor.execute(self.sql)
        return cursor

    def _write_local_data_files(self, cursor):
        """
        Takes a cursor, and writes results to a local file.

        :return: A dictionary where keys are filenames to be used as object
            names in GCS, and values are file handles to local files that
            contain the data for the GCS objects.
        """
        schema = list(map(lambda schema_tuple: schema_tuple[0], cursor.description))
        file_no = 0
        tmp_file_handle = NamedTemporaryFile(mode='w', delete=True)
        tmp_file_handles = {self.filename.format(file_no): tmp_file_handle}

        for row in cursor:
            # Convert datetime objects to utc seconds, and decimals to floats
            row = map(self.convert_types, row)
            row_dict = dict(zip(schema, row))

            # TODO validate that row isn't > 2MB. BQ enforces a hard row size of 2MB.
            json.dump(row_dict, tmp_file_handle)

            # Append newline to make dumps BigQuery compatible.
            tmp_file_handle.write('\n')

            # Stop if the file exceeds the file size limit.
            if tmp_file_handle.tell() >= self.approx_max_file_size_bytes:
                file_no += 1
                tmp_file_handle = NamedTemporaryFile(mode='w', delete=True)
                tmp_file_handles[self.filename.format(file_no)] = tmp_file_handle

        return tmp_file_handles

    def _write_local_schema_file(self, cursor):
        """
        Takes a cursor, and writes the BigQuery schema for the results to a
        local file system.

        :return: A dictionary where key is a filename to be used as an object
            name in GCS, and values are file handles to local files that
            contains the BigQuery schema fields in .json format.
        """
        schema = []
        for field in cursor.description:
            # See PEP 249 for details about the description tuple.
            field_name = field[0]
            field_type = self.type_map(field[1])
            # Always allow TIMESTAMP to be nullable. MySQLdb returns None types
            # for required fields because some MySQL timestamps can't be
            # represented by Python's datetime (e.g. 0000-00-00 00:00:00).
            field_mode = 'NULLABLE' if field[6] or field_type == 'TIMESTAMP' else 'REQUIRED'
            schema.append({
                'name': field_name,
                'type': field_type,
                'mode': field_mode,
            })

        self.logger.info('Using schema for %s: %s', self.schema_filename, schema)
        tmp_schema_file_handle = NamedTemporaryFile(mode='w', delete=True)
        json.dump(schema, tmp_schema_file_handle)
        return {self.schema_filename: tmp_schema_file_handle}

    def _upload_to_gcs(self, files_to_upload):
        """
        Upload all of the file splits (and optionally the schema .json file) to
        Google cloud storage.
        """
        hook = GoogleCloudStorageHook(google_cloud_storage_conn_id=self.google_cloud_storage_conn_id,
                                      delegate_to=self.delegate_to)
        for object, tmp_file_handle in files_to_upload.items():
            hook.upload(self.bucket, object, tmp_file_handle.name, 'application/json')

    @classmethod
    def convert_types(cls, value):
        """
        Takes a value from MySQLdb, and converts it to a value that's safe for
        JSON/Google cloud storage/BigQuery. Dates are converted to UTC seconds.
        Decimals are converted to floats.
        """
        if type(value) in (datetime, date):
            return time.mktime(value.timetuple())
        elif isinstance(value, Decimal):
            return float(value)
        else:
            return value

    @classmethod
    def type_map(cls, mysql_type):
        """
        Helper function that maps from MySQL fields to BigQuery fields. Used
        when a schema_filename is set.
        """
        d = {
            FIELD_TYPE.INT24: 'INTEGER',
            FIELD_TYPE.TINY: 'INTEGER',
            FIELD_TYPE.BIT: 'INTEGER',
            FIELD_TYPE.DATETIME: 'TIMESTAMP',
            FIELD_TYPE.DATE: 'TIMESTAMP',
            FIELD_TYPE.DECIMAL: 'FLOAT',
            FIELD_TYPE.NEWDECIMAL: 'FLOAT',
            FIELD_TYPE.DOUBLE: 'FLOAT',
            FIELD_TYPE.FLOAT: 'FLOAT',
            FIELD_TYPE.INT24: 'INTEGER',
            FIELD_TYPE.LONG: 'INTEGER',
            FIELD_TYPE.LONGLONG: 'INTEGER',
            FIELD_TYPE.SHORT: 'INTEGER',
            FIELD_TYPE.TIMESTAMP: 'TIMESTAMP',
            FIELD_TYPE.YEAR: 'INTEGER',
        }
        return d[mysql_type] if mysql_type in d else 'STRING'
<|endoftext|>


<commit_before>
import discord
from discord.ext import commands, tasks

import asyncio
import datetime

import parsedatetime

from utilities import checks

def setup(bot):
	bot.add_cog(Reminders(bot))

class Reminders(commands.Cog):

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

		self.calendar = parsedatetime.Calendar()
		# Patch https://github.com/bear/parsedatetime/commit/7a759c1f8ff7563f12ac2c1f2ea0b41452f61dec
		# until fix is released
		if "secss" in self.calendar.ptc.units["seconds"]:
			self.calendar.ptc.units["seconds"].append("secs")
			self.calendar.ptc.units["seconds"].remove("secss")
			self.calendar.ptc.units["seconds"].append('s')
		# Add mo as valid abbreviation for month
		self.calendar.ptc.units["months"].append("mo")

		self.current_timer = None
		self.new_reminder = asyncio.Event()
		self.restarting_timer = False
		self.timer.start().set_name("Reminders")

	def cog_unload(self):
		self.timer.cancel()

	async def initialize_database(self):
		await self.bot.connect_to_database()
		await self.bot.db.execute("CREATE SCHEMA IF NOT EXISTS reminders")
		await self.bot.db.execute(
			"""
			CREATE TABLE IF NOT EXISTS reminders.reminders (
				id				INT PRIMARY KEY GENERATED ALWAYS AS IDENTITY,
				user_id			BIGINT,
				channel_id		BIGINT,
				message_id		BIGINT,
				created_time	TIMESTAMPTZ DEFAULT NOW(),
				remind_time		TIMESTAMPTZ,
				reminder		TEXT,
				reminded		BOOL DEFAULT FALSE,
				failed			BOOL DEFAULT FALSE
			)
			"""
		)

	async def cog_check(self, ctx):
		return await checks.not_forbidden().predicate(ctx)

	@commands.group(name = "reminder", aliases = ["remind", "timer"],
					invoke_without_command = True, case_insensitive = True)
	async def reminder_command(self, ctx, *, reminder: commands.clean_content):
		'''Set reminders'''
		# TODO: Allow setting timezone (+ for time command as well)
		# Clean reminder input
		for prefix in ("me about ", "me to ", "me "):
			if reminder.startswith(prefix):
				reminder = reminder[len(prefix):]
		reminder = reminder.replace("from now", "")
		# Parse reminder
		now = datetime.datetime.now(datetime.timezone.utc)
		if not (matches := self.calendar.nlp(reminder, sourceTime = now)):
			raise commands.BadArgument("Time not specified")
		parsed_datetime, flags, start_pos, end_pos, matched_text = matches[0]
		parsed_datetime = parsed_datetime.replace(tzinfo = datetime.timezone.utc)
		if parsed_datetime < now:
			raise commands.BadArgument("Time is in the past")
		# Respond
		reminder = reminder[:start_pos] + reminder[end_pos + 1:]
		reminder = reminder.strip()
		response = await ctx.embed_reply(fields = (("Reminder", reminder or ctx.bot.ZWS),),
											footer_text = f"Set for {parsed_datetime.isoformat()}",
											timestamp = parsed_datetime)
		# Insert into database
		created_time = ctx.message.created_at.replace(tzinfo = datetime.timezone.utc)
		await self.bot.db.execute(
			"""
			INSERT INTO reminders.reminders (user_id, channel_id, message_id, created_time, remind_time, reminder)
			VALUES ($1, $2, $3, $4, $5, $6)
			""",
			ctx.author.id, ctx.channel.id, response.id, created_time, parsed_datetime, reminder
		)
		# Update timer
		if self.current_timer and parsed_datetime < self.current_timer:
			self.restarting_timer = True
			self.timer.restart()
			self.timer.get_task().set_name("Reminders")
		else:
			self.new_reminder.set()

	# TODO: reminders command / list subcommand
	# TODO: delete/cancel/remove subcommand
	# TODO: clear subcommand

	# R/PT0S
	@tasks.loop()
	async def timer(self):
		record = await self.bot.db.fetchrow(
			"""
			SELECT * FROM reminders.reminders
			WHERE reminded = FALSE AND failed = FALSE
			ORDER BY remind_time
			LIMIT 1
			"""
		)
		if not record:
			self.new_reminder.clear()
			return await self.new_reminder.wait()
		self.current_timer = record["remind_time"]
		if record["remind_time"] > (now := datetime.datetime.now(datetime.timezone.utc)):
			await asyncio.sleep((record["remind_time"] - now).total_seconds())
		if not (channel := self.bot.get_channel(record["channel_id"])):
			# TODO: Attempt to fetch channel?
			return await self.bot.db.execute("UPDATE reminders.reminders SET failed = TRUE WHERE id = $1", record["id"])
		user = self.bot.get_user(record["user_id"]) or self.bot.fetch_user(record["user_id"])
		# TODO: Handle user not found?
		embed = discord.Embed(color = self.bot.bot_color)
		try:
			message = await channel.fetch_message(record["message_id"])
			embed.description = f"[{record['reminder'] or 'Reminder'}]({message.jump_url})"
		except discord.NotFound:
			embed.description = record["reminder"] or "Reminder"
		embed.set_footer(text = "Reminder set")
		embed.timestamp = record["created_time"]
		try:
			await channel.send(user.mention, embed = embed)
		except discord.Forbidden:
			# TODO: Attempt to send without embed
			# TODO: Fall back to DM
			await self.bot.db.execute("UPDATE reminders.reminders SET failed = TRUE WHERE id = $1", record["id"])
		else:
			await self.bot.db.execute("UPDATE reminders.reminders SET reminded = TRUE WHERE id = $1", record["id"])

	@timer.before_loop
	async def before_timer(self):
		await self.initialize_database()
		await self.bot.wait_until_ready()

	@timer.after_loop
	async def after_timer(self):
		if self.restarting_timer:
			self.restarting_timer = False
		else:
			print(f"{self.bot.console_message_prefix}Reminders task cancelled @ {datetime.datetime.now().isoformat()}")

<commit_msg>[Discord] Improve clarity of Reminders cog initialization<commit_after>
import discord
from discord.ext import commands, tasks

import asyncio
import datetime

import parsedatetime

from utilities import checks

def setup(bot):
	bot.add_cog(Reminders(bot))

class Reminders(commands.Cog):

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

		self.calendar = parsedatetime.Calendar()
		# Patch https://github.com/bear/parsedatetime/commit/7a759c1f8ff7563f12ac2c1f2ea0b41452f61dec
		# until fix is released
		if "secss" in self.calendar.ptc.units["seconds"]:
			self.calendar.ptc.units["seconds"].remove("secss")
			self.calendar.ptc.units["seconds"].append("secs")
			self.calendar.ptc.units["seconds"].append('s')
		# Add mo as valid abbreviation for month
		self.calendar.ptc.units["months"].append("mo")

		self.current_timer = None
		self.new_reminder = asyncio.Event()
		self.restarting_timer = False
		self.timer.start().set_name("Reminders")

	def cog_unload(self):
		self.timer.cancel()

	async def initialize_database(self):
		await self.bot.connect_to_database()
		await self.bot.db.execute("CREATE SCHEMA IF NOT EXISTS reminders")
		await self.bot.db.execute(
			"""
			CREATE TABLE IF NOT EXISTS reminders.reminders (
				id				INT PRIMARY KEY GENERATED ALWAYS AS IDENTITY,
				user_id			BIGINT,
				channel_id		BIGINT,
				message_id		BIGINT,
				created_time	TIMESTAMPTZ DEFAULT NOW(),
				remind_time		TIMESTAMPTZ,
				reminder		TEXT,
				reminded		BOOL DEFAULT FALSE,
				failed			BOOL DEFAULT FALSE
			)
			"""
		)

	async def cog_check(self, ctx):
		return await checks.not_forbidden().predicate(ctx)

	@commands.group(name = "reminder", aliases = ["remind", "timer"],
					invoke_without_command = True, case_insensitive = True)
	async def reminder_command(self, ctx, *, reminder: commands.clean_content):
		'''Set reminders'''
		# TODO: Allow setting timezone (+ for time command as well)
		# Clean reminder input
		for prefix in ("me about ", "me to ", "me "):
			if reminder.startswith(prefix):
				reminder = reminder[len(prefix):]
		reminder = reminder.replace("from now", "")
		# Parse reminder
		now = datetime.datetime.now(datetime.timezone.utc)
		if not (matches := self.calendar.nlp(reminder, sourceTime = now)):
			raise commands.BadArgument("Time not specified")
		parsed_datetime, flags, start_pos, end_pos, matched_text = matches[0]
		parsed_datetime = parsed_datetime.replace(tzinfo = datetime.timezone.utc)
		if parsed_datetime < now:
			raise commands.BadArgument("Time is in the past")
		# Respond
		reminder = reminder[:start_pos] + reminder[end_pos + 1:]
		reminder = reminder.strip()
		response = await ctx.embed_reply(fields = (("Reminder", reminder or ctx.bot.ZWS),),
											footer_text = f"Set for {parsed_datetime.isoformat()}",
											timestamp = parsed_datetime)
		# Insert into database
		created_time = ctx.message.created_at.replace(tzinfo = datetime.timezone.utc)
		await self.bot.db.execute(
			"""
			INSERT INTO reminders.reminders (user_id, channel_id, message_id, created_time, remind_time, reminder)
			VALUES ($1, $2, $3, $4, $5, $6)
			""",
			ctx.author.id, ctx.channel.id, response.id, created_time, parsed_datetime, reminder
		)
		# Update timer
		if self.current_timer and parsed_datetime < self.current_timer:
			self.restarting_timer = True
			self.timer.restart()
			self.timer.get_task().set_name("Reminders")
		else:
			self.new_reminder.set()

	# TODO: reminders command / list subcommand
	# TODO: delete/cancel/remove subcommand
	# TODO: clear subcommand

	# R/PT0S
	@tasks.loop()
	async def timer(self):
		record = await self.bot.db.fetchrow(
			"""
			SELECT * FROM reminders.reminders
			WHERE reminded = FALSE AND failed = FALSE
			ORDER BY remind_time
			LIMIT 1
			"""
		)
		if not record:
			self.new_reminder.clear()
			return await self.new_reminder.wait()
		self.current_timer = record["remind_time"]
		if record["remind_time"] > (now := datetime.datetime.now(datetime.timezone.utc)):
			await asyncio.sleep((record["remind_time"] - now).total_seconds())
		if not (channel := self.bot.get_channel(record["channel_id"])):
			# TODO: Attempt to fetch channel?
			return await self.bot.db.execute("UPDATE reminders.reminders SET failed = TRUE WHERE id = $1", record["id"])
		user = self.bot.get_user(record["user_id"]) or self.bot.fetch_user(record["user_id"])
		# TODO: Handle user not found?
		embed = discord.Embed(color = self.bot.bot_color)
		try:
			message = await channel.fetch_message(record["message_id"])
			embed.description = f"[{record['reminder'] or 'Reminder'}]({message.jump_url})"
		except discord.NotFound:
			embed.description = record["reminder"] or "Reminder"
		embed.set_footer(text = "Reminder set")
		embed.timestamp = record["created_time"]
		try:
			await channel.send(user.mention, embed = embed)
		except discord.Forbidden:
			# TODO: Attempt to send without embed
			# TODO: Fall back to DM
			await self.bot.db.execute("UPDATE reminders.reminders SET failed = TRUE WHERE id = $1", record["id"])
		else:
			await self.bot.db.execute("UPDATE reminders.reminders SET reminded = TRUE WHERE id = $1", record["id"])

	@timer.before_loop
	async def before_timer(self):
		await self.initialize_database()
		await self.bot.wait_until_ready()

	@timer.after_loop
	async def after_timer(self):
		if self.restarting_timer:
			self.restarting_timer = False
		else:
			print(f"{self.bot.console_message_prefix}Reminders task cancelled @ {datetime.datetime.now().isoformat()}")

<|endoftext|>


<commit_before>"""
  Core Identity models for atmosphere.
Note:
  Multiple users can 'own' an identity (IdentityMembership - group.py)
"""

from datetime import timedelta

from django.db import models

from threepio import logger

class Identity(models.Model):
    """
    An Identity is the minimal set of credentials necessary
    to authenticate against a single provider
    """

    created_by = models.ForeignKey("AtmosphereUser")
    provider = models.ForeignKey("Provider")

    @classmethod
    def delete_identity(cls, username, provider_location):
        #Do not move up. ImportError.
        from core.models import AtmosphereUser, Group, Credential, Quota,\
            Provider, AccountProvider,\
            IdentityMembership, ProviderMembership

        provider = Provider.objects.get(location__iexact=provider_location)
        user = AtmosphereUser.objects.get(username=username)
        group = Group.objects.get(name=username)
        my_ids = Identity.objects.filter(
            created_by=user, provider=provider)
        for ident in my_ids:
            membership_set = ident.identitymembership_set.all()
            membership_set.delete()
            ident.delete()
        group.delete()
        user.delete()
        return

    def can_share(self, django_user):
        """
        You CAN share an identity IFF:
        0. You are a staff user
        1. You are the original owner of the identity
        2. You are the leader of a group who contains the owner of the identity
        """
        #This person leads a group, may be able to share.
        #Check 0
        if django_user.is_staff:
            return True
        #Check 1
        original_owner = self.created_by
        if original_owner == django_user:
            return True
        #Check 2
        shared = False
        leader_groups = django_user.group_set.get(leaders__in=[django_user])
        for group in leader_groups:
            id_member = g.identitymembership_set.get(identity=self)
            if not id_member:
                continue
            #ASSERT: You have SHARED access to the identity
            shared = True
            if original_owner in group.user_set.all():
                return True
        #User can't share.. Log the attempt for record-keeping
        if shared:
            logger.info("FAILED SHARE ATTEMPT: User:%s Identity:%s "
                        "Reason: You are not a leader of any group that "
                        "contains the actual owner of the identity (%s)."
                        % (django_user, self, original_owner))
        else:
            logger.info("FAILED SHARE ATTEMPT: User:%s Identity:%s "
                        % (django_user, self))

        return False

    def share(self, core_group, quota=None):
        """
        """
        from core.models import IdentityMembership, ProviderMembership
        existing_membership = IdentityMembership.objects.filter(
            member=core_group, identity=self)
        if existing_membership:
            return existing_membership[0]

        #User does not already have membership - Check for provider membership
        prov_membership = ProviderMembership.objects.filter(
            member=core_group, provider=self.provider)
        if not prov_membership:
            raise Exception("Cannot share identity membership before the"
                            " provider is shared")

        #Ready to create new membership for this group
        if not quota:
            quota = Quota.default_quota()
        allocation = Allocation.default_allocation()
        new_membership = IdentityMembership.objects.get_or_create(
            member=core_group, identity=self, quota=quota, allocation=allocation)[0]
        return new_membership

    def unshare(self, core_group):
        """
        Potential problem:
        1. User X creates/imports an openstack account (& is the owner),
        2. User X shares identitymembership with User Y,
        3. User X or Y tries to unshare IdentityMembership with the opposing
        user.

        Solution:
        ONLY unshare if this user is the original owner of the identity
        """
        from core.models import IdentityMembership
        existing_membership = IdentityMembership.objects.filter(
            member=core_group, identity=self)
        return existing_membership[0].delete()

    def get_membership(self):
        identity_members = self.identitymembership_set.all()
        group_names = [id_member.member for id_member in identity_members]
        #TODO: Add 'rules' if we want to hide specific users (staff, etc.)
        return group_names

    @classmethod
    def create_identity(cls, username, provider_location,
                        max_quota=False, account_admin=False, **kwarg_creds):
        """
        Create new User/Group & Identity for given provider_location
        NOTES:
        * kwargs prefixed with 'cred_' will be collected as credentials
        * Can assign optional flags:
          + max_quota - Assign the highest quota available, rather than
            default.
          + account_admin - Private Clouds only - This user should have ALL
            permissions including:
              * Image creation (Glance)
              * Account creation (Keystone)
              * Access to ALL instances launched over ALL users

          Atmosphere will run fine without an account_admin, but the above
          features will be disabled.
        """
        #Do not move up. ImportError.
        from core.models import Group, Credential, Quota,\
            Provider, AccountProvider,\
            IdentityMembership, ProviderMembership

        provider = Provider.objects.get(location__iexact=provider_location)

        credentials = {}
        for (c_key, c_value) in kwarg_creds.items():
            if 'cred_' not in c_key.lower():
                continue
            c_key = c_key.replace('cred_', '')
            credentials[c_key] = c_value

        (user, group) = Group.create_usergroup(username)

        #NOTE: This specific query will need to be modified if we want
        # 2+ Identities on a single provider

        id_membership = IdentityMembership.objects.filter(
            member__name=user.username,
            identity__provider=provider,
            identity__created_by__username=user.username)
        if not id_membership:
            #1. Create a Provider Membership
            p_membership = ProviderMembership.objects.get_or_create(
                provider=provider, member=group)[0]

            #2. Create an Identity Membership
            identity = Identity.objects.get_or_create(
                created_by=user, provider=provider)[0]
            #Two-tuple, (Object, created)
            id_membership = IdentityMembership.objects.get_or_create(
                identity=identity, member=group, quota=Quota.default_quota())
        #Either first in list OR object from two-tuple.. Its what we need.
        id_membership = id_membership[0]

        #ID_Membership exists.

        #3. Make sure that all kwargs exist as credentials
        # NOTE: Because we assume only one identity per provider
        #       We can add new credentials to
        #       existing identities if missing..
        # In the future it will be hard to determine when we want to
        # update values on an identity Vs. create a second, new
        # identity.
        for (c_key, c_value) in credentials.items():
            test_key_exists = Credential.objects.filter(
                identity=id_membership.identity,
                key=c_key)
            if test_key_exists:
                logger.info("Conflicting Key Error: Key:%s Value:%s "
                            "Replacement:%s" %
                            (c_key, c_value, test_key_exists[0].value))
                #No Dupes... But should we really throw an Exception here?
                continue
            Credential.objects.get_or_create(
                identity=id_membership.identity,
                key=c_key,
                value=c_value)[0]
        #4. Assign a different quota, if requested
        if max_quota:
            quota = Quota.max_quota()
            id_membership.quota = quota
            id_membership.save()
        if account_admin:
            admin = AccountProvider.objects.get_or_create(
                provider=id_membership.identity.provider,
                identity=id_membership.identity)[0]

        #5. Save the user to activate profile on first-time use
        user.save()
        #Return the identity
        return id_membership.identity

    def is_active(self):
        return self.provider.is_active()

    def creator_name(self):
        return self.created_by.username

    def get_credential(self, key):
        cred = self.credential_set.filter(key=key)
        return cred[0].value if cred else None

    def get_credentials(self):
        cred_dict = {}
        for cred in self.credential_set.all():
            cred_dict[cred.key] = cred.value
        return cred_dict

    def get_allocation(self):
        id_member = self.identitymembership_set.all()[0]
        return id_member.allocation

    def get_quota(self):
        id_member = self.identitymembership_set.all()[0]
        return id_member.quota

    def get_allocation_dict(self):
        id_member = self.identitymembership_set.all()[0]
        allocation_dict = id_member.get_allocation_dict()
        return allocation_dict

    def get_quota_dict(self):
        id_member = self.identitymembership_set.all()[0]
        #See core/models/group.py#IdentityMembership
        quota_dict = id_member.get_quota_dict()
        allocation_dict = self.get_allocation_dict()
        if allocation_dict:
            quota_dict.update({"allocation": allocation_dict})
        return quota_dict

    def json(self):
        return {
            'id': self.id,
            'creator': self.created_by.username,
            'provider': self.provider.json(),
            'credentials': [cred.json() for cred
                            in self.credential_set.order_by('key')],
        }

    def __unicode__(self):
        output = "%s %s - " % (self.provider, self.created_by.username)
        output += "Credentials {"
        for c in self.credential_set.order_by('key'):
            output += "%s, " % (c.key,)
        output = output[:-2] + "}"
        return output

    class Meta:
        db_table = "identity"
        app_label = "core"
        verbose_name_plural = "identities"
<commit_msg>Added Default Allocation to create_identity function so that users don't fall through cracks<commit_after>"""
  Core Identity models for atmosphere.
Note:
  Multiple users can 'own' an identity (IdentityMembership - group.py)
"""

from datetime import timedelta

from django.db import models

from threepio import logger

class Identity(models.Model):
    """
    An Identity is the minimal set of credentials necessary
    to authenticate against a single provider
    """

    created_by = models.ForeignKey("AtmosphereUser")
    provider = models.ForeignKey("Provider")

    @classmethod
    def delete_identity(cls, username, provider_location):
        #Do not move up. ImportError.
        from core.models import AtmosphereUser, Group, Credential, Quota,\
            Provider, AccountProvider,\
            IdentityMembership, ProviderMembership

        provider = Provider.objects.get(location__iexact=provider_location)
        user = AtmosphereUser.objects.get(username=username)
        group = Group.objects.get(name=username)
        my_ids = Identity.objects.filter(
            created_by=user, provider=provider)
        for ident in my_ids:
            membership_set = ident.identitymembership_set.all()
            membership_set.delete()
            ident.delete()
        group.delete()
        user.delete()
        return

    def can_share(self, django_user):
        """
        You CAN share an identity IFF:
        0. You are a staff user
        1. You are the original owner of the identity
        2. You are the leader of a group who contains the owner of the identity
        """
        #This person leads a group, may be able to share.
        #Check 0
        if django_user.is_staff:
            return True
        #Check 1
        original_owner = self.created_by
        if original_owner == django_user:
            return True
        #Check 2
        shared = False
        leader_groups = django_user.group_set.get(leaders__in=[django_user])
        for group in leader_groups:
            id_member = g.identitymembership_set.get(identity=self)
            if not id_member:
                continue
            #ASSERT: You have SHARED access to the identity
            shared = True
            if original_owner in group.user_set.all():
                return True
        #User can't share.. Log the attempt for record-keeping
        if shared:
            logger.info("FAILED SHARE ATTEMPT: User:%s Identity:%s "
                        "Reason: You are not a leader of any group that "
                        "contains the actual owner of the identity (%s)."
                        % (django_user, self, original_owner))
        else:
            logger.info("FAILED SHARE ATTEMPT: User:%s Identity:%s "
                        % (django_user, self))

        return False

    def share(self, core_group, quota=None):
        """
        """
        from core.models import IdentityMembership, ProviderMembership, Quota, Allocation
        existing_membership = IdentityMembership.objects.filter(
            member=core_group, identity=self)
        if existing_membership:
            return existing_membership[0]

        #User does not already have membership - Check for provider membership
        prov_membership = ProviderMembership.objects.filter(
            member=core_group, provider=self.provider)
        if not prov_membership:
            raise Exception("Cannot share identity membership before the"
                            " provider is shared")

        #Ready to create new membership for this group
        if not quota:
            quota = Quota.default_quota()
        allocation = Allocation.default_allocation()
        new_membership = IdentityMembership.objects.get_or_create(
            member=core_group, identity=self, quota=quota, allocation=allocation)[0]
        return new_membership

    def unshare(self, core_group):
        """
        Potential problem:
        1. User X creates/imports an openstack account (& is the owner),
        2. User X shares identitymembership with User Y,
        3. User X or Y tries to unshare IdentityMembership with the opposing
        user.

        Solution:
        ONLY unshare if this user is the original owner of the identity
        """
        from core.models import IdentityMembership
        existing_membership = IdentityMembership.objects.filter(
            member=core_group, identity=self)
        return existing_membership[0].delete()

    def get_membership(self):
        identity_members = self.identitymembership_set.all()
        group_names = [id_member.member for id_member in identity_members]
        #TODO: Add 'rules' if we want to hide specific users (staff, etc.)
        return group_names

    @classmethod
    def create_identity(cls, username, provider_location,
                        max_quota=False, account_admin=False, **kwarg_creds):
        """
        Create new User/Group & Identity for given provider_location
        NOTES:
        * kwargs prefixed with 'cred_' will be collected as credentials
        * Can assign optional flags:
          + max_quota - Assign the highest quota available, rather than
            default.
          + account_admin - Private Clouds only - This user should have ALL
            permissions including:
              * Image creation (Glance)
              * Account creation (Keystone)
              * Access to ALL instances launched over ALL users

          Atmosphere will run fine without an account_admin, but the above
          features will be disabled.
        """
        #Do not move up. ImportError.
        from core.models import Group, Credential, Quota,\
            Provider, AccountProvider, Allocation,\
            IdentityMembership, ProviderMembership

        provider = Provider.objects.get(location__iexact=provider_location)

        credentials = {}
        for (c_key, c_value) in kwarg_creds.items():
            if 'cred_' not in c_key.lower():
                continue
            c_key = c_key.replace('cred_', '')
            credentials[c_key] = c_value

        (user, group) = Group.create_usergroup(username)

        #NOTE: This specific query will need to be modified if we want
        # 2+ Identities on a single provider

        id_membership = IdentityMembership.objects.filter(
            member__name=user.username,
            identity__provider=provider,
            identity__created_by__username=user.username)
        if not id_membership:
            #1. Create a Provider Membership
            p_membership = ProviderMembership.objects.get_or_create(
                provider=provider, member=group)[0]

            default_allocation = Allocation.default_allocation()
            #2. Create an Identity Membership
            identity = Identity.objects.get_or_create(
                created_by=user, provider=provider)[0]
            #Two-tuple, (Object, created)
            id_membership = IdentityMembership.objects.get_or_create(
                identity=identity, member=group, allocation=default_allocation, quota=Quota.default_quota())
        #Either first in list OR object from two-tuple.. Its what we need.
        id_membership = id_membership[0]

        #ID_Membership exists.

        #3. Make sure that all kwargs exist as credentials
        # NOTE: Because we assume only one identity per provider
        #       We can add new credentials to
        #       existing identities if missing..
        # In the future it will be hard to determine when we want to
        # update values on an identity Vs. create a second, new
        # identity.
        for (c_key, c_value) in credentials.items():
            test_key_exists = Credential.objects.filter(
                identity=id_membership.identity,
                key=c_key)
            if test_key_exists:
                logger.info("Conflicting Key Error: Key:%s Value:%s "
                            "Replacement:%s" %
                            (c_key, c_value, test_key_exists[0].value))
                #No Dupes... But should we really throw an Exception here?
                continue
            Credential.objects.get_or_create(
                identity=id_membership.identity,
                key=c_key,
                value=c_value)[0]
        #4. Assign a different quota, if requested
        if max_quota:
            quota = Quota.max_quota()
            id_membership.quota = quota
            id_membership.allocation = None
            id_membership.save()
        if account_admin:
            admin = AccountProvider.objects.get_or_create(
                provider=id_membership.identity.provider,
                identity=id_membership.identity)[0]

        #5. Save the user to activate profile on first-time use
        user.save()
        #Return the identity
        return id_membership.identity

    def is_active(self):
        return self.provider.is_active()

    def creator_name(self):
        return self.created_by.username

    def get_credential(self, key):
        cred = self.credential_set.filter(key=key)
        return cred[0].value if cred else None

    def get_credentials(self):
        cred_dict = {}
        for cred in self.credential_set.all():
            cred_dict[cred.key] = cred.value
        return cred_dict

    def get_allocation(self):
        id_member = self.identitymembership_set.all()[0]
        return id_member.allocation

    def get_quota(self):
        id_member = self.identitymembership_set.all()[0]
        return id_member.quota

    def get_allocation_dict(self):
        id_member = self.identitymembership_set.all()[0]
        allocation_dict = id_member.get_allocation_dict()
        return allocation_dict

    def get_quota_dict(self):
        id_member = self.identitymembership_set.all()[0]
        #See core/models/group.py#IdentityMembership
        quota_dict = id_member.get_quota_dict()
        allocation_dict = self.get_allocation_dict()
        if allocation_dict:
            quota_dict.update({"allocation": allocation_dict})
        return quota_dict

    def json(self):
        return {
            'id': self.id,
            'creator': self.created_by.username,
            'provider': self.provider.json(),
            'credentials': [cred.json() for cred
                            in self.credential_set.order_by('key')],
        }

    def __unicode__(self):
        output = "%s %s - " % (self.provider, self.created_by.username)
        output += "Credentials {"
        for c in self.credential_set.order_by('key'):
            output += "%s, " % (c.key,)
        output = output[:-2] + "}"
        return output

    class Meta:
        db_table = "identity"
        app_label = "core"
        verbose_name_plural = "identities"
<|endoftext|>


<commit_before>"""
PyChain API library
~~~~~~~~~~~~~~~~~~~

Usage example::

    >>> import chain
    >>> chain.conf.api_key = 'YOUR-KEY'
    >>> chain.conf.api_secret = 'YOUR-SECRET'
    >>> chain.Webhook.create('https://your-server-url.com')

"""
from .requestor import rest
from .exceptions import (
    ChainException, APIConnectionError, APIError, AuthError,
    InvalidRequestError
)

class Webhook(object):
    _api_path_collection = 'webhooks'
    _api_path_single = 'webhooks/{webhook_id}'

    def __init__(self, id_):
        self.id_ = id_
        self.url = None

    def save(self):
        data = {
            'url': self.url
        }
        path = self._api_path_single.format(webhook_id=self.id_)
        return rest('put', path, data=data)

    def delete(self):
        path = self._api_path_single.format(webhook_id=self.id_)
        return rest('delete', path)

    @classmethod
    def construct_from(cls, attributes):
        webhook = cls(id_=attributes.get('id'))
        webhook.url = attributes.get('url')
        return webhook

    @classmethod
    def retrieve(cls, id_):
        path = cls._api_path_single.format(webhook_id=id_)
        json_resp = rest('get', path)
        return cls.construct_from(json_resp)

    @classmethod
    def create(cls, url, id_=None):
        data = {
            'url': url
        }
        if id_ is not None:
            data['id'] = id_
        json_resp = rest('post', cls._api_path_collection, data=data)
        return cls.construct_from(json_resp)

class WebhookEvent(object):
    _api_path_collection = 'webhooks/{webhook_id}/events'

    @classmethod
    def create(cls, webhook_id, event, block_chain, address, confirmations=None):
        data = {
            'webhook_id': webhook_id,
            'event': event,
            'block_chain': block_chain,
            'address': address
        }
        if confirmations is not None:
            data['confirmations'] = confirmations
        path = cls._api_path_collection.format(webhook_id=webhook_id)
        return rest('post', path, data=data)
<commit_msg>Add Webhook __repr__<commit_after>"""
PyChain API library
~~~~~~~~~~~~~~~~~~~

Usage example::

    >>> import chain
    >>> chain.conf.api_key = 'YOUR-KEY'
    >>> chain.conf.api_secret = 'YOUR-SECRET'
    >>> chain.Webhook.create('https://your-server-url.com')

"""
from .requestor import rest
from .exceptions import (
    ChainException, APIConnectionError, APIError, AuthError,
    InvalidRequestError
)

class Webhook(object):
    _api_path_collection = 'webhooks'
    _api_path_single = 'webhooks/{webhook_id}'

    def __init__(self, id_):
        self.id_ = id_
        self.url = None

    def __repr__(self):
        return '<chain.Webhook {}>'.format(self.id_)

    def save(self):
        data = {
            'url': self.url
        }
        path = self._api_path_single.format(webhook_id=self.id_)
        return rest('put', path, data=data)

    def delete(self):
        path = self._api_path_single.format(webhook_id=self.id_)
        return rest('delete', path)

    @classmethod
    def construct_from(cls, attributes):
        webhook = cls(id_=attributes.get('id'))
        webhook.url = attributes.get('url')
        return webhook

    @classmethod
    def retrieve(cls, id_):
        path = cls._api_path_single.format(webhook_id=id_)
        json_resp = rest('get', path)
        return cls.construct_from(json_resp)

    @classmethod
    def create(cls, url, id_=None):
        data = {
            'url': url
        }
        if id_ is not None:
            data['id'] = id_
        json_resp = rest('post', cls._api_path_collection, data=data)
        return cls.construct_from(json_resp)

class WebhookEvent(object):
    _api_path_collection = 'webhooks/{webhook_id}/events'

    @classmethod
    def create(cls, webhook_id, event, block_chain, address, confirmations=None):
        data = {
            'webhook_id': webhook_id,
            'event': event,
            'block_chain': block_chain,
            'address': address
        }
        if confirmations is not None:
            data['confirmations'] = confirmations
        path = cls._api_path_collection.format(webhook_id=webhook_id)
        return rest('post', path, data=data)
<|endoftext|>


<commit_before>from openerp.osv import orm, osv, fields
from datetime import datetime as dt, timedelta as td
from openerp.tools import DEFAULT_SERVER_DATETIME_FORMAT as DTF
from openerp.tools.translate import _
import logging

_logger = logging.getLogger(__name__)

class t4_clinical_apikey(orm.Model):
    _name = 't4.clinical.apikey'

    _columns = {
        'username': fields.char('Username', size=150, required=True),
        'database': fields.char('Database', size=150, required=True),

    }

class t4_clinical_api(orm.AbstractModel):
    _name = 't4.clinical.api'
    _inherit = 't4.clinical.api'

    # # # # # # # #
    #  ACTIVITIES #
    # # # # # # # #

    def getActivities(self, cr, uid, ids, context=None):
        domain = [('id', 'in', ids)] if ids else []
        activity_pool = self.pool['t4.activity']
        activity_ids = activity_pool.search(cr, uid, domain, context=context)
        activity_values = activity_pool.read(cr, uid, activity_ids, [], context=context)
        return activity_values

    def cancel(self, cr, uid, activity_id, context=None):
        activity_pool = self.pool['t4.activity']
        domain = [('id', '=', activity_id)]
        activity_ids = activity_pool.search(cr, uid, domain, context=context)
        if not activity_ids:
            raise osv.except_osv(_('Error!'), 'Activity ID not found: %s' % activity_id)
        return activity_pool.cancel(cr, uid, activity_id, context=context)

    def submit(self, cr, uid, activity_id, data, context=None):
        activity_pool = self.pool['t4.activity']
        domain = [('id', '=', activity_id)]
        activity_ids = activity_pool.search(cr, uid, domain, context=context)
        if not activity_ids:
            raise osv.except_osv(_('Error!'), 'Activity ID not found: %s' % activity_id)
        return activity_pool.submit(cr, uid, activity_id, data, context=context)

    def unassign(self, cr, uid, activity_id, context=None):
        activity_pool = self.pool['t4.activity']
        domain = [('id', '=', activity_id)]
        activity_ids = activity_pool.search(cr, uid, domain, context=context)
        if not activity_ids:
            raise osv.except_osv(_('Error!'), 'Activity ID not found: %s' % activity_id)
        return activity_pool.unassign(cr, uid, activity_id, context=context)

    def assign(self, cr, uid, activity_id, data, context=None):
        activity_pool = self.pool['t4.activity']
        user_pool = self.pool['res.users']
        user_id = uid
        domain = [('id', '=', activity_id)]
        activity_ids = activity_pool.search(cr, uid, domain, context=context)
        if not activity_ids:
            raise osv.except_osv(_('Error!'), 'Activity ID not found: %s' % activity_id)
        if data.get('user_id'):
            user_id = data['user_id']
            domain = [('id', '=', user_id)]
            user_ids = user_pool.search(cr, uid, domain, context=context)
            if not user_ids:
                raise osv.except_osv(_('Error!'), 'User ID not found: %s' % user_id)
        return activity_pool.assign(cr, uid, activity_id, user_id, context=context)

    def complete(self, cr, uid, activity_id, data, context=None):
        activity_pool = self.pool['t4.activity']
        domain = [('id', '=', activity_id)]
        activity_ids = activity_pool.search(cr, uid, domain, context=context)
        if not activity_ids:
            raise osv.except_osv(_('Error!'), 'Activity ID not found: %s' % activity_id)
        return activity_pool.complete(cr, uid, activity_id, data, context=context)

    # # # # # # #
    #  PATIENTS #
    # # # # # # #

    def getActivitiesForPatient(self, cr, uid, patient_id, activity_type, start_date=dt.now()+td(days=-30),
                                end_date=dt.now(), context=None):
        if activity_type == 'ews':
            model_pool = self.pool['t4.clinical.patient.observation.ews']
        ids = model_pool.search(cr, uid, [
            ('patient_id', '=', patient_id),
            ('state', '=', 'completed'),
            ('date_terminated', '>=', start_date.strftime(DTF)),
            ('date_terminated', '<=', end_date.strftime(DTF))], context=context)
        return model_pool.read(cr, uid, ids, [], context=context)<commit_msg>[FIX] getActivitiesforPatient allows any activity type now.<commit_after>from openerp.osv import orm, osv, fields
from datetime import datetime as dt, timedelta as td
from openerp.tools import DEFAULT_SERVER_DATETIME_FORMAT as DTF
from openerp.tools.translate import _
import logging

_logger = logging.getLogger(__name__)

class t4_clinical_apikey(orm.Model):
    _name = 't4.clinical.apikey'

    _columns = {
        'username': fields.char('Username', size=150, required=True),
        'database': fields.char('Database', size=150, required=True),

    }

class t4_clinical_api(orm.AbstractModel):
    _name = 't4.clinical.api'
    _inherit = 't4.clinical.api'

    # # # # # # # #
    #  ACTIVITIES #
    # # # # # # # #

    def getActivities(self, cr, uid, ids, context=None):
        domain = [('id', 'in', ids)] if ids else []
        activity_pool = self.pool['t4.activity']
        activity_ids = activity_pool.search(cr, uid, domain, context=context)
        activity_values = activity_pool.read(cr, uid, activity_ids, [], context=context)
        return activity_values

    def cancel(self, cr, uid, activity_id, context=None):
        activity_pool = self.pool['t4.activity']
        domain = [('id', '=', activity_id)]
        activity_ids = activity_pool.search(cr, uid, domain, context=context)
        if not activity_ids:
            raise osv.except_osv(_('Error!'), 'Activity ID not found: %s' % activity_id)
        return activity_pool.cancel(cr, uid, activity_id, context=context)

    def submit(self, cr, uid, activity_id, data, context=None):
        activity_pool = self.pool['t4.activity']
        domain = [('id', '=', activity_id)]
        activity_ids = activity_pool.search(cr, uid, domain, context=context)
        if not activity_ids:
            raise osv.except_osv(_('Error!'), 'Activity ID not found: %s' % activity_id)
        return activity_pool.submit(cr, uid, activity_id, data, context=context)

    def unassign(self, cr, uid, activity_id, context=None):
        activity_pool = self.pool['t4.activity']
        domain = [('id', '=', activity_id)]
        activity_ids = activity_pool.search(cr, uid, domain, context=context)
        if not activity_ids:
            raise osv.except_osv(_('Error!'), 'Activity ID not found: %s' % activity_id)
        return activity_pool.unassign(cr, uid, activity_id, context=context)

    def assign(self, cr, uid, activity_id, data, context=None):
        activity_pool = self.pool['t4.activity']
        user_pool = self.pool['res.users']
        user_id = uid
        domain = [('id', '=', activity_id)]
        activity_ids = activity_pool.search(cr, uid, domain, context=context)
        if not activity_ids:
            raise osv.except_osv(_('Error!'), 'Activity ID not found: %s' % activity_id)
        if data.get('user_id'):
            user_id = data['user_id']
            domain = [('id', '=', user_id)]
            user_ids = user_pool.search(cr, uid, domain, context=context)
            if not user_ids:
                raise osv.except_osv(_('Error!'), 'User ID not found: %s' % user_id)
        return activity_pool.assign(cr, uid, activity_id, user_id, context=context)

    def complete(self, cr, uid, activity_id, data, context=None):
        activity_pool = self.pool['t4.activity']
        domain = [('id', '=', activity_id)]
        activity_ids = activity_pool.search(cr, uid, domain, context=context)
        if not activity_ids:
            raise osv.except_osv(_('Error!'), 'Activity ID not found: %s' % activity_id)
        return activity_pool.complete(cr, uid, activity_id, data, context=context)

    # # # # # # #
    #  PATIENTS #
    # # # # # # #

    def getActivitiesForPatient(self, cr, uid, patient_id, activity_type, start_date=dt.now()+td(days=-30),
                                end_date=dt.now(), context=None):
        model_pool = self.pool['t4.clinical.patient.observation.'+activity_type]
        ids = model_pool.search(cr, uid, [
            ('patient_id', '=', patient_id),
            ('state', '=', 'completed'),
            ('date_terminated', '>=', start_date.strftime(DTF)),
            ('date_terminated', '<=', end_date.strftime(DTF))], context=context)
        return model_pool.read(cr, uid, ids, [], context=context)<|endoftext|>


<commit_before>import urllib2, json, time, sys
from datetime import date, datetime
from dateutil.rrule import rrule, DAILY
from optparse import OptionParser

parser = OptionParser()
parser.add_option("-f", dest="fahrenheit", action="store", default=False, type="string", help="Convert to FAHRENHEIT")
parser.add_option("-e", dest="end", action="store", default=False, type="string", help="START date")
parser.add_option("-s", dest="start", action="store", default=False, type="string", help="END date")
parser.add_option("-t", dest="token", action="store", default=False, type="string", help="Weather Underground TOKEN")
(options, args) = parser.parse_args()

if options.token:
    token = options.token
else:
    parser.print_help()
    sys.exit()

if options.start:
    start = options.start
else:
    parser.print_help()
    sys.exit()

if options.end:
    end = options.end
else:
    parser.print_help()
    sys.exit()

if options.fahrenheit:
	fahrenheit = True
else:
	fahrenheit = False

start = datetime.strptime(start,'%Y-%m-%d')
end = datetime.strptime(end,'%Y-%m-%d')

url = ""

if end < start:
	print "Error: end date " + str(end) + " occurs before start date " + str(start)
	sys.exit()

# d48122149ff66bca

for dt in rrule(DAILY, dtstart=start, until=end):

	total = 0.0
	temp = 0.0
  	count = 0

	wunderground_url ="http://api.wunderground.com/api/" + token + "/history_" + dt.strftime("%Y%m%d") +"/q/NY/New_York_City.json"

	try:
		url = urllib2.urlopen(wunderground_url)
		parsed_json = json.loads(url.read())
	except:
		print "Error reading URL " + wunderground_url
		print "Is your token correct?"
		url.close()
		sys.exit()

	try:
		for mean in parsed_json['history']['observations']:
			if fahrenheit:
				total += float(mean['tempi'])
			else:
				total += float(mean['tempm'])
  	  		count += 1

		temp = (total / count)
		print dt.strftime("%Y-%m-%d") + "," + str(temp)
	except:
		print "Error retrieving temperature records for start date " + str(start) + " end date " + str(end)
		url.close()

	time.sleep(10)<commit_msg>Update wunderground.py<commit_after>import urllib2, json, time, sys
from datetime import date, datetime
from dateutil.rrule import rrule, DAILY
from optparse import OptionParser

parser = OptionParser()
parser.add_option("-f", dest="fahrenheit", action="store", default=False, type="string", help="Convert to FAHRENHEIT")
parser.add_option("-e", dest="end", action="store", default=False, type="string", help="START date")
parser.add_option("-s", dest="start", action="store", default=False, type="string", help="END date")
parser.add_option("-t", dest="token", action="store", default=False, type="string", help="Weather Underground TOKEN")
(options, args) = parser.parse_args()

if options.token:
    token = options.token
else:
    parser.print_help()
    sys.exit()

if options.start:
    start = options.start
else:
    parser.print_help()
    sys.exit()

if options.end:
    end = options.end
else:
    parser.print_help()
    sys.exit()

if options.fahrenheit:
	fahrenheit = True
else:
	fahrenheit = False

start = datetime.strptime(start,'%Y-%m-%d')
end = datetime.strptime(end,'%Y-%m-%d')

url = ""

if end < start:
	print "Error: end date " + str(end) + " occurs before start date " + str(start)
	sys.exit()

for dt in rrule(DAILY, dtstart=start, until=end):

	total = 0.0
	temp = 0.0
  	count = 0

	wunderground_url ="http://api.wunderground.com/api/" + token + "/history_" + dt.strftime("%Y%m%d") +"/q/NY/New_York_City.json"

	try:
		url = urllib2.urlopen(wunderground_url)
		parsed_json = json.loads(url.read())
	except:
		print "Error reading URL " + wunderground_url
		print "Is your token correct?"
		url.close()
		sys.exit()

	try:
		for mean in parsed_json['history']['observations']:
			if fahrenheit:
				total += float(mean['tempi'])
			else:
				total += float(mean['tempm'])
  	  		count += 1

		temp = (total / count)
		print dt.strftime("%Y-%m-%d") + "," + str(temp)
	except:
		print "Error retrieving temperature records for start date " + str(start) + " end date " + str(end)
		url.close()

	time.sleep(10)
<|endoftext|>