milvus/internal/rootcoord/ddl_callbacks_collection_function_test.go

// Licensed to the LF AI & Data foundation under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package rootcoord

import (
	"context"
	"testing"

	"github.com/bytedance/mockey"
	"github.com/cockroachdb/errors"
	"github.com/stretchr/testify/mock"
	"github.com/stretchr/testify/suite"

	"github.com/milvus-io/milvus-proto/go-api/v3/commonpb"
	"github.com/milvus-io/milvus-proto/go-api/v3/milvuspb"
	"github.com/milvus-io/milvus-proto/go-api/v3/schemapb"
	"github.com/milvus-io/milvus/internal/metastore/model"
	"github.com/milvus-io/milvus/internal/mocks/streamingcoord/server/mock_broadcaster"
	mockrootcoord "github.com/milvus-io/milvus/internal/rootcoord/mocks"
	"github.com/milvus-io/milvus/pkg/v3/streaming/util/message"
	"github.com/milvus-io/milvus/pkg/v3/streaming/util/types"
	"github.com/milvus-io/milvus/pkg/v3/util/typeutil"
)

type DDLCallbacksCollectionFunctionTestSuite struct {
	suite.Suite
	core            *Core
	mockBroadcaster *mock_broadcaster.MockBroadcastAPI
}

func (suite *DDLCallbacksCollectionFunctionTestSuite) SetupTest() {
	suite.core = initStreamingSystemAndCore(suite.T())
	suite.mockBroadcaster = mock_broadcaster.NewMockBroadcastAPI(suite.T())
}

func (suite *DDLCallbacksCollectionFunctionTestSuite) TearDownTest() {
	suite.mockBroadcaster = nil
	suite.core.Stop()
	suite.core = nil
}

// Helper function to create a test collection
func (suite *DDLCallbacksCollectionFunctionTestSuite) createTestCollection() *model.Collection {
	return &model.Collection{
		DBID:         1,
		CollectionID: 100,
		Name:         "test_collection",
		Description:  "test collection",
		AutoID:       true,
		Fields: []*model.Field{
			{
				FieldID:          101,
				Name:             "id",
				IsPrimaryKey:     true,
				AutoID:           true,
				DataType:         schemapb.DataType_Int64,
				IsFunctionOutput: false,
			},
			{
				FieldID:          102,
				Name:             "vector",
				DataType:         schemapb.DataType_FloatVector,
				IsFunctionOutput: false,
			},
			{
				FieldID:          103,
				Name:             "text_field",
				DataType:         schemapb.DataType_VarChar,
				IsFunctionOutput: false,
			},
			{
				FieldID:          104,
				Name:             "output_field",
				DataType:         schemapb.DataType_FloatVector,
				IsFunctionOutput: true,
			},
		},
		Functions: []*model.Function{
			{
				ID:               1001,
				Name:             "test_function",
				Type:             schemapb.FunctionType_TextEmbedding,
				Description:      "test function",
				InputFieldIDs:    []int64{103},
				InputFieldNames:  []string{"text_field"},
				OutputFieldIDs:   []int64{104},
				OutputFieldNames: []string{"output_field"},
				Params: []*commonpb.KeyValuePair{
					{Key: "param1", Value: "value1"},
				},
			},
		},
		VirtualChannelNames: []string{"test_channel"},
		EnableDynamicField:  false,
		Properties:          []*commonpb.KeyValuePair{{Key: "key1", Value: "value1"}},
		SchemaVersion:       1,
	}
}

// Helper function to create a test function schema
func (suite *DDLCallbacksCollectionFunctionTestSuite) createTestFunctionSchema() *schemapb.FunctionSchema {
	return &schemapb.FunctionSchema{
		Name:             "new_function",
		Type:             schemapb.FunctionType_TextEmbedding,
		Description:      "new test function",
		InputFieldNames:  []string{"text_field"},
		OutputFieldNames: []string{"vector"},
		Params: []*commonpb.KeyValuePair{
			{Key: "param1", Value: "value1"},
		},
	}
}

func TestDDLCallbacksCollectionFunctionTestSuite(t *testing.T) {
	suite.Run(t, new(DDLCallbacksCollectionFunctionTestSuite))
}

// Test callAlterCollection function
func (suite *DDLCallbacksCollectionFunctionTestSuite) TestCallAlterCollection_Success() {
	ctx := context.Background()
	coll := suite.createTestCollection()
	dbName := "test_db"
	collectionName := "test_collection"

	// Create a test core with proper meta setup
	core := newTestCore(withHealthyCode())
	mockMeta := mockrootcoord.NewIMetaTable(suite.T())
	core.meta = mockMeta

	// Mock meta calls for getCacheExpireForCollection
	mockMeta.EXPECT().GetCollectionByName(mock.Anything, dbName, collectionName, typeutil.MaxTimestamp, mock.Anything).Return(coll, nil)
	mockMeta.EXPECT().ListAliases(mock.Anything, dbName, collectionName, typeutil.MaxTimestamp).Return([]string{}, nil)

	// Mock broadcaster
	suite.mockBroadcaster.EXPECT().Broadcast(mock.Anything, mock.MatchedBy(func(msg message.BroadcastMutableMessage) bool {
		// Verify the message structure
		return msg != nil
	})).Return(&types.BroadcastAppendResult{}, nil)

	err := callAlterCollection(ctx, core, suite.mockBroadcaster, coll, coll, dbName, collectionName)
	suite.NoError(err)
}

func (suite *DDLCallbacksCollectionFunctionTestSuite) TestCallAlterCollection_GetCacheExpireFailed() {
	ctx := context.Background()
	coll := suite.createTestCollection()
	dbName := "test_db"
	collectionName := "test_collection"

	// Create a test core with proper meta setup
	core := newTestCore(withHealthyCode())
	mockMeta := mockrootcoord.NewIMetaTable(suite.T())
	core.meta = mockMeta

	// Mock meta calls to return error
	mockMeta.EXPECT().GetCollectionByName(mock.Anything, dbName, collectionName, typeutil.MaxTimestamp, mock.Anything).Return(nil, errors.New("cache expire error"))

	err := callAlterCollection(ctx, core, suite.mockBroadcaster, coll, coll, dbName, collectionName)
	suite.Error(err)
	suite.Contains(err.Error(), "cache expire error")
}

func (suite *DDLCallbacksCollectionFunctionTestSuite) TestCallAlterCollection_BroadcastFailed() {
	ctx := context.Background()
	coll := suite.createTestCollection()
	dbName := "test_db"
	collectionName := "test_collection"

	// Create a test core with proper meta setup
	core := newTestCore(withHealthyCode())
	mockMeta := mockrootcoord.NewIMetaTable(suite.T())
	core.meta = mockMeta

	// Mock meta calls for getCacheExpireForCollection
	mockMeta.EXPECT().GetCollectionByName(mock.Anything, dbName, collectionName, typeutil.MaxTimestamp, mock.Anything).Return(coll, nil)
	mockMeta.EXPECT().ListAliases(mock.Anything, dbName, collectionName, typeutil.MaxTimestamp).Return([]string{}, nil)

	// Mock broadcaster to return error
	suite.mockBroadcaster.EXPECT().Broadcast(mock.Anything, mock.Anything).Return(nil, errors.New("broadcast error"))

	err := callAlterCollection(ctx, core, suite.mockBroadcaster, coll, coll, dbName, collectionName)
	suite.Error(err)
	suite.Contains(err.Error(), "broadcast error")
}

// Test alterFunctionGenNewCollection function
func (suite *DDLCallbacksCollectionFunctionTestSuite) TestAlterFunctionGenNewCollection_Success() {
	ctx := context.Background()
	coll := suite.createTestCollection()

	// Create a function schema to alter existing function
	fSchema := &schemapb.FunctionSchema{
		Name:             "test_function", // Same name as existing function
		Type:             schemapb.FunctionType_TextEmbedding,
		Description:      "updated test function",
		InputFieldNames:  []string{"text_field"},
		OutputFieldNames: []string{"vector"}, // Different output field
	}

	err := alterFunctionGenNewCollection(ctx, fSchema, coll)
	suite.NoError(err)

	// Verify function ID is preserved
	suite.Equal(int64(1001), fSchema.Id)

	// Verify input field IDs are set
	suite.Equal([]int64{103}, fSchema.InputFieldIds)

	// Verify output field IDs are set
	suite.Equal([]int64{102}, fSchema.OutputFieldIds)

	// Verify old output field is no longer marked as function output
	oldOutputField := coll.Fields[3] // output_field
	suite.False(oldOutputField.IsFunctionOutput)

	// Verify new output field is marked as function output
	newOutputField := coll.Fields[1] // vector
	suite.True(newOutputField.IsFunctionOutput)

	// Verify function is added to collection
	suite.Len(coll.Functions, 1) // Original + new
}

func (suite *DDLCallbacksCollectionFunctionTestSuite) TestAlterFunctionGenNewCollection_DiscardsClientFieldIds() {
	ctx := context.Background()
	coll := suite.createTestCollection()

	// Client injects bogus field IDs (output_field_ids points at the unrelated
	// "vector" field 102). Only the name-resolved IDs must persist, else a later
	// drop_function_field would delete field 102.
	fSchema := &schemapb.FunctionSchema{
		Name:             "test_function",
		Type:             schemapb.FunctionType_TextEmbedding,
		InputFieldNames:  []string{"text_field"},
		InputFieldIds:    []int64{999},
		OutputFieldNames: []string{"output_field"},
		OutputFieldIds:   []int64{102},
	}

	err := alterFunctionGenNewCollection(ctx, fSchema, coll)
	suite.NoError(err)
	suite.Equal([]int64{103}, fSchema.InputFieldIds)  // text_field only, 999 discarded
	suite.Equal([]int64{104}, fSchema.OutputFieldIds) // output_field only, 102 discarded
}

func (suite *DDLCallbacksCollectionFunctionTestSuite) TestAlterFunctionGenNewCollection_FunctionNotExists() {
	ctx := context.Background()
	coll := suite.createTestCollection()

	fSchema := &schemapb.FunctionSchema{
		Name: "non_existent_function",
	}

	err := alterFunctionGenNewCollection(ctx, fSchema, coll)
	suite.Error(err)
	suite.Contains(err.Error(), "function non_existent_function not exists")
}

func (suite *DDLCallbacksCollectionFunctionTestSuite) TestAlterFunctionGenNewCollection_InputFieldNotExists() {
	ctx := context.Background()
	coll := suite.createTestCollection()

	fSchema := &schemapb.FunctionSchema{
		Name:            "test_function",
		InputFieldNames: []string{"non_existent_field"},
	}

	err := alterFunctionGenNewCollection(ctx, fSchema, coll)
	suite.Error(err)
	suite.Contains(err.Error(), "function's input field non_existent_field not exists")
}

func (suite *DDLCallbacksCollectionFunctionTestSuite) TestAlterFunctionGenNewCollection_OutputFieldNotExists() {
	ctx := context.Background()
	coll := suite.createTestCollection()

	fSchema := &schemapb.FunctionSchema{
		Name:             "test_function",
		InputFieldNames:  []string{"text_field"},
		OutputFieldNames: []string{"non_existent_field"},
	}

	err := alterFunctionGenNewCollection(ctx, fSchema, coll)
	suite.Error(err)
	suite.Contains(err.Error(), "function's output field non_existent_field not exists")
}

// Test edge cases and error conditions
func (suite *DDLCallbacksCollectionFunctionTestSuite) TestAlterFunctionGenNewCollection_OldOutputFieldNotExists() {
	coll := suite.createTestCollection()

	// Modify the existing function to have a non-existent output field
	coll.Functions[0].OutputFieldNames = []string{"non_existent_output"}

	fSchema := &schemapb.FunctionSchema{
		Name:             "test_function",
		InputFieldNames:  []string{"text_field"},
		OutputFieldNames: []string{"vector"},
	}

	err := alterFunctionGenNewCollection(context.Background(), fSchema, coll)
	suite.Error(err)
	suite.Contains(err.Error(), "old version function's output field non_existent_output not exists")
}

// Test with empty collections and functions
func (suite *DDLCallbacksCollectionFunctionTestSuite) TestAlterFunctionGenNewCollection_EmptyCollection() {
	// Create collection with no functions
	coll := &model.Collection{
		Fields:    []*model.Field{},
		Functions: []*model.Function{},
	}

	fSchema := &schemapb.FunctionSchema{
		Name: "test_function",
	}

	err := alterFunctionGenNewCollection(context.Background(), fSchema, coll)
	suite.Error(err)
	suite.Contains(err.Error(), "function test_function not exists")
}

// Test max function ID calculation
func (suite *DDLCallbacksCollectionFunctionTestSuite) TestAddFunction_NextFunctionIDCalculation() {
	coll := suite.createTestCollection()

	// Add more functions to test max ID calculation
	coll.Functions = append(coll.Functions, &model.Function{
		ID:   2000,
		Name: "function2",
	})
	coll.Functions = append(coll.Functions, &model.Function{
		ID:   1500,
		Name: "function3",
	})

	nextFunctionID := int64(StartOfUserFunctionID)
	for _, f := range coll.Functions {
		nextFunctionID = max(nextFunctionID, f.ID+1)
	}

	suite.Equal(int64(2001), nextFunctionID)
}

// Test broadcastAlterCollectionForAlterFunction
func (suite *DDLCallbacksCollectionFunctionTestSuite) TestBroadcastAlterCollectionForAlterFunction() {
	suite.Run("success", func() {
		coll := suite.createTestCollection()

		mockMeta := mockrootcoord.NewIMetaTable(suite.T())
		mockMeta.EXPECT().GetCollectionByName(mock.Anything, "test_db", "test_collection", typeutil.MaxTimestamp, mock.Anything).Return(coll, nil)
		suite.core.meta = mockMeta

		req := &milvuspb.AlterCollectionFunctionRequest{
			DbName:         "test_db",
			CollectionName: "test_collection",
			FunctionSchema: &schemapb.FunctionSchema{
				Name:             "test_function",
				Type:             schemapb.FunctionType_TextEmbedding,
				InputFieldNames:  []string{"text_field"},
				OutputFieldNames: []string{"output_field"},
				// identity unchanged; only a whitelisted connection param ("url") is altered
				Params: []*commonpb.KeyValuePair{
					{Key: "param1", Value: "value1"},
					{Key: "url", Value: "http://new-endpoint"},
				},
			},
		}
		mocker := mockey.Mock(callAlterCollection).Return(nil).Build()
		defer mocker.UnPatch()

		err := suite.core.broadcastAlterCollectionForAlterFunction(context.Background(), req)
		suite.NoError(err)
	})

	suite.Run("altering non-whitelisted param rejected", func() {
		coll := suite.createTestCollection()

		mockMeta := mockrootcoord.NewIMetaTable(suite.T())
		mockMeta.EXPECT().GetCollectionByName(mock.Anything, "test_db", "test_collection", typeutil.MaxTimestamp, mock.Anything).Return(coll, nil)
		suite.core.meta = mockMeta

		req := &milvuspb.AlterCollectionFunctionRequest{
			DbName:         "test_db",
			CollectionName: "test_collection",
			FunctionSchema: &schemapb.FunctionSchema{
				Name:             "test_function",
				Type:             schemapb.FunctionType_TextEmbedding,
				InputFieldNames:  []string{"text_field"},
				OutputFieldNames: []string{"output_field"},
				Params: []*commonpb.KeyValuePair{
					{Key: "param1", Value: "changed"},
				},
			},
		}
		err := suite.core.broadcastAlterCollectionForAlterFunction(context.Background(), req)
		suite.ErrorContains(err, "cannot be altered")
	})

	suite.Run("external collection rejected", func() {
		coll := suite.createTestCollection()
		coll.Fields[2].ExternalField = "text_col"

		mockMeta := mockrootcoord.NewIMetaTable(suite.T())
		mockMeta.EXPECT().GetCollectionByName(mock.Anything, "test_db", "test_collection", typeutil.MaxTimestamp, mock.Anything).Return(coll, nil)
		suite.core.meta = mockMeta

		req := &milvuspb.AlterCollectionFunctionRequest{
			DbName:         "test_db",
			CollectionName: "test_collection",
			FunctionSchema: &schemapb.FunctionSchema{
				Name:             "test_function",
				Type:             schemapb.FunctionType_TextEmbedding,
				InputFieldNames:  []string{"text_field"},
				OutputFieldNames: []string{"output_field"},
			},
		}
		err := suite.core.broadcastAlterCollectionForAlterFunction(context.Background(), req)
		suite.ErrorContains(err, externalCollectionFunctionMutationUnsupportedMsg)
	})
}