add survey_publisher_2022 with budget / funding sources

compendium_v2/publishers/survey_publisher_2022.py

+import logging
+import click
+import enum
+import math
+
+from compendium_v2.environment import setup_logging
+from compendium_v2.config import load
+from compendium_v2 import db, survey_db
+from compendium_v2.db import model
+
+setup_logging()
+
+logger = logging.getLogger('survey-publisher-2022')
+
+BUDGET_QUERY = """
+SELECT DISTINCT ON (n.id, a.question_id)
+    n.abbreviation AS nren,
+    a.value AS budget
+FROM answers a
+JOIN nrens n ON a.nren_id = n.id
+JOIN questions q ON a.question_id = q.id
+JOIN sections s ON q.section_id = s.id
+JOIN compendia c ON s.compendium_id = c.id
+WHERE
+    a.question_id = 16402
+AND c.year = 2022
+ORDER BY n.id, a.question_id DESC
+"""
+
+FUNDING_SOURCES_TEMPLATE_QUERY = """
+SELECT DISTINCT ON (n.id, a.question_id)
+    n.abbreviation AS nren,
+    a.value AS value
+FROM answers a
+JOIN nrens n ON a.nren_id = n.id
+JOIN questions q ON a.question_id = q.id
+JOIN sections s ON q.section_id = s.id
+JOIN compendia c ON s.compendium_id = c.id
+WHERE
+    a.question_id = {}
+AND c.year = 2022
+ORDER BY n.id, a.question_id DESC
+"""
+
+
+class FundingSource(enum.Enum):
+    CLIENT_INSTITUTIONS = 16405
+    EUROPEAN_FUNDING = 16406
+    COMMERCIAL = 16407
+    OTHER = 16408
+    GOV_PUBLIC_BODIES = 16409
+
+
+def setup_db(config):
+    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)
+
+
+def transfer_budget():
+    with db.session_scope() as session, survey_db.session_scope() as survey:
+        rows = survey.execute(BUDGET_QUERY)
+        for row in rows:
+            nren = row['nren']
+            _budget = row['budget']
+            try:
+                budget = float(_budget.replace('"', '').replace(',', ''))
+            except ValueError:
+                logger.info(
+                    f'{nren} has no budget for 2022. Skipping. ({_budget}))')
+                continue
+
+            if budget > 200:
+                logger.info(
+                    f'{nren} has budget set to >200M EUR for 2022. ({budget})')
+
+            budget_entry = model.BudgetEntry(
+                nren=nren,
+                budget=budget,
+                year=2022,
+            )
+            session.merge(budget_entry)
+        session.commit()
+
+
+def transfer_funding_sources():
+    with db.session_scope() as session, survey_db.session_scope() as survey:
+        sources = {source.value: dict() for source in FundingSource}
+        nrens = set()
+
+        for source in FundingSource:
+            data = survey.execute(
+                FUNDING_SOURCES_TEMPLATE_QUERY.format(source.value))
+            for row in data:
+                nren = row['nren']
+                nrens.add(nren)
+                _value = row['value']
+                try:
+                    value = float(_value.replace('"', '').replace(',', ''))
+                except ValueError:
+                    name = source.name
+                    logger.info(
+                        f'{nren} has invalid value for {name}.'
+                        + f' ({_value}))')
+                    value = 0
+
+                sources[source.value][nren] = value
+
+        client_institutions = sources[FundingSource.CLIENT_INSTITUTIONS.value]
+        european_funding = sources[FundingSource.EUROPEAN_FUNDING.value]
+        gov_public_bodies = sources[FundingSource.GOV_PUBLIC_BODIES.value]
+        commercial = sources[FundingSource.COMMERCIAL.value]
+        other = sources[FundingSource.OTHER.value]
+
+        _data = [client_institutions, european_funding,
+                 gov_public_bodies, commercial, other]
+
+        for nren in nrens:
+
+            def _get_nren(source, nren):
+                try:
+                    return source[nren]
+                except KeyError:
+                    return 0
+            total = sum([_get_nren(source, nren) for source in _data])
+
+            if not math.isclose(total, 100, abs_tol=0.01):
+                logger.info(
+                    f'{nren} funding sources do not sum to 100%. ({total})')
+
+            funding_source = model.FundingSource(
+                nren=nren,
+                year=2022,
+                client_institutions=client_institutions.get(nren, 0),
+                european_funding=european_funding.get(nren, 0),
+                gov_public_bodies=gov_public_bodies.get(nren, 0),
+                commercial=commercial.get(nren, 0),
+                other=other.get(nren, 0),
+            )
+            session.merge(funding_source)
+        session.commit()
+
+
+@click.command()
+@click.option('--config', type=click.STRING, default='config.json')
+def cli(config):
+    app_config = load(open(config, 'r'))
+    setup_db(app_config)
+    transfer_budget()
+    transfer_funding_sources()
+
+
+if __name__ == "__main__":
+    cli()