Take draft status into account when importing

Author: bglearning

deepset_cloud_sdk/_service/pipeline_service.py

@@ -374,32 +374,59 @@ async def _create_index(self, name: str, pipeline_yaml: str) -> Response:
         )
 
     async def _overwrite_pipeline(self, name: str, pipeline_yaml: str) -> Response:
-        """Overwrite a pipeline in deepset AI Platform by creating a new version.
-        If creating a new version fails (e.g. pipeline doesn't exist), create the
-        pipeline instead.
+        """Overwrite a pipeline in deepset AI Platform.
+
+        Behavior:
+        - First try to fetch the latest version.
+        - If the pipeline doesn't exist (404), create it instead.
+        - If the latest version is a draft (is_draft == True), PATCH that version.
+        - Otherwise, create a new version via POST /pipelines/{name}/versions.
 
         :param name: Name of the pipeline.
         :param pipeline_yaml: Generated pipeline YAML string.
         """
-        # First try to create a new version of the existing pipeline
-        version_response = await self._api.post(
+        # First get the (last) version id if available
+        version_response = await self._api.get(
             workspace_name=self._workspace_name,
             endpoint=f"pipelines/{name}/versions",
-            json={"config_yaml": pipeline_yaml},
         )
 
-        if version_response.status_code == HTTPStatus.CREATED:
-            logger.debug("Created new version for pipeline %s.", name)
-            return version_response
-        # If creating a version fails, assume the pipeline doesn't exist and create it
-        logger.debug(
-            "Failed to create new version for pipeline %s (status %s). "
-            "Assuming pipeline does not exist and creating it instead.",
-            name,
-            version_response.status_code,
-        )
+        # If pipeline doesn't exist (404), create it instead
+        if version_response.status_code == HTTPStatus.NOT_FOUND:
+            logger.debug("Pipeline %s not found, creating new pipeline.", name)
+            response = await self._create_pipeline(name=name, pipeline_yaml=pipeline_yaml)
+            return response
+
+        version_body = version_response.json()
+        latest_version = version_body["data"][0]
+        version_id = latest_version["version_id"]
+        is_draft = latest_version.get("is_draft", False)
+
+        if is_draft:
+            # If the latest version is a draft, patch that version
+            logger.debug(
+                "Latest version %s of pipeline %s is a draft, patching existing version.",
+                version_id,
+                name,
+            )
+            response = await self._api.patch(
+                workspace_name=self._workspace_name,
+                endpoint=f"pipelines/{name}/versions/{version_id}",
+                json={"config_yaml": pipeline_yaml},
+            )
+        else:
+            # Otherwise, create a new version
+            logger.debug(
+                "Latest version %s of pipeline %s is not a draft, creating new version.",
+                version_id,
+                name,
+            )
+            response = await self._api.post(
+                workspace_name=self._workspace_name,
+                endpoint=f"pipelines/{name}/versions",
+                json={"config_yaml": pipeline_yaml, "is_draft": True},
+            )
 
-        response = await self._create_pipeline(name=name, pipeline_yaml=pipeline_yaml)
         return response
 
     async def _create_pipeline(self, name: str, pipeline_yaml: str) -> Response:

tests/integration/workflows/test_integration_pipeline_client.py

@@ -484,17 +484,23 @@ async def test_import_pipeline_with_overwrite_fallback_to_create_async(
     ) -> None:
         """Test overwriting a pipeline that doesn't exist, falling back to creation."""
         # Mock validation success
-        validation_route = respx.post("https://test-api-url.com/workspaces/test-workspace/pipeline_validations").mock(
+        validation_route = respx.post(
+            "https://test-api-url.com/workspaces/test-workspace/pipeline_validations"
+        ).mock(
             return_value=Response(status_code=HTTPStatus.NO_CONTENT)
         )
 
-        # Mock failed version creation (non-201 for POST /pipelines/{name}/versions)
-        version_create_route = respx.post(
+        # Mock 404 response for GET (resource not found)
+        version_check_route = respx.get(
             "https://test-api-url.com/workspaces/test-workspace/pipelines/test-pipeline-fallback/versions"
-        ).mock(return_value=Response(status_code=HTTPStatus.BAD_REQUEST))
+        ).mock(
+            return_value=Response(status_code=HTTPStatus.NOT_FOUND)
+        )
 
         # Mock successful creation
-        create_route = respx.post("https://test-api-url.com/workspaces/test-workspace/pipelines").mock(
+        create_route = respx.post(
+            "https://test-api-url.com/workspaces/test-workspace/pipelines"
+        ).mock(
             return_value=Response(status_code=HTTPStatus.OK, json={"id": "test-pipeline-id"})
         )
 
@@ -508,24 +514,20 @@ async def test_import_pipeline_with_overwrite_fallback_to_create_async(
 
         await test_async_client.import_into_deepset(sample_pipeline, pipeline_config)
 
-        # Verify all three endpoints were called
+        # Verify all three endpoints were called in sequence
         assert validation_route.called
-        assert version_create_route.called
+        assert version_check_route.called
         assert create_route.called
 
         # Check validation request
         validation_request = validation_route.calls[0].request
         assert validation_request.headers["Authorization"] == "Bearer test-api-key"
         validation_body = json.loads(validation_request.content)
         assert "query_yaml" in validation_body
-        # When overwrite=True, name should be excluded from validation payload (if your code does that)
-        # assert "name" not in validation_body
-
-        # Check attempted version creation request
-        version_create_request = version_create_route.calls[0].request
-        assert version_create_request.headers["Authorization"] == "Bearer test-api-key"
-        version_body = json.loads(version_create_request.content)
-        assert "config_yaml" in version_body
+
+        # Check GET attempt
+        version_check_request = version_check_route.calls[0].request
+        assert version_check_request.headers["Authorization"] == "Bearer test-api-key"
 
         # Check fallback creation
         create_request = create_route.calls[0].request

tests/unit/service/test_pipeline_service.py

@@ -587,7 +587,7 @@ async def test_import_index_with_overwrite_fallback_to_create(
     async def test_import_pipeline_with_overwrite_true(
         self, pipeline_service: PipelineService, index_pipeline: Pipeline, mock_api: AsyncMock
     ) -> None:
-        """Test importing a pipeline with overwrite=True creates a new version via POST endpoint."""
+        """Test importing a pipeline with overwrite=True patches latest draft version."""
         config = PipelineConfig(
             name="test_pipeline_overwrite",
             inputs=PipelineInputs(query=["retriever.query"]),
@@ -600,39 +600,99 @@ async def test_import_pipeline_with_overwrite_true(
         validation_response = Mock(spec=Response)
         validation_response.status_code = HTTPStatus.NO_CONTENT.value
 
-        # Mock successful "create new version" response
+        # Mock successful versions response, latest version is a draft
+        versions_response = Mock(status_code=HTTPStatus.OK.value)
+        versions_response.json.return_value = {
+            "data": [{"version_id": "42abcd", "is_draft": True}],
+        }
+
+        # Mock successful overwrite (PATCH) response
         overwrite_response = Mock(spec=Response)
-        overwrite_response.status_code = HTTPStatus.CREATED.value
+        overwrite_response.status_code = HTTPStatus.OK.value
+
+        mock_api.post.return_value = validation_response
+        mock_api.get.return_value = versions_response
+        mock_api.patch.return_value = overwrite_response
+
+        await pipeline_service.import_async(index_pipeline, config)
+
+        # validation + GET versions + PATCH draft version
+        assert mock_api.post.call_count == 1
+        assert mock_api.get.call_count == 1
+        assert mock_api.patch.call_count == 1
+
+        # Check validation call
+        validation_call = mock_api.post.call_args_list[0]
+        assert validation_call.kwargs["endpoint"] == "pipeline_validations"
+        assert "query_yaml" in validation_call.kwargs["json"]
+        # When overwrite=True, name should be excluded from validation payload
+        assert "name" not in validation_call.kwargs["json"]
+
+        # Check PATCH call
+        overwrite_call = mock_api.patch.call_args_list[0]
+        assert (
+            overwrite_call.kwargs["endpoint"]
+            == "pipelines/test_pipeline_overwrite/versions/42abcd"
+        )
+        assert "config_yaml" in overwrite_call.kwargs["json"]
+
+    @pytest.mark.asyncio
+    async def test_import_pipeline_with_overwrite_true_creates_new_version_when_not_draft(
+        self, pipeline_service: PipelineService, index_pipeline: Pipeline, mock_api: AsyncMock
+    ) -> None:
+        """Test importing a pipeline with overwrite=True creates a new version when latest version is not draft."""
+        config = PipelineConfig(
+            name="test_pipeline_overwrite",
+            inputs=PipelineInputs(query=["retriever.query"]),
+            outputs=PipelineOutputs(documents="meta_ranker.documents"),
+            strict_validation=False,
+            overwrite=True,
+        )
+
+        # Mock successful validation response
+        validation_response = Mock(spec=Response)
+        validation_response.status_code = HTTPStatus.NO_CONTENT.value
+
+        # Mock versions response, latest version is NOT a draft
+        versions_response = Mock(status_code=HTTPStatus.OK.value)
+        versions_response.json.return_value = {
+            "data": [{"version_id": "42abcd", "is_draft": False}],
+        }
+
+        # Mock successful "create new version" response
+        new_version_response = Mock(spec=Response)
+        new_version_response.status_code = HTTPStatus.CREATED.value
 
         # First POST is validation, second POST is "create new version"
-        mock_api.post.side_effect = [validation_response, overwrite_response]
+        mock_api.post.side_effect = [validation_response, new_version_response]
+        mock_api.get.return_value = versions_response
 
         await pipeline_service.import_async(index_pipeline, config)
 
-        # Should call validation endpoint first, then create-version endpoint
+        # validation + GET versions + POST versions (new version)
         assert mock_api.post.call_count == 2
-        # No GET/PATCH/PUT calls in the overwrite path anymore
-        assert mock_api.get.call_count == 0
+        assert mock_api.get.call_count == 1
         assert mock_api.patch.call_count == 0
-        assert mock_api.put.call_count == 0
 
         # Check validation call
         validation_call = mock_api.post.call_args_list[0]
         assert validation_call.kwargs["endpoint"] == "pipeline_validations"
         assert "query_yaml" in validation_call.kwargs["json"]
-        # When overwrite=True, name should be excluded from validation payload
         assert "name" not in validation_call.kwargs["json"]
 
-        # Check create-version call
-        overwrite_call = mock_api.post.call_args_list[1]
-        assert overwrite_call.kwargs["endpoint"] == "pipelines/test_pipeline_overwrite/versions"
-        assert "config_yaml" in overwrite_call.kwargs["json"]
+        # Check create-version POST call
+        create_version_call = mock_api.post.call_args_list[1]
+        assert (
+            create_version_call.kwargs["endpoint"]
+            == "pipelines/test_pipeline_overwrite/versions"
+        )
+        assert "config_yaml" in create_version_call.kwargs["json"]
 
     @pytest.mark.asyncio
     async def test_import_pipeline_with_overwrite_fallback_to_create(
         self, pipeline_service: PipelineService, index_pipeline: Pipeline, mock_api: AsyncMock
     ) -> None:
-        """Test importing a pipeline with overwrite=True that falls back to create when version creation fails."""
+        """Test importing a pipeline with overwrite=True that falls back to create when resource doesn't exist."""
 
         config = PipelineConfig(
             name="test_pipeline_fallback",
@@ -646,41 +706,35 @@ async def test_import_pipeline_with_overwrite_fallback_to_create(
         validation_response = Mock(spec=Response)
         validation_response.status_code = HTTPStatus.NO_CONTENT.value
 
-        # Mock non-201 response for POST /pipelines/{name}/versions (version creation fails)
-        version_fail_response = Mock(spec=Response)
-        version_fail_response.status_code = HTTPStatus.BAD_REQUEST.value
+        # Mock 404 response for GET (resource not found)
+        not_found_response = Mock(spec=Response)
+        not_found_response.status_code = HTTPStatus.NOT_FOUND.value
 
-        # Mock successful creation response for POST /pipelines
+        # Mock successful creation response
         create_response = Mock(spec=Response)
         create_response.status_code = HTTPStatus.CREATED.value
 
-        # POST calls: validation, create-version (fails), create-pipeline (fallback)
-        mock_api.post.side_effect = [validation_response, version_fail_response, create_response]
+        mock_api.post.side_effect = [validation_response, create_response]
+        mock_api.get.return_value = not_found_response
 
         await pipeline_service.import_async(index_pipeline, config)
 
-        # Should call validation endpoint, then POST to create new version (fails),
-        # then POST to create the pipeline
-        assert mock_api.post.call_count == 3
-        # No GET anymore; overwrite logic doesn't fetch versions
-        assert mock_api.get.call_count == 0
-        assert mock_api.patch.call_count == 0
-        assert mock_api.put.call_count == 0
+        # validation + GET (404) + POST create
+        assert mock_api.post.call_count == 2
+        assert mock_api.get.call_count == 1
 
         # Check validation call
         validation_call = mock_api.post.call_args_list[0]
         assert validation_call.kwargs["endpoint"] == "pipeline_validations"
         assert "query_yaml" in validation_call.kwargs["json"]
-        # When overwrite=True, name should be excluded from validation payload
         assert "name" not in validation_call.kwargs["json"]
 
-        # Check attempted version creation call
-        version_call = mock_api.post.call_args_list[1]
-        assert version_call.kwargs["endpoint"] == "pipelines/test_pipeline_fallback/versions"
-        assert "config_yaml" in version_call.kwargs["json"]
+        # Check GET versions attempt
+        get_call = mock_api.get.call_args_list[0]
+        assert get_call.kwargs["endpoint"] == "pipelines/test_pipeline_fallback/versions"
 
-        # Check fallback create-pipeline call
-        create_call = mock_api.post.call_args_list[2]
+        # Check fallback POST call
+        create_call = mock_api.post.call_args_list[1]
         assert create_call.kwargs["endpoint"] == "pipelines"
         assert create_call.kwargs["json"]["name"] == "test_pipeline_fallback"
         assert "query_yaml" in create_call.kwargs["json"]