neo4j_uploader.models

  1from pydantic import BaseModel, Field
  2from typing import Optional
  3
  4# Specify Google doctstring type for pdoc auto doc generation
  5__docformat__ = "google"
  6
  7class Neo4jConfig(BaseModel):
  8    """
  9    Object for specifying target local or hosted Neo4j database instance to upload to data to.
 10
 11    Args:
 12        neo4j_uri (str): The URI of the Neo4j instance to upload to.
 13        neo4j_password (str): The password for the Neo4j instance to upload to.
 14        neo4j_user (str): The username for the Neo4j instance to upload to.
 15        neo4j_database (str): The name of the Neo4j database to upload to. Default 'neo4j'.
 16        max_batch_size (int): Maximum number of nodes to upload in a single batch. Default 500.
 17        overwrite (bool): Overwrite existing nodes. Default False.
 18    """
 19    neo4j_uri : str
 20    neo4j_password: str
 21    neo4j_user : str = Field(default="neo4j")
 22    neo4j_database : str = Field(default="neo4j")
 23    max_batch_size : int = Field(default=500)
 24    overwrite : bool = False
 25
 26class Nodes(BaseModel):
 27    """Configuration object for uploading nodes to a Neo4j database.
 28
 29    Args:
 30        labels (list[str]): List of node labels to upload (ie Person, Place, etc).
 31        key (str): Unique key for each node.
 32        records (list[dict]): List of dictionary objects containing node data.
 33        exclude_keys (list[str]): List of keys to exclude from upload.
 34        dedupe (bool): Remove duplicate entries. Default True.
 35    """
 36    labels: list[str]
 37    key: str
 38    records: list[dict]
 39    exclude_keys: Optional[list[str]] = []
 40    dedupe : Optional[bool] = True
 41
 42class TargetNode(BaseModel):
 43    """Node specification object for uploading relationships to a Neo4j database.
 44
 45    Args:
 46        node_label (str): Optional Node label of the target node. Including a label is more performant than without.
 47        node_key (str): Target key or property name that identifies a unique node.
 48        record_key (str): Key within the relationship record, whose value will be mapped to the node_key.
 49    """
 50    node_label: Optional[str] = None
 51    node_key: str
 52    record_key: str
 53
 54class Relationships(BaseModel):
 55    """Configuration object for uploading relationships to a Neo4j database.
 56
 57    Args:
 58        type (str): Relationship type (ie FOLLOWS, WORKS_AT, etc).
 59        from_node (TargetNode): TargetNode object for the source node.
 60        to_node (TargetNode): TargetNode object for the target node.
 61        records (list[dict]): List of dictionary objects containing relationship data.
 62        exclude_keys (list[str]): List of keys to exclude from upload.
 63        auto_exclude_keys (bool): Automatically exclude keys used to reference nodes used in the from_node and to_node arguments. Default True.
 64        dedupe (bool): Remove duplicate entries. Default True.
 65    """
 66    type: str
 67    from_node: TargetNode
 68    to_node: TargetNode
 69    records : list[dict]
 70    exclude_keys: Optional[list[str]] = []
 71    auto_exclude_keys: Optional[bool] = True
 72    dedupe: Optional[bool] = True
 73
 74class GraphData(BaseModel):
 75    """Object representation of nodes and relationships specifications and records to upload to a specified Neo4j database.
 76
 77    Args:
 78        nodes (list[Nodes]): List of Nodes configuration objects for specifying nodes to upload.
 79        relationships (list[Relationships]): Optional List of Relationships configuration objects for specifying relationships to upload.
 80    """
 81    nodes: list[Nodes] = []
 82    relationships: Optional[list[Relationships]] = []
 83
 84class UploadResult(BaseModel):
 85    """Result object for uploading nodes to a Neo4j database.
 86
 87    Args:
 88        was_successful (bool): True if upload was successful.
 89        seconds_to_complete (float): Time taken to upload nodes in seconds.
 90        nodes_created (int): Number of nodes created.
 91        relationships_created (int): Number of relationships created.
 92        properties_set (int): Number of properties set.
 93        error_message (str): Error message if upload failed.
 94    """
 95    was_successful : bool
 96    seconds_to_complete: Optional[float] = None
 97    nodes_created: Optional[int] = None
 98    relationships_created: Optional[int] = None
 99    properties_set: Optional[int] = None
100    error_message: Optional[str] = None