feat: replicate command (#16375)

Co-authored-by: Jimmy Ma <[email protected]>

Author: benmoriceau

airbyte-api/server-api/src/main/openapi/config.yaml

@@ -6294,6 +6294,60 @@ paths:
                     type: string
                     enum: [pending, running, completed, cancelled]
 
+  /v1/commands/get:
+    post:
+      summary: Get the a command record
+      tags: [command, internal]
+      operationId: getCommand
+      requestBody:
+        content:
+          application/json:
+            schema:
+              type: object
+              title: CommandGetRequest
+              required: [id]
+              properties:
+                id:
+                  type: string
+      responses:
+        "200":
+          description: Success
+          content:
+            application/json:
+              schema:
+                type: object
+                title: CommandGetResponse
+                required:
+                  - id
+                  - workload_id
+                  - command_input
+                  - command_type
+                  - workspace_id
+                  - organization_id
+                  - created_at
+                  - updated_at
+                properties:
+                  id:
+                    type: string
+                  workload_id:
+                    type: string
+                  command_type:
+                    type: string
+                  command_input:
+                    type: object
+                  workspace_id:
+                    type: string
+                    format: uuid
+                  organization_id:
+                    type: string
+                    format: uuid
+                  created_at:
+                    type: string
+                    format: date-time
+                  updated_at:
+                    type: string
+                    format: date-time
+
   /v1/commands/cancel:
     post:
       summary: Get the status of command

airbyte-commons-temporal/src/main/kotlin/io/airbyte/commons/temporal/scheduling/ConnectorCommandWorkflow.kt

@@ -7,6 +7,7 @@ package io.airbyte.commons.temporal.scheduling
 import com.fasterxml.jackson.annotation.JsonSubTypes
 import com.fasterxml.jackson.annotation.JsonTypeInfo
 import com.fasterxml.jackson.databind.annotation.JsonDeserialize
+import io.airbyte.config.CatalogDiff
 import io.airbyte.config.ConnectorJobOutput
 import io.airbyte.config.StandardCheckConnectionInput
 import io.airbyte.config.StandardDiscoverCatalogInput
@@ -28,6 +29,7 @@ import java.util.UUID
   JsonSubTypes.Type(value = DiscoverCommandInput::class, name = ConnectorCommandInput.DISCOVER),
   JsonSubTypes.Type(value = DiscoverCommandApiInput::class, name = ConnectorCommandInput.DISCOVER_COMMAND),
   JsonSubTypes.Type(value = SpecCommandInput::class, name = ConnectorCommandInput.SPEC),
+  JsonSubTypes.Type(value = ReplicationCommandApiInput::class, name = ConnectorCommandInput.REPLICATION_COMMAND),
 )
 sealed interface ConnectorCommandInput {
   companion object {
@@ -36,6 +38,7 @@ sealed interface ConnectorCommandInput {
     const val DISCOVER = "discover"
     const val DISCOVER_COMMAND = "discover_command"
     const val SPEC = "spec"
+    const val REPLICATION_COMMAND = "replication_command"
   }
 
   val type: String
@@ -281,6 +284,57 @@ data class SpecCommandInput(
     }
 }
 
+@JsonDeserialize(builder = ReplicationCommandApiInput.Builder::class)
+data class ReplicationCommandApiInput(
+  val input: ReplicationApiInput,
+) : ConnectorCommandInput {
+  override val type: String = ConnectorCommandInput.REPLICATION_COMMAND
+
+  // This is duplicated of io.airbyte.workers.model.CheckConnectionInput to avoid dependency hell
+  @JsonDeserialize(builder = ReplicationApiInput.Builder::class)
+  data class ReplicationApiInput(
+    val connectionId: UUID,
+    val jobId: String,
+    val attemptId: Long,
+    val appliedCatalogDiff: CatalogDiff?,
+  ) {
+    class Builder
+      @JvmOverloads
+      constructor(
+        var connectionId: UUID? = null,
+        var jobId: String? = null,
+        var attemptId: Long? = null,
+        var appliedCatalogDiff: CatalogDiff? = null,
+      ) {
+        fun actorId(connectionId: UUID) = apply { this.connectionId = connectionId }
+
+        fun jobId(jobId: String) = apply { this.jobId = jobId }
+
+        fun attemptId(attemptId: Long) = apply { this.attemptId = attemptId }
+
+        fun appliedCatalogDiff(appliedCatalogDiff: CatalogDiff) = apply { this.appliedCatalogDiff = appliedCatalogDiff }
+
+        fun build() =
+          ReplicationApiInput(
+            connectionId = connectionId ?: throw IllegalArgumentException("actorId must be specified"),
+            jobId = jobId ?: throw IllegalArgumentException("jobId must be specified"),
+            attemptId = attemptId ?: throw IllegalArgumentException("attemptId must be specified"),
+            appliedCatalogDiff = appliedCatalogDiff,
+          )
+      }
+  }
+
+  class Builder
+    @JvmOverloads
+    constructor(
+      var input: ReplicationApiInput? = null,
+    ) {
+      fun input(input: ReplicationApiInput) = apply { this.input = input }
+
+      fun build() = ReplicationCommandApiInput(input = input ?: throw IllegalArgumentException("input must be specified"))
+    }
+}
+
 @WorkflowInterface
 interface ConnectorCommandWorkflow {
   @WorkflowMethod

airbyte-commons-temporal/src/test/kotlin/io/airbyte/commons/temporal/scheduling/ConnectorCommandWorkflowTest.kt

@@ -0,0 +1,38 @@
+/*
+ * Copyright (c) 2020-2025 Airbyte, Inc., all rights reserved.
+ */
+
+package io.airbyte.commons.temporal.scheduling
+
+import io.airbyte.commons.json.Jsons
+import io.airbyte.config.CatalogDiff
+import org.junit.jupiter.api.Assertions
+import org.junit.jupiter.api.Test
+import java.util.UUID
+
+class ConnectorCommandWorkflowTest {
+  @Test
+  fun `appliedCatalogDiff sets the appliedCatalogDiff property`() {
+    val replicationApiInput =
+      ReplicationCommandApiInput.ReplicationApiInput(
+        connectionId = UUID.randomUUID(),
+        jobId = "123",
+        attemptId = 1,
+        appliedCatalogDiff =
+          CatalogDiff().withAdditionalProperty(
+            "key",
+            "value",
+          ),
+      )
+
+    val replicationCommandApiInput =
+      ReplicationCommandApiInput(
+        input = replicationApiInput,
+      )
+
+    val serializedApiInput = Jsons.serialize(replicationCommandApiInput)
+    val deserializedApiInput = Jsons.deserialize(serializedApiInput, ReplicationCommandApiInput::class.java)
+
+    Assertions.assertEquals(replicationCommandApiInput, deserializedApiInput)
+  }
+}

airbyte-config/config-models/src/main/resources/types/ConnectorJobOutput.yaml

@@ -14,6 +14,7 @@ properties:
       - checkConnection
       - discoverCatalogId
       - spec
+      - replicate
   checkConnection:
     "$ref": StandardCheckConnectionOutput.yaml
   discoverCatalogId:
@@ -26,5 +27,7 @@ properties:
     description: A boolean indicating whether the configuration was updated during the job, e.g. if an AirbyteConfigControlMessage was received.
     type: boolean
     default: false
+  replicate:
+    "$ref": StandardSyncOutput.yaml
   failureReason:
     "$ref": FailureReason.yaml

airbyte-connector-sidecar/src/main/kotlin/io/airbyte/connectorSidecar/ConnectorMessageProcessor.kt

@@ -315,6 +315,7 @@ class ConnectorMessageProcessor(
         ConnectorJobOutput.OutputType.SPEC -> ConnectorCommand.SPEC
         ConnectorJobOutput.OutputType.CHECK_CONNECTION -> ConnectorCommand.CHECK
         ConnectorJobOutput.OutputType.DISCOVER_CATALOG_ID -> ConnectorCommand.DISCOVER
+        ConnectorJobOutput.OutputType.REPLICATE -> throw IllegalStateException("Cannot get connector command from output type $outputType")
       }
 
     fun getMessagesByType(

airbyte-server/src/main/kotlin/io/airbyte/server/apis/controllers/CommandApiController.kt

@@ -9,6 +9,8 @@ import io.airbyte.api.model.generated.CancelCommandRequest
 import io.airbyte.api.model.generated.CancelCommandResponse
 import io.airbyte.api.model.generated.CheckCommandOutputRequest
 import io.airbyte.api.model.generated.CheckCommandOutputResponse
+import io.airbyte.api.model.generated.CommandGetRequest
+import io.airbyte.api.model.generated.CommandGetResponse
 import io.airbyte.api.model.generated.CommandStatusRequest
 import io.airbyte.api.model.generated.CommandStatusResponse
 import io.airbyte.api.model.generated.DiscoverCommandOutputRequest
@@ -77,6 +79,27 @@ class CommandApiController(
     }
   }
 
+  @Post("/get")
+  @Secured(AuthRoleConstants.WORKSPACE_READER)
+  @ExecuteOn(AirbyteTaskExecutors.IO)
+  override fun getCommand(
+    @Body commandGetRequest: CommandGetRequest,
+  ): CommandGetResponse {
+    val command = commandService.get(commandGetRequest.id)
+    return CommandGetResponse().apply {
+      id(commandGetRequest?.id)
+      command?.let {
+        commandType(it.commandType)
+        commandInput(it.commandInput)
+        workspaceId(it.workspaceId)
+        workloadId(it.workloadId)
+        organizationId(it.organizationId)
+        createdAt(it.createdAt)
+        updatedAt(it.updatedAt)
+      }
+    }
+  }
+
   @Post("/output/discover")
   @Secured(AuthRoleConstants.WORKSPACE_READER)
   @ExecuteOn(AirbyteTaskExecutors.IO)

airbyte-server/src/main/kotlin/io/airbyte/server/services/CommandService.kt

@@ -428,6 +428,33 @@ class CommandService(
     }
   }
 
+  data class CommandModel(
+    val id: String,
+    val workloadId: String,
+    val commandType: String,
+    val commandInput: JsonNode,
+    val workspaceId: UUID,
+    val organizationId: UUID,
+    val createdAt: OffsetDateTime,
+    val updatedAt: OffsetDateTime,
+  )
+
+  fun get(commandId: String): CommandModel? =
+    commandsRepository
+      .findById(commandId)
+      .map { command ->
+        CommandModel(
+          id = command.id,
+          workloadId = command.workloadId,
+          commandType = command.commandType,
+          commandInput = command.commandInput,
+          workspaceId = command.workspaceId,
+          organizationId = command.organizationId,
+          createdAt = command.createdAt ?: OffsetDateTime.now(),
+          updatedAt = command.updatedAt ?: OffsetDateTime.now(),
+        )
+      }.orElse(null)
+
   fun getStatus(commandId: String): CommandStatus? =
     commandsRepository
       .findById(commandId)

airbyte-worker-models/src/main/kotlin/io/airbyte/workers/models/ReplicationApiInput.kt

@@ -0,0 +1,15 @@
+/*
+ * Copyright (c) 2020-2025 Airbyte, Inc., all rights reserved.
+ */
+
+package io.airbyte.workers.models
+
+import io.airbyte.config.CatalogDiff
+import java.util.UUID
+
+data class ReplicationApiInput(
+  val connectionId: UUID,
+  val jobId: String,
+  val attemptId: Long,
+  val appliedCatalogDiff: CatalogDiff?,
+)

airbyte-workers/src/main/java/io/airbyte/workers/temporal/sync/AsyncReplicationActivityImpl.java

@@ -201,7 +201,7 @@ public StandardSyncOutput getReplicationOutput(final ReplicationActivityInput re
       final var workerAndReplicationInput = getWorkerAndReplicationInput(replicationActivityInput);
       final WorkloadApiWorker worker = workerAndReplicationInput.worker;
 
-      final var output = worker.getOutput(workloadId);
+      final ReplicationOutput output = worker.getOutput(workloadId);
       return finalizeOutput(replicationActivityInput, output);
     } catch (final Exception e) {
       ApmTraceUtils.addActualRootCauseToTrace(e);