Skip to content

Commit

Permalink
Using new API in neo4j application
Browse files Browse the repository at this point in the history
  • Loading branch information
@oxisto
oxisto committed Dec 11, 2024
1 parent e6d1980 commit 684b2a6
Show file tree
Hide file tree
Showing 4 changed files with 26 additions and 156 deletions.
27 changes: 11 additions & 16 deletions ...de/fraunhofer/aisec/cpg/v2/Persistence.kt → ...ofer/aisec/cpg/persistence/Persistence.kt
View file
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,7 @@
* \______/ \__| \______/
*
*/
package de.fraunhofer.aisec.cpg.v2
package de.fraunhofer.aisec.cpg.persistence

import de.fraunhofer.aisec.cpg.TranslationResult
import de.fraunhofer.aisec.cpg.graph.Name
Expand All @@ -49,23 +49,14 @@ import kotlin.reflect.full.memberProperties
import kotlin.reflect.full.superclasses
import kotlin.reflect.full.withNullability
import kotlin.uuid.Uuid
import org.neo4j.driver.GraphDatabase
import org.neo4j.driver.Session
import org.neo4j.ogm.typeconversion.CompositeAttributeConverter
import org.slf4j.LoggerFactory

/**
* docker run \ --name neo4j-apoc \ -p 7474:7474 -p 7687:7687 \ -d \ -e NEO4J_AUTH=neo4j/password \
* -e NEO4JLABS_PLUGINS='["apoc"]' \ neo4j:5
*/
val dbUri = "neo4j://localhost"
val dbUser = "neo4j"
val dbPassword = "password"

val neo4jSession by lazy {
val driver = GraphDatabase.driver(dbUri, org.neo4j.driver.AuthTokens.basic(dbUser, dbPassword))
driver.session()
}

val labelCache: MutableMap<KClass<*>, Set<String>> = mutableMapOf()

val schemaPropertiesCache:
Expand All @@ -77,6 +68,7 @@ val log = LoggerFactory.getLogger("Persistence")
val edgeChunkSize = 10000
val nodeChunkSize = 10000

context(Session)
fun TranslationResult.persist() {
val b = Benchmark(Persistable::class.java, "Persisting translation result")

Expand Down Expand Up @@ -108,12 +100,13 @@ fun TranslationResult.persist() {
b.stop()
}

context(Session)
private fun List<Node>.persist() {
this.chunked(nodeChunkSize).map { chunk ->
val b = Benchmark(Persistable::class.java, "Persisting chunk of ${chunk.size} nodes")
val params =
mapOf("props" to chunk.map { mapOf("labels" to it::class.labels) + it.properties() })
neo4jSession.executeWrite { tx ->
this@Session.executeWrite { tx ->
tx.run(
"""
UNWIND ${"$"}props AS map
Expand All @@ -129,10 +122,11 @@ private fun List<Node>.persist() {
}
}

context(Session)
private fun Collection<Edge<*>>.persist() {
// Create an index for the "id" field of node, because we are "MATCH"ing on it in the edge
// creation. We need to wait for this to be finished
neo4jSession.executeWrite { tx ->
this@Session.executeWrite { tx ->
tx.run("CREATE INDEX IF NOT EXISTS FOR (n:Node) ON (n.id)").consume()
}

Expand All @@ -157,6 +151,7 @@ private fun Collection<Edge<*>>.persist() {
* Some of our relationships are not real "edges" (i.e., [Edge]) (yet). We need to handle these case
* separately (for now).
*/
context(Session)
private fun List<Node>.persistExtraRelationships() {
val relationships =
this.flatMap {
Expand Down Expand Up @@ -195,15 +190,15 @@ private fun List<Node>.persistExtraRelationships() {
)
}

relationships.chunked(10000).map { chunk -> createRelationships(chunk) }
relationships.chunked(10000).map { chunk -> this@Session.createRelationships(chunk) }
}

private fun createRelationships(
private fun Session.createRelationships(
props: List<Map<String, Any?>>,
) {
val b = Benchmark(Persistable::class.java, "Persisting chunk of ${props.size} relationships")
val params = mapOf("props" to props)
neo4jSession.executeWrite { tx ->
executeWrite { tx ->
tx.run(
"""
UNWIND ${'$'}props AS map
Expand Down
79 changes: 15 additions & 64 deletions cpg-neo4j/src/main/kotlin/de/fraunhofer/aisec/cpg_vis_neo4j/Application.kt
View file
Original file line number Diff line number Diff line change
Expand Up @@ -30,23 +30,20 @@ import de.fraunhofer.aisec.cpg.*
import de.fraunhofer.aisec.cpg.frontends.CompilationDatabase.Companion.fromFile
import de.fraunhofer.aisec.cpg.helpers.Benchmark
import de.fraunhofer.aisec.cpg.passes.*
import de.fraunhofer.aisec.cpg.persistence.persist
import java.io.File
import java.net.ConnectException
import java.nio.file.Paths
import java.util.concurrent.Callable
import kotlin.reflect.KClass
import kotlin.system.exitProcess
import org.neo4j.driver.exceptions.AuthenticationException
import org.neo4j.ogm.config.Configuration
import org.neo4j.driver.GraphDatabase
import org.neo4j.ogm.context.EntityGraphMapper
import org.neo4j.ogm.context.MappingContext
import org.neo4j.ogm.cypher.compiler.MultiStatementCypherCompiler
import org.neo4j.ogm.cypher.compiler.builders.node.DefaultNodeBuilder
import org.neo4j.ogm.cypher.compiler.builders.node.DefaultRelationshipBuilder
import org.neo4j.ogm.exception.ConnectionException
import org.neo4j.ogm.metadata.MetaData
import org.neo4j.ogm.session.Session
import org.neo4j.ogm.session.SessionFactory
import org.slf4j.Logger
import org.slf4j.LoggerFactory
import picocli.CommandLine
Expand Down Expand Up @@ -381,33 +378,14 @@ class Application : Callable<Int> {
* Pushes the whole translationResult to the neo4j db.
*
* @param translationResult, not null
* @throws InterruptedException, if the thread is interrupted while it try´s to connect to the
* neo4j db.
* @throws ConnectException, if there is no connection to bolt://localhost:7687 possible
*/
@Throws(InterruptedException::class, ConnectException::class)
fun pushToNeo4j(translationResult: TranslationResult) {
val bench = Benchmark(this.javaClass, "Push cpg to neo4j", false, translationResult)
log.info("Using import depth: $depth")
log.info(
"Count base nodes to save: " +
translationResult.components.size +
translationResult.additionalNodes.size
)

val sessionAndSessionFactoryPair = connect()

val session = sessionAndSessionFactoryPair.first
session.beginTransaction().use { transaction ->
if (!noPurgeDb) session.purgeDatabase()
session.save(translationResult.components, depth)
session.save(translationResult.additionalNodes, depth)
transaction.commit()
val session = connect()
with(session) {
if (!noPurgeDb) executeWrite { tx -> tx.run("MATCH (n) DETACH DELETE n").consume() }
translationResult.persist()
}

session.clear()
sessionAndSessionFactoryPair.second.close()
bench.addMeasurement()
session.close()
}

/**
Expand All @@ -420,41 +398,14 @@ class Application : Callable<Int> {
* @throws ConnectException, if there is no connection to bolt://localhost:7687 possible
*/
@Throws(InterruptedException::class, ConnectException::class)
fun connect(): Pair<Session, SessionFactory> {
var fails = 0
var sessionFactory: SessionFactory? = null
var session: Session? = null
while (session == null && fails < MAX_COUNT_OF_FAILS) {
try {
val configuration =
Configuration.Builder()
.uri("$PROTOCOL$host:$port")
.credentials(neo4jUsername, neo4jPassword)
.verifyConnection(VERIFY_CONNECTION)
.build()
sessionFactory = SessionFactory(configuration, *packages)

session = sessionFactory.openSession()
} catch (ex: ConnectionException) {
sessionFactory = null
fails++
log.error(
"Unable to connect to localhost:7687, " +
"ensure the database is running and that " +
"there is a working network connection to it."
)
Thread.sleep(TIME_BETWEEN_CONNECTION_TRIES)
} catch (ex: AuthenticationException) {
log.error("Unable to connect to localhost:7687, wrong username/password!")
exitProcess(EXIT_FAILURE)
}
}
if (session == null || sessionFactory == null) {
log.error("Unable to connect to localhost:7687")
exitProcess(EXIT_FAILURE)
}
assert(fails <= MAX_COUNT_OF_FAILS)
return Pair(session, sessionFactory)
fun connect(): org.neo4j.driver.Session {
val driver =
GraphDatabase.driver(
"$PROTOCOL$host:$port",
org.neo4j.driver.AuthTokens.basic(neo4jUsername, neo4jPassword)
)
driver.verifyConnectivity()
return driver.session()
}

/**
Expand Down
41 changes: 0 additions & 41 deletions cpg-neo4j/src/test/kotlin/de/fraunhofer/aisec/cpg/v2/TestPersistence.kt
View file

This file was deleted.

35 changes: 0 additions & 35 deletions cpg-neo4j/src/test/kotlin/de/fraunhofer/aisec/cpg_vis_neo4j/Neo4JTest.kt
View file
Original file line number Diff line number Diff line change
Expand Up @@ -25,14 +25,9 @@
*/
package de.fraunhofer.aisec.cpg_vis_neo4j

import de.fraunhofer.aisec.cpg.frontends.TestLanguageFrontend
import de.fraunhofer.aisec.cpg.graph.Name
import de.fraunhofer.aisec.cpg.graph.builder.translationResult
import de.fraunhofer.aisec.cpg.graph.declarations.ImportDeclaration
import de.fraunhofer.aisec.cpg.graph.functions
import kotlin.test.Test
import kotlin.test.assertEquals
import kotlin.test.assertNotNull
import org.junit.jupiter.api.Tag

@Tag("integration")
Expand All @@ -47,34 +42,4 @@ class Neo4JTest {

application.pushToNeo4j(translationResult)
}

@Test
fun testSimpleNameConverter() {
val result =
with(TestLanguageFrontend()) {
translationResult {
val import = ImportDeclaration()
import.name = Name("myname")
import.alias = Name("myname", Name("myparent"), "::")
additionalNodes += import
}
}

val app = Application()
app.pushToNeo4j(result)

val sessionAndSessionFactoryPair = app.connect()

val session = sessionAndSessionFactoryPair.first
session.beginTransaction().use { transaction ->
val imports = session.loadAll(ImportDeclaration::class.java)
assertNotNull(imports)

var loadedImport = imports.singleOrNull()
assertNotNull(loadedImport)
assertEquals("myname", loadedImport.alias?.localName)

transaction.commit()
}
}
}

0 comments on commit 684b2a6

Please sign in to comment.