[BUG]: Allow repairing partially created attached functions (#5981)

## Description of changes

_Summarize the changes made by this PR._

- Improvements & Bug fixes
    - The attach_function codepath did not allow repairing a nonready function that was a result of a previously failed attach. This diff fixes that.
    - This diff also adds code to the FinishCreateAttachedFunction method to ensure the invariant that there can only be one function attached to a collection.
- New functionality
    - ...

Many changes and tests here were ported from a change @rescrv had made.

## Test plan

_How are these changes tested?_

test_task_api.py has been edited for this change.

- [ ] Tests pass locally with `pytest` for python, `yarn test` for js, `cargo test` for rust

## Migration plan

_Are there any migrations, or any forwards/backwards compatibility changes needed in order to make sure this change deploys reliably?_

## Observability plan

_What is the plan to instrument and monitor this change?_

## Documentation Changes

_Are all docstrings for user-facing APIs updated if required? Do we need to make documentation changes in the_ [_docs section](https://github.com/chroma-core/chroma/tree/main/docs/docs.trychroma.com)?_

Author: tanujnay112

chromadb/test/distributed/test_task_api.py

@@ -384,3 +384,127 @@ def test_delete_orphaned_output_collection(basic_http_client: System) -> None:
     with pytest.raises(NotFoundError):
         # Try to use the function - it should fail since it's detached
         client.get_collection("output_collection")
+
+
+def test_partial_attach_function_repair(
+    basic_http_client: System,
+) -> None:
+    """Test creating and removing a function with the record_counter operator"""
+    client = ClientCreator.from_system(basic_http_client)
+    client.reset()
+
+    # Create a collection
+    collection = client.get_or_create_collection(
+        name="my_document",
+    )
+
+    # Create a task that counts records in the collection
+    attached_fn = collection.attach_function(
+        name="count_my_docs",
+        function=RECORD_COUNTER_FUNCTION,
+        output_collection="my_documents_counts",
+        params=None,
+    )
+
+    # Verify task creation succeeded
+    assert attached_fn is not None
+
+    collection2 = client.get_or_create_collection(
+        name="my_document2",
+    )
+
+    # Create a task that counts records in the collection
+    # This should fail
+    with pytest.raises(
+        ChromaError, match=r"Output collection \[my_documents_counts\] already exists"
+    ):
+        attached_fn = collection2.attach_function(
+            name="count_my_docs",
+            function=RECORD_COUNTER_FUNCTION,
+            output_collection="my_documents_counts",
+            params=None,
+        )
+
+    # Detach the function
+    assert (
+        collection.detach_function(attached_fn.name, delete_output_collection=True)
+        is True
+    )
+
+    # Create a task that counts records in the collection
+    attached_fn = collection2.attach_function(
+        name="count_my_docs",
+        function=RECORD_COUNTER_FUNCTION,
+        output_collection="my_documents_counts",
+        params=None,
+    )
+    assert attached_fn is not None
+
+
+def test_count_function_attach_and_detach_attach_attach(
+    basic_http_client: System,
+) -> None:
+    """Test creating and removing a function with the record_counter operator"""
+    client = ClientCreator.from_system(basic_http_client)
+    client.reset()
+
+    # Create a collection
+    collection = client.get_or_create_collection(
+        name="my_document",
+        metadata={"description": "Sample documents for task processing"},
+    )
+
+    # Create a task that counts records in the collection
+    attached_fn = collection.attach_function(
+        name="count_my_docs",
+        function=RECORD_COUNTER_FUNCTION,
+        output_collection="my_documents_counts",
+        params=None,
+    )
+
+    # Verify task creation succeeded
+    assert attached_fn is not None
+    initial_version = get_collection_version(client, collection.name)
+
+    # Add documents
+    collection.add(
+        ids=["doc_{}".format(i) for i in range(0, 300)],
+        documents=["test document"] * 300,
+    )
+
+    # Verify documents were added
+    assert collection.count() == 300
+
+    wait_for_version_increase(client, collection.name, initial_version)
+    # Give some time to invalidate the frontend query cache
+    sleep(60)
+
+    result = client.get_collection("my_documents_counts").get("function_output")
+    assert result["metadatas"] is not None
+    assert result["metadatas"][0]["total_count"] == 300
+
+    # Remove the task
+    success = collection.detach_function(
+        attached_fn.name, delete_output_collection=True
+    )
+
+    # Verify task removal succeeded
+    assert success is True
+
+    # Create a task that counts records in the collection
+    attached_fn = collection.attach_function(
+        name="count_my_docs",
+        function=RECORD_COUNTER_FUNCTION,
+        output_collection="my_documents_counts",
+        params=None,
+    )
+    assert attached_fn is not None
+
+    # Create a task that counts records in the collection
+    attached_fn = collection.attach_function(
+        name="count_my_docs",
+        function=RECORD_COUNTER_FUNCTION,
+        output_collection="my_documents_counts",
+        params=None,
+    )
+    assert attached_fn is not None

go/pkg/sysdb/coordinator/create_task_test.go

@@ -69,7 +69,7 @@ func (suite *AttachFunctionTestSuite) setupAttachFunctionMocks(ctx context.Conte
 	// Phase 1: Create attached function in transaction
 	// Check if any attached function exists for this collection
 	suite.mockMetaDomain.On("AttachedFunctionDb", mock.Anything).Return(suite.mockAttachedFunctionDb).Once()
-	suite.mockAttachedFunctionDb.On("GetByCollectionID", inputCollectionID).
+	suite.mockAttachedFunctionDb.On("GetAnyByCollectionID", inputCollectionID).
 		Return([]*dbmodel.AttachedFunction{}, nil).Once()
 
 	suite.mockMetaDomain.On("DatabaseDb", mock.Anything).Return(suite.mockDatabaseDb).Once()
@@ -165,7 +165,7 @@ func (suite *AttachFunctionTestSuite) TestAttachFunction_SuccessfulCreation() {
 	// Setup mocks that will be called within the transaction (using mock.Anything for context)
 	// Check if any attached function exists for this collection
 	suite.mockMetaDomain.On("AttachedFunctionDb", mock.Anything).Return(suite.mockAttachedFunctionDb).Once()
-	suite.mockAttachedFunctionDb.On("GetByCollectionID", inputCollectionID).
+	suite.mockAttachedFunctionDb.On("GetAnyByCollectionID", inputCollectionID).
 		Return([]*dbmodel.AttachedFunction{}, nil).Once()
 
 	// Look up database
@@ -283,7 +283,7 @@ func (suite *AttachFunctionTestSuite) TestAttachFunction_IdempotentRequest_Alrea
 
 			// Inside transaction: check for existing attached functions
 			suite.mockMetaDomain.On("AttachedFunctionDb", txCtx).Return(suite.mockAttachedFunctionDb).Once()
-			suite.mockAttachedFunctionDb.On("GetByCollectionID", inputCollectionID).
+			suite.mockAttachedFunctionDb.On("GetAnyByCollectionID", inputCollectionID).
 				Return([]*dbmodel.AttachedFunction{existingAttachedFunction}, nil).Once()
 
 			// Note: validateAttachedFunctionMatchesRequest uses dbmodel.GetFunctionNameByID (static lookup),
@@ -352,7 +352,7 @@ func (suite *AttachFunctionTestSuite) TestAttachFunction_RecoveryFlow() {
 
 	// Phase 1: Create attached function in transaction
 	suite.mockMetaDomain.On("AttachedFunctionDb", mock.Anything).Return(suite.mockAttachedFunctionDb).Once()
-	suite.mockAttachedFunctionDb.On("GetByCollectionID", inputCollectionID).
+	suite.mockAttachedFunctionDb.On("GetAnyByCollectionID", inputCollectionID).
 		Return([]*dbmodel.AttachedFunction{}, nil).Once()
 
 	suite.mockMetaDomain.On("DatabaseDb", mock.Anything).Return(suite.mockDatabaseDb).Once()
@@ -408,7 +408,7 @@ func (suite *AttachFunctionTestSuite) TestAttachFunction_RecoveryFlow() {
 
 			// Inside transaction: check for existing attached functions
 			suite.mockMetaDomain.On("AttachedFunctionDb", txCtx).Return(suite.mockAttachedFunctionDb).Once()
-			suite.mockAttachedFunctionDb.On("GetByCollectionID", inputCollectionID).
+			suite.mockAttachedFunctionDb.On("GetAnyByCollectionID", inputCollectionID).
 				Return([]*dbmodel.AttachedFunction{incompleteAttachedFunction}, nil).Once()
 
 			// Note: validateAttachedFunctionMatchesRequest uses dbmodel.GetFunctionNameByID (static lookup),
@@ -484,6 +484,7 @@ func (suite *AttachFunctionTestSuite) TestAttachFunction_IdempotentRequest_Param
 		OutputCollectionName:    outputCollectionName,
 		FunctionID:              existingOperatorID,
 		MinRecordsForInvocation: int64(MinRecordsForInvocation),
+		IsReady:                 true,
 		CreatedAt:               now,
 		UpdatedAt:               now,
 	}
@@ -497,7 +498,7 @@ func (suite *AttachFunctionTestSuite) TestAttachFunction_IdempotentRequest_Param
 
 			// Inside transaction: check for existing attached functions
 			suite.mockMetaDomain.On("AttachedFunctionDb", txCtx).Return(suite.mockAttachedFunctionDb).Once()
-			suite.mockAttachedFunctionDb.On("GetByCollectionID", inputCollectionID).
+			suite.mockAttachedFunctionDb.On("GetAnyByCollectionID", inputCollectionID).
 				Return([]*dbmodel.AttachedFunction{existingAttachedFunction}, nil).Once()
 
 			// Note: validateAttachedFunctionMatchesRequest uses dbmodel.GetFunctionNameByID (static lookup)

go/pkg/sysdb/coordinator/task.go

@@ -99,43 +99,36 @@ func (s *Coordinator) AttachFunction(ctx context.Context, req *coordinatorpb.Att
 	err := s.catalog.txImpl.Transaction(ctx, func(txCtx context.Context) error {
 		// Check if there's any active (ready, non-deleted) attached function for this collection
 		// We only allow one active attached function per collection
-		existingAttachedFunctions, err := s.catalog.metaDomain.AttachedFunctionDb(txCtx).GetByCollectionID(req.InputCollectionId)
+		existingAttachedFunctions, err := s.catalog.metaDomain.AttachedFunctionDb(txCtx).GetAnyByCollectionID(req.InputCollectionId)
 		if err != nil {
 			log.Error("AttachFunction: failed to check for existing attached function", zap.Error(err))
 			return err
 		}
-		if len(existingAttachedFunctions) > 0 {
-			if len(existingAttachedFunctions) > 1 {
-				log.Error("AttachFunction: collection has multiple attached functions")
-				return status.Errorf(codes.Internal, "data inconsistency: collection has %d attached functions, expected at most 1", len(existingAttachedFunctions))
-			}
-
-			existingAttachedFunction := existingAttachedFunctions[0]
-
-			// Same name - validate it matches our request (idempotency)
-			log.Info("AttachFunction: attached function exists with same name, validating parameters",
-				zap.String("attached_function_id", existingAttachedFunction.ID.String()))
 
-			matches, err := s.validateAttachedFunctionMatchesRequest(txCtx, existingAttachedFunction, req)
+		for _, attachedFunction := range existingAttachedFunctions {
+			matches, err := s.validateAttachedFunctionMatchesRequest(txCtx, attachedFunction, req)
 			if err != nil {
 				return err
 			}
-			if !matches {
-				functionName, err := dbmodel.GetFunctionNameByID(existingAttachedFunction.FunctionID)
+			if matches {
+				// If the attached function matches the request, use it
+				attachedFunctionID = attachedFunction.ID
+				return nil
+			}
+
+			if attachedFunction.IsReady {
+				log.Error("AttachFunction: collection already has an attached function", zap.String("name", attachedFunction.Name))
+				functionName, err := dbmodel.GetFunctionNameByID(attachedFunction.FunctionID)
 				if err != nil {
 					log.Error("AttachFunction: unknown function ID", zap.Error(err))
 					return err
 				}
 				return status.Errorf(codes.AlreadyExists,
 					"collection already has an attached function: name=%s, function=%s, output_collection=%s",
-					existingAttachedFunction.Name,
+					attachedFunction.Name,
 					functionName,
-					existingAttachedFunction.OutputCollectionName)
+					attachedFunction.OutputCollectionName)
 			}
-
-			// Validation passed, reuse the existing attached function ID (idempotent)
-			attachedFunctionID = existingAttachedFunction.ID
-			return nil
 		}
 
 		// Look up database_id
@@ -626,6 +619,17 @@ func (s *Coordinator) FinishCreateAttachedFunction(ctx context.Context, req *coo
 			return err
 		}
 
+		// 7. Validate that there is only one ready attached function for this collection
+		existingAttachedFunctions, err := s.catalog.metaDomain.AttachedFunctionDb(txCtx).GetByCollectionID(attachedFunction.InputCollectionID)
+		if err != nil {
+			log.Error("FinishCreateAttachedFunction: failed to get attached functions", zap.Error(err))
+			return err
+		}
+		if len(existingAttachedFunctions) > 1 {
+			log.Error("FinishCreateAttachedFunction: multiple attached functions found for collection", zap.String("collection_id", attachedFunction.InputCollectionID))
+			return common.ErrAttachedFunctionAlreadyExists
+		}
+
 		log.Info("FinishCreateAttachedFunction: successfully created output collection and set is_ready=true",
 			zap.String("attached_function_id", attachedFunctionID.String()),
 			zap.String("output_collection_id", collectionID.String()))

go/pkg/sysdb/metastore/db/dao/task.go

@@ -153,6 +153,21 @@ func (s *attachedFunctionDb) GetByCollectionID(inputCollectionID string) ([]*dbm
 	return attachedFunctions, nil
 }
 
+func (s *attachedFunctionDb) GetAnyByCollectionID(inputCollectionID string) ([]*dbmodel.AttachedFunction, error) {
+	var attachedFunctions []*dbmodel.AttachedFunction
+	err := s.db.
+		Where("input_collection_id = ?", inputCollectionID).
+		Where("is_deleted = ?", false).
+		Find(&attachedFunctions).Error
+
+	if err != nil {
+		log.Error("GetAnyByCollectionID failed", zap.Error(err), zap.String("input_collection_id", inputCollectionID))
+		return nil, err
+	}
+
+	return attachedFunctions, nil
+}
+
 func (s *attachedFunctionDb) SoftDelete(inputCollectionID string, name string) error {
 	// Update name and is_deleted in a single query
 	// Format: _deleted_<original_name>_<id>

go/pkg/sysdb/metastore/db/dbmodel/mocks/IAttachedFunctionDb.go

@@ -42,6 +42,7 @@ type IAttachedFunctionDb interface {
 	GetByID(id uuid.UUID) (*AttachedFunction, error)
 	GetAnyByID(id uuid.UUID) (*AttachedFunction, error) // TODO(tanujnay112): Consolidate all the getters.
 	GetByCollectionID(inputCollectionID string) ([]*AttachedFunction, error)
+	GetAnyByCollectionID(inputCollectionID string) ([]*AttachedFunction, error)
 	Update(attachedFunction *AttachedFunction) error
 	Finish(id uuid.UUID) error
 	SoftDelete(inputCollectionID string, name string) error