Skip to content

Commit 151680b

Browse files
tanujnay112tanujnay112
committed
[CHORE]: Consolidate all attached function getters
1 parent b6edf1a commit 151680b

File tree

15 files changed

+277
-494
lines changed

15 files changed

+277
-494
lines changed

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

Lines changed: 13 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -67,15 +67,10 @@ func (suite *AttachFunctionTestSuite) setupAttachFunctionMocks(ctx context.Conte
6767
databaseName := request.Database
6868
functionName := request.FunctionName
6969

70-
// Phase 0: No existing attached function
71-
suite.mockMetaDomain.On("AttachedFunctionDb", ctx).Return(suite.mockAttachedFunctionDb).Once()
72-
suite.mockAttachedFunctionDb.On("GetByName", inputCollectionID, attachedFunctionName).
73-
Return(nil, nil).Once()
74-
7570
// Phase 1: Create attached function in transaction
7671
suite.mockMetaDomain.On("AttachedFunctionDb", mock.Anything).Return(suite.mockAttachedFunctionDb).Once()
77-
suite.mockAttachedFunctionDb.On("GetByName", inputCollectionID, attachedFunctionName).
78-
Return(nil, nil).Once()
72+
suite.mockAttachedFunctionDb.On("GetAttachedFunctions", (*uuid.UUID)(nil), &attachedFunctionName, &inputCollectionID, false).
73+
Return([]*dbmodel.AttachedFunction{}, nil).Once()
7974

8075
suite.mockMetaDomain.On("DatabaseDb", mock.Anything).Return(suite.mockDatabaseDb).Once()
8176
suite.mockDatabaseDb.On("GetDatabases", tenantID, databaseName).
@@ -172,8 +167,8 @@ func (suite *AttachFunctionTestSuite) TestAttachFunction_SuccessfulCreation() {
172167
// Setup mocks that will be called within the transaction (using mock.Anything for context)
173168
// Check if attached function exists (idempotency check inside transaction)
174169
suite.mockMetaDomain.On("AttachedFunctionDb", mock.Anything).Return(suite.mockAttachedFunctionDb).Once()
175-
suite.mockAttachedFunctionDb.On("GetByName", inputCollectionID, attachedFunctionName).
176-
Return(nil, nil).Once()
170+
suite.mockAttachedFunctionDb.On("GetAttachedFunctions", (*uuid.UUID)(nil), &attachedFunctionName, &inputCollectionID, true).
171+
Return([]*dbmodel.AttachedFunction{}, nil).Once()
177172

178173
// Look up database
179174
suite.mockMetaDomain.On("DatabaseDb", mock.Anything).Return(suite.mockDatabaseDb).Once()
@@ -290,8 +285,8 @@ func (suite *AttachFunctionTestSuite) TestAttachFunction_IdempotentRequest_Alrea
290285

291286
// ===== Phase 1: Transaction checks if attached function exists =====
292287
suite.mockMetaDomain.On("AttachedFunctionDb", ctx).Return(suite.mockAttachedFunctionDb).Once()
293-
suite.mockAttachedFunctionDb.On("GetByName", inputCollectionID, attachedFunctionName).
294-
Return(existingAttachedFunction, nil).Once()
288+
suite.mockAttachedFunctionDb.On("GetAttachedFunctions", (*uuid.UUID)(nil), &attachedFunctionName, &inputCollectionID, false).
289+
Return([]*dbmodel.AttachedFunction{existingAttachedFunction}, nil).Once()
295290

296291
// Mock transaction call
297292
suite.mockTxImpl.On("Transaction", ctx, mock.AnythingOfType("func(context.Context) error")).
@@ -375,8 +370,8 @@ func (suite *AttachFunctionTestSuite) TestAttachFunction_RecoveryFlow() {
375370

376371
// Phase 1: Create attached function in transaction
377372
suite.mockMetaDomain.On("AttachedFunctionDb", mock.Anything).Return(suite.mockAttachedFunctionDb).Once()
378-
suite.mockAttachedFunctionDb.On("GetByName", inputCollectionID, attachedFunctionName).
379-
Return(nil, nil).Once()
373+
suite.mockAttachedFunctionDb.On("GetAttachedFunctions", (*uuid.UUID)(nil), &attachedFunctionName, &inputCollectionID, false).
374+
Return([]*dbmodel.AttachedFunction{}, nil).Once()
380375

381376
suite.mockMetaDomain.On("DatabaseDb", mock.Anything).Return(suite.mockDatabaseDb).Once()
382377
suite.mockDatabaseDb.On("GetDatabases", tenantID, databaseName).
@@ -428,10 +423,10 @@ func (suite *AttachFunctionTestSuite) TestAttachFunction_RecoveryFlow() {
428423

429424
// ========== SECOND ATTEMPT: Recovery Succeeds ==========
430425

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

436431
// Validate function matches
437432
suite.mockMetaDomain.On("FunctionDb", ctx).Return(suite.mockFunctionDb).Once()
@@ -518,8 +513,8 @@ func (suite *AttachFunctionTestSuite) TestAttachFunction_IdempotentRequest_Param
518513

519514
// ===== Phase 1: Transaction checks if task exists - finds task with different params =====
520515
suite.mockMetaDomain.On("AttachedFunctionDb", mock.Anything).Return(suite.mockAttachedFunctionDb).Once()
521-
suite.mockAttachedFunctionDb.On("GetByName", inputCollectionID, attachedFunctionName).
522-
Return(existingAttachedFunction, nil).Once()
516+
suite.mockAttachedFunctionDb.On("GetAttachedFunctions", (*uuid.UUID)(nil), &attachedFunctionName, &inputCollectionID, true).
517+
Return([]*dbmodel.AttachedFunction{existingAttachedFunction}, nil).Once()
523518

524519
// Validate function - returns DIFFERENT function name
525520
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)