Make authentication handling work with 6.x Java Driver

The main objective of this update is to adjust authentication handling to make it work with 6.x Java Driver.

Author: injectives

common/src/main/scala/org/neo4j/spark/service/SchemaService.scala

@@ -730,7 +730,7 @@ class SchemaService(
       s"spark_${entityType}_${constraintType.replace(s"$entityType ", "")}-CONSTRAINT_${entityIdentifier}_$dashSeparatedProps".quote()
     val props = keys.values.map(_.quote()).map("e." + _).mkString(", ")
     val asciiRepresentation: String = createCypherPattern(entityType, entityIdentifier)
-    session.executeWrite(
+    session.executeWriteWithoutResult(
       tx => {
         tx.run(
           s"CREATE CONSTRAINT $constraintName IF NOT EXISTS FOR $asciiRepresentation REQUIRE ($props) IS $constraintType"

common/src/main/scala/org/neo4j/spark/util/Neo4jOptions.scala

@@ -20,12 +20,9 @@ import org.apache.spark.internal.Logging
 import org.apache.spark.sql.SaveMode
 import org.apache.spark.sql.SparkSession
 import org.jetbrains.annotations.TestOnly
-import org.neo4j.connectors.authn.AuthenticationToken
 import org.neo4j.connectors.authn.AuthenticationTokenSupplierFactory
 import org.neo4j.connectors.authn.BearerAuthenticationToken
 import org.neo4j.connectors.authn.CustomAuthenticationToken
-import org.neo4j.connectors.authn.DisabledAuthenticationToken
-import org.neo4j.connectors.authn.ExpiringAuthenticationToken
 import org.neo4j.connectors.authn.KerberosAuthenticationToken
 import org.neo4j.connectors.authn.UserNameAndPasswordAuthenticationToken
 import org.neo4j.driver.Config.TrustStrategy
@@ -40,8 +37,9 @@ import java.time.Duration
 import java.util
 import java.util.ServiceLoader
 import java.util.UUID
+import java.util.concurrent.CompletableFuture
+import java.util.concurrent.CompletionStage
 import java.util.concurrent.TimeUnit
-import java.util.function.Supplier
 
 import scala.collection.JavaConverters._
 import scala.language.implicitConversions
@@ -498,47 +496,51 @@ case class Neo4jDriverOptions(
     (URI.create(urls.head.trim), resolved)
   }
 
-  // TODO this is partially intentionally not working as expect, changed to make it compile
-  // TODO missing here is also the keycloack support that comes from the authn commons
   private def createAuthTokenManager: AuthTokenManager = {
     if (auth == null || auth.isEmpty) {
       throw new IllegalArgumentException(s"Authentication type name is required")
     }
-    val token = createAuthTokenSupplier.get()
-    token match {
-      case bearerAuthenticationToken: BearerAuthenticationToken =>
-        AuthTokenManagers.bearer(() =>
-          AuthTokens.bearer(
-            bearerAuthenticationToken.getToken
-          ).expiringAt(bearerAuthenticationToken.getExpiresAt.toEpochMilli)
-        )
-      case customAuthenticationToken: CustomAuthenticationToken =>
-        AuthTokenManagers.basic(() =>
-          AuthTokens.custom(
-            customAuthenticationToken.getPrincipal,
-            customAuthenticationToken.getCredentials,
-            customAuthenticationToken.getRealm,
-            customAuthenticationToken.getScheme,
-            customAuthenticationToken.getParameters
-          )
-        )
-      case disabledAuthenticationToken: DisabledAuthenticationToken =>
-        AuthTokenManagers.basic(() => AuthTokens.none())
-      case kerberosAuthenticationToken: KerberosAuthenticationToken =>
-        AuthTokenManagers.basic(() => AuthTokens.kerberos(kerberosAuthenticationToken.getToken))
-      case userNameAndPasswordAuthenticationToken: UserNameAndPasswordAuthenticationToken =>
-        AuthTokenManagers.basic(() =>
-          AuthTokens.basic(
-            userNameAndPasswordAuthenticationToken.getUsername,
-            userNameAndPasswordAuthenticationToken.getPassword,
-            userNameAndPasswordAuthenticationToken.getRealm
-          )
-        )
-      case _ => throw new IllegalStateException("bam")
+    val token = createAuthTokenSupplier
+    val name = token.getName
+    val username = authParameters.get("username")
+    val password = authParameters.get("password")
+    val supplier = token.create(username.orNull, password.orNull, authParameters.asJava)
+
+    name match {
+      case "basic" =>
+        val token = supplier.get().asInstanceOf[UserNameAndPasswordAuthenticationToken]
+        new StaticAuthTokenManager(AuthTokens.basic(token.getUsername, token.getPassword))
+      case "bearer" | "keycloak" =>
+        AuthTokenManagers.bearer(() => {
+          val token = supplier.get().asInstanceOf[BearerAuthenticationToken]
+          val authToken = AuthTokens.bearer(token.getToken)
+          val exp = token.getExpiresAt
+          if (exp == null) {
+            authToken.expiringAt(Long.MaxValue)
+          } else {
+            authToken.expiringAt(exp.toEpochMilli)
+          }
+        })
+      case "custom" =>
+        val token = supplier.get().asInstanceOf[CustomAuthenticationToken]
+        new StaticAuthTokenManager(AuthTokens.custom(
+          token.getPrincipal,
+          token.getCredentials,
+          token.getRealm,
+          token.getScheme,
+          token.getParameters
+        ))
+      case "kerberos" =>
+        AuthTokenManagers.basic(() => {
+          val token = supplier.get().asInstanceOf[KerberosAuthenticationToken]
+          AuthTokens.kerberos(token.getToken)
+        })
+      case "none" =>
+        new StaticAuthTokenManager(AuthTokens.none())
     }
   }
 
-  private def createAuthTokenSupplier: Supplier[AuthenticationToken] = {
+  private def createAuthTokenSupplier: AuthenticationTokenSupplierFactory = {
     if (auth == null || auth.isEmpty) {
       throw new IllegalArgumentException(s"Authentication type name is required")
     }
@@ -562,9 +564,7 @@ case class Neo4jDriverOptions(
       )
     }
 
-    val username = authParameters.get("username")
-    val password = authParameters.get("password")
-    filteredSupplierFactories.head.create(username.orNull, password.orNull, authParameters.asJava)
+    filteredSupplierFactories.head
   }
 
 }
@@ -722,6 +722,11 @@ object Neo4jOptions {
   }
 }
 
+class StaticAuthTokenManager(authToken: AuthToken) extends AuthTokenManager {
+  override def getToken: CompletionStage[AuthToken] = CompletableFuture.completedStage(authToken)
+  override def handleSecurityException(authToken: AuthToken, exception: exceptions.SecurityException): Boolean = false
+}
+
 class CaseInsensitiveEnumeration extends Enumeration {
 
   def withCaseInsensitiveName(s: String): Value = {

common/src/test/scala/org/neo4j/spark/service/AuthenticationTest.scala

@@ -18,15 +18,17 @@ package org.neo4j.spark.service
 
 import org.junit.Test
 import org.junit.runner.RunWith
-import org.mockito.ArgumentMatchers
+import org.mockito.ArgumentCaptor
 import org.mockito.ArgumentMatchers._
 import org.mockito.Mockito.times
+import org.neo4j.driver.AuthTokenManager
+import org.neo4j.driver.AuthTokens
 import org.neo4j.driver.Config
 import org.neo4j.driver.GraphDatabase
-import org.neo4j.driver.internal.security.ExpirationBasedAuthTokenManager
 import org.neo4j.spark.util.DriverCache
 import org.neo4j.spark.util.Neo4jOptions
 import org.powermock.api.mockito.PowerMockito
+import org.powermock.core.classloader.annotations.PowerMockIgnore
 import org.powermock.core.classloader.annotations.PrepareForTest
 import org.powermock.modules.junit4.PowerMockRunner
 import org.testcontainers.shaded.com.google.common.io.BaseEncoding
@@ -36,6 +38,7 @@ import java.util
 
 @PrepareForTest(Array(classOf[GraphDatabase]))
 @RunWith(classOf[PowerMockRunner])
+@PowerMockIgnore(Array("javax.management.*"))
 class AuthenticationTest {
 
   @Test
@@ -56,8 +59,9 @@ class AuthenticationTest {
     driverCache.getOrCreate()
 
     PowerMockito.verifyStatic(classOf[GraphDatabase], times(1))
-    // was GraphDatabase.driver(any[URI](), ArgumentMatchers.eq(AuthTokens.custom("", token, "", "")), any(classOf[Config]))
-    GraphDatabase.driver(any[URI](), ArgumentMatchers.any[ExpirationBasedAuthTokenManager], any[Config]())
+    val managerCaptor = ArgumentCaptor.forClass(classOf[AuthTokenManager])
+    GraphDatabase.driver(any[URI](), managerCaptor.capture(), any[Config]())
+    assert(AuthTokens.custom("", token, "", "") == managerCaptor.getValue.getToken.toCompletableFuture.join())
   }
 
   @Test
@@ -77,7 +81,8 @@ class AuthenticationTest {
     driverCache.getOrCreate()
 
     PowerMockito.verifyStatic(classOf[GraphDatabase], times(1))
-    // was GraphDatabase.driver(any[URI](), ArgumentMatchers.eq(AuthTokens.bearer(token)), any())
-    GraphDatabase.driver(any[URI](), ArgumentMatchers.any[ExpirationBasedAuthTokenManager], any[Config]())
+    val managerCaptor = ArgumentCaptor.forClass(classOf[AuthTokenManager])
+    GraphDatabase.driver(any[URI](), managerCaptor.capture(), any[Config]())
+    assert(AuthTokens.bearer(token) == managerCaptor.getValue.getToken.toCompletableFuture.join())
   }
 }

spark-3/src/test/scala/org/neo4j/spark/DataSourceWriterTSE.scala

@@ -576,9 +576,11 @@ class DataSourceWriterTSE extends SparkConnectorScalaBaseTSE {
   @Test
   def `should throw an error because the node already exists`(): Unit = {
     SparkConnectorScalaSuiteIT.session()
-      .executeWrite(tx => tx.run("CREATE CONSTRAINT person_surname FOR (p:Person) REQUIRE p.surname IS UNIQUE"))
+      .executeWriteWithoutResult(tx =>
+        tx.run("CREATE CONSTRAINT person_surname FOR (p:Person) REQUIRE p.surname IS UNIQUE")
+      )
     SparkConnectorScalaSuiteIT.session()
-      .executeWrite(tx => tx.run("CREATE (p:Person{name: 'Andrea', surname: 'Santurbano'})"))
+      .executeWriteWithoutResult(tx => tx.run("CREATE (p:Person{name: 'Andrea', surname: 'Santurbano'})"))
 
     val ds = Seq(SimplePerson("Andrea", "Santurbano")).toDS()
 
@@ -602,16 +604,18 @@ class DataSourceWriterTSE extends SparkConnectorScalaBaseTSE {
       }
     } finally {
       SparkConnectorScalaSuiteIT.session()
-        .executeWrite(tx => tx.run("DROP CONSTRAINT person_surname"))
+        .executeWriteWithoutResult(tx => tx.run("DROP CONSTRAINT person_surname"))
     }
   }
 
   @Test
   def `should update the node that already exists`(): Unit = {
     SparkConnectorScalaSuiteIT.session()
-      .executeWrite(tx => tx.run("CREATE CONSTRAINT person_surname FOR (p:Person) REQUIRE p.surname IS UNIQUE"))
+      .executeWriteWithoutResult(tx =>
+        tx.run("CREATE CONSTRAINT person_surname FOR (p:Person) REQUIRE p.surname IS UNIQUE")
+      )
     SparkConnectorScalaSuiteIT.session()
-      .executeWrite(tx => tx.run("CREATE (p:Person{name: 'Federico', surname: 'Santurbano'})"))
+      .executeWriteWithoutResult(tx => tx.run("CREATE (p:Person{name: 'Federico', surname: 'Santurbano'})"))
 
     val ds = Seq(SimplePerson("Andrea", "Santurbano")).toDS()
 
@@ -635,7 +639,7 @@ class DataSourceWriterTSE extends SparkConnectorScalaBaseTSE {
     assertEquals("Andrea", nodeList.head.get("n").asNode().get("name").asString())
 
     SparkConnectorScalaSuiteIT.session()
-      .executeWrite(tx => tx.run("DROP CONSTRAINT person_surname"))
+      .executeWriteWithoutResult(tx => tx.run("DROP CONSTRAINT person_surname"))
   }
 
   @Test
@@ -766,7 +770,9 @@ class DataSourceWriterTSE extends SparkConnectorScalaBaseTSE {
   @Test
   def `should handle unusual column names`(): Unit = {
     SparkConnectorScalaSuiteIT.session()
-      .executeWrite(tx => tx.run("CREATE CONSTRAINT instrument_name FOR (i:Instrument) REQUIRE i.name IS UNIQUE"))
+      .executeWriteWithoutResult(tx =>
+        tx.run("CREATE CONSTRAINT instrument_name FOR (i:Instrument) REQUIRE i.name IS UNIQUE")
+      )
 
     val musicDf = Seq(
       (12, "John Bonham", "Drums", "f``````oo"),
@@ -791,7 +797,7 @@ class DataSourceWriterTSE extends SparkConnectorScalaBaseTSE {
       .save()
 
     SparkConnectorScalaSuiteIT.session()
-      .executeWrite(tx => tx.run("DROP CONSTRAINT instrument_name"))
+      .executeWriteWithoutResult(tx => tx.run("DROP CONSTRAINT instrument_name"))
 
     val musicDfCheck = ss.read.format(classOf[DataSource].getName)
       .option("url", SparkConnectorScalaSuiteIT.server.getBoltUrl)