diff --git a/cpg-neo4j/src/main/kotlin/de/fraunhofer/aisec/cpg/v2/Persistence.kt b/cpg-neo4j/src/main/kotlin/de/fraunhofer/aisec/cpg/persistence/Persistence.kt similarity index 95% rename from cpg-neo4j/src/main/kotlin/de/fraunhofer/aisec/cpg/v2/Persistence.kt rename to cpg-neo4j/src/main/kotlin/de/fraunhofer/aisec/cpg/persistence/Persistence.kt index d5b01e3d03..fd264b40ce 100644 --- a/cpg-neo4j/src/main/kotlin/de/fraunhofer/aisec/cpg/v2/Persistence.kt +++ b/cpg-neo4j/src/main/kotlin/de/fraunhofer/aisec/cpg/persistence/Persistence.kt @@ -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 @@ -49,7 +49,7 @@ 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 @@ -57,15 +57,6 @@ 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, Set> = mutableMapOf() val schemaPropertiesCache: @@ -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") @@ -108,12 +100,13 @@ fun TranslationResult.persist() { b.stop() } +context(Session) private fun List.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 @@ -129,10 +122,11 @@ private fun List.persist() { } } +context(Session) private fun Collection>.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() } @@ -157,6 +151,7 @@ private fun Collection>.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.persistExtraRelationships() { val relationships = this.flatMap { @@ -195,15 +190,15 @@ private fun List.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>, ) { 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 diff --git a/cpg-neo4j/src/main/kotlin/de/fraunhofer/aisec/cpg_vis_neo4j/Application.kt b/cpg-neo4j/src/main/kotlin/de/fraunhofer/aisec/cpg_vis_neo4j/Application.kt index ea2b225751..b286dfc02a 100644 --- a/cpg-neo4j/src/main/kotlin/de/fraunhofer/aisec/cpg_vis_neo4j/Application.kt +++ b/cpg-neo4j/src/main/kotlin/de/fraunhofer/aisec/cpg_vis_neo4j/Application.kt @@ -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 @@ -381,33 +378,14 @@ class Application : Callable { * 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() } /** @@ -420,41 +398,14 @@ class Application : Callable { * @throws ConnectException, if there is no connection to bolt://localhost:7687 possible */ @Throws(InterruptedException::class, ConnectException::class) - fun connect(): Pair { - 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() } /** diff --git a/cpg-neo4j/src/test/kotlin/de/fraunhofer/aisec/cpg/v2/TestPersistence.kt b/cpg-neo4j/src/test/kotlin/de/fraunhofer/aisec/cpg/v2/TestPersistence.kt deleted file mode 100644 index 9ec02f4896..0000000000 --- a/cpg-neo4j/src/test/kotlin/de/fraunhofer/aisec/cpg/v2/TestPersistence.kt +++ /dev/null @@ -1,41 +0,0 @@ -/* - * Copyright (c) 2024, Fraunhofer AISEC. All rights reserved. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - * - * $$$$$$\ $$$$$$$\ $$$$$$\ - * $$ __$$\ $$ __$$\ $$ __$$\ - * $$ / \__|$$ | $$ |$$ / \__| - * $$ | $$$$$$$ |$$ |$$$$\ - * $$ | $$ ____/ $$ |\_$$ | - * $$ | $$\ $$ | $$ | $$ | - * \$$$$$ |$$ | \$$$$$ | - * \______/ \__| \______/ - * - */ -package de.fraunhofer.aisec.cpg.v2 - -import de.fraunhofer.aisec.cpg_vis_neo4j.createTranslationResult -import org.junit.jupiter.api.Tag -import org.junit.jupiter.api.Test - -@Tag("integration") -class TestPersistence { - @Test - fun testPersist() { - val result = createTranslationResult() - - neo4jSession.executeWrite { tx -> tx.run("MATCH (n) DETACH DELETE n").consume() } - result.second.persist() - } -} diff --git a/cpg-neo4j/src/test/kotlin/de/fraunhofer/aisec/cpg_vis_neo4j/Neo4JTest.kt b/cpg-neo4j/src/test/kotlin/de/fraunhofer/aisec/cpg_vis_neo4j/Neo4JTest.kt index 96f477d58b..a77d2d0269 100644 --- a/cpg-neo4j/src/test/kotlin/de/fraunhofer/aisec/cpg_vis_neo4j/Neo4JTest.kt +++ b/cpg-neo4j/src/test/kotlin/de/fraunhofer/aisec/cpg_vis_neo4j/Neo4JTest.kt @@ -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") @@ -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() - } - } }