Skip to content

Commit 6cfa1f5

Browse files
tanujnay112tanujnay112
committed
[CHORE]: Consolidate all attached function getters
1 parent 4d9d369 commit 6cfa1f5

File tree

16 files changed

+338
-801
lines changed

16 files changed

+338
-801
lines changed

‎go/pkg/sysdb/coordinator/create_task_test.go‎

Lines changed: 11 additions & 64 deletions
Original file line numberDiff line numberDiff line change
@@ -18,15 +18,6 @@ import (
1818
"google.golang.org/protobuf/types/known/timestamppb"
1919
)
2020

21-
// testMinimalUUIDv7 is the test's copy of minimalUUIDv7 from task.go
22-
// UUIDv7 format: [timestamp (48 bits)][version (4 bits)][random (12 bits)][variant (2 bits)][random (62 bits)]
23-
// This UUID has all zeros for timestamp and random bits, making it the minimal valid UUIDv7.
24-
var testMinimalUUIDv7 = uuid.UUID{
25-
0x00, 0x00, 0x00, 0x00, 0x00, 0x00, // timestamp = 0 (bytes 0-5)
26-
0x70, 0x00, // version 7 (0x7) in high nibble, low nibble = 0 (bytes 6-7)
27-
0x80, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, // variant bits + rest = 0 (bytes 8-15)
28-
}
29-
3021
// MockMemberlistStore is a mock implementation of memberlist_manager.IMemberlistStore for testing
3122
type MockMemberlistStore struct {
3223
mock.Mock
@@ -57,50 +48,6 @@ type AttachFunctionTestSuite struct {
5748
coordinator *Coordinator
5849
}
5950

60-
// setupAttachFunctionMocks sets up all the mocks for an AttachFunction call (Phases 0 and 1)
61-
// Returns a function that can be called to capture the created attached function ID
62-
func (suite *AttachFunctionTestSuite) setupAttachFunctionMocks(ctx context.Context, request *coordinatorpb.AttachFunctionRequest, databaseID string, functionID uuid.UUID) func(*dbmodel.AttachedFunction) bool {
63-
inputCollectionID := request.InputCollectionId
64-
attachedFunctionName := request.Name
65-
outputCollectionName := request.OutputCollectionName
66-
tenantID := request.TenantId
67-
databaseName := request.Database
68-
functionName := request.FunctionName
69-
70-
// Phase 0: No existing attached function
71-
suite.mockMetaDomain.On("AttachedFunctionDb", ctx).Return(suite.mockAttachedFunctionDb).Once()
72-
suite.mockAttachedFunctionDb.On("GetAnyByName", inputCollectionID, attachedFunctionName).
73-
Return(nil, nil).Once()
74-
75-
// Phase 1: Create attached function in transaction
76-
suite.mockMetaDomain.On("AttachedFunctionDb", mock.Anything).Return(suite.mockAttachedFunctionDb).Once()
77-
suite.mockAttachedFunctionDb.On("GetAnyByName", inputCollectionID, attachedFunctionName).
78-
Return(nil, nil).Once()
79-
80-
suite.mockMetaDomain.On("DatabaseDb", mock.Anything).Return(suite.mockDatabaseDb).Once()
81-
suite.mockDatabaseDb.On("GetDatabases", tenantID, databaseName).
82-
Return([]*dbmodel.Database{{ID: databaseID, Name: databaseName}}, nil).Once()
83-
84-
suite.mockMetaDomain.On("FunctionDb", mock.Anything).Return(suite.mockFunctionDb).Once()
85-
suite.mockFunctionDb.On("GetByName", functionName).
86-
Return(&dbmodel.Function{ID: functionID, Name: functionName}, nil).Once()
87-
88-
suite.mockMetaDomain.On("CollectionDb", mock.Anything).Return(suite.mockCollectionDb).Once()
89-
suite.mockCollectionDb.On("GetCollections",
90-
[]string{inputCollectionID}, (*string)(nil), tenantID, databaseName, (*int32)(nil), (*int32)(nil), false).
91-
Return([]*dbmodel.CollectionAndMetadata{{Collection: &dbmodel.Collection{ID: inputCollectionID}}}, nil).Once()
92-
93-
suite.mockMetaDomain.On("CollectionDb", mock.Anything).Return(suite.mockCollectionDb).Once()
94-
suite.mockCollectionDb.On("GetCollections",
95-
[]string(nil), &outputCollectionName, tenantID, databaseName, (*int32)(nil), (*int32)(nil), false).
96-
Return([]*dbmodel.CollectionAndMetadata{}, nil).Once()
97-
98-
// Return a matcher function that can be used to capture attached function data
99-
return func(attachedFunction *dbmodel.AttachedFunction) bool {
100-
return true
101-
}
102-
}
103-
10451
func (suite *AttachFunctionTestSuite) SetupTest() {
10552
// Create all mocks - note: we manually control AssertExpectations
10653
// to avoid conflicts with automatic cleanup
@@ -172,8 +119,8 @@ func (suite *AttachFunctionTestSuite) TestAttachFunction_SuccessfulCreation() {
172119
// Setup mocks that will be called within the transaction (using mock.Anything for context)
173120
// Check if attached function exists (idempotency check inside transaction)
174121
suite.mockMetaDomain.On("AttachedFunctionDb", mock.Anything).Return(suite.mockAttachedFunctionDb).Once()
175-
suite.mockAttachedFunctionDb.On("GetAnyByName", inputCollectionID, attachedFunctionName).
176-
Return(nil, nil).Once()
122+
suite.mockAttachedFunctionDb.On("GetAttachedFunctions", (*uuid.UUID)(nil), &attachedFunctionName, &inputCollectionID, false).
123+
Return([]*dbmodel.AttachedFunction{}, nil).Once()
177124

178125
// Look up database
179126
suite.mockMetaDomain.On("DatabaseDb", mock.Anything).Return(suite.mockDatabaseDb).Once()
@@ -290,8 +237,8 @@ func (suite *AttachFunctionTestSuite) TestAttachFunction_IdempotentRequest_Alrea
290237

291238
// ===== Phase 1: Transaction checks if attached function exists =====
292239
suite.mockMetaDomain.On("AttachedFunctionDb", ctx).Return(suite.mockAttachedFunctionDb).Once()
293-
suite.mockAttachedFunctionDb.On("GetAnyByName", inputCollectionID, attachedFunctionName).
294-
Return(existingAttachedFunction, nil).Once()
240+
suite.mockAttachedFunctionDb.On("GetAttachedFunctions", (*uuid.UUID)(nil), &attachedFunctionName, &inputCollectionID, false).
241+
Return([]*dbmodel.AttachedFunction{existingAttachedFunction}, nil).Once()
295242

296243
// Mock transaction call
297244
suite.mockTxImpl.On("Transaction", ctx, mock.AnythingOfType("func(context.Context) error")).
@@ -375,8 +322,8 @@ func (suite *AttachFunctionTestSuite) TestAttachFunction_RecoveryFlow() {
375322

376323
// Phase 1: Create attached function in transaction
377324
suite.mockMetaDomain.On("AttachedFunctionDb", mock.Anything).Return(suite.mockAttachedFunctionDb).Once()
378-
suite.mockAttachedFunctionDb.On("GetAnyByName", inputCollectionID, attachedFunctionName).
379-
Return(nil, nil).Once()
325+
suite.mockAttachedFunctionDb.On("GetAttachedFunctions", (*uuid.UUID)(nil), &attachedFunctionName, &inputCollectionID, false).
326+
Return([]*dbmodel.AttachedFunction{}, nil).Once()
380327

381328
suite.mockMetaDomain.On("DatabaseDb", mock.Anything).Return(suite.mockDatabaseDb).Once()
382329
suite.mockDatabaseDb.On("GetDatabases", tenantID, databaseName).
@@ -428,10 +375,10 @@ func (suite *AttachFunctionTestSuite) TestAttachFunction_RecoveryFlow() {
428375

429376
// ========== SECOND ATTEMPT: Recovery Succeeds ==========
430377

431-
// Phase 0: GetByName returns incomplete attached function (with ErrAttachedFunctionNotReady, which AttachFunction handles)
378+
// Phase 0: GetAttachedFunctions returns incomplete attached function (with onlyReady=false to include all)
432379
suite.mockMetaDomain.On("AttachedFunctionDb", ctx).Return(suite.mockAttachedFunctionDb).Once()
433-
suite.mockAttachedFunctionDb.On("GetAnyByName", inputCollectionID, attachedFunctionName).
434-
Return(incompleteAttachedFunction, nil).Once()
380+
suite.mockAttachedFunctionDb.On("GetAttachedFunctions", (*uuid.UUID)(nil), &attachedFunctionName, &inputCollectionID, false).
381+
Return([]*dbmodel.AttachedFunction{incompleteAttachedFunction}, nil).Once()
435382

436383
// Validate function matches
437384
suite.mockMetaDomain.On("FunctionDb", ctx).Return(suite.mockFunctionDb).Once()
@@ -518,8 +465,8 @@ func (suite *AttachFunctionTestSuite) TestAttachFunction_IdempotentRequest_Param
518465

519466
// ===== Phase 1: Transaction checks if task exists - finds task with different params =====
520467
suite.mockMetaDomain.On("AttachedFunctionDb", mock.Anything).Return(suite.mockAttachedFunctionDb).Once()
521-
suite.mockAttachedFunctionDb.On("GetAnyByName", inputCollectionID, attachedFunctionName).
522-
Return(existingAttachedFunction, nil).Once()
468+
suite.mockAttachedFunctionDb.On("GetAttachedFunctions", (*uuid.UUID)(nil), &attachedFunctionName, &inputCollectionID, false).
469+
Return([]*dbmodel.AttachedFunction{existingAttachedFunction}, nil).Once()
523470

524471
// Validate function - returns DIFFERENT function name
525472
suite.mockMetaDomain.On("FunctionDb", mock.Anything).Return(suite.mockFunctionDb).Once()

‎go/pkg/sysdb/coordinator/heap_client_integration_test.go‎

Lines changed: 7 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -370,14 +370,16 @@ func (suite *HeapClientIntegrationTestSuite) TestPartialTaskCleanup_ThenRecreate
370370
suite.T().Logf("Cleanup completed, removed %d tasks", cleanupResp.CleanedUpCount)
371371

372372
// STEP 3: Verify task still exists and can be retrieved
373-
getResp, err := suite.sysdbClient.GetAttachedFunctionByName(ctx, &coordinatorpb.GetAttachedFunctionByNameRequest{
374-
InputCollectionId: collectionID,
375-
Name: taskName,
373+
getResp, err := suite.sysdbClient.GetAttachedFunctions(ctx, &coordinatorpb.GetAttachedFunctionsRequest{
374+
InputCollectionId: &collectionID,
375+
Name: &taskName,
376+
OnlyReady: func() *bool { b := true; return &b }(),
376377
})
377378
suite.NoError(err, "Task should still exist after cleanup")
378379
suite.NotNil(getResp)
379-
suite.Equal(taskResp.Id, getResp.AttachedFunction.Id)
380-
suite.T().Logf("Task still exists after cleanup: %s", getResp.AttachedFunction.Id)
380+
suite.Require().Len(getResp.AttachedFunctions, 1)
381+
suite.Equal(taskResp.Id, getResp.AttachedFunctions[0].Id)
382+
suite.T().Logf("Task still exists after cleanup: %s", getResp.AttachedFunctions[0].Id)
381383

382384
// STEP 4: Delete the task
383385
_, err = suite.sysdbClient.DetachFunction(ctx, &coordinatorpb.DetachFunctionRequest{

‎go/pkg/sysdb/coordinator/list_attached_functions_test.go‎

Lines changed: 12 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -84,7 +84,7 @@ func (suite *ListAttachedFunctionsTestSuite) TestListAttachedFunctions_Success()
8484
}
8585

8686
suite.mockMetaDomain.On("AttachedFunctionDb", ctx).Return(suite.mockAttachedFunctionDb).Once()
87-
suite.mockAttachedFunctionDb.On("GetByCollectionID", collectionID).Return(attachedFunctions, nil).Once()
87+
suite.mockAttachedFunctionDb.On("GetAttachedFunctions", (*uuid.UUID)(nil), (*string)(nil), &collectionID, true).Return(attachedFunctions, nil).Once()
8888

8989
functionOne := &dbmodel.Function{ID: functionID1, Name: "function-one"}
9090
functionTwo := &dbmodel.Function{ID: functionID2, Name: "function-two"}
@@ -106,8 +106,8 @@ func (suite *ListAttachedFunctionsTestSuite) TestListAttachedFunctions_Success()
106106
})).
107107
Return([]*dbmodel.Function{functionOne, functionTwo}, nil).Once()
108108

109-
req := &coordinatorpb.ListAttachedFunctionsRequest{InputCollectionId: collectionID}
110-
resp, err := suite.coordinator.ListAttachedFunctions(ctx, req)
109+
req := &coordinatorpb.GetAttachedFunctionsRequest{InputCollectionId: &collectionID, OnlyReady: func() *bool { b := true; return &b }()}
110+
resp, err := suite.coordinator.GetAttachedFunctions(ctx, req)
111111

112112
suite.Require().NoError(err)
113113
suite.Require().NotNil(resp)
@@ -128,10 +128,10 @@ func (suite *ListAttachedFunctionsTestSuite) TestListAttachedFunctions_EmptyResu
128128
collectionID := "test-collection"
129129

130130
suite.mockMetaDomain.On("AttachedFunctionDb", ctx).Return(suite.mockAttachedFunctionDb).Once()
131-
suite.mockAttachedFunctionDb.On("GetByCollectionID", collectionID).Return([]*dbmodel.AttachedFunction{}, nil).Once()
131+
suite.mockAttachedFunctionDb.On("GetAttachedFunctions", (*uuid.UUID)(nil), (*string)(nil), &collectionID, true).Return([]*dbmodel.AttachedFunction{}, nil).Once()
132132

133-
req := &coordinatorpb.ListAttachedFunctionsRequest{InputCollectionId: collectionID}
134-
resp, err := suite.coordinator.ListAttachedFunctions(ctx, req)
133+
req := &coordinatorpb.GetAttachedFunctionsRequest{InputCollectionId: &collectionID, OnlyReady: func() *bool { b := true; return &b }()}
134+
resp, err := suite.coordinator.GetAttachedFunctions(ctx, req)
135135

136136
suite.Require().NoError(err)
137137
suite.Require().NotNil(resp)
@@ -160,13 +160,13 @@ func (suite *ListAttachedFunctionsTestSuite) TestListAttachedFunctions_FunctionD
160160
}
161161

162162
suite.mockMetaDomain.On("AttachedFunctionDb", ctx).Return(suite.mockAttachedFunctionDb).Once()
163-
suite.mockAttachedFunctionDb.On("GetByCollectionID", collectionID).Return([]*dbmodel.AttachedFunction{attachedFunction}, nil).Once()
163+
suite.mockAttachedFunctionDb.On("GetAttachedFunctions", (*uuid.UUID)(nil), (*string)(nil), &collectionID, true).Return([]*dbmodel.AttachedFunction{attachedFunction}, nil).Once()
164164

165165
suite.mockMetaDomain.On("FunctionDb", ctx).Return(suite.mockFunctionDb).Once()
166166
suite.mockFunctionDb.On("GetByIDs", []uuid.UUID{functionID}).Return(nil, errors.New("db error")).Once()
167167

168-
req := &coordinatorpb.ListAttachedFunctionsRequest{InputCollectionId: collectionID}
169-
resp, err := suite.coordinator.ListAttachedFunctions(ctx, req)
168+
req := &coordinatorpb.GetAttachedFunctionsRequest{InputCollectionId: &collectionID, OnlyReady: func() *bool { b := true; return &b }()}
169+
resp, err := suite.coordinator.GetAttachedFunctions(ctx, req)
170170

171171
suite.Require().Error(err)
172172
suite.Nil(resp)
@@ -194,15 +194,15 @@ func (suite *ListAttachedFunctionsTestSuite) TestListAttachedFunctions_InvalidPa
194194
}
195195

196196
suite.mockMetaDomain.On("AttachedFunctionDb", ctx).Return(suite.mockAttachedFunctionDb).Once()
197-
suite.mockAttachedFunctionDb.On("GetByCollectionID", collectionID).Return([]*dbmodel.AttachedFunction{attachedFunction}, nil).Once()
197+
suite.mockAttachedFunctionDb.On("GetAttachedFunctions", (*uuid.UUID)(nil), (*string)(nil), &collectionID, true).Return([]*dbmodel.AttachedFunction{attachedFunction}, nil).Once()
198198

199199
functionModel := &dbmodel.Function{ID: functionID, Name: "function"}
200200

201201
suite.mockMetaDomain.On("FunctionDb", ctx).Return(suite.mockFunctionDb).Once()
202202
suite.mockFunctionDb.On("GetByIDs", []uuid.UUID{functionID}).Return([]*dbmodel.Function{functionModel}, nil).Once()
203203

204-
req := &coordinatorpb.ListAttachedFunctionsRequest{InputCollectionId: collectionID}
205-
resp, err := suite.coordinator.ListAttachedFunctions(ctx, req)
204+
req := &coordinatorpb.GetAttachedFunctionsRequest{InputCollectionId: &collectionID, OnlyReady: func() *bool { b := true; return &b }()}
205+
resp, err := suite.coordinator.GetAttachedFunctions(ctx, req)
206206

207207
suite.Require().Error(err)
208208
suite.Nil(resp)

0 commit comments

Comments
 (0)