budget.py

import csv
import logging
from collections import defaultdict
from typing import Any

from flask import Blueprint, jsonify, current_app
from sqlalchemy import values

from compendium_v2 import db, survey_db
from compendium_v2.background_task import xlsx_to_csv_sheet_parsing_task
from compendium_v2.db import model
from compendium_v2.survey_db import model as survey_model
from compendium_v2.routes import common

routes = Blueprint('budget', __name__)


@routes.before_request
def before_request():
    config = current_app.config['CONFIG_PARAMS']
    dsn_prn = config['SQLALCHEMY_DATABASE_URI']
    db.init_db_model(dsn_prn)
    dsn_survey = config['SURVEY_DATABASE_URI']
    survey_db.init_db_model(dsn_survey)


logger = logging.getLogger(__name__)

col_pal = ['#fd7f6f', '#7eb0d5', '#b2e061',
           '#bd7ebe', '#ffb55a', '#ffee65',
           '#beb9db', '#fdcce5', '#8bd3c7']

BUDGET_RESPONSE_SCHEMA = {
    '$schema': 'http://json-schema.org/draft-07/schema#',

    'definitions': {
        'budget': {
            'type': 'object',
            'properties': {
                'id': {'type': 'number'},
                'NREN': {'type': 'string'},
                'BUDGET': {'type': 'string'},
                'BUDGET_YEAR': {'type': 'integer'},
            },
            'required': ['id'],
            'additionalProperties': False
        }
    },

    'type': 'array',
    'items': {'$ref': '#/definitions/budget'}
}


@routes.route('/', methods=['GET'])
@common.require_accepts_json
def budget_view() -> Any:
    """
    handler for /api/budget/ requests

    response will be formatted as:

    .. asjson::
        compendium_v2.routes.data_entry.BUDGET_RESPONSE_SCHEMA

    :return:
    """

    def _extract_data(entry: model.BudgetEntry):
        return {
            'id': entry.id,
            'NREN': entry.nren,
            'BUDGET': entry.budget,
            'BUDGET_YEAR': entry.year,
        }

    with db.session_scope() as session:
        entries = sorted([_extract_data(entry)
                          for entry in session.query(model.BudgetEntry)],
                         key=lambda d: (d['BUDGET_YEAR'], d['NREN']))
    return jsonify(entries)


@routes.route('/migration', methods=['GET'])
@common.require_accepts_json
def db_budget_migration():
    with survey_db.session_scope() as survey_session, \
            db.session_scope() as session:
        _entries = session.query(model.BudgetEntry)

        inserted = defaultdict(dict)

        for entry in _entries:
            inserted[entry.nren][entry.year] = entry.budget

        data = survey_session.query(survey_model.Nrens)
        for nren in data:
            for budget in nren.budgets:
                abbrev = nren.abbreviation
                year = budget.year

                if inserted.get(abbrev, {}).get(year):
                    continue
                else:
                    inserted[abbrev][year] = True
                entry = model.BudgetEntry(
                    nren=abbrev, budget=budget.budget, year=year)
                session.add(entry)
        # Import the data to database
        xlsx_to_csv_sheet_parsing_task.parse_budget_xlsx_file()
        with open('compendium_v2/background_task/csv/BudgetCsvFile.csv',
                  newline='') as csvfile:
            reader = csv.reader(csvfile)

            for row in reader:
                if row is not None:
                    entry = session.query(
                        model.BudgetEntry).filter_by(nren=row[0],
                                                     year=row[
                                                         2])
                    dup_entry: model.BudgetEntry = entry.first()
                    if dup_entry:
                        dup_entry.budget = row[1]
                        entry.update({"budget": row[1]})

                    else:
                        print("add new")
                        print(row)
                        budget_entry = model.BudgetEntry(
                            nren=row[0], budget=row[1], year=row[2])
                        session.add(budget_entry)
        session.commit()

    return "Success"