🔍 Code Extractor

Generates a Python file containing code snippets and helper functions for interacting with a Neo4j graph database based on the provided schema information.

Purpose

This function creates a comprehensive Python module with pre-built functions for common Neo4j operations including CRUD operations for nodes, relationship queries, graph traversals, and database statistics. It analyzes the schema_info dictionary to generate type-specific functions for each node label and relationship type found in the database, making it easier for developers to interact with their specific Neo4j schema without writing Cypher queries from scratch.

Source Code

def generate_python_snippets(schema_info, output_dir):
    """Generate Python code snippets for working with the Neo4j database"""
    snippets_file = os.path.join(output_dir, "neo4j_python_snippets.py")
    limit=25
    
    snippets = """# Neo4j Python Snippets
# Generated from database schema

from neo4j import GraphDatabase

# Connection parameters
URI = "bolt://localhost:7687"
USERNAME = "neo4j"
PASSWORD = "password"

# Connect to Neo4j
driver = GraphDatabase.driver(URI, auth=(USERNAME, PASSWORD))

# Helper function for running queries
def run_query(query, params=None):
    with driver.session() as session:
        result = session.run(query, params or {})
        return [dict(record) for record in result]

"""
    
    # Add node label queries
    snippets += "# Node label queries\n"
    for label in sorted(schema_info["node_labels"]):
        properties = schema_info["node_properties"].get(label, [])
        count = schema_info["node_counts"].get(label, 0)
        
        snippets += f"""
def get_all_{label.lower()}(limit=100):
    \"\"\"Return {label} nodes (limited to {limit})\"\"\"
    query = \"\"\"
    MATCH (n:{label})
    RETURN n
    LIMIT $limit
    \"\"\"
    return run_query(query, {{"limit": limit}})

def get_{label.lower()}_by_id(id):
    \"\"\"Get a {label} node by its ID\"\"\"
    query = \"\"\"
    MATCH (n:{label} {{id: $id}})
    RETURN n
    \"\"\"
    result = run_query(query, {{"id": id}})
    return result[0] if result else None

def get_{label.lower()}_by_uid(uid):
    \"\"\"Get a {label} node by its UID\"\"\"
    query = \"\"\"
    MATCH (n:{label} {{UID: $uid}})
    RETURN n
    \"\"\"
    result = run_query(query, {{"uid": uid}})
    return result[0] if result else None

def create_{label.lower()}(properties):
    \"\"\"Create a new {label} node\"\"\"
    props_list = ', '.join([f"n.{{prop}} = ${{prop}}" for prop in properties.keys()])
    query = f\"\"\"
    CREATE (n:{label})
    SET {{props_list}}
    RETURN n
    \"\"\"
    result = run_query(query, properties)
    return result[0] if result else None

"""
    
    # Add relationship queries
    snippets += "# Relationship queries\n"
    for source, targets in schema_info["node_relationships"].items():
        for target, rels in targets.items():
            for rel in rels:
                rel_type = rel["type"]
                
                snippets += f"""
def get_{source.lower()}_with_{rel_type.lower()}_{target.lower()}(source_id, limit=100):
    \"\"\"Get {target} nodes connected to a {source} via {rel_type}\"\"\"
    query = \"\"\"
    MATCH (source:{source} {{id: $source_id}})-[r:{rel_type}]->(target:{target})
    RETURN target
    LIMIT $limit
    \"\"\"
    return run_query(query, {{"source_id": source_id, "limit": limit}})

def create_{rel_type.lower()}_relationship(source_id, target_id, properties=None):
    \"\"\"Create a {rel_type} relationship from {source} to {target}\"\"\"
    props = ""
    if properties:
        props_list = ', '.join([f"r.{{prop}} = ${{prop}}" for prop in properties.keys()])
        props = f"SET {{props_list}}"
    
    query = f\"\"\"
    MATCH (source:{source} {{id: $source_id}})
    MATCH (target:{target} {{id: $target_id}})
    CREATE (source)-[r:{rel_type}]->(target)
    {{props}}
    RETURN r
    \"\"\"
    
    params = {{"source_id": source_id, "target_id": target_id}}
    if properties:
        params.update(properties)
    
    result = run_query(query, params)
    return result[0] if result else None

"""
    
    # Add graph traversal examples
    snippets += "# Graph traversal examples\n"
    
    # Pick a path through the graph for an example
    if schema_info["node_relationships"]:
        path = []
        visited = set()
        
        # Start with the first node that has a relationship
        for source in schema_info["node_relationships"]:
            if schema_info["node_relationships"][source]:
                current = source
                path.append(current)
                visited.add(current)
                
                # Try to build a path of at least 3 nodes if possible
                while len(path) < 3:
                    next_node = None
                    for target in schema_info["node_relationships"].get(current, {}):
                        if target not in visited:
                            next_node = target
                            path.append(next_node)
                            visited.add(next_node)
                            current = next_node
                            break
                    
                    if not next_node:
                        break
                
                break
        
        # If we found a path, create an example traversal
        if len(path) >= 2:
            path_str = ' -> '.join(path)
            max_depth=5
            
            snippets += f"""
def traverse_path_example(start_id, max_depth=3):
    \"\"\"Example of traversing a path: {path_str}\"\"\"
    query = \"\"\"
    MATCH path = (start:{path[0]} {{id: $start_id}})-[*1..{max_depth}]->(n)
    RETURN path
    LIMIT 10
    \"\"\"
    return run_query(query, {{"start_id": start_id}})

"""
    
    # Add complex query examples
    snippets += "# Complex query examples\n"
    
    # Generate a complex query example if we have enough data
    if len(schema_info["node_labels"]) >= 2 and schema_info["relationship_types"]:
        # Find two node types with a relationship
        example_source = None
        example_target = None
        example_rel = None
        
        for source, targets in schema_info["node_relationships"].items():
            for target, rels in targets.items():
                if rels:
                    example_source = source
                    example_target = target
                    example_rel = rels[0]["type"]
                    break
            if example_source:
                break
        
        if example_source and example_target:
            snippets += f"""
def complex_query_example(search_term, limit=10):
    \"\"\"Example of a more complex query combining multiple node types\"\"\"
    query = \"\"\"
    MATCH (source:{example_source})-[r:{example_rel}]->(target:{example_target})
    WHERE source.name CONTAINS $search_term OR target.name CONTAINS $search_term
    RETURN source, r, target
    LIMIT $limit
    \"\"\"
    return run_query(query, {{"search_term": search_term, "limit": limit}})

"""
    
    # Add statistics example
    snippets += """
def get_database_statistics():
    \"\"\"Get general statistics about the database\"\"\"
    stats = {}
    
    # Node counts by label
    query = \"\"\"
    MATCH (n)
    RETURN labels(n) AS label, count(*) AS count
    \"\"\"
    label_counts = run_query(query)
    stats["node_counts"] = {record["label"][0]: record["count"] for record in label_counts if record["label"]}
    
    # Relationship counts by type
    query = \"\"\"
    MATCH ()-[r]->()
    RETURN type(r) AS type, count(*) AS count
    \"\"\"
    rel_counts = run_query(query)
    stats["relationship_counts"] = {record["type"]: record["count"] for record in rel_counts}
    
    return stats

# Don't forget to close the driver when done
# driver.close()
"""
    
    with open(snippets_file, "w") as f:
        f.write(snippets)

Parameters

Parameter Details

schema_info: A dictionary containing Neo4j database schema information with keys: 'node_labels' (list of node label names), 'node_properties' (dict mapping labels to property lists), 'node_counts' (dict mapping labels to node counts), 'node_relationships' (nested dict mapping source labels to target labels to relationship details), and 'relationship_types' (list of relationship type names). This structure is typically generated by a schema introspection function.

output_dir: String path to the directory where the generated Python snippets file should be written. The function will create a file named 'neo4j_python_snippets.py' in this directory. The directory must exist or be creatable.

Return Value

This function returns None (implicitly). Its primary effect is the side effect of writing a Python file to disk at the path '{output_dir}/neo4j_python_snippets.py' containing generated code snippets.

Dependencies

• os
• neo4j

Required Imports

import os

Usage Example

import os

# Prepare schema information (typically from a schema introspection function)
schema_info = {
    'node_labels': ['Person', 'Company'],
    'node_properties': {
        'Person': ['id', 'name', 'age'],
        'Company': ['id', 'name', 'industry']
    },
    'node_counts': {
        'Person': 150,
        'Company': 50
    },
    'node_relationships': {
        'Person': {
            'Company': [
                {'type': 'WORKS_AT', 'count': 120}
            ]
        }
    },
    'relationship_types': ['WORKS_AT']
}

# Generate snippets
output_directory = './generated_code'
os.makedirs(output_directory, exist_ok=True)
generate_python_snippets(schema_info, output_directory)

# The generated file can then be imported and used:
# from generated_code.neo4j_python_snippets import get_all_person, create_person
# persons = get_all_person(limit=10)

Best Practices

• Ensure the schema_info dictionary is complete and accurate before calling this function, as it drives all code generation
• Review and customize the generated connection parameters (URI, USERNAME, PASSWORD) in the output file before using it in production
• The generated file includes a driver connection that should be closed when done (driver.close())
• Generated functions use parameterized queries to prevent Cypher injection attacks
• The limit parameter in generated functions defaults to 100 to prevent accidentally loading large datasets
• Consider the generated file as a starting template that may need customization for specific use cases
• The function assumes node properties include 'id' and 'UID' fields for lookup functions
• Generated code uses context managers (with statements) for proper session management

Tags

Similar Components

AI-powered semantic similarity - components with related functionality:

• function generate_cypher_examples 70.8% similar

  Generates a comprehensive Cypher query examples file for interacting with a Neo4j graph database based on the provided schema information.

  From: /tf/active/vicechatdev/neo4j_schema_report.py

• function generate_neo4j_schema_report 64.9% similar

  Generates a comprehensive schema report of a Neo4j graph database, including node labels, relationships, properties, constraints, indexes, and sample data, outputting multiple report formats (JSON, HTML, Python snippets, Cypher examples).

  From: /tf/active/vicechatdev/neo4j_schema_report.py

• function generate_diagram_data 59.1% similar

  Transforms Neo4j schema information into a structured format suitable for graph visualization, creating separate node and edge data structures.

  From: /tf/active/vicechatdev/neo4j_schema_report.py

• function add_document_to_graph 49.0% similar

  Creates nodes and relationships in a Neo4j graph database for a processed document, including its text and table chunks, connecting it to a folder hierarchy.

  From: /tf/active/vicechatdev/offline_docstore_multi.py

• function generate_html_report 48.4% similar

  Generate HTML report from schema info

  From: /tf/active/vicechatdev/neo4j_schema_report.py