graph_builder.py

"""
Graph builder module for converting GDELT data to graph formats.

This module provides classes for building graphs from GDELT data in various formats,
including NetworkX, Neo4j, and st-link-analysis. It supports batch processing,
input validation, custom properties, and logging for robust and efficient graph construction.
"""
import pandas as pd
import networkx as nx
import logging
import json
from neo4j import GraphDatabase
from typing import Dict, List, Optional, Set, Tuple

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class GraphBuilder:
    """
    Base class for building graphs from GDELT data.

    Attributes:
        ENTITY_MAPPINGS (dict): Mapping of GDELT fields to node types and relationships.
        logger (Logger): Logger instance for tracking progress and errors.
    """
    ENTITY_MAPPINGS = {
        "V2EnhancedPersons": ("Person", "MENTIONED_IN"),
        "V2EnhancedOrganizations": ("Organization", "MENTIONED_IN"),
        "V2EnhancedLocations": ("Location", "LOCATED_IN"),
        "V2EnhancedThemes": ("Theme", "CATEGORIZED_AS"),
        "V2.1AllNames": ("Name", "MENTIONED_IN"),
        "V2.1Counts": ("Count", "MENTIONED_IN"),
        "V2.1Amounts": ("Amount", "MENTIONED_IN"),
    }

    def __init__(self):
        self.logger = logger

    def validate_input(self, df):
        """
        Validate input DataFrame for required columns.

        Args:
            df (pd.DataFrame): Input DataFrame containing GDELT data.

        Raises:
            ValueError: If any required column is missing.
        """
        required_columns = ["GKGRECORDID", "DATE", "SourceCommonName", "DocumentIdentifier"]
        for col in required_columns:
            if col not in df.columns:
                raise ValueError(f"Missing required column: {col}")

    def process_entities(self, row, custom_node_props=None, custom_edge_props=None):
        """
        Process entities from a row and return nodes and relationships.

        Args:
            row (pd.Series): A row of GDELT data.
            custom_node_props (dict, optional): Custom properties for nodes.
            custom_edge_props (dict, optional): Custom properties for edges.

        Returns:
            Tuple[List[Dict], List[Dict]]: Lists of nodes and relationships.
        """
        nodes = []
        relationships = []
        event_id = row["GKGRECORDID"]
        event_date = row["DATE"]
        event_source = row["SourceCommonName"]
        event_document_id = row["DocumentIdentifier"]
        event_quotations = row["V2.1Quotations"] if pd.notna(row["V2.1Quotations"]) else ""
        event_tone = float(row["tone"]) if pd.notna(row["tone"]) else 0.0

        # Add event node with custom properties
        event_props = {
            "date": event_date,
            "source": event_source,
            "document": event_document_id,
            "quotations": event_quotations,
            "tone": event_tone
        }
        if custom_node_props:
            event_props.update(custom_node_props)

        nodes.append({
            "id": event_id,
            "type": "event",
            "properties": event_props
        })

        # Process each entity type
        for field, (label, relationship) in self.ENTITY_MAPPINGS.items():
            if pd.notna(row[field]):
                entities = [e.strip() for e in row[field].split(';') if e.strip()]
                for entity in entities:
                    nodes.append({
                        "id": entity,
                        "type": label.lower(),
                        "properties": {"name": entity}
                    })
                    edge_props = {"created_at": event_date}
                    if custom_edge_props:
                        edge_props.update(custom_edge_props)
                    relationships.append({
                        "from": entity,
                        "to": event_id,
                        "type": relationship,
                        "properties": edge_props
                    })

        return nodes, relationships

    def validate_graph(self, G):
        """
        Validate the graph for consistency.

        Args:
            G (nx.Graph): The graph to validate.

        Raises:
            ValueError: If the graph is invalid.
        """
        if not isinstance(G, nx.Graph):
            raise ValueError("Input is not a valid NetworkX graph.")
        if len(G.nodes) == 0:
            raise ValueError("Graph has no nodes.")
        if len(G.edges) == 0:
            raise ValueError("Graph has no edges.")
        self.logger.info("Graph validation passed.")

class NetworkXBuilder(GraphBuilder):
    """
    Builder for NetworkX graphs.

    Attributes:
        directed (bool): Whether to create a directed graph.
    """
    def __init__(self, directed=False):
        super().__init__()
        self.directed = directed

    def build_graph(self, df, custom_node_props=None, custom_edge_props=None):
        """
        Build a NetworkX graph from the DataFrame.

        Args:
            df (pd.DataFrame): Input DataFrame containing GDELT data.
            custom_node_props (dict, optional): Custom properties for nodes.
            custom_edge_props (dict, optional): Custom properties for edges.

        Returns:
            nx.Graph: The constructed graph.
        """
        self.validate_input(df)
        G = nx.DiGraph() if self.directed else nx.Graph()

        try:
            for _, row in df.iterrows():
                nodes, relationships = self.process_entities(row, custom_node_props, custom_edge_props)

                # Add nodes
                for node in nodes:
                    G.add_node(node["id"], type=node["type"], **node["properties"])

                # Add relationships
                for rel in relationships:
                    G.add_edge(rel["from"], rel["to"], relationship=rel["type"], **rel["properties"])

            self.validate_graph(G)
            self.logger.info("NetworkX graph built successfully.")
            return G
        except Exception as e:
            self.logger.error(f"Error building NetworkX graph: {str(e)}")
            raise

class Neo4jBuilder(GraphBuilder):
    """
    Builder for Neo4j graphs.

    Attributes:
        driver (neo4j.Driver): Neo4j driver instance.
        logger (Logger): Logger instance for tracking progress and errors.
    """
    def __init__(self, uri, user, password):
        super().__init__()
        self.driver = GraphDatabase.driver(uri, auth=(user, password))

    def close(self):
        """Close the Neo4j driver."""
        self.driver.close()

    def build_graph(self, df, batch_size=1000, custom_node_props=None, custom_edge_props=None):
        """
        Build a Neo4j graph from the DataFrame with batch processing.

        Args:
            df (pd.DataFrame): Input DataFrame containing GDELT data.
            batch_size (int): Number of rows to process in each batch.
            custom_node_props (dict, optional): Custom properties for nodes.
            custom_edge_props (dict, optional): Custom properties for edges.
        """
        self.validate_input(df)

        with self.driver.session() as session:
            batch = []
            for _, row in df.iterrows():
                nodes, relationships = self.process_entities(row, custom_node_props, custom_edge_props)
                batch.append((nodes, relationships))
                if len(batch) >= batch_size:
                    session.execute_write(self._create_graph_elements_batch, batch)
                    batch = []
            if batch:
                session.execute_write(self._create_graph_elements_batch, batch)

            self.logger.info("Neo4j graph built successfully.")

    def _create_graph_elements_batch(self, tx, batch):
        """
        Create nodes and relationships in Neo4j in batches.

        Args:
            tx (neo4j.Transaction): Neo4j transaction.
            batch (List[Tuple[List[Dict], List[Dict]]]): Batch of nodes and relationships.
        """
        for nodes, relationships in batch:
            # Create nodes
            for node in nodes:
                query = f"""
                MERGE (n:{node['type']} {{id: $id}})
                SET n += $properties
                """
                tx.run(query, id=node["id"], properties=node["properties"])

            # Create relationships
            for rel in relationships:
                query = f"""
                MATCH (a {{id: $from_id}})
                MATCH (b {{id: $to_id}})
                MERGE (a)-[r:{rel['type']}]->(b)
                SET r += $properties
                """
                tx.run(query, from_id=rel["from"], to_id=rel["to"], properties=rel["properties"])

class StLinkBuilder(GraphBuilder):
    """
    Builder for st-link-analysis compatible graphs.

    Attributes:
        logger (Logger): Logger instance for tracking progress and errors.
    """
    def __init__(self):
        super().__init__()

    def build_graph(self, df, custom_node_props=None, custom_edge_props=None):
        """
        Build graph in st-link-analysis format.

        Args:
            df (pd.DataFrame): Input DataFrame containing GDELT data.
            custom_node_props (dict, optional): Custom properties for nodes.
            custom_edge_props (dict, optional): Custom properties for edges.

        Returns:
            Dict: Graph data in st-link-analysis format.
        """
        self.validate_input(df)
        all_nodes = []
        all_edges = []
        edge_counter = 0
        added_nodes = set()

        try:
            for _, row in df.iterrows():
                nodes, relationships = self.process_entities(row, custom_node_props, custom_edge_props)

                # Process nodes
                for node in nodes:
                    if node["id"] not in added_nodes:
                        stlink_node = {
                            "data": {
                                "id": str(node["id"]),
                                "label": node["type"].upper(),
                                **node["properties"]
                            }
                        }
                        all_nodes.append(stlink_node)
                        added_nodes.add(node["id"])

                # Process relationships/edges
                for rel in relationships:
                    edge_counter += 1
                    stlink_edge = {
                        "data": {
                            "id": f"e{edge_counter}",
                            "source": str(rel["from"]),
                            "target": str(rel["to"]),
                            "label": rel["type"],
                            **rel["properties"]
                        }
                    }
                    all_edges.append(stlink_edge)

            self.logger.info("st-link-analysis graph built successfully.")
            return {
                "nodes": all_nodes,
                "edges": all_edges
            }
        except Exception as e:
            self.logger.error(f"Error building st-link-analysis graph: {str(e)}")
            raise

    def write_json(self, graph_data, filename):
        """
        Write graph to JSON file with streaming.

        Args:
            graph_data (Dict): Graph data in st-link-analysis format.
            filename (str): Output file name.
        """
        try:
            with open(filename, 'w') as f:
                json.dump(graph_data, f, indent=2)
            self.logger.info(f"Graph data written to {filename} successfully.")
        except Exception as e:
            self.logger.error(f"Error writing JSON file: {str(e)}")
            raise