@wearesage/schema/adapters/neo4j.ts

A flexible schema definition and validation system for TypeScript with multi-database support

import { Type } from '../core/types';
import { DatabaseAdapter } from './interface';
import { MetadataRegistry } from '../core/MetadataRegistry';
import { SchemaReflector } from '../core/SchemaReflector';
import { RelationshipOptions } from '../core/types';
import neo4j, { Driver, Session, Transaction, Result, Record, Node, Relationship, Path, SessionMode } from 'neo4j-driver';


/**
 * Neo4j-specific relationship type decorator
 */
export function RelationshipType(type: string) {
  return function(target: any, propertyKey: string) {
    Reflect.defineMetadata('neo4j:relationshipType', type, target, propertyKey);
  };
}

/**
 * Shorthand decorator for Neo4j entities
 */
export function Neo4j<T extends { new (...args: any[]): any }>(target: T): T {
  return DatabaseAdapter('Neo4j')(target);
}

// Use real Neo4j driver types
type Neo4jNode = Node;
type Neo4jRelationship = Relationship;
type Neo4jPath = Path;
type Neo4jRecord = Record;
type Neo4jResult = Result;
type Neo4jTransaction = Transaction;
type Neo4jSession = Session;
type Neo4jDriver = Driver;

/**
 * Interface for Neo4j auth configuration
 */
export interface Neo4jAuthConfig {
  username: string;
  password: string;
}

/**
 * Implementation of DatabaseAdapter for Neo4j
 */
export class Neo4jAdapter implements DatabaseAdapter {
  readonly type = 'Neo4j';
  private registry: MetadataRegistry;
  private reflector: SchemaReflector;
  private driver!: Neo4jDriver;
  private initialized = false;
  
  /**
   * Constructor for Neo4j adapter
   * @param registry The metadata registry to use for schema information
   * @param connectionString The Neo4j connection string
   * @param auth The Neo4j authentication (username/password)
   */
  constructor(
    registry: MetadataRegistry, 
    private connectionString: string,
    private auth: Neo4jAuthConfig = { username: 'neo4j', password: 'password' }
  ) {
    this.registry = registry;
    this.reflector = new SchemaReflector(registry);
  }
  
  /**
   * Initialize the Neo4j connection
   */
  private async initialize(): Promise<void> {
    if (this.initialized) return;
    
    try {
      // Create the Neo4j driver with real neo4j-driver
      this.driver = neo4j.driver(
        this.connectionString,
        neo4j.auth.basic(this.auth.username, this.auth.password),
        {
          // Optional configuration
          maxConnectionLifetime: 3 * 60 * 60 * 1000, // 3 hours
          maxConnectionPoolSize: 50,
          connectionAcquisitionTimeout: 30 * 1000, // 30 seconds
          disableLosslessIntegers: true // Convert Neo4j integers to JavaScript numbers
        }
      );
      
      // Verify connectivity
      await this.driver.verifyConnectivity();
      this.initialized = true;
      console.log('[Neo4jAdapter] Successfully connected to Neo4j');
    } catch (error) {
      console.error('[Neo4jAdapter] Failed to connect to Neo4j:', error);
      throw error;
    }
  }
  
  /**
   * Get a Neo4j session with optional transaction mode
   * @param mode Session mode (READ or WRITE)
   * @returns A Neo4j session
   */
  private async getSession(mode: SessionMode = neo4j.session.WRITE): Promise<Neo4jSession> {
    await this.initialize();
    return this.driver.session({
      defaultAccessMode: mode,
      database: 'neo4j' // Use default database, can be configured if needed
    });
  }
  
  /**
   * Build a Cypher query to find an entity by its ID
   */
  private buildFindByIdQuery<T>(entityType: Type<T>): string {
    const labels = this.getEntityLabels(entityType);
    return `MATCH (n:${labels.join(':')} {id: $id}) RETURN n`;
  }
  
  /**
   * Resolve relationship target (function or string) to actual entity type
   */
  private resolveRelationshipTarget(target: Function | string): Type<any> {
    if (typeof target === 'string') {
      // Resolve string entity name through registry
      const registeredEntities = this.registry.getRegisteredEntities();
      const entityType = registeredEntities.find(e => e.name === target);
      if (!entityType) {
        throw new Error(`Entity '${target}' not found in registry. Available entities: ${registeredEntities.map(e => e.name).join(', ')}`);
      }
      return entityType;
    } else if (typeof target === 'function') {
      // Call function to get the entity type or string
      const result = target();
      if (typeof result === 'string') {
        return this.resolveRelationshipTarget(result);
      }
      return result;
    } else {
      throw new Error(`Invalid relationship target: ${target}`);
    }
  }

  /**
   * Get all Neo4j labels for an entity
   */
  private getEntityLabels<T>(entityType: Type<T>): string[] {
    // Get custom labels if defined, otherwise use the entity name
    const customLabels = Reflect.getMetadata('entity:labels', entityType);
    if (customLabels && Array.isArray(customLabels)) {
      return customLabels;
    }
    
    return [entityType.name];
  }
  
  /**
   * Get relationship type for a property
   */
  private getRelationshipType(entityType: Function, propertyKey: string): string {
    // Get custom relationship type if defined, otherwise use uppercase property name
    const customType = Reflect.getMetadata('neo4j:relationshipType', entityType.prototype, propertyKey);
    if (customType) {
      return customType;
    }
    
    // Look for relationship metadata
    const relationshipOptions = this.registry.getRelationshipMetadata(entityType, propertyKey);
    if (relationshipOptions?.name) {
      return relationshipOptions.name;
    }
    
    return propertyKey.toUpperCase();
  }
  
  /**
   * Convert Neo4j node to entity instance
   */
  private nodeToEntity<T>(entityType: Type<T>, node: any): T {
    const entity = new entityType();
    
    // Copy node properties to entity
    if (node && typeof node === 'object') {
      // Handle Neo4j Integer conversions
      const properties = node.properties ? { ...node.properties } : {};
      
      Object.keys(properties).forEach(key => {
        // Get the property value
        let value = properties[key];
        
        // Convert Neo4j date types to JavaScript types if needed
        if (neo4j.isDate && neo4j.isDate(value)) {
          value = value.toString(); // or value.toStandardDate() for Date object
        } else if (neo4j.isDateTime && neo4j.isDateTime(value)) {
          value = value.toString(); // or value.toStandardDate() for Date object
        } else if (neo4j.isLocalDateTime && neo4j.isLocalDateTime(value)) {
          value = value.toString();
        } else if (neo4j.isLocalTime && neo4j.isLocalTime(value)) {
          value = value.toString();
        } else if (neo4j.isTime && neo4j.isTime(value)) {
          value = value.toString();
        }
        
        // Assign the converted value to the entity
        (entity as any)[key] = value;
      });
    }
    
    return entity;
  }
  
  /**
   * Serialize value for Neo4j (only primitives allowed)
   */
  private serializeForNeo4j(value: any): any {
    // Handle null/undefined
    if (value === null || value === undefined) {
      return null;
    }
    
    // Handle primitives directly
    if (typeof value === 'string' || typeof value === 'number' || typeof value === 'boolean') {
      return value;
    }
    
    // Handle Date objects
    if (value instanceof Date) {
      return value.toISOString();
    }
    
    // Handle arrays
    if (Array.isArray(value)) {
      // For empty arrays, return empty array
      if (value.length === 0) {
        return [];
      }
      
      // For arrays with primitives, return as-is
      const firstItem = value[0];
      if (typeof firstItem === 'string' || typeof firstItem === 'number' || typeof firstItem === 'boolean') {
        return value;
      }
      
      // For complex arrays, serialize to JSON string
      return JSON.stringify(value);
    }
    
    // Handle empty objects
    if (typeof value === 'object' && Object.keys(value).length === 0) {
      return null; // Convert empty objects to null
    }
    
    // Handle complex objects - serialize to JSON string
    if (typeof value === 'object') {
      return JSON.stringify(value);
    }
    
    // Fallback: convert to string
    return String(value);
  }

  /**
   * Convert entity to Neo4j properties
   */
  private entityToProperties<T>(entity: T): { [key: string]: any } {
    const entityType = (entity as any).constructor as Type<T>;
    const entitySchema = this.reflector.getEntitySchema(entityType);
    const properties: { [key: string]: any } = {};
    
    // Convert each property with Neo4j serialization
    for (const [key, meta] of Object.entries(entitySchema.properties)) {
      if ((entity as any)[key] !== undefined) {
        const value = (entity as any)[key];
        const serializedValue = this.serializeForNeo4j(value);
        
        // Only include non-null values
        if (serializedValue !== null) {
          properties[key] = serializedValue;
        }
      }
    }
    
    return properties;
  }

  /**
   * Query for a single entity by criteria
   */
  async query<T>(entityType: Type<T>, criteria: object): Promise<T | null> {
    try {
      await this.initialize();
      const session = await this.getSession(neo4j.session.READ);
      
      try {
        const labels = this.getEntityLabels(entityType);
        
        // If criteria includes an ID, use the optimized findById query
        if ('id' in criteria) {
          const query = this.buildFindByIdQuery(entityType);
          const result = await session.run(query, { id: criteria.id });
          
          if (result.records.length === 0) {
            return null;
          }
          
          const node = result.records[0].get('n');
          return this.nodeToEntity(entityType, node);
        }
        
        // Otherwise, build a query based on criteria
        let whereClause = '';
        if (Object.keys(criteria).length > 0) {
          const conditions = Object.keys(criteria).map(key => `n.${key} = $${key}`).join(' AND ');
          whereClause = `WHERE ${conditions}`;
        }
        
        const query = `
          MATCH (n:${labels.join(':')})
          ${whereClause}
          RETURN n LIMIT 1
        `;
        
        const result = await session.run(query, criteria);
        
        if (result.records.length === 0) {
          return null;
        }
        
        const node = result.records[0].get('n');
        return this.nodeToEntity(entityType, node);
      } finally {
        await session.close();
      }
    } catch (error) {
      console.error('[Neo4jAdapter] Error in query:', error);
      throw error;
    }
  }
  
  /**
   * Query for multiple entities by criteria
   */
  async queryMany<T>(entityType: Type<T>, criteria: object): Promise<T[]> {
    try {
      await this.initialize();
      const session = await this.getSession(neo4j.session.READ);
      
      try {
        const labels = this.getEntityLabels(entityType);
        
        // Build the query based on criteria
        let whereClause = '';
        if (Object.keys(criteria).length > 0) {
          const conditions = Object.keys(criteria).map(key => `n.${key} = $${key}`).join(' AND ');
          whereClause = `WHERE ${conditions}`;
        }
        
        const query = `
          MATCH (n:${labels.join(':')})
          ${whereClause}
          RETURN n
        `;
        
        const result = await session.run(query, criteria);
        
        return result.records.map(record => {
          const node = record.get('n');
          return this.nodeToEntity(entityType, node);
        });
      } finally {
        await session.close();
      }
    } catch (error) {
      console.error('[Neo4jAdapter] Error in queryMany:', error);
      throw error;
    }
  }
  
  /**
   * Save an entity to Neo4j
   */
  async save<T>(entity: T): Promise<void> {
    try {
      await this.initialize();
      const session = await this.getSession(neo4j.session.WRITE);
      
      try {
        // Start a transaction for atomicity
        const tx = session.beginTransaction();
        
        try {
          const entityType = (entity as any).constructor as Type<T>;
          const entitySchema = this.reflector.getEntitySchema(entityType);
          const labels = this.getEntityLabels(entityType);
          
          // Validate the entity before saving
          const validation = this.reflector.validateEntity(entity);
          if (!validation.valid) {
            throw new Error(`Invalid entity: ${validation.errors.join(', ')}`);
          }
          
          // Extract properties
          const properties = this.entityToProperties(entity);
          
          // Create or update the node using individual property setting
          const propertySetters = Object.keys(properties)
            .filter(key => key !== 'id') // Don't set ID since it's in the MERGE
            .map(key => `n.${key} = $${key}`)
            .join(', ');
          
          const createNodeQuery = propertySetters.length > 0 
            ? `
            MERGE (n:${labels.join(':')} {id: $id})
            SET ${propertySetters}
            RETURN n
          `
            : `
            MERGE (n:${labels.join(':')} {id: $id})
            RETURN n
          `;
          
          const queryParams = { id: properties['id'], ...properties };
          
          console.log('🔍 Fixed query:', createNodeQuery);
          console.log('🔍 Fixed params:', JSON.stringify(queryParams, null, 2));
          
          await tx.run(createNodeQuery, queryParams);
          
          // Handle relationships
          for (const [key, relationshipMetaUntyped] of Object.entries(entitySchema.relationships || {})) {
            const relationshipMeta = relationshipMetaUntyped as RelationshipOptions;
            const relatedEntity = (entity as any)[key];
            if (!relatedEntity) continue;
            
            const relationshipType = this.getRelationshipType(entityType, key);
            
            // Handle different relationship cardinalities
            if (relationshipMeta.cardinality === 'one') {
              // Handle one-to-one or many-to-one relationship
              const relatedProperties = this.entityToProperties(relatedEntity);
              
              // Resolve the target entity type
              const targetType = this.resolveRelationshipTarget(relationshipMeta.target);
              const relatedLabels = this.getEntityLabels(targetType);
              
              // Create or update the related node first
              const relatedPropertySetters = Object.keys(relatedProperties)
                .filter(key => key !== 'id')
                .map(key => `m.${key} = $${key}`)
                .join(', ');
              
              const createRelatedNodeQuery = `
                MERGE (m:${relatedLabels.join(':')} {id: $id})
                SET ${relatedPropertySetters}
                RETURN m
              `;
              
              await tx.run(createRelatedNodeQuery, { id: relatedProperties['id'], ...relatedProperties });
              
              // Create the relationship
              const createRelationshipQuery = `
                MATCH (n:${labels.join(':')} {id: $nodeId})
                MATCH (m:${relatedLabels.join(':')} {id: $relatedId})
                MERGE (n)-[r:${relationshipType}]->(m)
                RETURN r
              `;
              
              await tx.run(createRelationshipQuery, {
                nodeId: properties['id'],
                relatedId: relatedProperties['id']
              });
            } else if (relationshipMeta.cardinality === 'many' && Array.isArray(relatedEntity)) {
              // Handle one-to-many or many-to-many relationship
              // First, delete existing relationships
              const deleteRelationshipsQuery = `
                MATCH (n:${labels.join(':')} {id: $nodeId})-[r:${relationshipType}]->()
                DELETE r
              `;
              
              await tx.run(deleteRelationshipsQuery, {
                nodeId: properties['id']
              });
              
              // Then create nodes and relationships for each related entity
              for (const related of relatedEntity) {
                const relatedProperties = this.entityToProperties(related);
                
                // Resolve the target entity type
                const targetType = this.resolveRelationshipTarget(relationshipMeta.target);
                const relatedLabels = this.getEntityLabels(targetType);
                
                // Create or update the related node
                const relatedPropertySetters = Object.keys(relatedProperties)
                  .filter(key => key !== 'id')
                  .map(key => `m.${key} = $${key}`)
                  .join(', ');
                
                const createRelatedNodeQuery = `
                  MERGE (m:${relatedLabels.join(':')} {id: $id})
                  SET ${relatedPropertySetters}
                  RETURN m
                `;
                
                await tx.run(createRelatedNodeQuery, { id: relatedProperties['id'], ...relatedProperties });
                
                // Create the relationship
                const createRelationshipQuery = `
                  MATCH (n:${labels.join(':')} {id: $nodeId})
                  MATCH (m:${relatedLabels.join(':')} {id: $relatedId})
                  MERGE (n)-[r:${relationshipType}]->(m)
                  RETURN r
                `;
                
                await tx.run(createRelationshipQuery, {
                  nodeId: properties['id'],
                  relatedId: relatedProperties['id']
                });
              }
            }
          }
          
          // Commit the transaction
          await tx.commit();
        } catch (error) {
          // Rollback transaction on error
          await tx.rollback();
          throw error;
        }
      } finally {
        await session.close();
      }
    } catch (error) {
      console.error('[Neo4jAdapter] Error in save:', error);
      throw error;
    }
  }
  
  /**
   * Delete an entity from Neo4j
   */
  async delete<T>(entityType: Type<T>, id: string | number): Promise<void> {
    try {
      await this.initialize();
      const session = await this.getSession(neo4j.session.WRITE);
      
      try {
        const labels = this.getEntityLabels(entityType);
        
        // Delete the node and all its relationships
        const query = `
          MATCH (n:${labels.join(':')} {id: $id})
          DETACH DELETE n
        `;
        
        await session.run(query, { id });
      } finally {
        await session.close();
      }
    } catch (error) {
      console.error('[Neo4jAdapter] Error in delete:', error);
      throw error;
    }
  }
  
  /**
   * Execute a raw Cypher query
   */
  async runNativeQuery<T = any>(query: string, params?: any): Promise<T> {
    try {
      await this.initialize();
      const session = await this.getSession();
      
      try {
        const result = await session.run(query, params);
        
        // Convert Neo4j result to a more usable format
        const convertedResult = {
          records: result.records.map(record => {
            const obj: { [key: string]: any } = {};
            
            for (let i = 0; i < record.keys.length; i++) {
              const key = record.keys[i];
              // Skip if key is a symbol
              if (typeof key === 'symbol') continue;
              const value = record.get(key);
              
              // Handle Neo4j Node objects
              if (value && neo4j.isNode && neo4j.isNode(value)) {
                obj[key] = {
                  ...(value.properties || {}),
                  _labels: value.labels,
                  _id: value.identity.toString()
                };
              } 
              // Handle Neo4j Relationship objects
              else if (value && neo4j.isRelationship && neo4j.isRelationship(value)) {
                // Check what properties are available for the relationship
                // Neo4j JS driver has changed property names in different versions
                const startId = 
                  (value.startNodeElementId && value.startNodeElementId.toString()) || 
                  (value.start && value.start.toString()) || 
                  "unknown";
                
                const endId = 
                  (value.endNodeElementId && value.endNodeElementId.toString()) || 
                  (value.end && value.end.toString()) || 
                  "unknown";
                
                obj[key] = {
                  ...(value.properties || {}),
                  _type: value.type,
                  _id: value.identity.toString(),
                  _startNodeId: startId,
                  _endNodeId: endId
                };
              }
              // Handle Neo4j Path objects
              else if (value && neo4j.isPath && neo4j.isPath(value)) {
                const segments = [];
                if (value.segments && Array.isArray(value.segments)) {
                  for (const segment of value.segments) {
                    segments.push({
                      start: {
                        ...(segment.start.properties || {}),
                        _labels: segment.start.labels,
                        _id: segment.start.identity.toString()
                      },
                      relationship: {
                        ...(segment.relationship.properties || {}),
                        _type: segment.relationship.type,
                        _id: segment.relationship.identity.toString()
                      },
                      end: {
                        ...(segment.end.properties || {}),
                        _labels: segment.end.labels,
                        _id: segment.end.identity.toString()
                      }
                    });
                  }
                }
                obj[key] = { segments };
              }
              // Handle other Neo4j types
              else if (value && typeof value === 'object' && value.constructor) {
                // Handle date/time types
                if ((neo4j.isDateTime && neo4j.isDateTime(value)) || 
                    (neo4j.isDate && neo4j.isDate(value)) || 
                    (neo4j.isTime && neo4j.isTime(value)) || 
                    (neo4j.isLocalDateTime && neo4j.isLocalDateTime(value)) || 
                    (neo4j.isLocalTime && neo4j.isLocalTime(value)) || 
                    (neo4j.isDuration && neo4j.isDuration(value))) {
                  obj[key] = value.toString();
                } else {
                  // For other objects, convert to plain object safely
                  obj[key] = typeof value === 'object' && value !== null ? Object.assign({}, value) : value;
                }
              } else {
                // For primitive values
                obj[key] = value;
              }
            }
            
            return obj;
          }),
          summary: {
            counters: result.summary.counters ? result.summary.counters.updates() : {},
            resultAvailableAfter: result.summary.resultAvailableAfter,
            resultConsumedAfter: result.summary.resultConsumedAfter,
            database: result.summary.database,
            query: {
              text: result.summary.query.text,
              parameters: result.summary.query.parameters
            }
          }
        };
        
        return convertedResult as unknown as T;
      } finally {
        await session.close();
      }
    } catch (error) {
      console.error('[Neo4jAdapter] Error in runNativeQuery:', error);
      throw error;
    }
  }
  
  /**
   * Close the Neo4j driver when done
   */
  async close(): Promise<void> {
    if (this.driver) {
      await this.driver.close();
      this.initialized = false;
      console.log('[Neo4jAdapter] Connection to Neo4j closed');
    }
  }
}