forked from airbytehq/airbyte-platform
-
Notifications
You must be signed in to change notification settings - Fork 0
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Enable the use of temporal as a generic queue (#9571)
- Loading branch information
Showing
18 changed files
with
260 additions
and
13 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
114 changes: 114 additions & 0 deletions
114
airbyte-commons-temporal-core/src/main/java/io/airbyte/commons/temporal/queue/Internal.kt
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,114 @@ | ||
| package io.airbyte.commons.temporal.queue | ||
|
|
||
| import com.fasterxml.jackson.annotation.JsonTypeInfo | ||
| import com.fasterxml.jackson.databind.annotation.JsonDeserialize | ||
| import com.google.common.annotations.VisibleForTesting | ||
| import io.airbyte.commons.temporal.annotations.TemporalActivityStub | ||
| import io.temporal.activity.ActivityInterface | ||
| import io.temporal.activity.ActivityMethod | ||
| import io.temporal.workflow.WorkflowInterface | ||
| import io.temporal.workflow.WorkflowMethod | ||
|
|
||
| /** | ||
| * Message abstraction. | ||
| * | ||
| * We wrap the actual message to be able to pass metadata around. | ||
| */ | ||
| @JsonDeserialize(builder = Message.Builder::class) | ||
| class Message<T : Any> constructor(data: T) { | ||
| // TODO this should be a data class, however, need to make the JsonTypeInfo annotation work | ||
|
|
||
| // This enables passing T around | ||
| @JsonTypeInfo(include = JsonTypeInfo.As.WRAPPER_ARRAY, use = JsonTypeInfo.Id.CLASS, property = "@bodyClass") | ||
| val data: T | ||
|
|
||
| init { | ||
| this.data = data | ||
| } | ||
|
|
||
| /** | ||
| * Builder for messages. | ||
| * | ||
| * A builder is needed in order to have data as non-nullable because the deserializer requires a no-arg constructor. | ||
| */ | ||
| class Builder<T : Any> | ||
| @JvmOverloads | ||
| constructor(data: T? = null) { | ||
| @JsonTypeInfo(include = JsonTypeInfo.As.WRAPPER_ARRAY, use = JsonTypeInfo.Id.CLASS, property = "@bodyClass") | ||
| var data: T? | ||
|
|
||
| init { | ||
| this.data = data | ||
| } | ||
|
|
||
| fun data(data: T) = apply { this.data = data } | ||
|
|
||
| fun build() = Message(data = data!!) | ||
| } | ||
|
|
||
| override fun equals(other: Any?): Boolean { | ||
| if (this === other) return true | ||
| if (javaClass != other?.javaClass) return false | ||
|
|
||
| other as Message<*> | ||
|
|
||
| if (data != other.data) return false | ||
|
|
||
| return true | ||
| } | ||
|
|
||
| override fun hashCode(): Int { | ||
| return data.hashCode() | ||
| } | ||
| } | ||
|
|
||
| /** | ||
| * Generic queue activity. | ||
| */ | ||
| @ActivityInterface | ||
| interface QueueActivity<T : Any> { | ||
| @ActivityMethod | ||
| fun consume(message: Message<T>) | ||
| } | ||
|
|
||
| /** | ||
| * Generic temporal workflow interface for a message queue. | ||
| */ | ||
| @WorkflowInterface | ||
| interface QueueWorkflow<T : Any> { | ||
| /** | ||
| * Submits a message to the queue. | ||
| */ | ||
| @WorkflowMethod | ||
| fun publish(message: Message<T>) | ||
| } | ||
|
|
||
| /** | ||
| * Generic temporal queue activity implementation. | ||
| */ | ||
| class QueueActivityImpl<T : Any>(private val messageConsumer: MessageConsumer<T>) : QueueActivity<T> { | ||
| override fun consume(message: Message<T>) { | ||
| messageConsumer.consume(message.data) | ||
| } | ||
| } | ||
|
|
||
| /** | ||
| * Generic temporal queue workflow implementation. | ||
| * | ||
| * This is open to simplify initialization when starting a workflow because temporal requires a type reference. | ||
| */ | ||
| open class QueueWorkflowImpl<T : Any> : QueueWorkflow<T> { | ||
| /** | ||
| * The consumer activity. | ||
| * | ||
| * This is lateinit because the TemporalActivityStub will initialize the value post creation using activities | ||
| * from the dependency injection container. | ||
| */ | ||
| @VisibleForTesting | ||
| @TemporalActivityStub(activityOptionsBeanName = "queueActivityOptions") | ||
| protected lateinit var activity: QueueActivity<T> | ||
|
|
||
| override fun publish(message: Message<T>) { | ||
| activity.consume(message) | ||
| } | ||
| } |
8 changes: 8 additions & 0 deletions
8
...-commons-temporal-core/src/main/java/io/airbyte/commons/temporal/queue/MessageConsumer.kt
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,8 @@ | ||
| package io.airbyte.commons.temporal.queue | ||
|
|
||
| /** | ||
| * MessageConsumer interface for a temporal queue. | ||
| */ | ||
| interface MessageConsumer<T : Any> { | ||
| fun consume(input: T) | ||
| } |
31 changes: 31 additions & 0 deletions
31
...-temporal-core/src/main/java/io/airbyte/commons/temporal/queue/TemporalMessageProducer.kt
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,31 @@ | ||
| package io.airbyte.commons.temporal.queue | ||
|
|
||
| import io.airbyte.commons.temporal.WorkflowClientWrapped | ||
| import io.temporal.client.WorkflowOptions | ||
|
|
||
| /** | ||
| * Generic message producer for a temporal based queue. | ||
| */ | ||
| class TemporalMessageProducer<T : Any>(private val workflowClientWrapped: WorkflowClientWrapped) { | ||
| /** | ||
| * Publish a message to the subject. | ||
| */ | ||
| fun publish( | ||
| subject: String, | ||
| message: T, | ||
| ) { | ||
| doPublish<QueueWorkflow<T>>(subject, message) | ||
| } | ||
|
|
||
| // This is a workaround to get a class with a generic. | ||
| // Temporal newWorkflowStub call requires a Class<T>, reified enables this. | ||
| // This is added as a private fun because of the visibility constraint from inline on member access. | ||
| private inline fun <reified W : QueueWorkflow<T>> doPublish( | ||
| subject: String, | ||
| message: T, | ||
| ) { | ||
| val workflowOptions = WorkflowOptions.newBuilder().setTaskQueue(subject).build() | ||
| val workflow = workflowClientWrapped.newWorkflowStub(W::class.java, workflowOptions) | ||
| workflow.publish(Message(message)) | ||
| } | ||
| } |
93 changes: 93 additions & 0 deletions
93
...e-commons-temporal-core/src/test/java/io/airbyte/commons/temporal/queue/BasicQueueTest.kt
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,93 @@ | ||
| package io.airbyte.commons.temporal.queue | ||
|
|
||
| import com.fasterxml.jackson.databind.annotation.JsonDeserialize | ||
| import io.airbyte.commons.temporal.WorkflowClientWrapped | ||
| import io.airbyte.metrics.lib.MetricClient | ||
| import io.temporal.activity.ActivityOptions | ||
| import io.temporal.client.WorkflowClient | ||
| import io.temporal.testing.TestWorkflowEnvironment | ||
| import io.temporal.worker.Worker | ||
| import io.temporal.workflow.Workflow | ||
| import org.junit.jupiter.api.AfterAll | ||
| import org.junit.jupiter.api.BeforeAll | ||
| import org.junit.jupiter.api.Test | ||
| import org.mockito.Mockito.mock | ||
| import org.mockito.Mockito.spy | ||
| import org.mockito.Mockito.verify | ||
| import java.time.Duration | ||
|
|
||
| // Payload for the Queue | ||
| @JsonDeserialize(builder = TestQueueInput.Builder::class) | ||
| data class TestQueueInput(val input: String) { | ||
| // Using a builder here to prove that we can use a payload with non-nullable fields. | ||
| data class Builder(var input: String? = null) { | ||
| fun input(input: String) = apply { this.input = input } | ||
|
|
||
| fun build() = TestQueueInput(input = input!!) | ||
| } | ||
| } | ||
|
|
||
| // The actual consumer | ||
| class TestConsumer : MessageConsumer<TestQueueInput> { | ||
| override fun consume(input: TestQueueInput) { | ||
| println(input) | ||
| } | ||
| } | ||
|
|
||
| // Test implementation, this is required to map the activities and for registering an implementation with temporal. | ||
| class TestWorkflowImpl : QueueWorkflowImpl<TestQueueInput>() { | ||
| init { | ||
| initializeActivity<QueueActivity<TestQueueInput>>() | ||
| } | ||
|
|
||
| // reified is a trick to be able to retrieve the type reference on a generic class to pass it to temporal. | ||
| private inline fun <reified W : QueueActivity<TestQueueInput>> initializeActivity() { | ||
| // Initializing the activity the official temporal way to be able to run the temporal standard tests. | ||
| // For an actual workflow, we'd want to inject the activities so this init wouldn't be needed. | ||
| this.activity = | ||
| Workflow.newActivityStub( | ||
| W::class.java, | ||
| ActivityOptions.newBuilder().setStartToCloseTimeout(Duration.ofSeconds(10)).build(), | ||
| ) | ||
| } | ||
| } | ||
|
|
||
| class BasicQueueTest { | ||
| companion object { | ||
| val QUEUE_NAME = "testQueue" | ||
|
|
||
| lateinit var activity: QueueActivityImpl<TestQueueInput> | ||
|
|
||
| lateinit var testEnv: TestWorkflowEnvironment | ||
| lateinit var worker: Worker | ||
| lateinit var client: WorkflowClient | ||
|
|
||
| @JvmStatic | ||
| @BeforeAll | ||
| fun setUp() { | ||
| testEnv = TestWorkflowEnvironment.newInstance() | ||
| worker = testEnv.newWorker(QUEUE_NAME) | ||
| worker.registerWorkflowImplementationTypes(TestWorkflowImpl::class.java) | ||
| client = testEnv.workflowClient | ||
|
|
||
| activity = spy(QueueActivityImpl(TestConsumer())) | ||
| worker.registerActivitiesImplementations(activity) | ||
| testEnv.start() | ||
| } | ||
|
|
||
| @JvmStatic | ||
| @AfterAll | ||
| fun tearDown() { | ||
| testEnv.close() | ||
| } | ||
| } | ||
|
|
||
| @Test | ||
| fun testRoundTrip() { | ||
| val producer = TemporalMessageProducer<TestQueueInput>(WorkflowClientWrapped(client, mock(MetricClient::class.java))) | ||
| val message = TestQueueInput("boom!") | ||
| producer.publish(QUEUE_NAME, message) | ||
|
|
||
| verify(activity).consume(Message(message)) | ||
| } | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.