feat(jdbc): implementation of WorkerTaskExecutionStorage

cli/src/main/resources/application.yml

@@ -227,6 +227,9 @@ kestra:
       multipleconditions:
         table: "multipleconditions"
         cls: io.kestra.core.models.triggers.multipleflows.MultipleConditionWindow
+      workertaskexecutions:
+        table: "workertaskexecutions"
+        cls: io.kestra.core.runners.WorkerTaskExecution
 
   elasticsearch:
     defaults:

core/src/main/java/io/kestra/core/queues/QueueService.java

@@ -7,10 +7,7 @@
 import io.kestra.core.models.templates.Template;
 import io.kestra.core.models.triggers.Trigger;
 import io.kestra.core.models.triggers.multipleflows.MultipleConditionWindow;
-import io.kestra.core.runners.WorkerInstance;
-import io.kestra.core.runners.WorkerTask;
-import io.kestra.core.runners.WorkerTaskResult;
-import io.kestra.core.runners.WorkerTaskRunning;
+import io.kestra.core.runners.*;
 
 import jakarta.inject.Singleton;
 
@@ -39,6 +36,8 @@ public String key(Object object) {
             return ((Trigger) object).uid();
         } else if (object.getClass() == MultipleConditionWindow.class) {
             return ((MultipleConditionWindow) object).uid();
+        } else if (object.getClass() == WorkerTaskExecution.class) {
+            return ((WorkerTaskExecution) object).getExecution().getId();
         } else {
             throw new IllegalArgumentException("Unknown type '" + object.getClass().getName() + "'");
         }

jdbc-mysql/src/main/java/io/kestra/runner/mysql/MysqlWorkerTaskExecutionStorage.java

@@ -0,0 +1,15 @@
+package io.kestra.runner.mysql;
+
+import io.kestra.core.runners.WorkerTaskExecution;
+import io.kestra.jdbc.runner.AbstractWorkerTaskExecutionStorage;
+import io.kestra.repository.mysql.MysqlRepository;
+import io.micronaut.context.ApplicationContext;
+import jakarta.inject.Singleton;
+
+@Singleton
+@MysqlQueueEnabled
+public class MysqlWorkerTaskExecutionStorage extends AbstractWorkerTaskExecutionStorage {
+    public MysqlWorkerTaskExecutionStorage(ApplicationContext applicationContext) {
+        super(new MysqlRepository<>(WorkerTaskExecution.class, applicationContext));
+    }
+}

jdbc-mysql/src/main/resources/migrations/mysql/V1__initial.sql

@@ -198,3 +198,9 @@ CREATE TABLE multipleconditions (
     INDEX namespace__flow_id__condition_id (namespace, flow_id, condition_id),
     INDEX start_date__end_date (start_date, end_date)
 ) ENGINE INNODB CHARACTER SET utf8mb4 COLLATE utf8mb4_unicode_ci;
+
+
+CREATE TABLE workertaskexecutions (
+    `key` VARCHAR(250) NOT NULL PRIMARY KEY,
+    `value` JSON NOT NULL
+) ENGINE INNODB CHARACTER SET utf8mb4 COLLATE utf8mb4_unicode_ci;

jdbc-mysql/src/test/java/io/kestra/runner/mysql/MysqlWorkerTaskExecutionStorageTest.java

@@ -0,0 +1,9 @@
+package io.kestra.runner.mysql;
+
+import io.kestra.core.models.triggers.multipleflows.AbstractMultipleConditionStorageTest;
+import io.kestra.jdbc.runner.AbstractWorkerTaskExecutionStorage;
+import io.kestra.jdbc.runner.AbstractWorkerTaskExecutionTest;
+
+class MysqlWorkerTaskExecutionStorageTest extends AbstractWorkerTaskExecutionTest {
+
+}

jdbc-mysql/src/test/resources/application.yml

@@ -41,3 +41,6 @@ kestra:
       multipleconditions:
         table: "multipleconditions"
         cls: io.kestra.core.models.triggers.multipleflows.MultipleConditionWindow
+      workertaskexecutions:
+        table: "workertaskexecutions"
+        cls: io.kestra.core.runners.WorkerTaskExecution

jdbc-postgres/src/main/java/io/kestra/runner/postgres/PostgresMultipleConditionStorage.java

@@ -8,7 +8,7 @@
 
 @Singleton
 @PostgresQueueEnabled
-public class PostgresMultipleConditionStorage  extends AbstractJdbcMultipleConditionStorage {
+public class PostgresMultipleConditionStorage extends AbstractJdbcMultipleConditionStorage {
     public PostgresMultipleConditionStorage(ApplicationContext applicationContext) {
         super(new PostgresRepository<>(MultipleConditionWindow.class, applicationContext));
     }

jdbc-postgres/src/main/java/io/kestra/runner/postgres/PostgresWorkerTaskExecutionStorage.java

@@ -0,0 +1,15 @@
+package io.kestra.runner.postgres;
+
+import io.kestra.core.runners.WorkerTaskExecution;
+import io.kestra.jdbc.runner.AbstractWorkerTaskExecutionStorage;
+import io.kestra.repository.postgres.PostgresRepository;
+import io.micronaut.context.ApplicationContext;
+import jakarta.inject.Singleton;
+
+@Singleton
+@PostgresQueueEnabled
+public class PostgresWorkerTaskExecutionStorage extends AbstractWorkerTaskExecutionStorage {
+    public PostgresWorkerTaskExecutionStorage(ApplicationContext applicationContext) {
+        super(new PostgresRepository<>(WorkerTaskExecution.class, applicationContext));
+    }
+}

jdbc-postgres/src/main/resources/migrations/postgres/V1__initial.sql

@@ -219,4 +219,11 @@ CREATE TABLE multipleconditions (
 );
 
 CREATE INDEX multipleconditions_namespace__flow_id__condition_id ON multipleconditions (namespace, flow_id, condition_id);
-CREATE INDEX multipleconditions_namespace__start_date__end_date ON multipleconditions (start_date, end_date);
+CREATE INDEX multipleconditions_namespace__start_date__end_date ON multipleconditions (start_date, end_date);
+
+
+/* ----------------------- workertaskexecutions ----------------------- */
+CREATE TABLE workertaskexecutions (
+    key VARCHAR(250) NOT NULL PRIMARY KEY,
+    value JSONB NOT NULL
+);

jdbc-postgres/src/test/java/io/kestra/runner/postgres/PostgresWorkerTaskExecutionStorageTest.java

@@ -0,0 +1,7 @@
+package io.kestra.runner.postgres;
+
+import io.kestra.jdbc.runner.AbstractWorkerTaskExecutionTest;
+
+class PostgresWorkerTaskExecutionStorageTest extends AbstractWorkerTaskExecutionTest {
+
+}

jdbc-postgres/src/test/resources/application.yml

@@ -45,3 +45,6 @@ kestra:
       multipleconditions:
         table: "multipleconditions"
         cls: io.kestra.core.models.triggers.multipleflows.MultipleConditionWindow
+      workertaskexecutions:
+        table: "workertaskexecutions"
+        cls: io.kestra.core.runners.WorkerTaskExecution

jdbc/src/main/java/io/kestra/jdbc/runner/AbstractWorkerTaskExecutionStorage.java

@@ -0,0 +1,55 @@
+package io.kestra.jdbc.runner;
+
+import io.kestra.core.runners.WorkerTaskExecution;
+import io.kestra.jdbc.AbstractJdbcRepository;
+import io.kestra.jdbc.repository.AbstractRepository;
+import org.jooq.DSLContext;
+import org.jooq.Field;
+import org.jooq.Record1;
+import org.jooq.SelectConditionStep;
+import org.jooq.impl.DSL;
+
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+
+public abstract class AbstractWorkerTaskExecutionStorage extends AbstractRepository {
+    protected AbstractJdbcRepository<WorkerTaskExecution> jdbcRepository;
+
+    public AbstractWorkerTaskExecutionStorage(AbstractJdbcRepository<WorkerTaskExecution> jdbcRepository) {
+        this.jdbcRepository = jdbcRepository;
+    }
+
+    public Optional<WorkerTaskExecution> get(String executionId) {
+        return this.jdbcRepository
+            .getDslContext()
+            .transactionResult(configuration -> {
+                SelectConditionStep<Record1<Object>> select = DSL
+                    .using(configuration)
+                    .select(DSL.field("value"))
+                    .from(this.jdbcRepository.getTable())
+                    .where(
+                        DSL.field(DSL.quotedName("key")).eq(executionId)
+                    );
+
+                return this.jdbcRepository.fetchOne(select);
+            });
+    }
+
+    public void save(List<WorkerTaskExecution> workerTaskExecutions) {
+        this.jdbcRepository
+            .getDslContext()
+            .transaction(configuration -> {
+                DSLContext context = DSL.using(configuration);
+
+                workerTaskExecutions.forEach(workerTaskExecution -> {
+                    Map<Field<Object>, Object> fields = this.jdbcRepository.persistFields(workerTaskExecution);
+                    this.jdbcRepository.persist(workerTaskExecution, context, fields);
+                });
+            });
+    }
+
+    public void delete(WorkerTaskExecution workerTaskExecution) {
+        this.jdbcRepository.delete(workerTaskExecution);
+    }
+}

jdbc/src/test/java/io/kestra/jdbc/runner/AbstractWorkerTaskExecutionTest.java

@@ -0,0 +1,56 @@
+package io.kestra.jdbc.runner;
+
+import io.kestra.core.models.executions.Execution;
+import io.kestra.core.models.executions.TaskRun;
+import io.kestra.core.tasks.flows.Flow;
+import io.kestra.core.runners.WorkerTaskExecution;
+import io.kestra.core.utils.IdUtils;
+import io.kestra.jdbc.JdbcTestUtils;
+import io.micronaut.test.extensions.junit5.annotation.MicronautTest;
+import jakarta.inject.Inject;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+
+import java.util.List;
+import java.util.Optional;
+
+import static org.hamcrest.MatcherAssert.assertThat;
+import static org.hamcrest.Matchers.is;
+
+@MicronautTest(transactional = false)
+public abstract class AbstractWorkerTaskExecutionTest {
+    @Inject
+    AbstractWorkerTaskExecutionStorage workerTaskExecutionStorage;
+
+    @Inject
+    JdbcTestUtils jdbcTestUtils;
+
+    @Test
+    void suite() throws Exception {
+
+        WorkerTaskExecution workerTaskExecution = WorkerTaskExecution.builder()
+            .execution(Execution.builder().id(IdUtils.create()).build())
+            .task(Flow.builder().type(Flow.class.getName()).id(IdUtils.create()).build())
+            .taskRun(TaskRun.builder().id(IdUtils.create()).build())
+            .build();
+
+        workerTaskExecutionStorage.save(List.of(workerTaskExecution));
+
+
+        Optional<WorkerTaskExecution> find = workerTaskExecutionStorage.get(workerTaskExecution.getExecution().getId());
+        assertThat(find.isPresent(), is(true));
+        assertThat(find.get().getExecution().getId(), is(workerTaskExecution.getExecution().getId()));
+
+
+        workerTaskExecutionStorage.delete(workerTaskExecution);
+
+        find = workerTaskExecutionStorage.get(workerTaskExecution.getExecution().getId());
+        assertThat(find.isPresent(), is(false));
+    }
+
+    @BeforeEach
+    protected void init() {
+        jdbcTestUtils.drop();
+        jdbcTestUtils.migrate();
+    }
+}