Module: Neo4jBolt

Extended by:

Neo4jBolt

Included in:

Neo4jBolt

Defined in:

lib/neo4j_bolt.rb,
lib/neo4j_bolt/version.rb

Defined Under Namespace

Modules: DriverRegistry Classes: ConstraintValidationFailedError, Error, ExpectedOneResultError, IntegerOutOfRangeError, Node, Relationship, SyntaxError, TransactionContext

Constant Summary

CONSTRAINT_INDEX_PREFIX =

"neo4j_bolt_"

LOAD_INITIAL_BATCH_SIZE =

5_000

MIN_INTEGER =

-(2**63)

MAX_INTEGER =

(2**63) - 1

TRANSACTION_CONTEXT_KEY =

:neo4j_bolt_transaction_contexts

VERSION =

"0.4.3"

Class Attribute Summary

• .bolt_host ⇒ Object

  Returns the value of attribute bolt_host.

• .bolt_port ⇒ Object

  Returns the value of attribute bolt_port.

• .bolt_verbosity ⇒ Object

  Returns the value of attribute bolt_verbosity.

Instance Method Summary

• #cleanup_neo4j ⇒ Object
• #dump_database(io, progress_io: nil, progress_color: "cyan") ⇒ Object
• #load_database_dump(io, force_append: false, progress_io: nil, progress_color: "cyan", initial_batch_size: LOAD_INITIAL_BATCH_SIZE) ⇒ Object
• #neo4j_query(query, data = {}, &block) ⇒ Object
• #neo4j_query_expect_one(query, data = {}) ⇒ Object
• #rollback ⇒ Object

  Explicit rollback is rollback-only: the surrounding block may finish, but its one real upstream transaction can no longer commit.

• #setup_constraints_and_indexes(constraints, indexes) ⇒ Object
• #transaction ⇒ Object
• #wait_for_neo4j ⇒ Object

Class Attribute Details

.bolt_host ⇒ Object

Returns the value of attribute bolt_host.

# File 'lib/neo4j_bolt.rb', line 357

def bolt_host
  @bolt_host
end

.bolt_port ⇒ Object

Returns the value of attribute bolt_port.

# File 'lib/neo4j_bolt.rb', line 357

def bolt_port
  @bolt_port
end

.bolt_verbosity ⇒ Object

Returns the value of attribute bolt_verbosity.

# File 'lib/neo4j_bolt.rb', line 358

def bolt_verbosity
  @bolt_verbosity
end

Instance Method Details

#cleanup_neo4j ⇒ Object

# File 'lib/neo4j_bolt.rb', line 471

def cleanup_neo4j
  DriverRegistry.cleanup
  nil
rescue Neo4j::Driver::Exceptions::Neo4jException => error
  raise_mapped_error(error)
end

#dump_database(io, progress_io: nil, progress_color: "cyan") ⇒ Object

# File 'lib/neo4j_bolt.rb', line 539

def dump_database(io, progress_io: nil, progress_color: "cyan")
  progress = ProgressReporter.new(progress_io, color: progress_color)
  dumped_nodes = 0
  dumped_relationships = 0

  transaction do
    identity_function = dump_identity_function
    total_nodes = neo4j_query_expect_one("MATCH (n) RETURN count(n) AS count")["count"]
    total_relationships = neo4j_query_expect_one("MATCH ()-[r]->() RETURN count(r) AS count")["count"]
    dump_ids = total_relationships.positive? ? {} : nil

    progress.update(
      progress_message("Dumping", dumped_nodes, dumped_relationships, total_nodes, total_relationships),
      force: true
    )

    neo4j_query(
      "MATCH (n) " \
      "RETURN #{identity_function}(n) AS identity, labels(n) AS labels, properties(n) AS properties " \
      "ORDER BY #{identity_function}(n)"
    ) do |row|
      dump_ids[row["identity"].to_s] = dumped_nodes if dump_ids
      io.puts "n #{JSON.generate(id: dumped_nodes, labels: row["labels"], properties: row["properties"])}"
      dumped_nodes += 1
      if (dumped_nodes % DUMP_PROGRESS_STEP).zero?
        progress.update(progress_message(
          "Dumping", dumped_nodes, dumped_relationships, total_nodes, total_relationships
        ))
      end
    end

    if total_relationships.positive?
      neo4j_query(
        "MATCH (from)-[r]->(to) " \
        "RETURN #{identity_function}(from) AS from_identity, " \
        "#{identity_function}(to) AS to_identity, type(r) AS type, properties(r) AS properties " \
        "ORDER BY #{identity_function}(r)"
      ) do |row|
        io.puts "r #{JSON.generate(
          from: dump_ids.fetch(row["from_identity"].to_s),
          to: dump_ids.fetch(row["to_identity"].to_s),
          type: row["type"],
          properties: row["properties"]
        )}"
        dumped_relationships += 1
        if (dumped_relationships % DUMP_PROGRESS_STEP).zero?
          progress.update(progress_message(
            "Dumping", dumped_nodes, dumped_relationships, total_nodes, total_relationships
          ))
        end
      end
    end
  end

  progress.finish("Dumped: #{dumped_nodes} nodes, #{dumped_relationships} relationships")
  nil
ensure
  progress&.close
end

#load_database_dump(io, force_append: false, progress_io: nil, progress_color: "cyan", initial_batch_size: LOAD_INITIAL_BATCH_SIZE) ⇒ Object

# File 'lib/neo4j_bolt.rb', line 599

def load_database_dump(io, force_append: false, progress_io: nil, progress_color: "cyan",
                       initial_batch_size: LOAD_INITIAL_BATCH_SIZE)
  raise Error, "load_database_dump cannot run inside a transaction" if transaction_context

  initial_batch_size = Integer(initial_batch_size)
  raise ArgumentError, "initial_batch_size must be positive" unless initial_batch_size.positive?

  unless force_append
    count = neo4j_query_expect_one("MATCH (n) RETURN count(n) AS count")["count"]
    raise Error, "There are nodes in this database, exiting now." unless count.zero?
  end

  progress = ProgressReporter.new(progress_io, color: progress_color)
  node_batches = Hash.new { |hash, key| hash[key] = [] }
  relationship_batches = Hash.new { |hash, key| hash[key] = [] }
  total_nodes, total_relationships = parse_dump(io, node_batches, relationship_batches, progress)

  temporary_token = SecureRandom.hex(12)
  temporary_label = "__neo4j_bolt_load_#{temporary_token}"
  temporary_id_property = "__neo4j_bolt_load_id_#{temporary_token}"
  temporary_index = "neo4j_bolt_load_#{temporary_token}"
  quoted_temporary_label = quote_identifier(temporary_label)
  quoted_temporary_id = quote_identifier(temporary_id_property)
  quoted_temporary_index = quote_identifier(temporary_index)
  needs_relationship_lookup = total_relationships.positive?
  loaded_nodes = 0
  loaded_relationships = 0
  index_created = false

  progress.update(
    progress_message("Loading", loaded_nodes, loaded_relationships, total_nodes, total_relationships,
                     batch_size: initial_batch_size),
    force: true
  )

  original_error = nil
  cleanup_error = nil
  begin
    node_batch_size = initial_batch_size
    node_batches.each_value do |nodes|
      node_batch_size = adaptive_each_slice(nodes, node_batch_size, progress, "node") do |slice, current_batch_size|
        labels = slice.first.fetch("labels")
        label_clause = labels.map { |label| ":#{quote_identifier(label)}" }.join
        label_clause += ":#{quoted_temporary_label}" if needs_relationship_lookup
        temporary_assignment = if needs_relationship_lookup
                                 "SET n.#{quoted_temporary_id} = item.id\n"
                               else
                                 ""
                               end
        count = neo4j_query_expect_one(<<~CYPHER, nodes: slice)["count_n"]
          UNWIND $nodes AS item
          CREATE (n#{label_clause})
          SET n = item.properties
          #{temporary_assignment}RETURN count(n) AS count_n
        CYPHER
        raise Error, "Expected #{slice.size} nodes, got #{count}." unless count == slice.size

        loaded_nodes += slice.size
        progress.update(progress_message(
          "Loading", loaded_nodes, loaded_relationships, total_nodes, total_relationships,
          batch_size: current_batch_size
        ))
      end
    end

    if needs_relationship_lookup
      progress.note("Building temporary relationship lookup index...")
      neo4j_query(
        "CREATE INDEX #{quoted_temporary_index} " \
        "FOR (n:#{quoted_temporary_label}) ON (n.#{quoted_temporary_id})"
      )
      index_created = true
      wait_for_load_index(temporary_index, progress)
      progress.update(
        progress_message("Loading", loaded_nodes, loaded_relationships, total_nodes, total_relationships,
                         batch_size: initial_batch_size),
        force: true
      )

      relationship_batch_size = initial_batch_size
      relationship_batches.each do |type, relationships|
        relationship_batch_size = adaptive_each_slice(
          relationships, relationship_batch_size, progress, "relationship"
        ) do |slice, current_batch_size|
          count = neo4j_query_expect_one(<<~CYPHER, relationships: slice)["count_r"]
            UNWIND $relationships AS item
            MATCH (from:#{quoted_temporary_label} {#{quoted_temporary_id}: item.from})
            MATCH (to:#{quoted_temporary_label} {#{quoted_temporary_id}: item.to})
            CREATE (from)-[r:#{quote_identifier(type)}]->(to)
            SET r = item.properties
            RETURN count(r) AS count_r
          CYPHER
          raise Error, "Expected #{slice.size} relationships, got #{count}." unless count == slice.size

          loaded_relationships += slice.size
          progress.update(progress_message(
            "Loading", loaded_nodes, loaded_relationships, total_nodes, total_relationships,
            batch_size: current_batch_size
          ))
        end
      end
    end
  rescue Exception => error # rubocop:disable Lint/RescueException -- preserve original failure through cleanup
    original_error = error
    raise
  ensure
    if index_created
      begin
        neo4j_query("DROP INDEX #{quoted_temporary_index} IF EXISTS")
      rescue Error => error
        cleanup_error ||= error
      end
    end

    if needs_relationship_lookup
      begin
        progress&.note("Cleaning temporary load metadata...")
        cleanup_load_metadata(
          quoted_temporary_label, quoted_temporary_id, initial_batch_size, total_nodes, progress
        )
      rescue Error => error
        cleanup_error ||= error
      end
    end

    raise cleanup_error if original_error.nil? && cleanup_error
  end

  progress.finish("Loaded: #{loaded_nodes} nodes, #{loaded_relationships} relationships")
  nil
ensure
  progress&.close
end

#neo4j_query(query, data = {}, &block) ⇒ Object

# File 'lib/neo4j_bolt.rb', line 384

def neo4j_query(query, data = {}, &block)
  validate_parameters!(data)
  log_query(query, data)

  context = transaction_context
  if context
    run_and_convert(context.transaction, query, data, &block)
  else
    DriverRegistry.with_driver do |driver|
      driver.session do |session|
        run_and_convert(session, query, data, &block)
      end
    end
  end
rescue Neo4j::Driver::Exceptions::Neo4jException => error
  mark_transaction_failed!
  raise_mapped_error(error)
rescue IntegerOutOfRangeError
  mark_transaction_failed!
  raise
rescue Exception # rubocop:disable Lint/RescueException -- preserve rollback-only on callback/local errors
  mark_transaction_failed!
  raise
end

#neo4j_query_expect_one(query, data = {}) ⇒ Object

Raises:

• (ExpectedOneResultError)

# File 'lib/neo4j_bolt.rb', line 409

def neo4j_query_expect_one(query, data = {})
  rows = neo4j_query(query, data)
  return rows.first if rows.size == 1

  mark_transaction_failed!
  raise ExpectedOneResultError, "Expected one result, but got #{rows.size}."
end

#rollback ⇒ Object

Explicit rollback is rollback-only: the surrounding block may finish, but its one real upstream transaction can no longer commit.

Raises:

• (Error)

# File 'lib/neo4j_bolt.rb', line 463

def rollback
  context = transaction_context
  raise Error, "rollback called outside a transaction" unless context

  context.rollback_only = true
  nil
end

#setup_constraints_and_indexes(constraints, indexes) ⇒ Object

# File 'lib/neo4j_bolt.rb', line 497

def setup_constraints_and_indexes(constraints, indexes)
  wanted_constraints = Set.new
  wanted_indexes = Set.new

  constraints.each do |entry|
    label, property = schema_entry(entry, "constraint")
    name = "#{CONSTRAINT_INDEX_PREFIX}#{label}_#{property}"
    wanted_constraints << name
    neo4j_query(
      "CREATE CONSTRAINT #{quote_identifier(name)} IF NOT EXISTS " \
      "FOR (n:#{quote_identifier(label)}) REQUIRE n.#{quote_identifier(property)} IS UNIQUE"
    )
  end

  indexes.each do |entry|
    label, property = schema_entry(entry, "index")
    name = "#{CONSTRAINT_INDEX_PREFIX}#{label}_#{property}"
    wanted_indexes << name
    neo4j_query(
      "CREATE INDEX #{quote_identifier(name)} IF NOT EXISTS " \
      "FOR (n:#{quote_identifier(label)}) ON (n.#{quote_identifier(property)})"
    )
  end

  neo4j_query("SHOW CONSTRAINTS").each do |row|
    name = row["name"]
    next unless name&.start_with?(CONSTRAINT_INDEX_PREFIX)
    next if wanted_constraints.include?(name)

    neo4j_query("DROP CONSTRAINT #{quote_identifier(name)}")
  end

  neo4j_query("SHOW INDEXES").each do |row|
    name = row["name"]
    next unless name&.start_with?(CONSTRAINT_INDEX_PREFIX)
    next if wanted_indexes.include?(name) || wanted_constraints.include?(name)

    neo4j_query("DROP INDEX #{quote_identifier(name)}")
  end
  nil
end

#transaction ⇒ Object

# File 'lib/neo4j_bolt.rb', line 417

def transaction
  context = transaction_context
  if context
    context.depth += 1
    begin
      yield
    rescue Exception # rubocop:disable Lint/RescueException -- rollback-only must survive every unwind
      context.rollback_only = true
      raise
    ensure
      context.depth -= 1
    end
    return nil
  end

  DriverRegistry.with_driver do |driver|
    driver.session do |session|
      upstream_transaction = session.begin_transaction
      context = TransactionContext.new(
        session: session, transaction: upstream_transaction, depth: 1, rollback_only: false
      )
      transaction_contexts[self.object_id] = context

      begin
        yield
      rescue Exception # rubocop:disable Lint/RescueException -- ensure rollback for non-StandardError too
        context.rollback_only = true
        raise
      ensure
        begin
          context.rollback_only ? upstream_transaction.rollback : upstream_transaction.commit
        rescue Neo4j::Driver::Exceptions::Neo4jException => error
          raise_mapped_error(error)
        ensure
          transaction_contexts.delete(self.object_id)
        end
      end
    end
  end
  nil
rescue Neo4j::Driver::Exceptions::Neo4jException => error
  raise_mapped_error(error)
end

#wait_for_neo4j ⇒ Object

# File 'lib/neo4j_bolt.rb', line 478

def wait_for_neo4j
  attempts = Integer(ENV.fetch("NEO4J_BOLT_WAIT_ATTEMPTS", "30"))
  delay = Float(ENV.fetch("NEO4J_BOLT_WAIT_DELAY", "1"))
  last_error = nil

  attempts.times do |attempt|
    begin
      neo4j_query("RETURN 1 AS ready")
      return true
    rescue Error => error
      last_error = error
      warn "Waiting for Neo4j (attempt #{attempt + 1}/#{attempts})..." if Neo4jBolt.bolt_verbosity.to_i.positive?
      sleep delay if attempt + 1 < attempts
    end
  end

  raise(last_error || Error.new("Neo4j did not become ready"))
end