def main():
    sess = GlobalDB.db().session

    # Fill in long_to_short and short_to_long dicts
    long_to_short_dict = {
        FILE_TYPE_DICT['appropriations']: {
            'budgetauthorityavailableamounttotal_cpe':
            'total_budgetary_resources_cpe',
            'budget_authority_available_cpe': 'total_budgetary_resources_cpe'
        }
    }
    for col in sess.query(FileColumn.name, FileColumn.name_short,
                          FileColumn.file_id).all():
        if not long_to_short_dict.get(col.file_id):
            long_to_short_dict[col.file_id] = {}
        long_to_short_dict[col.file_id][col.name] = col.name_short

    # Loop through the file types by name
    file_type_names = ['appropriations', 'program_activity', 'award_financial']
    for file_type_name in file_type_names:
        file_type_id = FILE_TYPE_DICT[file_type_name]
        logger.info('Starting to update the {} table'.format(
            FTI_TABLENAME_DICT[file_type_id]))

        # Load all published files with file_type_id matching the file_type_id
        file_columns = sess.query(FileColumn).filter(
            FileColumn.file_id == file_type_id).all()
        csv_schema = {f.name_short: f for f in file_columns}
        query = filenames_by_file_type(file_type_id)

        # Loop through the files and load each one individually
        for query_resp in query.all():
            filename = query_resp[0]
            submission_id = query_resp[1]

            # Only add rows to the DB if this submission has not yet been loaded into the certified table
            to_skip = sess.query(FTI_TABLE_DICT[file_type_id]).filter_by(
                submission_id=submission_id)
            if to_skip.first() is None:
                submission = sess.query(Submission).filter_by(
                    submission_id=submission_id).one()

                if submission.publish_status_id == PUBLISH_STATUS_DICT[
                        'published']:
                    # Copy data directly from the table
                    insert_from_table(file_type_id, submission_id)

                else:
                    # Insert data by file
                    insert_file(filename, submission_id, file_type_id,
                                csv_schema, long_to_short_dict[file_type_id])

                    # Populate tas and account_num
                    update_account_nums(FTI_TABLE_DICT[file_type_id],
                                        submission_id)
                    sess.commit()
            else:
                logger.info(
                    'Skipping file "{}"; Submission {} already loaded'.format(
                        filename, submission_id))
def test_update_account_nums_no_match(database, factory):
    """If a TAS doesn't share fields, we don't expect a match"""
    sess = database.session
    submission = SubmissionFactory(reporting_start_date=date(2010, 10, 10),
                                   reporting_end_date=date(2010, 10, 31))
    sess.add(submission)
    sess.flush()
    tas = TASFactory(internal_start_date=date(2010, 9, 1))
    # note these will have different fields
    model = factory(submission_id=submission.submission_id)
    assert model.account_num is None
    sess.add_all([tas, model])
    sess.commit()

    validationManager.update_account_nums(model.__class__,
                                          submission.submission_id)

    model = sess.query(model.__class__).one()  # we'll only have one entry
    assert model.account_num is None
def test_update_account_nums_has_match_open_ended(database, factory):
    """ If there are models which match the TAS (with an undefined end date), they should be modified """
    sess = database.session
    submission = SubmissionFactory(reporting_start_date=date(2010, 10, 1),
                                   reporting_end_date=date(2010, 10, 1))
    sess.add(submission)
    sess.flush()
    tas = TASFactory(internal_start_date=date(2010, 9, 1))
    model = factory(submission_id=submission.submission_id,
                    **tas.component_dict())
    assert model.account_num is None
    sess.add_all([tas, model])
    sess.commit()

    validationManager.update_account_nums(model.__class__,
                                          submission.submission_id)

    model = sess.query(model.__class__).one()  # we'll only have one entry
    assert model.account_num == tas.account_num
def test_update_account_nums_bad_dates(database, factory):
    """If the relevant TAS does not overlap the date of the submission, it
    should not be used"""
    sess = database.session
    submission = SubmissionFactory(reporting_start_date=date(2010, 10, 1),
                                   reporting_end_date=date(2010, 10, 1))
    sess.add(submission)
    sess.flush()
    tas = TASFactory(internal_start_date=date(2011, 1, 1))
    model = factory(submission_id=submission.submission_id,
                    **tas.component_dict())
    assert model.account_num is None
    sess.add_all([tas, model])
    sess.commit()

    validationManager.update_account_nums(model.__class__,
                                          submission.submission_id)

    model = sess.query(model.__class__).one()  # we'll only have one entry
    assert model.account_num is None