From e12e1a916998a70d72205a707b6a9a56e08d28e3 Mon Sep 17 00:00:00 2001 From: Joey French Date: Wed, 15 Oct 2025 20:35:15 -0500 Subject: [PATCH 1/5] Refactor Kuzu API to Graph API and enhance configuration - Renamed `kuzu_api` to `graph_api` across the codebase for clarity and consistency. - Updated environment variables in `.env.example` to include new Neo4j connection settings. - Adjusted deployment workflows to reflect the new API naming. - Enhanced documentation to guide users on the updated API structure and usage. - Added comprehensive tests for the new Graph API components. --- .env.example | 17 +- .github/workflows/deploy-kuzu-infra.yml | 2 +- .github/workflows/deploy-kuzu-volumes.yml | 6 +- .github/workflows/prod.yml | 2 +- .github/workflows/staging.yml | 2 +- .vscode/tasks.json | 6 +- CLAUDE.md | 2 +- README.md | 8 +- bin/entrypoint.sh | 15 +- bin/tools/package-scripts.sh | 2 +- compose.yaml | 84 ++++++- pyproject.toml | 3 +- robosystems/config/env.py | 16 +- robosystems/{kuzu_api => graph_api}/README.md | 32 +-- .../{kuzu_api => graph_api}/__init__.py | 0 .../{kuzu_api => graph_api}/__main__.py | 10 +- robosystems/{kuzu_api => graph_api}/app.py | 2 +- robosystems/graph_api/backends/__init__.py | 32 +++ robosystems/graph_api/backends/base.py | 69 ++++++ robosystems/graph_api/backends/kuzu.py | 105 ++++++++ robosystems/graph_api/backends/neo4j.py | 225 ++++++++++++++++++ robosystems/{kuzu_api => graph_api}/cli.py | 4 +- .../{kuzu_api => graph_api}/client/README.md | 10 +- .../client/__init__.py | 0 .../{kuzu_api => graph_api}/client/base.py | 0 .../{kuzu_api => graph_api}/client/client.py | 0 .../{kuzu_api => graph_api}/client/config.py | 0 .../client/exceptions.py | 0 .../{kuzu_api => graph_api}/client/factory.py | 2 +- .../client/sse_client.py | 0 .../client/sync_client.py | 0 .../{kuzu_api => graph_api}/core/__init__.py | 0 .../core/admission_control.py | 0 .../graph_api/core/backend_cluster_manager.py | 128 ++++++++++ .../core/cluster_manager.py | 8 +- .../core/connection_pool.py | 0 .../core/database_manager.py | 0 .../core/metrics_collector.py | 0 .../core/task_manager.py | 0 .../{kuzu_api => graph_api}/core/task_sse.py | 0 .../{kuzu_api => graph_api}/core/utils.py | 0 robosystems/{kuzu_api => graph_api}/main.py | 4 +- .../middleware/__init__.py | 0 .../middleware/auth.py | 0 .../middleware/request_limits.py | 0 .../models/__init__.py | 0 .../{kuzu_api => graph_api}/models/cluster.py | 0 .../models/database.py | 0 .../models/ingestion.py | 0 .../models/streaming.py | 0 .../routers/__init__.py | 0 .../routers/databases/__init__.py | 0 .../routers/databases/backup.py | 6 +- .../routers/databases/ingest.py | 4 +- .../routers/databases/management.py | 4 +- .../routers/databases/metrics.py | 4 +- .../routers/databases/query.py | 65 +++-- .../routers/databases/restore.py | 6 +- .../routers/databases/schema.py | 6 +- .../{kuzu_api => graph_api}/routers/health.py | 2 +- .../{kuzu_api => graph_api}/routers/info.py | 4 +- .../routers/metrics.py | 4 +- .../{kuzu_api => graph_api}/routers/tasks.py | 6 +- robosystems/middleware/graph/clusters.py | 2 +- .../middleware/graph/multitenant_utils.py | 4 +- robosystems/middleware/graph/repository.py | 2 +- robosystems/middleware/graph/router.py | 4 +- robosystems/middleware/mcp/client.py | 2 +- robosystems/middleware/mcp/factory.py | 2 +- robosystems/middleware/robustness/README.md | 4 +- .../operations/graph/entity_graph_service.py | 4 +- .../operations/graph/generic_graph_service.py | 2 +- .../graph/shared_repository_service.py | 6 +- .../operations/graph/subgraph_service.py | 2 +- robosystems/routers/graphs/copy/execute.py | 2 +- robosystems/routers/graphs/copy/strategies.py | 2 +- robosystems/routers/graphs/health.py | 4 +- robosystems/routers/graphs/info.py | 4 +- robosystems/routers/graphs/limits.py | 4 +- robosystems/tasks/billing/usage_collector.py | 2 +- robosystems/tasks/graph_operations/backup.py | 8 +- robosystems/tasks/sec_xbrl/ingestion.py | 6 +- robosystems/tasks/sec_xbrl/maintenance.py | 2 +- robosystems/tasks/sec_xbrl/orchestration.py | 2 +- tests/README.md | 2 +- .../client/test_base.py | 10 +- .../client/test_client_extended.py | 10 +- .../client/test_config.py | 2 +- .../client/test_factory.py | 30 +-- .../client/test_kuzu_client.py | 6 +- .../client/test_kuzu_exceptions.py | 2 +- .../client/test_sse_client.py | 84 ++++--- .../routers/databases/test_db_query.py | 82 ++++--- .../routers/databases/test_management.py | 36 +-- .../routers/test_health.py | 18 +- .../routers/test_info.py | 18 +- .../routers/test_metrics.py | 24 +- .../routers/test_tasks.py | 38 +-- .../test_auth_middleware.py | 46 ++-- .../test_cluster_server.py | 46 ++-- .../test_connection_pool.py | 4 +- .../test_database_manager.py | 54 ++--- .../test_databases_metrics.py | 4 +- tests/{kuzu_api => graph_api}/test_main.py | 78 +++--- .../test_schema_security.py | 14 +- .../{kuzu_api => graph_api}/test_task_sse.py | 6 +- tests/integration/test_circuit_breaker.py | 2 +- tests/integration/test_query_timeout.py | 6 +- tests/middleware/mcp/test_client.py | 4 +- tests/middleware/mcp/test_mcp_factory.py | 12 +- tests/operations/graph/test_entity_service.py | 2 +- .../graph/test_generic_graph_service.py | 20 +- .../graph/test_shared_repository_service.py | 32 +-- tests/tasks/sec_xbrl/conftest.py | 2 +- tests/tasks/sec_xbrl/test_ingestion.py | 2 +- tests/unit/graph_api/__init__.py | 0 tests/unit/graph_api/backends/__init__.py | 0 .../backends/test_backend_factory.py | 59 +++++ uv.lock | 16 +- 119 files changed, 1275 insertions(+), 473 deletions(-) rename robosystems/{kuzu_api => graph_api}/README.md (96%) rename robosystems/{kuzu_api => graph_api}/__init__.py (100%) rename robosystems/{kuzu_api => graph_api}/__main__.py (55%) rename robosystems/{kuzu_api => graph_api}/app.py (99%) create mode 100644 robosystems/graph_api/backends/__init__.py create mode 100644 robosystems/graph_api/backends/base.py create mode 100644 robosystems/graph_api/backends/kuzu.py create mode 100644 robosystems/graph_api/backends/neo4j.py rename robosystems/{kuzu_api => graph_api}/cli.py (96%) rename robosystems/{kuzu_api => graph_api}/client/README.md (97%) rename robosystems/{kuzu_api => graph_api}/client/__init__.py (100%) rename robosystems/{kuzu_api => graph_api}/client/base.py (100%) rename robosystems/{kuzu_api => graph_api}/client/client.py (100%) rename robosystems/{kuzu_api => graph_api}/client/config.py (100%) rename robosystems/{kuzu_api => graph_api}/client/exceptions.py (100%) rename robosystems/{kuzu_api => graph_api}/client/factory.py (99%) rename robosystems/{kuzu_api => graph_api}/client/sse_client.py (100%) rename robosystems/{kuzu_api => graph_api}/client/sync_client.py (100%) rename robosystems/{kuzu_api => graph_api}/core/__init__.py (100%) rename robosystems/{kuzu_api => graph_api}/core/admission_control.py (100%) create mode 100644 robosystems/graph_api/core/backend_cluster_manager.py rename robosystems/{kuzu_api => graph_api}/core/cluster_manager.py (99%) rename robosystems/{kuzu_api => graph_api}/core/connection_pool.py (100%) rename robosystems/{kuzu_api => graph_api}/core/database_manager.py (100%) rename robosystems/{kuzu_api => graph_api}/core/metrics_collector.py (100%) rename robosystems/{kuzu_api => graph_api}/core/task_manager.py (100%) rename robosystems/{kuzu_api => graph_api}/core/task_sse.py (100%) rename robosystems/{kuzu_api => graph_api}/core/utils.py (100%) rename robosystems/{kuzu_api => graph_api}/main.py (97%) rename robosystems/{kuzu_api => graph_api}/middleware/__init__.py (100%) rename robosystems/{kuzu_api => graph_api}/middleware/auth.py (100%) rename robosystems/{kuzu_api => graph_api}/middleware/request_limits.py (100%) rename robosystems/{kuzu_api => graph_api}/models/__init__.py (100%) rename robosystems/{kuzu_api => graph_api}/models/cluster.py (100%) rename robosystems/{kuzu_api => graph_api}/models/database.py (100%) rename robosystems/{kuzu_api => graph_api}/models/ingestion.py (100%) rename robosystems/{kuzu_api => graph_api}/models/streaming.py (100%) rename robosystems/{kuzu_api => graph_api}/routers/__init__.py (100%) rename robosystems/{kuzu_api => graph_api}/routers/databases/__init__.py (100%) rename robosystems/{kuzu_api => graph_api}/routers/databases/backup.py (94%) rename robosystems/{kuzu_api => graph_api}/routers/databases/ingest.py (99%) rename robosystems/{kuzu_api => graph_api}/routers/databases/management.py (96%) rename robosystems/{kuzu_api => graph_api}/routers/databases/metrics.py (95%) rename robosystems/{kuzu_api => graph_api}/routers/databases/query.py (60%) rename robosystems/{kuzu_api => graph_api}/routers/databases/restore.py (96%) rename robosystems/{kuzu_api => graph_api}/routers/databases/schema.py (98%) rename robosystems/{kuzu_api => graph_api}/routers/health.py (95%) rename robosystems/{kuzu_api => graph_api}/routers/info.py (83%) rename robosystems/{kuzu_api => graph_api}/routers/metrics.py (92%) rename robosystems/{kuzu_api => graph_api}/routers/tasks.py (97%) rename tests/{kuzu_api => graph_api}/client/test_base.py (97%) rename tests/{kuzu_api => graph_api}/client/test_client_extended.py (98%) rename tests/{kuzu_api => graph_api}/client/test_config.py (99%) rename tests/{kuzu_api => graph_api}/client/test_factory.py (91%) rename tests/{kuzu_api => graph_api}/client/test_kuzu_client.py (97%) rename tests/{kuzu_api => graph_api}/client/test_kuzu_exceptions.py (99%) rename tests/{kuzu_api => graph_api}/client/test_sse_client.py (87%) rename tests/{kuzu_api => graph_api}/routers/databases/test_db_query.py (80%) rename tests/{kuzu_api => graph_api}/routers/databases/test_management.py (89%) rename tests/{kuzu_api => graph_api}/routers/test_health.py (89%) rename tests/{kuzu_api => graph_api}/routers/test_info.py (89%) rename tests/{kuzu_api => graph_api}/routers/test_metrics.py (93%) rename tests/{kuzu_api => graph_api}/routers/test_tasks.py (90%) rename tests/{kuzu_api => graph_api}/test_auth_middleware.py (89%) rename tests/{kuzu_api => graph_api}/test_cluster_server.py (94%) rename tests/{kuzu_api => graph_api}/test_connection_pool.py (99%) rename tests/{kuzu_api => graph_api}/test_database_manager.py (93%) rename tests/{kuzu_api => graph_api}/test_databases_metrics.py (98%) rename tests/{kuzu_api => graph_api}/test_main.py (82%) rename tests/{kuzu_api => graph_api}/test_schema_security.py (94%) rename tests/{kuzu_api => graph_api}/test_task_sse.py (98%) create mode 100644 tests/unit/graph_api/__init__.py create mode 100644 tests/unit/graph_api/backends/__init__.py create mode 100644 tests/unit/graph_api/backends/test_backend_factory.py diff --git a/.env.example b/.env.example index 51dfdcc5..e0f0821f 100755 --- a/.env.example +++ b/.env.example @@ -125,11 +125,15 @@ CELERY_RESULT_BACKEND=redis://:valkey@valkey:6379/1 # API_KEY_CACHE_TTL=300 ## ============================================================================= -## KUZU GRAPH DATABASE CONFIGURATION +## GRAPH DATABASE CONFIGURATION (KUZU AND NEO4J) ## ============================================================================= +## Graph Backend Selection +## Options: kuzu (default), neo4j +# BACKEND_TYPE=kuzu + ## Kuzu API URL -KUZU_API_URL=http://kuzu:8001 +KUZU_API_URL=http://kuzu-api:8001 ## User Graph Limits USER_GRAPHS_DEFAULT_LIMIT=100 @@ -170,6 +174,15 @@ KUZU_MAX_DATABASES_PER_NODE=50 # INSTANCE_ID=unknown # CLUSTER_TIER=standard +## Neo4j Backend Configuration +# NEO4J_URI=bolt://neo4j-db:7687 +# NEO4J_USERNAME=neo4j +# NEO4J_PASSWORD= # Retrieved from AWS Secrets Manager in prod/staging +# NEO4J_ENTERPRISE=false # Enable multi-database support (requires Enterprise license) +# NEO4J_MAX_CONNECTION_POOL_SIZE=50 +# NEO4J_CONNECTION_ACQUISITION_TIMEOUT=60 +# NEO4J_MAX_CONNECTION_LIFETIME=3600 + ## ============================================================================= ## PERFORMANCE AND SCALING ## ============================================================================= diff --git a/.github/workflows/deploy-kuzu-infra.yml b/.github/workflows/deploy-kuzu-infra.yml index 64e7a2aa..61629d9e 100644 --- a/.github/workflows/deploy-kuzu-infra.yml +++ b/.github/workflows/deploy-kuzu-infra.yml @@ -42,7 +42,7 @@ on: required: false type: string default: "" - kuzu_api_rotation_code_key: + graph_api_rotation_code_key: description: "S3 key for Kuzu API Rotation Lambda deployment package" required: false type: string diff --git a/.github/workflows/deploy-kuzu-volumes.yml b/.github/workflows/deploy-kuzu-volumes.yml index 183296b6..b6fff15e 100644 --- a/.github/workflows/deploy-kuzu-volumes.yml +++ b/.github/workflows/deploy-kuzu-volumes.yml @@ -49,7 +49,7 @@ on: default: "" # Secrets Configuration - kuzu_api_secret_arn: + graph_api_secret_arn: description: "ARN of the Kuzu API secret for Lambda authentication" required: false type: string @@ -164,9 +164,9 @@ jobs: fi # Add Kuzu API secret if provided - if [ -n "${{ inputs.kuzu_api_secret_arn }}" ]; then + if [ -n "${{ inputs.graph_api_secret_arn }}" ]; then STACK_PARAMS="$STACK_PARAMS \ - ParameterKey=KuzuSecretArn,ParameterValue=${{ inputs.kuzu_api_secret_arn }}" + ParameterKey=KuzuSecretArn,ParameterValue=${{ inputs.graph_api_secret_arn }}" fi # Deploy or update the stack diff --git a/.github/workflows/prod.yml b/.github/workflows/prod.yml index 8a2f5616..fd88f4ef 100644 --- a/.github/workflows/prod.yml +++ b/.github/workflows/prod.yml @@ -418,7 +418,7 @@ jobs: # Lambda Configuration lambda_code_bucket: robosystems-${{ vars.ENVIRONMENT_PROD || 'prod' }}-deployment # Secrets Configuration - kuzu_api_secret_arn: ${{ needs.deploy-kuzu-infra.outputs.secret_arn }} + graph_api_secret_arn: ${{ needs.deploy-kuzu-infra.outputs.secret_arn }} secrets: ACTIONS_TOKEN: ${{ secrets.ACTIONS_TOKEN }} AWS_ACCESS_KEY_ID: ${{ secrets.AWS_ACCESS_KEY_ID }} diff --git a/.github/workflows/staging.yml b/.github/workflows/staging.yml index 8158fe01..00ca0376 100644 --- a/.github/workflows/staging.yml +++ b/.github/workflows/staging.yml @@ -419,7 +419,7 @@ jobs: # Lambda Configuration lambda_code_bucket: robosystems-${{ vars.ENVIRONMENT_STAGING || 'staging' }}-deployment # Secrets Configuration - kuzu_api_secret_arn: ${{ needs.deploy-kuzu-infra.outputs.secret_arn }} + graph_api_secret_arn: ${{ needs.deploy-kuzu-infra.outputs.secret_arn }} secrets: ACTIONS_TOKEN: ${{ secrets.ACTIONS_TOKEN }} AWS_ACCESS_KEY_ID: ${{ secrets.AWS_ACCESS_KEY_ID }} diff --git a/.vscode/tasks.json b/.vscode/tasks.json index e4c6b82b..d6d2d157 100644 --- a/.vscode/tasks.json +++ b/.vscode/tasks.json @@ -269,6 +269,8 @@ "pg", "valkey", "kuzu", + "neo4j", + "premium", "localstack", "api", "worker", @@ -288,7 +290,9 @@ "api", "worker", "beat", - "kuzu", + "kuzu-api", + "neo4j-db", + "neo4j-api", "pg-iam", "valkey", "grafana", diff --git a/CLAUDE.md b/CLAUDE.md index b4b654fd..2058b889 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -89,7 +89,7 @@ robosystems/ ├── adapters/ # External service integrations ├── config/ # Centralized configuration ├── security/ # Security implementations -├── kuzu_api/ # Graph database API service +├── graph_api/ # Graph database API service └── scripts/ # Utility and admin scripts ``` diff --git a/README.md b/README.md index 03007763..84bfc05a 100644 --- a/README.md +++ b/README.md @@ -121,7 +121,7 @@ just logs-follow worker # CloudWatch log search - **API-First Design**: All database access through REST APIs with no direct database connections - **Schema-Driven Operations**: All graph operations derive from curated schemas (RoboLedger, RoboInvestor, and more) -#### Kuzu API System (`/robosystems/kuzu_api/`) +#### Kuzu API System (`/robosystems/graph_api/`) The **Kuzu API** is a FastAPI microservice that runs alongside Kuzu databases on instances, providing: @@ -132,7 +132,7 @@ The **Kuzu API** is a FastAPI microservice that runs alongside Kuzu databases on - **Streaming Support**: NDJSON streaming for large query results - **Admission Control**: CPU/memory-based backpressure to prevent overload -#### Client-Factory System (`/robosystems/kuzu_api/client/`) +#### Client-Factory System (`/robosystems/graph_api/client/`) The client-factory layer provides intelligent routing between application code and Kuzu infrastructure: @@ -387,8 +387,8 @@ Each major system component has detailed documentation: ### Kuzu Graph Database System -- **`/robosystems/kuzu_api/README.md`**: Complete Kuzu API documentation -- **`/robosystems/kuzu_api/client/README.md`**: Client-factory system for intelligent routing +- **`/robosystems/graph_api/README.md`**: Complete Kuzu API documentation +- **`/robosystems/graph_api/client/README.md`**: Client-factory system for intelligent routing ### Middleware Components diff --git a/bin/entrypoint.sh b/bin/entrypoint.sh index 52f710cb..45286886 100755 --- a/bin/entrypoint.sh +++ b/bin/entrypoint.sh @@ -120,7 +120,7 @@ case $DOCKER_PROFILE in "kuzu-writer") echo "Starting Kuzu Writer API..." # max-databases will be loaded from tier configuration based on CLUSTER_TIER - exec uv run python -m robosystems.kuzu_api \ + exec uv run python -m robosystems.graph_api \ --node-type writer \ --repository-type entity \ --port ${KUZU_PORT:-8001} \ @@ -138,13 +138,24 @@ case $DOCKER_PROFILE in READONLY_FLAG="" fi # max-databases will be loaded from tier configuration based on CLUSTER_TIER - exec uv run python -m robosystems.kuzu_api \ + exec uv run python -m robosystems.graph_api \ --node-type ${KUZU_NODE_TYPE} \ --repository-type shared \ --port ${KUZU_PORT:-8002} \ --base-path ${KUZU_DATABASE_PATH:-/app/data/kuzu-dbs} \ ${READONLY_FLAG} ;; + "neo4j-writer") + echo "Starting Neo4j Graph API..." + # Graph API with Neo4j backend + # Backend type determined by BACKEND_TYPE env var (neo4j_community or neo4j_enterprise) + # Note: base-path is for metadata only (actual data stored in Neo4j database via Bolt) + exec uv run python -m robosystems.graph_api \ + --node-type writer \ + --repository-type entity \ + --port ${GRAPH_API_PORT:-8002} \ + --base-path /app/data/neo4j-metadata + ;; *) echo "Unknown profile: $DOCKER_PROFILE" exit 1 diff --git a/bin/tools/package-scripts.sh b/bin/tools/package-scripts.sh index 2adf4c59..f032bc4b 100755 --- a/bin/tools/package-scripts.sh +++ b/bin/tools/package-scripts.sh @@ -68,7 +68,7 @@ psycopg2-binary==2.9.9" package_lambda "valkey-rotation" "valkey_rotation.py" "boto3==1.34.14 redis==5.0.1" -package_lambda "kuzu-api-rotation" "kuzu_api_rotation.py" "boto3==1.34.14" +package_lambda "kuzu-api-rotation" "graph_api_rotation.py" "boto3==1.34.14" # Package snapshot Lambda functions (temporary - will be migrated to volume manager) package_lambda "kuzu-snapshot-creator" "kuzu_snapshot_creator.py" "boto3==1.34.14" diff --git a/compose.yaml b/compose.yaml index cef8560a..c9f6fbe8 100644 --- a/compose.yaml +++ b/compose.yaml @@ -66,9 +66,9 @@ services: restart: always profiles: ["beat", "robosystems", "all"] - # Kuzu - kuzu: - container_name: kuzu + # Kuzu Graph API - Graph API with Kuzu Backend + kuzu-api: + container_name: kuzu-api build: context: . dockerfile: Dockerfile @@ -76,6 +76,7 @@ services: environment: - DOCKER_PROFILE=kuzu-writer - ENVIRONMENT=dev + - BACKEND_TYPE=kuzu - KUZU_PORT=8001 - KUZU_DATABASE_PATH=/app/data/kuzu-dbs - KUZU_NODE_TYPE=writer @@ -97,6 +98,74 @@ services: restart: always profiles: ["kuzu", "dev", "robosystems", "all"] + # Neo4j Database - Graph Database Backend + neo4j-db: + container_name: neo4j-db + image: neo4j:5.25-community + environment: + - NEO4J_AUTH=neo4j/neo4jpassword + - NEO4J_PLUGINS=["apoc"] + - NEO4J_dbms_memory_pagecache_size=512M + - NEO4J_dbms_memory_heap_initial__size=512M + - NEO4J_dbms_memory_heap_max__size=1G + - NEO4J_apoc_export_file_enabled=true + - NEO4J_apoc_import_file_enabled=true + - NEO4J_apoc_import_file_use__neo4j__config=true + ports: + - "7474:7474" # HTTP Browser + - "7687:7687" # Bolt Protocol + volumes: + - ./data/neo4j/data:/data + - ./data/neo4j/logs:/logs + - ./data/neo4j/import:/var/lib/neo4j/import + - ./data/neo4j/plugins:/plugins + healthcheck: + test: + [ + "CMD-SHELL", + "wget --no-verbose --tries=1 --spider http://localhost:7474 || exit 1", + ] + interval: 10s + timeout: 5s + retries: 5 + start_period: 30s + restart: always + profiles: ["neo4j", "dev", "robosystems", "all"] + + # Neo4j Graph API - Graph API with Neo4j Backend + neo4j-api: + container_name: neo4j-api + build: + context: . + dockerfile: Dockerfile + command: /app/bin/entrypoint.sh + environment: + - DOCKER_PROFILE=neo4j-writer + - ENVIRONMENT=dev + - BACKEND_TYPE=neo4j_community + - NEO4J_URI=bolt://neo4j-db:7687 + - NEO4J_USERNAME=neo4j + - NEO4J_PASSWORD=neo4jpassword + - GRAPH_API_PORT=8002 + - CLUSTER_TIER=professional + env_file: + - .env + ports: + - "8002:8002" + volumes: + - ./robosystems:/app/robosystems # Mount code for hot reload + healthcheck: + test: ["CMD", "curl", "-f", "http://localhost:8002/health"] + interval: 30s + timeout: 10s + retries: 3 + start_period: 10s + depends_on: + neo4j-db: + condition: service_healthy + restart: always + profiles: ["neo4j", "dev", "robosystems", "all"] + # Robosystems App Frontend robosystems-app: container_name: robosystems-app @@ -201,14 +270,7 @@ services: --maxmemory-policy allkeys-lru --requirepass ${VALKEY_AUTH_TOKEN:-valkey} healthcheck: - test: - [ - "CMD", - "valkey-cli", - "-a", - "${VALKEY_AUTH_TOKEN:-valkey}", - "ping", - ] + test: ["CMD", "valkey-cli", "-a", "${VALKEY_AUTH_TOKEN:-valkey}", "ping"] interval: 5s timeout: 3s retries: 5 diff --git a/pyproject.toml b/pyproject.toml index cf814c9a..c0ed6424 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -44,7 +44,8 @@ dependencies = [ "python-multipart>=0.0.16,<1.0", "sse-starlette>=2.4.1,<3.0", # Graph database - "kuzu>=0.11.0,<0.12", + "kuzu==0.11.2", + "neo4j>=5.15.0,<6.0", # Database and ORM "alembic>=1.16.0,<2.0", "sqlalchemy>=2.0.0,<3.0", diff --git a/robosystems/config/env.py b/robosystems/config/env.py index 3b0455d1..1dcd8bd7 100644 --- a/robosystems/config/env.py +++ b/robosystems/config/env.py @@ -366,9 +366,12 @@ class EnvConfig: DATABASE_ECHO = get_bool_env("DATABASE_ECHO", False) # ========================================================================== - # DATABASE CONFIGURATION - KUZU GRAPH DATABASE + # DATABASE CONFIGURATION - GRAPH DATABASES (KUZU AND NEO4J) # ========================================================================== + # Graph Backend Selection + BACKEND_TYPE = get_str_env("BACKEND_TYPE", "kuzu") # Options: kuzu, neo4j + # Basic Kuzu configuration KUZU_API_URL = get_str_env("KUZU_API_URL", "http://localhost:8001") KUZU_API_KEY = get_secret_value("KUZU_API_KEY", "") @@ -377,6 +380,17 @@ class EnvConfig: KUZU_NODE_TYPE = get_str_env("KUZU_NODE_TYPE", "writer") KUZU_S3_BUCKET = get_str_env("KUZU_S3_BUCKET", "") + # Neo4j configuration + NEO4J_URI = get_str_env("NEO4J_URI", "bolt://localhost:7687") + NEO4J_USERNAME = get_str_env("NEO4J_USERNAME", "neo4j") + NEO4J_PASSWORD = get_secret_value("NEO4J_PASSWORD", "") + NEO4J_ENTERPRISE = get_bool_env("NEO4J_ENTERPRISE", False) + NEO4J_MAX_CONNECTION_POOL_SIZE = get_int_env("NEO4J_MAX_CONNECTION_POOL_SIZE", 50) + NEO4J_CONNECTION_ACQUISITION_TIMEOUT = get_int_env( + "NEO4J_CONNECTION_ACQUISITION_TIMEOUT", 60 + ) + NEO4J_MAX_CONNECTION_LIFETIME = get_int_env("NEO4J_MAX_CONNECTION_LIFETIME", 3600) + # User graph creation limits (safety valve) # User graph limits (from secrets for runtime control) USER_GRAPHS_DEFAULT_LIMIT = get_int_env( diff --git a/robosystems/kuzu_api/README.md b/robosystems/graph_api/README.md similarity index 96% rename from robosystems/kuzu_api/README.md rename to robosystems/graph_api/README.md index 837aa02c..71917c63 100644 --- a/robosystems/kuzu_api/README.md +++ b/robosystems/graph_api/README.md @@ -41,7 +41,7 @@ High-performance HTTP API server for Kuzu graph database cluster management. Fas ### Core Components ``` -kuzu_api/ +graph_api/ ├── app.py # FastAPI application factory ├── main.py # Server entry point ├── __main__.py # Module entry point @@ -335,7 +335,7 @@ Response: { ### Async Client ```python -from robosystems.kuzu_api.client import AsyncKuzuClient +from robosystems.graph_api.client import AsyncKuzuClient async with AsyncKuzuClient( base_url="http://kuzu-writer:8001", @@ -368,7 +368,7 @@ async with AsyncKuzuClient( ### Sync Client ```python -from robosystems.kuzu_api.client import KuzuClient +from robosystems.graph_api.client import KuzuClient client = KuzuClient( base_url="http://kuzu-writer:8001", @@ -385,7 +385,7 @@ data = client.query( ### Client Factory with Intelligent Routing ```python -from robosystems.kuzu_api.client.factory import get_kuzu_client +from robosystems.graph_api.client.factory import get_kuzu_client # Factory handles routing based on graph type and operation client = await get_kuzu_client( @@ -542,7 +542,7 @@ OTEL_EXPORTER_OTLP_ENDPOINT=http://172.17.0.1:4318 just start robosystems # Run API server locally -uv run python -m robosystems.kuzu_api \ +uv run python -m robosystems.graph_api \ --base-path ./data/kuzu-dbs \ --node-type writer \ --port 8001 @@ -560,19 +560,19 @@ docker run -d \ -e KUZU_NODE_TYPE=writer \ -e WRITER_TIER=standard \ robosystems-api:latest \ - python -m robosystems.kuzu_api + python -m robosystems.graph_api ``` ### CLI Tools ```bash # Server mode -python -m robosystems.kuzu_api --help +python -m robosystems.graph_api --help # Client CLI -python -m robosystems.kuzu_api cli health -python -m robosystems.kuzu_api cli query kg1a2b3c "MATCH (n) RETURN count(n)" -python -m robosystems.kuzu_api cli ingest kg1a2b3c /path/to/data.parquet +python -m robosystems.graph_api cli health +python -m robosystems.graph_api cli query kg1a2b3c "MATCH (n) RETURN count(n)" +python -m robosystems.graph_api cli ingest kg1a2b3c /path/to/data.parquet ``` ## Testing @@ -581,29 +581,29 @@ python -m robosystems.kuzu_api cli ingest kg1a2b3c /path/to/data.parquet ```bash # Run all tests -uv run pytest tests/kuzu_api/ -v +uv run pytest tests/graph_api/ -v # Run specific test categories -uv run pytest tests/kuzu_api/test_client.py -v -uv run pytest tests/kuzu_api/test_ingestion.py -v +uv run pytest tests/graph_api/test_client.py -v +uv run pytest tests/graph_api/test_ingestion.py -v ``` ### Integration Tests ```bash # Requires running Kuzu instance -uv run pytest tests/kuzu_api/ -m integration +uv run pytest tests/graph_api/ -m integration # Test with real S3 AWS_ENDPOINT_URL=http://localhost:4566 \ - uv run pytest tests/kuzu_api/test_s3_ingestion.py + uv run pytest tests/graph_api/test_s3_ingestion.py ``` ### Load Testing ```bash # Using locust for load testing -locust -f tests/kuzu_api/loadtest.py \ +locust -f tests/graph_api/loadtest.py \ --host http://localhost:8001 \ --users 100 \ --spawn-rate 10 diff --git a/robosystems/kuzu_api/__init__.py b/robosystems/graph_api/__init__.py similarity index 100% rename from robosystems/kuzu_api/__init__.py rename to robosystems/graph_api/__init__.py diff --git a/robosystems/kuzu_api/__main__.py b/robosystems/graph_api/__main__.py similarity index 55% rename from robosystems/kuzu_api/__main__.py rename to robosystems/graph_api/__main__.py index 47347b5b..7108be3f 100644 --- a/robosystems/kuzu_api/__main__.py +++ b/robosystems/graph_api/__main__.py @@ -3,9 +3,9 @@ Kuzu API module entry point. This module enables running the Kuzu API server or client as a module: - python -m robosystems.kuzu_api --help # Server help - python -m robosystems.kuzu_api --base-path /data ... # Start server - python -m robosystems.kuzu_api cli health # Client commands + python -m robosystems.graph_api --help # Server help + python -m robosystems.graph_api --base-path /data ... # Start server + python -m robosystems.graph_api cli health # Client commands """ import sys @@ -15,11 +15,11 @@ if len(sys.argv) > 1 and sys.argv[1] == "cli": # Remove 'cli' from args and run client sys.argv.pop(1) - from robosystems.kuzu_api.cli import main + from robosystems.graph_api.cli import main exit(main()) else: # Default to server - from robosystems.kuzu_api.main import main + from robosystems.graph_api.main import main exit(main()) diff --git a/robosystems/kuzu_api/app.py b/robosystems/graph_api/app.py similarity index 99% rename from robosystems/kuzu_api/app.py rename to robosystems/graph_api/app.py index 21f34d3b..b13d56f4 100644 --- a/robosystems/kuzu_api/app.py +++ b/robosystems/graph_api/app.py @@ -13,7 +13,7 @@ from fastapi.staticfiles import StaticFiles from fastapi.responses import HTMLResponse -from robosystems.kuzu_api.routers import ( +from robosystems.graph_api.routers import ( databases, health, info, diff --git a/robosystems/graph_api/backends/__init__.py b/robosystems/graph_api/backends/__init__.py new file mode 100644 index 00000000..1211e79e --- /dev/null +++ b/robosystems/graph_api/backends/__init__.py @@ -0,0 +1,32 @@ +from typing import Optional, Union +from robosystems.config import env +from .kuzu import KuzuBackend +from .neo4j import Neo4jBackend +from robosystems.logger import logger + + +_backend_instance: Optional[Union[KuzuBackend, Neo4jBackend]] = None + + +def get_backend() -> Union[KuzuBackend, Neo4jBackend]: + global _backend_instance + + if _backend_instance is None: + backend_type = env.BACKEND_TYPE + + if backend_type == "kuzu": + _backend_instance = KuzuBackend() + logger.info("Initialized Kuzu backend (Standard tier)") + elif backend_type == "neo4j_community": + _backend_instance = Neo4jBackend(enterprise=False) + logger.info("Initialized Neo4j Community backend (Professional/Enterprise tiers)") + elif backend_type == "neo4j_enterprise": + _backend_instance = Neo4jBackend(enterprise=True) + logger.info("Initialized Neo4j Enterprise backend (Premium tier)") + else: + raise ValueError(f"Unknown BACKEND_TYPE: {backend_type}") + + return _backend_instance + + +__all__ = ["get_backend", "KuzuBackend", "Neo4jBackend", "GraphBackend"] diff --git a/robosystems/graph_api/backends/base.py b/robosystems/graph_api/backends/base.py new file mode 100644 index 00000000..f5d96aef --- /dev/null +++ b/robosystems/graph_api/backends/base.py @@ -0,0 +1,69 @@ +from abc import ABC, abstractmethod +from typing import Dict, Any, List, Optional +from dataclasses import dataclass + + +@dataclass +class DatabaseInfo: + name: str + node_count: int + relationship_count: int + size_bytes: int + + +@dataclass +class ClusterTopology: + mode: str + leader: Optional[Dict[str, Any]] = None + followers: Optional[List[Dict[str, Any]]] = None + members: Optional[List[Dict[str, Any]]] = None + + +class GraphBackend(ABC): + @abstractmethod + async def execute_query( + self, + graph_id: str, + cypher: str, + parameters: Optional[Dict[str, Any]] = None, + database: Optional[str] = None, + ) -> List[Dict[str, Any]]: + pass + + @abstractmethod + async def execute_write( + self, + graph_id: str, + cypher: str, + parameters: Optional[Dict[str, Any]] = None, + database: Optional[str] = None, + ) -> List[Dict[str, Any]]: + pass + + @abstractmethod + async def create_database(self, database_name: str) -> bool: + pass + + @abstractmethod + async def delete_database(self, database_name: str) -> bool: + pass + + @abstractmethod + async def list_databases(self) -> List[str]: + pass + + @abstractmethod + async def get_database_info(self, database_name: str) -> DatabaseInfo: + pass + + @abstractmethod + async def get_cluster_topology(self) -> ClusterTopology: + pass + + @abstractmethod + async def health_check(self) -> bool: + pass + + @abstractmethod + def close(self) -> None: + pass diff --git a/robosystems/graph_api/backends/kuzu.py b/robosystems/graph_api/backends/kuzu.py new file mode 100644 index 00000000..aa4be5c1 --- /dev/null +++ b/robosystems/graph_api/backends/kuzu.py @@ -0,0 +1,105 @@ +from typing import Dict, Any, List, Optional +from pathlib import Path +import shutil + +from robosystems.middleware.graph.engine import Engine +from robosystems.logger import logger +from .base import GraphBackend, DatabaseInfo, ClusterTopology + + +class KuzuBackend(GraphBackend): + def __init__(self, data_path: str = "/data/kuzu-dbs"): + self.data_path = data_path + self._engines: Dict[str, Engine] = {} + Path(data_path).mkdir(parents=True, exist_ok=True) + + def _get_engine(self, graph_id: str) -> Engine: + if graph_id not in self._engines: + database_path = f"{self.data_path}/{graph_id}" + self._engines[graph_id] = Engine(database_path) + logger.debug(f"Created Kuzu engine for {graph_id}: {database_path}") + return self._engines[graph_id] + + async def execute_query( + self, + graph_id: str, + cypher: str, + parameters: Optional[Dict[str, Any]] = None, + database: Optional[str] = None, + ) -> List[Dict[str, Any]]: + engine = self._get_engine(graph_id) + return engine.execute_query(cypher, parameters) + + async def execute_write( + self, + graph_id: str, + cypher: str, + parameters: Optional[Dict[str, Any]] = None, + database: Optional[str] = None, + ) -> List[Dict[str, Any]]: + engine = self._get_engine(graph_id) + return engine.execute_query(cypher, parameters) + + async def create_database(self, database_name: str) -> bool: + self._get_engine(database_name) + logger.info(f"Created Kuzu database for {database_name}") + return True + + async def delete_database(self, database_name: str) -> bool: + if database_name in self._engines: + self._engines[database_name].close() + del self._engines[database_name] + + db_path = Path(f"{self.data_path}/{database_name}") + if db_path.exists(): + shutil.rmtree(db_path) + logger.info(f"Deleted Kuzu database for {database_name}") + return True + + async def list_databases(self) -> List[str]: + db_path = Path(self.data_path) + if not db_path.exists(): + return [] + return [d.name for d in db_path.iterdir() if d.is_dir()] + + async def get_database_info(self, database_name: str) -> DatabaseInfo: + engine = self._get_engine(database_name) + + node_count = 0 + relationship_count = 0 + + try: + node_result = engine.execute_single("MATCH (n) RETURN count(n) as count") + if node_result: + node_count = node_result.get("count", 0) + + rel_result = engine.execute_single("MATCH ()-[r]->() RETURN count(r) as count") + if rel_result: + relationship_count = rel_result.get("count", 0) + except Exception as e: + logger.warning(f"Failed to get database stats for {database_name}: {e}") + + db_path = Path(f"{self.data_path}/{database_name}") + size_bytes = ( + sum(f.stat().st_size for f in db_path.rglob("*") if f.is_file()) + if db_path.exists() + else 0 + ) + + return DatabaseInfo( + name=database_name, + node_count=node_count, + relationship_count=relationship_count, + size_bytes=size_bytes, + ) + + async def get_cluster_topology(self) -> ClusterTopology: + return ClusterTopology(mode="embedded", leader={"backend": "kuzu"}) + + async def health_check(self) -> bool: + return True + + def close(self) -> None: + for engine in self._engines.values(): + engine.close() + self._engines.clear() diff --git a/robosystems/graph_api/backends/neo4j.py b/robosystems/graph_api/backends/neo4j.py new file mode 100644 index 00000000..115e1308 --- /dev/null +++ b/robosystems/graph_api/backends/neo4j.py @@ -0,0 +1,225 @@ +from typing import Dict, Any, List, Optional +import json +import boto3 +from neo4j import AsyncGraphDatabase, AsyncDriver +from robosystems.logger import logger +from robosystems.config import env +from .base import GraphBackend, DatabaseInfo, ClusterTopology + + +class Neo4jBackend(GraphBackend): + def __init__(self, enterprise: bool = False): + self.enterprise = enterprise + self.bolt_url = env.NEO4J_URI + self.driver: Optional[AsyncDriver] = None + self._cluster_topology: Optional[Dict] = None + self._password: Optional[str] = None + self._username = env.NEO4J_USERNAME + + async def _ensure_connected(self): + if self.driver is None: + await self._connect() + + async def _connect(self): + if env.NEO4J_PASSWORD: + self._password = env.NEO4J_PASSWORD + else: + secrets = boto3.client("secretsmanager", region_name=env.AWS_REGION) + secret_value = secrets.get_secret_value( + SecretId=f"robosystems/{env.ENVIRONMENT}/neo4j" + ) + creds = json.loads(secret_value["SecretString"]) + self._password = creds["password"] + + self.driver = AsyncGraphDatabase.driver( + self.bolt_url, + auth=(self._username, self._password), + max_connection_lifetime=env.NEO4J_MAX_CONNECTION_LIFETIME, + max_connection_pool_size=env.NEO4J_MAX_CONNECTION_POOL_SIZE, + connection_acquisition_timeout=env.NEO4J_CONNECTION_ACQUISITION_TIMEOUT, + ) + logger.info(f"Connected to Neo4j at {self.bolt_url} (enterprise={self.enterprise})") + + def _get_database_name(self, graph_id: str, database: Optional[str] = None) -> str: + if database: + return database + elif self.enterprise: + return f"kg_{graph_id}_main" + else: + return "neo4j" + + async def execute_query( + self, + graph_id: str, + cypher: str, + parameters: Optional[Dict[str, Any]] = None, + database: Optional[str] = None, + ) -> List[Dict[str, Any]]: + await self._ensure_connected() + + db_name = self._get_database_name(graph_id, database) + + cypher_upper = cypher.strip().upper() + is_read = cypher_upper.startswith(("MATCH", "RETURN", "WITH", "CALL")) + + async with self.driver.session(database=db_name) as session: + if self.enterprise and is_read: + result = await session.run(cypher, parameters or {}) + else: + result = await session.run(cypher, parameters or {}) + + records = await result.data() + return records + + async def execute_write( + self, + graph_id: str, + cypher: str, + parameters: Optional[Dict[str, Any]] = None, + database: Optional[str] = None, + ) -> List[Dict[str, Any]]: + await self._ensure_connected() + + db_name = self._get_database_name(graph_id, database) + + async with self.driver.session(database=db_name) as session: + + async def _tx_function(tx): + result = await tx.run(cypher, parameters or {}) + return await result.data() + + records = await session.execute_write(_tx_function) + return records + + async def create_database(self, database_name: str) -> bool: + if not self.enterprise: + raise ValueError("Multi-database requires Neo4j Enterprise") + + await self._ensure_connected() + + async with self.driver.session(database="system") as session: + await session.run(f"CREATE DATABASE `{database_name}` IF NOT EXISTS") + + logger.info(f"Created Neo4j database: {database_name}") + return True + + async def delete_database(self, database_name: str) -> bool: + if not self.enterprise: + raise ValueError("Cannot delete database in Neo4j Community") + + if database_name == "neo4j": + raise ValueError("Cannot delete default 'neo4j' database") + + await self._ensure_connected() + + async with self.driver.session(database="system") as session: + await session.run(f"DROP DATABASE `{database_name}` IF EXISTS") + + logger.info(f"Deleted Neo4j database: {database_name}") + return True + + async def list_databases(self) -> List[str]: + await self._ensure_connected() + + if not self.enterprise: + return ["neo4j"] + + async with self.driver.session(database="system") as session: + result = await session.run("SHOW DATABASES") + records = await result.data() + databases = [record["name"] for record in records] + + return databases + + async def get_database_info(self, database_name: str) -> DatabaseInfo: + await self._ensure_connected() + + node_count = 0 + relationship_count = 0 + size_bytes = 0 + + try: + async with self.driver.session(database=database_name) as session: + node_result = await session.run("MATCH (n) RETURN count(n) as count") + node_data = await node_result.single() + if node_data: + node_count = node_data["count"] + + rel_result = await session.run("MATCH ()-[r]->() RETURN count(r) as count") + rel_data = await rel_result.single() + if rel_data: + relationship_count = rel_data["count"] + + if self.enterprise: + async with self.driver.session(database="system") as session: + size_result = await session.run( + f"SHOW DATABASE `{database_name}` YIELD sizeOnDisk" + ) + size_data = await size_result.single() + if size_data and "sizeOnDisk" in size_data: + size_bytes = size_data["sizeOnDisk"] + + except Exception as e: + logger.warning(f"Failed to get database stats for {database_name}: {e}") + + return DatabaseInfo( + name=database_name, + node_count=node_count, + relationship_count=relationship_count, + size_bytes=size_bytes, + ) + + async def get_cluster_topology(self) -> ClusterTopology: + await self._ensure_connected() + + if not self.enterprise: + return ClusterTopology(mode="single", leader={"url": self.bolt_url}) + + try: + async with self.driver.session(database="system") as session: + result = await session.run("CALL dbms.cluster.overview()") + records = await result.data() + + topology = [] + leader = None + followers = [] + + for record in records: + member = { + "id": record["id"], + "address": record["address"], + "role": record["role"], + "database": record.get("database", ""), + } + topology.append(member) + + if member["role"] == "LEADER": + leader = member + elif member["role"] == "FOLLOWER": + followers.append(member) + + return ClusterTopology( + mode="cluster", leader=leader, followers=followers, members=topology + ) + + except Exception as e: + logger.warning(f"Not running in cluster mode: {e}") + return ClusterTopology(mode="single", leader={"url": self.bolt_url}) + + async def health_check(self) -> bool: + try: + await self._ensure_connected() + + async with self.driver.session(database="system") as session: + await session.run("RETURN 1") + + return True + except Exception as e: + logger.error(f"Neo4j health check failed: {e}") + return False + + def close(self) -> None: + if self.driver: + import asyncio + + asyncio.create_task(self.driver.close()) diff --git a/robosystems/kuzu_api/cli.py b/robosystems/graph_api/cli.py similarity index 96% rename from robosystems/kuzu_api/cli.py rename to robosystems/graph_api/cli.py index e719b4d2..b64e69f8 100644 --- a/robosystems/kuzu_api/cli.py +++ b/robosystems/graph_api/cli.py @@ -2,8 +2,8 @@ Kuzu API Client - Command-line interface for Kuzu API Usage: - python -m robosystems.kuzu_api query "MATCH (c:Entity) RETURN c.name LIMIT 5" - python -m robosystems.kuzu_api health --url $KUZU_API_URL + python -m robosystems.graph_api query "MATCH (c:Entity) RETURN c.name LIMIT 5" + python -m robosystems.graph_api health --url $KUZU_API_URL This module provides a CLI interface using the new unified client structure. """ diff --git a/robosystems/kuzu_api/client/README.md b/robosystems/graph_api/client/README.md similarity index 97% rename from robosystems/kuzu_api/client/README.md rename to robosystems/graph_api/client/README.md index 60c32cf4..1cc4b4b3 100644 --- a/robosystems/kuzu_api/client/README.md +++ b/robosystems/graph_api/client/README.md @@ -124,7 +124,7 @@ circuit_breaker_timeout: 60 seconds ### Basic Usage ```python -from robosystems.kuzu_api.client.factory import KuzuClientFactory +from robosystems.graph_api.client.factory import KuzuClientFactory # Create client with automatic routing client = await KuzuClientFactory.create_client( @@ -179,7 +179,7 @@ async for chunk in result: ### Error Handling ```python -from robosystems.kuzu_api.client.exceptions import ( +from robosystems.graph_api.client.exceptions import ( KuzuTimeoutError, KuzuSyntaxError, ServiceUnavailableError @@ -308,7 +308,7 @@ Attempt 4: 4.0s + jitter (0-0.4s) ```python # Enable debug logging import logging -logging.getLogger("robosystems.kuzu_api").setLevel(logging.DEBUG) +logging.getLogger("robosystems.graph_api").setLevel(logging.DEBUG) # Logs include: - Routing decisions @@ -451,7 +451,7 @@ curl -X GET http://10.0.1.100:8001/health \ ```python from fastapi import APIRouter, Depends -from robosystems.kuzu_api.client.factory import KuzuClientFactory +from robosystems.graph_api.client.factory import KuzuClientFactory router = APIRouter() @@ -476,7 +476,7 @@ async def get_kuzu_client(graph_id: str): ```python from celery import shared_task -from robosystems.kuzu_api.client.factory import KuzuClientFactory +from robosystems.graph_api.client.factory import KuzuClientFactory @shared_task async def process_graph_data(graph_id: str, data_files: list): diff --git a/robosystems/kuzu_api/client/__init__.py b/robosystems/graph_api/client/__init__.py similarity index 100% rename from robosystems/kuzu_api/client/__init__.py rename to robosystems/graph_api/client/__init__.py diff --git a/robosystems/kuzu_api/client/base.py b/robosystems/graph_api/client/base.py similarity index 100% rename from robosystems/kuzu_api/client/base.py rename to robosystems/graph_api/client/base.py diff --git a/robosystems/kuzu_api/client/client.py b/robosystems/graph_api/client/client.py similarity index 100% rename from robosystems/kuzu_api/client/client.py rename to robosystems/graph_api/client/client.py diff --git a/robosystems/kuzu_api/client/config.py b/robosystems/graph_api/client/config.py similarity index 100% rename from robosystems/kuzu_api/client/config.py rename to robosystems/graph_api/client/config.py diff --git a/robosystems/kuzu_api/client/exceptions.py b/robosystems/graph_api/client/exceptions.py similarity index 100% rename from robosystems/kuzu_api/client/exceptions.py rename to robosystems/graph_api/client/exceptions.py diff --git a/robosystems/kuzu_api/client/factory.py b/robosystems/graph_api/client/factory.py similarity index 99% rename from robosystems/kuzu_api/client/factory.py rename to robosystems/graph_api/client/factory.py index c4f11be0..c2164a2b 100644 --- a/robosystems/kuzu_api/client/factory.py +++ b/robosystems/graph_api/client/factory.py @@ -19,7 +19,7 @@ from typing import Dict, Any from enum import Enum import redis.asyncio as redis -from robosystems.kuzu_api.client import KuzuClient +from robosystems.graph_api.client import KuzuClient from robosystems.config import env from robosystems.config.valkey_registry import ValkeyDatabase from robosystems.logger import logger diff --git a/robosystems/kuzu_api/client/sse_client.py b/robosystems/graph_api/client/sse_client.py similarity index 100% rename from robosystems/kuzu_api/client/sse_client.py rename to robosystems/graph_api/client/sse_client.py diff --git a/robosystems/kuzu_api/client/sync_client.py b/robosystems/graph_api/client/sync_client.py similarity index 100% rename from robosystems/kuzu_api/client/sync_client.py rename to robosystems/graph_api/client/sync_client.py diff --git a/robosystems/kuzu_api/core/__init__.py b/robosystems/graph_api/core/__init__.py similarity index 100% rename from robosystems/kuzu_api/core/__init__.py rename to robosystems/graph_api/core/__init__.py diff --git a/robosystems/kuzu_api/core/admission_control.py b/robosystems/graph_api/core/admission_control.py similarity index 100% rename from robosystems/kuzu_api/core/admission_control.py rename to robosystems/graph_api/core/admission_control.py diff --git a/robosystems/graph_api/core/backend_cluster_manager.py b/robosystems/graph_api/core/backend_cluster_manager.py new file mode 100644 index 00000000..62bdf237 --- /dev/null +++ b/robosystems/graph_api/core/backend_cluster_manager.py @@ -0,0 +1,128 @@ +import time +from datetime import datetime +from typing import Optional +import psutil +from fastapi import HTTPException, status + +from robosystems.graph_api.backends import get_backend +from robosystems.graph_api.models.database import QueryRequest, QueryResponse +from robosystems.graph_api.models.cluster import ( + ClusterHealthResponse, + ClusterInfoResponse, +) +from robosystems.logger import logger + + +class BackendClusterService: + def __init__(self): + self.backend = get_backend() + self.start_time = time.time() + self.last_activity: Optional[datetime] = None + logger.info( + f"Backend Cluster Service initialized with backend type: {type(self.backend).__name__}" + ) + + def get_uptime(self) -> float: + return time.time() - self.start_time + + async def execute_query(self, request: QueryRequest) -> QueryResponse: + start_time = time.time() + + try: + logger.debug(f"Executing query on {request.database}: {request.cypher[:100]}...") + + result = await self.backend.execute_query( + graph_id=request.database, + cypher=request.cypher, + parameters=request.parameters, + ) + + execution_time = (time.time() - start_time) * 1000 + self.last_activity = datetime.now() + + logger.info( + f"Query executed successfully on {request.database}: " + f"{len(result)} rows in {execution_time:.2f}ms" + ) + + columns = list(result[0].keys()) if result else [] + + return QueryResponse( + data=result, + columns=columns, + execution_time_ms=execution_time, + row_count=len(result), + database=request.database, + ) + + except HTTPException: + raise + except Exception as e: + execution_time = (time.time() - start_time) * 1000 + logger.error( + f"Query execution failed on {request.database}: {e} " + f"(after {execution_time:.2f}ms)" + ) + raise HTTPException( + status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, + detail=f"Query execution failed: {str(e)}", + ) + + async def get_cluster_health(self) -> ClusterHealthResponse: + try: + is_healthy = await self.backend.health_check() + uptime = self.get_uptime() + + cpu_usage = psutil.cpu_percent(interval=0.1) + memory_usage = psutil.virtual_memory().percent + + if not is_healthy: + status_str = "unhealthy" + elif cpu_usage > 90 or memory_usage > 90: + status_str = "critical" + elif cpu_usage > 75 or memory_usage > 75: + status_str = "warning" + else: + status_str = "healthy" + + return ClusterHealthResponse( + status=status_str, + uptime_seconds=uptime, + node_type="backend", + base_path="", + max_databases=0, + current_databases=0, + capacity_remaining=0, + read_only=False, + last_activity=(self.last_activity.isoformat() if self.last_activity else None), + ) + except Exception as e: + logger.error(f"Health check failed: {e}") + raise HTTPException( + status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, + detail=f"Health check failed: {str(e)}", + ) + + async def get_cluster_info(self) -> ClusterInfoResponse: + try: + await self.backend.get_cluster_topology() + databases = await self.backend.list_databases() + uptime = self.get_uptime() + + return ClusterInfoResponse( + node_id=f"backend-{type(self.backend).__name__}", + node_type="backend", + cluster_version="1.0.0", + base_path="", + max_databases=0, + databases=databases, + uptime_seconds=uptime, + read_only=False, + configuration=None, + ) + except Exception as e: + logger.error(f"Failed to get cluster info: {e}") + raise HTTPException( + status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, + detail=f"Failed to get cluster info: {str(e)}", + ) diff --git a/robosystems/kuzu_api/core/cluster_manager.py b/robosystems/graph_api/core/cluster_manager.py similarity index 99% rename from robosystems/kuzu_api/core/cluster_manager.py rename to robosystems/graph_api/core/cluster_manager.py index e6e00228..48aae7f7 100644 --- a/robosystems/kuzu_api/core/cluster_manager.py +++ b/robosystems/graph_api/core/cluster_manager.py @@ -19,12 +19,12 @@ from robosystems.config import env from .database_manager import KuzuDatabaseManager from .metrics_collector import KuzuMetricsCollector -from robosystems.kuzu_api.models.database import QueryRequest, QueryResponse -from robosystems.kuzu_api.models.cluster import ( +from robosystems.graph_api.models.database import QueryRequest, QueryResponse +from robosystems.graph_api.models.cluster import ( ClusterHealthResponse, ClusterInfoResponse, ) -from robosystems.kuzu_api.core.utils import ( +from robosystems.graph_api.core.utils import ( validate_database_name, validate_query_parameters, ) @@ -765,7 +765,7 @@ def get_cluster_health(self) -> ClusterHealthResponse: def get_cluster_info(self) -> ClusterInfoResponse: """Get detailed cluster information including all configuration parameters.""" from robosystems.config import env - from robosystems.kuzu_api.models.cluster import ( + from robosystems.graph_api.models.cluster import ( MemoryConfiguration, QueryConfiguration, AdmissionControlConfig, diff --git a/robosystems/kuzu_api/core/connection_pool.py b/robosystems/graph_api/core/connection_pool.py similarity index 100% rename from robosystems/kuzu_api/core/connection_pool.py rename to robosystems/graph_api/core/connection_pool.py diff --git a/robosystems/kuzu_api/core/database_manager.py b/robosystems/graph_api/core/database_manager.py similarity index 100% rename from robosystems/kuzu_api/core/database_manager.py rename to robosystems/graph_api/core/database_manager.py diff --git a/robosystems/kuzu_api/core/metrics_collector.py b/robosystems/graph_api/core/metrics_collector.py similarity index 100% rename from robosystems/kuzu_api/core/metrics_collector.py rename to robosystems/graph_api/core/metrics_collector.py diff --git a/robosystems/kuzu_api/core/task_manager.py b/robosystems/graph_api/core/task_manager.py similarity index 100% rename from robosystems/kuzu_api/core/task_manager.py rename to robosystems/graph_api/core/task_manager.py diff --git a/robosystems/kuzu_api/core/task_sse.py b/robosystems/graph_api/core/task_sse.py similarity index 100% rename from robosystems/kuzu_api/core/task_sse.py rename to robosystems/graph_api/core/task_sse.py diff --git a/robosystems/kuzu_api/core/utils.py b/robosystems/graph_api/core/utils.py similarity index 100% rename from robosystems/kuzu_api/core/utils.py rename to robosystems/graph_api/core/utils.py diff --git a/robosystems/kuzu_api/main.py b/robosystems/graph_api/main.py similarity index 97% rename from robosystems/kuzu_api/main.py rename to robosystems/graph_api/main.py index 6efcf13c..8f62dd66 100644 --- a/robosystems/kuzu_api/main.py +++ b/robosystems/graph_api/main.py @@ -9,8 +9,8 @@ import uvicorn -from robosystems.kuzu_api.app import create_app -from robosystems.kuzu_api.core import init_cluster_service +from robosystems.graph_api.app import create_app +from robosystems.graph_api.core import init_cluster_service from robosystems.middleware.graph.clusters import NodeType, RepositoryType from robosystems.logger import logger diff --git a/robosystems/kuzu_api/middleware/__init__.py b/robosystems/graph_api/middleware/__init__.py similarity index 100% rename from robosystems/kuzu_api/middleware/__init__.py rename to robosystems/graph_api/middleware/__init__.py diff --git a/robosystems/kuzu_api/middleware/auth.py b/robosystems/graph_api/middleware/auth.py similarity index 100% rename from robosystems/kuzu_api/middleware/auth.py rename to robosystems/graph_api/middleware/auth.py diff --git a/robosystems/kuzu_api/middleware/request_limits.py b/robosystems/graph_api/middleware/request_limits.py similarity index 100% rename from robosystems/kuzu_api/middleware/request_limits.py rename to robosystems/graph_api/middleware/request_limits.py diff --git a/robosystems/kuzu_api/models/__init__.py b/robosystems/graph_api/models/__init__.py similarity index 100% rename from robosystems/kuzu_api/models/__init__.py rename to robosystems/graph_api/models/__init__.py diff --git a/robosystems/kuzu_api/models/cluster.py b/robosystems/graph_api/models/cluster.py similarity index 100% rename from robosystems/kuzu_api/models/cluster.py rename to robosystems/graph_api/models/cluster.py diff --git a/robosystems/kuzu_api/models/database.py b/robosystems/graph_api/models/database.py similarity index 100% rename from robosystems/kuzu_api/models/database.py rename to robosystems/graph_api/models/database.py diff --git a/robosystems/kuzu_api/models/ingestion.py b/robosystems/graph_api/models/ingestion.py similarity index 100% rename from robosystems/kuzu_api/models/ingestion.py rename to robosystems/graph_api/models/ingestion.py diff --git a/robosystems/kuzu_api/models/streaming.py b/robosystems/graph_api/models/streaming.py similarity index 100% rename from robosystems/kuzu_api/models/streaming.py rename to robosystems/graph_api/models/streaming.py diff --git a/robosystems/kuzu_api/routers/__init__.py b/robosystems/graph_api/routers/__init__.py similarity index 100% rename from robosystems/kuzu_api/routers/__init__.py rename to robosystems/graph_api/routers/__init__.py diff --git a/robosystems/kuzu_api/routers/databases/__init__.py b/robosystems/graph_api/routers/databases/__init__.py similarity index 100% rename from robosystems/kuzu_api/routers/databases/__init__.py rename to robosystems/graph_api/routers/databases/__init__.py diff --git a/robosystems/kuzu_api/routers/databases/backup.py b/robosystems/graph_api/routers/databases/backup.py similarity index 94% rename from robosystems/kuzu_api/routers/databases/backup.py rename to robosystems/graph_api/routers/databases/backup.py index 3ce2b7f0..96ca1340 100644 --- a/robosystems/kuzu_api/routers/databases/backup.py +++ b/robosystems/graph_api/routers/databases/backup.py @@ -9,9 +9,9 @@ from fastapi import APIRouter, BackgroundTasks, Depends, HTTPException, Path from fastapi import status as http_status -from robosystems.kuzu_api.models.database import BackupRequest, BackupResponse -from robosystems.kuzu_api.core.cluster_manager import get_cluster_service -from robosystems.kuzu_api.core.task_manager import backup_task_manager +from robosystems.graph_api.models.database import BackupRequest, BackupResponse +from robosystems.graph_api.core.cluster_manager import get_cluster_service +from robosystems.graph_api.core.task_manager import backup_task_manager from robosystems.logger import logger router = APIRouter(prefix="/databases", tags=["Database Backup"]) diff --git a/robosystems/kuzu_api/routers/databases/ingest.py b/robosystems/graph_api/routers/databases/ingest.py similarity index 99% rename from robosystems/kuzu_api/routers/databases/ingest.py rename to robosystems/graph_api/routers/databases/ingest.py index 63071f04..a75698a6 100644 --- a/robosystems/kuzu_api/routers/databases/ingest.py +++ b/robosystems/graph_api/routers/databases/ingest.py @@ -25,8 +25,8 @@ ) from pydantic import BaseModel, Field -from robosystems.kuzu_api.core.cluster_manager import get_cluster_service -from robosystems.kuzu_api.core.connection_pool import get_connection_pool +from robosystems.graph_api.core.cluster_manager import get_cluster_service +from robosystems.graph_api.core.connection_pool import get_connection_pool from robosystems.logger import logger from robosystems.config.valkey_registry import ValkeyDatabase from robosystems.config import env diff --git a/robosystems/kuzu_api/routers/databases/management.py b/robosystems/graph_api/routers/databases/management.py similarity index 96% rename from robosystems/kuzu_api/routers/databases/management.py rename to robosystems/graph_api/routers/databases/management.py index 53a97490..16c097a5 100644 --- a/robosystems/kuzu_api/routers/databases/management.py +++ b/robosystems/graph_api/routers/databases/management.py @@ -8,13 +8,13 @@ from fastapi import APIRouter, Depends, HTTPException, Path from fastapi import status as http_status -from robosystems.kuzu_api.models.database import ( +from robosystems.graph_api.models.database import ( DatabaseCreateRequest, DatabaseCreateResponse, DatabaseListResponse, DatabaseInfo, ) -from robosystems.kuzu_api.core.cluster_manager import get_cluster_service +from robosystems.graph_api.core.cluster_manager import get_cluster_service from robosystems.middleware.graph.clusters import NodeType from robosystems.logger import logger diff --git a/robosystems/kuzu_api/routers/databases/metrics.py b/robosystems/graph_api/routers/databases/metrics.py similarity index 95% rename from robosystems/kuzu_api/routers/databases/metrics.py rename to robosystems/graph_api/routers/databases/metrics.py index 44e0d1c0..0eb2c765 100644 --- a/robosystems/kuzu_api/routers/databases/metrics.py +++ b/robosystems/graph_api/routers/databases/metrics.py @@ -9,8 +9,8 @@ from fastapi import APIRouter, Depends, HTTPException, Path from fastapi import status as http_status -from robosystems.kuzu_api.core.cluster_manager import get_cluster_service -from robosystems.kuzu_api.core.utils import validate_database_name +from robosystems.graph_api.core.cluster_manager import get_cluster_service +from robosystems.graph_api.core.utils import validate_database_name from robosystems.logger import logger router = APIRouter(prefix="/databases", tags=["Database Metrics"]) diff --git a/robosystems/kuzu_api/routers/databases/query.py b/robosystems/graph_api/routers/databases/query.py similarity index 60% rename from robosystems/kuzu_api/routers/databases/query.py rename to robosystems/graph_api/routers/databases/query.py index 02db6f09..fa742c43 100644 --- a/robosystems/kuzu_api/routers/databases/query.py +++ b/robosystems/graph_api/routers/databases/query.py @@ -10,13 +10,14 @@ import json from contextlib import contextmanager -from robosystems.kuzu_api.models.database import QueryRequest -from robosystems.kuzu_api.core.cluster_manager import get_cluster_service -from robosystems.kuzu_api.core.admission_control import ( +from robosystems.graph_api.models.database import QueryRequest +from robosystems.graph_api.core.cluster_manager import get_cluster_service +from robosystems.graph_api.core.admission_control import ( get_admission_controller, AdmissionDecision, ) from robosystems.logger import logger +from robosystems.config import env router = APIRouter(prefix="/databases", tags=["Database Queries"]) @@ -31,12 +32,22 @@ def track_connection(admission_controller, database_name): admission_controller.release_connection(database_name) +def _get_cluster_service_for_request(): + backend_type = env.BACKEND_TYPE + if backend_type in ["neo4j_community", "neo4j_enterprise"]: + from robosystems.graph_api.core.backend_cluster_manager import BackendClusterService + + return BackendClusterService() + return get_cluster_service() + + @router.post("/{graph_id}/query") async def execute_query( request: QueryRequest, graph_id: str = Path(..., description="Graph database identifier"), streaming: bool = False, - cluster_service=Depends(get_cluster_service), + database: str = None, + cluster_service=Depends(_get_cluster_service_for_request), ): """ Execute a Cypher query against a specific database with admission control. @@ -81,20 +92,32 @@ async def execute_query( database=graph_id, cypher=request.cypher, parameters=request.parameters ) - if not streaming: - # Standard response - note that execute_query now has internal limits - # to prevent memory issues (MAX_ROWS = 10000) - return cluster_service.execute_query(query_request) - - # Streaming response for large result sets - def generate_stream(): - for chunk in cluster_service.execute_query_streaming( - query_request, chunk_size=1000 - ): - yield json.dumps(chunk) + "\n" - - return StreamingResponse( - generate_stream(), - media_type="application/x-ndjson", - headers={"X-Streaming": "true", "Cache-Control": "no-cache"}, - ) + backend_type = env.BACKEND_TYPE + + if backend_type in ["neo4j_community", "neo4j_enterprise"]: + # Use async backend cluster service + if not streaming: + return await cluster_service.execute_query(query_request) + else: + # Streaming not yet implemented for Neo4j backend + raise HTTPException( + status_code=status.HTTP_501_NOT_IMPLEMENTED, + detail="Streaming not yet implemented for Neo4j backend", + ) + else: + # Use existing Kuzu cluster service (sync) + if not streaming: + return cluster_service.execute_query(query_request) + + # Streaming response for large result sets + def generate_stream(): + for chunk in cluster_service.execute_query_streaming( + query_request, chunk_size=1000 + ): + yield json.dumps(chunk) + "\n" + + return StreamingResponse( + generate_stream(), + media_type="application/x-ndjson", + headers={"X-Streaming": "true", "Cache-Control": "no-cache"}, + ) diff --git a/robosystems/kuzu_api/routers/databases/restore.py b/robosystems/graph_api/routers/databases/restore.py similarity index 96% rename from robosystems/kuzu_api/routers/databases/restore.py rename to robosystems/graph_api/routers/databases/restore.py index 4408450a..8596edd2 100644 --- a/robosystems/kuzu_api/routers/databases/restore.py +++ b/robosystems/graph_api/routers/databases/restore.py @@ -17,9 +17,9 @@ ) from fastapi import status as http_status -from robosystems.kuzu_api.models.database import RestoreResponse -from robosystems.kuzu_api.core.cluster_manager import get_cluster_service -from robosystems.kuzu_api.core.task_manager import restore_task_manager +from robosystems.graph_api.models.database import RestoreResponse +from robosystems.graph_api.core.cluster_manager import get_cluster_service +from robosystems.graph_api.core.task_manager import restore_task_manager from robosystems.logger import logger router = APIRouter(prefix="/databases", tags=["Database Backup"]) diff --git a/robosystems/kuzu_api/routers/databases/schema.py b/robosystems/graph_api/routers/databases/schema.py similarity index 98% rename from robosystems/kuzu_api/routers/databases/schema.py rename to robosystems/graph_api/routers/databases/schema.py index f21404dd..75263d40 100644 --- a/robosystems/kuzu_api/routers/databases/schema.py +++ b/robosystems/graph_api/routers/databases/schema.py @@ -10,13 +10,13 @@ from fastapi import APIRouter, Depends, HTTPException, Path from fastapi import status as http_status -from robosystems.kuzu_api.models.database import ( +from robosystems.graph_api.models.database import ( SchemaInstallRequest, SchemaInstallResponse, QueryRequest, ) -from robosystems.kuzu_api.core.cluster_manager import get_cluster_service -from robosystems.kuzu_api.core.utils import validate_database_name +from robosystems.graph_api.core.cluster_manager import get_cluster_service +from robosystems.graph_api.core.utils import validate_database_name from robosystems.logger import logger router = APIRouter(prefix="/databases", tags=["Database Schema"]) diff --git a/robosystems/kuzu_api/routers/health.py b/robosystems/graph_api/routers/health.py similarity index 95% rename from robosystems/kuzu_api/routers/health.py rename to robosystems/graph_api/routers/health.py index 98a76a68..c7e652e1 100644 --- a/robosystems/kuzu_api/routers/health.py +++ b/robosystems/graph_api/routers/health.py @@ -8,7 +8,7 @@ from fastapi import APIRouter, Depends from fastapi.responses import JSONResponse -from robosystems.kuzu_api.core.cluster_manager import get_cluster_service +from robosystems.graph_api.core.cluster_manager import get_cluster_service router = APIRouter(tags=["Cluster Health"]) diff --git a/robosystems/kuzu_api/routers/info.py b/robosystems/graph_api/routers/info.py similarity index 83% rename from robosystems/kuzu_api/routers/info.py rename to robosystems/graph_api/routers/info.py index b35d6a16..5a23904f 100644 --- a/robosystems/kuzu_api/routers/info.py +++ b/robosystems/graph_api/routers/info.py @@ -7,8 +7,8 @@ from fastapi import APIRouter, Depends -from robosystems.kuzu_api.models.cluster import ClusterInfoResponse -from robosystems.kuzu_api.core.cluster_manager import get_cluster_service +from robosystems.graph_api.models.cluster import ClusterInfoResponse +from robosystems.graph_api.core.cluster_manager import get_cluster_service router = APIRouter(tags=["Cluster Info"]) diff --git a/robosystems/kuzu_api/routers/metrics.py b/robosystems/graph_api/routers/metrics.py similarity index 92% rename from robosystems/kuzu_api/routers/metrics.py rename to robosystems/graph_api/routers/metrics.py index 82e5dca3..03949856 100644 --- a/robosystems/kuzu_api/routers/metrics.py +++ b/robosystems/graph_api/routers/metrics.py @@ -9,8 +9,8 @@ from typing import Dict, Any from fastapi import APIRouter, Depends -from robosystems.kuzu_api.core.cluster_manager import get_cluster_service -from robosystems.kuzu_api.core.admission_control import get_admission_controller +from robosystems.graph_api.core.cluster_manager import get_cluster_service +from robosystems.graph_api.core.admission_control import get_admission_controller router = APIRouter(tags=["Cluster Metrics"]) diff --git a/robosystems/kuzu_api/routers/tasks.py b/robosystems/graph_api/routers/tasks.py similarity index 97% rename from robosystems/kuzu_api/routers/tasks.py rename to robosystems/graph_api/routers/tasks.py index f881db85..c99ca069 100644 --- a/robosystems/kuzu_api/routers/tasks.py +++ b/robosystems/graph_api/routers/tasks.py @@ -10,11 +10,11 @@ from fastapi import status as http_status from sse_starlette.sse import EventSourceResponse -from robosystems.kuzu_api.core.task_sse import generate_task_sse_events, TaskType -from robosystems.kuzu_api.routers.databases.ingest import ( +from robosystems.graph_api.core.task_sse import generate_task_sse_events, TaskType +from robosystems.graph_api.routers.databases.ingest import ( task_manager as ingestion_task_manager, ) -from robosystems.kuzu_api.core.task_manager import ( +from robosystems.graph_api.core.task_manager import ( backup_task_manager, restore_task_manager, ) diff --git a/robosystems/middleware/graph/clusters.py b/robosystems/middleware/graph/clusters.py index 90be3003..47026ebf 100644 --- a/robosystems/middleware/graph/clusters.py +++ b/robosystems/middleware/graph/clusters.py @@ -7,7 +7,7 @@ correctly in production environments with multiple EC2 instances. For entity graphs, use: - from robosystems.kuzu_api.client import get_kuzu_client + from robosystems.graph_api.client import get_kuzu_client client = await get_kuzu_client(graph_id) This module defines cluster configurations for single-writer, multiple-reader architecture. diff --git a/robosystems/middleware/graph/multitenant_utils.py b/robosystems/middleware/graph/multitenant_utils.py index 615cd437..38f4fff6 100644 --- a/robosystems/middleware/graph/multitenant_utils.py +++ b/robosystems/middleware/graph/multitenant_utils.py @@ -599,8 +599,8 @@ async def ensure_database_with_schema( Raises: RuntimeError: If creation fails """ - from robosystems.kuzu_api.client import KuzuClient - from robosystems.kuzu_api.client.exceptions import KuzuClientError + from robosystems.graph_api.client import KuzuClient + from robosystems.graph_api.client.exceptions import KuzuClientError # Create Kuzu client headers = {} diff --git a/robosystems/middleware/graph/repository.py b/robosystems/middleware/graph/repository.py index a0052ef7..5a6aa895 100644 --- a/robosystems/middleware/graph/repository.py +++ b/robosystems/middleware/graph/repository.py @@ -19,7 +19,7 @@ from .base import GraphOperation from .engine import Repository -from robosystems.kuzu_api.client import KuzuClient +from robosystems.graph_api.client import KuzuClient from robosystems.logger import logger diff --git a/robosystems/middleware/graph/router.py b/robosystems/middleware/graph/router.py index ab59e3da..44fefd54 100644 --- a/robosystems/middleware/graph/router.py +++ b/robosystems/middleware/graph/router.py @@ -67,7 +67,7 @@ async def get_repository( return Repository(database_path) else: # Use the new enhanced client factory for all routing - from robosystems.kuzu_api.client.factory import get_kuzu_client + from robosystems.graph_api.client.factory import get_kuzu_client from .streaming_wrapper import add_streaming_support logger.debug(f"Using enhanced client factory for {graph_id}") @@ -144,7 +144,7 @@ def _create_test_repository(self, cluster_config: ClusterConfig): api_key = env.KUZU_API_KEY - from robosystems.kuzu_api.client import KuzuClient + from robosystems.graph_api.client import KuzuClient client = KuzuClient(base_url=cluster_config.alb_endpoint, api_key=api_key) client.graph_id = "test" diff --git a/robosystems/middleware/mcp/client.py b/robosystems/middleware/mcp/client.py index d493af76..72873528 100644 --- a/robosystems/middleware/mcp/client.py +++ b/robosystems/middleware/mcp/client.py @@ -16,7 +16,7 @@ from robosystems.logger import logger from robosystems.config import env -from robosystems.kuzu_api.client import KuzuClient +from robosystems.graph_api.client import KuzuClient from .exceptions import KuzuAPIError, KuzuQueryTimeoutError, KuzuQueryComplexityError diff --git a/robosystems/middleware/mcp/factory.py b/robosystems/middleware/mcp/factory.py index 3a589b44..077e4aba 100644 --- a/robosystems/middleware/mcp/factory.py +++ b/robosystems/middleware/mcp/factory.py @@ -29,7 +29,7 @@ async def create_kuzu_mcp_client( """ # If URL not provided, use KuzuClientFactory to discover the proper endpoint if not api_base_url: - from robosystems.kuzu_api.client.factory import KuzuClientFactory + from robosystems.graph_api.client.factory import KuzuClientFactory from robosystems.middleware.graph.multitenant_utils import MultiTenantUtils # Determine operation type based on graph diff --git a/robosystems/middleware/robustness/README.md b/robosystems/middleware/robustness/README.md index 7b26f3e6..11bf2633 100644 --- a/robosystems/middleware/robustness/README.md +++ b/robosystems/middleware/robustness/README.md @@ -162,13 +162,13 @@ tracker = HealthTracker() # Register service tracker.register_service( - name="kuzu_api", + name="graph_api", health_check=lambda: check_kuzu_health(), interval=30 # Check every 30 seconds ) # Get service status -status = tracker.get_status("kuzu_api") +status = tracker.get_status("graph_api") print(f"Service health: {status.health_score}%") print(f"Average latency: {status.avg_latency}ms") ``` diff --git a/robosystems/operations/graph/entity_graph_service.py b/robosystems/operations/graph/entity_graph_service.py index 51f636e3..958c8686 100644 --- a/robosystems/operations/graph/entity_graph_service.py +++ b/robosystems/operations/graph/entity_graph_service.py @@ -22,8 +22,8 @@ from ...models.iam import UserGraph, UserLimits from ...config import env from ...models.api import EntityCreate, EntityResponse -from ...kuzu_api.client import KuzuClient, get_kuzu_client_for_instance -from ...kuzu_api.client.exceptions import KuzuClientError +from ...graph_api.client import KuzuClient, get_kuzu_client_for_instance +from ...graph_api.client.exceptions import KuzuClientError from ...middleware.graph.allocation_manager import KuzuAllocationManager from ...middleware.graph.types import InstanceTier from ...exceptions import ( diff --git a/robosystems/operations/graph/generic_graph_service.py b/robosystems/operations/graph/generic_graph_service.py index 0484895d..3fc03edb 100644 --- a/robosystems/operations/graph/generic_graph_service.py +++ b/robosystems/operations/graph/generic_graph_service.py @@ -163,7 +163,7 @@ async def create_graph( logger.info(f"Creating database {graph_id} on Kuzu writer") # Use KuzuClient with proper API key - from ...kuzu_api.client import get_kuzu_client_for_instance + from ...graph_api.client import get_kuzu_client_for_instance kuzu_client = await get_kuzu_client_for_instance(cluster_info.private_ip) diff --git a/robosystems/operations/graph/shared_repository_service.py b/robosystems/operations/graph/shared_repository_service.py index d707082c..0c0ff5dd 100644 --- a/robosystems/operations/graph/shared_repository_service.py +++ b/robosystems/operations/graph/shared_repository_service.py @@ -16,7 +16,7 @@ from datetime import datetime, timezone from ...logger import logger -from ...kuzu_api.client import get_kuzu_client_for_instance +from ...graph_api.client import get_kuzu_client_for_instance from ...config import env @@ -54,7 +54,7 @@ async def create_shared_repository( try: # For shared repositories, connect to the shared master directly # The shared master is already running and registered in DynamoDB - from ...kuzu_api.client.factory import KuzuClientFactory + from ...graph_api.client.factory import KuzuClientFactory # Get the shared master URL from DynamoDB discovery shared_master_url = await KuzuClientFactory._get_shared_master_url() @@ -135,7 +135,7 @@ async def ensure_shared_repository_exists( # Try to get database info first try: - from ...kuzu_api.client.factory import KuzuClientFactory + from ...graph_api.client.factory import KuzuClientFactory # Use factory to get proper client for shared repository # Shared repositories automatically route to shared_master/shared_replica infrastructure diff --git a/robosystems/operations/graph/subgraph_service.py b/robosystems/operations/graph/subgraph_service.py index b76fee56..e16f796e 100644 --- a/robosystems/operations/graph/subgraph_service.py +++ b/robosystems/operations/graph/subgraph_service.py @@ -25,7 +25,7 @@ validate_subgraph_name, validate_parent_graph_id, ) -from ...kuzu_api.client.factory import get_kuzu_client_for_instance +from ...graph_api.client.factory import get_kuzu_client_for_instance from ...exceptions import GraphAllocationError from ...logger import logger diff --git a/robosystems/routers/graphs/copy/execute.py b/robosystems/routers/graphs/copy/execute.py index 1461a9eb..e612755b 100644 --- a/robosystems/routers/graphs/copy/execute.py +++ b/robosystems/routers/graphs/copy/execute.py @@ -21,7 +21,7 @@ from robosystems.database import get_db_session from robosystems.middleware.auth.dependencies import get_current_user -from robosystems.kuzu_api.client.factory import KuzuClientFactory +from robosystems.graph_api.client.factory import KuzuClientFactory from robosystems.middleware.graph.types import InstanceTier from robosystems.middleware.rate_limits import ( subscription_aware_rate_limit_dependency, diff --git a/robosystems/routers/graphs/copy/strategies.py b/robosystems/routers/graphs/copy/strategies.py index e71501af..2fb19756 100644 --- a/robosystems/routers/graphs/copy/strategies.py +++ b/robosystems/routers/graphs/copy/strategies.py @@ -9,7 +9,7 @@ from typing import Dict, Any import uuid -from robosystems.kuzu_api.client import KuzuClient +from robosystems.graph_api.client import KuzuClient from robosystems.security import SecurityAuditLogger, SecurityEventType from robosystems.logger import logger from robosystems.middleware.sse.event_storage import ( diff --git a/robosystems/routers/graphs/health.py b/robosystems/routers/graphs/health.py index 58c40f05..87dcf6bc 100644 --- a/robosystems/routers/graphs/health.py +++ b/robosystems/routers/graphs/health.py @@ -17,7 +17,7 @@ from robosystems.middleware.graph.dependencies import get_universal_repository_with_auth from robosystems.models.api.graph import DatabaseHealthResponse from robosystems.middleware.otel.metrics import endpoint_metrics_decorator -from robosystems.kuzu_api.client import KuzuClient +from robosystems.graph_api.client import KuzuClient from robosystems.logger import logger from robosystems.middleware.robustness import ( CircuitBreakerManager, @@ -34,7 +34,7 @@ async def _get_kuzu_client(graph_id: str) -> KuzuClient: """Get Kuzu client for the specified graph using factory for endpoint discovery.""" - from robosystems.kuzu_api.client.factory import KuzuClientFactory + from robosystems.graph_api.client.factory import KuzuClientFactory from robosystems.middleware.graph.multitenant_utils import MultiTenantUtils # Determine operation type based on graph diff --git a/robosystems/routers/graphs/info.py b/robosystems/routers/graphs/info.py index 8d8274c1..0940fdc9 100644 --- a/robosystems/routers/graphs/info.py +++ b/robosystems/routers/graphs/info.py @@ -18,7 +18,7 @@ from robosystems.middleware.graph.dependencies import get_universal_repository_with_auth from robosystems.models.api.graph import DatabaseInfoResponse from robosystems.middleware.otel.metrics import endpoint_metrics_decorator -from robosystems.kuzu_api.client import KuzuClient +from robosystems.graph_api.client import KuzuClient from robosystems.logger import logger from robosystems.middleware.robustness import ( CircuitBreakerManager, @@ -35,7 +35,7 @@ async def _get_kuzu_client(graph_id: str) -> KuzuClient: """Get Kuzu client for the specified graph using factory for endpoint discovery.""" - from robosystems.kuzu_api.client.factory import KuzuClientFactory + from robosystems.graph_api.client.factory import KuzuClientFactory from robosystems.middleware.graph.multitenant_utils import MultiTenantUtils # Determine operation type based on graph diff --git a/robosystems/routers/graphs/limits.py b/robosystems/routers/graphs/limits.py index 81a7a8f5..1762b6d4 100644 --- a/robosystems/routers/graphs/limits.py +++ b/robosystems/routers/graphs/limits.py @@ -17,7 +17,7 @@ ) from robosystems.middleware.graph.dependencies import get_universal_repository_with_auth from robosystems.middleware.otel.metrics import endpoint_metrics_decorator -from robosystems.kuzu_api.client import KuzuClient +from robosystems.graph_api.client import KuzuClient from robosystems.logger import logger from robosystems.middleware.robustness import ( CircuitBreakerManager, @@ -34,7 +34,7 @@ async def _get_kuzu_client(graph_id: str) -> KuzuClient: """Get Kuzu client for the specified graph using factory for endpoint discovery.""" - from robosystems.kuzu_api.client.factory import KuzuClientFactory + from robosystems.graph_api.client.factory import KuzuClientFactory from robosystems.middleware.graph.multitenant_utils import MultiTenantUtils # Determine operation type based on graph diff --git a/robosystems/tasks/billing/usage_collector.py b/robosystems/tasks/billing/usage_collector.py index 69a62ac5..7a86a3dc 100644 --- a/robosystems/tasks/billing/usage_collector.py +++ b/robosystems/tasks/billing/usage_collector.py @@ -119,7 +119,7 @@ def graph_usage_collector(self): async def collect_graph_metrics(graph_id: str) -> Dict: """Collect usage metrics for a specific graph using KuzuClientFactory.""" - from robosystems.kuzu_api.client.factory import KuzuClientFactory + from robosystems.graph_api.client.factory import KuzuClientFactory try: # Use factory to create client with proper authentication and routing diff --git a/robosystems/tasks/graph_operations/backup.py b/robosystems/tasks/graph_operations/backup.py index a711c2b2..e859b9b0 100644 --- a/robosystems/tasks/graph_operations/backup.py +++ b/robosystems/tasks/graph_operations/backup.py @@ -118,7 +118,7 @@ def create_graph_backup( # Call Kuzu API to create the backup logger.info(f"Calling Kuzu API to create backup for graph '{graph_id}'") - from robosystems.kuzu_api.client.factory import get_kuzu_client + from robosystems.graph_api.client.factory import get_kuzu_client # Get properly routed Kuzu client client = asyncio.run(get_kuzu_client(graph_id, operation_type="read")) @@ -548,7 +548,7 @@ def restore_graph_backup( # Call Kuzu API to restore the backup logger.info(f"Calling Kuzu API to restore database for graph '{graph_id}'") - from ...kuzu_api.client.factory import get_kuzu_client_sync + from ...graph_api.client.factory import get_kuzu_client_sync # Get properly routed Kuzu client client = get_kuzu_client_sync(graph_id, operation_type="write") @@ -589,7 +589,7 @@ def restore_graph_backup( if verify_after_restore: try: # Create a simple verification by checking if we can connect and query - from ...kuzu_api.client.factory import get_kuzu_client + from ...graph_api.client.factory import get_kuzu_client client = asyncio.run(get_kuzu_client(graph_id, operation_type="read")) @@ -1104,7 +1104,7 @@ def restore_graph_backup_sse( progress_tracker.emit_progress("Downloading backup from storage...", 40) # Use Kuzu API client for restoration - from ...kuzu_api.client.factory import get_kuzu_client + from ...graph_api.client.factory import get_kuzu_client client = asyncio.run(get_kuzu_client(graph_id, operation_type="write")) diff --git a/robosystems/tasks/sec_xbrl/ingestion.py b/robosystems/tasks/sec_xbrl/ingestion.py index f66dd365..234b0b2f 100644 --- a/robosystems/tasks/sec_xbrl/ingestion.py +++ b/robosystems/tasks/sec_xbrl/ingestion.py @@ -190,7 +190,7 @@ def ingest_sec_data( # IMPORTANT: We do NOT create or recreate databases here - only verify logger.info("Verifying SEC database exists with proper schema...") try: - from robosystems.kuzu_api.client.factory import KuzuClientFactory + from robosystems.graph_api.client.factory import KuzuClientFactory import asyncio async def verify_database(): @@ -696,7 +696,7 @@ async def _bulk_load_node_type( ignore_errors: bool, ) -> Dict[str, Any]: """Bulk load node type using S3 glob pattern via SYNC mode with direct S3 COPY.""" - from robosystems.kuzu_api.client.factory import KuzuClientFactory + from robosystems.graph_api.client.factory import KuzuClientFactory import boto3 from urllib.parse import urlparse @@ -901,7 +901,7 @@ async def _bulk_load_relationship_type( ignore_errors: bool, ) -> Dict[str, Any]: """Bulk load relationship type using S3 glob pattern via SYNC mode with direct S3 COPY.""" - from robosystems.kuzu_api.client.factory import KuzuClientFactory + from robosystems.graph_api.client.factory import KuzuClientFactory import boto3 from urllib.parse import urlparse diff --git a/robosystems/tasks/sec_xbrl/maintenance.py b/robosystems/tasks/sec_xbrl/maintenance.py index 68c19489..fb0da6e2 100644 --- a/robosystems/tasks/sec_xbrl/maintenance.py +++ b/robosystems/tasks/sec_xbrl/maintenance.py @@ -54,7 +54,7 @@ def reset_sec_database(confirm: bool = False) -> Dict: async def reset_database(): """Async function to reset the database.""" - from robosystems.kuzu_api.client.factory import KuzuClientFactory + from robosystems.graph_api.client.factory import KuzuClientFactory try: # Get a client for the SEC database diff --git a/robosystems/tasks/sec_xbrl/orchestration.py b/robosystems/tasks/sec_xbrl/orchestration.py index 0c5b3efc..99f0bd36 100644 --- a/robosystems/tasks/sec_xbrl/orchestration.py +++ b/robosystems/tasks/sec_xbrl/orchestration.py @@ -585,7 +585,7 @@ def cleanup_phase_connections(phase: str) -> Dict: # Clear Kuzu client factory connection pools try: - from robosystems.kuzu_api.client.factory import KuzuClientFactory + from robosystems.graph_api.client.factory import KuzuClientFactory if hasattr(KuzuClientFactory, "_connection_pools"): pool_count = len(KuzuClientFactory._connection_pools) diff --git a/tests/README.md b/tests/README.md index d4b190bf..a1fa7c7d 100644 --- a/tests/README.md +++ b/tests/README.md @@ -7,7 +7,7 @@ This directory contains tests for the RoboSystems Service application. The tests - `adapters/` - Tests for external service adapters (MCP, S3) - `processors/` - Tests for data transformation processors (XBRL, QuickBooks, schedules) - `integration/` - End-to-end integration tests -- `kuzu_api/` - Tests for Kuzu database cluster services +- `graph_api/` - Tests for Kuzu database cluster services - `middleware/` - Tests for middleware components (OpenTelemetry metrics) - `models/` - Tests for database models (IAM, graph, financial entities) - `operations/` - Tests for business logic services diff --git a/tests/kuzu_api/client/test_base.py b/tests/graph_api/client/test_base.py similarity index 97% rename from tests/kuzu_api/client/test_base.py rename to tests/graph_api/client/test_base.py index 0d25ffc2..13d64aad 100644 --- a/tests/kuzu_api/client/test_base.py +++ b/tests/graph_api/client/test_base.py @@ -4,9 +4,9 @@ from unittest.mock import patch import pytest -from robosystems.kuzu_api.client.base import BaseKuzuClient -from robosystems.kuzu_api.client.config import KuzuClientConfig -from robosystems.kuzu_api.client.exceptions import ( +from robosystems.graph_api.client.base import BaseKuzuClient +from robosystems.graph_api.client.config import KuzuClientConfig +from robosystems.graph_api.client.exceptions import ( KuzuAPIError, KuzuTransientError, KuzuClientError, @@ -312,7 +312,7 @@ def test_initialization_warns_missing_api_key_in_prod(self): mock_env.KUZU_API_KEY = None mock_env.ENVIRONMENT = "prod" - with patch("robosystems.kuzu_api.client.base.logger") as mock_logger: + with patch("robosystems.graph_api.client.base.logger") as mock_logger: BaseKuzuClient(base_url="http://localhost:8001") mock_logger.warning.assert_called_with("KuzuClient initialized without API key") @@ -322,7 +322,7 @@ def test_initialization_debug_missing_api_key_in_dev(self): mock_env.KUZU_API_KEY = None mock_env.ENVIRONMENT = "dev" - with patch("robosystems.kuzu_api.client.base.logger") as mock_logger: + with patch("robosystems.graph_api.client.base.logger") as mock_logger: BaseKuzuClient(base_url="http://localhost:8001") mock_logger.debug.assert_called_with( "KuzuClient initialized without API key (development mode)" diff --git a/tests/kuzu_api/client/test_client_extended.py b/tests/graph_api/client/test_client_extended.py similarity index 98% rename from tests/kuzu_api/client/test_client_extended.py rename to tests/graph_api/client/test_client_extended.py index 93185f60..a0c18b95 100644 --- a/tests/kuzu_api/client/test_client_extended.py +++ b/tests/graph_api/client/test_client_extended.py @@ -6,8 +6,8 @@ import pytest import httpx -from robosystems.kuzu_api.client.client import KuzuClient -from robosystems.kuzu_api.client.exceptions import ( +from robosystems.graph_api.client.client import KuzuClient +from robosystems.graph_api.client.exceptions import ( KuzuAPIError, KuzuTimeoutError, KuzuTransientError, @@ -306,7 +306,7 @@ async def __anext__(self): mock_sse_context = AsyncMock() mock_sse_context.aiter_sse.return_value = MockSSEIterator() - with patch("robosystems.kuzu_api.client.client.aconnect_sse") as mock_connect: + with patch("robosystems.graph_api.client.client.aconnect_sse") as mock_connect: mock_connect.return_value = mock_sse_context result = await client.ingest_with_sse( @@ -575,7 +575,7 @@ async def mock_iter(): mock_sse_context.__aenter__.return_value = mock_sse_context mock_sse_context.__aexit__.return_value = None - with patch("robosystems.kuzu_api.client.client.aconnect_sse") as mock_connect: + with patch("robosystems.graph_api.client.client.aconnect_sse") as mock_connect: mock_connect.return_value = mock_sse_context result = await client._monitor_ingestion_sse( @@ -608,7 +608,7 @@ async def mock_iter(): mock_sse_context.__aenter__.return_value = mock_sse_context mock_sse_context.__aexit__.return_value = None - with patch("robosystems.kuzu_api.client.client.aconnect_sse") as mock_connect: + with patch("robosystems.graph_api.client.client.aconnect_sse") as mock_connect: mock_connect.return_value = mock_sse_context result = await client._monitor_ingestion_sse(task_id="error-task") diff --git a/tests/kuzu_api/client/test_config.py b/tests/graph_api/client/test_config.py similarity index 99% rename from tests/kuzu_api/client/test_config.py rename to tests/graph_api/client/test_config.py index 088fe2f0..e4bb927f 100644 --- a/tests/kuzu_api/client/test_config.py +++ b/tests/graph_api/client/test_config.py @@ -4,7 +4,7 @@ from unittest.mock import patch import pytest -from robosystems.kuzu_api.client.config import KuzuClientConfig +from robosystems.graph_api.client.config import KuzuClientConfig class TestKuzuClientConfig: diff --git a/tests/kuzu_api/client/test_factory.py b/tests/graph_api/client/test_factory.py similarity index 91% rename from tests/kuzu_api/client/test_factory.py rename to tests/graph_api/client/test_factory.py index 9ac4e831..d9bb7315 100644 --- a/tests/kuzu_api/client/test_factory.py +++ b/tests/graph_api/client/test_factory.py @@ -4,7 +4,7 @@ from unittest.mock import MagicMock, patch, AsyncMock import pytest -from robosystems.kuzu_api.client.factory import ( +from robosystems.graph_api.client.factory import ( KuzuClientFactory, CircuitBreaker, get_kuzu_client, @@ -108,7 +108,7 @@ class TestKuzuClientFactory: @pytest.fixture def mock_env(self): """Mock environment variables.""" - with patch("robosystems.kuzu_api.client.factory.env") as mock: + with patch("robosystems.graph_api.client.factory.env") as mock: mock.KUZU_CONNECT_TIMEOUT = 5.0 mock.KUZU_READ_TIMEOUT = 30.0 mock.KUZU_CIRCUIT_BREAKER_THRESHOLD = 5 @@ -126,7 +126,7 @@ async def test_create_client_for_user_graph(self, mock_env): """Test creating client for user graph.""" # Mock allocation manager with patch( - "robosystems.kuzu_api.client.factory.KuzuAllocationManager" + "robosystems.graph_api.client.factory.KuzuAllocationManager" ) as MockAllocationManager: mock_manager = AsyncMock() MockAllocationManager.return_value = mock_manager @@ -140,7 +140,7 @@ async def test_create_client_for_user_graph(self, mock_env): mock_manager.find_database_location.return_value = mock_location # Mock KuzuClient creation - with patch("robosystems.kuzu_api.client.factory.KuzuClient") as MockClient: + with patch("robosystems.graph_api.client.factory.KuzuClient") as MockClient: mock_client = MagicMock() mock_client._graph_id = "kg123456" mock_client._database_name = "kg123456" @@ -162,7 +162,7 @@ async def test_create_client_for_shared_repository(self, mock_env): # Mock ALB health check with patch.object(KuzuClientFactory, "_check_alb_health", return_value=True): # Mock KuzuClient creation - with patch("robosystems.kuzu_api.client.factory.KuzuClient") as MockClient: + with patch("robosystems.graph_api.client.factory.KuzuClient") as MockClient: mock_client = MagicMock() mock_client._graph_id = "sec" mock_client._database_name = "sec" @@ -184,7 +184,7 @@ async def test_create_client_for_subgraph(self, mock_env): """Test creating client for subgraph.""" # Mock allocation for parent graph with patch( - "robosystems.kuzu_api.client.factory.KuzuAllocationManager" + "robosystems.graph_api.client.factory.KuzuAllocationManager" ) as MockAllocationManager: mock_manager = AsyncMock() MockAllocationManager.return_value = mock_manager @@ -198,7 +198,7 @@ async def test_create_client_for_subgraph(self, mock_env): mock_manager.find_database_location.return_value = mock_location # Mock httpx client creation - with patch("robosystems.kuzu_api.client.factory.httpx.AsyncClient"): + with patch("robosystems.graph_api.client.factory.httpx.AsyncClient"): client = await KuzuClientFactory.create_client("kg123456:subgraph1") assert client is not None @@ -209,7 +209,7 @@ async def test_create_client_for_subgraph(self, mock_env): async def test_create_client_shared_write_operation(self, mock_env): """Test shared repository write goes to master.""" # Mock KuzuClient creation to control the response - with patch("robosystems.kuzu_api.client.factory.KuzuClient") as MockClient: + with patch("robosystems.graph_api.client.factory.KuzuClient") as MockClient: mock_client = MagicMock() mock_client._graph_id = "sec" mock_client._database_name = "sec" @@ -229,7 +229,7 @@ async def test_create_client_alb_unhealthy_fallback(self, mock_env): # Mock ALB as unhealthy with patch.object(KuzuClientFactory, "_check_alb_health", return_value=False): # Mock KuzuClient creation - with patch("robosystems.kuzu_api.client.factory.KuzuClient") as MockClient: + with patch("robosystems.graph_api.client.factory.KuzuClient") as MockClient: mock_client = MagicMock() mock_client._graph_id = "sec" mock_client._database_name = "sec" @@ -252,7 +252,7 @@ async def test_create_client_no_allocation(self, mock_env): mock_env.ENVIRONMENT = "dev" with patch( - "robosystems.kuzu_api.client.factory.KuzuAllocationManager" + "robosystems.graph_api.client.factory.KuzuAllocationManager" ) as MockAllocationManager: mock_manager = AsyncMock() MockAllocationManager.return_value = mock_manager @@ -267,7 +267,7 @@ async def test_create_client_no_allocation(self, mock_env): async def test_create_client_with_tier_override(self, mock_env): """Test creating client with tier override.""" with patch( - "robosystems.kuzu_api.client.factory.KuzuAllocationManager" + "robosystems.graph_api.client.factory.KuzuAllocationManager" ) as MockAllocationManager: mock_manager = AsyncMock() MockAllocationManager.return_value = mock_manager @@ -281,7 +281,7 @@ async def test_create_client_with_tier_override(self, mock_env): mock_manager.find_database_location.return_value = mock_location # Mock KuzuClient creation - with patch("robosystems.kuzu_api.client.factory.KuzuClient") as MockClient: + with patch("robosystems.graph_api.client.factory.KuzuClient") as MockClient: mock_client = MagicMock() mock_client._graph_id = "kg123456" mock_client._database_name = "kg123456" @@ -333,7 +333,7 @@ async def test_get_kuzu_client(self): @pytest.mark.asyncio async def test_get_kuzu_client_for_instance(self): """Test get_kuzu_client_for_instance function.""" - with patch("robosystems.kuzu_api.client.factory.KuzuClient") as MockClient: + with patch("robosystems.graph_api.client.factory.KuzuClient") as MockClient: mock_client = MagicMock() MockClient.return_value = mock_client @@ -349,7 +349,7 @@ async def test_get_kuzu_client_for_instance(self): async def test_get_kuzu_client_for_sec_ingestion(self): """Test get_kuzu_client_for_sec_ingestion function.""" # Mock the environment to be development to avoid DynamoDB calls - with patch("robosystems.kuzu_api.client.factory.env") as mock_env: + with patch("robosystems.graph_api.client.factory.env") as mock_env: mock_env.is_development.return_value = True mock_env.KUZU_API_URL = "http://localhost:8001" mock_env.KUZU_API_KEY = "test-key" @@ -357,7 +357,7 @@ async def test_get_kuzu_client_for_sec_ingestion(self): mock_env.KUZU_READ_TIMEOUT = 30.0 # Mock KuzuClient creation - with patch("robosystems.kuzu_api.client.factory.KuzuClient") as MockClient: + with patch("robosystems.graph_api.client.factory.KuzuClient") as MockClient: mock_client = MagicMock() mock_client._graph_id = "sec" mock_client._database_name = "sec" diff --git a/tests/kuzu_api/client/test_kuzu_client.py b/tests/graph_api/client/test_kuzu_client.py similarity index 97% rename from tests/kuzu_api/client/test_kuzu_client.py rename to tests/graph_api/client/test_kuzu_client.py index 44888421..622405fd 100644 --- a/tests/kuzu_api/client/test_kuzu_client.py +++ b/tests/graph_api/client/test_kuzu_client.py @@ -5,9 +5,9 @@ import pytest import httpx -from robosystems.kuzu_api.client.client import KuzuClient -from robosystems.kuzu_api.client.config import KuzuClientConfig -from robosystems.kuzu_api.client.exceptions import ( +from robosystems.graph_api.client.client import KuzuClient +from robosystems.graph_api.client.config import KuzuClientConfig +from robosystems.graph_api.client.exceptions import ( KuzuTimeoutError, KuzuTransientError, KuzuSyntaxError, diff --git a/tests/kuzu_api/client/test_kuzu_exceptions.py b/tests/graph_api/client/test_kuzu_exceptions.py similarity index 99% rename from tests/kuzu_api/client/test_kuzu_exceptions.py rename to tests/graph_api/client/test_kuzu_exceptions.py index aac7b211..9fbc1716 100644 --- a/tests/kuzu_api/client/test_kuzu_exceptions.py +++ b/tests/graph_api/client/test_kuzu_exceptions.py @@ -2,7 +2,7 @@ import pytest -from robosystems.kuzu_api.client.exceptions import ( +from robosystems.graph_api.client.exceptions import ( KuzuAPIError, KuzuTransientError, KuzuTimeoutError, diff --git a/tests/kuzu_api/client/test_sse_client.py b/tests/graph_api/client/test_sse_client.py similarity index 87% rename from tests/kuzu_api/client/test_sse_client.py rename to tests/graph_api/client/test_sse_client.py index a5f69be3..d318fd61 100644 --- a/tests/kuzu_api/client/test_sse_client.py +++ b/tests/graph_api/client/test_sse_client.py @@ -7,7 +7,7 @@ import httpx -from robosystems.kuzu_api.client.sse_client import ( +from robosystems.graph_api.client.sse_client import ( KuzuIngestionSSEClient, monitor_ingestion_sync, ) @@ -99,7 +99,7 @@ async def test_start_and_monitor_ingestion_success(self, sse_client): ] with patch( - "robosystems.kuzu_api.client.sse_client.httpx.AsyncClient" + "robosystems.graph_api.client.sse_client.httpx.AsyncClient" ) as mock_client_class: mock_client = AsyncMock() mock_client.__aenter__.return_value = mock_client @@ -107,7 +107,9 @@ async def test_start_and_monitor_ingestion_success(self, sse_client): mock_client.post.return_value = mock_post_response mock_client_class.return_value = mock_client - with patch("robosystems.kuzu_api.client.sse_client.aconnect_sse") as mock_connect: + with patch( + "robosystems.graph_api.client.sse_client.aconnect_sse" + ) as mock_connect: mock_connect.return_value = MockEventSource(events) result = await sse_client.start_and_monitor_ingestion( @@ -146,7 +148,7 @@ async def test_start_and_monitor_ingestion_http_error(self, sse_client): mock_response.text = "Bad request" with patch( - "robosystems.kuzu_api.client.sse_client.httpx.AsyncClient" + "robosystems.graph_api.client.sse_client.httpx.AsyncClient" ) as mock_client_class: mock_client = AsyncMock() mock_client.__aenter__.return_value = mock_client @@ -172,7 +174,7 @@ async def test_start_and_monitor_ingestion_http_error(self, sse_client): async def test_start_and_monitor_ingestion_exception(self, sse_client): """Test handling of general exceptions.""" with patch( - "robosystems.kuzu_api.client.sse_client.httpx.AsyncClient" + "robosystems.graph_api.client.sse_client.httpx.AsyncClient" ) as mock_client_class: mock_client = AsyncMock() mock_client.__aenter__.return_value = mock_client @@ -202,14 +204,16 @@ async def test_monitor_via_sse_heartbeat_handling(self, sse_client): ] with patch( - "robosystems.kuzu_api.client.sse_client.httpx.AsyncClient" + "robosystems.graph_api.client.sse_client.httpx.AsyncClient" ) as mock_client_class: mock_client = AsyncMock() mock_client.__aenter__.return_value = mock_client mock_client.__aexit__.return_value = None mock_client_class.return_value = mock_client - with patch("robosystems.kuzu_api.client.sse_client.aconnect_sse") as mock_connect: + with patch( + "robosystems.graph_api.client.sse_client.aconnect_sse" + ) as mock_connect: mock_connect.return_value = MockEventSource(events) result = await sse_client._monitor_via_sse( @@ -244,17 +248,19 @@ async def test_monitor_via_sse_progress_logging(self, sse_client): ] with patch( - "robosystems.kuzu_api.client.sse_client.httpx.AsyncClient" + "robosystems.graph_api.client.sse_client.httpx.AsyncClient" ) as mock_client_class: mock_client = AsyncMock() mock_client.__aenter__.return_value = mock_client mock_client.__aexit__.return_value = None mock_client_class.return_value = mock_client - with patch("robosystems.kuzu_api.client.sse_client.aconnect_sse") as mock_connect: + with patch( + "robosystems.graph_api.client.sse_client.aconnect_sse" + ) as mock_connect: mock_connect.return_value = MockEventSource(events) - with patch("robosystems.kuzu_api.client.sse_client.time.time") as mock_time: + with patch("robosystems.graph_api.client.sse_client.time.time") as mock_time: # Simulate time progression for progress logging mock_time.side_effect = [0, 0, 0, 0, 35, 35, 35, 70, 70, 70] @@ -281,14 +287,16 @@ async def test_monitor_via_sse_failed_event(self, sse_client): ] with patch( - "robosystems.kuzu_api.client.sse_client.httpx.AsyncClient" + "robosystems.graph_api.client.sse_client.httpx.AsyncClient" ) as mock_client_class: mock_client = AsyncMock() mock_client.__aenter__.return_value = mock_client mock_client.__aexit__.return_value = None mock_client_class.return_value = mock_client - with patch("robosystems.kuzu_api.client.sse_client.aconnect_sse") as mock_connect: + with patch( + "robosystems.graph_api.client.sse_client.aconnect_sse" + ) as mock_connect: mock_connect.return_value = MockEventSource(events) result = await sse_client._monitor_via_sse( @@ -307,14 +315,16 @@ async def test_monitor_via_sse_error_event(self, sse_client): events = [MockSSEEvent("error", json.dumps({"error": "Stream interrupted"}))] with patch( - "robosystems.kuzu_api.client.sse_client.httpx.AsyncClient" + "robosystems.graph_api.client.sse_client.httpx.AsyncClient" ) as mock_client_class: mock_client = AsyncMock() mock_client.__aenter__.return_value = mock_client mock_client.__aexit__.return_value = None mock_client_class.return_value = mock_client - with patch("robosystems.kuzu_api.client.sse_client.aconnect_sse") as mock_connect: + with patch( + "robosystems.graph_api.client.sse_client.aconnect_sse" + ) as mock_connect: mock_connect.return_value = MockEventSource(events) result = await sse_client._monitor_via_sse( @@ -338,14 +348,16 @@ async def test_monitor_via_sse_invalid_json(self, sse_client): ] with patch( - "robosystems.kuzu_api.client.sse_client.httpx.AsyncClient" + "robosystems.graph_api.client.sse_client.httpx.AsyncClient" ) as mock_client_class: mock_client = AsyncMock() mock_client.__aenter__.return_value = mock_client mock_client.__aexit__.return_value = None mock_client_class.return_value = mock_client - with patch("robosystems.kuzu_api.client.sse_client.aconnect_sse") as mock_connect: + with patch( + "robosystems.graph_api.client.sse_client.aconnect_sse" + ) as mock_connect: mock_connect.return_value = MockEventSource(events) result = await sse_client._monitor_via_sse( @@ -379,14 +391,16 @@ async def test_monitor_via_sse_timeout(self, sse_client): ] with patch( - "robosystems.kuzu_api.client.sse_client.httpx.AsyncClient" + "robosystems.graph_api.client.sse_client.httpx.AsyncClient" ) as mock_client_class: mock_client = AsyncMock() mock_client.__aenter__.return_value = mock_client mock_client.__aexit__.return_value = None mock_client_class.return_value = mock_client - with patch("robosystems.kuzu_api.client.sse_client.aconnect_sse") as mock_connect: + with patch( + "robosystems.graph_api.client.sse_client.aconnect_sse" + ) as mock_connect: # Create an event source that yields events slowly async def slow_events(): for event in events: @@ -399,7 +413,7 @@ async def slow_events(): mock_event_source.aiter_sse = slow_events mock_connect.return_value = mock_event_source - with patch("robosystems.kuzu_api.client.sse_client.time.time") as mock_time: + with patch("robosystems.graph_api.client.sse_client.time.time") as mock_time: # Simulate time progression beyond timeout mock_time.side_effect = [0, 0.5, 1.5, 2.0] # Exceeds 1 second timeout @@ -416,14 +430,16 @@ async def slow_events(): async def test_monitor_via_sse_asyncio_timeout(self, sse_client): """Test handling of asyncio.TimeoutError.""" with patch( - "robosystems.kuzu_api.client.sse_client.httpx.AsyncClient" + "robosystems.graph_api.client.sse_client.httpx.AsyncClient" ) as mock_client_class: mock_client = AsyncMock() mock_client.__aenter__.return_value = mock_client mock_client.__aexit__.return_value = None mock_client_class.return_value = mock_client - with patch("robosystems.kuzu_api.client.sse_client.aconnect_sse") as mock_connect: + with patch( + "robosystems.graph_api.client.sse_client.aconnect_sse" + ) as mock_connect: mock_connect.side_effect = asyncio.TimeoutError() result = await sse_client._monitor_via_sse( @@ -439,14 +455,16 @@ async def test_monitor_via_sse_asyncio_timeout(self, sse_client): async def test_monitor_via_sse_unexpected_exception(self, sse_client): """Test handling of unexpected exceptions during monitoring.""" with patch( - "robosystems.kuzu_api.client.sse_client.httpx.AsyncClient" + "robosystems.graph_api.client.sse_client.httpx.AsyncClient" ) as mock_client_class: mock_client = AsyncMock() mock_client.__aenter__.return_value = mock_client mock_client.__aexit__.return_value = None mock_client_class.return_value = mock_client - with patch("robosystems.kuzu_api.client.sse_client.aconnect_sse") as mock_connect: + with patch( + "robosystems.graph_api.client.sse_client.aconnect_sse" + ) as mock_connect: mock_connect.side_effect = RuntimeError("Unexpected error") result = await sse_client._monitor_via_sse( @@ -465,14 +483,16 @@ async def test_monitor_via_sse_stream_ends_unexpectedly(self, sse_client): events = [] with patch( - "robosystems.kuzu_api.client.sse_client.httpx.AsyncClient" + "robosystems.graph_api.client.sse_client.httpx.AsyncClient" ) as mock_client_class: mock_client = AsyncMock() mock_client.__aenter__.return_value = mock_client mock_client.__aexit__.return_value = None mock_client_class.return_value = mock_client - with patch("robosystems.kuzu_api.client.sse_client.aconnect_sse") as mock_connect: + with patch( + "robosystems.graph_api.client.sse_client.aconnect_sse" + ) as mock_connect: mock_connect.return_value = MockEventSource(events) result = await sse_client._monitor_via_sse( @@ -500,14 +520,16 @@ async def test_monitor_via_sse_completion_without_records(self, sse_client): ] with patch( - "robosystems.kuzu_api.client.sse_client.httpx.AsyncClient" + "robosystems.graph_api.client.sse_client.httpx.AsyncClient" ) as mock_client_class: mock_client = AsyncMock() mock_client.__aenter__.return_value = mock_client mock_client.__aexit__.return_value = None mock_client_class.return_value = mock_client - with patch("robosystems.kuzu_api.client.sse_client.aconnect_sse") as mock_connect: + with patch( + "robosystems.graph_api.client.sse_client.aconnect_sse" + ) as mock_connect: mock_connect.return_value = MockEventSource(events) result = await sse_client._monitor_via_sse( @@ -533,7 +555,7 @@ def test_monitor_ingestion_sync_success(self): "duration_seconds": 10.0, } - with patch("robosystems.kuzu_api.client.sse_client.asyncio.run") as mock_run: + with patch("robosystems.graph_api.client.sse_client.asyncio.run") as mock_run: mock_run.return_value = expected_result result = monitor_ingestion_sync( @@ -553,7 +575,7 @@ def test_monitor_ingestion_sync_with_defaults(self): """Test synchronous wrapper with default parameters.""" expected_result = {"status": "completed", "task_id": "task-456"} - with patch("robosystems.kuzu_api.client.sse_client.asyncio.run") as mock_run: + with patch("robosystems.graph_api.client.sse_client.asyncio.run") as mock_run: mock_run.return_value = expected_result result = monitor_ingestion_sync( @@ -569,7 +591,7 @@ def test_monitor_ingestion_sync_failure(self): """Test synchronous wrapper handling failures.""" expected_result = {"status": "failed", "error": "Connection error"} - with patch("robosystems.kuzu_api.client.sse_client.asyncio.run") as mock_run: + with patch("robosystems.graph_api.client.sse_client.asyncio.run") as mock_run: mock_run.return_value = expected_result result = monitor_ingestion_sync( @@ -584,7 +606,7 @@ def test_monitor_ingestion_sync_failure(self): def test_monitor_ingestion_sync_asyncio_exception(self): """Test handling when asyncio.run raises an exception.""" - with patch("robosystems.kuzu_api.client.sse_client.asyncio.run") as mock_run: + with patch("robosystems.graph_api.client.sse_client.asyncio.run") as mock_run: mock_run.side_effect = RuntimeError("Event loop error") with pytest.raises(RuntimeError, match="Event loop error"): diff --git a/tests/kuzu_api/routers/databases/test_db_query.py b/tests/graph_api/routers/databases/test_db_query.py similarity index 80% rename from tests/kuzu_api/routers/databases/test_db_query.py rename to tests/graph_api/routers/databases/test_db_query.py index 56a69308..d7823b3f 100644 --- a/tests/kuzu_api/routers/databases/test_db_query.py +++ b/tests/graph_api/routers/databases/test_db_query.py @@ -6,23 +6,27 @@ from fastapi.testclient import TestClient import json -from robosystems.kuzu_api.app import create_app -from robosystems.kuzu_api.core.admission_control import AdmissionDecision +from robosystems.graph_api.app import create_app +from robosystems.graph_api.core.admission_control import AdmissionDecision class TestDatabaseQueryRouter: """Test cases for database query endpoints.""" @pytest.fixture - def client(self): + def client(self, monkeypatch): """Create a test client.""" + monkeypatch.setenv("BACKEND_TYPE", "kuzu") + app = create_app() - # Override the cluster service dependency - from robosystems.kuzu_api.core.cluster_manager import get_cluster_service + # Override the cluster service dependency factory function + from robosystems.graph_api.routers.databases.query import ( + _get_cluster_service_for_request, + ) mock_service = MagicMock() - app.dependency_overrides[get_cluster_service] = lambda: mock_service + app.dependency_overrides[_get_cluster_service_for_request] = lambda: mock_service return TestClient(app) @@ -52,13 +56,15 @@ def mock_query_response(self): def test_execute_query_success(self, client, mock_query_request, mock_query_response): """Test successful query execution.""" # Configure the mock service that was already injected - from robosystems.kuzu_api.core.cluster_manager import get_cluster_service + from robosystems.graph_api.routers.databases.query import ( + _get_cluster_service_for_request, + ) - mock_service = client.app.dependency_overrides[get_cluster_service]() + mock_service = client.app.dependency_overrides[_get_cluster_service_for_request]() mock_service.execute_query.return_value = mock_query_response with patch( - "robosystems.kuzu_api.routers.databases.query.get_admission_controller" + "robosystems.graph_api.routers.databases.query.get_admission_controller" ) as mock_get_admission: mock_admission = MagicMock() mock_admission.check_admission.return_value = (AdmissionDecision.ACCEPT, "OK") @@ -96,13 +102,15 @@ def test_execute_query_with_parameters(self, client): } # Configure the mock service that was already injected - from robosystems.kuzu_api.core.cluster_manager import get_cluster_service + from robosystems.graph_api.routers.databases.query import ( + _get_cluster_service_for_request, + ) - mock_service = client.app.dependency_overrides[get_cluster_service]() + mock_service = client.app.dependency_overrides[_get_cluster_service_for_request]() mock_service.execute_query.return_value = expected_response with patch( - "robosystems.kuzu_api.routers.databases.query.get_admission_controller" + "robosystems.graph_api.routers.databases.query.get_admission_controller" ) as mock_get_admission: mock_admission = MagicMock() mock_admission.check_admission.return_value = (AdmissionDecision.ACCEPT, "OK") @@ -126,7 +134,7 @@ def test_execute_query_with_parameters(self, client): def test_execute_query_admission_rejected(self, client, mock_query_request): """Test query rejected by admission control.""" with patch( - "robosystems.kuzu_api.routers.databases.query.get_admission_controller" + "robosystems.graph_api.routers.databases.query.get_admission_controller" ) as mock_get_admission: mock_admission = MagicMock() mock_admission.check_admission.return_value = ( @@ -156,13 +164,15 @@ def test_execute_query_streaming(self, client, mock_query_request): ] # Configure the mock service that was already injected - from robosystems.kuzu_api.core.cluster_manager import get_cluster_service + from robosystems.graph_api.routers.databases.query import ( + _get_cluster_service_for_request, + ) - mock_service = client.app.dependency_overrides[get_cluster_service]() + mock_service = client.app.dependency_overrides[_get_cluster_service_for_request]() mock_service.execute_query_streaming.return_value = streaming_chunks with patch( - "robosystems.kuzu_api.routers.databases.query.get_admission_controller" + "robosystems.graph_api.routers.databases.query.get_admission_controller" ) as mock_get_admission: mock_admission = MagicMock() mock_admission.check_admission.return_value = (AdmissionDecision.ACCEPT, "OK") @@ -202,13 +212,15 @@ def test_execute_query_empty_result(self, client): } # Configure the mock service that was already injected - from robosystems.kuzu_api.core.cluster_manager import get_cluster_service + from robosystems.graph_api.routers.databases.query import ( + _get_cluster_service_for_request, + ) - mock_service = client.app.dependency_overrides[get_cluster_service]() + mock_service = client.app.dependency_overrides[_get_cluster_service_for_request]() mock_service.execute_query.return_value = empty_response with patch( - "robosystems.kuzu_api.routers.databases.query.get_admission_controller" + "robosystems.graph_api.routers.databases.query.get_admission_controller" ) as mock_get_admission: mock_admission = MagicMock() mock_admission.check_admission.return_value = (AdmissionDecision.ACCEPT, "OK") @@ -249,13 +261,15 @@ def test_execute_query_complex_cypher(self, client): } # Configure the mock service that was already injected - from robosystems.kuzu_api.core.cluster_manager import get_cluster_service + from robosystems.graph_api.routers.databases.query import ( + _get_cluster_service_for_request, + ) - mock_service = client.app.dependency_overrides[get_cluster_service]() + mock_service = client.app.dependency_overrides[_get_cluster_service_for_request]() mock_service.execute_query.return_value = complex_response with patch( - "robosystems.kuzu_api.routers.databases.query.get_admission_controller" + "robosystems.graph_api.routers.databases.query.get_admission_controller" ) as mock_get_admission: mock_admission = MagicMock() mock_admission.check_admission.return_value = (AdmissionDecision.ACCEPT, "OK") @@ -276,13 +290,15 @@ def test_execute_query_shared_database( ): """Test query execution on shared database (e.g., SEC).""" # Configure the mock service that was already injected - from robosystems.kuzu_api.core.cluster_manager import get_cluster_service + from robosystems.graph_api.routers.databases.query import ( + _get_cluster_service_for_request, + ) - mock_service = client.app.dependency_overrides[get_cluster_service]() + mock_service = client.app.dependency_overrides[_get_cluster_service_for_request]() mock_service.execute_query.return_value = mock_query_response with patch( - "robosystems.kuzu_api.routers.databases.query.get_admission_controller" + "robosystems.graph_api.routers.databases.query.get_admission_controller" ) as mock_get_admission: mock_admission = MagicMock() mock_admission.check_admission.return_value = (AdmissionDecision.ACCEPT, "OK") @@ -301,13 +317,15 @@ def test_execute_query_shared_database( def test_execute_query_connection_tracking_on_error(self, client, mock_query_request): """Test that connections are released even when query fails.""" # Configure the mock service that was already injected - from robosystems.kuzu_api.core.cluster_manager import get_cluster_service + from robosystems.graph_api.routers.databases.query import ( + _get_cluster_service_for_request, + ) - mock_service = client.app.dependency_overrides[get_cluster_service]() + mock_service = client.app.dependency_overrides[_get_cluster_service_for_request]() mock_service.execute_query.side_effect = Exception("Query execution failed") with patch( - "robosystems.kuzu_api.routers.databases.query.get_admission_controller" + "robosystems.graph_api.routers.databases.query.get_admission_controller" ) as mock_get_admission: mock_admission = MagicMock() mock_admission.check_admission.return_value = (AdmissionDecision.ACCEPT, "OK") @@ -341,13 +359,15 @@ def test_execute_query_large_result_non_streaming(self, client): } # Configure the mock service that was already injected - from robosystems.kuzu_api.core.cluster_manager import get_cluster_service + from robosystems.graph_api.routers.databases.query import ( + _get_cluster_service_for_request, + ) - mock_service = client.app.dependency_overrides[get_cluster_service]() + mock_service = client.app.dependency_overrides[_get_cluster_service_for_request]() mock_service.execute_query.return_value = large_response with patch( - "robosystems.kuzu_api.routers.databases.query.get_admission_controller" + "robosystems.graph_api.routers.databases.query.get_admission_controller" ) as mock_get_admission: mock_admission = MagicMock() mock_admission.check_admission.return_value = (AdmissionDecision.ACCEPT, "OK") diff --git a/tests/kuzu_api/routers/databases/test_management.py b/tests/graph_api/routers/databases/test_management.py similarity index 89% rename from tests/kuzu_api/routers/databases/test_management.py rename to tests/graph_api/routers/databases/test_management.py index 7b669e6e..82571633 100644 --- a/tests/kuzu_api/routers/databases/test_management.py +++ b/tests/graph_api/routers/databases/test_management.py @@ -5,8 +5,8 @@ from fastapi import status from fastapi.testclient import TestClient -from robosystems.kuzu_api.app import create_app -from robosystems.kuzu_api.models.database import ( +from robosystems.graph_api.app import create_app +from robosystems.graph_api.models.database import ( DatabaseCreateResponse, DatabaseListResponse, DatabaseInfo, @@ -23,7 +23,7 @@ def client(self): app = create_app() # Override the cluster service dependency - from robosystems.kuzu_api.core.cluster_manager import get_cluster_service + from robosystems.graph_api.core.cluster_manager import get_cluster_service mock_service = MagicMock() app.dependency_overrides[get_cluster_service] = lambda: mock_service @@ -60,7 +60,7 @@ def mock_database_list_response(self, mock_database_info): def test_list_databases_success(self, client, mock_database_list_response): """Test successful database listing.""" # Configure the mock service that was already injected - from robosystems.kuzu_api.core.cluster_manager import get_cluster_service + from robosystems.graph_api.core.cluster_manager import get_cluster_service mock_service = client.app.dependency_overrides[get_cluster_service]() mock_service.db_manager.get_all_databases_info.return_value = ( @@ -90,7 +90,7 @@ def test_list_databases_empty(self, client): ) # Configure the mock service that was already injected - from robosystems.kuzu_api.core.cluster_manager import get_cluster_service + from robosystems.graph_api.core.cluster_manager import get_cluster_service mock_service = client.app.dependency_overrides[get_cluster_service]() mock_service.db_manager.get_all_databases_info.return_value = empty_response @@ -117,7 +117,7 @@ def test_create_database_entity_schema(self, client): ) # Configure the mock service that was already injected - from robosystems.kuzu_api.core.cluster_manager import get_cluster_service + from robosystems.graph_api.core.cluster_manager import get_cluster_service mock_service = client.app.dependency_overrides[get_cluster_service]() mock_service.read_only = False @@ -149,7 +149,7 @@ def test_create_database_shared_schema(self, client): ) # Configure the mock service that was already injected - from robosystems.kuzu_api.core.cluster_manager import get_cluster_service + from robosystems.graph_api.core.cluster_manager import get_cluster_service mock_service = client.app.dependency_overrides[get_cluster_service]() mock_service.read_only = False @@ -166,7 +166,7 @@ def test_create_database_shared_schema(self, client): def test_create_database_shared_without_repository_name(self, client): """Test creating shared database without repository name fails.""" # Configure the mock service that was already injected - from robosystems.kuzu_api.core.cluster_manager import get_cluster_service + from robosystems.graph_api.core.cluster_manager import get_cluster_service mock_service = client.app.dependency_overrides[get_cluster_service]() mock_service.read_only = False @@ -182,7 +182,7 @@ def test_create_database_shared_without_repository_name(self, client): def test_create_database_on_readonly_node(self, client): """Test that database creation fails on read-only nodes.""" # Configure the mock service that was already injected - from robosystems.kuzu_api.core.cluster_manager import get_cluster_service + from robosystems.graph_api.core.cluster_manager import get_cluster_service mock_service = client.app.dependency_overrides[get_cluster_service]() mock_service.read_only = True @@ -198,7 +198,7 @@ def test_create_database_on_readonly_node(self, client): def test_get_database_info_success(self, client, mock_database_info): """Test retrieving database information.""" # Configure the mock service that was already injected - from robosystems.kuzu_api.core.cluster_manager import get_cluster_service + from robosystems.graph_api.core.cluster_manager import get_cluster_service mock_service = client.app.dependency_overrides[get_cluster_service]() mock_service.db_manager.get_database_info.return_value = mock_database_info @@ -217,7 +217,7 @@ def test_get_database_info_not_found(self, client): from fastapi import HTTPException # Configure the mock service that was already injected - from robosystems.kuzu_api.core.cluster_manager import get_cluster_service + from robosystems.graph_api.core.cluster_manager import get_cluster_service mock_service = client.app.dependency_overrides[get_cluster_service]() mock_service.db_manager.get_database_info.side_effect = HTTPException( @@ -232,7 +232,7 @@ def test_get_database_info_not_found(self, client): def test_delete_database_success(self, client): """Test successful database deletion.""" # Configure the mock service that was already injected - from robosystems.kuzu_api.core.cluster_manager import get_cluster_service + from robosystems.graph_api.core.cluster_manager import get_cluster_service mock_service = client.app.dependency_overrides[get_cluster_service]() mock_service.read_only = False @@ -249,7 +249,7 @@ def test_delete_database_success(self, client): def test_delete_database_on_readonly_node(self, client): """Test that database deletion fails on read-only nodes.""" # Configure the mock service that was already injected - from robosystems.kuzu_api.core.cluster_manager import get_cluster_service + from robosystems.graph_api.core.cluster_manager import get_cluster_service mock_service = client.app.dependency_overrides[get_cluster_service]() mock_service.read_only = True @@ -262,14 +262,14 @@ def test_delete_database_on_readonly_node(self, client): def test_delete_shared_database_warning(self, client): """Test that deleting shared database logs warning.""" # Configure the mock service that was already injected - from robosystems.kuzu_api.core.cluster_manager import get_cluster_service + from robosystems.graph_api.core.cluster_manager import get_cluster_service mock_service = client.app.dependency_overrides[get_cluster_service]() mock_service.read_only = False mock_service.node_type = NodeType.SHARED_MASTER mock_service.db_manager.delete_database.return_value = None with patch( - "robosystems.kuzu_api.routers.databases.management.logger" + "robosystems.graph_api.routers.databases.management.logger" ) as mock_logger: response = client.delete("/databases/sec") @@ -288,7 +288,7 @@ def test_create_database_custom_schema(self, client): ) # Configure the mock service that was already injected - from robosystems.kuzu_api.core.cluster_manager import get_cluster_service + from robosystems.graph_api.core.cluster_manager import get_cluster_service mock_service = client.app.dependency_overrides[get_cluster_service]() mock_service.read_only = False @@ -330,7 +330,7 @@ def test_list_databases_multiple(self, client): ) # Configure the mock service that was already injected - from robosystems.kuzu_api.core.cluster_manager import get_cluster_service + from robosystems.graph_api.core.cluster_manager import get_cluster_service mock_service = client.app.dependency_overrides[get_cluster_service]() mock_service.db_manager.get_all_databases_info.return_value = list_response @@ -355,7 +355,7 @@ def test_get_database_info_with_unhealthy_status(self, client): ) # Configure the mock service that was already injected - from robosystems.kuzu_api.core.cluster_manager import get_cluster_service + from robosystems.graph_api.core.cluster_manager import get_cluster_service mock_service = client.app.dependency_overrides[get_cluster_service]() mock_service.db_manager.get_database_info.return_value = unhealthy_db_info diff --git a/tests/kuzu_api/routers/test_health.py b/tests/graph_api/routers/test_health.py similarity index 89% rename from tests/kuzu_api/routers/test_health.py rename to tests/graph_api/routers/test_health.py index 6889a406..b42dcbeb 100644 --- a/tests/kuzu_api/routers/test_health.py +++ b/tests/graph_api/routers/test_health.py @@ -5,7 +5,7 @@ from fastapi import status from fastapi.testclient import TestClient -from robosystems.kuzu_api.app import create_app +from robosystems.graph_api.app import create_app class TestHealthRouter: @@ -17,7 +17,7 @@ def client(self): app = create_app() # Override the cluster service dependency - from robosystems.kuzu_api.core.cluster_manager import get_cluster_service + from robosystems.graph_api.core.cluster_manager import get_cluster_service mock_service = MagicMock() app.dependency_overrides[get_cluster_service] = lambda: mock_service @@ -35,7 +35,7 @@ def mock_cluster_service(self): def test_health_check_success(self, client, mock_cluster_service): """Test successful health check.""" # Configure the mock service that was already injected - from robosystems.kuzu_api.core.cluster_manager import get_cluster_service + from robosystems.graph_api.core.cluster_manager import get_cluster_service mock_service = client.app.dependency_overrides[get_cluster_service]() mock_service.get_uptime.return_value = 3600 @@ -52,7 +52,7 @@ def test_health_check_success(self, client, mock_cluster_service): def test_health_check_with_memory_info(self, client, mock_cluster_service): """Test health check with memory information.""" # Configure the mock service that was already injected - from robosystems.kuzu_api.core.cluster_manager import get_cluster_service + from robosystems.graph_api.core.cluster_manager import get_cluster_service mock_service = client.app.dependency_overrides[get_cluster_service]() mock_service.get_uptime.return_value = 3600 @@ -81,7 +81,7 @@ def test_health_check_with_memory_info(self, client, mock_cluster_service): def test_health_check_without_psutil(self, client, mock_cluster_service): """Test health check when psutil is not available.""" # Configure the mock service that was already injected - from robosystems.kuzu_api.core.cluster_manager import get_cluster_service + from robosystems.graph_api.core.cluster_manager import get_cluster_service mock_service = client.app.dependency_overrides[get_cluster_service]() mock_service.get_uptime.return_value = 3600 @@ -102,7 +102,7 @@ def test_health_check_without_psutil(self, client, mock_cluster_service): def test_health_check_service_error(self, client): """Test health check when cluster service has an error.""" # Configure the mock service that was already injected - from robosystems.kuzu_api.core.cluster_manager import get_cluster_service + from robosystems.graph_api.core.cluster_manager import get_cluster_service mock_service = client.app.dependency_overrides[get_cluster_service]() mock_service.get_uptime.side_effect = Exception("Service unavailable") @@ -117,7 +117,7 @@ def test_health_check_service_error(self, client): def test_health_check_database_error(self, client): """Test health check when database manager has an error.""" # Configure the mock service that was already injected - from robosystems.kuzu_api.core.cluster_manager import get_cluster_service + from robosystems.graph_api.core.cluster_manager import get_cluster_service mock_service = client.app.dependency_overrides[get_cluster_service]() mock_service.get_uptime.return_value = 1000 @@ -133,7 +133,7 @@ def test_health_check_database_error(self, client): def test_health_check_zero_databases(self, client, mock_cluster_service): """Test health check with zero databases.""" # Configure the mock service that was already injected - from robosystems.kuzu_api.core.cluster_manager import get_cluster_service + from robosystems.graph_api.core.cluster_manager import get_cluster_service mock_service = client.app.dependency_overrides[get_cluster_service]() mock_service.get_uptime.return_value = 3600 @@ -149,7 +149,7 @@ def test_health_check_zero_databases(self, client, mock_cluster_service): def test_health_check_response_format(self, client, mock_cluster_service): """Test that health check response has expected format.""" # Configure the mock service that was already injected - from robosystems.kuzu_api.core.cluster_manager import get_cluster_service + from robosystems.graph_api.core.cluster_manager import get_cluster_service mock_service = client.app.dependency_overrides[get_cluster_service]() mock_service.get_uptime.return_value = 3600 diff --git a/tests/kuzu_api/routers/test_info.py b/tests/graph_api/routers/test_info.py similarity index 89% rename from tests/kuzu_api/routers/test_info.py rename to tests/graph_api/routers/test_info.py index adcd196e..e9933e56 100644 --- a/tests/kuzu_api/routers/test_info.py +++ b/tests/graph_api/routers/test_info.py @@ -5,8 +5,8 @@ from fastapi import status from fastapi.testclient import TestClient -from robosystems.kuzu_api.app import create_app -from robosystems.kuzu_api.models.cluster import ClusterInfoResponse +from robosystems.graph_api.app import create_app +from robosystems.graph_api.models.cluster import ClusterInfoResponse class TestInfoRouter: @@ -18,7 +18,7 @@ def client(self): app = create_app() # Override the cluster service dependency - from robosystems.kuzu_api.core.cluster_manager import get_cluster_service + from robosystems.graph_api.core.cluster_manager import get_cluster_service mock_service = MagicMock() app.dependency_overrides[get_cluster_service] = lambda: mock_service @@ -43,7 +43,7 @@ def mock_cluster_info(self): def test_get_cluster_info_success(self, client, mock_cluster_info): """Test successful cluster info retrieval.""" # Configure the mock service that was already injected - from robosystems.kuzu_api.core.cluster_manager import get_cluster_service + from robosystems.graph_api.core.cluster_manager import get_cluster_service mock_service = client.app.dependency_overrides[get_cluster_service]() mock_service.get_cluster_info.return_value = mock_cluster_info @@ -75,7 +75,7 @@ def test_get_cluster_info_shared_repository(self, client): ) # Configure the mock service that was already injected - from robosystems.kuzu_api.core.cluster_manager import get_cluster_service + from robosystems.graph_api.core.cluster_manager import get_cluster_service mock_service = client.app.dependency_overrides[get_cluster_service]() mock_service.get_cluster_info.return_value = mock_info @@ -102,7 +102,7 @@ def test_get_cluster_info_read_only_node(self, client): ) # Configure the mock service that was already injected - from robosystems.kuzu_api.core.cluster_manager import get_cluster_service + from robosystems.graph_api.core.cluster_manager import get_cluster_service mock_service = client.app.dependency_overrides[get_cluster_service]() mock_service.get_cluster_info.return_value = mock_info @@ -129,7 +129,7 @@ def test_get_cluster_info_no_databases(self, client): ) # Configure the mock service that was already injected - from robosystems.kuzu_api.core.cluster_manager import get_cluster_service + from robosystems.graph_api.core.cluster_manager import get_cluster_service mock_service = client.app.dependency_overrides[get_cluster_service]() mock_service.get_cluster_info.return_value = mock_info @@ -157,7 +157,7 @@ def test_get_cluster_info_at_capacity(self, client): ) # Configure the mock service that was already injected - from robosystems.kuzu_api.core.cluster_manager import get_cluster_service + from robosystems.graph_api.core.cluster_manager import get_cluster_service mock_service = client.app.dependency_overrides[get_cluster_service]() mock_service.get_cluster_info.return_value = mock_info @@ -184,7 +184,7 @@ def test_get_cluster_info_response_model_validation(self, client): ) # Configure the mock service that was already injected - from robosystems.kuzu_api.core.cluster_manager import get_cluster_service + from robosystems.graph_api.core.cluster_manager import get_cluster_service mock_service = client.app.dependency_overrides[get_cluster_service]() mock_service.get_cluster_info.return_value = mock_info diff --git a/tests/kuzu_api/routers/test_metrics.py b/tests/graph_api/routers/test_metrics.py similarity index 93% rename from tests/kuzu_api/routers/test_metrics.py rename to tests/graph_api/routers/test_metrics.py index 5acc93ed..7b0fd7d6 100644 --- a/tests/kuzu_api/routers/test_metrics.py +++ b/tests/graph_api/routers/test_metrics.py @@ -6,7 +6,7 @@ from fastapi.testclient import TestClient from datetime import datetime -from robosystems.kuzu_api.app import create_app +from robosystems.graph_api.app import create_app from robosystems.middleware.graph.clusters import NodeType @@ -19,7 +19,7 @@ def client(self): app = create_app() # Override the cluster service dependency - from robosystems.kuzu_api.core.cluster_manager import get_cluster_service + from robosystems.graph_api.core.cluster_manager import get_cluster_service mock_service = MagicMock() app.dependency_overrides[get_cluster_service] = lambda: mock_service @@ -120,7 +120,7 @@ async def test_get_metrics_complete( ): """Test getting complete metrics snapshot.""" # Configure the mock service that was already injected - from robosystems.kuzu_api.core.cluster_manager import get_cluster_service + from robosystems.graph_api.core.cluster_manager import get_cluster_service mock_service = client.app.dependency_overrides[get_cluster_service]() mock_service.node_id = "test-node-01" @@ -139,7 +139,7 @@ async def test_get_metrics_complete( # Setup admission controller with patch( - "robosystems.kuzu_api.routers.metrics.get_admission_controller" + "robosystems.graph_api.routers.metrics.get_admission_controller" ) as mock_get_admission: mock_admission = MagicMock() mock_admission.get_metrics.return_value = mock_admission_metrics @@ -215,7 +215,7 @@ async def test_get_metrics_minimal(self, client): } # Configure the mock service that was already injected - from robosystems.kuzu_api.core.cluster_manager import get_cluster_service + from robosystems.graph_api.core.cluster_manager import get_cluster_service mock_service = client.app.dependency_overrides[get_cluster_service]() mock_service.node_id = "empty-node" @@ -230,7 +230,7 @@ async def test_get_metrics_minimal(self, client): mock_service.metrics_collector = mock_collector with patch( - "robosystems.kuzu_api.routers.metrics.get_admission_controller" + "robosystems.graph_api.routers.metrics.get_admission_controller" ) as mock_get_admission: mock_admission = MagicMock() mock_admission.get_metrics.return_value = minimal_admission @@ -248,7 +248,7 @@ async def test_get_metrics_minimal(self, client): async def test_get_metrics_shared_node(self, client): """Test metrics for shared repository node.""" # Configure the mock service that was already injected - from robosystems.kuzu_api.core.cluster_manager import get_cluster_service + from robosystems.graph_api.core.cluster_manager import get_cluster_service mock_service = client.app.dependency_overrides[get_cluster_service]() mock_service.node_id = "shared-master-01" @@ -293,7 +293,7 @@ async def test_get_metrics_shared_node(self, client): mock_service.metrics_collector = mock_collector with patch( - "robosystems.kuzu_api.routers.metrics.get_admission_controller" + "robosystems.graph_api.routers.metrics.get_admission_controller" ) as mock_get_admission: mock_admission = MagicMock() mock_admission.get_metrics.return_value = {"requests_accepted": 10000} @@ -342,7 +342,7 @@ async def test_get_metrics_high_load(self, client): } # Configure the mock service that was already injected - from robosystems.kuzu_api.core.cluster_manager import get_cluster_service + from robosystems.graph_api.core.cluster_manager import get_cluster_service mock_service = client.app.dependency_overrides[get_cluster_service]() mock_service.node_id = "overloaded-node" @@ -359,7 +359,7 @@ async def test_get_metrics_high_load(self, client): mock_service.metrics_collector = mock_collector with patch( - "robosystems.kuzu_api.routers.metrics.get_admission_controller" + "robosystems.graph_api.routers.metrics.get_admission_controller" ) as mock_get_admission: mock_admission = MagicMock() mock_admission.get_metrics.return_value = high_load_admission @@ -379,7 +379,7 @@ async def test_get_metrics_high_load(self, client): async def test_get_metrics_error_handling(self, client): """Test metrics when some collectors fail.""" # Configure the mock service that was already injected - from robosystems.kuzu_api.core.cluster_manager import get_cluster_service + from robosystems.graph_api.core.cluster_manager import get_cluster_service mock_service = client.app.dependency_overrides[get_cluster_service]() mock_service.node_id = "error-node" @@ -402,7 +402,7 @@ async def test_get_metrics_error_handling(self, client): mock_service.metrics_collector = mock_collector with patch( - "robosystems.kuzu_api.routers.metrics.get_admission_controller" + "robosystems.graph_api.routers.metrics.get_admission_controller" ) as mock_get_admission: mock_admission = MagicMock() mock_admission.get_metrics.return_value = {"requests_accepted": 0} diff --git a/tests/kuzu_api/routers/test_tasks.py b/tests/graph_api/routers/test_tasks.py similarity index 90% rename from tests/kuzu_api/routers/test_tasks.py rename to tests/graph_api/routers/test_tasks.py index caed829c..c5e70a8f 100644 --- a/tests/kuzu_api/routers/test_tasks.py +++ b/tests/graph_api/routers/test_tasks.py @@ -6,8 +6,8 @@ from fastapi.testclient import TestClient import json -from robosystems.kuzu_api.app import create_app -from robosystems.kuzu_api.core.task_sse import TaskType +from robosystems.graph_api.app import create_app +from robosystems.graph_api.core.task_sse import TaskType class TestTasksRouter: @@ -19,7 +19,7 @@ def client(self): app = create_app() # Override the cluster service dependency - from robosystems.kuzu_api.core.cluster_manager import get_cluster_service + from robosystems.graph_api.core.cluster_manager import get_cluster_service mock_service = MagicMock() app.dependency_overrides[get_cluster_service] = lambda: mock_service @@ -99,7 +99,7 @@ async def test_list_tasks_all( ] with patch( - "robosystems.kuzu_api.routers.tasks.unified_task_manager" + "robosystems.graph_api.routers.tasks.unified_task_manager" ) as mock_manager: mock_manager.list_all_tasks = AsyncMock(return_value=all_tasks) @@ -122,7 +122,7 @@ async def test_list_tasks_with_status_filter( running_tasks = [mock_ingestion_task] with patch( - "robosystems.kuzu_api.routers.tasks.unified_task_manager" + "robosystems.graph_api.routers.tasks.unified_task_manager" ) as mock_manager: mock_manager.list_all_tasks = AsyncMock(return_value=running_tasks) @@ -141,7 +141,7 @@ async def test_list_tasks_with_type_filter( all_tasks = [mock_ingestion_task, mock_backup_task, mock_copy_task] with patch( - "robosystems.kuzu_api.routers.tasks.unified_task_manager" + "robosystems.graph_api.routers.tasks.unified_task_manager" ) as mock_manager: mock_manager.list_all_tasks = AsyncMock(return_value=all_tasks) @@ -175,7 +175,7 @@ async def test_list_tasks_with_limit(self, client): ] with patch( - "robosystems.kuzu_api.routers.tasks.unified_task_manager" + "robosystems.graph_api.routers.tasks.unified_task_manager" ) as mock_manager: mock_manager.list_all_tasks = AsyncMock(return_value=many_tasks) @@ -189,7 +189,7 @@ async def test_list_tasks_with_limit(self, client): async def test_get_task_status_success(self, client, mock_ingestion_task): """Test getting task status successfully.""" with patch( - "robosystems.kuzu_api.routers.tasks.unified_task_manager" + "robosystems.graph_api.routers.tasks.unified_task_manager" ) as mock_manager: mock_manager.get_task = AsyncMock(return_value=mock_ingestion_task) @@ -205,7 +205,7 @@ async def test_get_task_status_success(self, client, mock_ingestion_task): async def test_get_task_status_not_found(self, client): """Test getting status for non-existent task.""" with patch( - "robosystems.kuzu_api.routers.tasks.unified_task_manager" + "robosystems.graph_api.routers.tasks.unified_task_manager" ) as mock_manager: mock_manager.get_task = AsyncMock(return_value=None) @@ -233,7 +233,7 @@ async def test_get_task_statistics( ] with patch( - "robosystems.kuzu_api.routers.tasks.unified_task_manager" + "robosystems.graph_api.routers.tasks.unified_task_manager" ) as mock_manager: mock_manager.list_all_tasks = AsyncMock(return_value=all_tasks) @@ -264,7 +264,7 @@ async def test_get_task_statistics_by_type(self, client): ] with patch( - "robosystems.kuzu_api.routers.tasks.unified_task_manager" + "robosystems.graph_api.routers.tasks.unified_task_manager" ) as mock_manager: mock_manager.list_all_tasks = AsyncMock(return_value=tasks) @@ -280,7 +280,7 @@ async def test_get_task_statistics_by_type(self, client): async def test_get_task_statistics_empty(self, client): """Test statistics when no tasks exist.""" with patch( - "robosystems.kuzu_api.routers.tasks.unified_task_manager" + "robosystems.graph_api.routers.tasks.unified_task_manager" ) as mock_manager: mock_manager.list_all_tasks = AsyncMock(return_value=[]) @@ -299,14 +299,14 @@ def test_monitor_task_endpoint_exists(self, client): # SSE endpoints return EventSourceResponse which test client doesn't handle well # Just verify the endpoint exists with patch( - "robosystems.kuzu_api.routers.tasks.unified_task_manager" + "robosystems.graph_api.routers.tasks.unified_task_manager" ) as mock_manager: mock_manager.get_task_type.return_value = TaskType.INGESTION - with patch("robosystems.kuzu_api.routers.tasks.EventSourceResponse") as mock_sse: + with patch("robosystems.graph_api.routers.tasks.EventSourceResponse") as mock_sse: mock_sse.return_value = MagicMock() - with patch("robosystems.kuzu_api.routers.tasks.generate_task_sse_events"): + with patch("robosystems.graph_api.routers.tasks.generate_task_sse_events"): # This will fail with SSE but confirms endpoint is registered response = client.get("/tasks/test_123/monitor") # EventSourceResponse doesn't work with TestClient @@ -316,7 +316,7 @@ def test_monitor_task_endpoint_exists(self, client): @pytest.mark.asyncio async def test_unified_task_manager_get_task_from_redis(self): """Test UnifiedTaskManager getting task from Redis.""" - from robosystems.kuzu_api.routers.tasks import UnifiedTaskManager + from robosystems.graph_api.routers.tasks import UnifiedTaskManager manager = UnifiedTaskManager() mock_task = {"task_id": "test_123", "status": "running"} @@ -334,7 +334,7 @@ async def test_unified_task_manager_get_task_from_redis(self): @pytest.mark.asyncio async def test_unified_task_manager_get_task_fallback(self): """Test UnifiedTaskManager fallback to specific managers.""" - from robosystems.kuzu_api.routers.tasks import UnifiedTaskManager + from robosystems.graph_api.routers.tasks import UnifiedTaskManager manager = UnifiedTaskManager() mock_task = {"task_id": "ingest_123", "status": "completed"} @@ -357,7 +357,7 @@ async def test_unified_task_manager_get_task_fallback(self): @pytest.mark.asyncio async def test_unified_task_manager_list_all_tasks(self): """Test UnifiedTaskManager listing all tasks.""" - from robosystems.kuzu_api.routers.tasks import UnifiedTaskManager + from robosystems.graph_api.routers.tasks import UnifiedTaskManager manager = UnifiedTaskManager() tasks = [ @@ -392,7 +392,7 @@ async def test_unified_task_manager_list_all_tasks(self): def test_unified_task_manager_get_task_type(self): """Test UnifiedTaskManager determining task type from ID.""" - from robosystems.kuzu_api.routers.tasks import UnifiedTaskManager + from robosystems.graph_api.routers.tasks import UnifiedTaskManager manager = UnifiedTaskManager() diff --git a/tests/kuzu_api/test_auth_middleware.py b/tests/graph_api/test_auth_middleware.py similarity index 89% rename from tests/kuzu_api/test_auth_middleware.py rename to tests/graph_api/test_auth_middleware.py index e0a8605e..c268afab 100644 --- a/tests/kuzu_api/test_auth_middleware.py +++ b/tests/graph_api/test_auth_middleware.py @@ -1,4 +1,4 @@ -"""Tests for kuzu_api auth middleware.""" +"""Tests for graph_api auth middleware.""" import json import time @@ -8,7 +8,7 @@ from fastapi.responses import JSONResponse from starlette.datastructures import Headers -from robosystems.kuzu_api.middleware.auth import ( +from robosystems.graph_api.middleware.auth import ( KuzuAuthMiddleware, get_api_key_from_secrets_manager, clear_api_key_cache, @@ -49,7 +49,7 @@ async def test_middleware_exempt_paths(self, mock_app, mock_request): @pytest.mark.asyncio async def test_middleware_development_bypass(self, mock_app, mock_request): """Test that authentication is bypassed in development.""" - with patch("robosystems.kuzu_api.middleware.auth.env") as mock_env: + with patch("robosystems.graph_api.middleware.auth.env") as mock_env: mock_env.ENVIRONMENT = "dev" mock_env.KUZU_API_KEY = None @@ -63,7 +63,7 @@ async def test_middleware_development_bypass(self, mock_app, mock_request): @pytest.mark.asyncio async def test_middleware_valid_api_key_header(self, mock_app, mock_request): """Test successful authentication with valid API key in header.""" - with patch("robosystems.kuzu_api.middleware.auth.env") as mock_env: + with patch("robosystems.graph_api.middleware.auth.env") as mock_env: mock_env.ENVIRONMENT = "prod" mock_env.KUZU_API_KEY = None @@ -78,7 +78,7 @@ async def test_middleware_valid_api_key_header(self, mock_app, mock_request): @pytest.mark.asyncio async def test_middleware_valid_bearer_token(self, mock_app, mock_request): """Test successful authentication with Bearer token.""" - with patch("robosystems.kuzu_api.middleware.auth.env") as mock_env: + with patch("robosystems.graph_api.middleware.auth.env") as mock_env: mock_env.ENVIRONMENT = "staging" mock_env.KUZU_API_KEY = None @@ -93,7 +93,7 @@ async def test_middleware_valid_bearer_token(self, mock_app, mock_request): @pytest.mark.asyncio async def test_middleware_missing_api_key(self, mock_app, mock_request): """Test authentication failure with missing API key.""" - with patch("robosystems.kuzu_api.middleware.auth.env") as mock_env: + with patch("robosystems.graph_api.middleware.auth.env") as mock_env: mock_env.ENVIRONMENT = "prod" mock_env.KUZU_API_KEY = None @@ -109,7 +109,7 @@ async def test_middleware_missing_api_key(self, mock_app, mock_request): @pytest.mark.asyncio async def test_middleware_invalid_api_key(self, mock_app, mock_request): """Test authentication failure with invalid API key.""" - with patch("robosystems.kuzu_api.middleware.auth.env") as mock_env: + with patch("robosystems.graph_api.middleware.auth.env") as mock_env: mock_env.ENVIRONMENT = "staging" mock_env.KUZU_API_KEY = None @@ -126,7 +126,7 @@ async def test_middleware_invalid_api_key(self, mock_app, mock_request): @pytest.mark.asyncio async def test_middleware_rate_limiting(self, mock_app, mock_request): """Test rate limiting after multiple failed attempts.""" - with patch("robosystems.kuzu_api.middleware.auth.env") as mock_env: + with patch("robosystems.graph_api.middleware.auth.env") as mock_env: mock_env.ENVIRONMENT = "prod" mock_env.KUZU_API_KEY = None @@ -149,7 +149,7 @@ async def test_middleware_rate_limiting(self, mock_app, mock_request): @pytest.mark.asyncio async def test_middleware_rate_limit_expiry(self, mock_app, mock_request): """Test that rate limiting expires after lockout duration.""" - with patch("robosystems.kuzu_api.middleware.auth.env") as mock_env: + with patch("robosystems.graph_api.middleware.auth.env") as mock_env: mock_env.ENVIRONMENT = "prod" mock_env.KUZU_API_KEY = None @@ -179,7 +179,7 @@ async def test_middleware_resets_failed_attempts_on_success( self, mock_app, mock_request ): """Test that successful auth resets failed attempt counter.""" - with patch("robosystems.kuzu_api.middleware.auth.env") as mock_env: + with patch("robosystems.graph_api.middleware.auth.env") as mock_env: mock_env.ENVIRONMENT = "prod" mock_env.KUZU_API_KEY = None @@ -199,7 +199,7 @@ async def test_middleware_resets_failed_attempts_on_success( def test_middleware_initialization_with_env_key(self, mock_app): """Test middleware initialization with key from environment.""" - with patch("robosystems.kuzu_api.middleware.auth.env") as mock_env: + with patch("robosystems.graph_api.middleware.auth.env") as mock_env: mock_env.ENVIRONMENT = "prod" mock_env.KUZU_API_KEY = "env-api-key" @@ -209,12 +209,12 @@ def test_middleware_initialization_with_env_key(self, mock_app): def test_middleware_initialization_requires_key_in_prod(self, mock_app): """Test that middleware requires API key in production.""" - with patch("robosystems.kuzu_api.middleware.auth.env") as mock_env: + with patch("robosystems.graph_api.middleware.auth.env") as mock_env: mock_env.ENVIRONMENT = "prod" mock_env.KUZU_API_KEY = None with patch( - "robosystems.kuzu_api.middleware.auth.get_api_key_from_secrets_manager" + "robosystems.graph_api.middleware.auth.get_api_key_from_secrets_manager" ) as mock_get_key: mock_get_key.return_value = None @@ -257,7 +257,7 @@ def test_cleanup_failed_attempts(self, mock_app): class TestSecretsManagerIntegration: """Test cases for Secrets Manager integration.""" - @patch("robosystems.kuzu_api.middleware.auth.boto3.client") + @patch("robosystems.graph_api.middleware.auth.boto3.client") def test_get_api_key_from_secrets_manager_success(self, mock_boto_client): """Test successful API key retrieval from Secrets Manager.""" # Clear cache first @@ -269,7 +269,7 @@ def test_get_api_key_from_secrets_manager_success(self, mock_boto_client): "SecretString": json.dumps({"KUZU_API_KEY": "secret-key-123"}) } - with patch("robosystems.kuzu_api.middleware.auth.env") as mock_env: + with patch("robosystems.graph_api.middleware.auth.env") as mock_env: mock_env.ENVIRONMENT = "prod" api_key = get_api_key_from_secrets_manager(key_type="writer") @@ -278,7 +278,7 @@ def test_get_api_key_from_secrets_manager_success(self, mock_boto_client): SecretId="robosystems/prod/kuzu" ) - @patch("robosystems.kuzu_api.middleware.auth.boto3.client") + @patch("robosystems.graph_api.middleware.auth.boto3.client") def test_get_api_key_from_secrets_manager_not_found(self, mock_boto_client): """Test handling of missing secret in Secrets Manager.""" # Clear cache first @@ -292,13 +292,13 @@ def test_get_api_key_from_secrets_manager_not_found(self, mock_boto_client): {"Error": {"Code": "ResourceNotFoundException"}}, "GetSecretValue" ) - with patch("robosystems.kuzu_api.middleware.auth.env") as mock_env: + with patch("robosystems.graph_api.middleware.auth.env") as mock_env: mock_env.ENVIRONMENT = "staging" api_key = get_api_key_from_secrets_manager() assert api_key is None - @patch("robosystems.kuzu_api.middleware.auth.boto3.client") + @patch("robosystems.graph_api.middleware.auth.boto3.client") def test_get_api_key_from_secrets_manager_no_key_in_secret(self, mock_boto_client): """Test handling when secret exists but doesn't contain API key.""" # Clear cache first @@ -310,13 +310,13 @@ def test_get_api_key_from_secrets_manager_no_key_in_secret(self, mock_boto_clien "SecretString": json.dumps({"OTHER_KEY": "value"}) } - with patch("robosystems.kuzu_api.middleware.auth.env") as mock_env: + with patch("robosystems.graph_api.middleware.auth.env") as mock_env: mock_env.ENVIRONMENT = "prod" api_key = get_api_key_from_secrets_manager() assert api_key is None - @patch("robosystems.kuzu_api.middleware.auth.boto3.client") + @patch("robosystems.graph_api.middleware.auth.boto3.client") def test_get_api_key_caching(self, mock_boto_client): """Test that API key is cached after first retrieval.""" # Clear cache first @@ -328,7 +328,7 @@ def test_get_api_key_caching(self, mock_boto_client): "SecretString": json.dumps({"KUZU_API_KEY": "cached-key"}) } - with patch("robosystems.kuzu_api.middleware.auth.env") as mock_env: + with patch("robosystems.graph_api.middleware.auth.env") as mock_env: mock_env.ENVIRONMENT = "prod" # First call @@ -349,7 +349,7 @@ class TestAPIKeyGeneration: def test_create_api_key(self): """Test secure API key generation.""" with patch( - "robosystems.kuzu_api.middleware.auth.SecurityAuditLogger" + "robosystems.graph_api.middleware.auth.SecurityAuditLogger" ) as mock_logger: api_key, key_hash = create_api_key(prefix="test") @@ -374,7 +374,7 @@ def test_create_api_key_unique(self): # All keys should be unique assert len(keys) == 10 - @patch("robosystems.kuzu_api.middleware.auth.bcrypt") + @patch("robosystems.graph_api.middleware.auth.bcrypt") def test_create_api_key_bcrypt_hashing(self, mock_bcrypt): """Test that bcrypt is used for hashing.""" mock_salt = b"$2b$12$test_salt" diff --git a/tests/kuzu_api/test_cluster_server.py b/tests/graph_api/test_cluster_server.py similarity index 94% rename from tests/kuzu_api/test_cluster_server.py rename to tests/graph_api/test_cluster_server.py index 121c0ef0..ea285fb5 100644 --- a/tests/kuzu_api/test_cluster_server.py +++ b/tests/graph_api/test_cluster_server.py @@ -6,16 +6,16 @@ from unittest.mock import MagicMock, patch from fastapi.testclient import TestClient -from robosystems.kuzu_api.core.cluster_manager import ( +from robosystems.graph_api.core.cluster_manager import ( KuzuClusterService, validate_cypher_query, ) -from robosystems.kuzu_api.core.utils import ( +from robosystems.graph_api.core.utils import ( validate_database_name, validate_query_parameters, ) -from robosystems.kuzu_api.app import create_app -from robosystems.kuzu_api.models.database import QueryRequest, DatabaseCreateRequest +from robosystems.graph_api.app import create_app +from robosystems.graph_api.models.database import QueryRequest, DatabaseCreateRequest from robosystems.middleware.graph.clusters import NodeType, RepositoryType from robosystems.exceptions import ConfigurationError @@ -165,7 +165,7 @@ def teardown_method(self): """Clean up test fixtures.""" shutil.rmtree(self.temp_dir, ignore_errors=True) - @patch("robosystems.kuzu_api.core.cluster_manager.KuzuDatabaseManager") + @patch("robosystems.graph_api.core.cluster_manager.KuzuDatabaseManager") def test_initialization_entity_writer(self, mock_db_manager): """Test initialization of entity writer node.""" service = KuzuClusterService( @@ -184,7 +184,7 @@ def test_initialization_entity_writer(self, mock_db_manager): assert isinstance(service.start_time, float) mock_db_manager.assert_called_once_with(self.base_path, 100, read_only=False) - @patch("robosystems.kuzu_api.core.cluster_manager.KuzuDatabaseManager") + @patch("robosystems.graph_api.core.cluster_manager.KuzuDatabaseManager") def test_initialization_shared_writer(self, mock_db_manager): """Test initialization of shared repository writer.""" service = KuzuClusterService( @@ -199,7 +199,7 @@ def test_initialization_shared_writer(self, mock_db_manager): assert service.node_type == NodeType.WRITER assert service.repository_type == RepositoryType.SHARED - @patch("robosystems.kuzu_api.core.cluster_manager.KuzuDatabaseManager") + @patch("robosystems.graph_api.core.cluster_manager.KuzuDatabaseManager") def test_initialization_validation_errors(self, mock_db_manager): """Test initialization validation for invalid configurations.""" @@ -233,7 +233,7 @@ def test_initialization_validation_errors(self, mock_db_manager): ) assert service.repository_type == RepositoryType.SHARED - @patch("robosystems.kuzu_api.core.cluster_manager.KuzuDatabaseManager") + @patch("robosystems.graph_api.core.cluster_manager.KuzuDatabaseManager") def test_execute_query_success(self, mock_db_manager): """Test successful query execution.""" # Mock database manager instance @@ -287,7 +287,7 @@ def test_execute_query_success(self, mock_db_manager): ] assert response.execution_time_ms > 0 - @patch("robosystems.kuzu_api.core.cluster_manager.KuzuDatabaseManager") + @patch("robosystems.graph_api.core.cluster_manager.KuzuDatabaseManager") def test_execute_query_database_not_found(self, mock_db_manager): """Test query execution with non-existent database.""" from fastapi import HTTPException @@ -312,8 +312,8 @@ def test_execute_query_database_not_found(self, mock_db_manager): assert exc_info.value.status_code == 404 assert "not found" in str(exc_info.value.detail) - @patch("robosystems.kuzu_api.core.cluster_manager.KuzuDatabaseManager") - @patch("robosystems.kuzu_api.core.cluster_manager.ThreadPoolExecutor") + @patch("robosystems.graph_api.core.cluster_manager.KuzuDatabaseManager") + @patch("robosystems.graph_api.core.cluster_manager.ThreadPoolExecutor") def test_execute_query_timeout(self, mock_executor_class, mock_db_manager): """Test query execution timeout using ThreadPoolExecutor.""" from concurrent.futures import TimeoutError as FuturesTimeoutError @@ -367,7 +367,7 @@ def test_execute_query_timeout(self, mock_executor_class, mock_db_manager): # Verify timeout was used correctly mock_future.result.assert_called_once_with(timeout=1.0) - @patch("robosystems.kuzu_api.core.cluster_manager.KuzuDatabaseManager") + @patch("robosystems.graph_api.core.cluster_manager.KuzuDatabaseManager") def test_execute_query_large_result_set(self, mock_db_manager): """Test query execution with large result set (DoS protection).""" # Mock database manager instance before creating service @@ -410,8 +410,8 @@ def test_execute_query_large_result_set(self, mock_db_manager): # Should close result to free resources mock_result.close.assert_called_once() - @patch("robosystems.kuzu_api.core.cluster_manager.psutil") - @patch("robosystems.kuzu_api.core.cluster_manager.KuzuDatabaseManager") + @patch("robosystems.graph_api.core.cluster_manager.psutil") + @patch("robosystems.graph_api.core.cluster_manager.KuzuDatabaseManager") def test_get_cluster_health(self, mock_db_manager, mock_psutil): """Test cluster health check.""" service = KuzuClusterService( @@ -439,7 +439,7 @@ def test_get_cluster_health(self, mock_db_manager, mock_psutil): assert health.read_only is False assert health.uptime_seconds > 0 - @patch("robosystems.kuzu_api.core.cluster_manager.KuzuDatabaseManager") + @patch("robosystems.graph_api.core.cluster_manager.KuzuDatabaseManager") def test_get_cluster_info(self, mock_db_manager): """Test cluster information retrieval.""" service = KuzuClusterService( @@ -476,7 +476,7 @@ def setup_method(self): # Initialize a mock cluster service for all tests from robosystems.middleware.graph.clusters import NodeType - from robosystems.kuzu_api.core import cluster_manager + from robosystems.graph_api.core import cluster_manager mock_service = MagicMock() mock_service.node_type = NodeType.WRITER @@ -493,7 +493,7 @@ def teardown_method(self): shutil.rmtree(self.temp_dir, ignore_errors=True) # Reset the global cluster service - from robosystems.kuzu_api.core import cluster_manager + from robosystems.graph_api.core import cluster_manager cluster_manager._cluster_service = None @@ -523,7 +523,7 @@ def test_health_endpoint(self): def test_cluster_info_endpoint(self): """Test cluster info endpoint.""" - from robosystems.kuzu_api.models.cluster import ClusterInfoResponse + from robosystems.graph_api.models.cluster import ClusterInfoResponse # Mock cluster info response mock_info = ClusterInfoResponse( @@ -551,7 +551,7 @@ def test_cluster_info_endpoint(self): def test_execute_query_endpoint(self): """Test query execution endpoint.""" - from robosystems.kuzu_api.models.database import QueryResponse + from robosystems.graph_api.models.database import QueryResponse # Mock cluster service mock_response = QueryResponse( @@ -582,7 +582,7 @@ def test_execute_query_endpoint(self): def test_list_databases_endpoint(self): """Test database listing endpoint.""" - from robosystems.kuzu_api.models.database import DatabaseListResponse, DatabaseInfo + from robosystems.graph_api.models.database import DatabaseListResponse, DatabaseInfo # Mock cluster service mock_db_info = DatabaseInfo( @@ -620,7 +620,7 @@ def test_list_databases_endpoint(self): def test_create_database_endpoint_entity_writer(self): """Test database creation endpoint for entity writer.""" - from robosystems.kuzu_api.models.database import DatabaseCreateResponse + from robosystems.graph_api.models.database import DatabaseCreateResponse # Mock cluster service self.mock_service.read_only = False @@ -710,7 +710,7 @@ def test_delete_database_endpoint(self): def test_database_health_endpoint(self): """Test database health check endpoint.""" - from robosystems.kuzu_api.models.database import DatabaseInfo + from robosystems.graph_api.models.database import DatabaseInfo # Mock cluster service mock_db_info = DatabaseInfo( @@ -774,7 +774,7 @@ def test_ingest_data_endpoint_read_only(self): # Patch connection pool at module level before creating app with patch( - "robosystems.kuzu_api.core.connection_pool._connection_pool", MagicMock() + "robosystems.graph_api.core.connection_pool._connection_pool", MagicMock() ): app = create_app() client = TestClient(app) diff --git a/tests/kuzu_api/test_connection_pool.py b/tests/graph_api/test_connection_pool.py similarity index 99% rename from tests/kuzu_api/test_connection_pool.py rename to tests/graph_api/test_connection_pool.py index 70a6fdc0..2ff2e267 100644 --- a/tests/kuzu_api/test_connection_pool.py +++ b/tests/graph_api/test_connection_pool.py @@ -9,7 +9,7 @@ from unittest.mock import MagicMock, patch from datetime import datetime, timedelta, timezone -from robosystems.kuzu_api.core.connection_pool import ( +from robosystems.graph_api.core.connection_pool import ( KuzuConnectionPool, ConnectionInfo, initialize_connection_pool, @@ -643,7 +643,7 @@ def teardown_method(self): """Clean up test fixtures.""" shutil.rmtree(self.temp_dir, ignore_errors=True) # Reset global state - import robosystems.kuzu_api.core.connection_pool as pool_module + import robosystems.graph_api.core.connection_pool as pool_module pool_module._connection_pool = None diff --git a/tests/kuzu_api/test_database_manager.py b/tests/graph_api/test_database_manager.py similarity index 93% rename from tests/kuzu_api/test_database_manager.py rename to tests/graph_api/test_database_manager.py index 92c5994b..4450e64a 100644 --- a/tests/kuzu_api/test_database_manager.py +++ b/tests/graph_api/test_database_manager.py @@ -7,7 +7,7 @@ from unittest.mock import MagicMock, patch from datetime import datetime -from robosystems.kuzu_api.core.database_manager import ( +from robosystems.graph_api.core.database_manager import ( KuzuDatabaseManager, DatabaseCreateRequest, DatabaseInfo, @@ -95,7 +95,7 @@ def teardown_method(self): """Clean up test fixtures.""" shutil.rmtree(self.temp_dir, ignore_errors=True) - @patch("robosystems.kuzu_api.core.database_manager.initialize_connection_pool") + @patch("robosystems.graph_api.core.database_manager.initialize_connection_pool") def test_initialization(self, mock_init_pool): """Test manager initialization.""" mock_pool = MagicMock() @@ -119,7 +119,7 @@ def test_initialization(self, mock_init_pool): connection_ttl_minutes=30, ) - @patch("robosystems.kuzu_api.core.database_manager.initialize_connection_pool") + @patch("robosystems.graph_api.core.database_manager.initialize_connection_pool") @patch("kuzu.Database") @patch("kuzu.Connection") def test_create_database_success_entity_schema( @@ -173,7 +173,7 @@ def test_create_database_success_entity_schema( # Verify database was created successfully (connection pool manages connections) # New implementation uses connection pool, no direct database/connection storage - @patch("robosystems.kuzu_api.core.database_manager.initialize_connection_pool") + @patch("robosystems.graph_api.core.database_manager.initialize_connection_pool") @patch("kuzu.Database") @patch("kuzu.Connection") def test_create_database_read_only( @@ -209,7 +209,7 @@ def test_create_database_read_only( # Verify read-only database creation completed successfully # New implementation uses connection pool, no direct connection storage - @patch("robosystems.kuzu_api.core.database_manager.initialize_connection_pool") + @patch("robosystems.graph_api.core.database_manager.initialize_connection_pool") def test_create_database_capacity_exceeded(self, mock_init_pool): """Test database creation when capacity is exceeded.""" from fastapi import HTTPException @@ -233,7 +233,7 @@ def test_create_database_capacity_exceeded(self, mock_init_pool): assert exc_info.value.status_code == 507 # Insufficient Storage assert "Maximum database capacity reached" in str(exc_info.value.detail) - @patch("robosystems.kuzu_api.core.database_manager.initialize_connection_pool") + @patch("robosystems.graph_api.core.database_manager.initialize_connection_pool") def test_create_database_already_exists(self, mock_init_pool): """Test database creation when database already exists.""" from fastapi import HTTPException @@ -259,7 +259,7 @@ def test_create_database_already_exists(self, mock_init_pool): assert exc_info.value.status_code == 409 # Conflict assert "already exists" in str(exc_info.value.detail) - @patch("robosystems.kuzu_api.core.database_manager.initialize_connection_pool") + @patch("robosystems.graph_api.core.database_manager.initialize_connection_pool") @patch("kuzu.Database") def test_create_database_kuzu_error_cleanup(self, mock_db_class, mock_init_pool): """Test database creation error handling and cleanup.""" @@ -289,7 +289,7 @@ def test_create_database_kuzu_error_cleanup(self, mock_db_class, mock_init_pool) db_path = self.base_path / "test_fail.kuzu" assert not db_path.exists() - @patch("robosystems.kuzu_api.core.database_manager.initialize_connection_pool") + @patch("robosystems.graph_api.core.database_manager.initialize_connection_pool") def test_delete_database_success(self, mock_init_pool): """Test successful database deletion.""" mock_init_pool.return_value = MagicMock() @@ -316,7 +316,7 @@ def test_delete_database_success(self, mock_init_pool): # Verify directory was deleted assert not db_path.exists() - @patch("robosystems.kuzu_api.core.database_manager.initialize_connection_pool") + @patch("robosystems.graph_api.core.database_manager.initialize_connection_pool") def test_delete_database_not_found(self, mock_init_pool): """Test database deletion when database doesn't exist.""" from fastapi import HTTPException @@ -330,7 +330,7 @@ def test_delete_database_not_found(self, mock_init_pool): assert exc_info.value.status_code == 404 assert "not found" in str(exc_info.value.detail) - @patch("robosystems.kuzu_api.core.database_manager.initialize_connection_pool") + @patch("robosystems.graph_api.core.database_manager.initialize_connection_pool") def test_list_databases(self, mock_init_pool): """Test database listing.""" mock_init_pool.return_value = MagicMock() @@ -347,7 +347,7 @@ def test_list_databases(self, mock_init_pool): # Should return only .kuzu files, sorted assert databases == ["db1", "db2", "db3"] - @patch("robosystems.kuzu_api.core.database_manager.initialize_connection_pool") + @patch("robosystems.graph_api.core.database_manager.initialize_connection_pool") def test_list_databases_empty(self, mock_init_pool): """Test database listing when no databases exist.""" mock_init_pool.return_value = MagicMock() @@ -357,7 +357,7 @@ def test_list_databases_empty(self, mock_init_pool): assert databases == [] - @patch("robosystems.kuzu_api.core.database_manager.initialize_connection_pool") + @patch("robosystems.graph_api.core.database_manager.initialize_connection_pool") def test_get_database_info_success(self, mock_init_pool): """Test successful database info retrieval.""" mock_init_pool.return_value = MagicMock() @@ -383,7 +383,7 @@ def test_get_database_info_success(self, mock_init_pool): assert info.last_accessed is not None assert isinstance(datetime.fromisoformat(info.created_at), datetime) - @patch("robosystems.kuzu_api.core.database_manager.initialize_connection_pool") + @patch("robosystems.graph_api.core.database_manager.initialize_connection_pool") def test_get_database_info_not_found(self, mock_init_pool): """Test database info retrieval for non-existent database.""" from fastapi import HTTPException @@ -397,7 +397,7 @@ def test_get_database_info_not_found(self, mock_init_pool): assert exc_info.value.status_code == 404 assert "not found" in str(exc_info.value.detail) - @patch("robosystems.kuzu_api.core.database_manager.initialize_connection_pool") + @patch("robosystems.graph_api.core.database_manager.initialize_connection_pool") def test_get_all_databases_info(self, mock_init_pool): """Test retrieval of all databases info.""" mock_init_pool.return_value = MagicMock() @@ -437,7 +437,7 @@ def test_get_all_databases_info(self, mock_init_pool): assert response.node_capacity["capacity_remaining"] == 98 assert response.node_capacity["utilization_percent"] == 2.0 - @patch("robosystems.kuzu_api.core.database_manager.initialize_connection_pool") + @patch("robosystems.graph_api.core.database_manager.initialize_connection_pool") def test_health_check_all(self, mock_init_pool): """Test health check for all databases.""" mock_init_pool.return_value = MagicMock() @@ -473,7 +473,7 @@ def test_health_check_all(self, mock_init_pool): assert health_response.unhealthy_databases == 1 assert len(health_response.databases) == 2 - @patch("robosystems.kuzu_api.core.database_manager.initialize_connection_pool") + @patch("robosystems.graph_api.core.database_manager.initialize_connection_pool") def test_check_database_health_success(self, mock_init_pool): """Test successful database health check (file existence only).""" mock_init_pool.return_value = MagicMock() @@ -487,7 +487,7 @@ def test_check_database_health_success(self, mock_init_pool): assert is_healthy is True - @patch("robosystems.kuzu_api.core.database_manager.initialize_connection_pool") + @patch("robosystems.graph_api.core.database_manager.initialize_connection_pool") def test_check_database_health_failure(self, mock_init_pool): """Test database health check failure (file does not exist).""" mock_init_pool.return_value = MagicMock() @@ -498,7 +498,7 @@ def test_check_database_health_failure(self, mock_init_pool): assert is_healthy is False - @patch("robosystems.kuzu_api.core.database_manager.initialize_connection_pool") + @patch("robosystems.graph_api.core.database_manager.initialize_connection_pool") def test_check_database_health_not_found(self, mock_init_pool): """Test database health check for non-existent database.""" mock_init_pool.return_value = MagicMock() @@ -508,7 +508,7 @@ def test_check_database_health_not_found(self, mock_init_pool): assert is_healthy is False - @patch("robosystems.kuzu_api.core.database_manager.initialize_connection_pool") + @patch("robosystems.graph_api.core.database_manager.initialize_connection_pool") def test_apply_entity_schema(self, mock_init_pool): """Test entity schema application using dynamic schema loader.""" mock_init_pool.return_value = MagicMock() @@ -539,7 +539,7 @@ def test_apply_entity_schema(self, mock_init_pool): assert "CREATE NODE TABLE IF NOT EXISTS User" in all_calls_str assert "CREATE NODE TABLE IF NOT EXISTS Entity" in all_calls_str - @patch("robosystems.kuzu_api.core.database_manager.initialize_connection_pool") + @patch("robosystems.graph_api.core.database_manager.initialize_connection_pool") def test_apply_shared_schema(self, mock_init_pool): """Test shared schema application using dynamic schema loader.""" mock_init_pool.return_value = MagicMock() @@ -560,7 +560,7 @@ def test_apply_shared_schema(self, mock_init_pool): all_calls_str = " ".join(call_args) assert "CREATE NODE TABLE IF NOT EXISTS" in all_calls_str - @patch("robosystems.kuzu_api.core.database_manager.initialize_connection_pool") + @patch("robosystems.graph_api.core.database_manager.initialize_connection_pool") def test_apply_schema_unknown_type(self, mock_init_pool): """Test schema application with unknown schema type.""" mock_init_pool.return_value = MagicMock() @@ -572,7 +572,7 @@ def test_apply_schema_unknown_type(self, mock_init_pool): assert result is False - @patch("robosystems.kuzu_api.core.database_manager.initialize_connection_pool") + @patch("robosystems.graph_api.core.database_manager.initialize_connection_pool") def test_apply_fallback_entity_schema(self, mock_init_pool): """Test fallback entity schema application.""" mock_init_pool.return_value = MagicMock() @@ -595,7 +595,7 @@ def test_apply_fallback_entity_schema(self, mock_init_pool): assert "CREATE NODE TABLE IF NOT EXISTS User" in all_calls_str assert "CREATE REL TABLE IF NOT EXISTS HAS_USER" in all_calls_str - @patch("robosystems.kuzu_api.core.database_manager.initialize_connection_pool") + @patch("robosystems.graph_api.core.database_manager.initialize_connection_pool") def test_apply_fallback_entity_schema_with_error(self, mock_init_pool): """Test fallback entity schema application with database error.""" mock_init_pool.return_value = MagicMock() @@ -608,7 +608,7 @@ def test_apply_fallback_entity_schema_with_error(self, mock_init_pool): assert result is False - @patch("robosystems.kuzu_api.core.database_manager.initialize_connection_pool") + @patch("robosystems.graph_api.core.database_manager.initialize_connection_pool") def test_map_schema_type_to_kuzu(self, mock_init_pool): """Test schema type mapping to Kuzu types.""" mock_init_pool.return_value = MagicMock() @@ -630,7 +630,7 @@ def test_map_schema_type_to_kuzu(self, mock_init_pool): assert manager._map_schema_type_to_kuzu("UNKNOWN_TYPE") == "STRING" assert manager._map_schema_type_to_kuzu("") == "STRING" - @patch("robosystems.kuzu_api.core.database_manager.initialize_connection_pool") + @patch("robosystems.graph_api.core.database_manager.initialize_connection_pool") @patch("robosystems.schemas.loader.get_schema_loader") def test_apply_entity_schema_with_schema_loader_failure( self, mock_get_schema_loader, mock_init_pool @@ -650,7 +650,7 @@ def test_apply_entity_schema_with_schema_loader_failure( # Verify fallback was used (should have exactly 3 calls for minimal schema) assert mock_conn.execute.call_count == 3 - @patch("robosystems.kuzu_api.core.database_manager.initialize_connection_pool") + @patch("robosystems.graph_api.core.database_manager.initialize_connection_pool") def test_apply_shared_schema_with_different_repositories(self, mock_init_pool): """Test shared schema application with different repository types.""" mock_init_pool.return_value = MagicMock() @@ -669,7 +669,7 @@ def test_apply_shared_schema_with_different_repositories(self, mock_init_pool): result_other = manager._apply_shared_schema(mock_conn, "industry") assert result_other is True - @patch("robosystems.kuzu_api.core.database_manager.initialize_connection_pool") + @patch("robosystems.graph_api.core.database_manager.initialize_connection_pool") def test_close_all_connections(self, mock_init_pool): """Test closing all connections.""" mock_pool = MagicMock() diff --git a/tests/kuzu_api/test_databases_metrics.py b/tests/graph_api/test_databases_metrics.py similarity index 98% rename from tests/kuzu_api/test_databases_metrics.py rename to tests/graph_api/test_databases_metrics.py index a54dc8d4..ac8b8df7 100644 --- a/tests/kuzu_api/test_databases_metrics.py +++ b/tests/graph_api/test_databases_metrics.py @@ -1,4 +1,4 @@ -"""Tests for kuzu_api databases/metrics router.""" +"""Tests for graph_api databases/metrics router.""" import tempfile from pathlib import Path @@ -6,7 +6,7 @@ import pytest from fastapi import HTTPException, status -from robosystems.kuzu_api.routers.databases.metrics import get_database_metrics +from robosystems.graph_api.routers.databases.metrics import get_database_metrics from robosystems.middleware.graph.clusters import NodeType diff --git a/tests/kuzu_api/test_main.py b/tests/graph_api/test_main.py similarity index 82% rename from tests/kuzu_api/test_main.py rename to tests/graph_api/test_main.py index e518823f..cb818e41 100644 --- a/tests/kuzu_api/test_main.py +++ b/tests/graph_api/test_main.py @@ -1,4 +1,4 @@ -"""Tests for kuzu_api main module.""" +"""Tests for graph_api main module.""" import sys import tempfile @@ -12,9 +12,9 @@ class TestMain: """Test cases for main entry point.""" - @patch("robosystems.kuzu_api.main.uvicorn.run") - @patch("robosystems.kuzu_api.main.create_app") - @patch("robosystems.kuzu_api.main.init_cluster_service") + @patch("robosystems.graph_api.main.uvicorn.run") + @patch("robosystems.graph_api.main.create_app") + @patch("robosystems.graph_api.main.init_cluster_service") def test_main_basic_configuration( self, mock_init_cluster, mock_create_app, mock_uvicorn ): @@ -40,7 +40,7 @@ def test_main_basic_configuration( mock_env.get_kuzu_tier_config.side_effect = Exception("No tier config") with patch.object(sys, "argv", test_args): - from robosystems.kuzu_api.main import main + from robosystems.graph_api.main import main main() @@ -66,9 +66,9 @@ def test_main_basic_configuration( access_log=True, ) - @patch("robosystems.kuzu_api.main.uvicorn.run") - @patch("robosystems.kuzu_api.main.create_app") - @patch("robosystems.kuzu_api.main.init_cluster_service") + @patch("robosystems.graph_api.main.uvicorn.run") + @patch("robosystems.graph_api.main.create_app") + @patch("robosystems.graph_api.main.init_cluster_service") def test_main_read_only_mode(self, mock_init_cluster, mock_create_app, mock_uvicorn): """Test main with read-only mode.""" with tempfile.TemporaryDirectory() as tmpdir: @@ -88,7 +88,7 @@ def test_main_read_only_mode(self, mock_init_cluster, mock_create_app, mock_uvic mock_env.CLUSTER_TIER = None # No tier configured with patch.object(sys, "argv", test_args): - from robosystems.kuzu_api.main import main + from robosystems.graph_api.main import main main() @@ -97,9 +97,9 @@ def test_main_read_only_mode(self, mock_init_cluster, mock_create_app, mock_uvic call_args = mock_init_cluster.call_args[1] assert call_args["read_only"] is True - @patch("robosystems.kuzu_api.main.uvicorn.run") - @patch("robosystems.kuzu_api.main.create_app") - @patch("robosystems.kuzu_api.main.init_cluster_service") + @patch("robosystems.graph_api.main.uvicorn.run") + @patch("robosystems.graph_api.main.create_app") + @patch("robosystems.graph_api.main.init_cluster_service") def test_main_shared_master_node( self, mock_init_cluster, mock_create_app, mock_uvicorn ): @@ -126,7 +126,7 @@ def test_main_shared_master_node( mock_env.CLUSTER_TIER = None # No tier configured with patch.object(sys, "argv", test_args): - from robosystems.kuzu_api.main import main + from robosystems.graph_api.main import main main() @@ -139,9 +139,9 @@ def test_main_shared_master_node( repository_type=RepositoryType.SHARED, ) - @patch("robosystems.kuzu_api.main.uvicorn.run") - @patch("robosystems.kuzu_api.main.create_app") - @patch("robosystems.kuzu_api.main.init_cluster_service") + @patch("robosystems.graph_api.main.uvicorn.run") + @patch("robosystems.graph_api.main.create_app") + @patch("robosystems.graph_api.main.init_cluster_service") def test_main_with_tier_config( self, mock_init_cluster, mock_create_app, mock_uvicorn ): @@ -168,7 +168,7 @@ def test_main_with_tier_config( } with patch.object(sys, "argv", test_args): - from robosystems.kuzu_api.main import main + from robosystems.graph_api.main import main main() @@ -196,7 +196,7 @@ def test_main_invalid_node_repository_combination(self): with patch.object(sys, "argv", test_args): with pytest.raises(SystemExit): - from robosystems.kuzu_api.main import main + from robosystems.graph_api.main import main main() @@ -215,13 +215,13 @@ def test_main_invalid_shared_node_repository_combination(self): with patch.object(sys, "argv", test_args): with pytest.raises(SystemExit): - from robosystems.kuzu_api.main import main + from robosystems.graph_api.main import main main() - @patch("robosystems.kuzu_api.main.uvicorn.run") - @patch("robosystems.kuzu_api.main.create_app") - @patch("robosystems.kuzu_api.main.init_cluster_service") + @patch("robosystems.graph_api.main.uvicorn.run") + @patch("robosystems.graph_api.main.create_app") + @patch("robosystems.graph_api.main.init_cluster_service") def test_main_with_workers(self, mock_init_cluster, mock_create_app, mock_uvicorn): """Test main with multiple workers.""" with tempfile.TemporaryDirectory() as tmpdir: @@ -241,7 +241,7 @@ def test_main_with_workers(self, mock_init_cluster, mock_create_app, mock_uvicor mock_env.CLUSTER_TIER = None with patch.object(sys, "argv", test_args): - from robosystems.kuzu_api.main import main + from robosystems.graph_api.main import main main() @@ -250,9 +250,9 @@ def test_main_with_workers(self, mock_init_cluster, mock_create_app, mock_uvicor call_args = mock_uvicorn.call_args[1] assert call_args["workers"] == 4 - @patch("robosystems.kuzu_api.main.uvicorn.run") - @patch("robosystems.kuzu_api.main.create_app") - @patch("robosystems.kuzu_api.main.init_cluster_service") + @patch("robosystems.graph_api.main.uvicorn.run") + @patch("robosystems.graph_api.main.create_app") + @patch("robosystems.graph_api.main.init_cluster_service") def test_main_with_log_level(self, mock_init_cluster, mock_create_app, mock_uvicorn): """Test main with custom log level.""" with tempfile.TemporaryDirectory() as tmpdir: @@ -272,7 +272,7 @@ def test_main_with_log_level(self, mock_init_cluster, mock_create_app, mock_uvic mock_env.CLUSTER_TIER = None with patch.object(sys, "argv", test_args): - from robosystems.kuzu_api.main import main + from robosystems.graph_api.main import main main() @@ -281,9 +281,9 @@ def test_main_with_log_level(self, mock_init_cluster, mock_create_app, mock_uvic call_args = mock_uvicorn.call_args[1] assert call_args["log_level"] == "debug" - @patch("robosystems.kuzu_api.main.uvicorn.run") - @patch("robosystems.kuzu_api.main.create_app") - @patch("robosystems.kuzu_api.main.init_cluster_service") + @patch("robosystems.graph_api.main.uvicorn.run") + @patch("robosystems.graph_api.main.create_app") + @patch("robosystems.graph_api.main.init_cluster_service") def test_main_base_path_creation( self, mock_init_cluster, mock_create_app, mock_uvicorn ): @@ -304,7 +304,7 @@ def test_main_base_path_creation( mock_env.CLUSTER_TIER = None with patch.object(sys, "argv", test_args): - from robosystems.kuzu_api.main import main + from robosystems.graph_api.main import main main() @@ -312,9 +312,9 @@ def test_main_base_path_creation( assert base_path.exists() assert base_path.is_dir() - @patch("robosystems.kuzu_api.main.uvicorn.run") - @patch("robosystems.kuzu_api.main.create_app") - @patch("robosystems.kuzu_api.main.init_cluster_service") + @patch("robosystems.graph_api.main.uvicorn.run") + @patch("robosystems.graph_api.main.create_app") + @patch("robosystems.graph_api.main.init_cluster_service") def test_main_tier_config_exception_handling( self, mock_init_cluster, mock_create_app, mock_uvicorn ): @@ -338,7 +338,7 @@ def test_main_tier_config_exception_handling( mock_env.get_kuzu_tier_config.side_effect = Exception("Config error") with patch.object(sys, "argv", test_args): - from robosystems.kuzu_api.main import main + from robosystems.graph_api.main import main main() @@ -347,9 +347,9 @@ def test_main_tier_config_exception_handling( call_args = mock_init_cluster.call_args[1] assert call_args["max_databases"] == 150 # Falls back to CLI arg - @patch("robosystems.kuzu_api.main.uvicorn.run") - @patch("robosystems.kuzu_api.main.create_app") - @patch("robosystems.kuzu_api.main.init_cluster_service") + @patch("robosystems.graph_api.main.uvicorn.run") + @patch("robosystems.graph_api.main.create_app") + @patch("robosystems.graph_api.main.init_cluster_service") def test_main_shared_replica_node( self, mock_init_cluster, mock_create_app, mock_uvicorn ): @@ -374,7 +374,7 @@ def test_main_shared_replica_node( mock_env.CLUSTER_TIER = None with patch.object(sys, "argv", test_args): - from robosystems.kuzu_api.main import main + from robosystems.graph_api.main import main main() diff --git a/tests/kuzu_api/test_schema_security.py b/tests/graph_api/test_schema_security.py similarity index 94% rename from tests/kuzu_api/test_schema_security.py rename to tests/graph_api/test_schema_security.py index 55d28fb6..cdefa199 100644 --- a/tests/kuzu_api/test_schema_security.py +++ b/tests/graph_api/test_schema_security.py @@ -4,12 +4,12 @@ from unittest.mock import Mock, patch from fastapi.testclient import TestClient -from robosystems.kuzu_api.routers.databases.schema import ( +from robosystems.graph_api.routers.databases.schema import ( validate_ddl_statement, escape_identifier, ) -from robosystems.kuzu_api.app import create_app -from robosystems.kuzu_api.core import init_cluster_service +from robosystems.graph_api.app import create_app +from robosystems.graph_api.core import init_cluster_service from robosystems.middleware.graph.clusters import NodeType, RepositoryType @@ -124,7 +124,7 @@ def test_sql_injection_in_identifiers_blocked(self): @pytest.fixture def test_client(): """Create a test client with mocked cluster service.""" - from robosystems.kuzu_api.core import cluster_manager + from robosystems.graph_api.core import cluster_manager # Reset cluster service before test original_service = cluster_manager._cluster_service @@ -153,7 +153,7 @@ class TestSchemaEndpointSecurity: def test_path_traversal_in_database_name(self, test_client): """Test path traversal attempts in database names.""" with patch( - "robosystems.kuzu_api.core.cluster_manager.get_cluster_service" + "robosystems.graph_api.core.cluster_manager.get_cluster_service" ) as mock_get_service: mock_service = Mock() mock_service.read_only = False @@ -187,7 +187,7 @@ def test_path_traversal_in_database_name(self, test_client): def test_sql_injection_in_schema_endpoint(self, test_client): """Test SQL injection prevention in schema endpoints.""" - from robosystems.kuzu_api.core import cluster_manager + from robosystems.graph_api.core import cluster_manager # Access the actual cluster service that was initialized in the fixture cluster_service = cluster_manager._cluster_service @@ -203,7 +203,7 @@ def test_sql_injection_in_schema_endpoint(self, test_client): with patch.object(cluster_service, "execute_query", return_value=mock_result): # Test SQL injection in table info call with patch( - "robosystems.kuzu_api.routers.databases.schema.escape_identifier" + "robosystems.graph_api.routers.databases.schema.escape_identifier" ) as mock_escape: # Should call escape_identifier for table names mock_escape.side_effect = ValueError("Invalid identifier") diff --git a/tests/kuzu_api/test_task_sse.py b/tests/graph_api/test_task_sse.py similarity index 98% rename from tests/kuzu_api/test_task_sse.py rename to tests/graph_api/test_task_sse.py index 4b3488dd..b0263ac2 100644 --- a/tests/kuzu_api/test_task_sse.py +++ b/tests/graph_api/test_task_sse.py @@ -1,10 +1,10 @@ -"""Tests for kuzu_api task_sse module.""" +"""Tests for graph_api task_sse module.""" import json from unittest.mock import AsyncMock, patch import pytest -from robosystems.kuzu_api.core.task_sse import ( +from robosystems.graph_api.core.task_sse import ( TaskType, generate_task_sse_events, _get_progress_message, @@ -158,7 +158,7 @@ async def test_generate_task_sse_events_heartbeat(self): } events = [] - with patch("robosystems.kuzu_api.core.task_sse.time.time") as mock_time: + with patch("robosystems.graph_api.core.task_sse.time.time") as mock_time: # Simulate time passing for heartbeat mock_time.side_effect = [0, 0, 35, 35, 70] # Trigger heartbeat after 35 seconds diff --git a/tests/integration/test_circuit_breaker.py b/tests/integration/test_circuit_breaker.py index 1f223d18..f5253f59 100644 --- a/tests/integration/test_circuit_breaker.py +++ b/tests/integration/test_circuit_breaker.py @@ -5,7 +5,7 @@ import asyncio import pytest -from robosystems.kuzu_api.client.factory import CircuitBreaker +from robosystems.graph_api.client.factory import CircuitBreaker class TestCircuitBreaker: diff --git a/tests/integration/test_query_timeout.py b/tests/integration/test_query_timeout.py index 97041c66..ccb774fd 100644 --- a/tests/integration/test_query_timeout.py +++ b/tests/integration/test_query_timeout.py @@ -6,8 +6,8 @@ from concurrent.futures import ThreadPoolExecutor, TimeoutError as FuturesTimeoutError from unittest.mock import MagicMock, patch -from robosystems.kuzu_api.core.cluster_manager import KuzuClusterService -from robosystems.kuzu_api.models.database import QueryRequest +from robosystems.graph_api.core.cluster_manager import KuzuClusterService +from robosystems.graph_api.models.database import QueryRequest from robosystems.middleware.graph.clusters import NodeType, RepositoryType from fastapi import HTTPException @@ -46,7 +46,7 @@ def slow_operation(): # Verify the future is cancelled or still running assert future.cancelled() or future.running() - @patch("robosystems.kuzu_api.core.cluster_manager.KuzuDatabaseManager") + @patch("robosystems.graph_api.core.cluster_manager.KuzuDatabaseManager") def test_query_timeout_with_slow_query(self, mock_db_manager): """Test actual query timeout with simulated slow query.""" from robosystems.config import env diff --git a/tests/middleware/mcp/test_client.py b/tests/middleware/mcp/test_client.py index 3613af8b..e1fc2355 100644 --- a/tests/middleware/mcp/test_client.py +++ b/tests/middleware/mcp/test_client.py @@ -458,7 +458,7 @@ async def test_create_kuzu_mcp_client_default(self): patch("robosystems.middleware.mcp.client.KuzuClient"), patch("robosystems.middleware.mcp.client.httpx.AsyncClient"), patch( - "robosystems.kuzu_api.client.factory.KuzuClientFactory.create_client" + "robosystems.graph_api.client.factory.KuzuClientFactory.create_client" ) as mock_factory, patch( "robosystems.middleware.graph.multitenant_utils.MultiTenantUtils.is_shared_repository", @@ -491,7 +491,7 @@ async def test_create_kuzu_mcp_client_prod(self): patch("robosystems.middleware.mcp.client.KuzuClient"), patch("robosystems.middleware.mcp.client.httpx.AsyncClient"), patch( - "robosystems.kuzu_api.client.factory.KuzuClientFactory.create_client" + "robosystems.graph_api.client.factory.KuzuClientFactory.create_client" ) as mock_factory, patch( "robosystems.middleware.graph.multitenant_utils.MultiTenantUtils.is_shared_repository", diff --git a/tests/middleware/mcp/test_mcp_factory.py b/tests/middleware/mcp/test_mcp_factory.py index ae9b08fb..46bcf156 100644 --- a/tests/middleware/mcp/test_mcp_factory.py +++ b/tests/middleware/mcp/test_mcp_factory.py @@ -38,7 +38,7 @@ async def test_create_client_with_explicit_url(self, mock_client_class, mock_env @patch("robosystems.middleware.mcp.factory.env") @patch("robosystems.middleware.mcp.factory.KuzuMCPClient") - @patch("robosystems.kuzu_api.client.factory.KuzuClientFactory") + @patch("robosystems.graph_api.client.factory.KuzuClientFactory") @patch("robosystems.middleware.graph.multitenant_utils.MultiTenantUtils") async def test_create_client_with_discovery_shared_repo( self, mock_utils, mock_factory, mock_client_class, mock_env @@ -81,7 +81,7 @@ async def test_create_client_with_discovery_shared_repo( @patch("robosystems.middleware.mcp.factory.env") @patch("robosystems.middleware.mcp.factory.KuzuMCPClient") - @patch("robosystems.kuzu_api.client.factory.KuzuClientFactory") + @patch("robosystems.graph_api.client.factory.KuzuClientFactory") @patch("robosystems.middleware.graph.multitenant_utils.MultiTenantUtils") async def test_create_client_with_discovery_user_graph( self, mock_utils, mock_factory, mock_client_class, mock_env @@ -125,7 +125,7 @@ async def test_create_client_with_discovery_user_graph( @patch("robosystems.middleware.mcp.factory.env") @patch("robosystems.middleware.mcp.factory.KuzuMCPClient") - @patch("robosystems.kuzu_api.client.factory.KuzuClientFactory") + @patch("robosystems.graph_api.client.factory.KuzuClientFactory") @patch("robosystems.middleware.graph.multitenant_utils.MultiTenantUtils") async def test_create_client_url_discovery_fallback_base_url( self, mock_utils, mock_factory, mock_client_class, mock_env @@ -159,7 +159,7 @@ async def test_create_client_url_discovery_fallback_base_url( @patch("robosystems.middleware.mcp.factory.env") @patch("robosystems.middleware.mcp.factory.KuzuMCPClient") - @patch("robosystems.kuzu_api.client.factory.KuzuClientFactory") + @patch("robosystems.graph_api.client.factory.KuzuClientFactory") @patch("robosystems.middleware.graph.multitenant_utils.MultiTenantUtils") async def test_create_client_url_discovery_env_fallback( self, mock_utils, mock_factory, mock_client_class, mock_env @@ -194,7 +194,7 @@ async def test_create_client_url_discovery_env_fallback( @patch("robosystems.middleware.mcp.factory.env") @patch("robosystems.middleware.mcp.factory.KuzuMCPClient") - @patch("robosystems.kuzu_api.client.factory.KuzuClientFactory") + @patch("robosystems.graph_api.client.factory.KuzuClientFactory") @patch("robosystems.middleware.graph.multitenant_utils.MultiTenantUtils") async def test_create_client_url_discovery_final_fallback( self, mock_utils, mock_factory, mock_client_class, mock_env @@ -274,7 +274,7 @@ async def test_create_client_default_graph_id(self, mock_client_class, mock_env) @patch("robosystems.middleware.mcp.factory.logger") @patch("robosystems.middleware.mcp.factory.env") @patch("robosystems.middleware.mcp.factory.KuzuMCPClient") - @patch("robosystems.kuzu_api.client.factory.KuzuClientFactory") + @patch("robosystems.graph_api.client.factory.KuzuClientFactory") @patch("robosystems.middleware.graph.multitenant_utils.MultiTenantUtils") async def test_create_client_logs_discovery( self, mock_utils, mock_factory, mock_client_class, mock_env, mock_logger diff --git a/tests/operations/graph/test_entity_service.py b/tests/operations/graph/test_entity_service.py index 06cc2ad2..f9c5f2f8 100644 --- a/tests/operations/graph/test_entity_service.py +++ b/tests/operations/graph/test_entity_service.py @@ -233,7 +233,7 @@ async def test_install_entity_schema_unknown_extension(self, mocker): @pytest.mark.asyncio async def test_create_graph_metadata_node_error_suppression(self, mocker): """Test that GraphMetadata creation errors are suppressed.""" - from robosystems.kuzu_api.client.exceptions import KuzuClientError + from robosystems.graph_api.client.exceptions import KuzuClientError mock_kuzu_client = mocker.AsyncMock() mock_kuzu_client.query.side_effect = KuzuClientError("Duplicate node") diff --git a/tests/operations/graph/test_generic_graph_service.py b/tests/operations/graph/test_generic_graph_service.py index aa7d3ea8..18478869 100644 --- a/tests/operations/graph/test_generic_graph_service.py +++ b/tests/operations/graph/test_generic_graph_service.py @@ -119,7 +119,7 @@ async def test_create_graph_success( mock_alloc_class.return_value = mock_manager with patch( - "robosystems.kuzu_api.client.get_kuzu_client_for_instance" + "robosystems.graph_api.client.get_kuzu_client_for_instance" ) as mock_get_client: mock_get_client.return_value = mock_kuzu_client @@ -200,7 +200,7 @@ async def test_create_graph_with_custom_schema( mock_alloc_class.return_value = mock_manager with patch( - "robosystems.kuzu_api.client.get_kuzu_client_for_instance" + "robosystems.graph_api.client.get_kuzu_client_for_instance" ) as mock_get_client: mock_get_client.return_value = mock_kuzu_client @@ -420,7 +420,7 @@ def cancellation_callback(): mock_alloc_class.return_value = mock_manager with patch( - "robosystems.kuzu_api.client.get_kuzu_client_for_instance" + "robosystems.graph_api.client.get_kuzu_client_for_instance" ) as mock_get_client: mock_get_client.return_value = mock_kuzu_client @@ -472,7 +472,7 @@ async def test_create_graph_schema_installation_failure( mock_alloc_class.return_value = mock_manager with patch( - "robosystems.kuzu_api.client.get_kuzu_client_for_instance" + "robosystems.graph_api.client.get_kuzu_client_for_instance" ) as mock_get_client: mock_get_client.return_value = mock_kuzu_client @@ -520,7 +520,7 @@ async def test_create_graph_with_initial_data( mock_alloc_class.return_value = mock_manager with patch( - "robosystems.kuzu_api.client.get_kuzu_client_for_instance" + "robosystems.graph_api.client.get_kuzu_client_for_instance" ) as mock_get_client: mock_get_client.return_value = mock_kuzu_client @@ -581,7 +581,7 @@ async def test_create_graph_metadata_storage_failure( mock_alloc_class.return_value = mock_manager with patch( - "robosystems.kuzu_api.client.get_kuzu_client_for_instance" + "robosystems.graph_api.client.get_kuzu_client_for_instance" ) as mock_get_client: mock_get_client.return_value = mock_kuzu_client @@ -642,7 +642,7 @@ async def test_create_graph_credit_pool_failure( mock_alloc_class.return_value = mock_manager with patch( - "robosystems.kuzu_api.client.get_kuzu_client_for_instance" + "robosystems.graph_api.client.get_kuzu_client_for_instance" ) as mock_get_client: mock_get_client.return_value = mock_kuzu_client @@ -712,7 +712,7 @@ async def test_create_graph_with_custom_metadata( mock_alloc_class.return_value = mock_manager with patch( - "robosystems.kuzu_api.client.get_kuzu_client_for_instance" + "robosystems.graph_api.client.get_kuzu_client_for_instance" ) as mock_get_client: mock_get_client.return_value = mock_kuzu_client @@ -776,7 +776,7 @@ async def test_create_graph_different_tiers( mock_alloc_class.return_value = mock_manager with patch( - "robosystems.kuzu_api.client.get_kuzu_client_for_instance" + "robosystems.graph_api.client.get_kuzu_client_for_instance" ) as mock_get_client: mock_get_client.return_value = mock_kuzu_client @@ -901,7 +901,7 @@ async def test_full_graph_creation_lifecycle(self): mock_alloc_class.return_value = mock_manager with patch( - "robosystems.kuzu_api.client.get_kuzu_client_for_instance" + "robosystems.graph_api.client.get_kuzu_client_for_instance" ) as mock_get_client: mock_client = AsyncMock() mock_client.create_database.return_value = {"status": "created"} diff --git a/tests/operations/graph/test_shared_repository_service.py b/tests/operations/graph/test_shared_repository_service.py index 1330d3b7..5765790d 100644 --- a/tests/operations/graph/test_shared_repository_service.py +++ b/tests/operations/graph/test_shared_repository_service.py @@ -44,7 +44,7 @@ def mock_kuzu_client(self): async def test_create_shared_repository_sec(self, service, mock_kuzu_client): """Test successful creation of SEC shared repository.""" with patch( - "robosystems.kuzu_api.client.factory.KuzuClientFactory._get_shared_master_url" + "robosystems.graph_api.client.factory.KuzuClientFactory._get_shared_master_url" ) as mock_get_url: mock_get_url.return_value = "http://10.0.1.100:8080" @@ -78,7 +78,7 @@ async def test_create_shared_repository_all_types(self, service, mock_kuzu_clien for repo_name in valid_repos: with patch( - "robosystems.kuzu_api.client.factory.KuzuClientFactory._get_shared_master_url" + "robosystems.graph_api.client.factory.KuzuClientFactory._get_shared_master_url" ) as mock_get_url: mock_get_url.return_value = "http://10.0.1.100:8080" @@ -110,7 +110,7 @@ async def test_create_shared_repository_invalid_name(self, service): async def test_create_shared_repository_invalid_url_format(self, service): """Test handling of invalid shared master URL format.""" with patch( - "robosystems.kuzu_api.client.factory.KuzuClientFactory._get_shared_master_url" + "robosystems.graph_api.client.factory.KuzuClientFactory._get_shared_master_url" ) as mock_get_url: # Return invalid URL format mock_get_url.return_value = "invalid://url:format" @@ -132,7 +132,7 @@ async def test_create_shared_repository_unhealthy_database( } with patch( - "robosystems.kuzu_api.client.factory.KuzuClientFactory._get_shared_master_url" + "robosystems.graph_api.client.factory.KuzuClientFactory._get_shared_master_url" ) as mock_get_url: mock_get_url.return_value = "http://10.0.1.100:8080" @@ -153,7 +153,7 @@ async def test_create_shared_repository_unhealthy_database( async def test_create_shared_repository_client_error(self, service): """Test handling of Kuzu client errors.""" with patch( - "robosystems.kuzu_api.client.factory.KuzuClientFactory._get_shared_master_url" + "robosystems.graph_api.client.factory.KuzuClientFactory._get_shared_master_url" ) as mock_get_url: mock_get_url.return_value = "http://10.0.1.100:8080" @@ -177,7 +177,7 @@ async def test_create_shared_repository_creation_failure( ) with patch( - "robosystems.kuzu_api.client.factory.KuzuClientFactory._get_shared_master_url" + "robosystems.graph_api.client.factory.KuzuClientFactory._get_shared_master_url" ) as mock_get_url: mock_get_url.return_value = "http://10.0.1.100:8080" @@ -198,7 +198,7 @@ async def test_create_shared_repository_creation_failure( async def test_create_shared_repository_no_user(self, service, mock_kuzu_client): """Test creation without specifying created_by user.""" with patch( - "robosystems.kuzu_api.client.factory.KuzuClientFactory._get_shared_master_url" + "robosystems.graph_api.client.factory.KuzuClientFactory._get_shared_master_url" ) as mock_get_url: mock_get_url.return_value = "http://10.0.1.100:8080" @@ -223,7 +223,7 @@ async def test_url_parsing(self, service): for url, expected_ip in test_cases: with patch( - "robosystems.kuzu_api.client.factory.KuzuClientFactory._get_shared_master_url" + "robosystems.graph_api.client.factory.KuzuClientFactory._get_shared_master_url" ) as mock_get_url: mock_get_url.return_value = url @@ -267,7 +267,7 @@ async def test_repository_already_exists(self, mock_client): } with patch( - "robosystems.kuzu_api.client.factory.KuzuClientFactory.create_client" + "robosystems.graph_api.client.factory.KuzuClientFactory.create_client" ) as mock_factory: mock_factory.return_value = mock_client @@ -288,7 +288,7 @@ async def test_repository_does_not_exist_creates_it(self, mock_client): mock_client.get_database_info.side_effect = Exception("Database not found") with patch( - "robosystems.kuzu_api.client.factory.KuzuClientFactory.create_client" + "robosystems.graph_api.client.factory.KuzuClientFactory.create_client" ) as mock_factory: mock_factory.return_value = mock_client @@ -319,7 +319,7 @@ async def test_repository_exists_but_unhealthy(self, mock_client): } with patch( - "robosystems.kuzu_api.client.factory.KuzuClientFactory.create_client" + "robosystems.graph_api.client.factory.KuzuClientFactory.create_client" ) as mock_factory: mock_factory.return_value = mock_client @@ -346,7 +346,7 @@ async def test_repository_with_custom_url(self, mock_client): } with patch( - "robosystems.kuzu_api.client.factory.KuzuClientFactory.create_client" + "robosystems.graph_api.client.factory.KuzuClientFactory.create_client" ) as mock_factory: mock_factory.return_value = mock_client @@ -362,7 +362,7 @@ async def test_repository_creation_error_handling(self, mock_client): mock_client.get_database_info.side_effect = Exception("Not found") with patch( - "robosystems.kuzu_api.client.factory.KuzuClientFactory.create_client" + "robosystems.graph_api.client.factory.KuzuClientFactory.create_client" ) as mock_factory: mock_factory.return_value = mock_client @@ -388,7 +388,7 @@ async def test_full_repository_lifecycle(self): service = SharedRepositoryService() with patch( - "robosystems.kuzu_api.client.factory.KuzuClientFactory._get_shared_master_url" + "robosystems.graph_api.client.factory.KuzuClientFactory._get_shared_master_url" ) as mock_get_url: mock_get_url.return_value = "http://10.0.1.100:8080" @@ -419,7 +419,7 @@ async def test_full_repository_lifecycle(self): # Verify it exists with patch( - "robosystems.kuzu_api.client.factory.KuzuClientFactory.create_client" + "robosystems.graph_api.client.factory.KuzuClientFactory.create_client" ) as mock_factory: mock_factory.return_value = mock_client @@ -435,7 +435,7 @@ async def test_concurrent_repository_creation(self): service = SharedRepositoryService() with patch( - "robosystems.kuzu_api.client.factory.KuzuClientFactory._get_shared_master_url" + "robosystems.graph_api.client.factory.KuzuClientFactory._get_shared_master_url" ) as mock_get_url: mock_get_url.return_value = "http://10.0.1.100:8080" diff --git a/tests/tasks/sec_xbrl/conftest.py b/tests/tasks/sec_xbrl/conftest.py index 7374e8bb..5b08d2e4 100644 --- a/tests/tasks/sec_xbrl/conftest.py +++ b/tests/tasks/sec_xbrl/conftest.py @@ -110,7 +110,7 @@ def mock_kuzu_client(): @pytest.fixture def mock_kuzu_factory(mock_kuzu_client): """Mock KuzuClientFactory.""" - with patch("robosystems.kuzu_api.client.factory.KuzuClientFactory") as mock_factory: + with patch("robosystems.graph_api.client.factory.KuzuClientFactory") as mock_factory: mock_factory.create_client = AsyncMock(return_value=mock_kuzu_client) yield mock_factory diff --git a/tests/tasks/sec_xbrl/test_ingestion.py b/tests/tasks/sec_xbrl/test_ingestion.py index 5d8d7e56..83874170 100644 --- a/tests/tasks/sec_xbrl/test_ingestion.py +++ b/tests/tasks/sec_xbrl/test_ingestion.py @@ -30,7 +30,7 @@ def test_schema_types_structure(self): # Note: RELATIONSHIPS_NEEDING_IGNORE_ERRORS is no longer needed # since we always use ignore_errors=True for SEC data - @patch("robosystems.kuzu_api.client.factory.KuzuClientFactory") + @patch("robosystems.graph_api.client.factory.KuzuClientFactory") @patch("robosystems.adapters.s3.S3Client") def test_ingestion_flow_logic(self, mock_s3_class, mock_kuzu_factory): """Test the core ingestion flow logic.""" diff --git a/tests/unit/graph_api/__init__.py b/tests/unit/graph_api/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/unit/graph_api/backends/__init__.py b/tests/unit/graph_api/backends/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/unit/graph_api/backends/test_backend_factory.py b/tests/unit/graph_api/backends/test_backend_factory.py new file mode 100644 index 00000000..7997bea6 --- /dev/null +++ b/tests/unit/graph_api/backends/test_backend_factory.py @@ -0,0 +1,59 @@ +import pytest +from robosystems.graph_api.backends import get_backend, KuzuBackend, Neo4jBackend + + +def test_backend_factory_kuzu(monkeypatch, tmp_path): + monkeypatch.setenv("BACKEND_TYPE", "kuzu") + monkeypatch.setattr("robosystems.graph_api.backends._backend_instance", None) + monkeypatch.setattr( + "robosystems.graph_api.backends.kuzu.KuzuBackend.__init__", + lambda self: setattr(self, "data_path", str(tmp_path)) + or setattr(self, "_engines", {}), + ) + + backend = get_backend() + + assert isinstance(backend, KuzuBackend) + + +def test_backend_factory_neo4j_community(monkeypatch): + monkeypatch.setenv("BACKEND_TYPE", "neo4j_community") + monkeypatch.setattr("robosystems.graph_api.backends._backend_instance", None) + + backend = get_backend() + + assert isinstance(backend, Neo4jBackend) + assert backend.enterprise is False + + +def test_backend_factory_neo4j_enterprise(monkeypatch): + monkeypatch.setenv("BACKEND_TYPE", "neo4j_enterprise") + monkeypatch.setattr("robosystems.graph_api.backends._backend_instance", None) + + backend = get_backend() + + assert isinstance(backend, Neo4jBackend) + assert backend.enterprise is True + + +def test_backend_factory_invalid_type(monkeypatch): + monkeypatch.setenv("BACKEND_TYPE", "invalid_backend") + monkeypatch.setattr("robosystems.graph_api.backends._backend_instance", None) + + with pytest.raises(ValueError, match="Unknown BACKEND_TYPE"): + get_backend() + + +def test_backend_factory_singleton(monkeypatch, tmp_path): + monkeypatch.setenv("BACKEND_TYPE", "kuzu") + monkeypatch.setattr("robosystems.graph_api.backends._backend_instance", None) + monkeypatch.setattr( + "robosystems.graph_api.backends.kuzu.KuzuBackend.__init__", + lambda self: setattr(self, "data_path", str(tmp_path)) + or setattr(self, "_engines", {}), + ) + + backend1 = get_backend() + backend2 = get_backend() + + assert backend1 is backend2 diff --git a/uv.lock b/uv.lock index 24d46eef..057ebbef 100644 --- a/uv.lock +++ b/uv.lock @@ -1984,6 +1984,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/a9/82/0340caa499416c78e5d8f5f05947ae4bc3cba53c9f038ab6e9ed964e22f1/nbformat-5.10.4-py3-none-any.whl", hash = "sha256:3b48d6c8fbca4b299bf3982ea7db1af21580e4fec269ad087b9e81588891200b", size = 78454, upload-time = "2024-04-04T11:20:34.895Z" }, ] +[[package]] +name = "neo4j" +version = "5.28.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "pytz" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/50/69/4862fabc082f2447131aada5c91736155349d77ebf443af7f59553b7b789/neo4j-5.28.2.tar.gz", hash = "sha256:7d38e27e4f987a45cc9052500c6ee27325cb23dae6509037fe31dd7ddaed70c7", size = 231874, upload-time = "2025-07-30T06:04:34.669Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/04/00/1f74089c06aec1fac9390e2300a6a6b2381e0dac281783d64ccca9d681fd/neo4j-5.28.2-py3-none-any.whl", hash = "sha256:5c53b5c3eee6dee7e920c9724391aa38d7135a651e71b766da00533b92a91a94", size = 313156, upload-time = "2025-07-30T06:04:31.438Z" }, +] + [[package]] name = "nest-asyncio" version = "1.6.0" @@ -3295,6 +3307,7 @@ dependencies = [ { name = "kuzu" }, { name = "logger" }, { name = "lxml" }, + { name = "neo4j" }, { name = "opentelemetry-api" }, { name = "opentelemetry-exporter-otlp" }, { name = "opentelemetry-instrumentation-aiohttp-client" }, @@ -3364,10 +3377,11 @@ requires-dist = [ { name = "httpx-sse", specifier = ">=0.4.0" }, { name = "intuit-oauth", specifier = ">=1.2.0,<2.0" }, { name = "jupyter", marker = "extra == 'dev'", specifier = ">=1.1.0,<2.0" }, - { name = "kuzu", specifier = ">=0.11.0,<0.12" }, + { name = "kuzu", specifier = "==0.11.2" }, { name = "logger", specifier = ">=1.4,<2.0" }, { name = "lxml", specifier = ">=5.4.0,<6.0" }, { name = "moto", extras = ["s3"], marker = "extra == 'dev'", specifier = ">=5.1.6,<6.0" }, + { name = "neo4j", specifier = ">=5.15.0,<6.0" }, { name = "opentelemetry-api", specifier = ">=1.35.0,<2.0" }, { name = "opentelemetry-exporter-otlp", specifier = ">=1.35.0,<2.0" }, { name = "opentelemetry-instrumentation-aiohttp-client", specifier = ">=0.55b0,<1.0" }, From b48b73f570cc0e1e4ce8a4026f955a5acc340ff1 Mon Sep 17 00:00:00 2001 From: Joey French Date: Wed, 15 Oct 2025 21:28:31 -0500 Subject: [PATCH 2/5] Enhance Neo4j backend tests with environment variable configuration - Added environment variable settings for Neo4j URI, username, and password in the test cases for both community and enterprise backends. - Updated the test for invalid backend type to include the corresponding environment variable setting. --- .../graph_api/backends/test_backend_factory.py | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/tests/unit/graph_api/backends/test_backend_factory.py b/tests/unit/graph_api/backends/test_backend_factory.py index 7997bea6..e7161eb3 100644 --- a/tests/unit/graph_api/backends/test_backend_factory.py +++ b/tests/unit/graph_api/backends/test_backend_factory.py @@ -18,6 +18,13 @@ def test_backend_factory_kuzu(monkeypatch, tmp_path): def test_backend_factory_neo4j_community(monkeypatch): monkeypatch.setenv("BACKEND_TYPE", "neo4j_community") + monkeypatch.setenv("NEO4J_URI", "bolt://localhost:7687") + monkeypatch.setenv("NEO4J_USERNAME", "neo4j") + monkeypatch.setenv("NEO4J_PASSWORD", "password") + monkeypatch.setattr("robosystems.config.env.BACKEND_TYPE", "neo4j_community") + monkeypatch.setattr("robosystems.config.env.NEO4J_URI", "bolt://localhost:7687") + monkeypatch.setattr("robosystems.config.env.NEO4J_USERNAME", "neo4j") + monkeypatch.setattr("robosystems.config.env.NEO4J_PASSWORD", "password") monkeypatch.setattr("robosystems.graph_api.backends._backend_instance", None) backend = get_backend() @@ -28,6 +35,13 @@ def test_backend_factory_neo4j_community(monkeypatch): def test_backend_factory_neo4j_enterprise(monkeypatch): monkeypatch.setenv("BACKEND_TYPE", "neo4j_enterprise") + monkeypatch.setenv("NEO4J_URI", "bolt://localhost:7687") + monkeypatch.setenv("NEO4J_USERNAME", "neo4j") + monkeypatch.setenv("NEO4J_PASSWORD", "password") + monkeypatch.setattr("robosystems.config.env.BACKEND_TYPE", "neo4j_enterprise") + monkeypatch.setattr("robosystems.config.env.NEO4J_URI", "bolt://localhost:7687") + monkeypatch.setattr("robosystems.config.env.NEO4J_USERNAME", "neo4j") + monkeypatch.setattr("robosystems.config.env.NEO4J_PASSWORD", "password") monkeypatch.setattr("robosystems.graph_api.backends._backend_instance", None) backend = get_backend() @@ -38,6 +52,7 @@ def test_backend_factory_neo4j_enterprise(monkeypatch): def test_backend_factory_invalid_type(monkeypatch): monkeypatch.setenv("BACKEND_TYPE", "invalid_backend") + monkeypatch.setattr("robosystems.config.env.BACKEND_TYPE", "invalid_backend") monkeypatch.setattr("robosystems.graph_api.backends._backend_instance", None) with pytest.raises(ValueError, match="Unknown BACKEND_TYPE"): From c5bbea16ab1a1b28bc61b04c56fe12d29b27fe71 Mon Sep 17 00:00:00 2001 From: Joey French Date: Wed, 15 Oct 2025 21:44:02 -0500 Subject: [PATCH 3/5] Update Neo4j backend and environment configuration - Set default Neo4j connection parameters in `.env.example` for easier setup. - Refactored Neo4j backend methods to use asynchronous closing. - Updated `compose.yaml` to utilize environment variables for Neo4j password. - Renamed Kuzu-related middleware and classes to reflect the Graph API context. - Enhanced authentication middleware to support both Kuzu and Neo4j backends. --- .env.example | 8 ++++---- compose.yaml | 4 ++-- robosystems/graph_api/backends/base.py | 2 +- robosystems/graph_api/backends/kuzu.py | 2 +- robosystems/graph_api/backends/neo4j.py | 20 +++++++------------- robosystems/graph_api/middleware/__init__.py | 6 +++--- robosystems/graph_api/middleware/auth.py | 19 ++++++++++++------- 7 files changed, 30 insertions(+), 31 deletions(-) diff --git a/.env.example b/.env.example index e0f0821f..c8324b4a 100755 --- a/.env.example +++ b/.env.example @@ -175,10 +175,10 @@ KUZU_MAX_DATABASES_PER_NODE=50 # CLUSTER_TIER=standard ## Neo4j Backend Configuration -# NEO4J_URI=bolt://neo4j-db:7687 -# NEO4J_USERNAME=neo4j -# NEO4J_PASSWORD= # Retrieved from AWS Secrets Manager in prod/staging -# NEO4J_ENTERPRISE=false # Enable multi-database support (requires Enterprise license) +NEO4J_URI=bolt://neo4j-db:7687 +NEO4J_USERNAME=neo4j +NEO4J_PASSWORD=neo4jpassword # Retrieved from AWS Secrets Manager in prod/staging +NEO4J_ENTERPRISE=false # Enable multi-database support (requires Enterprise license) # NEO4J_MAX_CONNECTION_POOL_SIZE=50 # NEO4J_CONNECTION_ACQUISITION_TIMEOUT=60 # NEO4J_MAX_CONNECTION_LIFETIME=3600 diff --git a/compose.yaml b/compose.yaml index c9f6fbe8..6efc3282 100644 --- a/compose.yaml +++ b/compose.yaml @@ -103,7 +103,7 @@ services: container_name: neo4j-db image: neo4j:5.25-community environment: - - NEO4J_AUTH=neo4j/neo4jpassword + - NEO4J_AUTH=neo4j/${NEO4J_PASSWORD-neo4jpassword} - NEO4J_PLUGINS=["apoc"] - NEO4J_dbms_memory_pagecache_size=512M - NEO4J_dbms_memory_heap_initial__size=512M @@ -145,7 +145,7 @@ services: - BACKEND_TYPE=neo4j_community - NEO4J_URI=bolt://neo4j-db:7687 - NEO4J_USERNAME=neo4j - - NEO4J_PASSWORD=neo4jpassword + - NEO4J_PASSWORD=${NEO4J_PASSWORD} - GRAPH_API_PORT=8002 - CLUSTER_TIER=professional env_file: diff --git a/robosystems/graph_api/backends/base.py b/robosystems/graph_api/backends/base.py index f5d96aef..2a69c940 100644 --- a/robosystems/graph_api/backends/base.py +++ b/robosystems/graph_api/backends/base.py @@ -65,5 +65,5 @@ async def health_check(self) -> bool: pass @abstractmethod - def close(self) -> None: + async def close(self) -> None: pass diff --git a/robosystems/graph_api/backends/kuzu.py b/robosystems/graph_api/backends/kuzu.py index aa4be5c1..a519d90c 100644 --- a/robosystems/graph_api/backends/kuzu.py +++ b/robosystems/graph_api/backends/kuzu.py @@ -99,7 +99,7 @@ async def get_cluster_topology(self) -> ClusterTopology: async def health_check(self) -> bool: return True - def close(self) -> None: + async def close(self) -> None: for engine in self._engines.values(): engine.close() self._engines.clear() diff --git a/robosystems/graph_api/backends/neo4j.py b/robosystems/graph_api/backends/neo4j.py index 115e1308..cc78c525 100644 --- a/robosystems/graph_api/backends/neo4j.py +++ b/robosystems/graph_api/backends/neo4j.py @@ -21,7 +21,7 @@ async def _ensure_connected(self): await self._connect() async def _connect(self): - if env.NEO4J_PASSWORD: + if env.NEO4J_PASSWORD is not None and env.NEO4J_PASSWORD != "": self._password = env.NEO4J_PASSWORD else: secrets = boto3.client("secretsmanager", region_name=env.AWS_REGION) @@ -59,15 +59,10 @@ async def execute_query( db_name = self._get_database_name(graph_id, database) - cypher_upper = cypher.strip().upper() - is_read = cypher_upper.startswith(("MATCH", "RETURN", "WITH", "CALL")) - + # Use Neo4j driver's automatic routing for cluster mode + # For non-cluster mode, this has no effect async with self.driver.session(database=db_name) as session: - if self.enterprise and is_read: - result = await session.run(cypher, parameters or {}) - else: - result = await session.run(cypher, parameters or {}) - + result = await session.run(cypher, parameters or {}) records = await result.data() return records @@ -218,8 +213,7 @@ async def health_check(self) -> bool: logger.error(f"Neo4j health check failed: {e}") return False - def close(self) -> None: + async def close(self) -> None: if self.driver: - import asyncio - - asyncio.create_task(self.driver.close()) + await self.driver.close() + self.driver = None diff --git a/robosystems/graph_api/middleware/__init__.py b/robosystems/graph_api/middleware/__init__.py index 495b5027..f45831f0 100644 --- a/robosystems/graph_api/middleware/__init__.py +++ b/robosystems/graph_api/middleware/__init__.py @@ -1,8 +1,8 @@ """ -Middleware components for the Kuzu API server. +Middleware components for the Graph API server. """ -from .auth import KuzuAuthMiddleware +from .auth import GraphAuthMiddleware, KuzuAuthMiddleware from .request_limits import RequestSizeLimitMiddleware -__all__ = ["KuzuAuthMiddleware", "RequestSizeLimitMiddleware"] +__all__ = ["GraphAuthMiddleware", "KuzuAuthMiddleware", "RequestSizeLimitMiddleware"] diff --git a/robosystems/graph_api/middleware/auth.py b/robosystems/graph_api/middleware/auth.py index d0daf2a6..7024a7f1 100644 --- a/robosystems/graph_api/middleware/auth.py +++ b/robosystems/graph_api/middleware/auth.py @@ -1,8 +1,9 @@ """ -Authentication middleware for Kuzu API with environment-based security. +Authentication middleware for Graph API with environment-based security. Provides API key authentication for production/staging environments while allowing unrestricted access in development and from bastion hosts. +Supports both Kuzu and Neo4j backends. """ import json @@ -22,9 +23,9 @@ from robosystems.security import SecurityAuditLogger, SecurityEventType -class KuzuAuthMiddleware(BaseHTTPMiddleware): +class GraphAuthMiddleware(BaseHTTPMiddleware): """ - Authentication middleware for Kuzu API. + Authentication middleware for Graph API. Features: - API key authentication in production/staging @@ -32,6 +33,7 @@ class KuzuAuthMiddleware(BaseHTTPMiddleware): - Bypassed for requests from bastion hosts - Bypassed for health check endpoints - Rate limiting for failed auth attempts + - Works with both Kuzu and Neo4j backends """ # Endpoints that don't require authentication @@ -71,14 +73,14 @@ def __init__(self, app, api_key: Optional[str] = None, key_type: str = "writer") if self.auth_enabled and not self.api_key: logger.error( - f"Kuzu API key not configured for {self.key_type} in {self.environment} environment!" + f"Graph API key not configured for {self.key_type} in {self.environment} environment!" ) raise ValueError( f"KUZU_API_KEY must be set for {self.key_type} in production/staging" ) logger.info( - f"Kuzu Auth Middleware initialized - Environment: {self.environment}, " + f"Graph Auth Middleware initialized - Environment: {self.environment}, " f"Auth Enabled: {self.auth_enabled}, Key Type: {self.key_type}" ) @@ -117,8 +119,8 @@ async def dispatch(self, request: Request, call_next): def _validate_api_key(self, request: Request) -> None: """Validate API key from request headers.""" - # Check for API key in header - api_key = request.headers.get("X-Kuzu-API-Key") + # Check for API key in header (support both old and new header names) + api_key = request.headers.get("X-Graph-API-Key") or request.headers.get("X-Kuzu-API-Key") if not api_key: # Also check Authorization header auth_header = request.headers.get("Authorization", "") @@ -265,3 +267,6 @@ def create_api_key(prefix: str = "kuzu") -> tuple[str, str]: ) return api_key, key_hash + + +KuzuAuthMiddleware = GraphAuthMiddleware From 1196aee490ccf2f471f9077997aa3a6a4b1e4eef Mon Sep 17 00:00:00 2001 From: Joey French Date: Wed, 15 Oct 2025 21:51:46 -0500 Subject: [PATCH 4/5] Refactor Graph API backend imports and improve API key validation formatting - Added import for `GraphBackend` in the backend initialization file. - Reformatted API key retrieval in the authentication middleware for improved readability. --- robosystems/graph_api/backends/__init__.py | 1 + robosystems/graph_api/middleware/auth.py | 4 +++- 2 files changed, 4 insertions(+), 1 deletion(-) diff --git a/robosystems/graph_api/backends/__init__.py b/robosystems/graph_api/backends/__init__.py index 1211e79e..1a0d707a 100644 --- a/robosystems/graph_api/backends/__init__.py +++ b/robosystems/graph_api/backends/__init__.py @@ -1,5 +1,6 @@ from typing import Optional, Union from robosystems.config import env +from .base import GraphBackend from .kuzu import KuzuBackend from .neo4j import Neo4jBackend from robosystems.logger import logger diff --git a/robosystems/graph_api/middleware/auth.py b/robosystems/graph_api/middleware/auth.py index 7024a7f1..655ed63d 100644 --- a/robosystems/graph_api/middleware/auth.py +++ b/robosystems/graph_api/middleware/auth.py @@ -120,7 +120,9 @@ async def dispatch(self, request: Request, call_next): def _validate_api_key(self, request: Request) -> None: """Validate API key from request headers.""" # Check for API key in header (support both old and new header names) - api_key = request.headers.get("X-Graph-API-Key") or request.headers.get("X-Kuzu-API-Key") + api_key = request.headers.get("X-Graph-API-Key") or request.headers.get( + "X-Kuzu-API-Key" + ) if not api_key: # Also check Authorization header auth_header = request.headers.get("Authorization", "") From 5c2518259f8e35f1e5d0b5997da59243e3d33c4d Mon Sep 17 00:00:00 2001 From: Joey French Date: Wed, 15 Oct 2025 22:04:42 -0500 Subject: [PATCH 5/5] Refactor documentation and update Graph API to support multiple backends - Updated README.md to reflect support for both Kuzu and Neo4j graph databases. - Renamed Kuzu API references to Graph API for consistency across documentation. - Enhanced Graph API client and middleware documentation to clarify backend support. - Improved environment variable configurations for both Kuzu and Neo4j backends. - Added backend-specific details in the Graph API and client architecture sections. --- README.md | 66 ++-- robosystems/graph_api/README.md | 120 ++++-- robosystems/graph_api/backends/README.md | 455 +++++++++++++++++++++++ robosystems/graph_api/client/README.md | 97 +++-- robosystems/middleware/graph/README.md | 29 +- 5 files changed, 687 insertions(+), 80 deletions(-) create mode 100644 robosystems/graph_api/backends/README.md diff --git a/README.md b/README.md index 84bfc05a..2aa50b3f 100644 --- a/README.md +++ b/README.md @@ -2,7 +2,7 @@ RoboSystems is an enterprise-grade financial knowledge graph platform that transforms complex financial data into actionable intelligence through graph-based analytics and AI-powered insights. -- **Graph-Based Financial Intelligence**: Leverages Kuzu graph database technology to model complex financial relationships, enabling deep analysis of relationships between accounting, financial reporting, portfolio management, and public XBRL data +- **Graph-Based Financial Intelligence**: Leverages graph database technology (Kuzu or Neo4j) to model complex financial relationships, enabling deep analysis of relationships between accounting, financial reporting, portfolio management, and public XBRL data - **GraphRAG Architecture**: Knowledge graph-based retrieval-augmented generation for LLM-powered financial analysis over enterprise financial and operating data - **Model Context Protocol (MCP)**: Standardized server and [client](https://www.npmjs.com/package/@robosystems/mcp) for LLM integration with natural language querying - **Multi-Source Data Integration**: Seamlessly integrates QuickBooks accounting data, SEC XBRL filings (10-K, 10-Q), and custom financial datasets into a unified knowledge graph @@ -13,7 +13,7 @@ RoboSystems is an enterprise-grade financial knowledge graph platform that trans RoboSystems bridges the gap between raw financial data and actionable business intelligence by creating interconnected knowledge graphs that reveal hidden relationships, patterns, and insights that traditional databases miss. It's the backbone for next-generation financial applications that need to understand not just numbers, but the relationships and context behind them. -- **Multi-Tenant Graph Databases**: Create isolated Kuzu database instances with cluster-based scaling +- **Multi-Tenant Graph Databases**: Create isolated graph database instances (Kuzu or Neo4j) with cluster-based scaling - **AI Agent Interface**: Natural language financial analysis through Claude powered agents via Model Context Protocol (MCP) - **Entity Graph Creation**: Curated enterprise financial data schemas for defined use cases with RoboLedger, RoboInvestor and more - **Generic Graph Creation**: Custom schema definitions with custom node/relationship types @@ -41,7 +41,7 @@ just start This initializes the `.env` file and starts the complete RoboSystems stack with: -- Kuzu graph database +- Graph database (Kuzu by default, Neo4j optional) - PostgreSQL with automatic migrations - Valkey message broker - All development services @@ -111,32 +111,42 @@ just logs-follow worker # CloudWatch log search - **MCP Integration**: Model Context Protocol for AI-powered financial analytics - **Celery Workers** with priority queues for asynchronous processing -### Kuzu Graph Database System +### Graph Database System -**Kuzu** is a high-performance embedded graph database that powers RoboSystems' financial knowledge graph platform. This system provides multi-tenant graph databases with enterprise-grade scaling and reliability. +RoboSystems supports **pluggable graph database backends** to provide flexibility and choice for different deployment scenarios: -- **Cluster-Based Infrastructure**: Tiered instances (Standard/Enterprise/Premium) for different workload requirements -- **Multi-Tenant Isolation**: Each entity gets a dedicated database (`kg12345abc`) with complete data isolation -- **Shared Repositories**: Common databases for SEC filings, industry benchmarks, and economic indicators -- **API-First Design**: All database access through REST APIs with no direct database connections -- **Schema-Driven Operations**: All graph operations derive from curated schemas (RoboLedger, RoboInvestor, and more) +#### Supported Backends + +- **Kuzu** (Default): High-performance embedded graph database, ideal for Standard tier deployments +- **Neo4j Community**: Client-server architecture for Professional/Enterprise tiers with advanced features +- **Neo4j Enterprise**: Full enterprise features including multi-database support for Premium tier -#### Kuzu API System (`/robosystems/graph_api/`) +#### Graph API System (`/robosystems/graph_api/`) -The **Kuzu API** is a FastAPI microservice that runs alongside Kuzu databases on instances, providing: +The **Graph API** is a FastAPI microservice that provides a unified interface regardless of backend: -- **HTTP REST Interface**: High-performance API for all graph operations (port 8001) -- **Multi-Database Management**: Handles up to 10 databases per instance (Standard tier) -- **Connection Pooling**: Efficient resource management with max 3 connections per database +- **Backend Abstraction**: Consistent API whether using Kuzu or Neo4j +- **HTTP REST Interface**: High-performance API for all graph operations (port 8001 for Kuzu, 8002 for Neo4j) +- **Multi-Database Management**: Handles multiple databases per instance (backend-dependent) +- **Connection Pooling**: Efficient resource management with backend-optimized pooling - **Async Ingestion**: Queue-based data loading with S3 integration - **Streaming Support**: NDJSON streaming for large query results - **Admission Control**: CPU/memory-based backpressure to prevent overload +#### Infrastructure Design + +- **Cluster-Based Infrastructure**: Tiered instances (Standard/Enterprise/Premium) for different workload requirements +- **Multi-Tenant Isolation**: Each entity gets a dedicated database (`kg12345abc`) with complete data isolation +- **Shared Repositories**: Common databases for SEC filings, industry benchmarks, and economic indicators +- **API-First Design**: All database access through REST APIs with no direct database connections +- **Schema-Driven Operations**: All graph operations derive from curated schemas (RoboLedger, RoboInvestor, and more) + #### Client-Factory System (`/robosystems/graph_api/client/`) -The client-factory layer provides intelligent routing between application code and Kuzu infrastructure: +The client-factory layer provides intelligent routing between application code and graph database infrastructure: -- **Automatic Discovery**: Finds database instances via DynamoDB registry +- **Backend-Agnostic**: Works seamlessly with both Kuzu and Neo4j backends +- **Automatic Discovery**: Finds database instances via DynamoDB registry (Kuzu) or direct connection (Neo4j) - **Redis Caching**: Caches instance locations to reduce lookups - **Circuit Breakers**: Prevents cascading failures with automatic recovery - **Connection Reuse**: HTTP/2 connection pooling for efficiency @@ -192,8 +202,8 @@ The client-factory layer provides intelligent routing between application code a ### Data Layer -- **Kuzu Graph Database**: Financial knowledge graph with cluster-based scaling -- **DynamoDB**: Kuzu database allocation registry, instance and volume management +- **Graph Database**: Pluggable backend (Kuzu or Neo4j) for financial knowledge graphs with cluster-based scaling +- **DynamoDB**: Database allocation registry, instance and volume management - **PostgreSQL**: Primary relational database for identity and access management - **Valkey**: Message broker and caching (separate DBs for queues, cache, progress tracking) - **AWS S3**: Document storage and database synchronization @@ -203,8 +213,8 @@ The client-factory layer provides intelligent routing between application code a - **VPC**: AWS VPC with NAT Gateway, CloudTrail, and VPC Flow Logs - **API**: ECS Fargate ARM64/Graviton with auto-scaling and WAF - **Workers**: ECS Fargate ARM64/Graviton with auto-scaling -- **Kuzu Writers**: EC2 Graviton instances with DynamoDB registry and management lambdas -- **Kuzu Readers**: EC2 Graviton instances with load balancing for shared repositories +- **Graph Database Writers**: EC2 Graviton instances (Kuzu) or ECS containers (Neo4j) with DynamoDB registry and management lambdas +- **Graph Database Readers**: EC2 Graviton instances or ECS containers with load balancing for shared repositories - **Database & Cache**: RDS PostgreSQL + ElastiCache Valkey instances - **Observability**: Amazon Managed Prometheus + Grafana with AWS SSO - **Self-Hosted CI/CD**: GitHub Actions runner on dedicated infrastructure @@ -229,7 +239,7 @@ The client-factory layer provides intelligent routing between application code a - **Multi-Agent Architecture**: Intelligent routing to specialized agents based on query context - **Dynamic Agent Selection**: Automatic selection of the most appropriate agent for each task - **Parallel Query Processing**: Batch processing of multiple queries simultaneously -- **Context-Aware Responses**: GraphRAG-enabled agents with native kuzu graph database integration +- **Context-Aware Responses**: GraphRAG-enabled agents with native graph database integration - **Extensible Framework**: Support for custom agents with specific domain expertise ### Credit System @@ -298,12 +308,13 @@ All infrastructure is managed through CloudFormation templates in `/cloudformati - **`beat.yaml`**: Celery beat scheduler for periodic tasks and cron jobs - **`worker-monitor.yaml`**: Lambda function for monitoring worker health and queue depths -#### Kuzu Graph Database +#### Graph Database Infrastructure -- **`kuzu-infra.yaml`**: Base infrastructure for Kuzu clusters (security groups, roles, registries) +- **`kuzu-infra.yaml`**: Base infrastructure for graph database clusters (security groups, roles, registries) - **`kuzu-volumes.yaml`**: EBS volume management and snapshot automation -- **`kuzu-writers.yaml`**: Auto-scaling EC2 writer clusters with tiered instance types +- **`kuzu-writers.yaml`**: Auto-scaling EC2 writer clusters with tiered instance types (Kuzu backend) - **`kuzu-shared-replicas.yaml`**: ECS Fargate read replicas for shared repositories (SEC) +- **`neo4j-*.yaml`**: Neo4j-specific infrastructure templates (when using Neo4j backend) #### Observability @@ -385,9 +396,10 @@ Each major system component has detailed documentation: - **`/robosystems/models/api/README.md`**: Centralized Pydantic models for API validation - **`/robosystems/config/README.md`**: Configuration management and environment handling -### Kuzu Graph Database System +### Graph Database System -- **`/robosystems/graph_api/README.md`**: Complete Kuzu API documentation +- **`/robosystems/graph_api/README.md`**: Complete Graph API documentation (supports Kuzu and Neo4j backends) +- **`/robosystems/graph_api/backends/README.md`**: Backend abstraction layer and implementation details - **`/robosystems/graph_api/client/README.md`**: Client-factory system for intelligent routing ### Middleware Components diff --git a/robosystems/graph_api/README.md b/robosystems/graph_api/README.md index 71917c63..1846c929 100644 --- a/robosystems/graph_api/README.md +++ b/robosystems/graph_api/README.md @@ -1,6 +1,11 @@ -# Kuzu API +# Graph API -High-performance HTTP API server for Kuzu graph database cluster management. FastAPI-based microservice that runs alongside Kuzu databases on EC2 instances, providing REST endpoints for multi-tenant graph operations with enterprise-grade reliability and security. +High-performance HTTP API server for graph database cluster management with pluggable backend support. FastAPI-based microservice that provides REST endpoints for multi-tenant graph operations with enterprise-grade reliability and security. + +**Supported Backends:** +- **Kuzu** (Default): High-performance embedded graph database, ideal for Standard tier deployments +- **Neo4j Community**: Client-server architecture for Professional/Enterprise tiers with advanced features +- **Neo4j Enterprise**: Full enterprise features including multi-database support for Premium tier ## Table of Contents @@ -27,14 +32,17 @@ High-performance HTTP API server for Kuzu graph database cluster management. Fas │ GraphRouter Layer │ │ (Intelligent Routing Logic) │ ├─────────────────────────────────────────────────────────────┤ -│ KuzuClientFactory Layer │ +│ GraphClientFactory Layer │ │ (Circuit Breakers, Retry Logic) │ ├─────────────────────────────────────────────────────────────┤ -│ Kuzu API Layer │ -│ (FastAPI on Port 8001) │ +│ Graph API Layer │ +│ (FastAPI on Port 8001/8002 depending on backend) │ +├─────────────────────────────────────────────────────────────┤ +│ Backend Abstraction Layer │ +│ (Pluggable: Kuzu, Neo4j Community/Enterprise) │ ├─────────────────────────────────────────────────────────────┤ -│ Kuzu Database Engine │ -│ (Native Graph Database) │ +│ Graph Database Engine │ +│ (Kuzu Embedded or Neo4j Bolt) │ └─────────────────────────────────────────────────────────────┘ ``` @@ -46,6 +54,12 @@ graph_api/ ├── main.py # Server entry point ├── __main__.py # Module entry point │ +├── backends/ # Backend implementations +│ ├── base.py # Abstract backend interface +│ ├── kuzu.py # Kuzu backend implementation +│ ├── neo4j.py # Neo4j backend implementation +│ └── __init__.py # Backend factory +│ ├── client/ # Python clients │ ├── client.py # Async client implementation │ ├── sync_client.py # Synchronous client @@ -56,7 +70,7 @@ graph_api/ ├── core/ # Core services │ ├── cluster_manager.py # Cluster orchestration │ ├── database_manager.py # Database lifecycle management -│ ├── connection_pool.py # Connection pooling (max 3/DB) +│ ├── connection_pool.py # Connection pooling │ ├── admission_control.py # Backpressure management │ └── metrics_collector.py # Performance metrics │ @@ -74,7 +88,7 @@ graph_api/ │ └── tasks.py # Background task tracking │ ├── middleware/ -│ ├── auth.py # API key authentication +│ ├── auth.py # API key authentication (backend-agnostic) │ └── request_limits.py # Rate limiting │ └── models/ # Pydantic models @@ -84,17 +98,30 @@ graph_api/ └── cluster.py # Cluster configuration ``` -### Node Types +### Node Types (Kuzu Backend) + +When using the Kuzu backend, the system deploys different node types: - **Writer Nodes** (`writer`): Entity database read/write operations on EC2 - **Shared Master** (`shared_master`): Repository ingestion and writes on EC2 - **Shared Replica** (`shared_replica`): Read-only replicas on EC2 with ALB +### Backend Selection (Neo4j) + +When using Neo4j backend: + +- **Community Edition**: Single database, suitable for development and Professional/Enterprise tiers +- **Enterprise Edition**: Multi-database support for Premium tier with advanced features + ## Deployment Infrastructure ### CloudFormation Stack Architecture -The Kuzu API is deployed through a sophisticated multi-stack CloudFormation architecture: +The Graph API deployment architecture varies by backend: + +#### Kuzu Backend Infrastructure + +For Kuzu deployments, the system uses a sophisticated multi-stack CloudFormation architecture: ``` 1. Infrastructure Stack (kuzu-infra.yaml) @@ -121,6 +148,23 @@ The Kuzu API is deployed through a sophisticated multi-stack CloudFormation arch └─ Read-only EC2 instances (r7g.medium) ``` +#### Neo4j Backend Infrastructure + +For Neo4j deployments (future implementation): + +``` +1. Neo4j Database Stack (neo4j-db.yaml) + ├─ ECS Fargate Service or EC2 instances + ├─ EBS volumes for data persistence + ├─ Application Load Balancer + └─ Multi-AZ deployment for Enterprise + +2. Neo4j API Stack (neo4j-api.yaml) + ├─ ECS Fargate Service (Graph API) + ├─ Bolt connection to Neo4j database + └─ Health checks and auto-scaling +``` + ### Infrastructure Tiers #### Production Environment @@ -236,20 +280,22 @@ cfn-signal --success --stack ${STACK_NAME} ... ```http POST /databases -Authorization: X-Kuzu-API-Key: {api_key} +Authorization: X-Graph-API-Key: {api_key} Content-Type: application/json { "graph_id": "kg1a2b3c4d5", "schema_type": "entity" // entity|shared|custom } + +Note: Both X-Graph-API-Key and X-Kuzu-API-Key headers are supported for backward compatibility ``` #### Execute Query ```http POST /databases/{graph_id}/query -Authorization: X-Kuzu-API-Key: {api_key} +Authorization: X-Graph-API-Key: {api_key} Content-Type: application/json { @@ -265,7 +311,7 @@ Content-Type: application/json ```http POST /databases/{graph_id}/copy -Authorization: X-Kuzu-API-Key: {api_key} +Authorization: X-Graph-API-Key: {api_key} Content-Type: application/json { @@ -278,13 +324,15 @@ Content-Type: application/json "region": "us-east-1" } } + +Note: S3 bulk copy is Kuzu-specific. Neo4j uses alternative data loading methods. ``` This returns a task ID that can be monitored via Server-Sent Events: ```http GET /tasks/{task_id}/monitor -Authorization: X-Kuzu-API-Key: {api_key} +Authorization: X-Graph-API-Key: {api_key} ``` ### System Operations @@ -401,23 +449,34 @@ client = await get_kuzu_client( ### Environment Variables ```bash -# Node Configuration +# Backend Configuration +BACKEND_TYPE=kuzu # kuzu|neo4j_community|neo4j_enterprise + +# Node Configuration (Kuzu Backend) KUZU_NODE_TYPE=writer # writer|shared_master|shared_replica WRITER_TIER=standard # standard|enterprise|premium|shared KUZU_DATABASE_PATH=/data/kuzu-dbs # Storage location -KUZU_PORT=8001 # API port +KUZU_PORT=8001 # API port (8001 for Kuzu, 8002 for Neo4j) + +# Neo4j Configuration (Neo4j Backend) +NEO4J_URI=bolt://neo4j-db:7687 # Neo4j Bolt connection +NEO4J_USERNAME=neo4j # Neo4j username +NEO4J_PASSWORD= # Retrieved from Secrets Manager +NEO4J_ENTERPRISE=false # Enable multi-database support +GRAPH_API_PORT=8002 # API port for Neo4j backend # Performance Settings -KUZU_MAX_DATABASES_PER_NODE=10 # Tier-specific limit -KUZU_MAX_MEMORY_MB=14336 # Total memory allocation -KUZU_MEMORY_PER_DB_MB=2048 # Per-database memory +KUZU_MAX_DATABASES_PER_NODE=10 # Tier-specific limit (Kuzu) +KUZU_MAX_MEMORY_MB=14336 # Total memory allocation (Kuzu) +KUZU_MEMORY_PER_DB_MB=2048 # Per-database memory (Kuzu) KUZU_CHUNK_SIZE=1000 # Streaming chunk size KUZU_QUERY_TIMEOUT=30 # Query timeout seconds KUZU_MAX_QUERY_LENGTH=10000 # Max query characters KUZU_CONNECTION_POOL_SIZE=10 # Connections per database +NEO4J_MAX_CONNECTION_POOL_SIZE=50 # Neo4j connection pool size # Authentication -KUZU_API_KEY=kuzu_prod_... # Unified API key +KUZU_API_KEY= # Unified API key (both backends) # AWS Configuration AWS_DEFAULT_REGION=us-east-1 @@ -429,8 +488,8 @@ KUZU_CIRCUIT_BREAKERS_ENABLED=true # Enable circuit breakers KUZU_REDIS_CACHE_ENABLED=true # Enable Redis caching KUZU_RETRY_LOGIC_ENABLED=true # Enable automatic retries KUZU_HEALTH_CHECKS_ENABLED=true # Enable health checking -SHARED_REPLICA_ALB_ENABLED=false # Enable replica ALB routing -ALLOW_SHARED_MASTER_READS=true # Allow reads from master +SHARED_REPLICA_ALB_ENABLED=false # Enable replica ALB routing (Kuzu) +ALLOW_SHARED_MASTER_READS=true # Allow reads from master (Kuzu) ``` ### Schema Types @@ -445,6 +504,12 @@ ALLOW_SHARED_MASTER_READS=true # Allow reads from master All API requests require authentication via API key header: +```http +X-Graph-API-Key: graph_api_64_character_random_string +``` + +Or the legacy header (still supported): + ```http X-Kuzu-API-Key: kuzu_prod_64_character_random_string ``` @@ -713,12 +778,21 @@ aws autoscaling start-instance-refresh \ ## Known Limitations +### Kuzu Backend + 1. **Sequential Ingestion**: Files processed one at a time per database (Kuzu constraint) 2. **Connection Limit**: Maximum 3 concurrent connections per database 3. **Single Writer**: Only one write operation per database at a time 4. **No Cross-Database Queries**: Each query scoped to single database 5. **Volume Attachment**: One EBS volume per database (no striping) +### Neo4j Backend + +1. **Multi-Database**: Only available in Enterprise edition +2. **Connection Pooling**: Limited by Neo4j configuration +3. **Bolt Protocol**: Requires network connectivity to Neo4j instance +4. **Licensing**: Enterprise features require Neo4j Enterprise license + ## Contributing 1. Follow existing patterns in codebase diff --git a/robosystems/graph_api/backends/README.md b/robosystems/graph_api/backends/README.md new file mode 100644 index 00000000..228340cf --- /dev/null +++ b/robosystems/graph_api/backends/README.md @@ -0,0 +1,455 @@ +# Graph Database Backend Abstraction Layer + +This module provides a pluggable backend architecture for RoboSystems' graph database system, supporting multiple graph database technologies through a unified interface. + +## Overview + +The backend abstraction layer allows RoboSystems to support different graph database technologies while maintaining a consistent API and application logic. This enables: + +- **Technology Flexibility**: Switch between graph databases based on requirements +- **Multi-Tier Support**: Different backends for different subscription tiers +- **Future-Proofing**: Easy integration of new graph database technologies +- **Development Simplification**: Test with lightweight backends, deploy with production-grade systems + +## Supported Backends + +### Kuzu (Default) + +**Type**: Embedded graph database +**Best For**: Standard tier, high-performance single-instance deployments +**Status**: Production-ready + +**Key Features:** +- High-performance embedded database +- Low latency for local access +- Columnar storage for analytics +- COPY operations for bulk data loading +- Direct file system access + +**Infrastructure:** +- EC2-based writer instances +- DynamoDB registry for allocation +- EBS volumes for persistence +- Auto-scaling groups by tier + +### Neo4j Community + +**Type**: Client-server graph database +**Best For**: Professional/Enterprise tiers, multi-user access +**Status**: Development/Testing + +**Key Features:** +- Battle-tested graph database +- Bolt protocol for client-server communication +- APOC procedures for extended functionality +- Cypher query language +- Built-in visualization tools + +**Limitations:** +- Single database per instance +- No multi-database support +- Community license restrictions + +### Neo4j Enterprise + +**Type**: Client-server graph database with enterprise features +**Best For**: Premium tier, enterprise customers +**Status**: Future implementation + +**Key Features:** +- Multi-database support +- Advanced security features +- Clustering and replication +- Enterprise support +- Role-based access control + +**Requirements:** +- Neo4j Enterprise license +- Additional infrastructure costs + +## Architecture + +### Backend Interface + +All backends implement the abstract `GraphBackend` interface defined in `base.py`: + +```python +class GraphBackend(ABC): + """Abstract base class for graph database backends.""" + + @abstractmethod + async def execute_query( + self, + graph_id: str, + cypher: str, + parameters: Optional[Dict[str, Any]] = None, + database: Optional[str] = None, + ) -> List[Dict[str, Any]]: + """Execute a read query.""" + pass + + @abstractmethod + async def execute_write( + self, + graph_id: str, + cypher: str, + parameters: Optional[Dict[str, Any]] = None, + database: Optional[str] = None, + ) -> List[Dict[str, Any]]: + """Execute a write query.""" + pass + + @abstractmethod + async def create_database(self, database_name: str) -> bool: + """Create a new database.""" + pass + + @abstractmethod + async def delete_database(self, database_name: str) -> bool: + """Delete a database.""" + pass + + @abstractmethod + async def list_databases(self) -> List[str]: + """List all databases.""" + pass + + @abstractmethod + async def get_database_info(self, database_name: str) -> DatabaseInfo: + """Get database metadata.""" + pass + + @abstractmethod + async def get_cluster_topology(self) -> ClusterTopology: + """Get cluster topology information.""" + pass + + @abstractmethod + async def health_check(self) -> bool: + """Check backend health.""" + pass + + @abstractmethod + async def close(self) -> None: + """Close connections and cleanup.""" + pass +``` + +### Backend Factory + +The backend factory in `__init__.py` provides singleton access to the configured backend: + +```python +from robosystems.graph_api.backends import get_backend + +# Get the configured backend (singleton) +backend = get_backend() + +# Use the backend +results = await backend.execute_query( + graph_id="kg1a2b3c4d5", + cypher="MATCH (n:Entity) RETURN n LIMIT 10" +) +``` + +Backend selection is controlled by the `BACKEND_TYPE` environment variable: +- `kuzu` (default) +- `neo4j_community` +- `neo4j_enterprise` + +## Implementation Details + +### Kuzu Backend (`kuzu.py`) + +**Connection Management:** +- Direct file system access to database directories +- Connection pool per database (max 3 connections) +- Memory-mapped file access for performance + +**Database Operations:** +- Each database stored in separate directory +- COPY operations for bulk loading from S3 +- Streaming support via NDJSON + +**Limitations:** +- Single writer per database +- Sequential file processing +- No network-based queries + +### Neo4j Backend (`neo4j.py`) + +**Connection Management:** +- Async Neo4j Python driver +- Connection pooling via driver configuration +- Bolt protocol for all communication + +**Database Operations:** +- Community: Single `neo4j` database +- Enterprise: Multi-database with `kg_{graph_id}_main` naming +- Transaction support via driver + +**Query Routing:** +- Automatic routing by Neo4j driver (cluster mode) +- Session-based query execution +- No manual read/write detection needed + +## Configuration + +### Environment Variables + +```bash +# Backend Selection +BACKEND_TYPE=kuzu # kuzu|neo4j_community|neo4j_enterprise + +# Kuzu Configuration +KUZU_DATABASE_PATH=/data/kuzu-dbs # Database storage location +KUZU_CONNECTION_POOL_SIZE=3 # Connections per database +KUZU_QUERY_TIMEOUT=30 # Query timeout (seconds) + +# Neo4j Configuration +NEO4J_URI=bolt://neo4j-db:7687 # Bolt connection URI +NEO4J_USERNAME=neo4j # Neo4j username +NEO4J_PASSWORD= # Retrieved from Secrets Manager +NEO4J_ENTERPRISE=false # Enable multi-database support +NEO4J_MAX_CONNECTION_POOL_SIZE=50 # Connection pool size +NEO4J_CONNECTION_ACQUISITION_TIMEOUT=60 # Timeout for acquiring connection +NEO4J_MAX_CONNECTION_LIFETIME=3600 # Max connection lifetime +``` + +### Docker Configuration + +Development environment supports both backends via profiles: + +```bash +# Start with Kuzu (default) +docker-compose --profile kuzu up + +# Start with Neo4j +docker-compose --profile neo4j up + +# Both +docker-compose --profile kuzu --profile neo4j up +``` + +## Usage Patterns + +### Basic Query Execution + +```python +from robosystems.graph_api.backends import get_backend + +# Get backend instance +backend = get_backend() + +# Execute read query +results = await backend.execute_query( + graph_id="kg1a2b3c4d5", + cypher="MATCH (e:Entity) RETURN e.name as name", + parameters={} +) + +# Execute write query +await backend.execute_write( + graph_id="kg1a2b3c4d5", + cypher="CREATE (e:Entity {identifier: $id, name: $name})", + parameters={"id": "entity-123", "name": "New Corp"} +) +``` + +### Database Management + +```python +# Create database +await backend.create_database("kg1a2b3c4d5") + +# List databases +databases = await backend.list_databases() + +# Get database info +info = await backend.get_database_info("kg1a2b3c4d5") +print(f"Nodes: {info.node_count}, Relationships: {info.relationship_count}") + +# Delete database +await backend.delete_database("kg1a2b3c4d5") +``` + +### Health Checks + +```python +# Check backend health +is_healthy = await backend.health_check() + +if not is_healthy: + logger.error("Backend unhealthy!") + +# Get cluster topology +topology = await backend.get_cluster_topology() +print(f"Mode: {topology.mode}") +if topology.mode == "cluster": + print(f"Leader: {topology.leader}") + print(f"Followers: {len(topology.followers)}") +``` + +## Testing + +### Unit Tests + +```bash +# Test backend factory +pytest tests/unit/graph_api/backends/test_backend_factory.py + +# Test Kuzu backend +pytest tests/unit/graph_api/backends/test_kuzu_backend.py + +# Test Neo4j backend +pytest tests/unit/graph_api/backends/test_neo4j_backend.py +``` + +### Integration Tests + +```bash +# Requires running backend instances +pytest tests/integration/graph_api/ -m backend_integration + +# Test specific backend +BACKEND_TYPE=kuzu pytest tests/integration/graph_api/ +BACKEND_TYPE=neo4j_community pytest tests/integration/graph_api/ +``` + +## Performance Characteristics + +### Kuzu Backend + +**Strengths:** +- Extremely fast local queries (microsecond latency) +- Efficient bulk loading via COPY +- Low memory footprint for small databases +- Columnar storage for analytics + +**Limitations:** +- Single writer constraint +- File system I/O bound for large graphs +- No distributed queries +- Limited concurrent connections (max 3) + +### Neo4j Backend + +**Strengths:** +- Proven enterprise scalability +- Rich ecosystem (APOC, GDS, etc.) +- Advanced query optimization +- Clustering support (Enterprise) + +**Limitations:** +- Network latency for all operations +- Higher memory requirements +- Licensing costs (Enterprise) +- Connection pool management complexity + +## Migration Guide + +### Kuzu to Neo4j + +To migrate a graph from Kuzu to Neo4j: + +1. Export data from Kuzu: +```python +# Export to CSV/Parquet +await kuzu_backend.export_database( + database_name="kg1a2b3c4d5", + output_path="/tmp/export" +) +``` + +2. Load into Neo4j: +```cypher +// Use LOAD CSV or neo4j-admin import +LOAD CSV WITH HEADERS FROM 'file:///entities.csv' AS row +CREATE (e:Entity { + identifier: row.identifier, + name: row.name +}) +``` + +3. Update configuration: +```bash +BACKEND_TYPE=neo4j_community +``` + +### Neo4j to Kuzu + +Reverse process using Neo4j export and Kuzu COPY operations. + +## Troubleshooting + +### Common Issues + +#### Backend Factory Returns Wrong Backend + +**Symptom**: Getting KuzuBackend when expecting Neo4jBackend + +**Solution**: +- Verify `BACKEND_TYPE` environment variable +- Clear backend singleton: `_backend_instance = None` +- Restart application + +#### Connection Failures + +**Kuzu**: +- Check database path exists and is writable +- Verify no file system corruption +- Check available disk space + +**Neo4j**: +- Verify Neo4j service is running +- Check Bolt port (7687) is accessible +- Validate credentials in Secrets Manager + +#### Performance Issues + +**Kuzu**: +- Monitor I/O wait times +- Check EBS volume performance +- Consider SSD-backed storage + +**Neo4j**: +- Review connection pool settings +- Monitor network latency +- Check Neo4j heap memory allocation + +## Future Enhancements + +1. **Additional Backends**: + - Amazon Neptune + - ArangoDB + - TigerGraph + - Memgraph + +2. **Features**: + - Backend-specific query optimization + - Automatic backend selection based on workload + - Cross-backend data synchronization + - Performance benchmarking framework + +3. **Infrastructure**: + - Multi-region backend support + - Automatic failover between backends + - Backend-specific monitoring dashboards + +## Contributing + +When adding a new backend: + +1. Implement the `GraphBackend` interface +2. Add configuration in `env.py` +3. Update backend factory in `__init__.py` +4. Add comprehensive tests +5. Update documentation +6. Add Docker Compose profile + +## Support + +For backend-specific issues: +- **Kuzu**: Check `/robosystems/graph_api/README.md` +- **Neo4j**: Consult Neo4j documentation at https://neo4j.com/docs/ +- **General**: Review this README and abstract interface diff --git a/robosystems/graph_api/client/README.md b/robosystems/graph_api/client/README.md index 1cc4b4b3..5f3ed1b9 100644 --- a/robosystems/graph_api/client/README.md +++ b/robosystems/graph_api/client/README.md @@ -1,8 +1,12 @@ -# Kuzu API Client & Factory +# Graph API Client & Factory ## Overview -The Kuzu Client and Factory system provides the critical interface between the RoboSystems application (API and workers) and the Kuzu API running on EC2 instances. This layer handles intelligent routing, connection pooling, circuit breaking, and automatic failover to ensure reliable graph database operations at scale. +The Graph Client and Factory system provides the critical interface between the RoboSystems application (API and workers) and the Graph API running on infrastructure. This layer handles intelligent routing, connection pooling, circuit breaking, and automatic failover to ensure reliable graph database operations at scale. + +**Backend Support:** +- **Kuzu**: EC2-based instances with DynamoDB registry discovery +- **Neo4j**: Direct Bolt connection or service discovery ## Architecture @@ -11,26 +15,29 @@ The Kuzu Client and Factory system provides the critical interface between the R │ Application Layer │ │ (FastAPI Routes / Celery Workers) │ ├─────────────────────────────────────────────────────────────────┤ -│ KuzuClientFactory │ +│ GraphClientFactory │ │ (Intelligent Routing & Discovery) │ ├─────────────────────────────────────────────────────────────────┤ -│ KuzuClient │ +│ GraphClient │ │ (Async/Sync HTTP Client with Retry Logic) │ ├─────────────────────────────────────────────────────────────────┤ -│ EC2 Infrastructure │ -│ ┌──────────┐ ┌──────────┐ ┌──────────┐ ┌──────────┐ │ -│ │ Standard │ │Enterprise│ │ Premium │ │ Shared │ │ -│ │ Writers │ │ Writers │ │ Writers │ │ Master │ │ -│ └──────────┘ └──────────┘ └──────────┘ └──────────┘ │ -│ ┌──────────────┐ │ -│ │ Replica ALB │ │ -│ └──────────────┘ │ +│ Backend-Specific Infrastructure │ +│ │ +│ Kuzu Backend: │ Neo4j Backend: │ +│ ┌──────────────────┐ │ ┌──────────────────┐ │ +│ │ EC2 Instances │ │ │ Neo4j Database │ │ +│ │ - Standard │ │ │ - Community │ │ +│ │ - Enterprise │ │ │ - Enterprise │ │ +│ │ - Premium │ │ │ (Bolt Protocol) │ │ +│ │ - Shared Master │ │ └──────────────────┘ │ +│ │ - Replica ALB │ │ │ +│ └──────────────────┘ │ │ └─────────────────────────────────────────────────────────────────┘ ``` ## Key Components -### 1. KuzuClientFactory (`factory.py`) +### 1. GraphClientFactory (`factory.py`) The factory is responsible for intelligent routing decisions based on: @@ -38,9 +45,12 @@ The factory is responsible for intelligent routing decisions based on: - **Operation Type**: Read vs Write operations - **Environment**: Development, Staging, Production - **Tier**: Standard, Enterprise, Premium for user graphs +- **Backend Type**: Kuzu or Neo4j #### Routing Logic +##### Kuzu Backend + ```python # Shared Repositories (SEC, industry, economic) ├── Write Operations → Shared Master (always) @@ -55,18 +65,29 @@ The factory is responsible for intelligent routing decisions based on: └── Route to appropriate tier writer instance ``` +##### Neo4j Backend + +```python +# All Operations (User and Shared) +├── Development → Local Neo4j instance (bolt://neo4j-db:7687) +└── Production/Staging → Neo4j cluster or single instance + ├── Community → Single database endpoint + └── Enterprise → Multi-database support +``` + #### Key Features -- **Dynamic Discovery**: Automatically discovers EC2 instances from DynamoDB +- **Dynamic Discovery**: Automatically discovers instances (DynamoDB for Kuzu, service discovery for Neo4j) +- **Backend Abstraction**: Consistent interface regardless of backend - **Circuit Breakers**: Prevents cascading failures with configurable thresholds -- **Connection Pooling**: Reuses HTTP/2 connections for efficiency -- **Redis Caching**: Caches instance locations to reduce DynamoDB lookups +- **Connection Pooling**: Reuses connections for efficiency (HTTP/2 for Kuzu, Bolt for Neo4j) +- **Redis Caching**: Caches instance locations to reduce lookups - **Automatic Failover**: Falls back to alternative endpoints when primary unavailable - **Retry Logic**: Exponential backoff with jitter for transient errors -### 2. KuzuClient (`client.py`) +### 2. GraphClient (`client.py`) -Asynchronous HTTP client for interacting with Kuzu API endpoints. +Asynchronous HTTP client for interacting with Graph API endpoints (backend-agnostic). #### Core Operations @@ -210,18 +231,25 @@ except ServiceUnavailableError as e: ### Core Configuration ```bash -# API Endpoints +# Backend Selection +BACKEND_TYPE=kuzu # kuzu|neo4j_community|neo4j_enterprise + +# API Endpoints (Kuzu Backend) KUZU_API_URL=http://localhost:8001 # Default API URL (dev/fallback) KUZU_REPLICA_ALB_URL=http://alb.internal # Replica ALB endpoint KUZU_API_KEY=kuzu_prod_64chars... # Authentication key +# API Endpoints (Neo4j Backend) +NEO4J_URI=bolt://neo4j-db:7687 # Neo4j Bolt connection +GRAPH_API_PORT=8002 # Graph API port for Neo4j + # Feature Flags KUZU_RETRY_LOGIC_ENABLED=true # Enable automatic retries KUZU_CIRCUIT_BREAKERS_ENABLED=true # Enable circuit breakers KUZU_HEALTH_CHECKS_ENABLED=true # Enable health checking KUZU_REDIS_CACHE_ENABLED=true # Enable Redis caching -SHARED_REPLICA_ALB_ENABLED=true # Enable replica ALB routing -ALLOW_SHARED_MASTER_READS=true # Allow reads from master +SHARED_REPLICA_ALB_ENABLED=true # Enable replica ALB routing (Kuzu) +ALLOW_SHARED_MASTER_READS=true # Allow reads from master (Kuzu) # Performance Tuning KUZU_CLIENT_TIMEOUT=30 # Request timeout (seconds) @@ -229,18 +257,23 @@ KUZU_CLIENT_MAX_RETRIES=3 # Maximum retry attempts KUZU_CIRCUIT_BREAKER_THRESHOLD=5 # Failures before opening KUZU_CIRCUIT_BREAKER_TIMEOUT=60 # Seconds before reset KUZU_CACHE_TTL=300 # Cache TTL (seconds) +NEO4J_MAX_CONNECTION_POOL_SIZE=50 # Neo4j connection pool size ``` -### DynamoDB Configuration +### DynamoDB Configuration (Kuzu Backend) ```bash -# For instance discovery +# For instance discovery (Kuzu only) KUZU_INSTANCE_REGISTRY_TABLE=robosystems-kuzu-{env}-instance-registry KUZU_GRAPH_REGISTRY_TABLE=robosystems-kuzu-{env}-graph-registry + +# Note: Neo4j backend uses direct connection URIs instead of DynamoDB discovery ``` ## Instance Discovery Flow +### Kuzu Backend + 1. **Check Cache**: Redis cache with 5-minute TTL 2. **Query DynamoDB**: Find instance hosting the graph 3. **Health Check**: Verify instance is healthy @@ -249,7 +282,7 @@ KUZU_GRAPH_REGISTRY_TABLE=robosystems-kuzu-{env}-graph-registry ```python # Internal discovery flow (handled automatically) -1. KuzuClientFactory.create_client("kg1a2b3c4d5") +1. GraphClientFactory.create_client("kg1a2b3c4d5") 2. → Check Redis: kuzu:prod:location:kg1a2b3c4d5 3. → Query DynamoDB: GraphRegistry[graph_id=kg1a2b3c4d5] 4. → Get instance: i-1234567890 at 10.0.1.100 @@ -257,6 +290,22 @@ KUZU_GRAPH_REGISTRY_TABLE=robosystems-kuzu-{env}-graph-registry 6. → Cache location for 300 seconds ``` +### Neo4j Backend + +1. **Direct Connection**: Uses configured NEO4J_URI +2. **Database Selection**: Routes to correct database (Community: single, Enterprise: multi) +3. **Connection Pool**: Manages Bolt protocol connections +4. **Health Check**: Verifies Neo4j database availability + +```python +# Internal connection flow (handled automatically) +1. GraphClientFactory.create_client("kg1a2b3c4d5") +2. → Use NEO4J_URI from configuration +3. → Select database: kg_kg1a2b3c4d5_main (Enterprise) or neo4j (Community) +4. → Create Bolt connection via Graph API +5. → Pool connections for efficiency +``` + ## Circuit Breaker Pattern The circuit breaker prevents cascading failures: diff --git a/robosystems/middleware/graph/README.md b/robosystems/middleware/graph/README.md index 384c0a41..b310111e 100644 --- a/robosystems/middleware/graph/README.md +++ b/robosystems/middleware/graph/README.md @@ -1,17 +1,23 @@ # Graph Middleware -This middleware layer provides the core graph database abstraction and routing logic for the RoboSystems platform. +This middleware layer provides the core graph database abstraction and routing logic for the RoboSystems platform with support for multiple backend types. ## Overview The graph middleware: -- Routes graph operations to appropriate Kuzu clusters +- Routes graph operations to appropriate clusters (Kuzu or Neo4j) +- Provides backend-agnostic database abstraction - Manages database connections and pooling - Handles query execution with caching and queuing - Provides admission control and backpressure management - Integrates with the credit system for usage tracking +**Supported Backends:** +- **Kuzu**: Embedded graph database with EC2-based clusters +- **Neo4j Community**: Client-server architecture for Professional/Enterprise tiers +- **Neo4j Enterprise**: Multi-database support for Premium tier + ## Architecture ``` @@ -96,7 +102,7 @@ class ClusterConfig: ### 3. Graph Engine (`engine.py`) -Direct Kuzu database access with connection management. +Direct graph database access with connection management (Kuzu-specific). **Features:** @@ -108,9 +114,12 @@ Direct Kuzu database access with connection management. **Usage:** ```python +# Kuzu-specific direct access (development/legacy) engine = GraphEngine(database_path="/data/kuzu-dbs/kg1a2b3c") result = engine.execute_query("MATCH (c:Entity) RETURN c") engine.close() + +# Note: For production use the backend abstraction layer instead ``` ### 4. Repository Pattern (`repository.py`) @@ -222,7 +231,7 @@ Core type definitions and enums. ### 9. Database Allocation (`allocation_manager.py`) -DynamoDB-based allocation manager for Kuzu databases across writer instances. +DynamoDB-based allocation manager for graph databases across instances (Kuzu-specific). **Features:** @@ -279,10 +288,17 @@ routing = MultiTenantUtils.get_graph_routing("kg1a2b3c") Key environment variables: ```bash -# Routing Configuration +# Backend Configuration +BACKEND_TYPE=kuzu # kuzu|neo4j_community|neo4j_enterprise + +# Routing Configuration (Kuzu) KUZU_ACCESS_PATTERN=api_writer # Access pattern (api_writer/api_reader/direct_file) KUZU_API_URL= # Localhost endpoint for routing (dynamic lookup in prod) +# Routing Configuration (Neo4j) +NEO4J_URI=bolt://neo4j-db:7687 # Neo4j Bolt connection +NEO4J_ENTERPRISE=false # Enable multi-database support + # Queue Configuration QUERY_QUEUE_MAX_SIZE=1000 # Maximum queries in queue QUERY_QUEUE_MAX_CONCURRENT=50 # Max concurrent executions @@ -296,8 +312,9 @@ LOAD_SHEDDING_ENABLED=true # Enable load shedding # Performance CONNECTION_POOL_SIZE=10 # Database connection pool size QUERY_TIMEOUT=30 # Query timeout (seconds) +NEO4J_MAX_CONNECTION_POOL_SIZE=50 # Neo4j connection pool size -# Database Allocation +# Database Allocation (Kuzu) KUZU_MAX_DATABASES_PER_NODE=50 # Max databases per Kuzu instance # Multi-tenant Configuration