@@ -51,6 +51,65 @@ func (i *WorkspaceInstance) TableName() string {
5151 return "d_b_workspace_instance"
5252}
5353
54+ // FindStoppedWorkspaceInstancesInRange finds WorkspaceInstanceForUsage that have been stopped between from (inclusive) and to (exclusive).
55+ func FindStoppedWorkspaceInstancesInRange (ctx context.Context , conn * gorm.DB , from , to time.Time ) ([]WorkspaceInstanceForUsage , error ) {
56+ var instances []WorkspaceInstanceForUsage
57+ var instancesInBatch []WorkspaceInstanceForUsage
58+
59+ tx := queryWorkspaceInstanceForUsage (ctx , conn ).
60+ Where ("wsi.stoppingTime >= ?" , TimeToISO8601 (from )).
61+ Where ("wsi.stoppingTime < ?" , TimeToISO8601 (to )).
62+ Where ("wsi.stoppingTime != ?" , "" ).
63+ Where ("wsi.usageAttributionId != ?" , "" ).
64+ FindInBatches (& instancesInBatch , 1000 , func (_ * gorm.DB , _ int ) error {
65+ instances = append (instances , instancesInBatch ... )
66+ return nil
67+ })
68+ if tx .Error != nil {
69+ return nil , fmt .Errorf ("failed to list workspace instances: %w" , tx .Error )
70+ }
71+
72+ return instances , nil
73+ }
74+
75+ // FindRunningWorkspaceInstances finds WorkspaceInstanceForUsage that are running at the point in time the querty is executed.
76+ func FindRunningWorkspaceInstances (ctx context.Context , conn * gorm.DB ) ([]WorkspaceInstanceForUsage , error ) {
77+ var instances []WorkspaceInstanceForUsage
78+ var instancesInBatch []WorkspaceInstanceForUsage
79+
80+ tx := queryWorkspaceInstanceForUsage (ctx , conn ).
81+ Where ("wsi.stoppingTime = ?" , "" ).
82+ Where ("wsi.usageAttributionId != ?" , "" ).
83+ FindInBatches (& instancesInBatch , 1000 , func (_ * gorm.DB , _ int ) error {
84+ instances = append (instances , instancesInBatch ... )
85+ return nil
86+ })
87+ if tx .Error != nil {
88+ return nil , fmt .Errorf ("failed to list workspace instances: %w" , tx .Error )
89+ }
90+
91+ return instances , nil
92+ }
93+
94+ // FindWorkspaceInstancesByIds finds WorkspaceInstanceForUsage by Id.
95+ func FindWorkspaceInstancesByIds (ctx context.Context , conn * gorm.DB , workspaceInstanceIds []uuid.UUID ) ([]WorkspaceInstanceForUsage , error ) {
96+ var instances []WorkspaceInstanceForUsage
97+ var instancesInBatch []WorkspaceInstanceForUsage
98+
99+ tx := queryWorkspaceInstanceForUsage (ctx , conn ).
100+ Where ("wsi.id in ?" , workspaceInstanceIds ).
101+ Where ("wsi.usageAttributionId != ?" , "" ).
102+ FindInBatches (& instancesInBatch , 1000 , func (_ * gorm.DB , _ int ) error {
103+ instances = append (instances , instancesInBatch ... )
104+ return nil
105+ })
106+ if tx .Error != nil {
107+ return nil , fmt .Errorf ("failed to list workspace instances: %w" , tx .Error )
108+ }
109+
110+ return instances , nil
111+ }
112+
54113// ListWorkspaceInstancesInRange lists WorkspaceInstances between from (inclusive) and to (exclusive).
55114// This results in all instances which have existed in the specified period, regardless of their current status, this includes:
56115// - terminated
@@ -61,26 +120,11 @@ func ListWorkspaceInstancesInRange(ctx context.Context, conn *gorm.DB, from, to
61120 var instances []WorkspaceInstanceForUsage
62121 var instancesInBatch []WorkspaceInstanceForUsage
63122
64- tx := conn .WithContext (ctx ).
65- Table (fmt .Sprintf ("%s as wsi" , (& WorkspaceInstance {}).TableName ())).
66- Select ("wsi.id as id, " +
67- "ws.projectId as projectId, " +
68- "ws.type as workspaceType, " +
69- "wsi.workspaceClass as workspaceClass, " +
70- "wsi.usageAttributionId as usageAttributionId, " +
71- "wsi.creationTime as creationTime, " +
72- "wsi.startedTime as startedTime, " +
73- "wsi.stoppingTime as stoppingTime, " +
74- "wsi.stoppedTime as stoppedTime, " +
75- "ws.ownerId as ownerId, " +
76- "ws.id as workspaceId" ,
77- ).
78- Joins (fmt .Sprintf ("LEFT JOIN %s AS ws ON wsi.workspaceId = ws.id" , (& Workspace {}).TableName ())).
123+ tx := queryWorkspaceInstanceForUsage (ctx , conn ).
79124 Where (
80125 conn .Where ("wsi.stoppingTime >= ?" , TimeToISO8601 (from )).Or ("wsi.stoppingTime = ?" , "" ),
81126 ).
82- Where ("wsi.creationTime < ?" , TimeToISO8601 (to )).
83- Where ("wsi.startedTime != ?" , "" ).
127+ Where ("wsi.startedTime < ?" , TimeToISO8601 (to )).
84128 Where ("wsi.usageAttributionId != ?" , "" ).
85129 FindInBatches (& instancesInBatch , 1000 , func (_ * gorm.DB , _ int ) error {
86130 instances = append (instances , instancesInBatch ... )
@@ -93,6 +137,24 @@ func ListWorkspaceInstancesInRange(ctx context.Context, conn *gorm.DB, from, to
93137 return instances , nil
94138}
95139
140+ func queryWorkspaceInstanceForUsage (ctx context.Context , conn * gorm.DB ) * gorm.DB {
141+ return conn .WithContext (ctx ).
142+ Table (fmt .Sprintf ("%s as wsi" , (& WorkspaceInstance {}).TableName ())).
143+ Select ("wsi.id as id, " +
144+ "ws.projectId as projectId, " +
145+ "ws.type as workspaceType, " +
146+ "wsi.workspaceClass as workspaceClass, " +
147+ "wsi.usageAttributionId as usageAttributionId, " +
148+ "wsi.creationTime as creationTime, " +
149+ "wsi.startedTime as startedTime, " +
150+ "wsi.stoppingTime as stoppingTime, " +
151+ "wsi.stoppedTime as stoppedTime, " +
152+ "ws.ownerId as ownerId, " +
153+ "ws.id as workspaceId" ,
154+ ).
155+ Joins (fmt .Sprintf ("LEFT JOIN %s AS ws ON wsi.workspaceId = ws.id" , (& Workspace {}).TableName ()))
156+ }
157+
96158const (
97159 AttributionEntity_User = "user"
98160 AttributionEntity_Team = "team"
0 commit comments