Skip to content

Commit

Permalink
feat(run): add requester id in list model run response
Browse files Browse the repository at this point in the history
  • Loading branch information
joremysh committed Oct 23, 2024
1 parent e4d733d commit 2965a18
Show file tree
Hide file tree
Showing 8 changed files with 113 additions and 242 deletions.
2 changes: 1 addition & 1 deletion go.mod
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,7 @@ require (
github.com/grpc-ecosystem/go-grpc-middleware v1.4.0
github.com/grpc-ecosystem/grpc-gateway/v2 v2.21.0
github.com/iancoleman/strcase v0.3.0
github.com/instill-ai/protogen-go v0.3.3-alpha.0.20241012090311-e872dc0b511d
github.com/instill-ai/protogen-go v0.3.3-alpha.0.20241022025309-9afd9231a821
github.com/instill-ai/usage-client v0.3.0-alpha
github.com/instill-ai/x v0.4.0-alpha.0.20240923052503-fd243919e005
github.com/jackc/pgx/v5 v5.6.0
Expand Down
4 changes: 2 additions & 2 deletions go.sum
Original file line number Diff line number Diff line change
Expand Up @@ -229,8 +229,8 @@ github.com/hjson/hjson-go/v4 v4.0.0 h1:wlm6IYYqHjOdXH1gHev4VoXCaW20HdQAGCxdOEEg2
github.com/hjson/hjson-go/v4 v4.0.0/go.mod h1:KaYt3bTw3zhBjYqnXkYywcYctk0A2nxeEFTse3rH13E=
github.com/iancoleman/strcase v0.3.0 h1:nTXanmYxhfFAMjZL34Ov6gkzEsSJZ5DbhxWjvSASxEI=
github.com/iancoleman/strcase v0.3.0/go.mod h1:iwCmte+B7n89clKwxIoIXy/HfoL7AsD47ZCWhYzw7ho=
github.com/instill-ai/protogen-go v0.3.3-alpha.0.20241012090311-e872dc0b511d h1:jf2RQtRFNxnPMkjTD0AAqXDXO8lHYOrWU3Hrr+yGEzY=
github.com/instill-ai/protogen-go v0.3.3-alpha.0.20241012090311-e872dc0b511d/go.mod h1:rf0UY7VpEgpaLudYEcjx5rnbuwlBaaLyD4FQmWLtgAY=
github.com/instill-ai/protogen-go v0.3.3-alpha.0.20241022025309-9afd9231a821 h1:yDtTSAjeM7gkiHkL2XX4yil+VHOIm2kbyk0ZlHskRW8=
github.com/instill-ai/protogen-go v0.3.3-alpha.0.20241022025309-9afd9231a821/go.mod h1:rf0UY7VpEgpaLudYEcjx5rnbuwlBaaLyD4FQmWLtgAY=
github.com/instill-ai/usage-client v0.3.0-alpha h1:yY5eNn5zINqy8wpOogiNmrVYzJKnd1KMnMxlYBpr7Tk=
github.com/instill-ai/usage-client v0.3.0-alpha/go.mod h1:8lvtZulkhQ7t8alttb2KkLKYoCp5u4oatzDbfFlEld0=
github.com/instill-ai/x v0.4.0-alpha.0.20240923052503-fd243919e005 h1:6iUpJs545Ivvv2joJCQoawRPr2YRPjux22k7t+gWg6s=
Expand Down
4 changes: 2 additions & 2 deletions pkg/handler/mock_service_test.go

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

2 changes: 1 addition & 1 deletion pkg/handler/public.go
Original file line number Diff line number Diff line change
Expand Up @@ -1140,7 +1140,7 @@ func (h *PublicHandler) ListModelRuns(ctx context.Context, req *modelpb.ListMode
return resp, nil
}

func (h *PublicHandler) ListModelRunsByCreditOwner(ctx context.Context, req *modelpb.ListModelRunsByCreditOwnerRequest) (*modelpb.ListModelRunsByCreditOwnerResponse, error) {
func (h *PublicHandler) ListModelRunsByRequester(ctx context.Context, req *modelpb.ListModelRunsByRequesterRequest) (*modelpb.ListModelRunsByRequesterResponse, error) {

eventName := "ListModelRunsByRequester"

Expand Down
244 changes: 29 additions & 215 deletions pkg/mock/repository_mock.gen.go

Large diffs are not rendered by default.

22 changes: 16 additions & 6 deletions pkg/repository/repository.go
Original file line number Diff line number Diff line change
Expand Up @@ -71,7 +71,7 @@ type Repository interface {
ListModelRuns(ctx context.Context, pageSize, page int64, filter filtering.Filter, order ordering.OrderBy, requesterUID string, isOwner bool, modelUID string) (modelRuns []*datamodel.ModelRun, totalSize int64, err error)
CreateModelRun(ctx context.Context, modelRun *datamodel.ModelRun) (*datamodel.ModelRun, error)
UpdateModelRun(ctx context.Context, modelRun *datamodel.ModelRun) error
ListModelRunsByRequester(ctx context.Context, pageSize, page int64, filter filtering.Filter, order ordering.OrderBy, requesterUID string, startedTimeBegin, startedTimeEnd time.Time) (modelTriggers []*datamodel.ModelRun, totalSize int64, err error)
ListModelRunsByRequester(ctx context.Context, params *ListModelRunsByRequesterParams) (modelTriggers []*datamodel.ModelRun, totalSize int64, err error)
}

// DefaultPageSize is the default pagination page size when page size is not assigned
Expand Down Expand Up @@ -818,8 +818,17 @@ func (r *repository) UpdateModelRun(ctx context.Context, modelRun *datamodel.Mod
Updates(&modelRun).Error
}

func (r *repository) ListModelRunsByRequester(ctx context.Context, pageSize, page int64, filter filtering.Filter, order ordering.OrderBy,
requesterUID string, startedTimeBegin, startedTimeEnd time.Time) ([]*datamodel.ModelRun, int64, error) {
type ListModelRunsByRequesterParams struct {
PageSize int64
Page int64
Filter filtering.Filter
Order ordering.OrderBy
RequesterUID string
StartedTimeBegin time.Time
StartedTimeEnd time.Time
}

func (r *repository) ListModelRunsByRequester(ctx context.Context, params *ListModelRunsByRequesterParams) ([]*datamodel.ModelRun, int64, error) {

logger, _ := custom_logger.GetZapLogger(ctx)

Expand All @@ -830,10 +839,10 @@ func (r *repository) ListModelRunsByRequester(ctx context.Context, pageSize, pag
db := r.CheckPinnedUser(ctx, r.db, tableModelRun)

whereConditions := []string{"requester_uid = ? and create_time >= ? and create_time <= ?"}
whereArgs := []any{requesterUID, startedTimeBegin, startedTimeEnd}
whereArgs := []any{params.RequesterUID, params.StartedTimeBegin, params.StartedTimeEnd}

var expr *clause.Expr
if expr, err = r.transpileFilter(filter, tableModelRun); err != nil {
if expr, err = r.transpileFilter(params.Filter, tableModelRun); err != nil {
return nil, 0, err
}
if expr != nil {
Expand All @@ -852,6 +861,7 @@ func (r *repository) ListModelRunsByRequester(ctx context.Context, pageSize, pag
}

queryBuilder := db.Preload(clause.Associations).Where(where, whereArgs...)
order := params.Order
if order.Fields == nil || len(order.Fields) == 0 {
order.Fields = append(order.Fields, ordering.Field{
Path: "create_time",
Expand All @@ -864,7 +874,7 @@ func (r *repository) ListModelRunsByRequester(ctx context.Context, pageSize, pag
queryBuilder.Order(orderString)
}

if err = queryBuilder.Limit(int(pageSize)).Offset(int(pageSize * page)).Find(&modelRuns).Error; err != nil {
if err = queryBuilder.Limit(int(params.PageSize)).Offset(int(params.PageSize * params.Page)).Find(&modelRuns).Error; err != nil {
logger.Error("failed in querying model runs", zap.Error(err))
return nil, 0, err
}
Expand Down
35 changes: 28 additions & 7 deletions pkg/repository/repository_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -132,7 +132,7 @@ func MockNamespaceModel(t *testing.T, repo repository.Repository) *datamodel.Mod
return models[0]
}

func TestRepository_ListModelRunsByCreditOwner(t *testing.T) {
func TestRepository_ListModelRunsByRequester(t *testing.T) {
c := qt.New(t)

s, err := miniredis.Run()
Expand Down Expand Up @@ -168,24 +168,45 @@ func TestRepository_ListModelRunsByCreditOwner(t *testing.T) {
require.NoError(t, err)

t.Run("got no result in given time range", func(t *testing.T) {
resp, totalSize, err := repo.ListModelRunsByRequester(ctx, 10, 0, filtering.Filter{}, ordering.OrderBy{}, userUUID.String(),
now.Add(-2*time.Hour), now.Add(-1*time.Hour))
resp, totalSize, err := repo.ListModelRunsByRequester(ctx, &repository.ListModelRunsByRequesterParams{
PageSize: 10,
Page: 0,
Filter: filtering.Filter{},
Order: ordering.OrderBy{},
RequesterUID: userUUID.String(),
StartedTimeBegin: now.Add(-2 * time.Hour),
StartedTimeEnd: now.Add(-1 * time.Hour),
})
require.NoError(t, err)
require.Zero(t, totalSize)
require.Empty(t, resp)
})

t.Run("got 1 result in given time range", func(t *testing.T) {
resp, totalSize, err := repo.ListModelRunsByRequester(ctx, 10, 0, filtering.Filter{}, ordering.OrderBy{}, userUUID.String(),
now.Add(-1*time.Hour), now.Add(1*time.Hour))
resp, totalSize, err := repo.ListModelRunsByRequester(ctx, &repository.ListModelRunsByRequesterParams{
PageSize: 10,
Page: 0,
Filter: filtering.Filter{},
Order: ordering.OrderBy{},
RequesterUID: userUUID.String(),
StartedTimeBegin: now.Add(-1 * time.Hour),
StartedTimeEnd: now.Add(1 * time.Hour),
})
require.NoError(t, err)
require.Zero(t, totalSize-1)
require.Len(t, resp, 1)
})

t.Run("got no result with other's requester ID", func(t *testing.T) {
resp, totalSize, err := repo.ListModelRunsByRequester(ctx, 10, 0, filtering.Filter{}, ordering.OrderBy{}, uuid.Must(uuid.NewV4()).String(),
now.Add(-1*time.Hour), now.Add(1*time.Hour))
resp, totalSize, err := repo.ListModelRunsByRequester(ctx, &repository.ListModelRunsByRequesterParams{
PageSize: 10,
Page: 0,
Filter: filtering.Filter{},
Order: ordering.OrderBy{},
RequesterUID: userUUID.String(),
StartedTimeBegin: now.Add(-1 * time.Hour),
StartedTimeEnd: now.Add(1 * time.Hour),
})
require.NoError(t, err)
require.Zero(t, totalSize)
require.Empty(t, resp)
Expand Down
42 changes: 34 additions & 8 deletions pkg/service/service.go
Original file line number Diff line number Diff line change
Expand Up @@ -102,7 +102,7 @@ type Service interface {
CreateModelRun(ctx context.Context, triggerUID uuid.UUID, userUID uuid.UUID, modelUID uuid.UUID, version string, inputJSON []byte) (runLog *datamodel.ModelRun, err error)
UpdateModelRunWithError(ctx context.Context, runLog *datamodel.ModelRun, err error) *datamodel.ModelRun
ListModelRuns(ctx context.Context, req *modelpb.ListModelRunsRequest, filter filtering.Filter) (*modelpb.ListModelRunsResponse, error)
ListModelRunsByRequester(ctx context.Context, req *modelpb.ListModelRunsByCreditOwnerRequest) (*modelpb.ListModelRunsByCreditOwnerResponse, error)
ListModelRunsByRequester(ctx context.Context, req *modelpb.ListModelRunsByRequesterRequest) (*modelpb.ListModelRunsByRequesterResponse, error)
}

type service struct {
Expand Down Expand Up @@ -781,8 +781,9 @@ func (s *service) ListModelRuns(ctx context.Context, req *modelpb.ListModelRunsR
}

runnerIDMap := make(map[string]struct{})
for _, trigger := range runs {
runnerIDMap[trigger.RunnerUID.String()] = struct{}{}
for _, run := range runs {
runnerIDMap[run.RunnerUID.String()] = struct{}{}
runnerIDMap[run.RequesterUID.String()] = struct{}{}
}

runnerMap := make(map[string]*string)
Expand All @@ -799,6 +800,9 @@ func (s *service) ListModelRuns(ctx context.Context, req *modelpb.ListModelRunsR
for i, run := range runs {
pbModelRun = convertModelRunToPB(run)
pbModelRun.RunnerId = runnerMap[run.RunnerUID.String()]
if requesterID, ok := runnerMap[run.RequesterUID.String()]; ok && requesterID != nil {
pbModelRun.RequesterId = *requesterID
}

if CanViewPrivateData(run.RequesterUID.String(), requesterUID) {
pbModelRun.TaskInputs, pbModelRun.TaskOutputs, err = parseMetadataToStructArr(metadataMap, run)
Expand All @@ -819,10 +823,19 @@ func (s *service) ListModelRuns(ctx context.Context, req *modelpb.ListModelRunsR
}, nil
}

func (s *service) ListModelRunsByRequester(ctx context.Context, req *modelpb.ListModelRunsByCreditOwnerRequest) (*modelpb.ListModelRunsByCreditOwnerResponse, error) {
func (s *service) ListModelRunsByRequester(ctx context.Context, req *modelpb.ListModelRunsByRequesterRequest) (*modelpb.ListModelRunsByRequesterResponse, error) {
pageSize := s.pageSizeInRange(req.GetPageSize())
page := s.pageInRange(req.GetPage())

ns, err := s.GetRscNamespace(ctx, req.GetRequesterId())
if err != nil {
return nil, fmt.Errorf("invalid namespace: %w", err)
}

if err := s.checkNamespacePermission(ctx, ns); err != nil {
return nil, fmt.Errorf("checking namespace permissions: %w", err)
}

declarations, err := filtering.NewDeclarations([]filtering.DeclarationOption{
filtering.DeclareStandardFunctions(),
filtering.DeclareIdent("status", filtering.TypeString),
Expand Down Expand Up @@ -858,14 +871,23 @@ func (s *service) ListModelRunsByRequester(ctx context.Context, req *modelpb.Lis
return nil, fmt.Errorf("time range end time is earlier than start time")
}

runs, totalSize, err := s.repository.ListModelRunsByRequester(ctx, int64(pageSize), int64(page), filter, orderBy, requesterUID, startedTimeBegin, startedTimeEnd)
runs, totalSize, err := s.repository.ListModelRunsByRequester(ctx, &repository.ListModelRunsByRequesterParams{
PageSize: int64(pageSize),
Page: int64(page),
Filter: filter,
Order: orderBy,
RequesterUID: requesterUID,
StartedTimeBegin: startedTimeBegin,
StartedTimeEnd: startedTimeEnd,
})
if err != nil {
return nil, err
}

runnerIDMap := make(map[string]struct{})
for _, trigger := range runs {
runnerIDMap[trigger.RunnerUID.String()] = struct{}{}
for _, run := range runs {
runnerIDMap[run.RunnerUID.String()] = struct{}{}
runnerIDMap[run.RequesterUID.String()] = struct{}{}
}

runnerMap := make(map[string]*string)
Expand All @@ -883,10 +905,14 @@ func (s *service) ListModelRunsByRequester(ctx context.Context, req *modelpb.Lis
for i, run := range runs {
pbModelRun = convertModelRunToPB(run)
pbModelRun.RunnerId = runnerMap[run.RunnerUID.String()]
if requesterID, ok := runnerMap[run.RequesterUID.String()]; ok && requesterID != nil {
pbModelRun.RequesterId = *requesterID
}

pbModelRuns[i] = pbModelRun
}

return &modelpb.ListModelRunsByCreditOwnerResponse{
return &modelpb.ListModelRunsByRequesterResponse{
Runs: pbModelRuns,
TotalSize: int32(totalSize),
PageSize: pageSize,
Expand Down

0 comments on commit 2965a18

Please sign in to comment.