feat(core): document metrics

kestra-io · Mar 16, 2023 · 07a9472 · 07a9472
1 parent 3ed4718
commit 07a9472
Show file tree

Hide file tree

Showing 11 changed files with 145 additions and 5 deletions.
diff --git a/core/src/main/java/io/kestra/core/docs/ClassPluginDocumentation.java b/core/src/main/java/io/kestra/core/docs/ClassPluginDocumentation.java
@@ -22,6 +22,7 @@ public class ClassPluginDocumentation<T> {
     private String docDescription;
     private String docBody;
     private List<ExampleDoc> docExamples;
+    private List<MetricDoc> docMetrics;
     private Map<String, Object> defs = new TreeMap<>();
     private Map<String, Object> inputs = new TreeMap<>();
     private Map<String, Object> outputs = new TreeMap<>();
@@ -90,6 +91,20 @@ private ClassPluginDocumentation(JsonSchemaGenerator jsonSchemaGenerator, Regist
                 .collect(Collectors.toList());
         }
 
+        if (this.propertiesSchema.containsKey("$metrics")) {
+            List<Map<String, Object>> metrics = (List<Map<String, Object>>) this.propertiesSchema.get("$metrics");
+
+            this.docMetrics = metrics
+                .stream()
+                .map(r -> new MetricDoc(
+                    (String) r.get("name"),
+                    (String) r.get("type"),
+                    (String) r.get("unit"),
+                    (String) r.get("description")
+                ))
+                .collect(Collectors.toList());
+        }
+
         if (this.propertiesSchema.containsKey("properties")) {
             this.inputs = flatten(properties(this.propertiesSchema), required(this.propertiesSchema));
         }
@@ -157,4 +172,13 @@ public static class ExampleDoc {
         String title;
         String task;
     }
+
+    @AllArgsConstructor
+    @Getter
+    public static class MetricDoc {
+        String name;
+        String type;
+        String unit;
+        String description;
+    }
 }
diff --git a/core/src/main/java/io/kestra/core/docs/JsonSchemaGenerator.java b/core/src/main/java/io/kestra/core/docs/JsonSchemaGenerator.java
@@ -264,6 +264,20 @@ public CustomDefinition provideCustomSchemaDefinition(ResolvedType javaType, Sch
                     if (examples.size() > 0) {
                         collectedTypeAttributes.set("$examples", context.getGeneratorConfig().createArrayNode().addAll(examples));
                     }
+
+                    List<ObjectNode> metrics = Arrays
+                        .stream(pluginAnnotation.metrics())
+                        .map(metric -> context.getGeneratorConfig().createObjectNode()
+                            .put("name", metric.name())
+                            .put("type", metric.type())
+                            .put("unit", metric.unit())
+                            .put("description", metric.description())
+                        )
+                        .collect(Collectors.toList());
+
+                    if (metrics.size() > 0) {
+                        collectedTypeAttributes.set("$metrics", context.getGeneratorConfig().createArrayNode().addAll(metrics));
+                    }
                 }
             });
 
@@ -493,6 +507,9 @@ private void addMainRefProperties(JsonNode mainClassDef, ObjectNode objectNode)
         if (mainClassDef.has("$examples")) {
             objectNode.set("$examples", mainClassDef.get("$examples"));
         }
+        if (mainClassDef.has("$metrics")) {
+            objectNode.set("$metrics", mainClassDef.get("$metrics"));
+        }
     }
 
     private Object buildDefaultInstance(Class<?> cls) {

diff --git a/core/src/main/java/io/kestra/core/models/annotations/Metric.java b/core/src/main/java/io/kestra/core/models/annotations/Metric.java
@@ -0,0 +1,37 @@
+package io.kestra.core.models.annotations;
+
+import java.lang.annotation.Documented;
+import java.lang.annotation.ElementType;
+import java.lang.annotation.Inherited;
+import java.lang.annotation.Repeatable;
+import java.lang.annotation.Retention;
+import java.lang.annotation.Target;
+
+import static java.lang.annotation.RetentionPolicy.RUNTIME;
+
+@Documented
+@Inherited
+@Retention(RUNTIME)
+@Target({ElementType.TYPE, ElementType.ANNOTATION_TYPE})
+@Repeatable(Metrics.class)
+public @interface Metric {
+    /**
+     * The name of the metric
+     */
+    String name();
+
+    /**
+     * The type of the metric, should be 'counter' or 'timer'.
+     */
+    String type();
+
+    /**
+     * Optional unit, can be used for counter metric to denote the unit (records, bytes, ...)
+     */
+    String unit() default "";
+
+    /**
+     * Optional description
+     */
+    String description() default "";
+}
diff --git a/core/src/main/java/io/kestra/core/models/annotations/Metrics.java b/core/src/main/java/io/kestra/core/models/annotations/Metrics.java
@@ -0,0 +1,17 @@
+package io.kestra.core.models.annotations;
+
+import java.lang.annotation.Documented;
+import java.lang.annotation.ElementType;
+import java.lang.annotation.Inherited;
+import java.lang.annotation.Retention;
+import java.lang.annotation.Target;
+
+import static java.lang.annotation.RetentionPolicy.RUNTIME;
+
+@Documented
+@Inherited
+@Retention(RUNTIME)
+@Target({ElementType.TYPE, ElementType.ANNOTATION_TYPE})
+public @interface Metrics {
+    Metric[] value();
+}
diff --git a/core/src/main/java/io/kestra/core/models/annotations/Plugin.java b/core/src/main/java/io/kestra/core/models/annotations/Plugin.java
@@ -10,4 +10,6 @@
 @Target({ElementType.TYPE, ElementType.ANNOTATION_TYPE})
 public @interface Plugin {
     Example[] examples();
+
+    Metric[] metrics() default {};
 }
diff --git a/core/src/main/java/io/kestra/core/models/executions/metrics/Counter.java b/core/src/main/java/io/kestra/core/models/executions/metrics/Counter.java
@@ -16,9 +16,10 @@
 @Getter
 @NoArgsConstructor
 public final class Counter extends AbstractMetricEntry<Double> {
+    public static final String TYPE = "counter";
     @NotNull
     @JsonInclude
-    private final String type = "counter";
+    private final String type = TYPE;
 
     @NotNull
     @EqualsAndHashCode.Exclude

diff --git a/core/src/main/java/io/kestra/core/models/executions/metrics/Timer.java b/core/src/main/java/io/kestra/core/models/executions/metrics/Timer.java
@@ -17,9 +17,11 @@
 @Getter
 @NoArgsConstructor
 public class Timer extends AbstractMetricEntry<Duration> {
+    public static final String TYPE = "timer";
+
     @NotNull
     @JsonInclude
-    private final String type = "timer";
+    private final String type = TYPE;
 
     @NotNull
     @EqualsAndHashCode.Exclude

diff --git a/core/src/main/java/io/kestra/core/tasks/debugs/Return.java b/core/src/main/java/io/kestra/core/tasks/debugs/Return.java
@@ -1,5 +1,6 @@
 package io.kestra.core.tasks.debugs;
 
+import io.kestra.core.models.annotations.Metric;
 import io.swagger.v3.oas.annotations.media.Schema;
 import lombok.*;
 import lombok.experimental.SuperBuilder;
@@ -30,6 +31,10 @@
         @Example(
             code = "format: \"{{task.id}} > {{taskrun.startDate}}\""
         )
+    },
+    metrics = {
+        @Metric(name = "length", type = Counter.TYPE),
+        @Metric(name = "duration", type = Timer.TYPE)
     }
 )
 public class Return extends Task implements RunnableTask<Return.Output> {

diff --git a/core/src/main/resources/docs/task.hbs b/core/src/main/resources/docs/task.hbs
@@ -178,3 +178,13 @@ type: "{{cls}}"
 {{/if}}
 {{/each}}
 {{/if}}
+
+{{!-- {{                 Metrics                 }} --}}
+{{#if docMetrics}}
+    ## Metrics
+    {{#each docMetrics  as | metric | ~}}
+        ### `{{name}}`
+        * **Type:** =={{ type }}== {{#if unit }} ({{ unit }}) {{/if}}
+        {{#if description }} > {{ description }} {{/if}}
+    {{/each ~}}
+{{/if}}
diff --git a/core/src/test/java/io/kestra/core/docs/DocumentationGeneratorTest.java b/core/src/test/java/io/kestra/core/docs/DocumentationGeneratorTest.java
@@ -65,14 +65,17 @@ void bash() throws IOException {
     void returnDoc() throws IOException {
         PluginScanner pluginScanner = new PluginScanner(ClassPluginDocumentationTest.class.getClassLoader());
         RegisteredPlugin scan = pluginScanner.scan();
-        Class bash = scan.findClass(Return.class.getName()).orElseThrow();
+        Class returnTask = scan.findClass(Return.class.getName()).orElseThrow();
 
-        ClassPluginDocumentation<? extends Task> doc = ClassPluginDocumentation.of(jsonSchemaGenerator, scan, bash, Task.class);
+        ClassPluginDocumentation<? extends Task> doc = ClassPluginDocumentation.of(jsonSchemaGenerator, scan, returnTask, Task.class);
 
         String render = DocumentationGenerator.render(doc);
 
         assertThat(render, containsString("Debugging task that return"));
         assertThat(render, containsString("is mostly useful"));
+        assertThat(render, containsString("## Metrics"));
+        assertThat(render, containsString("### `length`\n" + "        * **Type:** ==counter== "));
+        assertThat(render, containsString("### `duration`\n" + "        * **Type:** ==timer== "));
     }
 
     @SuppressWarnings({"rawtypes", "unchecked"})

diff --git a/core/src/test/java/io/kestra/core/docs/JsonSchemaGeneratorTest.java b/core/src/test/java/io/kestra/core/docs/JsonSchemaGeneratorTest.java
@@ -6,7 +6,7 @@
 import io.kestra.core.models.tasks.VoidOutput;
 import io.kestra.core.models.triggers.AbstractTrigger;
 import io.kestra.core.runners.RunContext;
-import io.kestra.core.tasks.scripts.ScriptOutput;
+import io.kestra.core.tasks.debugs.Return;
 import io.kestra.core.Helpers;
 import io.swagger.v3.oas.annotations.media.Schema;
 import lombok.EqualsAndHashCode;
@@ -84,6 +84,7 @@ void flow() throws URISyntaxException {
             assertThat((String) bash.get("markdownDescription"), containsString("outputFiles.first"));
 
             var bashType = definitions.get("io.kestra.core.tasks.scripts.Bash-2");
+            assertThat(bashType, is(notNullValue()));
 
             var python = definitions.get("io.kestra.core.tasks.scripts.Python-1");
             assertThat((List<String>) python.get("required"), not(contains("exitOnFailed")));
@@ -138,6 +139,27 @@ void bash() throws URISyntaxException {
         });
     }
 
+    @SuppressWarnings("unchecked")
+    @Test
+    void returnTask() throws URISyntaxException {
+        Helpers.runApplicationContext((applicationContext) -> {
+            JsonSchemaGenerator jsonSchemaGenerator = applicationContext.getBean(JsonSchemaGenerator.class);
+
+            Map<String, Object> returnSchema = jsonSchemaGenerator.schemas(Return.class);
+            var definitions = (Map<String, Map<String, Object>>) returnSchema.get("definitions");
+            var returnTask = definitions.get("io.kestra.core.tasks.debugs.Return-1");
+            var metrics = (List<Object>) returnTask.get("$metrics");
+            assertThat(metrics.size(), is(2));
+
+            var firstMetric = (Map<String, Object>) metrics.get(0);
+            assertThat(firstMetric.get("name"), is("length"));
+            assertThat(firstMetric.get("type"), is("counter"));
+            var secondMetric = (Map<String, Object>) metrics.get(1);
+            assertThat(secondMetric.get("name"), is("duration"));
+            assertThat(secondMetric.get("type"), is("timer"));
+        });
+    }
+
     @SuppressWarnings("unchecked")
     @Test
     void testEnum() {