add sink lifecyle procedures (#328)

conker84 · web-flow · commit 9297f6683778 · 2020-08-31T22:29:39.000+02:00
diff --git a/common/src/main/kotlin/streams/config/StreamsConfig.kt b/common/src/main/kotlin/streams/config/StreamsConfig.kt
@@ -28,6 +28,8 @@ data class StreamsConfig(private val log: Log, private val dbms: DatabaseManagem
         const val SINK_ENABLED = "streams.sink.enabled"
         const val SINK_ENABLED_VALUE = false
         const val DEFAULT_PATH = "."
+        const val CHECK_APOC_TIMEOUT = "check.apoc.timeout"
+        const val CHECK_APOC_INTERVAL = "check.apoc.interval"
         private var afterInitListeners = mutableListOf<((MutableMap<String, String>) -> Unit)>()
 
         fun registerListener(after: (MutableMap<String, String>) -> Unit) {
@@ -54,8 +56,8 @@ data class StreamsConfig(private val log: Log, private val dbms: DatabaseManagem
     private fun loadConfiguration() {
         val properties = neo4jConfAsProperties()
 
-        val filteredValues = filterProperties(properties,
-                { key -> !SUPPORTED_PREFIXES.find { key.toString().startsWith(it) }.isNullOrBlank() })
+        val filteredValues = filterProperties(properties)
+                { key -> !SUPPORTED_PREFIXES.find { key.toString().startsWith(it) }.isNullOrBlank() }
 
         if (log.isDebugEnabled) {
             log.debug("Neo4j Streams Global configuration from neo4j.conf file: $filteredValues")
diff --git a/common/src/main/kotlin/streams/events/StreamsPluginStatus.kt b/common/src/main/kotlin/streams/events/StreamsPluginStatus.kt
@@ -0,0 +1,3 @@
+package streams.events
+
+enum class StreamsPluginStatus { RUNNING, STOPPED, UNKNOWN }
diff --git a/common/src/main/kotlin/streams/utils/Neo4jUtils.kt b/common/src/main/kotlin/streams/utils/Neo4jUtils.kt
@@ -37,6 +37,15 @@ object Neo4jUtils {
         }
     }
 
+    fun hasApoc(db: GraphDatabaseAPI): Boolean = try {
+         db.execute("RETURN apoc.version() AS version") {
+            it.columnAs<String>("version").next()
+            true
+        }
+    } catch (e: QueryExecutionException) {
+        false
+    }
+
     fun getLogService(db: GraphDatabaseAPI): LogService {
         return db.dependencyResolver
                 .resolveDependency(LogService::class.java)
diff --git a/common/src/main/kotlin/streams/utils/StreamsUtils.kt b/common/src/main/kotlin/streams/utils/StreamsUtils.kt
@@ -1,5 +1,8 @@
 package streams.utils
 
+import kotlinx.coroutines.runBlocking
+import kotlinx.coroutines.delay
+
 object StreamsUtils {
 
     const val UNWIND: String = "UNWIND \$events AS event"
@@ -22,4 +25,14 @@ object StreamsUtils {
         }
     }
 
+    fun blockUntilTrueOrTimeout(timeout: Long, delay: Long = 1000, action: () -> Boolean): Boolean = runBlocking {
+        val start = System.currentTimeMillis()
+        var success = action()
+        while (System.currentTimeMillis() - start < timeout && !success) {
+            delay(delay)
+            success = action()
+        }
+        success
+    }
+
 }
diff --git a/common/src/test/kotlin/streams/utils/Neo4jUtilsTest.kt b/common/src/test/kotlin/streams/utils/Neo4jUtilsTest.kt
@@ -25,4 +25,9 @@ class Neo4jUtilsTest {
         assertFalse { isEnterprise }
     }
 
+    @Test
+    fun `should not have APOC`() {
+        assertFalse { Neo4jUtils.hasApoc(db) }
+    }
+
 }
diff --git a/consumer/src/main/kotlin/streams/StreamsEventSink.kt b/consumer/src/main/kotlin/streams/StreamsEventSink.kt
@@ -1,9 +1,9 @@
 package streams
 
-import org.neo4j.configuration.Config
 import org.neo4j.kernel.internal.GraphDatabaseAPI
 import org.neo4j.logging.Log
 import streams.config.StreamsConfig
+import streams.events.StreamsPluginStatus
 
 abstract class StreamsEventSink(private val config: StreamsConfig,
                                 private val queryExecution: StreamsEventSinkQueryExecution,
@@ -24,6 +24,8 @@ abstract class StreamsEventSink(private val config: StreamsConfig,
 
     open fun printInvalidTopics() {}
 
+    abstract fun status(): StreamsPluginStatus
+
 }
 
 object StreamsEventSinkFactory {
diff --git a/consumer/src/main/kotlin/streams/StreamsEventSinkExtensionFactory.kt b/consumer/src/main/kotlin/streams/StreamsEventSinkExtensionFactory.kt
@@ -1,5 +1,8 @@
 package streams
 
+import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.GlobalScope
+import kotlinx.coroutines.launch
 import org.neo4j.dbms.api.DatabaseManagementService
 import org.neo4j.kernel.availability.AvailabilityGuard
 import org.neo4j.kernel.availability.AvailabilityListener
@@ -80,6 +83,7 @@ class StreamsEventSinkExtensionFactory : ExtensionFactory<StreamsEventSinkExtens
                         StreamsSinkProcedures.registerStreamsSinkConfiguration(streamsSinkConfiguration)
                         StreamsSinkProcedures.registerStreamsEventConsumerFactory(eventSink.getEventConsumerFactory())
                         StreamsSinkProcedures.registerStreamsEventSinkConfigMapper(eventSink.getEventSinkConfigMapper())
+                        StreamsSinkProcedures.registerStreamsEventSink(eventSink)
                     } catch (e: Exception) {
                         streamsLog.error("Error initializing the streaming sink:", e)
                     }
@@ -88,6 +92,28 @@ class StreamsEventSinkExtensionFactory : ExtensionFactory<StreamsEventSinkExtens
         }
 
         private fun initSinkModule(streamsTopicService: StreamsTopicService, streamsSinkConfiguration: StreamsSinkConfiguration) {
+            if (streamsSinkConfiguration.checkApocTimeout > -1) {
+                GlobalScope.launch(Dispatchers.IO) {
+                    val success = StreamsUtils.blockUntilTrueOrTimeout(streamsSinkConfiguration.checkApocTimeout, streamsSinkConfiguration.checkApocInterval) {
+                        val hasApoc = Neo4jUtils.hasApoc(db)
+                        if (!hasApoc && streamsLog.isDebugEnabled) {
+                            streamsLog.debug("APOC not loaded yet, next check in ${streamsSinkConfiguration.checkApocInterval} ms")
+                        }
+                        hasApoc
+                    }
+                    if (success) {
+                        initSink(streamsTopicService, streamsSinkConfiguration)
+                    } else {
+                        streamsLog.info("Streams Sink plugin not loaded as APOC are not installed")
+                    }
+                }
+            } else {
+                initSink(streamsTopicService, streamsSinkConfiguration)
+            }
+
+        }
+
+        private fun initSink(streamsTopicService: StreamsTopicService, streamsSinkConfiguration: StreamsSinkConfiguration) {
             streamsTopicService.clearAll()
             streamsTopicService.setAll(streamsSinkConfiguration.topics)
             eventSink.start()
diff --git a/consumer/src/main/kotlin/streams/StreamsSinkConfiguration.kt b/consumer/src/main/kotlin/streams/StreamsSinkConfiguration.kt
@@ -1,6 +1,8 @@
 package streams
 
 import streams.config.StreamsConfig
+import streams.extensions.toPointCase
+import streams.serialization.JSONUtils
 import streams.service.TopicUtils
 import streams.service.TopicValidationException
 import streams.service.Topics
@@ -10,8 +12,27 @@ data class StreamsSinkConfiguration(val enabled: Boolean = StreamsConfig.SINK_EN
                                     val proceduresEnabled: Boolean = StreamsConfig.PROCEDURES_ENABLED_VALUE,
                                     val topics: Topics = Topics(),
                                     val errorConfig: Map<String,Any?> = emptyMap(),
+                                    val checkApocTimeout: Long = -1,
+                                    val checkApocInterval: Long = 1000,
                                     val sourceIdStrategyConfig: SourceIdIngestionStrategyConfig = SourceIdIngestionStrategyConfig()) {
 
+    fun asMap(): Map<String, Any?> {
+        val configMap = JSONUtils.asMap(this)
+                .filterKeys { it != "topics" && it != "enabled" && it != "proceduresEnabled" && !it.startsWith("check") }
+                .mapKeys { it.key.toPointCase() }
+                .mapKeys {
+                    when (it.key) {
+                        "error.config" -> "streams.sink.errors"
+                        "procedures.enabled" -> "streams.${it.key}"
+                        else -> if (it.key.startsWith("streams.sink")) it.key else "streams.sink.${it.key}"
+                    }
+                }
+        val topicMap = this.topics.asMap()
+                .mapKeys { it.key.key }
+        val invalidTopics = mapOf("invalid_topics" to this.topics.invalid)
+        return (configMap + topicMap + invalidTopics)
+    }
+
     companion object {
         fun from(cfg: StreamsConfig, dbName: String, invalidTopics: List<String> = emptyList()): StreamsSinkConfiguration {
             val default = StreamsSinkConfiguration()
@@ -39,10 +60,19 @@ data class StreamsSinkConfiguration(val enabled: Boolean = StreamsConfig.SINK_EN
                     .filterKeys { it.startsWith("streams.sink.error") }
                     .mapKeys { it.key.substring("streams.sink.".length) }
 
+
             return default.copy(enabled = cfg.isSinkEnabled(dbName),
                     proceduresEnabled = cfg.hasProceduresEnabled(dbName),
                     topics = topics,
                     errorConfig = errorHandler,
+                    checkApocTimeout = cfg.config.getOrDefault("streams.${StreamsConfig.CHECK_APOC_TIMEOUT}",
+                            default.checkApocTimeout)
+                            .toString()
+                            .toLong(),
+                    checkApocInterval = cfg.config.getOrDefault("streams.${StreamsConfig.CHECK_APOC_INTERVAL}",
+                            default.checkApocInterval)
+                            .toString()
+                            .toLong(),
                     sourceIdStrategyConfig = sourceIdStrategyConfig)
         }
 
diff --git a/consumer/src/main/kotlin/streams/kafka/KafkaEventSink.kt b/consumer/src/main/kotlin/streams/kafka/KafkaEventSink.kt
@@ -20,6 +20,7 @@ import streams.StreamsEventSinkQueryExecution
 import streams.StreamsSinkConfiguration
 import streams.StreamsTopicService
 import streams.config.StreamsConfig
+import streams.events.StreamsPluginStatus
 import streams.service.errors.ErrorService
 import streams.service.errors.KafkaErrorService
 import streams.utils.KafkaValidationUtils.getInvalidTopicsError
@@ -77,6 +78,12 @@ class KafkaEventSink(private val config: StreamsConfig,
     }
 
     override fun start() { // TODO move to the abstract class
+        if (StreamsPluginStatus.RUNNING == status()) {
+            if (log.isDebugEnabled) {
+                log.debug("Kafka Sink is already started.")
+            }
+            return
+        }
         val streamsConfig = StreamsSinkConfiguration.from(config, db.databaseName())
         val topics = streamsTopicService.getTopics()
         val isWriteableInstance = Neo4jUtils.isWriteableInstance(db)
@@ -137,7 +144,7 @@ class KafkaEventSink(private val config: StreamsConfig,
                         }
                         TimeUnit.SECONDS.toMillis(1)
                     } else {
-                        val timeMillis = TimeUnit.MINUTES.toMillis(5)
+                        val timeMillis = TimeUnit.MINUTES.toMillis(3)
                         if (log.isDebugEnabled) {
                             log.debug("Not in a writeable instance, new check in $timeMillis millis")
                         }
@@ -166,6 +173,11 @@ class KafkaEventSink(private val config: StreamsConfig,
             }
         }, UninitializedPropertyAccessException::class.java)
     }
+    
+    override fun status(): StreamsPluginStatus = when (this::job.isInitialized && this.job.isActive) {
+        true -> StreamsPluginStatus.RUNNING
+        else -> StreamsPluginStatus.STOPPED
+    }
 
 }
 
diff --git a/consumer/src/main/kotlin/streams/procedures/StreamsSinkProcedures.kt b/consumer/src/main/kotlin/streams/procedures/StreamsSinkProcedures.kt
@@ -1,21 +1,36 @@
 package streams.procedures
 
+import org.apache.commons.lang3.exception.ExceptionUtils
+import org.neo4j.graphdb.GraphDatabaseService
+import org.neo4j.kernel.internal.GraphDatabaseAPI
 import org.neo4j.logging.Log
-import org.neo4j.procedure.*
+import org.neo4j.procedure.Context
+import org.neo4j.procedure.Description
+import org.neo4j.procedure.Mode
+import org.neo4j.procedure.Name
+import org.neo4j.procedure.Procedure
 import streams.StreamsEventConsumer
 import streams.StreamsEventConsumerFactory
+import streams.StreamsEventSink
 import streams.StreamsEventSinkConfigMapper
 import streams.StreamsSinkConfiguration
 import streams.config.StreamsConfig
+import streams.events.StreamsPluginStatus
+import streams.utils.Neo4jUtils
+import java.util.stream.Collectors
 import java.util.stream.Stream
 
 class StreamResult(@JvmField val event: Map<String, *>)
+class KeyValueResult(@JvmField val name: String, @JvmField val value: Any?)
 
 class StreamsSinkProcedures {
 
     @JvmField @Context
     var log: Log? = null
 
+    @JvmField @Context
+    var db: GraphDatabaseService? = null
+
     @Procedure(mode = Mode.READ, name = "streams.consume")
     @Description("streams.consume(topic, {timeout: <long value>, from: <string>, groupId: <string>, commit: <boolean>, partitions:[{partition: <number>, offset: <number>}]}) " +
             "YIELD event - Allows to consume custom topics")
@@ -35,6 +50,71 @@ class StreamsSinkProcedures {
         return data.map { StreamResult(mapOf("data" to it)) }.stream()
     }
 
+    @Procedure("streams.sink.start")
+    fun sinkStart(): Stream<KeyValueResult> {
+        checkEnabled()
+        return checkLeader {
+            try {
+                streamsEventSink?.start()
+                sinkStatus()
+            } catch (e: Exception) {
+                log?.error("Cannot start the Sink because of the following exception", e)
+                Stream.concat(sinkStatus(),
+                        Stream.of(KeyValueResult("exception", ExceptionUtils.getStackTrace(e))))
+            }
+        }
+    }
+
+    @Procedure("streams.sink.stop")
+    fun sinkStop(): Stream<KeyValueResult> {
+        checkEnabled()
+        return checkLeader {
+            try {
+                streamsEventSink?.stop()
+                sinkStatus()
+            } catch (e: Exception) {
+                log?.error("Cannot stopped the Sink because of the following exception", e)
+                Stream.concat(sinkStatus(),
+                        Stream.of(KeyValueResult("exception", ExceptionUtils.getStackTrace(e))))
+            }
+        }
+    }
+
+    @Procedure("streams.sink.restart")
+    fun sinkRestart(): Stream<KeyValueResult> {
+        val stopped = sinkStop().collect(Collectors.toList())
+        val hasError = stopped.any { it.name == "exception" }
+        if (hasError) {
+            return stopped.stream()
+        }
+        return sinkStart()
+    }
+
+    @Procedure("streams.sink.config")
+    fun sinkConfig(): Stream<KeyValueResult> {
+        checkEnabled()
+        return checkLeader {
+            streamsSinkConfiguration.asMap()
+                    .entries.stream()
+                    .map { KeyValueResult(it.key, it.value) }
+        }
+    }
+
+    @Procedure("streams.sink.status")
+    fun sinkStatus(): Stream<KeyValueResult> {
+        checkEnabled()
+        return checkLeader {
+            val value = (streamsEventSink?.status() ?: StreamsPluginStatus.UNKNOWN).toString()
+            Stream.of(KeyValueResult("status", value))
+        }
+    }
+
+    private fun checkLeader(lambda: () -> Stream<KeyValueResult>): Stream<KeyValueResult> = if (Neo4jUtils.isWriteableInstance(db as GraphDatabaseAPI)) {
+        lambda()
+    } else {
+        Stream.of(KeyValueResult("error", "You can use this procedure only in the LEADER or in a single instance configuration."))
+    }
+
     private fun readData(topic: String, procedureConfig: Map<String, Any>, consumerConfig: Map<String, String>): List<Any> {
         val cfg = procedureConfig.mapValues { if (it.key != "partitions") it.value else mapOf(topic to it.value) }
         val timeout = cfg.getOrDefault("timeout", 1000).toString().toLong()
@@ -80,6 +160,7 @@ class StreamsSinkProcedures {
         private lateinit var streamsSinkConfiguration: StreamsSinkConfiguration
         private lateinit var streamsEventConsumerFactory: StreamsEventConsumerFactory
         private lateinit var streamsEventSinkConfigMapper: StreamsEventSinkConfigMapper
+        private var streamsEventSink: StreamsEventSink? = null
 
         fun registerStreamsEventSinkConfigMapper(streamsEventSinkConfigMapper: StreamsEventSinkConfigMapper) {
             this.streamsEventSinkConfigMapper = streamsEventSinkConfigMapper
@@ -92,5 +173,9 @@ class StreamsSinkProcedures {
         fun registerStreamsEventConsumerFactory(streamsEventConsumerFactory: StreamsEventConsumerFactory) {
             this.streamsEventConsumerFactory = streamsEventConsumerFactory
         }
+
+        fun registerStreamsEventSink(streamsEventSink: StreamsEventSink) {
+            this.streamsEventSink = streamsEventSink
+        }
     }
 }
diff --git a/consumer/src/test/kotlin/integrations/kafka/KafkaEventSinkSimpleTSE.kt b/consumer/src/test/kotlin/integrations/kafka/KafkaEventSinkSimpleTSE.kt
diff --git a/consumer/src/test/kotlin/integrations/kafka/KafkaStreamsSinkProceduresTSE.kt b/consumer/src/test/kotlin/integrations/kafka/KafkaStreamsSinkProceduresTSE.kt
diff --git a/consumer/src/test/kotlin/streams/StreamsSinkConfigurationTest.kt b/consumer/src/test/kotlin/streams/StreamsSinkConfigurationTest.kt
diff --git a/doc/asciidoc/consumer/configuration.adoc b/doc/asciidoc/consumer/configuration.adoc
diff --git a/doc/asciidoc/procedures/index.adoc b/doc/asciidoc/procedures/index.adoc

Original file line number	Diff line number	Diff line change
`@@ -0,0 +1,3 @@`
	`1`	`+package streams.events`
	`2`	`+`
	`3`	`+enum class StreamsPluginStatus { RUNNING, STOPPED, UNKNOWN }`
Original file line number	Diff line number	Diff line change
`@@ -25,4 +25,9 @@ class Neo4jUtilsTest {`
`25`	`25`	`assertFalse { isEnterprise }`
`26`	`26`	`}`
`27`	`27`
	`28`	`+ @Test`
	`29`	+ fun `should not have APOC`() {
	`30`	`+ assertFalse { Neo4jUtils.hasApoc(db) }`
	`31`	`+ }`
	`32`	`+`
`28`	`33`	`}`