Add bind and configure service

sdk-common/src/main/java/dev/restate/sdk/endpoint/Endpoint.java

@@ -13,6 +13,7 @@
 import dev.restate.sdk.endpoint.definition.ServiceDefinitionFactories;
 import io.opentelemetry.api.OpenTelemetry;
 import java.util.*;
+import java.util.function.Consumer;
 import java.util.function.Function;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
@@ -46,11 +47,14 @@ public static class Builder {
      * Add a Restate service to the endpoint. This will automatically discover the generated factory
      * based on the class name.
      *
+     * <p>If you want to modify some of the service definition options, such as documentation,
+     * inactivity timeout, and so on, use {@link #bind(Object, Consumer)} instead.
+     *
      * <p>You can also manually instantiate the {@link ServiceDefinition} using {@link
      * #bind(ServiceDefinition)}.
      */
     public Builder bind(Object service) {
-      return this.bind(ServiceDefinitionFactories.discover(service).create(service, null));
+      return this.bind(service, ignored -> {});
     }
 
     /**
@@ -61,10 +65,45 @@ public Builder bind(Object service) {
      * <p>Look at the respective documentations of the HandlerRunner class in the Java or in the
      * Kotlin module.
      *
+     * <p>If you want to modify some of the service definition options, such as documentation,
+     * inactivity timeout, and so on, use {@link #bind(Object, HandlerRunner.Options, Consumer)}
+     * instead.
+     *
      * @see #bind(Object)
      */
     public Builder bind(Object service, HandlerRunner.Options options) {
-      return this.bind(ServiceDefinitionFactories.discover(service).create(service, options));
+      return this.bind(service, options, ignored -> {});
+    }
+
+    /**
+     * Same as {@link #bind(Object)} but allows to configure the {@link ServiceDefinition} before
+     * binding it.
+     *
+     * @see #bind(Object)
+     * @see ServiceDefinition.Configurator
+     */
+    public Builder bind(Object service, Consumer<ServiceDefinition.Configurator> configurator) {
+      return this.bind(
+          ServiceDefinitionFactories.discover(service)
+              .create(service, null)
+              .configure(configurator));
+    }
+
+    /**
+     * Same as {@link #bind(Object, HandlerRunner.Options)} but allows to configure the {@link
+     * ServiceDefinition} before binding it.
+     *
+     * @see #bind(Object, HandlerRunner.Options)
+     * @see ServiceDefinition.Configurator
+     */
+    public Builder bind(
+        Object service,
+        HandlerRunner.Options options,
+        Consumer<ServiceDefinition.Configurator> configurator) {
+      return this.bind(
+          ServiceDefinitionFactories.discover(service)
+              .create(service, options)
+              .configure(configurator));
     }
 
     /** Add a manual {@link ServiceDefinition} to the endpoint. */
@@ -153,6 +192,23 @@ public static Builder bind(Object service, HandlerRunner.Options options) {
     return new Builder().bind(service, options);
   }
 
+  /**
+   * @see Builder#bind(Object, Consumer)
+   */
+  public static Builder bind(Object object, Consumer<ServiceDefinition.Configurator> configurator) {
+    return new Builder().bind(object, configurator);
+  }
+
+  /**
+   * @see Builder#bind(Object, HandlerRunner.Options, Consumer)
+   */
+  public static Builder bind(
+      Object service,
+      HandlerRunner.Options options,
+      Consumer<ServiceDefinition.Configurator> configurator) {
+    return new Builder().bind(service, options, configurator);
+  }
+
   /**
    * @see Builder#bind(ServiceDefinition)
    */

sdk-common/src/main/java/dev/restate/sdk/endpoint/definition/HandlerDefinition.java

@@ -10,7 +10,9 @@
 
 import dev.restate.serde.Serde;
 import java.util.Collections;
+import java.util.HashMap;
 import java.util.Map;
+import java.util.function.Consumer;
 import org.jspecify.annotations.Nullable;
 
 public final class HandlerDefinition<REQ, RES> {
@@ -111,6 +113,83 @@ public HandlerDefinition<REQ, RES> withMetadata(Map<String, String> metadata) {
         runner);
   }
 
+  public HandlerDefinition<REQ, RES> configure(
+      Consumer<HandlerDefinition.Configurator> configurator) {
+    HandlerDefinition.Configurator configuratorObj =
+        new HandlerDefinition.Configurator(acceptContentType, documentation, metadata);
+    configurator.accept(configuratorObj);
+
+    return new HandlerDefinition<>(
+        name,
+        handlerType,
+        configuratorObj.acceptContentType,
+        requestSerde,
+        responseSerde,
+        configuratorObj.documentation,
+        configuratorObj.metadata,
+        runner);
+  }
+
+  public static final class Configurator {
+
+    private @Nullable String acceptContentType;
+    private @Nullable String documentation;
+    private Map<String, String> metadata;
+
+    public Configurator(
+        @Nullable String acceptContentType,
+        @Nullable String documentation,
+        Map<String, String> metadata) {
+      this.acceptContentType = acceptContentType;
+      this.documentation = documentation;
+      this.metadata = new HashMap<>(metadata);
+    }
+
+    public @Nullable String getAcceptContentType() {
+      return acceptContentType;
+    }
+
+    public void setAcceptContentType(@Nullable String acceptContentType) {
+      this.acceptContentType = acceptContentType;
+    }
+
+    public Configurator acceptContentType(@Nullable String acceptContentType) {
+      this.setAcceptContentType(acceptContentType);
+      return this;
+    }
+
+    public @Nullable String getDocumentation() {
+      return documentation;
+    }
+
+    public void setDocumentation(@Nullable String documentation) {
+      this.documentation = documentation;
+    }
+
+    public Configurator documentation(@Nullable String documentation) {
+      this.setDocumentation(documentation);
+      return this;
+    }
+
+    public Map<String, String> getMetadata() {
+      return metadata;
+    }
+
+    public void setMetadata(Map<String, String> metadata) {
+      this.metadata = metadata;
+    }
+
+    public Configurator addMetadata(String key, String value) {
+      this.metadata.put(key, value);
+      return this;
+    }
+
+    public Configurator metadata(Map<String, String> metadata) {
+      this.setMetadata(metadata);
+      return this;
+    }
+  }
+
   public static <REQ, RES> HandlerDefinition<REQ, RES> of(
       String handler,
       HandlerType handlerType,

sdk-common/src/main/java/dev/restate/sdk/endpoint/definition/ServiceDefinition.java

@@ -9,6 +9,7 @@
 package dev.restate.sdk.endpoint.definition;
 
 import java.util.*;
+import java.util.function.Consumer;
 import java.util.function.Function;
 import java.util.stream.Collectors;
 import org.jspecify.annotations.Nullable;
@@ -66,6 +67,74 @@ public ServiceDefinition withMetadata(Map<String, String> metadata) {
     return new ServiceDefinition(serviceName, serviceType, handlers, documentation, metadata);
   }
 
+  public ServiceDefinition configure(Consumer<Configurator> configurator) {
+    Configurator configuratorObj = new Configurator(handlers, documentation, metadata);
+    configurator.accept(configuratorObj);
+
+    return new ServiceDefinition(
+        serviceName,
+        serviceType,
+        configuratorObj.handlers,
+        configuratorObj.documentation,
+        configuratorObj.metadata);
+  }
+
+  public static final class Configurator {
+
+    private Map<String, HandlerDefinition<?, ?>> handlers;
+    private @Nullable String documentation;
+    private Map<String, String> metadata;
+
+    private Configurator(
+        Map<String, HandlerDefinition<?, ?>> handlers,
+        @Nullable String documentation,
+        Map<String, String> metadata) {
+      this.handlers = new HashMap<>(handlers);
+      this.documentation = documentation;
+      this.metadata = new HashMap<>(metadata);
+    }
+
+    public @Nullable String getDocumentation() {
+      return documentation;
+    }
+
+    public void setDocumentation(@Nullable String documentation) {
+      this.documentation = documentation;
+    }
+
+    public Configurator documentation(@Nullable String documentation) {
+      this.setDocumentation(documentation);
+      return this;
+    }
+
+    public Map<String, String> getMetadata() {
+      return metadata;
+    }
+
+    public void setMetadata(Map<String, String> metadata) {
+      this.metadata = metadata;
+    }
+
+    public Configurator addMetadata(String key, String value) {
+      this.metadata.put(key, value);
+      return this;
+    }
+
+    public Configurator metadata(Map<String, String> metadata) {
+      this.setMetadata(metadata);
+      return this;
+    }
+
+    public Configurator configureHandler(
+        String handlerName, Consumer<HandlerDefinition.Configurator> configurator) {
+      if (!handlers.containsKey(handlerName)) {
+        throw new IllegalArgumentException("Handler " + handlerName + " not found");
+      }
+      handlers.computeIfPresent(handlerName, (k, v) -> v.configure(configurator));
+      return this;
+    }
+  }
+
   @Override
   public boolean equals(Object object) {
     if (this == object) return true;

sdk-core/src/test/java/dev/restate/sdk/core/AssertUtils.java

@@ -83,10 +83,18 @@ public static EndpointManifestSchemaAssert assertThatDiscovery(Object... service
       builder.bind(svc);
     }
 
+    return assertThatDiscovery(builder);
+  }
+
+  public static EndpointManifestSchemaAssert assertThatDiscovery(Endpoint.Builder builder) {
+    return assertThatDiscovery(builder.build());
+  }
+
+  public static EndpointManifestSchemaAssert assertThatDiscovery(Endpoint endpoint) {
     return new EndpointManifestSchemaAssert(
         new EndpointManifest(
                 EndpointManifestSchema.ProtocolMode.BIDI_STREAM,
-                builder.build().getServiceDefinitions(),
+                endpoint.getServiceDefinitions(),
                 true)
             .manifest(),
         EndpointManifestSchemaAssert.class);

sdk-core/src/test/java/dev/restate/sdk/core/javaapi/CodegenDiscoveryTest.java

@@ -16,6 +16,7 @@
 import dev.restate.sdk.core.generated.manifest.Input;
 import dev.restate.sdk.core.generated.manifest.Output;
 import dev.restate.sdk.core.generated.manifest.Service;
+import dev.restate.sdk.endpoint.Endpoint;
 import org.junit.jupiter.api.Test;
 
 public class CodegenDiscoveryTest {
@@ -66,4 +67,20 @@ void workflowType() {
         .extractingHandler("run")
         .returns(Handler.Ty.WORKFLOW, Handler::getTy);
   }
+
+  @Test
+  void usingTransformer() {
+    assertThatDiscovery(
+            Endpoint.bind(
+                new CodegenTest.RawInputOutput(),
+                sd ->
+                    sd.documentation("My service documentation")
+                        .configureHandler(
+                            "rawInputWithCustomCt",
+                            hd -> hd.documentation("My handler documentation"))))
+        .extractingService("RawInputOutput")
+        .returns("My service documentation", Service::getDocumentation)
+        .extractingHandler("rawInputWithCustomCt")
+        .returns("My handler documentation", Handler::getDocumentation);
+  }
 }

sdk-core/src/test/kotlin/dev/restate/sdk/core/kotlinapi/CodegenDiscoveryTest.kt

@@ -8,12 +8,12 @@
 // https://github.com/restatedev/sdk-java/blob/main/LICENSE
 package dev.restate.sdk.core.kotlinapi
 
-import dev.restate.sdk.Context
 import dev.restate.sdk.core.AssertUtils.assertThatDiscovery
 import dev.restate.sdk.core.generated.manifest.Handler
 import dev.restate.sdk.core.generated.manifest.Input
 import dev.restate.sdk.core.generated.manifest.Output
 import dev.restate.sdk.core.generated.manifest.Service
+import dev.restate.sdk.kotlin.endpoint.*
 import org.assertj.core.api.Assertions
 import org.assertj.core.api.InstanceOfAssertFactories.type
 import org.junit.jupiter.api.Test
@@ -72,4 +72,21 @@ class CodegenDiscoveryTest {
         .extractingHandler("run")
         .returns(Handler.Ty.WORKFLOW) { obj -> obj.ty }
   }
+
+  @Test
+  fun usingTransformer() {
+    assertThatDiscovery(
+            endpoint {
+              bind(CodegenTest.RawInputOutput()) {
+                it.documentation = "My service documentation"
+                it.configureHandler("rawInputWithCustomCt") {
+                  it.documentation = "My handler documentation"
+                }
+              }
+            })
+        .extractingService("RawInputOutput")
+        .returns("My service documentation", Service::getDocumentation)
+        .extractingHandler("rawInputWithCustomCt")
+        .returns("My handler documentation", Handler::getDocumentation)
+  }
 }