0.4.0 (#2)

* Custom node key support added

* Relationship upload process updated to use tuple data and unwind.

* Dedupe Nodes and relationship dictionary input payloads added

* Upload now returns seconds to complete, nodes created, relationships created, and total properties set

* Params added to query generation

* Dedupe modified to CREATE new nodes if enabled

* Target neo4j database name option added

* Subprocessing functions moved to a utils file + nodes also using WITH + UNWIND statement

* Batch processing added

* Batch configuration from main upload function

* version updated 0.4.0

Author: jalakoo

README.md

@@ -35,7 +35,7 @@ upload(credentials, data, node_key="name)
 * `node_key` - The name of the key in the nodes dictionary that contains the node name. Defaults to "name"
 * `dedupe_nodes` - If True, will not upload duplicate Nodes. Only updates existing with specified properties. Defaults to True
 * `dedupe_relationships` - If True, will not created duplicate Relationships. Only updates existing with specified properties. Defaults to True
-* `should_overwrite` - If True, will overwrite existing nodes and relationships. Defaults to False
+* `should_overwrite` - If True, will clear existing constraints, and will overwrite any existing Nodes and Relationships. Defaults to False.
 
 
 ## Data Argument Format

neo4j_uploader/__init__.py

@@ -1,5 +1,6 @@
 from timeit import default_timer as timer
 from neo4j_uploader.n4j import execute_query, reset
+from neo4j_uploader.upload_utils import upload_nodes, upload_relationships
 from neo4j_uploader.logger import ModuleLogger
 import json
 
@@ -18,303 +19,16 @@ def stop_logging():
     ModuleLogger().info(f'Discontinuing logging')
     ModuleLogger().is_enabled = False
 
-def prop_subquery(record: dict, exclude_keys: list[str] = [])-> str:
-
-    # Embed any prop data within brackets { }
-    query = " {"
-
-    filtered_keys = [key for key in record.keys() if key not in exclude_keys]
-    sorted_keys = sorted(list(filtered_keys))
-
-    for idx, a_key in enumerate(sorted_keys):
-        value = record[a_key]
-
-        # Do not set properties with a None/Null/Empty value
-        if value is None:
-            continue
-        if isinstance(value, str):
-            if value.lower() == "none":
-                continue
-            if value.lower() == "null":
-                continue
-            if value.lower() == "empty":
-                continue
-
-        if idx!= 0:
-            query += ", "
-
-        # Using params better but requires creating a separate dictionary be created and passed up to the aggregate function call
-        if isinstance(value, str):
-            escaped_value = value.replace('"','\\"')
-            query += f'`{a_key}`:"{escaped_value}"'
-        else:
-            query += f'`{a_key}`:{value}'
-
-    # Close out query
-    query += "}"
-
-    return query
-
-class HashableDict:
-    def __init__(self, dictionary):
-        self.dictionary = dictionary
-
-    def __hash__(self):
-        return hash(frozenset(self.dictionary.items()))
-    
-def upload_node_records_query(
-    label: str,
-    nodes: list[dict],
-    dedupe : bool = True
-    ):
-    """
-    Generate Cypher Node update query.
-
-    Args:
-        label: Label of Nodes to generate
-
-        nodes: A list of dictionaries representing Node properties
-
-        dedupe: Remove duplicate entries
-    
-    Raises:
-        Exceptions if data is not in the correct format or if the upload fails.
-    """
-        
-    if nodes is None:
-        return None
-    if len(nodes) == 0:
-        return None
-    
-    query = ""
-    count = 0 # Using count to distinguish node variables
-
-    if dedupe == True:
-        nodes = [dict(t) for t in {tuple(n.items()) for n in nodes}]
-
-    for node_record in nodes:
-        count += 1
-        
-        # Add a newline if this is not the first node
-        if count != 1:
-            query += "\n"
-
-        # Convert contents into a subquery specifying node properties
-        subquery = prop_subquery(node_record)
-
-        # Cypher does not support labels with whitespaces
-        query += f"""MERGE (`{label}{count}`:`{label}`{subquery})"""
-
-    return query
-
-def upload_nodes(
-    neo4j_creds:(str, str, str),
-    nodes: dict,
-    dedupe : bool = True
-)-> (int, int):
-    """
-    Uploads a list of dictionary objects as nodes.
-
-    Args:
-        neo4j_credits: Tuple containing the hostname, username, and password of the target Neo4j instance
-
-        nodes: A dictionary of objects to upload. Each key is a unique node label and contains a list of records as dictionary objects.
-
-        dedupe: Remove duplicate entries. Default True
-    
-    Returns:
-        A tuple containing the number of nodes created and properties set.
-
-    Raises:
-        Exceptions if data is not in the correct format or if the upload fails.
-    """
-    if nodes is None:
-        return None
-    if len(nodes) == 0:
-        return None
-    
-    # For reporting
-    nodes_created = 0
-    props_set = 0
-
-    query = """"""
-    expected_count = 0
-    ModuleLogger().debug(f'Uploading node records: {nodes}')
-    for node_label, nodes_list in nodes.items():
-        # Process all similar labeled nodes together
-        query += upload_node_records_query(node_label, nodes_list, dedupe=dedupe)
-        expected_count += len(nodes_list)
-
-    ModuleLogger().debug(f'upload nodes query: {query}')
-
-    records, summary, keys = execute_query(neo4j_creds, query)
-
-    # Sample summary
-    # {'metadata': {'query': '<query>', 'parameters': {}, 'query_type': 'w', 'plan': None, 'profile': None, 'notifications': None, 'counters': {'_contains_updates': True, 'labels_added': 17, 'nodes_created': 17, 'properties_set': 78}, 'result_available_after': 73, 'result_consumed_after': 0}
-    nodes_created += summary.counters.nodes_created
-    props_set += summary.counters.properties_set
-    
-    ModuleLogger().info(f'Results from upload nodes: \n\tRecords: {records}\n\tSummary: {summary.__dict__}\n\tKeys: {keys}')
-    
-    return (nodes_created, props_set)
-
-
-def with_relationship_elements(
-    relationships: list[dict],
-    nodes_key: str = "_uid",
-    dedupe: bool = False
-    ) -> str:
-    """
-    Returns elements to be added into a batch relationship creation query.
-
-    Args:
-        relationships: A list of dictionary records for each relationship property.
-
-        nodes_key: The property key that uniquely identifies Nodes. 
-    
-    Raises:
-        Exceptions if data is not in the correct format or if the upload ungracefully fails.
-    """
-    # TODO: Possible cypher injection entry point?
-
-    result = []
-
-    if dedupe == True:
-        relationships = [dict(t) for t in {tuple(r.items()) for r in relationships}]
-    
-    for rel in relationships:
-
-        # Find from and to node key identifiers
-        from_key = f"_from_{nodes_key}"
-        to_key = f"_to_{nodes_key}"
-
-        # Get unique key of from and to nodes
-        from_node = rel.get(from_key, None)
-        if isinstance(from_node, str):
-            from_node = f"'{from_node}'"
-
-        to_node = rel.get(to_key, None)
-        if isinstance(to_node, str):
-            to_node = f"'{to_node}'"
-
-        # Validate we from and to nodes to work with
-        if from_node is None:
-            ModuleLogger().warning(f'{type} Relationship missing {from_key} property. Skipping relationship {rel}')
-            continue
-        if to_node is None:
-            ModuleLogger().warning(f'{type} Relationship missing {to_key} property. Skipping relationship {rel}')
-            continue
-
-        # string for props
-        props_string = prop_subquery(rel, exclude_keys=[from_key, to_key])
-
-        with_element = f"[{from_node},{to_node},{props_string}]"
-        result.append(with_element)
-    
-    return result
-
-
-
-def upload_relationships(
-    neo4j_creds:(str, str, str),
-    relationships: dict,
-    nodes_key: str = "_uid",
-    dedupe : bool = True
-)-> (int, int):
-    """
-    Uploads a list of dictionary objects as relationships.
-
-    Args:
-        neo4j_credits: Tuple containing the hostname, username, and password of the target Neo4j instance
-
-        nodes: A dictionary of objects to upload. Each key is a unique relationship type and contains a list of records as dictionary objects.
-
-        nodes_key: The property key that uniquely identifies Nodes.
-
-        dedupe: False means a new relationship will always be created for the from and to nodes. True if existing relationships should only be updated. Note that if several relationships already exist, all matching relationships will get their properties updated. Default True.
-    
-    Returns:
-        A tuple of relationships created, properties set
-
-    Raises:
-        Exceptions if data is not in the correct format or if the upload ungracefully fails.
-    """
-
-    # Final query needs to look something like this
-    # WITH [['1202109692044412','1204806322568817', {name:"fish"}],['1202109692044411','test', {}]] AS from_to_id_pairs
-
-    # UNWIND from_to_id_pairs as pair
-
-    # MATCH (fromNode { `gid`: pair[0]})
-    # MATCH (toNode { `gid`: pair[1]})
-
-    # CREATE (fromNode)-[r:`WITHIN`]->(toNode)
-    # SET r = pair[2]
-
-
-    # Validate
-    if relationships is None:
-        return None
-    if len(relationships) ==0:
-        return None
-    
-
-    ModuleLogger().debug(f'upload relationships source data: {relationships}')
-
-    # Upload counts
-    relationships_created = 0
-    props_set = 0
-
-    # Sort so we get a consistent output
-    filtered_keys = [key for key in relationships.keys()]
-    sorted_keys = sorted(list(filtered_keys))
-
-    # NOTE: Need to process each relationship type separately as batching this fails with no warning
-    for idx, rel_type in enumerate(sorted_keys):
-
-        rel_list = relationships[rel_type]
-        ModuleLogger().debug(f'Starting to process relationships type: {rel_type} ...')
-        
-        # Process all similar labeled nodes together
-        with_elements = with_relationship_elements(rel_list, nodes_key, dedupe=dedupe)
-
-        if len(with_elements) is None:
-            ModuleLogger().warning(f'Could not process relationships type {rel_type}. Check if data exsists and matches expected schema')
-            continue
-        
-        with_elements_str = ",".join(with_elements)
-
-        # Assemble final query
-        rel_upload_query = f"""WITH [{with_elements_str}] AS from_to_data\nUNWIND from_to_data AS tuple\nMATCH (fromNode {{`{nodes_key}`:tuple[0]}})\nMATCH (toNode {{`{nodes_key}`:tuple[1]}})"""
-
-        if dedupe == True:
-            rel_upload_query += f"\nMERGE (fromNode)-[r:`{rel_type}`]->(toNode)"
-        else:
-            rel_upload_query += f"\nCREATE (fromNode)-[r:`{rel_type}`]->(toNode)"
-        rel_upload_query +=f"\nSET r = tuple[2]"
-
-        records, summary, keys = execute_query(neo4j_creds, rel_upload_query)
-
-        # Sample summary result
-        # {'metadata': {'query': "<rel_upload_query>", 'parameters': {}, 'query_type': 'w', 'plan': None, 'profile': None, 'notifications': None, 'counters': {'_contains_updates': True, 'relationships_created': 1, 'properties_set': 2}, 'result_available_after': 209, 'result_consumed_after': 0}
-
-        # Can we have optionals yet?
-        relationships_created += summary.counters.relationships_created
-        props_set += summary.counters.properties_set
-
-        ModuleLogger().info(f'Results from uploading relationships type: {rel_type}: \n\tRecords: {records}\n\tSummary: {summary.__dict__}\n\tKeys: {keys}')
-
-    return (relationships_created, props_set)
-
-
 def upload(
-        neo4j_creds:(str, str, str), 
-        data: str | dict,
-        node_key : str = "_uid",
-        dedupe_nodes : bool = True,
-        dedupe_relationships : bool = True,
-        should_overwrite: bool = False
-        )-> (float, int, int, int):
+    neo4j_creds:(str, str, str), 
+    data: str | dict,
+    node_key : str = "_uid",
+    dedupe_nodes : bool = True,
+    dedupe_relationships : bool = True,
+    should_overwrite: bool = False,
+    database_name: str = 'neo4j',
+    max_batch_size: int = 500,
+    )-> (float, int, int, int):
     """
     Uploads a dictionary of records to a target Neo4j instance.
 
@@ -344,6 +58,9 @@ def upload(
         except Exception as e:
             raise Exception(f'Input data string not a valid JSON format: {e}')
 
+    if node_key is None or node_key == "":
+        raise Exception(f'node_key cannot be None or an empty string')
+    
     # Start clock
     start = timer()
 
@@ -355,19 +72,35 @@ def upload(
     if should_overwrite is True:
         reset(neo4j_creds)
 
-    nodes_created, node_props_set = upload_nodes(neo4j_creds, nodes, dedupe=dedupe_nodes)
+    nodes_created, node_props_set = upload_nodes(
+        neo4j_creds, 
+        nodes, 
+        node_key= node_key, 
+        dedupe=dedupe_nodes, 
+        database= database_name,
+        max_batch_size= max_batch_size)
+    
+    all_props_set = node_props_set
     relationships_created = 0,
-    relationship_props_set = 0
 
     # Upload relationship data next
     rels = data.get('relationships', None)
     if rels is not None and len(rels) > 0:
         ModuleLogger().info(f'Begin processing relationships: {rels}')
-        relationships_created, relationship_props_set = upload_relationships(neo4j_creds, rels, node_key, dedupe = dedupe_relationships)
+        relationships_created, relationship_props_set = upload_relationships(
+            neo4j_creds, 
+            rels, 
+            node_key, 
+            dedupe = dedupe_relationships, 
+            database=database_name,
+            max_batch_size=max_batch_size)
+        
+        all_props_set += relationship_props_set
 
-    # TODO: Verify uploads successful
     stop = timer()
     time_to_complete = round((stop - start), 4)
-    all_props_set = node_props_set + relationship_props_set
 
-    return (time_to_complete, nodes_created, relationships_created, all_props_set)
+    return time_to_complete, nodes_created, relationships_created, all_props_set
+
+def clear_db(creds: (str, str, str), database: str):
+    return reset(creds, database)