airbyte-internal-ops 0.9.1__py3-none-any.whl → 0.10.1__py3-none-any.whl

airbyte_ops_mcp/cloud_admin/api_client.py

@@ -1432,3 +1432,409 @@ def set_organization_connector_version_override(
         )
 
     return True
+
+
+def start_connector_rollout(
+    docker_repository: str,
+    docker_image_tag: str,
+    actor_definition_id: str,
+    updated_by: str,
+    rollout_strategy: Literal["manual", "automated", "overridden"],
+    config_api_root: str,
+    initial_rollout_pct: int | None = None,
+    final_target_rollout_pct: int | None = None,
+    client_id: str | None = None,
+    client_secret: str | None = None,
+    bearer_token: str | None = None,
+) -> dict[str, Any]:
+    """Start a connector rollout workflow.
+
+    This function calls the platform API to start a connector rollout workflow,
+    transitioning the rollout from INITIALIZED to WORKFLOW_STARTED state.
+
+    Args:
+        docker_repository: The docker repository (e.g., "airbyte/source-github")
+        docker_image_tag: The docker image tag (e.g., "1.2.0-rc.2")
+        actor_definition_id: The actor definition ID (UUID)
+        updated_by: The user ID performing the operation (UUID)
+        rollout_strategy: The rollout strategy - "manual", "automated", or "overridden"
+        config_api_root: The Config API root URL (e.g., CLOUD_CONFIG_API_ROOT)
+        initial_rollout_pct: Initial/step percentage for rollout progression (0-100).
+            For automated rollouts, this is the percentage increment per step.
+            Default is 25% if not specified.
+        final_target_rollout_pct: Maximum percentage of actors to pin (0-100).
+            The rollout will not exceed this percentage. Default is 50% if not specified.
+        client_id: The Airbyte Cloud client ID (required if no bearer_token)
+        client_secret: The Airbyte Cloud client secret (required if no bearer_token)
+        bearer_token: Pre-existing bearer token (takes precedence over client credentials)
+
+    Returns:
+        Dictionary containing the API response with the updated rollout data
+
+    Raises:
+        PyAirbyteInputError: If the API request fails
+    """
+    access_token = _get_access_token(client_id, client_secret, bearer_token)
+
+    endpoint = f"{config_api_root}/connector_rollout/manual_start"
+    payload: dict[str, Any] = {
+        "docker_repository": docker_repository,
+        "docker_image_tag": docker_image_tag,
+        "actor_definition_id": actor_definition_id,
+        "updated_by": updated_by,
+        "rollout_strategy": rollout_strategy,
+    }
+
+    if initial_rollout_pct is not None:
+        payload["initial_rollout_pct"] = initial_rollout_pct
+    if final_target_rollout_pct is not None:
+        payload["final_target_rollout_pct"] = final_target_rollout_pct
+
+    response = requests.post(
+        endpoint,
+        json=payload,
+        headers={
+            "Authorization": f"Bearer {access_token}",
+            "User-Agent": ops_constants.USER_AGENT,
+            "Content-Type": "application/json",
+        },
+        timeout=60,
+    )
+
+    if response.status_code == 404:
+        raise PyAirbyteInputError(
+            message=f"Rollout not found for {docker_repository}:{docker_image_tag}",
+            context={
+                "docker_repository": docker_repository,
+                "docker_image_tag": docker_image_tag,
+                "actor_definition_id": actor_definition_id,
+                "endpoint": endpoint,
+                "status_code": response.status_code,
+                "response": response.text,
+            },
+        )
+
+    if response.status_code == 422:
+        raise PyAirbyteInputError(
+            message=f"Invalid input for starting rollout: {response.text}",
+            context={
+                "docker_repository": docker_repository,
+                "docker_image_tag": docker_image_tag,
+                "actor_definition_id": actor_definition_id,
+                "rollout_strategy": rollout_strategy,
+                "endpoint": endpoint,
+                "status_code": response.status_code,
+                "response": response.text,
+            },
+        )
+
+    if response.status_code != 200:
+        raise PyAirbyteInputError(
+            message=f"Failed to start connector rollout: {response.status_code} {response.text}",
+            context={
+                "docker_repository": docker_repository,
+                "docker_image_tag": docker_image_tag,
+                "actor_definition_id": actor_definition_id,
+                "rollout_strategy": rollout_strategy,
+                "endpoint": endpoint,
+                "status_code": response.status_code,
+                "response": response.text,
+            },
+        )
+
+    return response.json()
+
+
+def progress_connector_rollout(
+    docker_repository: str,
+    docker_image_tag: str,
+    actor_definition_id: str,
+    rollout_id: str,
+    updated_by: str,
+    config_api_root: str,
+    target_percentage: int | None = None,
+    actor_ids: list[str] | None = None,
+    client_id: str | None = None,
+    client_secret: str | None = None,
+    bearer_token: str | None = None,
+) -> dict[str, Any]:
+    """Progress a connector rollout by pinning actors to the RC version.
+
+    This function calls the platform API to progress a connector rollout,
+    either by specifying a target percentage or specific actor IDs to pin.
+
+    Args:
+        docker_repository: The docker repository (e.g., "airbyte/source-github")
+        docker_image_tag: The docker image tag (e.g., "1.2.0-rc.2")
+        actor_definition_id: The actor definition ID (UUID)
+        rollout_id: The rollout ID (UUID)
+        updated_by: The user ID performing the operation (UUID)
+        config_api_root: The Config API root URL (e.g., CLOUD_CONFIG_API_ROOT)
+        target_percentage: Target percentage of actors to pin (1-100)
+        actor_ids: Specific actor IDs to pin (alternative to target_percentage)
+        client_id: The Airbyte Cloud client ID (required if no bearer_token)
+        client_secret: The Airbyte Cloud client secret (required if no bearer_token)
+        bearer_token: Pre-existing bearer token (takes precedence over client credentials)
+
+    Returns:
+        Dictionary containing the API response with the updated rollout data
+
+    Raises:
+        PyAirbyteInputError: If the API request fails
+    """
+    access_token = _get_access_token(client_id, client_secret, bearer_token)
+
+    endpoint = f"{config_api_root}/connector_rollout/manual_rollout"
+    payload: dict[str, Any] = {
+        "docker_repository": docker_repository,
+        "docker_image_tag": docker_image_tag,
+        "actor_definition_id": actor_definition_id,
+        "id": rollout_id,
+        "updated_by": updated_by,
+    }
+
+    if target_percentage is not None:
+        payload["target_percentage"] = target_percentage
+    if actor_ids is not None:
+        payload["actor_ids"] = actor_ids
+
+    response = requests.post(
+        endpoint,
+        json=payload,
+        headers={
+            "Authorization": f"Bearer {access_token}",
+            "User-Agent": ops_constants.USER_AGENT,
+            "Content-Type": "application/json",
+        },
+        timeout=60,
+    )
+
+    if response.status_code == 404:
+        raise PyAirbyteInputError(
+            message=f"Rollout not found: {rollout_id}",
+            context={
+                "rollout_id": rollout_id,
+                "docker_repository": docker_repository,
+                "docker_image_tag": docker_image_tag,
+                "endpoint": endpoint,
+                "status_code": response.status_code,
+                "response": response.text,
+            },
+        )
+
+    if response.status_code == 422:
+        raise PyAirbyteInputError(
+            message=f"Invalid input for rollout progression: {response.text}",
+            context={
+                "rollout_id": rollout_id,
+                "docker_repository": docker_repository,
+                "docker_image_tag": docker_image_tag,
+                "target_percentage": target_percentage,
+                "actor_ids": actor_ids,
+                "endpoint": endpoint,
+                "status_code": response.status_code,
+                "response": response.text,
+            },
+        )
+
+    if response.status_code != 200:
+        raise PyAirbyteInputError(
+            message=f"Failed to progress connector rollout: {response.status_code} {response.text}",
+            context={
+                "rollout_id": rollout_id,
+                "docker_repository": docker_repository,
+                "docker_image_tag": docker_image_tag,
+                "target_percentage": target_percentage,
+                "actor_ids": actor_ids,
+                "endpoint": endpoint,
+                "status_code": response.status_code,
+                "response": response.text,
+            },
+        )
+
+    return response.json()
+
+
+def finalize_connector_rollout(
+    docker_repository: str,
+    docker_image_tag: str,
+    actor_definition_id: str,
+    rollout_id: str,
+    updated_by: str,
+    state: Literal["succeeded", "failed_rolled_back", "canceled"],
+    config_api_root: str,
+    client_id: str | None = None,
+    client_secret: str | None = None,
+    bearer_token: str | None = None,
+    error_msg: str | None = None,
+    failed_reason: str | None = None,
+    retain_pins_on_cancellation: bool | None = None,
+) -> dict[str, Any]:
+    """Finalize a connector rollout by promoting, rolling back, or canceling.
+
+    This function calls the platform API to finalize a connector rollout.
+
+    Args:
+        docker_repository: The docker repository (e.g., "airbyte/source-github")
+        docker_image_tag: The docker image tag (e.g., "1.2.0-rc.2")
+        actor_definition_id: The actor definition ID (UUID)
+        rollout_id: The rollout ID (UUID)
+        updated_by: The user ID performing the operation (UUID)
+        state: The final state - "succeeded" (promote), "failed_rolled_back" (rollback),
+            or "canceled" (cancel)
+        config_api_root: The Config API root URL (e.g., CLOUD_CONFIG_API_ROOT)
+        client_id: The Airbyte Cloud client ID (required if no bearer_token)
+        client_secret: The Airbyte Cloud client secret (required if no bearer_token)
+        bearer_token: Pre-existing bearer token (takes precedence over client credentials)
+        error_msg: Optional error message (for failed/canceled states)
+        failed_reason: Optional failure reason (for failed/canceled states)
+        retain_pins_on_cancellation: If True, retain version pins when canceling
+
+    Returns:
+        Dictionary containing the API response with status field
+
+    Raises:
+        PyAirbyteInputError: If the API request fails
+    """
+    access_token = _get_access_token(client_id, client_secret, bearer_token)
+
+    endpoint = f"{config_api_root}/connector_rollout/manual_finalize"
+    payload: dict[str, Any] = {
+        "docker_repository": docker_repository,
+        "docker_image_tag": docker_image_tag,
+        "actor_definition_id": actor_definition_id,
+        "id": rollout_id,
+        "updated_by": updated_by,
+        "state": state,
+    }
+
+    if error_msg:
+        payload["error_msg"] = error_msg
+    if failed_reason:
+        payload["failed_reason"] = failed_reason
+    if retain_pins_on_cancellation is not None:
+        payload["retain_pins_on_cancellation"] = retain_pins_on_cancellation
+
+    response = requests.post(
+        endpoint,
+        json=payload,
+        headers={
+            "Authorization": f"Bearer {access_token}",
+            "User-Agent": ops_constants.USER_AGENT,
+            "Content-Type": "application/json",
+        },
+        timeout=60,
+    )
+
+    if response.status_code == 404:
+        raise PyAirbyteInputError(
+            message=f"Rollout not found: {rollout_id}",
+            context={
+                "rollout_id": rollout_id,
+                "docker_repository": docker_repository,
+                "docker_image_tag": docker_image_tag,
+                "endpoint": endpoint,
+                "status_code": response.status_code,
+                "response": response.text,
+            },
+        )
+
+    if response.status_code == 422:
+        raise PyAirbyteInputError(
+            message=f"Invalid input for rollout finalization: {response.text}",
+            context={
+                "rollout_id": rollout_id,
+                "docker_repository": docker_repository,
+                "docker_image_tag": docker_image_tag,
+                "state": state,
+                "endpoint": endpoint,
+                "status_code": response.status_code,
+                "response": response.text,
+            },
+        )
+
+    if response.status_code != 200:
+        raise PyAirbyteInputError(
+            message=f"Failed to finalize connector rollout: {response.status_code} {response.text}",
+            context={
+                "rollout_id": rollout_id,
+                "docker_repository": docker_repository,
+                "docker_image_tag": docker_image_tag,
+                "state": state,
+                "endpoint": endpoint,
+                "status_code": response.status_code,
+                "response": response.text,
+            },
+        )
+
+    return response.json()
+
+
+def get_actor_sync_info(
+    rollout_id: str,
+    config_api_root: str,
+    client_id: str | None = None,
+    client_secret: str | None = None,
+    bearer_token: str | None = None,
+) -> dict[str, Any]:
+    """Get actor sync info for a connector rollout.
+
+    This function calls the platform API to get sync statistics for actors
+    participating in a rollout. Unlike the SQL-based approach, this endpoint
+    filters sync stats to only include syncs that actually used the RC version
+    associated with the rollout.
+
+    Args:
+        rollout_id: The rollout ID (UUID)
+        config_api_root: The Config API root URL (e.g., CLOUD_CONFIG_API_ROOT)
+        client_id: The Airbyte Cloud client ID (required if no bearer_token)
+        client_secret: The Airbyte Cloud client secret (required if no bearer_token)
+        bearer_token: Pre-existing bearer token (takes precedence over client credentials)
+
+    Returns:
+        Dictionary containing:
+        - data.actorSelectionInfo: Info about actor selection (numActors, numPinnedToConnectorRollout, etc.)
+        - data.syncs: Dict mapping actor_id to sync stats (numConnections, numSucceeded, numFailed)
+
+    Raises:
+        PyAirbyteInputError: If the API request fails or rollout not found
+    """
+    access_token = _get_access_token(client_id, client_secret, bearer_token)
+
+    endpoint = f"{config_api_root}/connector_rollout/get_actor_sync_info"
+    payload = {"id": rollout_id}
+
+    response = requests.post(
+        endpoint,
+        json=payload,
+        headers={
+            "Authorization": f"Bearer {access_token}",
+            "User-Agent": ops_constants.USER_AGENT,
+            "Content-Type": "application/json",
+        },
+        timeout=60,
+    )
+
+    if response.status_code == 404:
+        raise PyAirbyteInputError(
+            message=f"Rollout not found: {rollout_id}",
+            context={
+                "rollout_id": rollout_id,
+                "endpoint": endpoint,
+                "status_code": response.status_code,
+                "response": response.text,
+            },
+        )
+
+    if response.status_code != 200:
+        raise PyAirbyteInputError(
+            message=f"Failed to get actor sync info: {response.status_code} {response.text}",
+            context={
+                "rollout_id": rollout_id,
+                "endpoint": endpoint,
+                "status_code": response.status_code,
+                "response": response.text,
+            },
+        )
+
+    return response.json()

airbyte_ops_mcp/cloud_admin/models.py

@@ -127,3 +127,102 @@ class OrganizationVersionOverrideResult(BaseModel):
         if self.success:
             return f"✓ {self.message}"
         return f"✗ {self.message}"
+
+
+class ConnectorRolloutStartResult(BaseModel):
+    """Result of a connector rollout start operation.
+
+    This model provides detailed information about the outcome of starting
+    a connector rollout workflow.
+    """
+
+    success: bool = Field(description="Whether the operation succeeded")
+    message: str = Field(description="Human-readable message describing the result")
+    docker_repository: str | None = Field(
+        default=None,
+        description="The docker repository (e.g., 'airbyte/source-github')",
+    )
+    docker_image_tag: str | None = Field(
+        default=None,
+        description="The docker image tag (e.g., '1.2.0-rc.2')",
+    )
+    actor_definition_id: str | None = Field(
+        default=None,
+        description="The actor definition ID (UUID)",
+    )
+    rollout_strategy: Literal["manual", "automated", "overridden"] | None = Field(
+        default=None,
+        description="The rollout strategy used",
+    )
+
+    def __str__(self) -> str:
+        """Return a string representation of the operation result."""
+        if self.success:
+            return f"OK {self.message}"
+        return f"FAILED {self.message}"
+
+
+class ConnectorRolloutProgressResult(BaseModel):
+    """Result of a connector rollout progress operation.
+
+    This model provides detailed information about the outcome of progressing
+    a connector rollout (pinning actors to the RC version).
+    """
+
+    success: bool = Field(description="Whether the operation succeeded")
+    message: str = Field(description="Human-readable message describing the result")
+    rollout_id: str | None = Field(
+        default=None,
+        description="The rollout ID that was progressed",
+    )
+    docker_repository: str | None = Field(
+        default=None,
+        description="The docker repository (e.g., 'airbyte/source-github')",
+    )
+    docker_image_tag: str | None = Field(
+        default=None,
+        description="The docker image tag (e.g., '1.2.0-rc.2')",
+    )
+    target_percentage: int | None = Field(
+        default=None,
+        description="The target percentage of actors to pin",
+    )
+
+    def __str__(self) -> str:
+        """Return a string representation of the operation result."""
+        if self.success:
+            return f"OK {self.message}"
+        return f"FAILED {self.message}"
+
+
+class ConnectorRolloutFinalizeResult(BaseModel):
+    """Result of a connector rollout finalization operation.
+
+    This model provides detailed information about the outcome of finalizing
+    a connector rollout (promote, rollback, or cancel).
+    """
+
+    success: bool = Field(description="Whether the operation succeeded")
+    message: str = Field(description="Human-readable message describing the result")
+    rollout_id: str | None = Field(
+        default=None,
+        description="The rollout ID that was finalized",
+    )
+    docker_repository: str | None = Field(
+        default=None,
+        description="The docker repository (e.g., 'airbyte/source-github')",
+    )
+    docker_image_tag: str | None = Field(
+        default=None,
+        description="The docker image tag (e.g., '1.2.0-rc.2')",
+    )
+    state: Literal["succeeded", "failed_rolled_back", "canceled"] | None = Field(
+        default=None,
+        description="The final state of the rollout",
+    )
+
+    def __str__(self) -> str:
+        """Return a string representation of the operation result."""
+        if self.success:
+            return f"OK {self.message}"
+        return f"FAILED {self.message}"