diff --git a/.gitignore b/.gitignore index 95e23038..89ceb3f1 100644 --- a/.gitignore +++ b/.gitignore @@ -61,6 +61,7 @@ wheels/ .installed.cfg *.egg MANIFEST +.python-version # IDE .vscode/ diff --git a/README.md b/README.md index 753d9395..687d7611 100644 --- a/README.md +++ b/README.md @@ -360,7 +360,7 @@ results = vector_store.search(query="supply chain", top_k=5) ### Graph Store & Triplet Store -> **Neo4j, FalkorDB support** • **SPARQL queries** • **RDF triplets** +> **Neo4j, FalkorDB, Amazon Neptune support** • **SPARQL queries** • **RDF triplets** ```python from semantica.graph_store import GraphStore @@ -370,6 +370,24 @@ from semantica.triplet_store import TripletStore graph_store = GraphStore(backend="neo4j", uri="bolt://localhost:7687", user="neo4j", password="password") graph_store.add_nodes([{"id": "n1", "labels": ["Person"], "properties": {"name": "Alice"}}]) +# Amazon Neptune Graph Store (OpenCypher via HTTP with IAM Auth) +neptune_store = GraphStore( + backend="neptune", + endpoint="your-cluster.us-east-1.neptune.amazonaws.com", + port=8182, + region="us-east-1", + iam_auth=True, # Uses AWS credential chain (boto3, env vars, or IAM role) +) + +# Node Operations +neptune_store.add_nodes([ + {"labels": ["Person"], "properties": {"id": "alice", "name": "Alice", "age": 30}}, + {"labels": ["Person"], "properties": {"id": "bob", "name": "Bob", "age": 25}}, +]) + +# Query Operations +result = neptune_store.execute_query("MATCH (p:Person) RETURN p.name, p.age") + # Triplet Store (Blazegraph, Jena, RDF4J) triplet_store = TripletStore(backend="blazegraph", endpoint="http://localhost:9999/blazegraph") triplet_store.add_triplet({"subject": "Alice", "predicate": "knows", "object": "Bob"}) diff --git a/cookbook/introduction/21_Amazon_Neptune_Store.ipynb b/cookbook/introduction/21_Amazon_Neptune_Store.ipynb new file mode 100644 index 00000000..03ae6971 --- /dev/null +++ b/cookbook/introduction/21_Amazon_Neptune_Store.ipynb @@ -0,0 +1,667 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Amazon Neptune Graph Store\n", + "\n", + "## Overview\n", + "\n", + "This notebook covers the Amazon Neptune Database integration in Semantica. Amazon Neptune is a fully managed graph database service that supports both property graphs (via OpenCypher/Gremlin) and RDF graphs (via SPARQL).\n", + "\n", + "### Key Features\n", + "\n", + "- **IAM Authentication**: Secure access using AWS SigV4 signatures via AuthManager\n", + "- **OpenCypher Support**: Query using standard OpenCypher syntax\n", + "- **Bolt Protocol**: Uses Neo4j Bolt driver for efficient binary communication\n", + "- **Native ~id Support**: Leverages Neptune's native element ID handling\n", + "- **Full CRUD Operations**: Create, read, update, delete nodes and relationships\n", + "- **Automatic Retry**: Built-in retry logic with exponential backoff for transient errors\n", + "\n", + "### Prerequisites\n", + "\n", + "- An Amazon Neptune Database cluster\n", + "- AWS credentials configured (boto3, environment variables, or IAM role)\n", + "- Network access to your Neptune cluster (VPC, security groups)\n", + "\n", + "---" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Installation\n", + "\n", + "```bash\n", + "# Install Semantica with Neptune support\n", + "pip install semantica\n", + "\n", + "# Required dependencies (installed automatically)\n", + "pip install boto3 neo4j\n", + "```" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "!pip install semantica" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Configuration\n", + "\n", + "Set your Neptune cluster endpoint and AWS credentials. Replace the placeholder values with your actual configuration." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import os\n", + "\n", + "# Neptune cluster configuration - REPLACE WITH YOUR VALUES\n", + "os.environ[\"NEPTUNE_ENDPOINT\"] = \"your-cluster.us-east-1.neptune.amazonaws.com\"\n", + "os.environ[\"NEPTUNE_PORT\"] = \"8182\"\n", + "os.environ[\"AWS_REGION\"] = \"us-east-1\"\n", + "\n", + "# AWS credentials (if using IAM Auth and not relying on IAM role or ~/.aws/credentials)\n", + "# os.environ[\"AWS_ACCESS_KEY_ID\"] = \"your-access-key-id\"\n", + "# os.environ[\"AWS_SECRET_ACCESS_KEY\"] = \"your-secret-access-key\"\n", + "# os.environ[\"AWS_SESSION_TOKEN\"] = \"your-session-token\"\n", + "\n", + "print(f\"Neptune Endpoint: {os.environ.get('NEPTUNE_ENDPOINT')}\")\n", + "print(f\"AWS Region: {os.environ.get('AWS_REGION')}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Step 1: Initialize Neptune Store\n", + "\n", + "Initialize a connection to your Amazon Neptune cluster with IAM authentication." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import os\n", + "from semantica.graph_store import GraphStore\n", + "\n", + "# Option 1: Using GraphStore factory (recommended)\n", + "neptune_store = GraphStore(\n", + " backend=\"neptune\",\n", + " endpoint=os.environ.get(\"NEPTUNE_ENDPOINT\"),\n", + " port=int(os.environ.get(\"NEPTUNE_PORT\", 8182)),\n", + " region=os.environ.get(\"AWS_REGION\", \"us-east-1\"),\n", + " iam_auth=True,\n", + ")\n", + "\n", + "# Connect to Neptune\n", + "neptune_store.connect()\n", + "print(\"Connected to Amazon Neptune!\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Development/Testing Without IAM Auth\n", + "\n", + "For development or testing environments where IAM authentication is not required (e.g., Neptune notebooks or VPC-only access), you can disable IAM signing:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# For dev/test environments without IAM authentication\n", + "neptune_store_dev = GraphStore(\n", + " backend=\"neptune\",\n", + " endpoint=os.environ.get(\"NEPTUNE_ENDPOINT\"),\n", + " port=int(os.environ.get(\"NEPTUNE_PORT\", 8182)),\n", + " region=os.environ.get(\"AWS_REGION\", \"us-east-1\"),\n", + " iam_auth=False, # Disable IAM signing for dev/test\n", + ")\n", + "neptune_store_dev.connect()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Authentication Options\n", + "\n", + "IAM Authentication (recommended for production) automatically uses the AWS credential chain:\n", + "1. Environment variables (AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY)\n", + "2. AWS credentials file (~/.aws/credentials)\n", + "3. IAM role (for EC2, Lambda, ECS)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Step 2: Node Operations\n", + "\n", + "### Creating Nodes\n", + "\n", + "Nodes represent entities in your graph. Each node can have:\n", + "- **ID**: A unique identifier (custom or auto-generated UUID)\n", + "- **Labels**: Categories/types (e.g., `Person`, `Company`)\n", + "- **Properties**: Key-value pairs (e.g., `{\"name\": \"Alice\", \"age\": 30}`)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Create a single node with custom ID (id in properties)\n", + "alice = neptune_store.create_node(\n", + " labels=[\"Person\"],\n", + " properties={\"id\": \"alice\", \"name\": \"Alice\", \"age\": 30, \"role\": \"Engineer\"}\n", + ")\n", + "print(f\"Created node: {alice}\")\n", + "\n", + "# Create a node with auto-generated UUID (no id in properties)\n", + "bob = neptune_store.create_node(\n", + " labels=[\"Person\"],\n", + " properties={\"name\": \"Bob\", \"age\": 25, \"role\": \"Designer\"}\n", + ")\n", + "print(f\"Created node with UUID: {bob['id']}\")\n", + "\n", + "# Create a company node with auto-generated ID\n", + "acme = neptune_store.create_node(\n", + " labels=[\"Company\"],\n", + " properties={\"name\": \"Acme Corp\", \"industry\": \"Technology\", \"founded\": 2010}\n", + ")\n", + "print(f\"Created company: {acme}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Creating Multiple Nodes (Batch)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Batch create nodes for better performance\n", + "# Include 'id' in properties for custom IDs\n", + "nodes_data = [\n", + " {\"labels\": [\"Person\"], \"properties\": {\"id\": \"charlie\", \"name\": \"Charlie\", \"age\": 35}},\n", + " {\"labels\": [\"Person\"], \"properties\": {\"id\": \"diana\", \"name\": \"Diana\", \"age\": 28}},\n", + " {\"labels\": [\"Location\"], \"properties\": {\"name\": \"San Francisco\", \"state\": \"CA\"}},\n", + "]\n", + "\n", + "created_nodes = neptune_store.create_nodes(nodes_data)\n", + "print(f\"Created {len(created_nodes)} nodes in batch\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Retrieving Nodes" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Get a specific node by ID\n", + "alice_node = neptune_store.get_node(node_id=\"alice\")\n", + "print(f\"Retrieved: {alice_node}\")\n", + "\n", + "# Get nodes by label\n", + "people = neptune_store.get_nodes(labels=[\"Person\"], limit=10)\n", + "print(f\"Found {len(people)} Person nodes:\")\n", + "for person in people:\n", + " print(f\" - {person.get('properties', {}).get('name')}\")\n", + "\n", + "# Get nodes by properties\n", + "engineers = neptune_store.get_nodes(\n", + " labels=[\"Person\"],\n", + " properties={\"role\": \"Engineer\"},\n", + " limit=5\n", + ")\n", + "print(f\"Found {len(engineers)} engineers\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Updating Nodes" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Update node properties (merge mode - default)\n", + "updated_alice = neptune_store.update_node(\n", + " node_id=\"alice\",\n", + " properties={\"age\": 31, \"department\": \"AI Research\"},\n", + " merge=True\n", + ")\n", + "print(f\"Updated Alice: {updated_alice}\")\n", + "\n", + "# Replace all properties (merge=False)\n", + "# WARNING: This removes properties not in the update\n", + "replaced = neptune_store.update_node(\n", + " node_id=\"charlie\",\n", + " properties={\"name\": \"Charlie\", \"age\": 36},\n", + " merge=False\n", + ")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Deleting Nodes" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Delete a node (with detach=True to also delete relationships)\n", + "deleted = neptune_store.delete_node(node_id=\"diana\", detach=True)\n", + "print(f\"Deleted diana: {deleted}\")\n", + "\n", + "# Without detach (fails if node has relationships)\n", + "# neptune_store.delete_node(node_id=\"alice\", detach=False)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Step 3: Relationship Operations\n", + "\n", + "### Creating Relationships\n", + "\n", + "Relationships connect nodes and represent connections between entities." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Create a relationship between Alice and Acme\n", + "works_at = neptune_store.create_relationship(\n", + " start_node_id=\"alice\",\n", + " end_node_id=acme[\"id\"],\n", + " rel_type=\"WORKS_AT\",\n", + " properties={\"since\": 2020, \"position\": \"Senior Engineer\"}\n", + ")\n", + "print(f\"Created relationship: {works_at}\")\n", + "\n", + "# Create a KNOWS relationship between people\n", + "knows_rel = neptune_store.create_relationship(\n", + " start_node_id=\"alice\",\n", + " end_node_id=bob[\"id\"],\n", + " rel_type=\"KNOWS\",\n", + " properties={\"since\": 2019}\n", + ")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Retrieving Relationships" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Get all relationships for a node\n", + "alice_rels = neptune_store.get_relationships(node_id=\"alice\", direction=\"both\")\n", + "print(f\"Alice has {len(alice_rels)} relationships\")\n", + "\n", + "# Get outgoing relationships only\n", + "outgoing = neptune_store.get_relationships(node_id=\"alice\", direction=\"out\")\n", + "\n", + "# Filter by relationship type\n", + "works_rels = neptune_store.get_relationships(\n", + " node_id=\"alice\",\n", + " rel_type=\"WORKS_AT\",\n", + " direction=\"out\"\n", + ")\n", + "print(f\"Alice's work relationships: {len(works_rels)}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Deleting Relationships" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Delete a specific relationship by ID\n", + "if works_at.get(\"id\"):\n", + " deleted = neptune_store.delete_relationship(rel_id=works_at[\"id\"])\n", + " print(f\"Deleted relationship: {deleted}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Step 4: OpenCypher Queries\n", + "\n", + "Amazon Neptune supports OpenCypher queries via the Bolt protocol. Execute complex graph patterns using standard Cypher syntax." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Simple query\n", + "results = neptune_store.execute_query(\n", + " \"MATCH (p:Person) RETURN p.name, p.age ORDER BY p.age\"\n", + ")\n", + "print(\"People in the graph:\")\n", + "for record in results.get(\"records\", []):\n", + " print(f\" - {record.get('p.name')}: {record.get('p.age')} years old\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Using parameters (safer and more efficient)\n", + "results = neptune_store.execute_query(\n", + " \"MATCH (p:Person) WHERE p.age > $min_age RETURN p.name, p.age\",\n", + " parameters={\"min_age\": 25}\n", + ")\n", + "print(f\"People over 25: {len(results.get('records', []))}\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Find relationships between nodes\n", + "results = neptune_store.execute_query(\"\"\"\n", + " MATCH (p:Person)-[r:WORKS_AT]->(c:Company)\n", + " RETURN p.name as employee, c.name as company, r.since as start_year\n", + "\"\"\")\n", + "for record in results.get(\"records\", []):\n", + " print(f\"{record['employee']} works at {record['company']} since {record['start_year']}\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Count and aggregate\n", + "results = neptune_store.execute_query(\"\"\"\n", + " MATCH (p:Person)\n", + " RETURN count(p) as total, avg(p.age) as avg_age, max(p.age) as max_age\n", + "\"\"\")\n", + "stats = results.get(\"records\", [{}])[0]\n", + "print(f\"Total: {stats.get('total')}, Avg Age: {stats.get('avg_age'):.1f}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Step 5: Graph Analytics\n", + "\n", + "### Get Neighbors\n", + "\n", + "Traverse the graph to find connected nodes." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Get immediate neighbors (depth=1)\n", + "neighbors = neptune_store.get_neighbors(\n", + " node_id=\"alice\",\n", + " direction=\"both\",\n", + " depth=1\n", + ")\n", + "print(f\"Alice's direct neighbors: {len(neighbors)}\")\n", + "\n", + "# Get neighbors up to 2 hops away\n", + "extended = neptune_store.get_neighbors(\n", + " node_id=\"alice\",\n", + " direction=\"out\",\n", + " depth=2\n", + ")\n", + "print(f\"Nodes within 2 hops: {len(extended)}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Shortest Path\n", + "\n", + "Find the shortest path between two nodes." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Find shortest path\n", + "path = neptune_store.shortest_path(\n", + " start_node_id=\"alice\",\n", + " end_node_id=\"charlie\",\n", + " max_depth=5\n", + ")\n", + "\n", + "if path:\n", + " print(\"Path found!\")\n", + " print(f\" Length: {path.get('length')}\")\n", + " print(f\" Nodes: {len(path.get('nodes', []))}\")\n", + " print(f\" Relationships: {len(path.get('relationships', []))}\")\n", + "else:\n", + " print(\"No path found between nodes\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Step 6: Graph Statistics\n", + "\n", + "Get comprehensive statistics about your graph." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Get graph statistics\n", + "stats = neptune_store.get_stats()\n", + "\n", + "print(\"Graph Statistics:\")\n", + "print(f\" Total nodes: {stats.get('node_count', 'N/A')}\")\n", + "print(f\" Total relationships: {stats.get('relationship_count', 'N/A')}\")\n", + "\n", + "print(\"\\nNode labels:\")\n", + "for label, count in stats.get('label_counts', {}).items():\n", + " print(f\" - {label}: {count}\")\n", + "\n", + "print(\"\\nRelationship types:\")\n", + "for rel_type, count in stats.get('relationship_type_counts', {}).items():\n", + " print(f\" - {rel_type}: {count}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Step 7: Connection Management\n", + "\n", + "Always close connections when done to free resources." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Check connection status\n", + "status = neptune_store.get_status()\n", + "print(f\"Connection status: {status}\")\n", + "\n", + "# Close the connection\n", + "neptune_store.close()\n", + "print(\"Connection closed\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Neptune-Specific Considerations\n", + "\n", + "### Native Element IDs\n", + "\n", + "Neptune uses native `~id` for element identification. Include `id` in properties to set a custom ID:\n", + "\n", + "```python\n", + "# Create a node with custom ID (include 'id' in properties)\n", + "node = neptune_store.create_node(\n", + " labels=[\"Person\"],\n", + " properties={\"id\": \"my-custom-id\", \"name\": \"Test\"}\n", + ")\n", + "\n", + "# Create a node with auto-generated UUID (omit 'id' from properties)\n", + "node = neptune_store.create_node(\n", + " labels=[\"Person\"],\n", + " properties={\"name\": \"Test\"}\n", + ")\n", + "\n", + "# The ID is used in id() function calls internally:\n", + "# MATCH (n) WHERE id(n) = 'my-custom-id' RETURN n\n", + "```\n", + "\n", + "### OpenCypher Considerations\n", + "\n", + "Amazon Neptune Database's OpenCypher implementation has some differences from Neo4j:\n", + "\n", + "1. **No `shortestPath()` function**: Use variable-length path patterns or `allShortestPaths()`\n", + "2. **Labels syntax**: Use `labels(n)` function to retrieve node labels\n", + "3. **Property updates**: Use `SET n += {props}` for merge behavior\n", + "\n", + "For the complete OpenCypher specification supported by Amazon Neptune Database, see the [AWS documentation](https://docs.aws.amazon.com/neptune/latest/userguide/access-graph-opencypher.html).\n", + "\n", + "### Amazon Neptune Analytics\n", + "\n", + "For analytical (OLAP) workloads such as graph algorithms, aggregations, and large-scale traversals, consider [Amazon Neptune Analytics](https://docs.aws.amazon.com/neptune-analytics/latest/userguide/what-is-neptune-analytics.html). Neptune Analytics complements Neptune Database by providing optimized performance for analytical queries while Neptune Database is optimized for transactional (OLTP) workloads.\n", + "\n", + "### Performance Tips\n", + "\n", + "1. **Use batch operations** for creating multiple nodes/relationships\n", + "2. **Use parameters** in queries to enable query caching\n", + "3. **Limit result sets** with `LIMIT` clause" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Summary\n", + "\n", + "This notebook covered the Amazon Neptune Graph Store integration:\n", + "\n", + "- **IAM Authentication**: Secure AWS SigV4 signing\n", + "- **CRUD Operations**: Full node and relationship management\n", + "- **OpenCypher Queries**: Standard graph query language\n", + "- **Graph Analytics**: Neighbors and shortest path algorithms\n", + "- **Statistics & Monitoring**: Graph metrics and status\n", + "\n", + "### Key Takeaways\n", + "\n", + "- Neptune uses native `~id` for element identification\n", + "- IAM authentication is recommended for production\n", + "- Bolt protocol provides efficient binary query interface\n", + "- Semantica abstracts Neptune-specific syntax differences\n", + "\n", + "### Next Steps\n", + "\n", + "- [Graph Store (Neo4j/FalkorDB)](09_Graph_Store.ipynb) - Compare with other backends\n", + "- [Building Knowledge Graphs](07_Building_Knowledge_Graphs.ipynb) - Build production KGs\n", + "- [Graph Analytics](10_Graph_Analytics.ipynb) - Advanced analytics algorithms" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "name": "python", + "version": "3.9.0" + } + }, + "nbformat": 4, + "nbformat_minor": 4 +} diff --git a/pyproject.toml b/pyproject.toml index a488ad2e..428979e8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -209,8 +209,12 @@ graph-falkordb = [ "falkordb>=1.0.0", "redis>=4.3.0" ] +graph-amazon-neptune = [ + "boto3>=1.24.0", + "neo4j>=5.0.0" +] graph-all = [ - "semantica[graph-neo4j,graph-falkordb]" + "semantica[graph-neo4j,graph-falkordb,graph-amazon-neptune]" ] parse-docling = [ "docling>=1.0.0" diff --git a/semantica/graph_store/__init__.py b/semantica/graph_store/__init__.py index b8e471b4..ea952b2e 100644 --- a/semantica/graph_store/__init__.py +++ b/semantica/graph_store/__init__.py @@ -1,48 +1,71 @@ """ Graph Store Module -This module provides comprehensive property graph database integration for the -Semantica framework, supporting multiple graph database backends including Neo4j -and FalkorDB for storing and querying knowledge graphs. +This module provides comprehensive property graph database integration for +the Semantica framework, supporting multiple graph database backends including +Neo4j and FalkorDB for storing and querying knowledge graphs. Algorithms Used: Graph Store Management: - - Store Registration: Store type detection, store factory pattern, configuration management, default store selection - - Backend Pattern: Unified interface for multiple backends (Neo4j, FalkorDB), backend instantiation, backend-specific operation delegation - - Store Selection: Default store resolution, store ID lookup, store validation + - Store Registration: Store type detection, store factory pattern, + configuration management, default store selection + - Backend Pattern: Unified interface for multiple backends (Neo4j, + FalkorDB), backend instantiation, backend-specific operation delegation + - Store Selection: Default store resolution, store ID lookup, + store validation Node and Relationship Operations: - - Node Creation: Single node insertion, batch node insertion, property validation, label management, backend delegation - - Node Retrieval: Pattern matching (label/property filtering), Cypher query construction, result extraction, node reconstruction - - Node Update: Property update, label modification, atomic update operations, conflict detection - - Node Deletion: Node matching, cascade deletion (optional), deletion operation delegation, result verification - - Relationship Creation: Single relationship insertion, batch insertion, property validation, type management + - Node Creation: Single node insertion, batch node insertion, + property validation, label management, backend delegation + - Node Retrieval: Pattern matching (label/property filtering), + Cypher query construction, result extraction, node reconstruction + - Node Update: Property update, label modification, atomic update + operations, conflict detection + - Node Deletion: Node matching, cascade deletion (optional), + deletion operation delegation, result verification + - Relationship Creation: Single relationship insertion, batch insertion, + property validation, type management - Relationship Retrieval: Pattern matching, path queries, traversal queries - Relationship Update: Property update, type modification - - Relationship Deletion: Relationship matching, deletion operation delegation + - Relationship Deletion: Relationship matching, deletion operation + delegation Graph Query Execution: - - Cypher Query: Full Cypher query language support for Neo4j and FalkorDB (OpenCypher) - - Pattern Matching: Node and relationship pattern matching, variable binding, path matching - - Graph Traversal: BFS/DFS traversal, shortest path algorithms, path finding + - Cypher Query: Full Cypher query language support for Neo4j and + FalkorDB (OpenCypher) + - Pattern Matching: Node and relationship pattern matching, variable + binding, path matching + - Graph Traversal: BFS/DFS traversal, shortest path algorithms, + path finding - Aggregation: COUNT, SUM, AVG, MIN, MAX operations, GROUP BY support - - Query Optimization: Query caching, execution plan analysis, index utilization + - Query Optimization: Query caching, execution plan analysis, + index utilization Graph Analytics: - - Centrality Algorithms: Degree centrality, betweenness centrality, PageRank, closeness centrality - - Community Detection: Label propagation, Louvain modularity, connected components - - Path Algorithms: Shortest path, all shortest paths, Dijkstra, A* pathfinding + - Centrality Algorithms: Degree centrality, betweenness centrality, + PageRank, closeness centrality + - Community Detection: Label propagation, Louvain modularity, + connected components + - Path Algorithms: Shortest path, all shortest paths, Dijkstra, + A* pathfinding - Similarity: Node similarity, Jaccard similarity, cosine similarity Store Backends: - - Neo4j Store: Official Neo4j Python driver, Bolt protocol communication, transaction support, multi-database support, APOC procedures - - FalkorDB Store: Redis-based graph database, sparse matrix representation, linear algebra queries, OpenCypher support, ultra-fast performance + - Neo4j Store: Official Neo4j Python driver, Bolt protocol + communication, transaction support, multi-database support, + APOC procedures + - FalkorDB Store: Redis-based graph database, sparse matrix + representation, linear algebra queries, OpenCypher support, + ultra-fast performance Bulk Operations: - - Batch Processing: Chunking algorithm (fixed-size batch creation), batch size optimization, memory management for large datasets - - Transaction Management: ACID transaction support, batch commits, rollback on failure - - Progress Tracking: Load progress calculation, elapsed time tracking, throughput calculation + - Batch Processing: Chunking algorithm (fixed-size batch creation), + batch size optimization, memory management for large datasets + - Transaction Management: ACID transaction support, batch commits, + rollback on failure + - Progress Tracking: Load progress calculation, elapsed time tracking, + throughput calculation Key Features: - Multi-backend property graph support (Neo4j, FalkorDB) @@ -79,33 +102,42 @@ Convenience Functions: - list_available_methods: List registered graph store methods Example Usage: - >>> from semantica.graph_store import GraphStore, create_node, create_relationship, execute_query + >>> from semantica.graph_store import GraphStore, create_node, \ + ... create_relationship, execute_query >>> # Using convenience functions - >>> node_id = create_node(labels=["Person"], properties={"name": "Alice", "age": 30}) - >>> rel_id = create_relationship(start_id=node1_id, end_id=node2_id, rel_type="KNOWS", properties={"since": 2020}) - >>> results = execute_query("MATCH (p:Person) WHERE p.age > 25 RETURN p.name") + >>> node_id = create_node(labels=["Person"], + ... properties={"name": "Alice", "age": 30}) + >>> rel_id = create_relationship(start_id=node1_id, end_id=node2_id, + ... rel_type="KNOWS", + ... properties={"since": 2020}) + >>> results = execute_query("MATCH (p:Person) WHERE p.age > 25 " + ... "RETURN p.name") >>> # Using classes directly >>> store = GraphStore(backend="neo4j", uri="bolt://localhost:7687") - >>> node_id = store.create_node(labels=["Person"], properties={"name": "Bob"}) + >>> node_id = store.create_node(labels=["Person"], + ... properties={"name": "Bob"}) >>> results = store.execute_query("MATCH (n) RETURN n LIMIT 10") Author: Semantica Contributors License: MIT """ -from .config import GraphStoreConfig, graph_store_config -from .falkordb_store import ( - FalkorDBStore, - FalkorDBClient, - FalkorDBGraph, +from .amazon_neptune import ( + AmazonNeptuneStore, + NeptuneAuthTokenManager, + NeptuneDriver, + NeptuneSession, + NeptuneTransaction, ) +from .config import GraphStoreConfig, graph_store_config +from .falkordb_store import FalkorDBClient, FalkorDBGraph, FalkorDBStore from .graph_store import ( + GraphAnalytics, GraphManager, GraphStore, NodeManager, QueryEngine, RelationshipManager, - GraphAnalytics, ) from .methods import ( create_node, @@ -125,11 +157,7 @@ from .methods import ( update_node, update_relationship, ) -from .neo4j_store import ( - Neo4jStore, - Neo4jDriver, - Neo4jTransaction, -) +from .neo4j_store import Neo4jDriver, Neo4jStore, Neo4jTransaction from .registry import MethodRegistry, method_registry __all__ = [ @@ -144,6 +172,12 @@ __all__ = [ "Neo4jStore", "Neo4jDriver", "Neo4jTransaction", + # Amazon Neptune + "AmazonNeptuneStore", + "NeptuneAuthTokenManager", + "NeptuneDriver", + "NeptuneSession", + "NeptuneTransaction", # FalkorDB "FalkorDBStore", "FalkorDBClient", @@ -171,4 +205,3 @@ __all__ = [ "MethodRegistry", "method_registry", ] - diff --git a/semantica/graph_store/amazon_neptune.py b/semantica/graph_store/amazon_neptune.py new file mode 100644 index 00000000..988fc470 --- /dev/null +++ b/semantica/graph_store/amazon_neptune.py @@ -0,0 +1,1798 @@ +""" +Amazon Neptune Store Module + +This module provides Amazon Neptune graph database integration with IAM authentication +using the Neo4j Bolt driver for property graph storage and OpenCypher querying +in the Semantica framework. + +Key Features: + - IAM authentication using AWS SigV4 signing via AuthManager + - OpenCypher query language support via Bolt protocol + - Node and relationship CRUD operations + - Graph analytics + - Batch operations with progress tracking + - Automatic retry with backoff for transient errors + - Connection recovery and token refresh + +Main Classes: + - AmazonNeptuneStore: Main Neptune store for graph operations + - NeptuneAuthTokenManager: IAM authentication handler using AuthManager interface + +Example Usage: + >>> from semantica.graph_store import AmazonNeptuneStore + >>> store = AmazonNeptuneStore( + ... endpoint="your-neptune-cluster.region.neptune.amazonaws.com", + ... port=8182, + ... region="us-east-1", + ... iam_auth=True + ... ) + >>> store.connect() + >>> node_id = store.create_node(labels=["Person"], properties={"name": "Alice"}) + >>> results = store.execute_query("MATCH (p:Person) RETURN p.name") + >>> store.close() + +Note: Amazon Neptune uses the Bolt protocol for OpenCypher queries. + The endpoint is: bolt://:8182 + +Author: Semantica Contributors +License: MIT +""" + +import json +import os +import sys +import uuid +from typing import Any, Dict, List, Optional, Union + +from ..utils.exceptions import ProcessingError, ValidationError +from ..utils.logging import get_logger +from ..utils.progress_tracker import get_progress_tracker + +# Optional boto3 for AWS credentials and SigV4 signing +try: + from types import SimpleNamespace + + import boto3 + from botocore.auth import SigV4Auth, _host_from_url + from botocore.awsrequest import AWSRequest + + BOTO3_AVAILABLE = True +except (ImportError, OSError): + BOTO3_AVAILABLE = False + boto3 = None + SigV4Auth = None + AWSRequest = None + SimpleNamespace = None + _host_from_url = None + +# Optional Neo4j driver for Bolt protocol +try: + from neo4j import GraphDatabase + from neo4j.api import basic_auth + from neo4j.auth_management import AuthManager + from neo4j.exceptions import ClientError, DatabaseError, ServiceUnavailable + + NEO4J_AVAILABLE = True +except (ImportError, OSError): + NEO4J_AVAILABLE = False + GraphDatabase = None + AuthManager = None + basic_auth = None + ServiceUnavailable = Exception + DatabaseError = Exception + ClientError = Exception + +# Optional backoff for retry logic +try: + import backoff + + BACKOFF_AVAILABLE = True +except (ImportError, OSError): + BACKOFF_AVAILABLE = False + backoff = None + + +# Retry configuration +NUM_RETRIES = 3 +RETRIABLE_ERROR_MESSAGES = [ + "Signature expired", + "Invalid authentication parameters", + "Operation terminated (out of memory)", + "Operation terminated (deadline exceeded)", + "Operation terminated (cancelled by user)", + "Database reset is in progress", + "Operation failed due to conflicting concurrent operations", + "Max number of request have breached", + "Max connection limit breached", + "Operation terminated (internal error)", + "Connection is closed", +] +NETWORK_ERRORS = (OSError,) +if NEO4J_AVAILABLE: + NETWORK_ERRORS = (OSError, ClientError) + + +# Module-level logger for NeptuneAuthTokenManager (kept outside class to avoid +# serialization issues when Neo4j driver inspects the AuthManager object) +_auth_logger = get_logger("neptune_auth_token_manager") + + +# Use AuthManager base class if available, otherwise just use object +_AuthManagerBase = AuthManager if NEO4J_AVAILABLE and AuthManager else object + + +class NeptuneAuthTokenManager(_AuthManagerBase): + """ + Custom AuthManager for Amazon Neptune using SigV4 signing. + Compatible with Neo4j Python Driver 5.x+ + + This class implements the AuthManager interface to provide automatic + token refresh and handling of security exceptions for Neptune IAM authentication. + + Note: This class intentionally does NOT store a logger as an instance + attribute because the Neo4j driver inspects AuthManager objects and + cannot serialize Logger types. + """ + + # Constants for SigV4 authentication + SCHEME = "basic" + REALM = "realm" + SERVICE_NAME = "neptune-db" + HTTP_METHOD_HDR = "HttpMethod" + DUMMY_USERNAME = "username" + AUTHORIZATION = "Authorization" + HOST = "Host" + X_AMZ_DATE = "X-Amz-Date" + X_AMZ_SECURITY_TOKEN = "X-Amz-Security-Token" + + def __init__( + self, + neptune_endpoint: str, + aws_region: str, + access_key: Optional[str] = None, + secret_key: Optional[str] = None, + session_token: Optional[str] = None, + ): + """ + Initialize Neptune Auth Token Manager. + + Args: + neptune_endpoint: Neptune endpoint URL (bolt://...) + aws_region: AWS region + access_key: AWS access key ID (optional, uses boto3 credentials + if not provided) + secret_key: AWS secret access key (optional, uses boto3 + credentials if not provided) + session_token: AWS session token for temporary credentials (optional) + + Note: This class intentionally stores only simple types (str, None) as instance + attributes because the Neo4j driver inspects AuthManager objects and cannot + serialize complex types like Logger or boto3.Session. + """ + if not BOTO3_AVAILABLE: + raise ProcessingError( + "boto3 is required for IAM authentication. " + "Install with: pip install boto3" + ) + + # Replace bolt protocol with https for signing + # Only store simple serializable types as instance attributes! + self.neptune_endpoint = neptune_endpoint.replace("bolt", "https") + self.aws_region = aws_region + self.cached_auth = None + + # Store credential strings (can be None - will use default chain) + self._access_key = access_key + self._secret_key = secret_key + self._session_token = session_token + + def get_auth(self): + """ + Return the current authentication information. + Uses caching to avoid regenerating SigV4 signatures unnecessarily. + + Returns: + Basic auth token containing SigV4 signature information + """ + cached = "None" if self.cached_auth is None else "cached" + _auth_logger.debug(f"get_auth() called, cached_auth is {cached}") + if self.cached_auth is None: + _auth_logger.info("Generating new SigV4 signed auth token...") + self.cached_auth = self._generate_sigv4_auth_token() + else: + _auth_logger.debug("Using cached SigV4 token") + + return self.cached_auth + + def handle_security_exception(self, auth, error) -> bool: + """ + Handle security exceptions by refreshing the token. + + Args: + auth: The authentication token that caused the exception + error: The security exception that occurred + + Returns: + True to retry with new token, False to propagate exception + """ + error_msg = str(error) if error else "Unknown error" + _auth_logger.warning( + f"Caught SecurityException: {error_msg} - regenerating token" + ) + + # Force token regeneration + self.refresh_token() + return True + + def refresh_token(self): + """ + Force refresh of the authentication token. + Invalidates cached token to force regeneration on next request. + """ + self.cached_auth = None + _auth_logger.info("Auth token forcibly refreshed due to signature expiration") + + def update_context(self, context): + """ + Update execution context for logging purposes. + + Note: We don't store the context as an instance attribute because + the Neo4j driver cannot serialize complex objects. + + Args: + context: Execution context (e.g., from serverless environment) + """ + # Note: Don't store context - Neo4j can't serialize it + # Instead, just log that we received it + _auth_logger.debug( + "Execution context received (not stored due to serialization constraints)" + ) + + def _generate_sigv4_auth_token(self): + """ + Generate a new SigV4 signed authentication token for Neptune. + + Returns: + Basic auth token containing SigV4 signature information + """ + if not BOTO3_AVAILABLE: + raise ProcessingError( + "boto3 is required for IAM authentication. " + "Install with: pip install boto3" + ) + + try: + _auth_logger.info( + f"Generating SigV4 token for endpoint: {self.neptune_endpoint}" + ) + + # Create AWS request for signing + request = AWSRequest(method="GET", url=self.neptune_endpoint, data=None) + host_value = _host_from_url(request.url) + _auth_logger.debug(f"Host header value: {host_value}") + request.headers.add_header("Host", host_value) + + # Get AWS credentials - match the sample implementation exactly + # Create fresh session each time (not stored as instance attribute) + credentials = boto3.Session().get_credentials() + if not credentials: + raise ProcessingError( + "AWS credentials not found. Configure credentials via " + "environment variables, ~/.aws/credentials, IAM role, " + "or provide explicitly." + ) + + # Create credentials namespace for SigV4Auth + # (matching sample implementation) + creds = SimpleNamespace( + access_key=credentials.access_key, + secret_key=credentials.secret_key, + token=credentials.token, + region=self.aws_region, + ) + + # Sign the request + SigV4Auth(creds, self.SERVICE_NAME, self.aws_region).add_auth(request) + + # Create auth info JSON from signed headers + auth_info_json = self._get_auth_info_json(request) + _auth_logger.info(f"Auth info JSON: {auth_info_json}") + + # Return basic auth token with dummy username and signed headers as password + return basic_auth(self.DUMMY_USERNAME, auth_info_json) + + except Exception as e: + raise ProcessingError( + f"Failed to generate SigV4 auth token for Neptune: {str(e)}" + ) from e + + def _get_auth_info_json(self, request) -> str: + """ + Convert signed request headers into JSON string for authentication. + + Args: + request: The signed AWS HTTP request + + Returns: + JSON string containing authentication information + """ + auth_info = { + self.AUTHORIZATION: request.headers.get(self.AUTHORIZATION), + self.HTTP_METHOD_HDR: request.method, + self.X_AMZ_DATE: request.headers.get(self.X_AMZ_DATE), + self.HOST: request.headers.get(self.HOST), + self.X_AMZ_SECURITY_TOKEN: request.headers.get(self.X_AMZ_SECURITY_TOKEN), + } + + return json.dumps(auth_info) + + +class NeptuneDriver: + """Neptune driver wrapper (analogous to Neo4jDriver).""" + + def __init__(self, driver: Any): + """Initialize Neptune driver wrapper.""" + self.driver = driver + self.logger = get_logger("neptune_driver") + + def session(self, database: Optional[str] = None) -> "NeptuneSession": + """ + Create a new session. + + Args: + database: Database name (ignored for Neptune, included for + API compatibility) + + Returns: + NeptuneSession instance + """ + if not NEO4J_AVAILABLE: + raise ProcessingError("Neo4j driver not available") + + try: + # Neptune doesn't support multiple databases, but accepts + # database parameter + session = self.driver.session() + return NeptuneSession(session) + except Exception as e: + raise ProcessingError(f"Failed to create session: {str(e)}") + + def verify_connectivity(self) -> bool: + """Verify connectivity to Neptune server.""" + if not NEO4J_AVAILABLE: + return False + + try: + # Simple query to verify connectivity + with self.driver.session() as session: + session.run("RETURN 1") + return True + except Exception as e: + self.logger.warning(f"Connectivity check failed: {e}") + return False + + def close(self) -> None: + """Close the driver.""" + if self.driver: + self.driver.close() + + +class NeptuneSession: + """Neptune session wrapper (analogous to Neo4jSession).""" + + def __init__(self, session: Any): + """Initialize Neptune session wrapper.""" + self.session = session + self.logger = get_logger("neptune_session") + + def run(self, query: str, parameters: Optional[Dict[str, Any]] = None) -> Any: + """ + Run an OpenCypher query. + + Args: + query: OpenCypher query string + parameters: Query parameters + + Returns: + Query result + """ + if not NEO4J_AVAILABLE: + raise ProcessingError("Neo4j driver not available") + + try: + result = self.session.run(query, parameters or {}) + return result + except Exception as e: + raise ProcessingError(f"Query execution failed: {str(e)}") + + def begin_transaction(self) -> "NeptuneTransaction": + """Begin a new transaction.""" + if not NEO4J_AVAILABLE: + raise ProcessingError("Neo4j driver not available") + + try: + tx = self.session.begin_transaction() + return NeptuneTransaction(tx) + except Exception as e: + raise ProcessingError(f"Failed to begin transaction: {str(e)}") + + def read_transaction(self, func: Any, **kwargs) -> Any: + """Execute a read transaction.""" + return self.session.execute_read(func, **kwargs) + + def write_transaction(self, func: Any, **kwargs) -> Any: + """Execute a write transaction.""" + return self.session.execute_write(func, **kwargs) + + def close(self) -> None: + """Close the session.""" + if self.session: + self.session.close() + + def __enter__(self): + """Context manager entry.""" + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + """Context manager exit.""" + self.close() + + +class NeptuneTransaction: + """Neptune transaction wrapper (analogous to Neo4jTransaction).""" + + def __init__(self, transaction: Any): + """Initialize Neptune transaction wrapper.""" + self.transaction = transaction + self.logger = get_logger("neptune_transaction") + + def run(self, query: str, parameters: Optional[Dict[str, Any]] = None) -> Any: + """ + Run an OpenCypher query within the transaction. + + Args: + query: OpenCypher query string + parameters: Query parameters + + Returns: + Query result + """ + if not NEO4J_AVAILABLE: + raise ProcessingError("Neo4j driver not available") + + try: + result = self.transaction.run(query, parameters or {}) + return result + except Exception as e: + raise ProcessingError(f"Transaction query failed: {str(e)}") + + def commit(self) -> None: + """Commit the transaction.""" + if self.transaction: + self.transaction.commit() + + def rollback(self) -> None: + """Rollback the transaction.""" + if self.transaction: + self.transaction.rollback() + + def __enter__(self): + """Context manager entry.""" + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + """Context manager exit.""" + if exc_type is not None: + self.rollback() + else: + self.commit() + + +class AmazonNeptuneStore: + """ + Amazon Neptune store for property graph storage and OpenCypher querying + using the Neo4j Bolt driver. + + Features: + • IAM authentication using AWS SigV4 signing via AuthManager + • OpenCypher query language support via Bolt protocol + • Node and relationship CRUD operations + • Transaction support with commit/rollback + • Graph analytics + • Automatic retry with backoff for transient errors + • Connection recovery and token refresh + • Performance optimization + + Note: Amazon Neptune uses the Bolt protocol for OpenCypher queries. + The endpoint is: bolt://:8182 + """ + + def __init__( + self, + endpoint: Optional[str] = None, + port: int = 8182, + region: Optional[str] = None, + iam_auth: bool = True, + access_key: Optional[str] = None, + secret_key: Optional[str] = None, + session_token: Optional[str] = None, + use_ssl: bool = True, + max_connection_pool_size: int = 50, + connection_timeout: float = 30.0, + **config, + ): + """ + Initialize Amazon Neptune store. + + Args: + endpoint: Neptune cluster endpoint + (e.g., 'cluster.region.neptune.amazonaws.com') + port: Neptune Bolt port (default: 8182) + region: AWS region (e.g., 'us-east-1') + iam_auth: Use IAM authentication (default: True) + access_key: AWS access key ID (optional) + secret_key: AWS secret access key (optional) + session_token: AWS session token for temporary credentials (optional) + use_ssl: Use SSL/TLS connection (default: True, required for Neptune) + max_connection_pool_size: Maximum number of connections in the pool + connection_timeout: Connection timeout in seconds + **config: Additional configuration options + """ + self.logger = get_logger("amazon_neptune_store") + self.config = config + self.progress_tracker = get_progress_tracker() + + # Ensure progress tracker is enabled + if not self.progress_tracker.enabled: + self.progress_tracker.enabled = True + + # Check dependencies + if not NEO4J_AVAILABLE: + raise ProcessingError( + "neo4j driver is required. Install with: pip install neo4j" + ) + + # Connection settings + self.endpoint = endpoint or config.get("endpoint") + self.port = int(port) if port else config.get("port", 8182) + self.region = ( + region + or config.get("region") + or os.getenv("AWS_REGION") + or os.getenv("AWS_DEFAULT_REGION", "us-east-1") + ) + self.iam_auth = iam_auth + self.use_ssl = use_ssl + self.max_connection_pool_size = max_connection_pool_size + self.connection_timeout = connection_timeout + + # IAM credentials + self._access_key = access_key + self._secret_key = secret_key + self._session_token = session_token + + # Driver and auth manager + self._driver = None + self._auth_manager: Optional[NeptuneAuthTokenManager] = None + self._connected = False + + @property + def bolt_uri(self) -> str: + """Get the Bolt URI for Neptune connection.""" + return f"bolt://{self.endpoint}:{self.port}/opencypher" + + def _is_retriable_error(self, e: Exception) -> bool: + """ + Determine if an exception is retriable based on error message patterns. + + Args: + e: Exception to evaluate + + Returns: + True if error should be retried, False otherwise + """ + is_retriable = False + err_msg = str(e) + + # Check for DatabaseError with authentication issues + if NEO4J_AVAILABLE and isinstance(e, DatabaseError): + is_retriable = any( + retriable_msg in err_msg for retriable_msg in RETRIABLE_ERROR_MESSAGES + ) + elif isinstance(e, NETWORK_ERRORS): + is_retriable = True + else: + is_retriable = any( + retriable_msg in err_msg for retriable_msg in RETRIABLE_ERROR_MESSAGES + ) + + type_name = type(e).__name__ + self.logger.debug( + f"Retry evaluation: [{type_name}] {err_msg} -> " + f"is_retriable={is_retriable}" + ) + return is_retriable + + def _is_non_retriable_error(self, e: Exception) -> bool: + """ + Determine if an exception is non-retriable (inverse of is_retriable_error). + + Args: + e: Exception to evaluate + + Returns: + True if error should not be retried, False otherwise + """ + return not self._is_retriable_error(e) + + def _reset_connection_if_needed(self, details: Dict = None): + """ + Reset the driver connection if the current exception indicates a + connection issue. Called by backoff decorator when retrying failed + operations. + + Args: + details: Backoff details dictionary (optional) + """ + e = sys.exc_info()[1] + if e is None: + return + + err_msg = str(e) + is_reconnectable = False + + if isinstance(e, NETWORK_ERRORS): + is_reconnectable = True + else: + is_reconnectable = any(msg in err_msg for msg in RETRIABLE_ERROR_MESSAGES) + + self.logger.info( + f"Connection issue detected: is_reconnectable={is_reconnectable}" + ) + + if is_reconnectable: + self.logger.info("Resetting connection due to connection issue") + self._recreate_driver() + + def _recreate_driver(self): + """Recreate the Neo4j driver.""" + if self._driver: + try: + self._driver.close() + except Exception as close_err: + self.logger.warning(f"Error closing driver: {close_err}") + self._driver = self._create_driver() + + def _create_driver(self): + """ + Create and configure Neo4j driver for Neptune connection. + Uses IAM authentication if iam_auth is True. + + Returns: + Configured Neo4j GraphDatabase driver instance + + Raises: + ProcessingError: If driver creation fails + """ + self.logger.info("Creating Neo4j driver for Neptune") + + try: + bolt_uri = self.bolt_uri + self.logger.info(f"Connecting to: {bolt_uri}") + + if self.iam_auth: + if not self.region: + raise ValidationError( + "AWS region is required for IAM authentication. " + "Set via 'region' parameter or AWS_REGION environment variable." + ) + + # Create auth manager for IAM authentication + self._auth_manager = NeptuneAuthTokenManager( + neptune_endpoint=bolt_uri, + aws_region=self.region, + access_key=self._access_key, + secret_key=self._secret_key, + session_token=self._session_token, + ) + + driver = GraphDatabase.driver( + bolt_uri, + auth=self._auth_manager, + encrypted=self.use_ssl, + max_connection_pool_size=self.max_connection_pool_size, + connection_timeout=self.connection_timeout, + ) + else: + driver = GraphDatabase.driver( + bolt_uri, + auth=None, + encrypted=self.use_ssl, + max_connection_pool_size=self.max_connection_pool_size, + connection_timeout=self.connection_timeout, + ) + + return driver + + except Exception as e: + self.logger.error(f"Failed to create driver: {str(e)}") + raise ProcessingError(f"Failed to create Neptune driver: {str(e)}") from e + + def _execute_with_retry(self, operation_func, *args, **kwargs): + """ + Execute an operation with retry logic. + + Args: + operation_func: Function to execute + *args: Arguments to pass to the function + **kwargs: Keyword arguments to pass to the function + + Returns: + Result of the operation + """ + if BACKOFF_AVAILABLE: + # Use backoff decorator for retry logic + @backoff.on_exception( + backoff.constant, + ( + (ServiceUnavailable, DatabaseError, OSError) + if NEO4J_AVAILABLE + else (OSError,) + ), + max_tries=NUM_RETRIES, + jitter=None, + giveup=self._is_non_retriable_error, + on_backoff=lambda details: self._reset_connection_if_needed(details), + interval=1, + ) + def _execute(): + return operation_func(*args, **kwargs) + + return _execute() + else: + # Simple retry without backoff library + last_error = None + for attempt in range(NUM_RETRIES): + try: + return operation_func(*args, **kwargs) + except Exception as e: + last_error = e + if self._is_non_retriable_error(e): + raise + self.logger.warning( + f"Attempt {attempt + 1}/{NUM_RETRIES} failed: {e}" + ) + self._reset_connection_if_needed() + raise last_error + + def _run_query( + self, query: str, parameters: Optional[Dict[str, Any]] = None + ) -> List[Dict[str, Any]]: + """ + Execute an OpenCypher query against Neptune. + + Args: + query: OpenCypher query string + parameters: Query parameters dictionary (optional) + + Returns: + List of dictionaries containing query results + """ + + def _execute(): + try: + with self._driver.session() as session: + result = session.run(query, parameters or {}) + return [dict(record) for record in result] + + # In some cases, a generic AttributeError is thrown prior to + # the Neo4j specific exception. This throws the root context + # of the AttributeError for more specific error handling. + except AttributeError as e: + self.logger.error(f"AttributeError caught: {str(e)}") + if e.__context__: + raise e.__context__ from e + raise + except ( + (DatabaseError, ServiceUnavailable) if NEO4J_AVAILABLE else Exception + ) as e: + self.logger.error(f"Query error: {str(e)}") + raise + + return self._execute_with_retry(_execute) + + def connect(self, **options) -> bool: + """ + Connect to Amazon Neptune (initializes Bolt driver). + + Args: + **options: Connection options + + Returns: + True if connected successfully + """ + if not self.endpoint: + raise ValidationError("Neptune endpoint is required") + + try: + # Create driver + self._driver = self._create_driver() + + # Test connectivity with a simple query + self._run_query("RETURN 1 as test") + + self._connected = True + self.logger.info(f"Connected to Amazon Neptune at {self.endpoint}") + return True + + except Exception as e: + self._connected = False + raise ProcessingError(f"Failed to connect to Neptune: {str(e)}") from e + + def close(self) -> None: + """Close connection to Neptune.""" + if self._driver: + try: + self._driver.close() + except Exception as e: + self.logger.warning(f"Error closing driver: {e}") + self._driver = None + self._connected = False + self.logger.info("Disconnected from Amazon Neptune") + + def _ensure_connected(self): + """Ensure the driver is connected, connecting if necessary.""" + if self._driver is None or not self._connected: + self.connect() + + def get_session(self, database: Optional[str] = None) -> NeptuneSession: + """ + Get or create a session. + + Args: + database: Database name (ignored for Neptune, included for + API compatibility) + + Returns: + NeptuneSession instance + """ + self._ensure_connected() + return NeptuneSession(self._driver.session()) + + def _generate_id(self) -> str: + """Generate a unique node/relationship ID.""" + return str(uuid.uuid4()) + + def _parse_record_value(self, value: Any) -> Any: + """ + Parse a value from a Neo4j record. + + Args: + value: Value from Neo4j record + + Returns: + Parsed value + """ + if value is None: + return None + + # Check if it's a Neo4j Node + if hasattr(value, "id") and hasattr(value, "labels"): + # It's a Node + return { + "id": value.get("~id", str(value.id) if hasattr(value, "id") else None), + "labels": list(value.labels) if hasattr(value, "labels") else [], + "properties": dict(value) if value else {}, + } + + # Check if it's a Neo4j Relationship + if hasattr(value, "type") and hasattr(value, "start_node"): + return { + "id": value.get("~id", str(value.id) if hasattr(value, "id") else None), + "type": value.type, + "start_node_id": ( + str(value.start_node.id) + if hasattr(value.start_node, "id") + else None + ), + "end_node_id": ( + str(value.end_node.id) if hasattr(value.end_node, "id") else None + ), + "properties": dict(value) if value else {}, + } + + # Check if it's a dict (Neptune returns dicts for nodes/relationships) + if isinstance(value, dict): + if "~id" in value or "~entityType" in value: + # This is a Neptune node or relationship + properties = {k: v for k, v in value.items() if not k.startswith("~")} + parsed = { + "id": value.get("~id"), + "labels": value.get("~labels", []), + "properties": properties, + } + # Handle relationships + if "~type" in value: + parsed["type"] = value.get("~type") + if "~start" in value: + parsed["start_node_id"] = value.get("~start") + if "~end" in value: + parsed["end_node_id"] = value.get("~end") + return parsed + return value + + return value + + def _parse_results(self, records: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + """ + Parse Neptune/Neo4j query results. + + Args: + records: List of record dictionaries + + Returns: + List of parsed records + """ + parsed_records = [] + for record in records: + parsed_record = {} + for key, value in record.items(): + parsed_record[key] = self._parse_record_value(value) + parsed_records.append(parsed_record) + return parsed_records + + def create_node( + self, + labels: List[str], + properties: Dict[str, Any], + **options, + ) -> Dict[str, Any]: + """ + Create a node in the graph. + + Uses Neptune's native ~id for node identification. If 'id' is provided + in properties, it will be used as the Neptune ~id. Otherwise, a UUID + is generated. + + If a node with the same ID already exists: + - By default (merge=True), uses MERGE to return existing node or create new + - With merge=False, uses CREATE which will fail if ID exists + + Args: + labels: Node labels + properties: Node properties (may include 'id' for custom ~id) + **options: Additional options including: + - merge (bool): If True (default), use MERGE; if False, use CREATE + + Returns: + Created node information including ID + """ + tracking_id = self.progress_tracker.start_tracking( + module="graph_store", + submodule="AmazonNeptuneStore", + message=f"Creating node with labels {labels}", + ) + + try: + self._ensure_connected() + + # Extract ID from properties or generate one + props_copy = dict(properties) + node_id = props_copy.pop("id", None) or self._generate_id() + use_merge = options.get("merge", True) + + label_str = ":".join(labels) if labels else "Node" + + # Build parameters + params = {"node_id": str(node_id)} + for key, value in props_copy.items(): + params[key] = value + + if use_merge: + # MERGE: Return existing node if ID matches, or create new + set_parts = [] + for key in props_copy.keys(): + set_parts.append(f"n.{key} = ${key}") + + if set_parts: + set_clause = ", ".join(set_parts) + query = ( + f"MERGE (n:{label_str} {{`~id`: $node_id}}) " + f"ON CREATE SET {set_clause} " + f"ON MATCH SET {set_clause} RETURN n" + ) + else: + query = f"MERGE (n:{label_str} {{`~id`: $node_id}}) RETURN n" + else: + # CREATE: Will fail if node with same ID exists + prop_parts = ["`~id`: $node_id"] + for key in props_copy.keys(): + prop_parts.append(f"{key}: ${key}") + prop_assignments = ", ".join(prop_parts) + query = f"CREATE (n:{label_str} {{{prop_assignments}}}) RETURN n" + + records = self._run_query(query, params) + parsed = self._parse_results(records) + + if parsed: + node_data = parsed[0].get("n", {}) + if isinstance(node_data, dict): + if "id" not in node_data: + node_data["id"] = node_id + if "labels" not in node_data: + node_data["labels"] = labels + + self.progress_tracker.stop_tracking( + tracking_id, + status="completed", + message=f"Created node with ID {node_id}", + ) + return node_data + + # Node was created but not returned - return what we know + self.progress_tracker.stop_tracking( + tracking_id, + status="completed", + message=f"Created node with ID {node_id}", + ) + return { + "id": node_id, + "labels": labels, + "properties": props_copy, + } + + except Exception as e: + self.progress_tracker.stop_tracking( + tracking_id, status="failed", message=str(e) + ) + raise ProcessingError(f"Failed to create node: {str(e)}") from e + + def create_nodes( + self, + nodes: List[Dict[str, Any]], + **options, + ) -> List[Dict[str, Any]]: + """ + Create multiple nodes in batch. + + Args: + nodes: List of node dictionaries with 'labels' and 'properties' + **options: Additional options + + Returns: + List of created node information + """ + tracking_id = self.progress_tracker.start_tracking( + module="graph_store", + submodule="AmazonNeptuneStore", + message=f"Creating {len(nodes)} nodes in batch", + ) + + try: + created_nodes = [] + + for node in nodes: + labels = node.get("labels", []) + properties = node.get("properties", {}) + + result = self.create_node(labels, properties, **options) + created_nodes.append(result) + + self.progress_tracker.stop_tracking( + tracking_id, + status="completed", + message=f"Created {len(created_nodes)} nodes", + ) + return created_nodes + + except Exception as e: + self.progress_tracker.stop_tracking( + tracking_id, status="failed", message=str(e) + ) + raise ProcessingError(f"Failed to create nodes: {str(e)}") from e + + def get_node( + self, + node_id: Union[int, str], + **options, + ) -> Optional[Dict[str, Any]]: + """ + Get a node by ID (using Neptune's native ~id). + + Args: + node_id: Node ID (Neptune's ~id value) + **options: Additional options + + Returns: + Node information or None if not found + """ + try: + self._ensure_connected() + + # Use id(n) function to match and return Neptune's native ~id + query = ( + "MATCH (n) WHERE id(n) = $id " + "RETURN n, labels(n) as labels, id(n) as node_id" + ) + + records = self._run_query(query, {"id": str(node_id)}) + parsed = self._parse_results(records) + + if parsed: + record = parsed[0] + node_data = record.get("n", {}) + # Use node_id from query (the ~id we set) as primary ID + returned_id = record.get("node_id", node_id) + properties = ( + node_data.get("properties", {}) + if isinstance(node_data, dict) + else {} + ) + return { + "id": returned_id, + "labels": record.get("labels", []), + "properties": properties, + } + return None + + except Exception as e: + raise ProcessingError(f"Failed to get node: {str(e)}") from e + + def get_nodes( + self, + labels: Optional[List[str]] = None, + properties: Optional[Dict[str, Any]] = None, + limit: int = 100, + **options, + ) -> List[Dict[str, Any]]: + """ + Get nodes matching criteria. + + Args: + labels: Filter by labels + properties: Filter by properties + limit: Maximum number of nodes to return + **options: Additional options + + Returns: + List of matching nodes + """ + try: + self._ensure_connected() + + # Build query + if labels: + label_str = ":".join(labels) + query = f"MATCH (n:{label_str})" + else: + query = "MATCH (n)" + + # Add property filters + params = {} + if properties: + conditions = [] + for key, value in properties.items(): + param_key = f"prop_{key}" + conditions.append(f"n.{key} = ${param_key}") + params[param_key] = value + query += " WHERE " + " AND ".join(conditions) + + # Include id(n) to get Neptune's native ~id + query += ( + f" RETURN n, labels(n) as labels, id(n) as node_id " f"LIMIT {limit}" + ) + + records = self._run_query(query, params) + parsed = self._parse_results(records) + + nodes = [] + for record in parsed: + node = record.get("n", {}) + # Use node_id from query (the ~id we set) as primary ID + returned_id = record.get("node_id") + properties_data = ( + node.get("properties", {}) if isinstance(node, dict) else {} + ) + nodes.append( + { + "id": returned_id, + "labels": record.get("labels", []), + "properties": properties_data, + } + ) + + return nodes + + except Exception as e: + raise ProcessingError(f"Failed to get nodes: {str(e)}") from e + + def update_node( + self, + node_id: Union[int, str], + properties: Dict[str, Any], + merge: bool = True, + **options, + ) -> Dict[str, Any]: + """ + Update a node's properties (using Neptune's native ~id). + + Args: + node_id: Node ID (Neptune's ~id value) + properties: Properties to update + merge: If True, merge properties; if False, replace + **options: Additional options + + Returns: + Updated node information + """ + try: + self._ensure_connected() + + if merge: + # SET n += $props merges properties + query = ( + "MATCH (n) WHERE id(n) = $id SET n += $props " + "RETURN n, labels(n) as labels, id(n) as node_id" + ) + else: + # SET n = $props replaces all properties + query = ( + "MATCH (n) WHERE id(n) = $id SET n = $props " + "RETURN n, labels(n) as labels, id(n) as node_id" + ) + + records = self._run_query(query, {"id": str(node_id), "props": properties}) + parsed = self._parse_results(records) + + if parsed: + record = parsed[0] + node_data = record.get("n", {}) + # Use node_id from query (the ~id we set) as primary ID + returned_id = record.get("node_id", node_id) + properties_data = ( + node_data.get("properties", {}) + if isinstance(node_data, dict) + else {} + ) + return { + "id": returned_id, + "labels": record.get("labels", []), + "properties": properties_data, + } + else: + raise ProcessingError(f"Node with ID {node_id} not found") + + except Exception as e: + raise ProcessingError(f"Failed to update node: {str(e)}") from e + + def delete_node( + self, + node_id: Union[int, str], + detach: bool = True, + **options, + ) -> bool: + """ + Delete a node (using Neptune's native ~id). + + Args: + node_id: Node ID (Neptune's ~id value) + detach: If True, delete relationships as well + **options: Additional options + + Returns: + True if deleted successfully + """ + try: + self._ensure_connected() + + if detach: + query = "MATCH (n) WHERE id(n) = $id DETACH DELETE n" + else: + query = "MATCH (n) WHERE id(n) = $id DELETE n" + + self._run_query(query, {"id": str(node_id)}) + return True + + except Exception as e: + raise ProcessingError(f"Failed to delete node: {str(e)}") from e + + def create_relationship( + self, + start_node_id: Union[int, str], + end_node_id: Union[int, str], + rel_type: str, + properties: Optional[Dict[str, Any]] = None, + **options, + ) -> Dict[str, Any]: + """ + Create a relationship between two nodes (using Neptune's native ~id). + + Uses Neptune's native ~id for relationship identification. If 'id' is + provided in properties, it will be used as the Neptune ~id. + + Args: + start_node_id: Start node ID (Neptune's ~id value) + end_node_id: End node ID (Neptune's ~id value) + rel_type: Relationship type + properties: Relationship properties (may include 'id' for custom ~id) + **options: Additional options + + Returns: + Created relationship information + """ + tracking_id = self.progress_tracker.start_tracking( + module="graph_store", + submodule="AmazonNeptuneStore", + message=f"Creating relationship [{rel_type}]", + ) + + try: + self._ensure_connected() + + properties = properties or {} + props_copy = dict(properties) + rel_id = props_copy.pop("id", None) or self._generate_id() + + # Build query using id() function for node matching and ~id for relationship + params = { + "start_id": str(start_node_id), + "end_id": str(end_node_id), + "rel_id": str(rel_id), + } + + # Build property assignments including ~id + prop_parts = ["`~id`: $rel_id"] + for key, value in props_copy.items(): + prop_parts.append(f"{key}: ${key}") + params[key] = value + + prop_assignments = ", ".join(prop_parts) + query = ( + f"MATCH (a), (b) WHERE id(a) = $start_id AND id(b) = $end_id " + f"CREATE (a)-[r:{rel_type} {{{prop_assignments}}}]->(b) RETURN r" + ) + + records = self._run_query(query, params) + parsed = self._parse_results(records) + + rel_data = { + "id": rel_id, + "type": rel_type, + "start_node_id": start_node_id, + "end_node_id": end_node_id, + "properties": props_copy, + } + + if parsed: + returned_rel = parsed[0].get("r", {}) + if isinstance(returned_rel, dict) and returned_rel.get("id"): + rel_data["id"] = returned_rel["id"] + + self.progress_tracker.stop_tracking( + tracking_id, + status="completed", + message=f"Created relationship with ID {rel_id}", + ) + return rel_data + + except Exception as e: + self.progress_tracker.stop_tracking( + tracking_id, status="failed", message=str(e) + ) + raise ProcessingError(f"Failed to create relationship: {str(e)}") from e + + def get_relationships( + self, + node_id: Optional[Union[int, str]] = None, + rel_type: Optional[str] = None, + direction: str = "both", + limit: int = 100, + **options, + ) -> List[Dict[str, Any]]: + """ + Get relationships matching criteria (using Neptune's native ~id). + + Args: + node_id: Filter by node ID (Neptune's ~id value) + rel_type: Filter by relationship type + direction: Direction ("in", "out", "both") + limit: Maximum number of relationships + **options: Additional options + + Returns: + List of matching relationships + """ + try: + self._ensure_connected() + + type_filter = f":{rel_type}" if rel_type else "" + params = {} + + if node_id is not None: + params["node_id"] = str(node_id) + if direction == "out": + query = ( + f"MATCH (a)-[r{type_filter}]->(b) " + f"WHERE id(a) = $node_id " + f"RETURN r, id(a) as start_id, id(b) as end_id " + f"LIMIT {limit}" + ) + elif direction == "in": + query = ( + f"MATCH (a)<-[r{type_filter}]-(b) " + f"WHERE id(a) = $node_id " + f"RETURN r, id(b) as start_id, id(a) as end_id " + f"LIMIT {limit}" + ) + else: + query = ( + f"MATCH (a)-[r{type_filter}]-(b) " + f"WHERE id(a) = $node_id " + f"RETURN r, id(a) as start_id, id(b) as end_id " + f"LIMIT {limit}" + ) + else: + query = ( + f"MATCH (a)-[r{type_filter}]->(b) " + f"RETURN r, id(a) as start_id, id(b) as end_id " + f"LIMIT {limit}" + ) + + records = self._run_query(query, params) + parsed = self._parse_results(records) + + relationships = [] + for record in parsed: + rel = record.get("r", {}) + relationships.append( + { + "id": rel.get("id") if isinstance(rel, dict) else None, + "type": ( + rel.get("type", rel_type) + if isinstance(rel, dict) + else rel_type + ), + "start_node_id": record.get("start_id"), + "end_node_id": record.get("end_id"), + "properties": ( + rel.get("properties", {}) if isinstance(rel, dict) else {} + ), + } + ) + + return relationships + + except Exception as e: + raise ProcessingError(f"Failed to get relationships: {str(e)}") from e + + def delete_relationship( + self, + rel_id: Union[int, str], + **options, + ) -> bool: + """ + Delete a relationship (using Neptune's native ~id). + + Args: + rel_id: Relationship ID (Neptune's ~id value) + **options: Additional options + + Returns: + True if deleted successfully + """ + try: + self._ensure_connected() + + # Use id(r) function to match Neptune's native ~id + query = "MATCH ()-[r]->() WHERE id(r) = $id DELETE r" + + self._run_query(query, {"id": str(rel_id)}) + return True + + except Exception as e: + raise ProcessingError(f"Failed to delete relationship: {str(e)}") from e + + def execute_query( + self, + query: str, + parameters: Optional[Dict[str, Any]] = None, + **options, + ) -> Dict[str, Any]: + """ + Execute an OpenCypher query. + + Args: + query: OpenCypher query string + parameters: Query parameters + **options: Additional options + + Returns: + Query results + """ + tracking_id = self.progress_tracker.start_tracking( + module="graph_store", + submodule="AmazonNeptuneStore", + message="Executing OpenCypher query", + ) + + try: + self._ensure_connected() + + records = self._run_query(query, parameters or {}) + parsed = self._parse_results(records) + + self.progress_tracker.stop_tracking( + tracking_id, + status="completed", + message=f"Query returned {len(parsed)} records", + ) + + return { + "success": True, + "records": parsed, + "keys": list(parsed[0].keys()) if parsed else [], + "metadata": {"query": query}, + } + + except Exception as e: + self.progress_tracker.stop_tracking( + tracking_id, status="failed", message=str(e) + ) + raise ProcessingError(f"Query execution failed: {str(e)}") from e + + def get_neighbors( + self, + node_id: Union[int, str], + rel_type: Optional[str] = None, + direction: str = "both", + depth: int = 1, + **options, + ) -> List[Dict[str, Any]]: + """ + Get neighboring nodes (using Neptune's native ~id). + + Args: + node_id: Starting node ID (Neptune's ~id value) + rel_type: Filter by relationship type + direction: Direction ("in", "out", "both") + depth: Traversal depth + **options: Additional options + + Returns: + List of neighboring nodes + """ + try: + self._ensure_connected() + + type_filter = f":{rel_type}" if rel_type else "" + + if direction == "out": + pattern = f"-[r{type_filter}*1..{depth}]->" + elif direction == "in": + pattern = f"<-[r{type_filter}*1..{depth}]-" + else: + pattern = f"-[r{type_filter}*1..{depth}]-" + + # Include id(neighbor) to get the ~id we set when creating node + query = ( + f"MATCH (start){pattern}(neighbor) " + f"WHERE id(start) = $node_id " + f"RETURN DISTINCT neighbor, labels(neighbor) as labels, " + f"id(neighbor) as neighbor_id" + ) + + records = self._run_query(query, {"node_id": str(node_id)}) + parsed = self._parse_results(records) + + neighbors = [] + for record in parsed: + node = record.get("neighbor", {}) + # Use neighbor_id from query (the ~id we set) as primary ID + neighbor_id = record.get("neighbor_id") + neighbors.append( + { + "id": neighbor_id + or (node.get("id") if isinstance(node, dict) else None), + "labels": record.get("labels", []), + "properties": ( + node.get("properties", {}) + if isinstance(node, dict) + else node if isinstance(node, dict) else {} + ), + } + ) + + return neighbors + + except Exception as e: + raise ProcessingError(f"Failed to get neighbors: {str(e)}") from e + + def shortest_path( + self, + start_node_id: Union[int, str], + end_node_id: Union[int, str], + rel_type: Optional[str] = None, + max_depth: int = 10, + **options, + ) -> Optional[Dict[str, Any]]: + """ + Find shortest path between two nodes (using Neptune's native ~id). + + Note: Neptune OpenCypher has limited support for shortestPath syntax. + This implementation uses a BFS-style approach compatible with Neptune. + + Args: + start_node_id: Starting node ID (Neptune's ~id value) + end_node_id: Ending node ID (Neptune's ~id value) + rel_type: Filter by relationship type + max_depth: Maximum path length + **options: Additional options + + Returns: + Shortest path information or None if not found + """ + try: + self._ensure_connected() + + type_filter = f":{rel_type}" if rel_type else "" + + # Neptune doesn't support named path patterns in shortestPath + # Use iterative depth search instead + for depth in range(1, max_depth + 1): + # Include id(start) and id(end) to get ~id we set + query = ( + f"MATCH (start)-[r{type_filter}*{depth}]-(end) " + f"WHERE id(start) = $start_id AND id(end) = $end_id " + f"RETURN start, end, {depth} as length, " + f"id(start) as start_node_id, id(end) as end_node_id, " + f"labels(start) as start_labels, " + f"labels(end) as end_labels LIMIT 1" + ) + + records = self._run_query( + query, + { + "start_id": str(start_node_id), + "end_id": str(end_node_id), + }, + ) + parsed = self._parse_results(records) + + if parsed: + record = parsed[0] + start_node = record.get("start", {}) + end_node = record.get("end", {}) + + # Use the ~id from id() function as primary ID + return { + "length": depth, + "start_node": { + "id": record.get("start_node_id") + or ( + start_node.get("id") + if isinstance(start_node, dict) + else None + ), + "labels": record.get("start_labels", []), + "properties": ( + start_node.get("properties", {}) + if isinstance(start_node, dict) + else start_node if isinstance(start_node, dict) else {} + ), + }, + "end_node": { + "id": record.get("end_node_id") + or ( + end_node.get("id") + if isinstance(end_node, dict) + else None + ), + "labels": record.get("end_labels", []), + "properties": ( + end_node.get("properties", {}) + if isinstance(end_node, dict) + else end_node if isinstance(end_node, dict) else {} + ), + }, + "found": True, + } + + # No path found within max_depth + return None + + except Exception as e: + self.logger.warning(f"Shortest path search failed: {str(e)}") + return None + + def create_index( + self, + label: str, + property_name: str, + index_type: str = "btree", + **options, + ) -> bool: + """ + Create an index on a property. + + Note: Neptune manages indexes differently than Neo4j. + This logs a message about Neptune's index management. + + Args: + label: Node label + property_name: Property to index + index_type: Index type + **options: Additional options + + Returns: + True (acknowledgement) + """ + self.logger.info( + f"Index creation for {label}.{property_name}: " + "Neptune manages indexes automatically. Consider using Neptune's " + "management console or AWS CLI for advanced index management." + ) + return True + + def get_stats(self) -> Dict[str, Any]: + """Get database statistics.""" + try: + self._ensure_connected() + + stats = {} + + # Node count + records = self._run_query("MATCH (n) RETURN count(n) as count") + stats["node_count"] = records[0]["count"] if records else 0 + + # Relationship count + records = self._run_query("MATCH ()-[r]->() RETURN count(r) as count") + stats["relationship_count"] = records[0]["count"] if records else 0 + + stats["backend"] = "amazon_neptune" + stats["endpoint"] = self.endpoint + stats["protocol"] = "bolt" + + return stats + + except Exception as e: + self.logger.warning(f"Failed to get stats: {str(e)}") + return {"status": "error", "message": str(e)} + + def get_status(self) -> Dict[str, Any]: + """ + Get Neptune connection status. + + Returns: + Status information + """ + return { + "backend": "amazon_neptune", + "endpoint": self.endpoint, + "port": self.port, + "region": self.region, + "iam_auth": self.iam_auth, + "use_ssl": self.use_ssl, + "protocol": "bolt", + "connected": self._connected, + } + + def update_auth_context(self, context) -> None: + """ + Update the authentication context for logging purposes. + + This is useful in serverless environments to update logging context. + + Args: + context: Execution context object + """ + if self._auth_manager: + self._auth_manager.update_context(context) + + def refresh_auth_token(self) -> None: + """ + Force refresh of the authentication token. + + Useful when you know the token is about to expire or has expired. + """ + if self._auth_manager: + self._auth_manager.refresh_token() diff --git a/semantica/graph_store/config.py b/semantica/graph_store/config.py index d779027f..387584cd 100644 --- a/semantica/graph_store/config.py +++ b/semantica/graph_store/config.py @@ -6,7 +6,8 @@ supporting multiple configuration sources including environment variables, confi and programmatic configuration. Supported Configuration Sources: - - Environment variables: GRAPH_STORE_DEFAULT_BACKEND, GRAPH_STORE_NEO4J_URI, GRAPH_STORE_FALKORDB_HOST, etc. + - Environment variables: GRAPH_STORE_DEFAULT_BACKEND, + GRAPH_STORE_NEO4J_URI, GRAPH_STORE_FALKORDB_HOST, etc. - Config files: YAML, JSON, TOML formats - Programmatic: Python API for setting graph store configurations @@ -44,7 +45,11 @@ from ..utils.logging import get_logger class GraphStoreConfig: - """Configuration manager for graph store module - supports .env files, environment variables, and programmatic config.""" + """ + Configuration manager for graph store module. + + Supports .env files, environment variables, and programmatic config. + """ def __init__(self, config_file: Optional[str] = None): """ @@ -124,6 +129,15 @@ class GraphStoreConfig: "GRAPH_STORE_FALKORDB_PORT": "falkordb_port", "GRAPH_STORE_FALKORDB_PASSWORD": "falkordb_password", "GRAPH_STORE_FALKORDB_GRAPH_NAME": "falkordb_graph_name", + # Amazon Neptune settings + "GRAPH_STORE_NEPTUNE_ENDPOINT": "neptune_endpoint", + "GRAPH_STORE_NEPTUNE_PORT": "neptune_port", + "GRAPH_STORE_NEPTUNE_REGION": "neptune_region", + "GRAPH_STORE_NEPTUNE_IAM_AUTH": "neptune_iam_auth", + "GRAPH_STORE_NEPTUNE_USE_SSL": "neptune_use_ssl", + "AWS_ACCESS_KEY_ID": "neptune_access_key", + "AWS_SECRET_ACCESS_KEY": "neptune_secret_key", + "AWS_SESSION_TOKEN": "neptune_session_token", } for env_var, config_key in env_mappings.items(): @@ -135,6 +149,7 @@ class GraphStoreConfig: "timeout", "max_retries", "falkordb_port", + "neptune_port", ]: try: self._config[config_key] = int(value) @@ -142,7 +157,11 @@ class GraphStoreConfig: self.logger.warning( f"Invalid integer value for {env_var}: {value}" ) - elif config_key in ["neo4j_encrypted"]: + elif config_key in [ + "neo4j_encrypted", + "neptune_iam_auth", + "neptune_use_ssl", + ]: self._config[config_key] = value.lower() in [ "true", "1", @@ -171,6 +190,15 @@ class GraphStoreConfig: "falkordb_port": 6379, "falkordb_password": None, "falkordb_graph_name": "default", + # Amazon Neptune defaults + "neptune_endpoint": None, + "neptune_port": 8182, + "neptune_region": None, + "neptune_iam_auth": True, + "neptune_use_ssl": True, + "neptune_access_key": None, + "neptune_secret_key": None, + "neptune_session_token": None, } for key, default_value in defaults.items(): @@ -269,6 +297,24 @@ class GraphStoreConfig: "graph_name": self._config.get("falkordb_graph_name"), } + def get_neptune_config(self) -> Dict[str, Any]: + """ + Get Amazon Neptune-specific configuration. + + Returns: + Neptune configuration dictionary + """ + return { + "endpoint": self._config.get("neptune_endpoint"), + "port": self._config.get("neptune_port"), + "region": self._config.get("neptune_region"), + "iam_auth": self._config.get("neptune_iam_auth"), + "use_ssl": self._config.get("neptune_use_ssl"), + "access_key": self._config.get("neptune_access_key"), + "secret_key": self._config.get("neptune_secret_key"), + "session_token": self._config.get("neptune_session_token"), + } + def reset(self) -> None: """Reset configuration to defaults.""" self._config.clear() @@ -278,4 +324,3 @@ class GraphStoreConfig: # Global configuration instance graph_store_config = GraphStoreConfig() - diff --git a/semantica/graph_store/graph_store.py b/semantica/graph_store/graph_store.py index ef15ff17..933b167a 100644 --- a/semantica/graph_store/graph_store.py +++ b/semantica/graph_store/graph_store.py @@ -34,7 +34,7 @@ License: MIT from typing import Any, Dict, List, Optional, Tuple, Union -from ..utils.exceptions import ProcessingError, ValidationError +from ..utils.exceptions import ValidationError from ..utils.logging import get_logger from ..utils.progress_tracker import get_progress_tracker from .config import graph_store_config @@ -214,7 +214,9 @@ class RelationshipManager: Returns: List of relationships """ - return self.backend.get_relationships(node_id, rel_type, direction, limit, **options) + return self.backend.get_relationships( + node_id, rel_type, direction, limit, **options + ) def delete( self, @@ -290,6 +292,7 @@ class QueryEngine: ) -> str: """Generate cache key for query.""" import hashlib + key_str = f"{query}:{str(parameters)}" return hashlib.md5(key_str.encode()).hexdigest() @@ -340,7 +343,9 @@ class GraphAnalytics: Returns: Path information or None """ - return self.backend.shortest_path(start_node_id, end_node_id, rel_type, max_depth, **options) + return self.backend.shortest_path( + start_node_id, end_node_id, rel_type, max_depth, **options + ) def get_neighbors( self, @@ -363,7 +368,9 @@ class GraphAnalytics: Returns: List of neighboring nodes """ - return self.backend.get_neighbors(node_id, rel_type, direction, depth, **options) + return self.backend.get_neighbors( + node_id, rel_type, direction, depth, **options + ) def degree_centrality( self, @@ -440,7 +447,7 @@ class GraphAnalytics: Component information """ backend_type = type(self.backend).__name__ - + if "Neo4j" in backend_type: query = """ CALL gds.wcc.stream({ @@ -452,16 +459,23 @@ class GraphAnalytics: """ params = {"label": labels[0] if labels else "*"} result = self.backend.execute_query(query, params) - return [{"component": r["componentId"], "nodes": r["nodes"]} for r in result] - + return [ + {"component": r["componentId"], "nodes": r["nodes"]} for r in result + ] + elif "NetworkX" in backend_type: import networkx as nx + G = self.backend.graph components = list(nx.connected_components(G)) - return [{"component": i, "nodes": list(c)} for i, c in enumerate(components)] - + return [ + {"component": i, "nodes": list(c)} for i, c in enumerate(components) + ] + else: - raise NotImplementedError(f"connected_components not implemented for {backend_type}") + raise NotImplementedError( + f"connected_components not implemented for {backend_type}" + ) class GraphManager: @@ -534,7 +548,11 @@ class GraphStore: self.progress_tracker.enabled = True # Determine backend - self.backend = backend or config.get("backend") or graph_store_config.get("default_backend", "neo4j") + self.backend = ( + backend + or config.get("backend") + or graph_store_config.get("default_backend", "neo4j") + ) self.config = config # Initialize store backend @@ -546,16 +564,25 @@ class GraphStore: """Initialize the appropriate store backend based on backend type.""" if self.backend == "neo4j": from .neo4j_store import Neo4jStore + neo4j_config = graph_store_config.get_neo4j_config() neo4j_config.update(self.config) self._store_backend = Neo4jStore(**neo4j_config) elif self.backend == "falkordb": from .falkordb_store import FalkorDBStore + falkordb_config = graph_store_config.get_falkordb_config() falkordb_config.update(self.config) self._store_backend = FalkorDBStore(**falkordb_config) + elif self.backend == "neptune" or self.backend == "amazon_neptune": + from .amazon_neptune import AmazonNeptuneStore + + neptune_config = graph_store_config.get_neptune_config() + neptune_config.update(self.config) + self._store_backend = AmazonNeptuneStore(**neptune_config) + else: raise ValidationError(f"Unknown backend: {self.backend}") @@ -621,7 +648,9 @@ class GraphStore: **options, ) -> List[Dict[str, Any]]: """Get nodes matching criteria.""" - return self._manager.nodes.get(labels=labels, properties=properties, limit=limit, **options) + return self._manager.nodes.get( + labels=labels, properties=properties, limit=limit, **options + ) def update_node( self, @@ -665,7 +694,9 @@ class GraphStore: **options, ) -> List[Dict[str, Any]]: """Get relationships.""" - return self._manager.relationships.get(node_id, rel_type, direction, limit, **options) + return self._manager.relationships.get( + node_id, rel_type, direction, limit, **options + ) def delete_relationship( self, @@ -723,7 +754,9 @@ class GraphStore: node_id, rel_type, direction, actual_depth, **options ) - def query(self, query: str, parameters: Optional[Dict[str, Any]] = None, **options) -> List[Dict[str, Any]]: + def query( + self, query: str, parameters: Optional[Dict[str, Any]] = None, **options + ) -> List[Dict[str, Any]]: """ Execute a query and return results (Compatibility method for ContextRetriever). @@ -771,42 +804,46 @@ class GraphStore: # Convert to GraphStore format (labels, properties) graph_nodes = [] for node in nodes: - # Extract label from type - labels = [node.get("type", "Entity")] - if isinstance(labels[0], str): - labels = [labels[0]] # Ensure list + # Extract labels - support both 'labels' array and 'type' string + labels = node.get("labels") + if not labels: + node_type = node.get("type", "Entity") + labels = [node_type] if isinstance(node_type, str) else node_type + if isinstance(labels, str): + labels = [labels] # Prepare properties props = node.get("properties", {}).copy() - + # Ensure ID is preserved if "id" in node and "id" not in props: props["id"] = node["id"] - + # Ensure content/text is preserved if "content" in node and "content" not in props: props["content"] = node["content"] if "text" in node and "text" not in props: props["text"] = node["text"] - graph_nodes.append({ - "labels": labels, - "properties": props - }) + graph_nodes.append({"labels": labels, "properties": props}) # Use batch creation - # Note: create_nodes expects dicts with 'labels' and 'properties' keys if passed directly? + # Note: create_nodes expects dicts with 'labels' and 'properties' + # keys if passed directly? # Let's check create_nodes signature implementation in manager. - # But here I'll assume create_nodes takes a list of such dicts or similar. + # But here I'll assume create_nodes takes a list of such dicts + # or similar. # Actually, let's look at create_nodes wrapper in this file: # def create_nodes(self, nodes: List[Dict[str, Any]], **options) # It passes to self._manager.nodes.create_batch(nodes) - + # If create_batch expects specific format, I should match it. # Assuming create_batch is smart enough or expects standard format. - # To be safe, let's look at NodeManager.create_batch if possible, but I can't easily. - # Standard expectation: List of dicts where each dict has labels and properties. - + # To be safe, let's look at NodeManager.create_batch if possible, + # but I can't easily. + # Standard expectation: List of dicts where each dict has labels + # and properties. + result = self.create_nodes(graph_nodes, **options) return len(result) @@ -827,17 +864,21 @@ class GraphStore: target_id = edge.get("target_id") rel_type = edge.get("type", "RELATED_TO") properties = edge.get("properties", {}).copy() - + # Preserve weight if "weight" in edge: properties["weight"] = edge["weight"] if source_id and target_id: try: - self.create_relationship(source_id, target_id, rel_type, properties, **options) + self.create_relationship( + source_id, target_id, rel_type, properties, **options + ) count += 1 except Exception as e: - self.logger.warning(f"Failed to add edge {source_id}->{target_id}: {e}") + self.logger.warning( + f"Failed to add edge {source_id}->{target_id}: {e}" + ) return count def build_from_conversations( @@ -872,27 +913,29 @@ class GraphStore: all_nodes = [] all_edges = [] seen_nodes = set() - + for conv in conversations: # Load conversation if string (file path) conv_data = conv if isinstance(conv, str): from pathlib import Path + from ..utils.helpers import read_json_file + conv_data = read_json_file(Path(conv)) nodes, edges = self._process_conversation_to_elements( - conv_data, + conv_data, extract_intents=extract_intents, - extract_sentiments=extract_sentiments + extract_sentiments=extract_sentiments, ) - + # Add unique nodes for node in nodes: if node["id"] not in seen_nodes: all_nodes.append(node) seen_nodes.add(node["id"]) - + all_edges.extend(edges) if link_entities: @@ -904,13 +947,8 @@ class GraphStore: edge_count = self.add_edges(all_edges) self.progress_tracker.stop_tracking(tracking_id, status="completed") - - return { - "statistics": { - "node_count": node_count, - "edge_count": edge_count - } - } + + return {"statistics": {"node_count": node_count, "edge_count": edge_count}} except Exception as e: self.progress_tracker.stop_tracking( @@ -930,92 +968,112 @@ class GraphStore: """ nodes = [] edges = [] - + # Process entities for entity in entities: entity_id = entity.get("id") or entity.get("entity_id") if entity_id: - nodes.append({ - "id": entity_id, - "type": entity.get("type", "entity"), - "properties": { - "content": entity.get("text") or entity.get("label") or entity_id, - **entity + nodes.append( + { + "id": entity_id, + "type": entity.get("type", "entity"), + "properties": { + "content": entity.get("text") + or entity.get("label") + or entity_id, + **entity, + }, } - }) + ) # Process relationships for rel in relationships: source = rel.get("source_id") target = rel.get("target_id") if source and target: - edges.append({ - "source_id": source, - "target_id": target, - "type": rel.get("type", "related_to"), - "weight": rel.get("confidence", 1.0), - "properties": rel - }) + edges.append( + { + "source_id": source, + "target_id": target, + "type": rel.get("type", "related_to"), + "weight": rel.get("confidence", 1.0), + "properties": rel, + } + ) node_count = self.add_nodes(nodes) edge_count = self.add_edges(edges) - + return {"statistics": {"node_count": node_count, "edge_count": edge_count}} - def _process_conversation_to_elements(self, conv_data: Dict[str, Any], **kwargs) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]]]: + def _process_conversation_to_elements( + self, conv_data: Dict[str, Any], **kwargs + ) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]]]: """Helper to process conversation into nodes and edges.""" nodes = [] edges = [] - + conv_id = conv_data.get("id") or f"conv_{hash(str(conv_data)) % 10000}" # Conversation node - nodes.append({ - "id": conv_id, - "type": "conversation", - "properties": { - "content": conv_data.get("content", "") or conv_data.get("summary", ""), - "timestamp": conv_data.get("timestamp") + nodes.append( + { + "id": conv_id, + "type": "conversation", + "properties": { + "content": conv_data.get("content", "") + or conv_data.get("summary", ""), + "timestamp": conv_data.get("timestamp"), + }, } - }) + ) name_to_id = {} - extract_entities = kwargs.get("extract_entities", True) # Default true if not passed? - # Actually ContextGraph defaults to True in init, but here we are static. + # Note: extract_entities option is available but not used in this + # implementation. Default true if not passed. ContextGraph defaults + # to True in init, but here we are static. # Let's assume True unless told otherwise or check config. - + # Extract entities for entity in conv_data.get("entities", []): entity_id = entity.get("id") or entity.get("entity_id") - entity_text = entity.get("text") or entity.get("label") or entity.get("name") or entity_id + entity_text = ( + entity.get("text") + or entity.get("label") + or entity.get("name") + or entity_id + ) entity_type = entity.get("type", "entity") # Generate ID if missing if not entity_id and entity_text: import hashlib - entity_hash = hashlib.md5(f"{entity_text}_{entity_type}".encode()).hexdigest()[:12] + + entity_hash = hashlib.md5( + f"{entity_text}_{entity_type}".encode() + ).hexdigest()[:12] entity_id = f"{entity_type.lower()}_{entity_hash}" if entity_id: if entity_text: name_to_id[entity_text] = entity_id - nodes.append({ - "id": entity_id, - "type": "entity", # Normalize type? - "properties": { - "content": entity_text, - "type": entity_type, - **entity + nodes.append( + { + "id": entity_id, + "type": "entity", # Normalize type? + "properties": { + "content": entity_text, + "type": entity_type, + **entity, + }, } - }) - + ) + # Edge: Conversation -> Entity - edges.append({ - "source_id": conv_id, - "target_id": entity_id, - "type": "mentions" - }) + edges.append( + {"source_id": conv_id, "target_id": entity_id, "type": "mentions"} + ) # Extract relationships for rel in conv_data.get("relationships", []): @@ -1029,43 +1087,54 @@ class GraphStore: target = name_to_id[rel.get("target")] if source and target: - edges.append({ - "source_id": source, - "target_id": target, - "type": rel.get("type", "related_to"), - "weight": rel.get("confidence", 1.0), - "properties": rel - }) - + edges.append( + { + "source_id": source, + "target_id": target, + "type": rel.get("type", "related_to"), + "weight": rel.get("confidence", 1.0), + "properties": rel, + } + ) + return nodes, edges - def _link_entities_elements(self, nodes: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + def _link_entities_elements( + self, nodes: List[Dict[str, Any]] + ) -> List[Dict[str, Any]]: """Link similar entities.""" edges = [] # Lazy import to avoid circular dependency try: from ..context.entity_linker import EntityLinker - linker = EntityLinker() # Use default config + + linker = EntityLinker() # Use default config except (ImportError, OSError): return [] entity_nodes = [n for n in nodes if n.get("type") == "entity"] for i, node1 in enumerate(entity_nodes): content1 = node1["properties"].get("content", "") - if not content1: continue - + if not content1: + continue + for node2 in entity_nodes[i + 1 :]: content2 = node2["properties"].get("content", "") - if not content2: continue - - similarity = linker._calculate_text_similarity(content1.lower(), content2.lower()) + if not content2: + continue + + similarity = linker._calculate_text_similarity( + content1.lower(), content2.lower() + ) if similarity >= linker.similarity_threshold: - edges.append({ - "source_id": node1["id"], - "target_id": node2["id"], - "type": "similar_to", - "weight": similarity - }) + edges.append( + { + "source_id": node1["id"], + "target_id": node2["id"], + "type": "similar_to", + "weight": similarity, + } + ) return edges @property @@ -1087,4 +1156,3 @@ class GraphStore: def analytics(self) -> GraphAnalytics: """Get analytics engine.""" return self._manager.analytics - diff --git a/tests/test_amazon_neptune.py b/tests/test_amazon_neptune.py new file mode 100644 index 00000000..40295aad --- /dev/null +++ b/tests/test_amazon_neptune.py @@ -0,0 +1,1702 @@ +""" +Unit tests for Amazon Neptune Graph Store with Neo4j Bolt driver. + +Tests cover: +- Initialization with and without IAM authentication +- Connection management using Bolt driver +- Node CRUD operations +- Relationship CRUD operations +- OpenCypher query execution +- Graph analytics (neighbors, shortest path) +- Statistics and status +- NeptuneAuthTokenManager SigV4 signing +- Retry logic and connection recovery +""" + +import unittest +from typing import Any, Dict, List, Optional +from unittest.mock import MagicMock, patch + + +class MockNeo4jRecord: + """Mock Neo4j record for testing.""" + + def __init__(self, data: Dict[str, Any]): + self._data = data + + def __iter__(self): + return iter(self._data.items()) + + def __getitem__(self, key): + return self._data[key] + + def keys(self): + return self._data.keys() + + def values(self): + return self._data.values() + + def items(self): + return self._data.items() + + def get(self, key, default=None): + return self._data.get(key, default) + + +class MockNeo4jResult: + """Mock Neo4j query result.""" + + def __init__(self, records: List[Dict[str, Any]]): + self._records = [MockNeo4jRecord(r) for r in records] + self._index = 0 + + def __iter__(self): + return iter(self._records) + + def single(self): + return self._records[0] if self._records else None + + def peek(self): + return self._records[0] if self._records else None + + +class MockNeo4jSession: + """Mock Neo4j session for testing.""" + + def __init__(self, driver): + self.driver = driver + self._closed = False + + def run(self, query: str, parameters: Optional[Dict] = None) -> MockNeo4jResult: + """Execute a query and return mock results.""" + parameters = parameters or {} + + # Store query for inspection + self.driver.last_query = query + self.driver.last_parameters = parameters + + query_lower = query.lower() + + # Simple RETURN 1 test query + if "return 1" in query_lower: + return MockNeo4jResult([{"test": 1}]) + + # CREATE node + if "create" in query_lower and "(n:" in query_lower: + return self._handle_create_node(query, parameters) + + # MERGE node + if "merge" in query_lower and "(n:" in query_lower: + return self._handle_create_node(query, parameters) + + # MATCH node by id + if "match" in query_lower and "id(n)" in query_lower: + return self._handle_get_node(query, parameters) + + # MATCH nodes by label + if "match" in query_lower and "return n" in query_lower: + return self._handle_get_nodes(query, parameters) + + # CREATE relationship + if "create" in query_lower and "-[r:" in query_lower: + return self._handle_create_relationship(query, parameters) + + # MATCH relationship + if "match" in query_lower and "-[r" in query_lower: + return self._handle_get_relationships(query, parameters) + + # DELETE + if "delete" in query_lower: + return self._handle_delete(query, parameters) + + # SET (update) + if "set" in query_lower: + return self._handle_update(query, parameters) + + # COUNT queries (stats) + if "count" in query_lower: + return self._handle_stats(query) + + # Default response + return MockNeo4jResult([{"n": {}}]) + + def _handle_create_node(self, query: str, parameters: Dict) -> MockNeo4jResult: + """Handle node creation queries.""" + self.driver.node_counter += 1 + node_id = parameters.get("node_id", f"node_{self.driver.node_counter}") + labels = ["TestLabel"] + properties = {k: v for k, v in parameters.items() if k != "node_id"} + + node = {"id": node_id, "labels": labels, "properties": properties} + self.driver.nodes[node_id] = node + + return MockNeo4jResult( + [{"n": {"~id": node_id, "~labels": labels, **properties}}] + ) + + def _handle_get_node(self, query: str, parameters: Dict) -> MockNeo4jResult: + """Handle get node by ID queries.""" + node_id = parameters.get("id") + if node_id and node_id in self.driver.nodes: + node = self.driver.nodes[node_id] + return MockNeo4jResult( + [ + { + "n": { + "~id": node["id"], + "~labels": node["labels"], + **node["properties"], + }, + "labels": node["labels"], + } + ] + ) + return MockNeo4jResult([]) + + def _handle_get_nodes(self, query: str, parameters: Dict) -> MockNeo4jResult: + """Handle get nodes queries.""" + results = [] + for node in self.driver.nodes.values(): + results.append( + { + "n": { + "~id": node["id"], + "~labels": node["labels"], + **node["properties"], + }, + "labels": node["labels"], + } + ) + return MockNeo4jResult(results) + + def _handle_create_relationship( + self, query: str, parameters: Dict + ) -> MockNeo4jResult: + """Handle relationship creation queries.""" + self.driver.rel_counter += 1 + rel_id = parameters.get("rel_id", f"rel_{self.driver.rel_counter}") + + rel = { + "id": rel_id, + "start_node_id": parameters.get("start_id"), + "end_node_id": parameters.get("end_id"), + "type": "TEST_REL", + "properties": {}, + } + self.driver.relationships[rel_id] = rel + + return MockNeo4jResult([{"r": {"~id": rel_id, "~type": "TEST_REL"}}]) + + def _handle_get_relationships( + self, query: str, parameters: Dict + ) -> MockNeo4jResult: + """Handle get relationships queries.""" + results = [] + for rel in self.driver.relationships.values(): + results.append( + { + "r": {"~id": rel["id"], "~type": rel["type"]}, + "start_id": rel["start_node_id"], + "end_id": rel["end_node_id"], + } + ) + return MockNeo4jResult(results) + + def _handle_delete(self, query: str, parameters: Dict) -> MockNeo4jResult: + """Handle delete queries.""" + node_id = parameters.get("id") + if node_id and node_id in self.driver.nodes: + del self.driver.nodes[node_id] + return MockNeo4jResult([]) + + def _handle_update(self, query: str, parameters: Dict) -> MockNeo4jResult: + """Handle update queries.""" + node_id = parameters.get("id") + if node_id and node_id in self.driver.nodes: + props = parameters.get("props", {}) + self.driver.nodes[node_id]["properties"].update(props) + node = self.driver.nodes[node_id] + return MockNeo4jResult( + [ + { + "n": { + "~id": node["id"], + "~labels": node["labels"], + **node["properties"], + }, + "labels": node["labels"], + } + ] + ) + return MockNeo4jResult([]) + + def _handle_stats(self, query: str) -> MockNeo4jResult: + """Handle statistics queries.""" + return MockNeo4jResult( + [{"count": len(self.driver.nodes) + len(self.driver.relationships)}] + ) + + def close(self): + """Close the session.""" + self._closed = True + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + self.close() + + +class MockNeo4jDriver: + """Mock Neo4j driver for testing.""" + + def __init__(self, *args, **kwargs): + self.nodes = {} + self.relationships = {} + self.node_counter = 0 + self.rel_counter = 0 + self.last_query = None + self.last_parameters = None + self._closed = False + + def session(self, **kwargs): + """Create a new session.""" + return MockNeo4jSession(self) + + def close(self): + """Close the driver.""" + self._closed = True + + def verify_connectivity(self): + """Verify driver connectivity.""" + pass + + +class MockGraphDatabase: + """Mock Neo4j GraphDatabase class.""" + + @staticmethod + def driver(uri, auth=None, **kwargs): + return MockNeo4jDriver(uri, auth=auth, **kwargs) + + +class MockBasicAuth: + """Mock basic_auth function.""" + + def __init__(self, username, password): + self.username = username + self.password = password + + +def mock_basic_auth(username, password): + """Mock basic_auth function.""" + return MockBasicAuth(username, password) + + +# Create a real AuthManager base class for testing (not MagicMock) +class MockAuthManager: + """Mock AuthManager base class for testing.""" + + pass + + +# Create mock modules +mock_neo4j = MagicMock() +mock_neo4j.GraphDatabase = MockGraphDatabase +mock_neo4j.auth_management = MagicMock() +mock_neo4j.auth_management.AuthManager = MockAuthManager +mock_neo4j.api = MagicMock() +mock_neo4j.api.basic_auth = mock_basic_auth +mock_neo4j.exceptions = MagicMock() +mock_neo4j.exceptions.Neo4jError = Exception +mock_neo4j.exceptions.ServiceUnavailable = Exception +mock_neo4j.exceptions.DatabaseError = Exception +mock_neo4j.exceptions.ClientError = Exception +mock_neo4j.exceptions.AuthError = Exception + +mock_boto3 = MagicMock() +mock_session = MagicMock() +mock_credentials = MagicMock() +mock_credentials.access_key = "test_access_key" +mock_credentials.secret_key = "test_secret_key" +mock_credentials.token = None +mock_session.get_credentials.return_value = mock_credentials +mock_boto3.Session.return_value = mock_session + + +class MockHeaders: + """Mock headers object that supports both dict-like access and add_header.""" + + def __init__(self): + self._headers = {} + + def add_header(self, key, value): + self._headers[key] = value + + def get(self, key, default=None): + return self._headers.get(key, default) + + def __setitem__(self, key, value): + self._headers[key] = value + + def __getitem__(self, key): + return self._headers[key] + + def __contains__(self, key): + return key in self._headers + + +class MockAWSRequest: + """Mock AWSRequest for testing botocore SigV4 signing.""" + + def __init__(self, method="GET", url="", data=None, headers=None): + self.method = method + self.url = url + self.data = data + self.headers = MockHeaders() + if headers: + for k, v in headers.items(): + self.headers[k] = v + + +class MockSigV4Auth: + """Mock SigV4Auth for testing botocore signing.""" + + def __init__(self, credentials, service_name, region): + self.credentials = credentials + self.service_name = service_name + self.region = region + + def add_auth(self, request): + """Add mock SigV4 authentication headers to the request.""" + from datetime import datetime, timezone + + # Generate a timestamp + t = datetime.now(timezone.utc) + amz_date = t.strftime("%Y%m%dT%H%M%SZ") + date_stamp = t.strftime("%Y%m%d") + + # Extract host from URL + import urllib.parse + + parsed = urllib.parse.urlparse(request.url) + host = parsed.netloc + + # Build mock authorization header + credential_scope = ( + f"{date_stamp}/{self.region}/{self.service_name}/aws4_request" + ) + authorization = ( + f"AWS4-HMAC-SHA256 " + f"Credential=test_access_key/{credential_scope}, " + f"SignedHeaders=content-type;host;x-amz-date, " + f"Signature=mocksignature123456789" + ) + + request.headers["Host"] = host + request.headers["X-Amz-Date"] = amz_date + request.headers["Authorization"] = authorization + + +def mock_host_from_url(url): + """Mock _host_from_url function.""" + import urllib.parse + + parsed = urllib.parse.urlparse(url) + return parsed.netloc + + +# Create mock botocore modules +mock_botocore = MagicMock() +mock_botocore_auth = MagicMock() +mock_botocore_auth.SigV4Auth = MockSigV4Auth +mock_botocore_auth._host_from_url = mock_host_from_url +mock_botocore_awsrequest = MagicMock() +mock_botocore_awsrequest.AWSRequest = MockAWSRequest + + +class TestAmazonNeptuneStoreInit(unittest.TestCase): + """Test AmazonNeptuneStore initialization.""" + + def test_init_with_iam_auth(self): + """Test initialization with IAM authentication enabled.""" + with patch.dict( + "sys.modules", + { + "neo4j": mock_neo4j, + "neo4j.auth_management": mock_neo4j.auth_management, + "neo4j.api": mock_neo4j.api, + "neo4j.exceptions": mock_neo4j.exceptions, + "boto3": mock_boto3, + "botocore": mock_botocore, + "botocore.auth": mock_botocore_auth, + "botocore.awsrequest": mock_botocore_awsrequest, + }, + ): + from importlib import reload + + from semantica.graph_store import amazon_neptune + + reload(amazon_neptune) + + store = amazon_neptune.AmazonNeptuneStore( + endpoint="test-cluster.us-east-1.neptune.amazonaws.com", + port=8182, + region="us-east-1", + iam_auth=True, + ) + + self.assertEqual( + store.endpoint, "test-cluster.us-east-1.neptune.amazonaws.com" + ) + self.assertEqual(store.port, 8182) + self.assertEqual(store.region, "us-east-1") + self.assertTrue(store.iam_auth) + + def test_init_without_iam_auth(self): + """Test initialization without IAM authentication.""" + with patch.dict( + "sys.modules", + { + "neo4j": mock_neo4j, + "neo4j.auth_management": mock_neo4j.auth_management, + "neo4j.api": mock_neo4j.api, + "neo4j.exceptions": mock_neo4j.exceptions, + "boto3": mock_boto3, + }, + ): + from importlib import reload + + from semantica.graph_store import amazon_neptune + + reload(amazon_neptune) + + store = amazon_neptune.AmazonNeptuneStore( + endpoint="test-cluster.us-east-1.neptune.amazonaws.com", + port=8182, + region="us-east-1", + iam_auth=False, + ) + + self.assertFalse(store.iam_auth) + + def test_init_default_port(self): + """Test initialization with default port.""" + with patch.dict( + "sys.modules", + { + "neo4j": mock_neo4j, + "neo4j.auth_management": mock_neo4j.auth_management, + "neo4j.api": mock_neo4j.api, + "neo4j.exceptions": mock_neo4j.exceptions, + "boto3": mock_boto3, + }, + ): + from importlib import reload + + from semantica.graph_store import amazon_neptune + + reload(amazon_neptune) + + store = amazon_neptune.AmazonNeptuneStore( + endpoint="test-cluster.us-east-1.neptune.amazonaws.com", + region="us-east-1", + ) + + self.assertEqual(store.port, 8182) + + def test_bolt_uri_property(self): + """Test the bolt_uri property.""" + with patch.dict( + "sys.modules", + { + "neo4j": mock_neo4j, + "neo4j.auth_management": mock_neo4j.auth_management, + "neo4j.api": mock_neo4j.api, + "neo4j.exceptions": mock_neo4j.exceptions, + "boto3": mock_boto3, + }, + ): + from importlib import reload + + from semantica.graph_store import amazon_neptune + + reload(amazon_neptune) + + store = amazon_neptune.AmazonNeptuneStore( + endpoint="test.neptune.amazonaws.com", + port=8182, + region="us-east-1", + iam_auth=False, + ) + + self.assertEqual( + store.bolt_uri, "bolt://test.neptune.amazonaws.com:8182/opencypher" + ) + + +class TestAmazonNeptuneStoreOperations(unittest.TestCase): + """Test AmazonNeptuneStore CRUD operations.""" + + def setUp(self): + """Set up test fixtures.""" + self.modules_patcher = patch.dict( + "sys.modules", + { + "neo4j": mock_neo4j, + "neo4j.auth_management": mock_neo4j.auth_management, + "neo4j.api": mock_neo4j.api, + "neo4j.exceptions": mock_neo4j.exceptions, + "boto3": mock_boto3, + "botocore": mock_botocore, + "botocore.auth": mock_botocore_auth, + "botocore.awsrequest": mock_botocore_awsrequest, + }, + ) + self.modules_patcher.start() + + from importlib import reload + + from semantica.graph_store import amazon_neptune + + reload(amazon_neptune) + + self.store = amazon_neptune.AmazonNeptuneStore( + endpoint="test-cluster.us-east-1.neptune.amazonaws.com", + port=8182, + region="us-east-1", + iam_auth=False, + ) + self.store.connect() + + def tearDown(self): + """Clean up test fixtures.""" + if hasattr(self, "store"): + self.store.close() + self.modules_patcher.stop() + + def test_connect_success(self): + """Test successful connection.""" + self.assertTrue(self.store._connected) + + def test_create_node_with_custom_id(self): + """Test creating a node with custom ID (id in properties).""" + node = self.store.create_node( + labels=["Person"], properties={"id": "alice", "name": "Alice", "age": 30} + ) + + self.assertIsNotNone(node) + self.assertEqual(node.get("id"), "alice") + + def test_create_node_with_generated_id(self): + """Test creating a node with auto-generated UUID (no id in properties).""" + node = self.store.create_node(labels=["Person"], properties={"name": "Bob"}) + + self.assertIsNotNone(node) + self.assertIn("id", node) + + def test_get_nodes(self): + """Test retrieving nodes by label.""" + self.store.create_node( + labels=["Person"], properties={"id": "alice", "name": "Alice"} + ) + + nodes = self.store.get_nodes(labels=["Person"], limit=10) + self.assertIsInstance(nodes, list) + + def test_execute_query(self): + """Test executing an OpenCypher query.""" + result = self.store.execute_query("MATCH (n) RETURN n LIMIT 10", parameters={}) + + self.assertIsNotNone(result) + self.assertIn("records", result) + + def test_execute_query_with_parameters(self): + """Test executing a parameterized query.""" + result = self.store.execute_query( + "MATCH (p:Person) WHERE p.age > $min_age RETURN p", + parameters={"min_age": 25}, + ) + + self.assertIsNotNone(result) + + def test_get_stats(self): + """Test getting graph statistics.""" + stats = self.store.get_stats() + + self.assertIsNotNone(stats) + self.assertIn("node_count", stats) + self.assertEqual(stats["protocol"], "bolt") + + def test_get_status(self): + """Test getting connection status.""" + status = self.store.get_status() + + self.assertIsNotNone(status) + self.assertEqual(status["backend"], "amazon_neptune") + self.assertEqual(status["protocol"], "bolt") + self.assertEqual(status["connected"], True) + + def test_close_connection(self): + """Test closing the connection.""" + self.store.close() + self.assertFalse(self.store._connected) + + +class TestAmazonNeptuneStoreIAMAuth(unittest.TestCase): + """Test IAM authentication functionality.""" + + def test_auth_manager_creation(self): + """Test that AuthManager is created correctly for IAM auth.""" + with patch.dict( + "sys.modules", + { + "neo4j": mock_neo4j, + "neo4j.auth_management": mock_neo4j.auth_management, + "neo4j.api": mock_neo4j.api, + "neo4j.exceptions": mock_neo4j.exceptions, + "boto3": mock_boto3, + "botocore": mock_botocore, + "botocore.auth": mock_botocore_auth, + "botocore.awsrequest": mock_botocore_awsrequest, + }, + ): + from importlib import reload + + from semantica.graph_store import amazon_neptune + + reload(amazon_neptune) + + store = amazon_neptune.AmazonNeptuneStore( + endpoint="test-cluster.us-east-1.neptune.amazonaws.com", + port=8182, + region="us-east-1", + iam_auth=True, + ) + store.connect() + + # Verify auth manager was created + self.assertIsNotNone(store._auth_manager) + self.assertEqual(store._auth_manager.aws_region, "us-east-1") + + store.close() + + def test_refresh_auth_token(self): + """Test manual token refresh.""" + with patch.dict( + "sys.modules", + { + "neo4j": mock_neo4j, + "neo4j.auth_management": mock_neo4j.auth_management, + "neo4j.api": mock_neo4j.api, + "neo4j.exceptions": mock_neo4j.exceptions, + "boto3": mock_boto3, + "botocore": mock_botocore, + "botocore.auth": mock_botocore_auth, + "botocore.awsrequest": mock_botocore_awsrequest, + }, + ): + from importlib import reload + + from semantica.graph_store import amazon_neptune + + reload(amazon_neptune) + + store = amazon_neptune.AmazonNeptuneStore( + endpoint="test-cluster.us-east-1.neptune.amazonaws.com", + port=8182, + region="us-east-1", + iam_auth=True, + ) + store.connect() + + # Refresh should not raise + store.refresh_auth_token() + + store.close() + + def test_update_auth_context(self): + """Test updating auth context (for Lambda).""" + with patch.dict( + "sys.modules", + { + "neo4j": mock_neo4j, + "neo4j.auth_management": mock_neo4j.auth_management, + "neo4j.api": mock_neo4j.api, + "neo4j.exceptions": mock_neo4j.exceptions, + "boto3": mock_boto3, + "botocore": mock_botocore, + "botocore.auth": mock_botocore_auth, + "botocore.awsrequest": mock_botocore_awsrequest, + }, + ): + from importlib import reload + + from semantica.graph_store import amazon_neptune + + reload(amazon_neptune) + + store = amazon_neptune.AmazonNeptuneStore( + endpoint="test-cluster.us-east-1.neptune.amazonaws.com", + port=8182, + region="us-east-1", + iam_auth=True, + ) + store.connect() + + # Mock Lambda context + mock_context = MagicMock() + # update_auth_context should not raise - it logs but doesn't store + # (Neo4j driver can't serialize complex objects as instance attributes) + store.update_auth_context(mock_context) + + # The method should complete without error + self.assertIsNotNone(store._auth_manager) + + store.close() + + +class TestAmazonNeptuneStoreGraphAnalytics(unittest.TestCase): + """Test graph analytics methods.""" + + def setUp(self): + """Set up test fixtures.""" + self.modules_patcher = patch.dict( + "sys.modules", + { + "neo4j": mock_neo4j, + "neo4j.auth_management": mock_neo4j.auth_management, + "neo4j.api": mock_neo4j.api, + "neo4j.exceptions": mock_neo4j.exceptions, + "boto3": mock_boto3, + }, + ) + self.modules_patcher.start() + + from importlib import reload + + from semantica.graph_store import amazon_neptune + + reload(amazon_neptune) + + self.store = amazon_neptune.AmazonNeptuneStore( + endpoint="test-cluster.us-east-1.neptune.amazonaws.com", + region="us-east-1", + iam_auth=False, + ) + self.store.connect() + + def tearDown(self): + """Clean up test fixtures.""" + if hasattr(self, "store"): + self.store.close() + self.modules_patcher.stop() + + def test_get_neighbors(self): + """Test getting neighbors of a node.""" + self.store.create_node( + labels=["Person"], properties={"id": "alice", "name": "Alice"} + ) + + neighbors = self.store.get_neighbors(node_id="alice", direction="both", depth=1) + + self.assertIsInstance(neighbors, list) + + def test_shortest_path(self): + """Test finding shortest path between nodes.""" + self.store.create_node( + labels=["Person"], properties={"id": "alice", "name": "Alice"} + ) + self.store.create_node( + labels=["Person"], properties={"id": "bob", "name": "Bob"} + ) + + path = self.store.shortest_path( + start_node_id="alice", end_node_id="bob", max_depth=5 + ) + + # Path may be None if no path exists, which is valid + self.assertTrue(path is None or isinstance(path, dict)) + + +class TestGraphStoreNeptuneBackend(unittest.TestCase): + """Test GraphStore with Neptune backend.""" + + def test_graphstore_neptune_initialization(self): + """Test GraphStore initialization with Neptune backend.""" + with patch.dict( + "sys.modules", + { + "neo4j": mock_neo4j, + "neo4j.auth_management": mock_neo4j.auth_management, + "neo4j.api": mock_neo4j.api, + "neo4j.exceptions": mock_neo4j.exceptions, + "boto3": mock_boto3, + }, + ): + from importlib import reload + + from semantica.graph_store import amazon_neptune, graph_store + + reload(amazon_neptune) + reload(graph_store) + + store = graph_store.GraphStore( + backend="neptune", + endpoint="test-cluster.us-east-1.neptune.amazonaws.com", + region="us-east-1", + iam_auth=False, + ) + + self.assertEqual(store.backend, "neptune") + + +class TestAmazonNeptuneStoreCRUD(unittest.TestCase): + """Test complete CRUD operations including update and delete.""" + + def setUp(self): + """Set up test fixtures.""" + self.modules_patcher = patch.dict( + "sys.modules", + { + "neo4j": mock_neo4j, + "neo4j.auth_management": mock_neo4j.auth_management, + "neo4j.api": mock_neo4j.api, + "neo4j.exceptions": mock_neo4j.exceptions, + "boto3": mock_boto3, + }, + ) + self.modules_patcher.start() + + from importlib import reload + + from semantica.graph_store import amazon_neptune + + reload(amazon_neptune) + + self.store = amazon_neptune.AmazonNeptuneStore( + endpoint="test-cluster.us-east-1.neptune.amazonaws.com", + port=8182, + region="us-east-1", + iam_auth=False, + ) + self.store.connect() + + def tearDown(self): + """Clean up test fixtures.""" + if hasattr(self, "store"): + self.store.close() + self.modules_patcher.stop() + + def test_update_node(self): + """Test updating a node's properties.""" + self.store.create_node( + labels=["Person"], + properties={"id": "update_test", "name": "Original", "age": 25}, + ) + + updated = self.store.update_node( + node_id="update_test", + properties={"name": "Updated", "city": "NYC"}, + merge=True, + ) + + self.assertIsNotNone(updated) + + def test_update_node_replace(self): + """Test replacing a node's properties (merge=False).""" + self.store.create_node( + labels=["Person"], + properties={"id": "replace_test", "name": "Original", "age": 25}, + ) + + updated = self.store.update_node( + node_id="replace_test", properties={"name": "Replaced"}, merge=False + ) + + self.assertIsNotNone(updated) + + def test_delete_node(self): + """Test deleting a node.""" + self.store.create_node( + labels=["Person"], properties={"id": "delete_test", "name": "ToDelete"} + ) + + result = self.store.delete_node(node_id="delete_test", detach=True) + self.assertTrue(result) + + def test_delete_node_without_detach(self): + """Test deleting a node without detaching relationships.""" + self.store.create_node( + labels=["Person"], properties={"id": "delete_test_2", "name": "ToDelete"} + ) + + result = self.store.delete_node(node_id="delete_test_2", detach=False) + self.assertTrue(result) + + def test_get_node_by_id(self): + """Test getting a specific node by ID.""" + self.store.create_node( + labels=["Person"], properties={"id": "get_test", "name": "GetMe"} + ) + + node = self.store.get_node(node_id="get_test") + self.assertTrue(node is None or isinstance(node, dict)) + + +class TestAmazonNeptuneStoreRelationships(unittest.TestCase): + """Test relationship CRUD operations.""" + + def setUp(self): + """Set up test fixtures.""" + self.modules_patcher = patch.dict( + "sys.modules", + { + "neo4j": mock_neo4j, + "neo4j.auth_management": mock_neo4j.auth_management, + "neo4j.api": mock_neo4j.api, + "neo4j.exceptions": mock_neo4j.exceptions, + "boto3": mock_boto3, + }, + ) + self.modules_patcher.start() + + from importlib import reload + + from semantica.graph_store import amazon_neptune + + reload(amazon_neptune) + + self.store = amazon_neptune.AmazonNeptuneStore( + endpoint="test-cluster.us-east-1.neptune.amazonaws.com", + port=8182, + region="us-east-1", + iam_auth=False, + ) + self.store.connect() + + # Create test nodes for relationship tests + self.store.create_node( + labels=["Person"], properties={"id": "rel_alice", "name": "Alice"} + ) + self.store.create_node( + labels=["Person"], properties={"id": "rel_bob", "name": "Bob"} + ) + + def tearDown(self): + """Clean up test fixtures.""" + if hasattr(self, "store"): + self.store.close() + self.modules_patcher.stop() + + def test_create_relationship(self): + """Test creating a relationship between nodes.""" + rel = self.store.create_relationship( + start_node_id="rel_alice", + end_node_id="rel_bob", + rel_type="KNOWS", + properties={"since": 2020}, + ) + + self.assertIsNotNone(rel) + self.assertEqual(rel["type"], "KNOWS") + self.assertEqual(rel["start_node_id"], "rel_alice") + self.assertEqual(rel["end_node_id"], "rel_bob") + + def test_create_relationship_with_custom_id(self): + """Test creating a relationship with a custom ID.""" + rel = self.store.create_relationship( + start_node_id="rel_alice", + end_node_id="rel_bob", + rel_type="FRIENDS", + properties={"id": "custom_rel_id", "level": "close"}, + ) + + self.assertIsNotNone(rel) + self.assertEqual(rel["id"], "custom_rel_id") + + def test_get_relationships_outgoing(self): + """Test getting outgoing relationships from a node.""" + self.store.create_relationship( + start_node_id="rel_alice", end_node_id="rel_bob", rel_type="KNOWS" + ) + + rels = self.store.get_relationships(node_id="rel_alice", direction="out") + + self.assertIsInstance(rels, list) + + def test_get_relationships_incoming(self): + """Test getting incoming relationships to a node.""" + self.store.create_relationship( + start_node_id="rel_alice", end_node_id="rel_bob", rel_type="KNOWS" + ) + + rels = self.store.get_relationships(node_id="rel_bob", direction="in") + + self.assertIsInstance(rels, list) + + def test_get_relationships_both(self): + """Test getting all relationships for a node.""" + rels = self.store.get_relationships(node_id="rel_alice", direction="both") + + self.assertIsInstance(rels, list) + + def test_get_relationships_by_type(self): + """Test filtering relationships by type.""" + rels = self.store.get_relationships(node_id="rel_alice", rel_type="KNOWS") + + self.assertIsInstance(rels, list) + + def test_get_all_relationships(self): + """Test getting all relationships without node filter.""" + rels = self.store.get_relationships(limit=10) + self.assertIsInstance(rels, list) + + def test_delete_relationship(self): + """Test deleting a relationship.""" + self.store.create_relationship( + start_node_id="rel_alice", + end_node_id="rel_bob", + rel_type="TEMPORARY", + properties={"id": "rel_to_delete"}, + ) + + result = self.store.delete_relationship(rel_id="rel_to_delete") + self.assertTrue(result) + + +class TestAmazonNeptuneStoreErrors(unittest.TestCase): + """Test error handling and edge cases.""" + + def setUp(self): + """Set up test fixtures.""" + self.modules_patcher = patch.dict( + "sys.modules", + { + "neo4j": mock_neo4j, + "neo4j.auth_management": mock_neo4j.auth_management, + "neo4j.api": mock_neo4j.api, + "neo4j.exceptions": mock_neo4j.exceptions, + "boto3": mock_boto3, + }, + ) + self.modules_patcher.start() + + from importlib import reload + + from semantica.graph_store import amazon_neptune + + reload(amazon_neptune) + self.amazon_neptune = amazon_neptune + + def tearDown(self): + """Clean up test fixtures.""" + self.modules_patcher.stop() + + def test_missing_endpoint(self): + """Test initialization with missing endpoint.""" + store = self.amazon_neptune.AmazonNeptuneStore( + endpoint=None, region="us-east-1", iam_auth=False + ) + # Connection should fail due to missing endpoint + with self.assertRaises(Exception): + store.connect() + + def test_create_node_without_labels(self): + """Test creating a node with empty labels.""" + store = self.amazon_neptune.AmazonNeptuneStore( + endpoint="test-cluster.us-east-1.neptune.amazonaws.com", + region="us-east-1", + iam_auth=False, + ) + store.connect() + + try: + # Empty labels should still work - defaults to "Node" + created_node = store.create_node(labels=[], properties={"name": "NoLabel"}) + self.assertIsNotNone(created_node) + finally: + store.close() + + def test_create_node_with_merge_false(self): + """Test creating a node with merge=False option.""" + store = self.amazon_neptune.AmazonNeptuneStore( + endpoint="test-cluster.us-east-1.neptune.amazonaws.com", + region="us-east-1", + iam_auth=False, + ) + store.connect() + + try: + created_node = store.create_node( + labels=["Person"], + properties={"name": "MergeFalse"}, + merge=False, + ) + self.assertIsNotNone(created_node) + finally: + store.close() + + def test_get_nodes_with_property_filter(self): + """Test getting nodes with property filters.""" + store = self.amazon_neptune.AmazonNeptuneStore( + endpoint="test-cluster.us-east-1.neptune.amazonaws.com", + region="us-east-1", + iam_auth=False, + ) + store.connect() + + try: + nodes = store.get_nodes( + labels=["Person"], properties={"name": "Alice", "age": 30}, limit=5 + ) + self.assertIsInstance(nodes, list) + finally: + store.close() + + def test_create_index(self): + """Test create_index (Neptune handles indexes automatically).""" + store = self.amazon_neptune.AmazonNeptuneStore( + endpoint="test-cluster.us-east-1.neptune.amazonaws.com", + region="us-east-1", + iam_auth=False, + ) + + # Should return True and log info about Neptune's index management + result = store.create_index(label="Person", property_name="name") + self.assertTrue(result) + + +class TestNeptuneWrapperClasses(unittest.TestCase): + """Test Neptune wrapper classes.""" + + def setUp(self): + """Set up test fixtures.""" + self.modules_patcher = patch.dict( + "sys.modules", + { + "neo4j": mock_neo4j, + "neo4j.auth_management": mock_neo4j.auth_management, + "neo4j.api": mock_neo4j.api, + "neo4j.exceptions": mock_neo4j.exceptions, + "boto3": mock_boto3, + }, + ) + self.modules_patcher.start() + + from importlib import reload + + from semantica.graph_store import amazon_neptune + + reload(amazon_neptune) + self.amazon_neptune = amazon_neptune + + def tearDown(self): + """Clean up test fixtures.""" + self.modules_patcher.stop() + + def test_neptune_driver_session(self): + """Test NeptuneDriver.session() method.""" + mock_driver = MockNeo4jDriver() + wrapper = self.amazon_neptune.NeptuneDriver(mock_driver) + + session = wrapper.session() + self.assertIsInstance(session, self.amazon_neptune.NeptuneSession) + + def test_neptune_driver_verify_connectivity(self): + """Test NeptuneDriver.verify_connectivity() method.""" + mock_driver = MockNeo4jDriver() + wrapper = self.amazon_neptune.NeptuneDriver(mock_driver) + + result = wrapper.verify_connectivity() + self.assertTrue(result) + + def test_neptune_driver_close(self): + """Test NeptuneDriver.close() method.""" + mock_driver = MockNeo4jDriver() + wrapper = self.amazon_neptune.NeptuneDriver(mock_driver) + + wrapper.close() + self.assertTrue(mock_driver._closed) + + def test_neptune_session_run(self): + """Test NeptuneSession.run() method.""" + mock_driver = MockNeo4jDriver() + mock_session = MockNeo4jSession(mock_driver) + wrapper = self.amazon_neptune.NeptuneSession(mock_session) + + result = wrapper.run("RETURN 1 as test") + self.assertIsNotNone(result) + + def test_neptune_session_close(self): + """Test NeptuneSession.close() method.""" + mock_driver = MockNeo4jDriver() + mock_session = MockNeo4jSession(mock_driver) + wrapper = self.amazon_neptune.NeptuneSession(mock_session) + + wrapper.close() + self.assertTrue(mock_session._closed) + + def test_neptune_session_context_manager(self): + """Test NeptuneSession context manager.""" + mock_driver = MockNeo4jDriver() + mock_session = MockNeo4jSession(mock_driver) + wrapper = self.amazon_neptune.NeptuneSession(mock_session) + + with wrapper as session: + self.assertIsNotNone(session) + self.assertTrue(mock_session._closed) + + +class TestAmazonNeptuneStoreBatchOperations(unittest.TestCase): + """Test batch operations.""" + + def setUp(self): + """Set up test fixtures.""" + self.modules_patcher = patch.dict( + "sys.modules", + { + "neo4j": mock_neo4j, + "neo4j.auth_management": mock_neo4j.auth_management, + "neo4j.api": mock_neo4j.api, + "neo4j.exceptions": mock_neo4j.exceptions, + "boto3": mock_boto3, + }, + ) + self.modules_patcher.start() + + from importlib import reload + + from semantica.graph_store import amazon_neptune + + reload(amazon_neptune) + + self.store = amazon_neptune.AmazonNeptuneStore( + endpoint="test-cluster.us-east-1.neptune.amazonaws.com", + port=8182, + region="us-east-1", + iam_auth=False, + ) + self.store.connect() + + def tearDown(self): + """Clean up test fixtures.""" + if hasattr(self, "store"): + self.store.close() + self.modules_patcher.stop() + + def test_create_nodes_batch(self): + """Test creating multiple nodes in batch.""" + nodes_data = [ + {"labels": ["Person"], "properties": {"id": "batch1", "name": "Alice"}}, + {"labels": ["Person"], "properties": {"id": "batch2", "name": "Bob"}}, + {"labels": ["Company"], "properties": {"id": "batch3", "name": "Acme"}}, + ] + + created = self.store.create_nodes(nodes_data) + + self.assertEqual(len(created), 3) + self.assertEqual(created[0]["id"], "batch1") + self.assertEqual(created[1]["id"], "batch2") + self.assertEqual(created[2]["id"], "batch3") + + def test_create_nodes_empty(self): + """Test creating empty batch.""" + created = self.store.create_nodes([]) + self.assertEqual(len(created), 0) + + +class TestAmazonNeptuneStoreRetryLogic(unittest.TestCase): + """Test retry logic and error handling.""" + + def setUp(self): + """Set up test fixtures.""" + self.modules_patcher = patch.dict( + "sys.modules", + { + "neo4j": mock_neo4j, + "neo4j.auth_management": mock_neo4j.auth_management, + "neo4j.api": mock_neo4j.api, + "neo4j.exceptions": mock_neo4j.exceptions, + "boto3": mock_boto3, + }, + ) + self.modules_patcher.start() + + from importlib import reload + + from semantica.graph_store import amazon_neptune + + reload(amazon_neptune) + self.amazon_neptune = amazon_neptune + + self.store = amazon_neptune.AmazonNeptuneStore( + endpoint="test-cluster.us-east-1.neptune.amazonaws.com", + port=8182, + region="us-east-1", + iam_auth=False, + ) + self.store.connect() + + def tearDown(self): + """Clean up test fixtures.""" + if hasattr(self, "store"): + self.store.close() + self.modules_patcher.stop() + + def test_is_retriable_error_signature_expired(self): + """Test that 'Signature expired' is retriable.""" + error = Exception("Signature expired") + result = self.store._is_retriable_error(error) + self.assertTrue(result) + + def test_is_retriable_error_connection_closed(self): + """Test that 'Connection is closed' is retriable.""" + error = Exception("Connection is closed") + result = self.store._is_retriable_error(error) + self.assertTrue(result) + + def test_is_retriable_error_not_retriable(self): + """Test that generic errors are not retriable.""" + error = Exception("Some other error") + result = self.store._is_retriable_error(error) + self.assertFalse(result) + + def test_is_non_retriable_error(self): + """Test _is_non_retriable_error method.""" + error = Exception("Some other error") + result = self.store._is_non_retriable_error(error) + self.assertTrue(result) + + def test_is_non_retriable_error_retriable(self): + """Test _is_non_retriable_error with retriable error.""" + error = Exception("Signature expired") + result = self.store._is_non_retriable_error(error) + self.assertFalse(result) + + +class TestAmazonNeptuneStoreSession(unittest.TestCase): + """Test session management.""" + + def setUp(self): + """Set up test fixtures.""" + self.modules_patcher = patch.dict( + "sys.modules", + { + "neo4j": mock_neo4j, + "neo4j.auth_management": mock_neo4j.auth_management, + "neo4j.api": mock_neo4j.api, + "neo4j.exceptions": mock_neo4j.exceptions, + "boto3": mock_boto3, + }, + ) + self.modules_patcher.start() + + from importlib import reload + + from semantica.graph_store import amazon_neptune + + reload(amazon_neptune) + self.amazon_neptune = amazon_neptune + + self.store = amazon_neptune.AmazonNeptuneStore( + endpoint="test-cluster.us-east-1.neptune.amazonaws.com", + port=8182, + region="us-east-1", + iam_auth=False, + ) + self.store.connect() + + def tearDown(self): + """Clean up test fixtures.""" + if hasattr(self, "store"): + self.store.close() + self.modules_patcher.stop() + + def test_get_session(self): + """Test getting a session.""" + session = self.store.get_session() + self.assertIsInstance(session, self.amazon_neptune.NeptuneSession) + + def test_get_session_with_database_param(self): + """Test getting a session with database parameter (ignored for Neptune).""" + session = self.store.get_session(database="ignored") + self.assertIsInstance(session, self.amazon_neptune.NeptuneSession) + + +class TestAmazonNeptuneStoreGraphAnalyticsDirections(unittest.TestCase): + """Test graph analytics with different directions.""" + + def setUp(self): + """Set up test fixtures.""" + self.modules_patcher = patch.dict( + "sys.modules", + { + "neo4j": mock_neo4j, + "neo4j.auth_management": mock_neo4j.auth_management, + "neo4j.api": mock_neo4j.api, + "neo4j.exceptions": mock_neo4j.exceptions, + "boto3": mock_boto3, + }, + ) + self.modules_patcher.start() + + from importlib import reload + + from semantica.graph_store import amazon_neptune + + reload(amazon_neptune) + + self.store = amazon_neptune.AmazonNeptuneStore( + endpoint="test-cluster.us-east-1.neptune.amazonaws.com", + region="us-east-1", + iam_auth=False, + ) + self.store.connect() + + def tearDown(self): + """Clean up test fixtures.""" + if hasattr(self, "store"): + self.store.close() + self.modules_patcher.stop() + + def test_get_neighbors_outgoing(self): + """Test getting outgoing neighbors.""" + self.store.create_node( + labels=["Person"], properties={"id": "alice", "name": "Alice"} + ) + + neighbors = self.store.get_neighbors(node_id="alice", direction="out", depth=1) + + self.assertIsInstance(neighbors, list) + + def test_get_neighbors_incoming(self): + """Test getting incoming neighbors.""" + self.store.create_node( + labels=["Person"], properties={"id": "alice", "name": "Alice"} + ) + + neighbors = self.store.get_neighbors(node_id="alice", direction="in", depth=1) + + self.assertIsInstance(neighbors, list) + + def test_get_neighbors_with_rel_type(self): + """Test getting neighbors filtered by relationship type.""" + self.store.create_node( + labels=["Person"], properties={"id": "alice", "name": "Alice"} + ) + + neighbors = self.store.get_neighbors( + node_id="alice", rel_type="KNOWS", direction="both", depth=2 + ) + + self.assertIsInstance(neighbors, list) + + def test_shortest_path_with_rel_type(self): + """Test shortest path with relationship type filter.""" + self.store.create_node( + labels=["Person"], properties={"id": "alice", "name": "Alice"} + ) + self.store.create_node( + labels=["Person"], properties={"id": "bob", "name": "Bob"} + ) + + path = self.store.shortest_path( + start_node_id="alice", end_node_id="bob", rel_type="KNOWS", max_depth=3 + ) + + # Path may be None if no path exists + self.assertTrue(path is None or isinstance(path, dict)) + + +class TestNeptuneAuthTokenManager(unittest.TestCase): + """Test NeptuneAuthTokenManager class directly.""" + + def test_auth_token_generation(self): + """Test that SigV4 auth token is generated.""" + with patch.dict( + "sys.modules", + { + "neo4j": mock_neo4j, + "neo4j.auth_management": mock_neo4j.auth_management, + "neo4j.api": mock_neo4j.api, + "neo4j.exceptions": mock_neo4j.exceptions, + "boto3": mock_boto3, + "botocore": mock_botocore, + "botocore.auth": mock_botocore_auth, + "botocore.awsrequest": mock_botocore_awsrequest, + }, + ): + from importlib import reload + + from semantica.graph_store import amazon_neptune + + reload(amazon_neptune) + + endpoint = ( + "bolt://test-cluster.us-east-1.neptune.amazonaws.com" ":8182/opencypher" + ) + auth_manager = amazon_neptune.NeptuneAuthTokenManager( + neptune_endpoint=endpoint, + aws_region="us-east-1", + access_key="AKIAIOSFODNN7EXAMPLE", + secret_key="wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY", + ) + + # get_auth should return a basic auth object + auth = auth_manager.get_auth() + self.assertIsNotNone(auth) + + def test_auth_token_caching(self): + """Test that auth token is cached.""" + with patch.dict( + "sys.modules", + { + "neo4j": mock_neo4j, + "neo4j.auth_management": mock_neo4j.auth_management, + "neo4j.api": mock_neo4j.api, + "neo4j.exceptions": mock_neo4j.exceptions, + "boto3": mock_boto3, + "botocore": mock_botocore, + "botocore.auth": mock_botocore_auth, + "botocore.awsrequest": mock_botocore_awsrequest, + }, + ): + from importlib import reload + + from semantica.graph_store import amazon_neptune + + reload(amazon_neptune) + + endpoint = ( + "bolt://test-cluster.us-east-1.neptune.amazonaws.com" ":8182/opencypher" + ) + auth_manager = amazon_neptune.NeptuneAuthTokenManager( + neptune_endpoint=endpoint, + aws_region="us-east-1", + ) + + # First call generates token, second returns cached + first_auth = auth_manager.get_auth() + second_auth = auth_manager.get_auth() + + self.assertIs(first_auth, second_auth) + + def test_token_refresh(self): + """Test token refresh functionality.""" + with patch.dict( + "sys.modules", + { + "neo4j": mock_neo4j, + "neo4j.auth_management": mock_neo4j.auth_management, + "neo4j.api": mock_neo4j.api, + "neo4j.exceptions": mock_neo4j.exceptions, + "boto3": mock_boto3, + "botocore": mock_botocore, + "botocore.auth": mock_botocore_auth, + "botocore.awsrequest": mock_botocore_awsrequest, + }, + ): + from importlib import reload + + from semantica.graph_store import amazon_neptune + + reload(amazon_neptune) + + endpoint = ( + "bolt://test-cluster.us-east-1.neptune.amazonaws.com" ":8182/opencypher" + ) + auth_manager = amazon_neptune.NeptuneAuthTokenManager( + neptune_endpoint=endpoint, + aws_region="us-east-1", + ) + + # Get initial token + auth_manager.get_auth() + + # Force refresh + auth_manager.refresh_token() + + # Cached token should be None + self.assertIsNone(auth_manager.cached_auth) + + # Next get_auth call should generate new token + new_auth = auth_manager.get_auth() + self.assertIsNotNone(new_auth) + + def test_handle_security_exception(self): + """Test security exception handling.""" + with patch.dict( + "sys.modules", + { + "neo4j": mock_neo4j, + "neo4j.auth_management": mock_neo4j.auth_management, + "neo4j.api": mock_neo4j.api, + "neo4j.exceptions": mock_neo4j.exceptions, + "boto3": mock_boto3, + "botocore": mock_botocore, + "botocore.auth": mock_botocore_auth, + "botocore.awsrequest": mock_botocore_awsrequest, + }, + ): + from importlib import reload + + from semantica.graph_store import amazon_neptune + + reload(amazon_neptune) + + endpoint = ( + "bolt://test-cluster.us-east-1.neptune.amazonaws.com" ":8182/opencypher" + ) + auth_manager = amazon_neptune.NeptuneAuthTokenManager( + neptune_endpoint=endpoint, + aws_region="us-east-1", + ) + + # Get initial token + auth = auth_manager.get_auth() + + # Simulate security exception + mock_error = Exception("Signature expired") + result = auth_manager.handle_security_exception(auth, mock_error) + + # Should return True to retry + self.assertTrue(result) + # Token should be cleared + self.assertIsNone(auth_manager.cached_auth) + + def test_endpoint_conversion(self): + """Test that bolt:// is converted to https:// for signing.""" + with patch.dict( + "sys.modules", + { + "neo4j": mock_neo4j, + "neo4j.auth_management": mock_neo4j.auth_management, + "neo4j.api": mock_neo4j.api, + "neo4j.exceptions": mock_neo4j.exceptions, + "boto3": mock_boto3, + "botocore": mock_botocore, + "botocore.auth": mock_botocore_auth, + "botocore.awsrequest": mock_botocore_awsrequest, + }, + ): + from importlib import reload + + from semantica.graph_store import amazon_neptune + + reload(amazon_neptune) + + endpoint = ( + "bolt://test-cluster.us-east-1.neptune.amazonaws.com" ":8182/opencypher" + ) + auth_manager = amazon_neptune.NeptuneAuthTokenManager( + neptune_endpoint=endpoint, + aws_region="us-east-1", + ) + + # Endpoint should be converted to https for signing + self.assertTrue(auth_manager.neptune_endpoint.startswith("https://")) + + +if __name__ == "__main__": + unittest.main()