batch_functions.py

from typing import Optional
from object_classes import ClinicalTrial


def load_trial_from_db(nct_id: str, db_connection) -> Optional[ClinicalTrial]:
    """Load trial data from database and return as ClinicalTrial object"""
    cursor = db_connection.cursor(dictionary=True)
    cursor.execute("SELECT * FROM trials WHERE nct_id = %s", (nct_id,))
    trial_data = cursor.fetchone()
    trial_id = trial_data['trial_id']
    if trial_data:
        trial = ClinicalTrial(**trial_data)

    #interventions
    cursor.execute("SELECT * FROM interventions WHERE trial_id_int = %s", (trial_id,))
    interventions_data = cursor.fetchall()
    if interventions_data:
        trial['interventions'] = interventions_data

    #arm groups
    cursor.execute("SELECT * FROM arm_groups WHERE trial_id_arm = %s", (trial_id,))
    arm_groups_data = cursor.fetchall()
    if arm_groups_data:
        trial['arm_groups'] = arm_groups_data

    #diseases
    cursor.execute("SELECT * FROM diseases WHERE trial_id_dis = %s", (trial_id,))
    diseases_data = cursor.fetchall()
    if diseases_data:
        trial['diseases'] = diseases_data

    #drugs and targets   
    cursor.execute("SELECT * FROM drugs WHERE trial_id_dru = %s", (trial_id,))
    drugs_data = cursor.fetchall()
    if drugs_data:
         trial['drugs'] = drugs_data
    for drug in drugs_data:
        drug_id = drug['drug_id']
        cursor.execute("SELECT * FROM targets WHERE drug_id_tar = %s", (drug_id,))
        targets_data = cursor.fetchall()
        if targets_data:
            drug['targets'] = targets_data
    
    return trial

def check_gene_names(trial: ClinicalTrial, hgnc_data) -> ClinicalTrial:
    if 'drugs' in trial and trial['drugs']:
        for drug in trial['drugs']:
            if 'targets' in drug and drug['targets']:
                for target in drug['targets']:
                    gene_symbol = target['gene']
                    gene_row = hgnc_data[hgnc_data['Approved symbol'] == gene_symbol]
                    if not gene_row.empty:
                        target['hgnc_id'] = gene_row['HGNC ID'].values[0]
                        target['gene_symbol'] = gene_symbol
                    else:
                        gene_row = hgnc_data[hgnc_data['Previous symbols'] == gene_symbol]
                        if not gene_row.empty:
                            target['hgnc_id'] = gene_row['HGNC ID'].values[0]
                            target['gene_symbol'] = gene_row['Approved symbol'].values[0]
                        else:
                            gene_row = hgnc_data[hgnc_data['Alias symbols'] == gene_symbol]
                            if not gene_row.empty:
                                target['hgnc_id'] = gene_row['HGNC ID'].values[0]
                                target['gene_symbol'] = gene_row['Approved symbol'].values[0]
                            else:
                                target['hgnc_id'] = None
                                target['gene_symbol'] = "UNKNOWN"
    return trial


from mysql.connector import Error
from typing import Dict, Any, List

def insert_trial_data(trial: Dict[str, Any], connection, cursor) -> str:
    """Insert trial and all related objects into database"""
    try:
        # Insert trial base data
        trial_sql = """
        INSERT INTO trials (
            nct_id, trial_type, brief_title, official_title,
            conditions, conditions_mesh_terms, keywords, study_type, 
            intervention_mesh_terms, eligibility_summary,
            trial_source, annotation_status, ai_summary, annotation_time_seconds
        ) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
        """
        cursor.execute(trial_sql, (
            trial['nct_id'],
            trial['trial_type'],
            trial['brief_title'],
            trial['official_title'], 
            ','.join(trial['conditions']),
            ','.join(trial['conditions_mesh_terms']),
            ','.join(trial['keywords']),
            trial['study_type'],
            ','.join(trial['intervention_mesh_terms']),
            trial.get('eligibility_summary'),
            trial.get('trial_source'),
            trial.get('annotation_status'),
            trial.get('ai_summary'),
            trial.get('annotation_time_seconds')
        ))
        trial_id = cursor.lastrowid


        # Insert arm groups
        for arm in trial['arm_groups']:
            arm_sql = """
            INSERT INTO arm_groups (
                trial_id_arm, arm_group_label, arm_group_type, arm_group_description, intervention_names
            ) VALUES (%s, %s, %s, %s, %s)
            """
            cursor.execute(arm_sql, (
                trial_id,
                arm['arm_group_label'],
                arm['arm_group_type'],
                arm['arm_group_description'],
                ','.join(arm['intervention_names'])
            ))

        # Insert interventions
        for intervention in trial['interventions']:
            intervention_sql = """
            INSERT INTO interventions (
                trial_id_int, intervention_type_int, intervention_name, intervention_description, arm_group_labels, other_names
            ) VALUES (%s, %s, %s, %s, %s, %s)
            """
            cursor.execute(intervention_sql, (
                trial_id,
                intervention['intervention_type_int'],
                intervention['intervention_name'],
                intervention['intervention_description'],
                ','.join(intervention['arm_group_labels']),
                ','.join(intervention['other_names'])
            ))

        # Insert diseases
        if 'diseases' in trial and trial['diseases']:
            for disease in trial['diseases']:
                disease_sql = """
                INSERT INTO diseases (
                    trial_id_dis, disease_type, disease_name, tissue_type,
                    disease_mesh_term, oncotree_name, oncotree_main_type
                ) VALUES (%s, %s, %s, %s, %s, %s, %s)
                """
            cursor.execute(disease_sql, (
                trial_id,
                disease['disease_type'],
                disease['disease_name'],
                disease['tissue_type'],
                disease.get('disease_mesh_term'),
                disease.get('Oncotree_name'),
                disease.get('Oncotree_mainType')
            ))

        # Insert drugs and their targets
        if 'drugs' in trial and trial['drugs']:
            for drug in trial['drugs']:
                drug_sql = """
                INSERT INTO drugs (
                    trial_id_dru, drug_name, intervention_type, drug_description,
                    drug_type, nci_thesaurus_concept_id, nci_thesaurus_preferred_term,
                    nci_thesaurus_definition, drug_mesh_term, drug_category,
                    drug_class, drug_delivery_route, drug_mechanism_of_action
                ) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
                """
                cursor.execute(drug_sql, (
                    trial_id,
                    drug['drug_name'],
                    drug['intervention_type'],
                    drug['drug_description'],
                    drug['drug_type'],
                    drug.get('NCI_thesaurus_concept_id'),
                    drug.get('NCI_thesaurus_preferred_term'),
                    drug.get('NCI_thesaurus_definition'),
                    drug.get('drug_mesh_term'),
                    drug.get('drug_category'),
                    drug.get('drug_class'),
                    drug.get('drug_delivery_route'),
                    drug.get('drug_mechanism_of_action')
                ))
                drug_id = cursor.lastrowid

                # Insert targets
                if 'targets' in drug and drug['targets']:
                    for target in drug['targets']:
                        target_sql = """
                        INSERT INTO targets (
                            drug_id_tar, gene, gene_symbol, hgnc_id, antigen_ligand,
                            binding_type, cytotoxicity, cytotoxicity_type, cytotoxicity_mechanism,
                            enzyme_product, epitope, variant, isoform,
                            hla_specific, gated, other_modifier
                        ) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
                        """
                        cursor.execute(target_sql, (
                            drug_id,
                            target['gene'],
                            target['gene_symbol'],
                            target['hgnc_id'],
                            target['antigen_ligand'],
                            target['binding_type'],
                            target.get('cytotoxicity'),
                            target.get('cytotoxicity_type'),
                            target.get('cytotoxicity_mechanism'),
                            target['enzyme_product'],
                            target['epitope'],
                            target['variant'],
                            target['isoform'],
                            target['HLA_specific'],
                            target['gated'],
                            target.get('other_modifier')
                        ))


        connection.commit()
        print(f"Successfully inserted trial {trial['nct_id']}")

    except Error as e:
        print(f"Error inserting trial data: {e}")
        nct_error = trial['nct_id']
        connection.rollback()
        return nct_error
    return ""


def update_gene_data(trial: Dict[str, Any], connection, cursor) -> str:
    try:
        #query drug id
        if 'drugs' in trial and trial['drugs']:
            for drug in trial['drugs']:
                drug_id = drug['drug_id']
                if 'targets' in drug and drug['targets']:
                    for target in drug['targets']:
                        target_sql = """
                        UPDATE targets SET gene_symbol = %s, hgnc_id = %s WHERE target_id = %s
                        """
                        cursor.execute(target_sql, (
                            target['gene_symbol'],
                            target['hgnc_id'],
                            target['target_id']
                        ))


        connection.commit()
        print(f"Successfully update gene names for trial {trial['nct_id']}")

    except Error as e:
        print(f"Error updating trial data: {e}")
        nct_error = trial['nct_id']
        connection.rollback()
        return nct_error
    return ""


def update_trial_status(trial: Dict[str, Any], connection, cursor) -> str:
   
    try:
        trial_id = trial['trial_id']
        trial_sql = """
        UPDATE trials SET annotation_status = %s WHERE trial_id = %s
        """

        cursor.execute(trial_sql, (
            trial['annotation_status'],
            trial_id
        ))

        connection.commit()
        print(f"Successfully update trial status for trial {trial['nct_id']}")

    except Error as e:
        print(f"Error updating trial data: {e}")
        nct_error = trial['nct_id']
        connection.rollback()
        return nct_error
    return ""