|
|
"""
|
|
|
Neo4j Database Manager
|
|
|
Handle graph database connections and operations
|
|
|
"""
|
|
|
|
|
|
from neo4j import GraphDatabase
|
|
|
from typing import Dict, List, Optional, Any
|
|
|
import yaml
|
|
|
import logging
|
|
|
|
|
|
logging.basicConfig(level=logging.INFO)
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
|
|
|
class DatabaseManager:
|
|
|
"""Manage Neo4j database connections and schema"""
|
|
|
|
|
|
def __init__(self, config_path: str = "config.yml"):
|
|
|
with open(config_path, 'r') as f:
|
|
|
self.config = yaml.safe_load(f)['neo4j']
|
|
|
|
|
|
self.driver = GraphDatabase.driver(
|
|
|
self.config['uri'],
|
|
|
auth=(self.config['username'], self.config['password'])
|
|
|
)
|
|
|
|
|
|
logger.info(f"Connected to Neo4j at {self.config['uri']}")
|
|
|
|
|
|
def close(self):
|
|
|
"""Close database connection"""
|
|
|
self.driver.close()
|
|
|
|
|
|
def __enter__(self):
|
|
|
return self
|
|
|
|
|
|
def __exit__(self, exc_type, exc_val, exc_tb):
|
|
|
self.close()
|
|
|
|
|
|
def execute_query(self, query: str, parameters: Optional[Dict] = None) -> List[Dict]:
|
|
|
"""Execute a Cypher query and return results"""
|
|
|
with self.driver.session() as session:
|
|
|
result = session.run(query, parameters or {})
|
|
|
return [record.data() for record in result]
|
|
|
|
|
|
def initialize_schema(self):
|
|
|
"""Initialize database schema with constraints and indexes"""
|
|
|
queries = [
|
|
|
|
|
|
"CREATE CONSTRAINT gene_id IF NOT EXISTS FOR (g:Gene) REQUIRE g.gene_id IS UNIQUE",
|
|
|
"CREATE CONSTRAINT mutation_id IF NOT EXISTS FOR (m:Mutation) REQUIRE m.mutation_id IS UNIQUE",
|
|
|
"CREATE CONSTRAINT patient_id IF NOT EXISTS FOR (p:Patient) REQUIRE p.patient_id IS UNIQUE",
|
|
|
"CREATE CONSTRAINT cancer_type_id IF NOT EXISTS FOR (c:CancerType) REQUIRE c.cancer_type_id IS UNIQUE",
|
|
|
|
|
|
|
|
|
"CREATE INDEX gene_symbol IF NOT EXISTS FOR (g:Gene) ON (g.symbol)",
|
|
|
"CREATE INDEX mutation_position IF NOT EXISTS FOR (m:Mutation) ON (m.chromosome, m.position)",
|
|
|
"CREATE INDEX patient_project IF NOT EXISTS FOR (p:Patient) ON (p.project_id)",
|
|
|
]
|
|
|
|
|
|
with self.driver.session() as session:
|
|
|
for query in queries:
|
|
|
try:
|
|
|
session.run(query)
|
|
|
logger.info(f"Executed: {query[:50]}...")
|
|
|
except Exception as e:
|
|
|
logger.warning(f"Schema query failed (may already exist): {e}")
|
|
|
|
|
|
logger.info("Database schema initialized")
|
|
|
|
|
|
def clear_database(self):
|
|
|
"""Clear all nodes and relationships (use with caution!)"""
|
|
|
query = "MATCH (n) DETACH DELETE n"
|
|
|
with self.driver.session() as session:
|
|
|
session.run(query)
|
|
|
logger.info("Database cleared")
|
|
|
|
|
|
|
|
|
class GeneRepository:
|
|
|
"""Repository for Gene nodes"""
|
|
|
|
|
|
def __init__(self, db_manager: DatabaseManager):
|
|
|
self.db = db_manager
|
|
|
|
|
|
def create_gene(self, gene_data: Dict) -> Dict:
|
|
|
"""Create a Gene node"""
|
|
|
query = """
|
|
|
MERGE (g:Gene {gene_id: $gene_id})
|
|
|
SET g.symbol = $symbol,
|
|
|
g.name = $name,
|
|
|
g.chromosome = $chromosome,
|
|
|
g.start_position = $start_position,
|
|
|
g.end_position = $end_position,
|
|
|
g.strand = $strand,
|
|
|
g.gene_type = $gene_type
|
|
|
RETURN g
|
|
|
"""
|
|
|
result = self.db.execute_query(query, gene_data)
|
|
|
return result[0]['g'] if result else {}
|
|
|
|
|
|
def get_gene_by_symbol(self, symbol: str) -> Optional[Dict]:
|
|
|
"""Find gene by symbol"""
|
|
|
query = """
|
|
|
MATCH (g:Gene {symbol: $symbol})
|
|
|
RETURN g
|
|
|
"""
|
|
|
result = self.db.execute_query(query, {'symbol': symbol})
|
|
|
return result[0]['g'] if result else None
|
|
|
|
|
|
def get_gene_mutations(self, gene_id: str) -> List[Dict]:
|
|
|
"""Get all mutations for a gene"""
|
|
|
query = """
|
|
|
MATCH (g:Gene {gene_id: $gene_id})<-[:AFFECTS]-(m:Mutation)
|
|
|
RETURN m
|
|
|
ORDER BY m.position
|
|
|
"""
|
|
|
result = self.db.execute_query(query, {'gene_id': gene_id})
|
|
|
return [r['m'] for r in result]
|
|
|
|
|
|
|
|
|
class MutationRepository:
|
|
|
"""Repository for Mutation nodes"""
|
|
|
|
|
|
def __init__(self, db_manager: DatabaseManager):
|
|
|
self.db = db_manager
|
|
|
|
|
|
def create_mutation(self, mutation_data: Dict, gene_id: str) -> Dict:
|
|
|
"""Create a Mutation node and link to Gene"""
|
|
|
query = """
|
|
|
MATCH (g:Gene {gene_id: $gene_id})
|
|
|
MERGE (m:Mutation {mutation_id: $mutation_id})
|
|
|
SET m.chromosome = $chromosome,
|
|
|
m.position = $position,
|
|
|
m.reference = $reference,
|
|
|
m.alternate = $alternate,
|
|
|
m.consequence = $consequence,
|
|
|
m.variant_type = $variant_type,
|
|
|
m.quality = $quality
|
|
|
MERGE (m)-[:AFFECTS]->(g)
|
|
|
RETURN m
|
|
|
"""
|
|
|
params = {**mutation_data, 'gene_id': gene_id}
|
|
|
result = self.db.execute_query(query, params)
|
|
|
return result[0]['m'] if result else {}
|
|
|
|
|
|
def link_mutation_to_patient(self, mutation_id: str, patient_id: str, properties: Optional[Dict] = None):
|
|
|
"""Create HAS_MUTATION relationship"""
|
|
|
query = """
|
|
|
MATCH (p:Patient {patient_id: $patient_id})
|
|
|
MATCH (m:Mutation {mutation_id: $mutation_id})
|
|
|
MERGE (p)-[r:HAS_MUTATION]->(m)
|
|
|
SET r.allele_frequency = $allele_frequency,
|
|
|
r.depth = $depth
|
|
|
RETURN r
|
|
|
"""
|
|
|
params = {
|
|
|
'patient_id': patient_id,
|
|
|
'mutation_id': mutation_id,
|
|
|
'allele_frequency': properties.get('allele_frequency', 0) if properties else 0,
|
|
|
'depth': properties.get('depth', 0) if properties else 0
|
|
|
}
|
|
|
self.db.execute_query(query, params)
|
|
|
|
|
|
def get_mutation_frequency(self, mutation_id: str) -> Dict:
|
|
|
"""Calculate mutation frequency across patients"""
|
|
|
query = """
|
|
|
MATCH (m:Mutation {mutation_id: $mutation_id})
|
|
|
MATCH (p:Patient)-[:HAS_MUTATION]->(m)
|
|
|
OPTIONAL MATCH (all:Patient)
|
|
|
WITH m, count(DISTINCT p) as patients_with_mutation, count(DISTINCT all) as total_patients
|
|
|
RETURN m.mutation_id as mutation_id,
|
|
|
patients_with_mutation,
|
|
|
total_patients,
|
|
|
toFloat(patients_with_mutation) / total_patients as frequency
|
|
|
"""
|
|
|
result = self.db.execute_query(query, {'mutation_id': mutation_id})
|
|
|
return result[0] if result else {}
|
|
|
|
|
|
|
|
|
class PatientRepository:
|
|
|
"""Repository for Patient nodes"""
|
|
|
|
|
|
def __init__(self, db_manager: DatabaseManager):
|
|
|
self.db = db_manager
|
|
|
|
|
|
def create_patient(self, patient_data: Dict) -> Dict:
|
|
|
"""Create a Patient node"""
|
|
|
query = """
|
|
|
MERGE (p:Patient {patient_id: $patient_id})
|
|
|
SET p.project_id = $project_id,
|
|
|
p.age = $age,
|
|
|
p.gender = $gender,
|
|
|
p.race = $race,
|
|
|
p.ethnicity = $ethnicity,
|
|
|
p.vital_status = $vital_status
|
|
|
RETURN p
|
|
|
"""
|
|
|
result = self.db.execute_query(query, patient_data)
|
|
|
return result[0]['p'] if result else {}
|
|
|
|
|
|
def link_patient_to_cancer_type(self, patient_id: str, cancer_type_id: str, properties: Optional[Dict] = None):
|
|
|
"""Create DIAGNOSED_WITH relationship"""
|
|
|
query = """
|
|
|
MATCH (p:Patient {patient_id: $patient_id})
|
|
|
MATCH (c:CancerType {cancer_type_id: $cancer_type_id})
|
|
|
MERGE (p)-[r:DIAGNOSED_WITH]->(c)
|
|
|
SET r.stage = $stage,
|
|
|
r.grade = $grade,
|
|
|
r.diagnosis_date = $diagnosis_date
|
|
|
RETURN r
|
|
|
"""
|
|
|
params = {
|
|
|
'patient_id': patient_id,
|
|
|
'cancer_type_id': cancer_type_id,
|
|
|
'stage': properties.get('stage') if properties else None,
|
|
|
'grade': properties.get('grade') if properties else None,
|
|
|
'diagnosis_date': properties.get('diagnosis_date') if properties else None
|
|
|
}
|
|
|
self.db.execute_query(query, params)
|
|
|
|
|
|
def get_patient_mutations(self, patient_id: str) -> List[Dict]:
|
|
|
"""Get all mutations for a patient"""
|
|
|
query = """
|
|
|
MATCH (p:Patient {patient_id: $patient_id})-[r:HAS_MUTATION]->(m:Mutation)-[:AFFECTS]->(g:Gene)
|
|
|
RETURN m, g, r.allele_frequency as allele_frequency, r.depth as depth
|
|
|
ORDER BY g.symbol
|
|
|
"""
|
|
|
result = self.db.execute_query(query, {'patient_id': patient_id})
|
|
|
return result
|
|
|
|
|
|
|
|
|
class CancerTypeRepository:
|
|
|
"""Repository for CancerType nodes"""
|
|
|
|
|
|
def __init__(self, db_manager: DatabaseManager):
|
|
|
self.db = db_manager
|
|
|
|
|
|
def create_cancer_type(self, cancer_data: Dict) -> Dict:
|
|
|
"""Create a CancerType node"""
|
|
|
query = """
|
|
|
MERGE (c:CancerType {cancer_type_id: $cancer_type_id})
|
|
|
SET c.name = $name,
|
|
|
c.tissue = $tissue,
|
|
|
c.disease_type = $disease_type
|
|
|
RETURN c
|
|
|
"""
|
|
|
result = self.db.execute_query(query, cancer_data)
|
|
|
return result[0]['c'] if result else {}
|
|
|
|
|
|
def get_common_mutations(self, cancer_type_id: str, limit: int = 10) -> List[Dict]:
|
|
|
"""Get most common mutations for a cancer type"""
|
|
|
query = """
|
|
|
MATCH (c:CancerType {cancer_type_id: $cancer_type_id})<-[:DIAGNOSED_WITH]-(p:Patient)
|
|
|
MATCH (p)-[:HAS_MUTATION]->(m:Mutation)-[:AFFECTS]->(g:Gene)
|
|
|
WITH m, g, count(DISTINCT p) as patient_count
|
|
|
RETURN m, g, patient_count
|
|
|
ORDER BY patient_count DESC
|
|
|
LIMIT $limit
|
|
|
"""
|
|
|
result = self.db.execute_query(query, {'cancer_type_id': cancer_type_id, 'limit': limit})
|
|
|
return result
|
|
|
|
|
|
def get_statistics(self, cancer_type_id: str) -> Dict:
|
|
|
"""Get statistics for a cancer type"""
|
|
|
query = """
|
|
|
MATCH (c:CancerType {cancer_type_id: $cancer_type_id})<-[:DIAGNOSED_WITH]-(p:Patient)
|
|
|
OPTIONAL MATCH (p)-[:HAS_MUTATION]->(m:Mutation)
|
|
|
WITH c, count(DISTINCT p) as total_patients, count(DISTINCT m) as total_mutations
|
|
|
RETURN c.name as cancer_type,
|
|
|
total_patients,
|
|
|
total_mutations,
|
|
|
CASE WHEN total_patients > 0
|
|
|
THEN toFloat(total_mutations) / total_patients
|
|
|
ELSE 0
|
|
|
END as avg_mutations_per_patient
|
|
|
"""
|
|
|
result = self.db.execute_query(query, {'cancer_type_id': cancer_type_id})
|
|
|
return result[0] if result else {}
|
|
|
|
