def getModelData(model):
    """Return records for the ModelFile or Dat tables.
    
    Args:
        model(str): either 'DAT' or 'MODEL'. If anything else will return 'MODEL'.
    
    Return:
        tuple(cols:header strings, rows: list of record data).
    """
    pm.connectDB()
    cols = []
    rows = []
    try:
        if model == 'DAT':
            mquery = pm.Dat.select()
            cols = ['timestamp', 'name', 'amendments', 'comments']
            for r in mquery:
                rows.append([r.timestamp.strftime("%Y-%m-%d %H:%M:%S"), r.name, r.amendments, r.comments])
        
        elif model == 'IED':
            mquery = pm.Ied.select()
            cols = ['timestamp', 'name', 'ref', 'amendments', 'comments']
            for r in mquery:
                rows.append([r.timestamp.strftime("%Y-%m-%d %H:%M:%S"), r.name, r.ref, r.amendments, r.comments])
        
        else:
            mquery = pm.ModelFile.select().where(pm.ModelFile.model_type == model)
            cols = ['timestamp', 'name', 'comments']
            for r in mquery:
                rows.append([r.timestamp.strftime("%Y-%m-%d %H:%M:%S"), r.name, r.comments])
    finally:
        pm.disconnectDB()
    
    return cols, rows
def getRunData():
    """Return records for the Run table.
    
    Return:
        tuple(cols:header strings, rows: list of record data).
    """
    pm.connectDB()
    cols = [
        'id', 'timestamp', 'run_hash', 'run_options', 'event_name', 
        'setup', 'comments', 'ief', 'tcf', 'initial_conditions', 'isis_results', 
        'tuflow_results', 'estry_results', 'event_duration', 'modeller', 
        'isis_version', 'tuflow_version', 'ief_dir', 'tcf_dir', 'log_dir', 
        'run_status', 'mb'
    ]
    rows = []
    try:
        rquery = pm.Run.select()
        rows = []
        for r in rquery:
            rows.append(
                [
                 r.id, r.timestamp.strftime("%Y-%m-%d %H:%M:%S"), r.run_hash,
                 r.run_options, r.event_name, r.setup, r.comments, r.ief, r.tcf, 
                 r.initial_conditions, r.isis_results, r.tuflow_results, r.estry_results, 
                 r.event_duration, r.modeller, 
                 r.isis_version, r.tuflow_version, r.ief_dir,
                 r.tcf_dir, r.log_dir, r.run_status, r.mb
                ]
            )
    finally:
        pm.disconnectDB()
    
    return cols, rows
def updateModelRow(updateDict, model_name):
    """Update values in the ModelFile table.
    
    Args:
        updateDict(dict): containing {fieldname: value} pairs for updating.
        model_name(str): the ModelFile_name value to query.
    """
    pm.connectDB()
    try:
        query = pm.ModelFile.update(**updateDict).where(pm.ModelFile.name == model_name)
        query.execute()
    finally:
        pm.disconnectDB()
def updateRunRow(updateDict, run_id):
    """Update values in the Run table.
    
    Args:
        updateDict(dict): containing {fieldname: value} pairs for updating.
        run_id(int): the Run.id value to query.
    """
    pm.connectDB()
    try:
        query = pm.Run.update(**updateDict).where(pm.Run.id == run_id)
        query.execute()
    finally:
        pm.disconnectDB()
def updateIedRow(updateDict, ied_name):
    """Update values in the Ied table.
    
    Args:
        updateDict(dict): containing {fieldname: value} pairs for updating.
        ied_name(str): the Ied.name value to query.
    """
    pm.connectDB()
    try:
        query = pm.Ied.update(**updateDict).where(pm.Ied.name == ied_name)
        query.execute()
    finally:
        pm.disconnectDB()
def deleteIedRow(ied_name, remove_orphans=True, connect_db=True):
    """Delete a record in the Ied table.
    
    Deletes the specified Ied record.
    
    Args:
        ied_name(str): the Ied.name to query against.
    """
    if connect_db:
        pm.connectDB()
    try:
        i = pm.Ied.get(pm.Ied.name == ied_name)
        i.delete_instance()
    finally:
        if connect_db:
            pm.disconnectDB()
def deleteDatRow(dat_name, connect_db=True):
    """Delete a record in the Dat table.
    
    Deletes the specified Dat record.
    
    Args:
        dat_name(str): the Dat.name to query against.
    """
    if connect_db:
        pm.connectDB()
    try:
        d = pm.Dat.get(pm.Dat.name == dat_name)
        d.delete_instance()
    finally:
        if connect_db:
            pm.disconnectDB()
def getRunRow(run_id):
    """Return the Run record associated with the given Run.id.
    
    Args:
        run_id(int): the Run.id to query against.
    
    Return:
        dict - containing {field: value} for the record.
    """    
    pm.connectDB()
    try:
        query = pm.Run.get(pm.Run.id == run_id)
        run = shortcuts.model_to_dict(query, recurse=False)
    finally:
        pm.disconnectDB()
    return run
def addDat(dat):
    """Add a new record to the Dat table.
    
    Args:
        dat(dict): containing the records to update.
    
    Return:
        peewee.Model - the newly created Dat record.
    """
    pm.connectDB()
    try:
        d, created = pm.Dat.get_or_create(
            name=dat['NAME'], amendments=dat['AMENDMENTS'], 
            comments=dat['COMMENTS']
        )
    finally:
        pm.disconnectDB()
    return d
def addRun(run, run_hash, ief_dir, tcf_dir, dat=None):
    """Add a new record to the Run table.
    
    Args:
        run(dict): containing the values to update the Run table with.
        run_hash(str): unique hash code for this run.
        ief_dir(str): the ief directory string - can be ''.
        tcf_dir(str): the tcf directory string - can be ''.
        dat=None(peewee.Model): the Dat record to reference against the 
            Run.dat foreign key.
    
    Return:
        peewee.Model - the newly created Run record.
    """
    pm.connectDB()
    
    try:
#         if dat is not None:
        r = pm.Run(dat=dat, run_hash=run_hash, setup=run['SETUP'], modeller=run['MODELLER'],
                ief=run['IEF'], tcf=run['TCF'], initial_conditions=run['INITIAL_CONDITIONS'], 
                isis_results=run['ISIS_RESULTS'], tuflow_results=run['TUFLOW_RESULTS'],
                estry_results=run['ESTRY_RESULTS'], event_duration=run['EVENT_DURATION'],
                comments=run['COMMENTS'], isis_version=run['ISIS_BUILD'], tuflow_version=run['TUFLOW_BUILD'],
                event_name=run['EVENT_NAME'], ief_dir=ief_dir, tcf_dir=tcf_dir,
                log_dir=run['LOG_DIR'], run_options=run['RUN_OPTIONS'], 
                run_status=run['RUN_STATUS'], mb=run['MB'])
#         else:
#             r = pm.Run(run_hash=run_hash, setup=run['SETUP'], modeller=run['MODELLER'],
#                     ief=run['IEF'], tcf=run['TCF'], initial_conditions=run['INITIAL_CONDITIONS'], 
#                     isis_results=run['ISIS_RESULTS'], tuflow_results=run['TUFLOW_RESULTS'],
#                     estry_results=run['ESTRY_RESULTS'], event_duration=run['EVENT_DURATION'],
#                     comments=run['COMMENTS'], isis_version=run['ISIS_BUILD'], tuflow_version=run['TUFLOW_BUILD'],
#                     event_name=run['EVENT_NAME'], ief_dir=ief_dir, tcf_dir=tcf_dir,
#                     log_dir=run['LOG_DIR'], run_options=run['RUN_OPTIONS'],
#                     run_status=run['RUN_STATUS'], mb=run['MB'])
        r.save()
    finally:
        pm.disconnectDB()
    return r
def deleteModelRow(model_name, remove_orphans=True, connect_db=True):
    """Delete a record in the ModelFile table.
    
    Deletes the specified ModelFile record and any foreign key associations.
    
    Args:
        model_name(str): the ModelFile.name to query against.
        remove_orphans=True(bool): if True will call deleteOrphanFiles after.
    """
    if connect_db:
        pm.connectDB()
    try:
        m = pm.ModelFile.get(pm.ModelFile.name == model_name)
        m.delete_instance(recursive=True)
        
        # Delete any orphaned subfiles
        if remove_orphans:
            deleteOrphanFiles(execute=True)
        
    finally:
        if connect_db:
            pm.disconnectDB()
def deleteOrphanFiles(run_id=-1, execute=True, connect_db=True):
    """Find any orphaned file references and delete them from the database.
    
    Checks the SubFile table for and SubFile.name values that aren't in the
    ModelFile_SubFile table and deletes them.
    
    Args:
        execute=True(bool): I don't think this does anything? Should probably
            always leave as default until further notice. DEBUG.
    """    
    if connect_db:
        pm.connectDB()
    try:
        # Delete any orphaned subfiles
        subs = pm.ModelFile_SubFile.select(pm.ModelFile_SubFile.sub_file_id)
        q = pm.SubFile.delete().where(~(pm.SubFile.name << subs))
        if execute:
            q.execute()
        # Need to clean up the subfile entries as well
        subs = pm.Run_SubFile.select(pm.Run_SubFile.sub_file_id)
        q = pm.SubFile.delete().where(~(pm.SubFile.name << subs))
        if execute:
            q.execute()
        
        if not run_id == -1:
            # Delete all of the run references subfile (Run_SubFile)
            rsquery = pm.Run_SubFile.delete().where(pm.Run_SubFile.run_id == run_id)
            if execute:
                rsquery.execute()
        
            # Delete all of the run referencd ied files (Run_Ied)
            iquery = pm.Run_Ied.delete().where(pm.Run_Ied.run_id == run_id)
            if execute:
                iquery.execute()
    
    finally:
        if connect_db:
            pm.disconnectDB()
def getSimpleQuery(table, value1, with_files, new_sub_only, new_model_only, run_id, value2=''):
    """Get the results of a query from the database.
    
    This is a bit of a beast of a function, but it seems hard to break it down 
    until I find a better/neater way of running the queryies.
    
    Basically returns header columns, row data list, tuple for using. It is a 
    bit of a black box at the moment. It is not possible to state which fields
    will be used, only the values to test them against.
    
    Args:
        table (str): table name in the db.
        value1(str): field value to check.
        with_files(bool): Used with queries on the ModelFile table. States whether
            to return associated SubFiles or not.
        new_sub_only(bool): If True and with_files==True it will return only
            the associated SubFiles that have new_file == True.
        new_model_only(bool): Same as new_sub_only, but for ModelFile.
        run_id(int): the Run.id value to compare. If -1 it will not be used.
        value2=''(str): optional second field value to check.
    """
    
    pm.connectDB()
    cols = []
    rows = []
    try:
        if table == 'DAT':
            
            cols = ['Date', 'Name', 'Amendments', 'Comments']
            rows = []
            
            # If run_id given - select all with that id
            if run_id != -1:
                query = (pm.Run
                            .select(pm.Run, pm.Dat)
                            .join(pm.Dat)
                            .where(pm.Run.id == run_id)
                         )
                cols = ['Run ID', 'Date', 'Name', 'Amendments', 'Comments']
                for r in query:
                    rows.append([r.id, r.timestamp.strftime("%Y-%m-%d %H:%M:%S"), r.dat.name, r.dat.amendments, r.dat.comments])
            else:
                query = pm.Dat.select()
                query = checkWildcard(query, pm.Dat.name, value1)
            
                cols = ['Date', 'Name', 'Amendments', 'Comments']
                for r in query:
                    rows.append([r.timestamp.strftime("%Y-%m-%d %H:%M:%S"), r.name, r.amendments, r.comments])
        
        elif table == 'IED':
            rows = []
            
            # If using a run_id we need to join the Run_ModelFile table too
            if run_id != -1:
                query = (pm.Run_Ied
                        .select(pm.Run_Ied, pm.Ied)
                        .join(pm.Ied)
                        .switch(pm.Run_Ied)
                        .where(pm.Run_Ied.run_id == run_id)
                        )
                cols = ['Run ID', 'Date', 'Name', 'Ref', 'Amendments', 'Comments']
                for r in query:
                    rows.append([r.run_id, r.timestamp.strftime("%Y-%m-%d %H:%M:%S"), r.ied.name, r.ied.ref, r.ied.amendments, r.ied.comments])
            else:
                query = pm.Ied.select()
                cols = ['Date', 'Name', 'Ref', 'Amendments', 'Comments']
                for r in query:
                    rows.append([r.timestamp.strftime("%Y-%m-%d %H:%M:%S"), r.name, r.ref, r.amendments, r.comments])
        
        elif table == 'RUN Options' or table == 'RUN Event':

            query = pm.Run.select()
            if table == 'RUN Event':
                query = checkWildcard(query, pm.Run.event_name, value1)
            else:
                query = checkWildcard(query, pm.Run.run_options, value1)
       
            cols = ['Run ID', 'Date', 'Event Name', 'Run Options', 'Comments', 'Status', 'MB']
            rows = []
            for r in query:
                rows.append([str(r.id), r.timestamp.strftime("%Y-%m-%d %H:%M:%S"), r.event_name, r.run_options, r.comments, r.run_status, r.mb])

        
        else:
            # Returning associated SubFile's as well so extra queries for the
            # value2 param will be needed
            if with_files:
               
                # If using a run_id we need to join the Run_ModelFile and Run_Subfile 
                # tables too
                if run_id != -1:
                    query = (pm.ModelFile_SubFile
                            .select(pm.ModelFile_SubFile, pm.ModelFile, pm.SubFile, pm.Run_ModelFile, pm.Run_SubFile)
                            .join(pm.SubFile)
                            .switch(pm.ModelFile_SubFile)
                            .join(pm.ModelFile)
                            .join(pm.Run_ModelFile)
                            .switch(pm.SubFile)
                            .join(pm.Run_SubFile)
                            )
                else:
                    query = (pm.ModelFile_SubFile
                            .select(pm.ModelFile_SubFile, pm.ModelFile, pm.SubFile)
                            .join(pm.SubFile)
                            .switch(pm.ModelFile_SubFile)
                            .join(pm.ModelFile)
                            )
                
                if not table == 'All Modelfiles':
                        query = query.where(pm.ModelFile.model_type == table)
                
                query = checkWildcard(query, pm.ModelFile.name, value1)
                query = checkWildcard(query, pm.SubFile.name, value2)
                
                if new_sub_only:
                    query = query.where(pm.ModelFile_SubFile.new_file == True)
                

                rows = []
                if run_id != -1:

                    if new_model_only:
                        query = query.where(pm.Run_ModelFile.new_file == True)
                    
                    # Filter model files by run id
                    query = query.where(pm.Run_ModelFile.run_id == run_id)
                    # Filter subfiles by run id
                    query = query .where(pm.Run_SubFile.run_id == run_id)
                    cols = ['Run ID', 'Modelfile Timestamp', 'Model Type', 'Modelfile', 'Modelfile New', 'Comments', 'Subfile', 'Subfile Timestamp', 'Subfile New']
                    for r in query:
                        rows.append([r.model_file.run_modelfile.run_id, r.timestamp.strftime("%Y-%m-%d %H:%M:%S"), r.model_file.model_type, r.model_file.name, r.model_file.run_modelfile.new_file, r.model_file.comments, r.sub_file.name, r.sub_file.timestamp.strftime("%Y-%m-%d %H:%M:%S"), r.new_file])
                
                else:
                    cols = ['Modelfile Timestamp', 'Model Type', 'Modelfile', 'Comments', 'Subfile', 'Subfile Timestamp', 'Subfile New']
                    for r in query:
                        rows.append([r.timestamp.strftime("%Y-%m-%d %H:%M:%S"), r.model_file.model_type, r.model_file.name, r.model_file.comments, r.sub_file.name, r.sub_file.timestamp.strftime("%Y-%m-%d %H:%M:%S"), r.new_file])
            
            # If not with_files then we don't need to join the SubFile and 
            # ModelFile_SubFile tables
            else:
                query = (pm.Run_ModelFile
                            .select(pm.Run_ModelFile, pm.Run, pm.ModelFile)
                            .join(pm.ModelFile)
                            .switch(pm.Run_ModelFile)
                            .join(pm.Run))

                if not table == 'All Modelfiles':
                    query = query.where(pm.ModelFile.model_type == table)
                
                query = checkWildcard(query, pm.ModelFile.name, value1)
                
                if new_model_only:
                    query = query.where(pm.Run_ModelFile.new_file == True)
                
                if run_id != -1:
                    query = query.where(pm.Run.id == run_id)
                
                    cols = ['Run ID', 'Modelfile Timestamp', 'Modelfile New', 'Model Type', 'Modelfile', 'Comments']
                    rows = []
                    for r in query:
                        rows.append([r.run_id, r.timestamp.strftime("%Y-%m-%d %H:%M:%S"), r.new_file, r.model_file.model_type, r.model_file.name, r.model_file.comments])
                else:
                    cols = ['Run ID', 'Modelfile Timestamp', 'Modelfile New', 'Model Type', 'Modelfile', 'Comments']
                    rows = []
                    for r in query:
                        rows.append([r.run_id, r.timestamp.strftime("%Y-%m-%d %H:%M:%S"), r.new_file, r.model_file.model_type, r.model_file.name, r.model_file.comments])
                    
    finally: 
        pm.disconnectDB()
    
    return cols, rows
def updateNewStatus():
    """Updates the Run_ModelFile.new_file and ModelFile_SubFile.new_file status flags.
    
    Updated the new_file status flag to True for all ModelFile and SubFile
    records that come first when ordering by name and timestamp.
    
    For Run_ModelFile the tables are ordered by ModelFile.name, 
    Run_ModelFile.timestamp and the first one updated to True.
    
    For ModelFile_SubFile the tables are ordered by ModelFile.model_type, 
    ModelFile_SubFile.sub_file_id, ModelFile_SubFile.timestamp and the first 
    one updated to True.
    """    
    pm.connectDB()
    try:
        query = (pm.ModelFile_SubFile
                    .select(pm.ModelFile_SubFile, pm.ModelFile, pm.SubFile)
                    #.distinct(pm.ModelFile_SubFile.sub_file_id)
                    .join(pm.SubFile)
                    .switch()#pm.ModelFile_SubFile)
                    .join(pm.ModelFile)
                    .order_by(
                            pm.ModelFile.model_type, 
                            pm.ModelFile_SubFile.sub_file_id, 
                            pm.ModelFile_SubFile.timestamp)
                    )
        
        query.execute()
    except:
        pm.disconnectDB()
        raise
    
    
    try:
        query2 = (pm.Run_ModelFile
                    .select(pm.Run_ModelFile, pm.Run, pm.ModelFile)
                    #.distinct(pm.Run_ModelFile.model_file_id)
                    .join(pm.ModelFile)
                    .switch(pm.Run_ModelFile)
                    .join(pm.Run)
                    .order_by(
                            pm.ModelFile.name, 
                            pm.Run_ModelFile.timestamp)
                    )
        
        query2.execute()
    except:
        pm.disconnectDB()
        raise
    
    try:
        found_names = []
        with pm.logit_db.atomic():
            for q in query:
                i = q.id
                n = q.sub_file_id
                if not n in found_names:
                    found_names.append(n)
                    q = pm.ModelFile_SubFile.update(new_file=True).where(pm.ModelFile_SubFile.id == i)
                    q.execute()
            
            run_found_names = []
            for q in query2:
                ri = q.id
                rn = q.model_file_id
                if not rn in run_found_names:
                    run_found_names.append(rn)
                    q = pm.Run_ModelFile.update(new_file=True).where(pm.Run_ModelFile.id == ri)
                    q.execute()
    finally:
        pm.disconnectDB()
def deleteRunRow(run_id, delete_recursive=False):
    """Delete a record in the Run table.
    
    Delete's a Run record. If delete_recursive it will delete any associated
    records in the other tables. This is done in a (semi)intelligent way:
    All ModelFile records associated with this run_id in the Run_ModelFile
    table will be deleted IF they are not refernced by another Run.id in the
    Run_ModelFile table. This is also True of the Dat reference.
    
    TODO:
        I don't think this is very well written at the moment and is definitely
            quite slow.
    
    Args:
        run_id(int): the Run.id value to query against.
        delete_recursive=False(bool): if True will delete any foreign key
            associations and remove any associated Dat and ModelFile records.
    """
    pm.connectDB()
    
    try:
        model_del = []
        ied_del = []
        dat = None
        r = None
        try:
            r = pm.Run.get(pm.Run.id == run_id)
        except Exception as err:
            logger.warning('Could not find entry for run_id = %s' % run_id)
            logger.exception(err)

        if delete_recursive:
            
            # If we have a dat in this run
            if r.dat is not None:
                run_d = pm.Dat.get(name=r.dat)
                
                # If that dat isn't referenced in any other runs
                if pm.Run.select().filter(dat=run_d.name).count() < 2:
                    dat = run_d.name
            
    #         # Get all the run_subfile's referenced
    #         sq = pm.Run_SubFile.select().where(pm.Run_SubFile.run_id == run_id)

            # Get all modelfile's and a count of how many times there are referenced
            mq = (pm.Run_ModelFile
                    .select(pm.Run_ModelFile.run_id, pm.Run_ModelFile.model_file_id, fn.Count(pm.Run_ModelFile.model_file_id).alias('count'))
                    .group_by(pm.Run_ModelFile.model_file_id)
                 )
            # If they are only referenced by this run add to the delete list
            # TODO: VERY slow. Must try harder
            for m in mq:
                if m.count < 2 and m.run_id == run_id:
                    model_del.append(m.model_file_id)
            
            # Do the same for ied files
            iq = (pm.Run_Ied
                    .select(pm.Run_Ied.run_id, pm.Run_Ied.ied_id, fn.Count(pm.Run_Ied.ied_id).alias('count'))
                    .group_by(pm.Run_Ied.ied_id)
                 )
            for i in iq:
                if i.count < 2 and i.run_id == run_id:
                    ied_del.append(i.ied_id)
            

        r.delete_instance(recursive=delete_recursive)
        
        with pm.logit_db.atomic():
            if dat is not None: 
                deleteDatRow(dat, connect_db=False)
            for m in model_del:
                deleteModelRow(m, remove_orphans=False, connect_db=False)
            for i in ied_del:
                deleteIedRow(i, remove_orphans=False, connect_db=False)

    finally:
        pm.disconnectDB()