shishantbiswas/coder
1
0
Fork 0
mirror of https://github.com/coder/coder.git synced 2026-06-02 20:48:20 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: batch connection logs to avoid DB lock contention (#23727)

Browse Source 
- Running 30k connections was generating a ton of lock contention in the
DB
This commit is contained in:
Jon Ayers
2026-04-03 15:47:26 -05:00
committed by GitHub
parent 333503f74e
commit a1d51f0dab
21 changed files with 2168 additions and 426 deletions
Download Patch File Download Diff File Expand all files Collapse all files
coderd/database/dbauthz/dbauthz.go
+7 -7
View File
@@ -1627,6 +1627,13 @@ func (q *querier) BatchUpdateWorkspaceNextStartAt(ctx context.Context, arg datab
return q.db.BatchUpdateWorkspaceNextStartAt(ctx, arg)
}
func (q *querier) BatchUpsertConnectionLogs(ctx context.Context, arg database.BatchUpsertConnectionLogsParams) error {
if err := q.authorizeContext(ctx, policy.ActionUpdate, rbac.ResourceConnectionLog); err != nil {
return err
}
return q.db.BatchUpsertConnectionLogs(ctx, arg)
}
func (q *querier) BulkMarkNotificationMessagesFailed(ctx context.Context, arg database.BulkMarkNotificationMessagesFailedParams) (int64, error) {
if err := q.authorizeContext(ctx, policy.ActionUpdate, rbac.ResourceNotificationMessage); err != nil {
return 0, err
@@ -7065,13 +7072,6 @@ func (q *querier) UpsertChatWorkspaceTTL(ctx context.Context, workspaceTtl strin
return q.db.UpsertChatWorkspaceTTL(ctx, workspaceTtl)
}
func (q *querier) UpsertConnectionLog(ctx context.Context, arg database.UpsertConnectionLogParams) (database.ConnectionLog, error) {
if err := q.authorizeContext(ctx, policy.ActionUpdate, rbac.ResourceConnectionLog); err != nil {
return database.ConnectionLog{}, err
}
return q.db.UpsertConnectionLog(ctx, arg)
}
func (q *querier) UpsertDefaultProxy(ctx context.Context, arg database.UpsertDefaultProxyParams) error {
if err := q.authorizeContext(ctx, policy.ActionUpdate, rbac.ResourceSystem); err != nil {
return err
coderd/database/dbauthz/dbauthz_test.go
+3 -4
View File
@@ -338,10 +338,9 @@ func (s *MethodTestSuite) TestAuditLogs() {
}
func (s *MethodTestSuite) TestConnectionLogs() {
s.Run("UpsertConnectionLog", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
ws := testutil.Fake(s.T(), faker, database.WorkspaceTable{})
arg := database.UpsertConnectionLogParams{Ip: defaultIPAddress(), Type: database.ConnectionTypeSsh, WorkspaceID: ws.ID, OrganizationID: ws.OrganizationID, ConnectionStatus: database.ConnectionStatusConnected, WorkspaceOwnerID: ws.OwnerID}
dbm.EXPECT().UpsertConnectionLog(gomock.Any(), arg).Return(database.ConnectionLog{}, nil).AnyTimes()
s.Run("BatchUpsertConnectionLogs", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) {
arg := database.BatchUpsertConnectionLogsParams{}
dbm.EXPECT().BatchUpsertConnectionLogs(gomock.Any(), arg).Return(nil).AnyTimes()
check.Args(arg).Asserts(rbac.ResourceConnectionLog, policy.ActionUpdate)
}))
s.Run("GetConnectionLogsOffset", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) {
coderd/database/dbgen/dbgen.go
+47 -3
View File
@@ -76,7 +76,7 @@ func AuditLog(t testing.TB, db database.Store, seed database.AuditLog) database.
}
func ConnectionLog(t testing.TB, db database.Store, seed database.UpsertConnectionLogParams) database.ConnectionLog {
log, err := db.UpsertConnectionLog(genCtx, database.UpsertConnectionLogParams{
arg := database.UpsertConnectionLogParams{
ID: takeFirst(seed.ID, uuid.New()),
Time: takeFirst(seed.Time, dbtime.Now()),
OrganizationID: takeFirst(seed.OrganizationID, uuid.New()),
@@ -89,7 +89,7 @@ func ConnectionLog(t testing.TB, db database.Store, seed database.UpsertConnecti
Int32: takeFirst(seed.Code.Int32, 0),
Valid: takeFirst(seed.Code.Valid, false),
},
Ip: pqtype.Inet{
IP: pqtype.Inet{
IPNet: net.IPNet{
IP: net.IPv4(127, 0, 0, 1),
Mask: net.IPv4Mask(255, 255, 255, 255),
@@ -117,9 +117,53 @@ func ConnectionLog(t testing.TB, db database.Store, seed database.UpsertConnecti
Valid: takeFirst(seed.DisconnectReason.Valid, false),
},
ConnectionStatus: takeFirst(seed.ConnectionStatus, database.ConnectionStatusConnected),
}
var disconnectTime sql.NullTime
if arg.ConnectionStatus == database.ConnectionStatusDisconnected {
disconnectTime = sql.NullTime{Time: arg.Time, Valid: true}
}
err := db.BatchUpsertConnectionLogs(genCtx, database.BatchUpsertConnectionLogsParams{
ID: []uuid.UUID{arg.ID},
ConnectTime: []time.Time{arg.Time},
OrganizationID: []uuid.UUID{arg.OrganizationID},
WorkspaceOwnerID: []uuid.UUID{arg.WorkspaceOwnerID},
WorkspaceID: []uuid.UUID{arg.WorkspaceID},
WorkspaceName: []string{arg.WorkspaceName},
AgentName: []string{arg.AgentName},
Type: []database.ConnectionType{arg.Type},
Code: []int32{arg.Code.Int32},
CodeValid: []bool{arg.Code.Valid},
Ip: []pqtype.Inet{arg.IP},
UserAgent: []string{arg.UserAgent.String},
UserID: []uuid.UUID{arg.UserID.UUID},
SlugOrPort: []string{arg.SlugOrPort.String},
ConnectionID: []uuid.UUID{arg.ConnectionID.UUID},
DisconnectReason: []string{arg.DisconnectReason.String},
DisconnectTime: []time.Time{disconnectTime.Time},
})
require.NoError(t, err, "insert connection log")
return log
// Query back the actual row from the database. On upsert
// conflict the DB keeps the original row's ID, so we can't
// rely on arg.ID. Match on the conflict key for rows with a
// connection_id, or by primary key for NULL connection_id.
rows, err := db.GetConnectionLogsOffset(genCtx, database.GetConnectionLogsOffsetParams{})
require.NoError(t, err, "query connection logs")
for _, row := range rows {
if arg.ConnectionID.Valid {
if row.ConnectionLog.ConnectionID == arg.ConnectionID &&
row.ConnectionLog.WorkspaceID == arg.WorkspaceID &&
row.ConnectionLog.AgentName == arg.AgentName {
return row.ConnectionLog
}
} else if row.ConnectionLog.ID == arg.ID {
return row.ConnectionLog
}
}
require.Failf(t, "connection log not found", "id=%s", arg.ID)
return database.ConnectionLog{} // unreachable
}
func Template(t testing.TB, db database.Store, seed database.Template) database.Template {
coderd/database/dbmetrics/querymetrics.go
+8 -8
View File
@@ -208,6 +208,14 @@ func (m queryMetricsStore) BatchUpdateWorkspaceNextStartAt(ctx context.Context,
return r0
}
func (m queryMetricsStore) BatchUpsertConnectionLogs(ctx context.Context, arg database.BatchUpsertConnectionLogsParams) error {
start := time.Now()
r0 := m.s.BatchUpsertConnectionLogs(ctx, arg)
m.queryLatencies.WithLabelValues("BatchUpsertConnectionLogs").Observe(time.Since(start).Seconds())
m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "BatchUpsertConnectionLogs").Inc()
return r0
}
func (m queryMetricsStore) BulkMarkNotificationMessagesFailed(ctx context.Context, arg database.BulkMarkNotificationMessagesFailedParams) (int64, error) {
start := time.Now()
r0, r1 := m.s.BulkMarkNotificationMessagesFailed(ctx, arg)
@@ -5024,14 +5032,6 @@ func (m queryMetricsStore) UpsertChatWorkspaceTTL(ctx context.Context, workspace
return r0
}
func (m queryMetricsStore) UpsertConnectionLog(ctx context.Context, arg database.UpsertConnectionLogParams) (database.ConnectionLog, error) {
start := time.Now()
r0, r1 := m.s.UpsertConnectionLog(ctx, arg)
m.queryLatencies.WithLabelValues("UpsertConnectionLog").Observe(time.Since(start).Seconds())
m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "UpsertConnectionLog").Inc()
return r0, r1
}
func (m queryMetricsStore) UpsertDefaultProxy(ctx context.Context, arg database.UpsertDefaultProxyParams) error {
start := time.Now()
r0 := m.s.UpsertDefaultProxy(ctx, arg)
coderd/database/dbmock/dbmock.go
+14 -15
View File
@@ -233,6 +233,20 @@ func (mr *MockStoreMockRecorder) BatchUpdateWorkspaceNextStartAt(ctx, arg any) *
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "BatchUpdateWorkspaceNextStartAt", reflect.TypeOf((*MockStore)(nil).BatchUpdateWorkspaceNextStartAt), ctx, arg)
}
// BatchUpsertConnectionLogs mocks base method.
func (m *MockStore) BatchUpsertConnectionLogs(ctx context.Context, arg database.BatchUpsertConnectionLogsParams) error {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "BatchUpsertConnectionLogs", ctx, arg)
ret0, _ := ret[0].(error)
return ret0
}
// BatchUpsertConnectionLogs indicates an expected call of BatchUpsertConnectionLogs.
func (mr *MockStoreMockRecorder) BatchUpsertConnectionLogs(ctx, arg any) *gomock.Call {
mr.mock.ctrl.T.Helper()
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "BatchUpsertConnectionLogs", reflect.TypeOf((*MockStore)(nil).BatchUpsertConnectionLogs), ctx, arg)
}
// BulkMarkNotificationMessagesFailed mocks base method.
func (m *MockStore) BulkMarkNotificationMessagesFailed(ctx context.Context, arg database.BulkMarkNotificationMessagesFailedParams) (int64, error) {
m.ctrl.T.Helper()
@@ -9442,21 +9456,6 @@ func (mr *MockStoreMockRecorder) UpsertChatWorkspaceTTL(ctx, workspaceTtl any) *
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpsertChatWorkspaceTTL", reflect.TypeOf((*MockStore)(nil).UpsertChatWorkspaceTTL), ctx, workspaceTtl)
}
// UpsertConnectionLog mocks base method.
func (m *MockStore) UpsertConnectionLog(ctx context.Context, arg database.UpsertConnectionLogParams) (database.ConnectionLog, error) {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "UpsertConnectionLog", ctx, arg)
ret0, _ := ret[0].(database.ConnectionLog)
ret1, _ := ret[1].(error)
return ret0, ret1
}
// UpsertConnectionLog indicates an expected call of UpsertConnectionLog.
func (mr *MockStoreMockRecorder) UpsertConnectionLog(ctx, arg any) *gomock.Call {
mr.mock.ctrl.T.Helper()
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpsertConnectionLog", reflect.TypeOf((*MockStore)(nil).UpsertConnectionLog), ctx, arg)
}
// UpsertDefaultProxy mocks base method.
func (m *MockStore) UpsertDefaultProxy(ctx context.Context, arg database.UpsertDefaultProxyParams) error {
m.ctrl.T.Helper()
coderd/database/modelmethods.go
+26
View File
@@ -10,6 +10,7 @@ import (
"time"
"github.com/google/uuid"
"github.com/sqlc-dev/pqtype"
"golang.org/x/exp/maps"
"golang.org/x/oauth2"
"golang.org/x/xerrors"
@@ -923,3 +924,28 @@ func WorkspaceIdentityFromWorkspace(w Workspace) WorkspaceIdentity {
func (r GetWorkspaceAgentAndWorkspaceByIDRow) RBACObject() rbac.Object {
return r.WorkspaceTable.RBACObject()
}
// UpsertConnectionLogParams contains the parameters for upserting a
// connection log entry. This struct is hand-maintained (not generated
// by sqlc) because the single-row UpsertConnectionLog query was
// removed in favor of BatchUpsertConnectionLogs, but the struct is
// still used as the canonical connection log event type throughout
// the codebase.
type UpsertConnectionLogParams struct {
ID uuid.UUID `db:"id" json:"id"`
OrganizationID uuid.UUID `db:"organization_id" json:"organization_id"`
WorkspaceOwnerID uuid.UUID `db:"workspace_owner_id" json:"workspace_owner_id"`
WorkspaceID uuid.UUID `db:"workspace_id" json:"workspace_id"`
WorkspaceName string `db:"workspace_name" json:"workspace_name"`
AgentName string `db:"agent_name" json:"agent_name"`
Type ConnectionType `db:"type" json:"type"`
Code sql.NullInt32 `db:"code" json:"code"`
IP pqtype.Inet `db:"ip" json:"ip"`
UserAgent sql.NullString `db:"user_agent" json:"user_agent"`
UserID uuid.NullUUID `db:"user_id" json:"user_id"`
SlugOrPort sql.NullString `db:"slug_or_port" json:"slug_or_port"`
ConnectionID uuid.NullUUID `db:"connection_id" json:"connection_id"`
DisconnectReason sql.NullString `db:"disconnect_reason" json:"disconnect_reason"`
Time time.Time `db:"time" json:"time"`
ConnectionStatus ConnectionStatus `db:"connection_status" json:"connection_status"`
}
coderd/database/querier.go
+1 -1
View File
@@ -65,6 +65,7 @@ type sqlcQuerier interface {
BatchUpdateWorkspaceAgentMetadata(ctx context.Context, arg BatchUpdateWorkspaceAgentMetadataParams) error
BatchUpdateWorkspaceLastUsedAt(ctx context.Context, arg BatchUpdateWorkspaceLastUsedAtParams) error
BatchUpdateWorkspaceNextStartAt(ctx context.Context, arg BatchUpdateWorkspaceNextStartAtParams) error
BatchUpsertConnectionLogs(ctx context.Context, arg BatchUpsertConnectionLogsParams) error
BulkMarkNotificationMessagesFailed(ctx context.Context, arg BulkMarkNotificationMessagesFailedParams) (int64, error)
BulkMarkNotificationMessagesSent(ctx context.Context, arg BulkMarkNotificationMessagesSentParams) (int64, error)
// Calculates the telemetry summary for a given provider, model, and client
@@ -991,7 +992,6 @@ type sqlcQuerier interface {
UpsertChatUsageLimitGroupOverride(ctx context.Context, arg UpsertChatUsageLimitGroupOverrideParams) (UpsertChatUsageLimitGroupOverrideRow, error)
UpsertChatUsageLimitUserOverride(ctx context.Context, arg UpsertChatUsageLimitUserOverrideParams) (UpsertChatUsageLimitUserOverrideRow, error)
UpsertChatWorkspaceTTL(ctx context.Context, workspaceTtl string) error
UpsertConnectionLog(ctx context.Context, arg UpsertConnectionLogParams) (ConnectionLog, error)
// The default proxy is implied and not actually stored in the database.
// So we need to store it's configuration here for display purposes.
// The functional values are immutable and controlled implicitly.
coderd/database/querier_test.go
+482 -197
View File
@@ -3566,9 +3566,11 @@ func connectionOnlyIDs[T database.ConnectionLog | database.GetConnectionLogsOffs
return ids
}
func TestUpsertConnectionLog(t *testing.T) {
func TestBatchUpsertConnectionLogs(t *testing.T) {
t.Parallel()
createWorkspace := func(t *testing.T, db database.Store) database.WorkspaceTable {
t.Helper()
u := dbgen.User(t, db, database.User{})
o := dbgen.Organization(t, db, database.Organization{})
tpl := dbgen.Template(t, db, database.Template{
@@ -3584,253 +3586,536 @@ func TestUpsertConnectionLog(t *testing.T) {
})
}
// zeroTime is the sentinel value that the SQL treats as "no
// connect/disconnect time provided".
zeroTime := time.Time{}
defaultIP := pqtype.Inet{
IPNet: net.IPNet{
IP: net.IPv4(127, 0, 0, 1),
Mask: net.IPv4Mask(255, 255, 255, 255),
},
Valid: true,
}
t.Run("SingleConnect", func(t *testing.T) {
t.Parallel()
db, _ := dbtestutil.NewDB(t)
ctx := context.Background()
ws := createWorkspace(t, db)
connID := uuid.New()
connectTime := dbtime.Now()
err := db.BatchUpsertConnectionLogs(ctx, database.BatchUpsertConnectionLogsParams{
ID: []uuid.UUID{uuid.New()},
ConnectTime: []time.Time{connectTime},
OrganizationID: []uuid.UUID{ws.OrganizationID},
WorkspaceOwnerID: []uuid.UUID{ws.OwnerID},
WorkspaceID: []uuid.UUID{ws.ID},
WorkspaceName: []string{ws.Name},
AgentName: []string{"agent"},
Type: []database.ConnectionType{database.ConnectionTypeSsh},
Code: []int32{0},
CodeValid: []bool{false},
Ip: []pqtype.Inet{defaultIP},
UserAgent: []string{""},
UserID: []uuid.UUID{uuid.Nil},
SlugOrPort: []string{""},
ConnectionID: []uuid.UUID{connID},
DisconnectReason: []string{""},
DisconnectTime: []time.Time{zeroTime},
})
require.NoError(t, err)
rows, err := db.GetConnectionLogsOffset(ctx, database.GetConnectionLogsOffsetParams{LimitOpt: 10})
require.NoError(t, err)
require.Len(t, rows, 1)
require.True(t, connectTime.Equal(rows[0].ConnectionLog.ConnectTime))
require.False(t, rows[0].ConnectionLog.DisconnectTime.Valid,
"disconnect_time should be NULL for a connect-only event")
})
t.Run("ConnectThenDisconnect", func(t *testing.T) {
t.Parallel()
db, _ := dbtestutil.NewDB(t)
ctx := context.Background()
ws := createWorkspace(t, db)
connectionID := uuid.New()
agentName := "test-agent"
// 1. Insert a 'connect' event.
connID := uuid.New()
connectTime := dbtime.Now()
connectParams := database.UpsertConnectionLogParams{
ID: uuid.New(),
Time: connectTime,
OrganizationID: ws.OrganizationID,
WorkspaceOwnerID: ws.OwnerID,
WorkspaceID: ws.ID,
WorkspaceName: ws.Name,
AgentName: agentName,
Type: database.ConnectionTypeSsh,
ConnectionID: uuid.NullUUID{UUID: connectionID, Valid: true},
ConnectionStatus: database.ConnectionStatusConnected,
Ip: pqtype.Inet{
IPNet: net.IPNet{
IP: net.IPv4(127, 0, 0, 1),
Mask: net.IPv4Mask(255, 255, 255, 255),
},
Valid: true,
},
}
log1, err := db.UpsertConnectionLog(ctx, connectParams)
// Insert connect.
err := db.BatchUpsertConnectionLogs(ctx, database.BatchUpsertConnectionLogsParams{
ID: []uuid.UUID{uuid.New()},
ConnectTime: []time.Time{connectTime},
OrganizationID: []uuid.UUID{ws.OrganizationID},
WorkspaceOwnerID: []uuid.UUID{ws.OwnerID},
WorkspaceID: []uuid.UUID{ws.ID},
WorkspaceName: []string{ws.Name},
AgentName: []string{"agent"},
Type: []database.ConnectionType{database.ConnectionTypeSsh},
Code: []int32{0},
CodeValid: []bool{false},
Ip: []pqtype.Inet{defaultIP},
UserAgent: []string{""},
UserID: []uuid.UUID{uuid.Nil},
SlugOrPort: []string{""},
ConnectionID: []uuid.UUID{connID},
DisconnectReason: []string{""},
DisconnectTime: []time.Time{zeroTime},
})
require.NoError(t, err)
// Insert disconnect for same connection.
disconnectTime := connectTime.Add(time.Second)
err = db.BatchUpsertConnectionLogs(ctx, database.BatchUpsertConnectionLogsParams{
ID: []uuid.UUID{uuid.New()},
ConnectTime: []time.Time{zeroTime},
OrganizationID: []uuid.UUID{ws.OrganizationID},
WorkspaceOwnerID: []uuid.UUID{ws.OwnerID},
WorkspaceID: []uuid.UUID{ws.ID},
WorkspaceName: []string{ws.Name},
AgentName: []string{"agent"},
Type: []database.ConnectionType{database.ConnectionTypeSsh},
Code: []int32{1},
CodeValid: []bool{true},
Ip: []pqtype.Inet{defaultIP},
UserAgent: []string{""},
UserID: []uuid.UUID{uuid.Nil},
SlugOrPort: []string{""},
ConnectionID: []uuid.UUID{connID},
DisconnectReason: []string{"test disconnect"},
DisconnectTime: []time.Time{disconnectTime},
})
require.NoError(t, err)
require.Equal(t, connectParams.ID, log1.ID)
require.False(t, log1.DisconnectTime.Valid, "DisconnectTime should not be set on connect")
// Check that one row exists.
rows, err := db.GetConnectionLogsOffset(ctx, database.GetConnectionLogsOffsetParams{LimitOpt: 10})
require.NoError(t, err)
require.Len(t, rows, 1)
// 2. Insert a 'disconnected' event for the same connection.
disconnectTime := connectTime.Add(time.Second)
disconnectParams := database.UpsertConnectionLogParams{
ConnectionID: uuid.NullUUID{UUID: connectionID, Valid: true},
WorkspaceID: ws.ID,
AgentName: agentName,
ConnectionStatus: database.ConnectionStatusDisconnected,
// Updated to:
Time: disconnectTime,
DisconnectReason: sql.NullString{String: "test disconnect", Valid: true},
Code: sql.NullInt32{Int32: 1, Valid: true},
// Ignored
ID: uuid.New(),
OrganizationID: ws.OrganizationID,
WorkspaceOwnerID: ws.OwnerID,
WorkspaceName: ws.Name,
Type: database.ConnectionTypeSsh,
Ip: pqtype.Inet{
IPNet: net.IPNet{
IP: net.IPv4(127, 0, 0, 1),
Mask: net.IPv4Mask(255, 255, 255, 254),
},
Valid: true,
},
}
log2, err := db.UpsertConnectionLog(ctx, disconnectParams)
require.NoError(t, err)
// Updated
require.Equal(t, log1.ID, log2.ID)
require.True(t, log2.DisconnectTime.Valid)
require.True(t, disconnectTime.Equal(log2.DisconnectTime.Time))
require.Equal(t, disconnectParams.DisconnectReason.String, log2.DisconnectReason.String)
rows, err = db.GetConnectionLogsOffset(ctx, database.GetConnectionLogsOffsetParams{})
require.NoError(t, err)
require.Len(t, rows, 1)
row := rows[0].ConnectionLog
require.True(t, connectTime.Equal(row.ConnectTime))
require.True(t, row.DisconnectTime.Valid)
require.True(t, disconnectTime.Equal(row.DisconnectTime.Time))
require.Equal(t, "test disconnect", row.DisconnectReason.String)
require.Equal(t, int32(1), row.Code.Int32)
})
t.Run("ConnectDoesNotUpdate", func(t *testing.T) {
t.Run("DuplicateConnectIsNoOp", func(t *testing.T) {
t.Parallel()
db, _ := dbtestutil.NewDB(t)
ctx := context.Background()
ws := createWorkspace(t, db)
connectionID := uuid.New()
agentName := "test-agent"
// 1. Insert a 'connect' event.
connID := uuid.New()
connectTime := dbtime.Now()
connectParams := database.UpsertConnectionLogParams{
ID: uuid.New(),
Time: connectTime,
OrganizationID: ws.OrganizationID,
WorkspaceOwnerID: ws.OwnerID,
WorkspaceID: ws.ID,
WorkspaceName: ws.Name,
AgentName: agentName,
Type: database.ConnectionTypeSsh,
ConnectionID: uuid.NullUUID{UUID: connectionID, Valid: true},
ConnectionStatus: database.ConnectionStatusConnected,
Ip: pqtype.Inet{
IPNet: net.IPNet{
IP: net.IPv4(127, 0, 0, 1),
Mask: net.IPv4Mask(255, 255, 255, 255),
},
Valid: true,
},
mkParams := func(ct time.Time, ip pqtype.Inet) database.BatchUpsertConnectionLogsParams {
return database.BatchUpsertConnectionLogsParams{
ID: []uuid.UUID{uuid.New()},
ConnectTime: []time.Time{ct},
OrganizationID: []uuid.UUID{ws.OrganizationID},
WorkspaceOwnerID: []uuid.UUID{ws.OwnerID},
WorkspaceID: []uuid.UUID{ws.ID},
WorkspaceName: []string{ws.Name},
AgentName: []string{"agent"},
Type: []database.ConnectionType{database.ConnectionTypeSsh},
Code: []int32{0},
CodeValid: []bool{false},
Ip: []pqtype.Inet{ip},
UserAgent: []string{""},
UserID: []uuid.UUID{uuid.Nil},
SlugOrPort: []string{""},
ConnectionID: []uuid.UUID{connID},
DisconnectReason: []string{""},
DisconnectTime: []time.Time{zeroTime},
}
}
log, err := db.UpsertConnectionLog(ctx, connectParams)
err := db.BatchUpsertConnectionLogs(ctx, mkParams(connectTime, defaultIP))
require.NoError(t, err)
// 2. Insert another 'connect' event for the same connection.
connectTime2 := connectTime.Add(time.Second)
connectParams2 := database.UpsertConnectionLogParams{
ConnectionID: uuid.NullUUID{UUID: connectionID, Valid: true},
WorkspaceID: ws.ID,
AgentName: agentName,
ConnectionStatus: database.ConnectionStatusConnected,
rows1, err := db.GetConnectionLogsOffset(ctx, database.GetConnectionLogsOffsetParams{LimitOpt: 10})
require.NoError(t, err)
require.Len(t, rows1, 1)
// Ignored
ID: uuid.New(),
Time: connectTime2,
OrganizationID: ws.OrganizationID,
WorkspaceOwnerID: ws.OwnerID,
WorkspaceName: ws.Name,
Type: database.ConnectionTypeSsh,
Code: sql.NullInt32{Int32: 0, Valid: false},
Ip: pqtype.Inet{
IPNet: net.IPNet{
IP: net.IPv4(127, 0, 0, 1),
Mask: net.IPv4Mask(255, 255, 255, 254),
},
Valid: true,
// Second connect with later time and different IP.
otherIP := pqtype.Inet{
IPNet: net.IPNet{
IP: net.IPv4(10, 0, 0, 1),
Mask: net.IPv4Mask(255, 255, 255, 255),
},
Valid: true,
}
origLog, err := db.UpsertConnectionLog(ctx, connectParams2)
err = db.BatchUpsertConnectionLogs(ctx, mkParams(connectTime.Add(time.Second), otherIP))
require.NoError(t, err)
require.Equal(t, log, origLog, "connect update should be a no-op")
// Check that still only one row exists.
rows, err := db.GetConnectionLogsOffset(ctx, database.GetConnectionLogsOffsetParams{})
rows2, err := db.GetConnectionLogsOffset(ctx, database.GetConnectionLogsOffsetParams{LimitOpt: 10})
require.NoError(t, err)
require.Len(t, rows, 1)
require.Equal(t, log, rows[0].ConnectionLog)
require.Len(t, rows2, 1)
// The LEAST logic should pick the earlier connect_time; IP and
// other fields are not updated on conflict.
require.True(t, connectTime.Equal(rows2[0].ConnectionLog.ConnectTime),
"connect_time should remain the original (earlier) value")
})
t.Run("DisconnectThenConnect", func(t *testing.T) {
t.Run("OrderIndependentConnectTime", func(t *testing.T) {
t.Parallel()
db, _ := dbtestutil.NewDB(t)
ctx := context.Background()
ws := createWorkspace(t, db)
connectionID := uuid.New()
agentName := "test-agent"
// Insert just a 'disconect' event
connID := uuid.New()
disconnectTime := dbtime.Now()
disconnectParams := database.UpsertConnectionLogParams{
ID: uuid.New(),
Time: disconnectTime,
OrganizationID: ws.OrganizationID,
WorkspaceOwnerID: ws.OwnerID,
WorkspaceID: ws.ID,
WorkspaceName: ws.Name,
AgentName: agentName,
Type: database.ConnectionTypeSsh,
ConnectionID: uuid.NullUUID{UUID: connectionID, Valid: true},
ConnectionStatus: database.ConnectionStatusDisconnected,
DisconnectReason: sql.NullString{String: "server shutting down", Valid: true},
Ip: pqtype.Inet{
IPNet: net.IPNet{
IP: net.IPv4(127, 0, 0, 1),
Mask: net.IPv4Mask(255, 255, 255, 255),
},
Valid: true,
},
connectTime := disconnectTime.Add(-5 * time.Second)
// Disconnect arrives first.
err := db.BatchUpsertConnectionLogs(ctx, database.BatchUpsertConnectionLogsParams{
ID: []uuid.UUID{uuid.New()},
ConnectTime: []time.Time{disconnectTime},
OrganizationID: []uuid.UUID{ws.OrganizationID},
WorkspaceOwnerID: []uuid.UUID{ws.OwnerID},
WorkspaceID: []uuid.UUID{ws.ID},
WorkspaceName: []string{ws.Name},
AgentName: []string{"agent"},
Type: []database.ConnectionType{database.ConnectionTypeSsh},
Code: []int32{0},
CodeValid: []bool{true},
Ip: []pqtype.Inet{defaultIP},
UserAgent: []string{""},
UserID: []uuid.UUID{uuid.Nil},
SlugOrPort: []string{""},
ConnectionID: []uuid.UUID{connID},
DisconnectReason: []string{"bye"},
DisconnectTime: []time.Time{disconnectTime},
})
require.NoError(t, err)
// Connect arrives second with the real (earlier) connect_time.
err = db.BatchUpsertConnectionLogs(ctx, database.BatchUpsertConnectionLogsParams{
ID: []uuid.UUID{uuid.New()},
ConnectTime: []time.Time{connectTime},
OrganizationID: []uuid.UUID{ws.OrganizationID},
WorkspaceOwnerID: []uuid.UUID{ws.OwnerID},
WorkspaceID: []uuid.UUID{ws.ID},
WorkspaceName: []string{ws.Name},
AgentName: []string{"agent"},
Type: []database.ConnectionType{database.ConnectionTypeSsh},
Code: []int32{0},
CodeValid: []bool{false},
Ip: []pqtype.Inet{defaultIP},
UserAgent: []string{""},
UserID: []uuid.UUID{uuid.Nil},
SlugOrPort: []string{""},
ConnectionID: []uuid.UUID{connID},
DisconnectReason: []string{""},
DisconnectTime: []time.Time{zeroTime},
})
require.NoError(t, err)
rows, err := db.GetConnectionLogsOffset(ctx, database.GetConnectionLogsOffsetParams{LimitOpt: 10})
require.NoError(t, err)
require.Len(t, rows, 1)
require.True(t, connectTime.Equal(rows[0].ConnectionLog.ConnectTime),
"LEAST should pick the earlier connect_time")
})
t.Run("DisconnectFieldsAreWriteOnce", func(t *testing.T) {
t.Parallel()
db, _ := dbtestutil.NewDB(t)
ctx := context.Background()
ws := createWorkspace(t, db)
connID := uuid.New()
disconnectTime := dbtime.Now()
mkDisconnect := func(reason string, code int32) database.BatchUpsertConnectionLogsParams {
return database.BatchUpsertConnectionLogsParams{
ID: []uuid.UUID{uuid.New()},
ConnectTime: []time.Time{disconnectTime},
OrganizationID: []uuid.UUID{ws.OrganizationID},
WorkspaceOwnerID: []uuid.UUID{ws.OwnerID},
WorkspaceID: []uuid.UUID{ws.ID},
WorkspaceName: []string{ws.Name},
AgentName: []string{"agent"},
Type: []database.ConnectionType{database.ConnectionTypeSsh},
Code: []int32{code},
CodeValid: []bool{true},
Ip: []pqtype.Inet{defaultIP},
UserAgent: []string{""},
UserID: []uuid.UUID{uuid.Nil},
SlugOrPort: []string{""},
ConnectionID: []uuid.UUID{connID},
DisconnectReason: []string{reason},
DisconnectTime: []time.Time{disconnectTime},
}
}
_, err := db.UpsertConnectionLog(ctx, disconnectParams)
err := db.BatchUpsertConnectionLogs(ctx, mkDisconnect("first reason", 1))
require.NoError(t, err)
firstRows, err := db.GetConnectionLogsOffset(ctx, database.GetConnectionLogsOffsetParams{})
// Second disconnect with different reason and code.
err = db.BatchUpsertConnectionLogs(ctx, mkDisconnect("second reason", 2))
require.NoError(t, err)
require.Len(t, firstRows, 1)
// We expect the connection event to be marked as closed with the start
// and close time being the same.
require.True(t, firstRows[0].ConnectionLog.DisconnectTime.Valid)
require.Equal(t, disconnectTime, firstRows[0].ConnectionLog.DisconnectTime.Time.UTC())
require.Equal(t, firstRows[0].ConnectionLog.ConnectTime.UTC(), firstRows[0].ConnectionLog.DisconnectTime.Time.UTC())
rows, err := db.GetConnectionLogsOffset(ctx, database.GetConnectionLogsOffsetParams{LimitOpt: 10})
require.NoError(t, err)
require.Len(t, rows, 1)
row := rows[0].ConnectionLog
require.Equal(t, "first reason", row.DisconnectReason.String,
"disconnect_reason should not be overwritten")
require.Equal(t, int32(1), row.Code.Int32,
"code should not be overwritten")
})
// Now insert a 'connect' event for the same connection.
// This should be a no op
connectTime := disconnectTime.Add(time.Second)
connectParams := database.UpsertConnectionLogParams{
ID: uuid.New(),
Time: connectTime,
OrganizationID: ws.OrganizationID,
WorkspaceOwnerID: ws.OwnerID,
WorkspaceID: ws.ID,
WorkspaceName: ws.Name,
AgentName: agentName,
Type: database.ConnectionTypeSsh,
ConnectionID: uuid.NullUUID{UUID: connectionID, Valid: true},
ConnectionStatus: database.ConnectionStatusConnected,
DisconnectReason: sql.NullString{String: "reconnected", Valid: true},
Code: sql.NullInt32{Int32: 0, Valid: false},
Ip: pqtype.Inet{
IPNet: net.IPNet{
IP: net.IPv4(127, 0, 0, 1),
Mask: net.IPv4Mask(255, 255, 255, 255),
},
Valid: true,
},
t.Run("ConnectAfterDisconnectIsNoOp", func(t *testing.T) {
t.Parallel()
db, _ := dbtestutil.NewDB(t)
ctx := context.Background()
ws := createWorkspace(t, db)
connID := uuid.New()
disconnectTime := dbtime.Now()
// Insert disconnect first.
err := db.BatchUpsertConnectionLogs(ctx, database.BatchUpsertConnectionLogsParams{
ID: []uuid.UUID{uuid.New()},
ConnectTime: []time.Time{disconnectTime},
OrganizationID: []uuid.UUID{ws.OrganizationID},
WorkspaceOwnerID: []uuid.UUID{ws.OwnerID},
WorkspaceID: []uuid.UUID{ws.ID},
WorkspaceName: []string{ws.Name},
AgentName: []string{"agent"},
Type: []database.ConnectionType{database.ConnectionTypeSsh},
Code: []int32{42},
CodeValid: []bool{true},
Ip: []pqtype.Inet{defaultIP},
UserAgent: []string{""},
UserID: []uuid.UUID{uuid.Nil},
SlugOrPort: []string{""},
ConnectionID: []uuid.UUID{connID},
DisconnectReason: []string{"server shutdown"},
DisconnectTime: []time.Time{disconnectTime},
})
require.NoError(t, err)
rows1, err := db.GetConnectionLogsOffset(ctx, database.GetConnectionLogsOffsetParams{LimitOpt: 10})
require.NoError(t, err)
require.Len(t, rows1, 1)
require.True(t, rows1[0].ConnectionLog.DisconnectTime.Valid)
require.Equal(t, "server shutdown", rows1[0].ConnectionLog.DisconnectReason.String)
require.Equal(t, int32(42), rows1[0].ConnectionLog.Code.Int32)
// Insert connect for same connection_id.
err = db.BatchUpsertConnectionLogs(ctx, database.BatchUpsertConnectionLogsParams{
ID: []uuid.UUID{uuid.New()},
ConnectTime: []time.Time{disconnectTime.Add(time.Second)},
OrganizationID: []uuid.UUID{ws.OrganizationID},
WorkspaceOwnerID: []uuid.UUID{ws.OwnerID},
WorkspaceID: []uuid.UUID{ws.ID},
WorkspaceName: []string{ws.Name},
AgentName: []string{"agent"},
Type: []database.ConnectionType{database.ConnectionTypeSsh},
Code: []int32{0},
CodeValid: []bool{false},
Ip: []pqtype.Inet{defaultIP},
UserAgent: []string{""},
UserID: []uuid.UUID{uuid.Nil},
SlugOrPort: []string{""},
ConnectionID: []uuid.UUID{connID},
DisconnectReason: []string{""},
DisconnectTime: []time.Time{zeroTime},
})
require.NoError(t, err)
rows2, err := db.GetConnectionLogsOffset(ctx, database.GetConnectionLogsOffsetParams{LimitOpt: 10})
require.NoError(t, err)
require.Len(t, rows2, 1)
row := rows2[0].ConnectionLog
require.True(t, row.DisconnectTime.Valid,
"disconnect_time should not be cleared by a later connect")
require.Equal(t, "server shutdown", row.DisconnectReason.String,
"disconnect_reason should not be cleared")
require.Equal(t, int32(42), row.Code.Int32,
"code should not be cleared")
})
t.Run("CodeZeroPreserved", func(t *testing.T) {
t.Parallel()
db, _ := dbtestutil.NewDB(t)
ctx := context.Background()
ws := createWorkspace(t, db)
connID := uuid.New()
now := dbtime.Now()
err := db.BatchUpsertConnectionLogs(ctx, database.BatchUpsertConnectionLogsParams{
ID: []uuid.UUID{uuid.New()},
ConnectTime: []time.Time{now},
OrganizationID: []uuid.UUID{ws.OrganizationID},
WorkspaceOwnerID: []uuid.UUID{ws.OwnerID},
WorkspaceID: []uuid.UUID{ws.ID},
WorkspaceName: []string{ws.Name},
AgentName: []string{"agent"},
Type: []database.ConnectionType{database.ConnectionTypeSsh},
Code: []int32{0},
CodeValid: []bool{true},
Ip: []pqtype.Inet{defaultIP},
UserAgent: []string{""},
UserID: []uuid.UUID{uuid.Nil},
SlugOrPort: []string{""},
ConnectionID: []uuid.UUID{connID},
DisconnectReason: []string{"normal"},
DisconnectTime: []time.Time{now},
})
require.NoError(t, err)
rows, err := db.GetConnectionLogsOffset(ctx, database.GetConnectionLogsOffsetParams{LimitOpt: 10})
require.NoError(t, err)
require.Len(t, rows, 1)
require.True(t, rows[0].ConnectionLog.Code.Valid, "code should be non-NULL")
require.Equal(t, int32(0), rows[0].ConnectionLog.Code.Int32,
"code=0 should be preserved, not treated as NULL")
})
t.Run("CodeNullWhenInvalid", func(t *testing.T) {
t.Parallel()
db, _ := dbtestutil.NewDB(t)
ctx := context.Background()
ws := createWorkspace(t, db)
connID := uuid.New()
now := dbtime.Now()
err := db.BatchUpsertConnectionLogs(ctx, database.BatchUpsertConnectionLogsParams{
ID: []uuid.UUID{uuid.New()},
ConnectTime: []time.Time{now},
OrganizationID: []uuid.UUID{ws.OrganizationID},
WorkspaceOwnerID: []uuid.UUID{ws.OwnerID},
WorkspaceID: []uuid.UUID{ws.ID},
WorkspaceName: []string{ws.Name},
AgentName: []string{"agent"},
Type: []database.ConnectionType{database.ConnectionTypeSsh},
Code: []int32{99},
CodeValid: []bool{false},
Ip: []pqtype.Inet{defaultIP},
UserAgent: []string{""},
UserID: []uuid.UUID{uuid.Nil},
SlugOrPort: []string{""},
ConnectionID: []uuid.UUID{connID},
DisconnectReason: []string{""},
DisconnectTime: []time.Time{zeroTime},
})
require.NoError(t, err)
rows, err := db.GetConnectionLogsOffset(ctx, database.GetConnectionLogsOffsetParams{LimitOpt: 10})
require.NoError(t, err)
require.Len(t, rows, 1)
require.False(t, rows[0].ConnectionLog.Code.Valid,
"code should be NULL when code_valid is false")
})
t.Run("NullConnectionIDEvents", func(t *testing.T) {
t.Parallel()
db, _ := dbtestutil.NewDB(t)
ctx := context.Background()
ws := createWorkspace(t, db)
now := dbtime.Now()
// Insert two web events with NULL connection_id (uuid.Nil →
// NULL via NULLIF) for the same workspace/agent.
for i := range 2 {
err := db.BatchUpsertConnectionLogs(ctx, database.BatchUpsertConnectionLogsParams{
ID: []uuid.UUID{uuid.New()},
ConnectTime: []time.Time{now.Add(time.Duration(i) * time.Second)},
OrganizationID: []uuid.UUID{ws.OrganizationID},
WorkspaceOwnerID: []uuid.UUID{ws.OwnerID},
WorkspaceID: []uuid.UUID{ws.ID},
WorkspaceName: []string{ws.Name},
AgentName: []string{"agent"},
Type: []database.ConnectionType{database.ConnectionTypeSsh},
Code: []int32{200},
CodeValid: []bool{true},
Ip: []pqtype.Inet{defaultIP},
UserAgent: []string{"Mozilla/5.0"},
UserID: []uuid.UUID{uuid.Nil},
SlugOrPort: []string{"web-terminal"},
ConnectionID: []uuid.UUID{uuid.Nil},
DisconnectReason: []string{""},
DisconnectTime: []time.Time{zeroTime},
})
require.NoError(t, err)
}
_, err = db.UpsertConnectionLog(ctx, connectParams)
rows, err := db.GetConnectionLogsOffset(ctx, database.GetConnectionLogsOffsetParams{LimitOpt: 10})
require.NoError(t, err)
require.Len(t, rows, 2,
"NULL connection_id rows should not conflict with each other")
})
secondRows, err := db.GetConnectionLogsOffset(ctx, database.GetConnectionLogsOffsetParams{})
require.NoError(t, err)
require.Len(t, secondRows, 1)
require.Equal(t, firstRows, secondRows)
t.Run("MultipleIndependentConnections", func(t *testing.T) {
t.Parallel()
db, _ := dbtestutil.NewDB(t)
ctx := context.Background()
ws := createWorkspace(t, db)
now := dbtime.Now()
// Upsert a disconnection, which should also be a no op
disconnectParams.DisconnectReason = sql.NullString{
String: "updated close reason",
Valid: true,
n := 5
ids := make([]uuid.UUID, n)
connectTimes := make([]time.Time, n)
orgIDs := make([]uuid.UUID, n)
ownerIDs := make([]uuid.UUID, n)
wsIDs := make([]uuid.UUID, n)
wsNames := make([]string, n)
agentNames := make([]string, n)
types := make([]database.ConnectionType, n)
codes := make([]int32, n)
codeValids := make([]bool, n)
ips := make([]pqtype.Inet, n)
userAgents := make([]string, n)
userIDs := make([]uuid.UUID, n)
slugOrPorts := make([]string, n)
connIDs := make([]uuid.UUID, n)
disconnectReasons := make([]string, n)
disconnectTimes := make([]time.Time, n)
for i := range n {
ids[i] = uuid.New()
connectTimes[i] = now.Add(time.Duration(i) * time.Second)
orgIDs[i] = ws.OrganizationID
ownerIDs[i] = ws.OwnerID
wsIDs[i] = ws.ID
wsNames[i] = ws.Name
agentNames[i] = "agent"
types[i] = database.ConnectionTypeSsh
codes[i] = 0
codeValids[i] = false
ips[i] = defaultIP
userAgents[i] = ""
userIDs[i] = uuid.Nil
slugOrPorts[i] = ""
connIDs[i] = uuid.New()
disconnectReasons[i] = ""
disconnectTimes[i] = zeroTime
}
_, err = db.UpsertConnectionLog(ctx, disconnectParams)
err := db.BatchUpsertConnectionLogs(ctx, database.BatchUpsertConnectionLogsParams{
ID: ids,
ConnectTime: connectTimes,
OrganizationID: orgIDs,
WorkspaceOwnerID: ownerIDs,
WorkspaceID: wsIDs,
WorkspaceName: wsNames,
AgentName: agentNames,
Type: types,
Code: codes,
CodeValid: codeValids,
Ip: ips,
UserAgent: userAgents,
UserID: userIDs,
SlugOrPort: slugOrPorts,
ConnectionID: connIDs,
DisconnectReason: disconnectReasons,
DisconnectTime: disconnectTimes,
})
require.NoError(t, err)
thirdRows, err := db.GetConnectionLogsOffset(ctx, database.GetConnectionLogsOffsetParams{})
rows, err := db.GetConnectionLogsOffset(ctx, database.GetConnectionLogsOffsetParams{LimitOpt: 10})
require.NoError(t, err)
require.Len(t, secondRows, 1)
// The close reason shouldn't be updated
require.Equal(t, secondRows, thirdRows)
require.Len(t, rows, n, "each unique connection_id should produce its own row")
})
}
coderd/database/queries.sql.go
+117 -114
View File
@@ -7338,6 +7338,123 @@ func (q *sqlQuerier) UpsertChatUsageLimitUserOverride(ctx context.Context, arg U
return i, err
}
const batchUpsertConnectionLogs = `-- name: BatchUpsertConnectionLogs :exec
INSERT INTO connection_logs (
id, connect_time, organization_id, workspace_owner_id, workspace_id,
workspace_name, agent_name, type, code, ip, user_agent, user_id,
slug_or_port, connection_id, disconnect_reason, disconnect_time
)
SELECT
u.id,
u.connect_time,
u.organization_id,
u.workspace_owner_id,
u.workspace_id,
u.workspace_name,
u.agent_name,
u.type,
-- Use the validity flag to distinguish "no code" (NULL) from a
-- legitimate zero exit code.
CASE WHEN u.code_valid THEN u.code ELSE NULL END,
u.ip,
NULLIF(u.user_agent, ''),
NULLIF(u.user_id, '00000000-0000-0000-0000-000000000000'::uuid),
NULLIF(u.slug_or_port, ''),
NULLIF(u.connection_id, '00000000-0000-0000-0000-000000000000'::uuid),
NULLIF(u.disconnect_reason, ''),
NULLIF(u.disconnect_time, '0001-01-01 00:00:00Z'::timestamptz)
FROM (
SELECT
unnest($1::uuid[]) AS id,
unnest($2::timestamptz[]) AS connect_time,
unnest($3::uuid[]) AS organization_id,
unnest($4::uuid[]) AS workspace_owner_id,
unnest($5::uuid[]) AS workspace_id,
unnest($6::text[]) AS workspace_name,
unnest($7::text[]) AS agent_name,
unnest($8::connection_type[]) AS type,
unnest($9::int4[]) AS code,
unnest($10::bool[]) AS code_valid,
unnest($11::inet[]) AS ip,
unnest($12::text[]) AS user_agent,
unnest($13::uuid[]) AS user_id,
unnest($14::text[]) AS slug_or_port,
unnest($15::uuid[]) AS connection_id,
unnest($16::text[]) AS disconnect_reason,
unnest($17::timestamptz[]) AS disconnect_time
) AS u
ON CONFLICT (connection_id, workspace_id, agent_name)
DO UPDATE SET
-- Pick the earliest real connect_time. The zero sentinel
-- ('0001-01-01') means the batch didn't know the connect_time
-- (e.g. a pure disconnect event), so we keep the existing value.
connect_time = CASE
WHEN EXCLUDED.connect_time = '0001-01-01 00:00:00Z'::timestamptz
THEN connection_logs.connect_time
WHEN connection_logs.connect_time = '0001-01-01 00:00:00Z'::timestamptz
THEN EXCLUDED.connect_time
ELSE LEAST(connection_logs.connect_time, EXCLUDED.connect_time)
END,
disconnect_time = CASE
WHEN connection_logs.disconnect_time IS NULL
THEN EXCLUDED.disconnect_time
ELSE connection_logs.disconnect_time
END,
disconnect_reason = CASE
WHEN connection_logs.disconnect_reason IS NULL
THEN EXCLUDED.disconnect_reason
ELSE connection_logs.disconnect_reason
END,
code = CASE
WHEN connection_logs.code IS NULL
THEN EXCLUDED.code
ELSE connection_logs.code
END
`
type BatchUpsertConnectionLogsParams struct {
ID []uuid.UUID `db:"id" json:"id"`
ConnectTime []time.Time `db:"connect_time" json:"connect_time"`
OrganizationID []uuid.UUID `db:"organization_id" json:"organization_id"`
WorkspaceOwnerID []uuid.UUID `db:"workspace_owner_id" json:"workspace_owner_id"`
WorkspaceID []uuid.UUID `db:"workspace_id" json:"workspace_id"`
WorkspaceName []string `db:"workspace_name" json:"workspace_name"`
AgentName []string `db:"agent_name" json:"agent_name"`
Type []ConnectionType `db:"type" json:"type"`
Code []int32 `db:"code" json:"code"`
CodeValid []bool `db:"code_valid" json:"code_valid"`
Ip []pqtype.Inet `db:"ip" json:"ip"`
UserAgent []string `db:"user_agent" json:"user_agent"`
UserID []uuid.UUID `db:"user_id" json:"user_id"`
SlugOrPort []string `db:"slug_or_port" json:"slug_or_port"`
ConnectionID []uuid.UUID `db:"connection_id" json:"connection_id"`
DisconnectReason []string `db:"disconnect_reason" json:"disconnect_reason"`
DisconnectTime []time.Time `db:"disconnect_time" json:"disconnect_time"`
}
func (q *sqlQuerier) BatchUpsertConnectionLogs(ctx context.Context, arg BatchUpsertConnectionLogsParams) error {
_, err := q.db.ExecContext(ctx, batchUpsertConnectionLogs,
pq.Array(arg.ID),
pq.Array(arg.ConnectTime),
pq.Array(arg.OrganizationID),
pq.Array(arg.WorkspaceOwnerID),
pq.Array(arg.WorkspaceID),
pq.Array(arg.WorkspaceName),
pq.Array(arg.AgentName),
pq.Array(arg.Type),
pq.Array(arg.Code),
pq.Array(arg.CodeValid),
pq.Array(arg.Ip),
pq.Array(arg.UserAgent),
pq.Array(arg.UserID),
pq.Array(arg.SlugOrPort),
pq.Array(arg.ConnectionID),
pq.Array(arg.DisconnectReason),
pq.Array(arg.DisconnectTime),
)
return err
}
const countConnectionLogs = `-- name: CountConnectionLogs :one
SELECT
COUNT(*) AS count
@@ -7753,120 +7870,6 @@ func (q *sqlQuerier) GetConnectionLogsOffset(ctx context.Context, arg GetConnect
return items, nil
}
const upsertConnectionLog = `-- name: UpsertConnectionLog :one
INSERT INTO connection_logs (
id,
connect_time,
organization_id,
workspace_owner_id,
workspace_id,
workspace_name,
agent_name,
type,
code,
ip,
user_agent,
user_id,
slug_or_port,
connection_id,
disconnect_reason,
disconnect_time
) VALUES
($1, $15, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14,
-- If we've only received a disconnect event, mark the event as immediately
-- closed.
CASE
WHEN $16::connection_status = 'disconnected'
THEN $15 :: timestamp with time zone
ELSE NULL
END)
ON CONFLICT (connection_id, workspace_id, agent_name)
DO UPDATE SET
-- No-op if the connection is still open.
disconnect_time = CASE
WHEN $16::connection_status = 'disconnected'
-- Can only be set once
AND connection_logs.disconnect_time IS NULL
THEN EXCLUDED.connect_time
ELSE connection_logs.disconnect_time
END,
disconnect_reason = CASE
WHEN $16::connection_status = 'disconnected'
-- Can only be set once
AND connection_logs.disconnect_reason IS NULL
THEN EXCLUDED.disconnect_reason
ELSE connection_logs.disconnect_reason
END,
code = CASE
WHEN $16::connection_status = 'disconnected'
-- Can only be set once
AND connection_logs.code IS NULL
THEN EXCLUDED.code
ELSE connection_logs.code
END
RETURNING id, connect_time, organization_id, workspace_owner_id, workspace_id, workspace_name, agent_name, type, ip, code, user_agent, user_id, slug_or_port, connection_id, disconnect_time, disconnect_reason
`
type UpsertConnectionLogParams struct {
ID uuid.UUID `db:"id" json:"id"`
OrganizationID uuid.UUID `db:"organization_id" json:"organization_id"`
WorkspaceOwnerID uuid.UUID `db:"workspace_owner_id" json:"workspace_owner_id"`
WorkspaceID uuid.UUID `db:"workspace_id" json:"workspace_id"`
WorkspaceName string `db:"workspace_name" json:"workspace_name"`
AgentName string `db:"agent_name" json:"agent_name"`
Type ConnectionType `db:"type" json:"type"`
Code sql.NullInt32 `db:"code" json:"code"`
Ip pqtype.Inet `db:"ip" json:"ip"`
UserAgent sql.NullString `db:"user_agent" json:"user_agent"`
UserID uuid.NullUUID `db:"user_id" json:"user_id"`
SlugOrPort sql.NullString `db:"slug_or_port" json:"slug_or_port"`
ConnectionID uuid.NullUUID `db:"connection_id" json:"connection_id"`
DisconnectReason sql.NullString `db:"disconnect_reason" json:"disconnect_reason"`
Time time.Time `db:"time" json:"time"`
ConnectionStatus ConnectionStatus `db:"connection_status" json:"connection_status"`
}
func (q *sqlQuerier) UpsertConnectionLog(ctx context.Context, arg UpsertConnectionLogParams) (ConnectionLog, error) {
row := q.db.QueryRowContext(ctx, upsertConnectionLog,
arg.ID,
arg.OrganizationID,
arg.WorkspaceOwnerID,
arg.WorkspaceID,
arg.WorkspaceName,
arg.AgentName,
arg.Type,
arg.Code,
arg.Ip,
arg.UserAgent,
arg.UserID,
arg.SlugOrPort,
arg.ConnectionID,
arg.DisconnectReason,
arg.Time,
arg.ConnectionStatus,
)
var i ConnectionLog
err := row.Scan(
&i.ID,
&i.ConnectTime,
&i.OrganizationID,
&i.WorkspaceOwnerID,
&i.WorkspaceID,
&i.WorkspaceName,
&i.AgentName,
&i.Type,
&i.Ip,
&i.Code,
&i.UserAgent,
&i.UserID,
&i.SlugOrPort,
&i.ConnectionID,
&i.DisconnectTime,
&i.DisconnectReason,
)
return i, err
}
const deleteCryptoKey = `-- name: DeleteCryptoKey :one
UPDATE crypto_keys
SET secret = NULL, secret_key_id = NULL
coderd/database/queries/connectionlogs.sql
+69 -49
View File
@@ -251,55 +251,75 @@ DELETE FROM connection_logs
USING old_logs
WHERE connection_logs.id = old_logs.id;
-- name: UpsertConnectionLog :one
-- name: BatchUpsertConnectionLogs :exec
INSERT INTO connection_logs (
id,
connect_time,
organization_id,
workspace_owner_id,
workspace_id,
workspace_name,
agent_name,
type,
code,
ip,
user_agent,
user_id,
slug_or_port,
connection_id,
disconnect_reason,
disconnect_time
) VALUES
($1, @time, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14,
-- If we've only received a disconnect event, mark the event as immediately
-- closed.
CASE
WHEN @connection_status::connection_status = 'disconnected'
THEN @time :: timestamp with time zone
ELSE NULL
END)
id, connect_time, organization_id, workspace_owner_id, workspace_id,
workspace_name, agent_name, type, code, ip, user_agent, user_id,
slug_or_port, connection_id, disconnect_reason, disconnect_time
)
SELECT
u.id,
u.connect_time,
u.organization_id,
u.workspace_owner_id,
u.workspace_id,
u.workspace_name,
u.agent_name,
u.type,
-- Use the validity flag to distinguish "no code" (NULL) from a
-- legitimate zero exit code.
CASE WHEN u.code_valid THEN u.code ELSE NULL END,
u.ip,
NULLIF(u.user_agent, ''),
NULLIF(u.user_id, '00000000-0000-0000-0000-000000000000'::uuid),
NULLIF(u.slug_or_port, ''),
NULLIF(u.connection_id, '00000000-0000-0000-0000-000000000000'::uuid),
NULLIF(u.disconnect_reason, ''),
NULLIF(u.disconnect_time, '0001-01-01 00:00:00Z'::timestamptz)
FROM (
SELECT
unnest(sqlc.arg('id')::uuid[]) AS id,
unnest(sqlc.arg('connect_time')::timestamptz[]) AS connect_time,
unnest(sqlc.arg('organization_id')::uuid[]) AS organization_id,
unnest(sqlc.arg('workspace_owner_id')::uuid[]) AS workspace_owner_id,
unnest(sqlc.arg('workspace_id')::uuid[]) AS workspace_id,
unnest(sqlc.arg('workspace_name')::text[]) AS workspace_name,
unnest(sqlc.arg('agent_name')::text[]) AS agent_name,
unnest(sqlc.arg('type')::connection_type[]) AS type,
unnest(sqlc.arg('code')::int4[]) AS code,
unnest(sqlc.arg('code_valid')::bool[]) AS code_valid,
unnest(sqlc.arg('ip')::inet[]) AS ip,
unnest(sqlc.arg('user_agent')::text[]) AS user_agent,
unnest(sqlc.arg('user_id')::uuid[]) AS user_id,
unnest(sqlc.arg('slug_or_port')::text[]) AS slug_or_port,
unnest(sqlc.arg('connection_id')::uuid[]) AS connection_id,
unnest(sqlc.arg('disconnect_reason')::text[]) AS disconnect_reason,
unnest(sqlc.arg('disconnect_time')::timestamptz[]) AS disconnect_time
) AS u
ON CONFLICT (connection_id, workspace_id, agent_name)
DO UPDATE SET
-- No-op if the connection is still open.
disconnect_time = CASE
WHEN @connection_status::connection_status = 'disconnected'
-- Can only be set once
AND connection_logs.disconnect_time IS NULL
THEN EXCLUDED.connect_time
ELSE connection_logs.disconnect_time
END,
disconnect_reason = CASE
WHEN @connection_status::connection_status = 'disconnected'
-- Can only be set once
AND connection_logs.disconnect_reason IS NULL
THEN EXCLUDED.disconnect_reason
ELSE connection_logs.disconnect_reason
END,
code = CASE
WHEN @connection_status::connection_status = 'disconnected'
-- Can only be set once
AND connection_logs.code IS NULL
THEN EXCLUDED.code
ELSE connection_logs.code
END
RETURNING *;
-- Pick the earliest real connect_time. The zero sentinel
-- ('0001-01-01') means the batch didn't know the connect_time
-- (e.g. a pure disconnect event), so we keep the existing value.
connect_time = CASE
WHEN EXCLUDED.connect_time = '0001-01-01 00:00:00Z'::timestamptz
THEN connection_logs.connect_time
WHEN connection_logs.connect_time = '0001-01-01 00:00:00Z'::timestamptz
THEN EXCLUDED.connect_time
ELSE LEAST(connection_logs.connect_time, EXCLUDED.connect_time)
END,
disconnect_time = CASE
WHEN connection_logs.disconnect_time IS NULL
THEN EXCLUDED.disconnect_time
ELSE connection_logs.disconnect_time
END,
disconnect_reason = CASE
WHEN connection_logs.disconnect_reason IS NULL
THEN EXCLUDED.disconnect_reason
ELSE connection_logs.disconnect_reason
END,
code = CASE
WHEN connection_logs.code IS NULL
THEN EXCLUDED.code
ELSE connection_logs.code
END;