chore: refactor agent stats streaming (#5112)

This commit is contained in:
Colin Adler
2022-11-18 16:46:53 -06:00
committed by GitHub
parent 13a4cfa670
commit ae38bbeab6
20 changed files with 534 additions and 310 deletions

View File

@ -30,31 +30,31 @@ func New() database.Store {
return &fakeQuerier{
mutex: &sync.RWMutex{},
data: &data{
apiKeys: make([]database.APIKey, 0),
agentStats: make([]database.AgentStat, 0),
organizationMembers: make([]database.OrganizationMember, 0),
organizations: make([]database.Organization, 0),
users: make([]database.User, 0),
gitAuthLinks: make([]database.GitAuthLink, 0),
groups: make([]database.Group, 0),
groupMembers: make([]database.GroupMember, 0),
auditLogs: make([]database.AuditLog, 0),
files: make([]database.File, 0),
gitSSHKey: make([]database.GitSSHKey, 0),
parameterSchemas: make([]database.ParameterSchema, 0),
parameterValues: make([]database.ParameterValue, 0),
provisionerDaemons: make([]database.ProvisionerDaemon, 0),
provisionerJobAgents: make([]database.WorkspaceAgent, 0),
provisionerJobLogs: make([]database.ProvisionerJobLog, 0),
provisionerJobResources: make([]database.WorkspaceResource, 0),
provisionerJobResourceMetadata: make([]database.WorkspaceResourceMetadatum, 0),
provisionerJobs: make([]database.ProvisionerJob, 0),
templateVersions: make([]database.TemplateVersion, 0),
templates: make([]database.Template, 0),
workspaceBuilds: make([]database.WorkspaceBuild, 0),
workspaceApps: make([]database.WorkspaceApp, 0),
workspaces: make([]database.Workspace, 0),
licenses: make([]database.License, 0),
apiKeys: make([]database.APIKey, 0),
agentStats: make([]database.AgentStat, 0),
organizationMembers: make([]database.OrganizationMember, 0),
organizations: make([]database.Organization, 0),
users: make([]database.User, 0),
gitAuthLinks: make([]database.GitAuthLink, 0),
groups: make([]database.Group, 0),
groupMembers: make([]database.GroupMember, 0),
auditLogs: make([]database.AuditLog, 0),
files: make([]database.File, 0),
gitSSHKey: make([]database.GitSSHKey, 0),
parameterSchemas: make([]database.ParameterSchema, 0),
parameterValues: make([]database.ParameterValue, 0),
provisionerDaemons: make([]database.ProvisionerDaemon, 0),
workspaceAgents: make([]database.WorkspaceAgent, 0),
provisionerJobLogs: make([]database.ProvisionerJobLog, 0),
workspaceResources: make([]database.WorkspaceResource, 0),
workspaceResourceMetadata: make([]database.WorkspaceResourceMetadatum, 0),
provisionerJobs: make([]database.ProvisionerJob, 0),
templateVersions: make([]database.TemplateVersion, 0),
templates: make([]database.Template, 0),
workspaceBuilds: make([]database.WorkspaceBuild, 0),
workspaceApps: make([]database.WorkspaceApp, 0),
workspaces: make([]database.Workspace, 0),
licenses: make([]database.License, 0),
},
}
}
@ -89,28 +89,28 @@ type data struct {
userLinks []database.UserLink
// New tables
agentStats []database.AgentStat
auditLogs []database.AuditLog
files []database.File
gitAuthLinks []database.GitAuthLink
gitSSHKey []database.GitSSHKey
groups []database.Group
groupMembers []database.GroupMember
parameterSchemas []database.ParameterSchema
parameterValues []database.ParameterValue
provisionerDaemons []database.ProvisionerDaemon
provisionerJobAgents []database.WorkspaceAgent
provisionerJobLogs []database.ProvisionerJobLog
provisionerJobResources []database.WorkspaceResource
provisionerJobResourceMetadata []database.WorkspaceResourceMetadatum
provisionerJobs []database.ProvisionerJob
templateVersions []database.TemplateVersion
templates []database.Template
workspaceBuilds []database.WorkspaceBuild
workspaceApps []database.WorkspaceApp
workspaces []database.Workspace
licenses []database.License
replicas []database.Replica
agentStats []database.AgentStat
auditLogs []database.AuditLog
files []database.File
gitAuthLinks []database.GitAuthLink
gitSSHKey []database.GitSSHKey
groupMembers []database.GroupMember
groups []database.Group
licenses []database.License
parameterSchemas []database.ParameterSchema
parameterValues []database.ParameterValue
provisionerDaemons []database.ProvisionerDaemon
provisionerJobLogs []database.ProvisionerJobLog
provisionerJobs []database.ProvisionerJob
replicas []database.Replica
templateVersions []database.TemplateVersion
templates []database.Template
workspaceAgents []database.WorkspaceAgent
workspaceApps []database.WorkspaceApp
workspaceBuilds []database.WorkspaceBuild
workspaceResourceMetadata []database.WorkspaceResourceMetadatum
workspaceResources []database.WorkspaceResource
workspaces []database.Workspace
deploymentID string
derpMeshKey string
@ -942,6 +942,52 @@ func (q *fakeQuerier) GetWorkspaceByID(_ context.Context, id uuid.UUID) (databas
return database.Workspace{}, sql.ErrNoRows
}
func (q *fakeQuerier) GetWorkspaceByAgentID(_ context.Context, agentID uuid.UUID) (database.Workspace, error) {
q.mutex.RLock()
defer q.mutex.RUnlock()
var agent database.WorkspaceAgent
for _, _agent := range q.workspaceAgents {
if _agent.ID == agentID {
agent = _agent
break
}
}
if agent.ID == uuid.Nil {
return database.Workspace{}, sql.ErrNoRows
}
var resource database.WorkspaceResource
for _, _resource := range q.workspaceResources {
if _resource.ID == agent.ResourceID {
resource = _resource
break
}
}
if resource.ID == uuid.Nil {
return database.Workspace{}, sql.ErrNoRows
}
var build database.WorkspaceBuild
for _, _build := range q.workspaceBuilds {
if _build.JobID == resource.JobID {
build = _build
break
}
}
if build.ID == uuid.Nil {
return database.Workspace{}, sql.ErrNoRows
}
for _, workspace := range q.workspaces {
if workspace.ID == build.WorkspaceID {
return workspace, nil
}
}
return database.Workspace{}, sql.ErrNoRows
}
func (q *fakeQuerier) GetWorkspaceByOwnerIDAndName(_ context.Context, arg database.GetWorkspaceByOwnerIDAndNameParams) (database.Workspace, error) {
q.mutex.RLock()
defer q.mutex.RUnlock()
@ -1801,8 +1847,8 @@ func (q *fakeQuerier) GetWorkspaceAgentByAuthToken(_ context.Context, authToken
defer q.mutex.RUnlock()
// The schema sorts this by created at, so we iterate the array backwards.
for i := len(q.provisionerJobAgents) - 1; i >= 0; i-- {
agent := q.provisionerJobAgents[i]
for i := len(q.workspaceAgents) - 1; i >= 0; i-- {
agent := q.workspaceAgents[i]
if agent.AuthToken == authToken {
return agent, nil
}
@ -1815,8 +1861,8 @@ func (q *fakeQuerier) GetWorkspaceAgentByID(_ context.Context, id uuid.UUID) (da
defer q.mutex.RUnlock()
// The schema sorts this by created at, so we iterate the array backwards.
for i := len(q.provisionerJobAgents) - 1; i >= 0; i-- {
agent := q.provisionerJobAgents[i]
for i := len(q.workspaceAgents) - 1; i >= 0; i-- {
agent := q.workspaceAgents[i]
if agent.ID == id {
return agent, nil
}
@ -1829,8 +1875,8 @@ func (q *fakeQuerier) GetWorkspaceAgentByInstanceID(_ context.Context, instanceI
defer q.mutex.RUnlock()
// The schema sorts this by created at, so we iterate the array backwards.
for i := len(q.provisionerJobAgents) - 1; i >= 0; i-- {
agent := q.provisionerJobAgents[i]
for i := len(q.workspaceAgents) - 1; i >= 0; i-- {
agent := q.workspaceAgents[i]
if agent.AuthInstanceID.Valid && agent.AuthInstanceID.String == instanceID {
return agent, nil
}
@ -1843,7 +1889,7 @@ func (q *fakeQuerier) GetWorkspaceAgentsByResourceIDs(_ context.Context, resourc
defer q.mutex.RUnlock()
workspaceAgents := make([]database.WorkspaceAgent, 0)
for _, agent := range q.provisionerJobAgents {
for _, agent := range q.workspaceAgents {
for _, resourceID := range resourceIDs {
if agent.ResourceID != resourceID {
continue
@ -1859,7 +1905,7 @@ func (q *fakeQuerier) GetWorkspaceAgentsCreatedAfter(_ context.Context, after ti
defer q.mutex.RUnlock()
workspaceAgents := make([]database.WorkspaceAgent, 0)
for _, agent := range q.provisionerJobAgents {
for _, agent := range q.workspaceAgents {
if agent.CreatedAt.After(after) {
workspaceAgents = append(workspaceAgents, agent)
}
@ -1913,7 +1959,7 @@ func (q *fakeQuerier) GetWorkspaceResourceByID(_ context.Context, id uuid.UUID)
q.mutex.RLock()
defer q.mutex.RUnlock()
for _, resource := range q.provisionerJobResources {
for _, resource := range q.workspaceResources {
if resource.ID == id {
return resource, nil
}
@ -1926,7 +1972,7 @@ func (q *fakeQuerier) GetWorkspaceResourcesByJobID(_ context.Context, jobID uuid
defer q.mutex.RUnlock()
resources := make([]database.WorkspaceResource, 0)
for _, resource := range q.provisionerJobResources {
for _, resource := range q.workspaceResources {
if resource.JobID != jobID {
continue
}
@ -1940,7 +1986,7 @@ func (q *fakeQuerier) GetWorkspaceResourcesByJobIDs(_ context.Context, jobIDs []
defer q.mutex.RUnlock()
resources := make([]database.WorkspaceResource, 0)
for _, resource := range q.provisionerJobResources {
for _, resource := range q.workspaceResources {
for _, jobID := range jobIDs {
if resource.JobID != jobID {
continue
@ -1956,7 +2002,7 @@ func (q *fakeQuerier) GetWorkspaceResourcesCreatedAfter(_ context.Context, after
defer q.mutex.RUnlock()
resources := make([]database.WorkspaceResource, 0)
for _, resource := range q.provisionerJobResources {
for _, resource := range q.workspaceResources {
if resource.CreatedAt.After(after) {
resources = append(resources, resource)
}
@ -1978,7 +2024,7 @@ func (q *fakeQuerier) GetWorkspaceResourceMetadataCreatedAfter(ctx context.Conte
defer q.mutex.RUnlock()
metadata := make([]database.WorkspaceResourceMetadatum, 0)
for _, m := range q.provisionerJobResourceMetadata {
for _, m := range q.workspaceResourceMetadata {
_, ok := resourceIDs[m.WorkspaceResourceID]
if !ok {
continue
@ -1993,7 +2039,7 @@ func (q *fakeQuerier) GetWorkspaceResourceMetadataByResourceID(_ context.Context
defer q.mutex.RUnlock()
metadata := make([]database.WorkspaceResourceMetadatum, 0)
for _, metadatum := range q.provisionerJobResourceMetadata {
for _, metadatum := range q.workspaceResourceMetadata {
if metadatum.WorkspaceResourceID == id {
metadata = append(metadata, metadatum)
}
@ -2006,7 +2052,7 @@ func (q *fakeQuerier) GetWorkspaceResourceMetadataByResourceIDs(_ context.Contex
defer q.mutex.RUnlock()
metadata := make([]database.WorkspaceResourceMetadatum, 0)
for _, metadatum := range q.provisionerJobResourceMetadata {
for _, metadatum := range q.workspaceResourceMetadata {
for _, id := range ids {
if metadatum.WorkspaceResourceID == id {
metadata = append(metadata, metadatum)
@ -2319,7 +2365,7 @@ func (q *fakeQuerier) InsertWorkspaceAgent(_ context.Context, arg database.Inser
TroubleshootingURL: arg.TroubleshootingURL,
}
q.provisionerJobAgents = append(q.provisionerJobAgents, agent)
q.workspaceAgents = append(q.workspaceAgents, agent)
return agent, nil
}
@ -2339,7 +2385,7 @@ func (q *fakeQuerier) InsertWorkspaceResource(_ context.Context, arg database.In
Icon: arg.Icon,
DailyCost: arg.DailyCost,
}
q.provisionerJobResources = append(q.provisionerJobResources, resource)
q.workspaceResources = append(q.workspaceResources, resource)
return resource, nil
}
@ -2354,7 +2400,7 @@ func (q *fakeQuerier) InsertWorkspaceResourceMetadata(_ context.Context, arg dat
Value: arg.Value,
Sensitive: arg.Sensitive,
}
q.provisionerJobResourceMetadata = append(q.provisionerJobResourceMetadata, metadatum)
q.workspaceResourceMetadata = append(q.workspaceResourceMetadata, metadatum)
return metadatum, nil
}
@ -2681,7 +2727,7 @@ func (q *fakeQuerier) UpdateWorkspaceAgentConnectionByID(_ context.Context, arg
q.mutex.Lock()
defer q.mutex.Unlock()
for index, agent := range q.provisionerJobAgents {
for index, agent := range q.workspaceAgents {
if agent.ID != arg.ID {
continue
}
@ -2689,7 +2735,7 @@ func (q *fakeQuerier) UpdateWorkspaceAgentConnectionByID(_ context.Context, arg
agent.LastConnectedAt = arg.LastConnectedAt
agent.DisconnectedAt = arg.DisconnectedAt
agent.UpdatedAt = arg.UpdatedAt
q.provisionerJobAgents[index] = agent
q.workspaceAgents[index] = agent
return nil
}
return sql.ErrNoRows
@ -2699,13 +2745,13 @@ func (q *fakeQuerier) UpdateWorkspaceAgentVersionByID(_ context.Context, arg dat
q.mutex.Lock()
defer q.mutex.Unlock()
for index, agent := range q.provisionerJobAgents {
for index, agent := range q.workspaceAgents {
if agent.ID != arg.ID {
continue
}
agent.Version = arg.Version
q.provisionerJobAgents[index] = agent
q.workspaceAgents[index] = agent
return nil
}
return sql.ErrNoRows

View File

@ -113,6 +113,7 @@ type sqlcQuerier interface {
GetWorkspaceBuildByWorkspaceIDAndBuildNumber(ctx context.Context, arg GetWorkspaceBuildByWorkspaceIDAndBuildNumberParams) (WorkspaceBuild, error)
GetWorkspaceBuildsByWorkspaceID(ctx context.Context, arg GetWorkspaceBuildsByWorkspaceIDParams) ([]WorkspaceBuild, error)
GetWorkspaceBuildsCreatedAfter(ctx context.Context, createdAt time.Time) ([]WorkspaceBuild, error)
GetWorkspaceByAgentID(ctx context.Context, agentID uuid.UUID) (Workspace, error)
GetWorkspaceByID(ctx context.Context, id uuid.UUID) (Workspace, error)
GetWorkspaceByOwnerIDAndName(ctx context.Context, arg GetWorkspaceByOwnerIDAndNameParams) (Workspace, error)
GetWorkspaceCountByUserID(ctx context.Context, ownerID uuid.UUID) (int64, error)

View File

@ -17,7 +17,7 @@ import (
)
const deleteOldAgentStats = `-- name: DeleteOldAgentStats :exec
DELETE FROM AGENT_STATS WHERE created_at < now() - interval '30 days'
DELETE FROM agent_stats WHERE created_at < NOW() - INTERVAL '30 days'
`
func (q *sqlQuerier) DeleteOldAgentStats(ctx context.Context) error {
@ -45,16 +45,17 @@ func (q *sqlQuerier) GetLatestAgentStat(ctx context.Context, agentID uuid.UUID)
}
const getTemplateDAUs = `-- name: GetTemplateDAUs :many
select
SELECT
(created_at at TIME ZONE 'UTC')::date as date,
user_id
from
FROM
agent_stats
where template_id = $1
group by
WHERE
template_id = $1
GROUP BY
date, user_id
order by
date asc
ORDER BY
date ASC
`
type GetTemplateDAUsRow struct {
@ -6145,6 +6146,55 @@ func (q *sqlQuerier) InsertWorkspaceResourceMetadata(ctx context.Context, arg In
return i, err
}
const getWorkspaceByAgentID = `-- name: GetWorkspaceByAgentID :one
SELECT
id, created_at, updated_at, owner_id, organization_id, template_id, deleted, name, autostart_schedule, ttl, last_used_at
FROM
workspaces
WHERE
workspaces.id = (
SELECT
workspace_id
FROM
workspace_builds
WHERE
workspace_builds.job_id = (
SELECT
job_id
FROM
workspace_resources
WHERE
workspace_resources.id = (
SELECT
resource_id
FROM
workspace_agents
WHERE
workspace_agents.id = $1
)
)
)
`
func (q *sqlQuerier) GetWorkspaceByAgentID(ctx context.Context, agentID uuid.UUID) (Workspace, error) {
row := q.db.QueryRowContext(ctx, getWorkspaceByAgentID, agentID)
var i Workspace
err := row.Scan(
&i.ID,
&i.CreatedAt,
&i.UpdatedAt,
&i.OwnerID,
&i.OrganizationID,
&i.TemplateID,
&i.Deleted,
&i.Name,
&i.AutostartSchedule,
&i.Ttl,
&i.LastUsedAt,
)
return i, err
}
const getWorkspaceByID = `-- name: GetWorkspaceByID :one
SELECT
id, created_at, updated_at, owner_id, organization_id, template_id, deleted, name, autostart_schedule, ttl, last_used_at

View File

@ -16,16 +16,17 @@ VALUES
SELECT * FROM agent_stats WHERE agent_id = $1 ORDER BY created_at DESC LIMIT 1;
-- name: GetTemplateDAUs :many
select
SELECT
(created_at at TIME ZONE 'UTC')::date as date,
user_id
from
FROM
agent_stats
where template_id = $1
group by
WHERE
template_id = $1
GROUP BY
date, user_id
order by
date asc;
ORDER BY
date ASC;
-- name: DeleteOldAgentStats :exec
DELETE FROM AGENT_STATS WHERE created_at < now() - interval '30 days';
DELETE FROM agent_stats WHERE created_at < NOW() - INTERVAL '30 days';

View File

@ -8,6 +8,35 @@ WHERE
LIMIT
1;
-- name: GetWorkspaceByAgentID :one
SELECT
*
FROM
workspaces
WHERE
workspaces.id = (
SELECT
workspace_id
FROM
workspace_builds
WHERE
workspace_builds.job_id = (
SELECT
job_id
FROM
workspace_resources
WHERE
workspace_resources.id = (
SELECT
resource_id
FROM
workspace_agents
WHERE
workspace_agents.id = @agent_id
)
)
);
-- name: GetWorkspaces :many
SELECT
workspaces.*, COUNT(*) OVER () as count