Movatterモバイル変換

[0]ホーム
URL:

Skip to content

Navigation Menu

Sign in
Appearance settings

Search code, repositories, users, issues, pull requests...

Provide feedback

We read every piece of feedback, and take your input very seriously.

Saved searches

Use saved searches to filter your results more quickly

 Sign up
Appearance settings

Commit50ba223

Browse files
authored
feat: add db query for setting interception ended_at field (#20437)
Adds UpdateAIBridgeInterceptionEnded query to mark interceptions asdone.Needed forcoder/internal#1051
1 parent6318520 commit50ba223

File tree

18 files changed

+196
-23
lines changed

18 files changed

+196
-23
lines changed

‎coderd/apidoc/docs.go‎

Lines changed: 4 additions & 0 deletions
Some generated files are not rendered by default. Learn more aboutcustomizing how changed files appear on GitHub.

‎coderd/apidoc/swagger.json‎

Lines changed: 4 additions & 0 deletions
Some generated files are not rendered by default. Learn more aboutcustomizing how changed files appear on GitHub.

‎coderd/database/db2sdk/db2sdk.go‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -962,7 +962,7 @@ func AIBridgeInterception(interception database.AIBridgeInterception, initiator
962962
// created_at ASC
963963
returnsdkToolUsages[i].CreatedAt.Before(sdkToolUsages[j].CreatedAt)
964964
})
965-
return codersdk.AIBridgeInterception{
965+
intc:= codersdk.AIBridgeInterception{
966966
ID:interception.ID,
967967
Initiator:MinimalUserFromVisibleUser(initiator),
968968
Provider:interception.Provider,
@@ -973,6 +973,10 @@ func AIBridgeInterception(interception database.AIBridgeInterception, initiator
973973
UserPrompts:sdkUserPrompts,
974974
ToolUsages:sdkToolUsages,
975975
}
976+
ifinterception.EndedAt.Valid {
977+
intc.EndedAt=&interception.EndedAt.Time
978+
}
979+
returnintc
976980
}
977981

978982
funcAIBridgeTokenUsage(usage database.AIBridgeTokenUsage) codersdk.AIBridgeTokenUsage {

‎coderd/database/dbauthz/dbauthz.go‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4711,6 +4711,13 @@ func (q *querier) UnfavoriteWorkspace(ctx context.Context, id uuid.UUID) error {
47114711
returnupdate(q.log,q.auth,fetch,q.db.UnfavoriteWorkspace)(ctx,id)
47124712
}
47134713

4714+
func (q*querier)UpdateAIBridgeInterceptionEnded(ctx context.Context,params database.UpdateAIBridgeInterceptionEndedParams) (database.AIBridgeInterception,error) {
4715+
iferr:=q.authorizeAIBridgeInterceptionAction(ctx,policy.ActionUpdate,params.ID);err!=nil {
4716+
return database.AIBridgeInterception{},err
4717+
}
4718+
returnq.db.UpdateAIBridgeInterceptionEnded(ctx,params)
4719+
}
4720+
47144721
func (q*querier)UpdateAPIKeyByID(ctx context.Context,arg database.UpdateAPIKeyByIDParams)error {
47154722
fetch:=func(ctx context.Context,arg database.UpdateAPIKeyByIDParams) (database.APIKey,error) {
47164723
returnq.db.GetAPIKeyByID(ctx,arg.ID)

‎coderd/database/dbauthz/dbauthz_test.go‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4617,4 +4617,13 @@ func (s *MethodTestSuite) TestAIBridge() {
46174617
db.EXPECT().ListAIBridgeToolUsagesByInterceptionIDs(gomock.Any(),ids).Return([]database.AIBridgeToolUsage{},nil).AnyTimes()
46184618
check.Args(ids).Asserts(rbac.ResourceSystem,policy.ActionRead).Returns([]database.AIBridgeToolUsage{})
46194619
}))
4620+
4621+
s.Run("UpdateAIBridgeInterceptionEnded",s.Mocked(func(db*dbmock.MockStore,faker*gofakeit.Faker,check*expects) {
4622+
intcID:= uuid.UUID{1}
4623+
params:= database.UpdateAIBridgeInterceptionEndedParams{ID:intcID}
4624+
intc:=testutil.Fake(s.T(),faker, database.AIBridgeInterception{ID:intcID})
4625+
db.EXPECT().GetAIBridgeInterceptionByID(gomock.Any(),intcID).Return(intc,nil).AnyTimes()// Validation.
4626+
db.EXPECT().UpdateAIBridgeInterceptionEnded(gomock.Any(),params).Return(intc,nil).AnyTimes()
4627+
check.Args(params).Asserts(intc,policy.ActionUpdate).Returns(intc)
4628+
}))
46204629
}

‎coderd/database/dbgen/dbgen.go‎

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1495,7 +1495,7 @@ func ClaimPrebuild(
14951495
returnclaimedWorkspace
14961496
}
14971497

1498-
funcAIBridgeInterception(t testing.TB,db database.Store,seed database.InsertAIBridgeInterceptionParams) database.AIBridgeInterception {
1498+
funcAIBridgeInterception(t testing.TB,db database.Store,seed database.InsertAIBridgeInterceptionParams,endedAt*time.Time) database.AIBridgeInterception {
14991499
interception,err:=db.InsertAIBridgeInterception(genCtx, database.InsertAIBridgeInterceptionParams{
15001500
ID:takeFirst(seed.ID,uuid.New()),
15011501
InitiatorID:takeFirst(seed.InitiatorID,uuid.New()),
@@ -1504,6 +1504,13 @@ func AIBridgeInterception(t testing.TB, db database.Store, seed database.InsertA
15041504
Metadata:takeFirstSlice(seed.Metadata,json.RawMessage("{}")),
15051505
StartedAt:takeFirst(seed.StartedAt,dbtime.Now()),
15061506
})
1507+
ifendedAt!=nil {
1508+
interception,err=db.UpdateAIBridgeInterceptionEnded(genCtx, database.UpdateAIBridgeInterceptionEndedParams{
1509+
ID:interception.ID,
1510+
EndedAt:*endedAt,
1511+
})
1512+
require.NoError(t,err,"insert aibridge interception")
1513+
}
15071514
require.NoError(t,err,"insert aibridge interception")
15081515
returninterception
15091516
}

‎coderd/database/dbmetrics/querymetrics.go‎

Lines changed: 7 additions & 0 deletions
Some generated files are not rendered by default. Learn more aboutcustomizing how changed files appear on GitHub.

‎coderd/database/dbmock/dbmock.go‎

Lines changed: 15 additions & 0 deletions
Some generated files are not rendered by default. Learn more aboutcustomizing how changed files appear on GitHub.

‎coderd/database/querier.go‎

Lines changed: 1 addition & 0 deletions
Some generated files are not rendered by default. Learn more aboutcustomizing how changed files appear on GitHub.

‎coderd/database/querier_test.go‎

Lines changed: 65 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7724,3 +7724,68 @@ func TestUpdateTaskWorkspaceID(t *testing.T) {
77247724
})
77257725
}
77267726
}
7727+
7728+
funcTestUpdateAIBridgeInterceptionEnded(t*testing.T) {
7729+
t.Parallel()
7730+
db,_:=dbtestutil.NewDB(t)
7731+
7732+
t.Run("NonExistingInterception",func(t*testing.T) {
7733+
t.Parallel()
7734+
ctx:=testutil.Context(t,testutil.WaitLong)
7735+
7736+
got,err:=db.UpdateAIBridgeInterceptionEnded(ctx, database.UpdateAIBridgeInterceptionEndedParams{
7737+
ID:uuid.New(),
7738+
EndedAt:time.Now(),
7739+
})
7740+
require.ErrorContains(t,err,"no rows in result set")
7741+
require.EqualValues(t, database.AIBridgeInterception{},got)
7742+
})
7743+
7744+
t.Run("OK",func(t*testing.T) {
7745+
t.Parallel()
7746+
ctx:=testutil.Context(t,testutil.WaitLong)
7747+
7748+
user:=dbgen.User(t,db, database.User{})
7749+
interceptions:= []database.AIBridgeInterception{}
7750+
7751+
for_,uid:=range []uuid.UUID{{1}, {2}, {3}} {
7752+
insertParams:= database.InsertAIBridgeInterceptionParams{
7753+
ID:uid,
7754+
InitiatorID:user.ID,
7755+
Metadata:json.RawMessage("{}"),
7756+
}
7757+
7758+
intc,err:=db.InsertAIBridgeInterception(ctx,insertParams)
7759+
require.NoError(t,err)
7760+
require.Equal(t,uid,intc.ID)
7761+
require.False(t,intc.EndedAt.Valid)
7762+
interceptions=append(interceptions,intc)
7763+
}
7764+
7765+
intc0:=interceptions[0]
7766+
endedAt:=time.Now()
7767+
// Mark first interception as done
7768+
updated,err:=db.UpdateAIBridgeInterceptionEnded(ctx, database.UpdateAIBridgeInterceptionEndedParams{
7769+
ID:intc0.ID,
7770+
EndedAt:endedAt,
7771+
})
7772+
require.NoError(t,err)
7773+
require.EqualValues(t,updated.ID,intc0.ID)
7774+
require.True(t,updated.EndedAt.Valid)
7775+
require.WithinDuration(t,endedAt,updated.EndedAt.Time,5*time.Second)
7776+
7777+
// Updating first interception again should fail
7778+
updated,err=db.UpdateAIBridgeInterceptionEnded(ctx, database.UpdateAIBridgeInterceptionEndedParams{
7779+
ID:intc0.ID,
7780+
EndedAt:endedAt.Add(time.Hour),
7781+
})
7782+
require.ErrorIs(t,err,sql.ErrNoRows)
7783+
7784+
// Other interceptions should not have ended_at set
7785+
for_,intc:=rangeinterceptions[1:] {
7786+
got,err:=db.GetAIBridgeInterceptionByID(ctx,intc.ID)
7787+
require.NoError(t,err)
7788+
require.False(t,got.EndedAt.Valid)
7789+
}
7790+
})
7791+
}

0 commit comments

Comments
 (0)
[8]ページ先頭
©2009-2025 Movatter.jp