notifications-api/app/commands.py

import uuid
from datetime import datetime, timedelta
from decimal import Decimal
from flask_script import Command, Option

from app import db
from app.dao.monthly_billing_dao import (
    create_or_update_monthly_billing,
    get_monthly_billing_by_notification_type,
    get_service_ids_that_need_billing_populated
)
from app.models import PROVIDERS, User, SMS_TYPE, EMAIL_TYPE
from app.dao.services_dao import (
    delete_service_and_all_associated_db_objects,
    dao_fetch_all_services_by_user
)
from app.dao.provider_rates_dao import create_provider_rates
from app.dao.users_dao import (delete_model_user, delete_user_verify_codes)


class CreateProviderRateCommand(Command):

    option_list = (
        Option('-p', '--provider_name', dest="provider_name", help='Provider name'),
        Option('-c', '--cost', dest="cost", help='Cost (pence) per message including decimals'),
        Option('-d', '--valid_from', dest="valid_from", help="Date (%Y-%m-%dT%H:%M:%S) valid from")
    )

    def run(self, provider_name, cost, valid_from):
        if provider_name not in PROVIDERS:
            raise Exception("Invalid provider name, must be one of ({})".format(', '.join(PROVIDERS)))

        try:
            cost = Decimal(cost)
        except:
            raise Exception("Invalid cost value.")

        try:
            valid_from = datetime.strptime('%Y-%m-%dT%H:%M:%S', valid_from)
        except:
            raise Exception("Invalid valid_from date. Use the format %Y-%m-%dT%H:%M:%S")

        create_provider_rates(provider_name, valid_from, cost)


class PurgeFunctionalTestDataCommand(Command):

    option_list = (
        Option('-u', '-user-email-prefix', dest='user_email_prefix', help="Functional test user email prefix."),
    )

    def run(self, user_email_prefix=None):
        if user_email_prefix:
            users = User.query.filter(User.email_address.like("{}%".format(user_email_prefix))).all()
            for usr in users:
                # Make sure the full email includes a uuid in it
                # Just in case someone decides to use a similar email address.
                try:
                    uuid.UUID(usr.email_address.split("@")[0].split('+')[1])
                except ValueError:
                    print("Skipping {} as the user email doesn't contain a UUID.".format(usr.email_address))
                else:
                    services = dao_fetch_all_services_by_user(usr.id)
                    if services:
                        for service in services:
                            delete_service_and_all_associated_db_objects(service)
                    else:
                        delete_user_verify_codes(usr)
                        delete_model_user(usr)


class CustomDbScript(Command):

    option_list = (
        Option('-n', '-name-of-db-function', dest='name_of_db_function', help="Function name of the DB script to run"),
    )

    def run(self, name_of_db_function):
        db_function = getattr(self, name_of_db_function, None)
        if callable(db_function):
            db_function()
        else:
            print('The specified function does not exist.')

    def backfill_notification_statuses(self):
        """
        This will be used to populate the new `Notification._status_fkey` with the old
        `Notification._status_enum`
        """
        LIMIT = 250000
        subq = "SELECT id FROM notification_history WHERE notification_status is NULL LIMIT {}".format(LIMIT)
        update = "UPDATE notification_history SET notification_status = status WHERE id in ({})".format(subq)
        result = db.session.execute(subq).fetchall()

        while len(result) > 0:
            db.session.execute(update)
            print('commit {} updates at {}'.format(LIMIT, datetime.utcnow()))
            db.session.commit()
            result = db.session.execute(subq).fetchall()

    def update_notification_international_flag(self):
        # 250,000 rows takes 30 seconds to update.
        subq = "select id from notifications where international is null limit 250000"
        update = "update notifications set international = False where id in ({})".format(subq)
        result = db.session.execute(subq).fetchall()

        while len(result) > 0:
            db.session.execute(update)
            print('commit 250000 updates at {}'.format(datetime.utcnow()))
            db.session.commit()
            result = db.session.execute(subq).fetchall()

        # Now update notification_history
        subq_history = "select id from notification_history where international is null limit 250000"
        update_history = "update notification_history set international = False where id in ({})".format(subq_history)
        result_history = db.session.execute(subq_history).fetchall()
        while len(result_history) > 0:
            db.session.execute(update_history)
            print('commit 250000 updates at {}'.format(datetime.utcnow()))
            db.session.commit()
            result_history = db.session.execute(subq_history).fetchall()

    def fix_notification_statuses_not_in_sync(self):
        """
        This will be used to correct an issue where Notification._status_enum and NotificationHistory._status_fkey
        became out of sync. See 979e90a.

        Notification._status_enum is the source of truth so NotificationHistory._status_fkey will be updated with
        these values.
        """
        MAX = 10000

        subq = "SELECT id FROM notifications WHERE cast (status as text) != notification_status LIMIT {}".format(MAX)
        update = "UPDATE notifications SET notification_status = status WHERE id in ({})".format(subq)
        result = db.session.execute(subq).fetchall()

        while len(result) > 0:
            db.session.execute(update)
            print('Committed {} updates at {}'.format(len(result), datetime.utcnow()))
            db.session.commit()
            result = db.session.execute(subq).fetchall()

        subq_hist = "SELECT id FROM notification_history WHERE cast (status as text) != notification_status LIMIT {}" \
            .format(MAX)
        update = "UPDATE notification_history SET notification_status = status WHERE id in ({})".format(subq_hist)
        result = db.session.execute(subq_hist).fetchall()

        while len(result) > 0:
            db.session.execute(update)
            print('Committed {} updates at {}'.format(len(result), datetime.utcnow()))
            db.session.commit()
            result = db.session.execute(subq_hist).fetchall()

    def link_inbound_numbers_to_service(self):
        update = """
        UPDATE inbound_numbers SET
        service_id = services.id,
        updated_at = now()
        FROM services
        WHERE services.sms_sender = inbound_numbers.number AND
        inbound_numbers.service_id is null
        """
        result = db.session.execute(update)
        db.session.commit()

        print("Linked {} inbound numbers to service".format(result.rowcount))


class PopulateMonthlyBilling(Command):
    option_list = (
        Option('-y', '-year', dest="year", help="Use for integer value for year, e.g. 2017"),
    )

    def run(self, year):
        service_ids = get_service_ids_that_need_billing_populated(
            start_date=datetime(2016, 5, 1), end_date=datetime(2017, 8, 16)
        )
        start, end = 1, 13
        if year == '2016':
            start = 4

        for service_id in service_ids:
            print('Starting to populate data for service {}'.format(str(service_id)))
            print('Starting populating monthly billing for {}'.format(year))
            for i in range(start, end):
                print('Population for {}-{}'.format(i, year))
                self.populate(service_id, year, i)

    def populate(self, service_id, year, month):
        create_or_update_monthly_billing(service_id, datetime(int(year), int(month), 1))
        sms_res = get_monthly_billing_by_notification_type(
            service_id, datetime(int(year), int(month), 1), SMS_TYPE
        )
        email_res = get_monthly_billing_by_notification_type(
            service_id, datetime(int(year), int(month), 1), EMAIL_TYPE
        )
        print("Finished populating data for {} for service id {}".format(month, str(service_id)))
        print('SMS: {}'.format(sms_res.monthly_totals))
        print('Email: {}'.format(email_res.monthly_totals))